执行摘要

• 一句话：将 expert_mask_gpu 所有权从 FusedMoE 层迁移到 StandardDispatcher
• 推荐动作：值得精读以学习如何识别和修复抽象泄漏。设计原则清晰，适合作为代码重构的范例。

功能与动机

expert_mask_gpu 的计算仅依赖于 dispatcher 的 local_expert_mapping 和 num_local_experts，但原先由 FusedMoE 层在 forward_impl 中计算，属于抽象泄漏。将其放入 dispatcher 中更符合职责分离原则，也便于未来 dispatcher 后端扩展时维护。

实现拆解

实现分为两步：

1. 在 StandardDispatcher.__init__() 中添加 self.expert_mask_gpu = None 和 self.num_local_experts = moe_runner_config.num_local_experts；在 StandardDispatcher.dispatch() 中，当启用 aiter（_use_aiter 为真）且 local_expert_mapping 已初始化时，计算 expert_mask_gpu。
2. 从 FusedMoE 层中移除 self.expert_mask_gpu = None 初始化及 forward_impl 中的计算块。更新五个量化文件（fp8.py, mxfp4.py, quark_w4a4_mxfp4_moe.py, unquant.py）中所有引用 layer.expert_mask_gpu 的地方为 layer.dispatcher.expert_mask_gpu。

无测试变更——纯重构，语义保持不变。

关键文件：

• python/sglang/srt/layers/moe/token_dispatcher/standard.py（模块 分发器；类别 source；类型 core-logic）: 核心变更：添加 expert_mask_gpu 属性和计算逻辑，并存储 num_local_experts。
• python/sglang/srt/layers/moe/fused_moe_triton/layer.py（模块 MoE层；类别 source；类型 core-logic）: 从 FusedMoE 层移除 expert_mask_gpu 的初始化和计算，精简代码。
• python/sglang/srt/layers/quantization/fp8.py（模块 量化层；类别 source；类型 core-logic）: 修改 aiter 路径下 FP8 量化 MoE 内核的 expert_mask 参数来源。

关键符号：未识别

关键源码片段

python/sglang/srt/layers/moe/token_dispatcher/standard.py

核心变更：添加 expert_mask_gpu 属性和计算逻辑，并存储 num_local_experts。

# standard.py
class StandardDispatcher(BaseDispatcher):
    def __init__(self, moe_runner_config: MoeRunnerConfig):
        # ... 其他初始化 ...
        self.num_local_experts = moe_runner_config.num_local_experts # 新增：保存 local expert 数量
        self.local_expert_mapping = None
        self.expert_mask_gpu = None # 新增：初始化属性
​
    def dispatch(self, hidden_states, topk_output):
        # ... 前面的 FP4 all-gather 和 local_expert_mapping 创建 ...
        if self.local_expert_mapping is not None:
            if _use_aiter:
                # 将 expert_mask 的计算从 FusedMoE 层移动到这里
                self.expert_mask_gpu = (
                    (self.local_expert_mapping >= 0) & 
                    (self.local_expert_mapping < self.num_local_experts)
                ).to(torch.int32).to(device="cuda")
            else:
                # 非 aiter 路径：直接使用 local_expert_mapping 重映射 topk_ids
                topk_output = topk_output._replace(
                    topk_ids=self.local_expert_mapping[topk_output.topk_ids]
                )
        return StandardDispatchOutput(...)

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

从 FusedMoE 层移除 expert_mask_gpu 的初始化和计算，精简代码。

# layer.py
class FusedMoE(torch.nn.Module):
    def __init__(self, ...):
        # ... 原有初始化 ...
        # 删除 : self.expert_mask_gpu = None
​
    def forward_impl(self, hidden_states, topk_output):
        # ...
        dispatch_output = self.dispatcher.dispatch(hidden_states, topk_output)
        # 删除以下完整块，因为 expert_mask 已在 dispatcher.dispatch() 中计算
        # if _use_aiter and self.dispatcher.local_expert_mapping is not None:
        # self.expert_mask_gpu = ...
        combine_input = self.run_moe_core(dispatch_output)
        # ...

评论区精华

PR 没有引发 review 讨论，提交者自行合并。从 commit message 可见动机明确，变更直白。

• 暂无高价值评论线程

风险与影响

• 风险：风险极低。这是一次纯重构，所有计算逻辑和数据流不变，仅改变张量的拥有者和访问路径。需要确认所有调用点已更新（已检查 5 处）。潜在风险：如果未来有其他后端或自定义 dispatcher 未继承 StandardDispatcher，可能需要类似调整，但这不属于本 PR 范围。
• 影响：对用户无影响，对系统行为无影响。对开发者而言，expert_mask_gpu 现在属于 dispatcher，使得 MoE 层更干净，dispatcher 职责更完整。影响范围限定在 aiter 路径下的量化 MoE 内核调用。团队内无需额外沟通，改动小且自包含。
• 风险标记：纯重构，风险低

关联脉络

• 暂无明显关联 PR