CUDA backend: one current stream on Windows, shim copies during another thread's graph capture, and low-memory mode moves weights off the GPU while autotuning (#24597)

### What is wrong today

**Windows stream state.** The CUDA delegate and its shims share a
per-thread "current stream" map. In the default Windows build
(`EXECUTORCH_BUILD_SHARED` off) it exists twice, in
`aoti_cuda_shims.dll` (not exported) and in the program. The shims read
the empty copy and make their own stream, so they race with generated
kernels. On every platform, a thread with no stream selected also leaked
a stream.

**Copies during another thread's graph capture.** With no caller stream,
a shim copy used `cudaMemcpy` on the legacy stream. CUDA rejects that
during a capture on a blocking stream like the delegate's, so the
process aborted. Reading back a strided output also failed.

**Export memory.** AOTInductor's final autotuning pass keeps the real
weights on the GPU, unused. A video policy needing about 60 GiB ran out
of memory on an 80 GB GPU.

### What this change does

- Move the current-stream map into the shared library `extension_cuda`,
in an internal namespace. With no stream selected, use
`cudaStreamPerThread`, the stream the delegate runs on.
- Copy with `cudaMemcpyAsync` on the current stream, and read strided
outputs back on the delegate's stream.
- Still on the legacy stream, among others: the allocator's copy and
`synchronize()`, mutable-state and load-time copies, and the capture
step's seed copies.
- In the existing `low_memory_mode`, move the weights to a file in
PyTorch's cache directory during that pass. It needs PyTorch 2.15 or a
nightly from 2026-09-24 or later; with older PyTorch the mode works as
before. It needs free space there, and PyTorch skips the move when the
last generated kernel runs on the CPU.

### Tested

- A100: the CUDA runtime and shim tests, and most CUDA Python tests
(PyTorch 2.14), pass. The new per-thread, strided read-back and
cross-thread copy tests fail without their fix.
- With the Windows layout emulated on Linux, `test_cuda_current_stream`
fails before the change and passes after. CI's Windows CUDA runtime job
runs it in that layout, but does not build the shim tests, so
`RefusedInsideGraphCapture` and `CopiesWhileAnotherThreadCaptures` run
only on Linux. Windows was not built locally.
- A100 with PyTorch 2.15: the low-memory tests pass and fail without the
change. In low-memory mode, a model with 512 MiB of weights peaks at
about 0.66 GiB of GPU memory instead of 1.16 GiB and matches eager. With
PyTorch 2.14, nothing changes.
- Pending: exporting the video policy in low-memory mode.

Author: shoumikhin

.ci/scripts/test_cuda_runtime_windows.ps1

@@ -35,6 +35,7 @@ $tests = @(
     "test_cuda_weight_cache",
     "test_cuda_output_copy",
     "test_cuda_sort_rand",
+    "test_cuda_current_stream",
     "test_cuda_guard",
     "test_cuda_stream_guard"
 )

.ci/scripts/wheel/test_shared_libraries.py

@@ -81,14 +81,14 @@
 _CUDA_BACKEND_SYMBOLS = ("executorch::backends::cuda::load_library",)
 _CUDA_EXTENSION_SYMBOLS = (
     "executorch::extension::cuda::getCallerStream",
+    "executorch::extension::cuda::internal::getCurrentCUDAStream",
     "executorch::extension::cuda::CudaAllocator::instance",
 )
 
