Intuition

The 2017 Transformer is still the architecture every state-of-the-art LLM and ViT runs on today. *Broadly speaking* — but "broadly speaking" hides a lot of work. Almost every component has been quietly replaced between 2017 and 2025. Not torn down, refined. Same skeleton, modernised organs. The slide pack is named *"ViT-5: Vision Transformers for The Mid-2020s"* — and the framing is exactly that. Nine specific upgrades, grouped into four themes: stability (Pre-Norm, LayerScale), normalisation (RMSNorm, QK-Norm), efficiency (Flash Attention, KV-cache, GQA), positional encoding (RoPE), and attention hygiene (Registers). Knowing what each one does and why is the entire exam.

Explanation

Why does any of this matter? When you train a 100-billion-parameter Transformer on a trillion tokens, every one of these is the difference between *"it diverges after 50k steps"* and *"it converges to GPT-4."* These are not nice-to-haves — they are the engineering reality of modern Transformers.

Change #1 — Post-Norm → Pre-Norm. *2017 original:* $x_{\text{out}}=\text{LN}(x+\text{Sublayer}(x))$ — LayerNorm *after* the residual addition. *Modern:* $x_{\text{out}}=x+\text{Sublayer}(\text{LN}(x))$ — LayerNorm *before* the sublayer, inside the residual branch. Why this matters: in Pre-Norm, the residual stream $x$ flows directly from input to output, *untouched by any normalisation*. This direct path is the residual stream — the highway gradients use to flow back to early layers. Post-Norm squashes every residual through a LayerNorm before passing on; gradients have to push through a normaliser at every layer, and very deep Transformers (24+ layers) become unstable to train without careful warmup. Exam line: *Pre-Norm improves training stability by giving gradients a direct, unnormalised path through the residual stream.* Every modern LLM (GPT, LLaMA, PaLM, Gemma) is Pre-Norm.

Change #2 — LayerNorm → RMSNorm. *LayerNorm* centres then scales: $\mu =\frac{1}{D}\sum x_i$; ${\sigma }^2=\frac{1}{D}\sum (x_i-\mu {)}^2$; $y_i=(x_i-\mu )/\sqrt{{\sigma }^2+\epsilon }$; output $={\gamma }_i{y}_i+{\beta }_i$. Two learnable params per dim. *RMSNorm* drops the centring step: $\text{rms}=\sqrt{\frac{1}{D}\sum x_i^2}$; $y_i=x_i/(\text{rms}+\epsilon )$; output $={\gamma }_i{y}_i$. **No mean subtraction, no bias $\beta$.** Slightly faster (fewer ops) and empirically performs as well as or better than LayerNorm. LLaMA, Gemma, and many recent LLMs use RMSNorm.

The geometric quiz — the slide-7 exam target. *Picture a 2D scatter of points $(x_1,x_2)$.* LayerNorm: centring ($x-\mu$) shifts every point so the projection onto the all-ones vector $(1,1)$ is zero — every point lands on the line $x_1+x_2=0$ (the anti-diagonal). Scaling ($÷\sigma$) puts each point at unit distance from the origin on that line. Result: all points lie on two specific symmetric points $(+a,-a)$ and $(-a,+a)$. RMSNorm: no centring. Scaling sends every point to the full unit circle (radius $\sqrt{D}$ in $D$ dims).

The geometric quiz answer in one box. *LayerNorm:* points → small circle on the $(D-1)$-dim zero-mean hyperplane (1 magnitude DOF + 1 mean DOF removed = 2 constraints). *RMSNorm:* points → full $D$-dim sphere of radius $\sqrt{D}$ (only 1 magnitude constraint). Memorise: *"LayerNorm collapses onto the anti-diagonal line; RMSNorm spreads around the full unit circle."* Dropping the mean doesn't hurt because the LLM has plenty of other ways to centre activations (bias terms in linear layers, learnable $\gamma$).

