refactor(executorch)!: share one caller stream, and ship the CUDA delegate in the runtime wheel

[shoumikhin]

What this does

The ExecuTorch backend carried its own private thread-local for the caller-selected
CUDA stream. This replaces it with ExecuTorch's shared
executorch::extension::cuda::CallerStreamGuard, so every CUDA-capable delegate in a
process reads one shared selection instead of one per backend.

That matters because a Torch-TensorRT export can split a graph across two delegates:
TensorRT claims what it can, and ExecuTorch's CUDA/AOTI delegate fills the gaps. Both
have to observe the same stream, or they cannot share a CUDA green context. Upstream's
own header names exactly that case as the reason the API exists.

The Python runtime now carries both GPU delegates

Reaching that shared library means building ExecuTorch with
EXECUTORCH_BUILD_CUDA=ON, and upstream then links its CUDA backend into the same
_portable_lib this wheel ships:

if(EXECUTORCH_BUILD_CUDA)
list(APPEND _dep_libs aoti_cuda_backend) # registers CudaBackendendif()
target_link_libraries(portable_libPRIVATE${_dep_libs})

That is worth keeping rather than working around. A program split across the two
delegates needs both registered in one runtime to run at all, and ExecuTorch's
published wheel registers only XnnpackBackend. So before this there was no Python
runtime anywhere that could execute such a program, even though exporting one was
already supported. That gap is now closed.

Keeping it means shipping everything the backend needs. The build produces two shared
libraries that ExecuTorch's own wheel does not publish, and the extensions carry a
DT_NEEDED on both:

library	role	shipped before
libextension_cuda.so	owns the caller-stream thread-local every CUDA-capable delegate reads	yes
libaoti_cuda_shims.so	the runtime half of the CUDA backend, called by the kernels compiled into a .pte	no

Shipping only the first left a CUDA backend with half its runtime missing, which
surfaces as an ImportError on a missing shared object. Both ship now, data_loader
gets the same $ORIGIN search path _portable_lib already had so neither extension
depends on module load order to find them, and the runtime test asserts CudaBackend
is registered beside the existing TensorRTBackend and XnnpackBackend assertions.

The static C++ runtime settings, and the existing check for them, now cover every
shared object the wheel ships rather than only the two Python extensions. A shared
library left on the host libstdc++ hands that dependency straight back to
_portable_lib.so through its own DT_NEEDED.

Breaking change

torch_tensorrt::executorch_backend::CudaStreamGuard is removed from the public
header, with no deprecated alias. Callers switch to:

#include<executorch/extension/cuda/caller_stream.h>
executorch::extension::cuda::CallerStreamGuard guard(stream);
module.forward(inputs);

This is deliberate rather than forced. An alias would have shared the same
thread-local, so keeping one would buy source compatibility at the cost of two names
for one primitive. The symbol shipped in v2.13.0 release candidates, so this is a
source break for any C++ consumer of it, taken knowingly.

The green-context path now works, and is verified

This was previously undocumented and untested, and the docs actively disclaimed it.
One .pte whose graph splits across the TensorRT delegate and ExecuTorch's CUDA/AOTI
delegate, driven by a single cuGreenCtxStreamCreate stream on an A100 with 108 SMs
and a green context holding 8 of them:

planned buffer[0] = 16384 bytes on CPU
planned buffer[1] = 32768 bytes on device_type 1
ordinary stream exit 0 first 8 values: 0.6722 0.6722 0.6722 ...
green context with 8 SM(s) exit 0 first 8 values: 0.6722 0.6722 0.6722 ...

against an eager reference of 0.672167. Reproduce by building the reference runner
with -DEXECUTORCH_BUILD_CUDA=ON and passing --green_context_sms=8.

Two things had to be fixed before that was even possible.

Planned buffers ignored their declared device. A two-delegate program could not run
in the reference runner at all:

cuda_backend.cpp:529] Tensor 0 has device_type=CUDA but its data pointer is not
backed by CUDA device memory (cudaMemoryType=0)
method.cpp:1525] CALL_DELEGATE execute failed at instruction 2: 0x12

