4.9 Kernel Fusion, torch.compile, CUDA Graphs & Compilers¶

Modern GPU hardware delivers extraordinary peak throughput on paper, but reaching even half that peak in practice requires eliminating the hidden tax that fragments every computation: kernel launch overhead, redundant memory round-trips, and the interpreter overhead of Python-level dispatch. This chapter is about the compiler and runtime machinery that hunts those inefficiencies down — op fusion, CUDA graphs, torch.compile, and the broader compiler ecosystem from XLA to TVM.

By the end you should be able to explain why fusing a ReLU into a matrix multiplication matters, how TorchDynamo captures a graph without sacrificing Python flexibility, when CUDA graphs pay off, and what the Inductor backend actually does to the fused computation before it hits the device.

Why Plain PyTorch Is Slow (and Why That Is Surprising)¶

PyTorch’s eager mode is a triumph of usability: every line runs immediately, errors surface instantly, and you can print() anywhere. But each call to an eager operator hides a sequence of events that accumulates wall-clock time even when the actual arithmetic takes microseconds:

1. Python interpreter overhead. The call dispatches through PyTorch’s C++ dispatcher, which inspects tensor dtypes, devices, and autograd metadata.
2. Kernel launch. cudaLaunchKernel submits work to the GPU’s hardware queue. Each launch costs on the order of a few to tens of microseconds of CPU time.
3. Global memory round-trip. For a sequence of elementwise ops — say, x = gelu(linear(x)) — each intermediate result is written to GPU DRAM and immediately read back by the next kernel, even though the values could have stayed in registers or L2 cache.

Consider a simple fused operation: y = relu(a * b + c) where all tensors are [4096, 4096] FP16 — \(4096 \times 4096 \times 2\,\text{B} = 32\,\text{MB}\) each. A naive eager implementation launches three kernels: t = a*b (reads two tensors, writes one), u = t + c (reads two, writes one), y = relu(u) (reads one, writes one). That is \(3 + 3 + 2 = 8\) tensor-sized transfers, or 256 MB of DRAM traffic. A fused kernel reads a, b, c once, keeps the intermediates in registers, and writes y once — 4 transfers, or 128 MB. Taking a round 2 TB/s of HBM bandwidth (an H100 PCIe part is about 2 TB/s; the SXM part is about 3.35 TB/s, and a well-written streaming kernel achieves roughly 70–85% of peak):

so fusion halves the traffic here. The win grows with the length of the chain: fusing \(n\) pointwise ops over a shared operand set pushes the traffic ratio toward \(n\), because the fused kernel’s cost stays pinned at (inputs + output) while the naive cost grows linearly in \(n\). That is the arithmetic ideal. Real speedups are smaller due to occupancy and scheduling, but the memory-traffic argument is real. See The Roofline Model & Performance Engineering for the full bandwidth/compute roofline framework.

Kernel Fusion: Mechanisms and Taxonomy¶

Kernel fusion means merging several GPU kernels into one so that intermediate values never leave the SM (Streaming Multiprocessor). There are several distinct flavours.

Horizontal (Producer–Consumer) Fusion¶

A producer kernel writes a tensor that is immediately consumed by the next kernel. After fusion, both live in a single kernel and the intermediate stays in registers or shared memory. Classic examples:

• Bias-add + activation (GELU, SiLU, ReLU)
• LayerNorm, which is three passes (mean, variance, normalize) collapsed into one
• Softmax: finding the row max, computing exponentials, and dividing, all in one pass (this is exactly what FlashAttention I: IO-Awareness & The Online Softmax exploits)

Vertical (Loop) Fusion¶

Multiple loops over the same tensor are merged. For example, computing both mean and var of a tensor in a single pass rather than two sequential reductions.

Operator Tiling & Blocking¶

A matmul followed by an elementwise activation can be fused by computing tiles of the matmul result into shared memory, immediately applying the activation to that tile, then writing to DRAM. This avoids the full DRAM write-then-read for the matmul output.

The Triton language (covered in depth in Writing GPU Kernels with Triton) makes writing such fused kernels tractable. TorchInductor, the backend behind torch.compile, generates Triton code automatically. For kernels that still need to be hand-authored, PyTorch’s newer Helion DSL adds an autotuned, Python-embedded layer on top of Triton: you write ordinary PyTorch-style tensor code and Helion searches over indexing, tiling, and pipelining strategies to compile it down to Triton (and, more recently, to TPU Pallas), narrowing the gap between hand-tuned kernels and what torch.compile fuses automatically.

