From 0cc18632c16ff69016be50ab267e7379b1f1ed55 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 Oct 2026 12:39:57 -0700 Subject: [PATCH 1/9] [2.15] bump the ExecuTorch delegate to 0.2.1 Since the #4821 backport, every .pte saved from this branch stores its engines as named data (blob tag TR03). The 0.2.0 delegate published with the first 2.15.0 release candidate predates that and refuses these files with only "failed to parse TensorRT blob". A delegate rebuilt from this branch at the same version would look identical to pip, so an environment that already has the old one would keep it. Bump the delegate to 0.2.1, so pip sees it as newer than 0.2.0 and replaces the old delegate wherever it installs one. The README's build recipe uses the new version. --- py/torch-tensorrt-executorch-runtime/README.md | 4 ++-- py/torch-tensorrt-executorch-runtime/version.txt | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/py/torch-tensorrt-executorch-runtime/README.md b/py/torch-tensorrt-executorch-runtime/README.md index 6d2b112cb46..c8a020973ea 100644 --- a/py/torch-tensorrt-executorch-runtime/README.md +++ b/py/torch-tensorrt-executorch-runtime/README.md @@ -230,7 +230,7 @@ render it for your machine before building. The template takes seven values: export CUDA_HOME=/usr/local/cuda-13.2 export CU_VERSION=cu132 export CHANNEL=nightly -export BUILD_VERSION=0.2.0.dev0 +export BUILD_VERSION=0.2.1.dev0 export TORCH_INSTALL_PATH="$(python -c 'import torch; print(torch.__path__[0])')" export TENSORRT_CUDA_VERSION_UPPER_BOUND=14.0 export TENSORRT_RTX_CUDA_VERSION_UPPER_BOUND=14.0 @@ -248,7 +248,7 @@ python -m pip install pyyaml patchelf tensorrt-cu13 \ --extra-index-url "https://download.pytorch.org/whl/test/${CU_VERSION}" \ --extra-index-url https://pypi.nvidia.com/ \ "executorch==1.6.0.dev20261009" -export TORCH_TENSORRT_EXECUTORCH_RUNTIME_VERSION="0.2.0.dev0+${CU_VERSION}" +export TORCH_TENSORRT_EXECUTORCH_RUNTIME_VERSION="0.2.1.dev0+${CU_VERSION}" python -m pip wheel --no-build-isolation --no-deps \ --wheel-dir dist py/torch-tensorrt-executorch-runtime ``` diff --git a/py/torch-tensorrt-executorch-runtime/version.txt b/py/torch-tensorrt-executorch-runtime/version.txt index 341cf11faf9..7dff5b89211 100644 --- a/py/torch-tensorrt-executorch-runtime/version.txt +++ b/py/torch-tensorrt-executorch-runtime/version.txt @@ -1 +1 @@ -0.2.0 \ No newline at end of file +0.2.1 \ No newline at end of file From 6b58ad696a7d687b0bed35a49d16cdc254446926 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 Oct 2026 12:41:03 -0700 Subject: [PATCH 2/9] [2.15] test that the ExecuTorch guide's release install is one command The guide on this branch installs a release with one command: pip install "torch_tensorrt[executorch]" Once this branch pins an ExecuTorch release, its PyPI wheel's executorch extra brings ExecuTorch's CUDA 13.2 build and the delegate built for CUDA 13.2, so naming either in the command would only let them drift apart. A new test checks that the guide's release install is exactly this command. The install command test also gains a case that refuses a command which names the nightly index without --pre. --- .../dynamo/executorch/test_executorch_pin.py | 47 ++++++++++++++++--- 1 file changed, 40 insertions(+), 7 deletions(-) diff --git a/tests/py/dynamo/executorch/test_executorch_pin.py b/tests/py/dynamo/executorch/test_executorch_pin.py index 62f67b96db0..97de4fe695a 100644 --- a/tests/py/dynamo/executorch/test_executorch_pin.py +++ b/tests/py/dynamo/executorch/test_executorch_pin.py @@ -936,15 +936,29 @@ def test_install_channel_guard_rejects_unsupported_nightly_recipes( @pytest.mark.unit @pytest.mark.parametrize( - ("path", "command", "allowed"), + ("path", "command", "error"), [ - ("docsrc/install.rst", 'pip install "torch-tensorrt[executorch]"', True), - ("README.md", 'pip install "torch-tensorrt[executorch]"', False), - ("docsrc/install.rst", 'pip install --pre "torch-tensorrt[executorch]"', False), + ("docsrc/install.rst", 'pip install "torch-tensorrt[executorch]"', None), + ( + "README.md", + 'pip install "torch-tensorrt[executorch]"', + "names no nightly channel", + ), + ( + "docsrc/install.rst", + 'pip install --pre "torch-tensorrt[executorch]"', + "names no nightly channel", + ), + ( + "docsrc/install.rst", + 'pip install "torch-tensorrt[executorch]" ' + "--index-url https://download.pytorch.org/whl/nightly/cu132", + "without --pre", + ), ], ) def test_only_user_docs_may_print_the_release_install( - tmp_path, monkeypatch, path, command, allowed + tmp_path, monkeypatch, path, command, error ): recipe = tmp_path / path recipe.parent.mkdir(parents=True, exist_ok=True) @@ -952,13 +966,32 @@ def test_only_user_docs_may_print_the_release_install( subprocess.run(["git", "init", "-q", str(tmp_path)], check=True) subprocess.run(["git", "-C", str(tmp_path), "add", path], check=True) monkeypatch.setattr(sys.modules[__name__], "REPO_ROOT", tmp_path) - if allowed: + if error is None: test_every_printed_install_instruction_names_the_nightly_channel() else: - with pytest.raises(AssertionError, match="names no nightly channel"): + with pytest.raises(AssertionError, match=error): test_every_printed_install_instruction_names_the_nightly_channel() +@pytest.mark.unit +def test_the_release_install_is_one_command() -> None: + """The guide installs this release from PyPI with the extra alone. + + The release's PyPI wheel is built for the CUDA version of PyTorch's PyPI wheel, and once the + pin is an ExecuTorch release its extra brings ExecuTorch's CUDA build and the delegate built + with it, so naming either would only let them drift apart. + """ + guide = REPO_ROOT / "docsrc/user_guide/runtime_performance/executorch.rst" + commands = [ + shlex.split(block.replace("\\\n", "")) + for block in re.findall( + r"^ *pip install (?:.*\\\n)*.*$", guide.read_text(encoding="utf-8"), re.M + ) + ] + releases = [argv for argv in commands if "--pre" not in argv] + assert releases == [["pip", "install", "torch_tensorrt[executorch]"]], commands + + @pytest.mark.unit @pytest.mark.parametrize("variable", [False, True]) def test_install_channel_check_detects_removed_suffix_fix( From 288b552c7f65583d333733d14bf0d4a0c8cae5cf Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 24 Oct 2026 15:17:08 -0700 Subject: [PATCH 3/9] [2.15] point the .pte saving section at the install guide The saving section still printed a nightly install command and said the integration supports Linux. On this branch that command takes main's nightly in a fresh environment, and with 2.15 installed it adds ExecuTorch but no delegate, because a release's extra leaves the delegate out. It also disagreed with the guide linked just above it. Replace the platform sentence and the command with a pointer to the guide's Installation section. --- .../runtime_performance/saving_models.rst | 19 +------------------ 1 file changed, 1 insertion(+), 18 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index c39faec4e97..d92ea744922 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -233,24 +233,7 @@ c) ExecuTorch (.pte) The ``executorch`` output format lowers the compiled module to an ExecuTorch ``.pte`` program, delegating the TensorRT engines to the Torch-TensorRT ExecuTorch -backend. This CUDA integration links CUDA 13 libraries, so it supports Linux with -any CUDA 13 PyTorch build and needs an ExecuTorch nightly wheel. Install from the -channel matching your CUDA, for example a CUDA 13.2 environment: - -.. code-block:: bash - - pip install --pre "torch_tensorrt[executorch]" \ - --extra-index-url https://download.pytorch.org/whl/nightly/cu132 - -Substitute the channel for your CUDA, such as ``cu134`` for CUDA 13.4. ExecuTorch -publishes a channel later than PyTorch does, so a very new CUDA minor may not have -ExecuTorch wheels yet. Keep PyTorch, ExecuTorch and Torch-TensorRT on the -same CUDA channel. ``--pre`` permits prereleases; it does not request an upgrade -of an installed package. An existing version that declares the extra can gain -its missing dependencies, and dependency resolution may change other packages. -If an older version does not declare the extra, deliberately install the intended -compatible Torch-TensorRT wheel first. A fresh virtual environment avoids changing -an existing working stack. +backend. To install it, follow the Installation section of :ref:`executorch_deployment`. There are two ways to produce a ``.pte``, and they suit different needs: From febbb04071e362250c0d7a6094aeada73ca317d4 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 27 Oct 2026 12:42:03 -0700 Subject: [PATCH 4/9] [2.15] fix contradictions in the ExecuTorch guide and READMEs The 2.15 guide and the pages it links disagreed with each other and with the code: - The saving section said each delegate of a coalesced program runs on its own CUDA stream, and that a zero-copy TensorRT delegate returns before its engine finishes. Both delegates use the calling thread's default stream, and the TensorRT delegate waits for its work. The section now says so, and names the one case that needs a sync: the CUDA backend running last. - Two pages described an overwritten coalesced weight file differently. Both now say it can give wrong answers or fail to load. - The comparison table said a .pt2 runs on Windows. AOTInductor packaging is Linux only. - The guide promised dynamic shapes in a section that has none, and its example list missed two examples. - A coalesced program's data file is passed when loading from Python, the JetPack build joins the builds without a delegate, and ATen, KV and SM are spelled out. - The saving page's format list no longer calls ExecuTorch output Linux-only. - The guide said a host tensor passed to a device-resident program only costs a copy, and that Python cannot run a program with unplanned outputs. ExecuTorch's Python runtime refuses the host tensor and allocates device memory for unplanned outputs, so the guide and the runtime README now say so. - Two save() error messages said the integration supports only Linux. It supports Windows x86-64 too. --- .../runtime_performance/executorch.rst | 83 ++++++++++--------- .../runtime_performance/saving_models.rst | 70 ++++++---------- .../README.md | 37 +++++---- py/torch_tensorrt/_compile.py | 5 +- 4 files changed, 93 insertions(+), 102 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/executorch.rst b/docsrc/user_guide/runtime_performance/executorch.rst index 6f00e262f9a..368ff151e6b 100644 --- a/docsrc/user_guide/runtime_performance/executorch.rst +++ b/docsrc/user_guide/runtime_performance/executorch.rst @@ -21,8 +21,8 @@ program, so the whole model stays on the GPU. .. note:: - The ExecuTorch delegate runs on Linux, x86_64 and aarch64, and on Windows x86_64. It - requires CUDA 13. Ordinary Torch-TensorRT builds keep their separate CUDA 12 support. + The prebuilt ExecuTorch delegate runs on Linux, x86_64 and aarch64, and on Windows x86_64. + It requires CUDA 13. Ordinary Torch-TensorRT builds keep their separate CUDA 12 support. ---- @@ -68,9 +68,9 @@ README shows. The delegate is built for the standard Linux wheels (x86-64 and aarch64), and for Windows x86-64. On aarch64 it needs glibc 2.35 or newer, the same as TensorRT itself. -The TensorRT-RTX and Python-only builds have no delegate. To export a coalesced program you also -need a CUDA toolkit, because ExecuTorch's CUDA backend compiles the leftover operators with -``nvcc``. +The JetPack, TensorRT-RTX and Python-only builds have no delegate. To export a coalesced program +you also need a CUDA toolkit, because ExecuTorch's CUDA backend compiles the leftover operators +with ``nvcc``. ---- @@ -106,16 +106,17 @@ The usual ``ir="dynamo"`` path, with two extra arguments to ``torch_tensorrt.sav retrace=False, ) -:ref:`executorch_save` covers the rest: dynamic shapes, several methods in one file, a -zero-copy KV cache, and the two-step ``torch_tensorrt.executorch.export()`` path for -programs that need work before they are written. +:ref:`executorch_save` covers the rest: several methods in one file, a zero-copy key-value +(KV) cache, and the two-step ``torch_tensorrt.executorch.export()`` path for programs that +need work before they are written. Coalescing the leftover operators onto the CUDA backend ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -TensorRT has no converter for every ATen operator, and by default the ones it rejects run on -the CPU, costing a copy each way on every call. Pass a ``CudaPartitioner`` and ExecuTorch's -CUDA backend compiles them instead, so nothing leaves the GPU: +TensorRT has no converter for every ATen operator (PyTorch's core operator set), and by +default the ones it rejects run on the CPU, costing a copy each way on every call. Pass a +``CudaPartitioner`` and ExecuTorch's CUDA backend compiles them instead, so nothing leaves +the GPU: .. code-block:: python @@ -140,9 +141,9 @@ where TensorRT cannot take ``erfinv``, the saved program's delegate list reads .. warning:: The CUDA backend names its external weight file per device, not per model, so saving two - coalesced programs into one directory overwrites the first one's weights. That program - still loads, still reports finding its weights, and returns a wrong answer with no error. - Give each export its own directory. + coalesced programs into one directory overwrites the first one's weights. The first program + can then return wrong answers with no error, or fail to load. Give each export its own + directory. .. _executorch_external_engine_data: @@ -223,20 +224,21 @@ copies are pure overhead: ), ) -Three things are easy to miss: +Three things to know: * **Both skip flags need** ``enable_non_cpu_memory_planning=True``, because the copies are inserted during device-aware memory planning. Asking for a skip with it off raises. * **Inputs must be unplanned**, through ``MemoryPlanningPass(alloc_graph_input=False)``. Otherwise the program reserves its own input buffer and the runtime fills it from the caller, which puts the copy straight back. -* **Leave outputs planned if Python will run the program**, so the program's device arena - owns the output. Add ``alloc_graph_output=False`` only for a C++ caller that supplies the - address with ``Module::set_output``. Python cannot. +* **The outputs can stay in the program's own GPU memory**, its device arena. To give each output + its own buffer instead, also pass ``alloc_graph_output=False``. Then each output needs a GPU + buffer before the call: in C++ you set it with ``Module::set_output``, and the Python runtime + allocates one for you. -The choice is baked into the file, so feed such a program CUDA tensors. A host tensor still -gives the right answer, but it copies on every call, which is the cost you exported to -avoid, and nothing warns you. +The choice is baked into the file, so feed such a program CUDA tensors. ExecuTorch's Python +runtime refuses a CPU tensor. From C++, a CPU tensor costs a copy on every call when the first +delegate to use that input is TensorRT, and fails when it is the CUDA backend. ---- @@ -261,14 +263,16 @@ registers the backend. Everything after it is ExecuTorch's own API: If the delegate cannot load, that import raises at once, instead of surfacing later as a program that will not load. -A coalesced program needs nothing extra: both backends are registered and the program records -which parts go where. Torch-TensorRT is only needed to export, not to run. +A coalesced program needs no extra import: both backends are registered and the program +records which parts go where. If its export wrote a data file such as ``aoti_cuda_blob.ptd``, +pass it with ``load_program(..., data_path=Path("aoti_cuda_blob.ptd"))``. Torch-TensorRT is only +needed to export, not to run. .. note:: ``torch_tensorrt.load(path, format="executorch")`` still works but is deprecated. It copies CUDA inputs to the CPU and supports embedded weights only. Use the Runtime API - above instead. A device-resident program requires it. + above instead. A device-resident program runs only through the Runtime API. ---- @@ -417,12 +421,12 @@ To put every delegate on a stream of your own, scope a guard over the whole exec :ref:`Running a coalesced .pte ` has the same rule for a decode loop. -That stream can be a green context, which holds a fixed number of streaming multiprocessors, +That stream can be a green context, which holds a fixed number of streaming multiprocessors (SMs), so the model stays inside that partition and the rest of the GPU is free. Create it with ``cuGreenCtxStreamCreate`` and scope the same guard over it. The limit rides the stream, so the green context does not have to be made current. -The reference runner has this built in. Pass the SM count: +The reference runner has this built in. Pass the number of SMs: .. code-block:: bash @@ -464,10 +468,11 @@ and drivers older than CUDA 12.5. The delegate logs the reason once. .. warning:: While a recording runs, no other thread may create or destroy a TensorRT engine or - execution context, or call a whole-device sync such as ``cudaDeviceSynchronize`` or - ``torch.cuda.synchronize()``. Every input shape change starts another recording, so - loading everything up front does not end the risk. Leave replay off when other threads - may do any of that. + execution context outside this delegate, or call a whole-device sync such as + ``cudaDeviceSynchronize`` or ``torch.cuda.synchronize()``. Every input shape change, and + every growth of the shared activation scratch, starts another recording, so loading + everything up front does not end the risk. Leave replay off when other threads may do any + of that. A C++ host can refuse a program's saved request for one load: @@ -525,10 +530,10 @@ and a device that has run a pooled engine must not be reset with ``cudaDeviceRes Shared engines ^^^^^^^^^^^^^^^ -Loading the same program twice in one process, say once per robot arm, used to hold its -weights in device memory twice. Handles built from the same engine bytes, on the same device -and with the same weight streaming request, now share one engine. Each still gets its own -execution context, buffers and lock. +Loading the same program twice in one process, say once per robot arm, keeps one copy of its +weights on the GPU. Handles built from the same engine bytes, on the same device and with the +same weight streaming request, share one engine. Each still gets its own execution context, +buffers and lock. This is on by default. On an 8 GB Jetson Orin Nano, loading one robot policy twice grew reported GPU memory by 806 MiB without sharing and by 122 MiB with it. The match is found by @@ -557,12 +562,12 @@ Comparison: .pte vs .pt2 * - libtorch at runtime - Not required - Required - * - Non-TRT ops + * - Operators TensorRT cannot run - ExecuTorch CUDA backend, or CPU fallback - Compiled by AOTInductor * - Platform - - Linux + Windows x86_64, CUDA 13 - - Linux + Windows + - Linux (x86_64, aarch64) and Windows x86_64, CUDA 13 + - Linux ---- @@ -570,6 +575,10 @@ Examples --------- * ``examples/torchtrt_executorch_example/export_static_shape.py`` writes a minimal ``.pte``. +* ``examples/torchtrt_executorch_example/export_dynamic_shape.py`` exports a model whose batch + size can change. +* ``examples/torchtrt_executorch_example/export_kv_cache_decode.py`` exports a decoder step + whose key-value cache stays in memory you own. * ``examples/torchtrt_executorch_example/export_coalesced.py`` splits one graph across the TensorRT delegate and ExecuTorch's CUDA backend. * ``examples/torchtrt_executorch_example/export_device_resident.py`` does the same with no diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index d92ea744922..800a30f58bd 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -22,7 +22,7 @@ specifying the `output_format` flag. Here are the options `output_format` will a * `exported_program` : This is the default. We perform transformations on the graphmodule first and use `torch.export.save` to save the module. * `torchscript` : We trace the graphmodule via `torch.jit.trace` and save it via `torch.jit.save`. * `PT2 Format` : This is a next generation runtime for PyTorch models, allowing them to run in Python and in C++ -* `executorch` : We lower the graphmodule to an ExecuTorch ``.pte`` program, delegating the TensorRT engines to the ExecuTorch backend. Linux-only; requires the ``executorch`` package. +* `executorch` : We lower the graphmodule to an ExecuTorch ``.pte`` program, delegating the TensorRT engines to the ExecuTorch backend. Runs on Linux (x86-64 and aarch64) and Windows x86-64; see :ref:`executorch_deployment` to install it. a) ExportedProgram ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -482,16 +482,10 @@ the third is refused where it is visible in the ``.pte``: surface as wrong results *or* as an illegal memory access. See :ref:`Running a coalesced .pte `. -* **Synchronizing that stream before reading a cache on the host**, for any - zero-copy ``.pte`` a runner drives on a caller stream, coalesced or not. A - delegate whose aliased outputs are threaded through it reflects each one into - its delegate output and so waits for the engine before returning; zero-copy - elides those outputs, so there is nothing to reflect and ``execute()`` returns - with the engine still running. A single-delegate decode loop owes the - synchronization as much as a coalesced program does. It is only the caller - stream that brings the duty: with none installed ``execute()`` synchronizes - before it returns, as it does whenever a delegate input or output is staged - through the host. +* **Synchronizing that stream before reading a cache on the host**, if the ``.pte`` is + coalesced and ExecuTorch's CUDA backend runs after the engine. The TensorRT delegate waits + for its engine before it returns, but the CUDA backend may return early. See + :ref:`Running a coalesced .pte `. * **Sharing one cache between methods.** Zero-copy is per method: it makes each method's engine write that method's buffer. Giving a prefill and a decode @@ -539,43 +533,29 @@ must be pointed at those data files to load them. .. warning:: - The CUDA backend names its external weight blob per-device (e.g. - ``aoti_cuda_blob.ptd``), not per-model, so saving two different coalesced - ``.pte`` into the same directory overwrites the blob and the first ``.pte`` - will fail to load. Save each coalesced model into its own directory. + The CUDA backend names its external weight file per device (for example + ``aoti_cuda_blob.ptd``), not per model, so saving two coalesced programs into + one directory overwrites the first one's weights. The first program can then + return wrong answers with no error, or fail to load. Save each coalesced model + into its own directory. .. _executorch_single_stream: -**Running a coalesced .pte: use a single CUDA stream** - -A coalesced ``.pte`` runs on more than one backend delegate (the TensorRT delegate -plus the CUDA backend). By default each backend enqueues its GPU work on its *own* -CUDA stream, and delegate execution is asynchronous -- a delegate returns after -*enqueuing* its work, not after it completes. Because separate CUDA streams are not -ordered relative to one another, at a delegate boundary the consuming delegate can -begin before the producing delegate's writes have finished, reading incomplete -data. This is a race: it is intermittent and can surface as wrong results or an -illegal memory access. - -The runtime does not impose a shared stream across delegates, so it is the -**runner's responsibility** to run all delegates on one CUDA stream. Create a -single stream and scope ``executorch::extension::cuda::CallerStreamGuard`` over it -for the duration of execution. That one guard reaches every CUDA-capable -delegate: they resolve a single shared ``libextension_cuda``, so the TensorRT -backend and the CUDA backend read the same caller-stream storage. All GPU work is -then enqueued in order and every cross-boundary dependency is satisfied, while -execution stays asynchronous. - -If the runner reads a delegate's outputs between calls (for example, an -autoregressive decode loop), synchronize the shared stream before reading: the -work may still be in flight when ``execute()`` returns, and a host-side copy on -the default stream will not wait for a non-blocking stream. A model that threads -its aliased outputs through the delegate is insulated from this in practice: -reflecting each aliased output into its delegate output makes the delegate wait -for the engine before it returns. Under -:ref:`zero-copy KV ` those outputs are elided, so there -is nothing to reflect and the delegate returns with the engine still running -- -the synchronization is then the only thing making a host read see the new values. +**Running a coalesced .pte on one CUDA stream** + +A coalesced ``.pte`` runs on two delegates: the TensorRT delegate and ExecuTorch's CUDA +backend. With no stream chosen, both use ``cudaStreamPerThread``, the calling thread's default +stream, so a program run from one thread is already in the right order. + +To use a stream of your own, put one ``executorch::extension::cuda::CallerStreamGuard`` around +the whole call. It reaches both delegates, because they share one ``libextension_cuda``. Do not +give the two delegates different streams: work on separate streams is not ordered, and the +result is a race that can show up as wrong values or an illegal memory access. + +The TensorRT delegate waits for its own work before it returns. ExecuTorch's CUDA backend can +return while its work is still queued, so when it runs last, ``execute()`` can return before the +outputs are ready. Synchronize the stream before you read an output on the host, for example +between the steps of a decode loop. **ExecuTorch lowering options** diff --git a/py/torch-tensorrt-executorch-runtime/README.md b/py/torch-tensorrt-executorch-runtime/README.md index c8a020973ea..d5235d3fce8 100644 --- a/py/torch-tensorrt-executorch-runtime/README.md +++ b/py/torch-tensorrt-executorch-runtime/README.md @@ -46,16 +46,15 @@ wheel's directory on its own when you link an imported library. Then ship the delegate yourself, or point a library path at it at run time. For device-resident exports with `alloc_graph_output=False`, C++ Module callers must provide a CUDA output tensor with `Module::set_output` before execution. -Python callers cannot, so that arrangement is C++ only. +ExecuTorch's Python runtime allocates one for each output itself. -If you want one program that runs from both, turn off allocation for the inputs -and leave it on for the outputs. The caller then supplies the input buffer, which -is the copy worth avoiding, and the program's own device arena owns the output. +The simplest arrangement turns off allocation for the inputs and leaves it on for the outputs. +The caller then supplies the input buffer, which is the copy worth avoiding, and the program's own +device arena owns the output. Measured on two architectures, that arrangement runs from Python and from C++ with -no boundary copies in the program and results identical to eager. Turning -allocation off for both works only from C++. Leaving it on for both is refused by the -export step, because a program that allocates its own inputs has nothing for a caller -to hand in. +no boundary copies in the program and results identical to eager. Leaving it on for both is +refused by the export step, because a program that allocates its own inputs has nothing for a +caller to hand in. Point CMake at both wheels. The example above calls `find_package(executorch)` as well, and that package lives in its own distribution. ExecuTorch is a @@ -164,8 +163,9 @@ If the delegate cannot be loaded, that import raises, rather than letting the failure surface later as a program that will not load. A program whose graph is split between TensorRT and ExecuTorch's own CUDA backend -needs nothing extra here: both backends are registered, and the program says which -parts go where. +needs no extra import: both backends are registered, and the program says which +parts go where. If its export wrote a data file such as `aoti_cuda_blob.ptd`, pass it +with `load_program(..., data_path=Path("aoti_cuda_blob.ptd"))`. ## Runtime libraries @@ -345,11 +345,13 @@ backend it cannot find. Loading and running a program is ExecuTorch's API, not this package's. Tensor placement, method lookup and output devices are all documented by ExecuTorch. A program exported with `skip_h2d_for_method_inputs` keeps its inputs on the -device, because nothing here copies them. +device, because nothing here moves them off the GPU. -Such a program does not check that promise, and it will not tell you when you -break it. Hand it a host tensor and it runs, returns the right answer, and -stages the copy for you, once per call. So a program you believe is copying +ExecuTorch's Python runtime refuses a host tensor for such an input. From C++, when TensorRT is +the first delegate to read the input, the program does not check that promise, and it will not +tell you when you break it. +Hand it a host tensor and it runs, returns the right answer, and stages the copy +for you, once per call. So a program you believe is copying nothing can be copying on every run with nothing to show it. The other direction is refused, so the leniency goes one way only. If the copies are why you exported this way, measure them rather than trusting the flag. @@ -360,10 +362,9 @@ extension: it is private and can change without notice. The two were measured against each other on every arrangement here, host and device, including a graph split with ExecuTorch's own CUDA backend, and they agree to the digit. -The choice that does matter is who owns the output buffer. Leave allocation on for -the graph outputs if Python will run the program, because a Python caller has no -way to hand one in. Turn it off only for a C++ consumer that supplies the address -itself. +The choice that does matter is who owns the output buffer. With allocation on for +the graph outputs, the program's device arena owns them. With it off, a C++ caller +supplies each address, and ExecuTorch's Python runtime allocates one for each output. ## Use diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index f7330d43f2e..3de2e759fb4 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -945,7 +945,7 @@ def save( if output_format == "executorch" and not _has_executorch_exir(): raise ImportError( "Saving with output_format='executorch' requires executorch.exir. " - "This CUDA integration supports Linux. Setup: " + "This CUDA integration supports Linux and Windows x86-64. Setup: " + executorch_install_command() ) if output_format == "executorch": @@ -1586,7 +1586,8 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None except ImportError: raise ImportError( "Could not import the ExecuTorch export integration for " - "output_format='executorch'. This CUDA integration supports Linux. " + "output_format='executorch'. This CUDA integration supports Linux and " + "Windows x86-64. " "Setup: " + executorch_install_command() ) # save() only ever lowers a single method, so keep rejecting the per-method From dd6a5b68903c2c6b660ec7368f97fae01ab42a84 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 26 Oct 2026 11:03:01 -0700 Subject: [PATCH 5/9] [2.15] update the ExecuTorch guide for #4833 and #4881 On this branch the guide was picked after the change that stops save() from retracing, so its examples still pass retrace=False and arg_inputs, and the text presents them as two extra arguments save needs. For the ExecuTorch format save() now keeps the compiled graph as it is and needs no example inputs. The examples now use one extra argument, output_format="executorch", as on main. The CUDA graph paragraph also said that engines on the shared activation scratch never replay. With #4833 they do, each with its own buffers for its inputs and outputs, and they record again when the pool moves to a larger buffer. The paragraph now says so. --- .../runtime_performance/executorch.rst | 33 +++++++------------ 1 file changed, 12 insertions(+), 21 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/executorch.rst b/docsrc/user_guide/runtime_performance/executorch.rst index 368ff151e6b..5245fbeed48 100644 --- a/docsrc/user_guide/runtime_performance/executorch.rst +++ b/docsrc/user_guide/runtime_performance/executorch.rst @@ -77,11 +77,9 @@ with ``nvcc``. Compile and Save ----------------- -The usual ``ir="dynamo"`` path, with two extra arguments to ``torch_tensorrt.save``: - -* ``output_format="executorch"`` selects the ``.pte`` writer. -* ``retrace=False`` keeps the compiled graph as it is. Re-exporting would drop the TensorRT - engines before ExecuTorch's partitioner sees them. +The usual ``ir="dynamo"`` path, with one extra argument to ``torch_tensorrt.save``: +``output_format="executorch"`` selects the ``.pte`` writer. For this format ``save`` keeps +the compiled graph as it is rather than re-exporting it, so it needs no example inputs. .. code-block:: python @@ -98,13 +96,7 @@ The usual ``ir="dynamo"`` path, with two extra arguments to ``torch_tensorrt.sav min_block_size=1, ) - torch_tensorrt.save( - trt_gm, - "model.pte", - output_format="executorch", - arg_inputs=example_inputs, - retrace=False, - ) + torch_tensorrt.save(trt_gm, "model.pte", output_format="executorch") :ref:`executorch_save` covers the rest: several methods in one file, a zero-copy key-value (KV) cache, and the two-step ``torch_tensorrt.executorch.export()`` path for programs that @@ -127,8 +119,6 @@ the GPU: trt_gm, "coalesced.pte", output_format="executorch", - arg_inputs=example_inputs, - retrace=False, partitioners=[ CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")]) ], @@ -212,8 +202,6 @@ copies are pure overhead: trt_gm, "device_resident.pte", output_format="executorch", - arg_inputs=example_inputs, - retrace=False, backend_config=ExecutorchBackendConfig( propagate_device_config=PropagateDeviceConfig( skip_h2d_for_method_inputs=True, @@ -448,8 +436,6 @@ launch many short kernels, where CPU launch work dominates each call. trt_gm, "model.pte", output_format="executorch", - arg_inputs=example_inputs, - retrace=False, use_cuda_graphs=True, ) @@ -461,9 +447,14 @@ It is not free. Each engine keeps a device buffer per input and output, and ever call copies each one in and out. With few kernels, or large inputs and outputs, those copies can cost more than the launches they save. -Some engines never replay and still pay the copies: changing shapes, aliased outputs, the -shared activation scratch below, a green-context stream, GPUs without stream-ordered memory, -and drivers older than CUDA 12.5. The delegate logs the reason once. +Some engines never replay: changing shapes, aliased outputs, a green-context stream, GPUs +without stream-ordered memory, and drivers older than CUDA 12.5. Only changing shapes still pay +the copies; the others take the ordinary path without the extra buffers. The delegate logs the +reason once for aliased outputs, missing stream-ordered memory and an old driver. + +Engines on the shared activation scratch below do replay, each with its own buffers for its +inputs and outputs. When the pool moves to a larger buffer, an engine records again, as it does +after a shape change. .. warning:: From d79dd00b916140582f483754c3b5735bddd61be4 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sat, 24 Oct 2026 08:14:45 -0700 Subject: [PATCH 6/9] [2.15] give the source-built extension_cuda ExecuTorch's c10 headers and core The delegate can build ExecuTorch's shared extension_cuda library from a checkout when nothing else provides it. Built standalone, as the delegate README's archive-only recipe does, that build failed twice. First it stopped compiling on c10/util/safe_numerics.h, because the include list it was given held only the checkout's parent folder, not ExecuTorch's own copy of c10. With that fixed, it stopped linking with "cannot find -lexecutorch_core": extension_cuda links executorch_core, and only an ExecuTorch build in the same project defines that target. The source-build branch now adds the c10 folder to the include list, as ExecuTorch's own build does, and, when no executorch_core target exists, defines one from the prebuilt core library. It is the same libexecutorch_core.a the delegate already links further down, and EXECUTORCH_CORE_LIBRARY, given on the command line or by a parent project, still overrides it. A project that adds ExecuTorch with add_subdirectory already has the target, so nothing changes there. Tested on this branch, Linux aarch64 with CUDA 13.2, following the archive-only recipe with no prebuilt extension_cuda: the build fails on the missing c10 header before this change, fails to find executorch_core with only the header half of it, and builds after it. A new unit test checks the branch keeps both pieces, and fails without this change. Backported from #4927 on main, where it landed with the ExecuTorch docs rewrite. --- .../torch_tensorrt/executorch/CMakeLists.txt | 19 ++++++++++++++++- .../dynamo/executorch/test_cmake_runtime.py | 21 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt index 53c6ac2e2fa..b01058cb1d9 100644 --- a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt +++ b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt @@ -161,7 +161,24 @@ elseif(_executorch_backend_tensorrt_source_root AND # than a copy of it, so this stays correct if ExecuTorch changes the target # and so a duplicate definition is a hard error instead of two libraries. if(NOT _common_include_directories) - set(_common_include_directories "${_executorch_include_parent}") + set(_common_include_directories + "${_executorch_include_parent}" + "${_executorch_backend_tensorrt_source_root}/runtime/core/portable_type/c10" + ) + endif() + # extension_cuda links executorch_core. Built standalone, no ExecuTorch build in this project + # defines it, so stand in the prebuilt core that the delegate links below. + if(NOT TARGET executorch_core) + if(NOT DEFINED EXECUTORCH_CORE_LIBRARY) + set(EXECUTORCH_CORE_LIBRARY + "${_executorch_backend_tensorrt_source_root}/cmake-out/libexecutorch_core.a" + CACHE FILEPATH "Path to the ExecuTorch core static library") + endif() + if(NOT EXISTS "${EXECUTORCH_CORE_LIBRARY}") + message(FATAL_ERROR "EXECUTORCH_CORE_LIBRARY='${EXECUTORCH_CORE_LIBRARY}' does not exist") + endif() + add_library(executorch_core STATIC IMPORTED) + set_target_properties(executorch_core PROPERTIES IMPORTED_LOCATION "${EXECUTORCH_CORE_LIBRARY}") endif() add_subdirectory( "${_executorch_backend_tensorrt_source_root}/extension/cuda" diff --git a/tests/py/dynamo/executorch/test_cmake_runtime.py b/tests/py/dynamo/executorch/test_cmake_runtime.py index cf6a5f07353..84ca84aacf0 100644 --- a/tests/py/dynamo/executorch/test_cmake_runtime.py +++ b/tests/py/dynamo/executorch/test_cmake_runtime.py @@ -438,6 +438,27 @@ def test_the_cuda_extension_alias_yields_to_a_consumers_own_target() -> None: ] +def test_the_source_built_cuda_extension_gets_executorch_headers_and_core() -> None: + """Built standalone from a checkout, extension_cuda failed to compile and then to link. + + Its sources include ExecuTorch's c10 headers, and it links executorch_core, which only an + ExecuTorch build in the same project defines. The standalone archive recipe has neither. + """ + source = ( + Path(__file__).resolve().parents[4] + / "cpp/src/torch_tensorrt/executorch/CMakeLists.txt" + ).read_text(encoding="utf-8") + branch = source.split("elseif(_executorch_backend_tensorrt_source_root AND")[1] + branch = branch.split( + '"${_executorch_backend_tensorrt_source_root}/extension/cuda"' + )[0] + assert "/runtime/core/portable_type/c10" in branch, branch + assert "if(NOT TARGET executorch_core)" in branch, branch + # A path the parent project set must survive the cache default. + assert "if(NOT DEFINED EXECUTORCH_CORE_LIBRARY)" in branch, branch + assert "add_library(executorch_core STATIC IMPORTED)" in branch, branch + + @pytest.mark.parametrize("embed", [True, False]) def test_the_run_path_opt_out_changes_what_the_consumer_links( tmp_path, linker_tools, embed From c50e66317a5933a7a2d0e0457f3d364225b4d7f7 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 10:43:34 -0700 Subject: [PATCH 7/9] [2.15] rewrite the ExecuTorch deployment docs, as #4927 did on main The ExecuTorch guide, the .pte saving section, the three ExecuTorch READMEs, the example scripts and a few docstrings had drifted from the code and from each other. #4927 rewrote them on main. The delegate and export code on this branch are the same as on main, so the same text applies here, with these changes for this branch: - The runtime wheel README builds version 0.2.1 of the delegate against this branch's ExecuTorch pin, 1.6.0.dev20261009, instead of main's 0.3.0 and 1.6.0.dev20261021. - The reference runner README checks out the ExecuTorch commit that this branch pins in MODULE.bazel. - The test that the guide's release install is one command stays. The guide's Python examples for load and process options and for clone_outputs use ExecuTorch 1.6.0's Python runtime. This branch's current development pin predates them; the release moves the pin to 1.6.0, which has them. What the rewrite changes: - One Installation section, linked from the other pages, with one command per channel. Releases: pip install "torch_tensorrt[executorch]". Nightly: pip install --pre "torch_tensorrt[executorch]" --index-url https://download.pytorch.org/whl/nightly/cu132. - A complete C++ project, with the commands that build and run it. - Sections on runtime options, weight streaming, engines in a .ptd file and optimization profiles. - The device-resident recipe leaves GPU outputs unplanned, and says what ExecuTorch's Python runtime does with host inputs. - save(), export() and TensorRTPartitioner docstrings now match the code: a target_device spec only labels a delegate's inputs and outputs, a process option set to False also overrides use_cuda_graphs, and generate_etrecord=True fails for a program with TensorRT engines. - The example scripts point at the guide instead of printing their own install commands. --- cpp/src/torch_tensorrt/executorch/README.md | 1286 +++++------------ docsrc/py_api/executorch.rst | 3 +- .../runtime_performance/executorch.rst | 684 +++++---- .../runtime_performance/saving_models.rst | 422 ++---- .../executorch_reference_runner/README.md | 419 ++---- .../load_model_device_resident.py | 15 +- .../export_coalesced.py | 7 +- .../export_device_resident.py | 53 +- .../export_dynamic_shape.py | 9 +- .../export_kv_cache_decode.py | 7 +- .../export_multi_profile.py | 16 +- .../export_static_shape.py | 9 +- .../README.md | 501 +++---- py/torch_tensorrt/_compile.py | 19 +- py/torch_tensorrt/executorch/_export.py | 10 +- py/torch_tensorrt/executorch/partitioner.py | 9 +- 16 files changed, 1322 insertions(+), 2147 deletions(-) diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index 0e944237300..f607cc7f1e8 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -1,9 +1,31 @@ # Torch-TensorRT ExecuTorch Backend -This package is included in `libtorchtrt.tar.gz` as -`torch_tensorrt/src/torch_tensorrt/executorch/`. It builds the TensorRT -backend delegate for ExecuTorch from source. -Building this backend requires CUDA Toolkit 12.5 or newer. +This directory is the source of the TensorRT delegate for ExecuTorch: the plugin that runs the +TensorRT engines stored in a `.pte` file. It ships in `libtorchtrt.tar.gz` as +`torch_tensorrt/src/torch_tensorrt/executorch/`, except in the JetPack and TensorRT-RTX builds. + +Most applications do not need to build it. The `torch-tensorrt-executorch-runtime` wheel ships it +prebuilt, with a CMake package. The +[ExecuTorch deployment guide](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html) +covers installing, exporting and running a model. This README covers building the delegate from +source, and what a C++ application has to know when it runs one. + +## Build it with your runner + +You need: + +- A CUDA Toolkit, 12.5 or newer. The prebuilt pieces need CUDA 13: the wheels, and the + `libextension_cuda.so` in `libtorchtrt.tar.gz` that the standalone archive below uses. +- A TensorRT SDK. Point `TensorRT_ROOT` at it. +- An ExecuTorch source checkout with its submodules (`git submodule update --init --recursive`), + at the commit your Torch-TensorRT release pins in its `MODULE.bazel`. `libtorchtrt.tar.gz` does + not carry that file, so look the commit up in the release's source on GitHub. +- PyTorch and PyYAML, installed in the Python environment that CMake finds. ExecuTorch's kernel + code generator uses both, and its CUDA backend looks for PyTorch when it is configured. + +Put ExecuTorch and Torch-TensorRT side by side. The paths below are the ones in +`libtorchtrt.tar.gz`. In a source checkout of Torch-TensorRT, the delegate is in +`cpp/src/torch_tensorrt/executorch` instead. ```text user_runner_project/ @@ -11,21 +33,19 @@ user_runner_project/ torch_tensorrt/ ``` -This backend requires ExecuTorch 1.4 or a source commit containing -`pytorch/executorch#20158` and `pytorch/executorch#20498`. The normal integration -path for a runner that already enables ExecuTorch's CUDA backend is to add both -ExecuTorch (with `EXECUTORCH_BUILD_CUDA=ON`) and this package, so the TensorRT -backend links ExecuTorch's shared `extension_cuda` target directly. A libtorch-free -runner should leave the full CUDA/AOTI backend disabled; this package then builds -only the minimal shared `extension_cuda` caller-stream library from the ExecuTorch -source checkout. Consumers on ELF platforms may instead set -`EXECUTORCH_EXTENSION_CUDA_LIBRARY` to a prebuilt shared library; the value is -checked to be a shared object, because a static copy would give each delegate its -own caller-stream state. Linking `executorch::backend_tensorrt` -makes the backend archive a dependency of your runner target, so you do not need a -separate backend build step. +Then add both to your runner's CMake build and link the delegate's target: ```cmake +set(EXECUTORCH_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/executorch") +set(BUILD_TESTING OFF CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_PYBIND OFF CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_CUDA ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_EXTENSION_MODULE ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP ON CACHE BOOL "" FORCE) + add_subdirectory("executorch") add_subdirectory("torch_tensorrt/src/torch_tensorrt/executorch") @@ -38,15 +58,126 @@ target_link_libraries(my_runner PRIVATE ) ``` -The backend archive is available as the `executorch_trt_backend` CMake target -and is written to `${CMAKE_BINARY_DIR}/lib/libexecutorch_trt_backend.a`. -`libextension_cuda` remains a shared runtime dependency so every CUDA-capable -delegate in the process observes the same caller-stream TLS instance. +- **Turn on ExecuTorch's CUDA backend, even if your programs use only TensorRT.** It builds the + shared `extension_cuda` library that the delegate links, and it registers the GPU memory + allocator that the program's GPU buffers need. Configuring it needs PyTorch, but the runner it + builds does not link libtorch. +- **Keep `extension_cuda` a shared library.** Every CUDA delegate in the process must read the same + caller stream, so the build refuses a static copy. +- `EXECUTORCH_SOURCE_DIR` (or `EXECUTORCH_ROOT`) tells the delegate where the ExecuTorch checkout + is. It reads the headers from there, and the build in the next paragraph needs it too. +- The extension options are what the CUDA backend and the `Module` class need. + +If you cannot turn on the CUDA backend, add ExecuTorch without it. The delegate then builds the +shared `extension_cuda` library itself from the checkout that `EXECUTORCH_SOURCE_DIR` names, or uses +a prebuilt one that you name with `EXECUTORCH_EXTENSION_CUDA_LIBRARY` (Linux only), such as +`torch_tensorrt/lib/libextension_cuda.so` from `libtorchtrt.tar.gz`. Nothing then registers the GPU +memory allocator, so your runner has to register ExecuTorch's `CudaAllocator` itself before it loads +a program, as `cpp/src/torch_tensorrt/executorch/RegisterCudaDeviceAllocator.cpp` in the +Torch-TensorRT repository does. Without that, loading fails with `No device allocator registered`. + +Linking `executorch::backend_tensorrt` builds the delegate's static archive, +`${CMAKE_BINARY_DIR}/lib/libexecutorch_trt_backend.a` (target `executorch_trt_backend`), as a +dependency of your runner. There is no separate step. + +`examples/executorch_reference_runner/` in the Torch-TensorRT repository is a complete runner +built this way. + +### Standalone archive + +Use this only when you need `libexecutorch_trt_backend.a` without a runner that adds ExecuTorch +with `add_subdirectory`. First build the ExecuTorch core runtime and its portable kernels: + +```bash +export EXECUTORCH_ROOT="${PWD}/executorch" +export TensorRT_ROOT=/path/to/extracted/TensorRT +export LD_LIBRARY_PATH="${TensorRT_ROOT}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + +cmake -S "${EXECUTORCH_ROOT}" -B "${EXECUTORCH_ROOT}/cmake-out" \ + -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_TESTING=OFF \ + -DEXECUTORCH_BUILD_PYBIND=OFF \ + -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ + -DEXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP=ON + +cmake --build "${EXECUTORCH_ROOT}/cmake-out" --target executorch_core portable_ops_lib -j +``` + +Then build the archive, from the folder that holds `executorch/` and `torch_tensorrt/`: + +```bash +cmake -S torch_tensorrt/src/torch_tensorrt/executorch -B build-torchtrt-executorch \ + -DEXECUTORCH_ROOT="${EXECUTORCH_ROOT}" \ + -DEXECUTORCH_EXTENSION_CUDA_LIBRARY="${PWD}/torch_tensorrt/lib/libextension_cuda.so" \ + -DTensorRT_ROOT="${TensorRT_ROOT}" + +cmake --build build-torchtrt-executorch --target executorch_trt_backend -j +``` -## Caller Stream API Migration +- `EXECUTORCH_EXTENSION_CUDA_LIBRARY` names the prebuilt `extension_cuda` library that + `libtorchtrt.tar.gz` ships. It links CUDA 13, so use a CUDA 13 toolkit with it. Leave it out to + build that library from the checkout instead. It then links the `libexecutorch_core.a` built + above. Set `EXECUTORCH_CORE_LIBRARY` if that file is somewhere else. +- Link the archive into your program with `-Wl,--whole-archive`, together with + `libextension_cuda.so`, ExecuTorch's portable kernels (`libportable_ops_lib.a`, also with + `-Wl,--whole-archive`, and `libportable_kernels.a`), the ExecuTorch core (`libexecutorch_core.a`), + TensorRT's `nvinfer`, the CUDA runtime and pthreads. The delegate and the kernels register + themselves from static initializers, and a plain static link drops them. +- The portable kernels include the two copy operators that a default export puts at the edges of + the program, `et_copy::_h2d_copy.out` and `et_copy::_d2h_copy.out`. Without them, loading fails + with `Error::OperatorMissing`, and a build with ExecuTorch logging on also logs + `kernel 'et_copy::_h2d_copy.out' not found`. +- Nothing in this setup registers the GPU memory allocator, so your program has to register + ExecuTorch's `CudaAllocator` itself, as described above. + +## Running programs from C++ + +The delegate registers itself as `TensorRTBackend` when its library loads. After that you use the +ordinary ExecuTorch C++ API, for example `executorch::extension::Module`. The rules below are the +ones a caller has to keep. + +### Threads + +- **Use one `Module` per thread.** Many threads, each with its own `Module`, work, including + threads that share one input buffer or one caller stream. ExecuTorch's `Module` is not thread + safe, so never call one `Module` from two threads at once, not even two different methods of it. + Mostly ExecuTorch then fails with `Error::InvalidState`, but sometimes a thread gets the other + thread's answer. + +### Buffers + +Every buffer you pass in must stay alive and unchanged until `execute()` returns. If another +thread frees an input during the call, the call can return a wrong answer with `Error::Ok`. Nothing +below your code can detect that. Python callers cannot hit this, because a tensor owns its memory. -`torch_tensorrt::executorch_backend::CudaStreamGuard` has been removed. Use -ExecuTorch's backend-neutral guard instead: +TensorRT needs every input and output address to be a multiple of 256 bytes, and it does not check. +A view that starts partway into its storage, such as `x[1:]` of a 1-D tensor, misses that. So the +delegate copies such a buffer through an aligned one of its own, and logs this once per engine. The +answer is right; the cost is one copy per buffer per call, and two for a buffer the engine updates +in place. `torch_tensorrt.save()` plans memory on 256-byte boundaries. ExecuTorch's own default is +16 bytes, so a program finalized with `to_executorch()` directly should pass +`MemoryPlanningPass(alignment=256)`, or its planned buffers are copied too. + +### CUDA streams + +- With no stream chosen, the delegate runs on `cudaStreamPerThread`, the default stream of the + calling thread. ExecuTorch's CUDA backend does the same, so a program split between the two runs + in the right order on one thread. +- To choose a stream, scope `executorch::extension::cuda::CallerStreamGuard` from + `` around the call. The guard applies to the calling + thread, and there it reaches every CUDA delegate. The stream can come from a CUDA green context, + which keeps the program inside a fixed set of the GPU's streaming multiprocessors (SMs). TensorRT + then warns that the engine was built without green context constraints, which it says may cause + crashes or undefined behavior. Torch-TensorRT has no option to build engines with them. +- The stream must be on the same device as the engine. +- `execute()` always waits for the engine to finish before it returns, on any stream. +- Calls on one delegate handle are serialized by the delegate. Destroying a handle while a call on + it runs is not allowed. + +`torch_tensorrt::executorch_backend::CudaStreamGuard` was removed in Torch-TensorRT 2.14. Use +`CallerStreamGuard` instead; it is the same idea, shared by every CUDA delegate: ```cpp #include @@ -55,405 +186,248 @@ executorch::extension::cuda::CallerStreamGuard guard(stream); module.forward(inputs); ``` -The old class is intentionally not kept as a deprecated alias: the goal is one -backend-neutral primitive and one shared TLS definition, so all CUDA-capable -delegates read the same caller-stream selection. (A deprecated `using` alias to -`executorch::extension::cuda::CallerStreamGuard` would have shared that same TLS, -so this removal is an API-simplification choice, not a correctness requirement.) -This is a source-breaking C++ change; downstream callers must switch to the new -type. - -### Coalesced programs with weights need one directory each - -A coalesced program's CUDA partition keeps its weights in a separate file written beside the -program, -under a fixed name. Exporting a second program into the same directory overwrites the first one's -weights, and the keys inside that file describe the graph's shape rather than its values, so the -first -program still loads, still reports finding weights, and returns a wrong answer with no error. - -Measured on two GPUs: wrong by 0.85, and bit identical across five runs, which is what makes it -read as -a working model rather than a broken one. Retraining and re-exporting the same architecture into one -directory is the ordinary way to hit this. - -Give each export its own directory. - -### A device-resident output reaches Python tagged for the device but backed by host memory - -Measured on two GPUs. A program exported device resident runs, and its values are right, but the -tensor -the Python bindings hand back reports the device while its pointer is ordinary host memory. Asked -through the driver, a genuine device tensor answers device memory on device zero, and this one -answers -unregistered host memory on device minus two. - -So from Python, `is_cuda` on that output proves nothing about residency, and neither does -anything built -on it. The same programs run correctly through a C++ consumer, where the caller supplies the output -address, so the runtime and this backend are doing their part. The structural reason is that the -type -carrying output metadata has no device field, so a generic runner cannot learn where a non-planned -output belongs, and a fix belongs upstream. - -Two consequences worth knowing. The default output clone in the bindings copies from that -pointer, which -is why an asynchronous run can return zeros. And ExecuTorch's own CUDA backend refuses such a buffer -while this backend accepts it, so a coalesced program's behaviour depends on which backend owns -the last -partition. Measured with the same script, the same machine and the same three operators: with -TensorRT -last it passes five times out of five, and with the CUDA backend last it fails five times out of -five. - -There is a workaround, and it is one export flag. Asking for the graph output to be planned, so the -program's own CUDA arena owns it instead of the bindings, turns both of those into five passes -out of -five. That is also the evidence that the two backends do not really disagree: hand either of -them real -device memory and both accept it. Only one of them checks. - -Which is also why this backend's leniency is not something to rely on. It never reads the device -tag at -all, and the devices differ: of the four tested, one reads pageable host memory through shared -page tables, two read it by faulting pages in, and one cannot read it at all, so the same program -that works on a discrete card can fail on an integrated one. - -### A split program runs only on the first GPU - -A program whose graph is split between this delegate and ExecuTorch's CUDA -backend appears to fail on any GPU but the first. On a four-GPU machine, with the -engine moved to the second device and every buffer also there, it failed ten runs -of ten with an illegal memory access, from Python and from C++ alike, while the -same program on the first device worked and this delegate alone on the second -device worked. The allocator was seen freeing on the first device while the engine -ran on the second, so only one half of the program moved. - -Read that as a warning rather than a settled result. The engine was moved by -rewriting the device recorded in the program, not by exporting for that device, so -it has not been reproduced the way a user would reach it. If you have more than one -GPU and a split program, prefer the first one until this is confirmed properly. - -### Loading the same program many times grows host memory - -A program whose graph is split across this delegate and ExecuTorch's CUDA backend grows the -process by roughly 390 to 440 kilobytes for every load, run and drop cycle. Measured on three -machines and two architectures as a slope fitted over hundreds of cycles, and the clearest -evidence is a count rather than a slope: one extra mapped shared object and four extra mapped -regions per cycle, with no variance at all, on the split program and none on a program carrying -only this delegate. - -The cause is on the CUDA backend's side: it opens a compiled library per load and never unmaps -it. A program carrying only this delegate does not, and the same measurement puts that at a -couple of hundred bytes per cycle, which is noise. Device memory does not grow at all, with a -slope of zero over hundreds of cycles on a probe proven to be reporting rather than dead. - -So a long-lived process that loads programs repeatedly will grow. Load once and keep the -program if you can. - -### A corrupted engine can return a wrong answer instead of an error - -Nothing checks the serialized engine against a digest, because the blob format carries none. The -header's -own fields are validated, and a corruption that breaks them is refused, but a corruption inside -the engine -bytes themselves is not seen. - -Measured with 24 probes that change one value inside the engine: 12 were refused, 12 were -accepted, and 8 -of those returned a wrong answer, one of them negative infinity. Each wrong value repeated -identically -across five runs, so this is deterministic rather than flaky. Corruption inside the delegate's own -metadata was ignored in all four probes. - -That is a property of the format rather than of this backend, and closing it means adding a -digest to the -format and to everything that writes one, which would not be readable by programs exported -before it. Until -then, treat a program file as trusted input: check it in transit, and do not run one from a -source you would -not run code from. - -### A buffer handed in has to outlive the call - -Measured, because it is the one misuse nothing reports. Freeing an input from another thread while -`execute` is running returned a plausible wrong answer in 253 of 300 runs: the right rank, the right -type, every value finite, and a status of success. On a linear model the recycled bytes produce a -scaled version of the right answer, which is exactly the shape a caller would not question. - -Nothing below this layer can catch it. The delegate is handed an address, and by the time the device -reads it the memory belongs to somebody else. So the rule is the caller's to keep: every buffer stays -alive and unchanged until the call returns. The Python path cannot reach this, because a tensor there -owns its storage. - -### A buffer off the 256-byte boundary is copied - -TensorRT needs every input and output address to be a multiple of 256 bytes, and it does not check. -A view that starts partway into its storage, such as `x[1:]` of a 1-D tensor, misses that. Bound as -it is, the engine returns a wrong answer, or its kernel faults and every later CUDA call in the -process fails. So the backend copies such a buffer through an aligned one of its own, and logs this -once per engine. The answer is right. The cost is one copy per buffer per call, and two for a -buffer the engine updates in place, which is copied in and then back. - -`torch_tensorrt.save()` plans memory on 256-byte boundaries, so buffers a program plans in GPU memory -are bound directly. ExecuTorch's own default is 16 bytes, so a program finalized with `to_executorch()` -directly should pass `MemoryPlanningPass(alignment=256)`, or its planned buffers take this path too. -On a device that reads host memory directly, an unaligned host buffer is copied the same way instead -of being read in place. - -### Running from several threads at once - -One program per thread works. Sharing one program across threads does not, and is refused -rather than allowed, because the runtime this plugs into does not permit it. - -Measured on a two-card A100 host, 100 inference calls per thread per attempt, with the threads -released into `execute()` by a barrier so they really do overlap (confirmed by an in-flight -counter reaching the thread count, and by the calls' own start and finish times overlapping for -about half of the busy time): - -| shape | attempts with a failure | -| --- | --- | -| one `Module` per thread, 2 threads | 0 of 20 | -| one `Module` per thread, 4 threads | 0 of 20 | -| one `Module` per thread, staggered starts instead of a barrier | 0 of 20 | -| one `Module` per thread, both threads reading one input buffer | 0 of 20 | -| one `Module` per thread, both inside a `CallerStreamGuard` on one shared stream | 0 of 20 | -| one `Module` per thread, loaded concurrently with nothing warmed up first | 0 of 10 | -| one `Module` shared by 2 threads | 20 of 20 | -| one `Module` shared by 4 threads | 19 of 20 | - -That is 26,000 answers compared against the expected value with no mismatch on the per-thread -shape. The backend's own `execute()` mutex is what makes it hold: two threads on one delegate -handle are serialized, and two threads on two handles run at the same time without touching each -other's state. - -Sharing one `Module` is not a backend problem and no backend change can fix it. ExecuTorch says so -itself, in `extension/module/module.h`: "This class is not thread-safe and performs no internal -synchronization. Calling execute concurrently on the same Module instance from multiple threads is -unsafe, regardless of whether share_memory_arenas is true or false." Note "regardless": two threads -on two different methods of one `Module` is unsafe too, not only two threads on the same method. -Mostly the runtime catches it and returns `Error::InvalidState` with "Inputs can not be set mid -execution", but not always: in the measurements above it also returned between 16 and 107 answers -per configuration that belonged to the other thread. - -### A crash when two threads load at the same time under the Python bindings - -Loading from more than one thread through the Python bindings crashes, roughly one attempt in six. -The same shape in C++ is clean (the sixth row of the table above), so this is specific to the -bindings and not to loading. - -The crash is a jump to a null address inside a stream flush, and the stack names every layer: - - TensorRTBackend::init - -> ET_LOG -> executorch::runtime::internal::logf -> vlogf (in libexecutorch.so) - -> std::cerr flush -> the sentry also flushes std::cout - -> std::cout's stream buffer -> address 0 - -ExecuTorch's logging writes to the process-wide `std::cerr`, and `std::cerr` is tied to -`std::cout`, so every log line flushes `std::cout` as well. The bindings swap what those two -streams point at, per call, so that output reaches a Python stream. When a second thread is -inside a log call while that swap happens, it flushes a stream buffer that is being replaced and -calls through a pointer that is no longer there. An earlier reading of this as a deadlock was -looking at the same stack with a different ending: whether it hangs on the interpreter lock or -crashes on the swapped buffer depends on where the second thread happens to be. - -Nothing in this backend can fix that, because a backend cannot log without going through -ExecuTorch's logging. The workaround is the same either way: load every program once on one -thread, then hand them out. The fix belongs in ExecuTorch, which should serialize its own logging -and should not redirect a process-wide stream while another thread may be writing through it. - -A standalone program using TensorRT the same way from several threads, with no ExecuTorch and no -Python, runs eighty thousand cycles cleanly. - -### Caller-stream contract for the TensorRT backend - -The upstream `CallerStreamGuard` documents the generic contract (per-thread, -nested scoping; the caller owns the stream for the guard's lifetime; the caller -manages host-data lifetime for async work). The TensorRT backend adds these -requirements, which previously lived on the removed `CudaStreamGuard`: - -- The selected stream must be on the TensorRT engine's device. -- Calls using one delegate handle must not overlap, and must not overlap with - its destruction; the backend serializes `execute()` calls with an internal - mutex, but destruction is not mutex-guarded. -- `execute()` always returns with the work finished, whatever the memory it was - given. The runtime this plugs into has no asynchronous execute: its own - `execute()` returns to mean the work is done, and callers read the outputs - straight afterwards. There was an opt-in that returned early, and it could not - be honoured safely, so nothing returns in flight any more. -- With no guard active, the backend falls back to `cudaStreamPerThread`. -- The backend always waits for the engine before returning, whatever stream it ran on. - ExecuTorch's runtime has no asynchronous execute: its `execute()` returns success to mean the - work is finished, every caller reads the outputs straight after it returns, and the API hands - back no event or future to wait on. There is therefore nobody an early return could be honest - with, and no option to ask for one. An earlier revision of this backend had such an option and - it is gone: on an idle GPU it returned zeros forty times out of forty, with the stream still - unready for a median of 142 microseconds, growing to 38 milliseconds on a larger graph. -- With the `use_shared_activation_scratch` backend option enabled, one buffer - per device backs the activation scratch of every execution context created - while it was on, so no two enqueues against it may overlap. The backend - enforces this itself: it holds a per-device lock from the claim on the buffer - through the enqueue and the completion event recorded on it, so two - `execute()` calls on one device are serialized at submission and the second's - stream waits on the first's enqueue. They may run on one stream or on two, and - they may be submitted concurrently from two threads, but they will not run - concurrently on the device, so the pool costs the parallelism between them. - Contexts created while the option was off keep their own scratch and are - unaffected, and so is a context whose engine needs no activation scratch under - any shape: that one is left out of the pool and serializes against nothing. -- The reference-runner smoke test runs inference inside a caller-stream guard on - the discrete-GPU CI configuration, and it takes the staging path there. The backend - binds a host pointer straight through only where the device both reads pageable host - memory and does so through shared host page tables. A discrete card reports the first - and not the second, because it serves pageable memory by faulting pages in one at a - time, which measured 33 times slower than one bulk copy on a loop that rewrites its - input each call. Either way that - configuration returns with the work already finished, because the backend always - waits. CI separately asserts that the runner resolves one shared - `libextension_cuda.so`. -- CUDA green-context streams work, and are the case this shared primitive exists - for: one `cuGreenCtxStreamCreate` stream drives both the TensorRT delegate and - ExecuTorch's CUDA/AOTI delegate, so both are confined to the same SM partition. - That both honour it is a property of the code, not of a passing run: the CUDA - delegate takes the per-thread stream when it initialises, then reads the caller's - selection again on every execute and routes every kernel and boundary copy through - it. Reading only the initialisation suggests the opposite, and a program split over - two separate streams still returns the right numbers, so a matching result settles - nothing either way. - Exercised by hand on an A100 with 108 SMs, using a `.pte` whose graph splits across - both delegates and a green context holding 8 of them. To reproduce, build the - reference runner, whose CMake configuration always enables the CUDA delegate, and - run it with `--green_context_sms=8`. - - CI runs the same option on the coalesced program, so this rests on more than the - hand measurement. A device that cannot provide the partition makes the runner exit - with its distinct status, and that case is skipped rather than failed, so a green - context is only exercised where the device has the SMs for one. - - -## Engine load errors - -Loading returns `Error::MemoryAllocationFailed` when TensorRT's last error contains -`OutOfMemory` or `out of memory`, ignoring case. This covers deserialization, -weight streaming budget application and execution context creation. Other TensorRT -load failures return `Error::InvalidProgram`. - -The log names the failed operation and preserves up to 1023 characters of the -TensorRT reason across short lines. Allocation failures include free and total -device memory in bytes, or the memory query error if those figures are unavailable. -These figures are not the failed allocation size. Loading does not clear a caller's -pending CUDA error to classify TensorRT failures. - -An engine serialized without its weights, for example compiled with -`immutable_weights=False, strip_engine_weights=True`, is refused with -`Error::InvalidProgram`, because this backend cannot refit it. To check, the backend -asks TensorRT how many weight bytes the plan still holds, and refuses the plan when that -is less than its total: - -- A stripped plan that had nothing to strip still holds all its weights, so it loads. -- A plan with no weights at all loads. -- A stripped plan saved again as it is still counts as stripped, so it is refused. -- A stripped plan saved again without `SerializationFlag::kEXCLUDE_WEIGHTS`, and not - refitted first, no longer counts as stripped, so this check does not catch it. Refit a - stripped plan before you save it again with its weights. - -For an ordinary engine with weights, the question makes TensorRT log a warning that the -engine was not built with `BuilderFlag::kSTRIP_PLAN`, each time the engine is -deserialized. The backend forwards it at Info level, and it is harmless. - -## Engines in a .ptd file - -A program exported with `external_engine_data=` keeps its engines in a `.ptd` file -next to the `.pte`. Load it with that file, for example -`Module module("model.pte", "model.ptd");`. Without the file, or with one that -lacks the engine, loading returns `Error::InvalidExternalData`. A program that shares -`aoti_cuda_blob.ptd` with the CUDA backend can fail first in that backend, with -`Error::NotFound`. - -The engine is GPU code, so treat the `.ptd` as trusted input, just like the `.pte`. -Each engine is stored under a key made from its sha256, but the key is only a name: -nothing checks the engine against it at load. - -## Shared engines - -Loading the same program twice in one process, for example once per robot arm, used to -deserialize its TensorRT engine twice, so its weights sat in device memory twice. Now handles -loaded from the same engine bytes, for the same device and the same weight streaming request, -share one engine. Each handle still creates its own execution context, I/O buffers and lock. From -C++, two handles of one engine can run at the same time on two threads, unless both handles -were loaded with shared activation scratch enabled and the engine needs scratch. Those two -handles then take turns, like all pooled calls on one device. From Python -they take turns anyway, because ExecuTorch's Python `execute` holds the interpreter lock for the -whole call. - -A second load still allocates its own context and buffers. In measurements on an 8 GB Jetson -Orin Nano, loading one robot policy twice increased CUDA-reported memory usage by 806 MiB for -the second load without sharing, compared with 122 MiB with sharing. Loading two methods that -carry the same engine in one program gave 742 MiB versus 120 MiB. These are whole-load changes in the board's shared -memory, not separate engine and context allocations. Serialized plan size is a different measure. -Inference speed was unchanged in those runs. To find a match, every load with sharing on hashes -the engine bytes, even a load that ends up sharing nothing. That takes about 0.15 s per GB of -engine on a Jetson AGX Thor and about 0.25 s per GB on the Orin Nano. - -- The weight streaming budget belongs to the engine, and TensorRT refuses to change it once a - context exists. Sharing uses the requested budget, not the value after TensorRT clamps or - ignores it. A different request gets its own engine. With no budget requested, the automatic - budget from the load that publishes first is kept. Racing loads can size that budget while - another copy occupies device memory. Pass an explicit budget for a predictable value. -- A match is found by hashing the engine bytes, because the bytes are freed after loading and - keeping a copy to compare would cost the memory being saved. The hash is not collision - resistant, so this relies on every program in the process being trusted input, as the - section on corrupted engines above already asks. -- The engine is freed when the last handle that uses it is destroyed. - -Set the `use_shared_engines` load option (a boolean, true by default) to false to keep one -module's engines private. Pass it through `Module::load`, like `weight_streaming_budget`. -It changes nothing for other modules loading at the same time. Private loads neither look up -nor publish an engine, and do not hash the bytes. A value of the wrong type is rejected with -`Error::InvalidArgument`. Passing this load-only key to `executorch::runtime::set_option` -also returns `Error::InvalidArgument`; pass it to `Module::load` instead. - -The ExecuTorch Python bindings do not expose load options, so every Python load hashes the -engine bytes and uses sharing. The C++ example below sets options before any method is loaded: +### Trusted input + +Treat every `.pte` and `.ptd` file as trusted input, like code. The engines in them run on the GPU, +and nothing checks them against a digest: the file format has none. A corrupted header is refused, +but corrupted engine bytes can load and return wrong answers. + +### Load errors + +- Loading returns `Error::MemoryAllocationFailed` when TensorRT runs out of memory while it + deserializes the engine, applies the weight streaming budget, or creates the execution context. + The log then shows the free and total GPU memory. Other TensorRT load failures return + `Error::InvalidProgram`, with TensorRT's reason in the log. +- An engine saved without its weights, for example compiled with + `immutable_weights=False, strip_engine_weights=True`, is refused with `Error::InvalidProgram`, + because this delegate cannot refit it. To check, the delegate asks TensorRT how many weight bytes + the engine still holds, and refuses it when that is less than its total. An engine that had + nothing to strip, or has no weights at all, loads. A stripped engine saved again as it is stays + refused. Saved again without `SerializationFlag::kEXCLUDE_WEIGHTS` and not refitted first, it no + longer counts as stripped and this check misses it, so refit before you save it again. For an + ordinary engine, TensorRT logs one harmless warning each time it is deserialized, saying it was + not built with `kSTRIP_PLAN`. + +### Known issues + +- **Coalesced programs with weights need one directory each.** A coalesced program is one split + between TensorRT and ExecuTorch's CUDA backend. The CUDA backend always names its weight file + `aoti_cuda_blob.ptd`, so a second export into the same directory overwrites the first one's + weights. The first program can then return wrong answers with no error, or fail to load. +- **Loading a coalesced program again and again grows host memory**, by about 400 KB per load, + run and drop. ExecuTorch's CUDA backend opens a compiled library on every load and never closes + it. GPU memory does not grow, and a program with only TensorRT engines does not grow. Load once + and keep the program if you can. +- **A coalesced program may fail on any GPU but the first**, with an illegal memory access. This + was seen only with a program whose recorded device was edited by hand, not one exported for that + device, so it is not confirmed. On a machine with several GPUs, prefer the first one for a + coalesced program. + +## Runtime options + +The delegate reads these options. A load option applies to one `Module::load` call. A process +option applies to every program loaded after it is set with `executorch::runtime::set_option`. +The [deployment guide](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#runtime-options) +shows how to set both, from C++ or from Python. + +| Option | Kind | Default | What it does | +| --- | --- | --- | --- | +| `use_cuda_graphs` | load and process, bool | the choice saved at export, else off | CUDA graph replay, below | +| `use_shared_activation_scratch` | process, bool | off | One scratch buffer per device, below | +| `use_shared_engines` | load, bool | on | Loads of the same engine share it, below | +| `weight_streaming_budget` | load, bytes as a decimal string or an int | the budget saved at export, else TensorRT's automatic one | GPU bytes for each engine's weights | + +A load option of the wrong type fails the load, and a process option of the wrong type fails +`set_option`. An int budget must be smaller than 2^31 bytes (2 GiB), so pass a larger one as a +string. `set_option` returns `Error::NotFound` when the delegate is not linked into the binary, and +`Error::InvalidArgument` for `use_shared_engines`, which is a load option only. It ignores keys it +does not read, including `weight_streaming_budget`. + +### CUDA graph replay + +With replay on, the delegate records an engine's kernel launches once as a CUDA graph, and then +replays the whole engine with one launch. It helps engines with fixed shapes that launch many +short kernels. + +Each engine reads the setting when it loads, so a later change does not affect engines already +loaded. The first match wins: + +1. The `use_cuda_graphs` load option. +2. A process option set to `false`. +3. The choice saved at export with `use_cuda_graphs=` in `torch_tensorrt.save()` or + `torch_tensorrt.executorch.export()`. +4. The process option, which is off if it was never set. + +How it runs: the first call with a set of input shapes runs normally, the second records the graph +and launches it, and later calls replay it. With the shared activation scratch on, a call does not +wait while another engine loads, is destroyed or records. It runs normally, and a later call +records. Each engine keeps one GPU buffer per input and output, and every call through the graph +copies the inputs in and the outputs out. For an engine with few kernels, or with large inputs and +outputs, these copies can cost more than they save. + +Some calls do not replay: + +- A change of input shapes, or of the optimization profile, records again. So does a growth of the + shared activation scratch below. A loop that changes shape or profile on every call never + replays, and still pays the copies. +- An engine with aliased outputs, a GPU without stream-ordered memory, or a driver older than CUDA + 12.5 never replays. The delegate logs the reason. +- A call on a green context stream, or on a stream from another CUDA context, takes the ordinary + path. The recorded graph is kept for the next call that can replay it. +- After three failed recordings for one set of shapes, the engine logs an error and stays on the + ordinary path until its shapes or its profile change. + +Engines on the shared activation scratch replay too. + +**When recording is unsafe.** The delegate keeps its own engine loads and destroys from overlapping +a recording. It cannot protect against the rest of the process. While a recording runs, these are +unsafe on any thread: + +- creating or destroying a TensorRT engine or execution context outside this delegate, +- a whole device sync, such as `cudaDeviceSynchronize()` or `torch.cuda.synchronize()`. Loading a + program with ExecuTorch's CUDA backend can make such a sync too. + +Recording can happen at any time in an engine's life, for example on its first larger shape. Leave +replay off when other threads may do any of these things. Also do not record a call to `execute()` +into a CUDA graph of your own. + +### Shared activation scratch + +Every TensorRT execution context keeps its own scratch memory for intermediate results. A model +lowered to many small engines pays for that many times, and can run out of GPU memory on the +number of engines alone. With the `use_shared_activation_scratch` process option on, the contexts +created after it share one buffer per device. The buffer grows to the largest size any call has +needed, so you save the sum of the separate sizes minus the largest one. An engine that needs no +scratch at all stays out of the pool. ```cpp -#include +#include -using namespace executorch::extension; using namespace executorch::runtime; -Error load_with_private_engines(Module& module) { - BackendOptions<1> options; - const Error stored = options.set_option("use_shared_engines", false); - if (stored != Error::Ok) { - return stored; - } - LoadBackendOptionsMap by_backend; - const Error mapped = by_backend.set_options("TensorRTBackend", options.view()); - if (mapped != Error::Ok) { - return mapped; - } - return module.load(by_backend); -} +BackendOptions<1> options; +options.set_option("use_shared_activation_scratch", true); +const Error error = set_option("TensorRTBackend", options.view()); ``` -Turn it off when two loads of one program must not share an engine, for example to compare -their memory, or when the loads run under different CUDA driver contexts on one device, which -sharing does not tell apart. +Set it before you load the methods that should use the pool. What it costs: + +- **No parallelism between pooled calls on one device.** The delegate holds a per-device lock from + claiming the buffer until the work is queued and marked. Two pooled calls can be submitted from + two threads, on one stream or two, but they never run at the same time on the device. +- **A pooled call waits for the pooled call before it.** Do not hold back the work of one pooled + call, for example with a host function on its stream, until a later pooled call returns. The + later call then never returns. +- **The buffer never shrinks.** The largest size ever needed stays allocated until the process + exits. Whether an engine asks for the size its current shapes need, or for the most its profile + allows, is decided when TensorRT builds it. +- **No CUDA graph capture.** See the capture section below. +- **No `cudaDeviceReset()`.** The pool keeps its buffer, an event and a stream for the life of the + process. Do not reset a device that ran a pooled engine. + +With `use_cuda_graphs` on as well, each replaying engine still keeps its own buffers for its inputs +and outputs. + +#### Capturing a call into a CUDA graph is not supported + +A pooled call whose stream is being captured fails with `Error::NotSupported`. The handoff from one +pooled call to the next waits on an event recorded outside the capture, and that wait fails with +`cudaErrorStreamCaptureIsolation` and invalidates the capture in every capture mode. A growth of +the buffer adds a `cudaMalloc` and a host wait, which invalidate a global or thread-local capture. +Without the refusal, you would learn of it only when `cudaStreamEndCapture` returns +`cudaErrorStreamCaptureInvalidated`. + +- The check runs before anything in the call that could break the capture. +- If the query fails for another reason, such as an earlier fault on the device, the call fails + with `Error::InvalidProgram` and that fault's message instead. +- Only pooled calls check, and only their own stream. A capture on another stream is not caught. + So do not overlap `execute()` with a global capture on any thread, or with a thread-local capture + on the calling thread. This holds for every stream, and also with the option off or for an engine + that needs no scratch: the call's final `cudaStreamSynchronize`, or the `cudaFree(nullptr)` on a + default stream, can invalidate the capture too. + +None of this makes capturing the delegate itself supported. + +#### How the pool grows + +The delegate asks each engine for its scratch size on every `execute()`, after the input shapes are +bound. An engine whose answer does not change with the shapes grows the pool at most once, on its +first run, so running the largest such engine first leaves one allocation. An engine whose answer +does change can grow it on any call that needs more than every call before it. Which kind an engine +is gets decided when TensorRT builds it, for example by the `kRUNTIME_ACTIVATION_RESIZE_10_10` +preview feature. + +A growth frees the old buffer before `execute()` returns, without holding the per-device lock: it +waits on the host for the last call that used the buffer, queues `cudaFreeAsync` on a stream the +pool owns, and synchronizes that stream. The host wait is a whole inference with no time limit, so +if work queued ahead of that call is released only after this thread returns, the call never +returns. Measured on an A100 with CUDA 13.0, with 512 MiB behind a 10 ms call and a host function +parked for 1500 ms on an unrelated stream: + +- A device-wide `cudaFree` waits for the whole device: 1500.8 to 1501.4 ms. +- The pool's own stream waits only for the buffer's work: 10.6 to 11.7 ms. +- The caller's stream does not work. A `cudaFreeAsync` of `cudaMalloc` memory gives the bytes back + only at the next synchronize of that stream, and a caller that keeps its data on the GPU may + never synchronize it. Measured, 4 GiB queued that way stayed lost for the life of the process, + and a growth sequence of 8, 16, 24 and 32 GiB ran out of memory on the fourth step. + +Where the device has no stream-ordered allocator, or the pool's stream could not be created, the +growth falls back to a device-wide `cudaFree`. + +A call whose shapes need no scratch still needs a buffer installed, because `enqueueV3` refuses a +context with none when its engine needs scratch for any shape. Against an empty pool such a call +allocates the one-byte minimum, not the engine's figure for its whole profile, because the pool +never shrinks: on one dynamic test engine that figure is 4 MiB, against 128 KiB for the next real +call. Each call also installs its own size rather than the pool's capacity. A size query that +failed then shows up, because TensorRT refuses a buffer smaller than the bound shapes need, and +the delegate fails the call with `Error::InvalidState` instead of running on a buffer it may have +freed. + +### Shared engines + +Loading the same program twice in one process, for example once per robot arm, keeps one +TensorRT engine in GPU memory instead of two. Loads share an engine when they have the same engine +bytes, the same device, and the same requested weight streaming budget. Each load still creates its +own execution context, buffers and lock, so two loads of one engine can run at the same time from +two threads, unless both use the shared activation scratch. + +Measured on an 8 GB Jetson Orin Nano, loading one robot policy a second time took 806 MiB more GPU +memory without sharing, and 122 MiB more with it. Inference speed did not change. To find a match, +every load hashes the engine bytes, which takes about 0.15 seconds per GB on a Jetson AGX Thor and +0.25 seconds per GB on the Orin Nano. + +- The hash does not resist deliberate collisions, so every program in the process must be trusted. +- The weight streaming budget belongs to the engine, so a load that asks for a different budget + gets its own engine. With no budget requested, the first load's automatic budget is kept. +- The engine is freed when the last load that uses it is destroyed. + +Set the `use_shared_engines` load option to `false` to keep one program's engines private, for +example to compare memory, or when loads run under different CUDA driver contexts on one device, +which sharing does not tell apart. A private load does not hash the bytes. + +### Engines in a .ptd file + +A program exported with `external_engine_data=` keeps its engines in a `.ptd` file +next to the `.pte`. Load it with that file, for example +`Module module("model.pte", "model.ptd");`. Without the file, or with one that lacks the engine, +loading returns `Error::InvalidExternalData`. A program whose engines share `aoti_cuda_blob.ptd` +with the CUDA backend can fail first in that backend, with `Error::NotFound`. Each engine is +stored under a key made from its SHA-256 hash, but the key is only a name: nothing checks the +engine against it. -## Optimization profiles +### Optimization profiles -An engine exported with several optimization profiles, for example one for single-token decode -and one for prefill, runs one profile at a time. A C++ host chooses it per thread with -`OptimizationProfileGuard` from `torch_tensorrt/executorch/TensorRTBackend.h`, scoped around the -call: +An engine exported with several optimization profiles, for example one for +decoding one token at a time and one for a long prompt, runs one profile at a time. A C++ host +picks it for each call with `OptimizationProfileGuard` from +`torch_tensorrt/executorch/TensorRTBackend.h`: ```cpp #include using torch_tensorrt::executorch_backend::OptimizationProfileGuard; -constexpr int32_t kDecodeProfile = 0; // indices follow the export-time profile order +constexpr int32_t kDecodeProfile = 0; // same order as the profiles at export constexpr int32_t kPrefillProfile = 1; { @@ -461,486 +435,26 @@ constexpr int32_t kPrefillProfile = 1; module.forward(prefill_inputs); } { - auto guard = OptimizationProfileGuard::automatic(); + auto guard = OptimizationProfileGuard::automatic(); // chosen from the input shapes module.forward(inputs); } ``` -- Without a guard every engine runs profile 0. An input outside its range fails with - `Error::InvalidArgument`, and when the engine has more profiles the message says how many and - points at the guard. -- A pinned index runs that profile, and only the shapes it covers. An index the engine lacks fails - the call. An engine with a single profile ignores a nonnegative pin, runs profile 0 and logs that - once, so a pin meant for another engine in the same method does not fail it. A negative index - always fails. +- Without a guard, every engine runs profile 0. An input outside profile 0's range fails with + `Error::InvalidArgument`, and when the engine has more profiles the message says so. +- A pinned index runs that profile, and an index the engine does not have fails the call. The one + exception is an engine with a single profile and a pin above 0: it runs profile 0 and logs that + once. A negative index always fails. - `automatic()` keeps the current profile while the inputs fit it, and otherwise takes the lowest - index that fits. Where ranges overlap the choice depends on earlier calls, so pin when it matters. -- The request belongs to the thread that made it, and a nested guard restores the enclosing one - when it ends. Every TensorRT delegate in the method reads the same index and resolves it against - its own profile list, so pin by index only when the engines were built from one list. -- Each handle keeps its own active profile, so handles of a shared engine switch independently. A - switch is issued on the call's stream, before the pooled scratch is claimed, and a switch drops - the handle's recorded CUDA graph (see [CUDA graph replay](#cuda-graph-replay)). -- Profile selection needs the header. The `torch-tensorrt-executorch-runtime` wheel ships none, - and on Windows its DLL does not export the guard, so an app that links the delegate from the - wheel always runs profile 0. Build against the source tree or the release tarball, which carries - the headers under `include/torch_tensorrt/executorch/`. Python has no way to choose a profile. - -## CUDA graph replay - -CUDA graph replay is off by default. When it is on, the delegate records an -engine's kernel launches once as a CUDA graph and then replays the whole engine -with one launch. This helps engines with fixed shapes that launch many short -kernels, where launch work on the CPU is a large share of each call. - -### Turning it on, and off - -Three settings choose it, all read when an engine loads, so a later change does -not affect engines already loaded. They apply in this order: - -1. The `use_cuda_graphs` load option passed to `Module::load` wins over every - other setting. It must be a boolean; any other type fails the load with - `Error::InvalidArgument`. -2. With no load option, an explicit false set through - `executorch::runtime::set_option` turns replay off for the process, even if - the program saved true. -3. Otherwise, the `use_cuda_graphs` compile spec saved at export chooses replay: - `b"1"` turns it on and `b"0"` turns it off. It wins over a process option set - to true and over the default when the process option was never set. Any other - value, or the key twice, fails the load with `Error::InvalidProgram`. -4. If the program has no saved choice, the process option applies. Replay stays - off when that option was never set. - -In Python, at export: - -```python -import torch_tensorrt - -torch_tensorrt.save( - trt_module, - "model.pte", - output_format="executorch", - use_cuda_graphs=True, -) -``` - -Both `save()` and `torch_tensorrt.executorch.export()` accept `use_cuda_graphs=True` -(on), `False` (off), or `None` (default, no baked choice). The choice applies to every -method's TensorRT delegates. Export rejects non-booleans and raw -`CompileSpec("use_cuda_graphs", ...)` entries; pass the keyword instead. - -In C++, a host that must not replay refuses it for one load: - -```cpp -#include - -using namespace executorch::extension; -using namespace executorch::runtime; - -Error load_without_cuda_graphs(Module& module) { - BackendOptions<1> options; - const Error stored = options.set_option("use_cuda_graphs", false); - if (stored != Error::Ok) { - return stored; - } - LoadBackendOptionsMap by_backend; - const Error mapped = by_backend.set_options("TensorRTBackend", options.view()); - if (mapped != Error::Ok) { - return mapped; - } - return module.load(by_backend); -} -``` - -A C++ host can also turn it on for every program that does not choose, before it -loads them: - -```cpp -BackendOptions<1> options; -options.set_option("use_cuda_graphs", true); -set_option("TensorRTBackend", options.view()); -``` - -### When recording is unsafe - -Recording can recur throughout an engine's lifetime. Every input shape change -and every optimization profile switch starts another warmup and recording cycle. A growth of the shared activation -scratch pool also drops the recorded graph: when any engine on a device needs -more than the pool holds, every engine on that device's pool records again. That -can happen long after warm-up, for example when a dynamic engine sees a larger -shape, so warming up every shape does not end the risk for those engines. -A failed recording is retried on the next call, up to three failures for that -set of shapes. A shape change or a profile switch resets that budget and a pool -growth does not, so an engine that used up its three failures stays on the -ordinary path until its shapes or its profile change. Loading everything before the first replay does not end the risk. - -A lock shared across this backend prevents recording from overlapping its own -engine and execution context creation and destruction. Loading or destroying a -program through this backend is protected from this recording race. A handle -must still not be destroyed while it is executing. An engine on the shared -activation scratch pool records while it holds its device's pool lock, so it -does not wait for this lock: while a load, a destroy, or another engine's -recording holds it, the engine runs that call with an ordinary launch and records -on a later call. - -The lock does not cover work outside this backend. Creating or destroying a -TensorRT engine or execution context elsewhere in the process remains unsafe -while recording runs. A whole device sync on any thread, such as -`cudaDeviceSynchronize` or `torch.cuda.synchronize()`, is also unsafe. The sync -itself can fail with `cudaErrorStreamCaptureUnsupported`, and the overlap can -crash the process inside the CUDA driver. Loading a program with ExecuTorch's -CUDA delegate can make such a sync too. - -Leave replay off when other threads may perform any of that work. Use an -explicit false process option to refuse a program's saved request, or a false -load option for one load. Replay also does not support wrapping `execute()` in -a caller's own CUDA capture. - -### What it costs - -Each engine keeps one stable device buffer per input and output, and every call -that runs through the graph copies each input in and each output out. For an -engine with few kernels, or with large inputs or outputs, these copies can cost -more than the launches they save. Shapes that change on every call never replay -and still pay the copies. Engines on the shared activation scratch pool pay this -too, so turning replay on for a model split into many small engines adds these -buffers back for every engine. For example, four pooled engines with a 4 MiB -input and a 4 MiB output each keep 32 MiB of these buffers next to the shared -scratch, and make 8 copies per round. - -### How it works - -The first call with a set of input shapes runs on the stable buffers with an -ordinary launch, the second records the graph on a private stream and launches -it on the caller's stream, and later calls replay. A shape change starts this -again, and so does an optimization profile switch at the same shapes, because a -graph replays the kernels of the profile it was recorded under. A handle keeps -one graph, so a loop that switches profile on every call, such as alternating -prefill and decode on one handle, records on every switch and never replays. -A new activation scratch buffer also starts recording again: an engine on the shared -activation scratch pool records against the buffer the pool hands it, and when a -larger request from any engine moves the pool to a new buffer, the engine's next -call drops the old graph. Recording then starts again on a later call, unless -the engine already used up its recording failures. Changing the caller stream -does not restart recording when the stream is in the same current CUDA context -and is not a green context stream. This includes the legacy default stream and -the default stream for each thread. A call with no caller stream replays on -`cudaStreamPerThread`. `execute()` waits for the stream before it returns, as it -always does. - -Engines with aliased outputs, GPUs without support for stream ordered memory, -and drivers older than CUDA 12.5 keep the ordinary path. Each reason is logged -once at load or first call. A call on a green context stream or a stream from -another context also takes the ordinary path, and the recorded graph is kept for -the next compatible call unless the scratch buffer, the shapes, or the profile -changed. Taking -the ordinary path does not guarantee that a stream from another context works -with the engine; the call can still fail. - -If allocating a graph buffer fails, the engine uses the ordinary path until an -input shape or the profile changes. Staging allocations can still fail with -`Error::MemoryAllocationFailed`; retained graph buffers are not reclaimed to -retry them. A failed recording is retried on the next call, and after three -failures for one set of shapes the engine logs an error and uses the ordinary -path until a shape or the profile changes. - -## Shared activation scratch - -A TensorRT execution context allocates its own activation scratch and holds it -for as long as the context lives, so a model lowered to N single-layer engines -pays N copies and can run out of device memory on the layer count alone. The -`use_shared_activation_scratch` backend option (a boolean, off by default) -instead backs a device's contexts from one buffer, grown to the largest -requirement any call on that device has asked for. Which contexts: those created -while the option was on, less any whose engine needs no scratch under any shape, -which are left out of the pool entirely. The two exceptions are spelled out -below. - -```cpp -#include - -executorch::runtime::Error enable_shared_activation_scratch() { - executorch::runtime::BackendOptions<1> options; - // Returns Error::InvalidArgument when the object has no room left. One key in a - // BackendOptions<1> always fits, so this cannot fail as written; the check is - // what keeps it honest when a second key is added beside it. - const executorch::runtime::Error stored = - options.set_option("use_shared_activation_scratch", true); - if (stored != executorch::runtime::Error::Ok) { - return stored; - } - // Nothing was set if this fails: every context still allocates its own scratch. - return executorch::runtime::set_option("TensorRTBackend", options.view()); -} -``` - -`Error::NotFound` means no backend is registered under that name, which is what a -binary that has not linked the backend archive gets. Nothing forces the check: the -free `executorch::runtime::set_option` is not `ET_NODISCARD`, so dropping its -return compiles. - -N per-context copies collapse to one, so the reclaimed memory is the sum of the N -requirements less the largest of them. Set the option before loading the methods -whose contexts should use the pool, and read the `use_shared_activation_scratch` -bullet of the caller-stream contract above: contexts sharing a buffer do not run -concurrently on the device. The pool never shrinks, so the largest scratch it was -ever asked for stays allocated until the process exits. - -Turning on `use_cuda_graphs` as well adds memory back: each replaying engine -keeps its own stable buffer for every input and output, as "What it costs" under -CUDA graph replay describes. - -A call that asks for nothing never grows a pool that already holds a buffer: it is -handed that buffer, whatever size it is. Against an empty pool it does allocate, -which the paragraph after next is about. It is not handed the pool's size with it -either: what the context is told it owns is the figure this call asked for, so a -zero becomes the one-byte minimum below and TensorRT is told the buffer is one -byte, whatever the pool's capacity is. Two of the things a zero from the per-shape -query can mean reach this point, and are indistinguishable where it is read: the -shapes bound to this call need none -- an empty batch inside a profile that admits -one -- or the query failed. Neither wants a buffer of its own. - -Such a call cannot be left with no buffer at all. `enqueueV3` refuses a -`kUSER_MANAGED` context with no device memory installed as soon as the engine -reports needing some under *any* shape, however little the shapes actually bound -need. So where the pool holds nothing yet, a zero allocates the smallest buffer -that satisfies that check, and the first call with a real requirement grows it. -Standing the engine's own figure in for the zero would satisfy the check too, but -that figure covers the whole profile: on the dynamic engine the tests use it is -4 MiB against the 128 KiB the next call needs, and the pool never shrinks. - -Only one of the two causes is safe to hand a buffer that small, and the install is -what separates them. `setDeviceMemoryV2` refuses a buffer smaller than the bound -shapes need, and it returns `void`: where the shapes genuinely need nothing the -expected size is zero and the one-byte buffer is accepted, but where the query -failed the engine expects what it always did, the install is refused, and the -context keeps the buffer it was last given -- which a growth may already have -freed, so the enqueue would read and write memory the pool no longer owns while -`enqueueV3` reports success. The backend reads the refusal back through a TensorRT -`IErrorRecorder` scoped to that one call and fails with `Error::InvalidState` -rather than enqueueing. Nothing reaches `enqueueV3` on an install TensorRT -rejected. - -What each call installs is its own requirement, not the capacity the pool holds, -which is larger whenever an earlier call asked for more. Handing a context the -larger figure would say it owns bytes that belong to whatever ran before it, and -the pool never clears the buffer between users -- activation scratch is written -before it is read, so a context that stays inside its own requirement never sees -those bytes, and one that runs past it sees another engine's activations rather -than zeros. The smaller figure also keeps the refusal above sharp: after a growth -the capacity may well cover what a failed query concealed, and an install sized -from it would be accepted. - -The third thing a zero can mean is an engine that needs no scratch under *any* -shape, and that one is settled before a call is ever made: such an engine is left -out of the pool entirely, so it takes no per-device lock and does not serialize -against the engines that do. - -The buffer grows when a call asks for more than every call before it did, and the -growth has to get rid of the buffer it replaces *before `execute()` returns*. -Not disposing of it inside the call is not an option: the pool grows monotonically, -so a run that let each retired buffer outlive its growth would end up holding the -sum of every size the pool was ever grown to rather than the largest of them. - -Every growth gets rid of it the same way, and none of the three steps reads -anything about the call the growth was made from: wait on the host for the enqueue -the handoff event names, queue a `cudaFreeAsync` on a stream the pool owns for that -device, and synchronize that stream, which is what returns the bytes. All three are -made with the per-device lock dropped, so none of them is serialized behind the -pool's own lock and none holds the next claimant off while it runs: two pooled calls -serialize at submission, as the caller-stream contract above says, and a growth is -not an exception to it. - -**The host wait is the thing to know about**, because it is what a growth pays -whenever the buffer it replaces has an enqueue recorded against it. It is a whole -inference, and it has no upper bound: whatever the enqueue is waiting for holds it -too, so if a host function ahead of that enqueue will be released only by this -thread once `execute()` returns, the growing call does not return and the thread -that would release it is the one inside the wait. - -The event it waits on is the handoff, and a call that reached its own enqueue has -recorded that enqueue there before it releases the claim, so the wait covers this -call's inference as well as the one before it. Only a growth pays that; a call with -nothing retired makes no CUDA call in the release at all. A growth whose retired -buffer never had an enqueue recorded against it does not pay it either: the -handoff carries no recording, so there is nothing to wait for and the disposal -goes straight to the free. A pooled call that returned between taking the buffer -and recording its enqueue -- a refused install, a refused `enqueueV3` -- is what -leaves the pool in that state. - -Read that cost against what `execute()` already does rather than on top of it. -Every call, pooled or not, ends by synchronizing the stream it ran on, so every -call already waits for its own engine and for whatever else the caller had queued -on that stream. What a growth adds is the wait for the *other* engine's enqueue, -the one against the buffer it is retiring, which may be on a stream this call never -touched. The pool's own disposal stream is what keeps the free itself off the -caller's stream, and the cases under `SharedScratchDisposalStream` are what pin -that; it is not visible in how long a call takes, because the caller-stream -synchronize dominates it. - -Sharing one buffer has an unbounded wait of its own, growth or no growth. A call -that synchronizes waits, through the handoff, for every pooled enqueue submitted -before it on that device. Park a host function ahead of one pooled call's enqueue -and release it only once a later pooled call has returned, and the later call is -the one that never returns. - -What a growth does *not* wait for is the rest of the device, and that is what the -pool's own stream buys. A device-wide `cudaFree` waits for everything queued on the -device rather than for the work that touched the buffer. Measured on an A100 with -CUDA 13.0 and driver 13.0, freeing a 512 MiB buffer behind a 10 ms enqueue with a -host function parked for 1500 ms on a stream nothing else in the measurement -submitted to: 1500.8-1501.4 ms device-wide against 10.6-11.7 ms on the pool's -stream, both with every byte back by the time the disposal returned. That wait is -not hypothetical: the backend's own tests deadlocked on the device-wide free once, -when a change of test order turned a case that parks a host function on its own -stream into the one that grew the pool. - -The stream has to be the pool's rather than the calling engine's. A `cudaFreeAsync` -of a `cudaMalloc`'d pointer -- which is what this pool holds -- hands the bytes -back at the next `cudaStreamSynchronize` of the stream it was queued on and at no -point before it. Measured on the same machine: it returns `cudaSuccess` and defers -the free, a stream `cudaStreamQuery` reports as drained still holds the bytes, -`cudaMemGetInfo` does not move, and a `cudaMalloc` of the same size fails with out -of memory until the synchronize. So whoever queues the free has to be able to -promise that synchronize, and only the pool can. Queuing it on the caller's stream -costs the disposal nothing, 0.0 ms in the same measurement, because all it does is -queue; the -whole 512 MiB was still resident when it returned, and the caller's own synchronize -is what gives it back. Where the caller makes one, that is no faster overall: -10.4-11.8 ms to the same point. Where it does not -- the caller-stream, -device-resident case, which is the one this option exists for -- the bytes never -come back at all: nothing this backend does will ever synchronize that stream, and -the stream is the caller's to destroy. Measured, a 4 GiB block queued -that way and then abandoned with the stream is gone for the life of the process, on -a device with tens of gigabytes still free: not at `cudaStreamDestroy`, not at -`cudaDeviceSynchronize`, not from another stream, and not at `cudaMemPoolTrimTo`. -Replaying the pool's own growth sequence on one such stream -- 8, 16, 24 and 32 GiB, -no synchronize between them -- four growths that all succeed when the disposal -returns the bytes fail on the fourth with out of memory, free stuck at 30.83 GiB -on a device that started with 78.83 GiB of it. - -That last is a property of the pointer rather than of stream-ordered frees in -general, so it is worth saying which one this pool holds. Replayed with -`cudaMallocAsync` on both sides of the same never-synchronized stream, all four -growths succeed: the free moves from 54.83 to 38.83 GiB for the 24 GiB request, -the retired 16 GiB reused with nobody having synchronized anything. Those bytes go -back to the device's memory pool rather than to the driver, and the next -allocation can take them from there. The pool allocates with `cudaMalloc`, so the -shape above is the one it is in, and the disposal stream is what that shape costs. - -Which calls do synchronize their stream is knowable, and the disposal does not ask. -It is a claim about what the caller has still to do rather than about anything the -pool holds, and three defects on this path in a row were a disposal getting that -claim wrong -- the last of them a call that failed after the growth and so had no -synchronization left to make, having recorded at the claim that it would make one. -So the disposal takes no argument: there is nothing to record when the buffer is -retired and nothing a later return can leave stale. - -A device-wide `cudaFree`, and the wait for the whole device that comes with it, is -still the fallback. It is what a growth makes where the device has no -stream-ordered allocator (`cudaDevAttrMemoryPoolsSupported`) and `cudaFreeAsync` -reports `cudaErrorNotSupported` instead of freeing -- and on any other code it -reports, since a queued free that did not happen is not a free -- and where the -pool's stream could not be created at all. A program that parks work across an -`execute()` has to keep such growths away from it, and the paragraph below on how -often the pool grows is what says whether a run order can do that. - -What an engine answers when asked how much it needs is decided when it is built, -not when it runs. The builder's `kRUNTIME_ACTIVATION_RESIZE_10_10` preview feature -makes an engine report what the shapes just bound need; without it, whether an -engine does that or reports its profile maximum depends on how TensorRT planned -it. Either way the pool can settle well above the live data, and nothing the -runtime does changes it. - -How often the pool grows follows from that. The backend asks afresh on every -`execute()`, after the input shapes are bound, so an engine whose answer does not -vary with the bound shapes grows the pool at most once, on its first run. For a -program built only from such engines, running the largest one first leaves the -pool with a single allocation. An engine whose answer does vary can grow it on any -call whose shapes need more than every call before them, so with one of those in -the program no run order bounds the number of allocations. One engine over a -`[1..4, 512, 512]` profile is either, depending on how it was built. - -### Capturing execute from the caller is not supported - -A pooled call whose selected stream is capturing is refused with -`Error::NotSupported`. The handoff between one enqueue and the next waits on an -event recorded outside the capture, and `cudaStreamWaitEvent` on such an event -fails with `cudaErrorStreamCaptureIsolation` and invalidates the capture under -every capture mode, `cudaStreamCaptureModeRelaxed` included. A growth adds more of -the same: its `cudaMalloc`, and the host wait its disposal of the replaced buffer -begins with, invalidate the capture under `Global` and `ThreadLocal`, and under -`Relaxed` are permitted but run uncaptured. Left to run, the caller learns of any -of it only when `cudaStreamEndCapture` hands back -`cudaErrorStreamCaptureInvalidated` and a null graph, long after the call that -caused it. - -The capture check runs before the wait on a previous enqueue, the `cudaMalloc` -that grows an input staging buffer, and the `cudaFree(nullptr)` that makes a -context current for a default stream. These calls precede the pool's own calls -and, outside `Relaxed`, can invalidate the capture. Only the device query and -the device switch run before the check, and a capture takes both. - -Where the query fails for a reason that is not a capture, the call is not told it -is capturing. `cudaStreamIsCapturing` also hands back a sticky fault left by -earlier work on the device -- measured on an A100 with CUDA 12.8, after an illegal -memory access it returns that fault with the capture status still -`cudaStreamCaptureStatusNone` -- and a call that meets one fails with -`Error::InvalidProgram` under that fault's own message instead. The one query -failure that does still mean a capture, `cudaErrorStreamCaptureImplicit`, is -refused with the rest. - -Only handles that claim pooled scratch check their selected stream. A capture -on another stream is not checked. Do not overlap `execute()` with a `Global` -capture on any thread or a `ThreadLocal` capture on the calling thread. This -applies to every stream, including streams created by the caller, with the option -off or with an option-on engine that needs no scratch. Those engines skip the -capture check, and their final `cudaStreamSynchronize` can fail and invalidate -the capture. On a default stream, `cudaFree(nullptr)` can fail earlier, before -enqueueing. A `ThreadLocal` capture on another thread is not affected by these -calls. This does not make capturing the delegate itself supported. - -### cudaDeviceReset() is not survivable - -The pool is not guarded against it: it holds its buffer, its handoff event and the -stream its growths dispose on for the process lifetime, and a reset destroys the -primary context under all three. The next pooled `execute()` on that device then -waits on a destroyed event and hands the engine a pointer that is no longer a -device allocation, and the next growth after that queues its free on a stream that -no longer exists. Guarding this would mean revalidating all three on every call, -and the check that would catch it is as expensive as the work it protects. Treat a -device the backend has run a pooled engine on as one that must not be reset. - -## Standalone Backend Archive - -Use this path only when you need `libexecutorch_trt_backend.a` without building -a runner that adds ExecuTorch with `add_subdirectory`. In that standalone mode, -build the ExecuTorch core runtime first: - -```bash -export EXECUTORCH_ROOT="${PWD}/executorch" -export TensorRT_ROOT=/path/to/extracted/TensorRT -export LD_LIBRARY_PATH="${TensorRT_ROOT}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" - -cmake -S "${EXECUTORCH_ROOT}" -B "${EXECUTORCH_ROOT}/cmake-out" \ - -DCMAKE_BUILD_TYPE=Release \ - -DBUILD_TESTING=OFF \ - -DEXECUTORCH_BUILD_PYBIND=OFF \ - -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ - -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ - -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ - -DEXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP=ON - -cmake --build "${EXECUTORCH_ROOT}/cmake-out" --target executorch_core -j -``` - -Then build the TensorRT backend archive from this package: - -```bash -cmake -S torch_tensorrt/src/torch_tensorrt/executorch -B build-torchtrt-executorch \ - -DEXECUTORCH_ROOT="${EXECUTORCH_ROOT}" \ - -DTensorRT_ROOT="${TensorRT_ROOT}" - -cmake --build build-torchtrt-executorch \ - --target executorch_trt_backend -j -``` + index that fits. Where ranges overlap, the choice depends on earlier calls, so pin when it + matters. +- The guard belongs to the thread that made it, and a nested guard restores the outer one when it + ends. Every TensorRT engine in the method reads the same index, so pin by index only when the + engines were built from one list of profiles. +- Each load keeps its own active profile, so loads that share an engine switch on their own. A + switch drops the load's recorded CUDA graph. +- The header is `cpp/include/torch_tensorrt/executorch/TensorRTBackend.h` in the source tree, and + `torch_tensorrt/include/torch_tensorrt/executorch/TensorRTBackend.h` in `libtorchtrt.tar.gz`. The + wheel ships no header. On Linux, an application can compile against the header, which also needs + the TensorRT headers, and link the wheel's delegate. On Windows, the wheel's library does not + export the guard. Python cannot choose a profile. diff --git a/docsrc/py_api/executorch.rst b/docsrc/py_api/executorch.rst index 13dfbd965a2..ac22676bae1 100644 --- a/docsrc/py_api/executorch.rst +++ b/docsrc/py_api/executorch.rst @@ -9,7 +9,8 @@ torch_tensorrt.executorch .. note:: - This module requires ExecuTorch and is only supported on Linux. + This module needs the ``executorch`` extra. It runs on Linux and on Windows x86-64. See + :ref:`executorch_deployment`. Overview -------- diff --git a/docsrc/user_guide/runtime_performance/executorch.rst b/docsrc/user_guide/runtime_performance/executorch.rst index 5245fbeed48..dc89b457635 100644 --- a/docsrc/user_guide/runtime_performance/executorch.rst +++ b/docsrc/user_guide/runtime_performance/executorch.rst @@ -3,26 +3,33 @@ ExecuTorch Deployment ===================== -**ExecuTorch** is PyTorch's runtime for edge devices. It loads one ``.pte`` file, holding the -program and every compiled payload, from a small C++ library that does not link libtorch. -That suits a robot, a drone, or any device where a full PyTorch install is too heavy. +`ExecuTorch `_ is PyTorch's runtime for edge devices. It runs a +model from one ``.pte`` file, with a small C++ library that does not need libtorch, PyTorch's own +C++ library. That suits a robot, a drone, or any device where a full PyTorch install is too heavy. -Torch-TensorRT writes that file. TensorRT compiles the operators it can convert into engines, -the engines are stored inside the ``.pte``, and a delegate hands each one back to TensorRT at -run time. Operators TensorRT will not take go to ExecuTorch's own CUDA backend in the same -program, so the whole model stays on the GPU. +Torch-TensorRT can write that file. TensorRT compiles the parts of the model it supports into +engines, and the engines are stored inside the ``.pte``. At run time, the TensorRT delegate (the +ExecuTorch plugin for those parts) hands each engine back to TensorRT. Operators that TensorRT +cannot compile can go to ExecuTorch's own CUDA backend in the same file, so the whole model stays +on the GPU. -**When to use ExecuTorch** +Use ExecuTorch when you want to: -* Deploying to an edge device or a C++ application with no Python and no libtorch. -* Shipping one file that holds the program, the TensorRT engines, and any fallback kernels. -* Driving the model from your own CUDA stream, including a green-context stream that pins the - work to part of the GPU. +* run the model from a C++ application with no Python and no libtorch, +* ship the program and its TensorRT engines in one file (a coalesced program, described below, + adds one data file), or +* run the model on your own CUDA stream, for example one that keeps it inside part of the GPU. .. note:: - The prebuilt ExecuTorch delegate runs on Linux, x86_64 and aarch64, and on Windows x86_64. - It requires CUDA 13. Ordinary Torch-TensorRT builds keep their separate CUDA 12 support. + The delegate runs on Linux (x86-64 and aarch64), and on Windows x86-64. The + prebuilt delegate needs CUDA 13. On Windows, the delegate loads the CUDA runtime from a CUDA 13 + toolkit install, so make sure ``CUDA_PATH`` or ``CUDA_HOME`` points at one. + + By default, TensorRT builds each engine for the operating system and the GPU it is exported on, + so export your programs on the kind of machine that will run them. On Windows, ExecuTorch's + CUDA backend cannot export a program, and its Python runtime there does not include that + backend, so a coalesced program (see below) is a Linux setup. ---- @@ -68,24 +75,28 @@ README shows. The delegate is built for the standard Linux wheels (x86-64 and aarch64), and for Windows x86-64. On aarch64 it needs glibc 2.35 or newer, the same as TensorRT itself. -The JetPack, TensorRT-RTX and Python-only builds have no delegate. To export a coalesced program -you also need a CUDA toolkit, because ExecuTorch's CUDA backend compiles the leftover operators -with ``nvcc``. +The JetPack, TensorRT-RTX and Python-only builds have no delegate. To export a coalesced program you also +need a CUDA toolkit, because ExecuTorch's CUDA backend compiles the leftover operators with +``nvcc``. ---- -Compile and Save ------------------ +Export a Model +--------------- -The usual ``ir="dynamo"`` path, with one extra argument to ``torch_tensorrt.save``: -``output_format="executorch"`` selects the ``.pte`` writer. For this format ``save`` keeps -the compiled graph as it is rather than re-exporting it, so it needs no example inputs. +Compile with ``torch_tensorrt.dynamo.compile``, then save with ``output_format="executorch"``: .. code-block:: python import torch import torch_tensorrt + + class MyModel(torch.nn.Module): + def forward(self, x): + return torch.relu(x) + 1 + + model = MyModel().eval().cuda() example_inputs = (torch.randn(2, 3, 4, 4, device="cuda"),) @@ -95,28 +106,43 @@ the compiled graph as it is rather than re-exporting it, so it needs no example arg_inputs=example_inputs, min_block_size=1, ) - torch_tensorrt.save(trt_gm, "model.pte", output_format="executorch") -:ref:`executorch_save` covers the rest: several methods in one file, a zero-copy key-value -(KV) cache, and the two-step ``torch_tensorrt.executorch.export()`` path for programs that -need work before they are written. +``save()`` writes the compiled graph as it is, so it needs no example inputs. + +:ref:`executorch_save` lists every option that ``save()`` takes for this format. It also shows +how to put several methods in one file, and how to let the engine of a large language model (LLM) +update its key-value (KV) cache in place. -Coalescing the leftover operators onto the CUDA backend -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Run the leftover operators on the CUDA backend +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -TensorRT has no converter for every ATen operator (PyTorch's core operator set), and by -default the ones it rejects run on the CPU, costing a copy each way on every call. Pass a -``CudaPartitioner`` and ExecuTorch's CUDA backend compiles them instead, so nothing leaves -the GPU: +TensorRT cannot convert every PyTorch operator. By default the ones it rejects run on +ExecuTorch's CPU kernels, with a copy to the CPU and back on every call, and an operator with no +CPU kernel, such as ``erfinv``, stops the program from loading. Pass a ``CudaPartitioner``, and +ExecuTorch's CUDA backend compiles them instead, so nothing leaves the GPU: .. code-block:: python from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + + class CoalescedModel(torch.nn.Module): + def forward(self, x): + return torch.cos(torch.erfinv(torch.tanh(x))) + + + coalesced_model = CoalescedModel().eval().cuda() + coalesced_inputs = (torch.randn(64, 64, device="cuda"),) + + coalesced_gm = torch_tensorrt.dynamo.compile( + torch.export.export(coalesced_model, coalesced_inputs), + arg_inputs=coalesced_inputs, + min_block_size=1, + ) torch_tensorrt.save( - trt_gm, + coalesced_gm, "coalesced.pte", output_format="executorch", partitioners=[ @@ -124,73 +150,28 @@ the GPU: ], ) -TensorRT partitions first and ``CudaPartitioner`` takes the rest. For ``cos(erfinv(tanh(x)))``, -where TensorRT cannot take ``erfinv``, the saved program's delegate list reads -``['TensorRTBackend', 'CudaBackend', 'TensorRTBackend']``. - -.. warning:: - - The CUDA backend names its external weight file per device, not per model, so saving two - coalesced programs into one directory overwrites the first one's weights. The first program - can then return wrong answers with no error, or fail to load. Give each export its own - directory. - -.. _executorch_external_engine_data: - -Keeping engines in a .ptd file -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -By default every TensorRT engine is stored inside the ``.pte``. Pass ``external_engine_data=`` -with a file name, without the ``.ptd`` extension, and the engines go to ``.ptd`` next to -the ``.pte`` instead: - -.. code-block:: python - - torch_tensorrt.save( - trt_gm, "model.pte", output_format="executorch", external_engine_data="model" - ) - -The program then needs that file to load. Pass the file itself, not its directory: - -.. code-block:: python - - from pathlib import Path - - import torch_tensorrt_executorch_runtime # noqa: F401 - from executorch.runtime import Runtime - - program = Runtime.get().load_program(Path("model.pte"), data_path=Path("model.ptd")) - -In C++, pass the file as the second argument of ``Module``: - -.. code-block:: cpp - - Module module("model.pte", "model.ptd"); +TensorRT takes what it can first, and the ``CudaPartitioner`` takes the rest. TensorRT cannot +take ``erfinv``, so this program's delegates are +``['TensorRTBackend', 'CudaBackend', 'TensorRTBackend']``. These docs call such a program +*coalesced*. -Without the file, the method fails to load. In C++ the error is ``Error::InvalidExternalData``. -In Python, ``load_method`` raises a ``RuntimeError`` that ends in ``error: 0x:24``, the same -error. The engines are GPU code, so treat the ``.ptd`` as trusted input, -just like the ``.pte``: nothing checks at load that its engines are the ones that were saved. +The CUDA backend compiles its part with ``nvcc``, the CUDA compiler, so the export needs a CUDA +toolkit. The CUDA backend also writes the weights of its part to a data file, +``aoti_cuda_blob.ptd``, next to the ``.pte``. Pass that file when you load the program, as shown +in `Run from Python`_ and `Run from C++`_. If the CUDA part has weights, loading fails without it. .. warning:: - The ``.ptd`` name does not follow the ``.pte`` name, and ``save()`` replaces any file with - that name. Say you save two programs into one directory with the same - ``external_engine_data=``. The ``.ptd`` then holds only the second program's engines, and the - first program fails to load. Name the file after its program, as above, or give each program - its own directory. - - A program with a ``CudaPartitioner`` part needs its own directory anyway (see the warning - above). The CUDA backend already writes ``aoti_cuda_blob.ptd`` for it. If you also want the - engines in a ``.ptd``, pass ``external_engine_data="aoti_cuda_blob"`` so they go into that same - file: ``load_program`` takes only one ``data_path``. + Every coalesced export names its weight file ``aoti_cuda_blob.ptd``, so a second export into + the same directory overwrites the first one's weights. The first model can then return wrong + answers with no error, or fail to load. Save each model into its own directory. -Keeping inputs and outputs on the GPU -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +Keep inputs and outputs on the GPU +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -By default ExecuTorch copies host to device before the first delegate and back after the -last, so a method is safe to call with CPU tensors. If your data is already on the GPU, those -copies are pure overhead: +By default, the program copies its inputs from the CPU to the GPU before the first delegate, and +its outputs back after the last one. That way you can call it with CPU tensors. If your data is +already on the GPU, those copies are wasted time. Turn them off at export: .. code-block:: python @@ -207,34 +188,66 @@ copies are pure overhead: skip_h2d_for_method_inputs=True, skip_d2h_for_method_outputs=True, ), - enable_non_cpu_memory_planning=True, - memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ), ), ) -Three things to know: +* **The two skip flags work only with** ``enable_non_cpu_memory_planning``, which is on by + default. With either flag set, turning it off makes the export raise an error. +* **Do not let the program allocate its inputs:** pass ``alloc_graph_input=False``. Otherwise + the program keeps its own input buffer, and every call fails with error ``0x10`` + (``NotSupported``). The log then says ``Planned-buffer copy needs host memory on both sides``. +* **Do not let it allocate its outputs either:** ``alloc_graph_output=False`` keeps the outputs + out of the program's memory. From Python, the runtime allocates GPU memory for each output, and + the next call does not overwrite an output you still hold, unless you pass + ``clone_outputs=False``. From C++, you must give the method your own GPU buffer once with + ``Module::set_output``, before the first run. If you keep outputs planned instead, the next call + overwrites them, so Python's ``execute()`` copies them on every call unless you pass + ``clone_outputs=False``. + +Pass CUDA tensors to a program exported this way. ExecuTorch's Python runtime refuses a CPU +tensor. From C++, a CPU tensor costs a copy on every call when the first delegate to use that +input is TensorRT, and fails when it is the CUDA backend. + +Keep the engines in a separate file +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +By default the engines are stored inside the ``.pte``, so the file is as large as the model. With +``external_engine_data``, they go to an ExecuTorch data file (``.ptd``) next to the ``.pte`` +instead. Give the name without the extension: + +.. code-block:: python + + torch_tensorrt.save( + trt_gm, + "small.pte", + output_format="executorch", + external_engine_data="small", + ) + +This writes a small ``small.pte`` and a ``small.ptd`` next to it. Pass the ``.ptd`` file itself, +not its directory, when you load the program. Without it, loading the method fails. The Python runtime takes one data file per +program, so for a coalesced program pass ``external_engine_data="aoti_cuda_blob"``: the engines +then go into the CUDA backend's file. + +Treat a ``.ptd`` file as trusted input, the same as a ``.pte``. The engines in it are GPU code. -* **Both skip flags need** ``enable_non_cpu_memory_planning=True``, because the copies are - inserted during device-aware memory planning. Asking for a skip with it off raises. -* **Inputs must be unplanned**, through ``MemoryPlanningPass(alloc_graph_input=False)``. - Otherwise the program reserves its own input buffer and the runtime fills it from the - caller, which puts the copy straight back. -* **The outputs can stay in the program's own GPU memory**, its device arena. To give each output - its own buffer instead, also pass ``alloc_graph_output=False``. Then each output needs a GPU - buffer before the call: in C++ you set it with ``Module::set_output``, and the Python runtime - allocates one for you. +.. warning:: -The choice is baked into the file, so feed such a program CUDA tensors. ExecuTorch's Python -runtime refuses a CPU tensor. From C++, a CPU tensor costs a copy on every call when the first -delegate to use that input is TensorRT, and fails when it is the CUDA backend. + The ``.ptd`` name does not follow the ``.pte`` name, and ``save()`` replaces any file with that + name. Save two programs into one directory with the same ``external_engine_data``, and the + ``.ptd`` holds only the second program's engines, so the first fails to load. Name the file + after its program, as above, or give each program its own directory. ---- -Python Inference ------------------ +Run from Python +---------------- -Import the delegate package once, before any program is loaded. That import is what -registers the backend. Everything after it is ExecuTorch's own API: +Import the delegate package once, before you load a program. The import is what registers the +delegate. Everything after it is ExecuTorch's own API: .. code-block:: python @@ -246,66 +259,93 @@ registers the backend. Everything after it is ExecuTorch's own API: program = Runtime.get().load_program(Path("model.pte")) forward = program.load_method("forward") - outputs = forward.execute((torch.ones(2, 3, 4, 4),)) + outputs = forward.execute((torch.randn(2, 3, 4, 4),)) + +If the delegate cannot load, the import raises an error right away, instead of a program failing +to load later. Torch-TensorRT itself is needed only to export, not to run. A program exported as +above takes CPU tensors. + +If the export wrote a data file, such as ``aoti_cuda_blob.ptd`` from a coalesced export, pass it. +The runtime does not look for it on its own, and if the file holds weights or engines, loading +fails without it: + +.. code-block:: python + + program = Runtime.get().load_program( + Path("coalesced.pte"), data_path=Path("aoti_cuda_blob.ptd") + ) + outputs = program.load_method("forward").execute((torch.ones(64, 64),)) + +A program that keeps its inputs and outputs on the GPU takes CUDA tensors and returns CUDA tensors: -If the delegate cannot load, that import raises at once, instead of surfacing later as a -program that will not load. +.. code-block:: python -A coalesced program needs no extra import: both backends are registered and the program -records which parts go where. If its export wrote a data file such as ``aoti_cuda_blob.ptd``, -pass it with ``load_program(..., data_path=Path("aoti_cuda_blob.ptd"))``. Torch-TensorRT is only -needed to export, not to run. + program = Runtime.get().load_program(Path("device_resident.pte")) + outputs = program.load_method("forward").execute( + (torch.ones(2, 3, 4, 4, device="cuda"),) + ) .. note:: - ``torch_tensorrt.load(path, format="executorch")`` still works but is deprecated. It - copies CUDA inputs to the CPU and supports embedded weights only. Use the Runtime API - above instead. A device-resident program runs only through the Runtime API. + ``torch_tensorrt.load(path, format="executorch")`` is deprecated. It copies CUDA inputs to the + CPU, so it runs only programs that take CPU inputs and have no data file. Use the runtime API + above instead. ---- -C++ Inference --------------- +Run from C++ +------------- -The wheels ship a prebuilt delegate and a CMake package, so a C++ application links them -without building anything: +The wheels ship a prebuilt delegate and a CMake package, so a C++ application can link it without +building anything. The steps below are for Linux. Put the program shown further down in +``main.cpp``, the ``model.pte`` from `Export a Model`_ in the same folder, and this next to them as +``CMakeLists.txt``: .. code-block:: cmake - find_package(executorch REQUIRED COMPONENTS backend_cuda extension_cuda kernels_optimized) + cmake_minimum_required(VERSION 3.28) + project(my_app CXX) + set(CMAKE_CXX_STANDARD 17) + + find_package(executorch REQUIRED COMPONENTS backend_cuda kernels_optimized) find_package(executorch_backend_tensorrt REQUIRED) + add_executable(my_app main.cpp) target_link_libraries(my_app PRIVATE executorch::runtime executorch::backend_cuda executorch::backend_tensorrt - executorch::extension_cuda executorch::kernels_optimized ) -* ``kernels_optimized`` supplies the ``et_copy`` operators that move data across the method - boundary. -* ``backend_cuda`` registers the device allocator those copies use, so even a TensorRT-only - program needs it. Without it the program loads and then the first instruction fails with - ``_h2d_copy: no device allocator registered``. -* ``extension_cuda`` provides ``CallerStreamGuard``, used below to choose the CUDA stream. +* ``kernels_optimized`` provides the copy operators that move data between the CPU and the GPU at + the edges of the program. Without it, loading fails with + ``kernel 'et_copy::_h2d_copy.out' not found``. +* ``backend_cuda`` registers the allocator for the program's GPU memory. A program that uses only + TensorRT needs it too. Without it, loading fails with ``No device allocator registered``. -The two packages live in two distributions, so point CMake at both. ExecuTorch is a -namespace package, so take its path from the distribution metadata, not from ``__file__``: +The two CMake packages are inside two Python packages, so point CMake at both. ExecuTorch's +location comes from its package metadata, because ``executorch.__file__`` is not set. Build and +run in the Python environment where you installed the packages: .. code-block:: bash - cmake -DCMAKE_PREFIX_PATH="$(TORCH_TENSORRT_SKIP_DELEGATE_REGISTRATION=1 python -c 'import importlib.metadata as m, torch_tensorrt_executorch_runtime as r, pathlib; print(str(pathlib.Path(str(m.distribution("executorch").locate_file("executorch"))) / "share" / "cmake") + ";" + str(pathlib.Path(r.__file__).parent))')" ... + ET_CMAKE=$(python -c 'import importlib.metadata as m, pathlib; print(pathlib.Path(str(m.distribution("executorch").locate_file("executorch"))) / "share" / "cmake")') + DELEGATE=$(TORCH_TENSORRT_SKIP_DELEGATE_REGISTRATION=1 python -c 'import pathlib, torch_tensorrt_executorch_runtime as r; print(pathlib.Path(r.__file__).parent)') + cmake -S . -B build -DCMAKE_PREFIX_PATH="${ET_CMAKE};${DELEGATE}" + cmake --build build + ./build/my_app -The environment variable keeps that a path query, so the delegate is not loaded just to -print a path. +The variable in the second line lets Python import the delegate package to read its path, without +loading the delegate. -CMake 3.28 or newer is required. Older versions write the ``$ORIGIN`` token in a runtime -search path incorrectly, so ``backend_cuda`` rejects them. CMake is not in these wheels, and -a freshly imaged Jetson has none. +You need CMake 3.28 or newer: on older versions the ExecuTorch package does not provide the +``executorch::`` targets used above. The runtime wheel README shows a way for CMake 3.19 to 3.27. +CMake is not in these wheels, so install it yourself. You do not need to find TensorRT: the +delegate finds it in the TensorRT wheel on its own. -There is no header to include. The delegate registers itself when its library loads, and the -rest is the ordinary ExecuTorch C++ API: +There is no header to include. The delegate registers itself when its library loads, and the rest +is the ordinary ExecuTorch C++ API: .. code-block:: cpp @@ -332,64 +372,95 @@ rest is the ordinary ExecuTorch C++ API: return 0; } -Linking also records the wheel's library directory in your binary, so the application finds -the delegate with no library path set. That suits a local build and not anything you -redistribute. To turn it off, set both of these before ``find_package`` and ship the -delegate yourself: +If the export wrote a data file, pass it to the constructor: + +.. code-block:: cpp + + Module module("coalesced.pte", "aoti_cuda_blob.ptd"); + +On Linux, linking the target also records the wheel's library directory in your binary, so the +application finds the delegate with no library path set. That is right for a local build, but not +for anything you ship to other machines. To turn it off, set both of these before +``find_package``, and then ship the delegate yourself: .. code-block:: cmake set(EXECUTORCH_BACKEND_TENSORRT_EMBED_RUNPATH OFF) set(CMAKE_SKIP_BUILD_RPATH ON) -You do not need to find TensorRT. The delegate locates it from the sibling wheel on its own, -and your application talks to the delegate rather than to TensorRT. +To build the delegate from source instead, follow the +`delegate README `_. +The same README lists the rules a C++ caller has to keep (threads, buffers, streams), the load +errors, and the known issues. ``examples/executorch_reference_runner/`` is a complete C++ +example. -Building the delegate from source -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +---- -The source ships inside ``libtorchtrt.tar.gz`` as -``torch_tensorrt/src/torch_tensorrt/executorch/``. Turn on ExecuTorch's CUDA backend and the -extensions a ``Module`` app needs, add the delegate beside ExecuTorch, and link its target: +Runtime Options +---------------- -.. code-block:: cmake +Some of the options below are saved in the ``.pte`` at export, and some are set when the program +runs. - set(EXECUTORCH_BUILD_CUDA ON CACHE BOOL "" FORCE) - set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON CACHE BOOL "" FORCE) - set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER ON CACHE BOOL "" FORCE) - set(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR ON CACHE BOOL "" FORCE) - set(EXECUTORCH_BUILD_EXTENSION_MODULE ON CACHE BOOL "" FORCE) - set(EXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP ON CACHE BOOL "" FORCE) - - add_subdirectory("executorch") - add_subdirectory("torch_tensorrt/src/torch_tensorrt/executorch") - - target_link_libraries(my_runner PRIVATE - executorch - executorch::backends - executorch::extensions - executorch::kernels - executorch::backend_tensorrt +Setting a runtime option +^^^^^^^^^^^^^^^^^^^^^^^^^ + +A *load option* applies to one program. In C++, pass it to ``Module::load`` before the first call: + +.. code-block:: cpp + + #include + + using namespace executorch::extension; + using namespace executorch::runtime; + + Module module("model.pte"); + BackendOptions<1> options; + options.set_option("use_cuda_graphs", false); + LoadBackendOptionsMap backend_options; + backend_options.set_options("TensorRTBackend", options.view()); + const Error error = module.load(backend_options); + +In Python, pass it to ``load_program``: + +.. code-block:: python + + program = Runtime.get().load_program( + Path("model.pte"), + backend_options={"TensorRTBackend": {"use_cuda_graphs": False}}, ) -Building it needs CUDA Toolkit 12.5 or newer. Keep ``libextension_cuda`` shared rather than -static, so every delegate in the process reads the same caller-stream state. +A *process option* applies to every program loaded after it is set. In C++: -``examples/executorch_reference_runner/`` is a complete C++ example. ``libtorchtrt.tar.gz`` -also carries it prebuilt as ``torch_tensorrt/bin/example_executorch_runner``. +.. code-block:: cpp ----- + #include + + using namespace executorch::runtime; + + BackendOptions<1> options; + options.set_option("use_shared_activation_scratch", true); + const Error error = set_option("TensorRTBackend", options.view()); -Runtime Performance --------------------- +``Error::NotFound`` from ``set_option`` means the delegate is not linked into your binary. + +In Python: + +.. code-block:: python + + Runtime.get().backend_registry.set_option( + "TensorRTBackend", {"use_shared_activation_scratch": True} + ) + +.. _executorch_single_stream: Choosing the CUDA stream ^^^^^^^^^^^^^^^^^^^^^^^^^ -With no stream chosen, both delegates use ``cudaStreamPerThread``, the calling thread's -default stream, so a coalesced program run from one thread is already ordered correctly. +With no stream chosen, both delegates run on ``cudaStreamPerThread``, the default stream of the +calling thread. So a coalesced program run from one thread is already in the right order. -To put every delegate on a stream of your own, scope a guard over the whole execution: +To run every delegate on a stream of your own, put one guard around the whole call: .. code-block:: cpp @@ -400,35 +471,54 @@ To put every delegate on a stream of your own, scope a guard over the whole exec cuda::CallerStreamGuard guard(stream); module.forward(input); -* One guard reaches every delegate, because they share one ``libextension_cuda``. -* The stream must be on the engine's device. -* The CUDA backend refuses a caller stream for a method that uses its own CUDA graphs. -* **Synchronize your stream before reading outputs.** The TensorRT delegate always waits for - its work, but ExecuTorch's CUDA backend can return once the work is only queued, so - ``execute()`` on a coalesced program can return early. - :ref:`Running a coalesced .pte ` has the same rule for a decode - loop. +The guard is in the ``extension_cuda`` library, so add it to the CMake recipe above. Your own CUDA +calls, such as creating the stream, also need the CUDA runtime library: -That stream can be a green context, which holds a fixed number of streaming multiprocessors (SMs), -so the model stays inside that partition and the rest of the GPU is free. Create it with -``cuGreenCtxStreamCreate`` and scope the same guard over it. The limit rides the stream, so -the green context does not have to be made current. +.. code-block:: cmake -The reference runner has this built in. Pass the number of SMs: + find_package(executorch REQUIRED COMPONENTS backend_cuda extension_cuda kernels_optimized) + find_package(CUDAToolkit REQUIRED) + target_link_libraries(my_app PRIVATE executorch::extension_cuda CUDA::cudart) + +* One guard reaches every delegate, because they all share one ``extension_cuda`` library. +* The stream must be on the same device as the engine. +* The CUDA backend refuses a caller stream for a method that uses its own CUDA graph option. +* **Synchronize the stream before you read GPU outputs.** The TensorRT delegate always waits for + its work to finish, but the CUDA backend can return while its work is still queued. + +That stream can belong to a CUDA green context. A green context owns a fixed number of the GPU's SMs +(streaming multiprocessors), so the model stays inside that part of the GPU and the rest is free for +other work. Create the stream with ``cuGreenCtxStreamCreate`` and put the same guard around the +call. ``cuGreenCtxStreamCreate`` is a CUDA driver function, so also link ``CUDA::cuda_driver``. The +green context does not need to be made current. TensorRT then warns that the engine was built +without green context constraints, which it says may cause crashes or undefined behavior. +Torch-TensorRT has no option to build engines with those constraints. + +The reference runner in ``examples/executorch_reference_runner/`` can do this for you. Build it as +its README shows (the prebuilt copy described below cannot run this program), then pass the number +of SMs: .. code-block:: bash example_executorch_runner --model_path=coalesced.pte --green_context_sms=8 -If a green context cannot be created it fails with a distinct status instead of falling back, -so a passing run really used one. ``0``, the default, uses an ordinary stream. +If the device cannot provide the partition, the runner stops with exit status 2 instead of using +an ordinary stream. ``0``, the default, uses an ordinary stream. The runner takes no data file, so +it runs a coalesced program only when the CUDA part has no weights, like the one +``export_coalesced.py`` writes, and it cannot run a program saved with ``external_engine_data``. + +A prebuilt copy of the runner ships as ``torch_tensorrt/bin/example_executorch_runner``, in +``libtorchtrt.tar.gz`` and in the Torch-TensorRT wheels for Linux, except the JetPack and +TensorRT-RTX builds. It has no CUDA backend, and its only CPU operators are the copies at the +edges of the program. So it runs only programs that TensorRT takes in full, not coalesced ones. CUDA graph replay ^^^^^^^^^^^^^^^^^^ -Replay is off by default. With it on, the delegate records an engine's kernel launches once -and then replays the whole engine with a single launch. It helps fixed-shape engines that -launch many short kernels, where CPU launch work dominates each call. +With replay on, the delegate records an engine's kernel launches once as a CUDA graph, and then +replays the whole engine with a single launch. This helps engines with fixed shapes that launch +many short kernels, where the CPU work of launching them is a large part of each call. It is off +by default. Turn it on at export: .. code-block:: python @@ -439,104 +529,128 @@ launch many short kernels, where CPU launch work dominates each call. use_cuda_graphs=True, ) -``save()`` and ``torch_tensorrt.executorch.export()`` both take ``use_cuda_graphs=True``, -``False``, or ``None`` for no baked choice. It applies to every method's TensorRT delegates -and is read when an engine loads. +``save()`` and ``torch_tensorrt.executorch.export()`` take ``True``, ``False``, or ``None`` (the +default, which saves no choice). The choice applies to every TensorRT engine in the program. At +run time, the first of these that is set decides: + +#. the ``use_cuda_graphs`` load option, +#. a process option set to ``False``, +#. the choice saved in the program, +#. a process option set to ``True``. + +With none of them set, replay is off. + +Replay is not free. Each engine keeps its own GPU buffer for every input and output, and every +replayed call copies the inputs in and the outputs out. With few kernels, or with large inputs +and outputs, those copies can cost more than the launches they save. -It is not free. Each engine keeps a device buffer per input and output, and every replayed -call copies each one in and out. With few kernels, or large inputs and outputs, those copies -can cost more than the launches they save. +Some calls never replay: -Some engines never replay: changing shapes, aliased outputs, a green-context stream, GPUs -without stream-ordered memory, and drivers older than CUDA 12.5. Only changing shapes still pay -the copies; the others take the ordinary path without the extra buffers. The delegate logs the -reason once for aliased outputs, missing stream-ordered memory and an old driver. +* An engine whose input shapes change on every call never replays, and it still pays the copies. +* An engine with aliased outputs, a GPU without stream-ordered memory, and a driver older than + CUDA 12.5 never replay. They take the ordinary path, without the extra copies. +* A call on a green context stream, or on a stream of another CUDA context, takes the ordinary + path for that call. -Engines on the shared activation scratch below do replay, each with its own buffers for its -inputs and outputs. When the pool moves to a larger buffer, an engine records again, as it does -after a shape change. +An engine on the shared activation scratch below also replays. It records again whenever the +shared buffer grows. .. warning:: - While a recording runs, no other thread may create or destroy a TensorRT engine or - execution context outside this delegate, or call a whole-device sync such as - ``cudaDeviceSynchronize`` or ``torch.cuda.synchronize()``. Every input shape change, and - every growth of the shared activation scratch, starts another recording, so loading - everything up front does not end the risk. Leave replay off when other threads may do any - of that. + While a graph is being recorded, no other thread may create or destroy a TensorRT engine or + execution context outside this delegate, or call a whole device sync such as + ``cudaDeviceSynchronize()`` or ``torch.cuda.synchronize()``. Loading a program with + ExecuTorch's CUDA backend can make such a sync too. Recording happens again after every shape + change, so loading everything first does not remove the risk. Leave replay off when other + threads may do any of these things. -A C++ host can refuse a program's saved request for one load: +Shared activation scratch +^^^^^^^^^^^^^^^^^^^^^^^^^^ -.. code-block:: cpp +Every TensorRT execution context keeps its own scratch memory for intermediate results, for as +long as it lives. A model split into many small engines pays for that many times, and can run out +of GPU memory on the number of engines alone. With the ``use_shared_activation_scratch`` process +option, all contexts on a device share one buffer instead. The buffer grows to the largest size +any call has needed, so you save the sum of the separate sizes minus the largest one. - #include +Set the option before you load the methods that should use the shared buffer. Methods loaded +before it keep their own scratch. - using namespace executorch::extension; - using namespace executorch::runtime; +The cost is parallelism. Two calls that share the buffer never run at the same time on one +device: the delegate takes a per-device lock from claiming the buffer to queuing the work. The +buffer never shrinks. Do not call ``cudaDeviceReset()`` on a device that ran an engine with the +shared buffer. Recording a call into your own CUDA graph is not supported, with or without the +shared buffer. - Error load_without_cuda_graphs(Module& module) { - BackendOptions<1> options; - const Error stored = options.set_option("use_cuda_graphs", false); - if (stored != Error::Ok) { - return stored; - } - LoadBackendOptionsMap by_backend; - const Error mapped = by_backend.set_options("TensorRTBackend", options.view()); - if (mapped != Error::Ok) { - return mapped; - } - return module.load(by_backend); - } +Shared engines +^^^^^^^^^^^^^^^ -Shared activation scratch -^^^^^^^^^^^^^^^^^^^^^^^^^^ +When the same program is loaded twice in one process, for example once per robot arm, both loads +share one TensorRT engine. The engine's weights are then in GPU memory once. Each load still gets +its own execution context and buffers. Loads share an engine when they have the same engine +bytes, the same device, and the same weight streaming budget. -Every TensorRT execution context holds its own activation scratch for as long as it lives, so -a model split into many small engines can run out of device memory on the engine count alone. -This option backs all of a device's contexts from one buffer, grown to the largest -requirement seen so far: +This is on by default. On an 8 GB Jetson Orin Nano, loading one robot policy a second time took +806 MiB more GPU memory without sharing, and 122 MiB more with it. To find a match, every load +hashes the engine bytes, which takes about 0.15 seconds per GB on a Jetson AGX Thor. The hash +does not resist deliberate collisions, so every program in the process must be trusted. -.. code-block:: cpp +To keep one program's engines private, pass the ``use_shared_engines`` load option as ``false``. - #include +Weight streaming +^^^^^^^^^^^^^^^^^ - using namespace executorch::runtime; +An engine compiled with ``enable_weight_streaming=True`` can keep part of its weights in CPU +memory and stream them to the GPU when needed. The budget is the number of bytes of weights that +stay on the GPU. Save one with ``weight_streaming_budget_per_engine``, or set one at run time +with the ``weight_streaming_budget`` load option, which takes the number of bytes as a decimal +string, or as an int below 2 GiB. The budget applies to each engine separately, so a program with N engines can keep up to N times +the budget on the GPU. With no budget, TensorRT chooses one from the free memory of the device. - Error enable_shared_activation_scratch() { - BackendOptions<1> options; - const Error stored = options.set_option("use_shared_activation_scratch", true); - if (stored != Error::Ok) { - return stored; - } - return set_option("TensorRTBackend", options.view()); - } +Optimization profiles +^^^^^^^^^^^^^^^^^^^^^^ -Set it before loading the methods that should use the pool. You save the sum of the separate -requirements less the largest one. ``Error::NotFound`` means the delegate is not linked in. +An engine can hold several optimization profiles: one set of weights, tuned for several ranges +of input shapes. The usual case is an LLM, with one profile for decoding one token at a time and +one for a long prompt. Declare the profiles with ``torch_tensorrt.Input(profiles=[...])``, as +:ref:`multi_optimization_profiles_tutorial` shows. A C++ application then picks the profile for +each call with ``OptimizationProfileGuard``: -The cost is parallelism. The backend holds a per-device lock from claiming the buffer through -the enqueue, so two calls on one device are serialized at submission. The pool never shrinks, -and a device that has run a pooled engine must not be reset with ``cudaDeviceReset()``. +.. code-block:: cpp -Shared engines -^^^^^^^^^^^^^^^ + #include + + using torch_tensorrt::executorch_backend::OptimizationProfileGuard; -Loading the same program twice in one process, say once per robot arm, keeps one copy of its -weights on the GPU. Handles built from the same engine bytes, on the same device and with the -same weight streaming request, share one engine. Each still gets its own execution context, -buffers and lock. + constexpr int32_t kDecodeProfile = 0; // same order as the profiles at export + constexpr int32_t kPrefillProfile = 1; + + { + OptimizationProfileGuard guard(kPrefillProfile); + module.forward(prefill_inputs); + } + { + OptimizationProfileGuard guard(kDecodeProfile); + module.forward(decode_inputs); + } -This is on by default. On an 8 GB Jetson Orin Nano, loading one robot policy twice grew -reported GPU memory by 806 MiB without sharing and by 122 MiB with it. The match is found by -hashing the engine bytes on every load, roughly 0.15 s per GB on a Jetson AGX Thor. To keep a -module's engines private, pass the ``use_shared_engines`` load option as ``false``, the same -way as ``use_cuda_graphs`` above. +To let each engine choose from the input shapes instead, create the guard with +``OptimizationProfileGuard::automatic()``. With no guard, every engine runs profile 0, and Python +cannot choose a profile. The header is +``cpp/include/torch_tensorrt/executorch/TensorRTBackend.h`` in the source tree, and +``torch_tensorrt/include/torch_tensorrt/executorch/TensorRTBackend.h`` in ``libtorchtrt.tar.gz``. +The wheel ships no header. On Linux, an application can still compile against the header from +the source, which also needs the TensorRT headers, and link the wheel's delegate. On Windows, the +wheel's library does not export the guard. ---- Comparison: .pte vs .pt2 ------------------------- +A ``.pt2`` file is the other way to run a Torch-TensorRT model without Python: AOTInductor, +PyTorch's ahead-of-time compiler, packages the whole model, and :ref:`aot_inductor` describes it. + .. list-table:: :widths: 30 35 35 :header-rows: 1 @@ -553,11 +667,11 @@ Comparison: .pte vs .pt2 * - libtorch at runtime - Not required - Required - * - Operators TensorRT cannot run - - ExecuTorch CUDA backend, or CPU fallback + * - Operators TensorRT rejects + - ExecuTorch CUDA backend, or the CPU - Compiled by AOTInductor * - Platform - - Linux (x86_64, aarch64) and Windows x86_64, CUDA 13 + - Linux (x86-64 and aarch64) and Windows x86-64, with CUDA 13 - Linux ---- @@ -565,14 +679,14 @@ Comparison: .pte vs .pt2 Examples --------- -* ``examples/torchtrt_executorch_example/export_static_shape.py`` writes a minimal ``.pte``. -* ``examples/torchtrt_executorch_example/export_dynamic_shape.py`` exports a model whose batch - size can change. -* ``examples/torchtrt_executorch_example/export_kv_cache_decode.py`` exports a decoder step - whose key-value cache stays in memory you own. -* ``examples/torchtrt_executorch_example/export_coalesced.py`` splits one graph across the - TensorRT delegate and ExecuTorch's CUDA backend. -* ``examples/torchtrt_executorch_example/export_device_resident.py`` does the same with no - copies at the method boundary. -* ``examples/executorch_reference_runner/`` loads and runs a ``.pte`` from Python and from - C++, including the green-context option. +In ``examples/torchtrt_executorch_example/``: + +* ``export_static_shape.py`` writes a minimal ``.pte``. +* ``export_dynamic_shape.py`` exports a model whose batch size can change. +* ``export_coalesced.py`` splits one graph between the TensorRT delegate and the CUDA backend. +* ``export_device_resident.py`` does the same with no copies at the edges of the program. +* ``export_kv_cache_decode.py`` exports a decoder step whose KV cache stays in memory you own. +* ``export_multi_profile.py`` exports one engine with two optimization profiles. + +``examples/executorch_reference_runner/`` loads and runs a ``.pte`` from Python and from C++, +including optimization profiles and the green context option. diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 800a30f58bd..f4676e03fe5 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -16,13 +16,13 @@ Dynamo IR ------------- The output type of `ir=dynamo` compilation of Torch-TensorRT is `torch.fx.GraphModule` object by default. -We can save this object in either `TorchScript` (`torch.jit.ScriptModule`), `ExportedProgram` (`torch.export.ExportedProgram`) or `PT2` formats by +We can save this object in either `TorchScript` (`torch.jit.ScriptModule`), `ExportedProgram` (`torch.export.ExportedProgram`), `PT2` or ExecuTorch formats by specifying the `output_format` flag. Here are the options `output_format` will accept * `exported_program` : This is the default. We perform transformations on the graphmodule first and use `torch.export.save` to save the module. * `torchscript` : We trace the graphmodule via `torch.jit.trace` and save it via `torch.jit.save`. * `PT2 Format` : This is a next generation runtime for PyTorch models, allowing them to run in Python and in C++ -* `executorch` : We lower the graphmodule to an ExecuTorch ``.pte`` program, delegating the TensorRT engines to the ExecuTorch backend. Runs on Linux (x86-64 and aarch64) and Windows x86-64; see :ref:`executorch_deployment` to install it. +* `executorch` : We lower the graphmodule to an ExecuTorch ``.pte`` program, with the TensorRT engines stored inside it. Requires the ``executorch`` extra; see :ref:`executorch_deployment`. a) ExportedProgram ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -227,24 +227,17 @@ c) ExecuTorch (.pte) .. seealso:: - :ref:`executorch_deployment` is the end to end guide: install the stack, run the - saved program from Python and from C++, and tune the runtime. This section covers - the saving step itself. + :ref:`executorch_deployment` is the end to end guide. It shows how to install the packages, + run the saved program from Python and from C++, and tune it at run time. This section covers + the saving step. -The ``executorch`` output format lowers the compiled module to an ExecuTorch -``.pte`` program, delegating the TensorRT engines to the Torch-TensorRT ExecuTorch -backend. To install it, follow the Installation section of :ref:`executorch_deployment`. +``output_format="executorch"`` writes an ExecuTorch ``.pte`` program. The TensorRT engines are +stored in it, and the TensorRT delegate runs them at run time. There are two ways to write one: -There are two ways to produce a ``.pte``, and they suit different needs: - -* **Use** ``torch_tensorrt.save()``. This is the default, and the right choice - whenever a ``.pte`` file is all that is needed. It runs the whole pipeline in one - call and writes the file. -* **Use** ``torch_tensorrt.executorch.export()`` only when the program has to be - changed before it is written to disk. It stops at the Edge program, hands it back - for inspection or customization, and leaves serialization to the caller. - -Start with ``save()``. Reach for ``export()`` when something below is required. +* ``torch_tensorrt.save()`` does everything in one call and writes the file. Start here. +* ``torch_tensorrt.executorch.export()`` stops before the file is written, and returns + ExecuTorch's Edge program, the form the program has before it becomes a file. Use it to inspect + or change the program first, or to put several methods in one file. Default: ``save()`` """"""""""""""""""" @@ -253,67 +246,102 @@ Default: ``save()`` import torch import torch_tensorrt - import torch_tensorrt.executorch model = MyModel().eval().cuda() inputs = [torch.randn((1, 3, 224, 224)).cuda()] trt_gm = torch_tensorrt.compile(model, ir="dynamo", arg_inputs=inputs) torch_tensorrt.save(trt_gm, "trt.pte", output_format="executorch") -``save`` writes both the ``.pte`` and any external ``.ptd`` tensor-data files. To keep the -TensorRT engines in a ``.ptd`` file instead of the ``.pte``, see -:ref:`executorch_external_engine_data`. +``save()`` writes the ``.pte``, and any ``.ptd`` data files next to it. It writes one method, +``forward``, plus any constant methods you pass. + +These keyword arguments of ``save()`` apply only to this format. With any other format they log a +warning and are ignored. With this format, an unknown keyword raises a ``TypeError``. + +* ``partitioners``: more partitioners, run after TensorRT's. For example, a ``CudaPartitioner`` + sends the operators TensorRT rejects to ExecuTorch's CUDA backend. See + :ref:`executorch_deployment` for that recipe. +* ``compile_specs``: extra compile specs for TensorRT's partitioner. You rarely need it: the + options below set their own specs, and passing one of those specs here raises a ``ValueError``. +* ``backend_config``: an ``ExecutorchBackendConfig`` for ExecuTorch's ``to_executorch()``, for + example to keep inputs and outputs on the GPU. ``save()`` plans its memory on the 256-byte + boundary TensorRT needs, on a copy of the config you pass. +* ``constant_methods``: extra methods that return a constant, such as + ``{"get_max_seq_len": 2048}`` for an LLM runner. +* ``transform_passes``: extra Edge transform passes to run before lowering. +* ``compile_config``: an ``EdgeCompileConfig``. The default turns off Edge IR validation + (``_check_ir_validity=False``), because the TensorRT engine placeholder does not pass it. A + config you pass is used as it is, so set ``_check_ir_validity=False`` on it yourself. +* ``use_cuda_graphs``: save CUDA graph replay on or off for every engine. +* ``weight_streaming_budget_per_engine``: how many bytes of each engine's weights stay on the GPU, + for an engine compiled with ``enable_weight_streaming=True``. +* ``external_engine_data``: a file name, without ``.ptd``, to store the engines in instead of the + ``.pte``. +* ``zero_copy_kv``: let an engine update its KV cache in place. See + :ref:`Zero-copy KV cache `. +* ``generate_etrecord``: also write an ETRecord, a debug record for the ExecuTorch Developer Tools + ``Inspector``, as ``_etrecord.bin``. It does not work for a program with TensorRT + engines: writing the record fails with a ``SerializeError``. + +:ref:`executorch_deployment` explains the run time side of ``use_cuda_graphs``, +``weight_streaming_budget_per_engine`` and ``external_engine_data``. + +For example: + +.. code-block:: python + + from executorch.exir import EdgeCompileConfig + + torch_tensorrt.save( + trt_gm, + "trt.pte", + output_format="executorch", + constant_methods={"get_max_seq_len": 2048}, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) Advanced: ``export()`` """""""""""""""""""""" -Stop at the standard Edge program boundary when the program needs work before it is -serialized: inspecting the delegated graph, adding metadata, controlling final memory -planning, or carrying more than one method. The caller then serializes it: +``torch_tensorrt.executorch.export()`` takes a TensorRT-compiled ``GraphModule``, an +``ExportedProgram`` that already holds engines, or a dict that maps method names to such +programs. It runs TensorRT's partitioner first, then yours in order, and returns ExecuTorch's +``EdgeProgramManager``. You then finalize and write the program yourself. Here +``extra_partitioners``, ``passes``, ``edge_config`` and ``backend_config`` stand for your own +values of the options above: .. code-block:: python + import torch_tensorrt.executorch + edge = torch_tensorrt.executorch.export( trt_gm, - arg_inputs=inputs, - retrace=False, partitioners=extra_partitioners, transform_passes=passes, compile_config=edge_config, constant_methods={"get_vocab_size": 256}, ) - # Inspect edge.exported_program() or apply additional Edge transforms here. + # Inspect edge.exported_program() or apply more Edge transforms here. program = edge.to_executorch(config=backend_config) with open("trt.pte", "wb") as output: program.write_to_file(output) program.write_tensor_data_to_file(".") +It takes the same options as ``save()``, except ``backend_config``, which goes to +``to_executorch()``. ``partitioners``, ``compile_specs`` and ``transform_passes`` can also be a dict +keyed by method name. + TensorRT needs every input and output on a 256-byte boundary, and ExecuTorch plans memory on -16-byte boundaries by default. ``save`` plans on 256 for you. Here, give ``backend_config`` a +16-byte boundaries by default. ``save()`` plans on 256 for you. Here, give ``backend_config`` a ``MemoryPlanningPass(alignment=256)``, or the delegate copies each misaligned buffer on every call. -``executorch.export`` accepts a TensorRT-compiled ``GraphModule``, an -engine-bearing ``ExportedProgram``, or a mapping of independently exported -methods. It always applies the TensorRT partitioner first, followed by caller -``partitioners`` in order, and returns ExecuTorch's native -``EdgeProgramManager``. Use ``save`` for the one-shot path where Torch-TensorRT -manages Edge lowering, finalization, and persistence. - **Several methods in one .pte** -Pass a mapping of method name to ``ExportedProgram`` to keep independent entry -points, such as a separate prefill and decode, in one program. Give each method -its own partitioner instances when a partitioner carries method-specific state: a -partitioner holds its compile specs from construction, so one instance whose specs -name a method would tag every method sharing it with that same name. Reusing such -an instance across methods is rejected. - -Sharing an instance whose specs name no method is not rejected, because some -backends are built to share one. A backend that instead reads its own method name -from its specs, such as the CUDA backend below, raises -``Could not find method name in compile specs`` during lowering, so give it one -instance per method carrying that method's name spec. +Pass a dict of method name to ``ExportedProgram`` to keep separate entry points, such as prefill +and decode, in one program. Give each method its own partitioner instances. The CUDA backend +reads its method name from its compile specs, so it needs one ``CudaPartitioner`` per method, +each with that method's name: .. code-block:: python @@ -324,292 +352,74 @@ instance per method carrying that method's name spec. {"prefill": prefill_program, "decode": decode_program}, partitioners={ "prefill": [ - CudaPartitioner( - [CudaBackend.generate_method_name_compile_spec("prefill")] - ) + CudaPartitioner([CudaBackend.generate_method_name_compile_spec("prefill")]) ], "decode": [ - CudaPartitioner( - [CudaBackend.generate_method_name_compile_spec("decode")] - ) + CudaPartitioner([CudaBackend.generate_method_name_compile_spec("decode")]) ], }, ) -Each method is validated before any program is rewritten, so an error in one -method leaves the others untouched. A method mapping preserves independent entry -points but does not by itself give them shared mutable state. +``export()`` refuses one partitioner instance shared by methods when its specs name a method. The +methods share no state: two methods do not share a buffer unless you plan it that way. .. warning:: - **The returned Edge program shares tensor storage with the programs you - passed in.** Only structure is copied: the graph, the graph signature, the - ``state_dict`` keys, and node metadata. Weights and every other tensor - payload are shared by reference. - - TensorRT engines are not shared. Export never deep copies an engine object, - since that would serialize and deserialize it, but every engine is decoded - once into a byte buffer the returned program owns. Plan for the bytes of a - multi-gigabyte engine to be resident twice while both programs are alive. + **The returned Edge program shares its weights with the programs you passed in.** Do not + change a weight in place in a transform pass, for example with ``weight.data.mul_(scale)``, + and do not change a source program after ``export()``. Either change shows up in the other + program too, with no error. Build a new tensor instead. - Two consequences to plan for: - - * A transform pass must not modify a shared payload in place. An in-place - edit such as ``weight.data.mul_(scale)`` changes the program you passed in, - and every other Edge program exported from it. Build a new tensor and - rebind it instead of mutating the existing one. - * Modifying a source program after calling ``export`` is also visible in the - Edge program. Finish preparing a program before exporting it. - - Neither case raises an error or a warning, so treat every shared payload as - read-only. + Engines are not shared. ``export()`` copies each engine into the Edge program once, so a large + engine is in host memory twice while both programs exist. .. _executorch_zero_copy_kv: **Zero-copy KV cache** -When a TensorRT engine has aliased I/O -- a KV cache it updates through an -aliased binding -- running the engine over the cache already is the update. -ExecuTorch does not know that, so by default it pays for the update twice per -execution: it hands the delegate an ``_h2d_copy`` staging copy of the buffer -instead of the buffer itself, then copies the engine's aliased output back into -the buffer afterwards. For a KV cache both copies are cache-sized, per token. - -``zero_copy_kv=True`` removes them, so the engine writes the caller's buffer -directly. Through the two-step ``export`` + ``to_executorch`` path it takes two -calls, one at each end of the Edge boundary: - -.. code-block:: python - - from torch_tensorrt.executorch import ( - check_zero_copy_kv, - export, - zero_copy_backend_config, - ) - - edge = export( - {"prefill": prefill_program, "decode": decode_program}, - zero_copy_kv=True, - ) - - # The argument is optional. Pass your own ExecutorchBackendConfig and - # zero_copy_backend_config composes onto it; every other field (memory - # planning, passes) is preserved, so give it MemoryPlanningPass(alignment=256) - # to keep the cache on the boundary TensorRT needs. - program = edge.to_executorch(zero_copy_backend_config()) - - # Neither call is enforced -- see the warning below. - check_zero_copy_kv(program) - -Two fields of that config are not merely carried. The engine writes the cache -wherever memory planning put it, so ``enable_non_cpu_memory_planning=False`` -- -which plans every tensor into a single host arena whatever device its -``TensorSpec`` asks for -- cannot be combined with zero-copy KV: -``to_executorch`` raises instead of writing a ``.pte`` whose every ``execute()`` -fails on a host pointer. That refusal reads the config object -``zero_copy_backend_config`` returned, so two ways of setting the field are -outside it. Building a new config out of the returned one with -``dataclasses.replace`` copies the field by value and the pass by reference, so -the pass goes on reading the config it was built for: call -``zero_copy_backend_config`` again on the derived config. And a -``memory_planning_pass`` of your own that does not already have an attribute of -that name never receives the field at all -- ``to_executorch`` assigns it onto -the planner rather than passing it -- so where the caches land is that planner's -own business. Neither has to be left to the runtime to discover: -``check_zero_copy_kv`` reads the arena planning actually chose -- on this path -it is yours to call, as the warning below says. That planner does have to record -the device of each arena it places, or the check refuses the program for saying -nothing about where the cache lives -- see below. And -``propagate_device_config.skip_h2d_for_method_inputs`` is refused outright: -``PropagateDevicePass`` refuses to un-stage a method input -whose placeholder does not have exactly one user, and a zero-copy cache always -has two, so preserving that option would hand back a configuration that cannot -finalize at all. It is refused wherever it is written -- in one -``PropagateDeviceConfig`` or in a per-method dict of them -- and on every value -that pass reads as on rather than only on ``True``, because it tests the field -for truth without ever resolving it per method, so even a dict of ``False`` is -on for every method. ``False`` and the empty dict are what it reads as off, and -those are carried unchanged. - -It is opt-in rather than automatic because the resulting ``.pte`` needs a -runtime that understands a delegate whose aliased outputs are elided. Producing -one silently would break a runner built before this feature. - -.. warning:: - - **Both calls are required, and nothing enforces them.** Exporting with - ``zero_copy_kv=True`` removes the copy-back; finalizing without - ``zero_copy_backend_config`` leaves the engine writing a per-call staging - copy that is discarded, so the cache never updates. For a KV cache that is - wrong output, not a crash -- and it is a ``.pte`` that was written without - complaint, since ``to_executorch`` belongs to ExecuTorch and finalizes - whatever config it is given. - - So finalizing a zero-copy export yourself is two obligations, on any manager - and by whatever route it was reached: pass ``zero_copy_backend_config``, and - hand the program that comes back to - ``torch_tensorrt.executorch.check_zero_copy_kv``, which reads the finalized - program and refuses one whose caches never reach the engine or ended up - planned somewhere it cannot write them -- while there is still no ``.pte``:: +An LLM's TensorRT engine can update its KV cache in place, through an input and an output that +alias the same buffer. By default ExecuTorch does not know that. It hands the engine a copy of +the cache, and then copies the engine's result back: two cache-sized copies for every token. +``zero_copy_kv=True`` removes both, so the engine writes the cache directly. - program = edge.transform(passes).to_executorch(zero_copy_backend_config()) - torch_tensorrt.executorch.check_zero_copy_kv(program) - - ``torch_tensorrt.save`` owns both ends and does both for you; it is the path - on which forgetting is not possible, and it is single-method only -- a - multi-method program has to take the two calls above. - -``torch_tensorrt.save`` finalizes the program itself, so a single -``zero_copy_kv=True`` covers both steps: +With ``save()``, one flag does it: .. code-block:: python - torch_tensorrt.save( - trt_gm, "decode.pte", output_format="executorch", - zero_copy_kv=True, - ) + torch_tensorrt.save(trt_gm, "decode.pte", output_format="executorch", zero_copy_kv=True) -It installs ``zero_copy_backend_config`` for you, so there is no need to hand it -one as ``backend_config`` as well: that installs the pass twice, which is -redundant rather than an error -- the second run finds the buffers already -un-staged. The two entry points are alternatives, not a pair. - -Three further responsibilities are the caller's. The first two nothing checks; -the third is refused where it is visible in the ``.pte``: - -* **One CUDA stream for every delegate**, if the ``.pte`` is coalesced. Getting - this wrong is a race, not a deterministic error: it is intermittent and can - surface as wrong results *or* as an illegal memory access. See - :ref:`Running a coalesced .pte `. - -* **Synchronizing that stream before reading a cache on the host**, if the ``.pte`` is - coalesced and ExecuTorch's CUDA backend runs after the engine. The TensorRT delegate waits - for its engine before it returns, but the CUDA backend may return early. See - :ref:`Running a coalesced .pte `. - -* **Sharing one cache between methods.** Zero-copy is per method: it makes each - method's engine write that method's buffer. Giving a prefill and a decode - method *the same* cache is a memory-planning question -- their mutable buffers - have to land at the same arena offsets -- which ExecuTorch's memory planner - owns and which is deployment-specific. Supply your own - ``memory_planning_pass`` for it; ``zero_copy_backend_config`` preserves it. - Such a planner has to record the device of each arena it places, which means - running ExecuTorch's ``apply_algo`` -- the only thing that writes - ``non_const_buffer_device`` -- and running it with - ``enable_non_cpu_memory_planning=True``, since that parameter defaults to - ``False`` and with it off ``apply_algo`` plans every spec into one CPU bucket - and records nothing either. Without that record the ``.pte`` reports every - planned buffer as CPU, and a runner that honours it backs the cache with host - memory the engine cannot write. ``check_zero_copy_kv`` refuses a program in - that state, under a message of its own rather than the host-arena one: what - it can read is that nothing says where the cache lives, not that your planner - put it among the host tensors. - -**Coalesced TensorRT + CUDA .pte** - -To run the ops TensorRT does not take on ExecuTorch's CUDA (AOTInductor) backend -instead of leaving them non-delegated, pass a ``CudaPartitioner`` via -``partitioners=``. It is appended after the TensorRT partitioner, so TensorRT -claims what it can and the ``CudaPartitioner`` picks up the rest as a catch-all: +With ``export()``, it takes two more calls, one at each end of the Edge program: .. code-block:: python - from executorch.backends.cuda.cuda_backend import CudaBackend - from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + from torch_tensorrt.executorch import check_zero_copy_kv, export, zero_copy_backend_config - torch_tensorrt.save( - trt_gm, "trt.pte", output_format="executorch", - partitioners=[ - CudaPartitioner( - [CudaBackend.generate_method_name_compile_spec("forward")] - ) - ], - ) - -This needs ExecuTorch's CUDA backend and a CUDA toolkit (nvcc/ptxas) at export -time, and produces a ``.pte`` that requires a CUDA runtime at load. Any external -CUDA weights are written as ``.ptd`` data file(s) next to the ``.pte``; the runtime -must be pointed at those data files to load them. + edge = export({"prefill": prefill_program, "decode": decode_program}, zero_copy_kv=True) + program = edge.to_executorch(zero_copy_backend_config()) + check_zero_copy_kv(program) .. warning:: - The CUDA backend names its external weight file per device (for example - ``aoti_cuda_blob.ptd``), not per model, so saving two coalesced programs into - one directory overwrites the first one's weights. The first program can then - return wrong answers with no error, or fail to load. Save each coalesced model - into its own directory. - -.. _executorch_single_stream: - -**Running a coalesced .pte on one CUDA stream** - -A coalesced ``.pte`` runs on two delegates: the TensorRT delegate and ExecuTorch's CUDA -backend. With no stream chosen, both use ``cudaStreamPerThread``, the calling thread's default -stream, so a program run from one thread is already in the right order. - -To use a stream of your own, put one ``executorch::extension::cuda::CallerStreamGuard`` around -the whole call. It reaches both delegates, because they share one ``libextension_cuda``. Do not -give the two delegates different streams: work on separate streams is not ordered, and the -result is a race that can show up as wrong values or an illegal memory access. - -The TensorRT delegate waits for its own work before it returns. ExecuTorch's CUDA backend can -return while its work is still queued, so when it runs last, ``execute()`` can return before the -outputs are ready. Synchronize the stream before you read an output on the host, for example -between the steps of a decode loop. - -**ExecuTorch lowering options** - -``torch_tensorrt.save`` takes these extra keyword arguments. They are only -consulted for the ``executorch`` format; passing them with any other -``output_format`` logs a warning and is otherwise ignored. Of the six below, -``constant_methods``, ``transform_passes`` and ``compile_config`` are forwarded -to ExecuTorch's ``to_edge_transform_and_lower(...)``, and so is -``generate_etrecord``, which ``save`` also reads itself to write the record -beside the ``.pte``. ``backend_config`` goes to ``to_executorch(...)`` instead, -and ``zero_copy_kv`` is read by ``save`` on both sides of that boundary. - -* ``constant_methods`` — a ``dict`` of extra constant methods to embed in the - ``.pte`` (e.g. ``{"get_max_seq_len": 2048}`` for an LLM runner). -* ``transform_passes`` — additional edge-dialect transform passes to run before - lowering. -* ``compile_config`` — an ``EdgeCompileConfig``. When omitted, Torch-TensorRT - supplies a default with ``_check_ir_validity=False`` (the TensorRT - ``execute_engine`` placeholder graph does not pass edge-IR validation). A - caller-supplied config is forwarded **verbatim**, so if you pass your own and - your graph carries TensorRT engines, set ``_check_ir_validity=False`` explicitly. -* ``backend_config`` — an ``ExecutorchBackendConfig`` forwarded to - ``to_executorch(...)``, with its ``MemoryPlanningPass`` aligned to 256 bytes on a copy. -* ``zero_copy_kv`` — a ``bool`` (default ``False``, single-method only). Lets the - TensorRT engine update an aliased KV cache in place. ``save`` owns both ends of - the Edge boundary, so this one argument covers what the two-call path spells - out; see :ref:`Zero-copy KV cache `. -* ``generate_etrecord`` — a ``bool`` (default ``False``). When ``True``, an - `ETRecord `_ is written - next to the ``.pte`` as ``_etrecord.bin`` (e.g. ``trt.pte`` → - ``trt_etrecord.bin``) for use with the ExecuTorch Developer Tools ``Inspector``. - -.. code-block:: python - - from executorch.exir import EdgeCompileConfig - - torch_tensorrt.save( - trt_gm, "trt.pte", output_format="executorch", - constant_methods={"get_max_seq_len": 2048}, - compile_config=EdgeCompileConfig(_check_ir_validity=False), - generate_etrecord=True, - ) - -The ETRecord sidecar can be parsed back and paired with a runtime ETDump in the -Developer Tools ``Inspector``: - -.. code-block:: python + **Nothing makes you call both.** If you finalize with a config that did not come from + ``zero_copy_backend_config``, ``to_executorch()`` still writes a program, but the engine then + writes a copy that is thrown away, so the cache never updates. That is wrong output, not a + crash. ``check_zero_copy_kv`` refuses such a program before you write the file. + ``save(..., zero_copy_kv=True)`` does both steps for you. - from executorch.devtools import Inspector - from executorch.devtools.etrecord import parse_etrecord +* ``zero_copy_backend_config(config)`` starts from your own config, if you pass one, and keeps + every field of it, so give that config ``MemoryPlanningPass(alignment=256)`` as above. It refuses ``skip_h2d_for_method_inputs``. The caches have to be planned in + GPU memory, so ``enable_non_cpu_memory_planning=False`` makes finalizing fail. +* If you build a new config from the returned one, for example with ``dataclasses.replace``, + call ``zero_copy_backend_config`` again on the new one. +* It is opt-in because the ``.pte`` needs a runtime that knows about it. A runner built before + this feature would break on it. +* To share one cache between two methods, such as prefill and decode, both methods' buffers have + to be planned at the same place. Pass your own ``memory_planning_pass`` for that. + ``zero_copy_backend_config`` keeps it. That planner must record which device each memory area + is on, which ExecuTorch's ``apply_algo`` does when it runs with + ``enable_non_cpu_memory_planning=True``. Otherwise ``check_zero_copy_kv`` refuses the program. - etrecord = parse_etrecord("trt_etrecord.bin") - inspector = Inspector(etdump_path="etdump.etdp", etrecord=etrecord) - inspector.print_data_tabular() +The :ref:`API reference ` describes every refusal in detail. Saving torch.compile models @@ -619,7 +429,7 @@ Saving torch.compile models lazily on the first call and live in memory only. There is no direct way to call ``torch_tensorrt.save()`` on a ``torch.compile``-compiled model. -**Workaround: switch to the AOT path** +**Use the AOT path instead** Replace ``torch.compile`` with ``torch_tensorrt.compile(ir="dynamo")`` to get a serializable ``torch.fx.GraphModule``: diff --git a/examples/executorch_reference_runner/README.md b/examples/executorch_reference_runner/README.md index c6eb1b709e0..e89d93f4e57 100644 --- a/examples/executorch_reference_runner/README.md +++ b/examples/executorch_reference_runner/README.md @@ -1,282 +1,206 @@ -# Torch-TensorRT ExecuTorch Reference Runner +# Torch-TensorRT ExecuTorch Reference Runners -This directory contains minimal C++ reference runners for loading and -executing Torch-TensorRT compiled models saved in ExecuTorch `.pte` format: +Small example programs that load and run a `.pte` file written by Torch-TensorRT. They are +reference code with fixed example inputs, so replace those with your real inputs. -| Target | Shows | -| ------ | ----- | -| `example_executorch_runner` (`main.cpp`) | The low-level `Program` / `Method` loading sequence | -| `example_executorch_multi_profile_runner` (`multi_profile_main.cpp`) | Selecting a TensorRT optimization profile per call through the high-level `Module` API | -| `example_executorch_multi_profile_benchmark` (`multi_profile_benchmark.cpp`) | What per-call profile switching costs | +| Program | Source | Shows | +| --- | --- | --- | +| `example_executorch_runner` | `main.cpp` | The low-level `Program` and `Method` loading sequence, and running on a green context stream | +| `kv_cache_decode_check` | `kv_cache_decode_check.cpp` | That a key-value (KV) cache in memory the caller owns keeps its updates from call to call | +| `example_executorch_multi_profile_runner` | `multi_profile_main.cpp` | Picking a TensorRT optimization profile for each call through the `Module` API | +| `example_executorch_multi_profile_benchmark` | `multi_profile_benchmark.cpp` | What picking the right profile is worth | +| `load_model.py` | Python | Loading and running a `.pte` with ExecuTorch's Python runtime | +| `load_model_device_resident.py` | Python | The same for a program that keeps its inputs and outputs on the GPU | -The `.pte` file contains an ExecuTorch program with embedded TensorRT engine -payloads. The runner links the TensorRT ExecuTorch backend, loads the `.pte` -with the ExecuTorch C++ runtime, prepares input tensors, calls `execute()`, and -prints output shapes and sample values. +The [ExecuTorch deployment guide](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html) +explains the concepts used here. -This is reference code. It fills all inputs with `1.0f`; replace that input -setup with your application's real input buffers. +## Get a model -## Input Model - -You can also generate a sample `.pte` from the Torch-TensorRT source tree: +The export scripts are in `examples/torchtrt_executorch_example/` of the Torch-TensorRT +repository. They need Torch-TensorRT with its `executorch` extra, installed as the guide's +[Installation](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#installation) +section shows. ```bash +# A minimal program, for example_executorch_runner and load_model.py. python examples/torchtrt_executorch_example/export_static_shape.py --model_path=model.pte -# Two-profile Gemma-3 engine for the multi-profile runner below. Needs -# transformers (pip install "transformers>=5.4.0") for the Gemma-3 definition; -# the executorch extra does not pull it in. Defaults to a mini Gemma-3 that -# needs no download and exports in about a minute, most of it spent serializing -# the engine into the .pte. Add --weights google/gemma-3-1b-it for the real 1B -# model -- but that .pte is about 2 GB and serialization scales with engine -# size, so budget hours rather than minutes for it. The export script documents -# the measured rate. -python examples/torchtrt_executorch_example/export_multi_profile.py --model_path=model_gemma3_multi_profile.pte -``` - -## Build The Reference Runner - -A normal reference runner build does not need separate steps for -`libexecutorch_core.a` and `libexecutorch_trt_backend.a`. The runner CMake adds -both ExecuTorch and the Torch-TensorRT ExecuTorch source package, and linking -`executorch::backend_tensorrt` makes the backend archive a dependency of -`example_executorch_runner`. - -The build also turns on ExecuTorch's CUDA backend, so the CUDA toolkit has to be -installed. That is needed even for a program that uses only the TensorRT delegate. -A Torch-TensorRT export marks its delegate inputs and outputs as CUDA memory, so -ExecuTorch places device copies around the delegate, and the allocator those copies -use is registered by the CUDA backend. Build without it and the runner loads the -program, initializes the TensorRT engine, then stops on the first instruction with -`_h2d_copy: no device allocator registered for device_type=1`. +# A program split between TensorRT and ExecuTorch's CUDA backend, for the green context example. +python examples/torchtrt_executorch_example/export_coalesced.py --model_path=coalesced.pte -The `libtorchtrt.tar.gz` package also includes a prebuilt reference runner: +# A one-layer decoder step, for kv_cache_decode_check. +python examples/torchtrt_executorch_example/export_kv_cache_decode.py --model_path=kv_cache_decode.pte -```text -torch_tensorrt/bin/example_executorch_runner +# One engine with two optimization profiles, for the two multi-profile programs. +python examples/torchtrt_executorch_example/export_multi_profile.py --model_path=model_gemma3_multi_profile.pte ``` +`export_coalesced.py` and `export_device_resident.py` also need a CUDA toolkit, because the CUDA +backend compiles its part with `nvcc`. Both write `aoti_cuda_blob.ptd` next to their `.pte`. Here +that is harmless, because neither CUDA part has weights, but give your own models one directory +each, as the guide explains. `export_multi_profile.py` needs `transformers>=5.4.0`, +which the extra does not install. By default it exports a small Gemma-3 with random weights, which +needs no download and takes about a minute. With `--weights google/gemma-3-1b-it` it exports the +real 1B model, whose `.pte` is about 2 GB and takes hours to write. That model is gated on Hugging +Face, so accept its license and log in first. + +## Build the C++ programs + +You need: + +- A CUDA toolkit. The build turns on ExecuTorch's CUDA backend, which this build needs even for a + program that uses only TensorRT: here the backend registers the allocator for the program's GPU + memory. +- PyTorch and PyYAML, installed in the Python environment that CMake finds. ExecuTorch's kernel + code generator uses both, and the CUDA backend needs PyTorch to configure, but the programs do + not link libtorch. +- A TensorRT SDK. Point `TensorRT_ROOT` at it. +- `libtorchtrt.tar.gz` from your Torch-TensorRT release, which carries the delegate source and + these examples. +- An ExecuTorch source checkout at the commit that your Torch-TensorRT release pins in its + `MODULE.bazel`. The tarball does not carry that file, so look it up in the release's source on + GitHub. The default below is the commit pinned in the same source tree as this README. + ```bash -# Get the ExecuTorch source snapshot this package is built against. Keep this in sync -# with the executorch commit pinned in MODULE.bazel. +# The ExecuTorch commit pinned in MODULE.bazel. EXECUTORCH_REF="${EXECUTORCH_REF:-a9665e25097b324a7116efa2bd6602287c9c064f}" -git clone --filter=blob:none --no-checkout \ - https://github.com/pytorch/executorch.git executorch +git clone --filter=blob:none --no-checkout https://github.com/pytorch/executorch.git executorch pushd executorch git fetch --depth 1 origin "${EXECUTORCH_REF}" git checkout FETCH_HEAD git submodule update --init --recursive --depth 1 popd -# download the libtorchtrt.tar.gz tar xvf libtorchtrt.tar.gz -export EXECUTORCH_SOURCE_DIR="${PWD}/executorch" -# tarball untared path export TORCH_TENSORRT_ROOT="${PWD}/torch_tensorrt" -export EXECUTORCH_BACKEND_TENSORRT_SOURCE_DIR="${TORCH_TENSORRT_ROOT}/src/torch_tensorrt/executorch" export TensorRT_ROOT=/path/to/extracted/TensorRT export LD_LIBRARY_PATH="${TensorRT_ROOT}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" cmake -S "${TORCH_TENSORRT_ROOT}/examples/executorch_reference_runner" \ -B build-executorch-reference-runner \ - -DEXECUTORCH_SOURCE_DIR="${EXECUTORCH_SOURCE_DIR}" \ - -DEXECUTORCH_BACKEND_TENSORRT_SOURCE_DIR="${EXECUTORCH_BACKEND_TENSORRT_SOURCE_DIR}" \ + -DEXECUTORCH_SOURCE_DIR="${PWD}/executorch" \ + -DEXECUTORCH_BACKEND_TENSORRT_SOURCE_DIR="${TORCH_TENSORRT_ROOT}/src/torch_tensorrt/executorch" \ -DTensorRT_ROOT="${TensorRT_ROOT}" -cmake --build build-executorch-reference-runner --target example_executorch_runner -j -``` - -This packaged reference-runner flow currently targets Linux x86_64 and SBSA, -matching the native backend's Bazel compatibility constraints. Windows and -JetPack packaging are outside this flow. - -Expected artifact: - -```text -build-executorch-reference-runner/example_executorch_runner -``` - -The build also creates the ExecuTorch core, the shared caller-stream extension, -and the TensorRT backend archive as dependencies: - -```text -build-executorch-reference-runner/executorch/libexecutorch_core.a -build-executorch-reference-runner/torch_tensorrt_executorch/executorch_extension_cuda/libextension_cuda.so -build-executorch-reference-runner/lib/libexecutorch_trt_backend.a -``` - -## Load And Run A `.pte` Model - -### Python - -In a fresh Linux CUDA 13 environment, install the `executorch` authoring -stack through the `[executorch]` extra. Substitute the channel for your CUDA, such -as `cu134` for CUDA 13.4, keeping PyTorch, ExecuTorch and Torch-TensorRT on the -same channel: - -```bash -pip install --pre "torch-tensorrt[executorch]" \ - --extra-index-url https://download.pytorch.org/whl/nightly/cu132 \ - --extra-index-url https://pypi.nvidia.com +cmake --build build-executorch-reference-runner -j \ + --target example_executorch_runner kv_cache_decode_check ``` -Both indexes are required. The extra's ExecuTorch floor names a development build the public index -does not carry, and that index ships the inference library only as source, which pip spends about -twenty minutes failing to build. -`--pre` allows prereleases; it does not request an upgrade. Released versions can -already declare the extra. If an older installation lacks it, first install the -intended compatible Torch-TensorRT wheel deliberately. Adding the extra may change -dependencies, so use a fresh environment to preserve an existing working stack. +This builds the runner and the KV cache check into `build-executorch-reference-runner/`, and +the delegate's static archive, `build-executorch-reference-runner/lib/libexecutorch_trt_backend.a`, +as a dependency. Naming the targets keeps the build from compiling every default ExecuTorch target +too. -The extra also installs the companion wheel, `torch-tensorrt-executorch-runtime`, so there is no -second command. That wheel ships just the TensorRT delegate, a single shared library that registers -itself with the ExecuTorch runtime from the `executorch` distribution rather than bundling a runtime -of its own, and loading a `.pte` through the delegate needs it. To build it from source instead, -follow `py/torch-tensorrt-executorch-runtime/README.md`. +For the two multi-profile programs, add +`example_executorch_multi_profile_runner` and `example_executorch_multi_profile_benchmark` to the +targets. To build from a source checkout instead of the tarball, point `TORCH_TENSORRT_ROOT` at the +checkout and leave out `-DEXECUTORCH_BACKEND_TENSORRT_SOURCE_DIR`: the same `CMakeLists.txt` then +finds the delegate source on its own. -The Python example uses ExecuTorch's Runtime API to back planned device arenas -with CUDA memory. Then load and run the model: +This flow covers Linux x86-64 and Arm servers (SBSA). Windows and JetPack are not part of it. +Except in the JetPack and TensorRT-RTX builds, `libtorchtrt.tar.gz` and the Torch-TensorRT wheels +for Linux also carry a prebuilt `torch_tensorrt/bin/example_executorch_runner`. It registers the +GPU allocator itself. It has no CUDA backend, and its only CPU operators are the copies at the edges +of the program. So it runs only programs that TensorRT takes in full, such as the one +`export_static_shape.py` writes. To find the copy in your Python environment: ```bash -python examples/executorch_reference_runner/load_model.py \ - --model_path=model.pte \ - --num_runs=1 +python -c 'import importlib.util, pathlib; print(pathlib.Path(importlib.util.find_spec("torch_tensorrt").origin).parent / "bin")' ``` -The legacy `torch_tensorrt.load(path, format="executorch")` entry point still -works, but emits a deprecation warning. Its `method_names` property, -`run(inputs, method="forward")`, and `forward(*inputs)` interface remain -supported for at least six months after the deprecation first ships. It still -copies CUDA inputs to CPU and supports embedded weights only. New applications -should use the Runtime API shown above; device-resident programs must use it -directly to keep their inputs on CUDA. - -### C++ - -Run the reference runner against a Torch-TensorRT compiled ExecuTorch model: +## Run a program from C++ ```bash -./build-executorch-reference-runner/example_executorch_runner --model_path=/path/to/model.pte --num_runs=1 +./build-executorch-reference-runner/example_executorch_runner --model_path=model.pte --num_runs=1 ``` -The runner demonstrates this C++ loading sequence: +`main.cpp` shows the low-level loading sequence: ```text executorch::runtime::runtime_init() FileDataLoader::from(model_path) Program::load(loader) program.method_meta(method_name) -allocate planned ExecuTorch memory +allocate the method's planned memory program.load_method(method_name, memory_manager) method.set_input(...) method.execute() method.get_outputs(...) ``` -Loading the method initializes the TensorRT ExecuTorch backend for any -Torch-TensorRT delegate subgraphs embedded in the `.pte`. Applications can -scope `executorch::extension::cuda::CallerStreamGuard` around execution to run -TensorRT and CUDA/AOTI delegates on one ordinary caller-owned CUDA stream. On the -discrete-GPU CI configuration, this runner's host-backed inputs and outputs take -the synchronized staging path, and that is the path a discrete card always takes: it -can read pageable host memory, but by faulting pages in one at a time rather than -through shared page tables, which is far slower than one bulk copy. Only an integrated -part that shares host page tables has its host storage bound directly. Either way the -call returns with the work finished. -The Python `torch_tensorrt` package is needed when exporting the `.pte`; it is not -needed by this native runner at inference time. - -## Running both delegates on one green-context stream - -The point of the shared caller stream is that TensorRT and the CUDA/AOTI delegate -can be driven by one caller-owned stream, including a CUDA green-context stream, so -both stay inside the same SM partition. To exercise that, build with the CUDA -delegate enabled and pass `--green_context_sms`: +Loading the method initializes the TensorRT delegate for each engine in the `.pte`. The runner does +not need the Python `torch_tensorrt` package, which is only needed to export. It takes no data file, +so it runs a program split with the CUDA backend only when that part has no weights, like the one +`export_coalesced.py` writes, and it cannot run a program saved with `external_engine_data`. -```bash -cmake -S "${TORCH_TENSORRT_ROOT}/examples/executorch_reference_runner" \ - -B build-executorch-reference-runner \ - -DEXECUTORCH_BUILD_CUDA=ON \ - -DEXECUTORCH_SOURCE_DIR="${EXECUTORCH_SOURCE_DIR}" \ - -DTensorRT_ROOT="${TensorRT_ROOT}" -cmake --build build-executorch-reference-runner --target example_executorch_runner -j +### Run on a green context stream -build-executorch-reference-runner/example_executorch_runner \ - --model_path=two_delegate.model.pte --num_runs=1 --green_context_sms=8 -``` +A CUDA green context owns a fixed number of the GPU's SMs (streaming multiprocessors). Running the +program on a stream of one keeps the whole model inside that part of the GPU. Both the TensorRT +delegate and ExecuTorch's CUDA backend honor it, so test it with the coalesced program: -`--green_context_sms=N` creates the caller stream inside a green context holding N -SMs. The runner refuses with a distinct status rather than falling back if one -cannot be created, and says how many the device has, so a -passing run always means a green context was really used. `N=0`, the default, uses -an ordinary stream. +```bash +./build-executorch-reference-runner/example_executorch_runner \ + --model_path=coalesced.pte --num_runs=1 --green_context_sms=8 +``` -Enabling `EXECUTORCH_BUILD_CUDA` does not make this runner depend on libtorch. It -needs `EXECUTORCH_BUILD_EXTENSION_TENSOR=ON`, which is set automatically, and the -result links no libtorch and no libc10. +`--green_context_sms=N` creates the stream inside a green context with at least N SMs. If the +device cannot provide that, the runner stops with exit status 2 and says how many SMs the device +has, instead of quietly using an ordinary stream. So a passing run really used a green context. +`0`, the default, uses an ordinary stream. -CI exercises the green-context option. The verification script runs the coalesced -program twice, once on an ordinary stream and once with `--green_context_sms=8`, and -treats the runner's distinct status for "no green context available" as a skip, so the -green run happens only where the device has the SMs for one. The other reference-runner -checks use the CUDA-enabled build with an ordinary stream and host-backed method inputs -and outputs. +TensorRT warns on such a run that the engine was built without green context constraints, which it +says may cause crashes or undefined behavior. Torch-TensorRT has no option to build engines with +those constraints. -## Caller-Owned KV-Cache Persistence Check +## Run a program from Python -`kv_cache_decode_check` is a small self-asserting runner for a caller-owned -KV-cache decode `.pte` (its aliased KV output is bound in place to the caller's -mutable buffer, which persists across `execute()` calls). +```bash +python examples/executorch_reference_runner/load_model.py --model_path=model.pte --num_runs=1 +``` -Export a minimal single-layer decode model: +For a program exported to keep its inputs and outputs on the GPU, `load_model_device_resident.py` +passes a CUDA input and checks that a CUDA output comes back with the right values: ```bash -python examples/torchtrt_executorch_example/export_kv_cache_decode.py \ - --model_path=kv_cache_decode.pte +python examples/torchtrt_executorch_example/export_device_resident.py --model_path=device_resident.pte +python examples/executorch_reference_runner/load_model_device_resident.py --model_path=device_resident.pte ``` -The same CMake build produces the check runner (`kv_cache_decode_check` -target). Run it: +## Check a KV cache the caller owns + +`kv_cache_decode_check` tests a decoder step whose KV cache stays in a buffer the caller owns. The +engine writes the cache in place, and the writes have to persist from one `execute()` call to the +next. ```bash ./build-executorch-reference-runner/kv_cache_decode_check --model_path=kv_cache_decode.pte ``` -It loads the method twice, explicitly zeroing the caller-owned KV arenas on each -load (ExecuTorch drops a mutated buffer's initial value, so the runner must zero -them itself), and runs a decode at `input_pos=1` once with no prior step and once -after a step at `input_pos=0`. -Because the causal attention at position 1 covers positions 0..1, the two logits -differ only if the KV written at position 0 persisted across `execute()` calls. -The runner prints `[kv-check] PASS` and returns 0 on success, or fails if the -two are identical (the update did not persist). It requires a CUDA device. - -That pair runs twice over, printing `caller stream: none` and then `caller -stream: own`. The second scopes a `CallerStreamGuard` over the decode loop on a -stream the runner creates, so the engine runs on that stream; the first leaves the -caller stream unset, so it runs on `cudaStreamPerThread` instead. Both are streams -a caller can put the backend on, and a `.pte` whose KV writes persist on one but -not the other is still broken, so both have to pass. `execute()` waits for the -engine before returning in either case. - -## Selecting An Optimization Profile - -A TensorRT engine can hold several optimization profiles: one weight set, one -engine, several kernel tunings, each valid over a different input-shape range. -Scope an `OptimizationProfileGuard` around the call to pick one. - -A profile is identified by its index in the list declared at export time. The -library defines no index constants; name them yourself to match the exporter, as -`export_multi_profile.py` declares decode first and prefill second: +It runs a decode step at position 1 twice: once on a fresh cache, and once after a step at +position 0. Attention at position 1 reads positions 0 and 1, so the two results differ only if +the cache write at position 0 persisted. The program zeroes the cache itself on each load, because +ExecuTorch does not keep the initial value of a buffer the program changes. It prints +`[kv-check] PASS` and returns 0 when the results differ, and fails when they are the same. + +It does all of this twice, printing `caller stream: none` and then `caller stream: own`. The first +pass runs on `cudaStreamPerThread`, and the second on a stream the program creates and selects with +`CallerStreamGuard`. Both have to pass. It needs a CUDA device. + +## Pick an optimization profile + +An engine can hold several optimization profiles: one set of weights, tuned for several input +shape ranges. Scope an `OptimizationProfileGuard` around a call to pick one: ```cpp +#include #include using torch_tensorrt::executorch_backend::OptimizationProfileGuard; -constexpr int32_t kDecodeProfile = 0; +constexpr int32_t kDecodeProfile = 0; // export_multi_profile.py declares decode first constexpr int32_t kPrefillProfile = 1; executorch::extension::Module module("model_gemma3_multi_profile.pte"); @@ -290,75 +214,50 @@ executorch::extension::Module module("model_gemma3_multi_profile.pte"); } ``` -The guard records an index for the calling thread and nothing else — it does not -inspect the `Module`, `Method`, or delegate handles, and does not call TensorRT. -Each TensorRT delegate reads it inside its own `execute()` and switches there. -Construct it on the thread that calls `forward()`. With no guard in scope, every -delegate runs profile 0. - -To have each delegate choose from the input shapes instead of being told an -index, use the named constructor: - -```cpp -auto profile_guard = OptimizationProfileGuard::automatic(); -``` - -The index reaches every TensorRT delegate in the method, and each resolves it -against its own profile list. Nothing makes index 1 mean the same thing in two -engines, so pin by index only when the `.pte` holds one TensorRT engine or when -its engines were compiled from a single profile list. An engine with just one -profile runs profile 0 and logs that the pin did nothing; a multi-profile engine -that lacks the index fails the execution. +A profile is named by its index in the list declared at export, so define your own constants that +match the exporter. The guard only records the index for the calling thread. Each TensorRT engine +reads it inside its own `execute()` and switches there, so create the guard on the thread that +calls `forward()`. With no guard, every engine runs profile 0. To let each engine pick from the +input shapes instead, use `auto profile_guard = OptimizationProfileGuard::automatic();`. -Build and run: +The index reaches every TensorRT engine in the method, and each looks it up in its own list of +profiles. So pin by index only when the `.pte` holds one engine, or when its engines were compiled +from one list of profiles. An engine with one profile runs profile 0 for any pin above 0 and logs +that the pin did nothing. An engine with several profiles that lacks the index fails the call, and +a negative index always fails. The delegate README in `cpp/src/torch_tensorrt/executorch/` has +the full rules. ```bash -cmake --build build-executorch-reference-runner --target example_executorch_multi_profile_runner -j ./build-executorch-reference-runner/example_executorch_multi_profile_runner \ --model_path=model_gemma3_multi_profile.pte ``` -After the correctness walkthrough it times decode on each profile, the same -comparison `examples/dynamo/multi_optimization_profiles.py` makes through the -Python runtime. Everything below is the ExecuTorch delegate, so the numbers are -not comparable with that example's or with the two-engine table in the -[tutorial](../../docsrc/tutorials/runtime_opt/multi_optimization_profiles.rst): -different runtime, and there a separately compiled single-profile engine -provides the baseline rather than a second profile of this one. It prints the -per-call latency of decode on each profile and of prefill, then how much faster -decode is on its own profile. - -The profile is pinned around each timing loop rather than per call, so profile -switches stay out of the measurement. Prefill appears once because the decode -profile does not accept a 128-token input at all — prefill has only one profile -it can run on. - -### What Selecting A Profile Is Worth - -`multi_profile_benchmark.cpp` times the same prefill/decode loop twice against -one engine: once with every call pinned to the prefill profile (it accepts -`seq == 1` too, so decode runs on prefill-tuned kernels), and once with each -phase pinned to its own profile. Both sides are the same engine, so what this -isolates is the choice of profile, not the cost of a separately compiled -single-profile build. +After checking that each profile gives correct results, the program times decode on each profile, +and prefill on the one profile that accepts its 128-token input. It then prints how much faster +decode is on its own profile. The profile is pinned around each timing loop, so switching is not +part of the measurement. These numbers come from the ExecuTorch delegate, so do not compare them +with the [multiple optimization profiles tutorial](https://docs.pytorch.org/TensorRT/tutorials/runtime_opt/multi_optimization_profiles.html), +which uses a different runtime and a different baseline. + +### What picking a profile is worth + +`multi_profile_benchmark.cpp` runs the same prefill and decode loop twice on one engine: once with +every call pinned to the prefill profile, which also accepts one token, and once with each phase +on its own profile. Both runs use the same engine, so the difference is only the choice of +profile. ```bash -cmake --build build-executorch-reference-runner --target example_executorch_multi_profile_benchmark -j ./build-executorch-reference-runner/example_executorch_multi_profile_benchmark \ --model_path=model_gemma3_multi_profile.pte ``` -A switch is charged to whichever call makes it, so switching pays off only once -the per-token decode gain, summed over a round's decode steps, exceeds the -switch cost. The mini Gemma-3 exported by default exercises the API; export with -`--weights google/gemma-3-1b-it` to see what the feature is worth. - -When comparing wall-clock rounds, keep blocks long (`--block_rounds=8`). With -short blocks the prefill-only configuration inherits the decode profile from the -preceding switching block and pays a switch it would never pay in production, -which inflates switching's margin. - -Read the `min` and `p10` columns. The two configurations are interleaved in -short blocks so that other tenants on the GPU perturb both equally, and since -interference only ever adds time, the low percentiles are the signal; the median -and `p90` tell you how busy the machine was, not what switching cost. +- The call that switches profiles pays for the switch. Switching wins only when the decode gain, + added up over a round's decode steps, is larger than that cost. +- The small default model only exercises the API. Export with `--weights google/gemma-3-1b-it` to + see what the feature is worth. +- To compare whole rounds, use long blocks (`--block_rounds=8`). With short blocks, the + prefill-only run inherits the decode profile from the block before it, and pays a switch it would + never pay in production. +- Read the `min` and `p10` (10th percentile) columns. The two runs alternate in short blocks, so + other work on the GPU slows both equally. Interference only ever adds time, so the low + percentiles are the signal. The median and `p90` show how busy the machine was. diff --git a/examples/executorch_reference_runner/load_model_device_resident.py b/examples/executorch_reference_runner/load_model_device_resident.py index 1337f749a70..b2d877f1363 100644 --- a/examples/executorch_reference_runner/load_model_device_resident.py +++ b/examples/executorch_reference_runner/load_model_device_resident.py @@ -5,11 +5,9 @@ The program this loads was exported with ``skip_h2d_for_method_inputs`` and ``skip_d2h_for_method_outputs``, so its boundary carries no copy operators and it -wants a CUDA input. A host input still gives the right answer, but it costs a -staging copy that the export existed to remove, so it defeats the point rather -than breaking the program. Where a host-backed buffer tagged for the device does -fail is in the runtime rather than in this delegate, on a device that reports it -cannot read pageable host memory. +wants a CUDA input. ExecuTorch's Python runtime refuses a host input. From C++, +a host input costs a copy on every call when TensorRT reads it first, and fails +when the CUDA backend does. That contract is what this script checks. ``export_device_resident.py`` already asserts the serialized program contains no boundary copy operators; this asserts @@ -63,11 +61,8 @@ outputs = forward.execute((x,)) y = outputs[0] -# This catches the export flags not taking, and that is all it catches. It does not prove the buffer is -# device memory: the tag and the pointer can disagree, and measured here they do, because the bindings -# hand back a buffer tagged cuda whose pointer is ordinary host memory. Proving residency needs the -# driver's own view of the pointer, which is more machinery than an example should carry, so the honest -# claim is the narrow one. The numbers below are still checked, and they are correct. +# This catches the export flags not taking. It does not ask the driver whether the pointer is device +# memory, which is more machinery than an example should carry; the numbers below are still checked. if not y.is_cuda: raise AssertionError( f"FATAL: output is tagged {y.device}, not cuda. The usual cause is that " diff --git a/examples/torchtrt_executorch_example/export_coalesced.py b/examples/torchtrt_executorch_example/export_coalesced.py index 8423a0db717..4981b19a819 100644 --- a/examples/torchtrt_executorch_example/export_coalesced.py +++ b/examples/torchtrt_executorch_example/export_coalesced.py @@ -42,10 +42,9 @@ Prerequisites ------------- -Install Torch-TensorRT with the ExecuTorch extra before running this example:: - - pip install -e ".[executorch]" \ - --extra-index-url https://download.pytorch.org/whl/nightly/cu132 +Install Torch-TensorRT with its ``executorch`` extra, as the Installation section of +the ExecuTorch deployment guide shows: +https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#installation ExecuTorch's CUDA backend also needs a CUDA toolkit (``nvcc``) at export time, for the AOTInductor compile. diff --git a/examples/torchtrt_executorch_example/export_device_resident.py b/examples/torchtrt_executorch_example/export_device_resident.py index 436976c3db4..fa8b8645fb0 100644 --- a/examples/torchtrt_executorch_example/export_device_resident.py +++ b/examples/torchtrt_executorch_example/export_device_resident.py @@ -23,32 +23,20 @@ already has its data on the GPU those copies are pure overhead, and they are what ``PropagateDeviceConfig`` turns off. -Two settings are needed, not one. Skipping the copy is not enough on its own: -memory planning allocates graph inputs and outputs by default, so the runtime -would still reserve its own buffer and fill it from the caller's memory with a -host memcpy, which is undefined for device memory and puts the copy straight -back. ``alloc_graph_input=False`` and ``alloc_graph_output=False`` are what stop -that, and ``enable_non_cpu_memory_planning`` is required for planning to run over -non-CPU tensors at all. - -Those two settings choose one of two valid arrangements, and it is worth knowing -which. With them off, as here, nobody allocates the boundary buffers and the -CALLER supplies their addresses, which is what a C++ consumer does and what -ExecuTorch's own reference runner is built for. The other arrangement leaves -planning on, so the program's own device arena owns them; ExecuTorch's own CUDA -example takes that one, calling ``to_executorch()`` with no planning overrides at -all. - -The first arrangement cannot work from Python as things stand, because the -bindings do not supply an address for a non-planned output and there is no way to -give them one. Measured: with ExecuTorch's CUDA backend owning the last partition -a program exported this way fails every time, and the C++ error says why, that the -data pointer cannot be null. It appears to work when the TensorRT delegate owns -the last partition only because that delegate does not check, and the buffer it is -handed is host memory. - -So run a program exported this way from C++, and export with planning left on if -Python is the target. +Skipping the copies is not enough on its own. Memory planning allocates the +graph inputs by default, and a program that owns its input buffer cannot take a +GPU tensor: setting one fails with error ``0x10`` (``NotSupported``), and the +log says "Error copying tensor data at input 0". So ``alloc_graph_input=False`` +is required. ``enable_non_cpu_memory_planning``, on by default, is what lets +planning place tensors in GPU memory at all. + +This script also passes ``alloc_graph_output=False``, so the runtime does not +copy the outputs: a C++ caller sets their memory once with +``Module::set_output``, and from Python the runtime allocates it and does not +overwrite an output the caller still holds. Planned outputs work too, but they +live in the program's own GPU memory, which the next call overwrites, and +Python's ``execute()`` copies them on every call unless it gets +``clone_outputs=False``. This script asserts on the serialized program rather than trusting the flags: the exported ``.pte`` must contain neither copy operator, and it must still carry @@ -57,10 +45,9 @@ Prerequisites ------------- -Install Torch-TensorRT with the ExecuTorch extra before running this example:: - - pip install -e ".[executorch]" \ - --extra-index-url https://download.pytorch.org/whl/nightly/cu132 +Install Torch-TensorRT with its ``executorch`` extra, as the Installation section of +the ExecuTorch deployment guide shows: +https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#installation ExecuTorch's CUDA backend also needs a CUDA toolkit (``nvcc``) at export time, for the AOTInductor compile. @@ -141,9 +128,9 @@ def main() -> None: skip_d2h_for_method_outputs=True, ), enable_non_cpu_memory_planning=True, - # Without these the runtime allocates its own buffer for the - # input and output it was just told not to copy, and fills the - # input from the caller's memory with a host memcpy. + # The runtime refuses to copy the caller's CUDA tensor into a + # planned input, so the input must not be planned. An unplanned + # output is not copied either. memory_planning_pass=MemoryPlanningPass( alloc_graph_input=False, alloc_graph_output=False ), diff --git a/examples/torchtrt_executorch_example/export_dynamic_shape.py b/examples/torchtrt_executorch_example/export_dynamic_shape.py index 37ab114abac..cb8c272daea 100644 --- a/examples/torchtrt_executorch_example/export_dynamic_shape.py +++ b/examples/torchtrt_executorch_example/export_dynamic_shape.py @@ -17,12 +17,9 @@ Prerequisites ------------- -Install Torch-TensorRT with the ExecuTorch extra before running this example:: - - pip install -e ".[executorch]" \ - --extra-index-url https://download.pytorch.org/whl/nightly/cu132 - -See https://pytorch.org/executorch/stable/getting-started-setup.html for details. +Install Torch-TensorRT with its ``executorch`` extra, as the Installation section of +the ExecuTorch deployment guide shows: +https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#installation """ # %% diff --git a/examples/torchtrt_executorch_example/export_kv_cache_decode.py b/examples/torchtrt_executorch_example/export_kv_cache_decode.py index 3586d3e3a78..24c9d81a6c8 100644 --- a/examples/torchtrt_executorch_example/export_kv_cache_decode.py +++ b/examples/torchtrt_executorch_example/export_kv_cache_decode.py @@ -29,10 +29,9 @@ Prerequisites ------------- -Install Torch-TensorRT with the ExecuTorch extra before running this example:: - - pip install -e ".[executorch]" \ - --extra-index-url https://download.pytorch.org/whl/nightly/cu132 +Install Torch-TensorRT with its ``executorch`` extra, as the Installation section of +the ExecuTorch deployment guide shows: +https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#installation """ import argparse diff --git a/examples/torchtrt_executorch_example/export_multi_profile.py b/examples/torchtrt_executorch_example/export_multi_profile.py index 4d1aa189231..db0a5c0bc4a 100644 --- a/examples/torchtrt_executorch_example/export_multi_profile.py +++ b/examples/torchtrt_executorch_example/export_multi_profile.py @@ -45,18 +45,12 @@ Prerequisites ------------- -Install Torch-TensorRT with the ExecuTorch extra before running this example. -The Gemma-3 definition comes from ``transformers``, which that extra does not -pull in:: - - pip install --pre -e ".[executorch]" "transformers>=5.4.0" \ - --extra-index-url https://download.pytorch.org/whl/nightly/cu132 +Install Torch-TensorRT with its ``executorch`` extra, as the ExecuTorch +deployment guide shows: +https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#installation -``--pre`` is required: the ExecuTorch companion package has only prerelease -versions on the nightly index, and without it pip falls back to the last -release, which pins an older torch and fails to resolve. - -See https://pytorch.org/executorch/stable/getting-started-setup.html for details. +The Gemma-3 definition comes from ``transformers``, which that extra does not +install, so also install ``transformers>=5.4.0``. """ # %% diff --git a/examples/torchtrt_executorch_example/export_static_shape.py b/examples/torchtrt_executorch_example/export_static_shape.py index 76e15372085..18d121d107b 100644 --- a/examples/torchtrt_executorch_example/export_static_shape.py +++ b/examples/torchtrt_executorch_example/export_static_shape.py @@ -13,12 +13,9 @@ Prerequisites ------------- -Install Torch-TensorRT with the ExecuTorch extra before running this example:: - - pip install -e ".[executorch]" \ - --extra-index-url https://download.pytorch.org/whl/nightly/cu132 - -See https://pytorch.org/executorch/stable/getting-started-setup.html for details. +Install Torch-TensorRT with its ``executorch`` extra, as the Installation section of +the ExecuTorch deployment guide shows: +https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#installation """ # %% diff --git a/py/torch-tensorrt-executorch-runtime/README.md b/py/torch-tensorrt-executorch-runtime/README.md index d5235d3fce8..db7e358ca7e 100644 --- a/py/torch-tensorrt-executorch-runtime/README.md +++ b/py/torch-tensorrt-executorch-runtime/README.md @@ -1,21 +1,79 @@ -# Torch-TensorRT ExecuTorch Runtime Wheel +# torch-tensorrt-executorch-runtime -This directory builds `torch-tensorrt-executorch-runtime`. The Linux wheel -contains one shared library, `lib/libexecutorch_backend_tensorrt.so`, holding the -TensorRT delegate and nothing else. The ExecuTorch runtime it registers with -comes from the `executorch` wheel. +This wheel is the TensorRT delegate for the ExecuTorch runtime: the plugin that runs the TensorRT +engines stored in a `.pte` file. It holds one shared library, +`lib/libexecutorch_backend_tensorrt.so`, and a CMake package for it. The ExecuTorch runtime that +the delegate plugs into comes from the `executorch` wheel. -The layout follows the one ExecuTorch uses for its own backends, so the delegate -is an out-of-tree sibling of them rather than a special case: +Its layout is close to the one ExecuTorch uses for its own delegates: -``` +```text executorch/ torch_tensorrt_executorch_runtime/ lib/libexecutorch_backend_cuda.so lib/libexecutorch_backend_tensorrt.so share/cmake/executorch-config.cmake lib/cmake/executorch_backend_tensorrt/executorch_backend_tensorrt-config.cmake ``` -Python users just import the package. A C++ app links it the same way it links -one of ExecuTorch's own backends: +The [ExecuTorch deployment guide](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html) +is the main documentation: how to install this wheel together with Torch-TensorRT, export a model, +and run it from Python and C++. To install, follow its +[Installation](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#installation) +section. + +## Version matching + +This wheel has to match its Torch-TensorRT wheel: the same PyTorch, ExecuTorch, CUDA and TensorRT, +built with the same C++ ABI (the compiler's binary conventions). It needs CUDA 13. Builds +exist for `cu132` and `cu134`. The Linux wheels are tagged `manylinux_2_28_x86_64` and +`manylinux_2_35_aarch64`, so aarch64 needs glibc 2.35 or newer. The Windows x86-64 wheel holds +`executorch_backend_tensorrt.dll` (a Windows shared +library, or DLL) instead. + +On the PyTorch indexes it pins the exact ExecuTorch build it was compiled against, including its +CUDA label, such as `+cu132`. The delegate needs `libexecutorch_extension_cuda.so`, which only +ExecuTorch's CUDA wheels ship, so a CPU build of ExecuTorch would install and then fail on import. +The exact pin rules that out. PyPI allows no label, so there it pins the version alone. On Linux, +PyPI's ExecuTorch is the CUDA build. + +It names PyTorch with no version, because ExecuTorch leaves that choice open too. Installed on its +own, it can get a PyTorch older than the ExecuTorch it pins, and that fails at the first run with a +confusing error about a tensor. Installing it as the deployment guide shows avoids that. + +## Use it from Python + +Import the package once, before you load a program. The import registers the delegate with +ExecuTorch. Everything after it is ExecuTorch's own runtime API: + +```python +from pathlib import Path + +import torch +import torch_tensorrt_executorch_runtime # noqa: F401 +from executorch.runtime import Runtime + +program = Runtime.get().load_program(Path("model.pte")) +forward = program.load_method("forward") +outputs = forward.execute((torch.ones((2, 3, 4, 4)),)) +``` + +ExecuTorch's own delegates register because they are linked into its Python extension. A delegate +in a separate wheel cannot join that link, and ExecuTorch has no hook to discover one, so this +package loads the delegate when it is imported. It replaces nothing: ExecuTorch keeps its own +delegates and kernels. + +- If the delegate cannot load, the import raises `DelegateCompatibilityError` and says what to + install. You get the error at the import, not later as a program that cannot find its backend. +- The import order does not matter: importing this package before or after `executorch.runtime` + works the same. +- Set `TORCH_TENSORRT_SKIP_DELEGATE_REGISTRATION=1` to import the package without loading the + delegate. That is for tools that need only its metadata, such as the path query below. Call + `register()` later if such a tool then needs the delegate. After a normal import it does nothing. + +Load programs through `executorch.runtime`, ExecuTorch's public Python API. Do not use the +underscore-prefixed loader inside its extension. It is private and can change without notice. + +## Use it from C++ + +A C++ application links the delegate the same way it links one of ExecuTorch's own delegates: ```cmake find_package(executorch REQUIRED COMPONENTS backend_cuda kernels_optimized) @@ -28,12 +86,37 @@ target_link_libraries(my_app PRIVATE ) ``` -The optimized-kernel library supplies the `et_copy` host/device copy operators. +Point CMake at both wheels. ExecuTorch's location comes from its package metadata, because +`executorch.__file__` is not set: + +```bash +ET_CMAKE=$(python -c 'import importlib.metadata as m, pathlib; print(pathlib.Path(str(m.distribution("executorch").locate_file("executorch"))) / "share" / "cmake")') +DELEGATE=$(TORCH_TENSORRT_SKIP_DELEGATE_REGISTRATION=1 python -c 'import pathlib, torch_tensorrt_executorch_runtime as r; print(pathlib.Path(r.__file__).parent)') +cmake -S . -B build -DCMAKE_PREFIX_PATH="${ET_CMAKE};${DELEGATE}" +``` -Linking that target also records the wheel's own library directory in your binary, -so it finds the delegate without a library path being set. That is right for an -application built against an installed wheel and wrong for anything you intend to -redistribute, so it can be turned off: +The deployment guide explains each target and shows a complete program. The +[delegate README](https://github.com/pytorch/TensorRT/blob/main/cpp/src/torch_tensorrt/executorch/README.md) +lists the rules a C++ caller has to keep (threads, buffers, streams), the load errors, and the +known issues. + +- **There is no header.** The delegate registers itself from a static initializer when its library + loads. The CMake target links the library with `--no-as-needed`, because your code references no + symbol in it, and the linker would otherwise drop it and leave the backend unregistered. +- **You do not need to find TensorRT.** The delegate finds it in the TensorRT wheel next to it. +- **CMake 3.28 or newer** is needed for the recipe above: on older versions the ExecuTorch package + does not provide the `executorch::` targets. CMake is not part of these wheels, so install it + yourself. The section below shows a way for older CMake. +- **No header for optimization profiles.** Choosing a profile needs + `OptimizationProfileGuard` from `TensorRTBackend.h`, which only the source tree and + `libtorchtrt.tar.gz` carry. On Linux you can compile against it, with the TensorRT headers, and + link this wheel's delegate. On Windows the wheel's library does not export the guard. Without it, + every engine runs its profile 0. + +On Linux, linking the target also records the wheel's library directory in your binary, so it +finds the delegate with no library path set. That is right for an application built against an +installed wheel, and wrong for anything you ship to other machines. To turn it off, set both of +these before `find_package`, then ship the delegate yourself or point a library path at it: ```cmake set(EXECUTORCH_BACKEND_TENSORRT_EMBED_RUNPATH OFF) @@ -41,190 +124,68 @@ set(CMAKE_SKIP_BUILD_RPATH ON) find_package(executorch_backend_tensorrt REQUIRED) ``` -Set both before `find_package`. The second is needed because CMake records the -wheel's directory on its own when you link an imported library. Then ship the -delegate yourself, or point a library path at it at run time. -For device-resident exports with `alloc_graph_output=False`, C++ Module callers -must provide a CUDA output tensor with `Module::set_output` before execution. -ExecuTorch's Python runtime allocates one for each output itself. - -The simplest arrangement turns off allocation for the inputs and leaves it on for the outputs. -The caller then supplies the input buffer, which is the copy worth avoiding, and the program's own -device arena owns the output. -Measured on two architectures, that arrangement runs from Python and from C++ with -no boundary copies in the program and results identical to eager. Leaving it on for both is -refused by the export step, because a program that allocates its own inputs has nothing for a -caller to hand in. - -Point CMake at both wheels. The example above calls `find_package(executorch)` -as well, and that package lives in its own distribution. ExecuTorch is a -namespace package, so `executorch.__file__` is `None` and has to be located -through its distribution metadata instead. The variable in front is what keeps -this a path query: importing this package loads the delegate, which raises when -the delegate cannot be loaded, and a path does not need it loaded. - -```bash -cmake -DCMAKE_PREFIX_PATH="$(TORCH_TENSORRT_SKIP_DELEGATE_REGISTRATION=1 python -c 'import importlib.metadata as m, torch_tensorrt_executorch_runtime as r, pathlib; print(str(pathlib.Path(str(m.distribution("executorch").locate_file("executorch"))) / "share" / "cmake") + ";" + str(pathlib.Path(r.__file__).parent))')" ... -``` - -CMake 3.28 or newer is required for the example above, not because of this -package but because the `backend_cuda` component it pairs with rejects anything -older: earlier versions write the `$ORIGIN` token in a runtime search path -incorrectly. +The second line is needed because CMake records the directory on its own when you link an imported +library. -Bring your own CMake. It is not part of these wheels, and a freshly imaged Jetson -has none at all, which is easy to miss because the Python side needs none. +### CMake older than 3.28 -You do not need to find TensorRT. The delegate needs it, and it records where to -look relative to its own location, so a linker that reads that record resolves it -from the sibling wheel without being told. That is also why the example links no -TensorRT target: your application uses the delegate's interface, not TensorRT's. - -This package itself needs only 3.19. On 3.19 through 3.27 you can still use it, -by asking ExecuTorch for no components and linking the variables it gives you -instead of its targets: +This package itself needs only CMake 3.19. On 3.19 to 3.27, find ExecuTorch without components and +link the variables it sets instead of its targets: ```cmake find_package(executorch REQUIRED) find_package(executorch_backend_tensorrt REQUIRED) -target_include_directories(app PRIVATE ${EXECUTORCH_INCLUDE_DIRS}) -target_compile_definitions(app PRIVATE ${EXECUTORCH_COMPILE_DEFINITIONS}) -target_link_libraries(app PRIVATE +target_include_directories(my_app PRIVATE ${EXECUTORCH_INCLUDE_DIRS}) +target_compile_definitions(my_app PRIVATE ${EXECUTORCH_COMPILE_DEFINITIONS}) +target_link_libraries(my_app PRIVATE ${EXECUTORCH_LIBRARIES} executorch::backend_tensorrt ) ``` -Without an imported target to carry them, the include directories and the compile -definitions have to be applied by hand as well, which is what the three variables -above are for. - -There is nothing to include. The wheel ships no header: the delegate registers -itself with ExecuTorch's backend registry from a static initializer inside the -shared library, and everything after that is ExecuTorch's own runtime API. The -one feature this leaves out is choosing an optimization profile, which needs -`OptimizationProfileGuard` from `torch_tensorrt/executorch/TensorRTBackend.h`. -Linked from this wheel, every engine runs its profile 0. Build against the -source tree or the release tarball to select profiles. The -CMake target links the library with `--no-as-needed`, because nothing in a -consumer references a symbol the delegate defines, and the default would drop it -and leave the backend unregistered. - -The wheel must use the same Python, PyTorch, ExecuTorch, CUDA, TensorRT, and -C++ ABI as its matching Torch-TensorRT wheel. This delegate requires CUDA 13. Nightly builds, and -release builds on the PyTorch indexes, cover `cu132` and `cu134`. The PyPI release wheels are the -`cu132` builds, the CUDA version of PyTorch's PyPI wheel. -Ordinary Torch-TensorRT release and JetPack builds keep their separate CUDA 12 support. - -## Install -The -[user guide](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#installation) -has the install commands. Installed that way, the `executorch` extra of a Torch-TensorRT wheel -brings this wheel (the one built together with it) and the exact CUDA build of ExecuTorch it links. -This is the nightly command. Swap `cu132` for `cu134` if you run CUDA 13.4: -```bash -python -m pip install --pre "torch-tensorrt[executorch]" \ - --index-url https://download.pytorch.org/whl/nightly/cu132 -``` -`--pre` lets pip choose the nightly builds, which are prereleases. -You never name PyTorch: it arrives as a dependency of Torch-TensorRT. This wheel asks -for it by name with no version, because ExecuTorch imports it and leaves the choice open, so installing this -wheel on its own lets the resolver pick a PyTorch older than the ExecuTorch it pins. That -combination installs, imports and loads a program, then fails at the first run with a message about -a tensor rather than about versions. Install it the way this section shows and the question does not -arise. Building this wheel -by hand is only needed to work on it, and those instructions are below. - -ExecuTorch can be an older nightly than the rest, because this wheel pins the exact build it -was compiled against. That is deliberate, and the next paragraph says why. - -A CUDA build of ExecuTorch is required, not only to build against. The delegate needs a -library that only ExecuTorch's CUDA wheels carry, so a processor-only build installs and -then fails on import. On the PyTorch indexes this wheel pins the exact ExecuTorch build it was -compiled against, label and all, which is what stops that happening quietly. -## Use it from Python - -Import the package once, anywhere before you load a program. Importing is what -registers the delegate, and nothing else about your code changes: - -```python -from pathlib import Path - -import torch -import torch_tensorrt_executorch_runtime # noqa: F401 -from executorch.runtime import Runtime - -program = Runtime.get().load_program(Path("model.pte")) -forward = program.load_method("forward") -outputs = forward.execute((torch.ones((2, 3, 4, 4)),)) -``` - -If the delegate cannot be loaded, that import raises, rather than letting the -failure surface later as a program that will not load. - -A program whose graph is split between TensorRT and ExecuTorch's own CUDA backend -needs no extra import: both backends are registered, and the program says which -parts go where. If its export wrote a data file such as `aoti_cuda_blob.ptd`, pass it -with `load_program(..., data_path=Path("aoti_cuda_blob.ptd"))`. - ## Runtime libraries -The wheel bundles no ExecuTorch, PyTorch, c10, TensorRT, or CUDA shared -libraries. The delegate carries origin-relative runtime search paths, exactly as -the build sets them: +The wheel bundles no ExecuTorch, PyTorch, TensorRT or CUDA libraries. The delegate finds them in +the wheels installed next to it, through these runtime search paths: - `$ORIGIN` - `$ORIGIN/../../executorch/lib` - `$ORIGIN/../../tensorrt_libs` - `$ORIGIN/../../nvidia/cu13/lib` -There is no `$ORIGIN/../torch/lib` entry, because the delegate links no torch, -and no `$ORIGIN/../../nvidia/cuda_runtime/lib` entry, because that is the CUDA 12 -layout and this package requires CUDA 13. +`$ORIGIN` is this package's `lib` directory. The delegate links no PyTorch, so there is no +`torch/lib` entry. For a TensorRT or CUDA installation outside these places, make its `lib` +directory visible to the dynamic loader, for example with `LD_LIBRARY_PATH`. -`$ORIGIN` is this package's `lib` directory; the three `../../` entries reach sibling -distributions, because `libexecutorch.so`, the TensorRT libraries, and the CUDA -runtime belong to other wheels. These packages are installed transitively with -the matching `torch-tensorrt` wheel. For a system TensorRT or CUDA installation -outside these standard locations, its `lib` directory must be available through -the system dynamic loader configuration or `LD_LIBRARY_PATH`. +NVIDIA publishes its CUDA runtime wheels for Linux only. On Windows the wheel does not depend on +one, and the delegate loads the CUDA runtime from a CUDA 13 toolkit install instead: when the +package is imported, it adds the `bin` directory of `CUDA_PATH` or `CUDA_HOME` to the DLL search +path. -The shared Linux build tags this wheel with `wheel tags` after the native -checks: `manylinux_2_28_x86_64` or `manylinux_2_35_aarch64`, with `py3-none`. -This updates wheel metadata and RECORD without bundling external libraries or -changing the delegate. Payload, dependency metadata, and installed-library -resolution are checked before the shared artifact is uploaded. +## Build from source -## Build +You need this only to work on the delegate itself. The steps below are for Linux. > [!IMPORTANT] -> Build this wheel with `--no-build-isolation`. The delegate links the -> prebuilt runtime out of the ExecuTorch wheel that is installed at build time, -> and it must use the exact PyTorch installation the matching Torch-TensorRT -> artifacts were built against. An isolated build may download a newer, -> ABI-incompatible PyTorch or ExecuTorch. - -The build shells out to Bazel to compile the delegate, so `bazelisk` or `bazel` -must be on `PATH`. TensorRT itself arrives through Bazel's `@tensorrt` external -repository, so no local SDK path is needed. - -The example below assumes Linux with a matching CUDA 13 PyTorch and -Torch-TensorRT installation. Substitute the channel for your CUDA throughout, such -as `cu134` for CUDA 13.4. - -On x86 the build also needs a developer toolset, the kind the release containers -carry. Without one, the C++ runtime helpers the delegate references are not -absorbed into the library, so the check that runs after linking refuses the result -and names the symbol version it found. That refusal is correct: the library would -demand a newer C++ runtime than its own platform tag promises. An ordinary -development machine usually has no such toolset, so build there through the same -container the release uses, or take the wheel a continuous integration run already -produced. Arm does not have this problem, and its tag is higher for that reason. - -The build reads the CUDA location out of the repository's `MODULE.bazel`, and the -copy checked in names one specific version. On a machine with a different CUDA, the -fetch fails saying that path does not exist. That file is generated from a template, so -render it for your machine before building. The template takes seven values: +> Build with `--no-build-isolation`. The delegate links the ExecuTorch runtime from the +> `executorch` wheel installed at build time, and it must use the same PyTorch as the matching +> Torch-TensorRT build. An isolated build may download a newer PyTorch or ExecuTorch that is not +> ABI compatible. + +What you need: + +- A matching CUDA 13 PyTorch and Torch-TensorRT installed. +- `bazelisk` or `bazel` on `PATH`. The build calls Bazel to compile the delegate, and Bazel fetches + TensorRT itself, so you need no local TensorRT SDK. +- On x86-64, a developer toolset like the one in the release containers. Without it, the C++ + runtime helpers the delegate uses are not built into the library, and the check after linking + refuses the result: the library would need a newer C++ runtime than its platform tag promises. + An ordinary development machine usually has no such toolset, so build in the release container, + or take a wheel that a continuous integration run built. aarch64 does not have this problem. + +First render `MODULE.bazel` for your machine. The checked-in copy names one CUDA path, and on a +machine with another CUDA the fetch fails, saying that path does not exist. The command below +overwrites the checked-in file, so do not commit the result. Substitute your own CUDA version: ```bash export CUDA_HOME=/usr/local/cuda-13.2 @@ -232,15 +193,17 @@ export CU_VERSION=cu132 export CHANNEL=nightly export BUILD_VERSION=0.2.1.dev0 export TORCH_INSTALL_PATH="$(python -c 'import torch; print(torch.__path__[0])')" -export TENSORRT_CUDA_VERSION_UPPER_BOUND=14.0 -export TENSORRT_RTX_CUDA_VERSION_UPPER_BOUND=14.0 +export TENSORRT_CUDA_VERSION_UPPER_BOUND=13.4 +export TENSORRT_RTX_CUDA_VERSION_UPPER_BOUND=13.4 envsubst < toolchains/ci_workspaces/MODULE.bazel.tmpl > MODULE.bazel ``` -Render it yourself rather than running `packaging/pre_build_script.sh`. That script is -for the release containers: it also installs packages, and on a development machine it -can replace the PyTorch you already have with one from another index, which leaves you -worse off than when you started. +Do not run `packaging/pre_build_script.sh` for this. It is for the release containers: it also +installs packages, and on a development machine it can replace your PyTorch with one from another +index. + +Then install the build dependencies and build the wheel. Use the ExecuTorch version that +`pyproject.toml` pins: ```bash python -m pip install pyyaml patchelf tensorrt-cu13 \ @@ -253,154 +216,56 @@ python -m pip wheel --no-build-isolation --no-deps \ --wheel-dir dist py/torch-tensorrt-executorch-runtime ``` -`TORCH_TENSORRT_EXECUTORCH_RUNTIME_VERSION` sets the companion's own version. Without it, the -version is the companion's base version plus a development suffix and Git revision. The shared -build keeps the main wheel's date and CUDA suffix while retaining the companion's independent -base version. Its CMake version also describes the companion, not the main wheel. - -A release build of the main wheel (`RELEASE=1`) reads the same variable: +- `tensorrt-cu13` is needed at build time so the wheel can record an exact `tensorrt-cu13==` + requirement, read from the installed package. It also installs from PyPI alone: there its + `tensorrt-cu13-libs` dependency is a source package that downloads the real wheel from NVIDIA's + index while it builds. NVIDIA's index only makes that step faster. +- The delegate compiles and links against the installed ExecuTorch wheel, which ships the headers, + the prebuilt runtime and a CMake package. It must be a CUDA wheel. No ExecuTorch source checkout + is involved. +- `TORCH_TENSORRT_EXECUTORCH_RUNTIME_VERSION` sets this wheel's version. Without it, the version + is the base version in `version.txt` plus `.dev0` and the Git revision. Nightly builds keep this + wheel's own base version, with the main wheel's date and CUDA suffix, such as + `0.2.1.dev20261023+cu132`. A release on the PyTorch indexes uses the base version and the CUDA + suffix, such as `0.2.1+cu132`; the PyPI release has none, such as `0.2.1`. +- A release build of the main wheel (`RELEASE=1`) reads the same variable and records the version + for its error messages. When `TORCH_TENSORRT_EXECUTORCH_RUNTIME_CHANNEL` names a place that + carries the pinned ExecuTorch build (`nightly` for a development build, `test` or `pypi` for a + release), the main wheel's `executorch` extra requires exactly this version of this wheel. A + Windows release wheel for PyPI records no channel, because PyPI's Windows ExecuTorch is the CPU + build. + +Install the result with PyTorch's indexes for your CUDA version. The ExecuTorch build it pins is +on one of them: a development build on the nightly index, a release on the test and stable +indexes. -- It records the version, so its error messages can name this wheel. -- If `TORCH_TENSORRT_EXECUTORCH_RUNTIME_CHANNEL` names a place that carries the pinned - ExecuTorch build (`nightly` for a development build, `test` or `pypi` for a release), its - `executorch` extra requires exactly that version. Otherwise the extra leaves this wheel out. - Windows release wheels for PyPI set no channel. - -`tensorrt-cu13` is needed at build time so the wheel can record an exact -`tensorrt-cu13==` requirement: that version is read from the installed -distribution, which the Bazel-provided libraries alone do not carry. NVIDIA's -index above lets pip take its wheel directly; without it, the PyPI `tensorrt-cu13` -source package downloads that wheel while it builds. - -The delegate compiles and links entirely against the installed ExecuTorch -wheel, which ships the headers, the prebuilt runtime, and a CMake package. A -CUDA wheel is required: the CPU wheel ships no CUDA extension, and ExecuTorch -releases up to 1.4.1 ship no linkable runtime at all. ExecuTorch is not built -from source for this wheel, so no source checkout or `EXECUTORCH_SOURCE_DIR` is -involved. +```bash +python -m pip install dist/torch_tensorrt_executorch_runtime-*.whl \ + --extra-index-url https://download.pytorch.org/whl/nightly/cu132 \ + --extra-index-url https://download.pytorch.org/whl/test/cu132 +``` ### Rebuilding and editable installs -After rebuilding a wheel, install it with `python -m pip install --no-deps ---force-reinstall` followed by its path. This replaces an installed wheel with -the same version. The delegate is a shared library, not a Python extension; -`build_ext --inplace` does not rebuild it. +To replace an installed wheel of the same version with a rebuilt one, use +`python -m pip install --no-deps --force-reinstall` and the wheel's path. The delegate is a shared +library, not a Python extension, so `build_ext --inplace` does not rebuild it. -For editable development, use the same matching dependencies and version setting: +For editable development, with the same dependencies and version setting: ```bash python -m pip install --no-build-isolation --no-deps \ --editable py/torch-tensorrt-executorch-runtime ``` -Repeat this command after native or CMake changes. Stop native consumers before -rebuilding and restart them afterward; rebuilding does not update a loaded -library. Python source changes are visible directly. For strict editable mode, -add `--config-settings editable_mode=strict` and keep the generated link directory. -The native library and generated CMake files live beside the source package in -its generated library directory. - -`TORCH_TENSORRT_EXECUTORCH_DEBUG`, `TORCH_TENSORRT_ALLOW_UNPINNED_EXECUTORCH`, -and `TORCH_TENSORRT_SKIP_DELEGATE_REGISTRATION` accept `1`, `true`, `yes`, or -`on`, ignoring case. All other values, including unset, empty, `0`, and `false`, -are false. - -## Registration - -Loading the delegate adds `TensorRTBackend` to the backend registry that the -installed ExecuTorch runtime owns. It replaces nothing: the stock runtime keeps -its own backends and kernels, and XNNPACK and CPU fallback regions behave -exactly as they do without this wheel. - -Registration happens in the delegate's static initializer, so the library has -to be loaded before a delegated program is loaded. Importing this package does -that, and nothing else, so an import is all a normal program needs. There is a -`register()` for the rare case where the import order is not yours to choose, and -calling it twice is harmless. - -```python -from pathlib import Path - -import torch -import torch_tensorrt_executorch_runtime # noqa: F401 -from executorch.runtime import Runtime - -program = Runtime.get().load_program(Path("model.pte")) -outputs = program.load_method("forward").execute((torch.randn(1, 3, 224, 224),)) -``` - -ExecuTorch's own delegates register because they are linked into its pybindings -extension, so loading that extension pulls them in. A delegate shipped in a -separate wheel cannot join that link, and ExecuTorch has no discovery hook for -out-of-tree backends, so this package performs the equivalent step at import -time. Set `TORCH_TENSORRT_SKIP_DELEGATE_REGISTRATION=1` to import it without -loading the delegate; that is for tooling that wants the metadata only, and the -path recipe above is exactly that case. - -Because registering is the point, an import that cannot register raises -`DelegateCompatibilityError` naming what to install, rather than returning a -module that registered nothing and leaving the program to fail later with a -backend it cannot find. - -Loading and running a program is ExecuTorch's API, not this package's. Tensor -placement, method lookup and output devices are all documented by ExecuTorch. A -program exported with `skip_h2d_for_method_inputs` keeps its inputs on the -device, because nothing here moves them off the GPU. - -ExecuTorch's Python runtime refuses a host tensor for such an input. From C++, when TensorRT is -the first delegate to read the input, the program does not check that promise, and it will not -tell you when you break it. -Hand it a host tensor and it runs, returns the right answer, and stages the copy -for you, once per call. So a program you believe is copying -nothing can be copying on every run with nothing to show it. The other -direction is refused, so the leniency goes one way only. If the copies are why -you exported this way, measure them rather than trusting the flag. - -Load programs through `executorch.runtime`, which is ExecuTorch's public Python -runtime API. Do not use the underscore-prefixed loader inside its pybindings -extension: it is private and can change without notice. The two were measured -against each other on every arrangement here, host and device, including a graph -split with ExecuTorch's own CUDA backend, and they agree to the digit. - -The choice that does matter is who owns the output buffer. With allocation on for -the graph outputs, the program's device arena owns them. With it off, a C++ caller -supplies each address, and ExecuTorch's Python runtime allocates one for each output. - -## Use - -A **CUDA** build of `executorch` is required at runtime, not just to build. The delegate carries a -`DT_NEEDED` on `libexecutorch_extension_cuda.so`, which only ExecuTorch's CUDA wheels ship, and -those live on PyTorch's indexes: a development build on the nightly one, a release on the test and -stable ones. `install_requires` names the version including its local -label, so only the CUDA build the delegate linked satisfies it. A label-free specifier admits any -label, which let a `+cpu` wheel resolve and then fail to load at import. Carrying the label rules -that out, because PEP 440 only ignores labels when the specifier omits them. It does bind the wheel -to one CUDA train, which is correct: the delegate links that train's runtime. A wheel for PyPI -cannot carry a label, so it names the version alone. On Linux, PyPI's ExecuTorch is the CUDA build; -on Windows it is the CPU build, so the extra of a Windows release wheel on PyPI leaves this wheel -out. - -To install the wheel built above, use the same CUDA indexes as the build (nightly for a -development ExecuTorch pin, test for a release). Build it for the -CUDA train you run on: the requirement names that train, because the delegate links its runtime. -Its exact development-version requirements already permit the required prereleases. -CUDA and TensorRT packages may resolve from PyPI or NVIDIA's index; the extra -index applies to the whole dependency solve. +Run this again after changing native code or CMake files. Restart any process that loaded the +delegate, because rebuilding does not update a loaded library. Python changes are visible at once. +For strict editable mode, add `--config-settings editable_mode=strict`. -```bash -python -m pip install dist/torch_tensorrt_executorch_runtime-*.whl \ - --extra-index-url https://download.pytorch.org/whl/nightly/cu132 \ - --extra-index-url https://download.pytorch.org/whl/test/cu132 -``` +### Environment variables -```python -from pathlib import Path +- `TORCH_TENSORRT_SKIP_DELEGATE_REGISTRATION`: import the package without loading the delegate. +- `TORCH_TENSORRT_ALLOW_UNPINNED_EXECUTORCH`: build against an ExecuTorch other than the pinned one. +- `TORCH_TENSORRT_EXECUTORCH_DEBUG`: build the delegate with debug information. -import torch -import torch_tensorrt_executorch_runtime # noqa: F401 -from executorch.runtime import Runtime - -program = Runtime.get().load_program(Path("model.pte")) -forward = program.load_method("forward") -outputs = forward.execute((torch.ones((2, 3, 4, 4)),)) -``` +Each accepts `1`, `true`, `yes` or `on`, in any case. Any other value, or no value, means off. diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index 3de2e759fb4..fcf737547d1 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -829,8 +829,9 @@ def save( parameter takes precedence. use_cuda_graphs (Optional[bool]): Bake CUDA graph replay on (True) or off (False) into every TensorRT delegate for ``output_format="executorch"``. - None leaves the choice unset so the runtime default applies. A load-time - option overrides this value. Non-booleans are rejected for ExecuTorch; + None leaves the choice unset so the runtime default applies. A load + option, or a process option set to False, overrides this value. + Non-booleans are rejected for ExecuTorch; other formats ignore non-None values with a warning. Use this keyword, not a raw ``CompileSpec("use_cuda_graphs", ...)``. kwargs: Additional format-specific kwargs. ``partitioners=``, @@ -839,8 +840,9 @@ def save( ``weight_streaming_budget_per_engine=``, ``zero_copy_kv=`` and ``external_engine_data=`` are only used with ``output_format="executorch"``; otherwise they are - ignored with a warning. Pass ``compile_specs=[CompileSpec("target_device", - b"cuda:")]`` to override the default target device (``cuda:0``). + ignored with a warning. A ``CompileSpec("target_device", ...)`` in + ``compile_specs=`` only labels the delegate's inputs and outputs; each + engine runs on the GPU it was built on. ``backend_config=`` takes an ``Optional[ExecutorchBackendConfig]`` and is forwarded to ``to_executorch(config=...)`` to customize ExecuTorch lowering (e.g. memory planning or device placement). @@ -848,9 +850,9 @@ def save( boundary TensorRT needs for every input and output. ``partitioners=`` appends caller partitioners after the TensorRT partitioner, so ops TensorRT does not take can be routed to another - delegate -- e.g. a ``CudaPartitioner`` for a coalesced TensorRT + - CUDA ``.pte``. See :ref:`the ExecuTorch save guide - ` for the ``CudaPartitioner`` recipe, its + delegate, for example a ``CudaPartitioner`` for a coalesced TensorRT + + CUDA ``.pte``. See :ref:`the ExecuTorch deployment guide + ` for the ``CudaPartitioner`` recipe, its export-time requirements (CUDA backend + nvcc), and the external ``.ptd`` weight caveats. ``constant_methods=``, ``transform_passes=`` and @@ -863,7 +865,8 @@ def save( TRT engines. ``generate_etrecord=True`` writes an ETRecord to ``_etrecord.bin`` next to the ``.pte`` for use with the - ExecuTorch Inspector. + ExecuTorch Inspector. It does not work for a program with + TensorRT engines: writing the record raises a ``SerializeError``. ``weight_streaming_budget_per_engine=`` takes an ``Optional[int]`` number of bytes of engine weights that may stay resident in GPU memory, with the rest streamed from host memory. It applies to diff --git a/py/torch_tensorrt/executorch/_export.py b/py/torch_tensorrt/executorch/_export.py index f14d344ce35..03cfe27cec7 100644 --- a/py/torch_tensorrt/executorch/_export.py +++ b/py/torch_tensorrt/executorch/_export.py @@ -593,7 +593,8 @@ def export( update discarded. ``torch_tensorrt.save(output_format="executorch", zero_copy_kv=True)`` owns both steps and needs only the one flag. generate_etrecord (bool): Ask ExecuTorch for an ETRecord for later debugging. - This copies the whole program, engines included. + This copies the whole program, engines included. Saving the record does not + work for a program with TensorRT engines: it raises a ``SerializeError``. weight_streaming_budget_per_engine (Optional[int]): Bytes of engine weights that may stay resident in GPU memory, with the rest streamed from host memory. It applies to **each** TensorRT engine separately, not as a total for the @@ -608,9 +609,10 @@ def export( program total split proportionally across engines. use_cuda_graphs (Optional[bool]): Bake CUDA graph replay on (True) or off (False) into every method's TensorRT compile specs. None leaves the choice - unset so the runtime default applies. A load-time option overrides this - value. Non-booleans and raw ``CompileSpec("use_cuda_graphs", ...)`` entries - are rejected; use this keyword instead. + unset so the runtime default applies. A load option, or a process option + set to False, overrides this value. Non-booleans and raw + ``CompileSpec("use_cuda_graphs", ...)`` entries are rejected; use this + keyword instead. external_engine_data (Optional[str]): Write every TensorRT engine to an ExecuTorch data file named ``.ptd`` instead of into the ``.pte``, which then holds only a key per engine. Pass the name without diff --git a/py/torch_tensorrt/executorch/partitioner.py b/py/torch_tensorrt/executorch/partitioner.py index 82b2e5ea02e..e57eac95495 100644 --- a/py/torch_tensorrt/executorch/partitioner.py +++ b/py/torch_tensorrt/executorch/partitioner.py @@ -186,11 +186,10 @@ class TensorRTPartitioner(Partitioner): # type: ignore[misc] each such node becomes its own partition so the backend can serialize the engine to the same format as the TRT runtime. - If `compile_specs` does not already contain a ``target_device`` entry, - one defaulting to ``cuda:0`` is auto-appended (mirroring CudaPartitioner). - Callers targeting a non-default GPU should pre-populate - ``compile_specs`` with the desired ``CompileSpec("target_device", - b"cuda:")`` to override the default. + If `compile_specs` has no ``target_device`` entry, each partition gets one + naming the GPU its own engine was built on, or ``cuda:0`` if that cannot be + read. A ``CompileSpec("target_device", b"cuda:")`` in + ``compile_specs`` sets it for every partition instead. Note: ``target_device`` is AOT metadata only -- it drives ExecuTorch's PropagateDevicePass tagging at export time. At runtime the C++ backend From c6bfd22d2b794def5c675b37e9b0ceda0dec25b9 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 16:57:40 -0700 Subject: [PATCH 8/9] [2.15] say fold_shared_kv_attention does not apply under torch.compile on PyTorch 2.15 #4921 came to this branch with a test that compiles a multi-query attention model with torch.compile and expects the new pass to fold its two matmuls. It fails here in every dynamo-lowering CI job: AssertionError: 'Folded 2 matmul(s)' not found in '... found nothing to fold in this graph ...' The code is the same as on main, where the test passes. The difference is PyTorch. With 2.15, the graph that torch.compile hands to Torch-TensorRT has each matmul already split into expand, reshape and bmm, so there is no matmul left to fold. With the 2.16 nightly the same graph still has aten.matmul. torch.export keeps aten.matmul on both, so the pass works there and its other tests pass. This branch requires PyTorch 2.15, so here the option never has an effect under torch.compile. The docs, docstrings and the pass's log message now say that instead of naming only use_distributed_mode_trace, and the test is skipped below PyTorch 2.16. Reproduced on Linux x86-64 with an A100 and TensorRT 11.4, from this branch's Python sources: the test file gives 1 failed, 21 passed with torch 2.15.0+cu130, and 22 passed with torch 2.16.0.dev20261028+cu130. After this change: 21 passed, 1 skipped with 2.15, and 22 passed with the nightly. --- .../compilation/compilation_settings.rst | 5 +++-- py/torch_tensorrt/dynamo/_compiler.py | 6 +++--- py/torch_tensorrt/dynamo/_settings.py | 2 +- .../lowering/passes/fold_shared_kv_attention.py | 14 ++++++-------- .../lowering/test_fold_shared_kv_attention.py | 6 ++++++ 5 files changed, 19 insertions(+), 14 deletions(-) diff --git a/docsrc/user_guide/compilation/compilation_settings.rst b/docsrc/user_guide/compilation/compilation_settings.rst index b3474cc4ef1..1c4af23f5e8 100644 --- a/docsrc/user_guide/compilation/compilation_settings.rst +++ b/docsrc/user_guide/compilation/compilation_settings.rst @@ -351,8 +351,9 @@ Graph Partitioning slower depending on the GPU, so measure it. In FP16 the result can round differently; ``use_fp32_acc=True`` reduces that. Only static shapes are rewritten, and ``scaled_dot_product_attention`` also needs - ``decompose_attention=True``. The setting has no effect under ``torch.compile`` - with ``use_distributed_mode_trace``. + ``decompose_attention=True``. The setting has no effect under ``torch.compile``, + because on PyTorch 2.15 ``torch.compile`` splits each matrix product into + ``expand``, ``reshape`` and ``bmm`` before this rewrite runs. ---- diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index ab05d926fa1..4fbb03f5741 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -227,7 +227,7 @@ def cross_compile_for_windows( foldable. Supported IDs: ``"attention_mask_arange"``. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. - fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. The setting has no effect under torch.compile with use_distributed_mode_trace. Default is False. + fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. The setting has no effect under torch.compile, because on PyTorch 2.15 torch.compile splits each matrix product into expand, reshape and bmm before this rewrite runs. Default is False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. **kwargs: Any, Returns: @@ -638,7 +638,7 @@ def compile( foldable. Supported IDs: ``"attention_mask_arange"``. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. - fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. The setting has no effect under torch.compile with use_distributed_mode_trace. Default is False. + fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. The setting has no effect under torch.compile, because on PyTorch 2.15 torch.compile splits each matrix product into expand, reshape and bmm before this rewrite runs. Default is False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. build_route (Optional[str]): TensorRT Global Performance Tuning build route string (space-separated "-knob=value" tokens). None uses the default route. This feature requires TensorRT with Global Performance Tuning enabled; currently unavailable on TensorRT-RTX / Windows. tune_build_routes (Optional[str]): Build-route expression for an autotuning sweep that runs each candidate in a spawned child process (e.g. "-slice_fusion=[on|off] -kgen:codegen:cuda_tile=[0|1|2|3]"). None disables tuning. @@ -2090,7 +2090,7 @@ def convert_exported_program_to_serialized_trt_engine( foldable. Supported IDs: ``"attention_mask_arange"``. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. - fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. The setting has no effect under torch.compile with use_distributed_mode_trace. Default is False. + fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. The setting has no effect under torch.compile, because on PyTorch 2.15 torch.compile splits each matrix product into expand, reshape and bmm before this rewrite runs. Default is False. **kwargs: Any, Returns: bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py index 045d0ff9aaf..ab0288afb28 100644 --- a/py/torch_tensorrt/dynamo/_settings.py +++ b/py/torch_tensorrt/dynamo/_settings.py @@ -163,7 +163,7 @@ class CompilationSettings: foldable. Supported IDs: ``"attention_mask_arange"``. Default is empty. attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False. - fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. The setting has no effect under torch.compile with use_distributed_mode_trace. Default is False. + fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. The setting has no effect under torch.compile, because on PyTorch 2.15 torch.compile splits each matrix product into expand, reshape and bmm before this rewrite runs. Default is False. fallback_data_dependent_ops (bool): If True, operators whose converters require a TensorRT output allocator (i.e. data-dependent output shapes, such as nonzero) are added to torch_executed_ops and run in PyTorch instead of being lowered into a TensorRT engine. This is useful when targeting runtimes that cannot consume a TensorRT output allocator. Default is False. build_route (Optional[str]): TensorRT Global Performance Tuning build route string (space-separated "-knob=value" tokens). None uses the default route. This feature requires TensorRT with Global Performance Tuning enabled; currently unavailable on TensorRT-RTX / Windows. tune_build_routes (Optional[str]): Build-route expression for an autotuning sweep that runs each candidate in a spawned child process (e.g. "-slice_fusion=[on|off] -kgen:codegen:cuda_tile=[0|1|2|3]"). None disables tuning. diff --git a/py/torch_tensorrt/dynamo/lowering/passes/fold_shared_kv_attention.py b/py/torch_tensorrt/dynamo/lowering/passes/fold_shared_kv_attention.py index d56d858856a..8698bc489b3 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/fold_shared_kv_attention.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/fold_shared_kv_attention.py @@ -17,9 +17,8 @@ Only matmuls with static shapes are rewritten. ``scaled_dot_product_attention`` becomes matmuls only with ``decompose_attention``. Under ``torch.compile``, -``use_distributed_mode_trace`` lowers matmuls to ``bmm`` before this pass, and -``torch.compile`` turns it on by itself for multi-process ``torch.distributed`` -jobs and for graphs that use vmap. +PyTorch 2.15 lowers matmuls to ``bmm`` before this pass, so it does not apply +there. """ import logging @@ -159,10 +158,9 @@ def fold_shared_kv_attention( "fold_shared_kv_attention is on but found nothing to fold in this graph. " "It needs exactly one key/value head, " "static shapes, and decompose_attention=True for " - "scaled_dot_product_attention. It does not apply under torch.compile " - "with use_distributed_mode_trace, which torch.compile also turns on by " - "itself for multi-process torch.distributed jobs and for graphs that use " - "vmap. This is logged once per graph, so a torch.cond branch or a graph " - "split off by a graph break can log it even when other attention folded." + "scaled_dot_product_attention. It does not apply under torch.compile, " + "which on PyTorch 2.15 lowers matmuls to bmm first. This is logged once " + "per graph, so a torch.cond branch can log it even when other attention " + "folded." ) return gm diff --git a/tests/py/dynamo/lowering/test_fold_shared_kv_attention.py b/tests/py/dynamo/lowering/test_fold_shared_kv_attention.py index c973f778a17..d35adbe0f1d 100644 --- a/tests/py/dynamo/lowering/test_fold_shared_kv_attention.py +++ b/tests/py/dynamo/lowering/test_fold_shared_kv_attention.py @@ -6,7 +6,9 @@ import torch import torch.nn.functional as F import torch_tensorrt +from packaging import version from parameterized import parameterized +from torch_tensorrt._utils import sanitized_torch_version from torch_tensorrt.dynamo import CompilationSettings from torch_tensorrt.dynamo.lowering import get_decompositions, post_lowering from torch_tensorrt.dynamo.lowering.passes.fold_shared_kv_attention import ( @@ -295,6 +297,10 @@ def test_compiled_multi_query_attention_matches_pytorch(self) -> None: compiled(*inputs).float(), expected, atol=1e-2, rtol=1e-2 ) + @unittest.skipIf( + version.parse(sanitized_torch_version()) < version.parse("2.16.dev"), + "torch.compile on PyTorch 2.15 lowers matmul to bmm before the pass runs", + ) def test_torch_compile_folds_multi_query_attention(self) -> None: inputs = _inputs(heads=8, kv_heads=1, dtype=torch.float16) module = RepeatKVAttention(8, 1, 16).eval() From 1d65d2ef3a5deba42a655682a7a9bac4e4f1fce9 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 28 Oct 2026 13:37:06 -0700 Subject: [PATCH 9/9] fix(dynamo): decompose attention TensorRT cannot build, and report why a build failed (#4928) (cherry picked from commit 6d97b2b9ebdcb8e3e0134f7ea31b807b6d2b8736) --- .../dynamo/conversion/_TRTInterpreter.py | 40 ++-- .../dynamo/conversion/aten_ops_converters.py | 90 ++++++++ .../passes/decompose_unsupported_attention.py | 36 +++- py/torch_tensorrt/logging.py | 28 ++- tests/ci/suites.py | 2 +- .../conversion/test_trt_build_errors.py | 60 ++++++ .../test_decompose_unsupported_attention.py | 201 ++++++++++++++++++ tools/ci-dashboard/ci_dashboard.py | 3 +- 8 files changed, 429 insertions(+), 31 deletions(-) create mode 100644 tests/py/dynamo/conversion/test_trt_build_errors.py diff --git a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py index a7c56e115cb..62295d647bd 100644 --- a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py +++ b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py @@ -580,25 +580,31 @@ def run( builder_config, self.compilation_settings.timing_cache_path ) - if ( - ENABLED_FEATURES.tensorrt_rtx - or self.compilation_settings.version_compatible - ): - # TODO: When TRT-RTX matures, change it to build_engine_with_config - serialized_engine = self.builder.build_serialized_network( - self.ctx.net, builder_config - ) - if serialized_engine is None: - raise RuntimeError( - "TensorRT build_serialized_network returned None; engine build failed." + with TRT_LOGGER.capture_errors() as build_errors: + if ( + ENABLED_FEATURES.tensorrt_rtx + or self.compilation_settings.version_compatible + ): + # TODO: When TRT-RTX matures, change it to build_engine_with_config + serialized_engine = self.builder.build_serialized_network( + self.ctx.net, builder_config ) - runtime = trt.Runtime(TRT_LOGGER) - cuda_engine = runtime.deserialize_cuda_engine(serialized_engine) - else: - cuda_engine = self.builder.build_engine_with_config( - self.ctx.net, builder_config + cuda_engine = ( + None + if serialized_engine is None + else trt.Runtime(TRT_LOGGER).deserialize_cuda_engine( + serialized_engine + ) + ) + else: + cuda_engine = self.builder.build_engine_with_config( + self.ctx.net, builder_config + ) + if cuda_engine is None: + raise RuntimeError( + "TensorRT could not build the engine" + + (": " + " | ".join(build_errors) if build_errors else ".") ) - assert cuda_engine _LOGGER.debug( f"CPU memory usage after engine building: {get_cpu_memory_usage()} MB" diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 3278cb73635..7cda4e4dc7a 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -5194,12 +5194,102 @@ def is_unsupported(operand: Argument) -> bool: return True +# Ops that only reshape, copy or cast a tensor. TensorRT can fold them away, so tensors reached through them +# from one producer may become one tensor in the network. +_PASS_THROUGH_OPS = { + torch.ops.aten.view.default, + torch.ops.aten._unsafe_view.default, + torch.ops.aten.reshape.default, + torch.ops.aten._reshape_copy.default, + torch.ops.aten.transpose.int, + torch.ops.aten.permute.default, + torch.ops.aten.expand.default, + torch.ops.aten.unsqueeze.default, + torch.ops.aten.squeeze.dim, + torch.ops.aten.squeeze.dims, + torch.ops.aten.clone.default, + torch.ops.aten._to_copy.default, +} + + +def _pass_through_source(arg: Argument) -> Argument: + while isinstance(arg, Node) and arg.target in _PASS_THROUGH_OPS: + arg = arg.args[0] + return arg + + +def _is_bf16(arg: Argument) -> bool: + val = arg.meta.get("val") if isinstance(arg, Node) else None + return getattr(val, "dtype", None) == torch.bfloat16 + + +def _shares_bf16_query_with_another_attention(node: Node) -> bool: + if not _is_bf16(node.args[0]): + return False + query = _pass_through_source(node.args[0]) + return any( + other is not node + and other.target == torch.ops.aten.scaled_dot_product_attention.default + and _is_bf16(other.args[0]) + and _pass_through_source(other.args[0]) is query + for other in node.graph.nodes + ) + + +def _output_quantized_to_fp8(node: Node) -> bool: + # Imported here: the lowering passes import this module. + from torch_tensorrt.dynamo.lowering.passes.annotate_fp8_sdpa import ( + _is_fp8_quantize_op, + ) + + users = list(node.users) + while users: + user = users.pop() + if _is_fp8_quantize_op(user): + return True + if user.target in _PASS_THROUGH_OPS: + users.extend(user.users) + return False + + +def _inputs_quantized_to_fp8(node: Node) -> bool: + from torch_tensorrt.dynamo.lowering.passes.annotate_fp8_sdpa import ( + _is_fp8_quantize_op, + ) + + # Directly, as annotate_fp8_sdpa requires for the fused FP8 attention. + return all( + isinstance(arg, Node) and _is_fp8_quantize_op(arg) for arg in node.args[:3] + ) + + def scaled_dot_product_attention_validator( node: Node, settings: Optional[CompilationSettings] = None ) -> bool: if not attention_capability_validator(node, settings): return False + # TensorRT 11.3 and 11.4 crash or fail while building two attention layers that read the same + # bfloat16 query (float16 builds). Seen on a Jetson AGX Thor and an H100, and declined on every + # version, because the decomposition is always correct. + if _shares_bf16_query_with_another_attention(node): + _LOGGER.debug( + "Decomposing attention %s: it reads the same bfloat16 query as another attention, which " + "TensorRT's attention layer cannot build.", + node.name, + ) + return False + # TensorRT 11.4 fails to build an attention layer whose output is quantized to FP8, unless its query, + # key and value are FP8 too (the fused FP8 attention that annotate_fp8_sdpa sets up). Seen on a Jetson + # AGX Thor and an H100, and declined on every version for the same reason. + if _output_quantized_to_fp8(node) and not _inputs_quantized_to_fp8(node): + _LOGGER.debug( + "Decomposing attention %s: its output is quantized to FP8, which TensorRT's attention layer " + "cannot build unless its inputs are FP8 too.", + node.name, + ) + return False + attn_mask = args_bounds_check(node.args, 3, None) is_causal = args_bounds_check(node.args, 5, False) if is_causal and attn_mask is not None: diff --git a/py/torch_tensorrt/dynamo/lowering/passes/decompose_unsupported_attention.py b/py/torch_tensorrt/dynamo/lowering/passes/decompose_unsupported_attention.py index af367a4b9f2..6bee8c70cbb 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/decompose_unsupported_attention.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/decompose_unsupported_attention.py @@ -9,6 +9,7 @@ scaled_dot_product_attention_validator, ) from torch_tensorrt.dynamo.lowering._decompositions import ( + get_decompositions, scaled_dot_product_attention_decomposition, ) from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( @@ -47,18 +48,29 @@ def _inline_traced_decomp( gm: torch.fx.GraphModule, node: torch.fx.Node, decomp_fn: Callable[..., Any], - use_fp32_acc: bool, + settings: CompilationSettings, ) -> None: """Replace ``node`` with the aten subgraph from tracing ``decomp_fn``.""" example_args = tuple(_example_val(a) for a in node.args) example_kwargs = {k: _example_val(v) for k, v in node.kwargs.items()} - example_kwargs["use_fp32_acc"] = use_fp32_acc + example_kwargs["use_fp32_acc"] = settings.use_fp32_acc def wrapped(*args: Any) -> Any: return decomp_fn(*args, **example_kwargs) - traced = make_fx(wrapped, tracing_mode="fake")(*example_args) + # Lowered with the same table as the rest of the graph: this pass runs after decomposition, and ops such + # as aten.transpose.int have no converter of their own, so they would be left to PyTorch. + traced = make_fx( + wrapped, + decomposition_table=get_decompositions( + settings.enable_experimental_decompositions, + settings.decompose_attention, + settings.use_distributed_mode_trace, + settings.use_fp32_acc, + ), + tracing_mode="fake", + )(*example_args) arg_nodes = list(node.args) env: Dict[torch.fx.Node, Any] = {} @@ -107,8 +119,10 @@ def decompose_unsupported_attention( if settings.decompose_attention: return gm - changed = False - for node in list(gm.graph.nodes): + # Decide for every node before changing the graph: whether a node is supported can depend on the other + # attention nodes, such as two that read the same query. + declined = [] + for node in gm.graph.nodes: if node.op != "call_function" or node.target not in _ATTENTION_FALLBACKS: continue @@ -118,16 +132,16 @@ def decompose_unsupported_attention( continue validator, decomp = _ATTENTION_FALLBACKS[node.target] - if validator(node, settings): - continue + if not validator(node, settings): + declined.append((node, decomp)) + for node, decomp in declined: logger.info( - "Attention converter declined %s; decomposing that node in place " - "(native IAttention kept for other attention nodes)", + "Attention converter declined %s; decomposing that node in place", node.name, ) - _inline_traced_decomp(gm, node, decomp, settings.use_fp32_acc) - changed = True + _inline_traced_decomp(gm, node, decomp, settings) + changed = bool(declined) if changed: gm = clean_up_graph_after_modifications(gm) diff --git a/py/torch_tensorrt/logging.py b/py/torch_tensorrt/logging.py index 325cdba6f1b..2c76ff96815 100644 --- a/py/torch_tensorrt/logging.py +++ b/py/torch_tensorrt/logging.py @@ -2,7 +2,9 @@ # SPDX-License-Identifier: BSD-3-Clause import logging -from typing import Any +import threading +from contextlib import contextmanager +from typing import Any, Iterator, List import tensorrt as trt import torch @@ -18,8 +20,32 @@ class _TRTLogger(trt.ILogger): # type: ignore[misc] def __init__(self) -> None: trt.ILogger.__init__(self) + self._error_captures: List[List[str]] = [] + self._lock = threading.Lock() + + @contextmanager + def capture_errors(self) -> Iterator[List[str]]: + """Collect the errors TensorRT reports, from any thread, while the block runs, so a failure can name them.""" + errors: List[str] = [] + with self._lock: + self._error_captures.append(errors) + try: + yield errors + finally: + with self._lock: + # By identity: two open captures can hold equal lists, for example both empty. + self._error_captures = [ + e for e in self._error_captures if e is not errors + ] def log(self, severity: trt.ILogger.Severity, msg: str) -> None: + if severity in ( + trt.ILogger.Severity.INTERNAL_ERROR, + trt.ILogger.Severity.ERROR, + ): + with self._lock: + for errors in self._error_captures: + errors.append(msg) # TODO: Move to match once py39 reaches EoL if severity == trt.ILogger.Severity.INTERNAL_ERROR: _LOGGER.critical(msg) diff --git a/tests/ci/suites.py b/tests/ci/suites.py index 5793d0d5916..21a69306f04 100644 --- a/tests/ci/suites.py +++ b/tests/ci/suites.py @@ -50,7 +50,7 @@ # worker builds its own TensorRT engine, so N workers ≈ N× peak GPU memory. # `-n auto` (= CPU count, up to ~48 on big runners) is a footgun here — it OOMs # the model/runtime suites on real 24 GB runners, and an OOM'd engine build fails -# `assert cuda_engine` and then *reruns*, doubling wall-clock. So workers are +# ("TensorRT could not build the engine") and then *reruns*, doubling wall-clock. So workers are # tiered by worst-case engine size, not by core count: _LIGHT = "8" # small single-op / graph tests (conversion, partitioning, hlo) _HEAVY = "4" # full compile+execute suites (lowering, runtime) — OOM'd at -n 8 diff --git a/tests/py/dynamo/conversion/test_trt_build_errors.py b/tests/py/dynamo/conversion/test_trt_build_errors.py new file mode 100644 index 00000000000..4a957bb28b0 --- /dev/null +++ b/tests/py/dynamo/conversion/test_trt_build_errors.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +"""A TensorRT build that fails says why, instead of raising a bare AssertionError.""" + +from unittest import mock + +import pytest +import tensorrt as trt +import torch +from torch_tensorrt.logging import TRT_LOGGER + + +def test_capture_errors_collects_only_errors_reported_inside_it(): + TRT_LOGGER.log(trt.ILogger.Severity.ERROR, "before") + with TRT_LOGGER.capture_errors() as outer: + TRT_LOGGER.log(trt.ILogger.Severity.WARNING, "a warning") + with TRT_LOGGER.capture_errors() as inner: + TRT_LOGGER.log(trt.ILogger.Severity.ERROR, "inside") + TRT_LOGGER.log(trt.ILogger.Severity.ERROR, "outer only") + TRT_LOGGER.log(trt.ILogger.Severity.ERROR, "after") + assert outer == ["inside", "outer only"] and inner == ["inside"] + # Two captures that hold equal lists when the first one ends. + with TRT_LOGGER.capture_errors() as first: + with TRT_LOGGER.capture_errors() as second: + pass + TRT_LOGGER.log(trt.ILogger.Severity.ERROR, "first only") + assert first == ["first only"] and second == [] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="builds a TensorRT network") +@pytest.mark.parametrize("version_compatible", [False, True]) +def test_a_refused_build_raises_tensorrt_s_own_reason(version_compatible): + import torch_tensorrt + + def refuse(builder, network, config): + TRT_LOGGER.log( + trt.ILogger.Severity.ERROR, + "Networks with FP8 Q/DQ layers require hardware with FP8 support.", + ) + return None + + model = torch.nn.Linear(8, 8).eval().cuda() + inputs = (torch.randn(2, 8, device="cuda"),) + program = torch.export.export(model, inputs) + # Both build calls: version-compatible and TensorRT-RTX builds use the second, and an RTX + # Builder may not have the first. + with ( + mock.patch.object(trt.Builder, "build_engine_with_config", refuse, create=True), + mock.patch.object(trt.Builder, "build_serialized_network", refuse), + ): + with pytest.raises(RuntimeError, match="require hardware with FP8 support"): + torch_tensorrt.dynamo.compile( + program, + arg_inputs=list(inputs), + min_block_size=1, + version_compatible=version_compatible, + cache_built_engines=False, + reuse_cached_engines=False, + ) diff --git a/tests/py/dynamo/lowering/test_decompose_unsupported_attention.py b/tests/py/dynamo/lowering/test_decompose_unsupported_attention.py index 39014ff17e4..17863e9b4e7 100644 --- a/tests/py/dynamo/lowering/test_decompose_unsupported_attention.py +++ b/tests/py/dynamo/lowering/test_decompose_unsupported_attention.py @@ -1,3 +1,6 @@ +import importlib.util +import platform +import unittest from unittest import mock import torch @@ -125,6 +128,204 @@ def forward(self, q, k, v): f"SDPA should have been decomposed for MLA shapes: {unexpected_ops_seen}", ) + @unittest.skipIf(not torch.cuda.is_available(), "builds a TensorRT engine") + def test_decomposed_attention_runs_in_tensorrt(self): + # The per-node fallback once left aten.transpose.int and aten._unsafe_view to + # PyTorch, splitting the model across engines. + import torch_tensorrt + + class MLA(nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention(q, k, v) + + inputs = ( + torch.randn(1, 2, 16, 64, device="cuda", dtype=torch.float16), + torch.randn(1, 2, 16, 64, device="cuda", dtype=torch.float16), + torch.randn(1, 2, 16, 32, device="cuda", dtype=torch.float16), + ) + trt_model = torch_tensorrt.dynamo.compile( + torch.export.export(MLA(), inputs), + arg_inputs=list(inputs), + min_block_size=1, + decompose_attention=False, + cache_built_engines=False, + reuse_cached_engines=False, + ) + torch.testing.assert_close( + trt_model(*inputs), MLA()(*inputs), rtol=1e-2, atol=1e-2 + ) + self.assertEqual( + [name for name, _ in trt_model.named_children()], ["_run_on_acc_0"] + ) + + def test_attentions_sharing_a_bf16_query_are_decomposed(self): + # Wan2.1's image-to-video cross-attention reads one query against text and image keys. + class TwoAttentions(nn.Module): + def forward(self, q, text, image): + q = q.transpose(1, 2) + return torch.nn.functional.scaled_dot_product_attention( + q, text, text + ) + torch.nn.functional.scaled_dot_product_attention(q, image, image) + + sdpa = torch.ops.aten.scaled_dot_product_attention.default + settings = CompilationSettings(min_block_size=1, decompose_attention=False) + for dtype, remaining in ((torch.bfloat16, 0), (torch.float16, 2)): + with self.subTest(dtype=dtype): + inputs = ( + torch.randn(1, 6, 2, 8, dtype=dtype), + torch.randn(1, 2, 5, 8, dtype=dtype), + torch.randn(1, 2, 3, 8, dtype=dtype), + ) + gm = torch.export.export(TwoAttentions(), inputs).module() + gm = decompose_unsupported_attention(gm, settings) + self.assertEqual( + sum(n.target == sdpa for n in gm.graph.nodes), remaining + ) + if dtype == torch.bfloat16: + torch.testing.assert_close( + gm(*inputs), TwoAttentions()(*inputs), rtol=1e-2, atol=1e-2 + ) + + def test_attentions_reading_one_tensor_in_different_dtypes_are_kept(self): + # Only one of the two queries is bfloat16, so they cannot become one tensor in TensorRT. + class TwoDtypes(nn.Module): + def forward(self, q, k16, kbf16): + cast = torch.ops.aten._to_copy.default + a = torch.nn.functional.scaled_dot_product_attention( + cast(q, dtype=torch.float16), k16, k16 + ) + b = torch.nn.functional.scaled_dot_product_attention( + cast(q, dtype=torch.bfloat16), kbf16, kbf16 + ) + return a.float() + b.float() + + inputs = ( + torch.randn(1, 2, 6, 8), + torch.randn(1, 2, 5, 8, dtype=torch.float16), + torch.randn(1, 2, 5, 8, dtype=torch.bfloat16), + ) + gm = torch.export.export(TwoDtypes(), inputs).module() + gm = decompose_unsupported_attention( + gm, CompilationSettings(min_block_size=1, decompose_attention=False) + ) + sdpa = torch.ops.aten.scaled_dot_product_attention.default + self.assertEqual(sum(n.target == sdpa for n in gm.graph.nodes), 2) + + @unittest.skipIf( + importlib.util.find_spec("modelopt") is None or platform.system() != "Linux", + "needs NVIDIA Model Optimizer's quantize op, which is only supported on Linux", + ) + def test_attention_whose_output_is_quantized_to_fp8_is_decomposed(self): + import modelopt.torch.quantization # noqa: F401 # registers tensorrt.quantize_op + + quantize = torch.ops.tensorrt.quantize_op.default + + def graph(fp8_inputs: bool) -> torch.fx.GraphModule: + g = torch.fx.Graph() + q, k, v, amax = (g.placeholder(n) for n in ("q", "k", "v", "amax")) + if fp8_inputs: + q, k, v = ( + g.call_function(quantize, (t, amax, 8, 4, False, False)) + for t in (q, k, v) + ) + out = g.call_function( + torch.ops.aten.scaled_dot_product_attention.default, (q, k, v) + ) + out = g.call_function(torch.ops.aten.transpose.int, (out, 1, 2)) + g.output(g.call_function(quantize, (out, amax, 8, 4, False, False))) + return torch.fx.GraphModule({}, g) + + for fp8_inputs, supported in ((False, False), (True, True)): + with self.subTest(fp8_inputs=fp8_inputs): + sdpa = next( + n + for n in graph(fp8_inputs).graph.nodes + if n.target == torch.ops.aten.scaled_dot_product_attention.default + ) + self.assertEqual( + scaled_dot_product_attention_validator(sdpa, None), supported + ) + + @unittest.skipIf(not torch.cuda.is_available(), "builds a TensorRT engine") + def test_attentions_sharing_a_query_build(self): + # TensorRT crashed or failed building these before the attentions were decomposed. + import torch_tensorrt + + class TwoAttentions(nn.Module): + def __init__(self, combine): + super().__init__() + self.combine = combine + + def forward(self, q, text, image): + return self.combine( + torch.nn.functional.scaled_dot_product_attention(q, text, text), + torch.nn.functional.scaled_dot_product_attention(q, image, image), + ) + + inputs = tuple( + torch.randn(1, 8, n, 128, device="cuda", dtype=torch.bfloat16) + for n in (256, 512, 257) + ) + for combine in (torch.add, lambda a, b: torch.cat([a, b], dim=-1)): + with self.subTest(combine=combine): + model = TwoAttentions(combine) + trt_model = torch_tensorrt.dynamo.compile( + torch.export.export(model, inputs), + arg_inputs=list(inputs), + min_block_size=1, + decompose_attention=False, + cache_built_engines=False, + reuse_cached_engines=False, + ) + torch.testing.assert_close( + trt_model(*inputs), model(*inputs), rtol=1e-2, atol=1e-2 + ) + # The decomposed attentions run in TensorRT, not in PyTorch between engines. + self.assertEqual( + [name for name, _ in trt_model.named_children()], ["_run_on_acc_0"] + ) + + @unittest.skipIf( + not torch.cuda.is_available() + or torch.cuda.get_device_capability() < (8, 9) + or importlib.util.find_spec("modelopt") is None + or platform.system() != "Linux", + "needs an FP8 GPU and NVIDIA Model Optimizer's quantize op (Linux only)", + ) + def test_attention_whose_output_is_quantized_to_fp8_builds(self): + # TensorRT 11.4 failed to build this before the attention was decomposed. + import modelopt.torch.quantization # noqa: F401 # registers tensorrt.quantize_op + import torch_tensorrt + + class AttentionThenFP8(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("amax", torch.tensor(4.0)) + self.out = nn.Linear(256, 256) + + def forward(self, q, k, v): + x = torch.nn.functional.scaled_dot_product_attention(q, k, v) + x = x.transpose(1, 2).reshape(1, 128, 256) + x = torch.ops.tensorrt.quantize_op(x, self.amax, 8, 4, False, False) + return self.out(x) + + model = AttentionThenFP8().cuda().half().eval() + inputs = tuple( + torch.randn(1, 4, 128, 64, device="cuda", dtype=torch.float16) + for _ in range(3) + ) + trt_model = torch_tensorrt.dynamo.compile( + torch.export.export(model, inputs), + arg_inputs=list(inputs), + min_block_size=1, + decompose_attention=False, + cache_built_engines=False, + reuse_cached_engines=False, + ) + torch.testing.assert_close( + trt_model(*inputs), model(*inputs), rtol=0.1, atol=0.1 + ) + if __name__ == "__main__": run_tests() diff --git a/tools/ci-dashboard/ci_dashboard.py b/tools/ci-dashboard/ci_dashboard.py index 0769f3f94f5..d5417074e01 100755 --- a/tools/ci-dashboard/ci_dashboard.py +++ b/tools/ci-dashboard/ci_dashboard.py @@ -1109,7 +1109,7 @@ def classify(status, conclusion): # converter gap, or a real regression (numerical/assertion)? First pattern wins, # so order is most-specific → most-generic. `kind` drives the tag colour: # infra = "probably not your bug" · env/gap = setup/feature hole · bug = real. -# NOTE: an OOM often surfaces AS an AssertionError ("assert cuda_engine"); the OOM +# NOTE: an OOM often surfaces as a failed engine build ("TensorRT could not build the engine"); the OOM # signatures below are listed first on purpose so they win over the generic assert. _FAIL_CATS = [ ( @@ -1119,6 +1119,7 @@ def classify(status, conclusion): re.compile( r"OutOfMemory|out of memory|createCaskHardwareInfo|\bCask\b|" r"build_serialized_network returned None|assert cuda_engine|cuda_engine\b|" + r"TensorRT could not build the engine|" r"catchCudaError|cudaError|cuInit|defaultAllocator|no CUDA GPUs|" r"CUDA error|cudaMalloc|device-side assert", re.I,