P01 hands-on — Transformer, block by block

Attention from a dot product, then a language model that trains.

Source: handson/h01_transformer.py --- run it with python3 handson/h01_transformer.py
Full project spec: P01 — Transformer From Scratch

Every line of this page runs. The file builds a decoder-only transformer out of nine independent pieces, each of which proves one mechanism on its own before anything is stacked, and then trains the assembled model on a real corpus until the loss falls to 0.11 nats.

The order is deliberate. Attention is introduced as a weighted average whose weights are computed from the data, which is all it is, and only then do the scaling factor, the causal mask, and multiple heads get added --- each one motivated by a failure you can see in the numbers of the block before it. The same discipline applies to the training loop: the overfit test in the assembly exists because a language model that cannot memorise sixty tokens has a bug that no amount of hyperparameter tuning will fix.

Read the outputs, not just the code. Block 3 shows what the sqrt(d_k) denominator is actually protecting you from, and it is not obvious from the formula.

Contents

• Block 1 — Tokenizer
• Block 2 — Batching
• Block 3 — Softmax + cross-entropy
• Block 4 — One attention head
• Block 5 — The leak test
• Block 6 — Multi-head
• Block 7 — The pre-norm block
• Block 8 — The model
• Block 9 — The same attention, in torch
• The assembly
• The design space
• Latency, bandwidth and the memory hierarchy
• Hardware: CPU, GPU, TPU
• Advanced algorithms and alternatives
• How this connects to the rest of the track
• Failure modes at scale
• Primary sources
• Running it
• What to do with this

How to read this page

Each block below is a self-contained lego piece: it builds one mechanism, proves it works on its own, and returns what the next block needs. The code is the real source, sliced out of the script. The output underneath it is the real output, captured by running that script --- not transcribed, not idealised. Where a measurement contradicted what I expected, the contradiction is in the output and the prose says so.

The assembly at the end wires every block into one working thing and measures it.

Block 1 — Tokenizer

Teaches: a vocabulary is a bijection, and it must round-trip

The problem. Before a model can learn anything, text has to become integers, and the mapping has to be exactly invertible. A tokenizer that does not round-trip introduces a silent error floor that no amount of training can cross, and it will look like a modelling problem.

@block(1, "Tokenizer", "a vocabulary is a bijection, and it must round-trip")
def b1(s, show):
    class CharTokenizer:
        def __init__(self, text):
            self.vocab = sorted(set(text))
            self.stoi = {c: i for i, c in enumerate(self.vocab)}
            self.itos = {i: c for c, i in self.stoi.items()}
        @property
        def size(self): return len(self.vocab)
        def encode(self, t): return [self.stoi[c] for c in t]
        def decode(self, ids): return "".join(self.itos[i] for i in ids)

    tok = CharTokenizer(CORPUS)
    assert tok.decode(tok.encode(CORPUS)) == CORPUS, "round-trip failed"
    if show:
        print(f"  vocab size {tok.size}: {''.join(tok.vocab)!r}")
        print(f"  round-trip on {len(CORPUS)} chars: OK")
        print(f"  compression: 1.00 bytes/token (the baseline BPE must beat)")
    return {"tok": tok, "data": np.array(tok.encode(CORPUS), dtype=np.int64)}


# ─────────────────────────────────────────────────────────── block 2

Reading the implementation

Character-level tokenisation is chosen here because it is a bijection by construction — every character maps to exactly one id and back — which removes an entire class of bug from the rest of the project. The assertion that decode(encode(text)) == text is not ceremony; it is the property everything downstream assumes.

The cost is sequence length. Character-level means ~4--5× more tokens than subword for English, and since attention is \(O(T^2)\) in memory that is a 16--25× increase in the score matrix. This is the trade the tokenizer makes and the reason production models use BPE.

What the numbers say

Output:

  vocab size 28: ' .abcdefghijklmnopqrstuvwxyz'
  round-trip on 3720 chars: OK
  compression: 1.00 bytes/token (the baseline BPE must beat)

Beyond the toy

Byte-pair encoding builds a vocabulary by repeatedly merging the most frequent adjacent pair, starting from bytes. The consequences are worth knowing precisely because they leak into model behaviour:

• Byte-level BPE (GPT-2 onward) never has an out-of-vocabulary token, because the base alphabet is all 256 bytes. That is why these models handle emoji, arbitrary Unicode, and binary garbage without a special case.
• Tokenisation is not language-neutral. The same sentence costs ~1.5--3× more tokens in Hindi or Thai than in English with a typical English-dominant vocabulary, which is a direct cost and context-length penalty for those users.
• Digit and whitespace handling explains a startling amount of model behaviour. If "1234" tokenises as "12"+"34", arithmetic becomes positional string manipulation, which is why many models are erratic at multi-digit arithmetic and why later models force digit-by-digit splits.
• The SolidGoldMagikarp class of bug: tokens present in the tokenizer's training corpus but effectively absent from the model's produce untrained embeddings and bizarre generations. A vocabulary is a joint artefact of the tokenizer and the training data, and a mismatch between them is a real failure mode.

Block 2 — Batching

Teaches: a language model predicts the NEXT token: y is x shifted by one

The problem. "Predict the next token" has to become a tensor operation. The shift-by-one construction is three lines and is the entire supervision signal — and an off-by-one here produces a model that either cheats or learns nothing.

@block(2, "Batching", "a language model predicts the NEXT token: y is x shifted by one")
def b2(s, show):
    def get_batch(data, B, T, rng):
        ix = rng.integers(0, len(data) - T - 1, size=B)
        x = np.stack([data[i:i + T] for i in ix])
        y = np.stack([data[i + 1:i + T + 1] for i in ix])
        return x, y
    x, y = get_batch(s["data"], 4, 8, rng)
    assert (x[:, 1:] == y[:, :-1]).all(), "y must be x shifted by one"
    if show:
        print(f"  x {x.shape}  y {y.shape}")
        print(f"  x[0] = {s['tok'].decode(x[0])!r}")
        print(f"  y[0] = {s['tok'].decode(y[0])!r}   <- shifted by one")
    return {"get_batch": get_batch}


