Transformer internals

This chapter dissects a single pre-norm decoder block: the residual stream, the two sub-layers, and the parameter accounting that lets you estimate any model size on the back of a napkin. It is the architectural substrate that every subsequent chapter builds on.

x' = x + Attention(LayerNorm(x))   # sub-layer 1
y  = x' + FFN(LayerNorm(x'))        # sub-layer 2

This chapter builds self-attention from scratch: why the formula is what it is, what each division and masking step does, and where the O(L²) cost comes from. Every attention variant in the rest of the transformer section — MQA, GQA, MLA, FlashAttention — is a direct response to the costs analyzed here.

Q = X W_Q     # (L, d_k)  — what each token is looking for
K = X W_K     # (L, d_k)  — what each token offers as a key
V = X W_V     # (L, d_v)  — what each token offers as a value

scores = Q Kᵀ / √d_k        # (L, L)   — how relevant is each key?
scores = scores + causal_mask # add −∞ to future positions
attn   = softmax(scores)    # (L, L)   — rows sum to 1
out    = attn @ V           # (L, d_v) — blend values by weight

The KV cache is the dominant VRAM cost at long context. This chapter traces the four-step evolution that shrinks it from O(H·d) per token (MHA) down to ~7% of that (MLA), explaining the mechanism and quality trade-off at each step. Know all four cold — this is a top-5 LLM interview topic.

The six families

RoPE in one paragraph

RoPE (Su 2021, arxiv 2104.09864) groups Q and K dimensions into 2D pairs and rotates each pair by an angle θ_i · pos, with θ_i = base−2i/d (default base 10000). Because rotations preserve dot products, Q·K after rotation depends only on relative position (pos_q − pos_k). This gives you relative-position semantics without learning a relative bias matrix.

