Modern Architecture Deltas

TL;DR — A 2026 LLM is the 2017 transformer with five substitutions. The norm moved before the sublayer (pre-norm) so gradients reach depth without a hand-tuned warmup; LayerNorm became RMSNorm because re-centering was doing no work and re-scaling was; the FFN became a 3-matrix SwiGLU gate at $8 d /3$ hidden width; multi-head attention became grouped-query attention so the KV cache shrank by the query-to-KV head ratio; and attention moved into an IO-aware fused kernel (FlashAttention) that is exact, not approximate — it removes memory traffic, not FLOPs. Each delta is a response to a specific thing that broke at scale, and each has a cost that shows up somewhere.

1. Why this exists

The 2017 transformer in The Transformer trains a 6-layer encoder–decoder on 2014-era translation data. Scale that recipe to 100+ layers and trillions of tokens and four things break, each in a different way:

Depth refuses to train. Post-LN transformers — the original design, with the norm applied after the residual add — diverge at depth unless you spend a long, fragile learning-rate warmup. Xiong et al. proved the cause: at initialization the expected gradients near the output layer grow with depth, so the early steps take far too large a step. That is a training-stability failure, and it is why essentially every modern model is pre-LN (Xiong et al. 2020, ICML).

Normalization is on the critical path of every layer. LayerNorm computes a mean and a variance, and at 100+ layers × every token × every step that reduction is measurable wall-clock. Zhang & Sennrich’s finding was that the re-centering half of LayerNorm is dispensable — only the re-scaling mattered — so half the reduction can be deleted for free (Zhang & Sennrich 2019, NeurIPS).

The FFN is the parameter budget and it is paying for a weaker function. The position-wise FFN holds roughly two-thirds of a transformer’s parameters. It spent them on act(xW_1)W_2; Shazeer showed that multiplying in a second, gate-controlled projection is strictly better at equal parameter count and equal FLOPs (Shazeer 2020).

Inference is memory-bandwidth-bound, not compute-bound. Once you serve the model, the throughput ceiling is how many bytes you must move per token. Two separate byte counts dominate: the KV cache (attention) and the attention score matrix (training and prefill). GQA attacks the first (Ainslie et al. 2023); FlashAttention attacks the second (Dao et al. 2022).

The pattern to internalise. None of these is “a better idea someone had”. Each is a named mechanism deleting a named cost. If you can recite the cost, the mechanism follows.

2. How it works

2.1 The residual stream, and why normalization position is a real choice

Everything below is easier if you hold one picture: the residual stream. The model is $L$ blocks, each of which adds to a running vector and never replaces it:

$x_{l + 1} = x_{l} + Sublayer (\cdot)$

Sublayers (attention, FFN) read from the stream and write back into it. Read capacity, not computation, is the right intuition: the stream is a shared workspace and each block appends a correction.

Normalization has to happen somewhere, and there are exactly two natural places. Post-norm (Vaswani 2017) normalizes the sum; pre-norm normalizes the input to the sublayer, leaving the sum alone:

$x_{l + 1} = LN (x_{l} + F_{l} (x_{l})) vs x_{l + 1} = x_{l} + F_{l} (LN (x_{l}))$

The difference is structural, not cosmetic. In post-norm, the identity path is broken at every layer: $x_{l}$ passes through an LN before it can reach $x_{l + 1}$ . LN divides by the standard deviation of its input, so $\partial x_{l + 1} / \partial x_{l}$ carries a factor from that LN’s Jacobian at every one of the $L$ layers. The product of $L$ such factors is what Xiong et al. show has the wrong behaviour at initialization: the expected gradient magnitude near the output grows with depth, so the step at step 0 is scaled wrong and the optimizer needs warmup to survive it. In pre-norm, $\partial x_{l + 1} / \partial x_{l}$ contains an exact identity term, so gradient reaches every layer without passing through a normalization Jacobian; Xiong et al. prove the consequence and remove warmup entirely for pre-LN with comparable results.

There is a second, forward-pass reason GPT-2 discovered independently. With the identity path intact, each block adds to the stream. If each block’s contribution has roughly unit variance and is roughly independent of the others, the stream’s variance after $L$ blocks is

$Var (x_{L}) \approx Var (x_{0}) + l = 0 \sum L - 1 Var (F_{l}) \approx L$

so the magnitude grows like $L$ : a 100-layer model ends up with outputs an order of magnitude larger than a 12-layer one, and the learning rate that suited one does not suit the other. GPT-2’s fix, still in modeling_gpt2.py, shrinks each residual branch’s output projection at initialization by $1/ 2 L$ — the 2 counts the two branches (attention, FFN) per block — so the accumulated variance stays $O (1)$ regardless of depth:

# ... scale the weights of residual layers at initialization by a factor of
# 1/√N where N is the # of residual layers.   -- GPT-2
std = self.config.initializer_range / math.sqrt(2 * self.config.n_layer)
init.normal_(module.c_proj.weight, mean=0.0, std=std)

Why this is an interview-grade distinction. “Pre-norm is more stable” is a fact. “Pre-norm keeps an identity term in $\partial x_{l + 1} / \partial x_{l}$ , so no normalization Jacobian enters the depth-accumulated product; post-norm puts an LN on the identity path at every layer” is the mechanism. The follow-up is always “then why does GPT-2 still scale the init?” — pre-norm fixes the backward path and does nothing about the forward variance accumulation.

2.2 RMSNorm: deleting the mean from LayerNorm

LayerNorm, written out, is

$LN (x) = γ ⊙ \frac{x - μ}{σ ^{2} + ϵ} + β, μ = \frac{1}{d} i = 1 \sum d x_{i}, σ^{2} = \frac{1}{d} i = 1 \sum d (x_{i} - μ)^{2}$

Symbol by symbol, with the justification:

$x \in R^{d}$ — one token’s activation in the residual stream. Normalization is per-token, not per-batch: a token’s normalization cannot depend on other examples in the batch, which is what makes the model usable at batch size 1 during decoding.
$μ$ — the mean over the $d$ features of that one token. Subtracting it re-centers the vector so its entries sum to zero.
$σ^{2}$ — the variance across the same $d$ features. Dividing by $σ^{2} + ϵ$ re-scales to unit standard deviation, making the sublayer’s input distribution independent of how large the residual stream happened to grow. This is the half that does the work.
$ϵ$ — a small constant ( $1 0^{- 5}$ or $1 0^{- 6}$ ), added inside the square root. Without it, a constant vector — or an all-zero position after masking — divides by zero and produces NaN. It costs a tiny scale bias; it buys numerical survival.
$γ \in R^{d}$ , $β \in R^{d}$ — learned per-feature gain and bias. Normalization discards the input’s scale and offset; $γ$ and $β$ let the network relearn which features should be large and where the zero point sits, so the layer is not locked into a rigid unit-variance representation.

RMSNorm asks: which of those terms is doing the work? Its claim is that re-centering is not, and that the re-scaling is:

$RMSNorm (x) = γ ⊙ \frac{x}{\frac{1}{d} \sum _{i = 1}^{d} x _{i}^{2} + ϵ}$

The denominator is the root mean square, $\frac{1}{d} \sum x_{i}^{2}$ , not the standard deviation $\frac{1}{d} \sum (x_{i} - μ)^{2}$ . They differ only by the mean: $RMS^{2} = σ^{2} + μ^{2}$ . So RMSNorm is invariant to re-scaling ( $RMSNorm (c x) = RMSNorm (x)$ , $c > 0$ ) but not to re-centering ( $RMSNorm (x + c 1) \neq = RMSNorm (x)$ ). That is the trade: one invariance given up for a cheaper reduction.
$β$ is gone. A bias followed by a rescale is partly redundant; Zhang & Sennrich found no benefit from it, and dropping it removes $d$ parameters per norm layer.
$γ$ stays, for the original reason: it restores the per-feature expressivity pure normalization removes. Removing $γ$ too would make the norm a fixed, parameterless transformation, and the network could no longer learn feature scales.

Why is that a win rather than mere simplification? A norm layer’s dominant cost is its reduction and the extra memory pass. LayerNorm needs mean and variance; RMSNorm needs one mean-of-squares. Zhang & Sennrich report 7%–64% runtime reduction at comparable quality. At 100+ norm layers that is real wall-clock, and it is largest exactly where the norm is a big fraction of a small layer’s work.

The $ϵ$ trap. In bf16, $\frac{1}{d} \sum x_{i}^{2}$ can underflow for small activations and overflow for large ones — a sum of squares over a wide dynamic range is precisely the operation that loses precision. Production kernels accumulate it in fp32 even when the activations are bf16, and place $ϵ$ so the RMS cannot swamp it. Getting this wrong produces loss spikes that look like a data problem.

Removing RMSNorm entirely does not give a small quality regression; it gives training divergence at depth, because the residual stream’s magnitude is then unconstrained. RMSNorm is not decoration.

2.3 SwiGLU: the FFN as a learned gate

The original FFN is a two-matrix bottleneck with a pointwise nonlinearity:

$FFN (x) = W_{2} act (W_{1} x)$

Symbols: $x \in R^{d}$ is one token’s stream vector; $W_{1} \in R^{d_{ff} \times d}$ expands to a wide hidden space ( $d_{ff} = 4 d$ in the 2017 paper); $act$ is elementwise; $W_{2} \in R^{d \times d_{ff}}$ projects back. The hidden width is where the parameters live — $2 d d_{ff} = 8 d^{2}$ of them at $d_{ff} = 4 d$ .

