From 212077f163739dcba33b1a53d13db2d8c3de9481 Mon Sep 17 00:00:00 2001 From: Cheng Date: Fri, 30 Jan 2026 13:18:41 +0900 Subject: [PATCH] Fallback to pinned host memory when managed memory is not supported (#3075) --- mlx/backend/cuda/allocator.cpp | 186 +++++++++++++++++++++------------ mlx/backend/cuda/allocator.h | 9 +- mlx/backend/cuda/cuda_utils.h | 1 + mlx/backend/cuda/fence.cpp | 10 +- 4 files changed, 128 insertions(+), 78 deletions(-) diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp index 0570ed40..d6daee82 100644 --- a/mlx/backend/cuda/allocator.cpp +++ b/mlx/backend/cuda/allocator.cpp @@ -22,6 +22,49 @@ constexpr int page_size = 16384; // Any allocations smaller than this will try to use the small pool constexpr int small_block_size = 8; +// 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; + +bool supports_managed_memory() { + static bool managed_memory = []() { + int device_count = gpu::device_count(); + for (int i = 0; i < device_count; ++i) { + auto& d = cu::device(i); + if (!d.managed_memory()) { + return false; + } +#if defined(_WIN32) + // Empirically on Windows if there is no concurrentManagedAccess the + // managed memory also does not work. + if (!d.concurrent_managed_access()) { + return false; + } +#endif + } + return true; + }(); + return managed_memory; +} + +inline void* unified_malloc(size_t size) { + void* data = nullptr; + if (supports_managed_memory()) { + CHECK_CUDA_ERROR(cudaMallocManaged(&data, size)); + } else { + CHECK_CUDA_ERROR(cudaMallocHost(&data, size)); + } + return data; +} + +inline void unified_free(void* data) { + if (supports_managed_memory()) { + CHECK_CUDA_ERROR(cudaFree(data)); + } else { + CHECK_CUDA_ERROR(cudaFreeHost(data)); + } +} + #if CUDART_VERSION >= 13000 inline cudaMemLocation cuda_mem_loc(int i) { cudaMemLocation loc; @@ -35,24 +78,20 @@ inline int cuda_mem_loc(int 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; - SmallSizePool::SmallSizePool() { auto num_blocks = small_pool_size / small_block_size; buffer_ = new Block[num_blocks]; - next_free_ = buffer_; - CHECK_CUDA_ERROR(cudaMallocManaged(&data_, small_pool_size)); - - int device_count = gpu::device_count(); - for (int i = 0; i < device_count; ++i) { - if (cu::device(i).concurrent_managed_access()) { - auto loc = cuda_mem_loc(i); - CHECK_CUDA_ERROR(cudaMemAdvise( - data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc)); + data_ = unified_malloc(small_pool_size); + if (supports_managed_memory()) { + int device_count = gpu::device_count(); + for (int i = 0; i < device_count; ++i) { + if (device(i).concurrent_managed_access()) { + auto loc = cuda_mem_loc(i); + CHECK_CUDA_ERROR(cudaMemAdvise( + data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc)); + } } } @@ -65,7 +104,7 @@ SmallSizePool::SmallSizePool() { } SmallSizePool::~SmallSizePool() { - CHECK_CUDA_ERROR(cudaFree(data_)); + unified_free(data_); delete[] buffer_; } @@ -99,39 +138,23 @@ CudaAllocator::CudaAllocator() : buffer_cache_( page_size, [](CudaBuffer* buf) { return buf->size; }, - [this](CudaBuffer* buf) { cuda_free(buf); }) { + [this](CudaBuffer* buf) { free_cuda_buffer(buf); }) { 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; - CHECK_CUDA_ERROR(cudaGetDeviceCount(&device_count)); - int curr; - CHECK_CUDA_ERROR(cudaGetDevice(&curr)); + int device_count = gpu::device_count(); + free_streams_.resize(device_count); + mem_pools_.resize(device_count); for (int i = 0; i < device_count; ++i) { - CHECK_CUDA_ERROR(cudaSetDevice(i)); - 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); + auto& d = device(i); + if (d.memory_pools()) { + free_streams_[i] = CudaStream(d); + CHECK_CUDA_ERROR(cudaDeviceGetDefaultMemPool(&mem_pools_[i], i)); + } } - CHECK_CUDA_ERROR(cudaSetDevice(curr)); -} - -void copy_to_managed(CudaBuffer& buf) { - // TODO maybe make this async on a i/o stream to avoid synchronizing the - // device on malloc/and free - void* new_data; - CHECK_CUDA_ERROR(cudaMallocManaged(&new_data, buf.size)); - buf.device = -1; - CHECK_CUDA_ERROR(cudaMemcpy(new_data, buf.data, buf.size, cudaMemcpyDefault)); - CHECK_CUDA_ERROR(cudaFree(buf.data)); - buf.data = new_data; } Buffer @@ -140,8 +163,6 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) { return Buffer{new CudaBuffer{nullptr, 0, -1}}; } - // Find available buffer from cache. - std::unique_lock lock(mutex_); if (size <= small_block_size) { size = 8; } else if (size < page_size) { @@ -154,6 +175,8 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) { device = -1; } + // Find available buffer from cache. + std::unique_lock lock(mutex_); CudaBuffer* buf = buffer_cache_.reuse_from_cache(size); if (!buf) { // If we have a lot of memory pressure try to reclaim memory from the cache. @@ -171,9 +194,13 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) { if (!buf) { void* data = nullptr; if (device == -1) { - CHECK_CUDA_ERROR(cudaMallocManaged(&data, size)); + data = unified_malloc(size); } else { - CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream)); + if (free_streams_[device]) { // supports memory pools + CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream)); + } else { + CHECK_CUDA_ERROR(cudaMalloc(&data, size)); + } } if (!data) { std::ostringstream msg; @@ -189,12 +216,14 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) { // 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; + if (p) { + 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; + } } } } @@ -206,9 +235,10 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) { if (get_cache_memory() > max_pool_size_) { buffer_cache_.release_cached_buffers(get_cache_memory() - max_pool_size_); } - // Copy to managed here if the buffer is not on the right device + lock.unlock(); + // Copy to unified memory here if the buffer is not on the right device. if (buf->device >= 0 && buf->device != device) { - copy_to_managed(*buf); + move_to_unified_memory(*buf, stream); } return Buffer{buf}; } @@ -232,7 +262,7 @@ void CudaAllocator::free(Buffer buffer) { if (get_cache_memory() < max_pool_size_) { buffer_cache_.recycle_to_cache(buf); } else { - cuda_free(buf); + free_cuda_buffer(buf); } } @@ -244,20 +274,48 @@ size_t CudaAllocator::size(Buffer buffer) const { return buf->size; } +void CudaAllocator::move_to_unified_memory( + CudaBuffer& buf, + cudaStream_t stream) { + if (buf.device == -1) { + return; + } + void* data = unified_malloc(buf.size); + cudaMemcpyKind kind = + supports_managed_memory() ? cudaMemcpyDefault : cudaMemcpyDeviceToHost; + if (stream) { + CHECK_CUDA_ERROR(cudaMemcpyAsync(data, buf.data, buf.size, kind, stream)); + } else { + CHECK_CUDA_ERROR(cudaMemcpy(data, buf.data, buf.size, kind)); + } + cuda_free(buf); + buf.data = data; + buf.device = -1; +} + // This must be called with mutex_ aquired -void CudaAllocator::cuda_free(CudaBuffer* buf) { +void CudaAllocator::free_cuda_buffer(CudaBuffer* buf) { if (scalar_pool_.in_pool(buf)) { scalar_pool_.free(buf); } else { - if (buf->device >= 0) { - CHECK_CUDA_ERROR(cudaFreeAsync(buf->data, free_streams_[buf->device])); - } else { - CHECK_CUDA_ERROR(cudaFree(buf->data)); - } + cuda_free(*buf); delete buf; } } +void CudaAllocator::cuda_free(CudaBuffer& buf) { + if (buf.device == -1) { + unified_free(buf.data); + } else { + cudaStream_t stream = free_streams_[buf.device]; + if (stream) { + CHECK_CUDA_ERROR(cudaFreeAsync(buf.data, stream)); + } else { + CHECK_CUDA_ERROR(cudaFree(buf.data)); + } + } +} + size_t CudaAllocator::get_active_memory() const { return active_memory_; } @@ -309,14 +367,8 @@ CudaAllocator& allocator() { } Buffer malloc_async(size_t size, CommandEncoder& encoder) { - auto buffer = allocator().malloc_async( + return allocator().malloc_async( size, encoder.device().cuda_device(), encoder.stream()); - if (size && !buffer.ptr()) { - std::ostringstream msg; - msg << "[malloc_async] Unable to allocate " << size << " bytes."; - throw std::runtime_error(msg.str()); - } - return buffer; } } // namespace cu @@ -332,9 +384,7 @@ void* Buffer::raw_ptr() { return nullptr; } auto& cbuf = *static_cast(ptr_); - if (cbuf.device != -1) { - copy_to_managed(cbuf); - } + cu::allocator().move_to_unified_memory(cbuf); return cbuf.data; } diff --git a/mlx/backend/cuda/allocator.h b/mlx/backend/cuda/allocator.h index 7f6ad525..4bbc8429 100644 --- a/mlx/backend/cuda/allocator.h +++ b/mlx/backend/cuda/allocator.h @@ -54,6 +54,10 @@ class CudaAllocator : public allocator::Allocator { void free(Buffer buffer) override; size_t size(Buffer buffer) const override; + // Replace the memory of |buf| with unified memory (managed memory or pinned + // host memory), and copy the data over. Pass |stream| to copy asynchronously. + void move_to_unified_memory(CudaBuffer& buf, cudaStream_t stream = nullptr); + size_t get_active_memory() const; size_t get_peak_memory() const; void reset_peak_memory(); @@ -64,7 +68,8 @@ class CudaAllocator : public allocator::Allocator { void clear_cache(); private: - void cuda_free(CudaBuffer* buf); + void free_cuda_buffer(CudaBuffer* buf); + void cuda_free(CudaBuffer& buf); CudaAllocator(); friend CudaAllocator& allocator(); @@ -77,7 +82,7 @@ class CudaAllocator : public allocator::Allocator { BufferCache buffer_cache_; size_t active_memory_{0}; size_t peak_memory_{0}; - std::vector free_streams_; + std::vector free_streams_; std::vector mem_pools_; SmallSizePool scalar_pool_; }; diff --git a/mlx/backend/cuda/cuda_utils.h b/mlx/backend/cuda/cuda_utils.h index d9d55763..4c60fec2 100644 --- a/mlx/backend/cuda/cuda_utils.h +++ b/mlx/backend/cuda/cuda_utils.h @@ -83,6 +83,7 @@ class CudaGraphExec : public CudaHandle { class CudaStream : public CudaHandle { public: + using CudaHandle::CudaHandle; explicit CudaStream(cu::Device& device); }; diff --git a/mlx/backend/cuda/fence.cpp b/mlx/backend/cuda/fence.cpp index a50c9111..dff5f744 100644 --- a/mlx/backend/cuda/fence.cpp +++ b/mlx/backend/cuda/fence.cpp @@ -29,15 +29,9 @@ void Fence::update(Stream s, const array& a, bool cross_device) { auto& cbuf = *static_cast(const_cast(a).buffer().ptr()); if (cbuf.device != -1) { - void* new_data; - CHECK_CUDA_ERROR(cudaMallocManaged(&new_data, cbuf.size)); - cbuf.device = -1; - auto& encoder = cu::device(s.device).get_command_encoder(s); + auto& encoder = cu::get_command_encoder(s); encoder.commit(); - CHECK_CUDA_ERROR(cudaMemcpyAsync( - new_data, cbuf.data, cbuf.size, cudaMemcpyDefault, encoder.stream())); - CHECK_CUDA_ERROR(cudaFreeAsync(cbuf.data, encoder.stream())); - cbuf.data = new_data; + cu::allocator().move_to_unified_memory(cbuf, encoder.stream()); } } fence->count++;