Modern architecture evolution (Part 2): 1M context, KV Cache, and MLA
Why does 128K or 1M context blow up GPU memory, and what do MHA, GQA, MLA, and Linear Attention actually save?
During generation, each new token attends to all previous tokens. To avoid recomputing past Keys (K) and Values (V) at every step, engines cache them in VRAM as the KV Cache. At 128K to 1M tokens, this KV Cache becomes the #1 memory bottleneck.
Modern architectures attack attention cost in three ways: 1) Compress KV cache size (GQA shares KV heads; DeepSeek MLA projects KV into a low-rank latent vector); 2) Sparse Attention (DSA/NSA attends only to selected top tokens); 3) Linear / Hybrid Attention (reduces O(N²) complexity to O(N) linear state updates).
Cache past K/V tensors to drop decode cost from O(N²) to O(N), then compress with GQA/MLA
In long-context inference, KV Cache memory and bandwidth quickly overtake static model weights.
Causal Masking Keeps Past K/V Immutable Across Decode Steps
Only the newest token computes Q/K/V; past keys and values are read from VRAM.
Under Standard MHA, Long-Context KV Cache Dwarfs Model Weights
KV footprint scales linearly with sequence length, layer count, and KV heads.
Group-Shared KV Heads (GQA) or Latent Compression (MLA) Slash VRAM
Combines architectural compression with KV quantization to fit 80GB GPUs.
First-principles breakdown: three steps
Pinpoint the Bottleneck: O(N²) & KV Cache Wall
Storing K and V across all layers and heads for long contexts quickly consumes more VRAM than the model weights themselves.
Route 1: Compress KV Cache (GQA → MLA)
GQA shares KV across query groups; MLA compresses KV into a tiny low-rank latent vector, cutting cache VRAM by >90%.
Route 2: Prune O(N²) Dot Products (Sparse & Linear)
Instead of attending to all 1M past tokens, select relevant blocks dynamically or interleave O(N) linear attention layers.
Drag context length from 4K to 1M in the simulator to watch standard MHA overflow 8×H100 GPUs, then switch to GQA, MLA, and Linear-Hybrid to cut VRAM.
Drag to 1M tokens: standard MHA consumes 491.5GB of KV Cache VRAM, whereas MLA and Linear Hybrid compress it to 22GB and 7.4GB.
KV Cache
Caching previously computed Key and Value vectors in GPU memory so each new token doesn’t recompute the entire prefix.
Classic space-for-time trade-off; the primary reason long-context inference is expensive.MHA → MQA → GQA → MLA
MHA = 1 KV per head; MQA = 1 shared KV; GQA = grouped KV; MLA = low-rank compressed latent KV cache.
All four acronyms solve the exact same problem: shrinking KV Cache VRAM without hurting quality.Sparse & Hybrid Linear Attention
Sparse attention scores only top relevant chunks; Linear attention maintains a fixed-size state, often mixed with full attention layers.
The core engineering enablers behind 1M+ token context windows.Every new attention acronym either shrinks KV Cache memory (GQA/MLA) or avoids O(N²) all-pairs dot products (Sparse/Linear attention).
Why do GQA and MLA drastically reduce the number of GPUs needed to serve long-context requests?
Explanations, cases, and exercises are authored by Reallier. Links provide further learning. The workplace activity and people are fictional.