Repository navigation
[ROCm][Perf][MiniMax-M3.1] Four-wave decode top-k for the NVFP4 indexer - #65087
Conversation
The decode top-k walked each row with one wave. Run four waves per row, each over every fourth window with the next window already loading, and merge their sorted top-16 lists by an exact rank over (key, index), which keeps AITER's order including ties, NaN and infinities. Prefill is unchanged. Extend the decode tests with 100 requests and the window edges where the waves meet. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Claude Sonnet 5.5 <noreply@anthropic.com> Signed-off-by: chao-xu-spec <chaox@amd.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: 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. 🚀 |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
/ci run |
|
✅ Triggered Buildkite CI #99595 for commit |
|
/ci run |
|
✅ @chao-xu-spec, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #99665 for commit |
Overview
On gfx950, the MiniMax-M3.1 NVFP4 indexer picks the 16 best blocks of every decode row with one wave per row. This PR gives each decode row four waves and merges their top-16 lists exactly, so the ordered block ids are unchanged. No other kernel, flag or call site changes.
Claims
pa_sparse_block_topkand to the old kernel in 4,500 sweep cases (772,650 rows); eager and CUDA-graph replay agree. In a replay of the real served calls 1 of 2,816 decode rows differs, in a call with duplicate page ids where the old kernel's own output changes from run to run. GSM8K and spec-decode accept length are unchanged.Validation
Microbench. MI355X (gfx950), ROCm 10.0.0, torch 2.13.0+rocm10.0.0,
vllm/vllm-openai-rocm:minimax-m31-flashwith the two top-k files of this PR on top of main (7b77ed5; the files are unchanged on main at 76c6f3f, where the test file was re-run); stock = main's files. CUDA-graph replay of 16 back-to-back calls on distinct score buffers, 1 GiB memset (cold L2) before every replay, median of 20 replays per process, 6 fresh processes per arm; per-call us, median of the 6 (process-to-process range at most 0.07 us ours, 0.21 us stock). 1 index head, 4 rows per request, output equal in every process. Scores are the real output of the unchanged scorer on captured queries and index-K pages of a served run (grid cells: re-scored captured requests clipped to the cell's lengths; live lengths uniform in [max/2, max]). Rows = requests, columns = max sequence length (graph width). Not a repo benchmark (none exists for this op); the method is in the details block.Stock / ours, us (speedup), real scores:
Served shapes (20 decode graph calls of a TP4 AgentX c16 run: real lengths, padded batch, max length 1M, 1-14 requests), weighted by call count: stock 18.91 us (18.91-18.92), ours 4.65 us (4.65-4.66) per call, 4.06x (per process 4.06-4.07x); per call 1.04x to 7.4x. The same shapes with random scores: 5.79 / 4.61 us, 1.255x (table and per-call numbers in the details block).
E2E. amd/MiniMax-M3.1-Flash-Quark-MXFP4 (MXFP4 experts, MXFP8 dense linears), TP4 EP1 on 4x MI355X, nvfp4 KV and indexer cache, DSpark with 3 draft tokens (synthetic acceptance 2.88), DBO, max-model-len 1M,
--max-num-seqs32 (c16) / 64 (c32). CI image plus an overlay of the two changed files (the four files of the decode top-k and its scorer are byte-identical to main's); stock = the image as is. InferenceX AgentX replay (agentx-harness 1.0.12, cc-traces-weka-062126), 3600 s per run, the CI client flags; 2 runs per arm in ABBA order (stock, ours, ours, stock), server restarted per run; ranges are min-max of the 2 runs. Noise floor = the larger of the stock-vs-stock range and the spread of 5 (c16) / 2 (c32) earlier stock runs; a change counts if it exceeds 2x the floor, has the same sign in both adjacent ABBA pairs and the runs do not overlap.Throughput, TPOT and interactivity gains (p50 and p90) are above the floor at both concurrencies (both ours runs above both stock runs). Mean output length is 975 (c16) / 1027 (c32); tok/s per mean output length gives +1.48% / +2.04%. In a kernel trace of the c16 server (profiler on, 12 windows of 64 steps per arm, rank 0) the decode top-k takes 21.2 us per call stock and 5.2 us ours (60 calls per decode step, 1.27 -> 0.31 ms of an 18 ms step); the scorer is unchanged (17.0 us per call at matched context in both). Dispatch: on all 4 ranks the ours runs launch
m3_index_topk_decode_indices_w16and never_w26; the stock runs launch only_w26.Accuracy. GSM8K, 1319 questions, 5-shot, temperature 0, chat completions, max_tokens 12288, concurrency 64; CI conc-64 recipe with real verification (block rejection sampling; the synthetic acceptance of the benchmark recipe degrades text). 2 servers per arm in ABBA order. vLLM's
tests/evals/gsm8k/gsm8k_eval.py(2 passes per server, 4 values per arm) and the CI harness (lm-eval, the CI-pinned version, 1 pass per server).tests/evals/gsm8kThe CI published 0.9666 strict for this recipe. Long-context check (my own script, since GSM8K stays under 131k tokens and never reaches the multi-wave path): a 5-digit key at a given depth in unique filler, 147k-971k prompt tokens, 98 of 98 correct on both arms.
Exactness and tests (kernel level, MI355X):
Reproduce and more microbench numbers
Random scores (N(0,1), a few outliers), stock / ours, us (speedup), same method:
The old kernel does a row of up to 213k tokens in one window and needs only a few dozen candidates per window on random scores, so this table is its best case. Real scores, served shapes (call = requests, padded rows, mean context; weight = share of calls; rows over CAP = rows whose old-kernel candidate stage overflows 64 candidates, emulated):
Details
Each decode row is one workgroup that walks the row's block scores in windows, keeps a running top 16 and ranks the candidates. The old kernel did this with a single wave over windows of 1,664 blocks (213k tokens); rows longer than that serialize on the one wave. Now the workgroup has four waves, each walking every fourth window of 1,024 blocks with its own candidates and running top 16, and the lists are merged by an exact rank over (score, block id), which reproduces AITER's order including ties, NaN and infinities. A wave with no live window pads its list with dead entries, and a row of at most one window skips the merge. The lane-bound threshold uses 16 bits instead of 12 and falls back to all 32 bits when too many keys pass. Prefill compiles to the same code as before. The kernel name for decode changes from
..._w26to..._w16.Why real scores matter: the old kernel keeps the keys above a bound and ranks at most 64 of them per window; when more pass, it rebuilds the candidates, and with many near-ties it runs an exact bit search on one wave. Random scores never reach that. Real index scores sit in a narrow range, so the 12-bit bound lets more than 64 keys through in about a third of the rows (emulated on the captured rows: 33% of rows with the old bound and window size, 6% with the 16-bit bound), and a call takes as long as its slowest row. Calls in which no row overflows run at the random-score time. I did not build a 16-bit-bound-only or four-wave-only variant, so I cannot say how much of the 4x comes from each.
Limits:
max_seq_len(1M in serving), not the row lengths, so the host cannot choose between the two kernels.benchmarks/kernels/.Not a duplicate: no open PR changes the top-k files or their test file (checked against main at 76c6f3f). #64198 (MXFP4 K/V and index cache) does not touch them. #64264 (NVFP4 kernel warmup at startup) and #64273 (monokernel decode on M3.1) are closed drafts that only carried the same bytes of these files as main. #64448 (decode CP for DSpark 7) is open and edits the indexer module and the CP push kernel, not the top-k kernel, so it is complementary.
AI assistance: AI assistance was used in this PR, for the kernel change, the added tests and benchmark scripts, the measurements and this description; the models are listed as
Co-authored-byon the commit. The numbers in this description come from the runs described above.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (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.