[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>

tests/config/test_omni_config.py

@@ -1090,7 +1090,10 @@ def test_diffusion_config_preserves_existing_coercion_hooks():
 
 
 def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
     monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "TEA_CACHE")
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
 
     cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
         diffusion_attention_backend="flash_attn",
@@ -1100,6 +1103,7 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
         kv_cache_skip_layers=[2],
         static_lora_scale=0.25,
         diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
         diffusers_load_kwargs=None,
         diffusers_call_kwargs=None,
     )
@@ -1116,7 +1120,10 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
     assert cfg.diffusers_call_kwargs == {}
 
 
-def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path):
+def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
     deploy_path = tmp_path / "dreamzero_diffusion_aliases.yaml"
     deploy_path.write_text(
         "\n".join(
@@ -1128,6 +1135,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
                 "    diffusion_attention_backend: flash_attn",
                 "    fa_deterministic: true",
                 "    diffusion_kv_mode: paged_scheduler",
+                "    diffusion_kv_max_rows_per_request: 2",
             ]
         )
     )
@@ -1141,6 +1149,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
     assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn"
     assert stage.diffusion_config.fa_deterministic is True
     assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+    assert stage.diffusion_config.diffusion_kv_max_rows_per_request == 2
 
 
 def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
@@ -1183,6 +1192,7 @@ def test_diffusion_config_field_classification_covers_current_fields():
         "prompt_embed_cache_size",
         "diffusion_kv_cache_dtype",
         "diffusion_kv_mode",
