执行摘要

• 一句话：添加 NPU 上 MoE 权重更新时形状恢复逻辑
• 推荐动作：建议阅读此 PR 的 patch，理解 NPU 上 MoE 权重布局的特殊处理。对于维护 NPU 后端的开发者，这是一个重要的参考。

功能与动机

For RL feature update_weights_from_disk/tensor/distributed on MoE models like Qwen3-MoE, the process_weights_after_loading method of UnquantizedFusedMoEMethod changes the param shape for MoE weights w13 and w2. To successfully update weights and maintain accuracy, we need to restore the canonical shapes before loading.

实现拆解

1. 导入 is_npu 并定义全局变量：在 python/sglang/srt/layers/moe/fused_moe_triton/layer.py 中导入 is_npu 函数，并定义 _is_npu = is_npu()，用于判断当前设备是否为 NPU。

2. 在 _weight_loader_impl 中增加 NPU 专有逻辑：在 UnquantizedFusedMoEMethod 分支内，当 _is_npu 为真时，对 w2_weight 和 w13_weight 分别检查参数形状是否与加载权重的形状匹配。如果不匹配，则对 param.data 进行 transpose(1,2).contiguous() 转置，使其恢复为加载前的形状，从而保证后续 maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load 能正确执行。

3. 无额外测试或配置变更：本次只修改了核心源码文件，未添加测试或配置，但通过已有的 CI 流程验证。

关键文件：

• python/sglang/srt/layers/moe/fused_moe_triton/layer.py（模块 MoE 层；类别 source；类型 core-logic；符号 is_npu, _is_npu, _weight_loader_impl）: 核心权重加载器，添加了 NPU 专属的转置逻辑，确保 update_weights 时形状正确。

关键符号：_weight_loader_impl

关键源码片段

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

核心权重加载器，添加了 NPU 专属的转置逻辑，确保 update_weights 时形状正确。

# python/sglang/srt/layers/moe/fused_moe_triton/layer.py
​
# 全局变量定义
_is_npu = is_npu() # NPU 检测，控制后续转置逻辑
​
# _weight_loader_impl 方法中，UnquantizedFusedMoEMethod 分支
if isinstance(method, UnquantizedFusedMoEMethod):
    if _is_npu:
        # 处理 w2_weight：如果 param 形状不匹配 loaded_weight，转置
        if weight_name.endswith(".experts.w2_weight"):
            if param.data.shape[1] != loaded_weight.shape[0]:
                param.data = param.data.transpose(1, 2).contiguous()
        # 处理 w13_weight：类似检查并转置
        if weight_name.endswith(".experts.w13_weight"):
            if param.data.shape[2] != loaded_weight.shape[1]:
                param.data = param.data.transpose(1, 2).contiguous()
    # 通用恢复逻辑
    method.maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load(
        layer=self,
        param=param,
        weight_name=weight_name,
    )

评论区精华

只有一个 reviewer（gemini-code-assist[bot]）提出了三点改进建议：

• PEP 8 空格修复（weight_name.endswith("...") : 多余空格）；
• 使用 elif 代替独立的 if 以提高可读性；
• 使用预期规范形状（从 layer 属性推导）代替 loaded_weight 维度进行形状比较，以提高鲁棒性。

作者拒绝了这些建议（Refuseed），未作修改。PR 最终由 sglang-npu-bot 批准合并。

• PEP 8 空格与鲁棒性改进建议 (style): 作者拒绝（Refuseed），未修改。PR 合并时未采纳。

风险与影响

• 风险：风险较低。变更仅在 _is_npu 为真且使用了 UnquantizedFusedMoEMethod 时生效，不影响 CUDA 或其他后端。唯一潜在风险是 transpose 操作如果形状检查不准确可能导致维度错误，但由于条件判断基于与 loaded_weight 的比较，实际运行中应安全。缺少单元测试，但集成测试可通过 RL update_weights 场景覆盖。
• 影响：影响范围：仅限 NPU 后端 + MoE 模型 + 使用 UnquantizedFusedMoEMethod 的 RL 训练更新权重场景。用户无感知功能变更，但修复后 RL update_weights 在 NPU 上的精度恢复正常。对系统团队是关键的兼容性修复。
• 风险标记：NPU 特定代码可能影响其他后端, 缺少测试覆盖

关联脉络

• 暂无明显关联 PR