🚧 Week 3 Day 4: Direct Paged Attention

🚧 This chapter is under review and may change.

In this chapter, we will build direct paged attention. The scheduler passes request-local block tables and context lengths to a GPU kernel, which reads K/V from the shared layer pool without gathering a dense batch first.

Prerequisite: Complete Week 3 Day 3’s paged storage and Week 2 Day 5’s online-softmax attention. The new concept here is translating logical K/V positions through a block table. Tiled FlashAttention comes only after this direct path works.

Paged KV Cache vs Paged Attention

These two ideas are related, but they are not the same:

1. Paged KV cache KV is stored in fixed-size pages.
2. Paged attention The attention path reads KV directly from those pages via metadata such as a page table.

You can implement the first one without the second one, but the real serving payoff comes when both are present.

The Metadata a Paged Runtime Needs

Once KV is paged, dense B x H x S x D tensors are no longer the natural runtime representation. Instead, the runtime should prepare metadata like:

block_table:  [B, max_pages_per_request]
context_lens: [B]

For the current layer being executed:

• block_table[b, i] gives the page id for request b’s current-layer logical page i
• context_lens[b] gives the valid token count for request b

This is the bridge between the scheduler and the attention kernel.

A production runtime often also carries write-side metadata such as slot_mapping. For this chapter, we keep the write side inside the cache and focus on the read-side metadata needed by attention.

Why block_table Matters

Suppose one layer cache for request A has:

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

Then the logical sequence positions map to physical storage like this:

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

The attention runtime does not need a fully gathered dense tensor if it already knows:

• which current-layer page each logical block lives in,
• how long the context is,
• and where the current query positions are.

That is exactly what block_table and context_lens encode.

The Paged Attention API

At this point, the runtime should grow a new attention entry point:

paged_attention(
    query,
    key_pages,
    value_pages,
    block_table,
    context_lens,
    page_size,
    scale=None,
    mask="causal",
)

With shapes like:

query:          B, H_q, L, D
key_pages[i]:   1, H_kv, page_size, D
value_pages[i]: 1, H_kv, page_size, D
block_table:    B, max_pages
context_lens:   B

The source length is no longer represented by one contiguous tensor dimension. The operator reconstructs it logically from the page table.

In this chapter, paged_attention should read pages directly from a GPU kernel. The runtime contract is now: model code and batching code pass pages plus metadata, and the attention kernel walks that metadata without first rebuilding dense K/V.

Prefill Metadata

During prefill, a chunk may span multiple pages. The runtime needs to know:

• which current-layer pages already existed,
• which new pages were allocated,
• how many valid tokens are in the tail page,
• how to map incoming K/V rows into page storage.

In this teaching implementation, the cache still owns the write-side bookkeeping. The attention path only needs the block table after the write is done.

Decode Metadata

During decode, each active request typically writes one token.

The runtime should be able to:

1. append the token’s K/V to the current tail page,
2. allocate a new page only if the tail page is full,
3. update the current layer cache’s context_len,
4. run attention over the full logical context using block_table

This is the point where decode stops paying the repeated dense-repack cost from Day 1.

Choose a Schedule for Each Query Shape

Before implementing the GPU path, separate decode from prefill. A single tile shape cannot keep the GPU busy for both a one-token query and a long prompt. Use these design rules:

1. Preserve the Week 2 BF16 model boundary and reuse its internal accumulation policy unchanged.
2. For short queries, expose parallelism across the cached context. Do not reserve most of a threadgroup for query rows that do not exist.
3. For prefill, begin with a direct page-walking schedule whose address calculation is easy to validate.
4. Use the readable equation written with mlx.core and the dense Week 2 attention kernel in your solution as correctness oracles for the new page-walking schedule.

Start with this dispatch plan and treat its thresholds as values to verify on your hardware:

Put the shape decision at the extension boundary rather than converting inputs or falling back to dense attention in Python. Benchmark values immediately below and above each threshold while keeping the model-facing paged_attention API unchanged.

How This Maps to tiny-llm

src/tiny_llm/attention.py

Add a new function:

def paged_attention(...):
    ...

In your solution, make it a correctness-first page-walking Metal kernel with online softmax:

1. use block_table[b] to find the physical pages for request b,
2. use context_lens[b] to ignore unused tail capacity,
3. visit K/V in small tiles instead of materializing dense K/V,
4. merge each tile into the output with online softmax.

