[Diffusion] Add native paged KV backend for diffusion workers (#6102)

Signed-off-by: Acerak01-fy <wfy2003324@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/diffusion_kv/config.py

@@ -37,3 +37,8 @@ def parse_diffusion_kv_cache_mode(value: object) -> DiffusionKVCacheMode:
             "use 'dense_legacy' or 'paged_scheduler'"
         )
     return mode
+
+
+def is_scheduler_paged_kv_mode(mode: DiffusionKVCacheMode) -> bool:
+    """Return whether a parsed cache mode uses Scheduler-owned paging."""
+    return mode is DiffusionKVCacheMode.PAGED_SCHEDULER

vllm_omni/diffusion/diffusion_kv/initialization.py

@@ -73,7 +73,13 @@ def initialize_diffusion_kv_control_plane(
         scheduler_kv_cache_config,
         vllm_config,
     )
-    executor.set_kv_cache_configs(worker_configs)
+    # Native sizing may auto-fit ``max_model_len=-1`` against the profiled
+    # memory budget.  Forward that final value so every Worker sizes its
+    # physical BlockTables from the same bound as the Scheduler.
+    resolved_max_model_len = vllm_config.model_config.max_model_len
+    if type(resolved_max_model_len) is not int or resolved_max_model_len <= 0:
+        raise ValueError("native Diffusion KV sizing must resolve a positive max_model_len")
+    executor.set_kv_cache_configs(worker_configs, resolved_max_model_len)
     return (
         scheduler_kv_cache_config,
         scheduler_block_size,

vllm_omni/diffusion/diffusion_kv/model_runner_backend.py

@@ -0,0 +1,647 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import copy
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.distributed import get_dcp_group
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheSpec
+from vllm.v1.worker import block_table as native_block_table
+from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache
+from vllm.v1.worker.gpu.block_table import BlockTables
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode
+from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
+from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
+    DiffusionPagedAttentionAdapter,
+    DiffusionPagedAttentionLayerAdapter,
+    DiffusionPagedAttentionRow,
+    DiffusionPagedAttentionRowBinding,
+    PreparedDiffusionPagedAttentionBatch,
+)
+from vllm_omni.platforms import current_omni_platform
+
+DiffusionKVIdentity = tuple[str, int | None, str | None]
+DiffusionKVSnapshot = tuple[object, ...]
+
+
+@dataclass(frozen=True)
+class _DiffusionKVRequestState:
+    generation: int
+    snapshot: DiffusionKVSnapshot
+    row_token_lens: tuple[tuple[DiffusionKVIdentity, int], ...]
+
+
+@dataclass(frozen=True)
+class _DiffusionKVRowInstall:
+    identity: DiffusionKVIdentity
+    token_len: int
+    block_ids: tuple[tuple[int, ...], ...]
+
+
+class DiffusionKVModelRunnerBackend:
+    """Native paged-KV state and BlockTable operations for a model runner."""
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        od_config: OmniDiffusionConfig,
+        device: torch.device,
+    ) -> None:
+        self.vllm_config = vllm_config
+        self.od_config = od_config
+        self.device = device
+        self.kv_cache_config: KVCacheConfig | None = None
+        self.kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
+        self.block_tables: BlockTables | None = None
+        self.paged_attention_adapter: DiffusionPagedAttentionAdapter | None = None
+        self._kv_cache_layer_adapters: dict[str, DiffusionPagedAttentionLayerAdapter] = {}
+        self._diffusion_kv_identity_to_row: dict[DiffusionKVIdentity, int] = {}
+        self._diffusion_kv_free_rows: list[int] = []
+        self._diffusion_kv_request_states: dict[str, _DiffusionKVRequestState] = {}
+
+    def register_kv_cache_layers(
+        self,
+        layers: Mapping[str, tuple[Any, KVCacheSpec]],
+    ) -> dict[str, KVCacheSpec]:
+        """Register native adapters and return their canonical cache specs."""
+
+        if self.kv_cache_config is not None:
+            raise RuntimeError("Diffusion KV cache layers cannot be changed after physical initialization")
+
+        if layers:
+            current_omni_platform.get_diffusion_kv_block_tables_cls()
+            if not callable(current_omni_platform.build_diffusion_kv_attn_metadata):
+                raise NotImplementedError(
+                    "Diffusion paged KV requires platform-native BlockTables and attention metadata hooks"
+                )
+
+        forward_context = self.vllm_config.compilation_config.static_forward_context
+        previous_adapters = self._kv_cache_layer_adapters
+        for layer_name, (layer, spec) in layers.items():
+            if not isinstance(spec, AttentionSpec):
+                raise TypeError(
+                    f"Diffusion KV layer {layer_name!r} produced unsupported spec {type(spec).__name__}; "
+                    "only native attention specs are supported"
+                )
+            if not layer.attn_backend.supports_paged_kv:
+                raise NotImplementedError(
+                    f"Diffusion paged KV layer {layer_name!r} requires an Omni backend with paged support; "
+                    f"selected {layer.attn_backend.get_name()}"
+                )
+            existing = forward_context.get(layer_name)
+            if existing is not None and existing is not layer and existing is not previous_adapters.get(layer_name):
+                raise RuntimeError(f"Diffusion KV layer name {layer_name!r} collides with the native forward context")
+
+        previous_use_non_causal = self.vllm_config.attention_config.use_non_causal
+        self.vllm_config.attention_config.use_non_causal = any(
+            bool(getattr(spec, "non_causal", False)) for _, spec in layers.values()
+        )
+        adapters: dict[str, DiffusionPagedAttentionLayerAdapter] = {}
+        try:
+            with set_current_vllm_config(self.vllm_config):
+                for layer_name, (layer, spec) in layers.items():
+                    adapters[layer_name] = DiffusionPagedAttentionLayerAdapter(
+                        layer_name=layer_name,
+                        layer=layer,
+                        spec=spec,
+                        vllm_config=self.vllm_config,
+                        device=self.device,
+                        ulysses_degree=self._get_paged_attention_ulysses_degree(layer),
+                    )
+        except Exception:
+            self.vllm_config.attention_config.use_non_causal = previous_use_non_causal
+            raise
+
+        geometries = {(adapter.num_heads, adapter.num_kv_heads, adapter.head_size) for adapter in adapters.values()}
+        if len(geometries) > 1:
+            self.vllm_config.attention_config.use_non_causal = previous_use_non_causal
+            raise ValueError(
+                "Native paged attention metadata builders require one rank-local attention geometry; "
+                f"got {sorted(geometries)!r}"
+            )
+        attention_geometry = next(iter(geometries)) if geometries else None
+        set_attention_geometry = (
+            getattr(self.vllm_config.model_config, "set_attention_geometry", None)
+            if attention_geometry is not None
+            else None
+        )
+        if attention_geometry is not None and not callable(set_attention_geometry):
+            self.vllm_config.attention_config.use_non_causal = previous_use_non_causal
+            raise RuntimeError(
+                "Diffusion native paged attention requires a model config that exposes set_attention_geometry"
+            )
+
+        updated_forward_context = dict(forward_context)
+        for layer_name in set(layers) | set(previous_adapters):
+            entry = updated_forward_context.get(layer_name)
+            layer_entry = layers.get(layer_name)
+            if entry is previous_adapters.get(layer_name) or (layer_entry is not None and entry is layer_entry[0]):
+                updated_forward_context.pop(layer_name, None)
+        updated_forward_context.update(adapters)
+
+        if attention_geometry is not None:
+            num_heads, num_kv_heads, head_size = attention_geometry
+            try:
+                set_attention_geometry(
+                    num_heads=num_heads,
+                    num_kv_heads=num_kv_heads,
+                    head_size=head_size,
+                )
+            except Exception:
+                self.vllm_config.attention_config.use_non_causal = previous_use_non_causal
+                raise
+
+        forward_context.clear()
+        forward_context.update(updated_forward_context)
+        self._kv_cache_layer_adapters = adapters
+        return {layer_name: adapter.spec for layer_name, adapter in adapters.items()}
+
+    def _get_paged_attention_ulysses_degree(self, layer: Any) -> int:
+        if bool(getattr(layer, "skip_sequence_parallel", False)):
+            return 1
+
+        parallel_config = getattr(self.od_config, "parallel_config", None)
+        ulysses_degree = int(getattr(parallel_config, "ulysses_degree", 1))
+        ring_degree = int(getattr(parallel_config, "ring_degree", 1))
+        allgather_degree = int(getattr(parallel_config, "allgather_degree", 1))
+        ulysses_mode = str(getattr(parallel_config, "ulysses_mode", "strict"))
+        if ring_degree > 1:
+            raise NotImplementedError(
+                "Diffusion paged attention does not support Ring or hybrid Ulysses+Ring attention"
+            )
+        if allgather_degree > 1:
+            raise NotImplementedError("Diffusion paged attention does not support AllGather-KV sequence parallelism")
+        if ulysses_degree > 1 and ulysses_mode != "strict":
+            raise NotImplementedError(
+                f"Diffusion paged attention currently supports only strict Ulysses; got ulysses_mode={ulysses_mode!r}"
+            )
+        return ulysses_degree
+
+    def _get_max_rows_per_request(self) -> int:
+        max_rows = getattr(self.od_config, "diffusion_kv_max_rows_per_request", None)
+        if type(max_rows) is not int or max_rows <= 0:
+            raise ValueError(
+                "paged_scheduler requires a positive integer diffusion_kv_max_rows_per_request adapter limit"
+            )
+        return max_rows
+
+    def _get_max_model_len(self) -> int:
+        max_len = getattr(self.vllm_config.model_config, "max_model_len", None)
+        if type(max_len) is not int or max_len <= 0:
+            raise ValueError("paged_scheduler requires a positive model_config.max_model_len")
+        return max_len
+
+    def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
+        """Allocate and bind native attention KV tensors for this rank."""
+        if self.kv_cache_config is not None:
+            raise RuntimeError("Diffusion native KV cache is already initialized")
+        if not self._kv_cache_layer_adapters:
+            raise RuntimeError("Diffusion KV cache layers must be registered before physical initialization")
+
+        configured_layers = {
+            layer_name for group in kv_cache_config.kv_cache_groups for layer_name in group.layer_names
+        }
+        registered_layers = set(self._kv_cache_layer_adapters)
+        if configured_layers != registered_layers:
+            raise ValueError(
+                "Rank-local Diffusion KVCacheConfig layer mismatch: "
+                f"expected={sorted(registered_layers)}, configured={sorted(configured_layers)}"
+            )
+
+        kv_cache_config = copy.deepcopy(kv_cache_config)
+        max_len = self._get_max_model_len()
+        max_rows_per_request = self._get_max_rows_per_request()
+        scheduler_config = self.vllm_config.scheduler_config
+        max_num_seqs = getattr(scheduler_config, "max_num_seqs", None)
+        if type(max_num_seqs) is not int or max_num_seqs <= 0:
+            raise ValueError("paged_scheduler requires a positive scheduler_config.max_num_seqs")
+
+        kv_cache_groups = kv_cache_config.kv_cache_groups
+        if not kv_cache_groups:
+            raise ValueError("paged_scheduler requires at least one native KV cache group")
+        block_sizes = [group.kv_cache_spec.block_size for group in kv_cache_groups]
+        if any(type(block_size) is not int or block_size <= 0 for block_size in block_sizes):
+            raise ValueError("native diffusion KV cache groups must have positive integer block sizes")
+        unaligned_max_num_blocks_per_group = [
+            group.kv_cache_spec.max_num_blocks_per_req(self.vllm_config, max_len) for group in kv_cache_groups
+        ]
+        if any(type(capacity) is not int or capacity <= 0 for capacity in unaligned_max_num_blocks_per_group):
+            raise ValueError("native diffusion KV cache groups must have positive row capacities")
+        max_num_blocks_per_group = [
+            native_block_table.get_block_table_width(capacity, block_size)
+            for capacity, block_size in zip(unaligned_max_num_blocks_per_group, block_sizes, strict=True)
+        ]
+
+        max_num_reqs = max_num_seqs * max_rows_per_request
+        max_num_batched_tokens = getattr(scheduler_config, "max_num_batched_tokens", None)
+        if type(max_num_batched_tokens) is not int or max_num_batched_tokens <= 0:
+            raise ValueError("scheduler_config.max_num_batched_tokens must be a positive integer")
+        parallel_config = self.vllm_config.parallel_config
+        cp_size = parallel_config.decode_context_parallel_size
+        cp_rank = get_dcp_group().rank_in_group if cp_size > 1 else 0
+        cp_interleave = parallel_config.cp_kv_cache_interleave_size
+
+        # Native metadata builders size per-row buffers from max_num_seqs.
+        # A diffusion public request can occupy several sequence/context rows.
+        scheduler_config.max_num_seqs = max_num_reqs
+        kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
+        previous_adapter_caches = {
+            layer_name: adapter.kv_cache for layer_name, adapter in self._kv_cache_layer_adapters.items()
+        }
+        try:
+            with set_current_vllm_config(self.vllm_config):
+                attn_groups, _, kernel_block_sizes = init_attn_backend(
+                    kv_cache_config,
+                    self.vllm_config,
+                    self.device,
+                )
+                block_tables_cls = current_omni_platform.get_diffusion_kv_block_tables_cls()
+                block_tables = block_tables_cls(
+                    block_sizes=block_sizes,
+                    max_num_reqs=max_num_reqs,
+                    max_num_batched_tokens=max_num_batched_tokens,
+                    max_num_blocks_per_group=max_num_blocks_per_group,
+                    device=self.device,
+                    kernel_block_sizes=kernel_block_sizes,
+                    cp_size=cp_size,
+                    cp_rank=cp_rank,
+                    cp_interleave=cp_interleave,
+                )
+                paged_attention_adapter = DiffusionPagedAttentionAdapter(
+                    vllm_config=self.vllm_config,
+                    device=self.device,
+                    kv_cache_config=kv_cache_config,
+                    block_tables=block_tables,
+                    attn_groups=attn_groups,
+                    layers=self._kv_cache_layer_adapters,
+                    resolve_row=self._resolve_paged_attention_row,
+                )
+                init_kv_cache(
+                    kv_caches,
+                    self.vllm_config.compilation_config.static_forward_context,
+                    kv_cache_config,
+                    attn_groups,
+                    self.device,
+                    self.vllm_config.cache_config.cache_dtype,
+                    kernel_block_sizes,
+                    self.vllm_config,
+                )
+        except Exception:
+            scheduler_config.max_num_seqs = max_num_seqs
+            for layer_name, previous_cache in previous_adapter_caches.items():
+                self._kv_cache_layer_adapters[layer_name].kv_cache = previous_cache
+            raise
+
+        self.kv_caches = kv_caches
+        self.block_tables = block_tables
+        self.paged_attention_adapter = paged_attention_adapter
+        self._diffusion_kv_identity_to_row.clear()
+        self._diffusion_kv_free_rows = list(range(max_num_reqs - 1, -1, -1))
+        self._diffusion_kv_request_states.clear()
+        self.kv_cache_config = kv_cache_config
+
+    def _validate_row(
+        self,
+        *,
+        identity: DiffusionKVIdentity,
+        token_len: int,
+        block_ids: tuple[list[int], ...],
+        seen_identities: set[DiffusionKVIdentity],
+    ) -> _DiffusionKVRowInstall:
+        if identity in seen_identities:
+            raise ValueError(f"Duplicate diffusion KV row identity: {identity!r}")
+        seen_identities.add(identity)
+        if type(token_len) is not int or token_len < 0:
+            raise ValueError(f"Diffusion KV row {identity!r} has invalid token length {token_len!r}")
+
+        assert self.kv_cache_config is not None
+        groups = self.kv_cache_config.kv_cache_groups
+        if not isinstance(block_ids, tuple) or len(block_ids) != len(groups):
+            actual_groups = len(block_ids) if isinstance(block_ids, (tuple, list)) else "invalid"
+            raise ValueError(f"Diffusion KV row {identity!r} has {actual_groups} block groups; expected {len(groups)}")
+
+        max_len = self._get_max_model_len()
+        normalized_groups: list[tuple[int, ...]] = []
+        for group_index, (group_block_ids, kv_cache_group) in enumerate(zip(block_ids, groups, strict=True)):
+            if not isinstance(group_block_ids, list):
+                raise ValueError(f"Diffusion KV row {identity!r} group {group_index} block IDs must be a list")
+            spec = kv_cache_group.kv_cache_spec
+            expected_count = spec.max_num_blocks_per_req(self.vllm_config, token_len)
+            if len(group_block_ids) != expected_count:
+                raise ValueError(
+                    f"Diffusion KV row {identity!r} group {group_index} has {len(group_block_ids)} blocks; "
+                    f"expected {expected_count} for {token_len} tokens"
+                )
+            row_capacity = spec.max_num_blocks_per_req(self.vllm_config, max_len)
+            if len(group_block_ids) > row_capacity:
+                raise ValueError(
+                    f"Diffusion KV row {identity!r} group {group_index} exceeds row capacity {row_capacity}"
+                )
+            normalized_ids: list[int] = []
+            for block_id in group_block_ids:
+                if type(block_id) is not int:
+                    raise ValueError(
+                        f"Diffusion KV row {identity!r} group {group_index} has non-integer block ID {block_id!r}"
+                    )
+                if not 0 <= block_id < self.kv_cache_config.num_blocks:
+                    raise ValueError(
+                        f"Diffusion KV row {identity!r} group {group_index} block ID {block_id} is outside "
+                        f"[0, {self.kv_cache_config.num_blocks})"
+                    )
+                normalized_ids.append(block_id)
+            normalized_groups.append(tuple(normalized_ids))
+
+        return _DiffusionKVRowInstall(
+            identity=identity,
+            token_len=token_len,
+            block_ids=tuple(normalized_groups),
+        )
+
+    def _prepare_install(
+        self,
+        metadata: DiffusionKVMetadata,
+    ) -> tuple[list[_DiffusionKVRowInstall], DiffusionKVSnapshot]:
+        if self.block_tables is None or self.kv_cache_config is None:
+            raise RuntimeError("paged_scheduler native KV cache is not initialized")
+        if type(metadata.request_id) is not str or not metadata.request_id:
+            raise ValueError("Diffusion KV metadata request_id must be a non-empty string")
+        if type(metadata.allocation_generation) is not int or metadata.allocation_generation < 0:
+            raise ValueError("Diffusion KV allocation_generation must be a non-negative integer")
+        if type(self.kv_cache_config.num_blocks) is not int or self.kv_cache_config.num_blocks <= 0:
+            raise ValueError("native diffusion KV cache config must provide a positive num_blocks")
+        if not metadata.sequences:
+            raise ValueError(f"Diffusion KV request {metadata.request_id!r} must contain at least one sequence")
+
+        seen_identities: set[DiffusionKVIdentity] = set()
+        installs: list[_DiffusionKVRowInstall] = []
+        sequence_snapshots: list[object] = []
+        for sequence in metadata.sequences:
+            if type(sequence.sequence_id) is not int or sequence.sequence_id < 0:
+                raise ValueError(f"Diffusion KV sequence_id must be a non-negative integer: {sequence.sequence_id!r}")
+            sequence_identity = (metadata.request_id, sequence.sequence_id, None)
+            sequence_install = self._validate_row(
+                identity=sequence_identity,
+                token_len=sequence.seq_len,
+                block_ids=sequence.block_ids,
+                seen_identities=seen_identities,
+            )
+            installs.append(sequence_install)
+            sequence_snapshots.append(
+                (
+                    sequence.sequence_id,
+                    sequence.prefix_len,
+                    sequence.target_len,
+                    sequence.seq_len,
+                    sequence_install.block_ids,
+                    sequence.context_ids,
+                )
+            )
+
+        context_snapshots: list[object] = []
+        for context in metadata.contexts:
+            if type(context.context_id) is not str or not context.context_id:
+                raise ValueError("Diffusion KV context_id must be a non-empty string")
+            context_identity = (metadata.request_id, None, context.context_id)
+            context_install = self._validate_row(
+                identity=context_identity,
+                token_len=context.num_tokens,
+                block_ids=context.block_ids,
+                seen_identities=seen_identities,
+            )
+            installs.append(context_install)
+            context_snapshots.append(
+                (
+                    context.context_id,
+                    context.cache_role,
+                    context.num_tokens,
+                    context_install.block_ids,
+                )
+            )
+
+        max_rows_per_request = self._get_max_rows_per_request()
+        if len(installs) > max_rows_per_request:
+            raise ValueError(
+                f"Diffusion KV request {metadata.request_id!r} requires {len(installs)} rows; "
+                f"adapter limit is {max_rows_per_request}"
+            )
+        snapshot: DiffusionKVSnapshot = (
+            metadata.request_id,
+            tuple(sequence_snapshots),
+            tuple(context_snapshots),
+        )
+        return installs, snapshot
+
+    def _rollback_staged_writes(self, rows: list[int], previous_num_blocks: object | None) -> None:
+        block_tables = self.block_tables
+        if block_tables is None:
+            return
+        for group_table in getattr(block_tables, "block_tables", ()):
+            clear_staged_writes = getattr(group_table, "clear_staged_writes", None)
+            if clear_staged_writes is not None:
+                clear_staged_writes()
+        if previous_num_blocks is None:
+            return
+        native_num_blocks = getattr(block_tables, "num_blocks", None)
+        if native_num_blocks is None or not hasattr(native_num_blocks, "np"):
+            return
+        native_num_blocks.np[:, rows] = previous_num_blocks
+        copy_to_uva = getattr(native_num_blocks, "copy_to_uva", None)
+        if copy_to_uva is not None:
+            copy_to_uva()
+
+    def _apply_rows(self, rows: list[int], installs: list[_DiffusionKVRowInstall]) -> None:
+        assert self.block_tables is not None
+        if self.paged_attention_adapter is not None:
+            self.paged_attention_adapter.invalidate_prepared_batches()
+        native_num_blocks = getattr(self.block_tables, "num_blocks", None)
+        previous_num_blocks = None
+        if native_num_blocks is not None and hasattr(native_num_blocks, "np"):
+            previous_num_blocks = native_num_blocks.np[:, rows].copy()
+        try:
+            for row, install in zip(rows, installs, strict=True):
+                self.block_tables.append_block_ids(
+                    row,
+                    tuple(list(group_ids) for group_ids in install.block_ids),
+                    overwrite=True,
+                )
+            self.block_tables.apply_staged_writes()
+        except Exception:
+            self._rollback_staged_writes(rows, previous_num_blocks)
+            raise
+
+    def install_diffusion_kv_metadata(self, metadata: DiffusionKVMetadata) -> bool:
+        """Install one Scheduler allocation snapshot into native Worker rows."""
+        if not is_scheduler_paged_kv_mode(
+            getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+        ):
+            raise ValueError("Dense diffusion execution must not install Diffusion KV metadata")
+
+        installs, snapshot = self._prepare_install(metadata)
+        current_state = self._diffusion_kv_request_states.get(metadata.request_id)
+        if current_state is not None:
+            if metadata.allocation_generation < current_state.generation:
+                raise ValueError(
+                    f"Stale Diffusion KV allocation generation {metadata.allocation_generation} for request "
+                    f"{metadata.request_id!r}; installed generation is {current_state.generation}"
+                )
+            if snapshot != current_state.snapshot:
+                raise ValueError(
+                    f"Conflicting Diffusion KV allocation snapshot for active request {metadata.request_id!r}; "
+                    "remove the request before replacing its rows"
+                )
+            if metadata.allocation_generation == current_state.generation:
+                return False
+            self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState(
+                generation=metadata.allocation_generation,
+                snapshot=snapshot,
+                row_token_lens=current_state.row_token_lens,
+            )
+            return False
+
+        existing_identities = [
+            install.identity for install in installs if install.identity in self._diffusion_kv_identity_to_row
+        ]
+        if existing_identities:
+            raise RuntimeError(f"Diffusion KV row registry contains orphaned identities: {existing_identities!r}")
+        if len(installs) > len(self._diffusion_kv_free_rows):
+            raise ValueError(
+                f"Diffusion KV request {metadata.request_id!r} requires {len(installs)} rows, but only "
+                f"{len(self._diffusion_kv_free_rows)} rows are free"
+            )
+
+        rows = list(reversed(self._diffusion_kv_free_rows[-len(installs) :])) if installs else []
+        if rows:
+            self._apply_rows(rows, installs)
+
+        if installs:
+            del self._diffusion_kv_free_rows[-len(installs) :]
+        self._diffusion_kv_identity_to_row.update(
+            (install.identity, row) for install, row in zip(installs, rows, strict=True)
+        )
+        self._diffusion_kv_request_states[metadata.request_id] = _DiffusionKVRequestState(
+            generation=metadata.allocation_generation,
+            snapshot=snapshot,
+            row_token_lens=tuple((install.identity, install.token_len) for install in installs),
+        )
+        return True
+
+    def get_diffusion_kv_row(
+        self,
+        request_id: str,
+        sequence_id: int | None,
+        context_id: str | None = None,
+    ) -> int:
+        """Resolve a Scheduler allocation identity to its native table row."""
+        if (
+            not is_scheduler_paged_kv_mode(
+                getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            )
+            or self.block_tables is None
+        ):
+            raise RuntimeError("paged_scheduler native BlockTables are not initialized")
+        identity = (request_id, None, context_id) if context_id is not None else (request_id, sequence_id, None)
+        try:
+            return self._diffusion_kv_identity_to_row[identity]
+        except KeyError as exc:
+            raise KeyError(f"No native diffusion KV row is installed for {identity!r}") from exc
+
+    def get_paged_attention_adapter(self) -> DiffusionPagedAttentionAdapter:
+        if self.paged_attention_adapter is None:
+            raise RuntimeError("paged_scheduler native attention adapter is not initialized")
+        return self.paged_attention_adapter
+
+    def prepare_paged_attention_batch(
+        self,
+        rows: Sequence[DiffusionPagedAttentionRow],
+    ) -> PreparedDiffusionPagedAttentionBatch:
+        """Build native page-table metadata for one Diffusion forward.
+
+        The scheduler allocation payload describes capacity, while the model
+        integration supplies the current ``query_len``/``kv_start_pos`` span
+        for each row.  Keeping this boundary explicit prevents the Worker from
+        guessing model-specific text/image layout.
+        """
+
+        return self.get_paged_attention_adapter().prepare_batch(rows)
+
+    def activate_paged_attention(self, batch: PreparedDiffusionPagedAttentionBatch):
+        """Return the context manager that exposes a prepared batch to Omni Attention."""
+
+        return self.get_paged_attention_adapter().activate(batch)
+
+    def _resolve_paged_attention_row(
+        self,
+        request_id: str,
+        sequence_id: int | None,
+        context_id: str | None,
+    ) -> DiffusionPagedAttentionRowBinding:
+        row_index = self.get_diffusion_kv_row(request_id, sequence_id, context_id)
+        identity = (request_id, None, context_id) if context_id is not None else (request_id, sequence_id, None)
+        request_state = self._diffusion_kv_request_states.get(request_id)
+        if request_state is not None:
+            for installed_identity, token_len in request_state.row_token_lens:
+                if installed_identity == identity:
+                    return DiffusionPagedAttentionRowBinding(
+                        row_index=row_index,
+                        max_seq_len=token_len,
+                    )
+        raise RuntimeError(f"Diffusion KV request state is missing logical length for {identity!r}")
+
+    def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int:
+        """Retire Worker rows without logically freeing Scheduler-owned blocks."""
+        if (
+            not is_scheduler_paged_kv_mode(
+                getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            )
+            or self.block_tables is None
+        ):
+            return 0
+        request_id_set = set(request_ids)
+        identities_and_rows = [
+            (identity, row)
+            for identity, row in self._diffusion_kv_identity_to_row.items()
+            if identity[0] in request_id_set
+        ]
+        rows = [row for _, row in identities_and_rows]
+        if rows:
+            assert self.kv_cache_config is not None
+            num_groups = len(self.kv_cache_config.kv_cache_groups)
+            installs = [
+                _DiffusionKVRowInstall(
+                    identity=identity,
+                    token_len=0,
+                    block_ids=tuple(tuple() for _ in range(num_groups)),
+                )
+                for identity, _ in identities_and_rows
+            ]
+            self._apply_rows(rows, installs)
+
+        for identity, _ in identities_and_rows:
+            del self._diffusion_kv_identity_to_row[identity]
+        for request_id in request_id_set:
+            self._diffusion_kv_request_states.pop(request_id, None)
+        self._diffusion_kv_free_rows.extend(rows)
+        self._diffusion_kv_free_rows.sort(reverse=True)
+        return len(rows)
+
+    def refresh_block_table_layout(self) -> None:
+        """Refresh native pointer tensors after a CuMem KV-cache wake-up."""
+        if (
+            not is_scheduler_paged_kv_mode(
+                getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            )
+            or self.block_tables is None
+        ):
+            return
+        if self.paged_attention_adapter is not None:
+            self.paged_attention_adapter.invalidate_prepared_batches()
+        self.block_tables.init_block_table_layout_tensors()

vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py

@@ -0,0 +1,778 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass, field, replace
+from typing import Any
+
+import torch
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.utils.math_utils import cdiv
+from vllm.v1.attention.backend import AttentionType
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
+from vllm.v1.attention.selector import get_attn_backend
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig
+from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
+from vllm.v1.worker.gpu.block_table import BlockTables
+
+from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
+    PagedPiecewisePlan,
+    build_paged_piecewise_plan,
+)
+from vllm_omni.diffusion.forward_context import override_paged_kv_adapter
+from vllm_omni.platforms import current_omni_platform
+
+
+@dataclass(frozen=True)
+class DiffusionPagedAttentionRowBinding:
+    """Worker row and logical length installed for one allocation identity."""
+
+    row_index: int
+    max_seq_len: int
+
+
+DiffusionKVRowResolver = Callable[
+    [str, int | None, str | None],
+    DiffusionPagedAttentionRowBinding,
+]
+
+
+class DiffusionPagedAttentionLayerAdapter(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.
+    """
+
+    def __init__(
+        self,
+        *,
+        layer_name: str,
+        layer: Any,
+        spec: AttentionSpec,
+        vllm_config: VllmConfig,
+        device: torch.device,
+        ulysses_degree: int = 1,
+    ) -> None:
+        if type(ulysses_degree) is not int or ulysses_degree <= 0:
+            raise ValueError(f"ulysses_degree must be a positive integer, got {ulysses_degree!r}")
+        num_heads = int(layer.num_heads)
+        num_kv_heads = int(spec.num_kv_heads)
+        if num_heads <= 0 or num_kv_heads <= 0 or num_heads % num_kv_heads != 0:
+            raise ValueError(
+                "Paged attention requires positive Q/KV heads with num_heads divisible by num_kv_heads: "
+                f"num_heads={num_heads}, num_kv_heads={num_kv_heads}"
+            )
+        if num_heads % ulysses_degree != 0 or num_kv_heads % ulysses_degree != 0:
+            raise ValueError(
+                "Paged attention requires Q/KV heads divisible by ulysses_degree: "
+                f"num_heads={num_heads}, num_kv_heads={num_kv_heads}, ulysses_degree={ulysses_degree}"
+            )
+        num_heads //= ulysses_degree
+        num_kv_heads //= ulysses_degree
+
+        attention_config = vllm_config.attention_config
+        previous_backend = attention_config.backend
+        previous_backend_per_kind = attention_config.backend_per_kind
+        try:
+            # This is a portable vLLM backend request. The active platform
+            # resolves it to its native implementation, such as FlashAttention
+            # on CUDA or AscendAttentionBackend on NPU.
+            attention_config.backend = AttentionBackendEnum.FLASH_ATTN
+            attention_config.backend_per_kind = {}
+            with set_current_vllm_config(vllm_config):
+                attn_backend = get_attn_backend(
+                    head_size=spec.head_size,
+                    dtype=vllm_config.model_config.dtype,
+                    kv_cache_dtype=vllm_config.cache_config.cache_dtype,
+                    attn_type=AttentionType.DECODER,
+                    num_heads=num_heads,
+                )
+        finally:
+            attention_config.backend = previous_backend
+            attention_config.backend_per_kind = previous_backend_per_kind
+        canonical_spec = replace(
+            spec,
+            num_kv_heads=num_kv_heads,
+            indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(),
+        )
+        self.layer_name = layer_name
+        self.spec = canonical_spec
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_size = int(canonical_spec.head_size)
+        self.head_size_v = int(getattr(canonical_spec, "head_size_v", canonical_spec.head_size))
+        if self.head_size_v != self.head_size:
+            raise ValueError(
+                "Diffusion native paged attention requires equal key/value head sizes; "
+                f"got head_size={self.head_size}, head_size_v={self.head_size_v}"
+            )
+        self.softmax_scale = float(layer.softmax_scale)
+        self.attn_backend = attn_backend
+        self.kv_cache: torch.Tensor | None = None
+        # Native attention backends consume either device tensors or host
+        # scalar copies of these scales. The paged diffusion path is currently
+        # unquantized, so initialize both contracts to identity.
+        self._q_scale = torch.ones(1, device=device, dtype=torch.float32)
+        self._k_scale = torch.ones(1, device=device, dtype=torch.float32)
+        self._v_scale = torch.ones(1, device=device, dtype=torch.float32)
+        self._q_scale_float = 1.0
+        self._k_scale_float = 1.0
+        self._v_scale_float = 1.0
+        with set_current_vllm_config(vllm_config):
+            self.impl = self._create_native_impl(vllm_config)
+
+    def get_attn_backend(self):
+        return self.attn_backend
+
+    def _create_native_impl(self, vllm_config: VllmConfig):
+        impl_cls = self.attn_backend.get_impl_cls()
+        impl = impl_cls(
+            num_heads=self.num_heads,
+            head_size=self.head_size,
+            scale=self.softmax_scale,
+            num_kv_heads=self.num_kv_heads,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype=vllm_config.cache_config.cache_dtype,
+            logits_soft_cap=None,
+            attn_type=AttentionType.DECODER,
+            kv_sharing_target_layer_name=None,
+        )
+        if not callable(getattr(impl, "forward", None)):
+            raise RuntimeError(
+                f"Native paged-attention backend {self.attn_backend.get_name()!r} has no forward implementation"
+            )
+        if not self.attn_backend.forward_includes_kv_cache_update and not callable(
+            getattr(impl, "do_kv_cache_update", None)
+        ):
+            raise RuntimeError(
+                f"Native paged-attention backend {self.attn_backend.get_name()!r} cannot update the KV cache"
+            )
+        if hasattr(impl, "vllm_flash_attn_version") and impl.vllm_flash_attn_version is None:
+            raise RuntimeError(f"Native paged-attention kernel is unavailable for diffusion layer {self.layer_name!r}")
+        return impl
+
+    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> AttentionSpec:
+        del vllm_config
+        return self.spec
+
+
+@dataclass(frozen=True)
+class DiffusionPagedAttentionRow:
+    """One logical BlockTable row participating in a paged attention call."""
+
+    request_id: str
+    query_len: int
+    seq_len: int
+    kv_start_pos: int = 0
+    sequence_id: int | None = None
+    context_id: str | None = None
+
+    def __post_init__(self) -> None:
+        if type(self.request_id) is not str or not self.request_id:
+            raise ValueError("Paged attention request_id must be a non-empty string")
+        if (self.sequence_id is None) == (self.context_id is None):
+            raise ValueError("Paged attention row requires exactly one of sequence_id or context_id")
+        if self.sequence_id is not None and (type(self.sequence_id) is not int or self.sequence_id < 0):
+            raise ValueError("Paged attention sequence_id must be a non-negative integer")
+        if self.context_id is not None and (type(self.context_id) is not str or not self.context_id):
+            raise ValueError("Paged attention context_id must be a non-empty string")
+        if type(self.query_len) is not int or self.query_len <= 0:
+            raise ValueError("Paged attention query_len must be a positive integer")
+        if type(self.seq_len) is not int or self.seq_len <= 0:
+            raise ValueError("Paged attention seq_len must be a positive integer")
+        if type(self.kv_start_pos) is not int or self.kv_start_pos < 0:
+            raise ValueError("Paged attention kv_start_pos must be a non-negative integer")
+        if self.kv_start_pos + self.query_len > self.seq_len:
+            raise ValueError(
+                "Paged attention write span exceeds seq_len: "
+                f"start={self.kv_start_pos}, query_len={self.query_len}, seq_len={self.seq_len}"
+            )
+
+    @property
+    def identity(self) -> tuple[str, int | None, str | None]:
+        return (self.request_id, self.sequence_id, self.context_id)
+
+
+@dataclass(frozen=True)
+class PreparedDiffusionPagedAttentionBatch:
+    """Native metadata shared by all paged attention layers in one forward."""
+
+    rows: tuple[DiffusionPagedAttentionRow, ...]
+    row_indices: torch.Tensor
+    query_start_loc: torch.Tensor
+    seq_lens: torch.Tensor
+    positions: torch.Tensor
+    block_tables: tuple[torch.Tensor, ...]
+    slot_mappings: torch.Tensor
+    attn_metadata: dict[str, Any]
+    slot_mappings_by_layer: dict[str, torch.Tensor]
+    num_tokens: int
+    _owner: object = field(repr=False, compare=False)
+    _generation: int = field(repr=False, compare=False)
+
+
+@dataclass(frozen=True, slots=True)
+class DiffusionPagedAttentionContext:
+    """One layer's native inputs for an Omni paged-backend invocation."""
+
+    layer: DiffusionPagedAttentionLayerAdapter
+    query: torch.Tensor
+    key_write: torch.Tensor
+    value_write: torch.Tensor
+    slot_mapping: torch.Tensor
+    native_metadata: Any
+    piecewise_plan: PagedPiecewisePlan | None
+    piecewise_native_metadata: tuple[Any, ...]
+    query_token_shape: tuple[int, ...]
+    query_has_head_dims: bool
+
+    def restore_output(self, output: torch.Tensor) -> torch.Tensor:
+        if self.query_has_head_dims:
+            return output.reshape(*self.query_token_shape, self.layer.num_heads, self.layer.head_size_v)
+        return output.reshape(*self.query_token_shape, self.layer.num_heads * self.layer.head_size_v)
+
+
+class DiffusionPagedAttentionAdapter:
+    """Translate diffusion rows into vLLM metadata for an Omni paged backend."""
+
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        device: torch.device,
+        kv_cache_config: KVCacheConfig,
+        block_tables: BlockTables,
+        attn_groups: list[list[Any]],
+        layers: Mapping[str, DiffusionPagedAttentionLayerAdapter],
+        resolve_row: DiffusionKVRowResolver,
+    ) -> None:
+        if not layers:
+            raise ValueError("Paged attention requires at least one native attention layer")
+        if len(attn_groups) != len(kv_cache_config.kv_cache_groups):
+            raise ValueError(
+                "Paged attention group mismatch: "
+                f"builders={len(attn_groups)}, cache_groups={len(kv_cache_config.kv_cache_groups)}"
+            )
+        self.vllm_config = vllm_config
+        self.device = torch.device(device)
+        self.kv_cache_config = kv_cache_config
+        self.block_tables = block_tables
+        self.attn_groups = attn_groups
+        self.layers = dict(layers)
+        self.resolve_row = resolve_row
+        self._owner = object()
+        self._prepare_generation = 0
+        self._active_batch: PreparedDiffusionPagedAttentionBatch | None = None
+        self._active_piecewise_plan: PagedPiecewisePlan | None = None
+        self._active_piecewise_native_metadata: tuple[dict[str, Any], ...] | None = None
+        self._causal_by_group = self._resolve_group_causality()
+        self._reorder_batch_threshold = self._resolve_reorder_batch_threshold()
+
+    def _resolve_group_causality(self) -> dict[int, bool]:
+        causal_by_group: dict[int, bool] = {}
+        for group_index, cache_group in enumerate(self.kv_cache_config.kv_cache_groups):
+            layer_causality = {
+                not bool(getattr(self.layers[layer_name].spec, "non_causal", False))
+                for layer_name in cache_group.layer_names
+            }
+            if len(layer_causality) != 1:
+                raise ValueError(f"Paged attention cache group {group_index} mixes causal and non-causal layers")
+            causal_by_group[group_index] = layer_causality.pop()
+        return causal_by_group
+
+    def _resolve_reorder_batch_threshold(self) -> int | None:
+        thresholds = [
+            group.get_metadata_builder(0).reorder_batch_threshold
+            for group_index, groups in enumerate(self.attn_groups)
+            if self._causal_by_group[group_index]
+            for group in groups
+        ]
+        concrete_thresholds = [threshold for threshold in thresholds if threshold is not None]
+        return min(concrete_thresholds, default=None)
+
+    def _validate_native_batch_order(self, rows: tuple[DiffusionPagedAttentionRow, ...]) -> None:
+        threshold = self._reorder_batch_threshold
+        if threshold is None or not any(self._causal_by_group.values()):
+            return
+
+        found_long_query = False
+        for row in rows:
+            if row.query_len > threshold:
+                found_long_query = True
+            elif found_long_query:
+                raise ValueError(
+                    "Causal paged attention rows must place native decode/short-query rows before "
+                    f"prefill/long-query rows (decode threshold={threshold})"
+                )
+
+    def _validate_row_capacity(
+        self,
+        rows: tuple[DiffusionPagedAttentionRow, ...],
+        row_indices: list[int],
+    ) -> None:
+        native_num_blocks = self.block_tables.num_blocks.np
+        blocks_per_kv_block = self.block_tables.blocks_per_kv_block
+        for row, row_index in zip(rows, row_indices, strict=True):
+            if type(row_index) is not int or not 0 <= row_index < self.block_tables.max_num_reqs:
+                raise ValueError(f"Paged attention row {row.identity!r} resolved to invalid Worker row {row_index!r}")
+            for group_index, (block_size, block_multiplier) in enumerate(
+                zip(
+                    self.block_tables.block_sizes,
+                    blocks_per_kv_block,
+                    strict=True,
+                )
+            ):
+                required_blocks = cdiv(
+                    row.seq_len,
+                    block_size * self.block_tables.cp_size,
+                )
+                required_kernel_blocks = required_blocks * block_multiplier
+                installed_kernel_blocks = int(native_num_blocks[group_index, row_index])
+                if required_kernel_blocks > installed_kernel_blocks:
+                    raise ValueError(
+                        f"Paged attention row {row.identity!r} requires {required_blocks} blocks in "
+                        f"cache group {group_index} for seq_len={row.seq_len}, but its installed Worker "
+                        f"row contains only {installed_kernel_blocks // block_multiplier} blocks"
+                    )
+
+    def _build_native_metadata(
+        self,
+        *,
+        query_lens: Sequence[int],
+        query_start_loc: torch.Tensor,
+        query_start_loc_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        positions: torch.Tensor,
+        block_tables: Sequence[torch.Tensor],
+        slot_mappings: torch.Tensor,
+        causal: bool | Mapping[int, bool],
+    ) -> dict[str, Any]:
+        dcp_local_seq_lens = None
+        if self.block_tables.cp_size > 1:
+            dcp_local_seq_lens = get_dcp_local_seq_lens(
+                seq_lens,
+                dcp_size=self.block_tables.cp_size,
+                dcp_rank=self.block_tables.cp_rank,
+                cp_kv_cache_interleave_size=self.block_tables.cp_interleave,
+            )
+        with set_current_vllm_config(self.vllm_config):
+            return current_omni_platform.build_diffusion_kv_attn_metadata(
+                attn_groups=self.attn_groups,
+                num_reqs=len(query_lens),
+                num_tokens=sum(query_lens),
+                query_start_loc_gpu=query_start_loc,
+                query_start_loc_cpu=query_start_loc_cpu,
+                max_query_len=max(query_lens),
+                seq_lens=seq_lens,
+                max_seq_len=max(int(seq_len) for seq_len in seq_lens_cpu),
+                block_tables=block_tables,
+                slot_mappings=slot_mappings,
+                kv_cache_config=self.kv_cache_config,
+                seq_lens_cpu=seq_lens_cpu,
+                seq_lens_cpu_upper_bound=seq_lens_cpu,
+                dcp_local_seq_lens=dcp_local_seq_lens,
+                positions=positions,
+                causal=causal,
+            )
+
+    def prepare_batch(
+        self,
+        rows: Sequence[DiffusionPagedAttentionRow],
+    ) -> PreparedDiffusionPagedAttentionBatch:
+        rows = tuple(rows)
+        if not rows:
+            raise ValueError("Paged attention batch must contain at least one row")
+        if self._active_batch is not None:
+            raise RuntimeError("Cannot prepare paged attention metadata during an active forward")
+        if len(rows) > self.block_tables.max_num_reqs:
+            raise ValueError(
+                f"Paged attention batch has {len(rows)} rows; Worker capacity is {self.block_tables.max_num_reqs}"
+            )
+
+        identities = [row.identity for row in rows]
+        if len(set(identities)) != len(identities):
+            raise ValueError("Paged attention batch contains duplicate row identities")
+        row_bindings = [self.resolve_row(row.request_id, row.sequence_id, row.context_id) for row in rows]
+        row_indices_list = [binding.row_index for binding in row_bindings]
+        if len(set(row_indices_list)) != len(row_indices_list):
+            raise ValueError("Paged attention batch resolves multiple inputs to the same Worker row")
+        for row, binding in zip(rows, row_bindings, strict=True):
+            if type(binding.max_seq_len) is not int or binding.max_seq_len < 0:
+                raise ValueError(
+                    f"Paged attention row {row.identity!r} resolved to invalid logical capacity {binding.max_seq_len!r}"
+                )
+            if row.seq_len > binding.max_seq_len:
+                raise ValueError(
+                    f"Paged attention row {row.identity!r} uses seq_len={row.seq_len}, but its installed "
+                    f"allocation has logical length {binding.max_seq_len}"
+                )
+        self._validate_native_batch_order(rows)
+        self._validate_row_capacity(rows, row_indices_list)
+
+        query_lens = [row.query_len for row in rows]
+        num_tokens = sum(query_lens)
+        if num_tokens > self.block_tables.max_num_batched_tokens:
+            raise ValueError(
+                f"Paged attention batch has {num_tokens} tokens; Worker capacity is "
+                f"{self.block_tables.max_num_batched_tokens}"
+            )
+        query_offsets = [0]
+        for query_len in query_lens:
+            query_offsets.append(query_offsets[-1] + query_len)
+
+        # Native BlockTables reuses persistent gather/slot buffers. Starting a
+        # new preparation invalidates every batch prepared before it, even if a
+        # later metadata builder raises.
+        self._prepare_generation += 1
+        generation = self._prepare_generation
+        row_indices = torch.tensor(row_indices_list, dtype=torch.int32, device=self.device)
+        query_start_loc_cpu = torch.tensor(query_offsets, dtype=torch.int32)
+        query_start_loc = query_start_loc_cpu.to(self.device)
+        seq_lens_cpu = torch.tensor([row.seq_len for row in rows], dtype=torch.int32)
+        seq_lens = seq_lens_cpu.to(self.device)
+        positions = torch.cat(
+            [
+                torch.arange(
+                    row.kv_start_pos,
+                    row.kv_start_pos + row.query_len,
+                    dtype=torch.int64,
+                    device=self.device,
+                )
+                for row in rows
+            ]
+        )
+
+        block_tables = self.block_tables.gather_block_tables(
+            row_indices,
+            num_reqs_padded=len(rows),
+        )
+        slot_mappings = self.block_tables.compute_slot_mappings(
+            row_indices,
+            query_start_loc,
+            positions,
+            num_tokens_padded=num_tokens,
+        )
+        causal: bool | Mapping[int, bool]
+        causal_values = set(self._causal_by_group.values())
+        causal = causal_values.pop() if len(causal_values) == 1 else self._causal_by_group
+        attn_metadata = self._build_native_metadata(
+            query_lens=query_lens,
+            query_start_loc=query_start_loc,
+            query_start_loc_cpu=query_start_loc_cpu,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            positions=positions,
+            block_tables=block_tables,
+            slot_mappings=slot_mappings,
+            causal=causal,
+        )
+        slot_mappings_by_layer = build_slot_mappings_by_layer(
+            slot_mappings,
+            self.kv_cache_config,
+        )
+        return PreparedDiffusionPagedAttentionBatch(
+            rows=rows,
+            row_indices=row_indices,
+            query_start_loc=query_start_loc,
+            seq_lens=seq_lens,
+            positions=positions,
+            block_tables=tuple(block_tables),
+            slot_mappings=slot_mappings,
+            attn_metadata=attn_metadata,
+            slot_mappings_by_layer=slot_mappings_by_layer,
+            num_tokens=num_tokens,
+            _owner=self._owner,
+            _generation=generation,
+        )
+
+    def invalidate_prepared_batches(self) -> None:
+        """Invalidate native buffer views after BlockTable state changes."""
+
+        if self._active_batch is not None:
+            raise RuntimeError("Cannot change paged attention BlockTables during an active forward")
+        self._prepare_generation += 1
+
+    @contextmanager
+    def activate(
+        self,
+        batch: PreparedDiffusionPagedAttentionBatch,
+    ) -> Iterator[DiffusionPagedAttentionAdapter]:
+        if batch._owner is not self._owner:
+            raise ValueError("Prepared paged attention batch belongs to a different adapter")
+        if batch._generation != self._prepare_generation:
+            raise ValueError("Prepared paged attention batch is stale after a newer batch preparation")
+        if self._active_batch is not None:
+            raise RuntimeError("Paged attention adapter already has an active forward")
+        self._active_batch = batch
+        self._active_piecewise_plan = None
+        self._active_piecewise_native_metadata = None
+        try:
+            with override_paged_kv_adapter(self), set_current_vllm_config(self.vllm_config):
+                yield self
+        finally:
+            self._active_piecewise_plan = None
+            self._active_piecewise_native_metadata = None
+            self._active_batch = None
+
+    @staticmethod
+    def _flatten_tensor(
+        tensor: torch.Tensor,
+        *,
+        num_heads: int,
+        head_size: int,
+        name: str,
+    ) -> tuple[torch.Tensor, tuple[int, ...], bool]:
+        if tensor.ndim >= 3 and tensor.shape[-2:] == (num_heads, head_size):
+            return tensor.reshape(-1, num_heads, head_size), tuple(tensor.shape[:-2]), True
+        hidden_size = num_heads * head_size
+        if tensor.ndim >= 1 and tensor.shape[-1] == hidden_size:
+            return tensor.reshape(-1, num_heads, head_size), tuple(tensor.shape[:-1]), False
+        raise ValueError(
+            f"Paged attention {name} must end in ({num_heads}, {head_size}) or ({hidden_size},); "
+            f"got shape={tuple(tensor.shape)}"
+        )
+
+    @staticmethod
+    def _validate_token_layout(
+        token_shape: tuple[int, ...],
+        batch: PreparedDiffusionPagedAttentionBatch,
+        *,
+        name: str,
+    ) -> None:
+        if len(token_shape) == 1:
+            if token_shape[0] == batch.num_tokens:
+                return
+            raise ValueError(
+                f"Paged attention {name} token count must match the prepared write batch: "
+                f"tokens={token_shape[0]}, prepared={batch.num_tokens}"
+            )
+        if len(token_shape) == 2:
+            batch_size, tokens_per_row = token_shape
+            query_lens = tuple(row.query_len for row in batch.rows)
+            if batch_size == len(batch.rows) and all(query_len == tokens_per_row for query_len in query_lens):
+                return
+            raise ValueError(
+                f"Paged attention batched {name} layout must match prepared rows for the current write: "
+                f"shape={token_shape}, row_query_lens={query_lens}"
+            )
+        raise ValueError(
+            f"Paged attention {name} supports packed [T, ...] or uniform batched [B, T, ...] token layouts; "
+            f"got token shape={token_shape}"
+        )
+
+    def _get_piecewise_plan(
+        self,
+        full_attn_spans: list[list[tuple[int, int]]],
+    ) -> PagedPiecewisePlan:
+        batch = self._active_batch
+        if batch is None:
+            raise RuntimeError("Piecewise paged attention requires an active prepared batch")
+        if any(row.kv_start_pos + row.query_len != row.seq_len for row in batch.rows):
+            invalid_rows = [row.identity for row in batch.rows if row.kv_start_pos + row.query_len != row.seq_len]
+            raise ValueError(
+                "Paged piecewise attention requires each query/write span to end at seq_len; "
+                f"invalid rows={invalid_rows!r}"
+            )
+
+        cached_plan = self._active_piecewise_plan
+        if cached_plan is not None:
+            frozen_spans = tuple(tuple(tuple(span) for span in row_spans) for row_spans in full_attn_spans)
+            if cached_plan.spans != frozen_spans:
+                raise ValueError("Paged piecewise attention metadata changed between layers in one active batch")
+            return cached_plan
+
+        plan = build_paged_piecewise_plan(
+            full_attn_spans,
+            query_offsets=[row.kv_start_pos for row in batch.rows],
+            query_lens=[row.query_len for row in batch.rows],
+            seq_lens=[row.seq_len for row in batch.rows],
+            device=self.device,
+        )
+        self._active_piecewise_plan = plan
+        return plan
+
+    def _get_piecewise_native_metadata(
+        self,
+        plan: PagedPiecewisePlan,
+    ) -> tuple[dict[str, Any], ...]:
+        """Build one native metadata object per segment and cache it by batch."""
+
+        batch = self._active_batch
+        if batch is None:
+            raise RuntimeError("Piecewise paged attention requires an active prepared batch")
+        cached_metadata = self._active_piecewise_native_metadata
+        if cached_metadata is not None:
+            return cached_metadata
+
+        metadata_by_segment: list[dict[str, Any]] = []
+        for packed_segment in plan.segments:
+            row_segments = packed_segment.row_segments
+            query_lens = [segment.q_end - segment.q_start for segment in row_segments]
+            seq_lens_values = [segment.kv_end for segment in row_segments]
+            query_offsets = [0]
+            for query_len in query_lens:
+                query_offsets.append(query_offsets[-1] + query_len)
+            query_start_loc_cpu = torch.tensor(query_offsets, dtype=torch.int32)
+            seq_lens_cpu = torch.tensor(seq_lens_values, dtype=torch.int32)
+            metadata_by_segment.append(
+                self._build_native_metadata(
+                    query_lens=query_lens,
+                    query_start_loc=query_start_loc_cpu.to(self.device),
+                    query_start_loc_cpu=query_start_loc_cpu,
+                    seq_lens=seq_lens_cpu.to(self.device),
+                    seq_lens_cpu=seq_lens_cpu,
+                    positions=batch.positions.index_select(0, packed_segment.query_indices),
+                    block_tables=batch.block_tables,
+                    slot_mappings=batch.slot_mappings.index_select(-1, packed_segment.query_indices),
+                    causal=(row_segments[0].mode == "causal"),
+                )
+            )
+        self._active_piecewise_native_metadata = tuple(metadata_by_segment)
+        return self._active_piecewise_native_metadata
+
+    def prepare_layer_context(
+        self,
+        layer_name: str,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        *,
+        omni_attn_metadata: Any | None = None,
+    ) -> DiffusionPagedAttentionContext:
+        if self._active_batch is None:
+            raise RuntimeError("Paged attention forward must run inside adapter.activate(batch)")
+        try:
+            layer = self.layers[layer_name]
+        except KeyError as exc:
+            raise KeyError(f"Unknown diffusion paged attention layer {layer_name!r}") from exc
+        batch = self._active_batch
+        full_attn_spans = self._validate_omni_attn_metadata(omni_attn_metadata)
+
+        query_flat, query_token_shape, query_has_head_dims = self._flatten_tensor(
+            query,
+            num_heads=layer.num_heads,
+            head_size=layer.head_size,
+            name="query",
+        )
+        key_flat, key_token_shape, _ = self._flatten_tensor(
+            key,
+            num_heads=layer.num_kv_heads,
+            head_size=layer.head_size,
+            name="key",
+        )
+        value_flat, value_token_shape, _ = self._flatten_tensor(
+            value,
+            num_heads=layer.num_kv_heads,
+            head_size=layer.head_size_v,
+            name="value",
+        )
+        self._validate_token_layout(query_token_shape, batch, name="query")
+        self._validate_token_layout(key_token_shape, batch, name="key")
+        self._validate_token_layout(value_token_shape, batch, name="value")
+        if query.device != self.device or key.device != self.device or value.device != self.device:
+            raise ValueError(
+                f"Paged attention Q/K/V must be on {self.device}; "
+                f"got query={query.device}, key={key.device}, value={value.device}"
+            )
+        if query.dtype != key.dtype or query.dtype != value.dtype:
+            raise ValueError(f"Paged attention Q/K/V dtypes must match; got {query.dtype}, {key.dtype}, {value.dtype}")
+        expected_dtype = self.vllm_config.model_config.dtype
+        if query.dtype != expected_dtype:
+            raise ValueError(
+                f"Paged attention Q/K/V dtype must match model activation dtype {expected_dtype}; got {query.dtype}"
+            )
+        if not bool(getattr(layer.spec, "non_causal", False)):
+            non_suffix_rows = [row.identity for row in batch.rows if row.kv_start_pos + row.query_len != row.seq_len]
+            if non_suffix_rows:
+                raise ValueError(
+                    "Causal paged attention requires each query/write span to end at seq_len; "
+                    f"invalid rows={non_suffix_rows!r}"
+                )
+
+        slot_mapping = batch.slot_mappings_by_layer.get(layer_name)
+        if slot_mapping is None:
+            raise KeyError(f"No native slot mapping was built for diffusion layer {layer_name!r}")
+        slot_mapping = slot_mapping.reshape(-1)[: batch.num_tokens]
+        if slot_mapping.numel() != batch.num_tokens:
+            raise ValueError(
+                f"Paged attention slot mapping for {layer_name!r} has {slot_mapping.numel()} entries; "
+                f"expected {batch.num_tokens}"
+            )
+
+        # vLLM's writer consumes slot_mapping directly, so a span that starts
+        # inside a physical block naturally writes only the requested suffix.
+        key_flat = key_flat.contiguous()
+        value_flat = value_flat.contiguous()
+        try:
+            native_metadata = batch.attn_metadata[layer_name]
+        except KeyError as exc:
+            raise KeyError(f"No native attention metadata was built for diffusion layer {layer_name!r}") from exc
+        piecewise_plan: PagedPiecewisePlan | None = None
+        piecewise_native_metadata: tuple[Any, ...] = ()
+        if full_attn_spans is not None:
+            piecewise_plan = self._get_piecewise_plan(full_attn_spans)
+            native_metadata_by_segment = self._get_piecewise_native_metadata(piecewise_plan)
+            try:
+                piecewise_native_metadata = tuple(
+                    segment_metadata[layer_name] for segment_metadata in native_metadata_by_segment
+                )
+            except KeyError as exc:
+                raise KeyError(
+                    f"No piecewise native attention metadata was built for diffusion layer {layer_name!r}"
+                ) from exc
+        return DiffusionPagedAttentionContext(
+            layer=layer,
+            query=query_flat.contiguous(),
+            key_write=key_flat,
+            value_write=value_flat,
+            slot_mapping=slot_mapping,
+            native_metadata=native_metadata,
+            piecewise_plan=piecewise_plan,
+            piecewise_native_metadata=piecewise_native_metadata,
+            query_token_shape=query_token_shape,
+            query_has_head_dims=query_has_head_dims,
+        )
+
+    @staticmethod
+    def _validate_omni_attn_metadata(
+        metadata: Any | None,
+    ) -> list[list[tuple[int, int]]] | None:
+        if metadata is None:
+            return None
+        full_attn_spans = getattr(metadata, "full_attn_spans", None)
+        attn_mask = getattr(metadata, "attn_mask", None)
+        unsupported_fields = [
+            field_name
+            for field_name in (
+                "joint_attn_mask",
+                "query_ranges",
+                "video_layout",
+            )
+            if getattr(metadata, field_name, None) is not None
+        ]
+        if attn_mask is not None:
+            if full_attn_spans is None:
+                unsupported_fields.append("attn_mask")
+            elif not isinstance(attn_mask, torch.Tensor) or attn_mask.ndim != 4:
+                unsupported_fields.append("attn_mask (piecewise paging requires a 4D tensor)")
+        extra = getattr(metadata, "extra", None)
+        if extra:
+            unsupported_fields.append(f"extra={sorted(extra)}")
+        if unsupported_fields:
+            raise NotImplementedError(
+                "Diffusion paged attention cannot translate Omni attention metadata fields "
+                f"{unsupported_fields!r} to native FlashAttention metadata"
+            )
+        return full_attn_spans