[Feat] minimax_m31: Run MiniMax-M3.1-Flash's decode FFN as one monokernel launch a layer by kevin-mii · Pull Request #56 · sgl-project/sgl-amd-mono

kevin-mii · 2026-10-24T04:00:27Z

Motivation

On SGLang's standard path, MiniMax-M3.1-Flash decodes at low concurrency with about seven launches per sparse layer after attention:

the attention all-reduce with residual add, norm and quant
the router, fused with the shared expert's gate/up
the shared expert's activation and down projection
the routed MoE's two stages
the MoE all-reduce

At c1 (TP4, MI350X, 150K context, DSpark with 7 drafts) that chain takes about 76 µs per layer, roughly half of an 8.66 ms step. vLLM's MiniMax-M3 ROCm path runs the same work as one persistent FlyDSL launch per layer (vllm/models/minimax_m3/amd/specialized_decode/mono), and MiniMax-M3.1-Flash has the same text architecture.

This PR brings that launch here for M3.1's decode and DSpark verify steps of up to 16 rows, through the monokernel interface (--monokernel-backend amd).

Modifications

minimax_m31/ vendors vLLM's FFN launch at vllm-project/vllm@a50259bb (Apache-2.0, attribution in the package docstring). The first commit copies it with only the imports rewritten, so each later commit is a reviewable change against upstream:
- The launch builds on the plugin's own common primitives (ops, plan, sync, peer, epoch, guard), not on vLLM's copies of them. The occupancy check loads the HIP runtime through common.peer.hip_runtime().
- Fail-stop as every other model's launches. Waits run under common.guard: SGLang's deadline (BuildContext.deadline_s), the peers' failure flags, and a stop at the null page on failure, so no token of a failed step is sent. SGLANG_AMD_MONO_FAULT_INJECT works. vLLM's own error words and step-end handshake are gone; vLLM reads them in a host executor that SGLang has no place for.
- No epoch limit. The attention seam's triple pairs carry 16-bit tags. vLLM resets the epoch from its host executor every ~2,000 steps; here the tags cycle 1 .. 65535 per parity on the device, and the last launch of a cycle zeroes that parity's receive half past its epoch barrier.
- The shared expert runs on the checkpoint's own MXFP8 values, from tensors of its own (a16w8.py). vLLM reads it as MXFP4 expert 128 of the routed tensors. On M3.1 that requantization has about 4.6× the output error of FP8 (relative 0.18–0.19 vs 0.04), and in a test build it halved DSpark's real acceptance at one request (3.6 → 2.0 tokens per step) while GSM8K still passed. Routed slots keep A16W4.
minimax_m31/spec.py: the forward runs SGLang's embedding, dense layers, input norms and attention, then one FFN launch per sparse layer, which hands its output and residual back to SGLang's residual stream. SGLang's final norm (with DSpark's last-layer capture) and logits close the step. Attention and KV stay SGLang's, so NVFP4 and FP8 KV and HiCache are untouched. The load hook keeps each rank's MXFP8 shared-expert shard. SGLANG_AMD_MONO_CHECK=1 runs SGLang's layer on a copy of each sparse layer's input in eager steps and logs the differences; NUMERICS comes from those runs.
minimax/spec.py: MiniMax-M3 and M3.1 share the MiniMaxM3Sparse* architectures, and SGLang registers one spec per architecture. This spec states only what both need, hands each checkpoint to its own spec, and re-checks that spec's requirements as SGLang's guard does. Unit tests pin that split to SGLang's Requirements and MonoSpec fields and the two models' device limits, so a field SGLang adds or a divergence fails CI. minimax/check.py holds the check-mode comparison both models log.
test/e2e/test_minimax_m31_server.py (new): the DSpark verify and decode servers' routing lines and chat GSM8K past the 16-row width, check mode against NUMERICS, and a stalled rank stopping the server without tokens.
README and CODEOWNERS: M3.1's row, launch, requirements and test command.

Accuracy tests

