Lecture 10 · Part 3

Inference

Percy takes a break from scaling laws to ask what happens after training. Prefill parallelizes like training, but generation emits one token at a time, so attention's arithmetic intensity is stuck below 1 and inference is memory-bound. Everything else follows from that one fact — the latency-throughput tradeoff, the war on KV-cache size, quantization and pruning, speculative sampling, and vLLM-style paging.

▶ Lecture video Executable lecture (.py) ↗

Inference is where the compute now goes

Training is a one-time cost; inference is a repeated one, incurred every single day. Percy’s calibration point: OpenAI is estimated to produce about 8.6 trillion tokens a day, while DeepSeek v4 was trained on 32 trillion — so in under four days of serving, the token count of a frontier training run goes by. And the agentic shift raises the stakes. A chatbot’s tokens are for a human to read, and humans read slowly; an agent’s tokens are mostly internal — thinking, tool calls, introspection — with only the final output meant for eyes. “You should really think about the number of tokens generated as really the compute spend,” and for an ambitious enough problem there is no ceiling on how much value more inference buys. Inference also sits inside the pipeline itself: generation-based evaluation and RL rollouts both need it. Hence a whole serving ecosystem — vLLM (the go-to), SGLang (strong for agentic workloads), TensorRT-LLM (very fast, narrower), llama.cpp (CPU inference).

But what does “fast” mean? Three metrics, for different settings:

Time-to-first-token (TTFT) — how long a user waits before anything appears.
Latency (seconds/token) — how fast tokens stream for one query.
Throughput (tokens/second) — how fast tokens appear across many queries, the number a batch-processing job cares about.

The one high-level bit to remember: in training you see all the tokens at once — the sequence is just another dimension of a big tensor that gets multiplied. Autoregressive generation forbids that. Tokens come out sequentially, one at a time, so you cannot parallelize over the sequence dimension, and keeping the accelerator busy gets hard.

The KV cache, and the arithmetic intensity of generation

Naive inference refeeds the whole growing history through the transformer for every new token: each forward pass is $O(T^2)$ , so generating $T$ tokens costs $O(T^3)$ . But the transformer is causal — appending a token changes nothing about earlier positions — so the keys and values of the prefix can be computed once and reused.

Notation, borrowed from Google’s scaling book (much of this lecture follows it): $B$ batch, $D$ model dimension, $F = 4D$ the MLP dimension, $N$ query heads, $K$ key/value heads, $H$ head dimension with $D = NH$ , and two sequence lengths — $S$ tokens conditioned on, $T$ tokens being generated. Training has $T = S$ ; generation has $T = 1$ .

First a lecture-2 reprise. Multiplying $X$ ( $B \times D$ ) by $W$ ( $D \times F$ ) in bf16 reads $2BD + 2DF$ bytes, does $2BDF$ FLOPs, writes $2BF$ bytes; when $B$ is much smaller than $D$ and $F$ the intensity simplifies to just $B$ . The H100’s accelerator intensity is $989 \times 10^{12} / 3.35 \times 10^{12} \approx 295$ , so this matmul is compute-bound iff $B$ exceeds 295. The extreme case $B = 1$ is a matrix–vector product with intensity 1 — you read a whole $D \times F$ matrix to do only $2DF$ FLOPs — and “this is basically what happens with inference.”

Now do the same accounting for a transformer block, matmuls only (everything else is few FLOPs and fusable). The MLP’s three projections (up, gate, down) cost $6BTDF$ FLOPs against $4BTD + 4BTF + 6DF$ bytes; assuming $BT$ is much smaller than $D$ and $F$ , the intensity is $BT$ . Attention is a different story: reading $Q$ , $K$ , $V$ and writing the output moves $4BSD + 4BTD$ bytes for $4BSTD$ FLOPs, giving intensity

$\frac{4BSTD}{4BSD + 4BTD} = \frac{ST}{S + T}.$

The batch dimension has vanished. Specialize both formulas to the two stages:

Latency, throughput, and the batch-size dial

Memory-boundedness simplifies everything: assuming compute and communication overlap perfectly, the time per step is just the memory that must move through HBM.

Code walkthrough — Transformer performance stats, from the trace

The trace builds these symbolically with sympy and substitutes a config:

num_params = 2*V*D + 3*D*F*L + (2*D*N*H + 2*D*K*H)*L
parameter_size = 2 * num_params                 # bf16
kv_cache_size_per_seq = S * (K*H) * L * 2 * 2   # key+value, bf16
memory = B * kv_cache_size_per_seq + parameter_size
latency = memory / memory_bandwidth             # seconds/token
throughput = B / latency                        # tokens/second

Llama 2 13B’s config: $S=1024$ , $D=5120$ , $F=13824$ , $N=K=40$ (no GQA), $H=128$ , $L=40$ , $V=32000$ , H100 bandwidth $3.35\times10^{12}$ bytes/s. That yields 13.0B parameters (sanity check: it is advertised as 13B), 26 GB of weights, and an 839 MB KV cache per sequence.

Worked example — Llama 2 13B on one H100, sweeping batch size

$B$	Memory	Latency	Throughput
1	26.9 GB	8.0 ms/token	124 tok/s
64	79.7 GB (just fits)	23.8 ms/token	~2,700 tok/s
256	241 GB	—	out of memory

Increasing $B$ worsens latency (a bigger KV cache to read every step) but improves throughput (the 26 GB of parameters are loaded once for the whole batch) — until the KV cache blows past 80 GB, with the throughput gains already diminishing.

Shrinking the KV cache without losing the model

Memory is the bottleneck and the KV cache can outweigh the parameters at large batch size, so the recurring game is: reduce the KV cache, carefully, without giving up accuracy.

Grouped-query attention (GQA) keeps $N$ query heads but only $K$ key/value heads ( $K = N$ is standard multi-head attention; $K = 1$ is multi-query attention, “which no one uses because it’s really bad”), cutting the cache by $N/K$ . Rerunning the Llama numbers with $K = 8$ (a 1:5 ratio): at $B = 64$ memory drops from 79.7 GB to 33.4 GB, latency improves to 10.0 ms/token, and throughput jumps to ~6,400 tok/s — reducing memory improves both metrics, since it’s only the batch dimension that puts them in tension. And the freed memory lets you push to $B = 256$ (65.6 GB, ~13,000 tok/s). The GQA paper’s evals show little accuracy loss — though DeepSeek’s later tables show MHA beating GQA, so “take everything that’s not just math with a grain of salt.”

Multi-head latent attention (MLA), from DeepSeek, compresses instead of sharing: project activations down to a $C$ -dimensional latent (DeepSeek v2: from $NH = 16384$ to $C = 512$ ), cache only that, and materialize keys and values on demand. One wrinkle: MLA is incompatible with RoPE, so 64 extra dimensions are added for it (576 total). Their tables show MLA matching or slightly beating full MHA. Cross-layer attention (CLA) shares KVs across layers, the way GQA shares them across heads, improving the accuracy-versus-cache-size Pareto frontier. Sliding-window attention attends only to the last $k$ tokens, making the cache independent of sequence length (effective context still grows with depth, since information propagates across layers) — but it hurts accuracy, so practical models interleave local layers with full-attention layers. DeepSeek v4’s 1M-context stack piles on compressed sparse attention: compress every $m$ tokens into one, cheaply score and select a top- $k$ subset, compress further. And beyond the transformer sit linear attention and state-space models (Mamba, gated DeltaNet), which replace the cache with a fixed-size state — more powerful than a sliding window, but a lossy summary of the past — plus non-autoregressive diffusion models.

Quantize, prune, distill

Two more lossy levers, more systems- than architecture-flavored. Quantization reduces the precision of the numbers themselves: bf16 is the inference default, and the menu runs down through fp8, int8, and int4. You can train with quantization in the loop (quantization-aware training — simulate the round-trip during the forward pass; effective but expensive) or quantize after the fact (post-training quantization — cheap). Naive per-tensor scale-and-zero-point works poorly; GPTQ uses Hessian information to push each layer’s quantization error into the not-yet-quantized weights; activation-aware quantization (AWQ) observes that some activation channels are large, so the weights they touch matter more — keep 0.1 to 1 percent of weights in high precision and crunch the rest, getting fp16-to-int3’s 4x memory reduction with a 3.2x speedup.

Model pruning is cruder: rip parts out of a big model, then fix it up. NVIDIA’s recipe identifies important layers, heads, and hidden dimensions on a small calibration set (~1,024 samples, judging by activation magnitudes), removes the unimportant ones, and distills the original model into the pruned one to heal it — taking a 15B model to 8B with little accuracy loss and far less training than starting over. That’s the general pattern: either define a faster architecture and train it from scratch, or initialize the fast architecture from a good model and repair it with distillation.

Speculative sampling: a lossless shortcut

Everything so far trades accuracy for speed. Speculative sampling (also called speculative decoding) is the elegant exception. The asymmetry it exploits: prefill processes tokens in parallel and is compute-bound, generation is one-at-a-time and memory-bound — so checking a proposed sequence is much faster than generating it.

Too few draft tokens wastes the target model’s parallel checking; too many means frequent rejections; the sweet spot in the original paper sits around three or four. Draft models are typically an order of magnitude smaller (70B target with an 8B draft, 8B with 1B) and ideally distilled toward the target — so all the compression tricks above reappear here: if the shrunken model is good enough, serve it; if not, it makes a fine draft model. Extensions like Medusa (draft multiple tokens in parallel) and EAGLE (draft from the target’s high-level features) keep improving the drafter.

Serving live traffic: continuous batching and PagedAttention

A live server is nothing like training’s dense token blocks: requests arrive at different times, share prefixes (system prompts, multiple samples per prompt), and have different lengths. Two systems ideas handle the mess. Continuous batching (introduced by Orca) schedules at the iteration level: decode one step for every sequence in the batch, eject sequences that finish, and slot newly arrived requests in immediately rather than waiting for the batch to drain. Ragged lengths are handled by selective batching — attention has to process each sequence separately anyway (each has its own cache), but the MLP layers, which carry most of the FLOPs, can concatenate everything into one mega-sequence, e.g. lengths 3, 9, and 5 becoming a single $17 \times H$ tensor.

PagedAttention (the core idea of the vLLM paper) fixes how the KV cache is stored. Reserving a contiguous max-length buffer per request wastes memory exactly the way files fragment a hard drive: internal fragmentation (you reserved 1,024 tokens and generated 50) and external fragmentation (unusable gaps between allocations). The operating-systems people recognized their own problem and applied their own solution — paging. Chop each sequence’s cache into fixed-size blocks that can live anywhere in memory, keep an index, and share blocks across requests: a common system prompt is cached once for everyone, multiple samples from one prompt share its prefix, and copy-on-write splits a block only when two continuations diverge.

Back of the chalk tray — what to remember

Inference is a repeated cost that now rivals training compute, and agents make tokens-generated synonymous with compute-spent; fast means TTFT, latency, or throughput depending on who’s asking.
Prefill is compute-bound; generation is memory-bound, because attention’s generation intensity $S/(S+1)$ never reaches 1 and batching can’t fix it — each sequence owns its KV cache.
Since generation is memory I/O, latency = bytes moved / bandwidth and throughput = B / latency: bigger batches buy throughput at the price of latency, until the KV cache exhausts HBM.
The recurring lossy game is shrinking the KV cache (GQA, MLA, CLA, sliding windows) or the model (quantization, pruning + distillation) without giving up accuracy.
Speculative sampling is lossless: a draft model guesses, the target model checks in parallel, and the math guarantees exact samples from the target.
Live serving borrows from operating systems — continuous batching for ragged arrivals, PagedAttention’s blocks and copy-on-write for a fragmented, shared KV cache.

Check yourself

