Understanding Attention

Chapter 1: One Query, Many Keys

Forget sequences for a moment. Attention, at its smallest, is one query and a handful of key-value pairs. Get this picture right and every later chapter is bookkeeping on top of it.

Each pair (k_j, v_j) has two jobs, and they are different jobs. The key decides how much this pair is going to count. The value is what actually gets poured into the output. Two separate vectors, two separate learned projections, two separate spaces.

Notice that the query never touches a value directly. It only ever meets keys. And the values never influence a weight. The two halves of the mechanism meet exactly once, at the final sum.

Move the Query, Watch the Mixture Move

On the left is key space, where q and the three keys live. The dot product is a projection: how far q reaches along k_j, scaled by the length of k_j. On the right is value space, a completely different space, where the three values sit at the corners of a triangle. The output is the weighted average, so it can never leave that triangle.

Do It by Hand

Set the query to q = [1, 1] and use the keys and values we will meet again in Chapter 7:

The raw dot products are q · k₁ = 1, q · k₂ = 1 and q · k₃ = 1 + 1 = 2. With d_k = 2 we divide by √2 = 1.4142, giving scores 0.7071, 0.7071 and 1.4142.

Exponentiate: e^0.7071 = 2.0281, e^0.7071 = 2.0281, e^1.4142 = 4.1133. The sum is 8.1695. Divide each by the sum:

Now mix the values. First coordinate: 0.2483·1 + 0.2483·0 + 0.5035·3 = 0.2483 + 1.5105 = 1.7588. Second coordinate: 0.2483·0 + 0.2483·2 + 0.5035·1 = 0.4966 + 0.5035 = 1.0001. So o = [1.7587, 1.0000] once you carry the full precision.

Two Knobs, Two Different Effects

Slide v₃ in the widget. The weights do not move at all. The triangle deforms and the output slides with it, because you changed what is being mixed, not how much of it. Now move the query instead. The triangle holds still and the output walks across it, because you changed the mixture coefficients.

That is the cleanest demonstration of why attention carries both K and V. A system with keys only can decide relevance but has nothing to deliver. A system with values only can deliver but cannot decide. The learned projections W_K and W_V are trained to fill those two distinct roles from the same token representation.

Chapter 2: Shapes First

Most attention bugs are shape bugs. Not "I forgot the formula" bugs. Shape bugs. So before any more mathematics, we pin down every dimension and give each one a letter, exactly as the handbook and the PyTorch documentation do.

For a single head and a single example we can drop the batch and head axes entirely and write:

Then the three products follow with no freedom left:

Why the Transpose Is Forced

Matrix multiplication only works when the inner dimensions agree. Q is L × d_h. To contract it against K we need something whose first dimension is d_h, and K is S × d_h, so we transpose it to d_h × S. The result is L × S, one number for each query-key pair. That is the entire reason for the T in QK^T.

The same argument fixes the second product. A is L × S, V is S × d_v, the S contracts away, and the output is L × d_v. One row per query, and the row lives in value space.

Explore the Shape Algebra

Three presets matter. In self-attention over one full sequence, L = S. In cross-attention a decoder of length L attends to an encoder memory of length S, and the two are unrelated. In cached decode a single new query attends to the whole prefix, so L = 1 while S keeps climbing. Switch between them and watch which boxes change shape.

Drag L to 1 and leave S at 12. The score box becomes a single row of 12 numbers. No triangle, no square, no n × n. That is one decode step of a language model, and it is the most common attention call ever executed in production, by a wide margin.

The Full Batched Layout

Real kernels want the batch and the heads out of the way, so the standard memory layout puts them first:

Attention then runs as a batched matrix multiply over the leading two axes. This is the layout torch.nn.functional.scaled_dot_product_attention expects, and getting it wrong by leaving the sequence axis in position 1 instead of position 2 is mistake number one in every attention codebase.

Second, modern grouped attention allows H_q ≠ H_kv. Standard multi-head attention usually has the same number of query and key and value heads. Multi-query attention uses one KV head. Grouped-query attention uses an intermediate number. Move the two head sliders apart in the widget and watch the KV boxes shrink while the query box does not.

Chapter 3: Where Q, K, V Come From

So far Q, K and V have arrived by magic. They do not. They are three linear projections of the same input, and understanding that is what stops you believing the folklore about queries asking questions.

The input to an attention layer is a matrix of token representations X, of shape n × d_model. One row per token, one column per feature. A standard self-attention layer applies three learned linear maps:

For one head, W_Q and W_K are d_model × d_h, and W_V is d_model × d_v. In multi-head attention each head conceptually has its own three matrices, or, in practice, one larger projection is computed and reshaped into heads. The Transformer paper defines multi-head attention in terms of separate learned projections per head, and the reshape trick is an implementation detail that computes the same thing faster.

The Score Matrix, Cell by Cell

Take query row q_i and key row k_j. Their dot product is one number:

The complete matrix product QK^T computes every one of those query-key dot products in parallel. Every row corresponds to one query position. Every column corresponds to one key position. That is the whole content of the score matrix, and a single matmul on a GPU does all L × S of them at once.

Click any cell in the grid below, or use the two sliders, and the widget expands that one cell into its term-by-term arithmetic. Four tokens, head width 3, so every score is a sum of three products you can check on paper.

