From 221a6273ce3212c79483df233b4511fdf8fbe6d0 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Tue, 1 Sep 2026 16:29:57 -0700 Subject: [PATCH] [Kernel] Avoid zero-bias allocation in fused softmax routing (#36811) --- .../sglang/kernels/ops/moe/moe_fused_gate.py | 46 ++++++++++++------- python/sglang/srt/layers/moe/topk.py | 7 +-- .../kernels/ops/moe/test_moe_fused_gate.py | 17 +++++++ 3 files changed, 47 insertions(+), 23 deletions(-) diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index cf1c56fda..640a6610b 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/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 - bias = tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0).to( - tl.float32 - ) # [BLOCK_N] + # Prefetch a real bias before the PDL wait. Plain softmax routing has no + # bias, so keep the zero value in registers rather than materializing and + # clearing a device tensor for every routing call. + 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,17 +286,24 @@ def moe_fused_gate( torch.float16, torch.bfloat16, ), "scores must be float32/float16/bfloat16" - # 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. - assert bias.dtype in ( - torch.float32, - 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" + 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. + assert bias.dtype in ( + torch.float32, + torch.float16, + torch.bfloat16, + ), "bias must be float32/float16/bfloat16" + 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" 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), diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index e7431a2f2..7fa5c8f9e 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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, diff --git a/test/registered/kernels/ops/moe/test_moe_fused_gate.py b/test/registered/kernels/ops/moe/test_moe_fused_gate.py index 49614597d..3e909fa95 100644 --- a/test/registered/kernels/ops/moe/test_moe_fused_gate.py +++ b/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),