Transformer Architecture and Its Variants

After tokenization, the model finally receives a sequence of ids. A frontier dense model turns those ids into predictions with a stack of nearly identical transformer blocks, and that block has barely changed since 2023 (Touvron et al. 2023; Bai et al. 2023). Almost everything inside it is a question that has been answered and closed: how to normalize, how to activate, how to encode position. One question stays open, and it is not in the math of the forward pass at all. It is a cost that lives in memory during generation, the size of the key-value cache, and it is the main reason an architect still redraws a working block. The useful reading is therefore double: the block as a settled forward pass, and the same block as the single open cost that drives attention variants from MHA through MQA and GQA to MLA.

One block, two jobs

A transformer block has to do two distinct jobs and do them in a way that survives being stacked a hundred times deep. It must let tokens exchange information, since a language model is useless if position $i$ cannot see position $j$, and it must transform each token's features, which is where most of the model's knowledge lives. Stacking those two operations naively does not work: a deep residual network is hard to keep trainable, the activation and normalization choices decide whether the gradient (the signal that says how to adjust each parameter to lower the loss) survives to the bottom layer, and attention, the one operation that mixes across positions, costs $O(n^2)$ in compute and leaves behind a cache that grows without bound during generation.

The right way to read a block is not "attention, then MLP." It is two sublayers that each read a shared residual stream, compute a delta, and write it back:

$$x = x + \text{sublayer}(\text{norm}(x))$$

The residual stream is the model's working memory, a vector of width $d_{\text{model}}$ carried from the embedding to the final layer. Each block gets one pass to read it and add to it. Depth is the number of such read-modify-write passes; width is how much the stream can hold at once. This framing makes the rest of the design fall out. Figure 8.1 draws one pre-norm block this way: the residual stream runs straight down as an identity path, and each sublayer branches off it through a norm, computes a delta, and adds the delta back.

The first sublayer is attention, the only place tokens exchange information, and the subject of the second half of this chapter. The second is a position-wise MLP, also called the feed-forward network or FFN, applied to each token independently, and it holds most of the parameters: it expands to roughly four times $d_{\text{model}}$ and projects back, costing about $8 \cdot d_{\text{model}}^2$ per block, dwarfing attention's projections. Knowledge lives mostly in the MLP; mixing lives entirely in attention.

Three questions the field has closed

Three choices keep a deep stack of these blocks trainable, and all three are settled. They are not live debates; they are the record of how the consensus stack was assembled, one near-free win at a time.

Normalization has two orthogonal axes. The form is LayerNorm (Ba et al. 2016), which re-centers and re-scales with a learned gain and bias, versus root mean square layer normalization (RMSNorm) (Zhang and Sennrich 2019), which only re-scales, dropping the mean subtraction and the bias. Mean-centering buys little, so RMSNorm is cheaper and now the default. The placement is post-norm, the original design with the norm after the residual add (Vaswani et al. 2017), versus pre-norm, with the norm inside the branch and an identity path straight down the residual stream. Pre-norm keeps a clean gradient highway from the top of the stack to the bottom, and it is what lets a deep model train without elaborate warmup (slowly ramping the learning rate up over the first steps of training) (Xiong et al. 2020).

Activation has a clear lineage: ReLU, simple but prone to dead units, then GeLU (Hendrycks and Gimpel 2016), the smooth BERT and GPT-2 default, then gated linear units, of which Swish-gated linear unit (SwiGLU) is the modern winner (Shazeer 2020). SwiGLU computes $(\text{Swish}(xW_1) \odot xV)\, W_2$, three matrices rather than two. The easy-to-miss detail: the hidden dimension is scaled by about $2/3$ so the gated FFN keeps the same parameter count as a plain $4\times$ FFN, which is what makes a like-for-like comparison fair.

Position is the third. Attention is permutation-invariant on its own, so the block must be told where tokens sit. rotary position embedding (RoPE) rotates the query and key vectors by an angle proportional to absolute position, in two-dimensional subspace pairs, so that the dot product $q \cdot k$ depends only on the relative offset $i - j$ (Su et al. 2021). A base-frequency parameter, theta, sets how fast the rotation turns and quietly governs how far context can later be extended. ALiBi takes the other route: it adds a static, head-specific linear penalty $-m \cdot (i - j)$ to the pre-softmax scores, with no learned positional parameters at all (Press et al. 2022). The contrast is rotations on vectors versus biases on scores. RoPE became the dense default.

One bracket choice frames the stack from outside. Embedding and unembedding sit at the two ends: the input embedding maps token ids to vectors, the LM head maps the final residual state to vocabulary logits (the raw per-token scores the model produces before softmax turns them into probabilities). Weight tying shares one matrix for both (Press and Wolf 2017). This is a genuine tension, not a default. Tying saves $\text{vocab} \cdot d$ parameters and regularizes small models, but once the vocabulary is large the saving is marginal, and most large modern models leave the two untied.

