Bases: ABC
Abstract class for diffusion attention backends.
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
supports_piecewise_spans: bool = False
class-attribute instance-attribute ¶
supports_prefix_kv_slicing: bool = False
abstractmethod staticmethod ¶
Get the list of supported head sizes for this backend.
classmethod ¶
indexes_kv_by_block_stride() -> bool
Whether this backend reads K/V pages by the runtime block stride.
Returning True means the physical cache layout has num_blocks as its outer stride, so native vLLM may safely use page-size padding when it unifies cache layouts across layers. Dense diffusion backends conservatively keep the default False; a paged backend should override this only when its kernel actually follows that layout.
Since vLLM 0.29 this no longer rides on AttentionSpec. It selects the block-outermost physical layout pinned on CacheConfig.kv_cache_layout -- see vllm_omni.diffusion.diffusion_kv.layout -- whose is_block_outermost is what upstream reads when it sizes KVCacheTensor regions.
classmethod ¶
Return whether this backend can consume a nontrivial attn_mask.
attention_spec is the resolved per-role config when the user picked a backend explicitly. Implementations that depend on kernel variant (for example FlashInfer cute-dsl vs fa2) must consult it so capability probes match a runnable configuration.
classmethod ¶
supports_multi_doc_packed_varlen() -> bool
Whether this backend keeps N-document packed boundaries isolated.
When True, the backend consumes AttentionMetadata.extra cu_seqlens as a genuine block-diagonal attention plan (a dedicated varlen kernel, not a padding-mask rebuild), so a caller may pack multiple real requests into one forward without attention crossing document boundaries. When False, callers packing more than one real document must run one forward per document; otherwise a backend that only supports a [real, pad] two-document contract, or that ignores cu_seqlens outright, will silently attend across request boundaries. Backends whose kernel selection is platform-dependent must consult current_omni_platform.
classmethod ¶
supports_packed_mask_free() -> bool
Whether [real, pad] packed layouts can run without attn_mask.
When True, models that pack a [real, pad] two-document layout and provide AttentionMetadata.packed_padding alongside the packed cu_seqlens/max_seqlen metadata may skip constructing the padding mask entirely. Backends whose mask-free behavior is platform-dependent must check current_omni_platform.
classmethod ¶
Raise if this backend's optional dependencies are missing.
Called during platform resolution, i.e. before model construction, so a backend that probes its kernel package lazily still reports the problem while the user can still act on it.
forward(
query: Tensor,
key: Tensor,
value: Tensor,
attn_metadata: T | None = None,
) -> Tensor
Dispatch to platform-specific forward implementation.
forward_cuda(
query: Tensor,
key: Tensor,
value: Tensor,
attn_metadata: T | None = None,
) -> Tensor
forward_hip(
query: Tensor,
key: Tensor,
value: Tensor,
attn_metadata: T | None = None,
) -> Tensor
forward_musa(
query: Tensor,
key: Tensor,
value: Tensor,
attn_metadata: T | None = None,
) -> Tensor
forward_npu(
query: Tensor,
key: Tensor,
value: Tensor,
attn_metadata: T | None = None,
) -> Tensor
forward_paged(paged_kv_context: Any) -> Tensor
Execute one Worker-prepared paged-KV attention call.
The context is intentionally opaque to the common attention layer. Backends opt in by setting supports_paged_kv on their backend class and implementing this method. Dense callers continue to use forward unchanged.
dataclass ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
packed_padding: PackedPaddingMetadata | None = None
class-attribute instance-attribute ¶
query_ranges: tuple[QueryRange, ...] | None = None
dataclass ¶
Producer-validated mask-free view of padding in a [real, pad] packing.
The cumulative-length tensors are canonical two-element [0, length] views. Consumers may use them without reading device scalars because the producer owns the packing.
dataclass ¶
dataclass ¶
Video-grid slices in a packed multimodal sequence.
A model that packs its sequence as [prefix | t*h*w video rows | padding] publishes this so backends can recover spatiotemporal locality; the prefix holds everything that is not video (text, visual conditions, audio). Publishing it also asserts that any attn_mask masks only the trailing padding, so prefix_len + t*h*w is the used length of the sequence.
prefix_len/latent_grid retain the original one-tail contract. A Ref2VA layout instead publishes used_len and every physical video span, so audio and image rows between videos remain dense.
Plain ints, so reading it never forces a device-to-host sync.