From b3c12b1cd2c1a6dacca3978fa6b0cc2e2a08a125 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 30 Aug 2026 06:33:17 +0000 Subject: [PATCH 1/3] Run CUDA delegates on the per-thread stream Each delegate handle created its own CUDA stream. Delegates in one program run one after another, so a stream apiece bought no ordering between them: a delegate could read an input while the delegate that produces it still had work queued on a different stream. A program that is entirely one backend does not notice while it stays on one thread, because the stream is then the same throughout. A program split across the CUDA and TensorRT backends does. The TensorRT delegate runs on the per-thread stream when no caller stream is installed, so the two backends sat on different streams with nothing between them. Such a program returned a different answer on almost every call: executing one loaded program forty times produced nine distinct output sums, one of them correct. Every handle now takes the per-thread stream, which makes the ordering the graph expresses the ordering the device sees, in either direction and between two CUDA delegates, with no event handshake needed because both backends end up on the same stream. A caller stream still takes precedence per execute. The guarantee is per thread, and that is the one thing to get right about this change. cudaStreamPerThread resolves to a different stream on each host thread, so delegates called from one thread are ordered and delegates called from different threads are not. That is a real loss for one case. use_shared_cuda_stream created a single stream held by the backend and handed it to every handle on every thread, so a caller running the encoder on one thread and the decoder on another did get ordering from it. The option is now accepted and ignored, and its log line says plainly that ordering is per thread and that such a caller has to order the calls itself. Both in-repo callers drive their methods from one thread and no longer set it. Two smaller consequences of one stream rather than many. Two independent programs submitted from one thread with no synchronization between them now run one after the other rather than overlapping, measured at about 2x on two equal single-block kernels; a caller stream per program restores the overlap. And a handle no longer owns a stream that destroy could free while the thread-local table still pointed at it, which was a real dangling-stream path before. Because one stream is now shared, a capture that an error abandons would otherwise leave that stream capturing for every later delegate on the thread. A scope guard ends the capture and also frees the buffers the attempt pinned and resets the method to warmup, so a retry captures from a clean state. Without that second half the retry appends a second set of static buffers, replay reads one set while the input copies target the other, and every execute returns the values captured at that moment with nothing reporting an error, which is worse than the loud failure it replaced. Test plan: backends/cuda/tests/test_coalesced_determinism.py exports a program split across both backends, checks the saved program contains both backends, runs it a hundred times on one loaded program, and requires every result to equal the first exactly. It passes with this change and fails without it at the second run. Each step of the model carries its own buffer: one constant read by every island collapses to the same placeholder name several times in the flattened graph, which is a separate export defect and would fail this test before it reached the run loop. Verified separately on an H100, since no test covers it: a capture step that fails after capture began leaves the stream clean and the method's buffers freed, and three replays after the retry return the current answer. With the guard restoring only the stream, all three returned the capture-time value instead. Measured on Linux aarch64, sm_110, on a program of twenty-five delegates: before, thirty-nine of forty runs were wrong; after, seven thousand six hundred consecutive runs were correct across several processes. Programs of a single delegate on either backend were already correct and stayed so over five hundred runs each. Median latency for that split program went from about 850 to about 940 microseconds, which is the ordering that was previously skipped. A single delegate program is unchanged. The three existing C++ test binaries for this backend pass. The Python suite for this backend has eighteen failures on this machine both with and without this change, so it introduces no regressions; those are an unrelated gap in matmul and convolution lowering on this architecture. cudaStreamPerThread has no alias in the HIP compatibility header, so this adds one, or the ROCm build of this backend would not compile. That alias was checked by hand against a stub of the ROCm definitions and not by a ROCm build: the ROCm and Windows jobs are skipped for pull requests from a fork, so nothing in CI compiled it here. Not covered: the test above needs the TensorRT delegate, which nothing in the repo installs, so it is collected and skipped in CI, and that skip is silent because the job clears pytest's reporting flags. Installing the delegate there would not fix it either, because that job runs on x86 where this reordering does not reproduce; real coverage needs a job on an architecture where it does. Also not covered: Windows, and the CUDA graph capture and replay paths beyond confirming that a program using them runs and agrees with itself over forty runs. --- backends/cuda/runtime/cuda_backend.cpp | 176 ++++++++++-------- backends/cuda/runtime/cuda_delegate_handle.h | 43 +---- .../cuda/tests/test_coalesced_determinism.py | 164 ++++++++++++++++ .../runtime/engine/muse_glimmer_engine.cpp | 4 +- extension/asr/runner/seq2seq_runner.cpp | 35 +--- extension/cuda/runtime_api.h | 2 + 6 files changed, 285 insertions(+), 139 deletions(-) create mode 100644 backends/cuda/tests/test_coalesced_determinism.py diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 4322d2da15a..5d2af640db4 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); } From 828c95c37ff90b429825571f004e8c3672795d26 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 1 Sep 2026 15:19:58 -0700 Subject: [PATCH 2/3] Refuse the deprecated stream option, and end the capture before freeing Review found several problems with the previous revision. The deprecated option was accepted and ignored. That option was the only thing ordering methods driven from different threads, so a caller who set it for that reason kept building, kept running, and got unordered device work and wrong numbers. The one signal was an Info log, and a default Release configure compiles ET_LOG out entirely, so in practice there was no signal at all. It now returns NotSupported when set to true, matching the shape the graph option already uses for ROCm, and InvalidArgument for a non-boolean, which the previous revision had stopped checking. The long explanation moved into a comment: formatted into the runtime's 256 byte log buffer that message came to 338 characters and was cut mid-word, losing exactly the sentence telling the caller what to do. The capture teardown ran in the wrong order. The tensor cleanup guard was declared after the capture guard, so it was destroyed first and issued its frees into a stream that was still capturing. Measured on an H100: a free on a capturing stream returns invalid argument and the block is not returned, 64 MiB in the probe. Ending the capture is now its own guard, declared after the cleanup so it runs before it. With that order the unwind ends the capture, deletes the tensors, then frees, and nothing leaks. Three more, all on paths an error takes: The guard was disarmed before the first graph launch, the output copies and the final synchronize, so a failure in any of them left the method in the replay phase with the static output pointers recorded while the cleanup guard deleted those same buffers. The disarm moved past the last failure point, and the AOTI-owned outputs are released from the cleanup list before the copies rather than after. A static input buffer was tracked only after its seeding copy, so a failed copy leaked it. It is tracked as soon as the allocation succeeds. On capture failure the guard rewound to a full warmup, so a capture that can never succeed returned an error on one call in four for the life of the process. It now disables graphs for that method, which fails once and then runs eagerly. Also: the guard destroyed the graph before its exec, the opposite order from ~CudaGraphState, so the three teardown copies had already drifted; the handle's stream field defaulted to null, which is the legacy default stream and a genuinely different stream from the per-thread one the comment claimed, and it is now initialised to that sentinel so the invariant holds by construction; a comment claimed cudaGetLastError only clears this call's error, which it does not; a nine line comment in the ASR runner survived the code it described and explained export flags that file does not touch; the CMake block defining CUDA_AVAILABLE for the ASR runner lost its last consumer and the top level already defines it; and the HIP alias moved up with the other value aliases with a note on why it has to stay a macro. Test plan: The determinism test now compares the first result against eager before the run to run loop. Comparing a program only against itself passes a delegate that reads a stale buffer the same way every call, which is the same class of bug this change fixes. Verified on an H100, none of it covered by CI: a free into a still-capturing stream returns invalid argument and leaks; with the guards in this order the unwind ends the capture first and the free succeeds one executable graph replayed from two threads on two streams returns no error and serializes, so the refusal to mix a caller stream with graphs is about the shared static buffers and not a CUDA rule; the comment says that now the deprecation message formats to 338 characters into a 256 byte buffer hipStreamPerThread is ((hipStream_t)2), a cast to a pointer type, so it cannot be an inline constexpr like its neighbours Both guard classes were extracted and compiled standalone with nvcc, and the handle's new default was checked to equal cudaStreamPerThread (0x2) where the old one was null. clang-format is clean on all four C++ files; black and usort are clean on the test. The full backend still has no CI job that builds it from a fork, so this is not a build of the backend itself. --- backends/cuda/runtime/cuda_backend.cpp | 157 +++++++++++------- backends/cuda/runtime/cuda_delegate_handle.h | 7 +- .../cuda/tests/test_coalesced_determinism.py | 16 +- extension/asr/runner/CMakeLists.txt | 16 -- extension/asr/runner/seq2seq_runner.cpp | 12 +- extension/cuda/runtime_api.h | 6 +- 6 files changed, 122 insertions(+), 92 deletions(-) diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 5d2af640db4..d6942c9e265 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -235,14 +235,25 @@ class ET_EXPERIMENTAL CudaBackend final "effect; ignoring it.", kSkipCopyOutputToCpuForMethod); } else if (std::strcmp(option.key, kUseSharedCudaStream) == 0) { - 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); + // Refused rather than ignored: this option was the only thing ordering + // methods driven from different threads, so silently dropping it would + // give such a caller unordered device work and wrong results. Methods + // now run on the calling thread's stream, which orders methods called + // from one thread but not across threads; a caller that needs that must + // order the calls itself. + if (auto* val = std::get_if(&option.value)) { + if (*val) { + ET_LOG( + Error, + "Option %s is deprecated and no longer orders methods across " + "threads. See the comment at this check for what to do instead.", + kUseSharedCudaStream); + return Error::NotSupported; + } + } else { + ET_LOG(Error, "Option %s must be a boolean.", kUseSharedCudaStream); + return Error::InvalidArgument; + } } else if (std::strcmp(option.key, kWeightSharingAcrossMethods) == 0) { if (auto* val = std::get_if(&option.value)) { set_weight_sharing_across_methods(*val); @@ -456,8 +467,10 @@ class ET_EXPERIMENTAL CudaBackend final const std::optional caller_stream = executorch::extension::cuda::getCallerStream(); - // A captured CUDA graph is bound to its capture stream and cannot be safely - // replayed on a different, caller-provided stream. + // Replaying a captured graph on a caller-provided stream is not itself a + // CUDA error, but the static buffers this path pins are shared by every + // replay, so two callers on two streams would race over them. Refused + // rather than synchronized, which predates this change. ET_CHECK_OR_RETURN_ERROR( !(caller_stream && handle->cuda_graph_state.phase != CudaGraphPhase::Disabled), @@ -602,42 +615,28 @@ class ET_EXPERIMENTAL CudaBackend final // 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. + // 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. + // + // Ending the capture itself is a separate guard, declared after the tensor + // cleanup below so that it runs before it: freeing a device buffer on a + // still-capturing stream fails with invalid argument and leaks the block. // - // Disarmed once the graph is instantiated and the state is consistent. + // Disarmed once the capture step has fully succeeded. 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. + // Free what this attempt pinned and disable graphs for this method, so + // a capture that cannot succeed costs one error rather than one on + // every fourth call for the life of the process. Eager execution is + // correct, just slower. for (void* ptr : state_->static_input_ptrs) { (void)cudaFree(ptr); } @@ -645,16 +644,17 @@ class ET_EXPERIMENTAL CudaBackend final 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; - } + // Same order as ~CudaGraphState: the exec depends on the graph. if (state_->graph_exec != nullptr) { (void)cudaGraphExecDestroy(state_->graph_exec); state_->graph_exec = nullptr; } - state_->phase = CudaGraphPhase::Warmup; - state_->warmup_remaining = kCudaGraphWarmupSteps; + if (state_->graph != nullptr) { + (void)cudaGraphDestroy(state_->graph); + state_->graph = nullptr; + } + state_->phase = CudaGraphPhase::Disabled; + state_->warmup_remaining = 0; (void)cudaGetLastError(); } // Before capture begins. From here a failure still unwinds the pinned @@ -662,16 +662,11 @@ class ET_EXPERIMENTAL CudaBackend final 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; @@ -697,15 +692,17 @@ class ET_EXPERIMENTAL CudaBackend final i, cudaGetErrorString(merr)); + // Tracked before the seeding copy, so a failed copy still unwinds + // through the guard instead of leaking this allocation. + handle->cuda_graph_state.static_input_ptrs.push_back(static_ptr); + handle->cuda_graph_state.static_input_nbytes.push_back(nbytes); + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpy( static_ptr, et_input->const_data_ptr(), nbytes, cudaMemcpyDeviceToDevice)); - handle->cuda_graph_state.static_input_ptrs.push_back(static_ptr); - handle->cuda_graph_state.static_input_nbytes.push_back(nbytes); - slim_inputs[i] = make_slimtensor_from_blob_with_etensor_metadata( static_ptr, et_input); continue; @@ -745,6 +742,43 @@ class ET_EXPERIMENTAL CudaBackend final } }); + // Ends a capture that an early return would otherwise leave running. + // Declared after `cleanup` so it is destroyed before it: a device buffer + // freed on a still-capturing stream fails with invalid argument and leaks + // the block. + // + // Leaving the stream capturing matters because handles share the per-thread + // stream: the next delegate on this thread would have its kernels captured + // instead of run, and later synchronizes would fail. + class EndCaptureGuard { + public: + ~EndCaptureGuard() { + if (stream_ == nullptr) { + return; + } + cudaGraph_t abandoned = nullptr; + if (cudaStreamEndCapture(stream_, &abandoned) == cudaSuccess) { + if (abandoned != nullptr) { + (void)cudaGraphDestroy(abandoned); + } + } else { + // Clears the sticky error so the next unrelated CUDA call on this + // thread does not inherit it. This clears whatever error is pending, + // not only the one from above. + (void)cudaGetLastError(); + } + } + void arm(cudaStream_t stream) { + stream_ = stream; + } + void disarm() { + stream_ = nullptr; + } + + private: + cudaStream_t stream_ = nullptr; + } end_capture_guard; + // Run the AOTI container. // NOTE: run() steals input handles (RAII wraps them at the start of // run_impl) and may replace output handles with its own. @@ -766,7 +800,7 @@ class ET_EXPERIMENTAL CudaBackend final Internal, "cudaStreamBeginCapture failed: %s", cudaGetErrorString(cerr)); - capture_guard.arm_capture(cuda_stream); + end_capture_guard.arm(cuda_stream); } AOTIRuntimeError error = handle->run( @@ -798,6 +832,9 @@ class ET_EXPERIMENTAL CudaBackend final // End capture → instantiate graph cudaError_t gerr = cudaStreamEndCapture(cuda_stream, &handle->cuda_graph_state.graph); + // The stream has left capture either way, so the guard must not end it + // again; the state guard below still unwinds what the attempt pinned. + end_capture_guard.disarm(); ET_CHECK_OR_RETURN_ERROR( gerr == cudaSuccess, Internal, @@ -814,15 +851,18 @@ class ET_EXPERIMENTAL CudaBackend final "cudaGraphInstantiate failed: %s", cudaGetErrorString(gerr)); - // Record static output pointers (stable under graph replay) + // Record static output pointers (stable under graph replay). Releasing + // them from slim_outputs here, before the copies below, keeps the cleanup + // guard from deleting buffers the AOTI runtime owns if one of those + // copies fails. for (size_t i = 0; i < n_outputs; i++) { SlimTensor* out = slim_outputs[i]; handle->cuda_graph_state.static_output_ptrs.push_back(out->data_ptr()); handle->cuda_graph_state.static_output_nbytes.push_back(out->nbytes()); + slim_outputs[i] = nullptr; } handle->cuda_graph_state.phase = CudaGraphPhase::Replay; - capture_guard.disarm(); ET_LOG( Info, "CUDA graph: captured and instantiated for '%s'", @@ -845,11 +885,12 @@ class ET_EXPERIMENTAL CudaBackend final handle->cuda_graph_state.static_output_nbytes[i], cudaMemcpyDeviceToDevice, cuda_stream)); - // Don't delete — static buffers are owned by the AOTI runtime. - slim_outputs[i] = nullptr; } ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamSynchronize(cuda_stream)); + // Last failure point is behind us, so the captured state is now the state + // the next call should replay from. + capture_guard.disarm(); return Error::Ok; } diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index f53c7498d40..d4af1ed0741 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -170,15 +170,16 @@ struct CudaGraphState { }; // CUDA-specific delegate handle that extends AOTIDelegateHandle. -// This consolidates CUDA stream management into a single location. struct CudaDelegateHandle : public aoti::AOTIDelegateHandle { // Extra AOTI metadata used to validate per-FQN weights before binding. AOTInductorModelContainerGetConstantDtypeFunc get_constant_dtype{nullptr}; // 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; + // holder destroys nothing. Initialised to that sentinel rather than null, + // because null is the legacy default stream, which is a different stream and + // would silently drop the per-thread ordering this handle relies on. + cudaStream_t cuda_stream = cudaStreamPerThread; // The stream this handle's work runs on. cudaStream_t get_cuda_stream() const { diff --git a/backends/cuda/tests/test_coalesced_determinism.py b/backends/cuda/tests/test_coalesced_determinism.py index 594d2317bbc..e9af8f0dbfd 100644 --- a/backends/cuda/tests/test_coalesced_determinism.py +++ b/backends/cuda/tests/test_coalesced_determinism.py @@ -24,7 +24,11 @@ def _build_coalesced_pte(outdir): - """Export a model split across the TensorRT and CUDA backends.""" + """Export a model split across the TensorRT and CUDA backends. + + Returns the program path, the inputs, and the eager result for those inputs, so + the caller can check the answer is right and not only self-consistent. + """ import torch_tensorrt import torch_tensorrt_executorch_runtime # noqa: F401 from executorch.backends.cuda.cuda_backend import CudaBackend @@ -60,6 +64,7 @@ def forward(self, x): inputs = (torch.randn(8, 256, device="cuda", generator=gen),) with torch.inference_mode(): + eager = model(*inputs).cpu() 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. @@ -84,7 +89,7 @@ def forward(self, x): backend_config=ExecutorchBackendConfig(), ) - return pte, inputs + return pte, inputs, eager class TestCoalescedDeterminism(unittest.TestCase): @@ -110,7 +115,7 @@ def test_repeated_execution_agrees(self): from executorch.runtime import Runtime with tempfile.TemporaryDirectory() as outdir: - pte, inputs = _build_coalesced_pte(outdir) + pte, inputs, eager = _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. @@ -146,6 +151,11 @@ def test_repeated_execution_agrees(self): first = first[0] if isinstance(first, (list, tuple)) else first reference = first.clone() + # Against eager with a tolerance, because a delegate that reads a stale + # buffer the same way on every call is self-consistent and still wrong, + # which the run-to-run check below cannot see. + torch.testing.assert_close(reference, eager, rtol=1e-3, atol=1e-3) + # 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): diff --git a/extension/asr/runner/CMakeLists.txt b/extension/asr/runner/CMakeLists.txt index b47cddaf48c..9b5a1fd1140 100644 --- a/extension/asr/runner/CMakeLists.txt +++ b/extension/asr/runner/CMakeLists.txt @@ -37,22 +37,6 @@ set_target_properties( extension_asr_runner PROPERTIES POSITION_INDEPENDENT_CODE ON ) -# If the project is configured to build with CUDA support, try to find a CUDA -# runtime (prefer the CUDAToolkit package). If found, expose a compile-time -# macro so sources can conditionally compile CUDA-aware code. -if(EXECUTORCH_BUILD_CUDA) - find_package(CUDAToolkit QUIET) - if(CUDAToolkit_FOUND) - target_compile_definitions(extension_asr_runner PUBLIC CUDA_AVAILABLE) - message(STATUS "CUDAToolkit found; defining CUDA_AVAILABLE for ASR runner") - else() - message( - STATUS - "CUDA requested (EXECUTORCH_BUILD_CUDA=ON) but no CUDA runtime found" - ) - endif() -endif() - install( TARGETS extension_asr_runner EXPORT ExecuTorchTargets diff --git a/extension/asr/runner/seq2seq_runner.cpp b/extension/asr/runner/seq2seq_runner.cpp index b79197451a4..fef57aa12bb 100644 --- a/extension/asr/runner/seq2seq_runner.cpp +++ b/extension/asr/runner/seq2seq_runner.cpp @@ -106,16 +106,8 @@ Error Seq2SeqRunner::load() { static_cast(method_names.count(kEncoderMethodName)), static_cast(method_names.count(kDecoderMethodName))); - // 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. - + // This runner drives its methods from one thread, so they are ordered by + // running on that thread's CUDA stream and nothing has to be configured here. 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 55f0634adc9..b150055ce5a 100644 --- a/extension/cuda/runtime_api.h +++ b/extension/cuda/runtime_api.h @@ -34,6 +34,10 @@ inline constexpr cudaStreamCaptureMode cudaStreamCaptureModeRelaxed = inline constexpr unsigned long long cudaGraphInstantiateFlagAutoFreeOnLaunch = hipGraphInstantiateFlagAutoFreeOnLaunch; +// A macro, unlike the aliases above, because hipStreamPerThread casts an +// integer to a pointer type and so is not a constant expression. +#define cudaStreamPerThread hipStreamPerThread + struct cudaPointerAttributes { cudaMemoryType type{}; int device = -1; @@ -137,8 +141,6 @@ inline cudaError_t cudaStreamBeginCapture( return hipStreamBeginCapture(stream, mode); } -#define cudaStreamPerThread hipStreamPerThread - inline cudaError_t cudaStreamCreate(cudaStream_t* stream) { return hipStreamCreate(stream); } From 374821d4e1bc9c6f4d0dd3cbe5b19872ac851417 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 1 Sep 2026 16:15:08 -0700 Subject: [PATCH 3/3] Guard the per-thread stream macro against a redefinition The HIP compatibility section defines cudaStreamPerThread unconditionally, so a HIP header that defines the CUDA spelling itself would warn. Guarded. --- extension/cuda/runtime_api.h | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/extension/cuda/runtime_api.h b/extension/cuda/runtime_api.h index b150055ce5a..e95bf36994c 100644 --- a/extension/cuda/runtime_api.h +++ b/extension/cuda/runtime_api.h @@ -35,8 +35,11 @@ inline constexpr unsigned long long cudaGraphInstantiateFlagAutoFreeOnLaunch = hipGraphInstantiateFlagAutoFreeOnLaunch; // A macro, unlike the aliases above, because hipStreamPerThread casts an -// integer to a pointer type and so is not a constant expression. +// integer to a pointer type and so is not a constant expression. Guarded in +// case a HIP header ever defines the CUDA spelling itself. +#ifndef cudaStreamPerThread #define cudaStreamPerThread hipStreamPerThread +#endif struct cudaPointerAttributes { cudaMemoryType type{};