perf: optimize prefill q8kv8 indexer for rubin by caitlinw-nvidia · Pull Request #17 · MiniMax-AI/MSA

caitlinw-nvidia · 2026-10-01T23:41:42Z

Perf: optimize prefill q8kv8 indexer for rubin
Adds Rubin (SM107) support to the Q8KV8 paged prefill indexer.

Changes

indexer_gemm.py

Use the SM107 FP8 MMA instruction (SM107MmaFP8Op, K=64) on Rubin: two MMAs per 128-token page instead of four. The op and the instruction depth are selected by (compute_capability, io_dtype) through uses_sm107_fp8_mma() / mma_inst_k(), so SM100/SM103 keep the generic F8F6F4 op at K=32 and an FP6/FP4 variant of the operator would not need to touch the selection.
Enable the existing tensor-memory load-and-reduce epilogue (LdRed32x32bOp MAX, already used on SM103) on SM107, replacing the fmax tree / enable use_tmem_load_reduce.
Register budgets: reallocated every warp to have 168 registers each for better performance.
K loader: prefetch the next page-table entry while the current page's TMA is issued, taking the dependent address load off the issue path.
Two Q stages on SM107 so the next task's Q tile lands while the current task's pages are consumed; the MMA reads its A operand from the current Q consumer stage.

plan.py, interface.py

Instead of 64 or 128 page chunks, dynamic choose page chunks between [8, 128] so a batch produces about one task per cluster. With the fixed rule a batch with few Q tiles made fewer tasks than clusters (one tile with 433 pages ran as 7 tasks on 7 of 106 clusters); now such tiles are split finer until the machine is full. Batches with enough tiles to fill the clusters already keep the old 128-page chunk.

Performance (Rubin / SM107)

Metric: the repository's formal benchmark, benchmarks/.../indexer/prefill/q8kv8/benchmark.py --suite full --verify: 128 production cases, CUDA Graph, cold cache, 5 warmup + 20 timed replays, median per case, production-weighted mean (66 weighted cases, 14,805 represented calls). Baseline = nv_dev (dbd31ee) compiled directly for SM107 and run on same Rubin GPU.

Performance (Rubin / SM107)

	Weighted mean E2E
`nv_dev`	54.30 us
this PR	49.97 us (-8.0%)

Query length	Cases	Weight	Median per-case change
short (<512 tokens)	14	3998	-20.7%
medium (512-4095)	21	3440	-8.4%
long (>=4096)	93	7367	-5.6%

Change (incremental, Rubin)	Weighted E2E gain
K=64 FP8 MMA instruction	-0.6%
TMEM load-and-reduce epilogue	-2.2%
Planner page-chunk sizing	-4.6%
Loader/MMA register budget 32 -> 168	-2.35%
Page-table prefetch	-1.6%
Two Q stages	-0.4%

Correctness (Rubin)

Suite	Result
`tests/.../indexer/prefill/q8kv8`: api, manifest, production, runtime, plan-reuse, CUDA Graph (H=1/2/4) + `tests/.../indexer/_common/topk_select`	80 passed
`test_real_cases.py`, full tier, H=1/2/4	192 passed
`test_all_real_cases.py` (all 10,562 prefill shapes)	512 passed
Benchmark `--verify`, full suite H=1 (two runs)	256/256 cases passed

For all 128 production cases, the outputs were also bitwise-identical to original baseline.

^{Need help on this PR? Tag @codesmith-bot with what you need. Autofix is disabled.}

caitlinw-nvidia · 2026-10-01T23:42:39Z

cc @hlu1 @wangshangsam

NVIDIA-JerryChen · 2026-10-05T15:13:34Z

lgtm

…x-AI#17) SM107 only; SM100/SM103 code paths unchanged: - SM107MmaFP8Op (K=64) for the FP8 UMMA - TMEM load-and-reduce epilogue on SM107 - 168 registers for every warp role - K loader prefetches the next page-table entry - two Q stages - planner caps the page chunk (8..128) to fill the clusters Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Caitlin Wang <caitlinw@nvidia.com>

…x-AI#17) SM107-specific changes: - SM107MmaFP8Op (K=64) for the FP8 UMMA - TMEM load-and-reduce epilogue on SM107 - 168 registers for every warp role - K loader prefetches the next page-table entry - two Q stages - planner caps the page chunk (8..128) to fill the clusters Changes on all architectures: - define the physical page index before the dynamic K-load loop on the non-SM107 path (port of MiniMax-AI/MSA bc5f83a) - keep the four persistent task-bucket loops rolled (cutlass.range(..., unroll=1)) instead of unrolling them at trace time, which mainly shrinks the generated code Rubin (SM107), full 128-case prefill-indexer suite, score + top-k per call, weighted mean versus dev e441976: - H=1: 54.71 -> 49.34 us (about 1.11x); short queries (<512) about -22% - H=2: 88.04 -> 82.15 us (about 1.07x) Top-k output is bitwise identical in all 128 cases. The prefill and top-k tests in test_q8_indexer.py pass on SM107; on SM100 (B200) they passed before the loop-rolling change. Suggested-by: Fan Shi <fans@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Caitlin Wang <caitlinw@nvidia.com>

