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

Open
fairydreaming wants to merge 4 commits into
base: master
Choose a base branch
from
+47 −12
Open

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

Changes from 1 commit
Commits
Show all changes
4 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
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
30 changes: 22 additions & 8 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 @@ -76,16 +76,30 @@ void ggml_cuda_op_top_k(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;

ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * nrows);
int * tmp_dst = temp_dst_alloc.get();
// process input in chunks to avoid excessive temporary buffers memory usage
const size_t nb01 = src0->nb[1];
const size_t chunk_size = 1LL << 26;
const int64_t nrows_per_chunk = chunk_size > nb01 ? chunk_size / nb01 : 1;
// make sure chunk_nrows can be safely cast to int below
GGML_ASSERT(nrows_per_chunk <= std::numeric_limits<int>::max());

for (int64_t i = 0; i < nrows; i+= nrows_per_chunk) {
Comment thread
fairydreaming marked this conversation as resolved.
Outdated
int64_t chunk_nrows = std::min(nrows_per_chunk, nrows - i);

ggml_cuda_pool_alloc<int> temp_dst_alloc(pool, ncols * chunk_nrows);

@ggerganov ggerganov Jun 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

I'm not sure how the cuda pool works exactly - just wondering if we actually need to have this allocation inside the loop and not one time before it?

@fairydreaming fairydreaming Jun 19, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

I'm not sure how the cuda pool works exactly - just wondering if we actually need to have this allocation inside the loop and not one time before it?

@ggerganov Good point, will try it out.

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);
if (shared_mem > max_shared_mem || ncols > 1024) {
argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, chunk_nrows, GGML_SORT_ORDER_DESC, stream);
} else {
argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, chunk_nrows, GGML_SORT_ORDER_DESC, stream);
}
CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), chunk_nrows,
cudaMemcpyDeviceToDevice, stream));

src0_d += ncols * chunk_nrows;
dst_d += k * chunk_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