Skip to content

[Feat] [AMD] minimax: Support MiniMax-M3.1-Flash on gfx950 with DSpark and fused kernels - #45936

Closed
kevin-mii wants to merge 3 commits into
sgl-project:minimax-m3.1-flashfrom
kevin-mii:kevin-mii/minimax-m3.1-flash-amd
Closed

kevin-mii wants to merge 3 commits into
sgl-project:minimax-m3.1-flashfrom
kevin-mii:kevin-mii/minimax-m3.1-flash-amd

Conversation

@kevin-mii

@kevin-mii kevin-mii commented Oct 23, 2026 •

Copy link
Copy Markdown
Collaborator

Motivation

Run MiniMax-M3.1-Flash on AMD Instinct MI350X / MI355X (gfx950) with DSpark speculative decoding, at TP 4 per server. The same paths serve AMD's Quark MXFP4 export of the model.

Dependency stack: this PR is stacked on #45925 (MiniMax-M3.1-Flash support) and targets its branch minimax-m3.1-flash, so the diff is AMD-only. Review #45925 first. Once it merges, retarget this PR to main.

Modifications

  • Model override (model_overrides/minimax_m3_hip.py): gfx950 defaults for the NVFP4 checkpoint and the Quark export.
    • Attention runs on the Triton backend, and the DSpark draft on Triton.
    • The MoE runner is AITER, with page size 128.
    • The KV format defaults to NVFP4 under DSpark.
    • The quick all-reduce level and AITER's tuned tables are set here too.
    • Guards: M3.1 NVFP4 on ROCm refuses TP above 4 (TP8 with attention DP produced runaway whitespace in GSM8K answers), and attention TP above the KV head count, which returned garbage on ROCm.
  • Weights
    • Routed experts are requantized from NVFP4 to MXFP4 at load and run on AITER's FlyDSL fused MoE, with tables tuned for the M3.1 TP4 shapes.
    • The MXFP8 attention, dense and shared-expert projections are converted to FP8 at load and run on AITER GEMMs with tuned block-scale and rowwise tables. Up to 128 tokens they use a per-channel copy; above that, block scaling.
    • A small-M FP8 GEMM serves the smallest batches.
  • KV cache and attention
    • Under DSpark, the cache is the packed NVFP4 KV and indexer-K pool (MSA-v1 layout, as on Blackwell), with Triton quantize/dequantize and a fused store.
    • Sparse decode runs on a gfx950 Gluon kernel. Prefill runs on a Gluon kernel from 2048 tokens and on AITER's Gluon paged attention below that. Indexer scoring is Triton.
    • SGLANG_MINIMAX_M3_AMD_KV_FORMAT=fp8 selects an FP8 cache instead, which is the default for the Quark export.
    • HiCache copies of the packed pages run on HIP.
  • Fusions
    • QK-norm, RoPE and the KV-cache write run as one kernel.
    • RMSNorm and SwiGLU write the FP8 input of the next GEMM directly.
    • The shared-expert output folds into the MoE reduction.
    • Up to 16 tokens, all-reduce + residual add + RMSNorm + FP8 quant run as one kernel on AITER's custom all-reduce, and so do router GEMV + top-k + expert sort.
  • All-reduce: AITER custom all-reduce handles messages up to 2.4 MiB. Above that, at TP 4, a JIT-built quick all-reduce takes over: INT8 below 64 MiB, INT4 above, bitwise equal to sgl-kernel's. SGLANG_ENABLE_MINIMAX_M3_AMD_QUICK_REDUCE_INT4=0 keeps INT8. Other world sizes, levels and dtypes fall back to sgl-kernel.
  • DSpark
    • Draft attention runs on the Triton backend.
    • The Markov step bias (w1 gather + w2 GEMV) is one kernel.
    • The draft's context K/V for all layers is written through one block-FP8 GEMM plus one norm/RoPE/store kernel.
    • The verify argmax runs without gathering the logits across GPUs.
    • The draft step is queued before the host waits for the previous step's sequence lengths.
  • AITER integration: two runtime patches, each applied only if AITER's function signature matches and otherwise left unpatched with a log line:
    • aiter_moe_addend folds an addend into AITER's stage-2 reduction.
    • moe_sorting_small hands AITER a presorted routing.
    • The tuned tables ship under layers/quantization/aiter_tuned_configs and are merged after AITER's defaults through its AITER_CONFIG_* variables.
  • ROCm code lives in _hip modules or behind is_hip / gfx950 checks, so CUDA paths are unchanged.
  • Indexer CP for DSpark (second commit, opt-in SGLANG_MINIMAX_M3_INDEXER_CP=1): the decode indexer's context partitioning now covers DSpark verify on the NVFP4 index cache.
    • Each rank scores a quarter of the index-K blocks for all four index heads in one tile; the merge picks exactly what the per-rank selector picks.
    • Candidates go straight into the head owner's IPC receive buffer instead of an all-gather. It engages from 12 verify requests.
    • Per layer at 64 requests x 110K context: 165.7 -> 84.0 us. AgentX c64 (3600 s, HiCache): 86,995 -> 90,140 tok/s/GPU.