Two things to notice while you click around. First, the grid is not symmetric. S_ij uses query i against key j, and S_ji uses query j against key i, which are different vectors from different projections. Second, plenty of cells are negative. A dot product is free to be negative and often is.

The Word "Similarity" Is a Half Truth

Calling S_ij a similarity is convenient but incomplete. Dot-product attention learns W_Q and W_K specifically so that the resulting compatibility is useful for the task. It need not behave like a generic semantic similarity metric at all.

A concrete illustration. A head that implements "copy the token that follows the previous occurrence of the current token" will give a high score to a key whose surface meaning has nothing in common with the query's token. It scores high because the two projections were trained to align on a positional and structural relationship, not on meaning. If you insist on reading the score as semantic similarity you will misread that head completely.

Chapter 4: Why Divide by √d_k

The Transformer does not feed QK^T to softmax. It feeds QK^T divided by √d_k. That divisor is the single most asked about symbol in the equation, and it has a short, honest answer.

The Variance Derivation

Here is the paper's own motivation, in full. Suppose the components of q and k are independent random variables with mean zero and variance one. Then the dot product is a sum of d_k products:

Each term q_rk_r has mean zero, because q_r and k_r are independent and both have mean zero. Its variance is E[q_r²k_r²] = E[q_r²]E[k_r²] = 1 · 1 = 1. Independent terms add their variances, so:

So the typical size of a raw score grows like √d_k. At d_k = 64 the standard deviation is 8. At d_k = 1024 it is 32. Dividing by √d_k pulls that illustrative variance back to approximately one, whatever the head width.

What Big Logits Do to Softmax

Take four raw scores from a d_k = 64 head: 8.0, 0.1, −0.3 and 0.5. Exponentiate: e^8.0 = 2980.958, e^0.1 = 1.1052, e^−0.3 = 0.7408, e^0.5 = 1.6487. The sum is 2984.4527, so the weights are:

That is a one-hot vector in all but name. Now divide the same scores by √64 = 8, giving 1.0, 0.0125, −0.0375 and 0.0625. Exponentiate those: 2.7183, 1.0126, 0.9632, 1.0645, sum 5.7585, weights:

The first key still wins, and wins clearly. But the other three still exist. The gradient story is the sharp version of this: the derivative of a softmax weight with respect to its own logit is a(1 − a). At a = 0.9988 that is 0.00117. At a = 0.4720 it is 0.2492, over two hundred times larger. A saturated softmax barely learns.

Sample It Yourself

The widget runs a real Monte Carlo. It draws thousands of q and k vectors with unit-variance components, computes the dot products, and histograms them. The warm histogram is unscaled, the teal one is divided by √d_k. Drag d_k and watch the warm distribution spread out while the teal one stands still.

The Same Experiment in PyTorch

This block runs as written. It draws twenty thousand query rows against four keys each, at four head widths, and reports the variance of the raw dot product, the variance after scaling, and the average largest softmax weight in both cases.

python
import torch

torch.manual_seed(0)
rows, keys = 20000, 4          # 20000 query rows, 4 keys each

print("  d_k   var(q.k)   var(scaled)   mean top weight   scaled")
for d_k in (4, 64, 256, 1024):
    q = torch.randn(rows, 1, d_k)      # unit variance, mean zero
    k = torch.randn(rows, keys, d_k)
    s = (q * k).sum(-1)                # raw dot products, one row per query
    scaled = s / d_k ** 0.5
    top = torch.softmax(s, dim=-1).max(-1).values.mean()
    top_scaled = torch.softmax(scaled, dim=-1).max(-1).values.mean()
    print("%5d  %9.2f  %12.2f  %16.4f  %7.4f"
          % (d_k, s.var().item(), scaled.var().item(), top.item(), top_scaled.item()))

  d_k   var(q.k)   var(scaled)   mean top weight   scaled
    4       4.03          1.01            0.6525   0.4972
   64      63.72          1.00            0.9105   0.5165
  256     255.52          1.00            0.9557   0.5172
 1024    1023.10          1.00            0.9782   0.5169

Read the last two columns. Unscaled, the average winning weight climbs from 0.65 to 0.98 as the head gets wider, purely because the head got wider. The model has not learned anything; the geometry did it. Scaled, that number sits at about 0.51 at every width, so the head's sharpness is something the model chooses through its weights rather than something the dimension imposes.

Chapter 5: Masks, or Which Connections Are Allowed

Attention as defined so far lets every query see every key. Often that is wrong. A language model must not let a token peek at its own future. A batch with padded short sequences must not attend to padding. A long document model may deliberately restrict each query to a window.

All of these are the same intervention: before softmax, add a mask or bias to the score matrix. For causal language modelling the standard additive mask is:

Because e^−∞ = 0, masked entries receive exactly zero probability after softmax, and the remaining entries renormalize among themselves. The original Transformer decoder implements autoregressive masking exactly this way: illegal pre-softmax connections are set to −∞.

Add the Mask, Then Normalize. Never the Other Way

Take the second query row of the example we keep returning to. Its three raw scaled scores are 0, 0.7071 and 0.7071. Exponentiate: 1.0000, 2.0281, 2.0281, sum 5.0562. The unmasked softmax is:

Query 2 is not allowed to see key 3. Do it correctly, by setting the third logit to −∞ first: the third exponential becomes 0, the sum becomes 3.0281, and the row is:

Now do it the tempting way, by running softmax first and then zeroing the illegal entry. You get [0.1978, 0.4011, 0], which sums to 0.5989. The row no longer sums to one, so the output vector is silently scaled down by about forty percent, and every residual connection downstream inherits the shrunken vector. The fix is to renormalize, which recovers 0.3302 and 0.6698, but a lot of code forgets that step and the bug is quiet.

Paint a Mask

The grid below is the allowed-connection pattern for the current mask. Switch between the four families and watch which cells survive. Toggle the view to see the same mask written as an additive floating-point bias rather than as Booleans.

The Non-Square Trap

Set L to 3 and S to 8 in the widget, then press Flip alignment. Both grids are "lower triangular" in some sense, and they mean completely different things. Aligned to the top left, query 0 sees only key 0. Aligned to the bottom right, query 0 sees keys 0 through 5, because it is really the sixth token of a sequence whose first five keys are already in the cache.

For cached decoding the bottom-right reading is the one that matches reality: a query block of length L sitting at the end of a prefix of length S should see everything up to and including its own position, which means an offset of S − L on every row.

Four Families, One Mechanism

Notice that only the first two are about correctness. The last two are about cost, and they change the model, because a connection that is masked out cannot carry information in that layer at all. Chapter 11 comes back to that distinction, because it is the line between FlashAttention and sparse attention.

Chapter 6: Row-wise Softmax and Aggregation

Two operations remain, and the second one is the one people skip. Softmax turns scores into weights. Then AV turns weights into an output. The first step gets all the attention. The second step is where information actually moves.

Look carefully at the index that is being summed over. It is r, the key index, and the sum runs across one row. So for each query i the weights are nonnegative and sum to one across the key positions:

Rows, Not Columns

Why row-wise? Because each query independently produces a distribution over its permitted key positions. That is the object we want: a set of mixing coefficients for one output. A column-wise softmax would normalize across queries instead, answering "which query cares most about this key," which is a different question and a different operation.

The consequence is concrete. Under column normalization a single query's coefficients no longer sum to one, so its output is not a weighted average of values at all. It is an unnormalized combination whose scale depends on how popular its keys happen to be with other queries. Flip the toggle in the widget and watch the row sums stop being one.

The Multiplication That Moves Information

The output at query position i is:

This final multiplication is easy to overlook, but it is where attention actually moves information. The score matrix only decides how much each value contributes. The values are the vectors that get mixed into the new representation.

If a particular query gives weights (0.7, 0.2, 0.1) to three values, its output is 0.7v₁ + 0.2v₂ + 0.1v₃. With our running values that is:

The output remains in the value-vector space, not in score space. Its dimension is d_v. Its units, if the values had units, would be the units of the values. Nothing about the score survives into the output except through the coefficients.

Try to push the output outside the triangle. You cannot. Softmax weights are nonnegative and sum to one, so the output is a convex combination, and convex combinations live inside the convex hull of their inputs. A single attention head can only ever produce an interpolation of the values available to it. Everything else a Transformer does, including the feed-forward block and the residual stream, exists partly to escape that limitation.

Chapter 7: The Worked Example, Every Number

This is the chapter that turns the equation into something you own. We compute one causal attention head completely by hand, three tokens, head width two, and every intermediate number stays on screen. If you can explain where each number comes from, the core attention equation is no longer a black box.

Three query positions and three key and value positions, so this is same-length self-attention and the score matrix is 3 × 3. Setting Q = K is a deliberate simplification: it means position 3's key is [1, 1], which is compatible with both earlier keys, and that shows up in the final weights.

Step 1: The Raw Dot Products

Every entry is a query row dotted with a key row. Row 1 is q₁ = [1, 0]. Against k₁ = [1, 0] that is 1·1 + 0·0 = 1. Against k₂ = [0, 1] it is 1·0 + 0·1 = 0. Against k₃ = [1, 1] it is 1·1 + 0·1 = 1. Carry on for the other two rows:

Step 2: Divide by √2

d_k = 2, so the divisor is √2 = 1.41421. Every entry gets smaller by that factor: 1 becomes 0.7071, 2 becomes 1.4142, and 0 stays 0.

Step 3: Impose the Causal Mask

Query 1 may use only key 1. Query 2 may use keys 1 and 2. Query 3 may use all three. Everything above the diagonal becomes −∞:

Step 4: Softmax, Row by Row

Row 1. Only one position is legal, so whatever its logit is, it takes all the mass. A₁ = [1, 0, 0]. This is worth sitting with: the first token of every causal sequence always attends entirely to itself, in every head of every layer, forever.

Row 2. The legal logits are [0, 0.7071]. Exponentials are e⁰ = 1.0000 and e^0.7071 = 2.0281, sum 3.0281. Divide: 1.0000/3.0281 = 0.3302 and 2.0281/3.0281 = 0.6698. So A₂ = [0.3302, 0.6698, 0].

Row 3. The logits are [0.7071, 0.7071, 1.4142]. Exponentials are 2.0281, 2.0281 and 4.1133, sum 8.1695. Divide: 2.0281/8.1695 = 0.2483 twice, and 4.1133/8.1695 = 0.5035. So A₃ = [0.2483, 0.2483, 0.5035].

Every row sums to one up to rounding, and every causally forbidden position has weight exactly zero. Read what the third row says: the third query assigns about half its attention weight to the third value and about a quarter to each earlier value. These numbers are not annotations we added. They are the normalized result of the learned-compatibility calculation in this toy setup.

