torch.compile for ML Performance Engineers

Concept-level, interview-deployable. Version anchor: PyTorch 2.10–2.12 era (2.10 shipped Jan 2026; two-month cadence since). Material changes in the last ~2 years are flagged inline: automatic dynamic shapes maturing (2.1→), compile caching becoming production-real (FXGraph/AOTAutograd/PGO caches + portable save_cache_artifacts, 2.6→), regional compilation, and vLLM-style piecewise compilation becoming the serving default. Everything deeper than the marked boundaries is compiler-team territory — knowing where that line is, is itself the staff signal.

1 · The Stack in One Diagram

One sentence per stage. Dynamo intercepts CPython bytecode execution, symbolically evaluates it, and emits an FX graph of tensor ops plus guards — runtime predicates under which that graph is valid. AOTAutograd traces the joint forward+backward, removes mutation (functionalization), and lowers to a normalized ATen op set. The backend turns ATen into machine-executable code: Inductor schedules/fuses and emits Triton (GPU) or C++/OpenMP (CPU); the openxla backend exports to StableHLO and hands the whole problem to XLA, where fusion/layout decisions belong to XLA, not PyTorch.

1.5 · What Each Stage Actually Looks Like

One function traced through every artifact — this is what you see on a real screen, and quoting any one of these from memory is worth a paragraph of description:

def ffn_out(x, w, b):          # [B,4096] @ [4096,4096] + bias, relu
    return torch.relu(x @ w + b)
cf = torch.compile(ffn_out, fullgraph=True)

Stage 0 — what Dynamo hooks (CPython bytecode, dis.dis(ffn_out)):

  2   LOAD_FAST     x        LOAD_FAST     w
      BINARY_OP     @ (MATMUL)
      LOAD_FAST     b        BINARY_OP     + (ADD)
      LOAD_GLOBAL   torch    LOAD_ATTR     relu
      CALL          1        RETURN_VALUE

Stage 1 — Dynamo's FX graph (TORCH_LOGS=graph_code):

def forward(self, L_x_: "bf16[8, 4096]", L_w_: "bf16[4096, 4096]", L_b_: "bf16[4096]"):
    matmul: "bf16[8, 4096]" = L_x_ @ L_w_
    add:    "bf16[8, 4096]" = matmul + L_b_;  matmul = None
    relu:   "bf16[8, 4096]" = torch.relu(add);  add = None
    return (relu,)

…and the guards that make it valid (TORCH_LOGS=guards, excerpt):

TENSOR_MATCH: check_tensor(L['x'], Tensor, DispatchKeySet(CUDA, BFloat16),
              torch.bfloat16, device=0, requires_grad=False,
              size=[8, 4096], stride=[4096, 1])     # ← batch dim 8 is SPECIALIZED (first call)
TENSOR_MATCH: check_tensor(L['w'], …, size=[4096, 4096], …)
ID_MATCH:     ___check_obj_id(G['torch'], 0x7f2…)   # torch itself hasn't been monkeypatched

Second call with batch 16 → this guard fails → recompile with size=[s0, 4096]: that size going symbolic in the log IS automatic dynamic shapes happening.

Stage 2 — AOTAutograd's ATen graph (TORCH_LOGS=aot_graphs, inference path):

def forward(self, arg0_1: "bf16[8, 4096]", arg1_1: "bf16[4096, 4096]", arg2_1: "bf16[4096]"):
    mm:      "bf16[8, 4096]" = torch.ops.aten.mm.default(arg0_1, arg1_1)
    add:     "bf16[8, 4096]" = torch.ops.aten.add.Tensor(mm, arg2_1)
    relu:    "bf16[8, 4096]" = torch.ops.aten.relu.default(add)
    return (relu,)

Everything is now namespaced aten.*, functional, and (in training) would carry the backward: threshold_backward + two transposed mms in the same joint graph.

Stage 3a — GPU: Inductor's generated Triton (TORCH_COMPILE_DEBUG=1 → output_code.py, abbreviated):