A gated linear unit replaces the single expansion with two, and multiplies them elementwise:

$SwiGLU (x) = W_{2} (gate Swish_{β} (x W_{1}) ⊙ value (x W_{3})), Swish_{β} (z) = z σ (β z)$

$x W_{1} \in R^{d_{ff}}$ — the gate pre-activation. It is the branch that decides, per hidden unit, how much of the value branch to let through.
$Swish_{β} (z) = z σ (β z)$ with $σ$ the logistic sigmoid and $β$ a learned or fixed slope (Shazeer sets $β = 1$ , which is exactly SiLU). This $β$ is the gate’s slope and is unrelated to the LayerNorm bias $β$ of §2.2; the two names coincide only because both are conventional. Why this function and not ReLU? Because it is smooth and non-monotone: it dips slightly negative for small negative $z$ and asymptotes to $z$ for large positive $z$ . The non-monotonicity is what lets the gate express “partially suppress this unit” instead of only “on” or “off”, and the smoothness means it has a usable gradient everywhere, including at the origin where ReLU’s gradient is discontinuous.
$x W_{3} \in R^{d_{ff}}$ — the value branch with no nonlinearity. It is linear so that the product’s nonlinearity comes entirely from the gate.
$⊙$ — elementwise product. This is the whole trick: the gate learns a multiplicative, input-dependent mask. A ReLU acts like a fixed mask on a single projection; a GLU lets the mask depend on the input through a different learned projection. That is strictly more expressive per parameter.
$W_{2} \in R^{d \times d_{ff}}$ — the down-projection back to model width. Remove it and the block writes a $d_{ff}$ -wide vector into the residual stream; the shapes no longer close.

What breaks if you remove the gate. Setting $Swish_{β} (x W_{1}) = 1$ reduces SwiGLU to $W_{2} (x W_{3})$ — a single linear layer, collapsing the FFN to one low-rank linear map. Removing the nonlinearity while keeping the product gives Shazeer’s Bilinear variant, $(x W) ⊙ (x V) W_{2}$ , a quadratic form that empirically underperforms the gated variants.

Where $8 d /3$ comes from. Three matrices instead of two means $3 d d_{ff}$ parameters. Holding the parameter count of the $4 d$ FFN constant: $3 d d_{ff} = 2 d (4 d) \Rightarrow d_{ff} = \frac{8 d}{3}$ . Shazeer states this explicitly — “to keep the number of parameters and the amount of computation constant, we reduce the number of hidden units $d_{ff}$ by a factor of $\frac{2}{3}$ ” — and implementations then round up to a hardware-friendly multiple (commonly 128, sometimes 256). That rounding is why real configs contain odd widths like 11008; at $d = 4096$ , $8 d /3 = 10922.67$ , which rounds to 11008.

MHA gives every one of the $h$ query heads its own key and value head. At training time that is what you want; at decode time it is a liability, because the KV cache is the thing you re-read on every generated token. The cache size is

$bytes = 2 \cdot L \cdot n_{k v} \cdot d_{h e a d} \cdot n_{c t x} \cdot B \cdot b$

$2$ — one copy for keys, one for values. (Keys and values are both cached; queries are not, because a past query is never needed again.)
$L$ — number of layers. Every layer has its own cache.
$n_{k v}$ — number of key/value heads. This is the only factor GQA changes.
$d_{h e a d}$ — per-head width (both $d_{k}$ and $d_{v}$ ; modern models use $d_{k} = d_{v}$ ).
$n_{c t x}$ — cached context length in tokens.
$B$ — batch size.
$b$ — bytes per element (2 for fp16/bf16, 1 for fp8).

Multi-Query Attention (Shazeer 2019) sets $n_{k v} = 1$ : all query heads share one key head and one value head. The cache shrinks by a factor of $h$ . Grouped-Query Attention (Ainslie et al. 2023) interpolates: split the $h$ query heads into $g$ groups and give each group its own KV head, so $n_{k v} = g$ :

$g = h \Rightarrow MHA, g = 1 \Rightarrow MQA, 1 < g < h \Rightarrow GQA$

$g$ is a hard design constant in the checkpoint — the number of KV heads is baked into the weight shapes, so you cannot change it at serving time.

Concretely, the arithmetic. For a 70B-class model with $L = 80$ , $h = 64$ query heads, $d_{h e a d} = 128$ , fp16:

