Transformer Block Anatomy, in depth: the residual stream, attention and MLP sublayers, shapes, parameter and FLOP counts, and the bugs that hide in a block

The block as an editor of the residual stream

The most useful mental model is the residual stream. Each token position carries a vector of width d (4,096 in our example) from the embedding layer to the output layer. Blocks never replace that vector; each sublayer reads a normalised copy, computes something, and adds the result back. In a pre-norm block the whole computation is two lines:

h   = x + Attention(Norm1(x))      # mix information across positions
out = h + MLP(Norm2(h))            # transform each position independently

This has three consequences. First, the gradient has a direct path from the loss to every layer through the additions, which is why deep pre-norm stacks train stably without warm-up tricks that post-norm needed. Second, every sublayer writes into a shared space, so later layers can read features written by much earlier ones. Third, the stream's magnitude grows with depth, which is one reason the final layer applies one more norm before the output projection.

The two sublayers have different jobs. Attention is the only place where positions exchange information: token 57 can read from tokens 0 to 57. The MLP acts on each position separately and holds most of the parameters; it is where much of the model's stored knowledge is thought to live.

Every step with shapes

Use B for batch size, T for sequence length, d for model width, h for query heads, h_kv for key and value heads, d_h = d / h for head width and d_ff for the MLP's hidden width. The example configuration is d = 4,096, h = 32, h_kv = 8, d_h = 128, d_ff = 14,336.

1. Norm. RMSNorm scales each position's vector by the reciprocal of its root mean square and multiplies by a learned gain of size d. Shape stays (B, T, d).
2. Projections. Q = xW_q gives (B, T, h, d_h) = (B, T, 32, 128). K and V use smaller matrices and give (B, T, 8, 128) each. That is grouped-query attention: each group of 4 query heads shares one key and value head, which cuts the KV cache by 4 times.
3. RoPE. Rotate pairs of dimensions in Q and K by angles that depend on position, so the dot product between a query and a key depends on their relative distance. V is not rotated.
4. Scores. Transpose to (B, heads, T, d_h), compute QK^T / sqrt(d_h) to get (B, 32, T, T), and add a causal mask of minus infinity above the diagonal so position i cannot see positions after i.
5. Softmax and mix. Softmax over the last axis, in float32, then multiply by V to get (B, 32, T, 128). Fused kernels such as FlashAttention compute this without ever storing the T x T matrix.
6. Output projection. Merge heads back to (B, T, 4,096) and multiply by W_o. Add to the residual stream.
7. MLP. Norm again, then compute two projections to width d_ff (gate and up), apply SiLU to the gate, multiply elementwise, and project back down to d. Add to the residual stream.

A complete block in PyTorch

The implementation below is self-contained and matches the description above. It favours clarity over speed but uses the fused attention kernel PyTorch provides.

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

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-5):
        super().__init__(); self.eps = eps; self.g = nn.Parameter(torch.ones(d))
    def forward(self, x):
        x32 = x.float()
        x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x32 * self.g.float()).type_as(x)

def rope(x, cos, sin):                       # x: (B, H, T, d_h); cos, sin: (T, d_h // 2)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    out = torch.stack((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1)
    return out.flatten(-2).type_as(x)

class Block(nn.Module):
    def __init__(self, d=4096, h=32, h_kv=8, d_ff=14336):
        super().__init__()
        self.h, self.h_kv, self.d_h = h, h_kv, d // h
        self.n1, self.n2 = RMSNorm(d), RMSNorm(d)
        self.wq = nn.Linear(d, h * self.d_h, bias=False)
        self.wk = nn.Linear(d, h_kv * self.d_h, bias=False)
        self.wv = nn.Linear(d, h_kv * self.d_h, bias=False)
        self.wo = nn.Linear(h * self.d_h, d, bias=False)
        self.w_gate = nn.Linear(d, d_ff, bias=False)
        self.w_up = nn.Linear(d, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d, bias=False)

    def attn(self, x, cos, sin):
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.h, self.d_h).transpose(1, 2)
        k = self.wk(x).view(B, T, self.h_kv, self.d_h).transpose(1, 2)
        v = self.wv(x).view(B, T, self.h_kv, self.d_h).transpose(1, 2)
        q, k = rope(q, cos, sin), rope(k, cos, sin)
        rep = self.h // self.h_kv               # share each K, V head across a query group
        k = k.repeat_interleave(rep, dim=1)
        v = v.repeat_interleave(rep, dim=1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.wo(y.transpose(1, 2).reshape(B, T, -1))

    def forward(self, x, cos, sin):
        x = x + self.attn(self.n1(x), cos, sin)
        return x + self.w_down(F.silu(self.w_gate(self.n2(x))) * self.w_up(self.n2(x)))

Two details are deliberate. The norm is computed in float32 and cast back, because the mean of squares overflows or loses precision in float16. And the reshape order is view to (B, T, heads, d_h) and then transpose; viewing directly to (B, heads, T, d_h) silently mixes positions and heads and still produces tensors of the right shape. The cos and sin tables are precomputed once per sequence length from the RoPE base frequency.

Worked example: counting parameters

With no biases, the parameter count is a sum of matrix sizes. For the example configuration:

Thirty-two blocks give about 6.98 billion parameters. A vocabulary of 128,256 tokens with untied input and output matrices adds 2 x 128,256 x 4,096 = about 1.05 billion, plus a final norm, for about 8.03 billion: the published size of Llama 3 8B, whose configuration this is. Two things stand out. The MLP holds about 81% of each block's parameters, so it dominates weight memory. And grouped-query attention makes K and V a quarter the size of Q and O.

The SwiGLU width is not arbitrary. A classic MLP with a 4d hidden layer has 8d^2 parameters; SwiGLU has three matrices, so a hidden width of about (8/3)d keeps the count similar. Here 14,336 / 4,096 = 3.5, a little wider than 8/3, which the Llama 3 designers chose and rounded to a hardware-friendly multiple.

Compute and memory per token

Each weight participates in one multiply and one add per token in the forward pass, so the matrix multiplies cost about 2 x N FLOPs per token, where N is the parameter count excluding the input embedding lookup. For this model that is roughly 2 x 7.5 billion = 15 GFLOPs per generated token. Training costs about three times the forward pass, which is the source of the familiar 6N FLOPs per token rule for estimating training compute.

Attention adds work that grows with context. Computing QK^T and multiplying by V costs about 4 x T x d FLOPs per layer for a token that attends to T earlier positions. At T = 8,192 that is about 134 MFLOPs per layer, or 4.3 GFLOPs across 32 layers: almost 30% on top of the matrix multiplies for the last token in an 8k window, which is why long context is not free even with fused kernels.

Memory during generation is dominated by the KV cache. Each token stores K and V for every layer: 2 x 32 layers x 8 heads x 128 dims x 2 bytes in bf16 = 131,072 bytes, 128 KiB per token. An 8,192-token sequence needs 1 GiB of cache, and a batch of 32 such sequences needs 32 GiB, more than half the size of the bf16 weights. Without grouped-query attention, with 32 KV heads, it would be four times larger. The KV cache page covers paging and quantisation.

Design choices and why they won

• Pre-norm over post-norm. Post-norm, the original 2017 layout, normalises after the residual add, so gradients pass through every norm on the way down and deep stacks need careful warm-up. Pre-norm keeps the residual path clean and trains reliably at depth.
• RMSNorm over LayerNorm. Dropping mean-centring and the bias saves a reduction and parameters with no measured quality loss in most reports.
• SwiGLU over GELU MLPs. Gated MLPs gave consistently better loss per parameter in published comparisons, at the cost of a third matrix.
• RoPE over learned absolute positions. Relative position falls out of the dot product, and the base frequency can be scaled to extend context.
• Grouped-query attention. Most of multi-head attention's quality at a fraction of its KV cache, which is what limits serving batch size.
• No biases. Removing linear-layer biases costs nothing measurable and simplifies sharding and quantisation.
• Optional QK-norm. Some recent models normalise Q and K per head before the dot product to stop attention logits from growing during training, a common cause of loss spikes.

Bugs that hide inside a block, and the tests that catch them

Block bugs are dangerous because a wrong block still trains: the loss goes down, just not as far. The common ones are a causal mask that is off by one or missing during training, RoPE applied to V or with the wrong pairing convention when loading weights from another codebase, softmax or norm statistics computed in half precision, the view-transpose mix-up above, repeating K and V heads in the wrong order, and dropout left on at inference. Four cheap tests catch nearly all of them:

def test_causality(block, cos, sin, T=16, d=4096):
    x = torch.randn(1, T, d)
    y1 = block(x, cos, sin)
    x2 = x.clone(); x2[:, -1] += 10.0            # change only the last token
    y2 = block(x2, cos, sin)
    assert torch.allclose(y1[:, :-1], y2[:, :-1], atol=1e-5)   # earlier positions unchanged

def test_incremental_matches_full(model, tokens):
    full = model(tokens)                          # one pass over the whole sequence
    step = decode_with_kv_cache(model, tokens)    # one token at a time
    assert torch.allclose(full, step, atol=1e-3)

The other two: compare logits against a trusted reference implementation on the same weights to within a small tolerance, and check that the parameter count matches the arithmetic in the table above. The incremental-decode test is the one that catches RoPE position offsets and KV-cache indexing errors, which a full-sequence test never exercises.

What to do next

1. Write the block above from memory, then run the causality test and a parameter-count assertion against your own arithmetic.
2. Load a small open model's weights into your block and match its logits to the reference implementation within tolerance.
3. Add a KV cache and the incremental-decode test; debug position offsets until full and incremental outputs agree.
4. Compute FLOPs per token and KV-cache bytes per token for the model you serve, and use them to size batch and context limits.
5. Read the component pages for pre-norm, RMSNorm, RoPE and SwiGLU, and try one swap at a time, such as LayerNorm for RMSNorm, on a small training run while watching the loss curve.