vllm_omni.diffusion.attention.layer ¶

attention = 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,
    backend_kwargs=backend_kwargs,
)

attn_backend = attn_backend_cls

attn_impl_cls = get_impl_cls()

backend_pref = None

causal = causal

gather_idx = gather_idx

layer_idx: int | None = _try_extract_layer_index(prefix)

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

qkv_layout = qkv_layout

ring_pg = ring_group

ring_runner = RingParallelAttention(
    sp_group, attn_backend_pref=backend_pref
)

role = role

role_category = role_category

scatter_idx = scatter_idx

sdpa_fallback = 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