Attention and the Transformer

By the end of this lesson you will be able to compute one head of attention by hand and check it against PyTorch; name the three projections behind the letters Q, K and V and find them in a checkpoint; say what a head is, why Qwen3-8B (Apache-2.0) has 32 query heads but only 8 key-value heads, and which of those two numbers sizes the cache; explain how position gets into a mechanism that is otherwise indifferent to order, with the rotation rates computed; count the parameters of the block that repeats thirty-six times; and say, with arithmetic from config.json, at what context length the quadratic term in attention overtakes everything else. This is the lesson the rest of the course leans on hardest: every memory budget, every long-context problem and every parallelism decision from Part 18 onwards traces back to it.

The problem attention solves

Before 2017 the dominant sequence models were recurrent: they read a sequence one step at a time, carrying a fixed-size state forward. Two things follow from that design. Information from far back has to survive many steps of being rewritten, and the steps cannot be computed in parallel because each depends on the one before.

The paper that changed this was “Attention Is All You Need” (Vaswani and colleagues, arXiv:1706.03762, submitted 12 June 2017). Its abstract proposes “a new simple network architecture, the Transformer, based solely on attention mechanisms, dispensing with recurrence and convolutions entirely”, and reports that the resulting models were “more parallelizable and requiring significantly less time to train”. That second property is why the architecture, rather than any particular model, is the thing that mattered: it fits the hardware from Part 1’s lesson on matrix multiplication.

The idea is direct. Instead of squeezing history through a fixed-size state, let every position look at every other position and take a weighted average of what it finds. The weights are computed from the content, so what a position looks at changes with what is there.

Queries, keys and values

Each position’s vector is multiplied by three learned matrices, giving a query (what this position is looking for), a key (what it offers to the others) and a value (what is passed along when the two match). One head of attention over n positions is then four lines of arithmetic:

q = x · W_q      k = x · W_k      v = x · W_v        one row per position, head_dim columns each
scores  = q · kᵀ / sqrt(head_dim)                    (n, n): row i is query i against every key
weights = softmax(scores + mask)                     one softmax per row, so each row sums to one
output  = weights · v                                (n, head_dim): a weighted average of the values

Three details in those lines carry the design. The division by sqrt(head_dim) is there because, in the paper’s words, “for large values of d_k, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients”; PyTorch’s scaled_dot_product_attention documents the same default, 1/sqrt(E). The mask is a matrix added to the scores before the softmax, and the Transformers documentation states the convention: a float mask holds 0.0 where a position may attend and -inf where it may not, because “a score plus -inf drops to zero after the softmax, so the position is excluded”. And the softmax runs one row at a time, so every position ends with weights that sum to one over the positions it was allowed to see.

Putting -inf above the diagonal is what makes a model causal: position seven attends to positions one to seven and nothing later. This is not a detail. It is what makes next-token prediction a valid training signal at every position at once, since no position can read its own answer. Here is the whole computation on four positions with an 8-wide head, then the same inputs handed to PyTorch’s implementation:

RunnableAll tracks

attention-by-hand.py - one causal head over four positions, then checked against PyTorch

"""One head of causal self-attention over four positions, by hand, then checked against PyTorch."""
import numpy as np
import torch
import torch.nn.functional as F

np.set_printoptions(precision=3, suppress=True, linewidth=120)
rng = np.random.default_rng(0)
n, d = 4, 8                                    # four positions; an 8-wide head (Qwen3 heads are 128 wide)
x = rng.standard_normal((n, d), dtype=np.float32)          # one input vector per position
W_q, W_k, W_v = (rng.standard_normal((d, d), dtype=np.float32) / np.sqrt(d) for _ in range(3))

q, k, v = x @ W_q, x @ W_k, x @ W_v            # three projections of the same input
scores = q @ k.T / np.sqrt(d)                  # (n, n): row i holds query i against every key
mask = np.triu(np.full((n, n), -np.inf, dtype=np.float32), k=1)   # -inf above the diagonal
masked = scores + mask                         # the float mask convention: 0 keeps, -inf removes
weights = np.exp(masked - masked.max(axis=1, keepdims=True))
weights /= weights.sum(axis=1, keepdims=True)  # softmax, one row at a time
out = weights @ v                              # (n, d): each row is a weighted average of the values

print("scores = q @ k.T / sqrt(d)\n", scores)
print("scores + causal mask\n", masked)
print("softmax per row = the attention weights\n", weights)
print("row sums:", weights.sum(axis=1), "  largest entry above the diagonal:", weights[np.triu_indices(n, 1)].max())
print("output = weights @ v, shape", out.shape)

ref = F.scaled_dot_product_attention(torch.from_numpy(q), torch.from_numpy(k), torch.from_numpy(v), is_causal=True)
print("max |difference| from torch scaled_dot_product_attention(is_causal=True):",
      f"{float((ref - torch.from_numpy(out)).abs().max()):.1e}")

Output — what you should see

attention-by-hand.py, numpy 2.5.2 and torch 2.14.0 on a CPU; the seed makes these exact numbers reproducible

