diff --git a/docs/source/backends-overview.md b/docs/source/backends-overview.md index 1d626a04ba4..5e3d575ca4b 100644 --- a/docs/source/backends-overview.md +++ b/docs/source/backends-overview.md @@ -22,7 +22,7 @@ Backends are the bridge between your exported model and the hardware it runs on. |--------------------------------------------------------------|---------------|---------------|---------------------------------| | [XNNPACK](backends/xnnpack/xnnpack-overview.md) | All | CPU | General-purpose, fallback | | [CUDA](backends/cuda/cuda-overview.md) | Linux/Windows | GPU | NVIDIA GPU acceleration | -| [TensorRT](backends/tensorrt/tensorrt-overview.md) | Linux | GPU | NVIDIA GPU acceleration via TensorRT | +| [TensorRT](backends/tensorrt/tensorrt-overview.md) | Linux/Windows | GPU | NVIDIA GPU acceleration via TensorRT | | [Core ML](backends/coreml/coreml-overview.md) | iOS, macOS | NPU/GPU/CPU | Apple devices, high performance | | [MLX](/backends/mlx/mlx-overview.md) | iOS (experimental), macOS | GPU | Apple Silicon GPU (MLX) | | [Vulkan](backends/vulkan/vulkan-overview.md) | Android, Linux, Windows | GPU | Android devices (mature); Desktops (experimental) | diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md index ef57aca7ae5..db7d0adfb75 100644 --- a/docs/source/backends/cuda/cuda-overview.md +++ b/docs/source/backends/cuda/cuda-overview.md @@ -7,24 +7,29 @@ The CUDA backend is the ExecuTorch solution for running models on NVIDIA GPUs. I - **Optimized GPU Execution**: Uses AOTInductor to generate highly optimized CUDA kernels for model operators - **Triton Kernel Support**: Leverages Triton for GEMM (General Matrix Multiply), convolution, and SDPA (Scaled Dot-Product Attention) kernels. - **Quantization Support**: INT4 weight quantization with tile-packed format for improved performance and reduced memory footprint -- **Cross-Platform**: Supports both Linux and Windows platforms +- **Cross-Platform**: Runs on Linux, and on Windows from C++. Programs for Windows are exported on Linux. - **Multiple Model Support**: Works with various models including LLMs, vision-language models, and audio models ## Target Requirements Below are the requirements for running a CUDA-delegated ExecuTorch model: -- **Hardware**: NVIDIA GPU with CUDA compute capability -- **CUDA Toolkit**: CUDA 11.x or later (CUDA 12.x recommended) -- **Operating System**: Linux or Windows -- **Drivers**: PyTorch-Compatible NVIDIA GPU drivers installed +- **Hardware**: NVIDIA GPU with compute capability 8.0 (Ampere) or newer. This is the floor the + prebuilt packages are built for. +- **CUDA**: 13 for the prebuilt packages. On Linux they bring the CUDA runtime with them. On + Windows it comes from a CUDA 13 toolkit install. A build from source also supports CUDA 12.6. +- **Operating System**: Linux, or Windows from C++ for programs exported on Linux (see + [CUDA on Windows](../../using-executorch-cpp.md#cuda-on-windows)) +- **Drivers**: An NVIDIA driver that supports the CUDA version you use: CUDA 13 for the prebuilt + packages. ## Development Requirements To develop and export models using the CUDA backend: -- **Python**: Python 3.8+ +- **Python**: Python 3.11 to 3.15 - **PyTorch**: PyTorch with CUDA support +- **CUDA Toolkit**: export compiles kernels with `nvcc`, the CUDA compiler, so it needs a CUDA toolkit - **ExecuTorch**: Install ExecuTorch with CUDA backend support ## Using the CUDA Backend @@ -83,7 +88,7 @@ For a complete working example, see the [CUDA export script](https://github.com/ ## Runtime Integration -To run the model on device, use the standard ExecuTorch runtime APIs. See [Running on Device](getting-started.md#running-on-device) for more information. +To run the model on device, use the standard ExecuTorch runtime APIs. See [Running on Device](../../getting-started.md#running-on-device) for more information. When building from source, pass `-DEXECUTORCH_BUILD_CUDA=ON` when configuring the CMake build to compile the CUDA backend. @@ -118,7 +123,7 @@ exec_program = et_program.to_executorch() ```python # CPU tensors. The runtime copies them to the GPU and copies results back. -outputs = method([torch.randn(4, 64)]) +outputs = method.execute([torch.randn(4, 64)]) ``` This is the simplest option and the right default. The cost is a copy in each direction on @@ -140,24 +145,22 @@ exec_program = et_program.to_executorch( skip_d2h_for_method_outputs=True, ), enable_non_cpu_memory_planning=True, - # Required alongside the skips. Memory planning reserves a buffer for graph inputs by - # default, and the runtime fills a planned input by copying the caller's memory into that - # buffer, which reintroduces the copy you just asked to skip. On device memory that copy is - # a host memcpy into a device pointer, which is undefined. - # Outputs stay planned, so the runtime allocates them on the device and hands back a - # pointer to its own buffer. See the note below for when to unplan them instead. - memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + # Keeps the method from reserving its own input and output memory. + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ), ) ) ``` ```python # Now the tensors must already be on the GPU. Passing CPU tensors is a caller error. -outputs = method([torch.randn(4, 64, device="cuda")]) +outputs = method.execute([torch.randn(4, 64, device="cuda")]) ``` Outputs also stay on the device, which is what you want when the next stage of your pipeline is -also on the GPU. +also on the GPU. From Python, each output comes back in GPU memory that the next call does not +overwrite, unless you pass `clone_outputs=False`. ### Things worth knowing @@ -170,24 +173,17 @@ it to `False` while asking for either skip raises a `ValueError`, because copy i during device-aware memory planning. **Inputs must be unplanned**, via `MemoryPlanningPass(alloc_graph_input=False)`. Without it the -program reserves its own input buffer and the runtime copies the caller's memory into it, so the -copy comes back at run time. On device memory that copy is a host memcpy into a device pointer, -which is undefined, so the program crashes rather than returning a wrong answer. - -**Outputs have two valid shapes, and which one you want depends on the language.** - -Leave outputs planned, which is the default, when the caller cannot supply device memory itself. -The planned buffer is allocated on the device because `enable_non_cpu_memory_planning` is on, and -`skip_d2h_for_method_outputs` means nothing copies it back, so the result stays on the device and -the caller receives a pointer to the runtime's own buffer. That buffer belongs to the method and is -reused by the next execution, so read or copy the result before running again. This is the shape -every Python caller wants, and it is what the exporters in this repository use, for example -`examples/models/gemma4_31b/export.py`. - -Set `alloc_graph_output=False` only when the caller supplies the output memory itself, which today -means a C++ caller handing the runtime a device tensor through `Module::set_output`. The Python -bindings allocate host memory for an unplanned output, so a device resident method exported this -way cannot run from Python. +program reserves its own input buffer on the device, and the runtime refuses to copy the caller's +tensor into it. The call fails with error `0x10` (`NotSupported`), and the log says +`Planned-buffer copy needs host memory on both sides`. + +**Leave outputs unplanned too**, via `alloc_graph_output=False`. Then each result lands in memory +you own, with no extra copy. A C++ caller must give each unplanned output a buffer once with +`Module::set_output`, before the first run: GPU memory when the output copy is skipped, host +memory otherwise. Pass each output's index to `set_output`; the C++ example below sets output 0. If +you keep outputs planned instead, they live in the method's own memory and the next run overwrites +them. Python's `execute()` then copies each output on every call, unless you pass +`clone_outputs=False`. **Per-method selection is on the outer config, not on the flags.** Pass a dict of `PropagateDeviceConfig` keyed by method name: @@ -221,6 +217,11 @@ device right after the scalar type. Nothing is copied: ```cpp auto input = from_blob(device_data, {rows, columns}, ScalarType::Float, DeviceType::CUDA); +auto output = from_blob(device_result, {rows, out_columns}, ScalarType::Float, DeviceType::CUDA); +// Once, before the first run. Every run then writes its result to device_result. +if (module.set_output(output) != Error::Ok) { + return 1; +} auto result = module.forward(input); ``` @@ -235,6 +236,7 @@ across, which costs one transfer: ```cpp auto host_input = make_tensor_ptr({rows, columns}, std::move(data)); auto input = clone_tensor_ptr_to(host_input, DeviceType::CUDA); +// Set the output with module.set_output first, as above. auto result = module.forward(input); ``` diff --git a/docs/source/backends/tensorrt/tensorrt-overview.md b/docs/source/backends/tensorrt/tensorrt-overview.md index d85b427d688..9b4ab7b43ca 100644 --- a/docs/source/backends/tensorrt/tensorrt-overview.md +++ b/docs/source/backends/tensorrt/tensorrt-overview.md @@ -1,87 +1,98 @@ # TensorRT Backend The TensorRT backend runs ExecuTorch models on NVIDIA GPUs with -[TensorRT](https://developer.nvidia.com/tensorrt), NVIDIA's inference library. At export -time the model graph is compiled into TensorRT engines. The engines are stored inside the -`.pte` file, and at run time the delegate hands each one back to TensorRT. +[TensorRT](https://developer.nvidia.com/tensorrt), NVIDIA's inference library. At export time, +TensorRT compiles the parts of the model it supports into engines (programs optimized for your +GPU), and the engines are stored in the `.pte` file. At run time, the TensorRT delegate hands each +engine back to TensorRT. The delegate is built and shipped by the [Torch-TensorRT](https://github.com/pytorch/TensorRT) -project, so it comes from its wheel and not from the ExecuTorch wheel. Everything after -installation is the normal ExecuTorch API: the same `.pte` file, the same `Module` class, the -same runtime. +project, not by ExecuTorch. You export with Torch-TensorRT, and you run with the normal ExecuTorch +runtime and its usual APIs. The +[Torch-TensorRT ExecuTorch guide](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html) +covers every option in detail. + +This page describes Torch-TensorRT 2.15 and newer. Features marked "newer than 2.15" need a later +Torch-TensorRT build, such as a +[nightly build](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#nightly-builds). ## Features -- **TensorRT engines inside a `.pte`.** One file holds the program and its engines, so the - runtime loads a model the usual way. -- **Works with the CUDA backend in one program.** Operators TensorRT cannot convert can be - compiled by the ExecuTorch [CUDA backend](../cuda/cuda-overview.md) instead of falling back - to the CPU. Both delegates live in the same `.pte` and run in the same method. -- **GPU resident inputs and outputs.** The copies around the method boundary can be turned - off, so a pipeline that already holds its data on the GPU does not pay for them. -- **CUDA graph replay.** An engine can be recorded once and replayed with a single launch, - which helps models that launch many short kernels. This needs a Torch-TensorRT build newer - than the 2026-10-06 nightly, see [Feature availability](#feature-availability). -- **Shared activation scratch.** Many engines in one program can share one scratch buffer - instead of each keeping its own. -- **Shared engines.** Loading the same program twice keeps one copy of the engine weights in - GPU memory. This also needs a build newer than the 2026-10-06 nightly. -- **Caller-owned CUDA stream.** The application can give every delegate the same stream, - including a CUDA green context stream, so the model stays inside one part of the GPU. +- **TensorRT engines inside a `.pte`.** One file holds the program and its engines. +- **Works with the CUDA backend in one program.** Operators that TensorRT cannot compile can go to + the ExecuTorch [CUDA backend](../cuda/cuda-overview.md) instead of the CPU, so the whole model + stays on the GPU. +- **Inputs and outputs that stay on the GPU.** The copies at the edges of the program can be + turned off, for a pipeline that already has its data on the GPU. +- **Your own CUDA stream.** The stream can also belong to a CUDA green context, which keeps the + model inside a fixed part of the GPU. TensorRT warns that this setup can crash, as + [Performance Options](#performance-options) explains. +- **CUDA graph replay.** The GPU work of an engine is recorded once and then replayed with a single + launch, which saves CPU time on every call. +- **Shared activation scratch.** Many engines can share one GPU buffer for their temporary results + (activations), which saves memory. +- **Shared engines.** Loading the same program twice keeps one copy of the engine weights in GPU + memory. +- **Weight streaming.** An engine can keep part of its weights in CPU memory. +- **Zero-copy KV cache.** The engine of a large language model (LLM) can update its key-value (KV) + cache in place. +- **Newer than 2.15:** engines in a separate data file, and optimization profiles (tunings for + different input sizes) chosen per call. ## Target Requirements -- **Hardware**: NVIDIA GPU. -- **Operating system**: Linux, on x86_64 and on aarch64. -- **CUDA**: 13.x. The delegate links CUDA 13, so a CUDA 12 build is not supported. -- **Drivers**: an NVIDIA driver that matches the CUDA version. +- **Hardware**: an NVIDIA GPU that CUDA 13 supports. The prebuilt ExecuTorch CUDA packages need + compute capability 8.0 (Ampere) or newer. +- **Operating system**: Linux on x86-64 with glibc (the GNU C library) 2.28 or newer, Linux on + aarch64 with glibc 2.35 or newer, or Windows on x86-64. +- **CUDA**: 13. The prebuilt delegate links the CUDA 13 runtime. +- **Drivers**: an NVIDIA driver that supports CUDA 13. + +On Linux, TensorRT and the CUDA runtime come as dependencies of the Python packages, so you do not +need a system install of either. On Windows, TensorRT comes the same way, but the delegate loads +the CUDA runtime from a CUDA 13 toolkit install. Install the toolkit, and make sure `CUDA_PATH` or +`CUDA_HOME` points at it. -TensorRT itself and the CUDA runtime arrive as dependencies of the wheels, so a system -install of either is not required. +By default, TensorRT builds each engine for the operating system, the GPU and the TensorRT version +it is exported with. So export your programs on the kind of machine that will run them, with the +same packages. + +On Windows, ExecuTorch's CUDA backend cannot export a program, and ExecuTorch's Python runtime there +does not include the CUDA backend. So +[sending unsupported operators to the CUDA backend](#sending-unsupported-operators-to-the-cuda-backend) +is a Linux setup. A program that uses only TensorRT still runs from Python on Windows, because there +the delegate registers the CUDA memory allocator itself. ## Development Requirements -One command installs everything needed to export a model. The `executorch` extra pulls in a -CUDA build of ExecuTorch, the delegate wheel, Torch-TensorRT, and PyTorch: +One command installs everything, in a fresh virtual environment: ```bash -pip install --pre "torch-tensorrt[executorch]" \ - --index-url https://download.pytorch.org/whl/nightly/cu132 \ - --extra-index-url https://pypi.org/simple \ - --extra-index-url https://pypi.nvidia.com +pip install "torch_tensorrt[executorch]" ``` -Swap `cu132` for the channel that matches your CUDA, for example `cu134` for CUDA 13.4. Keep -PyTorch, ExecuTorch, and Torch-TensorRT on the same channel. - -All three indexes matter. Without `--pre` pip takes the latest stable Torch-TensorRT from the -public index, which pairs with an older ExecuTorch release than this page describes. Without -NVIDIA's index the TensorRT package resolves to a source build that takes about twenty -minutes and then fails. - -A CUDA build of ExecuTorch is required, not only to build against. The delegate needs a -library that only the CUDA wheels carry, so a CPU-only build installs and then fails on -import. - -If you also want the CUDA backend to pick up the operators TensorRT rejects, install a CUDA -toolkit as well. That backend compiles those operators at export time with `nvcc`. - -### Feature availability +That gives you Torch-TensorRT, PyTorch, a CUDA build of ExecuTorch, and the TensorRT delegate, +which is the `torch-tensorrt-executorch-runtime` package. On Windows, add the CUDA package index, +because PyPI has only the CPU builds of PyTorch and ExecuTorch for Windows: -Two features on this page are newer than the 2026-10-06 Linux nightly of Torch-TensorRT and -its delegate package, which was the newest one when this page was written: +```bash +pip install "torch_tensorrt[executorch]" --extra-index-url https://download.pytorch.org/whl/cu132 +``` -- [CUDA graph replay](#cuda-graph-replay). With the 2026-10-06 build, `torch_tensorrt.save` - rejects `use_cuda_graphs` with a `TypeError`. -- [Shared engines](#shared-engines). The 2026-10-06 delegate keeps every engine private, so - loading a program twice still holds two copies of its weights. +The +[installation section](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#installation) +of the Torch-TensorRT guide has the other options. -Both are on the Torch-TensorRT main branch and in its 2.15 release branch. To use them, -install a wheel newer than 2026-10-06 once one is published, or build Torch-TensorRT from -source. Everything else on this page works with the 2026-10-06 build. +The CUDA build of ExecuTorch is required even to run a program: the delegate needs ExecuTorch's +`extension_cuda` library, which only the CUDA builds ship. To send the operators that TensorRT +rejects to the CUDA backend, you also need a CUDA toolkit, because that backend compiles them at +export time with `nvcc`, the CUDA compiler. PyTorch finds `nvcc` on your `PATH` or under +`CUDA_HOME`. ## Using the TensorRT Backend +The Python examples on this page run in order, and later ones reuse names from earlier ones. + ### Exporting a model Compile the model with Torch-TensorRT, then save it in ExecuTorch format: @@ -93,68 +104,83 @@ import torch_tensorrt class MyModel(torch.nn.Module): def forward(self, x): - return x + 1 + return torch.relu(x) + 1 with torch.no_grad(): model = MyModel().eval().cuda() - example_input = (torch.randn(2, 3, 4, 4).cuda(),) + example_inputs = (torch.randn(2, 3, 4, 4, device="cuda"),) - exported_program = torch.export.export(model, example_input) + exported_program = torch.export.export(model, example_inputs) trt_gm = torch_tensorrt.dynamo.compile( exported_program, - arg_inputs=example_input, + arg_inputs=example_inputs, min_block_size=1, ) - torch_tensorrt.save( trt_gm, "model.pte", output_format="executorch", - arg_inputs=example_input, - retrace=False, ) ``` -`retrace=False` is required. Re-exporting the compiled graph drops the engine, and the error -only shows up much later when the program is loaded. +The Torch-TensorRT page on +[saving models](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/saving_models.html#executorch-save) +describes the other options of `save`. It also shows how to put several methods in one file, and +how to let the engine update a KV cache in place. -### Letting the CUDA backend take what TensorRT rejects +### Sending unsupported operators to the CUDA backend -TensorRT has no converter for every ATen operator. By default the leftover operators run on -the CPU, which costs a copy in each direction. Pass a `CudaPartitioner` and they are compiled -by the ExecuTorch CUDA backend instead, so the whole model stays on the GPU: +TensorRT cannot convert every PyTorch operator. By default the ones it rejects run on ExecuTorch's +CPU kernels, with a copy in each direction on every call. If an operator such as `erfinv` has no +CPU kernel, the method fails to load. Pass a `CudaPartitioner`, and the ExecuTorch CUDA backend +compiles those operators instead: ```python from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner -torch_tensorrt.save( - trt_gm, - "coalesced.pte", - output_format="executorch", - arg_inputs=example_input, - retrace=False, - partitioners=[ - CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")]) - ], -) + +class CoalescedModel(torch.nn.Module): + def forward(self, x): + return torch.cos(torch.erfinv(torch.tanh(x))) + + +with torch.no_grad(): + 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( + coalesced_gm, + "coalesced.pte", + output_format="executorch", + partitioners=[ + CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")]) + ], + ) ``` -The TensorRT partitioner always runs first, and the `CudaPartitioner` picks up the rest. For -a model such as `cos(erfinv(tanh(x)))`, where TensorRT cannot take `erfinv`, the result is one -`.pte` whose delegate list reads `['TensorRTBackend', 'CudaBackend', 'TensorRTBackend']`. +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']`. +The Torch-TensorRT docs call such a program *coalesced*. -The CUDA backend writes its external weights to a `.ptd` file beside the `.pte` under a fixed -name. Save each coalesced model into its own directory. A second export into the same -directory overwrites the first model's weights, and the first model then still loads and -returns a wrong answer with no error. +The export writes the weights of the CUDA part to a data file, `aoti_cuda_blob.ptd`, next to the +`.pte`. Pass that file when you load the program: if the CUDA part has weights, loading fails +without it. The file has the same name for every model, so a second export into the same directory +overwrites the first one's weights, and the first model can then return wrong answers with no +error. Save each model into its own directory. ### Keeping inputs and outputs on the GPU -By default ExecuTorch inserts a host to device copy before the first delegate and a device to -host copy after the last one, so a method is safe to call with CPU tensors. If your data is -already on the GPU, turn those copies off: +By default, the program copies its inputs to the GPU before the first delegate, and its outputs +back after the last one, so you can call it with CPU tensors. If your data is already on the GPU, +turn those copies off. In the flag names, h2d means host (CPU) to device (GPU), and d2h the other +way: ```python from executorch.exir import ExecutorchBackendConfig @@ -165,45 +191,43 @@ torch_tensorrt.save( trt_gm, "device_resident.pte", output_format="executorch", - arg_inputs=example_input, - retrace=False, backend_config=ExecutorchBackendConfig( propagate_device_config=PropagateDeviceConfig( 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 points are easy to miss: - -- **Both skip flags need `enable_non_cpu_memory_planning=True`.** Copy insertion happens - during device aware memory planning. Setting it to `False` raises a `ValueError`. -- **Inputs must be unplanned**, with `MemoryPlanningPass(alloc_graph_input=False)`. Without - it the program reserves its own input buffer and the runtime copies the caller's memory - into it, which puts the copy straight back. -- **Leave the outputs planned when Python runs the program.** The program's own device arena - then holds the output, and Python returns a copy of it as a CUDA tensor. With - `alloc_graph_output=False` the caller provides the output memory instead. A C++ caller does - that with `Module::set_output`. The Python bindings do it correctly starting with ExecuTorch - 1.6.0.dev20261008, but the 2026-10-06 delegate wheel pins ExecuTorch 1.6.0.dev20260925, so - with that wheel keep the outputs planned. - -The choice is baked into the `.pte`. A program exported this way wants CUDA tensors, so pass -CUDA tensors. A host tensor works only when TensorRT is the first to read that input. On most -GPUs it then stages a copy on every call, which is the cost the export existed to remove. When -the CUDA backend reads it first, as it can in a coalesced program, it refuses the host memory -and the call fails. +- 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 a `ValueError`. +- `alloc_graph_input=False` is required. Without it 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`. +- `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, as the + [CUDA backend page](../cuda/cuda-overview.md#things-worth-knowing) shows. 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. The Python runtime refuses a CPU tensor. From +C++, what happens to a CPU tensor depends on which delegate reads that input first: + +- TensorRT: the call works, but copies the tensor on every call. +- The CUDA backend: the call fails. ## Runtime Integration ### Python -Import the delegate package once, anywhere before a program is loaded. The import is what -registers the backend. Nothing else about your code changes: +Import the delegate package once, before you load a program. The import registers the delegate. +ExecuTorch's own CUDA backend needs no import, and nothing else in your code changes: ```python from pathlib import Path @@ -217,57 +241,36 @@ forward = program.load_method("forward") outputs = forward.execute((torch.ones(2, 3, 4, 4),)) ``` -If the delegate cannot be loaded, that import raises straight away instead of leaving the -failure to show up later as a program that will not load. - -A coalesced program needs no extra registration. Both backends are registered, and the -program records which parts go where. Its weights are a different matter. If the part the -CUDA backend compiled has weights, they are in the `aoti_cuda_blob.ptd` file saved beside the -`.pte`, and the runtime does not look for that file on its own. Pass it when you load: +If the delegate cannot load, the import raises a `DelegateCompatibilityError` right away. If the +export wrote a data file, such as `aoti_cuda_blob.ptd` from the CUDA backend, pass it when you load: ```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),)) ``` -Without it, loading a program whose CUDA part has weights fails. - -### C++ against the installed wheels - -The wheels ship a prebuilt delegate and a CMake package, so a C++ application can link them -without building anything from source: - -```cmake -find_package(executorch REQUIRED COMPONENTS backend_cuda kernels_optimized) -find_package(executorch_backend_tensorrt REQUIRED) +A program that keeps its inputs and outputs on the GPU takes CUDA tensors and returns CUDA tensors: -target_link_libraries(my_app PRIVATE - executorch::runtime - executorch::backend_cuda - executorch::backend_tensorrt - executorch::kernels_optimized +```python +program = Runtime.get().load_program(Path("device_resident.pte")) +outputs = program.load_method("forward").execute( + (torch.ones(2, 3, 4, 4, device="cuda"),) ) ``` -`kernels_optimized` supplies the `et_copy` operators that move data across the method -boundary. `backend_cuda` is needed by any coalesced program, and also by a TensorRT only -program that still has those boundary copies, because it registers the device allocator they -use. +### C++ -The two packages live in two distributions, so point CMake at both. ExecuTorch is a namespace -package, so its path has to come from its distribution metadata: +The Python packages ship the prebuilt delegate and its CMake package, so a C++ application can link +it without building anything from source. The steps below are for Linux. On Windows, link the +same components and copy the DLLs beside your program, as +[CUDA on Windows](../../using-executorch-cpp.md#cuda-on-windows) shows, but export TensorRT programs +on Windows itself. -```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 needed for the form above, because the `backend_cuda` component -rejects older versions. CMake is not part of the wheels, so install it yourself. - -There is no header to include. The delegate registers itself with the backend registry from a -static initializer inside the shared library, and everything after that is the ordinary -ExecuTorch C++ API: +**1. Write the program.** Running a program needs no header from the delegate. The delegate +registers itself when its library loads, and the rest is the ordinary ExecuTorch C++ API. Save +this as `main.cpp`: ```cpp #include @@ -280,8 +283,7 @@ using namespace executorch::extension; int main() { Module module("model.pte"); - std::vector data(2 * 3 * 4 * 4, 1.0f); - auto input = make_tensor_ptr({2, 3, 4, 4}, std::move(data)); + auto input = make_tensor_ptr({2, 3, 4, 4}, std::vector(2 * 3 * 4 * 4, 1.0f)); const auto outputs = module.forward(input); if (!outputs.ok()) { @@ -294,198 +296,160 @@ int main() { } ``` -For a coalesced program whose CUDA part has weights, pass the data file to the constructor -as well, for the same reason as in Python: - -```cpp -Module module("coalesced.pte", "aoti_cuda_blob.ptd"); -``` - -Linking the target also records the wheel's own library directory in your binary, which is -right for an application built against an installed wheel and wrong for anything you plan to -redistribute. Turn it off with `EXECUTORCH_BACKEND_TENSORRT_EMBED_RUNPATH` and -`CMAKE_SKIP_BUILD_RPATH`, both set before `find_package`. - -### Building from source - -The delegate 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 uses, add the delegate next to ExecuTorch, and link the target it -provides: +**2. Write `CMakeLists.txt`.** ```cmake -set(EXECUTORCH_BUILD_CUDA ON) -set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON) -set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER ON) -set(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR ON) -set(EXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP ON) -set(EXECUTORCH_BUILD_EXTENSION_MODULE ON) - -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 -) -``` - -Turn on the CUDA backend even for a program that only uses TensorRT. It builds the shared -`extension_cuda` library the delegate links, and it registers the device allocator the -boundary copies use. Without it the delegate stops at configure time. It also needs PyTorch -installed in the Python environment that CMake finds. The other options are the extensions -the CUDA backend and the `Module` class depend on, and CMake stops with an error if one of -them is missing. - -Building the delegate needs CUDA Toolkit 12.5 or newer. `libextension_cuda` stays a shared -library on purpose, so that every CUDA-capable delegate in the process reads the same caller -stream. - -## Performance Options - -### Choosing the CUDA stream - -With no stream chosen, both delegates run on `cudaStreamPerThread`, the default stream of the -calling thread. A coalesced program run from one thread is therefore ordered correctly as it -is, with nothing to add. - -To run every delegate on a stream of your own, for example a green context stream, scope a -guard over the whole execution: - -```cpp -#include - -using namespace executorch::extension; - -cuda::CallerStreamGuard guard(stream); -module.forward(input); -``` - -The guard lives in the `extension_cuda` library, so add it to the link recipe above: - -```cmake -find_package(executorch REQUIRED COMPONENTS backend_cuda extension_cuda kernels_optimized) - -target_link_libraries(my_app PRIVATE executorch::extension_cuda) -``` - -One guard reaches every CUDA-capable delegate, because they all resolve the same shared -`libextension_cuda`. 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 graphs. - -### Green contexts - -Because both delegates honor the caller's stream, that stream can be a CUDA green context -stream. A green context holds a fixed number of streaming multiprocessors, so the whole model -is confined to that part of the GPU and the rest stays free for other work. Create the stream -with `cuGreenCtxStreamCreate` and scope the same guard over it. - -This has been exercised on a 108 SM card with a green context holding 8 of them, running a -program split across the TensorRT delegate and the CUDA backend. - -### CUDA graph replay - -This feature needs a Torch-TensorRT build newer than the 2026-10-06 nightly, see -[Feature availability](#feature-availability). - -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. It helps engines with -fixed shapes that launch many short kernels, where the CPU launch work is a large part of -each call. +cmake_minimum_required(VERSION 3.28) +project(my_app CXX) -Turn it on at export time: +find_package(executorch REQUIRED COMPONENTS backend_cuda kernels_optimized) +find_package(executorch_backend_tensorrt REQUIRED) -```python -torch_tensorrt.save( - trt_gm, - "model.pte", - output_format="executorch", - arg_inputs=example_input, - retrace=False, - use_cuda_graphs=True, +add_executable(my_app main.cpp) +target_link_libraries(my_app PRIVATE + executorch::runtime + executorch::backend_cuda + executorch::backend_tensorrt + executorch::kernels_optimized ) ``` -It is not free. Each engine keeps one stable device buffer per input and output, and every -call copies each input in and each output out. For an engine with few kernels, or with large -inputs and outputs, those copies can cost more than the launches they save. Shapes that -change on every call never replay and still pay the copies. - -Some engines never replay, even with replay on. An engine on the shared activation scratch -below always takes the ordinary path, so turning on both options gives no replay. Engines -with aliased outputs, GPUs without stream-ordered memory, and drivers older than CUDA 12.5 -also keep the ordinary path, as does a call on a green context stream. The delegate logs the -reason once. - -Recording is also not safe next to every kind of work. While a recording is running, creating -or destroying a TensorRT engine elsewhere in the process is unsafe, and so is a whole device -sync such as `cudaDeviceSynchronize` or `torch.cuda.synchronize()`. Leave replay off when -other threads may do any of that. A C++ host can refuse a program's saved request for one -load: +- `kernels_optimized` provides the CPU operators: the copies between the CPU and the GPU at the + edges of the program, and any operator that stays on the CPU. +- `backend_cuda` registers the allocator for the program's GPU memory, so a program that uses only + TensorRT needs it too. +- You need CMake 3.28 or newer. On older versions `find_package(executorch ...)` fails, because the + package provides the `executorch::` targets only from 3.28. CMake is not part of the packages, so + install it yourself. -```cpp -#include - -using namespace executorch::extension; -using namespace executorch::runtime; +**3. Build and run.** The two CMake packages are inside two Python packages, so point CMake at +both: -BackendOptions<1> options; -options.set_option("use_cuda_graphs", false); -LoadBackendOptionsMap by_backend; -by_backend.set_options("TensorRTBackend", options.view()); -module.load(by_backend); +```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}" +cmake --build build +./build/my_app ``` -### Shared activation scratch +`TORCH_TENSORRT_SKIP_DELEGATE_REGISTRATION=1` in the second command lets Python import the +delegate package to read its path, without loading the delegate. Run `./build/my_app` from the +folder that holds `model.pte`. -A TensorRT execution context allocates its own activation scratch and holds it for as long as -the context lives. A model lowered to many small engines pays that cost many times over, and -can run out of GPU memory on the engine count alone. This option backs all of a device's -contexts from one buffer instead, grown to the largest size any call has asked for: +If the export wrote a data file, pass it to the constructor: ```cpp -#include - -using namespace executorch::runtime; - -BackendOptions<1> options; -options.set_option("use_shared_activation_scratch", true); -set_option("TensorRTBackend", options.view()); +Module module("coalesced.pte", "aoti_cuda_blob.ptd"); ``` -Set it before loading the methods whose contexts should use the pool. The memory saved is the -sum of the separate scratch sizes less the largest of them. - -The cost is parallelism. Contexts that share one buffer do not run at the same time on the -device. The backend holds a per device lock from the claim on the buffer through the enqueue, -so two calls on one device are serialized at submission. The pool also never shrinks, and a -device that has run a pooled engine must not be reset with `cudaDeviceReset()`. - -### Shared engines +The Torch-TensorRT guide explains +[how to ship such an application](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#run-from-c) +to other machines, and links to the instructions for building the delegate from source. -This feature needs a Torch-TensorRT build newer than the 2026-10-06 nightly, see -[Feature availability](#feature-availability). - -When the same program is loaded twice, for example once per robot arm, handles loaded from -the same engine bytes, for the same device and the same weight streaming request, share one -engine. The engine weights are in GPU memory once. Each handle still gets its own execution -context and buffers. +## Performance Options -This is on by default. Measured on an 8 GB Jetson Orin Nano, loading one robot policy a second -time grew GPU memory by 806 MiB without sharing and by 122 MiB with it. Those numbers are the -cost of the second load, not the total. Finding a match means hashing the engine bytes on -every load, which costs roughly 0.15 s per GB on a Jetson AGX Thor. Pass the -`use_shared_engines` load option as `false` to keep one module's engines private. +Each option below links to its section of the Torch-TensorRT guide. The guide's section +[Setting a runtime option](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#setting-a-runtime-option) +shows how to set the options that apply at run time, from C++ or from Python. + +- [**Choosing the CUDA stream.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#choosing-the-cuda-stream) + By default, both delegates run on `cudaStreamPerThread`, the default stream of the calling + thread, so a program split between them runs in the right order on one thread. To use your own + stream, put one guard around the call. It reaches every delegate: + + ```cpp + #include + + using namespace executorch::extension; + + cuda::CallerStreamGuard guard(stream); + module.forward(input); + ``` + + The guard is in the `extension_cuda` library. Your own CUDA calls, such as creating the stream, + also need the CUDA runtime library from a CUDA 13 toolkit. In `CMakeLists.txt`, replace the + `find_package(executorch ...)` line from step 2 with the first line below, add the second line + after it, and add the third line after `add_executable`: + + ```cmake + 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) + ``` + + - 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, because the CUDA backend can return while + its work is still queued. + - The stream can belong to a CUDA green context. A green context owns a fixed number of the GPU's + streaming multiprocessors (SMs), so the model stays inside that part of the GPU. Creating its + stream with `cuGreenCtxStreamCreate` also needs `CUDA::cuda_driver`. 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. +- [**CUDA graph replay.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#cuda-graph-replay) + Turn it on when you save. It helps engines with fixed shapes that launch many short kernels: + + ```python + torch_tensorrt.save( + trt_gm, + "cuda_graphs.pte", + output_format="executorch", + use_cuda_graphs=True, + ) + ``` + + - It costs a copy of every input and output, and an engine whose input shapes change on every + call never replays. Calls on a green context stream, or on a stream of another CUDA context, + do not replay either, and the guide lists a few rarer cases. + - A program saved with replay on can still have it turned off at load, with the + `use_cuda_graphs` option. + - While a graph is being recorded, no other thread may: + - create or destroy a TensorRT engine or execution context (the object that runs an engine) + outside this delegate, + - call `cudaDeviceSynchronize()` or `torch.cuda.synchronize()`, + - or load a program that uses the CUDA backend, which can make such a sync. + + Recording happens again after every shape change, so leave replay off when other threads may + do these things. +- [**Shared activation scratch.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#shared-activation-scratch) + The `use_shared_activation_scratch` option, set for the whole process before the program loads, + lets all engines on a device share one buffer for their temporary results. That saves memory for + a model split into many engines, but two engines on the shared scratch never run at the same + time. The scratch never shrinks, and a device that used it must not be reset with + `cudaDeviceReset()`. +- [**Shared engines.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#shared-engines) + On by default: loading the same program twice, on the same device and with the same weight + streaming budget, keeps one copy of each engine in GPU memory. Engines are matched by a hash of + their bytes, which a crafted program could match on purpose, so every program in the process must + be trusted. To keep one program's engines private, pass the `use_shared_engines` option as + `false` when you load it. +- [**Weight streaming.**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#weight-streaming) + An engine compiled with `enable_weight_streaming=True` can keep part of its weights in CPU memory. + Set how many bytes stay on the GPU with `weight_streaming_budget_per_engine` in + `torch_tensorrt.save`, or with the `weight_streaming_budget` load option. The load option takes + the number of bytes as a decimal string, or as an integer below 2 GiB. +- [**Engines in a separate data file**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#keep-the-engines-in-a-separate-file) + (newer than 2.15). Pass `external_engine_data="engines"` to `torch_tensorrt.save` to write the + engines to `engines.ptd` instead of the `.pte`, and pass that file when you load the program. + The Python loader takes one data file, so for a program that also uses the CUDA backend, pass + `external_engine_data="aoti_cuda_blob"` to share its file. +- [**Optimization profiles**](https://docs.pytorch.org/TensorRT/user_guide/runtime_performance/executorch.html#optimization-profiles) + (newer than 2.15). One engine can be tuned for several ranges of input shapes, for example one + for decoding one token at a time and one for a long prompt. A C++ application picks the profile for + each call with `torch_tensorrt::executorch_backend::OptimizationProfileGuard`, from the + `torch_tensorrt/executorch/TensorRTBackend.h` header. Without the guard, every engine runs + profile 0, and Python cannot choose a profile. The Python packages do not carry that header, and + the guide explains where to get it. ## Examples - [Export a static shape model](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_static_shape.py) -- [Export a coalesced TensorRT and CUDA model](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_coalesced.py) +- [Export a model whose input shapes can change](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_dynamic_shape.py) +- [Export a model split between TensorRT and the CUDA backend](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_coalesced.py) - [Export a model that keeps its data on the GPU](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_device_resident.py) -- [C++ reference runner, including the green context option](https://github.com/pytorch/TensorRT/tree/main/examples/executorch_reference_runner) - -The Torch-TensorRT documentation covers the export options in more detail, including dynamic -shapes, programs with more than one method, and a zero copy KV cache for decoder models. +- [Export a decoder step whose KV cache stays in memory you own](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_kv_cache_decode.py) +- [Export one engine with two optimization profiles](https://github.com/pytorch/TensorRT/blob/main/examples/torchtrt_executorch_example/export_multi_profile.py) + (newer than 2.15) +- [C++ and Python reference runners](https://github.com/pytorch/TensorRT/tree/main/examples/executorch_reference_runner) + (the two multi-profile programs need a build newer than 2.15) diff --git a/docs/source/getting-started.md b/docs/source/getting-started.md index 3ffde9e9f9a..b03a1cc9759 100644 --- a/docs/source/getting-started.md +++ b/docs/source/getting-started.md @@ -37,8 +37,10 @@ pip install executorch torch \ | NVIDIA GPU, CUDA 13.2 | `cu132` | | NVIDIA GPU, CUDA 13.4 | `cu134` | -The CUDA packages are built for Linux, on x86_64 and ARM64. Use the `cpu` -variant on macOS and on Windows. If your CUDA version is not in the table, +The CUDA packages are built for Linux, on x86_64 and ARM64, and for Windows on +x86_64. On Windows they run CUDA programs from C++, and programs for the CUDA +backend are exported on Linux (see [CUDA on Windows](using-executorch-cpp.md#cuda-on-windows)). +Use the `cpu` variant on macOS. If your CUDA version is not in the table, choose the closest lower one with the same major version, because CUDA works across minor versions but not across major ones. There is no package for CUDA 12, so on CUDA 12 use the `cpu` variant or build from source. On CUDA 13.0 or 13.1 no @@ -54,8 +56,7 @@ To get a change that has landed on the `main` branch but is not in a release yet, use a nightly build. These are rebuilt every day. Put `nightly/` in front of the variant name, for example `nightly/cu132`, and add `--pre` to the command, otherwise pip skips development versions. Nightly builds cover the same -variants. CUDA 13.4 is the newest, and until the next release it is in nightly -builds only. +variants. The second index is needed because a bare `--index-url` replaces PyPI instead of adding to it, and some dependencies are published only on PyPI. diff --git a/exir/passes/propagate_device_config.py b/exir/passes/propagate_device_config.py index 6071209ffeb..119231f563a 100644 --- a/exir/passes/propagate_device_config.py +++ b/exir/passes/propagate_device_config.py @@ -27,9 +27,9 @@ class PropagateDeviceConfig: """Controls whether the runtime copies method inputs and outputs across the device boundary. - Skipping a copy also means the runtime must not reserve its own buffer for the tensor, or it - would fill that buffer from the caller's memory and the copy would come back. Memory planning - allocates graph inputs and outputs by default, so a program using either skip below needs:: + Memory planning reserves buffers for graph inputs and outputs by default. A program that skips + the input copy must use ``alloc_graph_input=False``. Unplanning the outputs as well, as below, + copies the least:: ExecutorchBackendConfig( propagate_device_config=PropagateDeviceConfig( @@ -43,7 +43,9 @@ class PropagateDeviceConfig: ) Without ``alloc_graph_input=False`` the runtime reserves device memory for an input the caller - already owns, then copies into it with a host memcpy, which is undefined for device memory. + already owns and refuses to copy into it, so every call fails with ``NotSupported``. Planned + outputs work, but they live in the method's memory, which the next run overwrites, and the + Python ``execute()`` copies them unless it gets ``clone_outputs=False``. """ # When True, method-level input tensors that feed directly into a device @@ -52,8 +54,9 @@ class PropagateDeviceConfig: # pre-staged on GPU. # # Pair with MemoryPlanningPass(alloc_graph_input=False), or the runtime reserves its own buffer - # for the input and copies the caller's memory into it, which is what this exists to avoid. - # A dict can be used to set per-method values, keyed by method name. + # for the input and refuses the caller's tensor. + # For per-method values, pass a dict of PropagateDeviceConfig, keyed by method name, as the + # propagate_device_config of ExecutorchBackendConfig. skip_h2d_for_method_inputs: Union[bool, Dict[str, bool]] = False # When True, device delegate outputs that are directly method outputs @@ -61,8 +64,11 @@ class PropagateDeviceConfig: # Useful for cross-method GPU pipelines where the next method consumes # GPU tensors directly. # - # Pair with MemoryPlanningPass(alloc_graph_output=False) for the same reason as the input flag. - # A dict can be used to set per-method values, keyed by method name. + # Planned outputs work, but they live in the method's memory, which the next run overwrites, + # and the Python runtime copies them. MemoryPlanningPass(alloc_graph_output=False) avoids + # that copy, and then a C++ caller must set each output with Module::set_output. + # For per-method values, pass a dict of PropagateDeviceConfig, keyed by method name, as the + # propagate_device_config of ExecutorchBackendConfig. skip_d2h_for_method_outputs: Union[bool, Dict[str, bool]] = False def __hash__(self) -> int: