执行摘要

• 一句话：为AMD平台Qwen3.5 MoE模型优化共享专家追加逻辑，用单个Triton内核融合4次内核启动以提升路由性能。
• 推荐动作：该PR值得精读，特别是对于关注高性能计算和MoE模型优化的工程师。重点关注新增的Triton内核设计，它展示了如何将多个独立操作融合为单次启动以减少开销，同时保留逐token权重的精度要求。此外，注意其平台特定性（AMD/AITER），这反映了项目中对不同硬件后端的差异化优化策略。

功能与动机

PR body明确指出，_append_shared_to_topk_output函数位于MoE路由的关键路径上，原先使用4次独立内核启动（2次逐元素操作+2次拼接）来追加共享专家ID和权重，带来了不必要的启动开销。Qwen3.5模型需要从学习的shared_expert_gate中获取逐token的共享权重，因此优化必须保留这种逐token的权重精度。

实现拆解

1. 新增融合内核函数：在python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py中，新增Triton JIT内核_fused_append_shared_experts_with_weights_kernel及其包装函数fused_append_shared_experts_with_weights。该内核在一个启动中完成：复制原始topk_ids和topk_weights，追加共享专家ID（基于N_BASE计算），并加载和追加逐token的共享权重。包装函数负责处理输入张量的形状转换（如将1D权重扩展为2D）、计算合适的块大小，并调用内核。
2. 更新模型调用点：修改python/sglang/srt/models/qwen2_moe.py中的_append_shared_to_topk_output方法，移除原先手动创建共享ID张量、扩展权重并拼接的逻辑，改为导入并调用新的fused_append_shared_experts_with_weights函数。这简化了代码，并将计算委托给优化后的内核。
3. 测试与验证配套：PR body提供了详细的准确性测试（GSM8K准确率保持在~0.95）和性能基准测试结果（在不同并发度下，吞吐量提升1.93%-4.16%，延迟指标改善1.81%-3.99%），表明优化有效且无回归。本次改动未包含直接的单元测试文件变更，但依赖现有的模型测试和CI流程进行验证。

关键文件：

• python/sglang/srt/layers/moe/fused_moe_triton/fused_moe_triton_kernels.py（模块 MoE层；类别 source；类型 core-logic；符号 _fused_append_shared_experts_with_weights_kernel, fused_append_shared_experts_with_weights）: 新增了融合Triton内核及其包装函数，是性能优化的核心实现。
• python/sglang/srt/models/qwen2_moe.py（模块 模型层；类别 source；类型 data-contract；符号 _append_shared_to_topk_output）: 修改了MoE模型中的关键路由方法，以使用新的融合内核，是优化的调用入口。

关键符号：_fused_append_shared_experts_with_weights_kernel, fused_append_shared_experts_with_weights, _append_shared_to_topk_output

关键源码片段

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

新增了融合Triton内核及其包装函数，是性能优化的核心实现。

@triton.jit
def _fused_append_shared_experts_with_weights_kernel(
    topk_ids_ptr,
    topk_weights_ptr,
    shared_weights_ptr,
    out_ids_ptr,
    out_weights_ptr,
    N_BASE,
    K: tl.constexpr,
    S: tl.constexpr,
    BLOCK_K: tl.constexpr,
    BLOCK_S: tl.constexpr,
):
    pid = tl.program_id(0) # 每个程序ID处理一个token行
    ids_row_ptr = pid * K
    out_row_ptr = pid * (K + S)
​
    # 加载并存储原始的topk IDs和权重
    offs_k = tl.arange(0, BLOCK_K)
    mask_k = offs_k < K
    ids = tl.load(topk_ids_ptr + ids_row_ptr + offs_k, mask=mask_k)
    ws = tl.load(topk_weights_ptr + ids_row_ptr + offs_k, mask=mask_k)
    tl.store(out_ids_ptr + out_row_ptr + offs_k, ids, mask=mask_k)
    tl.store(out_weights_ptr + out_row_ptr + offs_k, ws, mask=mask_k)
