Part III — Middle end

Memory planning

Lesson 05 fused chains of ops to cut HBM traffic — but fusion, and the graph itself, leave behind a litter of intermediate buffers: every value flowing along an edge needs storage somewhere. Hand each one its own allocation, free it when the op finishes, and you do the obvious thing — which is also the wrong thing. The naive plan makes peak memory the sum of every buffer that ever existed and drowns the runtime in per-op allocator calls. This lesson treats deciding where each buffer lives as a compile-time problem: figure out exactly when each value is alive, then let buffers that are never alive at the same time share the same bytes.

The previous step left this broken

Fusion (lesson 05) is a transform on the graph; it decided which ops become one kernel, but it said nothing about the bytes those kernels read and write. The optimized graph is still a dataflow DAG, and every edge carries a tensor that must occupy physical memory at the moment it is produced. Run the naive policy — allocate a fresh buffer for each value, free it after its last consumer — and two costs explode. First, peak memory: a careless allocator can hold dozens of intermediates live at once and your peak is the sum of all of them, even though most never overlap in time. Second, allocator traffic: thousands of cudaMalloc/cudaFree calls per step, each a synchronizing, fragmenting, microsecond-scale tax. The model fits the chip in FLOPs and runs fast in kernels, and still OOMs or stalls because nobody decided where the bytes go.

Linear position

Forced by: a graph of intermediate buffers, allocated naively, makes peak memory the sum of all buffers and floods the runtime with allocator calls.
New idea: memory planning as a compile-time pass — run liveness analysis to learn each buffer's lifetime (first def to last use), build the interference graph of buffers that are live at the same time, then assign storage so non-overlapping buffers share a slot (register-coloring style), do ops in place when safe, donate an input's storage to an output, and serve everything from one preallocated arena. Goal: collapse peak from sum-of-buffers toward max-concurrent, and cut allocator calls to ≈1.
Forces next: the plan decides where and when a buffer lives, but not its physical shape in memory — and the wrong byte layout makes every downstream kernel slow. That is layout (lesson 07).

The plan

Five moves. (1) Price the naive policy: per-tensor malloc/free, fragmentation, and why peak = sum, not max. (2) Define liveness — first-def to last-use — and read it off a worked graph with byte counts. (3) Reuse by graph coloring; the interference graph; in-place ops. (4) Buffer donation/aliasing (XLA input-output aliasing, JAX donated args) and the correctness guard that makes it legal. (5) Arena/static allocation, and how training (recompute, checkpointing) reshapes the plan — the activation-memory wall. Then drive the planner widget until peak collapses from sum toward max-concurrent.

1 · The naive policy and why peak = sum

The simplest allocator is "allocate on produce, free on last consume." It is correct and it is what eager mode does. It has two failure modes that a compiler, seeing the whole graph at once, can fix.

Allocator traffic. Every intermediate is a real cudaMalloc and a real cudaFree. A device malloc is not free: it can synchronize the stream, it fragments the heap, and it costs on the order of a few microseconds. A transformer layer fused down to, say, 20 kernels can still produce 30–50 live intermediates; across 80 layers that is thousands of allocator calls per forward step. Even at ≈2 µs each, a few thousand calls is several milliseconds of pure bookkeeping per step — wasted next to the kernels you worked so hard to fuse.

Peak memory. This is the expensive one. Define peak memory as the maximum, over the whole execution, of the total bytes live at any one instant. The naive policy does not control it; it merely reacts to it. Worse, a naive allocator that does not reuse freed blocks well makes your high-water mark drift toward the sum of every buffer that was ever allocated, because freed-but-not-recycled slabs sit in the heap fragmenting it. The whole point of planning is to recognize that two buffers whose lifetimes never overlap can occupy the same bytes — so the real lower bound on peak is the maximum concurrent live bytes, not the sum. The gap between those two numbers is what this pass recovers.

peak_naive ≈ Σ_{all buffers} size(b) ⟶ peak_planned = max_t Σ_{b live at t} size(b)

2 · Liveness: when is each buffer alive?

Liveness analysis is the classic compiler dataflow question, lifted to tensors: a buffer is live from the instant it is defined (the op that produces it) to the instant of its last use (the last op that reads it). Before its def and after its last use, its bytes are free for anyone. In SSA / value-semantics IR (lesson 03), this is trivial to read off — each value is defined exactly once, and you walk the DAG in execution order to find the last consumer of each edge.

