Feat/support sparse cp reuse cache

[MMadhatter]

based on #780

[LLLLKKKK]

🤖 AI Code Review

概述

本 PR 为 Sparse MLA FP8 路径添加 Context Parallel (CP) 支持，实现 prefill 阶段的 CP + KV cache 复用。主要新增 SparseMlaFp8CPOp、SparseMlaCpImpl，以及 CP 感知的 indexer 路径和动态 CP 激活逻辑。

优点

• CP 感知的 indexer 路径设计合理，apply_rope_and_rotate_q_k_cp、quant_q_k_cp、_get_topk_ragged_cp 职责清晰
• WriteCacheStoreOp 创建逻辑重构到 common.py 共享 helper，减少重复
• 测试覆盖了基本的 CP op 和 topk ragged CP 场景

建议改进

P0 - Bug

1. Pickle 向后不兼容 (ConfigInit.cc): ParallelismConfig pickle tuple 从 16 元素增加到 17（新增 use_ub_comm）。已有的序列化对象反序列化时会直接抛 "Invalid state!"。建议兼容处理：

   if (t.size() != 16 && t.size() != 17) throw std::runtime_error("Invalid state!");
   if (t.size() > 16) c.use_ub_comm = t[16].cast<bool>();

2. Metadata 清空操作作用在错误对象上 (flashmla_sparse_cp_impl.py): metadata = meta.tile_scheduler_metadata 后对 metadata 设置 .tile_scheduler_metadata = None 和 .num_splits = None，实际修改的是子对象而非 meta。self._fp8_kernel_metadata_q0 仍持有旧 metadata，layer_id == 0 的清理无效，layer_id > 0 使用过期 metadata。应改为 meta.tile_scheduler_metadata = None; meta.num_splits = None。

P1 - 重要

1. init_user_buffers_environment 移除了关键安全检查 (collective_torch.py): 旧代码检查 torch.distributed.is_initialized()、is_enabled() 和 is_cuda()，新代码只检查 use_ub_comm。建议恢复 distributed 初始化和 CUDA 设备检查。

2. num_heads 在 CP 启用时不除以 tp_size (mla_attention.py): 构造时根据 CP 配置决定是否除以 tp_size，但该模块同时用于 prefill 和 decode。Decode 阶段 CP 不活跃，使用完整 head_num 会导致 shape 不匹配。建议按 forward pass 类型动态决定。

3. kv_len vs total_kv_len 属性不一致 (indexer_op.py): quant_q_k_cp 设置 self.kv_len，但 _get_topk_ragged_cp 使用 self.total_kv_len（外部赋值）。调用顺序变化或外部赋值缺失会导致断言失败。建议统一属性。

4. 不相关的 MoE 修复混入本 PR (config_resolver.py): config.tp_size（考虑 CP 的 get_attn_tp_size()）与 config.parallelism_config.tp_size 语义不同，这是对 MoE routing 的行为变更，建议拆分到独立 PR。

P2 - 建议

1. attn_output is None 时静默替换为零张量 (mla_attention.py)，可能掩盖 attention 失败。建议加 debug log。
2. q_nope.shape[0] == 0 时返回未初始化的 torch.empty 张量 (flashmla_sparse_impl.py)，建议改用 torch.zeros。
3. RoPE 逻辑在 flashmla_sparse_cp_impl.py 和 indexer_op.py 中重复，建议提取共享工具函数。
4. barrier 被 import 但未使用。

总结

P0-1（pickle 兼容性）和 P0-2（metadata bug）需在合并前修复。P1 问题建议一并处理或明确 acknowledge。

[LLLLKKKK]

🤖 AI Code Review — PR #792 Feat/support sparse cp reuse cache

概述

为 Sparse MLA (DeepSeek-V3 DSA) 添加 Context Parallel (CP) prefill 支持。核心思路：all-gather KV 到所有 rank，restore 到逻辑顺序写入 KV cache，每个 rank 只处理自己的 Q 子集。同时引入 use_ub_comm 配置项控制 user-buffer 通信器初始化。规模：22 files, +1727/-147。

优点

