Skip to content

refactor(executorch): use shared caller stream - #4421

Closed
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:executorch-shared-caller-stream
Closed

refactor(executorch): use shared caller stream#4421
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:executorch-shared-caller-stream

Conversation

@shoumikhin

@shoumikhin shoumikhin commented Jul 22, 2026

Copy link
Copy Markdown
Contributor

The problem

An ExecuTorch program can mix delegates. One subgraph may run on TensorRT while
another runs on the CUDA/AOTI backend. If the application wants both to run on a
CUDA stream it owns, both delegates have to agree on which stream that is.

Today they cannot agree. Torch-TensorRT keeps its own private "caller stream"
value, and ExecuTorch's CUDA backend keeps a different one. Selecting a stream for
one backend leaves the other running somewhere else, which breaks ordering.

The fix

Delete Torch-TensorRT's private caller-stream state and read ExecuTorch's shared
one instead.

Before:

torch_tensorrt::executorch_backend::CudaStreamGuard guard(stream);
module.forward(inputs);

After:

#include <executorch/extension/cuda/caller_stream.h>

executorch::extension::cuda::CallerStreamGuard guard(stream);
module.forward(inputs);

TensorRT and other CUDA-capable ExecuTorch delegates now read the same selection.

Why one shared library matters

The selection lives in a thread_local variable inside a single shared library,
libextension_cuda.so. A shared library is a .so file that a program loads at
run time, and every part of the program that loads the same one sees the same
variable.

If a second copy of that variable ends up in the process, for example because
something linked a static archive instead, then the runner writes to one copy and
the delegate reads the other. Nothing crashes. The delegate silently falls back to
a different stream:

runner selected stream 0x1056930
delegate observed      0x2          <- cudaStreamPerThread, the fallback

So the build must guarantee exactly one shared library. This PR enforces that in
three ways:

  • CMake accepts a prebuilt override only after reading the file's ELF header and
    confirming it is a shared object. Checking the file name is not enough, because
    the linker happily links a static archive that has been renamed to .so, which
    produces exactly the duplicate above.
  • When building from an ExecuTorch source checkout, CMake calls
    add_subdirectory on ExecuTorch's own extension/cuda instead of re-declaring
    the target. This keeps the two builds identical, and if something else also
    declares the target, CMake fails loudly rather than producing two libraries.
  • CI inspects the built and the packaged runner and requires a real
    DT_NEEDED entry for libextension_cuda.so plus dynamic imports of both
    getCallerStream and CallerStreamGuard. A private copy satisfies those
    references at link time and leaves no import, so its absence is the signal.

What else changes

  • Use the selected stream for the TensorRT enqueue and for host staging copies.
  • Fall back to cudaStreamPerThread when no guard is active.
  • Read the selection once, then derive both the stream and whether the backend may
    return with work still in flight. Two separate reads could drift apart.
  • Package exactly one shared libextension_cuda.so.
  • Run the ExecuTorch backend C++ tests in CI, which were previously only built.
  • Keep the native reference runner free of libtorch, and run real inference inside
    a CallerStreamGuard.

Compatibility

Removing torch_tensorrt::executorch_backend::CudaStreamGuard is an intentional
source-level C++ API change. Native callers switch to CallerStreamGuard as shown
above. The replacement is behavior-preserving: an explicitly selected null stream
still counts as a caller selection, matching the old two-field encoding.

This PR validates ordinary CUDA streams. On the discrete-GPU CI configuration the
reference runner exercises the synchronized host-staging path. The device-resident
asynchronous path is not covered end to end by that runner. CUDA green-context
streams need context-aware completion-event handling and are not claimed as
supported here.

Python export-only usage is unchanged.

Testing

Verified on Linux x86_64 with an NVIDIA H100, CUDA 12.8, and TensorRT 11.0,
against the pinned ExecuTorch source checkout:

  • Every CMake selection branch, 10 cases: source build, prebuilt shared library,
    empty override, missing override, and rejection of a static archive, a linker
    script, and an executable renamed to .so. A shared object with no .so
    extension is correctly accepted, since the header decides, not the name.
  • The linkage assertions, checked against a deliberately broken runner that
    embeds a private copy. The gate rejects it, including after stripping, and
    accepts a correctly linked one.
  • The caller-stream unit tests, 6 cases, including an explicit null stream, an
    explicit cudaStreamPerThread, nesting, and per-thread isolation.
  • One shared thread-local across two independently linked shared libraries plus
    the executable, and the reverse case confirming a duplicate copy is detectable.
  • ET_CHECK_MSG survives -O2 -DNDEBUG, so release-build checks are not
    compiled out.
  • The backend builds for both runtime flavors. C++ formatting and shell syntax
    are clean.

Dependency

Native caller-stream support needs ExecuTorch 1.4 APIs. The version bump is owned
by a separate change and is temporarily included in this branch's diff against
main; those pin files disappear after rebasing once it lands.

@meta-cla meta-cla Bot added the cla signed label Jul 22, 2026
@github-actions github-actions Bot added component: tests Issues re: Tests component: build system Issues re: Build system component: api [C++] Issues re: C++ API labels Jul 22, 2026
@shoumikhin
shoumikhin force-pushed the executorch-shared-caller-stream branch 2 times, most recently from 764b1ca to 5a2ddd9 Compare July 22, 2026 20:52
@shoumikhin

Copy link
Copy Markdown
Contributor Author

CI note: the failing checks here are unrelated to this change and appear to be pre-existing trunk flakiness.

  • L1 dynamo compile tests fail only in test_hf_gqa_model.py::test_dynamic_head_dim_with_hf_model with numerical assert_close mismatches (rtol=1e-1). That is a Python Dynamo HuggingFace-GQA accuracy test; this PR touches only the C++ ExecuTorch delegate (cpp/src/torch_tensorrt/executorch/*) plus Bazel/CMake/BUCK build files and a verify script, none of which are on that code path.
  • The RTX - Python-only dynamo runtime tests failures are likewise in the Python runtime path, not the ExecuTorch delegate.

The branch is already based on the latest main, so these are not stale-base failures. Flagging as unrelated; happy to rerun once the trunk flakiness clears.

@shoumikhin
shoumikhin force-pushed the executorch-shared-caller-stream branch from 801cdcf to c7e3a09 Compare July 26, 2026 05:28
@shoumikhin shoumikhin changed the title executorch: use ExecuTorch's shared caller stream refactor: use ExecuTorch's shared caller stream in the TensorRT ExecuTorch backend Jul 26, 2026
@shoumikhin
shoumikhin force-pushed the executorch-shared-caller-stream branch from c7e3a09 to a2dc104 Compare July 28, 2026 22:05
@shoumikhin shoumikhin changed the title refactor: use ExecuTorch's shared caller stream in the TensorRT ExecuTorch backend refactor(executorch): use shared caller stream Jul 28, 2026
@shoumikhin

Copy link
Copy Markdown
Contributor Author

CI update after the rebase/force-push:

  • This branch is intentionally stacked on Pin ExecuTorch to the release/1.4 branch head #4434 until that version-pin PR merges; the four pin files will disappear from refactor(executorch): use shared caller stream #4421 after the final rebase onto main.

  • The current JetPack wheel failure is unrelated to this change. It fails in packaging/pre_build_script.sh before building Torch-TensorRT because the JetPack package index cannot resolve Torch's filelock dependency:

    ERROR: Could not find a version that satisfies the requirement filelock (from torch)
    ModuleNotFoundError: No module named 'torch'
    

The caller-stream-specific ExecuTorch static workflow and the remaining build matrix are still running.

@shoumikhin
shoumikhin marked this pull request as ready for review July 28, 2026 22:56
@shoumikhin
shoumikhin force-pushed the executorch-shared-caller-stream branch from a2dc104 to 27b28a0 Compare July 29, 2026 13:58
@narendasan
narendasan requested a review from cehongwang July 29, 2026 17:20
@shoumikhin
shoumikhin force-pushed the executorch-shared-caller-stream branch 3 times, most recently from b18ebc2 to 9e0294f Compare July 29, 2026 19:50
Replace the TensorRT ExecuTorch backend private caller-stream TLS with ExecuTorch CallerStreamGuard/getCallerStream so CUDA-capable delegates share one process-wide selection.

Link and package one shared extension_cuda instance, add ordinary caller-stream inference coverage in the reference runner, and verify both CMake-built and packaged runners consume the shared TLS without libtorch.

The ExecuTorch release/1.4 pin is intentionally owned by the preceding version-bump commit.
@shoumikhin shoumikhin closed this Jul 31, 2026
@shoumikhin
shoumikhin force-pushed the executorch-shared-caller-stream branch from 9e0294f to 6e51eef Compare July 31, 2026 06:35
@github-actions github-actions Bot added documentation Improvements or additions to documentation component: lowering Issues re: The lowering / preprocessing passes component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: converters Issues re: Specific op converters labels Jul 31, 2026
@shoumikhin

Copy link
Copy Markdown
Contributor Author

Pushed a revision that tightens the single-shared-library guarantee and removes some
unnecessary code. Summary of what changed since the last push, and why.

The shared-object check was not actually checking. The prebuilt-override path
validated the file name with a regex. The linker does not care about file names, so a
static archive renamed to libextension_cuda.so was accepted, linked as an archive,
and produced a binary with no DT_NEEDED entry and a private copy of the
caller-stream variable. That is exactly the duplicate this PR exists to prevent. It
now reads the ELF header and requires a shared object, handling both endianness and
both ELF classes.

Two libraries could be created silently. If a project added this package before
ExecuTorch, the imported target and ExecuTorch's real target coexisted with no
warning, and both libraries got built. The source path now calls add_subdirectory
on ExecuTorch's own extension/cuda instead of re-declaring the target, so a
duplicate is a hard CMake error. That also removes a hand-copied build rule that
could drift from upstream.

An empty override bricked a working configuration. if(DEFINED ...) is true for
an empty string, so passing -DEXECUTORCH_EXTENSION_CUDA_LIBRARY= failed the build
even with a valid ExecuTorch checkout present. Changed to a truthiness check.

Removed the one-line public wrapper. getTensorRTExecutionStream() was a new
public symbol whose body was getCallerStream().value_or(cudaStreamPerThread). It
also caused the backend to read the thread-local twice for one decision. The backend
now reads once and derives both answers from that read, which cannot drift.

Corrected the CI symbol assertions. The check used nm --defined-only, which
reads .symtab. That section is stripped from release binaries, so the assertion
passed vacuously. I measured which checks actually catch a deliberately broken
runner: the missing DT_NEEDED entry and the absence of dynamic imports both catch
it, and both survive stripping. CI now asserts those, requiring imports of
getCallerStream (referenced by the delegate) and CallerStreamGuard (referenced by
the runner).

Smaller build and CI surface. The unit test no longer links TensorRT, so the
libnvinfer lookup, the LD_LIBRARY_PATH export, the sandbox-disabling
--test_strategy=standalone, and the GPU env passthrough are all gone. I also
reverted the CMake minimum bump for the backend, since nothing there needs a newer
version; it is now 3.19 to match the ExecuTorch extension it builds. The reference
runner keeps 3.24 because it adds the full ExecuTorch tree.

Documentation. Softened a claim that the runner "validates stream selection". It
exercises guarded inference and checks output values; the linkage assertions are what
validate the shared-library requirement. Also fixed a stale output path.

Net effect is 27 fewer lines than the previous revision. Full verification, on Linux
x86_64 with an H100, CUDA 12.8, and TensorRT 11.0 against the pinned ExecuTorch
checkout, is listed in the updated description. All suites pass.

Two things I did not change, and want to flag rather than hide:

  • The device-resident asynchronous path still has no end-to-end test. The CI runner's
    tensors are host-backed, so it always takes the synchronized path.
  • A runtime self-check inside the runner would not help. The backend is a static
    archive whole-archived into the runner executable, so the runner and the delegate
    are one linkage unit and always agree, even in a deliberately broken build. I
    verified that, which is why the linkage assertions are the meaningful gate here
    rather than a runtime comparison.

@shoumikhin

Copy link
Copy Markdown
Contributor Author

Closing note: this pull request was auto-closed by GitHub, not intentionally.

I force-pushed an update from a shallow clone (git clone --depth 1). In a shallow
clone git does not have the parent commits, so amending produced a commit with no
parent at all. Pushing that left this branch with no history in common with main,
and GitHub closed the pull request because such a branch can never merge.

The branch has been repaired: the commit now sits on the correct parent and the file
contents are unchanged from what was under review here. GitHub refuses to reopen a
pull request whose branch was recreated, so the work continues in a new one.

Continued in #4454.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [C++] Issues re: C++ API component: build system Issues re: Build system component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: core Issues re: The core compiler component: evaluators Issues re: Specific op evaluators component: lowering Issues re: The lowering / preprocessing passes component: partitioning component: runtime component: tests Issues re: Tests documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant