[caitlinw-nvidia]

Overview

This PR makes two changes to the Q8KV8 paged prefill indexer (python/fmha_sm100/cute/):

1. Rubin (SM107) optimizations for Q8KV8PrefillIndexerSm100, ported from perf: optimize prefill q8kv8 indexer for rubin MiniMax-AI/MSA#17 (nv_dev). It includes the follow-up bc5f83a from nv_dev and unroll = 1 optimization.
2. A caller-owned score buffer for BatchPrefillIndexerQ8KV8Wrapper: plan(..., score_out=...) and a public run_scores(). [MiniMax-M3 Performance] Prefer CuTe for MiniMax M3 FP8 indexer prefill vllm#60930 uses these to run this scorer inside vLLM's shared-score-buffer / native-top-k indexer.

Changes

Kernel and planner (q8kv8_indexer_prefill.py, q8kv8_indexer_prefill_plan.py, q8_indexer_interface.py). SM107 changes only.

• SM107MmaFP8Op (K=64) for the FP8 UMMA instead of the generic F8F6F4 op (K=32).
• Two Q stages, so the next task's Q tile loads while the current pages are still being consumed.
• 168 registers for every warp role, instead of 192/32 for score warps vs. the rest.
• The K loader prefetches the next page-table entry while the current TMA is issued.
• TMEM load-and-reduce for the per-page max.
• Adaptive page chunk: the planner caps the page chunk (8–128 pages) so that short-query batches fill the clusters in one wave. This fixes the short-query, moderate-context shapes (e.g. Q=128 at 32K context) where a fixed 64-page chunk produced only a few tasks.

On all architectures:

• The four persistent task-bucket loops run as runtime loops (cutlass.range(..., unroll=1)) instead of being unrolled at trace time, which shrinks the generated code.
• physical_page is defined before the dynamic K-load loop on the non-SM107 path (port of bc5f83a) (to avoid hanging).

Wrapper API (q8_indexer_interface.py, README):

• plan(..., score_out=None) takes an 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 there's no recompile and no kernel change. The buffer is validated on the host: dtype, device, shape, strides, alignment.
• run_scores(q, k_cache) is a public scores-only call (no top-k). run() is unchanged.

Validation

Correctness:

• test_q8_indexer.py -k "prefill or topk_select", with fresh caches and a GPU idle check each run:
  • SM107 (VR200): 41 passed on this PR's head, including 4 new score_out tests.
  • SM100 (B200): 37 passed on the port before the loop-rolling change. Not yet rerun on this PR's head.
• New tests check that:
  • score_out at an 8,192-page pitch is bitwise identical to the wrapper-owned scores and leaves other entries untouched;
  • run_scores runs no top-k, and run()'s top-k is unchanged;
  • invalid buffers are rejected.

Kernel performance (Rubin SM107; 128-case prefill-indexer suite; score + top-k per call; CUDA-graph replay, 5 warmup + 20 timed; versus dev e441976):

	H=1 (attention TP4)	H=2 (attention TP2)
Weighted mean, port vs. dev	54.71 → 49.34 µs (1.109×)	88.04 → 82.15 µs (1.072×)
Short queries (< 512), median	−22.2%	−14.3%
Medium (512–4095), median	−6.8%	−5.1%
Long (≥ 4096), median	−6.3%	−5.3%
Loop rolling, additional	1.016×	1.006×
Top-k vs. dev	bitwise identical, 128/128	bitwise identical, 128/128

End to end: vLLM nightly 81198e9 with vllm-project/vllm#60684 and #60930 overlaid; MiniMax-M3-MXFP8, TP4, Rubin; --kv-cache-dtype fp8, indexer_kv_dtype: "fp8", minimax_m3_indexer_prefill_backend: "cute". Both arms logged prefill scorer: CuTe.

Workload	After vs. before
Prefill only: ISL 8192, OSL 1, concurrency 32, seeds 45–47, prefix caching off	+0.94% input tok/s (96,542 vs 95,644; per seed +0.89 to +0.97%), −0.9% TTFT

The end-to-end gain is small because fresh 8K prompts are the kernel's weakest case (long queries, short context), and the indexer is a small share of each prefill step. Short-chunk, long-context traffic (chunked prefill, prefix-cached multi-turn) is where the kernel gains concentrate.

Accuracy (NeMo-Skills, same configuration as the end-to-end runs; temperature 0, top-p 1, 32,768-token generation cap):

	Before	After	Paired exact test
GPQA Diamond	178/198 (89.90%)	173/198 (87.37%)	p = 0.33
MMLU	91.05%	90.78%	p = 0.14

No significant difference. The kernel produces bitwise-identical top-k to dev, so the differences are consistent with temperature-0 serving not being bit-reproducible.