minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[MoE][ROCm] Admit the unified Triton router on ROCm, including single-group routing (#38328)

Browse Source 
Co-authored-by: JohnQinAMD <yanyuan.qin@amd.com>
Co-authored-by: RuibinCheung <ruibzhan@amd.com>
Co-authored-by: Zhang, Jiejing <jiejing.zhang@amd.com>
This commit is contained in:
xiaobochen-amd
2026-09-12 16:21:28 -07:00
committed by GitHub
co-authored by JohnQinAMD RuibinCheung Zhang, Jiejing
parent a66451c058
commit 6657f7d844
4 changed files with 259 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/moe/moe_fused_gate.py
+1 -1
View File
@@ -88,7 +88,7 @@ def moe_fused_gate_jit(
@triton.jit
def _router_triton_kernel(
scores_ptr, # [M, N] fp32, GEMM output (raw logits)
scores_ptr, # [M, N] raw logits, fp32/fp16/bf16 (upcast to fp32 on load)
bias_ptr, # [N] fp32/fp16/bf16 (upcast to fp32 on load)
out_weights_ptr, # [M, K] fp32
out_indices_ptr, # [M, K] int32
python/sglang/srt/layers/attention/wave_ops/decode_attention.py
+52
View File
@@ -6,6 +6,7 @@ It supports page size = 1.
import functools
import logging
import torch
from wave_lang.kernel.lang.global_symbols import *
from wave_lang.kernel.wave.compile import WaveCompileOptions, wave_compile
from wave_lang.kernel.wave.constraints import GenericDot, MMAOperand, MMAType
@@ -23,6 +24,31 @@ import os
dump_generated_mlir = int(os.environ.get("WAVE_DUMP_MLIR", 0))
@functools.lru_cache(maxsize=None)
def _is_rocm10_or_newer() -> bool:
"""Return whether the runtime needs the ROCm 10 Wave decode workaround."""
hip_version = torch.version.hip
if hip_version is None:
return False
try:
hip_major_minor = tuple(int(part) for part in hip_version.split(".")[:2])
except ValueError:
return False
# torch 2.11's ROCm 10 build reports HIP 7.15.
return hip_major_minor >= (7, 15)
def _needs_triton_fallback(q, k_buffer, v_buffer) -> bool:
# Wave's paged decode kernel returns NaNs for this shape with ROCm 10 on
# gfx942. Keep Wave enabled for every other shape and older ROCm versions.
shape = (q.shape[1], k_buffer.shape[1], q.shape[2], v_buffer.shape[2])
if shape != (128, 1, 576, 512):
return False
return _is_rocm10_or_newer()
@functools.lru_cache(maxsize=4096)
def get_wave_kernel(
shape: paged_decode_attention_shape,
@@ -119,6 +145,32 @@ def decode_attention_wave(
num_seqs, num_query_heads, head_size = q.shape
_, num_kv_heads, _ = k_buffer.shape
_, _, head_size_kv = v_buffer.shape
if _needs_triton_fallback(q, k_buffer, v_buffer):
# The Wave and Triton intermediates contain the same number of values,
# but use different dimension orders. Reuse their storage so the
# fallback does not add an allocation to the decode path.
from sglang.kernels.ops.attention.decode_attention import (
decode_attention_fwd_grouped as triton_decode_attention_fwd_grouped,
)
triton_decode_attention_fwd_grouped(
q,
k_buffer,
v_buffer,
o,
b_req_idx,
req_to_token,
attn_logits.reshape(num_seqs, num_query_heads, max_kv_splits, head_size_kv),
attn_logits_max.reshape(num_seqs, num_query_heads, max_kv_splits),
num_kv_splits,
max_kv_splits,
sm_scale,
1.0,
logit_cap,
)
return
block_size = 32
shape = paged_decode_attention_shape(
num_query_heads,
python/sglang/srt/layers/moe/topk.py
+20 -10
View File
@@ -728,7 +728,6 @@ class TopK(BaseFusedOp):
num_token_non_padded: Optional[torch.Tensor] = None,
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
) -> TopKOutput:
from sglang.srt.hardware_backend.npu.moe.topk import fused_topk_npu
return fused_topk_npu(
@@ -1624,11 +1623,10 @@ def biased_grouped_topk_gpu(
# topk for routed experts only (shared experts are appended separately below)
topk_routed = topk - num_fused_shared_experts
if (
_is_cuda
and num_expert_group
and num_expert_group > 1
and envs.SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK.get()
):
(_is_cuda and num_expert_group and num_expert_group > 1)
# ROCm also admits single-group routing; CUDA's condition is unchanged.
or (_is_hip and num_expert_group)
) and envs.SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK.get():
# Opt-in: unified Triton router for DeepSeek-V3 grouped routing. Bit-exact
# with the flashinfer/AOT paths on DeepSeek-V3.2 e2e (validated); handles any
# experts-per-group (no <=32 cap). Off by default — see the env-var comment.
@@ -1636,18 +1634,26 @@ def biased_grouped_topk_gpu(
moe_fused_gate as jit_grouped_gate,
)
# The kernel wants the total width; select_experts passes a routed-only
# topk on the aiter path only (`num_routed_topk if _use_aiter else top_k`).
#
# True, not the caller's flag: an aiter runner skips the post-MoE multiply,
# so the routed weights must carry routed_scaling_factor and the shared
# slot must be 1.0. That flag describes the runner, not this kernel.
return jit_grouped_gate(
gating_output.to(dtype=torch.float32),
gating_output,
correction_bias.to(dtype=torch.float32),
topk,
topk + num_fused_shared_experts if _use_aiter else topk,
scoring_func="sigmoid",
num_fused_shared_experts=num_fused_shared_experts,
renormalize=renormalize,
routed_scaling_factor=(
routed_scaling_factor if routed_scaling_factor is not None else 1.0
),
apply_routed_scaling_factor_on_output=bool(
apply_routed_scaling_factor_on_output
apply_routed_scaling_factor_on_output=(
True
if (_use_aiter and num_fused_shared_experts > 0)
else bool(apply_routed_scaling_factor_on_output)
),
num_expert_group=num_expert_group,
topk_group=topk_group,
@@ -2210,6 +2216,10 @@ def _post_process_topk_ids(
recorder_topk_ids = topk_ids
_aiter_append = num_fused_shared_experts > 0 and _use_aiter
if _aiter_append and envs.SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK.get():
# That router emits the shared slots itself; appending again would write
# the shared id twice and evict a real routed expert.
_aiter_append = topk_ids.shape[-1] < topk_config.top_k
if _aiter_append and use_per_rank_shared_slots:
# Fused path: append shared experts AND apply the per-rank shared-slot
test/registered/kernel/moe/test_jit_grouped_topk.py
+186
View File
@@ -0,0 +1,186 @@
"""The unified Triton router, admitted on ROCm and for single-group routing.
Pins that the router selects what the torch reference selects, and that the
shared expert appears exactly once -- two places can emit it (this router, or
_post_process_topk_ids), and if both do, the id is written twice and evicts a
real routed expert while the model keeps producing plausible logits.
The reference is `biased_grouped_topk_impl`, not `select_experts` with the flag
off: on ROCm that is the aiter path, which casts the correction bias down to the
gating dtype, and GLM-5.2 keeps that bias where bf16 cannot separate neighbours.
It reorders routing on its own.
"""
import sys
import pytest
import torch
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_amd_ci(est_time=20, suite="jit-kernel-unit-test-amd")
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
E, TOPK_ROUTED, SHARED, SCALE = 256, 8, 1, 2.5
HIDDEN = 512
def _jit_routed(monkeypatch, logits, hidden, bias, groups):
"""Routed ids from select_experts with the unified router on, sorted."""
monkeypatch.setenv("SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK", "1")
from sglang.srt.layers.moe.topk import TopKConfig, select_experts
cfg = TopKConfig(
top_k=TOPK_ROUTED + SHARED,
renormalize=True,
use_grouped_topk=True,
num_expert_group=groups,
num_fused_shared_experts=SHARED,
topk_group=1,
scoring_func="sigmoid",
correction_bias=bias,
routed_scaling_factor=SCALE,
apply_routed_scaling_factor_on_output=False,
)
ids = select_experts(hidden, logits, cfg).topk_ids.long()
return ids, ids[ids < E].view(ids.shape[0], TOPK_ROUTED).sort(-1).values
def _reference_routed(logits, hidden, bias, groups):
"""Routed ids from the torch reference, sorted.
It overwrites the last slot with the shared id, so the routed experts are
what survives below E.
"""
from sglang.srt.layers.moe.topk import biased_grouped_topk_impl
_, ids = biased_grouped_topk_impl(
hidden_states=hidden,
gating_output=logits,
correction_bias=bias,
topk=TOPK_ROUTED + SHARED,
renormalize=True,
num_expert_group=groups,
topk_group=1,
num_fused_shared_experts=SHARED,
routed_scaling_factor=SCALE,
)
ids = ids.long()
return ids[ids < E].view(ids.shape[0], TOPK_ROUTED).sort(-1).values
def _inputs(tokens, dev="cuda"):
torch.manual_seed(0)
hidden = torch.randn(tokens, HIDDEN, dtype=torch.bfloat16, device=dev)
# A narrow band at a large offset, as GLM-5.2's bias is: near-equal values
# are what a router has to keep apart.
bias = (7.0 + 0.04 * torch.randn(E, device=dev)).float()
# fp32 logits: in bf16 the two round the sigmoid differently and near-equal
# rows would flip for that reason alone.
logits = torch.randn(tokens, E, dtype=torch.float32, device=dev)
return hidden, bias, logits
def _groups(request_groups: int) -> int:
"""Skip single-group cases where the gate does not admit them.
biased_grouped_topk_gpu admits one group on ROCm only; CUDA still requires
num_expert_group > 1, so a groups=1 case there never reaches the router and
the test would assert against a path it did not exercise.
"""
from sglang.srt.layers.moe import topk as topk_mod
if request_groups == 1 and not topk_mod._is_hip:
pytest.skip("single-group routing is admitted on ROCm only")
return request_groups
@pytest.mark.parametrize("groups", [1, 8])
@pytest.mark.parametrize("tokens", [6, 48, 256])
def test_jit_router_selects_what_the_reference_selects(monkeypatch, tokens, groups):
groups = _groups(groups)
hidden, bias, logits = _inputs(tokens)
want = _reference_routed(logits, hidden, bias, groups)
_, got = _jit_routed(monkeypatch, logits, hidden, bias, groups)
score = logits.sigmoid() + bias
for r in (want != got).any(-1).nonzero().flatten().tolist():
only_want = sorted(set(want[r].tolist()) - set(got[r].tolist()))
only_got = sorted(set(got[r].tolist()) - set(want[r].tolist()))
# Exact ties may break either way; nothing else may.
for x, y in zip(only_want, only_got):
assert score[r, x].item() == score[r, y].item(), (
f"row {r}: reference took {x} (score {score[r, x].item():.9f}) "
f"but the router took {y} (score {score[r, y].item():.9f})"
)
@pytest.mark.parametrize("groups", [1, 8])
def test_shared_expert_appears_exactly_once(monkeypatch, groups):
groups = _groups(groups)
hidden, bias, logits = _inputs(48)
ids, routed = _jit_routed(monkeypatch, logits, hidden, bias, groups)
assert ids.shape[-1] == TOPK_ROUTED + SHARED
shared = (ids >= E).sum(-1)
assert torch.equal(shared, torch.full_like(shared, SHARED)), (
f"shared expert appears {shared.tolist()} times a row, expected {SHARED}"
)
assert (routed[:, 1:] != routed[:, :-1]).all(), "a routed expert repeats"
@pytest.mark.parametrize("use_aiter", [True, False])
def test_router_is_asked_for_the_total_width(monkeypatch, use_aiter):
"""The width handed to the kernel, for both callers.
select_experts passes `num_routed_topk if _use_aiter else top_k`; the kernel
wants the total either way. Whichever GPU runs the suite fixes `_use_aiter`
and can only exercise one half, so pin the arithmetic here.
"""
from sglang.srt.layers.moe import topk as topk_mod
seen = {}
class _Captured(Exception):
pass
def _capture(scores, bias, topk, **kwargs):
seen["topk"] = topk
raise _Captured
monkeypatch.setenv("SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK", "1")
monkeypatch.setattr(topk_mod, "_use_aiter", use_aiter, raising=False)
monkeypatch.setattr(
"sglang.kernels.ops.moe.moe_fused_gate.moe_fused_gate", _capture
)
hidden, bias, logits = _inputs(4)
# The aiter caller hands over routed-only; everyone else the total.
topk_in = TOPK_ROUTED if use_aiter else TOPK_ROUTED + SHARED
with pytest.raises(_Captured):
topk_mod.biased_grouped_topk_gpu(
hidden_states=hidden,
gating_output=logits,
correction_bias=bias,
topk=topk_in,
renormalize=True,
# 8 groups, not 1: the arithmetic under test is the same either
# way, and only this value reaches the router on both platforms.
num_expert_group=8,
topk_group=1,
num_fused_shared_experts=SHARED,
routed_scaling_factor=SCALE,
)
assert seen["topk"] == TOPK_ROUTED + SHARED, (
f"_use_aiter={use_aiter}: caller passed topk={topk_in}, kernel was asked "
f"for {seen['topk']} slots, expected {TOPK_ROUTED + SHARED} "
f"(routed {TOPK_ROUTED} + shared {SHARED})"
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))