CUDA/HIP: support argsort/top-k ncols>1024 on ROCm via hipCUB

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 <hipcub/hipcub.hpp>
+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<typename T>
@@ -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;

server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh

@@ -2,15 +2,15 @@
 
 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,
                               const int        ncols,
                               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,

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;

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:

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