When Karpathy reproduced GPT-2 on an A100 (§6.7), one line of code, model = torch.compile(model), took a training step from about 300 milliseconds to about 130. The model was the same, the batch was the same, and the number of FLOPs was the same. So where did more than half of the time go?

It went to two things this chapter has already named. Most of a Transformer's operations other than matmuls are memory-bound (§11.2), and in eager PyTorch each one is a separate kernel that reads its inputs from HBM and writes its output back, so a chain of them pays for a round trip at every link. And each of those kernels is launched by Python, at a cost of microseconds of CPU time, which for small kernels leaves the GPU waiting. Kernel fusion attacks the first and CUDA graphs the second, and torch.compile does both. This section shows how much each one saves, builds a fused LayerNorm plus GELU, looks at what the compiler generates, and teaches you to read a profiler trace, which is how you find out which problem you have in the first place.

Why fusion saves memory traffic

Consider a chain of $k$ elementwise operations on a tensor of $n$ elements, say a bias add, a GELU, and a dropout. Run as separate kernels, each one reads $n$ elements and writes $n$, so the chain moves about $2kn$ elements. Fused into one kernel, the input is read once, all $k$ operations happen in registers, and the result is written once, $2n$ elements. Since the chain is memory-bound, its time is proportional to its bytes, and fusion makes it about $k$ times faster.

In plain words, an intermediate result that never leaves the chip costs nothing to store.

The general rule is just as simple. Divide a computation into groups, each running as one kernel. A group's traffic is the bytes of every tensor it reads that was produced outside the group (each counted once, however many of its operations read it), plus the bytes of every tensor it produces that is needed outside the group, by a later group or as a final output. Tensors that are produced and consumed inside the group stay in registers or shared memory. Compilers such as TorchInductor reason in exactly these terms when they decide what to fuse.

Not everything can be fused, and the limits come from the hardware.

• Elementwise into elementwise always works, since each output element depends only on the same element of the inputs.
• Reductions with their neighbors work when the reduction's extent fits in one program, so the block reduction of §11.4 can finish on chip. A row-wise softmax or LayerNorm over a few thousand features can absorb the elementwise operations before and after it, as the fused softmax of §11.6 did.
• Elementwise into a matmul's epilogue works because each output tile of $\mathbf{C}$ sits in registers before it is stored (§11.5), so a bias, an activation, a residual add, or a cast can be applied there. Prologues, elementwise operations applied to an operand as it is loaded, work too.
• Two matmuls in a row do not fuse in general, because the second needs whole rows of the first's output, which no single tile holds. Doing it anyway, for the specific case of attention, takes an algorithm, FlashAttention (§12.3), not a compiler pass.
• Global reductions, such as a sum over a whole tensor, need a grid-wide combination and end a fusion group (§11.4).

Here is an MLP sublayer with its bias, activation, residual add, and the next LayerNorm. Click the gaps between operations to fuse or split them, and watch the bytes, the kernel count, and the estimated time on an H100.

def plan_traffic(ops, sizes, groups, outputs):
    total = 0
    for op, (ins, outs) in ops.items():              # placeholder: ignores the groups entirely
        total += sum(sizes[t] for t in ins) + sum(sizes[t] for t in outs)
    return total

A fused LayerNorm and GELU

Put numbers on one common pair, a LayerNorm (§5.2) followed by a GELU, on activations of 16,384 tokens by 4,096 features in bfloat16, 128 MiB. Written as plain PyTorch tensor operations, the LayerNorm alone is a mean, a subtraction, a square, another mean, an add and a reciprocal square root, a multiply, a multiply by the learned gain $\boldsymbol\gamma$, and an add of the learned shift $\boldsymbol\beta$, and the tanh form of GELU,

is several more. Each full-size operation reads and writes the whole tensor. The cell counts the passes.

PyTorch's built-in F.layer_norm and F.gelu are already single kernels, so the realistic eager baseline is two kernels and four passes, and fusing them halves the time. A LayerNorm written out by hand in tensor operations, which is common in research code, is more than ten times slower than the fused kernel, which is the kind of gap torch.compile closes automatically.

The fused kernel gives one program to each row, loads the row once, computes the statistics in float32 registers, normalizes, applies the affine map and the GELU, and stores the result. One identity makes the GELU cheaper. Since $\tfrac12(1 + \tanh u) = \sigma(2u)$, where $\sigma$ is the logistic sigmoid, the tanh form of GELU is $h\,\sigma\big(2\sqrt{2/\pi}\,(h + 0.044715\,h^3)\big)$, one sigmoid instead of a tanh and two extra operations.

Rows too wide for one block would use the two-pass streaming approach of §11.6, or split the row across warps and merge their partial means and variances with Chan's associative formula from §11.4. Either way the statistics are accumulated in float32, as §11.4's error analysis requires. The cell checks the kernel's math in NumPy, row by row, against the unfused chain, and checks the sigmoid identity.

Fusion also helps the backward pass. A fused kernel can save only its input and the two per-row statistics, and recompute the normalized values and the GELU in the backward kernel, instead of saving every intermediate. Recomputing cheap elementwise operations costs far less than storing and reloading them, which is the same trade FlashAttention makes for the attention weights.

Launch overhead and CUDA graphs

Fusion attacks bytes. The other cost is the launch itself. In eager mode, the CPU runs Python, dispatches each operation through PyTorch, and calls the CUDA driver to enqueue a kernel, several microseconds per operation. The GPU executes kernels in order. If kernels are long, the CPU runs ahead and the queue stays full. If they are short, the GPU finishes each one before the CPU has enqueued the next, and sits idle in between. This is the overhead-bound regime of §11.2, and it is typical of inference at small batch, where a decode step is hundreds of tiny kernels (§6.5).

The cell simulates the queue for a decode step of 300 kernels of 1 to 5 microseconds each, with a launch cost of 6 microseconds per kernel, and then with all of them captured in one graph.

With small kernels, eager execution keeps the GPU busy only about half the time, and the step length is set by the CPU's launch rate, 300 launches of 6 microseconds. A CUDA graph removes the problem. You run the step once while the CUDA runtime records every kernel and its arguments, then replay the recording with a single launch. The GPU-side work is identical, and the CPU cost drops from hundreds of launches to one.

The price is rigidity. A graph replays exactly what it recorded, so shapes must be static, inputs and outputs must live at fixed addresses (hence the copies into static_x, and why serving systems preallocate the KV cache), and the captured region can contain no CPU synchronization and no Python control flow that depends on tensor values. Variable batch sizes are handled by capturing a graph per bucket of sizes. torch.compile(model, mode="reduce-overhead") does this capture for you.

def step_time(gpu_us, launch_us, graph_launch_us=None):
    busy = sum(gpu_us)
    return busy, 1.0                                 # placeholder: pretends launches are free

def min_kernel_us(launch_us):
    return 0.0

torch.compile and Inductor

You met torch.compile in §0.9. TorchDynamo captures your Python as graphs of PyTorch operations, AOTAutograd traces the matching backward graph, and TorchInductor compiles both (Ansel et al., 2024). Now we can say precisely what Inductor does on a GPU, and what it does not.

It fuses memory-bound work into generated Triton kernels. Inductor lowers each operation to a loop-level description, then its scheduler merges producers into consumers wherever the rules of the first part of this section allow, pointwise chains, reductions with their pointwise neighbors, and sibling operations that read the same input. Each fused group becomes one Triton kernel. In a profile you see them by name, with the fused operations listed in the name. triton_poi_fused_... is a pointwise group. triton_per_fused_... is a persistent reduction, in which one program holds the whole reduced row in registers, like the fused softmax of §11.6. triton_red_fused_... is a looped reduction for rows too long for that, which streams the row in tiles and carries a running result, the pattern §11.6 used for long softmax rows.

It leaves matmuls to libraries, by default. Matrix multiplications and convolutions are dispatched to cuBLAS and cuDNN. With mode="max-autotune", Inductor also benchmarks its own Triton matmul templates, which can absorb elementwise epilogues, against the library and keeps the fastest.

It chooses what to save for the backward pass. AOTAutograd partitions the joint forward and backward graph with a min-cut, deciding which intermediates to save and which cheap ones to recompute, so the fused backward kernels read fewer bytes.

It reduces overhead. The compiled graph runs as a short sequence of kernels without Python dispatch for each operation, and mode="reduce-overhead" wraps it in CUDA graphs.

It also pattern-matches a few important subgraphs, such as an attention written out as matmul, softmax, matmul, and replaces them with fused library kernels. What it does not do is invent a new algorithm. It will not discover the online softmax for you in a custom attention variant, and it cannot fuse across a graph break, the point where Dynamo falls back to Python because of something it cannot trace (a .item(), a print of a tensor, data-dependent control flow, an unsupported library call). Each break splits fusion groups and adds launches. You can see everything it generates.

