🚧 Week 3 Day 3: Paged KV Cache

🚧 This chapter is under review and may change.

In this chapter, we will design the paged KV cache, the storage abstraction behind paged attention. Continuous batching creates many request-owned caches with different lifetimes and sequence lengths. Storing each cache as one growing tensor makes every append depend on a contiguous allocation and makes batch construction revisit historical K/V.

Fixed-size pages separate a sequence’s logical order from its physical placement. The runtime can append, release, and reuse storage without moving a request’s complete history. This chapter changes the storage layout first and uses dense attention as a correctness checkpoint. Day 4 will read the pages directly.

📚 Readings

• vLLM Paged Attention Design
• Efficient Memory Management for Large Language Model Serving with PagedAttention

Why the Dense KV Layout Becomes Expensive

Right now, the mental model looks like this:

request A -> one dense KV tensor
request B -> one dense KV tensor
request C -> one dense KV tensor

Before attention, the runtime repacks them into:

keys:   [B, H, S_max, D]
values: [B, H, S_max, D]
mask:   [B, 1, L, S_max]

The trouble is that decode only adds a tiny amount of new information each step, but the dense layout keeps revisiting old KV.

For example, if a request already has 17 cached tokens and we decode 1 more token:

new useful work: append 1 token
dense repack view: rebuild 18 logical positions

For one request this is fine. For many live requests, the runtime spends more and more time moving previously computed KV instead of doing actual model work.

The Page Abstraction

Instead of storing each layer’s KV for a request as one long tensor, we divide storage into fixed-size pages:

key_pages:   pages with up to page_size token slots
value_pages: pages with up to page_size token slots

Each layer cache keeps a small page table:

page_ids = [12, 5, 3]
context_len = 10

That means:

page 12 -> tokens 0..3
page  5 -> tokens 4..7
page  3 -> tokens 8..9

The logical sequence is still length 10. The difference is that the runtime is no longer forced to represent it as one contiguous tensor.

The model owns one physical page pool per transformer layer. Request caches for the same layer share its pool, while every request-and-layer cache keeps its own page_ids, page_lens, and offset. A page id is therefore local to one layer, matching the K/V storage buffer that the attention kernel receives.

page_size is the physical page capacity. Unused tail slots are not part of the logical sequence; page_lens decides which prefix of each page is valid.

Why Fixed-Size Pages Help

The page abstraction gives us two immediate wins:

1. Appending a token usually updates only the current tail page in the pool.
2. Finished requests can return their pages to the layer’s shared free list.

This is the key memory-management idea behind paged attention systems such as vLLM.

Data Structures We Need

1. PagePool

The model should own one pool per layer, each with a free-page allocator and flat K/V page storage:

free_pages: available page ids for this layer
keys[page_id]:   physical key page
values[page_id]: physical value page

Requests share physical storage only when they are executing the same layer. Layer 0 and layer 1 may both use page ids [0, 1] because those ids address different buffers. Keeping the layer dimension outside page_id prevents a one-token write in one layer from copying or serializing the page storage for every other layer.

The backing slab should grow geometrically rather than by exactly one page. Keep logical num_pages separate from physical capacity so callers still see only allocated pages while pool growth copies old storage logarithmically many times.

Growing an exact-size slab from p to p + 1 pages copies approximately 1 + 2 + ... + p old pages, which is quadratic in the final page count. Starting with four pages and doubling capacity copies fewer than twice the eventual capacity across all growth events. Splitting the slab by transformer layer also prevents a new page in layer 0 from replacing the storage object used by every other layer. These two changes amortize allocator-copy work so most new-page allocations do not copy old pages, without changing the logical page table.

Implement this abstraction as TinyKvPagedPool.

2. PagedRequestCache

A layer cache for one request should track:

• page_ids
• page_lens
• offset
• page_size

Derived values:

• num_pages = len(page_ids)
• context_len = offset
• last_page_fill = page_lens[-1] when at least one page exists

Implement the request view as TinyKvPagedCache. Create it with a pool from the model; it should not allocate its own pool, because that would isolate one request from the shared page allocator.

Create one TinyKvPagedCache per transformer layer. Caches for different requests share a layer pool, but they do not share metadata: each cache owns its own page_ids, page_lens, and offset.