+        "diffusion_kv_max_rows_per_request",
     } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
     assert {
         "revision",

tests/diffusion/attention/test_piecewise_attn.py

@@ -23,8 +23,10 @@ import vllm_omni.diffusion.attention.backends.utils.piecewise_attn as piecewise_
 from vllm_omni.diffusion.attention.backends.abstract import QueryRange
 from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     Segment,
+    build_paged_piecewise_plan,
     build_segments,
     piecewise_attn,
+    run_paged_piecewise_plan,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -84,6 +86,72 @@ def test_build_segments_full_span_retains_global_kv_end():
     assert segments == [Segment(q_start=4, q_end=6, kv_end=7, mode="full")]
 
 
+def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
+    full_attn_spans = [[(2, 5)], [(4, 7)]]
+    plan = build_paged_piecewise_plan(
+        full_attn_spans,
+        query_offsets=[0, 2],
+        query_lens=[6, 6],
+        seq_lens=[6, 8],
+        device=DEVICE,
+    )
+    query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
+    key = query + 100
+    value = query + 200
+    calls = []
+
+    def run_segment(segment_query, segment_key, segment_value, segment_batch):
+        calls.append(
+            (
+                [segment.mode for segment in segment_batch.row_segments],
+                [segment.kv_end for segment in segment_batch.row_segments],
+                segment_batch.query_indices.tolist(),
+            )
+        )
+        torch.testing.assert_close(segment_key, segment_query + 100)
+        torch.testing.assert_close(segment_value, segment_query + 200)
+        return segment_query
+
+    output = run_paged_piecewise_plan(
+        query,
+        key,
+        value,
+        plan,
+        plan.segments,
+        run_segment,
+    )
+
+    torch.testing.assert_close(output, query)
+    assert calls == [
+        (["causal", "causal"], [2, 4], [0, 1, 6, 7]),
+        (["full", "full"], [5, 7], [2, 3, 4, 8, 9, 10]),
+        (["causal", "causal"], [6, 8], [5, 11]),
+    ]
+
+
+def test_paged_piecewise_runner_uses_runner_output_shape():
+    plan = build_paged_piecewise_plan(
+        [[(2, 5)]],
+        query_offsets=[0],
+        query_lens=[6],
+        seq_lens=[6],
+        device=DEVICE,
+    )
+    query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
+
+    output = run_paged_piecewise_plan(
+        query,
+        query,
+        query,
+        plan,
+        plan.segments,
+        lambda segment_query, *_args: segment_query[..., :3],
+    )
+
+    assert output.shape == (6, 2, 3)
+    torch.testing.assert_close(output, query[..., :3])
+
+
 @pytest.mark.parametrize("global_spans", SPAN_CASES)
 @pytest.mark.parametrize("q_range", Q_RANGE_CASES)
 @pytest.mark.parametrize("batch_size", BATCH_CASES)

tests/diffusion/diffusion_kv/test_block_tables.py

@@ -0,0 +1,772 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import math
+from contextlib import nullcontext
+from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+from vllm.v1.worker import block_table as native_block_table
+
+from vllm_omni.diffusion.diffusion_kv import model_runner_backend as model_runner_backend_module
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.metadata import (
+    DiffusionKVContextMetadata,
+    DiffusionKVMetadata,
+    DiffusionKVSequenceMetadata,
+)
+from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
+from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _registration_backend(parallel_config: object):
+    model_config = SimpleNamespace(set_attention_geometry=Mock())
+    vllm_config = SimpleNamespace(
+        attention_config=SimpleNamespace(use_non_causal=False),
+        compilation_config=SimpleNamespace(static_forward_context={}),
+        model_config=model_config,
+    )
+    backend = DiffusionKVModelRunnerBackend(
+        vllm_config=vllm_config,
+        od_config=SimpleNamespace(parallel_config=parallel_config),
+        device=torch.device("cpu"),
+    )
+    return backend, vllm_config, model_config
+
+
+def _parallel_config(**overrides):
+    values = dict(ulysses_degree=1, ring_degree=1, allgather_degree=1, ulysses_mode="strict")
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def test_cache_layer_registration_installs_adapters_with_ulysses_local_geometry(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    backend, vllm_config, model_config = _registration_backend(_parallel_config(ulysses_degree=2))
+    layer = SimpleNamespace(
+        num_heads=8,
+        softmax_scale=0.125,
+        skip_sequence_parallel=False,
+        attn_backend=SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN"),
+    )
+    unrelated = object()
+    vllm_config.compilation_config.static_forward_context.update({"layer-0": layer, "unrelated": unrelated})
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=8,
+        dtype=torch.float16,
+        non_causal=True,
+    )
+    constructor_calls = []
+
+    class FakeLayerAdapter:
+        def __init__(self, *, layer_name, layer, spec, ulysses_degree, **_kwargs) -> None:
+            constructor_calls.append((layer_name, layer, spec, ulysses_degree))
+            self.num_heads = layer.num_heads // ulysses_degree
+            self.num_kv_heads = spec.num_kv_heads // ulysses_degree
+            self.head_size = spec.head_size
+            self.spec = replace(spec, num_kv_heads=self.num_kv_heads)
+
+    monkeypatch.setattr(model_runner_backend_module, "DiffusionPagedAttentionLayerAdapter", FakeLayerAdapter)
+    monkeypatch.setattr(model_runner_backend_module, "set_current_vllm_config", lambda _config: nullcontext())
+
+    specs = backend.register_kv_cache_layers({"layer-0": (layer, spec)})
+
+    installed = vllm_config.compilation_config.static_forward_context["layer-0"]
+    assert isinstance(installed, FakeLayerAdapter)
+    assert vllm_config.compilation_config.static_forward_context["unrelated"] is unrelated
+    assert specs["layer-0"].num_kv_heads == 2
+    assert constructor_calls == [("layer-0", layer, spec, 2)]
+    model_config.set_attention_geometry.assert_called_once_with(
+        num_heads=4,
+        num_kv_heads=2,
+        head_size=8,
+    )
+    assert vllm_config.attention_config.use_non_causal is True
+
+
+def test_cache_layer_registration_rejects_backend_without_paged_support() -> None:
+    backend, _, _ = _registration_backend(_parallel_config())
+    layer = SimpleNamespace(
+        attn_backend=SimpleNamespace(supports_paged_kv=False, get_name=lambda: "SDPA"),
+    )
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=8,
+        dtype=torch.float16,
+    )
+
+    with pytest.raises(NotImplementedError, match="layer-0.*paged support.*SDPA"):
+        backend.register_kv_cache_layers({"layer-0": (layer, spec)})
+
+
+def test_cache_layer_registration_resolves_platform_paged_hooks(monkeypatch: pytest.MonkeyPatch) -> None:
+    backend, _, _ = _registration_backend(_parallel_config())
+
+    def unsupported_block_tables():
+        raise NotImplementedError("platform has no paged BlockTables")
+
+    monkeypatch.setattr(
+        model_runner_backend_module.current_omni_platform,
+        "get_diffusion_kv_block_tables_cls",
+        unsupported_block_tables,
+    )
+
+    with pytest.raises(NotImplementedError, match="no paged BlockTables"):
+        backend.register_kv_cache_layers({"layer-0": (object(), Mock(spec=FullAttentionSpec))})
+
+
+@pytest.mark.parametrize(
+    ("parallel_config", "message"),
+    [
+        (_parallel_config(ring_degree=2), "Ring"),
+        (_parallel_config(allgather_degree=2), "AllGather-KV"),
+        (_parallel_config(ulysses_degree=2, ulysses_mode="advanced_uaa"), "strict Ulysses"),
+    ],
+)
+def test_cache_layer_registration_rejects_unsupported_sp_modes(parallel_config, message: str) -> None:
+    backend, _, _ = _registration_backend(parallel_config)
+
+    with pytest.raises(NotImplementedError, match=message):
+        backend._get_paged_attention_ulysses_degree(SimpleNamespace(skip_sequence_parallel=False))
+
+
+def test_cache_layer_registration_ignores_sp_for_opted_out_layer() -> None:
+    backend, _, _ = _registration_backend(
+        _parallel_config(ulysses_degree=2, ring_degree=2, ulysses_mode="advanced_uaa")
+    )
+
+    assert backend._get_paged_attention_ulysses_degree(SimpleNamespace(skip_sequence_parallel=True)) == 1
+
+
+class _FakeSpec:
+    def __init__(self, block_size: int) -> None:
+        self.block_size = block_size
+        self.capacity_calls: list[tuple[object, int]] = []
+
+    def max_num_blocks_per_req(self, vllm_config: object, max_len: int) -> int:
+        self.capacity_calls.append((vllm_config, max_len))
+        return math.ceil(max_len / self.block_size)
+
+
+class _FakeStagedTable:
+    def __init__(self) -> None:
+        self.staged: list[tuple[int, tuple[int, ...]]] = []
+        self.clear_calls = 0
+
+    def clear_staged_writes(self) -> None:
+        self.staged.clear()
+        self.clear_calls += 1
+
+
+class _FakeNumBlocks:
+    def __init__(self, num_groups: int, max_num_reqs: int) -> None:
+        self.np = np.zeros((num_groups, max_num_reqs), dtype=np.int32)
+        self.copy_calls = 0
+
+    def copy_to_uva(self) -> None:
+        self.copy_calls += 1
+
+
+class _FakeBlockTables:
+    instances: list[_FakeBlockTables] = []
+
+    def __init__(self, **kwargs: object) -> None:
+        self.kwargs = kwargs
+        self.block_sizes = list(kwargs["block_sizes"])
+        self.max_num_reqs = int(kwargs["max_num_reqs"])
+        self.max_num_batched_tokens = int(kwargs["max_num_batched_tokens"])
+        self.cp_size = int(kwargs["cp_size"])
+        self.cp_rank = int(kwargs["cp_rank"])
+        self.cp_interleave = int(kwargs["cp_interleave"])
+        self.blocks_per_kv_block = [1] * len(self.block_sizes)
+        self.block_tables = [_FakeStagedTable() for _ in self.block_sizes]
+        self.num_blocks = _FakeNumBlocks(len(self.block_sizes), self.max_num_reqs)
+        self.append_calls: list[tuple[int, tuple[list[int], ...], bool]] = []
+        self.apply_calls = 0
+        self.layout_refreshes = 0
+        self.fail_on_append_call: int | None = None
+        self.fail_apply = False
+        self.committed: list[dict[int, tuple[int, ...]]] = [dict() for _ in self.block_sizes]
+        self.__class__.instances.append(self)
+
+    def append_block_ids(
+        self,
+        req_index: int,
+        new_block_ids: tuple[list[int], ...],
+        overwrite: bool,
+    ) -> None:
+        call_number = len(self.append_calls) + 1
+        if self.fail_on_append_call == call_number:
+            raise RuntimeError("fake append failure")
+        copied_ids = tuple(list(group_ids) for group_ids in new_block_ids)
+        self.append_calls.append((req_index, copied_ids, overwrite))
+        for group_index, group_ids in enumerate(copied_ids):
+            self.num_blocks.np[group_index, req_index] = len(group_ids)
+            if group_ids:
+                self.block_tables[group_index].staged.append((req_index, tuple(group_ids)))
+
+    def apply_staged_writes(self) -> None:
+        self.apply_calls += 1
+        if self.fail_apply:
+            raise RuntimeError("fake apply failure")
+        for group_index, group_table in enumerate(self.block_tables):
+            for row, group_ids in group_table.staged:
+                self.committed[group_index][row] = group_ids
+            group_table.staged.clear()
+        self.num_blocks.copy_to_uva()
+
+    def init_block_table_layout_tensors(self) -> None:
+        self.layout_refreshes += 1
+
+
+def _metadata(
+    request_id: str = "req-0",
+    generation: int = 1,
+    *,
+    seq_len: int = 9,
+    sequence_block_ids: tuple[list[int], ...] = ([1, 2, 3], [4, 5]),
+    contexts: tuple[DiffusionKVContextMetadata, ...] | None = None,
+) -> DiffusionKVMetadata:
+    if contexts is None:
+        contexts = (
+            DiffusionKVContextMetadata(
+                context_id="text",
+                cache_role="cross_attention",
+                num_tokens=5,
+                block_ids=([6, 7], [8]),
+            ),
+        )
+    return DiffusionKVMetadata(
+        request_id=request_id,
+        allocation_generation=generation,
+        sequences=(
+            DiffusionKVSequenceMetadata(
+                sequence_id=0,
+                prefix_len=4,
+                target_len=5,
+                seq_len=seq_len,
+                block_ids=sequence_block_ids,
+                context_ids=tuple(context.context_id for context in contexts),
+            ),
+        ),
+        contexts=contexts,
+    )
+
+
+def _install_native_modules(monkeypatch: pytest.MonkeyPatch, events: list[tuple]) -> None:
+    _FakeBlockTables.instances.clear()
+
+    def get_block_table_width(max_num_blocks, block_size):
+        events.append(("width", max_num_blocks, block_size))
+        return max_num_blocks + 1
+
+    def init_backend(kv_cache_config, vllm_config, device):
+        events.append(("backend", kv_cache_config, vllm_config, device))
+        group = SimpleNamespace(get_metadata_builder=lambda _index: SimpleNamespace(reorder_batch_threshold=None))
+        return [[group], [group]], "cg-support", [4, 8]
+
+    def init_cache(
+        runner_kv_caches,
+        forward_context,
+        kv_cache_config,
+        attn_groups,
+        device,
+        cache_dtype,
+        kernel_block_sizes,
+        vllm_config,
+    ):
+        events.append(("cache", kv_cache_config, attn_groups, kernel_block_sizes))
+        runner_kv_caches.extend(["cache-0", "cache-1"])
+        return {"layer-0": "tensor-0", "layer-1": "tensor-1"}
+
+    monkeypatch.setattr(model_runner_backend_module, "init_attn_backend", init_backend)
+    monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", init_cache)
+    monkeypatch.setattr(
+        model_runner_backend_module.current_omni_platform,
+        "get_diffusion_kv_block_tables_cls",
+        lambda: _FakeBlockTables,
+    )
+    monkeypatch.setattr(native_block_table, "get_block_table_width", get_block_table_width, raising=False)
+
+
+def _runner(
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    max_rows_per_request: int = 3,
+    max_num_seqs: int = 2,
+    max_len: int = 16,
+    max_num_batched_tokens: int = 32,
+    cp_size: int = 1,
+    cp_rank: int = 0,
+    cp_interleave: int = 1,
+) -> tuple[DiffusionKVModelRunnerBackend, _FakeBlockTables, list[tuple]]:
+    events: list[tuple] = []
+    _install_native_modules(monkeypatch, events)
+    od_config = SimpleNamespace(
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+        diffusion_kv_max_rows_per_request=max_rows_per_request,
+        max_num_seqs=max_num_seqs,
+    )
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=max_len),
+        scheduler_config=SimpleNamespace(
+            max_num_batched_tokens=max_num_batched_tokens,
+            max_num_seqs=max_num_seqs,
+        ),
+        parallel_config=SimpleNamespace(
+            decode_context_parallel_size=cp_size,
+            cp_kv_cache_interleave_size=cp_interleave,
+        ),
+        compilation_config=SimpleNamespace(static_forward_context={"layer-0": object(), "layer-1": object()}),
+        cache_config=SimpleNamespace(cache_dtype="auto"),
+    )
+    monkeypatch.setattr(
+        model_runner_backend_module,
+        "get_dcp_group",
+        lambda: SimpleNamespace(rank_in_group=cp_rank),
+    )
+    backend = DiffusionKVModelRunnerBackend(
+        vllm_config=vllm_config,
+        od_config=od_config,
+        device="cuda:1",
+    )
+    backend._kv_cache_layer_adapters = {
+        "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None),
+        "layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None),
+    }
+
+    config = SimpleNamespace(
+        num_blocks=32,
+        kv_cache_groups=[
+            SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=_FakeSpec(4)),
+            SimpleNamespace(layer_names=["layer-1"], kv_cache_spec=_FakeSpec(8)),
+        ],
+    )
+    backend.initialize_kv_cache(config)
+    return backend, _FakeBlockTables.instances[-1], events
+
+
+def test_initialize_builds_native_block_tables_from_rank_local_config(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, events = _runner(monkeypatch)
+
+    assert events[:2] == [("width", 4, 4), ("width", 2, 8)]
+    assert [event[0] for event in events[2:]] == ["backend", "cache"]
+    assert block_tables.kwargs == {
+        "block_sizes": [4, 8],
+        "max_num_reqs": 6,
+        "max_num_batched_tokens": 32,
+        "max_num_blocks_per_group": [5, 3],
+        "device": "cuda:1",
+        "kernel_block_sizes": [4, 8],
+        "cp_size": 1,
+        "cp_rank": 0,
+        "cp_interleave": 1,
+    }
+    assert runner.block_tables is block_tables
+    assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
+    assert runner.kv_caches == ["cache-0", "cache-1"]
+    assert runner.vllm_config.scheduler_config.max_num_seqs == 6
+
+
+def test_initialize_propagates_native_context_parallel_layout(monkeypatch: pytest.MonkeyPatch) -> None:
+    _, block_tables, _ = _runner(
+        monkeypatch,
+        cp_size=2,
+        cp_rank=1,
+        cp_interleave=4,
+    )
+
+    assert block_tables.kwargs["cp_size"] == 2
+    assert block_tables.kwargs["cp_rank"] == 1
+    assert block_tables.kwargs["cp_interleave"] == 4
+
+
+def test_initialize_failure_does_not_publish_partial_cache_state(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    events: list[tuple] = []
+    _install_native_modules(monkeypatch, events)
+    successful_init_cache = model_runner_backend_module.init_kv_cache
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=16),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=1),
+        parallel_config=SimpleNamespace(
+            decode_context_parallel_size=1,
+            cp_kv_cache_interleave_size=1,
+        ),
+        compilation_config=SimpleNamespace(static_forward_context={"layer-0": object(), "layer-1": object()}),
+        cache_config=SimpleNamespace(cache_dtype="auto"),
+    )
+    backend = DiffusionKVModelRunnerBackend(
+        vllm_config=vllm_config,
+        od_config=SimpleNamespace(
+            diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+            diffusion_kv_max_rows_per_request=2,
+        ),
+        device="cuda:1",
+    )
+    backend._kv_cache_layer_adapters = {
+        "layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-0"),
+        "layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-1"),
+    }
+    config = SimpleNamespace(
+        num_blocks=32,
+        kv_cache_groups=[
+            SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=_FakeSpec(4)),
+            SimpleNamespace(layer_names=["layer-1"], kv_cache_spec=_FakeSpec(8)),
+        ],
+    )
+
+    def fail_init_cache(runner_kv_caches, *_args, **_kwargs):
+        runner_kv_caches.append("partial-cache")
+        backend._kv_cache_layer_adapters["layer-0"].kv_cache = "partial-binding"
+        raise RuntimeError("cache binding failed")
+
+    monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", fail_init_cache)
+    with pytest.raises(RuntimeError, match="cache binding failed"):
+        backend.initialize_kv_cache(config)
+
+    assert backend.kv_cache_config is None
+    assert backend.kv_caches == []
+    assert backend.block_tables is None
+    assert backend.paged_attention_adapter is None
+    assert backend._kv_cache_layer_adapters["layer-0"].kv_cache == "placeholder-0"
+    assert backend.vllm_config.scheduler_config.max_num_seqs == 1
+
+    monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", successful_init_cache)
+    backend.initialize_kv_cache(config)
+    assert backend.kv_cache_config is not None
+    assert backend.kv_caches == ["cache-0", "cache-1"]
+    assert backend.vllm_config.scheduler_config.max_num_seqs == 2
+
+
+def test_initialize_rejects_config_for_different_registered_layers() -> None:
+    backend = DiffusionKVModelRunnerBackend(
+        vllm_config=SimpleNamespace(),
+        od_config=SimpleNamespace(),
+        device="cuda:1",
+    )
+    backend._kv_cache_layer_adapters = {
+        "expected-layer": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None)
+    }
+    config = SimpleNamespace(
+        kv_cache_groups=[SimpleNamespace(layer_names=["other-layer"])],
+    )
+
+    with pytest.raises(ValueError, match="layer mismatch"):
+        backend.initialize_kv_cache(config)
+
+
+def test_valid_sequence_and_context_install_into_native_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+
+    assert runner.install_diffusion_kv_metadata(_metadata()) is True
+
+    assert block_tables.append_calls == [
+        (0, ([1, 2, 3], [4, 5]), True),
+        (1, ([6, 7], [8]), True),
+    ]
+    assert block_tables.apply_calls == 1
+    assert runner.get_diffusion_kv_row("req-0", 0) == 0
+    assert runner.get_diffusion_kv_row("req-0", 0, "text") == 1
+    sequence_binding = runner._resolve_paged_attention_row("req-0", 0, None)
+    context_binding = runner._resolve_paged_attention_row("req-0", None, "text")
+    assert (sequence_binding.row_index, sequence_binding.max_seq_len) == (0, 9)
+    assert (context_binding.row_index, context_binding.max_seq_len) == (1, 5)
+
+
+def test_block_table_mutations_invalidate_prepared_attention_batches(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    runner, _, _ = _runner(monkeypatch)
+    assert runner.paged_attention_adapter is not None
+    runner.paged_attention_adapter.invalidate_prepared_batches = Mock()
+
+    runner.install_diffusion_kv_metadata(_metadata())
+    runner.paged_attention_adapter.invalidate_prepared_batches.assert_called_once_with()
+
+    runner.remove_diffusion_kv_requests(["req-0"])
+    assert runner.paged_attention_adapter.invalidate_prepared_batches.call_count == 2
+
+    runner.refresh_block_table_layout()
+    assert runner.paged_attention_adapter.invalidate_prepared_batches.call_count == 3
+
+
+def test_cache_layer_registration_is_frozen_after_physical_initialization(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    runner, _, _ = _runner(monkeypatch)
+
+    with pytest.raises(RuntimeError, match="cannot be changed"):
+        runner.register_kv_cache_layers({})
+
+
+def test_request_scoped_context_uses_one_row_across_sequences(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    first_sequence = _metadata().sequences[0]
+    metadata = replace(
+        _metadata(),
+        sequences=(
+            first_sequence,
+            replace(
+                first_sequence,
+                sequence_id=1,
+                block_ids=([9, 10, 11], [12, 13]),
+            ),
+        ),
+    )
+
+    assert runner.install_diffusion_kv_metadata(metadata) is True
+
+    assert block_tables.append_calls == [
+        (0, ([1, 2, 3], [4, 5]), True),
+        (1, ([9, 10, 11], [12, 13]), True),
+        (2, ([6, 7], [8]), True),
+    ]
+    assert runner.get_diffusion_kv_row("req-0", 0, "text") == 2
+    assert runner.get_diffusion_kv_row("req-0", 1, "text") == 2
+
+
+def test_row_block_count_uses_native_spec_semantics(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    spec = runner.kv_cache_config.kv_cache_groups[0].kv_cache_spec
+    spec.max_num_blocks_per_req = lambda _config, token_len: math.ceil(token_len / 8)
+
+    metadata = _metadata(
+        sequence_block_ids=([1, 2], [4, 5]),
+        contexts=(),
+    )
+
+    assert runner.install_diffusion_kv_metadata(metadata) is True
+    assert block_tables.append_calls == [(0, ([1, 2], [4, 5]), True)]
+
+
+def test_repeated_native_block_ids_are_allowed(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    metadata = _metadata(
+        sequence_block_ids=([1, 1, 3], [4, 4]),
+        contexts=(),
+    )
+
+    assert runner.install_diffusion_kv_metadata(metadata) is True
+    assert block_tables.append_calls == [(0, ([1, 1, 3], [4, 4]), True)]
+
+
+def test_generation_is_idempotent_stale_safe_and_conflict_safe(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    metadata = _metadata()
+    runner.install_diffusion_kv_metadata(metadata)
+
+    assert runner.install_diffusion_kv_metadata(metadata) is False
+    assert runner.install_diffusion_kv_metadata(replace(metadata, allocation_generation=2)) is False
+    assert block_tables.apply_calls == 1
+
+    with pytest.raises(ValueError, match="Stale Diffusion KV allocation generation"):
+        runner.install_diffusion_kv_metadata(metadata)
+    with pytest.raises(ValueError, match="Conflicting Diffusion KV allocation snapshot"):
+        runner.install_diffusion_kv_metadata(_metadata(generation=3, sequence_block_ids=([9, 2, 3], [4, 5])))
+    assert block_tables.apply_calls == 1
+
+
+@pytest.mark.parametrize(
+    ("metadata", "message"),
+    [
+        (_metadata(sequence_block_ids=([1, 2, 3],)), "block groups"),
+        (_metadata(sequence_block_ids=([1, 2], [4, 5])), "expected 3"),
+        (_metadata(sequence_block_ids=([1, 2, 32], [4, 5])), "outside"),
+        (
+            _metadata(
+                seq_len=17,
+                sequence_block_ids=([1, 2, 3, 4, 5], [6, 7, 8]),
+                contexts=(),
+            ),
+            "exceeds row capacity",
+        ),
+    ],
+)
+def test_invalid_group_count_block_count_range_and_row_capacity_are_atomic(
+    monkeypatch: pytest.MonkeyPatch,
+    metadata: DiffusionKVMetadata,
+    message: str,
+) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+
+    with pytest.raises(ValueError, match=message):
+        runner.install_diffusion_kv_metadata(metadata)
+
+    assert block_tables.append_calls == []
+    assert block_tables.apply_calls == 0
+    assert runner._diffusion_kv_identity_to_row == {}
+    assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
+
+
+def test_invalid_later_context_does_not_stage_earlier_valid_row(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    bad_context = DiffusionKVContextMetadata(
+        context_id="text",
+        cache_role="cross_attention",
+        num_tokens=5,
+        block_ids=([6, 32], [8]),
+    )
+
+    with pytest.raises(ValueError, match="outside"):
+        runner.install_diffusion_kv_metadata(_metadata(contexts=(bad_context,)))
+
+    assert block_tables.append_calls == []
+    assert all(not group_table.staged for group_table in block_tables.block_tables)
+
+
+def test_native_append_failure_rolls_back_staged_state_and_registry(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    block_tables.fail_on_append_call = 2
+
+    with pytest.raises(RuntimeError, match="fake append failure"):
+        runner.install_diffusion_kv_metadata(_metadata())
+
+    assert np.count_nonzero(block_tables.num_blocks.np) == 0
+    assert all(not group_table.staged for group_table in block_tables.block_tables)
+    assert runner._diffusion_kv_identity_to_row == {}
+    assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
+
+
+def test_native_apply_failure_rolls_back_staged_state_and_registry(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    block_tables.fail_apply = True
+
+    with pytest.raises(RuntimeError, match="fake apply failure"):
+        runner.install_diffusion_kv_metadata(_metadata())
+
+    assert block_tables.apply_calls == 1
+    assert np.count_nonzero(block_tables.num_blocks.np) == 0
+    assert all(not group_table.staged for group_table in block_tables.block_tables)
+    assert all(group_table.clear_calls == 1 for group_table in block_tables.block_tables)
+    assert runner._diffusion_kv_identity_to_row == {}
+    assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
+
+
+def test_adapter_and_global_row_capacity_are_enforced(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch, max_rows_per_request=2, max_num_seqs=1)
+    second_context = DiffusionKVContextMetadata(
+        context_id="image",
+        cache_role="cross_attention",
+        num_tokens=1,
+        block_ids=([9], [10]),
+    )
+
+    with pytest.raises(ValueError, match="adapter limit is 2"):
+        runner.install_diffusion_kv_metadata(_metadata(contexts=(_metadata().contexts[0], second_context)))
+    assert block_tables.append_calls == []
+
+    runner.install_diffusion_kv_metadata(_metadata())
+    append_count = len(block_tables.append_calls)
+    with pytest.raises(ValueError, match="only 0 rows are free"):
+        runner.install_diffusion_kv_metadata(_metadata("req-1"))
+    assert len(block_tables.append_calls) == append_count
+
+
+def test_empty_and_duplicate_sequence_metadata_are_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    empty = DiffusionKVMetadata(request_id="req-0", allocation_generation=0, sequences=())
+    duplicate_sequence = _metadata().sequences[0]
+    duplicate = DiffusionKVMetadata(
+        request_id="req-0",
+        allocation_generation=0,
+        sequences=(duplicate_sequence, duplicate_sequence),
+    )
+
+    with pytest.raises(ValueError, match="at least one sequence"):
+        runner.install_diffusion_kv_metadata(empty)
+    with pytest.raises(ValueError, match="Duplicate diffusion KV row identity"):
+        runner.install_diffusion_kv_metadata(duplicate)
+    assert block_tables.append_calls == []
+
+
+@pytest.mark.parametrize(
+    ("metadata", "message"),
+    [
+        (_metadata(generation=-1), "non-negative integer"),
+        (_metadata(request_id=""), "non-empty string"),
+        (_metadata(sequence_block_ids=([1, 2, True], [4, 5])), "non-integer block ID"),
+    ],
+)
+def test_request_generation_and_block_id_types_are_validated(
+    monkeypatch: pytest.MonkeyPatch,
+    metadata: DiffusionKVMetadata,
+    message: str,
+) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+
+    with pytest.raises(ValueError, match=message):
+        runner.install_diffusion_kv_metadata(metadata)
+
+    assert block_tables.append_calls == []
+    assert runner._diffusion_kv_request_states == {}
+
+
+def test_cleanup_is_idempotent_and_reuses_native_rows(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch, max_rows_per_request=2, max_num_seqs=1)
+    runner.install_diffusion_kv_metadata(_metadata())
+
+    assert runner.remove_diffusion_kv_requests(["missing"]) == 0
+    assert runner.remove_diffusion_kv_requests(["req-0", "req-0"]) == 2
+    assert block_tables.append_calls[-2:] == [(0, ([], []), True), (1, ([], []), True)]
+    assert block_tables.num_blocks.np[:, :2].tolist() == [[0, 0], [0, 0]]
+    assert runner.remove_diffusion_kv_requests(["req-0"]) == 0
+    with pytest.raises(KeyError, match="No native diffusion KV row"):
+        runner.get_diffusion_kv_row("req-0", 0)
+
+    runner.install_diffusion_kv_metadata(_metadata("req-1"))
+    assert runner.get_diffusion_kv_row("req-1", 0) == 0
+    assert runner.get_diffusion_kv_row("req-1", 0, "text") == 1
+
+
+def test_cached_step_does_not_reinstall_metadata() -> None:
+    runner = object.__new__(DiffusionModelRunner)
+    runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
+    runner.pipeline = object()
+    runner.install_diffusion_kv_metadata = Mock()
+    runner._supports_step_mode = Mock(return_value=False)
+    scheduler_output = DiffusionSchedulerOutput(
+        step_id=1,
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData(request_ids=["req-0"]),
+        finished_req_ids=set(),
+        num_running_reqs=1,
+        num_waiting_reqs=0,
+    )
+
+    with pytest.raises(ValueError, match="does not support step execution"):
+        runner.execute_stepwise(scheduler_output)
+    runner.install_diffusion_kv_metadata.assert_not_called()
+
+
+def test_dense_cleanup_and_wake_are_noops_but_paged_wake_refreshes(monkeypatch: pytest.MonkeyPatch) -> None:
+    runner, block_tables, _ = _runner(monkeypatch)
+    runner.refresh_block_table_layout()
+    assert block_tables.layout_refreshes == 1
+
+    dense_backend = DiffusionKVModelRunnerBackend(
+        vllm_config=SimpleNamespace(),
+        od_config=SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY),
+        device="cuda:1",
+    )
+    dense_backend.block_tables = Mock()
+    assert dense_backend.remove_diffusion_kv_requests(["req-0"]) == 0
+    dense_backend.refresh_block_table_layout()
+    dense_backend.block_tables.init_block_table_layout_tensors.assert_not_called()
+    with pytest.raises(ValueError, match="Dense diffusion execution"):
+        dense_backend.install_diffusion_kv_metadata(_metadata())

tests/diffusion/diffusion_kv/test_config.py

@@ -3,6 +3,7 @@
 
 import pytest
 
+from vllm_omni.config.omni_config import _DiffusionConfigProjection
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 
@@ -39,6 +40,8 @@ def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
     with pytest.raises(ValueError, match="does not support imported AR KV"):
         OmniDiffusionConfig.from_kwargs(
             diffusion_kv_mode="paged_scheduler",
+            diffusion_kv_max_rows_per_request=1,
+            max_num_batched_tokens=1,
             omni_kv_config={"need_recv_cache": True},
         )
 
@@ -47,6 +50,33 @@ def test_paged_scheduler_does_not_depend_on_model_registry() -> None:
     config = OmniDiffusionConfig.from_kwargs(
         model_class_name="FutureDiffusionModel",
         diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=1,
+        max_num_batched_tokens=1,
     )
 
     assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
+
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+def test_paged_scheduler_mode_is_platform_agnostic(config_cls) -> None:
+    config = config_cls.from_kwargs(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=1,
+    )
+
+    assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
+
+def test_paged_scheduler_requires_a_worker_row_limit() -> None:
+    with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"):
+        OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_scheduler")
+    with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"):
+        _DiffusionConfigProjection.from_kwargs(diffusion_kv_mode="paged_scheduler")
+
+
+@pytest.mark.parametrize("invalid_limit", [0, -1, True, 1.5])
+def test_worker_row_limit_must_be_a_positive_integer(invalid_limit: object) -> None:
+    with pytest.raises((TypeError, ValueError)):
+        OmniDiffusionConfig.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit)
+    with pytest.raises((TypeError, ValueError)):
+        _DiffusionConfigProjection.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit)

tests/diffusion/diffusion_kv/test_initialization.py

@@ -22,6 +22,7 @@ class _Executor:
         self.specs = [{"layer0": object()}, {"layer0": object()}]
         self.memory = [1024, 1024]
         self.configs = None
+        self.resolved_max_model_len = None
         self.profile_requests = None
 
     def get_kv_cache_specs(self):
@@ -31,8 +32,9 @@ class _Executor:
         self.profile_requests = profile_requests
         return self.memory
 
-    def set_kv_cache_configs(self, configs) -> None:
+    def set_kv_cache_configs(self, configs, resolved_max_model_len) -> None:
         self.configs = configs
+        self.resolved_max_model_len = resolved_max_model_len
 
 
 def _od_config(**overrides):
@@ -55,6 +57,7 @@ def _od_config(**overrides):
         ),
         gpu_memory_utilization=0.9,
         kv_cache_memory_bytes=None,
+        max_num_seqs=1,
         max_num_batched_tokens=64,
         num_gpus=1,
     )
@@ -64,9 +67,9 @@ def _od_config(**overrides):
 
 def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None:
     executor = _Executor()
-    od_config = _od_config(num_gpus=2)
+    od_config = _od_config(num_gpus=2, max_model_len=-1)
     vllm_config = SimpleNamespace(
-        model_config=SimpleNamespace(max_model_len=4096),
+        model_config=SimpleNamespace(max_model_len=256),
         max_in_flight_tokens=128,
     )
     worker_configs = [object(), object()]
@@ -91,6 +94,7 @@ def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None:
     build.assert_called_once_with(vllm_config, executor.specs, executor.memory)
     assert result == (scheduler_kv_cache_config, 16, 16, vllm_config)
     assert executor.configs == worker_configs
+    assert executor.resolved_max_model_len == vllm_config.model_config.max_model_len
     assert executor.profile_requests is profile_requests
 
 

tests/diffusion/diffusion_kv/test_native_block_tables_gpu.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from importlib import import_module, util
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
+
+
+def test_native_block_tables_slot_mapping_handles_non_aligned_tail() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for the native BlockTables slot-mapping smoke test")
+    module_name = "vllm.v1.worker.gpu.block_table"
+    try:
+        module_spec = util.find_spec(module_name)
+    except ModuleNotFoundError as exc:
+        if exc.name is None or not module_name.startswith(f"{exc.name}."):
+            raise
+        pytest.skip("installed vLLM does not provide GPU BlockTables")
+    if module_spec is None:
+        pytest.skip("installed vLLM does not provide GPU BlockTables")
+    block_table_module = import_module(module_name)
+    block_tables_cls = getattr(block_table_module, "BlockTables", None)
+    if block_tables_cls is None or not hasattr(block_tables_cls, "compute_slot_mappings"):
+        pytest.skip("installed vLLM does not provide the required BlockTables API")
+
+    device = torch.device("cuda", torch.accelerator.current_device_index())
+    block_tables = block_tables_cls(
+        block_sizes=[4],
+        max_num_reqs=1,
+        max_num_batched_tokens=16,
+        max_num_blocks_per_group=[4],
+        device=device,
+        kernel_block_sizes=[4],
+    )
+    block_tables.append_block_ids(0, ([1, 2, 3],), overwrite=True)
+    block_tables.apply_staged_writes()
+    idx_mapping = torch.tensor([0], dtype=torch.int32, device=device)
+    query_start_loc = torch.tensor([0, 9], dtype=torch.int32, device=device)
+    positions = torch.arange(9, dtype=torch.int64, device=device)
+
+    slot_mappings = block_tables.compute_slot_mappings(
+        idx_mapping,
+        query_start_loc,
+        positions,
+        num_tokens_padded=9,
+    )
+    torch.accelerator.synchronize(device)
+
+    assert slot_mappings.cpu().tolist() == [[4, 5, 6, 7, 8, 9, 10, 11, 12]]

tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py

@@ -0,0 +1,251 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from torch.nn import functional as F
+from vllm.config import AttentionConfig, CacheConfig, VllmConfig, set_current_vllm_config
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+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.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
+from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
+    DiffusionPagedAttentionAdapter,
+    DiffusionPagedAttentionLayerAdapter,
+    DiffusionPagedAttentionRow,
+    DiffusionPagedAttentionRowBinding,
+)
+from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
+
+_LAYER_NAME = "model.layers.0.attn"
+_NUM_HEADS = 2
+_HEAD_SIZE = 64
+_BLOCK_SIZE = 16
+
+
+class _SmokeDiffusionAttention(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.num_heads = _NUM_HEADS
+        self.num_kv_heads = _NUM_HEADS
+        self.head_size = _HEAD_SIZE
+        self.softmax_scale = _HEAD_SIZE**-0.5
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention")
+def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> None:
+    device = torch.device("cuda", torch.accelerator.current_device_index())
+    vllm_config = VllmConfig(
+        cache_config=CacheConfig(
+            block_size=_BLOCK_SIZE,
+            cache_dtype="float16",
+            enable_prefix_caching=False,
+        ),
+        attention_config=AttentionConfig(use_non_causal=True),
+    )
+    model_config = _DiffusionVllmModelConfig(
+        model=None,
+        dtype=torch.float16,
+        max_model_len=32,
+        original_max_model_len=32,
+        hf_config=SimpleNamespace(model_type="diffusion_paged_attention_smoke"),
+    )
+    vllm_config.model_config = model_config
+    model_config.set_attention_geometry(
+        num_heads=_NUM_HEADS,
+        num_kv_heads=_NUM_HEADS,
+        head_size=_HEAD_SIZE,
+    )
+    vllm_config.compilation_config.static_forward_context.clear()
+
+    diffusion_layer = _SmokeDiffusionAttention().to(device)
+    spec = FullAttentionSpec(
+        block_size=_BLOCK_SIZE,
+        num_kv_heads=_NUM_HEADS,
+        head_size=_HEAD_SIZE,
+        dtype=torch.float16,
+        non_causal=True,
+    )
+    with set_current_vllm_config(vllm_config):
+        native_layer = DiffusionPagedAttentionLayerAdapter(
+            layer_name=_LAYER_NAME,
+            layer=diffusion_layer,
+            spec=spec,
+            vllm_config=vllm_config,
+            device=device,
+        )
+    vllm_config.compilation_config.static_forward_context[_LAYER_NAME] = native_layer
+
+    canonical_spec = native_layer.spec
+    kv_cache_config = KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=canonical_spec.page_size_bytes * 4,
+                shared_by=[_LAYER_NAME],
+            )
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                layer_names=[_LAYER_NAME],
+                kv_cache_spec=canonical_spec,
+            )
+        ],
+    )
+    with set_current_vllm_config(vllm_config):
+        attn_groups, _, kernel_block_sizes = init_attn_backend(
+            kv_cache_config,
+            vllm_config,
+            device,
+        )
+        runner_kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
+        init_kv_cache(
+            runner_kv_caches,
+            vllm_config.compilation_config.static_forward_context,
+            kv_cache_config,
+            attn_groups,
+            device,
+            vllm_config.cache_config.cache_dtype,
+            kernel_block_sizes,
+            vllm_config,
+        )
+        block_tables = BlockTables(
+            block_sizes=[_BLOCK_SIZE],
+            max_num_reqs=1,
+            max_num_batched_tokens=32,
+            max_num_blocks_per_group=[2],
+            device=device,
+            kernel_block_sizes=kernel_block_sizes,
+        )
+
+    block_tables.append_block_ids(0, ([3, 1],), overwrite=True)
+    block_tables.apply_staged_writes()
+    adapter = DiffusionPagedAttentionAdapter(
+        vllm_config=vllm_config,
+        device=device,
+        kv_cache_config=kv_cache_config,
+        block_tables=block_tables,
+        attn_groups=attn_groups,
+        layers={_LAYER_NAME: native_layer},
+        resolve_row=lambda _request_id, _sequence_id, _context_id: DiffusionPagedAttentionRowBinding(
+            row_index=0,
+            max_seq_len=19,
+        ),
+    )
+    prefix_batch = adapter.prepare_batch(
+        [
+            DiffusionPagedAttentionRow(
+                request_id="req-0",
+                sequence_id=0,
+                query_len=17,
+                seq_len=17,
+            )
+        ]
+    )
+
+    torch.manual_seed(1)
+    query = torch.randn(19, _NUM_HEADS, _HEAD_SIZE, dtype=torch.float16, device=device)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+    omni_backend = FlashAttentionImpl(
+        num_heads=_NUM_HEADS,
+        head_size=_HEAD_SIZE,
+        softmax_scale=diffusion_layer.softmax_scale,
+        num_kv_heads=_NUM_HEADS,
+    )
+    with adapter.activate(prefix_batch):
+        prefix_context = adapter.prepare_layer_context(_LAYER_NAME, query[:17], key[:17], value[:17])
+        prefix_output = omni_backend.forward_paged(prefix_context)
+    prefix_slot_mappings = prefix_batch.slot_mappings.clone()
+
+    suffix_batch = adapter.prepare_batch(
+        [
+            DiffusionPagedAttentionRow(
+                request_id="req-0",
+                sequence_id=0,
+                query_len=2,
+                seq_len=19,
+                kv_start_pos=17,
+            )
+        ]
+    )
+    with adapter.activate(suffix_batch):
+        suffix_context = adapter.prepare_layer_context(_LAYER_NAME, query[17:], key[17:], value[17:])
+        suffix_output = omni_backend.forward_paged(suffix_context)
+    suffix_slot_mappings = suffix_batch.slot_mappings.clone()
+
+    piecewise_batch = adapter.prepare_batch(
+        [
+            DiffusionPagedAttentionRow(
+                request_id="req-0",
+                sequence_id=0,
+                query_len=19,
+                seq_len=19,
+            )
+        ]
+    )
+    mixed_mask = torch.ones(19, 19, dtype=torch.bool, device=device).tril()
+    mixed_mask[5:10, 5:10] = True
+    mixed_metadata = AttentionMetadata(
+        attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0),
+        full_attn_spans=[[(5, 10)]],
+    )
+    with adapter.activate(piecewise_batch):
+        piecewise_context = adapter.prepare_layer_context(
+            _LAYER_NAME,
+            query,
+            key,
+            value,
+            omni_attn_metadata=mixed_metadata,
+        )
+        piecewise_output = omni_backend.forward_paged(piecewise_context)
+
+    prefix_reference = (
+        F.scaled_dot_product_attention(
+            query[:17].transpose(0, 1).unsqueeze(0),
+            key[:17].transpose(0, 1).unsqueeze(0),
+            value[:17].transpose(0, 1).unsqueeze(0),
+            is_causal=False,
+            scale=diffusion_layer.softmax_scale,
+        )
+        .squeeze(0)
+        .transpose(0, 1)
+    )
+    suffix_reference = (
+        F.scaled_dot_product_attention(
+            query[17:].transpose(0, 1).unsqueeze(0),
+            key.transpose(0, 1).unsqueeze(0),
+            value.transpose(0, 1).unsqueeze(0),
+            is_causal=False,
+            scale=diffusion_layer.softmax_scale,
+        )
+        .squeeze(0)
+        .transpose(0, 1)
+    )
+    piecewise_reference = (
+        F.scaled_dot_product_attention(
+            query.transpose(0, 1).unsqueeze(0),
+            key.transpose(0, 1).unsqueeze(0),
+            value.transpose(0, 1).unsqueeze(0),
+            attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0),
+            scale=diffusion_layer.softmax_scale,
+        )
+        .squeeze(0)
+        .transpose(0, 1)
+    )
+    torch.accelerator.synchronize(device)
+
+    assert prefix_slot_mappings.cpu().tolist() == [list(range(3 * _BLOCK_SIZE, 4 * _BLOCK_SIZE)) + [_BLOCK_SIZE]]
+    assert suffix_slot_mappings.cpu().tolist() == [[_BLOCK_SIZE + 1, _BLOCK_SIZE + 2]]
+    torch.testing.assert_close(prefix_output, prefix_reference, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(suffix_output, suffix_reference, rtol=2e-2, atol=2e-2)
+    torch.testing.assert_close(piecewise_output, piecewise_reference, rtol=2e-2, atol=2e-2)

tests/diffusion/diffusion_kv/test_spec_discovery.py

@@ -25,6 +25,23 @@ class _Backend:
         return True
 
 
+class _RunnerKVBackend:
+    def __init__(self) -> None:
+        self.kv_cache_config = None
+
+    def register_kv_cache_layers(self, layers):
+        return {layer_name: spec for layer_name, (_layer, spec) in layers.items()}
+
+    def initialize_kv_cache(self, config) -> None:
+        configured_layers = {layer_name for group in config.kv_cache_groups for layer_name in group.layer_names}
+        if configured_layers != {"image_attention"}:
+            raise ValueError(
+                "Rank-local Diffusion KVCacheConfig layer mismatch: "
+                f"expected=['image_attention'], configured={sorted(configured_layers)}"
+            )
+        self.kv_cache_config = config
+
+
 def _attention(*, enabled: bool) -> Attention:
     attention = Attention.__new__(Attention)
     nn.Module.__init__(attention)
@@ -47,6 +64,7 @@ def _runner(attention: Attention) -> DiffusionModelRunner:
     pipeline = nn.Module()
     pipeline.image_attention = attention
     runner.pipeline = pipeline
+    runner.diffusion_kv_backend = _RunnerKVBackend()
     return runner
 
 
@@ -102,14 +120,21 @@ def test_worker_selects_its_rank_local_config() -> None:
     worker = object.__new__(DiffusionWorker)
     worker.rank = 1
     worker.od_config = SimpleNamespace(num_gpus=2)
+    worker.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=None),
+        cache_config=SimpleNamespace(num_gpu_blocks=None),
+    )
+    worker._maybe_get_memory_pool_context = lambda _tag: nullcontext()
     worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config))
