Handout — The Decoder-Only Transformer

The RNN handout ended with two structural complaints: the recurrence cannot be parallelized over time, and everything about the past has to fit through a fixed-size vector. Attention answers both with the same move — stop summarizing the past, and keep all of it addressable. This handout builds the decoder-only Transformer from that idea, then reads a production implementation of it line by line.

1From a summary to an addressable past

An RNN computes $h_t = f(h_{t-1}, x_t)$. Whatever position $t$ needs to know about position $t - 40$ has to have survived 40 lossy overwrites of a $d_h$-dimensional vector. The information may still be there; there is no mechanism for asking for it specifically.

Attention inverts the design. Keep a representation of every previous position, and let each position issue a query — "which earlier positions are relevant to me?" — answered by comparing that query against a key published by each position. The answer is a set of weights, used to average the values those positions also publish. Nothing is compressed away in advance, and the retrieval is learned.

Two consequences follow immediately, and they are exactly the RNN's two complaints, reversed. Any position can reach any other in one step — the path length between two tokens is $O(1)$ rather than $O(\text{distance})$, which is also why gradients no longer have to survive a long product of Jacobians. And all positions can be computed simultaneously, because position $t$'s output depends on the other positions' representations at the same layer, not on a state threaded through them. Training becomes a few large matrix multiplications, which is precisely what a GPU wants.

2Self-attention

2.1  Queries, keys, values

Let $X \in \mathbb{R}^{T \times d}$ hold one row per position. Three learned projections produce

and the output is

with $M$ the mask of §2.2. Read the shapes, because they carry the meaning: $QK^\top$ is $T \times T$ — one score for every ordered pair of positions — the softmax is taken over the last axis so each row is a distribution over positions, and multiplying by $V \in \mathbb{R}^{T \times d_v}$ produces $T \times d_v$, one output vector per position.

The same input supplies queries, keys and values, which is what "self" means. In a translation model's cross-attention, $Q$ comes from the decoder and $K, V$ from the encoder; the mechanism is identical.

Why $1/\sqrt{d_k}$

Not a tuning constant. If the entries of $q$ and $k$ are independent with unit variance, then $q \cdot k$ is a sum of $d_k$ such products and has variance $d_k$ — so its standard deviation grows as $\sqrt{d_k}$. Feed scores of that scale into a softmax and it saturates:

d = 64
q, k = torch.randn(10000, d), torch.randn(10000, d)
dots = (q * k).sum(-1)
print(dots.var().item(), (dots * d**-0.5).var().item())
# 64.03  ->  1.000        variance d, brought back to 1

# and what that does to a softmax over 8 positions:
# unscaled: max prob 0.9980, entropy 0.014
# scaled:   max prob 0.4125, entropy 1.708   (uniform would be 2.079)

An unscaled softmax puts 99.8% of its mass on one position before training has learned anything — and Handout #6 §5 Q2's other half applies: a saturated softmax has near-zero gradient, so the layer cannot learn its way out. Dividing by $\sqrt{d_k}$ keeps the scores $O(1)$ regardless of head dimension. It is a variance-control argument, of the same family as the initialization schemes.

2.2  The causal mask

A language model may not look ahead: $p(x_t \mid x_{<t})$ must not depend on $x_t$ or later. But $QK^\top$ computes all $T^2$ pairs, including future ones. The fix is to add $-\infty$ to the disallowed entries before the softmax, so they receive exactly zero weight:

$\exp(-\infty) = 0$, so masked positions contribute nothing and — importantly — the remaining weights still sum to 1, because the softmax normalizes over what is left. Adding a mask before the softmax is not the same as zeroing weights after it, and only the former gives a proper distribution.

2.3  Two ways to apply the mask

The equation says "add $M$", but in code there are two conventions, and you will meet both. An additive mask is a float tensor of $0$ and $-\infty$ that you add to the scores. A boolean mask is a tensor of True/False that you hand to masked_fill, which writes the sentinel in for you:

# additive: build the float tensor, then add
scores = q @ k.transpose(-2, -1) * d_k**-0.5
scores = scores + mask                       # mask: 0 where allowed, -inf where not

# boolean: keep a predicate, let masked_fill choose the sentinel
scores = q @ k.transpose(-2, -1) * d_k**-0.5
scores = scores.masked_fill(blocked, torch.finfo(scores.dtype).min)

