diff --git a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt index 53c6ac2e2f..b01058cb1d 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/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index 0e94423730..f607cc7f1e 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 13dfbd965a..ac22676bae 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 ecdd7765bb..dc89b45763 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 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 TensorRT-RTX and Python-only builds have no delegate. To export a coalesced program you also +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,27 +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: 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. +``save()`` writes the compiled graph as it is, so it needs no example inputs. -Coalescing the leftover operators onto the CUDA backend -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +: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. -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: +Run the leftover operators on the CUDA backend +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +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=[ @@ -123,73 +150,28 @@ CUDA backend compiles them instead, so nothing leaves 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. That program - still loads, still reports finding its weights, and returns a wrong answer with no error. - 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")) +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*. -In C++, pass the file as the second argument of ``Module``: - -.. code-block:: cpp - - Module module("model.pte", "model.ptd"); - -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. + 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. - 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``. +Keep inputs and outputs on the GPU +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -Keeping 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 @@ -206,33 +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 are easy to miss: +* **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. -* **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. +Treat a ``.ptd`` file as trusted input, the same as a ``.pte``. The engines in it are GPU code. + +.. warning:: -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 ``.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 @@ -244,64 +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, that import raises at once, instead of surfacing later as a -program that will not load. +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. -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. +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: + +.. code-block:: python + + 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 requires it. + ``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 @@ -328,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()); + +``Error::NotFound`` from ``set_option`` means the delegate is not linked into your binary. -Runtime Performance --------------------- +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 @@ -396,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, -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 SM count: + 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 @@ -435,98 +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: -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. +#. the ``use_cuda_graphs`` load option, +#. a process option set to ``False``, +#. the choice saved in the program, +#. a process option set to ``True``. -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. +With none of them set, replay is off. -.. 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. +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. -A C++ host can refuse a program's saved request for one load: +Some calls never replay: -.. code-block:: cpp +* 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. - #include +An engine on the shared activation scratch below also replays. It records again whenever the +shared buffer grows. - using namespace executorch::extension; - using namespace executorch::runtime; +.. warning:: - 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); - } + 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. Shared activation scratch ^^^^^^^^^^^^^^^^^^^^^^^^^^ -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: +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. -.. code-block:: cpp +Set the option before you load the methods that should use the shared buffer. Methods loaded +before it keep their own scratch. - #include +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. - using namespace executorch::runtime; +Shared engines +^^^^^^^^^^^^^^^ - 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()); - } +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. -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. +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. -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()``. +To keep one program's engines private, pass the ``use_shared_engines`` load option as ``false``. -Shared engines -^^^^^^^^^^^^^^^ +Weight streaming +^^^^^^^^^^^^^^^^^ + +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. + +Optimization profiles +^^^^^^^^^^^^^^^^^^^^^^ + +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``: + +.. code-block:: cpp -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. + #include -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. + using torch_tensorrt::executorch_backend::OptimizationProfileGuard; + + 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); + } + +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 @@ -543,22 +667,26 @@ Comparison: .pte vs .pt2 * - libtorch at runtime - Not required - Required - * - Non-TRT ops - - ExecuTorch CUDA backend, or CPU fallback + * - Operators TensorRT rejects + - ExecuTorch CUDA backend, or the CPU - Compiled by AOTInductor * - Platform - - Linux + Windows x86_64, CUDA 13 - - Linux + Windows + - Linux (x86-64 and aarch64) and Windows x86-64, with CUDA 13 + - Linux ---- Examples --------- -* ``examples/torchtrt_executorch_example/export_static_shape.py`` writes a minimal ``.pte``. -* ``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 c39faec4e9..f4676e03fe 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. Linux-only; requires the ``executorch`` package. +* `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,41 +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. 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: +``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: -.. 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. - -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()`` """"""""""""""""""" @@ -270,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 @@ -341,312 +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. - - 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. + **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. - 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. +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. -``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: +With ``save()``, one flag does it: .. 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) + torch_tensorrt.save(trt_gm, "decode.pte", output_format="executorch", zero_copy_kv=True) -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``:: - - 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: - -.. code-block:: python - - 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**, 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. - -* **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 - - torch_tensorrt.save( - trt_gm, "trt.pte", output_format="executorch", - partitioners=[ - CudaPartitioner( - [CudaBackend.generate_method_name_compile_spec("forward")] - ) - ], - ) + from torch_tensorrt.executorch import check_zero_copy_kv, export, zero_copy_backend_config -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 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. - -.. _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. - -**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 @@ -656,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 d8349cd7b3..e4c620ddfe 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:-5e379f610e034da507b48da61464679609517b52}" -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 1337f749a7..b2d877f136 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 8423a0db71..4981b19a81 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 436976c3db..fa8b8645fb 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 37ab114aba..cb8c272dae 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 3586d3e3a7..24c9d81a6c 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 4d1aa18923..db0a5c0bc4 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 76e1537208..18d121d107 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 216155361c..fd9b11ad88 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. -Python callers cannot, so that arrangement is C++ only. - -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. -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. - -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 nothing extra here: both backends are registered, and the program says which -parts go where. - ## 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.3.0.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,153 +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.3.0.dev20261023+cu132`. A release on the PyTorch indexes uses the base version and the CUDA + suffix, such as `0.3.0+cu132`; the PyPI release has none, such as `0.3.0`. +- 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 copies them. - -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 -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. 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. - -## 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 f7330d43f2..fcf737547d 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 @@ -945,7 +948,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 +1589,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 diff --git a/py/torch_tensorrt/executorch/_export.py b/py/torch_tensorrt/executorch/_export.py index f14d344ce3..03cfe27cec 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 82b2e5ea02..e57eac9549 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 diff --git a/tests/py/dynamo/executorch/test_cmake_runtime.py b/tests/py/dynamo/executorch/test_cmake_runtime.py index cf6a5f0735..84ca84aacf 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 diff --git a/tests/py/dynamo/executorch/test_executorch_pin.py b/tests/py/dynamo/executorch/test_executorch_pin.py index 62f67b96db..cb16ae85f4 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,10 +966,10 @@ 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()