Step 5: Multiply by V

Finally O = AV. Each output row is a weighted sum of value rows.

Row 1: 1·[1, 0] = [1, 0]. The first token copies its own value exactly, which follows from its one-hot weight row.

Row 2: 0.3302·[1, 0] + 0.6698·[0, 2] = [0.3302, 1.3396] when you round at four places, and [0.3302, 1.3395] when you carry full precision through.

Row 3: 0.2483·[1, 0] + 0.2483·[0, 2] + 0.5035·[3, 1]. First coordinate: 0.2483 + 0 + 1.5105 = 1.7588. Second coordinate: 0 + 0.4966 + 0.5035 = 1.0001. Carrying full precision gives [1.7587, 1.0000].

Step Through It Yourself

The showcase widget walks the five steps with every matrix on screen. Use Next and Back, or drag the step slider. Pick a row to see that row's arithmetic spelled out underneath, term by term.

The Same Numbers, Three Ways

Here is the implementation ladder. The same computation in plain Python with no libraries, in NumPy, and in PyTorch. All three print the same A and the same O, and all three agree with the hand arithmetic above. Run them in that order and the abstraction stops feeling like magic, because you wrote the bottom rung yourself.

Rung 1: plain Python, no imports beyond the standard library.

python
import math

Q = [[1, 0], [0, 1], [1, 1]]
K = [[1, 0], [0, 1], [1, 1]]
V = [[1, 0], [0, 2], [3, 1]]
d_k = 2

def dot(a, b):
    return sum(x * y for x, y in zip(a, b))

def softmax(row):
    m = max(row)
    e = [math.exp(v - m) for v in row]
    z = sum(e)
    return [x / z for x in e]

S  = [[dot(q, k) for k in K] for q in Q]
Sc = [[s / math.sqrt(d_k) for s in r] for r in S]
Sm = [[Sc[i][j] if j <= i else -math.inf for j in range(3)] for i in range(3)]
A  = [softmax(r) for r in Sm]
O  = [[sum(A[i][j] * V[j][c] for j in range(3)) for c in range(2)] for i in range(3)]

def show(name, M):
    print(name)
    for r in M:
        print("   " + "  ".join("%8.4f" % v for v in r))

show("QK^T", S)
show("QK^T / sqrt(2)", Sc)
show("A = softmax(masked)", A)
show("O = A V", O)

QK^T
     1.0000    0.0000    1.0000
     0.0000    1.0000    1.0000
     1.0000    1.0000    2.0000
QK^T / sqrt(2)
     0.7071    0.0000    0.7071
     0.0000    0.7071    0.7071
     0.7071    0.7071    1.4142
A = softmax(masked)
     1.0000    0.0000    0.0000
     0.3302    0.6698    0.0000
     0.2483    0.2483    0.5035
O = A V
     1.0000    0.0000
     0.3302    1.3395
     1.7587    1.0000

Two details in that code earn their place. Subtracting the row maximum before exponentiating is the standard numerically stable softmax: it cannot change the result, because a common factor cancels between numerator and denominator, and it prevents e^large from overflowing. And -math.inf works directly, because math.exp(-inf - m) is exactly 0.0.

Rung 2: NumPy, where the loops become array operations.

python
import numpy as np

Q = K = np.array([[1., 0.], [0., 1.], [1., 1.]])
V = np.array([[1., 0.], [0., 2.], [3., 1.]])

S = Q @ K.T / np.sqrt(Q.shape[-1])
S = np.where(np.tril(np.ones_like(S)) == 1, S, -np.inf)   # causal: keep j <= i
S = S - S.max(axis=-1, keepdims=True)                     # stable softmax
E = np.exp(S)
A = E / E.sum(axis=-1, keepdims=True)
O = A @ V

np.set_printoptions(precision=4, suppress=True)
print("A =\n", A)
print("O =\n", O)
print("row sums:", A.sum(axis=-1))

A =
 [[1.     0.     0.    ]
 [0.3302 0.6698 0.    ]
 [0.2483 0.2483 0.5035]]
O =
 [[1.     0.    ]
 [0.3302 1.3395]
 [1.7587 1.    ]]
row sums: [1. 1. 1.]

Rung 3: PyTorch, the transparent implementation. This is the handbook's own code, wrapped in a function and applied to our three tokens. It mirrors the mathematics rather than optimizing it.

python
import math
import torch

def attention(q, k, v, causal=False):
    # q, k, v: (..., L/S, d_head)
    scores = q @ k.transpose(-2, -1)
    scores = scores / math.sqrt(q.size(-1))
    if causal:
        L, S = q.size(-2), k.size(-2)
        # This simple mask assumes a square aligned case.
        mask = torch.triu(torch.ones(L, S, dtype=torch.bool, device=q.device), diagonal=1)
        scores = scores.masked_fill(mask, float("-inf"))
    weights = torch.softmax(scores, dim=-1)
    return weights @ v, weights

q = k = torch.tensor([[1., 0.], [0., 1.], [1., 1.]])
v = torch.tensor([[1., 0.], [0., 2.], [3., 1.]])

out, A = attention(q, k, v, causal=True)
torch.set_printoptions(precision=4, sci_mode=False)
print("A =")
print(A)
print("O =")
print(out)

