The cache data structure·10 min

A cache for every layer

0 / 1

# grade this stage (from your course directory)
npx theinference test --stage multi-layer-cache --watch

One KVCache holds one layer's history. GPT-2 has 12 attention layers, so the model needs 12 of them. ModelCache is that container: an array of 12 independent KVCaches, one per layer. Each layer appends to and reads from only its own cache — layer 5's keys never leak into layer 6's attention. This is the structure both prefill and decode write into for the rest of the course.

What to build

In src/kvcache.ts, ModelCache already initializes layers to 12 fresh KVCaches. Implement its clone() to return a new ModelCache whose 12 layers are each a clone() of the originals — no shared references, so copying the model cache copies every layer independently.

Then wire the adapter. It receives a list of layer indices; for each one it appends a single dummy token (one K row and one V row) to that layer's cache, leaving the other layers untouched. Finally it returns each layer's logical shape [12, len, 64] as a 12-element row — one row per layer, in layer order 0..11.

Your task

Implement ModelCache and runModelCacheShapes

Implement ModelCache.clone in src/kvcache.ts, then wire runModelCacheShapes(appends) in tests/adapters.ts: build a fresh ModelCache, and for each layer index in appends, append one dummy token (a 768-vector for K and one for V) to that layer's cache. Return 12 rows — for each layer l, its shape [12, mc.layers[l].length, 64].

tests/adapters.ts

import { ModelCache } from '../src/kvcache.ts'

export function runModelCacheShapes(appends: number[]): number[][] {
  const mc = new ModelCache()
  const k = new Array(768).fill(0)
  const v = new Array(768).fill(0)
  for (const layer of appends) mc.layers[layer]!.append(k, v)
  // one [n_heads, len, head_dim] = [12, len, 64] row per layer, 12 rows
  return mc.layers.map((c) => [12, c.length, 64])
}

Independence is the point

This is the container prefill and decode both write into. Keeping the 12 caches independent is exactly what lets each layer read only its own history — appending to layer 3 must leave layers 0–2 and 4–11 at their previous lengths. If a single append shows up in two layers' shapes, your layers array is sharing one KVCache instead of holding twelve.

Clone deep, layer by layer

ModelCache.clone() must call each layer's clone() — this.layers.map((l) => l.clone()) — not copy the array of references. A shallow copy hands the new model cache the same 12 KVCache objects, so an append to one corrupts the other. Chapter 6's prefix reuse copies the whole model cache before decoding a suffix; this is where that copy is made correct.

Each layer counts only its own appends

The shape row for layer l is [12, len, 64] where len is that layer's length — the number of tokens appended to layer l alone, not the total across all layers. If appends is [0, 0, 5], layer 0 ends at length 2, layer 5 at length 1, and every other layer at length 0.

How it's graded

The adapter runs the fixture's append plan and the grader compares all 12 returned shapes to the reference, integer-exact. Every layer's len must match — a wrong count on any single layer fails the step. No tolerance: this is pure structural bookkeeping.

Build the KV cache One token