From 2339bd39d661965b821bfbe652db8f536e0b3dcd Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Sun, 11 Oct 2026 02:15:18 -0700 Subject: [PATCH] [Bugfix][Mooncake] Use the EAGLE replay boundary when retaining sliding-window tails (#61012) Signed-off-by: Stu Cao Signed-off-by: Yifan Qiao Co-authored-by: Claude Fable 5.1 Co-authored-by: Yifan Qiao Co-authored-by: Codex --- .../unit/test_mooncake_store_coordinator.py | 44 ++++++++++++++++ .../v1/mooncake/store/coordinator.py | 13 +++-- vllm/v1/core/kv_cache_coordinator.py | 51 +++++++++++-------- vllm/v1/core/single_type_kv_cache_manager.py | 4 +- 4 files changed, 85 insertions(+), 27 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index ac52ca07ab..6896777226 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/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.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 48a36bc82f..2382948268 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/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, diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 019941a1bf..b6defc52e8 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/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 diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 77c6a4a5f7..218d59f430 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/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)