[None][perf] MiniMax-M3 MSA: split the KV axis for the decode indexer proxy

[hyukn]

@coderabbitai summary

Description

The MSA indexer proxy is a full-kv_len MQA max-score pass that scores every KV block so the top-k selection can pick 16. At the canonical decode point it costs 1576.5 us/step, 8.01% of the step, while sitting at 16.1% of its own KV-read byte floor and under 1% MFU, which makes it the largest device cost in the M3 decode step that is not already at a floor.

Why it is slow. The proxy runs on a persistent grid: get_grid_shape returns dim3(num_sm) regardless of how much work exists (fmha_tile_scheduler.hpp:126-129), and each CTA reads its slice out of packed_work_range (plan.cuh:386-388). Work items are batch * num_qo_heads * ceil(qo_len/qo_tile_size), and at decode _compute_pack_factor folds all four index heads into one packed qo head (api.py:110-120,557-558), so a batch of 16 yields 16 work items on a 152-SM device. The other CTAs pick up an empty range and exit, while the 16 busy ones each walk ceil(8700/256) = 34 KV tiles serially.

Why the head axis is the wrong lever. The binding quantity is serial depth per CTA, not CTA count. Setting pack_factor=1 raises the work-item count to 64 but leaves the chain at 34 iterations, because the causal clamp eff_kv = min(q_end + offset_q, kv_len) absorbs the whole q_end reduction once offset_q is 8699 (plan.cuh:136-146). It also re-reads K once per head, taking KV traffic from 35.6 MB to 142.6 MB for no reduction in wall time.

Changes

• Enable num_kv_splits for the decode indexer-proxy plan, which is the axis that actually shortens the per-CTA tile chain (plan.cuh:279-283). The proxy plan had it pinned to 1. Default is 4, with TRTLLM_MINIMAX_M3_MSA_PROXY_KV_SPLITS as an override; setting it to 1 restores the previous single-piece plan exactly.
• Prefill and mixed batches keep num_kv_splits=1. Those request out, so the split-KV O/LSE reduction does run there and would need its own parity evaluation before being enabled.
• Extend _MsaGraphSafePlan to mirror the split-KV worklists (kv_tile_begin_indices, kv_tile_end_indices, kv_split_indices, num_kv_splits_per_row) and both split-KV workspaces (workspace_o, workspace_lse) into graph-stable buffers, alongside the plan mirrors that already existed. Widths follow the vendor's own expressions: api.py:704-709 sizes packed_work_info and the three kv_tile_* worklists at one shared max_work_items = 131072 * max(num_kv_splits, 1), api.py:711 sizes workspace_o at total_qo_len * num_kv_splits * num_qo_heads * 128, and api.py:715 sizes workspace_lse at num_kv_splits * total_qo_len * num_qo_heads.
• Mirror rather than pin both workspaces. workspace_o goes through _alloc_workspace_buf, which reallocates when the requested size grows, and workspace_lse through _alloc_perplan_buf, which reallocates on every call (api.py:78-101), so neither address is stable enough to pin.
• Grow a mirror to the plan's actual width while steps are still eager. A CUDA graph only needs the address to be stable across replays, and before capture there are none, so the __init__ widths act as pre-allocation hints and the steady state allocates nothing. After capture the addresses are frozen and a short mirror is a hard error.

Python only, no rebuild required.

Performance

Kernel

Standalone proxy pass on GB300, CUDA-graph replay timings at the traced decode geometry (batch 16, kv_len ~8700, 4 index heads, page_size 128, bf16, OnlyScore):

num_kv_splits	busy CTAs	us	max_score vs splits=1
1	16/152	32.16	reference
2	32/152	21.95	bitwise-identical
4	64/152	15.81	bitwise-identical
8	128/152	15.78	bitwise-identical
16	128/152	15.78	bitwise-identical
32	128/152	15.74	bitwise-identical

2.04x, saturating at the predicted 34/4 = 8 cap. Time is linear in serial depth (ki = 2/34/64 gives 13.7/32.2/48.6 us), which is the signature of a latency chain rather than a bandwidth or FLOP floor. splits=4 reaches the same 15.8 us as splits=8 while halving the workspace_o mirror, so 4 is the default.

End to end

GB300, MiniMax-M3-NVFP4, TP1/PP1/EP1, ISL 8192 / OSL 1024, concurrency 16, num_prompts 32, total_token_throughput:

arm	tok/s	node
base	6130.02	nvl72d027-T06
same build, KV_SPLITS=1	6171.52	nvl72d010-T11
KV_SPLITS=4	6446.05	nvl72d104-T12
KV_SPLITS=4, under nsys	6339.14	nvl72d010-T18

+4.45% against the paired same-build KV_SPLITS=1 arm, and +5.16% against base. The two KV_SPLITS=1 arms agree to 0.68% across different nodes, so the delta is not node drift, and it is roughly 4x the 1.11% same-node repeat spread measured on this workload.

The added mirror refresh copy is 9.18 us/step against a 934 us/step kernel saving, and its cost is flat in width (8x the bytes costs 1.06x the time), so the wider mirrors are not a measurable cost. _build_step_plans() runs once per step from prepare(), not once per layer.

Accuracy

max_score is torch.equal to the num_kv_splits=1 output at every split width in the table above, measured across batch in {1, 2, 4, 8, 15, 16, 32} and kv_len in {256 ... 16384}, including consecutive decode lengths. Downstream top-k block selection is therefore identical and greedy decode output is unchanged.

This is expected structurally, not only observed. In OnlyScore mode out is None, so the split-KV O/LSE reduction is skipped entirely (api.py:919). Each per-k-tile maximum is written by exactly one split, seeded at -INFINITY over that split's own [kv_tile_begin, effective_end) range (mainloop:921-926,1322), so there is no cross-split combination to get wrong. Maximum is associative, commutative and exact in floating point, so partitioning the KV axis cannot perturb the result.

Test Coverage

• tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py covers the MSA backend and the graph-safe plan path; the split-KV mirrors are exercised through it whenever num_kv_splits > 1.
• Parity was measured directly against the num_kv_splits=1 plan over the batch and kv_len grid described above, comparing max_score with torch.equal and the resulting top-k block indices for exact equality.
• End-to-end serving runs at the canonical decode point, reported in the tables above, including a paired same-build A/B via TRTLLM_MINIMAX_M3_MSA_PROXY_KV_SPLITS on a single node.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.