-    configs = [object(), object()]
+    configs = [SimpleNamespace(num_blocks=4), SimpleNamespace(num_blocks=8)]
 
-    worker.set_kv_cache_configs(configs)
+    worker.set_kv_cache_configs(configs, 64)
 
     assert worker.installed is configs[1]
+    assert worker.vllm_config.model_config.max_model_len == 64
+    assert worker.vllm_config.cache_config.num_gpu_blocks == 8
     with pytest.raises(ValueError, match="rank count mismatch"):
-        worker.set_kv_cache_configs(configs[:1])
+        worker.set_kv_cache_configs(configs[:1], 64)
 
 
 def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None:

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -19,8 +19,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 _TRANSFORMER_MODULE = "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer"
 
-NUM_HEADS = 4
-NUM_KV_HEADS = 2
+NUM_HEADS = 32
+NUM_KV_HEADS = 8
 HEAD_DIM = 16
 IMAGE_TOKEN_LEN = 8
 SUFFIX_TOKEN_LEN = 3
@@ -35,8 +35,14 @@ SCALING = 1.0 / math.sqrt(HEAD_DIM)
 class MockAttention(nn.Module):
     def __init__(self, num_heads, head_size, causal=False, softmax_scale=None, num_kv_heads=None, **kwargs):
         super().__init__()
