diff --git a/tests/diffusion/ar_diffusion/test_capability_runner.py b/tests/diffusion/ar_diffusion/test_capability_runner.py index d753efbf6..01d8b2844 100644 --- a/tests/diffusion/ar_diffusion/test_capability_runner.py +++ b/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, []) diff --git a/tests/diffusion/ar_diffusion/test_kv_cache.py b/tests/diffusion/ar_diffusion/test_kv_cache.py index 84908fef1..e73713f8e 100644 --- a/tests/diffusion/ar_diffusion/test_kv_cache.py +++ b/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. diff --git a/tests/diffusion/ar_diffusion/test_paged_attention.py b/tests/diffusion/ar_diffusion/test_paged_attention.py index 7014a60c7..47e9b028e 100644 --- a/tests/diffusion/ar_diffusion/test_paged_attention.py +++ b/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]) diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py index efa5cc2c1..b178981ff 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py +++ b/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( diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py index 0ef639c98..669fe7a2f 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py +++ b/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, ) diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py index 2e7f72522..0ea57e39f 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py +++ b/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. diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py index ec14a9177..670046d0b 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py +++ b/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 diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py index 858cf843d..4180ae587 100644 --- a/vllm_omni/experimental/ar_diffusion/runner.py +++ b/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,