Transformers, in the order they actually need to be learned

Field notes, not a tutorial: every mechanism traced down to the tensor shape and the hardware number that makes it true, in the order a working engineer actually needs it, from matrix multiply to why a 175B-parameter run spikes at step 40,000. Every chapter carries the math, a worked numeric example, real production numbers from named frontier models, and the code.

00 · Prerequisites#

Skip these and every later chapter becomes memorization instead of understanding.

❓ Why this gate exists. Every transformer explanation assumes you can read $QK^T$ and know what shape falls out, that you know why softmax turns numbers into a distribution and why it can die, that “backprop” isn’t a magic word, and that you understand why a GPU is fast at one thing and slow at another. Without these, attention looks like an incantation instead of a weighted average, and every systems chapter later reads as folklore instead of arithmetic.

Linear algebra: the minimum that’s actually load-bearing#

• Matrix multiply as batched dot products. $(B,T,d) \times (d,d) \to (B,T,d)$ is “for every token, take a weighted sum of its features.” That’s a linear layer.
• Dot product as similarity. $a \cdot b = |a||b|\cos(\theta)$. Large when vectors point the same direction. This is the entire mechanism attention scores are built on.
• Shapes over symbols. The single highest-leverage skill for reading transformer code is tracking tensor shape through every line, not re-deriving the calculus. Chapter 4 exists because this is where real implementations die.

Probability: what you actually need#

• Softmax turns arbitrary real numbers into a probability distribution:

$\text{softmax}(x)_i = \frac{e^{x_i}}{\sum_j e^{x_j}}$

It’s temperature-sensitive: scale the inputs down and the distribution flattens toward uniform; scale up and it sharpens toward argmax.

• Cross-entropy loss is “negative log probability the model assigned to the correct next token.” Training a language model is nothing more than minimizing this, one token at a time.
• Distributions over the vocabulary. The model’s output at every position is a probability distribution over ~32k–256k tokens, not a single answer. Sampling strategy (greedy, top-k, nucleus) is a separate decision layered on top.

Backpropagation, the one-paragraph version#

Forward pass computes a scalar loss. Backward pass applies the chain rule, layer by layer in reverse, to get $\partial \text{loss}/\partial \text{weight}$ for every parameter. You never need to derive this by hand in practice, autodiff does it, but you need the intuition: gradients are signals traveling backward through the same graph the data traveled forward through, and anything that blocks that path (a saturated sigmoid, a dead ReLU, a badly scaled residual) blocks learning.

Take $y = f(g(x))$ where $g(x) = wx + b$ and $f(z) = \text{ReLU}(z)$:

$\frac{\partial y}{\partial w} = f'(g(x)) \cdot x$

If $g(x) < 0$, then $f'(g(x)) = 0$, so $\partial y/\partial w = 0$ always, regardless of $x$: a “dead” unit, no gradient will ever reach $w$ through this path again. Every stability trick in chapter 5 is a variation on keeping this backward signal alive.

GPU basics: memory vs. compute#

A GPU kernel is either compute-bound (limited by how many FLOPs the cores can do) or memory-bound (limited by how fast data moves between HBM and on-chip SRAM). The ratio that decides which regime you’re in is arithmetic intensity: FLOPs per byte moved.

$\text{intensity} = \frac{\text{FLOPs}}{\text{bytes moved}}$

A large matmul (an FFN projection) has intensity on the order of $d$, high, compute-bound. An elementwise op (softmax, LayerNorm, dropout, bias-add) has intensity around 1, low, memory-bound, no matter how “small” it looks on paper.

quantity	value (A100, bf16)
SRAM bandwidth	~19 TB/s
HBM bandwidth	~2 TB/s
tensor-core peak	~312 TFLOPS
ridge point	~156 FLOPs/byte

That order-of-magnitude bandwidth gap between SRAM and HBM is the entire reason “IO-aware” algorithms beat “FLOP-optimal” ones, and it’s the whole justification for FlashAttention in chapter 6: naive attention is memory-bound because it materializes and re-reads a full $T \times T$ matrix through slow HBM; a smarter algorithm that never leaves SRAM wins even though it does the identical math.

Linear algebra is the alphabet, probability is the grammar, backprop is how the sentence gets corrected after you say it wrong, and the SRAM/HBM gap is why you can’t just shout the whole book at once, you have to read it in chunks that fit in your hands, and how you chunk it is worth more than how fast you can shout.

01 · Tokenization & the input pipeline#

Treat this as a first-class system constraint, not a preprocessing footnote. It sets your compute budget, your vocabulary’s ceiling, and silently determines what the model can and can’t learn.

❓ Why this topic. Tokenization controls three things simultaneously: sequence length (and therefore compute cost, since attention is quadratic in tokens), vocabulary size (and therefore embedding-table and softmax-layer size), and how gracefully rare or unseen text degrades. Get it wrong and you pay for it at every layer above it, forever, because retokenizing means retraining from scratch.

Byte-Pair Encoding, worked by hand#

BPE starts from the smallest unit (bytes or characters) and greedily merges the most frequent adjacent pair, repeatedly, until the vocabulary hits its target size.

corpus (with word-boundary marker): "l o w </w>", "l o w e r </w>", "n e w e s t </w>", "w i d e s t </w>"

step 1: most frequent adjacent pair = ('e','s') → merge into 'es'
step 2: most frequent pair now = ('es','t') → merge into 'est'
step 3: most frequent pair now = ('l','o') → merge into 'lo'
...continue until vocab_size target reached

result: "lower" → ["lo", "w", "e", "r", "</w>"]

Real tokenizer training (GPT-2’s tiktoken, SentencePiece) runs this over billions of characters, producing tens of thousands of merges. Inference-time tokenization just replays the learned merge list in priority order against new text, a lookup, not a re-optimization.

Byte-level BPE: why GPT-2 tokenizes bytes, not characters#

Character-level BPE has an unsolved edge: Unicode has over 140,000 characters, so a naive character vocabulary is already huge before any merges happen, and any character not seen during tokenizer training becomes an unrepresentable gap. GPT-2 (Radford et al., 2019) instead starts from the 256 possible byte values as the base alphabet. Every string, in every language, every emoji, every binary blob, is representable as some sequence of bytes, so byte-level BPE has a mathematical guarantee: there is no out-of-vocabulary input, ever. The cost is that a single character in some scripts (CJK, Devanagari) can require 2–4 bytes and therefore 2–4 tokens even after merges, a documented reason non-English text costs measurably more tokens per unit of meaning than English on GPT-family tokenizers.

Vocabulary size across real systems#

model	tokenizer	vocab size
GPT-2 / GPT-3	byte-level BPE	50,257
GPT-4 (cl100k_base)	byte-level BPE	~100,256
Llama 2	SentencePiece BPE	32,000
Llama 3	byte-level BPE (tiktoken-based)	128,256
PaLM / Gemini family	SentencePiece	256,000

Larger vocabularies mean shorter sequences for the same text at the direct cost of a bigger embedding table and a bigger, slower final softmax layer, since that layer’s cost scales as $O(T \times V \times d)$. Llama 3’s jump from 32k to 128k vocabulary was explicitly justified by Meta as a ~15% average token-count reduction across languages, a real inference-throughput win purchased with a larger embedding matrix.

Special tokens and chat templates#

Production tokenizers reserve IDs for structural tokens that never appear in raw text: <|endoftext|> (document boundary, GPT-family), <|im_start|>/<|im_end|> (turn boundaries, ChatML-style formats), role markers for system/user/assistant. Instruction-tuned models are trained on a specific chat template, an exact sequence of these special tokens wrapping each turn, and at inference time the serving code must reproduce that exact template byte-for-byte. A model trained with <|im_start|>user and served with a hand-rolled "User: " string sees a token sequence it never trained on, and instruction-following quality degrades in a way that looks like a model regression but is actually a formatting bug.

⚠️ Gotchas

• Leading-whitespace sensitivity. GPT-style BPE tokenizes " the" (leading space) as a different token from "the". A prompt-formatting bug that adds or drops one space silently changes the entire downstream token sequence.
• Digit chunking breaks arithmetic. Pre-2023 GPT tokenizers chunked multi-digit numbers inconsistently, so the same “hundreds place” appeared in wildly different token contexts depending on total digit count, a documented major contributor to GPT-3/early-GPT-4 unreliability at multi-digit arithmetic.
• Train/serve tokenizer version skew. A serving stack pointed at a newer or older tokenizer file than the one the model trained against silently misaligns token IDs, producing fluent-looking garbage with no error thrown anywhere.

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")

text = "unbelievable"
tokens = tokenizer.tokenize(text)   # ['un', 'believ', 'able']
ids = tokenizer.encode(text)        # [403, 12881, 540]

print(tokens, ids)
print("tokens per char:", len(ids) / len(text))
print(" the" == tokenizer.decode(tokenizer.encode(" the")))  # whitespace round-trips, verify it

📚 Further reading

• Sennrich et al., 2016 -- Neural Machine Translation of Rare Words with Subword Units (BPE)
• Radford et al., 2019 -- Language Models are Unsupervised Multitask Learners (GPT-2, byte-level BPE)
• Kudo & Richardson, 2018 -- SentencePiece

02 · Core transformer architecture#

Embeddings, positional encoding, attention, feedforward, residual stream, LayerNorm: the five pieces every block repeats, and where the parameters actually go.