scores = q @ k.T / sqrt(d)
 [[-0.145 -0.057  0.286  1.299]
 [ 0.66  -0.645 -0.332 -1.322]
 [-0.14   0.677 -0.659  0.497]
 [-0.058 -0.255 -1.244 -0.77 ]]
scores + causal mask
 [[-0.145   -inf   -inf   -inf]
 [ 0.66  -0.645   -inf   -inf]
 [-0.14   0.677 -0.659   -inf]
 [-0.058 -0.255 -1.244 -0.77 ]]
softmax per row = the attention weights
 [[1.    0.    0.    0.   ]
 [0.787 0.213 0.    0.   ]
 [0.259 0.587 0.154 0.   ]
 [0.382 0.314 0.117 0.187]]
row sums: [1. 1. 1. 1.]   largest entry above the diagonal: 0.0
output = weights @ v, shape (4, 8)
max |difference| from torch scaled_dot_product_attention(is_causal=True): 1.7e-16

Read the weights row by row. Position 0 can see only itself, so its weight is 1. Position 3 spread itself over all four positions and its row still sums to one. Every entry above the diagonal is exactly zero, not small, because exp(-inf) is zero. The last line shows that this is the same computation PyTorch’s scaled_dot_product_attention performs with is_causal=True, to within rounding: the optimised backends in the last section differ in how they compute it, never in what.

In a real model the three matrices are tensors you can list. Qwen3-8B’s first block holds these, read from the JSON header of its first safetensors shard (2026-09-12); the names are the same in every block and every Qwen3 size, only the shapes change:

Tensor	Shape, stored as (out, in)	What it is
model.layers.0.self_attn.q_proj.weight	[4096, 4096]	W_q for all 32 query heads at once: 32 × 128 output rows
model.layers.0.self_attn.k_proj.weight	[1024, 4096]	W_k for the 8 key-value heads: 8 × 128 rows
model.layers.0.self_attn.v_proj.weight	[1024, 4096]	W_v, the same shape as W_k
model.layers.0.self_attn.o_proj.weight	[4096, 4096]	Projects the concatenated head outputs back into the residual stream
model.layers.0.self_attn.q_norm.weight	[128]	A per-head RMS normalisation of the query, applied before the rotation below
model.layers.0.self_attn.k_norm.weight	[128]	The same for the key

The q_norm and k_norm rows are a Qwen3 choice you will meet when the lab lists tensors: the modelling code normalises each query and key head before anything else happens to it. There is no bias on any of these, attention_bias is false in the configuration, so attention in one block is four matrices plus two 128-entry scales, which the block section below adds up.

!Important

The attention weights are a matrix with one row and one column per position, and Transformers will hand it back when the eager implementation is in use. The attentions output is documented as a tuple with one entry per layer, each of shape (batch_size, num_heads, sequence_length, sequence_length), holding the “attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads”. The last two dimensions both being sequence_length is the whole cost story of this lesson, and the lab prints one such matrix for one head.

Many heads, and fewer key-value heads

One set of projections can express one notion of what to look for. The paper’s answer is to run several in parallel, which “allows the model to jointly attend to information from different representation subspaces at different positions”. Those are heads, and the mechanism is a reshape: q_proj produces 4,096 numbers per position, the code views them as 32 heads of 128, each head runs the four lines above on its own 128-wide slice, and the 32 outputs are concatenated back to 4,096 and multiplied by o_proj. head_dim is its own config.json field, not hidden_size divided by the head count: the modelling code builds q_proj as [num_attention_heads × head_dim, hidden_size] and o_proj as the transpose of that, so the heads’ combined width can differ from the residual stream’s. In Qwen3-8B and Qwen3-1.7B (also Apache-2.0) the product happens to equal hidden_size; in Qwen3-0.6B (Apache-2.0, the lab’s reduced-path model) and Qwen3-4B (Apache-2.0) attention projects the stream up to a wider space and back, which is what the lab’s tensor listing will show. From the four config.json files, read 2026-09-12:

Model	hidden_size	num_attention_heads × head_dim	q_proj shape	o_proj shape
Qwen3-0.6B	1,024	16 × 128 = 2,048	[2048, 1024]	[1024, 2048]
Qwen3-1.7B	2,048	16 × 128 = 2,048	[2048, 2048]	[2048, 2048]
Qwen3-4B	2,560	32 × 128 = 4,096	[4096, 2560]	[2560, 4096]
Qwen3-8B	4,096	32 × 128 = 4,096	[4096, 4096]	[4096, 4096]

The block-parameters script below reads head_dim for exactly this reason; run on Qwen3-0.6B’s config.json it prints 596,049,920 parameters, with attention at 40% of the block rather than Qwen3-8B’s 22%.

