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