Config	$n_{k v}$	bytes/token	cache @ 32k ctx, batch 1
MHA	64	$2 \cdot 80 \cdot 64 \cdot 128 \cdot 2 = 2, 621, 440 \approx 2.5$ MiB	$\approx 80$ GiB
GQA-8	8	$2 \cdot 80 \cdot 8 \cdot 128 \cdot 2 = 327, 680 \approx 0.31$ MiB	$\approx 10.0$ GiB

(Binary units: 1 MiB $= 102 4^{2}$ B, 1 GiB $= 102 4^{3}$ B, so 10.0 GiB $= 1.0737 \times 1 0^{10}$ B, 10.74 GB.) That is an 8× reduction, and it is the difference between one long-context request fitting on a node and not fitting. The quality argument for interpolating rather than collapsing to MQA is Ainslie et al.’s empirical result: uptrained GQA reaches quality close to MHA at approximately MQA speed. They also supply the recipe most teams actually use — take an existing MHA checkpoint, mean-pool the KV heads within each group, and continue training on ~5% of the original pretraining compute, rather than pretraining a GQA model from scratch.

Why this is the highest-leverage serving knob. Decode is memory-bandwidth-bound: each generated token requires reading the weights and the whole KV cache out of HBM. Shrinking $n_{k v}$ by 8× shrinks one of the two terms that set the throughput ceiling, and unlike quantization it does not touch the weights, so it does not interact with your quantized checkpoint.

MLA, one paragraph. DeepSeek-V2’s Multi-head Latent Attention takes a different route: instead of sharing KV heads across queries, it compresses the KV state into a low-rank latent vector and caches that, decompressing per layer at use time. The paper reports a 93.3% KV-cache reduction versus DeepSeek 67B alongside a 5.76× maximum generation throughput increase (DeepSeek-V2). The trade is that the cache is no longer a plain $n_{k v} \cdot d_{h e a d}$ layout, so kernels must understand the compression — which is why MLA support lands in serving frameworks later than GQA support.

2.5 FlashAttention: what it changes, and what it does not

Standard attention computes $S = Q K^{⊤} / d_{k}$ , then $P = softmax (S)$ , then $O = P V$ . The $n \times n$ matrix $S$ (and $P$ ) is materialised in HBM. For $n = 128 k$ , $h = 64$ heads, fp16, that is

$12800 0^{2} \times 64 \times 2 bytes \approx 2.1 TB$

per attention layer. Long before the FLOPs become the problem, the memory traffic does.

FlashAttention’s insight is that the bottleneck is IO between HBM and on-chip SRAM, not arithmetic. It (1) tiles the computation — load blocks of $Q$ , $K$ , $V$ into SRAM and compute blockwise; (2) uses an online softmax carrying a running maximum $m$ and running normalizer $ℓ$ , so the softmax is computed incrementally across tiles without ever seeing a full row; and (3) recomputes the $S$ blocks in the backward pass rather than storing them, because a recompute in SRAM is cheaper than a read from HBM. The result is $O (n)$ memory instead of $O (n^{2})$ and far fewer HBM accesses.

FlashAttention is exact. It is not an approximation, not a sparsity pattern, not a low-rank projection. It computes the same mathematical function as the reference implementation, up to floating-point reassociation. It does not reduce the $O (n^{2})$ FLOP count. It reduces memory traffic and memory footprint.

The three things it does not fix: (1) the FLOPs — prefill at 100k context is still $O (n^{2})$ compute, and only sub-quadratic attention (sliding window, sparse, linear/hybrid) changes that; (2) the KV cache — it caches exactly what MHA/GQA caches and makes that cache usable at long context, without shrinking it, which is GQA/MLA’s job; (3) arbitrary attention bias — fused kernels support a subset of masks and biases, and an exotic additive bias or per-request custom mask can force a fallback to the slower math kernel, so check support before designing around one.

flowchart LR
    A["Q,K,V"] --> B["eager: write n×n scores to HBM,<br/>read back for softmax, write P, read back"]
    A --> C["FlashAttention: tile into SRAM,<br/>online softmax (running m, ℓ),<br/>n×n never hits HBM"]
    B --> D["O(n²) memory, 5+ HBM round trips"]
    C --> E["O(n) memory, exact same result"]

3. Implementation

Reference implementations of the four deltas, plus the kernel call. All of it is runnable against PyTorch 2.14 [as of 2026-10]; no CUDA required for the first three.

import torch
import torch.nn as nn
import torch.nn.functional as F


class RMSNorm(nn.Module):
    """gamma only: no mean subtraction, no beta."""
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.eps, self.weight = eps, nn.Parameter(torch.ones(d_model))

    def forward(self, x):
        x32 = x.float()                      # sum-of-squares needs fp32 range
        rms = torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x32 * rms).to(x.dtype) * self.weight