CUDA Graphs: Eliminating CPU Launch Overhead¶

Fusion addresses memory traffic. CUDA graphs address the other overhead: CPU-side kernel launch latency.

The Problem: CPU is the Bottleneck for Short Kernels¶

On modern hardware, launching a CUDA kernel costs roughly 5–20 µs of CPU time. For a large transformer forward pass with many small elementwise kernels, the CPU spends a significant fraction of time just scheduling work rather than doing it. On a batch size 1 inference request — common in interactive serving — GPU utilization can be surprisingly low because the CPU cannot feed kernels fast enough.

How CUDA Graphs Work¶

CUDA graphs record a sequence of GPU operations (kernel launches, memory copies, memsets) into an opaque graph object during a capture phase. Thereafter, the entire graph is replayed with a single CPU call — cudaGraphLaunch. The GPU receives the full work description at once and can schedule it optimally, without the CPU re-entering the driver for each kernel.

import torch

# Typical pattern: warmup, capture, then replay in a loop.

def build_cuda_graph(model, static_input):
    """Capture model forward pass into a CUDA graph."""
    # 1. Warmup: run eagerly a few times to warm caches / allocate memory
    for _ in range(3):
        _ = model(static_input)

    torch.cuda.synchronize()

    # 2. Capture phase: all CUDA work between begin/end is recorded
    g = torch.cuda.CUDAGraph()
    with torch.cuda.graph(g):
        static_output = model(static_input)

    torch.cuda.synchronize()
    return g, static_output


def replay_cuda_graph(g, static_input, new_input, static_output):
    """Replay the captured graph with new data by updating the static buffer."""
    # IMPORTANT: the graph is bound to the same memory addresses.
    # We copy new data into the static buffer before replaying.
    static_input.copy_(new_input)
    g.replay()              # Single CPU call — replays entire recorded sequence
    return static_output.clone()  # Clone before next replay overwrites it

The critical constraint: memory addresses must be identical across captures and replays. This means CUDA graphs work best for fixed-shape, fixed-batch workloads. vLLM and TensorRT-LLM both use CUDA graphs for the decode phase of inference (constant batch × 1 token per step) — see TensorRT-LLM, TGI & Other Serving Stacks.

When CUDA Graphs Break¶

• Dynamic control flow that depends on tensor values (e.g., early-exit logic) cannot be captured since the graph is a static DAG.
• Host-to-device synchronizations (e.g., tensor.item(), printing) inside the captured region stall or fail.
• Variable shapes: a graph captured with batch=32 cannot be replayed with batch=16 without recapture.

The vLLM serving system manages this by pre-capturing graphs for a discrete set of batch sizes and padding up to the closest one at runtime. Its V1 engine goes a step further with piecewise CUDA graphs: it runs torch.compile over the model, splits the compiled graph at the attention operator — whose shapes genuinely vary with each batch’s mix of sequence lengths — and captures CUDA graphs only for the fixed-shape pieces in between. The result is graph replay for the bulk of the decode step while attention remains a flexible, hand-written paged kernel. See vLLM: Architecture, PagedAttention & Internals.

torch.compile: The Full PyTorch Compiler Stack¶

torch.compile, introduced in PyTorch 2.0, is a JIT compilation pipeline that transforms eager Python code into optimized kernels with minimal user friction. It is not a single technology but a layered stack of three components: TorchDynamo, AOTAutograd, and TorchInductor.

TorchDynamo: Safe Graph Capture via Bytecode Rewriting¶

The hardest problem in compiling Python is that Python is dynamic: code can inspect its own stack frames, call arbitrary C extensions, use dynamic dispatch, and depend on Python objects in arbitrary ways. Prior approaches (TorchScript, torch.fx manual tracing) required the user to rewrite code into a compilable subset.

TorchDynamo takes a different approach: it installs a frame evaluation hook at the CPython bytecode level and speculatively traces through Python execution, building an FX Graph as it encounters PyTorch operations. When it encounters something it cannot trace — a Python print, a shape-dependent if, a call into a C extension — it records a graph break and falls back to eager execution for that segment.

