✨[Feature] a global switch for engine execution profiling, reachable from an AOTInductor package

[SrivastavaKshitij]

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

[abhinav-phi]

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]

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

Go ahead and put up a PR