PyTorch Compile Stack: Operator Fusion from TorchDynamo to TorchInductor

Summary: From torch.jit.trace to torch.compile, PyTorch's compile stack has been through three generations. This post clarifies what TorchScript / TorchDynamo / TorchInductor each do, measures torch.compile speedups with reproducible scripts, explains operator fusion, and covers which modes to use in production.

0. Experiment Goals

• Map the three-generation evolution: TorchScript → TorchDynamo → TorchInductor
• Understand TorchDynamo's graph capture mechanism (bytecode-hook capture, not tracing)
• Measure fusion gains with torch.compile and read the compile process via TORCH_LOGS
• Master production selection: mode="reduce-overhead", fullgraph, dynamic shapes

1. Three Generations: First, Who Does What

How it works:

• TorchScript (2019): torch.jit.script/trace converts a Python model to a static graph. Problem: static graphs do not support dynamic control flow (trace records one path), and Python compatibility is poor—community reception was lukewarm.
• TorchDynamo (2022): captures automatically via bytecode-level hooks on Python frames, intercepts only at the Python layer, carves out graphable subgraphs for backend compilation, and falls back to Eager elsewhere. No user code changes needed—this is the first layer of torch.compile.
• TorchInductor (2022+): the default backend of TorchDynamo, translating captured graphs into Triton kernels (GPU)—connecting the PyTorch compile stack to the Triton compiler (the link to the KB Triton Compiler post).

Python model
   │ torch.compile
   ▼
TorchDynamo (bytecode-level subgraph capture)
   │
   ▼
TorchInductor (graph → Triton kernels)
   │
   ▼
CUDA execution

2. One-Line Speedup: torch.compile Out of the Box

import torch, time

def model(x):
    # Typical MLP + softmax: many small ops
    for _ in range(4):
        x = torch.nn.functional.linear(x, torch.randn(512, 512, device="cuda"))
        x = torch.relu(x)
    return torch.softmax(x, dim=-1)

x = torch.randn(1024, 512, device="cuda")

# Eager mode (default)
for _ in range(5):
    model(x)
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(200):
    model(x)
torch.cuda.synchronize()
t_eager = (time.perf_counter() - t0) / 200

# Compiled mode
compiled = torch.compile(model)
for _ in range(5):
    compiled(x)                     # first call triggers compilation (warmup)
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(200):
    compiled(x)
torch.cuda.synchronize()
t_compile = (time.perf_counter() - t0) / 200

print(f"Eager: {t_eager*1000:.2f} ms | Compile: {t_compile*1000:.2f} ms | speedup {t_eager/t_compile:.2f}x")

Results (sample magnitude · A100):

Eager: 0.38 ms | Compile: 0.21 ms | speedup 1.8x

Why it is faster:

• In Eager mode every op (linear, relu, softmax) launches an independent kernel; one forward starts ~10 kernels, each with fixed launch overhead (~5-10μs).
• Operator fusion: TorchInductor fuses adjacent ops into one Triton kernel—linear+relu fused, softmax uses a reduce kernel; kernel count drops from 10 to 3-4. What is saved is kernel launch overhead and intermediate tensor reads/writes (memory bandwidth is the bottleneck for compute-light small models).
• More layers and more fragmented ops → bigger fusion gains; a single large GEMM already sits near bandwidth limits, so compile gains are limited—do not expect miracles on already-optimized large ops.

3. See Through the Compile Process: TORCH_LOGS

# Enable compile logs to watch Dynamo capture and Inductor codegen
TORCH_LOGS="+dynamo,+inductor" python3 bench.py 2>&1 | head -40
# Key output:
#   [graph break]  — where the Python layer could not be graphified (control flow/external calls)
#   [Generated Triton code] — the kernel source Inductor generated

Engineering notes:

• graph break is the enemy of compile gains: the more fragmented the graph, the fewer fusion opportunities. Common triggers: Python control flow (if/for depending on tensor values), print, dict operations, external-library calls.
• Use torch.compile(model, fullgraph=True) to force "compile whole graph or error"—useful during development to expose breaks; do not enable in production (dynamic control-flow models will fail outright).
• Use torch._dynamo.explain to list break locations and fix them one by one for maximum gain:

explain = torch._dynamo.explain(fx)(model, x)   # prints break reasons

4. Production Selection: Modes and Dynamic Shapes

# Recommended production combo
compiled = torch.compile(
    model,
    mode="reduce-overhead",   # CUDA graph trims scheduling overhead; throughput-first
    fullgraph=False,          # allow fallback for stability
)

Step-by-step:

• mode="reduce-overhead": additionally captures the whole compute graph once with CUDA Graph and replays it, saving per-iteration kernel dispatch overhead—big wins for small batch/short sequences; at large batch the dispatch share is small and gains flatten.
• mode="max-autotune": autotunes every op for the best kernel config—compile time grows to minutes, suited for "compile once, serve long" inference deployments; use sparingly in training iteration.
• Dynamic shapes: when input shapes vary, Dynamo recompiles per shape (compile-cache bloat). Use torch.compile(..., dynamic=True) or constrain shape ranges to reduce recompiles.
• Training vs inference: compile in training mainly saves forward+backward kernel launches (throughput +10-30%); inference deployment uses compile + torch.inference_mode() + CUDA Graph for the largest latency wins.

5. Traps and Boundaries

1. First call compiles slowly: warmup triggers compilation (seconds to minutes); pre-warm in production or cache (torch._inductor.cudagraph to disk) in CI.
2. Dynamic control-flow models: Python branches like if length > N in Transformers cause graph breaks and discount gains—move branches outside the model or accept partial compilation.
3. Custom ops: third-party CUDA kernels (some flash-attn wrappers) may not be recognized by Inductor; graphs break there.
4. Slightly higher memory: CUDA Graph capture reserves memory; watch out on memory-tight setups.
5. Version sensitivity: Dynamo/Inductor evolve fast; re-benchmark after PyTorch upgrades (±10% swings possible).

6. Engineering Takeaways

1. torch.compile is the default acceleration move: 1.5-2× for op-dense models, try with zero code changes first.
2. Fusion is the core mechanism: kernel launch overhead and intermediate tensor reads/writes are the savings.
3. graph break is the enemy: expose with fullgraph=True in dev, locate with torch._dynamo.explain, fix one by one.
4. mode="reduce-overhead" for throughput (CUDA Graph); pair with inference_mode() for latency.
5. Dynamic shapes bloat the compile cache: dynamic=True or constrain shapes.
6. Slow compile is normal: pre-warm/cache in CI; don't wait for compilation on the first production request.
7. Don't expect miracles on large ops: fusion gains come from fragmented ops; a single large GEMM already sits near bandwidth limits.

7. Further Reading

• PyTorch torch.compile docs
• TorchInductor design notes
• KB: 08. AI Compiler · Triton Compiler · Triton Learning Path