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);