Change #3 — LayerScale (Touvron et al., CaiT, 2021). In very deep Transformers, the residual contribution from each sublayer can either explode or vanish. LayerScale inserts a per-channel learnable diagonal scaling into the residual branch: $x_{\text{out}}=x+\text{diag}(\lambda )\cdot \text{Sublayer}(\text{LN}(x))$ where $\lambda \in {\mathbb{R}}^D$ is initialised to a tiny value (${\lambda }_{\text{init}}\approx 10^{-4}$). At the start of training, every sublayer contributes almost nothing — the network behaves like a stack of identity functions, and gradients flow trivially. As training proceeds, useful directions in $\lambda$ grow. Exam line: *LayerScale gives each sublayer a learnable "volume knob" per channel, initialised near zero so training begins from a near-identity network.*

Change #4 — QK-Norm (Henry et al., 2020). Attention has a hidden numerical hazard: $Q{K}^{\top }$ can grow large in magnitude (especially at long sequences or with large head dims), pushing softmax into saturation. The original paper's fix — divide by $\sqrt{d_k}$ — is *approximately correct but not enough at scale*. QK-Norm applies LayerNorm (or RMSNorm) to $Q$ and $K$ independently before the dot product: $Q^{\prime }=\text{LN}(Q)$, $K^{\prime }=\text{LN}(K)$, then $\text{softmax}(Q^{\prime }{K}^{\prime \top }/\sqrt{d_k})V$. Result: $Q^{\prime },K^{\prime }$ are unit-norm along their head dim, dot products bounded by $\pm 1$ per pair → softmax sees a well-conditioned input — no extreme values, no saturation, no exploding gradients.

Change #5 — Registers (Darcet et al., 2023: "Vision Transformers Need Registers"). Researchers training large ViTs noticed strange high-norm activations at uninformative patches — chunks of sky, blurred backgrounds, low-content regions. The model was using these patches as scratch space, dumping unrelated global information there. With no dedicated place to put global state, the network hijacks the least informative tokens and corrupts them. Attention maps then look chaotic.

The register fix. Append $k$ extra trainable tokens (typically 4–8) to the input sequence: $[\text{CLS}|{\text{patch}}_1|\dots |{\text{patch}}_N|{\text{reg}}_1|\dots |{\text{reg}}_k]$. These have no associated patch — just learnable vectors *appended at the input and discarded at the output*. The network now has dedicated scratch space; high-norm activations cluster on the register tokens rather than corrupting real patches. Attention maps become clean; dense-prediction downstream tasks improve. Slide line verbatim: *Registers = "garbage collector" tokens; they prevent attention peaks at blank areas by giving the model dedicated scratch space.*

Change #6 — Flash Attention (Dao et al., NeurIPS 2022). The single most consequential efficiency improvement to Transformers in the last 5 years. An EXACT attention algorithm — same outputs as standard attention, no approximation — that is dramatically faster and uses far less memory.

The attention bottleneck Flash Attention solves. Standard attention computes $\text{softmax}(Q{K}^{\top }/\sqrt{d_k})V$. The intermediate $Q{K}^{\top }$ has shape $N\times N$. For $N=8192$, this is a 64M-entry matrix. The bottleneck on modern GPUs isn't compute — it's moving the matrix between HBM (large, slow) and SRAM (small, ~100× faster). Standard attention writes the full $N\times N$ matrix to HBM, reads it back for softmax, writes again, reads for the value multiplication. Lots of slow memory traffic.

The Flash Attention trick — tile and never materialise. Process $Q,K,V$ in small blocks that fit in SRAM; compute softmax block-by-block using online streaming; never write the full $N\times N$ matrix to HBM. The mathematical trick is the *online softmax algorithm*: keep a running max $m$ and running denominator $\ell$, rescale the partial output when the max changes. Loop: load tile of $K,V$ → compute partial scores → update $(m,\ell ,O)$ → discard tile. **Memory cost drops from $\mathcal{O}(N^2)$ to $\mathcal{O}(N)$. Wall-clock 2–4× faster, sometimes more. Two exam emphases:** *(1) exact* — bitwise close to standard attention; *(2) IO-aware* — the speedup comes from minimising HBM ↔ SRAM data movement, not from doing less computation.