+        self.paged_kv_active = False
+        self.calls = []
+
+    def is_paged_kv_active(self):
+        return self.paged_kv_active
 
     def forward(self, query, key, value, attn_metadata=None, **kwargs):
+        self.calls.append((query, key, value))
         return query
 
 
@@ -126,6 +132,37 @@ def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
     assert mgr.state_dict() == {}
 
 
+@pytest.mark.parametrize(
+    ("paged_kv_active", "expected_kv_heads"),
+    [(False, NUM_HEADS), (True, NUM_KV_HEADS)],
+)
+def test_gqa_kv_stays_compressed_for_paged_attention(
+    paged_kv_active: bool,
+    expected_kv_heads: int,
+) -> None:
+    mgr = _make_cache_mgr()
+    mgr.attn.paged_kv_active = paged_kv_active
+    q_len = 3
+    key, value = _make_known_kv(q_len)
+
+    _call_mgr(
+        mgr,
+        bs=1,
+        q_len=q_len,
+        seq_len=q_len,
+        key_flat=key,
+        value_flat=value,
+        first_step=True,
+        num_image_tokens=0,
+        gen_timestep_scatter_index=_gen_timestep_index(1, 0),
+    )
+
+    _, key_input, value_input = mgr.attn.calls[-1]
+    expected_shape = (1, q_len, expected_kv_heads, HEAD_DIM)
+    assert key_input.shape == expected_shape
+    assert value_input.shape == expected_shape
+
+
 # ============================================================
 # Test 1: No AR KV — basic cache → reuse
 # ============================================================

