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);