# Module 16: Efficient and Long Attention

> FlashAttention and the GPU memory hierarchy, local and global attention layers, soft-capping, sandwich normalization and attention sinks.

- Part IV: The Modern Architecture · elective · Advanced · 4 hours
- By the end you can explain how attention is made cheaper without changing its result, and how long contexts are kept stable.
- Builds on: [Module 15: The LLaMA Block](https://transformers-and-math.drmhse.com/md/15-the-llama-block.md)
- Web page (interactive figures, labs and equation explainers): https://transformers-and-math.drmhse.com/modules/16-efficient-and-long-attention
- From *Transformers & Math* by Mike Chumba: https://transformers-and-math.drmhse.com

Attention's cost grows with the square of the sequence length (section 6.6), and its KV cache with every generated token (section 12.2). This module collects the techniques that make long contexts affordable and stable: FlashAttention, which computes exactly the same attention while moving far less data, and small architecture changes, as in Google's Gemma family, that attend locally in most layers, bound the size of numbers, and keep attention sinks.

## 16.1 Memory, Not Arithmetic: The GPU Hierarchy

A GPU has large, comparatively slow **HBM** (high-bandwidth memory, the "80 GB" on the box) and small, fast on-chip **SRAM**, about 20 MB on an NVIDIA A100. Arithmetic happens only on data brought into SRAM.

Dao et al. (2022) pointed out that standard attention is **memory-bound**: it spends most of its time moving the $N \times N$ score matrix, writing it to HBM, reading it back for the softmax, writing the result, and reading it again to multiply by $V$. FlashAttention computes the same result without ever storing the full matrix.

```mermaid
graph TD
    HBM["Slow HBM (A100: 80 GB, about 1.5–2 TB/s)"] <--> Tiles["Load one block of Q and one block of K, V into fast SRAM (about 20 MB, about 19 TB/s)"]
    Tiles --> OnlineSoftmax["Compute that block's scores and update a running softmax; the full N × N matrix is never stored"]
    OnlineSoftmax --> Accumulate["Add the block's contribution to the output O, kept in SRAM"]
    Accumulate --> WriteBack["Write the finished O to HBM once"]
```

> **In words.** Work through the score matrix in blocks that fit in fast on-chip memory, keeping a running maximum, sum and output per row. The full matrix is never stored.

## 16.2 FlashAttention and the Online Softmax

The softmax seems to need the whole row at once: its sum, and its maximum for safety (section 4.3). The online softmax keeps a running maximum and sum instead, correcting them as each block arrives.

Split a row into blocks $\mathbf{x}^{(1)}$ and $\mathbf{x}^{(2)}$, each with its maximum $m_i$ and sum $\ell_i = \sum_j e^{x^{(i)}_j - m_i}$. With overall maximum $m = \max(m_1, m_2)$, multiplying by $e^{m_i - m}$ rescales each sum to the common maximum, since $e^{x - m_i} \cdot e^{m_i - m} = e^{x - m}$:

$$
\ell = e^{m_1 - m} \ell_1 + e^{m_2 - m} \ell_2
$$

The outputs combine the same way. With $O_i$ the output from block $i$ alone, the combined output weights each by its share of the total sum:

$$
O_{\text{new}} = \operatorname{diag}\left(\frac{e^{m_1 - m} \ell_1}{\ell}\right) O_1 + \operatorname{diag}\left(\frac{e^{m_2 - m} \ell_2}{\ell}\right) O_2
$$

Many queries are processed at once, each row with its own $m$ and $\ell$; $\operatorname{diag}$ gives each row its own scale. More blocks repeat the same update against the running result.

> **Interactive figure** `flash-tiling`, on the web page: https://transformers-and-math.drmhse.com/modules/16-efficient-and-long-attention#figure-flash-tiling
>
> **Try:** Follow one row and load the tiles one at a time. Before the last tile, predict whether the running output can already equal the exact answer, and what happens to the running sum whenever a tile brings a new, larger maximum.
>
> **What happens:** Not until the last tile, since every key contributes to the exact weighted average. When a new maximum arrives, the old running sum and output are multiplied by $e^{m_{\text{old}} - m_{\text{new}}}$, which shrinks them to the new scale; that is the correction in the formulas above. After the last tile the running output matches the exact attention output to about seven decimal places.

> **In words.** A sixteen-by-sixteen score matrix in vertical tiles of keys. Each click loads a tile and updates one query row's running maximum, sum and output, which match the exact attention output once the last tile is in.

The result is exactly standard attention, up to rounding. Extra memory drops from $\mathcal{O}(N^2)$ to $\mathcal{O}(N)$, and HBM traffic falls several-fold, since the score matrix never goes there. The paper reports speedups often quoted as $3\times$ to $5\times$, and the smaller memory is much of what makes long contexts practical.

> **Measured** (lab 4c, `experiments/online_softmax.py`). The update is exact for every block size from $4$ to $64$ (largest difference $5.96 \times 10^{-7}$). It also matches this repository's own causal attention to $1.19 \times 10^{-7}$, including the causal mask. The mask needs care: the streaming version has to apply it using each token's position in the full sequence, not its position inside the current block.

> **Measured** (lab 4c). If the scores are small, dropping the rescaling (never subtracting any maximum) *also* gives the exact answer: the running ratio still computes the same expression. So the correction is about **stability**, not correctness. Multiply the scores by $40$ and the uncorrected version returns `inf` or `nan`, while the corrected one stays exact to $1.2 \times 10^{-7}$. This is section 4.3's lesson again: the algebra was never in doubt; the problem is the range of numbers a float can hold.

> **Not measured.** The streaming code here checks correctness, not speed. It does not test the $3$–$5\times$ wall-clock figure, which needs real tiling inside a single fused GPU kernel.

## 16.3 Local and Global Attention

Gemma is Google DeepMind's family of open models. The first Gemma (2024) had embedding scaling and a 256,000-token vocabulary; Gemma 2 (2024) added alternating attention (this section), sandwich normalization (16.4) and logit soft-capping (16.5), all on top of section 15.6's block.

### Alternating Local and Global Attention
Full attention's compute grows with $N^2$ and each layer's KV cache with $N$, both expensive at 8k to 128k tokens. Gemma 2 alternates two kinds of layer:
- **Local (sliding-window) attention**: each token attends only to the $W = 4096$ most recent tokens, so the layer's cache never exceeds $W$ entries.
- **Global attention** over the whole sequence.

```mermaid
graph TD
    Layer1["Layer 1: local window, W = 4096"] --> Layer2["Layer 2: global attention"]
    Layer2 --> Layer3["Layer 3: local window, W = 4096"]
    Layer3 --> Layer4["Layer 4: global attention"]
    Layer4 --> LayerN["... alternating to the top of the stack"]
```

> **In words.** Cheap local layers look at the last 4096 tokens; the global layers between them see the whole sequence, so information still travels from start to end.

The saving grows with context length and the share of local layers. With Gemma 2's 8192-token context, a local layer holds at most half a global layer's cache. Gemma 3 uses five local layers per global one with $W = 1024$; at 128k tokens that keeps about 17% of an all-global cache: $(5 \cdot 1024 + 131072) / (6 \cdot 131072) \approx 0.173$.

> **Interactive figure** `local-global`, on the web page: https://transformers-and-math.drmhse.com/modules/16-efficient-and-long-attention#figure-local-global
>
> **Try:** A token in a stack of local layers with window $W$ can see $W$ positions back in one layer. Before stacking layers, predict how far back the last token can reach after three local layers, and after one global layer.
>
> **What happens:** About $3W$ after three local layers, since each layer lets information hop one more window back. A single global layer reaches the very first token directly. Mixing the two keeps most layers cheap while still connecting the whole context.

> **In words.** On a scaled-down sequence, which tokens each query sees in local and global layers, and how far back the last token reaches as layers stack: one more window per local layer, the very start after the first global layer. The readout gives the cache relative to an all-global stack.

## 16.4 Sandwich Normalization and Embedding Scaling

### Sandwich Normalization: RMSNorm Before and After
Gemma 2 keeps Pre-LN (section 8.5) and adds a second RMSNorm on the sub-layer's output, before the addition, sandwiching the sub-layer:

$$
\mathbf{x}_{\text{mid}} = \operatorname{SubLayer}(\operatorname{RMSNorm}_{\text{pre}}(\mathbf{x}_l))
$$

$$
\mathbf{x}_{l+1} = \mathbf{x}_l + \operatorname{RMSNorm}_{\text{post}}(\mathbf{x}_{\text{mid}})
$$

The first norm steadies the sub-layer's input; the second limits what it adds, so no layer can push a huge update into the residual stream. The skip path $\mathbf{x}_l +$ is still a plain addition, keeping section 8.5's gradient highway intact.

### Scaling Embeddings by $\sqrt{d_{\text{model}}}$
GPT-2 and LLaMA feed the embedding row straight into the first layer. Gemma, like the 2017 Transformer, first multiplies it by the square root of the width:

$$
\mathbf{x}_0 = \left(\mathbf{e}_t^\top W_E\right) \times \sqrt{d_{\text{model}}}
$$

($\mathbf{e}_t^\top W_E$ is row $t$ of $W_E$, section 1.5.) Embedding entries start with standard deviation about $1/\sqrt{d_{\text{model}}}$; the scale brings them to about 1, matching what the layers add. Otherwise the token's own identity would be small beside the first layer's update and hard for later layers to recover.

## 16.5 Logit Soft-Capping and QK-Norm

Both the attention scores and the output logits can grow very large in training, causing loss spikes and, at sampling time, overconfidence. Gemma 2 bounds every logit within $\pm\kappa$ (kappa) using $\tanh$, which maps any number into $(-1, 1)$:

$$
\tilde{z} = \kappa \tanh\left(\frac{z}{\kappa}\right)
$$

Gemma 2 uses two caps:
- $\kappa = 50.0$ for the attention scores, $z = QK^\top / \sqrt{d_k}$.
- $\kappa = 30.0$ for the output logits over the vocabulary.

```mermaid
graph LR
    Logit["Raw logit z, any size"] --> Div["Divide by the cap κ, e.g. 30"]
    Div --> Tanh["tanh squashes it into -1 to 1"]
    Tanh --> Mult["Multiply by κ again"]
    Mult --> Capped["Capped logit, between -30 and +30"]
```

> **In words.** Divide by the cap, squash with tanh, multiply by the cap. Small logits pass almost unchanged; large ones bend smoothly toward the cap instead of being cut off.

> **Interactive figure** `softcap-curve`, on the web page: https://transformers-and-math.drmhse.com/modules/16-efficient-and-long-attention#figure-softcap-curve
>
> **Try:** With a cap of 30, predict what a logit of 30 comes out as, and whether a logit of 300 can still be told apart from a logit of 100 after capping.
>
> **What happens:** A logit of 30 becomes about 22.8, and both 100 and 300 come out within a hair of 30: the cap squashes very large values together. Unlike a hard clip, though, the curve never goes completely flat, so the gradient never quite reaches zero.

> **In words.** The soft cap plotted against the raw logit, beside a hard clip and the identity line. A slider sets the cap, and a marker follows a chosen logit to its capped value and slope.

### What the Cap Does to Small and Large Logits
- **Small logits pass through**, since $\tanh(u) \approx u$ near zero: with $\kappa = 30$, 5 becomes 4.95.
- **Large logits approach the cap**: 30 becomes 22.8 and 100 becomes 29.9.
- **The gradient never reaches zero.** The slope $1 - \tanh^2(z/\kappa)$ is always positive, whereas a hard clip has zero slope beyond the cap, so a clipped logit stops learning.

Gemma 3 replaced soft-capping with QK-norm: RMSNorm on the queries and keys before their dot product, which bounds the scores another way.

## 16.6 Attention Sinks and Streaming

To stream endless text, a natural idea is to keep only the most recent tokens in the KV cache. But as soon as the first token is evicted, perplexity shoots up, even when it is just the start marker `<s>`. The cause is the **attention sinks** of section 14.6 (Xiao et al., 2023): the spare attention parked there gets spread over other tokens and disrupts the layer. The fix is to always keep the first $4$ tokens alongside the sliding window; with that, the paper's models kept a stable perplexity over very long streams.

## Exercises

**Exercise (check).** Does FlashAttention change the result of attention? What does it change?

*Solution:* No: it computes exactly the same output, up to floating-point rounding. It changes how much data moves between slow and fast GPU memory, by never storing the full score matrix.

**Exercise (apply).** Merge two softmax blocks. Block 1 has maximum $m_1 = 2$ and sum $\ell_1 = 3$ (computed with its own maximum); block 2 has $m_2 = 5$ and $\ell_2 = 1.5$. What is the combined sum $\ell$ relative to the overall maximum?

*Answer:* 1.6494 (within 0.001)

*Solution:* The overall maximum is $m = 5$. Then $\ell = e^{2 - 5} \cdot 3 + e^{5 - 5} \cdot 1.5 = 3e^{-3} + 1.5 \approx 0.149 + 1.5 = 1.649$.

**Exercise (apply).** Gemma 3 uses five local layers (window 1,024) for every global layer. What share of an all-global stack's KV cache does it need at a context of 32,768 tokens?

*Answer:* 0.1927 (within 0.002)

*Solution:* Per six layers: $5 \times 1{,}024 + 32{,}768 = 37{,}888$ cached positions, against $6 \times 32{,}768 = 196{,}608$ for six global layers. The share is $37{,}888 / 196{,}608 \approx 0.193$.

**Exercise (apply).** With soft-capping at $\kappa = 30$, what does a logit of 30 become?

*Answer:* 22.85 (within 0.02)

*Solution:* $30 \tanh(30/30) = 30 \tanh(1) \approx 30 \times 0.7616 = 22.85$.

**Exercise (extend).** FlashAttention does the same number of multiply-adds as standard attention, yet runs several times faster. Explain where the time goes in standard attention and what FlashAttention changes.

*Solution:* Standard attention is memory-bound: it writes the full $N \times N$ score matrix to slow HBM, reads it back for the softmax, writes the weights, and reads them again to multiply by $V$. For long sequences that traffic dominates. FlashAttention processes the scores in tiles that fit in fast on-chip SRAM, keeps a running maximum, sum and output per row (the online softmax), and never writes the score matrix to HBM at all, so it reads and writes far less data for the same arithmetic.

### Core Academic References
- **Dao et al. (2022); Dao (2023)**: *FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness*. NeurIPS 2022; FlashAttention-2, 2023. [arXiv:2205.14135](https://arxiv.org/abs/2205.14135)
- **Milakov & Gimelshein (2018)**: *Online normalizer calculation for softmax*. [arXiv:1805.02867](https://arxiv.org/abs/1805.02867)
- **Gemma Team, Google DeepMind (2024)**: *Gemma 2: Improving Open Language Models at a Practical Size*. [arXiv:2408.00118](https://arxiv.org/abs/2408.00118)
- **Gemma Team, Google DeepMind (2025)**: *Gemma 3 Technical Report*. [arXiv:2503.19786](https://arxiv.org/abs/2503.19786)
- **Beltagy, Peters & Cohan (2020)**: *Longformer: The Long-Document Transformer* (sliding-window attention). [arXiv:2004.05150](https://arxiv.org/abs/2004.05150)
- **Xiao et al. (2023)**: *Efficient Streaming Language Models with Attention Sinks* (StreamingLLM). [arXiv:2309.17453](https://arxiv.org/abs/2309.17453)