tests/diffusion/test_diffusion_config_propagation.py

@@ -149,10 +149,18 @@ class TestCreateDefaultDiffusion:
         assert ea["enforce_eager"] is True
         assert ea["lora_path"] == "/tmp/lora"
 
-    def test_diffusion_kv_mode_roundtrip(self):
-        od = _roundtrip_diffusion_config(model="x", diffusion_kv_mode="paged_scheduler")
+    def test_diffusion_kv_mode_roundtrip(self, monkeypatch):
+        from vllm_omni.platforms import current_omni_platform
+
+        monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+        od = _roundtrip_diffusion_config(
+            model="x",
+            diffusion_kv_mode="paged_scheduler",
+            diffusion_kv_max_rows_per_request=2,
+        )
 
         assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+        assert od.diffusion_kv_max_rows_per_request == 2
 
     def test_diffusion_kv_sizing_fields_roundtrip(self, monkeypatch):
         monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda self: 29500)

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -12,6 +12,7 @@ import pytest
 from vllm_omni.diffusion import diffusion_engine as diffusion_engine_module
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -88,6 +89,50 @@ def test_emit_step_outputs_finalizes_finished_request_without_stream() -> None:
     assert engine.scheduler.get_request_state(request_id) is None
 
 
+def _make_kv_cleanup_engine(mode: DiffusionKVCacheMode) -> tuple[DiffusionEngine, Mock]:
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(diffusion_kv_mode=mode)
+    cleanup = Mock()
+    engine.executor = SimpleNamespace(remove_diffusion_kv_requests=cleanup)
+    return engine, cleanup
+
+
+def test_paged_terminal_output_clears_worker_rows() -> None:
+    engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER)
+    engine._finalize_finished_request = lambda request_id, *_args: request_id
+    engine._put_output = lambda *_args: None
+
+    engine._emit_finished_outputs({"req-0", "req-1"})
+
+    cleanup.assert_called_once()
+    assert set(cleanup.call_args.args[0]) == {"req-0", "req-1"}
+
+
+def test_paged_abort_clears_worker_rows_before_scheduler_free() -> None:
+    engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER)
+    engine.scheduler = SimpleNamespace(
+        get_request_state=lambda _request_id: None,
+        finish_requests=Mock(side_effect=AssertionError("unknown request must not be logically freed")),
+    )
+
+    engine._abort_requests(["req-idle", "req-idle"])
+
+    cleanup.assert_called_once_with(["req-idle"])
+    engine.scheduler.finish_requests.assert_not_called()
+
+
+def test_dense_terminal_and_abort_paths_skip_worker_row_cleanup() -> None:
+    engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.DENSE_LEGACY)
+    engine._finalize_finished_request = lambda request_id, *_args: request_id
+    engine._put_output = lambda *_args: None
+    engine.scheduler = SimpleNamespace(get_request_state=lambda _request_id: None)
+
+    engine._emit_finished_outputs({"req-0"})
+    engine._abort_requests(["req-idle"])
+
+    cleanup.assert_not_called()
+
+
 def test_init_accepts_custom_scheduler(monkeypatch: pytest.MonkeyPatch) -> None:
     od_config = SimpleNamespace(
         custom_pipeline_args=None,

tests/diffusion/test_diffusion_model_runner.py

@@ -11,6 +11,7 @@ import torch
 import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
 from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
 
@@ -193,6 +194,11 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
         parallel_config=SimpleNamespace(use_hsdp=False),
         streaming_output=False,
     )
+    runner.diffusion_kv_backend = DiffusionKVModelRunnerBackend(
+        vllm_config=runner.vllm_config,
+        od_config=runner.od_config,
+        device=runner.device,
+    )
     runner.kv_transfer_manager = SimpleNamespace(
         receive_kv_cache=lambda req, target_device=None: None,
         receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None,

tests/diffusion/test_diffusion_plugin_hooks.py

@@ -40,6 +40,12 @@ class TestPlatformDiffusionHooks:
         result = OmniPlatform.get_diffusion_model_runner_cls()
         assert result == "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner"
 
+    def test_get_diffusion_kv_block_tables_cls_default(self):
+        """Test the default paged-KV BlockTables implementation."""
+        from vllm.v1.worker.gpu.block_table import BlockTables
+
+        assert OmniPlatform.get_diffusion_kv_block_tables_cls() is BlockTables
+
     def test_oot_enum_exists(self):
         """Test OOT is a valid platform enum value."""
         assert OmniPlatformEnum.OOT.value == "oot"

tests/diffusion/test_diffusion_step_pipeline.py

@@ -17,6 +17,8 @@ import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
 from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -306,11 +308,14 @@ def _make_vllm_config():
     )
 
 
-def _make_runner():
+def _make_runner(
+    diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY,
+):
     runner = object.__new__(DiffusionModelRunner)
     runner.vllm_config = _make_vllm_config()
     runner.od_config = SimpleNamespace(
         cache_backend=None,
+        diffusion_kv_mode=diffusion_kv_mode,
         parallel_config=SimpleNamespace(use_hsdp=False),
         streaming_output=False,
     )
@@ -330,6 +335,7 @@ def _make_distributed_runner(mode: str, device: torch.device):
     runner.vllm_config = _make_vllm_config()
     runner.od_config = SimpleNamespace(
         cache_backend=None,
+        diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
         parallel_config=SimpleNamespace(use_hsdp=False),
         streaming_output=False,
     )
@@ -502,6 +508,81 @@ def test_step_profiler_reports_denoise_step_as_diffuse():
 class TestRunner:
     """DiffusionModelRunner.execute_stepwise"""
 
+    @pytest.mark.parametrize(
+        ("diffusion_kv_mode", "should_cleanup"),
+        [
+            (DiffusionKVCacheMode.DENSE_LEGACY, False),
+            (DiffusionKVCacheMode.PAGED_SCHEDULER, True),
+        ],
+    )
+    def test_scheduler_finished_cleanup_only_runs_for_paged_kv(
+        self,
+        mocker: MockerFixture,
+        diffusion_kv_mode: DiffusionKVCacheMode,
+        should_cleanup: bool,
+    ):
+        runner = _make_runner(diffusion_kv_mode)
+        cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests")
+        expected_output = object()
+        mocker.patch.object(runner, "_execute_stepwise_core", return_value=expected_output)
+        scheduler_output = DiffusionSchedulerOutput(
+            step_id=1,
+            scheduled_new_reqs=[],
+            scheduled_cached_reqs=CachedRequestData.make_empty(),
+            finished_req_ids={"finished-request"},
+            num_running_reqs=0,
+            num_waiting_reqs=0,
+        )
+
+        output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+
+        assert output is expected_output
+        if should_cleanup:
+            cleanup.assert_called_once_with(["finished-request"])
+        else:
+            cleanup.assert_not_called()
+
+    @pytest.mark.parametrize(
+        ("diffusion_kv_mode", "should_cleanup"),
+        [
+            (DiffusionKVCacheMode.DENSE_LEGACY, False),
+            (DiffusionKVCacheMode.PAGED_SCHEDULER, True),
+        ],
+    )
+    def test_terminal_cleanup_only_runs_for_paged_kv(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        mocker: MockerFixture,
+        diffusion_kv_mode: DiffusionKVCacheMode,
+        should_cleanup: bool,
+    ):
+        runner = _make_runner(diffusion_kv_mode)
+        cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests")
+        install = mocker.patch.object(runner, "install_diffusion_kv_metadata")
+        req = _make_step_request()
+        scheduler_output = _make_scheduler_output(req, step_id=0)
+        if should_cleanup:
+            scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata = DiffusionKVMetadata(
+                request_id=req.request_id,
+                allocation_generation=1,
+                sequences=(),
+            )
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+
+        first_output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+        assert first_output.get_request_output(req.request_id).finished is False
+        cleanup.assert_not_called()
+
+        output = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
+
+        assert output.get_request_output(req.request_id).finished is True
+        if should_cleanup:
+            install.assert_called_once_with(scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata)
+            cleanup.assert_called_once_with([req.request_id])
+        else:
+            install.assert_not_called()
+            cleanup.assert_not_called()
+
     def test_completes_request_and_clears_state(self, monkeypatch):
         runner = _make_runner()
         req = _make_step_request()

vllm_omni/config/omni_config.py

@@ -678,6 +678,7 @@ class _DiffusionConfigProjection:
     cache_config: Any = field(default_factory=dict)
     enable_cache_dit_summary: bool = False
     diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
+    diffusion_kv_max_rows_per_request: int | None = Field(default=None, ge=1, strict=True)
     enable_prompt_embed_cache: bool = False
     prompt_embed_cache_size: int = Field(default=32, ge=1)
     enable_session_state_manager: bool = False
@@ -816,6 +817,11 @@ class _DiffusionConfigProjection:
             )
 
         self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
+        if (
+            self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+            and self.diffusion_kv_max_rows_per_request is None
+        ):
+            raise ValueError("paged_scheduler requires diffusion_kv_max_rows_per_request to be set")
         self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps)
         self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers)
 

vllm_omni/diffusion/attention/backends/abstract.py

@@ -15,6 +15,11 @@ class AttentionBackend(ABC):
 
     accept_output_buffer: bool = False
     supports_piecewise_spans: bool = False
+    # A backend that supports this capability can consume the opaque paged-KV
+    # context prepared by the diffusion Worker data plane.  Keeping the
+    # capability on the backend class prevents a paged request from silently
+    # falling back to dense attention on an incompatible implementation.
+    supports_paged_kv: bool = False
     # The backend can represent a contiguous valid K/V prefix by slicing the
     # tensors instead of materializing a padding mask. Models may use this to
     # avoid a slower masked-attention plan when tail padding is not semantic.
@@ -221,6 +226,18 @@ class AttentionImpl(ABC, Generic[T]):
         else:
             raise NotImplementedError(f"No forward implementation for platform: {current_omni_platform}")
 
+    def forward_paged(self, paged_kv_context: Any) -> torch.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.
+        """
+
+        del paged_kv_context
+        raise NotImplementedError(f"{type(self).__name__} does not support paged KV attention")
+
     def forward_cuda(
         self,
         query: torch.Tensor,

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -15,6 +15,7 @@ from vllm_omni.diffusion.attention.backends.abstract import (
 from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
 from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     piecewise_attn,
+    run_paged_piecewise_plan,
 )
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
 from vllm_omni.platforms import current_omni_platform
@@ -25,6 +26,7 @@ logger = init_logger(__name__)
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
     supports_piecewise_spans: bool = True
+    supports_paged_kv: bool = True
 
     @classmethod
     def supports_packed_mask_free(cls) -> bool:
@@ -246,6 +248,73 @@ class FlashAttentionImpl(AttentionImpl):
         # (b s) h d -> b s h d
         return out.reshape(batch_size, q_len, *out.shape[1:])
 
+    def forward_paged(self, paged_kv_context) -> torch.Tensor:
+        """Run platform-native paged attention through the Omni backend.
+
+        ``DiffusionPagedAttentionAdapter`` owns BlockTables and prepares the
+        rank-local native cache context.  The adapter no longer performs the
+        attention call itself; this method is the backend-owned execution
+        boundary.  Its native layer wrapper keeps vLLM version-specific cache
+        and kernel details out of Omni's common ``Attention`` layer.
+        """
+
+        layer = getattr(paged_kv_context, "layer", None)
+        if layer is None:
+            raise TypeError("paged_kv_context must expose a native diffusion attention layer")
+        kv_cache = getattr(layer, "kv_cache", None)
+        if kv_cache is None:
+            raise RuntimeError(f"Native KV cache is not bound for diffusion layer {layer.layer_name!r}")
+        native_impl = getattr(layer, "impl", None)
+        if native_impl is None:
+            raise RuntimeError(f"Native attention implementation is not bound for diffusion layer {layer.layer_name!r}")
+        if not layer.attn_backend.forward_includes_kv_cache_update:
+            native_impl.do_kv_cache_update(
+                layer,
+                paged_kv_context.key_write,
+                paged_kv_context.value_write,
+                kv_cache,
+                paged_kv_context.slot_mapping,
+            )
+
+        def run_native_attention(
+            query: torch.Tensor,
+            key: torch.Tensor,
+            value: torch.Tensor,
+            native_metadata,
+        ) -> torch.Tensor:
+            output = torch.empty(
+                (query.shape[0], layer.num_heads, layer.head_size_v),
+                dtype=query.dtype,
+                device=query.device,
+            )
+            return native_impl.forward(
+                layer,
+                query,
+                key,
+                value,
+                kv_cache,
+                native_metadata,
+                output,
+            )
+
+        if paged_kv_context.piecewise_plan is not None:
+            output = run_paged_piecewise_plan(
+                paged_kv_context.query,
+                paged_kv_context.key_write,
+                paged_kv_context.value_write,
+                paged_kv_context.piecewise_plan,
+                paged_kv_context.piecewise_native_metadata,
+                run_native_attention,
+            )
+        else:
+            output = run_native_attention(
+                paged_kv_context.query,
+                paged_kv_context.key_write,
+                paged_kv_context.value_write,
+                paged_kv_context.native_metadata,
+            )
+        return paged_kv_context.restore_output(output)
+
     def forward_cuda(
         self,
         query: torch.Tensor,

vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py

@@ -15,6 +15,8 @@ Per segment:
 
 from __future__ import annotations
 
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Literal, NamedTuple
 
 import torch
@@ -123,3 +125,109 @@ def piecewise_attn(
     if len(outputs) == 1:
         return outputs[0]
     return torch.cat(outputs, dim=1)
+
+
+@dataclass(frozen=True, slots=True)
+class PagedPiecewiseSegment:
+    """One aligned segment across rows in a packed paged batch."""
+
+    row_segments: tuple[Segment, ...]
+    query_indices: torch.Tensor
+
+
+@dataclass(frozen=True, slots=True)
+class PagedPiecewisePlan:
+    """Piecewise segment batches for one packed paged-attention batch."""
+
+    spans: tuple[tuple[tuple[int, int], ...], ...]
+    segments: tuple[PagedPiecewiseSegment, ...]
+    num_query_tokens: int
+
+
+PagedPiecewiseRunner = Callable[
+    [torch.Tensor, torch.Tensor, torch.Tensor, object],
+    torch.Tensor,
+]
+
+
+def build_paged_piecewise_plan(
+    full_attn_spans: Sequence[Sequence[tuple[int, int]]],
+    query_offsets: Sequence[int],
+    query_lens: Sequence[int],
+    seq_lens: Sequence[int],
+    *,
+    device: torch.device | str | None = None,
+) -> PagedPiecewisePlan:
+    """Build packed indices for corresponding piecewise segments in each row."""
+
+    row_count = len(full_attn_spans)
+    if row_count == 0 or not (row_count == len(query_offsets) == len(query_lens) == len(seq_lens)):
+        raise ValueError("Paged piecewise inputs must have one entry per row")
+
+    spans = tuple(tuple(tuple(span) for span in row_spans) for row_spans in full_attn_spans)
+
+    packed_offsets = [0]
+    for row_spans, query_offset, query_len, seq_len in zip(spans, query_offsets, query_lens, seq_lens, strict=True):
+        previous_end = 0
+        for start, end in row_spans:
+            if start < previous_end or not 0 <= start < end <= seq_len:
+                raise ValueError("Paged piecewise spans must be sorted, non-overlapping, and within the sequence")
+            previous_end = end
+        if query_offset < 0 or query_len <= 0 or query_offset + query_len > seq_len:
+            raise ValueError("Paged piecewise query range must be within the sequence")
+        packed_offsets.append(packed_offsets[-1] + query_len)
+
+    segments_by_row = tuple(
+        tuple(build_segments(row_spans, query_offset, query_len))
+        for row_spans, query_offset, query_len in zip(spans, query_offsets, query_lens, strict=True)
+    )
+    if len({len(row_segments) for row_segments in segments_by_row}) != 1:
+        raise ValueError("Paged piecewise rows must produce aligned segments")
+
+    packed_segments = []
+    for row_segments in zip(*segments_by_row, strict=True):
+        if len({segment.mode for segment in row_segments}) != 1:
+            raise ValueError("Paged piecewise rows must use the same attention mode per segment")
+        query_indices = []
+        for row_index, (segment, query_offset) in enumerate(zip(row_segments, query_offsets, strict=True)):
+            local_start = segment.q_start - query_offset
+            local_end = segment.q_end - query_offset
+            query_indices.extend(range(packed_offsets[row_index] + local_start, packed_offsets[row_index] + local_end))
+        packed_segments.append(
+            PagedPiecewiseSegment(
+                row_segments=tuple(row_segments),
+                query_indices=torch.tensor(query_indices, dtype=torch.long, device=device),
+            )
+        )
+
+    return PagedPiecewisePlan(spans, tuple(packed_segments), packed_offsets[-1])
+
+
+def run_paged_piecewise_plan(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    plan: PagedPiecewisePlan,
+    segment_metadata: Sequence[object],
+    segment_runner: PagedPiecewiseRunner,
+) -> torch.Tensor:
+    """Run and scatter piecewise attention over packed Q/K/V write tensors."""
+
+    if query.shape[0] != plan.num_query_tokens:
+        raise ValueError(f"Paged piecewise plan has {plan.num_query_tokens} query tokens, got {query.shape[0]}")
+
+    output = None
+    for segment, metadata in zip(plan.segments, segment_metadata, strict=True):
+        indices = segment.query_indices
+        segment_output = segment_runner(
+            query.index_select(0, indices),
+            key.index_select(0, indices),
+            value.index_select(0, indices),
+            metadata,
+        )
+        if output is None:
+            output = segment_output.new_empty((query.shape[0], *segment_output.shape[1:]))
+        output.index_copy_(0, indices, segment_output)
+
+    assert output is not None
+    return output

vllm_omni/diffusion/attention/layer.py

@@ -24,7 +24,11 @@ from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
 from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
 from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
-from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+from vllm_omni.diffusion.forward_context import (
+    get_forward_context,
+    get_ulysses_mode,
+    is_forward_context_available,
+)
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
@@ -73,6 +77,11 @@ class Attention(nn.Module):
         self.role = role
         self.role_category = role_category
         self.qkv_layout = qkv_layout
+        # ``prefix`` is also the stable layer identity used by vLLM's native
+        # KV-cache metadata.  Keep it on the Omni layer so the active paged
+        # adapter can dispatch the already-resharded Q/K/V to the matching
+        # native cache tensor without replacing this Omni execution path.
+        self.prefix = prefix
         if paged_kv_cache_role == "":
             raise ValueError("paged_kv_cache_role must be non-empty when provided")
         self.paged_kv_cache_role = paged_kv_cache_role
@@ -311,19 +320,48 @@ class Attention(nn.Module):
     ) -> torch.Tensor:
         # Get the appropriate parallel strategy based on SP active state
         strategy = self._get_active_parallel_strategy()
+        paged_adapter = self._active_paged_kv_adapter()
+        use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None
+        if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False):
+            raise NotImplementedError(
+                f"Diffusion paged KV requires an Omni backend with paged support; "
+                f"selected {self.attn_backend.get_name()}"
+            )
+        if use_paged_attention and strategy is not self._no_parallel_strategy:
+            strategy_name = strategy.name
+            if self.use_ring or strategy_name == "ring":
+                raise NotImplementedError(
+                    "paged Scheduler KV is not supported with Ring attention; use strict Ulysses or no SP"
+                )
+            if strategy_name == "allgather_kv":
+                raise NotImplementedError("paged Scheduler KV is not supported with AllGather-KV sequence parallelism")
+            if strategy_name == "ulysses" and get_ulysses_mode(default="strict") != "strict":
+                raise NotImplementedError("paged Scheduler KV currently supports only strict Ulysses")
 
         # 1. Prepare inputs (Communication / Resharding)
         # For Ulysses: AllToAll Q/K/V; Slicing joint_q/k/v
         # For Ring: Concat joint_q
         query, key, value, attn_metadata, ctx = strategy.pre_attention(query, key, value, attn_metadata)
 
-        attn_metadata = self._with_kv_cache_dtype(attn_metadata)
-
-        # 2. Kernel Execution (Computation)
-        if self.use_ring and strategy is not self._no_parallel_strategy:
-            out = self._run_ring_attention(query, key, value, attn_metadata)
+        # 2. Kernel execution stays inside the selected Omni backend.  The
+        # Worker adapter only prepares the native page-table context after
+        # SP has produced rank-local Q/K/V tensors.
+        if use_paged_attention:
+            assert paged_adapter is not None
+            paged_kv_context = paged_adapter.prepare_layer_context(
+                self.prefix,
+                query,
+                key,
+                value,
+                omni_attn_metadata=attn_metadata,
+            )
+            out = self.attention.forward_paged(paged_kv_context)
         else:
-            out = self._run_local_attention(query, key, value, attn_metadata)
+            attn_metadata = self._with_kv_cache_dtype(attn_metadata)
+            if self.use_ring and strategy is not self._no_parallel_strategy:
+                out = self._run_ring_attention(query, key, value, attn_metadata)
+            else:
+                out = self._run_local_attention(query, key, value, attn_metadata)
 
         # 3. Post-processing (Reverse Communication)
         # For Ulysses: AllToAll Output, and AllGather Joint Output
@@ -331,6 +369,19 @@ class Attention(nn.Module):
 
         return out
 
+    @staticmethod
+    def _active_paged_kv_adapter():
+        """Return the opaque Worker adapter installed for this forward."""
+
+        if not is_forward_context_available():
+            return None
+        return getattr(get_forward_context(), "paged_kv_adapter", None)
+
+    def is_paged_kv_active(self) -> bool:
+        """Return whether this layer will use Scheduler-managed paged KV."""
+
+        return self.paged_kv_cache_role is not None and self._active_paged_kv_adapter() is not None
+
     def _run_local_attention(self, query, key, value, attn_metadata):
         self._assert_piecewise_compatible(attn_metadata)
 

vllm_omni/diffusion/data.py

@@ -714,6 +714,9 @@ class OmniDiffusionConfig:
 
     # Local Diffusion KV ownership and cache-layout mode.
     diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
+    # Maximum number of native BlockTable rows one public request can own
+    # (sequences plus independent contexts). The model adapter defines it.
+    diffusion_kv_max_rows_per_request: int | None = None
 
     # Optional override for the diffusion model runner class (import path).
     # Precedence in the worker: this override > the runner declared by the
@@ -990,6 +993,15 @@ class OmniDiffusionConfig:
         if not isinstance(self.diffusion_compile_dynamic, bool):
             raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}")
         self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
+        if self.diffusion_kv_max_rows_per_request is not None and (
+            type(self.diffusion_kv_max_rows_per_request) is not int or self.diffusion_kv_max_rows_per_request <= 0
+        ):
+            raise ValueError("diffusion_kv_max_rows_per_request must be a positive integer when set")
+        if (
+            self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+            and self.diffusion_kv_max_rows_per_request is None
+        ):
+            raise ValueError("paged_scheduler requires diffusion_kv_max_rows_per_request to be set")
         if self.kv_cache_memory_bytes is not None and self.kv_cache_memory_bytes < 0:
             raise ValueError("kv_cache_memory_bytes must be non-negative")
         if not 0.0 < self.gpu_memory_utilization <= 1.0:

vllm_omni/diffusion/diffusion_engine.py

@@ -29,7 +29,7 @@ from vllm_omni.diffusion.data import (
     DiffusionRequestAbortedError,
     OmniDiffusionConfig,
 )
-from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode
 from vllm_omni.diffusion.diffusion_kv.initialization import initialize_diffusion_kv_control_plane
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.io_support import (
@@ -634,12 +634,25 @@ class DiffusionEngine:
             if not task.future.done():
                 task.future.set_exception(exc)
 
+    def _remove_diffusion_kv_requests(self, request_ids: Iterable[str]) -> None:
+        """Clear terminal Worker rows while Scheduler owns the allocations."""
+
+        od_config = getattr(self, "od_config", None)
+        if od_config is None or not is_scheduler_paged_kv_mode(
+            getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+        ):
+            return
+        unique_request_ids = list(dict.fromkeys(request_ids))
+        if unique_request_ids:
+            self.executor.remove_diffusion_kv_requests(unique_request_ids)
+
     def _emit_finished_outputs(
         self,
         finished_ids: set[str],
         runner_output: BaseRunnerOutput | None = None,
         missing_result_error: str = "Diffusion execution finished without a final output",
     ) -> None:
+        self._remove_diffusion_kv_requests(finished_ids)
         for rid in finished_ids:
             if runner_output is not None:
                 _output = runner_output.get_request_output(rid)
@@ -659,6 +672,8 @@ class DiffusionEngine:
             self._emit_finished_outputs(finished_ids, runner_output)
             return
 
+        self._remove_diffusion_kv_requests(finished_ids)
+
         delivered_finished_req_ids: set[str] = set()
 
         # finished_ids may have some requests that are not scheduler in this round.
@@ -852,6 +867,7 @@ class DiffusionEngine:
                 sched_output = self.scheduler.schedule()
                 if sched_output.is_empty:
                     if target_request_id in sched_output.finished_req_ids:
+                        self._remove_diffusion_kv_requests([target_request_id])
                         return self._finalize_finished_request(target_request_id)
                     if not self.scheduler.has_requests():
                         raise RuntimeError("Diffusion scheduler has no runnable requests.")
@@ -881,6 +897,7 @@ class DiffusionEngine:
                 if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1:
                     raise ValueError("Sync func should receive one result at one time")
                 if target_request_id in finished_req_ids:
+                    self._remove_diffusion_kv_requests([target_request_id])
                     req_output = runner_output.get_request_output(target_request_id)
                     output = self._finalize_finished_request(
                         target_request_id,
@@ -1232,8 +1249,11 @@ class DiffusionEngine:
 
     def _abort_requests(self, request_ids: str | Iterable[str]) -> None:
         request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids)
+        request_ids = list(dict.fromkeys(request_ids))
 
-        for request_id in dict.fromkeys(request_ids):
+        self._remove_diffusion_kv_requests(request_ids)
+
+        for request_id in request_ids:
             if self.scheduler.get_request_state(request_id) is not None:
                 self.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
 

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

vllm_omni/diffusion/executor/abstract.py

@@ -142,13 +142,24 @@ class DiffusionExecutor(ABC):
             raise TypeError("determine_available_kv_memory must return list[int]")
         return result
 
-    def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig]) -> None:
-        """Send each native rank-local config back to all Workers."""
+    def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig], resolved_max_model_len: int) -> None:
+        """Send rank-local configs and the resolved model length to all Workers."""
 
         # The default control-plane RPC mode executes on every rank and has
         # rank 0 return the gathered rank statuses, so failures on nonzero
         # ranks are not silently dropped.
-        self.collective_rpc("set_kv_cache_configs", args=(kv_cache_configs,))
+        self.collective_rpc("set_kv_cache_configs", args=(kv_cache_configs, resolved_max_model_len))
+
+    def remove_diffusion_kv_requests(self, request_ids: list[str]) -> None:
+        """Clear request rows on every Worker after Scheduler retirement."""
+
+        unique_request_ids = list(dict.fromkeys(request_ids))
+        if not unique_request_ids:
+            return
+        self.collective_rpc(
+            "remove_diffusion_kv_requests",
+            args=(unique_request_ids,),
+        )
 
     @abstractmethod
     def shutdown(self) -> None:

vllm_omni/diffusion/forward_context.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import torch
 import vllm.ir
@@ -26,6 +26,11 @@ class ForwardContext:
     vllm_config: VllmConfig | None = None
     omni_diffusion_config: OmniDiffusionConfig | None = None
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None
+    # Active Worker-side paged KV adapter.  The adapter is installed only for
+    # the duration of a paged forward; dense forwards leave this as ``None``.
+    # Keep the field opaque here to avoid coupling the common context module to
+    # the diffusion_kv implementation.
+    paged_kv_adapter: Any | None = None
     split_text_embed_in_sp: bool = False
     denoise_step_idx: int | None = None
     denoise_timestep: float | None = None
@@ -209,6 +214,30 @@ def set_forward_context(
                 yield
 
 
+@contextmanager
+def override_paged_kv_adapter(adapter: Any | None):
+    """Temporarily expose a Worker paged-KV adapter to Omni Attention.
+
+    This is deliberately a small context override instead of a second global
+    forward context.  The model runner owns the outer context and the adapter
+    only replaces one opaque field while its prepared native metadata is live.
+    """
+
+    if _forward_context is None:
+        # Unit-level adapter users can prepare/activate metadata without an
+        # Omni model forward context.  In that case there is nothing to
+        # override and the adapter's explicit ``forward`` API remains usable.
+        yield
+        return
+
+    previous = _forward_context.paged_kv_adapter
+    _forward_context.paged_kv_adapter = adapter
+    try:
+        yield
+    finally:
+        _forward_context.paged_kv_adapter = previous
+
+
 def set_forward_context_denoise_step_idx(step_idx: int | None) -> None:
     """Set the current diffusion denoise step on the active ForwardContext."""
     if _forward_context is not None:

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -1134,7 +1134,7 @@ class ImageKVCacheManager(nn.Module):
                 joint_text_query = query[:, :0, :, :]
                 joint_text_key, joint_text_value = self._reuse_prompt_kv(key, value, seq_len, bs, shard_image_size)
 
-        if not keep_kv_compressed:
+        if not keep_kv_compressed and not self.attn.is_paged_kv_active():
             key = repeat_kv(key, repeat_num)
             value = repeat_kv(value, repeat_num)
             if self.sp_size > 1:

vllm_omni/diffusion/vllm_config.py

@@ -5,7 +5,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 import torch
@@ -80,6 +80,15 @@ class _DiffusionVllmModelConfig:
     logits_processors: tuple[str, ...] = ()
     use_mla: bool = False
     is_moe: bool = False
+    # Native vLLM attention layers/builders read these fields directly.  The
+    # diffusion model does not have a full vLLM ModelConfig, so keep the
+    # narrow compatibility surface here instead of manufacturing a second
+    # model configuration object in the Worker backend.
+    is_mm_prefix_lm: bool = False
+    rswa_window: int | None = None
+    _attention_num_heads: int | None = field(default=None, init=False, repr=False)
+    _attention_num_kv_heads: int | None = field(default=None, init=False, repr=False)
+    _attention_head_size: int | None = field(default=None, init=False, repr=False)
 
     # Needed for models that bundle things like LogitsProcessors, e.g., SenseNova
     head_dtype: torch.dtype | None = None
@@ -91,6 +100,56 @@ class _DiffusionVllmModelConfig:
     def is_nvfp4_quantized(self) -> bool:
         return self.quantization == "modelopt_fp4"
 
+    def set_attention_geometry(
+        self,
+        *,
+        num_heads: int,
+        num_kv_heads: int,
+        head_size: int,
+    ) -> None:
+        """Record rank-local geometry needed by native metadata builders.
+
+        Diffusion attention modules are loaded outside vLLM's model loader, so
+        their dimensions are not available when this lightweight config is
+        created.  The paged Worker backend fills them in after model loading
+        and before calling ``init_attn_backend``.
+        """
+
+        values = {
+            "num_heads": num_heads,
+            "num_kv_heads": num_kv_heads,
+            "head_size": head_size,
+        }
+        invalid = {name: value for name, value in values.items() if type(value) is not int or value <= 0}
+        if invalid:
+            raise ValueError(f"Diffusion attention geometry must contain positive integers: {invalid!r}")
+        self._attention_num_heads = num_heads
+        self._attention_num_kv_heads = num_kv_heads
+        self._attention_head_size = head_size
+
+    def _require_attention_geometry(self) -> tuple[int, int, int]:
+        geometry = (
+            self._attention_num_heads,
+            self._attention_num_kv_heads,
+            self._attention_head_size,
+        )
+        if any(value is None for value in geometry):
+            raise RuntimeError("Diffusion attention geometry is unavailable before paged cache layer registration")
+        num_heads, num_kv_heads, head_size = geometry
+        assert num_heads is not None and num_kv_heads is not None and head_size is not None
+        return num_heads, num_kv_heads, head_size
+
+    def get_num_attention_heads(self, parallel_config: Any) -> int:
+        del parallel_config
+        return self._require_attention_geometry()[0]
+
+    def get_num_kv_heads(self, parallel_config: Any) -> int:
+        del parallel_config
+        return self._require_attention_geometry()[1]
+
+    def get_head_size(self) -> int:
+        return self._require_attention_geometry()[2]
+
     @property
     def is_diffusion(self) -> bool:
         # This config is also installed as the current VllmConfig for the
@@ -164,6 +223,7 @@ def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiff
         vllm_config.cache_config.kv_cache_memory_bytes = getattr(od_config, "kv_cache_memory_bytes", None)
         # Prefix hashes/publication are intentionally deferred to the next PR.
         vllm_config.cache_config.enable_prefix_caching = False
+        vllm_config.scheduler_config.max_num_seqs = int(od_config.max_num_seqs)
         max_num_batched_tokens = getattr(od_config, "max_num_batched_tokens", None)
         if max_num_batched_tokens is not None:
             vllm_config.scheduler_config.max_num_batched_tokens = int(max_num_batched_tokens)

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -24,6 +24,7 @@ from vllm.logger import init_logger
 from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
 
+from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.cache.cachedit import cache_summary
 from vllm_omni.diffusion.cache.prompt_embed_cache import (
     install_prompt_embed_cache,
@@ -34,6 +35,11 @@ from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
+from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
+from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
+    DiffusionPagedAttentionRow,
+    PreparedDiffusionPagedAttentionBatch,
+)
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import (
@@ -144,9 +150,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         self.prompt_embed_cache: Any | None = None
         self.input_batch: InputBatch | None = None
         self.model_memory_usage = 0
-        # Rank-local native config generated by the Engine. Physical tensors
-        # and BlockTable binding remain the responsibility of the paged
-        # data-plane initialization path.
+        self.diffusion_kv_backend = DiffusionKVModelRunnerBackend(
+            vllm_config=vllm_config,
+            od_config=od_config,
+            device=device,
+        )
+        # Compatibility view for callers that inspect the installed config;
+        # physical ownership remains in ``diffusion_kv_backend``.
         self.kv_cache_config: KVCacheConfig | None = None
 
         # Cache for per-request stepwise state.
@@ -381,41 +391,60 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
     def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
         """Collect native specs from cache-enabled loaded attention modules."""
 
-        from vllm_omni.diffusion.attention.layer import Attention
-        from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
-
         if self.od_config.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
             return {}
         if self.pipeline is None:
             raise RuntimeError("Model must be loaded before collecting Diffusion KV cache specs")
 
-        specs: dict[str, KVCacheSpec] = {}
+        cache_layers: dict[str, tuple[Attention, KVCacheSpec]] = {}
         for layer_name, module in self.pipeline.named_modules():
             if not isinstance(module, Attention):
                 continue
             spec = module.get_kv_cache_spec(self.vllm_config)
             if spec is not None:
-                specs[layer_name] = spec
-        if not specs:
+                cache_layers[layer_name] = (module, spec)
+        if not cache_layers:
             raise RuntimeError(
                 "paged_scheduler Diffusion KV found no cache-enabled Attention modules "
                 f"in {type(self.pipeline).__name__}"
             )
-        return specs
+        return self.diffusion_kv_backend.register_kv_cache_layers(cache_layers)
 
     def set_kv_cache_config(self, kv_cache_config: KVCacheConfig) -> None:
-        """Retain the Engine-generated rank-local config for the data plane."""
+        """Physically initialize the Engine-generated rank-local config."""
 
-        configured_layers = {
-            layer_name for group in kv_cache_config.kv_cache_groups for layer_name in group.layer_names
-        }
-        expected_layers = set(self.get_kv_cache_spec())
-        if configured_layers != expected_layers:
-            raise ValueError(
-                "Rank-local Diffusion KVCacheConfig layer mismatch: "
-                f"expected={sorted(expected_layers)}, configured={sorted(configured_layers)}"
-            )
-        self.kv_cache_config = kv_cache_config
+        self.diffusion_kv_backend.initialize_kv_cache(kv_cache_config)
+        self.kv_cache_config = self.diffusion_kv_backend.kv_cache_config
+
+    def install_diffusion_kv_metadata(self, metadata: DiffusionKVMetadata) -> bool:
+        return self.diffusion_kv_backend.install_diffusion_kv_metadata(metadata)
+
+    def get_diffusion_kv_row(
+        self,
+        request_id: str,
+        sequence_id: int | None,
+        context_id: str | None = None,
+    ) -> int:
+        return self.diffusion_kv_backend.get_diffusion_kv_row(request_id, sequence_id, context_id)
+
+    def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int:
+        return self.diffusion_kv_backend.remove_diffusion_kv_requests(request_ids)
+
+    def refresh_diffusion_kv_block_table_layout(self) -> None:
+        self.diffusion_kv_backend.refresh_block_table_layout()
+
+    def prepare_paged_attention_batch(
+        self,
+        rows: list[DiffusionPagedAttentionRow] | tuple[DiffusionPagedAttentionRow, ...],
+    ) -> PreparedDiffusionPagedAttentionBatch:
+        """Prepare model-specific query/write spans against installed Worker rows."""
+
+        return self.diffusion_kv_backend.prepare_paged_attention_batch(rows)
+
+    def activate_paged_attention(self, batch: PreparedDiffusionPagedAttentionBatch):
+        """Expose a prepared paged batch to Omni Attention for one forward."""
+
+        return self.diffusion_kv_backend.activate_paged_attention(batch)
 
     def clear_prompt_embed_cache(self) -> None:
         """Evict all cached text-encoder outputs (e.g. between training epochs).
