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

Co-authored-by: JohnQinAMD <yanyuan.qin@amd.com>
Co-authored-by: RuibinCheung <ruibzhan@amd.com>
Co-authored-by: Zhang, Jiejing <jiejing.zhang@amd.com>

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

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,

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
+        (_is_cuda and num_expert_group and num_expert_group > 1)
-        and num_expert_group
+        # ROCm also admits single-group routing; CUDA's condition is unchanged.
-        and num_expert_group > 1
+        or (_is_hip and num_expert_group)
-        and envs.SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK.get()
+    ) 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

@@ -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"]))