class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff=None):
        super().__init__()
        # 8/3*d keeps parameter count equal to the 4d ReLU FFN; round to x128.
        d_ff = d_ff or 128 * ((int(8 * d_model / 3) + 127) // 128)
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)  # W_1
        self.w_up = nn.Linear(d_model, d_ff, bias=False)    # W_3
        self.w_down = nn.Linear(d_ff, d_model, bias=False)  # W_2

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class Block(nn.Module):
    """Pre-norm: norm the INPUT of each sublayer; the residual add stays bare.
    GQA: h query heads, g KV heads (g == h -> MHA, g == 1 -> MQA)."""
    def __init__(self, d_model, n_heads, n_kv_heads):
        super().__init__()
        assert n_heads % n_kv_heads == 0
        self.h, self.g, self.d_head = n_heads, n_kv_heads, d_model // n_heads
        self.n1, self.n2 = RMSNorm(d_model), RMSNorm(d_model)
        self.w_q = nn.Linear(d_model, n_heads * self.d_head, bias=False)
        self.w_k = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.w_v = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.w_o = nn.Linear(n_heads * self.d_head, d_model, bias=False)
        self.ffn = SwiGLU(d_model)

    def forward(self, x, is_causal=True):
        B, T, _ = x.shape
        y = self.n1(x)
        q = self.w_q(y).view(B, T, self.h, self.d_head).transpose(1, 2)
        k = self.w_k(y).view(B, T, self.g, self.d_head).transpose(1, 2)
        v = self.w_v(y).view(B, T, self.g, self.d_head).transpose(1, 2)
        # enable_gqa expands KV inside the kernel; doing it yourself with
        # repeat_interleave allocates the memory GQA exists to save.
        attn = F.scaled_dot_product_attention(
            q, k, v, is_causal=is_causal, enable_gqa=(self.h != self.g))
        x = x + self.w_o(attn.transpose(1, 2).contiguous().view(B, T, -1))
        return x + self.ffn(self.n2(x))

The knobs that matter, with realistic ranges

Knob	Typical range	Moving it up	Moving it down
`d_ff`	$8 d /3$ , rounded to ×128	More FFN parameters (dominant cost)	Cheaper, weaker
`n_kv_heads` $g$	1 (MQA) to $h$ (MHA); commonly $h /8$	Better quality, linearly bigger cache	Smaller cache, small quality cost
`eps` (RMSNorm)	$1 0^{- 6}$ (fp32) to $1 0^{- 5}$ (bf16)	More stable at low precision, slight scale bias	Underflow risk on small activations
Norm placement	pre-norm (default)	—	Post-norm needs warmup at depth

The kernel

F.scaled_dot_product_attention dispatches at runtime among three backends [as of PyTorch 2.14]: FlashAttention-2, Memory-Efficient Attention, and a plain C++ math implementation. All are enabled by default and the fastest compatible one is selected. Two behaviours that bite in production, both from the official docs:

Boolean mask semantics are inverted relative to nn.MultiheadAttention. In SDPA, True means “participate”; in MultiheadAttention’s key_padding_mask, True means “mask out”. Migrating without inverting the mask produces silently wrong attention.
Output can differ between backends. Fusing floating-point operations reassociates them; the docs state plainly that the output may differ depending on which kernel is chosen. This is why bit-exact reproducibility across GPU generations is not available for attention.

Force or forbid a specific backend when you need determinism or when a fused kernel’s mask support is missing:

from torch.nn.attention import sdpa_kernel, SDPBackend

with sdpa_kernel([SDPBackend.FLASH_ATTENTION]):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
# A backend that cannot run raises RuntimeError ("No available kernel. Aborting
# execution.") rather than warning and silently falling back -- use this during
# bring-up, not in production loops.

4. Numbers and trade-offs

Pre-norm vs post-norm: what you actually pay

The forward cost is identical — one normalization per sublayer either way. The difference is when you can stop tuning the learning rate: post-norm needs a warmup schedule whose length you must get right, and getting it wrong at 100 layers means divergence, whereas pre-norm lets Xiong et al. remove warmup with comparable results. You trade a little representational rigidity for a large reduction in hyperparameter risk. At 6 layers on a laptop, post-norm is fine.

RMSNorm vs LayerNorm: measured, not asserted

Zhang & Sennrich’s headline is 7%–64% faster runtime at comparable quality across RNN, seq2seq translation, image-caption retrieval and CIFAR-10 CNN experiments. The range’s shape matters: the saving is largest where the norm is a big fraction of a small layer’s work, and smallest where matmuls dominate. At 7B+ transformer scale it is a few percent of wall-clock — real, but not why RMSNorm won. It won for being no worse and strictly cheaper. A useful cross-check when someone claims a big norm speedup: divide the norm’s FLOPs by the layer’s total FLOPs. If the norm is 2% of the work, no implementation of it can give you 20%.

