Implement prefill
# grade this stage (from your course directory)
npx theinference test --stage build-prefill --watchTime to build the pass that fills the cache. prefill takes the prompt ids, runs one batched forward over positions 0..prompt_len-1 (embedding each token with wpe[position]), and writes every layer's K and V for every prompt position. It returns the last token's next-token logits — the prediction for the first token you'll generate. When this turns green, you have a populated cache that the decode loop can pick up from.
What to build
Implement prefill
Implement prefill(promptIds, w, mc) in src/kvcache.ts: embed each prompt token with embedToken(id, position, w) (so position i adds wpe[i]), run the full 12-layer stack under the full lower-triangular mask writing K/V into mc for positions 0..prompt_len-1, then return the last position's [50257] logits. You'll lean on embedToken and the same per-layer block you used for decodeStep. Then wire it into the runPrefill adapter, which constructs a fresh ModelCache, calls prefill, and returns the logits.
export function runPrefill(promptIds: number[]): number[] | Promise<number[]> {
// construct a fresh ModelCache, call prefill(promptIds, w, mc),
// and return the LAST token's next-token logits, length 50257
}
Prefill uses the full triangular mask (token i attends to positions ≤ i), not the trivial single-query decode mask. When it finishes, every layer's cache must hold exactly prompt_len positions — across all 12 layers that's [12, n_heads, prompt_len, 64]. A wrong sequence length breaks the very first decode step: the first generated token must get position prompt_len (= cache length), so an off-by-one here corrupts wpe on the handoff. Embed with wpe[0..prompt_len-1] and write one cache row per prompt position — no more, no less.
Your returned [50257] logits are compared against the recompute oracle (recomputeLogits over the same prompt) with allclose, atol 1e-3, rtol 1e-2. That's the honest envelope for float32 summation-order drift across 12 blocks — not slack. The cache shapes are checked integer-exact: every layer must be at length prompt_len.