[dsv4.1] Exact bf16 consumer top-k adapted from DeepSelect

[DarkSharpness]

Motivation

The DeepSeek-V4.1 two-level sparse indexer's consumer step picks the top-512 of up to 16384 bf16 sparse logits per row and writes them through the per-row physical block table (topk_transform_bf16_small). The kernel behind it was a placeholder: a single 13-bit fp16-derived histogram (approximate at the selection boundary), 1024 threads x 16 elements, and untuned.

Modifications

Replace it with an exact selection adapted from DeepSelect. DeepSelect's kernel starts by selecting the top-k of a 32 KiB "init window" that lives entirely in registers; our consumer rows (<= 16384 bf16) are exactly that window, so our kernel agent took that part of the algorithm, dropped the TMA/streaming machinery around it, rewrote it in plain SIMT and specialised it for 16384-wide rows at k = 512. The result was then reviewed and wired into the existing host API here.

What is kept from DeepSelect:

• two radix passes on the raw bf16 bytes: a 256-bin histogram of the high byte, then of the low byte among the elements sharing the pivot's high byte, which locates the k-th largest value exactly;
• the "distorted" value order (x ^ (x < 0 ? 0xFFFF : 0x8000): negatives reversed and below positives, -0 below +0), applied once per lane while walking the histogram instead of once per element;
• the packed hit test: __hgt2_mask / __heq2_mask on bf16x2, prmt + dp4a to turn the flags into one bit per element, one warp prefix plus one shared atomic per warp for the block-wide (gt, eq) prefix, and an equal quota so exactly k elements are taken.

What is specific to this port:

• one CTA of 512 threads per row, each thread keeping a contiguous slice of up to 32 scores (4 x 16 B vectors) in registers for the whole kernel; 40 registers per thread, 3 CTAs per SM, no spills;
• the negative half of the high-byte histogram is offset by 16 words so that +x and -x with the same exponent (the common case for centred scores) do not hit the same smem bank on every atomic;
• the partial last vector of a row is padded with a negative NaN, which sits below -inf in the distorted order and fails every ordered compare, so it is never counted nor selected;
• a -0 pivot is handled separately (the float census sees +0 == -0 while the histogram ranked them apart);
• the selected indices are staged in shared memory (a union with the histograms, 9 KiB in total) and written with one coalesced page-transformed copy.

The host API is unchanged: per-row seq_lens (prefetched ahead of the PDL wait, like the streaming top-k), the fused page-table transform, -1 past min(k, seq_len). k is now at most 2048 (was 4096). NaN scores are never selected (no ordered compare matches them): a row with n positive NaNs yields its top (k - n) real scores and -1 in the remaining slots, derived from the census totals so nothing reads stale shared memory; a negative NaN orders below -inf and is simply never picked.

Accuracy

test/registered/kernels/ops/attention/test_topk_bf16.py (100 cases): the existing configurations plus odd row widths (partial last vector), heavy ties / constant / all -inf rows, signed zeros around the pivot, denormals and full-range bf16 bit patterns, and rows with positive / negative NaNs below and above k. Every case checks the selected values as a multiset against torch.topk (elements of equal value may swap) and the page transform through the inverse table. test/registered/attention/unittests/dsv4/test_dsv41_sparse_indexer.py passes.

Performance

B200, marker.do_bench (CUDA graph, L2-rotated inputs), k = 512, page transform included, random bf16 scores, permuted block table. Old kernel -> this PR, in us:

valid len	rows = 1	rows = 32	rows = 128	rows = 512
16384	5.24 -> 4.93	5.51 -> 5.18	7.09 -> 6.58	15.94 -> 12.38
9000	4.36 -> 4.06	4.71 -> 4.36	5.08 -> 5.07	13.03 -> 9.03
4096	3.97 -> 3.49	4.32 -> 3.81	4.74 -> 4.32	10.83 -> 7.36

Beyond ~450 rows the GPU is saturated (3 CTAs/SM) and time is linear in rows: 36 us at 2048 rows and 134 us at 8192 rows for full-width rows.

k = 2048 stays supported but is not the target shape: it is 10-15% slower than the old kernel at small batch (the epilogue runs on 512 threads instead of 1024) and 10% faster at 512 rows.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings as needed.
• Provide throughput / latency benchmark results.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #34764516423
Latest PR Test (Extra): ❌ Run #34764516294
Latest PR Test (AMD ROCm 10): ❌ Run #34764516371