Multi-optimization profile for executorch

[cehongwang]

Multi-Optimization-Profile Support in the ExecuTorch Backend

Targets	torch_tensorrt.executorch (Python export), cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp (C++ runtime), TR01 blob format
Related	ExecuTorch design overview: other/executorch.md

1. Problem

A TensorRT engine can carry N optimization profiles — one weight set, several
kernel tunings, each valid over a different input-shape range. This matters for
bimodal workloads like LLM inference:

Phase	input_ids	Notes
Prefill	[B, 32..2048]	long sequence, one shot
Decode	[B, 1]	one token at a time, latency-critical

One profile spanning [1, 2048] picks kernels wrong for both phases; separate
decode/prefill profiles let TensorRT tune each path and the runtime switch per
call.

Profiles are declared with
torch_tensorrt.Input(profiles=[{min_shape, opt_shape, max_shape}, ...]) — an
ordered list whose index is the profile index — and selected at inference
by integer index (or "auto"). There are no profile names.

The standard runtime (Python TRTEngine / C++ torch::classes::tensorrt::Engine)
already does this: it reconstructs per-profile bounds by index from the
deserialized engine (core/runtime/TRTEngine.cpp:633) and switches via
set_active_profile / auto_select_profile (TRTEngine.cpp:687-782). Selection
is purely runtime; a freshly loaded engine defaults to index 0, auto-select off
(TRTEngine.h:297-299).

The ExecuTorch export path does not support this, for three reasons:

1. No libtorch / TRT C++ runtime at inference. output_format="executorch"
   (other/executorch.md §1) produces a .pte run by a small C++ runner with only
   ExecuTorch + TensorRT + CUDA; TRTEngine::set_active_profile does not exist there.

2. The backend hardcodes profile 0. initialize_input_profiles()
   (TensorRTBackend.cpp:170-191) reads bounds from profile 0 only:

   bounds.min = handle.engine->getProfileShape(name.c_str(), 0, kMIN);
   bounds.max = handle.engine->getProfileShape(name.c_str(), 0, kMAX);

   A multi-profile engine loads, but only profile 0 is ever used.

3. No Python at inference, so no context manager. A .pte runs via
   method.execute(); selection must be expressible through a C++ runtime hook.

This design extends multi-profile support to the ExecuTorch path so a
multi-profile-compiled graph can be saved as a .pte and run with the right
profile picked at inference — by index.

2. Goals / Non-goals

Goals

• Make the C++ backend multi-profile aware: discover all N profiles at
  init(), validate against the active profile, switch via
  setOptimizationProfileAsync.
• Index-based semantics matching the standard runtime: manual pin by index
  (default 0), opt-in lazy first-fit auto-select, sticky "keep the active profile
  if it still fits".
• A runtime hook to pin a profile by index (and toggle auto-select) without
  re-loading the .pte.
• Unchanged authoring API (Input(profiles=[...])).
• Zero TR01 and zero Python-export change; existing .ptes load identically.

Non-goals

• No general "compile ATen inside EXIR to TRT" path (other/executorch.md §4.2:
  engines come from Dynamo, ExecuTorch only wraps them).
• No cross-partition profile coordination; each engine is independent.
• No profile-aware CUDA Graph caching (the backend uses no CUDA Graphs).
• No profile names; index-only.
• No compile-time active-profile default — selection is runtime state, not baked
  into the artifact (see §5, §7).

3. Design summary

The ExecuTorch runtime reimplements the standard runtime's index-based logic
(setup_optimization_profiles, profile_fits, auto_select_profile,
set_active_profile) without libtorch. Selection is entirely runtime, via two
paths:

Priority	Mechanism	Where it lives	When to use
1	Manual pin by index	C++ hook on TensorRTBackend (§6.3)	Per-call switching (LLM prefill/decode)
2	Auto-select by input shape	C++ execute() (§6.2), enabled via the hook	Shapes drive selection

Opt-in: with no runtime call a .pte behaves as today (active_profile_index == 0, auto off). A single-profile .pte additionally has num_optimization_profiles == 1 and never switches — byte-for-byte identical to current behavior.