These choices, assembled, are the whole settled block. Vaswani et al. shipped the original transformer with post-norm and a plain FFN (Vaswani et al. 2017). Because post-norm diverges at depth without heavy warmup, pre-norm replaced it (Xiong et al. 2020); RMSNorm then shed LayerNorm's mean subtraction at no measured quality cost (Zhang and Sennrich 2019). SwiGLU edged out GeLU under the $2/3$ rescaling (Shazeer 2020), and on the positional axis RoPE displaced ALiBi as the default (Su et al. 2021). By the time of LLaMA and Qwen, the consensus had set: pre-norm, RMSNorm, a SwiGLU FFN, RoPE, and no bias terms anywhere (Touvron et al. 2023; Bai et al. 2023). Strikingly little has changed in the dense block since 2023; the action moved to data, scale, and post-training.

The one cost that stays open

Now attention itself, the sublayer that mixes across positions and the source of the chapter's live constraint. Scaled dot-product attention projects each token to a query, key, and value, then computes (Vaswani et al. 2017)

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

Softmax turns each row of scores into positive weights that sum to one; the $\sqrt{d}$ divisor keeps the logits from growing with dimension and saturating it. A causal mask zeros the upper triangle so a token attends only to the past. The cost is $O(n^2)$ in both compute and the score matrix, which is the pressure behind every downstream efficiency technique.

Multi-head attention splits $d$ into $h$ heads so attention can attend to several subspaces at once, concatenating their outputs and projecting once more (Vaswani et al. 2017). Heads are cheap in parameters but not free at inference, because each head carries its own keys and values, and those must be kept.

Keeping them is not a detail; it is what the rest of the chapter costs. To see why, follow how the block actually turns a prompt into text, because a language model generates one token at a time and does it in two passes with very different shapes. The first pass is prefill: it reads the entire prompt at once, every position attending to every earlier one, and computes the keys and values for all prompt tokens together. The second pass is decode: it emits the answer one token at a time, feeding each token just produced back in to produce the next. Decode is autoregressive: the output of one step is the input of the next, so the sequence grows by one position per step.

Decode is where keeping the keys and values pays off. A token's key and value never change once computed, so rather than recompute attention over the whole sequence at every step, which would be quadratic work repeated for every token, the model caches them and each step does the attention math for only the single new token against the cache of all the past ones. The cache is the price of not recomputing. That is why these vectors must be kept, and why their cost is measured in memory rather than compute. The serving consequences of the prefill-decode split, the latencies it sets and the scheduling it forces, are the subject of Chapter 31; here it is enough that decode is what gives the cache a reason to exist.

That cost is the spine of the chapter. The cache size is

$$2 \cdot \text{layers} \cdot \text{heads} \cdot d_{\text{head}} \cdot \text{seq} \cdot \text{batch} \cdot \text{dtype}$$

It grows linearly with context length and batch size, and at long context it overtakes the model weights themselves. Figure 8.4 draws that linear growth and the crossover with a flat weights line, with the slope of each variant set by how many KV heads it keeps. This single fact is what explains MQA, GQA, and MLA. The question each answers is: how much of the KV cache can you delete before quality breaks?

Try this for a 7B-class model: plot each variant's KV cache against context and read off where it overtakes the weights. Change n_kv to see how sharing heads pushes the crossover out.

import numpy as np
import matplotlib.pyplot as plt

layers, d_head, dtype, params = 32, 128, 2, 7e9  # fp16 bytes, 7B weights
weights_gb = params * dtype / 1e9
seq = np.linspace(0, 200_000, 400)

for name, n_kv in [("MHA", 32), ("GQA", 8), ("MQA", 1)]:
    per_tok = 2 * layers * n_kv * d_head * dtype  # bytes per token
    cache_gb = per_tok * seq / 1e9
    crossover = weights_gb * 1e9 / per_tok
    print(f"{name:3s}: {per_tok/1024:6.0f} KB/token, crosses weights at {crossover:>7.0f} tokens")
    plt.plot(seq / 1000, cache_gb, label=name)

plt.axhline(weights_gb, ls="--", color="gray", label=f"weights ({weights_gb:.0f} GB)")
plt.xlabel("context length (1000s of tokens)"); plt.ylabel("KV cache (GB)")
plt.legend(); plt.title("KV cache vs context, batch=1"); plt.show()

The first answer cuts the most. Multi-query attention shares a single key/value head across all query heads, cutting the cache by a factor of $h$ (Shazeer 2019). The cut is dramatic, but collapsing to one KV head can destabilize training and drop quality non-gracefully. The second answer interpolates. Grouped-query attention shares K and V within $g$ groups of query heads, recovering MHA ($g = h$) at one extreme and MQA ($g = 1$) at the other, and it can be uptrained cheaply from an existing MHA checkpoint (Ainslie et al. 2023). GQA recovers most of the quality at most of the savings and became the default in the LLaMA and Qwen families (Touvron et al. 2023; Bai et al. 2023). Figure 8.5 shows the mechanism: the cache holds one K and V per distinct KV head, so the savings come entirely from how many query heads are made to share each one.

Multi-head latent attention, introduced in DeepSeek-V2, attacks the same cost from a different angle (DeepSeek-AI 2024). Instead of sharing heads, it projects K and V down into a low-rank latent vector that is cached in place of the full per-head K and V, decompressing on the fly, with a decoupled RoPE component carried alongside. The pitch is GQA-class cache savings, or better, at MHA-class quality. It is the second branch of the cache-compression frontier that Figure 8.6 lays out.

The third branch is younger and attacks a different term. Sharing and compressing both shrink what each cached token costs, but every past token is still attended, so the $O(n^2)$ score computation survives untouched. Sparse attention drops that premise: each query attends to a selected subset of past tokens rather than to all of them, cutting attention compute and long-context cost together. Bolting a sparsity pattern onto a model trained with full attention has historically cost quality, so the current designs learn the selection during training. DeepSeek's native sparse attention (NSA) trains a hierarchical token selection jointly with the rest of the model (Yuan et al. 2025), and its successor, DeepSeek sparse attention (DSA), shipped in the production model DeepSeek-V3.2-Exp in September 2025, uses a lightweight indexer to pick which tokens each query reads (DeepSeek-AI 2025). What sparsity does to serving and to long-context behavior is Chapter 35's territory; here it matters as the frontier's third branch, the first to attack the quadratic itself.

FlashAttention belongs to this story only as a reminder of what these variants do not solve. It keeps attention exact and makes the $O(n^2)$ score matrix affordable by never writing it to high-bandwidth memory, but it does not shrink the KV cache, which is a memory footprint, not a compute kernel (Dao et al. 2022). The serving-time mechanics of that kernel and of cache paging belong to Chapter 33 and Chapter 32.

Architecture knobs that carry cost

Every choice above is a balance with a knee. Set side by side, they are the dials an architect actually turns.

• Norm placement, pre versus post. Post-norm can reach slightly better final loss when it trains at all, but it is unstable at depth without heavy warmup and careful initialization. Pre-norm trades a sliver of peak quality for a robustly trainable gradient highway, and wins in practice (Xiong et al. 2020).
• Norm form, LayerNorm versus RMSNorm. Mean-centering buys little. RMSNorm drops it for lower cost and one fewer parameter set, a near-free win, now default (Zhang and Sennrich 2019).
• Activation choice. SwiGLU's gain is small but consistent, and effectively free under the $2/3$ hidden-dimension rescaling (Shazeer 2020). The cost is a third weight matrix and slightly more irregular shapes for kernels and parallelism.
• Positional scheme as a one-way door. The base scheme is chosen at run start and is expensive to change. It constrains how far context can later be extended and how well the model extrapolates. RoPE's theta base is a small knob with large long-context consequences, the mechanisms for which are Chapter 35.
• Weight tying. A real tension: parameter saving and small-model regularization against the freedom and marginal cost at scale that lead large models to untie (Press and Wolf 2017).
• Depth versus width. At a fixed parameter budget, deeper means more sequential composition but harder training and parallelism, with pipeline bubbles and serial latency, while wider means more per-step compute, friendlier to tensor parallelism, and easier to keep stable. This is co-designed with the systems layer in Chapter 10, not resolved in isolation, and the quality-per-FLOP side of it is set by Chapter 5.
• KV cache versus quality, the dominant axis. MHA gives full quality at full cache, GQA most quality at a fraction of cache, MQA the smallest cache at a quality and stability risk, and MLA claims to break the trade through latent compression. The right point depends on target context length and batch, since the cache scales with both.

Building the block, and how it breaks

Assembling the consensus dense block is mostly bookkeeping once the choices above are made: pre-norm with RMSNorm on each sublayer's input, a SwiGLU FFN with the hidden dimension rescaled by $2/3$, RoPE applied to queries and keys before the dot product, GQA grouping for the KV heads, and no bias terms anywhere. Since LLaMA, nearly every open dense model carries this shape (Touvron et al. 2023).

The failure modes are worth naming because each bites at a different time. Post-norm at depth blows up or stalls early in training, the classic symptom pre-norm was invented to fix (Xiong et al. 2020). A positional scheme or a too-small RoPE theta chosen for short context caps long-context behavior later, and the cost surfaces only at the first context-window bump. Forgetting the $2/3$ rescaling when switching to a gated FFN silently inflates or shrinks the parameter budget and breaks like-for-like comparisons. Untied, unregularized output logits can drift without bound, which is why z-loss guards them, with mechanics in Chapter 5. And collapsing too far toward MQA can destabilize training and drop quality non-gracefully, which is the reason GQA exists (Ainslie et al. 2023): the cheapest cache is not safely the best one.

Further reading