[Fangzhou-Ai]

Purpose

The NVFP4 sparse decode kernel of MiniMax-M3.1 (ops/nvfp4_sparse_decode, ROCm gfx950, XCD-aware grid) picks its number of splits with decode_splits, a cost model that spreads the work over all CUs of the device (256 on MI355X). The XCD-aware grid does not do that: it places a request's rows and all of its splits on one XCD (32 CUs). The model therefore mispredicts: at 5 concurrent requests (DSpark 7: 8 rows per request) it picks 6 splits (the s6 variant, 15.4-15.7 us per layer) where 4 splits take 11.9 us, and at 1-4 requests it picks 8 splits (12.4-13.3 us) where 4 splits take 11.2-12.4 us. s6 appears only at exactly 5 requests, which is a common decode batch in the AgentX c16 workload (about 15% of request decode time).

This PR adds _xcd_splits, used only when the heuristic would otherwise run (explicit splits= and the splits == 1 large-grid path are unchanged), for the shape it was measured on (top-k 16, 8 rows per request = DSpark 7); every other shape keeps decode_splits:

• up to 3 requests per XCD, i.e. up to 8 requests on 4 XCDs and 17-24 requests: 4 splits;
• 2 requests per XCD (9-16 requests): 2 splits.

Not covered by any open PR that I could find (the open MiniMax PRs that touch sparse decode code are the NVIDIA MSA ones and #56451, which changes DSpark verification, not the split choice).

Measurements (MI355X, gfx950; 4 ranks, amd/MiniMax-M3.1-Flash-Quark-MXFP4, TP4, DSpark 7, NVFP4 KV and indexer)

Kernel level (sweep of splits x XCD mode, N = 1..64 requests, 30K / 150K / 400K contexts, per-layer us): the kernel time does not depend on the context length; one split costs ~21 us (16 blocks serial, about 0.8 us per block plus ~8 us fixed). Best split counts: 4 splits up to 8 requests (11.2-12.4 us), 2 splits for 9-16 requests (15.3-16 us), 4 splits for 17-24 requests (21 us). The old cost model gave 8 splits at 1-4 requests (12.4-13.3 us) and 6 at 5 requests (15.7 us). Disabling the XCD-aware mode would be 0.3-0.6 us better only at 9-10 requests, so the mode is kept.

Per-layer kernel time re-measured on the rebased head (graph replay, 8 rows per request, 150K context, L2 flushed; "old" forces decode_splits, "new" is the default path of this PR). Relative difference between old and new outputs is 7e-4 to 8e-4 where the split count differs and exactly 0 where it does not:

N	old splits	old us	new splits	new us
1	8	13.0	4	11.5
2	8	13.0	4	11.7
3	8	13.1	4	11.8
4	8	13.2	4	11.8
5	6	15.9	4	11.9
6	4	12.1	4	12.1
8	4	12.5	4	12.4
12	2	16.1	2	15.7
16	2	16.2	2	15.9
24	4	21.9	4	21.3

N=6 and above pick the same split count as before; the 0.1-0.6 us differences there are run-to-run noise.

Decode step, this change alone (p50 of the gap between streamed chunks with N concurrent 150K-token streams, 3 passes, DBO off, VLLM_ROCM_USE_SKINNY_GEMM=0; the same-code rows below measure server-to-server noise, about +-0.6%):

N	1	2	3	4	5	6	8	12	16
random-token prompts	+0.3%	+0.2%	-1.2%	-1.8%	-1.4%	-1.1%	-0.2%	+0.5%	-0.6%
real-text prompts	-0.3%	-0.8%	-	-0.2%	-1.1%	-	-0.6%	+0.7%	-0.1%

The gain is -1.1% to -1.8% at 3-6 requests (about 0.15-0.25 ms per decode step); everything else is within the noise band. I therefore claim only the kernel-level change and this small step-level gain at 3-6 requests.

Correctness

• The splits only change the fp32 summation order of the softmax merge, so outputs are not bitwise equal to the old split counts. New default splits vs the unsplit kernel: relative error < 3e-3 (observed about 1e-3 on bf16 outputs) and bitwise repeatable, for N = 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 16, 17, 20, 24, 32, 40 requests x 8 and 4 rows x 30K / 150K / 500K contexts. The 21-case subset in this PR is added to tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py.
• tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py on the rebased head: 101 passed, 0 failed on MI355X (80 existing cases + 21 new: 1 _xcd_splits table test and 20 default-splits-vs-unsplit cases; vs an fp64 reference < 5e-3). The tests ran against an overlay of the installed tree with this PR's host.py and test file (main's minimax_m3/amd code needs a newer model_executor/models/utils.py than the machine's installed build); this branch differs from main only in those two files.
• End-to-end sanity ("What is 17*23? Answer briefly." -> 391) on a tree with this change stacked with other attention changes.

Limitations

• Tuned and measured for DSpark 7 (8 rows per request), top-k 16, one KV head per rank (TP4), on MI355X only; other row counts and 2-4 KV heads per rank (TP2, TP1) keep the old cost model (the table over-splits there, measured).
• Step-level effect is small and limited to 3-6 requests.

AI assistance: this change was developed with Claude Code; the measurements above were run on MI355X (gfx950).

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[Fangzhou-Ai]

I need some roofline analysis to proceed with this PR

[Fangzhou-Ai]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99751 for commit 9b1a2b1a2772.

[olka-amd]

outputs are not bitwise equal to the old split counts

If outputs differ let's run A/B gsm8k with and without this change (reported end-to-end sanity check is shallow).
Would be also great to run speed-bench to see if acceptance rate is affected in any way.

[Fangzhou-Ai]

GSM8K A/B: no accuracy change

Full GSM8K test set (n=1319), side by side on amd/MiniMax-M3.1-Flash-Quark-MXFP4 (MI355X gfx950, TP4, NVFP4 KV and indexer, DSpark 7 with real draft acceptance, DBO off, VLLM_ROCM_USE_SKINNY_GEMM=0). Both arms use the same server script and eval command; the only difference is this PR's host.py. Eval: chat mode (zero-shot, reasoning allowed, max_tokens 6000, "#### N" final line), concurrency 16, temperature 0.

arm	strict	flexible
this PR	0.9393	0.9386
stock main (dd4e401)	0.9409	0.9401

Per question (strict): 19 correct only with the PR, 21 correct only on main (McNemar exact p = 0.87). The 0.16 pp gap is well inside the sampling noise (SE about 0.65 pp per arm). Only 100 of 1319 outputs are textually identical, as expected: a different split count changes the fp32 softmax-merge order (about 1e-3 relative on bf16 outputs, see the Correctness section), which shifts greedy decoding for some questions, but the changes split evenly in both directions.

Sanity prompt ("What is 17*23?") returns 391 on both.

AI assistance: the runs were driven with Claude Code; I reviewed the numbers.

🤖 Generated with Claude Code

[micah-wil]

LGTM!

[micah-wil]

/ci run

[github-actions]

✅ Triggered Buildkite CI #100111 for commit 4712c5e81b5d.

[micah-wil]

Actually I just noticed this test file doesn't actually run anywhere, can we add it to MiniMax-M3 Attention Kernels?

[Fangzhou-Ai]

Done in 61d5ade: the file is now in the :amd: (MI355 DPX) MiniMax-M3 Attention Kernels step (source_file_dependencies and the pytest command in the mirror.amd block of .buildkite/test_areas/kernels.yaml). I did not add it to the B200 command: the module skips itself off ROCm / gfx950, so it could only ever run on the MI355 lane. The repo's "new tests are wired into Buildkite CI" pre-commit check passes. The file is 125 tests, about 1 minute on MI355X.

[peymanr]

PR Pundit review

(Testing PR Pundit Code Review agent. Not all comments may be relevant. This is a one-time automated pass: please address what is useful, no need to request a re-review. Please reply to each point, whether it applies or not: we harvest these replies, and both kinds are enormously helpful for tuning the agent. Reach out to @peymanr for suggestions.)

#		Finding	Where
1	🟠	Decode at TP=1/2 may over-split: this PR keys _xcd_splits on request count, ignoring num_kv_heads	vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/host.py:102-121
2	🟠	Split-selection regressions ship unnoticed: new tests only check numerical agreement with splits=1, not the chosen split count	tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py:259

1. 🟠 Should fix: Decode at TP=1/2 may over-split: this PR keys _xcd_splits on request count, ignoring num_kv_heads

vllm/models/minimax_m3/amd/ops/nvfp4_sparse_decode/host.py:102-121

_xcd_splits is keyed only on the request count, but the XCD grid puts all num_kv_heads and all splits of a request on one XCD, so per-XCD work grows with the KV heads per rank. The table was measured with one KV head per rank (the new test passes GROUP query heads), yet it also runs at 2 or 4 KV heads per rank (TP=2 or TP=1 with the config's 4 KV heads), where it overrides decode_splits. Example: 8 requests at 4 KV heads launch 64x4x4=1024 workgroups where the cost model picks 1 split (256). Whether this is slower was not measured; consider passing num_kv_heads into _xcd_splits, or returning None when it is above 1.

Who is affected: MiniMax-M3 NVFP4 sparse decode with DSpark 7 (topk=16, decode_query_len=8) on 8-XCD gfx950, run at TP=1 or TP=2 so each rank holds 2-4 KV heads.

Checked against the code by an independent pass.

Related past reviews: #56148 on vllm/v1/attention/ops/triton_unified_attention.py, #60990 on vllm/models/minimax_m3/common/sparse_attention.py, #21716 on tests/compile/test_fusion_attn.py

2. 🟠 Should fix: Split-selection regressions ship unnoticed: new tests only check numerical agreement with splits=1, not the chosen split count

tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py:259

test_default_splits_match_unsplit only checks output is finite, close to splits=1, and repeatable, so it would still pass if the call site fed _xcd_splits the wrong request count (rows instead of rows // decode_query_len) or the call-site hunk were reverted. It runs only at num_kv_heads=1 with no zero-length (padding) requests, leaving the shapes where the tuned choice may misfit uncovered. Assert the chosen split count (or patch _xcd_splits to record its arguments), and add a parametrization with num_kv_heads>1 and padded requests at qlen 8.

Who is affected: Only test coverage: CI on gfx950 would not catch a wiring regression at the host.py call site; runtime behaviour is unaffected.

Checked against the code by an independent pass.

Related past reviews: #58737 on tests/v1/attention/test_mamba_align_block_table.py, #61837 on tests/models/step5/test_amd_csa_paged.py, #52231 on tests/v1/determinism/test_mx_linear_batch_invariant.py

_{Shared from a pilot of PR Pundit, an agentic code reviewer that builds and tests the change and reads the project's past reviews. Reviewed at 4712c5e8; findings checked and posted by @peymanr; points already raised on this PR are left out.}

[Fangzhou-Ai]

Thanks, both findings are valid; fixed in 61d5ade.

1. num_kv_heads (valid, real regression). I swept explicit split counts at 1, 2 and 4 KV heads per rank (1-24 requests, DSpark 7, top-k 16, HIP-graph replay, per call), comparing the PR default with the old cost model:

KV heads	requests	PR default	old cost model
1	1-5, 9, 17	9.9-18.1 us	11.0-20.5 us (the PR is 8-14% faster)
1	6, 8, 12, 16, 24	same as old (+0.1-0.2%)
2	16	21.3 us	18.3 us (+16% for the PR)
4	8	21.5 us	18.4 us (+17%)
4	16	37.0 us	30.3 us (+22%)
4	24	55.5 us	42.8 us (+30%)

(Not a uniform loss, so the numbers above are the worst cases: at 2 KV heads the default is equal or faster up to 9 requests, by up to 20% at 1-2 requests, and equal or slower from 12; at 4 KV heads it is 28% faster at 1 request, within 3% of the old model at 2-6 and 12, 4.5% faster at 9, and slower at 8, 16, 17 and 24. I did not tune those shapes.)

So the table is right at one KV head per rank (TP4, what it was measured on), and over-splits at 2 and 4 (TP2, TP1). _xcd_splits now takes num_kv_heads and returns None unless it is 1, so TP2 / TP1 keep the old behaviour exactly. Retuning TP2 / TP1 is a separate piece of work; I did not widen this PR.

2. Tests (valid). Added test_default_splits_are_chosen_per_xcd: it records the arguments of _get_kernel through the real minimax_m3_nvfp4_sparse_decode_fly entry point and asserts the chosen split count (the table values at one KV head, the cost model or 1 split at 2 and 4), over 1/2/4 KV heads x 1, 5, 8, 9, 16, 17, 24, 40 requests with every third request padding (seq_len 0), and checks zero rows for padding and agreement with the unsplit kernel. test_xcd_splits_table covers the new argument. Checked that it catches the wiring bugs you describe: feeding rows instead of rows // decode_query_len fails 3 cases, removing the KV-head guard fails 14. The file passes on MI355X (125 tests) and is now in the AMD CI step.

[Fangzhou-Ai]

@olka-amd SPEED-Bench A/B on acceptance; the GSM8K A/B is the earlier comment above (0.9393 vs 0.9409 strict, McNemar p = 0.87).

Setup. SPEED-Bench Qualitative split (880 prompts, 11 categories, built with the NeMo-Skills prepare.py), vllm bench serve --dataset-name speed_bench --num-prompts -1 --max-concurrency 16 --temperature 0 --speed-bench-output-len 1024. Same server flags as the GSM8K run (MI355X, TP4, NVFP4 KV and indexer, DSpark 7 with real draft acceptance, block rejection, DBO off). Two arms from the same tree, only host.py differs: this PR at 61d5ade and stock main (dd4e401). Two full passes per arm on the same server, so the pass-to-pass spread gives a same-code noise level. Acceptance is read from exact deltas of the server's vllm:spec_decode_num_* counters around each pass. All 880 requests succeeded in every pass.

arm	pass	acceptance rate	acceptance length	position 0 / 1 / 2 / 3 / 4 / 5 / 6 (%)
this PR	1	31.52%	3.206	71.4 / 49.2 / 34.4 / 24.6 / 18.0 / 13.3 / 9.8
this PR	2	31.22%	3.186	71.1 / 48.7 / 34.0 / 24.2 / 17.7 / 13.1 / 9.7
main	1	31.45%	3.201	71.3 / 49.2 / 34.4 / 24.5 / 17.8 / 13.2 / 9.9
main	2	31.37%	3.196	71.1 / 49.0 / 34.3 / 24.4 / 17.7 / 13.2 / 9.9

Both passes combined: acceptance length 3.196 (this PR) vs 3.199 (main). The difference between arms (0.003) is smaller than the pass-to-pass spread inside one arm (0.021 for this PR, 0.005 for main), and every position agrees to within 0.3 pp, so I see no effect on acceptance.

This measures acceptance only; the two arms ran on different GPU sets (0-3 and 4-7), so I am not using it for speed. The decode-step effect is in the PR description.

One caveat on the numbers I had before: the GSM8K server logs showed mean acceptance length 4.63 (PR) vs 4.71 (main). That run was not paired (the generated texts differ between arms, 1219 of 1319 answers are not identical) and had no repeat, so I would not read it as an acceptance change; the controlled run above is the better measurement.

AI assistance: the runs were driven with Claude Code.

🤖 Generated with Claude Code

[Fangzhou-Ai]

/ci run

[github-actions]

✅ Triggered Buildkite CI #100347 for commit 2a4a5f0f2703.

[Fangzhou-Ai]

The MI355 DPX MiniMax-M3 Attention Kernels failure was in the new test_default_splits_are_chosen_per_xcd (all 24 cases, assert 0 == 8), not in the kernel.

That runner is a DPX-partitioned MI355X (DPX/NPS2 in the log), so each device reports fewer than 8 XCDs. _xcd_count correctly returns 0 there (the XCD-aware grid is only enabled for 8 XCDs on gfx950, the topology the table was tuned on) and the kernel falls back to the cost model, but the test hard-asserted 8.

c3f095e:

• skips that test unless the device has 8 XCDs;
• adds test_default_splits_without_the_xcd_grid, so partitioned devices still cover the default path (plain grid, cost model, agrees with the unsplit kernel);
• adds test_xcd_grid_needs_the_tuned_topology, which pins when the grid is on (8 XCDs on gfx950 only) without needing a GPU.

Checked locally on an SPX MI355X (8 XCDs): the per-XCD test runs and passes, as do the new tests. With the HIP XCD query forced to 4, the per-XCD test skips and the fallback tests pass.

[micah-wil]

/ci run

[github-actions]

✅ Triggered Buildkite CI #100393 for commit 97088ed51d1a.