4. TR01 blob format: no change

Everything the backend needs is reconstructable from the deserialized engine:

• num_profiles — from engine->getNbOptimizationProfiles().
• Per-profile per-input (min, max) — from engine->getProfileShape(name, p, kMIN | kMAX).

Since selection is index-based and made at runtime, there are no names and no
policy to persist. TR01 metadata, TENSORRT_MAGIC ("TR01"), and HEADER_FORMAT
are untouched.

5. Python export

5.1 No code change to the export path

No Python-side transformation is required.

• TensorRTPartitioner — no change. The active profile is runtime execution
  state, not a graph property. Unlike target_device (which the partitioner sets
  because PropagateDevicePass bakes it into the .pte's extra_tensor_info),
  profile selection has no AOT effect — it changes no nodes, lowering, engine
  bytes, or memory planning (outputs are pre-allocated at the max envelope
  regardless). Nothing for the partitioner to do.
• TensorRTBackend.preprocess() — no change. Builds the TR01 blob as today;
  all N profiles are inside the serialized engine (§4).
• _replace_execute_engine_for_executorch — no change. Profile-count
  agnostic.

All profile behavior lives in the C++ backend (§6).

5.2 All engines share the same profile count by construction

The runtime hooks (§6.3) assume every TRT delegate in a .pte has the same
number of optimization profiles, in the same order. No extra enforcement is
needed to guarantee this — it is a direct consequence of how profiles are
specified at compile time:

1. The user passes a single optimization_profiles=[prefill, decode, ...] list
   to torch_tensorrt.dynamo.compile.
2. That one list is applied to every engine partition the graph is split
   into. For each partition, the interpreter calls add_optimization_profile
   once per list entry.
3. Therefore every engine reports the same getNbOptimizationProfiles(), and
   profile index i refers to the same shape regime (e.g. "decode") in every
   engine.

Because there is exactly one profile list and it is applied uniformly, a single
integer index is meaningful across all delegates in the .pte.

5.3 Enforcing the invariant at save time

Because the invariant is what makes the whole-method set_active_profile
(§6.3) atomic and index-meaningful, we validate it explicitly in
_save_as_executorch rather than trust it implicitly:

• After the engines are available (each carries
  settings.optimization_profiles in its serialized_metadata, and the
  profile count is also recoverable from the engine bytes), assert that every
  TRT partition reports the same profile count.
• On mismatch, raise at save time with an actionable message naming the
  offending partitions and their counts, rather than letting a runner discover
  it via a confusing per-delegate index error at inference.

This is the one small Python-side addition; it is validation only (no graph or
blob mutation) and is a no-op for the common single-engine .pte.

6. C++ backend changes

Add the same multi-profile members/methods the standard runtime has, adapted to
EngineHandle.

6.1 EngineHandle extensions + initialize_input_profiles()

Extend the handle (cpp/include/torch_tensorrt/executorch/TensorRTBackend.h):

struct EngineHandle {
  // ... existing fields ...

  int32_t num_profiles = 1;
  int32_t active_profile_index = 0;   // default 0; runtime hook can move it
  bool auto_select_profiles = false;  // default off; runtime hook can enable

  // Per input, the dims that vary within or across profiles, each with
  // per-profile (min, max). An empty entry = fully-static / shape-inference
  // input, skipped during matching.
  struct ProfileDynamicDim {
    int32_t dim_index;
    std::vector<std::pair<int64_t, int64_t>> profile_ranges;  // size == num_profiles
  };
  std::vector<std::vector<ProfileDynamicDim>> profile_dynamic_dims;  // outer: input index
};

Rewrite initialize_input_profiles() (TensorRTBackend.cpp:170-191) to enumerate
all profiles instead of hardcoding 0, building profile_dynamic_dims: gather
(min, max) for every dim across every profile, keeping only dims that vary within
a profile (min != max) or differ across profiles (a dim with the same fixed
extent everywhere can't distinguish profiles; TensorRT validates it at
setInputShape). Shape-tensor inputs are still rejected.

handle.num_profiles = handle.engine->getNbOptimizationProfiles();
if (handle.num_profiles > 1) {
  // build handle.profile_dynamic_dims[input][dim] = per-profile (min,max),
  // skipping dims that can't distinguish profiles.
}

init() leaves the defaults (active_profile_index = 0, auto_select_profiles = false) and does not parse compile_specs — matching the standard runtime's load
behavior.

6.2 execute() bounds check + auto-select (TensorRTBackend.cpp:398-479)

Today the per-input check compares against profile 0. Two changes:

1. Validate against active_profile_index.
2. Auto-select when num_profiles > 1 && auto_select_profiles:
   • Sticky: keep the active profile if it still fits all inputs (avoids
     thrashing when shapes alternate).
   • Else lazy first-fit: scan 0..N-1, pick the first whose [min, max]
     contain every input, then setOptimizationProfileAsync(p, stream).
     Overlapping profiles resolve to the lowest matching index — declare the more
     specific regime (e.g. decode) first, or pin manually.
   • If none fit → Error::InvalidArgument pointing at the pin hook (§6.3).

Invariants:

• The switch runs on the same CudaStreamGuard-scoped stream as enqueueV3 (no
  default-stream sync).
• After a switch, output shapes are re-inferred via inferShapes — which
  execute() already does each call, so no extra bookkeeping.

6.3 Runtime pin: new index-based backend hook

Selection is runtime, so the primary API is a runner-callable pin that bypasses
method.execute(). Backends are addressable via the Backend registry
(register_backend at the bottom of TensorRTBackend.cpp), so add a small
index-based side-channel.

Invariant: all TRT delegates in a .pte share the same profile count.
Every engine is built with the same number of optimization profiles, in the
same order (profile index i means the same regime — e.g. "decode" — across
all engines). This makes a single integer index globally meaningful, so the
whole-method pin can be atomic and strict: it validates the index against
the common profile count once and either applies to every delegate or fails
without touching any of them. §5.3 covers how the invariant is enforced at
export time.

// New header: cpp/include/torch_tensorrt/executorch/TensorRTBackend.h
namespace torch_tensorrt::executorch_backend {

// Pin the active profile (by index) on all TensorRT delegates owned by `method`.
// - profile_index >= 0: pin that index on every delegate (disables auto-select).
// All delegates share one profile count (see §5.3), so the index is validated
// once up front: on an out-of-range index NO delegate is modified and
// Error::InvalidArgument is returned (atomic, all-or-nothing).
::executorch::runtime::Error set_active_profile(
    ::executorch::runtime::Method& method,
    int32_t profile_index);

// Enable/disable shape-based auto-selection on all TRT delegates of `method`.
::executorch::runtime::Error set_auto_select_profiles(
    ::executorch::runtime::Method& method,
    bool enabled);

// Per-delegate variant: pin only ONE engine to a different index than the rest.
// The shared-count invariant guarantees the index range is identical to what the
// whole-method call accepts; this overload exists only to let a multi-engine
// `.pte` run different engines on different profiles simultaneously.
::executorch::runtime::Error set_active_profile(
    ::executorch::runtime::DelegateHandle* handle,
    int32_t profile_index);

// RAII guard for scoped profile selection, modeled on PyTorch's device/stream
// guards (c10::DeviceGuard, c10::cuda::CUDAStreamGuard): the ctor pins the
// profile, the dtor restores the previous selection on scope exit.
class OptimizationProfileGuard {
 public:
  OptimizationProfileGuard(::executorch::runtime::Method& method, int32_t profile_index);
  ~OptimizationProfileGuard();  // restores prior (active_profile_index, auto_select) state
  OptimizationProfileGuard(const OptimizationProfileGuard&) = delete;
  OptimizationProfileGuard& operator=(const OptimizationProfileGuard&) = delete;
};

}  // namespace torch_tensorrt::executorch_backend

set_active_profile walks method's delegates, finds each EngineHandle* owned
by TensorRTBackend, validates the index, and calls the same switch helper as the
auto-select path. Runner pattern:

constexpr int32_t DECODE_IDX = 0, PREFILL_IDX = 1;
{
    torch_tensorrt::executorch_backend::OptimizationProfileGuard guard(method, PREFILL_IDX);
    method.execute();   // uses profile 1 (prefill)
}
{
    torch_tensorrt::executorch_backend::OptimizationProfileGuard guard(method, DECODE_IDX);
    for (int i = 0; i < num_tokens; ++i) {
        method.set_input(input_ids, 0);
        method.execute();   // uses profile 0 (decode)
    }
}

Opt-in and additive: runners that only call method.execute() get profile 0 /
auto off — no code change.

7. Alternatives considered

Approach	Why rejected
Persist profile names in TR01	Index-only; bounds/count reconstructable by index (§4).
Export-time default active profile (TR01 field or CompileSpec)	The active profile is runtime state with no AOT effect (unlike target_device, §5); baking it would conflate partitioning with execution and diverge from the standard runtime.
Profile index as a hidden scalar input	Needs re-partitioning, leaks a TRT concept into the model signature, breaks EP shape invariants.
One .pte per profile	2× artifact/load, no shared weights; loses "one engine, N profiles".
Pure auto-select (no pin hook)	Fails for overlapping ranges shape alone can't disambiguate.

8. Backward compatibility

1. Existing .ptes load unchanged. A single-profile engine reports
   getNbOptimizationProfiles() == 1 → num_profiles == 1, active_profile_index == 0, auto off, empty profile_dynamic_dims — today's profile-0 path.
2. New .ptes on old runners. An old libexecutorch_trt_backend.a reads
   profile 0 only — no crash, just no switching.
3. Single-profile compile flow unchanged. No Input(profiles=[...]) → same
   code path as today.

9. Portability considerations

• No CUDA-graph bookkeeping. The backend calls enqueueV3 directly; a switch
  just needs setOptimizationProfileAsync + the inferShapes it already runs.
• AOT vs runtime split (see §5): target_device is AOT (baked into the
  .pte); profile selection is a runtime call with no serialized footprint — hence
  partitioner vs C++-hook layering.
• Multi-engine .pte. Every engine is built with the same profile count and
  order (§5.3), so a single index is meaningful across all of them. The
  whole-method set_active_profile validates the index once against that common
  count and applies atomically (all delegates switch, or none do and it returns
  Error::InvalidArgument). The per-DelegateHandle variant (§6.3) still lets a
  caller pin one engine to a different index than the rest when engines must run
  different regimes concurrently.

10. Implementation plan

Ordered for isolated landing/testing:

1. Multi-profile bounds. Extend EngineHandle (§6.1); rewrite
   initialize_input_profiles() to build profile_dynamic_dims. execute() still
   uses index 0. Existing tests untouched.
2. Auto-select in execute(). Add profile_fits + auto_select
   (sticky-then-first-fit). Test a two-profile engine ([1,1] and [1,MAX])
   alternating inputs + an overlap-ordering case.
3. Save-time invariant check. Add the shared-profile-count assertion to
   _save_as_executorch (§5.3). Unit test: a hand-constructed multi-engine graph
   with mismatched counts raises a clear error; matching counts pass.
4. Runtime hooks. set_active_profile / set_auto_select_profiles /
   OptimizationProfileGuard + build registration. Reference-runner test: pin 1 then 0, assert
   the switch and correct outputs; also test set_auto_select_profiles(true) and an
   out-of-range index returning Error::InvalidArgument without switching any
   delegate.
5. Docs & example. Multi-profile export/run pair in
   examples/torchtrt_executorch_example/; update other/executorch.md.

Steps 1–3 are independently useful and don't depend on the hooks.

11. Open questions

1. Should we enable auto-selection by default? Need a benchmark to decide

[narendasan]

What does 3.2 When to use: Shapes drive selection mean?

[cehongwang]

That is the auto-select feature right? Select the profile based on input shape

[narendasan]

Opt-in: with no runtime call a .pte behaves as today (active_profile_index == 0, auto off). A single-profile .pte additionally has num_optimization_profiles == 1 and never switches — byte-for-byte identical to current behavior.

Do you think auto should be default?

[cehongwang]

Well that introduces extra latency. We can benchmark the latency how much the impact is and decide whether it should be default.

[narendasan]

5.2 Shared profile count comes for free from RFC 0001

What does this mean?

[cehongwang]

That's the previous RFC for Torch-TensorRT multi-optimization-profile enablement. I reworded to remove any ambiguity.

[narendasan]

6.3

I like this but lets call this a OptimizationProfileGuard or OptProfileGuard, similar to StreamGuards or other similar mechanisms in PyTorch

[narendasan]

https://docs.pytorch.org/docs/2.13/accelerator/guard.html

[shoumikhin]

Review of RFC #4417: Multi-Optimization-Profile Support in ExecuTorch

Thanks for putting this RFC together. I support the feature and think we should keep moving forward with it.

The problem is real, especially for LLMs. Prefill and decode are very different workloads:

Prefill: process a long prompt, such as 1,000 tokens
Decode:  process one new token at a time

Using separate TensorRT optimization profiles is like giving the engine two gears:

Prefill profile: tuned for large inputs
Decode profile:  tuned for tiny inputs

Keeping both profiles in one engine is useful because the profiles can share the same weights. Discovering the profiles from the deserialized TensorRT engine also makes sense, and I agree that the TR01 format probably does not need to change.

My recommendation is to approve the direction but revise the runtime-control design before implementation.

Main issue: users normally run an ExecuTorch Module

Most ExecuTorch users run a model through the high-level Module API:

executorch::extension::Module module("model.pte");
auto result = module.forward(inputs);

They normally do not obtain a raw Method or a DelegateHandle.

The RFC currently proposes APIs that walk a Method and find all of its TensorRT delegate handles. ExecuTorch does not expose those handles through its public API, so that implementation cannot be built safely as written.

I would not solve this with a global registry of handles keyed by method name. Different modules commonly have a method named forward, so the names are not unique. Such a registry would also introduce difficult ownership and cleanup rules when modules unload.

Suggested API: follow CallerStreamGuard

ExecuTorch already has a good model for execution-scoped settings: CallerStreamGuard.

For public/OSS context, this shared caller-stream API is newer than the currently released ExecuTorch API. It comes from work being upstreamed in ExecuTorch PRs #20158 and #20498. The useful part here is the pattern: one setting for the current thread is read by every relevant delegate when it runs.

A caller can write:

CallerStreamGuard stream_guard(stream);
auto result = module.forward(inputs);

The guard stores a value for the current thread. CUDA-capable delegates read it when they execute. It does not inspect the Module, Method, or delegate handles.

Optimization-profile selection can follow the same design. I would keep the name OptimizationProfileGuard, as suggested in the discussion, but change its API from the RFC's OptimizationProfileGuard(Method&, int) to an execution-scoped OptimizationProfileGuard(index) that does not need a raw Method:

executorch::extension::Module module("model.pte");

// Run a long prompt.
{
  OptimizationProfileGuard profile_guard(PREFILL_PROFILE);
  auto result = module.forward(prefill_inputs);
}

// Generate one token at a time.
{
  OptimizationProfileGuard profile_guard(DECODE_PROFILE);
  auto result = module.forward(decode_inputs);
}

The guard should only store an execution request such as:

Default           Use the normal default, currently profile 0
Pinned(index)     Use one exact profile
Auto              Choose a profile from the input shapes

It should not call TensorRT directly in its constructor or destructor.

CallerStreamGuard and OptimizationProfileGuard do different jobs and can be used together. The stream guard says where GPU work runs. The profile guard says which TensorRT profile to use. The profile switch should use the stream that TensorRTBackend::execute() already selected:

CallerStreamGuard stream_guard(stream);
OptimizationProfileGuard profile_guard(PREFILL_PROFILE);
auto result = module.forward(prefill_inputs);

Each TensorRT delegate should apply the request inside TensorRTBackend::execute():

1. Lock the engine.
2. Wait for earlier work on that execution context.
3. Read the current thread's profile request.
4. Collect all runtime input shapes.
5. Choose the profile index.
6. Validate every input rank and dimension against that profile's bounds.
7. Switch the TensorRT context on the actual execution stream.
8. Set all TensorRT input shapes.
9. Bind addresses and run TensorRT.

This keeps the TensorRT operation where the engine, mutex, in-flight event, CUDA stream, and input shapes are already available.

The guard should be created on the same thread that calls module.forward() or module.execute(). ExecuTorch already documents that concurrent execution on the same Module instance is unsupported, so this feature does not need to make one shared module thread-safe.

Profile selection must happen before input binding

TensorRT requires setOptimizationProfileAsync() to happen before setInputShape() for dynamic inputs.

The implementation therefore needs a small pre-pass over the inputs. It cannot choose a profile halfway through the current per-input binding loop.

The current clear rank/bounds check should remain. For a single-profile engine, it should behave exactly as it does today. For a multi-profile engine, it should check the selected profile and return a useful input error before enqueue if a rank or dimension is invalid.

In simplified form:

collect_all_input_shapes();
select_profile_index();
validate_all_inputs_against_selected_profile();
switch_profile_on_execution_stream();
set_all_input_shapes();
bind_all_addresses();
enqueue();

Static TensorRT delegates should remain valid

A partitioned model can reasonably contain:

Dynamic TensorRT delegate: 2 profiles
Static TensorRT delegate:  1 profile
CUDA delegate:              no TensorRT profiles

That should not be rejected just because the delegates do not all report the same profile count.

A practical rule could be:

Auto mode:
  A one-profile TensorRT delegate stays on profile 0.

Pinned profile 0:
  A one-profile TensorRT delegate stays on profile 0.

Pinned profile N, where N > 0:
  A multi-profile delegate uses N.
  A truly static one-profile delegate does not need to switch.
  A dynamic one-profile delegate that cannot honor N returns an error.

Non-TensorRT delegates:
  Ignore the request.

If profile indices are intended to represent the same global regimes across several dynamic TensorRT delegates, the export path should preserve and validate that relationship explicitly. Equal profile counts alone do not prove that profile index 1 has the same meaning in every independently compiled partition.

This simpler API intentionally drops the RFC's per-DelegateHandle feature for pinning different TensorRT engines inside one forward() to different profiles. Normal Module users cannot get that handle. It also conflicts with the idea that one profile index means one mode for the whole method. I think it is fine to leave this out of the first version. If we need it later, it should use a stable delegate ID or a generic ExecuTorch execution-options API, not a raw handle from the user.

Do not promise physical all-or-nothing switching

The RFC currently describes whole-method switching as atomic.

Imagine a method with three TensorRT delegates:

Delegate A switches successfully
Delegate B switches successfully
Delegate C fails

TensorRT does not provide a transaction that can undo the earlier switches. One profile request can apply to the whole module.forward() call, but we cannot promise that several TensorRT engines switch all at once or not at all.

I suggest changing the guarantee to:

One execution sees one consistent profile-selection request. Each TensorRT delegate applies that request when it executes and may return an error if it cannot honor it.

Clarify the automatic-selection rule

The RFC currently combines two rules:

Keep the active profile if it still fits.
Otherwise, choose the lowest matching profile.

This is history-dependent when profiles overlap.

For example:

Profile 0: exactly 1 token
Profile 1: 1 to 2048 tokens

After a prefill call, profile 1 is active. A one-token decode input also fits profile 1, so the sticky rule keeps profile 1 instead of returning to profile 0.

That may be intentional, but it is different from saying that overlapping profiles always resolve to the lowest matching index.

Please choose and document one policy. For example:

Preferred-first:
  Find the preferred matching profile on every call.
  Keep the active profile only when it is still the preferred match.

Alternatively, keep the current sticky policy but describe it accurately:

Sticky-first-fit:
  Keep the active profile if it fits.
  Only scan from profile 0 when the active profile no longer fits.

For prefill/decode workloads, explicit pinning may be more predictable than automatic selection.

Authoring and export details

The RFC should consistently use the landed authoring API:

Input(profiles=[...])

rather than a separate compile(optimization_profiles=[...]) argument.

It should also name the prerequisite multi-profile authoring/runtime work and the minimum Torch-TensorRT version or commit that contains it. Without that prerequisite, this is not a C++-only change.

The proposed save-time validation refers to settings.optimization_profiles, but that is not currently a serialized source of profile count and order. The RFC should choose one real source of truth:

1. Persist a small canonical profile signature in engine metadata, or
2. Deserialize and inspect every TensorRT engine during ExecuTorch export.

The first option looks simpler. It would let us check that profile index 1 means the same thing in every partition instead of guessing from the profile count. It can use the existing serialized metadata field without changing the TR01 wire format.

Validation should live at the public ExecuTorch export/lowering boundary, before final memory planning and serialization, rather than being tied to an older private save helper.

Output-buffer safety

For the normal authoring path, Input(profiles=[...]) creates an export range that covers the union of all profile ranges.

For example:

Decode profile maximum sequence: 1
Prefill profile maximum sequence: 2048
Exported Dim maximum sequence:    2048

That is the correct starting point for ExecuTorch upper-bound memory planning.

The RFC should state this rule clearly:

Every selectable TensorRT profile must remain inside the bounded torch.export input domain used to plan the .pte, and every bounded symbolic output must fit the resulting ExecuTorch output capacity.

Please add an ExecuTorch test where a nonzero profile produces the largest output:

Profile 0 output: [1, 1, 16]
Profile 1 output: [1, 64, 16]

Running profile 1 through Module::forward() should successfully return [1, 64, 16] without exceeding the planned output buffer.

Data-dependent outputs that require TensorRT's output allocator can remain unsupported and should be documented separately.

Compatibility wording

I agree that existing single-profile .pte files should continue to work on a new runner.

A new multi-profile .pte may also load on an old runner, but the old runner only understands profile 0. An input that is valid only for profile 1 will fail.

Please call that file/load compatibility, not full runtime compatibility.

Adding fields to the public EngineHandle struct changes its C++ layout. This is fine in the normal source/static build when the header and backend archive are rebuilt together. It is not safe to mix a new header with an old prebuilt backend, or the other way around. If possible, I would keep EngineHandle private.

Named methods are a useful fallback, not a replacement

ExecuTorch can also represent separate named methods:

module.execute("prefill", prefill_inputs);
module.execute("decode", decode_inputs);

That is a useful fallback and should be mentioned in the alternatives section.

However, separate methods will generally use separate TensorRT engines and may duplicate engine state or weights. The multi-profile design is still valuable because it provides one engine with shared weights and several tuned shape regimes.

Tests requested before approval

There is already useful multi-profile coverage for the normal Torch-TensorRT runtime, but it does not exercise the ExecuTorch .pte backend.

Please add an end-to-end ExecuTorch test using the high-level Module API that covers:

• Manual profile 0 selection.
• Manual nonzero-profile selection.
• Automatic profile selection.
• Prefill -> decode -> prefill transitions.
• Overlapping profile ranges and the documented selection policy.
• A shape accepted only by a nonzero profile.
• A static TensorRT delegate together with a multi-profile TensorRT delegate.
• TensorRT and CUDA delegates in one .pte.
• An invalid profile index.
• A nonzero profile producing the largest output.
• A real serialize/load round trip before execution.
• Old-runner behavior being limited to profile 0.

The existing reference runner and tests can remain useful, but the primary example for this feature should show Module::forward() or Module::execute(), because that is the API most users call.

Recommendation

I support the feature and the overall direction.

Before approval, I recommend replacing the Method&/delegate-walking hooks with an execution-scoped, thread-local OptimizationProfileGuard used around Module::forward() or Module::execute(), following the existing CallerStreamGuard pattern.

In simple terms:

The multiple gears are useful. Instead of reaching inside the Module and trying to move every gear manually, let the guard announce which gear this execution wants. Each TensorRT delegate can then switch safely when it naturally runs.

With that API change, one clear auto-selection rule, clear export rules, and real ExecuTorch runtime tests, this looks like a good path forward.