They compute the same thing, and the difference is engineering, not mathematics. Both are common: nn.MultiheadAttention and F.scaled_dot_product_attention accept either form, while Hugging Face passes an additive mask into attention nearly everywhere — which is why §6.4's line reads attn_weights + attention_mask.

The memory column matters less than it looks, because masks are almost always broadcast from $[1, 1, T, T]$ rather than materialized per batch and head. The dtype column matters more. An additive mask has to be built in the score dtype: a fp32 mask added to fp16 scores silently upcasts the whole $T \times T$ score tensor to fp32, and casting finfo(float32).min down to fp16 overflows to -inf. A boolean mask has no dtype to get wrong.

So: purely structural masking — causal, padding — is a little cleaner and cheaper as boolean. The moment any continuous bias enters, build one additive tensor and add everything into it, because that is the only representation the soft terms fit in. Real libraries do both: Hugging Face's masking_utils.py composes causal, padding and sliding-window masks as boolean predicate functions (and_masks, or_masks) and then, in eager_mask, converts once at the boundary with torch.where(mask, 0.0, torch.finfo(dtype).min) — predicates where composition happens, floats where arithmetic happens. And on modern PyTorch, prefer handing the mask to F.scaled_dot_product_attention in either form: it applies the sentinel correctly and dispatches to a fused kernel that never materializes the $T \times T$ scores at all (§2.5).

2.4  Multiple heads

One attention operation produces one weighted average per position — one "kind" of lookup. Real text needs several at once (the syntactic subject, the coreferent, the nearby modifier). So run $h$ attention operations in parallel on $d/h$-dimensional slices, concatenate, and project:

Note the parameter budget is unchanged: $h$ heads of width $d/h$ have the same total projection size as one head of width $d$. Heads are close to free, and they are the reason attention maps are interpretable per head rather than in aggregate.

2.5  What it costs

The $T^2$ terms are the whole story of Transformer engineering. Note which resource each one costs: the $O(T^2 d)$ compute is often not the binding constraint at moderate $T$, because $Td^2$ dominates while $T < d$; the $O(T^2)$ memory for the score matrix is what actually breaks first, and that is what FlashAttention removes — it computes the same result without ever materializing the $T \times T$ matrix. And Handout #1 §1.4's big-O exercises on $Ax$ and $AB$ were preparation for exactly this table.

3The block

Attention alone is a weighted average — it mixes information across positions but applies no per-position nonlinear computation. A Transformer layer therefore alternates two sublayers, and this alternation is the actual architecture:

1. Attention — mixes across positions. This is the only place in the whole model where positions talk to each other.
2. A position-wise MLP — the same small network applied independently to every position, which does the per-token computation. It is typically $4\times$ wider than $d$ and holds roughly two thirds of the model's parameters.

Each is wrapped in a residual connection and a normalization. The residual is not decoration: $y = x + f(x)$ means $\partial y/\partial x = I + \partial f/\partial x$, so there is an identity path along which gradients reach the bottom of a 100-layer stack undiminished. That is the RNN handout §3.3's additive-path idea, transplanted from depth-in-time to depth-in-layers.

Pre-norm versus post-norm. The original paper normalized after the residual add, $\operatorname{Norm}(x + f(x))$. Every large model now normalizes before the sublayer, $x + f(\operatorname{Norm}(x))$, because that keeps the residual path free of any normalization at all — so the identity path is exactly the identity, and deep stacks train without a learning-rate warmup babysitting them. You will see this in the code as two norm calls placed at the top of each sublayer, and it is worth recognizing on sight since it is one of the few architectural details that visibly changed after 2017.

4Position information

Here is a fact worth sitting with: attention is permutation-equivariant. Shuffle the rows of $X$ and the outputs shuffle identically — the mechanism contains no notion of order. Order has to be injected, or the model is a bag of words.

