All About Attention

Attention lets each token combine information from other tokens. This guide starts with a weighted lookup, builds the matrix formulation, and then follows the implementation choices that affect memory and speed.

Before you start: be comfortable with matrix multiplication, vectors and basic probability. Read sections 1–3 first; the later sections are reference material you can return to as you implement each idea. Model examples describe particular architectures, not requirements for every transformer.

1. Attention as a soft lookup

Attention is a weighted lookup. A query represents what a position is looking for. Each candidate position provides a key for comparison and a value containing the information it can contribute.

Compare the query with every allowed key to get a score for each position. Softmax converts those scores into non-negative weights that sum to one. The output is the weighted sum of the corresponding values. Unlike a dictionary lookup, attention can combine information from several positions.

Concretely, with $N$ keys $k_1, \ldots, k_N \in \mathbb{R}^d$, matching values $v_1, \ldots, v_N$, and a query $q \in \mathbb{R}^d$:

$$\text{output} = \sum_{i=1}^N \alpha_i \, v_i, \qquad \alpha_i = \frac{\exp(\text{score}(q, k_i))}{\sum_j \exp(\text{score}(q, k_j))}$$

Each weight $\alpha_i$ tells us how much value $v_i$ contributes. If one score is much larger than the others, its weight approaches one and the output approaches that value. Equal scores give equal weights, so the output is the mean of the values. The values can have a different dimension from the queries and keys.

In self-attention, learned projections produce queries, keys and values from the same sequence of hidden states. In cross-attention, queries come from one sequence while keys and values come from another. A hidden state is the vector representation of a token at that point in the model.

Check your understanding Why use a weighted sum rather than picking the single best-matching key? What does the soft version buy us that argmax does not? Show answer Hide answer

Soft weights let an output combine evidence from several positions. They also provide a differentiable route from the loss to the scores: changing a score changes the mixture continuously.

A hard argmax index is locally constant except at ties, so ordinary backpropagation through that selection does not tell the model how to improve the scores. Hard-selection models can use other estimators or training strategies; softmax is one convenient approach, not the only possible one.

2. Scaled dot-product attention

Dot-product attention uses $\text{score}(q,k)=q\cdot k$. If each row of $Q$ is a query and each row of $K$ is a key, the matrix product $QK^\top$ computes every query-key score at once. This maps well to optimized matrix-multiplication routines.

A dot product depends on both vector direction and magnitude. A larger magnitude is not automatically greater confidence. Other score functions, including cosine similarity, make different choices about which information to retain.

Why scale the scores? Suppose all query and key components are independent, with mean zero and variance one. For $d_k$ components, the dot product is a sum of $d_k$ independent terms, each with variance one. Its variance is therefore $d_k$, and its standard deviation is $\sqrt{d_k}$.

Larger score differences can make softmax concentrate on one key, leaving small derivatives that make score updates harder. Dividing by $\sqrt{d_k}$ controls the score scale under this model:

$$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V$$

Here $d_k$ is the number of features in each query or key for one head. Softmax is applied separately to each query's row of scores, across the key positions. Multiplication by $V$ then forms one weighted value sum per query.

For $N_q$ queries, $N_k$ keys and value dimension $d_v$, the shapes are $Q:(N_q,d_k)$, $K:(N_k,d_k)$ and $V:(N_k,d_v)$. The score matrix has shape $(N_q,N_k)$ and the output has shape $(N_q,d_v)$.

The scaled scores have variance one under the assumptions above. Learned scales and correlations can change that during training. See Attention Is All You Need, section 3.2.1.

Check your understanding What would happen if we used $QK^\top$ without the $\sqrt{d}$ scaling? Walk through the consequences in concrete numbers for a 128-dim head. Show answer Hide answer

Under the independent, zero-mean, unit-variance initialization model, the raw score standard deviation is $\sqrt{128}\approx11.3$. This does not determine the largest score or its gap to the runner-up in a particular row.

For an explicit two-key example, take scores $[6,0]$. Softmax assigns about 0.998 to the first key. Divide those same scores by $\sqrt{128}$ and its weight becomes about 0.630. These are illustrative scores, not a measurement of a random 4096-key attention row.