A =
tensor([[1.0000, 0.0000, 0.0000],
        [0.3302, 0.6698, 0.0000],
        [0.2483, 0.2483, 0.5035]])
O =
tensor([[1.0000, 0.0000],
        [0.3302, 1.3395],
        [1.7587, 1.0000]])

The code exposes the exact conceptual stages: QK^T, then scale, then mask, then softmax, then AV. What it does not expose is how a production kernel fuses those operations, handles numerically stable blockwise softmax, chooses hardware backends, represents causal bias in non-square cases, or supports grouped query heads. Chapters 10 through 12 are about exactly those omissions.

Chapter 8: Self, Cross, Causal, and Many Heads

The attention equation never said Q, K and V had to come from the same sequence. That is a choice made by the layer around it, and swapping that choice gives you a different named mechanism with the same arithmetic.

Two Streams, Three Configurations

In self-attention, all three projections read the same sequence, or more precisely the same layer input:

The Transformer encoder uses self-attention in this form. A causal decoder also uses self-attention, but masks future positions.

In cross-attention, the queries come from the decoder side while keys and values come from the encoder outputs:

If the target stream has length L and the source has length S, the score matrix is L × S, exactly as Chapter 2 promised.

Many Heads, and What They Actually Cost

A single attention head produces one learned compatibility-and-aggregation pattern. Multi-head attention performs several attention operations in parallel using different learned projections, concatenates the resulting head outputs, and applies an output projection:

Here is the arithmetic that surprises people. If d_model is held fixed and each head has dimension about d_model/H, then increasing the number of heads does not simply multiply the full-dimensional attention cost by H, because each head is narrower. The original Transformer used this design specifically to keep total cost comparable to a single full-dimensional head.

Count it. With d_model = 512 and H = 8, each head is 64 wide. The score computation for one head is L · S · 64 multiply-adds, and there are 8 heads, giving L · S · 512. With H = 1 and a head of width 512 you get L · S · 512 as well. Same number. The head count moved the work around; it did not add work.

Concat, Then Project

Each head produces an L × d_v output. Concatenating H of them along the feature axis gives L × H·d_v, which is L × d_model when d_v = d_model/H. Then W_O, of shape d_model × d_model, mixes across heads.

W_O is not decoration. Without it the head outputs would sit in disjoint slices of the residual stream and could never combine. It is the one place where information from different heads meets inside the layer.

Chapter 9: Position, Then Training versus Decoding

Plain noncausal self-attention, with no position-dependent signal anywhere, is permutation-equivariant. Permute the input rows and the outputs are permuted the same way, unchanged in value. It has no inherent notion that one token is "third" or that another is "five positions away."

The original Transformer therefore added positional encodings to the token embeddings. Modern decoder-only LLMs often use Rotary Position Embedding, or RoPE, which rotates query and key components by position-dependent angles. In a simplified two dimensional block:

Now set q_m = R_mq and k_n = R_nk and compute their dot product. Rotations are orthogonal, so R_m^TR_n = R_n−m, and:

The absolute positions cancel and the attention interaction contains relative displacement information. RoFormer presents RoPE as encoding absolute position through rotations while introducing explicit relative-position dependence in self-attention.

Check It by Hand

Take q = [1, 0], k = [0, 1] and θ = 30 degrees. Rotating q by m steps gives [cos(30m), sin(30m)]. Rotating k by n steps gives [−sin(30n), cos(30n)]. Their dot product is

At m = 2 and n = 1 that is sin(30°) = 0.5. At m = 7 and n = 6, wildly different absolute positions but the same displacement, it is sin(30°) = 0.5 again. Drag both sliders together in the widget and the score readout does not move.

Causal Masking Parallelises Training

An autoregressive language model factorizes a token sequence as a product of conditionals:

When training a decoder-only Transformer on a full sequence, the model can process many positions in parallel, because a causal mask prevents the representation at position t from using future tokens. The masked self-attention at each layer enforces the allowed dependency pattern. For query position t, attention can use keys and values from positions 1 through t, but not t+1 through T.

This is one reason training and generation feel so different computationally. During training, the full known sequence is available and large matrix multiplications process all query positions together. During generation, token t+1 does not exist until token t has been generated, so decoding is sequential across generated tokens.

Prefill and Decode Are Serving Phases

Consider a causal decoder layer processing a prefix of length T. There are two regimes, and every serving system in production distinguishes them.

In the prefill or training-like pass, many query positions are processed together. Dense attention forms interactions across the relevant query-key pairs, typically through large batched matrix operations.

In incremental decoding, after the prefix has been processed, the model generates a new token. Its new hidden state produces a new query, key and value at each layer. Crucially, the keys and values for earlier positions do not need to be recomputed from scratch: because of causal attention, earlier token representations were already determined without access to future tokens. That single observation leads to the KV cache, which is Chapter 10.

Chapter 10: The KV Cache and the Head-Count Decision

Chapter 9 ended on an observation: during causal decoding, the keys and values of earlier positions never change. They were computed without access to future tokens, and no future token can alter them. So do not recompute them. Store them.

That is the KV cache. Each layer keeps the key and value vectors produced for prior tokens. When a new token arrives, the layer computes the new token's query, key and value, appends its key and value to the cache, and uses the new query to attend over all cached keys and values.

Watch the Cache Reproduce the Worked Example

