wmantly/llama.cpp-xrip
1
0 Fork 0
Code Issues Pull Requests Actions 29 Packages Projects Releases Wiki Activity

CUDA: use radix TOP_K when CUB DeviceTopK is unavailable

Browse Source 
CUDA builds with CCCL < 3.2 have no cub::DeviceTopK, so ggml_top_k falls back to a full argsort + copy (~12 bytes of temp per cell). For sparse-attention top-k over long contexts this OOMs small cards (200K cols x 128 rows ~= 190MB+ of temp per layer). Reuse the exact radix selection from #27466 (merged for HIP) on CUDA when CUB_TOP_K_AVAILABLE is off: O(nrows) temp, same selected set. Also gate next_power_of_2, unused on that path now.
This commit is contained in:
Rhonstin authored and xrip committed 2026-09-21 20:54:01 +03:00
1 parent 23cbecae91
commit f56cf8ff73
1 file changed
+19 -2
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/top-k.cu
+19 -2
View File
@@ -10,6 +10,12 @@ using namespace cub;
# endif // CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2
#endif // GGML_CUDA_USE_CUB
// Exact radix top-k selection when CUB DeviceTopK is unavailable (CCCL < 3.2):
// O(nrows) temp instead of the O(ncols*nrows) full-argsort fallback. Selects the same set.
#if !defined(CUB_TOP_K_AVAILABLE)
#define GGML_CUDA_TOP_K_RADIX
#endif
#ifdef CUB_TOP_K_AVAILABLE
static void top_k_cub(ggml_cuda_pool & pool,
@@ -36,7 +42,7 @@ static void top_k_cub(ggml_cuda_pool & pool,
ncols, k, env));
}
#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE
#elif defined(GGML_CUDA_USE_CUB) && !defined(GGML_CUDA_TOP_K_RADIX) // CUB_TOP_K_AVAILABLE
static int next_power_of_2(int x) {
int n = 1;
@@ -48,7 +54,7 @@ static int next_power_of_2(int x) {
#endif // CUB_TOP_K_AVAILABLE
#if !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP)
#if defined(GGML_CUDA_TOP_K_RADIX) || (!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);
@@ -232,6 +238,17 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
for (int i = 0; i < nrows; i++) {
top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream);
}
#elif defined(GGML_CUDA_TOP_K_RADIX)
// exact radix selection: ~nrows*blocks*256 ints of temp (MBs, not 100s of MBs)
if (ncols > 1024) {
top_k_radix_cuda(pool, src0_d, dst_d, (int) ncols, (int) nrows, (int) k, stream);
} else {
ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows);
int * tmp_dst = temp_dst_alloc.get();
argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream);
CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows,
cudaMemcpyDeviceToDevice, stream));
}
#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE
// Fall back to argsort + copy
const int ncols_pad = next_power_of_2(ncols);