diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index d5536a1a4f..1760ca14d0 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -3641,6 +3641,104 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(): assert [len(blocks) for blocks in computed_blocks.blocks] == [2, 8] +def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt(): + """An identical resend and a longer sibling resume at DIFFERENT positions. + + How far a lookup matches depends on who is asking. A resend of the same + prompt caps its lookup at ``num_tokens - 1`` (the last token is recomputed + for logits), while a sibling whose prompt merely starts with this one caps + above the prompt. The two coincide unless the prompt length is an exact + multiple of the alignment -- there they differ by one alignment unit, and + under the EAGLE drop BOTH are reachable. + + Retaining only the higher one leaves the resend with every retained state + above every candidate its lookup can produce, and the reconciled hit + collapses to 0 -- the same zero-hit failure sparse retention already fixes + at unaligned prompt lengths. + """ + block_size = 32 + num_spec = 3 + kv_cache_config = KVCacheConfig( + num_blocks=100, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float16, + ), + ), + KVCacheGroupSpec( + ["mamba_mtp"], + MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + num_speculative_blocks=num_spec, + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=block_size, + retention_interval=0, + use_eagle=True, + ) + + # 128 tokens, an exact multiple of the 32-token alignment. A longer sibling + # matches 128 and drops to 96; this prompt's own resend caps at 127, matches + # 96 and drops to 64. Both states must survive retention. + token_ids = [i for i in range(4) for _ in range(block_size)] + req0 = make_request("0", token_ids, block_size, sha256) + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0) + assert num_computed_tokens == 0 + # Prefill in block-aligned chunks the way the align-mode scheduler does: a + # state only materializes as a chunk's running-state block, so a + # single-shot prefill could not retain the lower one. + for chunk_end in (32, 64, 96, 128): + blocks = manager.allocate_slots( + req0, + chunk_end - req0.num_computed_tokens, + num_computed_tokens, + computed_blocks, + num_lookahead_tokens=num_spec, + ) + assert blocks is not None + req0.num_computed_tokens = chunk_end + + # Block ``i`` ends at token ``(i + 1) * 32``, so positions 64 and 96 are + # mamba blocks 1 and 2. + pool = manager.block_pool + expected_mamba_cached = {1, 2} + for i in range(4): + cached = pool.get_cached_block(req0.block_hashes[i], kv_cache_group_ids=[1]) + if i in expected_mamba_cached: + assert cached is not None, f"mamba hash {i} should be cached" + else: + assert cached is None, f"mamba hash {i} should not be cached" + manager.free(req0) + + # The identical resend: full attention matches blocks 0-2 (96 tokens, capped + # by num_tokens - 1) and the EAGLE drop caps the candidate at 64. Without + # the lower state retained the reconciled hit would be 0. + req1 = make_request("1", token_ids, block_size, sha256) + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1) + assert num_computed_tokens == 2 * block_size + assert [len(blocks) for blocks in computed_blocks.blocks] == [2, 2] + + # The longer sibling resumes one alignment unit higher, off the same prompt. + longer = make_request("2", token_ids + [9] * block_size, block_size, sha256) + computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(longer) + assert num_computed_tokens == 3 * block_size + + def test_block_lookup_cache_single_block_per_key(): cache = BlockHashToBlockMap() key0 = BlockHashWithGroupId(b"hash0") diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 1ca4037e0e..3911372359 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -192,7 +192,7 @@ def test_circular_buffer_allocates_one_block_for_the_request_lifetime(): ) assert manager.allocate_new_blocks(request_id, num_tokens, num_tokens) == [] - manager.cache_blocks(request_id, 1024, replay_boundary=0) + manager.cache_blocks(request_id, 1024, replay_boundaries=(0,)) manager.remove_skipped_blocks(request_id, 1024) assert manager.get_num_common_prefix_blocks(request_id) == 0 assert manager.get_num_skipped_tokens(1024) == 0 diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 0755a69299..a495f997bf 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -307,27 +307,26 @@ class KVCacheCoordinator(ABC): for manager in self.single_type_managers ) - def get_replay_boundary(self, request: Request) -> int: - """Return the position a later request replaying this prompt resumes at. + def get_replay_boundaries(self, request: Request) -> tuple[int, ...]: + """Positions a later request replaying this prompt can resume at. - A cache hit is the shortest hit across all groups, so this is a - model-level position: every group has to retain state here, whether or - not it is the group that drops. Groups differ only in how much they - keep around it -- EAGLE groups also keep the block above, which they + A hit is the shortest across all groups, so every group retains state + at each position; EAGLE groups also keep the block above, which they match and drop back from (see ``reachable_block_mask``). - Under EAGLE that block must exist, so the boundary sits one alignment - unit below the prompt's last aligned position; every group's block size - divides the alignment, so the block above always fits in the prompt. + Two positions are reachable: a resend of the identical prompt is capped + at ``num_tokens - 1`` (its last token is recomputed for logits), a + longer sibling matches the final aligned block. They differ only on a + block-aligned prompt, where retaining just the higher one collapses the + resend's hit to 0. The alignment is the scheduler block size, not the + finer hash granularity, which would over-estimate the reach. """ if not self.eagle_group_ids: - return request.num_prompt_tokens - 1 - aligned = ( - request.num_prompt_tokens - // self.scheduler_block_size - * self.scheduler_block_size - ) - return max(aligned - self.scheduler_block_size, 0) + return (request.num_prompt_tokens - 1,) + block = self.scheduler_block_size + resend = (request.num_prompt_tokens - 1) // block * block + extension = request.num_prompt_tokens // block * block + return tuple(sorted({max(resend - block, 0), max(extension - block, 0)})) def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: """ @@ -339,7 +338,7 @@ class KVCacheCoordinator(ABC): that need to be cached (including tokens that are already cached). """ - replay_boundary = self.get_replay_boundary(request) + boundaries = self.get_replay_boundaries(request) for manager in self.single_type_managers: # Only cache tokens with finalized KV. The last num_reprefillable_tokens # tokens can be re-prefilled during multi-module MTP. @@ -350,7 +349,7 @@ class KVCacheCoordinator(ABC): request, num_tokens_to_cache, retention_interval=self.retention_interval, - replay_boundary=replay_boundary, + replay_boundaries=boundaries, ) def free(self, request_id: str) -> None: @@ -775,7 +774,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: cached_num_computed_tokens = self._align_cacheable(num_computed_tokens) - replay_boundary = self.get_replay_boundary(request) + boundaries = self.get_replay_boundaries(request) for manager in self.single_type_managers: num_tokens_to_cache = cached_num_computed_tokens # EAGLE groups match one block past each aligned boundary and drop @@ -802,7 +801,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): request, num_tokens_to_cache, retention_interval=self.retention_interval, - replay_boundary=replay_boundary, + replay_boundaries=boundaries, ) def find_longest_cache_hit( diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index a3cc1088e2..15b35d9f7b 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -452,7 +452,7 @@ class SingleTypeKVCacheManager(ABC): num_tokens: int, retention_interval: int | None = None, *, - replay_boundary: int, + replay_boundaries: Sequence[int], ) -> None: """ Cache the blocks for the request. @@ -465,6 +465,8 @@ class SingleTypeKVCacheManager(ABC): keeps dense checkpointing; ``0`` keeps only the latest replay boundary; a positive multiple of ``scheduler_block_size`` keeps a tail once per that-sized segment. Only SWA acts on it. + replay_boundaries: Positions a later request replaying this prompt + can resume at, from ``get_replay_boundaries``. """ num_cached_blocks = self.num_cached_block.get(request.request_id, 0) num_full_blocks = num_tokens // self.block_size @@ -473,9 +475,9 @@ class SingleTypeKVCacheManager(ABC): return # Token boundaries whose reachable tail must be retained under sparse - # retention: the replay boundary (``num_prompt - 1``, capped by - # ``get_computed_blocks``) and any detected shared-prefix junction. - reachable_boundaries = [replay_boundary] + # retention: every position a replaying sibling can resume at (see + # ``get_replay_boundaries``) and any detected shared-prefix junction. + reachable_boundaries = [*replay_boundaries] if request.shared_prefix_boundary: reachable_boundaries.append(request.shared_prefix_boundary) @@ -818,13 +820,13 @@ class FullAttentionManager(SingleTypeKVCacheManager): num_tokens: int, retention_interval: int | None = None, *, - replay_boundary: int, + replay_boundaries: Sequence[int], ) -> None: super().cache_blocks( request, num_tokens, retention_interval=retention_interval, - replay_boundary=replay_boundary, + replay_boundaries=replay_boundaries, ) hash_block_size = self.block_pool.hash_block_size if self.block_size == hash_block_size: @@ -1228,7 +1230,7 @@ class CircularBufferManager(FullAttentionManager): num_tokens: int, retention_interval: int | None = None, *, - replay_boundary: int, + replay_boundaries: Sequence[int], ) -> None: return @@ -1940,14 +1942,14 @@ class MambaManager(SingleTypeKVCacheManager): num_tokens: int, retention_interval: int | None = None, *, - replay_boundary: int, + replay_boundaries: Sequence[int], ) -> None: num_cached_blocks_before = self.num_cached_block.get(request.request_id, 0) super().cache_blocks( request, num_tokens, retention_interval=retention_interval, - replay_boundary=replay_boundary, + replay_boundaries=replay_boundaries, ) num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0) if self.mamba_cache_mode == "align": @@ -2090,7 +2092,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager): num_tokens: int, retention_interval: int | None = None, *, - replay_boundary: int, + replay_boundaries: Sequence[int], ) -> None: # We do not cache blocks for cross-attention to be shared between # requests, so this method is not relevant.