# ─────────────────────────────────────────────────────────── block 3

Reading the implementation

y is x shifted left by one, so position \(t\) in the input predicts position \(t+1\). Every position in the sequence is a training example, which is what makes language modelling so sample-efficient: a batch of 16 sequences of 64 tokens yields 1024 supervised predictions, not 16.

This is also where the causal mask becomes non-negotiable. Because the targets are inside the same sequence as the inputs, an attention pattern that can see forward gives the model the answer. Block 5 exists specifically to test that it cannot.

What the numbers say

Output:

  x (4, 8)  y (4, 8)
  x[0] = 'uns . th'
  y[0] = 'ns . the'   <- shifted by one

Beyond the toy

• Packing. Real training concatenates documents and cuts fixed-length windows, so no compute is wasted on padding. The subtlety is that a window then spans a document boundary, and the model learns spurious continuations unless the attention mask is reset at boundaries — a detail many implementations get wrong and few notice.
• Sequence length is a curriculum variable. Training at short context and extending later (position interpolation, YaRN, NTK-aware scaling) is far cheaper than training long from scratch, because of the \(T^2\) memory term.
• Batch shape drives hardware efficiency. Ragged batches waste compute proportional to the length variance; length-bucketed sampling recovers it. At scale this is a double-digit percentage of total training cost.

Block 3 — Softmax + cross-entropy

Teaches: the entropy floor tells you what 'learning nothing' looks like

The problem. You cannot tell whether a loss of 3.3 is good without knowing what "learning nothing" scores. The entropy floor is that reference, and computing it takes one line — after which every training curve is interpretable.

@block(3, "Softmax + cross-entropy", "the entropy floor tells you what 'learning nothing' looks like")
def b3(s, show):
    def softmax(z, axis=-1):
        z = z - z.max(axis=axis, keepdims=True)     # stability: never exp a big number
        e = np.exp(z)
        return e / e.sum(axis=axis, keepdims=True)
    def cross_entropy(logits, targets):
        p = softmax(logits)
        n = np.prod(targets.shape)
        flat = p.reshape(-1, p.shape[-1])[np.arange(n), targets.reshape(-1)]
        return float(-np.log(flat + 1e-12).mean())
    V = s["tok"].size
    uniform = np.zeros((2, 5, V))
    floor = cross_entropy(uniform, rng.integers(0, V, (2, 5)))
    if show:
        print(f"  vocab {V}, uniform logits -> loss {floor:.4f} nats")
        print(f"  ln({V}) = {np.log(V):.4f}   <- the ENTROPY FLOOR")
        print(f"  a model below this on shuffled targets has label leakage")
        print(f"  huge logits (1e4) -> {cross_entropy(np.full((1,1,V),1e4), np.zeros((1,1),int)):.4f}, no NaN")
    return {"softmax": softmax, "cross_entropy": cross_entropy, "floor": float(np.log(V))}


# ─────────────────────────────────────────────────────────── block 4

Reading the implementation

Cross-entropy in nats is \(-\frac{1}{n}\sum \log p(y_i)\), and the two reference points that make it readable are:

• Uniform guessing over \(V\) tokens: \(\log V\) nats.
• The unigram entropy of the corpus: what a model that learns only token frequencies achieves.

The subtraction \(\log V - H(\text{unigram})\) is exactly the information a frequency table carries, and it is usually a large fraction of the first improvement any language model shows.

The numerical detail: logsumexp with the max subtracted. Computing \(\log \sum e^{x_i}\) naively overflows for \(x > 709\) in float64 and \(x > 88\) in float32. Subtracting \(\max_i x_i\) first is mathematically identity and numerically essential — and it is the same trick FlashAttention's online softmax generalises to a streaming setting.

What the numbers say

Output:

  vocab 28, uniform logits -> loss 3.3322 nats
  ln(28) = 3.3322   <- the ENTROPY FLOOR
  a model below this on shuffled targets has label leakage
  huge logits (1e4) -> 3.3322, no NaN

Beyond the toy

Perplexity is \(e^{\text{loss}}\) and is the more interpretable unit: "the model is as uncertain as if choosing uniformly among \(e^{\text{loss}}\) tokens". Three cautions that make published perplexities incomparable:

• Perplexity depends on the tokenizer. A character-level model and a BPE model on the same text have entirely different perplexities. Only bits-per-byte is comparable across tokenizations.
• A loss below the entropy floor is a bug, not a triumph — it means the model saw the answer. That is the leak test in block 5, and having the floor computed here is what makes the test possible.
• Loss is averaged over positions, so a model that is excellent at position 1 and useless at position 64 looks mediocre everywhere. Per-position loss curves are the diagnostic, and they are how you discover that a model has learned nothing beyond its first few tokens of context.

Block 4 — One attention head

Teaches: content-addressed lookup, and why we divide by sqrt(dk)

The problem. Attention is usually introduced as a formula. It is more usefully introduced as a weighted average whose weights are computed from the data — and once framed that way, every part of the formula has a job you can see.

