Week 2 — Transformer Memory and Attention Economics
Attention · KV Cache · Transformers
Advanced
I — Long-context architecture
Overview
Attention is the highest-leverage mechanism in the course. This week builds Q/K/V, the O(N²) cost, KV caching, MHA/GQA/latent attention, and the three independent compression axes the report identifies: entry size, sequence length, and layers.
What You Will Learn
- Derive scaled dot-product attention and its complexity.
- Write the KV-cost decomposition as bytes/entry × entries × layers.
- Compare dense, sliding-window and top-K sparse attention.
- Implement causal multi-head attention without a framework attention op.
Core Concepts
Scaled dot-product attention
\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\!\left(\frac{QK^{\mathsf T}}{\sqrt{d_k}}\right)V
The three compression axes
\text{KV cost}\approx(\text{bytes/entry})\times(\text{sequence entries})\times(\text{cached layers})
KV caching
Keys/values already computed are deterministic and do not change; caching them turns per-step cost from O(N²) recompute into O(N) append + attend.
Prerequisites
Work these pages on this site before the lecture.
Lecture notes
Scaled dot-product attention is O(N^2 d) compute and O(N^2) attention weights if materialized. FlashAttention never writes the full N\times N matrix; it tiles SRAM so the algorithm is still quadratic but the IO is linear in the tiles. That is why “we implemented attention” and “we can run 32K” are different claims.
The course’s three compression axes are independent:
- Entry size — GQA / MQA / MLA shrink bytes per token (shared KV heads, or a latent KV).
- Sequence entries — sliding window or top-K sparse attention drop how many keys a query sees.
- Cached layers — cross-layer reuse (week 4) drops how many layers store a unique KV.
Build the simulator first with dense / window / top-K masks. Do not start from a fused kernel.
Required readings
- Attention Is All You Need
- FlashAttention-2
- GQA (Ainslie et al.)
- DeepSeek-V2 / MLA
- Longformer · BigBird
Related on this site
Lab / Implementation
Build an attention simulator with dense, sliding-window and top-K modes, instrumented for FLOPs, KV bytes and peak memory. Implement causal MHA from scratch and match a reference.
Not on this site (paper reading required)
- FlashAttention (IO-aware)
- GQA / MQA / MLA
- sparse-attention kernels
Mastery Check
- DERIVE: the O(N²) vs O(N·n_win) argument and the 1/√d_k variance rationale
- IMPLEMENT: causal MHA + KV cache, three-mode attention simulator
- BENCHMARK: crossover N per mode; parity vs reference
- DEBUG: mask-after-softmax, wrong scale, head/sequence transpose
- EXPLAIN: why KV caching changes autoregressive inference economics