[mdavis36]

Overview

Lets MiniMax M3.1's NVFP4 index cache (#64111) be sharded over TP ranks with indexer CP, which #61929 and #63671 support only for FP8. Depends on #63671 (itself on #61929): only 7 commits are new here (this PR's commits), and the branch also merges #63671's latest head and main. Like #61929, it needs vllm-project/MSA#46 (FMHA_SM100_SRC_DIR) until the MSA pin moves past it.

Not a duplicate: no open PR shards the NVFP4 index cache. #59980 (M3, FP8) and #64448 (ROCm) split decode scoring over a replicated cache, as main's opt-in indexer_decode_cp does.

Claims

• With --attention-config '{"indexer_kv_dtype":"nvfp4","indexer_cache_sharding":"on"}' (off by default), each TP rank stores only its rows of every NVFP4 index page: 36 B instead of 72 B per token per sparse layer at TP2, 18 B at TP4. Supported: TP2 and TP4 at --block-size 128 (256 cannot serve yet, see below), indexer_prefill_cp='gather'.
• Decode rounds index-K as the stock Q8KV4 decode indexer does on each arch (to e4m3 on SM100/SM103, exact on SM107), so the sharded CuTe scorer's page scores match the stock kernel's up to FP32 summation order.
• At P256/TP2 with eight-query decode (DSpark k=7), a rank's shard pages are stock NVFP4 pages, so decode scores them with the stock Q8KV4 decode indexer instead. On GB200 and GB300 that makes the per-layer decode indexer 1.35-1.4x faster than with the CuTe scorer at 16 requests and 1.6x faster at 64 (see Validation), and its scores are bitwise identical to the stock scores of the unsharded pages. VLLM_MINIMAX_M3_ICP_Q8KV4_DECODE=0 keeps the CuTe scorer. Serving cannot reach this path yet, so only kernel tests cover it: at --block-size 256, [Attention] Add P256 MiniMax M3 indexer CP layout #63671 lets NVFP4 main attention keep 256-token pages, but [Feature] Add MiniMax M3 indexer context-parallel kernels MSA#46's Q8KV4 prefill accepts only 128-token pages (found by reading the code, not run).
• The gather prefill stages each rank's shard pages to FP8 with the stock NVFP4 prefill's kernel, then all-gathers them.
• End to end on GB300 at TP2 and --block-size 128 (so the CuTe scorer), one 1-hour agentic run at concurrency 16 per variant: sharding gives 9.8% more KV tokens at the same prefix-cache hit rate, but ITL rises 25% and throughput falls 2%. An MRCR accuracy check is inconclusive: pooled −1.1 points (95% CI −2.2 to +0.1). See Validation.
• With sharding off, behavior is unchanged apart from the shared paths listed under Details.

Validation

pytest tests/kernels/attention/test_minimax_m3.py
pytest tests/kernels/attention/test_minimax_m3_indexer_prefill_sharded.py
pytest tests/distributed/test_minimax_m3_indexer_sharding.py  # 4 GPUs: TP2, TP4
pytest tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py  # needs a source build

These ran before the branch merged main 3d23773f1, with vllm-project/MSA#46 (f380d0e) and the branch's Python, either over a prebuilt vLLM ("prebuilt") or over the CUDA 13.4 nightly image for main 3966924d5 with _C_stable_libtorch rebuilt from the branch ("image"):

Arch	Base	Build	Tests	Passed	Skipped	Failed
SM100, SM103	#63671 at a76900d6e	image	test_fused_minimax_m3_qknorm_rope_kv_insert.py	168	0	0
SM100, SM103	#63671 at a76900d6e	image	test_minimax_m3.py -k 'sharded or q8kv4 or gather_nvfp4'	335	0	8 (a)
SM100, SM103	#63671 at a76900d6e	prebuilt	test_minimax_m3.py -k 'sharded or q8kv4 or gather_nvfp4 or gather_prefill_pages or sharding_policy'	334	6	8 (a)
SM100, SM103	#63671 at 54a5c7113	prebuilt	same	342	6	0
SM107	#63671 at 54a5c7113	prebuilt	test_minimax_m3.py	481	55	14 (b)
SM107	#63671 at 54a5c7113	prebuilt	test_minimax_m3_indexer_prefill_sharded.py	164	0	0
SM107	#63671 at 54a5c7113	prebuilt	test_minimax_m3_indexer_sharding.py	2	0	0

(a) test_sharded_candidate_merge_deduplicates_and_forces, from #61929's dc68fa728 (its test indexed a CPU tensor with a CUDA mask), not this PR. #63671's 9fcb12446, which the branch now merges, fixes the test.
(b) main's stock NVFP4 indexer tests (test_nvfp4_indexer_impl_matches_reference, test_nvfp4_indexer_decode_cuda_graph_replay), which fail the same way without this PR: their reference rounds K to e4m3, which the stock kernel does not on SM107.
The 6 prebuilt skips on SM100/SM103 are the QMUL4 gather variants, which that Triton cannot build; they run and pass in the image.

Decode indexer time per layer and rank in µs (scoring, plus candidate publish and merge for ICP), at TP2, --block-size 256 and eight queries per request:

GPU	Requests	Context	Replicated stock	ICP, CuTe scorer	ICP, stock Q8KV4 scorer
GB200	16	128K	44.1	78.0	57.6
GB200	16	168K	53.7	88.3	63.7
GB200	64	128K	127.6	206.9	125.1
GB200	64	168K	154.1	239.8	147.6
GB300	16	128K	44.6	78.2	57.8
GB300	16	168K	53.1	88.7	65.7
GB300	64	128K	125.9	207.2	125.4
GB300	64	168K	159.2	240.0	147.9

These come from one GPU with both ranks simulated: each kernel runs as its own CUDA graph with a cold L2, and each time is the median of 200 runs, averaged over two runs that agreed within 2%. NVLink latency and rank skew are not modeled, and timing ICP's three graphs separately adds a few µs against the stock path's one.

The new >65535-page gather test fails on the previous kernel (RuntimeError: Triton Error [CUDA]: invalid argument) and passes with this PR.

End to end. MiniMax-M3.1-Flash on GB300 in the image above, one TP2 server per run: --block-size 128, NVFP4 KV and index cache, DSpark k=7, the V2 model runner, indexer_prefill_cp='gather', and sharding off or on. "On" decodes with the CuTe scorer (P128 shards are not stock pages) and, unlike "off", skips the fused MXFP8 QKV epilogue. Each server's log shows the expected mode, and none logged a traceback or error.

Performance, from one 1-hour run per variant of InferenceX's agentic coding trace (AgentX, thinking on) at concurrency 16, with synthetic draft acceptance 4.04:

	Off	On	On/off
KV cache size (tokens)	5,492,736	6,030,720	1.098
Total throughput (tok/s/GPU)	55,344	54,171	0.979
ITL mean / P90 (ms)	3.82 / 5.70	4.77 / 7.17	1.25 / 1.26
TTFT mean / P90 (ms)	753 / 1,554	771 / 1,504	1.02 / 0.97
Request latency P90 (s)	9.69	11.44	1.18
Prompt tokens read from the prefix cache	97.20%	97.26%

The prefix-cache hit rate is the same in both runs, so this workload gains nothing from the extra capacity. In both runs the same 3 requests (of 2,578 and 2,535 measured) returned no content text, which the client counts as errors.

Accuracy, from MRCR with tests/evals/mrcr's loader and scorer on openai/mrcr, at temperature 0, max_tokens 4096 and real draft verification, on the same samples in both variants:

Samples	Prompt tokens	Off	On	On − off (95% CI)	Identical responses
600 (16-64K bins, concurrency 8)	16-65K	0.9889	0.9850	−0.0040 (−0.0122, +0.0038)	84%
150 (128-256K bin, concurrency 4)	133-260K	0.9328	0.8956	−0.0372 (−0.0868, +0.0105)	69%
750 pooled		0.9777	0.9671	−0.0106 (−0.0225, +0.0008)	81%

No difference is significant (pooled sign test p = 0.12: 63 samples better, 78 worse), but the interval does not exclude a 2-point drop, so the check is inconclusive. In the long arm, 12 samples hit the 4096-token limit in only one variant (8 with sharding on, 4 with it off). They make up 90% of that arm's difference; the other 138 differ by −0.004.

Not run yet, so this stays a draft:

• serving at --block-size 256 (blocked, see Claims) and at TP4, and repeats of the single end-to-end runs above;
• the GPU tests and end-to-end runs above, repeated since the merge of main 3d23773f1;
• in CI: the multi-GPU test file has no CI step, and CI covers SM100 only.

Details

• Motivation: KV capacity. Without sharding every rank stores the whole NVFP4 index cache; sharded, TP2 halves and TP4 quarters it, and the freed memory goes to KV tokens.
• Tradeoffs:
  • The gather prefill all-gathers FP8 rows (128 B rather than 72 B per token), and like the FP8 gather of [Attention] Add MiniMax M3 indexer CP for prefill and decode #61929 it neither chunks nor reserves its per-layer transient. Main's extensible KV cache, on by default with the V2 runner, now leaves less free memory for that transient than the old 0.92 gpu_memory_utilization default did.
  • A sharded cache disables the fused MXFP8 QKV epilogue, which writes whole NVFP4 index pages.
  • Sharded prefill ranks pages with MSA's fmha OnlyScore and sparse_topk_select instead of the stock Q8KV8 prefill indexer, so its selection can differ from the replicated path at exact ties.
• Why a second decode scorer: only at P256/TP2 is a rank's shard a whole stock NVFP4 page, so only there can the stock Q8KV4 decode indexer, which already serves the replicated cache, score it. It needs full eight-query tiles, because the candidate publish reads contiguous [tokens, heads, columns] scores. Every other shape keeps the CuTe scorer.
• Policy: at TP4, indexer CP yields to the opt-in indexer_decode_cp, which keeps a replicated cache. local_score prefill and KV connectors (including CPU offload) are not supported with NVFP4 sharding yet. dspark is now accepted for sharding.
• Shared paths: the C++ insert's idx_block_size runs on every M3.1 deployment; the NVFP4 gather kernel now puts pages on grid x, since one sharded staging call can exceed grid y's 65,535 limit, and main's unsharded prefill uses that kernel too; [Attention] Add MiniMax M3 indexer CP for prefill and decode #61929's FP8 scorer epilogue is factored into fold/store helpers shared with the NVFP4 kernel; the Q8KV4 path calls MSA's private _run_scores, as main's decode CP does.

AI assistance (Claude Code) was used to write this change and its tests.

---

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.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[mergify]

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

[mergify]

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

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

[mergify]

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

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