PyTorch Inductor, in depth: lowering, the loop-level IR, fusion scheduling, Triton code generation, autotuning and AOT packaging

Why a fusing compiler pays off

Most operations in a modern network are not matmuls. Layer norms, activations, residual adds, dropout, softmax, casts and optimizer updates are elementwise or reduction operations that do very little arithmetic per byte. In eager PyTorch each one is a separate kernel: it reads its inputs from GPU memory (HBM), computes, and writes its output back, and the next kernel reads it again. For these operations the GPU spends its time moving bytes, not doing math, so the runtime is roughly total bytes moved divided by memory bandwidth.

Inductor attacks that cost directly. If five elementwise operations become one kernel, the intermediate tensors live in registers and never touch HBM. The arithmetic is unchanged; the traffic shrinks. Everything else in Inductor, from its IR to its scheduler, exists to find those opportunities safely and to emit kernels good enough that fusing them is a net win. The general principle is explained in kernel fusion; Inductor is a compiler that applies it automatically to whatever graph it is handed.

Decomposition and lowering to loop-level IR

Inductor receives an FX graph of ATen operations. Its first step is decomposition: many ATen operators are rewritten in terms of simpler ones, so aten.gelu or a batch-norm becomes arithmetic and reductions on primitives. This keeps the set of operations the backend must handle small, and it exposes the pieces of composite operators to fusion.

Next comes lowering. Each remaining operator is mapped to Inductor's loop-level IR. The important IR nodes are Pointwise, which describes an output element as a Python function of an index, and Reduction, which describes a reduction over one or more dimensions. These nodes do not hold data; they hold an inner_fn that says how to compute one element from loads of other buffers. Because the description is a function, chaining two pointwise operations simply composes their functions, and nothing is materialised.

Something must eventually be written to memory. Inductor decides which values to realize into a Buffer: graph outputs, values needed by the backward pass, inputs to operations that need a real tensor such as an external matmul, and values read so many times that recomputing them would cost more than storing them. Operations Inductor has no lowering for fall back to the eager ATen kernel through a FallbackKernel, which is correct but breaks fusion around it.

The scheduler and fusion

After lowering, the graph is a list of buffers, each with a recorded set of memory reads and writes expressed as index formulas. The scheduler wraps each buffer in a node, builds a dependency graph from those reads and writes, and then greedily merges nodes into fused groups.

Two kinds of merge matter. Vertical fusion joins a producer with its consumer, so the intermediate never leaves registers: a bias add feeding a GELU feeding a residual add becomes one kernel. Horizontal fusion joins independent nodes that read the same input or iterate over the same shape, so one pass over memory serves several outputs, as when the mean and variance of a normalisation are computed together. A reduction can also absorb pointwise work before it, and pointwise work after it can be fused as an epilogue when the shapes line up.

Candidate pairs are scored mainly by the memory traffic the merge would save, and merges that would create a cycle in the dependency graph, mismatch iteration sizes, or exceed a size limit are rejected. The heuristics are tuned for common patterns and change between releases, so the reliable way to know what fused is to read the output rather than to predict it. Fusion is not free: a bigger kernel uses more registers, can lower occupancy, and can recompute a value in several places. Inductor usually wins that trade for memory-bound code; for some reductions with awkward layouts it does not, and that shows up in profiles as a fused kernel slower than the eager pair.

Generated Triton and C++ kernels

Each fused group is handed to a code generator. On NVIDIA and AMD GPUs this is Triton: Inductor writes Python source decorated with @triton.jit, and Triton compiles it to machine code. Kernel names encode their shape: triton_poi_ for pointwise, triton_red_ for a looped reduction, triton_per_ for a persistent reduction where the whole reduced row fits in one block, and triton_tem_ for template kernels such as a GEMM. The rest of the name lists the fused operators. Below is an abridged pointwise kernel of the kind Inductor emits for a bias add, GELU and residual add; exact output varies by PyTorch version.