MI350X (gfx950), SGLang af6d50a85e (sgl-project/sglang#45936 on #45925) on a ROCm 10 image (torch 2.11, Triton 3.8, FlyDSL 0.4.0, HIP 7.15), MiniMaxAI/MiniMax-M3.1-Flash, TP4:

test/kernels/test_minimax_m31_ffn.py: 6 passed. Every stage at 1, 3, 8 and 16 rows (TP=1 build): the routed experts against AITER's fused MoE on SGLang's post-load MXFP4 layout, the shared expert against torch on its MXFP8 values. A launch whose wait times out stops (fails with a guard that does not stop), and a tag cycle ends with its receive half zeroed and every launch matching one at epoch 0 (fails without the clear).
test/unit: 140 passed. One contract failure against this SGLang branch: deepseek_v41/spec.py reads mxfp8_aiter_preshuffled, which that branch lacks; unrelated to this PR.
test/e2e/test_minimax_m31_server.py: 6 passed. Both servers' startup lines show their tested batches on the launch; chat GSM8K past the 16-row width scores 0.985 (DSpark verify) and 0.975 (decode); every checked launch is within NUMERICS; a rank stalled by SGLANG_AMD_MONO_FAULT_INJECT stops all four GPUs with faults at the null page and no tokens sent.
Chat GSM8K (sglang.test.run_eval, 200 questions, greedy, DSpark verify server, real acceptance): 0.975 with --monokernel-backend amd, 0.975 on the standard path. 8-shot completion prompts score under 0.1 on either path, so the gate is chat.
Check mode, two runs of 16 GSM8K prompts (49,476 and 44,004 launches), worst case: residual rel 0.042, cosine 0.99831; MLP output on the rows routed alike rel 0.28, cosine 0.66; 6.0–6.1% of rows routed to other experts than SGLang's.

Speed tests

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

	decode tok/s
standard path	390.8
`--monokernel-backend amd`	422.3 (+8.1%)

Chat GSM8K above (4 requests at a time, real acceptance): 24.8 s vs 26.8 s on the standard path, 1,312 vs 1,272 output tok/s.

The port to common and the guard leave the launch's time unchanged (TP=1 build, 1 / 4 / 8 / 16 rows: 26.3 / 45.7 / 65.8 / 108.8 µs against 27.0 / 45.9 / 65.7 / 108.1 µs vendored).

On an earlier revision, AgentX c1 (900 s, TP4, DSpark with 7 drafts at golden AL 4.04, SemiAnalysisAI/InferenceX#4445's c1 cell) measured 6,013 → 6,151 tok/s/GPU, p90 449.5 → 470.2 and median 495.2 → 548.0 tok/s/user. c1 screens at 900 s move by several percent between identical runs.

Not covered

MI355X, and TP other than 4.
A TP4 kernel-level test. TP4 is covered end to end.
Steps of more than 16 rows (DSpark at 3+ requests), which run SGLang's path.
On a ROCm 10 image, #45936's bf16 GEMM table carries hipBLASLt indices that fit its ROCm 7.2 image only, so verify-graph capture fails until those rows are left out (AITER_CONFIG_GEMM_BF16). These tests ran with them left out on both paths.
SGLang could select a spec per checkpoint, which would let minimax/spec.py drop its re-check of the requirements.
The launch reads at about 3.3 TB/s at 8 rows on MI350X. Its front stages (norm → router → route → table) cost about 12–16 µs before any routed weight is read, and a routed tile's loads wait for the previous tile's tail. Overlapping them across tiles by unrolling was 6% slower end to end; that's a follow-up.

🤖 Generated with Claude Code

vLLM's MiniMax-M3 specialized decode runs each sparse layer's post-attention work (attention TP all-reduce, residual add, Gemma norm, router, top-4, shared and routed A16W4 experts, MoE TP all-reduce) as one persistent FlyDSL launch for up to 16 rows. MiniMax-M3.1-Flash has the same text architecture, and SGLang's M3.1 path keeps its routed experts in the layout the kernel reads. Copied at vllm-project/vllm@a50259b (Apache-2.0) with only the imports rewritten to this package, so later commits show each change against upstream. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ROCm's Python wheels ship libamdhip64 only under its versioned soname, outside the loader path, so ctypes.CDLL("libamdhip64.so") fails in SGLang's ROCm 10 images and the residency check refused every width. Take the library torch loaded from /proc/self/maps, falling back to the plain soname. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…other kernel test_source_digest requires each launcher to close over SOURCES = kernel_sources(__name__), which digests the module and every plugin and aiter FlyDSL module it imports. vLLM's sources.py hashed two directories instead; drop it for the package's own digest. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

vLLM's launch reads the shared expert as MXFP4 expert 128 of the routed tensors. M3.1's checkpoint stores the shared expert in MXFP8 (SGLang serves it as an FP8 MLP), and requantizing it to MXFP4 costs about 4.6x the output error of FP8 (rel 0.18-0.19 vs 0.04 on layers 3, 30 and 59). In a test build that drift halved DSpark's real acceptance at one request (3.6 -> 2.0 tokens a step) while GSM8K stayed at 97%, since the target still verifies. The launch now takes the shared expert as w13s / s13s / w2s / s2s in MXFP8 (a16w8.py): a16w4's lane-to-K mapping with 32 FP8 bytes a lane a step and the lane's group-32 scale as an f32, dequantized by v_cvt_scalef32_pk_bf16_fp8. The up stage's shared-expert CTAs and the down stage's slot 0 take this path; routed slots keep A16W4. test_minimax_m31_ffn checks every stage against AITER's fused MoE (routed) and torch (shared, on its MXFP8 values) at 1, 3, 8 and 16 rows: 4 passed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… steps The spec's forward runs the embedding, the dense layers and each sparse layer's input norm and attention on SGLang's modules, then one FFN launch from the unreduced o_proj partial to the reduced MoE output; the next layer's input norm applies the residual add, and SGLang's final norm and logits processor close the step. Attention and the KV cache stay SGLang's, so the NVFP4 and FP8 KV formats and HiCache need nothing from the launch. The load hook keeps this rank's shard of each shared expert's MXFP8 tensors (rows of gate / up, columns of down), and the build lays them out for the launch's MXFP8 path, so the shared expert runs on the checkpoint's values. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…3 or M3.1 MiniMax-M3 and MiniMax-M3.1-Flash both load as MiniMaxM3Sparse*, and SGLang registers one spec an architecture, so both now resolve to minimax.spec. It states only what both need (bf16, SGLANG_USE_AITER) in its requirements and hands each checkpoint to its own spec for the rest: the M3.1 one for the NVFP4-expert checkpoint, MiniMax-M3's otherwise. M3's requirements and features are checked the way SGLang's guard checks a spec's. SGLang's per-runner backend recheck reads only the outer requirements, so M3's aiter-backend requirement is checked once, at startup. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Rewriting vLLM's import paths to this package moved them across isort's sections; no code change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…import The shared spec asserted at import that MiniMax-M3 and M3.1 declare the same device limits. sgl-amd-mono#9 gives M3 TP 2 and 8, so merging it would make the plugin's MiniMax spec fail to import and serve neither model. A unit test now holds the invariant, so such a change fails CI instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…endored copy vcommon/ held vLLM's copies of the DeepSeek-V4.1 mono primitives, which were taken from this plugin's common/ and had drifted from it. The launch now imports common.ops, common.plan, common.sync and common.peer: every name it used is there, the same or equivalent (the clock is common.ops.clock, helpers are @flyc.jit as common's are, and the gfx942 branches are not needed on gfx950). The occupancy check loads the HIP runtime through common.peer's loader, made public as hip_runtime(), instead of a second lookup. test_minimax_m31_ffn: 4 passed. The TP=1 launch times match the vendored build (1, 4, 8 and 16 rows: 26.9, 45.9, 65.3 and 109.3 us). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…r wrap a stale tag The vendored launch kept vLLM's own fault scheme, without the host executor vLLM pairs it with, so in SGLang: - a timed-out wait set every rank's error word and the launch ran on; no host code read the word, so the step's tokens went out computed from data that never came. Waits were bounded by a fixed 10 s, not SGLang's deadline, and SGLANG_AMD_MONO_FAULT_INJECT did nothing. - the triple pairs' 16-bit tags repeat after EPOCH_LIMIT launches (126,976, about 2,100 decode steps at 60 layers). vLLM's executor resets the epoch between steps before then; here nothing did, so every later launch set the error word at refuse_wrap and a stale pair could carry a current tag. The waits now run under common.guard, as every other model's do: SGLang's deadline (BuildContext.deadline_s) in the epoch's guard words, the peers' failure flags past the two parity halves, a stop at the null page on failure, and SGLANG_AMD_MONO_FAULT_INJECT. The epoch barrier is common.epoch's. The error line, the step-end handshake and the injected delay are gone, with the runner's final / inject arguments and its unused reset(). SGLang gives a spec no host hook between steps, so the tags cycle on the device instead: 1 .. 65535 per parity, never the 0 a cleared half holds, and the last launch of a parity's cycle zeroes this rank's receive half past its epoch barrier (no peer writes it again before this rank's next launch has pushed). No epoch limit remains. test_minimax_m31_ffn: 6 passed. test_a_timed_out_wait_stops_the_launch fails with a guard that does not stop; test_a_tag_cycle_ends_with_its_receive_half_cleared fails without the clear. TP=1 launch times are unchanged (1, 4, 8, 16 rows: 26.3, 45.7, 65.8, 108.8 us). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

SGLang registers one spec an architecture, so the MiniMax spec re-checks each checkpoint's Requirements as SGLang's guard would and takes the remaining fields from MiniMax-M3's. Two tests now pin that split to SGLang's own structs: every Requirements field is re-checked or agreed on (dtypes, which SGLang checks on the shared spec), and every MonoSpec field is either handed to the checkpoint's spec or agreed on by both models. A field SGLang adds, or a model that diverges on a field the shared spec does not dispatch, fails these instead of going unchecked. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ffn.py's two common.plan imports merge into one and its first-party imports sort by module; sync_acquire went unused with the old epoch barrier. The longer hip_runtime() call in common/peer.py wraps at 88 columns. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…n bounds SGLANG_AMD_MONO_CHECK=1 now runs, in eager steps, SGLang's layer on a copy of each sparse layer's input (ResidualStream.snapshot, on a stream of its own) and logs the launch's MLP output and residual against it, as MiniMax-M3's check mode does. The comparison moves to minimax/check.py, which both models share. NUMERICS replaces the MiniMax-M3 bounds M3.1 borrowed with ones from two check runs on 16 GSM8K prompts (49,476 and 44,004 launches), worst case: residual rel 0.042 / cosine 0.99831, MLP output on rows routed alike rel 0.28 / cosine 0.66, 6.0-6.1% of rows routed to other experts. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Like every other model's: the DSpark verify server (2 requests of 8 rows at the largest width) and the decode server (16 rows) show their tested batches in the startup line and pass chat GSM8K past the largest width (0.975 / 0.985 and 0.975; the standard path scores 0.975; 8-shot completion prompts score under 0.1 on either path, so the gate is chat). Check mode holds every launch to NUMERICS, and a rank stalled by SGLANG_AMD_MONO_FAULT_INJECT stops the server with faults at the null page only and no tokens sent; that test turns crash dumps off, as #55 did for GLM's. 6 passed (MI350X, TP4, sgl-project/sglang#45936). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The model row and design note, its shared-expert copy, the SGLang PRs it needs, a launch command, the minimax / minimax_m31 packages, and its test command; the MiniMax-M3 one leaves M3.1's tests out. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

README.md's design bullet took main's GLM-5.3-Flash tail launch (#23) and this branch's MiniMax-M3.1-Flash sentence; nothing else conflicted. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

## Motivation **Summary.** After #24, the score stage of the GLM-5.3 indexer read the keys at about 0.7 TB/s, and at 5 to 8 rows past 128k keys it took about half of an indexer launch. This PR scores each group of 16 keys on one wave over all 32 heads, with the reads of the next groups in flight. The scores do not change: in one launch, they are bitwise equal to the scores of #24 on the same inputs. On 8x MI355X at TP8, in InferenceX agentic coding with MTP, interactivity p50 is 1.43x, 1.46x and 1.23x that of the normal path with MTP at concurrency 1, 2 and 4 (#24: 1.15x, 1.17x and 1.11x). **On main after #24.** This PR is two commits on main (`240d587`, which has #24 and #56): `e405d6b` is `e0d7655` rebased with the same changed lines (and `e0d7655` is `9733ff4` rebased), and `3419a0e` is the fix from the review. The results below were measured at `9733ff4` unless marked. **Cause.** In the paged score stage of #24: - A CTA scored a tile of 64 keys with its 8 waves: 4 groups of 16 keys, each for the two halves of the 32 heads. The two waves of a group read the same keys, so each key was read two times. - The two halves were added through LDS, with two block barriers a tile. - The block-table read and the key reads of a tile started only after the tile before it ended. Thus each tile waited for two memory round trips: about 2.1 us a tile, 0.7 TB/s at 256k keys. - Each key also compared its cache slot with the slots of all S rows of the launch (the new-key lookup of #15). This cost grows with the width (S = 2, 4 or 8). ## Modifications - `glm53/gfx950/kernel.py`, paged score stage only: - A wave scores a group of 16 keys over all 32 heads: two MFMA chains on one read of the keys. The head sums stay in the lanes, so the stage has no LDS reduction and no barrier between groups. - A wave walks every 8th group of its CTA's tiles; the tiles of a CTA do not change. The key reads of the next 2 groups are in flight, and their block-table reads go one step further ahead. The first block-table reads of a row overlap the staging of its query. - The query fragments and the head weights of the row stay in registers. The values of a row (bound, request, counts) are wave-uniform. - The new-key lookup of #15 runs only for each wave's group that holds any of the last S keys of a row: a launch writes only those keys (a verify row's draft chain, a decode row's own key). A wave's groups are `SCORE_STRIDE` (32768) keys apart, so a wave has at most one such group; the walk asserts `S + 15 < SCORE_STRIDE` (review fix). - Each score is computed in the same order as before: the same MFMA chain, the same sums in the same order. - `test/e2e/test_glm53_server.py` (review fix): the parallel check case of #24 passes 98304 keys (rows of about 108k keys at widths 2 and 4), where each wave's walk refills its buffers and looks blocks ahead. Before, the longest case (about 36k keys) gave each wave at most 2 groups, so the refills never ran. - No new check. With a wrong refill offset (a deliberate error, not in this PR: a refilled buffer reads the keys 16 positions on), the parallel check case fails on the existing stream-cosine bound: 22 violations, lowest cosine 0.914 (bound 0.95). `layer.topk.rel` rises to 0.66 (0.052 without the error). So once a case reaches the walk's refills, the existing bounds catch an error there. The bitwise check needs a debug build (not in this PR). ## Accuracy Tests **Setup.** 8x MI355X (gfx950) at TP8; zai-org/GLM-5.3 @ `ba32f8c1` (FP8); the image and the commits of #24; the server flags of `test/e2e/glm53_server.py`; a new server for each test. **Test 1: bitwise scores in one launch.** A debug build (not in this PR) also runs the per-tile score code of #24 in the same launch, on the same inputs, and writes those scores to a second scratch buffer. A script compares the two buffers. | Steps | Rows | Keys | Keys bitwise equal | |---|---|---|---| | Decode: 1 to 8 rows, 8k to 256k keys | 256 | 29,961,984 | 29,961,984 | | MTP verify: 1 and 2 requests, 8k to 256k keys | 192 | 29,631,744 | 29,631,744 | The exact top-k check of #24 also found 0 wrong rows in both runs. We did not compare two servers: on this image, two prefills of the same prompt give different bits. After a prefill of 131072 tokens, the scores of the first decode step correlate at 0.9988 between two passes (one prefill, cached) and at 0.9457 between two prefills in the same server, with the same build. Between a server with #24 and a server with this PR, they correlate at 0.9460. **Test 2: E2E tests at TP8 (review fix).** We ran `test/e2e/test_glm53_server.py` as written at `3419a0e` on main `240d587`. All 7 tests passed in one run (1214 s): - `TestMonokernelDecode`: average KL 0.018 to 0.022 for each batch (bound 0.06); GSM8K (sgl-eval, thinking on) 0.99. - `TestMonokernelMtp`: GSM8K 0.99, accept length 3.28. - `TestMonokernelCheck`: no violations. `layer.topk_exact` (#24) and `layer.topk_own` are 0 in all 954 indexer checks. - Decode steps: 1425 checks, stream cosine 0.9942 or more, `layer.topk.rel` 0.014 or less. - Verify steps: 1350 checks, stream cosine 0.9976 or more, `layer.topk.rel` 0.017 or less. - Past the parallel top-k, with rows past 98304 keys: 1200 checks, stream cosine 0.9968 or more, `layer.topk.rel` 0.052 or less. - `TestMonokernelFault`: passed. Since #55, the test sets `HSA_DISABLE_COREDUMP_ON_EXCEPTION=1` and turns off SGLang's crash dumps itself. Two earlier runs also passed all 7 tests: at `9733ff4`, before the review fix, and at `9b48912`, the fix on #24's branch before #56 merged. ## Speed Tests and Profiling **Setup.** The setup of #24: one node with 8x MI355X, TP8, image `lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261019`, SGLang `e45cbce599`, AITER `2bb549102b`, the InferenceX GLM recipe flags, and a new server for each point on an idle node. The arms differ only in `--monokernel-backend amd` and the checkout. **Kernel timeline: #24 to this PR.** The debug build of #24, random prompts, 1 to 8 rows decoding together, the last launch of each of 6 passes on each of the 8 ranks (medians). The estimated step is 18 indexer launches and 57 launches without the indexer. | Keys per row | Rows | Score stage (us) | Indexer launch (us) | Estimated step (ms) | |---|---|---|---|---| | 8k | 1 | 5 -> 5 (0.9x) | 97 -> 96 | 4.84 -> 4.80 (1.01x) | | 8k | 2 | 12 -> 9 (1.3x) | 108 -> 106 | 4.98 -> 4.96 (1.00x) | | 8k | 4 | 18 -> 14 (1.3x) | 136 -> 134 | 6.61 -> 6.59 (1.00x) | | 8k | 8 | 29 -> 24 (1.2x) | 206 -> 202 | 10.65 -> 10.61 (1.00x) | | 32k | 1 | 11 -> 5 (2.0x) | 119 -> 114 | 5.20 -> 5.12 (1.01x) | | 32k | 4 | 45 -> 29 (1.5x) | 182 -> 172 | 7.45 -> 7.26 (1.03x) | | 32k | 8 | 89 -> 62 (1.4x) | 285 -> 261 | 12.10 -> 11.67 (1.04x) | | 64k | 1 | 15 -> 6 (2.5x) | 129 -> 133 | 5.33 -> 5.49 (0.97x) | | 64k | 4 | 63 -> 35 (1.8x) | 206 -> 182 | 7.87 -> 7.47 (**1.05x**) | | 64k | 8 | 126 -> 74 (1.7x) | 334 -> 281 | 12.95 -> 11.95 (**1.08x**) | | 128k | 1 | 25 -> 8 (3.1x) | 144 -> 126 | 5.60 -> 5.32 (**1.05x**) | | 128k | 2 | 50 -> 19 (2.7x) | 174 -> 144 | 6.19 -> 5.59 (**1.11x**) | | 128k | 4 | 102 -> 41 (2.5x) | 251 -> 191 | 8.71 -> 7.58 (**1.15x**) | | 128k | 8 | 203 -> 87 (2.3x) | 417 -> 306 | 14.56 -> 12.44 (**1.17x**) | | 256k | 1 | 44 -> 12 (3.7x) | 175 -> 148 | 6.20 -> 5.73 (**1.08x**) | | 256k | 2 | 88 -> 27 (3.2x) | 224 -> 166 | 7.05 -> 6.05 (**1.17x**) | | 256k | 4 | 180 -> 60 (3.0x) | 343 -> 225 | 10.37 -> 8.33 (**1.25x**) | | 256k | 5 | 227 -> 78 (2.9x) | 453 -> 305 | 14.94 -> 12.30 (**1.21x**) | The launches without the indexer run the same code in both builds; between the two runs they differ by -1.8% to +3.4%. The differences at 8k keys and at 64k keys with 1 row are in this noise. The score stage's time is its window, from the first start to the last end of its tasks. With this PR, a tile's start and end stamps can come from two waves, so the time of one task is not meaningful, but the window is. **AgentX with MTP.** InferenceX agentic coding (`agentx-harness` 1.0.11, corpus `23f152f6`), 3600 s a point, the same simulated accept length (2.99) in all arms. | Concurrency | Interactivity p50: normal path with MTP / #15 / #24 / this PR | p90 | Total throughput (tok/s/GPU) | Requests done | Input p50 | |---|---|---|---|---|---| | 1 | 204.5 / 0.95x / 1.15x / **292.1 (1.43x)** | 178.5 / 0.94x / 1.14x / **216.1 (1.21x)** | 1921 / 0.98x / 1.11x / **2386 (1.24x)** | 234 / 231 / 246 / 261 | 225k / 221k / 285k / 300k | | 2 | 199.7 / 0.95x / 1.17x / **292.0 (1.46x)** | 149.6 / 0.91x / 1.03x / **184.9 (1.24x)** | 2044 / 0.99x / 1.15x / **2418 (1.18x)** | 377 / 374 / 409 / 427 | 97k / 96k / 111k / 116k | | 4 | 180.7 / 0.86x / 1.11x / **222.6 (1.23x)** | 128.6 / 0.83x / 0.96x / **146.3 (1.14x)** | 2961 / 0.97x / 1.06x / **3657 (1.24x)** | 501 / 493 / 536 / 582 | 91k / 91k / 89k / 112k | Time of a verify step for each generated token (ms), from the server log intervals without prefill, concurrency 1 and 2 (normal path with MTP / #15 / #24 / this PR): | Requests in the step (verify rows) | 0-64k | 64-128k | 128-192k | 192-320k | 320k or more | |---|---|---|---|---|---| | 1 (4) | 4.61 / 3.21 / 3.15 / **3.06** | 4.66 / 3.50 / 3.29 / **3.12** | 4.74 / 3.93 / 3.49 / **3.21** | 4.86 / 4.92 / 4.00 / **3.33** | 4.95 / 5.65 / 4.39 / **3.48** | | 2 (8) | 5.51 / 4.92 / 4.79 / **4.61** | 5.57 / 5.39 / 5.02 / **4.70** | 5.61 / 5.73 / 5.31 / **4.85** | 5.82 / 7.51 / 6.20 / **5.04** | too few steps | With 2 requests past 192k keys, #24 was still slower than the normal path; this PR is 1.15x faster there. **Fixed length without MTP.** Prompts of 131072 and 262144 tokens, 1024 new tokens, 6 prompts for each concurrent request, the radix cache flushed before each concurrency. Decode step without prefill (ms, median of the server log intervals): | Keys | Batch | Normal path | #24 | This PR | |---|---|---|---|---| | 128k | 1 | 11.04 | 6.47 | **6.25** | | 128k | 4 | 12.94 | 9.28 | **8.16** | | 128k | 8 | 15.52 | 14.78 | **12.52** | | 256k | 1 | 11.36 | 7.01 | **6.49** | | 256k | 4 | 13.55 | 10.93 | **8.75** | | 256k | 5 | 15.92 | 15.40 | **12.61** | At 256k keys, the #24 column is the same selection code on `e952b22` (without #15); #24 itself was measured at 128k. ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change. - [x] The SGLang commit or PR this was tested against is named above. - [x] Accuracy and speed results are attached when kernel or forward code changes. - [ ] A new model adds one line to `.github/CODEOWNERS` that names its owner. ## Review and Merge Process 1. The `lint` and `Check PR title` checks must pass. 2. Ask a code owner from `.github/CODEOWNERS` to review. 3. Merge with "Squash and merge": the PR title becomes the commit title on `main`.

…le, read B operands from LDS (#59) ## 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, 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 #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](https://claude.com/claude-code) --------- Co-authored-by: Kevin Mi <kevin.mi@radixark.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

## Motivation Without CUDA graphs, every launch of the DeepSeek-V4.1, GLM-5.3, GLM-5.3-Flash and MiniMax-M3 runners goes through FlyDSL's JIT dispatch (`JitFunction.__call__`): it binds the signature, checks the globals and builds the full cache key at each launch. For GLM-5.3-Flash's launchers (50 integer arguments), that costs 182 µs of host time a launch, so an eager decode step is host-bound. This PR reuses AITER's compiled-dispatch helper: it keeps the first real launch's `CompiledFunction` and calls it after that. A replayed CUDA graph does not call the launchers, so decode with CUDA graphs (the default) does not change. ## Modifications - These launches use `_run_compiled` from `aiter.ops.flydsl.kernels.tensor_shim`, with the runtime addresses, scalars and stream passed positionally: - DeepSeek-V4.1: K1, K2, FFN and K2i. - GLM-5.3: the layer. - GLM-5.3-Flash: the layer and the head. - MiniMax-M3: the launch. MiniMax-M3.1 already caches its compiled FFN callable. - `compile_only` in the four runners uses AITER's `_preload_compiled`. AITER handles the no-dispatch preload and restores the environment. - The GLM runners import both helpers where they use them. AITER probes the GPU when it is imported, and `test/unit` imports these modules on CPU. - `test/kernels/test_compiled_dispatch.py` checks the dispatch that the runners rely on: - a preload never runs the kernel; - the first real call runs it once; - later calls skip JIT dispatch, also inside a graph. - AITER requirement: `_preload_compiled`, from ROCm/aiter#5346 (`354e278e5`) onwards. ROCm/aiter#6946 is not needed, because this repository never calls `_run_compiled` while `COMPILE_ONLY` is set. ## Accuracy Tests The kernels and their arguments do not change, and the outputs are bitwise equal to main's: - GLM-5.3-Flash: the eager step of layers 3-43 at TP4, 4 and 8 rows, on every rank. - DeepSeek-V4.1 K1 through `DSV41MonoRunner.k1`: at TP2; widths 1, 6, 24 and 48; fold, settled and index. The outputs match after the first and third eager call and after a graph replay, in 24/24 cases. E2E on MI355X: - `test_glm53_flash_server.py` (TP4): 5/5. GSM8K past the largest width scores 0.995. - `test_glm53_server.py` (FP8, TP8): 8/8. GSM8K scores 0.99 for decode and 0.99 with MTP, with an accept length of 3.3-3.4. - `test_minimax_m3_server.py` (TP4, with sgl-project/sglang#41998): 7/9. GSM8K scores 0.925 with EAGLE3 and 0.885 for decode. The 2 failing cases also fail on main: the test expects the spec name `minimax_m3_hip`, but the spec has been named `minimax_m3_family_hip` since #56. After the GLM import change, CPU-only `test/unit` gives the same result as on main, and GLM-5.3-Flash stays bitwise equal to main at TP4. ## Speed Tests and Profiling MI355X at TP4, with a fresh server for each arm and two rounds. The load is random 2048 tokens in and 256 out. The table gives median TPOT in ms, main -> this PR: | Model | CUDA graphs, c1 / c4 | `--disable-cuda-graph`, c1 / c4 | | --- | --- | --- | | GLM-5.3-Flash | 2.520 -> 2.523 / 3.032 -> 3.029 | 11.28 -> 2.72 / 11.75 -> 3.19 | | GLM-5.3 (FP8) | 6.812 -> 6.812 / 9.351 -> 9.342 | 20.06 -> 7.05 / 20.43 -> 9.39 | | MiniMax-M3 + EAGLE3 | 1.430 -> 1.426 / 4.118 -> 4.174 | 6.29 -> 3.31 / 9.58 -> 6.38 | - **With CUDA graphs**, TPOT does not change (within ±1.4% at every covered width). The gain applies only to steps that run without a graph. - **Host time** of a GLM-5.3-Flash launch goes from 182 µs to 14.7 µs. An eager step of layers 3-43 goes from 7.50 ms to 2.54 ms, which equals the step's graph replay time. Setup: image `rocm/sgl-dev:v0.5.22-rocm10-mi35x-20261026`, with SGLang `5ef99013ac`, AITER `285579b62` and FlyDSL 0.4.0. ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change. - The dispatch: `test_compiled_dispatch.py`. - DeepSeek-V4.1 FFN and K2: `test_deepseek_v41_kernels.py`. K2i: `test_deepseek_v41_k2i.py`. - GLM-5.3-Flash: `test_glm53_flash_kernels.py`. - MiniMax-M3 `compile_only`: `test_minimax_m3_compile.py`. - Only the e2e tests reach the GLM-5.3 and MiniMax-M3 launches. - No committed test reaches `DSV41MonoRunner.k1`, because `check_k1.py` calls the K1 launcher directly. K1 was checked bitwise against main, as described above. - [x] The SGLang commit or PR this was tested against is named above. - [x] Accuracy and speed results are attached. - [ ] A new model adds one line to `.github/CODEOWNERS` that names its owner. (No new model.) --------- Co-authored-by: Jin Pan <jin.pan@amd.com>

Kevin Mi and others added 7 commits October 23, 2026 23:12

[Chore] minimax_m31: Sort the vendored imports as the isort hook does

449046c

Rewriting vLLM's import paths to this package moved them across isort's sections; no code change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

kevin-mii requested review from HaiShaw, alexnails and zijiecode as code owners October 24, 2026 04:00

Kevin Mi and others added 10 commits October 25, 2026 01:25

[Chore] minimax_m31: Apply the lint job's isort and ruff-format fixes

0ece66f

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Merge main into kevin/minimax-m31-ffn

89aaccb

README.md's design bullet took main's GLM-5.3-Flash tail launch (#23) and this branch's MiniMax-M3.1-Flash sentence; nothing else conflicted. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

HaiShaw merged commit de763e7 into main Oct 25, 2026
2 checks passed

HaiShaw deleted the kevin/minimax-m31-ffn branch October 25, 2026 03:31

jhinpan mentioned this pull request Oct 25, 2026

[Perf] [MI355X] glm53: Score the paged indexer keys per wave #42

Merged

4 of 5 tasks

kevin-mii mentioned this pull request Oct 25, 2026

[Perf] minimax_m31: Run o_proj in the launch, route on the norm's scale, read B operands from LDS #59

Merged

This was referenced Oct 26, 2026

[Fix] glm53_flash: Give a step's first launch half the watchdog as its wait deadline (preview) MarloweAI/sgl-amd-mono-upstream#10

Closed

[Fix] glm53_flash: Give a step's first launch half the watchdog as its wait deadline #61

Merged

jhinpan mentioned this pull request Oct 27, 2026

[Perf] mono: Reuse AITER compiled FlyDSL dispatch #66

Merged

4 of 5 tasks

eppaneamd mentioned this pull request Oct 30, 2026

[Feat] [MI355X] qwen38_flash_next: Serve TP2 decode of the MXFP4 model #93

Draft

3 of 5 tasks

kevin-mii mentioned this pull request Nov 2, 2026

[Feat] [AMD] minimax: Support MiniMax-M3.1-Flash on MI350X / MI355X (gfx950) sgl-project/sglang#46693

Open

Navigation Menu

Repository navigation

[Feat] minimax_m31: Run MiniMax-M3.1-Flash's decode FFN as one monokernel launch a layer - #56

[Feat] minimax_m31: Run MiniMax-M3.1-Flash's decode FFN as one monokernel launch a layer#56
HaiShaw merged 17 commits into
mainfrom
kevin/minimax-m31-ffn

kevin-mii commented Oct 24, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

kevin-mii commented Oct 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy tests

Speed tests

Not covered

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kevin-mii commented Oct 24, 2026 •

edited

Loading