Fix CUDA termination crash via deterministic worker teardown and lazy stream evaluation

mlx/backend/cuda/allocator.cpp

@@ -67,11 +67,12 @@ inline void* unified_malloc(size_t size) {
   return data;
 }
 
+// Errors are ignored since memory is freed in destructors, which can not throw.
 inline void unified_free(void* data) {
   if (supports_managed_memory()) {
-    CHECK_CUDA_ERROR(cudaFree(data));
+    cudaFree(data);
   } else {
-    CHECK_CUDA_ERROR(cudaFreeHost(data));
+    cudaFreeHost(data);
   }
 }
 
@@ -319,14 +320,15 @@ void CudaAllocator::free_async(CudaBuffer& buf, cudaStream_t stream) {
   if (buf.device == -1) {
     unified_free(buf.data);
   } else {
-    // Free asynchronously when memory pools is supported.
+    // Free asynchronously when memory pools is supported. Errors are ignored
+    // since memory is freed in destructors, which can not throw.
     if (mem_pools_[buf.device]) {
       if (!stream) {
         stream = free_streams_[buf.device];
       }
-      CHECK_CUDA_ERROR(cudaFreeAsync(buf.data, stream));
+      cudaFreeAsync(buf.data, stream);
     } else {
-      CHECK_CUDA_ERROR(cudaFree(buf.data));
+      cudaFree(buf.data);
     }
   }
 }

mlx/backend/cuda/cuda_utils.h

@@ -26,11 +26,11 @@ class CudaHandle {
   }
 
   ~CudaHandle() {
-    // Skip if there was an error to avoid throwing in the destructors
-    if (cudaPeekAtLastError() != cudaSuccess) {
-      return;
+    // Errors are ignored since destructors can not throw, and destroying can
+    // fail when the CUDA runtime is shutting down.
+    if (handle_ != nullptr) {
+      Destroy(handle_);
     }
-    reset();
   }
 
   CudaHandle(const CudaHandle&) = delete;

mlx/backend/cuda/device.cpp

@@ -206,15 +206,17 @@ CommandEncoder::CommandEncoder(Device& d)
     : device_(d),
       stream_(d),
       graph_(d),
-      worker_(std::make_shared<Worker>(d)),
+      worker_(std::make_unique<Worker>(d)),
       graph_cache_("MLX_CUDA_GRAPH_CACHE_SIZE", /* default_capacity */ 400) {
   std::tie(max_ops_per_graph_, max_mb_per_graph_) = get_graph_limits(d);
-  worker_->start();
 }
 
 CommandEncoder::~CommandEncoder() {
-  synchronize();
-  worker_->stop();
+  try {
+    synchronize();
+  } catch (...) {
+    // Synchronizing can fail when the CUDA runtime is shutting down.
+  }
 }
 
 void CommandEncoder::add_completed_handler(std::function<void()> task) {
@@ -490,6 +492,8 @@ void CommandEncoder::commit() {
 
 void CommandEncoder::synchronize() {
   CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_));
+  // Commit first so the handler below runs after all pending tasks.
+  commit();
   auto p = std::make_shared<std::promise<void>>();
   std::future<void> f = p->get_future();
   add_completed_handler([p = std::move(p)]() { p->set_value(); });

mlx/backend/cuda/device.h

@@ -146,7 +146,7 @@ class CommandEncoder {
   Device& device_;
   CudaStream stream_;
   CudaGraph graph_;
-  std::shared_ptr<Worker> worker_;
+  std::unique_ptr<Worker> worker_;
   int node_count_{0};
   bool in_concurrent_{false};
   std::vector<cudaGraphNode_t> from_nodes_;

mlx/backend/cuda/event.cu

@@ -232,8 +232,8 @@ AtomicEvent::AtomicEvent(Device& d) {
     cuda_free = cudaFree;
     coherent_ = false;
   }
-  buf_ = std::shared_ptr<void>(
-      buf, [cuda_free](void* buf) { CHECK_CUDA_ERROR(cuda_free(buf)); });
+  // Errors are ignored since the deleter may run after CUDA is shut down.
+  buf_ = std::shared_ptr<void>(buf, [cuda_free](void* buf) { cuda_free(buf); });
   if (coherent_) {
     *ptr() = 0;
   } else {

mlx/backend/cuda/worker.cpp

@@ -3,82 +3,81 @@
 #include "mlx/backend/cuda/worker.h"
 #include "mlx/backend/cuda/device.h"
 
+#include <thread>
+
 namespace mlx::core::cu {
 
 Worker::Worker(Device& d)
     : signal_stream_(d),
-      signal_event_(d, cudaEventDisableTiming | cudaEventBlockingSync) {}
-
-Worker::~Worker() = default;
-
-void Worker::start() {
-  // Note that |shared_from_this| can not be called in constructor.
-  worker_ = std::thread(&Worker::thread_fn, shared_from_this());
-  // Detach the thread and let it free itself after finishing tasks.
+      signal_event_(d, cudaEventDisableTiming | cudaEventBlockingSync),
+      state_(std::make_shared<State>()) {
+  // Detach the thread and let it free the state after finishing tasks.
   // This is to avoid deadlock when joining threads on exit on Windows:
   // https://developercommunity.visualstudio.com/t/1654756
-  worker_.detach();
+  std::thread(&State::thread_fn, state_).detach();
 }
 
-void Worker::stop() {
+Worker::~Worker() {
   {
-    std::lock_guard lock(mtx_);
-    stop_ = true;
+    std::lock_guard lock(state_->mtx);
+    state_->stop = true;
   }
-  cond_.notify_one();
+  state_->cond.notify_one();
 }
 
 void Worker::add_task(std::function<void()> task) {
   pending_tasks_.push_back(std::move(task));
 }
 
-void Worker::signal(void* data) {
-  auto w = static_cast<Worker*>(data);
-  {
-    std::lock_guard lock(w->mtx_);
-    w->signaled_batch_++;
-  }
-  w->cond_.notify_one();
-}
-
 void Worker::commit(cudaStream_t stream) {
   // Move pending tasks into tasks
   if (pending_tasks_.empty()) {
     return;
   }
   {
-    std::lock_guard lock(mtx_);
+    std::lock_guard lock(state_->mtx);
     // Move pending tasks into ready tasks
-    worker_tasks_[++committed_batch_] = std::move(pending_tasks_);
+    state_->worker_tasks[++committed_batch_] = std::move(pending_tasks_);
   }
   signal_event_.record(stream);
   signal_event_.wait(signal_stream_);
-  CHECK_CUDA_ERROR(cudaLaunchHostFunc(signal_stream_, signal, this));
+  CHECK_CUDA_ERROR(
+      cudaLaunchHostFunc(signal_stream_, State::signal, state_.get()));
+}
+
+// static
+void Worker::State::signal(void* data) {
+  auto state = static_cast<State*>(data);
+  {
+    std::lock_guard lock(state->mtx);
+    state->signaled_batch++;
+  }
+  state->cond.notify_one();
 }
 
-void Worker::thread_fn() {
+void Worker::State::thread_fn() {
   uint64_t current_batch = 0;
   while (true) {
     Tasks tasks;
     {
-      std::unique_lock<std::mutex> lk(mtx_);
-      cond_.wait(lk, [this, current_batch] {
-        return this->signaled_batch_ > current_batch || this->stop_;
+      std::unique_lock<std::mutex> lk(mtx);
+      cond.wait(lk, [this, current_batch] {
+        return this->signaled_batch > current_batch || this->stop;
       });
-      if (stop_) {
+      if (stop) {
         return;
       }
-      current_batch = signaled_batch_;
-      auto end = worker_tasks_.upper_bound(current_batch);
-      for (auto it = worker_tasks_.begin(); it != end; ++it) {
+      current_batch = signaled_batch;
+      auto end = worker_tasks.upper_bound(current_batch);
+      for (auto it = worker_tasks.begin(); it != end; ++it) {
         if (tasks.empty()) {
           tasks = std::move(it->second);
         } else {
           std::move(
               it->second.begin(), it->second.end(), std::back_inserter(tasks));
         }
       }
-      worker_tasks_.erase(worker_tasks_.begin(), end);
+      worker_tasks.erase(worker_tasks.begin(), end);
     }
     // Make sure tasks are cleared before the next wait
     for (int i = 0; i < tasks.size(); ++i) {

mlx/backend/cuda/worker.h

@@ -9,22 +9,18 @@
 #include <map>
 #include <memory>
 #include <mutex>
-#include <thread>
 
 namespace mlx::core::cu {
 
 // Run tasks in worker thread, synchronized with cuda stream.
-class Worker : public std::enable_shared_from_this<Worker> {
+class Worker {
  public:
   explicit Worker(Device& d);
   ~Worker();
 
   Worker(const Worker&) = delete;
   Worker& operator=(const Worker&) = delete;
 
-  void start();
-  void stop();
-
   // Add a pending |task| that will run when consumed or commited.
   void add_task(std::function<void()> task);
 
@@ -33,27 +29,31 @@ class Worker : public std::enable_shared_from_this<Worker> {
   void commit(cudaStream_t stream);
 
  private:
-  static void signal(void*);
+  using Tasks = std::vector<std::function<void()>>;
+
+  // State shared with the detached worker thread, which may outlive the Worker
+  // and free the state after the CUDA runtime is gone.
+  struct State {
+    static void signal(void* data);
+    void thread_fn();
 
-  void thread_fn();
-  std::mutex mtx_;
-  std::condition_variable cond_;
+    std::mutex mtx;
+    std::condition_variable cond;
+    uint64_t signaled_batch{0};
+    bool stop{false};
+    std::map<uint64_t, Tasks> worker_tasks;
+  };
 
   uint64_t committed_batch_{0};
-  uint64_t signaled_batch_{0};
 
   // Cuda stream and event for signaling kernel completion.
   CudaStream signal_stream_;
   CudaEvent signal_event_;
 
-  bool stop_{false};
-
   // Tasks are put in |pending_tasks_| first, and then moved to
-  // |worker_tasks_| when end_batch() is called.
-  using Tasks = std::vector<std::function<void()>>;
+  // |state_->worker_tasks| when commit() is called.
   Tasks pending_tasks_;
-  std::map<uint64_t, Tasks> worker_tasks_;
-  std::thread worker_;
+  std::shared_ptr<State> state_;
 };
 
 } // namespace mlx::core::cu

mlx/compile.cpp

@@ -3,6 +3,7 @@
 #include <atomic>
 #include <cstdlib>
 #include <map>
+#include <optional>
 #include <shared_mutex>
 #include <sstream>
 #include <unordered_map>
@@ -302,9 +303,10 @@ std::uintptr_t get_function_address(const std::function<T(U...)>& fun) {
 class CompileCache {
  public:
   struct CacheEntry {
-    CacheEntry(Stream stream, bool shapeless)
+    CacheEntry(std::optional<Stream> stream, bool shapeless)
         : stream(stream), shapeless(shapeless) {};
-    Stream stream;
+    // The default stream when the function was traced, if there was one.
+    std::optional<Stream> stream;
     bool shapeless;
     std::vector<array> inputs;
     std::vector<array> outputs;
@@ -370,7 +372,8 @@ class CompileCache {
     // Loop over entries and check:
     // - Default stream and device match the entry's default stream
     // - Inputs match i.e. shapes and types must be equal.
-    auto stream = default_stream(default_device());
+    // The default stream is only created when the function actually uses it.
+    auto stream = default_stream_if_exists(default_device());
     for (CacheEntry& entry : entries) {
       // Check that the default stream and device match
       if (entry.stream != stream) {
@@ -1152,6 +1155,8 @@ ArrayFnWithExtra compile(
       // Trace to build the graph
       std::tie(entry.inputs, entry.outputs, entry.extra) =
           compile_trace(fun, inputs, shapeless);
+      // Tracing may have created the default stream.
+      entry.stream = default_stream_if_exists(default_device());
 
       // DFS the graph and get a tape, and a map of array id to (parent,
       // position in parent inputs)

mlx/stream.cpp

@@ -50,6 +50,10 @@ Stream default_stream(Device d) {
   return s.value();
 }
 
+std::optional<Stream> default_stream_if_exists(Device d) {
+  return default_stream_storage(d);
+}
+
 void set_default_stream(Stream s) {
   if (!gpu::is_available() && s.device == Device::gpu) {
     throw std::invalid_argument(

mlx/stream.h

@@ -2,6 +2,7 @@
 
 #pragma once
 
+#include <optional>
 #include <tuple>
 #include <vector>
 
@@ -29,6 +30,9 @@ struct MLX_API ThreadLocalStream : public Stream {
 /** Get the default stream of current thread for the given device. */
 MLX_API Stream default_stream(Device d);
 
+/** Get the default stream of current thread if it exists. */
+std::optional<Stream> default_stream_if_exists(Device d);
+
 /** Make the stream the default for its device on current thread. */
 MLX_API void set_default_stream(Stream s);
 

mlx/transforms.cpp

@@ -80,14 +80,16 @@ thread_local int detail::RetainGraph::tracing_counter{0};
 array eval_impl(std::vector<array> outputs, bool async) {
   std::deque<array> tape;
 
-  // Make an effort to choose a good output stream
-  Stream stream = default_stream(default_device());
-  for (auto& o : outputs) {
-    if (o.status() == array::Status::unscheduled && o.has_primitive()) {
-      stream = o.primitive().stream();
-      break;
+  // Make an effort to choose a good output stream, and only create the default
+  // stream when there is no other choice.
+  Stream stream = [&outputs]() {
+    for (auto& o : outputs) {
+      if (o.status() == array::Status::unscheduled && o.has_primitive()) {
+        return o.primitive().stream();
+      }
     }
-  }
+    return default_stream(default_device());
+  }();
 
   // Map of array id that needs fence and stream it's computed on
   std::unordered_map<uintptr_t, std::pair<uint32_t, bool>> needs_fence;