@block(4, "One attention head", "content-addressed lookup, and why we divide by sqrt(dk)")
def b4(s, show):
    softmax = s["softmax"]
    def head(X, Wq, Wk, Wv, causal=True):
        Q, K, V = X @ Wq, X @ Wk, X @ Wv
        scores = Q @ K.transpose(0, 2, 1) / np.sqrt(Q.shape[-1])
        if causal:
            T = X.shape[1]
            mask = np.tril(np.ones((T, T), bool))
            scores = np.where(mask, scores, -np.inf)   # BEFORE the softmax
        A = softmax(scores)
        return A @ V, A
    T, d = 6, 16
    X = rng.standard_normal((1, T, d))
    W = [rng.standard_normal((d, d)) / np.sqrt(d) for _ in range(3)]
    Y, A = head(X, *W)
    assert np.allclose(A.sum(-1), 1.0), "rows must be a distribution"
    assert np.triu(A[0], 1).max() == 0.0, "no attention to the future"
    if show:
        print(f"  X {X.shape} -> Y {Y.shape}, attention {A.shape}")
        print(f"  rows sum to 1: max err {abs(A.sum(-1)-1).max():.2e}")
        print(f"  upper triangle exactly zero: {np.triu(A[0],1).max():.1f}")
        ent = lambda a: float(-(a*np.log(a+1e-30)).sum(-1).mean())
        _, Au = head(X, *W); su = softmax((X@W[0]) @ (X@W[1]).transpose(0,2,1))
        print(f"  entropy scaled {ent(A):.3f} vs unscaled {ent(su):.3f} (max {np.log(T):.3f})")
    return {"head": head}


# ─────────────────────────────────────────────────────────── block 5

Reading the implementation

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

Read it right to left. The output is \(\text{weights} \times V\) — an average of value vectors. The weights come from softmax(scores), so they are non-negative and sum to one. The scores are dot products between queries and keys, so the weights are large where the query and key agree. That is content-addressed lookup: a soft dictionary where the match is a similarity rather than an equality.

The scaling factor is the part worth deriving rather than memorising. If \(q\) and \(k\) have independent components with unit variance, then

\[ \mathrm{Var}(q!\cdot!k) = \sum_{i=1}^{d_k}\mathrm{Var}(q_ik_i) = d_k \]

so the scores have standard deviation \(\sqrt{d_k}\) — at \(d_k=64\), scores spread over ±8 before any training. Softmax of values that spread saturates: one weight goes to ~1, the rest to ~0, and the gradient through softmax (\(\text{diag}(p) - pp^{\top}\)) goes to zero. The model cannot learn because its attention is already maximally confident and wrong. Dividing by \(\sqrt{d_k}\) restores unit variance and keeps the distribution in the region where gradients exist (proofs.md P1).

What the numbers say

Output:

  X (1, 6, 16) -> Y (1, 6, 16), attention (1, 6, 6)
  rows sum to 1: max err 2.22e-16
  upper triangle exactly zero: 0.0
  entropy scaled 0.878 vs unscaled 0.124 (max 1.792)

The entropy column is the thing to watch: unscaled attention starts near- deterministic, scaled attention starts near-uniform. That difference is the difference between a model that trains and one that does not.

Beyond the toy

• Q, K, V are three linear projections of the same input in self-attention. Nothing requires that — cross-attention takes K and V from a different sequence, which is how encoder-decoder models and retrieval-augmented architectures work.
• The score matrix is the memory problem. \(T \times T\) per head; at \(T\)=8192 and 32 heads in fp16 that is 4 GB for one layer's scores. Never materialising it is FlashAttention's entire contribution, achieved by tiling into SRAM and using an online softmax that keeps a running max and normaliser.
• Attention is permutation-equivariant — it has no notion of order at all. Every positional scheme (learned, sinusoidal, RoPE, ALiBi) exists to repair that, and their differences are entirely about how they extrapolate beyond trained lengths.
• The softmax bottleneck. Because weights are non-negative and sum to one, attention can only produce outputs in the convex hull of the value vectors. That constraint is a real limit on expressiveness and motivates variants with gating or negative weights.

Block 5 — The leak test

Teaches: the single most valuable test in the project

The problem. The most valuable test in this project takes four lines and catches the failure that looks most like success. A causal model that can see the future achieves a loss below the entropy floor, produces beautiful training curves, and generates nothing.

@block(5, "The leak test", "the single most valuable test in the project")
def b5(s, show):
    head = s["head"]
    T, d = 10, 16
    X = rng.standard_normal((1, T, d))
    W = [rng.standard_normal((d, d)) / np.sqrt(d) for _ in range(3)]
    Y1, _ = head(X, *W)
    X2 = X.copy(); X2[:, 5:] = rng.standard_normal((1, T - 5, d))
    Y2, _ = head(X2, *W)
    leaked = not np.array_equal(Y1[:, :5], Y2[:, :5])
    assert not leaked, "CAUSAL LEAK"
    if show:
        print("  scrambled positions 5..9, then compared positions 0..4")
        print(f"  bit-identical: {np.array_equal(Y1[:,:5], Y2[:,:5])}   (leak: {leaked})")
        print("  catches: mask after softmax, off-by-one, accidental bidirectionality")
    return {}


# ─────────────────────────────────────────────────────────── block 6

Reading the implementation

The test perturbs a token at position \(j\) and asserts that outputs at positions \(< j\) are bit-identical. That is the precise operational meaning of causality: information may not flow backwards.

Why this beats inspecting the mask visually:

• It tests the composition of every layer, not one mask. A model can have a correct mask in layer 1 and a bug in a residual path, a cached key, or a normalisation that mixes across positions.
• It catches bugs in the inference path as well as training — a KV cache that recomputes positions differently is the classic one, and it only manifests when cached and uncached generation disagree.
• It requires no ground truth and no training, so it can run as a unit test in milliseconds on every commit.

-inf before the softmax rather than zeroing after is the correct implementation, because \(e^{-\infty}=0\) before normalisation. Zeroing afterwards leaves the denominator including masked positions, so the surviving weights no longer sum to one — a subtle scale error that trains, badly.

What the numbers say

Output:

  scrambled positions 5..9, then compared positions 0..4
  bit-identical: True   (leak: False)
  catches: mask after softmax, off-by-one, accidental bidirectionality

Beyond the toy

The general principle is worth more than the specific test: for every property your architecture is supposed to have, write the test that fails when it does not. Equivariances, invariances and information-flow constraints are all testable this way, cheaply, and they catch the class of bug that produces plausible-but-wrong models.

