The LLaMA Block

Between 2020 and 2024 the decoder-only Transformer of Modules 7 and 8 picked up changes that LLaMA 1 to 3, Mistral, Qwen, Gemma and almost every large open model now share:

• RoPE replaces the sinusoidal position vector that was added to the embeddings.
• RMSNorm replaces LayerNorm.
• SwiGLU replaces the ReLU or GELU feed-forward layer. (Gemma uses GeGLU, the same design with GELU in place of Swish.)
• Grouped-Query Attention (GQA) lets several query heads share one key head and one value head, which shrinks the KV cache.

The first three change what the model computes; the last changes the memory needed to generate. Section 15.6 sets all four beside the 2017 block. FlashAttention, which computes the same attention more cheaply, is in Module 16.

15.1 Rotary Position Embedding (RoPE)Run this section

Rotary Position Embedding (RoPE) (Su et al., 2021) does not add a position vector at the bottom of the model, as section 7.1 did. Inside every attention layer it rotates each query and key by an angle set by the token's position. In the dot product the two rotations combine into one rotation by the difference of positions, so the score depends on how far apart the tokens are, not where they sit.

In words. Query and key are each rotated by their own position's angle; in the dot product the rotations reduce to one rotation by the difference, so only the gap between the tokens matters.

The 2D Rotation Matrix

Picture a two-number vector $\mathbf{x} = [x_1, x_2]^\top$ as an arrow in a plane. Following the RoPE paper, this section writes vectors as columns. The matrix below turns the arrow counter-clockwise by $m\theta$, for position $m$ and fixed frequency $\theta$:

A real query or key of $d$ numbers (say 128) is split into $d/2$ neighboring pairs, each rotated on its own, so the full matrix has $2 \times 2$ rotations along its diagonal:

Each pair has its own frequency, shrinking geometrically:

The first pair turns one radian per token and the last about ten thousand times more slowly: fast pairs separate neighbors, slow pairs separate positions thousands apart. This is section 7.1's schedule, applied over the head width $d$ rather than $d_{\text{model}}$.

In words. A query pair and a key pair as arrows, rotated by position times frequency. Sliders move both positions and the frequency; the angle between the arrows and their dot product depend only on the position difference.

Why Only the Distance Survives

Rotate a query at position $m$ and a key at position $n$ (the tilde marks rotated vectors):

By $(AB)^\top = B^\top A^\top$, the rotations meet in the middle of the dot product:

Two facts finish it. The transpose of a rotation runs it backwards, $R_{\theta, m}^\top = R_{\theta, -m}$ (negating the angle flips each $\sin$, exactly as transposing does). And consecutive rotations add their angles. So back by $m$ then forward by $n$ is one turn by $n - m$:

Substituting that into the dot product:

The score depends on the two vectors' content and the gap $m - n$, nothing else ($g$ names "some function of these"). Move a passage 500 tokens later and every score inside it is unchanged. That does not mean the model works at any length, as section 15.2 explains.

Measured (lab 4a, experiments/rope.py). The identity holds up to float32 rounding. For a single 2D pair, the score at a fixed offset stays the same as the pair slides to different positions, to within $4.77 \times 10^{-7}$. Over a full 32-dimensional head with the published $10000^{-2i/d}$ frequencies (the same schedule, counting $i$ from zero), it holds to about $10^{-5}$ at offsets $0, 1, 7, 63$ and $200$.

Measured (lab 4a). Now use a learned table of absolute positions instead. Take the same two tokens at the same distance apart and move them to different absolute positions, and the score for that pair changes by $5.78$. That is more than $10^5$ times larger than RoPE's rounding error. The table also has no entry at all for positions past the length it was trained on, while a rotation can be computed for any position and keeps every score consistent.

15.2 Stretching RoPE, and ALiBiRun this section

Train on $4096$ tokens and run at $32{,}768$, and beyond position 4096 the slow pairs reach angles $m\theta_i$ never seen in training, so quality drops. (The fast pairs have gone round the circle many times, so every angle is familiar.) Four common fixes:

1. Linear position interpolation. Shrink positions into the trained range, $m' = m \cdot \frac{L_{\text{train}}}{L_{\text{target}}}$. Every angle is familiar, but neighbors now differ by smaller angles, blurring the fast pairs that encode local order, so some fine-tuning is usually needed.
2. Raising the base, as LLaMA 3 did:
   $$\theta_{\text{base}} = 10{,}000 \longrightarrow 500{,}000$$

   What does this say?introduces 1 notation
   The first pair is unchanged (its exponent is zero) and the slowest stretch most, giving room for distant positions while keeping local detail. LLaMA 3.1, reaching 128k tokens, adds a further rescaling of the frequencies.
3. NTK-aware scaling picks the new base from the extension factor $\lambda$ (new length over old):
   $$\theta_{\text{base}}' = \theta_{\text{base}} \cdot \lambda^{d / (d - 2)}$$

   What does this say?
   The exponent makes the slowest pair slow by exactly $\lambda$, as interpolation would, while leaving the fastest unchanged.
4. YaRN (Peng et al., 2023) makes the split explicit per pair: pairs with wavelengths much shorter than the trained context are left alone, longer ones are interpolated fully, and those between are blended. It also scales the logits up slightly to counter the flatter attention of longer contexts. Qwen2.5 uses it for long-context versions.

In words. Each RoPE pair's wavelength for a model trained on 4,096 tokens and run at 32,768. Unchanged, slow pairs reach unseen angles. Interpolation stretches every wavelength, crowding neighbors in the fast pairs. NTK-aware scaling stretches only the slow end, though some middle pairs still overshoot. YaRN stretches slow pairs fully, leaves fast ones alone and blends between, so no pair reaches an unseen angle.

A Different Route: ALiBi

ALiBi (Press et al., 2022) adds no position information to the vectors. It subtracts from each score a penalty proportional to how far back the key is:

The bracket holds one entry per key up to the query's position $i$: $0$ for itself, $-1$ for the token before, and so on. Each head has its own fixed slope $m$, a geometric sequence such as $\frac{1}{2}, \frac{1}{4}, \dots, \frac{1}{256}$ for eight heads, so some heads focus nearby and others look far back. Depending only on distance, ALiBi degrades gracefully beyond its training length. BLOOM and MPT use it; most later models chose RoPE, whose extensions recover long-context quality with modest fine-tuning.

15.3 RMSNormRun this section

Zhang & Sennrich (2019) argued that LayerNorm's mean subtraction (section 7.3) is not what keeps training stable; the division that fixes the scale is. Dropping the mean, models trained about as well. The result is RMSNorm.

In words. Divide the vector by its own root mean square, then multiply by learned scales. Nothing is subtracted or added.

The RMSNorm Formula

The root mean square is what it says: the square root of the mean of the squares, a measure of size like the standard deviation but without subtracting the mean.

• No mean: one fewer pass over the numbers, often quoted as saving 10–15% of the normalization time.
• No shift $\boldsymbol{\beta}$: half the parameters.

Measured (lab 4b, experiments/norm_and_gqa.py). The parameter claim is exact. LayerNorm has $2d$ parameters and RMSNorm has $d$, so at the tested width the saving is exactly $d_{\text{model}} = 512$. The two layers' outputs have a correlation coefficient of $r = 0.9990$ (1 would mean identical up to scale and shift).

Measured (lab 4b). The lab does not test the speed claim. PyTorch's nn.LayerNorm runs as a single optimized GPU routine (a "fused kernel"), while the experiment's RMSNorm is six separate operations, each a separate trip through memory, and it times $2.8\times$ slower. That compares how the two were written, not the pass saved, so it neither supports nor contradicts the $10$–$15\%$ figure.

15.4 SwiGLURun this section

A gated linear unit (GLU) (Dauphin et al., 2017) makes two projections of the input and multiplies them, so one acts as a gate on the other. Shazeer (2020) tested gated variants in the Transformer's feed-forward layer; LLaMA uses SwiGLU:

• $\operatorname{Swish}(z) = z \cdot \sigma(\beta z) = \frac{z}{1 + e^{-\beta z}}$ with $\beta = 1$ (also called SiLU; this $\beta$ is unrelated to LayerNorm's shift), using the sigmoid of section 3.6. It passes large positive inputs almost unchanged, sends large negative ones to nearly zero, and bends smoothly between: $\operatorname{Swish}(2) \approx 1.76$, $\operatorname{Swish}(-2) \approx -0.24$.
• $\otimes$ is element-by-element multiplication.

In words. The input is projected twice; one copy goes through Swish, the two are multiplied element by element, and the result is projected back to model width. The multiplication makes it a gate: one branch controls how much of the other survives.

In words. One SwiGLU hidden unit: the Swish curve, with ReLU for comparison, marks the gate input, and the gate's output multiplies the up branch's value. Sliders set both inputs.

Keeping the Parameter Count Equal

SwiGLU has three matrices instead of two, so to keep the total the same the inner width shrinks from $4 d_{\text{model}}$ to:

Three matrices of $d_{\text{model}} \times \frac{8}{3}d_{\text{model}}$ hold $8 d_{\text{model}}^2$ weights, the same as two of $d_{\text{model}} \times 4d_{\text{model}}$, so parameters and arithmetic are unchanged. At this equal size SwiGLU still reached lower perplexity than ReLU and GELU in Shazeer's experiments. LLaMA rounds up to a convenient multiple: $d_{\text{ff}} = 11008$ for $d_{\text{model}} = 4096$ in the 7B model.

15.5 Grouped-Query AttentionRun this section

In multi-head attention (MHA, section 6.6), each of the $h_q$ query heads has its own key and value head, so $h_{kv} = h_q$. Every generated token must read the whole KV cache (section 12.2), which grows with $h_{kv}$ and takes much of each step's time.

Grouped-Query Attention (GQA) (Ainslie et al., 2023) splits the query heads into groups that each share one key head and one value head. The extreme, one pair shared by all, is multi-query attention (MQA).

In words. With four query heads: MHA gives each its own keys and values, GQA shares each pair between two queries, and MQA shares one pair among all four.

How Much Memory GQA Saves

LLaMA 3 8B has $h_q = 32$ and $h_{kv} = 8$, four query heads per key-value head. The cache holds one key and value per token, layer and key-value head, $2 \times n_{\text{layers}} \times h_{kv} \times d_{\text{head}} \times L \times n_{\text{bytes}}$ bytes for $L$ tokens (section 23.1), so $h_{kv} = 8$ instead of 32 gives:

• a cache a quarter the MHA size, a $75\%$ reduction;
• a quarter of the reading per step, so up to $4\times$ faster generation when reading the cache is the bottleneck;
• quality close to MHA and better than MQA, in Ainslie et al.'s experiments.

In words. LLaMA 3 8B's 32 query heads joined to the key-value heads they read. At 8,192 tokens the cache is 4 GiB with 32 key-value heads, 1 GiB with 8, and 128 MiB with one.

Measured (lab 4b). The arithmetic is exact. At $32$ query heads, $8$ KV heads, $d_{\text{head}} = 128$, $L = 8192$ and fp16, the cache falls from $4.00$ GiB to $1.00$ GiB. That is exactly the $75.0\%$ stated in this section, and the same figure section 23.1 works out for LLaMA-3 8B.

Measured (lab 4b). The structural claim is stronger than "nearly as good". Set $h_{kv} = h_q$ and GQA reproduces MHA bit for bit, with a largest difference of exactly $0.0$. GQA is a generalization of MHA, not an approximation of it: MHA is the special case with one query per group. That is why an existing MHA model can be converted to GQA, as Ainslie et al. did, by averaging the key and value heads within each group and then training briefly. Parameter counts fall as expected, $16{,}448 \to 9{,}280$, as the number of KV heads drops from $8$ to $1$.

15.6 The 2017 Block and the 2024 Block, Side by SideRun this section

The whole change from the 2017 block of section 7.6 to a 2024 LLaMA-style block:

None is large on its own, and the shape is unchanged: residual stream, alternating attention and per-token network, causal mask, next-token loss. Each detail makes training more stable, long contexts better behaved, or generation cheaper.

In words. The same input through small 2017-style and 2024-style models, both trained on Shakespeare, side by side, stage by stage. The 2024 model has no position stage, since it rotates queries and keys inside attention, and its heads share key-value heads in pairs.

Exercises

Core Academic References

• Su et al. (2021): RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864
• Press, Smith & Lewis (2022): Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (ALiBi). ICLR 2022. arXiv:2108.12409
• Peng et al. (2023): YaRN: Efficient Context Window Extension of Large Language Models. arXiv:2309.00071
• Zhang & Sennrich (2019): Root Mean Square Layer Normalization (RMSNorm). NeurIPS 2019. arXiv:1910.07467
• Dauphin et al. (2017): Language Modeling with Gated Convolutional Networks (GLU). ICML 2017. arXiv:1612.08083
• Shazeer (2020): GLU Variants Improve Transformer (SwiGLU). arXiv:2002.05202
• Ainslie et al. (2023): GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP 2023. arXiv:2305.13245
• Touvron et al. (2023): LLaMA: Open and Efficient Foundation Language Models. arXiv:2302.13971
• Grattafiori et al. (2024): The Llama 3 Herd of Models. arXiv:2407.21783