[MoE] Consolidate ungrouped + grouped gate/topk onto one Triton router (#26771) — faster than AOT on B200/H100/H200, at parity with flashinfer (#29771)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
d8a4f7a7aa
commit
a2d7eb303e
@@ -1,210 +0,0 @@
|
||||
import itertools
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.grouped_topk import grouped_topk as jit_grouped_topk
|
||||
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=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
|
||||
|
||||
|
||||
CORRECTNESS_CASES = get_ci_test_range(
|
||||
full_range=list(
|
||||
itertools.product(
|
||||
[1, 17, 128],
|
||||
[16, 32, 64, 128, 192, 256, 384, 512],
|
||||
[1, 2, 3, 4, 5, 6, 7, 8],
|
||||
)
|
||||
),
|
||||
ci_range=[
|
||||
(1, 16, 3), # smallest non-power-of-two topk
|
||||
(17, 128, 6), # Nemotron-3-Nano shape that exposed the bug
|
||||
(128, 192, 8), # Hunyuan-3 shape, power-of-two topk sanity case
|
||||
(33, 512, 7), # largest expert-count tier with non-power-of-two topk
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _make_inputs(num_tokens: int, num_experts: int, seed: int):
|
||||
torch.manual_seed(seed)
|
||||
hidden_states = torch.empty((num_tokens, 1), dtype=torch.float32, device="cuda")
|
||||
gating_output = torch.randn(
|
||||
(num_tokens, num_experts), dtype=torch.float32, device="cuda"
|
||||
)
|
||||
correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda") * 0.1
|
||||
return hidden_states, gating_output, correction_bias
|
||||
|
||||
|
||||
def _scatter_by_expert(
|
||||
weights: torch.Tensor, ids: torch.Tensor, num_experts: int
|
||||
) -> torch.Tensor:
|
||||
dense = torch.zeros(
|
||||
(weights.shape[0], num_experts), dtype=torch.float32, device=weights.device
|
||||
)
|
||||
dense.scatter_(1, ids.long(), weights)
|
||||
return dense
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens,num_experts,topk", CORRECTNESS_CASES)
|
||||
def test_grouped_topk_renormalize_matches_reference(
|
||||
num_tokens: int, num_experts: int, topk: int
|
||||
) -> None:
|
||||
hidden_states, gating_output, correction_bias = _make_inputs(
|
||||
num_tokens, num_experts, seed=1000 + num_experts * 10 + topk
|
||||
)
|
||||
scaling_factor = 2.826 if (num_experts, topk) == (192, 8) else 1.0
|
||||
|
||||
topk_weights, topk_ids = jit_grouped_topk(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
1,
|
||||
1,
|
||||
topk,
|
||||
True,
|
||||
scaling_factor,
|
||||
)
|
||||
ref_weights, ref_ids = biased_grouped_topk_impl(
|
||||
hidden_states,
|
||||
gating_output,
|
||||
correction_bias,
|
||||
topk,
|
||||
True,
|
||||
1,
|
||||
1,
|
||||
routed_scaling_factor=scaling_factor,
|
||||
apply_routed_scaling_factor_on_output=True,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
_scatter_by_expert(topk_weights, topk_ids, num_experts),
|
||||
_scatter_by_expert(ref_weights, ref_ids, num_experts),
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
topk_weights.sum(dim=-1),
|
||||
torch.full((num_tokens,), scaling_factor, dtype=torch.float32, device="cuda"),
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("topk", [3, 5, 6, 7])
|
||||
def test_grouped_topk_non_power_of_two_renormalize(topk: int) -> None:
|
||||
hidden_states, gating_output, correction_bias = _make_inputs(
|
||||
num_tokens=64, num_experts=128, seed=2000 + topk
|
||||
)
|
||||
|
||||
topk_weights, topk_ids = jit_grouped_topk(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
1,
|
||||
1,
|
||||
topk,
|
||||
True,
|
||||
1.0,
|
||||
)
|
||||
ref_weights, ref_ids = biased_grouped_topk_impl(
|
||||
hidden_states,
|
||||
gating_output,
|
||||
correction_bias,
|
||||
topk,
|
||||
True,
|
||||
1,
|
||||
1,
|
||||
routed_scaling_factor=1.0,
|
||||
apply_routed_scaling_factor_on_output=True,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
_scatter_by_expert(topk_weights, topk_ids, 128),
|
||||
_scatter_by_expert(ref_weights, ref_ids, 128),
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
topk_weights.sum(dim=-1),
|
||||
torch.ones((64,), dtype=torch.float32, device="cuda"),
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
|
||||
|
||||
def test_grouped_topk_negative_choice_scores_match_reference() -> None:
|
||||
hidden_states, gating_output, correction_bias = _make_inputs(
|
||||
num_tokens=64, num_experts=128, seed=23758
|
||||
)
|
||||
correction_bias.fill_(-2.0)
|
||||
|
||||
topk_weights, topk_ids = jit_grouped_topk(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
1,
|
||||
1,
|
||||
6,
|
||||
True,
|
||||
1.0,
|
||||
)
|
||||
ref_weights, ref_ids = biased_grouped_topk_impl(
|
||||
hidden_states,
|
||||
gating_output,
|
||||
correction_bias,
|
||||
6,
|
||||
True,
|
||||
1,
|
||||
1,
|
||||
routed_scaling_factor=1.0,
|
||||
apply_routed_scaling_factor_on_output=True,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
_scatter_by_expert(topk_weights, topk_ids, 128),
|
||||
_scatter_by_expert(ref_weights, ref_ids, 128),
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
|
||||
|
||||
def test_grouped_topk_without_renormalize_matches_reference() -> None:
|
||||
hidden_states, gating_output, correction_bias = _make_inputs(
|
||||
num_tokens=64, num_experts=128, seed=3006
|
||||
)
|
||||
|
||||
topk_weights, topk_ids = jit_grouped_topk(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
1,
|
||||
1,
|
||||
6,
|
||||
False,
|
||||
1.0,
|
||||
)
|
||||
ref_weights, ref_ids = biased_grouped_topk_impl(
|
||||
hidden_states,
|
||||
gating_output,
|
||||
correction_bias,
|
||||
6,
|
||||
False,
|
||||
1,
|
||||
1,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
_scatter_by_expert(topk_weights, topk_ids, 128),
|
||||
_scatter_by_expert(ref_weights, ref_ids, 128),
|
||||
rtol=1e-5,
|
||||
atol=1e-6,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -254,5 +254,219 @@ def test_moe_fused_gate_shapes_and_dtypes() -> None:
|
||||
)
|
||||
|
||||
|
||||
def _reference_softmax(
|
||||
gating: torch.Tensor, topk: int, renormalize: bool
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Plain-softmax topk reference (the AOT ``topk_softmax`` semantics)."""
|
||||
num_experts = gating.size(1)
|
||||
probs = torch.softmax(gating.float(), dim=-1)
|
||||
work = gating.float().clone()
|
||||
arange = torch.arange(num_experts, device=gating.device).unsqueeze(0)
|
||||
M = gating.size(0)
|
||||
idx = torch.empty(M, topk, dtype=torch.int32, device=gating.device)
|
||||
wgt = torch.empty(M, topk, dtype=torch.float32, device=gating.device)
|
||||
for k in range(topk):
|
||||
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)
|
||||
idx[:, k] = winner
|
||||
wgt[:, k] = probs.gather(1, winner.long().unsqueeze(1)).squeeze(1)
|
||||
work.scatter_(1, winner.long().unsqueeze(1), float("-inf"))
|
||||
if renormalize:
|
||||
wgt = wgt / wgt.sum(dim=1, keepdim=True)
|
||||
return wgt, idx
|
||||
|
||||
|
||||
@pytest.mark.parametrize("M", [1, 200, 1024])
|
||||
@pytest.mark.parametrize("num_experts,topk", [(128, 4), (256, 8), (512, 6)])
|
||||
@pytest.mark.parametrize("renormalize", [True, False])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
def test_moe_fused_gate_softmax_matches_aot(
|
||||
M: int, num_experts: int, topk: int, renormalize: bool, dtype: torch.dtype
|
||||
) -> None:
|
||||
"""Triton softmax path matches the AOT ``topk_softmax`` it replaces in fused_topk."""
|
||||
sgl_kernel = pytest.importorskip("sgl_kernel")
|
||||
torch.manual_seed(num_experts * 13 + topk)
|
||||
gating = torch.randn(M, num_experts, dtype=dtype, device=DEVICE) * 2.0
|
||||
zero_bias = torch.zeros(num_experts, dtype=torch.float32, device=DEVICE)
|
||||
|
||||
tri_w, tri_i = moe_fused_gate(
|
||||
gating, zero_bias, topk=topk, scoring_func="softmax", renormalize=renormalize
|
||||
)
|
||||
ref_w, ref_i = _reference_softmax(gating, topk, renormalize)
|
||||
|
||||
aot_w = torch.empty(M, topk, dtype=torch.float32, device=DEVICE)
|
||||
aot_i = torch.empty(M, topk, dtype=torch.int32, device=DEVICE)
|
||||
sgl_kernel.topk_softmax(aot_w, aot_i, gating, renormalize)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
dense_tri = _scatter_by_expert(tri_w, tri_i, num_experts)
|
||||
torch.testing.assert_close(
|
||||
dense_tri, _scatter_by_expert(ref_w, ref_i, num_experts), rtol=1e-3, atol=1e-3
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
dense_tri, _scatter_by_expert(aot_w, aot_i, num_experts), rtol=1e-3, atol=1e-3
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("M", [1, 200, 1024])
|
||||
@pytest.mark.parametrize("num_experts,topk", [(128, 4), (256, 8)])
|
||||
@pytest.mark.parametrize("renormalize", [True, False])
|
||||
@pytest.mark.parametrize("with_bias", [True, False])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
def test_moe_fused_gate_sigmoid_matches_aot(
|
||||
M: int,
|
||||
num_experts: int,
|
||||
topk: int,
|
||||
renormalize: bool,
|
||||
with_bias: bool,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
"""Triton sigmoid path matches the AOT ``topk_sigmoid`` it replaces in fused_topk."""
|
||||
sgl_kernel = pytest.importorskip("sgl_kernel")
|
||||
torch.manual_seed(num_experts * 17 + topk)
|
||||
gating = torch.randn(M, num_experts, dtype=dtype, device=DEVICE) * 2.0
|
||||
bias = (
|
||||
torch.randn(num_experts, dtype=torch.float32, device=DEVICE) * 0.5
|
||||
if with_bias
|
||||
else torch.zeros(num_experts, dtype=torch.float32, device=DEVICE)
|
||||
)
|
||||
|
||||
tri_w, tri_i = moe_fused_gate(
|
||||
gating, bias, topk=topk, scoring_func="sigmoid", renormalize=renormalize
|
||||
)
|
||||
aot_w = torch.empty(M, topk, dtype=torch.float32, device=DEVICE)
|
||||
aot_i = torch.empty(M, topk, dtype=torch.int32, device=DEVICE)
|
||||
sgl_kernel.topk_sigmoid(
|
||||
aot_w, aot_i, gating, renormalize, bias if with_bias else None
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
_scatter_by_expert(tri_w, tri_i, num_experts),
|
||||
_scatter_by_expert(aot_w, aot_i, num_experts),
|
||||
rtol=1e-3,
|
||||
atol=1e-3,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_experts,num_expert_group,topk_group,topk",
|
||||
[
|
||||
(256, 8, 4, 8), # DeepSeek-V3
|
||||
(128, 8, 4, 6),
|
||||
(256, 4, 2, 8),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
def test_moe_fused_gate_grouped_matches_production_impl(
|
||||
num_experts: int,
|
||||
num_expert_group: int,
|
||||
topk_group: int,
|
||||
topk: int,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
"""Grouped Triton routing must match the definitional biased_grouped_topk_impl.
|
||||
|
||||
The kernel adds DeepSeek-V3 grouped routing (per-group top-2-sum group scores,
|
||||
keep topk_group groups, then top-k within). biased_grouped_topk_impl is the
|
||||
eager reference the production grouped path is defined against.
|
||||
"""
|
||||
M = 256
|
||||
torch.manual_seed(num_experts * 7 + num_expert_group * 13 + topk)
|
||||
gating = torch.randn(M, num_experts, dtype=dtype, device=DEVICE) * 2.0
|
||||
bias = torch.randn(num_experts, dtype=torch.float32, device=DEVICE) * 0.5
|
||||
hidden = torch.randn(M, 16, dtype=dtype, device=DEVICE)
|
||||
|
||||
tri_w, tri_i = moe_fused_gate(
|
||||
gating,
|
||||
bias,
|
||||
topk=topk,
|
||||
scoring_func="sigmoid",
|
||||
renormalize=True,
|
||||
num_expert_group=num_expert_group,
|
||||
topk_group=topk_group,
|
||||
)
|
||||
ref_w, ref_i = biased_grouped_topk_impl(
|
||||
hidden,
|
||||
gating,
|
||||
bias,
|
||||
topk,
|
||||
True,
|
||||
num_expert_group=num_expert_group,
|
||||
topk_group=topk_group,
|
||||
num_fused_shared_experts=0,
|
||||
routed_scaling_factor=1.0,
|
||||
apply_routed_scaling_factor_on_output=False,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
torch.testing.assert_close(
|
||||
_scatter_by_expert(tri_w, tri_i, num_experts),
|
||||
_scatter_by_expert(ref_w, ref_i, num_experts),
|
||||
rtol=1e-3,
|
||||
atol=1e-3,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_experts,num_expert_group,topk_group,topk,num_fused_shared_experts",
|
||||
[
|
||||
(256, 8, 4, 8, 0), # DeepSeek-V3
|
||||
(256, 8, 4, 9, 1), # DeepSeek-V3 + one fused shared expert
|
||||
],
|
||||
)
|
||||
def test_grouped_dispatch_flag_matches_default(
|
||||
num_experts: int,
|
||||
num_expert_group: int,
|
||||
topk_group: int,
|
||||
topk: int,
|
||||
num_fused_shared_experts: int,
|
||||
) -> None:
|
||||
"""The opt-in SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK dispatch must match the
|
||||
default grouped path (flashinfer/AOT) that biased_grouped_topk_gpu selects when
|
||||
the flag is off. This covers the wiring, not just the raw kernel — validated
|
||||
bit-exact on DeepSeek-V3.2 e2e; here we assert parity against the default path.
|
||||
"""
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.moe.topk import biased_grouped_topk_gpu
|
||||
|
||||
M = 256
|
||||
torch.manual_seed(num_experts * 3 + num_expert_group * 5 + topk)
|
||||
# fp32 gating: both the default (flashinfer upcasts to fp32) and the Triton
|
||||
# dispatch (also upcasts) operate on the same fp32 scores, so no bf16
|
||||
# borderline-expert divergence is expected.
|
||||
gating = torch.randn(M, num_experts, dtype=torch.float32, device=DEVICE) * 2.0
|
||||
bias = torch.randn(num_experts, dtype=torch.float32, device=DEVICE) * 0.5
|
||||
hidden = torch.randn(M, 16, dtype=torch.float32, device=DEVICE)
|
||||
|
||||
kwargs = dict(
|
||||
num_expert_group=num_expert_group,
|
||||
topk_group=topk_group,
|
||||
num_fused_shared_experts=num_fused_shared_experts,
|
||||
routed_scaling_factor=2.5,
|
||||
apply_routed_scaling_factor_on_output=False,
|
||||
)
|
||||
with envs.SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK.override(False):
|
||||
def_w, def_i = biased_grouped_topk_gpu(
|
||||
hidden, gating, bias, topk, True, **kwargs
|
||||
)
|
||||
with envs.SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK.override(True):
|
||||
jit_w, jit_i = biased_grouped_topk_gpu(
|
||||
hidden, gating, bias, topk, True, **kwargs
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Compare routed experts only (shared-expert slot ids are placeholders the
|
||||
# downstream fusion overwrites; the routed selection + weights are what matter).
|
||||
topk_routed = topk - num_fused_shared_experts
|
||||
torch.testing.assert_close(
|
||||
_scatter_by_expert(def_w[:, :topk_routed], def_i[:, :topk_routed], num_experts),
|
||||
_scatter_by_expert(jit_w[:, :topk_routed], jit_i[:, :topk_routed], num_experts),
|
||||
rtol=1e-3,
|
||||
atol=1e-3,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
|
||||
Reference in New Issue
Block a user