执行摘要

• 一句话：修复 Waterfill 与动态 EPLB 在 fused shared experts 下的兼容性
• 推荐动作：此 PR 是 DeepEP+EPLB 兼容路径的关键修复，维护者应快速合入。代码设计清晰（通过分离 recorder ID 避免统计污染），可作为处理类似混合专家 ID 空间的参考实现。

功能与动机

PR body 明确指出需要修复 fused shared expert MoE weight views 被在线 EPLB 更新错误修改，以及确保 EPLB 统计仅跟踪逻辑路由专家 ID。MMLU 测试显示不加修复时可能存在 NCCL 超时。

实现拆解

1. 修正权重视图 (python/sglang/srt/models/deepseek_v2.py)：在 get_moe_weights 中对返回的权重张量进行切片，只取前 num_local_experts - num_fused_shared_experts 行，避免 EPLB 在线更新影响到 fused shared expert 的权重。
2. 分离 recorder topk_ids (python/sglang/srt/layers/moe/topk.py)：修改 _post_process_topk_ids 返回值，新增 recorder_topk_ids。当使用 DeepEP 后端且 fused shared experts > 0 时，recorder_topk_ids 只包含经过 EPLB 重映射的 routed 列，排除后续 DeepEP 插入的 shared slot。其他情况下与 topk_ids 保持一致。
3. 变更调用点 (python/sglang/srt/layers/moe/topk.py:select_experts)：将 on_select_experts 的输入从原始的 topk_ids 改为新的 recorder_topk_ids，确保 EPLB 统计基于正确的逻辑路由 ID。
4. 单元测试 (test/registered/unit/eplb/test_deepep_waterfill_eplb.py)：新增 4 个测试方法，覆盖 fused shared expert 排除、保持完整形状、DeepEP 与非 DeepEP 后端下的 recorder ID 正确性。

关键文件：

• python/sglang/srt/layers/moe/topk.py（模块 路由逻辑；类别 source；类型 core-logic；符号 _post_process_topk_ids, select_experts）: 核心逻辑变更：修改 _post_process_topk_ids 返回类型并分离 recorder_topk_ids，变更调用 select_experts 以使用正确的统计 ID。
• python/sglang/srt/models/deepseek_v2.py（模块 模型定义；类别 source；类型 data-contract；符号 DeepseekV2MoE.get_moe_weights）: 权重视图修复：get_moe_weights 切片排除 fused shared expert，避免 EPLB 在线更新破坏共享专家权重。
• test/registered/unit/eplb/test_deepep_waterfill_eplb.py（模块 测试；类别 test；类型 test-coverage；符号 _FakeExpertParam, TestDeepEPWaterfillEPLB, test_deepseek_moe_get_moe_weights_excludes_fused_shared_slot, test_deepseek_moe_get_moe_weights_keeps_full_shape_without_fusion）: 新增完整单元测试套件，验证 get_moe_weights 切片正确性以及 recorder_topk_ids 在不同后端下的行为。

关键符号：DeepseekV2MoE.get_moe_weights, _post_process_topk_ids, select_experts

关键源码片段

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

核心逻辑变更：修改 _post_process_topk_ids 返回类型并分离 recorder_topk_ids，变更调用 select_experts 以使用正确的统计 ID。