Accuracy Tests

On MI350X with ROCm 7.2.0, torch 2.9.1+rocm7.2, Triton 3.7.0 and AITER acf8fdf93 (image rocm/sgl-dev:v0.5.20-rocm720-mi35x-20260926), with sglang-kernel 0.4.9 built for gfx950:

  • Publication revision af6d50a85e
    • GPU unit suite: 23 gfx950 test files, 321 passed, 6 skipped, 828 subtests.
    • Four-GPU all-reduce tests: 4 passed.
    • CPU model tests: 33 passed, 24 subtests.
    • pre-commit (every hook in .pre-commit-config.yaml) passes on all changed files.
  • Four-GPU TP4 serving smoke, using the cookbook's AMD command:
GSM8K (sgl-eval, max 8192 tokens) NVFP4 KV (default) FP8 KV
1,319 questions, greedy 96.13 95.75
600 questions, temperature 1.0, top-p 0.95, seed 0 95.83
  • Image checks: image-understanding checks pass, including rendered text in PNG and JPEG, shape counting, two-image order, small text in a 2048×1536 image, a 37×53 image, and 16 concurrent image requests.
  • Real DSpark acceptance during GSM8K: about 3.7 tokens per step with NVFP4 KV (mean of decode-log samples).

Serving used MiniMaxAI/MiniMax-M3.1-Flash revision 9e7ebacac350497ad4a1272a2e817639bf11561c with SGLANG_USE_AITER=1 and these flags:

--tp 4 --reasoning-parser minimax-m3 --tool-call-parser minimax-m3 --speculative-algorithm DSPARK
--cuda-graph-backend-prefill breakable --fp8-gemm-backend aiter --chunked-prefill-size 32768
--max-running-requests 128 --cuda-graph-max-bs-decode 128 --mem-fraction-static 0.8

Greedy GSM8K on this stack varies between runs by about one point. For example, 96.29 / 96.59 at this revision with a pre-release draft checkpoint, and 95.9–96.9 on earlier revisions. AITER's stage-2 MoE kernels accumulate with atomics, so outputs are not bitwise deterministic run to run.

Speed Tests and Profiling

Four MI350X GPUs, TP 4, NVFP4 KV, DSpark acceptance simulated at 3 tokens per step (SGLANG_SIMULATE_ACC_LEN=3, match-expected, real-draft-token). Real-acceptance throughput was not measured for this PR.

Decode throughput is output tokens/s per server, median over 2 crossed rounds × 2 GPU sets with 3 repeats each. Every column uses 150K-token distinct contexts, except 64 requests, which uses 64K.

Concurrent requests 1 2 4 8 16 32 64
Output tok/s 376 640 1,016 1,629 2,515 3,978 5,765

Prefill is the latency of one extend on a cached prefix, median of 9 repeats over 2 crossed rounds × 2 GPU sets:

Cached prefix + new tokens 0 + 32K 100K + 2K 100K + 8K 200K + 8K 100K + 32K
Latency (ms) 551 91 191 220 630

These numbers were measured on a branch with the same AMD code. It differs from this PR only by #45925's _draft_context() wrapper around the prefill KV injection, which is a no-op without attention DP. They describe this stack's throughput, not gains attributable to individual changes.

Remaining validation and follow-ups

Re-validated at c9b6408 (rebased onto minimax-m3.1-flash 65db545, plus indexer CP), MI350X, ROCm 10:

  • GSM8K 400 with DSpark 7, SGLANG_MINIMAX_M3_INDEX_TOPK_FREQ=1, indexer CP and HiCache: 0.9675. Accept length 4.765 over decode batches with 12+ running requests. Long shared-prefix prompts reuse the cached prefix.
  • test_minimax_indexer_cp (134 subtests), test_minimax_indexer_cp_ipc (4 GPUs), test_minimax_verify_topk_idx (76), test_minimax_m31_hicache, test_minimax_m31_model (24).
  • These runs set AITER_CONFIG_GEMM_BF16 to a copy of the bf16 table without its 27 hipBLASLt rows: their solution indices belong to one hipBLASLt build, and on ROCm 10 they fail in verify graph capture (HIPBLAS_STATUS_INVALID_VALUE).

Not executed at this revision:

  • MI355X. Earlier revisions of this AMD work were served on MI355X with both the NVFP4 and Quark MXFP4 checkpoints.
  • The Quark MXFP4 checkpoint.
  • The HiCache high-throughput recipe.
  • An AgentX ladder.
  • ROCm 10 images.
  • The sglang-kernel 0.4.10 ROCm wheel pinned in pyproject.toml; validation used a gfx950 build of 0.4.9.
  • CUDA. Its paths are untouched by this diff.
  • PD disaggregation and attention DP.

The FP8 GEMM path launches AITER's split-K block-scale assembly kernel into an output that RMSNorm has already zeroed. This needs a local copy of the launcher, because AITER always memsets the output first. ROCm/aiter#6786 (draft) adds an out_zeroed argument. Remove the local launcher once the pinned AITER includes it.

Checklist

  • pre-commit passes on all changed files.
  • Registered regression tests cover the gfx950 kernels, numerics, padding and graph behavior (test/registered/amd, test/registered/kernels, test/registered/unit, test/registered/spec).
  • Accuracy and throughput results identify workload, hardware, checkpoint revision and tested revisions.
  • Comments and docstrings reviewed during cleanup.
  • Retarget to main after [Feat] minimax: Support MiniMax-M3.1-Flash #45925 merges.

🤖 Generated with Claude Code


CI States

Latest PR Test (Base): ⏳ Run #40273493287
Latest PR Test (Extra): ❌ Run #40273493187
Latest PR Test (AMD ROCm 10): ⏳ Run #40273493120

@kevin-mii
kevin-mii requested a review from yuan-luo as a code owner October 23, 2026 01:51
@github-actions github-actions Bot added quant LLM Quantization amd hicache Hierarchical Caching for SGLang blackwell SM100/SM120 jit-kernel labels Oct 23, 2026
@github-actions

This comment has been minimized.

