S03The Model·251 lines

The Transformer Forward Pass

A modern decoder block is seven operations. Write them once in NumPy and every kernel optimisation later has a reference to check against.

RMSNorm
RoPE
grouped-query attention
SwiGLU
residual stream

Motivation

The problem

You are about to spend seventeen chapters making model.forward() faster. First you need a version that is unambiguously correct: something to diff a fused kernel against when the fused kernel produces plausible-looking garbage.

A modern decoder block is seven operations, and the entire Llama-family architecture is that block repeated. No encoder, no cross-attention, no learned position embeddings, no biases. What follows fits comfortably in 250 lines of NumPy.

Core idea

The solution

Here is the whole block. Read it once; the rest of the chapter is commentary.

one decoder blockpython

def block(x, w, pos, cache=None):
    # ---- attention sublayer -------------------------------------
    h = rms_norm(x, w.attn_norm)
    q = h @ w.wq                       # [T, n_heads * head_dim]
    k = h @ w.wk                       # [T, n_kv_heads * head_dim]
    v = h @ w.wv
    q, k = rope(q, pos), rope(k, pos)  # position enters here, only here
    a = grouped_query_attention(q, k, v, causal=True)
    x = x + a @ w.wo                   # residual

    # ---- feed-forward sublayer ----------------------------------
    h = rms_norm(x, w.ffn_norm)
    x = x + (silu(h @ w.w1) * (h @ w.w3)) @ w.w2   # SwiGLU
    return x

Four design choices in there are worth understanding, because each one changes what the engine has to do:

Pre-norm, not post-norm. The normalisation happens on the way into each sublayer, and the residual stream itself is never normalised. That makes deep stacks trainable. For us it means the residual stream is the block's only state; everything else is a pure function of it.
RMSNorm, not LayerNorm. No mean subtraction, no bias. One reduction, one multiply.
RoPE, not learned positions. Position is applied by rotating pairs of dimensions in q and k. Nothing is added to the residual stream, and attention sees only relative position, which is what context extension tricks like YaRN build on. RoPE on its own still degrades past its trained length; those tricks exist because it does.
SwiGLU, not GELU-MLP. Three matmuls instead of two: a gate, an up-projection, and a down-projection.

RoPE is applied to q and k, never v

A KV cache works because of this. Attention scores depend on the relative position of q and k, which rotation encodes. Values carry no positional information, so a cached v is valid forever, and caching K post-rotation means you never recompute that rotation again.

Mechanics

How it works

Diagramdecoder block anatomy

Chapters S05 through S08 are all about the dashed tap on the right: k and v are the only tensors worth keeping between forward passes.

Grouped-query attention, and why it exists

Full multi-head attention gives every query head its own key and value heads. The expense lands in exactly the wrong place: the KV cache scales with the number of KV heads, and cache size limits how many requests you can run at once.

GQA shares one KV head across a group of query heads. Llama-3-70B uses 64 query heads and 8 KV heads, an 8× smaller cache, at a quality cost small enough that everyone now does it. Multi-query attention (MQA) is the extreme case: one KV head total.

gqa is a repeat, not a rewritepython

def grouped_query_attention(q, k, v, n_heads, n_kv_heads, causal=True):
    head_dim = q.shape[-1] // n_heads
    q = q.reshape(-1, n_heads, head_dim)
    k = k.reshape(-1, n_kv_heads, head_dim)
    v = v.reshape(-1, n_kv_heads, head_dim)

    reps = n_heads // n_kv_heads        # e.g. 4 query heads per kv head
    k = np.repeat(k, reps, axis=1)      # broadcast, do not copy in a real kernel
    v = np.repeat(v, reps, axis=1)

    scores = np.einsum("qhd,khd->hqk", q, k) / np.sqrt(head_dim)
    if causal:
        scores += causal_mask(q.shape[0], k.shape[0])
    return np.einsum("hqk,khd->qhd", softmax(scores), v).reshape(-1, n_heads * head_dim)

np.repeat is a teaching device, not a kernel

Materialising the repeated K and V undoes the entire memory saving. Production kernels index the shared KV head directly from the attention inner loop: a stride calculation in FlashAttention, a block-table lookup in PagedAttention.

Where the FLOPs and the bytes go

Two numbers describe every operation in the block: how much arithmetic it does, and how many bytes it must move to do it. Their ratio is arithmetic intensity, and it decides whether the GPU is working or waiting.

During prefill, sequence length is large, weights are reused across hundreds of tokens, and intensity is high. During decode, sequence length is 1: you read the entire model to produce one token, and intensity collapses to about 1 FLOP per byte on hardware that wants 300. The simulator below makes this concrete.

Explore

Try it

Simulatorblock anatomy explorerclick an operation; drag seq len from 1 to 4096

d_model2048q heads32kv heads (GQA)8seq len1 (decode)layers22

params / block: 45.1M
params / model: 992.0M
head dim: 64
GQA ratio: 4:1

one decoder block · click an operation

