Putting It Together: An Annotated End-to-End Forward Pass

A small decoder-only Transformer combining the whole course:

@dataclass
class Config:
    vocab_size: int = 1000
    d_model: int = 256
    n_layers: int = 4
    n_heads: int = 8            # query heads
    n_kv_heads: int = 2         # -> GQA group size 4
    d_head: int = 32
    d_ff: int = 512             # dense-layer FFN width
    n_experts: int = 8
    n_experts_active: int = 2   # top-k
    n_shared_experts: int = 1
    d_expert: int = 128         # fine-grained: narrower than d_ff
    first_k_dense: int = 1      # layer 0 is dense
    rope_theta: float = 10000.0
    max_seq_len: int = 512
    norm_eps: float = 1e-6

This is DeepSeek-shaped in miniature: MoE with a shared expert, dense first layer, fine-grained experts. Swap n_kv_heads = 8 and you have MHA; set n_experts = 0 and you have Llama-3-shaped.

And one block:

Note the two residual additions taking their input from before each norm. That is pre-norm (module 06): the residual stream is never normalized in place.

Illustrative instrumented shape trace, B = 1, T = 6, through block 1 (MoE). The default script uses (B,T)=(2,16) and prints a shorter smoke-test trace:

token ids                        (1, 6)
after embedding                  (1, 6, 256)
after RMSNorm (pre-norm)         (1, 6, 256)

  --- attention sublayer ---
Q  (H=8 query heads)             (1, 8, 6, 32)
K  (H_kv=2 kv heads)             (1, 2, 6, 32)    <- CACHED (small)
V                                (1, 2, 6, 32)    <- CACHED (small)
after QK-Norm + RoPE           Q (1, 8, 6, 32)   K (1, 2, 6, 32)
K,V expanded to H heads          (1, 8, 6, 32)    <- NOT cached
attention scores                 (1, 8, 6, 6)
attention output                 (1, 8, 6, 32)
heads concatenated               (1, 6, 256)
after W_o                        (1, 6, 256)
after residual add               (1, 6, 256)

  --- FFN sublayer ---
after RMSNorm 2                  (1, 6, 256)
router logits                    (6, 8)
top-2 expert indices             (6, 2)
   token 0 -> experts [0, 7]     token 1 -> experts [0, 7]
MoE output                       (1, 6, 256)
after residual add               (1, 6, 256)

final logits                     (1, 6, 1000)

Five things this trace shows that prose cannot:

1. The residual stream is (B, T, 256) at every checkpoint. Everything in between is a temporary excursion. This is module 06's shape invariant, made visible.
2. K and V are (1, 2, 6, 32) while Q is (1, 8, 6, 32). That 4× gap is exactly GQA's saving, and it is the pre-expansion tensors that get cached.
3. The expansion to 8 heads is not cached. repeat_interleave runs per step on the small tensors — this is why GQA saves cache and not just parameters.
4. The score matrix is (1, 8, 6, 6) — H × T × T. Grow T to 32,768 and this term is what explodes (modules 10–11).
5. Router logits are (6, 8), not (1, 6, 8). Routing is flattened across batch and time because it is per token, not per sequence (module 12).

Router collapse, visible at initialization

Over 12 routing slots (6 tokens × top-2), expert usage at random init was:

expert:  0   1   2   3   4   5   6   7
count:   6   1   0   2   0   0   0   3

Expert 0 takes half of all slots; four experts are never chosen. This is exactly the collapse dynamic module 12 described — visible in a randomly-initialised toy model, before any training has amplified it. This is what the load-balancing loss is fighting.

The core of GroupedQueryAttention.forward:

def forward(self, x, cos, sin, cache=None):
    B, T, _ = x.shape

    # 1. project. Note W_k/W_v are SMALLER: H_kv heads, not H.
    q = self.W_q(x).view(B, T, self.H,    self.d_head).transpose(1, 2)
    k = self.W_k(x).view(B, T, self.H_kv, self.d_head).transpose(1, 2)
    v = self.W_v(x).view(B, T, self.H_kv, self.d_head).transpose(1, 2)

    # 2. QK-Norm (module 06) — bounds score magnitude. BEFORE RoPE.
    q = self.q_norm(q)
    k = self.k_norm(k)

    # 3. RoPE (module 05) — Q and K only, never V.
    q = apply_rope(q, cos, sin)
    k = apply_rope(k, cos, sin)

    # 4. cache the SMALL tensors (module 11)
    if cache is not None:
        k, v = cache.update(k, v)

    # 5. explicit temporary copies to match query heads; not cached
    k = k.repeat_interleave(self.group_size, dim=1)
    v = v.repeat_interleave(self.group_size, dim=1)

    # 6. scaled dot-product with causal mask (modules 03, 08)
    #    dispatches to FlashAttention where available (module 11)
    past = k.shape[2] - T
    allowed = (torch.arange(k.shape[2], device=x.device)[None, :]
               <= past + torch.arange(T, device=x.device)[:, None])
    out = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed)

    # 7. concat heads and project (module 04)
    out = out.transpose(1, 2).reshape(B, T, self.H * self.d_head)
    return self.W_o(out)