The TensorRT delegate tolerates host-backed input by staging it. The CUDA delegate
verifies the storage really is device or managed memory. The tensor handed between the
two comes from a memory-planned buffer, and the runner allocated every planned buffer
on the host. ExecuTorch already reports the device per planned buffer through
MethodMeta::memory_planned_buffer_device(), and its own executor_runner allocates
accordingly; this follows that pattern. Before this change the same .pte aborted, so
the new coverage is not vacuous.

The runner could not host the CUDA delegate.EXECUTORCH_BUILD_CUDA was FORCEd
off, so -D on the command line was ignored. It is now overridable. The old comment
gave libtorch as the reason to keep it off, which is not accurate: measured on Linux
x86_64, enabling it needs one companion option and the resulting runner links no
libtorch and no libc10.

Two other defects fixed

The end-to-end assertions accepted wrong results. Fed deliberate garbage, both
passed:

output[0] shape=[999] numel=1 dtype=6
first 8 values: 9 9 2.0000 9 9 9 9 9
shape check: exit 0 value check: exit 0

The shape was never compared, and one correct value anywhere satisfied the value
check, so a stream-ordering regression returning stale output would pass. They now
assert the shape is exactly [2,3,4,4] and that every printed value is 2.0000,
verified in four directions including the garbage above.

The wheel linked a library it did not ship.native/CMakeLists.txt links
extension_cuda, which is a shared library upstream, so both extensions gain a
DT_NEEDED on libextension_cuda.so. It was absent from out_shared_libs and from
setup.py, so importing the wheel would have failed on a missing shared object. Now
installed to lib/ beside the extensions, collected, and copied so the existing
$ORIGIN rpath resolves it. Note ExecuTorch installs it to CMAKE_INSTALL_LIBDIR,
which is lib64 on some distributions, so relying on that alone would have missed it.

The driver API is loaded, not linked

--green_context_sms needs cuGreenCtx* from libcuda. The release build image ships
neither libcuda nor a stub, so nothing can link it there, and the runner has both a
CMake and a Bazel build. The entry points are resolved with dlopen when the flag is
used, so neither build system carries a CUDA driver dependency and a machine without a
driver gets a clear error from the flag rather than a link failure for everyone.

What is verified, and what is not

Verified on an A100: both stream modes on a genuine two-delegate program, both
delegate archives referencing getCallerStream, a single shared
libextension_cuda.so resolved, and the runner linking no libtorch.

Not verified: the device-resident asynchronous return path, so how a green context
interacts with the backend's internal completion event remains untested. The
green-context result is also hand-verified rather than in CI, because the CI
configuration builds the runner without the CUDA delegate. Both limits are stated in
the docs rather than glossed.

executorch-runtime-build and py-core fail on main as well, on all ten matrix
rows, so neither is caused by this change. No ExecuTorch test has run in CI on any
commit; #4525 is what fixes that.

[shoumikhin]

CI caught a real mistake in my previous push, now fixed.

I had removed the LD_LIBRARY_PATH setup around the Bazel test step, on the assumption
that only the caller-stream test needed TensorRT at run time. That was wrong.
test_executorch_binding_names also depends on @tensorrt//:nvinfer, so it failed to
start:

//tests/cpp/executorch:test_caller_stream PASSED
//tests/cpp/executorch:test_executorch_blob_header PASSED
//tests/cpp/executorch:test_executorch_binding_names FAILED
error while loading shared libraries: libnvinfer.so.11:
cannot open shared object file: No such file or directory

The underlying reason is that Bazel's cc_import provides lib/libnvinfer.so, but a
binary linked against it records a DT_NEEDED entry for the versioned soname
libnvinfer.so.11, which is not in the test runfiles. I reproduced this locally
against the same TensorRT SDK and got the identical message, then confirmed that
pointing LD_LIBRARY_PATH at the SDK's lib directory resolves it.

So the test step now locates the versioned library Bazel materialized and passes it
through:

tensorrt_lib_dir="$( find -L "$(bazel info output_base)/external" \ -path '*tensorrt*/lib/libnvinfer.so.*' -printf '%h\n'2>/dev/null | sort -u | head -n1)"
bazel test //tests/cpp/executorch:executorch_backend_tests \
--compilation_mode opt --config=linux --test_output=errors \
--test_env=LD_LIBRARY_PATH="${tensorrt_lib_dir}"

This is still simpler than before. It keeps the sandbox enabled, since
--test_strategy=standalone is no longer needed, and it drops the
CUDA_VISIBLE_DEVICES passthrough, since none of these tests run CUDA work. The glob
only matches versioned sonames, which is what the loader actually looks for.

Worth noting from the same run: test_caller_stream passed, and the full
bazel build //:libtorchtrt completed successfully with the reworked CMake, so the
ELF check and the switch to ExecuTorch's own extension/cuda build are exercised and
working.

[shoumikhin]

Second attempt at the test-runtime fix. My previous one was wrong in a way the logs
made obvious.

I set LD_LIBRARY_PATH to only the TensorRT directory, which replaced the value
instead of adding to it. That fixed libnvinfer.so.11 and immediately broke
libcudart.so.13, so the run went from one failing test to two:

test_executorch_blob_header PASSED
test_caller_stream FAILED <- regressed, previously passed
test_executorch_binding_names FAILED
error while loading shared libraries: libcudart.so.13

Both libraries have the same root cause. Bazel's cc_import ships the unversioned
libnvinfer.so and libcudart.so, but a binary linked against them records
DT_NEEDED entries for the versioned sonames, which are not in the test runfiles.

The step now locates the directory holding each versioned soname and appends to
LD_LIBRARY_PATH, keeping the toolchain entries that were already there:

bazel_external="$(bazel info output_base)/external"for_sonamein libnvinfer.so libcudart.so;do
_dir="$( find -L "${bazel_external}" -path "*/lib*/${_soname}.*" -printf '%h\n'2>/dev/null | sort -u | head -n1)"if [[ -z"${_dir}" ]];thenecho"Could not locate a versioned ${_soname} under ${bazel_external}">&2exit 1
fiexport LD_LIBRARY_PATH="${_dir}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"done

I verified this locally by reproducing both failure modes rather than reasoning about
them. Copying the two libraries somewhere off the default search path and running the
loader with --inhibit-cache, so the system library cache cannot mask the problem:

no LD_LIBRARY_PATH -> libnvinfer.so.11: cannot open shared object file
TensorRT directory only -> libcudart.so.12: cannot open shared object file
both directories appended -> both libraries loaded OK

The middle line is the failure this run hit. The guard also names the specific library
if either directory cannot be found, instead of failing later with a loader error.

This is still simpler than the original version of the step: the sandbox stays enabled,
since --test_strategy=standalone is not needed, and the CUDA_VISIBLE_DEVICES
passthrough stays removed, since none of these tests run CUDA work.

[shoumikhin]

CI is green on the ExecuTorch gate. All three C++ tests pass:

//tests/cpp/executorch:test_caller_stream PASSED in 0.1s
//tests/cpp/executorch:test_executorch_binding_names PASSED in 0.1s
//tests/cpp/executorch:test_executorch_blob_header PASSED in 0.1s
Executed 3 out of 3 tests: 3 tests pass.

The same job also exercised the rest of this change end to end:

• Built libextension_cuda.so through ExecuTorch's own extension/cuda CMake, which
  is the definition this change now reuses instead of re-declaring the target.
• Verified both the CMake-built and the packaged runner carry a real DT_NEEDED entry
  for libextension_cuda.so and import the caller-stream symbols rather than defining
  private copies.
• Ran real inference in both runners inside a CallerStreamGuard, each producing the
  expected values:

output[0] shape=[2,3,4,4] numel=96 dtype=6
first 8 values: 2.0000 2.0000 2.0000 2.0000 2.0000 2.0000 2.0000 2.0000

Remaining failures on this commit are pre-existing on main and unrelated to this
change:

Check	Here	On main
Python-only dynamo runtime, two variants	fail	fail
RTX Python-only dynamo runtime, two variants	fail	fail / cancelled
L0 core python	fail	flaky

