[ganyi1996ppo]

Overview

Adds --kv-cache-dtype mxfp4 with --attention-config '{"indexer_kv_dtype": "mxfp4"}' for MiniMax-M3.1 on ROCm gfx950 (MI355X, TP4): K, V and the indexer's key cache are stored as MXFP4 (e2m1 values, one e8m0 scale per 32), 68 B per token and head. Reference implementation: ROCm/ATOM#2770.

Draft; not runnable yet. #64111 has merged, so this is now rebased onto main (16d667348d) as 5 commits (4 feature commits plus a test-only type fix). It still waits on three open AITER PRs (#7056, #7060, #7061), merged #7057, and an AITER pin bump (see Details). Placeholder to claim the work and get early design feedback.

AI-assisted (Claude Code). No human has reviewed the diff line by line yet; this stays a draft until that is done.

Claims

• New cache dtype mxfp4 (uint8 storage, KVQuantMode.MXFP4) and a gfx950 attend for MiniMax-M3 sparse attention. Per head a 128-token page is [128 x 64 B e2m1][128 x 4 B e8m0] = 8704 B; the main K/V spec is num_head_slots=2*Hkv, state_content_bytes=68, the index spec head_size=68.
• The fused QK-norm / RoPE kernel (AITER fused_qknorm_idxrqknorm) writes the MXFP4 K, V and index-K pages and writes the main q as e4m3 into q_out.
• Index scores come from AITER's FlyDSL MXFP4 scorers (decode and prefill); the top-k is the one the NVFP4 indexer uses; the attend runs AITER's pa_decode_mxfp4 on a one-block-per-page table, for decode and prefill.
• The MXFP4 decode scorer's work map is one persistent buffer built for the model-length block bound, so FULL decode CUDA graphs replay a valid map.
• Not covered: Mono decode (stays off), indexer CP, TP other than 4, the MTP/DSpark draft layers (the existing check raises).

Validation

Rebased onto main 16d667348d (live main tip when this was pushed); git range-diff against the previous head d24ae8f7f2 shows the 4 feature commits patch-identical; the 5th commit (d6dbccaac4) only touches tests/models/minimax_m3/test_mxfp4_kv_cache.py. Checked on a CPU-only host, no GPU used:

Check	Result
pytest tests/models/minimax_m3/test_mxfp4_kv_cache.py tests/kernels/attention/test_minimax_m3_mxfp4.py	9 passed, 1 skipped (the kernel test file skips: needs gfx950 FlyDSL kernels)
pre-commit run (all hooks, including mypy) on the 14 changed files, at d6dbccaac4	pass. The earlier head ffd9f4a06d had 3 mypy errors in tests/models/minimax_m3/test_mxfp4_kv_cache.py (dynamic attributes on the fake aiter module, an Optional decode.max_seq_len); fixed in the 5th commit by typing the fake module as Any and asserting decode is not None, assertions unchanged.

CI at d6dbccaac4: pre-commit and pre-run-check pass; docs build was still pending when this was written.

Not validated after the rebase: any GPU run. The kernel tests, end-to-end serving, CUDA-graph capture, accuracy and speed were not re-run on this head. Results quoted in earlier versions of this description were measured on the pre-rebase head d24ae8f7f2 on an ROCm image and are not repeated as evidence for this head. CI has no gfx950 job that can run these ops: the kernel tests skip until the AITER pin carries them.

For orientation only, ATOM's own measurements of the same kernels (ATOM#2770, MiniMax-M3.1 TP4, MXFP4 KV and index against the NVFP4 path): c32 +7.4% throughput and +26.2% P90 interactivity, c48 +8.1% and +32.3%; KV capacity 13.4M against 7.3M tokens for fp8. These are ATOM results, not vLLM results.

Details

Prerequisites. ROCm/aiter#7057 (fused write kernel: MXFP4 K/V and index-K, e4m3 q_out; this PR follows its final API, in which q_fp8_out and q_scale are gone) has merged; the remaining ROCm/aiter#7056 (FlyDSL pa_decode_mxfp4), #7060 (index decode scorer) and #7061 (index prefill scorers) need to merge, then an AITER pin bump in docker/Dockerfile.rocm_base, then this PR. The pin is not bumped here. AITER is imported inside functions behind VLLM_ROCM_USE_AITER=1; without the ops the layer raises a message naming those PRs.

Commits (on 16d667348d):

1. 99556ddca5 Add the mxfp4 KV cache dtype and its page geometry (no behavior change on its own)
2. 507ac36aef MXFP4 K/V and index cache on gfx950 (indexer, attend, write, tests, docs)
3. 06fdab1145 Pass the e4m3 main q as q_out to the aiter MXFP4 write
4. ffd9f4a06d Keep the MXFP4 decode work map stable for CUDA-graph replay
5. d6dbccaac4 Fix mypy errors in test_mxfp4_kv_cache.py (test-only)

Why not a duplicate. No open or merged PR stores MiniMax-M3 K/V or index keys as MXFP4. #63547 is fp8 KV, #45838 is MXFP4 experts, #59739 (mxfp4_mla) and #64120 (mxfp4_ds_mla) are MLA/DSv4 caches. The dtype string mxfp4 is a distinct exact-match name from mxfp4_mla and mxfp4_ds_mla; KVQuantMode.MXFP4 = 16 avoids the 15 #64120 takes.

Design.

• It mirrors the NVFP4 hooks (is_mxfp4 dispatch in amd/model.py, a sibling indexer and sparse backend), and reuses its metadata builders, FlyDSL top-k and the existing sparse_pa block-table kernels, which gain a pages_per_block argument (1 for a page-128 cache).
• AITER's pa_sparse_block_topk is not used: its build only accepts pages_per_block == 8, and it expects scores already pinned by the scorer, which the MXFP4 scorers do not do. A small Triton kernel writes AITER's init/local sentinels (1e30 / 1e29) on each row's causal blocks, and the existing top-k reads only a row's causal blocks, so the scorer's unwritten entries are never seen.
• The index query is bf16 and the main query e4m3 (the configuration ATOM measured).
• It needs one KV head per rank (TP4 on the 4-KV-head checkpoint); the checks raise a message with the reason otherwise.

Size. About 850 lines of non-test code, most of it the two new files (indexer_mxfp4.py, sparse_attention_mxfp4.py) that parallel the NVFP4 ones, plus about 700 lines of tests. Docs: one line in docs/features/quantization/quantized_kvcache.md.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[mergify]

Documentation preview: https://vllm--64198.org.readthedocs.build/en/64198/

[mergify]

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

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

[ganyi1996ppo]

Pushed 91d7ee1 (one fast-forward commit on top of d6dbcca): the MXFP4 sparse attend now calls aiter's pa_decode(out, q, k_data, v_data, ctx, bt, scale, 1, key_scale=, value_scale=, work_plan=None, max_context_length=) instead of aiter.ops.flydsl.pa_decode_mxfp4; max_partitions=1 is dropped (a plan-free call already runs one task per row). rocm_mxfp4_aiter_unavailable_reason probes pa_decode instead of the removed module, and the mxfp4_pages docstrings are updated. Only vllm/models/minimax_m3/amd/{sparse_attention_mxfp4,indexer_mxfp4,ops/mxfp4_pages}.py changed.

Why: ROCm/aiter#7056 (head f017c23591) merges the standalone pa_decode_mxfp4 entry into pa_decode at the review's request; without this commit the availability probe reports the MXFP4 ops unavailable and the sparse path cannot run once #7056 lands. Depends on ROCm/aiter#7056.

Tested on gfx950 against aiter#7056 f017c23591 (same PA files), from an archive of this commit overlaid on the vllm nightly rocm100 image (0.32.0rc1, torch 2.13; python overlay, not a source build): tests/kernels/attention/test_minimax_m3_mxfp4.py -k sparse_attend, 6 passed; ruff and ruff format clean. Not run: the rest of that file (index/write/layer tests), model serving, CI.