​
    # 加载并存储共享专家的IDs和权重
    offs_s = tl.arange(0, BLOCK_S)
    mask_s = offs_s < S
    shared_ids = tl.cast(N_BASE + offs_s, ids.dtype) # 共享专家ID基于N_BASE偏移计算
    shared_ws = tl.load(shared_weights_ptr + pid * S + offs_s, mask=mask_s) # 逐token加载权重
    tl.store(out_ids_ptr + out_row_ptr + K + offs_s, shared_ids, mask=mask_s)
    tl.store(out_weights_ptr + out_row_ptr + K + offs_s, shared_ws, mask=mask_s)

python/sglang/srt/models/qwen2_moe.py

修改了MoE模型中的关键路由方法，以使用新的融合内核，是优化的调用入口。

def _append_shared_to_topk_output(
    self,
    topk_output: StandardTopKOutput,
    hidden_states: torch.Tensor,
) -> StandardTopKOutput:
    """Append shared expert ids and weights to topk output before fused MoE."""
    if not self.enable_shared_expert_fusion:
        return topk_output
    shared_weights = self._get_shared_expert_weights(hidden_states)
    if shared_weights is None:
        return topk_output
​
    # 导入并使用新的融合函数，替代原先的多步操作
    from sglang.srt.layers.moe.fused_moe_triton.fused_moe_triton_kernels import (
        fused_append_shared_experts_with_weights,
    )
​
    fused_topk_ids, fused_topk_weights = fused_append_shared_experts_with_weights(
        topk_output.topk_ids,
        topk_output.topk_weights,
        shared_weights,
        self.num_fused_shared_experts,
        N=self.num_experts, # 传递共享专家的基础ID
    )
    return StandardTopKOutput(
        topk_weights=fused_topk_weights,
        topk_ids=fused_topk_ids,
        router_logits=topk_output.router_logits,
    )

评论区精华

Review评论非常简短，两位审核者（kkHuang-amd和HaiShaw）均快速批准。HaiShaw的评论“LGTM, only AMD/AITER path relevant.”点明了本次优化的适用范围仅限于AMD/AITER路径，暗示了平台特定性，但未引发深入的技术讨论或争议。

• 优化适用范围确认 (design): 无争议，优化被认可为针对特定平台的有效改进。

风险与影响

• 风险：1. 平台兼容性风险：优化针对AMD平台（AITER后端），可能未在其他硬件（如NVIDIA）上充分测试，存在平台特定回归风险。
  2. 数值精度风险：新内核涉及数据类型转换（如shared_weights.to(topk_weights.dtype)）和扩展操作，虽然基准测试显示准确性无变化，但极端输入下可能引入细微数值差异。
  3. 内核健壮性风险：新增的Triton内核使用动态块大小（triton.next_power_of_2），对于非标准形状输入（如k或s非常大）可能影响性能或正确性。
  4. 依赖变更风险：qwen2_moe.py新增了从fused_moe_triton_kernels模块的导入，增加了模块间耦合，若内核函数接口未来变更，可能破坏模型代码。
• 影响：1. 性能影响：直接提升Qwen3.5 MoE模型在AMD平台上的推理性能，减少内核启动开销，在高并发场景下吞吐量提升可达4.16%，降低端到端延迟。
  2. 代码影响：简化了_append_shared_to_topk_output的实现逻辑，用单个函数调用替代多步张量操作，提高可读性和维护性。
  3. 用户影响：最终用户（使用AMD硬件运行Qwen3.5）将体验到更快的推理速度，无需更改任何配置或代码。
  4. 团队影响：为AMD平台优化树立了范例，可能鼓励类似针对关键路径的融合内核优化。
• 风险标记：平台特定优化, 内核健壮性依赖, 缺少单元测试

关联脉络

• PR #20736 [AMD] Enable share expert fusion with router experts for Qwen3.5 BF16 & FP8: 同属AMD平台Qwen3.5 MoE模型优化系列，该PR启用了共享专家融合，而本次PR进一步优化了融合路径中的内核启动效率。
• PR #22725 [Misc] Use cache_once for is_arch_support_pdl in sgl-kernel: 同为性能优化类PR，涉及内核层优化，但针对不同模块（sgl-kernel vs. MoE）。
• PR #22606 [serving] replace O(n²) stream_buffer string concat with integer offset: 同为性能优化PR，通过减少操作复杂度（O(n²)到O(1)）提升性能，与本次内核融合减少启动开销的思路类似。