The important change from dense attention is the K/V address calculation. Dense attention can advance through dense K/V by pointer arithmetic. Week 3 must translate each logical key position through block_table first:

logical key position -> logical page -> physical page id -> slot in page

After that lookup, the online-softmax update is the same recurrence as Week 2 Day 4. Keep the page-walking schedule simple enough that block-table and tail-page boundary errors are visible. Day 5 will tile its inner matrix work while preserving this address calculation.

One-token decode needs a different work decomposition. A 64-row prefill tile would leave almost every query row idle, so dispatch short queries to a vector-oriented kernel that partitions the context across SIMD groups and merges their partial online-softmax states. Do not run decode through a fixed 32-row scalar prefill tile.

The page pool should therefore expose contiguous physical storage:

key_pages:   P, H_kv, page_size, D
value_pages: P, H_kv, page_size, D

A Python list of page tensors is convenient for teaching the allocator, but a GPU kernel needs a single buffer so page_id can be turned into an address.

src/tiny_llm/qwen3_week3.py

The attention module should call the paged runtime directly:

metadata = cache.update_and_fetch_paged(...)
x = paged_attention(...)

Week 3 cache handles are expected to provide paged metadata. If a dense cache is passed to the Week 3 model, that is a programming error rather than a signal to silently fall back to dense attention.

src/tiny_llm/batch.py

The scheduler now needs to prepare runtime metadata instead of only dense K/V:

• per-layer page tables for each active request
• padded batch block_table
• context_lens

This is where continuous batching and paged attention finally connect. On Day 1, batching worked by repacking tensors. Here, batching should work by reusing page tables and updating only the new slots.

Implementation Order

Use this implementation order:

1. paged storage
2. block_table / context_lens plumbing
3. correctness-first page-walking GPU attention
4. model and batch dispatch

Each step has a direct correctness check before the next abstraction is added.

What Must Hold, and What Breaks If It Doesn’t

These are the invariants worth checking in tests:

1. context_len equals the number of written logical token positions. If it is too small, attention skips written K/V; if it is too large, attention reads unwritten tail slots. Either case makes paged output diverge from the dense baseline.
2. block_table reconstructs the same logical K/V order as the dense baseline. A wrong mapping can pair a query with the wrong token’s K/V and change the output even when every page contains valid data. Reordering complete pages can change a causal-prefix result because it changes which K/V pairs each query can see. By contrast, one-token decode over all positions in complete pages is permutation-invariant to their order when each K/V pair moves together.
3. The allocator gives each page to only one live cache handle unless sharing is explicit. If two live handles alias a page, a write for one request overwrites K/V that the other request can still attend to.
4. Releasing a request returns every page owned by every layer cache exactly once. Missing a page leaks pool capacity; returning one twice raises the pool’s already-free error instead of completing cleanup.
5. Decode allocates a new page only when the tail page overflows. Allocating earlier strands writable tail slots and inflates the used-page count. This course pool grows instead of reporting exhaustion, so that waste can force backing storage to grow and copy earlier, increasing memory pressure.

Task 1: Add Batch Metadata

src/tiny_llm/paged_kv_cache.py
src/tiny_llm/kv_cache.py
src/tiny_llm/batch.py

Modify TinyKvPagedCache.block_table, context_lens, paged_metadata, and update_and_fetch_paged in src/tiny_llm/paged_kv_cache.py. Then update Request.try_prefill and _step in src/tiny_llm/batch.py to carry those arrays for every active request.

Extend the batch cache and scheduler so they can prepare:

• block_table
• context_lens

for all active requests.

Task 2: Define paged_attention

src/tiny_llm/attention.py
src/extensions/src/paged_attention.cpp
src/extensions/src/paged_attention.metal

Modify these exact starter functions:

• paged_attention in src/tiny_llm/attention.py;
• tiny_llm_ext::paged_attention, PagedAttention::eval_cpu, and PagedAttention::eval_gpu in src/extensions/src/paged_attention.cpp;
• paged_attention_decode and paged_attention_scalar_f32 in src/extensions/src/paged_attention.metal.

This checkpoint also turns the already-readable quantized token lookup into the Week 3 one-dispatch path. Modify QuantizedEmbedding.__call__ in src/tiny_llm/embedding.py, tiny_llm_ext::quantized_embedding plus QuantizedEmbedding::eval_cpu/eval_gpu in src/extensions/src/quantized_matmul.cpp, and quantized_embedding_w4a16_g128 in src/extensions/src/quantized_matmul.metal. The starter declarations, bindings, stubs, and build registrations for both operations already exist and remain fail-closed until you replace them.

