[AMD] Use Triton softmax routing for Qwen3.5 on gfx950 (#39986)

python/sglang/srt/layers/moe/topk.py

@@ -110,6 +110,7 @@ from sglang.srt.utils import (
     get_compiler_backend,
     is_cpu,
     is_cuda,
+    is_gfx95_supported,
     is_hip,
     is_musa,
     is_npu,
@@ -125,6 +126,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 _is_cuda = is_cuda()
 _is_hip = is_hip()
+_is_gfx95 = is_gfx95_supported()
 _is_cpu = is_cpu()
 _is_cpu_amx_available = cpu_has_amx_support()
 _is_xpu = is_xpu()
@@ -151,6 +153,30 @@ _RENORMALIZE_SUM_EPSILON = 1e-20
 _skip_hip_pad_mask = get_bool_env_var("SGLANG_MORI_NO_PAD_MASK", "False")
 
 
+def _use_rocm_triton_softmax_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    correction_bias: Optional[torch.Tensor],
+    num_fused_shared_experts: int,
+    packed_out: Optional[torch.Tensor],
+) -> bool:
+    """Use the lower-latency Triton router for Qwen3.5 decode-sized rows."""
+    return (
+        _use_aiter
+        and _is_gfx95
+        and hidden_states.shape[1] == 4096
+        and hidden_states.dtype == torch.bfloat16
+        and gating_output.shape[0] <= 128
+        and gating_output.shape[1] == 512
+        and gating_output.dtype == torch.bfloat16
+        and topk == 10
+        and correction_bias is None
+        and num_fused_shared_experts == 0
+        and packed_out is None
+    )
+
+
 if _is_cuda:
     try:
         from flashinfer.fused_moe import fused_topk_deepseek as _fused_topk_deepseek
@@ -991,7 +1017,15 @@ def fused_topk(
     topk_ids = torch.empty(M, topk, dtype=torch.int32, device=hidden_states.device)
 
     if scoring_func == "softmax":
-        if _use_aiter:
+        use_rocm_triton = _use_rocm_triton_softmax_topk(
+            hidden_states,
+            gating_output,
+            topk,
+            correction_bias,
+            num_fused_shared_experts,
+            packed_out,
+        )
+        if _use_aiter and not use_rocm_triton:
             # Use fused_topk instead of topk_softmax to auto dispatch to the correct kernel
             topk_weights, topk_ids = aiter_fused_topk(
                 hidden_states,
@@ -1018,7 +1052,7 @@ def fused_topk(
                 num_token_non_padded=num_token_non_padded,
             )
         # ===== END TO BE REFACTORED ====
-        elif _is_cuda:
+        elif _is_cuda or use_rocm_triton:
             # Unified Triton router (subsumes the AOT topk_softmax CUDA kernel).
             from sglang.kernels.ops.moe.moe_fused_gate import (
                 moe_fused_gate as _jit_moe_fused_gate,

test/registered/amd/test_qwen35_moe_softmax_topk.py

@@ -0,0 +1,127 @@
+"""ROCm coverage for the decode-sized Qwen3.5 MoE softmax router."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.layers.moe import topk as topk_module
+from sglang.srt.utils import is_gfx95_supported
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
+
+
+@unittest.skipUnless(
+    torch.cuda.is_available() and torch.version.hip and is_gfx95_supported(),
+    "requires AMD gfx95",
+)
+class TestQwen35MoeSoftmaxTopK(CustomTestCase):
+    def test_triton_dispatch_matches_aiter(self):
+        from aiter.fused_moe import fused_topk as aiter_fused_topk
+
+        for num_tokens in (4, 12, 128):
+            with self.subTest(num_tokens=num_tokens):
+                torch.manual_seed(num_tokens)
+                hidden_states = torch.randn(
+                    num_tokens, 4096, device="cuda", dtype=torch.bfloat16
+                )
+                router_logits = torch.randn(
+                    num_tokens, 512, device="cuda", dtype=torch.bfloat16
+                )
+                ref_weights = torch.empty(
+                    num_tokens, 10, device="cuda", dtype=torch.float32
+                )
+                ref_ids = torch.empty(num_tokens, 10, device="cuda", dtype=torch.int32)
+                ref_weights, ref_ids = aiter_fused_topk(
+                    hidden_states,
+                    router_logits,
+                    10,
+                    True,
+                    topk_ids=ref_ids,
+                    topk_weights=ref_weights,
+                )
+
+                with (
+                    patch.object(topk_module, "_use_aiter", True),
+                    patch.object(topk_module, "_is_gfx95", True),
+                    patch.object(
+                        topk_module,
+                        "aiter_fused_topk",
+                        side_effect=AssertionError("AITER top-k should be bypassed"),
+                        create=True,
+                    ),
+                ):
+                    weights, ids = topk_module.fused_topk(
+                        hidden_states,
+                        router_logits,
+                        topk=10,
+                        renormalize=True,
+                    )
+
+                torch.testing.assert_close(ids, ref_ids, rtol=0, atol=0)
+                torch.testing.assert_close(weights, ref_weights, rtol=1e-5, atol=1e-6)
+
+    def test_dispatch_envelope_is_narrow(self):
+        hidden_states = torch.empty(128, 4096, device="cuda", dtype=torch.bfloat16)
+        logits = torch.empty(128, 512, device="cuda", dtype=torch.bfloat16)
+        packed = torch.empty(1, device="cuda")
+        with (
+            patch.object(topk_module, "_use_aiter", True),
+            patch.object(topk_module, "_is_gfx95", True),
+        ):
+            self.assertTrue(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states, logits, 10, None, 0, None
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    torch.empty(129, 4096, device="cuda", dtype=torch.bfloat16),
+                    torch.empty(129, 512, device="cuda", dtype=torch.bfloat16),
+                    10,
+                    None,
+                    0,
+                    None,
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states, logits.float(), 10, None, 0, None
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states[:, :2048], logits, 10, None, 0, None
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states, logits, 8, None, 0, None
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states,
+                    logits,
+                    10,
+                    torch.empty(512, device="cuda"),
+                    0,
+                    None,
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states, logits, 10, None, 1, None
+                )
+            )
+            self.assertFalse(
+                topk_module._use_rocm_triton_softmax_topk(
+                    hidden_states, logits, 10, None, 0, packed
+                )
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()