← Inference From First Principles

Compilation and Execution Graphs

Capturing a decode step as a graph: what Dynamo's guards promise, what a CUDA graph removes, what an Inductor-compiled kernel fuses, what each costs to build, and why the first claim in this chapter's earlier draft was wrong.

One decode step of GPT-2 took 16.5 milliseconds when PyTorch ran it operation by operation, and 3.4 milliseconds when the same GPU work was recorded once and replayed. The tokens were identical, all 129 of them. Four fifths of the eager step was not the GPU computing anything; it was the host deciding what to launch next.

The previous draft of this chapter reported a different result: that compilation could not be measured at all, because TorchInductor “could not find a C++ compiler” on this machine. That was wrong. Visual Studio 2022 is installed; its compiler was simply not on the PATH of the shell that ran the test. The error was real, the conclusion drawn from it was not, and the chapter had been recording an environment accident as a fact about the machine. With the compiler’s environment loaded, Inductor works on the CPU, though not in the way I first expected, and on the GPU it does not work because a different dependency is missing. Section by section, this chapter separates three things that the word “compile” covers, measures what each does to a decode step, and counts what each costs before it pays.

Where this sits

Chapter 7 measured a launch floor: about 8 to 13 microseconds per small GPU operation on the book’s machine on a good day. Chapter 9 found that a decode step is far from the hardware’s limits, and Chapter 11 turned that into three ceilings (launch, bandwidth, compute) with a step bounded by whichever is highest. Chapter 10 built the static cache, a preallocated buffer written in place. Chapter 17 found that the cost of a draft model’s step tracked its layer count more closely than its parameter count, which is consistent with a launch-bound step.

All of those chapters described the launch ceiling and worked around it. This chapter attacks it directly and measures how much is there to remove. It also raises the question every later chapter on serving has to answer: what has to stay the same between one step and the next for a recorded computation to be valid?

The Inference Lab piece added here is the fixed-shape decode step, with its state held in device buffers; a GraphedStep that captures it into a CUDA graph; a set of tools for counting graphs, breaks and recompilations; and the commands for measuring compile cost. The fixed-shape step is the form a graph-based version of Chapter 15’s scheduler or Chapter 17’s speculative loop would need, since a graph wants a step that does not allocate or change shape; neither has been built that way.

For the Pi series, there is a reproducibility lesson. A compiled artefact depends on the library version, the compiler, the hardware and the shapes it saw. Chapter 2’s manifest named the model and the weights; a compiled graph is a second artefact derived from them, with its own provenance.

Three things called compilation

They are different mechanisms with different requirements, and conflating them is what produced the chapter’s first error.

What it does What it needs What breaks it
Dynamo capture Traces Python into graphs of tensor operations, protected by guards Nothing beyond PyTorch; a backend decides what to do with the graph Code it cannot trace (data-dependent branches, .item(), print); guards that fail
Inductor Generates and builds fused kernels from a graph On CPU a C++ compiler; on CUDA, Triton Operations it cannot lower; a missing toolchain
CUDA graphs Records the GPU work of a whole step and replays it A CUDA GPU; no compiler Any change in shape, address or control flow between capture and replay

PyTorch’s documentation describes the first and the third as follows (established background; PyTorch 2.6 documentation read on 21 October 2026). Dynamo encodes its assumptions about locals and globals as guards, checks them on every call, and recompiles when one fails, up to a configured limit. A graph break splits the program into compiled pieces joined by ordinary Python, and the FAQ says to keep breaks to a minimum. A CUDA graph replay “reads from and writes to the same (virtual) memory addresses” on every replay, control flow that depends on data is prohibited, every tensor must keep its size and layout, and operations that synchronise the CPU with the GPU are prohibited during capture; the documentation says the main benefit is eliding CPU overhead. The two are independent: a Dynamo graph can be run eagerly, which is how the first experiment below counts graphs without a compiler, and a CUDA graph can be captured from plain eager code, which is how the second works.

Capture, and what it promises

A guard is a promise about the future, and the cheapest way to see what the promises are is to count how many times Dynamo has to make a new one. A small backend does this: it receives each graph, counts the call nodes, notes whether any shapes were symbolic, and runs the graph eagerly. No compiler is involved.

First, what stops a graph from forming. For four toy functions, with torch._dynamo.explain:

Function Graphs Breaks Reason Dynamo reports
A matmul, relu, sum 1 0
if x.sum() > 0: 2 1 generic_jump TensorVariable()
x / x.abs().max().item() 2 1 Tensor.item
a print inside the function 2 1 builtin: print

Every break is the same situation seen from different angles: Python needs a value that the graph has not computed yet, so the graph ends, Python runs, and a second graph starts. The cost of a break is the cost of two graphs and the lost chance to fuse across the boundary; it is not an error. But a loop that breaks on every token does so on every token.

Second, what shape changes cost. A function was called with twelve different batch sizes:

Setting Compilations for 12 batch sizes
Default 2 (the second with symbolic shapes)
dynamic=False 8, then the limit
dynamic=True 2

PyTorch’s documentation says it assumes static shapes by default, tries a dynamic recompile when a size change forces one, and that dynamic=False recompiles for each distinct size. The measurement agrees. With dynamic=False, the eighth compilation hit torch._dynamo.config.cache_size_limit (8 on this install), and Dynamo logged a warning and stopped compiling new shapes; I did not time the calls after that, but a limit is a cliff, and a serving system that sees nine batch sizes has found it. Under dynamic=True the function still compiled twice, not once; I did not establish why.

Third, the decode steps of earlier chapters. Chapter 10’s step, with a cache that grows by concatenation, was compiled and called 20 times: it compiled 8 times and reached the limit. Its cache tensors change length every call, and the guard on their shape fails every call. The fixed-shape step that this chapter builds was compiled and called 20 times, with the position advancing, and compiled once: one graph, no breaks, 716 operation nodes (506 by explain’s own count). The difference between the two is the whole subject of the next section.

Making the step graphable

To be recorded and replayed, a step must do the same work at every position. Chapter 10’s forward_cached does not: it slices the cache to its current length, builds a mask of a different size, and reads the position from a Python integer. The fixed-shape version moves everything that varies into device buffers:

  • the token being fed is a [B, 1] tensor on the device, and the step writes its own argmax back into it;
  • the position is a one-element tensor on the device, which the step advances;
  • the output tokens are collected in a [B, max_len] buffer, written at the position;
  • the cache is Chapter 10’s StaticCache, written with index_copy_ at the position;
  • attention reads the whole buffer and masks the slots beyond the position, instead of slicing.

The core of the step:

state.k[layer].index_copy_(2, state.pos, k)          # write at the device-side position
state.v[layer].index_copy_(2, state.pos, v)
scores = (q @ state.k[layer].transpose(-2, -1)) / math.sqrt(dh)   # [B, h, 1, max_len]
scores = scores.masked_fill(state.positions > state.pos, float("-inf"))
...
nxt = logits.argmax(-1)
state.out.index_copy_(1, state.pos, nxt)             # record the token
state.tok.copy_(nxt)                                  # feed it to the next step
state.pos.add_(1)                                     # advance, on the device

Nothing on the host depends on the values. That has a price: attention now touches max_len cache slots every step, however many are valid. At batch 1 and max_len 256 on this model that costs nothing measurable (the fixed-shape eager step, at 16.5 ms, is not slower than Chapter 10’s concatenating one, at 18.1), but the price scales with the buffer and with the batch. Unit tests check what the design promises: the fixed-shape decoding equals Chapter 10’s static-cache decoding token for token; garbage written into every unreached cache slot changes nothing; and two sequences decoded together give each its own answer.

The GraphedStep captures one call of this step:

stream.wait_stream(torch.cuda.current_stream())        # warm up on a side stream
with torch.cuda.stream(stream):
    for _ in range(3): decode_step(state, weights, config)
torch.cuda.current_stream().wait_stream(stream)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
    logits = decode_step(state, weights, config)       # recorded, not executed for real
# replay: graph.replay() runs the same kernels on the same buffers

Warm-up and capture both execute the step, so they advance the position and write the cache. The class saves every buffer first and restores it afterwards; a test checks that the state after construction is the state before. Without that, the first replayed token would be the fifth, which looks like a model error and is not one.

What a CUDA graph removes

The measurement: GPT-2 in float32, a 16-token prompt, 128 decode steps, max_len 256, five rounds with the methods in a rotating order. Medians per step:

Batch Chapter 10 concat, eager Fixed-shape, eager CUDA graph Graph over fixed-shape eager
1 18.1 ms 16.5 ms 3.37 ms 4.9×
4 19.6 17.1 4.46 3.8×
16 19.5 17.4 5.32 3.3×
64 20.6 17.4 11.22 1.5×

Three things to read from it. First, the eager step takes about 17 milliseconds whatever the batch, from 1 sequence to 64. The GPU’s arithmetic is growing 64-fold and the step time is barely moving, so the GPU is not what the step is waiting for. Second, the graph’s time does move with the batch (3.4 ms to 11.2 ms), because it is now the GPU’s own time. At batch 64 the GPU work is long enough that eager execution’s launch gaps are partly hidden behind it, and the graph’s advantage falls to 1.5 times. This is Chapter 11’s picture exactly: a step’s time is the largest of the ceilings, and at small batch the launch ceiling is the largest. Third, one fifth of the eager step at batch 1 (3.37 of 16.5 ms) is what remains when the launches are gone, which is about twice the bandwidth floor: the model’s 498 MB of float32 weights at Chapter 11’s measured 304 GB/s would take 1.64 ms. A graph does not make a step run at the hardware’s limit; it makes it stop running at the host’s.

The eager timings above run without synchronising the host and the device at every step, because the token is chosen on the device. A real engine reads each token back to check for a stop condition, so I measured that too: reading the token after each step took the eager step to 17.4, 18.3, 18.6 and 20.0 ms across the four batch sizes, and the graph to 3.33, 4.61, 5.55 and 11.36 ms. Reading back costs the graph nothing noticeable: its replay finishes the GPU work, then the read returns at once. The graph’s advantage survives the round trip.

The numbers are specific to a small model. GPT-2 has 124 million parameters and 12 layers, and the step is dominated by its fixed count of small operations. A larger model has the same kind of operations on larger tensors, and the share of launch overhead falls; the Llama pair of Chapter 17 (28 and 16 layers) would show a smaller proportional gain on the target and a larger one on the draft. I have not measured it, and the previous sentence is a prediction, not a result.

What a graph costs

  • Time to build: warm-up plus capture took 0.18 seconds at batch 1. The saving is 13.1 ms per step, so a graph pays for itself in 14 steps.
  • Memory: at batch 8, the state and cache took 144 MiB, and the captured graph held about 10 MiB beyond that for the intermediates it allocated during capture.
  • A graph per shape. A captured graph has one batch size and one buffer length. A server with batches from 1 to 64 needs either a graph for each size it will see (memory for each) or padding to a few sizes (wasted work). Nothing in this chapter chooses between them.

What breaks it silently

Two hazards, both tested.

A graph keeps reading the memory it was captured with. After capturing, I replaced the state’s input token by a new tensor, as a careless refactor might. The replayed graph went on reading the old buffer, and the 5 tokens it produced were exactly those of an untouched run. The documentation’s remedy is to copy new data into the captured tensors in place, and the lesson for an engine is that anything the graph reads must be a long-lived buffer with a stable address. No error is raised.

A graph does not check its own bounds. Replaying the graph past the end of a 24-slot cache, an out-of-range index_copy_ triggered a device-side assertion and the process died. It cannot be caught in Python (the CUDA context is already in an error state), so the experiment runs in a child process. An eager loop would have raised a catchable exception at the same point. Moving the loop on to the device removes the host’s chance to check the position; the host has to check before it replays.

What Inductor fuses

Inductor takes the graph and produces new code. On CUDA it generates Triton kernels, and Triton is not installed on this machine: torch.compile on a GPU tensor fails with “Cannot find a working triton installation”. That experiment is NOT_RUN, and no Windows-specific build of Triton was installed to make it run. On the CPU, Inductor generates C++ and needs a compiler, which is where the previous draft went wrong; the repository’s helper locates Visual Studio’s vcvars64.bat, imports the variables it sets into the process, and cl appears on PATH. The toy MLP of the original example then compiles, in 3 to 8 seconds across three trials, with identical output.

What it could not compile

The real decode step did not compile. The linker failed on the generated kernel with unresolved external symbol aoti_torch_check. I bisected it with small functions:

Function Compiles with Inductor on CPU, torch 2.6, Windows
w[idx], an embedding-style gather no
buffer.index_copy_(2, pos, new) no
masked_fill then softmax yes
argmax yes
layer_norm, matmul, tanh yes
torch.where(mask, a, b) yes

Indexed loads and stores generate a bounds check that calls a symbol this build’s kernels cannot link. I report it as an observation of one library version on one operating system, not a general property of Inductor; other releases or platforms may differ. It is exactly the kind of fact that an “it compiles” claim hides.

To measure anything, I split the step. The embedding lookup, the mask construction and the final bookkeeping (writing the token, advancing the position) stay eager; they are a handful of small operations. Everything else, the twelve layers of attention and MLP, goes through torch.compile, with the cache write rewritten as a full-buffer torch.where that Inductor can lower. A unit test checks the split step equals the indexed step (logit difference below 1e-5 over 15 steps on a small model) and that the compute core is a single graph with no breaks.

