Attention Mechanism

Watch a token gather the context it needs and see how attention changes its representation.

The attention mechanism in transformers lets every position in a sequence build a new representation of itself by taking a weighted average of information from other positions. Each position produces a query vector, every position offers a key vector and a value vector, all obtained by multiplying the same hidden states by three learned matrices. A query is compared with every eligible key by a dot product, the scores are divided by √d_k and turned into weights with softmax, and those weights mix the value vectors. In one line, Attention(Q, K, V) = softmax(QKᵀ/√d_k + M)V, where M is a mask that forbids some pairs. That is all the operation does: arithmetic on matrices, with no lookup table and no rules about grammar. What the matrices contain is learned during training.

This lesson works through that formula with numbers small enough to verify by hand (jump to the complete numeric calculation), extends it to a whole sequence, and then covers the variants and costs you meet in real systems: causal and padding masks, cross-attention, multiple heads, quadratic score matrices and the key/value cache. The surrounding model, with its embeddings, feed-forward layers, residual connections and output head, is the subject of Transformer Basics; here we stay inside the attention operation.

Compare two sentences: "The bank approved the loan" and "The boat reached the bank." After tokenization and the embedding lookup, the token bank is mapped to the same vector in both, because an embedding table is a function of the token ID alone (see Embeddings Explained). Yet a financial institution and a river edge need different representations before a model can continue the text sensibly. Something has to move information from approved and loan, or from boat and reached, into the vector for bank.

Attention is the operation that does this routing inside a transformer. After one attention step the vector at the position of bank is no longer the raw embedding; it is a blend that includes information gathered from other positions, and the blend differs between the two sentences. Stack many layers and each position keeps being refined using the already-refined states of the others, so representations become progressively more context-dependent.

Be careful not to read too much into this. A single attention layer does not "resolve" ambiguity by itself: the mixed vector is passed through a feed-forward network, combined with residual connections, and refined by later layers, and the model may still get the sense wrong. Attention provides the mechanism for exchanging information between positions; whether a trained model uses it well for a given case is an empirical question.

Strip away the names and attention is a three-step recipe applied to one position at a time. First, compute a relevance score between this position and each position it is allowed to look at. Second, normalise the scores into weights that are non-negative and sum to 1. Third, produce the output as the weighted sum of one vector per eligible position. If the weights were all equal the output would be a plain average of the inputs; learned, input-dependent weights make it a selective average.

Two things to hold on to. The weights are computed from the data, so the same layer routes differently for different inputs, which is what separates attention from a fixed convolution window or a fixed mixing matrix. And the weights are an intermediate quantity of the computation, not a report of why the model finally answered as it did; we return to that limit in a dedicated section below.

Suppose the layer receives hidden states X ∈ R^(n×d_model): one row per token position, n positions, d_model numbers per position. The layer owns three learned matrices and applies them to every row:

Q = X · W_Q     W_Q ∈ R^(d_model × d_k)    →  Q ∈ R^(n × d_k)
K = X · W_K     W_K ∈ R^(d_model × d_k)    →  K ∈ R^(n × d_k)
V = X · W_V     W_V ∈ R^(d_model × d_v)    →  V ∈ R^(n × d_v)

Toy sizes: n = 3, d_model = 8, d_k = d_v = 4
X: [3, 8]   W_Q: [8, 4]   Q: [3, 4]   (row i of Q = x_i · W_Q)

Row i of Q is what position i "asks"; row j of K is how position j "advertises" itself for matching; row j of V is the content position j contributes if it is attended to. Quote marks matter: these are the usual mnemonics, but mechanically they are three different linear views of the same vector. Training shapes them so that useful comparisons become large dot products, and keys and values are separate projections so that "what makes me a good match" and "what I hand over" can differ.

The figure also shows why the later formula can be written with whole matrices: nothing here depends on one token at a time, so all n queries are scored against all n keys in a single matrix product, which maps very well onto GPUs. In the common case d_k = d_v, and in multi-head attention they equal d_model / h, as shown below.

The attention formula computes, for each query, a softmax-weighted average of the value rows, with the weights coming from scaled query-key dot products.

Attention(Q, K, V) = softmax( Q Kᵀ / √d_k  +  M ) V