import torch
import torch._dynamo as dynamo

# Demonstrate graph breaks with explain()
def my_func(x):
    y = torch.sin(x)          # traced
    if x.shape[0] > 10:       # graph break: dynamic control flow
        y = y * 2
    return torch.cos(y)       # traced (in new subgraph)

explanation = dynamo.explain(my_func)(torch.randn(5))
print(explanation.graphs)          # Two subgraphs
print(explanation.break_reasons)  # "Data-dependent control flow"

Each subgraph between graph breaks is independently compiled and cached. The guard system records the assumptions that were true during tracing (e.g., x.shape[0] == 5, x.dtype == torch.float32) and re-traces if they change. This recompilation can be expensive, which is why you want to minimize graph breaks in hot paths.

AOTAutograd: Differentiation Before Compilation¶

For training, we need not just the forward pass but also the backward pass to be compiled and fused. AOTAutograd (“Ahead-Of-Time Autograd”) uses the functorch dispatcher to trace through the autograd engine itself at compile time, producing a single joint FX graph representing both forward and backward. This joint graph is then handed to the backend, which can fuse across the forward/backward boundary — for example, fusing an activation function with its gradient computation.

import torch
from torch._functorch.aot_autograd import aot_function

def fn(x, w):
    """A simple layer: linear + sigmoid."""
    return torch.sigmoid(x @ w)

# AOTAutograd decomposes this into a joint graph.
# During compilation, it generates:
#   forward:  z = x @ w; y = sigmoid(z)
#   backward: dy/dz = y * (1 - y); ...
# The sigmoid and its gradient can be fused into a single kernel.
compiled_fn = aot_function(fn, fw_compiler=lambda g, _: g, bw_compiler=lambda g, _: g)

TorchInductor: The Default Backend¶

TorchInductor is the default lowering backend. It takes the fused FX graph and generates either Triton (for CUDA/ROCm GPUs) or C++ (for CPU). Its key optimizations:

Loop fusion and tiling. Inductor represents computation as loops over tensor elements and applies polyhedral-style analysis to identify which loops can be fused and tiled for cache locality.

Pointwise fusion. Sequences of elementwise ops are automatically merged into a single Triton kernel. A transformer block’s bias-add, GELU, and dropout might collapse into one kernel.

Reduction scheduling. Reductions (softmax, LayerNorm, mean) are split into a tile-wise pass followed by a global reduction, matching the two-pass structure that fits GPU occupancy constraints.

Epilogue fusion. Many cuBLAS/cuDNN kernels support an “epilogue” — a post-matmul elementwise op applied inside the same kernel. Inductor exploits this to fuse bias-add into gemm calls for free.

import torch

# The simplest possible torch.compile usage
model = MyTransformerBlock(d_model=4096, n_heads=32)
model = torch.compile(model)  # <-- that's it

# First call triggers JIT compilation (~30s for a large model).
# Subsequent calls use the compiled version.
x = torch.randn(8, 512, 4096, device='cuda', dtype=torch.bfloat16)
y = model(x)  # compiled

compile() Options and Modes¶

torch.compile exposes several modes trading compilation time for runtime speed:

# "default": balanced; good for most cases
model = torch.compile(model)

# "reduce-overhead": enables CUDA graphs internally for fixed shapes
model = torch.compile(model, mode="reduce-overhead")

# "max-autotune": exhaustive kernel search via Triton autotuner
# Compilation can take many minutes but delivers best throughput
model = torch.compile(model, mode="max-autotune")

# "max-autotune-no-cudagraphs": autotune without CUDA graph capture
# Safer for models with dynamic shapes or side effects
model = torch.compile(model, mode="max-autotune-no-cudagraphs")

# Dynamic shapes: recompile less aggressively when shapes change
model = torch.compile(model, dynamic=True)

# Inspect what happened: graph breaks, subgraphs, guards
torch._dynamo.explain(model)(x)

Ahead-of-Time Compilation: torch.export and AOTInductor¶

torch.compile is just-in-time: the first call for every new shape signature pays tracing and code generation, and the artifact dies with the Python process. For deployment you usually want the opposite — build once, offline, then run with no tracing and ideally no Python interpreter at all. PyTorch’s answer is a two-step path that sits alongside torch.compile rather than replacing it:

1. torch.export captures a whole-graph, graph-break-free representation: an ExportedProgram holding an ATen-level FX graph plus its symbolic shape constraints and a graph signature naming every parameter and buffer the graph consumes. Unlike Dynamo’s JIT mode it refuses to fall back to eager — a data-dependent branch makes export raise an error instead of silently splitting the model. For a shippable artifact that strictness is the feature.
2. AOTInductor lowers that exported program to a self-contained package: the same Triton/CUDA kernels Inductor would have generated, plus a compiled .so wrapper, loadable from Python or linked directly into a C++ server.

import torch

model = MyTransformerBlock(d_model=4096, n_heads=32)
model = model.to(device="cuda", dtype=torch.bfloat16).eval()
example = (torch.randn(1, 512, 4096, device="cuda", dtype=torch.bfloat16),)

# Declare which dimensions may vary; everything else is frozen at export time.
batch = torch.export.Dim("batch", min=1, max=64)
ep = torch.export.export(model, example, dynamic_shapes={"x": {0: batch}})

# Lower to a standalone package (a .pt2 file containing the compiled kernels).
path = torch._inductor.aoti_compile_and_package(ep)
runner = torch._inductor.aoti_load_package(path)   # no re-tracing at load time
y = runner(example[0])

Rule of thumb: use torch.compile for training and research loops where flexibility wins, and torch.export + AOTInductor when you need a fixed, auditable, Python-free binary — the niche TensorRT and ONNX Runtime occupy, but staying inside PyTorch’s own operator semantics so exported numerics match what you trained.

Diagnosing and Reducing Graph Breaks¶

Graph breaks prevent whole-graph compilation and limit how much Inductor can fuse. You can diagnose them with:

# Set environment variable to log every graph break with a traceback
TORCH_LOGS=graph_breaks python train.py

import torch._dynamo
# Or programmatically:
torch._dynamo.config.verbose = True

# Alternatively, use the explain() API:
explanation = torch._dynamo.explain(model)(example_input)
for reason in explanation.break_reasons:
    print(reason)

Common causes and fixes:

Worked Example: Compiling a Transformer Block¶

The Broader Compiler Ecosystem¶

torch.compile is PyTorch-specific. The LLM ecosystem also involves several other compiler stacks, each with different tradeoffs.

XLA: The TPU Compiler (and PyTorch/XLA)¶

XLA (Accelerated Linear Algebra) is Google’s compiler for TPUs, and the foundation of JAX. XLA takes a computation expressed as an HLO (High-Level Optimizer) graph and applies:

• Operation fusion: similar to Inductor, but implemented in LLVM-based passes
• Layout optimization: permutes tensor dimension orders to maximize memory access patterns on TPU systolic arrays
• Rematerialization: drops and recomputes activations to reduce memory, analogous to gradient checkpointing (see Memory-Efficient Training: Checkpointing, Offloading & LoRA Math)

PyTorch/XLA exposes PyTorch’s API running on XLA backends. The compilation model is “lazy execution”: ops accumulate in an HLO graph, which is compiled and dispatched when a value is explicitly requested (e.g., at loss.item() or a barrier).

The key difference from eager PyTorch: you must minimize the frequency of HLO graph compilations (“compilation triggers”), because each compilation is expensive. Shape dynamism is therefore particularly costly on XLA.

TVM / Apache TVM¶

TVM is an end-to-end ML compiler that takes an ONNX or Relay IR graph and applies a search-based optimization (AutoTVM, Ansor/AutoScheduler) to find near-optimal kernel implementations across diverse hardware (CUDA, Metal, Vulkan, ARM, RISC-V). Unlike Inductor which generates Triton, TVM generates device-specific low-level code by searching a space of loop transformations (tiling, vectorization, loop unrolling, operator fusion).

TVM’s strength is portability: the same model can be compiled for edge devices, CPUs, and GPUs with a single pipeline. Its relative weakness compared to hand-tuned CUDA kernels for transformers is that the search space may miss tricks specific to GPU tensor cores (WMMA/warp-level matrix multiply accumulate).

ONNX Runtime and TensorRT¶