3. Tail-Append Logic

When new K/V arrives for one layer:

1. look at that layer cache’s last page
2. if there is room, append only the new slice into the tail page
3. otherwise allocate a new page and continue writing
4. update cache metadata such as page_lens and offset

This replaces the dense-cache pattern of repeatedly concatenating along the sequence dimension.

Make write cost proportional to the appended slice

MLX arrays are functional and lazily evaluated. Writing pages[page_id, :, start:end, :] = values may build an update whose output is the entire page tensor; a small slice in Python does not guarantee a slice-sized device update.

Implement paged_cache_update as a small extension primitive in your solution. Its output aliases the existing page buffer, and its Metal grid covers only H * new_tokens * D elements. Page storage is request state, so this mutation boundary is explicit and safe as long as the cache owns its page and attention depends on the returned array. Full-buffer copies remain only when geometric capacity grows.

Complete every learner-extension integration point before rebuilding:

• create src/extensions/src/paged_attention.cpp for the primitive and src/extensions/src/paged_attention.metal for its kernel,
• register those C++ and Metal sources in their respective lists in src/extensions/CMakeLists.txt,
• declare paged_cache_update in src/extensions/src/tiny_llm_ext.h, and
• register its Python binding in src/extensions/bindings.cpp.

Then rebuild:

pdm run build-ext

Test this behavior through the cache interface: append across a tail-page boundary, grow the slab, release and reuse page ids, and compare the gathered logical sequence with TinyKvFullCache.

Prefill with Pages

Suppose page_size = 4 and one prefill chunk contains 6 tokens:

chunk = [t0 t1 t2 t3 t4 t5]

One possible layout is:

page 7 <- [t0 t1 t2 t3]
page 2 <- [t4 t5]        # 2 valid tokens, 2 unused slots of capacity

That layer cache’s metadata becomes:

page_ids = [7, 2]
context_len = 6

The important property is that a later decode token can be appended to page 2 without touching page 7.

Decode with Pages

During decode, each live request adds one token at a time.

With paged storage:

1. compute one-token k and v
2. check whether the tail page still has space
3. write into that page if possible
4. allocate a new page only when the old one is full

So if page_size = 4 and context_len = 9:

page_ids = [12, 5, 3]

Appending token 9 only updates the last page instead of rebuilding all earlier KV.

Correctness Checkpoint: Gather Pages for Dense Attention

The cleanest first implementation is paged storage with dense gather.

That means:

• pages in each layer pool are the source of truth,
• layer caches stop owning one monolithic K/V tensor,
• layer caches only track page metadata,
• attention still receives dense K/V reconstructed from pages.

This checkpoint isolates the storage and lifecycle work before adding indirect GPU reads:

• page allocation and reuse can be tested independently;
• the gathered sequence can be compared directly with TinyKvFullCache;
• copy counters establish the cost that direct page traversal should remove.

How This Maps to tiny-llm

src/tiny_llm/paged_kv_cache.py

Add:

• TinyKvPagedPool
• TinyKvPagedCache

Keep TinyKvFullCache in src/tiny_llm/kv_cache.py as a baseline and test oracle.

The chapter’s execution path is:

1. write new K/V into the layer cache’s tail page or newly allocated pages,
2. gather the layer cache’s pages back into dense K/V,
3. feed that dense K/V into the readable dense attention equation.

This chapter changes the storage model while preserving the dense attention equation as a correctness oracle.

src/tiny_llm/batch.py

Requests should own per-layer cache handles instead of long dense K/V tensors.

The scheduler should still:

• perform chunked prefill,
• hold active requests,
• free cache pages when a slot finishes.

The difference is that freeing a request now means releasing all pages owned by its layer caches back to the pool.

Add a small rewind(n) lifecycle hook. Rewind lets a caller remove the newest logical tokens without rebuilding the retained prefix. It frees whole pages that are no longer needed and shortens the valid length of the final remaining page. The optional speculative-decoding chapter will use this operation when drafted tokens are rejected.

Design Questions

Before implementing, make sure the following are clear:

1. What page size should this repo use for teaching?
2. How do we represent the free-page allocator?
3. How do we prove that paged storage reconstructs the same logical KV as TinyKvFullCache?
4. How do request cache handles share a layer pool while keeping their own page metadata?
5. When do we materialize page writes to avoid MLX lazy-graph growth?
6. How do we grow physical capacity without copying all old pages on every allocation?

Task 1: Design PagePool

src/tiny_llm/paged_kv_cache.py

Modify TinyKvPagedPool.__init__, allocate_page, write_page_slice, and free_page in this task. For the slice-sized device write, replace the Week 3 Day 3 stubs tiny_llm_ext::paged_cache_update, PagedCacheUpdate::eval_cpu, and PagedCacheUpdate::eval_gpu in src/extensions/src/paged_attention.cpp, and implement paged_cache_update_kernel in src/extensions/src/paged_attention.metal. Their declaration, binding, source/Metal files, and CMake registration already exist; do not create a second paged-cache API.

Design layer-owned page pools that:

• own a free-page allocator,
• store flat fixed-size K/V pages,
• allocates and frees page ids,
• supports writing a chunk into page storage,
• grows backing capacity geometrically,
• updates only the appended physical slice between growth events,
• is shared by all request caches for that layer, but not by other layers.

Test logical size and physical capacity separately. Allocating the fifth page, for example, may create capacity for eight pages, but key_pages and value_pages exposed to the attention runtime should contain only the five allocated page ids.

Task 2: Design PagedRequestCache

src/tiny_llm/paged_kv_cache.py

Modify TinyKvPagedCache.__init__, update_and_fetch, release, and rewind. Use TinyKvPagedPool.write_page_slice from Task 1 for every physical append.

Replace the “one layer cache = one dense KV tensor” model with:

• page_ids
• context_len
• append logic over fixed-size pages
• release() for returning pages on request completion
• rewind(n) for dropping the newest n logical tokens

Task 3: Add a Dense-Gather Compatibility Path

src/tiny_llm/paged_kv_cache.py
src/tiny_llm/qwen3_week3.py

Modify TinyKvPagedCache.gather_dense in src/tiny_llm/paged_kv_cache.py, plus Qwen3ModelWeek3.__init__, Qwen3ModelWeek3.create_kv_cache, and Qwen3MultiHeadAttention.__call__ in src/tiny_llm/qwen3_week3.py. This checkpoint deliberately does not implement paged_attention; Day 4 owns that function.

Build a compatibility path that reconstructs dense K/V from pages and compares it against TinyKvFullCache.

This gives us a correctness check before we change the attention path itself. Instantiate the Week 3 model with enable_paged_attention=False in this chapter so its attention reads the gathered dense tensors. Day 4 switches the same model to page-table metadata and the paged kernel.

Run that cumulative checkpoint through the normal generation and benchmark entry points:

pdm run main --solution tiny_llm --loader week3 \
  --disable-paged-attention --model qwen3-0.6b

pdm run bench --solution tiny_llm --loader week3 \
  --disable-paged-attention --model qwen3-0.6b

In the next chapter, we will take the next step: instead of gathering dense K/V before attention, we will pass runtime metadata such as block_table directly into a paged attention path.

What Paging Changes

Apple silicon’s unified memory removes the discrete-device transfer boundary, but it does not remove allocation, fragmentation, or copying inside the GPU-visible heap. Fixed-size pages still let a server reuse freed capacity, grow requests without reserving their maximum sequence length, and batch requests with different context lengths. These are useful lifecycle mechanisms, but a fixed-batch trace measures KV-storage headroom rather than admission capacity. Claiming that more requests can be admitted requires a separate memory-capped sweep.

Report fragmentation with an aligned numerator and denominator. The benchmark finds the snapshot with the largest sum of unused slots in the final live page of every request/layer cache, then divides that sum by all token slots in live pages at the same snapshot. It reports the unused-slot bytes as well. Unused physical pool capacity is excluded from that fraction and remains visible in the separate live-page and capacity-page counters.

pdm run test --week 3 --day 3

Your feedback is greatly appreciated. Join our Discord community.
Found an issue? Open an issue or pull request at github.com/skyzh/tiny-llm.
tiny-llm-book © 2025 by Alex Chi Z is licensed under CC BY-NC-SA 4.0.