[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
-        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