minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
jacky.cheng
2026-09-21 12:02:29 -07:00
committed by GitHub
parent f702a0be29
commit 90b3f8544c
2 changed files with 163 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/topk.py
+36 -2
View File
@@ -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
+127
View File
@@ -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()