The dynamo jobs stop before running any test with HTTP 403 FORBIDDEN for channel pkgs/main, a package-index problem in the job image. The L0 failure is
test_libtorchtrt_linkage.py, which fails on a ctypes open of
libnvinfer_plugin.so.11 that is missing from the runner image; this change touches no
files under tests/py.

Note that a large number of pytorch/executorch checks are also attached to this
commit SHA because it appears in that repository's CI as well. Those are not from this
pull request. Filtering by originating repository, the counts are 65 checks from
pytorch/TensorRT and 100 from pytorch/executorch.

[cehongwang]

This means that eventually enqueueV3 is happening on the caller stream, right? Is it going to be a problem @narendasan

[shoumikhin]

Good question, and the short answer is yes, but that is not new here.

enqueueV3 already ran on the caller-selected stream before this change. On the base
commit the same function had:

cudaStream_t stream = g_user_stream_set ? g_user_stream : cudaStreamPerThread; // line 382
...
if (!ctx->enqueueV3(stream)) { // line 571

This change only swaps where that stream value is read from. It was a thread-local
private to this delegate, and it is now ExecuTorch's shared one. The
enqueueV3 call, the staging copies, the synchronize, and the completion event are all
untouched.

What happens with no guard active

Nothing changes. getCallerStream() returns empty, so the stream is
cudaStreamPerThread and execute() still synchronizes before returning, exactly as
before. Existing callers see identical behavior.

When execute() can return with work still in flight

This is the part worth being precise about. The decision is:

constbool must_sync = output_staged_to_host || input_staged_from_host || !caller_stream_set;

execute() skips the end-of-call synchronize only when all of these hold:

1. the caller scoped a CallerStreamGuard, and
2. every input is already GPU-accessible, and
3. every output is already GPU-accessible.

If any tensor is host-backed, the backend staged it through a temporary device buffer,
so it synchronizes before returning to keep the "results are ready on return" contract.
And with no guard at all it always synchronizes. So the asynchronous path only happens
when the caller explicitly asked for stream semantics and no host memory is involved.

Ordering and lifetime safety

Within execute(), stream is assigned once and never reassigned, and every CUDA
operation uses that same variable: the host-to-device copy, enqueueV3, the
device-to-host copy, the synchronize, and the event record. There is no mixing of the
caller's stream with the default stream.

For the asynchronous path, the backend records a completion event on that stream and
waits on it in two places before it could disturb TensorRT state:

• at the start of the next execute(), before any setInputShape or
  setTensorAddress, since those are host-side calls on the execution context;
• in the handle destructor, before freeing staging buffers and releasing the execution
  context.

It waits on the event rather than the stream, so teardown stays correct even if the
caller has already destroyed their stream.

This is the same pattern ExecuTorch's CUDA backend uses

The caller-stream primitive exists precisely so several CUDA delegates can share one
caller-provided stream. ExecuTorch's own CUDA/AOTI backend reads the same value and
installs it for its whole execution:

// backends/cuda/runtime/cuda_backend.cppconst std::optional<cudaStream_t> caller_stream =
executorch::extension::cuda::getCallerStream();
...
setCurrentCUDAStream(caller_stream.value_or(handle->get_cuda_stream()), 0);

If the two delegates read different stream values, a program that mixes them cannot
order their work, which is the problem this change fixes.

Known limits, stated plainly

• The selected stream must be on the TensorRT engine's device. That is not validated in
  code today; a mismatch surfaces as an enqueueV3 failure. Pre-existing, and this
  change makes the error message name the likely cause.
• CUDA green-context streams need context-aware completion-event handling and are not
  claimed as supported. This change narrows the previous claim rather than widening it.
• The device-resident asynchronous path is not covered end to end by CI, because the
  reference runner's tensors are host-backed and therefore take the synchronized path.

Happy to gate the asynchronous return behind an explicit backend option instead, if you
would prefer that a backend-neutral stream selection never changes when execute()
returns. That would be a small follow-up and would leave the default fully synchronous.

[shoumikhin]

Pushed three commits that close the gap between what this PR enables and what it
demonstrates.

The primary use case now runs, and is verified. One .pte whose graph splits
across the TensorRT delegate and ExecuTorch's CUDA/AOTI delegate, driven by a single
cuGreenCtxStreamCreate stream, on an A100 with 108 SMs and a green context holding
8 of them:

planned buffer[0] = 16384 bytes on CPU
planned buffer[1] = 32768 bytes on device_type 1
ordinary stream exit 0 first 8 values: 0.6722 0.6722 0.6722 ...
green context with 8 SM(s) exit 0 first 8 values: 0.6722 0.6722 0.6722 ...

against an eager reference of 0.672167 for the runner's all-ones input.

Two things had to be fixed before that was possible.

1. A two-delegate program could not run in the reference runner at all:

cuda_backend.cpp:529] Tensor 0 has device_type=CUDA but its data pointer is not
backed by CUDA device memory (cudaMemoryType=0)
method.cpp:1525] CALL_DELEGATE execute failed at instruction 2: 0x12

The TensorRT delegate tolerates host-backed input by staging it. The CUDA delegate
verifies the storage really is device or managed memory. The tensor passed between
the two delegates comes from a memory-planned buffer, and the runner allocated every
planned buffer on the host. ExecuTorch already reports the device per planned buffer
via MethodMeta::memory_planned_buffer_device(), and its own executor_runner
allocates accordingly; the runner now follows that pattern. This is also the reason
the coverage is not vacuous: before the change, the same .pte aborted.

2. EXECUTORCH_BUILD_CUDA was FORCEd off, so -D on the command line was ignored
   and only the TensorRT delegate could ever be built here. It is now overridable.

The old comment's reason for keeping it off was wrong. It cited libtorch.
Measured on Linux x86_64, enabling the CUDA delegate needs one companion option,
EXECUTORCH_BUILD_EXTENSION_TENSOR=ON, now set automatically, and the resulting
runner links no libtorch and no libc10. Its full shared-library set is
libaoti_cuda_shims, libcudart, libcurand, libextension_cuda, libnvinfer,
plus libc and libstdc++.

Docs corrected. They previously told readers green-context streams were outside
the validated support matrix, while upstream's own header names exactly that case as
the reason the API exists. The docs now say it works, show how to reproduce it, and
keep the two real limits explicit: it is not in CI, because the CI configuration
builds the runner without the CUDA delegate, and it took the synchronized path, so
device-resident asynchronous return and its interaction with the internal completion
event are still uncovered.

Also retitled, since removing CudaStreamGuard is source-breaking for C++ consumers
rather than a plain refactor.

Note this PR's CI still cannot run its own checks until #4523 lands: the
executorch-runtime-build job has failed on every commit since it was introduced, and
executorch-runtime-test is gated on it.

[shoumikhin]

Heads up that this grew since your approval, so please re-check rather than take the
existing one as covering it.

Getting ExecuTorch's shared caller-stream library requires EXECUTORCH_BUILD_CUDA=ON,
and upstream then links its CUDA backend into the same _portable_lib this wheel
ships. The wheel was carrying that backend with half its runtime missing: the
extensions have a DT_NEEDED on both libextension_cuda.so and
libaoti_cuda_shims.so, ExecuTorch's published wheel provides neither, and only the
first was being shipped. That is an ImportError waiting to happen.

Rather than back the CUDA backend out, this keeps it and ships the rest, because it
closes a real gap: ExecuTorch's published wheel registers only XnnpackBackend, so
until now no Python runtime anywhere could execute a .pte split across the TensorRT
and CUDA delegates, even though exporting one was already supported. The runtime test
now asserts CudaBackend is registered.

Also here: data_loader gets the same $ORIGIN search path _portable_lib already
had, so neither extension depends on module load order to find those libraries, and
the static C++ runtime check now covers every shared object the wheel ships instead of
only the two Python extensions.

The title and description are updated to say all of this. Happy to split the CUDA
delegate part into its own PR if you would rather review it separately.

One thing outside this PR: the ExecuTorch wheel build is currently red on main for an
unrelated reason, a libstdc++ link-order bug. #4534 fixes that, and this PR needs it
before its build can get far enough to prove any of the above.