Llama 4 From Scratch

Llama-style models are decoder-only transformers — masked self-attention over a stack of identical blocks — with four choices that separate them from the 2018 GPT recipe: pre-norm RMSNorm, rotary position embeddings, a SwiGLU feed-forward, and grouped-query attention. Llama 4 adds two: a mixture-of-experts feed-forward with a shared expert, and interleaved attention layers that drop positional embeddings entirely.

The idea

Every block does two things, each wrapped in a residual add:

Pre-norm keeps a clean gradient path to the input, which is what makes deep stacks trainable.

• RMSNorm divides by the root-mean-square of the activations and applies a learned per-channel gain. No mean subtraction, no bias — cheaper than LayerNorm.
• RoPE rotates each query/key pair by an angle proportional to its position, so $q^\top k$ depends on the relative distance between tokens rather than absolute indices.
• GQA lets many query heads share fewer key/value heads. The KV cache — the memory that limits long-context inference — shrinks by exactly the sharing ratio.
• SwiGLU replaces the two-matrix MLP with three: $\text{down}\big(\text{Swish}(\text{gate}(x)) \odot \text{up}(x)\big)$, with the hidden width scaled to keep parameters comparable.
• MoE swaps the single FFN for $N$ expert FFNs and a router that sends each token to the top-$k$. Compute scales with $k$; capacity scales with $N$. Llama 4 keeps one always-on shared expert, so common features need not be relearned by every routed expert.
• iRoPE interleaves RoPE layers with layers that use no positional embedding at all, letting attention generalize past the trained window, with temperature scaling of attention logits at inference.

Worked example

Take a representative block: $d_{model} = 4096$, 32 query heads with $d_{head} = 128$, and 8 KV heads.

• The query projection is $4096 \times 4096$. The key and value projections are $4096 \times (8 \times 128) = 4096 \times 1024$ each — a quarter the size, not equal.
• KV cache per token per layer is $2 \times 8 \times 128 = 2048$ values. Full multi-head attention would cache $2 \times 32 \times 128 = 8192$ values. That is a 4× reduction, the ratio of query to KV heads.
• Across 32 layers, one token costs $32 \times 2048 = 65{,}536$ values. In fp16 that is 128 KB per token, so an 8192-token sequence occupies about 1 GB of cache.
• RoPE rotates pairs of dimensions within each head using frequencies $\theta_i = 10000^{-2i/d_{head}}$.

In code

import torch

class RMSNorm(torch.nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.g = torch.nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        return self.g * x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

class SwiGLU(torch.nn.Module):
    def __init__(self, d, hidden):
        super().__init__()
        self.gate = torch.nn.Linear(d, hidden, bias=False)
        self.up = torch.nn.Linear(d, hidden, bias=False)
        self.down = torch.nn.Linear(hidden, d, bias=False)

    def forward(self, x):
        return self.down(torch.nn.functional.silu(self.gate(x)) * self.up(x))

No biases anywhere: RMSNorm's gain subsumes what a bias would supply.

Check yourself

1. Why does RoPE make the attention score depend on relative rather than absolute position?
2. With 32 query heads and 8 KV heads, how much smaller is the KV cache, and what is the cost?
3. In an MoE feed-forward, why can total parameters grow far faster than per-token compute?

Key takeaways

• The block is pre-norm RMSNorm → attention and pre-norm RMSNorm → SwiGLU, each with a residual add.
• RoPE and GQA are the two choices that most affect long-context memory and behavior.
• Llama 4 scales capacity through MoE and context through iRoPE, not by making attention denser.