iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][KV Connector] Propagate EAGLE state across merged Mooncake store groups (#49069)

Browse Source 
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: xijiade.aihemaiti <3146335281@qq.com>
This commit is contained in:
Yifan Qiao
2026-08-02 19:32:26 -07:00
committed by GitHub
co-authored by xijiade.aihemaiti
parent e42c230e4f
commit 5c4fe4b17e
2 changed files with 118 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
+80
View File
@@ -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
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+38 -39
View File
@@ -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),