$x = x + \text{Attention}(\text{Norm}(x))$ $x = x + \text{FeedForward}(\text{Norm}(x))$

Stack this block N times (32 for a 7B-class model, 96 for GPT-3 175B, 118 for PaLM 540B), wrap an embedding layer at the bottom and an unembedding (linear + softmax) at the top. Everything past this point in the field is scale, engineering, and variations on these five pieces, not a new architecture.

What frontier labs actually shipped#

Every figure below is from the model’s own paper or technical report, except the GPT-4 row, which OpenAI has never disclosed and is marked accordingly. The point of laying them side by side: nothing below is a different architecture, it’s the same five-piece block with different choices for norm, activation, position scheme, attention variant, and whether the FFN is dense or a mixture of experts.

model	year	total / active params	layers	d_model	Q / KV heads	norm	FFN	position	context
GPT-3 175B	2020	175B / 175B	96	12,288	96 / 96	LayerNorm, pre	GELU, dense	learned absolute	2,048
LLaMA 2 70B	2023	70B / 70B	80	8,192	64 / 8 (GQA)	RMSNorm, pre	SwiGLU, dense	RoPE (θ=10,000)	4,096
LLaMA 3.1 405B	2024	405B / 405B	126	16,384	128 / 8 (GQA)	RMSNorm, pre	SwiGLU, dense	RoPE (θ=500,000)	128,000
Mistral 7B	2023	7.3B / 7.3B	32	4,096	32 / 8 (GQA)	RMSNorm, pre	SwiGLU, dense	RoPE + sliding window	8,192 (w=4,096)
Mixtral 8x7B	2023	46.7B / ~12.9B	32	4,096	32 / 8 (GQA)	RMSNorm, pre	SwiGLU, MoE top-2-of-8	RoPE	32,768
DeepSeek-V3	2024	671B / 37B	61	7,168	128 / MLA (compressed)	RMSNorm, pre	SwiGLU, MoE (256+1 shared, top-8)	RoPE	128,000
PaLM 540B	2022	540B / 540B	118	18,432	48 / 1 (MQA)	LayerNorm, pre, parallel block	SwiGLU, dense	RoPE	2,048
GPT-4*	2023	~1.8T / ~280B (reported)	unconfirmed	unconfirmed	unconfirmed	unconfirmed	MoE, ~16 experts, top-2 (reported)	unconfirmed	8K–128K (product tiers)

Embeddings + positional encoding#

Token embeddings give the model “what,” positional encodings give it “where.” Self-attention alone is permutation-invariant, so without positional information “dog bites man” and “man bites dog” would produce identical attention patterns.

$x = \text{TokenEmbedding}(\text{token}) \cdot \sqrt{d_{\text{model}}} + \text{PositionalEncoding}(\text{position})$

The original transformer scales the token embedding by $\sqrt{d_{\text{model}}}$ before adding positional encoding, a detail that’s easy to drop: without it, the fixed-magnitude sinusoidal positional signal can dominate the learned token signal early in training. Original sinusoidal scheme (Vaswani et al., 2017):

$PE(pos, 2i) = \sin(pos / 10000^{2i/d}) \qquad PE(pos, 2i{+}1) = \cos(pos / 10000^{2i/d})$

Modern frontier models (LLaMA, Mistral, GPT-NeoX, PaLM) mostly replaced this with RoPE, which rotates Q/K vectors as a function of position instead of adding a separate vector, see chapter 8 for why that generalizes to longer sequences than the model trained on.

Feedforward (MLP) layer#

Attention moves information between token positions; the feedforward layer processes each position independently, and per current interpretability research is where a large share of a model’s factual “memory” actually lives.

$\text{FFN}(x) = W_2 \cdot \text{activation}(W_1 x) \qquad W_1: d{\to}4d,\ W_2: 4d{\to}d$

model	activation	note
GPT-1, original transformer	ReLU	simple, but “dies” for negative inputs (chapter 0)
GPT-2, GPT-3, BERT	GELU	smooth, non-zero gradient everywhere
PaLM, LLaMA, Mistral	SwiGLU	gated, extra parameter matrix, best quality-per-FLOP in ablations

$\text{SwiGLU}(x) = \big(\text{Swish}(xW_1) \odot xV\big) \cdot W_2 \qquad \text{Swish}(z) = z \cdot \sigma(z)$

SwiGLU needs three weight matrices instead of two (an extra gate $V$), so at equal parameter budget LLaMA-family models shrink the hidden expansion to roughly 2.7× instead of 4× to keep total FFN parameter count comparable, then the gating mechanism recovers, and in the original ablations (Shazeer, 2020) exceeds, the quality of a plain 4× GELU FFN.

Mixture of Experts: the other axis of scale#

Every model so far activates its entire FFN for every token, a dense model. A Mixture of Experts (MoE) layer instead replaces one FFN with $N$ parallel FFNs (“experts”) and a small learned router that picks only $k$ of them per token, so total parameter count grows with $N$ while compute per token stays proportional only to $k$. This decouples “how much the model knows” (total parameters) from “how much compute answering any single token costs” (active parameters).

$\text{router}(x) = \text{softmax}(W_r x) \qquad \text{top\_k} = \text{indices of the } k \text{ highest router scores}$ $\text{MoE}(x) = \sum_{i \in \text{top\_k}} \text{router}(x)_i \cdot \text{FFN}_i(x)$

model	routing	total params	active params	notable detail
Mixtral 8x7B (Jiang et al., 2024)	top-2 of 8 experts, per layer	46.7B	~12.9B	same base architecture as dense Mistral 7B, just the FFN swapped for MoE
DeepSeek-V3 (DeepSeek-AI, 2024)	top-8 of 256 routed + 1 always-on shared expert	671B	37B	fine-grained experts plus a shared expert that captures common knowledge, so routed experts specialize more cleanly
GPT-4 (reported, unconfirmed)	top-2 of ~16 experts	~1.8T	~280B	never confirmed by OpenAI

⚠️ The systems problem MoE actually creates: load balancing. If the router is left to learn freely, it tends to collapse onto a small favorite subset of experts, “expert collapse,” wasting the rest of the parameter budget. The classic fix (Fedus et al., 2021, Switch Transformer) adds an auxiliary load-balancing loss term, but that loss competes with the main language-modeling loss and can hurt quality if weighted too heavily. DeepSeek-V3’s reported fix is auxiliary-loss-free balancing: each expert gets a learned bias added directly to its routing score, nudged up when under-used and down when over-used, balancing load without touching the training objective’s gradient directly.

A dense FFN is one generalist doing every task. MoE is a large team of specialists with a receptionist (the router) who sends each request to only two or three of them, the team’s total expertise can be enormous while the cost of handling any one request stays small, as long as the receptionist doesn’t keep sending everything to the same two people.

RMSNorm: what LLaMA, Mistral, DeepSeek, and PaLM actually run instead of LayerNorm#

LayerNorm does two things: it re-centers activations to zero mean, then rescales to unit variance, each with its own learned gain and bias.

$\text{LayerNorm}(x) = \gamma \cdot \frac{x - \text{mean}(x)}{\text{std}(x)} + \beta$ $\text{RMSNorm}(x) = \gamma \cdot \frac{x}{\text{RMS}(x)} \qquad \text{RMS}(x) = \sqrt{\text{mean}(x^2)}$

Zhang & Sennrich (2019) showed the re-centering step contributes little to LayerNorm’s benefit, the rescaling alone does almost all the stabilizing work, and dropping it removes an entire reduction and a set of parameters from every norm call in the network. At the depth and width of a 70B+ model, called twice per layer for hundreds of layers across a full training run, that’s a real, measured throughput win with no quality loss in the original ablations. This is why RMSNorm, not LayerNorm, is what actually ships in the production-config table above: LLaMA 1/2/3, Mistral, Mixtral, DeepSeek-V3, and PaLM all use RMSNorm; GPT-3 is the outlier still running the original LayerNorm.

Encoder-only, decoder-only, encoder-decoder: the same block, three wirings#

architecture	attention pattern	examples	typical use
encoder-only	full bidirectional, no causal mask	BERT, RoBERTa	classification, embeddings, understanding a fixed input
decoder-only	causal mask, sees only itself and the past	GPT-family, LLaMA, PaLM, Claude, most frontier LLMs	open-ended generation, chat, one architecture for everything via prompting
encoder-decoder	bidirectional encoder, causal decoder cross-attending to it	original transformer, T5, BART	sequence-to-sequence: translation, summarization

The field’s consolidation onto decoder-only for frontier LLMs is a practical bet, not a proof of superiority: a single causal model can imitate encoder-decoder behavior by concatenating “input, then output” in one sequence, trading a small architectural inefficiency for a much simpler, more uniform pretraining and scaling story. Encoder-only models remain the efficient choice specifically when you never need to generate, only classify or embed.

Residual stream: not a skip connection, the memory highway#

$x_{l+1} = x_l + F(x_l)$

Frontier-level framing: treat the residual stream as a shared communication channel every layer reads from and writes to, additively. Nothing is erased between layers, only added. Anthropic’s “A Mathematical Framework for Transformer Circuits” (Elhage et al., 2021) builds its entire analysis around this: because every layer’s output is summed into a common space, you can meaningfully decompose the final output as a sum of independent per-layer, even per-head, contributions, this decomposability is what makes circuit-level interpretability tractable at all.

