From 083538a0f6cb4597921b7a88abb8f0ce35db0551 Mon Sep 17 00:00:00 2001 From: deano Date: Thu, 9 Jul 2026 17:12:33 +0300 Subject: [PATCH] CUDA/HIP: support argsort/top-k with ncols > 1024 on ROCm via hipCUB MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The GPU argsort/top-k path uses a single-block bitonic sort (block_dims = ncols_pad), so it caps at 1024 threads/block. On CUDA, ncols > 1024 is handled by the CUB device-sort path; on HIP that path was compiled out (GGML_CUDA_USE_CUB is CUDA-only), so ggml_backend_cuda_supports_op reported TOP_K/ARGSORT as unsupported for ncols > 1024 and they fell back to the CPU reference — 275 of the test-backend-ops cases (full-vocab sampling, large sorts) on every ROCm build. Enable the existing device-sort path on HIP via hipCUB (rocPRIM), which provides the same DeviceRadixSort / DeviceSegmentedSort / DeviceSegmentedRadixSort API: * common.cuh: define GGML_CUDA_USE_HIPCUB on HIP. * argsort.cu/.cuh: include and compile argsort_f32_i32_cuda_cub on HIP too; hipCUB has no CCCL strided iterator, so the init_offsets segment path is used (already the CUB fallback). * top-k.cu: on HIP, route ncols > 1024 through the device sort + slice-first-k, mirroring the CUB branch; the fast partial-bitonic top-k still handles the common ncols <= 1024 decode/verify case unchanged. * ggml-cuda.cu: supports_op returns true for TOP_K/ARGSORT when either CUB or hipCUB is available. test-backend-ops -o TOP_K and -o ARGSORT: 0 not-supported / 0 fail on gfx1201 (R9700) and gfx1151 (Strix Halo) — the 275 previously-unsupported ncols>1024 cases (up to ne=[262144]) now run on-GPU and are bit-for-bit vs the CPU reference. Requires the rocm hipcub + rocprim headers (ship with ROCm). --- .../llama.cpp/ggml/src/ggml-cuda/argsort.cu | 12 +++++++++--- .../llama.cpp/ggml/src/ggml-cuda/argsort.cuh | 4 ++-- .../llama.cpp/ggml/src/ggml-cuda/common.cuh | 7 +++++++ .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 8 +++++--- .../deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu | 17 +++++++++++++++++ 5 files changed, 40 insertions(+), 8 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu index fe0810981..33d0dfb81 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu @@ -6,6 +6,12 @@ # define STRIDED_ITERATOR_AVAILABLE # endif using namespace cub; +#elif defined(GGML_CUDA_USE_HIPCUB) +// hipCUB exposes the CUB device-sort API (DeviceRadixSort / DeviceSegmentedSort / +// DeviceSegmentedRadixSort) over rocPRIM. No strided-iterator / CCCL support, so +// STRIDED_ITERATOR_AVAILABLE stays undefined and the init_offsets path is used. +# include +using namespace hipcub; #endif // GGML_CUDA_USE_CUB static __global__ void init_indices(int * indices, const int ncols, const int nrows) { @@ -26,7 +32,7 @@ static __global__ void init_offsets(int * offsets, const int ncols, const int nr } #endif // STRIDED_ITERATOR_AVAILABLE -#ifdef GGML_CUDA_USE_CUB +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const float * x, int * dst, @@ -138,7 +144,7 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, } } } -#endif // GGML_CUDA_USE_CUB +#endif // GGML_CUDA_USE_CUB || GGML_CUDA_USE_HIPCUB // Bitonic sort implementation template @@ -317,7 +323,7 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { enum ggml_sort_order order = (enum ggml_sort_order) dst->op_params[0]; -#ifdef GGML_CUDA_USE_CUB +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) const int ncols_pad = next_power_of_2(ncols); const size_t shared_mem = ncols_pad * sizeof(int); const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh index 22b7306f2..018dc9fb2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh @@ -2,7 +2,7 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst); -#ifdef GGML_CUDA_USE_CUB +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const float * x, int * dst, @@ -10,7 +10,7 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const int nrows, ggml_sort_order order, cudaStream_t stream); -#endif // GGML_CUDA_USE_CUB +#endif // GGML_CUDA_USE_CUB || GGML_CUDA_USE_HIPCUB void argsort_f32_i32_cuda_bitonic(const float * x, int * dst, const int ncols, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh index 9df1461a5..319b8742c 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh @@ -107,6 +107,13 @@ # define GGML_CUDA_USE_CUB #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +// HIP has no CUB, but rocPRIM ships a drop-in CUB-API layer via hipCUB. Enable it +// so argsort / top-k support ncols > 1024 (the single-block bitonic path caps at +// 1024 threads/block); without it those ops fall back to the CPU reference on ROCm. +#if defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +# define GGML_CUDA_USE_HIPCUB +#endif // defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + #ifdef __CUDA_ARCH_LIST__ constexpr bool ggml_cuda_has_arch_impl(int) { return false; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index feb9104a9..f248e0faf 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5276,10 +5276,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return ggml_is_contiguous_rows(op->src[0]); case GGML_OP_TOP_K: case GGML_OP_ARGSORT: -#ifndef GGML_CUDA_USE_CUB - return op->src[0]->ne[0] <= 1024; -#else +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) return true; +#else + // No device-wide sort backend: the single-block bitonic path caps at + // 1024 threads/block, so only ncols <= 1024 is supported on GPU. + return op->src[0]->ne[0] <= 1024; #endif case GGML_OP_SUM_ROWS: case GGML_OP_MEAN: diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu index 5edc1899e..5ab6da929 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu @@ -686,6 +686,23 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows, cudaMemcpyDeviceToDevice, stream)); #else // GGML_CUDA_USE_CUB + // ncols > 1024 exceeds the single-block bitonic cap (1024 threads/block). + // Route it through the device-wide sort (hipCUB / rocPRIM on ROCm) + slice + // the first k, mirroring the CUB branch above; keep the fast partial-bitonic + // top-k for ncols <= 1024 (the common decode/verify case). + if (ncols > 1024) { +#ifdef GGML_CUDA_USE_HIPCUB + ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * nrows); + int * tmp_dst = temp_dst_alloc.get(); + argsort_f32_i32_cuda_cub(pool, 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)); +#else + GGML_UNUSED(pool); + GGML_ABORT("top-k: ncols > 1024 requires CUB/hipCUB (rocPRIM)"); +#endif // GGML_CUDA_USE_HIPCUB + return; + } // Dedicated partial bitonic top-k: keeps only a kpad-wide window instead of // fully sorting the row then discarding all but the first k indices. GGML_UNUSED(pool);