[AMD] gfx950 small-batch MoE: expert-count gate for the small sort and FP8 block-scale small-M kernel

[chuyeh]

Motivation

Qwen3.5-397B-A17B-FP8 on MI355X is slower than B200 at low concurrency in the AgentX benchmark (InferenceX PR #3602). One factor is the MoE sort from #36559: past its compact range, its distributed kernel scans every 64-expert chunk per pair, so with Qwen3.5's 513 experts it is slower than AITER's sort for mid-size decode batches (6 to 23 tokens). At TP8 with 16 verify tokens, going back to AITER's sort saves 1.7 ms per decode cycle.

This PR does two things:

1. Fix the slowdown: past the compact range, MoE: small-batch sorting path with fused mxfp8 quantisation #36559's distributed kernel is only used up to 256 experts, and larger models go back to AITER. MiniMax-M3 (129 experts), which the kernel was written for, keeps it because it is faster there.
2. Add FP8 block-scale support to the [AMD] Small-M MXFP4 fused-MoE kernel for gfx950 (Qwen) #40204 small-M MoE kernel, which needs no sort and is much faster at small batches but only supported MXFP4.

With this PR, the interactivity improves by about 20% at concurrency 1 and 4.

Modifications

• moe_sorting_small.py: past the compact kernel's range (min(64, 2 * block_size) (token, expert) pairs), the distributed kernel only runs for up to 256 experts; above that, the sort and the fused MXFP8 quantization go back to AITER (+13 / -3).
• smallm_moe.hip: p1_body / p2_body take an F8 template flag. Shuffled FP8 weights have the same byte layout as fp4x2, so the load path is shared. e4m3 is converted to bf16 at scale 1.0 and each 128x128 block's fp32 scale is applied to the partial sum (cvt_scalef32_pk_bf16_fp8 only applies the scale's exponent). New kernels cover TP4 (intermediate 256) and TP8 (intermediate 128); the MXFP4 kernels are unchanged.
• smallm_moe_gfx950/__init__.py: takes the FP8 path for shuffled float8_e4m3fn w13 / w2 with contiguous fp32 block scales, up to 20 tokens at TP4 and 48 at TP8. SGLANG_ROCM_SMALLM_MOE=0 still turns the kernel off.
• moe_runner/aiter.py: passes w13_scale / w2_scale to smallm_moe_supported (+2 lines).
• Tests (MI35x suite): TestSmallMMoeFp8Gfx950 checks both shapes against a dequantized reference at 1 / 3 / 8 tokens and the cap, and that cap+1, per-channel or non-contiguous scales, and unshuffled weights fall back. New test_moe_sorting_small.py compares the sort and the fused MXFP8 quant with AITER on the compact path (Qwen3.5 shape) and the distributed path (MiniMax-M3 shape, 13 to 51 tokens), and checks the fallback at 513 experts.

Accuracy Tests

All runs use image rocm/sgl-dev:v0.5.20-rocm720-mi35x-20260927 (same image as lmsysorg/sglang-rocm with that tag), InferenceX recipe flags, GSM8K 1319 questions.

Model	Setup	Baseline	This PR
Qwen/Qwen3.5-397B-A17B-FP8	TP4, real draft acceptance, sglang.test.few_shot_gsm8k	0.933	0.941
amd/Qwen3.5-397B-A17B-MXFP4-AttnFP8-V2 (sort change)	TP2, lm-eval 5-shot, thinking off	0.977	0.979

The FP8 kernel's rel L2 against the dequantized reference is at most 7.4e-4 over 5 seeds in the unit test. The sort output matches AITER's layout in every path, so the sort change only affects speed.

Benchmarking and Profiling

Kernel microbenchmarks

MI355X, one GPU, CUDA graph replay, Qwen3.5-397B shapes (E=513, 11 slots per token, hidden 4096).

FP8 MoE layer, AITER (per_128x128) vs the FP8 small-M kernel, in µs:

Tokens	TP4 AITER	TP4 kernel	Change	TP8 AITER	TP8 kernel	Change
1	61.3	26.6	−57%	42.3	22.7	−46%
4	66.5	43.4	−35%	46.1	30.1	−35%
8	74.0	63.7	−14%	57.6	41.4	−28%
16	106.3	97.5	−8%	75.3	60.4	−20%
20	117.3	115.3	−2% (cap)	77.3	68.7	−11%
24	130.1	133.1	AITER used	84.3	78.0	−7%

TP8 caps, 60 layers rotated so the weights are not cache resident:

Tokens	AITER	Kernel	Change
24	68.7	65.2	−5%
28	78.2	75.8	−3%
32	83.5	82.9	~0
36	119.0	92.0	−23%
48	124.6	111.6	−10% (cap)

AITER moves to a slower configuration from 36 tokens, so the TP8 cap is 48 rather than 32.

MoE sort past the compact range, in µs, block size as AITER picks it. MXFP rows include the MXFP8 quant in both arms (AITER runs it as a second launch). Up to the compact range the sort is unchanged, and the FP8 small-M path does not use the sort:

Model	Tokens	Pairs	AITER	#36559	This PR
Qwen3.5 MXFP4 (sort + quant)	7	77	10.5	10.3	10.5
	16	176	11.7	15.9	11.7
	23	253	13.5	15.5	13.5
Qwen3.5 FP8 (sort only)	16	176	6.9	15.7	6.9
	23	253	6.8	15.4	6.8
MiniMax-M3 MXFP8 (E=129, top-4 + shared, sort + quant)	16	80	12.0	8.1	8.1
	35	175	21.2	13.1	13.1
	51	255	24.6	13.0	13.0

AITER picks block size 32 for the MXFP shapes and mostly 16 for Qwen3.5 FP8, where the compact kernel covers only 1 to 2 tokens.

AgentX benchmark

Setup:

• Hardware: MI355X, TP4.
• Model and recipe: Qwen/Qwen3.5-397B-A17B-FP8 with the InferenceX recipe qwen3.5-fp8-mi355x-sglang-agentic, image rocm/sgl-dev:v0.5.20-rocm720-mi35x-20260927.
• Server flags: as in the recipe. FP8 KV cache, EAGLE 3 / 1 / 4 with simulated acceptance 3.39, admission and decode CUDA graph batch at 2x concurrency.
• Client: the CI AgentX command, aiperf profile --scenario inferencex-agentx-mvp on semianalysis_cc_traces_weka_062126_256k, 3600 s per run.
• Order: one run at a time on the same GPUs (this PR c1, baseline c1, baseline c4, this PR c4).

Baseline is the image as shipped: AITER's FP8 MoE with #36559's sort (SGLANG_ROCM_SMALLM_MOE=0). With this PR, every decode batch at these concurrencies runs the FP8 small-M kernel and skips the sort, so the gain here comes from the kernel.

Concurrency 1:

Metric	Baseline	This PR	Change
ITL p50 / p90 (ms)	3.69 / 4.07	3.01 / 3.50	−18% / −14%
TTFT p50 / p90 (ms)	389 / 861	373 / 859	−4% / ~0
Requests completed in 3600 s	557	628	+13%
Output throughput (tok/s)	150	177	+18%
Interactivity p50 / p90 (tok/s/user)	271 / 283	332 / 349	+23% / +23%

Concurrency 4:

Metric	Baseline	This PR	Change
ITL p50 / p90 (ms)	4.10 / 5.32	3.44 / 4.34	−16% / −18%
TTFT p50 / p90 (ms)	551 / 1096	519 / 947	−6% / −14%
Requests completed in 3600 s	926	992	+7%
Output throughput (tok/s)	243	263	+8%
Interactivity p50 / p90 (tok/s/user)	244 / 278	291 / 333	+19% / +20%

Each EAGLE verify pass sends 4 tokens per running request through the MoE. At TP4 the kernel therefore covers running batches of 5 or fewer, which is concurrency 4 and below here; concurrency 8 and up mostly uses AITER.

Decode step time

Time for one EAGLE draft + verify cycle at a fixed batch size, same server flags and simulated acceptance.

TP4, image as shipped vs this PR, mean of two runs with the GPU sets swapped, in ms:

Batch size	Baseline	This PR	Change
1	11.89	9.51	−20%
4	16.51	13.08	−21%
8	16.35	16.41	~0 (past the cap)

TP8, this PR with SGLANG_ROCM_SMALLM_MOE=1 vs 0 (kernel vs AITER's MoE and sort), two runs per arm in ABBA order. Run-to-run noise is up to 0.23 ms:

Batch size	Verify tokens	Change per cycle
4	16	−0.78 ms
6	24	−0.51 ms
7	28	−0.30 ms
8	32	−0.07 ms (noise)

With the kernel off at batch size 4 (176 pairs), #36559's sort takes 33.00 ms per cycle and this PR's 31.28 ms, the 1.7 ms in Motivation.

Checklist

• Format your code according to the Format code with pre-commit guide.
• Add unit tests.
• Update documentation.
• Provide accuracy and speed benchmark results.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ✅ Run #37556056254
Latest PR Test (Extra): ❌ Run #37556056047
Latest PR Test (AMD ROCm 10): ❌ Run #37556056190

[yichiche]

@zcnrex Hi, this is the follow-up of #36559, can you help review?

[yichiche]

base-c-test-8-gpu-b300 / base-c-test-8-gpu-b300 (0) server is off. Ref: #42463

[yichiche]

• Affected Scope: NVIDIA behavior is unchanged. It does not run AiterRunner or the small-sort wrapper.
• AMD Guard: The expert cutoff is the AITER small-sort wrapper, with no gfx check. The FP8 kernel runs only on gfx950, and only from AiterRunner.run.
• Unit Test Quality: The tests use unittest, check real results, and stay under 500s. Coverage is mixed at sorting_impl_wrapper, which does not fail this row.
• Accuracy evidence: The PR body reports an accuracy result.
• Performance evidence: The PR body reports a performance result.
• Flags: No new environment variable.

LGTM.

[chuyeh]

/rerun-failed-ci