sglang 0.5.3rc2__py3-none-any.whl → 0.5.4.post1__py3-none-any.whl

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)

sglang/srt/layers/moe/rocm_moe_utils.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from enum import IntEnum
-from functools import cache
 from typing import Optional
 
 import torch

sglang/srt/layers/moe/router.py

@@ -11,7 +11,7 @@ _is_hip = is_hip()
 
 
 @triton.jit
-def fused_moe_router_kernel(
+def fused_moe_router_cudacore_kernel(
     input_ptr,  # input (bs, hidden_dim)
     moe_router_weight_ptr,  # input (num_experts, hidden_dim)
     topk_weights_ptr,  # output (bs, topk)
@@ -114,7 +114,7 @@ def fused_moe_router_kernel(
     # assert not moe_renormalize, "moe weight renormalization not implemented"
 
 
-def fused_moe_router_impl(
+def fused_moe_router_cudacore(
     x: torch.Tensor,
     router_weight: torch.Tensor,
     topk: int,
@@ -138,7 +138,7 @@ def fused_moe_router_impl(
         ),
     }
 
-    fused_moe_router_kernel[(bs,)](
+    fused_moe_router_cudacore_kernel[(bs,)](
         x,
         router_weight,
         topk_weights,
@@ -157,7 +157,7 @@ def fused_moe_router_impl(
 
 
 @triton.jit
-def fused_moe_router_large_bs_kernel(
+def fused_moe_router_tensorcore_kernel(
     a_ptr,  # input (bs, hidden_dim)
     b_ptr,  # input (num_experts, hidden_dim)
     topk_weights_ptr,  # output (bs, topk)
@@ -167,12 +167,15 @@ def fused_moe_router_large_bs_kernel(
     topk: tl.constexpr,  # only support topk <= 2
     moe_softcapping: tl.constexpr,
     moe_renormalize: tl.constexpr,  # not supported
+    correction_bias_ptr,
+    is_correction_bias: tl.constexpr,
     K: tl.constexpr,
     BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
     stride_am: tl.constexpr,
     stride_bn: tl.constexpr,
+    dp_attn_workaround_flag: tl.constexpr,
 ):
 
     # 1. get block id
@@ -217,6 +220,20 @@ def fused_moe_router_large_bs_kernel(
         exped = tl.exp(2 * logits_scaled)
         logits_softcapped = (exped - 1) / (exped + 1) * moe_softcapping
 
+    # Add bias after softcapping
+    if is_correction_bias:
+        bias = tl.load(
+            correction_bias_ptr + tl.arange(0, BLOCK_SIZE_N)[None, :],
+            mask=expert_mask.T,
+            other=0.0,
+        )
+        logits_softcapped = logits_softcapped + bias
+
+    if dp_attn_workaround_flag:
+        logits_softcapped = tl.where(
+            logits_softcapped != logits_softcapped, -1e9, logits_softcapped
+        )
+
     # 5. top1
     arange_block_size_n = tl.arange(0, BLOCK_SIZE_N)[None, :]
     cond_top1 = arange_block_size_n < num_experts
@@ -266,7 +283,7 @@ def fused_moe_router_large_bs_kernel(
         )
 
 
-def fused_moe_router_large_bs_impl(
+def fused_moe_router_tensorcore(
     x: torch.Tensor,
     router_weight: torch.Tensor,
     topk: int,
@@ -274,6 +291,7 @@ def fused_moe_router_large_bs_impl(
     BLOCK_SIZE_M: int,
     BLOCK_SIZE_N: int,
     BLOCK_SIZE_K: int,
+    correction_bias: Optional[torch.Tensor] = None,
 ):
     assert len(x.shape) == 2 and x.shape[1] == router_weight.shape[1]
     bs, hidden_dim = x.shape
@@ -285,10 +303,17 @@ def fused_moe_router_large_bs_impl(
 
     topk_weights = torch.empty((bs, topk), dtype=torch.float32, device=x.device)
     topk_ids = torch.empty((bs, topk), dtype=torch.int32, device=x.device)
+    is_correction_bias = correction_bias is not None
 
     grid = (triton.cdiv(bs, BLOCK_SIZE_M) * triton.cdiv(num_experts, BLOCK_SIZE_N),)
 
-    fused_moe_router_large_bs_kernel[grid](
+    # TODO(ch-wan): temporary workaround for dp attention. We should support masked
+    # router to skip padded tokens.
+    from sglang.srt.layers.dp_attention import is_dp_attention_enabled
+
+    dp_attn_workaround_flag = is_dp_attention_enabled()
+
+    fused_moe_router_tensorcore_kernel[grid](
         a_ptr=x,
         b_ptr=router_weight,
         topk_weights_ptr=topk_weights,
@@ -299,11 +324,14 @@ def fused_moe_router_large_bs_impl(
         moe_softcapping=moe_softcapping,
         moe_renormalize=False,
         K=hidden_dim,
+        correction_bias_ptr=correction_bias,
+        is_correction_bias=is_correction_bias,
         BLOCK_SIZE_M=BLOCK_SIZE_M,
         BLOCK_SIZE_N=BLOCK_SIZE_N,
         BLOCK_SIZE_K=BLOCK_SIZE_K,
         stride_am=hidden_dim,
         stride_bn=hidden_dim,
+        dp_attn_workaround_flag=dp_attn_workaround_flag,
     )
 
     return topk_weights, topk_ids
@@ -316,6 +344,7 @@ def fused_moe_router_shim(
     topk,
     renormalize,
     correction_bias: Optional[torch.Tensor] = None,
+    enable_deterministic_inference: bool = False,
 ):
     assert not renormalize
     assert (
@@ -324,16 +353,22 @@ def fused_moe_router_shim(
     )
     bs, hidden_dim = hidden_states.shape
     num_experts = gating_output.shape[0]
+
     BLOCK_SIZE_M = 32
-    BLOCK_SIZE_N = 16
-    BLOCK_SIZE_K = 256
+
+    BLOCK_SIZE_N = max(num_experts, 16)
+    BLOCK_SIZE_K = (
+        256 if num_experts < 256 else 64
+    )  # if experts are large, need to use smaller k block or shared memory OOM
+
     if (
-        bs >= 512
-        and topk <= 2
-        and num_experts <= BLOCK_SIZE_N
+        (bs >= 512 or num_experts > 8)
         and hidden_dim % BLOCK_SIZE_K == 0
+        # we keep using single kernel to avoid non-deterministic behavior
+        and not enable_deterministic_inference
     ):
-        return fused_moe_router_large_bs_impl(
+        # if large batch size or large expert, use kernel that uses tensorcore in matmul
+        return fused_moe_router_tensorcore(
             x=hidden_states,
             router_weight=gating_output,
             topk=topk,
@@ -341,9 +376,11 @@ def fused_moe_router_shim(
             BLOCK_SIZE_M=BLOCK_SIZE_M,
             BLOCK_SIZE_N=BLOCK_SIZE_N,
             BLOCK_SIZE_K=BLOCK_SIZE_K,
+            correction_bias=correction_bias,
         )
     else:
-        return fused_moe_router_impl(
+        # if smaller, use kernel that does not use tensorcore in matmul
+        return fused_moe_router_cudacore(
             x=hidden_states,
             router_weight=gating_output,
             topk=topk,
@@ -380,11 +417,10 @@ class FusedMoeRouter:
             renormalize=False,
         )
 
-    def forward_vllm(
+    def forward_torch(
         self,
         x: torch.Tensor,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        # g, _ = self.router_linear.forward(x)
         g = x.float() @ self.router_linear.weight.T.float()
 
         g = torch.tanh(g.float() / self.moe_softcapping) * self.moe_softcapping

sglang/srt/layers/moe/token_dispatcher/__init__.py

@@ -12,12 +12,18 @@ from sglang.srt.layers.moe.token_dispatcher.deepep import (
     DeepEPConfig,
     DeepEPDispatcher,
     DeepEPLLCombineInput,
-    DeepEPLLOutput,
+    DeepEPLLDispatchOutput,
     DeepEPNormalCombineInput,
-    DeepEPNormalOutput,
+    DeepEPNormalDispatchOutput,
+)
+from sglang.srt.layers.moe.token_dispatcher.mooncake import (
+    MooncakeCombineInput,
+    MooncakeDispatchOutput,
+    MooncakeEPDispatcher,
 )
 from sglang.srt.layers.moe.token_dispatcher.standard import (
     StandardCombineInput,
+    StandardDispatcher,
     StandardDispatchOutput,
 )
 
@@ -30,12 +36,16 @@ __all__ = [
     "DispatchOutput",
     "DispatchOutputFormat",
     "DispatchOutputChecker",
+    "MooncakeCombineInput",
+    "MooncakeDispatchOutput",
+    "MooncakeEPDispatcher",
+    "StandardDispatcher",
     "StandardDispatchOutput",
     "StandardCombineInput",
     "DeepEPConfig",
     "DeepEPDispatcher",
-    "DeepEPNormalOutput",
-    "DeepEPLLOutput",
+    "DeepEPNormalDispatchOutput",
+    "DeepEPLLDispatchOutput",
     "DeepEPLLCombineInput",
     "DeepEPNormalCombineInput",
 ]

sglang/srt/layers/moe/token_dispatcher/base.py

@@ -9,9 +9,9 @@ import torch
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.token_dispatcher import (
         DeepEPLLCombineInput,
-        DeepEPLLOutput,
+        DeepEPLLDispatchOutput,
         DeepEPNormalCombineInput,
-        DeepEPNormalOutput,
+        DeepEPNormalDispatchOutput,
         StandardCombineInput,
         StandardDispatchOutput,
     )
@@ -28,22 +28,28 @@ class DispatchOutputChecker:
     ) -> TypeGuard[StandardDispatchOutput]:
         return dispatch_output.format.is_standard()
 
+    @staticmethod
+    def format_is_triton_kernels(
+        dispatch_output: DispatchOutput,
+    ) -> TypeGuard[StandardDispatchOutput]:
+        return dispatch_output.format.is_standard()
+
     @staticmethod
     def format_is_deepep_normal(
         dispatch_output: DispatchOutput,
-    ) -> TypeGuard[DeepEPNormalOutput]:
+    ) -> TypeGuard[DeepEPNormalDispatchOutput]:
         return dispatch_output.format.is_deepep_normal()
 
     @staticmethod
     def format_is_deepep_ll(
         dispatch_output: DispatchOutput,
-    ) -> TypeGuard[DeepEPLLOutput]:
+    ) -> TypeGuard[DeepEPLLDispatchOutput]:
         return dispatch_output.format.is_deepep_ll()
 
     @staticmethod
     def format_is_deepep(
         dispatch_output: DispatchOutput,
-    ) -> TypeGuard[Union[DeepEPNormalOutput, DeepEPLLOutput]]:
+    ) -> TypeGuard[Union[DeepEPNormalDispatchOutput, DeepEPLLDispatchOutput]]:
         return dispatch_output.format.is_deepep()
 
 
@@ -73,7 +79,7 @@ class DispatchOutputFormat(Enum):
 class DispatchOutput(Protocol):
     """Protocol for dispatch outputs in different formats."""
 
-    # TODO: add hidden_states to the protocol
+    hidden_states: torch.Tensor
 
     @property
     def format(self) -> DispatchOutputFormat: ...