The softmax derivative is $\partial\alpha_i/\partial s_j=\alpha_i(\delta_{ij}-\alpha_j)$. As one weight approaches one, its own derivative and the others become small. Scaling can reduce this concentration; it does not guarantee that all heads avoid saturation or that omitting it makes an entire layer permanently inactive.

A worked example, $T=3$, $d_k=2$

Use three query positions and three key/value positions, with two features per vector. This example has no mask or dropout. The matrices are

$$ Q = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix}, \quad K = \begin{bmatrix} 1 & 0 \\ 0 & 1 \\ 1 & 1 \end{bmatrix}, \quad V = \begin{bmatrix} 2 & 0 \\ 0 & 2 \\ 1 & 1 \end{bmatrix}. $$

Compute $S = QK^\top$ row by row. Row 0 (query $[1,0]$) dotted with each key gives $[1, 0, 1]$. Row 1 ($[0,1]$) gives $[0, 1, 1]$. Row 2 ($[1,1]$) gives $[1, 1, 2]$. So

$$ S = \begin{bmatrix} 1 & 0 & 1 \\ 0 & 1 & 1 \\ 1 & 1 & 2 \end{bmatrix}, \qquad S/\sqrt{2} \approx \begin{bmatrix} 0.71 & 0 & 0.71 \\ 0 & 0.71 & 0.71 \\ 0.71 & 0.71 & 1.41 \end{bmatrix}. $$

Take row 0 the rest of the way. Exponentiate: $e^{0.71}, e^{0}, e^{0.71} \approx 2.03, 1, 2.03$. Sum $= 5.06$. Softmax weights $\alpha_0 \approx [0.40, 0.20, 0.40]$. Multiply by $V$:

$$ \text{out}_0 = 0.40 \cdot [2,0] + 0.20 \cdot [0,2] + 0.40 \cdot [1,1] = [1.20, 0.80]. $$

Rows 0 and 2 each contribute about 40%, and row 1 contributes about 20%. The output is a blend of their values, rounded here to $[1.20,0.80]$. Repeat the same calculation for each query row. Batches and heads add separate copies of this calculation; masks determine which keys each query may use.

A temperature $\tau>0$ changes the distribution by dividing scores by $\tau$ before softmax. Lower temperatures concentrate weight on larger scores; higher temperatures spread it more evenly. Change the slider to see the effect on these five fixed scores:

As $\tau$ approaches zero, a unique highest score receives weight approaching one. If several scores tie for the maximum, they share that limiting weight. As $\tau$ grows, the weights approach $1/5$ each.

The displayed entropy measures how spread out the weights are: $-\sum_i p_i\ln p_i$, with $0\ln0$ treated as zero. It is measured in nats because the logarithm is natural. For five choices it ranges from zero to $\ln5\approx1.61$.

Standard scaled dot-product attention uses no additional temperature after dividing by $\sqrt{d_k}$. That division is algebraically equivalent to a temperature of $\sqrt{d_k}$ applied to the raw dot products. Its variance justification relies on the assumptions above; it does not keep every trained model's score variance fixed.

QK-norm and cosine attention

The variance argument describes an initialization model. Learned query/key norms and correlations can change score distributions during training. Large score differences can concentrate softmax weights; numerical behavior also depends on dtype, quantization and implementation.

Normalizing queries and keys can control score magnitudes, often with a learned scale to retain flexibility. The exact normalization differs across architectures. It is one tool for training stability, not a guarantee across all precisions or training configurations.

3. Multi-head attention

Multi-head attention learns several query, key and value projections. In the common setup here, the model width $D$ is split across $H$ heads, each with $d_k=d_v=D/H$ projected features. Each head computes its own attention weights and output. Concatenate the head outputs, then apply the learned output projection $W^O$:

$$\text{MultiHead}(X) = \text{Concat}(\text{head}_1, \ldots, \text{head}_H) W^O$$

with $\text{head}_h = \text{Attention}(X W^Q_h, X W^K_h, X W^V_h)$. Why not one big head with the full $D$? Each head can learn a different attention pattern for the same token. The heads use separate learned projections, so they can combine different information.

