执行摘要

• 一句话：修复Mooncake connector中Mamba状态跨TP复制错误
• 推荐动作：建议所有使用Mooncake connector的用户升级，尤其是混合模型场景；开发人员可学习per-group replication factor的设计模式。

功能与动机

PR描述指出混合模型中Mamba状态在TP>1时被错误去重：由于其状态是按head/dim分片的，而Mooncake connector假设所有rank持有相同KV字节，导致每个rank只存储部分边界块，加载时加载到错误数据。这造成了严重的精度退化，如gsm8k warm accuracy从0.87降至0.72。

实现拆解

1. 导入新接口：在worker.py中导入UniformTypeKVCacheSpecs、MambaSpec、MLAAttentionSpec等，用于正确识别spec类型。
2. 添加per-group复制因子计算：_compute_group_tp_replication_factors方法遍历每个group的spec，递归展开UniformTypeKVCacheSpecs，根据spec类型返回复制因子（Mamba=1，MLA=tp_size，GQA=tp_size/num_kv_head）。
3. 替换model-wide属性：移除原来的put_step、head_or_tp_rank等全局变量，在MooncakeStoreWorker.__init__中计算per-group复制因子，并调用_init_lookup_key_prefixes按组初始化key命名空间。
4. 调整发送线程：KVCacheStoreSendingThread接受group_put_steps序列，存储时每个组使用各自的put_step，而不是单个值。
5. 修改lookup逻辑：将exists检查从全局的ranks_per_candidate改为per-group计算，每个组的key前缀不同。
6. 更新测试：新增6个单元测试覆盖Mamba分片组、混合组场景，并调整已有测试适配新接口。

关键文件：

• vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py（模块 KV连接器；类别 source；类型 core-logic；符号 _init_lookup_key_prefixes, _spec_tp_replication_factor, _compute_group_tp_replication_factors, rank_namespaces）: 核心实现文件，将复制策略从model-wide改为per-group，添加了组识别、复制因子计算和key命名空间调整。
• tests/v1/kv_connector/unit/test_mooncake_store_worker.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_tp_sharded_group_saves_every_block_on_every_rank, _refresh_group_tp_replication_factors, test_lookup_key_prefixes_expand_tp_sharded_groups_per_rank, test_group_tp_replication_factors_mixed_mla_gqa_mamba）: 新增6个测试用例，覆盖Mamba分片组、混合组复制因子、lookup边界检查等场景。
• tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py（模块 集成测试；类别 test；类型 test-coverage）: e2e测试中调整参数名（put_step → group_put_steps），适配新接口。

关键符号：_compute_group_tp_replication_factors, _spec_tp_replication_factor, _init_lookup_key_prefixes, rank_namespaces

关键源码片段

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

新增6个测试用例，覆盖Mamba分片组、混合组复制因子、lookup边界检查等场景。

# 测试 Mamba 分片组时每个 rank 必须保存所有块
def test_tp_sharded_group_saves_every_block_on_every_rank():
    """Sharded ranks must write every block because peers hold different bytes."""
    store = MagicMock()
    store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
    store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys)
    thread = _make_store_sending_thread(store, tp_rank=0, put_step=2)
    # 覆盖 group_put_steps 为 1，模拟 Mamba 组的无去重行为
    thread.group_put_steps = [1]
​
    thread.add_stored_request("req-a")
    thread._handle_request(
        ReqMeta(
            req_id="req-a",
            token_len_chunk=64,
            block_ids=([0, 1, 2, 3],),
            block_hashes=[b"a0", b"a1", b"a2", b"a3"],
            can_save=True,
        )
    )
​
    keys = store.batch_is_exist.call_args.args[0]
    # 预期保存所有 4 个块（不跨 rank 去重）
    assert len(keys) == 4

评论区精华

• GirasoleY 建议将复制因子缓存为类变量，作者在后续commit中实现。
• Dao007forever 对_spec_tp_replication_factor中UniformTypeKVCacheSpecs处理提供简化建议（假设只有一层），作者采纳并重构。

• Dao007forever 建议类型注解使用Sequence[int]，作者权衡后保留Sequence以保持通用性。

• 简化 UniformTypeKVCacheSpecs 处理逻辑 (design): 按建议简化，使用 flat 的 any-Mamba / all-MLA / GQA 检查。

• 类型注解 Sequence vs tuple (style): 保留 Sequence 作为输入参数类型。
• 缓存复制因子为类变量 (performance): 实现缓存。

风险与影响

• 风险：核心变更可能影响MLA/GQA的dedup路径，但回归测试覆盖了DeepSeek-V4-Flash（fp8 KV），结果显示warm精度不低于cold，字节一致性验证通过；Mamba组引入潜在风险：如果group spec类型识别错误，可能导致复制因子错误，但测试覆盖混合场景。
• 影响：影响使用Mooncake connector的混合模型用户，修复了严重精度bug；纯MLA/GQA模型无影响（回归通过）；系统层面改进了架构灵活性，允许多种组类型共存。
• 风险标记：核心路径变更, 影响MLA/GQA dedup, 回归测试覆盖DeepSeek

关联脉络

• 暂无明显关联 PR