• CP 路径与非 CP 路径分离清晰，通过 support_prefill_cp() / support_parallelism_config() 做能力过滤
• 测试覆盖较好：3 个新测试文件，包含 tp_size=1 下 CP vs non-CP 输出一致性验证
• plan() 中做了 bounds check 并提前报错，避免 forward 时 OOB
• write_cache_store 逻辑提取到 common.py，减少重复代码

建议改进

P0 - Bug

1. Pickle 兼容性: __setstate__ 严格等于会破坏向后兼容
   ConfigInit.cc 第 989 行 if (t.size() != 17) — 旧版本序列化的 tuple (size=16) 反序列化时直接抛异常。
   建议: 改为 if (t.size() < 16) 做最低版本容错，新字段用 if (t.size() > 16) c.use_ub_comm = t[16].cast<bool>();

2. FMHAType.CP_SPARSE_FLASHMLA 未在 C++ enum 中定义
   flashmla_sparse_cp_impl.py 第 1165 行 return FMHAType.CP_SPARSE_FLASHMLA，但 ConfigModules.h 的 FMHAType enum 中没有这个值，运行时会抛 AttributeError。
   建议: 在 ConfigModules.h 中添加 CP_SPARSE_FLASHMLA 并在 pybind 中注册。

P1 - 重要

3. and False 硬编码禁用 fast_path — 调试残留？
   attn_factory.py 第 46 行 and False 强制禁用所有 prefill fast path。如果是有意为之（CP 不支持 fast path），应改为 and not parallelism_config.prefill_cp_config.is_enabled() 并加注释。

4. config_resolver.py 语义变更影响 MoE 路由
   is_tp_equal_ep 从 config.parallelism_config.tp_size 改为 config.tp_size（即 get_attn_tp_size()）。CP 启用时 get_attn_tp_size() 返回 1，会导致 is_tp_equal_ep 返回 1 == ep_size，可能错误改变 MoE all-gather/reduce-scatter 策略。请确认是否有意为之。

5. init_user_buffers_environment 移除了安全检查
   原代码检查 torch.distributed.is_initialized() 和 is_cuda()，新代码只检查 use_ub_comm。非 CUDA 环境或分布式未初始化时设置 use_ub_comm=True 会 crash。

6. 共享 attn_inputs.input_lengths 临时替换有竞态风险
   SparseMlaCpImpl.__init__ 和 prepare 中 swap attn_inputs.input_lengths，如果有其他消费者同时读取会得到错误值。建议浅拷贝 attn_inputs 后修改副本。

P2 - 建议

7. _apply_input_bmm 返回 None 时下游未完全处理 — q_nope.shape[0] == 0 时返回 None，但 fmha_impl.forward 中 q[self.total_local_ids] 可能对 None 操作。

8. _get_topk_ragged_cp docstring 说返回 (topk0, topk1) tuple，实际返回单个 tensor — 请更新 docstring。

9. plan() 中多次 .item() / .cpu().tolist() 触发 CUDA sync — per-request 级别，高 QPS 下有影响。

10. Python 3.9+ tuple[...] 与 Tuple[...] 混用 — indexer_op.py 和 flashmla_sparse_cp_impl.py 中不一致，建议统一。

总结

代码结构合理，CP 路径设计清晰，测试覆盖较好。主要阻塞项是 pickle 兼容性（P0 #1）和 FMHAType.CP_SPARSE_FLASHMLA 未定义（P0 #2），这两个会导致运行时错误。and False 硬编码和 config_resolver.py 语义变更也需要确认意图。建议修复 P0 后再合入。

[LLLLKKKK]

🤖 AI Code Review — PR #792

PR 概述

Title: Feat/support sparse cp reuse cache
Author: MMadhatter
规模: 24(GitHub) files, ~+2700/-100

核心目标

在 PR #780（Sparse MLA CP）基础上，为 Sparse MLA Context Parallel 添加 KV cache reuse 支持。

---

Review 意见

1. PR description 过于简略 [P2]
   仅 "based on Feat/support sparse cp #780"。应说明 reuse cache 在 CP 场景下的具体实现策略。

整体评价

在 #780 基础上扩展 reuse cache 支持，改动高度相关。建议与 #780 协调合入顺序（#780 先合入）。

✅ LGTM ready to ci