Transformers and Attention

Compute attention, mask it, split heads, add positions and stack the block. PyTorch and NumPy from the lessons.

One pre-LN decoder block#

Task	Code
Apply causal mask	`scores.masked_fill(mask_slice == 0, float("-inf"))`
Softmax over keys	`F.softmax(scores, dim=-1)`
Merge heads	`x.transpose(1, 2).contiguous().view(b, t, h * dh)`
Add positions	`x = tok + pos.unsqueeze(0)`
Fill sinusoidal table	`pe[:, 0::2] = torch.sin(angle)`
Rotate Q and K	`apply_rope(x, pos, base=10000)`
Fused kernel with GQA	`scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)`
Tie LM head	`self.lm_head.weight = self.tok_embed.weight`

self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=True)Q, K, V in one matmul

causal_mask = torch.tril(torch.ones(max_context_length, max_context_length))

self.register_buffer("causal_mask", causal_mask, persistent=False)no grad

qkv = self.qkv_proj(x)(B, T, 3D)

q, k, v = qkv.chunk(3, dim=-1)three (B, T, D) views

q = self._split_heads(q)(B, H, T, d_head)

k = self._split_heads(k)

v = self._split_heads(v)

scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_head)

mask_slice = self.causal_mask[:t, :t]top left (T, T)

scores = scores.masked_fill(mask_slice == 0, float("-inf"))

weights = F.softmax(scores, dim=-1)rows sum to 1

context = self._merge_heads(torch.matmul(weights, v))(B, T, D)

out = self.out_proj(context)mixes the heads

Setting	Default	Note
Sinusoidal and RoPE base	10000	raise RoPE base for longer context
mlp_expansion	4	MLP is D to 4D to D
d_head	64	64 or 128 in production
Init std	0.02	Linear and Embedding weights
Residual projection scale	1 / sqrt(2 * num_layers)	out_proj and fc2
GPT-2 small d_model	768	12 heads, 12 layers, about 124M
KV bytes per token per layer	`2 * d_head * dtype_size`	per head, K and V
Compute optimal ratio	20	tokens per parameter

Symptom	Cause	Fix
`ValueError` on d_model and n_heads	D is not divisible by H	pick H that divides D
`seq_len` exceeds max_context_length	T is larger than the mask	crop to the last context_length tokens
Model cheats on next token	no causal mask	mask the upper triangle with `-inf`
Softmax gives `inf`	`exp` of a large number	subtract the row max first
Learning stalls	scores not scaled	divide by `sqrt(d_head)`
Deep stack diverges	post-LN on the residual path	use pre-LN
Weight tying saves nothing	weights copied, not shared	assign the same parameter
Bad output past trained length	positions do not extrapolate	use RoPE with base scaling

Why divide by sqrt(d_k)?: Large dot products saturate softmax and gradients vanish.
Why -inf in the mask?: exp(-inf) = 0, so future positions get zero weight.
Why add positions?: Attention alone is permutation invariant.
Why GQA?: It cuts KV cache memory by N/G with nearly full quality.
Is Flash Attention approximate?: No. A running (max, sum) keeps it exact.