Self-attention and positional structure

• Token — a chunk of text (a word piece) with an integer id; the unit the model reads and writes.
• Embedding — a learned vector (list of numbers) for each token id, the model’s starting representation of that token.
• Hidden state (X, residual stream) — the vector the model currently holds for each position; it is rewritten layer by layer.
• n — sequence length: how many tokens are in the context.
• d_model — the width of the hidden state, e.g. 4096; every position is a vector of this many numbers.
• d_k — the width of a query/key vector inside one head, e.g. 128.
• h — the number of attention heads.
• Matrix, shape [a, b] — a table with a rows and b columns. A “projection” is multiplying by a learned matrix to get a new vector.
• Dot product (a · b) — multiply two equal-length vectors element by element and add the results: [1,2]·[3,4] = 1·3 + 2·4 = 11. It is large when the vectors point the same way.
• Softmax — a function that turns a list of arbitrary numbers into positive weights that sum to 1.
• Logit — a raw, un-normalised score before softmax.
• Weight (parameter) — a number learned during training and then frozen. W_Q, W_K, W_V, W_O are weight matrices.
• Causal mask — a rule that stops a position from looking at later positions.
• HBM / SRAM — a GPU’s large-but-slow main memory / its small-but-very-fast on-chip memory.
• FLOPs — floating-point operations; a measure of arithmetic work.
• Autoregressive — generating one token at a time, each conditioned on all the previous ones.
• MHA / MQA / GQA — multi-head / multi-query / grouped-query attention: h query heads that have h, 1, or g key/value heads respectively.
• KV cache — stored keys and values of all past tokens, reused at every decode step.
• MLA (multi-head latent attention) — an attention variant that caches a compressed latent vector instead of full K and V.
• Prefill / decode — processing the whole prompt in one parallel pass / generating one token at a time.
• KiB / GiB — 1,024 bytes / 1,024^3 bytes.
• Permutation-equivariant — shuffling the inputs just shuffles the outputs the same way (the operation has no notion of order).

Attention does not read text. By the time it runs, the text has been through three steps:

1. Tokenize — "the cat sat" becomes integer ids, e.g. [464, 3797, 3332].
2. Embed — each id indexes a row in a learned table E of shape [vocab_size, d_model], giving one vector per token.
3. Stack — the vectors form a matrix X of shape [n, d_model] (one row per position).

X  (n = 3 positions, d_model = 4)

            dim0  dim1  dim2  dim3
 "the"  ->   1     0     1     0
 "cat"  ->   0     2     0     1
 "sat"  ->   1     1     0     1

(Throughout this page we use these tiny made-up numbers so every step can be checked by hand. Real models use d_model of 2,048-16,384 and the numbers are learned, not chosen.)

A transformer has many layers. The first layer’s X is embeddings (plus position information, see below). Every later layer’s X is the previous layer’s output. So when this page says “the hidden state”, it means “the current vector the model holds for this position”. Attention’s job is to let each position look at the other positions and pull in useful information, so that, for example, the vector for “it” can absorb what “it” refers to.

Each row of X is turned into three different vectors by multiplying with three weight matrices:

Q = X · W_Q        shapes:  [n, d_model] · [d_model, d_k]  ->  [n, d_k]
K = X · W_K                 [n, d_model] · [d_model, d_k]  ->  [n, d_k]
V = X · W_V                 [n, d_model] · [d_model, d_v]  ->  [n, d_v]     (d_v = d_k in practice)

• W_Q, W_K, W_V are learned parameters: they start random and are adjusted by gradient descent during training. They are the same for every position (the same “lens” is applied to all tokens) and fixed at inference.
• Q (query): “what am I looking for?” — computed from the position that is asking.
• K (key): “what do I contain / advertise?” — computed from every position that can be looked at.
• V (value): “what information do I hand over if someone attends to me?” — also from every position.

Why three different matrices instead of using X directly?

A position needs to play three different roles, and the information useful for each role differs:

• Matching (Q vs K) needs a representation tuned for comparison, e.g. “I am a verb looking for my subject” vs “I am a noun in subject position”.
• Content to pass on (V) needs a representation tuned for what the reader should receive, which may be different from what is useful for matching.

If Q, K and V were all just X, a token would match best with itself (a vector’s dot product with itself is its largest), and matching and content could not specialise. Three separate projections let the model learn “what to search for”, “how to be found”, and “what to deliver” independently.

Concrete numbers

With d_model = 4 and d_k = 2, our toy weights are:

W_Q = [[1,0],[0,1],[1,0],[0,1]]     W_K = [[0,1],[1,0],[0,1],[1,0]]     W_V = [[1,0],[0,2],[0,1],[1,0]]