Work a concrete chain. Take the fused residual block from lesson 05, executed as a linear schedule of six ops, each producing one buffer on a tensor of 64 MB (call that one unit, U = 64 MB):

step   op                produces   reads          buffer lifetime  (• live)
  0    x   = input        x  (64MB)  —              x : 0 ──────────► 3
  1    a   = matmul(x,W)   a  (64MB)  x             a : 1 ──► 2
  2    b   = bias(a)       b  (64MB)  a             b : 2 ──► 3
  3    c   = gelu(b)+x     c  (64MB)  b, x          c : 3 ──────────► 5
  4    d   = norm(c)       d  (64MB)  c             d : 4 ──► 5
  5    y   = matmul(d,W2)  y  (64MB)  d             y : 5 ─►

Naive accounting: six buffers ever produced, so a sloppy allocator drifts toward 6 × 64 = 384 MB. Now read the concurrent live set step by step. At step 3 the live buffers are {x, b, c} — x is still needed for the residual add, b is being consumed, c is being produced — that is 3 × 64 = 192 MB, the busiest instant. Everywhere else only two are live. So the true peak is 192 MB, not 384 MB: planning can halve it just by recognizing that a dies before d is born, so they can share bytes. The naive number counts buffers that were never alive together.

3 · Reuse by coloring, and in-place ops

Once you have lifetimes, assigning storage is exactly graph coloring — the same algorithm a CPU compiler uses for register allocation, scaled up from registers to tensor-sized slots. Build the interference graph: one node per buffer, an edge between two buffers whose lifetimes overlap. Two buffers that interfere must get different storage; two that do not (their lifetimes are disjoint) may share. A valid storage assignment is a coloring where no edge connects two same-colored nodes, and the number of colors you need — weighted by size — sets your peak.

interference

Two buffers interfere if their lifetimes overlap at any instant. Interfering buffers cannot share storage; non-interfering ones can. Same relation a register allocator colors.

reuse / coloring

Assign each buffer to a slot so no two interfering buffers share a slot. Greedy best-fit over a lifetime sweep is the practical heuristic; optimal slot-packing is NP-hard, so compilers approximate.

in-place

If an op's only remaining consumer of its input is the op itself, the output can overwrite the input's buffer — zero extra bytes. A degenerate, free case of reuse.

In the worked chain above, a (lives 1→2) and d (lives 4→5) do not interfere — a is dead by the time d is born — so they get the same slot. Likewise b and y. After coloring you need only three physical slots of 64 MB, which is 192 MB — exactly the max-concurrent bound. The allocator went from six allocations to three slots carved once.

In-place is the cheapest reuse of all. An elementwise op like relu or add_ can write its result over its input buffer if that input has no other live consumer — the bytes are recycled with zero new storage. Fusion (lesson 05) creates lots of these opportunities, because a fused epilogue's intermediate scalars never need their own HBM buffer at all. The catch is the only remaining consumer test: overwrite a buffer that something later still reads and you have silently corrupted the math. That guard is the recurring preserve-semantics motif — every reuse decision is legal only if you can prove no live value is clobbered.

4 · Donation and aliasing: reuse an input's storage

The graph's inputs are buffers too, and often the caller does not need them after the call. Buffer donation (a.k.a. input-output aliasing) lets the output of the program reuse the storage of an input that is dead afterward — the single biggest win for the parameter-update step, where a new weight tensor can be written straight over the old one instead of doubling parameter memory.

XLA input-output aliasing

XLA lets you mark that output i may alias input j; the compiled executable writes the result in place over the donated input's buffer. Standard for the optimizer's params ← params − lr·grad update.

JAX donated args

jax.jit(f, donate_argnums=(0,)) promises the caller will not use argument 0 again, so XLA may reuse its buffer for an output. Donating the params + optimizer state can roughly halve update-step memory.

the correctness guard

Donation is legal only if the donated buffer has no later use anywhere — the host must not read it after the call. Violate it and you read overwritten garbage; frameworks therefore make donation an explicit, opt-in promise.

The mechanism and the guard are the same as in-place, one level up: aliasing input j to output i is sound exactly when j is provably dead after the call. Because the compiler cannot always see the host's later use, donation is an explicit contract: in JAX you pass donate_argnums, in XLA you set the aliasing config, and you accept that touching the donated value afterward is undefined. Done right on a 7B model's update step, donating params and the two Adam moments can save tens of gigabytes — the difference between fitting and not.

5 · Arena allocation, and the training wall

