From 0dfab41182aaaa2284650461190ab7f0831b3c3c Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Wed, 3 Dec 2025 01:13:43 +0000 Subject: [PATCH 1/3] try not to fail when there should be memory available --- mlx/backend/cuda/allocator.cpp | 36 ++++++++++++++++++++++++++-------- mlx/backend/cuda/allocator.h | 1 + 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp index 5f8ecd715b..5596a4dda9 100644 --- a/mlx/backend/cuda/allocator.cpp +++ b/mlx/backend/cuda/allocator.cpp @@ -92,7 +92,8 @@ CudaAllocator::CudaAllocator() [this](CudaBuffer* buf) { cuda_free(buf); }) { size_t free, total; CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total)); - memory_limit_ = total * 0.9; + memory_limit_ = total * 0.95; + free_limit_ = total - memory_limit_; max_pool_size_ = memory_limit_; int device_count = 0; @@ -152,14 +153,26 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) { if (size <= small_block_size) { buf = scalar_pool_.malloc(); } - lock.unlock(); if (!buf) { - cudaError_t err; void* data = nullptr; - if (device == -1) { - CHECK_CUDA_ERROR(cudaMallocManaged(&data, size)); - } else { - CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream)); + while (!data) { + lock.unlock(); + + if (device == -1) { + CHECK_CUDA_ERROR(cudaMallocManaged(&data, size)); + } else { + CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream)); + } + + if (!data) { + // Try to release memory from the cache to defrag + if (get_cache_memory() == 0) { + break; + } else { + buffer_cache_.release_cached_buffers(size); + } + } + lock.lock(); } if (!data) { std::ostringstream msg; @@ -168,11 +181,18 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) { } buf = new CudaBuffer{data, size, device}; } - lock.lock(); } active_memory_ += buf->size; peak_memory_ = std::max(active_memory_, peak_memory_); + // If the OS reports that not enough memory is free, try to clear some memory + // from the cache. This prevents graph / kernel execution failing from OOM + size_t free, total; + CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total)); + if (free < free_limit_ && get_cache_memory() > 0) { + buffer_cache_.release_cached_buffers(free_limit_); + } + // Maintain the cache below the requested limit. if (get_cache_memory() > max_pool_size_) { buffer_cache_.release_cached_buffers(get_cache_memory() - max_pool_size_); diff --git a/mlx/backend/cuda/allocator.h b/mlx/backend/cuda/allocator.h index 133d4f8525..94707a9bd2 100644 --- a/mlx/backend/cuda/allocator.h +++ b/mlx/backend/cuda/allocator.h @@ -71,6 +71,7 @@ class CudaAllocator : public allocator::Allocator { std::mutex mutex_; size_t memory_limit_; + size_t free_limit_; size_t max_pool_size_; BufferCache buffer_cache_; size_t active_memory_{0}; From b3d6566bda1f147331685b6a549e3f70e76b84c0 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Thu, 4 Dec 2025 15:02:03 +0000 Subject: [PATCH 2/3] speed up mem check --- mlx/backend/cuda/allocator.cpp | 80 ++++++++++++++++++---------------- mlx/backend/cuda/allocator.h | 2 + 2 files changed, 45 insertions(+), 37 deletions(-) diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp index 5596a4dda9..f059a1b7fb 100644 --- a/mlx/backend/cuda/allocator.cpp +++ b/mlx/backend/cuda/allocator.cpp @@ -20,6 +20,19 @@ constexpr int page_size = 16384; // Any allocations smaller than this will try to use the small pool constexpr int small_block_size = 8; +#if CUDART_VERSION >= 13000 +cudaMemLocation cuda_mem_loc(int i) { + cudaMemLocation loc; + loc.type = cudaMemLocationTypeDevice; + loc.id = i; + return loc; +} +#else +int cuda_mem_loc(int i) { + return i; +} +#endif // CUDART_VERSION >= 13000 + // The small pool size in bytes. This should be a multiple of the host page // size and small_block_size. constexpr int small_pool_size = 4 * page_size; @@ -35,13 +48,7 @@ SmallSizePool::SmallSizePool() { int device_count = 0; CHECK_CUDA_ERROR(cudaGetDeviceCount(&device_count)); for (int i = 0; i < device_count; ++i) { -#if CUDART_VERSION >= 13000 - cudaMemLocation loc; - loc.type = cudaMemLocationTypeDevice; - loc.id = i; -#else - int loc = i; -#endif // CUDART_VERSION >= 13000 + auto loc = cuda_mem_loc(i); CHECK_CUDA_ERROR( cudaMemAdvise(data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc)); } @@ -90,10 +97,10 @@ CudaAllocator::CudaAllocator() page_size, [](CudaBuffer* buf) { return buf->size; }, [this](CudaBuffer* buf) { cuda_free(buf); }) { - size_t free, total; - CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total)); - memory_limit_ = total * 0.95; - free_limit_ = total - memory_limit_; + size_t free; + CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total_memory_)); + memory_limit_ = total_memory_ * 0.95; + free_limit_ = total_memory_ - memory_limit_; max_pool_size_ = memory_limit_; int device_count = 0; @@ -105,6 +112,10 @@ CudaAllocator::CudaAllocator() cudaStream_t s; CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking)); free_streams_.push_back(s); + + cudaMemPool_t mem_pool; + CHECK_CUDA_ERROR(cudaDeviceGetDefaultMemPool(&mem_pool, i)); + mem_pools_.push_back(mem_pool); } CHECK_CUDA_ERROR(cudaSetDevice(curr)); } @@ -153,26 +164,13 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) { if (size <= small_block_size) { buf = scalar_pool_.malloc(); } + lock.unlock(); if (!buf) { void* data = nullptr; - while (!data) { - lock.unlock(); - - if (device == -1) { - CHECK_CUDA_ERROR(cudaMallocManaged(&data, size)); - } else { - CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream)); - } - - if (!data) { - // Try to release memory from the cache to defrag - if (get_cache_memory() == 0) { - break; - } else { - buffer_cache_.release_cached_buffers(size); - } - } - lock.lock(); + if (device == -1) { + CHECK_CUDA_ERROR(cudaMallocManaged(&data, size)); + } else { + CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream)); } if (!data) { std::ostringstream msg; @@ -181,18 +179,26 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) { } buf = new CudaBuffer{data, size, device}; } + lock.lock(); + + // If any cuda memory pool has too much reserved memory, clear some + // memory from the cache. This prevents graph / kernel execution failing + // from OOM + if (get_cache_memory() > 0) { + for (auto p : mem_pools_) { + size_t used = 0; + CHECK_CUDA_ERROR(cudaMemPoolGetAttribute( + p, cudaMemPoolAttrReservedMemCurrent, &used)); + if (used > (total_memory_ - free_limit_)) { + buffer_cache_.release_cached_buffers(free_limit_); + break; + } + } + } } active_memory_ += buf->size; peak_memory_ = std::max(active_memory_, peak_memory_); - // If the OS reports that not enough memory is free, try to clear some memory - // from the cache. This prevents graph / kernel execution failing from OOM - size_t free, total; - CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total)); - if (free < free_limit_ && get_cache_memory() > 0) { - buffer_cache_.release_cached_buffers(free_limit_); - } - // Maintain the cache below the requested limit. if (get_cache_memory() > max_pool_size_) { buffer_cache_.release_cached_buffers(get_cache_memory() - max_pool_size_); diff --git a/mlx/backend/cuda/allocator.h b/mlx/backend/cuda/allocator.h index 94707a9bd2..7f6ad52549 100644 --- a/mlx/backend/cuda/allocator.h +++ b/mlx/backend/cuda/allocator.h @@ -72,11 +72,13 @@ class CudaAllocator : public allocator::Allocator { std::mutex mutex_; size_t memory_limit_; size_t free_limit_; + size_t total_memory_; size_t max_pool_size_; BufferCache buffer_cache_; size_t active_memory_{0}; size_t peak_memory_{0}; std::vector free_streams_; + std::vector mem_pools_; SmallSizePool scalar_pool_; }; From db9a71b9a53f45eacbc2637b2db8b2163709f5b8 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Sat, 6 Dec 2025 14:19:23 +0000 Subject: [PATCH 3/3] comment --- mlx/backend/cuda/allocator.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp index f059a1b7fb..13638c1afa 100644 --- a/mlx/backend/cuda/allocator.cpp +++ b/mlx/backend/cuda/allocator.cpp @@ -21,14 +21,14 @@ constexpr int page_size = 16384; constexpr int small_block_size = 8; #if CUDART_VERSION >= 13000 -cudaMemLocation cuda_mem_loc(int i) { +inline cudaMemLocation cuda_mem_loc(int i) { cudaMemLocation loc; loc.type = cudaMemLocationTypeDevice; loc.id = i; return loc; } #else -int cuda_mem_loc(int i) { +inline int cuda_mem_loc(int i) { return i; } #endif // CUDART_VERSION >= 13000