2.10 Modern Architecture Improvements & Design Choices¶

The Transformer described by Vaswani et al. in 2017 was a landmark, but the architecture used in today’s production LLMs — Llama 4, Qwen3, DeepSeek-V3, Gemma 3 — bears only a family resemblance to that original design. A practitioner opening a modern model config file for the first time will encounter terms like rms_norm, silu, rope_theta, num_key_value_heads, qk_norm, and no_bias. Each of these represents a deliberate, empirically validated architectural decision made in pursuit of better training stability, improved compute efficiency, or stronger final model quality.

This chapter is your annotated map of those decisions. We will examine each improvement from first principles — why it was introduced, what it fixes, how it works mathematically, how it appears in code, and finally what its name is in a real config.json and which library ships the fused kernel for it. We conclude with a “modern transformer recipe” summarizing which choices are nearly universal versus still debated. For the transformer block basics, see The Transformer Block: Norms, Residuals, MLPs & Activations; for positional encodings in depth, see Positional Encodings: Sinusoidal, Learned, RoPE & ALiBi; and for GQA/MQA mechanics, see Multi-Head Attention, MQA, GQA & MLA.

Every choice surveyed here is committed to in The Stack-100M Architecture, where the capstone freezes one config — RMSNorm pre-norm, SwiGLU, GQA with 8 query / 2 KV heads, RoPE with interleaved position-free layers, QK-norm, tied embeddings, no biases — and reproduces its parameter count to the last norm vector. Read this chapter for the menu; read that one for the order.

RMSNorm: Cheaper, Equally Stable Normalization¶

The problem with LayerNorm¶

The original Transformer used LayerNorm (Ba et al., “Layer Normalization”, 2016):

where \(\mu = \frac{1}{d}\sum_i x_i\) and \(\sigma^2 = \frac{1}{d}\sum_i (x_i - \mu)^2\). This requires two passes over the feature vector: one to compute the mean, one to compute the variance. More importantly, it requires two learned parameter vectors: \(\gamma\) (scale) and \(\beta\) (shift). The \(\beta\) term encodes an explicit re-centering.

RMSNorm derivation¶

Zhang and Sennrich (“Root Mean Square Layer Normalization”, NeurIPS 2019) asked: is the re-centering operation actually necessary? They ablated LayerNorm into its components and found that the re-scaling (via \(\gamma\)) drives almost all of LayerNorm’s benefit, while mean subtraction contributes little to final performance but accounts for roughly a third of LayerNorm’s compute.

Root Mean Square Normalization (RMSNorm) drops mean subtraction entirely:

There is no \(\beta\) parameter. One division by the RMS, then a learned element-wise scale. This is approximately 10–30% faster than LayerNorm on modern hardware because there is no mean subtraction kernel and no second reduction pass.

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    """
    Root Mean Square Layer Normalization.
    Used in Llama, Mistral, Qwen, DeepSeek, Gemma, and most modern LLMs.

    Parameters
    ----------
    dim : int
        The feature dimension (hidden size d_model).
    eps : float
        Small constant for numerical stability. 1e-5 is common; some
        models (Gemma) use 1e-6.
    """

    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        # gamma: learnable scale, initialized to ones so RMSNorm is
        # initially the identity (modulo the normalization step).
        self.weight = nn.Parameter(torch.ones(dim))

    def _norm(self, x: torch.Tensor) -> torch.Tensor:
        # x: (..., dim). Compute RMS over last dimension.
        # rsqrt = 1/sqrt for numerical efficiency.
        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Cast to float32 for the norm computation for precision,
        # then cast back. This pattern is used in Llama reference code.
        output = self._norm(x.float()).type_as(x)
        return output * self.weight


# --- Quick sanity check ---
torch.manual_seed(42)
x = torch.randn(2, 8, 512)   # batch=2, seq=8, d_model=512
norm = RMSNorm(dim=512)
y = norm(x)
# Verify unit-RMS along the feature dimension (before scaling by weight)
raw_normed = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + 1e-5)
print(f"RMS of raw_normed row 0: {raw_normed[0,0].pow(2).mean().sqrt():.4f}")  # ~1.000

Pre-norm vs post-norm¶

The original Transformer (and BERT) used post-norm: normalization applied after the residual addition. GPT-2 already moved the norm to the input of each sub-block — pre-norm — and added a single final norm before the LM head; nearly all modern models kept that placement:

Post-norm (Vaswani 2017, BERT):  x → Sublayer(x) + x → Norm → output
Pre-norm  (GPT-2 onward):        x → Norm(x) → Sublayer → + x → output

Pre-norm yields more stable gradients at large scale — the residual path is always clean, so gradients flow through the skip connection without being divided by a normalization operation. The trade-off is that pre-norm models can exhibit “representation collapse” at depth if not carefully initialized, which motivates other tricks like scaled initialization.

Two further placements are now common at the frontier, both of which keep the clean residual highway but re-bound what each block adds to it. Sandwich norm (Gemma ⅔) keeps the pre-norm and adds a second RMSNorm on the sublayer output before the residual add: x + Norm_post(Sublayer(Norm_pre(x))). Reordered / output norm (OLMo 2) drops the pre-norm and normalizes only the branch output: x + Norm(Sublayer(x)). Both cost one extra cheap reduction per sublayer and both were adopted specifically because they suppress loss spikes in long runs (see Training Stability, Loss Spikes & Debugging Large Runs). A different axis entirely — running the attention and MLP sublayers in parallel off a single norm, as PaLM and GPT-NeoX do — is derived in The Transformer Block: Norms, Residuals, MLPs & Activations.

SwiGLU: A Better MLP Activation¶

From ReLU to GELU to SwiGLU¶

The original Transformer used a plain two-layer MLP with ReLU:

BERT switched to GELU (Gaussian Error Linear Unit), which is smoother and empirically outperforms ReLU. Modern LLMs take this further with Gated Linear Units (GLU) and specifically their SiLU-gated variant — SwiGLU — introduced by Noam Shazeer (“GLU Variants Improve Transformer”, 2020).

