diff --git a/python/sglang/jit_kernel/csrc/moe/moe_fused_gate.cuh b/python/sglang/jit_kernel/csrc/moe/moe_fused_gate.cuh index 6476a3be2..373921df6 100644 --- a/python/sglang/jit_kernel/csrc/moe/moe_fused_gate.cuh +++ b/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) { diff --git a/python/sglang/jit_kernel/moe_fused_gate.py b/python/sglang/jit_kernel/moe_fused_gate.py index d0daad0a3..23dfa1dc7 100644 --- a/python/sglang/jit_kernel/moe_fused_gate.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 4318218ee..d373df266 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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 ( diff --git a/test/registered/jit/benchmark/bench_moe_fused_gate.py b/test/registered/jit/benchmark/bench_moe_fused_gate.py new file mode 100644 index 000000000..05dd5faec --- /dev/null +++ b/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() diff --git a/test/registered/jit/test_moe_fused_gate.py b/test/registered/jit/test_moe_fused_gate.py new file mode 100644 index 000000000..6716be582 --- /dev/null +++ b/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"]))