[jiacao-amd]

Overview

Tune gfx950 MXFP8 linear tiles for MiniMax-M3.1 decode and small prefill batches (16 < M <= 256, local N > 1280). Keep the main configuration for TP2 QKV at 128 < M <= 192, where the tuned tile regresses.

Claims

• Tile-only change; reuse the existing Triton GEMM and preserve the fallback outside the tuned range.
• One mid-M table owns the overlapping gfx950 17 <= M <= 64, N > 1280 range with [ROCm] MXFP8 linear: one tile row for 17-64 token decode batches #64195; integrate that overlap here rather than stacking selectors. This PR also tunes M65–256; [ROCm] MXFP8 linear: one tile row for 17-64 token decode batches #64195's additional small-N coverage is outside its scope.
• Keep the measured M48 MiniMax TP4 QKV tile in the table. Remove the V4.1-specific N=1792 bucket; normal V4.1 serving uses the separate block-scaled GEMM.

Validation

MiniMax-M3.1 GEMMs: main ffdf26702 vs 4bfce4e2e

MI355X, same image and inputs; 36 cases covering all six TP4/TP2 shapes at M48/64/96/128/192/256. The 34 changed tiles reduce GEMM latency by 4.64–58.15%. M48 TP4 QKV keeps its main tile. TP2 QKV now also keeps the main tile for 128 < M <= 192.

The initial M192 TP2 QKV measurement regressed 23.296 → 26.608 µs (+14.22%). That shape/range is excluded from the tuned path; M128/129/160/192/193 boundary checks confirm the intended dispatch. The other measured buckets are retained.

All 36 GEMM timings and measurement method

Values are main → PR, µs/GEMM; these are kernel timings.

Linear (N, K)	M48	M64	M96	M128	M192	M256
TP4 QKV (2560, 6144)	10.903 → 10.888	17.566 → 11.468	19.979 → 13.728	20.162 → 13.798	20.453 → 18.515	20.534 → 18.824
TP4 output (6144, 2048)	12.495 → 6.243	13.992 → 6.334	9.148 → 8.724	9.300 → 8.764	16.743 → 12.705	16.841 → 13.096
TP4 shared gate/up (1536, 6144)	9.379 → 7.353	9.961 → 7.501	16.243 → 11.439	17.960 → 11.546	20.216 → 13.649	20.044 → 13.782
TP4 shared down (6144, 768)	4.914 → 3.992	5.215 → 4.122	4.974 → 4.549	5.074 → 4.668	6.833 → 6.263	6.956 → 6.498
TP2 QKV (4992, 6144)	19.526 → 14.003	26.550 → 14.136	20.795 → 19.558	21.030 → 19.582	22.902 → 22.890	41.073 → 35.953
TP2 dense gate/up (12288, 6144)	42.104 → 22.029	53.383 → 22.342	42.124 → 36.388	42.721 → 39.517	62.752 → 59.103	67.605 → 62.010

Prequantized E4M3 activations/weights with E8M0 scales and BF16 output, seed 4671. Each graph rotates ceil(512 MiB / (N*K*33/32)) weight copies. Seven alternating rounds, 40 graph replays per round; report the median device-event time per GEMM. Quantization is outside the timed region.

Runtime: vllm/vllm-openai-rocm@sha256:9cef62361792bd9cc621725c3af799993b6a68b6b62e3778dcf9b65ff6fcd872. The image's baseline native GEMM source is byte-identical to main ffdf26702; the candidate source and tests match 4bfce4e2e.

Revalidated the final source on MI355X: 296 kernel tests passed, 0 skipped; six two-stream graph checks passed. Pre-commit passed.

.venv/bin/python -m pytest -q tests/kernels/quantization/test_rocm_mxfp8_linear.py

900 s AgentX: main 380c8e438 vs tiles only

4× MI355X, TP4, AgentX 1.0.12; 900 s per arm after warmup. Model: amd/MiniMax-M3.1-Flash-Quark-MXFP4 at 0a00e0e39fc00499b23efccc18a8ea63fce5c44c. InferenceX #4671 at bed892b82, DSpark K7, synthetic AL 4.04, seed 42.

Values are baseline → candidate. p90 interactivity uses output_token_throughput_per_user.p10; total throughput includes cached input tokens.

Candidate: 99ec093d5. One pair per concurrency; norm+quant fusion is already in the baseline.

Comparison	p90 interactivity (tok/s/user)	Output tok/s/GPU	Total tok/s/GPU
C4	287.74 → 325.33 (+13.06%)	61.82 → 64.67 (+4.61%)	7,358.65 → 7,458.75 (+1.36%)
C16	135.71 → 144.60 (+6.55%)	165.79 → 164.52 (-0.77%)	19,727.25 → 19,617.63 (-0.56%)
C32	89.83 → 87.81 (-2.26%)	355.57 → 361.78 (+1.75%)	47,170.00 → 47,906.26 (+1.56%)

ITL and TTFT (baseline → candidate)

Comparison	Mean ITL (ms)	p90 ITL (ms)	Mean TTFT (ms)	p90 TTFT (ms)
C4	2.293 → 2.209	3.470 → 3.054	379.025 → 384.964	930.524 → 936.523
C16	5.267 → 5.338	7.364 → 6.903	538.874 → 538.037	868.419 → 1,053.297
C32	7.937 → 7.625	11.125 → 11.382	745.069 → 681.098	1,347.779 → 1,254.675

