From 6657f7d8449f7c739906b6f8ebcb016dddcbceef Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Sun, 13 Sep 2026 07:21:28 +0800 Subject: [PATCH] [MoE][ROCm] Admit the unified Triton router on ROCm, including single-group routing (#38328) Co-authored-by: JohnQinAMD Co-authored-by: RuibinCheung Co-authored-by: Zhang, Jiejing --- .../sglang/kernels/ops/moe/moe_fused_gate.py | 2 +- .../attention/wave_ops/decode_attention.py | 52 +++++ python/sglang/srt/layers/moe/topk.py | 30 ++- .../kernel/moe/test_jit_grouped_topk.py | 186 ++++++++++++++++++ 4 files changed, 259 insertions(+), 11 deletions(-) create mode 100644 test/registered/kernel/moe/test_jit_grouped_topk.py diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index 624c331b2..f6933e6f0 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/moe_fused_gate.py @@ -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 diff --git a/python/sglang/srt/layers/attention/wave_ops/decode_attention.py b/python/sglang/srt/layers/attention/wave_ops/decode_attention.py index c76bee9af..2cdf29976 100644 --- a/python/sglang/srt/layers/attention/wave_ops/decode_attention.py +++ b/python/sglang/srt/layers/attention/wave_ops/decode_attention.py @@ -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, diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index f2f17c1ef..c090ef227 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -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 diff --git a/test/registered/kernel/moe/test_jit_grouped_topk.py b/test/registered/kernel/moe/test_jit_grouped_topk.py new file mode 100644 index 000000000..08bb8866c --- /dev/null +++ b/test/registered/kernel/moe/test_jit_grouped_topk.py @@ -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"]))