The original solution added a fixed sinusoidal vector to each input embedding (you saw it in the RNN example's PositionalEncoding class). Two problems: it is absolute, so what the model learns about "three tokens back" at position 10 is not shared with position 1000, and it degrades badly past the lengths seen in training.

RoPE (rotary position embeddings) takes a different route. Rather than adding anything to the input, it rotates $q$ and $k$ by an angle proportional to their position, in each of $d/2$ two-dimensional subspaces at different frequencies. The point is what this does to the dot product — and it is checkable:

q, k = torch.randn(8), torch.randn(8)
for m, n in [(3, 1), (10, 8), (100, 98), (5, 3)]:
    cos, sin = rope_cos_sin(torch.tensor([m, n]), dim=8)
    q_m = apply(q, cos[0], sin[0])
    k_n = apply(k, cos[1], sin[1])
    print(m, n, (q_m * k_n).sum().item(), q_m.norm().item())

# 3   1   0.142312   1.821554
# 10  8   0.142312   1.821554
# 100 98  0.142312   1.821554
# 5   3   0.142312   1.821554

Four different absolute positions, one identical dot product — because in every case $m - n = 2$. Rotating $q$ by $m\theta$ and $k$ by $n\theta$ makes their inner product a function of $(m-n)$ alone, so relative position falls out of the geometry rather than being learned. And the norm is unchanged (1.8216 before and after), because a rotation is orthogonal — nothing is scaled, so nothing has to be renormalized.

One more consequence, visible in the same construction: because position enters only through $q$ and $k$, the rotation must be applied inside every attention layer rather than once at the input. §6.3 is where you will see that in the code.

5Inference, and the KV cache

The RNN handout §7.4 showed PyTorch's example regenerating the whole prefix for every new Transformer token — $O(T^2)$ total, and honest about it. Production code does not do that.

Observe: when you append token $T{+}1$, the keys and values of positions $1..T$ do not change. They are functions of those positions' representations, which the causal mask guarantees are unaffected by anything later. So cache them. Each new token computes one query, one key and one value, appends the latter two to the cache, and attends over $T{+}1$ cached keys. Per-token cost drops from $O(T^2)$ recompute to $O(T d)$, and total generation from $O(T^3)$ to $O(T^2)$.

The cache is not free, and its size is the dominant memory cost of serving long contexts: $2 \cdot L \cdot T \cdot n_{kv} \cdot d_{head}$ values. For a Llama-3-8B-shaped model ($L = 32$, $n_{kv} \cdot d_{head} = 8 \times 128$) in bf16:

That right-hand column is why grouped-query attention exists. GQA gives every query head its own $W_Q$ but shares one key/value pair across a group of them — here, 4 query heads per KV head — cutting the cache by exactly that factor with little measured quality loss. It is a change made for the memory bandwidth of inference, not for modeling reasons, and §6.4 is where you will find it as a single function call.

6Reading real code: Hugging Face's Llama

The other files in this series were small on purpose. This one is not: modeling_llama.py is ~500 lines inside a library of hundreds of thousands, with decorators, dispatch tables, type unions and deprecation shims. It is the most-read Transformer implementation in the world, and being able to work in a file like this is the actual professional skill.

6.1  LlamaRMSNorm — a simplification, not a new idea

modeling_llama.py · LlamaRMSNorm.forward (verbatim)

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    input_dtype = hidden_states.dtype
    hidden_states = hidden_states.to(torch.float32)
    variance = hidden_states.pow(2).mean(-1, keepdim=True)
    hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
    return self.weight * hidden_states.to(input_dtype)

Compare LayerNorm: subtract the mean, divide by the standard deviation, scale and shift. RMSNorm drops the mean subtraction and the bias — it only rescales each vector to unit root-mean-square. Verify that it really does not center:

x = torch.randn(4, 16) + 3.0        # deliberately off-center
rms(x).mean(-1)                     # [0.954, 0.945, 0.961, 0.938]  <- NOT zero
F.layer_norm(x, (16,)).mean(-1)     # [0.0, 0.0, 0.0, 0.0]
rms(x).pow(2).mean(-1).sqrt()       # [1.0, 1.0, 1.0, 1.0]          <- unit RMS

Fewer operations, fewer parameters, empirically as good — one of many places where the 2017 design turned out to have a load-bearing part and a decorative one.

The two dtype lines are the more instructive detail. The whole model may run in bfloat16, which has about three decimal digits of precision; squaring and summing 4096 such numbers is precisely where that runs out. So this function upcasts to float32 for the reduction and casts back at the end. Numerical-stability casts around reductions are everywhere in production ML code, they are never in the paper, and noticing them is a good sign you are reading rather than skimming.

6.2  LlamaDecoderLayer.forward — §3, exactly

modeling_llama.py · LlamaDecoderLayer.forward (verbatim; the self_attn call's argument list elided)

    residual = hidden_states
    hidden_states = self.input_layernorm(hidden_states)
    # Self Attention
    hidden_states, _ = self.self_attn(...)
    hidden_states = residual + hidden_states

    # Fully Connected
    residual = hidden_states
    hidden_states = self.post_attention_layernorm(hidden_states)
    hidden_states = self.mlp(hidden_states)
    hidden_states = residual + hidden_states
    return hidden_states

Ten lines for the entire block, and the structure is completely legible: stash the residual, normalize, apply the sublayer, add back. Twice. Note that this is pre-norm — input_layernorm is applied to the input of the sublayer, and nothing normalizes the residual stream itself, so residual + ... is an unobstructed identity path from the last layer to the first (§3).

The name post_attention_layernorm is a genuine trap: it is after the attention sublayer in sequence, but it is the pre-norm of the MLP sublayer. The name describes its position in the file, not its role in the architecture, and it has confused a lot of readers into thinking Llama is post-norm. Read the code, not the identifier.

6.3  LlamaAttention.forward — shapes, RoPE, cache, dispatch

modeling_llama.py · LlamaAttention.forward (verbatim; signature and the interface call's kwargs elided)

    input_shape = hidden_states.shape[:-1]
    hidden_shape = (*input_shape, -1, self.head_dim)

    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

    cos, sin = position_embeddings
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

    if past_key_values is not None:
        key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)

    attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
        self.config._attn_implementation, eager_attention_forward
    )

    attn_output, attn_weights = attention_interface(...)

    attn_output = attn_output.reshape(*input_shape, -1).contiguous()
    attn_output = self.o_proj(attn_output)
    return attn_output, attn_weights

Six things, and every one of them is a section of this handout appearing as one or two lines.

• The view().transpose(1, 2) idiom is multi-head attention (§2.4), and it is worth decoding once and for all. Input is (B, T, d). q_proj gives (B, T, h·d_head). The view splits that last axis into (B, T, h, d_head) — no data moves, it is a reinterpretation — and the transpose makes it (B, h, T, d_head), so that the last two axes are the $T \times d_{head}$ matrix each head attends over, and the batch and head axes ride along as batch dimensions of the matmul. The heads are never separate objects; they are a reshape. That is how "$h$ parallel attentions" becomes one batched matmul.
• RoPE is applied here, per layer, to $q$ and $k$ only — never to $v$, never to the residual stream — which is §4's last paragraph made concrete. If you were looking for where positions enter the model, this is the only place.
• The cache is one line. past_key_values.update(...) appends this step's keys and values and hands back the full history, so the very same forward serves training (no cache; $T$ positions at once) and generation (cache; one position at a time). §5's optimization required no separate code path — the only difference is whether an object was passed in.
• The dispatch table. ALL_ATTENTION_FUNCTIONS.get_interface(...) selects FlashAttention, SDPA, or the fallback eager_attention_forward depending on configuration. This is the seam where a mathematically identical operation is swapped for a faster implementation, and the reason the default has the name "eager": it is the readable one, kept precisely so there is something to read.
• The mask is not here. It arrives as an argument, prepared once for the whole model rather than rebuilt per layer. Worth noting because a first-time reader looking for §2.2's triangular matrix will not find it in the attention class.
• o_proj is the $W_O$ of §2.4, and the reshape before it is the concatenation of heads. Two lines for the second half of the multi-head formula.

6.4  eager_attention_forward — the formula itself

modeling_llama.py · eager_attention_forward (verbatim, body only)

    key_states = repeat_kv(key, module.num_key_value_groups)
    value_states = repeat_kv(value, module.num_key_value_groups)

    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
    if attention_mask is not None:
        attn_weights = attn_weights + attention_mask

    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
    attn_output = torch.matmul(attn_weights, value_states)
    attn_output = attn_output.transpose(1, 2).contiguous()

    return attn_output, attn_weights

This is §2.1's equation, and you should be able to point at each symbol: matmul(query, key.transpose(2,3)) is $QK^\top$, * scaling is $1/\sqrt{d_k}$ (set as self.scaling = self.head_dim**-0.5 in __init__), + attention_mask is $+M$ — additive, before the softmax, exactly as §2.2 insisted, and the additive-vs-boolean choice of §2.3 resolved this way for the whole library — and matmul(attn_weights, value_states) is the multiplication by $V$.

Two details the equation does not show. dtype=torch.float32 inside the softmax is another stability upcast, for the same reason as §6.1. And repeat_kv is grouped-query attention:

modeling_llama.py · repeat_kv (verbatim, docstring trimmed)

    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

GQA's saving is in the cache (§5): only $n_{kv}$ keys and values are stored. But the matmul needs one key per query head, so they are broadcast back up here — and note expand, which creates a view with stride 0 rather than copying, followed by a reshape that does have to materialize. The if n_rep == 1: return early exit means plain multi-head attention pays nothing for this function's existence. Small code, and it tells you where the designers expected the cost to be.

6.5  LlamaMLP — three matrices, not two

modeling_llama.py · LlamaMLP.forward (verbatim)

def forward(self, x):
    down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
    return down_proj

A textbook Transformer MLP is $W_2\,\sigma(W_1 x)$ — two matrices. This one has three, and the middle * is an elementwise product of two different projections of the same input. That is a gated linear unit (with SiLU as act_fn, making it SwiGLU), and it is the same gating idea as the RNN handout §3.3: one branch computes a value, the other computes a multiplicative, input-dependent gate on it.

The parameter arithmetic is worth doing, because it explains a config value that otherwise looks arbitrary. For Llama-3-8B, $d = 4096$ and $d_{ff} = 14336$: $3 \cdot d \cdot d_{ff} = 176{,}160{,}768$ parameters per layer, where a two-matrix MLP would use $2 \cdot d \cdot d_{ff} = 117{,}440{,}512$. Since $14336 \approx \tfrac{2}{3} \times 4 \times 4096$, the third matrix was paid for by narrowing $d_{ff}$ from the traditional $4d$ — keeping the parameter count roughly fixed and spending it on gating instead of width. That is why intermediate_size is a strange non-round number in these configs.

6.6  LlamaForCausalLM — where Handout #7 §2 reappears

modeling_llama.py · LlamaForCausalLM (verbatim, excerpts)

class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
    ...
        hidden_states = outputs.last_hidden_state
        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :])

        loss = None
        if labels is not None:
            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)