• Prevents vanishing gradients: the gradient has a direct additive path back to every earlier layer, derivative exactly 1 along the skip connection.
• Lets shallow and deep layers contribute independently: a layer can learn to write near-zero and effectively skip itself for a given input.

Pre-norm vs. post-norm#

$\text{post-norm: } x + \text{LayerNorm}(F(x)) \qquad \text{pre-norm: } x + F(\text{LayerNorm}(x))$

Full treatment in chapter 5. The short version: pre-norm normalizes the input to each sublayer instead of its output, leaving the residual stream’s own additive path completely unobstructed by any normalization for backward gradients.

Think of the whole stack as an assembly line with a conveyor belt (the residual stream). Each station (layer) reads what’s on the belt, adds something, and passes it on, never removing what’s already there. Attention is the station that lets items on the belt talk to each other; the feedforward layer reshapes one item in isolation, using what’s effectively a large lookup table baked into its weights.

import torch, torch.nn as nn, math

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=4096):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe)

    def forward(self, x):                       # x: (B, T, d_model)
        return x + self.pe[:x.size(1)]

class SwiGLU(nn.Module):
    def __init__(self, d_model, mult=2.7):
        super().__init__()
        hidden = int(d_model * mult)
        self.w1 = nn.Linear(d_model, hidden)     # gate projection
        self.v  = nn.Linear(d_model, hidden)     # value projection
        self.w2 = nn.Linear(hidden, d_model)

    def forward(self, x):
        return self.w2(nn.functional.silu(self.w1(x)) * self.v(x))

📚 Further reading

• Vaswani et al., 2017 -- Attention Is All You Need
• Shazeer, 2020 -- GLU Variants Improve Transformer
• Elhage et al., 2021 -- A Mathematical Framework for Transformer Circuits (Anthropic)
• Zhang & Sennrich, 2019 -- Root Mean Square Layer Normalization
• Touvron et al., 2023 -- LLaMA / Llama 2
• Dubey et al., 2024 -- The Llama 3 Herd of Models
• Jiang et al., 2023 / 2024 -- Mistral 7B; Mixtral of Experts
• Fedus et al., 2021 -- Switch Transformers
• DeepSeek-AI, 2024 -- DeepSeek-V3 Technical Report

03 · Attention mechanics, deep dive#

Q, K, V; the 1/√d scaling; multi-head splitting; masking; and the attention variants every frontier serving stack now uses.

❓ Why this topic. This is the one genuinely new idea transformers introduced. Without it there is no way for token 500 to directly pull information from token 3; every prior sequence architecture (RNN, LSTM) relayed information step by step through a bottlenecked hidden state, which is why they struggled with long-range dependencies.

Q, K, V, worked by hand#

Every token produces three vectors via learned linear projections: a query (what am I looking for), a key (what do I contain, advertised for others to match against), a value (what I actually hand over if selected).

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

Tiny numeric example, 2 tokens, $d_k=2$:

Q = [[1, 0],      K = [[1, 0],      V = [[10, 0],
     [0, 1]]           [0, 1]]           [ 0,10]]

QK^T = [[1, 0],
        [0, 1]]

scale by 1/√2 ≈ 0.707:  [[0.707, 0],
                          [0, 0.707]]

softmax(row-wise): [[0.67, 0.33],
                     [0.33, 0.67]]     ← still a blend, not a hard copy, because scaling kept softmax soft

output = weights @ V = [[6.7, 3.3],
                         [3.3, 6.7]]

Why the scaling exists, precisely#

Dot products of two random $d$-dimensional vectors with unit-variance components have variance proportional to $d$. As $d_k$ grows, raw scores $QK^T$ grow large in magnitude, pushing softmax toward a near-one-hot regime where its gradient is nearly zero, and learning stalls.

$\text{Var}(q \cdot k) \approx d_k \qquad \text{Var}(q \cdot k / \sqrt{d_k}) \approx 1$

worked: d_k = 64 (typical head dim), unscaled dot products routinely land in [-40, 40]
  softmax([40, 0]) ≈ [1.0000000, 0.0000000]     -- dead gradient
  softmax([40/8, 0/8]) = softmax([5, 0]) ≈ [0.9933, 0.0067]  -- peaked, but not saturated dead

Multi-head splitting#

Instead of one attention computation over the full $d_{\text{model}}$, split into $h$ heads, each operating on $d_{\text{head}} = d_{\text{model}}/h$ dimensions, run attention independently per head, concatenate, project back.

$\text{MultiHead}(x) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\, W_O$

Anthropic’s “In-context Learning and Induction Heads” (Olsson et al., 2022) traced a specific two-head circuit found across many models: a previous-token head that copies information about the token immediately before the current position, feeding an induction head one or more layers later that searches the context for the last place the current token appeared and copies whatever followed it. That circuit alone explains a large share of in-context learning, why a model can complete "A: 1, B: 2, A:" with "1" correctly having seen the mapping exactly once in-context, with no gradient update at all.

Grouped-query and multi-query attention#

Standard multi-head attention gives every head its own K and V projections. At inference time this means the KV cache (chapter 7) is as large as if every head were independent. Multi-query attention (MQA) shares a single K/V pair across all query heads; grouped-query attention (GQA) is the practical middle ground, splitting query heads into groups that each share one K/V pair.

GQA: n_kv_heads groups, each of size (n_heads / n_kv_heads)
LLaMA 2 70B: 64 query heads, 8 KV heads  →  KV cache shrinks 8×  vs. full multi-head

The quality cost measured in the GQA paper (Ainslie et al., 2023) was small relative to the memory win, which is why essentially every frontier open-weight model released since (Llama 2/3, Mistral, Gemma) uses GQA rather than full multi-head attention.

Multi-head Latent Attention (MLA): DeepSeek’s answer to the same problem#

GQA reduces the KV cache by sharing K/V across groups of query heads, a discrete, coarse lever. DeepSeek-V2 (DeepSeek-AI, 2024) takes a different, continuous approach: instead of caching full-size K and V vectors per head, it projects each token down into one small shared latent vector, and reconstructs the per-head K and V from that latent on the fly during attention.

$c_{kv} = W_{\text{down}}\, x \qquad K_i = W_{\text{up},K,i}\, c_{kv} \qquad V_i = W_{\text{up},V,i}\, c_{kv}$

What’s actually cached is $c_{kv}$, one small vector per token, not K and V for every head. The DeepSeek-V2 paper reports roughly a 93% reduction in KV cache size versus standard multi-head attention at equivalent (or better) quality, and a reported 5.76× increase in maximum generation throughput compared to their own prior dense 67B model. This is the mechanism that lets DeepSeek-V3 serve a 671B-parameter (37B active) model with a KV cache footprint closer to a much smaller dense model’s, a good illustration of a broader pattern in frontier engineering: once a bottleneck is well understood mathematically, teams stop working around it (GQA) and start redesigning the mechanism that causes it (MLA).

Masking, all three kinds, and how they compose#

Three distinct mechanisms, frequently conflated, and a real batch usually needs more than one active on the same score matrix at once:

• Causal (look-ahead) mask: token $i$ must not see tokens $>i$, or next-token prediction becomes a trivial copy. Implemented by setting disallowed score positions to a large negative value before softmax.
• Padding mask: batching sequences of different lengths requires padding; the mask prevents real tokens from attending to padding positions.
• Packed-sequence (document) mask: production pretraining rarely pads at all, instead concatenating many short documents back-to-back into one fixed-length sequence, “sequence packing.” Token $i$ may attend causally within its own document, but must not attend across a document boundary into a different, unrelated document sitting earlier in the same packed sequence. Skipping this is a real, easy-to-miss bug: training still runs, loss still looks fine, and the model quietly learns to condition on unrelated neighboring documents.

These compose by logical AND: a packed, batched, causal training example needs causal ∩ padding ∩ document-boundary all applied to the same score matrix before a single softmax.

Attention dropout#

Standard dropout applied directly to the post-softmax attention weights, before the @ V matmul: randomly zero a fraction of attention weights and rescale the rest during training, disabled entirely at inference. The gotcha is entirely operational: a model left in train mode during evaluation or generation silently applies stochastic dropout to attention weights, producing run-to-run nondeterminism that looks like a model or sampling bug rather than a mode flag.

Query is the question you’re asking, key is the label on everyone else’s forehead, value is what they actually hand you if picked. Multi-head is running that exercise several times in parallel, each time asking a different kind of question. GQA is several people in the audience agreeing to share one notetaker instead of each keeping private notes. The causal mask is the rule that you can only call on people who spoke before you.

import torch, torch.nn.functional as F

def attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = Q @ K.transpose(-2, -1) / (d_k ** 0.5)   # (B, h, T, T)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    weights = F.softmax(scores, dim=-1)
    return weights @ V

def causal_mask(T, device):
    return torch.tril(torch.ones(T, T, device=device)).bool()

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.qkv = torch.nn.Linear(d_model, 3 * d_model)
        self.out = torch.nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        B, T, C = x.shape
        qkv = self.qkv(x).chunk(3, dim=-1)
        Q, K, V = [t.view(B, T, self.n_heads, self.d_head).transpose(1, 2) for t in qkv]
        out = attention(Q, K, V, mask)                 # (B, h, T, d_head)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.out(out)

