diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 21586f29a..97c04a694 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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, diff --git a/test/registered/amd/test_qwen35_moe_softmax_topk.py b/test/registered/amd/test_qwen35_moe_softmax_topk.py new file mode 100644 index 000000000..c6ed7eb75 --- /dev/null +++ b/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()