The key and value projections are narrower. num_key_value_heads is 8 in all four models, so k_proj and v_proj produce 8 heads of 128, and in Qwen3-8B the modelling code’s repeat_kv repeats each of them four times to serve the 32 query heads: query head i reads key-value head i // 4. That is grouped-query attention, from the paper of that name, which describes it as “a generalization of multi-query attention which uses an intermediate (more than one, less than number of query heads) number of key-value heads” and reports that it “achieves quality close to multi-head attention with comparable speed to MQA”. The reason it matters is not the parameter saving, which is modest; it is that keys and values are the two things stored for every token of context:

	Qwen3-8B as shipped	Qwen3-8B with 32 key-value heads	Qwen3-1.7B as shipped
Query heads × head_dim	32 × 128 = 4,096	32 × 128 = 4,096	16 × 128 = 2,048
Key-value heads × head_dim	8 × 128 = 1,024	32 × 128 = 4,096	8 × 128 = 1,024
Query heads sharing one key-value head	4	1	2
Attention parameters per block	41,943,296	67,109,120	12,583,168
Key + value numbers stored per token, per layer	2 × 1,024 = 2,048	2 × 4,096 = 8,192	2 × 1,024 = 2,048

That is arithmetic from config.json, not a measurement; the middle column is hypothetical. The last row is the one the rest of the course uses: the next lesson multiplies it by the layer count and the bytes per number to price a token of context, and Part 4 turns that price into which model fits your machine.

✕Common mistake

Reading num_attention_heads: 32 and assuming the cache holds 32 heads’ worth of keys and values per token is a common error in do-it-yourself memory arithmetic, and it overestimates by a factor of four on these models. The field that sizes the cache is num_key_value_heads.

Position has to be added deliberately

The dot products know nothing about where a key sits. The causal mask says which keys exist for a query, not how far away each one is, and the paper is explicit that something has to be added: “Since our model contains no recurrence and no convolution, in order for the model to make use of the order of the sequence, we must inject some information about the relative or absolute position of the tokens in the sequence.” The 2017 answer was a sinusoidal pattern added to the input embeddings.

The answer every model in this course uses is rotary position embedding, introduced in “RoFormer: Enhanced Transformer with Rotary Position Embedding” (Su and colleagues, arXiv:2104.09864). The abstract describes the method as one that “encodes the absolute position with a rotation matrix and meanwhile incorporates the explicit relative position dependency in self-attention formulation”, and lists the properties that follow: “the flexibility of sequence length, decaying inter-token dependency with increasing relative distances, and the capability of equipping the linear self-attention with relative position encoding”. The mechanism, as the Transformers code computes it:

inv_freq[i] = rope_theta ^ (−2i / head_dim)     i = 0 … 63 for a 128-wide head: 64 rates, in radians per token
angle[i] at position p = p × inv_freq[i]
pair i of a query, (q[i], q[i + 64]), is rotated by angle[i]; the key at its own position is rotated the same way

A rotation changes no lengths, and the dot product of two vectors rotated by angles a and b depends only on a − b. With a = m × inv_freq[i] for a query at position m and b = n × inv_freq[i] for a key at position n, each pair’s contribution to the score depends on (m − n) × inv_freq[i]: the distance, not the positions. Nothing is added to the values and nothing to the embeddings. The script computes the 64 rates for Qwen3-8B’s rope_theta of 1,000,000 and then tests the claim:

RunnableAll tracks

rope-relative.py - the rotation rates of one head, and the proof that a score depends on distance only

"""Rotary position embedding on one 128-wide head: the angle rates, and the proof that a score depends on distance only."""
import numpy as np

head_dim, rope_theta = 128, 1_000_000.0                   # Qwen3-8B's head_dim and rope_theta
pairs = np.arange(0, head_dim, 2) / head_dim              # 0, 2/128, 4/128, ... one entry per pair of dimensions
inv_freq = 1.0 / rope_theta ** pairs                      # radians per token for each pair, as transformers computes it