The visualization splits the projected features into $H$ blocks of $d_k = D/H$ each. Each projection can use all input features; a head is not restricted to a slice of the original input. Try increasing $H$ while keeping $D$ fixed:

For standard multi-head attention with $d_k = d_v = D/H$, the projection parameter count stays fixed as $H$ changes. The two attention matrix multiplications also retain the same leading arithmetic cost: $H N^2 (D/H) = N^2 D$ each. This does not make extra heads free in practice: softmax work grows with $H$, a materialized score tensor has $H N^2$ entries per example, and kernel efficiency depends on the shapes. See Attention Is All You Need, section 3.2.

Induction heads: a studied copying pattern

Olsson et al. (2022) study a copying pattern: after seeing A B ... A, a model predicts B. In a simple two-layer attention-only circuit:

1. A previous-token head copies information about the preceding token into the current position's representation. The position containing the earlier B therefore carries information about its preceding A.
2. An induction head at the later A matches that preceding-token information, attends to the earlier B position, and promotes B as the next token.

The learned weights implement the circuit; the prompt's particular A B pair is carried in activations, without updating weights during inference. This is evidence for a mechanism of in-context learning, not proof that all such learning in large models reduces to this circuit.

Check your understanding Are more heads always better? With model width $D=4096$, compare 32 heads ($d_k=128$), 64 heads ($d_k=64$) and 128 heads ($d_k=32$). What would you measure before choosing? Show answer Hide answer

Not necessarily. With fixed $D = 4096$, those choices give head dimensions 128, 64 and 32. More heads provide more separate attention patterns, but each query/key comparison uses a smaller learned representation.

The scaled-dot-product variance argument explains the $1/\sqrt{d_k}$ factor; it does not establish an optimal head count or a minimum useful head dimension.

Treat head count as a tradeoff to measure. Compare validation quality and runtime under the same training budget, and check memory at your intended sequence length. Keep the number of key/value heads in view too: the GQA section below explains why query-head count alone does not determine the KV-cache size.

4. Position encoding

Self-attention needs a way to distinguish token order. Without position-dependent masks or features, reordering the input simply reorders its output. This property is called permutation equivariance. Position features let the model distinguish sequences such as "dog bites person" and "person bites dog".

Sinusoidal encodings add a fixed vector of sine and cosine values to each token embedding. Different positions receive different vectors. Learned absolute encodings instead learn one vector per position in a table. Using positions beyond that table requires extending or replacing it. Neither approach alone guarantees good results on sequences longer than those used in training.

Rotary Position Embedding (RoPE) changes queries and keys rather than adding a vector to token embeddings. It groups an even number $d$ of their features into pairs. At position $p$, pair $i$ rotates by angle $p\theta_i$, where $\theta_i=b^{-2i/d}$ and $i=0,\ldots,d/2-1$. Here $b>1$ is the chosen frequency base. Each rotation preserves the length of that pair.

For fixed query and key vectors, rotating them at positions $p$ and $q$ makes their dot product depend on the difference $p-q$. Moving both by the same offset leaves that dot product unchanged. The vectors still carry token information, so this is a property of the positional transformation, not a claim that content stops mattering. See RoFormer.

Move the slider to compare rotation rates. Small pair indices rotate faster; large indices rotate more slowly. The diagram illustrates these frequencies, not a claim that individual pairs must learn particular syntactic roles.

The period is the number of position steps needed for a full rotation: $2\pi/\theta_i$. The slowest pair therefore has period $2\pi/\theta_{d/2-1}$. The chosen base and rotary dimension determine it; the training sequence length does not. Computing rotations at larger positions is possible, but using those positions effectively may require further training.

ALiBi adds a distance penalty to attention scores. For a causal query at position $i$, a key at $j\le i$ gets an added score of $-m_h(i-j)$, where $m_h>0$ is the slope for head $h$. A key ten positions away receives a larger penalty than one position away. Content scores still matter, so the closest key need not win. The paper reports useful length extrapolation in its experiments; longer sequences still need evaluation for a particular model and task.

Extending the context length

• Position Interpolation replaces position $p$ with $p/s$ for an extension factor $s>1$, then fine-tunes the model. This fits a longer sequence into the original position range and reduces the angle change between neighboring tokens.
• Base scaling changes $b$, affecting lower frequencies more while $\theta_0=1$ remains unchanged. Its quality depends on the model and extension recipe.
• YaRN adjusts different rotation frequencies differently and also rescales attention. Its results depend on the evaluated models and training setup.
• LongRoPE searches for interpolation settings that can vary across dimensions and positions. It extends context length in stages and includes adjustments to recover short-context performance.

Compare short-context quality as well as long-context retrieval after extending a model. A larger configured maximum is not evidence that the model can reliably use every token in it.

Check your understanding What does position interpolation do to RoPE's fastest and slowest pairs? Show answer Hide answer

Replacing $p$ by $p/s$ changes every angle from $\theta_i p$ to $\theta_i p/s$.

With $d=128$ and $b=10000$, the fastest pair is $i=0$: its increment is 1 radian per position, or 0.125 after interpolation by $s=8$.

The slowest pair is $i=63$: its increment is $10000^{-126/128}\approx1.155\times10^{-4}$ radians, or $1.443\times10^{-5}$ after interpolation. Both increments are divided by eight, but the fastest pair still changes much more between neighboring positions.

These calculations describe the change in position features. To judge whether the change helps, measure the adapted model on both short and long sequences.

5. Pre-LN vs Post-LN

Layer normalization (LN) rescales the features at each token position using their mean and variance, followed by learned scale and bias. A residual connection adds a sublayer's output back to its input. Let $x$ be that input and $F$ the attention or feed-forward sublayer. The two common arrangements are:

$$\text{Post-LN:}\quad y=\text{LN}(x+F(x))$$

$$\text{Pre-LN:}\quad y=x+F(\text{LN}(x))$$

Post-LN normalizes after the addition. Pre-LN normalizes the input to $F$ and leaves a direct path from $x$ to the output.

That direct path contributes the identity matrix $I$ to the derivative. The full derivative, or Jacobian, also includes the path through normalization and $F$: $I+J_F J_{\text{LN}}$. It describes how changes in the input affect the output. The direct path helps explain Pre-LN's behavior, but it does not guarantee unchanged gradients across many blocks.

Xiong et al. study gradients at initialization and the role of learning-rate warmup, which gradually raises the learning rate at the start of training. Their results support Pre-LN's optimization benefits in the studied setting. Stability still depends on initialization, optimizer settings and the rest of the architecture.

Stabilising attention training at scale

• Logit scale: logits are the scores passed to softmax. Large differences can concentrate its weights and make derivatives small. For a positive cap $c$, replacing logit $z$ with $c\tanh(z/c)$ smoothly bounds its magnitude; Gemma 2 is one documented example. The cap and its effect must be evaluated as part of the training recipe.
• Query/key norms: normalization and learned scaling can control attention-score magnitudes. They do not independently guarantee stability for the whole model.
• muP and muTransfer: muP, or maximal update parameterization, coordinates how parameters are initialized and scaled as model width changes. muTransfer uses this setup to transfer tuned hyperparameters from smaller models to larger ones. It requires a consistent training setup. Follow the authors' implementation and checks.
• Output z-loss: for output logits $z_j$, let $Z=\sum_j e^{z_j}$ be the denominator of softmax. Adding $\lambda(\log Z)^2$ to the training loss penalizes a large magnitude of $\log Z$; $\lambda$ controls the penalty strength. PaLM uses this auxiliary loss; it is separate from attention normalization and does not bound every individual logit by itself.

Measure the quantity a technique is meant to control, such as score magnitude or gradient size, alongside validation loss and training stability.

6. Which positions can attend to which?

Two separate choices matter: where the queries, keys and values come from, and which query-key pairs are allowed. A mask removes disallowed pairs from the attention calculation:

• Bidirectional self-attention: queries, keys and values come from the same sequence. Tokens can use both earlier and later valid positions; padding or task-specific masks may still apply.
• Causal self-attention: query position $i$ uses only allowed keys at $j\le i$. Future score entries are masked before softmax, usually with $-\infty$ in an additive mask.
• Cross-attention: queries come from one sequence and keys/values from another. Query and source lengths may differ. Fully observed source tokens need no causal restriction, but source padding and task-specific constraints still need masking.

Check the API's mask convention: in PyTorch MultiheadAttention, a true boolean padding-mask entry means ignore that key. Do not infer a boolean mask's meaning merely from its shape. Choose the attention pattern to match the information available at prediction time.

7. The KV cache

A key/value (KV) cache stores the keys and values already computed for earlier tokens at each attention layer. During causal inference, those tokens cannot attend to newly appended tokens. Their cached representations therefore remain reusable when model weights, earlier inputs and position handling stay unchanged.

Processing the prompt is called prefill; generating subsequent tokens is called decoding. Each decoding step computes the new token's query, key and value. Its query attends to the old cached keys and its new key, and the new key/value pair is added to the cache. See the Transformers caching guide.

The cache saves recomputation, but each new query still reads the earlier keys and values. With model width fixed, attending to $T$ stored positions takes work proportional to $T$, written $O(T)$. Generating $N$ tokens after a prompt of length $P$ costs $O(PN+N^2)$ for these attention operations. This count excludes the initial prefill and other model operations.

For ordinary cached attention with equal key and value head widths, the storage per request is

$$2\cdot T\cdot n_{\text{layers}}\cdot H_{KV}\cdot d_{\text{head}}\cdot b,$$

Here $T$ is the number of stored tokens, $n_{\text{layers}}$ the number of attention layers, $H_{KV}$ the number of key/value heads, $d_{\text{head}}$ their width, and $b$ the bytes per element. The factor 2 counts keys and values.

For 80 layers, 8 KV heads, width 128 and two-byte storage:

Stored tokens	KV-cache size
1	327,680 bytes = 320 KiB
32,768	10 GiB
131,072	40 GiB

One KiB is 1,024 bytes; one GiB is $2^{30}$ bytes. These totals cover only the stored keys and values, not model weights, other activations, memory-management overhead or extra requests. Splitting the cache across devices or storing values at lower precision changes the allocation per device.

Paged attention

Suppose a server reserves enough cache for the longest allowed sequence for every request. Short requests leave much of that reservation unused. PagedAttention instead stores the cache in fixed-size memory blocks. A mapping records which block contains each range of token positions, so blocks can be allocated as a sequence grows. A partly filled block can still have unused space.

Compatible requests can share cached prefix blocks. This saves duplicate storage and may avoid repeated prefix computation; it does not automatically eliminate each request's later attention reads or guarantee a fixed bandwidth saving. Share only prefixes whose model state and positional context produce the same cached values.

Speculative decoding and the cache

Speculative decoding uses a smaller or cheaper draft model to propose several tokens. The target model checks those candidates together, allowing some steps to be accepted without a separate target-model call for each one. Checking several candidates still performs more work than checking one, even when both fit in a single batched call.

Keep the cache entries for the accepted continuation and remove entries that belong only to rejected candidates. The cache implementation determines how to truncate or reuse its blocks.

Speed depends on how many candidates are accepted, how much drafting costs and how efficiently the target model checks a batch. To preserve the target model's sampling distribution, use the method's acceptance and correction probabilities, rather than an informal plausibility check. See Chen et al..

8. GQA, MQA, and MLA

A model can keep several query heads while sharing their keys and values. The names describe how that sharing works:

• Multi-head attention (MHA): each query head has its own key/value head.
• Multi-query attention (MQA) shares one key head and one value head across all query heads.
• Grouped-query attention (GQA) shares each KV head across a group of query heads. For 64 query heads and 8 KV heads, each group contains 8 query heads.

Keep token count, layer count, head width and storage precision fixed. Reducing 64 KV heads to 8 then uses one eighth of the cache. Reducing 8 to 1 uses one eighth again. These ratios describe cache storage. The effects on quality and speed depend on training, task and implementation.

Check your understanding Are the savings from 8 KV heads to 1 necessarily too small to matter? Show answer Hide answer

No. They remove 87.5% of the eight-head cache under the same storage assumptions. In the 32,768-token example above, that is a reduction from 10 GiB to 1.25 GiB per request. Whether it is worthwhile depends on quality, total memory pressure, parallelization and measured throughput.

Multi-head Latent Attention: caching a compressed representation