TORCH_LOGS="graph_breaks,recompiles" python train.py     # where and why graphs break or recompile
TORCH_LOGS="output_code" python train.py                 # print the generated Triton kernels
TORCH_COMPILE_DEBUG=1 python train.py                    # dump graphs, kernels, and fusion decisions to disk

Reading the generated code once is worth the effort. You will find kernels very much like the ones in §11.6, with tl.program_id, masked loads, and the operations of several PyTorch calls in one body, which demystifies the compiler and shows you exactly which operations were fused.

Reading a profiler trace

Everything so far assumed you know where the time goes. You find out with a profiler. torch.profiler records CPU operations and GPU kernels, with their timings, shapes, and the links between each launch and the kernel it started.

The exported trace shows time running left to right, with rows for each CPU thread (Python functions, aten:: operators, and cudaLaunchKernel calls) and rows for each GPU stream (the kernels themselves, with names like sm90_xmma_gemm_bf16... for cuBLAS matmuls, flash_fwd_kernel for attention, triton_poi_fused_... for compiled pointwise groups, vectorized_elementwise_kernel for eager elementwise operations, and kernels whose names begin with nccl for communication). Arrows connect each launch to its kernel. Reading a trace means recognizing a handful of shapes.

To decide whether a single kernel is good, compute its achieved throughput. Take its duration from the trace and its FLOPs and bytes from the shapes, and compare FLOP/s with the compute roof and bytes per second with the bandwidth roof (§11.2). A LayerNorm moving 256 MiB in 85 microseconds achieves about 3.2 TB/s, near the H100's 3.35, and is done. A GEMM at 300 TFLOP/s on large square bfloat16 inputs has headroom. For a deeper look at one kernel, NVIDIA's Nsight Compute reports the same comparison per kernel, as percentages of peak for compute and for each level of memory, along with occupancy, stalls, and bank conflicts. Nsight Systems shows the whole-program timeline across CPU, GPU, and network.

def summarize_trace(events, top=3):
    busy = sum(e["dur"] for e in events)               # placeholder: double-counts overlaps
    start = min(e["start"] for e in events); end = max(e["start"] + e["dur"] for e in events)
    return {"wall": end - start, "busy": busy, "utilization": busy / (end - start), "largest_gap": 0, "top": []}

A workflow for finding the bottleneck

Put the pieces together into a routine you can follow on any slow model.

1. Measure the step time and the MFU (§11.2). If MFU is already near what similar runs achieve, the remaining gains are in the algorithm, not the kernels.
2. Take a trace of a few steady-state steps. Check first for GPU idle time. Gaps mean overhead, synchronization, or input stalls, and no kernel optimization will help until they are gone.
3. Rank kernels by total time. Group them into matmuls, attention, communication, and everything else. If "everything else" is a large share, compile or fuse.
4. Check the biggest kernels against their roofs. Compute achieved FLOP/s or bytes per second. A kernel far below its roof is worth investigating in Nsight Compute, one near its roof can only be helped by changing what it computes.
5. Change one thing, and measure end to end. Many optimizations that speed up one kernel move the bottleneck somewhere else.

In PyTorch

The profiler works on the CPU too. The block profiles an MLP written with separate bias, GELU, and residual operations, and prints the operators sorted by time. Compare the share of the elementwise operators (aten::gelu, aten::add) with that of the matrix products (aten::mm). On a GPU the matmuls run on tensor cores and shrink far more than the elementwise work does, so the elementwise share grows, which is exactly the work that fusion removes.

Read more

• Ansel et al. (2024). PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. Dynamo, AOTAutograd, and Inductor, with the design of Inductor's fusion and its measured speedups.
• He (2022). Making Deep Learning Go Brrrr From First Principles. Why fusion helps, why overhead matters, and how to tell which regime you are in.
• Ivanov et al. (2021). Data Movement Is All You Need: A Case Study on Optimizing Transformers. Measures how much of a Transformer's time goes to memory-bound operations and how far fusion can reduce it.
• PyTorch. Accelerating PyTorch with CUDA Graphs. How graph capture works in PyTorch, its constraints, and measured gains.
• PyTorch. Profiling your PyTorch Module. The profiler API, schedules, and trace export.
• NVIDIA. Nsight Compute documentation. The per-kernel profiler, including the speed-of-light and roofline sections.

Watch next

▶

▶

▶