📚 Further reading

• Olsson et al., 2022 -- In-context Learning and Induction Heads (Anthropic)
• Shazeer, 2019 -- Fast Transformer Decoding: One Write-Head is All You Need (MQA)
• Ainslie et al., 2023 -- GQA: Training Generalized Multi-Query Transformer Models
• DeepSeek-AI, 2024 -- DeepSeek-V2 (Multi-head Latent Attention)

04 · Implementation & shape bugs#

The gap between “I understand attention” and “my attention layer runs correctly” is entirely shape bugs. This is where most first implementations die, silently.

❓ Why this topic. Transformer math is simple; transformer tensors are not, every operation juggles batch, sequence, head, and feature dimensions simultaneously, and PyTorch will happily broadcast two mismatched tensors into a nonsense shape without raising an error.

The shape contract to memorize#

tensor	shape
input	(B, T, d_model)
Q, K, V after projection	(B, T, d_model)
after head split	(B, h, T, d_head)
attention scores	(B, h, T, T)
attention output	(B, h, T, d_head) → (B, T, d_model)
causal mask	(T, T), broadcasts to (B, h, T, T)
padding mask	(B, 1, 1, T), broadcasts over heads and query positions

Bug stories: symptom → root cause → fix#

symptom	root cause	fix
RuntimeError on .view() after transpose	memory no longer contiguous after .transpose(1,2)	call .contiguous() before .view(), or use .reshape() deliberately
no error, but padding tokens leak into predictions	padding mask shaped (B,T) broadcast against (B,h,T,T) masks the wrong axis	reshape mask to (B,1,1,T) explicitly, never rely on implicit broadcast
suspiciously good training loss, garbage generation	off-by-one in causal mask, future tokens leaking	zero a future position directly, confirm earlier logits are unchanged (chapter 9)
trains, but underperforms a known-good baseline	heads split with wrong axis order, interleaving features	view(B,T,h,d_head) then transpose(1,2), never reshape the last dim directly into (h, d_head) blindly
trains, but noticeably less stable	scaled by √d_model instead of √d_head after splitting heads	divide by √d_head, computed after the split
works at training length, breaks at longer eval length	positional table precomputed for a fixed max_len, silently indexes out of bounds	assert T ≤ max_len explicitly, or use RoPE (chapter 8)

Shape bugs are like mailing a letter with the right words in the wrong envelope: the postal system (broadcasting rules) still “delivers” it, just to the wrong address, silently. Nothing crashes; the model just quietly learns something other than what you intended.

def multi_head_attention_checked(x, qkv_proj, out_proj, n_heads, mask=None):
    B, T, C = x.shape
    assert C % n_heads == 0, f"d_model={C} not divisible by n_heads={n_heads}"
    d_head = C // n_heads

    qkv = qkv_proj(x)
    assert qkv.shape == (B, T, 3 * C)

    Q, K, V = qkv.chunk(3, dim=-1)
    Q, K, V = [t.view(B, T, n_heads, d_head).transpose(1, 2) for t in (Q, K, V)]
    assert Q.shape == (B, n_heads, T, d_head)

    scores = Q @ K.transpose(-2, -1) / (d_head ** 0.5)   # scale by d_head, not d_model
    assert scores.shape == (B, n_heads, T, T)

    if mask is not None:
        assert mask.dim() >= 2   # never (B,T) applied raw against (B,h,T,T)
        scores = scores.masked_fill(mask == 0, float("-inf"))

    weights = scores.softmax(dim=-1)
    out = (weights @ V).transpose(1, 2).contiguous().view(B, T, C)
    return out_proj(out)

05 · Training stability & optimization#

Pre-norm vs. post-norm, residual scaling, initialization, the optimizer, learning-rate schedule, mixed precision, gradient clipping: the difference between a run that finishes and one that NaNs at 2am.

❓ Why this topic. Architecture papers show the forward pass. They don’t show that the exact same architecture, with slightly wrong initialization or no gradient clipping, diverges at scale in ways invisible at small scale.

Pre-norm vs. post-norm, why it matters this much#

$\text{post-norm: } x + \text{LayerNorm}(F(x)) \qquad \text{pre-norm: } x + F(\text{LayerNorm}(x))$

In post-norm, gradients must flow back through a LayerNorm at every single layer on their way to earlier layers, and LayerNorm’s gradient can attenuate signal at depth. Past roughly 24–48 layers, post-norm transformers become difficult to train without very careful warmup, a known sensitivity in the original GPT-2 setup. Pre-norm leaves the residual stream’s own additive path unobstructed by normalization, which is why essentially every frontier LLM since GPT-2 uses pre-norm to reach depths of 40–120+ layers.

Residual / output scaling#

Because pre-norm lets every layer add unnormalized output straight into the stream, the stream’s variance grows with depth. GPT-2’s implementation scaled the output projection weights of attention and FFN by $1/\sqrt{2 N_{\text{layers}}}$ at initialization specifically to counteract this growth. Skipping this is a common, quiet reason deep pre-norm models are harder to train than “just use pre-norm” implies.

Initialization#

• Standard choice: weights ~ $\mathcal{N}(0, 0.02^2)$ for most linear layers (GPT-2/3 convention).
• Embedding and unembedding layers are frequently tied (shared weights).
• Bad init doesn’t always crash, it can silently cost thousands of steps while the network “un-does” a poor starting point, real, wasted compute that never shows up as an error.

The optimizer: AdamW#

$m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t \qquad v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2$ $\theta_t = \theta_{t-1} - \text{lr} \cdot \frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon} \; - \; \text{lr} \cdot \lambda \cdot \theta_{t-1}$

AdamW’s contribution (Loshchilov & Hutter, 2019) is separating weight decay from the gradient-based update: plain Adam applies decay through the adaptive denominator, which interacts badly with per-parameter scaling; AdamW applies decay directly to the weights, decoupled. Essentially every frontier model trains with AdamW, not vanilla Adam.

Learning-rate schedule: warmup + cosine decay#

$\text{warmup } (0 \to w):\ \text{lr}(t) = \text{lr}_{\max} \cdot t/w$ $\text{decay } (w \to T):\ \text{lr}(t) = \text{lr}_{\min} + \tfrac{1}{2}(\text{lr}_{\max}-\text{lr}_{\min})\left(1+\cos\!\left(\pi \tfrac{t-w}{T-w}\right)\right)$

Warmup exists because Adam’s second-moment estimate $v_t$ starts at zero and is extremely noisy over the first few hundred steps. A large learning rate against that noisy denominator early on is a well-documented source of early divergence; ramping up linearly gives the estimate time to stabilize. Cosine decay afterward is empirically motivated: it consistently outperforms linear or step decay across GPT-3, Chinchilla, and LLaMA’s published setups.

Softmax precision: why attention runs in fp32 even inside a bf16 model#

Mixed precision doesn’t mean every operation runs in low precision. The $QK^T$ matmul and the softmax that follows are sensitive to numeric range, and softmax’s exp() can overflow before the surrounding matmuls would. Standard practice, and what fused kernels like FlashAttention do internally, is to accumulate the score matmul and run softmax itself in fp32 regardless of overall training precision, then cast the resulting attention weights back down. Skipping this and running softmax naively in fp16 is a documented source of attention becoming unexpectedly “sharper” than intended, a subtle quality regression with no crash to point at.

The mask-fill value interacts with this directly: a mask filled with -inf is safe once softmax runs in fp32, but that same fill combined with a naive fp16 softmax reproduces the all--inf-row NaN failure from chapter 3.

Mixed precision#

Training in fp16 or bf16 roughly halves memory and can nearly double throughput, but fp16 has a narrow exponent range: gradients or activations can underflow to zero or overflow to inf. Two standard fixes: loss scaling (multiply the loss by a large constant before backward, divide gradients by the same constant before the optimizer step) and preferring bf16 over fp16 when hardware supports it (same exponent range as fp32, no loss scaling needed, at the cost of less mantissa precision).

Past bf16: DeepSeek-V3’s FP8 training#

DeepSeek-V3’s technical report (DeepSeek-AI, 2024) documents training the bulk of a 671B-parameter model with FP8 mixed precision, one exponent step narrower again than bf16, one of the first public demonstrations that 8-bit floating point is viable for frontier-scale pretraining rather than only post-training quantization (chapter 7). Doing this safely needed fine-grained per-tile and per-block scaling factors, rather than one global loss-scale constant, so different regions of a tensor with very different magnitude ranges don’t all get squeezed through the same narrow FP8 range, plus keeping select precision-sensitive operations in higher precision. The reported payoff was a substantial reduction in memory and a meaningful increase in training throughput relative to bf16.

Multi-token prediction as an auxiliary training objective#

DeepSeek-V3 adds a secondary, auxiliary head trained to predict two tokens ahead at each position simultaneously (multi-token prediction, MTP), present throughout training. The reported benefits are twofold: a denser training signal per token, and the same extra head can be repurposed at inference as a free draft model for speculative decoding (chapter 7).

Gradient clipping#

$\text{if } \|g\| > \text{max\_norm}:\ g \leftarrow g \cdot \frac{\text{max\_norm}}{\|g\|}$

Transformer loss surfaces have occasional sharp cliffs; one bad batch can produce a gradient 100× normal magnitude, and without clipping that single step can permanently destabilize training.

Loss spikes at scale: what actually happened#

⚠️ Documented production incidents

• PaLM (540B, Google, 2022). ~20 loss spikes reported. The fix that worked: restart from a checkpoint ~100 steps before the spike, skip the ~200–500 surrounding data batches, rather than changing hyperparameters, implying at least some spikes were data-dependent.
• OPT-175B (Meta, 2022). The public training logbook documents repeated hardware failures, loss divergences, and manual interventions across the run, a rare, transparent look at how non-smooth a real frontier-scale run actually is.
• GLM-130B (Tsinghua/Zhipu, 2022). Reported that embedding-layer gradient shrinkage plus attention-score clipping was needed after standard gradient clipping alone proved insufficient at their scale.

The takeaway that generalizes: frequent checkpointing and real-time gradient-norm monitoring are standard practice at frontier scale, not paranoia.

Pre-norm is normalizing your inputs before you think about them; post-norm is thinking first and cleaning up the mess after. Warmup is easing the clutch out slowly instead of dumping it; gradient clipping is a seatbelt, it doesn’t prevent the sharp turn, it just stops that one turn from throwing you out of the car.

import torch, math

scaler = torch.cuda.amp.GradScaler()   # for fp16; unnecessary for bf16

def lr_at(step, warmup, total, lr_max, lr_min=1e-5):
    if step < warmup:
        return lr_max * step / warmup
    progress = (step - warmup) / max(1, total - warmup)
    return lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * progress))

for step, batch in enumerate(dataloader):
    for g in optimizer.param_groups:
        g["lr"] = lr_at(step, warmup=2000, total=total_steps, lr_max=3e-4)

    optimizer.zero_grad()
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        logits = model(batch.input_ids)
        loss = torch.nn.functional.cross_entropy(
            logits.view(-1, logits.size(-1)), batch.labels.view(-1)
        )

    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    if not torch.isfinite(grad_norm):
        optimizer.zero_grad()          # skip the step instead of corrupting weights
        continue

    scaler.step(optimizer)
    scaler.update()

📚 Further reading

• Loshchilov & Hutter, 2019 -- Decoupled Weight Decay Regularization (AdamW)
• Chowdhery et al., 2022 -- PaLM
• Zeng et al., 2022 -- GLM-130B
• Meta AI, 2022 -- OPT-175B logbook
• DeepSeek-AI, 2024 -- DeepSeek-V3 Technical Report (FP8 training, multi-token prediction)

06 · Systems & efficiency#

FlashAttention, activation checkpointing, memory layout, kernel fusion, and the parallelism strategies that make 100B+ parameter training possible at all.

❓ Why this topic. At scale, the bottleneck stops being “does the math work” and becomes “does it fit in memory, and does it run fast enough to finish before the compute budget runs out.”

The roofline model#

Every kernel’s achievable throughput is capped by $\min(\text{peak FLOPs},\ \text{intensity} \times \text{peak bandwidth})$. Elementwise ops and naive attention sit in the memory-bound region, more FLOPs/s hardware doesn’t help them; FlashAttention’s win is moving the same math rightward on this chart, not making it “faster math.”

FlashAttention: IO-aware, not FLOP-optimal#

Naive attention materializes the full $(T,T)$ score matrix in HBM, applies softmax, then reads it back for the second matmul. For $T=8192$ that’s a 64M-entry matrix per head, per batch item, written and read multiple times. FlashAttention (Dao et al., 2022) tiles the computation: it processes Q, K, V in blocks small enough to fit in on-chip SRAM, uses an online-softmax trick to compute a numerically correct softmax incrementally without ever materializing the full matrix, and writes only the final output to HBM. Memory traffic drops from $O(T^2)$ to $O(T)$, which is why it enabled the jump from ~2k–4k token context windows to 32k–128k+ ones.

Activation checkpointing#

Backprop normally requires keeping every layer’s activations in memory from the forward pass. For a 96-layer model that’s often the dominant training memory cost, larger than the parameters themselves. Activation checkpointing trades compute for memory: store only a subset of activations and recompute the rest during the backward pass. Typical cost: ~20–30% more compute time for often 60–80% less activation memory.

Memory layout#

• Contiguous vs. non-contiguous tensors. .transpose() and .permute() change how a tensor is viewed without moving data. Downstream ops assuming contiguity fail or misbehave until .contiguous() physically copies the data, at a real bandwidth cost.
• Row-major layout and cache locality. Reshaping the last axis (head-splitting) is cheap; reshaping an early axis often forces a copy.

Kernel fusion#

Every separate op (add, LayerNorm, GELU) normally means a full read-compute-write round trip to HBM per op. Kernel fusion combines several elementwise ops into a single custom kernel so data is read once, kept in SRAM through several operations, and written back once, the exact optimization behind torch.compile and NVIDIA’s fused kernels.

Parallelism strategies at 100B+ scale#

strategy	splits	note
data parallel	batch across GPUs, full model replicated	simplest, but model must fit on one device
tensor parallel (Megatron-LM)	individual weight matrices, within a layer	needs fast interconnect (NVLink)
pipeline parallel	different layers on different devices	introduces “bubble” idle time
ZeRO stage 1	optimizer state sharded	DeepSpeed; ~4× memory reduction on Adam states
ZeRO stage 2	+ gradients sharded	further reduction, more communication
ZeRO stage 3 / FSDP	+ parameters sharded	each GPU holds only a shard of weights at rest

Real numbers: GPT-3 175B (2020) trained on a cluster reported as equivalent to thousands of V100-class GPUs over multiple weeks; PaLM 540B (2022) trained on 6,144 TPU v4 chips across two Pods, reporting 46.2% model FLOPs utilization. PaLM’s block also departs from the sequential pattern used everywhere else in this chapter: it computes attention and FFN in parallel from the same normalized input ($x + \text{Attn}(\text{Norm}(x)) + \text{FFN}(\text{Norm}(x))$, both branches reading the identical normalized activation), which the paper reports as roughly a 15% training speedup at that scale, at a cost the paper describes as a small, acceptable quality hit versus the sequential form.

DeepSeek-V3’s technical report is unusually transparent about total cost for a frontier-scale run: 671B total parameters (37B active per token via MoE) trained on 14.8T tokens, reported at roughly 2.79M H800 GPU-hours total across pretraining, context extension, and post-training, an estimated total training cost around $5.6M. Whether or not that specific figure holds up to scrutiny industry-wide, the report’s own breakdown, and the FP8 training and MLA choices in chapters 3 and 5, are a rare public look at exactly which systems decisions a frontier lab attributes its efficiency to.

metric	value
TPU v4 chips, PaLM 540B	6,144
PaLM model FLOPs utilization	46.2%
H800 GPU-hours, DeepSeek-V3	~2.79M
FlashAttention memory	O(T), vs O(T²) naive

import torch
from torch.utils.checkpoint import checkpoint

class CheckpointedBlock(torch.nn.Module):
    def __init__(self, block):
        super().__init__()
        self.block = block

    def forward(self, x):
        return checkpoint(self.block, x, use_reentrant=False)

# fused/flash attention directly (PyTorch 2.x dispatches to a FlashAttention-style kernel)
out = torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=True)

📚 Further reading

• Dao et al., 2022 -- FlashAttention
• Rajbhandari et al., 2020 -- ZeRO
• Shoeybi et al., 2019 -- Megatron-LM
• Chowdhery et al., 2022 -- PaLM (Pathways, MFU reporting, parallel block)
• DeepSeek-AI, 2024 -- DeepSeek-V3 Technical Report (training cost breakdown)

07 · Inference & serving#

KV cache, prefill vs. decode, batching, speculative decoding, quantization: training-time tricks don’t help here, it’s a different optimization problem.

❓ Why this topic. Training is throughput-oriented: process huge batches, per-sample latency doesn’t matter. Inference, especially interactive chat, is latency-oriented: a user is waiting for the next token, right now.

KV cache, with real numbers#

Without caching, generating token $T{+}1$ naively re-runs the full forward pass over all $T{+}1$ tokens. The KV cache stores each layer’s K and V tensors for every previous token, so each new token only computes its own new Q, K, V.

$\text{cache size} = 2 \times n_{\text{layers}} \times n_{\text{kv\_heads}} \times d_{\text{head}} \times T \times \text{batch} \times \text{bytes}$

LLaMA-2 7B, full multi-head (32 layers, 32 heads, d_head=128), T=4096, batch=1, fp16:
  ≈ 2.1 GB   -- per single request

same, with GQA (8 KV heads instead of 32):
  ≈ 0.54 GB  -- 4× smaller, why GQA matters at serving time

Prefill vs. decode#

• Prefill: processing the entire input prompt at once, highly parallel, compute-bound, fills the KV cache.
• Decode: generating one token at a time, can’t be parallelized across tokens, reads the growing KV cache from HBM, strongly memory-bandwidth-bound rather than compute-bound.

Serving systems like vLLM and TensorRT-LLM report “time to first token” (prefill) and “tokens/sec” (decode) as separate metrics for exactly this reason.

Continuous batching#

Because decode is memory-bandwidth-bound, serving one request at a time badly under-utilizes GPU compute. Continuous batching (Orca, 2022; popularized by vLLM) adds new requests into a running batch as soon as a slot frees up, at token-generation-step granularity, commonly cited as a 2–4× throughput improvement over static batching.

Speculative decoding#

