[Refactor][CI] Reorganize test tree, and move slow e2e to weekly (#4951)

Signed-off-by: wangyu <410167048@qq.com>

tests/diffusion/ar_diffusion/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

tests/diffusion/ar_diffusion/test_capability_runner.py

@@ -0,0 +1,317 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU contracts for capability-driven AR-Diffusion sessions."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.experimental.ar_diffusion.capability import (
+    ARDiffusionCrossAttentionKVSpec,
+    ARDiffusionKVBranchSpec,
+    ARDiffusionKVCacheSpec,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+BLOCK = 16
+POS = "positive"
+NEG = "negative"
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def lingbot_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
+    """Single-kv_branch causal DMD: 3 latent frames/block + sink/window + text KV."""
+    return ARDiffusionKVCacheSpec(
+        num_layers=2,
+        num_kv_heads=4,  # already TP-local
+        head_size=64,
+        tokens_per_frame=BLOCK,
+        frames_per_block=3,
+        window_frames=5,
+        sink_frames=1,
+        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
+        session_capacity=capacity,
+        cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 8),),
+    )
+
+
+def dreamzero_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
+    return ARDiffusionKVCacheSpec(
+        num_layers=2,
+        num_kv_heads=4,
+        head_size=64,
+        tokens_per_frame=BLOCK,
+        frames_per_block=4,
+        window_frames=6,
+        kv_branches=(ARDiffusionKVBranchSpec(POS, 0), ARDiffusionKVBranchSpec(NEG, 1)),
+        session_capacity=capacity,
+        cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 8),),
+    )
+
+
+class CapablePipeline:
+    def __init__(self, spec: ARDiffusionKVCacheSpec) -> None:
+        self.spec = spec
+        self.bound_state = None
+        self.binds: list[str] = []
+        self.resets: list[str] = []
+        self.closes: list[str] = []
+
+    def ar_diffusion_kv_cache_spec(self) -> ARDiffusionKVCacheSpec:
+        return self.spec
+
+    @contextmanager
+    def bind_ar_diffusion_state(self, session_id: str, state):
+        assert self.bound_state is None
+        self.bound_state = state
+        self.binds.append(session_id)
+        try:
+            yield
+        finally:
+            self.bound_state = None
+
+    def reset_ar_diffusion_session(self, session_id: str) -> None:
+        self.resets.append(session_id)
+
+    def close_ar_diffusion_session(self, session_id: str) -> None:
+        self.closes.append(session_id)
+
+
+class WarmupPipeline(CapablePipeline):
+    def __init__(self, spec: ARDiffusionKVCacheSpec, requests: list[object]) -> None:
+        super().__init__(spec)
+        self.requests = requests
+
+    def ar_diffusion_warmup_requests(self, session_id: str):
+        assert session_id == ARDiffusionModelRunner._WARMUP_SID
+        return iter(self.requests)
+
+
+class BatchCapablePipeline(CapablePipeline):
+    supports_request_batch = True
+
+
+def make_runner(
+    pipeline: object,
+    *,
+    available_bytes: int = 1 << 28,
+    step_execution: bool = False,
+) -> ARDiffusionModelRunner:
+    runner = object.__new__(ARDiffusionModelRunner)
+    runner.od_config = SimpleNamespace(
+        max_num_seqs=1,
+        dtype=torch.float32,
+        enforce_eager=True,
+        step_execution=step_execution,
+    )
+    runner.device = torch.device("cpu")
+    runner.pipeline = pipeline
+    runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True)
+    runner.kv_cache = None
+    runner._ar_diffusion_capability = None
+    runner._ar_diffusion_kv_cache_spec = None
+    runner._sessions = OrderedDict()
+    runner._session_capacity = 0
+    runner._perf_e2e_times = []
+    runner._preallocate_kv_cache(available_bytes=available_bytes)
+    return runner
+
+
+def commit_one_frame(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str):
+    state = runner._get_or_create_session(session_id)
+    ctx = state.get_kv_caches(kv_branch, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    state.commit_paged_context(kv_branch)
+    return state
+
+
+def test_ar_runner_rejects_pipeline_without_capability():
+    runner = object.__new__(ARDiffusionModelRunner)
+    runner.od_config = SimpleNamespace(max_num_seqs=1)
+    runner.pipeline = object()
+    with pytest.raises(TypeError, match="SupportsARDiffusionPipeline"):
+        runner._preallocate_kv_cache(available_bytes=1 << 20)
+
+
+def test_lingbot_like_single_branch_session_reuse_reset_and_close():
+    pipeline = CapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline)
+    kv = runner.kv_cache
+    assert kv is not None
+    assert kv.num_local_kv_branches == 1
+    assert kv.frames_per_block == 3
+    assert kv.spec.window_chunks == 5
+    assert kv.spec.sink_chunks == 1
+    assert kv.cross_attention_lengths == {"text": 8}
+
+    first = commit_one_frame(runner, "s1", "main")
+    assert runner._get_or_create_session("s1") is first
+    k = torch.randn(1, 8, 4, 64)
+    v = torch.randn(1, 8, 4, 64)
+    first.populate_cross_attention("main", "text", [(k, v)] * first.num_layers)
+    assert first.get_cross_attention_kv("main", "text")[0]["k"].shape == k.shape
+
+    runner.reset_session("s1")
+    assert "s1" not in runner._sessions
+    assert "s1" not in kv._cross_sessions
+    assert pipeline.resets == ["s1"]
+    second = runner._get_or_create_session("s1")
+    assert second is not first
+
+    runner.close_session("s1")
+    assert "s1" not in runner._sessions
+    assert pipeline.closes == ["s1"]
+
+
+def test_lingbot_like_sink_survives_sliding_window_eviction():
+    runner = make_runner(CapablePipeline(lingbot_like_spec()))
+    kv = runner.kv_cache
+    assert kv is not None
+
+    for _ in range(8):
+        state = commit_one_frame(runner, "s1", "main")
+
+    table = kv.block_table(state.adapter("main"))
+    assert table[0] != kv.null_block_id
+    assert table[1] == kv.null_block_id
+
+    ctx = state.get_kv_caches("main", seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    visible, _ = ctx.video_block_table(torch.device("cpu"))
+    assert visible[0] == table[0]
+    assert len(visible) == 6  # sink + recent window, including current scratch
+
+
+def test_dreamzero_like_two_branches_are_independent():
+    runner = make_runner(CapablePipeline(dreamzero_like_spec()))
+    kv = runner.kv_cache
+    assert kv is not None and kv.num_local_kv_branches == 2
+    state = commit_one_frame(runner, "s1", POS)
+    assert len(kv.window_block_ids(state.adapter(POS))) == 1
+    assert kv.window_block_ids(state.adapter(NEG)) == []
+
+
+def test_lru_eviction_releases_blocks_and_notifies_pipeline():
+    pipeline = CapablePipeline(lingbot_like_spec(capacity=8))
+    runner = make_runner(pipeline)
+    kv = runner.kv_cache
+    assert kv is not None
+    assert runner._session_capacity == 1
+    assert kv.session_capacity == 1
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+    old = commit_one_frame(runner, "old", "main")
+    k = torch.randn(1, 8, 4, 64)
+    old.populate_cross_attention("main", "text", [(k, k)] * old.num_layers)
+    assert kv.manager.block_pool.get_num_free_blocks() < free_total
+
+    runner._get_or_create_session("new")
+
+    assert tuple(runner._sessions) == ("new",)
+    assert pipeline.closes == ["old"]
+    assert "old" not in kv._cross_sessions
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
+    pipeline = CapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline)
+    kv = runner.kv_cache
+    assert kv is not None
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    def boom(self, req, kv_prefetch_job=None):
+        state = pipeline.bound_state
+        ctx = state.get_kv_caches("main", seq_len=BLOCK, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(torch.device("cpu"))
+        raise RuntimeError("layer exploded")
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
+    request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
+
+    with pytest.raises(RuntimeError, match="layer exploded"):
+        runner.execute_model(request)
+
+    assert pipeline.bound_state is None
+    assert pipeline.closes == ["broken"]
+    assert not runner._sessions
+    assert not kv._adapters
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
+    pipeline = CapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline)
+    kv = runner.kv_cache
+    assert kv is not None
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    def return_after_allocation(self, req, kv_prefetch_job=None):
+        state = pipeline.bound_state
+        ctx = state.get_kv_caches("main", seq_len=BLOCK, commit_current=True)[0].forward_ctx
+        ctx.ensure_video_slots(torch.device("cpu"))
+        return object()
+
+    def synchronize_boom(device):
+        raise RuntimeError("asynchronous kernel failed")
+
+    monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
+    monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
+    runner.device = torch.device("cuda")
+    request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
+
+    with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
+        runner.execute_model(request)
+
+    assert pipeline.bound_state is None
+    assert pipeline.closes == ["broken"]
+    assert not runner._sessions
+    assert not kv._adapters
+    assert not runner._perf_e2e_times
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_ar_runner_rejects_step_and_request_batch_modes():
+    with pytest.raises(ValueError, match="step_execution=True"):
+        make_runner(CapablePipeline(lingbot_like_spec()), step_execution=True)
+    with pytest.raises(ValueError, match="request-batch execution"):
+        make_runner(BatchCapablePipeline(lingbot_like_spec()))
+
+
+def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints():
+    runner = object.__new__(ARDiffusionModelRunner)
+    with pytest.raises(RuntimeError, match="request-batch execution"):
+        runner.execute_model_batch(None, None)
+    with pytest.raises(RuntimeError, match="step execution"):
+        runner.execute_stepwise(None)
+
+
+def test_model_specific_warmup_provider_is_consumed(monkeypatch):
+    requests = [object(), object()]
+    pipeline = WarmupPipeline(lingbot_like_spec(), requests)
+    runner = make_runner(pipeline)
+    seen: list[object] = []
+    monkeypatch.setattr(runner, "execute_model", seen.append)
+
+    runner._warmup_ar_rollout()
+
+    assert seen == requests
+    assert pipeline.closes == [runner._WARMUP_SID]
+
+
+def test_pipeline_without_warmup_provider_is_safely_skipped(monkeypatch):
+    pipeline = CapablePipeline(lingbot_like_spec())
+    runner = make_runner(pipeline)
+    execute = SimpleNamespace(called=False)
+
+    def fail_if_called(request):
+        execute.called = True
+
+    monkeypatch.setattr(runner, "execute_model", fail_if_called)
+    runner._warmup_ar_rollout()
+    assert execute.called is False