Bars are multiply-accumulates per forward pass at the current sequence length. Total 90.2M per block, 2.0G per token for the whole model.

selected · rms_norm

out shape [1, 2048]

params 2.0k · flops 4.1k

No mean subtraction and no bias: RMSNorm is LayerNorm with the parts that did not matter removed. It is cheap, but it reads the whole residual stream, so it is bandwidth-bound.

grouped-query attention

Each shade is one KV head shared by 4 query heads. KV cache: 2048 B per token per layer, versus 8192 B for full multi-head, a 4.0× saving that costs almost no quality.

arithmetic intensity — why decode is bandwidth-bound

1.0seq 1

1.0seq 8

1.0seq 64

1.0seq 512

FLOPs per byte of weights read

An H100 needs roughly 300 FLOPs per byte to saturate its arithmetic units. At sequence length 1 you get about 1. The weights are read at full speed while the multipliers idle. That gap is why batching (S09) and speculative decoding (S14) are the two highest-leverage optimisations in the entire course.

Two experiments worth running. First, drag kv heads from 32 down to 1: the per-token cache cost falls by 32× while parameter count barely moves. That is GQA earning its keep. Second, drag seq len from 1 to 512 and watch arithmetic intensity climb. Every batching chapter later in this course is an attempt to get that intensity without making users wait.

Build it

Implementation

RoPE deserves its own look, because the implementation you write first is almost never the one the weights expect.

code/s03_transformer.py (excerpt)python

def rope(x, positions, theta=10000.0):
    """Rotate pairs of dimensions by an angle proportional to position."""
    T, D = x.shape
    half = D // 2

    # frequency per dimension pair: low dims rotate fast, high dims slow
    inv_freq = 1.0 / (theta ** (np.arange(0, half) * 2.0 / D))
    angles = positions[:, None] * inv_freq[None, :]     # [T, half]
    cos, sin = np.cos(angles), np.sin(angles)

    # NeoX layout: first half pairs with second half.
    x1, x2 = x[:, :half], x[:, half:]
    out = np.empty_like(x)
    out[:, :half] = x1 * cos - x2 * sin
    out[:, half:] = x2 * cos + x1 * sin
    return out

Two RoPE layouts exist and they are not compatible

The NeoX (half-split) layout pairs dimension i with i + D/2 (above). The GPT-J / interleaved layout pairs 2i with 2i+1. Both are called "RoPE". Pick the wrong one and the model still runs, still produces grammatical English, and is subtly incoherent. It is the single most demoralising bug in this course. HF-format Llama and Qwen checkpoints use NeoX (the converter permutes wq/wk to make it so); original Meta and GGUF Llama weights are interleaved. Check the config and the file format, do not guess.

The theta base is the other trap. Long-context models raise it (Llama-3 uses 500,000 instead of 10,000) to slow the rotation of the low-frequency dimensions, stretching their wavelengths so distant positions stay distinguishable. The highest-frequency pair does not depend on theta at all. Load a long-context model with the default base and quality falls off a cliff past a few thousand tokens.

Run it locally

Builds a small randomly-initialised Llama-style model, runs a forward pass, and prints the shape, mean and standard deviation of the residual stream after every sublayer. Includes a numerical check that GQA with n_kv_heads == n_heads matches plain multi-head attention exactly.

$ python code/s03_transformer.py

Expect: A per-layer activation table with stable statistics, an assertion pass on the GQA equivalence check, and a parameter-count breakdown showing the FFN dominating.

Only NumPy is required — setup instructions.

Production notes

In production

picoLM writes this block in ~340 lines of C with pre-computed RoPE sine/cosine tables. The trigonometry is hoisted out of the inner loop because sinf in a hot loop is slower than the matmul around it.
vLLM replaces every line here with a fused kernel: RMSNorm fused with the residual add, QKV as one packed matmul, RoPE fused into the attention prologue.
quant.cpp auto-detects architectural variants at load time (QK-norm in Qwen3, NeoX versus interleaved RoPE, dual-FFN) because the same GGUF loader has to serve seven architectures.

Exercises

1
Fuse q_proj, k_proj and v_proj into one matmul against a concatenated weight, then split the result. Confirm the output is bit-identical and measure the speedup.
2
Implement the interleaved RoPE layout alongside NeoX and write a test that shows they produce different attention scores for the same weights. Then you will recognise the failure when you hit it.
3
Add QK-norm (an RMSNorm applied to q and k before RoPE, as Qwen3 does) behind a config flag. Verify output is unchanged with the flag off — then flip it on with all-ones weights and observe that the output changes anyway, because RMSNorm rescales each vector by its own RMS.

Continue

What's next

The block ends with logits. S04 turns them into a token using temperature, top-k, top-p, min-p and repetition penalties, and shows why the sampler, which costs microseconds, is where most perceived quality problems live.

Check yourself

Questions

Answer before you look. A wrong answer tells you more than a right one, because the explanation points at the part of the chapter to reread.

Check yourselfQuestion 1 of 6

Why is RoPE applied to q and k but never to v?

Score 0/6