vllm_omni.diffusion.attention.layer ¶

allow_fp32_fallback = allow_fp32_fallback

attention = self.attn_impl_cls(
    num_heads=num_heads,
    head_size=head_size,
    softmax_scale=softmax_scale,
    causal=causal,
    num_kv_heads=num_kv_heads,
    qkv_layout=qkv_layout,
    prefix=prefix,
    backend_kwargs=backend_kwargs,
    role=role,
    backend_explicit=self.backend_explicit,
)

attn_backend: type[AttentionBackend] | None = (
    attn_backend_cls
)

attn_impl_cls = self.attn_backend.get_impl_cls()

attn_spec = spec

backend_explicit = False

backend_pref = None

causal = causal

gather_idx = gather_idx

head_size = head_size

layer_idx: int | None = _try_extract_layer_index(prefix)

num_heads = num_heads

num_kv_heads = (
    num_kv_heads if num_kv_heads is not None else num_heads
)

paged_kv_cache_dtype = paged_kv_cache_dtype

paged_kv_cache_role = paged_kv_cache_role

parallel_strategy = build_parallel_attention_strategy(
    scatter_idx=scatter_idx,
    gather_idx=gather_idx,
    use_sync=use_sync,
    causal=causal,
)

prefix = prefix

qkv_layout = qkv_layout

ring_pg = None

ring_runner = None

role = role

role_category = role_category

scatter_idx = scatter_idx

sdpa_fallback: AttentionImpl | None = (
    SDPABackend.get_impl_cls()(
        num_heads=num_heads,
        head_size=head_size,
        softmax_scale=softmax_scale,
        causal=causal,
        num_kv_heads=num_kv_heads,
        qkv_layout=qkv_layout,
    )
)

skip_sequence_parallel = skip_sequence_parallel

softmax_scale = softmax_scale

use_ring = False

use_sync = use_sync