Skip to content

vllm_omni.diffusion.diffusion_kv.paged_attention_adapter

DiffusionKVRowResolver module-attribute 

DiffusionKVRowResolver = Callable[
    [str, int | None, str | None],
    DiffusionPagedAttentionRowBinding,
]

DiffusionPagedAttentionAdapter

Translate diffusion rows into vLLM metadata for an Omni paged backend.

attn_groups instance-attribute 

attn_groups = attn_groups

block_tables instance-attribute 

block_tables = block_tables

device instance-attribute 

device = torch.device(device)

kv_cache_config instance-attribute 

kv_cache_config = kv_cache_config

layers instance-attribute 

layers = dict(layers)

resolve_row instance-attribute 

resolve_row = resolve_row

vllm_config instance-attribute 

vllm_config = vllm_config

activate 

activate(
    batch: PreparedDiffusionPagedAttentionBatch,
) -> Iterator[DiffusionPagedAttentionAdapter]

invalidate_prepared_batches 

invalidate_prepared_batches() -> None

Invalidate native buffer views after BlockTable state changes.

prepare_batch 

prepare_batch(
    rows: Sequence[DiffusionPagedAttentionRow],
) -> PreparedDiffusionPagedAttentionBatch

prepare_layer_context 

prepare_layer_context(
    layer_name: str,
    query: Tensor,
    key: Tensor,
    value: Tensor,
    *,
    omni_attn_metadata: Any | None = None,
) -> DiffusionPagedAttentionContext

DiffusionPagedAttentionContext dataclass

One layer's native inputs for an Omni paged-backend invocation.

key_write instance-attribute 

key_write: Tensor

layer instance-attribute 

layer: DiffusionPagedAttentionLayerAdapter

native_metadata instance-attribute 

native_metadata: Any

output_scatter_indices class-attribute instance-attribute 

output_scatter_indices: Tensor | None = None

piecewise_native_metadata instance-attribute 

piecewise_native_metadata: tuple[Any, ...]

piecewise_plan instance-attribute 

piecewise_plan: PagedPiecewisePlan | None

query instance-attribute 

query: Tensor

query_has_head_dims instance-attribute 

query_has_head_dims: bool

query_token_shape instance-attribute 

query_token_shape: tuple[int, ...]

slot_mapping instance-attribute 

slot_mapping: Tensor

value_write instance-attribute 

value_write: Tensor

restore_output 

restore_output(output: Tensor) -> Tensor

DiffusionPagedAttentionLayerAdapter

Bases: AttentionLayerBase

Register a diffusion layer with vLLM's native cache machinery.

This object deliberately does not subclass vllm.Attention. The latter owns a second execution path and would bypass Omni's sequence parallel pre/post hooks. The wrapper only supplies the small AttentionLayerBase contract needed by init_attn_backend and keeps the platform-native attention implementation/cache view available to the diffusion adapter.

attn_backend instance-attribute 

attn_backend = attn_backend

head_size instance-attribute 

head_size = int(canonical_spec.head_size)

head_size_v instance-attribute 

head_size_v = int(
    getattr(
        canonical_spec,
        "head_size_v",
        canonical_spec.head_size,
    )
)

impl instance-attribute 

impl = self._create_native_impl(vllm_config)

kv_cache instance-attribute 

kv_cache: Tensor | None = None

layer_name instance-attribute 

layer_name = layer_name

num_heads instance-attribute 

num_heads = num_heads

num_kv_heads instance-attribute 

num_kv_heads = num_kv_heads

softmax_scale instance-attribute 

softmax_scale = float(layer.softmax_scale)

spec instance-attribute 

spec = canonical_spec

get_attn_backend 

get_attn_backend()

get_kv_cache_spec 

get_kv_cache_spec(vllm_config: VllmConfig) -> AttentionSpec

DiffusionPagedAttentionMetadata dataclass

Runner-owned row layouts for one request-level denoise loop.

denoise_rows instance-attribute 

denoise_rows: tuple[DiffusionPagedAttentionRow, ...]

prefill_rows instance-attribute 

prefill_rows: tuple[DiffusionPagedAttentionRow, ...]

rows_for_step 

rows_for_step(
    step_idx: int | None,
) -> tuple[DiffusionPagedAttentionRow, ...]

DiffusionPagedAttentionRow dataclass

One logical BlockTable row participating in a paged attention call.

context_id class-attribute instance-attribute 

context_id: str | None = None

identity property 

identity: tuple[str, int | None, str | None]

kv_start_pos class-attribute instance-attribute 

kv_start_pos: int = 0

query_len instance-attribute 

query_len: int

request_id instance-attribute 

request_id: str

seq_len instance-attribute 

seq_len: int

sequence_id class-attribute instance-attribute 

sequence_id: int | None = None

DiffusionPagedAttentionRowBinding dataclass

Worker row and logical length installed for one allocation identity.

block_ids class-attribute instance-attribute 

block_ids: tuple[tuple[int, ...], ...] = ()

max_seq_len instance-attribute 

max_seq_len: int

row_index instance-attribute 

row_index: int

DiffusionPagedAttentionRuntime

Activate Runner-prepared rows lazily as the denoise loop advances.

adapter instance-attribute 

adapter = adapter

metadata instance-attribute 

metadata = metadata

activate 

activate() -> Iterator[DiffusionPagedAttentionRuntime]

close 

close() -> None

ensure_active 

ensure_active(
    step_idx: int | None,
) -> DiffusionPagedAttentionAdapter

PreparedDiffusionPagedAttentionBatch dataclass

Native metadata shared by all paged attention layers in one forward.

attn_metadata instance-attribute 

attn_metadata: dict[str, Any]

block_tables instance-attribute 

block_tables: tuple[Tensor, ...]

num_tokens instance-attribute 

num_tokens: int

positions instance-attribute 

positions: Tensor

query_start_loc instance-attribute 

query_start_loc: Tensor

row_indices instance-attribute 

row_indices: Tensor

rows instance-attribute 

rows: tuple[DiffusionPagedAttentionRow, ...]

seq_lens instance-attribute 

seq_lens: Tensor

slot_mappings instance-attribute 

slot_mappings: Tensor

slot_mappings_by_layer instance-attribute 

slot_mappings_by_layer: dict[str, Tensor]