@@ -664,17 +693,25 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             request_id=req.request_id,
             metadata=diffusion_kv_metadata,
         )
-        runner_output = self._execute_request_list(
-            [req],
-            od_config=self.od_config,
-            allow_single_output=True,
-            require_request_batch_support=False,
-            kv_prefetch_job=kv_prefetch_job,
-            record_name="pipeline_forward",
-        )
-        output = runner_output.runner_outputs[0].result
-        assert output is not None
-        return output
+        installed_request = False
+        if diffusion_kv_metadata is not None:
+            self.install_diffusion_kv_metadata(diffusion_kv_metadata)
+            installed_request = True
+        try:
+            runner_output = self._execute_request_list(
+                [req],
+                od_config=self.od_config,
+                allow_single_output=True,
+                require_request_batch_support=False,
+                kv_prefetch_job=kv_prefetch_job,
+                record_name="pipeline_forward",
+            )
+            output = runner_output.runner_outputs[0].result
+            assert output is not None
+            return output
+        finally:
+            if installed_request:
+                self.remove_diffusion_kv_requests([req.request_id])
 
     def profile_run(self, requests: list[OmniDiffusionRequest]) -> None:
         """Run the maximum per-rank request batch for memory profiling.
@@ -746,14 +783,23 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                 request_id=new_req.req.request_id,
                 metadata=new_req.diffusion_kv_metadata,
             )
-        reqs = [nr.req for nr in scheduler_output.scheduled_new_reqs]
-        return self._execute_request_list(
-            reqs,
-            od_config=od_config,
-            allow_single_output=False,
-            require_request_batch_support=True,
-            record_name="pipeline_forward_batch",
-        )
+        installed_request_ids: list[str] = []
+        try:
+            for new_req in scheduler_output.scheduled_new_reqs:
+                if new_req.diffusion_kv_metadata is not None:
+                    self.install_diffusion_kv_metadata(new_req.diffusion_kv_metadata)
+                    installed_request_ids.append(new_req.request_id)
+            reqs = [nr.req for nr in scheduler_output.scheduled_new_reqs]
+            return self._execute_request_list(
+                reqs,
+                od_config=od_config,
+                allow_single_output=False,
+                require_request_batch_support=True,
+                record_name="pipeline_forward_batch",
+            )
+        finally:
+            if installed_request_ids:
+                self.remove_diffusion_kv_requests(installed_request_ids)
 
     # ------------------------------------------------------------------
     # Step-wise execution
@@ -884,6 +930,40 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         if self.od_config.cache_backend not in (None, "none"):
             raise ValueError("Step mode does not support cache_backend yet.")
 
+        # Scheduler metadata is installed only for newly admitted requests;
+        # cached requests continue to use the row installed on their first
+        # step. Paged rows are retired here for scheduler-finished requests so
+        # a subsequent wave can reuse the native table slots. Dense execution
+        # must not issue Worker KV cleanup side effects.
+        if (
+            getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            is DiffusionKVCacheMode.PAGED_SCHEDULER
+            and scheduler_output.finished_req_ids
+        ):
+            self.remove_diffusion_kv_requests(list(scheduler_output.finished_req_ids))
+        installed_request_ids: list[str] = []
+        try:
+            for new_req in scheduler_output.scheduled_new_reqs:
+                if new_req.diffusion_kv_metadata is not None:
+                    self.install_diffusion_kv_metadata(new_req.diffusion_kv_metadata)
+                    installed_request_ids.append(new_req.request_id)
+            return self._execute_stepwise_core(
+                scheduler_output,
+                record_output_peak_memory=record_output_peak_memory,
+            )
+        except Exception:
+            if installed_request_ids:
+                self.remove_diffusion_kv_requests(installed_request_ids)
+            raise
+
+    def _execute_stepwise_core(
+        self,
+        scheduler_output: DiffusionSchedulerOutput,
+        *,
+        record_output_peak_memory: bool,
+    ) -> BatchRunnerOutput:
+        """Run the denoise step after metadata admission and row installation."""
+
         use_hsdp = self.od_config.parallel_config.use_hsdp
         grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
         with grad_context:
@@ -1005,7 +1085,16 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                             state.peak_memory_mb,
                         )
 
+                terminal_request_ids = [
+                    runner_output.request_id for runner_output in runner_output_list if runner_output.finished
+                ]
                 self._update_states_after(states, input_batch, pipeline_interrupted)
+                if (
+                    getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+                    is DiffusionKVCacheMode.PAGED_SCHEDULER
+                    and terminal_request_ids
+                ):
+                    self.remove_diffusion_kv_requests(terminal_request_ids)
 
                 return BatchRunnerOutput.from_list(runner_output_list)
 

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -188,6 +188,11 @@ class DiffusionWorker:
         # requests, which only carry their request_id in subsequent ticks.
         self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
         self.stage_id = getattr(od_config, "stage_id", 0)
+        if self.od_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER:
+            logger.warning_once(
+                "paged_scheduler initializes native paged KV storage, but no production diffusion model uses the "
+                "paged-attention adapter yet; model attention remains on the dense path."
+            )
         self.init_device()
         # Create model runner — one decision chain, in precedence order:
         #   1. explicit od_config.diffusion_model_runner_cls (user override),
@@ -431,16 +436,38 @@ class DiffusionWorker:
             )
         ]
 
-    def set_kv_cache_configs(self, kv_cache_configs: list[KVCacheConfig]) -> None:
-        """Install this rank's native config without allocating tensors yet."""
+    def set_kv_cache_configs(
+        self,
+        kv_cache_configs: list[KVCacheConfig],
+        resolved_max_model_len: int,
+    ) -> None:
+        """Select this rank's config and initialize its physical KV pages."""
 
         assert self.model_runner is not None
