[AR-Diffusion] Page in kernel-legal blocks so the default resolution can run (#6481)

Signed-off-by: linzhenpl07 <linzhenpl07@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>

tests/diffusion/ar_diffusion/test_capability_runner.py

@@ -376,10 +376,17 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline():
 
 def test_budget_reduced_capacity_drives_runner_lru():
     pipeline = CapablePipeline(tiny_spec(capacity=2))
-    # One tiny LingBot-like session requires 128 bytes; two require 192.
+    # One tiny LingBot-like session requires 1,808 bytes; two require 2,592.
+    #
+    # These read 128 and 192 while the paging unit was the frame. tiny_spec
+    # declares one token per frame, so it used to page a single token at a
+    # time -- a block size no attention kernel accepts. Paging at 16 makes
+    # each page sixteen times larger. The block *counts* are unchanged, and
+    # so is what this test is about: a budget that fits fewer sessions than
+    # the pipeline asked for must drive the runner's LRU.
     runner = make_runner(
         pipeline,
-        available_bytes=128,
+        available_bytes=2048,
         gpu_memory_fraction=1.0,
     )
     kv = runner.kv_cache
@@ -665,3 +672,106 @@ def test_pipeline_without_warmup_provider_is_safely_skipped(monkeypatch):
     monkeypatch.setattr(runner, "execute_model", fail_if_called)
     runner._warmup_ar_rollout()
     assert execute.called is False
+
+
+# ── paging unit vs eviction unit ────────────────────────────────────────────
+
+
+def test_a_frame_that_is_already_a_legal_block_keeps_its_paging():
+    """Every resolution that runs today must page exactly as it does today."""
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    for tokens_per_frame in (16, 320, 1440, 4096):
+        assert tokens_per_frame % 16 == 0
+        assert paging_block_size(tokens_per_frame) == tokens_per_frame
+
+
+def test_a_frame_the_kernel_would_reject_pages_at_the_finest_legal_unit():
+    """832x480 is the checkpoint's own default and the kernel rejects it.
+
+    (480/16) x (832/16) = 1560 tokens per frame, and 1560 % 16 == 8, so a
+    frame-sized block is not something FlashAttention's paged kernel accepts.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    assert 1560 % 16 == 8
+    assert paging_block_size(1560) == 16
+    # And whatever it returns is always something the kernel takes.
+    for tokens_per_frame in range(1, 200):
+        assert paging_block_size(tokens_per_frame) % 16 == 0
+
+
+# ── the kernel's constraint is data, not a constant ─────────────────────────
+
+
+def _multiple_of(base):
+    from vllm.v1.attention.backend import MultipleOf
+
+    return MultipleOf(base)
+
+
+def test_the_dispatching_module_answers_in_vllms_own_vocabulary():
+    """AR-Diffusion calls flash_attn itself, so no backend can be asked.
+
+    The answer has to come from the module that picks the kernel, and it has to
+    have the shape vLLM uses, so the policy that consumes it does not care
+    where it came from.
+    """
+    from vllm.v1.attention.backend import MultipleOf
+
+    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import supported_kernel_block_sizes
+
+    advertised = supported_kernel_block_sizes()
+    assert advertised, "a kernel that accepts nothing cannot be paged for"
+    for entry in advertised:
+        assert isinstance(entry, int | MultipleOf)
+
+
+def test_a_kernel_that_takes_one_fixed_size_gets_that_size():
+    """hpc_attn advertises [64] -- not a multiple, a single legal value.
+
+    A frame of 1560 tokens is not 64, and 1560 % 64 == 24, so the frame cannot
+    be the block and the only thing left to page at is 64. Assuming multiples
+    of 16 here would hand the kernel a block size it rejects.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    assert 1560 % 64 == 24
+    assert paging_block_size(1560, [64]) == 64
+    assert paging_block_size(64, [64]) == 64
+    assert paging_block_size(128, [64]) == 64  # a legal multiple is still not the advertised size
+
+
+def test_a_kernel_with_large_pages_still_keeps_a_whole_frame_when_it_fits():
+    """FlashInfer offers 128-and-up pages on Blackwell and not on Hopper.
+
+    The same checkpoint on the two cards must therefore page differently, which
+    is exactly what a constant cannot express.
+    """
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    hopper = [_multiple_of(16)]
+    blackwell = [_multiple_of(16), 128, 256]
+    assert paging_block_size(1440, hopper) == 1440  # a legal multiple of 16
+    assert paging_block_size(1440, blackwell) == 1440
+    assert paging_block_size(1560, hopper) == 16
+    assert paging_block_size(1560, blackwell) == 16  # finest legal unit, not the largest
+
+
+def test_whatever_comes_back_is_always_something_the_kernel_accepts():
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    for advertised, check in (
+        ([_multiple_of(16)], lambda n: n % 16 == 0),
+        ([64], lambda n: n == 64),
+        ([_multiple_of(32), 48], lambda n: n % 32 == 0 or n == 48),
+    ):
+        for tokens_per_frame in range(1, 300):
+            assert check(paging_block_size(tokens_per_frame, advertised))
+
+
+def test_a_kernel_that_advertises_nothing_is_an_error_not_a_guess():
+    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size
+
+    with pytest.raises(ValueError, match="no legal block size"):
+        paging_block_size(1560, [])

tests/diffusion/ar_diffusion/test_kv_cache.py

@@ -625,6 +625,74 @@ def test_non_contiguous_branch_indices_rejected():
         )
 
 
