🚧 Week 3 Day 5: Paged FlashAttention

🚧 This chapter is under review and may change.

In this chapter, you will replace the correctness-first implementation for the supported BF16 long-prefill hot case with paged FlashAttention. The operator still translates logical K/V positions through block_table, but it now stages page-backed tiles on chip and combines them with online softmax. Short decode and shapes outside the optimized region keep their correct Day 4 fallbacks.

Day 5 starts differently from an ordinary correctness checkpoint: the public behavior tests may already pass with your Day 4 implementation. That is intentional. Your work is to make the optimized path real and reachable while preserving the same observable behavior. The manual control-flow trace at the end of the chapter is the completion feedback for that performance change; it is guidance, not a hidden grading requirement.

This is a required chapter. FlashAttention belongs here rather than in Week 2 because the serving model’s real K/V source is now the page pool. Building a dense-only kernel first would create a second attention path and then require students to relearn its memory schedule around page translation.

Prerequisites

This chapter combines four prerequisites:

• Week 2 Day 5 introduced the online-softmax recurrence.
• Week 2 Day 6 introduced the cooperative 32×32 tile built from BF16 8×8 SIMD-matrix fragments.
• Week 3 Day 3 introduced physical pages and block tables.
• Week 3 Day 4 introduced direct page-walking attention and the decode schedule.

No new model dtype is introduced here. Preserve the Week 2 precision contract at the paged_attention boundary.

Why Optimize the Paged Path

A conventional attention expression materializes a score matrix with shape L × S. A page-walking implementation can avoid gathering K/V and still make that intermediate too large. Paged FlashAttention does both:

1. it resolves each K/V tile through block_table instead of gathering a dense cache;
2. it keeps only a query tile, one K/V tile, and online-softmax state on chip;
3. it writes the normalized output once after all visible pages are consumed.

The algorithm remains numerically equivalent attention under the course’s BF16 rounding tolerance. Only the order of loads and reductions changes.

Keep the Day 4 Interface

Do not add a second model-facing operator. Continue to call:

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

The reference solution puts this shape dispatch inside the extension:

The completed Week 3 model therefore has one paged-attention contract and an optimized BF16 long-prefill region. The public tests do not grade a kernel name, dispatch threshold, tile shape, private route flag, helper, or source layout. A behaviorally equivalent implementation, including one built with an external low-level library, is valid.

Task 1: Tile Queries and Paged K/V

The reference walkthrough begins paged_attention_mma_bf16_d128 in src/extensions/src/paged_attention.metal. Keep paged_attention_decode, paged_attention_scalar_f32, and paged_attention_scalar_bf16 from Day 4 unchanged; they remain the short-query and generic controls.

Use eight SIMD groups to cover a 64-row query block. Each SIMD group owns eight query rows and represents matrix operands as 8×8 fragments. Stage 32 logical K/V positions per iteration.

For every logical key row in a tile:

logical_position = tile_start + row
logical_page     = logical_position / page_size
slot             = logical_position % page_size
physical_page    = block_table[batch, logical_page]
address          = pages[physical_page, kv_head, slot, :]

Resolve the physical page while staging the tile. The matrix multiply should not know whether two adjacent logical rows came from adjacent physical pages.

The Qwen path uses 128-token pages and a 32-token K/V tile. An aligned tile is therefore physically contiguous even when the logical sequence as a whole is not. Assign each thread contiguous elements through a cooperative block loader so adjacent lanes issue coalesced reads. Keep a generic loader for a tile that crosses a page boundary. The reference reuses the course-owned CooperativeTileLoader and direct simdgroup_matrix fragments from Week 2 so those mechanisms remain visible. You may choose a different internal helper or low-level library as long as the operator retains direct page-table semantics and equivalent public behavior.

Tail cases are required. A query block, K/V tile, final page, or context may be partially full, and physical page ids need not be consecutive.

Task 2: Compute Tiled Online Softmax

Continue modifying paged_attention_mma_bf16_d128 in src/extensions/src/paged_attention.metal. This task fills the tiled online-softmax body; it does not add another public function.

For each query tile, maintain one running maximum, one running sum, and an unnormalized output accumulator per row. For each K/V tile:

1. compute Q @ Kᵀ with the Week 2 SIMD-matrix fragments;
2. apply scale and causal bounds;
3. merge the tile maximum into the running maximum;
4. rescale the previous sum and output accumulator;
5. compute exponentials for the current scores and update the running sum;
6. multiply the tile probabilities by V and update the output accumulator.

After the final visible tile, divide each output row by its running sum and store it using the model-facing dtype.

Multiply the attention scale by log2(e) once and use fast::exp2 for online-softmax rescaling inside the hot tile loop. This is mathematically equivalent to natural exponentials and avoids repeating a base conversion.

The causal offset is context_len - L. A key at logical position s is visible to query row l when:

s <= l + context_len - L

Skip a whole K/V tile when its first key is beyond the last visible key for the query block. This is both a correctness rule and an important causal-prefill optimization.

Task 3: Validate the Page Boundary

Complete the long-query selection in PagedAttention::eval_gpu in src/extensions/src/paged_attention.cpp, then test paged_attention_mma_bf16_d128 against the Day 4 kernels. Keep tiny_llm_ext::paged_attention and the Python paged_attention signature unchanged.

Use the GPU-debugging ladder from Week 2 Day 3:

