[faresobeid]

Speeds up select_indexed_blocks (Qwen3.8-Flash-Next QSA selection) by ~4x at real shapes. The selection stays the same: the scores are bitwise equal to main's, and so is the set of selected blocks, ties included. The output order is now deterministic (ascending block index).

Problem

• selection_score_kernel ran one query per CTA. The 4 indexer heads were padded to 16 MMA rows, and every query re-read all of its visible keys (~550 GB of key reads at 131K).
• radix_select_kernel wrote picks with atomic_add, so the output order changed from run to run. TileLang also put two __syncthreads() around every histogram atomic, because digit was a shared alloc_var. A single thread scanned the 256 bins serially, and partial ties at the threshold were resolved by a single thread looping over every column.

Change

• Score kernel. One CTA now takes 16 consecutive queries x 4 heads = 64 MMA rows, over the union of their key ranges (absolute key tiles, so documents inside a tile are handled). The heads are combined as (relu(h0)+relu(h1)) + (relu(h2)+relu(h3)), the same order main's reduce_sum used, so the scores are bitwise equal.
• Radix select.
  • Vectorized 16-byte loads. The score rows are padded to a multiple of 4 floats (score_workspace, allocated once per call and reused for every chunk).
  • No barriers inside the histogram loop.
  • Parallel threshold-bin search with T.cumsum.
  • The final pass compacts in column order with a block prefix sum. It takes the greater blocks plus the lowest-index tied blocks, which is the rule main already used. There are no atomics, so the output is sorted and deterministic. The sentinel is still num_blocks.
• API: select_indexed_blocks is unchanged. compute_selection_scores returns the same thing (a view into the padded buffer).

Numbers

H200. Single-document micro-batch. Indexer inputs from prime-rl's IndexedGatedAttention with random weights on locally correlated hidden states, the same setup as #22. "Layer" is one QSA attention module fwd+bwd, no AC. CP rows emulate the last CP rank on one GPU (keys and KV of the other ranks are precomputed; communication is not included).

queries x keys	select main	select this PR	speedup	layer main -> PR	selection share of layer
16K x 16K	2.8 ms	0.7 ms	3.9x	54 -> 52 ms	5% -> 1%
64K x 64K	32.0 ms	7.2 ms	4.4x	326 -> 300 ms	10% -> 2%
131K x 131K	126.8 ms	29.7 ms	4.3x	847 -> 748 ms	15% -> 4%
CP8 last rank: 16K x 131K	27.2 ms	7.1 ms	3.8x	137 -> 115 ms	20% -> 6%
CP8 last rank at 1M: 131K x 1M	2186 ms	516 ms	4.2x	3511 -> 1830 ms	62% -> 28%

Per 1 GiB chunk at CP8 16K x 131K (8192 rows):

	main	this PR
score kernel	9.30 ms	1.74 ms
radix select	4.61 ms	1.85 ms

• On main, the score kernel reads keys at an effective 4.7 TB/s, mostly from L2. At 1M the 64 MB key array no longer fits in L2.
• The radix select now reads the fp32 score table 5 times at about 2.7 TB/s. That table traffic is most of what is left. Removing it would need a fused score+select, which this PR does not do.

Validation

• Against main's kernels on 18 cases: score bits on all visible entries, the selected set per row, the exact output against a sorted torch reference (stable sort, lowest index first among ties), and repeat equality. All equal. The cases cover odd block counts (37, 1021, 1023), up to 5 documents, query counts that are not a multiple of 16, heavy ties (zeroed and duplicated keys), and topk 64 and 512.
• Through prime-rl's IndexedGatedAttention (20K tokens, 2 documents): sorted indexer indices, the attention output, d(hidden) and d(q_proj.weight) are bitwise equal to main (feat(indexed_attention): windowed GQA forward + deterministic key-tile dK/dV (Triton) #22's plan sorts the indices anyway).
• New tests/indexed_attention/test_selection.py (GPU): exact output against the reference, including ties, plus repeat equality. 6 passed on H200.

Notes

• TileLang prints a false-positive data-race warning for the threshold-bin write in T.Parallel (only one bin satisfies the condition) and for the score store.
• Using this in prime-rl needs a wheel release and a pin bump.