Change #7 — RoPE (Rotary Position Embedding, Su et al., 2021). For token at position $m$, embedding dim $d$, frequencies ${\theta }_i=1/10000^{2i/d}$: rotation matrix on each pair of dimensions: $R_m=\left(\begin{array}{cc}\cos (m{\theta }_i)& -\sin (m{\theta }_i)\\ \sin (m{\theta }_i)& \cos (m{\theta }_i)\end{array}\right)$. Rotated query/key: ${\tilde{q}}_m=R_m{q}_m$, ${\tilde{k}}_n=R_n{k}_n$. Critical property: ${\tilde{q}}_m^{\top }{\tilde{k}}_n=f(m-n)$ — the attention score depends only on *relative* position. Three advantages over learned absolute PEs: extrapolates to longer sequences than training (no fixed-size lookup table); encodes relative position directly (depends only on $m-n$); multiplicative, not additive (rotation inside attention, not added to input embedding). Variants: 2D-RoPE for images, M-RoPE for video.

Change #8 — KV Caching. A *pure-inference* optimisation. When generating autoregressively, naïvely you re-compute attention over the entire history at every step — $\mathcal{O}(N^2)$ work per token, $\mathcal{O}(N^3)$ for length-$N$ generation. **Observation: the $K$ and $V$ projections for past tokens never change** — once you compute $K_1,V_1$, they're fixed. So *store them in a cache*; at each new step: compute $Q,K,V$ only for the new token, append the new $K,V$ to the cache, compute attention of new $Q$ against all cached $K$s and $V$s. Per-token work: $\mathcal{O}(N^2)\to \mathcal{O}(N)$. Total generation: $\mathcal{O}(N^3)\to \mathcal{O}(N^2)$. This is the reason chatbot streaming is fast. *Trade-off:* memory — the cache grows linearly with sequence length, per layer, per head, per batch element. For long contexts and large models, KV-cache memory dominates.

Change #9 — Grouped-Query Attention (GQA). In MHA you have $H$ query heads, $H$ key heads, $H$ value heads → KV cache is proportional to $H$. For a 70B-param LLM with $H=64$, the KV cache becomes enormous at long contexts. Two extremes: *MHA* — $H$ separate $K,V$ heads (max capacity, max cache). *MQA (Multi-Query Attention)* — all query heads share one $K$ and one $V$ head (massive savings, factor of $H$, but quality drops). GQA — the middle ground: divide the $H$ query heads into $G$ groups; each group shares one $K$ head and one $V$ head. $H$ query heads, $G$ key heads, $G$ value heads with $G<H$. **KV cache is $H/G$ times smaller** than full MHA. Quality nearly as good as MHA. LLaMA-2, Gemma, Mistral, and most recent open LLMs use GQA.

The GQA trade-off in one line. *MHA:* max quality, max KV cache. *MQA:* min KV cache, some quality loss. *GQA:* tunable middle ground via $G$. Typical: $H=64$, $G=8$ → $8\times$ smaller cache.

The modern Transformer block — putting it all together. Self-attention sublayer: $h=\text{RMSNorm}(x)$ (Pre-Norm + RMSNorm) → $Q,K,V$ projections (GQA: K, V have fewer heads) → $Q=\text{LN}(Q),K=\text{LN}(K)$ (QK-Norm) → apply RoPE to $Q,K$ → append to KV-cache → Flash Attention → $x=x+\text{diag}({\lambda }_1)\cdot \text{attn}$ (LayerScale). MLP sublayer: same RMSNorm + LayerScale wrap around SwiGLU. For ViTs: add register tokens to the input, discard at output. Use 2D-RoPE for image patches. Every single line is different from the 2017 original.

Definitions

