Skip redundant moe_sum_reduce for single-expert routing on XPU (#22660)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py

@@ -790,11 +790,14 @@ def _fused_moe_kernel_sequence(
                     routed_scaling_factor,
                 )
     elif _is_xpu:
-        moe_sum_reduce(
+        if topk == 1 and routed_scaling_factor == 1.0 and not _use_intermediate:
-            intermediate_cache3.view(*intermediate_cache3.shape),
+            pass  # we wrote directly into out_hidden_states
-            out_hidden_states,
+        else:
-            routed_scaling_factor,
+            moe_sum_reduce(
-        )
+                intermediate_cache3.view(*intermediate_cache3.shape),
+                out_hidden_states,
+                routed_scaling_factor,
+            )
     else:
         if _has_vllm_ops:
             vllm_ops.moe_sum(

test/registered/moe/test_fused_moe.py

@@ -4,6 +4,7 @@ import torch
 from tqdm import tqdm
 
 from sglang.srt.layers.activation import SiluAndMul
+from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
 from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_moe
 from sglang.srt.layers.moe.topk import TopKConfig, select_experts
 from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
@@ -66,6 +67,7 @@ class TestFusedMOE(CustomTestCase):
         w2_scale=None,
         a1_scale=None,
         a2_scale=None,
+        routed_scaling_factor=None,
     ):
         set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
 
@@ -100,11 +102,16 @@ class TestFusedMOE(CustomTestCase):
                     a[mask] @ w1_compute[i].transpose(0, 1)
                 ) @ w2_compute[i].transpose(0, 1)
 
-        return (
+        result = (
             out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
         ).sum(dim=1)
+        if routed_scaling_factor is not None:
+            result = result * routed_scaling_factor
+        return result
 
-    def _test_case(self, m, n, k, e, topk, dtype, use_fp8_w8a8=False):
+    def _test_case(
+        self, m, n, k, e, topk, dtype, use_fp8_w8a8=False, routed_scaling_factor=None
+    ):
         rtol, atol = self.get_tolerance(dtype)
 
         if use_fp8_w8a8:
@@ -183,8 +190,15 @@ class TestFusedMOE(CustomTestCase):
                 topk_config=TopKConfig(top_k=topk, renormalize=False),
             )
 
-            triton_output = fused_moe(a, w1, w2, topk_output)
+            moe_runner_config = MoeRunnerConfig(
-            torch_output = self.torch_naive_moe(a, w1, w2, score, topk)
+                routed_scaling_factor=routed_scaling_factor
+            )
+            triton_output = fused_moe(
+                a, w1, w2, topk_output, moe_runner_config=moe_runner_config
+            )
+            torch_output = self.torch_naive_moe(
+                a, w1, w2, score, topk, routed_scaling_factor=routed_scaling_factor
+            )
             torch.testing.assert_close(
                 triton_output, torch_output, rtol=rtol, atol=atol
             )
@@ -239,6 +253,25 @@ class TestFusedMOE(CustomTestCase):
                                             empty_gpu_cache()
                                         pbar.update(1)
 
+    def test_single_expert_routing(self):
+        # Cover the topk == 1 fast path (e.g. Llama-4-Scout num_experts_per_tok=1),
+        # where the second kernel writes directly into the output when
+        # routed_scaling_factor == 1.0 and the reduction is otherwise applied.
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+        for routed_scaling_factor in [1.0, 1.5]:
+            for m in [1, 33]:
+                with self.subTest(m=m, routed_scaling_factor=routed_scaling_factor):
+                    self._test_case(
+                        m,
+                        n=128,
+                        k=128,
+                        e=8,
+                        topk=1,
+                        dtype=torch.bfloat16,
+                        routed_scaling_factor=routed_scaling_factor,
+                    )
+                    empty_gpu_cache()
+
 
 if __name__ == "__main__":
     unittest.main()