Prhub

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

原始 PR 作者 ivanium 合并时间 2026-08-03 10:32 文件变更 2 提交数 1 评论 5 代码增减 +118 / -39

执行摘要

修复 Mooncake 合并组 EAGLE 位传播,恢复外部前缀缓存命中

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 部署上复现了该失败。

值得精读。核心看点是“避免镜像实现漂移”——直接复用核心 SpecGroup 并在合并组上传播 EAGLE 位,使外部存储与核心缓存协调器保持同一分组语义;测试的 store-to-lookup 往返设计(先按 store mask 填 exists 集,再走对外查找)也值得借鉴。合入前请重点复核 walrus 更改,并协调 #48361 的合入顺序;建议合入后用真实 DeepSeek-V4 MTP P/D 部署跑一次前缀命中率回归。

讨论亮点

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。

实现拆解

  1. 变更入口vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.pyMooncakeStoreCoordinator。此前它是核心 HybridKVCacheCoordinator 的独立镜像,用本地元组表示分组,eagle_group_ids 直接从原始 is_eagle_group 标注计算;本 PR 将其改为复用核心的 SpecGroup 表示。
  2. 合并与传播_verify_and_split_kv_cache_groups 中,等 spec 组归并时若新成员是 EAGLE 组,用 _replace(use_eagle=True) 提升整个合并组;use_eagle=True 且无任何标注时保守 fallback 全开。eagle_group_ids 改为从合并后的 SpecGroup.use_eagle 展开到全部成员组 ID,保证 store_masklookup_mask 的 save/lookup 语义一致。
  3. 命中逻辑对齐_find_hit_blocks 的单组与多组路径中 drop_eagle_block 从“基于下标集合 eagle_attn_group_indices”改为“apply_eagle and group_eagle”,使每个成员组都执行相同 EAGLE 尾块丢弃;结尾 full-attn 块截断处由 assert 改为 walrus 条件跳过(行为宽松化)。
  4. 测试配套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)。
  5. 配置与部署:无新增配置项或部署改动;合入前需 rebase 解决与最新 coordinator 代码的冲突。
文件 模块 状态 重要度
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py 存储协调器 modified 7.05
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py 单元测试 modified 6.39

关键符号

_verify_and_split_kv_cache_groups _find_hit_blocks

关键源码片段

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py core-logic

核心修复文件:将本地分组元组替换为核心 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 test-coverage

新增两个回归测试,覆盖三组与 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

评论区精华

EAGLE 传播修复的正确性验证 正确性

claude[bot] 逐行对比核心 HybridKVCacheCoordinator.verify_and_split_kv_cache_groups,确认 SpecGroup 语义与 fallback 行为一致;xijiaat 在 B300 上用 DeepSeek-V4-Flash 实测复现 #45785 并验证修复。

结论:修复有效:pre-fix hit=0,post-fix hit=512。 · 已解决

tail-truncation 从 assert 改为 walrus 跳过 设计

claude[bot] 注意到 _find_hit_blocks 末尾截断循环中 assert full_blks is not None 被放宽为条件跳过,认为该组 ID 总在同一个迭代中被填充,行为看似 benign,但改变了 fail-fast 语义,值得二次确认。

结论:维护者批准合入,该变化被判定为 benign。 · 已解决

后续将 EAGLE 组单独拆分 设计

Dao007forever 建议后续在 coordinator 中尝试把 EAGLE 组独立拆分,以便清理当前传播逻辑。

结论:作为 TODO 记录,不在本 PR 处理。 · 待处理

Mergify 冲突需 rebase other

mergify[bot] 提示 PR 与最新 coordinator 代码存在 merge conflicts。

结论:合入前需 rebase 解决。 · 已解决

风险与影响

核心缓存命中路径变更:find_longest_cache_hitstore_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 验证

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论