Single-token cached attention·35 min

Single-step cached attention

0 / 1

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

This is the centerpiece. You'll implement cached attention for one new token at one layer, then drive it across a sequence one token at a time. Everything later — prefill, full decode, the generate loop, prefix caching — is this step composed and repeated. Get it exactly right and the rest is plumbing.

You reuse the provided non-cached GPT-2 (./gpt2.ts) as both your building blocks (qkv, lnRows) and your oracle. The cached attention output for the last token must match what the recompute path produces for that same position, within tolerance.

What to build

Implement two functions in src/kvcache.ts.

attnStep(lnHidden, w, layer, cache) is the new-token step. lnHidden is ln_1(hidden) for the one new token (a length-768 vector). Project it through qkv to get its Q/K/V (each 768, all 12 heads concatenated). Append the new K and V to the cache — its length is now old + 1. Then, per head (head_dim 64), dot the new Q against every cached K (all cache.length rows, including the one you just appended), scale by 1/sqrt(64), softmax, and take the weighted sum of the cached V. Concatenate the 12 heads back to 768 and apply the layer's attn.c_proj. Return that 768-vector — the attention contribution a, without the residual.

layerAttnIncremental(X, w, layer) walks a fresh cache over a whole sequence. X is the layer's input rows. For each position t: layer-norm the row, call attnStep, and push X[t] + a (the residual add) onto the output. Return all per-position attention-sublayer outputs.

Your task

Implement runStepAttention

Wire the adapter runStepAttention(seqIds): embed the ids (embed from ./src/gpt2.ts adds wte[id] + wpe[pos]), run layerAttnIncremental over layer 0, and return the last position's attention-sublayer output (length 768). The position lives in the embedding here — embed indexes wpe[pos] for pos = 0..n-1 — so your attnStep works on already-embedded rows and does not touch wpe itself.

tests/adapters.ts

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

export async function runStepAttention(seqIds: number[]): Promise<number[]> {
  const w = await loadWeights(WEIGHTS);
  const out = layerAttnIncremental(embed(seqIds, w), w, 0);
  return out[out.length - 1]!; // last position's attn-sublayer output [768]
}

The three classic bugs

Wrong absolute position. Here positions come from embed indexing wpe[pos] over 0..n-1; in a hand-rolled decode loop the new token's position is cache.length before append. Use the length after append (or reset to 0) and wpe shifts under every token.
Attending over oldLen instead of oldLen+1. Append K/V first, then loop j over the full cache.length. Loop the old length and the new token never attends to itself — every score row is missing its last entry.
Re-applying a triangular mask. The full forward pass masks the upper triangle, but in this single-token step the new token is last and legitimately attends to all cached positions. Drag a causal mask in here and you'll zero out valid past positions.

How it's graded

Your returned vector is compared against the recompute oracle — attnSublayer from ./gpt2.ts run at layer 0 — at the last position, with allclose, atol 1e-4, rtol 1e-3. Per-layer attention tolerance: the cached and recompute paths do the same math but accumulate the softmax-weighted sum in a different order, so they agree within that envelope, not bit-for-bit. A difference beyond tolerance is a real bug, not float drift.

References

One token Prove cached == recompute (one layer)