# 关键变更：_post_process_topk_ids 新增 recorder_topk_ids 返回值
# 用于 EPLB 统计时排除 fused shared expert 的 slot
​
def _post_process_topk_ids(
    topk_ids: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_config: TopKConfig,
    router_logits: torch.Tensor,
    layer_id: int,
    num_token_non_padded: Optional[torch.Tensor] = None,
    expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    # ... 省略无关部分
    # 声明 recorder 变量
    recorder_topk_ids = None
    if _is_cuda:
        if num_fused_shared_experts > 0 and is_deepep_class_backend():
            shared_cols = topk_ids[:, -num_fused_shared_experts:]
            routed_cols = topk_ids[:, :-num_fused_shared_experts]
            # EPLB 仅对 routed 列进行重映射
            routed_cols = _biased_grouped_topk_postprocess(
                routed_cols, expert_location_dispatch_info, num_token_non_padded
            )
            topk_ids = torch.cat([routed_cols, shared_cols], dim=-1)
            # recorder 只保留重映射后的 routed 列，排除 shared slot
            recorder_topk_ids = routed_cols
        else:
            topk_ids = _biased_grouped_topk_postprocess(
                topk_ids, expert_location_dispatch_info, num_token_non_padded
            )
    if recorder_topk_ids is None:
        recorder_topk_ids = topk_ids
    # ... 后续 DeepEP remap 仅修改 topk_ids，不影响 recorder_topk_ids
    return topk_ids, topk_weights, recorder_topk_ids
​
# 调用点调整：
def select_experts(...) -> StandardTopKOutput:
    # ...
    topk_ids, topk_weights, recorder_topk_ids = _post_process_topk_ids(...)
    # 使用 recorder_topk_ids 传入统计，避免 fused shared slot 污染负载均衡数据
    get_global_expert_distribution_recorder().on_select_experts(
        topk_ids=recorder_topk_ids
    )

python/sglang/srt/models/deepseek_v2.py

权重视图修复：get_moe_weights 切片排除 fused shared expert，避免 EPLB 在线更新破坏共享专家权重。

# DeepseekV2MoE.get_moe_weights 修改：只返回 routed expert 部分的权重
def get_moe_weights(self):
    # EPLB 仅重平衡物理路由专家，fused shared expert 必须保持不变
    num_local_experts_for_eplb = (
        self.experts.num_local_experts - self.num_fused_shared_experts
    )
    return [
        x.data[:num_local_experts_for_eplb] # 切片丢弃 fused shared 参数
        for name, x in self.experts.named_parameters()
        if name not in ["correction_bias"]
        and filter_moe_weight_param_global_expert(name, x, self.experts.num_local_experts)
    ]

test/registered/unit/eplb/test_deepep_waterfill_eplb.py

新增完整单元测试套件，验证 get_moe_weights 切片正确性以及 recorder_topk_ids 在不同后端下的行为。

# 测试 fused shared expert 情况下 get_moe_weights 返回切片后的权重
class TestDeepEPWaterfillEPLB(CustomTestCase):
    def test_deepseek_moe_get_moe_weights_excludes_fused_shared_slot(self):
        experts = _FakeExpertParam() # 5 个本地 expert，权重 shape (5,2)
        moe = SimpleNamespace(num_fused_shared_experts=1, experts=experts)
        shared_before = experts.weight.data[-1].clone() # 保存 fused shared 权重
​
        weights = DeepseekV2MoE.get_moe_weights(moe)
        # 应只返回 routed 部分：5-1=4 个 expert
        self.assertEqual(len(weights), 1)
        self.assertEqual(weights[0].shape, (4, 2))
​
        # 修改返回权重的最后一行为零，验证不影响原始 fused shared 权重
        weights[0][-1].zero_()
        self.assertTrue(torch.equal(experts.weight.data[-2], torch.zeros(2)))
        self.assertTrue(torch.equal(experts.weight.data[-1], shared_before))
​
    # ... 其他测试方法类似，验证 recorder_topk_ids 的正确性

评论区精华

仅有 gemini-code-assist[bot] 自动代码审查，无人工讨论。审阅者 ch-wan 直接批准，未提出异议。

• 代码审查自动评论 (other): 无需处理，bot 评论仅为信息性。

风险与影响

• 风险：
  1. 回归风险低：变更仅影响 DeepEP 后端 + fused shared experts + EPLB 的组合路径，非该路径的行为不受影响。
  2. 测试覆盖风险：测试全部在 CPU 上通过 mock 和 patching 执行，未在真实 GPU 上验证 DeepEP 调度流程，可能遗漏运行时内存/布局问题。
  3. 兼容性：_post_process_topk_ids 返回类型从 2 元组改为 3 元组，所有调用点均已更新（仅 select_experts 一处），无兼容断裂。
  4. 性能影响：增加了一次切片和广播操作，开销可忽略。
     - 影响：直接使 DeepEP Waterfill 用户能够安全启用动态 EPLB，修复了在线权重更新可能导致 fused shared expert 偏移或统计错误的缺陷。影响范围限于使用 --enable-ep-moe 且开启 num_fused_shared_experts > 0 的 DeepSeek-V2 类模型。间接为未来 EPLB 统一统计逻辑奠定基础。
     - 风险标记：Fused shared experts 路径, 测试无 GPU 端集成验证, 返回类型变更需确保调用点更新

关联脉络

• PR #27329 [LoRA] Experimental fast LoRA path with experimental_sgl_trtllm MoE backend for FP8 and NVFP4 models: 同为 MoE 性能相关 PR，涉及 fast LoRA 和 trtllm 后端，但与本 PR 无直接依赖关系。
• PR #27166 Reland "Support NextN = 2/4 in DSV32": 同为 DeepEP 后端改进，涉及 DSV32 和 DG 原生路径，与本 PR 在 DeepEP 动态 EPLB 上有潜在关联。