What it bought

One thread, GPT-2 in float32, 10 steps per measurement, six interleaved rounds:

Eager Inductor Speed-up
First process (empty compile cache) 42.5 ms 35.6 ms 1.19×
Second process (warm disk cache) 46.2 ms 38.4 ms 1.20×

(The two processes are separate runs, and the machine’s speed differs between them, so the ratios, not the absolute times, are comparable.) The compiled code produced the same tokens over 40 steps; the largest logit difference to eager was 0.001, from fusion reordering float32 arithmetic.

A speed-up of one fifth is what fusion alone can do here. The step moves about 0.5 GB of weights through one core in 42 ms, which is the shape of a step limited by memory traffic, not by operation count. Compilation fuses elementwise operations (layer-norm, gelu, the residual adds) so that intermediates are not written to memory and read back, and it replaces a string of small calls with one. It cannot make the matmuls read fewer weights. I did not measure the CPU’s single-thread bandwidth, so “bandwidth-bound” is what the arithmetic suggests, not something this chapter shows.

What it costs to start

The first call of the compiled core took 25.7 seconds with an empty cache, and 8.4 seconds in a new process whose disk cache was already populated. Against the saving of 6.9 ms per step (cold) and 7.8 ms (warm), compilation pays for itself after 3,715 steps cold and 1,077 steps warm. A process that serves a few thousand tokens before restarting never recovers the cold cost; one that serves millions does, many times over. This is the number a serving system needs and the one most easily omitted from a benchmark, which usually starts its clock after the warm-up.

The warm number is for the same machine, the same library version and the same shapes. A compile cache is an artefact, and what it is valid for (the library release, the compiler, the CPU or GPU, the shapes) is part of what Chapter 2’s manifest would need to record if the cache were shipped. I did not measure a cache moved between machines or versions.

What the numbers do not show

  • One small model. GPT-2’s 12 layers of small operations are the best case for removing launch overhead. Larger models have more arithmetic per launch.
  • Greedy decoding at batch sizes 1 to 64 with no scheduler. A real engine’s step also includes sampling choices, stop checks, admission and block management. A graph removes the model call’s overhead and nothing else.
  • No Inductor on the GPU. The most important compiler comparison for this book’s GPU measurements is the one that could not run. Whether Triton-generated kernels give more than the CPU’s 1.2 times on a launch-bound step is unmeasured.
  • A compiled step that is not quite the eager step. The split step uses a full-buffer where for the cache write, which does extra work, on both the eager and the compiled side of the comparison.
  • Compile times on a desktop. They were taken on a machine running other programs, once cold and once warm.
  • One library release. Dynamo’s limits and behaviours (the recompile limit, the automatic dynamic shapes, the link error) are those of PyTorch 2.6.0. The documentation read was the 2.6 pages; the “stable” URLs now point to a later release.
  • The reduce-overhead and max-autotune modes of torch.compile, which combine Inductor with CUDA graphs and with autotuning, were not run.
  • Graphs with a variable batch or length, and graphs together with speculative decoding or prefix caching.

Inference Lab increment

The Lab now has:

  • DecodeState and decode_step: a fixed-shape greedy step whose token, position, output and cache live in device buffers.
  • decode_core and decode_step_split: the same step with the indexed work kept eager so that a compiler that cannot lower indexing can still be used.
  • GraphedStep: capture with warm-up on a side stream, state restoration, and replay.
  • CountingBackend, and the capture and capture_step commands, for counting graphs, breaks and recompilations.
  • load_msvc_environment: finds Visual Studio and brings cl into the process.
  • break_even_calls and amortised_seconds: the arithmetic of whether a compile pays.
  • env, graphs, inductor_ops and inductor commands.

The rule this adds: a recorded computation is valid only while its guards hold, and the work of an engine is to make the guards hold by construction. Buffers have stable addresses, shapes come from a small set, and nothing the host decides depends on data inside the step. The speed-up is the reward for keeping the contract and the price of breaking it is either a recompilation or, for a CUDA graph, silently wrong tokens.

What you should now be able to do

  • Say what Dynamo’s guards protect, and use a counting backend to find out how many times a function recompiles for a given sequence of inputs.
  • Explain why Chapter 10’s step recompiles every call and the fixed-shape step compiles once.
  • Write a decode step with no host-visible dependence on position, capture it into a CUDA graph, and verify the graph against eager decoding token for token.
  • Name the contract a CUDA graph imposes and describe what happens when each clause is violated.
  • Compute from a compile time and two step times how many calls a compilation needs, and say what that implies for a deployment.
  • Distinguish a result about a library on one operating system from a result about a technique.