1. compare Day 4 page-walking attention with the readable equation written with mlx.core;
2. compare paged FlashAttention with the Day 4 path;
3. trace the model-to-kernel route, then design a matched operator benchmark before making a speed claim.

The behavior fixtures cover:

• a context contained in one page;
• a partial query block and context tail with non-consecutive physical pages;
• poisoned unused pages and tail slots, which must not affect the result;
• batched GQA where multiple query heads map to one K/V head;
• causal and non-causal calls, including an explicit scale;
• short decode and generic BF16 head-dimension fallbacks;
• a model-level long-prefill comparison through the public Week 3 path;
• output shape and dtype, plus numerical agreement with dense attention.

These are public input/output checks. They do not inspect allocator choices, page identifiers chosen by an implementation, kernel names, or routing state.

Force mx.eval immediately after each operator so compilation, dispatch, and addressing failures are reported at the responsible call.

pdm run test --week 3 --day 5

If this command is green before you begin Day 5, you have confirmed that Day 4’s correctness boundary is intact. Continue with the optimized implementation and use the trace below to confirm that the supported hot case no longer takes the scalar fallback.

Task 4: Integrate and Measure

Verify the existing dispatch in Qwen3MultiHeadAttention.__call__ and the shape selection inside PagedAttention::eval_gpu. Task 4 adds no new extension function.

The Week 3 model should use the tiled paged path automatically for supported long prefills. Short queries continue through the vector paged-decode schedule. Neither path gathers a dense K/V tensor. Canonical Week 3 uses MLX quantized projections, but its cache, paged attention, batching, and scheduling remain course-owned. This hybrid course path is not the full-MLX baseline.

The checked continuous-serving trace is useful system context, but it is not a matched scalar-versus-tiled Day 5 experiment. The current model-free runner also lacks that isolated long-BF16 comparison. Do not claim a Day 5 speedup from these rows. A separate benchmark follow-up should hold inputs, page tables, precision, warmup, synchronization, and every non-attention mechanism fixed while changing only the scalar-versus-tiled schedule.

Run the cumulative system check on your solution:

pdm run bench-serving-progression --solution tiny_llm --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/task367-final-main/raw/learner-week3-serving.json

The table below is checked reference evidence. Reproduce it separately with --solution ref and --json-output benchmark_results/task367-final-main/raw/week3-serving-final-main.json.

FlashAttention is expected to matter more as prefill grows. Treat that as a hypothesis until a matched operator benchmark measures it. It should not replace the Day 4 decode schedule: a one-token query has no query-tile reuse.

On the checked M4 Pro trace, all three course rows share the same projection seam:

Relative to dense serving, direct paging is 5.4% lower on prefill, 31.6% higher on output/request throughput, 82.3% higher on decode, 47.4% lower on measured peak KV storage, and 99.76% lower on avoidable logical copy volume. The compatibility row’s page-pool counter excludes its temporary dense staging allocation, so it is not a total peak. These are cumulative Week 3 system results; they do not isolate the Day 5 prefill schedule from paging, direct decode, allocation, or scheduling, and they do not credit the MLX projection seam to paged attention.

The existing 8K static sweep is another cumulative diagnostic. It does not isolate scalar versus tiled paged attention, measure request turnover, page reuse, or capacity. The performance appendix records the matched serving and long-context measurements. Long-context decode remains a Day 4 vector kernel workload; do not credit a prefill schedule with a decode gain.

pdm run bench-course-progression --solution tiny_llm --offline --suite course \
  --variant week2 --variant week3 --variant mlx --model qwen3-4b \
  --input-len 8192 --output-len 2 --prefill-logits last \
  --warmup 1 --repeats 4 --cooldown-seconds 1 \
  --json-output benchmark_results/task367-final-main/raw/learner-week3-8k.json

This command measures your Week 2 and Week 3 course rows while retaining MLX as the library baseline. The checked table below is reference evidence; reproduce it separately with --solution ref and --json-output benchmark_results/task367-final-main/raw/week3-8k-final-main.json.

The cumulative Week 3 row is 43.1% faster than the cumulative Week 2 row and reaches 72.5% of full MLX at this shape. This remains a static diagnostic: it does not measure request turnover, page reuse, admission capacity, the projection seam, or the Day 5 schedule causally. Its decode row is the Day 4 vector schedule, not evidence for the tiled prefill kernel. Full method and raw samples are in benchmark_results/task367-final-main/task367-final-main-benchmark-ledger.md.

Trace the implementation, not its names. Copy the prompt below into GPT or Claude after your behavior tests pass. This is learner guidance and is not part of automated grading.

Trace the actual control flow for Week 3 long-prefill paged attention. First determine whether the implementation stays in this repository or crosses into an external low-level library. For a repository implementation, start at the Python model call and follow native dispatch into Metal. For an external implementation, trace from the Python model call to the library boundary and record equivalent boundary evidence, including the call site, arguments, and selected backend. Determine whether BF16 queries with L > 8 and D == 128 reach the learner's optimized tiled or FlashAttention-equivalent implementation rather than the correctness fallback. Verify that short decode and unsupported or generic shapes retain correct fallback behavior. Cite file and line evidence for repository code and equivalent boundary evidence for external code, flag dead or unreachable paths, judge control flow rather than function or kernel names, and report findings only—do not modify files.

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.