@triton.jit
def triton_poi_fused_add_gelu_0(in_out_ptr0, in_ptr0, in_ptr1, xnumel, XBLOCK: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 4096                      # column index, used for the bias
    tmp0 = tl.load(in_out_ptr0 + (x2), xmask).to(tl.float32)    # matmul output
    tmp1 = tl.load(in_ptr0 + (x0), xmask).to(tl.float32)         # bias
    tmp2 = tmp0 + tmp1
    tmp3 = 0.5 * tmp2 * (1.0 + libdevice.erf(tmp2 * 0.7071067811865476))
    tmp4 = tl.load(in_ptr1 + (x2), xmask).to(tl.float32)         # residual
    tmp5 = tmp3 + tmp4
    tl.store(in_out_ptr0 + (x2), tmp5, xmask)

Notice what the IR became: index arithmetic (x0 = xindex % 4096 is the broadcast of the bias), loads with a mask for the ragged final block, fp32 upcasting for the math, and a single store. XBLOCK is a compile-time constant chosen by a small set of heuristic or autotuned configurations. Reductions add an RBLOCK for the reduced dimension and a loop over it unless the row is small enough to be persistent.

On CPU the same scheduled groups become C++ with OpenMP parallel loops and explicit vectorisation, compiled by the system compiler. The C++ path matters for inference on servers without GPUs and for debugging, since the same IR drives both backends.

Matmuls, templates and autotuning

Matmuls and convolutions are treated differently. By default Inductor calls the vendor library (cuBLAS or cuDNN on NVIDIA) as an extern kernel, because those libraries are very hard to beat. The cost is that an extern kernel is a fusion barrier: its output must be realised, and the pointwise work after it becomes a separate kernel that reads it back, as in the example above.

With mode="max-autotune" (or torch._inductor.config.max_autotune = True), Inductor also generates Triton GEMM templates for each matmul, benchmarks a set of tile configurations alongside the library call on your actual shapes, and keeps the fastest. A template can carry a fused epilogue, so a matmul plus bias plus activation may become one kernel. Compile time grows substantially because every candidate is compiled and timed, so this mode suits long-running training jobs and stable serving shapes, not notebooks. A related knob, coordinate_descent_tuning, searches block sizes for pointwise and reduction kernels as well. Results are cached on disk, so the second run of the same shapes is far cheaper.

The wrapper, CUDA graphs and AOTInductor

Kernels alone are not a program. Inductor also emits wrapper code that allocates output buffers, reuses freed buffers when lifetimes allow, launches kernels in order and calls extern kernels. In the default Python wrapper each launch still pays Python and driver overhead, which is invisible for large batches and dominant for small ones such as single-token LLM decoding.

That is what mode="reduce-overhead" addresses: it records the wrapper's launches into CUDA graphs, using a mechanism PyTorch calls CUDA graph trees, and replays them with one launch. The price is that graph replay needs stable memory addresses, so inputs are copied into static buffers, outputs may be overwritten by the next replay unless you clone them, and dynamic shapes produce one graph per distinct shape. For deployment without Python, AOTInductor compiles a torch.export program ahead of time into a package that a C++ or Python runtime loads:

import torch

ep = torch.export.export(model.eval(), (example_input,))
path = torch._inductor.aoti_compile_and_package(ep, package_path="model.pt2")

runner = torch._inductor.aoti_load_package(path)   # also loadable from C++
out = runner(example_input)

Worked example: counting bytes saved

Take the transformer MLP tail y = gelu(x @ W + b) + r with x @ W producing an 8192 by 4096 fp16 activation. One such tensor is 8192 x 4096 x 2 bytes = 64 MiB. Count HBM traffic after the matmul, ignoring the 8 KiB bias.

Eager moves 448 MiB; the fused kernel moves 192 MiB, a 2.3x reduction. On a GPU sustaining around 2 TB/s of effective bandwidth that is roughly 235 microseconds against 100, per layer, per step, before counting the two extra kernel launches eager pays. With max-autotune and an epilogue-fused template the 64 MiB write and re-read of the matmul output disappear too. Do the same arithmetic for your own hot paths before switching modes: it tells you the best case, and the profiler tells you how close you got.

Reading what Inductor generated

Inductor is unusually inspectable, and reading its output is the fastest way to debug performance. The useful switches are:

• TORCH_LOGS=output_code prints every generated kernel and the wrapper, with the source line each fused operator came from.
• TORCH_COMPILE_DEBUG=1 writes a debug directory per compiled graph containing the readable FX graph, the IR before and after fusion, and the final output code, so you can see which nodes the scheduler merged and which it refused.
• TORCH_TRACE=/some/dir plus the tlparse tool produces an HTML report of the whole compile, including graph breaks, recompiles and per-graph compile time.
• The PyTorch profiler shows generated kernels by name, so triton_red_fused_... in a trace maps straight back to the output code. See reading GPU traces.

A practical loop: profile, take the top five Inductor kernels by time, find each in the output code, and ask whether it is near the memory bound computed as in the worked example. A kernel far below the bound usually has a bad layout (non-contiguous reads), a reduction over a strided dimension, or a fallback op splitting what should have been one kernel.

Failure modes

• Fallback ops break fusion. An operator with no lowering runs eagerly and forces its inputs and outputs into memory. Look for direct ATen calls or extern_kernels. in the output code around slow regions, or for FallbackKernel in the IR dumps that TORCH_COMPILE_DEBUG=1 writes.
• Numerics differ slightly from eager. Fused kernels compute in fp32 and round once, and reductions may sum in a different order. Compare with tolerances appropriate to the dtype, never with exact equality.
• Compile time explodes under max-autotune or many shapes. Each new shape can trigger new kernels and new benchmarks. Mark dynamic dimensions, pad to buckets, and keep the cache warm across runs.
• CUDA graph outputs are overwritten. Under reduce-overhead, holding a reference to an output across steps can give you the next step's values. Clone outputs you keep.
• Memory grows instead of shrinking. Recomputation and CUDA graph static pools can raise peak memory; measure it, do not assume fusion saves it.
• A fused kernel is slower than eager. Usually register pressure or an awkward reduction. Confirm in the profiler, then report it with a minimal repro rather than disabling compilation globally.

Trade-offs

What to do next

1. Compile one model with default settings and record step time, peak memory and compile time.
2. Run with TORCH_LOGS=output_code once and find your three hottest generated kernels.
3. For each, compute bytes moved and compare with your GPU's bandwidth to see how close it is to the bound.
4. Search the output code for fallback and extern calls inside hot regions, and rewrite or report them.
5. Try reduce-overhead if steps are short and launch-bound; clone any outputs you keep.
6. Try max-autotune for long jobs with fixed shapes, and persist the cache directory between runs.
7. Add a numerics check against eager with dtype-appropriate tolerances to your CI.
8. For serving without Python, prototype an AOTInductor package and benchmark it against the JIT path.