Skip to content

vllm_omni.diffusion.attention.layer

logger module-attribute 

logger = init_logger(__name__)

Attention

Bases: Module

allow_fp32_fallback instance-attribute 

allow_fp32_fallback = allow_fp32_fallback

attention instance-attribute 

attention: AttentionImpl | Module

attn_backend instance-attribute 

attn_backend: type[AttentionBackend] | None

attn_impl_cls instance-attribute 

attn_impl_cls = self.attn_backend.get_impl_cls()

attn_spec instance-attribute 

attn_spec = spec

backend_explicit instance-attribute 

backend_explicit = False

backend_pref instance-attribute 

backend_pref = None

causal instance-attribute 

causal = causal

gather_idx instance-attribute 

gather_idx = gather_idx

head_size instance-attribute 

head_size = head_size

layer_idx instance-attribute 

layer_idx: int | None = _try_extract_layer_index(prefix)

num_heads instance-attribute 

num_heads = num_heads

num_kv_heads instance-attribute 

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

paged_kv_cache_dtype instance-attribute 

paged_kv_cache_dtype = paged_kv_cache_dtype

paged_kv_cache_role instance-attribute 

paged_kv_cache_role = paged_kv_cache_role

parallel_strategy instance-attribute 

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

prefix instance-attribute 

prefix = prefix

qkv_layout instance-attribute 

qkv_layout = qkv_layout

ring_pg instance-attribute 

ring_pg = None

ring_runner instance-attribute 

ring_runner = None

role instance-attribute 

role = role

role_category instance-attribute 

role_category = role_category

scatter_idx instance-attribute 

scatter_idx = scatter_idx

sdpa_fallback instance-attribute 

sdpa_fallback: SDPAImpl | None

skip_sequence_parallel instance-attribute 

skip_sequence_parallel = skip_sequence_parallel

softmax_scale instance-attribute 

softmax_scale = softmax_scale

use_ring instance-attribute 

use_ring = False

use_sync instance-attribute 

use_sync = use_sync

forward 

forward(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    attn_metadata: AttentionMetadata | None = None,
) -> Tensor

get_kv_cache_spec 

get_kv_cache_spec(
    vllm_config: VllmConfig,
) -> KVCacheSpec | None

Return native rank-local geometry for an opted-in paged cache.

is_paged_kv_active 

is_paged_kv_active() -> bool

Return whether this layer will use Scheduler-managed paged KV.

resolve_execution_path 

resolve_execution_path(
    context: ExecutionContext,
    query: Tensor,
    key: Tensor,
    value: Tensor,
    attn_metadata: AttentionMetadata | None,
) -> ExecutionPathResult

Compose backend capabilities with outer Attention boundaries.