tests/diffusion/ar_diffusion/test_engine_routing.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the generic engine-backend dispatcher + AR-Diffusion runner wiring.
+
+``DiffusionEngine.resolve_engine_class`` is a generic dispatcher (``"default"`` / a
+``DiffusionEngine`` subclass / an import-path string). DreamZero selects the
+AR-Diffusion engine via its deploy config's ``engine_backend`` qualname — no
+DreamZero/ar_diffusion knowledge lives in the public base, so the routing check
+here is simply that the ``ARDiffusionEngine`` qualname resolves correctly.
+"""
+
+from dataclasses import fields
+from types import SimpleNamespace
+
+import pytest
+from vllm.utils.import_utils import resolve_obj_by_qualname
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.experimental.ar_diffusion.engine import (
+    AR_DIFFUSION_MODEL_RUNNER_CLS,
+    ARDiffusionEngine,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _cfg(backend):
+    return SimpleNamespace(engine_backend=backend)
+
+
+def test_default_resolves_to_diffusion_engine():
+    assert DiffusionEngine.resolve_engine_class(_cfg("default")) is DiffusionEngine
+
+
+def test_subclass_type_is_returned():
+    assert DiffusionEngine.resolve_engine_class(_cfg(ARDiffusionEngine)) is ARDiffusionEngine
+
+
+def test_ar_diffusion_qualname_resolves():
+    # How DreamZero's deploy config selects the engine: a full import-path string.
+    cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"))
+    assert cls is ARDiffusionEngine
+
+
+def test_non_engine_type_raises():
+    with pytest.raises(TypeError):
+        DiffusionEngine.resolve_engine_class(_cfg(dict))
+
+
+def test_qualname_not_an_engine_raises():
+    with pytest.raises(TypeError):
+        DiffusionEngine.resolve_engine_class(_cfg("builtins.dict"))
+
+
+def test_bad_qualname_raises():
+    with pytest.raises(ValueError):
+        DiffusionEngine.resolve_engine_class(_cfg("not.a.real.module.NoSuchEngine"))
+
+
+def test_ar_diffusion_engine_is_diffusion_engine_subclass():
+    assert issubclass(ARDiffusionEngine, DiffusionEngine)
+
+
+def test_config_engine_backend_field_defaults_to_default():
+    field = {f.name: f for f in fields(OmniDiffusionConfig)}["engine_backend"]
+    assert field.default == "default"
+
+
+# --- AR-Diffusion worker/runner wiring -----------------------------------------------
+
+
+def test_ar_diffusion_model_runner_cls_resolves():
+    cls = resolve_obj_by_qualname(AR_DIFFUSION_MODEL_RUNNER_CLS)
+    assert cls.__name__ == "ARDiffusionModelRunner"
+
+
+def test_engine_class_declares_its_runner():
+    """Routing lives on the engine class (review: hsliuustc0106) — engines
+    never mutate od_config; the worker resolves the runner off the class."""
+    assert ARDiffusionEngine.default_diffusion_model_runner_cls == AR_DIFFUSION_MODEL_RUNNER_CLS
+    assert DiffusionEngine.default_diffusion_model_runner_cls is None
+
+
+def _select_runner(od_config, platform_default="PLATFORM"):
+    """Mirror of the worker's runner-selection chain (override > engine > platform)."""
+    runner_override = getattr(od_config, "diffusion_model_runner_cls", None)
+    if isinstance(runner_override, str) and runner_override:
+        return runner_override
+    engine_cls = DiffusionEngine.resolve_engine_class(od_config)
+    engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None)
+    if isinstance(engine_runner, str) and engine_runner:
+        return engine_runner
+    return platform_default
+
+
+def test_runner_selection_chain():
+    ar = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine"
+    # engine-declared runner wins over platform default
+    assert (
+        _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend=ar))
+        == AR_DIFFUSION_MODEL_RUNNER_CLS
+    )
+    # explicit override wins over the engine's declaration
+    assert (
+        _select_runner(SimpleNamespace(diffusion_model_runner_cls="my.custom.Runner", engine_backend=ar))
+        == "my.custom.Runner"
+    )
+    # default engine -> platform default
+    assert _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend="default")) == "PLATFORM"
+
+
+def test_dreamzero_pipeline_rejects_non_ar_diffusion_engine():
+    """Review (hsliuustc0106): a stale config with engine_backend='default'
+    must fail at init, not mid-forward on the first KV access."""
+    from types import SimpleNamespace
+
+    import pytest
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
+        DreamZeroPipeline(od_config=SimpleNamespace(engine_backend="default"))
+    with pytest.raises(ValueError, match="requires the AR-Diffusion engine"):
+        DreamZeroPipeline(od_config=SimpleNamespace())
+
+
+def test_runner_rejects_multi_sequence_config():
+    """Review (hsliuustc0106): the paged path writes batch index 0 only, so
+    max_num_seqs > 1 must be rejected at init."""
+    from types import SimpleNamespace
+
+    import pytest
+
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    fake = SimpleNamespace(od_config=SimpleNamespace(max_num_seqs=4))
+    with pytest.raises(ValueError, match="max_num_seqs=1"):
+        ARDiffusionModelRunner._preallocate_kv_cache(fake)
+
+
+def test_execute_model_accepts_kv_prefetch_job_kwarg():
+    """Regression: the worker passes ``kv_prefetch_job`` to every runner
+    (added with the KV prefetch path, #4448); the AR-Diffusion override must
+    accept and forward it or DreamZero crashes on the warm-up dummy run."""
+    import inspect
+
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    for cls in (DiffusionModelRunner, ARDiffusionModelRunner):
+        params = inspect.signature(cls.execute_model).parameters
+        assert "kv_prefetch_job" in params, cls

