[CUDA] Optimize TopK for wide last axes

[Tianlei Wu (tianleiwu)]

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

[Copilot]

🟡 Changes recommended

Boundary launch failures, signed-zero instability, and missing dispatch guards must be addressed.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Optimizes CUDA TopK for wide last-axis workloads using Hybrid and streaming small-K implementations.

Changes:

• Adds stable Hybrid TopK for K ≤ 256.
• Adds streaming TopK for K ≤ 32.
• Adds FP32, FP16, and BF16 regression coverage.

File summaries

File	Description
onnxruntime/test/providers/cpu/math/topk_op_test.cc	Adds stability and boundary tests.
onnxruntime/core/providers/cuda/math/topk_impl.cuh	Adds streaming dispatch and kernels.
onnxruntime/core/providers/cuda/math/topk_impl_f32.cu	Enables Hybrid for FP32.
onnxruntime/core/providers/cuda/math/topk_impl_f16.cu	Enables Hybrid for FP16.
onnxruntime/core/providers/cuda/math/topk_impl_bf16.cu	Enables Hybrid for BF16.
onnxruntime/core/providers/cuda/math/topk_hybrid.cuh	Implements partitioned Hybrid TopK.

Review details

Suppressed comments (2)

onnxruntime/core/providers/cuda/math/topk_hybrid.cuh:309

• The Hybrid predicate does not check whether this device supports cooperative kernel launches. Occupancy calculation can succeed on a non-cooperative device, after which cudaLaunchCooperativeKernel returns cudaErrorNotSupported instead of falling back to the existing TopK implementation. Reject such devices before selecting the reduction kernel.

  constexpr int block_size = ReductionBlockSize<K>();
  void* reduction_kernel = ReductionKernel<K, block_size>(factors);
  int blocks_per_sm = 0;
  const cudaError_t result = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
      &blocks_per_sm, reduction_kernel, block_size, 0);

onnxruntime/core/providers/cuda/math/topk_hybrid.cuh:322

• IsSupported accepts every positive last-axis length, so ordinary small TopK calls now pay four scratch allocations and two kernel launches even though this optimization is intended for wide axes; a shape such as [1, 16] previously used one small bitonic kernel. Add a benchmark-derived minimum-dimension gate (the new tests and streaming path use 4096 as the crossover boundary) to avoid regressing common small-axis workloads.

inline bool IsSupported(const CudaKernel* kernel, int64_t rows, int64_t dimension, int64_t k) {
  if (rows <= 0 || rows > kernel->GetDeviceProp().maxGridSize[1] ||
      dimension <= 0 || dimension > std::numeric_limits<int>::max() || k <= 0 || k > kMaxK) {
    return false;

• Files reviewed: 6/6 changed files
• Comments generated: 4
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Baiju Meswani (baijumeswani)]

This PR adds two CUDA fast paths for TopK on wide, contiguous last axes: a streaming SmallK path and a partitioned Hybrid path. It also adds stable tie-breaking, FP32/FP16/BF16 coverage, plugin BF16 numeric limits, and the synchronization needed for shared warp-merge sorting.

The implementation is well structured. Keeping each fast path behind a small IsSupported / Run interface makes the dispatch easy to follow and preserves the existing fallback. The reduction sizing and stable index ordering are carefully handled, and the extra comments around shared-memory aliasing make the synchronization requirements clear.

I found two correctness issues and two smaller quality gaps. The inline comments below include concrete suggestions. I would address the signed-zero and cooperative-launch items before approval.

[Baiju Meswani (baijumeswani)]

Open comment can be addressed as a follow up.