For inference, ONNX Runtime (with its CUDA Execution Provider) and NVIDIA TensorRT both accept ONNX graphs and apply fusion, layer calibration, and quantization:

• TensorRT builds a network plan from ONNX, fuses conv-bn-relu chains, inserts FP8/INT8 quantization nodes, and selects from a library of hand-written CUDA kernels. Extremely fast for fixed shapes and supported op patterns.
• ONNX Runtime is more flexible but less aggressive. It is the substrate for HuggingFace optimum.

TensorRT-LLM (covered in TensorRT-LLM, TGI & Other Serving Stacks) wraps TensorRT with transformer-specific plugins (FlashAttention, paged KV cache, in-flight batching).

The Compiler Stack for LLM Training vs. Inference¶

Practical Integration: Before and After torch.compile¶

Below is a complete, runnable before/after comparison. We profile a small GPT-like block to measure the actual gains across eager, compiled, and compiled-with-CUDA-graphs modes.

"""
torch_compile_demo.py — Measure eager vs compiled transformer block.

Requirements: PyTorch >= 2.0, CUDA GPU
Run: python torch_compile_demo.py
"""
import torch
import torch.nn as nn
from torch.utils.benchmark import Timer

# ─── Model ────────────────────────────────────────────────────────────────────

class FeedForward(nn.Module):
    """Standard FFN with SiLU gating (SwiGLU style)."""
    def __init__(self, d_model: int, ffn_dim: int):
        super().__init__()
        self.gate_proj = nn.Linear(d_model, ffn_dim, bias=False)
        self.up_proj   = nn.Linear(d_model, ffn_dim, bias=False)
        self.down_proj = nn.Linear(ffn_dim, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU: gate * silu(up)
        # torch.compile will fuse the silu + elementwise multiply
        return self.down_proj(
            torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)
        )