Multiplying row by row (row i of the result = row i of X times the matrix):

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

Check one entry by hand: Q[0] = X[0] · W_Q = [1,0,1,0] · W_Q. Column 0 of W_Q is [1,0,1,0], so the first output is 1·1 + 0·0 + 1·1 + 0·0 = 2. Column 1 is [0,1,0,1], giving 0+0+0+0 = 0. So Q[0] = [2, 0]. ✓.

The whole operation is one line:

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

Here is every piece, in order of execution.

Step 1 — Scores: how well does each query match each key? (Q Kᵀ)

Kᵀ is K with rows and columns swapped (shape [d_k, n]). The product Q Kᵀ has shape [n, n]; entry (i, j) is the dot product of query i with key j: how strongly position i wants to read from position j.

S = Q Kᵀ  =  [[ 0, 6, 4],        row i = position i as the asker
              [ 6, 0, 3],        column j = position j as the one being looked at
              [ 4, 3, 4]]

Check S[0][1] = Q[0] · K[1] = [2,0]·[3,0] = 6. The dot product is the right comparison because it is large when two vectors point in a similar direction and near zero or negative when they do not, and it is cheap (one matrix multiplication scores all pairs at once).

Step 2 — Scale (/ sqrt(d_k))

Divide every score by sqrt(d_k) = sqrt(2) ≈ 1.414:

[[0.00, 4.24, 2.83],
 [4.24, 0.00, 2.12],
 [2.83, 2.12, 2.83]]

Why this is needed is derived in the next concept. For now: it keeps scores in a sensible range as d_k grows.

Step 3 — Mask (+ M): forbid looking at the future

See the next-but-one concept. For now: entries above the diagonal get -∞.

[[0.00,  -inf,  -inf],
 [4.24,  0.00,  -inf],
 [2.83,  2.12,  2.83]]

Step 4 — Softmax each row: turn scores into weights

Softmax takes a list of scores z_1 … z_n and returns weights:

softmax(z)_j  =  exp(z_j) / ( exp(z_1) + exp(z_2) + … + exp(z_n) )

• exp(x) (e to the power x) makes every number positive and exaggerates differences (bigger scores get disproportionately more weight).
• Dividing by the sum makes the weights add up to 1, so they behave like “what fraction of my attention goes to each position”.
• exp(-∞) = 0, which is why -∞ in the mask means “weight exactly zero”.

Row 2 of our example: scores [2.83, 2.12, 2.83] -> exp = [16.9, 8.3, 16.9] -> sum = 42.2 -> weights [0.40, 0.20, 0.40].

A = softmax(rows)  =  [[1.000, 0.000, 0.000],
                       [0.986, 0.014, 0.000],
                       [0.401, 0.198, 0.401]]

Row i of A says how position i distributes its attention over positions 0..i.

Step 5 — Weighted sum of values (A · V)

Each output row is the attention-weighted average of the V rows:

out[2] = 0.401·V[0] + 0.198·V[1] + 0.401·V[2]
       = 0.401·[1,1] + 0.198·[1,4] + 0.401·[2,2]
       = [1.401, 1.994]

out = [[1.000, 1.000],
       [1.000, 1.042],
       [1.401, 1.994]]

Position 2 (“sat”) has gathered a blend of information from “the”, “cat”, and itself. In a real model that output goes through further layers, so the next layer’s representation of “sat” now contains information about its context.

Verifying with code

import numpy as np

X  = np.array([[1,0,1,0],[0,2,0,1],[1,1,0,1]], float)
WQ = np.array([[1,0],[0,1],[1,0],[0,1]], float)
WK = np.array([[0,1],[1,0],[0,1],[1,0]], float)
WV = np.array([[1,0],[0,2],[0,1],[1,0]], float)

Q, K, V = X @ WQ, X @ WK, X @ WV                       # step 0: projections
d_k = Q.shape[-1]
scores = Q @ K.T / np.sqrt(d_k)                        # steps 1-2
mask = np.triu(np.ones((3, 3), bool), k=1)             # True above the diagonal = future positions
scores = np.where(mask, -np.inf, scores)               # step 3
e = np.exp(scores - scores.max(-1, keepdims=True))     # step 4 (subtract the row max for numerical safety)
A = e / e.sum(-1, keepdims=True)
out = A @ V                                            # step 5
print(A.round(3)); print(out.round(3))

Subtracting the row maximum before exp does not change the softmax result (it cancels in the ratio) but prevents overflow for large scores. Every real implementation does this.

Intuition in one sentence

Q and K decide where to look (the weights A); V decides what to take from there; the output is the weighted average of what was taken. The weights come from the content itself, which is why attention is called content-based routing.

The claim

