[CUDA] Optimize TopK for wide last axes (#32404)
### Description
Optimize CUDA TopK for wide contiguous last axes. This adds a streaming
small-K implementation and a stable Hybrid implementation ported from
ONNX Runtime GenAI for float, float16, and bfloat16.
Each fast path lives in its own header and exposes the same
`IsSupported` / `Run` pair, so `topk_impl.cuh` only holds the dispatch
order and the pre-existing fallbacks:
| Header | Path |
|---|---|
| `math/topk_hybrid.cuh` | Partitioned + cooperative-reduction Hybrid
TopK |
| `math/topk_smallk.cuh` | Streaming small-K TopK |
| `cu_inc/topk_warp_sort.cuh` | Shared warp-level sort/pack primitives |
The Hybrid path handles sorted largest-element selection with K up to
256. It partitions each row, performs stable packed-key block radix
sorts, and cooperatively reduces partition candidates. It dynamically
selects partition sizes, padded K specializations, and one- to
three-step reduction plans.
The streaming small-K path handles K up to 32 and remains as a fallback
for smallest-element selection, dimensions beyond Hybrid's 256-partition
limit, and grids that cannot satisfy cooperative-launch residency. Other
unsupported configurations continue through the existing CUDA TopK
implementations.
Stable ties are ordered by ascending source index. The regression test
covers padded-K boundaries from 4 through 256 and one-, two-, and
three-step reductions.
### Motivation and Context
The existing CUDA TopK implementation assigns one block to each row and
scans a wide axis multiple times. This underutilizes the GPU for LLM
sampling workloads with few rows and large vocabularies.
On an NVIDIA H200 with CUDA 13, for one FP32 row with dimension 248,320
and K=16:
| Implementation | Mean device time |
|---|---:|
| Existing ORT CUDA TopK | 1.36 ms |
| Hybrid TopK | 25.33 us |
The Hybrid measurement contains 100 calls captured with Nsight Systems:
partition selection averaged 16.13 us, cooperative reduction 8.02 us,
and output writing 1.18 us.
### Dispatch Policy
The Hybrid path is selected only for dimensions at least 8,192. The
minimum K rises with the row count: K=8 for one or two rows, K=32 for
three or four rows, and K=64 above four rows. Smaller workloads remain
on the existing select-based paths.
This conservative policy comes from H200 Nsight Systems sweeps across K
values 1, 2, 4, 8, 16, 32, and 64. The main matrix covered 280 FP32/FP16
cells over rows 1, 2, 4, 8, and 16 and dimensions 8,192 through 248,320.
BF16 boundary checks and rows 3, 5, 6, and 7 brought the selected-policy
validation to 164 cells. Every selected cell was more than 2% faster
than the existing path; the worst ratio was 0.979x. Crossover points may
differ on other GPU architectures.
### Additional Fixes
- The small-K partial pass maps rows onto `grid.y`, so `IsSupported` now
rejects row counts beyond `maxGridSize[1]` and falls back to RadixTopK
instead of failing the launch.
- `WarpMergeSorter`'s CUB temp storage aliases the caller's shared
score/index buffers in `hybrid_topk::ReducePartitions` (they are union
members). The shared-memory loads and the write-back are now fenced
against CUB's own traffic.
- The plugin build maps `BFloat16` onto `nv_bfloat16`, which had no
`NumericLimits` specialization and silently fell back to the
`std::numeric_limits` primary template (0). Kernels that pad with
`Lowest()` therefore ranked the padding above every negative input;
`TopKOperator.NthElementBFloat16_NegativeVals` failed in the CUDA plugin
configuration before this fix.
### Testing
- `onnxruntime_provider_test --gtest_filter=TopKOperator.*`
- monolithic CUDA build: 70/70 passed
- CUDA EP plugin build (`onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON`): 70/70
passed
- `onnxruntime_provider_test
--gtest_filter=*BFloat16*:*BF16*:*Bfloat16*` on the plugin build: 58/58
passed, covering the shared `NumericLimits` change
- CUDA TopK correctness matrix: 24/24 cases passed across FP32/FP16,
largest/smallest, varied rows, dimensions, and K
- Stable large-K matrix: 7/7 previously failing block-merge cases passed
- Signed-zero stability: mixed `-0.0`/`+0.0` cases passed through Hybrid
and SmallK
- CUDA Graph capture and replay: 4/4 runs passed with preallocated I/O
- `lintrunner` and `git diff --check`
