🏭 AOTInductor: Ahead-of-Time Compilation for Deployment

Table of Contents

• 1. Motivation: JIT vs. AOT Compilation
  • 1.1 The JIT torch.compile Deployment Problem
  • 1.2 What AOTInductor Provides
• 2. Architecture and Position in the Stack
  • 2.1 Where AOTInductor Sits
  • 2.2 The Compilation Pipeline
• 3. The Export Step: torch.export
  • 3.1 ExportedProgram
  • 3.2 What export Guarantees
  • 3.3 Constraints and Dynamic Shapes
• 4. Inductor AOT Lowering
  • 4.1 From ExportedProgram to Inductor IR
  • 4.2 Kernel Compilation
  • 4.3 The C++ Wrapper
• 5. The Shared Library ABI
  • 5.1 .so Structure
  • 5.2 The C ABI Contract
  • 5.3 AOTIModelContainerRunner
• 6. Dynamic Shapes in AOTInductor
  • 6.1 Dim Objects and Constraints
  • 6.2 Shape Baking vs. Dynamic Dispatch
  • 6.3 Runtime Shape Checks
• 7. Python and C++ Runtime APIs
  • 7.1 aot_compile and aoti_load_package
  • 7.2 aoti_compile_and_package
  • 7.3 C++ Inference without Python
• 8. CUDA Graphs Integration
• 9. Limitations and Tradeoffs
  • 9.1 vs. JIT torch.compile
  • 9.2 Unsupported Patterns
• References

1. Motivation: JIT vs. AOT Compilation

1.1 The JIT torch.compile Deployment Problem

torch.compile is a just-in-time compiler: the first call to a compiled function triggers Dynamo tracing, AOTAutograd joint tracing, and Inductor kernel generation. This incurs three costs that are acceptable in training but problematic in production inference:

Beyond latency, many production serving stacks use C++ inference servers (TorchServe backend, Triton Inference Server, or custom gRPC services) that load models as shared libraries. Shipping a Python process per replica is expensive.

1.2 What AOTInductor Provides

AOTInductor (ahead-of-time Inductor) compiles a model before deployment into a standalone shared library (.so on Linux, .dll on Windows) that:

• Contains all compiled Triton/CUBIN kernels as embedded binary data
• Exposes a pure C ABI — callable from C++ with only libtorch (no Python)
• Requires no recompilation at runtime — no guards, no Dynamo, no Python overhead
• Supports deterministic latency from the first call

The tradeoff: shape flexibility must be declared upfront (via torch.export.Dim), and any input that violates the declared constraints is undefined behavior (not a graceful recompilation).

2. Architecture and Position in the Stack

2.1 Where AOTInductor Sits

AOTInductor replaces the JIT half of torch.compile with an offline compilation step:

flowchart TD
    subgraph JIT["JIT Path (torch.compile)"]
        J1["Python model\nfirst call"]
        J2["TorchDynamo\nbytecode interception"]
        J3["AOTAutograd\njoint tracing"]
        J4["TorchInductor JIT\nTriton kernels compiled at runtime"]
        J5["Execution (every call):\ncheck guards → run kernels"]
        J1 --> J2 --> J3 --> J4 --> J5
    end

    subgraph AOT["AOT Path (AOTInductor)"]
        A1["torch.export.export()\noffline, before deployment"]
        A2["ExportedProgram\npure FX graph + constraints"]
        A3["TorchInductor AOT\ngenerate Triton + C++ wrapper"]
        A4["model.so\nself-contained shared library"]
        A5["Production inference:\nload .so, call C ABI"]
        A1 --> A2 --> A3 --> A4 --> A5
    end

2.2 The Compilation Pipeline

