Skip to content

ggml : process data in smaller chunks in CUDA ggml_top_k() implementation to reduce temporary buffers memory usage #24776

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

Merged
fairydreaming merged 6 commits into from
Jul 9, 2026
Merged

ggml : process data in smaller chunks in CUDA ggml_top_k() implementation to reduce temporary buffers memory usage #24776

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c1a79a6
ggml : process data in smaller chunks in CUDA ggml_top_k() implementa…
sszymczy Jun 18, 2026
ca88122
ggml : allocate tmp_dst only only once before the loop
sszymczy Jun 19, 2026
4d0e838
chore : whitespaces
fairydreaming Jun 19, 2026
8521ec4
ggml : use chunked processing in both CUDA CUB top-k and argsort impl…
sszymczy Jun 25, 2026
807a9f7
chore : separate argsort_f32_i32_cuda_bitonic() call from return stat…
fairydreaming Jun 30, 2026
f887bc2
chore : replace ternary operators with min/max
sszymczy Jun 30, 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
35 changes: 30 additions & 5 deletions ggml/src/ggml-cuda/argsort.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,20 @@ static __global__ void init_offsets(int * offsets, const int ncols, const int nr
#endif // STRIDED_ITERATOR_AVAILABLE

#ifdef GGML_CUDA_USE_CUB

// returns the suggested maximum number of rows to process during one argsort_f32_i32_cuda_cub() call
int argsort_f32_i32_cuda_cub_chunk_nrows(const size_t nb01, const int64_t nrows) {
Comment thread
JohannesGaessler marked this conversation as resolved.
Resolved
// perform argsort in chunks up to approximately this size (currently 64MB)
// to avoid excessive temporary buffers memory usage
const int chunk_bytes = 1 << 26;

// calculate how many rows will fit in one chunk (must be at least one)
const int chunk_nrows = chunk_bytes > nb01 ? chunk_bytes / nb01 : 1;

// limit the resulting amount to total nrows
return nrows < chunk_nrows ? nrows : chunk_nrows;

@JohannesGaessler JohannesGaessler Jun 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Preferably use std::max and std::min here.

@fairydreaming fairydreaming Jun 30, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Preferably use std::max and std::min here.

@JohannesGaessler Done, there were two more ternary operators so I included them as well.

}

void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
const float * x,
int * dst,
Expand Down Expand Up @@ -254,11 +268,22 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const size_t shared_mem = ncols_pad * sizeof(int);
const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb;

if (shared_mem > max_shared_mem || ncols > 1024) {
ggml_cuda_pool & pool = ctx.pool();
argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, nrows, order, stream);
} else {
argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream);
// early return if we can use bitonic argsort
if (shared_mem <= max_shared_mem && ncols <= 1024) {
return argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream);
Comment thread
fairydreaming marked this conversation as resolved.
Resolved
}

const int chunk_nrows = argsort_f32_i32_cuda_cub_chunk_nrows(src0->nb[1], nrows);

ggml_cuda_pool & pool = ctx.pool();

for (int64_t i = 0; i < nrows; i += chunk_nrows) {
int iter_nrows = chunk_nrows < nrows - i ? chunk_nrows : nrows - i;

argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, iter_nrows, order, stream);

src0_d += ncols * iter_nrows;
dst_d += ncols * iter_nrows;
}
#else
argsort_f32_i32_cuda_bitonic(src0_d, (int *) dst_d, ncols, nrows, order, stream);
Expand Down
1 change: 1 addition & 0 deletions ggml/src/ggml-cuda/argsort.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

#ifdef GGML_CUDA_USE_CUB
int argsort_f32_i32_cuda_cub_chunk_nrows(const size_t nb01, const int64_t nrows);
void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
const float * x,
int * dst,
Expand Down
23 changes: 16 additions & 7 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 @@ -75,17 +75,26 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
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;
const bool use_bitonic = shared_mem <= max_shared_mem && ncols <= 1024;
const int chunk_nrows = argsort_f32_i32_cuda_cub_chunk_nrows(src0->nb[1], nrows);

ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows);
ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * chunk_nrows);
int * tmp_dst = temp_dst_alloc.get();

if (shared_mem > max_shared_mem || ncols > 1024) {
argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream);
} else {
argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream);
for (int64_t i = 0; i < nrows; i += chunk_nrows) {
int iter_nrows = chunk_nrows < nrows - i ? chunk_nrows : nrows - i;

if (use_bitonic) {
Comment thread
JohannesGaessler marked this conversation as resolved.
Resolved
argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, iter_nrows, GGML_SORT_ORDER_DESC, stream);
} else {
argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, iter_nrows, GGML_SORT_ORDER_DESC, stream);
}
CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), iter_nrows,
cudaMemcpyDeviceToDevice, stream));

src0_d += ncols * iter_nrows;
dst_d += k * iter_nrows;
}
CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows,
cudaMemcpyDeviceToDevice, stream));
#else // GGML_CUDA_USE_CUB
ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows);
int * tmp_dst = temp_dst_alloc.get();
Expand Down
Loading