17. A modern architecture: Qwen3

From GPT-2 to a 2025 model

The transformer block hasn’t changed shape since GPT-2: embed, then $L$ blocks of attention and MLP on a residual stream, then norm and head. What changed is the ingredients. Almost every open model since Llama (2023), including Qwen, Mistral and DeepSeek, made the same set of substitutions:

Dense Qwen3 is the first real model your engine will run. Here’s the contract for the 0.6B checkpoint, from its official config.json:

Notice that the query projection has width $16 \times 128 = 2{,}048$, which is twice the hidden size. Code that assumes head_dim = hidden_size / heads reshapes the projections wrongly and crashes (or worse, doesn’t). The book’s tiny default config deliberately has the same mismatch, so the bug surfaces in the tests.

RMSNorm

LayerNorm subtracts the mean and divides by the standard deviation. RMSNorm (Zhang and Sennrich, 2019) skips the mean and divides by the root mean square:

$$ \operatorname{RMSNorm}(x)_i = \gamma_i, \frac{x_i}{\sqrt{\frac1D\sum_j x_j^2 + \epsilon}} . $$

For $x = [3, 4]$: the mean square is $(9 + 16)/2 = 12.5$, the RMS is 3.5355, and the output is $[0.8485, 1.1314]$ before $\gamma$. The vector’s direction is unchanged, and its RMS becomes 1. It’s one reduction instead of two, with no bias, and models train just as well.

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

    def forward(self, x):
        """x / sqrt(mean(x^2) + eps) * weight, reduced in FP32.  (Your engine: Chapter 17)"""
        x32 = x.float()
        normalized = x32 * torch.rsqrt(x32.square().mean(-1, keepdim=True) + self.eps)
        return self.weight * normalized.to(x.dtype)

Two parity details matter when loading real weights:

1. Compute the mean of squares in FP32 even for BF16 inputs (Chapter 13).
2. The cast order is part of the contract. Qwen3 normalizes in FP32, casts back to the input dtype, then multiplies by $\gamma$. Gemma and the Qwen3.5 family use $(1+\gamma)$ with $\gamma$ initialized at zero, multiplying before the cast (you’ll meet that “zero-centered” variant in Chapter 29). In BF16, the difference changes the last bits of every layer’s output.

Rotary position embeddings

GPT-2 added a learned position vector to each token. RoPE (Su et al., 2021) instead rotates the query and key vectors by an angle proportional to their position, right before the attention scores are computed.

Take one pair of coordinates and rotate it by angle $\phi$:

$$ R(\phi) = \begin{pmatrix} \cos\phi & -\sin\phi \ \sin\phi & \cos\phi \end{pmatrix}. $$

Rotate a query at position $p$ by $p,\omega$ and a key at position $s$ by $s,\omega$. Their dot product is

$$ \big(R(p\omega),q\big)\cdot\big(R(s\omega),k\big) = q^{\top} R\big((s-p),\omega\big), k , $$

because rotations compose and $R(a)^\top = R(-a)$. The score depends only on the distance $s - p$, not on absolute positions. Shift both tokens by 10 and nothing changes. For example, with $q = k = [1, 0]$ and $\omega = 0.35$, two tokens 2 apart score $\cos(0.7) = 0.7648$, at positions (0, 2) or (10, 12) alike.

A head of dimension $d$ has $d/2$ coordinate pairs, each rotating at its own frequency:

$$ \omega_i = \theta^{-2i/d}, \qquad i = 0, \ldots, d/2 - 1 . $$

With $\theta = 10^6$ and $d = 128$, the fastest pair turns 1 radian per token and the slowest about $1.2 \times 10^{-6}$ radians, a wavelength of about 5 million tokens. Fast pairs resolve nearby order; slow pairs let distant tokens still be distinguished. Raising $\theta$ (Qwen3 uses $10^6$, the original paper $10^4$) stretches every wavelength, which helps long contexts.

The pairing convention is part of the contract

Which coordinates form a pair? The original paper paired adjacent coordinates $(0,1), (2,3), \ldots$ Llama, Qwen and most Hugging Face models pair coordinate $i$ with $i + d/2$, the split-half (“rotate half”) convention. Both are valid RoPEs, and a model is trained with exactly one. Using the other produces correct shapes and fluent-looking garbage. Our implementation uses split-half:

def rope_cos_sin(positions, rotary_dim, theta, dtype=torch.float32):
    """cos/sin tables for absolute positions [T] or [B, T] -> [B or 1, 1, T, rotary_dim].  (Your engine: Chapter 17)

    Frequencies theta^(-2i/rotary_dim) for i in [0, rotary_dim/2), repeated for the two halves.
    """
    if rotary_dim % 2:
        raise ValueError("RoPE needs an even number of rotated dimensions")
    inv_freq = theta ** (-torch.arange(0, rotary_dim, 2, device=positions.device, dtype=torch.float32) / rotary_dim)
    pos = positions.float() if positions.ndim == 2 else positions.float()[None]
    angles = pos[..., None] * inv_freq                    # [B, T, rotary_dim/2]
    angles = torch.cat((angles, angles), dim=-1)[:, None]  # [B, 1, T, rotary_dim]
    return angles.cos().to(dtype), angles.sin().to(dtype)


def rotate_half(x):
    a, b = x.chunk(2, dim=-1)
    return torch.cat((-b, a), dim=-1)


def apply_rope(x, cos, sin):
    """Rotate the first cos.shape[-1] features of x (split-half pairing); pass the rest through.  (Your engine: Chapter 17)

    Pairs are (i, i + rotary_dim/2), NOT adjacent features (2i, 2i+1): that is the checkpoint's
    convention and getting it wrong still produces plausible-looking shapes.
    """
    rotary = cos.shape[-1]
    x_rot, x_pass = x[..., :rotary], x[..., rotary:]
    rotated = x_rot * cos + rotate_half(x_rot) * sin
    return torch.cat((rotated, x_pass), dim=-1) if x_pass.shape[-1] else rotated

// Split-half pairing: feature i rotates with feature i + d/2.
inline void rope(float* x, size_t d, size_t pos, float theta) {
    for (size_t i = 0; i < d / 2; ++i) {
        float freq = std::pow(theta, -2.f * i / d), s = std::sin(pos * freq), c = std::cos(pos * freq);
        float a = x[i], b = x[i + d / 2];
        x[i] = a * c - b * s;
        x[i + d / 2] = a * s + b * c;
    }
}

#![allow(unused)]
fn main() {
/// Rotate one head vector in place for absolute position `pos`, pairing feature i with
/// feature i + d/2 (the split-half convention used by Qwen, Llama and the book's Python code).
pub fn rope(x: &mut [f32], pos: usize, theta: f32) {
    let half = x.len() / 2;
    for i in 0..half {
        let freq = theta.powf(-2.0 * i as f32 / x.len() as f32);
        let (sin, cos) = (pos as f32 * freq).sin_cos();
        let (a, b) = (x[i], x[i + half]);
        x[i] = a * cos - b * sin;
        x[i + half] = a * sin + b * cos;
    }
}
}

apply_rope rotates only the first cos.shape[-1] features and passes the rest through. Qwen3 rotates all 128, but Flash-Next’s attention rotates only the first 64 of 256 (partial RoPE, Chapter 29). The same function handles both.

Where RoPE sits in the computation matters for caching: Q and K are rotated before K is cached. A cached key already contains its position. Never rotate cached keys again; each new token rotates only its own Q and K, at its own position.

Attention with grouped KV heads and Q/K normalization

Qwen3’s attention is your Chapter 5 attention with three additions:

1. Separate Q, K, V projections, no biases, with output widths $H_q D_h$, $H_{kv} D_h$, $H_{kv} D_h$.
2. Q/K norm: an RMSNorm over each head’s 128 features, with a learned $\gamma$ shared across heads, applied to Q and K before RoPE. This keeps attention logits from growing without bound during training, a known source of instability.
3. GQA: 16 query heads share 8 KV heads; your causal_attention handles the mapping.

class Qwen3Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        if cfg.num_attention_heads % cfg.num_key_value_heads:
            raise ValueError("Query heads must be a multiple of KV heads")
        self.cfg = cfg
        d, hq, hkv = cfg.head_dim, cfg.num_attention_heads, cfg.num_key_value_heads
        self.q_proj = nn.Linear(cfg.hidden_size, hq * d, bias=False)
        self.k_proj = nn.Linear(cfg.hidden_size, hkv * d, bias=False)
        self.v_proj = nn.Linear(cfg.hidden_size, hkv * d, bias=False)
        self.o_proj = nn.Linear(hq * d, cfg.hidden_size, bias=False)
        self.q_norm = RMSNorm(d, cfg.rms_norm_eps)
        self.k_norm = RMSNorm(d, cfg.rms_norm_eps)

    def forward(self, x, positions, rope, cache=None, layer=0, rows=None):
        """Project, per-head RMSNorm on q and k, RoPE, cache, GQA attention, output projection.  (Your engine: Chapter 17)"""
        b, t, _ = x.shape
        c = self.cfg
        q = self.q_proj(x).view(b, t, c.num_attention_heads, c.head_dim)
        k = self.k_proj(x).view(b, t, c.num_key_value_heads, c.head_dim)
        v = self.v_proj(x).view(b, t, c.num_key_value_heads, c.head_dim)
        q = self.q_norm(q).transpose(1, 2)     # norm over head_dim, before RoPE
        k = self.k_norm(k).transpose(1, 2)
        v = v.transpose(1, 2)
        cos, sin = rope
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        key_positions = None
        if cache is not None:
            k, v, key_positions = cache.update(layer, k, v, positions, rows)   # keys are cached already rotated
        y = causal_attention(q, k, v, positions, key_positions)
        return self.o_proj(y.transpose(1, 2).reshape(b, t, -1))

SwiGLU

The GPT-2 MLP was expand → GELU → contract. Qwen3’s MLP has two expansions, one of which gates the other (Shazeer, 2020):

$$ \operatorname{SwiGLU}(x) = W_{\text{down}}\big(\operatorname{SiLU}(W_{\text{gate}},x) \odot W_{\text{up}},x\big), \qquad \operatorname{SiLU}(z) = z,\sigma(z). $$

The gate decides, feature by feature, how much of the “up” signal passes. If a gate pre-activation is 0, $\operatorname{SiLU}(0) = 0$ and the feature is shut off. If it’s 2, $\operatorname{SiLU}(2) = 1.7616$, and with an up value of 3 the feature passes $5.2848$. Gated MLPs consistently outperform ungated ones at the same parameter count. To keep that count comparable with a $4D$ GELU MLP, the intermediate size is usually about $\frac{8}{3}D$ (Qwen3-0.6B uses $3D$).

class SwiGLU(nn.Module):
    def __init__(self, width, hidden):
        super().__init__()
        self.gate_proj = nn.Linear(width, hidden, bias=False)
        self.up_proj = nn.Linear(width, hidden, bias=False)
        self.down_proj = nn.Linear(hidden, width, bias=False)

    def forward(self, x):
        """down(silu(gate(x)) * up(x)).  (Your engine: Chapter 17)"""
        return self.down_proj(nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))

The whole model

class Qwen3Layer(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.self_attn = Qwen3Attention(cfg)
        self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.mlp = SwiGLU(cfg.hidden_size, cfg.intermediate_size)

    def forward(self, x, positions, rope, cache=None, layer=0, rows=None):
        x = x + self.self_attn(self.input_layernorm(x), positions, rope, cache, layer, rows)
        return x + self.mlp(self.post_attention_layernorm(x))

class Qwen3(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.model = Qwen3Backbone(cfg)
        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
        if cfg.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

    @property
    def context_limit(self):
        return self.cfg.max_position_embeddings

    def forward(self, ids, cache=None, positions=None, rows=None, return_hidden=False):
        """ids [B, T] -> logits [B, T, vocab]. Same cache/position contract as GPT.forward.  (Your engine: Chapter 17)"""
        if positions is None:
            start = cache.length if cache is not None else 0
            if start + ids.shape[1] > self.context_limit:
                raise ValueError("Sequence exceeds max_position_embeddings")
            positions = torch.arange(start, start + ids.shape[1], device=ids.device)
        x = self.model.embed_tokens(ids)
        rope = rope_cos_sin(positions, self.cfg.head_dim, self.cfg.rope_theta, x.dtype)
        for i, layer in enumerate(self.model.layers):
            x = layer(x, positions, rope, cache, i, rows)
        hidden = self.model.norm(x)
        return hidden if return_hidden else self.lm_head(hidden)

    def cache_spec(self):
        c = self.cfg
        return c.num_hidden_layers, c.num_key_value_heads, c.head_dim

    def new_cache(self, batch, capacity):
        if capacity > self.context_limit:
            raise ValueError("Requested cache exceeds the context limit")
        c, p = self.cfg, self.model.embed_tokens.weight
        return KVCache(c.num_hidden_layers, batch, c.num_key_value_heads, capacity, c.head_dim, p.device, p.dtype)

The parameter budget of Qwen3-0.6B, counted from this code:

A quarter of this small model is the embedding table. In bigger models, the 28 × 15.7M of layers dominates.

Why parity tests matter now

From now on, your models must match someone else’s trained weights exactly, and small discrepancies don’t announce themselves. A wrong RoPE pairing, an epsilon of 1e-5 instead of 1e-6, LayerNorm instead of RMSNorm, or norm-then-RoPE swapped to RoPE-then-norm all produce a model that runs, generates English-like text after loading, and is subtly or badly wrong. The milestone tests therefore compare your Qwen3 against Hugging Face’s on small random configurations. Random weights exercise every code path, and a 1e-5 match on random weights is strong evidence of a correct implementation.

Build it

uv pip install -r optional-requirements.txt     # for the comparison with Transformers
pytest tests/test_ch17_qwen3.py
python run.py cache --impl engine                # your Qwen3, cached vs uncached

Stretch exercises

1. ★ Verify numerically that RoPE preserves vector norms, and that the score at distance 0 is just the unrotated dot product. Where: experiments/ch17.py (create it), importing engine.qwen3.rope_cos_sin and apply_rope.
2. ★★ Implement the adjacent-pair RoPE convention and show that loading Qwen3 weights with it changes the logits (by how much?). Where: add an adjacent-pair variant of apply_rope in engine/qwen3.py and select it in Qwen3Attention.forward.
3. ★★ Implement YaRN or linear RoPE scaling (multiply positions by a factor, or interpolate frequencies) and explain what it does to the wavelength table. Where: rope_cos_sin in engine/qwen3.py.
4. ★★★ Write a fused Triton kernel that applies Q/K-norm and RoPE to the Q and K projections in one pass, and plug it into Qwen3Attention behind a backend switch. Where: new engine/kernels/triton_rope.py, selected by Qwen3Attention.forward in engine/qwen3.py.

Check your understanding

1. Why can a correct-looking reshape still be wrong for Qwen3’s attention projections?
2. What happens to a key, in order, before it enters the cache?
3. What property of rotations makes a shared position shift cancel out in a RoPE dot product?
4. Why does SwiGLU have three weight matrices where GPT-2’s MLP has two?
5. Why are random-weight parity tests enough to establish the architecture, before you ever load the real checkpoint?

Going deeper

• BALLM has no Qwen3 chapter, but Raschka’s companion repository does: ch05/11_qwen3 builds Qwen3 from scratch in the book’s style, and ch05/07_gpt_to_llama walks the GPT-2 → Llama changes one by one.
• Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding (2021); Zhang and Sennrich, Root Mean Square Layer Normalization (2019); Shazeer, GLU Variants Improve Transformer (2020).
• Qwen Team, Qwen3 Technical Report (2025); the official modeling_qwen3.py in Transformers.