@triton.jit
def triton_poi_fused_add_relu_0(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    x0 = xindex % 4096
    tmp0 = tl.load(in_ptr0 + xindex)          # mm result
    tmp1 = tl.load(in_ptr1 + x0)              # bias (broadcast)
    tmp2 = tmp0 + tmp1
    tmp3 = triton_helpers.maximum(0, tmp2)    # relu — fused into the SAME kernel
    tl.store(out_ptr0 + xindex, tmp3)
# call site: extern_kernels.mm(arg0_1, arg1_1, out=buf0)   ← GEMM stays cuBLAS
#            triton_poi_fused_add_relu_0.run(buf0, arg2_1, buf1, 32768)

The reading that matters: kernel name says what fused (fused_add_relu — one kernel, one HBM round-trip for both ops) and the GEMM went to cuBLAS, not Triton. If you saw fused_add_0 and a separate fused_relu_1, fusion missed — that grep is the whole §5 inspection method.

Stage 3b — TPU: the StableHLO/XLA artifact (XLA_FLAGS=--xla_dump_to=…, post-optimization excerpt):

HloModule ffn_out, entry_computation_layout={(bf16[8,4096]{1,0}, bf16[4096,4096]{1,0}, bf16[4096]{0})->bf16[8,4096]{1,0}}

%fused_computation (p0: bf16[8,4096], p1: bf16[4096]) -> bf16[8,4096] {
  %broadcast = bf16[8,4096] broadcast(bf16[4096] %p1), dimensions={1}
  %add = bf16[8,4096] add(bf16[8,4096] %p0, %broadcast)
  ROOT %maximum = bf16[8,4096] maximum(%add, bf16[] constant(0))
}
ENTRY %main {
  %dot = bf16[8,4096] dot(%x, %w), lhs_contracting_dims={1}, rhs_contracting_dims={0}
  ROOT %fusion = bf16[8,4096] fusion(%dot, %b), kind=kLoop, calls=%fused_computation
}

Same story, XLA's dialect: dot stays a dot, add+maximum became a fusion(kind=kLoop). This HLO text is exactly the artifact class a TPU perf pipeline diffs between compiler versions — being able to read it is the TPU-side equivalent of reading output_code.py.

2 · Dynamo: Bytecode Capture, Guards, Graph Breaks

Why bytecode-level capture. torch.jit.trace records the ops one concrete execution ran — it silently bakes in the branch taken and the shapes seen; control flow is gone. torch.jit.script re-parses Python into TorchScript and dies on the Python it can't express. JAX traces pure functions of abstract values and simply requires you to write traceable code. Dynamo takes the fourth path: hook CPython's frame evaluation (PEP 523), symbolically execute the bytecode, extract the tensor-op subgraph, and compile as much as it can prove safe — falling back to real Python for the rest. That's the design contract: arbitrary Python runs correctly; the compiled fraction determines the speedup.

Guards, conceptually. A captured graph is a claim: "this graph is what your code does, provided these predicates hold" — tensor dtypes/devices/shape-classes, branch-relevant Python values, module identity. Guards are checked on every entry to the compiled function; pass → run compiled artifact, fail → recompile (new specialization) or fall back. Guards are the price of capturing an untyped dynamic language soundly: JAX doesn't need them because it forbids the things guards check for.

Graph breaks. When Dynamo hits something it can't soundly trace — data-dependent control flow on tensor values (if x.max() > 0:), calls into opaque code (NumPy on tensor data, arbitrary C extensions, print(tensor)), side effects it can't functionalize — it ends the current graph, runs the offending code in eager Python, and starts a new graph after. The cost is structural, not just overhead: each fragment is a separate compile unit, so fusion cannot cross the break, cudagraph capture is defeated (the eager region re-introduces per-launch CPU work), and guard checks multiply per fragment. One break in a decode inner loop can cost more than Inductor's fusion gained.

A realistic break and its fix (the classic: a Python-value decision on tensor data in the hot path):

# BREAKS: .item() pulls a value to Python → data-dependent branch → graph break every step
max_len = (input_ids != pad_id).sum(-1).max().item()
if max_len < 512:
    logits = model(input_ids[:, :max_len])

# FIXED: keep the decision out of the compiled region — bucket outside, compile per bucket
bucket = select_bucket(seq_lens)          # plain Python, outside compiled fn
logits = compiled_fns[bucket](input_ids)  # branch-free tensor code inside

Finding breaks: torch._dynamo.explain(fn)(args) reports every break with its reason and source line; TORCH_LOGS=graph_breaks,recompiles streams them in production runs; torch.compile(fullgraph=True) turns any break into a hard error — the correct setting for serving code, where a silent break is a silent performance cliff.

3 · Dynamic Shapes — and the TPU Collision

The mechanism. First compile specializes on the exact shapes seen. When a dimension changes on a later call, Dynamo (since 2.1, "automatic dynamic") recompiles that dimension as symbolic (a SymInt) rather than concrete — the graph then serves a family of shapes, with guards reduced to range/divisibility constraints. Dimensions get re-specialized anyway when the code forces it: a value used in Python control flow, a reshape whose arithmetic needs a concrete value, or an op whose lowering wants divisibility (e.g., a kernel that requires s % 128 == 0 will guard on it). torch._dynamo.mark_dynamic(t, dim) declares intent up front and skips the specialize-first-then-generalize dance; the PGO cache (2.6+) persists these dynamism decisions across restarts so the dance isn't repeated per process.

Recompilation storms. Symptom: throughput craters intermittently while the process "works"; every new (batch, seq-len) signature pays seconds-to-minutes of compile. Diagnosis: TORCH_LOGS=recompiles prints each recompile with the guard that failed — the single most useful log line in the stack; torch._dynamo.config.cache_size_limit (default 8 entries per code object) is the tripwire: exceeding it silently falls back to eager, converting a compile-time problem into a permanent 2× eager tax. Fix order: mark the honest dynamic dims → bucket the rest → raise the cache limit only when the shape set is finite and known.

The serving consequence. Production LLM serving does not let shapes float: engines enumerate a bucket set (batch × padded-length, plus cudagraph capture per bucket — §7), compile all of it at warm-up, and route requests by padding up. This is the same reasoning as the geometry discipline in kernel design: a bounded, known shape set converts an unbounded compile problem into a fixed warm-up cost.

4 · AOTAutograd in Two Paragraphs

Training can't be compiled from the forward graph alone: autograd normally builds the backward dynamically at runtime, which would leave the backward eager and uncompiled. AOTAutograd traces forward and backward ahead of time as one joint graph, then partitions it into compilable fw/bw halves — choosing what to save vs recompute at the partition boundary (this is where activation-memory-vs-recompute decisions physically live, e.g. min-cut partitioning).

Along the way it functionalizes — rewrites in-place mutation (add_, views) into pure ops — because both Inductor and XLA reason about SSA-style functional graphs, and it decomposes composite ops into a small normalized ATen/prims set so backends implement dozens of primitives, not thousands of ops. That's the whole concept load.

5 · Inductor at Concept Level

Inductor takes the ATen graph and (1) schedules — orders ops, picks loop structures and tilings; (2) fuses — merges producer/consumer elementwise/reduction chains into single kernels, and epilogues into matmuls; (3) codegens Triton for GPU (C++/OpenMP for CPU), calling vendor libs or generated templates for GEMMs; (4) autotunes — benchmarks candidate configs (block sizes, num_warps; and template-vs-cublas choices) at compile time, cached thereafter.

Why fusion is the money. Roofline logic: an unfused chain x.mul(s).add(b).relu() at bf16 reads and writes the full tensor per op — 3 reads + 3 writes = 6N·2 bytes of HBM for 3N FLOPs, AI ≈ 0.25 — hopelessly memory-bound; fused, it's 1 read + 1 write (2N·2 bytes) for the same FLOPs — 3× fewer HBM round-trips, ≈3× faster, mechanically, because time = bytes ÷ bandwidth in this regime. On decode-shaped workloads almost everything outside the GEMMs is in this regime, which is why compile's win on inference is mostly "fusion killed the elementwise traffic," not "faster matmuls."

Cudagraphs (mode="reduce-overhead"): Inductor wraps compiled regions in CUDA graph capture/replay, collapsing thousands of per-kernel launches (~5–10 µs of CPU each) into one graph replay. Mechanism: decode steps at small batch are launch-bound — GPU idle between tiny kernels waiting on the host — and replay removes the host from the loop. Constraints follow from capture: static shapes and stable memory addresses per bucket (hence per-bucket capture in serving), no eager islands inside (a graph break defeats it).

Inspecting output: TORCH_COMPILE_DEBUG=1 dumps per-graph directories (post-fusion IR, generated Triton in output_code.py); TORCH_LOGS=output_code streams the kernels; the generated Triton is readable and greppable — the fastest way to answer "did my ops fuse?" is to count kernels and look for your op names in one kernel vs three.

6 · The TPU Path: torch_xla, openxla, and the JAX Contrast

Two capture routes exist on TPU. Legacy lazy-tensor tracing: torch_xla tensors record ops into a lazy IR; at a barrier (xm.mark_step()) the accumulated graph compiles and runs. It re-traces every step (Python overhead per step) and turns any value-inspection (.item(), printing) into a synchronization + possible recompile. The current route — torch.compile(backend="openxla"): Dynamo captures once (guards, no per-step re-trace), AOTAutograd normalizes, the graph exports to StableHLO, and XLA owns everything after — fusion, layout assignment, memory planning. Inductor's entire role is played by XLA; there is no Triton in the picture.

What the engineer actually experiences differently: compile times are whole-program and large (minutes for big models vs seconds/kernel for Triton) but cached per shape signature; shape constraints are absolute (§3); and debugging swaps output_code.py for HLO artifacts — XLA_FLAGS=--xla_dump_to=/tmp/hlo, then read post-optimization HLO to see fusion decisions (the same artifact class a TPU perf-regression pipeline diffs). Op coverage gaps surface as CPU fallbacks that are catastrophic on TPU (device→host→device per op) rather than merely slow.

7 · torch.compile in Production Inference

How vLLM-class engines actually use it. Not one monolithic compile: piecewise compilation — the model is compiled in regions with attention (paged, custom-kernel, dynamic by nature) deliberately left outside the compiled/captured regions, called as a registered custom op between them. Per shape bucket: compile the pieces, then capture cudagraphs per bucket over the compiled regions; decode dispatch becomes bucket-select + graph replay. Warm-up runs the full bucket lattice before serving traffic (compile + capture + allocator steady-state). Caching across restarts (mature since ~2.6): FXGraph/AOTAutograd/Triton caches keyed remotely, plus portable torch.compiler.save_cache_artifacts() bundles shipped with the container image — turning a 20-minute cold warm-up into tens of seconds; the PGO cache additionally persists dynamism decisions. Restart cost is a fleet-availability number, so cache hit rate on warm-up is a monitored SLI, not a nicety.

Quantization interaction. Two modes: (a) compiled quant ops — int8/fp8 matmuls as ATen ops or custom ops with registered meta functions; Inductor fuses the dequant/activation epilogues around them — compile helps, killing the elementwise traffic quantization otherwise adds; (b) hand-written kernel suites (FP8 attention, fused MoE dispatch, MXFP4 GEMMs) where the kernel already fuses everything internally — compile's job shrinks to gluing between them, and a naive torch.compile over the whole model can fight them: decompositions splitting a pattern the kernel wanted whole, or a custom op without proper fake-tensor registration causing graph breaks around every call. The fix is registration discipline (torch.library.custom_op + fake impl), not abandoning compile.

The compiler-first, custom-kernel-last procedure — the decision list to say out loud:

8 · Failure-Mode Bestiary

9 · Interview Deployment Sheet

Version note, honestly stated: behavior described is the 2.10–2.12 era; automatic dynamic shapes (2.1+), mature multi-level caching incl. PGO and portable artifacts (2.6+), and Python 3.14/free-threaded support (2.10) are the recent material shifts. Anything guard-implementation-level or scheduler-source-level was deliberately excluded — deeper than this document is compiler-team territory, and saying so in the room is part of the answer.