Implemented executorch's multi-optimization profile

[cehongwang]

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[shoumikhin]

I reviewed the multi-profile runtime, policy tests, and reference documentation. I left one runtime correctness concern and two nonblocking test/documentation suggestions inline.

[shoumikhin]

Once setOptimizationProfileAsync(profile, stream) succeeds, several later failures can return before enqueueV3() and before the existing completion event is recorded. The profile switch may therefore remain in flight while the next execute() or destroy() reconfigures or destroys the same IExecutionContext. Could you add cleanup for every post-switch error path, either by synchronizing the stream or recording completion that the next execution and teardown will wait for? Delaying the profiles.active assignment alone would not address the context-lifetime race.

[shoumikhin]

Nonblocking test suggestion: could we add (1) a two-input table where profile 0 fits input 0 but not input 1, verifying auto-selection skips it, and (2) a three-profile table where the active profile no longer fits and two lower profiles do, verifying the rescan selects the lowest matching index? The implementation handles both today, but the current tests cover only one-input tables and a rescan with one fitting alternative.

[shoumikhin]

Nit: kPrefillProfile and kDecodeProfile are example-local constants and are not defined or exported by the public API, so this copied snippet does not compile as written. Could you define them in the snippet, or use profile indices 1 and 0 with a note that the indices follow the export-time profile order?

[shoumikhin]

Re-reviewed after the fixup commit. The three earlier comments are all addressed, thanks.

The core design holds up well. I checked the stream ordering against the contract in NvInferRuntime.h (the switch, the H2D copies, and the enqueue all use the one stream, so the required happens-before comes for free with no host sync), confirmed a fresh IExecutionContext really does start on profile 0, confirmed the default kSTATIC allocation means a switch never has to allocate, and confirmed every read and write of profiles.active is under the handle mutex. The selection policy matches _TRTEngine._auto_select_profile and TRTEngine::auto_select_profile exactly.

Two things I would like resolved before merge.

1. executorch-static-build is currently failing on this PR because the two new example .cpp files are not packaged. The job log shows Cannot find source file: multi_profile_main.cpp. Same job passes on the base commit, so it is this change. Because cmake aborts the whole configure, this also breaks the pre-existing runner for anyone unpacking the release tarball.

2. The mark_inflight refactor dropped an error return the previous code had, so a faulted inference can now report success on the skip-sync path.

The rest is smaller: a bare OptimizationProfileSelection.h on every consumer's include path, the static-vs-dynamic pin inconsistency, and the -1 sentinel in the public API, which is cheap to change now and expensive after release.

One request on scope. This is about 1760 added lines mixing the runtime feature, the in-flight-event refactor, two example programs, an export script, and a rewrite of the Python dynamo example. That last one touches the Python runtime example and is a separate concern from the C++ delegate. Splitting would make each piece much easier to review.

Validation note: static review only, no GPU run. The TensorRT semantics above come from the bundled header docs, and I did not reproduce the benchmark numbers.

[shoumikhin]

These two new runners are not packaged into the release tarball, and this is already failing CI.

examples/executorch_reference_runner/BUILD has a source_files filegroup that decides what ships:

filegroup(
    name = "source_files",
    srcs = ["CMakeLists.txt", "README.md", "main.cpp"],
)

That feeds executorch_reference_runner_pkg_files -> executorch_source_package -> libtorchtrt_tar. The PR does not touch it, so the tarball gets the new CMakeLists.txt but neither multi_profile_main.cpp nor multi_profile_benchmark.cpp.

The executorch-static-build job on this commit shows exactly that:

CMake Error at CMakeLists.txt:65 (add_executable):
  Cannot find source file:
    multi_profile_main.cpp
CMake Error at CMakeLists.txt:75 (add_executable):
  Cannot find source file:
    multi_profile_benchmark.cpp
CMake Generate step failed.  Build files cannot be regenerated correctly.

The same job passes on the base commit, so this is from this change. Worth noting the blast radius is wider than the two new targets: cmake aborts the whole configure step and emits no build files, so example_executorch_runner cannot be built either.

Could you add both files to the filegroup?

    srcs = [
        "CMakeLists.txt",
        "README.md",
        "main.cpp",
        "multi_profile_benchmark.cpp",
        "multi_profile_main.cpp",
    ],

The two new backend headers were correctly added to //cpp:executorch_backend_source_files, so this is the one spot that was missed. A require_tar_entry line per new file in verify-executorch-reference-runner.sh would also catch this at packaging time rather than at configure time.

[shoumikhin]

This refactor drops an error that the previous version reported.

