40. Transformer Decoder Block, End to End

Problem

Every autoregressive transformer, regardless of scale, is built from the same two-sublayer block repeated N times: causal self-attention and a position-wise MLP, each wrapped in a residual connection with pre-norm (GPT-style: normalize, run the sublayer, add the residual):

x = x + SelfAttention(LayerNorm1(x))
x = x + MLP(LayerNorm2(x))

Unlike wiring together a pre-built nn.MultiheadAttention, this problem builds the attention sublayer itself from its component projections: a fused QKV linear layer, a manual head-split/attend/merge, and an output projection. Getting the plumbing right end to end — where the causal mask gets applied, that the residual branches off before normalization, that reshapes preserve the right axes — is what makes assembling any fancier architecture on top of this tractable.

Source: src/40_transformer_decoder_block.py

class TransformerDecoderBlock(nn.Module):
    def __init__(self, d_model: int, num_heads: int, d_ff: int): ...
    def forward(self, x: torch.Tensor) -> torch.Tensor: ...  # (B, T, d_model) -> (B, T, d_model)

>>> torch.manual_seed(0)
>>> block = TransformerDecoderBlock(d_model=16, num_heads=2, d_ff=32)
>>> x = torch.randn(2, 5, 16)
>>> block(x).shape
torch.Size([2, 5, 16])

Step-by-Step Approach

1. Normalize the input with ln1, then project it through qkv_proj in one matmul and split the result into q, k, v with .chunk(3, dim=-1).
2. Reshape each of q, k, v into heads: (B, T, d_model) -> (B, H, T, d_head).
3. Compute causal attention scores, add a (T, T) upper-triangular -inf mask, softmax, then multiply by v.
4. Merge heads back to (B, T, d_model) and apply attn_out_proj.
5. Add this attention output to the original, unnormalized x as a residual connection — not to ln1(x).
6. Repeat the pattern for the MLP sublayer: x = x + self.mlp(self.ln2(x)).
7. Return the final x, which has passed through both residual sublayers.

The key insight is pre-norm placement: normalization happens on the branch going into each sublayer, while the residual addition always uses the value before that normalization. Get this backwards (normalize after adding, or normalize the residual branch itself) and the block still runs and produces a right-shaped output, but it's not the architecture being asked for and gradient flow through the residual stream degrades.

Reference solution

import math
import torch
import torch.nn as nn


class TransformerDecoderBlock(nn.Module):
    def __init__(self, d_model: int, num_heads: int, d_ff: int):
        super().__init__()
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.d_head = d_model // num_heads

        self.ln1 = nn.LayerNorm(d_model)
        self.qkv_proj = nn.Linear(d_model, 3 * d_model)
        self.attn_out_proj = nn.Linear(d_model, d_model)

        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
        )

    def _self_attention(self, x: torch.Tensor) -> torch.Tensor:
        B, T, d_model = x.shape
        qkv = self.qkv_proj(x)  # one matmul instead of three
        q, k, v = qkv.chunk(3, dim=-1)

        def split_heads(t):
            return t.view(B, T, self.num_heads, self.d_head).transpose(1, 2)

        q, k, v = split_heads(q), split_heads(k), split_heads(v)

        scores = q @ k.transpose(-1, -2) / math.sqrt(self.d_head)
        causal_mask = torch.triu(torch.full((T, T), float("-inf"), device=x.device), diagonal=1)
        scores = scores + causal_mask
        weights = torch.softmax(scores, dim=-1)
        out = weights @ v

        out = out.transpose(1, 2).contiguous().view(B, T, d_model)
        return self.attn_out_proj(out)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # pre-norm residual: normalize -> sublayer -> add; the residual
        # branches off the *unnormalized* x, not ln(x)
        x = x + self._self_attention(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x

Key Functions & Tricks

• nn.Linear(d_model, 3 * d_model) then .chunk(3, dim=-1) — one fused matmul for Q/K/V instead of three separate ones, a standard efficiency trick in real transformer implementations.
• torch.triu(..., diagonal=1) — builds the causal mask; diagonal=1 keeps self-attention to the current position allowed.
• nn.LayerNorm(d_model) applied before each sublayer, not after — the pre-norm convention essentially every modern autoregressive transformer uses over the original post-norm design.
• Residual add (x = x + sublayer(ln(x))) — the pattern that keeps gradients flowing cleanly through arbitrarily many stacked blocks.
• Causality unit test via input perturbation — changing only the last token and checking earlier outputs are unchanged is a cheap, direct way to verify a causal mask is actually wired in correctly, independent of any numeric oracle.

How to Recognize This Pattern