A small, fast “draft” model proposes several tokens ahead; the large target model verifies all of them in a single parallel forward pass, exactly the operation prefill is already good at. Any accepted prefix is free; the first disagreement is corrected and drafting resumes. Because verification is parallel, this produces a real wall-clock speedup (commonly 2–3×) with output distribution mathematically identical to plain decoding.

Quantization for serving#

Since decode is memory-bandwidth-bound, the fastest way to speed it up is often moving fewer bytes per weight. GPTQ (Frantar et al., 2023) and AWQ (Lin et al., 2023) are the two dominant post-training quantization methods, both reducing model memory footprint roughly 4× (INT4) with typically small, measured quality degradation.

⚠️ The production gotcha: KV cache memory, not weights, is usually the real capacity limit. PagedAttention (Kwon et al., 2023, the core idea behind vLLM) treats the KV cache like OS virtual memory pages instead of one contiguous allocation per sequence, eliminating fragmentation, reported in the paper as up to 24× higher throughput over naive serving in some configurations, primarily by fixing memory waste.

Prefill is reading a whole letter someone handed you before replying. Decode is dictating your reply one word at a time, with everything already said held in short-term memory (the KV cache). Speculative decoding is a fast assistant guessing your next few words while you nod or correct them. Quantization is writing your notes in shorthand so you can flip pages faster, at a small cost in precision.

cache_k, cache_v = [[] for _ in range(n_layers)], [[] for _ in range(n_layers)]

def prefill(prompt_tokens):
    x = embed(prompt_tokens)
    for l in range(n_layers):
        q, k, v = project(x, layer=l)
        cache_k[l], cache_v[l] = k, v          # (B, T_prompt, d_head) per layer
        x = block_forward(x, q, k, v, layer=l)
    return x

def decode_step(prev_token):
    x = embed(prev_token)                       # single new token
    for l in range(n_layers):
        q, k_new, v_new = project(x, layer=l)   # q,k_new,v_new: (B, 1, d_head)
        cache_k[l] = torch.cat([cache_k[l], k_new], dim=1)
        cache_v[l] = torch.cat([cache_v[l], v_new], dim=1)
        x = block_forward(x, q, cache_k[l], cache_v[l], layer=l)
    return x  # logits over vocab for the next token

📚 Further reading

• Kwon et al., 2023 -- PagedAttention (vLLM)
• Leviathan et al., 2023 -- Fast Inference from Transformers via Speculative Decoding
• Frantar et al., 2023 -- GPTQ
• Lin et al., 2023 -- AWQ

08 · Long context scaling#

Attention’s quadratic cost is the wall. Every long-context technique is a different way of climbing it, and the field’s context-length numbers have moved fast.

Context length, how fast the field actually moved#

model	year	context
GPT-3	2020	2,048 tokens
GPT-3.5-turbo-16k	2023	16,384 tokens
GPT-4-32k	2023	32,768 tokens
Claude 2	2023	100,000 tokens
Claude 3 / later	2024+	200,000 tokens
Gemini 1.5 Pro	2024	1,000,000 tokens (up to 10M research config)

That jump from thousands to millions of tokens in roughly four years is the stack of everything in this chapter and chapter 6 applied simultaneously.

Where the quadratic cost comes from#

$\text{scores} = QK^T \to \text{shape } (T,T) \to O(T^2) \text{ memory and compute}$

FlashAttention removes the memory blowup but not the compute blowup, the number of pairwise dot products is still quadratic.

Sparse attention: don’t compute all pairs#

• Sliding window (Longformer): each token attends only to a fixed-size local window, cost drops to $O(Tw)$, linear.
• Sliding window + global tokens + random (BigBird): adds a small number of global tokens plus random connections, preserving expressiveness while staying near-linear.
• Local + occasional full attention (Mistral-style): alternate layers between cheap local attention and expensive full attention so information still propagates globally across depth.

Mistral 7B (Jiang et al., 2023) uses sliding-window attention with window $W=4096$ and a rolling buffer KV cache: only the last $W$ tokens’ K/V ever need to be cached, so cache size stays fixed regardless of total sequence length. Because attention at layer $k$ can indirectly reach back roughly $k \times W$ tokens (each layer’s window sees information already mixed in by the previous layer’s window), a 32-layer model with $W=4096$ has a theoretical receptive field over 100k tokens deep in the network, without ever computing a full $T \times T$ attention matrix.

RoPE (Rotary Position Embeddings)#

Instead of adding a positional vector to the embedding, RoPE rotates the Q and K vectors by an angle proportional to their position.

$q'_m = R(m\theta)\, q_m \qquad q'_m \cdot k'_n \text{ depends only on } (m-n)$

The dot product between a rotated query and rotated key depends only on their relative position, not absolute position, which generalizes much better to sequence lengths longer than training.

Extending RoPE past training length: NTK-aware scaling and YaRN#

A model trained with RoPE at 4k tokens degrades sharply if fed 32k tokens directly. Position interpolation rescales positions down to fit the trained range, but blurs fine-grained relative position information. NTK-aware scaling instead reparameterizes RoPE’s base frequency so high-frequency (local) rotations stay close to their original values while low-frequency (long-range) rotations stretch to cover the new length. YaRN (Peng et al., 2023) combines this with a temperature adjustment to attention logits.

LLaMA 3.1’s own context extension from 8K (Llama 3) to 128K is a real, published instance of this family of techniques: Meta increased RoPE’s base frequency $\theta$ from 10,000 to 500,000, stretching the rotation period so the same relative-position math stays well-behaved at 16× the original trained length, without changing the architecture at all.

ALiBi (Attention with Linear Biases)#

Add a fixed, non-learned penalty to attention scores proportional to distance between positions, down-weighting far-away tokens before softmax by construction.

$\text{score}(i,j) = q_i \cdot k_j - m|i-j|$

Because the bias has no learned parameters tied to a specific max length, ALiBi (Press et al., 2022, used in BLOOM and MPT) extrapolates gracefully to sequences longer than training length, often better than raw RoPE without extra tricks, at the cost of a strong locality bias.

Full attention is everyone in a stadium trying to privately talk to everyone else simultaneously, it doesn’t scale. Sliding window is only talking to the people in your row, but the person next to you already talked to the row behind them, so information still travels. RoPE is agreeing that what matters is how many seats apart you are, not which absolute seat numbers you hold. NTK-scaling is stretching the seat spacing at the back of a stadium built for 4,000 so it can seat 32,000 without redesigning the front rows.

import torch

def rope_freqs(dim, max_len, base=10000.0):
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(max_len).float()
    freqs = torch.outer(t, inv_freq)               # (max_len, dim/2)
    return torch.cat([freqs, freqs], dim=-1)        # (max_len, dim)

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)

def apply_rope(x, freqs):                           # x: (B, h, T, d_head)
    cos, sin = freqs.cos(), freqs.sin()
    return x * cos + rotate_half(x) * sin

📚 Further reading

• Su et al., 2021 -- RoFormer (RoPE)
• Press et al., 2022 -- ALiBi
• Beltagy et al., 2020 -- Longformer
• Peng et al., 2023 -- YaRN
• Jiang et al., 2023 -- Mistral 7B (sliding window, rolling buffer cache)
• Dubey et al., 2024 -- The Llama 3 Herd of Models (RoPE θ scaling to 500,000)

09 · Debugging & profiling#

A working forward pass is not evidence of a correct implementation. These are the checks that catch what “it runs without crashing” doesn’t.

Shape checks#

See chapter 4 in full. Assert every intermediate shape, especially anywhere a transpose, view, or mask broadcast happens.

Overfit a single tiny batch first#

Before any real training run: take one batch, disable dropout, train on just that batch for a few hundred steps. Loss should go to near zero. If it doesn’t, there is a correctness bug that no amount of hyperparameter tuning on the full dataset will fix. This is the single highest-value debugging step in the entire pipeline, and the most commonly skipped one.

Verify the causal mask directly, don’t trust the training curve#

A leaked causal mask still shows a decreasing, plausible-looking training loss, because the model correctly learns to exploit the leak. The tell: training loss looks unreasonably good, and validation-time generation is much worse than the loss would predict.

Deterministic mask semantics test, for all three mask types#

def test_causal_mask(model, tokenizer):
    ids = tokenizer.encode("the quick brown fox jumps")
    x = torch.tensor([ids])
    logits_a = model(x)
    x2 = x.clone(); x2[0, -1] = tokenizer.encode("zzz")[0]   # corrupt only the LAST token
    logits_b = model(x2)
    # every earlier position must be bit-for-bit identical: nothing may see the future
    assert torch.allclose(logits_a[0, :-1], logits_b[0, :-1]), "causal mask is leaking"

def test_document_mask(model, doc_a_ids, doc_b_ids):
    packed = torch.cat([doc_a_ids, doc_b_ids]).unsqueeze(0)
    boundary = doc_a_ids.size(0)
    out_packed = model(packed, doc_boundaries=[boundary])
    out_alone = model(doc_b_ids.unsqueeze(0))
    assert torch.allclose(out_packed[0, boundary:], out_alone[0], atol=1e-4), "doc mask leaking across documents"

fp32 vs. mixed-precision loss curves: a real divergence-diagnosis technique#