The final move turns the plan into one allocation. Once coloring has assigned every buffer to a slot and sized the slots, the compiler computes the offset of each slot inside a single contiguous slab — the arena — preallocates it once at the start, and serves every "allocation" as a pointer into the arena at a known offset. No per-op malloc, no fragmentation, no synchronization: allocator traffic drops from thousands of calls to one. This is what XLA, TensorRT, and TVM's graph executor all do — the buffer assignment is baked into the compiled artifact.

1. livenessWalk the schedule; record each buffer's [def, last-use] interval.

2. interferenceEdge between buffers whose intervals overlap; in-place / donation are special dead-input cases.

3. color + size slotsGreedy best-fit assigns non-interfering buffers to shared slots; slot size = max buffer assigned to it.

4. arenaLay slots out at fixed offsets in one slab; preallocate once; every buffer is now a known offset.

Training breaks the tidy picture and is where planning earns its keep. The backward pass needs the forward activations, so every activation produced in the forward is live until its gradient is computed — lifetimes stretch across the whole forward and backward, and the concurrent live set balloons. This is the activation-memory wall: for a deep model, saved activations, not parameters, dominate peak memory (see cs336 · 05 the memory ledger for the per-parameter accounting). The compiler's lever here is rematerialization — drop some activations and recompute them in the backward, trading FLOPs for bytes. That makes save-vs-recompute a memory-planning decision the partitioner makes (a min-cut on the lifetime graph), which lesson 12 develops fully when autodiff builds the joint forward+backward graph. For now the point stands: liveness + coloring + arena collapse peak from sum toward max-concurrent, and recompute pushes max-concurrent itself down.

6 · Drive it: collapse peak from sum to max-concurrent

The widget lays out a fixed schedule of buffers as horizontal bars on a timeline — each bar spans a buffer's lifetime from def to last use. The vertical axis is storage slots. Turn on reuse (color) and watch non-overlapping bars drop into a shared slot while overlapping bars are forced apart; turn on in-place and watch elementwise outputs fold onto their inputs for free. The KPIs track the only numbers that matter: peak bytes, number of allocations, and bytes saved versus naive.

Memory planner: liveness, reuse, in-place

Each bar is a buffer's lifetime. Naive gives every buffer its own slot — peak = sum. Turn on reuse and buffers whose lifetimes don't overlap share a slot (peak → max-concurrent); turn on in-place to fold elementwise outputs onto their inputs. The knob that proves it: with everything off, peak is the tall sum; flip both on and peak collapses to the single busiest instant.

reuse (color non-overlapping) in-place elementwise buffer size: 64 MB

Peak live bytes

—

Slots allocated

—

Bytes saved vs naive

—

Max-concurrent bound

—

buffer lifetime bar in-place (folded onto input) slot boundary mode: naive

The lesson is the gap between the two horizontal bands. With reuse off, each buffer claims its own slot and the stack is as tall as the buffer count. With reuse on, the bars compact into the smallest number of slots that respects interference — and that number is fixed by the single busiest column of the timeline, the max-concurrent live set. No transform changed the math; the planner only noticed which buffers were never alive together.

Failure modes & checklist

Failure modes

In-place over a still-live input. Overwriting a buffer that a later op (or a residual branch) still reads. Signal: numerically wrong outputs that change when you toggle the in-place pass — the math depends on a clobbered value.
Donating an input the host reuses. Marking an argument donated, then reading it after the call. Signal: garbage in the supposedly-untouched input; nondeterministic results across runs.
Counting sum instead of max. Reporting peak as the total of all intermediates and concluding the model can't fit. Signal: your estimate is 2–4× the profiler's actual peak — you summed buffers that never coexisted.
Ignoring the activation wall in training. Planning only the forward and being shocked by OOM in the backward. Signal: forward fits comfortably, backward OOMs — activations stayed live across the whole pass.
Over-fusing into giant live buffers. Fusion that keeps a huge intermediate alive longer can raise max-concurrent. Signal: fewer kernels but higher peak — fusion and planning fighting each other.

Checklist

Compute liveness first. [def, last-use] per buffer off the SSA schedule; everything else depends on it.
Peak = max-concurrent. Size the arena to the busiest instant, not the sum of buffers.
Color non-overlapping buffers together. Greedy best-fit over a lifetime sweep; share slots where intervals are disjoint.
Prove dead-before-reuse. In-place and donation are legal only when the donor buffer has no later use.
Donate the update step. Alias params/optimizer-state outputs to their inputs (XLA aliasing / JAX donate_argnums).
Allocate one arena. Bake fixed offsets into the artifact; drive allocator calls toward one.

