执行摘要

• 一句话：修复 DeepSeek-V4 MegaMoE 和 FakeKVManager 两处启动崩溃
• 推荐动作：建议合并。修改简洁、定位准确，已在 8×B200 上验证。MegaMoE 入口统一类型转换的设计比修复上游各路径更健壮。

功能与动机

当使用 SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE=1 和 --disaggregation-transfer-backend fake --ep-dispatch-algorithm fake 启动 DeepSeek-V4-Pro 时，服务器在初始化期间因两个独立错误崩溃：TVM 内核的 dtype 不匹配错误和 FakeKVManager 缺少 kv_args 属性的 AttributeError。

实现拆解

1. 修复 MegaMoE 的 topk_ids dtype 不匹配：在 python/sglang/srt/layers/moe/mega_moe.py 的 _run_mega_routed 函数中，将 topk_ids 和 topk_weights 在传入 mega_moe_pre_dispatch 前显式转换为 int32 和 float32。这确保 JIT 内核的输入类型契约得到满足，并且与空 token 分支的类型一致。
2. 修复 FakeKVManager 缺失 kv_args 属性：在 python/sglang/srt/disaggregation/fake/conn.py 的 FakeKVManager.__init__ 中添加 self.kv_args = args，与 CommonKVManager 的模式保持一致，供解聚解码事件循环访问 self.kv_manager.kv_args.state_types。

关键文件：

• python/sglang/srt/layers/moe/mega_moe.py（模块 MoE 层；类别 source；类型 core-logic；符号 _run_mega_routed）: 修复 MegaMoE 预分派时 topk_ids 类型不匹配的核心逻辑，确保 JIT 内核输入类型正确。
• python/sglang/srt/disaggregation/fake/conn.py（模块 解聚层；类别 source；类型 core-logic；符号 FakeKVManager）: 修复 FakeKVManager 未存储 kv_args 属性的问题，对齐 CommonKVManager 行为。

关键符号：_run_mega_routed

关键源码片段

python/sglang/srt/layers/moe/mega_moe.py

修复 MegaMoE 预分派时 topk_ids 类型不匹配的核心逻辑，确保 JIT 内核输入类型正确。

def _run_mega_routed(...):
    # ... 前面的代码保持不变 ...
    if num_tokens > 0:
        # 显式转换为 int32 / float32，确保与 JIT 内核（mega_moe_pre_dispatch）期望的类型一致
        topk_ids_in = topk_ids.to(torch.int32)
        topk_weights_in = topk_weights.to(torch.float32)
    else:
        # 空 token 分支已使用 int32 / float32，保持一致
        topk_ids_in = hidden_states.new_empty((0, top_k), dtype=torch.int32)
        topk_weights_in = hidden_states.new_empty((0, top_k), dtype=torch.float32)
​
    if use_fp4_acts:
        deep_gemm.mega_moe_pre_dispatch(
            hidden_states, topk_ids_in, topk_weights_in, ..., use_fp4_acts=True,
        )
    else:
        mega_moe_pre_dispatch(
            hidden_states, topk_ids_in, topk_weights_in, ..., quant_group_size=32,
        )

评论区精华

Review 中 gemini-code-assist[bot] 建议同时将 topk_weights 转换为 float32 以保持与空分支一致，该建议被采纳。ShangmingCai 确认了 fake 后端的改动，但指出自己不是 MegaMoE 专家，需要其他 reviewer 审核。JoeLee314 将提交压缩为一个后请求重新审批，ShangmingCai 表示需等待其他 reviewer 对 MegaMoE 部分的审查。

• topk_weights 也应转换为 float32 (correctness): 该建议被采纳，最终补丁同时包含 topk_ids 和 topk_weights 的类型转换。
• MegaMoE 部分需其他 reviewer 审查 (other): 最终由 Fridge003 批准合并。

风险与影响

• 风险：
  • dtype 转换安全：int64→int32 转换对于专家索引值（通常远小于 2^31）是安全的；float32 转换可能损失 float16/bfloat16 权重精度，但 JIT 内核已期望 float32，且空分支已使用 float32，实际一致性更好。
  • 回归风险：两处修改均位于特定配置路径（MegaMoE + fake 解聚），不影响默认路径。编译器和运行时测试已验证启动成功。
  • 无性能影响：.to(int32) 在 GPU 上为无操作（当值已适合 int32 时），不影响前向计算。
• 影响：
  • 用户影响：修复了 DeepSeek-V4-Pro 在启用 MegaMoE 和 fake 解聚模式时的启动崩溃，用户可正常使用该配置。
  • 系统影响：无，修改仅影响启动路径，不改变运行时行为。
  • 团队影响：低，单行改动易于理解和维护。
  • 风险标记：特定配置路径, 缺少测试覆盖

关联脉络

• 暂无明显关联 PR