+        assert self.vllm_config is not None
         if len(kv_cache_configs) != self.od_config.num_gpus:
             raise ValueError(
                 "Diffusion KVCacheConfig rank count mismatch: "
                 f"expected={self.od_config.num_gpus}, got={len(kv_cache_configs)}"
             )
-        self.model_runner.set_kv_cache_config(kv_cache_configs[self.rank])
+        if not 0 <= self.rank < len(kv_cache_configs):
+            raise ValueError(f"Diffusion Worker rank {self.rank} has no rank-local KVCacheConfig")
+        if type(resolved_max_model_len) is not int or resolved_max_model_len <= 0:
+            raise ValueError("resolved Diffusion KV max_model_len must be a positive integer")
+
+        # Native cache sizing may resolve an explicit ``-1`` model-length
+        # sentinel to the capacity that actually fits the profiled pool.
+        self.vllm_config.model_config.max_model_len = resolved_max_model_len
+        kv_cache_config = kv_cache_configs[self.rank]
+        self.vllm_config.cache_config.num_gpu_blocks = kv_cache_config.num_blocks
+        with self._maybe_get_memory_pool_context("kv_cache"):
+            self.model_runner.set_kv_cache_config(kv_cache_config)
+
+    def remove_diffusion_kv_requests(self, request_ids: list[str]) -> int:
+        """Clear Worker-local rows without freeing Scheduler-owned blocks."""
+
+        assert self.model_runner is not None, "Model runner not initialized"
+        return self.model_runner.remove_diffusion_kv_requests(request_ids)
 
     def init_lora_manager(self) -> None:
         """Initialize the LoRA manager for this worker."""
@@ -730,6 +757,8 @@ class DiffusionWorker:
         allocator = CuMemAllocator.get_instance()
         allocator.wake_up(tags)
         current_omni_platform.synchronize()
+        if self.model_runner is not None and (tags is None or "kv_cache" in tags):
+            self.model_runner.refresh_diffusion_kv_block_table_layout()
         if len(self._sleep_saved_buffers) and self.model_runner is not None:
             model = self.model_runner.pipeline
             for name, buffer in model.named_buffers():

vllm_omni/platforms/interface.py

@@ -184,6 +184,28 @@ class OmniPlatform(Platform):
         """
         return "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner"
 
+    @classmethod
+    def get_diffusion_kv_block_tables_cls(cls) -> type:
+        """Return the platform's native paged-KV BlockTables implementation."""
+        from vllm.v1.worker.gpu.block_table import BlockTables
+
+        return BlockTables
+
+    @classmethod
+    def build_diffusion_kv_attn_metadata(cls, **kwargs: Any) -> dict[str, Any]:
+        """Build native attention metadata for the diffusion paged path.
+
+        The common CUDA/ROCm/XPU path uses vLLM's builder. Platforms with a
+        backend-specific metadata contract can override this hook without
+        making the diffusion adapter import their optional runtime package.
+        ``seq_lens_cpu`` is an adapter-only convenience value and is removed
+        before calling the upstream builder.
+        """
+        from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
+
+        kwargs.pop("seq_lens_cpu", None)
+        return build_attn_metadata(**kwargs)
+
     @classmethod
     def init_diffusion_worker_vllm_config(
         cls,

vllm_omni/platforms/npu/platform.py

@@ -80,6 +80,27 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
 
         init_ascend_config(vllm_config)
 
+    @classmethod
+    def get_diffusion_kv_block_tables_cls(cls) -> type:
+        from vllm_ascend.worker.v2.block_table import AscendBlockTables
+
+        return AscendBlockTables
+
+    @classmethod
+    def build_diffusion_kv_attn_metadata(cls, **kwargs: Any) -> dict[str, Any]:
+        """Build the Ascend metadata required by the native NPU backend."""
+        from vllm_ascend.attention.attention_v1 import AscendAttentionState
+        from vllm_ascend.worker.v2.attn_utils import build_attn_metadata
+
+        kwargs = dict(kwargs)
+        seq_lens_cpu = kwargs.pop("seq_lens_cpu")
+        kwargs["seq_lens_np"] = seq_lens_cpu.detach().cpu().numpy()
+        # The diffusion adapter always supplies a paged cache and the current
+        # K/V write span. ChunkedPrefill is Ascend's cache-backed FIA state for
+        # both multi-token updates and single-token updates in this path.
+        kwargs["attn_state"] = AscendAttentionState.ChunkedPrefill
+        return build_attn_metadata(**kwargs)
+
     @classmethod
     def init_diffusion_model_runner_runtime(cls, vllm_config: Any, od_config: Any, device: torch.device) -> None:
         from vllm_ascend.ascend_forward_context import set_mc2_mask, set_mc2_tokens_capacity