From 2ac18eddb94d35ae94d701f259a1a824911f2f7b Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 28 Jan 2026 10:43:22 +0900 Subject: [PATCH] [CUDA] Fallback Event impl when there is no hardware cpu/gpu coherency (#3070) --- mlx/backend/cuda/device.cpp | 6 ++ mlx/backend/cuda/device.h | 12 ++++ mlx/backend/cuda/event.cu | 136 ++++++++++++++++++++++++------------ mlx/backend/cuda/event.h | 27 +++---- 4 files changed, 124 insertions(+), 57 deletions(-) diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp index 8735cdbe..12033d94 100644 --- a/mlx/backend/cuda/device.cpp +++ b/mlx/backend/cuda/device.cpp @@ -42,6 +42,12 @@ Device::Device(int device) : device_(device) { &concurrent_managed_access_, cudaDevAttrConcurrentManagedAccess, device_)); + CHECK_CUDA_ERROR(cudaDeviceGetAttribute( + &host_native_atomic_, cudaDevAttrHostNativeAtomicSupported, device_)); + CHECK_CUDA_ERROR(cudaDeviceGetAttribute( + &managed_memory_, cudaDevAttrManagedMemory, device_)); + CHECK_CUDA_ERROR(cudaDeviceGetAttribute( + &memory_pools_, cudaDevAttrMemoryPoolsSupported, device_)); } Device::~Device() { diff --git a/mlx/backend/cuda/device.h b/mlx/backend/cuda/device.h index 41d6cf8f..527827af 100644 --- a/mlx/backend/cuda/device.h +++ b/mlx/backend/cuda/device.h @@ -164,12 +164,24 @@ class Device { bool concurrent_managed_access() const { return concurrent_managed_access_ == 1; } + bool host_native_atomic() const { + return host_native_atomic_ == 1; + } + bool managed_memory() const { + return managed_memory_ == 1; + } + bool memory_pools() const { + return memory_pools_ == 1; + } private: int device_; int compute_capability_major_; int compute_capability_minor_; int concurrent_managed_access_; + int host_native_atomic_; + int managed_memory_; + int memory_pools_; std::string device_name_; cublasLtHandle_t cublaslt_handle_{nullptr}; cudnnHandle_t cudnn_handle_{nullptr}; diff --git a/mlx/backend/cuda/event.cu b/mlx/backend/cuda/event.cu index e18a5d70..7334c17a 100644 --- a/mlx/backend/cuda/event.cu +++ b/mlx/backend/cuda/event.cu @@ -174,62 +174,94 @@ class CopyableCudaEvent { // AtomicEvent implementations /////////////////////////////////////////////////////////////////////////////// -__host__ __device__ void event_wait(AtomicEvent::Atomic* ac, uint64_t value) { - uint64_t current; - while ((current = ac->load()) < value) { - ac->wait(current); +__host__ __device__ void event_wait(uint32_t* ptr, uint32_t value) { + cuda::atomic_ref ac(*ptr); + uint32_t current; + while ((current = ac.load()) < value) { + ac.wait(current); } } -__host__ __device__ void event_signal(AtomicEvent::Atomic* ac, uint64_t value) { - ac->store(value); - ac->notify_all(); +__host__ __device__ void event_signal(uint32_t* ptr, uint32_t value) { + cuda::atomic_ref ac(*ptr); + ac.store(value); + ac.notify_all(); } -__global__ void event_wait_kernel(AtomicEvent::Atomic* ac, uint64_t value) { - event_wait(ac, value); +__global__ void event_wait_kernel(uint32_t* ptr, uint32_t value) { + event_wait(ptr, value); } -__global__ void event_signal_kernel(AtomicEvent::Atomic* ac, uint64_t value) { - event_signal(ac, value); +__global__ void event_signal_kernel(uint32_t* ptr, uint32_t value) { + __threadfence_system(); + event_signal(ptr, value); + __threadfence_system(); } -bool supports_concurrent_managed_access() { - static bool concurrent_managed_access = []() { +auto check_gpu_coherency() { + static auto coherency = []() { int device_count = gpu::device_count(); + bool concurrent_managed_access = true; + bool host_native_atomic = true; for (int i = 0; i < device_count; ++i) { - if (!cu::device(i).concurrent_managed_access()) { - return false; - } + auto& d = cu::device(i); + concurrent_managed_access &= d.concurrent_managed_access(); + host_native_atomic &= d.host_native_atomic(); } - return true; + return std::make_tuple(concurrent_managed_access, host_native_atomic); }(); - return concurrent_managed_access; + return coherency; } AtomicEvent::AtomicEvent() { - if (!supports_concurrent_managed_access()) { - throw std::runtime_error( - "Device does not support synchronization in managed memory."); + void* buf; + cudaError_t (*cuda_free)(void*); + // There are 3 kinds of systems we are implementing for: + // 1. concurrentManagedAccess == true + // => use cuda::atom_ref on managed memory + // 2. hostNativeAtomicSupported == true + // => use cuda::atom_ref on pinned host memory + // 2. no hardware cpu/gpu coherency + // => use cuda::atom_ref on device memory + auto [concurrent_managed_access, host_native_atomic] = check_gpu_coherency(); + if (concurrent_managed_access) { + CHECK_CUDA_ERROR(cudaMallocManaged(&buf, sizeof(uint32_t))); + cuda_free = cudaFree; + coherent_ = true; + } else if (host_native_atomic) { + CHECK_CUDA_ERROR(cudaMallocHost(&buf, sizeof(uint32_t))); + cuda_free = cudaFreeHost; + coherent_ = true; + } else { + CHECK_CUDA_ERROR(cudaMalloc(&buf, sizeof(uint32_t))); + cuda_free = cudaFree; + coherent_ = false; + } + buf_ = std::shared_ptr( + buf, [cuda_free](void* buf) { CHECK_CUDA_ERROR(cuda_free(buf)); }); + if (coherent_) { + *ptr() = 0; + } else { + CHECK_CUDA_ERROR(cudaMemset(buf, 0, sizeof(uint32_t))); } - buf_ = std::shared_ptr( - new Buffer{allocator().malloc(sizeof(Atomic))}, [](Buffer* ptr) { - allocator().free(*ptr); - delete ptr; - }); - *static_cast(buf_->raw_ptr()) = 0; } -void AtomicEvent::wait(uint64_t value) { +void AtomicEvent::wait(uint32_t value) { nvtx3::scoped_range r("cu::AtomicEvent::wait"); - event_wait(atomic(), value); + if (coherent_) { + event_wait(ptr(), value); + } else { + while (!is_signaled(value)) { + std::this_thread::yield(); + } + } } -void AtomicEvent::wait(cudaStream_t stream, uint64_t value) { - event_wait_kernel<<<1, 1, 0, stream>>>(atomic(), value); +void AtomicEvent::wait(cudaStream_t stream, uint32_t value) { + event_wait_kernel<<<1, 1, 0, stream>>>(ptr(), value); } -void AtomicEvent::wait(Stream s, uint64_t value) { +void AtomicEvent::wait(Stream s, uint32_t value) { nvtx3::scoped_range r("cu::AtomicEvent::wait(s)"); if (s.device == mlx::core::Device::cpu) { scheduler::enqueue(s, [*this, value]() mutable { wait(value); }); @@ -241,22 +273,26 @@ void AtomicEvent::wait(Stream s, uint64_t value) { } } -void AtomicEvent::signal(uint64_t value) { +void AtomicEvent::signal(uint32_t value) { nvtx3::scoped_range r("cu::AtomicEvent::signal"); - event_signal(atomic(), value); + if (coherent_) { + event_signal(ptr(), value); + } else { + signal(signal_stream(), value); + } } -void AtomicEvent::signal(cudaStream_t stream, uint64_t value) { - event_signal_kernel<<<1, 1, 0, stream>>>(atomic(), value); +void AtomicEvent::signal(cudaStream_t stream, uint32_t value) { + event_signal_kernel<<<1, 1, 0, stream>>>(ptr(), value); } -void AtomicEvent::signal(Stream s, uint64_t value) { +void AtomicEvent::signal(Stream s, uint32_t value) { nvtx3::scoped_range r("cu::AtomicEvent::signal(s)"); if (s.device == mlx::core::Device::cpu) { // Signal through a GPU stream so the atomic is updated in GPU - updating // the atomic in CPU sometimes does not get GPU notified. - static CudaStream stream(device(mlx::core::Device::gpu)); - scheduler::enqueue(s, [*this, value]() mutable { signal(stream, value); }); + scheduler::enqueue( + s, [*this, value]() mutable { signal(signal_stream(), value); }); } else { auto& encoder = get_command_encoder(s); encoder.commit(); @@ -265,14 +301,26 @@ void AtomicEvent::signal(Stream s, uint64_t value) { } } -bool AtomicEvent::is_signaled(uint64_t value) const { - nvtx3::scoped_range r("cu::AtomicEvent::is_signaled"); - return atomic()->load() >= value; +bool AtomicEvent::is_signaled(uint32_t val) const { + return value() >= val; } -uint64_t AtomicEvent::value() const { +uint32_t AtomicEvent::value() const { nvtx3::scoped_range r("cu::AtomicEvent::value"); - return atomic()->load(); + if (coherent_) { + cuda::atomic_ref ac(*ptr()); + return ac.load(); + } else { + uint32_t val; + CHECK_CUDA_ERROR( + cudaMemcpy(&val, ptr(), sizeof(uint32_t), cudaMemcpyDeviceToHost)); + return val; + } +} + +const CudaStream& AtomicEvent::signal_stream() { + static CudaStream stream(device(0)); + return stream; } } // namespace cu diff --git a/mlx/backend/cuda/event.h b/mlx/backend/cuda/event.h index 342e6ae3..68aa4d83 100644 --- a/mlx/backend/cuda/event.h +++ b/mlx/backend/cuda/event.h @@ -54,25 +54,26 @@ class CudaEvent { // CudaEvent so the latter should always be preferred when possible. class AtomicEvent { public: - using Atomic = cuda::atomic; - AtomicEvent(); - void wait(uint64_t value); - void wait(cudaStream_t stream, uint64_t value); - void wait(Stream s, uint64_t value); - void signal(uint64_t value); - void signal(cudaStream_t stream, uint64_t value); - void signal(Stream s, uint64_t value); - bool is_signaled(uint64_t value) const; - uint64_t value() const; + void wait(uint32_t value); + void wait(cudaStream_t stream, uint32_t value); + void wait(Stream s, uint32_t value); + void signal(uint32_t value); + void signal(cudaStream_t stream, uint32_t value); + void signal(Stream s, uint32_t value); + bool is_signaled(uint32_t value) const; + uint32_t value() const; private: - Atomic* atomic() const { - return static_cast(buf_->raw_ptr()); + const CudaStream& signal_stream(); + + uint32_t* ptr() const { + return static_cast(buf_.get()); } - std::shared_ptr buf_; + bool coherent_; + std::shared_ptr buf_; }; } // namespace mlx::core::cu