SwiGLU: equal-parameter comparison, not equal-width

The comparison that matters is at equal parameters and equal FLOPs, which is why $8 d /3$ exists. A SwiGLU at $d_{ff} = 4 d$ versus a ReLU FFN at $d_{ff} = 4 d$ is not a comparison of activations — the SwiGLU has 50% more parameters and of course wins. Shazeer’s ablations hold parameters constant; quote those.

GQA: exact, derivable memory numbers

Use the formula, not a rule of thumb: the number of layers and d_head differ per model and dominate the answer. The 70B table in §2.4 gives $0.31$ MB/token for GQA-8 and $2.5$ MB/token for MHA, exactly $64/8 = 8 \times$ . Against a 140 GB fp16 weight footprint, a single 32k-context request costs 10 GB of cache — and at batch 16 it costs 163 GB, larger than the weights. That crossover is why GQA is not optional for long-context serving.

FlashAttention: which axis moves

Change	FLOPs	Peak memory	HBM traffic	Wall-clock
FlashAttention vs eager attention	unchanged	$O (n^{2}) \to O (n)$	large reduction	faster
GQA vs MHA	unchanged	unchanged	KV read $\div (h / g)$	faster decode
Quantization (W4A16)	unchanged	weights $\div 4$	weight read $\div 4$	faster decode

These compose, and they attack different terms in the same bandwidth equation, so their speedups multiply roughly rather than overlap. The paper’s own end-to-end numbers: 15% on BERT-large at seq len 512, 3× on GPT-2 at 1K, 2.4× on long-range arena at 1K–4K, plus 0.7 better perplexity on GPT-2 and 6.4 points on long-document classification from the longer context it made affordable (FlashAttention). The small BERT number is the honest one to remember: at short sequence lengths, where the $n \times n$ matrix is small, FlashAttention buys little.

Version note. FlashAttention-2 (Dao 2023, 2307.08691) improves work partitioning and reports ~2× over v1; FlashAttention-3 (Shah et al. 2024, 2407.08608) targets Hopper specifically with fp8 support. The PyPI package flash-attn is at 2.8.3.post1 and requires Python ≥ 3.9 [as of 2026-10]. The FA3/Blackwell generation moves fast enough that you should treat any specific kernel version claim as stale within months. [UNVERIFIED: exact FA4/Blackwell release status.]

5. Pitfalls

Pre-norm with an unscaled residual init at depth. Symptom: the final pre-head norm is doing a lot of work, training loss is fine but the model needs an unusually low learning rate, and loss spikes appear when you raise it. The residual stream magnitude grows like $L$ ; scale each residual branch’s output projection by $1/ 2 L$ at initialization, as GPT-2 does.

RMSNorm’s $ϵ$ inside the square root but applied after rsqrt. Symptom: NaN or Inf losses that appear only under bf16 autocast and disappear in fp32. Compute the mean-of-squares in fp32 and place $ϵ$ inside the denominator, not as a post-hoc clamp.

Dropping $γ$ from RMSNorm “since it’s just a norm”. Symptom: the model trains but plateaus well below the baseline on every task, with no instability to explain it. Without a learned gain the layer cannot re-adjust per-feature scale after normalization.

SwiGLU at the ReLU-FFN hidden size. Symptom: the FFN has 50% more parameters than the baseline, so every comparison drawn from it is invalid. Use $8 d /3$ when the claim is “same budget”.

SwiGLU weight-naming mismatch. Symptom: a shape error on w_down, or worse, a silent load of the gate into the value slot and a loss worse than the released model reports. Published checkpoints use gate_proj/up_proj/down_proj (w1/w3/w2); map names explicitly and assert on shapes.

Assuming n_kv_heads is tunable at load time. Symptom: RuntimeError on weight shapes, or enable_gqa doing nothing because the checkpoint is plain MHA. $g$ is baked into the checkpoint; you can uptrain MHA into GQA, but you cannot flip it at serving time.

Materialising the expanded K/V instead of using enable_gqa. Symptom: peak memory is unchanged after moving to a GQA checkpoint, because repeat_interleave allocated the $h$ -head tensors you were trying to avoid. Confirm the fused kernel path is actually taken.

Treating FlashAttention as an approximation. Symptom: you avoid it out of fear and pay the memory cost for nothing, or you “A/B it for quality”, see a difference within noise, and conclude it is approximate. The differences are floating-point reassociation. A/B it for throughput and memory, never quality.

Expecting bit-exact reproducibility with a fused kernel. Symptom: the same seed produces a slightly different loss on a different GPU or PyTorch version. Documented behaviour, not a bug; SDPBackend.MATH is the escape hatch at a large throughput cost.