* [Perf][Attention] Run every sparse-prefill combine on the v1 combine Fold the nv_dev v1 combine (src/blackwell_prefill/combine.py) into src/sm100/fwd/combine.py and make it the public combine() for all callers: the Blackwell port route, the original FP8/BF16 path, the NVFP4 KV path and the Q8KV4 prefill. src/sm100/fwd/combine.py now has two sections, one per kernel: - Original combine: SparseAttentionForwardCombineF32F8 and its launcher combine_f32f8, unchanged apart from the names. It implements the FP32 and FP8 E4M3 STG.128 fake layouts as well as BF16/FP16. - v1 combine: SparseAttentionForwardCombine (moved verbatim) and its launcher combine(), the default. It takes BF16/FP16 partials only: the v1 kernel implements just that STG.128 fake layout. interface.py picks the launcher by partial dtype (_select_combine) on the two paths that accept a caller-chosen partial_dtype, the original FP8/BF16 path and the NVFP4 KV path. The Blackwell port route and Q8KV4 always write BF16 partials and call combine() directly. vLLM always writes BF16 partials, so production runs the v1 kernel. Other changes: - v1 launch: min_blocks_per_mp=0 (was 3 with output_scale + PDL), and a 3-stage O_partial ring on SM107 (2 elsewhere). - The Blackwell port route now accepts output_scale and passes it to the combine instead of falling back to the original path. - Delete src/blackwell_prefill/combine.py; drop the per-caller combine selection in interface.py and prefill_q8kv4/interface.py. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * perf(indexer): port Rubin Q8KV8 prefill indexer optimizations (MiniMax-AI#17) SM107-specific changes: - SM107MmaFP8Op (K=64) for the FP8 UMMA - TMEM load-and-reduce epilogue on SM107 - 168 registers for every warp role - K loader prefetches the next page-table entry - two Q stages - planner caps the page chunk (8..128) to fill the clusters Changes on all architectures: - define the physical page index before the dynamic K-load loop on the non-SM107 path (port of MiniMax-AI/MSA bc5f83a) - keep the four persistent task-bucket loops rolled (cutlass.range(..., unroll=1)) instead of unrolling them at trace time, which mainly shrinks the generated code Rubin (SM107), full 128-case prefill-indexer suite, score + top-k per call, weighted mean versus dev e441976: - H=1: 54.71 -> 49.34 us (about 1.11x); short queries (<512) about -22% - H=2: 88.04 -> 82.15 us (about 1.07x) Top-k output is bitwise identical in all 128 cases. The prefill and top-k tests in test_q8_indexer.py pass on SM107; on SM100 (B200) they passed before the loop-rolling change. Suggested-by: Fan Shi <fans@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Caitlin Wang <caitlinw@nvidia.com> * feat(indexer): let Q8KV8 prefill callers own the score buffer BatchPrefillIndexerQ8KV8Wrapper always allocated its own score buffer and its only public call ran the TopK as well. A caller that selects pages itself (vLLM's shared TopK over decode and prefill rows) needs the scores written into its own buffer, without the wrapper's TopK. - plan(..., score_out=None): optional caller-owned fp32 buffer of shape [total_q * num_heads, >= max_pages] with contiguous pages per row; a wider row pitch is accepted, so a slice of a larger buffer works. The kernel layout is already dynamic, so no recompile and no kernel change. - run_scores(q, k_cache): public scores-only call; run() is unchanged. Tests: score_out matches the wrapper-owned scores bitwise at an 8192-page pitch, leaves other entries untouched, runs no TopK, and run() TopK is unchanged; invalid score_out buffers are rejected. Prefill and TopK tests in test_q8_indexer.py pass on SM107 (41 passed). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Caitlin Wang <caitlinw@nvidia.com> --------- Signed-off-by: Caitlin Wang <caitlinw@nvidia.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

perf: optimize prefill q8kv8 indexer for rubin

e1639ac

NVIDIA-JerryChen merged commit 9ae7751 into MiniMax-AI:nv_dev Oct 5, 2026
1 check passed

caitlinw-nvidia mentioned this pull request Oct 12, 2026

[Perf][Prefill] Combine opt for rubin + prefill q8kv8 indexer opt vllm-project/MSA#36

Merged

Navigation Menu

Repository navigation

perf: optimize prefill q8kv8 indexer for rubin - #17

perf: optimize prefill q8kv8 indexer for rubin#17
NVIDIA-JerryChen merged 1 commit into
MiniMax-AI:nv_devfrom
caitlinw-nvidia:rubin-fp8-prefill-indexer

caitlinw-nvidia commented Oct 1, 2026 •

edited by blacksmith-sh Bot

Loading

Uh oh!

caitlinw-nvidia commented Oct 1, 2026

Uh oh!

NVIDIA-JerryChen commented Oct 5, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

caitlinw-nvidia commented Oct 1, 2026 • edited by blacksmith-sh Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Performance (Rubin / SM107)

Correctness (Rubin)

Uh oh!

caitlinw-nvidia commented Oct 1, 2026

Uh oh!

NVIDIA-JerryChen commented Oct 5, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

caitlinw-nvidia commented Oct 1, 2026 •

edited by blacksmith-sh Bot

Loading