Order matters and is easy to get wrong:

• QK-Norm before RoPE. Raschka's Qwen3 implementation does this, and swapping them changes the model.
• RoPE on Q and K only. Position should shape who attends to whom, not the content retrieved (module 05).
• Cache before expansion. Caching the expanded tensors would discard the entire GQA benefit — a genuine and common bug.
• Query i sees keys through past+i. This offset matters for cached chunks containing several new tokens; a square-prefill shortcut is insufficient. A valid single-token append can see every key currently in its cache.

def forward(self, x):
    B, T, D = x.shape
    flat = x.reshape(-1, D)                       # (B*T, D) — per-token routing

    logits = self.gate(flat)                      # (N, n_experts)
    topk_w, topk_idx = F.softmax(logits, dim=-1).topk(self.top_k, dim=-1)

    out = torch.zeros_like(flat)
    for e, expert in enumerate(self.experts):
        tok, slot = (topk_idx == e).nonzero(as_tuple=True)
        if tok.numel() == 0:
            continue                              # expert unused this batch
        out.index_add_(0, tok, topk_w[tok, slot, None] * expert(flat[tok]))

    for expert in self.shared:                    # always active (module 12)
        out = out + expert(flat)

    return out.view(B, T, D), logits

The gather/scatter structure is the point: each expert only ever sees the tokens routed to it. That is where the compute saving comes from. This loop is pedagogical — production uses grouped GEMMs or fused kernels — but the semantics are identical.

Measured from the running model:

total parameters      : 3,967,488
routed-expert params  : 2,359,296     (59.5% of the model)
shared-expert params  :   294,912
active per token      : 2,198,016     (55.4%)

Two observations:

Routed experts are 59.5% of parameters — the module-07 fact (FFN dominates) amplified by MoE. Exactly what MoE is designed to exploit.

Activation ratio is only 55.4%, versus DeepSeek V3's 5.5%. The toy model is far denser because it has 8 experts activating 2 (25%) while DeepSeek has 256 activating 9 (3.5%). Sparsity requires many experts — with 8 experts you simply cannot be very sparse. This is a concrete argument for fine-grained experts (module 12) that the arithmetic makes obvious.

logits, router_logits = model(ids)

# next-token objective (module 13) — the one-position shift
optimizer.zero_grad(set_to_none=True)
ce = F.cross_entropy(
    logits[:, :-1].reshape(-1, V),      # predictions at 0..T-2
    ids[:, 1:].reshape(-1),             # actual next tokens from the input document
)

# load-balancing, one term per MoE layer (module 12)
aux = sum(load_balancing_loss(rl, rl.topk(k, -1).indices, n_experts)
          for rl in router_logits)

(ce + aux).backward()
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()

Measured:

cross-entropy 7.0136   (expected approx ln(V) = 6.9078 at init)
aux loss      0.0342
grad-norm     7.987

The cross-entropy check is the single most valuable sanity test you can run. At initialization a language model should be uniformly uncertain over V tokens, giving loss near ln(V) under small-logit initialization. Here the CPU PyTorch 2.8 run gives 7.014 versus ln(1000)=6.908; these are diagnostics, not exact cross-platform output requirements.

If your initial loss is 200 rather than 7, your initialization is broken — which is exactly what happened in the first version of this file before proper weight init was added. Near-zero loss on diverse untrained targets warrants checking the causal mask and target shift; it is not proof of leakage on trivial data.

@torch.no_grad()
def generate(self, prompt_ids, max_new_tokens=20, temperature=0.8):
    import math
    if not isinstance(max_new_tokens, int) or max_new_tokens < 0:
        raise ValueError("max_new_tokens must be a nonnegative integer")
    if not math.isfinite(temperature) or temperature < 0:
        raise ValueError("invalid temperature")
    if prompt_ids.ndim != 2 or min(prompt_ids.shape) <= 0:
        raise ValueError("nonempty unpadded (batch, tokens) prompts required")
    if prompt_ids.shape[1] + max_new_tokens > self.cfg.max_seq_len:
        raise ValueError("requested output exceeds context window")
    if max_new_tokens == 0:
        return prompt_ids.clone()
    self.eval()
    caches = [KVCache() for _ in self.blocks]

    # --- PREFILL: whole prompt, one pass, causal ---
    logits, _ = self.forward(prompt_ids, caches, pos_offset=0)
    pos = prompt_ids.shape[1]
    next_id = self._sample(logits[:, -1], temperature)   # ONLY the last position

    out = [prompt_ids, next_id]

    # --- DECODE: one token per pass, reading the cache ---
    for _ in range(max_new_tokens - 1):
        logits, _ = self.forward(next_id, caches, pos_offset=pos)
        pos += 1
        next_id = self._sample(logits[:, -1], temperature)
        out.append(next_id)

    return torch.cat(out, dim=1)

Three details worth pinning down:

logits[:, -1] — video 84's point. During prefill the model produces logits for all 5 positions, but positions 0–3 predict tokens you already have. Only the last is a genuine prediction.

pos_offset — RoPE angles depend on absolute position. At decode step 3 the single input token is at position prompt_len + 2, not position 0. Forgetting this offset is a common and very confusing bug: generation looks plausible for a few tokens and then degrades.

Offset-aware masking permits keys through past+i; at T=1, an append-only cache contains no future keys. Other one-query layouts can still need masking.

Verifying the cache is correct

The most important test in the file, because a broken cache produces output that looks fine:

=== cached vs uncached equivalence ===
  max abs difference: 6.85e-07

Run 8 tokens through in one pass and record the final logits. Then prefill 7, decode the 8th through the cache, and compare. Agreement to 6.85e-07 is floating-point noise — the paths are equivalent.

Always write this test. A cache bug — wrong offset, caching post-expansion, stale positions — degrades quality subtly rather than crashing.

The structure is the same. Every difference is a hyperparameter, except MLA replacing GQA — and MLA slots into the same position in the block.

Production additions this file omits:

• FlashAttention — F.scaled_dot_product_attention already dispatches to it on suitable hardware
• PagedAttention — cache in fixed blocks rather than one growing tensor
• Grouped GEMM for experts instead of the Python loop
• Expert / tensor / pipeline parallelism
• Quantization
• Speculative decoding

None change the mathematics. All of them are module 11 and 14 material.

This decoder-only exercise complements two earlier runnable labs: frequency-logged BPE training and padded encoder-decoder translation. The translation lab adds variable-length batches, all four attention-mask contracts, disk artifact reloads, and separate teacher-forced/autoregressive evaluation. Its held-out toy combinations are not a real translation benchmark.

Run train_tiny_decoder.py from the repository root:

python content/courses/transformers/code/train_tiny_decoder.py
python site/test_transformer_corrections.py

The first script trains this actual Transformer on two short synthetic documents. The same token 5 must predict token 3 or 4 depending on its earlier context. Inputs and labels come from one shifted sequence, unlike an independent-random- target gradient smoke test. It checks falling training loss, context-dependent continuation, an in-memory checkpoint, and equality after the next restored optimizer update. It makes no held-out quality claim. The ambiguous first continuation after the shared BOS prevents a deterministic zero-loss solution on every position; not every remaining error is a bug.

The regression suite checks every position under full prefill, token-by-token decode and multiple cached chunk partitions, for batch sizes one/two and both dense and MoE modes. It also checks top-1 router gradients, balanced slot loss, zero-generation budgets, context overflow, empty prompts, cache offsets and EOS. The teaching implementation supports unpadded equal-length batches; it does not silently interpret any vocabulary ID as padding. A serving-grade padded batch requires per-request positions, valid-key masks and cache lengths.

Failure-injection exercise and answers: resetting RoPE offsets breaks cached equivalence; upper-left rectangular causal masking hides valid prefix keys; selected-only top-1 softmax removes task router gradients; unshifted labels turn next-token learning into input reconstruction. A one-step random loss check detects none of these reliably, which is why the contracts are tested separately.

The miniature model is pedagogical, not an exact DeepSeek replica: it uses GQA rather than MLA, full-softmax routing rather than DeepSeek's routing recipe, and omits training infrastructure, MTP and production dispatch kernels. The explicit expansion and dense offset mask favor clarity over fused-kernel speed.