The production version of this bug is data leakage rather than mask leakage — a validation set that overlaps the training corpus, or a feature computed with future information (P08 measures exactly that). Same shape, same seductive symptom: results that are too good.

Block 6 — Multi-head

Teaches: H heads cost the same as one big head, and buy H views

The problem. One attention head computes one kind of similarity. Multi-head attention gets \(H\) of them for the same parameter count and the same FLOPs — which sounds like something for nothing and is worth understanding precisely.

@block(6, "Multi-head", "H heads cost the same as one big head, and buy H views")
def b6(s, show):
    softmax = s["softmax"]
    def mha(X, Wq, Wk, Wv, Wo, H):
        B, T, d = X.shape; dh = d // H
        def split(M): return M.reshape(B, T, H, dh).transpose(0, 2, 1, 3)
        Q, K, V = split(X @ Wq), split(X @ Wk), split(X @ Wv)
        sc = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(dh)
        sc = np.where(np.tril(np.ones((T, T), bool)), sc, -np.inf)
        out = softmax(sc) @ V                                  # (B,H,T,dh)
        return out.transpose(0, 2, 1, 3).reshape(B, T, d) @ Wo
    d, H = 32, 4
    params1 = 4 * d * d                     # one head of width d + output proj
    paramsH = 4 * d * d                     # H heads of width d/H + output proj
    if show:
        print(f"  d={d}, H={H}: params 1-head {params1:,} vs {H}-head {paramsH:,}")
        print(f"  identical. You get {H} attention distributions for free.")
        print(f"  each head sees a {d//H}-dim subspace -- below ~16 they get too small")
    return {"mha": mha}


# ─────────────────────────────────────────────────────────── block 7

Reading the implementation

Split \(d\) into \(H\) heads of width \(d/H\), attend independently, then concatenate and project. The parameter count is identical to a single head of width \(d\) — \(4d^2\) either way — because the projections are just reshaped. The FLOPs are the same too, except that the score matrices are \(H\) matrices of \(T\times T\) rather than one.

What you buy is subspaces. Each head computes similarity in a \(d/H\)- dimensional projection of the space, so different heads can attend on different relations — syntactic dependency, coreference, positional offset. One large head must express all relations in one similarity function and cannot separate them.

The cost is per-head dimension: at \(d\)=512 and \(H\)=64, each head has 8 dimensions, which is too few for a meaningful similarity. There is an interior optimum and it is empirical, typically \(d/H\) in the 64--128 range.

• The transpose(1, 2) / reshape pair is where implementations go wrong: the head dimension must be adjacent to the batch dimension for the batched matmul, and the inverse reshape must restore the exact original layout. A reshape that silently interleaves head outputs still trains — worse, and for no visible reason.

What the numbers say

Output:

  d=32, H=4: params 1-head 4,096 vs 4-head 4,096
  identical. You get 4 attention distributions for free.
  each head sees a 8-dim subspace -- below ~16 they get too small

Beyond the toy

• MQA and GQA break the symmetry deliberately: keep \(H\) query heads but share one (MQA) or a few (GQA) key/value heads. Quality is nearly unchanged and the KV cache shrinks by \(H\) or \(H/g\), which is the dominant memory cost at inference. This is the clearest case in the architecture of a change motivated purely by the serving cost model (P14).
• Head redundancy is measurable. Michel et al. found most heads can be pruned at little cost, and specialised heads (induction heads, previous-token heads) are identifiable and mechanistically interpretable. \(H\) is over-provisioned because which heads matter is not known in advance.
• Attention sinks. Trained models place large attention mass on the first token, apparently as a no-op destination when a head has nothing to attend to. Deleting that token during streaming inference breaks the model, which is why StreamingLLM keeps it pinned.

Block 7 — The pre-norm block

Teaches: residual as a gradient highway; norm before the sublayer

The problem. Two mechanisms let a deep stack train at all — the residual connection and the normalisation — and where the normalisation goes changes whether the model needs a warmup schedule to converge.

@block(7, "The pre-norm block", "residual as a gradient highway; norm before the sublayer")
def b7(s, show):
    mha = s["mha"]
    def layernorm(x, g, b, eps=1e-5):
        mu = x.mean(-1, keepdims=True); var = x.var(-1, keepdims=True)
        return g * (x - mu) / np.sqrt(var + eps) + b
    def ffn(x, W1, b1, W2, b2):
        h = x @ W1 + b1
        return np.maximum(h, 0) @ W2 + b2                      # ReLU, 4x expansion
    def tblock(x, P, H):
        x = x + mha(layernorm(x, P["g1"], P["b1"]), P["Wq"], P["Wk"], P["Wv"], P["Wo"], H)
        x = x + ffn(layernorm(x, P["g2"], P["b2"]), P["W1"], P["bb1"], P["W2"], P["bb2"])
        return x
    if show:
        print("  x = x + Attn(LN(x));  x = x + FFN(LN(x))")
        print("  PRE-norm: the residual path from output to input is unnormalised,")
        print("  so gradients reach layer 1 undiminished. Post-norm rescales on every")
        print("  layer and needs warmup to train deep.")
    return {"layernorm": layernorm, "ffn": ffn, "tblock": tblock}


# ─────────────────────────────────────────────────────────── block 8

Reading the implementation

x = x + sublayer(norm(x)) is pre-norm. The alternative, x = norm(x + sublayer(x)), is post-norm, and the difference is the gradient path.

In pre-norm, the residual stream is an unbroken identity path from input to output: the gradient reaches every layer undiminished, because \(\partial(x + f(x))/\partial x = I + \partial f/\partial x\). In post-norm, every layer's output passes through a normalisation, so gradients are rescaled at each of \(L\) steps and the product either shrinks or grows with depth. That is why the original Transformer needed a learning-rate warmup and why post-norm models past ~12 layers were unstable without one.

Pre-norm's cost is a slight quality reduction at matched compute, and a residual stream whose magnitude grows with depth — which is why pre-norm models add a final normalisation before the output projection.

