diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2456f7dcc621..7423bd48fd8f 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5259,6 +5259,11 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_SUM: return ggml_is_contiguous_rows(op->src[0]); case GGML_OP_TOP_K: +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_HIP_TOPK_HIPCUB) + return true; +#else + return op->src[0]->ne[0] <= 1024; +#endif case GGML_OP_ARGSORT: #ifndef GGML_CUDA_USE_CUB return op->src[0]->ne[0] <= 1024; diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index 9681cd293338..ef8385035b3b 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -1,6 +1,10 @@ #include "argsort.cuh" #include "top-k.cuh" +#ifdef GGML_HIP_TOPK_HIPCUB +# include +#endif // GGML_HIP_TOPK_HIPCUB + #ifdef GGML_CUDA_USE_CUB # include # if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2) @@ -48,6 +52,89 @@ static int next_power_of_2(int x) { #endif // CUB_TOP_K_AVAILABLE +#ifdef GGML_HIP_TOPK_HIPCUB + +// The HIP backend has no CUB, so the shared-memory bitonic argsort is the only +// top-k path and it caps rows at 1024 elements. Vocabulary-width top-k (e.g. +// the DFlash2 candidate selector, k = 16 over 248,320 logits) therefore falls +// back to the CPU, which costs a device-to-host copy of the whole logits +// tensor on every draft step. hipCUB's segmented radix sort keeps it on the +// GPU. + +static __global__ void top_k_hipcub_init_indices(int * indices, const int ncols, const int nrows) { + const int col = blockIdx.x * blockDim.x + threadIdx.x; + const int row = blockIdx.y; + + if (col < ncols && row < nrows) { + indices[row * ncols + col] = col; + } +} + +static __global__ void top_k_hipcub_init_offsets(int * offsets, const int ncols, const int nrows_offset) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + + if (i < nrows_offset) { + offsets[i] = i * ncols; + } +} + +static void top_k_hipcub(ggml_cuda_pool & pool, + const float * src, + int * dst, + const int ncols, + const int nrows, + const int k, + cudaStream_t stream) { + const size_t n = (size_t) ncols * nrows; + + ggml_cuda_pool_alloc keys_in_alloc (pool, n); + ggml_cuda_pool_alloc keys_out_alloc(pool, n); + ggml_cuda_pool_alloc vals_in_alloc (pool, n); + ggml_cuda_pool_alloc vals_out_alloc(pool, n); + + float * keys_in = keys_in_alloc .get(); + float * keys_out = keys_out_alloc.get(); + int * vals_in = vals_in_alloc .get(); + int * vals_out = vals_out_alloc.get(); + + static const int block_size = 256; + + const dim3 grid_size((ncols + block_size - 1) / block_size, nrows); + top_k_hipcub_init_indices<<>>(vals_in, ncols, nrows); + + const int nrows_offset = nrows + 1; + ggml_cuda_pool_alloc offsets_alloc(pool, nrows_offset); + int * offsets = offsets_alloc.get(); + const dim3 offset_grid((nrows_offset + block_size - 1) / block_size); + top_k_hipcub_init_offsets<<>>(offsets, ncols, nrows_offset); + + CUDA_CHECK(cudaMemcpyAsync(keys_in, src, n * sizeof(float), cudaMemcpyDeviceToDevice, stream)); + + size_t temp_storage_bytes = 0; + CUDA_CHECK(hipcub::DeviceSegmentedRadixSort::SortPairsDescending( + nullptr, temp_storage_bytes, keys_in, keys_out, vals_in, vals_out, n, nrows, offsets, offsets + 1, 0, + sizeof(float) * 8, stream)); + + ggml_cuda_pool_alloc temp_storage_alloc(pool, temp_storage_bytes); + CUDA_CHECK(hipcub::DeviceSegmentedRadixSort::SortPairsDescending( + temp_storage_alloc.get(), temp_storage_bytes, keys_in, keys_out, vals_in, vals_out, n, nrows, offsets, + offsets + 1, 0, sizeof(float) * 8, stream)); + + // the op only requires the k largest indices, in any order + CUDA_CHECK(cudaMemcpy2DAsync(dst, k * sizeof(int), vals_out, ncols * sizeof(int), k * sizeof(int), nrows, + cudaMemcpyDeviceToDevice, stream)); +} + +// process at most this many bytes of keys+values per segmented sort call +static int top_k_hipcub_chunk_nrows(const int ncols, const int nrows) { + const size_t chunk_bytes = 1 << 26; // 64 MiB + const size_t row_bytes = (size_t) ncols * (sizeof(float) + sizeof(int)) * 2; + + return (int) std::min((int64_t) std::max(chunk_bytes / std::max(row_bytes, 1), 1), (int64_t) nrows); +} + +#endif // GGML_HIP_TOPK_HIPCUB + void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const float * src0_d = (const float *) src0->data; @@ -96,6 +183,22 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { dst_d += k * iter_nrows; } #else // GGML_CUDA_USE_CUB +# ifdef GGML_HIP_TOPK_HIPCUB + // the bitonic argsort below needs the whole row in shared memory + if (ncols > 1024) { + const int chunk_nrows = top_k_hipcub_chunk_nrows(ncols, nrows); + + for (int64_t i = 0; i < nrows; i += chunk_nrows) { + const int iter_nrows = std::min((int64_t) chunk_nrows, nrows - i); + + top_k_hipcub(pool, src0_d, dst_d, ncols, iter_nrows, k, stream); + + src0_d += (size_t) ncols * iter_nrows; + dst_d += (size_t) k * iter_nrows; + } + return; + } +# endif // GGML_HIP_TOPK_HIPCUB ggml_cuda_pool_alloc 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); diff --git a/ggml/src/ggml-hip/CMakeLists.txt b/ggml/src/ggml-hip/CMakeLists.txt index 47f16f56c470..c2b1ce276123 100644 --- a/ggml/src/ggml-hip/CMakeLists.txt +++ b/ggml/src/ggml-hip/CMakeLists.txt @@ -47,6 +47,14 @@ find_package(hip REQUIRED) find_package(hipblas REQUIRED) find_package(rocblas REQUIRED) +# hipCUB gives the top-k op a device-side path for rows wider than the shared +# memory bitonic argsort can handle (1024). Without it, vocabulary-width top-k +# is not supported by this backend and falls back to the CPU. +option(GGML_HIP_TOPK_HIPCUB "ggml: use hipCUB for wide-row top-k" ON) +if (GGML_HIP_TOPK_HIPCUB) + find_package(hipcub) +endif() + if (GGML_HIP_RCCL) find_package(rccl REQUIRED) endif() @@ -151,3 +159,9 @@ if (GGML_HIP_RCCL) endif() target_link_libraries(ggml-hip PRIVATE ggml-base hip::host roc::rocblas roc::hipblas) + +if (GGML_HIP_TOPK_HIPCUB AND hipcub_FOUND) + message(STATUS "HIP: using hipCUB for wide-row top-k") + target_compile_definitions(ggml-hip PRIVATE GGML_HIP_TOPK_HIPCUB) + target_link_libraries(ggml-hip PRIVATE hip::hipcub) +endif()