Napkin math in the lecture’s style — the formulas above and a calculator are all you need.

Problem 10.1 intensity accounting

A model is generating (so $T = 1$ ) with $S = 2048$ tokens of context on an H100 (accelerator intensity 295). What is the attention arithmetic intensity? What is the MLP intensity, and how many concurrent requests would make the MLP compute-bound? Why can’t the same batching trick save attention?

Show solution

Attention: $ST/(S+T) = 2048/2049 \approx 1$ — hopelessly below 295, and note that $B$ does not appear, so it stays there at any batch size. MLP: intensity is $BT = B$ , so you need $B \gt 295$ concurrent requests to cross the threshold. Batching helps the MLP because all sequences share the same weight matrices — read them once, use them $B$ times. In attention, each sequence has its own KV cache: $B$ is a batching dimension that appears in both operands and the output, so doubling $B$ doubles the bytes moved along with the FLOPs. It’s $B$ independent matrix–vector products, and more of them doesn’t raise the intensity of any one.

Problem 10.2 KV cache budgeting

An 8B-parameter model served in bf16 (16 GB of weights) uses GQA with $K = 8$ KV heads, $H = 128$ , $L = 32$ layers, and serves sequences of $S = 4096$ tokens. How big is the KV cache per sequence, and how many concurrent sequences fit on one 80 GB H100?

Show solution

Per sequence: $S \cdot KH \cdot L \cdot 2 \cdot 2 = 4096 \cdot 1024 \cdot 32 \cdot 4 \approx 537$ MB. Budget left after weights: $80 - 16 = 64$ GB, so $64 \times 10^9 / 5.37 \times 10^8 \approx 119$ sequences. Note the cache term dominates the weights well before that: past $B \approx 30$ , the server is spending more HBM on cache than on the model itself — which is why every KV-cache-reduction trick translates directly into batch size and throughput.

Problem 10.3 latency-throughput tradeoff

For the model of 10.2 on an H100 (bandwidth $3.35 \times 10^{12}$ bytes/s), compute the theoretical latency and throughput at $B = 32$ and at $B = 64$ . Which setting would you pick for a chatbot, and which for overnight document processing?

Show solution

$B = 32$ : memory per step $= 16 + 32 \times 0.537 = 33.2$ GB, latency $= 33.2 \times 10^9 / 3.35 \times 10^{12} \approx 9.9$ ms/token, throughput $= 32 / 0.0099 \approx 3{,}230$ tok/s. $B = 64$ : memory $= 50.4$ GB, latency $\approx 15.0$ ms/token, throughput $\approx 4{,}260$ tok/s. Doubling the batch made latency 1.5x worse but throughput only 1.3x better — the parameter reads are already amortized and the growing cache dominates. The chatbot wants $B = 32$ (or less): users feel latency and TTFT. The overnight job wants $B = 64$ (or as large as memory allows): nobody is waiting on any single sequence, only on the whole pile.

Problem 10.4 speculative acceptance

A target model over two tokens A and B has $q(A) = 0.6$ , $q(B) = 0.4$ ; the draft model has $p(A) = 0.8$ , $p(B) = 0.2$ . Walk one step of speculative sampling and verify the output distribution is exactly $q$ .

Show solution

The draft proposes A with probability 0.8; A is accepted with probability $\min(1, 0.6/0.8) = 0.75$ . It proposes B with probability 0.2; B is accepted with probability $\min(1, 0.4/0.2) = 1$ . On rejecting A (probability $0.8 \times 0.25 = 0.2$ ), sample from the residual $\max(q - p, 0)$ normalized: A’s residual is $\max(0.6 - 0.8, 0) = 0$ , B’s is $0.2$ , so the residual puts all mass on B. Totals: $P[A] = 0.8 \times 0.75 = 0.6 = q(A)$ and $P[B] = 0.2 \times 1 + 0.2 \times 1 = 0.4 = q(B)$ . Exact — the draft model only determines speed (how often you get tokens accepted in cheap bursts), never the distribution.