Refactor Triton-kernel MoE runner integration (#11795)

test/srt/test_triton_fused_moe.py

@@ -5,11 +5,10 @@ import torch.nn.functional as F
 from tqdm import tqdm
 
 from sglang.srt.layers.activation import SiluAndMul
-from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import (
-    triton_kernel_moe_forward,
-)
-from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
-from sglang.srt.layers.moe.topk import TopK
+from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
+from sglang.srt.layers.moe.moe_runner.triton_kernels import TritonKernelsQuantInfo
+from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
+from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -55,6 +54,7 @@ class TestFusedMOE(CustomTestCase):
         w2,
         score,
         topk,
+        return_per_expert: bool = False,
     ):
         B, D = a.shape
         a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
@@ -78,9 +78,14 @@ class TestFusedMOE(CustomTestCase):
                     a[mask] @ w1_compute[i].transpose(0, 1)
                 ) @ w2_compute[i].transpose(0, 1)
 
-        return (
-            out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
-        ).sum(dim=1)
+        weighted = out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(
+            out.dtype
+        )
+
+        if return_per_expert:
+            return weighted
+
+        return weighted.sum(dim=1)
 
     def _test_case(self, m, n, k, e, topk, dtype):
         rtol, atol = self.get_tolerance(dtype)
@@ -99,20 +104,43 @@ class TestFusedMOE(CustomTestCase):
             renormalize=False,
             use_grouped_topk=False,
         )
-        topk_op.use_triton_kernels = True
+        topk_op.topk_config.output_format = TopKOutputFormat.TRITON_KERNEL
         triton_topk_output = topk_op.forward_cuda(
             hidden_states=a,
             router_logits=score,
         )
 
-        moe_runner_config = MoeRunnerConfig(
-            inplace=False,
+        quant_info = TritonKernelsQuantInfo(w13_weight=w1_tri, w2_weight=w2_tri)
+
+        dispatch_output = StandardDispatchOutput(
+            hidden_states=a, topk_output=triton_topk_output
         )
-        triton_output = triton_kernel_moe_forward(
-            a, w1_tri, w2_tri, triton_topk_output, moe_runner_config
+
+        torch_per_expert = self.torch_naive_moe(
+            a, w1, w2, score, topk, return_per_expert=True
+        )
+        torch_combined = torch_per_expert.sum(dim=1)
+
+        def run_runner(config):
+            runner = MoeRunner(MoeRunnerBackend.TRITON_KERNELS, config)
+            result = runner.run(dispatch_output, quant_info)
+            return result.hidden_states
+
+        # Combined output (no_combine=False)
+        non_fused_config = MoeRunnerConfig(inplace=False)
+        non_fused_output = run_runner(non_fused_config)
+        torch.testing.assert_close(
+            non_fused_output, torch_combined, rtol=rtol, atol=atol
+        )
+
+        # Per-expert output (no_combine=True)
+        non_fused_no_combine_config = MoeRunnerConfig(
+            inplace=False, no_combine=True, top_k=topk
+        )
+        non_fused_no_combine_output = run_runner(non_fused_no_combine_config)
+        torch.testing.assert_close(
+            non_fused_no_combine_output, torch_per_expert, rtol=rtol, atol=atol
         )
-        torch_output = self.torch_naive_moe(a, w1, w2, score, topk)
-        torch.testing.assert_close(triton_output, torch_output, rtol=rtol, atol=atol)
 
     def test_various_configurations(self):
         m_values = [1, 32, 64, 256]