class DecoderLayer(nn.Module):
    def __init__(self, d_model: int = 2048, n_heads: int = 16, ffn_dim: int = 8192):
        super().__init__()
        self.norm1 = nn.RMSNorm(d_model)
        self.norm2 = nn.RMSNorm(d_model)
        # Use nn.MultiheadAttention for simplicity; FlashAttention in practice
        self.attn  = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ffn   = FeedForward(d_model, ffn_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Pre-norm residual connections
        h = self.norm1(x)
        h, _ = self.attn(h, h, h, need_weights=False)
        x = x + h

        h = self.norm2(x)
        h = self.ffn(h)
        x = x + h
        return x


# ─── Setup ────────────────────────────────────────────────────────────────────

device = torch.device("cuda")
dtype  = torch.bfloat16

# Input: batch=4, seq_len=1024, d_model=2048
x = torch.randn(4, 1024, 2048, device=device, dtype=dtype)

# Build models
model_eager    = DecoderLayer().to(device=device, dtype=dtype).eval()
model_compiled = torch.compile(
    DecoderLayer().to(device=device, dtype=dtype).eval(),
    mode="reduce-overhead",   # enables CUDA graphs internally
)
model_maxauto  = torch.compile(
    DecoderLayer().to(device=device, dtype=dtype).eval(),
    mode="max-autotune",      # exhaustive Triton kernel search
)

# ─── Warmup (critical for fair benchmarking) ──────────────────────────────────

with torch.inference_mode():
    # Eager warmup
    for _ in range(5):
        _ = model_eager(x)

    # Compiled warmup: first few calls trigger tracing + compilation
    for _ in range(5):
        _ = model_compiled(x)

    # max-autotune warmup: may take 30–120s on first run
    for _ in range(5):
        _ = model_maxauto(x)

torch.cuda.synchronize()

# ─── Benchmark ────────────────────────────────────────────────────────────────

def bench(fn, label: str, n: int = 200):
    """Run fn n times, return median latency in ms."""
    with torch.inference_mode():
        t = Timer(
            stmt="fn(x)",
            globals={"fn": fn, "x": x},
            label=label,
        )
        result = t.blocked_autorange(min_run_time=2.0)
    ms = result.median * 1e3
    print(f"{label:40s}  {ms:.3f} ms  ({1000/ms:.0f} iter/s)")
    return ms

t_eager    = bench(model_eager,    "Eager (baseline)")
t_compiled = bench(model_compiled, "torch.compile reduce-overhead")
t_maxauto  = bench(model_maxauto,  "torch.compile max-autotune")

print(f"\nSpeedup (reduce-overhead vs eager):  {t_eager/t_compiled:.2f}×")
print(f"Speedup (max-autotune vs eager):     {t_eager/t_maxauto:.2f}×")

# ─── Memory check ─────────────────────────────────────────────────────────────
# torch.compile should not increase peak memory significantly (no caching of fwd activations)
print(f"\nPeak CUDA memory allocated: {torch.cuda.max_memory_allocated()/1e9:.2f} GB")

A typical run on an A100 80 GB with the above configuration produces output along these lines:

Eager (baseline)                          6.812 ms  (147 iter/s)
torch.compile reduce-overhead             4.231 ms  (236 iter/s)
torch.compile max-autotune                3.947 ms  (253 iter/s)

Speedup (reduce-overhead vs eager):  1.61×
Speedup (max-autotune vs eager):     1.73×

The max-autotune mode finds better tile sizes for the Triton GEMM kernels, squeezing out additional throughput at the cost of a substantially longer first-run compilation.

Training with torch.compile and Autograd¶

Using torch.compile during training requires a few extra considerations:

import torch
import torch.nn as nn
from torch.optim import AdamW

# ─── Training setup ───────────────────────────────────────────────────────────

model = DecoderLayer(d_model=2048, n_heads=16, ffn_dim=8192)
model = model.to(device="cuda", dtype=torch.bfloat16)

# Compile the model BEFORE wrapping with optimizer.
# torch.compile sees a clean module, producing better graphs.
model = torch.compile(model)

optimizer = AdamW(model.parameters(), lr=1e-4, fused=True)  # fused AdamW for extra speed

x     = torch.randn(4, 1024, 2048, device="cuda", dtype=torch.bfloat16)
label = torch.randn(4, 1024, 2048, device="cuda", dtype=torch.bfloat16)

# Standard training loop — no changes required
for step in range(100):
    optimizer.zero_grad(set_to_none=True)  # set_to_none avoids a memset kernel
    out  = model(x)
    loss = nn.functional.mse_loss(out, label)
    loss.backward()            # AOTAutograd's compiled backward runs here
    optimizer.step()

    if step % 10 == 0:
        print(f"step {step}  loss {loss.item():.4f}")

Two practical notes:

1. set_to_none=True in zero_grad avoids a separate memset kernel per parameter — a nice interaction with torch.compile.
2. fused=True in AdamW uses a fused CUDA kernel for the Adam update, reducing the kernel count from O(number of parameters) to a handful.

Compile Time, Caching, and Regional Compilation¶

Compiling a 32-layer model as one graph can take minutes, and by default that work is thrown away when the process exits. Three practices make compilation affordable on a real training run:

• Regional compilation. Compile one decoder block rather than the whole model — for blk in model.blocks: blk.compile() — and Dynamo reuses the cached artifact for every structurally identical layer. Compile time drops roughly by the number of layers, and you forfeit only fusion across block boundaries, which for a pre-norm transformer is one residual add.
• Persistent caching. Point TORCHINDUCTOR_CACHE_DIR at fast local storage so Inductor’s FX-graph and Triton caches survive across processes; a warm restart then skips code generation almost entirely. This is what makes crash-restart on a long job cheap — see Checkpointing, Fault Tolerance & Long-Running Jobs.
• Keep shapes constant. Every new (batch, seq_len) signature is a recompile. Pack and pad training batches to a fixed shape and drop the ragged final batch, or the last step of each epoch silently retriggers the pipeline. Worse, once Dynamo exceeds its recompile budget (torch._dynamo.config.cache_size_limit, renamed recompile_limit in newer releases; 8 by default) it gives up and falls back to eager permanently — which presents as a mysterious mid-run slowdown rather than an error.

For the ~100M model built in Part XIV, compiling the decoder block against a fixed [B, 1024] packed batch shape is the single largest free throughput win available in the training loop; the loop that does it is in The Pretraining Run: A Complete Single-GPU Training Loop.

For distributed training, torch.compile composes with FSDP2 and DDP. The recommended order is to compile before wrapping with FSDP, so that Dynamo captures the model graph without FSDP’s all-gather hooks in the way. Consult Distributed Training I: Data Parallelism, DDP, ZeRO & FSDP for the full FSDP setup.

Graph Breaks Deep-Dive: What Actually Stops Compilation¶

Understanding what triggers graph breaks is essential for tuning models that compile well.

Data-Dependent Control Flow¶

# BAD: Dynamo cannot trace through this — graph break every call
def forward_bad(x):
    if x.max() > 1.0:    # x.max() requires synchronizing CPU/GPU
        x = x / x.max()  # dynamic branch
    return x

# GOOD: use tensor operations throughout
def forward_good(x):
    max_val = x.max()
    # Soft clamp: x / max(x.max(), 1.0) — always the same graph
    return x / torch.clamp(max_val, min=1.0)

Python Lists of Tensors¶

# BAD: list indexing can be dynamic; Dynamo may break here
def forward_bad(tensors: list):
    return [t * 2 for t in tensors]  # Python list comprehension → break

# GOOD: stack into a single tensor
def forward_good(tensors: list):
    stacked = torch.stack(tensors)   # Single op, fully traced
    return stacked * 2

Shape Guards and Recompilation¶

When Dynamo traces with input shape [4, 1024, 2048], it records a guard: x.shape == (4, 1024, 2048). A subsequent call with [8, 1024, 2048] triggers a full recompile. You can mitigate this by marking dimensions dynamic:

import torch._dynamo

# Mark batch dimension as dynamic before compilation
model = torch.compile(model, dynamic=True)

# More surgical control with torch._dynamo.mark_dynamic
# mark_dynamic must be called on the input tensor BEFORE it enters a
# compiled region — calling it from inside a @torch.compile'd function
# raises "Attempt to trace forbidden callable".
torch._dynamo.mark_dynamic(x, 0)   # dim 0 is dynamic

@torch.compile
def forward(x):
    return model(x)

With dynamic=True, Dynamo emits symbolic shapes rather than concrete values and produces a more general compiled artifact that handles variable batch sizes without recompilation.

Interview Corner¶

Key Takeaways¶

Further Reading¶

• TorchDynamo — Ansel et al., “TorchDynamo: Towards a More Flexible Python-based ML Framework,” PyTorch Blog / MLSys 2023. The primary design document for the bytecode tracing approach.
• PyTorch 2.0 — Meta AI, “PyTorch 2.0: Our Journey to the Next Generation of Production-Ready ML Frameworks,” NeurIPS 2023 system track. Introduces the full compile stack.
• torch.compile documentation — pytorch.org/docs/stable/torch.compile.html; the official guide to modes, backends, and debugging.
• XLA: The TensorFlow Compiler — Leary & Wang, “XLA: TensorFlow, Compiled,” TensorFlow Dev Summit 2017. The foundational XLA design.
• TVM — Chen et al., “TVM: An Automated End-to-End Optimizing Compiler for Deep Learning,” OSDI 2018. The origin of the search-based compilation approach.
• Triton — Tillet et al., “Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations,” MAPL 2019. The foundation for Inductor’s code generation.
• FlashAttention — Dao et al., “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,” NeurIPS 2022. The canonical example of hand-fused kernel engineering that torch.compile integrates with but cannot yet automatically reproduce.
• TensorRT-LLM — NVIDIA open-source project (github.com/NVIDIA/TensorRT-LLM). Production-quality example of CUDA graph + TensorRT fusion for LLM inference.

Exercises¶

1. The chapter opens by listing three hidden costs of eager PyTorch: Python interpreter/dispatcher overhead, kernel launch latency, and global memory round-trips. Kernel fusion and CUDA graphs are the two main tools introduced. For each tool, say which of the three costs it primarily attacks and which it leaves untouched. Then explain why mode="reduce-overhead" in torch.compile combines both.

2. The chapter estimates that the fused version of y = relu(a * b + c) on [4096, 4096] FP16 tensors takes about \(64\,\mu\text{s}\) versus \(128\,\mu\text{s}\) naive. Redo the calculation from scratch for [8192, 8192] FP16 tensors on the same 2 TB/s device, counting every operand read and every result write, and report the naive time, the fused time, and the speedup. Then say what would have to change for the fused speedup to approach \(3\times\) instead.

3. A batch-size-1 decode step launches 140 small kernels, each doing negligible arithmetic (memory-bound, ~1 µs of actual GPU work) but costing 8 µs of CPU launch time. The CPU launches kernels serially and the GPU cannot start a kernel until the CPU has launched it. (a) Estimate the per-step wall-clock time in eager mode. (b) After wrapping the decode step in a CUDA graph, the entire sequence replays with a single 8 µs CPU call and the GPU then runs the 140 kernels back-to-back. Estimate the new per-step time and the speedup. © Which sentence in the chapter explains why this workload is a good fit for CUDA graphs?

4. Consider this forward method. Using the chapter’s rules, mark every line where TorchDynamo would insert a graph break, name the cause, and rewrite the function so it compiles into a single graph.

```python
def forward(self, x, scale_list):
    y = torch.sin(x)
    if y.mean() > 0:
        y = y * 2.0
    for s in scale_list:          # scale_list is a Python list of floats
        y = y + s
    print("debug:", y.shape)
    return torch.relu(y)
```

def forward(self, x, scale_list):
    y = torch.sin(x)
    # Replace the data-dependent branch with a tensor select.
    cond = y.mean() > 0                 # still a tensor, but...
    y = torch.where(cond, y * 2.0, y)   # ...no Python branch -> one graph
    for s in scale_list:                # floats are compile-time constants
        y = y + s
    # print removed (or guarded outside the compiled region)
    return torch.relu(y)

5. Implement a small benchmark harness, in the style of the chapter’s torch_compile_demo.py, that measures the incremental contribution of CUDA graphs on top of Inductor fusion. Compile the same FeedForward module (from the chapter) three ways — eager, mode="max-autotune-no-cudagraphs", and mode="reduce-overhead" — warm each up, and print the two speedup ratios. Explain what the difference between the last two modes isolates.

import torch
import torch.nn as nn
from torch.utils.benchmark import Timer

class FeedForward(nn.Module):
    """SwiGLU FFN, as in the chapter."""
    def __init__(self, d_model: int, ffn_dim: int):
        super().__init__()
        self.gate_proj = nn.Linear(d_model, ffn_dim, bias=False)
        self.up_proj   = nn.Linear(d_model, ffn_dim, bias=False)
        self.down_proj = nn.Linear(ffn_dim, d_model, bias=False)

    def forward(self, x):
        return self.down_proj(
            torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)
        )

device, dtype = torch.device("cuda"), torch.bfloat16
x = torch.randn(4, 1024, 2048, device=device, dtype=dtype)

def make(mode=None):
    m = FeedForward(2048, 8192).to(device=device, dtype=dtype).eval()
    return m if mode is None else torch.compile(m, mode=mode)

model_eager  = make(None)
model_noCG   = make("max-autotune-no-cudagraphs")  # fusion only
model_reduce = make("reduce-overhead")             # fusion + CUDA graphs

# Warmup: eager needs a few, compiled modes need enough to finish tracing +
# (for reduce-overhead) CUDA graph capture.
with torch.inference_mode():
    for m in (model_eager, model_noCG, model_reduce):
        for _ in range(8):
            _ = m(x)
torch.cuda.synchronize()

def bench(fn, label):
    with torch.inference_mode():
        t = Timer(stmt="fn(x)", globals={"fn": fn, "x": x}, label=label)
        ms = t.blocked_autorange(min_run_time=2.0).median * 1e3
    print(f"{label:36s} {ms:.3f} ms")
    return ms

t_eager  = bench(model_eager,  "eager")
t_noCG   = bench(model_noCG,   "inductor fusion (no cudagraphs)")
t_reduce = bench(model_reduce, "fusion + cuda graphs")

print(f"\nfusion vs eager:            {t_eager / t_noCG:.2f}x")
print(f"cuda-graphs incremental:    {t_noCG / t_reduce:.2f}x")
print(f"total (reduce vs eager):    {t_eager / t_reduce:.2f}x")

6. The chapter claims AOTAutograd’s joint forward+backward graph enables a fusion “unavailable in eager mode”: fusing an activation with its gradient. Take y = sigmoid(z). (a) Derive \(\partial y / \partial z\) and show it can be written using only the forward output y, not z. (b) Explain concretely, in terms of memory traffic and kernel count, what cross-boundary fusion saves here versus eager mode, and why eager mode cannot do it.