Multi-head Latent Attention (MLA) stores a smaller learned representation from which key/value information can be obtained. This stored vector is called a latent. DeepSeek-V3 caches a 512-element latent plus a separate 64-element positional key at each layer.

Some projection matrices can be combined with the query and output projections, avoiding the need to reconstruct and cache every full key/value head during decoding. The separate positional component is still needed. Actual speed depends on how the implementation carries out these operations.

At two bytes per element, a layer storing those 576 elements uses 1,152 bytes per token. The eight-KV-head example above stores $2\cdot8\cdot128=2048$ elements, or 4,096 bytes per layer. Across 61 MLA layers versus 80 such GQA layers, the totals are 70,272 versus 327,680 bytes per token. These are cache-storage calculations, excluding auxiliary prediction modules and allocator overhead; they do not compare model quality or latency. See the DeepSeek-V3 report.

Check your understanding Does a smaller KV cache always mean faster inference? Show answer Hide answer

No. Less cache can reduce memory pressure and improve the batch size you can serve, but latency also depends on projection dimensions, kernel efficiency, context length and hardware. Compare memory in bytes using the same precision, then benchmark the complete implementation on the workload you intend to serve.

9. Sliding window and local attention

A causal sliding window lets each query use only its own position and a limited number of recent positions. If the window contains at most $W$ keys, a sequence of $N$ queries needs at most $NW$ query-key comparisons instead of $N^2$. At fixed head width, this changes attention work from $O(N^2)$ to $O(NW)$.

Stacking layers extends the possible flow of information. If each layer can look back $W-1$ positions, $L$ layers can connect positions up to $L(W-1)$ apart through intermediate states. That path does not give the final query direct access to every old token or guarantee that a detail survives.

Compare the masks below. The local-plus-prefix example keeps the first two tokens visible alongside a recent window. It illustrates a connectivity pattern, not the full implementation of a particular sparse-attention model.

Choose the window by measuring quality, memory and runtime on the intended workload. Include tasks that require distant details; a smaller window may save memory while losing information those tasks need.

Attention sinks

StreamingLLM (2023) observed that initial tokens can attract substantial attention despite limited semantic importance. Retaining their cached keys and values alongside a recent window improved streaming behavior in the evaluated pretrained models.

This preserves useful computation with a bounded cache; it does not retain the text of all evicted tokens or make arbitrary old facts directly retrievable. A sink is an observed attention pattern. A global token is a connectivity choice in a mask. They are not interchangeable definitions, and neither is a requirement of every sparse method.

Check your understanding If attention weights sum to less than one, is the weighted output necessarily unbounded? Show answer Hide answer

No. For nonnegative weights with $\sum_i \alpha_i \le 1$, the triangle inequality gives

$$\left\|\sum_i \alpha_i v_i\right\| \le \sum_i \alpha_i\|v_i\| \le \max_i\|v_i\|.$$

For example, append a zero-valued option with logit zero to softmax. Real-token weights become $e^{s_i}/(1+\sum_j e^{s_j})$, so they sum to less than one. The remaining mass multiplies zero; this preserves the bound on the weighted values.

Whether such a change improves training or reduces attention sinks is an empirical question. The inequality alone establishes neither outcome, and it does not bound the entire network after other projections and layers.

10. Sparse attention

Sparse attention computes only a selected subset of query-key pairs. A local window is one sparse pattern; other methods add connections to distant positions:

• BigBird (2020) combines local, random and global connections. Its expressivity results apply under the paper's assumptions; they do not promise identical outputs or quality to a particular full-attention model.
• Reformer (2020) uses locality-sensitive hashing to place similar vectors into candidate groups for attention. It also uses reversible layers, whose inputs can be reconstructed during the backward pass, to reduce stored activations. This candidate selection avoids scoring every query-key pair.
• Native Sparse Attention (2025) combines compressed blocks, selected blocks containing original token representations, and a local window. A learned gate combines these paths. Selection uses attention information from the compressed path; block sizes and selection budgets remain design choices. These settings still need to be chosen and evaluated.

Fewer connections do not automatically mean lower latency: selecting candidates, gathering memory and executing small or irregular operations also cost time. Compare quality and measured prefill/decoding performance, not just the number of nonzero mask entries.

