Build your own KV cache
Prove the win·25 min

Prove the work reduction

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

You've built a cache that reproduces recompute's logits at every step. Now prove it was worth it — not with a stopwatch, but with a count. A clock measures your laptop's mood: thermal throttling, garbage collection, whatever else is running. The work itself is exact. So this step grades a deterministic operation count: the cached path must do O(n) work per step, the recompute baseline O(n²). The gap between the two curves is the proof, and it's the same gap whether you run it on a phone or a datacenter GPU.

Two profiles, same prompt (promptLen=4, nSteps=4)
step k   n = promptLen + k     recompute n(n+1)/2     cached  n + 1 positions
  0          4                       10                     5
  1          5                       15                     6
  2          6                       21                     7
  3          7                       28                     8
                                      ^ quadratic            ^ linear

What to build

Your task

Implement cachedProfile and wire runWorkCount

Implement cachedProfile(promptLen, nSteps) in src/kvcache.ts: at step k (0-indexed) the cache already holds promptLen + k positions and you attend the one new query over all of them plus itself, so the per-step count is promptLen + k + 1. Return the length-nSteps array.

You already wrote recomputeProfile back in Chapter 0. Wire both into runWorkCount in tests/adapters.ts: return two rows — cached first, recompute second — so the grader can hold the linear and quadratic curves side by side.

tests/adapters.ts
import { cachedProfile, recomputeProfile } from '../src/kvcache.ts'

export function runWorkCount(promptLen: number, nSteps: number): number[][] {
  return [
    cachedProfile(promptLen, nSteps),    // O(n): step k touches promptLen + k + 1 positions
    recomputeProfile(promptLen, nSteps), // O(n²): step k touches n(n+1)/2 pairs, n = promptLen + k
  ]
}
The +1 is the new token attending to itself

At step k the cache holds promptLen + k past positions, and the freshly computed query attends over all of them plus its own new K/V row — that's the + 1. Drop it and step 0 reports promptLen instead of promptLen + 1; the whole curve shifts and the grader rejects it. This mirrors recomputeProfile's n = promptLen + k: same indexing convention, different cost.

The contrast is the artifact

Returning both rows isn't bookkeeping — it's the whole point of the step. Recompute's row grows quadratically because every step re-attends the entire sequence; cached's row grows by exactly one each step because the history is already stored and only the new token's K/V is computed. Two arrays, one linear and one quadratic, from the same prompt: that's the speedup made legible.

Why a count and not a timer

Wall-clock speedups are real but flaky — they drift with hardware, load, and the JIT. An operation count is a property of the algorithm, identical on every machine, so the grader can demand an integer-exact match. You'll still print an advisory wall-clock number in the next step; it just never decides pass or fail.

How it's graded

runWorkCount(promptLen, nSteps) returns a [2, nSteps] array. Row 0 (cached) is compared integer-exact to the O(n) profile; row 1 (recompute) integer-exact to the O(n²) profile. No tolerance — these are closed-form integer counts, so the only ways to fail are an off-by-one, swapped rows, or a wrong formula.