The memory bill
# grade this stage (from your course directory)
npx theinference test --stage the-memory-bill --watchCaching trades compute for memory: you stop recomputing K and V, but now you have to store them. This step makes that cost concrete. Given GPT-2 124M's config and a sequence length, compute the exact byte footprint of the cache — and the per-token growth that makes long context expensive.
total = 2 · n_layers · seqLen · n_heads · head_dim · dtypeBytes
^ ^ ^
K and V 12 heads 64 dims
(the leading 2)
GPT-2 124M: n_layers = 12 n_heads = 12 head_dim = 64
perToken = total / seqLen = 2 · n_layers · n_heads · head_dim · dtypeBytes
What to build
Implement kvBytes
Implement kvBytes(seqLen, dtypeBytes) in src/kvcache.ts returning { total, perToken }, where total = 2 · n_layers · seqLen · n_heads · head_dim · dtypeBytes for GPT-2 124M (n_layers=12, n_heads=12, head_dim=64) and perToken = total / seqLen. Then wire program.ts memory <seqLen> <dtypeBytes> to print exactly:
total=<bytes> perToken=<bytes>
All integer math — no floats, no rounding.
export function kvBytes(seqLen: number, dtypeBytes: number): { total: number; perToken: number } {
// total = 2 * 12 * seqLen * 12 * 64 * dtypeBytes
// perToken = total / seqLen
}
For GPT-2 124M fp32 (dtypeBytes=4) at seqLen=1024: 2·12·1024·12·64·4 = 75,497,472 bytes ≈ 72 MiB. Per token: 2·12·12·64·4 = 73,728 bytes. This is why context length blows up memory — linear in seqLen, per layer, per head. fp16 (dtypeBytes=2) halves both.
The cache stores both K and V — that's the 2 at the front of the formula. Drop it and every number is exactly half right, which is the easiest way to fail this step. The dtypeBytes is a parameter, not a constant: the same cache is 4 bytes/element in fp32 and 2 in fp16.
Each (seqLen, dtypeBytes) case in fixtures/mem_cases.json is checked integer-exact against the closed-form target — no tolerance. Pure integer arithmetic either lands on the byte count or it doesn't.