Q  ∈ R^(n_q × d_k)    one query row per querying position
K  ∈ R^(n_k × d_k)    one key row per position that can be attended to
V  ∈ R^(n_k × d_v)    one value row per such position
Q Kᵀ ∈ R^(n_q × n_k)  entry (i, j) = q_i · k_j   (raw score, a dot product)
√d_k                  scalar scale factor
M  ∈ R^(n_q × n_k)    0 where attention is allowed, -∞ where it is forbidden
softmax               applied along the LAST axis: each row sums to 1
output ∈ R^(n_q × d_v)

Three details deserve slow reading. First, the dot product q·k is large and positive when the two vectors point the same way, near zero when orthogonal, negative when opposed, so it serves as a learned similarity. It is unbounded and not a probability. Second, softmax runs along each row, i.e. over the keys for one fixed query: w_ij = exp(s_ij) / Σ_j' exp(s_ij'). Each row is a distribution over positions; the columns are not normalised. Third, adding -∞ to a score makes exp(-∞) = 0, so a forbidden position receives exactly zero weight, and the remaining weights renormalise among themselves.

Why divide by √d_k? If the entries of a query and a key are independent with zero mean and unit variance, their dot product has variance d_k: it is a sum of d_k terms, each of unit variance. Its typical magnitude therefore grows like √d_k. Large scores push softmax towards a one-hot distribution, where gradients through the softmax become tiny and learning stalls. Dividing by √d_k returns the scores to roughly unit variance regardless of head size. This is the argument given in the original transformer paper (Vaswani et al., 2017, arXiv:1706.03762); it relies on the independence assumption, which holds only approximately for trained vectors, so treat it as the design rationale rather than a guarantee.

The figure below lets you grow d_k and watch softmax saturate when the scale is skipped. The scores are an illustrative pattern, modelled so that unscaled values grow like √d_k.

Take the smallest case that still shows the whole mechanism: one query attending to two positions, with d_k = 2. All numbers are illustrative toy values chosen so the arithmetic is visible, not taken from a trained model.

q  = [1, 0]
k1 = [1, 0]    v1 = [2, 0]
k2 = [0, 1]    v2 = [0, 4]
d_k = 2,  √d_k = 1.41421

1. 1

   1. Raw scores: dot products

   s1 = q·k1 = 1·1 + 0·0 = 1 and s2 = q·k2 = 1·0 + 0·1 = 0. Scores: [1, 0].
2. 2

   2. Scale by √d_k

   [1, 0] / 1.41421 = [0.7071, 0].
3. 3

   3. Exponentiate

   exp(0.7071) = 2.0281 and exp(0) = 1. Their sum is 3.0281.
4. 4

   4. Normalise (softmax weights)

   w1 = 2.0281 / 3.0281 = 0.6698, w2 = 1 / 3.0281 = 0.3302. They sum to 1, as a distribution must.
5. 5

   5. Weighted sum of values

   out = 0.6698·[2, 0] + 0.3302·[0, 4] = [1.3396, 0] + [0, 1.3208] = [1.3396, 1.3208] using the four-digit weights. Carrying full precision gives [1.3395, 1.3210]; the last-digit difference is only rounding of the intermediate weights.

Notice what the output is not: it is neither v1 nor v2 but a point between them, pulled towards v1 because the query aligned better with k1. Now change the query. With q = [0, 1] the raw scores swap to [0, 1], the weights become [0.3302, 0.6698], and the output becomes [0.6605, 2.6790], now dominated by v2. With q = [1, 1] both scores equal 1/√2, so the weights are [0.5, 0.5] and the output is the plain average [1, 2]. Same keys, same values, different query, different mixture: that dependence on the query is the whole point.

Try these queries yourself: the figure below keeps the keys and values fixed and recomputes the scores, softmax weights and output as you move q.

Also notice how gentle the weighting is. The raw scores differ by 1 and the weights are still 67% / 33%. Softmax amplifies score differences exponentially: a gap of 1/√2 multiplies the odds by e^0.7071 ≈ 2.03. Had we skipped the scaling and used the raw gap of 1, the odds would be e ≈ 2.72, giving weights 0.7311 / 0.2689. In real heads with d_k of 64 or 128, the scaling divides by 8 or about 11.3, which is exactly what keeps sums of dozens of products from making softmax needlessly sharp.

Real layers do not loop over queries. Stacking every query as a row of Q turns the calculation into matrix products. Let the sequence have n = 3 positions and (toy values again) Q = K = [[1,0],[0,1],[1,1]] and V = [[2,0],[0,4],[1,1]], with d_k = 2.