11. Linear attention

One family of linear attention methods changes the similarity rule so that history can be summarized in a fixed-size state. A feature map $\phi$ turns each query or key into an $r$-dimensional vector. Use the nonnegative similarity $\phi(q_i)^\top\phi(k_j)$ in place of the exponential dot-product score, then normalize the weights.

With value width $d_v$, store a matrix $S$ that accumulates key/value outer products and a vector $z$ that accumulates key features:

$$ S = \sum_j \phi(k_j)v_j^\top \in \mathbb{R}^{r \times d_v}, $$

$$ z = \sum_j \phi(k_j) \in \mathbb{R}^{r}, $$

$$ o_i = \frac{\phi(q_i)^\top S}{\phi(q_i)^\top z}. $$

The numerator combines the stored values according to the current query. The denominator is the sum of that query's similarities to the stored keys.

For causal attention, start with zero $S$ and $z$, then add the current key/value contribution before computing output $o_i$. The sums then contain only positions $j\le i$. The stored state contains $rd_v+r$ scalars per head. Each update and read costs $O(rd_v)$ at fixed widths, independent of how many positions have already been processed. Producing $N$ outputs still requires $N$ steps. See Transformers are RNNs, section 3.

Check the denominator: if similarities are $[2,1]$ and scalar values are $[10,4]$, the normalized output is $(2\cdot10+1\cdot4)/(2+1)=8$, not 24. A zero denominator needs explicit numerical handling. Strictly positive query and key features give a positive denominator when at least one key is present in exact arithmetic. An empty key set, underflow or overflow still needs explicit handling.

Kimi Linear (2025) is a different recurrent design. Its Kimi Delta Attention uses learned gates to control state updates and is combined with MLA layers. Because the model mixes recurrent and full-attention components, the fixed-size state above does not describe its total memory use. See the Kimi Linear report.

Check your understanding Does linear attention necessarily give the same answer as softmax attention? Show answer Hide answer

No. A different kernel changes the weighting rule; a finite feature approximation introduces approximation error. Normalizing its weights does not restore softmax equivalence.

Recurrent state compresses history, so test retrieval and long-context behavior alongside throughput. Quality depends on the architecture, training and task; linear scaling alone does not prove either good or bad recall.

12. FlashAttention

FlashAttention computes softmax attention in small blocks so that less intermediate data must be written to and read from GPU memory. It preserves the mathematical attention operation; it does not replace the weights with a sparse or linear approximation. Floating-point results can still differ because the operations are ordered differently.

A straightforward implementation materializes the $N \times N$ score and probability matrices in high-bandwidth memory (HBM). A single $4096 \times 4096$ FP16 matrix occupies 32 MiB per example per head. Writing and rereading these intermediates can be expensive. Whether memory traffic or arithmetic limits a particular workload depends on its shapes, hardware and implementation.

The online softmax recurrence

Online softmax updates its result as blocks arrive. Here "online" means incremental computation, not an internet connection. Hold a block of $B_q$ queries and process blocks of $B_k$ keys. The scaled scores for key block $j$ form a matrix $s^{(j)}$ of shape $(B_q,B_k)$.

For each query row, keep a running maximum $m$, a sum of exponentials $\ell$, and an unnormalized weighted-value sum $O$. Initialize $m=-\infty$, $\ell=0$ and $O=0$. The equations below assume each processed row has at least one finite score in each block; masking needs the additional handling described below.

For each query row, take the maximum and sum across the current block's key positions:

$$ m^{(j)} = \max\!\left(m^{(j-1)}, \; \max(s^{(j)})\right), \qquad \ell^{(j)} = e^{m^{(j-1)} - m^{(j)}} \, \ell^{(j-1)} + \sum e^{s^{(j)} - m^{(j)}}. $$

Update the weighted-value sum using the values $V^{(j)}$ for that key block:

$$ O^{(j)} = e^{m^{(j-1)} - m^{(j)}} \, O^{(j-1)} + e^{s^{(j)} - m^{(j)}} \, V^{(j)}. $$

