[orangezed]

Stacked on #10 (FLARE_CANDIDATE_MERGE_FUSED): the first commit here is #10's. Review only the last commit.

📌 Description

FLARE_CANDIDATE_MERGE_RADIX=1 (default off). It is read when a merge is prepared, so set it identically for export and serving. The DSA candidate merge uses the b12x row top-k (run_row_topk, a 1024-thread single-CTA radix over 10-bit FP16 coarse bins) for only two things: the cutoff (tl.min of the top-K values, when more than K candidates are valid) and the compacted ids when at most K are valid. This PR replaces the row top-k, the threshold kernel and the cutoff-count kernel with one new Triton kernel, _radix_counts:

• It runs one program per row (16 warps, row in registers). It finds the K-th largest order key exactly with 8-bit MSD tl.histogram passes. Keys are FP32 bits mapped to the same unsigned total order the b12x radix uses, and the passes exit early once the K-th key's digit bin holds exactly the keys still needed.
• It takes the cutoff with the same tl.min (Triton minnumf) over that top-K set, with NaN padding neutral. It writes the per-512-tile greater/tie counts and the -1 fill.
• _select is unchanged, so the selected ids and their order are identical to the default. Rows with at most K valid candidates copy their compacted ids, as the row top-k emits them.
• Launches per merge go from 6 to 4: _count, _compact, _radix_counts, _select. It has its own native contract, flare.candidate_merge_radix. Default query, contract and cache keys are unchanged. Merges wider than 16,384 (TP8 x top-2048; 16-bit count fields, row in registers) silently keep the default contract, which is covered by a test.

Relation to #10 (FLARE_CANDIDATE_MERGE_FUSED): #10 folds the threshold kernel into the counts kernel (6 -> 5 launches, about -0.75 us). It keeps the row top-k, which is the expensive part at long context. This PR removes the row top-k as well, so it supersedes #10 when both are set: RADIX takes precedence over FUSED. #10 remains useful on its own, for example above the 16,384 width limit. #10's CUDA test unsets FLARE_CANDIDATE_MERGE_RADIX for isolation; that is a one-line change.

CUDA graph, us per merge (TP8 width 8 x 2048, top-k 2048, RTX PRO 6000 Max-Q, 200 calls per graph, median of 30, 3-run medians):

valid candidates	distribution	1 row: off -> radix	4 rows	8 rows
2,048 (<= K)	any	6.22 -> 4.04	6.58 -> 4.26	6.97 -> 4.61
16,384 (context >= 32K)	normal	20.0 -> 12.3	21.0 -> 12.75	21.8 -> 13.65
16,384	narrow FP16 band	23.8 -> 13.5	24.6 -> 14.0	26.4 -> 14.9
16,384	integer ties	23.8 -> 14.4	25.2 -> 14.9	26.1 -> 15.8

At 4,096 valid candidates the gain is smaller: 10.7 -> 10.5 (normal) to 15.0 -> 13.3 (ties) at 1 row. Profiler split at 1 row, long context: the row top-k is 13.7 / 15-32 / 18-33 us (normal / band / ties), the threshold 1.5 and the counts 1.3. _radix_counts is 8.5 / 9.8 / 10.7 us.

🧪 Tests

On this rebased branch (single RTX PRO 6000 Max-Q):

• Bitwise check, default vs RADIX (and FUSED): 12,188 cases, 0 differences in selected ids and order. Coverage:
  • top-k 2048, rows 1/2/4/8, widths 16384/8192/3000;
  • top-k 1024 and 512 at widths 8192/3000;
  • 12 valid counts each, including K-1..K+2;
  • distinct, integer ties, narrow FP16 band, FP16-bucket ties, all-equal, 90 % -inf, +-0, +NaN, -NaN, 90 % NaN and +-inf, spread and packed slots;
  • scratch arena 0xFF and output -777 before every replay.
• CUDA graph replays: 10,000 x rows 1/4/8 x 4 inputs (ties and narrow band at 8,192 valid, 4,096, 1,500), with arena and output poisoned before every replay and a comparison every replay: 0 mismatches in 120,000 replays.
• GPU tests test_candidate_merge_radix_cuda.py (new: 5 shapes including width 3000 / top-k 1024 / top-k 512, specials, poisoned scratch, CUDA graph, off vs on bitwise), test_candidate_merge_fused_cuda.py and test_candidate_merge_cuda.py were run with FLARE_RUN_CUDA_TESTS=1: 15 passed with the switch off, 15 passed with it on.
• CPU tests (tests/recipes/test_kimi3*, test_kimi31*, tests/components/test_candidate_merge*, test_collective*): 5 F / 148 P on Optionally fold the DSA candidate-merge threshold into the cutoff counts #10 and on this branch, with an identical failure set (pre-existing).

On our pre-rebase tree (release branch + these patches), the same gates also passed, plus 5,280 top-2048 cases at width 16384, and the switch has been in our deployment since then.

Topology assumptions

Generic (single GPU, any TP). The width limit 16,384 matches TP8 x top-2048. Wider merges keep the default path.

Reviewer Notes

• The ordering key and tl.min semantics were chosen to reproduce the row top-k's result bit for bit, including NaN and signed zero. If the row top-k's tie order ever changes, this contract has to follow it; the CUDA test compares off vs on bitwise.