The Holy Grail ATTENTION FROM FIRST PRINCIPLES
one value. It’s not differentiable (the gradient of “did this key exactly match” is
zero almost everywhere).
Now suppose you have a “soft” version: every key has a similarity score to your
query, and the returned value is a weighted average of all the values, weighted
by similarity. This is differentiable: small changes in the query produce small
changes in the similarity scores, which produce small changes in the weighted
average.
Attention is exactly this. The “dictionary” is the sequence of tokens (each to-
ken contributes a key and a value). The “query” is the token we’re computing
the contextual representation for. The output is a weighted sum of the values,
weighted by how similar each key is to the query.
That’s it. The whole operation is “compute similarities, softmax-weight them,
weighted-sum the values.” Everything else in this chapter is filling in the de-
tails.
6.3 Queries, keys, and values, derived
Concretely. We have a sequence of token representations, each of dimension
d_model:
X in R^(s x d_model)
where s is the sequence length. From X, we want to produce three things: - A
query for each position — what this position is “looking for.” - A key for each
position — what this position “advertises about itself” to others. - A value for
each position — what this position “contributes” if attended to.
We produce them with three independent learned linear projections:
Q = X W_Q # shape (s, d_k)
K = X W_K # shape (s, d_k)
V = X W_V # shape (s, d_v)
W_Q, W_K, W_V are learned parameter matrices. d_k is the dimension of queries and
keys (they have to match because we’re going to take their dot product). d_v is
the dimension of values; in practice it’s almost always equal to d_k. The whole
point of these projections is to give the model the ability to project the same
input into three different spaces — one for asking, one for being asked, one for
being read.
Why three different projections? Because the same token might play three dif-
ferent roles. As a query, it might “ask” about its semantic context. As a key,
it might “advertise” its syntactic role. As a value, it might “contribute” its lexi-
69
The Holy Grail ATTENTION FROM FIRST PRINCIPLES
cal content. Letting the model learn three different views of the same input is
strictly more expressive than forcing them all to be the same.
6.4 The dot-product similarity
Now we have queries and keys. We need a similarity function between a query
vector and a key vector. The simplest reasonable choice is the dot product:
similarity(q, k) = q . k = Sigma_i q_i k_i
The dot product is large and positive when q and k point in the same direc-
tion, large and negative when they point in opposite directions, and zero when
they’re orthogonal. It’s a linear function of both arguments, which is differen-
tiable, and it’s cheap (one matmul).
For the whole sequence at once, we compute the matrix of all query-key dot
products:
QK^T in R^(s x s)
Entry (i, j) is the dot product of the i-th query and the j-th key. This is the
attention score matrix before normalization. Reading position i’s row tells
you “how much position i should attend to each other position.”
The cost of this matmul is O(s2 . d_k), and the result is a matrix of size s x s.
This is the source of the famous O(s2) cost of attention — both compute
and memory are quadratic in sequence length. We’ll come back to this in §6.10.
6.5 Why sqrtd_k? — the variance argument, derived
Here is the part that everyone gets wrong on the whiteboard. The actual atten-
tion formula divides the dot products by sqrtd_k:
scores = Q K^T / sqrtd_k
Why? It’s a numerical-stability argument. Suppose q and k are vectors of di-
mension d_k whose entries are independent random variables with mean 0 and
variance 1. (This is roughly what learned projections produce, especially early
in training.) The dot product is:
q . k = Sigma_i q_i k_i
This is a sum of d_k independent products of mean-0 unit-variance random
variables. Each product has mean 0 (since q_i and k_i are independent) and
variance 1 (since Var(XY) = E[X2]E[Y2] = 1 for independent unit-variance vari-
ables). The sum of d_k such products has mean 0 and variance d_k, which means
standard deviation sqrtd_k.
70
The Holy Grail ATTENTION FROM FIRST PRINCIPLES
So as d_k grows, the magnitude of the dot products grows as sqrtd_k. For typical
d_k = 64 or 128, the dot products are around +/-8 to +/-11. This is bad for soft-
max. Softmax is very sensitive at this scale: when the inputs are large, almost
all the probability mass concentrates on the single largest input. The gradient
through softmax in this regime is very small (tiny softmax probabilities have
tiny gradients). Training stalls.
The fix is to divide by sqrtd_k, which exactly cancels the variance growth:
Var(q . k / sqrtd_k) = d_k / d_k = 1
Now the scores have variance 1 regardless of d_k, softmax stays in a healthy
regime, and gradients flow.
This is one of the cleanest “small detail that makes the whole thing work” stories
in deep learning. It’s also a favorite interview question: “Why does scaled dot-
product attention divide by sqrtd_k?” — and “to keep the variance of the dot
products at 1 so softmax doesn’t saturate” is the right answer.
6.6 Softmax over the sequence axis
We have raw scores Q K^T / sqrtd_k of shape (s, s). We want to turn each
row into a probability distribution over positions, so that each query “votes”
for which positions to attend to in a way that sums to 1.
A = softmax(QK^T / sqrtd_k)
The softmax is applied along the last axis (the key axis), so each row of A is a
probability distribution. Entry A[i, j] is “the fraction of position i’s attention
that goes to position j.”
Two consequences of softmax that come up later:
• Each row sums to 1. The total amount of attention is conserved. If a
query attends more to one position, it must attend less to others.
• It’s smooth but peaked. Small differences in scores produce small
differences in attention weights, but if one score is much larger than the
others, almost all the weight goes to that one. This is what makes atten-
tion act like a soft lookup: usually it focuses on a few positions, but the
focus is differentiable.
6.7 Multi-head attention
Now we have one learned attention operation. We could just use it. But the
original transformer paper observed that letting the model run several in-
71
The Holy Grail ATTENTION FROM FIRST PRINCIPLES
dependent attention operations in parallel and concatenating their
outputs is strictly more expressive at the same parameter count.
The construction:
1. Pick a number of heads H. Typical values: 8, 16, 32, 64.
2. Split the model dimension d_model into H chunks of size d_h = d_model / H.
3. Run H independent attention operations, each with its own W_Q^h, W_K^h,
W_V^h, each producing an output of dimension d_h.
4. Concatenate the outputs from all heads: (s, H x d_h) = (s, d_model).
5. Apply one more linear projection W_O to mix the heads: (s, d_model) ->
(s, d_model).
The intuition: different heads can learn to attend to different things. Some
heads pick up syntactic relationships (“this verb’s subject is over there”), some
pick up coreference (“this pronoun refers to that noun”), some pick up posi-
tional patterns (“look at the previous token”), some attend uniformly. The
multi-head construction gives the model the freedom to specialize.
In practice, you don’t run H separate matmuls. You compute one big matmul
that produces all heads at once and reshape:
# x shape: (N, S, D)
qkv = qkv_proj(x) # (N, S, 3 * D)
qkv = [Link](N, S, 3, H, D_h) # (N, S, 3, H, D_h)
q, k, v = [Link](dim=2) # each (N, S, H, D_h)
q = [Link](1, 2) # (N, H, S, D_h)
k = [Link](1, 2) # (N, H, S, D_h)
v = [Link](1, 2) # (N, H, S, D_h)
This is where the canonical attention shape (N, H, S, D_h) from Chapter 1 comes
from. The H dimension exists so that all heads run as one batched operation.
6.8 Causal masking — autoregressive attention
For an autoregressive language model — a model that generates one token at a
time, conditioned on all previous tokens — there’s a constraint: position i can
only attend to positions j <= i. It can’t see the future, because at inference
time the future doesn’t exist yet, and at training time we’re trying to teach the
model to predict the future from the past.
We enforce this with a causal mask: before the softmax, we set every entry
of the score matrix where j > i to -inf. After softmax, those entries become
exactly 0, so position i puts zero attention weight on any position to its right.
mask = [Link]([Link](S, S), diagonal=1).bool() # upper triangle, excluding diagonal
scores = scores.masked_fill(mask, float('-inf'))
72
The Holy Grail ATTENTION FROM FIRST PRINCIPLES
attn = [Link](dim=-1)
The mask is the same for every batch element and every head, so it’s broadcast
across the leading dims. For very long sequences this is the dominant memory
cost, which is one of the reasons FlashAttention (Chapter 25) doesn’t material-
ize the full attention matrix at all.
A subtle but important point: the causal mask is what makes prefix
caching possible. Because each token only attends to its leftward context,
the K and V vectors for past tokens never need to be recomputed when a new
token is added. This is the foundation of the KV cache, which is the foundation
of efficient autoregressive serving. We’ll see this in Chapter 22.
6.9 Padding masks
When you batch multiple sequences of different lengths together, you pad the
shorter ones to the length of the longest. You don’t want the model to attend to
those padding positions — they have no meaning. The fix is a padding mask:
a per-sequence boolean mask that marks the padding positions, and the same
masked_fill to -inf trick before softmax.
# pad_mask shape: (N, S), True where padding
# scores shape: (N, H, S, S)
scores = scores.masked_fill(pad_mask[:, None, None, :], float('-inf'))
The mask is applied along the key axis (the last S). Padding positions can’t
be attended to. Whether a query at a padding position “produces” attention is
moot — the loss is masked at those positions too, so the gradients don’t care.
In practice, padding masks and causal masks are combined into a single
mask before being added to the scores. Modern attention implementations
(FlashAttention, PyTorch’s scaled_dot_product_attention) take both as separate
arguments and fuse them in the kernel.
6.10 The complexity story — O(s2) and where it bites
The cost of attention is dominated by the attention score matmul:
QK^T: (s, d_k) @ (d_k, s) -> (s, s) cost: O(s2 . d_k) per head
softmax: (s, s) cost: O(s2)
attention @ V: (s, s) @ (s, d_v) -> (s, d_v) cost: O(s2 . d_v) per head
Total compute per head: O(s2 . d_k + s2 . d_v) = O(s2 . d_k) (since d_k ~ d_v).
Total memory per head: O(s2 + s . d_k). The dominant term is the s x s score
matrix.
73
The Holy Grail ATTENTION FROM FIRST PRINCIPLES
This is the famous O(s2) complexity of attention. Both compute and
memory grow quadratically in sequence length. It’s why long contexts are ex-
pensive, why frontier model context windows have only crept up over time, and
why every research direction in efficient attention is some attempt to break this
s2.
For comparison, the linear-in-s parts of a transformer block (the projections,
the MLP) have cost O(s . d2). So at small s, the linear parts dominate, and the
model is “compute-bound” in the matmul sense. At large s, the attention part
dominates, and the model is “attention-bound.” The crossover happens around
s ~ d_model, which for a typical d_model = 4096 is around 4k tokens. Above that,
attention starts to eat the budget.
This is why so many papers chase sub-quadratic attention alternatives —
sliding window attention (longformer, big-bird), linear attention, ring atten-
tion, state-space models (Mamba, Chapter 39). None of them have completely
replaced full softmax attention for the highest-quality models, but they all have
niches.
The other reason O(s2) matters: memory. The attention matrix is the largest
single tensor in a transformer forward pass at long sequence lengths. For N=1,
H=32, S=8192, the score matrix in fp16 is 1 x 32 x 8192 x 8192 x 2 bytes ~
4.3 GB. For S=32768 it’s 70 GB. Materializing this tensor on every layer is what
makes long-context inference expensive — and what FlashAttention solved by
not materializing it at all.
6.11 Naive attention in 20 lines
Putting it all together, here is scaled dot-product multi-head attention in Py-
Torch, with no optimizations:
import torch
import [Link] as nn
import [Link] as F
import math
class NaiveAttention([Link]):
def __init__(self, d_model, num_heads):
super().__init__()
assert d_model % num_heads == 0
self.d_h = d_model // num_heads
self.h = num_heads
self.qkv_proj = [Link](d_model, 3 * d_model, bias=False)
self.out_proj = [Link](d_model, d_model, bias=False)
def forward(self, x, causal=True):
# x: (N, S, D)
74
The Holy Grail ATTENTION FROM FIRST PRINCIPLES
N, S, D = [Link]
H, D_h = self.h, self.d_h
qkv = self.qkv_proj(x) # (N, S, 3D)
qkv = [Link](N, S, 3, H, D_h).transpose(2, 0) # (3, N, S, H, D_h)
q, k, v = qkv[0], qkv[1], qkv[2] # each (N, S, H, D_h)
q = [Link](1, 2) # (N, H, S, D_h)
k = [Link](1, 2)
v = [Link](1, 2)
scores = q @ [Link](-2, -1) / [Link](D_h) # (N, H, S, S)
if causal:
mask = [Link]([Link](S, S, device=[Link]), diagonal=1).bool()
scores = scores.masked_fill(mask, float('-inf'))
attn = [Link](dim=-1) # (N, H, S, S)
out = attn @ v # (N, H, S, D_h)
out = [Link](1, 2).contiguous().view(N, S, D) # (N, S, D)
return self.out_proj(out)
Read this code line by line. Every modern attention implementation, no matter
how heavily optimized, is doing the same thing as these 20 lines. FlashAtten-
tion does it without materializing the (N, H, S, S) tensor. PagedAttention does
it with a memory layout that supports sharing prefixes. GQA does it with fewer
key/value heads than query heads. MLA does it with K and V compressed into
a low-rank latent. They are all variations of this snippet.
If you can write this snippet from memory, with the sqrtd_h divisor in the right
place, you can pass any “explain attention” interview question.
6.12 Forward pointers
This is the spine of the rest of the book. Every later chapter is some refinement
of the operation we just built:
• Chapter 7 wraps the attention block in residuals, normalization, and an
FFN to make a full transformer.
• Chapter 22 introduces the KV cache — the realization that during au-
toregressive decoding, the K and V for past tokens never change, so we
can store them and only compute one new K/V per step. This is the foun-
dation of efficient inference.
• Chapter 24 (PagedAttention) stores the KV cache in fixed-size blocks
like a virtual memory system, enabling prefix sharing and efficient batch-
ing.
• Chapter 25 (FlashAttention) rewrites the attention kernel to fuse QK^T
-> softmax -> attention @ V into a single tile-based operation that never
75
The Holy Grail ATTENTION FROM FIRST PRINCIPLES
materializes the s x s matrix. This single optimization is the most impor-
tant kernel-level improvement in LLM inference.
• Chapter 33 introduces GQA, MQA, and MLA — three ways to com-
press the K and V sides so the KV cache is smaller.
• Chapter 36 introduces disaggregated prefill/decode, the realiza-
tion that the prefill and decode phases of attention have such different
computational profiles that they should run on different GPU pools.
• Chapter 39 (state-space models) rejects attention’s O(s2) entirely
and replaces it with a different sequence operator.
Every one of these is “the same operation as in §6.11, but with one specific cost
optimized away.” Hold the naive implementation in your head as the baseline.
6.13 The mental model
Eight points to take into Chapter 7:
1. Attention is a soft, differentiable, learnable lookup table.
2. Three projections — query, key, value — give the model three views of the
same input.
3. Dot products measure similarity; the sqrtd_k divisor keeps softmax in
a healthy regime.
4. Softmax over the sequence axis turns scores into a probability dis-
tribution.
5. Multi-head lets the model learn multiple attention patterns in parallel
at no extra parameter cost.
6. Causal masking is what makes autoregressive generation work and
what makes the KV cache possible.
7. O(s2) in both compute and memory is the cost. Every later optimization
is trying to dodge it.
8. Naive attention is 20 lines. Every “fancy” implementation is the same
operation with one cost removed.
In Chapter 7 we wrap this into a full transformer.
Read it yourself
• The original paper: Vaswani et al., Attention Is All You Need (2017). Read
it cover to cover. It’s only ten pages and it’s the most important paper of
the last decade in ML.
76
The Holy Grail ATTENTION FROM FIRST PRINCIPLES
• Jay Alammar, The Illustrated Transformer — the visual companion to
the original paper. The diagrams alone are worth your time.
• Lilian Weng’s blog post Attention? Attention! — a long, rigorous walk
through every attention variant pre-2018.
• Andrej Karpathy’s Let’s build GPT YouTube video — the “build attention
from scratch” version, with code.
• The PyTorch source for [Link].scaled_dot_product_attention
— read the docstring, then the implementation in aten/src/ATen/native/transformers/.
Practice
1. Write scaled dot-product attention in PyTorch from memory. Compare
to §6.11. Don’t peek.
2. Why does softmax saturate when its inputs are large? Compute soft-
max([10, 11]) and softmax([100, 101]) — they should be the same in theory
but different in practice. Why?
3. Derive the gradient of softmax(z)_i with respect to z_j. (Answer: soft-
max(z)_i (delta_ij - softmax(z)_j).) You will find this useful when read-
ing FlashAttention later.
4. For d_model = 4096 and S = 8192 and N = 1 and H = 32, compute the size of
the Q K^T attention score tensor in fp16. (Answer: ~4.3 GB.)
5. Why does causal masking enable a KV cache during decoding? Walk
through a step of generation in your head and identify exactly what
doesn’t change.
6. The naive attention in §6.11 has an if causal: branch that allocates the
mask every forward pass. Why is this inefficient, and how would you fix
it?
7. Stretch: Implement multi-head attention from scratch in NumPy (no
PyTorch). Run it on a small toy input and verify the output shape and the
row-sums-to-1 property.
77
The transformer end to end
“A transformer is a residual stream with two operations bolted
onto it.”
In Chapter 6 we built attention from scratch. In this chapter we wrap it into
a full transformer block, stack the blocks into a transformer, and answer the
architectural questions you will be asked in interviews:
• What are the three transformer architectures, and why did decoder-only
win?
• Why pre-norm instead of post-norm?
• LayerNorm vs RMSNorm — what’s the difference and why does anyone
care?
• What is a position encoding, and why is RoPE everywhere now?
• What is the “residual stream,” and why is it the right mental model?
By the end of this chapter you will be able to read any modern open-source
LLM’s code and immediately recognize what each block is doing.
Outline:
1. The transformer block, end to end.
2. Pre-norm vs post-norm — and why pre-norm won.
3. LayerNorm and RMSNorm.
4. The FFN: from 4d to SwiGLU and the parameter count.
5. The residual stream as the central mental model.
6. Position encodings: absolute, learned, sinusoidal, RoPE.
7. Stacking blocks: a full transformer.
8. The three architectures: encoder, decoder, encoder-decoder.
9. Why decoder-only won.
10. The full pseudo-code of a Llama-style transformer.
78