From ff2b436762b3c412ed8edfcef401e32629ad4c4c Mon Sep 17 00:00:00 2001 From: Inovello Date: Wed, 9 Sep 2026 07:59:16 -0700 Subject: [PATCH] CUDA: use radix-select TOP_K in the CUB fallback for wide rows When CUB has no DeviceTopK (every CCCL older than 3.2), ggml_cuda_op_top_k sorts every key of every row with a segmented radix sort and copies the first k. The qwen4exp QSA indexer at 131k context spent 5.1 ms per token in that sort. Make the radix-select kernel available to CUB builds without DeviceTopK and use it for rows of >= 8192 columns since on an RTX 3090 (CUDA 12.0, CUB 2.0.1) the segmented sort was still faster at 4096 columns and the radix-select became faster from 8192 onwards. GGML_CUDA_TOPK_ARGSORT=1 will force the old path for an A/B comparison. Using test-backend-ops, TOP_K scored 525/525 on both paths. Standalone at >= 8192 columns resulted in a median 2.25x faster over 46 cases. End to end at 131k context resulted in top-k kernels going from 5.1 to 0.25 ms/token and decode on cached 128-token continuations improving by 13 to 18 percent on every one of six paired runs. Assisted-by: Claude --- ggml/src/ggml-cuda/top-k.cu | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index c7a0c831788d..cf4122b27c68 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -1,6 +1,8 @@ #include "argsort.cuh" #include "top-k.cuh" +#include + #ifdef GGML_CUDA_USE_CUB # include # if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2) @@ -48,7 +50,8 @@ static int next_power_of_2(int x) { #endif // CUB_TOP_K_AVAILABLE -#if !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP) +// radix-select top-k: used by HIP for wide rows and by CUB builds without DeviceTopK (CCCL < 3.2) +#if !defined(CUB_TOP_K_AVAILABLE) && (defined(GGML_CUDA_USE_CUB) || defined(GGML_USE_HIP)) static __device__ __forceinline__ uint32_t top_k_float_to_ordered(float value) { const uint32_t bits = __float_as_uint(value); @@ -208,7 +211,7 @@ static void top_k_radix_cuda( src, dst, states, ncols, k, blocks_per_row); } -#endif // !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP) +#endif // !defined(CUB_TOP_K_AVAILABLE) && (defined(GGML_CUDA_USE_CUB) || defined(GGML_USE_HIP)) void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; @@ -233,6 +236,15 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream); } #elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE + // No DeviceTopK in this CCCL. Sorting every key of every row and copying the first k is expensive for wide rows, + // so use the radix-select top-k for rows of at least 8192 columns. The threshold is conservative: on an RTX 3090 + // the segmented sort was still faster at 4096 columns and the radix-select faster from 8192 columns on. + // GGML_CUDA_TOPK_ARGSORT=1 forces the argsort path (for A/B comparisons within one binary). + static const bool force_argsort = getenv("GGML_CUDA_TOPK_ARGSORT") != nullptr && atoi(getenv("GGML_CUDA_TOPK_ARGSORT")) != 0; + if (!force_argsort && ncols >= 8192) { + top_k_radix_cuda(pool, src0_d, dst_d, ncols, nrows, k, stream); + return; + } // Fall back to argsort + copy const int ncols_pad = next_power_of_2(ncols); const size_t shared_mem = ncols_pad * sizeof(int);