Build your own KV cache
The cache data structure·18 min

Build the KV cache

0 / 1
# grade this stage (from your course directory)
npx theinference test --stage build-the-cache --watch

Time to build the store. In src/kvcache.ts you implement the KVCache class: append(kRow, vRow) adds one position, length reports how many are stored, and clone() returns an independent deep copy. Each row is one token's K (or V) — a 768-vector, all 12 heads concatenated — so after k appends the buffer is the logical tensor [12, k, 64] and length == k. Get this exact and the rest of the course threads through it untouched.

What to build

Fill in the three members of KVCache in src/kvcache.ts. The fields k and v are arrays of rows (number[][]); length is already wired to k.length.

  • append(kRow, vRow) — push copies of the incoming rows onto k and v. Grows seq by exactly one.
  • clone() — return a new KVCache whose k and v are row-by-row copies, sharing no array references with the original.

Then wire the adapter below. It replays an append log and returns the stored K so the grader can compare it byte-for-byte to the reference.

Your task

Implement KVCache and runCacheAppend

Implement KVCache.append and KVCache.clone in src/kvcache.ts, then wire runCacheAppend(k, v) in tests/adapters.ts: construct a fresh KVCache, append each row pair in order, and return the stored K (cache.k) — an exact copy of the K rows you appended, in order.

tests/adapters.ts
import { KVCache } from '../src/kvcache.ts'

export function runCacheAppend(k: number[][], v: number[][]): number[][] {
  const cache = new KVCache()
  for (let i = 0; i < k.length; i++) cache.append(k[i]!, v[i]!)
  return cache.k // [seq, 768] — exact copy of the appended K rows
}
Append concatenates along the seq axis

Each append adds one row. After k appends, length == k and the logical buffer is [12, k, 64]. The grading is an exact match — atol=0 — so the K you return must equal the K you were handed, in the same order, with no rounding, reordering, or transposition.

Copy the rows — don't store the caller's arrays

append should store copies (kRow.slice()), not the passed-in arrays. If you keep references and the caller later mutates or reuses those arrays, your cache silently changes underneath you. The same discipline makes clone() correct: copy every row so the clone and original share no memory — Chapter 6's prefix reuse depends on it.

Watch the off-by-one

The classic structural bug is growth that doesn't match length: appending twice, appending nothing, or letting k and v drift to different lengths. After n appends both buffers must have exactly n rows. This step grades that invariant directly so it can never surface later inside attention.

How it's graded

The adapter replays the fixture's append log and the grader compares your returned K to the reference with allclose(atol=0) — bit-exact — and checks the row count is integer-exact. No tolerance: this is plain bookkeeping, so the only way to fail is a layout, copy, or off-by-one bug.