• Residual stream — The unbroken identity path in a Pre-Norm Transformer; gradients flow directly through it from output back to input.
• Pre-Norm — $x_{\text{out}}=x+\text{Sublayer}(\text{LN}(x))$. LayerNorm placed before the sublayer; residual stream is never normalised. Modern default.
• Post-Norm — $x_{\text{out}}=\text{LN}(x+\text{Sublayer}(x))$. 2017 original; needs careful warmup at depth.
• RMSNorm — Drop-the-mean variant of LayerNorm: $y=\gamma \cdot x/\sqrt{\frac{1}{D}\sum x_i^2}$. No mean subtraction, no bias $\beta$. Cheaper and as effective.
• LayerScale — Per-channel learnable diagonal $\lambda$ on each sublayer's residual contribution, initialised $\sim 10^{-4}$. Makes deep nets train as near-identity at init.
• QK-Norm — Apply LayerNorm (or RMSNorm) to $Q$ and $K$ separately before the attention dot product. Prevents softmax saturation at long sequences / large head dims.
• Registers — Extra learnable tokens prepended to the input with no positional encoding; act as global scratchpad. Without them, the model corrupts uninformative patches. Discarded at output.
• Flash Attention — Exact, IO-aware attention algorithm. Tiles $Q,K,V$ into SRAM-sized blocks; computes streaming online softmax; never materialises the $N\times N$ matrix in HBM. Memory $\mathcal{O}(N^2)\to \mathcal{O}(N)$. 2–4× faster.
• Online softmax — Stream the softmax: keep running max $m$, denominator $\ell$, output $O$; rescale $O,\ell$ when $m$ grows. Mathematically equivalent to standard softmax, computable tile-by-tile.
• RoPE (Rotary Position Embedding) — Rotate pairs of dimensions of $Q$ and $K$ by angle $m{\theta }_i$ proportional to position. Attention dot product depends only on relative position $m-n$.
• KV-cache — Cache of $K,V$ for past tokens during autoregressive generation. Per-step compute $\mathcal{O}(t^2)\to \mathcal{O}(t)$; total generation $\mathcal{O}(N^3)\to \mathcal{O}(N^2)$. Memory grows linearly.
• MHA / GQA / MQA — Multi-Head Attention: $H$ separate $K,V$ heads. Grouped-Query: $G<H$ shared groups. Multi-Query: $G=1$ (one $K,V$ for all heads). KV cache ratio: $H/G$.

Formulas

• $$\text{Pre-Norm:}{x}_{\text{out}}=x+\text{Sublayer}(\text{LN}(x))\text{(modern)}$$
• $$\text{Post-Norm:}{x}_{\text{out}}=\text{LN}(x+\text{Sublayer}(x))\text{(2017 original)}$$
• $$\text{LayerNorm:}{y}_i={\gamma }_i\cdot \frac{x_i-\mu }{\sqrt{{\sigma }^2+\epsilon }}+{\beta }_i$$
• $$\text{RMSNorm:}{y}_i={\gamma }_i\cdot \frac{x_i}{\sqrt{\frac{1}{D}{\sum }_j{x}_j^2+\epsilon }}(\text{no mean, no bias})$$
• $$\text{LayerScale:}{x}_{\text{out}}=x+\text{diag}(\lambda )\cdot \text{Sublayer}(\text{LN}(x)),{\lambda }_{\text{init}}\approx 10^{-4}$$
• $$\text{QK-Norm:}\text{softmax}\left(\frac{\text{LN}(Q)\cdot \text{LN}(K{)}^{\top }}{\sqrt{d_k}}\right)V$$
• $$\text{RoPE rotation:}{R}_m=\left(\begin{array}{cc}\cos (m{\theta }_i)& -\sin (m{\theta }_i)\\ \sin (m{\theta }_i)& \cos (m{\theta }_i)\end{array}\right),{\theta }_i=1/10000^{2i/d}$$
• $$\text{RoPE relative property:}{\tilde{q}}_m^{\top }{\tilde{k}}_n=f(m-n)$$
• $$\text{Flash Attention memory:}\mathcal{O}(N^2)\to \mathcal{O}(N)$$
• $$\text{GQA KV-cache reduction:}\text{factor}=H/G$$

