This notebook compares three attention variants used in modern LLMs:
We’ll implement each from scratch, visualize attention patterns, and benchmark memory/speed tradeoffs.
| # | Concept | Where in this notebook | Reference |
|---|---|---|---|
| 3 | Linear projection | Q, K, V = x @ W | Attention paper §3.2.1 |
| 4 | Dot-product similarity | scores = Q · K^T | Attention paper §3.2.1 |
| 5 | Scaling | scores / sqrt(head_dim) | Attention paper §3.2.1 |
| 6 | Softmax | weights = softmax(scores) | Attention paper §3.2.1 |
| 7 | Weighted sum | output = weights @ V | Attention paper §3.2.1 |
| 8 | Causal mask | mask future positions with -inf | Attention paper §3.1 |
| 9 | Multi-head attention | split into n_heads independent patterns | Attention paper §3.2.2 |
For the full concept stack (1–15) and how they compose into the transformer, see src/ai_playground/models/transformer.py.
See also: docs/PAPERS.md § Attention Variants · The Annotated Transformer
Attention in transformers, step-by-step (3Blue1Brown) animates exactly the query/key/value mechanics implemented below — more in Videos.
Master MHA, GQA, and MQA — what n_kv_heads buys in cache memory and what it trades in quality.
🏗️ Chapter milestone: M2 — choose your model’s attention shape against a KV-cache budget — the exercise at the end of this chapter; see The Project.
One session, ending in 🏗️ M2 — choose your n_kv_heads against the KV-cache budget. Cards, prerequisites, and done-when tests: the Session Guide.
# Setup — works from a repo checkout (via ../src) or standalone on Colab
import sys
sys.path.insert(0, '../src')
try:
import ai_playground # noqa: F401
except ImportError:
%pip install -q git+https://github.com/ggreg/ai_playground.git
%config InlineBackend.figure_format = 'retina'
from ai_playground.plotting import apply_plot_style
apply_plot_style()import sys
sys.path.insert(0, '../src')
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import time
import math
torch.manual_seed(42)
device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
print(f'Using device: {device}')Using device: mpsThe classic attention mechanism (§3.2.2 of the paper). For n_heads=8 and head_dim=64: - Q, K, V each have shape (batch, seq, n_heads, head_dim) - KV cache stores 2 * n_heads * head_dim values per token per layer
This is memory-intensive because every head has independent K and V projections.
Concepts implemented below: - Concept 3 (linear projection): Q = x @ W_q, K = x @ W_k, V = x @ W_v — three learned matrices project embeddings into query, key, and value spaces - Concept 9 (multi-head): .view(batch, seq, n_heads, head_dim) splits the single projection into n_heads independent attention patterns
def attention_manual(q, k, v, mask=None):
"""Standard scaled dot-product attention, written explicitly.
Implements concepts 4-8 from the first-principles list.
Reference: "Attention Is All You Need" §3.2.1
https://arxiv.org/abs/1706.03762
"""
# q, k, v: (batch, n_heads, seq_len, head_dim)
head_dim = q.size(-1)
# Concept 5 (scaling): precompute 1/sqrt(head_dim) to keep dot products
# from growing with dimension, which would saturate softmax.
scale = 1.0 / math.sqrt(head_dim)
# Concept 4 (dot-product similarity): Q @ K^T measures how relevant
# each key position is to each query position.
# (batch, n_heads, seq_q, head_dim) @ (batch, n_heads, head_dim, seq_k)
# -> (batch, n_heads, seq_q, seq_k)
scores = torch.matmul(q, k.transpose(-2, -1)) * scale
# Concept 8 (causal mask): set future positions to -inf so softmax
# gives them zero weight — token t can only attend to tokens 0..t.
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
# Concept 6 (softmax): normalize scores into probabilities that sum to 1
# for each query position. Higher score → more attention.
weights = F.softmax(scores, dim=-1)
# Concept 7 (weighted sum): blend value vectors by attention weights.
# This is how information flows — each token's output is a weighted
# combination of all values it attends to.
# (batch, n_heads, seq_q, seq_k) @ (batch, n_heads, seq_k, head_dim)
# -> (batch, n_heads, seq_q, head_dim)
output = torch.matmul(weights, v)
return output, weights
# Concept 8 (causal mask): lower triangular matrix of ones
def causal_mask(seq_len, device='cpu'):
return torch.tril(torch.ones(seq_len, seq_len, device=device)).unsqueeze(0).unsqueeze(0)
# Demo
batch, seq_len, n_heads, head_dim = 2, 32, 8, 64
dim = n_heads * head_dim # 512
x = torch.randn(batch, seq_len, dim, device=device)
# Concept 3 (linear projection): three separate weight matrices
# project the same input into Q, K, V spaces
W_q = torch.randn(dim, dim, device=device) * 0.02
W_k = torch.randn(dim, dim, device=device) * 0.02
W_v = torch.randn(dim, dim, device=device) * 0.02
# Concept 9 (multi-head): project then reshape into n_heads independent heads.
# Each head operates on head_dim dimensions (512 / 8 = 64 per head).
q = (x @ W_q).view(batch, seq_len, n_heads, head_dim).transpose(1, 2)
k = (x @ W_k).view(batch, seq_len, n_heads, head_dim).transpose(1, 2)
v = (x @ W_v).view(batch, seq_len, n_heads, head_dim).transpose(1, 2)
mask = causal_mask(seq_len, device)
out, weights = attention_manual(q, k, v, mask)
print(f'MHA output shape: {out.shape}') # (2, 8, 32, 64)
print(f'Attention weights shape: {weights.shape}') # (2, 8, 32, 32)
print(f'KV cache per layer per token: {2 * n_heads * head_dim} values = {2 * n_heads * head_dim * 2} bytes (fp16)')MHA output shape: torch.Size([2, 8, 32, 64])
Attention weights shape: torch.Size([2, 8, 32, 32])
KV cache per layer per token: 1024 values = 2048 bytes (fp16)MQA uses a single K and V head shared across ALL query heads. This reduces KV cache by n_headsx but can hurt quality slightly.
Used in: PaLM, Falcon, StarCoder
Paper: Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019)
# MQA: K and V have only 1 head
n_kv_heads_mqa = 1
W_k_mqa = torch.randn(dim, head_dim, device=device) * 0.02 # Only 1 head!
W_v_mqa = torch.randn(dim, head_dim, device=device) * 0.02
q_mqa = (x @ W_q).view(batch, seq_len, n_heads, head_dim).transpose(1, 2)
k_mqa = (x @ W_k_mqa).view(batch, seq_len, 1, head_dim).transpose(1, 2) # (batch, 1, seq, head_dim)
v_mqa = (x @ W_v_mqa).view(batch, seq_len, 1, head_dim).transpose(1, 2)
# Broadcasting: the single KV head is shared across all Q heads
# PyTorch broadcasts (batch, 1, seq, head_dim) against (batch, n_heads, seq, head_dim)
out_mqa, weights_mqa = attention_manual(q_mqa, k_mqa, v_mqa, mask)
print(f'MQA output shape: {out_mqa.shape}') # Same as MHA!
print(f'KV cache per layer per token: {2 * 1 * head_dim} values = {2 * 1 * head_dim * 2} bytes (fp16)')
print(f'KV cache reduction vs MHA: {n_heads}x')MQA output shape: torch.Size([2, 8, 32, 64])
KV cache per layer per token: 128 values = 256 bytes (fp16)
KV cache reduction vs MHA: 8xGQA is the middle ground: groups of Q heads share K/V heads. With n_kv_heads=4 and n_heads=8, each KV head serves 2 Q heads.
Used in: LLaMA 2/LLaMA 3, Mistral, Gemma
Paper: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)
n_kv_heads_gqa = 4
n_rep = n_heads // n_kv_heads_gqa # 2 Q heads per KV group
W_k_gqa = torch.randn(dim, n_kv_heads_gqa * head_dim, device=device) * 0.02
W_v_gqa = torch.randn(dim, n_kv_heads_gqa * head_dim, device=device) * 0.02
q_gqa = (x @ W_q).view(batch, seq_len, n_heads, head_dim).transpose(1, 2)
k_gqa = (x @ W_k_gqa).view(batch, seq_len, n_kv_heads_gqa, head_dim).transpose(1, 2)
v_gqa = (x @ W_v_gqa).view(batch, seq_len, n_kv_heads_gqa, head_dim).transpose(1, 2)
# Expand KV heads: (batch, n_kv, seq, dim) -> (batch, n_heads, seq, dim)
# Each KV head is repeated n_rep times
k_expanded = k_gqa[:, :, None, :, :].expand(batch, n_kv_heads_gqa, n_rep, seq_len, head_dim)
k_expanded = k_expanded.reshape(batch, n_heads, seq_len, head_dim)
v_expanded = v_gqa[:, :, None, :, :].expand(batch, n_kv_heads_gqa, n_rep, seq_len, head_dim)
v_expanded = v_expanded.reshape(batch, n_heads, seq_len, head_dim)
out_gqa, weights_gqa = attention_manual(q_gqa, k_expanded, v_expanded, mask)
print(f'GQA output shape: {out_gqa.shape}')
print(f'KV cache per layer per token: {2 * n_kv_heads_gqa * head_dim} values')
print(f'KV cache reduction vs MHA: {n_heads / n_kv_heads_gqa}x')GQA output shape: torch.Size([2, 8, 32, 64])
KV cache per layer per token: 512 values
KV cache reduction vs MHA: 2.0xEach heatmap below shows concept 6 (softmax) output: the attention weights for one head. Rows are query positions, columns are key positions. Bright = high attention.
The causal mask (concept 8) is visible as the black upper triangle — tokens can’t attend to future positions.
fig, axes = plt.subplots(1, 3, figsize=(18, 5))
for ax, w, title in [
(axes[0], weights[0, 0], 'MHA (head 0)'),
(axes[1], weights_mqa[0, 0], 'MQA (shared KV)'),
(axes[2], weights_gqa[0, 0], 'GQA (group 0)'),
]:
im = ax.imshow(w.detach().cpu().numpy(), cmap='viridis', aspect='auto')
ax.set_title(title, fontsize=14)
ax.set_xlabel('Key position')
ax.set_ylabel('Query position')
plt.colorbar(im, ax=ax, fraction=0.046)
plt.suptitle('Attention Weight Patterns (Causal Mask)', fontsize=16, y=1.02)
plt.tight_layout()
plt.show()
The KV cache is often the memory bottleneck during inference. Let’s compute how much memory each variant needs for realistic model sizes.
During autoregressive generation, we store K and V for every past token so we don’t recompute them. This is the tradeoff: memory for speed — without the cache, generation is O(n²); with it, each step is O(n). See inference/generate.py for the implementation, and PagedAttention for how serving systems manage this memory efficiently.
configs = [
('LLaMA-7B (MHA)', 32, 32, 128, 32), # n_heads, n_kv_heads, head_dim, n_layers
('LLaMA-2-7B (GQA)', 32, 4, 128, 32),
('LLaMA-2-7B (MQA)', 32, 1, 128, 32),
('LLaMA-3-8B', 32, 8, 128, 32),
('Mistral-7B', 32, 8, 128, 32),
]
seq_lengths = [1024, 4096, 16384, 32768, 131072]
print(f'{"Model":<22} | ' + ' | '.join(f'{s:>7}' for s in seq_lengths))
print('-' * 80)
for name, nh, nkv, hd, nl in configs:
kv_per_token = 2 * nl * nkv * hd * 2 # 2 for K+V, *2 for fp16 bytes
sizes = [kv_per_token * s / (1024**2) for s in seq_lengths] # in MB
print(f'{name:<22} | ' + ' | '.join(f'{s:>5.0f}MB' for s in sizes))
print(f'\n(Batch size = 1, FP16)')Model | 1024 | 4096 | 16384 | 32768 | 131072
--------------------------------------------------------------------------------
LLaMA-7B (MHA) | 512MB | 2048MB | 8192MB | 16384MB | 65536MB
LLaMA-2-7B (GQA) | 64MB | 256MB | 1024MB | 2048MB | 8192MB
LLaMA-2-7B (MQA) | 16MB | 64MB | 256MB | 512MB | 2048MB
LLaMA-3-8B | 128MB | 512MB | 2048MB | 4096MB | 16384MB
Mistral-7B | 128MB | 512MB | 2048MB | 4096MB | 16384MB
(Batch size = 1, FP16)Let’s measure actual forward pass time for each variant using our model implementation.
from ai_playground.models import Transformer, TransformerConfig
def benchmark_config(n_kv_heads, label, seq_len=512, batch=4, runs=20):
cfg = TransformerConfig(dim=512, n_layers=6, n_heads=8, n_kv_heads=n_kv_heads, max_seq_len=seq_len)
model = Transformer(cfg).to(device).eval()
x = torch.randint(0, cfg.vocab_size, (batch, seq_len), device=device)
# Warmup
with torch.no_grad():
for _ in range(3):
_ = model(x)
# Benchmark
if device == 'cuda':
torch.cuda.synchronize()
times = []
with torch.no_grad():
for _ in range(runs):
t0 = time.perf_counter()
_ = model(x)
if device == 'cuda':
torch.cuda.synchronize()
times.append(time.perf_counter() - t0)
avg_ms = sum(times) / len(times) * 1000
n_params = sum(p.numel() for p in model.parameters()) / 1e6
kv_cache_mb = 2 * cfg.n_layers * cfg.kv_heads * cfg.head_dim * seq_len * 2 / 1024**2
print(f'{label:<20} | {n_params:>6.1f}M params | {avg_ms:>8.2f}ms | KV cache: {kv_cache_mb:.1f}MB')
return avg_ms
print(f'{"Variant":<20} | {"Params":>13} | {"Latency":>10} | KV Cache')
print('-' * 75)
t_mha = benchmark_config(8, 'MHA (8 KV heads)') # Full MHA
t_gqa = benchmark_config(4, 'GQA (4 KV heads)') # GQA
t_gqa2 = benchmark_config(2, 'GQA (2 KV heads)') # More aggressive GQA
t_mqa = benchmark_config(1, 'MQA (1 KV head)') # Full MQA
print(f'\nSpeedup MQA vs MHA: {t_mha / t_mqa:.2f}x')Variant | Params | Latency | KV Cache
---------------------------------------------------------------------------
MHA (8 KV heads) | 47.9M params | 21.96ms | KV cache: 6.0MB
GQA (4 KV heads) | 46.3M params | 21.52ms | KV cache: 3.0MB
GQA (2 KV heads) | 45.6M params | 21.18ms | KV cache: 1.5MB
MQA (1 KV head) | 45.2M params | 20.71ms | KV cache: 0.8MB
Speedup MQA vs MHA: 1.06xServing memory is KV-cache memory. For your M1 config, compute KV bytes/token (2 x n_layers x n_kv_heads x head_dim x 4 in fp32) for MHA, GQA, and MQA variants, and choose the largest n_kv_heads that keeps batch 8 x 512 tokens under 8 MB of cache. Rerun this chapter’s quality comparison at your sizes before deciding — cheaper cache is not free. Update n_kv_heads in checkpoints/myllm/config.json.
Accept when: the recomputed cache size meets the budget and the config still passes M1’s parameter check. Part of The Project: Serve Your Own LLM.
Session S1.3 — verify: uv run pytest tests/milestones/test_m02_kv_budget.py.
n_heads. PaperFor the full reading list: docs/PAPERS.md
02_positional_encodings.ipynb — comparing RoPE, ALiBi, and sinusoidal embeddings