Part VII — Synthesis

Capstone — one block, source to kernel

Lesson 18 laid the stacks side by side and showed they are the same five stages — capture, represent, optimize, lower, codegen — wearing different product names. But a Rosetta stone is still a map, and a map is not the territory. The whole track has insisted on one claim: the passes are not a menu you pick from, they are one chain, each link forced by the wall the previous one hit. This last lesson proves it by doing the thing end to end. We take a single line — y = gelu(x @ w + b) with a dynamic batch dimension, plus its backward — and walk it from Python source to a tuned Triton kernel, naming the FILE# that owns each transform and tallying a scoreboard (kernels, HBM bytes, peak memory) the whole way. Then we point outward: the frontiers, the honest limits, and where to go next in the library.

The previous step left this broken

Lesson 18 gave you the column-by-column map of torch.compile, XLA, TVM, MLIR, and TensorRT against the five stages — enough to recognize any stack in the wild. What it could not do is connect the links for one concrete program so you feel each pass arrive as a forced consequence rather than a labeled box. You have seen fusion (05) and memory planning (06) and autotuning (10) each in isolation, with their own widget and their own toy. You have not yet watched a single tensor expression fall through all of them in order, with one running scoreboard, and seen the eager baseline of ≈5 kernels and ≈10× traffic collapse into 1–2 kernels at ≈2× traffic. Without that single traced example, the track is a pile of correct passes and not yet a pipeline.

Linear position

Forced by: the passes only earn their keep as a connected pipeline; seen separately they look optional, and the reader cannot yet narrate how a real line of PyTorch becomes a kernel.
New idea: walk one block — y = gelu(x @ w + b), dynamic batch, with backward — through every stage as a forced chain, watching the scoreboard (kernels, HBM bytes, peak memory, speedup-vs-eager) move at each link; then restate the three motifs and survey the frontiers and limits.
Forces next: nothing in this track — the chain is closed. The remaining walls (megakernels, distributed compilers, learned cost models, dynamic shapes, sparsity, the long tail of ops) point out of the compiler and into the rest of the library.

The plan

Four moves. (1) The end-to-end trace: source → Dynamo → FX/IR → canonicalize → op-coverage/decompose → autodiff/partition → fuse → quantize → memory-plan → layout → schedule → codegen → autotune → runtime/CUDA-graph → debug/verify → dynamic-shape guard → optional shard, each tagged with the FILE# that owns it, with the scoreboard collapsing as we go. (2) The three motifs restated one last time. (3) The frontiers, honestly bounded. (4) Where to go next in the library. Drive the scoreboard widget: turn the passes on in order and watch eager (≈5 kernels, ≈10× bytes) walk down to fully-compiled (1–2 kernels, ≈2× bytes).

1 · The trace — one block, every stage

The block is the smallest interesting thing a transformer does: a linear layer with a bias and a GELU activation, on a batch whose size is not known at compile time.

import torch

def block(x, w, b):           # x: (B, 1024)  w: (1024, 4096)  b: (4096,)
    return torch.nn.functional.gelu(x @ w + b)   #  y: (B, 4096)

cblock = torch.compile(block, dynamic=True)      # B is a symbolic dim
y = cblock(x, w, b); y.sum().backward()          # forward + backward

Run it eagerly first, so we have a number to beat (01). Take B = 2,048, bf16 (2 bytes). The intermediate x@w and the bias-added tensor and the GELU output are each 2,048 · 4,096 · 2 ≈ 16.8 MB. Eager runs this as a chain of separate kernels: the matmul, then a bias-add, then the GELU (which itself is several elementwise ops — the tanh-approx GELU is roughly cube → scale → add → tanh → scale → multiply). Counting the matmul plus the elementwise ops conservatively, that is ≈5 kernels, each reading its ≈16.8 MB input from HBM and writing ≈16.8 MB back. The matmul reads x and w once, but the elementwise tail re-reads and re-writes the activation four-plus times — ≈10× the activation's bytes crossing HBM when one read and one write would do. That is the eager baseline: ≈5 kernels, ≈10× activation traffic, peak memory = sum of all the live intermediates. Now lower it.

source → capture · 02Dynamo interprets the CPython bytecode and records the straight-line tensor ops into an FX graph. Because dynamic=True, the batch dim is captured as a symbolic s0 and a guard is installed (dtype, rank, s0 ≥ 2) so the cached graph is reused for any batch — no graph break here. Scoreboard: unchanged; we just have a graph now.

graph → typed IR · 03The FX graph becomes a typed SSA dataflow graph: every value carries shape (s0, 4096), dtype bf16, device cuda. SSA / value semantics is what makes the next rewrites local and safe. Scoreboard: unchanged; this is the substrate the passes run on.

