cs336 / lessons/18 · inferencelesson 19 / 20

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.

The previous step left this broken
Training spends 6ND FLOPs once and is over. Serving runs the forward pass billions of times, interactively, with a latency SLA. The catch: autoregressive generation produces tokens one at a time — token t+1 cannot start until token t is sampled — and each step reloads the model's entire weight set from HBM to compute a single new column of activations. A 14 GB model on a 3.3 TB/s GPU can physically stream its weights only ~235 times per second. (That ~235 is the weights-only ceiling — each step also reads the KV cache from HBM, which pulls the real number somewhat lower; §2.) That ceiling, not arithmetic, is why generation feels slow and costs what it costs — and it forces an entirely separate engineering stack to claw back.
Linear position
Forced by: a finished, aligned model is worthless unless each generated token is cheap in latency and dollars — and naive generation wastes the GPU it runs on.
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.
The plan
Five moves. (1) Split generation into prefill and decode, and show decode is memory-bound — callback to the roofline of lesson 6. (2) Derive the KV-cache size and prove it, not compute, caps batch and context; watch GQA divide it. (3) Batch for throughput: static vs continuous, and the latency it costs. (4) Speculative decoding: a small draft guesses, the big model checks k tokens in one pass. (5) Quantization: weight-only INT8/INT4 attacks the bandwidth wall directly, KV-quant attacks the cache. Then drive the decode-arithmetic widget: every knob is a real serving decision.

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.

PrefillDecode
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 / stepweights once, amortized over Pall weights + KV cache, every step
Arithmetic intensityhigh → compute-bound≈ 1 (batch 1) → memory-bound
Bottlenecktensor-core FLOP/sHBM bandwidth (TB/s)
Runsonce per requestonce per output token
Callback — the roofline of lesson 6
Lesson 06 split every op into compute-bound (sits at the FLOP ceiling, the chip is happy) and memory-bound (sits at the bandwidth ceiling, the chip starves). Training is dominated by big GEMMs — compute-bound. Decode is the canonical memory-bound op: it loads ~2N bytes to do ~2N FLOPs. That is why the entire inference stack below is, at heart, a campaign to raise decode's arithmetic intensity — push it rightward across the roofline ridge.

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:

KV bytes = 2 · L · nkv · dhead · T · batch · byteselt

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:

1 token, MHA= 512 KB
× 32,768 context≈ 17.2 GB for one sequence
weights themselves≈ 14 GB (bf16)
an 80 GB H100 holds~3 such long-context sequences (66 GB ÷ 17.2), then it is full

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 batchingContinuous batching
Batch changesonly between batchesevery decode step
A finished sequenceholds its slot until the longest finishesfrees its slot immediately
GPU utilization (variable lengths)low — padded to the maxhigh — slots refilled continuously
Throughput on real trafficbaseline~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.

1 · draftsmall model proposes k tokens, k cheap steps
2 · verifytarget scores all k in ONE forward pass (weights loaded once)
3 · acceptkeep the longest prefix the target agrees with; resample the first reject
4 · repeateach target pass yields 1…k+1 tokens instead of exactly 1

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.

LeverWhat it attacksRough winCost / risk
GQA (lesson 3)KV bytes ÷ g~4× more concurrency (g=4)tiny quality drop, baked in at train time
Continuous batchingGPU idle time~2–3× throughputper-user latency variance
Speculative decodingtarget passes / token~2–3× latency (small batch)needs a good draft; gains fade at high batch
Weight INT8 / INT4weight HBM traffic~2× / ~4× decode speedINT4 needs GPTQ/AWQ; small quality cost
KV-cache INT8cache bytes ÷ 2~2× context / batchKV 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.

Decode arithmetic — the bandwidth wall and the KV cap
A single 80 GB H100 at ~3.3 TB/s, ~990 bf16 TFLOP/s. Decode tokens/s is a weight-bandwidth-bound estimate: bytes/step = (weights at the chosen precision + KV read) × batch; tokens/s ≈ bandwidth / bytes-per-step × batch. The KV meter is weights + cache vs 80 GB; it turns red on overflow. Push batch up to watch aggregate throughput rise while per-user tokens/s falls, then drop GQA groups g or precision to shrink the wall.
KV cache / token
KV cache total
Decode tokens/s (aggregate)
Per-user tokens/s
Prefill FLOP/byte
Decode FLOP/byte
HBM 80 GB

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

Try it
Open the widget with the default 7B-class config. (1) Set context to 32,768, batch to 8, GQA g to 1 (MHA), precision bf16 — watch the HBM meter overflow red. Now drag g to 4: the KV cache quarters and it fits. That is lesson 3's payback in one motion. (2) Reset, then sweep batch from 1 to 64 and read both tokens/s KPIs: aggregate climbs, per-user falls — the throughput/latency trade. (3) Switch weight precision bf16 → INT4 at batch 1 and watch decode tokens/s jump ~3× at this context (it approaches 4× only at short context, where weights dominate the per-step bytes) while the FLOP/byte readout stays memory-bound. In one sentence each: why did INT4 help decode but barely help prefill, and which single knob would you turn first to serve 64 concurrent 32K-context users?

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.

Takeaway
Serving an autoregressive model is two workloads: prefill runs the whole prompt in one compute-bound pass and fills the KV cache; decode emits one token per step, reloading every weight and the whole cache from HBM for only ~2N FLOPs — so decode is memory-bound (callback to lesson 6's roofline) and is where all the latency and cost live. The KV cache, 2·L·nkv·dhead·T·batch·bytes, not compute, caps how many requests you can hold — a single 32K bf16-MHA sequence on a 7B model is ~17 GB, more than the weights — which is exactly why GQA (lesson 3, cache ÷ g) is a serving win and why vLLM pages it. Four levers raise decode's arithmetic intensity or shrink its bytes: continuous batching (~2–3× throughput, at the cost of per-user latency), speculative decoding (a distilled draft proposes k, the target verifies in one pass — same distribution, ~2–3× faster at low batch), and quantization (weight INT8/INT4 cuts the bandwidth wall ~2–4×; KV-INT8 doubles the context you can hold). Throughput and latency always trade; the deep implementation lives in vLLM, SGLang, and gpu_kernels/12–15. This closes the build — lesson 19 ties the whole budget together.

Interview prompts