Part V — Alignment & inference
Inference — serving the model you paid for
Lesson 15 finished the model: a base net pretrained on a budget, SFT'd into an assistant, then nudged toward human preferences under a KL anchor. The weights are frozen and good. But a frontier model nobody can afford to run is a museum piece. The economics of generating a token are nothing like the economics of training one, and the chip you optimized for training (lesson 6) is now mostly idle — starved on memory bandwidth, not FLOPs. This lesson derives why, and the four levers that get your tokens-per-dollar back.
New idea: the prefill (compute-bound) vs decode (memory-bound) asymmetry; the KV cache as the central object and the real cap on batch/context (size = 2·L·nkv·dhead·T·batch·bytes, shrunk by GQA — payback of lesson 3); and the four levers — continuous batching, speculative decoding, weight/KV quantization — that trade off latency, throughput, and quality.
Forces next: nothing new — this is the last build stage. Lesson 19 walks the whole budget end to end, every decision a forced consequence of the last.
1 · Two phases: prefill is compute-bound, decode is memory-bound
A single request is two completely different workloads stitched together. Naming them is the whole lesson.
Prefill processes the entire prompt in one parallel forward pass. If the prompt is P tokens, you push a (1, P, d) tensor through the model once: every position is computed simultaneously, the matmuls are big and dense, and the GPU's tensor cores run near peak. Prefill's job is to compute the logits for the last prompt position (the first generated token) and — crucially — to populate the KV cache: the key and value vectors for all P positions, at every layer, so future steps never recompute them. Prefill is compute-bound: it does ≈ 2N·P FLOPs against a one-time weight load, so its arithmetic intensity is high — it sits on the right side of the lesson-06 roofline.
Decode generates the rest, one token per step. At step t you feed a single token — a (batch, 1, d) tensor — through the model. To do that you must load every weight matrix from HBM once (a 14 GB model = 14 GB of traffic) and read the entire KV cache, but you do only ≈ 2N FLOPs per sequence. With batch size 1 the arithmetic intensity is roughly 2N / (2N) = 1 FLOP per byte — catastrophically below the H100 ridge point of ~300 FLOP/byte (lesson 6). The tensor cores idle while the memory bus is pinned. Decode is memory-bound, and decode is where essentially all the wall-clock and the cost live — a 500-token answer is 1 prefill pass and 499 decode steps.
| Prefill | Decode | |
|---|---|---|
| Input shape | (1, P, d) — whole prompt at once | (batch, 1, d) — one token / step |
| FLOPs | ≈ 2N·P (large) | ≈ 2N per sequence per step (tiny) |
| Bytes moved / step | weights once, amortized over P | all weights + KV cache, every step |
| Arithmetic intensity | high → compute-bound | ≈ 1 (batch 1) → memory-bound |
| Bottleneck | tensor-core FLOP/s | HBM bandwidth (TB/s) |
| Runs | once per request | once per output token |
2 · The KV cache: the object that actually caps you
Why cache K and V at all? Because attention at step t needs the keys and values of every prior token. Recomputing them each step would make decode O(T²) in compute. Instead you store them once and append one column per step. The cache is what makes decode O(T) per step — but it is also the thing that fills up your GPU. Derive its size:
The 2 is one tensor for K, one for V; L layers; nkv key/value heads (this is where GQA bites — see below); dhead per-head width; T cached positions; batch concurrent sequences; byteselt = 2 for bf16, 1 for INT8. Plug in a 7B-class model — the exact config from lesson 3: L = 32, dhead = 128, bf16. With full multi-head attention (nkv = h = 32) one token costs 2 · 32 · 32 · 128 · 2 = 524,288 bytes = 512 KB. So:
Read that again: a single 32K-context request eats more memory than the model weights, and the KV cache — not the FLOPs, not the parameters — is what decides how many requests you can serve concurrently. Batch size is capped by (HBM − weights) / (KV per sequence). That is the bottleneck the entire serving stack is built around. And because the cache scales as ∝ T, at the 32K–128K windows of lesson 15 (Long context) the KV cache — not the weights — becomes the dominant serving cost, which is exactly why GQA/MLA plus a paged KV cache are what make long-context serving affordable at all.
GQA pays back lesson 3
Lesson 03 introduced grouped-query attention as a swap with no training-time motivation — "its payoff lives in lesson 18." Here is the payoff. GQA shares each K/V head across a group of g query heads, so nkv = h / g instead of h. The cache divides by exactly g. For the same 7B-class model at g = 4 (nkv = 8): 2 · 32 · 8 · 128 · 2 = 131,072 bytes = 128 KB per token, and the 32K sequence drops from ~17.2 GB to ~4.3 GB. Same query heads, same quality to within a hair — 4× more concurrent users. GQA is a serving decision disguised as an architecture decision, and the multiply on your throughput is the group factor g.
MLA: compressing the cache
GQA shrinks the cache by sharing — fewer K/V heads to store. The frontier moves down the very same axis but with a different verb: compress. Multi-head Latent Attention (MLA), the attention variant in DeepSeek-V2 and V3, does not store per-head K and V at all. It projects each token down to a single low-rank latent vector ct and caches only that; at attention time it projects ct back up into the per-head keys and values on the fly. The slogan for the whole section: GQA caches fewer heads; MLA caches a smaller thing per head.
The win is that the cached latent is far narrower than even the nkv shared heads of GQA, so the per-token cache is smaller again — DeepSeek reports a cache a small fraction of MHA's, below GQA at comparable quality. The price is extra compute every step: the up-projection from ct to K and V is matmul work that GQA simply skips. But recall §1 — decode is bandwidth-bound, not compute-bound: the tensor cores are idle while the memory bus is pinned. Spending spare FLOPs to move fewer bytes is exactly the trade the whole stack wants, so the projection is close to free in wall-clock while the smaller cache buys real concurrency. MLA is GQA's logic — make the thing you reload each step smaller — pushed from "store less" to "store a compressed code."
This matters most at long context. The cache grows as ∝ T, so the longer the window the more the KV cache dominates the memory budget and the more a smaller per-token footprint is worth — see lesson 03 (GQA) for the architecture and lesson 15 (Long context) for why T-scaling is the regime where MLA earns its keep.
The cache also grows one token per step per sequence, so storing it as one contiguous block per request wastes huge amounts to fragmentation. PagedAttention (vLLM) borrows the OS virtual-memory trick — chop the cache into fixed-size blocks behind a page table, so memory packs tightly and prompts can share prefix blocks — and radix-tree prefix sharing (SGLang) goes further; see gpu_kernels/14 · paged & radix layouts and the vLLM track.
3 · Batching: static vs continuous, and the throughput/latency trade
Decode at batch 1 wastes the GPU: you load 14 GB of weights to update one sequence. Load them once and update B sequences in the same pass and the weight traffic is shared — arithmetic intensity rises ~B×, throughput (tokens/s across all users) rises with it, right up until you re-cross the ridge into compute-bound or run out of KV memory. Batching is the single biggest throughput lever in decode. But how you batch matters enormously.
Static batching groups B requests, runs them lockstep, and the whole batch finishes when the longest sequence does. A request that wants 20 tokens sitting next to one that wants 2,000 wastes 1,980 steps of a GPU slot it cannot release, and new arrivals wait for the entire batch to drain. Utilization craters under realistic, variable-length traffic.
Continuous (in-flight) batching works at the granularity of a single decode step: the moment any sequence emits its end-of-sequence token, it leaves the batch and a queued request takes its slot on the very next step. The batch composition changes every iteration. This is the headline scheduler trick in vLLM and TGI, and it routinely doubles-to-triples throughput over static batching on real workloads with no quality change. It is also what makes prefill/decode scheduling interesting — a fresh request needs a (compute-bound) prefill spliced in among the (memory-bound) decodes, and the scheduler chooses chunk sizes to keep latency smooth.
| Static batching | Continuous batching | |
|---|---|---|
| Batch changes | only between batches | every decode step |
| A finished sequence | holds its slot until the longest finishes | frees its slot immediately |
| GPU utilization (variable lengths) | low — padded to the max | high — slots refilled continuously |
| Throughput on real traffic | baseline | ~2–3× baseline |
The catch is universal and worth stating as a law: throughput and latency trade against each other. Bigger batches amortize the weight load over more users, so tokens/s/GPU climbs — but any individual user's tokens/s can fall, because their sequence now shares bandwidth and may queue behind a prefill. Interactive chat optimizes for low per-user latency (small batches, fast first token); a batch summarization job optimizes for throughput (huge batches, who cares about latency). Same model, opposite serving config. There is no single "fast."
4 · Speculative decoding: guess cheap, verify in bulk
Decode is bandwidth-bound, which means the GPU has FLOPs to spare — verifying several tokens at once costs almost the same as verifying one (you still load the weights once). Speculative decoding exploits exactly that slack. A small, cheap draft model autoregressively proposes the next k tokens (it is fast precisely because it is small). The big target model then runs one forward pass over all k proposals in parallel — like a mini-prefill — and checks each against its own distribution.
A rejection-sampling acceptance rule guarantees the output distribution is identical to plain target sampling — speculative decoding changes speed, never quality. If the draft agrees a fraction α of the time and proposes k, you emit on the order of (1 − αk+1)/(1 − α) tokens per target pass; at α ≈ 0.7, k = 4 that is ~2–3× fewer target passes, i.e. ~2–3× lower latency. The art is a draft that is fast yet agrees often — typically a distilled version of the target, which is exactly what the distillation track produces (self-speculative variants like Medusa instead bolt extra heads onto the target). The speedup is biggest where decode hurts most — small batch — and shrinks once batching already saturates the GPU.
5 · Quantization: attack the bandwidth wall directly
Decode's cost is bytes-of-weights moved per token. So the most direct lever is obvious: move fewer bytes. Weight-only quantization stores the weights in INT8 or INT4 and dequantizes them in-kernel just before the matmul. Activations and the math can stay in bf16 — the point is purely to shrink the HBM→SM traffic, which is the binding constraint. INT8 halves the weight bytes (≈2× decode speedup, near-lossless); INT4 quarters them (≈4× in the bandwidth limit, with a small, recoverable quality cost using schemes like GPTQ/AWQ that protect salient weights). Because decode is bandwidth-bound, this is one of the rare wins that improves both memory footprint and speed at once.
This is also where the params-vs-active-FLOPs motif from lesson 10 returns, now on the serving side. Decode cost scales with the bytes of weights you must load per token — and a Mixture-of-Experts model loads only its active experts per token, not all of them. So the same trick that let MoE buy parameters without buying training FLOPs also lets it serve a far larger total parameter count at the bandwidth cost of a small dense model. Capability tracks total params; both training FLOPs and decode bandwidth track only the active ones.
The second target is the cache. KV-cache quantization stores K and V in INT8 instead of bf16, halving the per-token cache bytes from §2 and composing with GQA — the ÷g and the ÷2 multiply. Mind the floor: KV is sensitive, and 4-bit KV needs per-channel scaling to avoid attention degrading. Quantization is also the on-ramp to the distillation/compression track, where quantization-aware training, distillation into a smaller dense model, and the speculative draft live in full.
| Lever | What it attacks | Rough win | Cost / risk |
|---|---|---|---|
| GQA (lesson 3) | KV bytes ÷ g | ~4× more concurrency (g=4) | tiny quality drop, baked in at train time |
| Continuous batching | GPU idle time | ~2–3× throughput | per-user latency variance |
| Speculative decoding | target passes / token | ~2–3× latency (small batch) | needs a good draft; gains fade at high batch |
| Weight INT8 / INT4 | weight HBM traffic | ~2× / ~4× decode speed | INT4 needs GPTQ/AWQ; small quality cost |
| KV-cache INT8 | cache bytes ÷ 2 | ~2× context / batch | KV is precision-sensitive at 4-bit |
6 · Where the deep treatment lives
This lesson is the map; the territory is three tracks. The production serving engines — paged KV, continuous batching, prefix caching, the scheduler — are the vLLM track and the SGLang track. The kernels underneath — fused attention-decode, paged/radix layouts, INT4 dequant-fused GEMMs, CUDA-graph capture to kill launch overhead at batch 1 — are gpu_kernels/13 (SGLang kernel stack) and gpu_kernels/14 (paged & radix layouts), with quantization kernels continuing into gpu_kernels/15. Compression — distillation, quantization-aware training, the speculative draft — is the distillation track. Read this to know which lever and why; read those to implement it.
7 · Drive the decode arithmetic
Every knob below is a real serving decision. Model size and weight precision set the per-token weight traffic (the bandwidth wall). Context T, batch, and GQA groups set the KV cache. Watch three things. First, KV cache fills the 80 GB budget long before compute does — and GQA / KV-quant claw it back. Second, raising batch climbs aggregate throughput (tokens/s) while per-user tokens/s falls — the throughput/latency trade made concrete. Third, the FLOP/byte readout shows prefill sitting comfortably above the ridge (compute-bound) while decode sits far below it (memory-bound) — the asymmetry that defines the whole lesson. The configuration that "breaks": long context × large batch × MHA overflows HBM (the meter turns red); shrink it with GQA or KV-quant.
Notice the asymmetry the FLOP/byte KPIs expose: prefill processes the whole prompt so it does thousands of FLOPs per weight byte and sits compute-bound; decode does ~2N FLOPs against a full weight reload and sits at single digits — far under the ~300 ridge. Every lever in this lesson is an attempt to drag that decode number rightward: batching shares the byte load, quantization shrinks it, speculation does more useful work per load.
Failure modes & checklist
Failure modes
- Optimizing decode FLOPs. Buying a faster-FLOP GPU to speed up generation. Decode is bandwidth-bound; you needed more TB/s or fewer bytes, not more TFLOP/s. Signal: a chip with 2× the FLOPs gives near-zero decode speedup.
- Forgetting the KV cache in the memory budget. Sizing a deployment by weights alone, then OOM-ing under long-context concurrent traffic. Signal: fits at 2K context, crashes at 32K with the same batch.
- Static batching in production. One 2,000-token request stalls a whole batch of short ones. Signal: p99 latency spikes and GPU utilization sits at 30% under bursty, variable-length load.
- A draft model that's too big or too divergent. Speculative decoding where the draft is slow or rarely accepted — overhead with no payoff. Signal: acceptance rate < 50% or end-to-end latency worse than plain decode.
- Quantizing the KV cache too aggressively. 4-bit KV without per-channel scaling silently degrades long-context attention. Signal: quality fine on short prompts, falls apart deep into long contexts.
Checklist
- Separate prefill and decode in your mental model and your scheduler; they are different workloads.
- Budget HBM as weights + KV cache, and size KV at 2·L·nkv·dhead·T·batch·bytes for your worst-case context.
- Pick GQA group g from the serving budget (lesson 3) — it divides the cache directly.
- Use continuous batching for any variable-length traffic; it is nearly free throughput.
- Choose your point on the throughput/latency curve deliberately — chat ≠ batch jobs.
- Reach for weight INT8/INT4 first on a bandwidth-bound decode; add speculative decoding for low-batch latency.
Checkpoint
Where this points next
The model is built, aligned, and now servable — paged KV cache, GQA-shrunk and INT8-quantized, packed into continuous batches, latency cut with a distilled speculative draft. There is no new wall: this is the last engineering stage in the chain. What remains is to see that it was a chain — that the tokenizer choice (01), the architecture swaps (03), the FSDP/TP split (08–09), the scaling-law size (11), and the serving levers here are all the same decision asked at different points: what is the most capable model this budget can buy, and serve? Lesson 17, the capstone, walks one concrete budget end to end and shows every earlier choice as a forced consequence of the last.
Interview prompts
- Why is decode memory-bound while prefill is compute-bound? (§1 — prefill does ≈2N·P FLOPs against one weight load over P parallel positions (high intensity); decode does ≈2N FLOPs per sequence but reloads all weights + KV each step, so intensity ≈ 1 FLOP/byte at batch 1, far below the ~300 ridge. All the cost is in the 499 decode steps, not the 1 prefill.)
- Derive the KV-cache size and explain what it caps. (§2 — 2·L·n_kv·d_head·T·batch·bytes; the 2 is K+V. For a 7B model (L=32, d_head=128, bf16, MHA) it's 512 KB/token, ~17 GB at 32K context — more than the weights. Concurrent batch is capped by (HBM − weights)/(KV per sequence); the cache, not FLOPs, is the limit.)
- How does GQA help inference and by how much? (§2 — it shares each K/V head across g query heads, so n_kv = h/g and the cache divides by exactly g; g=4 turns 512 KB/token into 128 KB, quadrupling concurrency, with query heads and quality essentially unchanged. It's a serving decision made at training time.)
- Static vs continuous batching — what changes and why does it matter? (§3 — static finishes when the longest sequence does and holds slots idle; continuous swaps a finished sequence out and a queued one in every decode step, refilling slots, giving ~2–3× throughput on variable-length traffic. It also requires splicing compute-bound prefills among memory-bound decodes.)
- How does speculative decoding speed things up without changing the output? (§4 — a small draft proposes k tokens; the target verifies all k in one parallel pass (cheap because decode had spare FLOPs); a rejection-sampling rule keeps the longest agreed prefix and guarantees the target's exact distribution. ~2–3× fewer target passes at α≈0.7, k=4; the draft is typically a distilled target.)
- Why does weight-only quantization help decode specifically? (§5 — decode is bound by weight bytes moved per token, so INT8/INT4 halving/quartering the weight traffic gives ~2×/~4× decode speedup; dequant happens in-kernel and math stays bf16. It barely helps compute-bound prefill, because prefill isn't bandwidth-limited.)
- State the throughput/latency trade in one rule. (§3 — bigger batches amortize the weight load over more users so aggregate tokens/s rises, but each user's tokens/s falls and may queue behind a prefill; chat optimizes per-user latency with small batches, batch jobs optimize aggregate throughput with large ones. There is no single "fast.")