torch.compile and CUDA graphs

01 · Concept

Concept

You profile single-stream decode on Qwen3.8-27B and the trace is mostly holes. The GPU does a burst of work, idles, does another burst, idles again. Utilization sits low, the power draw never approaches the card’s limit, and moving to a faster accelerator changes almost nothing. Nothing is broken. You have hit a bottleneck that is not on the GPU at all.

Lesson 7.3 established the first half of the story: decode has terrible arithmetic intensity. To emit one token the machine streams a very large quantity of weights through the arithmetic units and performs a small number of operations on each byte, so bandwidth, not compute, sets the pace. This lesson adds the second half. Every operation in eager mode carries host-side cost before any arithmetic happens: a Python frame, a dispatch through the operator registry, argument checking, and finally a launch that hands a kernel to the driver.

Count them. A single decoder layer in this model performs, roughly, two normalizations, the projections and the internals of its token mixer, three FFN matrix multiplies plus a gating multiply, and a handful of elementwise operations and reshapes — call it a few dozen kernels, and more in the DeltaNet layers, whose short convolution and recurrent state update are made of several small steps each. Take thirty as an illustrative average across 64 layers and you are issuing on the order of two thousand launches to produce one token. At a few microseconds of host time per launch, that is several milliseconds per token spent purely on dispatch. Turn it into a ceiling: eight milliseconds of host work per token caps you near 125 tokens per second regardless of what the GPU can do.

Now compare the two ceilings honestly, because the naive version of this argument is wrong. At bf16, streaming the whole 27B-parameter weight set for every token puts the bandwidth ceiling in the tens of tokens per second on a top-end accelerator — below the launch ceiling. So at bf16 batch 1, launch overhead is not yet the binding constraint; it is a tax of a few tens of percent riding on top of a bandwidth wall. Here is the turn: quantize to 4-bit (lesson 7.9) and the bytes you must move per token fall by roughly a factor of four, so the bandwidth ceiling rises by roughly a factor of four. The launch ceiling does not move at all. The bottleneck changes hands. This is why “compile it” is optional advice for a bf16 27B model and mandatory advice for a quantized one, and why small models running on fast cards are almost always launch-bound.

The first remedy is compilation. torch.compile traces the Python bytecode of the forward pass into a graph, guards that graph with assumptions about the tensors it saw (dtypes, shapes, device, and more), and lowers it to generated kernels. Two things get cheaper. Chains of elementwise work — normalization scaling, the SiLU gate multiply, residual adds — fuse into single kernels, which cuts both the launch count and the round trips to memory that made those operations expensive in the first place. And the whole traced region issues as compiled code rather than as re-interpreted Python.

Guards are also where compilation goes wrong. When a guard fails the graph is recompiled, and decode changes the sequence length on every single step. A model that recompiles once per token is dramatically slower than eager and will look, in a naive benchmark, like compilation is useless. The corrected approach is to make the varying dimension dynamic so one graph covers many lengths, keep the shapes that must stay static in a small set of buckets, warm up until the recompile counter stops moving, and only then measure. Measuring the first call after compiling — the wrong turn almost everyone takes once — times the compiler rather than the model.

The second remedy attacks what is left. Even a well-fused graph still issues a sequence of launches, and that sequence is identical on every decode step. CUDA graph capture records it once and replays the whole thing as a single submission, so the host issues one command instead of hundreds. The requirements are strict and they explain a great deal of serving-engine design. Shapes must be static, memory addresses must be static, and there can be no host synchronization or CPU-dependent control flow inside the captured region, because a replay does not re-run your Python — it re-executes recorded launches against recorded pointers.

Hold the accounting rather than the flags. Eager spends its time telling the GPU what to do; compilation reduces how much telling is needed, and graph capture reduces how often you have to tell it. Neither buys you a faster matrix multiply. They buy back the gaps in the trace — and the gaps are the whole reason a 27B model can leave a very expensive card mostly idle.