Checkpoint

Try it

Open the widget with everything off and read the peak: six buffers at 64 MB each in their own slots, ≈384 MB, with the naive band as tall as the buffer count. Now flip reuse on and watch the stack compress to three slots (192 MB) — that is the max-concurrent bound, reached at the busiest column where {x, b, c} are live together. Confirm the "bytes saved" KPI reads ≈192 MB (the sum-minus-max gap). Then flip in-place on and watch one more slot fold away as the elementwise output overwrites its input. Finally, set buffer size to 128 MB and predict the new peak before you read it: every number doubles, because reuse changes how many slots, not how big each one is.

Where this points next

You now decide where each buffer lives and when its bytes get recycled, so peak memory tracks the busiest instant instead of the sum, and the runtime makes one allocation instead of thousands. But the plan has been treating a buffer as an opaque blob of N bytes. It is not — a tensor has a physical layout: the order its dimensions are stored in memory, whether it is row- or column-major, padded, or tiled to match tensor-core fragments. Two plans with identical byte budgets can be worlds apart in speed, because the wrong layout forces every consuming kernel into strided, uncoalesced access or inserts transposes the planner never accounted for. Deciding each buffer's in-memory shape — and propagating that choice through the graph to cancel transposes — is the next forced step: lesson 07, Layout & data formats.

Takeaway

Fusion left a graph dense with intermediate buffers; allocating each naively makes peak memory the sum of every buffer that ever existed and floods the runtime with allocator calls. Memory planning fixes both at compile time. Liveness analysis gives each buffer a lifetime — first def to last use — and buffers whose lifetimes never overlap don't interfere, so they can share storage. Assigning storage is then graph coloring, the same algorithm as register allocation: the worked six-op chain drops from ≈384 MB (sum) to 192 MB, the max-concurrent bound set by the busiest instant where {x, b, c} coexist. In-place ops recycle an input's bytes when it has no other live consumer; donation/aliasing (XLA input-output aliasing, JAX donate_argnums) does the same for program inputs — both legal only under the no-later-use guard that is the preserve-semantics motif. An arena bakes the whole assignment into one preallocated slab, cutting allocator calls to ≈1. In training the backward pass keeps activations live across the whole pass — the activation-memory wall — so recompute becomes a planning knob (lesson 12). Net: peak collapses from sum toward max-concurrent. What planning still doesn't fix is each buffer's physical shape — the layout — which forces lesson 07.

Interview prompts

Why does naive per-tensor allocation make peak memory the sum of buffers, and what's the real lower bound? (§1 — a naive allocator can hold non-overlapping buffers in distinct, never-recycled slabs, so the high-water mark drifts toward Σ size(b); the true lower bound is max over time of the bytes live at one instant — max-concurrent — and the gap is what planning recovers.)
Define liveness for a tensor buffer and how you'd compute it. (§2 — a buffer is live from its defining op to its last consuming op; in SSA / value-semantics IR walk the schedule in execution order and record [def, last-use] per edge.)
How is buffer reuse the same problem as register allocation? (§3 — build an interference graph with an edge between buffers whose lifetimes overlap, then color so no two interfering buffers share a slot; it's graph coloring scaled from registers to tensor-sized slots, optimal packing NP-hard so greedy best-fit is used.)
When is an in-place op legal, and how does it differ from donation? (§3–4 — in-place is legal when the op's input has no other live consumer, so the output overwrites it; donation/aliasing is the same idea for a program input the caller won't reuse — both require the donor buffer be provably dead afterward.)
What are XLA input-output aliasing and JAX donated args, and what's the correctness guard? (§4 — they let an output reuse a dead input's buffer — e.g. the optimizer's in-place params update — legal only if the donated buffer has no later use; frameworks make it an explicit opt-in promise because the compiler can't see the host's later reads.)
Why does training blow up peak memory more than inference, and what's the compiler's lever? (§5 — the backward needs forward activations, so their lifetimes stretch across the whole pass and max-concurrent balloons — the activation-memory wall; the lever is rematerialization, dropping and recomputing activations to trade FLOPs for bytes, decided by a min-cut partitioner.)
You fused more aggressively and peak memory went up. How? (§Failure modes — fusion that keeps a large intermediate alive longer raises the max-concurrent live set; fusion and memory planning can fight, so the planner must price the lifetime cost of a fusion choice.)