Build your own KV cache
Single-token cached attention·30 min

Prove cached == recompute (one layer)

0 / 1
# grade this stage (from your course directory)
npx theinference test --stage prove-one-layer --watch

Last step you returned one position. Now you prove the whole layer. This is the "the optimization changes nothing" guarantee in miniature: walk a length-n sequence one token at a time from an empty cache, and at every position the cached attention output must equal what a full recompute produces there.

If this holds at all n positions for one layer, it holds for twelve, and it holds across a whole generation — which is exactly the headline you'll cash out in later chapters. The point of the cache is that it's an identity, not an approximation.

What to build

You already have layerAttnIncremental from the previous step — it walks a fresh cache across the sequence and returns every position's attention-sublayer output. This step just exposes all of those rows instead of the last one, and the grader checks every single position against recompute.

No new model code is required: the adapter embeds the ids and runs layerAttnIncremental over layer 0. If the previous step passed, the work here is wiring plus reading the warning below for what failure modes look like.

Your task

Implement runLayerEquiv

Wire runLayerEquiv(seqIds): embed the ids, run layerAttnIncremental over layer 0 with a single fresh cache, and return all per-position attention-sublayer outputs — an n × 768 array, one row per token. Each row is the incremental cached decode output for that position; the grader compares the whole stack to the recompute oracle.

tests/adapters.ts
import { embed } from "./src/gpt2.ts";
import { layerAttnIncremental, loadWeights } from "./src/kvcache.ts";

export async function runLayerEquiv(seqIds: number[]): Promise<number[][]> {
  const w = await loadWeights(WEIGHTS);
  // one fresh cache, decoded incrementally across all n positions
  return layerAttnIncremental(embed(seqIds, w), w, 0); // [n][768]
}
Read the failure shape — it names the bug

The pattern of which positions diverge tells you exactly what's wrong:

  • Only the first position matches, later ones drift more and more — your position counter isn't advancing. You're attending over a stale or short cache, so every step after the first sees the wrong history. (The classic: looping j over oldLen instead of oldLen+1.)
  • Every position is off by a constant — your wpe indexing base is wrong (off-by-one or reset to 0), so every token is embedded at the wrong absolute position.
  • First position fine, second onward wrong but stable — you re-applied a causal mask and are masking out legitimate past positions. Use the shape of the divergence to localize before you change anything.
How it's graded

Every position's row is compared to the recompute oracle (attnSublayer at layer 0) with allclose, atol 1e-4, rtol 1e-3. One divergent position fails the whole step — there's nowhere for a position bug to hide. The tolerance accounts only for float reassociation between the cached and recompute summation orders; anything larger is a genuine error.

Single-step cached attentionWhy only allclose