Q Kᵀ  (3×3, entry (i,j) = q_i · k_j)        divided by √2 = 1.41421
        k1  k2  k3                                    k1      k2      k3
  q1  [  1   0   1 ]                          q1  [0.7071  0      0.7071]
  q2  [  0   1   1 ]                          q2  [0       0.7071 0.7071]
  q3  [  1   1   2 ]                          q3  [0.7071  0.7071 1.4142]

softmax per ROW (each row sums to 1):
  w1 = [0.4011, 0.1978, 0.4011]     exp: 2.0281, 1, 2.0281   sum 5.0562
  w2 = [0.1978, 0.4011, 0.4011]
  w3 = [0.2483, 0.2483, 0.5035]     exp: 2.0281, 2.0281, 4.1133   sum 8.1695

output = W · V  (3×2), one row per querying position:
  o1 = 0.4011·[2,0] + 0.1978·[0,4] + 0.4011·[1,1] ≈ [1.2033, 1.1922]
  o2 ≈ [0.7967, 2.0056]
  o3 ≈ [1.0000, 1.4965]

Read the matrices row by row. Row i of the score matrix holds the scores of the i-th querying position against every key; row i of the weight matrix is that position's distribution over sources; row i of the output is its new vector, ready to replace its input. The computation therefore produces n outputs in one go, and each of the n² entries is a pair of positions. That square is both the source of attention's expressive power (any position can read from any other) and of its cost.

The formula never says where Q, K and V come from; that choice defines the variant. In self-attention all three are projections of the same sequence, so the sequence exchanges information with itself. In cross-attention the queries are projected from one sequence while keys and values are projected from another, so K and V have n_k rows that may differ from the n_q rows of Q. The score matrix is then n_q × n_k, not square, and the output still has one row per query.

The classic setting is an encoder–decoder translation model. The encoder reads the source sentence once and produces a matrix of contextual vectors. At every decoder layer, the target tokens generated so far first run self-attention among themselves, then cross-attention in which the decoder's states form the queries and the encoder's fixed outputs form the keys and values. The decoder is effectively asking "which source positions should inform the next target word" and reading from the encoder's output. Decoder-only language models use only masked self-attention; the same cross-attention idea reappears wherever one modality or stream conditions another, for example when a generator attends to text conditioning or to image features.

A language model is trained to predict each token from the ones before it. For efficiency, training does not run one position at a time: the whole sequence goes through the network in a single pass and every position produces a prediction at once, with the target for position i being the token at i+1. If position i could attend to position i+1, it would simply copy the answer from the input, the loss would collapse without the model learning anything useful, and at generation time, when future tokens do not exist yet, the behaviour would not transfer. The fix is a causal mask: position i may attend only to positions j ≤ i.

        j=1   j=2   j=3   j=4
 i=1  [   0    -∞    -∞    -∞ ]
 i=2  [   0     0    -∞    -∞ ]
 i=3  [   0     0     0    -∞ ]
 i=4  [   0     0     0     0  ]

allowed(i, j) = (j <= i)       M[i][j] = 0 if allowed else -∞

The mask is added to the scaled scores before softmax, so the upper triangle gets weight exactly zero. In the three-position example above, row 1 collapses to [1, 0, 0] and its output equals v1 = [2, 0]: with only itself visible, a position can only copy its own value. Row 2 becomes [0.3302, 0.6698, 0] (the same arithmetic as the two-position walkthrough) and row 3 is unchanged. Zero weight is the whole guarantee, and it holds in the later layers too: because no layer can move information backwards in position, the output at position i depends only on inputs at positions up to i. That also means you can test for leakage directly: change a later token and confirm every earlier output stays bit-identical.

Bidirectional (encoder-style) attention uses no causal mask, so each position can see every unmasked position on both sides. This suits tasks where the full input is available, such as classification or producing representations for retrieval, and is why encoder-only models can use right-hand context that a decoder-only model cannot. Which of these a model uses, and how the combination works in practice, is covered in Transformer Basics.

