diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 29cba8b5ada..8f72bc6bb21 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/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 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 get_shared_cuda_stream() const { - std::lock_guard 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 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,14 @@ class ET_EXPERIMENTAL CudaBackend final "effect; ignoring it.", kSkipCopyOutputToCpuForMethod); } else if (std::strcmp(option.key, kUseSharedCudaStream) == 0) { - if (auto* val = std::get_if(&option.value)) { - if (*val) { - create_shared_cuda_stream(); - } - } else { - ET_LOG(Error, "Option %s must be a boolean.", kUseSharedCudaStream); - return Error::InvalidArgument; - } + ET_LOG( + Info, + "Runtime backend option '%s' is DEPRECATED and has no effect. Methods " + "now run on the calling thread's stream, which orders methods called " + "from one thread. It does not order methods called from different " + "threads, which this option did, so a caller driving methods from more " + "than one thread must order those calls itself.", + kUseSharedCudaStream); } else if (std::strcmp(option.key, kWeightSharingAcrossMethods) == 0) { if (auto* val = std::get_if(&option.value)) { set_weight_sharing_across_methods(*val); @@ -432,30 +405,19 @@ 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 on one thread share that thread's stream, so one delegate's + // output is ordered against the next one's read, which a stream per handle + // left unordered. cudaStreamPerThread is a different stream on each host + // thread, so this orders delegates called from the same thread and not + // delegates called from different ones. The TensorRT delegate falls back to + // the same stream, in its executorch backend, so a split program on one + // thread is ordered too. + handle->cuda_stream = 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. @@ -638,6 +600,85 @@ class ET_EXPERIMENTAL CudaBackend final std::vector slim_inputs(n_inputs); std::vector slim_outputs(n_outputs); + // Undoes a capture attempt that an early return would otherwise abandon. + // + // The stream matters because handles share the per-thread stream: one left + // capturing means the next delegate on this thread has its kernels captured + // instead of run, and later synchronizes fail. The handle matters just as + // much. The buffers this attempt pinned would otherwise stay in their + // vectors with the phase still at warmup and no steps left, so the next + // call captures again and appends a second set. Replay then reads the + // second set while the input copies target the first, and every execute + // returns whatever those buffers held at capture time, with nothing + // reporting an error. + // + // Disarmed once the graph is instantiated and the state is consistent. + class CaptureGuard { + public: + ~CaptureGuard() { + if (state_ == nullptr) { + return; + } + // Only if capture actually began: the guard is armed before the buffers + // are pinned, so it also covers failures that happen before that point. + if (stream_ != nullptr) { + cudaGraph_t abandoned = nullptr; + const cudaError_t err = cudaStreamEndCapture(stream_, &abandoned); + if (err == cudaSuccess) { + if (abandoned != nullptr) { + (void)cudaGraphDestroy(abandoned); + } + } else { + // Only the status this destructor produced, so an error belonging + // to another user of this thread's stream is left where it was. + (void)cudaGetLastError(); + } + } + + // Free what this attempt pinned and put the method back where it was, + // so the next call retries from a clean state instead of capturing on + // top of this one. + for (void* ptr : state_->static_input_ptrs) { + (void)cudaFree(ptr); + } + state_->static_input_ptrs.clear(); + state_->static_output_ptrs.clear(); + state_->static_input_nbytes.clear(); + state_->static_output_nbytes.clear(); + if (state_->graph != nullptr) { + (void)cudaGraphDestroy(state_->graph); + state_->graph = nullptr; + } + if (state_->graph_exec != nullptr) { + (void)cudaGraphExecDestroy(state_->graph_exec); + state_->graph_exec = nullptr; + } + state_->phase = CudaGraphPhase::Warmup; + state_->warmup_remaining = kCudaGraphWarmupSteps; + (void)cudaGetLastError(); + } + // Before capture begins. From here a failure still unwinds the pinned + // buffers. + void arm(cuda::CudaGraphState* state) { + state_ = state; + } + // Once capture is running, so the destructor also ends it. + void arm_capture(cudaStream_t stream) { + stream_ = stream; + } + void disarm() { + state_ = nullptr; + } + + private: + cudaStream_t stream_ = nullptr; + cuda::CudaGraphState* state_ = nullptr; + } capture_guard; + + if (is_capture_step) { + capture_guard.arm(&handle->cuda_graph_state); + } + // Process input tensors: wrap the GPU-resident ETensor buffers directly. for (size_t i = 0; i < n_inputs; i++) { auto* et_input = &(args[i]->toTensor()); @@ -725,6 +766,7 @@ class ET_EXPERIMENTAL CudaBackend final Internal, "cudaStreamBeginCapture failed: %s", cudaGetErrorString(cerr)); + capture_guard.arm_capture(cuda_stream); } AOTIRuntimeError error = handle->run( @@ -780,6 +822,7 @@ class ET_EXPERIMENTAL CudaBackend final } handle->cuda_graph_state.phase = CudaGraphPhase::Replay; + capture_guard.disarm(); ET_LOG( Info, "CUDA graph: captured and instantiated for '%s'", @@ -864,11 +907,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 +936,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 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. diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index 32144ce139e..f53c7498d40 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/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 create_cuda_stream() { - cudaStream_t stream; - cudaError_t err = cudaStreamCreate(&stream); - if (err != cudaSuccess) { - return nullptr; - } - return std::shared_ptr( - new cudaStream_t(stream), CudaStreamDeleter()); -} - // Phases of the CUDA graph lifecycle for a delegate handle. // // The transition flow is: @@ -198,22 +175,14 @@ 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. - std::shared_ptr cuda_stream; + // The per-thread stream. Nothing owns it: the value is a fixed sentinel the + // driver resolves to a different stream on each host thread, so releasing the + // holder destroys nothing. + cudaStream_t cuda_stream = nullptr; - // Get the raw CUDA stream pointer for use in CUDA API calls. - // Returns nullptr if no stream is set. + // The stream this handle's work runs on. cudaStream_t get_cuda_stream() const { - return cuda_stream ? *cuda_stream : nullptr; - } - - // Check if this handle has a valid CUDA stream. - bool has_cuda_stream() const { - return cuda_stream != nullptr && *cuda_stream != nullptr; + return cuda_stream; } // CUDA graph state (warmup, capture, replay, static buffers) diff --git a/backends/cuda/tests/test_coalesced_determinism.py b/backends/cuda/tests/test_coalesced_determinism.py new file mode 100644 index 00000000000..594d2317bbc --- /dev/null +++ b/backends/cuda/tests/test_coalesced_determinism.py @@ -0,0 +1,164 @@ +# 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 os +import tempfile +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 + # One buffer per step rather than one shared across all of them. A + # single constant read by every island becomes the same placeholder + # name several times over in the flattened graph, which is a separate + # export defect and would fail here before reaching the run loop. + for i in range(depth): + self.register_buffer("mix%d" % i, 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 * getattr(self, "mix%d" % i) + 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"}, + ) + + 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(), + ) + + return pte, inputs + + +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): + + from executorch.runtime import Runtime + + with tempfile.TemporaryDirectory() as outdir: + pte, inputs = _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") + ] + # load_program takes one path, so a second file would be dropped and the + # failure would look like a runtime bug. + self.assertLessEqual(len(weights), 1, "expected at most one weights file") + 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 + differing = int((out != reference).sum()) + largest = float((out - reference).abs().max()) + self.assertTrue( + torch.equal(out, reference), + "run %d disagreed with run 0 in %d element(s), largest " + "difference %g" % (run, differing, largest), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.cpp b/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.cpp index 8ffd20aa82e..28969164e2b 100644 --- a/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.cpp +++ b/examples/models/muse-glimmer/runtime/engine/muse_glimmer_engine.cpp @@ -297,9 +297,7 @@ Result> build_muse_glimmer_module( /*share_memory_arenas=*/share_memory_arenas); #ifdef EXECUTORCH_BUILD_CUDA - executorch::runtime::BackendOptions<3> cuda_opts; - ET_CHECK_OK_OR_RETURN_ERROR( - cuda_opts.set_option("use_shared_cuda_stream", true)); + executorch::runtime::BackendOptions<2> cuda_opts; ET_CHECK_OK_OR_RETURN_ERROR( cuda_opts.set_option("weight_sharing_across_methods", true)); if (config.enable_cuda_graph) { diff --git a/extension/asr/runner/seq2seq_runner.cpp b/extension/asr/runner/seq2seq_runner.cpp index 35d430a6b28..b79197451a4 100644 --- a/extension/asr/runner/seq2seq_runner.cpp +++ b/extension/asr/runner/seq2seq_runner.cpp @@ -16,8 +16,6 @@ #include #include #include -#include -#include #include #include #include @@ -108,30 +106,15 @@ Error Seq2SeqRunner::load() { static_cast(method_names.count(kEncoderMethodName)), static_cast(method_names.count(kDecoderMethodName))); -#ifdef CUDA_AVAILABLE - // IMPORTANT: Set backend options BEFORE loading methods. - // The backend's init() is called during load_method(), which creates CUDA - // streams. We must configure shared stream mode before any init() calls. - // - // Keep encoder/decoder outputs on device and pass decoder logits directly - // into the sampler. With device memory planning, delegate inputs/outputs are - // GPU-resident and graph-level et_copy ops handle host<->device transfers; - // the export-time skip_d2h_for_method_outputs / skip_h2d_for_method_inputs - // flags elide the unnecessary copies. A shared CUDA stream is still required - // to guarantee correct ordering across methods when outputs stay on GPU. - executorch::runtime::BackendOptions<1> backend_options; - ET_CHECK_OK_OR_RETURN_ERROR( - backend_options.set_option("use_shared_cuda_stream", true)); - - const auto opt_err = - executorch::runtime::set_option("CudaBackend", backend_options.view()); - if (opt_err != ::executorch::runtime::Error::Ok) { - ET_LOG( - Error, - "Failed to set CUDA backend options: %d", - static_cast(opt_err)); - } -#endif + // Encoder and decoder outputs stay on device and the decoder logits go + // straight into the sampler. With device memory planning the delegate inputs + // and outputs are GPU-resident and graph-level et_copy ops handle host to + // device transfers, while the export-time skip_d2h_for_method_outputs and + // skip_h2d_for_method_inputs flags elide the copies that are not needed. + // Ordering across methods comes from every method running on the calling + // thread's stream, so nothing has to be configured here. That holds because + // this runner drives all three methods from one thread; a caller spreading + // them across threads would have to order them itself. ET_CHECK_OK_OR_RETURN_ERROR(module_->load_method(kEncoderMethodName)); encoder_method_loaded_ = true; diff --git a/extension/cuda/runtime_api.h b/extension/cuda/runtime_api.h index bae5c6a79bf..55f0634adc9 100644 --- a/extension/cuda/runtime_api.h +++ b/extension/cuda/runtime_api.h @@ -137,6 +137,8 @@ inline cudaError_t cudaStreamBeginCapture( return hipStreamBeginCapture(stream, mode); } +#define cudaStreamPerThread hipStreamPerThread + inline cudaError_t cudaStreamCreate(cudaStream_t* stream) { return hipStreamCreate(stream); }