tests/diffusion/ar_diffusion/test_kv_cache.py

@@ -0,0 +1,305 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the AR-Diffusion KV cache helpers (Phase 1, PR-2).
+
+Covers the request adapter, the chunk-window spec/manager (registration + the
+eviction policy), and the pool builder — exercised against the installed vLLM
+KV stack on CPU (block bookkeeping only, no GPU tensors).
+"""
+
+import pytest
+import torch
+from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
+from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionRequestAdapter,
+    ChunkWindowManager,
+    ChunkWindowSpec,
+    allocate_kv_pool_with_views,
+    build_kv_manager,
+    compute_num_blocks,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+BLOCK = 16
+
+
+def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
+    return ChunkWindowSpec(
+        block_size=BLOCK,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * chunk_size,
+        chunk_size=chunk_size,
+        window_chunks=window_chunks,
+        sink_chunks=sink_chunks,
+        reset_at_boundary=reset_at_boundary,
+    )
+
+
+# --- ChunkWindowSpec registration -------------------------------------------
+
+
+def test_spec_registration_resolves_to_chunk_window_manager():
+    # Without explicit registration the MRO walk would fall back to the parent
+    # SlidingWindowManager; assert the subclass manager wins.
+    spec = make_spec()
+    assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager
+
+
+def test_spec_kind_is_sliding_window():
+    assert get_kv_cache_spec_kind(make_spec()) == KVCacheSpecKind.SLIDING_WINDOW
+
+
+def test_spec_rejects_inconsistent_window():
+    with pytest.raises(ValueError):
+        ChunkWindowSpec(
+            block_size=BLOCK,
+            num_kv_heads=4,
+            head_size=64,
+            dtype=torch.float16,
+            sliding_window=99,  # != window_chunks * chunk_size
+            chunk_size=BLOCK,
+            window_chunks=2,
+        )
+
+
+# --- eviction policy (pure) -------------------------------------------------
+
+
+def test_sliding_replace_keeps_window():
+    # window = 2 chunks * 16 = 32. Base sliding formula keeps `window` tokens;
+    # the snap is to chunk boundaries.
+    def skip(n):
+        return chunk_window_skipped_tokens(n, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
+
+    assert skip(32) == 0  # nothing past the window yet
+    assert skip(48) == 16  # one chunk fell out of the window
+    assert skip(64) == 32
+
+
+def test_sliding_replace_snaps_to_chunk_boundary():
+    # A non-chunk-aligned overflow must snap down so a chunk is never half-evicted.
+    skip = chunk_window_skipped_tokens(50, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
+    assert skip % 16 == 0 and skip == 16
+
+
+def test_sink_chunks_protected():
+    # sink = 1 chunk (16 tokens) is never skipped.
+    skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
+    # raw overflow = 80 - 32 - 16 = 32 -> snapped 32
+    assert skip == 32
+
+
+def test_reset_at_boundary_drops_completed_past_sink():
+    skip = chunk_window_skipped_tokens(48, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=True)
+    # completed = 48; sink = 16 -> drop 32
+    assert skip == 32
+
+
+# --- ARDiffusionKVConfig ------------------------------------------------------------
+
+
+def test_kv_config_sliding_window_property():
+    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=3).sliding_window == 48
+    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=None).sliding_window is None
+
+
+# --- ARDiffusionRequestAdapter ------------------------------------------------------
+
+
+def test_adapter_advances_per_chunk_not_per_step():
+    a = ARDiffusionRequestAdapter("r0", chunk_size=16)
+    assert a.num_computed_tokens == 0
+    assert a.num_tokens == 16  # in-flight chunk
+    a.on_chunk_committed()
+    assert a.num_computed_tokens == 16
+    assert a.num_tokens == 32
+
+
+def test_adapter_accounts_for_prefill_prefix():
+    a = ARDiffusionRequestAdapter("r0", chunk_size=16, prefill_prefix_tokens=4)
+    assert a.num_computed_tokens == 4
+    assert a.num_prompt_tokens == 4
+    assert a.num_tokens == 20
+
+
+def test_adapter_status_is_vllm_enum():
+    assert isinstance(ARDiffusionRequestAdapter("r0", chunk_size=16).status, RequestStatus)
+
+
+# --- pool / manager ---------------------------------------------------------
+
+
+def test_compute_num_blocks():
+    # 1 MiB budget at fraction 0.5 with 16 KiB pages -> 32 blocks.
+    assert compute_num_blocks(1 << 20, 0.5, 16 << 10) == 32
+
+
+def test_paged_pool_layout_exposes_flat_slot_views():
+    kv_pools, k_pools, v_pools = allocate_kv_pool_with_views(
+        num_blocks=4,
+        block_size=BLOCK,
+        num_layers=1,
+        num_kv_heads=4,
+        head_dim=64,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64)
+    assert k_pools[0].shape == (4 * BLOCK, 4, 64)
+    assert v_pools[0].shape == (4 * BLOCK, 4, 64)
+
+    k_pools[0][BLOCK + 3].fill_(7)
+    v_pools[0][2 * BLOCK + 5].fill_(11)
+    assert torch.equal(kv_pools[0][0, 1, 3], k_pools[0][BLOCK + 3])
+    assert torch.equal(kv_pools[0][1, 2, 5], v_pools[0][2 * BLOCK + 5])
+
+
+def test_build_manager_allocate_free_roundtrip():
+    """End-to-end: a ARDiffusionRequestAdapter drives a real KVCacheManager.
+
+    This is the adapter conformance check — if the adapter were missing an
+    attribute the manager reads, allocate_slots/free would raise here.
+    """
+    spec = make_spec()
+    mgr = build_kv_manager(spec, ["layer0"], num_blocks=16, max_model_len=1024)
+    free_before = mgr.block_pool.get_num_free_blocks()
+
+    adapter = ARDiffusionRequestAdapter("req-0", chunk_size=BLOCK)
+    blocks = mgr.allocate_slots(adapter, num_new_tokens=BLOCK, full_sequence_must_fit=True)
+    assert blocks is not None
+    assert mgr.block_pool.get_num_free_blocks() < free_before
+
+    mgr.free(adapter)
+    assert mgr.block_pool.get_num_free_blocks() == free_before
+
+
+def test_cross_attn_pool_deducted_from_self_attn_budget():
+    """The cross-attn pool is allocated directly; its bytes are subtracted from
+    the self-attn paged-pool budget so the two together stay within the free
+    memory budget (review: zwhzzz0821)."""
+    L = 512
+    avail = 1 << 30  # 1 GiB
+    kv = ARDiffusionKVCache(
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2, gpu_memory_fraction=0.5),
+        num_layers=2,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=avail,
+        kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)),
+        session_capacity=1,
+        cross_attention_lengths={"text": L},
+        device=torch.device("cpu"),
+    )
+    cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2  # K+V, pos+neg, layers
+    expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
+    assert kv.num_blocks == expected
+    assert expected > 1 * (2 + 1) + 2  # above the local-slot floor, so the cross deduction is what's tested
+
+
+def _make_kv(
+    *,
+    local_branches,
+    num_frame_per_block=2,
+    window_chunks=9,
+    max_scratch_tokens_per_branch=0,
+):
+    kv_branches = (
+        (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0))
+        if local_branches == 1
+        else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1))
+    )
+    return ARDiffusionKVCache(
+        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
+        num_layers=1,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 16,  # tiny -> the floor binds
+        device=torch.device("cpu"),
+        kv_branches=kv_branches,
+        session_capacity=1,
+        frames_per_block=num_frame_per_block,
+        max_scratch_tokens_per_branch=max_scratch_tokens_per_branch,
+    )
+
+
+def test_pool_floor_is_branch_aware():
+    """CFG-parallel rank (one local kv_branch) sizes for one window + in-flight chunk;
+    a single-process run (both kv_branches) sizes for two. Scratch scales the same way."""
+    one = _make_kv(local_branches=1)
+    two = _make_kv(local_branches=2)
+
+    assert one.managed_num_blocks == 1 * (9 + 2) + 2  # 13
+    assert two.managed_num_blocks == 2 * (9 + 2) + 2  # 24
+    assert one.scratch_num_blocks == one.scratch_blocks_per_kv_branch
+    assert two.scratch_num_blocks == 2 * two.scratch_blocks_per_kv_branch
+    assert one.scratch_blocks_per_kv_branch == 2
+    assert one.num_blocks_total == 13 + one.scratch_blocks_per_kv_branch
+
+
+def test_scratch_capacity_is_derived_from_declared_geometry(monkeypatch):
+    monkeypatch.delenv("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", raising=False)
+    kv = _make_kv(
+        local_branches=1,
+        num_frame_per_block=6,
+        max_scratch_tokens_per_branch=BLOCK + 1,
+    )
+    assert kv.scratch_blocks_per_kv_branch == 8  # six video blocks + two auxiliary blocks
+
+
+def test_scratch_env_override_cannot_reduce_declared_minimum(monkeypatch):
+    monkeypatch.setenv("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "1")
+    kv = _make_kv(
+        local_branches=1,
+        num_frame_per_block=6,
+        max_scratch_tokens_per_branch=BLOCK + 1,
+    )
+    assert kv.scratch_blocks_per_kv_branch == 8
+
+
+def test_scratch_maps_to_slot_zero_with_one_local_branch():
+    """A CFG-parallel rank runs exactly one kv_branch: whichever CFG side it is,
+    its scratch lands in the rank's single slot (no dead second slot)."""
+    one = _make_kv(local_branches=1)
+    assert one.scratch_block_ids("positive", 0, 2) == one.scratch_block_ids("negative", 0, 2)
+
+    two = _make_kv(local_branches=2)
+    assert two.scratch_block_ids("positive", 0, 2) != two.scratch_block_ids("negative", 0, 2)
+
+
+def test_scratch_exhaustion_still_raises():
+    one = _make_kv(local_branches=1)
+    cap = one.scratch_blocks_per_kv_branch
+    with pytest.raises(RuntimeError, match="scratch blocks exhausted"):
+        one.scratch_block_ids("positive", 0, cap + 1)
+
+
+def test_non_contiguous_branch_indices_rejected():
+    with pytest.raises(ValueError, match="contiguous"):
+        ARDiffusionKVCache(
+            ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2),
+            num_layers=1,
+            num_kv_heads=4,
+            head_size=64,
+            dtype=torch.float32,
+            block_size=BLOCK,
+            max_model_len=4096,
+            available_bytes=1 << 16,
+            kv_branches=(ARDiffusionKVBranchSpec("main", 1),),
+            session_capacity=1,
+        )

