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

Per-step identity over a generation

0 / 1

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

One correct decode step isn't enough. The guarantee that matters is that across a whole generation, the cached path reproduces the slow recompute path's logits at every step. A position bug can hide in step 5 even when step 0 is perfect. So here you drive a real run — prefill the prompt, then walk a known continuation token by token — and check the prediction at each step against a from-scratch recompute. One divergent step fails. There's nowhere for the bug to hide.

Record the prediction, then consume the token

prefill(promptIds)                         →  logits  (prediction for cont[0])
for k in 0 .. contIds.length - 1:
    record logits                          →  row k  (must match recompute(prompt + cont[:k]))
    logits = decodeStep(contIds[k], promptIds.length + k, ...)

Read the loop carefully: at step k you record the current logits first — that's the model's prediction before it sees contIds[k] — then you feed contIds[k] in to advance the cache. So row k is the next-token distribution conditioned on prompt + cont[:k], exactly what a full recompute of that prefix at its last position produces. The token you decode at step k sits at absolute position promptIds.length + k.

What to build

Your task

Implement runPerStepLogits

The adapter receives promptIds and a continuation contIds. Prefill promptIds into a fresh ModelCache. Then for each continuation token: record the current logits (the prediction), then decodeStep that token at position promptIds.length + k to advance the cache. Return one logit row per continuation token — contIds.length rows, each length 50257. Row k must equal the recompute of [...promptIds, ...contIds.slice(0, k)] at its last position.

tests/adapters.ts

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

export async function runPerStepLogits(promptIds: number[], contIds: number[]): Promise<number[][]> {
  const w = await loadWeights(/* weights path */);
  const mc = new ModelCache();
  let logits = prefill(promptIds, w, mc);     // prediction for contIds[0]
  const rows: number[][] = [];
  for (let k = 0; k < contIds.length; k++) {
    rows.push(logits);                        // record BEFORE consuming the token
    logits = decodeStep(contIds[k], promptIds.length + k, w, mc);
  }
  return rows;
}

Drift means your position counter reset

If step 0 passes but the rows drift further off as k grows, your decode position counter isn't tracking the absolute index — it must be promptIds.length + k, not k and not 0. A counter that resets each step keeps the cache growing correctly while the wpe[position] embedding goes wrong, so early rows look fine and later ones rot. If even the first decode step (row 1) is off, the bug is in the prefill→decode handoff: the cache must hold promptIds.length rows per layer before the first decodeStep.

How it's graded

Every row is compared against the recompute oracle at that prefix with allclose, atol 1e-3, rtol 1e-2. Every step must pass — one divergent row fails the whole step. This is the headline of the course: cached GPT-2 ≡ recompute GPT-2, logits allclose at every position of a real generation. The float32 envelope is the same 1e-3 as before; anything well outside it is a position or caching bug, not rounding.

Full-model cached decode step Greedy decode end-to-end