Batching puts several sequences into one tensor, which forces them to share a length. A batch holding a 3-token and a 2-token sequence is padded to length 3 with a filler token. The filler carries no information, and letting real tokens attend to it would make the result depend on how much padding the batch happened to contain. A padding mask blocks columns: for the short sequence, key position 3 is forbidden for every query. This is a different thing from the causal mask. Causal masking is a property of the task (no looking at the future) and is the same triangle for every sequence; padding masking is a property of the batch (which positions exist) and differs per sequence.

causal                padding (key valid)         allowed = causal AND valid
[1 0 0]               [1 1 0]  (same for           [1 0 0]
[1 1 0]       AND     [1 1 0]   every row)    =    [1 1 0]
[1 1 1]               [1 1 0]                      [1 1 0]   <- row 3 is a padded query:
                                                            its output is discarded

Two practical hazards recur. A query row whose entire set of keys is masked has all scores equal to -∞, and softmax over that yields NaN (0/0); implementations either avoid fully masked rows, substitute a very large negative number instead of true infinity, or zero such rows afterwards. And mask conventions differ between libraries: some APIs take a boolean mask where True means "may attend", others where True means "blocked", and some take an additive float mask exactly like M above. PyTorch's two attention entry points illustrate this at the time of writing; read the documentation of the specific function before passing a mask, and test it with the leakage check described earlier. Other structural masks follow the same idea of removing particular pairs: sliding-window masks restrict each position to nearby keys, and packed-sequence masks stop documents concatenated in one row from attending across their boundaries.

A single attention operation produces one weighting per position, so it must use the same blend for everything it wants to read. Multi-head attention runs h attention operations in parallel, each with its own learned W_Q, W_K, W_V, and recombines them. Usually the model dimension is split evenly: d_head = d_model / h, so the total work is comparable to one full-width head.

X                         [n, 8]
head 1:  Q1,K1,V1 = X·W_Q1, X·W_K1, X·W_V1      each W: [8, 4]   → each Q1,K1,V1: [n, 4]
         scores1 = Q1·K1ᵀ / √4   → [n, n]    (scale uses d_head = 4, so divide by 2)
         out1 = softmax(scores1 + M) · V1       → [n, 4]
head 2:  same with its own W_Q2, W_K2, W_V2      → out2: [n, 4]
concat(out1, out2)        [n, 8]
output = concat · W_O     W_O: [8, 8]   → [n, 8]   (same shape as X, ready for the residual add)

Parameters (no biases): W_Q, W_K, W_V, W_O = 4 × (8×8) = 256

The benefit is that each head computes its own score matrix, so one head can weight positions one way while another weights them differently, and W_O mixes the results. Interpretability studies do find heads that look specialised, for instance attending to the previous token or to a matching earlier occurrence. But there is no guarantee that any given head has a single human-readable job, such as "the grammar head". Roles are learned, can be distributed over several heads, change between training runs, and many heads can be removed with little measurable loss. Treat a nameable head as an interesting observation about one trained model, not as part of the specification of the mechanism.

Because the weight matrix is right there, it is tempting to plot it and read it as an explanation: the model answered "refund" because it attended heavily to "charged twice". The weights show how a layer combined the value vectors it was given, and that is a real fact about the computation. They do not by themselves show which input caused the output, for several reasons.

The value vectors being averaged have different magnitudes and directions, so a small weight on a large or distinctive value can matter more than a large weight on a nearly neutral one. Weights from one head are mixed with those of other heads through W_O, passed through feed-forward layers and added to the residual stream, where each position's own earlier content is also present. And at layer 2, the vector at position j already contains information gathered from other positions, so "attends to j" no longer means "uses token j". Experiments that perturb or remove inputs and measure the change in the output (ablation, gradient-based attribution, activation patching) answer the causal question more directly than raw weights do. Use attention maps for debugging intuition and hypothesis generation, and verify any claim about importance with an intervention.

Full attention scores every query against every key, so each head in each layer computes n² scores. At n = 1,000 that is 1,000,000 pairs; at n = 2,000 it is 4,000,000. Doubling the length quadruples the pair count, and multiplying by 10 multiplies it by 100. The arithmetic of the score matrix (n²·d_k multiply-adds) and of applying the weights to the values (n²·d_v) scales the same way.