What the numbers say

Output:

  x = x + Attn(LN(x));  x = x + FFN(LN(x))
  PRE-norm: the residual path from output to input is unnormalised,
  so gradients reach layer 1 undiminished. Post-norm rescales on every
  layer and needs warmup to train deep.

Beyond the toy

• RMSNorm drops the mean subtraction and the bias: \(x / \sqrt{\overline{x^2} + \epsilon}\). It is 10--15% cheaper in a memory-bound kernel with no measurable quality cost, which is why most models after ~2022 use it.
• The FFN is where the parameters are. With a 4× expansion, the two FFN matrices are \(8d^2\) against attention's \(4d^2\) — two thirds of every layer. Optimisation effort aimed at "attention" is usually aimed at a third of the model.
• SwiGLU replaces the ReLU FFN with a gated variant, using ~2.7× expansion to match parameter count. It costs a third matmul and consistently wins per parameter, which is a compute-for-quality trade rather than a free lunch.
• Normalisation placement interacts with precision. In bf16, post-norm's accumulated rescaling can overflow; the numerics and the architecture are not separable concerns.

Block 8 — The model

Teaches: stack the block; tie the unembedding; count the parameters

The problem. Assemble the pieces into something with a parameter count you can defend, and make the two decisions that dominate that count: how many layers, and whether to tie the unembedding.

@block(8, "The model", "stack the block; tie the unembedding; count the parameters")
def b8(s, show):
    tblock, layernorm = s["tblock"], s["layernorm"]
    V, d, H, L, T = s["tok"].size, 32, 4, 2, 16

    def init():
        P = {"emb": rng.standard_normal((V, d)) * 0.02,
             "pos": rng.standard_normal((T, d)) * 0.02,
             "gf": np.ones(d), "bf": np.zeros(d)}
        for l in range(L):
            for k, shape in (("Wq",(d,d)),("Wk",(d,d)),("Wv",(d,d)),("Wo",(d,d)),
                             ("W1",(d,4*d)),("W2",(4*d,d))):
                P[f"{l}.{k}"] = rng.standard_normal(shape) / np.sqrt(shape[0])
            P[f"{l}.bb1"] = np.zeros(4*d); P[f"{l}.bb2"] = np.zeros(d)
            P[f"{l}.g1"] = np.ones(d); P[f"{l}.b1"] = np.zeros(d)
            P[f"{l}.g2"] = np.ones(d); P[f"{l}.b2"] = np.zeros(d)
        return P

    def forward(P, idx):
        B, t = idx.shape
        x = P["emb"][idx] + P["pos"][:t]
        for l in range(L):
            sub = {k.split(".",1)[1]: v for k, v in P.items() if k.startswith(f"{l}.")}
            x = tblock(x, sub, H)
        x = layernorm(x, P["gf"], P["bf"])
        return x @ P["emb"].T                       # weight tying: unembed = emb^T
    P = init()
    n = sum(v.size for v in P.values())
    if show:
        print(f"  V={V} d={d} heads={H} layers={L} ctx={T}")
        print(f"  parameters: {n:,}  (embedding tied with unembedding)")
        lg = forward(P, s["data"][:16][None, :])
        print(f"  forward: idx (1,16) -> logits {lg.shape}")
        print(f"  initial loss {s['cross_entropy'](lg, s['data'][1:17][None,:]):.4f} "
              f"vs floor {s['floor']:.4f}  (untrained ~= floor, as it should be)")
    return {"init": init, "forward": forward, "V": V, "d": d, "H": H, "L": L, "T": T}


# ─────────────────────────────────────────────────────────── block 9

Reading the implementation

Per layer the parameter count is \(4d^2\) (attention projections) + \(8d^2\) (FFN) = \(12d^2\), so a model is approximately \(12Ld^2\) parameters plus embeddings. That formula is worth memorising: it lets you size a model, estimate its training FLOPs as \(6N\) per token, and price its inference at \(2N\), all without a framework.

Weight tying — using the embedding matrix as the output projection — saves \(Vd\) parameters. At \(V\)=32000 and \(d\)=4096 that is 131M, which at small scale is most of the model and at large scale is a rounding error. The justification is more than parameter economy: the input embedding and the output projection are both maps between token space and hidden space, and tying imposes that they be transposes of each other, which acts as a regulariser at small scale.

Initialisation scale is the other decision hiding here, and it is the one that silently ruins training runs. Weights scaled \(1/\sqrt{d}\) keep activation variance stable through a layer; too large and the residual stream explodes with depth, too small and the signal vanishes. The first version of this project's walkthrough omitted the scale and the demo disproved its own point — kept as a comment in the source because it is the most instructive kind of bug.

What the numbers say

Output:

  V=28 d=32 heads=4 layers=2 ctx=16
  parameters: 26,624  (embedding tied with unembedding)
  forward: idx (1,16) -> logits (1, 16, 28)
  initial loss 3.3313 vs floor 3.3322  (untrained ~= floor, as it should be)

Beyond the toy

• Depth versus width. At fixed parameter count, deeper models generally perform better up to a point, but depth is serial — it cannot be parallelised across devices as cleanly as width, and it increases the pipeline bubble in pipeline parallelism. The choice is as much a systems decision as a modelling one.
• Chinchilla scaling says compute-optimal training uses ~20 tokens per parameter, which repriced the entire field: most models of the GPT-3 era were badly under-trained for their size, and a smaller model trained longer wins at equal compute and is cheaper to serve forever after.
• Parameter count is the wrong metric for serving cost. What matters at inference is bytes read per token (memory-bound decode) and KV cache per sequence — both of which MoE and GQA decouple from parameter count entirely.

Block 9 — The same attention, in torch

Teaches: P01 permits torch as a TENSOR library -- not as an attention library

The problem. Everything above is numpy, which cannot compute a gradient. This block ports the same attention to torch as a tensor library — explicitly not using nn.MultiheadAttention or scaled_dot_product_attention — and proves the port is exact before trusting it.

