Efficient and Long Attention

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.

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 SoftmaxRun this section

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}$:

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:

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.

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 AttentionRun this section

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.

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$.

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:

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{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-NormRun this section

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)$:

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.

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.

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

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
• Milakov & Gimelshein (2018): Online normalizer calculation for softmax. arXiv:1805.02867
• Gemma Team, Google DeepMind (2024): Gemma 2: Improving Open Language Models at a Practical Size. arXiv:2408.00118
• Gemma Team, Google DeepMind (2025): Gemma 3 Technical Report. arXiv:2503.19786
• Beltagy, Peters & Cohan (2020): Longformer: The Long-Document Transformer (sliding-window attention). arXiv:2004.05150
• Xiao et al. (2023): Efficient Streaming Language Models with Attention Sinks (StreamingLLM). arXiv:2309.17453