Repository navigation
✨[Feature] a global switch for engine execution profiling, reachable from an AOTInductor package #4679
Description
Bug Description
Per-engine execution profiling cannot be turned on for a model deployed as an AOTInductor
package — which, for us, is the only form the model is ever measured in.
TRTEngine::profile_execution gates all of the per-phase timing in
core/runtime/execute_engine.cpp, and the only way to set it is enable_profiling() on an
engine object, registered in core/runtime/register_jit_hooks.cpp as a method on the torchbind
class:
.def("enable_profiling", &TRTEngine::enable_profiling)
.def("set_profile_format", &TRTEngine::set_profile_format)
.def("disable_profiling", &TRTEngine::disable_profiling)Every other runtime switch has a global form alongside it in the same file —
set_multi_device_safe_mode, set_cudagraphs_mode, set_logging_level — but profiling has
none.
An AOTInductor package deserializes its own engine objects inside the compiled artifact. A
caller never holds one, so there is nothing to call enable_profiling() on, and the packaged
artifact is exactly the configuration whose performance one wants to explain.
To Reproduce
docker run --rm --gpus all --ipc=host -v "$PWD":/w -w /w \
nvcr.io/nvidia/pytorch:26.07-py3 python repro.py
repro.py
import torch
import torch.nn as nn
import torch_tensorrt
ROWS, COLS = 8, 16
PROFILING_HINTS = ("profil", "timing")
class Model(nn.Module):
"""Any model at all; the gap is in how the artifact is reached, not what it computes."""
def forward(self, x: torch.Tensor) -> torch.Tensor:
return torch.sigmoid(x * 2.0 + 1.0)
def main() -> None:
model = Model().eval().cuda()
x = torch.randn(ROWS, COLS, device="cuda")
exported = torch.export.export(model, (x,))
gm = torch_tensorrt.dynamo.compile(
exported,
inputs=(x,),
min_block_size=1,
pass_through_build_failures=True,
)
# The eager path: an engine object is in hand, so profiling is reachable.
eager_engine = next(
mod for _, mod in gm.named_children() if getattr(mod, "serialized_engine", None)
)
eager_engine.setup_engine()
print(
f"eager path, engine has enable_profiling: {hasattr(eager_engine.engine, 'enable_profiling')}"
)
global_switches = sorted(
n for n in dir(torch.ops.tensorrt) if any(h in n.lower() for h in PROFILING_HINTS)
)
print(f"global ops matching {PROFILING_HINTS}: {global_switches}")
ep = torch_tensorrt.dynamo.export(gm, arg_inputs=[x])
package = torch._inductor.aoti_compile_and_package(ep)
runner = torch._inductor.aoti_load_package(package)
print(f"aoti path, runner ran: {tuple(runner(x).shape)}")
engines_on_runner = sorted(
n for n in dir(runner) if "engine" in n.lower() and not n.startswith("__")
)
print(f"engine handles on the loaded package: {engines_on_runner}")
reproduced = not global_switches and not engines_on_runner
print(f"\nreproduced: {reproduced}")
if reproduced:
print(
"Profiling is reachable only through an engine object, and an AOTInductor "
"package never hands one out -- so the form the model is actually measured in "
"is the one form that cannot be profiled."
)
if __name__ == "__main__":
main()
output
eager path, engine has enable_profiling: True
global ops matching ('profil', 'timing'): []
... AOTInductor compile and the fake-class warnings it emits, elided ...
aoti path, runner ran: (8, 16)
engine handles on the loaded package: []
reproduced: True
Profiling is reachable only through an engine object, and an AOTInductor package never hands one out -- so the form the model is actually measured in is the one form that cannot be profiled.
Expected behavior
A global switch, in the shape of the ones already next to it in
core/runtime/register_jit_hooks.cpp:
torch.ops.tensorrt.set_profile_execution(True)or an environment variable read at engine construction, so that engines an artifact
deserializes for itself pick it up.
What we needed out of it, on a 14-engine model, was per-engine and per-phase host wall time:
engine calls total ms workspace inputs outputs enqueue
_run_on_acc_19_engine 25 8636.1 0.6 0.0 0.0 0.0
_run_on_acc_17_engine 25 4410.5 0.1 90.0 18.1 1567.0
ALL 27587.5
The split is what makes it useful: enqueueV3 is asynchronous, so the enqueue column is launch
cost rather than GPU time, while the allocations and the input packing are synchronous and
theirs is real. With that we could say which three engines held half the total, and rule out
the per-call workspace allocation (≤0.6 ms over 25 calls) as a cause. Without it, an
AOTInductor artifact is one number.
A dump entry point (torch.ops.tensorrt.dump_engine_timing()) would pair naturally with it,
since a packaged artifact has no object to read the counters off either.
Environment
Build information about Torch-TensorRT can be found by turning on debug messages
- Pytorch NGC container : 26.07-py3
Activity
abhinav-phi commentedon Sep 4, 2026
Hi! I'd like to work on this.
Plan: I'll design the feature against the described behavior, implement it with the project's conventions, add unit/integration tests, and update the docs.
Could you please assign this issue to me? Thanks!
/assign take
SrivastavaKshitij commentedon Sep 15, 2026
Go ahead and put up a PR