[Kernel] Avoid zero-bias allocation in fused softmax routing (#36811)

python/sglang/kernels/ops/moe/moe_fused_gate.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Tuple
+from typing import TYPE_CHECKING, Optional, Tuple
 
 import torch
 import triton
@@ -109,6 +109,7 @@ def _router_triton_kernel(
     HAS_SOFTCAP: tl.constexpr,  # tanh softcapping (softmax only)
     RENORMALIZE: tl.constexpr,
     APPLY_SCALE: tl.constexpr,  # apply_routed_scaling_factor_on_output
+    HAS_BIAS: tl.constexpr,
     USE_PDL: tl.constexpr,
     stride_sm,
     stride_sn,
@@ -126,10 +127,13 @@ def _router_triton_kernel(
     mask_m = offs_m < M
     mask_n = offs_n < N
 
-    # prefetch bias before PDL wait
+    # Prefetch a real bias before the PDL wait. Plain softmax routing has no
-    bias = tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0).to(
+    # bias, so keep the zero value in registers rather than materializing and
-        tl.float32
+    # clearing a device tensor for every routing call.
-    )  # [BLOCK_N]
+    if HAS_BIAS:
+        bias = tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
+    else:
+        bias = tl.zeros([BLOCK_N], dtype=tl.float32)
 
     if USE_PDL:
         tl.extra.cuda.gdc_wait()
@@ -254,7 +258,7 @@ def _router_triton_kernel(
 @debug_kernel_api
 def moe_fused_gate(
     scores: torch.Tensor,
-    bias: torch.Tensor,
+    bias: Optional[torch.Tensor],
     topk: int,
     scoring_func: str = "sigmoid",
     num_fused_shared_experts: int = 0,
@@ -282,6 +286,12 @@ def moe_fused_gate(
         torch.float16,
         torch.bfloat16,
     ), "scores must be float32/float16/bfloat16"
+    assert scores.ndim == 2, "scores must be 2D"
+    if bias is None:
+        assert (
+            scoring_func.lower() == "softmax"
+        ), "bias is required for non-softmax routing"
+    else:
         # The kernel loads the bias and upcasts it to fp32 in-register (see
         # _router_triton_kernel), so a non-fp32 bias (DeepSeek-V4 stores the
         # correction bias in bf16) needs no host-side cast/copy.
@@ -290,9 +300,10 @@ def moe_fused_gate(
             torch.float16,
             torch.bfloat16,
         ), "bias must be float32/float16/bfloat16"
-    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 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"
     if routed_scaling_factor is None:
         routed_scaling_factor = 1.0
@@ -349,7 +360,7 @@ def moe_fused_gate(
     extra = {"launch_pdl": True} if use_pdl else {}
     _router_triton_kernel[grid](
         scores,
-        bias,
+        bias if bias is not None else scores,
         weights,
         indices,
         M,
@@ -369,6 +380,7 @@ def moe_fused_gate(
         HAS_SOFTCAP=bool(moe_softcapping != 0.0),
         RENORMALIZE=bool(renormalize),
         APPLY_SCALE=bool(apply_routed_scaling_factor_on_output),
+        HAS_BIAS=bias is not None,
         USE_PDL=use_pdl,
         stride_sm=scores.stride(0),
         stride_sn=scores.stride(1),

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

@@ -986,14 +986,9 @@ def fused_topk(
                 moe_fused_gate as _jit_moe_fused_gate,
             )
 
-            zero_bias = torch.zeros(
-                gating_output.shape[1],
-                dtype=torch.float32,
-                device=gating_output.device,
-            )
             topk_weights, topk_ids = _jit_moe_fused_gate(
                 gating_output,
-                zero_bias,
+                None,
                 topk,
                 scoring_func="softmax",
                 renormalize=renormalize,

test/registered/kernels/ops/moe/test_moe_fused_gate.py

@@ -343,6 +343,23 @@ def test_moe_fused_gate_softmax_matches_aot(
     )
 
 
+@pytest.mark.parametrize("M", [1, 8, 32])
+def test_moe_fused_gate_softmax_none_bias_matches_zero_bias(M: int) -> None:
+    torch.manual_seed(M)
+    scores = torch.randn(M, 256, dtype=torch.float32, device=DEVICE)
+    zero_bias = torch.zeros(256, dtype=torch.float32, device=DEVICE)
+
+    none_w, none_i = moe_fused_gate(
+        scores, None, topk=8, scoring_func="softmax", renormalize=True
+    )
+    zero_w, zero_i = moe_fused_gate(
+        scores, zero_bias, topk=8, scoring_func="softmax", renormalize=True
+    )
+
+    torch.testing.assert_close(none_w, zero_w, rtol=0, atol=0)
+    torch.testing.assert_close(none_i, zero_i, rtol=0, atol=0)
+
+
 _SIGMOID_AOT_CASES = get_ci_test_range(
     [
         (1, 128, 4, True, True, torch.float32),