From ce4d0a62ef7b3bc102c89c56dbc8c8d324cbb4ff Mon Sep 17 00:00:00 2001 From: Cheng Date: Tue, 27 Jan 2026 11:19:20 +0900 Subject: [PATCH] Do not require ConcurrentManagedAccess when not used (#3062) --- mlx/backend/cuda/allocator.cpp | 25 ++++--- mlx/backend/cuda/conv.cpp | 2 +- mlx/backend/cuda/cublas_utils.cpp | 2 +- mlx/backend/cuda/device.cpp | 71 +++++++++++-------- mlx/backend/cuda/device.h | 18 ++--- mlx/backend/cuda/event.cu | 18 +++++ .../cuda/scaled_dot_product_attention.cpp | 4 +- 7 files changed, 91 insertions(+), 49 deletions(-) diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp index d767654f..0570ed40 100644 --- a/mlx/backend/cuda/allocator.cpp +++ b/mlx/backend/cuda/allocator.cpp @@ -3,7 +3,9 @@ #include "mlx/backend/cuda/allocator.h" #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/utils.h" +#include "mlx/backend/gpu/device_info.h" #include "mlx/memory.h" +#include "mlx/scheduler.h" #include "mlx/utils.h" #include @@ -45,12 +47,13 @@ SmallSizePool::SmallSizePool() { CHECK_CUDA_ERROR(cudaMallocManaged(&data_, small_pool_size)); - int device_count = 0; - CHECK_CUDA_ERROR(cudaGetDeviceCount(&device_count)); + int device_count = gpu::device_count(); for (int i = 0; i < device_count; ++i) { - auto loc = cuda_mem_loc(i); - CHECK_CUDA_ERROR( - cudaMemAdvise(data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc)); + if (cu::device(i).concurrent_managed_access()) { + auto loc = cuda_mem_loc(i); + CHECK_CUDA_ERROR(cudaMemAdvise( + data_, small_pool_size, cudaMemAdviseSetAccessedBy, loc)); + } } auto curr = next_free_; @@ -294,10 +297,14 @@ void CudaAllocator::clear_cache() { } CudaAllocator& allocator() { - // By creating the |allocator_| on heap, the destructor of CudaAllocator - // will not be called on exit and buffers in the cache will be leaked. This - // can save some time at program exit. - static CudaAllocator* allocator_ = new CudaAllocator; + static auto* allocator_ = []() { + // Ensure scheduler is created before allocator. + scheduler::scheduler(); + // By creating the |allocator_| on heap, the destructor of CudaAllocator + // will not be called on exit and buffers in the cache will be leaked. This + // can save some time at program exit. + return new CudaAllocator(); + }(); return *allocator_; } diff --git a/mlx/backend/cuda/conv.cpp b/mlx/backend/cuda/conv.cpp index 6d01da15..3b86bf28 100644 --- a/mlx/backend/cuda/conv.cpp +++ b/mlx/backend/cuda/conv.cpp @@ -103,7 +103,7 @@ std::optional build_conv_graph( const std::vector& dilation) { auto compute_dtype = (dtype == float16 || dtype == bfloat16) ? float32 : dtype; - DnnGraph graph(encoder.device().cudnn_handle(), dtype, compute_dtype); + DnnGraph graph(encoder.device().get_cudnn_handle(), dtype, compute_dtype); auto x_ = graph.tensor_nchw("X", 'x', x); auto w_ = graph.tensor_nchw("W", 'w', w); diff --git a/mlx/backend/cuda/cublas_utils.cpp b/mlx/backend/cuda/cublas_utils.cpp index 108f56c8..1176bd49 100644 --- a/mlx/backend/cuda/cublas_utils.cpp +++ b/mlx/backend/cuda/cublas_utils.cpp @@ -98,7 +98,7 @@ void CublasMatmulBase::init_base( M_ = a_rows; N_ = b_cols; scale_type_ = scale_type; - handle_ = device.lt_handle(); + handle_ = device.get_cublaslt_handle(); pref_ = cublas_utils::get_preference(device); heuristic_.state = CUBLAS_STATUS_NOT_INITIALIZED; diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp index 10a647b7..8735cdbe 100644 --- a/mlx/backend/cuda/device.cpp +++ b/mlx/backend/cuda/device.cpp @@ -3,6 +3,7 @@ #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/jit_module.h" #include "mlx/backend/cuda/worker.h" +#include "mlx/backend/gpu/device_info.h" #include "mlx/utils.h" #include @@ -37,31 +38,19 @@ Device::Device(int device) : device_(device) { &compute_capability_major_, cudaDevAttrComputeCapabilityMajor, device_)); CHECK_CUDA_ERROR(cudaDeviceGetAttribute( &compute_capability_minor_, cudaDevAttrComputeCapabilityMinor, device_)); - // Validate the requirements of device. - int attr = 0; CHECK_CUDA_ERROR(cudaDeviceGetAttribute( - &attr, cudaDevAttrConcurrentManagedAccess, device_)); - if (attr != 1) { - throw std::runtime_error( - fmt::format( - "Device {} does not support synchronization in managed memory.", - device_)); - } - - // The cublasLt handle is used by matmul. - make_current(); - CHECK_CUBLAS_ERROR(cublasLtCreate(<_)); - // The cudnn handle is used by Convolution. - CHECK_CUDNN_ERROR(cudnnCreate(&cudnn_)); - - // Initialize the jit module cache here ensures it is not - // unloaded before any evaluation is done - get_jit_module_cache(); + &concurrent_managed_access_, + cudaDevAttrConcurrentManagedAccess, + device_)); } Device::~Device() { - CHECK_CUDNN_ERROR(cudnnDestroy(cudnn_)); - CHECK_CUBLAS_ERROR(cublasLtDestroy(lt_)); + if (cudnn_handle_) { + CHECK_CUDNN_ERROR(cudnnDestroy(cudnn_handle_)); + } + if (cublaslt_handle_) { + CHECK_CUBLAS_ERROR(cublasLtDestroy(cublaslt_handle_)); + } } void Device::make_current() { @@ -82,6 +71,22 @@ CommandEncoder& Device::get_command_encoder(Stream s) { return it->second; } +cublasLtHandle_t Device::get_cublaslt_handle() { + if (!cublaslt_handle_) { + make_current(); + CHECK_CUBLAS_ERROR(cublasLtCreate(&cublaslt_handle_)); + } + return cublaslt_handle_; +} + +cudnnHandle_t Device::get_cudnn_handle() { + if (!cudnn_handle_) { + make_current(); + CHECK_CUDNN_ERROR(cudnnCreate(&cudnn_handle_)); + } + return cudnn_handle_; +} + CommandEncoder::CaptureContext::CaptureContext(CommandEncoder& enc) : enc(enc) { enc.device().make_current(); if (!use_cuda_graphs()) { @@ -491,13 +496,23 @@ void CommandEncoder::synchronize() { f.wait(); } -Device& device(mlx::core::Device device) { - static std::unordered_map devices; - auto it = devices.find(device.index); - if (it == devices.end()) { - it = devices.try_emplace(device.index, device.index).first; - } - return it->second; +Device& device(int cuda_device) { + static auto devices = []() { + std::vector devices; + int device_count = gpu::device_count(); + for (int i = 0; i < device_count; ++i) { + devices.emplace_back(i); + } + // Initialize the jit module cache here ensures it is not unloaded before + // any evaluation is done. + get_jit_module_cache(); + return devices; + }(); + return devices.at(cuda_device); +} + +Device& device(mlx::core::Device d) { + return device(d.index); } CommandEncoder& get_command_encoder(Stream s) { diff --git a/mlx/backend/cuda/device.h b/mlx/backend/cuda/device.h index 1c008b4e..81fd1d22 100644 --- a/mlx/backend/cuda/device.h +++ b/mlx/backend/cuda/device.h @@ -142,6 +142,7 @@ class Device { explicit Device(int device); ~Device(); + Device(Device&&) = default; Device(const Device&) = delete; Device& operator=(const Device&) = delete; @@ -149,6 +150,8 @@ class Device { void make_current(); CommandEncoder& get_command_encoder(Stream s); + cublasLtHandle_t get_cublaslt_handle(); + cudnnHandle_t get_cudnn_handle(); int cuda_device() const { return device_; @@ -159,24 +162,23 @@ class Device { int compute_capability_minor() const { return compute_capability_minor_; } - cublasLtHandle_t lt_handle() const { - return lt_; - } - cudnnHandle_t cudnn_handle() const { - return cudnn_; + bool concurrent_managed_access() const { + return concurrent_managed_access_ == 1; } private: int device_; int compute_capability_major_; int compute_capability_minor_; + int concurrent_managed_access_; std::string device_name_; - cublasLtHandle_t lt_; - cudnnHandle_t cudnn_; + cublasLtHandle_t cublaslt_handle_{nullptr}; + cudnnHandle_t cudnn_handle_{nullptr}; std::unordered_map encoders_; }; -Device& device(mlx::core::Device device); +Device& device(int cuda_device); +Device& device(mlx::core::Device d); CommandEncoder& get_command_encoder(Stream s); // Return an execution policy that does not sync for result. diff --git a/mlx/backend/cuda/event.cu b/mlx/backend/cuda/event.cu index 6880a904..e18a5d70 100644 --- a/mlx/backend/cuda/event.cu +++ b/mlx/backend/cuda/event.cu @@ -3,6 +3,7 @@ #include "mlx/backend/cuda/allocator.h" #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/event.h" +#include "mlx/backend/gpu/device_info.h" #include "mlx/event.h" #include "mlx/scheduler.h" @@ -193,7 +194,24 @@ __global__ void event_signal_kernel(AtomicEvent::Atomic* ac, uint64_t value) { event_signal(ac, value); } +bool supports_concurrent_managed_access() { + static bool concurrent_managed_access = []() { + int device_count = gpu::device_count(); + for (int i = 0; i < device_count; ++i) { + if (!cu::device(i).concurrent_managed_access()) { + return false; + } + } + return true; + }(); + return concurrent_managed_access; +} + AtomicEvent::AtomicEvent() { + if (!supports_concurrent_managed_access()) { + throw std::runtime_error( + "Device does not support synchronization in managed memory."); + } buf_ = std::shared_ptr( new Buffer{allocator().malloc(sizeof(Atomic))}, [](Buffer* ptr) { allocator().free(*ptr); diff --git a/mlx/backend/cuda/scaled_dot_product_attention.cpp b/mlx/backend/cuda/scaled_dot_product_attention.cpp index 108b0244..f20bd8c2 100644 --- a/mlx/backend/cuda/scaled_dot_product_attention.cpp +++ b/mlx/backend/cuda/scaled_dot_product_attention.cpp @@ -269,7 +269,7 @@ void sdpa_cudnn( bool output_logsumexp, Stream s) { auto& encoder = cu::get_command_encoder(s); - auto handle = encoder.device().cudnn_handle(); + auto handle = encoder.device().get_cudnn_handle(); malloc_with_same_layout(encoder, o, q); @@ -327,7 +327,7 @@ void sdpa_backward_cudnn( array& d_v, Stream s) { auto& encoder = cu::get_command_encoder(s); - auto handle = encoder.device().cudnn_handle(); + auto handle = encoder.device().get_cudnn_handle(); malloc_with_same_layout(encoder, d_q, q); malloc_with_same_layout(encoder, d_k, k);