Algorithmic complexity and memory use are separate questions. A straightforward implementation materialises the n × n score matrix: in 16-bit floats (2 bytes per entry) that is 2 MB per head per layer at n = 1,000 and 8 MB at n = 2,000, so a layer with 32 heads holds 64 MB versus 256 MB (decimal megabytes, one sequence). Optimised kernels such as FlashAttention (Dao et al., 2022, arXiv:2205.14135) compute exact attention in tiles that stay in fast on-chip memory and never write the full score matrix to main GPU memory, so their peak memory grows linearly in n rather than quadratically. The number of pair computations is unchanged, though, so time still grows quadratically; the saving is in memory traffic and peak memory. Approximate schemes (sliding windows, sparse patterns, low-rank or linear approximations) reduce the pair count itself and trade some modelling quality or generality for it. Real serving cost also depends on the feed-forward layers, batching and hardware, so quadratic pair growth explains why long contexts hurt without being the only term; see Context Window for the budgeting view.

Autoregressive generation produces one token at a time. Naively, to generate token t+1 you would rerun the model over all t tokens. But with a causal mask, the keys and values of earlier positions do not change when a new token is appended: k_j and v_j for j ≤ t depend only on inputs up to j. So the server stores each layer's K and V rows in a KV cache. For a new token it computes only that token's query, key and value, appends the new key and value to the cache, and attends from the single new query to all cached keys.

q_new, k_new, v_new = x_new · W_Q, x_new · W_K, x_new · W_V     # shapes [1, d_head]
K_cache = append(K_cache, k_new)                                 # [t+1, d_head]
V_cache = append(V_cache, v_new)                                 # [t+1, d_head]
scores  = q_new · K_cacheᵀ / √d_head                             # [1, t+1]  (no mask needed: the cache holds only the past)
out     = softmax(scores) · V_cache                              # [1, d_head]

The query is not cached, because it is used once, for the new position. What the new query must still do is read the whole history: each step costs a score row of length t+1, so per-step attention work grows linearly with the context, and generating T tokens costs about T²/2 pairs in total rather than recomputing T full forward passes. Decode steps tend to be limited by memory bandwidth, since the whole cache must be read for each token while only a little arithmetic is done with it.

The price is memory. Per token the cache stores a K and a V vector in every layer for every KV head: bytes/token = 2 · layers · kv_heads · d_head · bytes_per_value. With an illustrative configuration of 32 layers, 8 KV heads, d_head = 128 and 2-byte values this is 2·32·8·128·2 = 131,072 bytes = 128 KiB per token. For an 8,000-token sequence that is 1,048,576,000 bytes (about 1.05 GB, 0.98 GiB). The same model with 32 KV heads (one per query head, classic multi-head attention) would need 4× as much: 512 KiB per token, about 4.19 GB at 8,000 tokens. Grouped-query attention, which shares each K/V head among several query heads, and multi-query attention (a single shared K/V head) exist precisely to shrink this cache; the 8-versus-32 comparison is that effect. These configurations are illustrative, not those of any particular model, and the cache is per sequence: concurrent users multiply it.

Cache management (allocation, eviction, sharing prefixes across requests) belongs to the serving layer and is covered in AI Inference Infrastructure.

Most attention bugs fall into a few categories, and each has a direct test.

Mask errors are the most damaging because they are silent. An inverted boolean convention, a causal mask applied to an encoder, or a missing padding mask changes results without raising an error. Test by perturbation: alter the last token and assert that all earlier outputs are unchanged under a causal mask; pad a sequence with extra filler and assert the outputs for real tokens do not move. Check that every weight row sums to 1 and that masked entries are exactly 0.

Numerical instability arises when scores are large. The textbook softmax computes exp(s) directly, which overflows for scores around 88 in 32-bit floats and far sooner in 16-bit formats. The stable version subtracts the row maximum first, which leaves the result unchanged but keeps every exponent at or below zero: for scores [1000, 999] it returns [0.7311, 0.2689]. Low-precision training additionally calls for accumulating softmax in higher precision. Dropping the 1/√d_k factor is another route to saturated softmax and unstable training.

Position handling. Attention scores depend only on content of the vectors; a pure attention layer cannot tell "dog bites man" from "man bites dog" unless position information is injected, through added positional embeddings or rotations of queries and keys. A model used with sequences longer than it was trained on can therefore degrade even though attention is formally defined for any length, because the positional scheme was never exercised there.