-# The AOTI shim layer and the stream-guard state that lives with it. This is the state that was
-# genuinely duplicated: extracting the shims with a PUBLIC whole-archive replayed the extraction at
-# every consumer's link, so the guard's thread_local and these shims landed in three shipped binaries
-# at once, and a stream selected through one copy was invisible to the other two. The row above cannot
-# catch that, because its symbol only ever existed in one unconditionally shared library.
+# The AOTI shim layer. This is what was once genuinely duplicated: extracting the shims with a PUBLIC
+# whole-archive replayed the extraction at every consumer's link, so these shims, and the current-stream
+# thread_local that then lived with them, landed in three shipped binaries at once, and a stream selected
+# through one copy was invisible to the other two. That state now lives in the extension library above.
 _AOTI_SHIM_SYMBOLS = (
     "aoti_torch_empty_strided",
     "aoti_torch_delete_tensor_object",

.github/workflows/cuda.yml

@@ -782,12 +782,13 @@ jobs:
           -v -o "addopts="
 
         cmake --preset llm-release-cuda -DEXECUTORCH_BUILD_TESTS=ON
-        cmake --build cmake-out --target test_cuda_allocator test_cuda_allocator_exit test_cuda_mutable_state test_cuda_weight_cache test_cuda_output_copy test_cuda_sort_rand test_cuda_kv_cache test_cuda_guard test_cuda_stream_guard test_step_plan test_cuda_executor -j$(nproc)
+        cmake --build cmake-out --target test_cuda_allocator test_cuda_allocator_exit test_cuda_mutable_state test_cuda_weight_cache test_cuda_output_copy test_cuda_sort_rand test_cuda_current_stream test_cuda_kv_cache test_cuda_guard test_cuda_stream_guard test_step_plan test_cuda_executor -j$(nproc)
         ctest --test-dir cmake-out -R test_cuda_allocator --output-on-failure -V
         ctest --test-dir cmake-out -R test_cuda_mutable_state --output-on-failure -V
         ctest --test-dir cmake-out -R test_cuda_weight_cache --output-on-failure -V
         ctest --test-dir cmake-out -R test_cuda_output_copy --output-on-failure -V
         ctest --test-dir cmake-out -R test_cuda_sort_rand --output-on-failure -V
+        ctest --test-dir cmake-out -R test_cuda_current_stream --output-on-failure -V
         ctest --test-dir cmake-out -R test_cuda_kv_cache --output-on-failure -V
         ctest --test-dir cmake-out -R test_step_plan --output-on-failure -V
         # CudaExecutor end to end on a toy decoder, against its eager greedy

backends/aoti/common_shims_slim.cpp

@@ -10,6 +10,9 @@
 #include <executorch/runtime/platform/abort.h>
 #include <executorch/runtime/platform/log.h>
 
+// Keep this file free of static and thread_local state: a Windows build
+// without EXECUTORCH_BUILD_SHARED links it twice (see backends/cuda).
+
 namespace executorch {
 namespace backends {
 namespace aoti {

backends/aoti/slim/core/storage.h

@@ -178,17 +178,25 @@ struct DeviceTraits<c10::DeviceType::CUDA> {
           static_cast<int>(dst_device.index()));
     }
 
-    // Plain cudaMemcpy is host-synchronous on the default stream, which a
-    // green context would not confine. When a caller stream is active, copy
-    // on it asynchronously and synchronize it to preserve blocking
-    // semantics; otherwise fall back to the plain synchronous copy.
+    // Plain cudaMemcpy runs on the legacy default stream, which a green
+    // context would not confine, and which fails while another thread is
+    // capturing a CUDA graph. So copy on the caller stream, or else on the
+    // current stream.
     const auto caller_stream = executorch::extension::cuda::getCallerStream();
     if (caller_stream) {
       ET_CUDA_CHECK(
           cudaMemcpyAsync(dst, src, nbytes, direction, *caller_stream));
       ET_CUDA_CHECK(cudaStreamSynchronize(*caller_stream));
     } else {
-      ET_CUDA_CHECK(cudaMemcpy(dst, src, nbytes, direction));
+      auto current = executorch::backends::cuda::getCurrentCUDAStream(
+          src_device.is_cuda() ? src_device.index() : dst_device.index());
+      ET_CHECK_MSG(current.ok(), "Failed to get current CUDA stream");
+      ET_CUDA_CHECK(
+          cudaMemcpyAsync(dst, src, nbytes, direction, current.get()));
+      // cudaMemcpy does not wait for a copy between device buffers.
+      if (direction != cudaMemcpyDeviceToDevice) {
+        ET_CUDA_CHECK(cudaStreamSynchronize(current.get()));
+      }
     }
   }
 };

backends/aoti/slim/cuda/guard.cpp

@@ -9,76 +9,11 @@
 #include <executorch/backends/aoti/slim/cuda/guard.h>
 #include <executorch/runtime/platform/log.h>
 #include <limits>
-#include <optional>
-#include <unordered_map>
 
-namespace executorch::backends::cuda {
-
-namespace {
-// Thread-local stream storage (private to this file)
-thread_local std::unordered_map<DeviceIndex, cudaStream_t> current_streams_;
-} // namespace
-
-Error setCurrentCUDAStream(cudaStream_t stream, DeviceIndex device_index) {
-  if (device_index == -1) {
-    // Get current device if not specified
-    // CUDA API returns int, explicit cast to DeviceIndex (int8_t) following
-    // ATen
-    int tmp_device = -1;
-    ET_CUDA_CHECK_OR_RETURN_ERROR(cudaGetDevice(&tmp_device));
-    device_index = static_cast<DeviceIndex>(tmp_device);
-  }
-
-  current_streams_[device_index] = stream;
-  return Error::Ok;
-}
-
-Result<cudaStream_t> getCurrentCUDAStream(DeviceIndex device_index) {
-  if (device_index == -1) {
-    // CUDA API returns int, explicit cast to DeviceIndex (int8_t) following
-    // ATen
-    int tmp_device = -1;
-    ET_CUDA_CHECK_OR_RETURN_ERROR(cudaGetDevice(&tmp_device));
-    device_index = static_cast<DeviceIndex>(tmp_device);
-  }
+// Keep this file free of static and thread_local state: a Windows build
+// without EXECUTORCH_BUILD_SHARED links it twice (see backends/cuda).
 
-  auto it = current_streams_.find(device_index);
-  if (it != current_streams_.end()) {
-    return it->second;
-  }
-
-  cudaStream_t stream;
-  ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamCreate(&stream));
-  setCurrentCUDAStream(stream, device_index);
-  return stream;
-}
-
-std::optional<cudaStream_t> peekCurrentCUDAStream(DeviceIndex device_index) {
-  if (device_index == -1) {
-    int tmp_device = -1;
-    if (cudaGetDevice(&tmp_device) != cudaSuccess) {
-      return std::nullopt;
-    }
-    device_index = static_cast<DeviceIndex>(tmp_device);
-  }
-
-  auto it = current_streams_.find(device_index);
-  if (it == current_streams_.end()) {
-    return std::nullopt;
-  }
-  return it->second;
-}
-
-void clearCurrentCUDAStream(DeviceIndex device_index) {
-  if (device_index == -1) {
-    int tmp_device = -1;
-    if (cudaGetDevice(&tmp_device) != cudaSuccess) {
-      return;
-    }
-    device_index = static_cast<DeviceIndex>(tmp_device);
-  }
-  current_streams_.erase(device_index);
-}
+namespace executorch::backends::cuda {
 
 CUDAStreamGuard::CUDAStreamGuard(CUDAStreamGuard&& other) noexcept
     : device_guard_(std::move(other.device_guard_)),

backends/aoti/slim/cuda/guard.h

@@ -8,10 +8,9 @@
 
 #pragma once
 
-#include <optional>
-
 #include <executorch/backends/aoti/slim/c10/core/Device.h>
 #include <executorch/backends/aoti/slim/c10/cuda/Exception.h>
+#include <executorch/extension/cuda/current_stream.h>
 #include <executorch/extension/cuda/device_guard.h>
 #include <executorch/extension/cuda/runtime_api.h>
 #include <executorch/runtime/core/error.h>
@@ -20,6 +19,10 @@
 namespace executorch::backends::cuda {
 
 using ::executorch::extension::cuda::CUDAGuard;
+using ::executorch::extension::cuda::internal::clearCurrentCUDAStream;
+using ::executorch::extension::cuda::internal::getCurrentCUDAStream;
+using ::executorch::extension::cuda::internal::peekCurrentCUDAStream;
+using ::executorch::extension::cuda::internal::setCurrentCUDAStream;
 
 using executorch::runtime::Error;
 using executorch::runtime::Result;
@@ -29,46 +32,6 @@ using executorch::runtime::Result;
 // calls, and explicit static_cast handles inbound conversions from CUDA APIs.
 using executorch::backends::aoti::slim::c10::DeviceIndex;
 
-/**
- * Set the current CUDA stream for the specified device.
- *
- * @param stream The CUDA stream to set as current
- * @param device_index The device index (-1 to use current device)
- * @return Error code indicating success or failure
- */
-Error setCurrentCUDAStream(cudaStream_t stream, DeviceIndex device_index = -1);
-
-/**
- * Get the current CUDA stream for the specified device.
- * If no stream has been set, creates a new stream and sets it as current.
- *
- * @param device_index The device index (-1 to use current device)
- * @return Result containing the current stream on success, or an error code on
- * failure
- */
-Result<cudaStream_t> getCurrentCUDAStream(DeviceIndex device_index = -1);
-
-/**
- * The CUDA stream registered for the specified device, or std::nullopt if none
- * is set. Unlike getCurrentCUDAStream, it never creates one, so it can snapshot
- * the current selection without side effects. Also returns std::nullopt if the
- * current device cannot be queried (device_index -1), so nullopt does not
- * distinguish "no stream set" from "device query failed".
- *
- * @param device_index The device index (-1 to use current device)
- */
-std::optional<cudaStream_t> peekCurrentCUDAStream(
-    DeviceIndex device_index = -1);
-
-/**
- * Clears any CUDA stream registered for the specified device, restoring the
- * "no stream selected" state. Best-effort: if device_index is -1 and the
- * current device cannot be queried, it silently does nothing.
- *
- * @param device_index The device index (-1 to use current device)
- */
-void clearCurrentCUDAStream(DeviceIndex device_index = -1);
-
 /**
  * RAII guard that sets the current CUDA device and stream, restoring both on
  * destruction. This is useful for temporarily switching to a different device

backends/cuda/CMakeLists.txt

@@ -310,10 +310,9 @@ if(_cuda_is_msvc_toolchain)
 else()
   # The whole-archive pair is PRIVATE so it applies to this library's own link
   # and does not replay at the link of anything that links this one. In PUBLIC
-  # scope the archive was extracted again by each consumer, so the translation
-  # unit holding SlimTensor's per-device current stream thread_local landed in
-  # three shipped binaries, and one copy per process is what makes that state
-  # agree. The MSVC branch above uses PRIVATE for the same reason.
+  # scope the archive was extracted again by each consumer, so its objects
+  # landed in three shipped binaries. The MSVC branch above uses PRIVATE for the
+  # same reason.
   target_link_libraries(
     aoti_cuda_shims
     PRIVATE cuda_platform -Wl,--whole-archive aoti_common_shims_slim
@@ -337,10 +336,9 @@ endif()
 # source install keeps its own paths.
 if(EXECUTORCH_BUILD_SHARED)
   if(_cuda_is_windows_msvc)
-    # In the shared layout the delegate is its own DLL and imports the stream
-    # guard functions (get/set/peek/clearCurrentCUDAStream) from this one. They
-    # carry no AOTI_SHIM_EXPORT, so the annotations alone would not export them;
-    # export everything, as every other shipped DLL does.
+    # In the shared layout the delegate is its own DLL and links this one.
+    # Export everything, as every other shipped DLL does, not only the shims
+    # marked AOTI_SHIM_EXPORT.
     set_target_properties(
       aoti_cuda_shims PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS ON
     )
@@ -426,9 +424,11 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM)
 endif()
 
 if(_cuda_is_msvc_toolchain)
-  # cuda_backend.cpp uses SlimTensor CUDA utilities (e.g. getCurrentCUDAStream)
-  # from aoti_common_shims_slim via headers; propagate the static lib so final
-  # MSVC links (e.g. parakeet_runner) can resolve those C++ symbols.
+  # The shims DLL compiles these sources itself (aoti_common_shims_slim_obj),
+  # and cuda_backend.cpp still needs the SlimTensor headers and definitions that
+  # come with the static library. So the program carries a second copy of that
+  # code, which must hold no state; the current stream lives in extension_cuda
+  # for that reason.
   target_link_libraries(
     aoti_cuda_backend
     PRIVATE aoti_cuda_shims
@@ -508,6 +508,11 @@ if(BUILD_TESTING)
     )
     target_compile_definitions(test_cuda_sort_rand PRIVATE CUDA_AVAILABLE=1)
   endif()
+  et_cxx_test(
+    test_cuda_current_stream SOURCES runtime/test/test_cuda_current_stream.cpp
+    EXTRA_LIBS aoti_cuda_backend
+  )
+  target_compile_definitions(test_cuda_current_stream PRIVATE CUDA_AVAILABLE=1)
 
   if(EXECUTORCH_BUILD_EXTENSION_LLM)
     et_cxx_test(

backends/cuda/cuda_backend.py

@@ -1359,6 +1359,18 @@ def get_aoti_compile_options(
         options["max_autotune"] = max_autotune
         if autotune_at_compile_time is not None:
             options["triton.autotune_at_compile_time"] = autotune_at_compile_time
+        # The final autotuning pass picks each kernel's launch settings by
+        # running it on random tensors shaped like its inputs, weights included,
+        # so in low-memory mode the real weights wait in a file meanwhile instead
+        # of taking GPU memory next to their random stand-ins. The threshold of
+        # 0.0 moves them for every model, because low-memory mode is for models
+        # that are short of GPU memory, and moving a small model costs little.
+        # TODO: drop the hasattr once every torch the repo installs has the option.
+        if cls._is_low_memory_mode(compile_specs) and hasattr(
+            torch._inductor.config.aot_inductor,
+            "autotune_offload_constants_min_device_fraction",
+        ):
+            options["aot_inductor.autotune_offload_constants_min_device_fraction"] = 0.0
         # Add platform-specific options
 
         if platform == "windows":

backends/cuda/runtime/cuda_backend.cpp

@@ -695,8 +695,7 @@ class ET_EXPERIMENTAL CudaBackend final
         NotSupported,
         "CUDA graph is not supported together with a caller-provided CUDA stream.");
 
-    // Snapshot the prior selection without creating one (peek, not get), so the
-    // restore is exact and we don't leak a stream just to snapshot.
+    // Snapshot the prior selection with peek, not get, so the restore is exact.
     std::optional<cudaStream_t> prev_stream;
     if (caller_stream) {
       prev_stream = peekCurrentCUDAStream(0);