[mdavis36]

Summary

This replaces the histogram TopK behind the decode sparse_topk_select_publish with a top-16 that ranks each (token, head) row in one warp, in the candidate merge's order. It is one commit on f380d0e, the head of vllm-project#46 when this was opened, and targets that PR's branch, so merging it adds the commit to that PR.

In indexer context parallelism (ICP), each rank selects a local top-16 per (token, head) row and publishes it to the row's owner rank, which merges the candidates in vLLM. On f380d0e, publish takes longer than the merge. The histogram kernel runs one 512-thread CTA per row, but a decode row has one local column per page of context (1,344 at 168K tokens with P128 pages) and keeps 16. Its threads hold only a few columns each, while every row still pays for block-wide barriers, a block scan and merge sorts. With only two or three of those CTAs per SM, a decode step of thousands of rows runs in waves.

The API, row format, tagged/legacy protocol, PDL, index_scale/index_offset and which pages are forced (force_begin/force_end) are unchanged, so callers need no changes. Published rows differ from f380d0e only in the tie and special-value cases below. Prefill (sparse_topk_select_prefill) keeps the histogram kernel.

Changes

SparseTopKPublishWarpKernel ranks one row per warp, 4 rows per 128-thread CTA. Each column gets a 64-bit key: its order-preserving score bits above its inverted column index. Keys are distinct and sort in the candidate merge's order: score descending, forced pages as +inf, NaN as -inf, -0 as +0, then the lower column (so the lower page id). As ec980f2 does for prefill, ranking in the merge's order lets independently selected rank candidates compose exactly: each rank's top-16 holds every candidate the merge would take from it.

• Threshold: each lane takes its largest key. Every key at or above the 16th largest of the 32 lane maxima survives. At least 16 keys do, so the top 16 are among them; for random scores it is about 20 per row. Survivors take shared-memory slots with atomics, but each is ranked by counting the survivors above it, so the result does not depend on slot order.
• Fallback: if more than 64 keys survive, the warp falls back to 16 exact rounds over the row.
• Loads: each lane holds 8, 16 or 24 columns in registers, picked from the score buffer's width, so up to 768 columns of a row load while its valid-page count is still in flight. Columns past 768 are re-read in 768-column chunks on every pass over the row: twice normally (lane maxima, then survivors), and once per round in the fallback.
• Publish: the receive slot's address and, in tagged mode, the row's previous generation are loaded up front, so their latency hides behind the selection. The legacy ready-flag wait still runs just before the write.
• Left in place: the histogram kernel's decode-only paths (its 512-thread register cache, kCached) are now unused. Removing them is left for a follow-up to keep this diff small.

Behaviour changes vs f380d0e

• Exact ties at the cutoff, including -0 vs +0, go to the lower id. On f380d0e, rows of up to 128 columns give them to the higher id, and longer rows go by shared-memory atomic order and rank +0 above -0.
• Forced pages rank as +inf instead of FLT_MAX, as in the merge. A real +inf now ties with them (lower id first) instead of outranking them, and a real FLT_MAX ranks below them instead of tying.
• NaN ranks as -inf. On f380d0e, a NaN from arithmetic ranked above every score, forced pages included.

Not changed: prefill

sparse_topk_select_prefill keeps the histogram kernel. Its docstring promises cutoff ties to the lower id, matching the merge, and since ec980f2 its final candidates are ranked with a stable key on rows of any length. It still departs from the merge's order in three cases:

• Rows of more than 128 columns are binned by the fp16 histogram first, which puts +0 above -0.
• If more than 128 candidates share the cutoff's top 30 bits, exact ties go by shared-memory atomic order.
• On every row, forced pages score FLT_MAX, and a NaN from arithmetic ranks first.

Performance

No timings are included here, and GB200 and GB300 have not been timed. Publish kernel time (nsys) and the decode chain of score, publish and merge (CUDA events) were each measured back to back against f380d0e:

• Shapes: TP2 and TP4, P128 and P256, 16 and 64 requests of 8 query tokens, 168K tokens of context. Rows have 1,344 columns at P128 and 672 at P256.
• Data: random FP8 queries and random NVFP4 index pages in M3.1's cache layout, not model activations.
• Chain: [Attention] Shard MiniMax M3.1's NVFP4 index cache with indexer CP vllm-project/vllm#64631's NVFP4 shard scorer, then the publish kernel under test, then the current merge.
• Not modelled: every simulated rank runs on one GPU, so the cross-GPU exchange and rank skew are left out.

Publish kernel time drops in every configuration, by more at 64 requests than at 16, and the chain is faster by about the time saved in publish. A separate run at 512K tokens with P256 (2,048-column rows), which also used the warp merge under Related at TP2, saved about as much per layer. Longer rows have not been timed: 4,096 columns at 512K tokens with P128 or at 1M with P256, and 8,192 at 1M with P128. Since columns past 768 are re-read on every pass, those are the shapes to check.

Testing

tests/smoke/test_sparse_topk_publish.py is new: 552 cases that check published rows exactly against a reference in tagged and legacy modes. Score buffers of 17 to 5,000 columns cover the 8-, 16- and 24-column register paths and rows streamed past 768 columns, with valid counts from 0 to past the buffer. Cases cover cutoff ties, ±0, ±inf/NaN, scale/offset, forced pages, the fallback, slot reuse and PDL. Buffers wider than 5,000 columns (the op accepts up to 12,287) are not tested.

pytest tests/smoke/test_sparse_topk_publish.py \
       tests/smoke/test_sparse_topk_prefill.py    # 634 passed on GB200 and on GB300

• SM107: the same two files plus test_sparse_topk_forced.py, test_sparse_topk_large_boundary_bin.py and test_index_score_sharded.py: 725 passed.
• vLLM, at mdavis36/vllm@b60e562982 ([Attention] Shard MiniMax M3.1's NVFP4 index cache with indexer CP vllm-project/vllm#64631 plus the warp merge under Related), with vllm/third_party/fmha_sm100 replaced by this branch's python/fmha_sm100, since the MSA that vLLM pins predates sparse_topk_select_publish:
  • pytest tests/kernels/attention/test_minimax_m3.py -k 'warp_matches_legacy or candidate_merge or select_publish or sharded or q8kv4' gives 370 passed and 8 failed on GB200, GB300 and SM107.
  • tests/distributed/test_minimax_m3_indexer_sharding.py passes at TP2 and TP4 on SM107.
  • The 8 failures are the eight parametrizations of test_sharded_candidate_merge_deduplicates_and_forces ("indices should be either on cpu or on the same device"). It is a merge test that never calls the publish kernel, and it fails the same way without this change: since wolfcomos/vllm@dc68fa728, it indexes a CPU scores tensor with the CUDA nonempty mask.

Alternatives considered

• The histogram kernel with 128-thread CTAs: faster than with 512 in event timing, but well short of the warp kernel.
• Up to 1,536 columns of a row in registers (48 per lane) instead of re-reading past 768: slower at 64 requests.

Related

A one-warp-per-row Triton candidate merge for vLLM is on mdavis36/vllm:md/icp-merge-warp, on top of vllm-project/vllm#64631 and independent of this PR. Its auto setting uses it everywhere except TP4 decode, where the current merge is faster.