[abdel-darwish-27]

Problem

With CCCL >= 3.2 (GGML_CUDA_CUB_3DOT2=ON), ggml_cuda_op_top_k runs cub::DeviceTopK::MaxPairs in a serial per-row loop. Decode-shaped calls (nrows == 1) are fine, but prompt processing hands this op hundreds of rows at once: profiling one 141K-token prefill of a qwen4exp model (512-row ubatches x 12 QSA layers) showed 903,702 DeviceTopK invocations, each launching 3-4 kernels. DeepSeek V3.2 / V4 indexer top-k hits the same path.

Change

Keep DeviceTopK for nrows <= 4 and dispatch larger batches to the existing row-parallel argsort+copy path, until DeviceSegmentedTopK exists (NVIDIA/cccl#6391 — the existing TODO in this file). The argsort fallback previously only compiled when CCCL < 3.2; it is now reachable in both configurations, so next_power_of_2 moves out of the #elif.

Results

End-to-end prompt processing, Qwen3.8-Flash-Next UD-IQ4_XS on 2x RTX A6000 (sm_86), CUDA 12.1 + CCCL 3.2, master e4b9af00 vs master + this change (cold full prefills; multiple reps; contention-free samples only, foreign GPU activity sampled at 1 Hz):

prompt tokens	per-row DeviceTopK	batched dispatch
129,584	563-580 tok/s	712-713 tok/s (+25%)
61,767	675-698 tok/s	872-874 tok/s (+27%)
30,802	758-781 tok/s	921-923 tok/s (+20%)

At the shapes covered by test-backend-ops perf (nrows <= 16) the two paths are within noise of each other (e.g. ne=[32768,16],k=16: 257.8 vs 259.9 us/run), so the dispatch changes nothing measurable for small batches; the win comes from prompt-processing-scale row counts (hundreds of rows per op), where the serial launch overhead dominates.

test-backend-ops -o TOP_K passes (the suite covers nrows 1-16, which exercises both dispatch branches). Also verified the CCCL < 3.2 configuration still compiles (stock CUDA 12.1).

[ggml-gh-bot]

Hi @abdel-darwish-27, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

• Multiple open PRs from a new contributor: We limit new contributors (those without a previously merged PR) to 1 open PR at a time. You currently have 2 open PRs.

• AI-generated content: While code is allowed to be generated by AI, please write the PR description and commit messages on your own without the help of AI.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.