From 1d8afd3802bf5c3f1ed67e92c563aa5742050b09 Mon Sep 17 00:00:00 2001 From: tobymao Date: Tue, 8 Sep 2026 09:20:26 -0700 Subject: [PATCH] [BugFix] Retain both replay boundaries so an EAGLE resend of a block-aligned prompt still hits #53945 fixed sparse-retention reuse under EAGLE/MTP by moving the replay boundary one block below the prompt's last aligned position, which is where a lookup lands after the EAGLE drop. That is right for a sibling whose prompt merely *starts* with this one, but not for a resend of the identical prompt: a lookup is capped at ``num_tokens - 1`` because the last token must be recomputed to obtain logits, so an identical resend matches one block lower and then drops from there. The two coincide unless the prompt length is an exact multiple of the scheduler block size. There they differ by one block, and retaining only the higher one leaves the resend with every retained state above every candidate its lookup can produce -- the reconciled hit is 0, the same zero-hit failure sparse retention already avoids at unaligned prompt lengths. Return both positions from ``get_replay_boundaries`` and retain both. This only ever adds a position, so hit length cannot regress: swept over 136 prompt lengths (identical resend and longer sibling, hash block 16 and 64), 8 cases go from a 0-token hit to a real one and none get shorter. Note the alignment stays the scheduler block size rather than the finer hash granularity, on purpose: fine-grained hits extend into the first non-full block only when a partial tail was registered there, so a hash-granular boundary would over-estimate the reach and name a position above the real candidate (measured: that regresses lengths 160/176/224/240 from a real hit to 0). Signed-off-by: tobymao --- tests/v1/core/test_prefix_caching.py | 98 +++++++++++++++++++ .../core/test_single_type_kv_cache_manager.py | 2 +- vllm/v1/core/kv_cache_coordinator.py | 39 ++++---- vllm/v1/core/single_type_kv_cache_manager.py | 22 +++-- 4 files changed, 130 insertions(+), 31 deletions(-) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index d5536a1a4fa6..1760ca14d0d6 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 1ca4037e0e16..39113723591b 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 0755a692995f..a495f997bf3b 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -307,27 +307,26 @@ def allocate_new_blocks( 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 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: 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 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: 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 @@ def _align_cacheable(self, num_tokens: int) -> int: 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 @@ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: 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 a3cc1088e2c1..15b35d9f7b57 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 @@ def cache_blocks( 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 @@ def cache_blocks( 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 @@ def cache_blocks( 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 @@ def cache_blocks( 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 @@ def cache_blocks( num_tokens: int, retention_interval: int | None = None, *, - replay_boundary: int, + replay_boundaries: Sequence[int], ) -> None: return @@ -1940,14 +1942,14 @@ def cache_blocks( 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 @@ def cache_blocks( 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.