Three observations to close the loop with earlier handouts.

_tied_weights_keys is the RNN handout §7.2's self.decoder.weight = self.encoder.weight, expressed declaratively — same idea, same gradient accumulation, stated as configuration rather than as an assignment. With $|\Sigma| = 128256$ and $d = 4096$ that is 525M parameters shared rather than duplicated.

logits_to_keep is an optimization you can now derive yourself. During generation you need the distribution at the last position only, but lm_head projects to $|\Sigma| = 128{,}256$ — so computing it for all $T$ positions would produce a $T \times 128256$ tensor to throw away. Slicing first turns the single largest matmul in the model into a $1 \times d$ one. During training, logits_to_keep=0 keeps everything, because every position is a training example (§2.2's callout).

And loss_function(logits=..., labels=...) is Handout #7 §2 behind a library interface — including the shift. Hugging Face's convention is that you pass labels equal to input_ids and the shift is applied internally. If you shift them yourself as well, you get an off-by-two, a model that appears to train, and a mystery. Handout #7 §2.2 was about doing it by hand; this is the trap in its library form, and the two are the same trap.

6.7  Now do something to it

Show all answers

7Check yourself

Show all answers

8Optional further reading

Strictly optional — nothing here is assumed in class or on a quiz.

• Vaswani et al., Attention Is All You Need (2017) — read it after this handout, and notice how much of what you now consider standard (pre-norm, RMSNorm, RoPE, GQA, decoder-only) is not in it.
• Karpathy, nanoGPT/model.py — ~300 lines, the same architecture with none of the library machinery. The natural companion to §6: read it as the version of modeling_llama.py with the plumbing deleted.
• Alammar, The Illustrated Transformer — if §2 needed more pictures.
• Su et al., RoFormer (RoPE) and Ainslie et al., GQA — the two papers behind §4 and §5, both short and both readable now that you have seen the code.