执行摘要

• 一句话：为 XPU 添加 fused_topk 支持
• 推荐动作：建议合入。PR 改动清晰、聚焦，审核通过。开发者在其他平台（NPU）已有类似实现的先例，XPU 的对应路径是合理的扩展。可考虑后续添加 XPU 端到端测试以确保正确性。

功能与动机

XPU 平台上需要高效的 fused topk 实现来处理 MoE 层的专家路由，以减少 token 排序和选择的开销。PR 描述中提到 'Support fused_topk for xpu'，旨在避免因使用默认 torch native 实现而导致的性能瓶颈。

实现拆解

1. 新增 forward_xpu 方法：在 TopK 类（python/sglang/srt/layers/moe/topk.py）中新增 forward_xpu 方法，该方法接收与 forward 相同的参数（hidden_states、router_logits 等），返回 TopKOutput。
2. 条件判断逻辑：方法内部先设置 self.topk_config.torch_native = True（默认走 torch native 路径），然后根据条件 self.topk_config.top_k <= 8 and router_logits.shape[1] <= 256 决定是否将 torch_native 设为 False，从而切换到 fused 实现。该条件限制了 fused 路径仅在 top-k 较小且专家数适中时启用，以避免过大的专家数导致性能回退。
3. 调用通用函数：最终通过 select_experts 函数将参数传递给底层实现，select_experts 会根据 topk_config.torch_native 的值选择不同的 fused kernel 或 torch native 实现。
4. 无其他文件改动：变更集中在单个文件，未涉及测试、配置或部署配套。

关键文件：

• python/sglang/srt/layers/moe/topk.py（模块 MoE 层；类别 source；类型 core-logic；符号 forward_xpu）: 核心变更文件，新增 forward_xpu 方法用于 XPU 平台的 fused topk 支持。

关键符号：forward_xpu

关键源码片段

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

核心变更文件，新增 forward_xpu 方法用于 XPU 平台的 fused topk 支持。

# python/sglang/srt/layers/moe/topk.py
​
    def forward_xpu(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
        *,
        num_token_non_padded: Optional[torch.Tensor] = None,
        expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
    ) -> TopKOutput:
        # 默认使用 torch native 路径
        self.topk_config.torch_native = True
        # [NOTE] XPU device support for topk kernels
        # - support 'topk_softmax' and 'topk_sigmoid'
        # - support up to 8 top-k and 256 experts
        # 当 top_k <= 8 且专家数 <= 256 时，切换到 fused 实现以获得更好性能
        self.topk_config.torch_native = not (
            self.topk_config.top_k <= 8 and router_logits.shape[1] <= 256
        )
​
        # 进入统一的 select_experts 函数，由内部逻辑根据 torch_native 选择实现
        return select_experts(
            hidden_states=hidden_states,
            layer_id=self.layer_id,
            router_logits=router_logits,
            topk_config=self.topk_config,
            num_token_non_padded=num_token_non_padded,
            expert_location_dispatch_info=expert_location_dispatch_info,
        )

评论区精华

审核者 mingfeima 提出了一个代码建议：将原先的注释和条件逻辑简化并改为单行条件判断（self.topk_config.torch_native = not (self.topk_config.top_k <= 8 and router_logits.shape[1] <= 256)），并更新注释，该建议已经被采纳。最终审核者 approving 了该 PR，认为 'generally LGTM'。没有其他争论点或未解决的问题。

• 简化条件判断和注释 (style): 采纳建议，最终代码使用了简化的条件判断。

风险与影响

• 风险：风险较低。变更仅对 XPU 设备生效，不影响其他平台。条件判断 top_k <= 8 and num_experts <= 256 确保 fused 路径仅在预期范围内启用，避免因超大专家数导致性能退化。但缺少直接针对 XPU 的测试覆盖，如果未来条件发生变化或模型配置超出该范围，可能回退到 torch native 路径，导致性能不一致。
• 影响：影响范围小，仅针对 XPU 设备上的 MoE topk 操作。对用户而言，在 XPU 上运行模型时，当 top-k ≤ 8 且专家数 ≤ 256 时，topk 操作将使用优化的 fused 实现，可能带来性能提升。对系统无其他影响。团队需确保未来对 topk 逻辑的修改不破坏 XPU 路径。
• 风险标记：缺少测试覆盖

关联脉络

• PR #25775 [Perf][Qwen3.5] Add case 512 to topkGatingSoftmaxKernelLauncher: 同属 MoE topk 性能优化，涉及 fused kernel 路径的扩展。