Compilers That Fuse: XLA, torch.compile, CUDA Graphs

Why automate fusion at all

A transformer has hundreds of ops and ships in dozens of shapes; hand-fusing each path is hopeless. An operator-fusion compiler does it mechanically. It takes the computational graph your framework already builds, finds chains of ops it can safely fuse, generates one kernel per region, and schedules the whole thing. You get most of the bandwidth wins from the previous guide for free — and crucially the backward graph from autodiff gets fused too, where naive frameworks otherwise leak the most memory traffic.

XLA: ahead-of-time compilation

XLA (Accelerated Linear Algebra) takes the whole graph ahead of time, for a fixed set of input shapes, and compiles it into a tightly fused, optimized program for the target accelerator. Because it sees the entire computation up front, it can make global decisions — layout assignment, fusion boundaries, buffer reuse — that a kernel-at-a-time runtime cannot. The cost is recompilation when shapes change, which is why XLA-based stacks lean on fixed or bucketed shapes and padding. It is the backbone of JAX and a major TensorFlow path, and it targets GPUs and TPUs from the same graph.

Q_{\text{unfused}} \approx 2kn \;\gg\; 2n \approx Q_{\text{fused}}

Why XLA's tightly fused program is fast: chaining k element-wise ops unfused round-trips every intermediate through HBM (~2kn bytes); fusion keeps them in registers and moves only ~2n.

torch.compile: capture the eager graph

PyTorch is eager — ops run as Python executes them — so there is no graph to compile. torch.compile manufactures one. TorchDynamo traces your Python at runtime, capturing it into graphs and gracefully falling back to eager whenever it hits something it cannot trace (a so-called graph break). The captured graph is handed to a backend — by default Inductor, which fuses pointwise ops, generates Triton kernels for GPUs, and calls vendor libraries for the big GEMMs. You keep Python's flexibility and get compiled speed on the parts that matter.

model = torch.compile(model)        # first call traces + compiles
for batch in loader:                # later calls hit the cached graph
    loss = model(batch).loss
    loss.backward()
# watch for 'graph break' logs: each break splits the fused region

One line opts in; the win comes from keeping fused regions large, so minimize graph breaks.

CUDA Graphs: killing launch overhead

Even perfectly fused kernels pay a fixed CPU cost to launch each one — a few microseconds of driver work. When kernels are tiny (small batches, LLM decode of one token at a time), launch overhead can exceed the compute and the GPU sits idle between launches, CPU-bound. CUDA Graphs capture an entire sequence of launches once, freeze it into a replayable graph, and let the GPU fire the whole sequence with a single submission. The per-kernel CPU cost vanishes. It pairs naturally with the compilers above — capture the fused program, then replay it.

T_{\text{eager}} \approx N\ell + \sum_{k=1}^{N} t_k \;\xrightarrow{\text{CUDA Graph}}\; T_{\text{graph}} \approx \ell + \sum_{k=1}^{N} t_k

Eager launch cost scales with kernel count N; a captured CUDA Graph replays all N kernels from one launch, collapsing N·ℓ down to a single ℓ.

What compilers cannot do

Be honest about the ceiling. A compiler fuses within the bounds you give it: dynamic shapes trigger recompiles or graph breaks; control flow that depends on tensor values resists capture; and the compiler optimizes one device's program, not the communication between devices. None of these tools place a single byte across the network. To go beyond one accelerator you need collectives and a fabric — the next guide.