From 7475a10048148e3fd41f89712881ac18f8330809 Mon Sep 17 00:00:00 2001 From: cehongwang Date: Mon, 24 Aug 2026 21:46:31 +0000 Subject: [PATCH 1/2] feat(executorch): multi-optimization-profile support for the TensorRT delegate Lets an exported program carry several TensorRT optimization profiles and pick one per execution, either pinned via OptimizationProfileGuard or chosen automatically from the input shapes. Squashed from four commits (initial implementation plus three rounds of review fixups) so the rebase onto main resolves the overlapping runtime changes once. --- .../verify-executorch-reference-runner.sh | 15 +- core/runtime/TRTEngine.cpp | 25 +- core/runtime/TRTEngine.h | 2 +- cpp/BUILD | 31 +- .../executorch/OptimizationProfileSelection.h | 210 ++++++++++ .../executorch/TensorRTBackend.h | 187 ++++++--- .../executorch/TensorRTBindingNames.h | 14 + .../executorch/TensorRTBlobHeader.h | 46 ++- .../torch_tensorrt/executorch/CMakeLists.txt | 8 + .../torch_tensorrt/executorch/EngineHandle.h | 73 ++++ .../executorch/TensorRTBackend.cpp | 269 +++++++++++-- .../dynamo/multi_optimization_profiles.py | 106 +++-- examples/executorch_reference_runner/BUILD | 5 + .../CMakeLists.txt | 22 ++ .../executorch_reference_runner/README.md | 132 ++++++- .../multi_profile_benchmark.cpp | 361 ++++++++++++++++++ .../multi_profile_main.cpp | 342 +++++++++++++++++ .../export_multi_profile.py | 291 ++++++++++++++ .../dynamo/runtime/_TRTEngine.py | 23 +- tests/cpp/executorch/BUILD | 18 + .../test_optimization_profile_selection.cpp | 302 +++++++++++++++ .../test_multi_optimization_profiles.py | 13 + 22 files changed, 2324 insertions(+), 171 deletions(-) create mode 100644 cpp/include/torch_tensorrt/executorch/OptimizationProfileSelection.h create mode 100644 cpp/src/torch_tensorrt/executorch/EngineHandle.h create mode 100644 examples/executorch_reference_runner/multi_profile_benchmark.cpp create mode 100644 examples/executorch_reference_runner/multi_profile_main.cpp create mode 100644 examples/torchtrt_executorch_example/export_multi_profile.py create mode 100644 tests/cpp/executorch/test_optimization_profile_selection.cpp diff --git a/.github/scripts/verify-executorch-reference-runner.sh b/.github/scripts/verify-executorch-reference-runner.sh index 7af243e54a..04156941aa 100755 --- a/.github/scripts/verify-executorch-reference-runner.sh +++ b/.github/scripts/verify-executorch-reference-runner.sh @@ -285,6 +285,12 @@ require_tar_entry "torch_tensorrt/bin/example_executorch_runner" require_tar_entry "torch_tensorrt/lib/libextension_cuda.so" require_tar_entry "torch_tensorrt/examples/executorch_reference_runner/kv_cache_decode_check.cpp" require_tar_entry "torch_tensorrt/BUILD" +# Every source the packaged CMakeLists.txt names must ship. A missing one aborts +# the configure step below for all targets, not just the one that needs it, so +# check them here to fail at packaging with a clear message instead. +require_tar_entry "torch_tensorrt/examples/executorch_reference_runner/main.cpp" +require_tar_entry "torch_tensorrt/examples/executorch_reference_runner/multi_profile_main.cpp" +require_tar_entry "torch_tensorrt/examples/executorch_reference_runner/multi_profile_benchmark.cpp" export TORCH_TENSORRT_ROOT="${verify_root}/torch_tensorrt" export TORCHTRT_EXECUTORCH_SOURCE_DIR="${TORCH_TENSORRT_ROOT}/src/torch_tensorrt/executorch" @@ -305,7 +311,14 @@ fi cmake "${cmake_args[@]}" -build_targets=(example_executorch_runner) +# The multi-profile targets are built but not run: they need a Gemma-3 .pte and a +# GPU. Compiling them here is what keeps the packaged sources from rotting, since +# nothing else in CI touches them. +build_targets=( + example_executorch_runner + example_executorch_multi_profile_runner + example_executorch_multi_profile_benchmark +) if [[ -n "${kv_model_path}" ]]; then build_targets+=(kv_cache_decode_check) fi diff --git a/core/runtime/TRTEngine.cpp b/core/runtime/TRTEngine.cpp index f516d619f9..5f90fe94f7 100644 --- a/core/runtime/TRTEngine.cpp +++ b/core/runtime/TRTEngine.cpp @@ -720,15 +720,34 @@ void TRTEngine::set_active_profile(int64_t profile_index) { } void TRTEngine::set_active_profile_with_stream(int64_t profile_index, const c10::cuda::CUDAStream& stream) { - if (num_optimization_profiles <= 1) { + // An index this engine does not have cannot be honored. A single-profile engine + // is the one case where that is not a mistake to fail on: it has profile 0 and + // nothing to switch to, so a pin aimed at some other engine is ignored with a + // warning rather than raising. A multi-profile engine handed an index it lacks + // could have done something different, and a negative index is a computed value + // gone wrong on any engine, so both stay hard errors -- silently running on + // mistuned kernels is worse than stopping. Same split as the ExecuTorch backend + // (kPinIgnoredSingleProfile vs kRequestedProfileUnavailable). + // + // Reachable only by driving the engine directly; the Python wrapper validates + // the index against the profile count first. + if (num_optimization_profiles <= 1 && profile_index > 0) { + LOG_WARNING( + "Ignoring optimization profile index " << profile_index << ": this engine has " << num_optimization_profiles + << " optimization profile(s), so it stays on profile " + << active_profile_index << "."); return; } + TORCHTRT_CHECK( + profile_index >= 0 && profile_index < num_optimization_profiles, + "Optimization profile index " << profile_index << " is out of range: this engine has " + << num_optimization_profiles << " optimization profile(s)."); if (profile_index == active_profile_index) { return; } - // setOptimizationProfileAsync returns false for an out-of-range index; the - // index is validated upstream in TorchTensorRTModule.resolve_profile_index. + // The index is in range by the check above, so a false return here is TensorRT + // refusing the switch for some other reason. TORCHTRT_CHECK( exec_ctx()->setOptimizationProfileAsync(static_cast(profile_index), stream.stream()), "Failed to switch to optimization profile index " << profile_index); diff --git a/core/runtime/TRTEngine.h b/core/runtime/TRTEngine.h index cec6d8ab78..4f7d642088 100644 --- a/core/runtime/TRTEngine.h +++ b/core/runtime/TRTEngine.h @@ -327,7 +327,7 @@ struct TRTEngine : torch::CustomClassHolder { // method names are exposed via torchbind in register_jit_hooks.cpp // (``num_optimization_profiles``, ``_active_profile_index``, // ``_auto_select_profiles``, ``set_active_profile``). Index validation lives - // in the runtime-agnostic TorchTensorRTModule.resolve_profile_index. + // in the runtime-agnostic TorchTensorRTModule.set_optimization_profile. int64_t num_optimization_profiles = 1; // cuda_engine->getNbOptimizationProfiles() int64_t active_profile_index = 0; // profile currently loaded in exec_ctx bool auto_select_profiles = false; // opt-in shape-based selection (per call) diff --git a/cpp/BUILD b/cpp/BUILD index 2f1edf5b3c..ea581783e9 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -141,9 +141,33 @@ cc_library( ], ) +cc_library( + name = "tensorrt_executorch_optimization_profile_selection", + hdrs = [ + "include/torch_tensorrt/executorch/OptimizationProfileSelection.h", + ], + strip_include_prefix = "include", + # The header includes , so it cannot build where the deps below + # resolve to an empty list. + target_compatible_with = select({ + ":linux_x86_64": [], + ":sbsa": [], + "//conditions:default": ["@platforms//:incompatible"], + }), + deps = select({ + ":linux_x86_64": ["@tensorrt//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + "//conditions:default": [], + }), +) + cc_library( name = "tensorrt_executorch_backend", srcs = [ + # Private, deliberately not in hdrs: EngineHandle grows fields as the + # backend gains features and is never installed, so nothing outside this + # library may depend on its layout. + "src/torch_tensorrt/executorch/EngineHandle.h", "src/torch_tensorrt/executorch/TensorRTBackend.cpp", ], hdrs = [ @@ -161,18 +185,19 @@ cc_library( deps = [ ":tensorrt_executorch_binding_names", ":tensorrt_executorch_blob_header", + ":tensorrt_executorch_optimization_profile_selection", ":tensorrt_executorch_weight_streaming_budget", ] + select({ ":linux_x86_64": [ + "@cuda//:cudart", "@executorch//:executorch_headers", "@executorch//:extension_cuda", - "@cuda//:cudart", "@tensorrt//:nvinfer", ], ":sbsa": [ + "@cuda//:cudart", "@executorch//:executorch_headers", "@executorch//:extension_cuda", - "@cuda//:cudart", "@tensorrt_sbsa//:nvinfer", ], "//conditions:default": [], @@ -184,6 +209,7 @@ filegroup( name = "executorch_backend_source_files", srcs = [ "src/torch_tensorrt/executorch/CMakeLists.txt", + "src/torch_tensorrt/executorch/EngineHandle.h", "src/torch_tensorrt/executorch/README.md", "src/torch_tensorrt/executorch/TensorRTBackend.cpp", "src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp", @@ -204,6 +230,7 @@ filegroup( filegroup( name = "executorch_api_headers", srcs = [ + "include/torch_tensorrt/executorch/OptimizationProfileSelection.h", "include/torch_tensorrt/executorch/TensorRTBackend.h", "include/torch_tensorrt/executorch/TensorRTBindingNames.h", "include/torch_tensorrt/executorch/TensorRTBlobHeader.h", diff --git a/cpp/include/torch_tensorrt/executorch/OptimizationProfileSelection.h b/cpp/include/torch_tensorrt/executorch/OptimizationProfileSelection.h new file mode 100644 index 0000000000..ab8d6a09f9 --- /dev/null +++ b/cpp/include/torch_tensorrt/executorch/OptimizationProfileSelection.h @@ -0,0 +1,210 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +/** + * @file OptimizationProfileSelection.h + * @brief Which TensorRT optimization profile an execution runs under. + * + * Kept free of ExecuTorch, CUDA, and the engine itself so the policy can be + * exercised without a GPU; reporting the outcome is left to the caller. + */ +#pragma once + +#include + +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +/// @brief The [min, max] dim envelope one optimization profile allows for one input. +struct InputProfileBounds { + nvinfer1::Dims min{}; ///< Smallest shape the profile accepts. + nvinfer1::Dims max{}; ///< Largest shape the profile accepts. +}; + +/// @brief Everything a profile decision depends on, read from the engine once at init(). +struct ProfileTable { + /** + * @brief Bounds indexed [profile][input]. + * + * The outer size is the engine's optimization profile count, which is at + * least 1; a single-profile engine keeps exactly one row and never switches. + */ + std::vector> bounds; + /// @brief The profile currently loaded into the execution context. + int32_t active = 0; + + /// @return The engine's optimization profile count. + int32_t size() const { + return static_cast(bounds.size()); + } +}; + +/// @brief What the calling thread asked for, as resolved from OptimizationProfileGuard. +enum class ProfileRequest { + kUnset, ///< No guard in scope. + kPinned, ///< An exact index. + kAuto, ///< Choose from the input shapes. +}; + +/** + * @brief The outcome of one profile decision. + * + * Its own enum rather than executorch's Error so that this header stays + * independent of executorch and can be tested separately. + * + * Two axes: whether execution continues, and which message the caller prints. + * The two failure values stay apart rather than being merged and re-derived from + * the request kind, because the empty-table guard in select_profile() returns + * kNoProfileMatchesInputs for every request kind -- so one merged value would put + * the message back at the mercy of which branches each request can reach. + */ +enum class ProfileSelection { + kOk, ///< The selected profile is usable. + /** + * @brief Succeeded, but the pin could not be honored and profile 0 was used instead. + * + * Distinct from kOk so the caller can warn that the pin did nothing here. + */ + kPinIgnoredSingleProfile, + /** + * @brief A pinned index this engine does not have and cannot substitute for. + * + * Fatal, on the same split the standard runtime uses in + * TRTEngine::set_active_profile_with_stream: an engine that had profiles to + * choose between must not quietly run on one the caller did not ask for. + * OptimizationProfileGuard cannot validate anything -- it never sees an + * engine, by design -- so execute() is the only place an ExecuTorch caller's + * bad index can be caught at all. + */ + kRequestedProfileUnavailable, + /// @brief Auto-selection ran out of profiles. + kNoProfileMatchesInputs, +}; + +/** + * @brief Whether one input shape falls inside one profile's envelope. + * + * @param dims The runtime shape of the input. + * @param bounds The profile's [min, max] envelope for that input. + * @return true when the ranks match and every extent is in range. + */ +inline bool dims_fit(const nvinfer1::Dims& dims, const InputProfileBounds& bounds) { + if (dims.nbDims != bounds.min.nbDims) { + return false; + } + for (int d = 0; d < dims.nbDims; ++d) { + if (dims.d[d] < bounds.min.d[d] || dims.d[d] > bounds.max.d[d]) { + return false; + } + } + return true; +} + +/** + * @brief Whether one profile accepts every input shape of an execution. + * + * A profile index or bounds row that cannot describe these inputs answers "does + * not fit" rather than indexing past the end of the table. Neither is reachable + * from the backend, which builds one row of num_inputs bounds per profile and + * only ever stores an index this function approved -- but this is installed + * public API, and on the auto path the answer also happens to be the useful one: + * a ProfileTable carrying a stale ProfileTable::active rescans from 0 instead of + * crashing. + * + * @param table The engine's profile bounds. + * @param profile Index of the profile to test. + * @param input_dims Runtime shape of each input, in binding order. + * @return true when the profile accepts all of them. + */ +inline bool profile_fits(const ProfileTable& table, int32_t profile, const std::vector& input_dims) { + if (profile < 0 || profile >= table.size()) { + return false; + } + const auto& bounds = table.bounds[static_cast(profile)]; + if (bounds.size() < input_dims.size()) { + return false; + } + for (size_t i = 0; i < input_dims.size(); ++i) { + if (!dims_fit(input_dims[i], bounds[i])) { + return false; + } + } + return true; +} + +/** + * @brief Resolves one thread's profile request against one engine. + * + * @param table The engine's profile bounds and currently loaded profile. + * @param request What the calling thread asked for. + * @param index The pinned profile index; read only for ProfileRequest::kPinned. + * @param input_dims Runtime shape of each input, in binding order. + * @param[out] selected The profile to run, written only when the result allows + * execution to continue (ProfileSelection::kOk or + * ProfileSelection::kPinIgnoredSingleProfile). + * @return What the caller should do, and which message it should print. + */ +inline ProfileSelection select_profile( + const ProfileTable& table, + ProfileRequest request, + int32_t index, + const std::vector& input_dims, + int32_t& selected) { + // init() rejects an engine reporting no profiles, so this is unreachable in the + // backend. Checked here so the policy is safe to call on its own rather than on + // the strength of a guard in another translation unit. + if (table.bounds.empty()) { + return ProfileSelection::kNoProfileMatchesInputs; + } + + if (request == ProfileRequest::kUnset) { + selected = 0; + return ProfileSelection::kOk; + } + + if (request == ProfileRequest::kAuto) { + // Sticky first-fit: keep the loaded profile while it still fits, so shapes + // that alternate between two equally valid profiles don't thrash the + // context. Only rescan from 0 once it stops fitting. Overlapping profiles + // therefore resolve by history, not by lowest index; pin explicitly when + // that matters. + if (profile_fits(table, table.active, input_dims)) { + selected = table.active; + return ProfileSelection::kOk; + } + for (int32_t p = 0; p < table.size(); ++p) { + if (profile_fits(table, p, input_dims)) { + selected = p; + return ProfileSelection::kOk; + } + } + return ProfileSelection::kNoProfileMatchesInputs; + } + + if (index >= 0 && index < table.size()) { + selected = index; + return ProfileSelection::kOk; + } + + // A single-profile engine has no choice to get wrong: profile 0 is the only + // thing it can run, whether or not its shapes are dynamic. So a pin aimed at a + // multi-profile sibling in the same method must not fail it. An engine with + // several profiles is different -- substituting one would be a guess -- so an + // index it lacks stays an error there. + if (index > 0 && table.size() == 1) { + selected = 0; + return ProfileSelection::kPinIgnoredSingleProfile; + } + + return ProfileSelection::kRequestedProfileUnavailable; +} + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index b33d712d40..719dbb57e8 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -3,9 +3,14 @@ * * This source code is licensed under the BSD-style license found in the * LICENSE file in the root directory of this source tree. + */ + +/** + * @file TensorRTBackend.h + * @brief ExecuTorch backend delegate that runs TensorRT engines serialized by + * torch_tensorrt. * - * ExecuTorch backend delegate that runs TensorRT engines serialized by - * torch_tensorrt. The processed blob uses the standalone wire format from + * The processed blob uses the standalone TR01 wire format from * py/torch_tensorrt/executorch/serialization.py and is parsed directly here. * This runtime path intentionally does not depend on the legacy * Torch-TensorRT C++ runtime or libtorch. @@ -17,14 +22,15 @@ #include +#include "torch_tensorrt/executorch/OptimizationProfileSelection.h" + +#include #include -#include -#include -#include namespace torch_tensorrt { namespace executorch_backend { +/// @brief Deletes TensorRT interface objects, which are freed with `delete`. struct TRTDeleter { template void operator()(T* p) const { @@ -32,85 +38,146 @@ struct TRTDeleter { } }; +/// @brief Owning pointer to a TensorRT interface object. template using TRTUniquePtr = std::unique_ptr; +/// @brief Forwards TensorRT diagnostics to the ExecuTorch log. class TRTLogger : public nvinfer1::ILogger { public: void log(Severity severity, const char* msg) noexcept override; }; -struct InputProfileBounds { - nvinfer1::Dims min{}; - nvinfer1::Dims max{}; -}; - -struct EngineHandle { - TRTLogger logger; - TRTUniquePtr runtime; - TRTUniquePtr engine; - TRTUniquePtr exec_ctx; - std::vector input_binding_names; - std::vector output_binding_names; - std::vector input_profile_bounds; - std::vector cached_input_ptrs; - std::vector cached_input_sizes; - std::vector cached_output_ptrs; - std::vector cached_output_sizes; - size_t num_inputs = 0; - size_t num_outputs = 0; - // Per output binding [0..num_outputs): index into input_binding_names of the - // input it aliases (in-place KV-cache / user alias), or -1 for a normal output. - // Built at init from the blob's aliased_io. The KV buffers are threaded by - // ExecuTorch as caller-owned mutable-buffer delegate args (input AND aliased - // output): execute() binds each aliased TRT output binding to its aliased - // input's caller-provided pointer (in-place) and reflects the result into the - // delegate output EValue, which ExecuTorch's write-back copy_ then reads. - std::vector output_aliased_input_idx; - // Per input binding [0..num_inputs): true if any output aliases this input, so - // its in-place (KV/user) update must land in the caller-owned storage. Built at - // init from aliased_io; execute() uses it to reject a non-device-resident - // aliased input instead of silently staging its update into delegate scratch. - std::vector input_is_alias_target; - size_t num_aliased_outputs = 0; - int device_id = 0; - bool unified_memory = false; - std::mutex mu; - // Makes the skip-sync fast path safe to reuse: TensorRT forbids reconfiguring or - // destroying an execution context while one of its enqueues is in flight, so when - // execute() returns without an end sync it records this event; the next execute() - // and the destructor wait on it before touching exec_ctx. One event/flag pair - // suffices because a handle runs on a single thread at a time. - cudaEvent_t inflight_event = nullptr; - bool inflight_pending = false; - - ~EngineHandle(); -}; - +/** + * @brief The delegate ExecuTorch calls to run a TensorRT engine. + * + * Registered under the backend id `TensorRT`; a `.pte` produced by + * torch_tensorrt.save(output_format="executorch") dispatches to it. + */ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { public: + /// @return Whether a usable CUDA device and TensorRT runtime are present. bool is_available() const override; + /// @brief Deserializes one engine from its processed blob into a handle. ::executorch::runtime::Result<::executorch::runtime::DelegateHandle*> init( ::executorch::runtime::BackendInitContext& context, ::executorch::runtime::FreeableBuffer* processed, ::executorch::runtime::ArrayRef<::executorch::runtime::CompileSpec> compile_specs) const override; - // Runs the engine. With an executorch::extension::cuda::CallerStreamGuard active and - // no host staging required, this may return while the enqueue is still in flight on - // the selected stream, so the caller must keep device buffers alive and unmodified - // past return, order any other stream against this one, and synchronize the stream - // before reading device-resident outputs. The selected stream must be on the engine's - // device, and calls on one handle must not overlap each other or its destruction. - // Note that other CUDA delegates sharing the same guard may instead synchronize before - // returning, so do not assume results are ready on return from this one. + /** + * @brief Binds `args` and enqueues the engine, selecting the optimization + * profile the calling thread's OptimizationProfileGuard asked for and the + * CUDA stream its executorch::extension::cuda::CallerStreamGuard selected. + * + * With a CallerStreamGuard active and no host staging required, this may + * return while the enqueue is still in flight on the selected stream, so the + * caller must keep device buffers alive and unmodified past return, order any + * other stream against this one, and synchronize the stream before reading + * device-resident outputs. The selected stream must be on the engine's device, + * and calls on one handle must not overlap each other or its destruction. + * Note that other CUDA delegates sharing the same guard may instead + * synchronize before returning, so do not assume results are ready on return + * from this one. + */ ::executorch::runtime::Error execute( ::executorch::runtime::BackendExecutionContext& context, ::executorch::runtime::DelegateHandle* handle, ::executorch::runtime::Span<::executorch::runtime::EValue*> args) const override; + /// @brief Waits for any work still in flight, then releases the handle. void destroy(::executorch::runtime::DelegateHandle* handle) const override; }; +/** + * @brief Selects, for the calling thread, which TensorRT optimization profile + * the delegate runs; scope it around Module::forward() / Module::execute(). + * + * A profile is identified by its index in the export-time profile list, so name + * them to match whatever the exporter declared: + * + * @code + * constexpr int32_t kDecodeProfile = 0; // export order: decode first, + * constexpr int32_t kPrefillProfile = 1; // then prefill + * + * executorch::extension::Module module("model.pte"); + * { + * OptimizationProfileGuard profile_guard(kPrefillProfile); + * auto result = module.forward(prefill_inputs); + * } + * @endcode + * + * The guard records a request for the current thread and does nothing else: it + * never inspects the Module, Method, or delegate handles, and never calls + * TensorRT. Each TensorRT delegate reads the request inside its own execute(), + * where the engine, its lock, and the execution stream are already available, + * and switches there. Without a guard every delegate runs profile 0. + * + * Composes with executorch::extension::cuda::CallerStreamGuard, which is + * orthogonal: the stream guard says where the GPU work runs, this one says which + * profile it runs under. A switch is issued on whichever stream execute() + * selected. + * + * Contract: construct the guard on the thread that calls forward()/execute() + * (ExecuTorch does not support concurrent execution of one Module anyway). + * Nested guards restore the enclosing request on scope exit. + * + * One execution sees one consistent request, but several TensorRT engines in a + * method apply it independently as they run. TensorRT offers no way to undo a + * switch, so if a later engine rejects the request (a pinned index it does not + * have, or no profile matching its inputs) it returns an error with earlier + * engines already switched. + * + * @warning The index is delivered to every TensorRT delegate in the method, and + * each one resolves it against its own profile list. Nothing makes index 1 mean + * the same thing in two engines: if a `.pte` contains two engines compiled from + * different profile lists, one index can select prefill in one and decode in the + * other. Pin by index only when the engines were built from a single profile + * list, or when the `.pte` holds one TensorRT engine. An engine with a single + * profile is the benign case -- it runs profile 0 and logs that the pin did + * nothing -- while a multi-profile engine that lacks the index fails the + * execution. + */ +class OptimizationProfileGuard { + public: + /** + * @brief Pin an exact profile by its export-time index. + * + * An index this engine does not have is reported by execute(), not here, since + * the guard never sees the engine; that is deliberate, so a computed index + * (say -1 from a failed lookup) surfaces as an error rather than quietly + * meaning something else. + * + * @param profile_index Position in the export-time profile list. + */ + explicit OptimizationProfileGuard(int32_t profile_index); + + /// @brief Rejected so that OptimizationProfileGuard(true) cannot become index 1. + OptimizationProfileGuard(bool) = delete; + + /** + * @brief Have each delegate choose from the runtime input shapes instead of + * being told an index. + * + * Named rather than a sentinel index so it cannot collide with a computed one: + * + * @code + * auto profile_guard = OptimizationProfileGuard::automatic(); + * @endcode + */ + static OptimizationProfileGuard automatic(); + + ~OptimizationProfileGuard(); + OptimizationProfileGuard(const OptimizationProfileGuard&) = delete; + OptimizationProfileGuard& operator=(const OptimizationProfileGuard&) = delete; + + private: + struct AutoTag {}; + explicit OptimizationProfileGuard(AutoTag); + + ProfileRequest prev_request_; + int32_t prev_index_; +}; + } // namespace executorch_backend } // namespace torch_tensorrt diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBindingNames.h b/cpp/include/torch_tensorrt/executorch/TensorRTBindingNames.h index 5fe0375014..d72265d37d 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBindingNames.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBindingNames.h @@ -1,3 +1,17 @@ +/** + * @file TensorRTBindingNames.h + * @brief Re-exports the shared binding-name helpers into the ExecuTorch backend. + * + * The delegate has to map ExecuTorch's positional arguments onto TensorRT's + * named bindings, and the rules for that are shared with the standard runtime. + * Rather than restate them, this header pulls + * torch_tensorrt::core::runtime's helpers in, so both runtimes agree by + * construction. The include is spelled two ways because the installed tarball + * and the in-repo build lay the core runtime out differently. + * + * Everything here lives in a `detail` namespace: it is shared with the standard + * runtime, not part of the ExecuTorch backend's public API. + */ #pragma once #if __has_include("torch_tensorrt/core/runtime/TensorRTBindingNames.h") diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h b/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h index ce1dfaa9b9..0aab48991e 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h @@ -1,3 +1,13 @@ +/** + * @file TensorRTBlobHeader.h + * @brief The TR01 processed-blob layout the ExecuTorch delegate is handed. + * + * Written by py/torch_tensorrt/executorch/serialization.py: a 32-byte fixed + * header carrying the magic `TR01` and four offset/size fields, then JSON + * metadata, then the serialized TensorRT engine. Deliberately standalone, so + * loading a `.pte` needs neither the legacy Torch-TensorRT C++ runtime nor + * libtorch. + */ #pragma once #include @@ -18,18 +28,42 @@ struct AliasedBinding { std::string kind; // "kv_cache_update" (TRT-enforced) or "user" }; +/// @brief Where the metadata and engine live inside a TR01 blob, and what the metadata said. struct TensorRTBlobHeader { - uint32_t metadata_offset = 0; - uint32_t metadata_size = 0; - uint32_t engine_offset = 0; - uint64_t engine_size = 0; + uint32_t metadata_offset = 0; ///< Byte offset of the JSON metadata; never inside the header. + uint32_t metadata_size = 0; ///< Length of the JSON metadata in bytes. + uint32_t engine_offset = 0; ///< Byte offset of the serialized engine; 16-byte aligned. + uint64_t engine_size = 0; ///< Length of the serialized engine in bytes. + /// @brief TensorRT input binding names, in the order the metadata lists them. std::vector input_binding_names; + /// @brief TensorRT output binding names, in the order the metadata lists them. std::vector output_binding_names; + /// @brief Aliased output->input binding pairs declared by the metadata. std::vector aliased_io; - bool hardware_compatible = false; - int device_id = 0; + bool hardware_compatible = false; ///< Engine built hardware-compatible; false when unstated. + int device_id = 0; ///< Device the engine was built for; 0 when unstated. + /** + * @brief Start of the serialized engine inside a blob. + * + * @param blob Start of the processed blob `h` was parsed from. + * @param h Header parsed from that blob. + * @return Pointer to TensorRTBlobHeader::engine_size bytes of engine data. + */ static const void* engine_data(const void* blob, const TensorRTBlobHeader& h); + + /** + * @brief Validates a processed blob and reads its header. + * + * Checks the `TR01` magic, that the metadata and the engine both lie within + * `size`, that the engine is 16-byte aligned, and that the metadata ends at or + * before the engine; then reads the binding names and flags out of the JSON. + * + * @param data Start of the processed blob. + * @param size Bytes available at `data`. + * @param[out] out Filled in on success; left in an unspecified state otherwise. + * @return false when the blob is not a well-formed TR01 blob. + */ static bool parse(const void* data, std::size_t size, TensorRTBlobHeader& out); }; diff --git a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt index 1b503c567d..b5eb617140 100644 --- a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt +++ b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt @@ -173,6 +173,14 @@ add_library(torchtrt_executorch_backend INTERFACE) add_library(torchtrt::executorch_backend ALIAS torchtrt_executorch_backend) add_dependencies(torchtrt_executorch_backend executorch_trt_backend) +# The archive is linked by file below rather than by target, so the include path +# does not come along with it. Carry it here: a runner that scopes CudaStreamGuard +# or OptimizationProfileGuard needs the public header. +target_include_directories(torchtrt_executorch_backend + INTERFACE + "${CMAKE_CURRENT_LIST_DIR}/../../../include" +) + if(MSVC) target_link_libraries(torchtrt_executorch_backend INTERFACE diff --git a/cpp/src/torch_tensorrt/executorch/EngineHandle.h b/cpp/src/torch_tensorrt/executorch/EngineHandle.h new file mode 100644 index 0000000000..366b7f73b4 --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/EngineHandle.h @@ -0,0 +1,73 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + * + * Private state of a TensorRT ExecuTorch delegate. + * + * This header is deliberately not installed. EngineHandle grows fields as the + * backend gains features, so keeping it out of the public API means a new + * header can never disagree about its layout with an already-built backend + * archive. + */ +#pragma once + +#include "torch_tensorrt/executorch/OptimizationProfileSelection.h" +#include "torch_tensorrt/executorch/TensorRTBackend.h" + +#include +#include + +#include +#include +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { + +struct EngineHandle { + TRTLogger logger; + TRTUniquePtr runtime; + TRTUniquePtr engine; + TRTUniquePtr exec_ctx; + std::vector input_binding_names; + std::vector output_binding_names; + ProfileTable profiles; + std::vector cached_input_ptrs; + std::vector cached_input_sizes; + std::vector cached_output_ptrs; + std::vector cached_output_sizes; + size_t num_inputs = 0; + size_t num_outputs = 0; + // Per output binding [0..num_outputs): index into input_binding_names of the + // input it aliases (in-place KV-cache / user alias), or -1 for a normal output. + // Built at init from the blob's aliased_io. The KV buffers are threaded by + // ExecuTorch as caller-owned mutable-buffer delegate args (input AND aliased + // output): execute() binds each aliased TRT output binding to its aliased + // input's caller-provided pointer (in-place) and reflects the result into the + // delegate output EValue, which ExecuTorch's write-back copy_ then reads. + std::vector output_aliased_input_idx; + // Per input binding [0..num_inputs): true if any output aliases this input, so + // its in-place (KV/user) update must land in the caller-owned storage. Built at + // init from aliased_io; execute() uses it to reject a non-device-resident + // aliased input instead of silently staging its update into delegate scratch. + std::vector input_is_alias_target; + size_t num_aliased_outputs = 0; + int device_id = 0; + bool unified_memory = false; + std::mutex mu; + // Makes the skip-sync fast path safe to reuse: TensorRT forbids reconfiguring or + // destroying an execution context while one of its enqueues is in flight, so when + // execute() returns without an end sync it records this event; the next execute() + // and the destructor wait on it before touching exec_ctx. One event/flag pair + // suffices because a handle runs on a single thread at a time. + cudaEvent_t inflight_event = nullptr; + bool inflight_pending = false; + + ~EngineHandle(); +}; + +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 8408c13e88..8442303492 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -6,6 +6,7 @@ */ #include "torch_tensorrt/executorch/TensorRTBackend.h" +#include "EngineHandle.h" #include "torch_tensorrt/executorch/TensorRTBindingNames.h" #include "torch_tensorrt/executorch/TensorRTBlobHeader.h" #include "torch_tensorrt/executorch/WeightStreamingBudget.h" @@ -62,8 +63,33 @@ Error check_registration() { return kRegistrationResult; } +// The profile request in effect for this thread. `g_profile_index` is read only +// when the request is kPinned. Read by execute(), never by the guard itself. +thread_local ProfileRequest g_profile_request = ProfileRequest::kUnset; +thread_local int32_t g_profile_index = 0; + } // namespace +OptimizationProfileGuard::OptimizationProfileGuard(int32_t profile_index) + : prev_request_(g_profile_request), prev_index_(g_profile_index) { + g_profile_request = ProfileRequest::kPinned; + g_profile_index = profile_index; +} + +OptimizationProfileGuard::OptimizationProfileGuard(AutoTag) + : prev_request_(g_profile_request), prev_index_(g_profile_index) { + g_profile_request = ProfileRequest::kAuto; +} + +OptimizationProfileGuard OptimizationProfileGuard::automatic() { + return OptimizationProfileGuard(AutoTag{}); +} + +OptimizationProfileGuard::~OptimizationProfileGuard() { + g_profile_request = prev_request_; + g_profile_index = prev_index_; +} + void TRTLogger::log(Severity severity, const char* msg) noexcept { if (severity <= Severity::kERROR) { ET_LOG(Error, "TensorRT: %s", msg); @@ -74,12 +100,12 @@ void TRTLogger::log(Severity severity, const char* msg) noexcept { EngineHandle::~EngineHandle() { cudaSetDevice(device_id); - // A fast-path execute() may have returned with its enqueue still in flight on the - // caller's stream, still using exec_ctx and the cached staging buffers. Wait on + // An execute() may have returned with GPU work still in flight on the caller's + // stream, still using exec_ctx and the cached staging buffers. Wait on // the recorded completion event before destroying the context or freeing the // buffers. We wait on the event, not the stream, so this stays valid even if the - // caller already destroyed the stream. Non-skip executes synchronized inline, so - // inflight_pending is false there. Fall back to a device sync if no event exists. + // caller already destroyed the stream. Executes that synchronized inline cleared + // inflight_pending. Fall back to a device sync if no event exists. if (inflight_event != nullptr) { if (inflight_pending) { cudaError_t err = cudaEventSynchronize(inflight_event); @@ -176,21 +202,103 @@ Error initialize_input_profiles(EngineHandle& handle) { } } - handle.input_profile_bounds.reserve(handle.num_inputs); - for (const auto& name : handle.input_binding_names) { - InputProfileBounds bounds; - bounds.min = handle.engine->getProfileShape(name.c_str(), 0, nvinfer1::OptProfileSelector::kMIN); - bounds.max = handle.engine->getProfileShape(name.c_str(), 0, nvinfer1::OptProfileSelector::kMAX); - if (bounds.min.nbDims < 0 || bounds.max.nbDims < 0) { - ET_LOG(Error, "TensorRTBackend::init: getProfileShape failed for input '%s'", name.c_str()); - return Error::InvalidProgram; + const int32_t num_profiles = handle.engine->getNbOptimizationProfiles(); + if (num_profiles < 1) { + ET_LOG(Error, "TensorRTBackend::init: engine reports %d optimization profiles", num_profiles); + return Error::InvalidProgram; + } + + handle.profiles.bounds.resize(static_cast(num_profiles)); + for (int32_t p = 0; p < num_profiles; ++p) { + auto& bounds_for_profile = handle.profiles.bounds[static_cast(p)]; + bounds_for_profile.reserve(handle.num_inputs); + for (const auto& name : handle.input_binding_names) { + InputProfileBounds bounds; + bounds.min = handle.engine->getProfileShape(name.c_str(), p, nvinfer1::OptProfileSelector::kMIN); + bounds.max = handle.engine->getProfileShape(name.c_str(), p, nvinfer1::OptProfileSelector::kMAX); + if (bounds.min.nbDims < 0 || bounds.max.nbDims < 0) { + ET_LOG(Error, "TensorRTBackend::init: getProfileShape failed for input '%s' in profile %d", name.c_str(), p); + return Error::InvalidProgram; + } + bounds_for_profile.push_back(bounds); } - handle.input_profile_bounds.push_back(bounds); } return Error::Ok; } +std::string dims_to_string(const nvinfer1::Dims& dims) { + std::string out = "("; + for (int d = 0; d < dims.nbDims; ++d) { + if (d > 0) { + out += ", "; + } + out += std::to_string(dims.d[d]); + } + return out + ")"; +} + +// Auto-selection failing says nothing on its own about which input is out of +// range, and the offending shapes and every profile's envelope are already in +// hand. Print them so the fix does not need a rebuild with extra logging. +void log_no_profile_matches(const EngineHandle& handle, const std::vector& input_dims) { + ET_LOG( + Error, + "TensorRTBackend::execute: none of the engine's %d optimization profiles accept the input shapes; " + "fix the shapes or pin a profile with OptimizationProfileGuard", + handle.profiles.size()); + for (size_t i = 0; i < input_dims.size(); ++i) { + std::string ranges; + for (int32_t p = 0; p < handle.profiles.size(); ++p) { + const auto& bounds = handle.profiles.bounds[static_cast(p)][i]; + ranges += " profile " + std::to_string(p) + ": [" + dims_to_string(bounds.min) + ", " + + dims_to_string(bounds.max) + "]"; + } + ET_LOG( + Error, + " input '%s' is %s;%s", + handle.input_binding_names[i].c_str(), + dims_to_string(input_dims[i]).c_str(), + ranges.c_str()); + } +} + +// Redundant with profile_fits() on the auto path, which already established the +// selected profile accepts these shapes. It is the only bounds check on the +// pinned and unset paths, where select_profile() never tests fit. +Error validate_input_dims(const EngineHandle& handle, int32_t profile, const std::vector& input_dims) { + const auto& bounds = handle.profiles.bounds[static_cast(profile)]; + for (size_t i = 0; i < input_dims.size(); ++i) { + const char* name = handle.input_binding_names[i].c_str(); + const nvinfer1::Dims& dims = input_dims[i]; + if (dims.nbDims != bounds[i].min.nbDims) { + ET_LOG( + Error, + "TensorRTBackend::execute: input '%s' rank %d does not match profile %d rank %d", + name, + dims.nbDims, + profile, + bounds[i].min.nbDims); + return Error::InvalidArgument; + } + for (int d = 0; d < dims.nbDims; ++d) { + if (dims.d[d] < bounds[i].min.d[d] || dims.d[d] > bounds[i].max.d[d]) { + ET_LOG( + Error, + "TensorRTBackend::execute: input '%s' dim %d is %ld, outside profile %d bounds [%ld, %ld]", + name, + d, + static_cast(dims.d[d]), + profile, + static_cast(bounds[i].min.d[d]), + static_cast(bounds[i].max.d[d])); + return Error::InvalidArgument; + } + } + } + return Error::Ok; +} + bool is_cuda_accessible_ptr(const void* ptr) { if (ptr == nullptr) { return false; @@ -204,6 +312,33 @@ bool is_cuda_accessible_ptr(const void* ptr) { return attrs.type == cudaMemoryTypeDevice || attrs.type == cudaMemoryTypeManaged; } +// Marks the work just enqueued on `stream` as still in flight, so the next execute() +// and ~EngineHandle wait for it before they reconfigure or free exec_ctx. Recording +// over an already-recorded event just moves the marker forward, so callers can mark +// repeatedly as they enqueue more. If the event cannot be armed, drain instead: the +// caller has no other way to know the work is outstanding. +// +// A failure here usually means the enqueue itself faulted and left a sticky async +// error, so the result is propagated rather than logged and dropped: reporting Ok +// would blame the fault on some later, unrelated operator. +Error mark_inflight(EngineHandle& engine, cudaStream_t stream) { + const cudaError_t err = cudaEventRecord(engine.inflight_event, stream); + engine.inflight_pending = (err == cudaSuccess); + if (err == cudaSuccess) { + return Error::Ok; + } + ET_LOG(Error, "TensorRTBackend::execute: cudaEventRecord failed: %s", cudaGetErrorString(err)); + // A sticky fault surfaces again here, which is what turns this into a reported + // failure. A clean drain means the work really did finish, so the record failing + // cost us the marker but not the result. + const cudaError_t drain = cudaStreamSynchronize(stream); + if (drain == cudaSuccess) { + return Error::Ok; + } + ET_LOG(Error, "TensorRTBackend::execute: cudaStreamSynchronize failed: %s", cudaGetErrorString(drain)); + return Error::InvalidProgram; +} + } // namespace // --------------------------------------------------------------------------- @@ -267,8 +402,8 @@ Result TensorRTBackend::init( } // Created while device_id is current so the event belongs to the engine's device. - // It orders a later execute()/teardown after a skip-sync enqueue (see execute() - // and ~EngineHandle). Blocking-sync so the host yields instead of busy-spinning. + // It orders a later execute()/teardown after whatever execute() left running on the + // stream (see mark_inflight). Blocking-sync so the host yields, not busy-spins. cuda_err = cudaEventCreateWithFlags(&handle->inflight_event, cudaEventDisableTiming | cudaEventBlockingSync); if (cuda_err != cudaSuccess) { ET_LOG(Error, "TensorRTBackend::init: cudaEventCreateWithFlags failed: %s", cudaGetErrorString(cuda_err)); @@ -639,15 +774,18 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } // ------------------------------------------------------------------ - // 1. Bind input shapes and addresses + // 1. Collect the input shapes and settle on an optimization profile + // + // TensorRT requires setOptimizationProfileAsync() to precede setInputShape() + // for dynamic inputs, so the shapes are gathered and the profile chosen, + // validated, and switched up front rather than inside the binding loop below. // ------------------------------------------------------------------ // Device pointer each input binding was bound to; aliased outputs reuse the // pointer of the input they alias so their update lands in-place. std::vector input_bind_ptrs(num_inputs, nullptr); size_t arg_idx = 0; // running index into delegate args + std::vector input_dims(num_inputs); for (size_t i = 0; i < num_inputs; ++i) { - const std::string& name = engine->input_binding_names[i]; - EValue* arg = args[arg_idx++]; TORCHTRT_ET_CHECK_NOT_NULL( arg, Error::InvalidArgument, "TensorRTBackend::execute: input arg %zu is not a tensor", i); @@ -655,30 +793,72 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* ET_LOG(Error, "TensorRTBackend::execute: input %zu is not a tensor", i); return Error::InvalidArgument; } - - exec_aten::Tensor et_in = arg->toTensor(); - nvinfer1::Dims dims = to_trt_dims(et_in); - if (dims.nbDims > nvinfer1::Dims::MAX_DIMS) { - ET_LOG(Error, "TensorRTBackend::execute: input '%s' rank exceeds TensorRT limit", name.c_str()); + input_dims[i] = to_trt_dims(arg->toTensor()); + if (input_dims[i].nbDims > nvinfer1::Dims::MAX_DIMS) { + ET_LOG( + Error, + "TensorRTBackend::execute: input '%s' rank exceeds TensorRT limit", + engine->input_binding_names[i].c_str()); return Error::InvalidArgument; } + } - const auto& bounds = engine->input_profile_bounds[i]; - if (dims.nbDims != bounds.min.nbDims) { + // Snapshot the thread-local once: reading it again below would tie the log + // messages to the fact that only kPinned can reach them. + const ProfileRequest request = g_profile_request; + const int32_t requested_index = g_profile_index; + + int32_t profile = 0; + switch (select_profile(engine->profiles, request, requested_index, input_dims, profile)) { + case ProfileSelection::kOk: + break; + case ProfileSelection::kPinIgnoredSingleProfile: + ET_LOG( + Info, + "TensorRTBackend::execute: ignoring the pin on profile %d; this engine has one profile, " + "so it runs profile 0", + requested_index); + break; + case ProfileSelection::kRequestedProfileUnavailable: ET_LOG( Error, - "TensorRTBackend::execute: input '%s' rank %d does not match profile rank %d", - name.c_str(), - dims.nbDims, - bounds.min.nbDims); + "TensorRTBackend::execute: OptimizationProfileGuard requested profile %d but this engine has %d profile(s)", + requested_index, + engine->profiles.size()); + return Error::InvalidArgument; + case ProfileSelection::kNoProfileMatchesInputs: + log_no_profile_matches(*engine, input_dims); return Error::InvalidArgument; + } + + Error profile_err = validate_input_dims(*engine, profile, input_dims); + if (profile_err != Error::Ok) { + return profile_err; + } + if (profile != engine->profiles.active) { + if (!ctx->setOptimizationProfileAsync(profile, stream)) { + ET_LOG(Error, "TensorRTBackend::execute: setOptimizationProfileAsync(%d) failed", profile); + return Error::InvalidState; } - for (int d = 0; d < dims.nbDims; ++d) { - if (dims.d[d] < bounds.min.d[d] || dims.d[d] > bounds.max.d[d]) { - ET_LOG(Error, "TensorRTBackend::execute: input '%s' dim %d is outside profile bounds", name.c_str(), d); - return Error::InvalidArgument; - } + // The switch enqueues copies of the new profile's weights/scratch, and TensorRT + // forbids reconfiguring or destroying a context while they run. The binding loop + // below can still fail and return, and the tail marks only on the path that skips + // the sync, so mark them here rather than leaving it to either. + const Error mark_err = mark_inflight(*engine, stream); + engine->profiles.active = profile; + if (mark_err != Error::Ok) { + return mark_err; } + ET_LOG(Info, "TensorRTBackend::execute: switched to optimization profile %d", profile); + } + + // ------------------------------------------------------------------ + // 2. Bind input shapes and addresses + // ------------------------------------------------------------------ + for (size_t i = 0; i < num_inputs; ++i) { + exec_aten::Tensor et_in = args[i]->toTensor(); + const std::string& name = engine->input_binding_names[i]; + const nvinfer1::Dims& dims = input_dims[i]; if (!ctx->setInputShape(name.c_str(), dims)) { ET_LOG(Error, "TensorRTBackend::execute: setInputShape failed for '%s'", name.c_str()); @@ -752,7 +932,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } // ------------------------------------------------------------------ - // 2. Infer output shapes (requires all input shapes to be set first) + // 3. Infer output shapes (requires all input shapes to be set first) // ------------------------------------------------------------------ { const int32_t io_size = engine->engine->getNbIOTensors(); @@ -765,7 +945,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } // ------------------------------------------------------------------ - // 3. Bind output addresses + // 4. Bind output addresses // ExecuTorch pre-allocates output tensors at the maximum shape for // dynamic models. After inferShapes() TRT knows the actual output // dims, so update the ExecuTorch TensorImpl's sizes before computing @@ -906,7 +1086,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } // ------------------------------------------------------------------ - // 4. Enqueue inference on the current CUDA stream + // 5. Enqueue inference on the current CUDA stream // ------------------------------------------------------------------ if (!ctx->enqueueV3(stream)) { ET_LOG( @@ -941,9 +1121,11 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // buffer, which the caller may reuse once we return), or no caller stream is // active (preserve the historical "results ready on return" behavior). // Otherwise (caller stream + all I/O device-resident) leave the work enqueued so - // it composes with the caller's later GPU work, and record inflight_event so the - // next execute() and the destructor wait before reusing/freeing exec_ctx. The D2H - // copies live in the must_sync branch: an output staged to host always sets + // it composes with the caller's later GPU work, and mark it so the next execute() + // and the destructor wait before reusing/freeing exec_ctx. Only that path marks: + // draining is the stronger guarantee of the two, so recording an event the sync + // below would immediately consume would be pure overhead. The D2H copies live in + // the must_sync branch: an output staged to host always sets // output_staged_to_host, so outputs_needing_copy is empty on the skip path. // An aliased reflect enqueues the engine's in-place update into the delegate // output EValue on `stream`; ExecuTorch's buffer-mutation copy_ reads that EValue @@ -990,7 +1172,12 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* engine->inflight_pending = false; return Error::InvalidProgram; } - engine->inflight_pending = true; + } + cuda_err = cudaStreamSynchronize(stream); + engine->inflight_pending = false; + if (cuda_err != cudaSuccess) { + ET_LOG(Error, "TensorRTBackend::execute: cudaStreamSynchronize failed: %s", cudaGetErrorString(cuda_err)); + return Error::InvalidProgram; } return Error::Ok; } diff --git a/examples/dynamo/multi_optimization_profiles.py b/examples/dynamo/multi_optimization_profiles.py index df16f9b06c..c10006b531 100644 --- a/examples/dynamo/multi_optimization_profiles.py +++ b/examples/dynamo/multi_optimization_profiles.py @@ -22,9 +22,11 @@ select the active profile per call (by index, or ``"auto"``). This example compiles `google/gemma-3-1b-it -`_ **twice** -- once with a single -profile and once with separate prefill/decode profiles -- and compares the decode -and prefill latency of the two engines. +`_ **once** into a two-profile +engine and then runs the same engine two ways: every call on the prefill profile +(which accepts ``seq == 1`` as well, so it is what a conventional single-profile +engine gives you) versus each phase on its own profile. One engine, one set of +weights; the only difference is which profile is active when the call runs. .. note:: @@ -44,10 +46,9 @@ # Imports and Setup # ^^^^^^^^^^^^^^^^^^ # -# The HuggingFace attention path needs a TensorRT-friendly SDPA lowering. The -# reusable LLM helpers ``register_sdpa`` (a Gemma-3-specific SDPA pass) and -# ``export_llm`` live under ``tools/llm`` in the Torch-TensorRT repo, so we add -# that directory to ``sys.path``. +# ``export_llm``, a reusable helper that traces a decoder over a dynamic +# sequence length, lives under ``tools/llm`` in the Torch-TensorRT repo, so we +# add that directory to ``sys.path``. import sys import timeit @@ -74,8 +75,10 @@ # ^^^^^^^^^^^^^^ # # Load with ``use_cache=False`` (this example recomputes over the full sequence -# rather than using a KV cache, which keeps the export simple) and the ``sdpa`` -# attention implementation, then register the Gemma-3 SDPA lowering pass. +# rather than using a KV cache, which keeps the export simple). The ``sdpa`` +# attention implementation makes HuggingFace emit +# ``scaled_dot_product_attention``, which Torch-TensorRT converts to a single +# TensorRT attention layer. def load_model(): from transformers import AutoModelForCausalLM @@ -91,9 +94,6 @@ def load_model(): .cuda() .to(torch.float16) ) - from torchtrt_ext import register_sdpa - - register_sdpa.enable_sdpa_converter(MODEL_ID, model.config) return model @@ -140,21 +140,27 @@ def make_inputs(seq_len: int): ] # %% -# Export Once, Compile Twice -# ^^^^^^^^^^^^^^^^^^^^^^^^^^^ +# Export Once, Compile Once +# ^^^^^^^^^^^^^^^^^^^^^^^^^^ # -# ``export_llm`` traces the model over a dynamic ``seq`` range. We reuse the -# exported program for both the single-profile baseline (tuned at the prefill -# length, the conventional choice) and the multi-profile engine. +# ``export_llm`` traces the model over a dynamic ``seq`` range, and one compile +# turns that into one engine holding both profiles. No separate single-profile +# build is needed for the baseline: the prefill profile already accepts +# ``seq == 1``, so running every call on it reproduces what a single-profile +# engine does, without a second compile or a second set of weights to keep +# honest. from utils import export_llm # noqa: E402 example_ids, _ = make_inputs(PREFILL_SEQ) with torch.inference_mode(): exported = export_llm(model, example_ids, min_seq_len=1, max_seq_len=MAX_SEQ) +print("Compiling multi-profile engine (decode + prefill) ...") # ``offload_module_to_cpu`` must stay False here: it is currently incompatible # with the multi-profile ``Input(profiles=...)`` path (CPU/CUDA device mismatch). -common = dict( +trt_model = torch_tensorrt.dynamo.compile( + exported, + arg_inputs=multi_profile_inputs, use_fp32_acc=True, disable_tf32=True, offload_module_to_cpu=False, @@ -163,17 +169,6 @@ def make_inputs(seq_len: int): device=DEVICE, ) -print("Compiling single-profile engine (tuned at prefill length) ...") -bench_ids, bench_pos = make_inputs(PREFILL_SEQ) -trt_single = torch_tensorrt.dynamo.compile( - exported, inputs=[bench_ids, bench_pos], **common -) - -print("Compiling multi-profile engine (decode + prefill) ...") -trt_multi = torch_tensorrt.dynamo.compile( - exported, arg_inputs=multi_profile_inputs, **common -) - # %% # Correctness @@ -194,10 +189,10 @@ def logits(out): ref_decode = logits(model(decode_ids, position_ids=decode_pos)) ref_prefill = logits(model(prefill_ids, position_ids=prefill_pos)) - with optimization_profile(trt_multi, DECODE_IDX): - trt_decode = logits(trt_multi(decode_ids, decode_pos)) - with optimization_profile(trt_multi, PREFILL_IDX): - trt_prefill = logits(trt_multi(prefill_ids, prefill_pos)) + with optimization_profile(trt_model, DECODE_IDX): + trt_decode = logits(trt_model(decode_ids, decode_pos)) + with optimization_profile(trt_model, PREFILL_IDX): + trt_prefill = logits(trt_model(prefill_ids, prefill_pos)) def top1_match(a, b): @@ -212,9 +207,10 @@ def top1_match(a, b): # Latency Comparison # ^^^^^^^^^^^^^^^^^^^ # -# Time each regime on each engine. For the multi-profile engine we pin the -# matching profile around the loop (the realistic serving pattern). We report the -# min over several rounds to reduce noise. +# Decode is timed twice against the one engine: once on the prefill profile and +# once on its own. The profile is pinned around the whole loop rather than per +# call, which is the realistic serving pattern and keeps profile switches out of +# the measurement. We report the min over several rounds to reduce noise. def benchmark(run, iters: int = 50, warmup: int = 20, rounds: int = 3) -> float: for _ in range(warmup): run() @@ -230,28 +226,27 @@ def benchmark(run, iters: int = 50, warmup: int = 20, rounds: int = 3) -> float: with torch.inference_mode(): - single_decode = benchmark(lambda: trt_single(decode_ids, decode_pos)) - single_prefill = benchmark(lambda: trt_single(prefill_ids, prefill_pos)) - with optimization_profile(trt_multi, DECODE_IDX): - multi_decode = benchmark(lambda: trt_multi(decode_ids, decode_pos)) - with optimization_profile(trt_multi, PREFILL_IDX): - multi_prefill = benchmark(lambda: trt_multi(prefill_ids, prefill_pos)) + with optimization_profile(trt_model, PREFILL_IDX): + decode_on_prefill = benchmark(lambda: trt_model(decode_ids, decode_pos)) + prefill_on_prefill = benchmark(lambda: trt_model(prefill_ids, prefill_pos)) + with optimization_profile(trt_model, DECODE_IDX): + decode_on_decode = benchmark(lambda: trt_model(decode_ids, decode_pos)) # %% -# Results. Decode is the win: the multi-profile engine dedicates a *static* -# profile (``seq`` pinned to 1) to decode, so TensorRT specializes that path -# instead of serving it from kernels tuned for the long prefill length. Prefill -# is unchanged (both engines tune it at the same ``opt``). +# Results. Decode is the win: the decode profile pins ``seq`` to 1, so TensorRT +# specializes that path instead of serving it from kernels tuned for the long +# prefill length. 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, so it +# is the same call in both scenarios. print("\nPer-call latency (ms), batch=1") -print(f"{'regime':<20}{'single-profile':>16}{'multi-profile':>16}{'speedup':>10}") -print("-" * 62) -print( - f"{f'decode (seq={DECODE_SEQ})':<20}{single_decode:>16.3f}" - f"{multi_decode:>16.3f}{single_decode / multi_decode:>9.2f}x" -) +print(f"{'call':<24}{'active profile':>18}{'ms':>10}") +print("-" * 52) +print(f"{f'decode (seq={DECODE_SEQ})':<24}{'prefill':>18}{decode_on_prefill:>10.3f}") +print(f"{f'decode (seq={DECODE_SEQ})':<24}{'decode':>18}{decode_on_decode:>10.3f}") +print(f"{f'prefill (seq={PREFILL_SEQ})':<24}{'prefill':>18}{prefill_on_prefill:>10.3f}") print( - f"{f'prefill (seq={PREFILL_SEQ})':<20}{single_prefill:>16.3f}" - f"{multi_prefill:>16.3f}{single_prefill / multi_prefill:>9.2f}x" + f"\nGiving decode its own profile: {decode_on_prefill / decode_on_decode:.2f}x " + f"faster per token ({decode_on_prefill - decode_on_decode:+.3f} ms)" ) # %% @@ -262,6 +257,9 @@ def benchmark(run, iters: int = 50, warmup: int = 20, rounds: int = 3) -> float: # ``profiles=[{min_shape, opt_shape, max_shape}, ...]`` # (one per dynamic model input -- here ``input_ids`` and ``position_ids``). # - One export + one engine; each profile gets its own TensorRT kernel tuning. +# - A profile whose range covers the other regime doubles as the baseline: +# pinning every call to the prefill profile shows what a single-profile engine +# would do, with no second compile and no second set of weights. # - Select at runtime by **index** (``optimization_profile(m, i)``) or let # ``"auto"`` pick the first profile that fits the input shapes. # - Dedicating a static ``seq == 1`` profile to decode lets TensorRT tune that diff --git a/examples/executorch_reference_runner/BUILD b/examples/executorch_reference_runner/BUILD index 62cd8d6d47..4e3071b1ea 100644 --- a/examples/executorch_reference_runner/BUILD +++ b/examples/executorch_reference_runner/BUILD @@ -2,6 +2,9 @@ load("@rules_cc//cc:defs.bzl", "cc_binary") package(default_visibility = ["//visibility:public"]) +# Decides what ships in libtorchtrt.tar.gz. CMakeLists.txt references every +# runner source, so a file missing here aborts the whole configure step for +# anyone building from the tarball, not just its own target. filegroup( name = "source_files", srcs = [ @@ -10,6 +13,8 @@ filegroup( "kv_cache_decode_check.cpp", "load_model.py", "main.cpp", + "multi_profile_benchmark.cpp", + "multi_profile_main.cpp", ], ) diff --git a/examples/executorch_reference_runner/CMakeLists.txt b/examples/executorch_reference_runner/CMakeLists.txt index 3c30ef81af..33eb4fe380 100644 --- a/examples/executorch_reference_runner/CMakeLists.txt +++ b/examples/executorch_reference_runner/CMakeLists.txt @@ -81,3 +81,25 @@ target_link_libraries( executorch::kernels torchtrt::executorch_backend CUDA::cudart) + +add_executable(example_executorch_multi_profile_runner multi_profile_main.cpp) +target_link_libraries( + example_executorch_multi_profile_runner + PRIVATE + executorch + executorch::backends + executorch::extensions + executorch::kernels + torchtrt::executorch_backend) + +add_executable( + example_executorch_multi_profile_benchmark multi_profile_benchmark.cpp +) +target_link_libraries( + example_executorch_multi_profile_benchmark + PRIVATE + executorch + executorch::backends + executorch::extensions + executorch::kernels + torchtrt::executorch_backend) diff --git a/examples/executorch_reference_runner/README.md b/examples/executorch_reference_runner/README.md index 1b1ccba454..5b2396e2fc 100644 --- a/examples/executorch_reference_runner/README.md +++ b/examples/executorch_reference_runner/README.md @@ -1,7 +1,13 @@ # Torch-TensorRT ExecuTorch Reference Runner -This directory contains a minimal C++ reference runner for loading and executing -Torch-TensorRT compiled models saved in ExecuTorch `.pte` format. +This directory contains minimal C++ reference runners for loading and +executing Torch-TensorRT compiled models saved in ExecuTorch `.pte` format: + +| 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 | The `.pte` file contains an ExecuTorch program with embedded TensorRT engine payloads. The runner links the TensorRT ExecuTorch backend, loads the `.pte` @@ -17,6 +23,14 @@ You can also generate a sample `.pte` from the Torch-TensorRT source tree: ```bash python examples/torchtrt_executorch_example/export_static_shape.py --model_path=model.pte + +# Two-profile Gemma-3 engine for the multi-profile runner below. 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 @@ -207,3 +221,117 @@ 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. + +## 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: + +```cpp +#include + +using torch_tensorrt::executorch_backend::OptimizationProfileGuard; + +constexpr int32_t kDecodeProfile = 0; +constexpr int32_t kPrefillProfile = 1; + +executorch::extension::Module module("model_gemma3_multi_profile.pte"); +{ + OptimizationProfileGuard profile_guard(kPrefillProfile); + auto result = module.forward(prefill_inputs); +} +{ + OptimizationProfileGuard profile_guard(kDecodeProfile); + auto result = module.forward(decode_inputs); +} +``` + +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. + +Build and run: + +```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: + +``` +Per-call latency (ms), batch=1 +call active profile ms +---------------------------------------------------- +decode (seq=1) prefill 6.415 +decode (seq=1) decode 4.981 +prefill (seq=128) prefill 8.438 + +Giving decode its own profile: 1.29x faster per token (+1.434 ms) +``` + +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, which is what a +single-profile engine gives you), and once with each phase pinned to its own +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 +``` + +On the real `google/gemma-3-1b-it` (exported with `--weights +google/gemma-3-1b-it`) on an idle A40, decode is **1.29x faster** on its own +profile (6.42 ms down to 4.97 ms per token) while a switch costs ~3.6 ms, +charged to whichever call switches. That breaks even after about five decode +steps. End to end, one prefill plus 16 decode steps drops from 112.0 ms to +96.2 ms (14.1% faster), and a 64-step round from 420.6 ms to 336.2 ms (20.1%). + +Both numbers shrink with the model. The mini Gemma-3 exported by default is +small enough that decode gains only 0.02 ms (1.12x) against a 0.48 ms switch, so +it takes ~46 decode steps to break even and a 16-step round is actually 4-5% +slower with switching. Use it to exercise the API, and `--weights` 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. diff --git a/examples/executorch_reference_runner/multi_profile_benchmark.cpp b/examples/executorch_reference_runner/multi_profile_benchmark.cpp new file mode 100644 index 0000000000..aa37660098 --- /dev/null +++ b/examples/executorch_reference_runner/multi_profile_benchmark.cpp @@ -0,0 +1,361 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + * + * What per-call optimization profile selection is worth. + * + * Runs the same prefill/decode serving loop two ways against the one + * two-profile engine from + * examples/torchtrt_executorch_example/export_multi_profile.py: + * + * prefill-only - every call pinned to the prefill profile. The prefill + * profile accepts seq == 1, so decode runs on kernels + * TensorRT tuned for a 128-token prompt. + * switching - prefill pinned to the prefill profile and each decode step + * pinned to the decode profile, whose seq is pinned to 1. + * + * One engine, one set of weights, one export; the only difference is which + * profile is loaded when the call runs. Decode is where the difference should + * show, since that is the phase whose kernels the prefill profile mistunes. + * + * Measurement notes: + * - The two configurations are interleaved in short blocks so that drift and + * any other tenant on the GPU hit both roughly equally. + * - Interference can only add time, so the low percentiles are the signal. + * min and p10 are what to read; the tail says how busy the machine was. + * - The first call of each block is discarded: it inherits whichever profile + * the previous block left loaded, so it can carry a switch the block is + * not meant to be measuring. + * + * Usage: + * example_executorch_multi_profile_benchmark --model_path=model.pte + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include + +using executorch::extension::Module; +using executorch::runtime::EValue; +using torch_tensorrt::executorch_backend::OptimizationProfileGuard; + +namespace { + +constexpr int32_t kDecodeProfile = 0; +constexpr int32_t kPrefillProfile = 1; + +using Clock = std::chrono::steady_clock; + +const char* get_flag(int argc, char** argv, const char* flag, const char* def) { + const size_t n = strlen(flag); + for (int i = 1; i < argc; ++i) { + if (strncmp(argv[i], flag, n) == 0 && argv[i][n] == '=') { + return argv[i] + n + 1; + } + } + return def; +} + +// atoi() returns 0 for garbage, so an unusable value would otherwise pass for a +// deliberate 0 and surface much later as an empty sample set or a nan in the +// summary. Reject anything below `min` here, where the flag name is still known. +bool get_int_flag(int argc, char** argv, const char* flag, int def, int min, int& out) { + const char* raw = get_flag(argc, argv, flag, nullptr); + out = raw == nullptr ? def : atoi(raw); + if (out < min) { + ET_LOG(Error, "%s must be at least %d, got '%s'", flag, min, raw == nullptr ? "" : raw); + return false; + } + return true; +} + +// One [1, seq] index tensor. The dtype comes from the .pte's method signature +// rather than being assumed: the backend binds tensor pointers straight to +// TensorRT without converting, so a mismatch here is silent corruption. +// +// Token ids stay inside [1, vocab). An id past the end of the embedding table +// is not an error in TensorRT -- it gathers zero -- so it would not perturb the +// timings measured here, but it would make this workload unrepresentative of +// the one it is meant to stand in for. +class IndexTensor { + public: + IndexTensor(int32_t seq, exec_aten::ScalarType dtype, bool positions, int64_t vocab) + : sizes_{1, seq}, + dim_order_{0, 1}, + strides_{seq, 1}, + data_(static_cast(seq) * (dtype == exec_aten::ScalarType::Long ? 8 : 4)), + impl_(dtype, 2, sizes_.data(), data_.data(), dim_order_.data(), strides_.data()) { + for (int32_t i = 0; i < seq; ++i) { + const int64_t v = positions ? i : (1 + (static_cast(i) * 7919) % (vocab - 1)); + if (dtype == exec_aten::ScalarType::Long) { + reinterpret_cast(data_.data())[i] = v; + } else { + reinterpret_cast(data_.data())[i] = static_cast(v); + } + } + } + + EValue evalue() { + return EValue(exec_aten::Tensor(&impl_)); + } + + private: + std::vector sizes_; + std::vector dim_order_; + std::vector strides_; + std::vector data_; + exec_aten::TensorImpl impl_; +}; + +// The (input_ids, position_ids) pair for one sequence length, held so the +// EValue vector handed to forward() can be reused without reallocating. +class Step { + public: + Step(int32_t seq, exec_aten::ScalarType dtype, int64_t vocab) + : ids_(seq, dtype, false, vocab), + positions_(seq, dtype, true, vocab), + args_{ids_.evalue(), positions_.evalue()} {} + + const std::vector& args() const { + return args_; + } + + private: + IndexTensor ids_; + IndexTensor positions_; + std::vector args_; +}; + +struct Stats { + size_t n = 0; + double min = 0.0; + double p10 = 0.0; + double p25 = 0.0; + double median = 0.0; + double p90 = 0.0; +}; + +double percentile(const std::vector& sorted, double q) { + return sorted[static_cast(q * static_cast(sorted.size() - 1))]; +} + +Stats summarize(std::vector samples) { + Stats s; + if (samples.empty()) { + return s; + } + std::sort(samples.begin(), samples.end()); + s.n = samples.size(); + s.min = samples.front(); + s.p10 = percentile(samples, 0.10); + s.p25 = percentile(samples, 0.25); + s.median = percentile(samples, 0.50); + s.p90 = percentile(samples, 0.90); + return s; +} + +void print_stats(const char* label, const Stats& s) { + printf( + " %-30s n=%-5zu min=%8.3f p10=%8.3f p25=%8.3f median=%8.3f p90=%8.3f\n", + label, + s.n, + s.min, + s.p10, + s.p25, + s.median, + s.p90); +} + +// One forward under `profile`, timed end to end. Inputs are host tensors, so +// execute() stages H2D and synchronizes before returning; the interval covers +// the whole round trip. Returns milliseconds, or a negative value on failure. +double timed_forward(Module& module, const std::vector& args, int32_t profile) { + const auto t0 = Clock::now(); + double elapsed_ms = 0.0; + { + OptimizationProfileGuard profile_guard(profile); + auto result = module.forward(args); + elapsed_ms = std::chrono::duration(Clock::now() - t0).count(); + if (!result.ok()) { + ET_LOG(Error, "forward() failed: 0x%" PRIx32, static_cast(result.error())); + return -1.0; + } + } + return elapsed_ms; +} + +struct WorkloadResult { + std::vector prefill_ms; + std::vector decode_ms; + double wall_ms = 0.0; +}; + +// `rounds` iterations of one prefill followed by `decode_steps` decode steps. +bool run_block( + Module& module, + const Step& prefill, + const Step& decode, + int32_t prefill_profile, + int32_t decode_profile, + int rounds, + int decode_steps, + WorkloadResult& out) { + const auto start = Clock::now(); + for (int r = 0; r < rounds; ++r) { + const double p = timed_forward(module, prefill.args(), prefill_profile); + if (p < 0.0) { + return false; + } + if (r != 0) { // first call of a block inherits the previous block's profile + out.prefill_ms.push_back(p); + } + for (int d = 0; d < decode_steps; ++d) { + const double t = timed_forward(module, decode.args(), decode_profile); + if (t < 0.0) { + return false; + } + if (r != 0 || d != 0) { + out.decode_ms.push_back(t); + } + } + } + out.wall_ms += std::chrono::duration(Clock::now() - start).count(); + return true; +} + +void compare(const char* what, const Stats& prefill_only, const Stats& switching) { + const double d_min = prefill_only.min - switching.min; + const double d_p10 = prefill_only.p10 - switching.p10; + printf( + " %-30s %+8.3f ms (min) %+8.3f ms (p10) %.2fx (min)\n", + what, + d_min, + d_p10, + switching.min > 0.0 ? prefill_only.min / switching.min : 0.0); +} + +} // namespace + +int main(int argc, char** argv) { + executorch::runtime::runtime_init(); + + const char* model_path = get_flag(argc, argv, "--model_path", "model_gemma3_multi_profile.pte"); + int prefill_seq = 0; + int blocks = 0; + int block_rounds = 0; + int decode_steps = 0; + int warmup = 0; + if (!get_int_flag(argc, argv, "--prefill_seq", 128, 1, prefill_seq) || + !get_int_flag(argc, argv, "--blocks", 10, 1, blocks) || + // run_block() discards the first round of each block as warm-in, so a + // single round would leave prefill with no samples at all. + !get_int_flag(argc, argv, "--block_rounds", 3, 2, block_rounds) || + !get_int_flag(argc, argv, "--decode_steps", 16, 1, decode_steps) || + !get_int_flag(argc, argv, "--warmup", 20, 0, warmup)) { + return 1; + } + + Module module(model_path); + + // Take the input dtype from the method itself rather than assuming it: + // Torch-TensorRT may narrow int64 indices to int32 during lowering. + const auto meta = module.method_meta("forward"); + if (!meta.ok()) { + ET_LOG(Error, "could not read method_meta: 0x%" PRIx32, static_cast(meta.error())); + return 1; + } + const auto input0 = meta->input_tensor_meta(0); + if (!input0.ok()) { + ET_LOG(Error, "could not read input 0 metadata"); + return 1; + } + const exec_aten::ScalarType dtype = input0->scalar_type(); + + // The method returns the last position's logits, so the width of its output is + // the vocabulary the token ids below have to stay inside. + const auto output0 = meta->output_tensor_meta(0); + if (!output0.ok() || output0->sizes().empty()) { + ET_LOG(Error, "could not read output 0 metadata"); + return 1; + } + const auto logits_sizes = output0->sizes(); + const int64_t vocab = logits_sizes[logits_sizes.size() - 1]; + if (vocab < 2) { + ET_LOG(Error, "logits width %" PRId64 " is too small to draw token ids from", vocab); + return 1; + } + + Step prefill(prefill_seq, dtype, vocab); + Step decode(1, dtype, vocab); + + printf("model : %s\n", model_path); + printf("inputs : 2 x [1, seq] %s\n", dtype == exec_aten::ScalarType::Long ? "int64" : "int32"); + printf( + "workload : %d interleaved blocks x %d rounds x (1 prefill seq=%d + %d decode seq=1) per config\n", + blocks, + block_rounds, + prefill_seq, + decode_steps); + printf("units : milliseconds per module.forward(); read min/p10, the tail is machine noise\n\n"); + + for (int i = 0; i < warmup; ++i) { + if (timed_forward(module, prefill.args(), kPrefillProfile) < 0.0 || + timed_forward(module, decode.args(), kDecodeProfile) < 0.0) { + return 1; + } + } + + WorkloadResult prefill_only; + WorkloadResult switching; + for (int b = 0; b < blocks; ++b) { + if (!run_block( + module, prefill, decode, kPrefillProfile, kPrefillProfile, block_rounds, decode_steps, prefill_only) || + !run_block(module, prefill, decode, kPrefillProfile, kDecodeProfile, block_rounds, decode_steps, switching)) { + return 1; + } + } + + const Stats po_prefill = summarize(prefill_only.prefill_ms); + const Stats po_decode = summarize(prefill_only.decode_ms); + const Stats sw_prefill = summarize(switching.prefill_ms); + const Stats sw_decode = summarize(switching.decode_ms); + + printf("prefill-only (every call on the prefill profile)\n"); + print_stats("prefill (seq=128)", po_prefill); + print_stats("decode (seq=1)", po_decode); + printf("\nswitching (each phase on its own profile)\n"); + print_stats("prefill (seq=128)", sw_prefill); + print_stats("decode (seq=1)", sw_decode); + + printf("\nwhat switching bought (positive = switching is faster)\n"); + compare("decode", po_decode, sw_decode); + compare("prefill", po_prefill, sw_prefill); + printf( + " %-30s prefill-only=%9.1f ms switching=%9.1f ms %+.1f%%\n", + "wall time (see note)", + prefill_only.wall_ms, + switching.wall_ms, + // Same orientation as compare(): positive means switching came out ahead. + 100.0 * (prefill_only.wall_ms - switching.wall_ms) / prefill_only.wall_ms); + printf( + "\nnote: wall time is contention-prone, and interleaving charges each prefill-only block\n" + " one switch back that a single-profile engine would never pay, so it reads a little\n" + " kinder to switching than reality. multi_profile_main.cpp times a clean request.\n"); + + return 0; +} diff --git a/examples/executorch_reference_runner/multi_profile_main.cpp b/examples/executorch_reference_runner/multi_profile_main.cpp new file mode 100644 index 0000000000..35a9c928f0 --- /dev/null +++ b/examples/executorch_reference_runner/multi_profile_main.cpp @@ -0,0 +1,342 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + * + * Selecting a TensorRT optimization profile per call, through the high-level + * ExecuTorch Module API. + * + * Pairs with examples/torchtrt_executorch_example/export_multi_profile.py, + * which writes a two-profile Gemma-3 engine taking two [1, seq] index tensors + * (input_ids and position_ids) and returning the last position's logits: + * + * profile 0 -> decode, seq == 1 + * profile 1 -> prefill, seq in [1, 256], tuned at 128 + * + * Ends with per-call latency for decode on each profile, the same comparison + * examples/dynamo/multi_optimization_profiles.py makes through the Python + * runtime. For latency distributions see multi_profile_benchmark.cpp. + * + * Usage: + * example_executorch_multi_profile_runner --model_path=model.pte + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include + +using executorch::extension::Module; +using executorch::runtime::Error; +using executorch::runtime::EValue; +using torch_tensorrt::executorch_backend::OptimizationProfileGuard; + +namespace { + +constexpr int32_t kDecodeProfile = 0; +constexpr int32_t kPrefillProfile = 1; +constexpr int32_t kPrefillSeq = 128; +constexpr int32_t kMaxSeq = 256; + +// Timing loop at the end, matching examples/dynamo/multi_optimization_profiles.py. +constexpr int kWarmup = 20; +constexpr int kIters = 50; +constexpr int kRounds = 3; + +using Clock = std::chrono::steady_clock; + +const char* get_flag(int argc, char** argv, const char* flag, const char* def) { + const size_t n = strlen(flag); + for (int i = 1; i < argc; ++i) { + if (strncmp(argv[i], flag, n) == 0 && argv[i][n] == '=') { + return argv[i] + n + 1; + } + } + return def; +} + +// One [1, seq] index tensor. The dtype comes from the .pte's method signature +// rather than being assumed: the backend binds tensor pointers straight to +// TensorRT without converting, so a mismatch here is silent corruption. +// +// Token ids stay inside [1, vocab): TensorRT stores zero for an out-of-bounds +// gather instead of failing, so an id past the end of the embedding table would +// quietly substitute a zero embedding and make the predictions below +// meaningless. `vocab` comes from the model, so this holds for the small +// exported default as much as for the real 1B one. +class IndexTensor { + public: + IndexTensor(int32_t seq, exec_aten::ScalarType dtype, bool positions, int64_t vocab) + : sizes_{1, seq}, + dim_order_{0, 1}, + strides_{seq, 1}, + data_(static_cast(seq) * (dtype == exec_aten::ScalarType::Long ? 8 : 4)), + impl_(dtype, 2, sizes_.data(), data_.data(), dim_order_.data(), strides_.data()) { + for (int32_t i = 0; i < seq; ++i) { + const int64_t v = positions ? i : (1 + (static_cast(i) * 7919) % (vocab - 1)); + if (dtype == exec_aten::ScalarType::Long) { + reinterpret_cast(data_.data())[i] = v; + } else { + reinterpret_cast(data_.data())[i] = static_cast(v); + } + } + } + + EValue evalue() { + return EValue(exec_aten::Tensor(&impl_)); + } + + private: + std::vector sizes_; + std::vector dim_order_; + std::vector strides_; + std::vector data_; + exec_aten::TensorImpl impl_; +}; + +// Owns the (input_ids, position_ids) pair for one sequence length. +class Step { + public: + Step(int32_t seq, exec_aten::ScalarType dtype, int64_t vocab) + : ids_(seq, dtype, false, vocab), + positions_(seq, dtype, true, vocab), + args_{ids_.evalue(), positions_.evalue()} {} + + const std::vector& args() const { + return args_; + } + + private: + IndexTensor ids_; + IndexTensor positions_; + std::vector args_; +}; + +// The method returns the last position's logits, so the argmax is the token the +// model would emit next. Printing it makes it obvious when a profile switch +// changes shapes but not results. +void print_prediction(const char* label, const std::vector& outputs) { + if (outputs.empty() || !outputs[0].isTensor()) { + return; + } + exec_aten::Tensor t = outputs[0].toTensor(); + // const_data_ptr() is an unchecked cast, so reading a dtype we did not plan + // for walks the buffer at the wrong stride and runs off the end. Name the ones + // handled and skip anything else rather than printing corrupt numbers. + const exec_aten::ScalarType dtype = t.scalar_type(); + if (dtype != exec_aten::ScalarType::Float && dtype != exec_aten::ScalarType::Half && + dtype != exec_aten::ScalarType::BFloat16) { + ET_LOG(Info, "%s: logits dtype %d not handled by this example; skipping", label, static_cast(dtype)); + return; + } + double best = -1e30; + int64_t best_idx = -1; + for (int64_t i = 0; i < t.numel(); ++i) { + double v = 0.0; + if (dtype == exec_aten::ScalarType::Half) { + v = static_cast(t.const_data_ptr()[i]); + } else if (dtype == exec_aten::ScalarType::BFloat16) { + v = static_cast(t.const_data_ptr()[i]); + } else { + v = static_cast(t.const_data_ptr()[i]); + } + if (v > best) { + best = v; + best_idx = i; + } + } + fprintf(stderr, "%-28s logits=[", label); + for (ssize_t d = 0; d < t.dim(); ++d) { + fprintf(stderr, "%d%s", static_cast(t.size(d)), d + 1 < t.dim() ? "," : ""); + } + fprintf(stderr, "] next_token=%" PRId64 "\n", best_idx); +} + +// Runs one forward with whatever profile guard the caller has in scope. +bool run_guarded(Module& module, const char* label, const Step& step) { + auto result = module.forward(step.args()); + + if (!result.ok()) { + ET_LOG(Error, "%s: forward() failed: 0x%" PRIx32, label, static_cast(result.error())); + return false; + } + print_prediction(label, result.get()); + return true; +} + +// Runs one forward with `profile` pinned. The guard applies to every TensorRT +// delegate this thread executes while it is in scope. It stores the request +// only; each delegate switches inside its own execute(), on the stream that +// execute() already selected. +bool run(Module& module, const char* label, int32_t profile, const Step& step) { + OptimizationProfileGuard profile_guard(profile); + return run_guarded(module, label, step); +} + +// Same, but lets each delegate choose from the input shapes. +bool run_auto(Module& module, const char* label, const Step& step) { + auto profile_guard = OptimizationProfileGuard::automatic(); + return run_guarded(module, label, step); +} + +// Mean milliseconds per forward, best of kRounds. The profile is pinned around +// the whole loop rather than per call, which is the realistic serving pattern +// and keeps profile switches out of the measurement. Inputs are host tensors, +// so each forward() stages H2D and synchronizes before returning and the +// interval covers the whole round trip. Negative on failure. +double benchmark(Module& module, const Step& step, int32_t profile) { + OptimizationProfileGuard profile_guard(profile); + + for (int i = 0; i < kWarmup; ++i) { + if (!module.forward(step.args()).ok()) { + return -1.0; + } + } + double best = std::numeric_limits::infinity(); + for (int round = 0; round < kRounds; ++round) { + const auto start = Clock::now(); + for (int i = 0; i < kIters; ++i) { + if (!module.forward(step.args()).ok()) { + return -1.0; + } + } + const double ms = std::chrono::duration(Clock::now() - start).count(); + best = std::min(best, ms / kIters); + } + return best; +} + +// Decode is timed twice against the one engine: once on the prefill profile, +// which accepts seq == 1 and so runs it on kernels tuned for a kPrefillSeq +// prompt (what a single-profile engine gives you), and once on its own profile. +// Prefill appears once because the decode profile does not accept a kPrefillSeq +// input at all, so prefill has only one profile it can run on. +bool report_latency(Module& module, const Step& prefill, const Step& decode) { + const double decode_on_prefill = benchmark(module, decode, kPrefillProfile); + const double decode_on_decode = benchmark(module, decode, kDecodeProfile); + const double prefill_on_prefill = benchmark(module, prefill, kPrefillProfile); + if (decode_on_prefill < 0.0 || decode_on_decode < 0.0 || prefill_on_prefill < 0.0) { + ET_LOG(Error, "latency benchmark: forward() failed"); + return false; + } + + fprintf(stderr, "\nPer-call latency (ms), batch=1\n"); + fprintf(stderr, "%-24s%18s%10s\n", "call", "active profile", "ms"); + fprintf(stderr, "----------------------------------------------------\n"); + fprintf(stderr, "%-24s%18s%10.3f\n", "decode (seq=1)", "prefill", decode_on_prefill); + fprintf(stderr, "%-24s%18s%10.3f\n", "decode (seq=1)", "decode", decode_on_decode); + char prefill_label[32]; + snprintf(prefill_label, sizeof(prefill_label), "prefill (seq=%d)", kPrefillSeq); + fprintf(stderr, "%-24s%18s%10.3f\n", prefill_label, "prefill", prefill_on_prefill); + fprintf( + stderr, + "\nGiving decode its own profile: %.2fx faster per token (%+.3f ms)\n", + decode_on_prefill / decode_on_decode, + decode_on_prefill - decode_on_decode); + return true; +} + +} // namespace + +int main(int argc, char** argv) { + executorch::runtime::runtime_init(); + + const char* model_path = get_flag(argc, argv, "--model_path", "model_gemma3_multi_profile.pte"); + Module module(model_path); + + const auto meta = module.method_meta("forward"); + if (!meta.ok()) { + ET_LOG(Error, "could not read method_meta: 0x%" PRIx32, static_cast(meta.error())); + return 1; + } + const auto input0 = meta->input_tensor_meta(0); + if (!input0.ok()) { + ET_LOG(Error, "could not read input 0 metadata"); + return 1; + } + const exec_aten::ScalarType dtype = input0->scalar_type(); + + // The method returns the last position's logits, so the width of its output is + // the vocabulary the token ids below have to stay inside. + const auto output0 = meta->output_tensor_meta(0); + if (!output0.ok() || output0->sizes().empty()) { + ET_LOG(Error, "could not read output 0 metadata"); + return 1; + } + const auto logits_sizes = output0->sizes(); + const int64_t vocab = logits_sizes[logits_sizes.size() - 1]; + if (vocab < 2) { + ET_LOG(Error, "logits width %" PRId64 " is too small to draw token ids from", vocab); + return 1; + } + + Step prefill(kPrefillSeq, dtype, vocab); + Step long_prefill(kMaxSeq, dtype, vocab); + Step decode(1, dtype, vocab); + + // Long prompt: pin the prefill profile, whose kernels TensorRT tuned at a + // 128-token sequence. + bool ok = run(module, "pinned prefill (seq=128)", kPrefillProfile, prefill); + + // One token at a time: pin the decode profile, whose seq is pinned to 1 so + // TensorRT could specialize it instead of serving it from prefill kernels. + for (int token = 0; ok && token < 3; ++token) { + ok = run(module, "pinned decode (seq=1)", kDecodeProfile, decode); + } + + // Back to prefill at the profile's upper bound, to show the switch is per + // call and not one-way. + ok = ok && run(module, "pinned prefill (seq=256)", kPrefillProfile, long_prefill); + + // Auto-selection reads the input shapes instead. It is sticky: once the + // prefill profile is loaded a seq == 1 input still fits it, so this stays on + // profile 1 rather than dropping back to decode. Pin when that matters. + ok = ok && run_auto(module, "auto (seq=1)", decode); + + // With no guard in scope every delegate runs profile 0, which here accepts + // seq == 1 only. + if (ok) { + auto result = module.forward(decode.args()); + if (!result.ok()) { + ET_LOG(Error, "no guard: forward() failed: 0x%" PRIx32, static_cast(result.error())); + ok = false; + } else { + print_prediction("no guard (seq=1)", result.get()); + } + } + + // A pinned index the engine does not have is an input error, reported before + // anything is enqueued. + if (ok) { + OptimizationProfileGuard profile_guard(99); + auto result = module.forward(decode.args()); + if (result.ok()) { + ET_LOG(Error, "expected profile 99 to be rejected"); + ok = false; + } else { + fprintf(stderr, "%-28s rejected as expected\n", "pinned profile 99"); + } + } + + // Correctness is settled by here; what remains is what the choice is worth. + ok = ok && report_latency(module, prefill, decode); + + if (!ok) { + return 1; + } + ET_LOG(Info, "Multi-profile run completed."); + return 0; +} diff --git a/examples/torchtrt_executorch_example/export_multi_profile.py b/examples/torchtrt_executorch_example/export_multi_profile.py new file mode 100644 index 0000000000..e5a1247b2f --- /dev/null +++ b/examples/torchtrt_executorch_example/export_multi_profile.py @@ -0,0 +1,291 @@ +""" +.. _executorch_export_multi_profile: + +Saving a Multi-Optimization-Profile Gemma-3 Model in ExecuTorch Format (.pte) +============================================================================= + +Autoregressive LLMs run in two very different shape *regimes* that share one set +of weights: + +- **prefill**: the prompt is processed in one shot, so the sequence length + ``seq`` is large, and +- **decode**: tokens are generated one at a time, so ``seq == 1``. + +A single dynamic range ``seq in [1, max]`` works, but TensorRT can only tune +kernels for **one** ``opt`` point. Tuning for the prefill length leaves decode -- +the latency-critical, most-frequently-executed phase -- running on kernels +picked for a sequence it never sees. + +``torch_tensorrt.Input(profiles=[...])`` declares **N optimization profiles** on +a single input. The engine is built **once** (a single ``torch.export`` over the +union of all profiles) and each profile gets its own TensorRT kernel tuning: + +- profile ``0`` -> **decode**: ``seq`` pinned to 1 (a fully static profile) +- profile ``1`` -> **prefill**: ``seq`` in ``[1, MAX_SEQ]``, tuned at ``PREFILL_SEQ`` + +Run the result with ``examples/executorch_reference_runner``, which selects a +profile per call with ``OptimizationProfileGuard``, and measure what that +selection is worth with ``example_executorch_multi_profile_benchmark``. + +By default this exports a **randomly initialized mini Gemma-3**: the real +architecture (sliding-window and full attention, the Gemma-3 SDPA lowering) at a +few million parameters, so the whole export takes about a minute and needs no +download or Hugging Face account. Only the shapes matter for demonstrating +optimization profiles, and the weights never leave the engine. + +Pass ``--weights google/gemma-3-1b-it`` for the real 1B model. That is the +configuration the latency numbers in the runner README were measured on, and it +takes considerably longer: the ``.pte`` serialization step costs roughly 3.6 +seconds per megabyte of engine, so its ~2 GB engine is a couple of hours. + +.. note:: + + ``google/gemma-3-1b-it`` is **gated**: accept its license on the Hugging Face + Hub and authenticate (``hf auth login`` or the ``HF_TOKEN`` environment + variable) first, or point ``--weights`` at an ungated mirror of the same + architecture. A CUDA GPU is required either way. + +Prerequisites +------------- +Install Torch-TensorRT with the ExecuTorch extra before running this example:: + + pip install -e ".[executorch]" + +See https://pytorch.org/executorch/stable/getting-started-setup.html for details. +""" + +# %% +# Imports and Setup +# ^^^^^^^^^^^^^^^^^^ +# +# ``export_llm``, a reusable helper that traces a decoder over a dynamic +# sequence length, lives under ``tools/llm`` in the Torch-TensorRT repo, so we +# add that directory to ``sys.path``. + +import argparse +import sys +import time +from pathlib import Path + +import torch +import torch_tensorrt + +_start = time.time() + + +def stamp(phase: str) -> None: + """Each phase's cost, since export time is the first thing people ask about.""" + print(f"[{time.time() - _start:6.1f}s] {phase}", flush=True) + + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT / "tools" / "llm")) + +MODEL_ID = "google/gemma-3-1b-it" +DEVICE = torch.device("cuda:0") + +# The two regimes, matching examples/dynamo/multi_optimization_profiles.py. +MAX_SEQ = 256 # largest prompt the engine must support +PREFILL_SEQ = 128 +DECODE_SEQ = 1 +DECODE_IDX, PREFILL_IDX = 0, 1 + +# Gemma-3 shrunk to a few million parameters: the real layer structure, only +# narrower and shallower. ``sliding_window`` keeps the 1B model's 512, which is +# wider than MAX_SEQ, so as in the real model the window never binds over the +# exported range and every layer attends to the whole prefix. Narrowing it below +# MAX_SEQ would make the sliding layers genuinely windowed, and the engine would +# then need ``attn_bias_is_causal=False`` to keep the mask instead of assuming +# plain causality. +MINI_CONFIG = dict( + vocab_size=2048, + hidden_size=320, + intermediate_size=640, + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=1, + head_dim=80, + max_position_embeddings=512, + sliding_window=512, + layer_types=["sliding_attention", "sliding_attention", "full_attention"], +) + +parser = argparse.ArgumentParser() +parser.add_argument( + "--model_path", + default="model_gemma3_multi_profile.pte", + help="Path to save the .pte file", +) +parser.add_argument( + "--weights", + default=None, + help=( + "Hugging Face repo to load pretrained weights from, e.g. " + f"{MODEL_ID}. Omit to export a randomly initialized mini Gemma-3, " + "which needs no download and exports in about a minute." + ), +) +args = parser.parse_args() + + +# %% +# The Exported Method +# ^^^^^^^^^^^^^^^^^^^^ +# +# The wrapper fixes the ``.pte``'s method signature to two ``[1, seq]`` inputs +# and one output, and returns only the **last** position's logits -- the row a +# sampler actually reads. That keeps the output shape static at ``[1, vocab]`` +# whatever ``seq`` is, so ExecuTorch plans one small buffer instead of one sized +# for ``MAX_SEQ``, and a large device-to-host copy does not end up dominating +# the very latency this example is meant to measure. +class NextTokenLogits(torch.nn.Module): + def __init__(self, model: torch.nn.Module) -> None: + super().__init__() + self.model = model + + def forward( + self, input_ids: torch.Tensor, position_ids: torch.Tensor + ) -> torch.Tensor: + out = self.model(input_ids=input_ids, position_ids=position_ids) + return out.logits[:, -1, :] + + +# %% +# Build the Model +# ^^^^^^^^^^^^^^^^ +# +# Either way the model runs in fp16 with ``use_cache=False`` (this example +# recomputes over the full sequence rather than using a KV cache, which keeps +# the export simple). ``attn_implementation="sdpa"`` makes HuggingFace emit +# ``scaled_dot_product_attention``, which Torch-TensorRT converts to a single +# TensorRT attention layer; no SDPA lowering pass is needed. +def build_model() -> torch.nn.Module: + from transformers import Gemma3ForCausalLM, Gemma3TextConfig + + with torch.no_grad(): + if args.weights: + from transformers import AutoModelForCausalLM + + model = AutoModelForCausalLM.from_pretrained( + args.weights, + use_cache=False, + attn_implementation="sdpa", + ignore_mismatched_sizes=True, + ) + else: + config = Gemma3TextConfig( + use_cache=False, attn_implementation="sdpa", **MINI_CONFIG + ) + model = Gemma3ForCausalLM(config) + model = model.eval().cuda().to(torch.float16) + + params = sum(p.numel() for p in model.parameters()) + stamp( + f"model built: Gemma-3 ({args.weights or 'mini, random init'}), {params / 1e6:.1f}M params" + ) + return model + + +try: + model = build_model() +except Exception as e: # no GPU, or gated/unauthenticated --weights + print(f"Skipping example: could not build the model ({type(e).__name__}: {e}).") + print("A CUDA GPU is required. With --weights, accept the model license and") + print("authenticate (hf auth login / HF_TOKEN), or use an ungated mirror.") + sys.exit(0) + + +# %% +# Declaring the Optimization Profiles +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +# +# ``profiles`` is an ordered list and the list index *is* the optimization +# profile index selected at runtime. There are no profile names. Both model +# inputs are dynamic over ``seq``, so each gets a profiled ``Input`` with +# identical profiles. +# +# The ranges overlap at ``seq == 1``: a decode-sized input is valid under both +# profiles. That overlap is why auto-selection is history-dependent (it keeps +# the loaded profile while it still fits) and why prefill/decode serving should +# pin a profile explicitly rather than rely on auto. +profiles = [ + {"min_shape": (1, 1), "opt_shape": (1, 1), "max_shape": (1, 1)}, # decode + { + "min_shape": (1, 1), + "opt_shape": (1, PREFILL_SEQ), + "max_shape": (1, MAX_SEQ), + }, # prefill +] +multi_profile_inputs = [ + torch_tensorrt.Input(dtype=torch.int64, profiles=profiles), # input_ids + torch_tensorrt.Input(dtype=torch.int64, profiles=profiles), # position_ids +] + +# %% +# Export Bounds Must Cover Every Profile +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +# +# ExecuTorch plans the ``.pte``'s memory from the ``torch.export`` input domain, +# not from the TensorRT profiles, and it plans for the upper bound. So the +# exported ``Dim`` has to span the *union* of all profile ranges -- here +# ``[1, MAX_SEQ]``, the prefill maximum -- or a profile accepting a larger input +# than the plan allows for would overrun its buffer. +from utils import export_llm # noqa: E402 + +vocab = model.config.get_text_config().vocab_size +example_ids = torch.randint( + 1, vocab, (1, PREFILL_SEQ), dtype=torch.int64, device=DEVICE +) +with torch.inference_mode(): + exported = export_llm( + NextTokenLogits(model), example_ids, min_seq_len=1, max_seq_len=MAX_SEQ + ) +stamp("torch.export done") + +# %% +# Compile Once +# ^^^^^^^^^^^^^ +# +# One export, one compile, one engine holding both profiles. Nothing about the +# profiles is chosen here beyond their bounds; which one runs is a runtime +# decision made per call by the C++ runner. +# +# ``offload_module_to_cpu`` must stay False: it is currently incompatible with +# the multi-profile ``Input(profiles=...)`` path (CPU/CUDA device mismatch). +print("Compiling multi-profile engine (decode + prefill) ...") +with torch.inference_mode(): + trt_gm = torch_tensorrt.dynamo.compile( + exported, + arg_inputs=multi_profile_inputs, + use_fp32_acc=True, + disable_tf32=True, + offload_module_to_cpu=False, + min_block_size=1, + require_full_compilation=True, + device=DEVICE, + ) +stamp("TensorRT engine built") + + +# %% +# Save as ExecuTorch .pte format +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +# All profiles live inside the serialized engine, so the TR01 blob format is +# unchanged and the runtime rediscovers count and bounds at load. This step +# scales with the size of the engine, and dominates the export for large models. +position_ids = torch.arange(PREFILL_SEQ, device=DEVICE).unsqueeze(0) +torch_tensorrt.save( + trt_gm, + args.model_path, + output_format="executorch", + arg_inputs=(example_ids, position_ids), + retrace=False, +) +stamp("saved .pte") + +size_mb = Path(args.model_path).stat().st_size / 1e6 +print(f"\nSaved {args.model_path} ({size_mb:.1f} MB) with {len(profiles)} profiles.") +print(f" profile {DECODE_IDX} (decode): seq == {DECODE_SEQ}") +print( + f" profile {PREFILL_IDX} (prefill): seq in [1, {MAX_SEQ}], tuned at {PREFILL_SEQ}" +) diff --git a/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py b/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py index 5f2473f587..2ec360d339 100644 --- a/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py +++ b/py/torch_tensorrt/dynamo/runtime/_TRTEngine.py @@ -829,8 +829,29 @@ def _set_active_profile_with_stream( the enqueue stream). Used by auto-selection, which switches on ``_engine_stream`` before ``execute_async_v3``. Mirrors the C++ runtime. """ - if self.num_optimization_profiles <= 1: + # An index this engine does not have cannot be honored. A single-profile + # engine is the one case where that is not a mistake to fail on: it has + # profile 0 and nothing to switch to, so a pin aimed at some other engine + # is ignored with a warning rather than raising. A multi-profile engine + # handed an index it lacks could have done something different, and a + # negative index is a computed value gone wrong on any engine, so both + # stay errors -- silently running on mistuned kernels is worse than + # stopping. Kept identical to the C++ runtime, which this mirrors. + # + # Reachable only by calling the engine directly; optimization_profile() + # validates the index first. + if self.num_optimization_profiles <= 1 and profile_index > 0: + logger.warning( + f"Ignoring optimization profile index {profile_index}: this engine has " + f"{self.num_optimization_profiles} optimization profile(s), so it stays on " + f"profile {self._active_profile_index}." + ) return + if not 0 <= profile_index < self.num_optimization_profiles: + raise ValueError( + f"Optimization profile index {profile_index} is out of range: this engine " + f"has {self.num_optimization_profiles} optimization profile(s)." + ) if profile_index == self._active_profile_index: return self.context.set_optimization_profile_async(profile_index, stream.cuda_stream) diff --git a/tests/cpp/executorch/BUILD b/tests/cpp/executorch/BUILD index 17d2820bf2..4cb35b5622 100644 --- a/tests/cpp/executorch/BUILD +++ b/tests/cpp/executorch/BUILD @@ -9,6 +9,7 @@ test_suite( ":test_executorch_binding_names", ":test_executorch_blob_header", ":test_executorch_weight_streaming_budget", + ":test_optimization_profile_selection", ], ) @@ -47,3 +48,20 @@ cc_test( "@googletest//:gtest_main", ], ) + +# Asserts against the profile-selection policy only, so unlike the runtime tests +# it needs neither a GPU nor a TensorRT engine. It does need TensorRT headers, +# hence the same platform gate the backend carries. +cc_test( + name = "test_optimization_profile_selection", + srcs = ["test_optimization_profile_selection.cpp"], + target_compatible_with = select({ + "//cpp:linux_x86_64": [], + "//cpp:sbsa": [], + "//conditions:default": ["@platforms//:incompatible"], + }), + deps = [ + "//cpp:tensorrt_executorch_optimization_profile_selection", + "@googletest//:gtest_main", + ], +) diff --git a/tests/cpp/executorch/test_optimization_profile_selection.cpp b/tests/cpp/executorch/test_optimization_profile_selection.cpp new file mode 100644 index 0000000000..c5e25189ef --- /dev/null +++ b/tests/cpp/executorch/test_optimization_profile_selection.cpp @@ -0,0 +1,302 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + * + * Self-check for the optimization-profile selection policy. The policy is the + * only non-obvious part of multi-profile support and depends on nothing but the + * profile bounds table, so it runs here without a GPU, a TensorRT engine, or an + * ExecuTorch method. + * + * Applying the decision is not covered here or anywhere else yet, because all of + * it needs a live engine: setOptimizationProfileAsync, writing profiles.active + * back, the ordering that puts the switch before setInputShape, and + * mark_inflight. The one automated live-engine job, + * .github/scripts/verify-executorch-reference-runner.sh, exports a + * single-profile static model and installs no guard, so it never switches. + * example_executorch_multi_profile_runner does exercise all of it and returns + * nonzero on failure, so pointing that job at a multi-profile .pte is the cheap + * way to close the gap. + */ + +#include "torch_tensorrt/executorch/OptimizationProfileSelection.h" + +#include "gtest/gtest.h" + +#include +#include +#include + +namespace torch_tensorrt { +namespace executorch_backend { +namespace { + +nvinfer1::Dims dims(std::initializer_list extents) { + nvinfer1::Dims out{}; + out.nbDims = static_cast(extents.size()); + int i = 0; + for (int64_t extent : extents) { + out.d[i++] = extent; + } + return out; +} + +InputProfileBounds bounds(std::initializer_list min, std::initializer_list max) { + return InputProfileBounds{dims(min), dims(max)}; +} + +// An LLM-shaped engine: one [1, seq] input, profile 0 decodes a single token and +// profile 1 covers prefill. The ranges overlap at seq == 1 on purpose, which is +// what makes the sticky rule observable. +ProfileTable decode_and_prefill() { + ProfileTable table; + table.bounds = { + {bounds({1, 1}, {1, 1})}, // profile 0: decode + {bounds({1, 1}, {1, 2048})}, // profile 1: prefill + }; + return table; +} + +std::vector decode_input() { + return {dims({1, 1})}; +} + +std::vector prefill_input() { + return {dims({1, 512})}; +} + +TEST(ExecuTorchOptimizationProfileSelection, UnsetRequestRunsProfileZeroWhateverTheShapesSay) { + ProfileTable table = decode_and_prefill(); + table.active = 1; + int32_t selected = -1; + + EXPECT_EQ(select_profile(table, ProfileRequest::kUnset, 0, prefill_input(), selected), ProfileSelection::kOk); + EXPECT_EQ(selected, 0); +} + +TEST(ExecuTorchOptimizationProfileSelection, PinnedRequestTakesTheIndexVerbatim) { + ProfileTable table = decode_and_prefill(); + int32_t selected = -1; + + EXPECT_EQ(select_profile(table, ProfileRequest::kPinned, 1, prefill_input(), selected), ProfileSelection::kOk); + EXPECT_EQ(selected, 1); + + table.active = 1; + EXPECT_EQ(select_profile(table, ProfileRequest::kPinned, 0, decode_input(), selected), ProfileSelection::kOk); + EXPECT_EQ(selected, 0); +} + +TEST(ExecuTorchOptimizationProfileSelection, AutoPicksTheOnlyProfileThatFits) { + ProfileTable table = decode_and_prefill(); + int32_t selected = -1; + + EXPECT_EQ(select_profile(table, ProfileRequest::kAuto, 0, prefill_input(), selected), ProfileSelection::kOk); + EXPECT_EQ(selected, 1); +} + +// Auto is sticky where the profiles overlap: a one-token input still fits the +// prefill profile, so a decode step after a prefill step stays on profile 1 +// rather than dropping back to the lowest matching index. Documented behavior, +// and the reason prefill/decode workloads should pin instead. +TEST(ExecuTorchOptimizationProfileSelection, AutoKeepsTheActiveProfileWhereRangesOverlap) { + ProfileTable table = decode_and_prefill(); + table.active = 1; + int32_t selected = -1; + + EXPECT_EQ(select_profile(table, ProfileRequest::kAuto, 0, decode_input(), selected), ProfileSelection::kOk); + EXPECT_EQ(selected, 1); +} + +TEST(ExecuTorchOptimizationProfileSelection, AutoRescansOnceTheActiveProfileStopsFitting) { + ProfileTable table = decode_and_prefill(); + table.active = 0; + int32_t selected = -1; + + EXPECT_EQ(select_profile(table, ProfileRequest::kAuto, 0, prefill_input(), selected), ProfileSelection::kOk); + EXPECT_EQ(selected, 1); +} + +// A profile has to fit *every* input, not just the first one it is asked about. +// Here profile 0 accepts the one-token input_ids but not the longer second input, +// so auto has to keep looking rather than stop at the first partial match. +TEST(ExecuTorchOptimizationProfileSelection, AutoSkipsAProfileThatFitsOnlySomeInputs) { + ProfileTable table; + table.bounds = { + {bounds({1, 1}, {1, 1}), bounds({1, 1}, {1, 1})}, // profile 0: second input too narrow + {bounds({1, 1}, {1, 1}), bounds({1, 1}, {1, 128})}, // profile 1: fits both + }; + const std::vector inputs{dims({1, 1}), dims({1, 64})}; + int32_t selected = -1; + + EXPECT_EQ(select_profile(table, ProfileRequest::kAuto, 0, inputs, selected), ProfileSelection::kOk); + EXPECT_EQ(selected, 1); +} + +// The rescan is a first-fit from 0, so when the active profile stops fitting and +// several others would serve, the lowest matching index wins. +TEST(ExecuTorchOptimizationProfileSelection, RescanTakesTheLowestOfSeveralMatchingProfiles) { + ProfileTable table; + table.bounds = { + {bounds({1, 1}, {1, 64})}, // profile 0: fits + {bounds({1, 1}, {1, 256})}, // profile 1: fits too + {bounds({1, 512}, {1, 2048})}, // profile 2: active, no longer fits + }; + table.active = 2; + const std::vector short_input{dims({1, 32})}; + int32_t selected = -1; + + EXPECT_EQ(select_profile(table, ProfileRequest::kAuto, 0, short_input, selected), ProfileSelection::kOk); + EXPECT_EQ(selected, 0); +} + +// A shape no profile covers is an input error, not a silent clamp. +TEST(ExecuTorchOptimizationProfileSelection, AutoRejectsShapeNoProfileCovers) { + ProfileTable table = decode_and_prefill(); + const std::vector too_long{dims({1, 4096})}; + int32_t selected = -1; + + EXPECT_EQ( + select_profile(table, ProfileRequest::kAuto, 0, too_long, selected), ProfileSelection::kNoProfileMatchesInputs); +} + +TEST(ExecuTorchOptimizationProfileSelection, AutoRejectsRankThatDoesNotMatchTheProfile) { + ProfileTable table = decode_and_prefill(); + const std::vector wrong_rank{dims({1, 1, 1})}; + int32_t selected = -1; + + EXPECT_EQ( + select_profile(table, ProfileRequest::kAuto, 0, wrong_rank, selected), ProfileSelection::kNoProfileMatchesInputs); +} + +TEST(ExecuTorchOptimizationProfileSelection, PinningPastTheEndOfAMultiProfileEngineIsRejected) { + ProfileTable table = decode_and_prefill(); + int32_t selected = -1; + + EXPECT_EQ( + select_profile(table, ProfileRequest::kPinned, 2, decode_input(), selected), + ProfileSelection::kRequestedProfileUnavailable); + EXPECT_EQ( + select_profile(table, ProfileRequest::kPinned, -3, decode_input(), selected), + ProfileSelection::kRequestedProfileUnavailable); +} + +// A .pte may mix a multi-profile engine with a single-profile one. Pinning a +// nonzero profile for the former must not fail the latter, which has profile 0 +// and nothing to switch to. Reported as kPinIgnoredSingleProfile rather than +// kOk so execute() can say the pin did nothing here. +TEST(ExecuTorchOptimizationProfileSelection, SingleProfileEngineToleratesAPinItCannotHonor) { + ProfileTable static_engine; + static_engine.bounds = {{bounds({1, 16}, {1, 16})}}; + const std::vector fixed_input{dims({1, 16})}; + int32_t selected = -1; + + EXPECT_EQ( + select_profile(static_engine, ProfileRequest::kPinned, 1, fixed_input, selected), + ProfileSelection::kPinIgnoredSingleProfile); + EXPECT_EQ(selected, 0); +} + +// Whether the single profile's shapes are dynamic makes no difference: profile 0 +// is still the only thing the engine can run, so it is tolerated the same way. +TEST(ExecuTorchOptimizationProfileSelection, SingleProfileToleranceDoesNotDependOnDynamicShapes) { + ProfileTable dynamic_engine; + dynamic_engine.bounds = {{bounds({1, 1}, {1, 2048})}}; + int32_t selected = -1; + + EXPECT_EQ( + select_profile(dynamic_engine, ProfileRequest::kPinned, 1, decode_input(), selected), + ProfileSelection::kPinIgnoredSingleProfile); + EXPECT_EQ(selected, 0); +} + +// The tolerance covers indices the engine merely lacks, not nonsense ones. A +// negative index is the shape a failed lookup returns, so it has to be reported +// even here, where profile 0 would otherwise be a tempting substitute. This is +// what the removed kAutoSelectProfile == -1 sentinel used to swallow. +TEST(ExecuTorchOptimizationProfileSelection, SingleProfileEngineStillRejectsANegativeIndex) { + ProfileTable table; + table.bounds = {{bounds({1, 16}, {1, 16})}}; + const std::vector fixed_input{dims({1, 16})}; + int32_t selected = -1; + + EXPECT_EQ( + select_profile(table, ProfileRequest::kPinned, -1, fixed_input, selected), + ProfileSelection::kRequestedProfileUnavailable); +} + +// The tolerance stops at one profile. With several, substituting profile 0 would +// be a guess about which regime the caller wanted, so this stays an error. +TEST(ExecuTorchOptimizationProfileSelection, MultiProfileEngineDoesNotSubstituteForAMissingIndex) { + ProfileTable table; + table.bounds = { + {bounds({1, 1}, {1, 1})}, + {bounds({1, 1}, {1, 128})}, + }; + int32_t selected = -1; + + EXPECT_EQ( + select_profile(table, ProfileRequest::kPinned, 2, decode_input(), selected), + ProfileSelection::kRequestedProfileUnavailable); +} + +// An engine with no inputs has nothing to constrain the choice, so every profile +// trivially fits and auto keeps the loaded one. +TEST(ExecuTorchOptimizationProfileSelection, AutoHandlesAnEngineWithNoInputs) { + ProfileTable table; + table.bounds = {{}, {}}; + table.active = 1; + int32_t selected = -1; + + EXPECT_EQ(select_profile(table, ProfileRequest::kAuto, 0, {}, selected), ProfileSelection::kOk); + EXPECT_EQ(selected, 1); +} + +// A table with no profiles at all is malformed; init() rejects such an engine +// before execute() ever runs. Guarded here anyway so the policy never indexes an +// empty bounds vector on the strength of a check in another translation unit. +TEST(ExecuTorchOptimizationProfileSelection, EmptyTableIsRejectedRatherThanIndexed) { + ProfileTable empty; + int32_t selected = -1; + + EXPECT_EQ( + select_profile(empty, ProfileRequest::kUnset, 0, decode_input(), selected), + ProfileSelection::kNoProfileMatchesInputs); + EXPECT_EQ( + select_profile(empty, ProfileRequest::kAuto, 0, decode_input(), selected), + ProfileSelection::kNoProfileMatchesInputs); + EXPECT_EQ( + select_profile(empty, ProfileRequest::kPinned, 0, decode_input(), selected), + ProfileSelection::kNoProfileMatchesInputs); +} + +// The other two ways a hand-built table can point outside itself. The backend +// cannot produce either, but the header is installed, so the policy treats both +// as "this profile does not fit" instead of reading past the end: an `active` +// naming a profile the table does not have, and a bounds row describing fewer +// inputs than the engine was handed. +TEST(ExecuTorchOptimizationProfileSelection, AutoRescansPastAnActiveIndexTheTableDoesNotHave) { + ProfileTable table = decode_and_prefill(); + table.active = 7; + int32_t selected = -1; + + EXPECT_EQ(select_profile(table, ProfileRequest::kAuto, 0, prefill_input(), selected), ProfileSelection::kOk); + EXPECT_EQ(selected, 1); +} + +TEST(ExecuTorchOptimizationProfileSelection, ProfileWithFewerBoundsThanInputsDoesNotFit) { + ProfileTable table; + table.bounds = { + {bounds({1, 1}, {1, 128})}, // profile 0: describes only the first input + {bounds({1, 1}, {1, 128}), bounds({1, 1}, {1, 128})}, // profile 1: complete + }; + const std::vector two_inputs{dims({1, 8}), dims({1, 8})}; + int32_t selected = -1; + + EXPECT_EQ(select_profile(table, ProfileRequest::kAuto, 0, two_inputs, selected), ProfileSelection::kOk); + EXPECT_EQ(selected, 1); +} + +} // namespace +} // namespace executorch_backend +} // namespace torch_tensorrt diff --git a/tests/py/dynamo/runtime/test_multi_optimization_profiles.py b/tests/py/dynamo/runtime/test_multi_optimization_profiles.py index 6dab7cd987..75f4c4c3d4 100644 --- a/tests/py/dynamo/runtime/test_multi_optimization_profiles.py +++ b/tests/py/dynamo/runtime/test_multi_optimization_profiles.py @@ -240,6 +240,19 @@ def test_out_of_range_index_raises(self): with optimization_profile(self.trt_gm, 99): self.trt_gm(self.decode_in) + def test_engine_level_out_of_range_index_raises(self): + # Driving the engine directly skips optimization_profile()'s validation, + # so the engine has to reject an index it cannot honor rather than warn + # and keep running on whatever profile is loaded. These engines have two + # profiles, so there is no single-profile tolerance to fall back on. + # Only the exception type differs by runtime: ValueError from the Python + # engine, RuntimeError from the C++ TORCHTRT_CHECK. + for e in self._trt_engines(self.trt_gm): + with self.assertRaises((ValueError, RuntimeError)): + e.set_active_profile(e.num_optimization_profiles) + with self.assertRaises((ValueError, RuntimeError)): + e.set_active_profile(-1) + def test_non_int_profile_raises(self): with self.assertRaises(TypeError): with optimization_profile(self.trt_gm, "decode"): From 789aa3ba20d723df64f8bc088bfae9c7c0287353 Mon Sep 17 00:00:00 2001 From: cehongwang Date: Wed, 26 Aug 2026 21:41:50 +0000 Subject: [PATCH 2/2] Revised again --- .../executorch/TensorRTBackend.h | 25 +++++- .../torch_tensorrt/executorch/CMakeLists.txt | 5 +- .../torch_tensorrt/executorch/EngineHandle.h | 4 + .../executorch/TensorRTBackend.cpp | 80 +++++++++---------- .../dynamo/multi_optimization_profiles.py | 12 ++- .../multi_profile_benchmark.cpp | 8 +- .../test_multi_optimization_profiles.py | 62 +++++++++++++- 7 files changed, 140 insertions(+), 56 deletions(-) diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 719dbb57e8..0d7d7d8f27 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -137,8 +137,25 @@ class TensorRTBackend final : public ::executorch::runtime::BackendInterface { * profile is the benign case -- it runs profile 0 and logs that the pin did * nothing -- while a multi-profile engine that lacks the index fails the * execution. + * + * @warning Name the guard. A discarded temporary is destroyed at the end of the + * full-expression that made it, restoring the enclosing request before + * forward() is ever called, so the execution runs profile 0 rather than the one + * asked for: + * + * @code + * OptimizationProfileGuard(kPrefillProfile); // no-op, guard already dead + * OptimizationProfileGuard::automatic(); // no-op, guard already dead + * OptimizationProfileGuard guard(kPrefillProfile); // correct + * auto guard = OptimizationProfileGuard::automatic(); // correct + * @endcode + * + * Both mistakes are compiler warnings rather than a silently mistuned + * execution: the pinning constructor and automatic() are each [[nodiscard]] + * individually, because GCC applies a class-level [[nodiscard]] only to + * returned values, not to a discarded constructor temporary. */ -class OptimizationProfileGuard { +class [[nodiscard]] OptimizationProfileGuard { public: /** * @brief Pin an exact profile by its export-time index. @@ -150,7 +167,7 @@ class OptimizationProfileGuard { * * @param profile_index Position in the export-time profile list. */ - explicit OptimizationProfileGuard(int32_t profile_index); + [[nodiscard]] explicit OptimizationProfileGuard(int32_t profile_index); /// @brief Rejected so that OptimizationProfileGuard(true) cannot become index 1. OptimizationProfileGuard(bool) = delete; @@ -164,8 +181,10 @@ class OptimizationProfileGuard { * @code * auto profile_guard = OptimizationProfileGuard::automatic(); * @endcode + * + * @return A guard that must be bound to a name; discarding it selects nothing. */ - static OptimizationProfileGuard automatic(); + [[nodiscard]] static OptimizationProfileGuard automatic(); ~OptimizationProfileGuard(); OptimizationProfileGuard(const OptimizationProfileGuard&) = delete; diff --git a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt index b5eb617140..7e9208c15c 100644 --- a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt +++ b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt @@ -174,8 +174,9 @@ add_library(torchtrt::executorch_backend ALIAS torchtrt_executorch_backend) add_dependencies(torchtrt_executorch_backend executorch_trt_backend) # The archive is linked by file below rather than by target, so the include path -# does not come along with it. Carry it here: a runner that scopes CudaStreamGuard -# or OptimizationProfileGuard needs the public header. +# does not come along with it. Carry it here: a runner that scopes +# OptimizationProfileGuard needs the public header. (The caller-stream guard now +# comes from ExecuTorch, not from this header.) target_include_directories(torchtrt_executorch_backend INTERFACE "${CMAKE_CURRENT_LIST_DIR}/../../../include" diff --git a/cpp/src/torch_tensorrt/executorch/EngineHandle.h b/cpp/src/torch_tensorrt/executorch/EngineHandle.h index 366b7f73b4..e7ce02ee75 100644 --- a/cpp/src/torch_tensorrt/executorch/EngineHandle.h +++ b/cpp/src/torch_tensorrt/executorch/EngineHandle.h @@ -65,6 +65,10 @@ struct EngineHandle { // suffices because a handle runs on a single thread at a time. cudaEvent_t inflight_event = nullptr; bool inflight_pending = false; + // A pin this engine cannot honor is a property of the caller's guard, not of the + // call, so it would otherwise be reported identically on every execute(). One + // engine, one report: a decode loop must not turn it into a log flood. + bool pin_ignored_reported = false; ~EngineHandle(); }; diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 8442303492..5d6b690884 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -813,11 +813,14 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* case ProfileSelection::kOk: break; case ProfileSelection::kPinIgnoredSingleProfile: - ET_LOG( - Info, - "TensorRTBackend::execute: ignoring the pin on profile %d; this engine has one profile, " - "so it runs profile 0", - requested_index); + if (!engine->pin_ignored_reported) { + engine->pin_ignored_reported = true; + ET_LOG( + Info, + "TensorRTBackend::execute: ignoring the pin on profile %d; this engine has one profile, " + "so it runs profile 0 (reported once per engine)", + requested_index); + } break; case ProfileSelection::kRequestedProfileUnavailable: ET_LOG( @@ -849,7 +852,8 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* if (mark_err != Error::Ok) { return mark_err; } - ET_LOG(Info, "TensorRTBackend::execute: switched to optimization profile %d", profile); + // Debug, not Info: a decode/prefill loop switches on nearly every call. + ET_LOG(Debug, "TensorRTBackend::execute: switched to optimization profile %d", profile); } // ------------------------------------------------------------------ @@ -1134,43 +1138,33 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* const bool aliased_reflect_pending = !aliased_reflects.empty(); const bool must_sync = output_staged_to_host || input_staged_from_host || aliased_reflect_pending || !caller_stream_set; - if (must_sync) { - Error copy_err = Error::Ok; - for (auto& output : outputs_needing_copy) { - exec_aten::Tensor et_out = args[output.first]->toTensor(); - cuda_err = - cudaMemcpyAsync(et_out.mutable_data_ptr(), output.second, et_out.nbytes(), cudaMemcpyDeviceToHost, stream); - if (cuda_err != cudaSuccess) { - ET_LOG( - Error, - "TensorRTBackend::execute: D2H copy failed for output %zu: %s", - output.first, - cudaGetErrorString(cuda_err)); - // The enqueue already succeeded, so the engine is still running on the - // stream. Drain below before returning, or the next call mutates a live - // execution context, which TensorRT forbids. - copy_err = Error::InvalidProgram; - break; - } - } - cuda_err = cudaStreamSynchronize(stream); - engine->inflight_pending = false; - if (cuda_err != cudaSuccess) { - ET_LOG(Error, "TensorRTBackend::execute: cudaStreamSynchronize failed: %s", cudaGetErrorString(cuda_err)); - return Error::InvalidProgram; - } - if (copy_err != Error::Ok) { - return copy_err; - } - } else { - cuda_err = cudaEventRecord(engine->inflight_event, stream); + if (!must_sync) { + // Return with the enqueue still running, which is the contract documented on + // CallerStreamGuard use in TensorRTBackend.h: the caller synchronizes its own + // stream before reading device-resident outputs. Marking is what makes that + // safe -- it is the only thing telling the next execute() and ~EngineHandle to + // wait on the completion event before they touch exec_ctx -- so this path must + // neither sync (that would quietly make every call blocking) nor skip the mark + // (that would let the next call reconfigure a live context). + return mark_inflight(*engine, stream); + } + + Error copy_err = Error::Ok; + for (auto& output : outputs_needing_copy) { + exec_aten::Tensor et_out = args[output.first]->toTensor(); + cuda_err = + cudaMemcpyAsync(et_out.mutable_data_ptr(), output.second, et_out.nbytes(), cudaMemcpyDeviceToHost, stream); if (cuda_err != cudaSuccess) { - // Could not arm the completion marker; drain now so a later execute() or the - // destructor never reconfigures or frees exec_ctx while this enqueue runs. - ET_LOG(Error, "TensorRTBackend::execute: cudaEventRecord failed: %s", cudaGetErrorString(cuda_err)); - (void)cudaStreamSynchronize(stream); - engine->inflight_pending = false; - return Error::InvalidProgram; + ET_LOG( + Error, + "TensorRTBackend::execute: D2H copy failed for output %zu: %s", + output.first, + cudaGetErrorString(cuda_err)); + // The enqueue already succeeded, so the engine is still running on the + // stream. Drain below before returning, or the next call mutates a live + // execution context, which TensorRT forbids. + copy_err = Error::InvalidProgram; + break; } } cuda_err = cudaStreamSynchronize(stream); @@ -1179,7 +1173,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* ET_LOG(Error, "TensorRTBackend::execute: cudaStreamSynchronize failed: %s", cudaGetErrorString(cuda_err)); return Error::InvalidProgram; } - return Error::Ok; + return copy_err; } // --------------------------------------------------------------------------- diff --git a/examples/dynamo/multi_optimization_profiles.py b/examples/dynamo/multi_optimization_profiles.py index c10006b531..2e2dec334c 100644 --- a/examples/dynamo/multi_optimization_profiles.py +++ b/examples/dynamo/multi_optimization_profiles.py @@ -123,9 +123,15 @@ def make_inputs(seq_len: int): # - index ``0`` -> **decode**: ``seq`` pinned to 1 (a fully static profile) # - index ``1`` -> **prefill**: ``seq`` in ``[1, MAX_SEQ]``, tuned at ``PREFILL_SEQ`` # -# Profile order matters for auto-selection: the profiles overlap at ``seq == 1`` -# and auto-selection picks the *first* profile whose ``[min, max]`` accepts the -# input, so declaring ``decode`` first lets it win the ``seq == 1`` overlap. +# Profile order matters for auto-selection, but only on a rescan. Auto-selection +# is sticky: it keeps the active profile whenever that profile still accepts the +# input (switching invalidates the captured CUDA graph and forces shape +# re-inference), and only when it does not does it rescan and take the *lowest* +# index that fits. So declaring ``decode`` first wins the ``seq == 1`` overlap on +# a rescan, but it does not claw ``seq == 1`` back from ``prefill``: prefill's +# range is ``[1, MAX_SEQ]``, which still fits, so an alternating +# prefill/decode loop stays on prefill. Pin the profile (below) to make decode +# run under its own. profiles = [ {"min_shape": (1, 1), "opt_shape": (1, 1), "max_shape": (1, 1)}, # decode { diff --git a/examples/executorch_reference_runner/multi_profile_benchmark.cpp b/examples/executorch_reference_runner/multi_profile_benchmark.cpp index aa37660098..d89e7c4be1 100644 --- a/examples/executorch_reference_runner/multi_profile_benchmark.cpp +++ b/examples/executorch_reference_runner/multi_profile_benchmark.cpp @@ -335,11 +335,15 @@ int main(int argc, char** argv) { const Stats sw_prefill = summarize(switching.prefill_ms); const Stats sw_decode = summarize(switching.decode_ms); + // --prefill_seq is a flag, so the label has to carry the value that actually ran. + char prefill_label[32]; + snprintf(prefill_label, sizeof(prefill_label), "prefill (seq=%d)", prefill_seq); + printf("prefill-only (every call on the prefill profile)\n"); - print_stats("prefill (seq=128)", po_prefill); + print_stats(prefill_label, po_prefill); print_stats("decode (seq=1)", po_decode); printf("\nswitching (each phase on its own profile)\n"); - print_stats("prefill (seq=128)", sw_prefill); + print_stats(prefill_label, sw_prefill); print_stats("decode (seq=1)", sw_decode); printf("\nwhat switching bought (positive = switching is faster)\n"); diff --git a/tests/py/dynamo/runtime/test_multi_optimization_profiles.py b/tests/py/dynamo/runtime/test_multi_optimization_profiles.py index 75f4c4c3d4..fba0d2b666 100644 --- a/tests/py/dynamo/runtime/test_multi_optimization_profiles.py +++ b/tests/py/dynamo/runtime/test_multi_optimization_profiles.py @@ -10,9 +10,11 @@ from torch_tensorrt.runtime import optimization_profile # Profiles are an ordered list; the list index is the optimization-profile -# index selected at runtime. Order is meaningful for lazy auto-selection: the -# decode profile ([1, 1]) and prefill profile ([1, 64]) overlap at seq=1, so we -# declare decode FIRST (index 0) to make it win the overlap (first-working). +# index selected at runtime. Order matters on a rescan: the decode profile +# ([1, 1]) and prefill profile ([1, 64]) overlap at seq=1, and a rescan takes the +# lowest fitting index, so declaring decode FIRST (index 0) lets it win that +# overlap. Auto-selection is sticky, though, so it only rescans once the active +# profile stops fitting -- prefill accepts seq=1 and would keep it. DECODE_IDX = 0 PREFILL_IDX = 1 @@ -46,6 +48,19 @@ def _make_profiles_input(): ) +def _compile_mlp_single_profile(model, **kwargs): + """Compile with a static input, so the engine has exactly one profile.""" + example = torch.randn(4, 1, 16, dtype=torch.float16, device="cuda") + ep = torch.export.export(model, (example,)) + return torch_tensorrt.dynamo.compile( + ep, + arg_inputs=[example], + min_block_size=1, + enabled_precisions={torch.float16}, + **kwargs, + ) + + def _compile_mlp(model, **kwargs): inp = _make_profiles_input() example = torch.randn(4, 48, 16, dtype=torch.float16, device="cuda") @@ -253,6 +268,47 @@ def test_engine_level_out_of_range_index_raises(self): with self.assertRaises((ValueError, RuntimeError)): e.set_active_profile(-1) + def test_single_profile_engine_tolerates_a_nonzero_pin(self): + # A pin reaches every engine in the module, so an index meant for a + # multi-profile sibling lands on single-profile engines too. Those have + # profile 0 and nothing to switch to: they must warn and stay on 0 rather + # than fail an execution that was never about them (the C++ and Python + # runtimes and the ExecuTorch backend all make this same exception). + # Only reachable by driving the engine directly: set_optimization_profile + # range-checks the index against this engine's profile count first, which + # the tail of this test pins down. + trt_gm = _compile_mlp_single_profile(self.model) + static_in = torch.randn(4, 1, 16, dtype=torch.float16, device="cuda") + engines = self._trt_engines(trt_gm) + self.assertGreaterEqual(len(engines), 1) + + for e in engines: + self.assertEqual(e.num_optimization_profiles, 1) + if isinstance(e, TRTEngine): + # Python runtime: the warning is a Python log record. The C++ + # runtime logs through LOG_WARNING, which assertLogs cannot see. + with self.assertLogs("torch_tensorrt", level="WARNING") as logs: + e.set_active_profile(1) + self.assertTrue( + any("Ignoring optimization profile" in m for m in logs.output) + ) + else: + e.set_active_profile(1) + self.assertEqual(e._active_profile_index, 0) + + # Tolerating the pin must leave the engine runnable, on profile 0. + ref = self.model(static_in) + out = trt_gm(static_in) + self.assertTrue(torch.allclose(out, ref, atol=1e-2)) + for e in self._trt_engines(trt_gm): + self.assertEqual(e._active_profile_index, 0) + + # The tolerance is engine-level only: the public entry point still + # range-checks, so a user pin of 1 here is a mistake, not a no-op. + with self.assertRaises(ValueError): + with optimization_profile(trt_gm, 1): + trt_gm(static_in) + def test_non_int_profile_raises(self): with self.assertRaises(TypeError): with optimization_profile(self.trt_gm, "decode"):