The Problem
A naive autoregressive decoder does O(N²) work to generate N tokens: at each step it recomputes attention over the full prefix. For a 4K-token response that is 16M attention operations, most of them redundant. Every hidden state of a prefix token is deterministic once computed — you only need to run the new token's query against the cached keys and values of everything before.
On top of that, attention itself moves a lot of data. Standard attention materializes an N×N score matrix, N×d softmax output, N×d final output — too many reads and writes to HBM. For N≥2K, attention becomes memory-bound before it becomes FLOP-bound. Classic attention kernels underuse modern GPUs by 4–10×.
Two optimizations pushed frontier inference from "slow" to "fast": the KV cache, standard in autoregressive decoders since the original Transformer, and Flash Attention (Dao et al., 2022):
- KV cache. Store the K and V vectors of every prefix token. Each new token's attention is one query against the cached keys. Inference reduces from
O(N²) to O(N) per generation step. - Flash Attention. Tile the attention computation so the full N×N matrix never hits HBM. All of softmax + matmul happens in SRAM. 2–4× wall-clock speedup on A100; 5–10× on H100 with FP8.
By 2026 both are universal. Every production inference stack (vLLM, TensorRT-LLM, SGLang, llama.cpp) assumes them. Every frontier model ships with Flash Attention enabled.
Unlock the full lesson
You've read the first 2 sections. The rest of this lesson covers The Concept, Build, Check Yourself, Check Yourself, Key Terms & Next — plus a hands-on lab, quiz, and project artifact.
Create a free account — the opening phases of 29 of 35 courses are free, no credit card.