Full-model cached decode & per-step identity·30 min

Full-model cached decode step

0 / 1

# grade this stage (from your course directory)
npx theinference test --stage full-decode-step --watch

In Chapter 2 you proved one cached attention sublayer matches recompute. In Chapter 3 you prefilled a prompt into every layer's cache. Now you connect them end to end: take one new token, thread it through all 12 layers — each reading and growing its own K/V cache — apply the final LayerNorm, project against the tied embedding, and read off the next-token logits. This is a full forward pass, but the only piece that changed from your llm-inference model is attention. Everything else runs untouched.

One cached decode step

token id, position  →  embedToken (wte + wpe[position])
                    →  blockStep × 12   (each: ln_1 → attnStep(reads+grows cache) → ln_2 → mlp)
                    →  ln_f
                    →  · wteᵀ           (tied unembed)
                    →  logits [50257]

Only attention changed — reuse the rest

Reuse your llm-inference MLP, LayerNorm, and unembed exactly as they are — pull lnRows, mlp, and the wte projection from gpt2.ts. The single new idea is that the cache threads through layer by layer: blockStep calls attnStep, which reads the layer's whole K/V history, appends the new token's K/V, and attends. Layer l reads and grows mc.layers[l] only. Don't share one cache across layers.

What to build

Your task

Implement runDecodeStep

Implement decodeStep(id, position, w, mc) in src/kvcache.ts: embed the token at the absolute position, run it through all 12 blockSteps (each growing mc.layers[l]), then ln_f and the tied unembed → length-50257 logits.

Then wire the adapter. It receives a promptIds array and one nextId. Prefill promptIds into a fresh ModelCache, then decode the single token nextId at position promptIds.length (the next absolute index), and return that token's next-token logits.

tests/adapters.ts

import { prefill, decodeStep, ModelCache, loadWeights } from "../src/kvcache.ts";

export async function runDecodeStep(promptIds: number[], nextId: number): Promise<number[]> {
  const w = await loadWeights(/* weights path */);
  const mc = new ModelCache();
  prefill(promptIds, w, mc);                 // populate every layer's cache
  return decodeStep(nextId, promptIds.length, w, mc); // decode at position prompt_len
}

The position is an absolute index, not a step counter

GPT-2 adds wpe[position], so position must be the token's absolute index in the sequence. After prefilling promptIds, the new token sits at promptIds.length — not at 0, and not at 1. If this step passes alone but later steps drift, your decode position counter isn't tracking the absolute index prompt_len + step. If this first decode step is off by itself, check the prefill→decode handoff: the cache must already hold promptIds.length rows per layer before you decode, and the new position continues from there.

How it's graded

Your length-50257 logits are compared against a full recompute of [...promptIds, nextId] at the same position with allclose, atol 1e-3, rtol 1e-2. The cached path and the recompute path do the same math in a different order, so float32 summation-order drift across 12 layers is the only expected difference — 1e-3 is the honest envelope of that drift, not slack. If you're far outside it, it's a real bug (wrong position, a shared cache, or a missing ln_f), not rounding.

What does prefill leave behind?Per-step identity over a generation