Part III — Middle end
Operator fusion — the central win
Lesson 04 handed you a graph that is finally clean: constants folded, dead nodes gone, common subexpressions shared, everything in one canonical form so the pattern matcher can see it. But canonicalization only rearranges and shrinks the node list — it does not change the one fact that has dogged us since lesson 01: each surviving node is still its own kernel, and each kernel pays its own full round trip to HBM. The graph is tidy and the GPU is still starving. This is the lesson that fixes it, and it is the single highest-leverage pass in the whole stack — everything before it just sets it up, everything after it cleans up the mess it makes.
norm(gelu(x @ w + b)) is still five nodes — matmul, add, gelu, then mean/var, then the normalize — and the eager runtime executes each as a separate kernel launch over the whole tensor. Every one of those memory-bound nodes reads its input from HBM, does a few FLOPs per element, and writes its output back to HBM, only for the next node to read it straight back in. The cleanups removed wasted nodes; they did nothing about the bandwidth bill of the nodes that survive. For the elementwise/reduction work that makes up most of a model's op count, that bill — bytes moved, not FLOPs — is the wall-clock. We are still paying it once per op when the whole chain could pay it once.New idea: operator fusion — merge a region of the graph into a single kernel that loads its inputs from HBM once, does all the math in registers/SRAM, and writes the result once. A memory-bound chain of length n goes from ≈2n HBM round trips to ≈2. The kinds: vertical (producer→consumer chains), horizontal (siblings sharing an input), epilogue (fold pointwise work into a GEMM/conv tile), and the algorithm-level case, FlashAttention.
Forces next: fusion decides which intermediate buffers actually exist and how long each one lives — the buffer that was a kernel output is now a register that never touches HBM. That upends allocation, so you must now plan memory deliberately (lesson 06).
dropout(bias(x)) → norm, eager ≈6× bytes vs fused ≈2×. (3) Horizontal fusion of siblings and epilogue fusion into a matmul tile; point at FlashAttention as the algorithm-level limit. (4) The fusion decision and where it stops — register/shared-memory pressure, the compute-bound boundary, wasteful recompute. (5) How real compilers actually group: Inductor's pointwise+reduction Triton kernels, XLA fusion clustering, TVM compute_at. Then drive the widget until over-fusion turns the pressure meter red.1 · The lever: FLOPs per byte loaded
One kernel is one GPU program launched from the host. Every launch carries two costs: a few microseconds of launch overhead (≈3–10 µs to queue the grid), and — usually dominant — an HBM round trip, where the kernel reads its input tensor from HBM, does a handful of FLOPs per element, and writes the result back. For the elementwise and norm ops that make up most of a model's node count, the FLOPs are trivial and the bytes are everything.
Put a number on the round trip. A residual-stream activation of shape (B, T, d) = (8, 4096, 4096) in bf16 is 8 · 4096 · 4096 · 2 ≈ 268 MB. On an H100 at ≈3.3 TB/s, just reading it costs 268 MB / 3.3 TB/s ≈ 81 µs; a read-plus-write op costs ≈162 µs. A bias add is two FLOPs per element — ≈0.27 GFLOP — which the tensor cores finish in under a microsecond. The op is ≈99% waiting on memory. This is the lesson-01 roofline restated as wall-clock: a memory-bound op's time is bytes / bandwidth, and its FLOPs are free.
The quantity that decides whether an op is starving is its arithmetic intensity — FLOPs performed per byte moved across HBM. Below the hardware's ridge point (≈300 FLOP/byte on an H100 for bf16) the op is memory-bound; the chip's ≈1,000 TFLOP/s of tensor cores sit idle while it waits on ≈3.3 TB/s of bandwidth. Almost every non-GEMM op lives far below that ridge. Fusion is the lever that raises arithmetic intensity directly: do more math per byte loaded by doing many ops on the data while it is hot in registers, instead of bouncing it to HBM between each one. The full first-principles version of this argument lives in gpu_kernels · 10 (kernel first principles) and the roofline in cs336 · 07 (kernels & FlashAttention); here we lift it from one op to a region of the graph.
2 · Vertical fusion: the producer-consumer chain
Vertical fusion (also called producer-consumer fusion) merges a chain where each op consumes the output of the one before it. A producer writes a value; its consumer immediately reads it. If the consumer is the producer's only user, the producer never needs to write its result to HBM at all — the value can stay in a register and flow straight into the next computation.
Take the canonical residual epilogue norm(dropout(bias(x))) on a tensor of M bytes. Run eagerly, that is three nodes and three round trips. The arithmetic, op by op:
| execution | HBM reads | HBM writes | launches | total traffic |
|---|---|---|---|---|
| eager: bias | read x (M) | write t1 (M) | 1 | 2M |
| eager: dropout | read t1 (M) | write t2 (M) | 1 | 2M |
| eager: norm | read t2 (M) | write y (M) | 1 | 2M |
| eager total | 3M | 3M | 3 | ≈6M |
| fused (one kernel) | read x (M) | write y (M) | 1 | ≈2M |
Fusion rewrites the three nodes as one kernel: load each tile of x from HBM into registers/SRAM once, apply the bias, the dropout, and the normalize while the data is hot, write the final y once. Three launches collapse to one; ≈6M bytes of traffic collapse to ≈2M. For a memory-bound chain that is a ≈3× speedup, and it buys it without doing one FLOP less — the math is identical, only the byte movement changed. (The intermediates t1 and t2 simply cease to exist as HBM tensors; hold that thought, it is exactly what forces lesson 06.) This is the same ≈6×→2× collapse cs336/07 derives for dropout(bias(x)) → norm; the deeper kernel-level walkthrough is in gpu_kernels · 19 (anatomy of a fused kernel), where a fused RMSNorm reads x twice instead of five separate kernels reading it three or four times.
The general statement: a memory-bound chain of n pointwise ops costs ≈2n · M bytes eagerly (each op a read and a write) and ≈2M fused — a saving that grows with chain length. Bias+activation, residual-add+norm, scale+shift+cast — these short chains are the bread and butter of fusion, and they appear by the hundred in a transformer.
3 · Horizontal, epilogue, and the algorithm-level limit
Horizontal fusion merges sibling ops that share an input but do not feed each other — for example three separate projections that all read the same hidden state, or the query/key/value matmuls reading one residual stream. Run separately, each reloads the shared input from HBM; fused, the input is loaded once and the siblings run side by side over it. The win here is the read that would otherwise be paid three times, plus three launches becoming one. Inductor does this when it sees independent consumers of one buffer; XLA's clustering will pull siblings into the same fusion node.
Epilogue fusion is the highest-value structured case. A GEMM (matmul) is compute-bound — it already keeps its output tile hot in registers as it accumulates over the K dimension. The bias-add, the GELU, the residual-add, the cast that follow the matmul are each memory-bound on their own. Epilogue fusion does them while the matmul's output tile is still in registers, before it is ever written to HBM. The pointwise tail rides for free on a trip the matmul was going to make anyway: zero extra round trips for the activation. cublasLt exposes this as a fused epilogue; Inductor and XLA generate it directly. The deeper kernel-level template — prologue, load, on-chip math, epilogue — is exactly the structure in gpu_kernels · 19.
The headline case is FlashAttention. Naive attention materializes the full (B·h·T·T) score matrix in HBM — O(T²) memory and traffic, OOM at modest context. FlashAttention fuses scores → softmax → value-weighting into one kernel using the online-softmax recurrence, so the scores live only in SRAM and are never written down, dropping attention memory to O(T). It is fusion that required inventing a new algorithm, not just merging existing nodes, which is why it gets its own derivation. Do not re-derive it here — the full online-softmax recurrence, the tiling, and the naive-vs-flash complexity table are in cs336 · 07 (kernels and FlashAttention). For this track, the point is structural: fusion ranges from a trivial bias+GELU merge all the way up to FlashAttention, and the same currency — round trips saved — measures both.
4 · The fusion decision: when to stop
Fusion is not free past a point, and "fuse everything" is a real and common mistake. A fused kernel's working set — the inputs, the live intermediates, the accumulators it keeps on-chip — must fit in the SM's finite resources: ≈64 K 32-bit registers per SM and ≈228 KB of shared memory on an H100, both shared across all resident blocks. Register pressure is the term for how close a kernel comes to that limit. Fuse too long a chain, or fuse ops with large per-thread state, and the compiler runs out of registers and spills them to local memory (which is backed by HBM) — silently reintroducing the very traffic fusion was meant to remove, and crushing occupancy because fewer blocks fit per SM. The fused kernel ends up slower than the unfused chain.
Three boundaries tell the compiler where to cut a fusion group (the set of nodes assigned to one kernel):
| boundary | why it stops fusion | signal if you ignore it |
|---|---|---|
| register / SMEM pressure | Working set exceeds the SM's registers or shared memory; the compiler spills to HBM-backed local memory and occupancy collapses. | Register-spill warnings; a "fused" kernel no faster (or slower) than the chain it replaced. |
| compute-bound op (big GEMM) | A matmul is already FLOP-limited; fusing its output into a downstream consumer can force the consumer onto the GEMM's tiling and starve it. Fuse the epilogue into the GEMM, not the GEMM into something else. | The matmul kernel still dominates the profile; fusing around it changed nothing or regressed. |
| reduction → consumer with reuse | If a value is read by many downstream ops, fusing it into all of them recomputes it once per consumer. When the recompute cost exceeds the saved HBM round trip, materialize it once instead. | FLOP count balloons; an op that was memory-bound becomes needlessly compute-bound from redundant recompute. |
That third boundary is the subtle one. Fusion's whole trick — keep the value in a register instead of HBM — assumes the value has one consumer. When a producer feeds several consumers and you fuse it into each, you compute it once per consumer instead of once total: you traded a saved read for duplicated FLOPs. Sometimes that is a win (recompute is cheap, the op was memory-bound anyway — this is exactly FlashAttention's backward recompute trade). Sometimes it is a loss. The compiler weighs bytes saved by not materializing against FLOPs added by recomputing, which is why fusion and the recompute/checkpointing decision (lesson 12) are the same knob seen from two sides. And the rule that bounds all of it is the track's first motif: every fusion must preserve semantics — the fused kernel must compute bit-for-equivalent results (mind float non-associativity when fusion reorders a reduction).
5 · How real compilers group
None of this is hand-applied per kernel — it is a pass. Each stack does fusion grouping a little differently, but all answer the same question: which adjacent nodes can share a kernel without breaking resources or semantics?
kFusion) per cluster that lowers to a single kernel. Epilogue and elementwise fusion are the bulk of its wins.compute_at: you (or the auto-scheduler) place a producer's computation inside the consumer's loop nest, so the intermediate is consumed in-place rather than written out. Relay's fusion pass picks fusable regions before scheduling.The artifact at the end is the same in every case: a node-level graph annotated with fusion groups, each group destined to become exactly one kernel. The pass that draws those group boundaries — deciding which round trips to eliminate and where pressure forces a cut — is the highest-leverage pass in the compiler. Drive it yourself below.
6 · Drive it: draw the fusion boundaries
The widget is the canonical block matmul → bias → gelu → add → norm. Click the cut buttons on the edges between ops to split the chain into fusion groups; every contiguous run of ops with no cut between them becomes one kernel. Watch the three numbers the fusion pass actually optimizes — #kernels, HBM bytes, peak live buffers — and the register-pressure meter. Start with all cuts in (pure eager, one kernel per op) and remove them to fuse. Find the sweet spot. Then fuse everything into one group and watch the meter go red as the working set spills.
The shape of the result is the whole lesson: traffic falls steeply as you merge the first few memory-bound ops, the matmul's epilogue absorbs bias+gelu+add for almost nothing, and then — past the sweet spot — pressure climbs until one over-stuffed group spills and the curve bends back up. Fusion is the central win, but it is a win with a ceiling, and the ceiling is on-chip resources.
Failure modes & checklist
Failure modes
- Fusing a compute-bound op to "save a kernel." A big GEMM is FLOP-limited; fusing the matmul into a consumer can starve its tiling. Signal: you fused everything and the profiler still shows the matmul dominating — it was never bandwidth-bound, so fusion bought nothing.
- Over-fusion → register spill. One group's working set exceeds the SM's registers/SMEM, spilling to HBM-backed local memory and dropping occupancy. Signal: register-spill warnings; the "fused" kernel is no faster, or slower, than the chain.
- Fusing a multi-consumer producer into all consumers. Recomputes the value once per consumer; FLOPs balloon. Signal: a memory-bound op becomes compute-bound and total FLOPs jumped with no traffic win.
- Reordering a reduction across the float-associativity line. Fusing can change the summation order; results drift. Signal: bit-level mismatches vs eager that grow with reduction length — preserve-semantics was violated.
- Forgetting the freed intermediates still need a plan. Fusion deletes HBM tensors but creates register/scratch pressure and changes liveness. Signal: peak memory didn't drop as expected because allocation wasn't re-planned (lesson 06).
Checklist
- Profile first. Fuse memory-bound chains (intensity below the ridge). Don't fuse to "save kernels" on compute-bound ops.
- Count round trips, not ops. A memory-bound chain's time ≈ bytes / bandwidth; the goal is one read + one write per group.
- Prefer epilogue fusion into the GEMM. Fold bias/activation/cast onto the matmul tile while it's hot — near-zero extra traffic.
- Watch the pressure meter. Stop fusing before the working set spills; bigger groups aren't always better.
- Weigh recompute vs reuse. Fuse a multi-consumer producer only when recompute FLOPs cost less than the saved round trip.
- Re-plan memory after fusing. Fusion changes which buffers exist and how long they live — hand off to liveness/allocation (lesson 06).
Checkpoint
bias, gelu, and add so they fuse into the matmul's epilogue: watch #kernels drop and traffic fall sharply while the pressure meter stays green. (3) Hit Fuse everything: the meter goes red and the verdict flips to a spill — confirm HBM bytes tick up from the sweet spot because spilled registers go back to HBM. (4) Predict before sliding: does doubling the tensor size change the number of kernels at a fixed grouping? (No — grouping sets kernel count; size scales the bytes each kernel moves.) Find the grouping with the lowest HBM bytes that stays green; that is what a real fusion pass is searching for.Where this points next
Fusion is the central win, and it just broke your memory picture on purpose. Before this pass, every node in the graph had a clear HBM buffer with an obvious lifetime: allocate on write, free after the last read. After fusion, the intermediates inside a group — the t1, t2 of section 2 — no longer exist as HBM tensors at all; they are registers that live for a few instructions. Meanwhile the groups that do write outputs now have lifetimes that overlap in new ways, and an over-fused group quietly demands more on-chip scratch. The question "how much memory does this graph need, and which buffers can share storage?" no longer has a node-by-node answer — fusion rewrote it. You cannot allocate naively (peak = sum of every buffer) and you cannot ignore the lifetimes fusion created. That forces a deliberate compile-time pass over buffer liveness and reuse: 06 · Memory planning.
dropout(bias(x)) → norm goes from ≈6× the tensor's bytes to ≈2×, a ≈3× speedup with the math unchanged. The kinds — vertical/producer-consumer (elementwise chains), horizontal (siblings sharing an input), epilogue (fold pointwise work onto a hot GEMM tile for ≈0 extra traffic), and the algorithm-level FlashAttention (cs336/07) — all spend the same currency, round trips eliminated. But fusion has a ceiling: stop before a group's working set blows register/shared-memory pressure and spills, before you fuse around a compute-bound GEMM, and before recomputing a multi-consumer value costs more than the round trip it saves — always preserving semantics. Real stacks ship this as a grouping pass (Inductor's pointwise+reduction Triton kernels, XLA fusion clustering, TVM compute_at). And because fusion erases some intermediate buffers while reshaping the lifetimes of the rest, it forces the next pass: memory planning.Interview prompts
- Why does fusion give a ≈3× win on
dropout(bias(x)) → normbut barely help a big matmul? (§1–2 — the chain is memory-bound, so time ≈ bytes/bandwidth; fusion cuts ≈6M→2M traffic. A GEMM is compute-bound — its time is FLOPs, and the epilogue's traffic is tiny relative to the matmul, so merging saves little.) - Distinguish vertical, horizontal, and epilogue fusion with an example of each. (§2–3 — vertical: producer→consumer chain (bias→gelu) kept in registers; horizontal: siblings sharing an input (Q,K,V from the residual) loaded once; epilogue: bias+activation done on the matmul's output tile while it's still hot.)
- What is register pressure and how does it bound fusion? (§4 — how close a kernel comes to the SM's finite registers/SMEM; over-fuse and the compiler spills registers to HBM-backed local memory and occupancy collapses, so the fused kernel can be slower than the chain. It's the ceiling on how much you can merge.)
- When should a compiler NOT fuse a producer into its consumers? (§4 — when the producer feeds multiple consumers and fusing recomputes it once per consumer; if the added recompute FLOPs exceed the saved HBM round trip, materialize it once instead. Same knob as recompute/checkpointing.)
- How does Inductor decide fusion groups, and what does it emit? (§5 — it classifies nodes as pointwise/reduction, greedily groups compatible adjacent ops (chains, pointwise+reduction, matmul epilogues, horizontal pointwise siblings), and emits one Triton kernel per group.)
- Is FlashAttention "just" operator fusion? (§3 — it's algorithm-level fusion: it required the online-softmax recurrence so QKᵀ→softmax→·V fuse without ever materializing the O(T²) scores in HBM. Peephole node-merging can't do that; see cs336/07.)
- Why does fusion force a memory-planning pass next? (§Where-next — it deletes intermediate HBM buffers (now registers) and reshapes the lifetimes of the rest, so naive per-node allocation (peak = sum) is wrong; liveness and buffer reuse must be re-planned (lesson 06).)