TRANSFORMER FUNDAMENTALS

Text becomes integers before it becomes anything else. A tokenizer (BPE for most modern models) maps a string to a list of integers in [0, vocab_size). Llama 3 uses a 128,256-entry vocabulary; GPT-2 used 50,257. Nothing about this step is learned during pretraining — the merges were fit once on a corpus and frozen.

Those integers index into an embedding table of shape [vocab_size, d_model]. For Llama-3-8B that is [128256, 4096] = 525 million parameters, just to turn integers into vectors. The lookup produces a tensor of shape [batch, seq, d_model].

That tensor is the residual stream, and it is the single most useful mental object in the architecture. Its width never changes. Every block in the network reads the stream, computes something, and adds the result back:

x = x + attention(norm(x))
x = x + mlp(norm(x))

Note the +. Layers do not transform the stream, they contribute to it. This is why you can delete a layer from a trained transformer and often get degraded-but-coherent output: you removed one contribution from a sum of 32, not a link in a chain.

For inference, the shape [batch, seq, d_model] is where your intuition should live. During prefill seq is the whole prompt — maybe 2,000. During decode seq is 1. Same weights, same code path, wildly different arithmetic. That gap is the entire subject of Module 4.

A "hidden state" is just one slice of this stream: a single d_model-length vector at one position, at one layer. When people say a model "represents" something, they mean a direction in this 4096-dimensional space.

Every position produces three vectors from its residual-stream slice, via three learned projections:

Q = x @ W_q     "what am I looking for?"
K = x @ W_k     "what do I offer?"
V = x @ W_v     "what do I pass on if selected?"

The computation is four steps, and the shapes matter more than the formula:

scores  = Q @ K.T                    [S, S]     every query vs every key
scores  = scores / sqrt(head_dim)    [S, S]     keep logits out of softmax saturation
scores  = scores + causal_mask       [S, S]     -inf above the diagonal
weights = softmax(scores, dim=-1)    [S, S]     each ROW sums to 1
out     = weights @ V                [S, d_h]   weighted average of value vectors

Read row i of weights: it is a probability distribution over positions 0..i saying how much position i cares about each earlier position. The output at i is that distribution applied to the value vectors. Attention is retrieval by soft lookup — dot-product similarity picks what to read, and V is what gets read.

The 1/sqrt(head_dim) is not cosmetic. If Q and K entries are roughly unit-variance and independent, their dot product over head_dim=128 terms has variance ~128, standard deviation ~11.3. Feed logits of that magnitude into a softmax and you get a near-one-hot distribution with vanishing gradients. Dividing by sqrt(128) = 11.31 puts the variance back at 1.

The inference-critical observation is asymmetry. Q comes only from the current token. K and V come from every token so far. When you generate token 501, you need a fresh Q for position 501, but you need the same K and V for positions 0–500 that you computed on the previous step. They have not changed and they never will — this is precisely the fact that Module 2 exploits.

The causal mask sets scores[i][j] = -inf for all j > i, so softmax gives those entries exactly zero weight. Implementations add a matrix of 0 and -inf (or -1e9) rather than slicing, because a dense add is faster on a GPU than irregular indexing.

The training reason for this is the obvious one: it lets you compute the loss for all S positions in a single forward pass without any position cheating by looking at its own answer. That is teacher forcing, and it is why training is so much more hardware-efficient than generation.

The inference reason is the one people skip, and it is the load-bearing fact of this entire course. Causality means the representation of token i cannot be influenced by anything that comes after it. So when you append token 501 to the sequence:

• the K and V vectors for positions 0–500 are bit-for-bit identical to what they were before;
• the only new work is computing K, V, and Q for position 501;
• everything else can be read from storage.

Without causality this would be false. In a bidirectional encoder like BERT, adding a token at the end changes the representation of every earlier token, so there is nothing to cache. That is the real reason decoder-only models won for generation: not that they are more expressive, but that the causal constraint turns an O(n²) recomputation into an O(n) append.

Worth being precise about what is not cached. Q is not cached, because the query for a position is used once, at the step where that position is the newest token, and then never again. Attention outputs are not cached either. Only K and V — hence "KV cache" and not "QKV cache."

One attention pattern per layer is too few — a token usually needs to track several relationships at once. So the d_model vector is split into n_heads independent head_dim-sized subspaces, attention runs in each, and the results are concatenated.

The efficiency point that trips people up: heads are not implemented as a loop. The projection W_q is a single [4096, 4096] matrix. You do one matmul, then reshape and transpose the result into [B, n_heads, S, head_dim]. The heads only exist as a view.

Here is the full trace for Llama-3-8B (d_model=4096, n_heads=32, n_kv_heads=8, head_dim=128, ffn=14336), with batch B and sequence S:

Notice k_proj and v_proj output 1024, not 4096. That is grouped-query attention: 8 KV heads instead of 32, each shared by 4 query heads. It shrinks the KV cache by 4× on this model and is the single most consequential architectural decision for inference cost. Module 6 covers what it trades away.

Notice also [B, 32, S, S]. At S = 8192 and B = 1 in fp16 that one intermediate is 32 × 8192 × 8192 × 2 = 4.3 GB — for a single layer, for a single sequence. Never materialising that matrix is what FlashAttention is for (Module 7).

After attention, each position independently passes through a feed-forward network. "Independently" is exact — there is no mixing across positions here, which makes it trivially parallel and, at decode time, a plain matrix-vector product.

The classic transformer used W_2 @ relu(W_1 @ x) with an expansion factor of 4. Llama-family models use SwiGLU, which has three matrices instead of two:

gate = x @ W_gate        [B, S, 14336]
up   = x @ W_up          [B, S, 14336]
h    = silu(gate) * up   [B, S, 14336]     elementwise
out  = h @ W_down        [B, S, 4096]

where silu(z) = z * sigmoid(z). The gate branch acts as a learned, input-dependent filter on the up branch. Because there are three matrices, the hidden width is chosen smaller than 4× to keep the parameter count comparable — 14336 here, which is 3.5× d_model.

Now count parameters per layer for Llama-3-8B:

attention:  W_q  4096 x 4096  = 16,777,216
            W_k  4096 x 1024  =  4,194,304
            W_v  4096 x 1024  =  4,194,304
            W_o  4096 x 4096  = 16,777,216
                               ------------
                                41,943,040

mlp:        W_gate 4096 x 14336 = 58,720,256
            W_up   4096 x 14336 = 58,720,256
            W_down 14336 x 4096 = 58,720,256
                                 ------------
                                 176,160,768

The MLP is 4.2× the attention block. Across 32 layers: 1.34 B attention parameters, 5.64 B MLP parameters. Add the 525 M embedding and 525 M unembedding and you get 8.03 B — the "8B" on the label.

For inference this ranking matters. People assume attention dominates because it is the interesting part and because of the O(S²) term. At decode time with a modest context, it does not: you move 176 MB of MLP weights per layer versus 42 MB of attention weights, and the attention arithmetic touches a KV cache that is usually far smaller than either.

The original transformer put LayerNorm after each block: x = norm(x + attn(x)). This is post-norm, and it is hard to train deep — the residual path is interrupted by a normalization at every layer, so gradients get rescaled 2L times on the way back.

Every modern LLM uses pre-norm: x = x + attn(norm(x)). The normalization is applied to the input of the block; the residual stream itself is never touched. There is now an unbroken identity path from the final layer to the embedding, which is why 80-layer models train at all. A single final norm is applied before the LM head.

Most current models also replace LayerNorm with RMSNorm:

LayerNorm:  y = g * (x - mean(x)) / sqrt(var(x) + eps) + b
RMSNorm:    y = g * x / sqrt(mean(x^2) + eps)

RMSNorm drops the mean subtraction and the bias. It works about as well and is measurably cheaper — one pass over the vector instead of two, and one fewer parameter tensor.

Why an inference course cares: normalization layers are memory-bound with essentially zero arithmetic intensity. An RMSNorm over a [1, 4096] decode vector reads 4096 values, does a handful of FLOPs each, and writes 4096 values back. If you launch it as its own kernel you pay a full round trip to HBM for nothing. This is the canonical case for kernel fusion — folding the norm into the matmul that follows it so the data never leaves on-chip memory. Modules 4 and 7 return to this. A transformer at decode time is a long chain of these tiny bandwidth-bound operations glued between big matmuls, and the glue is not free.

Attention as defined is permutation-invariant — shuffle the tokens and the scores follow them unchanged. Position has to be injected explicitly.

The original approach added a learned or sinusoidal position vector to the embedding. Almost every current model instead uses RoPE (Rotary Position Embedding), which does something cleverer: it rotates Q and K, in-place, by an angle that depends on position.

Split each head_dim=128 vector into 64 pairs of coordinates. Treat each pair as a point in 2D and rotate it by m * theta_i, where m is the absolute position and theta_i = base^(-2i/head_dim) with base typically 10000 (Llama 3 uses 500000 to extend context). Low-index pairs rotate fast, high-index pairs rotate slowly — a positional code across many frequencies.

The property that makes it work falls out of rotation algebra: for a rotation R,

(R_m q) · (R_n k) = q · R_(n-m) k

The score between positions m and n depends only on n - m. You get relative position from an operation applied using absolute position. No extra parameters, nothing added to the residual stream.

Three inference consequences:

1. RoPE is applied to Q and K before they are cached. The cache stores rotated keys. Get the position index wrong on a cache hit and you silently corrupt the attention pattern — a classic prefix-caching bug (Module 6).
2. V is not rotated. Only Q and K participate in the dot product that needs position.
3. Context extension is a RoPE trick. Methods like position interpolation and YaRN rescale the frequencies so a model trained at 8k can attend over 128k. This is why long-context support often arrives without full retraining — and why long-context quality frequently lags long-context capability.

The final residual stream is normalized and projected to vocabulary size by the LM head, a [4096, 128256] matrix. Output shape: [B, S, 128256]. Those are logits — unnormalized scores, one per vocabulary entry.

That tensor is enormous relative to everything else. At B=1, S=4096 in fp32 it is 4096 × 128256 × 4 = 2.1 GB. Serving implementations therefore compute the LM head only for positions they actually need — during decode, that is the single last position, turning a [4096, 128256] matmul into a matrix-vector product.

Training runs one forward pass over the whole sequence, computes cross-entropy at every position against the next token, and updates. Every position is supervised in parallel. The GPU is doing S positions' worth of arithmetic per weight load, and arithmetic intensity is high.

Inference cannot do that, because position i+1's input is position i's output. You run the forward pass, take the logits at the last position, sample a token, append it, and run again. The loop is serial by construction.

Here is the asymmetry stated as a ratio. Generating 500 tokens from a 500-token prompt requires:

• Training-style work: one pass, 1000 positions, weights read once.
• Generation: one prefill pass over 500 positions, then 500 decode passes each producing one token. Weights read 501 times.

Same model, same total tokens, 500× the memory traffic. The prefill pass is doing hundreds of positions per weight load and the GPU is busy. Each decode pass does one position per weight load and the GPU is idle waiting on HBM.

That is the whole problem. Everything from here is an attempt to get more useful work out of each weight load.