[Bugfix][Mooncake] Use the EAGLE replay boundary when retaining sliding-window tails (#61012)

Signed-off-by: Stu Cao <stu@primeintellect.ai>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>

tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py

@@ -3,6 +3,7 @@
 
 from math import lcm
 
+import pytest
 import torch
 
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator import (  # noqa: E501
@@ -422,6 +423,49 @@ def test_store_mask_retention_interval_zero_keeps_only_replay_boundary():
     assert masks[1] == [i == 11 for i in range(16)]
 
 
+def test_store_mask_retention_interval_zero_eagle_keeps_lookup_tail():
+    """With EAGLE the replay boundary sits one scheduler block lower, as in
+    core's ``replay_boundaries``: the lookup matches that block and drops
+    back from it. The retained SWA tail must span the ``need`` contiguous
+    blocks the lookup requires (cdiv(8 - 1, 8) + 1 = 2 here); the non-EAGLE
+    boundary leaves only one of them inside the saved prefix."""
+    coord = _make_coord(
+        _retention_groups(), hash_block_size=8, use_eagle=True, retention_interval=0
+    )
+    # num_prompt=100 -> boundary 64 -> chunks 7, 8 (8 is the EAGLE peek).
+    masks = coord.store_mask(96, num_prompt_tokens=100)
+    assert masks[0] is None
+    assert masks[1] == [i in (7, 8) for i in range(12)]
+    # Aligned num_prompt=96 -> boundaries 32 (resend) and 64 (longer sibling).
+    masks = coord.store_mask(96, num_prompt_tokens=96)
+    assert masks[1] == [i in (3, 4, 7, 8) for i in range(12)]
+
+
+@pytest.mark.parametrize("num_prompt_tokens", [100, 96])
+def test_store_mask_retention_interval_zero_eagle_round_trip(num_prompt_tokens):
+    """Everything store_mask keeps is in the store; the EAGLE lookup of the
+    same prompt must then reach the full block below the one it drops (64 of
+    96). The unaligned prompt hit 0 before the fix; the aligned one already
+    hit."""
+    groups = _retention_groups()
+    coord = _make_coord(groups, hash_block_size=8, use_eagle=True, retention_interval=0)
+    hashes = _hashes(96 // coord.hash_block_size)
+    exists: set[tuple[int, bytes]] = set()
+    store_masks = coord.store_mask(96, num_prompt_tokens=num_prompt_tokens)
+    for gid, (group, mask) in enumerate(zip(groups, store_masks, strict=True)):
+        group_hashes = coord.block_hashes_for_spec(hashes, group.kv_cache_spec)
+        for chunk_id, block_hash in enumerate(group_hashes):
+            if mask is None or mask[chunk_id]:
+                exists.add((gid, bytes(block_hash)))
+
+    _masks, hit = coord.find_longest_cache_hit(
+        hashes,
+        max_length=96,
+        cached_block_pool=ExternalCachedBlockPool(coord.hash_block_size, exists),
+    )
+    assert hit == 64
+
+
 def test_store_mask_retention_interval_keeps_segment_and_replay_tails():
     """Sparse segment tails (interval=64 -> chunks 7,15) plus the replay
     boundary tail (num_prompt=100 -> chunk 11) coexist."""

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

@@ -11,6 +11,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
 )
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_coordinator import replay_boundaries
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
     KVCacheBlock,
@@ -341,6 +342,15 @@ class MooncakeStoreCoordinator:
             f"aligned_token_len ({aligned_token_len}) must be a multiple of "
             f"{mask_alignment}"
         )
+        reachable_boundaries = (
+            ()
+            if num_prompt_tokens is None
+            else replay_boundaries(
+                num_prompt_tokens,
+                self.lcm_block_size,
+                use_eagle=bool(self.eagle_group_ids),
+            )
+        )
         masks: list[list[bool] | None] = []
         for g_idx, g in enumerate(self.kv_cache_groups):
             spec = _unwrap_spec(g.kv_cache_spec)
@@ -352,9 +362,6 @@ class MooncakeStoreCoordinator:
             manager_cls = KVCacheSpecRegistry.get_manager_class(spec)
             assert manager_cls is not None
             use_eagle = g_idx in self.eagle_group_ids
-            reachable_boundaries = (
-                () if num_prompt_tokens is None else (num_prompt_tokens - 1,)
-            )
             mask = manager_cls.reachable_block_mask(
                 start_block=start_chunk,
                 end_block=end_chunk,

vllm/v1/core/kv_cache_coordinator.py

@@ -62,6 +62,30 @@ def _validate_prefix_cache_retention_interval(
         )
 
 
+def replay_boundaries(
+    num_prompt_tokens: int, scheduler_block_size: int, use_eagle: bool
+) -> tuple[int, ...]:
+    """Positions a later request replaying a prompt can resume at.
+
+    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``).
+
+    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 use_eagle:
+        return (num_prompt_tokens - 1,)
+    block = scheduler_block_size
+    resend = (num_prompt_tokens - 1) // block * block
+    extension = num_prompt_tokens // block * block
+    return tuple(sorted({max(resend - block, 0), max(extension - block, 0)}))
+
+
 class KVCacheCoordinator(ABC):
     """Coordinate the KV cache of different KV cache groups."""
 
@@ -319,27 +343,6 @@ class KVCacheCoordinator(ABC):
             for manager in self.single_type_managers
         )
 
-    def get_replay_boundaries(self, request: Request) -> tuple[int, ...]:
-        """Positions a later request replaying this prompt can resume at.
-
-        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``).
-
-        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,)
-        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 get_num_cacheable_tokens(
         self, num_computed_tokens: int, kv_cache_group_id: int
     ) -> int:
@@ -358,7 +361,11 @@ class KVCacheCoordinator(ABC):
                 (including tokens that are already cached).
 
         """
-        boundaries = self.get_replay_boundaries(request)
+        boundaries = replay_boundaries(
+            request.num_prompt_tokens,
+            self.scheduler_block_size,
+            use_eagle=bool(self.eagle_group_ids),
+        )
         for group_id, manager in enumerate(self.single_type_managers):
             if not manager.enable_caching:
                 continue

vllm/v1/core/single_type_kv_cache_manager.py

@@ -492,7 +492,7 @@ class SingleTypeKVCacheManager(ABC):
                 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``.
+                can resume at, from ``replay_boundaries``.
 
         """
         if not self.kv_cache_spec.prefix_cacheable:
@@ -505,7 +505,7 @@ class SingleTypeKVCacheManager(ABC):
 
         # Token boundaries whose reachable tail must be retained under sparse
         # retention: every position a replaying sibling can resume at (see
-        # ``get_replay_boundaries``) and any detected shared-prefix junction.
+        # ``replay_boundaries``) and any detected shared-prefix junction.
         reachable_boundaries = [*replay_boundaries]
         if request.shared_prefix_boundary:
             reachable_boundaries.append(request.shared_prefix_boundary)