SwiGLU mechanics¶

A Gated Linear Unit multiplies two linear projections together, one of which passes through a nonlinearity acting as a soft gate:

where \(\text{SiLU}(z) = z \cdot \sigma(z)\) (also called Swish), and \(\otimes\) denotes element-wise multiplication. The SiLU function is smooth, non-monotone, and has a non-zero gradient for negative inputs, all of which help gradient flow.

The gating mechanism is important: \(xV\) produces a “content” projection and \(\text{SiLU}(xW)\) produces a soft gate. The gate dynamically suppresses or amplifies each dimension of the content, giving the MLP a multiplicative interaction it lacked before.

Because SwiGLU has three weight matrices (\(W\), \(V\), \(W_2\)) instead of two, to keep parameter count equal to a standard 4x MLP, the hidden dimension is reduced to \(\frac{2}{3} \cdot 4d = \frac{8d}{3}\). In practice most models round to a multiple of 256 for hardware alignment; Llama 2 70B uses an intermediate size of 28,672 for a model dimension of 8,192.

import torch
import torch.nn as nn
import torch.nn.functional as F

class SwiGLUMLP(nn.Module):
    """
    SwiGLU feed-forward network used in Llama, Qwen, DeepSeek, etc.

    For a model with hidden_dim = d, the intermediate_dim is typically
    floor(8*d/3) rounded up to a multiple of 256.
    """

    def __init__(self, hidden_dim: int, intermediate_dim: int):
        super().__init__()
        # Three linear projections; no bias (see section on no-bias later)
        self.gate_proj = nn.Linear(hidden_dim, intermediate_dim, bias=False)
        self.up_proj   = nn.Linear(hidden_dim, intermediate_dim, bias=False)
        self.down_proj = nn.Linear(intermediate_dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # gate: apply SiLU to create a soft gate
        gate = F.silu(self.gate_proj(x))
        # up: linear content projection
        up   = self.up_proj(x)
        # Element-wise product (gated content), then project back down
        return self.down_proj(gate * up)


# Demonstration: compare SwiGLU vs vanilla FFN parameter counts
d = 4096  # typical hidden size

# Vanilla FFN with 4x expansion: 2 * d * 4d = 8d^2
vanilla_params = 2 * d * (4 * d)

# SwiGLU FFN with 8/3 expansion: 3 * d * (8d/3) = 8d^2 — same!
intermediate = int(8 * d / 3)
swiglu_params = 3 * d * intermediate
print(f"Vanilla 4x FFN params: {vanilla_params:,}")
print(f"SwiGLU  8/3 FFN params: {swiglu_params:,}")
# Both are approximately 8*d^2 — SwiGLU is iso-parameter with better quality.

RoPE: Rotary Position Embeddings¶

Rotary Position Embedding (RoPE), introduced by Su et al. (“RoFormer: Enhanced Transformer with Rotary Position Embedding”, 2021), is now the dominant positional encoding for autoregressive LLMs. See Positional Encodings: Sinusoidal, Learned, RoPE & ALiBi for the full derivation; here we focus on why it was adopted in the modern recipe and its practical configuration.

Why RoPE won¶

RoPE encodes position by rotating query and key vectors in 2D subspaces. The critical insight is that the dot product \(q_m^\top k_n\) naturally becomes a function of the relative offset \((m - n)\), not of absolute positions — this is the relative position property that matters for generalization. Unlike sinusoidal encodings, RoPE requires no separate embedding table and requires no modification to the value vectors. Unlike learned absolute position embeddings, it generalizes beyond training context length (with appropriate scaling, see below).

The rotation for position \(m\) applied to a 2D subspace of the query vector:

For a \(d_k\)-dimensional head, this is repeated for \(d_k/2\) rotation pairs, each with a different base frequency:

The rope_theta hyperparameter (the base) controls the frequency range. GPT-2 had no RoPE. Llama 1 used rope_theta=10000. Llama 3 extended this to rope_theta=500000 to improve long-context behavior by spreading frequencies more broadly, making it easier to extrapolate to new positions at inference time. RoPE’s dominance is now being probed at the long-context frontier: Llama 4 (2025) uses iRoPE, interleaving a minority of position-free “NoPE” attention layers among the RoPE layers to support its 10M-token context window.

import torch
import torch.nn as nn

def precompute_freqs_cis(dim: int, max_seq_len: int, theta: float = 10000.0):
    """
    Precompute complex exponentials for RoPE.
    Returns a tensor of shape (max_seq_len, dim//2) with dtype=complex64.
    The 'cis' notation: cis(x) = e^{ix} = cos(x) + i*sin(x).
    """
    # Frequencies: theta^{-2i/dim} for i in [0, dim/2)
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
    # Positions: [0, 1, 2, ..., max_seq_len-1]
    t = torch.arange(max_seq_len)
    # Outer product: shape (max_seq_len, dim//2)
    freqs = torch.outer(t, freqs)
    # Convert to complex: e^{i * freqs}
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis


def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor,
                     freqs_cis: torch.Tensor):
    """
    Apply RoPE to query and key tensors.
    xq, xk: (batch, seq_len, n_heads, head_dim)
    freqs_cis: (seq_len, head_dim//2) complex tensor
    """
    # Reshape to complex view: treat pairs of reals as complex numbers
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

    # Broadcast freqs_cis over batch and head dimensions
    # freqs_cis shape: (seq_len, head_dim//2) -> (1, seq_len, 1, head_dim//2)
    freqs_cis = freqs_cis.unsqueeze(0).unsqueeze(2)

    # Multiply in complex space = rotation in 2D pairs
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)

    return xq_out.type_as(xq), xk_out.type_as(xk)


# Demonstration
torch.manual_seed(0)
batch, seq_len, n_heads, head_dim = 2, 16, 8, 64
xq = torch.randn(batch, seq_len, n_heads, head_dim)
xk = torch.randn(batch, seq_len, n_heads, head_dim)

freqs = precompute_freqs_cis(head_dim, seq_len, theta=500000.0)  # Llama 3 theta
xq_rot, xk_rot = apply_rotary_emb(xq, xk, freqs)
print(f"RoPE applied. Output shape: {xq_rot.shape}")
# Verify relative position property: dot product depends only on relative offset
q0 = xq_rot[0, 3, 0]  # position 3, head 0
k5 = xk_rot[0, 8, 0]  # position 8, head 0
# The dot product q0.k5 encodes offset=5, not absolute positions 3 and 8.

Grouped Query Attention (GQA)¶

Standard Multi-Head Attention (MHA) maintains separate \(K\) and \(V\) projection matrices for each of \(H\) heads. During autoregressive decoding, the key-value (KV) cache grows as \(O(L \cdot H \cdot d_k)\) per layer — for a 70B model with 64 heads and a 128K context, this is on the order of tens of gigabytes. See Multi-Head Attention, MQA, GQA & MLA for the full treatment; here we focus on the design decision and its practical configuration.

Grouped Query Attention (GQA), introduced in Ainslie et al. (“GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints”, 2023), is a generalization that interpolates between Multi-Head Attention and Multi-Query Attention (MQA, which has a single KV head for all query heads):

MHA:  H_q = H_kv                  (all heads have own KV)
GQA:  H_q = G * H_kv  for G > 1   (G query heads share one KV head)
MQA:  H_kv = 1                     (all query heads share one KV)

The near-universal choice is 8 KV heads, almost regardless of query-head count: Llama 3 8B uses 32 query / 8 KV heads (\(G=4\)), Llama 2 70B and Llama 3 70B use 64 query / 8 KV heads (\(G=8\)), and the Qwen 2.5 / Qwen3 families follow the same pattern. This reduces KV cache memory by a factor of \(G\) while introducing only a small quality degradation versus MHA. Eight is not a coincidence: under tensor parallelism the KV heads are split across ranks, so \(H_{kv} \ge \text{TP degree}\) lets each GPU of a standard 8-GPU node own exactly one KV head with no replication and no extra all-gather (see Distributed Training II: Tensor, Pipeline, Sequence & Expert Parallelism). Picking \(H_{kv} < \text{TP}\) forces the KV cache to be duplicated on every rank and silently gives back the memory you were trying to save.

import torch
import torch.nn as nn
import math

class GroupedQueryAttention(nn.Module):
    """
    Grouped Query Attention (GQA) as used in Llama 2/3, Qwen, Mistral.

    n_heads: number of query heads (H_q)
    n_kv_heads: number of key/value heads (H_kv), must divide n_heads evenly
    """

    def __init__(self, d_model: int, n_heads: int, n_kv_heads: int):
        super().__init__()
        assert n_heads % n_kv_heads == 0, "n_heads must be divisible by n_kv_heads"
        self.n_heads    = n_heads
        self.n_kv_heads = n_kv_heads
        self.n_rep      = n_heads // n_kv_heads   # how many Q heads per KV head
        self.head_dim   = d_model // n_heads

        # Query projects to n_heads * head_dim
        self.q_proj = nn.Linear(d_model, n_heads    * self.head_dim, bias=False)
        # Key/Value project to n_kv_heads * head_dim (smaller!)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)

    def _repeat_kv(self, x: torch.Tensor) -> torch.Tensor:
        """Expand KV from n_kv_heads to n_heads by repeating."""
        # x: (batch, seq, n_kv_heads, head_dim)
        if self.n_rep == 1:
            return x
        # Repeat along the head dimension n_rep times
        return x.unsqueeze(3).expand(
            *x.shape[:2], self.n_kv_heads, self.n_rep, self.head_dim
        ).reshape(*x.shape[:2], self.n_heads, self.head_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        head_dim = self.head_dim

        # Project to Q, K, V
        q = self.q_proj(x).view(B, T, self.n_heads,    head_dim)
        k = self.k_proj(x).view(B, T, self.n_kv_heads, head_dim)
        v = self.v_proj(x).view(B, T, self.n_kv_heads, head_dim)

        # Expand K and V to match n_heads
        k = self._repeat_kv(k)  # (B, T, n_heads, head_dim)
        v = self._repeat_kv(v)

        # Transpose to (B, n_heads, T, head_dim) for attention
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Scaled dot-product attention
        scale = math.sqrt(head_dim)
        attn = (q @ k.transpose(-2, -1)) / scale         # (B, n_heads, T, T)
        attn = torch.softmax(attn, dim=-1)
        out  = attn @ v                                   # (B, n_heads, T, head_dim)

        # Merge heads and project
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(out)

QK-Norm: Stabilizing Attention Logits at Scale¶

The problem: logit explosion¶

In deep, wide models trained for many tokens, the dot products \(q_i \cdot k_j\) can grow to very large values. Once the logits are large in magnitude, the softmax saturates: one token gets weight ~1 and all others get weight ~0. This “attention collapse” degrades the model’s ability to attend to multiple positions, and the large pre-softmax logits create numerical instability, especially in bf16 where the dynamic range is narrow. See also Numerical Computing, Floating Point & Precision for why bf16 overflow is a real concern.

QK-norm: normalize Q and K before attention¶

The fix — introduced for ViT-22B (Dehghani et al., “Scaling Vision Transformers to 22 Billion Parameters”, 2023), studied systematically as an instability remedy by Wortsman et al. (“Small-scale proxies for large-scale Transformer training instabilities”, 2023), and now standard in Gemma 3, OLMo 2, Chameleon and Qwen3 — is to apply RMSNorm to the query and key vectors before computing the attention scores:

After normalization each head’s query and key have RMS \(\approx \gamma\), so \(\|q\| \approx \|k\| \approx \sqrt{d_k}\,\gamma\) and by Cauchy–Schwarz \(|q \cdot k| \le d_k \gamma_q \gamma_k\); after the \(1/\sqrt{d_k}\) scale the logits are \(O(\sqrt{d_k}\,\gamma_q\gamma_k)\) — bounded by the learned scales rather than by whatever magnitude the projections happened to drift to over a trillion tokens. Because \(\gamma\) is learned per channel, expressivity is largely retained; only the unbounded growth is removed.

Two implementation details matter. First, order versus RoPE: the convention (and what HuggingFace’s Qwen3Attention and Gemma3Attention do) is project → normalize → rotate. Either order gives the same bound, since rotation preserves vector norms, but the two are not the same function once \(\gamma\) is per-channel — rotation mixes channel pairs — so a converter that swaps them will silently produce a different model. Second, the norm is over head_dim, not d_model: one tiny \(\gamma \in \mathbb{R}^{d_k}\) shared across heads (Qwen3, Gemma 3) or one per head (some variants). The shared version is what the code below implements, and it adds only \(2 d_k\) parameters per layer.

class QKNormAttention(nn.Module):
    """
    Attention with per-head QK normalization, as in Gemma 3, OLMo 2 and Qwen3.
    (Gemma 2 used logit soft-capping instead; Gemma 3 replaced it with this.)
    In HuggingFace these modules are literally named `q_norm` / `k_norm`.
    Prevents logit explosion during long training runs.
    """

    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.n_heads  = n_heads
        self.head_dim = d_model // n_heads
        self.q_proj   = nn.Linear(d_model, d_model, bias=False)
        self.k_proj   = nn.Linear(d_model, d_model, bias=False)
        self.v_proj   = nn.Linear(d_model, d_model, bias=False)
        self.o_proj   = nn.Linear(d_model, d_model, bias=False)

        # One RMSNorm per head for Q and K; per-head head_dim
        self.q_norm = RMSNorm(self.head_dim)
        self.k_norm = RMSNorm(self.head_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim)
        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim)

        # Apply QK-norm per head (RMSNorm operates on head_dim)
        q = self.q_norm(q)   # normalizes the last dimension
        k = self.k_norm(k)

        # Standard scaled dot-product attention from here
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        attn = torch.softmax((q @ k.transpose(-2,-1)) / math.sqrt(self.head_dim), dim=-1)
        out  = (attn @ v).transpose(1,2).contiguous().view(B, T, C)
        return self.o_proj(out)

No Biases, Tied/Untied Embeddings, and Initialization¶

Dropping biases¶

GPT-2 had bias terms in every linear projection and LayerNorm. Modern models like Llama, Qwen, and Mistral remove biases from all linear layers. The motivation is empirical: at large scale biases do not meaningfully improve loss (they represent a negligible fraction of parameters), but they complicate optimizer state memory (Adam maintains a first and second moment for every parameter, so biases add to optimizer memory with little benefit). For a 7B model, removing biases saves a few hundred MB of optimizer state — not huge, but free.

There is also a theoretical argument: when using pre-RMSNorm (which has no bias itself), the preceding biases in linear projections are redundant, as RMSNorm can represent any affine output of a linear layer with a bias by adjusting its scale. A stability argument closed the case: biases are the one part of the model that receives gradient regardless of the input, so they drift monotonically over long runs and are a recurring source of activation outliers that later break INT8/FP8 quantization (see Quantization I: Post-Training Quantization). Qwen2 kept a bias on the Q/K/V projections specifically to help length extrapolation; Qwen3 removed it and added QK-norm instead — the clearest single data point that the field now prefers normalizing over biasing.

Tied vs untied embeddings¶

The embedding table maps token IDs to vectors of dimension \(d\) and has shape \((V, d)\) where \(V\) is the vocabulary size. The output “unembedding” (also called the LM head or logit projection) maps from \(d\) back to \(V\), also shape \((V, d)\).

Tied embeddings (Press & Wolf, 2017) share the same weight matrix for both. This saves \(V \times d\) parameters — for a vocabulary of 32,000 and \(d=4096\), that is ~131M. Untied embeddings use separate matrices. The choice:

• Tied: fewer parameters, simpler to implement, forces the embedding space to be simultaneously useful for input representation and output scoring.
• Untied: more expressive; the input and output embedding spaces can specialize, and the LM head can develop the large-norm “unembedding” directions that tying would push back into the input representation.

The decision is essentially a budget fraction question, and the field splits cleanly by scale rather than by year. Tie when \(V d\) is a large slice of the model — GPT-2, all Gemma models, Llama 3.2 1B/3B, and Qwen3 0.6B–4B tie; untie when it is a rounding error — the Llama ½/3 models at 7B and above, and Qwen3 8B+, untie. A workable rule of thumb: if \(V d\) exceeds roughly 10% of your parameter budget, tie. At \(V=128{,}256\) and \(d=3072\) the table is 394M parameters, over 12% of a 3B model — tie. At \(d=8192\) in a 70B model it is 1.05B, under 2% — untie and buy the expressivity.

Logit Soft-Capping¶

The problem: logit skew in final layer¶

During pretraining, the final linear projection from hidden state to vocabulary logits can develop large-magnitude outputs for a small set of tokens. When the model becomes very confident, the cross-entropy gradient for non-predicted tokens effectively vanishes, slowing learning. At inference time, extreme logits after temperature scaling create numerical issues.

Soft-capping¶

Gemma 2 introduces logit soft-capping: a differentiable function that squashes logits toward a cap value \(c\) while preserving sign and relative ordering for logits well within the cap:

For small \(|z_i| \ll c\), \(\tanh(z_i/c) \approx z_i/c\), so \(\hat{z}_i \approx z_i\) — linear behavior near zero. For large \(|z_i|\), \(\hat{z}_i \to \pm c\) asymptotically. With \(c=30\) (Gemma 2’s setting), logits are prevented from exceeding 30 in magnitude.

Gemma 2 also applies soft-capping to per-layer attention logits (pre-softmax), with \(c=50\):

Currency note: soft-capping lost. Gemma 3 explicitly replaced Gemma 2’s soft-capping with QK-norm, citing Dehghani et al. (2023), Wortsman et al. (2023) and the Chameleon team. The reason is a systems reason rather than a quality one: a \(\tanh\) applied to the full \(T \times T\) logit matrix cannot be expressed inside a FlashAttention kernel, which by construction never materializes that matrix (see FlashAttention I: IO-Awareness & The Online Softmax). Enabling attention soft-capping in HuggingFace therefore forces the eager attention path, which is exactly the throughput you were trying to buy with long context. QK-norm, by contrast, is two cheap elementwise ops on \(Q\) and \(K\) before the kernel, so it composes with FlashAttention for free.

The other standard guard on the final logits is the z-loss, an auxiliary penalty \(\beta_z \cdot \operatorname{mean}\big(\log\sum_v e^{z_v}\big)^2\) with \(\beta_z \approx 10^{-4}\) (PaLM, ST-MoE, OLMo 2). It applies pressure on the log-partition function during training instead of clamping at inference, so it changes no forward-pass math and costs nothing at serving time — which is why most 2025–2026 models use z-loss rather than logit capping. It is derived, fused with cross-entropy, and chunked for memory in Training Stability, Loss Spikes & Debugging Large Runs. Learn soft-capping because you will meet it in Gemma 2 checkpoints and in interviews; reach for QK-norm plus z-loss when you build.

import torch

def soft_cap(logits: torch.Tensor, cap: float) -> torch.Tensor:
    """
    Differentiable logit soft-capping (Gemma 2).

    Maps logits smoothly toward [-cap, +cap]. For |z| << cap, behavior
    is approximately linear; for |z| >> cap, asymptotes to ±cap.

    Args
    ----
    logits : (..., vocab_size) raw pre-softmax logits
    cap    : scalar, e.g. 30.0 for final logits or 50.0 for attention logits

    Returns
    -------
    Capped logits of same shape as input.
    """
    return cap * torch.tanh(logits / cap)


# Demonstration
z = torch.tensor([-100.0, -30.0, -10.0, 0.0, 10.0, 30.0, 100.0])
capped = soft_cap(z, cap=30.0)
print("Raw    :", z.tolist())
print("Capped :", [f"{v:.2f}" for v in capped.tolist()])
# Output:  Capped: [-29.92, -22.85, -9.65, 0.00, 9.65, 22.85, 29.92]
# Note: 100 -> 29.92, 30 -> 22.85 (compressed, not clipped)

Attention Sinks and Sink Tokens¶

The “attention sink” phenomenon¶

Xiao et al. (“Efficient Streaming LLMs with Attention Sinks”, 2023) made an empirical observation: attention maps in trained autoregressive models show that some tokens — almost always the first few tokens in the context (especially BOS, Beginning-Of-Sequence) — receive anomalously high attention weights regardless of their content relevance. These are called attention sinks.

Why does this happen? Attention weights must sum to 1 via softmax. When no other token is relevant, the model needs somewhere to “dump” the probability mass so it can attend to nothing useful without the softmax distribution becoming uniform (which would average together all values and potentially corrupt the output). The initial tokens, having been seen in every context, become the designated garbage-collector for probability mass.

Implications for model design and inference¶

The attention sink phenomenon has two practical consequences:

1. Long-context window extension via StreamingLLM: if you want to process infinite-length streams, you can evict old KV cache entries safely — as long as you keep the first few tokens’ KV cache (the sinks). Dropping the sink tokens causes catastrophic loss spikes.

2. Deliberate sink token design: some models prepend a learnable or fixed “register” token that is never attended to in the output but acts as a sink for attention. This frees the model from overloading the BOS token.

3. Learned sink logits, no token at all: the cleanest 2025 formulation drops the sink token and instead gives each head a learned scalar \(s_h\) that participates in the softmax denominator but has no corresponding value vector — i.e. \(\text{softmax}\) over \([\,q\!\cdot\!k_1, \ldots, q\!\cdot\!k_t, s_h\,]\), with the last slot’s weight simply discarded. The head can now attend to nothing (all real weights near zero) without distorting any value, which is the “off-by-one softmax” idea made trainable. OpenAI’s gpt-oss models (2025) ship this as a per-head sinks parameter, and it is what GptOssAttention implements in HuggingFace transformers. Concretely, the softmax denominator becomes \(e^{s_h} + \sum_j e^{q\cdot k_j}\), costing one extra scalar per head and nothing at inference.

Depth vs Width Trade-offs¶

Why we go deep¶

Scaling laws (Kaplan et al., “Scaling Laws for Neural Language Models”, 2020; Hoffmann et al., “Training Compute-Optimal Large Language Models” / Chinchilla, 2022) largely treat architecture shape as secondary to total parameter count and compute budget. See Scaling Laws: Kaplan, Chinchilla & Beyond for the full story. But within a fixed parameter budget, how to allocate parameters between depth (number of layers) and width (hidden dimension, number of heads, MLP expansion) matters.

The dominant empirical finding is that depth is worth more than width up to a point. Deeper models can represent functions of exponentially higher complexity (via function composition) than wide-but-shallow ones. However, deeper models are harder to parallelize in pipeline parallel training (more micro-batches needed to fill the pipeline bubble) and have slower sequential KV-cache generation at inference.

Practical aspect ratios¶

The ratio \(L / d_{\text{model}}\) tends to be consistent across generations:

Modern 7B-class models favor roughly 32 layers with \(d=4096\), yielding an attention head dimension of 128 (with 32 heads). Larger models scale \(d\) and \(L\) roughly in proportion, holding \(d_k = 128\) fixed and adding heads: \(d_k\) below 64 wastes tensor-core tiles (which want \(\ge 64\) along the contraction dimension) and starves each head of capacity, while \(d_k\) above 256 is unsupported by most FlashAttention builds. So the practical knobs are \(L\), \(H_q\) and \(H_{kv}\), with \(d = H_q \cdot d_k\) falling out.

Below ~1B parameters the trade-off tilts noticeably toward depth: MobileLLM (Liu et al., 2024) ablated shape at fixed parameter count for sub-billion models and found deeper-and-thinner consistently wins, which is why the capstone’s Stack-100M chooses \(d=512\) with 30 layers (\(L/d \approx 0.059\), six times “deeper” by this metric than a 7B model). See The Stack-100M Architecture for that derivation in full.

def llama_param_count(d, L, n_heads, n_kv, d_ff, vocab, head_dim=128, tied=True):
    # Q + O are full width; K + V are GQA-width
    attn  = 2 * d * (n_heads * head_dim) + 2 * d * (n_kv * head_dim)
    mlp   = 3 * d * d_ff          # gate, up, down (SwiGLU)
    norms = 2 * d                 # 2 RMSNorm per layer (gamma only)
    per_layer = attn + mlp + norms
    emb = vocab * d if tied else 2 * vocab * d
    return per_layer * L + emb + d   # + final RMSNorm

n = llama_param_count(d=3072, L=28, n_heads=24, n_kv=8,
                      d_ff=8192, vocab=128256, tied=True)
print(f"{n:,}")   # 3,212,749,824  -> ~3.21B, matches Llama 3.2 3B

The Modern Transformer Recipe¶

We now have enough pieces to assemble a complete reference. The table below summarizes the consensus choices circa 2024–2026, distinguishing “near-universal” (adopted by nearly all recent models) from “model-specific” (used by some, absent in others).

"""
Putting it all together: a from-scratch ModernTransformerBlock.
This implements the full recipe: pre-RMSNorm, SwiGLU, GQA with RoPE,
no biases, QK-norm optional. Runnable as a standalone unit test.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps
                               ).type_as(x) * self.weight


class SwiGLUFFN(nn.Module):
    def __init__(self, d_model: int, intermediate: int):
        super().__init__()
        self.gate = nn.Linear(d_model, intermediate, bias=False)
        self.up   = nn.Linear(d_model, intermediate, bias=False)
        self.down = nn.Linear(intermediate, d_model,  bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))


class ModernTransformerBlock(nn.Module):
    """
    A single Transformer block implementing the modern (2024-era) recipe:
      - Pre-RMSNorm on both attention and MLP sublayers
      - GQA-style attention (n_kv_heads <= n_heads)
      - SwiGLU FFN
      - No biases anywhere
      - Optional QK-norm

    Does NOT include RoPE application here for brevity; in a full model
    RoPE would be applied to Q and K after projection.
    """

    def __init__(
        self,
        d_model: int,
        n_heads: int,
        n_kv_heads: int,
        intermediate: int,
        qk_norm: bool = False,
        eps: float = 1e-5,
    ):
        super().__init__()
        assert d_model % n_heads == 0
        assert n_heads % n_kv_heads == 0

        self.n_heads    = n_heads
        self.n_kv_heads = n_kv_heads
        self.n_rep      = n_heads // n_kv_heads
        self.head_dim   = d_model // n_heads

        # Pre-norm before attention
        self.attn_norm = RMSNorm(d_model, eps)
        # Attention projections — no bias
        self.q_proj = nn.Linear(d_model, n_heads    * self.head_dim, bias=False)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(d_model, d_model,                    bias=False)

        # Optional QK-norm
        self.qk_norm = qk_norm
        if qk_norm:
            self.q_norm = RMSNorm(self.head_dim, eps)
            self.k_norm = RMSNorm(self.head_dim, eps)

        # Pre-norm before MLP
        self.ffn_norm = RMSNorm(d_model, eps)
        self.ffn      = SwiGLUFFN(d_model, intermediate)

        # Scaled residual initialization for stability in deep models
        # Scale output projections by 1/sqrt(2*n_layers); caller should
        # set this post-construction, e.g.:
        #   for block in blocks:
        #       block.o_proj.weight.data.mul_(1.0 / math.sqrt(2 * n_layers))
        #       block.ffn.down.weight.data.mul_(1.0 / math.sqrt(2 * n_layers))

    def _repeat_kv(self, x: torch.Tensor) -> torch.Tensor:
        if self.n_rep == 1:
            return x
        return x.unsqueeze(3).expand(
            *x.shape[:2], self.n_kv_heads, self.n_rep, self.head_dim
        ).reshape(*x.shape[:2], self.n_heads, self.head_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        # --- Attention sublayer (pre-norm) ---
        h = self.attn_norm(x)
        q = self.q_proj(h).view(B, T, self.n_heads,    self.head_dim)
        k = self.k_proj(h).view(B, T, self.n_kv_heads, self.head_dim)
        v = self.v_proj(h).view(B, T, self.n_kv_heads, self.head_dim)

        if self.qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)

        k = self._repeat_kv(k)
        v = self._repeat_kv(v)

        # (B, H, T, d_k)
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Causal mask for autoregressive decoding
        mask = torch.tril(torch.ones(T, T, device=x.device)).bool()
        attn_logits = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attn_logits = attn_logits.masked_fill(~mask, float('-inf'))
        attn_w = torch.softmax(attn_logits, dim=-1)

        out = (attn_w @ v).transpose(1, 2).contiguous().view(B, T, C)
        x = x + self.o_proj(out)          # residual connection

        # --- MLP sublayer (pre-norm) ---
        x = x + self.ffn(self.ffn_norm(x))  # residual connection

        return x


# --- Integration test ---
torch.manual_seed(0)
block = ModernTransformerBlock(
    d_model=512, n_heads=8, n_kv_heads=2, intermediate=1365,  # 8/3 * 512 ≈ 1365
    qk_norm=True
)
x = torch.randn(2, 16, 512)  # batch=2, seq=16, d_model=512
y = block(x)
print(f"Block output shape: {y.shape}")          # (2, 16, 512)
print(f"Output norm (should be finite): {y.norm().item():.3f}")
param_count = sum(p.numel() for p in block.parameters())
print(f"Block parameter count: {param_count:,}")  # 2,753,152 -> ~2.75M
# Check by hand: attention = 512*512 (Q) + 2*512*128 (K,V, GQA) + 512*512 (O)
#              = 655,360;  SwiGLU = 3*512*1365 = 2,096,640;
#              norms = 2*512 + 2*64 (qk_norm) = 1,152.  Total 2,753,152.

The Recipe in Practice: Config Files, Fused Kernels & Reference Repos¶

We opened by promising that a modern config.json would stop being alphabet soup. Here is the actual configuration of Llama 3.2 3B — the model we sized by hand earlier — with every field now decodable:

{
  "architectures": ["LlamaForCausalLM"],
  "hidden_size": 3072,
  "intermediate_size": 8192,
  "num_hidden_layers": 28,
  "num_attention_heads": 24,
  "num_key_value_heads": 8,
  "hidden_act": "silu",
  "rms_norm_eps": 1e-05,
  "rope_theta": 500000.0,
  "rope_scaling": {
    "rope_type": "llama3",
    "factor": 32.0,
    "low_freq_factor": 1.0,
    "high_freq_factor": 4.0,
    "original_max_position_embeddings": 8192
  },
  "max_position_embeddings": 131072,
  "attention_bias": false,
  "tie_word_embeddings": true,
  "vocab_size": 128256,
  "torch_dtype": "bfloat16"
}

Read it against this chapter: rms_norm_eps is the \(\epsilon\) of our RMSNorm (pre-norm placement is implied by the architecture class, not by a flag); hidden_act: "silu" with intermediate_size \(\ne 4d\) tells you it is SwiGLU, since a plain FFN would have no third matrix and would sit at \(4 \times 3072 = 12288\); num_key_value_heads: 8 against 24 query heads is GQA with \(G=3\); attention_bias: false is the no-bias decision; tie_word_embeddings: true is the small-model budget call (\(128256 \times 3072 = 394\)M would otherwise be paid twice); rope_theta: 500000 plus rope_scaling is long-context RoPE, trained at 8192 and extended 32× by frequency-dependent interpolation (see Long-Context Pretraining & Context Extension). There is no q_norm field because QK-norm in transformers is expressed as submodules in the attention class (Qwen3, Gemma 3, OLMo 2) rather than a boolean.

Instantiating that recipe — no download required — is three lines:

from transformers import LlamaConfig, LlamaForCausalLM, AutoModelForCausalLM

# Route 1: build the architecture from the recipe, randomly initialized.
cfg = LlamaConfig(
    hidden_size=3072, intermediate_size=8192, num_hidden_layers=28,
    num_attention_heads=24, num_key_value_heads=8,      # GQA, G=3
    hidden_act="silu", rms_norm_eps=1e-5,               # SwiGLU + RMSNorm
    rope_theta=500000.0, attention_bias=False,          # RoPE, no biases
    tie_word_embeddings=True, vocab_size=128256,
)
model = LlamaForCausalLM(cfg)
print(f"{sum(p.numel() for p in model.parameters()):,}")  # ~3.21B, matching our hand count

# Route 2: load real weights and read the same knobs back off the config.
# model = AutoModelForCausalLM.from_pretrained(
#     "meta-llama/Llama-3.2-3B",
#     attn_implementation="flash_attention_2",  # or "sdpa" (default) / "eager"
#     dtype="bfloat16",
# )

That attn_implementation argument is where this chapter meets Part IV. "eager" is the hand-rolled softmax we wrote above; "sdpa" dispatches to torch.nn.functional.scaled_dot_product_attention; "flash_attention_2" calls the flash-attn package’s kernel directly. Only the first can express logit soft-capping — which is precisely why Gemma 3 dropped it.

Where the fused kernels live¶

The from-scratch modules in this chapter are for understanding. In a real training loop every one of them has a fused counterpart, and using them is not optional at scale:

The last row matters more than it looks: it never materializes the \((B, T, V)\) logit tensor, which at \(BT=16384\) and \(V=128256\) would be 8.4 GB in fp32 — often the single largest tensor in a small-model training step. enable_gqa=True similarly avoids materializing the repeated K and V, which our teaching code does materialize. Liger-Kernel is written in Triton (see Writing GPU Kernels with Triton); torch.compile will fuse the simpler ones (RMSNorm, SwiGLU) for you without any extra dependency.

import torch
import torch.nn.functional as F

# The production form of the GQA forward pass: no _repeat_kv materialization,
# causal masking inside the kernel, and an IO-aware fused backend under the hood.
B, T, Hq, Hkv, dk = 2, 128, 24, 8, 128
q = torch.randn(B, Hq,  T, dk)      # (B, H_q,  T, d_k)
k = torch.randn(B, Hkv, T, dk)      # (B, H_kv, T, d_k)  <-- 3x smaller
v = torch.randn(B, Hkv, T, dk)

out = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True)
print(out.shape)                    # (2, 24, 128, 128) -- broadcast over the group

Reference implementations worth reading¶

When you need a correct, complete, trainable version of this recipe rather than a teaching one, these are the canonical open-source starting points, in rough order of increasing machinery: nanoGPT (Karpathy — the minimal GPT-2 training loop, the ancestor of Building a GPT From Scratch); litgpt (Lightning AI — clean single-file implementations of ~20 modern architectures); torchtitan (PyTorch’s own reference for pretraining Llama-family models with FSDP2 + tensor/pipeline parallelism); torchtune (PyTorch’s fine-tuning library, with the architectures expressed as composable builder functions); nanotron (HuggingFace’s minimal 3D-parallel pretrainer); and Megatron-LM / DeepSpeed for the industrial path, covered in Megatron-LM, DeepSpeed & Parallelism in Practice. All six implement exactly the components of this chapter; the differences between them are parallelism and ergonomics, not architecture.

Further Reading¶

• Zhang and Sennrich, “Root Mean Square Layer Normalization,” NeurIPS 2019 — the RMSNorm paper.
• Noam Shazeer, “GLU Variants Improve Transformer,” arXiv 2020 — SwiGLU and GeGLU.
• Su et al., “RoFormer: Enhanced Transformer with Rotary Position Embedding,” arXiv 2021 — original RoPE paper.
• Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints,” EMNLP 2023 — GQA, including a method for uptraining existing MHA checkpoints.
• Touvron et al., “Llama 2: Open Foundation and Fine-Tuned Chat Models,” arXiv 2023 — a comprehensive blueprint of the modern recipe in an openly released model.
• Team Gemma, “Gemma 2: Improving Open Language Models at a Practical Size,” Google DeepMind 2024 — logit soft-capping, sandwich normalization, and interleaved sliding-window attention; and “Gemma 3 Technical Report,” 2025, which swaps soft-capping for QK-norm.
• Wortsman et al., “Small-scale proxies for large-scale Transformer training instabilities,” 2023 — QK-norm, z-loss and logit-growth diagnostics studied at a scale you can actually reproduce.
• Press and Wolf, “Using the Output Embedding to Improve Language Models,” EACL 2017 — the original weight-tying argument.
• Liu et al., “MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases,” ICML 2024 — the deep-and-thin aspect-ratio result for small models.
• Xiao et al., “Efficient Streaming LLMs with Attention Sinks,” ICLR 2024 — attention sink phenomenon and StreamingLLM.
• Kaplan et al., “Scaling Laws for Neural Language Models,” arXiv 2020 — depth vs width trade-offs in the context of scaling.
• Meta AI, Llama 3 model card and technical report, 2024 — updated RoPE theta and architectural decisions at 70B/405B scale.

Exercises¶

1. (Conceptual) LayerNorm applies \(\frac{x-\mu}{\sqrt{\sigma^2+\epsilon}}\cdot\gamma + \beta\), whereas RMSNorm applies \(\frac{x}{\text{RMS}(x)}\cdot\gamma\) with no \(\beta\). (a) Which two operations present in LayerNorm does RMSNorm discard, and which one does it keep? (b) Zhang & Sennrich justified the discard empirically — what did their ablation find? © In the chapter’s RMSNorm.forward, why is the input cast to float32 before computing the norm and cast back afterward?

2. (Quantitative) You are configuring a SwiGLU FFN for a model with hidden dimension \(d = 4096\). A vanilla two-matrix FFN would use a 4x expansion (intermediate \(= 4d\)). (a) Compute the parameter count of that vanilla FFN. (b) Using the iso-parameter rule for SwiGLU’s three matrices, compute the target intermediate dimension, then round it to the nearest multiple of 256. © Compute the SwiGLU FFN’s parameter count at that rounded dimension and confirm it is close to the vanilla count.

3. (Quantitative) A decoder-only model has \(L = 48\) layers, \(H_q = 40\) query heads, head dimension \(d_k = 128\), and uses GQA with \(H_{kv} = 8\) KV heads. It runs in bf16 (2 bytes/element). (a) Compute the KV-cache size per token per layer. (b) Compute the total KV cache for a context of 16,384 tokens across all layers, in GiB. © What would the same total be with full MHA (\(H_{kv} = H_q = 40\)), and what is the reduction factor?

4. (Conceptual) Llama 1 used rope_theta = 10000; Llama 3 raised it to rope_theta = 500000 specifically to improve long-context behavior. (a) In the frequency formula \(\theta_i = \text{base}^{-2i/d_k}\), what happens to the set of rotation frequencies as the base increases from 10,000 to 500,000? (b) Why does that change help a model extrapolate to positions beyond its training length? © RoPE is described as giving a “relative position” property — state precisely what quantity the attention dot product \(q_m^\top k_n\) ends up depending on, and why that property is desirable.

5. (Implementation) Gemma 2 applies logit soft-capping in two places: on the attention logits (pre-softmax, \(c=50\)) and on the final vocabulary logits (\(c=30\)), using \(\hat z = c\cdot\tanh(z/c)\). Starting from the chapter’s ModernTransformerBlock, modify its forward so that attention logits are soft-capped at \(c=50\). Then (a) explain why the cap must be applied before the causal -inf mask, and (b) verify by hand what an uncapped logit of \(75\) becomes after soft-capping at \(c=50\).

import torch
import math

def soft_cap(x: torch.Tensor, cap: float) -> torch.Tensor:
    # Smoothly squashes x toward [-cap, +cap]; ~linear for |x| << cap.
    return cap * torch.tanh(x / cap)

def forward(self, x: torch.Tensor, attn_logit_cap: float = 50.0) -> torch.Tensor:
    B, T, C = x.shape
    # --- Attention sublayer (pre-norm) ---
    h = self.attn_norm(x)
    q = self.q_proj(h).view(B, T, self.n_heads,    self.head_dim)
    k = self.k_proj(h).view(B, T, self.n_kv_heads, self.head_dim)
    v = self.v_proj(h).view(B, T, self.n_kv_heads, self.head_dim)

    if self.qk_norm:
        q = self.q_norm(q)
        k = self.k_norm(k)

    k = self._repeat_kv(k)
    v = self._repeat_kv(v)

    q = q.transpose(1, 2)
    k = k.transpose(1, 2)
    v = v.transpose(1, 2)

    attn_logits = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)

    # >>> soft-cap the attention logits (Gemma 2, c=50) BEFORE masking <<<
    attn_logits = soft_cap(attn_logits, attn_logit_cap)

    mask = torch.tril(torch.ones(T, T, device=x.device)).bool()
    attn_logits = attn_logits.masked_fill(~mask, float('-inf'))
    attn_w = torch.softmax(attn_logits, dim=-1)

    out = (attn_w @ v).transpose(1, 2).contiguous().view(B, T, C)
    x = x + self.o_proj(out)

    # --- MLP sublayer (pre-norm) ---
    x = x + self.ffn(self.ffn_norm(x))
    return x