Derivations

Why Pre-Norm is more stable. Under Pre-Norm, $x_{\text{out}}=x+\text{Sublayer}(\text{LN}(x))$. The residual $x$ is *never* normalised — it passes through untouched. Variance of the residual stream grows additively across layers (bounded growth). Under Post-Norm, $x_{\text{out}}=\text{LN}(x+\text{Sublayer}(x))$. The combined signal is normalised every layer; for very deep stacks, the LayerNorm interacts with the residual in ways that produce ill-conditioned Jacobians — gradients can explode or vanish at depth. Pre-Norm trains 100+-layer Transformers without warmup; Post-Norm needs careful LR warmup even at 24 layers.

The geometric LN-vs-RMSNorm derivation. LayerNorm in ${\mathbb{R}}^D$: $x\mapsto x-\bar{x}1$ projects onto the hyperplane $H_0=\{v:\sum v_i=0\}$ (dimension $D-1$); $÷\sigma$ normalises onto the unit sphere within $H_0$. So output $\in S^{D-2}\subset H_0$ — a $(D-2)$-sphere embedded in a $(D-1)$-hyperplane. RMSNorm: skip the projection; $÷\text{rms}$ scales so $\Vert x\Vert =\sqrt{D}$. Output $\in S^{D-1}$ (the full sphere in ${\mathbb{R}}^D$). *Two DOFs removed by LayerNorm vs one by RMSNorm.*

RoPE preserves relative position. After rotating $q$ by angle $m\theta$ and $k$ by angle $n\theta$: ${\tilde{q}}^{\top }\tilde{k}=q^{\top }{R}_m^{\top }{R}_{n}k=q^{\top }{R}_{n-m}k$. So the dot product depends only on the difference $n-m$, not absolute positions. This is why RoPE generalises to longer sequences than seen at training — absolute position never appears.

