Lecture 10: Inference

What this lecture covers

Training a language model is a one-time cost, but inference (generating responses from a fixed model) happens over and over, so its efficiency matters enormously in practice. The lecture asks a simple question: given a trained transformer, how do you generate text quickly? It builds up the answer from first principles, starting with why token-by-token generation is fundamentally harder to make fast than training, then surveys the techniques the field uses to fix that: shrinking the KV cache, changing the architecture entirely, quantizing and pruning weights, speculating ahead with a small model, and scheduling requests cleverly.

This is lecture 10 of Stanford's CS336, sitting between the scaling-laws lectures. After watching, you should be able to explain why generation is memory-limited rather than compute-limited, compute the arithmetic intensity of a matrix multiplication or attention layer, and describe several concrete methods (GQA, MLA, quantization, speculative decoding, paged attention) used to make real inference systems fast.

Key ideas

• Arithmetic intensity: the ratio of flops to bytes transferred for an operation; if it exceeds the accelerator's own flops-to-bandwidth ratio you are compute-limited, otherwise you are memory-limited.
• KV cache: cached key and value vectors from previous tokens that let a transformer avoid recomputing the whole sequence at every generation step, turning quadratic per-token cost into roughly linear cost.
• Prefill vs. generation: prefill encodes the prompt in parallel and is compute-limited (efficient); generation produces one token at a time and is memory-limited (slow).
• Group query attention (GQA): uses fewer key/value heads than query heads to shrink the KV cache while keeping most of the model's expressiveness.
• Multi-head latent attention (MLA): projects keys and values into a much lower-dimensional space instead of reducing the number of heads, another way to shrink the cache.
• Quantization: storing weights and activations in lower precision (for example INT8 or INT4 instead of BF16) to cut memory traffic, at some risk to accuracy.
• Speculative decoding: a small draft model generates several tokens quickly, and the large target model checks them all in one parallel pass, accepting or rejecting them so the output is statistically identical to sampling from the target model alone.
• Continuous batching and paged attention: serving-system techniques that let heterogeneous, arriving-and-departing requests share GPU memory and compute efficiently, borrowed from operating-systems ideas like virtual memory.

Walkthrough

Why inference is a distinct, important problem (0:05)

The lecture opens by pointing out that inference underlies far more than chatbots: evaluation, test-time reasoning, and reinforcement-learning-based training all require generating tokens. Efficiency metrics are introduced: time to first token (how long before generation starts), latency (how fast tokens stream after that), and throughput (tokens generated per second across all users, useful for batch workloads). The key structural fact driving the rest of the lecture is that training can process a whole sequence in parallel, while generation is inherently sequential because each new token depends on everything before it.

Arithmetic intensity and why generation is memory-limited (8:10)

Using a simple matrix multiplication example, the lecture defines arithmetic intensity as flops divided by bytes transferred, and shows it depends on batch size. Comparing that ratio to a GPU's own flops-per-byte-of-bandwidth tells you whether an operation is compute-limited (good) or memory-limited (bad). A batch size of one, which is what happens during single-token generation, gives an arithmetic intensity of about one, which is far below what an H100 needs to be compute-bound. This is the mathematical reason generation is slow.

Counting flops and bytes for MLP and attention layers (15:22)

The lecture works through the same accounting for a full transformer, separating MLP layers from attention layers, using S for prompt length and T for the number of tokens being processed. MLP arithmetic intensity scales with B*T, so it can be improved simply by batching more requests together. Attention, however, has intensity that depends only on sequence length, not batch size, because every sequence carries its own private KV cache. This explains why attention is the harder bottleneck during generation: batching more users does not help it the way it helps the MLP layers.

Latency, throughput, and the batch-size trade-off (27:44)

Using a worked Llama 2 13B example on an H100, the lecture shows concretely how memory usage, latency, and throughput change as batch size grows from 1 to 16 to 256: latency worsens, throughput improves but with diminishing returns, and eventually the KV cache simply does not fit in GPU memory. This produces the core trade-off of serving systems: small batches favor low latency, large batches favor high throughput, and simply running multiple independent model copies is a trivial but effective way to scale throughput without added latency.

Shrinking the KV cache (38:00)

Since the KV cache is the main driver of memory limits, the lecture surveys architectural changes that shrink it: group query attention (fewer key/value heads), multi-head latent attention from DeepSeek (projecting keys and values to a smaller dimensional space), cross-layer attention (sharing key/value projections across layers), and local or sliding-window attention (only keeping a fixed window of recent tokens, sometimes mixed with occasional full-attention layers as in Character AI's architecture). Each trades some accuracy for a smaller, faster cache.

Beyond the transformer: state space models and diffusion (52:15)

The lecture then looks at architectures designed from scratch for efficient inference. State space models (S4, Mamba, and later linear-attention variants used by models like MiniMax) replace the growing KV cache with a constant-size state, at the cost of struggling on tasks that need exact long-range recall unless mixed with occasional full-attention layers. Diffusion language models generate all tokens in parallel and iteratively refine them, sidestepping autoregressive generation entirely; the lecture shows a coding demo from a diffusion model running dramatically faster than transformer-based generation.

Quantization, pruning, and speculative decoding (1:04:26)

Quantization reduces numeric precision (BF16 down to INT8 or INT4) to cut memory traffic, with techniques like LLM.int8() and activation-aware quantization handling outlier values specially to preserve accuracy. Pruning removes unimportant layers, heads, or hidden units identified via a calibration set, then distills the original model into the pruned one to recover quality. Speculative decoding is presented as a way to get speed without losing accuracy: a small draft model proposes several tokens, the large target model verifies them in one parallel (prefill-like) pass, and a rejection-sampling-style acceptance rule guarantees the output distribution exactly matches sampling from the target model alone.

Serving real traffic: continuous batching and paged attention (1:17:36)

Finally, the lecture addresses live serving, where requests arrive and finish at different times and have different lengths. Continuous batching inserts new requests as soon as GPU capacity frees up rather than waiting for a batch to fully finish. Paged attention, the idea behind vLLM, borrows virtual memory from operating systems, dividing the KV cache into fixed-size blocks that can be allocated non-contiguously, reducing fragmentation and enabling copy-on-write sharing of prompts across requests.

Before you watch

• Be comfortable with the transformer flop-counting and memory notation from the course's first lecture and assignment one.
• Review group query attention and multi-head latent attention if they were introduced in earlier architecture lectures, since this lecture revisits them specifically for inference.
• Some familiarity with GPU memory bandwidth versus compute throughput helps the arithmetic-intensity discussion land quickly.

Check your understanding

1. Why does a batch size of one give an arithmetic intensity of about one, and why is that bad for GPU utilization?
2. Explain why MLP layers become more efficient with larger batches but attention layers do not.
3. How does group query attention reduce the KV cache, and what is the accuracy trade-off?
4. Why does speculative decoding produce exact samples from the target model rather than an approximation?
5. What problem does paged attention solve, and what operating-system idea does it borrow?