Long-range and distraction effects. Softmax is a competition: adding many irrelevant tokens spreads weight over them, diluting the contribution of the one value that mattered. Models also often retrieve less reliably from some regions of a long context than others. These are empirical behaviours of trained models, not consequences of the formula, so measure them: place a known fact at several depths of a long prompt and record accuracy by depth and length. Measure latency and memory against sequence length, check whether you are on a fused attention path, and when output quality is the problem, evaluate interventions (shorter context, retrieval of only relevant passages, reordering) on a test set rather than trusting visual impressions of attention maps.

The following NumPy function implements the whole operation: scaled scores, boolean masking, stable softmax and the weighted sum, with shape assertions. It is runnable (NumPy installed) and reproduces the numbers above. It is written for clarity; production kernels fuse these steps, work in tiles, and handle batches and heads.

import numpy as np

def stable_softmax(scores, axis=-1):
    shifted = scores - scores.max(axis=axis, keepdims=True)   # largest entry becomes 0
    e = np.exp(shifted)
    return e / e.sum(axis=axis, keepdims=True)

def attention(Q, K, V, mask=None):
    """Q: [n_q, d_k]  K: [n_k, d_k]  V: [n_k, d_v]  mask: [n_q, n_k] bool, True = allowed."""
    assert Q.shape[1] == K.shape[1] and K.shape[0] == V.shape[0]
    scores = Q @ K.T / np.sqrt(Q.shape[1])                    # [n_q, n_k]
    if mask is not None:
        assert mask.shape == scores.shape
        scores = np.where(mask, scores, -np.inf)
    weights = stable_softmax(scores)                          # each row sums to 1
    return weights @ V, weights                               # [n_q, d_v], [n_q, n_k]

Q = K = np.array([[1., 0.], [0., 1.], [1., 1.]])
V = np.array([[2., 0.], [0., 4.], [1., 1.]])
causal = np.tril(np.ones((3, 3), dtype=bool))

out, w = attention(Q, K, V, causal)
print(np.round(w, 4))     # [[1. 0. 0.] [0.3302 0.6698 0.] [0.2483 0.2483 0.5035]]
print(np.round(out, 4))   # [[2. 0.] [0.6605 2.679] [1. 1.4965]]

out1, w1 = attention(np.array([[1., 0.]]), K[:2], V[:2])
print(np.round(w1, 4), np.round(out1, 4))   # [[0.6698 0.3302]] [[1.3395 1.321]]

A multi-head version reshapes the projected [n, d_model] matrices into [h, n, d_head], runs the same function per head (a batched matrix product in practice), concatenates along the feature axis and multiplies by W_O. Nothing else changes, which is why the single-head version is the right thing to understand first.

Attention is the one step in a transformer layer that moves information between positions. Everything else acts on each position independently: the feed-forward network transforms each position's vector, normalisation rescales it, and residual connections add each sublayer's output back to its input so the original information keeps flowing. A decoder-only layer applies, roughly, normalised masked attention plus a residual, then a feed-forward sublayer plus a residual, and the stack is repeated; after the last layer a projection turns each position's vector into logits over the vocabulary. The exact placement of normalisation and the other details vary by model family. Attention is therefore necessary to the transformer's behaviour but is not the whole transformer: much of the parameter count typically sits in the feed-forward layers, and a stack of attention layers without them, or without residual paths, would not train or behave the same way.

With that boundary clear, the full picture is in Transformer Basics, and the way a stack like this produces text one token at a time is the subject of Next-Token Prediction. Those pieces, together with the retrieval, memory and tools that decide what reaches the context window, make up an LLM application.

1. 1

   Scale scores by √d_k before softmax. Without it, dot products grow with d_k and softmax saturates toward one-hot weights, which starves gradients.

2. 2

   Masks are applied to the scaled scores before softmax. Use a causal mask for decoders, a padding mask for batches, and guard rows where every key is masked, because softmax over all -∞ yields NaN.

3. 3

   Mask bugs are silent: an inverted boolean, a missing padding mask or a causal mask on an encoder changes outputs without raising an error, so test them directly.

4. 4

   Attention weights are not an explanation. Value vectors differ in magnitude, so a small weight on a distinctive value can matter more than a large one on a bland value.

5. 5

   Full attention computes n² scores per head per layer; the cost shows up as compute and, if the n × n matrix is materialised, as memory.

6. 6

   KV caching removes recomputation of earlier keys and values during generation, but the cache costs 2 · layers · kv_heads · d_head · bytes per token and every step still reads the whole history.