sglang 0.5.4__py3-none-any.whl → 0.5.4.post2__py3-none-any.whl

sglang/srt/layers/moe/moe_runner/deep_gemm.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, List, Optional
 
 import torch
 
+from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.moe.moe_runner.base import (
     MoeQuantInfo,
     MoeRunnerConfig,
@@ -15,14 +16,31 @@ from sglang.srt.layers.moe.moe_runner.base import (
     register_pre_permute,
 )
 from sglang.srt.layers.moe.utils import MoeRunnerBackend
-from sglang.srt.utils import dispose_tensor
+from sglang.srt.utils import ceil_div, dispose_tensor, get_bool_env_var, is_hip, is_npu
+from sglang.srt.utils.offloader import get_offloader
 
 if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher.deepep import (
+        DeepEPLLCombineInput,
+        DeepEPLLDispatchOutput,
+        DeepEPNormalCombineInput,
+        DeepEPNormalDispatchOutput,
+    )
     from sglang.srt.layers.moe.token_dispatcher.standard import (
         StandardCombineInput,
         StandardDispatchOutput,
     )
 
+_is_hip = is_hip()
+_is_npu = is_npu()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+
+if not (_is_npu or _is_hip):
+    from sgl_kernel import silu_and_mul
+
+
+_MASKED_GEMM_FAST_ACT = get_bool_env_var("SGLANG_MASKED_GEMM_FAST_ACT")
+
 
 # TODO(kaixih@nvidia): ideally we should merge this logic into
 # `fill_gateup_input_triton_kernel` to directly generate e8m0 scale.
@@ -40,13 +58,23 @@ def _cast_to_e8m0_with_rounding_up(x: torch.Tensor) -> torch.Tensor:
     return new_x.transpose(1, 2).contiguous().transpose(1, 2)
 
 
+def copy_list_to_gpu_no_ce(arr: List[int]):
+    from sgl_kernel.elementwise import copy_to_gpu_no_ce
+
+    tensor_cpu = torch.tensor(arr, dtype=torch.int32, device="cpu")
+    tensor_gpu = torch.empty_like(tensor_cpu, device="cuda")
+    copy_to_gpu_no_ce(tensor_cpu, tensor_gpu)
+    return tensor_gpu
+
+
 @dataclass
 class DeepGemmRunnerInput(RunnerInput):
     hidden_states: torch.Tensor
     hidden_states_scale: torch.Tensor
-    masked_m: torch.Tensor
-    expected_m: int
     use_masked_gemm: bool
+    masked_m: Optional[torch.Tensor] = None
+    expected_m: Optional[int] = None
+    m_indices: Optional[torch.Tensor] = None
 
     @property
     def runner_backend(self) -> MoeRunnerBackend:
@@ -84,20 +112,100 @@ class DeepGemmRunnerCore(MoeRunnerCore):
         running_state: dict,
     ) -> DeepGemmRunnerOutput:
 
-        if runner_input.use_masked_gemm:
-            hidden_states = self._run_masked_gemm(
-                runner_input,
-                quant_info,
-                running_state,
+        if not runner_input.use_masked_gemm:
+            hidden_states = self._run_contiguous_gemm(
+                runner_input, quant_info, running_state
             )
         else:
-            hidden_states = self._run_contiguous_gemm(
-                runner_input,
-                quant_info,
-                running_state,
+            hidden_states = self._run_masked_gemm(
+                runner_input, quant_info, running_state
             )
         return DeepGemmRunnerOutput(hidden_states=hidden_states)
 
+    def _run_contiguous_gemm(
+        self,
+        runner_input: DeepGemmRunnerInput,
+        quant_info: DeepGemmMoeQuantInfo,
+        running_state: dict,
+    ) -> torch.Tensor:
+
+        from sglang.srt.layers.moe.ep_moe.kernels import tma_align_input_scale
+        from sglang.srt.layers.quantization.fp8_kernel import (
+            sglang_per_token_group_quant_fp8,
+        )
+
+        hidden_states = runner_input.hidden_states
+        hidden_states_scale = runner_input.hidden_states_scale
+        all_tokens = running_state["all_tokens"]
+        hidden_states_device = running_state["hidden_states_device"]
+        hidden_states_dtype = running_state["hidden_states_dtype"]
+        hidden_states_shape = running_state["hidden_states_shape"]
+        m_indices = runner_input.m_indices
+
+        N = quant_info.w13_weight.size(1)
+        K = hidden_states_shape[1]
+        scale_block_size = 128
+
+        w13_weight_fp8 = (
+            quant_info.w13_weight,
+            quant_info.w13_scale,
+        )
+        w2_weight_fp8 = (quant_info.w2_weight, quant_info.w2_scale)
+
+        gateup_output = torch.empty(
+            (all_tokens, N),
+            device=hidden_states_device,
+            dtype=torch.bfloat16,
+        )
+        if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
+            hidden_states_scale = tma_align_input_scale(hidden_states_scale)
+        deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig(
+            (hidden_states, hidden_states_scale),
+            w13_weight_fp8,
+            gateup_output,
+            m_indices,
+        )
+
+        dispose_tensor(hidden_states)
+        dispose_tensor(hidden_states_scale)
+
+        down_input = torch.empty(
+            (
+                all_tokens,
+                N // 2,
+            ),
+            device=gateup_output.device,
+            dtype=torch.bfloat16,
+        )
+        silu_and_mul(gateup_output.view(-1, N), down_input)
+        del gateup_output
+
+        down_input_fp8, down_input_scale = sglang_per_token_group_quant_fp8(
+            down_input,
+            scale_block_size,
+            column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
+            scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
+            scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
+        )
+        del down_input
+
+        down_output = torch.empty(
+            (all_tokens, K),
+            device=hidden_states_device,
+            dtype=torch.bfloat16,
+        )
+        if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
+            down_input_scale = tma_align_input_scale(down_input_scale)
+
+        deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig(
+            (down_input_fp8, down_input_scale),
+            w2_weight_fp8,
+            down_output,
+            m_indices,
+        )
+
+        return down_output
+
     def _run_masked_gemm(
         self,
         runner_input: DeepGemmRunnerInput,
@@ -109,6 +217,9 @@ class DeepGemmRunnerCore(MoeRunnerCore):
         from sglang.srt.layers.moe.ep_moe.kernels import (
             silu_and_mul_masked_post_quant_fwd,
         )
+        from sglang.srt.layers.quantization.fp8_kernel import (
+            sglang_per_token_group_quant_8bit,
+        )
 
         hidden_states = runner_input.hidden_states
         hidden_states_scale = runner_input.hidden_states_scale
@@ -122,15 +233,16 @@ class DeepGemmRunnerCore(MoeRunnerCore):
 
         hidden_states_device = running_state["hidden_states_device"]
 
-        if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
-            b, s_mn, s_k = hidden_states_scale.shape
-            assert (
-                s_mn % 4 == 0 and s_k % 4 == 0
-            ), f"scales must be aligned to 4, but got ({b}, {s_mn}, {s_k})"
-
         # GroupGemm-0
         if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
-            hidden_states_scale = _cast_to_e8m0_with_rounding_up(hidden_states_scale)
+            if hidden_states_scale.dtype != torch.int:
+                b, s_mn, s_k = hidden_states_scale.shape
+                assert (
+                    s_mn % 4 == 0 and s_k % 4 == 0
+                ), f"scales must be aligned to 4, but got ({b}, {s_mn}, {s_k})"
+                hidden_states_scale = _cast_to_e8m0_with_rounding_up(
+                    hidden_states_scale
+                )
         else:
             hidden_states_scale = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(
                 hidden_states_scale
@@ -149,35 +261,49 @@ class DeepGemmRunnerCore(MoeRunnerCore):
             expected_m,
         )
         dispose_tensor(hidden_states)
+        dispose_tensor(hidden_states_scale)
 
         # Act
-        down_input = torch.empty(
-            (
-                gateup_output.shape[0],
-                gateup_output.shape[1],
-                gateup_output.shape[2] // 2,
-            ),
-            device=hidden_states_device,
-            dtype=torch.float8_e4m3fn,
-        )
         scale_block_size = 128
-        down_input_scale = torch.empty(
-            (
-                gateup_output.shape[0],
-                gateup_output.shape[1],
-                gateup_output.shape[2] // 2 // scale_block_size,
-            ),
-            device=hidden_states_device,
-            dtype=torch.float32,
-        )
-        silu_and_mul_masked_post_quant_fwd(
-            gateup_output,
-            down_input,
-            down_input_scale,
-            scale_block_size,
-            masked_m,
-            scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
-        )
+        if _MASKED_GEMM_FAST_ACT:
+            down_input, down_input_scale = sglang_per_token_group_quant_8bit(
+                x=gateup_output,
+                dst_dtype=torch.float8_e4m3fn,
+                group_size=scale_block_size,
+                masked_m=masked_m,
+                column_major_scales=True,
+                scale_tma_aligned=True,
+                scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
+                fuse_silu_and_mul=True,
+                enable_v2=True,
+            )
+        else:
+            down_input = torch.empty(
+                (
+                    gateup_output.shape[0],
+                    gateup_output.shape[1],
+                    gateup_output.shape[2] // 2,
+                ),
+                device=hidden_states_device,
+                dtype=torch.float8_e4m3fn,
+            )
+            down_input_scale = torch.empty(
+                (
+                    gateup_output.shape[0],
+                    gateup_output.shape[1],
+                    gateup_output.shape[2] // 2 // scale_block_size,
+                ),
+                device=hidden_states_device,
+                dtype=torch.float32,
+            )
+            silu_and_mul_masked_post_quant_fwd(
+                gateup_output,
+                down_input,
+                down_input_scale,
+                scale_block_size,
+                masked_m,
+                scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
+            )
         del gateup_output
 
         # GroupGemm-1
@@ -198,18 +324,9 @@ class DeepGemmRunnerCore(MoeRunnerCore):
             masked_m,
             expected_m,
         )
-        del down_input
 
         return down_output
 
-    def _run_contiguous_gemm(
-        self,
-        runner_input: DeepGemmRunnerInput,
-        quant_info: DeepGemmMoeQuantInfo,
-        running_state: dict,
-    ) -> torch.Tensor:
-        pass
-
     @property
     def runner_backend(self) -> MoeRunnerBackend:
         return MoeRunnerBackend.DEEP_GEMM
@@ -222,6 +339,7 @@ def pre_permute_standard_to_deep_gemm(
     runner_config: MoeRunnerConfig,
     running_state: dict,
 ) -> DeepGemmRunnerInput:
+
     from sglang.srt.layers.moe.ep_moe.kernels import moe_ep_deepgemm_preprocess
 
     hidden_states, topk_output = dispatch_output
@@ -257,9 +375,9 @@ def pre_permute_standard_to_deep_gemm(
     return DeepGemmRunnerInput(
         hidden_states=hidden_states,
         hidden_states_scale=hidden_states_scale,
+        use_masked_gemm=True,
         masked_m=masked_m,
         expected_m=expected_m,
-        use_masked_gemm=True,
     )
 
 
@@ -302,3 +420,170 @@ def post_permute_deep_gemm_to_standard(
     return StandardCombineInput(
         hidden_states=output,
     )
+
+
+@register_pre_permute("deepep_ll", "deep_gemm")
+def pre_permute_deepep_ll_to_deep_gemm(
+    dispatch_output: DeepEPLLDispatchOutput,
+    quant_info: DeepGemmMoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+    running_state: dict,
+) -> DeepGemmRunnerInput:
+
+    hidden_states, hidden_states_scale, topk_ids, topk_weights, masked_m, expected_m = (
+        dispatch_output
+    )
+
+    running_state["topk_ids"] = topk_ids
+    running_state["topk_weights"] = topk_weights
+    running_state["hidden_states_shape"] = hidden_states.shape
+    running_state["hidden_states_dtype"] = hidden_states.dtype
+    running_state["hidden_states_device"] = hidden_states.device
+
+    return DeepGemmRunnerInput(
+        hidden_states=hidden_states,
+        hidden_states_scale=hidden_states_scale,
+        use_masked_gemm=True,
+        masked_m=masked_m,
+        expected_m=expected_m,
+    )
+
+
+@register_post_permute("deep_gemm", "deepep_ll")
+def post_permute_deep_gemm_to_deepep_ll(
+    runner_output: DeepGemmRunnerOutput,
+    quant_info: DeepGemmMoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+    running_state: dict,
+) -> DeepEPLLCombineInput:
+
+    from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPLLCombineInput
+
+    return DeepEPLLCombineInput(
+        hidden_states=runner_output.hidden_states,
+        topk_ids=running_state["topk_ids"],
+        topk_weights=running_state["topk_weights"],
+    )
+
+
+@register_pre_permute("deepep_normal", "deep_gemm")
+def pre_permute_deepep_normal_to_deep_gemm(
+    dispatch_output: DeepEPNormalDispatchOutput,
+    quant_info: DeepGemmMoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+    running_state: dict,
+) -> DeepGemmRunnerInput:
+
+    from sglang.srt.layers.moe.ep_moe.kernels import ep_scatter
+
+    (
+        hidden_states,
+        hidden_states_scale,
+        topk_ids,
+        topk_weights,
+        num_recv_tokens_per_expert,
+    ) = dispatch_output
+    assert runner_config.activation == "silu"
+
+    all_tokens = sum(num_recv_tokens_per_expert)
+    running_state["all_tokens"] = all_tokens
+
+    K = hidden_states.shape[1]
+
+    hidden_states_shape = hidden_states.shape
+    hidden_states_device = hidden_states.device
+    hidden_states_dtype = hidden_states.dtype
+
+    running_state["hidden_states_shape"] = hidden_states_shape
+    running_state["hidden_states_device"] = hidden_states_device
+    running_state["hidden_states_dtype"] = hidden_states_dtype
+    running_state["topk_ids"] = topk_ids
+    running_state["topk_weights"] = topk_weights
+
+    input_tensor = torch.empty(
+        (all_tokens, K),
+        device=hidden_states.device,
+        dtype=hidden_states.dtype,
+    )
+    if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
+        # TODO check whether need `zeros`
+        input_tensor_scale = torch.zeros(
+            (ceil_div(K // 128, 4), all_tokens),
+            device=hidden_states.device,
+            dtype=torch.int,
+        ).transpose(0, 1)
+    else:
+        input_tensor_scale = torch.empty(
+            (all_tokens, K // 128),
+            device=hidden_states.device,
+            dtype=torch.float32,
+        )
+    m_indices = torch.empty(all_tokens, device=hidden_states.device, dtype=torch.int32)
+    output_index = torch.empty_like(topk_ids)
+
+    if get_offloader().forbid_copy_engine_usage:
+        num_recv_tokens_per_expert_gpu = copy_list_to_gpu_no_ce(
+            num_recv_tokens_per_expert
+        )
+    else:
+        num_recv_tokens_per_expert_gpu = torch.tensor(
+            num_recv_tokens_per_expert,
+            dtype=torch.int32,
+            pin_memory=True,
+            device="cpu",
+        ).cuda(non_blocking=True)
+    expert_start_loc = torch.empty_like(num_recv_tokens_per_expert_gpu)
+
+    ep_scatter(
+        hidden_states,
+        hidden_states_scale,
+        topk_ids,
+        num_recv_tokens_per_expert_gpu,
+        expert_start_loc,
+        input_tensor,
+        input_tensor_scale,
+        m_indices,
+        output_index,
+        scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
+    )
+    dispose_tensor(hidden_states)
+    dispose_tensor(hidden_states_scale)
+
+    running_state["output_index"] = output_index
+
+    return DeepGemmRunnerInput(
+        hidden_states=input_tensor,
+        hidden_states_scale=input_tensor_scale,
+        use_masked_gemm=False,
+        m_indices=m_indices,
+    )
+
+
+@register_post_permute("deep_gemm", "deepep_normal")
+def post_permute_deep_gemm_to_deepep_normal(
+    runner_output: DeepGemmRunnerOutput,
+    quant_info: DeepGemmMoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+    running_state: dict,
+) -> DeepEPNormalCombineInput:
+
+    from sglang.srt.layers.moe.ep_moe.kernels import ep_gather
+    from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPNormalCombineInput
+
+    hidden_states = runner_output.hidden_states
+    topk_ids = running_state["topk_ids"]
+    topk_weights = running_state["topk_weights"]
+    output_index = running_state["output_index"]
+
+    gather_out = torch.empty(
+        running_state["hidden_states_shape"],
+        device=running_state["hidden_states_device"],
+        dtype=torch.bfloat16,
+    )
+    ep_gather(hidden_states, topk_ids, topk_weights, output_index, gather_out)
+
+    return DeepEPNormalCombineInput(
+        hidden_states=gather_out,
+        topk_ids=running_state["topk_ids"],
+        topk_weights=running_state["topk_weights"],
+    )

sglang/srt/layers/moe/moe_runner/runner.py

@@ -11,6 +11,7 @@ from sglang.srt.layers.moe.moe_runner.base import (
 )
 from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmRunnerCore
 from sglang.srt.layers.moe.moe_runner.triton import TritonRunnerCore
+from sglang.srt.layers.moe.moe_runner.triton_kernels import TritonKernelsRunnerCore
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend
 
 if TYPE_CHECKING:
@@ -31,6 +32,8 @@ class MoeRunner:
 
         if runner_backend.is_triton():
             self.runner_core = TritonRunnerCore(config)
+        elif runner_backend.is_triton_kernels():
+            self.runner_core = TritonKernelsRunnerCore(config)
         elif runner_backend.is_deep_gemm():
             self.runner_core = DeepGemmRunnerCore(config)
         else:

sglang/srt/layers/moe/moe_runner/triton_kernels.py

@@ -0,0 +1,194 @@
+"""Triton kernels MoE runner backend skeleton."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.layers.moe.moe_runner.base import (
+    MoeQuantInfo,
+    MoeRunnerConfig,
+    MoeRunnerCore,
+    RunnerInput,
+    RunnerOutput,
+    register_post_permute,
+    register_pre_permute,
+)
+from sglang.srt.layers.moe.utils import MoeRunnerBackend
+
+if TYPE_CHECKING:
+    from triton_kernels.matmul_ogs import PrecisionConfig
+    from triton_kernels.routing import GatherIndx, RoutingData, ScatterIndx
+
+    from sglang.srt.layers.moe.token_dispatcher.standard import (
+        StandardCombineInput,
+        StandardDispatchOutput,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Runner IO dataclasses
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class TritonKernelsRunnerInput(RunnerInput):
+    """Input bundle passed to the triton-kernels runner core."""
+
+    hidden_states: torch.Tensor
+    routing_data: "RoutingData"
+    gather_indx: "GatherIndx"
+    scatter_indx: "ScatterIndx"
+
+    @property
+    def runner_backend(self) -> MoeRunnerBackend:
+        return MoeRunnerBackend.TRITON_KERNELS
+
+
+@dataclass
+class TritonKernelsRunnerOutput(RunnerOutput):
+    """Output bundle returned from the triton-kernels runner core."""
+
+    hidden_states: torch.Tensor
+
+    @property
+    def runner_backend(self) -> MoeRunnerBackend:
+        return MoeRunnerBackend.TRITON_KERNELS
+
+
+@dataclass
+class TritonKernelsQuantInfo(MoeQuantInfo):
+    """Quantization payload consumed by the triton-kernels backend."""
+
+    w13_weight: torch.Tensor
+    w2_weight: torch.Tensor
+    w13_bias: Optional[torch.Tensor] = None
+    w2_bias: Optional[torch.Tensor] = None
+    w13_precision_config: Optional[PrecisionConfig] = None
+    w2_precision_config: Optional[PrecisionConfig] = None
+    global_num_experts: int = -1
+
+
+# ---------------------------------------------------------------------------
+# Runner core
+# ---------------------------------------------------------------------------
+
+
+class TritonKernelsRunnerCore(MoeRunnerCore):
+    """Execute MoE experts via the external triton_kernels package."""
+
+    def run(
+        self,
+        runner_input: TritonKernelsRunnerInput,
+        quant_info: TritonKernelsQuantInfo,
+        running_state: dict,
+    ) -> TritonKernelsRunnerOutput:
+        from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import (
+            triton_kernel_fused_experts,
+            triton_kernel_fused_experts_with_bias,
+        )
+
+        hidden_states = runner_input.hidden_states
+
+        common_kwargs = dict(
+            routing_data=runner_input.routing_data,
+            gather_indx=runner_input.gather_indx,
+            scatter_indx=None if self.config.no_combine else runner_input.scatter_indx,
+            inplace=False,
+            activation=self.config.activation,
+            apply_router_weight_on_input=self.config.apply_router_weight_on_input,
+            global_num_experts=quant_info.global_num_experts,
+        )
+
+        has_bias = quant_info.w13_bias is not None or quant_info.w2_bias is not None
+
+        if has_bias:
+            assert (
+                quant_info.w13_bias is not None and quant_info.w2_bias is not None
+            ), "Bias execution requires both w13_bias and w2_bias"
+            output = triton_kernel_fused_experts_with_bias(
+                hidden_states=hidden_states,
+                w1=quant_info.w13_weight,
+                w1_pcg=quant_info.w13_precision_config,
+                b1=quant_info.w13_bias,
+                w2=quant_info.w2_weight,
+                w2_pcg=quant_info.w2_precision_config,
+                b2=quant_info.w2_bias,
+                gemm1_alpha=self.config.gemm1_alpha,
+                gemm1_clamp_limit=self.config.gemm1_clamp_limit,
+                **common_kwargs,
+            )
+        else:
+            output = triton_kernel_fused_experts(
+                hidden_states=hidden_states,
+                w1=quant_info.w13_weight,
+                w2=quant_info.w2_weight,
+                **common_kwargs,
+            )
+
+        if self.config.no_combine:
+            tokens = runner_input.hidden_states.shape[0]
+            hidden = runner_input.hidden_states.shape[-1]
+            total_rows = output.shape[0]
+            top_k = total_rows // tokens
+            output = output.view(tokens, top_k, hidden)
+
+        return TritonKernelsRunnerOutput(hidden_states=output)
+
+    @property
+    def runner_backend(self) -> MoeRunnerBackend:
+        return MoeRunnerBackend.TRITON_KERNELS
+
+
+# ---------------------------------------------------------------------------
+# Permute / fused hooks
+# ---------------------------------------------------------------------------
+
+
+@register_pre_permute("standard", "triton_kernel")
+def pre_permute_standard_to_triton_kernels(
+    dispatch_output: "StandardDispatchOutput",
+    quant_info: TritonKernelsQuantInfo,
+    runner_config: MoeRunnerConfig,
+    running_state: dict,
+) -> TritonKernelsRunnerInput:
+    from sglang.srt.layers.moe.topk import TopKOutputChecker
+
+    hidden_states = dispatch_output.hidden_states
+    topk_output = dispatch_output.topk_output
+
+    assert TopKOutputChecker.format_is_triton_kernels(
+        topk_output
+    ), "Triton-kernel runner expects TritonKernelTopKOutput"
+
+    routing_data, gather_indx, scatter_indx = topk_output
+
+    return TritonKernelsRunnerInput(
+        hidden_states=hidden_states,
+        routing_data=routing_data,
+        gather_indx=gather_indx,
+        scatter_indx=scatter_indx,
+    )
+
+
+@register_post_permute("triton_kernel", "standard")
+def post_permute_triton_kernels_to_standard(
+    runner_output: TritonKernelsRunnerOutput,
+    quant_info: TritonKernelsQuantInfo,
+    runner_config: MoeRunnerConfig,
+    running_state: dict,
+) -> StandardCombineInput:
+    from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+
+    hidden_states = runner_output.hidden_states
+
+    if (
+        runner_config.routed_scaling_factor is not None
+        and runner_config.routed_scaling_factor != 1.0
+        and not runner_config.no_combine
+    ):
+        hidden_states.mul_(runner_config.routed_scaling_factor)
+
+    return StandardCombineInput(hidden_states=hidden_states)