Prefix caching·30 min

Save and reuse a prefix cache

0 / 1

# grade this stage (from your course directory)
npx theinference test --stage save-and-reuse --watch

This is the production flex. You already have prefill (populate every layer's cache for a prompt) and decodeStep (one token through all 12 layers, growing the cache). Prefix caching is those two pieces composed with discipline: prefill a shared prefix once and hand back the cache, then for each suffix reuse that cache instead of re-prefilling. Done right, the cached path reproduces a from-scratch recompute of the full prefix + suffix prompt — same logits, a fraction of the work.

What the two functions do

savePrefixCache(prefixIds, w):
  mc = new ModelCache()
  prefill(prefixIds, w, mc)          // positions 0 .. P-1 populated, all 12 layers
  return mc                          // the saved cache (length P)

generateWithPrefix(saved, prefixLen, suffixIds, w):
  mc = saved.clone()                 // COPY — never mutate the saved cache
  for i in 0 .. suffixIds.length-1:
    logits = decodeStep(suffixIds[i], prefixLen + i, w, mc)   // positions continue at prefixLen
  return logits                      // last suffix token's next-token logits

What to build

Your task

Implement savePrefixCache and generateWithPrefix

In src/kvcache.ts, implement two functions and wire them through runPrefixReuse in the adapters:

savePrefixCache(prefixIds, w) — make a fresh ModelCache, prefill the prefix into it, and return it.
generateWithPrefix(saved, prefixLen, suffixIds, w) — clone the saved cache, then decode the suffix into the copy with decodeStep, giving the suffix's i-th token absolute position prefixLen + i. Return the last suffix token's next-token logits.

The adapter prefills the prefix once, saves it, and reuses it for the suffix. Its output must equal a from-scratch recompute of prefix + suffix.

tests/adapters.ts

export async function runPrefixReuse(prefixIds: number[], suffixIds: number[]): Promise<number[]> {
  const w = await weights();
  const saved = savePrefixCache(prefixIds, w);
  // reuse the saved cache; positions continue at prefixIds.length
  return generateWithPrefix(saved, prefixIds.length, suffixIds, w);
}

The two bugs this step exists to catch

Both are silent — wrong numbers, no crash:

Resetting the suffix's positions to 0 instead of prefixLen. The suffix's i-th token must get absolute position prefixLen + i. Start the suffix at 0 and GPT-2's wpe[position] term double-uses the prefix's positions; the grader runs a position-reset variant and that variant must fail recompute, proving the rule is enforced.
Mutating the shared saved cache. decodeStep appends the suffix's K/V to whatever cache you pass it. If you decode directly into saved, the next suffix inherits this suffix's tokens. Always work on saved.clone() — the same reason HF prompt caching does copy.deepcopy(prompt_cache) per reuse. The saved cache must come out exactly as long as the prefix.

How it's graded

Your reused-cache logits are compared against a from-scratch recompute of the full prefix + suffix prompt with allclose, atol 1e-3, rtol 1e-2 — the standard envelope for a 12-block float32 forward pass, where summation-order drift compounds. A separate position-reset variant (suffix starting at position 0) is run too, and it must fail the comparison: if it passes, you haven't actually continued the positions, and the step doesn't count.

Reuse the prefix When does prefix reuse break?