执行摘要
- 一句话:修复 Mooncake 合并组 EAGLE 位传播,恢复外部前缀缓存命中
- 推荐动作:值得精读。核心看点是“避免镜像实现漂移”——直接复用核心
SpecGroup 并在合并组上传播 EAGLE 位,使外部存储与核心缓存协调器保持同一分组语义;测试的 store-to-lookup 往返设计(先按 store mask 填 exists 集,再走对外查找)也值得借鉴。合入前请重点复核 walrus 更改,并协调 #48361 的合入顺序;建议合入后用真实 DeepSeek-V4 MTP P/D 部署跑一次前缀命中率回归。
功能与动机
PR body 明确指出根因:DeepSeek V4 只在包含 MTP 层的缓存组上标注 EAGLE 组,coordinator 合并等 spec 组后对合并组应用 EAGLE last-block drop,而 save-side mask 之前只对单独标注的成员移位。外部 lookup 要求 hash 在每个成员组中都存在,因此不一致的 mask 导致其他 SWA 组的 proof-run chunks 缺失,外部前缀命中降为 0。issue #45785 在真实 DeepSeek-V4-Flash 部署上复现了该失败。
实现拆解
- 变更入口:
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py 的 MooncakeStoreCoordinator。此前它是核心 HybridKVCacheCoordinator 的独立镜像,用本地元组表示分组,eagle_group_ids 直接从原始 is_eagle_group 标注计算;本 PR 将其改为复用核心的 SpecGroup 表示。
- 合并与传播:
_verify_and_split_kv_cache_groups 中,等 spec 组归并时若新成员是 EAGLE 组,用 _replace(use_eagle=True) 提升整个合并组;use_eagle=True 且无任何标注时保守 fallback 全开。eagle_group_ids 改为从合并后的 SpecGroup.use_eagle 展开到全部成员组 ID,保证 store_mask 与 lookup_mask 的 save/lookup 语义一致。
- 命中逻辑对齐:
_find_hit_blocks 的单组与多组路径中 drop_eagle_block 从“基于下标集合 eagle_attn_group_indices”改为“apply_eagle and group_eagle”,使每个成员组都执行相同 EAGLE 尾块丢弃;结尾 full-attn 块截断处由 assert 改为 walrus 条件跳过(行为宽松化)。
- 测试配套:
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py 新增 test_eagle_flag_propagates_to_all_merged_swa_groups(三组形态,断言 eagle_group_ids == {1, 2}、两组 mask 相等、往返命中 64)与 test_dsv4_five_group_eagle_store_lookup_round_trip(DSV4-Flash 真实五组形态,往返命中 512)。
- 配置与部署:无新增配置项或部署改动;合入前需 rebase 解决与最新 coordinator 代码的冲突。
关键文件:
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py(模块 存储协调器;类别 source;类型 core-logic;符号 _verify_and_split_kv_cache_groups, _find_hit_blocks, eagle_group_ids, SpecGroup): 核心修复文件:将本地分组元组替换为核心 SpecGroup,合并等 spec 组时传播 EAGLE 位,并据此派生 eagle_group_ids,保证 store/lookup mask 一致。
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py(模块 单元测试;类别 test;类型 test-coverage;符号 test_eagle_flag_propagates_to_all_merged_swa_groups, test_dsv4_five_group_eagle_store_lookup_round_trip): 新增两个回归测试,覆盖三组与 DSV4-Flash 真实五组形态的 store-to-lookup 往返,锁定修复行为。
关键符号:_verify_and_split_kv_cache_groups, _find_hit_blocks
关键源码片段
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
核心修复文件:将本地分组元组替换为核心 SpecGroup,合并等 spec 组时传播 EAGLE 位,并据此派生 eagle_group_ids,保证 store/lookup mask 一致。
# 合并等 spec 的 KV cache 组并传播 EAGLE 位。
# 复用核心 vLLM 的 SpecGroup,避免外部存储侧分组语义
# 与 HybridKVCacheCoordinator 持续漂移。
def _verify_and_split_kv_cache_groups(self) -> None:
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
for idx, group in enumerate(attention_groups):
if group.spec == spec:
# 相同 spec 的组归并到同一个 SpecGroup。
# 关键修复点:只要该合并组内任一成员是 EAGLE 组,
# 就将 use_eagle 提升到合并组级别。
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(SpecGroup(spec, [i], manager_cls, g.is_eagle_group))
# Full attention 优先(与上游收敛排序一致)。
attention_groups.sort(key=lambda g: not isinstance(g.spec, FullAttentionSpec))
# use_eagle 开启但没有任何组标注时,保守地全部启用。
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
# 修复的核心:eagle_group_ids 不再直接取原始单组的 is_eagle_group,
# 而是从“合并后的 SpecGroup.use_eagle”推导出全部成员组 ID。
# 这样 store_mask 与 lookup_mask 对每个成员组都做相同的 EAGLE 尾块
# 移位,与合并组命中检查的 drop 行为保持一致。
self.eagle_group_ids = {
gid for g in attention_groups if g.use_eagle for gid in g.group_ids
}
# 多组命中路径:按合并后的 SpecGroup 逐个查找。
# drop_eagle_block 改为由 apply_eagle 与合并组自身的 use_eagle
# 共同决定,保证同 spec 的每个成员组都执行相同的 EAGLE 尾块丢弃。
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:
curr_hit_length = min(curr_hit_length, hit_length_by_group[first_group_id])
continue
drop_eagle_block = apply_eagle and group_eagle and idx not in eagle_verified
_max_length = curr_hit_length
# Mamba 循环状态不允许跨已验证前缀取抹除尾块的余量(见 #43559)。
if drop_eagle_block and not isinstance(spec, MambaSpec):
eagle_margin = (
self.hash_block_size
if self.enable_partial_hash_hits
and manager_cls.supports_fine_grained_hash_lookup
and spec.block_size > self.hash_block_size
else spec.block_size
)
_max_length = min(curr_hit_length + eagle_margin, max_length)
hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit(
block_hashes=block_hashes,
max_length=_max_length,
kv_cache_group_ids=group_ids,
block_pool=cast(BlockPool, cached_block_pool),
kv_cache_spec=spec,
drop_eagle_block=drop_eagle_block,
alignment_tokens=alignment_tokens,
)
if drop_eagle_block:
eagle_verified.add(idx)
elif _new_hit_length < curr_hit_length:
eagle_verified.clear()
curr_hit_length = _new_hit_length
for gid, blocks in zip(group_ids, hit_blocks, strict=True):
hit_blocks_by_group[gid] = blocks
hit_length_by_group[gid] = _new_hit_length
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
新增两个回归测试,覆盖三组与 DSV4-Flash 真实五组形态的 store-to-lookup 往返,锁定修复行为。
def test_eagle_flag_propagates_to_all_merged_swa_groups():
"""回归场景:MTP 与 PD 组合下外部存储前缀命中率为 0。
DSV4 会把 SWA 层拆成多个共享同一 spec 的 KV cache 组,
只有包含 MTP 层的组被标注为 is_eagle_group。lookup 侧合并
等 spec 组后对合并组统一做 EAGLE 尾块丢弃,并要求 hash 在
每个成员组都存在;因此 save 侧 mask 也必须对每个成员组做
相同的 EAGLE 移位,否则非标注组永远不存 proof-run 块,
外部查找命中固定为 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)
# 修复后两个同 spec 的 SWA 组都被标记为 EAGLE 组。
assert coord.eagle_group_ids == {1, 2}
# save 侧:两个 SWA 组必须产出完全相同的(EAGLE 移位后)mask。
masks = coord.store_mask(128, num_prompt_tokens=130)
assert masks[1] == masks[2]
# 往返验证:store_mask 保留的块全部放入外部存储,
# 合并组的 EAGLE lookup 必须能命中非零前缀(修复前为 0)。
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
评论区精华
claude[bot]:逐行对比核心 HybridKVCacheCoordinator.verify_and_split_kv_cache_groups,确认 SpecGroup 复用语义与 fallback 行为一致、eagle_group_ids 派生是正确的修复点,且 apply_eagle 门控的 drop_eagle_block 与旧实现行为等价;但指出尾部截断循环里 assert full_blks is not None 被放宽为 walrus 条件跳过,改变了 fail-fast 语义,值得二次确认。
xijiaat:在 B300(DeepSeek-V4-Flash、v0.26.0)上独立复现 #45785 并端到端验证本修复:pre-fix 命中 0,post-fix 命中 512;并将验证固化为五组往返回归测试提交到 PR 分支(ivanium/vllm#52)。
Dao007forever:LGTM,建议后续在 coordinator 中加 TODO,尝试把 EAGLE 组单独拆分,以便后续清理。
mergify[bot]:合入前存在 merge conflicts,需要 rebase。
- EAGLE 传播修复的正确性验证 (correctness): 修复有效:pre-fix hit=0,post-fix hit=512。
- tail-truncation 从 assert 改为 walrus 跳过 (design): 维护者批准合入,该变化被判定为 benign。
- 后续将 EAGLE 组单独拆分 (design): 作为 TODO 记录,不在本 PR 处理。
- Mergify 冲突需 rebase (other): 合入前需 rebase 解决。
风险与影响
- 风险:核心缓存命中路径变更:
find_longest_cache_hit 与 store_mask 是外部存储读写的关键路径,eagle_group_ids 语义变化会同时影响 save 与 lookup 两侧的全部缓存组,需确认没有其他调用方依赖旧的“仅标注组”集合语义。失败语义变化:_find_hit_blocks 尾部截断从 assert full_blks is not None 改为 walrus 条件跳过,若未来出现 full-attn 组未在当轮填充的情况,会把显式异常变成静默错误,增大排查难度(claude[bot] 已提示)。与 #48361 冲突:两个 PR 修改同一 Mooncake lookup 循环附近代码,合入顺序不当可能导致回归,需协调 rebase。验证覆盖:合并时 PR 仍为 draft,只有单元测试;xijiaat 提供了真实 B300 部署验证与测试补充,但该验证基于 v0.26.0 手动应用补丁,正式合入前建议用完整测试套件再确认一次。
- 影响:用户与业务:DeepSeek-V4 风格 SWA+MTP 的 Mooncake 外部存储用户从“前缀命中 0”恢复为可用命中,显著降低跨机重复 prefill 开销;非 EAGLE 配置行为不变,不改变模型输出。系统:仅影响 Mooncake store coordinator 的 mask 语义,不涉及 kernel 或模型权重路径;修复同时消除外部存储侧与核心 coordinator 的分组语义漂移,降低后续维护成本。团队:需与 #48361 保持合入协调;新增的 DSV4 五组往返测试为同模块后续改动提供了可依赖的回归基线。
- 风险标记:核心缓存命中路径变更, 失败语义从 assert 改为静默跳过, 与 #48361 需协调合入, 待真实模型 E2E 验证
关联脉络
- PR #50498 (feat): optionally disable lookup on PD decode: 同属 Mooncake store 外部缓存命中功能线,改动相邻模块的 store worker/scheduler/connector,可对照观察 store mask 与 lookup 逻辑的演进;本 PR 让被 enable_lookup 关闭的 lookup 语义在 EAGLE 场景下也能正确收敛。
参与讨论