After the last key block $J$, divide each row of $O^{(J)}$ by its scalar denominator $\ell^{(J)}$. If a new block raises the maximum, multiply the old sums by $e^{m^{(j-1)}-m^{(j)}}$ so old and new contributions use the same scale. Subtracting the running maximum keeps the exponentials from overflowing for finite scores.

Masked blocks: skip the update for a query row when the current block has no valid keys. Otherwise an initial all-masked block would evaluate $-\infty-(-\infty)$ and produce NaN. If no block has a valid key for that query, there is no normalized attention distribution. Define an explicit output policy, such as a zero row, and test that the implementation follows it.

Tiling means dividing the computation into these smaller blocks. Keeping working data in fast on-chip memory, including SRAM, avoids storing the full $N\times N$ score matrix in the larger high-bandwidth memory (HBM). Blocks may still be read more than once. The amount of data movement saved depends on head dimension and available on-chip memory. FlashAttention computes exact attention up to floating-point rounding, rather than approximating the attention pattern. See the FlashAttention paper.

Implementation evolution

FlashAttention-2 changes how work is divided across the GPU and reduces work outside matrix multiplication. FlashAttention-3 targets NVIDIA Hopper GPUs, overlaps computation with data movement, and adds an eight-bit floating-point (FP8) path with techniques to control quantization error. Their reported speedups depend on the benchmark configuration; they are not a fixed multiplier for every model. See the FlashAttention-3 paper.

For short sequences, the overhead of launching and coordinating a GPU kernel can outweigh the benefit of tiling. Activation checkpointing saves memory by recomputing some intermediate values during the backward pass. Combining it with an attention kernel that already recomputes intermediate values can add further work. Measure peak memory and runtime together before choosing the configuration.

Check your understanding If FlashAttention computes exact attention, why can its output differ from a simple implementation? Show answer Hide answer

Here, exact means that the algorithm does not approximate the attention pattern. Reordering floating-point operations can still change rounding, so mathematical equivalence does not guarantee identical bits. Check outputs with tolerances appropriate to the dtype and workload. PyTorch documents backend-dependent numerical differences.

Kernel support is a separate question. Device, dtype, head dimension and mask layout affect which implementation can run. A custom score or sparsity pattern may require a different kernel or a fallback; online softmax itself is not restricted to causal masks. Check backend support and benchmark your actual shapes before assuming a fused path is being used.

13. Subtracting attention maps and excluding self-matches

Differential Transformer forms two softmax attention maps and subtracts a learned multiple of one from the other before applying them to values. Unlike one softmax map, the resulting weights need not be nonnegative or sum to one. The paper motivates this as a way to reduce attention on irrelevant context and reports gains on its evaluated tasks. Those results do not establish a universal improvement or a fixed compute penalty across implementations. See the Differential Transformer paper.

Excluding self-matches means preventing each query from attending to the key at its own position, by masking the diagonal of the score matrix. Set forbidden logits to $-\infty$, not zero: a zero logit still receives probability after softmax. In causal attention, excluding the diagonal leaves the first token with no valid keys, so an implementation needs an explicit policy for that row.

There is no general identity $q_i \cdot k_i = \|q_i\|^2$: queries and keys have separate learned projections. Nor does the residual connection make self-attention redundant; the value and output projections can transform the token's representation. Treat diagonal masking as an architectural experiment, not a guaranteed retrieval improvement.

14. Choosing and testing an implementation

Start from a tested implementation and change one constraint at a time:

• Correctness: verify shapes, scaling, masks and all-masked-row behavior on small examples before optimizing.
• Quality: evaluate the tasks and context lengths you care about. Include retrieval tests if distant details matter; short-context loss alone cannot establish that behavior.
• Memory: measure peak training memory and decoding KV-cache size separately. GQA, MLA, sparse patterns and recurrent state make different tradeoffs.
• Speed: benchmark prefill and decoding separately at realistic batch sizes. Record hardware, dtype, software version and the selected kernel.
• Complexity: keep a simple reference implementation for comparisons. An optimization is useful only if its quality, maintenance cost and deployment requirements fit your project.

For the exercises here, implement standard scaled dot-product attention first, then compare one variant against it. Keep correctness, numerical error, memory and latency as separate measurements so you can explain what changed and why.