vllm-omni/tests/diffusion/diffusion_kv/test_block_tables.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni 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]

    # vLLM 0.29 dropped attn_groups and cache_dtype from init_kv_cache: the
    # layout now rides on the resolved CacheConfig.kv_cache_layout instead.
    def init_cache(
        runner_kv_caches,
        forward_context,
        kv_cache_config,
        device,
        kernel_block_sizes,
        vllm_config,
    ):
        events.append(("cache", kv_cache_config, 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, sequence_binding.block_ids) == (
        0,
        9,
        ((1, 2, 3), (4, 5)),
    )
    assert (context_binding.row_index, context_binding.max_seq_len, context_binding.block_ids) == (
        1,
        5,
        ((6, 7), (8,)),
    )


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())