[orangezed]

📌 Description

FLARE_CANDIDATE_MERGE_FUSED=1 (default off) is read when a merge is prepared, and must be set identically for export and serving. It replaces the _threshold launch plus the _counts launch of the DSA candidate merge with one _threshold_counts launch. Every (row, tile) program computes the cutoff (the minimum of the K largest scores) itself; tile 0 publishes it and writes the -1 fill that _select overwrites. That is 6 -> 5 launches per merge, over 24 DSA layers. The fused variant is a separate native contract (flare.candidate_merge_fused), so the default query, compile payload and cache keys are unchanged.

CUDA graph us/merge (TP8 width 16384, RTX PRO 6000, 3 runs):

• 1500 valid candidates: 6.30 -> 5.59 (1 row), 6.62 -> 5.86 (4 rows), 6.95 -> 6.21 (8 rows).
• 16384 valid: 20.00 -> 19.18, 21.49 -> 20.70, 21.99 -> 21.24.

Exactness: integer-identical output (selected IDs and their order). This was checked on 800 random cases: rows 1..16, width 8192/16384, valid 0..width, distinct and tied scores, poisoned scratch. It was also checked over 10,000 graph replays per case at rows 1/4/8, with 0 differences.

🧪 Tests

• New tests/components/test_candidate_merge_fused_cuda.py (3 tests, GPU). On CPU it was collected and skipped (actually run). Existing candidate-merge tests: 6 passed off and on, GPU, pre-rebase tree. The CPU failure set is identical to the release branch.

Topology assumptions

Generic (single GPU).

Reviewer Notes

• A single-program finalize (threshold + counts + select) was also exact, but 12 us slower when more than 2048 candidates are valid, so it is not used.

[orangezed]

GPU validation of this branch as published (base release/k31-exl3-tp8-20261108 1d4007f, with b12x from FlashInfer release/k31-exl3-tp8-20261108), one RTX PRO 6000 Max-Q. No changes were needed.

• Validator (outside the repo), OFF = threshold + counts + select, ON = fused finalize, through compile_forward + CUDA graph, scratch arena poisoned (0xFF) before every call:
  • 200 random cases (rows 1/2/4/8/16, width 16384, valid 0..width, distinct and tied scores, 2 seeds): output arrays integer-identical, 0 bad cases.
  • CUDA-graph replays with the scratch poisoned: 2,000 replays each at rows 1/4/8, for both full (16384 valid, with ties) and short (1500 valid) inputs. 0 mismatches against OFF.
• FLARE_RUN_CUDA_TESTS=1 pytest tests/components/test_candidate_merge_fused_cuda.py tests/components/test_candidate_merge_cuda.py: 9 passed with FLARE_CANDIDATE_MERGE_FUSED unset, and 9 passed with it set to 1.
• Timing, CUDA graph, us per merge (200 merges per graph), TP8 width 16384, off -> on:

rows	1500 valid	16384 valid	description (1500 / 16384)
1	6.07 -> 5.40	18.75 -> 17.93	6.30 -> 5.59 / 20.00 -> 19.18
4	6.20 -> 5.54	19.77 -> 19.09	6.62 -> 5.86 / 21.49 -> 20.70
8	6.70 -> 5.98	20.73 -> 20.03	6.95 -> 6.21 / 21.99 -> 21.24
16	7.24 -> 6.49	21.58 -> 20.92	(not quoted)

The saving is the same as in the description: about 0.7 us per merge. The absolute numbers are about 0.2-1.3 us lower in this run.