kevin-mii pushed a commit to sgl-project/sgl-amd-mono that referenced this pull request Oct 25, 2026
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>
HaiShaw pushed a commit to sgl-project/sgl-amd-mono that referenced this pull request Oct 25, 2026
…rnel launch a layer (#56)

## 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](https://claude.com/claude-code)

---------

Co-authored-by: Kevin Mi <kevin.mi@radixark.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
kevin-mii and others added 2 commits October 25, 2026 09:22
Serves MiniMax-M3.1-Flash (minimax_m3_nvfp4) and AMD's Quark MXFP4 export of it
with DSpark on gfx950, ROCm 7.2 / Triton 3.7, at TP 4 per server.

- Weights: the routed experts are requantized to MXFP4 at load and run on AITER's
  FlyDSL fused MoE with tables tuned for the M3.1 TP4 shapes. The MXFP8 attention,
  dense and shared-expert projections are converted to FP8 at load (a per-channel
  copy for up to 128 tokens, block-scaled above) and run on AITER GEMMs with tuned
  tables, plus a small-M FP8 GEMM for the smallest batches.
- KV cache: the packed NVFP4 KV and indexer-K pool under DSPARK (fp8 with
  SGLANG_MINIMAX_M3_AMD_KV_FORMAT=fp8, and for the Quark export). Sparse decode
  and prefill run on Gluon kernels, indexer scoring on Triton.
- Fusions: QK-norm, RoPE and the KV-cache write in one kernel; RMSNorm and SwiGLU
  write the FP8 input of the next GEMM; the shared-expert output folds into the
  MoE reduction; up to 16 tokens, all-reduce + residual add + RMSNorm + FP8 quant
  and router GEMV + top-k + expert sort each run as one kernel.
- All-reduce: AITER custom all-reduce up to 2.4 MiB, the JIT INT8/INT4 quick
  all-reduce above it at TP 4.
- DSpark: Triton draft attention, the markov step bias in one kernel, the draft's
  context K/V written through one block-FP8 GEMM, the verify argmax without the
  logits TP gather, and the draft queued before the host waits for seq lens.
- Guards: M3.1 NVFP4 on ROCm refuses TP above 4 and attention TP above the KV
  head count.

Validated on 4x MI350X with MiniMax-M3.1-Flash and the cookbook command: GSM8K
(1319, greedy) 96.29 with the NVFP4 KV cache and 96.44 with fp8, 96.33 sampled
(temperature 1.0, top-p 0.95, 600 questions); image understanding checks pass.

Co-authored-by: Zijie Chen <300606707+zijiecode@users.noreply.github.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… cache

Indexer CP (SGLANG_MINIMAX_M3_INDEXER_CP=1, gfx950 TP4) covered bf16/fp8 index K and
chain EAGLE drafts; M3.1's DSpark verify on the packed NVFP4 index cache fell back to
the native selector.

- The fused verify indexer kernel takes WORLD/RANK: each rank scores blocks
  local * 4 + rank for all four index heads in one 32-row tile and nominates its best
  16 per split; the merge selects per owned head exactly as the native selector.
- DSpark qualifies under the static verify schedule.
- Candidates go straight into the head owner's uncached IPC receive buffer, each word
  tagged with a per-row call counter (SGLANG_MINIMAX_M3_INDEXER_CP_IPC, default on),
  replacing the per-layer candidate all-gather.
- CP engages from 12 verify requests with the IPC exchange (48 with the gathers) and
  falls back past the receive buffers' 4096 rows.
- The gather path's merge ignores negative-id keys: AITER's custom all-gather
  zero-fills its output during CUDA-graph warmup.

Verify indexer per layer, 64 requests x 110K context on MI350X: 165.7 -> 84.0 us.
AgentX c64 (3600 s, HiCache): 86,995 -> 90,140 tok/s/GPU. GSM8K and acceptance unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@kevin-mii
kevin-mii force-pushed the kevin-mii/minimax-m3.1-flash-amd branch from af6d50a to c9b6408 Compare October 25, 2026 09:29
@kevin-mii kevin-mii changed the title [AMD] Support MiniMax-M3.1-Flash on gfx950 with DSpark and fused kernels [Feat] [AMD] minimax: Support MiniMax-M3.1-Flash on gfx950 with DSpark and fused kernels Oct 25, 2026
@kevin-mii

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Oct 25, 2026
- logits_processor_hip reads the TP group from get_parallel(), not the deprecated
  get_tp_group() that test_runtime_context rejects.
- FutureMap declares seq_lens_pull at class level, like _relay, so a map built without
  __init__ (the post-fold relay test) still publishes.
- The indexer-CP gate test expects DSpark enabled under the static verify schedule
  only, which the DSpark CP commit allows.

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

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

amd blackwell SM100/SM120 hicache Hierarchical Caching for SGLang jit-kernel quant LLM Quantization run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants