diff --git a/cpp/BUILD b/cpp/BUILD index 23509b9164..e64723af79 100644 --- a/cpp/BUILD +++ b/cpp/BUILD @@ -107,14 +107,29 @@ cc_library( ], ) +cc_library( + name = "tensorrt_executorch_optimization_profile_selection", + hdrs = [ + "src/torch_tensorrt/executorch/OptimizationProfileSelection.h", + ], + strip_include_prefix = "src/torch_tensorrt/executorch", + deps = select({ + ":linux_x86_64": ["@tensorrt//:nvinfer"], + ":sbsa": ["@tensorrt_sbsa//:nvinfer"], + "//conditions:default": [], + }), +) + cc_library( name = "tensorrt_executorch_backend", srcs = [ + "src/torch_tensorrt/executorch/EngineHandle.h", "src/torch_tensorrt/executorch/TensorRTBackend.cpp", ], hdrs = [ "include/torch_tensorrt/executorch/TensorRTBackend.h", ], + strip_include_prefix = "include", # Build the TensorRT backend as a static library. The final application # links this target together with the ExecuTorch runtime it was compiled # against, avoiding any runtime plugin/dlopen dependency. @@ -123,19 +138,19 @@ cc_library( ":sbsa": [], "//conditions:default": ["@platforms//:incompatible"], }), - strip_include_prefix = "include", deps = [ ":tensorrt_executorch_binding_names", ":tensorrt_executorch_blob_header", + ":tensorrt_executorch_optimization_profile_selection", ] + select({ ":linux_x86_64": [ - "@executorch//:executorch_headers", "@cuda//:cudart", + "@executorch//:executorch_headers", "@tensorrt//:nvinfer", ], ":sbsa": [ - "@executorch//:executorch_headers", "@cuda//:cudart", + "@executorch//:executorch_headers", "@tensorrt_sbsa//:nvinfer", ], "//conditions:default": [], @@ -147,6 +162,8 @@ filegroup( name = "executorch_backend_source_files", srcs = [ "src/torch_tensorrt/executorch/CMakeLists.txt", + "src/torch_tensorrt/executorch/EngineHandle.h", + "src/torch_tensorrt/executorch/OptimizationProfileSelection.h", "src/torch_tensorrt/executorch/README.md", "src/torch_tensorrt/executorch/TensorRTBackend.cpp", "src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp", diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 31383f9e17..52d4587e3f 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -17,10 +17,8 @@ #include +#include #include -#include -#include -#include namespace torch_tensorrt { namespace executorch_backend { @@ -40,39 +38,6 @@ class TRTLogger : public nvinfer1::ILogger { 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; - 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(); -}; - class TensorRTBackend final : public ::executorch::runtime::BackendInterface { public: bool is_available() const override; @@ -119,5 +84,55 @@ class CudaStreamGuard { bool prev_set_; }; +// Pass instead of an index to have each delegate pick a profile from the runtime +// input shapes rather than being told one. +inline constexpr int32_t kAutoSelectProfile = -1; + +// 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: +// +// 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); +// } +// +// 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 CudaStreamGuard, 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. +class OptimizationProfileGuard { + public: + // profile_index: an exact profile to pin, or kAutoSelectProfile. + explicit OptimizationProfileGuard(int32_t profile_index); + ~OptimizationProfileGuard(); + OptimizationProfileGuard(const OptimizationProfileGuard&) = delete; + OptimizationProfileGuard& operator=(const OptimizationProfileGuard&) = delete; + + private: + int32_t prev_index_; + bool prev_set_; +}; + } // namespace executorch_backend } // namespace torch_tensorrt diff --git a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt index 1b0546b545..dcb9c8c0c4 100644 --- a/cpp/src/torch_tensorrt/executorch/CMakeLists.txt +++ b/cpp/src/torch_tensorrt/executorch/CMakeLists.txt @@ -97,6 +97,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..cc8e899800 --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/EngineHandle.h @@ -0,0 +1,59 @@ +/* + * 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 "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; + 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/OptimizationProfileSelection.h b/cpp/src/torch_tensorrt/executorch/OptimizationProfileSelection.h new file mode 100644 index 0000000000..eba7e59c06 --- /dev/null +++ b/cpp/src/torch_tensorrt/executorch/OptimizationProfileSelection.h @@ -0,0 +1,135 @@ +/* + * 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. + * + * 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 { + +// The [min, max] dim envelope one optimization profile allows for one input. +struct InputProfileBounds { + nvinfer1::Dims min{}; + nvinfer1::Dims max{}; +}; + +// Everything a profile decision depends on, read from the engine once at init(). +struct ProfileTable { + // 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; + // True when every input dim is pinned to one extent in every profile. Such an + // engine accepts one shape only, so a request for a profile it does not have + // changes nothing it can do. + bool all_inputs_static = true; + // The profile currently loaded into the execution context. + int32_t active = 0; + + int32_t size() const { + return static_cast(bounds.size()); + } +}; + +// 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 +}; + +enum class ProfileSelection { + // Created this enum to decouple the profile header from executorch so that we can test it seperately + kOk, + // A pinned index this engine does not have, and it is not static enough for + // that to be harmless. + kRequestedProfileUnavailable, + // Auto-selection ran out of profiles. + kNoProfileMatchesInputs, +}; + +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; +} + +inline bool profile_fits(const ProfileTable& table, int32_t profile, const std::vector& input_dims) { + const auto& bounds = table.bounds[static_cast(profile)]; + for (size_t i = 0; i < input_dims.size(); ++i) { + if (!dims_fit(input_dims[i], bounds[i])) { + return false; + } + } + return true; +} + +// Resolves one thread's profile request against one engine. `index` is read +// only for ProfileRequest::kPinned. +inline ProfileSelection select_profile( + const ProfileTable& table, + ProfileRequest request, + int32_t index, + const std::vector& input_dims, + int32_t& selected) { + 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; + } + + // An engine that accepts exactly one shape has nothing to switch, so a request + // aimed at its multi-profile siblings in the same method is satisfied by + // profile 0 rather than failing the whole execution. A dynamic engine that + // lacks the index is a real mismatch and is reported. + if (index > 0 && table.size() == 1 && table.all_inputs_static) { + selected = 0; + return ProfileSelection::kOk; + } + + return ProfileSelection::kRequestedProfileUnavailable; +} + +} // 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 b2e3b08232..2fb5143660 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" @@ -51,6 +52,11 @@ using ::executorch::runtime::Span; namespace { thread_local cudaStream_t g_user_stream = nullptr; thread_local bool g_user_stream_set = false; +// The profile request in effect for this thread: an exact index to pin, +// kAutoSelectProfile, or unset (profile 0). Read by execute(), never by the +// guard itself. +thread_local int32_t g_profile_request = 0; +thread_local bool g_profile_request_set = false; } // namespace CudaStreamGuard::CudaStreamGuard(cudaStream_t stream) : prev_stream_(g_user_stream), prev_set_(g_user_stream_set) { @@ -63,6 +69,17 @@ CudaStreamGuard::~CudaStreamGuard() { g_user_stream_set = prev_set_; } +OptimizationProfileGuard::OptimizationProfileGuard(int32_t profile_index) + : prev_index_(g_profile_request), prev_set_(g_profile_request_set) { + g_profile_request = profile_index; + g_profile_request_set = true; +} + +OptimizationProfileGuard::~OptimizationProfileGuard() { + g_profile_request = prev_index_; + g_profile_request_set = prev_set_; +} + void TRTLogger::log(Severity severity, const char* msg) noexcept { if (severity <= Severity::kERROR) { ET_LOG(Error, "TensorRT: %s", msg); @@ -73,12 +90,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); @@ -175,21 +192,77 @@ 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; + } + for (int d = 0; d < bounds.min.nbDims; ++d) { + if (bounds.min.d[d] != bounds.max.d[d]) { + handle.profiles.all_inputs_static = false; + } + } + bounds_for_profile.push_back(bounds); } - handle.input_profile_bounds.push_back(bounds); } return Error::Ok; } +// Turns this thread's guard state into the request the policy understands. +ProfileRequest current_profile_request() { + if (!g_profile_request_set) { + return ProfileRequest::kUnset; + } + return g_profile_request == kAutoSelectProfile ? ProfileRequest::kAuto : ProfileRequest::kPinned; +} + +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; @@ -203,6 +276,20 @@ 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. +void mark_inflight(EngineHandle& engine, cudaStream_t stream) { + const cudaError_t err = cudaEventRecord(engine.inflight_event, stream); + if (err != cudaSuccess) { + ET_LOG(Error, "TensorRTBackend::execute: cudaEventRecord failed: %s", cudaGetErrorString(err)); + (void)cudaStreamSynchronize(stream); + } + engine.inflight_pending = (err == cudaSuccess); +} + } // namespace // --------------------------------------------------------------------------- @@ -257,8 +344,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)); @@ -393,8 +480,13 @@ 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. // ------------------------------------------------------------------ + std::vector input_dims(num_inputs); for (size_t i = 0; i < num_inputs; ++i) { EValue* arg = args[i]; TORCHTRT_ET_CHECK_NOT_NULL(arg, Error::InvalidArgument, "TensorRTBackend::execute: input %zu is not a tensor", i); @@ -402,31 +494,61 @@ 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(); - const std::string& name = engine->input_binding_names[i]; - 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) { + int32_t profile = 0; + switch (select_profile(engine->profiles, current_profile_request(), g_profile_request, input_dims, profile)) { + case ProfileSelection::kOk: + 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)", + g_profile_request, + engine->profiles.size()); return Error::InvalidArgument; + case ProfileSelection::kNoProfileMatchesInputs: + 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", + engine->profiles.size()); + 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, so mark them now instead of relying on the + // enqueue at the tail to do it. + mark_inflight(*engine, stream); + engine->profiles.active = profile; + 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()); @@ -479,7 +601,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(); @@ -492,7 +614,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 @@ -566,7 +688,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( @@ -576,15 +698,17 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* "cudaStreamPerThread is invalid while a green context is current."); return Error::InvalidState; } + mark_inflight(*engine, stream); - // The engine work is now in flight on `stream`. Decide whether to wait for it: + // The engine work is now in flight on `stream` and marked as such. Decide whether + // to wait for it here: // must_sync = an output is staged to host (the caller reads the D2H result on // return), an input was staged from host (its async H2D read the caller's host // 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 + // Otherwise (caller stream + all I/O device-resident) leave the work enqueued so it + // composes with the caller's later GPU work; the marker already tells the next + // execute() and the destructor to wait before reusing/freeing exec_ctx. 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. const bool must_sync = output_staged_to_host || input_staged_from_host || !g_user_stream_set; @@ -599,6 +723,9 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* "TensorRTBackend::execute: D2H copy failed for output %zu: %s", output.first, cudaGetErrorString(cuda_err)); + // Earlier iterations are still copying into the caller's output tensors. + (void)cudaStreamSynchronize(stream); + engine->inflight_pending = false; return Error::InvalidProgram; } } @@ -608,17 +735,6 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* ET_LOG(Error, "TensorRTBackend::execute: cudaStreamSynchronize failed: %s", cudaGetErrorString(cuda_err)); return Error::InvalidProgram; } - } else { - cuda_err = cudaEventRecord(engine->inflight_event, 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; - } - engine->inflight_pending = true; } 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/CMakeLists.txt b/examples/executorch_reference_runner/CMakeLists.txt index 2bd3544d67..7b6cdefadb 100644 --- a/examples/executorch_reference_runner/CMakeLists.txt +++ b/examples/executorch_reference_runner/CMakeLists.txt @@ -61,3 +61,25 @@ target_link_libraries( executorch::extensions executorch::kernels torchtrt::executorch_backend) + +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 a518353bb6..d3d6ec638c 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,13 @@ 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 1.9 GB and serialization runs at +# roughly 3.7 s/MB, so budget hours rather than minutes for it. +python examples/torchtrt_executorch_example/export_multi_profile.py --model_path=model_gemma3_multi_profile.pte ``` ## Build The Reference Runner @@ -103,3 +116,105 @@ Loading the method initializes the TensorRT ExecuTorch backend for any Torch-TensorRT delegate subgraphs embedded in the `.pte`. The Python `torch_tensorrt` package is needed when exporting the `.pte`; it is not needed by this native runner at inference time. + +## 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. Only +`kAutoSelectProfile` is a library constant; name the indices 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()`. Pass `kAutoSelectProfile` +instead of an index to choose from the input shapes; with no guard in scope, +every delegate runs profile 0. + +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..4f170e535f --- /dev/null +++ b/examples/executorch_reference_runner/multi_profile_benchmark.cpp @@ -0,0 +1,322 @@ +/* + * 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; +} + +int get_int_flag(int argc, char** argv, const char* flag, int def) { + const char* raw = get_flag(argc, argv, flag, nullptr); + return raw == nullptr ? def : atoi(raw); +} + +// 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. +class IndexTensor { + public: + IndexTensor(int32_t seq, exec_aten::ScalarType dtype, bool positions) + : 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) % 9000); + 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) + : ids_(seq, dtype, false), positions_(seq, dtype, true), 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"); + const int prefill_seq = get_int_flag(argc, argv, "--prefill_seq", 128); + const int blocks = get_int_flag(argc, argv, "--blocks", 10); + const int block_rounds = get_int_flag(argc, argv, "--block_rounds", 3); + const int decode_steps = get_int_flag(argc, argv, "--decode_steps", 16); + const int warmup = get_int_flag(argc, argv, "--warmup", 20); + + 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(); + + Step prefill(prefill_seq, dtype); + Step decode(1, dtype); + + 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, + 100.0 * (switching.wall_ms - prefill_only.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..659ea4e72c --- /dev/null +++ b/examples/executorch_reference_runner/multi_profile_main.cpp @@ -0,0 +1,297 @@ +/* + * 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::kAutoSelectProfile; +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. +class IndexTensor { + public: + IndexTensor(int32_t seq, exec_aten::ScalarType dtype, bool positions) + : 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) % 9000); + 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) + : ids_(seq, dtype, false), positions_(seq, dtype, true), 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(); + double best = -1e30; + int64_t best_idx = -1; + for (int64_t i = 0; i < t.numel(); ++i) { + const double v = t.scalar_type() == exec_aten::ScalarType::Half + ? static_cast(t.const_data_ptr()[i]) + : 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 under `profile`, which is either an exact index to pin or +// kAutoSelectProfile. +bool run(Module& module, const char* label, int32_t profile, const Step& step) { + // 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. + OptimizationProfileGuard profile_guard(profile); + 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; +} + +// 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(); + + Step prefill(kPrefillSeq, dtype); + Step long_prefill(kMaxSeq, dtype); + Step decode(1, dtype); + + // 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(module, "auto (seq=1)", kAutoSelectProfile, 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..e3251a634e --- /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 a ~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/tests/cpp/executorch/BUILD b/tests/cpp/executorch/BUILD index aee13cbcfd..50fc55269d 100644 --- a/tests/cpp/executorch/BUILD +++ b/tests/cpp/executorch/BUILD @@ -7,6 +7,7 @@ test_suite( tests = [ ":test_executorch_binding_names", ":test_executorch_blob_header", + ":test_optimization_profile_selection", ], ) @@ -27,3 +28,14 @@ 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. +cc_test( + name = "test_optimization_profile_selection", + srcs = ["test_optimization_profile_selection.cpp"], + 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..0828ccbee8 --- /dev/null +++ b/tests/cpp/executorch/test_optimization_profile_selection.cpp @@ -0,0 +1,206 @@ +/* + * 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. + */ + +#include "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 + }; + table.all_inputs_static = false; + 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 + }; + table.all_inputs_static = false; + 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.all_inputs_static = false; + 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 static one. Pinning a nonzero +// profile for the former must not fail the latter, which has one shape and so +// nothing to switch. +TEST(ExecuTorchOptimizationProfileSelection, StaticEngineToleratesAPinItCannotHonor) { + ProfileTable static_engine; + static_engine.bounds = {{bounds({1, 16}, {1, 16})}}; + static_engine.all_inputs_static = true; + 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::kOk); + EXPECT_EQ(selected, 0); +} + +// A single-profile *dynamic* engine genuinely cannot honor a nonzero pin, so it +// reports instead of quietly running the wrong regime. +TEST(ExecuTorchOptimizationProfileSelection, SingleProfileDynamicEngineRejectsANonzeroPin) { + ProfileTable dynamic_engine; + dynamic_engine.bounds = {{bounds({1, 1}, {1, 2048})}}; + dynamic_engine.all_inputs_static = false; + int32_t selected = -1; + + EXPECT_EQ( + select_profile(dynamic_engine, ProfileRequest::kPinned, 1, decode_input(), selected), + ProfileSelection::kRequestedProfileUnavailable); +} + +} // namespace +} // namespace executorch_backend +} // namespace torch_tensorrt