def rope(x, pos):
    """Rotate the 64 pairs of x by pos × inv_freq, in the rotate_half convention transformers uses."""
    angle = np.concatenate([pos * inv_freq, pos * inv_freq])
    x1, x2 = x[:head_dim // 2], x[head_dim // 2:]
    return x * np.cos(angle) + np.concatenate([-x2, x1]) * np.sin(angle)

for i in (0, 1, 16, 32, 48, 63):
    print(f"pair {i:2d}: {inv_freq[i]:.3e} rad per token   one full turn every {2 * np.pi / inv_freq[i]:>13,.0f} tokens")

rng = np.random.default_rng(0)
q, k = rng.standard_normal(head_dim), rng.standard_normal(head_dim)
print(f"\nunrotated q·k = {q @ k:.6f}")
for m, n in ((5, 2), (1005, 1002), (30005, 30002), (5, 3), (5, 0)):
    print(f"query at {m:5d}, key at {n:5d}, distance {m - n}:   rotated q·k = {rope(q, m) @ rope(k, n):.6f}")

Output — what you should see

rope-relative.py, numpy 2.5.2

pair  0: 1.000e+00 rad per token   one full turn every             6 tokens
pair  1: 8.058e-01 rad per token   one full turn every             8 tokens
pair 16: 3.162e-02 rad per token   one full turn every           199 tokens
pair 32: 1.000e-03 rad per token   one full turn every         6,283 tokens
pair 48: 3.162e-05 rad per token   one full turn every       198,692 tokens
pair 63: 1.241e-06 rad per token   one full turn every     5,063,256 tokens

unrotated q·k = 0.346338
query at     5, key at     2, distance 3:   rotated q·k = 4.196789
query at  1005, key at  1002, distance 3:   rotated q·k = 4.196789
query at 30005, key at 30002, distance 3:   rotated q·k = 4.196789
query at     5, key at     3, distance 2:   rotated q·k = 3.846083
query at     5, key at     0, distance 5:   rotated q·k = 3.428509

Two things to read off. First, the three distance-3 pairs give the same score to six decimal places whether the query sits at position 5, 1,005 or 30,005, and a different distance gives a different score: the rotation has turned absolute positions into relative ones. Second, the rates span six orders of magnitude. Pair 0 completes a turn every six tokens and resolves local order; pair 63 turns once every five million tokens and is nearly constant across any context the model will see, which is the abstract’s “decaying inter-token dependency with increasing relative distances” made concrete. rope_theta sets that spread, and it is why one head can express both “the previous token” and “somewhere in this document”.

Two configuration fields control this, and misreading them is a common long-context mistake in the course. Qwen3-8B declares rope_theta as 1000000 and rope_scaling as null: the model in its native state, which the card describes as supporting “context lengths of up to 32,768 tokens”. The same card describes reaching 131,072 tokens with YaRN by adding one block to config.json:

the rope_scaling block the Qwen3-8B model card gives for YaRN (read 2026-09-12)

{
  "rope_scaling": {
    "rope_type": "yarn",
    "factor": 4.0,
    "original_max_position_embeddings": 32768
  }
}

YaRN rescales the slow rates so that factor times the positions fit into the range of angles the model was trained on, and the card’s rule is that the factor follows the length you need, not the maximum: “It is also recommended to modify the factor as needed. For example, if the typical context length for your application is 65,536 tokens, it would be better to set factor as 2.0”, that is, factor = target length ÷ 32,768. The card gives the same setting as llama-server flags, --rope-scaling yarn --rope-scale 4 --yarn-orig-ctx 32768, where --rope-scale carries the factor and changes with it (2 for a 65,536-token application) while --yarn-orig-ctx stays at the trained 32,768; llama.cpp v0.4.0’s common/arg.cpp documents --rope-scale N as “RoPE context scaling factor, expands context by a factor of N”. The Part 6 lesson on llama-server is where the context flags get set. The card’s caveat is that every framework implements “static YaRN, which means the scaling factor remains constant regardless of input length, potentially impacting performance on shorter texts”, so the field is a decision:

Context length you will ask the engine for	rope_scaling	factor	What goes wrong otherwise
Up to 32,768 tokens	Leave it null, as shipped	none	The card: “If the average context length does not exceed 32,768 tokens, we do not recommend enabling YaRN in this scenario, as it may potentially degrade model performance”
32,769 to 65,536 tokens	Add the card’s block with the factor lowered, or --rope-scale 2	2.0	Positions past the trained range get angles the model never saw during training; the memory is spent on the context and what you lose is quality on the long inputs you allocated it for
65,537 to 131,072 tokens	Add the card’s block as printed, or --rope-scale 4	4.0	The same, over a wider range of positions
Above 131,072 tokens	Not described by the card	none given	Measure before trusting it

The Part 7 reality check measures what a context allocated without the scaling change actually recalls, and Part 24 budgets an agent’s context against these limits.

The block that repeats

A transformer is one block design, stacked. Qwen3-8B declares num_hidden_layers as 36, so this block occurs thirty-six times, each copy with its own weights.

One transformer block, as the Qwen3 models build it

1. 1NormaliseRMS normalisation of the incoming vector (input_layernorm). config.json gives rms_norm_eps as 1e-06.
2. 2Self-attention32 query heads in groups of 4 sharing 8 key-value heads, each 128 dimensions wide, with per-head query and key norms and rotary positions applied before the scores.
3. 3Add to the residual streamThe attention output is added to the block’s input, not substituted for it.
4. 4NormaliseA second RMS normalisation (post_attention_layernorm), before the feed-forward sub-layer.
5. 5Feed-forwardA gated network that widens each position from hidden_size 4096 to intermediate_size 12288 and back, with the SiLU activation named by hidden_act.
6. 6Add to the residual streamThe feed-forward output is added back. The result is the input to the next block.

Normalisation happens before each sub-layer, and each sub-layer's output is added back to the stream rather than replacing it. Thirty-six of these in sequence is the whole of Qwen3-8B between the embedding lookup and the final normalisation that feeds the output projection.

The modelling code’s forward pass for one block is four lines, and they name every tensor in the block:

Pseudocode — not a real command

h = x + attention(rmsnorm_1(x))          # input_layernorm, then self_attn, then add to the stream
y = h + mlp(rmsnorm_2(h))                # post_attention_layernorm, then mlp, then add again
mlp(z)     = down_proj( silu(gate_proj(z)) × up_proj(z) )     # the gated feed-forward network
rmsnorm(z) = z / sqrt(mean(z²) + rms_norm_eps) × g            # g is the learned scale, one per dimension

The residual stream is the x + and the h +: each sub-layer adds to its input rather than replacing it, so there is a path from the embedding to the output that passes through no matrix at all. Part 1’s lesson on why depth needs residual connections measured what that path does for the gradient; the mental model it left is the one to keep: a channel that each block reads from and writes an increment into, so that depth is a sequence of small edits rather than a chain of transformations.

Normalisation is rmsnorm, placed before each sub-layer rather than after it as the 2017 paper did (“a residual connection around each of the two sub-layers, followed by layer normalization”). Part 1 gave the formula. Two things to add: rms_norm_eps, 1e-06 here, is the constant that stops a division by zero, and the modelling code computes the mean of squares in FP32 whatever the weights’ precision, one of the places a BF16 model is quietly more precise than its file suggests.

The feed-forward network processes each position on its own, “applied to each position separately and identically” in the paper’s words, and it is where the information attention gathered gets transformed. The paper’s version was two matrices with a ReLU between them. Qwen3 uses the gated form from “GLU Variants Improve Transformer”, whose abstract describes a gated linear unit as “the component-wise product of two linear projections, one of which is first passed through” a nonlinearity, here the SiLU named by hidden_act. That is three matrices rather than two, gate_proj, up_proj and down_proj, each 4096 by 12288, and that is where the parameters are. Count them:

RunnableAll tracks

block-parameters.py - count one block from config.json fields, then the whole model

"""Count one transformer block's parameters from config.json fields, then the whole model. Pass a config.json path to use another model."""
import json
import sys

cfg = {"hidden_size": 4096, "intermediate_size": 12288, "num_hidden_layers": 36,
       "num_attention_heads": 32, "num_key_value_heads": 8, "head_dim": 128,
       "vocab_size": 151936, "tie_word_embeddings": False}          # Qwen/Qwen3-8B config.json, read 2026-09-12
if len(sys.argv) > 1:
    with open(sys.argv[1], encoding="utf-8") as fh:
        cfg = json.load(fh)

h, f, L = cfg["hidden_size"], cfg["intermediate_size"], cfg["num_hidden_layers"]
q_heads, kv_heads, d = cfg["num_attention_heads"], cfg["num_key_value_heads"], cfg["head_dim"]
tensors = [                                   # name in the safetensors header, shape as stored (out, in)
    ("self_attn.q_proj", (q_heads * d, h)), ("self_attn.k_proj", (kv_heads * d, h)),
    ("self_attn.v_proj", (kv_heads * d, h)), ("self_attn.o_proj", (h, q_heads * d)),
    ("self_attn.q_norm + k_norm", (2, d)),
    ("mlp.gate_proj", (f, h)), ("mlp.up_proj", (f, h)), ("mlp.down_proj", (h, f)),
    ("input_layernorm + post_attention_layernorm", (2, h)),
]
block = attention = 0
for name, shape in tensors:
    count = shape[0] * shape[1]
    block += count
    attention += count if name.startswith("self_attn") else 0
    print(f"{name:44s} {str(list(shape)):>15s} {count:>14,}")
mlp = block - attention - 2 * h
print(f"{'one block':44s} {'':15s} {block:>14,}   attention {attention / block:.1%}   feed-forward {mlp / block:.1%}")

embed = cfg["vocab_size"] * h
head = 0 if cfg.get("tie_word_embeddings") else embed
total = L * block + embed + head + h
print(f"{L} blocks {L * block:,} + embedding {embed:,} + output head {head:,} + final norm {h:,}")
print(f"= {total:,} parameters = {total / 1e9:.2f} B")

Output — what you should see

block-parameters.py with no argument (Qwen3-8B)

self_attn.q_proj                                [4096, 4096]     16,777,216
self_attn.k_proj                                [1024, 4096]      4,194,304
self_attn.v_proj                                [1024, 4096]      4,194,304
self_attn.o_proj                                [4096, 4096]     16,777,216
self_attn.q_norm + k_norm                           [2, 128]            256
mlp.gate_proj                                  [12288, 4096]     50,331,648
mlp.up_proj                                    [12288, 4096]     50,331,648
mlp.down_proj                                  [4096, 12288]     50,331,648
input_layernorm + post_attention_layernorm         [2, 4096]          8,192
one block                                                       192,946,432   attention 21.7%   feed-forward 78.3%
36 blocks 6,946,071,552 + embedding 622,329,856 + output head 622,329,856 + final norm 4,096
= 8,190,735,360 parameters = 8.19 B

Three checks fall out of that output. The feed-forward network is 78% of a block and attention 22%, which is why the mixture-of-experts models in the next lesson replace the feed-forward sub-layer and leave attention alone. Thirty-six blocks come to 6,946,071,552, the card’s “6.95B” non-embedding count; add the embedding matrix and the untied output head, 622,329,856 each, and the total is 8.19 billion, the card’s “8.2B”. And run against Qwen3-1.7B’s config.json, where tie_word_embeddings is true so the output head costs nothing extra, the script prints 1,720,574,976, the exact figure the lab’s forward-pass script reports after loading the weights:

RunnableAll tracks

the same count for Qwen3-1.7B, once the lab has downloaded it

python block-parameters.py ~/llm-course/models/qwen3-1.7b/config.json

Output — what you should see

block-parameters.py with Qwen3-1.7B's config.json, last three lines; the full run prints self_attn.q_proj [2048, 2048] 4,194,304 first

one block                                                        50,336,000   attention 25.0%   feed-forward 75.0%
28 blocks 1,409,408,000 + embedding 311,164,928 + output head 0 + final norm 2,048
= 1,720,574,976 parameters = 1.72 B

Decoder-only, and what got dropped

The original paper described an encoder-decoder model for translation: an encoder that reads the source with no mask, a decoder that writes the target with a causal mask, and a third kind of attention, cross-attention, in which the decoder’s queries meet the encoder’s keys and values. Every model in this course is decoder-only: one stack, one causal mask, next-token prediction, and no separate encoder. Summarising, translating and answering are all handled by putting the source text into the same context as the answer.

The choice has a runtime consequence you will measure in Part 5. Because every position’s row can be computed at once, reading a prompt is one large matrix multiplication per projection, the prefill; because generation adds one row at a time and that row needs the keys and values of every earlier position, decode stores them rather than recomputing them, and that store is the KV cache the next lesson sizes; on the unified-memory tracks it competes with the weights for the same pool. The mask is also removable: the Transformers documentation describes is_causal=False as switching to “bidirectional attention, where every token attends to every other token”, so that you can “use decoder-only models as text encoders, for example, to generate embeddings”, which is the relationship between generative and embedding models that the embeddings lesson set out.

Attention and the block structure survive from 2017 unchanged. What changed is the detail, and a configuration file is a record of which choices a model made:

Component	The 2017 paper	Qwen3-8B	Where it shows in config.json
Model shape	Encoder-decoder with cross-attention	Decoder-only	architectures: ["Qwen3ForCausalLM"]
Position	Sinusoidal encodings added to the input embeddings	Rotation applied to queries and keys	rope_theta, rope_scaling
Normalisation	Layer normalisation after each sub-layer	RMS normalisation before each sub-layer, plus per-head query and key norms	rms_norm_eps
Feed-forward	Two matrices with a ReLU between	Three matrices, gated, with SiLU	intermediate_size, hidden_act
Heads	As many key-value heads as query heads	8 key-value heads for 32 query heads	num_attention_heads, num_key_value_heads
Output projection	Shares the embedding matrix	A separate matrix of the same size	tie_word_embeddings: false

Why attention cost grows with context

Everything in the block except the score matrix is linear in the number of tokens: double the tokens and each projection and the feed-forward network do twice the work. The score matrix is not. Each of n queries meets each of n keys, per head, per layer, so both its work and its memory grow with n²:

per block, for a prompt of n tokens (FLOPs = 2 × multiply-adds, as in Part 1):
  linear work    = 2 × parameters_per_block × n            projections and the feed-forward network
  attention work = 4 × n² × hidden_size                    q · kᵀ and weights · v, summed over the heads
  score matrix   = n² × num_attention_heads entries        what eager attention holds in memory
the two kinds of work are equal when  n = parameters_per_block / (2 × hidden_size)

For Qwen3-8B, with 192,946,432 parameters per block and a hidden_size of 4096, that crossover is at 23,553 tokens. Arithmetic from the configuration, not a measurement:

Prompt tokens	Linear work per block	Attention work per block	Attention’s share of the block’s work	Score matrix per block, BF16 (the FP32 softmax doubles it)
1,024	0.40 TFLOP	0.02 TFLOP	4%	64 MiB
4,096	1.58 TFLOP	0.27 TFLOP	15%	1 GiB
32,768	12.6 TFLOP	17.6 TFLOP	58%	64 GiB
131,072	50.6 TFLOP	281 TFLOP	85%	1,024 GiB

Two readings. At the prompt lengths of ordinary chat the quadratic term is a rounding error, and prefill costs what Part 1 already priced: two FLOPs per parameter per token. At 32,768 tokens it is more than everything else combined, and at the 131,072 tokens the YaRN extension allows, it is most of the work, which is why time to first token on a long document grows faster than the document does. The memory column is the sharper limit: eager attention needs at least 64 GiB of scratch space per block at 32,768 tokens, and the Transformers eager path takes the softmax in FP32, which doubles it; the causal mask means half of every one of those matrices is -inf that was computed anyway. Time it on your own machine; the ratios are the point, not the milliseconds:

RunnableAll tracks

attention-cost.py - time one layer's eager attention as the context doubles

"""Time one layer's eager attention as the context doubles: the score matrix is n × n per head."""
import time
import torch

heads, head_dim = 8, 128                       # 8 of Qwen3-8B's 32 heads, so the largest score matrix is 512 MiB;
                                               # the process peaks near 1.4 GiB because masked_fill and softmax each take a copy
torch.manual_seed(0)

def eager_attention(n):
    q, k, v = (torch.randn(1, heads, n, head_dim) for _ in range(3))
    causal = torch.ones(n, n, dtype=torch.bool).triu(1)
    t0 = time.perf_counter()
    scores = (q @ k.transpose(-2, -1)) / head_dim ** 0.5          # (1, heads, n, n)
    scores = scores.masked_fill(causal, float("-inf"))
    out = torch.softmax(scores, dim=-1) @ v                         # (1, heads, n, head_dim)
    return scores.numel(), time.perf_counter() - t0

eager_attention(256)                           # warm-up, not timed
previous = None
for n in (512, 1024, 2048, 4096):
    entries, elapsed = eager_attention(n)
    ratio = "" if previous is None else f"   {elapsed / previous:3.1f}× the previous run"
    print(f"n = {n:5d}   scores {entries:>12,} entries = {entries * 4 / 2**20:6.1f} MiB FP32   {elapsed * 1000:7.1f} ms{ratio}")
    previous = elapsed

Output — what you should see

attention-cost.py, torch 2.14.0 on the author's CPU test box; your milliseconds will differ, the ratios should stay near four

n =   512   scores    2,097,152 entries =    8.0 MiB FP32       4.7 ms
n =  1024   scores    8,388,608 entries =   32.0 MiB FP32      19.7 ms   4.2× the previous run
n =  2048   scores   33,554,432 entries =  128.0 MiB FP32      63.0 ms   3.2× the previous run
n =  4096   scores  134,217,728 entries =  512.0 MiB FP32     240.9 ms   3.8× the previous run

That is the problem the optimised backends solve. FlashAttention’s abstract describes “an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM”: the same softmax, computed a tile at a time, so that the n × n matrix never exists in memory. The Transformers documentation puts the trade in two sentences, “Basic attention scales poorly because it materializes the full attention matrix in memory, creating bottlenecks that slow down inference. Optimized implementations rearrange the math to reduce memory traffic for faster, more affordable inference”, and at the pinned version lists flash_attention_3, flash_attention_2, flex_attention and sdpa, plus paged variants of three of them and of eager; plain eager, the reference implementation, appears only in its examples as the value passed to attn_implementation in from_pretrained. llama.cpp’s --flash-attn is the same idea in a different engine, and Part 6 shows what it does to the memory of a long context.

Because those backends never build the matrix, they cannot hand it back. The documentation’s custom-attention example ends by returning attn_output, attn_weights with the note that the weights “are optional here”. This is why the lab loads the model with attn_implementation="eager" before asking for output_attentions=True: it is the one implementation that computes the scores explicitly, it is fine for a five-token prompt on a small model, and it is exactly the trade the optimised backends exist to avoid.

Where this returns: Part 5 measures prefill and decode separately, Part 17 reuses the cache across prompts with prefix caching, and Part 18 splits the sequence itself across machines with sequence parallelism when the cache, not the weights, is what will not fit.

Work through one causal attention row

At the third position of a sequence, a causal decoder may attend to positions one, two and three. It must not attend to position four, even during training when the complete sequence is available in memory. Training can calculate many positions in parallel precisely because the mask prevents future information from entering earlier predictions.

Suppose the unmasked attention scores for those first three positions are equal. Softmax assigns each one third of the available attention mass; the output is the average of their value vectors. If the third score increases, its value contributes more. The attention weights do not themselves become words: the weighted vector still passes through projections, residual paths and subsequent layers before the vocabulary logits are produced.

This example gives you a useful check for an implementation: changing a future token must not change an earlier position’s logits in evaluation mode under a causal mask. Changing an earlier token may change later logits. Keep dropout disabled and compare within numerical tolerance. Attention visualisations can suggest relationships, but a colourful heatmap alone does not establish which information caused the final answer.

Recap

One head of attention is four lines: project to queries, keys and values; score every query against every key and divide by the square root of the head width; add -inf above the diagonal and take a softmax per row; average the values with those weights. Heads run it in parallel, each over its own 128-wide slice of the projected queries, keys and values, and grouped-query attention gives 32 query heads only 8 key-value heads to share. Rotary embeddings rotate queries and keys by position-dependent angles at 64 rates, so a score depends on distance; rope_theta sets the rates and rope_scaling is how a context is extended past the trained range. The block is normalise, attend, add, normalise, feed forward, add, with 78% of its 193 million parameters in the gated feed-forward network, thirty-six times in Qwen3-8B. Linear work grows with the prompt; attention’s score matrix grows with its square, overtakes the linear work at about 23,500 tokens on this model, and is never materialised by the optimised backends, which is why reading the weights back needs the eager implementation.

Check your understanding

Question 1. Qwen3-8B declares num_attention_heads 32, num_key_value_heads 8 and head_dim 128. What does that arrangement mean?

• The model has 40 heads in total
• Thirty-two query heads are divided into eight groups of four, and the heads in a group read one shared key head and one shared value head
• Eight layers use attention and thirty-two do not
• The key and value vectors are 8 dimensions wide

Show the answer and why

Answer: Thirty-two query heads are divided into eight groups of four, and the heads in a group read one shared key head and one shared value head

This is grouped-query attention: k_proj and v_proj are [1024, 4096], eight heads of 128, and repeat_kv serves each of them to four query heads (query head i reads key-value head i // 4). The stored keys and values per token per layer are therefore 2 × 1,024 numbers, not 2 × 4,096.

Question 2. A prompt to Qwen3-8B grows from 4,096 tokens to 32,768, eight times longer. By how much does the linear work per block grow, and by how much does the attention-score work?

• Both grow eight times
• Linear work eight times, attention-score work sixty-four times
• Linear work sixty-four times, attention-score work eight times
• Neither grows, because the KV cache holds the earlier tokens

Show the answer and why

Answer: Linear work eight times, attention-score work sixty-four times

Linear work is 2 × parameters × n; attention work is 4 × n² × hidden_size, so eight times the tokens is sixty-four times the scores. On this model the two are equal at about 23,500 tokens, so at 32,768 the quadratic term has already overtaken everything else, and an eager score matrix would need 64 GiB per block.

Question 3. Which statements about the residual connections in a transformer block are correct? Select all that apply.

• Each sub-layer’s output is added to its input rather than replacing it
• They provide a path along which gradients reach early layers
• They are what make the attention mask causal
• They let each block be read as an edit to a running stream rather than a full transformation

Show the answer and why

Answer: Each sub-layer’s output is added to its input rather than replacing it, They provide a path along which gradients reach early layers, They let each block be read as an edit to a running stream rather than a full transformation

Causality comes from adding -inf above the diagonal of the scores before the softmax, not from residuals. The residual stream is what makes deep stacks trainable, which Part 1 measured on a fifty-layer stack, and is the reason a block can be understood as writing an increment.

Question 4. Each line builds a mask for the eager implementation, which adds the mask to the scores before the softmax. Which one is the bug?

causal = torch.tril(torch.ones(n, n))
(a) mask = torch.where(causal.bool(), 0.0, float("-inf"))
(b) mask = causal
(c) mask = torch.where(causal.bool(), 0.0, -1e9)
(d) mask = torch.zeros(n, n).masked_fill(~causal.bool(), float("-inf"))

• (a)
• (b)
• (c)
• (d)

Show the answer and why

Answer: (b)

A float mask is added to the scores, so 0.0 keeps a position and -inf removes it. Line (b) adds 1.0 to the allowed positions and 0.0 to the forbidden ones: nothing is removed, every position reads the future, and the softmax rows still sum to one so nothing looks wrong. (a) and (d) build the same correct mask two ways; (c) uses a large finite negative, which rounds to the same zero weight after the softmax.

Question 5. You start Qwen3-8B under llama-server with --ctx-size 65536 and leave rope_scaling as null. What happens?

• Every engine refuses, because max_position_embeddings is 40960
• The engine allocates the cache, and positions beyond the trained range are rotated by angles the model never saw in training, so the memory is spent and long-input quality is what suffers
• Rotary embeddings extrapolate automatically, so nothing changes
• The score matrix stays the same size, so the run costs no more than at 32,768

Show the answer and why

Answer: The engine allocates the cache, and positions beyond the trained range are rotated by angles the model never saw in training, so the memory is spent and long-input quality is what suffers

RoPE has no built-in limit: the rotation is defined for any position, so nothing in the model stops the run; what happens next is the engine's choice. llama.cpp v0.4.0 logs "n_ctx_seq (65536) > n_ctx_train (40960) -- possible training context overflow" and continues. vLLM 0.28.0 reads the same field and refuses with "User-specified max_model_len (65536) is greater than the derived max_model_len (max_position_embeddings=40960 or model_max_length=None in model's config.json)" unless VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 is set, and its warning for that override says that with RoPE "positions exceeding derived_max_model_len lead to nan". Neither rescales the rates for you: the card's fix for a 65,536-token application is YaRN with factor 2.0 and original_max_position_embeddings 32768, and Part 7 measures what recall looks like without it. The score matrix, meanwhile, is four times larger at twice the length.

Question 6. Why does the lab load the model with attn_implementation="eager" before asking for output_attentions=True?

• Because eager is faster
• Because the optimised backends compute the softmax a tile at a time and never materialise the n × n matrix, so there is no matrix to return
• Because eager is the only implementation that supports a causal mask
• Because output_attentions is deprecated elsewhere

Show the answer and why

Answer: Because the optimised backends compute the softmax a tile at a time and never materialise the n × n matrix, so there is no matrix to return

FlashAttention-style backends tile the computation to avoid the memory traffic of the full matrix; the documentation describes the returned attention weights as optional for such a backend. Eager computes the scores explicitly, which is the cost the others avoid, and is fine for one short prompt on a small model.

Check my answersClear