When a run is unstable and it’s unclear whether the cause is a logic bug or a precision issue, run the identical config for a few hundred steps in full fp32 alongside the mixed-precision run. If fp32 is stable and mixed precision spikes, the cause is numeric range, look at softmax precision, loss scaling, and mask fill values before touching the architecture. If both spike at the same step on the same data, look at the data batch and gradient norm instead.

Throughput vs. memory: profile both, they trade off#

• Memory profiling tells you if you’re OOM-limited and where the peak actually occurs, often not where you’d guess, activation memory during backward frequently dominates.
• Throughput profiling tells you whether you’re compute-bound or memory-bandwidth-bound per kernel, which tells you whether fusing kernels, using FlashAttention, or increasing batch size is the right lever.
• A model that’s slower after “optimization” is a common outcome of adding compute-saving tricks to a workload that was already memory-bandwidth-bound.

The checklist, in order#

1. Shapes assert correctly through every layer.
2. Model overfits a single batch to ~zero loss.
3. Causal, padding, and document-boundary masks each verified directly with a deterministic test.
4. Gradient norms logged and finite every step.
5. If unstable: fp32 run compared against the mixed-precision run to isolate precision bugs from logic bugs.
6. Memory profiled to find the actual peak, not the assumed one.
7. Throughput profiled to confirm compute-bound vs. memory-bound before optimizing.

batch = next(iter(dataloader))
model.train()
for step in range(300):
    optimizer.zero_grad()
    logits = model(batch.input_ids)
    loss = F.cross_entropy(logits.view(-1, vocab_size), batch.labels.view(-1))
    loss.backward()
    grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1e9)  # measuring only
    optimizer.step()
    if step % 20 == 0:
        print(f"step {step:>3}  loss {loss.item():.4f}  grad_norm {grad_norm:.2f}")
# expect loss → ~0 well before step 300; a plateau means a bug, not a hyperparameter issue

10 · Model behavior & failure modes#

Hallucination, miscalibration, and sycophancy aren’t bugs to patch, they’re direct consequences of what the model was actually trained to do, and at frontier scale they get debugged from the inside: which internal concept activated, in which layer, under how much uncertainty, not just whether the output text looks wrong.

How hallucination actually forms#

The model doesn’t have a “fact database” it fails to look up. It has a residual stream carrying a superposition of weak and strong signals accumulated across layers. When genuine evidence for a claim is weak or absent, strong priors baked into the feedforward layers’ weights dominate instead, and the model completes the pattern with something fluent and specific rather than something honest about its own uncertainty. Interpretability research increasingly locates this as forming in mid-layers, rather than as a last-token output error.

⚠️ This is not hypothetical

• Mata v. Avianca (2023). A lawyer used ChatGPT for legal research and submitted a brief citing six court cases that did not exist. The attorneys were sanctioned; the incident became the reference case for LLM hallucination risk in professional use.
• Google Bard’s public demo (February 2023). Bard incorrectly stated the James Webb Space Telescope took “the very first pictures of a planet outside our solar system,” an error caught immediately by astronomers and widely reported as contributing to a sharp single-day drop in Alphabet’s market capitalization.

Confidence calibration#

$\text{ECE} = \sum_{\text{bins}} \frac{|\text{bin}|}{N} \left|\, \text{accuracy(bin)} - \text{confidence(bin)} \,\right|$

Base language models are reasonably well-calibrated on raw next-token prediction. The miscalibration problem shows up specifically after RLHF/instruction-tuning: OpenAI’s own GPT-4 technical report documented that RLHF measurably degrades calibration relative to the pretrained base model.

Sycophancy: a closely related, separately documented failure#

Anthropic’s “Discovering Language Model Behaviors with Model-Written Evaluations” (Perez et al., 2022) documented that RLHF-tuned models tend to shift their stated answers toward whatever position the user expresses in the prompt, even on factual questions, more strongly as model scale and RLHF strength increase.

Internal mechanisms to inspect, not just outputs#

• Middle-layer concept formation. Azaria & Mitchell (2023) train a linear probe on mid-layer hidden states that predicts whether a generated statement is true or false substantially better than chance, often before the sentence finishes generating, direct mechanistic evidence that “this doesn’t check out” is represented internally well before the output layer.
• Residual stream priors dominating weak evidence. When a fact isn’t strongly represented anywhere, the residual stream still must produce something at the final layer, and what it produces is whichever direction the FFN’s learned priors push hardest.
• Attention under weak or diffuse evidence. When no token in context strongly matches the query, attention weight spreads thin across many weakly relevant positions instead of concentrating. A diffuse attention pattern at the position generating a claim is itself a usable, nearly-free signal.
• Logit lens per-layer trace. Running the logit lens (chapter 11) across every layer for a hallucinated span frequently shows the wrong answer only stabilizing in the last third of the network.
• Confidence calibration as an internal-vs-external mismatch. The model’s internal state often already “knows” better than its output confidence does.

What actually helps in practice#

• Retrieval grounding (chapter 11): give the model real evidence to attend to, so a strong in-context signal competes against the FFN prior.
• Uncertainty-aware decoding: lower temperature and add explicit abstention thresholds keyed to entropy or attention-diffusion signals.
• Explicitly training refusal: fine-tune on examples where the correct response is an honest “I don’t know,” since nothing in next-token pretraining rewards abstention.
• Concept suppression / steering. Inference-Time Intervention (Li et al., NeurIPS 2023) finds a “truthfulness” direction via linear probes, then shifts activations along it at the specific attention heads found most predictive, at inference time, reporting a substantial truthfulness-benchmark improvement with no retraining.
• Monitoring activations in production. A lightweight trained probe alongside generation, cheap relative to a second full forward pass, reading the more-honest internal signal directly.

The useful question at frontier scale isn’t “does this output look wrong.” It’s “which internal concept activated, in which layer, under how much attention diffusion.” At that point hallucination prevention is a systems and instrumentation problem, not a prompting problem.

Code-level places to intervene#

# 1. retrieval-gated generation: only answer from evidence actually retrieved
def retrieval_gated_generate(model, query, retriever, min_score=0.4):
    docs, scores = retriever.search(query)
    if max(scores, default=0) < min_score:
        return "I don't have reliable information to answer that."
    context = "\n".join(docs)
    return model.generate(prompt=f"Context:\n{context}\n\nQuestion: {query}")

# 2. confidence gate using entropy + attention diffusion
def confidence_gate(logits, attn_weights, entropy_threshold=3.0, attn_peak_threshold=0.15):
    probs = logits.softmax(dim=-1)
    ent = -(probs * probs.clamp_min(1e-12).log()).sum(-1)
    peak_attn = attn_weights.max(dim=-1).values.mean()   # how concentrated is attention
    return (ent.mean() < entropy_threshold) and (peak_attn > attn_peak_threshold)

# 3. instrument a mid-layer probe (Azaria & Mitchell-style truthfulness probe)
activations = {}
def hook(module, inp, out):
    activations["mid_layer"] = out.detach()
