CUDA: enable the CUB path on HIP via hipCUB

.github/workflows/build-cuda-ubuntu.yml

@@ -85,7 +85,7 @@ jobs:
         id: depends
         run: |
           sudo apt-get update
-          sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev libssl-dev rocwmma-dev
+          sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev hipcub-dev libssl-dev rocwmma-dev
 
       - name: ccache
         uses: ggml-org/ccache-action@v1.2.21

.github/workflows/hip-quality-check.yml

@@ -49,7 +49,7 @@ jobs:
         id: depends
         run: |
           sudo apt-get update
-          sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev libssl-dev python3
+          sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev hipcub-dev libssl-dev python3
 
       - name: ccache
         uses: ggml-org/ccache-action@v1.2.21

ggml/src/ggml-cuda/argsort.cu

@@ -1,11 +1,10 @@
 #include "argsort.cuh"
 
 #ifdef GGML_CUDA_USE_CUB
-#    include <cub/cub.cuh>
 #    if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 1)
 #        define STRIDED_ITERATOR_AVAILABLE
 #        include <cuda/iterator>
-#    endif
+#    endif  // CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 1
 using namespace cub;
 #endif  // GGML_CUDA_USE_CUB
 

ggml/src/ggml-cuda/common.cuh

@@ -107,10 +107,6 @@
 #define GGML_CUDA_CC_IS_QY2(cc)      (cc >= GGML_CUDA_CC_QY2 && cc < GGML_CUDA_CC_PH1)
 #define GGML_CUDA_CC_IS_PH1(cc)      (cc >= GGML_CUDA_CC_PH1)
 
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
-#    define GGML_CUDA_USE_CUB
-#endif  // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
-
 // PDL host-side support (cudaLaunchKernelEx) requires CUDART >= 11.8.
 // However, this has been bugged in CTK < 12.3 for MSVC builds, see
 // https://github.com/ggml-org/llama.cpp/pull/22522#discussion_r3302393293

ggml/src/ggml-cuda/cumsum.cu

@@ -4,10 +4,6 @@
 #include "ggml-cuda/common.cuh"
 #include "ggml.h"
 