Forgetting SDPA’s inverted boolean-mask convention. Symptom: training loss looks excellent and generation is incoherent, because the “allowed” mask was inverted and the model attended to exactly what it should not see.

6. Using it in a project

This chapter is mostly not actionable when you call an API. It becomes actionable the moment you train from scratch or self-host.

Fine-tuning an open-weight checkpoint. None of these are decisions — pre-norm, RMSNorm, SwiGLU and the GQA head counts are all fixed by the checkpoint. You choose the checkpoint, and the two questions worth asking are: what is my KV cache per token at my target context? (compute it from the config) and does my serving stack have a fused kernel for this architecture? An architecture your framework falls back to the math kernel on will be several times slower than its benchmarks suggest.

Training a small model from scratch. Adopt all five. Pre-norm and RMSNorm are better defaults, SwiGLU at $8 d /3$ is free at equal budget, and FlashAttention is a one-line change through SDPA.

Serving long context. GQA/MLA first — largest single memory lever, composes with everything else. Then prefix caching for shared system prompts, then KV quantization.

When NOT to use each of these:

Pre-norm: when reproducing a specific 2017-era result, or on a shallow model where post-norm’s instability never manifests. Pre-norm is not a strict superset — post-norm’s optimization properties can be preferable at small depth; the case for pre-norm strengthens with depth.
RMSNorm: when you need re-centering invariance (some retrieval and embedding setups care that the representation is mean-centred), or when a tuned LayerNorm $ϵ$ cannot be re-validated.
SwiGLU: when the deployment target’s kernels have no gated-FFN path — some edge and NPU runtimes only optimize the 2-matrix FFN, and a gate on the slow path can be slower overall than a plain FFN on the fast path. Also at very small widths, where the third matmul’s launch overhead is not amortized.
GQA: on short-context, high-batch workloads where the KV cache is not the bottleneck and MQA’s quality cost is not worth paying. Do not pick it if you need per-head KV differentiation for interpretability work.
FlashAttention: when you need bit-exact reproducibility across hardware, or you use an attention variant (custom learned bias, unusual mask) the fused kernels do not support. The math backend is exact and portable; it is just slow.

7. How you would prove it worked

Every claim above is an ablation with a stated falsifier.

Ablation 1 — pre-norm vs post-norm at depth. Two identical 24-layer models, one each, trained at a fixed learning rate with no warmup. Metric: loss vs step. Baseline: the post-norm run. Prediction: post-norm diverges or spikes; pre-norm trains. Falsified if post-norm trains fine — in which case the model is too shallow or the LR too low for the effect to appear, and you have not tested the claim. A useful extra signal is the per-layer gradient norm at step 0: roughly flat for pre-norm, growing toward the output for post-norm.

Ablation 2 — RMSNorm vs LayerNorm. Swap only the norm layer; hold $d$ , depth, data, schedule, seed fixed. Metrics: tokens/sec and loss at fixed step count. Baseline: LayerNorm. Prediction: loss within noise, tokens/sec higher (a few percent at transformer scale, more at small scale). Falsified if a consistent loss gap appears in either direction — then the change was not free.

Ablation 3 — SwiGLU vs ReLU FFN at equal parameters. Build three configs: ReLU at $d_{ff} = 4 d$ ; SwiGLU at $d_{ff} = 4 d$ (more parameters — keep it only as a sanity check that parameters help); SwiGLU at $d_{ff} = 8 d /3$ . Metric: loss at matched parameter count and matched tokens. Falsified as evidence for gating if the $8 d /3$ SwiGLU does not beat the ReLU baseline — and if you only ran the $4 d$ SwiGLU, you have proven that more parameters help, which nobody doubted.

Ablation 4 — GQA-8 vs MHA on a 7B-class model. Metrics: measured KV bytes/token vs the formula; decode tokens/sec at batch 1, 8, 32; eval accuracy. Baseline: MHA. Derivable prediction: the cache is exactly $h / g = 8 \times$ smaller, and decode throughput gains most at small batch (bandwidth-bound) and least at large batch (compute-bound). Falsified if accuracy drops consistently beyond your eval’s noise floor across tasks — publish the drop if you see it, since MQA’s known weakness is exactly this.

Ablation 5 — FlashAttention on/off. Metrics: peak allocated memory and tokens/sec at $n \in {1 k, 4 k, 16 k}$ . Baseline: the eager/math path. Prediction: the memory gap widens with $n$ ; the throughput gap is small at 1k and large at 16k. Falsified if quality differs beyond numerical noise on a fixed seed — which should not happen, and if it does, check the mask convention or scale argument, not the kernel.

