vllm-omni/tests/diffusion/ar_diffusion/test_paged_attention.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""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 import paged_attention as paged_attention_module
from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV
from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ChunkWindowManager
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,
    chunk_size=BLOCK,
    sink_chunks=0,
    reset_at_boundary=False,
    dtype=torch.float32,
    device=torch.device("cpu"),
    reuse_history_staging=False,
):
    """Build a cache. ``chunk_size`` defaults to the block size, but the two are
    independent -- the shipped 832x480 gives 1560 tokens per frame against
    16-token blocks, so a frame is 97.5 blocks."""
    cfg = ARDiffusionKVConfig(
        enable=True,
        chunk_size=chunk_size,
        window_chunks=window_chunks,
        sink_chunks=sink_chunks,
        reset_at_boundary=reset_at_boundary,
        reuse_history_staging=reuse_history_staging,
    )
    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 _gpu_flash_attn_usable() -> bool:
    if not torch.cuda.is_available():
        return False
    if torch.version.hip is not None:
        for module in ("aiter", "flash_attn"):
            try:
                imported = __import__(module, fromlist=["flash_attn_varlen_func"])
                if getattr(imported, "flash_attn_varlen_func", None) is not None:
                    return True
            except ImportError:
                pass
        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,
    chunk_size: int = BLOCK,
) -> tuple[torch.Tensor, torch.Tensor]:
    span = n_chunks * chunk_size
    ctx = st.get_kv_caches(kv_branch, seq_len=span, commit_current=True)[0].forward_ctx
    ctx.ensure_video_slots(device)
    k = torch.randn(1, span, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
    v = torch.randn(1, span, 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.parametrize("history_chunks", [0, 1, 2, 3])
@pytest.mark.parametrize("window_chunks", [2, 4])
def test_staged_reuse_refreshes_the_current_blocks_at_their_live_offset(monkeypatch, history_chunks, window_chunks):
    """A second probe of the same AR block restages its current K/V where the table actually holds it.

    The padded table always ends in at least one action-capacity block, and while the window is still
    growing in unused window capacity too, so "the last current_blocks entries" is padding. The refresh
    has to start right after the visible history: empty (window growing), partial, full, and after a slide.
    """
    monkeypatch.setenv(KV_GATHER_ENV, "1")
    device = torch.device("cpu")
    dtype = torch.float32
    kv, st = make_state(dtype=dtype, device=device, window_chunks=window_chunks, reuse_history_staging=True)
    for buffer in kv.history_staging[0]:
        buffer.fill_(-7)
    if history_chunks:
        _commit_video_span(kv, st, kv_branch=POS, n_chunks=history_chunks, dtype=dtype, device=device)

    ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
    ctx.max_video_tokens = 2 * BLOCK
    ctx.ensure_video_slots(device)
    key_cache, value_cache = kv.key_cache(0), kv.value_cache(0)

    def probe(step: int):
        # Each probe of the block writes new current K/V, then stages.
        kv._k_pools[0][ctx.current_video_slot_mapping] = torch.full((BLOCK, N_HEADS, HEAD_DIM), float(step))
        kv._v_pools[0][ctx.current_video_slot_mapping] = torch.full((BLOCK, N_HEADS, HEAD_DIM), -float(step))
        # Store the metadata on the context the way prepare() does.
        (ctx.block_table, ctx.query_start_loc, ctx.seq_lens, ctx.max_query_len, ctx.max_seq_len) = (
            ctx.build_block_table(action_len=0, query_len=BLOCK, device=device)
        )
        block_table, max_seq_len = ctx.block_table, ctx.max_seq_len
        ctx._prepare_history_staging(0)
        n_blocks = max_seq_len // BLOCK
        block_ids = block_table[0, :n_blocks].to(torch.long)
        stage_k, stage_v = ctx.history_staging(0)
        # The custom op narrows the manager-owned buffers before calling this helper.
        stage_k, stage_v = stage_k[:max_seq_len], stage_v[:max_seq_len]
        paged_attention_module._stage_window(
            stage_k,
            stage_v,
            key_cache,
            value_cache,
            block_ids,
            n_blocks,
            BLOCK,
            first_block=ctx.stage_first_block if ctx.reuse_history else 0,
        )
        full_k = key_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
        full_v = value_cache.index_select(0, block_ids).reshape(n_blocks * BLOCK, N_HEADS, HEAD_DIM)
        for buffer in kv.history_staging[0]:
            assert (buffer[max_seq_len:] == -7).all()
        return stage_k, stage_v, full_k, full_v, n_blocks

    stage_k, stage_v, full_k, full_v, n_blocks = probe(1)
    assert ctx.reuse_history is False
    assert torch.equal(stage_k, full_k) and torch.equal(stage_v, full_v)

    stage_k, stage_v, full_k, full_v, n_blocks = probe(2)
    assert ctx.reuse_history is True
    visible_history_blocks = min(history_chunks, 2 - 1)  # window of 2 blocks minus the current one
    assert ctx.stage_first_block == visible_history_blocks
    # The padded table is wider than the live window, so the end-of-table guess is not the offset.
    assert n_blocks - 1 != ctx.stage_first_block
    assert torch.equal(stage_k, full_k) and torch.equal(stage_v, full_v)


@pytest.mark.parametrize("gather_enabled", [False, True])
def test_staging_is_only_allocated_for_the_gather_path(monkeypatch, gather_enabled):
    if gather_enabled:
        monkeypatch.setenv(KV_GATHER_ENV, "1")
    else:
        monkeypatch.delenv(KV_GATHER_ENV, raising=False)
    device = torch.device("cpu")
    kv, st = make_state(device=device, reuse_history_staging=True)
    # The manager owns the pairs and allocates them only when their consumer is on.
    assert bool(kv.history_staging) is gather_enabled
    ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
    ctx.ensure_video_slots(device)
    (ctx.block_table, ctx.query_start_loc, ctx.seq_lens, ctx.max_query_len, ctx.max_seq_len) = ctx.build_block_table(
        action_len=0, query_len=BLOCK, device=device
    )
    ctx._prepare_history_staging(0)
    assert ctx.staging_enabled is gather_enabled
    stage_k, _ = ctx.history_staging(0)
    assert (stage_k is not None) is gather_enabled


@pytest.mark.parametrize("window_chunks", [2, 4])
def test_layer_inputs_preserve_static_staging_tensors(monkeypatch, window_chunks):
    """Compiled inputs must retain the manager's static-address annotations, including with spare capacity."""
    monkeypatch.setenv(KV_GATHER_ENV, "1")
    device = torch.device("cpu")
    kv, st = make_state(device=device, window_chunks=window_chunks, reuse_history_staging=True)
    stage_key, stage_value = kv.history_staging[0]
    for _ in range(2):
        ctx = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)[0].forward_ctx
        ctx.max_video_tokens = 2 * BLOCK
        ctx.prepare(device, action_len=0, query_len=BLOCK)
        inputs = ctx.layer_inputs(0)
        assert inputs.stage_key is stage_key
        assert inputs.stage_value is stage_value
        assert inputs.max_seq_len == 3 * BLOCK
        assert stage_key.shape[0] == (window_chunks + 1) * BLOCK


@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU 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):
    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_k is not None else None,
        action_v[0] if action_v is not None 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")
    kv, 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
    assert i0.key_pool is kv._k_pools[0]
    assert i0.value_pool is kv._v_pools[0]
    assert i1.key_pool is kv._k_pools[1]
    assert i1.value_pool is kv._v_pools[1]


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_mutable_arguments_cannot_be_elided_as_defaults():
    # Older PyTorch ADInplaceOrView handlers index positional mutable inputs
    # directly. Default-valued trailing inputs can disappear before that handler.
    schema = torch.ops.vllm_omni.ar_diffusion_paged_write_attn.default._schema
    mutable = [arg for arg in schema.arguments if arg.alias_info is not None and arg.alias_info.is_write]
    assert {arg.name for arg in mutable} == {"key_pool", "value_pool", "stage_key", "stage_value"}
    assert all(not arg.has_default_value() for arg in mutable)