The online softmax trick (Flash Attention's mathematical core). Want $\text{softmax}(s{)}_j=e^{s_j}/{\sum }_k{e}^{s_k}$ without storing the full $s$. Process $s$ in tiles. Keep running max $m^{(t)}$ and running sum ${\ell }^{(t)}={\sum }_{k\le t}{e}^{s_k-m^{(t)}}$ and running output $O^{(t)}$. When a new tile produces a larger max $m^{(t+1)}>m^{(t)}$: rescale $\ell$ and $O$ by $e^{m^{(t)}-m^{(t+1)}}$. Add the new tile's contributions. The final output equals the standard-softmax output, byte-for-byte (up to floating-point reduction order).

KV-cache asymptotics. Without cache: at step $t$, compute attention over all $t$ tokens → $\mathcal{O}(t^2)$ per step (full-attention recompute) → ${\sum }_{t=1}^N{t}^2=\mathcal{O}(N^3)$ total generation. With cache: at step $t$, compute $Q$ for new token + attention of 1 query against $t$ cached keys → $\mathcal{O}(t)$ per step → ${\sum }_{t=1}^{N}t=\mathcal{O}(N^2)$ total. *Per-token: $\mathcal{O}(N)\to \mathcal{O}(1)$; total: cubic to quadratic.*

GQA cache scaling. KV-cache size $\propto$ #KV heads $\times$ sequence length. MHA: $H$ KV heads. GQA: $G<H$ KV heads. Ratio: $H/G$. For LLaMA-2-70B with $H=64,G=8$: $8\times$ smaller cache. At long contexts (32k+), this is the difference between *fits on one GPU* and *needs model parallelism*.

Examples

• LLaMA 3 stack. Pre-Norm + RMSNorm + RoPE + GQA + Flash Attention + SwiGLU MLP. No LayerScale (decoder-only LLMs are not as deep as some ViT variants).
• Geometric quiz in 2D. Points $(2,4),(3,5),(-1,1)$. LayerNorm: $\mu =3,\sigma =1$ for first → $(-1,1)/1=(-1,1)$. Second: $\mu =4,\sigma =1\to (-1,1)$. Third: $\mu =0,\sigma =1\to (-1,1)$. **All three points map to the same $(-1,1)$** on the anti-diagonal! (Modulo the learnable $\gamma ,\beta$, which can recover individuality.) RMSNorm: scale to RMS 1 — each point goes to a different location on the unit circle.
• GQA in LLaMA-2-70B. $H=64$ query heads grouped into $G=8$ KV groups → KV cache shrinks $8\times$ with negligible quality loss.
• Flash Attention v2 memory. At sequence length 32k, FP16 attention: standard requires $32000^2\cdot 2\approx 2$ GB just for the attention matrix per layer. Flash Attention requires $\mathcal{O}(N)\cdot d\cdot 2\approx 8$ MB per layer. >100× memory savings.
• LayerScale init. ${\lambda }_{\text{init}}=10^{-4}$ means the sublayer initially contributes $10^{-4}$ of its output to the residual — effectively zero. After training, $\lambda$ grows to $\mathcal{O}(0.1\text{–}1)$ on useful channels.
• Registers visual. Pre-register: heatmap of patch attention norms shows 5–10 spikes on patches of blank sky in a ViT-L. Post-register: spikes are absorbed by the 4 register tokens; patch attention is uniform on the object.
• RoPE wavelengths. For $d=64$, ${\theta }_0=1$ (period $2\pi$), ${\theta }_{31}=1/10000$ (period $20,000\pi$). Different pairs encode position at different frequencies — fine-to-coarse positional decomposition.

Diagrams

• Pre-Norm vs Post-Norm block. Two side-by-side block diagrams. Pre-Norm: residual line shown as a straight unbroken arrow from input to output; sublayer + LN forms a side branch. Post-Norm: residual goes into a final LN block; the LN is *on* the residual path.
• LayerNorm vs RMSNorm geometric. 2D scatter: LayerNorm collapses points onto the anti-diagonal $x_1+x_2=0$ at unit distance from origin (2 points: $(+a,-a)$ and $(-a,+a)$). RMSNorm: same points spread around the full unit circle.
• LayerScale architecture. A standard residual block with a diag(λ) element between the sublayer and the residual add; annotation: ${\lambda }_{\text{init}}\approx 10^{-4}$.
• QK-Norm placement. Q and K projection → LN block → dot product → softmax → ×V.
• Registers in a ViT sequence. Input: $[\text{CLS}|{\text{patch}}_1\dots {\text{patch}}_N|{\text{reg}}_1\dots {\text{reg}}_4]$; output: take CLS for classification; registers discarded.
• Flash Attention tiling. Show the $Q,K,V$ matrices partitioned into blocks; each iteration loads one $K,V$ tile into SRAM, computes a partial output, updates running stats, moves on. Annotate: *full $N\times N$ matrix never lives in HBM.*
• RoPE rotation. A pair of embedding dimensions $(q_{2i},q_{2i+1})$ shown as a 2D vector; rotation by angle $m{\theta }_i$ shown as an arc.
• GQA visualisation. $H=8$ query heads → $G=4$ groups of 2 query heads; each group shares one $K$ head and one $V$ head. Tree diagram showing the sharing.
• KV-cache buildup. Three autoregressive steps: at $t=1$ cache has $[K_1,V_1]$; at $t=2$ cache has $[K_1,K_2;V_1,V_2]$; at $t=3$ adds $K_3,V_3$. Only $Q_t$ is computed fresh each step.

Edge cases

• Pre-Norm output magnitude can drift with depth (the unbroken residual accumulates); a final LayerNorm before the classification head is still required.
• Flash Attention requires GPU support — Ampere+ for v1, Hopper+ for v3. Older GPUs cannot run it.
• **GQA with $G=1$ (MQA limit)** loses quality on tasks requiring fine-grained discrimination. GQA with $G\in \{4,8\}$ is the empirical sweet spot.
• KV-cache memory dominates at long contexts. LLaMA-2-70B at 32k context: KV cache > 20 GB; the model is held by KV cache memory, not by weights or activations.
• LayerScale on shallow models is unnecessary — only useful for very deep nets where sublayer outputs can blow up the residual.
• **RoPE base frequency $10000$** can be insufficient for very long contexts; LLaMA-2 increased it; some long-context models use even larger bases.
• Register tokens have no PE — adding positional embedding to them defeats the point (they're position-free scratchpad).

Common mistakes

• Calling RMSNorm *"LayerNorm without bias"* — also drops the mean subtraction. Both ops are removed.
• Treating RoPE as *additive* — it's multiplicative (rotation matrix applied to $Q,K$).
• Conflating Flash Attention with KV-cache — different problems (memory layout vs incremental compute).
• Stating registers are positional — they have no positional encoding; that's the point.
• Saying Flash Attention is *approximate* — it's exact. Same outputs as standard attention.
• Writing the Pre-Norm equation as $\text{LN}(x+\text{Sublayer}(x))$ — that's Post-Norm. Pre-Norm puts LN on the sublayer's input.
• Treating GQA as a quality-degraded MQA — it's the sweet spot between MHA (max quality) and MQA (min cache).
• Confusing the *online softmax* trick with *online learning* — they're unrelated; "online" here means *streaming* (process one tile at a time).

Shortcuts

• Nine modern upgrades: Pre-Norm, RMSNorm, LayerScale, QK-Norm, Registers, Flash Attention, RoPE, KV-cache, GQA.
• Pre-Norm = unbroken residual; Post-Norm normalises the combined signal.
• RMSNorm = LayerNorm minus mean minus bias. Same final shape, cheaper.
• Geometric quiz answer: LayerNorm → small circle on anti-diagonal; RMSNorm → full sphere.
• RoPE is multiplicative (rotation). Old PE is additive. RoPE attention depends only on relative position.
• Flash Attention is exact + IO-aware. Speed comes from avoiding HBM↔SRAM traffic, not from doing less compute.
• **KV-cache: $\mathcal{O}(N^3)\to \mathcal{O}(N^2)$** generation; memory grows linearly.
• **GQA reduces KV cache by $H/G$** (typical 8×). LLaMA-2/3, Gemma, Mistral all use GQA.

Proofs / Algorithms

RoPE encodes relative position. ${\tilde{q}}_m=R_{m}q$, ${\tilde{k}}_n=R_{n}k$. Then ${\tilde{q}}_m^{\top }{\tilde{k}}_n=q^{\top }{R}_m^{\top }{R}_{n}k=q^{\top }{R}_{n-m}k$ because $R_m^{\top }=R_{-m}$ and rotation matrices compose by adding angles. So the dot product depends only on the *difference* $n-m$, hence on *relative* position. Absolute positions never appear in the score — RoPE generalises beyond training-time lengths.

Online softmax correctness. Given $s=(s_1,\dots ,s_N)$, define $m=\max s$, $\ell =\sum e^{s_k-m}$, $\text{softmax}(s{)}_j=e^{s_j-m}/\ell$. Processing in tiles: maintain $(m^{(t)},{\ell }^{(t)})$ after seeing the first $t$ elements. When element $s_{t+1}$ arrives: $m^{(t+1)}=\max (m^{(t)},s_{t+1})$; ${\ell }^{(t+1)}={\ell }^{(t)}{e}^{m^{(t)}-m^{(t+1)}}+e^{s_{t+1}-m^{(t+1)}}$. The rescaling factor $e^{m^{(t)}-m^{(t+1)}}$ corrects the previously accumulated terms for the new max. After processing all $N$ elements, ${\ell }^{(N)}=\ell$ and the output equals the standard softmax bytewise.

GQA KV-cache reduction. Cache size $=L\times H_{KV}\times N\times d_k\times 2$ ($L$ layers, $H_{KV}$ KV heads, $N$ tokens, $d_k$ per-head dim, factor 2 for $K$ and $V$). MHA: $H_{KV}=H$. GQA: $H_{KV}=G$. Ratio: $H/G$. For $H=64,G=8$: 8× smaller cache, all other factors equal.