# Per 2D pair (x, y), at position p, with base frequency θ:
[x']   [cos(θp)  -sin(θp)] [x]
[y'] = [sin(θp)   cos(θp)] [y]

YaRN — extending RoPE to long context

YaRN (Peng 2023) categorizes RoPE frequencies into three bands and rescales them differently:

• High-freq dims (fast oscillation, position-sensitive): minimal interpolation — they extrapolate.
• Low-freq dims (slow, semantic): full positional interpolation (compress positions into the trained range).
• Mid-freq dims: smoothly interpolated between the two.

Plus an attention-temperature factor that compensates for the lower entropy at long context. Result: extend a 4k pretrained model to 128k with a brief continued-training pass on long-context data.

ALiBi vs RoPE — pick one and defend

Why a cache exists

For token t, attention reads K and V for all tokens ≤ t. None of those K, V vectors depend on t's query — they were computed when each past token was processed. Recomputing them per step would make decode O(L²) per token; caching makes it O(L) per token (reading the cache).

The size formula

cache_size = 2 · L_layers · n_kv_heads · d_head · seq_len · bytes_per_value

The 2 is for K + V. n_kv_heads is where MHA → MLA shrinks the bill. bytes_per_value is 2 for bf16/fp16, 1 for int8 quantized cache.

For a deeper dive on paged KV, prefix sharing and the prefill vs decode distinction, see LLM inference.

Standard FFN — the original

y = W_2 · σ(W_1 · x + b_1) + b_2     # σ ∈ {ReLU, GELU}

Hidden dim 4d. Two matrices, one nonlinearity. GPT-2/3, BERT.

SwiGLU — the 2024+ default

SwiGLU (Shazeer 2020) introduces a gate:

y = W_2 · ( Swish(W_1 · x) ⊙ (W_3 · x) )
   where Swish(z) = z · sigmoid(z)

Three matrices. To keep parameter count constant relative to a 4d standard FFN, hidden dim is reduced to ~2.67d (= 8d/3). Used by Llama, Mistral, DeepSeek, PaLM, Qwen.

Why gating works (intuition)

The gate Swish(W_1·x) learns which dimensions of the value W_3·x to suppress per-token. It's a soft, data-dependent feature selector — close in spirit to LSTM gates and to multiplicative interactions. Empirically the gain over plain ReLU/GELU is small but consistent and free, so everyone adopted it.

The forward pass

Per token x:

1. Router: logits = W_r · x over N experts.
2. Top-k: pick the k experts with highest logits; softmax their logits → gate weights g_i.
3. Expert forward: each picked expert i computes FFN_i(x).
4. Combine: y = Σ_i g_i · FFN_i(x).

Routing variants

• Top-1 (Switch Transformer, Fedus 2021): one expert per token. Cheapest; quality dip vs top-2.
• Top-2 (Mixtral): two experts. Mixtral-8x7B → 47B total params, ~13B active per token.
• Top-K with shared experts (DeepSeek-MoE, arxiv 2401.06066): some experts always activate, routed experts specialize. DeepSeek V3 has 1 shared + 256 routed, top-8 routed.
• Expert Choice (Zhou 2022): experts pick tokens, not vice versa — guarantees load balance by construction. Tradeoff: tokens not chosen are dropped or sent to a default.

Load balancing — get the formula right

Without intervention the router collapses onto a few favorite experts. Three remedies, in chronological order:

1. Switch Transformer auxiliary loss

L_aux = α · N · Σ_{i=1}^{N} f_i · P_i

• N = number of experts
• f_i = fraction of tokens routed to expert i (per batch)
• P_i = mean router probability for expert i (per batch)
• α = scalar coefficient (typically 0.01)
• Why the N multiplier? Under perfectly uniform routing (f_i = P_i = 1/N), the loss equals 1 regardless of N — cleanly normalized.

2. Router Z-loss (memorize alongside)

(ST-MoE / PaLM, also DeepSeek) Penalize the log-partition of router logits to keep them bounded:

L_z = (1/B) · Σ_i ( log Σ_j exp(x_{ij}) )²

Where x_ij is the router logit for expert j on token i. Without it, router logits drift in bf16/fp8 and softmax becomes numerically unstable.

3. Auxiliary-loss-free balancing (DeepSeek V3)

Drop the auxiliary loss entirely. Maintain a per-expert bias b_i that's added to the router logits at routing time only (not used for the gate weights). When an expert is over-subscribed, decrement its bias; under-subscribed, increment. The bias is a hyperparam-light PI controller on the routing distribution, with no gradient interference from an auxiliary loss term.

When MoE pays off

MoE buys parameter capacity for free FLOPs at training time, but the full model must fit in memory at inference (all experts must be loaded; only k run per token). On a single GPU MoE is rarely worth it. Across many GPUs with expert parallelism (each GPU holds different experts), MoE wins because compute scales with active params not total params.

Pre-norm vs post-norm

RMSNorm — LayerNorm minus the mean

RMSNorm (Zhang & Sennrich 2019) drops the mean-subtraction step of LayerNorm:

RMSNorm(x) = x / sqrt(mean(x²) + ε) · γ

Faster (no mean), no bias parameter, slightly better empirically. Used in Llama, Mistral, DeepSeek. LayerNorm is now legacy in LLMs.

Sandwich norm and QK-Norm

• Sandwich norm / double norm (Gemma 2, OLMo): LN both before and after a sub-layer. Adds stability at scale.
• QK-Norm (OLMo, some Gemma variants): apply LN to Q and K before the dot product. Bounds attention score magnitudes, prevents softmax saturation, key for training very large/wide models.

Why residuals are non-negotiable

Without residuals, the gradient through L sub-layers is a product of L Jacobians — vanishes or explodes. With residuals, ∂(x + f(x))/∂x = I + ∂f/∂x, preserving an identity component. This is the architectural enabler that lets transformers go to L=80, 100, or beyond.

The four families

Byte fallback — the robustness trick

When the tokenizer encounters a character it can't segment (rare emoji, unusual script), byte fallback emits the raw UTF-8 bytes as tokens instead of an <unk>. Critical for production: the model can always generate something, never crashes on input.

Vocab size — the tradeoff

Larger vocab → fewer tokens per text → faster inference per character; but larger embedding table and more compute per softmax. tiktoken cl100k_base (GPT-3.5/4) uses ~100k vocab; o200k (GPT-4o) uses 200k. Llama 3 jumped to 128k from Llama 2's 32k for the same reason.

BPE vs Unigram — when each

BPE is deterministic — given a tokenizer, every text has one segmentation. Unigram supports subword regularization: at training time, sample different segmentations of the same text per epoch as data augmentation. Used in some multilingual setups; BPE is more common.

The problem TP alone doesn't solve

With tensor parallelism (TP), Q/K/V projections are sharded across GPUs, but LayerNorm, residuals and dropout all see the full hidden dim and the full sequence on each GPU. At long context the activations from these regions blow VRAM. The shared region between sub-layers is what TP can't help.

Sequence parallelism

Sequence parallelism (Korthikanti 2022, arxiv 2205.05198) shards the sequence dimension during the LN/dropout regions where TP can't help:

• TP region (attention/FFN matmul): full sequence, sharded hidden.
• SP region (LN, dropout, residual): sharded sequence, full hidden.
• Boundaries: all-gather (going into TP region) and reduce-scatter (going out).

Net effect: dramatic activation memory reduction at the cost of two extra collective ops per layer. Combined with TP it gives much higher effective context per GPU.

Context parallelism — ring attention

Context parallelism shards the sequence across GPUs inside attention itself. Each GPU holds a chunk of the sequence; K and V chunks rotate around a ring (Ring Attention, Liu 2023, arxiv 2310.01889) so every Q chunk eventually sees every K,V chunk. With FlashAttention-style online softmax, you accumulate the partial outputs without materializing the full score matrix anywhere. This is what enables >1M context training and inference.

Common interview questions

1. "Why divide attention scores by √d_k?" → Without scaling, dot product variance grows with d_k → softmax saturates → vanishing gradients.
2. "Walk through MHA → MQA → GQA → MLA." → Reduces KV cache size step-by-step. MLA goes furthest by storing only a low-rank latent and reconstructing K,V on the fly.
3. "Why RoPE over learned positional?" → Encodes relative position naturally via rotation in 2D subspaces; extrapolates better with PI/YaRN tricks.
4. "Pre-norm vs post-norm — pick one and defend." → Pre-norm: gradients flow through residual cleanly, stable for depth, easier to train at scale. Post-norm: slightly better quality if stable, used in original Transformer, harder to train deep.
5. "Walk through one MoE forward pass." → Token x → router → softmax over E experts → top-k selection (and gating weights) → forward through k experts → weighted sum → output. Plus aux load-balance loss during training.
6. "Why does DeepSeek V3 use shared + routed experts?" → Shared experts capture common knowledge across all tokens; routed experts specialize on niches. Reduces redundant routed-expert capacity.
7. "How does QK-Norm help?" → LN on Q and K before dot product bounds attention score magnitudes; prevents softmax saturation at large scales; improves training stability for very deep / wide models.
8. "What changes if you switch from BPE to Unigram LM?" → Unigram supports probabilistic segmentation (sample different segmentations per training pass) → subword regularization. BPE is deterministic.

Transformer quiz — readiness check

1. Walk through one self-attention forward pass with shapes.
   Show answer

   x: (B, T, d). Q = xW_Q, K = xW_K, V = xW_V → each (B, T, d_k). scores = Q K^T / √d_k → (B, T, T). Apply causal mask (set upper triangle to -∞). attn = softmax(scores, dim=-1). out = attn V → (B, T, d_v).

2. Memory complexity of attention?
   Show answer

   O(T²) for the attention matrix. FLOPs O(T² d). FlashAttention reduces memory to O(T) via tiling + online softmax — never materializes the full attn matrix in HBM.

3. Difference between MHA, MQA, GQA, MLA in KV cache size.
   Show answer

   Per token per layer: MHA: 2·n_heads·d_head. MQA: 2·d_head (8× smaller for 8-head MHA). GQA(g): 2·g·d_head (intermediate). MLA: 2·d_latent (much smaller — DeepSeek-V2 latent ~512 vs MHA's ~8192).

4. Why does MLA need decoupled RoPE?
   Show answer

   RoPE is non-linear in position — you cannot absorb a position-rotated K reconstruction into Q's projection. So MLA splits Q,K into content (low-rank, no RoPE, absorbable) and rotary (small dim, RoPE-applied, MQA-style shared) parts. Final score = Q^C K^C + Q^R K^R.

5. Sinusoidal vs RoPE vs ALiBi — when each?
   Show answer

   Sinusoidal: original; can extrapolate weakly. Learned absolute: GPT-2/3; bad extrapolation. ALiBi: linear attention bias; good extrapolation; no positional embedding. RoPE: rotation in 2D subspaces; relative position via dot product; current default. RoPE + YaRN: best for long-context extension.

6. Why scale FFN hidden dim 4d?
   Show answer

   Empirically optimal width-to-depth ratio at fixed parameters. 4d gives the FFN ~2/3 of the block's parameters; the remaining 1/3 are attention. SwiGLU uses ~2.67d × 3 matrices to keep param count constant.

7. Switch Transformer aux loss formula?
   Show answer

   L_aux = α · N · Σ f_i · P_i where N = num experts, f_i = fraction of tokens routed to expert i, P_i = mean router prob for expert i, α ≈ 0.01. The N multiplier normalizes so uniform routing → loss = 1.

8. What is auxiliary-loss-free balancing (DeepSeek V3)?
   Show answer

   Drop the aux loss; instead maintain per-expert bias added to router logits at routing time only. Decrement bias for over-subscribed experts, increment for under-subscribed. Hyperparam-free PI controller; avoids the gradient interference aux losses introduce.

9. What does router Z-loss do?
   Show answer

   L_z = (1/B) Σ (log Σ exp(x_ij))² where x_ij are router logits. Penalizes the partition function — keeps router logits bounded so softmax stays stable in bf16/fp8. Used by ST-MoE / PaLM / DeepSeek.

10. Why softmax not sigmoid in attention?
    Show answer

    Each row should sum to 1 — attention is a distribution over keys. Sigmoid would let multiple keys be "fully" attended. Some recent work explores sigmoid attention for length generalization but softmax is the default.

11. Why do tokenizers matter for non-English?
    Show answer

    BPE merges trained on English-heavy corpora produce many tokens per word for non-English languages → cost asymmetry (more tokens to encode same content) and worse inductive bias for those languages. Fix: multilingual training, byte fallback, larger vocab.

12. What's the difference between byte-level BPE and SentencePiece?
    Show answer

    Byte-level BPE (GPT-2): operates on UTF-8 bytes; vocab includes all 256 bytes; handles any input. SentencePiece (Kudo): operates on raw text (no whitespace pretokenization); supports BPE or Unigram; standard for Llama / Mistral.

13. How does Mixtral 8x7B work?
    Show answer

    Each FFN block has 8 expert FFNs; router selects top-2 per token; output is gate-weighted sum of 2 experts. Total params 47B; active per token ~13B (because only 2/8 experts run + the shared attention). Inference is faster than 47B dense; quality close to it.

14. Why does the residual connection matter for deep transformers?
    Show answer

    Without residuals, gradient through L sublayers is a product of L Jacobians → vanishes or explodes. Residual ∂(x + f(x))/∂x = I + ∂f/∂x preserves a clean identity gradient path. Lets you train 100+ layer networks.

15. What's QK-Norm and when is it used?
    Show answer

    LayerNorm applied to Q and K before the dot product. Bounds attention score magnitudes → prevents softmax saturation → improves training stability at very large scale. Used in OLMo and some recent Gemma variants.

16. Why is GPT-style decoder-only the dominant architecture in 2026?
    Show answer

    (1) Generative tasks need autoregressive decoding. (2) Encoder-decoder requires task split (BERT-style filling vs GPT-style generation); decoder-only with causal attention does both. (3) Scaling has worked exceptionally well; switching costs aren't worth it. (4) Inference can use KV cache straightforwardly.

17. What does sequence parallelism shard, and what's the comm pattern?
    Show answer

    Shards the sequence dim during LN/dropout/residual regions (where TP can't help). All-gather + reduce-scatter on TP/SP boundaries. Cuts activation memory significantly. Combined with TP, gives much higher effective context per GPU.

18. What is "expert choice" routing vs token-choice?
    Show answer

    Token-choice (standard): each token picks top-k experts. Can produce imbalance. Expert-choice (Zhou 2022): each expert picks top-k tokens (capacity per expert is fixed). Guarantees load balance — no aux loss needed. Tradeoff: tokens not picked are dropped or routed to a default.

19. Why pre-norm enables deeper networks?
    Show answer

    Pre-norm: x + Sublayer(LN(x)) — residual is unnormalized, so gradients flow through it cleanly via the identity path. Post-norm: LN(x + Sublayer(x)) — gradient passes through the LN's reciprocal-std term, which can shrink or amplify; harder to train deep without careful warmup.

20. How does YaRN extend RoPE to long context?
    Show answer

    Frequency-categorized RoPE rescaling: high-freq RoPE dims (position-sensitive, fast oscillation) get less interpolation; low-freq dims (semantic, slow oscillation) get full interpolation. Plus an attention temperature factor. Allows extending pretrained 4k models to 128k with brief continued training.