@pytest.mark.parametrize("reuse_history_staging", [False, True])
def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(monkeypatch, reuse_history_staging):
    """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")
    monkeypatch.setenv(KV_GATHER_ENV, "1")
    kv, st = make_state(num_layers=2, window_chunks=2, reuse_history_staging=reuse_history_staging)
    if reuse_history_staging:
        # Hold the host-side staging offset fixed while table/slot values change.
        _commit_video_span(kv, st, kv_branch=POS, n_chunks=2, dtype=torch.float32, device=device)

    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 use one compiled function; the tensor-valued layer index
        # must not specialize it.
        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()


# ── a frame that is not a whole number of blocks ────────────────────────────

# 24 tokens per chunk against 16-token blocks leaves 8 over, the same remainder
# the shipped 832x480 produces (1560 % 16 == 8). Every fixture above holds
# chunk_size == BLOCK, so nothing there can reach this path.
RAGGED_CHUNK = 24


@pytest.mark.parametrize("commit_current", [False, True])
@pytest.mark.parametrize("action_len", [0, 3])
def test_a_ragged_chunk_still_matches_the_dense_reference(commit_current, action_len):
    """Attention reads a sequence, not a set of blocks.

    With one committed chunk of 24 tokens the history stops 8 slots into its
    last block. The committing path writes straight after it and is fine. The
    scratch path used to restart at slot zero of a fresh region, which left
    those 8 slots unwritten -- and since the kernel consumes the block table as
    one contiguous run, it read them as if they were tokens and shifted every
    token after them by 8. Shapes stayed correct throughout, so only the values
    showed it.
    """
    torch.manual_seed(0)
    device = torch.device("cpu")
    dtype = torch.float32
    # window_chunks=2 with one chunk of history keeps the whole sequence
    # resident, so the window does not also round up here; that case is
    # test_a_ragged_window_keeps_the_block_table_a_fixed_shape below.
    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=RAGGED_CHUNK)

    history_k, history_v = _commit_video_span(
        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=RAGGED_CHUNK
    )

    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
    ctx.ensure_video_slots(device)
    assert ctx.start_offset == RAGGED_CHUNK % BLOCK == 8

    current_k = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
    current_v = torch.randn(1, RAGGED_CHUNK, 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, RAGGED_CHUNK + 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,
    )

    new_k = torch.cat([history_k, current_k], dim=1)
    new_v = torch.cat([history_v, current_v], dim=1)
    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)


def test_a_ragged_scratch_chunk_is_physically_continuous_with_its_history():
    """The write targets themselves must leave no hole.

    This is the property the values above depend on, asserted directly so a
    regression names the cause rather than showing a numeric mismatch.
    """
    device = torch.device("cpu")
    kv, st = make_state(device=device, window_chunks=2, chunk_size=RAGGED_CHUNK)
    _commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=torch.float32, device=device, chunk_size=RAGGED_CHUNK)

    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=False)[0].forward_ctx
    ctx.ensure_video_slots(device)

    history_tail_block = ctx.history_block_ids[-1]
    # The chunk starts inside the history's own last block, not at a fresh one.
    assert ctx.current_video_block_ids[0] == history_tail_block
    first_slot = int(ctx.current_video_slot_mapping[0])
    assert first_slot == history_tail_block * BLOCK + 8

    # And the action region must not be charged for that managed block.
    ctx.ensure_action_slots(3, device)
    assert ctx._scratch_blocks_used == len(ctx.current_video_block_ids) - 1
    assert ctx.action_scratch_block_ids[0] not in ctx.current_video_block_ids


def test_a_ragged_window_keeps_the_block_table_a_fixed_shape():
    """The table's shape must not track how full the window is.

    With a sink the visible window is 2*24 + 24 = 72 tokens, which is 4.5
    blocks. Deriving the width by flooring that would understate the capacity,
    the max() against the live block count would take over, and the shape would
    change as the window filled -- recompiling the graph. The same floor would
    also let max_seq_len come out under the kv_len actually passed.
    """
    device = torch.device("cpu")
    kv, st = make_state(device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK)
    # The window is deliberately not a whole number of blocks.
    max_video_tokens = kv.spec.sliding_window + kv.spec.sink_chunks * kv.spec.chunk_size
    assert max_video_tokens % BLOCK != 0

    widths: set[int] = set()
    kv_lens: list[int] = []
    for _ in range(6):
        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
        ctx.ensure_video_slots(device)
        block_table, _, seq_lens, _, max_seq_len = ctx.build_block_table(
            action_len=0, query_len=RAGGED_CHUNK, device=device
        )
        widths.add(int(block_table.shape[1]))
        # A bound handed to the kernel has to actually bound the sequence.
        assert int(seq_lens[0]) <= max_seq_len
        kv_lens.append(int(seq_lens[0]))
        st.commit_paged_context(POS)

    assert len(widths) == 1, f"block table width varied across ticks: {sorted(widths)}"
    # The window has to have actually filled, or none of the above was tested.
    assert max(kv_lens) > min(kv_lens)
    assert max(kv_lens) <= ctx.max_video_blocks * BLOCK
    # Whole blocks are read, so each of the two boundaries -- the sink's end and
    # the recent window's start -- can contribute up to BLOCK - 1 tokens the
    # window no longer needs.
    assert max(kv_lens) <= max_video_tokens + 2 * (BLOCK - 1)


def _poison_pools(kv: ARDiffusionKVCache) -> None:
    for pool in (*kv._k_pools, *kv._v_pools):
        pool.fill_(float("nan"))


def _write_prepared_chunk(
    kv: ARDiffusionKVCache, ctx, *, dtype: torch.dtype, device: torch.device
) -> tuple[torch.Tensor, torch.Tensor]:
    """Write fresh K/V for a prepared chunk, poisoning the rest of its blocks first.

    Blocks are handed out whole and a reused block keeps whatever it held, so a
    slot read without having been written would not look unusual. Every slot of
    the chunk's blocks that this chunk does not write is set to NaN first --
    except the leading slots of a block the history still occupies -- so reading
    one turns the attention output NaN.
    """
    mapping = ctx.current_video_slot_mapping
    first_slot = int(mapping[0])
    for index, block in enumerate(ctx.current_video_block_ids):
        start = first_slot if index == 0 else block * BLOCK
        kv._k_pools[0][start : (block + 1) * BLOCK] = float("nan")
        kv._v_pools[0][start : (block + 1) * BLOCK] = float("nan")
    k = torch.randn(ctx.seq_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
    v = torch.randn(ctx.seq_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
    kv._k_pools[0][mapping] = k
    kv._v_pools[0][mapping] = v
    return k, v


def _window_positions(end: int, *, sink: int, window: int, written, history: int, resident: set[int]) -> list[int]:
    """Token positions attention must read at ``end``.

    Every token of the current chunk, plus every written history token in a
    block that is still resident and holds a sink token or one of the window's
    most recent tokens.
    """
    start = end - window
    blocks = {p // BLOCK for p in range(min(sink, end))} | {p // BLOCK for p in range(max(start, 0), end)}
    return [p for p in sorted(written) if p >= history or (p // BLOCK in blocks and p // BLOCK in resident)]


def _model_block_indices(kv, adapter) -> dict[int, int]:
    """Each resident block's index in model positions, by block id.

    Compaction removes the evicted gap after the sink's blocks and shifts the
    rest of the table down by it, so a table index past the sink is
    ``compacted_tokens`` behind the position it holds.
    """
    block_size = kv.block_size
    sink_blocks = -(-(kv.spec.sink_chunks * kv.spec.chunk_size) // block_size)
    shift = adapter.compacted_tokens // block_size
    return {
        int(block): index if index < sink_blocks else index + shift
        for index, block in enumerate(kv.block_table(adapter))
        if block != kv.null_block_id
    }


@pytest.mark.parametrize(
    ("sink_chunks", "window_chunks", "reset_at_boundary"),
    [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)],
)
def test_attention_after_eviction_reads_every_kept_token_and_nothing_unwritten(
    sink_chunks, window_chunks, reset_at_boundary
):
    """Once the window slides, the table must still cover the window exactly.

    The sink and the recent window are separate token ranges, and each can
    straddle a block edge on its own: 24-token chunks against 16-token blocks
    put a boundary 8 tokens past an edge on alternating chunks. Counting blocks
    off the ends of the resident list dropped the block holding the window's
    first tokens whenever both boundaries straddled, and the returned length
    ran past the last written slot. Neither changes a shape, so this compares
    values against dense attention over the tokens the window keeps -- after
    eviction, with every unwritten slot poisoned, on both the scratch and the
    committing forward of every chunk.

    The reset case keeps only the sink across a boundary, so the recent window
    reaches back over blocks that are no longer resident. Those must be skipped
    by position rather than read.

    Eight chunks are enough for compaction to drop the evicted gap from the
    table, after which storage positions run behind model positions. Removing
    a gap that is not a whole number of chunks moves where eviction snaps, and
    in the reset case the window then keeps tokens a boundary should have
    dropped.
    """
    torch.manual_seed(0)
    device, dtype = torch.device("cpu"), torch.float32
    kv, st = make_state(
        device=device,
        window_chunks=window_chunks,
        sink_chunks=sink_chunks,
        reset_at_boundary=reset_at_boundary,
        chunk_size=RAGGED_CHUNK,
    )
    _poison_pools(kv)
    sink, window = sink_chunks * RAGGED_CHUNK, int(kv.spec.sliding_window)
    committed: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
    evicted = compacted = False

    for chunk in range(8):
        history = int(st.adapter(POS).absolute_num_computed_tokens)
        end = history + RAGGED_CHUNK
        for commit_current in (False, True):
            ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
            ctx.ensure_video_slots(device)
            k, v = _write_prepared_chunk(kv, ctx, dtype=dtype, device=device)
            written = dict(committed)
            written.update({history + i: (k[i], v[i]) for i in range(RAGGED_CHUNK)})

            query = torch.randn(1, RAGGED_CHUNK, 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=0, query_len=RAGGED_CHUNK, device=device
            )
            resident = set(_model_block_indices(kv, st.adapter(POS)).values())
            positions = _window_positions(
                end, sink=sink, window=window, written=written, history=history, resident=resident
            )
            label = f"chunk {chunk}, commit_current={commit_current}"
            assert int(seq_lens[0]) == len(positions), f"{label}: kv_len does not match the written slots it covers"
            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,
            )
            ref = _dense_attention(
                query,
                torch.stack([written[p][0] for p in positions]).unsqueeze(0),
                torch.stack([written[p][1] for p in positions]).unsqueeze(0),
            )
            torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5, msg=lambda m, label=label: f"{label}: {m}")
        committed = written
        st.commit_paged_context(POS)
        compacted = compacted or st.adapter(POS).compacted_tokens > 0
        evicted = evicted or compacted or kv.null_block_id in kv.block_table(st.adapter(POS))

    assert evicted, "the window never slid, so nothing after eviction was tested"
    assert compacted, "the table was never compacted, so storage and model positions never diverged"


@pytest.mark.parametrize(
    ("sink_chunks", "window_chunks", "reset_at_boundary"),
    [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)],
)
def test_compaction_does_not_change_which_tokens_stay_resident(
    monkeypatch, sink_chunks, window_chunks, reset_at_boundary
):
    """Compaction renumbers the block table; it must not change what eviction keeps.

    Storage positions run ``compacted_tokens`` behind model positions, and
    eviction snaps to chunk boundaries counted in storage positions. Removing a
    gap that is not a whole number of chunks moves every later snap, so eviction
    keeps a different set of tokens -- in the reset case, tokens a boundary
    should have dropped, which the window then reads. Attention checked against
    what is resident cannot see that, so the oracle here is the same rollout
    with compaction switched off.
    """

    def resident_blocks_per_tick(compact: bool) -> tuple[list[set[int]], int]:
        with monkeypatch.context() as patch:
            if not compact:
                patch.setattr(ChunkWindowManager, "compact_block_table", lambda self, request_id: 0)
            kv, st = make_state(
                window_chunks=window_chunks,
                sink_chunks=sink_chunks,
                reset_at_boundary=reset_at_boundary,
                chunk_size=RAGGED_CHUNK,
            )
            ticks = []
            for _ in range(12):
                ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
                ctx.ensure_video_slots(torch.device("cpu"))
                ticks.append(set(_model_block_indices(kv, st.adapter(POS)).values()))
                st.commit_paged_context(POS)
            return ticks, st.adapter(POS).compacted_tokens

    expected, _ = resident_blocks_per_tick(compact=False)
    actual, compacted_tokens = resident_blocks_per_tick(compact=True)
    assert compacted_tokens > 0, "the table was never compacted, so nothing was tested"
    for tick, (want, got) in enumerate(zip(expected, actual)):
        assert got == want, f"tick {tick}: compaction kept blocks {sorted(got - want)} and dropped {sorted(want - got)}"


def test_history_staging_holds_a_ragged_window_and_restages_it_whole(monkeypatch):
    """The staging buffers must fit the table width this module pads to.

    A sink or window that is not a whole number of blocks spans one more block
    than its token count suggests, so sizing staging from tokens left it short
    and the first prepared forward refused to run. And because such a chunk
    shares its first block with the history's tail, the history a probe sees is
    not left untouched by the next probe, so it is restaged rather than reused.
    """
    monkeypatch.setenv(KV_GATHER_ENV, "1")
    device = torch.device("cpu")
    kv, st = make_state(
        device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK, reuse_history_staging=True
    )
    capacity = int(kv.history_staging[0][0].shape[0])
    for chunk in range(6):
        for probe in range(2):
            ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=False)[0].forward_ctx
            ctx.prepare(device=device, action_len=0, query_len=RAGGED_CHUNK)
            assert int(ctx.max_seq_len) <= capacity
            assert not ctx.reuse_history, f"chunk {chunk}, probe {probe}: a ragged chunk reused the staged history"
        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
        ctx.prepare(device=device, action_len=0, query_len=RAGGED_CHUNK)
        st.commit_paged_context(POS)


@pytest.mark.parametrize("commit_current", [False, True])
def test_action_tokens_after_a_partly_written_video_block_are_refused(commit_current):
    """Action K/V follows the video blocks in the table, and the kernel reads it as one run.

    A 24-token chunk from an empty history ends 8 slots into its second 16-token
    block, so the run would read those 8 unwritten slots as the first action
    tokens and never reach the last ones. Every shape stays right, which is why
    this is refused rather than computed.
    """
    _, st = make_state(window_chunks=2, chunk_size=RAGGED_CHUNK)
    ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx
    with pytest.raises(ValueError, match="partly written video block"):
        ctx.build_block_table(action_len=3, query_len=RAGGED_CHUNK + 3, device=torch.device("cpu"))


def test_the_shipped_geometry_reads_exactly_the_window_on_every_tick():
    """832x480 with the checkpoint's own sink and window, slot by slot.

    A 9-frame sink and a 9-frame window are each 8 tokens past a 16-token edge
    but 28080 tokens together -- a whole number of blocks -- so a width taken
    from their sum was one block short on every tick once the window had
    filled. Twelve ticks of three frames, checking every slot the kernel reads
    against the token that slot is supposed to hold. No attention is computed:
    at this size the bookkeeping is the whole question.
    """
    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size

    frame, frames_per_tick = (480 // 16) * (832 // 16), 3
    block = paging_block_size(frame)
    seq_len = frames_per_tick * frame
    cfg = ARDiffusionKVConfig(enable=True, chunk_size=frame, window_chunks=9, sink_chunks=9)
    kv = ARDiffusionKVCache(
        cfg,
        num_layers=1,
        num_kv_heads=1,
        head_size=8,
        dtype=torch.float16,
        block_size=block,
        max_model_len=1 << 20,
        available_bytes=1 << 27,
        kv_branches=(ARDiffusionKVBranchSpec(POS, 0),),
        session_capacity=1,
        frames_per_block=frames_per_tick,
        max_scratch_tokens_per_branch=seq_len,
        device=torch.device("cpu"),
    )
    st = ARDiffusionKVState(kv, "s1", {POS: kv.begin_request("r-pos")}, num_layers=1)
    sink, window = 9 * frame, int(kv.spec.sliding_window)
    holds: dict[int, int] = {}  # slot -> token position last written there
    widths: set[int] = set()
    filled = compacted = False

    for tick in range(12):
        history = int(st.adapter(POS).absolute_num_computed_tokens)
        end = history + seq_len
        ctx = st.get_kv_caches(POS, seq_len=seq_len, commit_current=True)[0].forward_ctx
        ctx.ensure_video_slots(torch.device("cpu"))
        for offset, slot in enumerate(ctx.current_video_slot_mapping.tolist()):
            holds[slot] = history + offset
        table, _, seq_lens, _, _ = ctx.build_block_table(action_len=0, query_len=seq_len, device=torch.device("cpu"))
        widths.add(int(table.shape[1]))
        kv_len = int(seq_lens[0])

        position_of_block = {b: i * block for b, i in _model_block_indices(kv, st.adapter(POS)).items()}
        read = []
        for b in table[0].tolist()[: -(-kv_len // block)]:
            assert b in position_of_block, f"tick {tick}: read block {b}, which holds no token of this session"
            read.extend((b * block + o, position_of_block[b] + o) for o in range(block))
        read = read[:kv_len]
        wrong = [(slot, want) for slot, want in read if holds.get(slot) != want]
        assert not wrong, f"tick {tick}: {len(wrong)} slots read hold a token other than the one their position implies"

        start = end - window
        kept = set(range(min(sink, end))) | set(range(max(start, 0), end))
        missed = kept - {want for _, want in read}
        assert not missed, f"tick {tick}: {len(missed)} tokens the window keeps were not read"

        filled = filled or end > sink + window
        st.commit_paged_context(POS)
        compacted = compacted or st.adapter(POS).compacted_tokens > 0

    assert filled, "the window never filled, so the case that fails was never reached"
    assert compacted, "the table was never compacted, so storage and model positions never diverged"
    assert len(widths) == 1, f"block table width varied across ticks: {sorted(widths)}"


class _CountedBlockIds:
    """Wraps a vLLM block manager's ``get_blocks`` and counts the block ids read.

    Every path to a request's block ids goes through ``get_blocks``, so a copy of
    the whole table is counted in full even if only a few entries are indexed.
    """

    def __init__(self, manager):
        self.reads = 0
        self._get_blocks = manager.get_blocks

    def __call__(self, request_id):
        from vllm.v1.core.kv_cache_manager import KVCacheBlocks

        blocks = self._get_blocks(request_id)
        return KVCacheBlocks(tuple([_CountedBlock(block, self) for block in group] for group in blocks.blocks))


class _CountedBlock:
    def __init__(self, block, counter: _CountedBlockIds):
        self._block = block
        self._counter = counter

    @property
    def block_id(self) -> int:
        self._counter.reads += 1
        return self._block.block_id


def test_reading_the_window_touches_only_the_blocks_it_can_keep(monkeypatch):
    """However long the block table is, reading the window must not cost more.

    Evicted positions stay in the table as null entries until compaction drops
    them a whole number of chunks at a time -- up to two frames' worth of blocks
    at 832x480 -- so the table can still be far longer than the window. Only the
    sink's blocks and the recent window's blocks can hold a kept token, so those
    are all the read may touch. Compaction is switched off here so the table
    grows to many times the window, and reads are counted where the ids leave
    vLLM's block manager.
    """
    monkeypatch.setattr(ChunkWindowManager, "compact_block_table", lambda self, request_id: 0)
    device = torch.device("cpu")
    kv, st = make_state(device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK)
    reads: list[int] = []
    table_lengths: list[int] = []
    for _ in range(60):
        ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx
        ctx.ensure_video_slots(device)
        counted = _CountedBlockIds(kv.manager)
        kv.manager.get_blocks = counted
        try:
            ctx.build_block_table(action_len=0, query_len=RAGGED_CHUNK, device=device)
        finally:
            del kv.manager.get_blocks
        reads.append(counted.reads)
        table_lengths.append(len(kv.block_table(st.adapter(POS))))
        st.commit_paged_context(POS)

    assert table_lengths[-1] > 10 * ctx.max_video_blocks, "the table never outgrew the window, so nothing was tested"
    assert max(reads) <= ctx.max_video_blocks, (
        f"read {max(reads)} table entries for a {ctx.max_video_blocks}-block window"
    )


def test_the_checkpoints_own_default_resolution_can_build_a_cache():
    """832x480 is what LingBot World v2 ships as its default, and it could not run.

    A frame-sized block made the resolution a kernel-compatibility question:
    1560 tokens per frame is not a multiple of 16, so FlashAttention's paged
    kernel rejected it and the default resolution had no realtime path at all.
    """
    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size

    tokens_per_frame = (480 // 16) * (832 // 16)
    assert tokens_per_frame == 1560
    assert tokens_per_frame % 16 == 8, "the whole point is that a frame is not a legal block"

    block_size = paging_block_size(tokens_per_frame)
    cfg = ARDiffusionKVConfig(enable=True, chunk_size=tokens_per_frame, window_chunks=2)
    kv = ARDiffusionKVCache(
        cfg,
        num_layers=1,
        num_kv_heads=N_HEADS,
        head_size=HEAD_DIM,
        dtype=torch.float32,
        block_size=block_size,
        max_model_len=1 << 16,
        available_bytes=1 << 28,
        kv_branches=(ARDiffusionKVBranchSpec(POS, 0),),
        session_capacity=1,
        frames_per_block=2,
        max_scratch_tokens_per_branch=block_size,
        device=torch.device("cpu"),
    )

    assert kv.block_size % 16 == 0
    # The eviction unit is still the frame; only the paging unit changed.
    assert kv.spec.chunk_size == tokens_per_frame
    assert kv.blocks_per_frame == -(-tokens_per_frame // block_size) == 98


def test_the_shipped_resolution_geometry_matches_dense_attention():
    """The real number, not a scaled-down stand-in: 1560 tokens per frame.

    This is the case the issue reported failing at ~4e-02 on the scratch path.
    It runs on CPU because the reference is dense attention, not a kernel.
    """
    torch.manual_seed(0)
    device = torch.device("cpu")
    dtype = torch.float32
    tokens_per_frame = (480 // 16) * (832 // 16)
    assert tokens_per_frame == 1560

    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=tokens_per_frame)
    history_k, history_v = _commit_video_span(
        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=tokens_per_frame
    )

    # The non-committing path is the one that was wrong: four of the five
    # forwards per generated block take it.
    ctx = st.get_kv_caches(POS, seq_len=tokens_per_frame, commit_current=False)[0].forward_ctx
    ctx.ensure_video_slots(device)
    assert ctx.start_offset == 8, "1560 % 16 == 8 is the whole reason this case exists"

    current_k = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
    current_v = torch.randn(1, tokens_per_frame, 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]

    query = torch.randn(1, tokens_per_frame, 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=0, query_len=tokens_per_frame, 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,
    )
    ref = _dense_attention(
        query,
        torch.cat([history_k, current_k], dim=1),
        torch.cat([history_v, current_v], dim=1),
    )
    torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5)


@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
@pytest.mark.parametrize("commit_current", [False, True])
def test_the_shipped_resolution_geometry_on_the_real_kernel(commit_current):
    """The shipped 832x480 geometry through FlashAttention's paged kernel.

    The CPU test above proves the addressing is right against a dense
    reference. This one proves the real kernel agrees, at the real 1560
    tokens per frame, through the production path -- host prep followed by
    the fused write+attend op, which is what a forward actually calls.
    """
    torch.manual_seed(0)
    device = torch.device("cuda")
    dtype = torch.float16
    tokens_per_frame = (480 // 16) * (832 // 16)
    assert tokens_per_frame == 1560

    kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=tokens_per_frame)
    history_k, history_v = _commit_video_span(
        kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=tokens_per_frame
    )

    layer_ctx = st.get_kv_caches(POS, seq_len=tokens_per_frame, commit_current=commit_current)[0]
    ctx = layer_ctx.forward_ctx
    current_k = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
    current_v = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)
    query = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device)

    ctx.prepare(device=device, action_len=0, query_len=tokens_per_frame)
    assert ctx.start_offset == 8, "1560 % 16 == 8 is the whole reason this case exists"
    assert kv.block_size == 16

    inputs = layer_ctx.to_layer_inputs()
    paged = paged_write_attn(inputs, query[0], current_k[0], current_v[0], None, None, HEAD_DIM**-0.5).unsqueeze(0)

    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 :]
    ref = _dense_attention(query, new_k, new_v)

    torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2)


@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
@pytest.mark.parametrize("history_chunks", [0, 1, 3])
@pytest.mark.parametrize("action_len", [0, 3])
def test_contiguous_kv_gather_path_matches_paged_path_gpu(monkeypatch, history_chunks, action_len):
    """VLLM_OMNI_AR_DIFFUSION_KV_GATHER=1 gathers the visible blocks and runs varlen FA3 without a block table.

    Covers an empty, partial and full window plus a partially filled action block:
    the gather reads the tail-padding null block, so its (zeroed) rows must be masked.
    """

    torch.manual_seed(0)
    device = torch.device("cuda")
    dtype = torch.float16
    kv, st = make_state(dtype=dtype, device=device, window_chunks=2)
    if history_chunks:
        _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=False)[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)
    ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1])
    inputs = layer_ctx.to_layer_inputs()

    def run() -> torch.Tensor:
        return paged_write_attn(
            inputs,
            query[0],
            current_k[0],
            current_v[0],
            action_k[0] if action_k is not None else None,
            action_v[0] if action_v is not None else None,
            HEAD_DIM**-0.5,
        ).unsqueeze(0)

    monkeypatch.delenv(KV_GATHER_ENV, raising=False)
    paged = run()
    monkeypatch.setenv(KV_GATHER_ENV, "1")
    gathered = run()
    assert torch.isfinite(gathered).all()
    # Same kernel family on the same K/V: only accumulation order differs.
    torch.testing.assert_close(gathered, paged, rtol=2e-3, atol=2e-3)