Before, on the skip-sync path, a failed cudaEventRecord logged, drained, and returned an error:

  } else {
    cuda_err = cudaEventRecord(engine->inflight_event, stream);
    if (cuda_err != cudaSuccess) {
      ...
      (void)cudaStreamSynchronize(stream);
      engine->inflight_pending = false;
      return Error::InvalidProgram;   // <-- gone now
    }

Now mark_inflight returns void and also discards the return code of its own fallback cudaStreamSynchronize, so with must_sync == false execute() goes on to return Error::Ok.

Why it matters: the likely reason cudaEventRecord fails right after enqueueV3 is a sticky asynchronous CUDA error from the enqueue itself. In that case we now report success for a call whose inference actually faulted, and the error resurfaces later attributed to some unrelated operator. Rare path, but it used to be reported and now is not.

Could you have it return an Error and propagate at both call sites?

Error mark_inflight(EngineHandle& engine, cudaStream_t stream) {
  const cudaError_t rec = cudaEventRecord(engine.inflight_event, stream);
  engine.inflight_pending = (rec == cudaSuccess);
  if (rec == cudaSuccess) {
    return Error::Ok;
  }
  ET_LOG(Error, "TensorRTBackend::execute: cudaEventRecord failed: %s", cudaGetErrorString(rec));
  return cudaStreamSynchronize(stream) == cudaSuccess ? Error::Ok : Error::InvalidProgram;
}

Minor related note: on the must_sync path the event recorded at the tail is synchronized away immediately after, so that record is wasted work. Harmless, just noting it since the helper is now unconditional.

[shoumikhin]

This puts a bare OptimizationProfileSelection.h on the include path of every app that links the backend.

strip_include_prefix here is the header's own directory, so the header ends up with no directory prefix at all. Every other header target in this repo strips to a directory:

strip_include_prefix = "include"   # -> torch_tensorrt/executorch/TensorRTBlobHeader.h
strip_include_prefix = "include"   # -> torch_tensorrt/executorch/TensorRTBindingNames.h

Since this target is in deps of tensorrt_executorch_backend, the virtual include dir propagates transitively, so a downstream app with its own file of that name can shadow ours (or vice versa). Bazel-only hygiene, no behavior change, so low priority, but the rest of the project deliberately avoids this.

If you do change it, two things to watch. The bare spelling is what makes one #include work in both build systems, because the header lives in src/ and the CMake build only puts cpp/include on the path (cpp/src/torch_tensorrt/executorch/CMakeLists.txt:37), so switching to strip_include_prefix = "src" alone would fix Bazel and break CMake. And the two existing include sites would need updating too:

• cpp/src/torch_tensorrt/executorch/EngineHandle.h
• tests/cpp/executorch/test_optimization_profile_selection.cpp

Simplest version is probably to move the header to cpp/include/torch_tensorrt/executorch/, include it as "torch_tensorrt/executorch/OptimizationProfileSelection.h" in both places, and strip to include like the siblings. It is already effectively public since the test depends on it.

[shoumikhin]

This tolerance treats two engines that are equally unable to honor the pin differently.

Pinning index 1:

single-profile STATIC  engine -> silently runs profile 0, returns kOk
single-profile DYNAMIC engine -> kRequestedProfileUnavailable

Both have exactly one profile, so neither has an index 1. The comment says the point is not to fail an innocent single-profile sibling when the guard was aimed at a multi-profile one, but that applies just as much to the dynamic sibling, which still fails. It also does not help two multi-profile engines with different counts (say 3 and 2, pin index 2), since neither is size() == 1.

The two existing runtimes each pick one rule and stick to it:

	out-of-range pin, single-profile engine
Python _TorchTensorRTModule.set_optimization_profile	raises ValueError, always
C++ TRTEngine::set_active_profile_with_stream	silently no-ops for all single-profile engines

Could we match one of them? If you keep the tolerance, applying it to all single-profile engines regardless of static or dynamic, plus a warning log, would at least make an ineffective pin visible rather than silent.

Narrow case in practice (one .pte mixing a static engine with a multi-profile one, and a nonzero pin), so not blocking, but the asymmetry will be hard to explain later.

[shoumikhin]

Could you document the multi-delegate hazard here?

The guard sets one thread-local that every TensorRT delegate in the method reads, so if a .pte has two engines whose profile lists differ, index 1 can mean prefill in one and decode in the other. The comment below says the delegates see one consistent request, which is true of the integer but not of its meaning.

For contrast, the other two runtimes both target something specific:

# Python: targets a module object, can be scoped to one submodule
with optimization_profile(trt_gm, 1): ...

and the C++ runtime keeps active_profile_index per engine instance.

I think the thread-local is the right call here given the ExecuTorch BackendInterface. Its official set_option channel is process-global, which would be worse under concurrency. So this is not a redesign request, just a docs one so a user with two engines is not surprised.

One idea worth a thought, not for this PR: BackendExecutionContext::get_method_name() is available inside execute(), so if prefill and decode were exported as two methods the profile could be chosen from the method name with no ambient state at all.

[shoumikhin]

Any dtype that is not Half reads through a float* here, including BFloat16.

const_data_ptr<T>() is an unchecked static_cast in the portable tensor type, so on a bf16 tensor this walks 4 bytes per element through a 2-bytes-per-element buffer and reads roughly twice past the end. No assert, just wrong numbers and an out-of-bounds read.

Latent today since export_multi_profile.py does .to(torch.float16), so the paired model never hits it. Still worth guarding, especially as IndexTensor just above carries a comment about dtype mismatch being silent corruption. A BFloat16 branch plus a log-and-skip for anything unexpected would close it.

[shoumikhin]

Two small flag edge cases in the benchmark.

--block_rounds=1 means the r != 0 guard drops every prefill sample, so prefill stats always come out n=0. The default is 3 so this only bites someone who passes 1, but the empty result is silent rather than explained.

--blocks=0, which is also what atoi returns for garbage input, leaves wall_ms == 0, so the wall-time percentage at the end computes 0.0 / 0.0 and prints nan. percentile() would also underflow size() - 1 to SIZE_MAX on an empty vector; it is currently shielded only by the empty check in summarize.

Rejecting non-positive parsed values up front would handle all three.

[shoumikhin]

These two files quote the same measurement at different rates.

Here: 1.9 GB and roughly 3.7 s/MB.
export_multi_profile.py:38: roughly 3.6 seconds per megabyte and ~2 GB.

Worth making them agree, or dropping the per-MB rate from one of the two. The nearby ~3.6 ms switch cost also reads confusingly next to 3.6 s/MB, since they are unrelated quantities that happen to share a number.

[shoumikhin]

EngineHandle.h in srcs rather than hdrs deviates from the rest of the repo, where .cpp goes in srcs and .h in hdrs without exception.

It works, and putting a deliberately private header in srcs is a legitimate Bazel idiom that matches the "not installed" note in the file. Just unexpected for a reader, so a one-line comment saying it is intentionally private would help.

[shoumikhin]

Typo: "seperately" -> "separately".

[narendasan]

@cehongwang as these headers now constitute a C++ api can you put in doxygen annotations so we can render docs?

[shoumikhin]

Thanks for the updates. The -1 sentinel, all_inputs_static, the header move, the
packaging, the dtype guard in print_prediction, and the flag validation are all
addressed, and kPinIgnoredSingleProfile resolves the static-vs-dynamic asymmetry
cleanly.

What I checked by building rather than reading, in case it saves you time:

• Built and ran the policy test against TensorRT 10.15 and CUDA 12.8 headers: 16/16
  pass, no warnings under -Wall -Wextra.
• Traced nested guards: an inner automatic() correctly restores an outer pinned
  index, and the stale g_profile_index under kAuto is never read.
• Checked the new CMake ../../../include in both layouts. It resolves to
  cpp/include in the source tree and torch_tensorrt/include in the tarball, so it
  is right, and it is genuinely needed since the old runner included no
  torch_tensorrt header.
• Confirmed the D2H drain at TensorRTBackend.cpp:787-790 fixes a real pre-existing
  bug. The base code returned without draining while inflight_pending was false, so
  the destructor could cudaFree the staging buffers underneath live copies. Nice
  catch.

Of the inline comments, only the token-id one looks worth blocking on; it will hit the
first person who follows the README end to end. The TRTEngine.cpp throw-versus-warn
question is next, and the rest are follow-up material.

Two asks, both small:

1. Could you fill in the PR description? It is still the template.

2. Please note there that EngineHandle is no longer in the installed public header.
   It shipped in v2.13.0:

   $ git show v2.13.0:cpp/include/torch_tensorrt/executorch/TensorRTBackend.h \
       | grep -nE "^struct EngineHandle"
   48:struct EngineHandle {
   

   The move itself is right and the reasoning in EngineHandle.h is sound, and
   InputProfileBounds is still reachable since it moved to the installed
   OptimizationProfileSelection.h. It is only EngineHandle that downstream code
   can no longer name, so it is worth a line for whoever writes the release notes.

[shoumikhin]

The default export path feeds token ids past the end of the vocabulary.

const int64_t v = positions ? i : (1 + (static_cast<int64_t>(i) * 7919) % 9000);

export_multi_profile.py:101 gives the mini Gemma-3 vocab_size=2048, but this
formula produces ids up to 8976:

seq=  1  max_id=    1  ids >= 2048:   0/1
seq=128  max_id= 8976  ids >= 2048: 101/128
seq=256  max_id= 8976  ids >= 2048: 199/256

So the first call in main, the pinned seq=128 prefill, already has 101 out-of-range
ids, and the seq=256 prefill has 199. Decode is unaffected since seq=1 only ever
produces id 1.

TensorRT stores zero for an out-of-bounds gather (NvInfer.h: "Zero will be stored
for OOB access"), so this does not crash. It silently substitutes zero embeddings for
most of the prompt, which means the next_token values this runner prints for the
default model are meaningless. That is worse than a crash in one way: nothing tells
the user their getting-started run was garbage.

This is the documented first-run flow. The README exports the mini model by default,
then runs this binary against it. The real google/gemma-3-1b-it vocab is ~262k, so
--weights is unaffected, which is probably why it went unnoticed.

The exporter already handles this correctly with torch.randint(1, vocab, ...)
(line 236). The runners just need ids that cannot exceed the smallest vocab they
might be pointed at. Same line in multi_profile_benchmark.cpp:99.

[shoumikhin]

On a multi-profile engine this turns a hard failure into a warning.

Before, an out-of-range index skipped the <= 1 guard, reached
setOptimizationProfileAsync, got false back, and TORCHTRT_CHECK threw. Now it
warns and returns, so set_active_profile(99) on a 2-profile engine goes from raising
to silently continuing on whatever profile was already loaded, which means silently
mistuned kernels.

Making the single-profile case non-silent is a genuine improvement. Would you consider
keeping the throw for an out-of-range index on a multi-profile engine, and warning only
where the engine could not have done anything differently (one profile, any nonzero
index)?

Reachability is limited: set_optimization_profile validates first and raises
ValueError, so only a caller driving the engine directly can hit this. I also checked
the warning cannot spam a hot loop, since every per-call caller is gated on
num_optimization_profiles > 1 and passes an index it already checked with
profile_fits.

One small thing while you are here: this fixed the .cpp comment that pointed at
TorchTensorRTModule.resolve_profile_index, but the identical reference survives at
TRTEngine.h:300. That name has never existed as code (git log -S finds it only in
those two comments); the real validator is
TorchTensorRTModule.set_optimization_profile.

[shoumikhin]

select_profile can index bounds out of range on a malformed table.

The guard above checks only the outer vector, so bounds.empty() being false says
nothing about active naming a valid row:

  if (table.bounds.empty()) {
    return ProfileSelection::kNoProfileMatchesInputs;
  }
  ...
    if (profile_fits(table, table.active, input_dims)) {

profile_fits then does table.bounds[static_cast<size_t>(profile)] (line 91) with no
range check. A size-2 table with active = 5 segfaults. The same applies if a bounds
row is shorter than input_dims, since bounds[i] is indexed over input_dims.size()
rather than the row length.

Neither is reachable through execute() today. profiles.active is only written at
TensorRTBackend.cpp:596 with a value select_profile already validated, and
initialize_input_profiles builds exactly num_inputs bounds per row. So this is
latent, not a live bug, and I would not hold the PR for it.

Raising it because of the comment right above the guard:

Checked here so the policy is safe to call on its own rather than on the strength of
a guard in another translation unit.

That is the bar this header sets for itself, and it is installed public API, so "on its
own" includes callers you do not control. The empty-table case got defense in depth;
the two sibling cases that actually crash did not. Either range-check active and the
row length, or narrow that comment to state the precondition. Fine as a follow-up.

[shoumikhin]

This points at tests that do not exist yet.

 * Deliberately out of scope, because all of it needs a live engine: applying the
 * decision (setOptimizationProfileAsync), writing profiles.active back, the
 * ordering that puts the switch before setInputShape, and mark_inflight. Those
 * belong to the end-to-end ExecuTorch tests.

There are no end-to-end ExecuTorch tests covering any of it. The only automated
live-engine path is verify-executorch-reference-runner.sh, which exports via
export_static_shape.py (single-profile, static) and runs example_executorch_runner,
which installs no guard. So if (profile != engine->profiles.active) at
TensorRTBackend.cpp:586 is never true there. tests/py/dynamo/executorch/ is
composition-only and says so, and the multi-profile tests under
tests/py/dynamo/runtime/ cover the standard runtime, not this backend.

The policy tests themselves are good and I confirmed all 16 pass. The gap is only that
everything the policy hands off to is unverified.

I am not asking for a GPU test in this PR. Could the comment just say the applied path
is not covered yet, rather than deferring to tests that were never written? If you do
want cheap coverage later, multi_profile_main.cpp already asserts profile 99 is
rejected and returns nonzero on failure, and the verify job has a GPU, so running it
against the mini .pte would cover most of this.

[shoumikhin]

The wall-time row uses the opposite sign convention from its own heading.

compare() computes prefill_only - switching (line 235), so positive means switching
won, matching the "positive = switching is faster" heading on line 324. The wall-time
percentage computes switching - prefill_only, so when switching is faster it prints a
negative number under that heading. Negating it would make the block consistent. The
note below is about contention bias, so it does not cover this.