+# --- eviction against the block table ---------------------------------------
+
+
+class _RecordingManager(ChunkWindowManager):
+    """Capture the block range eviction asks for, without a real block pool.
+
+    ``remove_skipped_blocks`` reads three attributes and calls one method. A
+    pool would add nothing: the bug under test is the conversion from the
+    spec's units into the block table's, and that is entirely visible here.
+    """
+
+    def __init__(self, spec):
+        self.kv_cache_spec = spec
+        self.sliding_window = spec.sliding_window
+        self.block_size = spec.block_size
+        self.freed: tuple[int, int] | None = None
+
+    def _remove_blocks_in_range(self, request_id, first_block, last_block):
+        self.freed = (first_block, last_block)
+
+
+def test_sink_is_preserved_when_a_frame_spans_many_blocks():
+    # The shipped geometry: 832x480 is 1560 tokens per frame, 1560 % 16 == 8 is
+    # not a legal block size, so paging falls back to 16 and a frame spans 98
+    # blocks. Every other eviction test uses chunk_size == BLOCK, where a frame
+    # is a block and a sink counted in frames indexes the block table correctly
+    # by accident -- which is why this went unnoticed.
+    chunk_size, sink_chunks, window_chunks = 1560, 9, 9
+    spec = ChunkWindowSpec(
+        block_size=16,
+        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,
+    )
+    manager = _RecordingManager(spec)
+
+    # Far enough past sink + window that eviction has definitely started.
+    manager.remove_skipped_blocks("req", total_computed_tokens=30 * chunk_size)
+    assert manager.freed is not None
+    first_freed, last_freed = manager.freed
+
+    # The assertion is on the token boundary rather than a block count, so it
+    # stays meaningful if the paging unit changes again.
+    sink_tokens = sink_chunks * chunk_size
+    assert first_freed * spec.block_size >= sink_tokens, (
+        f"eviction starts at token {first_freed * spec.block_size}, inside the {sink_tokens}-token sink"
+    )
+    # And it must not over-protect: at most one straddling block beyond the sink.
+    assert (first_freed - 1) * spec.block_size < sink_tokens
+    # Only wholly-skipped blocks are released.
+    skipped = manager.get_num_skipped_tokens(30 * chunk_size)
+    assert last_freed * spec.block_size <= sink_tokens + skipped
+
+
+def test_sink_conversion_is_identity_when_a_frame_is_a_block():
+    # Every geometry that paged one frame per block must be untouched.
+    spec = make_spec(chunk_size=BLOCK, window_chunks=2, sink_chunks=1)
+    manager = _RecordingManager(spec)
+    manager.remove_skipped_blocks("req", total_computed_tokens=8 * BLOCK)
+    assert manager.freed is not None
+    first_freed, _ = manager.freed
+    assert first_freed == spec.sink_chunks
+
+
 def test_staged_window_reuse_matches_a_full_gather():
     """Reusing the staged history must be byte-identical to re-gathering the whole window.
 

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -23,6 +23,7 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
 )
 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]
@@ -36,10 +37,26 @@ NEG = "negative"
 
 
 def make_state(
-    *, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu"), reuse_history_staging=False
+    *,
+    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=BLOCK, window_chunks=window_chunks, reuse_history_staging=reuse_history_staging
+        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,
@@ -109,11 +126,13 @@ def _commit_video_span(
     n_chunks: int,
     dtype: torch.dtype,
     device: torch.device,
+    chunk_size: int = BLOCK,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    ctx = st.get_kv_caches(kv_branch, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
+    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, 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)
+    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)
@@ -547,6 +566,644 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(monkeypa
         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])

vllm_omni/experimental/ar_diffusion/kv_cache/manager.py

@@ -42,6 +42,7 @@ from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
     chunk_slot_mapping,
     pool_write_chunk,
     resident_block_ids,
+    visible_window_blocks,
 )
 
 _log = init_logger(__name__)
@@ -282,7 +283,12 @@ class ARDiffusionKVCache:
         # forward needs one block per current frame plus space for any
         # model-declared action/state tokens that coexist with video KV.
         declared_scratch_blocks = (max_scratch_tokens_per_branch + block_size - 1) // block_size
-        minimum_scratch_blocks = self.frames_per_block + declared_scratch_blocks
+        # frames_per_block counts *frames*; convert to blocks rather than
+        # assuming one block per frame. A frame spans ceil(chunk_size /
+        # block_size) blocks once the two sizes are allowed to differ.
+        blocks_per_frame = -(-config.chunk_size // block_size)
+        self.blocks_per_frame = blocks_per_frame
+        minimum_scratch_blocks = self.frames_per_block * blocks_per_frame + declared_scratch_blocks
         override = os.environ.get("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH")
         override_blocks = int(override) if override is not None else 0
         if override_blocks < 0:
@@ -313,8 +319,23 @@ class ARDiffusionKVCache:
         )
 
         def _required_managed_blocks(capacity: int) -> int:
-            resident_per_session = config.sink_chunks + config.window_chunks
-            return self.num_local_kv_branches * (capacity * resident_per_session + self.frames_per_block) + 2
+            """Managed blocks needed to hold ``capacity`` resident sessions.
+
+            Every term here counts *frames*, so each is converted to blocks.
+            Before block size and frame size were allowed to differ this
+            conversion was the identity and the frame counts could be used
+            directly; they cannot once a frame spans several blocks.
+
+            Converting rather than rewriting keeps the budget the same number
+            of *tokens* it always was, so a finer block size does not inflate
+            the requirement -- it only stops the final block of each term from
+            rounding up a whole frame's worth of memory.
+            """
+            resident_frames = config.sink_chunks + config.window_chunks
+            per_session_blocks = resident_frames * self.blocks_per_frame
+            in_flight_blocks = self.frames_per_block * self.blocks_per_frame
+            headroom_blocks = 2 * self.blocks_per_frame
+            return self.num_local_kv_branches * (capacity * per_session_blocks + in_flight_blocks) + headroom_blocks
 
         # reuse_history_staging keeps one contiguous K and V buffer per layer,
         # sized to the padded visible window (sink + window plus the
@@ -327,9 +348,15 @@ class ARDiffusionKVCache:
         self.history_staging_tokens = 0
         if config.reuse_history_staging:
             if contiguous_kv_gather_enabled():
-                self.history_staging_tokens = (
-                    self.spec.sliding_window + config.sink_chunks * config.chunk_size + block_size
+                # The same width build_block_table pads to: a sink or window
+                # that is not a whole number of blocks spans one more block.
+                window_blocks = visible_window_blocks(
+                    chunk_size=config.chunk_size,
+                    block_size=block_size,
+                    sink_tokens=config.sink_chunks * config.chunk_size,
+                    window_tokens=self.spec.sliding_window,
                 )
+                self.history_staging_tokens = (window_blocks + 1) * block_size
                 self.history_staging_reserved_bytes = int(
                     2 * num_layers * self.history_staging_tokens * num_kv_heads * head_size * dtype.itemsize
                 )
@@ -649,6 +676,16 @@ class ARDiffusionKVCache:
     def block_table(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
         return list(self.manager.get_block_ids(adapter.request_id)[0])
 
+    def block_ids_at(self, adapter: ARDiffusionRequestAdapter, indices: Sequence[int]) -> list[int]:
+        """Block ids at ``indices`` of the request's block table, without copying it.
+
+        The table keeps a null entry for every evicted position, so it grows with
+        the session. :meth:`block_table` reads every entry; this reads only the
+        ones asked for, straight from the manager's own block list.
+        """
+        blocks = self.manager.get_blocks(adapter.request_id).blocks[0]
+        return [blocks[index].block_id for index in indices]
+
     def chunk_write_slots(self, adapter: ARDiffusionRequestAdapter) -> torch.Tensor:
         """Slot mapping for the in-flight chunk — the K/V write target."""
         return chunk_slot_mapping(

vllm_omni/experimental/ar_diffusion/kv_cache/paged.py

@@ -17,6 +17,7 @@ Cosmos port) reuses unchanged. Three concerns live here:
 
 from __future__ import annotations
 
+import math
 from collections.abc import Sequence
 from dataclasses import dataclass
 
@@ -196,6 +197,19 @@ def chunk_window_skipped_tokens(
     return (skipped // chunk_size) * chunk_size
 
 
+def visible_window_blocks(*, chunk_size: int, block_size: int, sink_tokens: int, window_tokens: int) -> int:
+    """Most blocks a sink plus a recent window can span; see ``max_video_blocks``.
+
+    Each range is converted on its own: the sink ends at a fixed position, and
+    the recent range starts on a chunk boundary, so at most
+    ``block_size - gcd(chunk_size, block_size)`` tokens past a block edge.
+    """
+    worst_start_offset = block_size - math.gcd(chunk_size, block_size)
+    sink_blocks = -(-sink_tokens // block_size)
+    recent_blocks = -(-(window_tokens + worst_start_offset) // block_size)
+    return sink_blocks + recent_blocks
+
+
 class ChunkWindowManager(SlidingWindowManager):
     """``SlidingWindowManager`` that evicts at chunk boundaries.
 
@@ -219,14 +233,26 @@ class ChunkWindowManager(SlidingWindowManager):
         Call only after commit/eviction, then shift the request's storage
         position by the returned amount. Physical pages and model positions
         are unchanged, including any allocated but not yet committed tail.
+
+        The gap starts where :meth:`remove_skipped_blocks` starts freeing: after
+        the sink's blocks, which are ``sink_chunks`` only while a frame is a
+        block. Its removal is also rounded down to a whole number of chunks.
+        Eviction snaps to chunk boundaries and reads its position in storage
+        coordinates, and paged attention reads each block's offset there too;
+        shifting storage by whole chunks *and* whole blocks leaves both exactly
+        as they were in model positions. When a frame is a block, every gap
+        already is one.
         """
         if self.enable_caching:
             raise RuntimeError("AR block-table compaction requires prefix caching to be disabled")
         blocks = self.req_to_blocks.get(request_id, [])
-        start = self.kv_cache_spec.sink_chunks
+        spec = self.kv_cache_spec
+        start = -(-(spec.sink_chunks * spec.chunk_size) // self.block_size)
         end = start
         while end < len(blocks) and blocks[end] == self._null_block:
             end += 1
+        blocks_per_step = spec.chunk_size // math.gcd(spec.chunk_size, self.block_size)
+        end = start + (end - start) // blocks_per_step * blocks_per_step
         if end == start:
             return 0
         del blocks[start:end]
@@ -241,17 +267,32 @@ class ChunkWindowManager(SlidingWindowManager):
         total_computed_tokens: int,
         num_prompt_tokens: int | None = None,
     ) -> None:
-        """Free the middle gap while preserving the leading attention sink."""
+        """Free the middle gap while preserving the leading attention sink.
+
+        ``sink_chunks`` counts frames while the block table is indexed in
+        blocks, so the sink is converted before it is used as an index. The two
+        were interchangeable while a frame was a block; this commit is what
+        separates them.
+
+        Neither boundary is guaranteed to land on a block edge -- 832x480 is
+        1560 tokens per frame against a 16-token block -- so a block can hold
+        the end of one frame and the start of the next. The sink end rounds up
+        so a straddling block stays protected, and the freed range ends on a
+        rounded-down boundary so only wholly-skipped blocks are released. Both
+        are exact when the sizes divide, which is every geometry that paged one
+        frame per block before.
+        """
         del num_prompt_tokens
         num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens)
         if num_skipped_tokens <= 0:
             return
-        sink_blocks = self.kv_cache_spec.sink_chunks
-        num_skipped_blocks = num_skipped_tokens // self.block_size
+        sink_tokens = self.kv_cache_spec.sink_chunks * self.kv_cache_spec.chunk_size
+        first_live_block = -(-sink_tokens // self.block_size)
+        last_freed_block = (sink_tokens + num_skipped_tokens) // self.block_size
         self._remove_blocks_in_range(
             request_id,
-            sink_blocks,
-            sink_blocks + num_skipped_blocks,
+            first_live_block,
+            last_freed_block,
         )
 
 

vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py

@@ -6,12 +6,15 @@ from __future__ import annotations
 
 import weakref
 from dataclasses import dataclass, field
-from typing import Any, ClassVar, NamedTuple
+from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple
 
 import torch
 
 from vllm_omni.experimental.ar_diffusion.kv_cache.config import contiguous_kv_gather_enabled
-from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping, visible_window_blocks
+
+if TYPE_CHECKING:
+    from vllm.v1.attention.backend import MultipleOf
 
 _LAYER_IDX_TENSORS: dict[int, torch.Tensor] = {}
 
@@ -91,6 +94,21 @@ class ARDiffusionPagedForwardContext:
     action_slot_mapping: torch.Tensor | None = None
     query_len: int = 0
     kv_len: int = 0
+    # Tokens already committed when this forward began. Recorded rather than
+    # derived from the block count: the last block of the history is only
+    # partly written when a chunk is not a whole number of blocks, so
+    # blocks * block_size overstates it and would make attention read past
+    # what was written.
+    _history_tokens: int = 0
+    # Scratch blocks the current video tokens occupy. Not the same as
+    # len(current_video_block_ids): when the history ends mid-block, the first
+    # entry there is the history's own tail block, which is managed, not
+    # scratch. Action K/V is placed after this, so it must not count it.
+    _scratch_blocks_used: int = 0
+    # Slot offset of this chunk's first token inside its first block. The
+    # chunk's blocks then hold positions ``history - offset + j * block_size``,
+    # which is how the read path places them among the history's blocks.
+    _current_offset: int = 0
     _allocated_video: bool = False
     _committed: bool = False
     _action_len: int = 0
@@ -106,14 +124,77 @@ class ARDiffusionPagedForwardContext:
     def block_size(self) -> int:
         return int(self.kv_cache.block_size)
 
+    @property
+    def chunk_size(self) -> int:
+        """Tokens in one chunk -- one latent frame. The eviction unit."""
+        return int(self.kv_cache.spec.chunk_size)
+
+    @property
+    def start_offset(self) -> int:
+        """Where in its first block this forward's tokens begin.
+
+        Non-zero exactly when the committed history does not end on a block
+        boundary, which is the normal case once a frame is not a whole number
+        of blocks: 1560 tokens per frame against 16-token blocks leaves the
+        history 8 slots into its last block on every odd-numbered chunk.
+        """
+        return int(self.adapter.num_computed_tokens) % self.block_size
+
+    @property
+    def sink_tokens(self) -> int:
+        """Leading tokens the window keeps however far it has slid."""
+        return int(self.kv_cache.spec.sink_chunks) * self.chunk_size
+
+    @property
+    def window_tokens(self) -> int:
+        """Most recent tokens the window keeps, this forward's own included."""
+        return self.max_video_tokens - self.sink_tokens
+
+    @property
+    def max_video_blocks(self) -> int:
+        """Most blocks the visible window can span, which fixes the table width.
+
+        The window is two token ranges, the sink and the most recent tokens, and
+        each is converted to blocks on its own because each can straddle a block
+        edge on its own. Rounding the two ranges' *sum* up once was a block short
+        whenever both straddled: at 832x480 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 it was short on every tick once the window had
+        filled.
+
+        The sink ends at a fixed position, so its blocks are just its tokens
+        rounded up. The recent range starts at a whole number of chunks, so its
+        start sits at most ``block_size - gcd(chunk_size, block_size)`` tokens
+        past a block edge, and it can need that much more room. A chunk that is
+        a whole number of blocks never straddles, so every geometry that paged
+        one frame per block keeps exactly the width it had.
+
+        The width is a constant of the configuration, which is what keeps the
+        block table a fixed shape as the window fills.
+        """
+        return visible_window_blocks(
+            chunk_size=self.chunk_size,
+            block_size=self.block_size,
+            sink_tokens=self.sink_tokens,
+            window_tokens=self.window_tokens,
+        )
+
     @property
     def num_current_video_blocks(self) -> int:
-        if self.seq_len % self.block_size != 0:
-            raise AssertionError(
-                "AR-Diffusion paged attention expects frame-aligned seq_len "
-                f"(multiple of block_size={self.block_size}), got {self.seq_len}"
-            )
-        return self.seq_len // self.block_size
+        """Blocks this forward's tokens occupy, counted from where they start.
+
+        A chunk need not be a whole number of blocks, so the count depends on
+        the offset it begins at, not only on its length: thirty tokens
+        starting at position thirty span three sixteen-token blocks, not two.
+        Rounding up leaves the tail of the last block unwritten, which is only
+        safe because nothing reads past ``kv_len`` -- see
+        :meth:`video_block_table`.
+
+        Both paths count the same way. The committing path writes straight
+        after the history; the scratch path is made to line up with it by
+        :meth:`ensure_video_slots`.
+        """
+        return -(-(self.start_offset + self.seq_len) // self.block_size)
 
     def ensure_video_slots(self, device: torch.device) -> None:
         """Allocate/write targets for the current video tokens, once per KV branch."""
@@ -121,8 +202,10 @@ class ARDiffusionPagedForwardContext:
             return
 
         n_blocks = self.num_current_video_blocks
+        self._history_tokens = int(self.adapter.num_computed_tokens)
         if self.commit_current:
             start = int(self.adapter.num_computed_tokens)
+            self._current_offset = start % self.block_size
             self.kv_cache.allocate_token_slots(self.adapter, self.seq_len)
             table = self.kv_cache.block_table(self.adapter)
             start_block = start // self.block_size
@@ -132,8 +215,27 @@ class ARDiffusionPagedForwardContext:
                 compute_slot_mapping(table, positions, self.block_size), device
             )
         else:
-            self.current_video_block_ids = self.kv_cache.scratch_block_ids(self.kv_branch, 0, n_blocks)
-            positions = torch.arange(self.seq_len, dtype=torch.long)
+            # The scratch region cannot simply start at slot zero. The kernel
+            # reads the history and this chunk as one contiguous run, so if the
+            # history stopped mid-block, starting here at zero would leave the
+            # rest of that block unwritten and the kernel would read those dead
+            # slots as if they were tokens, shifting the whole sequence.
+            #
+            # Instead the chunk begins where the history left off, spilling its
+            # first few tokens into the free tail of the history's own last
+            # block. That block is already this session's -- blocks are
+            # allocated whole -- and those slots hold nothing yet; the
+            # committing forward overwrites the same slots with the clean K/V
+            # later, so the committed state is unchanged either way.
+            # A history whose tail block is no longer resident has nothing to
+            # spill into; the window then starts at this chunk and zero is right.
+            offset = self.start_offset if self.history_block_ids else 0
+            self._current_offset = offset
+            self._scratch_blocks_used = n_blocks - (1 if offset else 0)
+            scratch_ids = self.kv_cache.scratch_block_ids(self.kv_branch, 0, self._scratch_blocks_used)
+            tail_block = [self.history_block_ids[-1]] if offset else []
+            self.current_video_block_ids = tail_block + scratch_ids
+            positions = torch.arange(offset, offset + self.seq_len, dtype=torch.long)
             self.current_video_slot_mapping = _to_device_async(
                 compute_slot_mapping(self.current_video_block_ids, positions, self.block_size), device
             )
@@ -152,7 +254,9 @@ class ARDiffusionPagedForwardContext:
             return
 
         action_blocks = (action_len + self.block_size - 1) // self.block_size
-        scratch_offset = 0 if self.commit_current else len(self.current_video_block_ids)
+        # Count scratch blocks, not entries: the video list may lead with the
+        # history's tail block, which lives in the managed pool.
+        scratch_offset = 0 if self.commit_current else self._scratch_blocks_used
         self.action_scratch_block_ids = self.kv_cache.scratch_block_ids(
             self.kv_branch,
             scratch_offset,
@@ -165,23 +269,75 @@ class ARDiffusionPagedForwardContext:
         self._action_len = action_len
 
     def video_block_table(self, device: torch.device) -> tuple[list[int], int]:
+        """Blocks the attention reads, and how many of their slots it reads.
+
+        The kernel reads the listed blocks as one contiguous run and stops after
+        the returned length, so two things must hold: every block holding a
+        token the window keeps is listed, in token order, and only the last
+        listed block may be partly written.
+
+        Blocks are chosen by the token positions they hold rather than counted
+        off the ends of the resident list. The window is the sink ``[0, S)``
+        plus the most recent tokens, and each of those can straddle a block
+        edge on its own. Taking a fixed number of blocks off the tail dropped
+        the block holding the window's first tokens exactly when both
+        straddled, and a length capped by a block count rather than counted
+        from what was written then ran past the last written slot.
+
+        Whole blocks are read, so each boundary block contributes up to
+        ``block_size - 1`` tokens the window has already let go of. They are
+        real K/V at their own positions; a slightly wider window is a far
+        smaller error than a truncated one.
+        """
         self.ensure_video_slots(device)
-        if self.max_video_tokens % self.block_size != 0:
-            raise AssertionError(
-                "AR-Diffusion paged attention requires max_video_tokens to be block-aligned, "
-                f"got max_video_tokens={self.max_video_tokens}, block_size={self.block_size}"
-            )
-        all_video_blocks = self.history_block_ids + self.current_video_block_ids
-        max_video_blocks = self.max_video_tokens // self.block_size
-        sink_blocks = int(self.kv_cache.spec.sink_chunks)
-        if len(all_video_blocks) <= max_video_blocks:
-            visible_video_blocks = all_video_blocks
-        else:
-            tail_blocks = max_video_blocks - sink_blocks
-            visible_video_blocks = all_video_blocks[:sink_blocks]
-            if tail_blocks:
-                visible_video_blocks += all_video_blocks[-tail_blocks:]
-        video_len = len(visible_video_blocks) * self.block_size
+        block_size = self.block_size
+        history = self._history_tokens
+        end = history + self.seq_len
+        sink_end = self.sink_tokens
+        recent_start = end - self.window_tokens
+
+        # Position of each block's first slot. History comes from the block
+        # table, where an evicted block is a null placeholder, so a block's index
+        # is its position -- in storage positions, which is what
+        # num_computed_tokens counts once compaction has dropped an evicted gap.
+        # Compaction shifts storage by whole chunks and whole blocks, so the sink,
+        # the window and every block offset sit exactly where they would in model
+        # positions. Reading the table after ensure_video_slots means a block
+        # this chunk's own allocation evicted is never listed. The current
+        # chunk's blocks follow from where its first token was placed; its first
+        # block is the history's own tail block whenever the history stopped
+        # part way through one, and keeps the history's position.
+        first_position: dict[int, int] = {}
+        #
+        # Only two index ranges can hold a token the window keeps: the sink's
+        # blocks, and the blocks from the recent window's start onward. Evicted
+        # positions stay in the table as null entries until compaction drops
+        # them a whole number of chunks at a time, so the table can still be
+        # much longer than the window; reading just these two ranges, without
+        # copying the table, keeps this at the window's size.
+        history_blocks = -(-history // block_size)
+        sink_blocks = min(-(-sink_end // block_size), history_blocks)
+        recent_first_block = max(max(recent_start, 0) // block_size, sink_blocks)
+        indices = [*range(sink_blocks), *range(recent_first_block, history_blocks)]
+        for index, block in zip(indices, self.kv_cache.block_ids_at(self.adapter, indices)):
+            if block != self.kv_cache.null_block_id:
+                first_position.setdefault(block, index * block_size)
+        current_start = history - self._current_offset
+        for index, block in enumerate(self.current_video_block_ids):
+            first_position.setdefault(int(block), current_start + index * block_size)
+
+        visible_video_blocks: list[int] = []
+        video_len = 0
+        for block, first in sorted(first_position.items(), key=lambda item: item[1]):
+            if first >= sink_end and first + block_size <= recent_start:
+                continue  # wholly between the sink and the recent window
+            if video_len % block_size:
+                raise RuntimeError(
+                    "AR-Diffusion paged attention would read past a partly written block: "
+                    f"block {visible_video_blocks[-1]} is followed by block {block}."
+                )
+            visible_video_blocks.append(block)
+            video_len += min(block_size, end - first)
         return visible_video_blocks, video_len
 
     def build_block_table(
@@ -200,6 +356,21 @@ class ARDiffusionPagedForwardContext:
         ``ceil(seq_lens/block_size)`` entries, so padding is never read.
         """
         video_blocks, video_len = self.video_block_table(device)
+        # Action K/V is listed after the video blocks and the kernel reads the
+        # table as one run, so the video run has to end on a block edge. If it
+        # stops part way through its last block, the run reads that block's
+        # unwritten slots as the first action tokens and never reaches the last
+        # ones -- a wrong attention with every shape intact. A frame that is a
+        # whole number of blocks always ends on an edge, so this refuses only
+        # geometries that could not page at all before the paging unit was
+        # separated from the frame.
+        if action_len > 0 and video_len % self.block_size:
+            raise ValueError(
+                "AR-Diffusion paged attention cannot place action tokens after a partly written video block: "
+                f"the video run ends {video_len % self.block_size} slots into block {video_blocks[-1]}, so the "
+                "kernel would read its unwritten slots as action tokens. Action tokens need tokens_per_frame "
+                f"to be a whole number of blocks (block_size={self.block_size})."
+            )
         self.ensure_action_slots(action_len, device)
         action_blocks = self.action_scratch_block_ids if action_len > 0 else []
         block_ids = video_blocks + action_blocks
@@ -208,12 +379,19 @@ class ARDiffusionPagedForwardContext:
 
         # Fixed capacity: full visible video window + one action-capacity block.
         action_capacity_blocks = max(1, (action_len + self.block_size - 1) // self.block_size)
-        width = max(self.max_video_tokens // self.block_size + action_capacity_blocks, len(block_ids))
+        # Both of these count in blocks. Deriving them from the token count by
+        # flooring would understate the capacity whenever the window is not a
+        # whole number of blocks: the width would fall back on len(block_ids)
+        # and start tracking how full the window is -- the exact shape churn
+        # the fixed width exists to prevent -- and max_seq_len could come out
+        # below the kv_len actually being passed, by up to block_size - 1.
+        capacity_blocks = self.max_video_blocks + action_capacity_blocks
+        width = max(capacity_blocks, len(block_ids))
         padded = block_ids + [0] * (width - len(block_ids))
 
         self.query_len = int(query_len)
         self.kv_len = int(video_len + action_len)
-        max_seq_len = int(self.max_video_tokens + action_capacity_blocks * self.block_size)
+        max_seq_len = int(capacity_blocks * self.block_size)
         # Built on the host once per AR block; the copies go through the same
         # pinned, non-blocking path as the slot mappings above so the CPU does
         # not wait for the stream to drain three times before the first layer.
@@ -281,7 +459,11 @@ class ARDiffusionPagedForwardContext:
             int(self.seq_len),
         )
         same_session = state.adapter is not None and state.adapter() is self.adapter
-        reuse = same_session and state.signature == signature
+        # A chunk that is not a whole number of blocks shares its first block with the history's tail (see
+        # ensure_video_slots), and the attention op restages the current chunk in whole blocks of its own,
+        # so such a forward restages the whole window instead.
+        block_aligned = self.start_offset == 0 and self.seq_len % self.block_size == 0
+        reuse = same_session and state.signature == signature and block_aligned
         state.adapter = weakref.ref(self.adapter)
         state.signature = signature
         self.staging_enabled = True
@@ -461,6 +643,36 @@ def _resolve_fa_version(head_size: int) -> int:
     return version
 
 
+def supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    """Block sizes the kernel this module dispatches to will accept.
+
+    Same shape as vLLM's ``AttentionBackend.get_supported_kernel_block_sizes``
+    -- a plain int is that exact size, ``MultipleOf(b)`` is any positive
+    multiple of ``b`` -- but answered here rather than read off a backend,
+    because AR-Diffusion does not go through backend selection. It calls
+    ``flash_attn_varlen_func`` itself, choosing between vLLM's CUDA build and
+    ROCm's AITER a few lines below, and only this module knows which.
+
+    It lives next to that choice so there is one place to change. The
+    constraint is a property of the kernel, not of the card, and the kernels
+    reachable from here agree on 16 today: vLLM's CUDA FlashAttention, ROCm
+    AITER, and upstream ``flash_attn`` all advertise ``MultipleOf(16)``. Other
+    backends in the same tree do not -- ``hpc_attn`` accepts only 64, and
+    FlashInfer advertises pages of 128 or more solely on Blackwell -- so a
+    caller must treat this as data to be queried, never as the number 16.
+
+    vLLM's FlashAttention backend does advertise a single 128-token page when FA4
+    runs its dedicated head-size-256 kernel on SM100/SM110. That kernel is not
+    reachable from here: ``get_flash_attn_version`` selects it only when the
+    caller passes ``supports_fa4_hd256=True``, and :func:`_resolve_fa_version`
+    does not, so a head size of 256 falls back to FA2, which pages at any
+    multiple of 16. Passing that flag later has to change this answer as well.
+    """
+    from vllm.v1.attention.backend import MultipleOf
+
+    return [MultipleOf(16)]
+
+
 def _rocm_flash_attn_varlen_func():
     """Resolve a ROCm varlen kernel, preferring AITER when available.
 

vllm_omni/experimental/ar_diffusion/kv_cache/state.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Model-facing session state backed by paged AR-Diffusion KV storage."""
 
 from __future__ import annotations

vllm_omni/experimental/ar_diffusion/runner.py

@@ -5,8 +5,9 @@ from __future__ import annotations
 import dataclasses
 import time
 from collections import OrderedDict
-from collections.abc import Iterator
+from collections.abc import Iterator, Sequence
 from contextlib import contextmanager
+from typing import TYPE_CHECKING
 
 import torch
 from vllm.logger import init_logger
@@ -29,9 +30,61 @@ from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVStat
 from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest
 from vllm_omni.platforms import current_omni_platform
 
+if TYPE_CHECKING:
+    from vllm.v1.attention.backend import MultipleOf
+
 logger = init_logger(__name__)
 
 
+def paging_block_size(tokens_per_frame: int, supported: Sequence[int | MultipleOf] | None = None) -> int:
+    """Choose the paging unit, which is not the eviction unit.
+
+    AR-Diffusion evicts by frame; the attention kernel tiles by a fixed number
+    of tokens. Using the frame as the block binds every output resolution to
+    that kernel constraint, so the resolutions that run today do so by
+    arithmetic accident -- and the LingBot World v2 checkpoint's own default
+    832x480 is not one of them:
+
+        (480/16) x (832/16) = 30 x 52 = 1560 tokens per frame, 1560 % 16 = 8
+
+    A frame that is already a legal block stays one, so every resolution that
+    works today keeps precisely the paging it has now. Anything else pages at
+    the finest legal unit, which costs a longer block table and gives back the
+    rounding waste of a block one whole frame wide -- at 832x480 that block is
+    over a gigabyte.
+
+    What counts as legal is asked, not assumed. ``supported`` carries vLLM's
+    own vocabulary: a plain int is that exact size, ``MultipleOf(b)`` is any
+    positive multiple of ``b``. The kernels AR-Diffusion reaches agree on
+    multiples of 16, but others in the same tree do not -- ``hpc_attn`` takes
+    only 64, and FlashInfer offers pages of 128 and up on Blackwell alone. A
+    constant here would be a silently wrong block size on the first machine
+    whose kernel differs, and a wrong block size is not an error: it is a
+    kernel reading tokens that were never written.
+    """
+    if supported is None:
+        # Imported here, not at module scope: paged_attention registers a
+        # custom op on import, and the tests that pop it from sys.modules to
+        # re-import it depend on nothing else holding it first.
+        from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import supported_kernel_block_sizes
+
+        supported = supported_kernel_block_sizes()
+    if not supported:
+        raise ValueError("The attention kernel advertises no legal block size.")
+
+    def legal(size: int) -> bool:
+        return size > 0 and any(
+            size == entry if isinstance(entry, int) else size % entry.base == 0 for entry in supported
+        )
+
+    if legal(tokens_per_frame):
+        return tokens_per_frame
+    finest = min(entry if isinstance(entry, int) else entry.base for entry in supported)
+    if not legal(finest):
+        raise ValueError(f"No legal paging unit among {supported!r}.")
+    return finest
+
+
 def resolve_ar_diffusion_kv_config(od_config: OmniDiffusionConfig) -> ARDiffusionKVConfig:
     """Resolve optional deployment overrides and enable AR-Diffusion KV."""
     raw = getattr(od_config, "ar_diffusion_kv_config", None)
@@ -139,7 +192,7 @@ class ARDiffusionModelRunner(DiffusionModelRunner):
             num_kv_heads=spec.num_kv_heads,
             head_size=spec.head_size,
             dtype=self.od_config.dtype,
-            block_size=spec.tokens_per_frame,
+            block_size=paging_block_size(spec.tokens_per_frame),
             max_model_len=spec.max_model_len,
             available_bytes=self._available_memory_bytes() if available_bytes is None else available_bytes,
             kv_branches=spec.kv_branches,