If q and k have d_k components that are independent with mean 0 and variance 1 (roughly true at initialisation), then q · k has variance d_k.

Derivation

q · k = q_1 k_1 + q_2 k_2 + … + q_{d_k} k_{d_k}.

• Each term q_i k_i has mean 0 and variance Var(q_i)·Var(k_i) = 1·1 = 1.
• The variance of a sum of independent terms is the sum of their variances: Var(q · k) = d_k.
• So the standard deviation of a raw score is sqrt(d_k): 8 for d_k = 64, about 11 for d_k = 128.

Dividing by sqrt(d_k) divides the standard deviation by sqrt(d_k), bringing it back to 1 whatever the head size.

import numpy as np
rng = np.random.default_rng(0)
for d in (4, 64, 512):
    q, k = rng.standard_normal((20000, d)), rng.standard_normal((20000, d))
    s = (q * k).sum(1)
    print(d, round(s.std(), 2), round((s / np.sqrt(d)).std(), 2))
# 4   2.0   1.0
# 64  8.03  1.0
# 512 22.54 1.0      <- raw std grows like sqrt(d); scaled std stays at 1

Why large scores are a problem — softmax saturation

Softmax of small-spread scores is smooth; softmax of large-spread scores is nearly one-hot:

softmax([2, 1, 0])     # [0.665, 0.245, 0.090]   smooth: several positions get weight
softmax([16, 8, 0])    # [1.000, 0.000, 0.000]   same scores x8: all weight on one position

When softmax is saturated, its gradient is almost zero (changing a score barely changes a weight that is already 0 or 1), so almost no learning signal flows back to W_Q and W_K, and attention gets stuck on whatever it picked at initialisation. The scaling factor is therefore a variance-control trick, not a tunable hyperparameter. (Some recent models add QK-norm, normalising q and k before the dot product, to control the same problem more directly at large scale.)

Why it is needed

A language model is trained to predict token t+1 from tokens 1..t. If position t could attend to position t+1, it could simply copy the answer, learn nothing useful, and be unusable at generation time (when future tokens do not exist yet). So every position may only attend to itself and earlier positions.

The mask

M is an [n, n] matrix added to the scores: 0 where attention is allowed and -∞ where it is forbidden.

M = [[ 0, -inf, -inf, -inf],       row i can see columns 0..i only  (lower triangle incl. diagonal)
     [ 0,   0,  -inf, -inf],
     [ 0,   0,    0,  -inf],
     [ 0,   0,    0,    0 ]]

After softmax (exp(-inf) = 0) the forbidden entries become exact zeros, and each row still sums to 1 over the allowed positions. In our worked example row 0 became [1, 0, 0] (the first token can only attend to itself).

Consequences

• Training is parallel: one forward pass computes the loss for all n positions at once, because the mask guarantees each position’s output depends only on its past.
• Generation is sequential: at inference you produce one token at a time.
• KV caching works (see below), a direct consequence of the mask.
• Encoder-only models (BERT) use no causal mask: every token sees all others, which suits understanding tasks but cannot generate left to right.
• Padding masks are a second, separate kind of mask that hides padding tokens in a batch; real implementations combine both.

What it is

Instead of one attention with width d_model, run h smaller attentions (“heads”) in parallel, each with its own W_Q, W_K, W_V of width d_k = d_model / h, then concatenate the results and mix them with a final matrix W_O.

for each head i in 1..h:
    Q_i = X W_Q^(i),  K_i = X W_K^(i),  V_i = X W_V^(i)          each [n, d_k]
    head_i = softmax( Q_i K_iᵀ / sqrt(d_k) + M ) V_i               [n, d_k]

MultiHead(X) = concat(head_1 … head_h) · W_O                        [n, h·d_k] · [d_model, d_model] -> [n, d_model]

Why split into heads?

One softmax per position produces one weighted average, i.e. one pattern of “where to look”. But a token often needs several kinds of information at once: the subject for agreement, the previous word for local context, an earlier mention of the same entity. Separate heads each compute their own attention pattern, in their own lower-dimensional “subspace” of the representation, so they can specialise. They run in parallel, and W_O lets the model recombine what each head found.

The compute and parameter cost is about the same as a single wide head: h projections of width d_model/h have the same total size as one projection of width d_model. The gain is expressiveness, not cost.

Parameter count

Per attention layer: W_Q, W_K, W_V, W_O are each d_model × d_model (when heads are concatenated), so 4 · d_model² parameters. For d_model = 4096: 4 × 16.8M ≈ 67M per layer.

Shapes in practice: split one big matmul

Implementations do not loop over heads. They compute one big projection and then reshape:

Q = X · W_Q                [n, d_model]
Q.view(n, h, d_k)          split the last dimension into h chunks of d_k
Q.transpose -> [h, n, d_k] put heads in a batch dimension, so all heads run as one batched matmul
scores = Q · Kᵀ            [h, n, n]   <- one n×n attention matrix PER HEAD

Note the [h, n, n] score tensor: memory grows with heads and with n², which motivates the next concepts.

import math, torch, torch.nn as nn, torch.nn.functional as F

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.h, self.d_k = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)   # W_Q, W_K, W_V packed into one matrix
        self.out = nn.Linear(d_model, d_model, bias=False)       # W_O

    def forward(self, x):                                        # x: [B, n, d_model]  (B = batch size)
        B, n, _ = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)                   # each [B, n, d_model]
        # split heads: [B, n, d_model] -> [B, h, n, d_k]
        q, k, v = (t.view(B, n, self.h, self.d_k).transpose(1, 2) for t in (q, k, v))

        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k)   # [B, h, n, n]
        mask = torch.triu(torch.ones(n, n, dtype=torch.bool, device=x.device), diagonal=1)
        scores = scores.masked_fill(mask, float("-inf"))         # block the future
        weights = F.softmax(scores, dim=-1)                      # rows sum to 1
        ctx = weights @ v                                        # [B, h, n, d_k]

        ctx = ctx.transpose(1, 2).reshape(B, n, -1)              # merge heads: [B, n, d_model]
        return self.out(ctx)                                     # mix heads with W_O

x = torch.randn(2, 5, 64)
print(CausalSelfAttention(64, 8)(x).shape)                       # torch.Size([2, 5, 64])

Mapping to the maths: self.qkv is X W_Q, X W_K, X W_V in one matmul; the view/transpose pair splits heads; the three lines with scores, mask, softmax are exactly steps 1-4; weights @ v is step 5; self.out is W_O. In production you would replace the manual middle with F.scaled_dot_product_attention(q, k, v, is_causal=True), which dispatches to a FlashAttention kernel and gives the same result without storing the n × n matrix.

The cost

The score matrix Q Kᵀ is [n, n] per head: computing it costs n² · d_k multiply-adds, and weights @ V costs the same again. Doubling the context length quadruples the attention work and (naively) the memory for the score matrix.

