From 6f342283e2a14caf7cd2df2eb1972bf2c3996035 Mon Sep 17 00:00:00 2001
From: Alessio Pollero <alessio.pollero@gmail.com>
Date: Tue, 15 Sep 2026 11:14:36 +0400
Subject: [PATCH 1/4] Fix CUDA termination crash via deterministic worker
 teardown and lazy stream evaluation

---
 mlx/backend/cuda/cuda_utils.h |  8 ++---
 mlx/backend/cuda/device.cpp   |  4 +--
 mlx/backend/cuda/device.h     |  2 +-
 mlx/backend/cuda/worker.cpp   | 67 +++++++++++++++++------------------
 mlx/backend/cuda/worker.h     | 32 ++++++++---------
 mlx/transforms.cpp            | 16 +++++----
 tests/scheduler_tests.cpp     | 10 ++++++
 7 files changed, 74 insertions(+), 65 deletions(-)

diff --git a/mlx/backend/cuda/cuda_utils.h b/mlx/backend/cuda/cuda_utils.h
index f8a234ee65..95fb8f2dd2 100644
--- a/mlx/backend/cuda/cuda_utils.h
+++ b/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;
diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp
index e7d8f2620d..22da93d193 100644
--- a/mlx/backend/cuda/device.cpp
+++ b/mlx/backend/cuda/device.cpp
@@ -206,15 +206,13 @@ 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();
 }
 
 void CommandEncoder::add_completed_handler(std::function<void()> task) {
diff --git a/mlx/backend/cuda/device.h b/mlx/backend/cuda/device.h
index 198f0b5ad8..cf525e97ed 100644
--- a/mlx/backend/cuda/device.h
+++ b/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_;
diff --git a/mlx/backend/cuda/worker.cpp b/mlx/backend/cuda/worker.cpp
index 6266d6099b..86e8a4cc55 100644
--- a/mlx/backend/cuda/worker.cpp
+++ b/mlx/backend/cuda/worker.cpp
@@ -3,74 +3,73 @@
 #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 {
@@ -78,7 +77,7 @@ void Worker::thread_fn() {
               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) {
diff --git a/mlx/backend/cuda/worker.h b/mlx/backend/cuda/worker.h
index 54a4456ace..5d220f1e7a 100644
--- a/mlx/backend/cuda/worker.h
+++ b/mlx/backend/cuda/worker.h
@@ -9,12 +9,11 @@
 #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();
@@ -22,9 +21,6 @@ class Worker : public std::enable_shared_from_this<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. The thread may free it at
+  // any time after the Worker is gone, so it must not own CUDA resources.
+  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
diff --git a/mlx/transforms.cpp b/mlx/transforms.cpp
index 1a8a5f9dc5..77d81e11e9 100644
--- a/mlx/transforms.cpp
+++ b/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());
+  }();
 
   struct FenceInfo {
     int stream_index;
diff --git a/tests/scheduler_tests.cpp b/tests/scheduler_tests.cpp
index 8a98d35eb9..532e81a383 100644
--- a/tests/scheduler_tests.cpp
+++ b/tests/scheduler_tests.cpp
@@ -114,6 +114,16 @@ TEST_CASE("test thread unsafe stream") {
   CHECK_EQ(expected, actual);
 }
 
+TEST_CASE("test eval does not create default stream") {
+  auto s = new_thread_unsafe_stream(default_device());
+  size_t num_streams = get_streams().size();
+
+  std::thread t([&] { eval(arange(10, s)); });
+  t.join();
+
+  CHECK_EQ(get_streams().size(), num_streams);
+}
+
 TEST_CASE("test thread local stream") {
   auto s = new_thread_local_stream(default_device());
   int result = sum(arange(10, s)).item<int>();

From 16722fb492fc5a3c69b17014304903da6653cec4 Mon Sep 17 00:00:00 2001
From: Alessio Pollero <alessio.pollero@gmail.com>
Date: Wed, 16 Sep 2026 14:04:23 +0400
Subject: [PATCH 2/4] Avoid implicit default stream in compile cache and ignore
 CUDA errors on teardown

---
 mlx/backend/cuda/allocator.cpp | 12 +++++++-----
 mlx/backend/cuda/device.cpp    |  8 +++++++-
 mlx/backend/cuda/event.cu      |  4 ++--
 mlx/backend/cuda/worker.h      |  4 ++--
 mlx/compile.cpp                | 11 ++++++++---
 mlx/stream.cpp                 |  4 ++++
 mlx/stream.h                   |  4 ++++
 tests/scheduler_tests.cpp      | 21 ++++++++++++++++++++-
 8 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp
index 1cab6b4081..707624ea5f 100644
--- a/mlx/backend/cuda/allocator.cpp
+++ b/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);
     }
   }
 }
diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp
index 22da93d193..7662fb2583 100644
--- a/mlx/backend/cuda/device.cpp
+++ b/mlx/backend/cuda/device.cpp
@@ -212,7 +212,11 @@ CommandEncoder::CommandEncoder(Device& d)
 }
 
 CommandEncoder::~CommandEncoder() {
-  synchronize();
+  try {
+    synchronize();
+  } catch (...) {
+    // Synchronizing can fail when the CUDA runtime is shutting down.
+  }
 }
 
 void CommandEncoder::add_completed_handler(std::function<void()> task) {
@@ -488,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(); });
diff --git a/mlx/backend/cuda/event.cu b/mlx/backend/cuda/event.cu
index d3b6f97f5d..d1f27e1230 100644
--- a/mlx/backend/cuda/event.cu
+++ b/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 {
diff --git a/mlx/backend/cuda/worker.h b/mlx/backend/cuda/worker.h
index 5d220f1e7a..451cbf3b9c 100644
--- a/mlx/backend/cuda/worker.h
+++ b/mlx/backend/cuda/worker.h
@@ -31,8 +31,8 @@ class Worker {
  private:
   using Tasks = std::vector<std::function<void()>>;
 
-  // State shared with the detached worker thread. The thread may free it at
-  // any time after the Worker is gone, so it must not own CUDA resources.
+  // 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();
diff --git a/mlx/compile.cpp b/mlx/compile.cpp
index bb17c44962..c3ccd737d1 100644
--- a/mlx/compile.cpp
+++ b/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)
diff --git a/mlx/stream.cpp b/mlx/stream.cpp
index b78ee67d67..08f4d6830f 100644
--- a/mlx/stream.cpp
+++ b/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(
diff --git a/mlx/stream.h b/mlx/stream.h
index 3099ee2682..bae55751cc 100644
--- a/mlx/stream.h
+++ b/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);
 
diff --git a/tests/scheduler_tests.cpp b/tests/scheduler_tests.cpp
index 532e81a383..cc2228e0f3 100644
--- a/tests/scheduler_tests.cpp
+++ b/tests/scheduler_tests.cpp
@@ -118,7 +118,26 @@ TEST_CASE("test eval does not create default stream") {
   auto s = new_thread_unsafe_stream(default_device());
   size_t num_streams = get_streams().size();
 
-  std::thread t([&] { eval(arange(10, s)); });
+  std::thread t([&] {
+    async_eval(arange(10, s));
+    eval(arange(10, s));
+  });
+  t.join();
+
+  CHECK_EQ(get_streams().size(), num_streams);
+}
+
+TEST_CASE("test compile does not create default stream") {
+  auto s = new_thread_unsafe_stream(default_device());
+  size_t num_streams = get_streams().size();
+
+  std::function<std::vector<array>(const std::vector<array>&)> fun =
+      [s](const std::vector<array>& inputs) {
+        return std::vector<array>{abs(inputs[0], s)};
+      };
+  auto cfun = compile(fun);
+
+  std::thread t([&] { eval(cfun({array({-1, 2})})); });
   t.join();
 
   CHECK_EQ(get_streams().size(), num_streams);

From 4ef62a8bfa21c7907b6bcfaaf86126332fe16c7c Mon Sep 17 00:00:00 2001
From: Alessio Pollero <alessio.pollero@gmail.com>
Date: Mon, 28 Sep 2026 07:28:31 +0200
Subject: [PATCH 3/4] Revert Worker changes and export default_stream_if_exists

---
 mlx/backend/cuda/device.cpp |  6 ++--
 mlx/backend/cuda/device.h   |  2 +-
 mlx/backend/cuda/worker.cpp | 67 +++++++++++++++++++------------------
 mlx/backend/cuda/worker.h   | 32 +++++++++---------
 mlx/stream.h                |  2 +-
 5 files changed, 55 insertions(+), 54 deletions(-)

diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp
index 7662fb2583..6bc16bf7ba 100644
--- a/mlx/backend/cuda/device.cpp
+++ b/mlx/backend/cuda/device.cpp
@@ -206,9 +206,10 @@ CommandEncoder::CommandEncoder(Device& d)
     : device_(d),
       stream_(d),
       graph_(d),
-      worker_(std::make_unique<Worker>(d)),
+      worker_(std::make_shared<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() {
@@ -217,6 +218,7 @@ CommandEncoder::~CommandEncoder() {
   } catch (...) {
     // Synchronizing can fail when the CUDA runtime is shutting down.
   }
+  worker_->stop();
 }
 
 void CommandEncoder::add_completed_handler(std::function<void()> task) {
@@ -492,8 +494,6 @@ 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(); });
diff --git a/mlx/backend/cuda/device.h b/mlx/backend/cuda/device.h
index cf525e97ed..198f0b5ad8 100644
--- a/mlx/backend/cuda/device.h
+++ b/mlx/backend/cuda/device.h
@@ -146,7 +146,7 @@ class CommandEncoder {
   Device& device_;
   CudaStream stream_;
   CudaGraph graph_;
-  std::unique_ptr<Worker> worker_;
+  std::shared_ptr<Worker> worker_;
   int node_count_{0};
   bool in_concurrent_{false};
   std::vector<cudaGraphNode_t> from_nodes_;
diff --git a/mlx/backend/cuda/worker.cpp b/mlx/backend/cuda/worker.cpp
index 86e8a4cc55..6266d6099b 100644
--- a/mlx/backend/cuda/worker.cpp
+++ b/mlx/backend/cuda/worker.cpp
@@ -3,73 +3,74 @@
 #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),
-      state_(std::make_shared<State>()) {
-  // Detach the thread and let it free the state after finishing tasks.
+      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.
   // This is to avoid deadlock when joining threads on exit on Windows:
   // https://developercommunity.visualstudio.com/t/1654756
-  std::thread(&State::thread_fn, state_).detach();
+  worker_.detach();
 }
 
-Worker::~Worker() {
+void Worker::stop() {
   {
-    std::lock_guard lock(state_->mtx);
-    state_->stop = true;
+    std::lock_guard lock(mtx_);
+    stop_ = true;
   }
-  state_->cond.notify_one();
+  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(state_->mtx);
+    std::lock_guard lock(mtx_);
     // Move pending tasks into ready tasks
-    state_->worker_tasks[++committed_batch_] = std::move(pending_tasks_);
+    worker_tasks_[++committed_batch_] = std::move(pending_tasks_);
   }
   signal_event_.record(stream);
   signal_event_.wait(signal_stream_);
-  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();
+  CHECK_CUDA_ERROR(cudaLaunchHostFunc(signal_stream_, signal, this));
 }
 
-void Worker::State::thread_fn() {
+void Worker::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 {
@@ -77,7 +78,7 @@ void Worker::State::thread_fn() {
               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) {
diff --git a/mlx/backend/cuda/worker.h b/mlx/backend/cuda/worker.h
index 451cbf3b9c..54a4456ace 100644
--- a/mlx/backend/cuda/worker.h
+++ b/mlx/backend/cuda/worker.h
@@ -9,11 +9,12 @@
 #include <map>
 #include <memory>
 #include <mutex>
+#include <thread>
 
 namespace mlx::core::cu {
 
 // Run tasks in worker thread, synchronized with cuda stream.
-class Worker {
+class Worker : public std::enable_shared_from_this<Worker> {
  public:
   explicit Worker(Device& d);
   ~Worker();
@@ -21,6 +22,9 @@ class 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);
 
@@ -29,31 +33,27 @@ class Worker {
   void commit(cudaStream_t stream);
 
  private:
-  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();
+  static void signal(void*);
 
-    std::mutex mtx;
-    std::condition_variable cond;
-    uint64_t signaled_batch{0};
-    bool stop{false};
-    std::map<uint64_t, Tasks> worker_tasks;
-  };
+  void thread_fn();
+  std::mutex mtx_;
+  std::condition_variable cond_;
 
   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
-  // |state_->worker_tasks| when commit() is called.
+  // |worker_tasks_| when end_batch() is called.
+  using Tasks = std::vector<std::function<void()>>;
   Tasks pending_tasks_;
-  std::shared_ptr<State> state_;
+  std::map<uint64_t, Tasks> worker_tasks_;
+  std::thread worker_;
 };
 
 } // namespace mlx::core::cu
diff --git a/mlx/stream.h b/mlx/stream.h
index bae55751cc..7bf0b508b5 100644
--- a/mlx/stream.h
+++ b/mlx/stream.h
@@ -31,7 +31,7 @@ struct MLX_API ThreadLocalStream : public Stream {
 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);
+MLX_API 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);

From 0fc55218698f092e104a3869a874a76929d47061 Mon Sep 17 00:00:00 2001
From: Cheng <git@zcbenz.com>
Date: Mon, 28 Sep 2026 17:13:31 +0800
Subject: [PATCH 4/4] nit

---
 mlx/backend/cuda/allocator.cpp | 4 +---
 mlx/backend/cuda/cuda_utils.h  | 3 +--
 mlx/backend/cuda/event.cu      | 1 -
 mlx/compile.cpp                | 8 ++------
 mlx/stream.cpp                 | 2 +-
 mlx/stream.h                   | 2 +-
 6 files changed, 6 insertions(+), 14 deletions(-)

diff --git a/mlx/backend/cuda/allocator.cpp b/mlx/backend/cuda/allocator.cpp
index 707624ea5f..5f4b9f2d55 100644
--- a/mlx/backend/cuda/allocator.cpp
+++ b/mlx/backend/cuda/allocator.cpp
@@ -67,7 +67,6 @@ 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()) {
     cudaFree(data);
@@ -320,8 +319,7 @@ 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. Errors are ignored
-    // since memory is freed in destructors, which can not throw.
+    // Free asynchronously when memory pools is supported.
     if (mem_pools_[buf.device]) {
       if (!stream) {
         stream = free_streams_[buf.device];
diff --git a/mlx/backend/cuda/cuda_utils.h b/mlx/backend/cuda/cuda_utils.h
index 95fb8f2dd2..083d19e6ff 100644
--- a/mlx/backend/cuda/cuda_utils.h
+++ b/mlx/backend/cuda/cuda_utils.h
@@ -26,8 +26,7 @@ class CudaHandle {
   }
 
   ~CudaHandle() {
-    // Errors are ignored since destructors can not throw, and destroying can
-    // fail when the CUDA runtime is shutting down.
+    // Skip error check to avoid throwing in destructors.
     if (handle_ != nullptr) {
       Destroy(handle_);
     }
diff --git a/mlx/backend/cuda/event.cu b/mlx/backend/cuda/event.cu
index d1f27e1230..bcea005443 100644
--- a/mlx/backend/cuda/event.cu
+++ b/mlx/backend/cuda/event.cu
@@ -232,7 +232,6 @@ AtomicEvent::AtomicEvent(Device& d) {
     cuda_free = cudaFree;
     coherent_ = false;
   }
-  // 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;
diff --git a/mlx/compile.cpp b/mlx/compile.cpp
index c3ccd737d1..3cab4c05d7 100644
--- a/mlx/compile.cpp
+++ b/mlx/compile.cpp
@@ -304,8 +304,7 @@ class CompileCache {
  public:
   struct CacheEntry {
     CacheEntry(std::optional<Stream> stream, bool shapeless)
-        : stream(stream), shapeless(shapeless) {};
-    // The default stream when the function was traced, if there was one.
+        : stream(stream), shapeless(shapeless) {}
     std::optional<Stream> stream;
     bool shapeless;
     std::vector<array> inputs;
@@ -372,8 +371,7 @@ 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.
-    // The default stream is only created when the function actually uses it.
-    auto stream = default_stream_if_exists(default_device());
+    auto stream = peek_default_stream(default_device());
     for (CacheEntry& entry : entries) {
       // Check that the default stream and device match
       if (entry.stream != stream) {
@@ -1155,8 +1153,6 @@ 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)
diff --git a/mlx/stream.cpp b/mlx/stream.cpp
index 08f4d6830f..902015cbf0 100644
--- a/mlx/stream.cpp
+++ b/mlx/stream.cpp
@@ -50,7 +50,7 @@ Stream default_stream(Device d) {
   return s.value();
 }
 
-std::optional<Stream> default_stream_if_exists(Device d) {
+std::optional<Stream> peek_default_stream(Device d) {
   return default_stream_storage(d);
 }
 
diff --git a/mlx/stream.h b/mlx/stream.h
index 7bf0b508b5..331af60a61 100644
--- a/mlx/stream.h
+++ b/mlx/stream.h
@@ -31,7 +31,7 @@ struct MLX_API ThreadLocalStream : public Stream {
 MLX_API Stream default_stream(Device d);
 
 /** Get the default stream of current thread if it exists. */
-MLX_API std::optional<Stream> default_stream_if_exists(Device d);
+MLX_API std::optional<Stream> peek_default_stream(Device d);
 
 /** Make the stream the default for its device on current thread. */
 MLX_API void set_default_stream(Stream s);