Skip to content

vllm_omni.diffusion.attention.selector

Diffusion attention backend selector.

This module resolves diffusion attention backends from: 1. per-role AttentionConfig 2. platform default

logger module-attribute 

logger = init_logger(__name__)

get_attn_backend_for_role 

get_attn_backend_for_role(
    role: str,
    head_size: int,
    attention_config: AttentionConfig | None = None,
    role_category: str | None = None,
) -> tuple[type[AttentionBackend], AttentionSpec | None]

Get attention backend for a specific attention role.

Lookup precedence
  1. attention_config.per_role[role] — exact match
  2. attention_config.per_role[role_category] — category fallback
  3. attention_config.default — global default
  4. Platform default — hardware-specific

Parameters:

Name Type Description Default
role str

Attention role string (e.g. "self", "cross", "joint", "ltx2.audio_to_video")

 required
head_size int

Head size for attention computation

 required
attention_config AttentionConfig | None

The AttentionConfig from OmniDiffusionConfig. If None, falls back to platform default behavior.

 None
role_category str | None

Optional category for fallback (e.g. "cross" for "ltx2.audio_to_video")

 None

Returns:

Type Description
type[AttentionBackend]

Tuple of (backend_class, AttentionSpec or None).
AttentionSpec | None

AttentionSpec is None when using platform default without explicit config.