canonicalize / fold / CSE · 04Cheap, always-valid cleanups: constant-fold the GELU constants, CSE any repeated subexpression, DCE dead nodes, put the graph in one normal form so fusion patterns match. Scoreboard: a node or two vanish; no kernel-count change yet — these passes set up the big win.

cover the op vocabulary · 14Before any rewrite can match, every op must be in a vocabulary the back end can generate. The GELU is decomposed to core prims (the tanh-approx chain) — canonicalization at the op level — so it fuses; the matmul stays a single recognized op. Had the block contained a custom or exotic op, this is where it would graph-break back to eager, get a registered decomposition / torch.library custom op with a meta kernel for shape propagation, or fall back to a vendor library (cuBLAS/cuDNN). Scoreboard: no number moves if coverage holds — but one uncovered op mid-chain would shatter the fusion region below.

autodiff + functionalize + partition · 12AOTAutograd functionalizes (removes in-place/aliasing — pays back the SSA motif from 03), builds the backward graph by reverse-mode AD (one vjp per op), then the partitioner min-cuts the joint forward/backward graph deciding what to save vs recompute. The same middle-end passes below now run on both halves. Scoreboard: more nodes (a backward appeared), but the same fusion/planning logic applies.

fuse · 05The highest-leverage pass. The bias-add and the whole GELU chain are epilogue-fused into the matmul: while the matmul's output tile is still hot in registers, add the bias and apply GELU, then write once. The ≈5 kernels collapse to ≈1 (the fused matmul+epilogue); ≈10× activation traffic drops to ≈2× (read inputs, write output once). This is the move that pays for the whole pipeline.

quantize (optional) · 15The one pass that deliberately spends precision for bytes. Insert QDQ nodes, pick per-channel scales by calibration (PTQ), then fuse the dequant into the matmul epilogue so the GEMM runs in int8/fp8 and only widens at the end. Scoreboard: model size and weight traffic drop ≈2× (int8 vs bf16), lighting int8 tensor cores — the only honest cost is a small, budgeted accuracy drift. This is the single biggest remaining attack on the bytes-moved currency after fusion.

memory-plan · 06Liveness analysis over the fused graph: the separate intermediates the eager run materialized no longer exist, and any remaining buffers whose lifetimes don't overlap share storage. Peak memory drops from sum-of-intermediates toward max-concurrent. Scoreboard: peak memory falls sharply.

layout · 07Assign physical layouts so the matmul feeds the tensor cores (tiled/blocked fragments, aligned/padded) and propagate to cancel adjacent transposes. Scoreboard: no kernel-count change, but the one kernel now runs near peak FLOP/s instead of strided-and-slow.

schedule · 08Pick the loop nest: tile sizes (BLOCK_M, BLOCK_N, BLOCK_K), reorder, vectorize, stage to SRAM. Same math, but the tiled schedule turns a memory-thrashing matmul into one that reuses each loaded tile across many FLOPs. Scoreboard: est. speedup climbs as the schedule improves locality.

codegen · 09Inductor emits Triton for the fused kernel — you write the tile program, Triton's MLIR compiler maps warps/threads and emits PTX → SASS. Trading a little peak perf for portability and far less codegen surface. Scoreboard: the kernel is now real machine code; the artifact is cached (compile once, run many).

autotune · 10The tile/warps/stages config is shape-dependent. Triton autotunes over a config list (or a learned cost model prunes candidates), keyed by shape/dtype, and caches the winner. Scoreboard: est. speedup reaches its ceiling — the last ~20–40% comes from here.

runtime + CUDA-graph · 16The tuned kernel is still inert — the runtime allocates its buffers from a caching pool (realizing the 06 plan without hot-path cudaMalloc), dispatches on a stream, and amortizes launch overhead. CUDA-graph capture/replay records the whole launch sequence once and replays it with a single driver call — the real fix for lesson 01's launch-bound ceiling. For Python-free serving, torch.export → AOTInductor emits a .so. Scoreboard: launch overhead collapses from N driver calls to ≈1; wall-time drops most on small/launch-bound shapes.

debug + verify trust · 17Fusion (05), the float reassociation from canonicalize (04), quantization (15), and graph-replay (16) all legitimately move bits — so the artifact is checked against an fp32 reference with an rtol/atol budget that distinguishes acceptable drift from a real bug. If it's wrong, the minifier bisects to the offending pass and TORCH_LOGS surfaces graph-break / recompile storms. Scoreboard: no KPI moves — this is the pass that lets you trust the other numbers.

