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 8b9d1e46b3..7460e45d17 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -567,3 +567,83 @@ def test_lookup_with_eagle_hybrid_full_plus_mamba_no_overrun(): # FullAttn matches 4 blocks, eagle pops 1 -> 48 verified tokens. The # Mamba group must serve its state@48 snapshot, not peek to state@64. assert hit == 48 + + +def test_eagle_flag_propagates_to_all_merged_swa_groups(): + """Regression for MTP x PD external-store 0% prefix hit. + + DSV4 splits SWA layers into several KV cache groups sharing one spec, and + only the group containing the MTP layer is annotated ``is_eagle_group``. + The lookup merges equal-spec groups, applies the eagle drop to the merged + group, and requires each chunk hash to exist in EVERY member group — so + the save-side masks must eagle-shift every member, not just the annotated + one. Without propagation the non-annotated groups never store the eagle + proof-run chunks and every external lookup returns 0. + """ + swa = _swa(block_size=16, sliding_window=32) + groups = [ + KVCacheGroupSpec(["L0"], _full(64)), + KVCacheGroupSpec(["L1"], swa), + KVCacheGroupSpec(["L2"], swa, is_eagle_group=True), + ] + coord = _make_coord(groups, hash_block_size=16, use_eagle=True) + assert coord.eagle_group_ids == {1, 2} + + # Save side: both SWA groups must produce identical (eagle-shifted) masks. + masks = coord.store_mask(128, num_prompt_tokens=130) + assert masks[1] == masks[2] + + # Round trip: everything store_mask kept is in the store; the eagle + # lookup must then serve a non-zero hit (it was 0 before the fix). + hs = _hashes(128 // 16) + exists = set() + for g_idx, g in enumerate(groups): + ghashes = chunk_hashes_for_block_size(hs, 16, g.kv_cache_spec.block_size) + mask = masks[g_idx] + for i in range(128 // g.kv_cache_spec.block_size): + if mask is None or mask[i]: + exists.add((g_idx, bytes(ghashes[i]))) + _masks, hit = coord.find_longest_cache_hit( + hs, max_length=128, cached_block_pool=ExternalCachedBlockPool(16, exists) + ) + assert hit == 64 + + +def test_dsv4_five_group_eagle_store_lookup_round_trip(): + """Cover the five KV groups observed with DeepSeek-V4-Flash + MTP. + + The two 64-token SWA groups have identical specs, but only one owns the + EAGLE layer. Saving each group through its store mask must still leave a + prefix that the merged SWA lookup can consume. + """ + swa_64_sw128 = _swa(block_size=64, sliding_window=128) + groups = [ + KVCacheGroupSpec(["full_mla"], _full(block_size=256)), + KVCacheGroupSpec(["swa"], swa_64_sw128), + KVCacheGroupSpec(["mtp"], swa_64_sw128, is_eagle_group=True), + KVCacheGroupSpec(["c4_state"], _swa(block_size=4, sliding_window=8)), + KVCacheGroupSpec(["c128_state"], _swa(block_size=8, sliding_window=128)), + ] + coord = _make_coord(groups, hash_block_size=4, use_eagle=True) + token_len = 768 + hashes = _hashes(token_len // coord.hash_block_size) + + # Mirror MooncakeStoreWorker's aligned save: only keys selected by each + # group's store mask are visible to the external lookup. + exists: set[tuple[int, bytes]] = set() + store_masks = coord.store_mask(token_len) + 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=token_len, + cached_block_pool=ExternalCachedBlockPool(coord.hash_block_size, exists), + ) + + # The final 256-token segment has no lookahead block, so EAGLE falls back + # to the previous aligned boundary instead of consuming all 768 tokens. + assert hit == 512 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 2fc71ee488..de5c846354 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 @@ -10,13 +10,11 @@ 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 SpecGroup from vllm.v1.core.kv_cache_utils import ( BlockHash, KVCacheBlock, ) -from vllm.v1.core.single_type_kv_cache_manager import ( - SingleTypeKVCacheManager, -) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheGroupSpec, @@ -90,11 +88,6 @@ class MooncakeStoreCoordinator: self.use_eagle = use_eagle # Mirror vLLM core's KVCacheCoordinator.retention_interval. self.retention_interval = retention_interval - self.eagle_group_ids = { - i for i, g in enumerate(kv_cache_groups) if g.is_eagle_group - } - if use_eagle and not self.eagle_group_ids: - self.eagle_group_ids = set(range(len(kv_cache_groups))) self._verify_and_split_kv_cache_groups() def align_lookup_length(self, length: int) -> int: @@ -109,34 +102,37 @@ class MooncakeStoreCoordinator: """Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but dispatches via spec_manager_map (we don't allocate managers). """ - attention_groups: list[ - tuple[KVCacheSpec, list[int], type[SingleTypeKVCacheManager]] - ] = [] + attention_groups: list[SpecGroup] = [] for i, g in enumerate(self.kv_cache_groups): spec = _unwrap_spec(g.kv_cache_spec) manager_cls = KVCacheSpecRegistry.get_manager_class(spec) assert manager_cls is not None, ( f"No manager registered for KVCacheSpec {spec}" ) - for existing_spec, group_ids, existing_cls in attention_groups: - if existing_spec == spec: - assert manager_cls is existing_cls - group_ids.append(i) + for idx, group in enumerate(attention_groups): + if group.spec == spec: + assert manager_cls is group.manager_cls + group.group_ids.append(i) + if g.is_eagle_group and not group.use_eagle: + attention_groups[idx] = group._replace(use_eagle=True) break else: - attention_groups.append((spec, [i], manager_cls)) + attention_groups.append( + SpecGroup(spec, [i], manager_cls, g.is_eagle_group) + ) # Full attention first (matches upstream convergence ordering). - self.attention_groups = sorted( - attention_groups, - key=lambda x: not isinstance(x[0], FullAttentionSpec), - ) - self.eagle_attn_group_indices: set[int] = { - i - for i, (_, group_ids, _) in enumerate(self.attention_groups) - if any(self.kv_cache_groups[gid].is_eagle_group for gid in group_ids) + attention_groups.sort(key=lambda g: not isinstance(g.spec, FullAttentionSpec)) + # Conservatively flag all groups when use_eagle is set but none is flagged. + if self.use_eagle and not any(g.use_eagle for g in attention_groups): + attention_groups = [g._replace(use_eagle=True) for g in attention_groups] + self.attention_groups = attention_groups + # Per-group eagle bits. SpecGroup carries use_eagle for the whole + # merged spec group, so the per-group store/lookup masks agree with + # the merged-group hit check, which applies the eagle drop to every + # group sharing the spec. + self.eagle_group_ids = { + gid for g in attention_groups if g.use_eagle for gid in g.group_ids } - if self.use_eagle and not self.eagle_attn_group_indices: - self.eagle_attn_group_indices = set(range(len(self.attention_groups))) def find_longest_cache_hit( self, @@ -287,25 +283,24 @@ class MooncakeStoreCoordinator: """Mirrors HybridKVCacheCoordinator.find_longest_cache_hit but dispatches via spec_manager_map (we don't allocate managers). - When ``apply_eagle`` is False, ignore ``eagle_attn_group_indices`` — + When ``apply_eagle`` is False, ignore each group's ``use_eagle`` — used by ``load_mask`` to avoid popping a second block on top of the one already removed by the lookup. """ - eagle_indices = self.eagle_attn_group_indices if apply_eagle else set() alignment_tokens = ( self.hash_block_size if self.enable_partial_hash_hits else self.lcm_block_size ) if len(self.attention_groups) == 1: - spec, group_ids, manager_cls = self.attention_groups[0] + spec, group_ids, manager_cls, group_eagle = self.attention_groups[0] hit_blocks, hit_length = manager_cls.find_longest_cache_hit( block_hashes=block_hashes, # type: ignore[arg-type] max_length=max_length, kv_cache_group_ids=group_ids, block_pool=cast(BlockPool, cached_block_pool), kv_cache_spec=spec, - drop_eagle_block=(0 in eagle_indices), + drop_eagle_block=apply_eagle and group_eagle, alignment_tokens=alignment_tokens, ) num_groups = len(self.kv_cache_groups) @@ -320,14 +315,16 @@ class MooncakeStoreCoordinator: hit_length_by_group: list[int] = [0] * num_groups is_simple_hybrid = len(self.attention_groups) == 2 and isinstance( - self.attention_groups[0][0], FullAttentionSpec + self.attention_groups[0].spec, FullAttentionSpec ) eagle_verified: set[int] = set() while True: curr_hit_length = hit_length - for idx, (spec, group_ids, manager_cls) in enumerate(self.attention_groups): + for idx, (spec, group_ids, manager_cls, group_eagle) in enumerate( + self.attention_groups + ): first_group_id = group_ids[0] cached = hit_blocks_by_group[first_group_id] if isinstance(spec, FullAttentionSpec) and cached is not None: @@ -336,7 +333,9 @@ class MooncakeStoreCoordinator: ) continue - drop_eagle_block = idx in eagle_indices and idx not in eagle_verified + drop_eagle_block = ( + apply_eagle and group_eagle and idx not in eagle_verified + ) _max_length = curr_hit_length # No eagle peek margin for a recurrent (Mamba) group: its finder # never drops a block, so a widened bound would match past the @@ -377,14 +376,14 @@ class MooncakeStoreCoordinator: # Truncate full-attention hit_blocks to final converged length; # other specs already trim themselves inside their hit logic. cdiv keeps # the partial tail block when hit_length is not block-aligned. - spec0, group_ids0, _ = self.attention_groups[0] - if isinstance(spec0, FullAttentionSpec): - num_blocks = cdiv(hit_length, spec0.block_size) - for gid in group_ids0: - full_blks = hit_blocks_by_group[gid] + first_group = self.attention_groups[0] + if isinstance(first_group.spec, FullAttentionSpec): + num_blocks = cdiv(hit_length, first_group.spec.block_size) + for group_id in first_group.group_ids: + full_blks = hit_blocks_by_group[group_id] assert full_blks is not None del full_blks[num_blocks:] - hit_length_by_group[gid] = hit_length + hit_length_by_group[group_id] = hit_length return ( tuple(blks if blks is not None else [] for blks in hit_blocks_by_group),