[BugFix] Retain both replay boundaries so an EAGLE resend of a block-aligned prompt still hits (#54713)

Signed-off-by: tobymao <toby.mao@gmail.com>

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")

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

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(

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.