[Bugfix][KV Connector] Propagate EAGLE state across merged Mooncake store groups (#49069)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: xijiade.aihemaiti <3146335281@qq.com>
Yifan Qiao committed
Aug 2, 2026 at 19:32 UTC
5c4fe4b17eaf42eda9698d9dccb4b2f2e0561641
2 files changed
+118
-39
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
+80
@@ -567,3 +567,83 @@ def test_lookup_with_eagle_hybrid_full_plus_mamba_no_overrun():
567
# FullAttn matches 4 blocks, eagle pops 1 -> 48 verified tokens. The
568
# Mamba group must serve its state@48 snapshot, not peek to state@64.
569
assert hit == 48
570
+
571
+
572
+def test_eagle_flag_propagates_to_all_merged_swa_groups():
573
+ """Regression for MTP x PD external-store 0% prefix hit.
574
+
575
+ DSV4 splits SWA layers into several KV cache groups sharing one spec, and
576
+ only the group containing the MTP layer is annotated ``is_eagle_group``.
577
+ The lookup merges equal-spec groups, applies the eagle drop to the merged
578
+ group, and requires each chunk hash to exist in EVERY member group — so
579
+ the save-side masks must eagle-shift every member, not just the annotated
580
+ one. Without propagation the non-annotated groups never store the eagle
581
+ proof-run chunks and every external lookup returns 0.
582
+ """
583
+ swa = _swa(block_size=16, sliding_window=32)
584
+ groups = [
585
+ KVCacheGroupSpec(["L0"], _full(64)),
586
+ KVCacheGroupSpec(["L1"], swa),
587
+ KVCacheGroupSpec(["L2"], swa, is_eagle_group=True),
588
+ ]
589
+ coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
590
+ assert coord.eagle_group_ids == {1, 2}
591
+
592
+ # Save side: both SWA groups must produce identical (eagle-shifted) masks.
593
+ masks = coord.store_mask(128, num_prompt_tokens=130)
594
+ assert masks[1] == masks[2]
595
+
596
+ # Round trip: everything store_mask kept is in the store; the eagle
597
+ # lookup must then serve a non-zero hit (it was 0 before the fix).
598
+ hs = _hashes(128 // 16)
599
+ exists = set()
600
+ for g_idx, g in enumerate(groups):
601
+ ghashes = chunk_hashes_for_block_size(hs, 16, g.kv_cache_spec.block_size)
602
+ mask = masks[g_idx]
603
+ for i in range(128 // g.kv_cache_spec.block_size):
604
+ if mask is None or mask[i]:
605
+ exists.add((g_idx, bytes(ghashes[i])))
606
+ _masks, hit = coord.find_longest_cache_hit(
607
+ hs, max_length=128, cached_block_pool=ExternalCachedBlockPool(16, exists)
608
+ )
609
+ assert hit == 64
610
+
611
+
612
+def test_dsv4_five_group_eagle_store_lookup_round_trip():
613
+ """Cover the five KV groups observed with DeepSeek-V4-Flash + MTP.
614
+
615
+ The two 64-token SWA groups have identical specs, but only one owns the
616
+ EAGLE layer. Saving each group through its store mask must still leave a
617
+ prefix that the merged SWA lookup can consume.
618
+ """
619
+ swa_64_sw128 = _swa(block_size=64, sliding_window=128)
620
+ groups = [
621
+ KVCacheGroupSpec(["full_mla"], _full(block_size=256)),
622
+ KVCacheGroupSpec(["swa"], swa_64_sw128),
623
+ KVCacheGroupSpec(["mtp"], swa_64_sw128, is_eagle_group=True),
624
+ KVCacheGroupSpec(["c4_state"], _swa(block_size=4, sliding_window=8)),
625
+ KVCacheGroupSpec(["c128_state"], _swa(block_size=8, sliding_window=128)),
626
+ ]
627
+ coord = _make_coord(groups, hash_block_size=4, use_eagle=True)
628
+ token_len = 768
629
+ hashes = _hashes(token_len // coord.hash_block_size)
630
+
631
+ # Mirror MooncakeStoreWorker's aligned save: only keys selected by each
632
+ # group's store mask are visible to the external lookup.
633
+ exists: set[tuple[int, bytes]] = set()
634
+ store_masks = coord.store_mask(token_len)
635
+ for gid, (group, mask) in enumerate(zip(groups, store_masks, strict=True)):
636
+ group_hashes = coord.block_hashes_for_spec(hashes, group.kv_cache_spec)
637
+ for chunk_id, block_hash in enumerate(group_hashes):
638
+ if mask is None or mask[chunk_id]:
639
+ exists.add((gid, bytes(block_hash)))
640
+
641
+ _masks, hit = coord.find_longest_cache_hit(
642
+ hashes,
643
+ max_length=token_len,
644
+ cached_block_pool=ExternalCachedBlockPool(coord.hash_block_size, exists),
645
+ )
646
+
647
+ # The final 256-token segment has no lookahead block, so EAGLE falls back
648
+ # to the previous aligned boundary instead of consuming all 768 tokens.
649
+ assert hit == 512
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+38
-39
@@ -10,13 +10,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
10
)
11
from vllm.utils.math_utils import cdiv
12
from vllm.v1.core.block_pool import BlockPool
13
+from vllm.v1.core.kv_cache_coordinator import SpecGroup
14
from vllm.v1.core.kv_cache_utils import (
15
BlockHash,
16
KVCacheBlock,
17
)
17
-from vllm.v1.core.single_type_kv_cache_manager import (
18
- SingleTypeKVCacheManager,
19
-)
18
from vllm.v1.kv_cache_interface import (
19
FullAttentionSpec,
20
KVCacheGroupSpec,
@@ -90,11 +88,6 @@ class MooncakeStoreCoordinator:
88
self.use_eagle = use_eagle
89
# Mirror vLLM core's KVCacheCoordinator.retention_interval.
90
self.retention_interval = retention_interval
93
- self.eagle_group_ids = {
94
- i for i, g in enumerate(kv_cache_groups) if g.is_eagle_group
95
- }
96
- if use_eagle and not self.eagle_group_ids:
97
- self.eagle_group_ids = set(range(len(kv_cache_groups)))
91
self._verify_and_split_kv_cache_groups()
92
93
def align_lookup_length(self, length: int) -> int:
@@ -109,34 +102,37 @@ class MooncakeStoreCoordinator:
102
"""Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but
103
dispatches via spec_manager_map (we don't allocate managers).
104
"""
112
- attention_groups: list[
113
- tuple[KVCacheSpec, list[int], type[SingleTypeKVCacheManager]]
114
- ] = []
105
+ attention_groups: list[SpecGroup] = []
106
for i, g in enumerate(self.kv_cache_groups):
107
spec = _unwrap_spec(g.kv_cache_spec)
108
manager_cls = KVCacheSpecRegistry.get_manager_class(spec)
109
assert manager_cls is not None, (
110
f"No manager registered for KVCacheSpec {spec}"
111
)
121
- for existing_spec, group_ids, existing_cls in attention_groups:
122
- if existing_spec == spec:
123
- assert manager_cls is existing_cls
124
- group_ids.append(i)
112
+ for idx, group in enumerate(attention_groups):
113
+ if group.spec == spec:
114
+ assert manager_cls is group.manager_cls
115
+ group.group_ids.append(i)
116
+ if g.is_eagle_group and not group.use_eagle:
117
+ attention_groups[idx] = group._replace(use_eagle=True)
118
break
119
else:
127
- attention_groups.append((spec, [i], manager_cls))
120
+ attention_groups.append(
121
+ SpecGroup(spec, [i], manager_cls, g.is_eagle_group)
122
+ )
123
# Full attention first (matches upstream convergence ordering).
129
- self.attention_groups = sorted(
130
- attention_groups,
131
- key=lambda x: not isinstance(x[0], FullAttentionSpec),
132
- )
133
- self.eagle_attn_group_indices: set[int] = {
134
- i
135
- for i, (_, group_ids, _) in enumerate(self.attention_groups)
136
- if any(self.kv_cache_groups[gid].is_eagle_group for gid in group_ids)
124
+ attention_groups.sort(key=lambda g: not isinstance(g.spec, FullAttentionSpec))
125
+ # Conservatively flag all groups when use_eagle is set but none is flagged.
126
+ if self.use_eagle and not any(g.use_eagle for g in attention_groups):
127
+ attention_groups = [g._replace(use_eagle=True) for g in attention_groups]
128
+ self.attention_groups = attention_groups
129
+ # Per-group eagle bits. SpecGroup carries use_eagle for the whole
130
+ # merged spec group, so the per-group store/lookup masks agree with
131
+ # the merged-group hit check, which applies the eagle drop to every
132
+ # group sharing the spec.
133
+ self.eagle_group_ids = {
134
+ gid for g in attention_groups if g.use_eagle for gid in g.group_ids
135
}
138
- if self.use_eagle and not self.eagle_attn_group_indices:
139
- self.eagle_attn_group_indices = set(range(len(self.attention_groups)))
136
137
def find_longest_cache_hit(
138
self,
@@ -287,25 +283,24 @@ class MooncakeStoreCoordinator:
283
"""Mirrors HybridKVCacheCoordinator.find_longest_cache_hit but
284
dispatches via spec_manager_map (we don't allocate managers).
285
290
- When ``apply_eagle`` is False, ignore ``eagle_attn_group_indices`` —
286
+ When ``apply_eagle`` is False, ignore each group's ``use_eagle`` —
287
used by ``load_mask`` to avoid popping a second block on top of the
288
one already removed by the lookup.
289
"""
294
- eagle_indices = self.eagle_attn_group_indices if apply_eagle else set()
290
alignment_tokens = (
291
self.hash_block_size
292
if self.enable_partial_hash_hits
293
else self.lcm_block_size
294
)
295
if len(self.attention_groups) == 1:
301
- spec, group_ids, manager_cls = self.attention_groups[0]
296
+ spec, group_ids, manager_cls, group_eagle = self.attention_groups[0]
297
hit_blocks, hit_length = manager_cls.find_longest_cache_hit(
298
block_hashes=block_hashes, # type: ignore[arg-type]
299
max_length=max_length,
300
kv_cache_group_ids=group_ids,
301
block_pool=cast(BlockPool, cached_block_pool),
302
kv_cache_spec=spec,
308
- drop_eagle_block=(0 in eagle_indices),
303
+ drop_eagle_block=apply_eagle and group_eagle,
304
alignment_tokens=alignment_tokens,
305
)
306
num_groups = len(self.kv_cache_groups)
@@ -320,14 +315,16 @@ class MooncakeStoreCoordinator:
315
hit_length_by_group: list[int] = [0] * num_groups
316
317
is_simple_hybrid = len(self.attention_groups) == 2 and isinstance(
323
- self.attention_groups[0][0], FullAttentionSpec
318
+ self.attention_groups[0].spec, FullAttentionSpec
319
)
320
eagle_verified: set[int] = set()
321
322
while True:
323
curr_hit_length = hit_length
324
330
- for idx, (spec, group_ids, manager_cls) in enumerate(self.attention_groups):
325
+ for idx, (spec, group_ids, manager_cls, group_eagle) in enumerate(
326
+ self.attention_groups
327
+ ):
328
first_group_id = group_ids[0]
329
cached = hit_blocks_by_group[first_group_id]
330
if isinstance(spec, FullAttentionSpec) and cached is not None:
@@ -336,7 +333,9 @@ class MooncakeStoreCoordinator:
333
)
334
continue
335
339
- drop_eagle_block = idx in eagle_indices and idx not in eagle_verified
336
+ drop_eagle_block = (
337
+ apply_eagle and group_eagle and idx not in eagle_verified
338
+ )
339
_max_length = curr_hit_length
340
# No eagle peek margin for a recurrent (Mamba) group: its finder
341
# never drops a block, so a widened bound would match past the
@@ -377,14 +376,14 @@ class MooncakeStoreCoordinator:
376
# Truncate full-attention hit_blocks to final converged length;
377
# other specs already trim themselves inside their hit logic. cdiv keeps
378
# the partial tail block when hit_length is not block-aligned.
380
- spec0, group_ids0, _ = self.attention_groups[0]
381
- if isinstance(spec0, FullAttentionSpec):
382
- num_blocks = cdiv(hit_length, spec0.block_size)
383
- for gid in group_ids0:
384
- full_blks = hit_blocks_by_group[gid]
379
+ first_group = self.attention_groups[0]
380
+ if isinstance(first_group.spec, FullAttentionSpec):
381
+ num_blocks = cdiv(hit_length, first_group.spec.block_size)
382
+ for group_id in first_group.group_ids:
383
+ full_blks = hit_blocks_by_group[group_id]
384
assert full_blks is not None
385
del full_blks[num_blocks:]
387
- hit_length_by_group[gid] = hit_length
386
+ hit_length_by_group[group_id] = hit_length
387
388
return (
389
tuple(blks if blks is not None else [] for blks in hit_blocks_by_group),