The transformer block

Attention lets a token gather information; it does not, on its own, make a network. The transformer block is the repeating unit that stacks around attention to make a working model: a normalization, the attention sublayer, a residual connection, another normalization, a feed-forward sublayer, another residual. A modern block (Qwen3, Llama, gpt-oss) differs from the 2017 original in four specific, deliberate ways, and this chapter derives all four because each is load-bearing for the memory and speed arguments later. The feed-forward became a gated MLP (SwiGLU). LayerNorm became RMSNorm. Absolute position embeddings became rotary (RoPE), which I derive as an actual rotation. And multi-head attention became grouped-query attention (GQA), whose entire justification is KV-cache pressure on cards exactly like the 16GB one. By the end I will assemble a block from these parts and match it, dimension for dimension, against a real Qwen3 config.

Theory

The residual stream is the backbone

Before the parts, the skeleton. A transformer block does not replace its input with the sublayer output; it adds the sublayer output to its input:

$$x^{\prime }=x+\mathrm{Attn}(\mathrm{Norm}(x)),x^{\prime \prime }=x^{\prime }+\mathrm{MLP}(\mathrm{Norm}(x^{\prime })).\text{\hspace{1em}(4.1)}$$

That "$x+$" is the residual connection, and reading equation (4.1) as an addition rather than a transformation is the single most useful frame for the whole architecture. The input $x$ flows down the stack essentially untouched (the "residual stream"), and each sublayer reads from it, computes something, and writes an increment back by addition. The stream is a running sum; a 36-layer model is 72 sublayers each adding their two cents. This matters for two reasons that recur all book. First, gradients: because the forward pass adds, the backward pass through $x+f(x)$ has Jacobian $I+f^{\prime }$, so the identity term $I$ guarantees a clean gradient path straight down the stream even when $f^{\prime }$ is small, this is why deep transformers train at all, and it is the residual restatement of the "identity term keeps gradients alive" idea from the autograd chapter. Second, the norm placement in equation (4.1) puts Norm inside the sublayer, before attention and before the MLP, never on the residual stream itself. This is pre-normalization, and it is what modern models use because it keeps the residual stream un-normalized and the training stable; the original transformer put the norm after the addition (post-norm) and was famously twitchy to train without warmup.

RMSNorm

Every sublayer's input gets normalized first, to keep activations at a controlled scale so the next matmul does not see wildly varying magnitudes. LayerNorm did this by subtracting the mean and dividing by the standard deviation across the feature dimension, then applying a learned scale and shift. RMSNorm drops the mean-centering and the shift entirely, keeping only the scaling by root-mean-square:

$$\mathrm{RMSNorm}(x{)}_i=\frac{x_i}{\sqrt{\frac{1}{d}{\sum }_{j=1}^d{x}_j^2+ϵ}}\cdot g_i,\text{\hspace{1em}(4.2)}$$

where $g\in {\mathbb{R}}^d$ is a learned per-feature gain and $ϵ$ (around $10^{-6}$) guards the square root. The denominator is the root-mean-square of the feature vector, hence the name. The justification for dropping the mean subtraction is partly empirical (it works as well) and partly efficiency: RMSNorm computes one reduction (a sum of squares) instead of two (mean, then variance), and skips the subtraction, which at the scale of every-token-every-layer is a real saving. There are no bias terms and no re-centering, so on disk an RMSNorm layer is a single $d$-vector of gains, which you will see when I read a config. The important property for later: like all normalizations it is applied per-token across features, so it does not mix information between positions (that is attention's job) and it does not change the sequence length or the memory shape.

The gated MLP: SwiGLU

The feed-forward sublayer is where most of a transformer's parameters live, and its job is a per-token nonlinear transformation: take each token's $d_{\text{model}}$-vector, blow it up to a wide intermediate dimension $d_{\text{ff}}$, apply a nonlinearity, and project back down. The classic version was two matrices with a ReLU or GELU between them, $\mathrm{MLP}(x)=W_{\text{down}}\sigma (W_{\text{up}}x)$. Modern models use a gated variant, SwiGLU, which splits the up-projection into two parallel paths and uses one to gate the other.

RoPE: position as rotation

Attention as I derived it in the last chapter is permutation-equivariant: shuffle the input tokens and the outputs shuffle the same way, because nothing in $Q{K}^{\top }$ knows position. A language model obviously needs to know order, so position information must be injected. Rotary Position Embedding (RoPE) does this in the most elegant way anyone has found: instead of adding a position vector to the embedding, it rotates the query and key vectors by an angle proportional to their position, so that the dot product between a query at position $m$ and a key at position $n$ depends only on their relative offset $m-n$.

GQA: sharing keys and values to shrink the KV cache

The last modern change is the one most tied to my 16GB constraint. At inference, an autoregressive model caches the keys and values of all past tokens so it does not recompute them for every new token (the KV cache, derived fully in Part II). The size of that cache is proportional to the number of key/value heads, and for long contexts it becomes the dominant memory cost, often larger than the weights. Grouped-query attention attacks this directly: keep all $h$ query heads (you want the representational capacity of many query views, from the multi-head argument in the last chapter), but let several query heads share one key head and one value head.

The through-line for these four changes: RMSNorm and SwiGLU are efficiency and capacity refinements to the sublayers, RoPE is a cleaner way to inject position that gives relative-position awareness for free, and GQA is a direct memory optimization aimed squarely at the KV cache. Three of the four are about doing more with less memory, which is the whole ethos of running this on one card.

Tooling

The tool is again Hugging Face transformers, this time the model config and the block module. A model's config.json (loaded as a PretrainedConfig) is the exact parameterization of everything above: hidden_size is $d_{\text{model}}$, intermediate_size is $d_{\text{ff}}$ (the $2/3$-adjusted SwiGLU width from equation 4.4), num_attention_heads is $h$, num_key_value_heads is $h_{kv}$ (equation 4.10's GQA lever), head_dim is $d_{\text{head}}$, num_hidden_layers is $N$, rms_norm_eps is the $ϵ$ of equation (4.2), rope_theta is the base of equation (4.8), and hidden_act names the SwiGLU nonlinearity (silu). The block module (Qwen3DecoderLayer) is equation (4.1) in code: an input RMSNorm, the self-attention with RoPE and GQA, a residual add, a post-attention RMSNorm, the SwiGLU MLP, and another residual add. Because the config is the parameterization, I can instantiate my own block from the same config and check that its parameter shapes and count match the real layer exactly, if my SwiGLU width, my head split, and my norm shapes all agree, I have almost certainly assembled the same block.

Lab

The goal is to assemble a transformer block from the four derived parts (RMSNorm, SwiGLU, RoPE, GQA attention) in raw PyTorch, parameterized entirely by a real Qwen3 config, and verify that its parameter shapes and total count match the reference Qwen3DecoderLayer dimension for dimension. The artifact is a JSON report of every submodule's shape alongside the reference, which proves the block I built is the block Qwen3 ships.

uv init labs/transformer-block
cd labs/transformer-block
uv add torch transformers

"""Assemble a transformer block from parts and match a Qwen3 config.

Artifact: artifacts/block_shapes.json comparing my block to the real layer.
"""
import json
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoConfig, AutoModelForCausalLM

MODEL = "Qwen/Qwen3-0.6B"
OUT = Path("artifacts")
OUT.mkdir(exist_ok=True)

class RMSNorm(nn.Module):                                    # eq 4.2
    def __init__(self, d, eps):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d))
        self.eps = eps
    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return x * rms * self.g

class SwiGLU(nn.Module):                                     # eq 4.3
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.gate = nn.Linear(d_model, d_ff, bias=False)
        self.up = nn.Linear(d_model, d_ff, bias=False)
        self.down = nn.Linear(d_ff, d_model, bias=False)
    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

class GQAAttention(nn.Module):                               # eqs 3.x + 4.9/4.10
    def __init__(self, cfg):
        super().__init__()
        self.h = cfg.num_attention_heads
        self.h_kv = cfg.num_key_value_heads
        self.d_head = cfg.head_dim
        d = cfg.hidden_size
        self.q_proj = nn.Linear(d, self.h * self.d_head, bias=False)
        self.k_proj = nn.Linear(d, self.h_kv * self.d_head, bias=False)
        self.v_proj = nn.Linear(d, self.h_kv * self.d_head, bias=False)
        self.o_proj = nn.Linear(self.h * self.d_head, d, bias=False)
        # (RoPE tables omitted for the shape-matching lab; see eq 4.5-4.8)

class Block(nn.Module):                                      # eq 4.1
    def __init__(self, cfg):
        super().__init__()
        self.input_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.attn = GQAAttention(cfg)
        self.post_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.mlp = SwiGLU(cfg.hidden_size, cfg.intermediate_size)

def shape_map(module) -> dict:
    return {n: list(p.shape) for n, p in module.named_parameters()}

def main() -> None:
    cfg = AutoConfig.from_pretrained(MODEL)
    mine = Block(cfg)

    ref_model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.float32)
    ref_layer = ref_model.model.layers[0]

    mine_params = sum(p.numel() for p in mine.parameters())
    ref_params = sum(p.numel() for p in ref_layer.parameters())

    report = {
        "model": MODEL,
        "config": {
            "hidden_size": cfg.hidden_size,
            "intermediate_size": cfg.intermediate_size,
            "num_attention_heads": cfg.num_attention_heads,
            "num_key_value_heads": cfg.num_key_value_heads,
            "head_dim": cfg.head_dim,
            "gqa_group_size": cfg.num_attention_heads // cfg.num_key_value_heads,
            "ffn_ratio": round(cfg.intermediate_size / cfg.hidden_size, 3),
            "rope_theta": cfg.rope_theta,
        },
        "my_block_params": mine_params,
        "ref_layer_params": ref_params,
        "my_shapes": shape_map(mine),
        "ref_shapes": shape_map(ref_layer),
    }
    (OUT / "block_shapes.json").write_text(json.dumps(report, indent=2))

    print(f"GQA group size g = {report['config']['gqa_group_size']} "
          f"(KV cache is 1/{report['config']['gqa_group_size']} of MHA)")
    print(f"FFN ratio d_ff/d_model = {report['config']['ffn_ratio']} "
          f"(SwiGLU, ~3x hidden per eq 4.4)")
    print(f"My block params:  {mine_params:,}")
    print(f"Ref layer params: {ref_params:,}")
    print(f"Artifact: {(OUT / 'block_shapes.json').resolve()}")

if __name__ == "__main__":
    main()

uv run python assemble_block.py

What you should see. The script writes artifacts/block_shapes.json listing every parameter tensor of my hand-built block next to the real Qwen3 layer's. The projection shapes line up: q_proj maps hidden_size to num_attention_heads * head_dim while k_proj and v_proj map to the smaller num_key_value_heads * head_dim, and you can read the GQA asymmetry straight off those two shapes, the K/V projections are literally $g$ times narrower than Q, which is equation (4.10) made visible in the weights. The printed summary reports the GQA group size (2 for the 0.6B, meaning half-size KV cache), the FFN ratio near 3 (the $2/3$-adjusted SwiGLU width of equation 4.4, not a clean power of two), and the two parameter totals within a fraction of a percent of each other, the small gap explained by the q/k norms and RoPE buffers I left out. Keep the JSON: it is the receipt that "a Qwen3 layer" is exactly RMSNorm + GQA-attention-with-RoPE + RMSNorm + SwiGLU wired residually, four derivations you can now point at.