Repository navigation
[AMD] Small-M MXFP4 fused-MoE kernel for gfx950 (Qwen) - #40204
Merged
Merged
Conversation
Two HIP kernels for the MXFP4 MoE at 1-40 tokens per rank on MI355X (hidden 4096, per-rank intermediate 256/512, 10 or 11 expert slots), built with hipcc on first use and launched through ctypes from the aiter MoE runner. Same weight/scale layouts as aiter fused_moe, bf16 activations, on by default on gfx950 (SGLANG_ROCM_SMALLM_MOE=0 turns it off), falls back to aiter on any build/load/launch failure and above 40 tokens. Adds a registered AMD unit test against a dequantized torch reference.
zijiecode
requested review from
BBuf,
DarkSharpness,
Edwardf0t1,
Fridge003,
HaiShaw,
HydraQYH,
Ying1123,
celve,
ch-wan,
ispobock,
merrymercy and
yuan-luo
as code owners
September 18, 2026 16:56
zijiecode
force-pushed
the
pr/smallm-moe-gfx950
branch
from
September 20, 2026 03:25
b6f7500 to
312b0a0
Compare
ROCm 10 images ship a second libamdhip64 in _rocm_sdk_devel; the unversioned
CDLL("libamdhip64.so") picked it and every launch on a torch stream failed with
709 (context is destroyed). Initialize CUDA first and bind to torch's copy via
find_loaded_library(). Unit test passes on rocm720 and rocm10 MI35x images.
zijiecode
force-pushed
the
pr/smallm-moe-gfx950
branch
from
September 20, 2026 03:26
312b0a0 to
175fef3
Compare
This was referenced Sep 20, 2026
The intermediate-512 (TP2) shape crosses over to aiter's flydsl path much earlier (~20 tokens vs ~44), so it stays on aiter for now and gets its own cap in a follow-up. The kernel entry points for 512 remain in the source.
HaiShaw
approved these changes
Sep 24, 2026
This was referenced Sep 24, 2026
kevin-mii
added a commit
to zcnrex/sglang
that referenced
this pull request
Sep 24, 2026
Conflict in moe_runner/aiter.py: main (sgl-project#40204) added the gfx950 small-M MXFP4 fused-MoE early return just before the fused_moe call; this branch picks the fused_moe activation (popping one passed via fused_moe_kwargs) at the same spot. Kept main's early return first and this branch's activation selection after it, since only the fused_moe call reads it. The small-M path requires hidden 4096 and topk 10/11, so MiniMax-M3 (hidden 6144, topk 5) never takes it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
For decode batches of 1 to ~40 tokens per rank, the MXFP4 MoE is the biggest cost per step on MI355X, and AITER
fused_moeonly uses a small part of the HBM bandwidth there. This PR adds a HIP kernel pair for these shapes; above 40 tokens AITER is still used.Modifications
python/sglang/kernels/ops/moe/smallm_moe_gfx950/(__init__.py,smallm_moe.hip): built withhipcconfirst use and launched through ctypes, no build-time or wheel change.
v_cvt_scalef32_pk_bf16_fp4, MFMA16x16x32 bf16,
silu(g) * uwritten as bf16. Phase 2: down GEMM per (token, 64-row slice) withv_dot2_f32_bf16,top-k weights applied in fp32.
moe_runner/aiter.py: one guarded call beforefused_moe, taken for FP4x2 weights, bf16 hidden 4096, per-rankintermediate 256 (the TP4 shape), 10 or 11 slots, silu,
tok <= 40, and none ofexpert_mask/doweight_stage1/bias /
a1_scale/no_combine/num_local_tokens. Everything else is unchanged. The TP2 shape (intermediate 512)crosses over to AITER much earlier (~20 tokens) and is left to a follow-up.
SGLANG_ROCM_SMALLM_MOE=0turns it off.test/registered/amd/test_smallm_moe_gfx950.py: dequantized torch reference (rel L2 < 5e-3at 1 / 3 / 16 / 40 tokens), fallback above 40 and for the TP2 shape, the off switch, graph replay == eager. Passes on
v0.5.19-rocm720-mi35x-20260915andv0.5.18-rocm10-mi35x-20260902.Accuracy Tests
amd/Qwen3.5-397B-A17B-MXFP4-AttnFP8-V2(revision e17e5f0), MI355X TP4, InferenceX recipe (FP8 KV cache, EAGLE MTP3 steps / 4 draft tokens, real draft), lm-eval
local-chat-completionswith chat template, 5-shot GSM8K, 16384-tokenoutput budget, temperature 0.
fused_moe)GPQA diamond (198 questions, repeat 8), same server, thinking on, temperature 0.6 / top_p 0.95 / top_k 20, 32768-token
output budget:
fused_moe)Benchmarking and Profiling
MI355X, TP4 / EP1, InferenceX MI355X recipe for this model (image
lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260915,--kv-cache-dtype fp8_e4m3 --mem-fraction-static 0.80, EAGLE MTP 3 / 1 / 4,--speculative-draft-model-quantization quark_mxfp4in both configurations), tree = maine4cbb28e+ #39901 + #39902 (+ this PR), full 3600 s aiperf runs onthe same GPUs, one configuration after the other.
Concurrency 1:
Concurrency 8:
Kernel microbenchmark (MI355X, HIP-graph replay with the MALL flushed between calls, Qwen3.5-397B TP4 shape,
E=513, 11 slots):fused_moe(µs)Checklist
CI States
Latest PR Test (Base): ⏳ Run #35960012293
Latest PR Test (Extra): ❌ Run #35960011935
Latest PR Test (AMD ROCm 10): ⏳ Run #35960012296