n = 8,192:    n² =      67 million entries per head per layer
n = 131,072:  n² =   17,179 million entries per head per layer  (~34 GB in fp16 for ONE head's matrix)

(The rest of a transformer layer, the projections and feed-forward network, scales linearly with n, so at short contexts the feed-forward network dominates and at long contexts attention does.)

FlashAttention — same maths, far less memory traffic

On a GPU, the bottleneck is usually not arithmetic but moving data between HBM (big, slow) and SRAM (tiny, very fast). Naive attention writes the full [n, n] score matrix to HBM, reads it back for softmax, writes the weights, and reads them again for · V: four trips of n² data.

FlashAttention splits Q, K, V into tiles that fit in SRAM and computes the output tile by tile, never storing the full n × n matrix. The difficulty is softmax: it needs the sum over the whole row before it can normalise. The trick is the online softmax, which keeps a running maximum m and running sum l and corrects earlier partial results when a new tile arrives:

for each tile of keys/values (K_j, V_j):
    s      = Q K_jᵀ / sqrt(d_k)
    m_new  = max(m_old, rowmax(s))
    l      = l · exp(m_old − m_new) + rowsum(exp(s − m_new))       # rescale old sum, add new
    acc    = acc · exp(m_old − m_new) + exp(s − m_new) · V_j       # rescale old output, add new
    m_old  = m_new
output = acc / l

Result: identical output (not an approximation), memory O(n) instead of O(n²), and typically 2-4x faster because of far less HBM traffic. FLOPs are unchanged. FlashAttention-2/3 further optimise parallelism and use newer GPU features.

Other mitigations

Approach	Idea	Cost / limitation
Sliding-window attention (Longformer, Mistral)	Each token attends only to the last w tokens (plus some global tokens): O(n·w)	Cannot directly retrieve beyond the window in one layer (information can still travel through stacked layers)
Sparse attention	Attend to selected blocks chosen by a small indexer or fixed pattern	Needs training/adaptation; quality depends on selection
Linear attention / state-space models (Mamba, RWKV, DeltaNet)	Replace softmax attention with a recurrence keeping a fixed-size state: O(n)	Weaker exact recall; often used in hybrids with a few full-attention layers
KV-cache compression, eviction, paging	Reduce the memory of stored keys/values at serving time	Eviction can drop needed tokens
GQA / MQA	Share K/V across heads to shrink the cache (see below)	Slight quality trade

The problem

Look at the formula: nothing in softmax(Q Kᵀ) V depends on where a token sits. If you shuffle the rows of X, the output rows shuffle the same way; “dog bites man” and “man bites dog” would give the same bag-of-words result. Order must be injected explicitly.

Sinusoidal encodings (original Transformer)

Add a fixed vector PE(pos) to each embedding, where dimension pairs use sine and cosine at geometrically spaced frequencies:

PE(pos, 2i)   = sin( pos / 10000^(2i / d_model) )
PE(pos, 2i+1) = cos( pos / 10000^(2i / d_model) )

• pos is the position index (0, 1, 2, …); i indexes dimension pairs.
• Low i -> high frequency (changes quickly with position, like a fast clock hand); high i -> very low frequency (like an hour hand). Together they give each position a unique “fingerprint”, and the offset between two positions is a simple linear function of the encodings.
• No parameters to learn; some extrapolation, but superseded.

Learned absolute embeddings (GPT-2, BERT)

A trainable table with one vector per position index, added to token embeddings. Simple, but the model has never seen positions beyond the training length, so it cannot extrapolate.

RoPE (rotary position embedding) — rotate Q and K

Modern models (Llama, Qwen, Mistral, Gemma) use RoPE. Instead of adding a position vector to the input, RoPE rotates each query and key vector by an angle proportional to its position, just before the dot product. Take pairs of dimensions (x_{2i}, x_{2i+1}) and treat each pair as a point in a 2-D plane; rotate it by angle pos · θ_i:

[ x'_{2i}   ]   [ cos(pos·θ_i)  −sin(pos·θ_i) ] [ x_{2i}   ]
[ x'_{2i+1} ] = [ sin(pos·θ_i)   cos(pos·θ_i) ] [ x_{2i+1} ]          θ_i = 10000^(−2i/d_k)

The key property: the dot product of a rotated query at position m and a rotated key at position n depends only on m − n, not on m and n separately, because rotating one vector by a and another by b and taking their dot product equals the dot product with relative rotation a − b. Verify numerically with one 2-D pair and θ = 0.5:

import math, numpy as np
def rot(v, pos, theta=0.5):
    c, s = math.cos(pos*theta), math.sin(pos*theta)
    return np.array([c*v[0] - s*v[1], s*v[0] + c*v[1]])

q = k = np.array([1.0, 0.0])
for pq, pk in [(5, 3), (10, 8), (2, 0)]:        # all have offset 2
    print(pq, pk, rot(q, pq) @ rot(k, pk))      # 0.5403 every time (= cos(2 * 0.5))

Positions 5 vs 3, 10 vs 8 and 2 vs 0 all give the same score because the offset is 2 in each case. That is precisely “relative position”, and it is what lets the model learn “the token 3 positions back” patterns that work anywhere in the sequence. Different pairs use different θ_i (fast to slow rotation), so the model can read both fine local offsets and coarse long-range distance.

Context extension. Beyond the trained length, rotation angles reach values never seen in training. Fixes rescale the angles: linear position interpolation (divide positions by a factor), NTK-aware scaling (change the base 10000), and YaRN (different scaling per frequency plus a temperature adjustment), usually followed by brief fine-tuning on long text.

ALiBi — a distance penalty on the scores

ALiBi adds no position vectors at all; it subtracts a penalty proportional to distance from the attention scores of each head:

score(i, j) = q_i · k_j / sqrt(d_k)  −  m_h · (i − j)          m_h = a fixed slope per head

Nearby tokens are favoured; each head has a different slope (steep = local, shallow = long-range). Good length extrapolation, but RoPE has been the more widely adopted choice.

Summary

Scheme	Where applied	Absolute or relative	Extrapolates past training length
Sinusoidal	Added to input embeddings	Absolute (relative recoverable)	Somewhat
Learned absolute	Added to input embeddings	Absolute	No
RoPE	Rotates Q and K inside each attention layer	Relative	With scaling (YaRN, NTK)
ALiBi	Bias added to scores	Relative	Yes, naturally

The problem it solves

During generation the model keeps K and V for every past token in the KV cache (see the next concept). In standard multi-head attention (MHA) every one of the h heads has its own K and V, so the cache stores h key vectors and h value vectors per token per layer. At long contexts and large batch sizes that cache, not the weights, is what fills GPU memory, and every decode step must read all of it from memory, so it also limits speed. The idea of MQA/GQA is that the queries need to differ per head (that is where the different “what am I looking for” questions live), but the keys and values can be shared by several heads at little cost in quality.

Three variants

• MHA (multi-head) — h query heads, h KV heads.
• MQA (multi-query, Shazeer 2019) — h query heads, exactly 1 KV head shared by all.
• GQA (grouped-query, Ainslie et al. 2023) — h query heads, g KV heads, where g divides h. Query heads are split into g groups of h/g heads; all heads in a group use the same K and V. MHA is the special case g = h, MQA is g = 1.

MHA (h=8, g=8)   Q1  Q2  Q3  Q4  Q5  Q6  Q7  Q8
                 K1  K2  K3  K4  K5  K6  K7  K8       cache holds 8 K + 8 V per token
                 V1  V2  V3  V4  V5  V6  V7  V8

GQA (h=8, g=2)   Q1  Q2  Q3  Q4 | Q5  Q6  Q7  Q8
                 K1 (shared)    | K2 (shared)          cache holds 2 K + 2 V per token
                 V1             | V2

MQA (h=8, g=1)   Q1  Q2  Q3  Q4  Q5  Q6  Q7  Q8
                 K1 (shared by all)                    cache holds 1 K + 1 V per token
                 V1

Symbols used below: L = number of layers, h = query heads, g = KV heads, d_k = width of one head, n = tokens in the context, b = bytes per number (2 for fp16/bf16). The factor 2 counts K and V.

KV cache bytes = 2 · L · g · d_k · n · b          (per sequence)

Shape walkthrough

Take d_model = 16, h = 8, d_k = d_model / h = 2, and n = 5 tokens.

Tensor	MHA (g=8)	GQA (g=2)	MQA (g=1)
W_Q	[16, 16]	[16, 16]	[16, 16]
W_K, W_V	[16, 16] each	[16, 4] each	[16, 2] each
Q after head split	[8, 5, 2]	[8, 5, 2]	[8, 5, 2]
K, V as cached	[8, 5, 2]	[2, 5, 2]	[1, 5, 2]
K, V after repeat	[8, 5, 2]	[8, 5, 2]	[8, 5, 2]
Floats cached per token per layer	32	8	4

The repeat is a view/broadcast at compute time. Only the compact [g, n, d_k] tensors are stored in the cache, which is where the saving comes from. W_K and W_V also get smaller, so the model has slightly fewer parameters (often compensated by making the FFN a little wider).

Runnable GQA attention (NumPy, causal)

import numpy as np

def softmax(z):
    z = z - z.max(-1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(-1, keepdims=True)

def gqa(x, Wq, Wk, Wv, h, g):
    n, d_model = x.shape
    d_k = d_model // h
    q = (x @ Wq).reshape(n, h, d_k).transpose(1, 0, 2)   # [h, n, d_k]
    k = (x @ Wk).reshape(n, g, d_k).transpose(1, 0, 2)   # [g, n, d_k]  <- this is what gets cached
    v = (x @ Wv).reshape(n, g, d_k).transpose(1, 0, 2)   # [g, n, d_k]
    k = np.repeat(k, h // g, axis=0)                     # [h, n, d_k]  heads 0..h/g-1 use KV head 0, etc.
    v = np.repeat(v, h // g, axis=0)
    s = q @ k.transpose(0, 2, 1) / np.sqrt(d_k)          # [h, n, n]
    s = np.where(np.triu(np.ones((n, n), bool), 1), -np.inf, s)   # causal mask
    out = softmax(s) @ v                                 # [h, n, d_k]
    return out.transpose(1, 0, 2).reshape(n, d_model)    # merge heads

rng = np.random.default_rng(0)
n, d_model, h = 5, 16, 8
d_k = d_model // h
x = rng.standard_normal((n, d_model))
Wq = rng.standard_normal((d_model, h * d_k))
for g in (8, 2, 1):
    Wk = rng.standard_normal((d_model, g * d_k))
    Wv = rng.standard_normal((d_model, g * d_k))
    y = gqa(x, Wq, Wk, Wv, h, g)
    print(f"g={g}: K+V projection params {Wk.size + Wv.size}, cache floats per token {2 * g * d_k}, out {y.shape}")
# g=8: K+V projection params 512, cache floats per token 32, out (5, 16)
# g=2: K+V projection params 128, cache floats per token 8, out (5, 16)
# g=1: K+V projection params 64, cache floats per token 4, out (5, 16)

Line by line: queries keep all h heads; K and V are projected to only g heads (the Wk/Wv widths are g · d_k); np.repeat(..., h // g, axis=0) copies each KV head h/g times so that query head i uses KV head i // (h/g); after that the code is ordinary multi-head attention. Output shape is identical for every g, so the layer is a drop-in replacement.

Computed cache sizes (fp16, b = 2)

Shapes below are those of Llama-3-8B (32 layers, 32 query heads, 8 KV heads, d_k = 128) and Llama-3-70B (80 layers, 64 query heads, 8 KV heads, d_k = 128). The “MHA” and “MQA” rows are hypothetical variants of the same models, shown for comparison.

Model shape	Variant	KV heads	Per token	8,192 tokens	131,072 tokens
8B	MHA (hypothetical)	32	512 KiB	4 GiB	64 GiB
8B	GQA (actual)	8	128 KiB	1 GiB	16 GiB
8B	MQA (hypothetical)	1	16 KiB	0.125 GiB	2 GiB
70B	MHA (hypothetical)	64	2,560 KiB	20 GiB	320 GiB
70B	GQA (actual)	8	320 KiB	2.5 GiB	40 GiB
70B	MQA (hypothetical)	1	40 KiB	0.31 GiB	5 GiB

Check by hand for the 8B GQA row: 2 · 32 · 8 · 128 · 2 bytes = 131,072 bytes = 128 KiB per token; times 8,192 tokens = 1 GiB. These are per sequence: a batch of 16 requests multiplies them by 16.

Trade-offs and gotchas

• Quality. MQA loses a little quality and can train less stably; GQA with g around 4-8 recovers nearly all of MHA’s quality, which is why most current open models use it. Ainslie et al. also showed that an existing MHA checkpoint can be converted (“uptrained”) by mean-pooling its K/V heads into groups and training briefly.
• Speed. Decode is memory-bandwidth-bound, so reading a h/g-times smaller cache per step directly speeds up long-context generation and allows larger batches.
• Tensor parallelism. KV heads must be spread across GPUs; if g is smaller than the number of GPUs, KV heads get replicated and some of the saving is lost.
• Not the same as sharing queries. Queries stay per-head; only K/V are shared. DeepSeek’s MLA (multi-head latent attention) goes further by caching a compressed latent instead of K/V; see Mixture of experts and architecture families.

Interpretability research (mechanistic analysis of trained transformers) finds that heads develop narrow, identifiable jobs:

• Previous-token heads — attend to position i−1. A basic building block: they let a token know what came directly before it.
• Induction heads — implement “if A B appeared earlier and I now see A, predict B”. Example:

context:   ... Harry Potter ... Harry  ->  predict "Potter"
step 1:    a previous-token head copies "what token came before me" into each position
           (so the position of "Potter" now carries the info "I follow Harry")
step 2:    at the second "Harry", an induction head matches its query against keys that say
           "I follow Harry", attends to "Potter", and copies it forward to the output

This two-head circuit is the mechanistic basis of in-context copying and much of few-shot learning; it appears in training at a distinct sudden drop in loss.

• Syntactic heads — attend along grammatical links (verb to its subject, noun to its determiner).
• Name-mover / duplicate-token heads — in “John gave Mary a book; Mary thanked ___”, specific heads move the right name to the output.
• Attention sinks — many heads place a large share of their weight on the first token (often a beginning-of-sequence token) when they have nothing useful to attend to. Softmax weights must sum to 1, so the head needs somewhere to put its weight, and the first token acts as a harmless dump. Removing that token breaks the model, and StreamingLLM keeps a few sink tokens plus a recent window to run indefinitely long streams. Some recent models add learned “sink” parameters or gating so heads can output “nothing”.

Heads are not uniformly useful: many are redundant and can be pruned with small quality loss.

Two kinds of head

Attention is content-based in principle: the weights A = softmax(Q Kᵀ / sqrt(d_k) + M) depend on the actual tokens in this input, because Q and K are computed from them. But when researchers average or inspect the attention maps of trained heads over many inputs, a surprising fraction of heads turn out to produce almost the same map for every input. For those heads the learned W_Q/W_K are effectively computing a positional rule, not a content match. A head whose map barely changes with the input is called input-independent or fixed-pattern; a head whose map changes with what the tokens are is content-based.

The common patterns, as concrete attention maps

Take the 5-token sequence A B C A B (rows = the position that is asking, columns = the position it reads from, causal so only columns <= row). All matrices below are printed from the code in the next section.

Previous-token head — row i puts weight 1 on column i-1 (the first token has nothing earlier, so it reads itself):

[[1. 0. 0. 0. 0.]
 [1. 0. 0. 0. 0.]
 [0. 1. 0. 0. 0.]
 [0. 0. 1. 0. 0.]
 [0. 0. 0. 1. 0.]]

This is a one-step shift of the sequence. Output row i is V[i-1].

First-token (sink) head — every row puts its weight on column 0:

[[1. 0. 0. 0. 0.]
 [1. 0. 0. 0. 0.]
 [1. 0. 0. 0. 0.]
 [1. 0. 0. 0. 0.]
 [1. 0. 0. 0. 0.]]

Every output row equals V[0]. This is the “attention sink”: the head has nothing useful to read and parks its weight on a harmless token (see What attention heads learn).

Uniform (average) head — row i spreads weight equally over columns 0..i:

[[1.   0.   0.   0.   0.  ]
 [0.5  0.5  0.   0.   0.  ]
 [0.33 0.33 0.33 0.   0.  ]
 [0.25 0.25 0.25 0.25 0.  ]
 [0.2  0.2  0.2  0.2  0.2 ]]

Output row i is the running mean of V[0..i], a cheap summary of the whole prefix.

Induction head — content-based: position i looks for an earlier occurrence of its own token and attends to the token that came after it. With A B C A B: position 3 (A) finds the earlier A at position 0 and reads position 1 (B); position 4 (B) finds the earlier B at position 1 and reads position 2 (C), which is exactly what “predict the next token by copying” needs:

[[1. 0. 0. 0. 0.]      position 0: no earlier match, reads itself
 [0. 1. 0. 0. 0.]      position 1 (B): no earlier B, reads itself
 [0. 0. 1. 0. 0.]      position 2 (C): no earlier C, reads itself
 [0. 1. 0. 0. 0.]      position 3 (A): earlier A at 0 -> read position 1 (B)
 [0. 0. 1. 0. 0.]]     position 4 (B): earlier B at 1 -> read position 2 (C)

Change the input to A B C D E and this map becomes the identity: the map depends on the content, which is the defining difference from the three heads above. Those three would print the same matrices for any five tokens.

How a fixed replacement is computed

1. Run the model on a calibration set of many varied inputs and record the attention matrix A of the head under study for each.
2. Average them to get A_bar (shape [n, n], or as a function of relative offset i - j so it works for any n).
3. Check how far each input’s map is from the average (for instance the maximum per-entry standard deviation, or the relative error of the head’s output).
4. If the head is close to constant, replace the softmax(Q Kᵀ ...) computation by the constant A_bar (or a cheap operation it implies: a shift, a prefix mean, a lookup of token 0). Then no Q, no K, no n x n score matrix is needed for that head, only A_bar V.

import numpy as np
rng = np.random.default_rng(0)
n, d = 6, 8
def softmax(z):
    z = z - z.max(-1, keepdims=True); e = np.exp(z); return e / e.sum(-1, keepdims=True)
causal = np.triu(np.ones((n, n), bool), 1)

def prev_token_head(x):                     # toy head: strong positional bonus + weak content term
    s = 0.5 * (x @ x.T) / np.sqrt(d) + 6.0 * np.eye(n, k=-1)
    return softmax(np.where(causal, -np.inf, s))

maps = np.stack([prev_token_head(rng.standard_normal((n, d))) for _ in range(2000)])
A_bar = maps.mean(0)                        # step 2: the fixed replacement
print(A_bar.round(2)[:3])
# [[1.   0.   0.   0.   0.   0.  ]
#  [0.98 0.02 0.   0.   0.   0.  ]
#  [0.   0.98 0.02 0.   0.   0.  ]]
print("max per-entry std:", maps.std(0).max().round(2))        # 0.03  -> nearly constant

errs = []
for _ in range(200):                        # step 3: output error when swapping in A_bar
    x, v = rng.standard_normal((n, d)), rng.standard_normal((n, d))
    real = prev_token_head(x) @ v
    errs.append(np.linalg.norm(real - A_bar @ v) / np.linalg.norm(real))
print("relative error:", np.mean(errs).round(3))               # 0.026  -> about 3 percent

For contrast, repeating the same procedure for a head whose scores are x W xᵀ with a random W (a purely content-based toy head) gives a maximum per-entry standard deviation of about 0.37 and a relative output error of about 0.61: the averaged map is a poor stand-in. These numbers come from this toy setup and are illustrative, not measurements of a real model; real calibration must be done on the real model and data, and the per-head error then judged by downstream quality, not output error alone.

Why it matters

Heads doing real content-based retrieval (induction, name movers) are the ones that resist replacement. Practically:

• KV eviction strategies such as “keep the first few tokens plus a recent window” work far better than the O(n²) framing suggests, because many heads only use sinks and recent context.
• Hybrid architectures keep a few full-attention layers for retrieval and use cheaper layers (sliding window, linear attention) elsewhere.
• Head pruning and sparsification methods score heads by how input-dependent they are and keep full attention only for the ones that need it.
• Long-context failures (the “lost in the middle” effect) often come from the few retrieval heads having to find a needle among many positions, not from the cheap heads.

Pitfalls

• A head can look fixed on average but switch behaviour on rare inputs (a head that is a sink almost always but retrieves when a particular pattern appears); average-map error hides that.
• Position-dependent maps depend on n and on RoPE-style relative offsets; express the fixed pattern in relative offsets, not absolute indices.
• Replacing heads after training changes the activations downstream layers were trained on, so brief fine-tuning is usually needed to recover quality.