[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;
+      float final_max = (lane_id < num_warps) ? warp_maxs[lane_id] : -FLT_MAX;
-      int final_expert = (lane_id < kWarpsPerToken) ? warp_experts[lane_id] : -1;
+      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 (

test/registered/jit/benchmark/bench_moe_fused_gate.py

@@ -0,0 +1,98 @@
+import torch
+from sgl_kernel import kimi_k2_moe_fused_gate as aot_kimi_k2_gate
+from sgl_kernel import moe_fused_gate as aot_moe_fused_gate
+
+from sglang.jit_kernel.benchmark import marker
+from sglang.jit_kernel.benchmark.utils import create_random
+from sglang.jit_kernel.moe_fused_gate import moe_fused_gate, moe_fused_gate_jit
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=20, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+
+
+TOPK = 8
+SCALE = 2.5
+# AOT moe_fused_gate requires experts_per_group <= 32, so split experts into
+# groups of 32 and select every group (topk_group == num_expert_group) to get a
+# flat top-k. The 384-expert (3x128) layout uses the dedicated Kimi-K2 kernel.
+AOT_GROUP_SIZE = 32
+
+
+@torch.compile
+def torch_router(scores, bias, topk, scoring_func):
+    """Reference PyTorch router: scoring + bias + top-k + renorm + scale."""
+    if scoring_func == "sigmoid":
+        activated = scores.sigmoid()
+    else:
+        activated = torch.nn.functional.softplus(scores).sqrt()
+    biased = activated + bias.unsqueeze(0)
+    _, ids = torch.topk(biased, k=topk, dim=-1)
+    weights = activated.gather(1, ids)
+    weights = weights / weights.sum(dim=-1, keepdim=True)
+    return weights * SCALE, ids.to(torch.int32)
+
+
+@marker.parametrize("scoring_func", ["sigmoid", "sqrtsoftplus"])
+@marker.parametrize("num_experts", [128, 256, 384, 512], [256, 384])
+@marker.parametrize("num_tokens", [1, 4, 16, 64, 512, 1024, 8192], [16, 1024])
+@marker.benchmark("provider", ["triton", "jit", "aot", "torch"])
+def benchmark(num_tokens: int, num_experts: int, scoring_func: str, provider: str):
+    torch.manual_seed(0)
+    scores = create_random(num_tokens, num_experts, dtype=torch.float32)
+    bias = create_random(num_experts, dtype=torch.float32)
+
+    common = dict(
+        topk=TOPK,
+        scoring_func=scoring_func,
+        renormalize=True,
+        routed_scaling_factor=SCALE,
+        apply_routed_scaling_factor_on_output=True,
+    )
+    if provider == "triton":
+        return marker.do_bench(
+            moe_fused_gate, input_args=(scores, bias), input_kwargs=common
+        )
+    if provider == "jit":
+        return marker.do_bench(
+            moe_fused_gate_jit, input_args=(scores, bias), input_kwargs=common
+        )
+    if provider == "torch":
+        return marker.do_bench(
+            torch_router, input_args=(scores, bias, TOPK, scoring_func)
+        )
+    if provider == "aot":
+        # The AOT CUDA kernels only implement sigmoid scoring.
+        if scoring_func != "sigmoid":
+            marker.skip("AOT kernel supports sigmoid only")
+        if num_experts == 384:  # 3 groups of 128 -> dedicated Kimi-K2 kernel
+            return marker.do_bench(
+                aot_kimi_k2_gate,
+                input_args=(scores, bias),
+                input_kwargs=dict(
+                    topk=TOPK,
+                    renormalize=True,
+                    routed_scaling_factor=SCALE,
+                    apply_routed_scaling_factor_on_output=True,
+                ),
+            )
+        num_group = max(num_experts // AOT_GROUP_SIZE, 1)
+        return marker.do_bench(
+            aot_moe_fused_gate,
+            input_args=(
+                scores,
+                bias,
+                num_group,
+                num_group,
+                TOPK,
+                0,  # num_fused_shared_experts
+                SCALE,
+                True,  # apply_routed_scaling_factor_on_output
+            ),
+        )
+    raise ValueError(f"unknown provider: {provider}")
+
+
+if __name__ == "__main__":
+    benchmark.run()

test/registered/jit/test_moe_fused_gate.py

@@ -0,0 +1,258 @@
+"""Correctness tests for the Triton :func:`moe_fused_gate` router.
+
+The Triton kernel is a drop-in reimplementation of the CUDA fused gate for the
+ungrouped case (``num_expert_group == 1``). We validate it three ways:
+
+* against an explicit, definition-based torch reference (documents the math),
+* against the CUDA JIT kernel it mirrors (:func:`moe_fused_gate_jit`), and
+* against the production ``biased_grouped_topk_impl`` for the sigmoid / no-shared
+  path that the kernel actually replaces in ``topk.py``.
+
+Comparisons are order-independent: weights are scattered back to a dense
+``[M, num_experts + num_shared]`` layout so the per-row column order (and any
+tie-break choice) does not matter.
+"""
+
+from __future__ import annotations
+
+import sys
+from typing import Tuple
+
+import pytest
+import torch
+
+from sglang.jit_kernel.moe_fused_gate import moe_fused_gate, moe_fused_gate_jit
+from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.srt.layers.moe.topk import biased_grouped_topk_impl
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+DEVICE = "cuda"
+
+
+def _scatter_by_expert(
+    weights: torch.Tensor, indices: torch.Tensor, num_columns: int
+) -> torch.Tensor:
+    """Scatter (weight, id) pairs into a dense ``[M, num_columns]`` tensor.
+
+    Makes the comparison independent of the per-row slot order, so the test does
+    not depend on how ties between equal scores are broken.
+    """
+    dense = torch.zeros(
+        (weights.shape[0], num_columns), dtype=torch.float32, device=weights.device
+    )
+    dense.scatter_(1, indices.long(), weights.float())
+    return dense
+
+
+def _reference_gate(
+    scores: torch.Tensor,
+    bias: torch.Tensor,
+    topk: int,
+    scoring_func: str,
+    num_fused_shared_experts: int,
+    renormalize: bool,
+    routed_scaling_factor: float,
+    apply_routed_scaling_factor_on_output: bool,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Definition-based eager reference matching the CUDA fused-gate semantics."""
+    if scoring_func == "sigmoid":
+        activated = scores.sigmoid()
+    else:
+        activated = torch.nn.functional.softplus(scores).sqrt()
+
+    biased = activated + bias.unsqueeze(0)
+    num_experts = scores.size(1)
+    num_routed = topk - num_fused_shared_experts
+
+    # Top-k_routed by biased score; lowest expert id wins on ties (matches kernel).
+    bs = biased.size(0)
+    work = biased.clone()
+    arange = torch.arange(num_experts, device=scores.device).unsqueeze(0)
+    routed_idx = torch.empty(bs, num_routed, dtype=torch.int32, device=scores.device)
+    routed_wgt = torch.empty(bs, num_routed, dtype=torch.float32, device=scores.device)
+    for k in range(num_routed):
+        vals, _ = work.max(dim=1, keepdim=True)
+        lane = torch.where(work == vals, arange, num_experts + 1)
+        winner = lane.min(dim=1).values.to(torch.int32)
+        routed_idx[:, k] = winner
+        routed_wgt[:, k] = activated.gather(1, winner.long().unsqueeze(1)).squeeze(1)
+        work.scatter_(1, winner.long().unsqueeze(1), float("-inf"))
+
+    routed_sum = routed_wgt.sum(dim=1, keepdim=True)
+    weights = torch.empty(bs, topk, dtype=torch.float32, device=scores.device)
+    indices = torch.empty(bs, topk, dtype=torch.int32, device=scores.device)
+    weights[:, :num_routed] = routed_wgt
+    indices[:, :num_routed] = routed_idx
+    if num_fused_shared_experts > 0:
+        weights[:, num_routed:] = routed_sum / routed_scaling_factor
+        for j in range(num_fused_shared_experts):
+            indices[:, num_routed + j] = num_experts + j
+
+    if renormalize:
+        norm = torch.where(routed_sum > 0.0, routed_sum, torch.ones_like(routed_sum))
+        weights = weights / norm
+    if apply_routed_scaling_factor_on_output:
+        weights = weights * routed_scaling_factor
+    return weights, indices
+
+
+def _make_inputs(M: int, num_experts: int, seed: int):
+    torch.manual_seed(seed)
+    scores = torch.randn(M, num_experts, dtype=torch.float32, device=DEVICE) * 2.0
+    bias = torch.randn(num_experts, dtype=torch.float32, device=DEVICE) * 0.5
+    return scores, bias
+
+
+_NUM_EXPERTS = get_ci_test_range([128, 256, 384, 512], [128, 384, 512])
+_M = get_ci_test_range([1, 7, 64, 256, 1024], [1, 64, 1024])
+
+
+@pytest.mark.parametrize("M", _M)
+@pytest.mark.parametrize("num_experts", _NUM_EXPERTS)
+@pytest.mark.parametrize("topk", [4, 6, 8])
+@pytest.mark.parametrize("scoring_func", ["sigmoid", "sqrtsoftplus"])
+@pytest.mark.parametrize("num_shared", [0, 1])
+@pytest.mark.parametrize("renormalize", [True, False])
+@pytest.mark.parametrize("apply_scale", [True, False])
+def test_moe_fused_gate_matches_reference(
+    M: int,
+    num_experts: int,
+    topk: int,
+    scoring_func: str,
+    num_shared: int,
+    renormalize: bool,
+    apply_scale: bool,
+) -> None:
+    scores, bias = _make_inputs(M, num_experts, seed=num_experts * 100 + topk)
+    scale = 2.5
+
+    kwargs = dict(
+        topk=topk,
+        scoring_func=scoring_func,
+        num_fused_shared_experts=num_shared,
+        renormalize=renormalize,
+        routed_scaling_factor=scale,
+        apply_routed_scaling_factor_on_output=apply_scale,
+    )
+    triton_w, triton_i = moe_fused_gate(scores, bias, **kwargs)
+    ref_w, ref_i = _reference_gate(scores, bias, **kwargs)
+    torch.cuda.synchronize()
+
+    num_columns = num_experts + num_shared
+    torch.testing.assert_close(
+        _scatter_by_expert(triton_w, triton_i, num_columns),
+        _scatter_by_expert(ref_w, ref_i, num_columns),
+        rtol=1e-4,
+        atol=1e-5,
+    )
+
+
+@pytest.mark.parametrize(
+    "M,num_experts,topk,num_shared,scoring_func",
+    [
+        # DeepSeek-V4-ish: sqrtsoftplus, ungrouped
+        (8192, 256, 6, 0, "sqrtsoftplus"),
+        (8192, 384, 6, 0, "sqrtsoftplus"),
+        # Kimi-K2 family: 384 experts, sigmoid, ungrouped
+        (8192, 384, 8, 0, "sigmoid"),
+        # Generic large MoE with a fused shared expert
+        (8192, 512, 8, 1, "sigmoid"),
+    ],
+)
+def test_moe_fused_gate_matches_cuda_jit(
+    M: int, num_experts: int, topk: int, num_shared: int, scoring_func: str
+) -> None:
+    """Triton output must match the CUDA JIT kernel it reimplements."""
+    scores, bias = _make_inputs(M, num_experts, seed=123)
+
+    kwargs = dict(
+        topk=topk,
+        scoring_func=scoring_func,
+        num_fused_shared_experts=num_shared,
+        renormalize=True,
+        routed_scaling_factor=2.5,
+        apply_routed_scaling_factor_on_output=True,
+    )
+    triton_w, triton_i = moe_fused_gate(scores, bias, **kwargs)
+    cuda_w, cuda_i = moe_fused_gate_jit(scores, bias, **kwargs)
+    torch.cuda.synchronize()
+
+    num_columns = num_experts + num_shared
+    torch.testing.assert_close(
+        _scatter_by_expert(triton_w, triton_i, num_columns),
+        _scatter_by_expert(cuda_w, cuda_i, num_columns),
+        rtol=1e-4,
+        atol=1e-5,
+    )
+
+
+@pytest.mark.parametrize("num_experts,topk", [(256, 6), (384, 8), (512, 8)])
+@pytest.mark.parametrize("apply_scale", [True, False])
+def test_moe_fused_gate_matches_production_impl(
+    num_experts: int, topk: int, apply_scale: bool
+) -> None:
+    """Match production ``biased_grouped_topk_impl`` on the path it replaces.
+
+    The kernel supersedes the ungrouped sigmoid CUDA path in ``topk.py``; that
+    reference hardcodes sigmoid and no fused shared expert, and the production
+    path always renormalizes (the impl only applies the scaling factor when
+    ``renormalize`` is set), so we compare on the renormalized path.
+    """
+    M = 128
+    scores, bias = _make_inputs(M, num_experts, seed=7)
+    scale = 2.5
+
+    triton_w, triton_i = moe_fused_gate(
+        scores,
+        bias,
+        topk=topk,
+        scoring_func="sigmoid",
+        renormalize=True,
+        routed_scaling_factor=scale,
+        apply_routed_scaling_factor_on_output=apply_scale,
+    )
+    hidden_states = torch.empty((M, 1), dtype=torch.float32, device=DEVICE)
+    ref_w, ref_i = biased_grouped_topk_impl(
+        hidden_states,
+        scores,
+        bias,
+        topk,
+        True,
+        num_expert_group=1,
+        topk_group=1,
+        routed_scaling_factor=scale,
+        apply_routed_scaling_factor_on_output=apply_scale,
+    )
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(
+        _scatter_by_expert(triton_w, triton_i, num_experts),
+        _scatter_by_expert(ref_w, ref_i, num_experts),
+        rtol=1e-4,
+        atol=1e-5,
+    )
+
+
+def test_moe_fused_gate_shapes_and_dtypes() -> None:
+    """Output shapes/dtypes and renormalized weights for a DeepSeek-V4 config."""
+    M, N, K = 64, 256, 8
+    scores, bias = _make_inputs(M, N, seed=0)
+
+    w, i = moe_fused_gate(scores, bias, topk=K, scoring_func="sqrtsoftplus")
+    assert w.shape == (M, K)
+    assert i.shape == (M, K)
+    assert w.dtype == torch.float32
+    assert i.dtype == torch.int32
+
+    # Selected expert ids are valid (no shared experts here).
+    assert (i >= 0).all() and (i < N).all()
+    # Renormalized weights sum to 1 per row.
+    torch.testing.assert_close(
+        w.sum(dim=1), torch.ones(M, device=DEVICE), rtol=1e-4, atol=1e-5
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))