✨[Feature] a global switch for engine execution profiling, reachable from an AOTInductor package
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3k
- Forks
- 410
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 78
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
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with core/runtime/register_jit_hooks.cpp to compare the existing global runtime switches, then trace profile_execution through core/runtime/execute_engine.cpp and run repro.py. Done means profiling can be enabled through a global operation or construction-time setting and is honored by engines deserialized inside an AOTInductor package, with the resulting per-engine timing information accessible for inspection.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- backend, machine-learning, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 52/100