Add a paged attention interface whose inputs come from the paged runtime rather than a dense reconstructed S dimension. Preserve the Week 2 precision contract without adding a new model dtype or conversion at the serving layer.

Walk every request’s block table while keeping online-softmax state:

running_max = max(previous_max, page_max)
running_sum = previous_sum * exp(previous_max - running_max) + page_sum
output = previous_output * exp(previous_max - running_max) + page_output

After all visible pages are consumed, divide output by running_sum. This is the key idea that lets the kernel avoid materializing dense K/V while still producing the same result as dense attention.

Implement two correctness-first GPU dispatches:

1. For L <= 8, partition logical context positions across SIMD groups and merge their partial (max, sum, output) states in threadgroup memory. The initial schedule uses 32 SIMD groups per query. Resolve the physical page once, then let group g visit slots g, g + 32, g + 64, and so on within that page; do not divide and reload block_table for every token.
2. For longer queries, assign query rows to a direct page-walking schedule and resolve every K/V tile through block_table. When a tile is aligned and cannot cross a page boundary, share its one physical page id across the whole tile. Favor inspectable ownership over the final tiled performance schedule.

Compare small deterministic fixtures with the readable equation written with mlx.core and the dense Week 2 attention path before tuning the page-walking schedule.

For the final Qwen decode schedule, specialize BF16 D = 128: each lane owns four contiguous dimensions of Q, K, V, and the output. After all context positions are visited, transpose the 32 partial output vectors through one compact 32×32 threadgroup tile. Each SIMD group then reduces four dimensions with simd_sum. This organizes the reduction in 4.25 KiB of scratch instead of storing one full partial vector per scalar output thread. Keep a generic BF16 specialization for other head dimensions so the optimization cannot silently reinterpret D = 32 as D = 128.

Your solution’s boundary

MLX remains the array runtime for shapes, reshapes, transposes, contiguous storage, dtype conversion, allocation, and custom-primitive dispatch. The attention implementation itself must remain in your solution: do not call mx.fast.scaled_dot_product_attention, reuse an MLX attention/Steel kernel, or reconstruct dense K/V and express the paged operator as MLX matmul plus softmax. MLX SDPA may appear only in tests and benchmarks as an external correctness oracle and performance baseline.

Both prefill and decode read page storage through this interface. Do not add a dense-only special case: Day 5 optimizes this same paged contract.

Task 3: Dispatch from the Model

src/tiny_llm/qwen3_week3.py

Modify Qwen3MultiHeadAttention.__call__, Qwen3ModelWeek3.__init__, and Qwen3ModelWeek3.__call__ to select the paged path and enable the custom embedding only at this cumulative checkpoint.

Update the model so it can route to paged attention when the cache provides paged runtime metadata.

Append K/V to the page pool and pass its metadata to attention for every query shape. Long queries use the direct paged-prefill schedule from this chapter; short queries use the vector paged-decode schedule. Neither path changes cache dtype or gathers a dense K/V tensor.

This creates the Day 4 routing policy:

prefill or long chunk -> direct page-walking attention
decode or short chunk -> paged vector attention

Day 5 replaces the long-query schedule with paged FlashAttention without changing this model-facing policy. --disable-paged-attention is a Day 4 dense-gather teaching ablation, not the completed serving path.

Task 4: Connect It to Continuous Batching

src/tiny_llm/batch.py

Modify Request.try_prefill, Request.decode_done, _step, and batch_generate. Request cleanup must call TinyKvPagedCache.release for every layer cache.

Update request admission, slot reuse, and request removal so that:

• finished requests free their pages,
• in this teaching implementation, that means freeing pages from every layer cache,
• new requests allocate from the corresponding layer pool,
• active decode steps reuse page metadata instead of rebuilding dense K/V.

After this chapter, the serving stack has the right structure for a real high-throughput runtime: paging is no longer just a storage trick, but part of the execution model itself.

Measure the Direct Page Walk

The goal of this lab is to decide when direct page traversal is useful. Paged attention is not automatically a faster attention operator: it trades regular, contiguous K/V access for flexible allocation and removes the dense repack that would otherwise happen before attention. Your measurements must include both sides of that trade.

Record three operator baselines on the same machine:

1. the dense Week 2 attention path in your solution, including any required K/V gather,
2. your direct paged-attention path,
3. the MLX attention path as a production-library baseline.

Use the same Qwen3-4B decode shape for all three paths. The dense control must include its required page-to-dense gather; the direct path reads the same page metadata; the MLX row measures its fused attention operator on the already gathered tensor:

pdm run bench-week3-attention --offline --contexts 128 1024 \
  --page-size 128 --warmup 5 --iterations 60 --repeats 4 \
  --cooldown-seconds 1 \
  --json-output benchmark_results/m4-pro-qwen3-4b-week3-attention-mlx-0.32.0.json

Each value is the median of four balanced fresh-process medians, with 60 synchronized calls after five warmups per process:

Direct traversal is 1.9% slower than dense-plus-gather at 128 tokens, but 40.7% faster at 1,024 tokens. MLX remains faster at both shapes. The checked BF16 outputs match the readable dense equation within the documented 2e-2 tolerance.

Checkpoint 1: Establish a Correct Direct Path

Implement the simplest page-walking kernel first. Verify that it:

• reads K/V through block_table without constructing a dense K/V tensor,
• ignores unused slots in the final page,
• matches dense attention for several page boundaries and context lengths,
• supports grouped-query attention when H_q != H_kv.

The correctness schedule may be slower than dense attention. At this checkpoint, the useful result is a trustworthy baseline and a working runtime interface.

Checkpoint 2: Design the Decode Schedule

Optimize for the one-token decode shape instead of treating page traversal as a serial loop. Work through these changes one at a time and benchmark after each one:

1. assign the lanes of a SIMD group to adjacent elements of a head so K/V loads can be coalesced,
2. load a page-table entry once and reuse it for all positions in that page,
3. keep the query and online-softmax state in registers across page tiles,
4. combine partial dot products with SIMD reductions instead of threadgroup scratch memory and repeated barriers,
5. specialize the L = 1 decode case so it does not carry prefill control flow,
6. prepare the batch’s page metadata once per scheduler step rather than once per layer or attention head.

Also benchmark the write path separately. A fast page-reading kernel cannot recover time lost to a functional whole-cache update before every layer.

For each change, explain which cost it targets: memory traffic, synchronization, address calculation, or dispatch overhead. Keep a change only when the measured result supports the explanation.

Optimize the paged path in your solution for the Qwen head dimension of 128, but keep one grouped-query schedule rather than duplicating the online-softmax recurrence for individual GQA ratios. This keeps the relationship among page traversal, head mapping, and reduction visible in one kernel.

Checkpoint 3: Evaluate the Serving System

Operator latency alone does not capture the purpose of paging. Run an end-to-end workload with requests entering and leaving the batch, then report:

• time per decode step and aggregate tokens per second,
• peak KV-cache memory and the number of live requests admitted,
• bytes or time spent gathering and repacking K/V,
• paged-attention latency relative to the dense Week 2 path in your solution and MLX.

Keep the dense path as a teaching ablation so you can measure when contiguous attention is faster. The completed serving route stays paged: it eliminates repacks, reuses pages across scheduler steps, and leaves more measured KV headroom in the fixed-batch trace. Proving that it admits more concurrent requests requires a memory-capped admission sweep. Day 5 optimizes its long-prefill schedule rather than routing around the page-table contract.

Use the paired serving runner rather than a preallocated static request:

pdm run bench-serving-progression --offline --repeats 4 \
  --model qwen3-4b --num-seqs 16 --batch-size 4 \
  --min-input-len 128 --max-input-len 1024 \
  --min-output-len 32 --max-output-len 128 --prefill-step 128 \
  --warmup 1 --cooldown-seconds 1 \
  --json-output benchmark_results/m4-pro-qwen3-4b-week3-serving-mlx-0.32.0.json

It compares Week 2 dense batch reconstruction, Week 3 paged storage with the dense-gather compatibility path, and Week 3 direct paged attention. It resets page capacity after warmup and reports prefill, output, and decode throughput alongside peak KV bytes, copy volume, page reuse, and tail fragmentation. The direct path’s four-process medians are 679.56 prefill tok/s, 41.88 output tok/s, 82.11 decode tok/s, and 0.558 requests/s. Its synchronized decode calls take 38.27/39.83/43.46 ms at median/p95/max; the completion gaps, which include intervening scheduler and prefill work, are 39.13/224.70/240.99 ms.

pdm run test --week 3 --day 4

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.