dynamic-shape guard · 11The tuned kernel is guarded by the symbolic batch s0: any B reuses the same compiled, tuned kernel (one compile, not a compile-storm), padding or bucketing only at the edges. Scoreboard: the speedup holds across batch sizes instead of re-paying compile.

shard (optional) · 13If w is too big for one device, annotate a sharding spec over a device mesh; the compiler propagates shardings, inserts the collectives (e.g. one all-reduce over the partitioned K dim), and overlaps comm with compute. Scoreboard: per-device memory drops; a communication term appears, scheduled to hide behind the matmul.

Read the chain top to bottom and the linear-thinking spine is visible in one frame: capture exists because optimization needs the whole program (01→02); the typed IR exists because rewriting must be safe (02→03); canonicalization exists to make fusion's patterns match (03→04→05); fusion creates and destroys buffers, which forces memory planning (05→06); planned buffers need a physical shape, which forces layout (06→07); the laid-out graph still only says what, which forces scheduling (07→08); a schedule is a plan, which forces codegen (08→09); codegen can emit anything, which forces search (09→10); every pass assumed static shapes, which forces dynamic-shape handling (10→11); we only compiled inference, which forces autodiff (→12); the joint graph outgrows one device, which forces distributed compilation (→13). And then shipping forces its own chain: a rewrite can only match ops the back end can generate, which forces op coverage (→14); fp16 still leaves bytes on the table, which forces quantization (→15); a compiled artifact is inert and launch-bound, which forces the runtime and CUDA-graph replay (→16); and every one of those passes can move bits, which forces a way to debug and trust the result (→17). Nothing here is optional. The scoreboard collapse — ≈5 kernels → 1–2, ≈10× bytes → ≈2×, peak memory from sum-of-intermediates toward one fused tile — is the whole track expressed as four numbers.

2 · The three motifs, one last time

preserve semantics

Every link above is a provable no-op on the math. Canonicalization (04), functionalization (12), the fusion rewrite (05), and the dynamic-shape guard (11) are all bounded by this invariant — it is what a compiler may legally do. The one place it bites: float non-associativity (04). The compiler is fast because it is correct first.

abstraction ↔ specialization

Lowering spends portability to buy speed. The block starts as portable framework ops (03), descends through linalg/loop-nest schedules (08), and bottoms out as a tensor-core-specific Triton/PTX kernel locked to one chip (09). Every rung down is more speed, less portability — and the stack (18) is just a choice of how far down to commit.

memory traffic is the cost

The scoreboard's headline number is bytes, not FLOPs. Fusion (05), memory planning (06), and layout (07) all attack HBM traffic — the roofline currency from 01 and gpu_kernels. The ≈10×→≈2× collapse is the speedup; the matmul's FLOPs barely changed. For everything but the GEMM itself, value = FLOPs per byte loaded.

3 · Frontiers — where the chain is still being built

The pipeline you just traced is real and shipping, but it is not finished. Each motif above has an open edge.

megakernelsFusion (05) merges a handful of ops; the frontier is fusing the whole model — or a whole decode step — into one persistent kernel that never returns to the host between layers, killing launch overhead and inter-op HBM traffic entirely. Hard because it strains register/SRAM budgets and the scheduler; promising for low-latency inference.

compilers for distributedSharding (13) is moving from a bolt-on to a first-class compiler concern: Pallas/Mosaic (tile kernels for TPU/GPU under JAX), Shardy (the successor sharding dialect to GSPMD), DTensor + torch.compile. The goal: write single-device semantics, let the compiler emit the SPMD program and tune the collectives.

ML-for-compilersAutotuning (10) is increasingly run by learned cost models and learned search policies that predict latency from features and prune the space — the compiler optimizing itself. Still brittle across new hardware and op shapes, but the direction is clear: less benchmarking, more prediction.

dynamic shapes & sparsityThe hardest open problem (11) stays open. Symbolic shapes cut recompiles but still pad/bucket; sparsity (MoE routing, structured/unstructured sparse weights, ragged batches) defeats the static-shape assumptions fusion and tiling lean on. Compiling a kernel whose work depends on data is mostly unsolved.

the long tail of opsCompilers nail the common path — matmul + elementwise + norm + attention. The thousandth op (a custom loss, an exotic conv, a new sampler) often graph-breaks (02) back to eager or needs a hand-written kernel. Coverage, not peak perf on the hot path, is where compilers quietly lose.