Reporting template (example — replace with your own numbers). “We moved a 7B-class decoder from MHA to GQA-8 by uptraining on 5% of the pretraining tokens. Measured KV memory fell from 0.63 MB/token (example) to 0.08 MB/token (example) — 8.0×, matching the formula — and decode throughput at batch 4 rose 1.7× (example). MMLU moved within noise (−0.4, n=5 seeds, example). FlashAttention was on throughout; it affects neither number.”

8. Interview drill

“Why does everyone use pre-norm instead of post-norm?” Strong: post-norm puts LayerNorm on the residual identity path, so $\partial x_{l + 1} / \partial x_{l}$ accumulates $L$ normalization Jacobians; Xiong et al. show the expected gradient near the output grows with depth at initialization, which is why post-norm needs tuned warmup. Pre-norm normalizes only the sublayer’s input, leaving an exact identity term in the Jacobian. Wrong: “pre-norm converges faster” — the controlled result is about stability and removing warmup, not a better final loss at matched budget. Mentioning GPT-2’s $1/ 2 L$ residual init completes the answer.

“What does RMSNorm drop relative to LayerNorm, and why is that allowed?” Strong: mean subtraction and the bias $β$ ; the RMS rescale and the gain $γ$ stay. Zhang & Sennrich hypothesise — with ablation support — that re-centering is dispensable and the rescaling is what stabilizes training. The gain is one reduction instead of two: 7%–64% faster on their workloads. Wrong: “it’s faster and just as good” — true, but says nothing about what was deleted or why that is safe. Also wrong: that RMSNorm removes the need for a gain.

“Derive the $8 d /3$ hidden size.” Strong: a GLU FFN has three matrices, so $3 d d_{ff}$ parameters; the ReLU FFN has two at hidden $4 d$ , equal to $8 d^{2}$ . Setting $3 d d_{ff} = 8 d^{2}$ gives $d_{ff} = 8 d /3$ , rounded to a multiple of 128. Wrong: “it’s 2/3 of 4d because there are three matrices instead of two” — the right ratio as numerology. What matters is that parameter count is held constant, because that makes the comparison fair.

“Why does adding a gate improve the FFN?” Strong: the gate makes the hidden-unit mask input-dependent. A ReLU is a fixed nonlinearity on one projection; a GLU multiplies $Swish (x W_{1})$ with $(x W_{3})$ , so how much of each value-branch unit passes is decided by a different learned projection of the same input — more expressive per parameter, and visible in Shazeer’s equal-parameter ablations. Wrong: “it adds a nonlinearity” — the FFN already had one. The gate adds multiplicative, input-conditioned control.

“How does GQA reduce memory, exactly?” Strong: the cache is $2 \cdot L \cdot n_{k v} \cdot d_{h e a d} \cdot n_{c t x} \cdot B \cdot b$ , and GQA reduces only $n_{k v}$ , from $h$ to $g$ — so the saving is exactly $h / g$ . Weights and FLOPs are untouched. Follow-up: “why not MQA ( $g = 1$ )?” → the largest saving, but Ainslie et al. report quality degradation; GQA interpolates, and their uptraining recipe (~5% of pretraining compute, mean-pool the KV heads) recovers MHA-level quality. Wrong: “GQA makes attention cheaper” — it changes the memory read during decode, not the FLOPs.

“Is FlashAttention approximate?” Strong: no — exact, up to floating-point reassociation. It removes the materialized $n \times n$ matrix from HBM via tiling and online softmax, and recomputes blocks in the backward pass. It reduces neither the $O (n^{2})$ FLOPs nor the KV cache. Wrong: “it trades a bit of quality for speed” — that describes sparse or linear attention. Second wrong answer: “it reduces FLOPs” — it reduces IO.

“Then what limits long-context inference?” Strong: two separate things. Prefill compute is $O (n^{2})$ FLOPs, made affordable — not smaller — by FlashAttention; decode is bandwidth-bound and dominated by re-reading the KV cache, addressed by GQA/MLA and KV quantization. Quality at long context is a third, separate problem: fitting the tokens does not mean using them well. Wrong: naming one bottleneck for both phases.

“What is MoE and when is it the wrong choice?” Strong: a Mixture-of-Experts replaces the dense FFN with $N$ expert FFNs and a router sending each token to the top- $k$ (usually 1–2), decoupling total parameters from FLOPs per token. It is wrong when you are memory-bound at small batch: every expert must be resident, so a 671B-parameter model with 37B active parameters still needs memory for 671B, and a single-request workload never amortizes it. Wrong: treating active-parameter count as the memory footprint.

Next: BERT vs GPT vs Encoder-Decoder — the deltas above describe how a modern block is wired; the next chapter is about what objective you train it with, and why one of the three attention patterns took the general-purpose crown.

Keyboard shortcuts

Henry's Notebook