flowchart LR
    EP["ExportedProgram
(ATen FX graph +
shape constraints)"]
    DEC["Decomposition
ATen → prim ops
(same as JIT path)"]
    LR["Inductor Lowering
prim FX → Inductor IR
Pointwise / Reduction / ExternKernel"]
    SCH["Scheduler
fusion planning"]
    TG["Triton Codegen
GPU: .cubin embedded
CPU: .cpp compiled"]
    CW["C++ Wrapper\ngeneration
aoti_torch_* C ABI calls"]
    LIB["Link: .so / .dll
(kernels + wrapper)"]

    EP --> DEC --> LR --> SCH --> TG --> CW --> LIB

The decomposition and Inductor lowering stages are identical to the JIT path — AOTInductor is simply Inductor invoked in a mode that writes artifacts to disk rather than executing them inline.

3. The Export Step: torch.export

3.1 ExportedProgram

torch.export.export(model, args, kwargs, dynamic_shapes) traces the model through a strict version of TorchDynamo with fullgraph=True and produces an ExportedProgram:

import torch
from torch.export import export, Dim

batch = Dim("batch", min=1, max=512)
ep = export(
    model,
    args=(x,),
    dynamic_shapes={"x": {0: batch}},
)

ExportedProgram contains: - graph_module — a torch.fx.GraphModule with ATen-level ops; no Python control flow - graph_signature — maps placeholder nodes to inputs/buffers/parameters - range_constraints — {symbol: ValueRanges(min, max)} for each dynamic dim - equality_constraints — [(InputDim(arg, i), InputDim(arg, j))] asserting dims are equal across inputs

3.2 What export Guarantees

torch.export enforces a stronger invariant than torch.compile:

3.3 Constraints and Dynamic Shapes

The dynamic_shapes argument maps each input argument (by name or position) to a dict from dimension index to a Dim object:

# Two inputs with a shared dynamic batch dimension
batch = Dim("batch", min=1, max=1024)
seq   = Dim("seq",   min=1, max=2048)

ep = export(
    model,
    args=(x, y),
    dynamic_shapes={
        "x": {0: batch, 1: seq},
        "y": {0: batch},          # same Dim object → equality constraint
    },
)

Sharing the same Dim object across inputs asserts that those dimensions are equal at runtime. Inductor can use this to avoid generating shape-check code and to produce tighter loop bounds.

Dimensions not mentioned in dynamic_shapes are static — their value is baked into the compiled kernel. Attempting to call the compiled artifact with a different value for a static dimension is undefined behavior.

4. Inductor AOT Lowering

4.1 From ExportedProgram to Inductor IR

AOTInductor calls torch._inductor.aot_compile(exported_program, inputs, options), which:

1. Decomposes ATen ops to the prim level using the same decomposition table as the JIT path
2. Lowers each prim op to an Inductor IR node via the lowerings registry (identical @register_lowering dispatch)
3. Runs the Scheduler — fusion planning, Δ-traffic scoring, same algorithm as JIT

The key difference is that all symbolic shapes come from the ExportedProgram’s range constraints rather than from Dynamo’s ShapeEnv. At lowering time, each dynamic dimension is a sympy.Symbol with known bounds; static dimensions are replaced by their concrete integer values.

4.2 Kernel Compilation

For GPU targets, each scheduled kernel is: 1. Generated as Triton Python source (same triton_heuristics / @triton.jit as JIT) 2. Compiled with triton.compile() to PTX then to CUBIN (.cubin binary) 3. The CUBIN is embedded in the output .so as a __attribute__((section(".rodata"))) byte array

For CPU targets, the generated C++ source is compiled with the system compiler (GCC/Clang) and linked into the .so.

Because compilation happens offline, AOTInductor can afford a more exhaustive autotuning sweep than JIT (where tuning latency would show up on the first inference request).

4.3 The C++ Wrapper

Inductor generates a C++ file (model.cpp) containing:

// Auto-generated by AOTInductor
#include <torch/csrc/inductor/aoti_runtime/interface.h>

extern "C" {

AOTIRuntimeError AOTInductorModelContainerCreate(
    AOTInductorModelContainerHandle* container_handle,
    size_t num_models,
    bool is_cpu,
    const char* cubin_dir) {
    // Initializes CUDA context, loads CUBINs
    *container_handle = new AOTInductorModelContainer(num_models, is_cpu, cubin_dir);
    return AOTI_RUNTIME_SUCCESS;
}

AOTIRuntimeError AOTInductorModelContainerRun(
    AOTInductorModelContainerHandle container_handle,
    AtenTensorHandle* input_handles,
    size_t num_inputs,
    AtenTensorHandle* output_handles,
    size_t num_outputs,
    AOTInductorStreamHandle stream_handle,
    AOTInductorProxyExecutorHandle proxy_executor_handle) {
    // Validates shapes, runs kernel sequence
    auto* model = get_model(container_handle);
    model->run(input_handles, output_handles, (cudaStream_t)stream_handle);
    return AOTI_RUNTIME_SUCCESS;
}

} // extern "C"

The AtenTensorHandle is an opaque pointer to a at::Tensor — the caller allocates tensors with torch::empty(...) or equivalent and passes pointers. The .so never calls back into Python.

5. The Shared Library ABI

5.1 .so Structure

The compiled .so contains:

The .so links against libtorch_cuda.so and libcuda.so but not against libpython. A serving process only needs libtorch and the CUDA runtime.

5.2 The C ABI Contract

The C ABI is defined in torch/csrc/inductor/aoti_runtime/interface.h:

typedef void* AOTInductorModelContainerHandle;
typedef void* AtenTensorHandle;
typedef void* AOTInductorStreamHandle;

AOTIRuntimeError AOTInductorModelContainerCreate(
    AOTInductorModelContainerHandle* container_handle,
    size_t num_models,
    bool is_cpu,
    const char* cubin_dir);

AOTIRuntimeError AOTInductorModelContainerRun(
    AOTInductorModelContainerHandle container_handle,
    AtenTensorHandle* input_handles,
    size_t num_inputs,
    AtenTensorHandle* output_handles,
    size_t num_outputs,
    AOTInductorStreamHandle stream_handle,
    AOTInductorProxyExecutorHandle proxy_executor_handle);

AOTIRuntimeError AOTInductorModelContainerDelete(
    AOTInductorModelContainerHandle container_handle);

num_models enables a multi-runner container — the container holds multiple independent copies of the model’s state (weights, scratch buffers), enabling concurrent inference requests on different CUDA streams without serialization.

5.3 AOTIModelContainerRunner

PyTorch provides a C++ helper class that wraps the raw ABI:

#include <torch/csrc/inductor/aoti_runner/model_container_runner_cuda.h>

// Load
torch::inductor::AOTIModelContainerRunnerCuda runner(
    "path/to/model.so",
    /*num_runners=*/4,           // concurrent request capacity
    /*device_str=*/"cuda",
    /*cubin_dir=*/""             // empty = embedded CUBINs
);

// Inference
std::vector<at::Tensor> inputs = {input_tensor};
std::vector<at::Tensor> outputs = runner.run(inputs);

runner.run() is thread-safe when num_runners > 1: it acquires one of the pre-allocated model slots from a fixed pool, runs the inference, and releases the slot. Requests that arrive when all slots are busy block (or can be queued by the caller).

6. Dynamic Shapes in AOTInductor

6.1 Dim Objects and Constraints

torch.export.Dim objects declare symbolic dimensions with optional bounds:

from torch.export import Dim

# Unbounded (Inductor will not specialize on this value)
batch = Dim("batch")

# Bounded — Inductor can use bounds for range checks and tile alignment
batch = Dim("batch", min=1, max=512)

# Derived — seq must be a multiple of 8 (useful for attention heads)
seq = Dim("seq", min=8, max=2048)

Under the hood, each Dim becomes a sympy.Symbol in ShapeEnv with the declared bounds stored as var_to_range. Inductor’s scheduler uses these bounds to decide whether symbolic guard expressions can be simplified (e.g., if batch >= 1 is guaranteed, no runtime check is needed for a loop that iterates batch times).

6.2 Shape Baking vs. Dynamic Dispatch

When Inductor lowers an op with a symbolic dimension s:

• Triton constexpr arguments — static dims become tl.constexpr parameters, enabling compile-time tile selection. Dynamic dims become regular integer arguments passed at each kernel launch.
• Index expressions — tl.arange(0, BLOCK_SIZE) uses BLOCK_SIZE as constexpr; offsets involving s become pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) with s as a runtime kernel argument.
• CUBIN portability — a single CUBIN can handle all values of s within [min, max] because s appears only in arithmetic, not in tiling decisions (which are constexpr).

6.3 Runtime Shape Checks

For dynamic dimensions declared with bounds, the C++ wrapper validates bounds at each call:

// Auto-generated shape validation code in model.cpp
void validate_inputs(AtenTensorHandle* inputs) {
    int64_t batch = inputs[0]->size(0);
    AOTI_TORCH_CHECK(batch >= 1 && batch <= 512,
        "Input dim 0 (batch) out of range: ", batch);
    // static dims are baked in and not checked
}

Violations raise an error (not a recompilation). There are no guards, no fallback — the caller is responsible for respecting the declared constraints.

7. Python and C++ Runtime APIs

7.1 aot_compile and aoti_load_package

Compilation (Python, offline):

import torch
from torch.export import export, Dim

# Step 1: export
batch = Dim("batch", min=1, max=64)
ep = export(model, (x,), dynamic_shapes={"x": {0: batch}})

# Step 2: AOT compile → .so
so_path = torch._inductor.aot_compile(
    ep,
    (x,),
    options={
        "max_autotune": True,      # exhaustive kernel autotuning
        "cuda_graphs": True,       # embed CUDA graph capture in .so
    },
)
print(so_path)  # e.g., /tmp/aot_inductor/model.so

Inference (Python, online):

# Load the compiled artifact back into Python
compiled_model = torch._inductor.aoti_load_package(so_path)

# Call like a normal callable — returns tensors
output = compiled_model(x_new)

aoti_load_package wraps AOTIModelContainerRunnerCuda in a Python callable, handling tensor handle conversion transparently.

7.2 aoti_compile_and_package

For distribution, aoti_compile_and_package creates a single .pt2 archive containing the .so and metadata:

package_path = torch._inductor.aoti_compile_and_package(
    ep,
    args=(x,),
    package_path="model.pt2",
    options={"max_autotune": True},
)

# Load from the package
compiled_model = torch._inductor.aoti_load_package("model.pt2")

The .pt2 format allows shipping a single file, version-checking the PyTorch ABI, and including multiple device variants (CPU + CUDA) in one bundle.

7.3 C++ Inference without Python

A production C++ server using only libtorch:

#include <torch/torch.h>
#include <torch/csrc/inductor/aoti_runner/model_container_runner_cuda.h>

int main() {
    // Load compiled model
    torch::inductor::AOTIModelContainerRunnerCuda runner("model.so", /*num_runners=*/2);

    // Prepare input (must match shape/dtype constraints)
    auto input = torch::randn({32, 512}, torch::TensorOptions()
                                            .dtype(torch::kFloat32)
                                            .device(torch::kCUDA));

    // Run inference (thread-safe for concurrent calls)
    auto outputs = runner.run({input});
    std::cout << outputs[0].sizes() << std::endl;

    return 0;
}

Compile against libtorch:

g++ -std=c++17 inference.cpp \
    -I$(python -c "import torch; print(torch.utils.cmake_prefix_path)")/../../include \
    -L$(python -c "import torch; print(torch.utils.cmake_prefix_path)")/../../lib \
    -ltorch -ltorch_cuda -lc10_cuda \
    -o inference

batch = Dim("batch", min=1, max=32)
seq   = Dim("seq",   min=1, max=2048)
ep = export(model, (x,), dynamic_shapes={"x": {0: batch, 1: seq}})
# hidden_dim=8192 is NOT in dynamic_shapes → baked in statically
so_path = torch._inductor.aot_compile(ep, (x,), options={"max_autotune": True})

8. CUDA Graphs Integration

CUDA graphs capture a sequence of GPU operations into a single replayable DAG, eliminating CPU-side kernel launch overhead (typically 5–20 μs per kernel) and enabling GPU-side scheduling optimizations.

AOTInductor can embed CUDA graph capture and replay inside the .so:

so_path = torch._inductor.aot_compile(
    ep, (x,),
    options={"cuda_graphs": True}
)

When cuda_graphs=True: 1. On the first run, the wrapper executes all kernels normally and records them into a cudaGraph_t 2. On subsequent runs (same shapes), it calls cudaGraphLaunch(graph, stream) — a single CPU call replaces \(N\) individual kernel launches 3. If a run has different shapes (within the dynamic range), CUDA graphs are bypassed and kernels run normally

The option max-autotune-no-cudagraphs sometimes outperforms max-autotune (which includes CUDA graphs) for models where the copy overhead exceeds the launch overhead savings. This is model-dependent and should be benchmarked.

9. Limitations and Tradeoffs

9.1 vs. JIT torch.compile

9.2 Unsupported Patterns

AOTInductor inherits torch.export’s restrictions plus its own:

References