Run CUDA delegates on the per-thread stream

backends/cuda/CMakeLists.txt

@@ -455,4 +455,5 @@ if(BUILD_TESTING)
     EXTRA_LIBS aoti_cuda_backend
   )
   target_compile_definitions(test_cuda_weight_cache PRIVATE CUDA_AVAILABLE=1)
+
 endif()

backends/cuda/runtime/cuda_backend.cpp

@@ -139,33 +139,6 @@ class ET_EXPERIMENTAL CudaBackend final
     return method_in_csv(method_name, cuda_graph_method_);
   }
 
-  // Create the shared CUDA stream. Called when use_shared_cuda_stream option
-  // is set to true. The presence of shared_cuda_stream_ indicates shared mode.
-  void create_shared_cuda_stream() {
-    std::lock_guard<std::mutex> guard(cuda_stream_mutex_);
-    if (shared_cuda_stream_ != nullptr) {
-      return; // Already created
-    }
-    shared_cuda_stream_ = cuda::create_cuda_stream();
-    if (shared_cuda_stream_ == nullptr) {
-      ET_LOG(Error, "Failed to create shared CUDA stream");
-      return;
-    }
-    ET_LOG(Info, "Created shared CUDA stream: %p", *shared_cuda_stream_);
-  }
-
-  // Get the shared CUDA stream. Returns nullptr if not in shared mode.
-  std::shared_ptr<cudaStream_t> get_shared_cuda_stream() const {
-    std::lock_guard<std::mutex> guard(cuda_stream_mutex_);
-    return shared_cuda_stream_;
-  }
-
-  // Check if we're using shared CUDA stream mode.
-  bool is_using_shared_cuda_stream() const {
-    std::lock_guard<std::mutex> guard(cuda_stream_mutex_);
-    return shared_cuda_stream_ != nullptr;
-  }
-
   // Enable the legacy dense-blob per-FQN cache. New FQN artifacts use
   // their FQN-addressed data keys automatically.
   void set_weight_sharing_across_methods(bool enabled) {
@@ -262,14 +235,17 @@ class ET_EXPERIMENTAL CudaBackend final
             "effect; ignoring it.",
             kSkipCopyOutputToCpuForMethod);
       } else if (std::strcmp(option.key, kUseSharedCudaStream) == 0) {
-        if (auto* val = std::get_if<bool>(&option.value)) {
-          if (*val) {
-            create_shared_cuda_stream();
-          }
-        } else {
+        if (std::get_if<bool>(&option.value) == nullptr) {
           ET_LOG(Error, "Option %s must be a boolean.", kUseSharedCudaStream);
           return Error::InvalidArgument;
         }
+        // Every method already runs on the per-thread stream, which is what
+        // this asked for, so there is nothing left to switch on.
+        ET_LOG(
+            Info,
+            "Runtime backend option '%s' is DEPRECATED and no longer has any "
+            "effect; every method already shares the per-thread stream.",
+            kUseSharedCudaStream);
       } else if (std::strcmp(option.key, kWeightSharingAcrossMethods) == 0) {
         if (auto* val = std::get_if<bool>(&option.value)) {
           set_weight_sharing_across_methods(*val);
@@ -432,30 +408,16 @@ class ET_EXPERIMENTAL CudaBackend final
           load_constants_legacy(handle, named_data_map, weights_blob_key));
     }
 
-    // Use shared CUDA stream if enabled via options, otherwise create one.
-    // A shared stream ensures proper ordering across multiple methods
-    // (e.g., encoder, decoder, sampler) when using skip-copy optimization.
-    if (is_using_shared_cuda_stream()) {
-      // Shared stream mode: all handles share the same stream.
-      handle->cuda_stream = get_shared_cuda_stream();
-      ET_LOG(
-          Info,
-          "Using shared CUDA stream %p for method %s",
-          handle->get_cuda_stream(),
-          method_name.c_str());
-    } else {
-      // Per-handle stream mode: each handle owns its own stream.
-      handle->cuda_stream = cuda::create_cuda_stream();
-      if (handle->cuda_stream == nullptr) {
-        delete handle;
-        return Error::Internal;
-      }
-      ET_LOG(
-          Info,
-          "Created new CUDA stream %p for method %s",
-          handle->get_cuda_stream(),
-          method_name.c_str());
-    }
+    // Handles share one stream so that one delegate's output is ordered against
+    // the next one's read, which a stream per handle left unordered. The
+    // TensorRT delegate falls back to the same stream, so a split program is
+    // ordered too.
+    handle->cuda_stream = std::make_shared<cudaStream_t>(cudaStreamPerThread);
+    ET_LOG(
+        Info,
+        "Using the per-thread CUDA stream %p for method %s",
+        handle->get_cuda_stream(),
+        method_name.c_str());
 
     // Initialize CUDA graph state if enabled for this method.
     if (should_use_cuda_graph_for_method(method_name)) {
@@ -487,7 +449,7 @@ class ET_EXPERIMENTAL CudaBackend final
     handle->get_num_outputs(handle->container_handle, &n_outputs);
 
     // Run on the caller-selected stream when one is active on this thread (e.g.
-    // a CUDA green-context stream), otherwise the handle's own stream. Every
+    // a CUDA green-context stream), otherwise the per-thread stream. Every
     // kernel and boundary copy reads getCurrentCUDAStream, so installing the
     // choice here routes the whole execution; restore the prior selection on
     // return so a caller stream does not linger for later work on this thread.
@@ -711,6 +673,35 @@ class ET_EXPERIMENTAL CudaBackend final
     ET_CHECK_OK_OR_RETURN_ERROR(cuda_stream_ret.error());
     cudaStream_t cuda_stream = cuda_stream_ret.get();
 
+    // Ends a capture that an early return would otherwise abandon. Every handle
+    // shares the per-thread stream, so a stream left capturing does not just
+    // break this method: the next delegate on this thread has its kernels
+    // captured instead of run, and later synchronizes fail. Disarmed once the
+    // real end has run.
+    class CaptureGuard {
+     public:
+      ~CaptureGuard() {
+        if (stream_ == nullptr) {
+          return;
+        }
+        cudaGraph_t abandoned = nullptr;
+        if (cudaStreamEndCapture(stream_, &abandoned) == cudaSuccess &&
+            abandoned != nullptr) {
+          (void)cudaGraphDestroy(abandoned);
+        }
+        (void)cudaGetLastError();
+      }
+      void arm(cudaStream_t stream) {
+        stream_ = stream;
+      }
+      void disarm() {
+        stream_ = nullptr;
+      }
+
+     private:
+      cudaStream_t stream_ = nullptr;
+    } capture_guard;
+
     if (is_capture_step) {
       // ----- CUDA graph CAPTURE -----
       ET_LOG(
@@ -725,6 +716,7 @@ class ET_EXPERIMENTAL CudaBackend final
           Internal,
           "cudaStreamBeginCapture failed: %s",
           cudaGetErrorString(cerr));
+      capture_guard.arm(cuda_stream);
     }
 
     AOTIRuntimeError error = handle->run(
@@ -756,6 +748,7 @@ class ET_EXPERIMENTAL CudaBackend final
       // End capture → instantiate graph
       cudaError_t gerr =
           cudaStreamEndCapture(cuda_stream, &handle->cuda_graph_state.graph);
+      capture_guard.disarm();
       ET_CHECK_OR_RETURN_ERROR(
           gerr == cudaSuccess,
           Internal,
@@ -864,11 +857,6 @@ class ET_EXPERIMENTAL CudaBackend final
 
     mutable_state_forget_handle(handle);
 
-    // The CUDA stream is managed by shared_ptr in the handle.
-    // It will be automatically destroyed when the last handle using it
-    // is destroyed. Just reset our reference.
-    handle->cuda_stream.reset();
-
     // NOTE: AOTInductorModelContainerDelete does not work correctly with
     // multiple .so files. Deleting one container frees shared resources,
     // which causes segmentation faults when attempting to delete other
@@ -898,14 +886,6 @@ class ET_EXPERIMENTAL CudaBackend final
   mutable std::mutex cuda_graph_method_mutex_;
   std::string cuda_graph_method_;
 
-  // Shared CUDA stream for all methods. When set (non-null), all methods use
-  // the same stream to ensure proper ordering across methods that hand off
-  // GPU-resident tensors (e.g. encoder -> decoder -> sampler). Created when
-  // use_shared_cuda_stream option is set to true. Managed via shared_ptr so
-  // it's automatically cleaned up when last handle is destroyed.
-  mutable std::mutex cuda_stream_mutex_;
-  std::shared_ptr<cudaStream_t> shared_cuda_stream_ = nullptr;
-
   // Whether to enable cross-method caching for legacy dense-blob artifacts.
   // Toggled by the kWeightSharingAcrossMethods runtime backend option. Default
   // OFF; versioned FQN artifacts do not consult this option.

backends/cuda/runtime/cuda_delegate_handle.h

@@ -64,29 +64,6 @@ struct CudaWeightStorage {
   CudaWeightStorage& operator=(const CudaWeightStorage&) = delete;
 };
 
-// Shared CUDA stream wrapper with proper RAII cleanup.
-// This ensures the stream is destroyed when all handles using it are destroyed.
-struct CudaStreamDeleter {
-  void operator()(cudaStream_t* stream) const {
-    if (stream != nullptr && *stream != nullptr) {
-      (void)cudaStreamDestroy(*stream);
-    }
-    delete stream;
-  }
-};
-
-// Creates a new shared CUDA stream.
-// Returns nullptr on failure.
-inline std::shared_ptr<cudaStream_t> create_cuda_stream() {
-  cudaStream_t stream;
-  cudaError_t err = cudaStreamCreate(&stream);
-  if (err != cudaSuccess) {
-    return nullptr;
-  }
-  return std::shared_ptr<cudaStream_t>(
-      new cudaStream_t(stream), CudaStreamDeleter());
-}
-
 // Phases of the CUDA graph lifecycle for a delegate handle.
 //
 // The transition flow is:
@@ -198,11 +175,8 @@ struct CudaDelegateHandle : public aoti::AOTIDelegateHandle {
   // Extra AOTI metadata used to validate per-FQN weights before binding.
   AOTInductorModelContainerGetConstantDtypeFunc get_constant_dtype{nullptr};
 
-  // CUDA stream for this handle, support both shared mode and single mode.
-  // In shared mode, all cuda delegate handles share the same stream (e.g., for
-  // skip-copy optimization), they will all hold a reference to the same
-  // shared_ptr. The stream is automatically destroyed when the last handle is
-  // destroyed. In single mode, every cuda delegate handle has its own stream.
+  // The per-thread stream, which the runtime owns rather than this handle, so
+  // releasing the holder does not destroy it.
   std::shared_ptr<cudaStream_t> cuda_stream;
 
   // Get the raw CUDA stream pointer for use in CUDA API calls.

backends/cuda/tests/test_coalesced_determinism.py

@@ -0,0 +1,163 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""A coalesced program must give the same answer every time.
+
+A program split across two backends hands buffers from one delegate to the next. If
+the delegates do not agree on a stream, a delegate can read an output before the
+work that writes it has run, and the program returns a different answer on each
+call. Nothing else in the suite runs a program that crosses backends, so nothing
+else would notice.
+
+Needs the TensorRT delegate, so it skips when torch_tensorrt or its ExecuTorch
+runtime is absent.
+"""
+
+import unittest
+
+import torch
+
+
+def _build_coalesced_pte(outdir):
+    """Export a model split across the TensorRT and CUDA backends."""
+    import torch_tensorrt
+    import torch_tensorrt_executorch_runtime  # noqa: F401
+    from executorch.backends.cuda.cuda_backend import CudaBackend
+    from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
+    from executorch.exir import ExecutorchBackendConfig
+
+    class Model(torch.nn.Module):
+        def __init__(self, dim=256, depth=6):
+            super().__init__()
+            self.depth = depth
+            self.register_buffer("shared", torch.randn(dim))
+            self.scales = torch.nn.ParameterList(
+                [torch.nn.Parameter(torch.randn(dim)) for _ in range(depth)]
+            )
+
+        def forward(self, x):
+            x = torch.relu(x)
+            for i in range(self.depth):
+                y = x * self.scales[i]
+                y = torch.relu(y)
+                y = y * self.shared
+                x = x + y
+            return x
+
+    torch.manual_seed(0)
+    model = Model().eval().cuda()
+    gen = torch.Generator(device="cuda").manual_seed(0)
+    inputs = (torch.randn(8, 256, device="cuda", generator=gen),)
+
+    with torch.inference_mode():
+        exported = torch.export.export(model, inputs)
+        # Withhold one operator from TensorRT so the graph has to split, which is
+        # what puts a delegate boundary in the middle of the data flow.
+        graph = torch_tensorrt.dynamo.compile(
+            exported,
+            inputs=list(inputs),
+            enabled_precisions={torch.float32},
+            min_block_size=1,
+            truncate_double=True,
+            torch_executed_ops={"torch.ops.aten.mul.Tensor"},
+        )
+
+    import os
+
+    pte = os.path.join(outdir, "coalesced.pte")
+    spec = CudaBackend.generate_method_name_compile_spec("forward")
+    torch_tensorrt.save(
+        graph,
+        pte,
+        output_format="executorch",
+        retrace=False,
+        arg_inputs=list(inputs),
+        partitioners=[CudaPartitioner([spec])],
+        backend_config=ExecutorchBackendConfig(),
+    )
+
+    islands = sum(
+        1
+        for node in graph.graph.nodes
+        if node.op == "call_module"
+        and ("_run_on_acc" in node.name or "_run_on_gpu" in node.name)
+    )
+    return pte, inputs, islands
+
+
+class TestCoalescedDeterminism(unittest.TestCase):
+    def setUp(self):
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA not available")
+        try:
+            import torch_tensorrt
+            import torch_tensorrt_executorch_runtime  # noqa: F401
+        except (ImportError, OSError) as error:
+            # OSError as well, since a shared library that fails to load raises that
+            # rather than ImportError.
+            self.skipTest("TensorRT delegate not installed: %s" % error)
+        # Importing is not enough: saving in this format needs the C++ runtime, and
+        # without it the export raises rather than this test skipping.
+        if not getattr(
+            torch_tensorrt.ENABLED_FEATURES, "torch_tensorrt_runtime", False
+        ):
+            self.skipTest("the TensorRT delegate is installed without its runtime")
+
+    def test_repeated_execution_agrees(self):
+        import os
+        import tempfile
+
+        from executorch.runtime import Runtime
+
+        with tempfile.TemporaryDirectory() as outdir:
+            pte, inputs, islands = _build_coalesced_pte(outdir)
+            # Read the program that will actually run, rather than the graph it
+            # came from: the island count is fixed before the ExecuTorch lowering,
+            # so it cannot say whether both backends ended up in the file.
+            with open(pte, "rb") as f:
+                written = f.read()
+            self.assertIn(
+                b"TensorRTBackend",
+                written,
+                "no TensorRT delegate in the saved program",
+            )
+            self.assertIn(
+                b"CudaBackend", written, "no CUDA delegate in the saved program"
+            )
+
+            weights = [
+                os.path.join(outdir, name)
+                for name in sorted(os.listdir(outdir))
+                if name.endswith(".ptd")
+            ]
+            runtime = Runtime.get()
+            program = (
+                runtime.load_program(pte, data_path=weights[0])
+                if weights
+                else runtime.load_program(pte)
+            )
+            method = program.load_method("forward")
+
+            host_inputs = [t.cpu() for t in inputs]
+            first = method.execute(host_inputs)
+            first = first[0] if isinstance(first, (list, tuple)) else first
+            reference = first.clone()
+
+            # Exact equality, not a tolerance: a delegate reading a buffer early
+            # produces a different answer, not a slightly different one.
+            for run in range(1, 100):
+                out = method.execute(host_inputs)
+                out = out[0] if isinstance(out, (list, tuple)) else out
+                self.assertTrue(
+                    torch.equal(out, reference),
+                    "run %d of a %d-delegate program disagreed with run 0: "
+                    "sum %r versus %r"
+                    % (run, islands, float(out.sum()), float(reference.sum())),
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()