-#ifdef GGML_CUDA_USE_CUB
-#   include <cub/cub.cuh>
-#endif // GGML_CUDA_USE_CUB
-
 template<typename T, int BLOCK_SIZE>
 static __global__ void cumsum_cub_kernel(
         const T * __restrict__ src,
@@ -195,18 +191,18 @@ static void cumsum_cub(ggml_cuda_pool & pool,
     size_t tmp_size = 0;
 
     // Query how much temp storage CUDA UnBound (CUB) needs
-    cub::DeviceScan::InclusiveSum(nullptr,   // d_temp_storage (null = just query size)
-                                  tmp_size,  // reference to size (will be set by CUB)
-                                  src,       // input pointer
-                                  dst,       // output pointer
-                                  ne,        // number of elements
-                                  stream     // CUDA stream to use
-    );
+    CUDA_CHECK(cub::DeviceScan::InclusiveSum(nullptr,   // d_temp_storage (null = just query size)
+                                             tmp_size,  // reference to size (will be set by CUB)
+                                             src,       // input pointer
+                                             dst,       // output pointer
+                                             ne,        // number of elements
+                                             stream     // CUDA stream to use
+    ));
 
     ggml_cuda_pool_alloc<uint8_t> tmp_alloc(pool, tmp_size);
 
     // Perform the inclusive scan
-    cub::DeviceScan::InclusiveSum((void *) tmp_alloc.get(), tmp_size, src, dst, ne, stream);
+    CUDA_CHECK(cub::DeviceScan::InclusiveSum((void *) tmp_alloc.get(), tmp_size, src, dst, ne, stream));
 }
 #endif // GGML_CUDA_USE_CUB
 

ggml/src/ggml-cuda/mean.cu

@@ -2,7 +2,6 @@
 #include "reduce_rows.cuh"
 
 #ifdef GGML_CUDA_USE_CUB
-#include <cub/cub.cuh>
 using namespace cub;
 #endif  // GGML_CUDA_USE_CUB
 
@@ -47,10 +46,10 @@ void ggml_cuda_op_mean(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
         size_t           tmp_size = 0;
         ggml_cuda_pool & pool     = ctx.pool();
 
-        DeviceReduce::Sum(nullptr, tmp_size, src0_d, dst_d, ncols, stream);
+        CUDA_CHECK(DeviceReduce::Sum(nullptr, tmp_size, src0_d, dst_d, ncols, stream));
 
         ggml_cuda_pool_alloc<uint8_t> tmp_alloc(pool, tmp_size);
-        DeviceReduce::Sum(tmp_alloc.ptr, tmp_size, src0_d, dst_d, ncols, stream);
+        CUDA_CHECK(DeviceReduce::Sum(tmp_alloc.ptr, tmp_size, src0_d, dst_d, ncols, stream));
 
         // Divide by ncols
         divide_by_count<float><<<1, 1, 0, stream>>>(dst_d, ncols);

ggml/src/ggml-cuda/ssm-scan.cu

@@ -1,15 +1,13 @@
+#include "ssm-scan.cuh"
+
 #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
 #define USE_CUB
 #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
 
 #ifdef USE_CUB
-#include <cub/cub.cuh>
 using namespace cub;
 #endif // USE_CUB
 
-#include "ssm-scan.cuh"
-
-
 // Minimum number of tokens to use SSD (State Space Duality) matmul path instead of scan path.
 // For n_tok <= this threshold, the scan kernel is used (lower overhead for short sequences).
 #define SSM_SSD_MIN_TOKENS 128

ggml/src/ggml-cuda/sum.cu

@@ -1,19 +1,17 @@
 #include "sum.cuh"
 #include "sumrows.cuh"
+#include <cstdint>
 
 #ifdef GGML_CUDA_USE_CUB
-#include <cub/cub.cuh>
 using namespace cub;
 #endif  // GGML_CUDA_USE_CUB
 
-#include <cstdint>
-
 void sum_f32_cuda(ggml_cuda_pool & pool, const float * x, float * dst, const int64_t ne, cudaStream_t stream) {
 #ifdef GGML_CUDA_USE_CUB
     size_t tmp_size = 0;
-    DeviceReduce::Sum(nullptr,       tmp_size, x, dst, ne, stream);
+    CUDA_CHECK(DeviceReduce::Sum(nullptr,       tmp_size, x, dst, ne, stream));
     ggml_cuda_pool_alloc<uint8_t> tmp_alloc(pool, tmp_size);
-    DeviceReduce::Sum(tmp_alloc.ptr, tmp_size, x, dst, ne, stream);
+    CUDA_CHECK(DeviceReduce::Sum(tmp_alloc.ptr, tmp_size, x, dst, ne, stream));
 #else
     // Use (inefficient) sum_rows implementation as a fallback.
     // For AMD there is rocPRIM which could be used as a drop-in replacement via hipcub but this would require C++11 -> C++14.

ggml/src/ggml-cuda/top-k.cu

@@ -2,7 +2,6 @@
 #include "top-k.cuh"
 
 #ifdef GGML_CUDA_USE_CUB
-#    include <cub/cub.cuh>
 #    if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2)
 #        define CUB_TOP_K_AVAILABLE
 #        include <cuda/iterator>

ggml/src/ggml-cuda/vendors/cuda.h

@@ -15,6 +15,11 @@
 #define FP8_AVAILABLE
 #endif // CUDART_VERSION >= 11080
 
+#if CUDART_VERSION >= 11070
+#define GGML_CUDA_USE_CUB
+#include <cub/cub.cuh>
+#endif // CUDART_VERSION >= 11070
+
 #if CUDART_VERSION >= 12080
 #include <cuda_fp4.h>
 #endif // CUDART_VERSION >= 12080

ggml/src/ggml-cuda/vendors/hip.h

@@ -10,6 +10,11 @@
 #include <rccl/rccl.h>
 #endif // GGML_USE_NCCL
 
+#if HIP_VERSION >= 60100000
+#define GGML_CUDA_USE_CUB
+#include <hipcub/hipcub.hpp>
+namespace cub = hipcub;
+#endif // HIP_VERSION >= 60100000
 
 #define CUBLAS_GEMM_DEFAULT HIPBLAS_GEMM_DEFAULT
 #define CUBLAS_GEMM_DEFAULT_TENSOR_OP HIPBLAS_GEMM_DEFAULT
@@ -118,6 +123,9 @@
 #define cudaStreamPerThread hipStreamPerThread
 #define cudaStreamSynchronize hipStreamSynchronize
 #define cudaStreamWaitEvent hipStreamWaitEvent
+#define cudaStreamIsCapturing hipStreamIsCapturing
+#define cudaStreamCaptureStatus hipStreamCaptureStatus
+#define cudaStreamCaptureStatusNone hipStreamCaptureStatusNone
 #define cudaGraphExec_t hipGraphExec_t
 #define cudaGraphNode_t hipGraphNode_t
 #define cudaKernelNodeParams hipKernelNodeParams

tests/test-backend-ops.cpp

@@ -9371,6 +9371,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2049, 2, 1, 3}, k));
     }
 
+    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {262144, 8192, 1, 1}, 1024));
+    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {1048576, 512, 1, 1}, 2048));
+
     // exhaustive top_k tests
     //for (int i = 1; i < 9999; ++i) {
     //    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {i, 2, 1, 3}, rand() % i + 1));