Attention Mechanics and KV Cache
Attention is the core token-mixing mechanism in transformer LLMs. The KV cache is the key serving trick that makes autoregressive generation feasible.
Self-Attention In One Sentence
Each token asks a learned question, compares it against learned addresses for previous tokens, and takes a weighted sum of their learned content.
Q, K, V
Given hidden state h, linear projections produce:
Q = h W_Q
K = h W_K
V = h W_V
Then:
scores = QK^T / sqrt(d_head)
weights = softmax(mask(scores))
output = weights V
The causal mask forces position t to see only positions <= t.
Multi-Head Attention
Multi-head attention splits the representation into multiple heads. Each head can learn a different relation pattern.
Shape intuition:
h: [B, T, d_model]
Q/K/V: [B, n_heads, T, d_head]
attention weights: [B, n_heads, T, T]
The heads are concatenated and projected back to d_model.
Why Attention Is Expensive
For sequence length T, full attention forms a T x T attention matrix per head.
Training/prefill cost scales roughly:
O(T^2 * d)
This is why long-context models are expensive, and why Long Context and Efficient Sequence Models matters.
Autoregressive Inference
During generation, the model produces one token at a time:
x_1 -> x_2 -> x_3 -> ...
Without caching, each new token would recompute keys and values for all previous tokens. The KV cache stores prior K and V tensors so the model only computes new query/key/value projections for the latest token and attends over cached history.
KV Cache Tradeoff
The KV cache saves compute but costs memory.
Memory grows with:
- Batch size.
- Sequence length.
- Number of layers.
- Number of KV heads.
- Head dimension.
- Numeric precision.
This is why inference can become memory-bandwidth-bound rather than FLOP-bound.
MHA, MQA, And GQA
MHA: Multi-Head Attention
Each query head has its own key and value heads.
Pros:
- Maximum attention expressivity.
Cons:
- Large KV cache.
MQA: Multi-Query Attention
Many query heads share one key/value head.
Pros:
- Much smaller KV cache.
- Faster decoding.
Cons:
- Can reduce quality if too aggressive.
Primary source: Fast Transformer Decoding: One Write-Head is All You Need.
GQA: Grouped-Query Attention
Query heads are partitioned into groups; each group shares key/value heads.
Pros:
- Compromise between MHA quality and MQA efficiency.
Primary source: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.
FlashAttention
FlashAttention is not a new attention definition. It computes exact attention with IO-aware tiling so less data is moved between GPU memory levels.
Primary source: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.
Key idea:
- Avoid materializing the full attention matrix in slow memory.
- Tile computation to use SRAM efficiently.
- Get exact attention with lower memory and often better speed.
Research Questions
- Can we reduce attention memory without damaging retrieval or reasoning?
- Which long-context tasks truly need full attention?
- How do MQA/GQA trade off quality versus serving cost at scale?
- How much of long-context performance is architecture versus data curriculum?