Skip to content

(WIP) CUDA/HIP: GPU partial top-k for ROCm (fix full-sort/CPU fallback) #27860

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
zihaomu wants to merge 2 commits into from
Closed

(WIP) CUDA/HIP: GPU partial top-k for ROCm (fix full-sort/CPU fallback) #27860

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
35c6be1
CUDA/HIP: GPU partial top-k for ROCm (fix full-sort/CPU fallback)
zihaomu Aug 28, 2026
400f0be
CUDA/HIP: switch ROCm top-k to capture-safe radix sort-pairs
zihaomu Aug 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions ggml/src/ggml-cuda/ggml-cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2570,6 +2570,21 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) {
}
}

#if defined(GGML_USE_HIP)
// [TAG_TOP_K_HIP_CUDA_GRAPHS]
// Multi-row top-k uses hipCUB DeviceSegmentedRadixSort. Instantiating a HIP graph that
// contains a large segmented sort overflows the ROCm runtime's graph builder (stack
// overflow inside libamdhip64). Single-row top-k (decode) uses DeviceRadixSort and is
// graph-safe, so only skip graphs for the nrows>1 case (prefill/batched, which does not
// rely on CUDA graphs anyway).
if (node->op == GGML_OP_TOP_K && ggml_nrows(node->src[0]) > 1) {
use_cuda_graph = false;
#ifndef NDEBUG
GGML_LOG_DEBUG("%s: disabling CUDA graphs due to multi-row TOP_K on ROCm\n", __func__);
#endif
}
#endif // GGML_USE_HIP

if (!use_cuda_graph) {
break;
}
Expand Down Expand Up @@ -5259,6 +5274,14 @@ 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_USE_HIP)
// capture-safe radix SortPairs handles any ncols / nrows on ROCm
return true;
#elif !defined(GGML_CUDA_USE_CUB)
return op->src[0]->ne[0] <= 1024;
#else
return true;
#endif
case GGML_OP_ARGSORT:
#ifndef GGML_CUDA_USE_CUB
return op->src[0]->ne[0] <= 1024;
Expand Down
56 changes: 56 additions & 0 deletions ggml/src/ggml-cuda/top-k.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,26 @@ static int next_power_of_2(int x) {

#endif // CUB_TOP_K_AVAILABLE

#if defined(GGML_USE_HIP)
#include <hipcub/hipcub.hpp>
// hipCUB has no DeviceTopK, so on ROCm top-k is implemented by sorting (score, index) pairs in
// descending order with DeviceRadixSort (single row) / DeviceSegmentedRadixSort (multi-row) and
// keeping the first k of each row. These radix sorts are stream-capture-safe. The helpers below
// fill the per-element value (original column index) and the per-row segment offsets.
static __global__ void ggml_cuda_topk_iota_rows(int * idx, const int ncols, const int64_t n) {
const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
idx[i] = (int) (i % ncols);
}
}
static __global__ void ggml_cuda_topk_init_offsets(int * off, const int ncols, const int n1) {
const int s = blockIdx.x * blockDim.x + threadIdx.x;
if (s < n1) {
off[s] = s * ncols;
}
}
#endif // GGML_USE_HIP

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;
Expand All @@ -70,6 +90,42 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
for (int i = 0; i < nrows; i++) {
top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream);
}
#elif defined(GGML_USE_HIP) // CUB_TOP_K_AVAILABLE
// ROCm top-k via capture-safe radix sort-pairs (see doc/PLAN_hip_partial_topk.md).
{
const int block = 256;
const int64_t n = ncols * nrows;
ggml_cuda_pool_alloc<float> keys_alloc(pool, n);
ggml_cuda_pool_alloc<int> vals_in_alloc(pool, n);
ggml_cuda_pool_alloc<int> vals_out_alloc(pool, n);
float * keys = keys_alloc.get();
int * vals_in = vals_in_alloc.get();
int * vals_out = vals_out_alloc.get();

CUDA_CHECK(cudaMemcpyAsync(keys, src0_d, n * sizeof(float), cudaMemcpyDeviceToDevice, stream));
ggml_cuda_topk_iota_rows<<<(n + block - 1) / block, block, 0, stream>>>(vals_in, (int) ncols, n);

size_t tmp_bytes = 0;
if (nrows == 1) {
CUDA_CHECK(hipcub::DeviceRadixSort::SortPairsDescending(nullptr, tmp_bytes, keys, keys,
vals_in, vals_out, (int) ncols, 0, (int) sizeof(float) * 8, stream));
ggml_cuda_pool_alloc<uint8_t> tmp_alloc(pool, tmp_bytes);
CUDA_CHECK(hipcub::DeviceRadixSort::SortPairsDescending(tmp_alloc.get(), tmp_bytes, keys, keys,
vals_in, vals_out, (int) ncols, 0, (int) sizeof(float) * 8, stream));
} else {
ggml_cuda_pool_alloc<int> off_alloc(pool, nrows + 1);
int * off = off_alloc.get();
ggml_cuda_topk_init_offsets<<<(nrows + 1 + block - 1) / block, block, 0, stream>>>(off, (int) ncols, (int) (nrows + 1));
CUDA_CHECK(hipcub::DeviceSegmentedRadixSort::SortPairsDescending(nullptr, tmp_bytes, keys, keys,
vals_in, vals_out, (int) n, (int) nrows, off, off + 1, 0, (int) sizeof(float) * 8, stream));
ggml_cuda_pool_alloc<uint8_t> tmp_alloc(pool, tmp_bytes);
CUDA_CHECK(hipcub::DeviceSegmentedRadixSort::SortPairsDescending(tmp_alloc.get(), tmp_bytes, keys, keys,
vals_in, vals_out, (int) n, (int) nrows, off, off + 1, 0, (int) sizeof(float) * 8, stream));
}
// keep the first k sorted indices of each row
CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), vals_out, ncols * sizeof(int),
k * sizeof(int), nrows, cudaMemcpyDeviceToDevice, stream));
}
#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE
// Fall back to argsort + copy
const int ncols_pad = next_power_of_2(ncols);
Expand Down