Torch Inductor AOT Backend Guide | NVIDIA AITune Documentation

The Torch Inductor AOT backend compiles models Ahead-of-Time using PyTorch’s AOT Inductor (torch._inductor.aoti_compile_and_package). The result is a self-contained .pt2 artifact that can be saved, loaded, and executed without Python-interpreter overhead.

Requires PyTorch ≥ 2.6.

Overview

AOT Compilation: Model compiled once, loaded as a native artifact at inference time
Portable Artifact: .pt2 file contains everything needed to run inference
Dynamic Shapes: Batch and spatial dimensions automatically detected and marked dynamic
No Python Overhead: Inference runs through a compiled runner with no Python graph tracing
Save / Load: Artifact persists across sessions via ait.save / ait.load

Quick Start

1 import aitune.torch as ait
2 from aitune.torch.backend import TorchInductorAotBackend
3 from aitune.torch.tune_strategy import OneBackendStrategy
4 
5 backend = TorchInductorAotBackend()
6 strategy = OneBackendStrategy(backend)
7 
8 model = ait.Module(model, "my-model", strategy=strategy)
9 ait.tune(model, input_data, batch_sizes=[1, 2])
10 
11 # Persist the compiled artifact
12 ait.save(model, "my_model.ait")

Loading in a later session:

1 import aitune.torch as ait
2 from aitune.torch.backend import TorchInductorAotBackend
3 from aitune.torch.tune_strategy import OneBackendStrategy
4 
5 model = ait.Module(original_model, "my-model", strategy=OneBackendStrategy(TorchInductorAotBackend()))
6 ait.load(model, "my_model.ait")

Configuration Options

TorchInductorAotBackendConfig

1 @dataclass
2 class TorchInductorAotBackendConfig(BackendConfig):
3     inductor_configs: dict[str, Any] | None = None

inductor_configs

Pass any key from torch._inductor.config directly to the compiler:

1 from aitune.torch.backend import TorchInductorAotBackend, TorchInductorAotBackendConfig
2 
3 # Default — no extra inductor options
4 backend = TorchInductorAotBackend()
5 
6 # Enable max-autotune kernel search
7 config = TorchInductorAotBackendConfig(inductor_configs={"max_autotune": True})
8 backend = TorchInductorAotBackend(config=config)
9 
10 # Coordinate-descent kernel tuning
11 config = TorchInductorAotBackendConfig(
12     inductor_configs={
13         "max_autotune": True,
14         "coordinate_descent_tuning": True,
15     }
16 )
17 backend = TorchInductorAotBackend(config=config)

See all available keys:

1 import torch
2 torch._inductor.list_options()

Dynamic Shapes

By default, dynamic shapes are inferred automatically from the data samples passed to ait.tune.

Batch axis: detected when the same tensor dimension varies proportionally with batch_size
Spatial / sequence axes: detected when a dimension varies independently of batch size

1 # Two samples with different spatial sizes → H and W detected as dynamic
2 data_224 = torch.randn(3, 224, 224, device="cuda")
3 data_256 = torch.randn(3, 256, 256, device="cuda")
4 
5 # batch_sizes=[1] keeps samples separate (no cross-shape stacking)
6 ait.tune(module, [data_224, data_256], batch_sizes=[1])

When input samples have different spatial sizes, use batch_sizes=[1] to prevent the data loader from stacking tensors of mismatched shapes.

The backend uses torch.export.Dim.AUTO for spatial / sequence axes, letting PyTorch infer valid ranges and divisibility constraints from the model automatically.

When recorded samples do not cover the full production range, provide explicit bounded dimensions on ait.Module. See User-provided dynamic shapes and the runnable ResNet dynamic-shapes example.

Save and Load

1 # After tuning
2 ait.save(model, "model.ait")
3 
4 # In a new process — original weights still needed for Module wrapping,
5 # but inference runs entirely through the compiled runner
6 model = ait.Module(original_model, "my-model",
7                    strategy=OneBackendStrategy(TorchInductorAotBackend()))
8 ait.load(model, "model.ait")
9 result = model(input_tensor)

After ait.tune completes, the original module is offloaded to CPU to free GPU memory. The compiled .pt2 runner is fully self-contained for inference.

Comparison with Torch Inductor JIT Backend

Feature	AOT Backend	JIT Backend
Compilation time	At `tune()` call	At first inference call
Artifact persistence	Yes (`.pt2` file)	No
Python overhead	None at inference	Minimal (compiled graph)
Dynamic shapes	Inferred or explicitly configured	Configurable via `dynamic=`
Save / Load	Supported	Supported (recompiles)
Requires PyTorch	≥ 2.6	Any

Next Steps

Compare with Torch Inductor JIT Backend for JIT compilation
Compare with TensorRT Backend for maximum performance
Learn about Tune Strategies
Review Deployment Guide