From 2167f5f4b46b660fdec6fbe0f60207408f6fd2d8 Mon Sep 17 00:00:00 2001 From: delgon Date: Sun, 20 Sep 2026 10:24:14 +0200 Subject: [PATCH] ggml-cuda : round growing pool requests up to a power of two ggml_cuda_pool_leg reuses a cached block only when it is >= the request, so a cached block smaller than the request can serve neither this request nor any larger later one. A caller whose request grows every step therefore misses the cache and strands its predecessor once per step, for as long as it grows. qsa-prefill.cu is such a caller: it sizes packed K/V scratch from the KV cache, which gains one ubatch per prefill step (8 MiB per 8192-token ubatch). Round the request up to a power of two when the pool's largest block sits just below it, which is that pattern's signature, so the series reuses one block per octave. Nothing is freed, so no address the pool has handed out is invalidated; that matters because pool scratch is returned to the pool during CUDA graph capture while its pointer is already baked into a captured kernel node. The OOM retry falls back to the unrounded size, so it can never fail where it would previously have succeeded. f = 2 is the optimum, not just the obvious choice: a series growing to M costs about f*M live plus M/(f-1) stranded, and f + 1/(f-1) is minimised at f = 2. Assisted-by: Claude Opus 5 (Claude Code) Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-cuda/ggml-cuda.cu | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 255fcf8294b5..2c2889c5f533 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -424,6 +424,8 @@ const ggml_cuda_device_info & ggml_cuda_info() { // buffer pool for cuda (legacy) struct ggml_cuda_pool_leg : public ggml_cuda_pool { static const int MAX_BUFFERS = 256; + // far below any KV-sized scratch, so small buffers are never rounded + static const size_t POOL_POW2_MIN_SIZE = 4ull*1024*1024; int device; struct ggml_cuda_buffer { @@ -463,6 +465,7 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { #endif size_t best_diff = 1ull << 36; int ibest = -1; + size_t largest_unusable = 0; for (int i = 0; i < MAX_BUFFERS; ++i) { ggml_cuda_buffer& b = buffer_pool[i]; if (b.ptr != nullptr) { @@ -470,7 +473,9 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { ++nnz; if (b.size > max_size) max_size = b.size; #endif - if (b.size >= size) { + if (b.size < size) { + if (b.size > largest_unusable) largest_unusable = b.size; + } else { size_t diff = b.size - size; if (diff < best_diff) { best_diff = diff; @@ -497,6 +502,20 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { void * ptr; size_t look_ahead_size = (size_t) (1.05 * size); look_ahead_size = 256 * ((look_ahead_size + 255)/256); + const size_t exact_size = look_ahead_size; + // Nothing cached can serve this request, and a cached block smaller than it + // never will, so a caller that grows its request step by step strands one block + // per step. A largest cached block just below the request is that signature: + // round up to a power of two so the series reuses one block per octave instead. + if (size > POOL_POW2_MIN_SIZE && largest_unusable >= size - size/4) { + size_t pow2 = POOL_POW2_MIN_SIZE; + while (pow2 < size && pow2 <= SIZE_MAX/2) { + pow2 *= 2; + } + if (pow2 >= size) { + look_ahead_size = pow2; + } + } ggml_cuda_set_device(device); cudaError_t err = ggml_cuda_device_malloc(&ptr, look_ahead_size, device); if (err == cudaErrorMemoryAllocation) { @@ -506,6 +525,7 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { device, look_ahead_size/1024.0/1024.0, cached_bytes/1024.0/1024.0); CUDA_CHECK(cudaDeviceSynchronize()); clear_pool(); + look_ahead_size = exact_size; // retry without the rounding err = ggml_cuda_device_malloc(&ptr, look_ahead_size, device); if (err == cudaSuccess) { GGML_LOG_DEBUG(GGML_CUDA_NAME " pool[%d]: retry succeeded\n", device);