Gunjan Dhanuka — Learning Notes

torch.compile (TorchDynamo / TorchInductor)

Properties1
tagsgpu, torch-compile, torchinductor, torchdynamo, aotautograd, triton, kernel-fusion, codegen, llm-systems

3 min read

torch.compile (TorchDynamo / TorchInductor)

Definition

torch.compile (PyTorch 2.x) is a just-in-time graph compiler that traces eager PyTorch, fuses memory-bound op chains, and generates Triton kernels automatically — no kernel authoring. Its default backend is TorchInductor. Most Triton kernels running in production LLM code today were not written by a human; they were emitted by Inductor. It is the first optimization to reach for: a one-line @torch.compile that captures most of the fusion wins you would otherwise hand-write.

How it works

The pipeline has three stages:

  1. TorchDynamo hooks CPython frame evaluation to safely capture an FX graph of your eager code, falling back to eager (“graph break”) on anything it can’t trace — so it’s always correct.
  2. AOTAutograd traces an ahead-of-time backward graph, and PrimTorch canonicalizes ~2000+ ATen ops down to ~250 primitives (a small, uniform target for codegen).
  3. TorchInductor lowers that graph and, for GPU, emits Triton kernels — fusing pointwise + reduction chains, choosing persistent vs. non-persistent reductions, and picking tile sizes / autotuning automatically. (On CPU it emits OpenMP C++.)

The headline capability is fusion. Inductor spots memory-bound op chains (bias → GELU → dropout, RMSNorm, softmax) and fuses each into one Triton kernel that reads inputs once, computes in registers, and writes once — the same DRAM-traffic win you engineer by hand, derived automatically from torch.softmax. For large GEMMs it is selective: it typically defers to cuBLAS/cuDNN, and under max-autotune benchmarks cuBLAS vs. CUTLASS vs. a Triton template and picks the winner, optionally fusing the pointwise epilogue in. Inspect the generated Triton with TORCH_LOGS="output_code".

Limits. Inductor fuses what’s adjacent in the graph; it generally will not discover a fundamentally new algorithm like FlashAttention’s online softmax — that requires IO-aware restructuring the graph doesn’t express, which is why attention is dispatched to a hand-written fused kernel. The relationship to hand-written Triton is layered, not competitive: you can drop a user-defined @triton.jit kernel into a compiled region and Inductor treats it as an opaque node, fusing around it. Reach for hand-Triton only for the residual hotspots Inductor can’t express or beats poorly (novel attention variants, fused MoE, quantized matmuls, fused-linear-cross-entropy).

Why it matters

torch.compile handles the breadth of a model’s memory-bound glue with zero authoring effort and correctness guaranteed by the tracer, freeing engineers to hand-write only the handful of kernels that define the efficiency frontier. It is the default for production LLM training/inference: measure eager, then measure compiled, before writing any kernel.

Taught in

See also

Graph View

Backlinks