model.blocks[len(model.blocks) // 2].register_forward_hook(hook)
# ... after the forward pass:
truth_score = truthfulness_probe(activations["mid_layer"][:, -1, :])   # small logistic-regression head

# 4. groundedness check: does the answer span actually overlap retrieved evidence
def groundedness_score(answer_tokens, retrieved_tokens):
    overlap = set(answer_tokens.tolist()) & set(retrieved_tokens.tolist())
    return len(overlap) / max(1, len(set(answer_tokens.tolist())))

# 5. penalize unsupported continuations during preference optimization (DPO-style)
loss = dpo_loss(policy_logps_chosen, policy_logps_rejected,
                 ref_logps_chosen, ref_logps_rejected, beta=0.1)

Entropy as a cheap (partial) uncertainty proxy#

$H(p) = -\sum_i p_i \log p_i$

Low entropy doesn’t guarantee correctness, a model can be confidently wrong, exactly the hallucination and sycophancy failure modes above. High entropy reliably indicates genuine uncertainty.

import torch

def token_entropy(logits):                    # logits: (T, vocab_size)
    probs = torch.softmax(logits, dim=-1)
    return -(probs * torch.log(probs.clamp_min(1e-12))).sum(dim=-1)   # (T,)

entropy = token_entropy(logits)
flagged = entropy > entropy.mean() + 2 * entropy.std()   # candidate low-confidence tokens

📚 Further reading

• OpenAI, 2023 -- GPT-4 Technical Report
• Perez et al., 2022 -- Discovering Language Model Behaviors with Model-Written Evaluations (Anthropic)
• Mata v. Avianca, S.D.N.Y. 2023
• Azaria & Mitchell, 2023 -- The Internal State of an LLM Knows When It’s Lying
• Li et al., 2023 -- Inference-Time Intervention (NeurIPS)
• Burns et al., 2022 -- Discovering Latent Knowledge in Language Models Without Supervision

11 · Frontier research directions#

Activation steering, monosemantic features, uncertainty quantification, retrieval grounding, logit lens, and reasoning-model training: where research is actively trying to fix what chapter 10 describes.

Activation steering#

$h' = h + \alpha v$

Many human-interpretable concepts (sentiment, refusal, specific topics, even something as narrow as “the Golden Gate Bridge,” per Anthropic’s 2024 “Golden Gate Claude” demonstration) are represented as roughly linear directions in activation space. Find the direction (often the difference of mean activations between contrastive prompt pairs) and amplify or suppress the concept without touching any weights. Applied specifically to truthfulness, this is exactly the mechanism behind Inference-Time Intervention (chapter 10): steer along a learned “truthful” direction to directly suppress hallucination.

Sparse autoencoders and monosemantic features#

A single neuron typically fires for many unrelated concepts at once, superposition: the model packs more features than it has neurons by representing each as a sparse combination of many neurons. Anthropic’s “Towards Monosemanticity” (2023) and “Scaling Monosemanticity” (2024, applied to Claude 3 Sonnet) train a sparse autoencoder on a layer’s activations to decompose them into a much larger set of sparse, individually interpretable features, directions that reliably correspond to single, human-nameable concepts (one documented feature fired specifically for the Golden Gate Bridge, another for sycophantic praise, another for code containing security vulnerabilities).

Uncertainty quantification, beyond raw entropy#

• Temperature scaling: a single learned scalar dividing the logits before softmax, fit post-hoc to recalibrate confidence without changing the model’s ranking of answers.
• Semantic entropy (Kuhn et al., 2023): cluster sampled generations by semantic equivalence first, then measure entropy over the clusters, separating “many ways to phrase the same true answer” from “genuinely different, competing answers.”
• Self-consistency: sample the same prompt N times at nonzero temperature; agreement across samples is a practical uncertainty signal, at N× the inference compute.

Retrieval grounding#

A retrieval system (dense embedding search, or classic BM25) fetches relevant documents at inference time and prepends them to the context, so the answer can come from in-context attention over real text rather than pattern completion over parametric memory (Lewis et al., 2020, RAG). It doesn’t eliminate hallucination but substantially reduces the “no evidence, so the prior wins” case.

Logit lens & probing#

The logit lens (nostalgebraist, 2020): take the hidden state at an intermediate layer, skip straight to the final unembedding matrix, and read off what token it would predict right now. This often reveals the model’s evolving best guess forming layer by layer, sometimes flipping answers partway through. Probing generalizes this: train a small (often linear) classifier on frozen intermediate activations to test whether a specific piece of information is linearly decodable at that layer.

Process supervision and reasoning verification#

Instead of only rewarding a correct final answer, train a process reward model to score each intermediate reasoning step, catching a wrong step even when the model accidentally lands on the right final answer. OpenAI’s “Let’s Verify Step by Step” (Lightman et al., 2023) showed process supervision outperforming outcome supervision on math reasoning.

DeepSeek-R1 (DeepSeek-AI, 2025) is a real, fully published counter-example to “you need expensive step-level supervision to get this”: DeepSeek-R1-Zero was trained with large-scale reinforcement learning directly on the base model, with no supervised fine-tuning cold start at all, using only an outcome-based reward (is the final answer correct, is the output in the right format). Reasoning behaviors, self-verification, reflection, long chains of thought, emerged purely from that RL signal. The training method, GRPO (Group Relative Policy Optimization, Shao et al., 2024, from DeepSeekMath), removes PPO’s separate value/critic network: for a given prompt, sample a group of completions, score each with the reward function, and compute each one’s advantage relative to the group’s own mean and standard deviation, a much cheaper baseline than training a full critic network. The full DeepSeek-R1 release adds a small amount of cold-start supervised data before RL, mainly to fix R1-Zero’s readability and language-mixing issues, then applies further RL on top. Published in full, unlike OpenAI’s o1/o3, whose training methodology has not been disclosed in comparable detail.

Activation steering is nudging someone’s train of thought mid-sentence. Sparse autoencoders are finally getting a large enough dictionary to translate the model’s internal shorthand into words a human recognizes, instead of only reading its lips. Process supervision is grading a math student on every line of work instead of only the boxed answer at the bottom. GRPO is grading that student against the average of their own classmates’ attempts, instead of paying for a separate tutor (the critic network) to grade every step.

def logit_lens(model, hidden_states_per_layer, layer_idx):
    h = hidden_states_per_layer[layer_idx]        # (B, T, d_model), pre-final-norm
    h = model.final_norm(h)                        # models expect a final norm before unembedding
    logits = h @ model.unembed.weight.T             # (B, T, vocab_size)
    return logits.argmax(dim=-1)                    # the layer's "current best guess" per token

# run across all layers to watch the prediction evolve/flip depth-wise
for l, h in enumerate(hidden_states_per_layer):
    guess = logit_lens(model, hidden_states_per_layer, l)
    print(l, tokenizer.decode(guess[0, -1]))

📚 Further reading

• Bricken et al., 2023 -- Towards Monosemanticity (Anthropic)
• Templeton et al., 2024 -- Scaling Monosemanticity (Anthropic, Claude 3 Sonnet)
• Lewis et al., 2020 -- Retrieval-Augmented Generation
• Kuhn et al., 2023 -- Semantic Uncertainty
• Lightman et al., 2023 -- Let’s Verify Step by Step (OpenAI)
• Shao et al., 2024 -- DeepSeekMath (GRPO)
• DeepSeek-AI, 2025 -- DeepSeek-R1

12 · End-to-end production pipeline#

Every chapter above, in the order it actually executes for a single user request, with a rough latency budget.

1. Tokenization: user prompt → subword tokens → integer IDs (chapter 1)
2. Embedding + positional info: IDs → vectors, position injected via RoPE or learned/sinusoidal encoding (chapters 2, 8)
3. Prefill: full prompt processed in one parallel forward pass through every transformer block (chapters 2, 3, 6, 7)
4. KV cache populated from prefill, stored per layer (chapter 7)
5. Final LayerNorm + unembedding: last hidden state → logits over the vocabulary
6. Sampling: logits → temperature scaling → top-k / nucleus filtering → next token sampled
7. Decode loop: new token → embed → attend against cached K/V plus one new K/V pair → next logits → repeat, batched via continuous batching (chapter 7), optionally accelerated with speculative decoding
8. Stop condition: EOS token, max length, or stop sequence reached → detokenize IDs back to text

stage	typical bottleneck	dominant lever
tokenization	negligible (μs–ms)	tokenizer vocab efficiency
prefill	compute-bound	FlashAttention, batching prompts
first token latency	prefill length	prompt length, KV cache prep
decode (per token)	memory-bandwidth-bound	continuous batching, quantization, speculative decoding
concurrent capacity	KV cache memory	PagedAttention, GQA/MQA/MLA

⚠️ Where this pipeline actually breaks in production

• Prefill latency dominates for long prompts, decode throughput dominates for long generations, they need separate optimization and often separate monitoring dashboards.
• KV cache memory pressure under concurrent load is usually the real capacity limit, not raw compute.
• A tokenizer version mismatch between training and serving silently degrades quality with no obvious error anywhere in the stack.

13 · Minimal working transformer#

Every chapter’s mechanism, assembled into a runnable decoder-only model with a generation loop.

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads, self.d_head = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.out = nn.Linear(d_model, d_model)

    def forward(self, x, causal_mask):
        B, T, C = x.shape
        Q, K, V = self.qkv(x).chunk(3, dim=-1)
        Q, K, V = [t.view(B, T, self.n_heads, self.d_head).transpose(1, 2) for t in (Q, K, V)]
        scores = (Q @ K.transpose(-2, -1)) / (self.d_head ** 0.5)
        scores = scores.masked_fill(causal_mask == 0, float("-inf"))
        attn = scores.softmax(dim=-1)
        out = (attn @ V).transpose(1, 2).contiguous().view(B, T, C)
        return self.out(out)

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.norm1, self.norm2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.ff = nn.Sequential(
            nn.Linear(d_model, 4 * d_model), nn.GELU(), nn.Linear(4 * d_model, d_model)
        )

    def forward(self, x, causal_mask):
        x = x + self.attn(self.norm1(x), causal_mask)   # pre-norm, chapter 5
        x = x + self.ff(self.norm2(x))
        return x

class TinyGPT(nn.Module):
    def __init__(self, vocab_size, d_model=256, n_heads=8, n_layers=6, max_len=1024):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_len, d_model)
        self.blocks = nn.ModuleList([TransformerBlock(d_model, n_heads) for _ in range(n_layers)])
        self.norm_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)
        self.head.weight = self.tok_emb.weight            # weight tying

    def forward(self, idx):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        mask = torch.tril(torch.ones(T, T, device=idx.device)).bool()
        for block in self.blocks:
            x = block(x, mask)
        return self.head(self.norm_f(x))

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=50):
        for _ in range(max_new_tokens):
            logits = self(idx)[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, top_k)
                logits[logits < v[:, [-1]]] = float("-inf")
            probs = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_id], dim=1)
        return idx

⚠️ What this omits, on purpose

• No KV cache in generate(), it recomputes the full forward pass every step, correct but quadratic-total instead of linear-total. Add chapter 7’s cache before using this for anything beyond a toy.
• No RoPE, uses learned absolute positions for simplicity, swap in chapter 8’s apply_rope for real long-context behavior.
• No mixed precision, gradient clipping, warmup schedule, or activation checkpointing in the training loop shown, chapters 5 and 6 have the production version of that loop.
• No GQA/MLA, every head gets its own K/V, fine at this toy scale, a real memory problem at the scale described in chapter 7.

---

Zero → Frontier Engineering, log 01. A practical reference, not a tutorial: assumes chapter 0’s prerequisites and gets denser from there. Named production incidents, papers, and numbers throughout are cited inline per chapter under “further reading.”