The clearest proof that caching is exact rather than approximate is to run it on the three token example from Chapter 7. Feed one token at a time, keep the cache, and compare the rows.

python
import math
import torch

def step(q_new, k_cache, v_cache):
    # q_new: (1, d_head). No mask: everything in the cache is already past.
    scores = (q_new @ k_cache.transpose(-2, -1)) / math.sqrt(q_new.size(-1))
    w = torch.softmax(scores, dim=-1)
    return w @ v_cache, w

Q = K = torch.tensor([[1., 0.], [0., 1.], [1., 1.]])
V = torch.tensor([[1., 0.], [0., 2.], [3., 1.]])

k_cache = torch.empty(0, 2)
v_cache = torch.empty(0, 2)
torch.set_printoptions(precision=4, sci_mode=False)

for t in range(3):
    k_cache = torch.cat([k_cache, K[t:t + 1]], dim=0)   # append this token's K
    v_cache = torch.cat([v_cache, V[t:t + 1]], dim=0)   # append this token's V
    o, w = step(Q[t:t + 1], k_cache, v_cache)
    print("t=%d  cache=%d rows  weights=%s  out=%s"
          % (t, k_cache.size(0),
             [round(x, 4) for x in w[0].tolist()],
             [round(x, 4) for x in o[0].tolist()]))

print("cache elements after 3 tokens:", k_cache.numel() + v_cache.numel())

t=0  cache=1 rows  weights=[1.0]  out=[1.0, 0.0]
t=1  cache=2 rows  weights=[0.3302, 0.6698]  out=[0.3302, 1.3395]
t=2  cache=3 rows  weights=[0.2483, 0.2483, 0.5035]  out=[1.7587, 1.0]
cache elements after 3 tokens: 12

Every row matches Chapter 7 exactly. Notice two things. There is no mask anywhere in that loop, because the cache contains only past positions, so causality is enforced by what is in the cache rather than by a triangle. And the weight row grows by one entry per step, which is the cost that never goes away.

How Big Does the Cache Get

For standard attention, an approximate cache-element count per token is:

The factor 2 accounts for K and V. In bytes, multiply by the storage bytes per element. Take a concrete model: 32 layers, d_h = 128, two-byte elements. With 32 KV heads that is 2 · 32 · 32 · 128 · 2 = 524288 bytes, which is 512 KiB for every single token.

python
layers, d_h, bytes_per_elem, tokens = 32, 128, 2, 8192

print("KV heads   cache / token      cache for 8192 tokens")
for h_kv in (32, 8, 1):
    per_token = 2 * layers * h_kv * d_h * bytes_per_elem   # bytes
    total = per_token * tokens
    print("%8d   %8.0f KiB %18.3f GiB"
          % (h_kv, per_token / 1024, total / 1024 ** 3))

KV heads   cache / token      cache for 8192 tokens
      32        512 KiB              4.000 GiB
       8        128 KiB              1.000 GiB
       1         16 KiB              0.125 GiB

Four gigabytes of cache for one sequence of 8192 tokens, before you have loaded a single weight. That is why the number of key and value heads stopped being an architecture detail and became an inference-systems decision.

MHA, MQA, GQA: Reduce the KV Heads, Not the Query Heads

Multi-head attention has separate key and value projections per attention head. If there are H query heads there are typically H key and value heads as well, and the cache formula above uses H_kv = H.

Multi-query attention was proposed by Shazeer for faster incremental Transformer decoding. It keeps multiple query heads but shares a single key head and a single value head across them, which greatly shrinks cached key and value tensors and memory-bandwidth demands.

Grouped-query attention generalizes MQA by using an intermediate number of key and value heads: more than one, but fewer than the query heads. Query heads are divided into groups that share a KV head. Ainslie and colleagues introduced GQA as a way to approach MQA-like inference efficiency while retaining quality closer to MHA in their experiments.

The code below shows how a grouped layer is actually implemented. The KV tensors are stored with H_kv heads and expanded to H_q heads on the fly, or handed to a kernel that understands grouping directly. Both paths give bit-identical results.

python
import torch
import torch.nn.functional as F

B, L, S, Hq, Hkv, dh = 1, 4, 6, 8, 2, 16
torch.manual_seed(0)

q = torch.randn(B, Hq, L, dh)
k = torch.randn(B, Hkv, S, dh)
v = torch.randn(B, Hkv, S, dh)

group = Hq // Hkv                      # 4 query heads share one KV head
k_rep = k.repeat_interleave(group, dim=1)   # (B, Hq, S, dh)
v_rep = v.repeat_interleave(group, dim=1)

out_manual = F.scaled_dot_product_attention(q, k_rep, v_rep)
out_gqa = F.scaled_dot_product_attention(q, k, v, enable_gqa=True)

print("q     ", tuple(q.shape))
print("k, v  ", tuple(k.shape), "  KV heads =", Hkv)
print("k_rep ", tuple(k_rep.shape), " after repeat_interleave")
print("out   ", tuple(out_gqa.shape))
print("group size Hq / Hkv =", group)
print("max abs difference:", (out_manual - out_gqa).abs().max().item())
print("KV elements cached per token: MHA", 2 * Hq * dh, " GQA", 2 * Hkv * dh)

q      (1, 8, 4, 16)
k, v   (1, 2, 6, 16)   KV heads = 2
k_rep  (1, 8, 6, 16)  after repeat_interleave
out    (1, 8, 4, 16)
group size Hq / Hkv = 4
max abs difference: 0.0
KV elements cached per token: MHA 256  GQA 64

