[JIT Kernel] Triton moe fused gate (#25835)

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: ziyi.xu <ziyi.xu@radixark.ai>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

python/sglang/jit_kernel/csrc/moe/moe_fused_gate.cuh

@@ -63,6 +63,10 @@ __global__ void moe_fused_gate_kernel_small_token(const MoEFusedGateParams __gri
   uint32_t tid = threadIdx.x;
   uint32_t warp_id = tid / kWarpSize;
   uint32_t lane_id = tid % kWarpSize;
+  // Actual warps launched (<= kWarpsPerToken). num_experts that need fewer than
+  // kWarpsPerToken warps leave the upper warp_maxs/warp_experts slots unwritten,
+  // so the cross-warp reduction below must only read the launched warps.
+  const uint32_t num_warps = blockDim.x / kWarpSize;
 
   extern __shared__ float shared_mem[];
   float* shared_scores = shared_mem;
@@ -116,8 +120,8 @@ __global__ void moe_fused_gate_kernel_small_token(const MoEFusedGateParams __gri
     __syncthreads();
 
     if (warp_id == 0) {
-      float final_max = (lane_id < kWarpsPerToken) ? warp_maxs[lane_id] : -FLT_MAX;
-      int final_expert = (lane_id < kWarpsPerToken) ? warp_experts[lane_id] : -1;
+      float final_max = (lane_id < num_warps) ? warp_maxs[lane_id] : -FLT_MAX;
+      int final_expert = (lane_id < num_warps) ? warp_experts[lane_id] : -1;
 
 #pragma unroll
       for (int offset = 16; offset > 0; offset /= 2) {

python/sglang/jit_kernel/moe_fused_gate.py

@@ -4,8 +4,11 @@ import logging
 from typing import TYPE_CHECKING, Tuple
 
 import torch
+import triton
+import triton.language as tl
 
-from sglang.jit_kernel.utils import cache_once, load_jit
+from sglang.jit_kernel.utils import cache_once, is_arch_support_pdl, load_jit
+from sglang.kernel_api_logging import debug_kernel_api
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module
@@ -37,7 +40,7 @@ def can_use_moe_fused_gate() -> bool:
         return False
 
 
-def moe_fused_gate(
+def moe_fused_gate_jit(
     input: torch.Tensor,
     bias: torch.Tensor,
     topk: int,
@@ -80,3 +83,169 @@ def moe_fused_gate(
     )
 
     return output, indices
+
+
+@triton.jit
+def _router_triton_kernel(
+    scores_ptr,  # [M, N] fp32, GEMM output (raw logits)
+    bias_ptr,  # [N]    fp32
+    out_weights_ptr,  # [M, K] fp32
+    out_indices_ptr,  # [M, K] int32
+    M,
+    routed_scaling_factor,
+    N: tl.constexpr,
+    K: tl.constexpr,  # total topk (includes fused shared experts)
+    K_ROUTED: tl.constexpr,  # K - num_fused_shared_experts
+    BLOCK_N: tl.constexpr,  # >= N, power of 2
+    BLOCK_K: tl.constexpr,  # >= K, power of 2
+    SCORING_FUNC: tl.constexpr,  # 0 = sigmoid, 1 = sqrtsoftplus
+    RENORMALIZE: tl.constexpr,
+    APPLY_SCALE: tl.constexpr,  # apply_routed_scaling_factor_on_output
+    USE_PDL: tl.constexpr,
+    stride_sm,
+    stride_sn,
+    stride_wm,
+    stride_wk,
+    stride_im,
+    stride_ik,
+) -> None:
+    pid = tl.program_id(0)
+    if pid >= M:
+        return
+
+    offs_n = tl.arange(0, BLOCK_N)
+    mask_n = offs_n < N
+    # prefetch bias before PDL wait
+    bias = tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
+
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    row_ptr = scores_ptr + pid * stride_sm + offs_n * stride_sn
+    scores = tl.load(row_ptr, mask=mask_n, other=0.0).to(tl.float32)
+
+    if SCORING_FUNC == 0:
+        # sigmoid(x) = 1 / (1 + exp(-x))
+        activated = tl.sigmoid(scores)
+    else:
+        # sqrt(softplus(x)) = sqrt(log1p(exp(x))); guard against overflow when x is large
+        sp = tl.where(
+            scores > 20.0,
+            scores,  # log1p(exp(big)) = big
+            tl.log(1.0 + tl.exp(scores)),
+        )
+        activated = tl.sqrt(sp)
+    biased = activated + bias
+
+    biased = tl.where(mask_n, biased, -float("inf"))
+    offs_k = tl.arange(0, BLOCK_K)
+    mask_k_total = offs_k < K
+    mask_k_routed = offs_k < K_ROUTED
+    selected_vals = tl.zeros([BLOCK_K], dtype=tl.float32)
+    selected_idx = tl.zeros([BLOCK_K], dtype=tl.int32)
+
+    cur = biased
+    for k in tl.static_range(K_ROUTED):
+        max_val = tl.max(cur, axis=0)
+        is_max = cur == max_val
+        lane_id = tl.where(is_max, offs_n, N + 1)
+        win_lane = tl.min(lane_id, axis=0).to(tl.int32)
+        win_activated = tl.sum(tl.where(offs_n == win_lane, activated, 0.0), axis=0)
+        slot = offs_k == k
+        selected_vals = tl.where(slot, win_activated, selected_vals)
+        selected_idx = tl.where(slot, win_lane, selected_idx)
+        cur = tl.where(offs_n == win_lane, -float("inf"), cur)
+
+    routed_sum = tl.sum(tl.where(mask_k_routed, selected_vals, 0.0), axis=0)
+
+    # Fill fused-shared-expert slots: weight = routed_sum / routed_scaling_factor,
+    # id = num_experts + (slot - K_ROUTED).
+    if K_ROUTED < K:
+        is_shared = (offs_k >= K_ROUTED) & mask_k_total
+        shared_weight = routed_sum / routed_scaling_factor
+        shared_idx = N + (offs_k - K_ROUTED)
+        selected_vals = tl.where(is_shared, shared_weight, selected_vals)
+        selected_idx = tl.where(is_shared, shared_idx, selected_idx)
+
+    if USE_PDL:
+        tl.extra.cuda.gdc_launch_dependents()
+
+    if RENORMALIZE:
+        norm = tl.where(routed_sum > 0.0, routed_sum, 1.0)
+        selected_vals = selected_vals / norm
+    if APPLY_SCALE:
+        selected_vals = selected_vals * routed_scaling_factor
+
+    out_w_ptr = out_weights_ptr + pid * stride_wm + offs_k * stride_wk
+    out_i_ptr = out_indices_ptr + pid * stride_im + offs_k * stride_ik
+    tl.store(out_w_ptr, selected_vals, mask=mask_k_total)
+    tl.store(out_i_ptr, selected_idx, mask=mask_k_total)
+
+
+@debug_kernel_api
+def moe_fused_gate(
+    scores: torch.Tensor,
+    bias: torch.Tensor,
+    topk: int,
+    scoring_func: str = "sigmoid",
+    num_fused_shared_experts: int = 0,
+    renormalize: bool = True,
+    routed_scaling_factor: float = 1.0,
+    apply_routed_scaling_factor_on_output: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Triton fused router: scoring + bias + topk + (optional) renorm/scale.
+
+    Mirrors the semantics of :func:`moe_fused_gate_jit` (the CUDA JIT kernel)
+    for the ungrouped case (``num_expert_group == 1``). The first argument is
+    named ``scores`` (raw GEMM logits) to match the existing call sites.
+    """
+    scoring_func_int = _SCORING_FUNC_MAP.get(scoring_func.lower())
+    assert (
+        scoring_func_int is not None
+    ), f"Unknown scoring_func '{scoring_func}', must be one of {list(_SCORING_FUNC_MAP.keys())}"
+    assert scores.dtype == torch.float32, "scores must be float32"
+    assert bias.dtype == torch.float32, "bias must be float32"
+    assert scores.ndim == 2, "scores must be 2D"
+    assert bias.ndim == 1, "bias must be 1D"
+    assert scores.size(1) == bias.size(0), "scores and bias must have same num_experts"
+    assert topk > num_fused_shared_experts, "topk must be > num_fused_shared_experts"
+
+    M, N = scores.shape
+    K = topk
+    K_routed = topk - num_fused_shared_experts
+
+    weights = torch.empty((M, K), dtype=torch.float32, device=scores.device)
+    indices = torch.empty((M, K), dtype=torch.int32, device=scores.device)
+
+    BLOCK_N = triton.next_power_of_2(N)  # 256 -> 256, 384 -> 512
+    BLOCK_K = triton.next_power_of_2(K)  # 6 -> 8, 8 -> 8
+    grid = (M,)
+    use_pdl = is_arch_support_pdl()
+    extra = {"launch_pdl": True} if use_pdl else {}
+    # A single warp keeps the per-row reductions cheap to synchronize.
+    _router_triton_kernel[grid](
+        scores,
+        bias,
+        weights,
+        indices,
+        M,
+        float(routed_scaling_factor),
+        N=N,
+        K=K,
+        K_ROUTED=K_routed,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+        SCORING_FUNC=scoring_func_int,
+        RENORMALIZE=bool(renormalize),
+        APPLY_SCALE=bool(apply_routed_scaling_factor_on_output),
+        USE_PDL=use_pdl,
+        stride_sm=scores.stride(0),
+        stride_sn=scores.stride(1),
+        stride_wm=weights.stride(0),
+        stride_wk=weights.stride(1),
+        stride_im=indices.stride(0),
+        stride_ik=indices.stride(1),
+        num_warps=1,
+        **extra,
+    )
+    return weights, indices

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

@@ -24,6 +24,7 @@ from typing import (
     NamedTuple,
     Optional,
     Protocol,
+    Tuple,
     TypeGuard,
     runtime_checkable,
 )
@@ -178,11 +179,6 @@ if _is_cuda:
     except ImportError:
         fused_topk_deepseek = None
 
-    try:
-        from sgl_kernel import kimi_k2_moe_fused_gate
-    except ImportError as e:
-        pass
-
 if _is_cuda or _is_hip or _is_xpu:
     from sgl_kernel import topk_softmax
 
@@ -1027,7 +1023,7 @@ def biased_topk_jit_kernel_impl(
     num_token_non_padded: Optional[torch.Tensor] = None,
     expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
     apply_routed_scaling_factor_on_output: Optional[bool] = False,
-):
+) -> Tuple[torch.Tensor, torch.Tensor]:
     assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
 
     if _use_aiter and scoring_func == "sqrtsoftplus" and num_fused_shared_experts == 0:
@@ -1355,8 +1351,7 @@ def biased_grouped_topk_gpu(
     num_fused_shared_experts: int = 0,
     routed_scaling_factor: Optional[float] = None,
     apply_routed_scaling_factor_on_output: Optional[bool] = False,
-):
-
+) -> Tuple[torch.Tensor, torch.Tensor]:
     num_tokens = gating_output.shape[0]
     num_experts = gating_output.shape[1]
     experts_per_group = (
@@ -1473,8 +1468,8 @@ def biased_grouped_topk_gpu(
             True,
             apply_routed_scaling_factor_on_output,
         )
+        return topk_weights, topk_ids
     else:
-        # Use optimized path for Kimi K2 (384 experts with num_expert_group=1)
         num_experts = gating_output.shape[1]
         if _is_cuda and num_experts == 384 and num_expert_group == 1:
             # ===== TO BE REFACTORED ====
@@ -1501,12 +1496,18 @@ def biased_grouped_topk_gpu(
                     apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
                 )
             # ===== END TO BE REFACTORED ====
-            return kimi_k2_moe_fused_gate(
+            from sglang.jit_kernel.moe_fused_gate import moe_fused_gate as jit_gate
+
+            return jit_gate(
                 gating_output.to(dtype=torch.float32),
                 correction_bias,
                 topk=topk,
+                scoring_func="sigmoid",
+                num_fused_shared_experts=num_fused_shared_experts,
                 renormalize=renormalize,
-                routed_scaling_factor=routed_scaling_factor,
+                routed_scaling_factor=(
+                    routed_scaling_factor if routed_scaling_factor is not None else 1.0
+                ),
                 apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
             )
         elif (