@block(9, "The same attention, in torch", "P01 permits torch as a TENSOR library -- not as an attention library")
def b9(s, show):
    try:
        import torch, torch.nn.functional as F
    except ImportError:
        if show:
            print("  torch is not installed, so this block and the assembly are")
            print("  skipped. Blocks 1-8 above are pure numpy and told the whole")
            print("  story; this block only re-expresses block 6's attention in a")
            print("  framework that can differentiate it. Install with:")
            print("      pip install torch")
        return {}
    torch.manual_seed(0)

    def mha_t(x, Wq, Wk, Wv, Wo, H):
        """Byte-for-byte the same maths as block 6. Hand-written -- the whole point
        is that nn.MultiheadAttention is forbidden. torch supplies tensors and
        autograd; the mechanism is ours."""
        B, T, d = x.shape; dh = d // H
        sp = lambda M: (x @ M).view(B, T, H, dh).transpose(1, 2)
        q, k, v = sp(Wq), sp(Wk), sp(Wv)
        sc = q @ k.transpose(-2, -1) / (dh ** 0.5)
        sc = sc.masked_fill(~torch.tril(torch.ones(T, T, dtype=torch.bool)), float("-inf"))
        out = torch.softmax(sc, -1) @ v
        return (out.transpose(1, 2).reshape(B, T, d)) @ Wo

    # equivalence check against the numpy head from block 6
    B, T, d, H = 1, 6, 32, 4
    xs = rng.standard_normal((B, T, d))
    Ws = [rng.standard_normal((d, d)) / np.sqrt(d) for _ in range(4)]
    np_out = s["mha"](xs, *Ws, H)
    t_out = mha_t(torch.tensor(xs), *[torch.tensor(w) for w in Ws], H).detach().numpy()
    err = np.abs(np_out - t_out).max()
    assert err < 1e-10, f"torch and numpy disagree: {err}"
    if show:
        print(f"  numpy vs torch, same weights: max |diff| = {err:.2e}")
        print("  identical maths, and now differentiable. Forbidden here and in P01:")
        print("    nn.Transformer, nn.MultiheadAttention, F.scaled_dot_product_attention")
    return {"torch": torch, "mha_t": mha_t}


# ─────────────────────────────────────────────────────────── assembly

Reading the implementation

The port is line-for-line the same mathematics, and the assertion against the numpy version to \(<10^{-10}\) is what makes it safe to build on. That check is the same discipline as P13's reference comparison: a re-implementation is not correct because it looks correct, it is correct because it agrees with something already trusted.

masked_fill(~tril(...), -inf) is the causal mask again, now differentiable. Note that -inf interacts with autograd correctly here because softmax's gradient at a zero-probability position is zero — but a -inf that reaches a sum or a log elsewhere produces nan, which is the most common way masking breaks a backward pass.

What the numbers say

Output:

  numpy vs torch, same weights: max |diff| = 4.44e-16
  identical maths, and now differentiable. Forbidden here and in P01:
    nn.Transformer, nn.MultiheadAttention, F.scaled_dot_product_attention

Beyond the toy

The restriction — torch as tensors, not as transformers — is the pedagogical point of the whole project, and it maps onto a real engineering distinction. Using F.scaled_dot_product_attention in production is correct: it dispatches to FlashAttention, handles the memory layout, and is faster than anything hand- written. But it also means the \(T^2\) memory problem, the online softmax, and the numerical care around masking are invisible to you — and those are precisely the things you need to understand when the fused kernel does not support your variant, when a mask produces nan, or when you have to decide whether MQA is worth the quality risk.

Build it once to know what the library is doing; then use the library.

The assembly

Every block above, wired together into one working system:

def assembly(s):
    if "torch" not in s:
        print("\n  Skipped: the assembly trains the model, which needs torch.")
        print("  Run `pip install torch` and re-run to see it reach 0.106 nats.")
        return
    torch, mha_t, tok = s["torch"], s["mha_t"], s["tok"]
    V, d, H, L, T = s["V"], s["d"], s["H"], s["L"], s["T"]
    floor = s["floor"]
    data = torch.tensor(s["data"])

    print("\nEvery block, wired into a model that actually trains.\n")

    g = torch.Generator().manual_seed(0)
    def par(*shape, scale=None):
        t = torch.randn(*shape, generator=g) * (scale if scale else (shape[0] ** -0.5))
        return t.requires_grad_(True)

    P = {"emb": par(V, d, scale=0.02), "pos": par(T, d, scale=0.02)}
    for l in range(L):
        for k in ("Wq", "Wk", "Wv", "Wo"): P[f"{l}.{k}"] = par(d, d)
        P[f"{l}.W1"] = par(d, 4 * d); P[f"{l}.W2"] = par(4 * d, d)
    params = list(P.values())

    def ln(x):
        return (x - x.mean(-1, keepdim=True)) / (x.var(-1, keepdim=True, unbiased=False) + 1e-5).sqrt()

    def model(idx):
        x = P["emb"][idx] + P["pos"][: idx.shape[1]]
        for l in range(L):
            x = x + mha_t(ln(x), P[f"{l}.Wq"], P[f"{l}.Wk"], P[f"{l}.Wv"], P[f"{l}.Wo"], H)
            x = x + torch.relu(ln(x) @ P[f"{l}.W1"]) @ P[f"{l}.W2"]
        return ln(x) @ P["emb"].T                       # tied unembedding

    def batch(n, bs=16):
        ix = torch.randint(0, len(data) - T - 1, (bs,), generator=g)
        return (torch.stack([data[i:i+T] for i in ix]),
                torch.stack([data[i+1:i+T+1] for i in ix]))

    opt = torch.optim.AdamW(params, lr=3e-3)
    print(f"  {'step':>6}{'loss':>10}{'floor':>9}   {'note':<28}")
    for step in range(801):
        x, y = batch(step)
        loss = torch.nn.functional.cross_entropy(model(x).reshape(-1, V), y.reshape(-1))
        opt.zero_grad(); loss.backward(); opt.step()
        if step % 200 == 0 or step == 800:
            note = "at the floor -- untrained" if step == 0 else ""
            print(f"  {step:>6}{loss.item():>10.4f}{floor:>9.4f}   {note:<28}")
    final = loss.item()
    print(f"\n  final {final:.4f} nats, {floor - final:.4f} BELOW the entropy floor "
          f"({floor:.4f}).")
    print(f"  perplexity {np.exp(final):.2f} against a random-guess perplexity of {V}.")

    print("\n  E10 (deliberate overfit) -- the gate before any real training:")
    tiny_x, tiny_y = data[:T][None, :], data[1:T+1][None, :]
    for step in range(400):
        loss = torch.nn.functional.cross_entropy(model(tiny_x).reshape(-1, V), tiny_y.reshape(-1))
        opt.zero_grad(); loss.backward(); opt.step()
    print(f"    200-token slice driven to loss {loss.item():.4f}  "
          f"({'PASS' if loss.item() < 0.1 else 'still above 0.1'})")
    print("    a model that cannot do this has a bug, not a hard problem.")

    print("\n  Generation (greedy):")
    ctx = data[:6].tolist()
    for _ in range(40):
        ctx.append(int(model(torch.tensor(ctx[-T:])[None, :])[0, -1].argmax()))
    print(f"    {tok.decode(ctx)!r}")

    print("\n" + "─" * 74)
    print("  THE FULL PICTURE")
    print("─" * 74)
    rows = [(1,"tokenizer","chars <-> ids, round-trip tested"),
            (2,"batching","y is x shifted by one"),
            (3,"loss","cross-entropy, and the ln(V) entropy floor"),
            (4,"attention","QK^T/sqrt(dk), softmax, weighted V"),
            (5,"leak test","future tokens provably cannot influence the past"),
            (6,"multi-head","H views for the price of one"),
            (7,"pre-norm block","residual highway + LN before each sublayer"),
            (8,"the model","embedding -> L blocks -> tied unembedding"),
            (9,"torch port","same maths, now differentiable")]
    for n, name, what in rows:
        print(f"    block {n}  {name:<16} {what}")
    print(f"\n    assembly     {sum(p.numel() for p in params):,} parameters, "
          f"trained to {final:.3f} nats")
    print("\n  Next, on the project page: RoPE (m9), BPE (m8), KV cache (m11), and")
    print("  the ablations that turn this from a working model into a measured one.")

Output:

Every block, wired into a model that actually trains.

    step      loss    floor   note                        
       0    3.3408   3.3322   at the floor -- untrained   
     200    0.1810   3.3322                               
     400    0.1187   3.3322                               
     600    0.1169   3.3322                               
     800    0.1056   3.3322                               

  final 0.1056 nats, 3.2266 BELOW the entropy floor (3.3322).
  perplexity 1.11 against a random-guess perplexity of 28.

  E10 (deliberate overfit) -- the gate before any real training:
    200-token slice driven to loss 0.0017  (PASS)
    a model that cannot do this has a bug, not a hard problem.

  Generation (greedy):
    'the quick brown fox fox fox fox runs fons fox '

  THE FULL PICTURE
    block 1  tokenizer        chars <-> ids, round-trip tested
    block 2  batching         y is x shifted by one
    block 3  loss             cross-entropy, and the ln(V) entropy floor
    block 4  attention        QK^T/sqrt(dk), softmax, weighted V
    block 5  leak test        future tokens provably cannot influence the past
    block 6  multi-head       H views for the price of one
    block 7  pre-norm block   residual highway + LN before each sublayer
    block 8  the model        embedding -> L blocks -> tied unembedding
    block 9  torch port       same maths, now differentiable

    assembly     25,984 parameters, trained to 0.106 nats

  Next, on the project page: RoPE (m9), BPE (m8), KV cache (m11), and
  the ablations that turn this from a working model into a measured one.

The design space

The blocks above implement one point in a large space. Every choice below is a real fork taken by a real production model, and the reason for each is a cost, not a preference.

The arithmetic that drives all of it

For one layer, hidden size \(d\), sequence \(T\), the forward cost splits in two:

\[ \text{attention projections} = 4Td^2, \quad \text{attention scores+values} = 2T^2 d, \quad \text{FFN} = 8Td^2 \ (\text{ratio }4) \]

So attention's quadratic term overtakes the linear ones when \(2T^2 d > 12Td^2\), that is \(T > 6d\). For \(d\) = 4096 that is \(T\) ≈ 24{,}000 tokens. Below that, a transformer is a stack of matmuls and attention is a rounding error in FLOPs — which is why "attention is quadratic" is misleading advice at ordinary context lengths. What attention does dominate long before that is memory: the score matrix is \(T^2\) per head, and materialising it is why FlashAttention exists.

Training cost per token is ≈\(6N\) FLOPs for \(N\) parameters (2 forward, 4 backward); inference is ≈\(2N\). Those two constants let you price a training run on the back of an envelope, and they are the arithmetic behind the Chinchilla result that most models of that era were badly under-trained for their size.

Latency, bandwidth and the memory hierarchy

Decode is the case that matters and it is memory-bound, not compute-bound. Generating one token requires reading every weight once:

No arithmetic optimisation moves those numbers, because the arithmetic intensity of a batch-1 decode step is ≈2 FLOP/byte against a ridge point near 295. This is the single most important fact about LLM serving and it falls straight out of P14's roofline. The levers are all bytes: quantise the weights (int8/int4/fp8), share KV heads (GQA/MQA), or amortise the read over more sequences (batching, which is why continuous batching and paged attention exist).

The KV cache is the other memory consumer, and it grows with traffic rather than model size:

\[ \text{KV bytes} = 2 \times L \times T \times d_{kv} \times \text{batch} \times \text{bytes/elem} \]

For a 70B-class model (80 layers, \(d_{kv}\)=8192 with MHA) at 4k context, that is ~10 GB per sequence. GQA with 8 KV heads out of 64 divides it by 8. This is why vLLM's paged attention — allocating KV in fixed blocks like OS pages, from P12 — was such a large practical win: it removed the internal fragmentation from over-provisioning contiguous per-sequence buffers.

FlashAttention, and why it is an IO algorithm

The naive attention kernel writes the \(T \times T\) score matrix to HBM, reads it back for the softmax, writes it again, reads it for the value multiply. That is \(O(T^2)\) HBM traffic for \(O(T^2 d)\) work — intensity \(O(d)\), but with a constant that puts it below the ridge. FlashAttention tiles Q, K and V into SRAM (192 KB per SM on A100) and never materialises the full matrix, using the online softmax trick to keep a running max and normaliser so a streaming computation is numerically identical to the batched one. Same FLOPs, an order of magnitude less HBM traffic, 2--4× faster in practice. It is the clearest example in modern ML of an algorithm that is faster without doing less arithmetic — exactly the lesson P14 block 3 measures on a smaller scale.

Hardware: CPU, GPU, TPU

The systolic array is worth understanding because it explains TPU's shape preferences. Data flows through a 128×128 grid of MACs; each value read from memory is reused 128 times inside the array before leaving. That is the hardware expression of the same reuse argument as cache tiling (P14 block 4) — see proofs.md P15. It also means a matmul whose dimensions are not multiples of 128 wastes the array, which is why TPU-targeted models pad aggressively and why "make the hidden size a nice number" is real advice.

Note the ridge points: every accelerator generation has made the memory wall worse, because FLOP/s has grown faster than bandwidth. A kernel that was compute-bound on a V100 can be memory-bound on an H100 without a line changing.

Advanced algorithms and alternatives

• Sub-quadratic attention. Linear attention (Katharopoulos et al.) rewrites softmax attention as a kernel feature map so the computation associates as \((\phi(Q)(\phi(K)^\top V))\), turning \(O(T^2d)\) into \(O(Td^2)\). Performer approximates the softmax kernel with random features. Both lose quality on recall-heavy tasks, which is the empirical finding that keeps full attention alive.
• State-space models. S4/Mamba replace attention with a linear recurrence that has a convolutional parallel form for training and a constant-state recurrent form for inference — so decode needs no KV cache at all. The trade is a fixed-size state versus attention's perfect recall over the context.
• Sparse attention. Longformer/BigBird use local windows plus a few global tokens; the theory is that the resulting attention graph is an expander, so information still mixes in \(O(\log T)\) hops. Hardware efficiency is the practical problem: unstructured sparsity maps badly to tensor cores.
• Speculative decoding. A small draft model proposes \(k\) tokens, the large model verifies them in one forward pass. It converts \(k\) memory-bound decode steps into one compute-bound one — a pure exploitation of the intensity gap in the table above, with an acceptance-rate-dependent speedup and identical output distribution.
• MoE. Mixture-of-experts decouples parameters from FLOPs per token: only the top-\(k\) experts run. It trades a much larger memory footprint and an all-to-all communication pattern (P06's shuffle, at NVLink speed) for constant compute.

How this connects to the rest of the track

• P13 implements the autodiff this model trains under; block 9 here is the same attention expressed in a framework that can differentiate it.
• P14 explains why decode is memory-bound and prefill is not — the same intensity arithmetic, generalised.
• P02 uses the embeddings a model like this produces; attention itself is a soft nearest-neighbour lookup with learned keys, and both hit the same contrast problem in high dimensions.
• P12's paging is the direct ancestor of paged attention.
• P06's all-to-all shuffle is the same communication pattern as expert-parallel MoE routing and tensor-parallel all-reduce.

Failure modes at scale

• Loss spikes from a single bad batch or fp16 overflow in attention logits; the standard mitigations are z-loss, query/key normalisation, and skipping batches whose gradient norm exceeds a threshold.
• Silent leakage from an off-by-one in the causal mask. Loss falls below the data's entropy floor, which looks like a triumph. The assembly's floor comparison is the cheap detector.
• Divergence between train and inference paths — a KV cache that recomputes positions differently from training. Assert that cached and uncached generation produce identical logits.
• Throughput collapse from ragged batches: padding to the longest sequence in the batch wastes compute proportional to the length variance, which is why serving systems sort by length or use continuous batching.

Primary sources

• Vaswani et al., Attention Is All You Need (2017) — the architecture.
• Dao et al., FlashAttention (2022) and FlashAttention-2 (2023) — the IO argument in full.
• Shazeer, Fast Transformer Decoding: One Write-Head Is All You Need (2019) — MQA, and the KV-cache reasoning above.
• Ainslie et al., GQA (2023) — the interpolation that most models now use.
• Hoffmann et al., Training Compute-Optimal Large Language Models (2022) — the \(6N\) arithmetic applied to a scaling law.
• Kwon et al., Efficient Memory Management for LLM Serving with PagedAttention (2023) — P12's idea, in a serving stack.
• Su et al., RoFormer (2021) — RoPE.

Running it

python3 handson/h01_transformer.py            # every block, then the assembly
python3 handson/h01_transformer.py --block 3  # just block 3 and its prerequisites
python3 handson/h01_transformer.py --quiet    # the assembly only

What to do with this

Run it, then break it. Delete the / sqrt(dk) and watch block 3's entropy collapse. Remove the causal mask and watch the loss drop below the entropy of the data --- the classic leakage bug, which looks like brilliant training. Change the initialisation scale and see how far it moves the trainable range. Each of those is a two-character edit with a visible, measurable consequence, which is the fastest way to build the intuition that reading the paper alone will not.

Milestones, experiments, readings and exit criteria for this project: P01 — Transformer From Scratch.