MLA: Compress What You Cache

Multi-head Latent Attention, introduced in DeepSeek-V2, approaches KV cache reduction differently from MQA and GQA. Instead of primarily sharing key and value heads, MLA uses low-rank joint compression of key and value information. A compressed latent is formed from the hidden state,

and key and value content is then derived through learned up-projections. The core motivation is that the cached state can be much smaller than storing conventional per-head keys and values.

Chapter 11: What "Quadratic" Means, and FlashAttention

"Attention is quadratic" is one sentence doing the work of three. Pulling those three apart is the difference between understanding why FlashAttention exists and repeating a slogan about it.

For dense self-attention with sequence length n and representation and head dimensions held fixed, the original Transformer analysis gives a per-layer self-attention term of order n²d. But three separate quantities hide inside that:

These distinctions explain how two statements can both be true. Dense attention still performs quadratic pairwise interactions. And an exact attention implementation can avoid materializing the entire quadratic attention matrix in high-bandwidth memory and run much faster. That second statement is the core of FlashAttention.

The Memory Hierarchy Is the Point

FlashAttention is an IO-aware exact attention algorithm. Its key insight is to account explicitly for the GPU memory hierarchy: transfers between high-bandwidth memory and faster on-chip SRAM can be a dominant cost. Instead of materializing the full attention matrix in HBM, FlashAttention tiles the computation so that blocks of Q, K and V and intermediate statistics can be processed on chip, reducing HBM reads and writes.

The mathematical target remains scaled dot-product attention. The algorithm uses an online, blocked softmax strategy so that exact results can be assembled without storing all scores at once in HBM.

The multiply-add counter does not move when you switch strategies. That is the whole point. The tiled path performs the same arithmetic on the same numbers and produces the same answer. What collapses is the peak intermediate held in HBM, from the full n × n matrix down to one tile, and with it the traffic between the two memory levels.

Online Softmax, and Why It Is Exact

The trick that makes tiling possible is a running softmax. Softmax needs the maximum and the sum of exponentials over the whole row, which you do not have while you are still streaming tiles. So carry three running quantities: the maximum seen so far, the running sum of exponentials, and the running weighted sum of values. When a new tile arrives with a larger maximum, rescale the two accumulators by e^{old max − new max} and carry on.

The code below does exactly that against a full-row reference, on a real 512 key row, and reports the difference.

python
import math
import torch

torch.manual_seed(0)
S, d = 512, 64
q = torch.randn(d)
K = torch.randn(S, d)
V = torch.randn(S, d)
scale = 1.0 / math.sqrt(d)

# Reference: materialise the whole score row, then softmax, then mix.
ref = torch.softmax(K @ q * scale, dim=-1) @ V

# Tiled online softmax: never hold more than one tile of scores.
tile = 64
m = torch.tensor(float("-inf"))   # running max
l = torch.tensor(0.0)             # running sum of exponentials
acc = torch.zeros(d)              # running weighted sum of values
peak = 0
for start in range(0, S, tile):
    s = K[start:start + tile] @ q * scale     # one tile of scores
    peak = max(peak, s.numel())
    m_new = torch.maximum(m, s.max())
    correction = torch.exp(m - m_new)
    p = torch.exp(s - m_new)
    l = l * correction + p.sum()
    acc = acc * correction + p @ V[start:start + tile]
    m = m_new
out = acc / l

print("scores held at once: full row %d, tiled %d" % (S, peak))
print("max abs difference:", (ref - out).abs().max().item())

scores held at once: full row 512, tiled 64
max abs difference: 5.960464477539063e-08

Six times ten to the minus eight is float32 rounding, not an approximation error. The two paths compute the same function. One of them just never held more than 64 scores at a time.

Generations: Optimize the Bottleneck You Actually Have

The durable concept is IO-aware exact attention. Later FlashAttention versions refine how that computation maps to newer GPU hardware.

FlashAttention-2 improved parallelism and work partitioning, including better distribution of work across thread blocks and warps and fewer non-matrix-multiply operations.

FlashAttention-4, introduced in 2026 for Blackwell-class GPUs, addresses a different hardware balance. Tensor-core throughput had grown faster than some surrounding resources, so bottlenecks shifted toward softmax and other non-matmul work, shared-memory traffic, and pipeline structure. Its contribution is therefore not a new definition of attention but a hardware-aware algorithm and kernel co-design for a changed machine.

The evergreen version of the lesson, in five lines:

Benchmark numbers from individual FlashAttention generations are intentionally omitted here, because they are hardware, shape, precision and software-version dependent. A speedup quoted without all four of those is not a fact you can carry anywhere.

Sparse and Local Attention Change Something Else Entirely

Dense attention connects every query to every key, producing quadratically many pairwise interactions for same-length self-attention. A completely different family of approaches changes the connectivity pattern itself.

Longformer combines local windowed attention with task-motivated global attention, and was designed so that attention cost scales linearly with sequence length under its fixed-window pattern. BigBird combines sparse components such as local, global and random connections, also reducing the full quadratic dependency while retaining important theoretical properties under its design.

Sparse attention can reduce arithmetic work by not computing all pairs. The trade-off is that the model no longer has unrestricted direct connectivity in a single layer, so pattern design matters. A token pair that is never scored in layer 7 can only communicate through some multi-hop path across layers, if one exists at all.