tests/diffusion/ar_diffusion/test_kv_orchestration.py

@@ -0,0 +1,108 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for ARDiffusionKVCache — the engine-level KV orchestration body (Phase 1)."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
+from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+DIMS = dict(
+    num_layers=2,
+    num_kv_heads=4,
+    head_size=64,
+    dtype=torch.float16,
+    block_size=16,
+    kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1)),
+    session_capacity=2,
+)
+
+
+def make_cache(*, chunk_size=16, window_chunks=2, available_bytes=1 << 24):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks)
+    return ARDiffusionKVCache(cfg, max_model_len=4096, available_bytes=available_bytes, **DIMS)
+
+
+def test_requires_enabled_config():
+    with pytest.raises(ValueError):
+        ARDiffusionKVCache(ARDiffusionKVConfig(enable=False), max_model_len=256, available_bytes=1 << 20, **DIMS)
+
+
+def test_requires_bounded_window():
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=16, window_chunks=None)
+    with pytest.raises(ValueError):
+        ARDiffusionKVCache(cfg, max_model_len=256, available_bytes=1 << 20, **DIMS)
+
+
+def test_full_request_lifecycle_and_eviction():
+    """begin -> per-chunk allocate/slots/commit over a long rollout -> free.
+
+    Exercises the orchestrator end-to-end and asserts the chunk window bounds
+    memory (pool plateaus) and frees cleanly.
+    """
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    adapter = kv.begin_request("req-0")
+    free_after = []
+    for k in range(10):
+        block_table = kv.allocate_chunk(adapter)
+        slots = kv.chunk_write_slots(adapter)
+        # Slots target real blocks in this chunk's table; length == chunk_size.
+        assert len(slots) == kv.spec.chunk_size
+        used = {int(s) // kv.block_size for s in slots}
+        assert kv.null_block_id not in used
+        assert used <= set(block_table)
+        # Resident window stays bounded by window + the in-flight chunk.
+        assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
+        free_after.append(kv.manager.block_pool.get_num_free_blocks())
+        kv.commit_chunk(adapter)
+
+    # Pool memory plateaus once the window is full (eviction recycles blocks).
+    assert free_after[-1] == free_after[kv.spec.window_chunks]
+
+    kv.end_request(adapter)
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total
+
+
+def test_num_computed_advances_per_chunk():
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    a = kv.begin_request("r")
+    assert a.num_computed_tokens == 0
+    kv.allocate_chunk(a)
+    kv.commit_chunk(a)
+    assert a.num_computed_tokens == 16
+    kv.end_request(a)
+
+
+def test_state_close_frees_both_branch_blocks():
+    """ARDiffusionKVState.close() returns both CFG kv_branches' pool blocks to the pool.
+
+    This is the primitive the runner's LRU eviction relies on: when a session is
+    evicted, close() must free the blocks both adapters hold, or session churn
+    leaks pool ownership (review P1).
+    """
+    from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+    kv = make_cache(chunk_size=16, window_chunks=2)
+    free_total = kv.manager.block_pool.get_num_free_blocks()
+
+    pos = kv.begin_request("bde__s")
+    neg = kv.begin_request("bde__s__neg")
+    state = ARDiffusionKVState(
+        kv,
+        "s",
+        {"positive": pos, "negative": neg},
+        num_layers=kv.num_layers,
+    )
+    for _ in range(3):
+        for adapter in (pos, neg):
+            kv.allocate_chunk(adapter)
+            kv.commit_chunk(adapter)
+    # Both kv_branches hold resident blocks now.
+    assert kv.manager.block_pool.get_num_free_blocks() < free_total
+
+    state.close()
+    # All blocks returned — no leak across an evicted session.
+    assert kv.manager.block_pool.get_num_free_blocks() == free_total

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -0,0 +1,409 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for AR-Diffusion paged self-attention contexts."""
+
+from __future__ import annotations
+
+import subprocess
+from importlib.util import find_spec
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+    ARDiffusionPagedLayerInputs,
+    ar_diffusion_paged_attention,
+    paged_write_attn,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+POS = "positive"
+NEG = "negative"
+
+
+def make_state(*, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu")):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=num_layers,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=dtype,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 26,
+        kv_branches=(ARDiffusionKVBranchSpec(POS, 0), ARDiffusionKVBranchSpec(NEG, 1)),
+        session_capacity=2,
+        frames_per_block=2,
+        max_scratch_tokens_per_branch=BLOCK,
+        device=device,
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, "s1", {POS: pos, NEG: neg}, num_layers=num_layers)
+
+
+def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * (HEAD_DIM**-0.5)
+    probs = torch.softmax(scores, dim=-1).to(value.dtype)
+    return torch.einsum("bhqk,bkhd->bqhd", probs, value)
+
+
+def _cuda_flash_attn_usable() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    try:
+        spec = find_spec("vllm.vllm_flash_attn")
+        if spec is None or spec.origin is None:
+            return True
+        fa2_so = Path(spec.origin).parent / "_vllm_fa2_C.abi3.so"
+        linked = subprocess.check_output(["ldd", str(fa2_so)], text=True, timeout=5)
+    except Exception:
+        return True
+    if "libcudart.so.13" not in linked:
+        return True
+    try:
+        out = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
+            text=True,
+            timeout=5,
+        )
+        driver_major = int(out.splitlines()[0].split(".")[0])
+    except Exception:
+        return True
+    return driver_major >= 580
+
+
+def _commit_video_span(
+    kv: ARDiffusionKVCache,
+    st: ARDiffusionKVState,
+    *,
+    kv_branch: str,
+    n_chunks: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    ctx = st.get_kv_caches(kv_branch, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = v[0]
+    st.commit_paged_context(kv_branch)
+    return k, v
+
+
+def test_paged_context_allocates_lazily_and_commits_after_forward():
+    _, st = make_state()
+
+    contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+    assert st.adapter(POS).completed_chunks == 0
+    assert ctx.current_video_slot_mapping is None
+
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.adapter(POS).completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 1
+
+    st.commit_paged_context(POS)
+    assert st.adapter(POS).completed_chunks == 1
+    assert st._committed[POS] == BLOCK
+
+
+def test_scratch_video_and_action_blocks_do_not_commit():
+    kv, st = make_state()
+
+    ctx = st.get_kv_caches(POS, seq_len=2 * BLOCK, commit_current=False)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    ctx.ensure_action_slots(3, torch.device("cpu"))
+
+    assert ctx.current_video_block_ids == kv.scratch_block_ids(POS, 0, 2)
+    assert ctx.action_scratch_block_ids == kv.scratch_block_ids(POS, 2, 1)
+    st.commit_paged_context(POS)
+    assert st.adapter(POS).completed_chunks == 0
+    assert st._committed[POS] == 0
+
+
+def test_pipeline_kv_get_paged_path_has_no_gather_backend():
+    kv, st = make_state()
+    assert not hasattr(kv, "gather_window_all_layers")
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    pipeline._ar_diffusion_kv_state = st
+    contexts = pipeline._kv_get(MagicMock(), False, seq_len=BLOCK, update_kv_cache=False)
+
+    assert len(contexts) == 1
+    assert isinstance(contexts[0], ARDiffusionPagedLayerContext)
+
+
+@pytest.mark.parametrize("history_chunks", [0, 1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, commit_current):
+    torch.manual_seed(0)
+    device = torch.device("cpu")
+    dtype = torch.float32
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+
+    history_k_parts: list[torch.Tensor] = []
+    history_v_parts: list[torch.Tensor] = []
+    if history_chunks:
+        k, v = _commit_video_span(
+            kv,
+            st,
+            kv_branch=POS,
+            n_chunks=history_chunks,
+            dtype=dtype,
+            device=device,
+        )
+        history_k_parts.append(k)
+        history_v_parts.append(v)
+
+    ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=commit_current)[0].forward_ctx
+    ctx.ensure_video_slots(device)
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0]
+    kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0]
+
+    action_k = action_v = None
+    if action_len:
+        ctx.ensure_action_slots(action_len, device)
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        kv._k_pools[0][ctx.action_slot_mapping] = action_k[0]
+        kv._v_pools[0][ctx.action_slot_mapping] = action_v[0]
+
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table(
+        action_len=action_len,
+        query_len=query.shape[1],
+        device=device,
+    )
+    paged = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=seq_lens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+
+    if history_k_parts:
+        history_k = torch.cat(history_k_parts, dim=1)
+        history_v = torch.cat(history_v_parts, dim=1)
+    else:
+        history_k = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        history_v = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)
+
+    before = st.adapter(POS).completed_chunks
+    st.commit_paged_context(POS)
+    assert st.adapter(POS).completed_chunks == before + (1 if commit_current else 0)
+
+
+@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required")
+@pytest.mark.parametrize("history_chunks", [1, 3])
+@pytest.mark.parametrize("action_len", [0, 3])
+@pytest.mark.parametrize("commit_current", [False, True])
+def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
+    pytest.importorskip("vllm.vllm_flash_attn")
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.float16
+    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
+
+    history_k, history_v = _commit_video_span(
+        kv,
+        st,
+        kv_branch=POS,
+        n_chunks=history_chunks,
+        dtype=dtype,
+        device=device,
+    )
+
+    layer_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=commit_current)[0]
+    ctx = layer_ctx.forward_ctx
+    current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+    action_k = action_v = None
+    if action_len:
+        action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+        action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
+
+    # The production path: once-per-forward host prep, then the fused
+    # write+attend custom op consuming the NamedTuple payload.
+    ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1])
+    inputs = layer_ctx.to_layer_inputs()
+    assert isinstance(inputs, ARDiffusionPagedLayerInputs)
+    paged = paged_write_attn(
+        inputs,
+        query[0],
+        current_k[0],
+        current_v[0],
+        action_k[0] if action_len else None,
+        action_v[0] if action_len else None,
+        HEAD_DIM**-0.5,
+    ).unsqueeze(0)
+
+    # Direct python-fn call on the same (already written) pools must be
+    # bit-exact: identical kernel, identical inputs.
+    direct = ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(0),
+        kv.value_cache(0),
+        block_table=ctx.block_table,
+        query_start_loc=ctx.query_start_loc,
+        seq_lens=ctx.seq_lens,
+        max_query_len=ctx.max_query_len,
+        max_seq_len=ctx.max_seq_len,
+        softmax_scale=HEAD_DIM**-0.5,
+        causal=False,
+    )
+    assert torch.equal(paged, direct)
+
+    new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :]
+    new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :]
+    if action_len:
+        new_k = torch.cat([new_k, action_k], dim=1)
+        new_v = torch.cat([new_v, action_v], dim=1)
+    ref = _dense_attention(query, new_k, new_v)
+
+    torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)
+
+
+def test_block_table_padded_to_fixed_width():
+    """Shapes must be constant across window growth: only values change."""
+    device = torch.device("cpu")
+    kv, st = make_state(window_chunks=2)
+
+    ctx1 = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx1.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(POS)
+
+    ctx2 = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    ctx2.prepare(device=device, action_len=0, query_len=BLOCK)
+    st.commit_paged_context(POS)
+
+    # 1-block vs 2-block visible history: same table width, same max_seq_len.
+    assert ctx1.block_table.shape == ctx2.block_table.shape
+    assert ctx1.max_seq_len == ctx2.max_seq_len
+    expected_width = kv.spec.sliding_window // kv.block_size + 1
+    assert ctx1.block_table.shape == (1, expected_width)
+    # Real lengths live in seq_lens, not the padded table.
+    assert int(ctx1.seq_lens[0]) == BLOCK
+    assert int(ctx2.seq_lens[0]) == 2 * BLOCK
+
+
+def test_prepare_is_idempotent_and_layers_share_metadata():
+    device = torch.device("cpu")
+    _, st = make_state(num_layers=2)
+    contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
+    fctx = contexts[0].forward_ctx
+    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+    table = fctx.block_table
+    fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+    assert fctx.block_table is table  # memoized, not rebuilt
+
+    i0, i1 = contexts[0].to_layer_inputs(), contexts[1].to_layer_inputs()
+    # 0-dim tensors (NOT python ints) so dynamo doesn't install per-layer
+    # value guards on the shared block code object.
+    assert isinstance(i0.layer_idx, torch.Tensor) and int(i0.layer_idx) == 0
+    assert isinstance(i1.layer_idx, torch.Tensor) and int(i1.layer_idx) == 1
+    # All layers share the same metadata tensor objects.
+    assert i0.block_table is i1.block_table
+    assert i0.seq_lens is i1.seq_lens
+    assert i0.video_slots is i1.video_slots
+
+
+def test_layer_inputs_before_prepare_raises():
+    _, st = make_state()
+    layer_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0]
+    with pytest.raises(RuntimeError, match="before prepare"):
+        layer_ctx.to_layer_inputs()
+
+
+def test_custom_op_registration_idempotent():
+    import importlib
+    import sys
+
+    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
+    mod = "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention"
+    saved = sys.modules.pop(mod)
+    try:
+        importlib.import_module(mod)  # re-registration must not raise
+    finally:
+        sys.modules[mod] = saved
+    assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn")
+
+
+def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
+    """The op must trace as one opaque node: fullgraph OK, and changed tensor
+    VALUES (new slots / block ids) must not trigger recompilation."""
+    import torch._dynamo
+
+    device = torch.device("cpu")
+    kv, st = make_state(num_layers=2, window_chunks=2)
+
+    def run_one_forward(commit):
+        contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=commit)
+        fctx = contexts[0].forward_ctx
+        fctx.prepare(device=device, action_len=0, query_len=BLOCK)
+        q = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        k = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        v = torch.randn(BLOCK, N_HEADS, HEAD_DIM)
+        # Both layers through ONE compiled fn: layer_idx is a tensor, so a
+        # different layer must NOT recompile (all 40 DiT blocks share the
+        # block-forward code object in production).
+        for layer_ctx in contexts:
+            out = compiled(layer_ctx.to_layer_inputs(), q, k, v)
+        st.commit_paged_context(POS)
+        return out
+
+    torch._dynamo.reset()
+    try:
+        from torch._dynamo.testing import CompileCounter
+
+        counter = CompileCounter()
+
+        def fn(inputs, q, k, v):
+            return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0
+
+        compiled = torch.compile(fn, backend=counter, fullgraph=True)
+
+        run_one_forward(commit=True)  # history grows between calls ->
+        run_one_forward(commit=True)  # block-table VALUES change, shapes don't
+        run_one_forward(commit=False)
+
+        assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"
+    finally:
+        # Leave a clean dynamo state for later suites in the same pytest process
+        # (e.g. model_executor transformers models).
+        torch._dynamo.reset()

tests/diffusion/ar_diffusion/test_pipeline_bridge.py

@@ -0,0 +1,280 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the ARDiffusionKVState paged-attention pipeline bridge."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionKVCache,
+    ARDiffusionKVConfig,
+    ARDiffusionPagedLayerContext,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+BLOCK = 16
+N_HEADS = 4
+HEAD_DIM = 64
+POS = "positive"
+NEG = "negative"
+
+
+def make_state(num_layers=1, window_chunks=4, cross_attn_length=0, shared_local_index=False):
+    cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
+    kv = ARDiffusionKVCache(
+        cfg,
+        num_layers=num_layers,
+        num_kv_heads=N_HEADS,
+        head_size=HEAD_DIM,
+        dtype=torch.float32,
+        block_size=BLOCK,
+        max_model_len=4096,
+        available_bytes=1 << 24,
+        kv_branches=(
+            ARDiffusionKVBranchSpec(POS, 0),
+            ARDiffusionKVBranchSpec(NEG, 0 if shared_local_index else 1),
+        ),
+        session_capacity=2,
+        cross_attention_lengths={"text": cross_attn_length} if cross_attn_length else None,
+        device=torch.device("cpu"),
+    )
+    pos = kv.begin_request("r-pos")
+    neg = kv.begin_request("r-neg")
+    return kv, ARDiffusionKVState(kv, "s1", {POS: pos, NEG: neg}, num_layers=num_layers)
+
+
+def _prepare_and_commit(st: ARDiffusionKVState, kv_branch: str, n_chunks: int) -> None:
+    ctx = st.get_kv_caches(kv_branch, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+    st.commit_paged_context(kv_branch)
+
+
+def test_get_returns_paged_layer_contexts_when_nothing_committed():
+    _, st = make_state(num_layers=2)
+    for kv_branch in (POS, NEG):
+        contexts = st.get_kv_caches(kv_branch, seq_len=BLOCK, commit_current=False)
+        assert len(contexts) == 2
+        assert all(isinstance(ctx, ARDiffusionPagedLayerContext) for ctx in contexts)
+        assert contexts[0].history_block_ids == []
+
+
+def test_managed_paged_context_commits_only_after_forward():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(POS, seq_len=2 * BLOCK, commit_current=True)
+    ctx = contexts[0].forward_ctx
+
+    assert st.adapter(POS).completed_chunks == 0
+    assert kv.window_block_ids(st.adapter(POS)) == []
+    ctx.ensure_video_slots(torch.device("cpu"))
+    assert st.adapter(POS).completed_chunks == 0
+    assert len(ctx.current_video_block_ids) == 2
+
+    st.commit_paged_context(POS)
+
+    assert st.adapter(POS).completed_chunks == 2
+    assert st._committed[POS] == 2 * BLOCK
+    assert len(kv.window_block_ids(st.adapter(POS))) == 2
+
+
+def test_scratch_paged_context_does_not_commit():
+    kv, st = make_state(num_layers=1)
+    contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
+    ctx = contexts[0].forward_ctx
+    ctx.ensure_video_slots(torch.device("cpu"))
+
+    assert ctx.current_video_block_ids == kv.scratch_block_ids(POS, 0, 1)
+    st.commit_paged_context(POS)
+
+    assert st.adapter(POS).completed_chunks == 0
+    assert st._committed[POS] == 0
+    assert kv.window_block_ids(st.adapter(POS)) == []
+
+
+def test_eviction_bounds_resident_window():
+    kv, st = make_state(num_layers=1, window_chunks=3)
+    _prepare_and_commit(st, POS, 3)
+    assert len(kv.window_block_ids(st.adapter(POS))) == 3
+
+    _prepare_and_commit(st, POS, 2)
+    ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
+    visible_blocks, video_len = ctx.video_block_table(torch.device("cpu"))
+    assert len(visible_blocks) == 3
+    assert video_len == 3 * BLOCK
+
+
+def test_branches_are_independent():
+    kv, st = make_state()
+    _prepare_and_commit(st, POS, 1)
+    assert len(kv.window_block_ids(st.adapter(POS))) == 1
+    assert kv.window_block_ids(st.adapter(NEG)) == []
+
+
+def test_reset_clears_session_window():
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    _prepare_and_commit(st, POS, 2)
+    _prepare_and_commit(st, NEG, 1)
+    assert st._committed[POS] == 2 * BLOCK and st._committed[NEG] == BLOCK
+    free_before = kv.manager.block_pool.get_num_free_blocks()
+
+    st.reset()
+
+    assert st._committed == {POS: 0, NEG: 0}
+    assert st._paged_pending == {POS: None, NEG: None}
+    assert kv.window_block_ids(st.adapter(POS)) == []
+    assert kv.window_block_ids(st.adapter(NEG)) == []
+    assert kv.manager.block_pool.get_num_free_blocks() > free_before
+    _prepare_and_commit(st, POS, 1)
+    assert len(kv.window_block_ids(st.adapter(POS))) == 1
+
+
+def test_get_cross_attention_kv_raises_before_populate():
+    _, st = make_state(cross_attn_length=8)
+    for kv_branch in (POS, NEG):
+        with pytest.raises(RuntimeError, match="before it was populated"):
+            st.get_cross_attention_kv(kv_branch, "text")
+
+
+def test_get_cross_attention_kv_returns_session_pool_dicts_when_populated():
+    L = 8
+    kv, st = make_state(num_layers=2, cross_attn_length=L)
+    written = []
+    for _ in range(2):
+        k = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        v = torch.randn(1, L, N_HEADS, HEAD_DIM)
+        written.append((k, v))
+    st.populate_cross_attention(POS, "text", written)
+
+    out = st.get_cross_attention_kv(POS, "text")
+    assert len(out) == 2
+    for i, (k, v) in enumerate(written):
+        assert out[i]["is_init"] is True
+        assert out[i]["k"].shape == (1, L, N_HEADS, HEAD_DIM)
+        assert torch.equal(out[i]["k"], k)
+        assert torch.equal(out[i]["v"], v)
+
+
+def test_window_reset_can_keep_named_cross_attention_cache():
+    kv, st = make_state(num_layers=1, cross_attn_length=8)
+    k = torch.randn(1, 8, N_HEADS, HEAD_DIM)
+    v = torch.randn(1, 8, N_HEADS, HEAD_DIM)
+    st.populate_cross_attention(POS, "text", [(k, v)])
+    old_adapter = st.adapter(POS)
+
+    st.reset(keep_cross_attention=("text",))
+
+    assert st.adapter(POS) is not old_adapter
+    assert torch.equal(st.get_cross_attention_kv(POS, "text")[0]["k"], k)
+    assert "s1" in kv._cross_sessions
+
+    st.reset()
+    assert "s1" not in kv._cross_sessions
+    assert not st.is_cross_attention_populated(POS, "text")
+
+
+def test_cross_attention_uses_logical_branch_when_local_slot_is_shared():
+    _, st = make_state(num_layers=2, cross_attn_length=8, shared_local_index=True)
+    assert st.kv_cache.num_local_kv_branches == 1
+    pos_layers = [
+        (torch.full((1, 8, N_HEADS, HEAD_DIM), 1.0), torch.full((1, 8, N_HEADS, HEAD_DIM), 2.0)) for _ in range(2)
+    ]
+    neg_layers = [
+        (torch.full((1, 8, N_HEADS, HEAD_DIM), 3.0), torch.full((1, 8, N_HEADS, HEAD_DIM), 4.0)) for _ in range(2)
+    ]
+
+    st.populate_cross_attention(POS, "text", pos_layers)
+    st.populate_cross_attention(NEG, "text", neg_layers)
+
+    for layer in st.get_cross_attention_kv(POS, "text"):
+        assert torch.count_nonzero(layer["k"] != 1.0) == 0
+        assert torch.count_nonzero(layer["v"] != 2.0) == 0
+    for layer in st.get_cross_attention_kv(NEG, "text"):
+        assert torch.count_nonzero(layer["k"] != 3.0) == 0
+        assert torch.count_nonzero(layer["v"] != 4.0) == 0
+
+
+def test_cross_attention_population_is_transactional():
+    kv, st = make_state(num_layers=2, cross_attn_length=8)
+    original = [
+        (torch.full((1, 8, N_HEADS, HEAD_DIM), value), torch.full((1, 8, N_HEADS, HEAD_DIM), -value))
+        for value in (1.0, 2.0)
+    ]
+    replacement = (
+        torch.full((1, 8, N_HEADS, HEAD_DIM), 9.0),
+        torch.full((1, 8, N_HEADS, HEAD_DIM), -9.0),
+    )
+
+    def interrupted_population():
+        yield replacement
+        raise RuntimeError("projection failed")
+
+    with pytest.raises(RuntimeError, match="projection failed"):
+        st.populate_cross_attention(POS, "text", interrupted_population())
+    assert not st.is_cross_attention_populated(POS, "text")
+    assert "s1" not in kv._cross_sessions
+
+    st.populate_cross_attention(POS, "text", original)
+
+    with pytest.raises(RuntimeError, match="projection failed"):
+        st.populate_cross_attention(POS, "text", interrupted_population())
+    out = st.get_cross_attention_kv(POS, "text")
+    for layer_index, layer in enumerate(out):
+        assert torch.equal(layer["k"], original[layer_index][0])
+        assert torch.equal(layer["v"], original[layer_index][1])
+
+    for invalid_layers in (original[:1], [*original, replacement]):
+        with pytest.raises(ValueError, match="expected 2 layers"):
+            st.populate_cross_attention(POS, "text", invalid_layers)
+        out = st.get_cross_attention_kv(POS, "text")
+        for layer_index, layer in enumerate(out):
+            assert torch.equal(layer["k"], original[layer_index][0])
+            assert torch.equal(layer["v"], original[layer_index][1])
+
+    invalid_shape = (torch.empty(1, 7, N_HEADS, HEAD_DIM), original[0][1])
+    with pytest.raises(ValueError, match="expected k/v shape"):
+        st.populate_cross_attention(POS, "text", [invalid_shape, original[1]])
+    out = st.get_cross_attention_kv(POS, "text")
+    assert torch.equal(out[0]["k"], original[0][0])
+
+
+def test_closed_state_cannot_repopulate_cross_attention():
+    kv, st = make_state(num_layers=1, cross_attn_length=8)
+    layer = (torch.ones(1, 8, N_HEADS, HEAD_DIM), torch.ones(1, 8, N_HEADS, HEAD_DIM))
+    st.close()
+
+    with pytest.raises(RuntimeError, match="is closed"):
+        st.populate_cross_attention(POS, "text", [layer])
+    assert "s1" not in kv._cross_sessions
+
+
+def test_kv_create_is_noop_engine_owns_allocation():
+    from unittest.mock import MagicMock
+
+    from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
+
+    p = DreamZeroPipeline.__new__(DreamZeroPipeline)
+    p._ar_diffusion_kv_state = object()
+    state = MagicMock()
+    p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
+    state.create_kv_caches.assert_not_called()
+
+
+def test_get_kv_cache_requires_frame_aligned_seqlen():
+    _, st = make_state()
+    with pytest.raises(AssertionError, match="frame-aligned"):
+        st.get_kv_caches(POS, seq_len=BLOCK + 1, commit_current=True)
+
+
+def test_prepare_one_branch_allocates_nothing_for_the_other():
+    """CFG-parallel laziness: a rank prepares only the kv_branch it runs."""
+    kv, st = make_state(num_layers=1, window_chunks=4)
+    pos_ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+    neg_ctx = st.get_kv_caches(NEG, seq_len=BLOCK, commit_current=True)[0].forward_ctx
+
+    pos_ctx.prepare(device=torch.device("cpu"), action_len=0, query_len=BLOCK)
+
+    assert pos_ctx._allocated_video
+    assert not neg_ctx._allocated_video
+    assert kv.window_block_ids(st.adapter(NEG)) == []

tests/diffusion/ar_diffusion/test_slot_mapping.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Slot-mapping / block-table tests for the AR-Diffusion engine (Phase 1, Step 3)."""
+
+import pytest
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache import (
+    ARDiffusionRequestAdapter,
+    ChunkWindowSpec,
+    build_kv_manager,
+    chunk_slot_mapping,
+    compute_slot_mapping,
+    resident_block_ids,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+BLOCK = 16
+
+
+def make_spec(*, window_chunks=2):
+    return ChunkWindowSpec(
+        block_size=BLOCK,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=window_chunks * BLOCK,
+        chunk_size=BLOCK,
+        window_chunks=window_chunks,
+    )
+
+
+def test_slot_mapping_matches_block_offsets():
+    # block table [5, 2, 9], block_size 4. slot = block*4 + (pos % 4).
+    slots = compute_slot_mapping([5, 2, 9], [0, 1, 4, 5, 8, 9], block_size=4)
+    assert slots.tolist() == [20, 21, 8, 9, 36, 37]
+
+
+def test_slot_mapping_rejects_bad_block_size():
+    with pytest.raises(ValueError):
+        compute_slot_mapping([1], [0], block_size=0)
+
+
+def test_chunk_slot_mapping_targets_current_chunk():
+    # After 2 committed chunks (num_computed=32), the in-flight chunk maps to the
+    # block covering positions 32..47 — i.e. block table index 2.
+    block_ids = [7, 4, 3]  # block index 2 -> physical block 3
+    slots = chunk_slot_mapping(block_ids, num_computed_tokens=32, chunk_size=4, block_size=16)
+    # positions 32..35 -> block_index 2 -> physical 3 -> slots 48..51
+    assert slots.tolist() == [48, 49, 50, 51]
+
+
+def test_resident_block_ids_excludes_null():
+    assert resident_block_ids([0, 1, 0, 2, 3], null_block_id=0) == [1, 2, 3]
+
+
+def test_blocktable_build_single_request():
+    """Integration: a real KVCacheManager block table -> a valid slot mapping.
+
+    Every slot for the in-flight chunk must land inside a real (non-null) block
+    that the manager actually allocated for the request.
+    """
+    spec = make_spec()
+    mgr = build_kv_manager(spec, ["l0"], num_blocks=16, max_model_len=1024)
+    null_id = mgr.block_pool.null_block.block_id
+    adapter = ARDiffusionRequestAdapter("req", chunk_size=BLOCK)
+
+    mgr.allocate_slots(adapter, num_new_tokens=BLOCK)
+    block_ids = mgr.get_block_ids(adapter.request_id)[0]
+    real = set(resident_block_ids(block_ids, null_id))
+    assert real, "expected at least one real block after allocation"
+
+    slots = chunk_slot_mapping(block_ids, adapter.num_computed_tokens, spec.chunk_size, spec.block_size)
+    # Each slot resolves to a real allocated block, none to null_block.
+    blocks_used = {int(s) // spec.block_size for s in slots}
+    assert blocks_used <= real
+    assert null_id not in blocks_used
+    assert len(slots) == spec.chunk_size