diff --git a/tests/diffusion/ar_diffusion/test_capability_runner.py b/tests/diffusion/ar_diffusion/test_capability_runner.py index 70166189b..d753efbf6 100644 --- a/tests/diffusion/ar_diffusion/test_capability_runner.py +++ b/tests/diffusion/ar_diffusion/test_capability_runner.py @@ -323,12 +323,18 @@ def test_lingbot_like_sink_survives_sliding_window_eviction(): kv = runner.kv_cache assert kv is not None - for _ in range(8): + state = commit_one_frame(runner, "s1", "main") + sink_block = kv.block_table(state.adapter("main"))[0] + for _ in range(7): state = commit_one_frame(runner, "s1", "main") - table = kv.block_table(state.adapter("main")) - assert table[0] != kv.null_block_id - assert table[1] == kv.null_block_id + adapter = state.adapter("main") + table = kv.block_table(adapter) + assert len(table) == 6 # one sink + five recent frames + assert table[0] == sink_block + assert kv.null_block_id not in table + assert adapter.absolute_num_computed_tokens == 8 * BLOCK + assert adapter.num_computed_tokens == 6 * BLOCK ctx = state.get_kv_caches("main", seq_len=BLOCK, commit_current=False)[0].forward_ctx visible, _ = ctx.video_block_table(torch.device("cpu")) diff --git a/tests/diffusion/ar_diffusion/test_kv_orchestration.py b/tests/diffusion/ar_diffusion/test_kv_orchestration.py index 4482fa6fb..0c1c870bf 100644 --- a/tests/diffusion/ar_diffusion/test_kv_orchestration.py +++ b/tests/diffusion/ar_diffusion/test_kv_orchestration.py @@ -6,6 +6,7 @@ import torch from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig +from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] DIMS = dict( @@ -108,3 +109,104 @@ def test_state_close_frees_both_branch_blocks(): state.close() # All blocks returned — no leak across an evicted session. assert kv.manager.block_pool.get_num_free_blocks() == free_total + + +@pytest.mark.parametrize( + "frames_per_call,block_size,max_model_len,num_calls", + [(1, 16, 113, 1000), (3, 16, 113, 1000), (3, 1560, 1 << 20, 241)], + ids=["single-frame", "three-frame", "lingbot-two-minutes"], +) +@pytest.mark.parametrize("sink_chunks,reset_at_boundary", [(0, False), (2, False), (0, True), (2, True)]) +def test_rollout_keeps_metadata_bounded( + frames_per_call, block_size, max_model_len, num_calls, sink_chunks, reset_at_boundary +): + """Long rollouts retain exact global-position K/V with bounded storage indices.""" + kv = ARDiffusionKVCache( + ARDiffusionKVConfig( + enable=True, + chunk_size=block_size, + window_chunks=3, + sink_chunks=sink_chunks, + reset_at_boundary=reset_at_boundary, + ), + num_layers=1, + num_kv_heads=1, + head_size=4, + dtype=torch.float32, + block_size=block_size, + max_model_len=max_model_len, + available_bytes=1 << 24, + kv_branches=(ARDiffusionKVBranchSpec("main", 0),), + session_capacity=1, + frames_per_block=frames_per_call, + device=torch.device("cpu"), + ) + adapter = kv.begin_request("long-rollout") + state = ARDiffusionKVState(kv, "long-rollout", {"main": adapter}, num_layers=1) + free_total = kv.manager.block_pool.get_num_free_blocks() + pool_shape = kv.key_cache(0).shape + retained_frames = sink_chunks + (0 if reset_at_boundary else 3) + for call_index in range(num_calls): + start = call_index * frames_per_call * block_size + stop = start + frames_per_call * block_size + if frames_per_call == 1: + kv.allocate_chunk(adapter) + slots = kv.chunk_write_slots(adapter) + else: + context = state.get_kv_caches("main", seq_len=stop - start, commit_current=True)[0].forward_ctx + context.ensure_video_slots(torch.device("cpu")) + slots = context.current_video_slot_mapping + assert len(kv.block_table(adapter)) <= retained_frames + frames_per_call + values = torch.arange(start, stop, dtype=torch.float32).view(-1, 1, 1).expand(-1, 1, 4) + kv.key_cache(0).flatten(0, 1)[slots] = values + kv.value_cache(0).flatten(0, 1)[slots] = -values + if frames_per_call == 1: + kv.commit_chunk(adapter) + else: + state.commit_paged_context("main") + assert state._committed["main"] == stop + sink_end = min(stop, sink_chunks * block_size) + tail_start = stop if reset_at_boundary else max(sink_end, stop - 3 * block_size) + expected = torch.cat((torch.arange(sink_end), torch.arange(tail_start, stop))).float() + expected = expected.view(-1, 1, 1).expand(-1, 1, 4) + resident = kv.window_block_ids(adapter) + torch.testing.assert_close(kv.key_cache(0)[resident].flatten(0, 1), expected) + torch.testing.assert_close(kv.value_cache(0)[resident].flatten(0, 1), -expected) + assert adapter.completed_chunks * block_size == stop + assert adapter.absolute_num_computed_tokens == stop + assert len(kv.block_table(adapter)) <= retained_frames + assert adapter.num_computed_tokens <= retained_frames * block_size + assert kv.key_cache(0).shape == pool_shape + assert free_total - kv.manager.block_pool.get_num_free_blocks() == min(stop // block_size, retained_frames) + state.reset() + fresh = state.adapter("main") + assert fresh.compacted_tokens == fresh.absolute_num_computed_tokens == fresh.num_computed_tokens == 0 + assert kv.block_table(fresh) == [] + assert kv.manager.block_pool.get_num_free_blocks() == free_total + kv.allocate_chunk(fresh) + assert len(kv.chunk_write_slots(fresh)) == block_size + state.close() + assert kv.manager.block_pool.get_num_free_blocks() == free_total + + +def test_compaction_offsets_are_request_local(): + kv = make_cache() + free_total = kv.manager.block_pool.get_num_free_blocks() + first, second = kv.begin_request("first"), kv.begin_request("second") + for index in range(8): + kv.allocate_chunk(first) + kv.commit_chunk(first) + if index in (0, 4): + kv.allocate_chunk(second) + kv.commit_chunk(second) + assert first.absolute_num_computed_tokens == 8 * 16 + assert first.compacted_tokens == 6 * 16 + assert second.absolute_num_computed_tokens == 2 * 16 + assert second.compacted_tokens == 0 + assert first.num_computed_tokens == second.num_computed_tokens == 2 * 16 + assert set(kv.block_table(first)).isdisjoint(kv.block_table(second)) + kv.end_request(first) + kv.end_request(second) + group = kv.manager.coordinator.single_type_managers[0] + assert not group.req_to_blocks and not group.num_cached_block + assert kv.manager.block_pool.get_num_free_blocks() == free_total diff --git a/vllm_omni/experimental/ar_diffusion/capability.py b/vllm_omni/experimental/ar_diffusion/capability.py index 145261513..484562227 100644 --- a/vllm_omni/experimental/ar_diffusion/capability.py +++ b/vllm_omni/experimental/ar_diffusion/capability.py @@ -81,6 +81,7 @@ class ARDiffusionKVCacheSpec: sink_frames: int = 0 reset_at_boundary: bool = False cross_attention: tuple[ARDiffusionCrossAttentionKVSpec, ...] = () + # Initial vLLM storage horizon; extended to fit an in-flight span. max_model_len: int = 1 << 20 max_scratch_tokens_per_branch: int = 0 model_owned_state_bytes_per_session: int = 0 diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py index f4cb416f6..48ede6473 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py @@ -15,6 +15,7 @@ from __future__ import annotations import inspect import os from collections.abc import Collection, Iterable, Sequence +from typing import cast import torch from vllm.logger import init_logger @@ -30,6 +31,7 @@ from vllm_omni.diffusion.diffusion_kv.layout import build_kv_cache_tensor from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ( + ChunkWindowManager, ChunkWindowSpec, allocate_kv_pool_with_views, chunk_slot_mapping, @@ -49,10 +51,10 @@ class ARDiffusionRequestAdapter: real ``KVCacheManager`` against this adapter so the surface cannot silently drift across vLLM versions. - An AR-Diffusion request advances one *chunk* at a time: ``allocate_slots`` is called - once per chunk and ``num_computed_tokens`` advances only when a chunk is - committed (:meth:`on_chunk_committed`), so the ``T`` denoise steps of a chunk - reuse the same slots. + ``completed_chunks`` and ``absolute_num_computed_tokens`` track the model + timeline. vLLM sees compacted storage positions through ``num_computed_tokens`` + and ``num_tokens``. A chunk's denoise steps reuse the same slots; only a + successful commit advances the timeline and compacts evicted metadata. """ def __init__( @@ -66,6 +68,7 @@ class ARDiffusionRequestAdapter: self._chunk_size = chunk_size self._prefill = prefill_prefix_tokens self._completed_chunks = 0 + self.compacted_tokens = 0 # Filled only when cross-request prefix reuse is enabled (Phase 3). self.block_hashes: list = [] self.skip_reading_prefix_cache = True @@ -79,14 +82,19 @@ class ARDiffusionRequestAdapter: self.num_in_flight_tokens = 0 @property - def num_computed_tokens(self) -> int: - """Persistent KV already materialized (committed chunks + prefill).""" + def absolute_num_computed_tokens(self) -> int: + """Cumulative model position, unaffected by storage compaction.""" return self._prefill + self._completed_chunks * self._chunk_size + @property + def num_computed_tokens(self) -> int: + """Committed position in the compact vLLM block table.""" + return self.absolute_num_computed_tokens - self.compacted_tokens + @property def num_tokens(self) -> int: - """Total tokens once the in-flight chunk is committed.""" - return self._prefill + (self._completed_chunks + 1) * self._chunk_size + """Storage position once the in-flight chunk is committed.""" + return self.num_computed_tokens + self._chunk_size @property def num_prompt_tokens(self) -> int: @@ -540,12 +548,9 @@ class ARDiffusionKVCache: def allocate_chunk(self, adapter: ARDiffusionRequestAdapter) -> list[int]: """Allocate a chunk's blocks (evicting out-of-window blocks first). - Returns the request's full block table (incl. null_block placeholders). + Returns the request's compact storage block table. """ - blocks = self.manager.allocate_slots(adapter, num_new_tokens=self.spec.chunk_size) - if blocks is None: - raise RuntimeError("AR-Diffusion KV pool exhausted while allocating a chunk") - table = self.block_table(adapter) + table = self.allocate_token_slots(adapter, self.spec.chunk_size) resident = resident_block_ids(table, self.null_block_id) _log.debug( "AR-Diffusion allocate_chunk: req=%s chunk=%d table_len=%d resident=%d free=%d", @@ -561,6 +566,10 @@ class ARDiffusionKVCache: """Allocate managed blocks for an in-flight video span without committing it.""" if num_tokens <= 0: raise ValueError(f"num_tokens must be positive, got {num_tokens}") + # vLLM clips slot allocation at max_model_len. Admit the entire + # in-flight span in compact storage coordinates, independent of the + # model's absolute frame positions. + self.manager.max_model_len = max(self.manager.max_model_len, adapter.num_computed_tokens + num_tokens) blocks = self.manager.allocate_slots(adapter, num_new_tokens=num_tokens) if blocks is None: raise RuntimeError("AR-Diffusion KV pool exhausted while allocating paged attention slots") @@ -618,6 +627,8 @@ class ARDiffusionKVCache: adapter.num_computed_tokens, num_prompt_tokens=adapter.num_prompt_tokens, ) + window_manager = cast(ChunkWindowManager, self.manager.coordinator.single_type_managers[0]) + adapter.compacted_tokens += window_manager.compact_block_table(adapter.request_id) _log.debug("AR-Diffusion commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks) # -- pool-backed K/V access -------------------------------------------- diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py index ddc1937a2..325bbaf16 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py @@ -197,6 +197,28 @@ class ChunkWindowManager(SlidingWindowManager): reset_at_boundary=spec.reset_at_boundary, ) + def compact_block_table(self, request_id: str) -> int: + """Remove the evicted gap after the sink; return its size in tokens. + + 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. + """ + 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 + end = start + while end < len(blocks) and blocks[end] == self._null_block: + end += 1 + if end == start: + return 0 + del blocks[start:end] + if request_id in self.num_cached_block: + cached = self.num_cached_block[request_id] + self.num_cached_block[request_id] = min(cached, start) + max(0, cached - end) + return (end - start) * self.block_size + def remove_skipped_blocks( self, request_id: str,