Be honest about the limits. Peak performance still needs humans — the fastest FlashAttention and GEMM kernels are hand-tuned; the compiler gets you 80–95% of the way, and the last few percent on a hot kernel is still a person with a profiler. Debugging compiled code is hard — a graph break, a silent recompile storm, or a fused kernel that's numerically off-by-a-float sends you reading generated Triton and guard logs, not your Python. And the deepest limit: the compiler cannot fix a bad algorithm. It will faithfully, beautifully lower an O(T²) attention into the best kernel that O(T²) can be — but it will not invent FlashAttention's online softmax for you. Algorithmic wins live above the compiler; it makes a good algorithm fast, it does not make a bad one good.

4 · Drive the whole pipeline

The widget below is the scoreboard from §1 made live. The block flows left→right through every stage; click a stage to toggle its pass on or off. The KPIs — #kernels, HBM bytes (× the activation), peak memory, est. speedup vs eager — recompute from the same back-of-envelope formulas the earlier lessons used. Start from the eager baseline (all off) and turn the passes on in order: watch fusion (05) do the heavy lifting on kernels and traffic, planning (06) drop peak memory, and layout/schedule/autotune (07/08/10) climb the speedup. Try turning fusion off while leaving codegen on, to see that emitting Triton without fusing buys you almost nothing — the order is the point.

End-to-end scoreboard: y = gelu(x @ w + b)

Toggle each pass. Turning the chain on in order walks the scoreboard from eager (≈5 kernels, ≈10× activation bytes) to fully-compiled (1–2 kernels, ≈2× bytes, near-peak speed). The knob that proves the chain: enable codegen without fusion — kernels barely move; fusion is where the win lives.

kernels launched

—

HBM traffic (× activation)

—

peak memory

—

est. speedup vs eager

—

The shape of the walk is the lesson: nothing happens for the first three stages (capture, IR, canonicalize just prepare), then fusion drops the floor out from under kernels and bytes, planning halves peak memory, and the back-end stages multiply the speedup. The eager-to-compiled gap is not one magic pass — it is a chain, and the chain only works in order.

Failure modes & checklist

Failure modes

Expecting a speedup with nothing to fuse. A single big GEMM is already compute-bound; torch.compile can't beat cuBLAS on it. Signal: you compiled a one-matmul function and saw ≈1× — there was no memory-bound chain (05) to collapse.
Recompile storms from dynamic shapes. Forgetting dynamic=True (11) so every new batch size triggers a fresh compile. Signal: first call after each shape change stalls for seconds; the guard log shows repeated specialization.
Graph breaks killing the win. A data-dependent .item() or an unsupported op splits the region (02), so fusion (05) only sees fragments. Signal: the profiler still shows many small kernels; torch._dynamo reports breaks.
Trusting the compiler to fix the algorithm. Compiling an O(T²) attention and expecting FlashAttention. Signal: memory still scales as T² — the compiler lowered your algorithm faithfully; it didn't replace it.
Compile-time cost ignored. Autotuning (10) and codegen (09) cost real seconds-to-minutes; pay it once and amortize, or it dominates short-lived jobs. Signal: a script that runs the model twice is slower compiled than eager.

Checklist

Price eager first. Count kernels and HBM round trips (01) so every pass has a number to beat.
Confirm the chain captured. Check for graph breaks (02) before blaming the optimizer — fusion can't span a break.
Make fusion the goal. Everything before it (03, 04) just sets up the one pass (05) that moves the scoreboard.
Read peak, not just time. Memory planning (06) and the partitioner (12) decide whether it fits at all.
Declare dynamism. Use symbolic shapes (11) to compile once and reuse across batch/seq.
Cache the tuned artifact. Pay search (10) once; key it by shape/dtype and reuse.

Checkpoint

Try it

