[CUDA] Thread safety (#3367)

mlx/backend/cuda/conv.cpp

@@ -39,7 +39,7 @@ struct ConvCacheKey {
 };
 
 auto& conv_cache() {
-  static LRUBytesKeyCache<
+  static thread_local LRUBytesKeyCache<
       ConvCacheKey,
       std::pair<ConvBackendType, std::optional<DnnGraph>>>
       cache("MLX_CUDA_CONV_CACHE_SIZE", /* default_capacity */ 128);
@@ -103,7 +103,7 @@ std::optional<DnnGraph> build_conv_graph(
     const std::vector<int64_t>& dilation) {
   auto compute_dtype =
       (dtype == float16 || dtype == bfloat16) ? float32 : dtype;
-  DnnGraph graph(encoder.device().get_cudnn_handle(), dtype, compute_dtype);
+  DnnGraph graph(get_cudnn_handle(encoder.device()), dtype, compute_dtype);
   auto x_ = graph.tensor_nchw("X", 'x', x);
   auto w_ = graph.tensor_nchw("W", 'w', w);
 
@@ -139,7 +139,7 @@ std::optional<DnnGraph> build_conv_graph(
   if (dtype == float32 && !env::enable_tf32()) {
     graph.deselect_numeric_notes({fe::NumericalNote_t::TENSOR_CORE});
   }
-  CHECK_CUDNN_FE_ERROR(graph.build());
+  CHECK_CUDNN_ERROR(graph.build());
   return graph;
 }
 
@@ -252,6 +252,10 @@ void register_args(
 
 } // namespace
 
+void init_cudnn_conv_cache() {
+  conv_cache();
+}
+
 void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
   nvtx3::scoped_range r("Convolution::eval_gpu");
   if (out_.size() == 0) {
@@ -289,7 +293,7 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
       std::tie(in, wt, out) =
           prepare_args(encoder, backend_type, in, wt, out, groups_, s);
       register_args(encoder, backend_type, in, wt, out, out_);
-      CHECK_CUDNN_FE_ERROR(graph->encode_capturing(
+      CHECK_CUDNN_ERROR(graph->encode_capturing(
           encoder,
           {
               {'x', gpu_ptr<void>(in)},
@@ -371,7 +375,7 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
 
   if (graph) {
     register_args(encoder, backend_type, in, wt, out, out_);
-    CHECK_CUDNN_FE_ERROR(graph->encode_capturing(
+    CHECK_CUDNN_ERROR(graph->encode_capturing(
         encoder,
         {
             {'x', gpu_ptr<void>(in)},

mlx/backend/cuda/cublas_utils.cpp

@@ -2,44 +2,13 @@
 
 #include "mlx/backend/cuda/cublas_utils.h"
 #include "mlx/backend/cuda/cuda.h"
+#include "mlx/backend/gpu/device_info.h"
 #include "mlx/utils.h"
 
 namespace mlx::core {
+
 namespace cublas_utils {
 
-namespace {
-
-struct CublasPreference {
-  CublasPreference(cu::Device& device) {
-    // The recommended cublas workspace size is 4 MiB for pre-Hopper and 32 MiB
-    // for Hopper+:
-    // https://docs.nvidia.com/cuda/cublas/#cublassetworkspace
-    uint64_t MiB = 1024 * 1024;
-    uint64_t workspace_size =
-        device.compute_capability_major() >= 9 ? 32 * MiB : 4 * MiB;
-
-    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&pref_));
-    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceSetAttribute(
-        pref_,
-        CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
-        &workspace_size,
-        sizeof(uint64_t)));
-  }
-
-  ~CublasPreference() {
-    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceDestroy(pref_));
-  }
-
-  cublasLtMatmulPreference_t pref_{nullptr};
-};
-
-} // namespace
-
-cublasLtMatmulPreference_t get_preference(cu::Device& device) {
-  static CublasPreference pref(device);
-  return pref.pref_;
-}
-
 cublasLtMatrixLayout_t create_matrix_layout(
     cudaDataType_t type,
     uint64_t rows,
@@ -70,6 +39,59 @@ cublasLtMatrixLayout_t create_matrix_layout(
 
 } // namespace cublas_utils
 
+namespace {
+
+auto& cublas_handles_cache() {
+  struct CublasHandles {
+    ~CublasHandles() {
+      if (handle) {
+        CHECK_CUBLAS_ERROR(cublasLtDestroy(handle));
+        CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceDestroy(pref));
+      }
+    }
+    cublasLtHandle_t handle{nullptr};
+    cublasLtMatmulPreference_t pref{nullptr};
+  };
+  static thread_local std::vector<CublasHandles> cache(gpu::device_count());
+  return cache;
+}
+
+auto get_cublas_handles(cu::Device& device) {
+  auto& storage = cublas_handles_cache().at(device.cuda_device());
+  if (!storage.handle) {
+    // Create cublasLt handle.
+    device.make_current();
+    CHECK_CUBLAS_ERROR(cublasLtCreate(&storage.handle));
+    // The recommended cublas workspace size is 4 MiB for pre-Hopper and 32
+    // MiB for Hopper+:
+    // https://docs.nvidia.com/cuda/cublas/#cublassetworkspace
+    uint64_t MiB = 1024 * 1024;
+    uint64_t workspace_size =
+        device.compute_capability_major() >= 9 ? 32 * MiB : 4 * MiB;
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&storage.pref));
+    CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceSetAttribute(
+        storage.pref,
+        CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
+        &workspace_size,
+        sizeof(uint64_t)));
+  }
+  return std::make_tuple(storage.handle, storage.pref);
+}
+
+} // namespace
+
+void check_cublas_error(const char* name, cublasStatus_t err) {
+  if (err != CUBLAS_STATUS_SUCCESS) {
+    // TODO: Use cublasGetStatusString when it is widely available.
+    throw std::runtime_error(
+        fmt::format("{} failed with code: {}.", name, static_cast<int>(err)));
+  }
+}
+
+void init_cublas_handles_cache() {
+  cublas_handles_cache();
+}
+
 CublasMatmulBase::~CublasMatmulBase() {
   CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(a_desc_));
   CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(b_desc_));
@@ -98,8 +120,7 @@ void CublasMatmulBase::init_base(
   M_ = a_rows;
   N_ = b_cols;
   scale_type_ = scale_type;
-  handle_ = device.get_cublaslt_handle();
-  pref_ = cublas_utils::get_preference(device);
+  std::tie(handle_, pref_) = get_cublas_handles(device);
   heuristic_.state = CUBLAS_STATUS_NOT_INITIALIZED;
 
   CHECK_CUBLAS_ERROR(

mlx/backend/cuda/cublas_utils.h

@@ -1,17 +1,15 @@
 // Copyright © 2025 Apple Inc.
 #pragma once
 
-#include <cublasLt.h>
 #include "mlx/array.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/dtype_utils.h"
 
+#include <cublasLt.h>
+
 namespace mlx::core {
 namespace cublas_utils {
 
-// Get the shared cublas preference for a device
-cublasLtMatmulPreference_t get_preference(cu::Device& device);
-
 cublasLtMatrixLayout_t create_matrix_layout(
     cudaDataType_t type,
     uint64_t rows,
@@ -42,6 +40,12 @@ inline cudaDataType_t dtype_to_cublas_type(Dtype dtype, std::string_view tag) {
 
 } // namespace cublas_utils
 
+void check_cublas_error(const char* name, cublasStatus_t err);
+
+#define CHECK_CUBLAS_ERROR(cmd) check_cublas_error(#cmd, (cmd))
+
+void init_cublas_handles_cache();
+
 class CublasMatmulBase {
  public:
   virtual ~CublasMatmulBase();

mlx/backend/cuda/cuda_utils.h

@@ -2,23 +2,17 @@
 
 #pragma once
 
-#include <cublasLt.h>
 #include <cuda.h>
 #include <cuda_runtime.h>
-#include <cudnn.h>
 
 namespace mlx::core {
 
 // Throw exception if the cuda API does not succeed.
-void check_cublas_error(const char* name, cublasStatus_t err);
 void check_cuda_error(const char* name, cudaError_t err);
 void check_cuda_error(const char* name, CUresult err);
-void check_cudnn_error(const char* name, cudnnStatus_t err);
 
 // The macro version that prints the command that failed.
-#define CHECK_CUBLAS_ERROR(cmd) check_cublas_error(#cmd, (cmd))
 #define CHECK_CUDA_ERROR(cmd) check_cuda_error(#cmd, (cmd))
-#define CHECK_CUDNN_ERROR(cmd) check_cudnn_error(#cmd, (cmd))
 
 // Base class for RAII managed CUDA resources.
 template <typename Handle, cudaError_t (*Destroy)(Handle)>

mlx/backend/cuda/cudnn_utils.cpp

@@ -2,6 +2,7 @@
 
 #include "mlx/backend/cuda/cudnn_utils.h"
 #include "mlx/backend/cuda/device.h"
+#include "mlx/backend/gpu/device_info.h"
 
 namespace mlx::core {
 
@@ -47,8 +48,48 @@ inline auto nhwc_to_nchw(const array& x) {
   return std::make_tuple(std::move(shape), std::move(strides));
 }
 
+auto& cudnn_handles_cache() {
+  struct CudnnHandle {
+    ~CudnnHandle() {
+      if (handle) {
+        CHECK_CUDNN_ERROR(cudnnDestroy(handle));
+      }
+    }
+    cudnnHandle_t handle{nullptr};
+  };
+  static thread_local std::vector<CudnnHandle> cache(gpu::device_count());
+  return cache;
+}
+
 } // namespace
 
+void check_cudnn_error(const char* name, cudnnStatus_t err) {
+  if (err != CUDNN_STATUS_SUCCESS) {
+    throw std::runtime_error(
+        fmt::format("{} failed: {}.", name, cudnnGetErrorString(err)));
+  }
+}
+
+void check_cudnn_error(const char* name, fe::error_t err) {
+  if (!err.is_good()) {
+    throw std::runtime_error(
+        fmt::format("{} failed: {}.", name, err.get_message()));
+  }
+}
+
+cudnnHandle_t get_cudnn_handle(cu::Device& device) {
+  auto& storage = cudnn_handles_cache().at(device.cuda_device());
+  if (!storage.handle) {
+    device.make_current();
+    CHECK_CUDNN_ERROR(cudnnCreate(&storage.handle));
+  }
+  return storage.handle;
+}
+
+void init_cudnn_handles_cache() {
+  cudnn_handles_cache();
+}
+
 fe::error_t DnnGraph::prepare() {
   RETURN_IF_ERROR(validate());
   try {
@@ -109,7 +150,7 @@ fe::error_t DnnGraph::encode_capturing(
 
 void* DnnGraph::prepare_workspace(cu::CommandEncoder& encoder) {
   int64_t workspace_size = 0;
-  CHECK_CUDNN_FE_ERROR(get_workspace_size(workspace_size));
+  CHECK_CUDNN_ERROR(get_workspace_size(workspace_size));
   return allocate_workspace(encoder, workspace_size);
 }
 

mlx/backend/cuda/cudnn_utils.h

@@ -21,14 +21,16 @@ class CommandEncoder;
 
 namespace fe = cudnn_frontend;
 
-#define CHECK_CUDNN_FE_ERROR(cmd)                                    \
-  do {                                                               \
-    auto error = cmd;                                                \
-    if (!error.is_good()) {                                          \
-      throw std::runtime_error(                                      \
-          fmt::format("{} failed: {}.", #cmd, error.get_message())); \
-    }                                                                \
-  } while (0)
+void check_cudnn_error(const char* name, cudnnStatus_t err);
+void check_cudnn_error(const char* name, fe::error_t err);
+
+#define CHECK_CUDNN_ERROR(cmd) check_cudnn_error(#cmd, (cmd))
+
+cudnnHandle_t get_cudnn_handle(cu::Device& device);
+
+void init_cudnn_handles_cache();
+void init_cudnn_conv_cache();
+void init_cudnn_sdpa_cache();
 
 // Return pointer alignment of |x|'s data.
 inline uint8_t get_alignment(const array& x) {

mlx/backend/cuda/device.cpp

@@ -1,7 +1,6 @@
 // Copyright © 2025 Apple Inc.
 
 #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"
@@ -55,14 +54,7 @@ Device::Device(int device) : device_(device) {
       &memory_pools_, cudaDevAttrMemoryPoolsSupported, device_));
 }
 
-Device::~Device() {
-  if (cudnn_handle_) {
-    CHECK_CUDNN_ERROR(cudnnDestroy(cudnn_handle_));
-  }
-  if (cublaslt_handle_) {
-    CHECK_CUBLAS_ERROR(cublasLtDestroy(cublaslt_handle_));
-  }
-}
+Device::~Device() = default;
 
 void Device::make_current() {
   // We need to set/get current CUDA device very frequently, cache it to reduce
@@ -76,30 +68,6 @@ void Device::make_current() {
   }
 }
 
-CommandEncoder& Device::get_command_encoder(Stream s) {
-  auto it = encoders_.find(s.index);
-  if (it == encoders_.end()) {
-    it = encoders_.try_emplace(s.index, *this).first;
-  }
-  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()) {
@@ -243,6 +211,10 @@ CommandEncoder::CommandEncoder(Device& d)
   std::tie(max_ops_per_graph_, max_mb_per_graph_) = get_graph_limits(d);
 }
 
+CommandEncoder::~CommandEncoder() {
+  synchronize();
+}
+
 void CommandEncoder::add_completed_handler(std::function<void()> task) {
   worker_.add_task(std::move(task));
 }
@@ -571,18 +543,17 @@ void CommandEncoder::synchronize() {
 }
 
 Device& device(int cuda_device) {
-  static auto devices = []() {
-    std::vector<Device> devices;
+  // The devices are leak intentionally as user code may still be accessing
+  // device after main thread teardown.
+  static auto* devices = []() {
+    auto* devices = new std::vector<Device>;
     int device_count = gpu::device_count();
     for (int i = 0; i < device_count; ++i) {
-      devices.emplace_back(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);
+  return devices->at(cuda_device);
 }
 
 Device& device(mlx::core::Device d) {
@@ -590,7 +561,18 @@ Device& device(mlx::core::Device d) {
 }
 
 CommandEncoder& get_command_encoder(Stream s) {
-  return device(s.device).get_command_encoder(s);
+  auto& encoders = get_command_encoders();
+  auto it = encoders.find(s.index);
+  if (it == encoders.end()) {
+    throw std::runtime_error(
+        fmt::format("There is no Stream(gpu, {}) in current thread.", s.index));
+  }
+  return it->second;
+}
+
+std::unordered_map<int, CommandEncoder>& get_command_encoders() {
+  static thread_local std::unordered_map<int, CommandEncoder> encoders;
+  return encoders;
 }
 
 } // namespace mlx::core::cu

mlx/backend/cuda/device.h

@@ -8,10 +8,6 @@
 #include "mlx/backend/cuda/worker.h"
 #include "mlx/stream.h"
 
-#include <cublasLt.h>
-#include <cuda.h>
-#include <cudnn.h>
-
 #include <unordered_map>
 
 namespace mlx::core::cu {
@@ -35,6 +31,7 @@ class CommandEncoder {
   };
 
   explicit CommandEncoder(Device& d);
+  ~CommandEncoder();
 
   CommandEncoder(const CommandEncoder&) = delete;
   CommandEncoder& operator=(const CommandEncoder&) = delete;
@@ -170,10 +167,6 @@ class Device {
   // Make this device the current cuda device, this method is thread-safe.
   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_;
   }
@@ -205,13 +198,12 @@ class Device {
   int managed_memory_;
   int memory_pools_;
   std::string device_name_;
-  cublasLtHandle_t cublaslt_handle_{nullptr};
-  cudnnHandle_t cudnn_handle_{nullptr};
-  std::unordered_map<int, CommandEncoder> encoders_;
 };
 
-Device& device(int cuda_device);
-Device& device(mlx::core::Device d);
-CommandEncoder& get_command_encoder(Stream s);
+MLX_API Device& device(int cuda_device);
+MLX_API Device& device(mlx::core::Device d);
+MLX_API CommandEncoder& get_command_encoder(Stream s);
+
+std::unordered_map<int, CommandEncoder>& get_command_encoders();
 
 } // namespace mlx::core::cu

mlx/backend/cuda/eval.cpp

@@ -2,6 +2,8 @@
 
 #include "mlx/backend/gpu/eval.h"
 #include "mlx/backend/cuda/allocator.h"
+#include "mlx/backend/cuda/cublas_utils.h"
+#include "mlx/backend/cuda/cudnn_utils.h"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/primitives.h"
 #include "mlx/scheduler.h"
@@ -18,7 +20,16 @@ void init() {
 }
 
 void new_stream(Stream s) {
-  cu::get_command_encoder(s);
+  // Make sure the handles get destroyed after CommandEncoder.
+  init_cublas_handles_cache();
+  init_cudnn_handles_cache();
+  init_cudnn_conv_cache();
+  init_cudnn_sdpa_cache();
+  // Create CommandEncoder.
+  assert(s.device == Device::gpu);
+  auto& encoders = cu::get_command_encoders();
+  auto& d = cu::device(s.device);
+  encoders.try_emplace(s.index, d);
 }
 
 void eval(array& arr) {

mlx/backend/cuda/jit_module.cpp

@@ -9,6 +9,7 @@
 #include <cstdlib>
 #include <filesystem>
 #include <fstream>
+#include <shared_mutex>
 
 #include <fmt/format.h>
 #include <nvrtc.h>
@@ -439,20 +440,28 @@ CUfunction JitModule::get_kernel(
   return get_kernel_and_dims(kernel_name, std::move(configure_kernel)).first;
 }
 
-std::unordered_map<std::string, JitModule>& get_jit_module_cache() {
-  static std::unordered_map<std::string, JitModule> map;
-  return map;
-}
-
 JitModule& get_jit_module(
     const mlx::core::Device& device,
     const std::string& name,
     const KernelBuilder& builder,
-    bool cache) {
-  auto& map = get_jit_module_cache();
-  auto it = map.find(name);
-  if (it == map.end()) {
-    it = map.try_emplace(name, cu::device(device), name, builder, cache).first;
+    bool use_disk_cache) {
+  // The cache are leak intentionally as user code may still be running JIT
+  // compiled code after main thread teardown.
+  static auto* cache = new std::unordered_map<std::string, JitModule>;
+  static auto* mtx = new std::shared_mutex;
+
+  {
+    std::shared_lock rlock(*mtx);
+    if (auto it = cache->find(name); it != cache->end()) {
+      return it->second;
+    }
+  }
+
+  std::unique_lock wlock(*mtx);
+  auto it = cache->find(name);
+  if (it == cache->end()) {
+    auto& d = cu::device(device);
+    it = cache->try_emplace(name, d, name, builder, use_disk_cache).first;
   }
   return it->second;
 }

mlx/backend/cuda/jit_module.h

@@ -91,11 +91,12 @@ class JitModule {
       Device& device,
       const std::string& module_name,
       const KernelBuilder& builder,
-      bool cache);
+      bool use_disk_cache);
   ~JitModule();
 
   JitModule(const JitModule&) = delete;
   JitModule& operator=(const JitModule&) = delete;
+
   CUfunction get_kernel(
       const std::string& kernel_name,
       std::function<void(CUfunction)> configure_kernel = nullptr);
@@ -109,8 +110,6 @@ class JitModule {
       kernels_;
 };
 
-std::unordered_map<std::string, JitModule>& get_jit_module_cache();
-
 JitModule& get_jit_module(
     const mlx::core::Device& device,
     const std::string& name,

mlx/backend/cuda/scaled_dot_product_attention.cpp

@@ -176,13 +176,13 @@ inline BytesKey<SDPACacheKey> build_sdpa_cache_key(
 }
 
 auto& sdpa_cache() {
-  static LRUBytesKeyCache<SDPACacheKey, DnnGraph> cache(
+  static thread_local LRUBytesKeyCache<SDPACacheKey, DnnGraph> cache(
       "MLX_CUDA_SDPA_CACHE_SIZE", /* default_capacity */ 256);
   return cache;
 }
 
 auto& sdpa_backward_cache() {
-  static LRUBytesKeyCache<SDPACacheKey, DnnGraph> cache(
+  static thread_local LRUBytesKeyCache<SDPACacheKey, DnnGraph> cache(
       "MLX_CUDA_SDPA_BACKWARD_CACHE_SIZE", /* default_capacity */ 64);
   return cache;
 }
@@ -249,10 +249,10 @@ DnnGraph build_sdpa_graph(
     graph.tensor(stats_, STATS, *stats)->set_output(true);
   }
 
-  CHECK_CUDNN_FE_ERROR(graph.prepare());
+  CHECK_CUDNN_ERROR(graph.prepare());
   graph.select_behavior_notes(
       {fe::BehaviorNote_t::SUPPORTS_CUDA_GRAPH_NATIVE_API});
-  CHECK_CUDNN_FE_ERROR(graph.build());
+  CHECK_CUDNN_ERROR(graph.build());
   return graph;
 }
 
@@ -298,15 +298,20 @@ DnnGraph build_sdpa_backward_graph(
   graph.tensor(d_k_, D_K, d_k)->set_output(true);
   graph.tensor(d_v_, D_V, d_v)->set_output(true);
 
-  CHECK_CUDNN_FE_ERROR(graph.prepare());
+  CHECK_CUDNN_ERROR(graph.prepare());
   graph.select_behavior_notes(
       {fe::BehaviorNote_t::SUPPORTS_CUDA_GRAPH_NATIVE_API});
-  CHECK_CUDNN_FE_ERROR(graph.build());
+  CHECK_CUDNN_ERROR(graph.build());
   return graph;
 }
 
 } // namespace
 
+void init_cudnn_sdpa_cache() {
+  sdpa_cache();
+  sdpa_backward_cache();
+}
+
 bool supports_sdpa_cudnn(
     const array& q,
     const array& k,
@@ -357,7 +362,7 @@ void sdpa_cudnn(
     bool output_logsumexp,
     Stream s) {
   auto& encoder = cu::get_command_encoder(s);
-  auto handle = encoder.device().get_cudnn_handle();
+  auto handle = get_cudnn_handle(encoder.device());
 
   malloc_with_same_layout(encoder, o, q);
 
@@ -440,7 +445,7 @@ void sdpa_cudnn(
     variant_pack[STATS] = gpu_ptr<void>(*stats);
   }
 
-  CHECK_CUDNN_FE_ERROR(graph.encode_graph(encoder, std::move(variant_pack)));
+  CHECK_CUDNN_ERROR(graph.encode_graph(encoder, std::move(variant_pack)));
 }
 
 void sdpa_backward_cudnn(
@@ -459,7 +464,7 @@ void sdpa_backward_cudnn(
     array& d_v,
     Stream s) {
   auto& encoder = cu::get_command_encoder(s);
-  auto handle = encoder.device().get_cudnn_handle();
+  auto handle = get_cudnn_handle(encoder.device());
 
   malloc_with_same_layout(encoder, d_q, q);
   malloc_with_same_layout(encoder, d_k, k);
@@ -522,7 +527,7 @@ void sdpa_backward_cudnn(
     variant_pack[SINKS] = gpu_ptr<void>(*sinks);
   }
 
-  CHECK_CUDNN_FE_ERROR(graph.encode_graph(encoder, std::move(variant_pack)));
+  CHECK_CUDNN_ERROR(graph.encode_graph(encoder, std::move(variant_pack)));
 }
 
 // Defined in scaled_dot_product_attention.cu file.

mlx/backend/cuda/utils.cpp

@@ -10,14 +10,6 @@
 
 namespace mlx::core {
 
-void check_cublas_error(const char* name, cublasStatus_t err) {
-  if (err != CUBLAS_STATUS_SUCCESS) {
-    // TODO: Use cublasGetStatusString when it is widely available.
-    throw std::runtime_error(
-        fmt::format("{} failed with code: {}.", name, static_cast<int>(err)));
-  }
-}
-
 void check_cuda_error(const char* name, cudaError_t err) {
   if (err != cudaSuccess) {
     throw std::runtime_error(
@@ -33,13 +25,6 @@ void check_cuda_error(const char* name, CUresult err) {
   }
 }
 
-void check_cudnn_error(const char* name, cudnnStatus_t err) {
-  if (err != CUDNN_STATUS_SUCCESS) {
-    throw std::runtime_error(
-        fmt::format("{} failed: {}.", name, cudnnGetErrorString(err)));
-  }
-}
-
 const char* dtype_to_cuda_type(const Dtype& dtype) {
   switch (dtype) {
     case bool_:

tests/scheduler_tests.cpp

@@ -66,7 +66,7 @@ TEST_CASE("test default stream in threads") {
 }
 
 TEST_CASE("test access stream in other thread") {
-  if (!metal::is_available()) {
+  if (!gpu::is_available()) {
     return;
   }