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