Open the scoreboard with everything off and read the eager baseline: ≈5 kernels, ≈10× activation traffic. Now turn passes on one at a time in spine order and note exactly which KPI each one moves — you should see capture/IR/canonicalize (02–04) move nothing, fusion (05) drop kernels from ≈5 to ≈1 and bytes from ≈10× to ≈2×, planning (06) roughly halve peak memory, and layout/schedule/autotune (07/08/10) climb the speedup from ≈1× toward its ceiling. Then do the experiment that proves order matters: reset to eager, turn on only codegen (09) and autotune (10) without fusion — confirm kernels and bytes barely budge. Finally, predict before you click: if you enable sharding (13) on a single device, does the speedup go up? (No — it adds a communication term with nothing to overlap; sharding pays off only when w doesn't fit one device.)

Where this points next

This is the end of the chain. You can now look at any pass in torch.compile, XLA, TVM, or MLIR and name the inefficiency that forced it, sketch the five-stage stack from memory, and trace a line of PyTorch from source to a tuned kernel with a running cost in your head. The track is closed — the link from here points out, to the rest of the library where each frontier lives in depth. For the kernels themselves — how a fused Triton or CUDA kernel is actually written, the memory hierarchy, tensor cores — go to gpu_kernels and triton_kernels. For how compiled kernels get wired into a production serving engine with paged KV and continuous batching, vllm and sglang. For the full training stack the compiler serves — parallelism, scaling, the systems half — cs336. And for shrinking the model the compiler then lowers, distillation. Back to the start any time at the track index.

Takeaway

A neural network is a hardware-agnostic graph; a GPU rewards a few large, fused, specialized kernels — and the wide gap between them is closed by one chain of semantics-preserving rewrites, each link forced by the wall the last one hit. Walking y = gelu(x @ w + b) with a dynamic batch from source to kernel makes the chain concrete: capture (02) because optimization needs the whole program, a typed SSA IR (03) so rewrites are safe, canonicalize/fold/CSE (04) to make patterns match, autodiff + functionalize + partition (12) to build and optimize the backward, fusion (05) — the highest-leverage pass, everything before it just sets it up — to collapse the bias+GELU into the matmul epilogue, memory planning (06) to cut peak, layout (07) for tensor cores, scheduling (08) for the loop nest, codegen to Triton (09), autotuning (10) for the config, all guarded by symbolic dynamic shapes (11) and optionally sharded across a device mesh (13) — then made shippable: op coverage so every node lowers (14), quantization for a further ≈2× on size and weight traffic (15), a runtime with CUDA-graph replay that collapses N launches to ≈1 (16), and a debug/verify step that earns trust against an fp32 reference (17). The scoreboard tells the story in four numbers: ≈5 kernels → 1–2, ≈10× activation bytes → ≈2× (and ≈2× smaller again with int8), peak from sum-of-intermediates toward one fused tile, and a multi-× speedup. The three motifs hold throughout — preserve semantics bounds what is legal, abstraction↔specialization is what lowering spends, and memory traffic is the real cost. The frontiers (megakernels, distributed compilers, learned cost models, dynamic shapes, sparsity, the long tail) are open, and the honest limits stand: peak perf still needs humans, compiled code is hard to debug, and the compiler will lower a bad algorithm faithfully but never fix it. That last line is the whole field in one sentence — a compiler makes a good algorithm fast; it does not make a slow one smart.

Interview prompts

Walk y = gelu(x @ w + b) from Python to a kernel and name what each stage does. (§1 — Dynamo capture (02) → typed FX/SSA IR (03) → canonicalize/fold/CSE (04) → AOTAutograd autodiff+functionalize+partition (12) → epilogue-fuse bias+GELU into the matmul (05) → memory-plan (06) → layout for tensor cores (07) → schedule tiles (08) → codegen Triton (09) → autotune (10) → guard with symbolic batch (11) → optionally shard (13).)
Which single pass moves the scoreboard most, and why is everything before it "just setup"? (§1–2 — fusion (05): for memory-bound chains wall-clock ≈ bytes/bandwidth, and fusing the bias+GELU epilogue into the matmul turns ≈5 kernels/≈10× traffic into ≈1 kernel/≈2×; capture/IR/canonicalize only exist to make that rewrite legal and matchable.)
You compiled a function that is one big matmul and saw ≈1× speedup. Why? (§Failure modes — a single GEMM is already compute-bound; there's no memory-bound chain to fuse, so the compiler can't beat cuBLAS. The win comes from collapsing the elementwise tail, which this function doesn't have.)
Where does autodiff fit in the pipeline, and what does the partitioner decide? (§1 link 12 — AOTAutograd functionalizes then builds the backward by reverse-mode AD; the partitioner min-cuts the joint graph to choose save-vs-recompute, trading backward FLOPs for activation memory. The same middle-end passes then optimize both halves.)
How does the compiler avoid a recompile per batch size, and what does that cost? (§1 link 11 — symbolic shapes carry the batch as s0 with guards, so one compiled+tuned kernel serves all batches; the cost is some padding/bucketing waste and looser specialization than a per-shape compile.)
Name three things a compiler still can't do well. (§3 — reach hand-tuned peak on hot kernels, give clean debugging across graph breaks/recompiles/float drift, and fix a bad algorithm (it will faithfully lower an O(T²) attention but won't invent online softmax). Sparsity and the long tail of ops are also largely open.)
Why is "the same five stages with different bets" the right mental model for every stack? (§2, lesson 18 — torch.compile/XLA/TVM/MLIR/TensorRT all capture→represent→optimize→lower→codegen; they differ only in IR choices and where on the abstraction↔specialization ladder they commit, so the trace above maps onto any of them.)