Key terms

  • Graph capture: recording operations as a graph.
  • Guard: a condition checked on each call that the captured graph still applies.
  • Graph break: a point where capture stops and Python runs.
  • Recompilation: capture and compilation repeated because a guard failed.
  • Dynamic shape: a dimension treated symbolically, so that one graph serves many sizes.
  • Inductor: PyTorch’s compiler backend, producing C++ on CPU and Triton kernels on CUDA.
  • CUDA graph: a recorded sequence of GPU work, replayed with one launch.
  • Cold start: the time before the first fast call, including compilation.
  • Break-even: the number of calls after which a compile has paid for itself.

Reader exercises

  1. Add a data-dependent branch to the fixed-shape step (for example, stop writing the cache when a token equals the end-of-text token). What does Dynamo report, and what does CUDA graph capture do?
  2. Capture graphs for batch sizes 1, 2, 4 and 8 and pad other sizes up. What is the total graph memory, and what fraction of the padded steps is wasted work at a uniform batch distribution from 1 to 8?
  3. Measure the graph with max_len 256, 512 and 1024. How does the cost of attending to the whole buffer grow, and at what context does slicing win?
  4. Install a Triton build for your platform and repeat the Inductor measurement on the GPU. How does the speed-up compare to the CPU’s 1.2 times, and does the indexed decode step compile?
  5. Measure how the cold and warm compile times depend on the number of layers, by compiling 2-, 6- and 12-layer prefixes of the model.
  6. Apply a GraphedStep to the draft model of Chapter 17. How much does it change the draft-to-target cost ratio, and so the break-even acceptance rate?
  7. Change the library version or the compiler and observe whether the disk compile cache is reused. What does the cache key contain?
  8. Use TORCH_LOGS=recompiles on Chapter 10’s step and read the reason for each recompilation.

References

Evidence in this chapter

Claim Category
Guards, graph breaks, recompilation to a limit, and CUDA graph constraints behave as documented Established background (PyTorch 2.6 documentation)
A plain function is 1 graph; a data-dependent branch, .item() and print each break the graph once Book result, torch._dynamo.explain, 4 toy functions
12 batch sizes: 2 compilations by default, 8 with dynamic=False (limit reached) Book result, torch 2.6
Chapter 10’s step compiled 8 times in 20 calls; the fixed-shape step once Book result
Fixed-shape eager and CUDA graph decoding equal Chapter 10’s tokens (129) Book result
Graph replay 3.37 ms against 16.5 ms eager at batch 1 (4.9x), 1.5x at batch 64 Book result, 5 interleaved rounds
About four fifths of the eager step at batch 1 is not GPU work Book result, inferred from the graph’s time
A graph pays for itself in 14 steps and holds about 10 MiB at batch 8 Book result
A replaced input tensor is ignored by the graph; replay past the buffer kills the process Book result, tested
Inductor on CPU: 1.19x to 1.20x on the compute of the step; 25.7 s cold, 8.4 s disk-warm; break-even 3,715 and 1,077 steps Book result, one thread, torch 2.6, Windows
Inductor on this build cannot link kernels with indexed loads or stores Book result, observation of one release and platform
Previous claim “no C++ compiler on this machine” Corrected: the compiler existed and was not on PATH
Inductor on CUDA, reduce-overhead, larger models, variable-shape graphs NOT_RUN

Implementation status

  • Manuscript: content/books/inference/18-chapter.md
  • Runnable examples: examples/ch18/compile.py (original toy MLP; its “compiler unavailable” conclusion is superseded), examples/ch18/graph_lab.py (fixed-shape step, CUDA graphs, capture tools and experiments)
  • Tests: tests/test_ch18_compile.py, tests/test_ch18_graph_lab.py (fixed-shape decoding against Chapter 10, stale cache slots never read, batch independence, split step equal to the indexed step, single graph for the step and its core, recompilation counts, graph breaks, environment parsing, break-even arithmetic; CUDA tests for replay equality and the stale-input hazard run only when a GPU is present)
  • Evidence: research/ch18/sources.md, research/ch18/results.json, research/evidence-register.md
  • Verification command: python -m unittest discover -s tests -v

The chapter is a draft with draft = false. It makes no claim about GPU compilers, larger models or other library versions. Chapter 19 begins Part 4 with mixture-of-experts inference, where the question of which computation happens at a given step stops being fixed.