Chapter 12: Code, Traps, and the Complete Mental Model

You can now derive attention, mask it, cache it, split it into heads and explain why a kernel tiles it. This chapter is the part you keep next to the keyboard: the production call, the seven ways people break it, the eight things people believe about it that are not true, and the six verbs that hold the whole thing together.

The Production Call, and Its One Sharp Edge

Current PyTorch exposes scaled dot-product attention as:

signature
torch.nn.functional.scaled_dot_product_attention(
    query, key, value,
    attn_mask=None,
    dropout_p=0.0,
    is_causal=False,
    scale=None,
    enable_gqa=False,
)

The documentation defines query length L, source and key length S, optional Boolean or additive masks, default scaling 1/√d, causal bias behavior, dropout, and experimental grouped-query support. Depending on inputs and backend, PyTorch can dispatch to optimized implementations rather than literally executing the pedagogical Python of Chapter 7.

Here is the call on our three tokens, done two ways, so you can see the mask polarity for yourself.

python
import torch
import torch.nn.functional as F

q = k = torch.tensor([[1., 0.], [0., 1.], [1., 1.]])
v = torch.tensor([[1., 0.], [0., 2.], [3., 1.]])

# SDPA wants (batch, heads, seq, head_dim): add two leading axes.
qh, kh, vh = (t[None, None] for t in (q, k, v))

out_causal = F.scaled_dot_product_attention(qh, kh, vh, is_causal=True)

# The same thing with an explicit Boolean mask.
# In PyTorch SDPA, True in attn_mask means the position PARTICIPATES.
allowed = torch.tril(torch.ones(3, 3, dtype=torch.bool))
out_mask = F.scaled_dot_product_attention(qh, kh, vh, attn_mask=allowed)

torch.set_printoptions(precision=4, sci_mode=False)
print("is_causal=True  ->")
print(out_causal[0, 0])
print("attn_mask=tril  ->")
print(out_mask[0, 0])
print("max abs difference:", (out_causal - out_mask).abs().max().item())

is_causal=True  ->
tensor([[1.0000, 0.0000],
        [0.3302, 1.3395],
        [1.7587, 1.0000]])
attn_mask=tril  ->
tensor([[1.0000, 0.0000],
        [0.3302, 1.3395],
        [1.7587, 1.0000]])
max abs difference: 0.0

Same numbers as the hand calculation, one more time, now through the production API. Four implementation lessons follow from that signature.

Implementation APIs evolve. The mathematics in this lesson is evergreen; production code should always be checked against the documentation for the installed framework version.

The Seven Common Implementation Mistakes

Eight Misconceptions, Corrected

A good technical explanation removes these shortcuts early, because later systems concepts build directly on the distinctions.

Are Attention Weights Explanations?

Attention weights are visually compelling. They form distributions over input positions, so it is tempting to read a high weight as "the model used this token because it mattered." That conclusion is not generally justified without additional analysis.

Jain and Wallace tested attention weights as explanations across several NLP models and found cases where learned attention was weakly related to gradient-based importance measures, and where very different attention distributions could yield similar predictions. They argued that standard attention weights should not automatically be treated as faithful explanations.

Wiegreffe and Pinter challenged a categorical rejection, arguing that the answer depends on the definition of explanation and on appropriate baselines and tests. They proposed diagnostic conditions under which attention can still be informative.

The Complete Mental Model: Six Verbs

Attention is easiest to remember as a sequence of six ideas. This is the checklist. Every attention implementation you will ever read does these six things in this order, and every bug you will ever chase is one of them done wrong.

Everything Else Modifies One Surrounding Dimension

If these distinctions are clear, you do not merely know the attention equation. You have the conceptual foundation required to understand why modern LLM architectures and inference systems keep changing the structures around it.

Where to Go Next

References

Bahdanau, Cho, Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. ICLR 2015. arXiv:1409.0473
Luong, Pham, Manning. Effective Approaches to Attention-based Neural Machine Translation. EMNLP 2015. arXiv:1508.04025
Vaswani et al. Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762
Shazeer. Fast Transformer Decoding: One Write-Head is All You Need. 2019. arXiv:1911.02150
Beltagy, Peters, Cohan. Longformer: The Long-Document Transformer. 2020. arXiv:2004.05150
Zaheer et al. Big Bird: Transformers for Longer Sequences. NeurIPS 2020. arXiv:2007.14062
Su et al. RoFormer: Enhanced Transformer with Rotary Position Embedding. 2021. arXiv:2104.09864
Dao, Fu, Ermon, Rudra, Re. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135
Ainslie et al. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023. arXiv:2305.13245
Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. 2023. arXiv:2307.08691
DeepSeek-AI et al. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. 2024. arXiv:2405.04434
Zadouri et al. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling. 2026. arXiv:2603.05451
Jain, Wallace. Attention is not Explanation. NAACL-HLT 2019.
Wiegreffe, Pinter. Attention is not not Explanation. EMNLP-IJCNLP 2019.
PyTorch documentation, torch.nn.functional.scaled_dot_product_attention, consulted for API shapes, masks, causality, scaling, dropout and grouped-query behavior.

"What I cannot create, I do not understand."
Richard Feynman. You have now created attention four times on this page, in arithmetic, in plain Python, in NumPy and in PyTorch.