[kevin-mii]

Motivation

At c1 with DSpark verifying 8 rows, MiniMax-M3.1-Flash's FFN launch (#56) takes about 60 µs a sparse layer at TP4 on MI350X, and SGLang runs o_proj in front of it as two more kernels (about 9 µs). Timing the launch's stages per CTA showed where the time went:

• The router waited for the norm to write and flag XN before it could read its 384 columns, one more hop in the chain in front of the routed weights.
• Each lane of an up unit or a down slot loaded its own column's B operand (XN, MID) from global memory. A slot of cnt columns read its last column 16 - cnt more times, and at 8 rows and 15-20 experts a slot has one or two columns. Those repeats were most of a CTA's vector memory traffic: 24 KB of XN a wave a unit against 12 KB of weights. Weights served from the MALL timed the same as from HBM.
• Once X came from LDS, an up unit waited on its own weights for about 3 µs of its 5.6, while the CTA's loads sat idle during its compute, reduction and stores.
• o_proj's per-token FP8 quant and GEMM ran as separate launches; ATOM's full-layer MiniMax-M3 kernel (and vLLM's port of it, [ROCm][Perf] Replace MiniMax-M3 MonoKernel with ATOM full-layer implementation vllm-project/vllm#63398) runs o_proj inside its launch.

Modifications

All in minimax_m31/ (and the byte count on a16w4's loads). Every output is bit-identical to main with SGLang's o_proj in front of it, at 1, 3, 8, 9 and 16 rows.

• Route on the norm's scale. The norm publishes each token's rstd (a pair) before it writes XN, and the router normalizes its own columns from SRES. The table no longer reaches XN through the router, so routed CTAs wait for XRDY after the table.
• B operands from LDS. XN goes to LDS once a CTA, after XRDY, for steps of up to 8 rows (16 don't fit); the up stage reads it there. A down slot's wave first copies only the slot's own MID columns to LDS (a 16-byte chunk a lane), then reads B there. The MID rows are the picks', in the LDS pool XN used.
• An up unit's weights load during the unit before. A routed unit issues its next unit's weights after its X reads (the first unit's go out after the table); past the CTA's last unit they read through a zero-sized buffer.
• o_proj in the launch. The push CTAs run it from the attention rows on SGLang's rowwise-FP8 copy of the MXFP8 weight (_ptpc_weight, _ptpc_scale): AITER's per-token FP8 quant (amax / 448, v_rcp_f32, clamp, RNE), then smallm_fp8_gemm's waves, K steps, accumulator chains and wave order for the row count, so the partial is SGLang's to the bit. The spec makes o_proj return its input while SGLang's attention runs, and refuses a layer whose index values add to o_proj's output or whose o_proj has no rowwise-FP8 copy. test_contract pins the attributes it reads.

Accuracy tests

MI350X (gfx950), SGLang af6d50a85e (as in #56) on a ROCm 10 image, MiniMaxAI/MiniMax-M3.1-Flash, TP4:

• out and res_out at 1, 3, 8, 9 and 16 rows (TP=1 build): identical to main fed SGLang's o_proj output (AITER's quant + smallm_fp8_gemm), bit for bit.
• test/kernels/test_minimax_m31_ffn.py: 8 passed. New: the launch's o_proj partial equals smallm_fp8_gemm's at 8 and 16 rows (its two wave splits), and is checked against a native reference before the existing stage checks.
• test/unit: 288 passed. One contract failure against this SGLang branch: deepseek_v41/spec.py reads mxfp8_aiter_preshuffled, which that branch lacks; unrelated (as on [Feat] minimax_m31: Run MiniMax-M3.1-Flash's decode FFN as one monokernel launch a layer #56).
• test/e2e/test_minimax_m31_server.py: 6 passed. Chat GSM8K past the 16-row width scores 0.98 (DSpark verify) and 0.975 (decode); 48,792 checked launches are all within NUMERICS; a stalled rank still stops the server without tokens. The check's residual difference is the same distribution as on main (median 0.0054, p99.9 0.021-0.024 over three runs of this branch), but its single worst launch varies from run to run (0.029, 0.047 and 0.077 against the 0.08 bound, in a different layer each time), so that bound is close to flaky on any revision.
• Chat GSM8K (sglang.test.run_eval, 200 questions, greedy, DSpark verify server, real acceptance): 0.975 with this PR, 0.98 on main (one question; std 0.156), acceptance 4.8-5.0 and 4.7-5.1 tokens per step.

Speed tests

c1 decode on the same stack (as in #56): DSpark acceptance held at 3 tokens per step (SGLANG_SIMULATE_ACC_LEN=3), 11K-token prompt, 1,024 greedy tokens, median of 4, arms alternated:

	decode tok/s
main	423.9, 428.9
without the o_proj commit	447.0, 453.2 (+5.6%)
this PR	455.1, 461.2 (+7.4%)

The launch with o_proj: TP4 FfnRunner, 8 rows, every rank routing alike, the router bias steering the rows to a set number of experts. Weights (o_proj's too) rotate through four copies, so they stream from HBM as a server's layers do (one copy stays in the 256 MB MALL). main and the first three commits time SGLang's quant + GEMM + the launch; this PR times the launch alone. Median of 5 × 200 launches:

distinct experts	main	without the o_proj commit	this PR
8	54.7 µs	48.1 µs	44.8 µs (−18%)
14	59.6 µs	54.6 µs	51.8 µs (−13%)
18	66.0 µs	61.2 µs	58.2 µs (−12%)

main's launch alone at 20 experts (62.8 µs) matches the 62.3 µs a launch measured in the server.

Not covered

• MI355X, and TP other than 4.
• Steps of 9-16 rows keep XN in global memory (16 rows of it don't fit in LDS with the rest).
• The attention side is still SGLang's: the next layer's add-RMSNorm, qkv GEMM, q / k norms, RoPE, cache writes and the sparse decode run as about 7 kernels a layer (about 40 µs). ATOM's MiniMax-M3 layer kernel runs them in its launch, but on its own KV and index cache layouts; folding them here is a follow-up.
• Tried and dropped, each measured at TP4 on the same bench: polling every slot's flags once a CTA and staging all of MID together (slower: the down stage no longer overlaps the up stage's tail); prefetching the next down slot's weights or flags (no gain: each slot's acquire waits on every load in flight); loading only the live lanes' B operands (slower).

🤖 Generated with Claude Code