Configuration and reproduction

Both arms use one image rebuilt from main 380c8e438 (pinned when this validation started), including Python and C++.

C16/C32 use NVFP4 KV/indexer, V2 + DBO and a 32K prefill budget. Both arms include the same text-only DBO startup-guard fix (supports_multimodal_inputs). C4 uses FP8 KV/indexer, mono decode, an 8K prefill budget and no DBO.

Use the corresponding #4671 serving lane. Set CONC to the table's concurrency and TOKENIZER to the pinned model snapshot:

agentx profile --url http://127.0.0.1:8000 \
  --model amd/MiniMax-M3.1-Flash-Quark-MXFP4 \
  --tokenizer "$TOKENIZER" \
  --concurrency "$CONC" --benchmark-duration 900 \
  --tokenizer-trust-remote-code \
  --server-metrics http://127.0.0.1:8000/metrics \
  --public-dataset semianalysis_cc_traces_weka_062126

GSM8K accuracy

Original scoring, all 1,319 questions per run; TP4/C16, 5-shot, temperature 0, seed 42, reasoning effort low, 4096-token output limit, and real DSpark K7 block rejection.

Repeated evaluation: two runs per arm, using the same model revision and serving configuration above and one image rebuilt from main 380c8e438, including the shared DBO startup fix. Candidate: 99ec093d5 (tiles only).

Configuration	Run 1 correct / 1319	Run 2 correct / 1319	Mean accuracy	Mean delta
Baseline	1,265	1,269	96.06%	—
Tiles only	1,273	1,267	96.29%	+0.23 pp

All four runs completed without API errors. Prefix-cache resets succeeded before every run; token-limit answers remain in the original scores.

Earlier single-pair GSM8K result: 1,274 → 1,267 correct (96.59% → 96.06%); invalid answers 0 → 0.

Details

Tile selection uses M, local N and K; BLOCK_K divides K. No new kernel or model interface. This follow-up changes no MiniMax TP4 tile, so the TP4 serving and GSM8K results above retain their original attribution. The new TP2 exclusion uses the existing main fallback.

AI assistance: OpenAI Codex and Claude Code.

Pull Request Checklist

• PR checklist used; AI assistance disclosed.
• Design fit, tests, documentation and CI coverage reviewed.
• Requested performance validation and model evaluation complete.
• Attribute the observed GSM8K accuracy difference.
• Hosted CI complete.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jiacao-amd.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[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.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jiacao-amd.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[ChuanLi1101]

Thanks for the detailed write-up. A few technical asks before approving:

1. Please re-run tests/kernels/quantization/test_rocm_mxfp8_linear.py on MI355X against this head. The PR states GPU correctness was not revalidated after the rebase.

2. Please post a tile-only AgentX (or equivalent serving) delta vs current main on the same image. The historical C16 numbers combine tiles with the old norm fusion, so they cannot be used to attribute this change.

3. Please reconcile the mid-M selector with [ROCm] MXFP8 linear: one tile row for 17-64 token decode batches #64195 for the overlapping 17 <= M <= 64 range, or document which table wins if both land. Landing both without a single owner will make later retuning hard.

4. The special case m_bucket == 48 and 1536 < N <= 2560 keeps the 16x16 tile. Please either fold it into the table or add a short comment with the measured shapes that require it, so it is not accidental.

5. The earlier DeepSeek wq_a|wkv M=64 regression (12.0 → 12.6 µs) should be rechecked on current main. If it still regresses, call that out explicitly or exclude those shapes from the mid-M path.

[jiacao-amd]

Thanks @ChuanLi1101 Addressed in 99ec093:

1. Re-ran tests/kernels/quantization/test_rocm_mxfp8_linear.py on MI355X against this head: 296 passed, 0 skipped. Six two-stream graph checks also passed.

2. Added tile-only 900s AgentX comparisons at C4/C16/C32 against main 380c8e438, pinned when validation started. Each pair uses the same rebuilt image, with only the tile patch toggled. The PR body includes all throughput and latency deltas, including the C16 throughput and C32 p90 declines.

3. Consolidated the overlapping mid-M dispatch into one table, using [ROCm] MXFP8 linear: one tile row for 17-64 token decode batches #64195's tile for the V4.1 wqa/wkv bucket. Its additional small-N coverage remains outside this PR.

4. Folded the M48 exception into the table and documented the measured shapes.

5. Rechecked the DeepSeek wqa/wkv M64 case against the pinned main: 9.315 → 7.823 µs. The earlier regression is resolved.

The description is updated with the results and reproduction details. Hosted CI remains pending the repository's pre-run gate.

[github-actions]

✅ @jiacao-amd, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[jiacao-amd]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99795 for commit 753f28ea9d53.

ℹ️ This PR is 188 commits behind upstream main. Merged upstream main into this branch as 753f28ea9d53 before starting CI. Pull the branch before pushing more changes.

[ChuanLi1101]

Looks good — prior asks and okorzh's threads are addressed (V4.1/1792 cleanup, M3.1 per-GEMM table, TP2 M192 fallback). LGTM.

[simondanielsson]

Thanks for the fix!

[simondanielsson]

Question: Will this affect other models as well? Ideally we'd only add this for M31 shapes