Project 09 — KV Cache Manager

★ Replace “one big slot per request” with paged blocks, and get prefix sharing almost for free.

1. What you build#

A block-based KV memory manager — the PagedAttention idea — with a free list, per-request block tables, reference counting, copy-on-write, hash-based prefix sharing, and LRU eviction of cached-but-unused blocks. First as a standalone simulator driven by a workload trace, then integrated into your Project 08 engine.

Diagram — Life of a physical block#

stateDiagram-v2
    direction LR
    [*] --> Free
    Free --> InUse: allocate, ref 1
    InUse --> Shared: prefix hit or fork, ref + 1
    Shared --> InUse: release, ref - 1
    InUse --> Cached: released while full and hashed, ref 0
    InUse --> Free: released while partial
    Cached --> InUse: prefix hit
    Cached --> Free: LRU eviction

2. Why it matters#

In Project 08 every slot reserved max_len tokens of KV, so a 50-token chat held memory for 2,048. Concurrency is capped by memory, so wasted memory is wasted revenue. Paging fixes the waste, and once KV lives in shareable, hashable blocks, prefix caching becomes a lookup instead of a feature.

It is virtual memory, re-invented for attention. If you know OS paging, you already know the design; the project is about seeing how exact the analogy is.

3. Read first#

• V.10 — PagedAttention
• V.11 — Prefix caching
• X.05 — Fragmentation and allocators
• X.06 — OOM
• II.06 — Memory allocation and virtual memory

4. Spec#

Pool:        num_blocks physical blocks, block_size tokens each (default 16)
             storage  K, V : [n_layer, num_blocks, n_kv_head, block_size, head_dim]
Block:       ref_count, hash (None until full), last_used
Table:       request_id -> [physical block ids]      (logical -> physical)

API
  can_allocate(n) / allocate(n) / free(request)
  append_token(request)            -> allocates a new block at boundaries
  fork(parent, child)              -> share all blocks, ref_count += 1
  write(request, block_idx)        -> copy-on-write if ref_count > 1
  lookup_prefix(token_ids)         -> longest run of cached full blocks
  evict()                          -> LRU over blocks with ref_count == 0

Prefix hash:  h_i = hash(h_{i-1}, tokens in block i)    (chain — a block's identity includes
                                                         everything before it)

5. Milestones#

1. Allocator. Free list, allocate, free. Property-test it: random alloc/free sequences never double-allocate or leak; free + used == total always.
2. Block tables and append. Model requests growing one token at a time.
3. Simulator. Replay a trace (arrival, prompt length, output length) against (a) slot allocation and (b) paged allocation with the same total memory. Report peak concurrency and wasted tokens.
4. Prefix sharing. Full blocks get a chained hash and enter a hash → block map. A new request reuses the longest matching prefix and only prefills the remainder.
5. Ref counts and eviction. Freed-but-hashed blocks stay cached with ref_count == 0 and are evicted LRU only when the free list is empty.
6. Copy-on-write. fork a request (beam search, n > 1); writing into a shared last block copies it first.
7. Integrate. Swap Project 08’s SlotKV for the paged pool. Attention gathers K/V through the block table. Equivalence test must still pass.

6. Starter skeleton#

var ErrOutOfBlocks = errors.New("out of KV blocks")

type BlockPool struct {
	blockSize int
	ref       []int             // reference count per block
	hash      []uint64          // chained content hash per block; 0 = not sealed
	free      []int             // blocks with nothing in them
	cached    map[uint64]int    // hash -> block id
	lru       *list.List        // cached block ids, least recently used at the front
	lruPos    map[int]*list.Element
	tables    map[string][]int  // request id -> block table
}

func (p *BlockPool) take() (int, error) {
	if n := len(p.free); n > 0 {
		b := p.free[n-1]
		p.free = p.free[:n-1]
		return b, nil
	}
	for e := p.lru.Front(); e != nil; e = e.Next() { // LRU over unreferenced cached blocks
		if b := e.Value.(int); p.ref[b] == 0 {
			delete(p.cached, p.hash[b])
			p.hash[b] = 0
			p.lru.Remove(e)
			delete(p.lruPos, b)
			return b, nil
		}
	}
	return 0, ErrOutOfBlocks
}

// Admit builds the request's block table and returns how many prompt tokens were
// already cached (skip prefilling those).
func (p *BlockPool) Admit(rid string, tokenIDs []int32) (hit int, err error) {
	var table []int
	var h uint64
	for i := 0; i+p.blockSize <= len(tokenIDs); i += p.blockSize {
		h = chainHash(h, tokenIDs[i:i+p.blockSize]) // depends on ALL tokens before it
		b, ok := p.cached[h]
		if !ok {
			break
		}
		p.ref[b]++
		p.lru.MoveToBack(p.lruPos[b])
		table, hit = append(table, b), hit+p.blockSize
	}
	need := (len(tokenIDs) - hit + p.blockSize - 1) / p.blockSize
	for i := 0; i < need; i++ {
		b, err := p.take()
		if err != nil {
			p.tables[rid] = table
			p.Release(rid) // give back what we took
			return 0, err
		}
		p.ref[b] = 1
		table = append(table, b)
	}
	p.tables[rid] = table
	return hit, nil
}

func (p *BlockPool) Release(rid string) {
	for _, b := range p.tables[rid] {
		if p.ref[b]--; p.ref[b] == 0 && p.hash[b] == 0 {
			p.free = append(p.free, b) // unhashed (partial) blocks go straight back
		}
	}
	delete(p.tables, rid)
}

// Yours to add: Seal(rid, blockIdx) — when a block fills, compute its chained hash,
// set p.hash[b], and insert it into p.cached and p.lru so later requests can find it.

7. What to measure#

8. Done when#

• Allocator invariants hold under randomized testing.
• The simulator shows the concurrency gain of paging on a heavy-tailed trace.
• Integrated engine passes the Project 08 equivalence test.
• A second request with a shared prefix skips that prefill, and you measured the TTFT.
• Copy-on-write is covered by a test where two forks diverge without corrupting each other.
• You can explain why the prefix hash must be chained.

9. Common pitfalls#

Hashing a block by its own tokens only. Two different prompts with the same 16 tokens at different positions would share KV. KV depends on the whole prefix — chain the hash.

Hash collisions treated as impossible. A collision here leaks one user’s context into another’s output. Use a strong hash, or verify tokens on hit.

Sharing the partially filled last block. Only full, immutable blocks are safe to share.

Ref count leaks on preemption or cancellation paths. Test every exit.

Cross-tenant sharing. Prefix-cache timing is a side channel (XI.09). Salt the hash per tenant.

10. Stretch goals#

• Radix tree over token sequences instead of a flat hash map, with LRU at tree nodes — the SGLang RadixAttention design (VIII.12).
• Second tier: spill evicted blocks to CPU memory or disk and bring them back. Compare restore time against recompute (XIII.07). The break-even is the interesting result.
• Preemption policies: LIFO vs largest-first; measure wasted work.
• Feed the simulator a real multi-turn chat trace and report hit rate by turn number.

11. Interview questions this project answers#

1. What problem does PagedAttention solve that a contiguous KV cache cannot?
2. How does a block table map logical to physical KV, and what is the OS analogue?
3. How does prefix caching fall out of paging?
4. When is copy-on-write needed?
5. How do you choose the block size?
6. What is the security risk of a shared prefix cache?

12. Next#

Project 10 — Quantized inference