Learn LLM Serving

This course is designed for systems engineers who want to understand how large language models (LLMs) work.

As a systems engineer, I am always curious about how things work internally and how to optimize them. I found it difficult to understand LLM inference because most open-source serving projects are highly optimized with CUDA kernels and other low-level techniques. It is hard to see the whole picture in a codebase with hundreds of thousands of lines. I therefore decided to implement an LLM serving project from scratch using only array and matrix operations. The goal was to understand what it takes to load an LLM’s parameters and perform the mathematical operations that generate text.

You can think of this course as an LLM counterpart to the Needle project from CMU’s Deep Learning Systems course.

Prerequisites

You should understand the basics of deep learning and be familiar with PyTorch. We recommend the following resources:

CMU Introduction to Machine Learning — covers the fundamentals of machine learning.
CMU Deep Learning Systems — teaches you how to build a framework like PyTorch from scratch.

Environment Setup

This course uses MLX, an array and machine learning framework for Apple silicon. For many learners, an Apple silicon device is easier to access than an NVIDIA GPU. In principle, you could also complete the course with PyTorch or NumPy, but the test infrastructure does not support them as implementation backends. Instead, the tests compare your implementation with trusted MLX operations and model implementations to verify correctness.

Course Structure

This course is divided into four weeks. We will serve Qwen3 MLX models, optimize the serving path, and use it to build a small coding agent.

Week 1: Serve Qwen3 using array and matrix operations written in Python.
Week 2: Measure the cached model, implement the selected C++ and Metal kernels, and re-profile after each change.
Week 3: Add further optimizations and batch requests for high-throughput serving.
Week 4: Reuse the serving stack in a local coding agent with tools, sessions, and evaluation.

Course Roadmap: What Depends on What

The course supports two different goals: implementing the cumulative serving stack, or studying and running a later checkpoint without completing all earlier exercises. These are not the same path.

Tiny-LLM roadmap. The cumulative interface and state path runs from Week 1 through the seven Week 2 days and Week 3 into Week 4. Week 2 Days 3 through 7 show optional MLX operator off-ramps that preserve the course interfaces; they are different from the full-MLX model baseline. Week 4 keeps the course prerequisite of setup plus Weeks 1 through 3, while its deterministic scripted-model tests for Days 1 through 7 can run after setup. Day 8 joins the scripted sequence to the real-model path, and Day 9 continues the Week 4 sequence.

On a narrow screen, scroll the roadmap horizontally; when the roadmap is focused, the left and right arrow keys move through it without changing chapters. Its labels stay at their readable desktop size.

Solid arrows in the diagram are interface and state prerequisites. They do not mean that you must hand-write every earlier optimization. A dashed border marks a custom operator that you may replace locally with its MLX equivalent while keeping the surrounding course interface. The reference and full-MLX lanes let you observe a completed system, but they do not fill in unfinished functions in src/tiny_llm.

Your goal	Start here	What earlier implementation is required?
Build the whole serving system	Week 1, then follow the solid arrows	Each week uses interfaces and mechanisms established by the previous week.
Skip a Week 2 kernel optimization	Keep that day’s course interface and wire the corresponding MLX operator at the seam	The earlier model, state, and interface work still needs to exist. This is a local code choice, not a CLI flag.
Read or experiment with a later week	Open that chapter and use `tiny_llm_ref`	None in your learner tree. Run the supplied reference tests or reference loader.
Compare with the production-library baseline	Use `--solution mlx`	None, but this runs the full MLX model and bypasses the course implementation.
Run the Week 4 Days 1–7 deterministic tests before finishing the serving stack	After setup, run the supplied scripted-model tests	The tests do not need a working serving implementation. The course still assumes setup plus Weeks 1–3 before Week 4; follow the Week 4 days in order, and Day 8’s real-model bridge needs the Week 3 model/tokenizer/KV-cache boundary.

The cumulative dependencies are deliberate:

Week 1 → Week 2: Week 2 starts from the readable Qwen3 model and replaces costs one measured mechanism at a time: first the generation algorithm and KV cache, then quantized and fused kernels. Days 1–2 establish state and a repeatable measurement; Days 3–5 follow the dominant cost, Day 6 is an optional operator lab, and Day 7 closes with a conditional schedule decision.
Week 2 → Week 3: Week 3 selects MLX quantized projections, but it keeps course-owned normalization, activation, cache, attention, paging, batching, and scheduling. This is an explicit operator seam, not “use the MLX model for Week 2.”
Week 3 → Week 4: Week 4 remains the next cumulative course week. Its Days 1–7 tests can exercise control flow with deterministic scripted models after setup, even before the serving stack works. Day 8 reconnects that harness to the real tokenizer and KV cache, so that checkpoint needs a working Week 3 path.

Is Week 2 required for Week 3? The Week 2 interfaces are; every Week 2 optimization is not. The current Week 3 starter reuses the Week 2 model shell, dense-cache contract, packed-weight plumbing, normalization, activation, attention, and matrix-fragment interfaces. You may preserve those interfaces and substitute MLX operators for custom optimization work, but starting Week 3 is not as simple as selecting --solution mlx. That flag selects the complete MLX model and bypasses the course-owned paging, batching, attention, and scheduler surfaces that Week 3 teaches. Skipping the entire Week 2 implementation would require a supplied hybrid starting checkpoint; that checkpoint does not exist today.

Week 2 operator off-ramps

Week 2 separates the mechanism you need later from the kernel you are invited to optimize. If your goal is to continue into Week 3 rather than implement every Metal kernel, you can make these explicit local substitutions:

Week 2 day	Keep in the course stack	Optional MLX substitution
Days 1–2	Dense KV-cache state, the Week 2 model boundary, and the matched measurement method	None; these are state and methodology rather than replaceable operators.
Day 3	Packed-weight containers, quantized embedding/model wiring, and the `quantized_linear` interface	Route projections through `mx.quantized_matmul` instead of the custom matrix-vector kernel.
Day 4	The Week 2 norm, position, and activation call sites	Use the corresponding MLX RMSNorm/RoPE operators and an MLX SiLU-based SwiGLU composition instead of the custom fused kernels.
Day 5	The quantized-projection interface and matrix-shaped dispatch boundary	Keep using the Day 3 MLX projection seam instead of implementing the SIMD-matrix schedule.
Day 6 (optional)	The dense-cache attention interface and its shape/mask adapter	Use `mx.fast.scaled_dot_product_attention` instead of the supplied bounded decode-attention branch.
Day 7	The Day 5 unsplit projection fallback and measured dispatch boundary	Keep the unsplit path rather than implementing Split-K where your measurement does not support it.

Only the quantized-projection seam is already selected by canonical Week 3. The Day 4 and optional Day 6 alternatives require you to wire the MLX call at the existing course interface; there is no --use-mlx-for-day command. These off-ramps let you study later mechanisms, but they do not complete the skipped day’s custom-kernel exercises, implementation-specific tests, or performance claims.

To run a completed checkpoint without solving it first:

# Build the supplied reference extension once after setup or a clean checkout.
pdm run build-ext-ref

# Run one supplied reference test group.
pdm run test-refsol --week 3 --day 1

# Run a completed course model.
pdm run main --solution ref --loader week3

# Run the separate full-MLX baseline.
pdm run main --solution mlx

--solution ref runs the supplied implementation end to end. --solution mlx runs MLX end to end. Neither command composes “earlier weeks from the reference or MLX, this week’s TODOs from my learner tree.” Per-operator substitution is a manual code edit that preserves the course interface; it is not a third solution mode. If you want to implement a later week in src/tiny_llm, its earlier interface and state prerequisites must already work; the repository does not currently provide a one-command hybrid checkpoint.

Choose a Model for Your Mac

The table below is a conservative starting point for recent Apple-silicon Mac mini and MacBook unified-memory sizes up to 64 GB. Across those machines, the available tiers are 8, 16, 18, 24, 32, 36, 48, and 64 GB.¹ Each entry is recommended / maximum for that week’s course path. The recommendation is the checkpoint to use while completing the exercises; the maximum is the largest course-supported checkpoint worth trying with short prompts and the chapter’s default batch settings.

Unified memory	Week 1	Week 2	Week 3	Week 4
8 GB	0.6B / 0.6B	0.6B / 1.7B²	0.6B / 1.7B	0.6B / 1.7B
16 GB	0.6B / 1.7B	4B / 8B²	4B / 8B	4B / 8B
18 GB	0.6B / 1.7B	4B / 8B²	4B / 8B	4B / 8B
24 GB	0.6B / 1.7B	4B / 8B²	4B / 8B	4B / 8B
32 GB	4B / 8B	4B / 8B	4B / 30B-A3B³	4B / 30B-A3B³
36 GB	4B / 8B	4B / 8B	4B / 30B-A3B³	4B / 30B-A3B³
48 GB	4B / 8B	4B / 8B	4B / 30B-A3B³	4B / 30B-A3B³
64 GB	4B / 8B	4B / 8B	4B / 30B-A3B³	4B / 30B-A3B³

Week 1 reads an official 4-bit checkpoint but materializes its linear and embedding weights in BF16. On an 8 GB Mac, keep the required path at 0.6B. On a 16–24 GB Mac, use 0.6B for the required work and treat 1.7B as an upper-end experiment. Week 2 Days 1–2 retain that dense BF16 model; Day 3 keeps weights packed for the quantized-matvec checkpoint. Weeks 3 and 4 inherit that packed path. More memory still helps after reaching the largest supported model because prompt length, batch size, KV caches, compilation, macOS, and other applications all share the same pool. These ceilings are therefore planning guidance, not a guarantee that every workload will avoid memory pressure.

How to Use This Book

The tiny-llm book is a hands-on guide rather than a textbook that explains every concept from first principles. We link to the resources that the authors found useful while implementing the project instead of repeating their explanations. Each chapter provides a sequence of tasks, supporting readings, and implementation hints.

The book also standardizes terminology and notation across those resources so that they map cleanly to the codebase. For example, we use consistent symbols for tensor dimensions and explain what H, L, and E mean at the point of use.

About the Authors

This course is created by Chi and Connor.

Chi is a systems software engineer at Neon (now acquired by Databricks), focusing on storage systems. Fascinated by large language models, he created this course to explore how LLM inference works.

Connor is a software engineer at PingCAP, developing the TiKV distributed key-value database. Curious about the internals of LLMs, he joined the project to practice building a high-performance LLM serving system from scratch and helped develop the course for the community.

Community

You can join skyzh’s Discord server to study with the tiny-llm community.

Get Started

Follow the instructions in Setting Up the Environment, then begin building tiny-llm.

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.

Apple lists these tiers across the M2 Mac mini, M3 Pro and M3 Max MacBook Pro, M4 Mac mini, and M5 MacBook Air specifications. Higher-memory configurations are outside this table. ↩
Week 2 Days 1–2 use the dense Week 1 loader, so keep using the Week 1 recommendation until the packed quantized-matvec path is complete on Day 3. The larger Week 2 entries apply after that checkpoint. ↩ ↩2 ↩3 ↩4
30B-A3B requires the optional Week 3 MoE implementation. In Week 4, select the Week 3 loader. Use batch size one and a short context when approaching this ceiling; 4B remains the required-course target. ↩ ↩2 ↩3 ↩4 ↩5 ↩6 ↩7 ↩8

Setting Up the Environment

To follow this course, you need a Mac with Apple silicon. The project uses PDM for dependency and environment management.

Install PDM

Follow the official installation guide to install PDM.

Clone the Repository

git clone https://github.com/skyzh/tiny-llm

The repository is organized as follows:

src/tiny_llm/ -- your implementation
src/tiny_llm_ref/ -- the reference implementation
tests/ -- unit tests for your implementation
tests_refsol/ -- unit tests for the reference implementation
book/ -- the book source

Reference implementations are available if you get stuck during the course.

Install Dependencies

cd tiny-llm
# This creates a virtual environment and installs all dependencies.
pdm install -v

Check the Installation

pdm run check-installation
# The reference solution should pass all Week 1 tests.
pdm run test-refsol -- -- -k week_1

Run Unit Tests

Your code is in src/tiny_llm. You can run the unit tests with:

pdm run test

Download the Model Parameters

We use the official 4-bit Qwen3 MLX model files. The default model is Qwen/Qwen3-0.6B-MLX-4bit, which is small enough for the dequantized Python implementation in Week 1. If your device has more memory, you can also try larger Qwen3 models.

Follow the Hugging Face CLI guide to install the hf command-line tool.

The model parameters are hosted on Hugging Face. After authenticating the CLI with your credentials, download them with:

hf auth login
hf download Qwen/Qwen3-0.6B-MLX-4bit
# Optional larger models:
hf download Qwen/Qwen3-1.7B-MLX-4bit
hf download Qwen/Qwen3-4B-MLX-4bit

Then, you can run:

pdm run main --solution ref --loader week1

The command should load the reference model and print generated text.

In Week 2, we will write C++ and Metal kernels. The required additional tools are covered at the end of Week 1.

Week 1: From Matmul to Text

This week, we will start with basic array and matrix operations and use them to turn Qwen3 model parameters into a model that generates text. We will implement the neural network layers used by Qwen3 with MLX’s array APIs.

We will use Qwen/Qwen3-0.6B-MLX-4bit. The course model uses BF16 weights and activations by default, so start with the 0.6B model before trying larger Qwen3 models. The required model path runs on the GPU. Small operator fixtures may use a different dtype as a readable correctness reference; they do not define the model-storage dtype.

Numerically sensitive operations may promote arithmetic to FP32 and cast the result back to BF16. Week 1 favors readable array expressions, even when that means materializing an FP32 intermediate. Week 2 replaces those full-tensor promotions with kernels that keep model-sized storage in BF16 and accumulate in FP32 registers.

What We Will Cover

Attention, multi-head attention, grouped-query attention, and multi-query attention
Positional encodings and RoPE
Using mx.fast.rms_norm for Qwen3’s per-head Q/K normalization, then implementing RMSNorm ourselves
Implementing the MLP, combining the attention components, and building the complete Transformer model
Loading Qwen3 model parameters and generating text

What We Will Not Cover

To make the journey as interesting as possible, we will skip a few things for now:

Quantization and dequantization internals. These will be covered in Week 2. For now, we use a provided helper to dequantize the Qwen3 weights before passing them to our layer implementations.
Low-level implementations of operations such as softmax, exponentiation, and logarithms. These operations are simple enough that using the MLX versions does not detract from the learning objectives.
Tokenization. We use the mlx_lm tokenizer rather than implementing one from scratch.
Decoding model-weight files. We use mlx_lm to load the model, then transfer its weights into our layer implementations.

Basic Matrix APIs

MLX’s Python API is designed to be familiar to NumPy users. If you are new to array programming, start with NumPy: the absolute basics for beginners.

You can also refer to the MLX Operations API for more details.

Qwen3 Models

You can run Qwen3 with MLX or vLLM. The readings below provide context for what we will build. By the end of the week, you will be able to use Qwen3 as a causal language model to generate text.

Reference implementations of Qwen3 are available in Hugging Face Transformers, vLLM, and mlx-lm. Use them to explore the model’s internals and compare them with this week’s implementation.

📚 Readings

Week 1 Day 1: Attention and Multi-Head Attention

The starter provides softmax through MLX so that Day 1 can focus on the attention data flow. Your required work is to complete scaled_dot_product_attention_simple and SimpleMultiHeadAttention in src/tiny_llm/attention.py, plus the linear helper in src/tiny_llm/basics.py. Other attention functions in the starter are for later days and are not part of this chapter.

Start by running the focused Task 1 tests. The command refreshes the supplied Day 1 test in tests/ before running it:

pdm run test --week 1 --day 1 -- -k task_1

The supplied softmax cases pass in the untouched starter, while the attention cases fail. This expected red checkpoint shows the behavior that your attention implementation must add.

An attention layer processes an input sequence and weighs the relevance of its different positions when producing each output. Attention is a key building block of Transformer models.

📚 Reading: Transformer Architecture

We use Qwen3, a decoder-only model, for text generation. The model takes a sequence of token IDs, maps them to embeddings, and produces logits for the next token at each sequence position. The generation loop will later use the final position’s logits to choose the next token ID.

📚 Reading: LLM Inference, the Decode Phase

An attention layer takes a query, a key, and a value. In a basic implementation, all three have the same shape: N.. x L x D.

N.. represents zero or more batch dimensions. Within each batch, L is the sequence length and D is the embedding dimension for one attention head.

For example, a sequence of 1,024 tokens with a head dimension of 512 is represented by a tensor of shape N.. x 1024 x 512.

Task 1: Implement `scaled_dot_product_attention_simple`

In this task, we will implement scaled dot-product attention. We assume that the input tensors Q, K, and V have the same shape. Later chapters will introduce attention variants whose input shapes differ.

src/tiny_llm/attention.py

📚 Readings

Annotated Transformer
PyTorch Scaled Dot Product Attention API (assume enable_gqa=False, assume dim_k=dim_v=dim_q and H_k=H_v=H_q)
MLX Scaled Dot Product Attention API (assume dim_k=dim_v=dim_q and H_k=H_v=H_q)
Attention is All You Need

Implement scaled_dot_product_attention_simple using the formula below. The function takes query, key, and value tensors with the same shape, plus an optional additive mask M.

$Attention = softmax (\frac{Q K ^{T}}{d _{k}} + M) V$

Here, $\frac{1}{d _{k}}$ is the default scale factor. Callers may supply a different scale factor.

L is seq_len, in PyTorch API it's S (source len)
D is head_dim

key: N.. x L x D
value: N.. x L x D
query: N.. x L x D
output: N.. x L x D
scale = 1/sqrt(D) if not specified

Use the supplied softmax helper for the required exercise. As an optional, ungraded bonus, replace its MLX call with your own numerically stable implementation: subtract the maximum value along axis, exponentiate the shifted values, then divide by their sum along the same axis. Preserve the helper’s public API and output behavior.

When this function is called from multi-head attention, the tensors will usually have these shapes:

key: 1 x H x L x D
value: 1 x H x L x D
query: 1 x H x L x D
output: 1 x H x L x D
mask: 1 x H x L x L

The function itself operates on the last two dimensions and must support any number of leading batch dimensions. The mask only needs a shape that can broadcast to the attention-score shape.

Run the Task 1 checkpoint again after implementing the function:

pdm run test --week 1 --day 1 -- -k task_1

When this checkpoint turns green, your attention function supports arbitrary leading batch dimensions, optional masks, and default or explicit scaling.

Task 2: Implement `SimpleMultiHeadAttention`

In this task, we will implement the multi-head attention layer.

src/tiny_llm/attention.py

📚 Readings

Annotated Transformer
PyTorch MultiHeadAttention API (assume dim_k=dim_v=dim_q and H_k=H_v=H_q)
MLX MultiHeadAttention API (assume dim_k=dim_v=dim_q and H_k=H_v=H_q)
The Illustrated GPT-2 (Visualizing Transformer Language Models) helps you better understand what key, value, and query are.

Implement SimpleMultiHeadAttention. The layer projects batches of query, key, and value vectors with the Q, K, and V weight matrices, then passes the projections to the attention function from Task 1. Finally, it applies the output projection O.

First, implement the linear function in basics.py. It takes a tensor of shape N.. x I, a weight matrix of shape O x I, and an optional bias vector of shape O. Its output has shape N.. x O, where I is the input dimension and O is the output dimension.

Use the focused linear tests as your next checkpoint:

pdm run test --week 1 --day 1 -- -k test_task_2_linear

Before you implement linear, this checkpoint is red. When it turns green, linear supports optional bias across the tested precisions and devices.

For SimpleMultiHeadAttention, the input tensors query, key, and value have shape N x L x E, where E is the embedding dimension for one token. The Q, K, and V projections each map E to H x D: H heads, each with dimension D. Reshape that final projection dimension into separate H and D dimensions.

You now have a tensor of shape N x L x H x D for each projection. Before applying attention, transpose each one to N x H x L x D.

This treats each attention head as an independent batch, allowing attention to be calculated separately for each head across sequence dimension L.
Leaving H after L would cause the matrix multiplication to mix the head and sequence dimensions. Each head must attend only to token relationships within its own subspace.

The attention function produces one output per head. Transpose the result back to N x L x H x D, reshape it to N x L x (H x D), and apply the output projection.

E is hidden_size or embed_dim or dims or model_dim
H is num_heads
D is head_dim
L is seq_len, in PyTorch API it's S (source len)

w_q/w_k/w_v: (H x D) x E
output/input: N x L x E
w_o: E x (H x D)

Run the Task 2 checkpoint after implementing the layer:

pdm run test --week 1 --day 1 -- -k task_2

When this checkpoint turns green, your layer projects query, key, and value tensors into independent attention heads and recombines their outputs through the final projection.

You can run all tests for the day with:

pdm run test --week 1 --day 1

When the full Day 1 suite turns green, you have a standalone multi-head attention layer that projects Q/K/V, evaluates each head independently, and recombines the result. Day 3 will generalize this attention mechanism to grouped-query attention for Qwen3; the SimpleMultiHeadAttention layer built here is a standalone exercise, not the model’s exact call path.

Week 1 Day 2: Positional Encodings and RoPE

The Day 2 starter already declares RoPE(dims, seq_len, base=10000, traditional=False). Its constructor and call method are empty. You will fill in those two methods: cache one table of position-dependent angles, then use it to rotate the last dimension of an input shaped (N, L, H, D). Day 3 will apply the non-traditional form to Qwen3’s query and key heads before attention.

📚 Readings

Task 1: Implement Traditional Rotary Positional Encoding

You will need to modify the following file:

src/tiny_llm/positional_encoding.py

Start by building the frequency table in RoPE.__init__. Let M = D // 2. Pair index i, where 0 <= i < M, has the angular rate below; multiplying it by a token position gives the angle for that pair.

angular_rate[i] = base ** (-i / (D // 2))
angle[position, i] = position * angular_rate[i]

Use mlx.core operations such as arange, power, outer, cos, and sin to precompute the cosine and sine of those angles for every position from 0 through seq_len - 1. The two tables have shape (seq_len, M). Implement the operator yourself with these array operations; mx.fast.rope is the supplied test’s correctness oracle, not the implementation for this exercise. For this lesson, assume that D is even, so M pairs cover the whole head dimension.

If offset is not provided, apply positions 0 through L - 1 to the input sequence. Otherwise, select positions from the supplied slice. For example, with offset=slice(5, 10), the input sequence must have length 5, and its first token uses the frequency for position 5. Reshape the selected (L, M) basis to broadcast across the batch and head axes.

For Week 1, you only need to support offset=None and a single slice. We will implement list[slice] for continuous batching later. For now, assume that every item in a batch uses the same offset.

x: (N, L, H, D)
cos/sin_freqs: (MAX_SEQ_LEN, D // 2)

Traditional RoPE interprets adjacent values along head dimension D as complex-number pairs. If D = 8, then x[0] and x[1] form one pair, x[2] and x[3] form another, and so on. Both values in a pair use the same frequency from cos_freqs and sin_freqs.

output[0] = x[0] * cos_freqs[0] + x[1] * -sin_freqs[0]
output[1] = x[0] * sin_freqs[0] + x[1] * cos_freqs[0]
output[2] = x[2] * cos_freqs[1] + x[3] * -sin_freqs[1]
output[3] = x[2] * sin_freqs[1] + x[3] * cos_freqs[1]
...and so on

You can implement this operation by reshaping x to (N, L, H, D // 2, 2) and applying the formula to each pair. Stack the real and imaginary results back along the pair axis, restore the original shape, and return the result in x.dtype.

📚 Readings

Run the focused command once before editing. The empty starter returns no array, so the comparison should fail when it tries to inspect the result. Run the same command again after implementing Task 1; when it passes, your cached basis, position selection, adjacent pairing, and dtype restoration work together.

pdm run test --week 1 --day 2 -- -k task_1

Task 2: Implement Non-Traditional `RoPE`

Qwen3 uses a non-traditional arrangement of RoPE pairs. Keep the same cached frequencies and position-selection logic. When traditional is false, split the head dimension into two halves, then pair corresponding values from the halves. Let x1 = x[..., :HALF_DIM] and x2 = x[..., HALF_DIM:].

output[0] = x1[0] * cos_freqs[0] + x2[0] * -sin_freqs[0]
output[HALF_DIM] = x1[0] * sin_freqs[0] + x2[0] * cos_freqs[0]
output[1] = x1[1] * cos_freqs[1] + x2[1] * -sin_freqs[1]
output[HALF_DIM + 1] = x1[1] * sin_freqs[1] + x2[1] * cos_freqs[1]
...and so on

Implement this form by selecting the first and second halves of x directly, applying the rotations, concatenating the results, and returning the original dtype. The constructor’s default is non-traditional because that is the layout the Qwen3 attention block will use.

📚 Readings

vLLM implementation of RoPE

This focused command should now pass with the half-split layout while reusing the same angles and offset handling:

pdm run test --week 1 --day 2 -- -k task_2

Finally, run both layouts together:

pdm run test --week 1 --day 2

Once that command passes, RoPE is ready for Day 3 to rotate Qwen3 query and key heads with one shared slice. Per-request list[slice] offsets remain a later continuous-batching problem.

Week 1 Day 3: Grouped Query Attention (GQA)

On Day 3, we will implement grouped-query attention (GQA). Qwen3 uses GQA to reduce the computational and memory costs of the key (K) and value (V) projections. In multi-head attention (MHA), every query (Q) head has a corresponding K and V head. With GQA, groups of Q heads share K and V heads. Multi-query attention (MQA) is the special case in which every Q head shares a single K/V head pair.

Readings

Task 1: Implement `scaled_dot_product_attention_grouped`

You will need to modify the following file:

src/tiny_llm/attention.py

In this task, we will implement grouped scaled dot-product attention, which forms the core of GQA.

Implement scaled_dot_product_attention_grouped in src/tiny_llm/attention.py. It is similar to standard scaled dot-product attention, but it supports a number of query heads that is a multiple of the number of key/value heads.

The main process is the same as standard scaled dot-product attention. The difference is that K and V heads are shared across multiple Q heads. Instead of H_q separate K and V heads, there are H K and V heads, each shared by n_repeats = H_q // H query heads.

Reshape query, key, and value so that K and V can be broadcast to the query heads in their respective groups during the matrix multiplications.

Separate the H and n_repeats dimensions in query.
Add a dimension of size 1 for n_repeats in key and value so that they broadcast across each group.

Then perform scaled dot-product attention: matrix multiplication, scaling, optional masking, softmax, and a final matrix multiplication. Broadcasting handles the head sharing without materializing repeated K and V tensors.

Using broadcasting instead of repeating K and V is more efficient because it avoids creating copies of the same data.

Finally, reshape the result to the expected output shape.

N.. is zero or more dimensions for batches
H_q is the number of query heads
H is the number of key/value heads (H_q must be divisible by H)
L is the query sequence length
S is the key/value sequence length
D is the head dimension

query: N.. x H_q x L x D
key: N.. x H x S x D
value: N.. x H x S x D
mask: N.. x H_q x L x S
output: N.. x H_q x L x D

In addition to grouped heads, this function supports different query and key/value sequence lengths: Q uses length L, while K and V use length S. H_q = H gives ordinary multi-head attention, while H = 1 gives multi-query attention. The leading N.. batch shape may contain any number of positive-sized dimensions, and D is an independent head dimension rather than necessarily hidden_size // H_q. An array mask is additive and must be forwarded to the attention scores.

You can test your implementation by running the following command:

pdm run test --week 1 --day 3 -- -k task_1

Task 2: Causal Masking

Readings

Writing an LLM from scratch, part 9 – causal attention

In this task, we will add causal masking to grouped attention.

Causal masking prevents attention from reading future tokens. When mask is set to the string "causal", apply a causal mask.

The additive causal mask has shape (L, S), where L is the query sequence length and S is the key/value sequence length. Allowed positions contain 0, and masked positions contain -inf, and the returned mask must use the requested dtype. Causal attention requires S >= L; reject inputs with more query positions than key/value positions. When S is greater than L, shift the diagonal by S - L so that the queries correspond to the final L positions in the key/value sequence. For example, if L = 3 and S = 5, the mask is:

0   0   0   -inf -inf
0   0   0   0    -inf
0   0   0   0    0

Implement causal_mask in src/tiny_llm/attention.py, then use it in scaled_dot_product_attention_grouped. Note that our shifted diagonal for L != S differs from the default behavior of some attention APIs.

You can test your implementation by running the following command:

pdm run test --week 1 --day 3 -- -k task_2

Task 3: Qwen3 Grouped Query Attention

In this task, we will implement Qwen3’s grouped-query attention. Modify the following file:

src/tiny_llm/qwen3_week1.py

Qwen3MultiHeadAttention implements attention for Qwen3. Follow this pseudocode:

x: B, L, E
q = linear(x, wq) -> B, L, H_q, D
k = linear(x, wk) -> B, L, H, D
v = linear(x, wv) -> B, L, H, D
q = rms_norm(q, q_norm)
k = rms_norm(k, k_norm)
q = rope(q, offset=slice(0, L))
k = rope(k, offset=slice(0, L))
(transpose as needed)
x = scaled_dot_product_attention_grouped(q, k, v, scale, mask) -> B, H_q, L, D  # use float32
(transpose as needed)
x = linear(x, wo) -> B, L, E

Qwen3 attention has no Q/K/V projection biases. Apply RMSNorm to each Q and K head before applying non-traditional RoPE; the order matters when the normalization weights are nonuniform. We will implement the reusable RMSNorm layer on Day 4, so call mx.fast.rms_norm directly for q_norm and k_norm today.

The grouped attention arithmetic must run in float32: cast the projected, normalized, and rotated Q and K tensors and the projected V tensor to float32 before calling scaled_dot_product_attention_grouped. Cast its result back to the input model dtype before the final output projection. Forward None, "causal", or an additive array mask unchanged.

You can test your implementation by running the following command:

pdm run test --week 1 --day 3 -- -k task_3

At the end of the day, you should be able to pass all tests of this day:

pdm run test --week 1 --day 3

Week 1 Day 4: RMSNorm and the Multilayer Perceptron

On Day 4, we will implement two important components of the Qwen3 Transformer architecture: RMSNorm and the multilayer perceptron (MLP), also known as the feed-forward network. RMSNorm is a normalization technique with less computational overhead than traditional layer normalization. The MLP applies nonlinear transformations after the attention block.

Task 1: Implement `RMSNorm`

In this task, we will implement the RMSNorm layer.

src/tiny_llm/layer_norm.py

Day 3 used mx.fast.rms_norm directly so that the GQA chapter could stay focused on attention. This task implements the same normalization rule as a reusable layer. From this point on, the Transformer block, final model normalization, and Q/K normalization path can use your RMSNorm implementation.

📚 Readings

Root Mean Square Layer Normalization
Qwen3 layers implementation in mlx-lm (includes RMSNorm) - See RMSNorm.

RMSNorm is defined as:

$y = \frac{x}{mean ( x ^{2} ) + ϵ} \cdot weight$

where:

x is the input tensor.
weight is a learned scaling parameter.
epsilon (eps) is a small constant, such as 1e-5 or 1e-6, added for numerical stability.
mean(x^2) is the mean of the squared elements along the final dimension.

Apply normalization independently to each feature vector along the input’s final dimension. Cast the input to float32 for the normalization calculation, including the mean, to preserve precision when the original values use float16 or bfloat16. Cast the normalized value back to the input dtype before applying weight. This matches the low-precision path used by MLX’s fast RMSNorm kernels: normalization statistics are accumulated in float32, while the final scaling happens in the model dtype.

D is the embedding dimension.

x: N.. x D
weight: D
output: N.. x D

You can test your implementation by running:

pdm run test --week 1 --day 4 -- -k task_1

Task 2: Implement the MLP Block

In this task, we will implement the MLP block named Qwen3MLP.

src/tiny_llm/qwen3_week1.py

The original Transformer uses a simple position-wise feed-forward network (FFN) in each block. It consists of two linear transformations with a ReLU activation between them.

Modern Transformer architectures, including Qwen3, often use more advanced FFN variants. Qwen3 uses SwiGLU, a gated linear unit (GLU) variant.

A plain FFN can be abstracted as:

h = activation(W_up(x))
out = W_down(h)

A GLU keeps the same expand-then-project-back shape but adds another projection that gates the intermediate features before W_down. This gives the MLP a learned, input-dependent way to control which intermediate channels matter, rather than applying an activation only to the features produced by W_up.

SwiGLU is the GLU variant used by Qwen3:

u = W_up(x)
g = SiLU(W_gate(x))
out = W_down(g * u)

📚 Readings

SwiGLU combines a GLU with the SiLU (sigmoid linear unit) activation function:

A GLU gates one linear projection of the input with another, using element-wise multiplication to control which features pass through.
SiLU is a smooth, non-monotonic activation function. Unlike ReLU, it has no zero-gradient region across all negative inputs and can produce nonzero outputs for negative values.

First, implement silu in basics.py. It takes a tensor of shape N.. x I and returns a tensor with the same shape:

$SiLU (x) = x * sigmoid (x) = \frac{x}{1 + e ^{- x}}$ Compute the sigmoid part in a numerically stable way:

if x >= 0:
    sigmoid(x) = 1 / (1 + exp(-x))
else:
    sigmoid(x) = exp(x) / (1 + exp(x))

The negative branch is algebraically equivalent to the direct sigmoid formula, but it prevents exp(-x) from becoming exp(large positive) when x is a large negative value. In vector code, first compute z = exp(-abs(x)). Use z / (1 + z) for negative inputs and 1 / (1 + z) otherwise. Do not rewrite the negative branch as 1 - 1 / (1 + z): in low precision, the fraction can round to 1, and the subtraction then incorrectly produces zero.

Then implement Qwen3MLP. Qwen3’s MLP contains:

A gate projection ( $W_{g a t e}$ )
An up projection ( $W_{u p}$ )
SiLU applied to the gate projection’s output
An element-wise product of the activated gate output and the up-projection output
A final down projection ( $W_{d o w n}$ )

This can be expressed as:

$MLP (x) = W_{d o w n} (SiLU (W_{g a t e} (x)) ⊙ W_{u p} (x))$ where $⊙$ denotes element-wise multiplication. Qwen3’s MLP projections do not use biases.

N.. is zero or more dimensions for batches
E is hidden_size (embedding dimension of the model)
I is intermediate_size (dimension of the hidden layer in MLP)
L is the sequence length

input: N.. x L x E
w_gate: I x E
w_up: I x E
w_down: E x I
output: N.. x L x E

You can test your implementation by running:

pdm run test --week 1 --day 4 -- -k task_2

At the end of the day, you should be able to pass all tests of this day:

pdm run test --week 1 --day 4

Week 1 Day 5: The Qwen3 Model

On Day 5, we will combine the components from the previous chapters into the complete Qwen3 model.

Model-level tests require the corresponding model files. Start with the default 0.6B model; download the larger models only if you want to test them as well:

hf download Qwen/Qwen3-0.6B-MLX-4bit
# Optional larger models:
hf download Qwen/Qwen3-1.7B-MLX-4bit
hf download Qwen/Qwen3-4B-MLX-4bit

The deterministic checkpoint tests do not need downloaded model files. The default 0.6B model is required for the final Day 5 integration check; only checks for the optional larger models are skipped when those files are unavailable. A skip means that the model-backed check did not run, not that it passed.

Task 1: Implement `Qwen3TransformerBlock`

src/tiny_llm/qwen3_week1.py

📚 Readings

Qwen3 uses the following Transformer block structure:

  input
/ |
| input_layernorm (RMSNorm)
| |
| Qwen3MultiHeadAttention
\ |
  Add (residual)
/ |
| post_attention_layernorm (RMSNorm)
| |
| MLP
\ |
  Add (residual)
  |
output

Read head_dim independently from the model configuration rather than deriving it from hidden_size and the number of attention heads. Match the reference block’s output dtype.

Run the tests for this task with:

pdm run test --week 1 --day 5 -- -k task_1

Task 2: Implement `Embedding`

src/tiny_llm/embedding.py

📚 Readings

LLM Embeddings Explained: A Visual and Intuitive Guide

The embedding layer maps token IDs (integers) to vectors of length embedding_dim. In this task, you will implement that lookup operation.

Embedding::__call__
weight: vocab_size x embedding_dim
Input: N.. (tokens)
Output: N.. x embedding_dim (vectors)

This can be implemented with array indexing.

Token IDs may have any number of leading dimensions. The lookup appends embedding_dim to that shape and preserves the BF16 model-facing dtype.

When input and output embeddings are tied, Qwen3 also uses the embedding weight as a linear projection from hidden vectors back to vocabulary logits.

Embedding::as_linear
weight: vocab_size x embedding_dim
Input: N.. x embedding_dim
Output: N.. x vocab_size

The projection accepts the same arbitrary leading dimensions, replacing only the final embedding_dim axis with vocab_size and preserving BF16.

Run the tests for this task with:

# Deterministic lookup and projection checks always run. The downloaded model adds integration parity.
pdm run test --week 1 --day 5 -- -k task_2

Task 3: Implement `Qwen3ModelWeek1`

Now that we have built all the Qwen3 components, we can implement Qwen3ModelWeek1.

src/tiny_llm/qwen3_week1.py

You will not implement the process of reading model parameters from tensor files. Instead, load the model with mlx_lm, then transfer its parameters into our implementation. The Qwen3ModelWeek1 constructor therefore accepts an MLX model.

The Qwen3 model has the following layers:

input
| (tokens: N..)
Embedding
| (N.. x hidden_size); note that hidden_size == embedding_dim
Qwen3TransformerBlock
| (N.. x hidden_size)
Qwen3TransformerBlock
| (N.. x hidden_size)
...
|
RMSNorm 
| (N.. x hidden_size)
Embedding.as_linear OR linear (lm_head)
| (N.. x vocab_size)
output

Read the number of layers, hidden size, head dimension, and other configuration values from mlx_model.args, whose type is defined by ModelArgs. The loaded weights are available through mlx_model.model; use the Qwen3 implementation and model metadata to identify the corresponding layer names.

Instantiate every configured Transformer block in model order, run all of them, and apply the final RMSNorm before the vocabulary projection. head_dim, the query and key/value head counts, intermediate size, RMS epsilon, maximum positions, and RoPE theta are separate configuration fields; map each one directly.

By this point, you have implemented RMSNorm. Replace the temporary Day 3 calls to mx.fast.rms_norm with RMSNorm(head_dim, q_norm, eps=...) and RMSNorm(head_dim, k_norm, eps=...). They implement the same formula; the built-in calls existed only to keep the GQA chapter focused on attention.

Different Qwen3 model variants map hidden vectors back to vocabulary logits in different ways. Some tie the input and output embeddings and use Embedding.as_linear; others have a separate lm_head linear layer. Select the strategy with mlx_model.args.tie_word_embeddings: if it is True, use Embedding.as_linear; otherwise, load and use lm_head.

The model takes a sequence of token IDs and returns unnormalized logits for every sequence position. On Day 6, we will use the final position’s logits to select the next token and generate a response.

The MLX models used in this course have quantized weights. Dequantize each linear or embedding layer before loading it into tiny-llm by using the provided quantize.dequantize_linear function, then store the readable Week 1 weight as BF16. Model activations and layer outputs should remain BF16. A readable attention or normalization expression may compute in FP32 for stability, but it must cast its model-facing result back to BF16.

Pass mask="causal" to every Transformer block. For a one-token sequence the mask has no effect; for longer sequences, it prevents each position from attending to future tokens.

Run the tests for this task with:

# The required 0.6B integration check exercises Task 3; download it if it is not already cached.
hf download Qwen/Qwen3-0.6B-MLX-4bit
pdm run test --week 1 --day 5 -- -k task_3

You may also download the optional 1.7B and 4B models from the commands at the top of this chapter. Their checks skip when the corresponding files are absent.

At the end of the day, you should be able to pass all tests of this day:

pdm run test --week 1 --day 5

Before treating this as a pass, confirm that the deterministic Task 2 checks ran and that the required 0.6B Task 3 integration check passed rather than skipped.

Week 1 Day 6: Generating the Response: Prefill and Decode

On Day 6, we will implement response generation for an LLM chatbot. The implementation is short, but it exercises much of the code from the previous days. Use this chapter to integrate and debug the complete Week 1 model.

Task 1: Implement `simple_generate`

src/tiny_llm/generate.py

simple_generate takes a model, tokenizer, prompt, and optional sampler, then streams the generated response to standard output. Its optional max_tokens argument defaults to 256 and limits the number of newly generated tokens. Generation has two phases: prefill and decode.

First, implement the nested _step function. It takes a one-dimensional array of token IDs, adds the batch dimension, and passes the result to the model. The model returns unnormalized logits over the vocabulary for every sequence position.

y: S (before adding a batch dimension)
model input: 1 x S
output_logits: 1 x S x vocab_size

You only need the last token’s logits to decide the next token. Therefore, you need to select the last token’s logits from the output logits.

logits = output_logits[:, -1, :]

You may normalize these logits into log probabilities with the log-sum-exp trick. This normalization does not change the result of argmax, but the sampler introduced on Day 7 expects log probabilities. If sampler is None, use mx.argmax along the final, vocabulary dimension. Otherwise, pass the (1, vocab_size) log-probability array to sampler, which returns one token ID in an integer array with shape (1,). Selecting the highest-scoring token at every step is called greedy decoding.

📚 The Log-Sum-Exp Trick
📚 Decoding Strategies in Large Language Models
📚 Tokenizer definition

With _step complete, implement the rest of simple_generate. Begin by encoding the prompt into a one-dimensional token array with tokenizer.encode(prompt, add_special_tokens=False). main.py has already formatted the chat prompt, so the tokenizer must not add another set of special tokens. If encoding produces no tokens, reject the prompt before calling the model.

Generate tokens in a loop until the model emits tokenizer.eos_token_id or max_tokens new tokens have been produced. Append each new token to the token array so that the next model call receives the complete sequence. An EOS token already inside the prompt is context and does not stop generation; only a newly generated EOS does.

Before the loop, bind one streaming detokenizer with detokenizer = tokenizer.detokenizer, then call detokenizer.reset() once. Keep and reuse that same stateful object throughout generation: feed every non-EOS output token to detokenizer.add_token(...) and print each detokenizer.last_segment as it becomes available. With the locked mlx-lm 0.31.3 dependency, each separate access to the tokenizer.detokenizer property creates a fresh streaming detokenizer, so repeatedly accessing the property would discard the buffered text. On either termination path, call detokenizer.finalize() on the saved object and print its final last_segment so buffered text is not lost. The function returns None after streaming the response.

An example of the sequences provided to the _step function is as below:

tokenized_prompt: [1, 2, 3, 4, 5, 6]
prefill: _step(model, [1, 2, 3, 4, 5, 6]) # returns 7
decode: _step(model, [1, 2, 3, 4, 5, 6, 7]) # returns 8
decode: _step(model, [1, 2, 3, 4, 5, 6, 7, 8]) # returns 9
...

In Week 2, we will accelerate decoding with a key-value cache so that the model does not recompute the entire sequence at every step.

First run the deterministic checkpoint, which does not download or load a model:

pdm run test --week 1 --day 6

Then complete the required product check with the default 0.6B model:

hf download Qwen/Qwen3-0.6B-MLX-4bit
pdm run main --solution tiny_llm --loader week1 --model qwen3-0.6b \
  --prompt "Give me a short introduction to large language model"

The first command downloads the model once; later runs use the cached copy. The product command should produce a reasonable explanation of large language models. Replace --solution tiny_llm with --solution ref to run the reference solution.

If downloaded, you can also try the larger models; these are optional demonstrations, not completion requirements:

pdm run main --solution tiny_llm --loader week1 --model qwen3-1.7b \
  --prompt "Give me a short introduction to large language model"
pdm run main --solution tiny_llm --loader week1 --model qwen3-4b \
  --prompt "Give me a short introduction to large language model"

Week 1 Day 7: Sampling and Preparing for Week 2

The Week 1 model can now produce normalized log probabilities and generate a greedy response. Today you will complete its single-request sampling path, run that path in the full Qwen3 loop, and prepare the tools needed for Week 2’s custom Metal extensions.

Task 1: Sampling

The starter already handles temp=0 with greedy decoding. You own the nonzero-temperature path: temperature, top-k, and top-p (nucleus) sampling.

src/tiny_llm/sampler.py

📚 mlx-lm sampler implementation

For one active request, the sampler receives normalized log probabilities with shape (1, vocab_size) and returns one token-ID array with shape (1,) and dtype uint32. The product does not pass multiple rows to this sampler.

When filtering, first copy the input. MLX indexed assignment mutates the array object, so masking the caller’s array directly would corrupt the log probabilities that the generation loop supplied.

Temperature Sampling

When temp=0, use greedy decoding. When temp is greater than 0, sample the next token from the log-probability distribution. A higher temperature flattens the distribution, making lower-probability tokens more likely and increasing output variety.

To implement temperature sampling, divide the log probabilities by the temperature and pass them to mx.random.categorical with axis=-1.

Top-k Sampling

Top-k sampling keeps only the k tokens with the highest log probabilities. Apply this filter before top-p and temperature scaling.

Use mx.argpartition to find the indices outside the top k, mask their log probabilities with -mx.inf, then apply temperature sampling.

Top-p (Nucleus) Sampling

Top-p sampling keeps the smallest high-probability set of tokens whose cumulative probability reaches or exceeds p. Apply this filter before temperature scaling.

One implementation uses mx.argsort to order the log probabilities from highest to lowest, applies exp to recover probabilities, and applies cumsum to compute cumulative probability. Keep a token when the cumulative probability before it is less than p; this includes the token that crosses the threshold. Mask the remaining log probabilities with -mx.inf.

The complete non-greedy order is:

copy the input;
apply top-k;
apply top-p, retaining the threshold-crossing token;
divide by temperature; and
draw with mx.random.categorical(..., axis=-1).

None or a non-positive value disables its corresponding filter. top_k == vocab_size and top_p == 1 leave the vocabulary unfiltered. The existing implementation raises ValueError when top_k is larger than the vocabulary; inputs outside these boundaries are not part of this lesson’s contract.

Run the deterministic, no-download checkpoint first:

pdm run copy-test --week 1 --day 7
pdm run test --week 1 --day 7 -- -q --tb=short

The distribution cases observe returned-token support and frequencies; the seeded case checks same-seed repeatability and different-seed divergence. An unseeded draw should not be expected to return one particular token.

After the focused checkpoint passes, verify the completed sampler in the full Week 1 product loop:

pdm run main --solution tiny_llm --loader week1 --model qwen3-0.6b --sampler-temp 0.5
pdm run main --solution tiny_llm --loader week1 --model qwen3-0.6b --sampler-temp 0.5 --sampler-top-k 10
pdm run main --solution tiny_llm --loader week1 --model qwen3-0.6b --sampler-temp 0.5 --sampler-top-p 0.9

These commands require the cached 0.6B model. If it is missing, download it and rerun them:

hf download Qwen/Qwen3-0.6B-MLX-4bit

Larger models remain optional and are not a Day 7 completion gate.

Task 2: Prepare for Week 2

Week 2 Days 1 and 2 introduce KV caching in Python, so you can begin them before the custom-extension toolchain is ready. Starting on Day 3, the C++ and Metal work requires full Xcode, its command-line tools, the Metal compiler, and CMake 3.27 or newer.

Install Xcode:

Install full Xcode from the Mac App Store or Apple Developer downloads. Full Xcode bundles its command-line tools; the standalone package described in Installing the command-line tools is an alternative, so you do not normally need to run xcode-select --install after installing Xcode.
Launch Xcode and Install Components:

After installation, launch Xcode at least once. It may prompt you to install additional macOS components; please do so (this is usually the default option).
Verify the Active Xcode Path:

Check which developer directory the command-line tools use:
```
xcode-select --print-path
```
If it does not point to full Xcode, switch it as described in Configuring command-line tools settings:
```
sudo xcode-select --switch /Applications/Xcode.app/Contents/Developer
```
Adjust the path if Xcode is installed elsewhere.
Resolve First-Launch or License Prompts if Needed:

Launch Xcode once. Only if the tools report an incomplete first launch or license problem, follow the reported recovery step, such as:
```
sudo xcodebuild -runFirstLaunch
sudo xcodebuild -license accept
```
Verify the Metal Compiler:
```
xcrun metal --version
```
With Xcode 26, the Metal toolchain may be a separate component. If the command reports that it is missing, use the conditional component workflow described in Downloading and installing additional Xcode components, then verify the compiler again:
```
xcodebuild -downloadComponent MetalToolchain
xcrun metal --version
```
Install and Verify CMake 3.27 or Newer:
```
brew install cmake
cmake --version
```

(This instruction is graciously provided by Liu Jinyi.)

Test the installation by compiling the code in src/extensions, which contains an axpby function adapted from the official MLX extension tutorial:

pdm run build-ext
pdm run build-ext-test

It should print c correct: True. The other exported extension names are fail-closed starter stubs labeled with the Week 2 or Week 3 checkpoint that implements them; this setup check calls only axpby.

If you are new to C++ or Metal, try a few small exercises before the custom kernel work on Day 3. For example, implement element-wise operations such as exp, sin, and cos, then use them in place of the corresponding MLX operations in your model implementation.

That completes Week 1: you now have a single-request Python inference loop that loads Qwen3, computes logits, samples tokens, and streams a response. Week 2 first adds KV caching in Python, then begins the custom Metal kernel path.

🚧 Week 2: A Step Closer to vLLM

Week 1 leaves you with a readable Qwen3 model that can generate text. This week, you will turn it into a measured single-request serving path. After each change, you will rerun the same synchronized workload, find where time now goes, and use that evidence to choose what to change next. Instead of collecting unrelated kernels, you will build an optimization story you can explain.

Days 1–5 form the main route. Day 6 is an optional lab for an operator that matters to a workload you choose. On Day 7, you will test Split-K where a short-shape measurement suggests it may help, then make a final keep-or-reject decision on the fixed product workload.

Begin with Day 1: KV Cache, where you will stop recomputing the entire prefix for every generated token.

⏱️ Time commitment. Days 3–5 introduce custom Metal kernels and may take substantially longer than Week 1 Days 6–7. You may skip Day 6. On Day 7, your Split-K implementation must be correct, but the course does not require an absolute speed or a universal crossover point.

Week 2 keeps BF16 for dense weights, quantization scales and biases, activations, projections, KV-cache entries, and model-facing kernel outputs. Packed W4 weight codes use uint32. Reductions, dot products, and online-softmax state accumulate in FP32 before returning BF16. Week 3 inherits these interfaces and precision boundaries.

Measure, Change, and Measure Again

Use the same four-part loop at each checkpoint:

Check correctness. Run the focused supplied test before timing.
Describe the workload. Record model, checkpoint, phase, token counts, prefill-logit mode, warmups, iterations, software, and device.
Find the next cost. Name the dominant operator category and make one bounded hypothesis before editing.
Rerun and decide. Repeat the identical product and attribution workload, then record keep, reject, or inconclusive, along with evidence that would change your conclusion.

The checked example makes that loop visible. Cached decode begins with projection work dominant; packed W4 exposes the pointwise category; fused model kernels leave projections as the next target; and the SIMD schedule shrinks prefill projection time. Split-K helps the measured 32-token shape but does not improve the fixed 128-token product control.

The decisions below use different metrics and denominators, so read each card as a bounded comparison rather than adding the percentages together. Packed W4, fused pointwise kernels, and SIMD prefill are kept for their measured controls; decode attention stays optional and inconclusive; Split-K is conditional at 32 tokens and rejected for the fixed 128-token workload.

Decision cards for one M4 Pro, macOS 27, Qwen3-4B, MLX 0.32.0 observation: keep packed W4 decode, fused pointwise kernels, and SIMD-matrix prefill; treat decode attention as optional and inconclusive; retain Split-K conditionally at 32 tokens but reject it for the fixed 128-token product workload.

You can complete this loop with the synchronized benchmark and portable attribution runner. Apple GPU capture and gpudebug appear only in the optional macOS 27 lab; neither is a prerequisite.

Daily Checkpoints

KV cache: make decode incremental and compare it with the Week 1 model on a matched workload.
Discover: learn to synchronize a measurement, attribute the cached model, and choose one bounded optimization. In the checked run, dense projections became the next target.
Packed W4 matvec: keep weights packed while you optimize decode projections, then re-profile. In the checked run, normalization, position, and activation work became visible next.
Fused model kernels: implement RMSNorm, RoPE, and SwiGLU one at a time. Keep each change only after a matched measurement, then re-profile prefill before choosing Day 5.
SIMD-matrix prefill: replace the matrix-shaped projection schedule chosen from the fixed 128-token prefill attribution.
Optional operator lab: choose a secondary operator category for one explicit workload. The supplied branch studies bounded decode attention, but an equivalent evidence-led operator experiment is also valid.
Conditional Split-K and final decision: try an under-filled 32-token projection while preserving the unsplit Day 5 fallback. Finish by rerunning the fixed 128×129 product workload and deciding whether to keep the change.

What Is Supplied and What You Own

The starter gives you model loading, the extension build system, benchmark and attribution runners, correctness tests, Python reference equations, stable checkpoint interfaces, and a compact checked M4 Pro evidence file. You will build the cache transition, integrate packed weights, implement the custom operators, and turn each measurement into a decision record.

The completed course path uses your implementations, not MLX replacements, for the operators it asks you to build. If you want to reach the later serving mechanisms without implementing one custom kernel, keep the course interface and connect the corresponding MLX operator locally. This substitution stays inside your course model. It is different from --solution mlx, which runs the separate full-MLX model.

Check Your Progress

Run the canonical selector after each day:

Course day	Test command
Day 1	`pdm run test --week 2 --day 1`
Day 2	`pdm run test --week 2 --day 2`
Day 3	`pdm run test --week 2 --day 3`
Day 4	`pdm run test --week 2 --day 4`
Day 5	`pdm run test --week 2 --day 5`
Day 6 (optional)	`pdm run test --week 2 --day 6`
Day 7	`pdm run test --week 2 --day 7`

When a command runs a model, benchmark, profile, capture, or reducer, pass --solution tiny_llm exactly as the chapter shows. Some command-line tools otherwise default to the completed reference, so omitting it may measure code you did not write.

Bring Forward Work from the Earlier Day Order

An earlier course order put decode attention on Day 5 and SIMD-matrix prefill on Day 6. If your checkout contains work from that order, you can keep it. First complete the current Day 5 SIMD gate, then use the optional Day 6 gate to check your retained attention implementation. The ordinary --week 2 --day 5 and --week 2 --day 6 commands above are the only selectors you need. Old Day 5 and Day 6 bookmarks now redirect to the corresponding canonical lessons.

What the Gates Check

Most required gates check public behavior: checkpoint and workload identity, operator results, fallbacks, synchronized output, and the decision-record schema. You may organize most internals differently. Course-ownership and extension-integration witnesses intentionally preserve explicit source and header seams. The gates do not grade device timings or require the optional gpudebug tooling.

Read the checked example as one machine’s optimization story, not a portable speed claim. Its absolute measurements come from one M4 Pro running macOS 27 with Qwen3-4B, a fixed 128-token prompt and 129-output-token product control, and n=2 balanced product samples. Six of eight captures exposed full shader/counter detail. The pre-SIMD prefill capture did not expose a shader ranking, while the Split-K capture exposed only static dispatch. The example marks the missing data unavailable instead of guessing.

Continue to Week 3

By the end of Week 2, your model decodes one token at a time from a dense KV cache, chooses separate prefill and decode projection schedules, and keeps its weights quantized. Week 3 keeps these model, cache, precision, and operator interfaces while adding paging and batching. You do not need the optional Day 6 attention branch to continue, and Day 7 begins from Day 5’s unsplit SIMD path.

The performance evidence ledger shows the checked causal example and its limits.

🚧 Week 2 Day 1: KV Cache

Your Week 1 Qwen model already generates by rerunning the full prefix. Day 1 keeps that path intact while you complete four separate Week 2 shells:

src/tiny_llm/kv_cache.py::TinyKvFullCache stores one layer’s dense K/V;
src/tiny_llm/qwen3_week2.py::Qwen3ModelWeek2 threads cache state and offsets through the model;
Qwen3ModelWeek2.create_kv_cache creates one cache per layer and request;
src/tiny_llm/generate.py prefills once, then sends only the new token.

Together, these pieces make prefill populate the cache and make decode send only the new token. The starter already supplies the Week 1 operators and the model-loading boundary. Start with the focused learner gate:

pdm run test --week 2 --day 1

When it passes, run the kv-cache checkpoint shown in Task 4. That live call puts the cache into the generation loop instead of exercising it only as an isolated data structure.

Each attention layer can then reuse the keys and values from previous tokens instead of recomputing the entire prefix at every step.

This is the foundation of Week 2 decode optimization. Week 3 will change how the cache is stored and shared, but the reuse starts here. Without it, every generated token reruns all model layers over an ever-growing prefix and can overwhelm gains from faster individual kernels.

📚 Readings

KV Caching Explained: Optimizing Transformer Inference Efficiency

First, make the repeated work concrete. Week 1 supplied the full sequence to the model on every step:

tokenized_prompt: [1, 2, 3, 4, 5, 6]
prefill: _step(model, [1, 2, 3, 4, 5, 6]) # returns 7
decode:  _step(model, [1, 2, 3, 4, 5, 6, 7]) # returns 8
decode:  _step(model, [1, 2, 3, 4, 5, 6, 7, 8]) # returns 9
...

x: B, L, E
q = linear(x, wq) -> B, L, H_q, D
k = linear(x, wk) -> B, L, H, D
v = linear(x, wv) -> B, L, H, D
q = rms_norm(q, q_norm)
k = rms_norm(k, k_norm)
q = rope(q, offset=slice(offset, offset + L))
k = rope(k, offset=slice(offset, offset + L))
(transpose as needed)
x = scaled_dot_product_attention_grouped(q, k, v, scale, mask) -> B, L, H_q, D
# q/k/v and the returned model tensor are BF16; the Python `mlx.core` expression may use FP32 intermediates
(transpose as needed)
x = linear(x, wo) -> B, L, E

The attention mechanism is computed as:

$Attention (Q, K, V) = softmax (\frac{Q K ^{T}}{d _{k}} + M) V$

Consider two consecutive decoding steps with L = S = 3 and L = S = 4. Assume that each attention head has dimension D = 4:

L = 3
Q        x  K^T     =         
1 1 1 1     1 2 3      1x1  -inf -inf
2 2 2 2     1 2 3      2x1  2x2  -inf
3 3 3 3     1 2 3      3x1  3x2  3x3
            1 2 3

L = 4
Q        x  K^T       =
1 1 1 1     1 2 3 4      1x1  -inf -inf -inf
2 2 2 2     1 2 3 4      2x1  2x2  -inf -inf
3 3 3 3     1 2 3 4      3x1  3x2  3x3  -inf
4 4 4 4     1 2 3 4      4x1  4x2  4x3  4x4

The leading 3 x 3 block of QK^T is identical in both steps. The causal mask prevents earlier queries from attending to the new token, so those outputs do not change either. Only the new query row can produce a new output; recomputing the earlier rows, softmax values, and products with V is wasted work.

Instead, cache the previous keys and values and compute only the projections for incoming tokens:

K in cache:
1 1 1 1
2 2 2 2

[a b c d] represent cached values

L = 1, S = 3
Q        x  K^T       =         
            (⬇️ is K not transposed)
            [1 1 1 1]      
            [2 2 2 2]      
3 3 3 3      3 3 3 3      3x1 3x2 3x3

L = 1, S = 4
Q        x  K^T       = 
            (⬇️ is K not transposed)
            [1 1 1 1]      
            [2 2 2 2]      
            [3 3 3 3]
4 4 4 4      4 4 4 4      4x1 4x2 4x3 4x4

Task 1: Implement the Key-Value Cache

src/tiny_llm/kv_cache.py

Each Transformer layer owns a key-value cache. Its update_and_fetch method:

Accepts the newly computed K and V for the incoming tokens.
Appends them along the sequence dimension.
Returns the complete cached K and V, the updated offset, and the mask.

For now, pass mask through unchanged and leave mask_length unused. Week 3 will use both when requests share a batch.

You may implement this in kv_cache.py as TinyKvFullCache:

L_new = number of incoming tokens

update_and_fetch(key, value, mask_length, mask) -> key, value, offset, mask

key:   B, H, L_new, D
value: B, H, L_new, D

if self.key_values is None:
    self.key_values = (key, value)
else:
    cached_key, cached_value = self.key_values
    self.key_values = (
        concat(cached_key, key, axis=2),
        concat(cached_value, value, axis=2),
    )

self.offset += L_new
key, value = self.key_values  # B, H, offset, D

return key, value, self.offset, mask

Keep this first cache deliberately simple and dense. Each mx.concat allocates a larger buffer and copies the previous K/V contents. Across a token-by-token decode of length S, those copies add up to O(S²) bytes even though the cache avoids O(S²) prefix recomputation. The reference cache records that traffic as growth_copy_bytes so the profiler can separate it from attention. Week 3 replaces repeated concatenation with preallocated pages for serving.

Task 2: Build the Cached Week 2 Model

src/tiny_llm/qwen3_week2.py

Keep the Week 1 Python model and its full-prefix generation loop unchanged. Build the separate qwen3_week2.py model with the same dense weights and the Week 1 mlx.core RMSNorm, RoPE, SwiGLU, and attention equations. Change only the state flow: the Week 2 model accepts a cache and an offset, while Week 1 continues to recompute the full prefix. Every later Week 2 chapter starts from this baseline.

Give each layer its own cache.
Add an offset argument to the model. It is the number of tokens already in the cache, and therefore the position of the first incoming token.
The argument should match the cache’s current sequence length. Assertions can make this invariant explicit.
The caller and cache both track the offset to make consistency checks easier.

Example computation flow:

x: B, L, E
q = linear(x, wq) -> B, L, H_q, D
k = linear(x, wk) -> B, L, H, D
v = linear(x, wv) -> B, L, H, D
q = rms_norm(q, q_norm)
k = rms_norm(k, k_norm)
q = rope(q, offset=slice(offset, offset + L))
k = rope(k, offset=slice(offset, offset + L))
transpose q, k, v to B, H, L, D
k, v = cache.update_and_fetch(k, v)  # k/v: B, H, S, D; q: B, H_q, L, D
x = scaled_dot_product_attention_grouped(q, k, v, scale, mask) -> B, H_q, L, D
# q/k/v and the returned model tensor are BF16; attention arithmetic is still the Week 1 `mlx.core` path
transpose and reshape x to B, L, H_q * D
x = linear(x, wo) -> B, L, E

Here, L is the number of incoming query tokens and S is the total cached sequence length after the update. This matches the Week 1 GQA convention: L is the query length, while S is the key/value source length. During single-token decoding, L = 1 and S grows by one on each call.

The linear layers, RMSNorm, RoPE, SwiGLU, and attention remain the Week 1 Python implementations at this checkpoint. Save packed weights and fast kernels for later checkpoints so this measurement isolates one algorithmic change. The model still uses BF16 storage; “Week 1 Python” describes the implementation style, not a return to an FP32 model.

Task 3: Create Request-Scoped Caches

src/tiny_llm/qwen3_week2.py

Implement create_kv_cache so each request receives one cache handle per Transformer layer. Pass the matching cache through each block, and keep the caller’s offset equal to the cache’s logical length.

The Day 1 test checks this request-scoped lifecycle together with the cache and model work from the earlier tasks.

Task 4: Connect the Serving Loop

src/tiny_llm/generate.py

Send the complete prompt on the first model call to prefill the cache. On each later call, send only the token produced by the preceding step and the number of tokens already cached. Week 3 moves this same lifecycle into the continuous-batching scheduler.

For example:

tokenized_prompt: [1, 2, 3, 4, 5, 6]
prefill: _step(model, [1, 2, 3, 4, 5, 6], 0)  # returns 7
decode:  _step(model, [7], 6)  # returns 8
decode:  _step(model, [8], 7)  # returns 9
...

You can test your solution with:

pdm run main --solution tiny_llm --loader week2 \
  --week2-checkpoint kv-cache --model qwen3-4b

You can also run the same loop with the reference solution:

pdm run main --solution tiny_llm_ref --loader week2 \
  --week2-checkpoint kv-cache --model qwen3-4b

Integrate and Measure

Finish Day 1 with a matched Week 1 versus cached Week 2 observation. The runner uses fresh processes, applies the same Qwen3-4B 128×129 workload to both rows, and writes the configuration beside the result:

pdm run bench-week2-progression --offline --solution tiny_llm --repeats 2 \
  --variant week1 --variant week2-kv-cache \
  --model qwen3-4b --input-len 128 --output-len 129 --warmup 2 \
  --json-output week2-day1-cache.json

Keep this JSON as Day 2’s baseline. Its useful result is the matched observation and recorded workload identity, not a speedup claim for another model, prompt length, output length, or device.

Day 1 changes the generation algorithm by removing full-prefix recomputation, so measure it with the end-to-end benchmark rather than inventing a shader-level limiter from a GPU trace. On Day 2, attribute this exact cached workload and turn the observation into a falsifiable next change. Begin Day 3 only after that evidence names dense projections.

🚧 Week 2 Day 2: Benchmarking and Profiling

Day 1 leaves you with a cached BF16 model and a working kv-cache checkpoint. Day 2 does not add another model operator. The supplied benchmark and portable attribution runners own request generation, warmups, synchronization, phase timing, and cache release. Your job is to freeze one like-for-like workload, identify its dominant operator category, and write the short decision that chooses the next change.

Start with the focused benchmark-lifecycle check:

pdm run test --week 2 --day 2

When it passes, record one matched tiny_llm/MLX pair and one attribution result, then write the short decision that follows from them. Those portable JSON records are the Day 2 checkpoint. Metal capture remains optional and never gates the next chapter.

Benchmark the Cached Model

Before changing the model, make the comparison trustworthy. Prefill processes many prompt tokens at once, while decode usually processes one token per request. At this checkpoint, decode repeatedly reads dense BF16 projection weights. Because a change can help one phase while hurting the other, benches/bench.py reports them separately:

prefill tokens per second: prompt tokens divided by prefill time;
decode tokens per second: generated tokens after the first token divided by decode time.

The first generated token is part of prefill. Leaving it out of decode keeps prompt length from distorting the decode number.

Decide what prefill should return before comparing implementations. Prompt scoring needs logits for every position; serving needs only the final prompt logit. Use --prefill-logits all for the former and --prefill-logits last for the latter. The runner applies one choice to your solution and MLX alike, so the two rows do the same work.

Keep the Week 2 generation algorithm matched too. Both sides use a KV cache: prefill the prompt once, then pass only the newly generated token on each decode step. A cached MLX baseline against a full-prefix solution would compare two different algorithms instead of locating the next optimization target.

Record a Matched Baseline

Use the same model, prompt length, output length, device, and warmup count for your solution and MLX:

pdm run bench --solution tiny_llm --loader week2 \
  --week2-checkpoint kv-cache --model qwen3-4b \
  --num-seqs 1 --min-input-len 128 --max-input-len 128 \
  --min-output-len 65 --max-output-len 65 --warmup 2 \
  --prefill-logits last

pdm run bench --solution mlx --loader week2 --model qwen3-4b \
  --num-seqs 1 --min-input-len 128 --max-input-len 128 \
  --min-output-len 65 --max-output-len 65 --warmup 2 \
  --prefill-logits last

Use --solution tiny_llm_ref with the same arguments when you want to compare your solution with the reference solution instead of MLX.

Or run the cumulative ladder in fresh processes:

pdm run bench-week2-progression --offline --repeats 2 \
  --solution tiny_llm \
  --variant week2-kv-cache --variant mlx \
  --model qwen3-4b --input-len 128 --output-len 129 --warmup 2 \
  --prefill-logits last --json-output week2-baseline.json

Benchmark on an otherwise idle machine. Stop other CPU- and GPU-intensive workloads, keep power mode and ambient conditions fixed, and wait for a stable temperature before comparing runs. Repeat each command, report the median, and record the hardware, MLX and mlx-lm versions, prefill-logit mode, and exact model. After a dependency upgrade, remeasure MLX instead of carrying the old baseline forward.

Synchronize Lazy Work

MLX builds computation graphs lazily. Timing only the Python call measures graph construction instead of GPU execution, so every timed iteration must evaluate its output:

start = perf_counter()
output = function()
mx.eval(output)
elapsed = perf_counter() - start

The benchmark must also call the cache release hook after warmups and timed runs. That lets caches return owned or shared resources even when a run fails; the focused Day 2 test covers both paths.

Attribute the Cached Model

Next, attribute the same cached-decode workload. Keep the learner solution, model, decode phase, and 128-token context fixed:

pdm run profile-week2-kernels --solution tiny_llm --model qwen3-4b \
  --case kv-cache:decode:128 --warmup 4 --iterations 12 \
  --json-output week2-day2-attribution.json

The result identifies its source, checkpoint, phase, token count, prompt rule, software, host, category medians, and category shares without depending on a private function name or Metal symbol. On the checked M4 Pro run, dense projections accounted for 81.5% of attributed cached-decode time. That bounded observation selected packed W4 projections for Day 3; another device or shape may point somewhere else.

Turn the observation into a decision with three sentences:

“Dense projections dominate this exact cached-decode workload.”
“Packing W4 weights and changing only the selected projection path should reduce that category and improve matched decode.”
“I will reject or revise the hypothesis if projection time does not fall or complete-model decode regresses under the same workload.”

Substitute the category you observed for the checked example. Your required work ends with the benchmark, attribution, and decision record. The macOS 27 capture lab is optional; no trace, gpudebug output, screenshot, or device-specific counter gates Day 3.

Why Quantize: The Decode Roofline

The measurement now has a hardware reason to test. LLM decode is typically memory-bandwidth bound: each token reads the model’s weights while doing relatively little work with them. Use the dimensions in the official Qwen3-4B configuration to calculate the ideal bound:

Qwen3-4B dimensions:
  hidden size        h = 2,560
  MLP size           i = 9,728
  query width        q = 4,096
  key/value width   kv = 1,024
  layers             L = 36
  vocabulary         V = 151,936

Projection weights per layer:
  Q and O: 2 × h × q       =  20,971,520
  K and V: 2 × h × kv      =   5,242,880
  MLP:     3 × h × i       =  74,711,040
  total per layer          = 100,925,440

All transformer layers: L × 100,925,440 = 3,633,315,840
Tied vocabulary head:    V × h           =   388,956,160
Total streamed weights:                    4,022,272,000

FLOPs per token: 2 × 4,022,272,000 = 8.045 GFLOPs

Count the tied embedding matrix once as the vocabulary projection. The single-row embedding lookup, normalization weights, activations, KV reads, and attention work are omitted, so the result is an upper bound for linear layers rather than a prediction of complete-model throughput. A dense FP16 or BF16 weight occupies two bytes:

4,022,272,000 weights × 2 bytes = 8.045 GB per token
arithmetic intensity = 8.045 GFLOPs / 8.045 GB = 1.0 FLOP/byte

FP16 and BF16 divide their 16 bits differently: FP16 gives more bits to the significand, while BF16 gives more bits to the exponent. That affects numerical range and precision, but not this bandwidth calculation. The course uses BF16 for activations and outputs.

Dense weight format	Bits per weight	Bytes per weight	Streamed weight bytes per token	Weight arithmetic intensity
FP16	16	2	8.045 GB	1.0 FLOP/byte
BF16	16	2	8.045 GB	1.0 FLOP/byte

This is the baseline to improve: both dense formats must stream roughly 8 GB of projection weights to generate one token. Save the matched benchmark result, then continue to Day 3, where the model keeps weights packed, replaces the live projection path, and reruns the same benchmark.

Optional: Inspect a Week 2 Capture on macOS 27

The synchronized product benchmark and portable operator-attribution runner are sufficient for every required Week 2 checkpoint. This page is an optional deeper look at the same evidence loop for learners with macOS 27 and /usr/bin/gpudebug. It is never an acceptance gate.

The checked example used Qwen3-4B on an Apple M4 Pro at source commit add389b747793e910f0506f5720dd0aac373d126, macOS 27 build 26A428, gpudebug 1.0, MLX 0.32.0, and mlx-lm 0.31.3. Its product control used a 128-token prompt, 129 output tokens, final-row prefill logits, seed 0, two warmups, and two balanced fresh-process samples. Its attribution cases used four warmups and twelve synchronized iterations. These identities bound the example; they are not a portable timing baseline.

1. Prove Correctness First

Choose one checkpoint, phase, and token count. Run its focused test before capturing it. For the Day 4 decode example:

pdm run build-ext
pdm run test --week 2 --day 4

pdm run profile-week2-kernels --solution tiny_llm --model qwen3-4b \
  --case swiglu:decode:128 --warmup 4 --iterations 12 \
  --json-output out/swiglu-decode-128.attribution.json

The second command is the portable evidence path. Read its checkpoint, workload, dominant category, and category shares before opening a GPU trace.

2. Capture One Synchronized Region

Create out/ first and choose output names that do not exist. The helper refuses to overwrite a trace, metadata file, or manifest.

mkdir -p out

MTL_CAPTURE_ENABLED=1 pdm run capture-week2 \
  --solution tiny_llm --model qwen3-4b \
  --checkpoint swiglu --phase decode --tokens 128 \
  --trace out/swiglu-decode-128.gputrace \
  --metadata out/swiglu-decode-128.capture.json \
  --manifest out/swiglu-decode-128.trace-manifest.sha256

The helper compiles and warms the exact shape outside the capture, then captures one synchronized model region. The metadata records source, model, checkpoint, phase, token count, prompt rule, software, host, and a canonical workload identity. The path-sorted manifest hashes every file inside the .gputrace package; treat the package and manifest as one evidence object.

3. Replay and Reduce

Run the serialized profile and save the JSON stream. You may also collect timeline, shader, or command queries into a second JSON-lines file.

gpudebug --json -t out/swiglu-decode-128.gputrace --timeout 1800 \
  -c 'profile run --gpu-state default --exec serial' \
  > out/swiglu-decode-128.profile.jsonl

pdm run reduce-week2-gpudebug \
  --capture-metadata out/swiglu-decode-128.capture.json \
  --manifest out/swiglu-decode-128.trace-manifest.sha256 \
  --profile-jsonl out/swiglu-decode-128.profile.jsonl \
  --commands-jsonl out/swiglu-decode-128.commands.jsonl \
  --output out/swiglu-decode-128.gpudebug.json

If you did not collect command queries, omit --commands-jsonl. Missing timeline, shader, command, or counter trees must remain explicitly unavailable; do not replace them with zero and do not infer occupancy. In the checked pre-SIMD 128-token prefill capture, the replay exposed timeline counters but no shader ranking. In the checked 32-token Split-K capture, only static dispatch presence was available and no occupancy conclusion was drawn.

4. Write a Bounded Decision

Use three sentences:

identify the dominant category for this exact checkpoint and workload;
name the next bounded change and the same-workload result that would support it;
state the result that would falsify the hypothesis or make you revert it.

For example: “At swiglu:decode:128, packed projections dominate this M4 Pro capture and the portable attribution. I will change only the selected projection schedule and rerun the identical workload. I will revert or choose another category if projection time does not fall or the complete-model phase regresses.” This is a reasoning record, not a claim that another device has the same bottleneck.

5. Preserve the Compact Result, Then Clean Up

Keep the capture metadata, manifest, portable attribution, and reduced result until you have checked their matching workload identity. Raw trace packages can be enormous; after preserving the compact evidence you need, remove only the exact trace package and raw streams you created:

rm -rf -- out/swiglu-decode-128.gputrace
rm -f -- out/swiglu-decode-128.profile.jsonl \
  out/swiglu-decode-128.commands.jsonl

The repository includes a compact checked M4 Pro result at benchmark_results/m4-pro-qwen3-4b-week2-gpudebug-macos27-mlx-0.32.0.json. Learners without macOS 27 can use it to practice reading identity, availability, dominant categories, and keep/reject decisions. They do not need to reproduce its exact kernel names, timings, or Metal schedule.

🚧 Week 2 Day 3: Quantize the Model

Day 2 leaves you with a synchronized dense BF16 baseline. The Day 3 starter already supplies the packed-weight container and its QuantizedWeights.from_mlx_layer loader, the extension declaration and binding, fail-closed C++/Metal stubs, and cumulative model switches. Your work is to:

dequantize selected embedding rows without expanding the full table;
define the lazy quantized-matmul primitive and its validation boundary;
implement the readable Metal matrix control and the decode-shaped SIMD matvec; and
wire packed projections and the tied output head into the live cached model.

Begin with the Python wrapper gate. Then build and exercise the GPU operator:

pdm run build-ext
pdm run test --week 2 --day 3 -- -k task_1
pdm run test --week 2 --day 3 -- -k gpu

Finish with the complete Day 3 gate and the quantized-matvec model checkpoint. The result is complete only when the cached model dispatches through your quantized path; packed storage and an isolated fast kernel are intermediate steps.

📚 Readings

Debug Metal Without a CPU Twin

You do not need a C++ CPU twin. Bring up the operator with this three-level validation ladder:

Write the equation in Python with mlx.core. This is the semantic oracle.
Translate it into a deliberately simple Metal kernel, usually with one thread responsible for one output element.
Optimize the validated Metal kernel with SIMD groups, vectorized loads, or SIMD-group matrix operations.

Compare each level with the one immediately above it. Full-model text output is too indirect to diagnose an optimized kernel.

Make Failures Small and Synchronous

Use deterministic fixtures whose expected values are easy to inspect: zeros, ones, ramps, identity-like weights, and a fixed random seed. Exercise a small aligned shape and then a tail shape. For example, test 8 and 10 rows for an 8-row tile, or sequence lengths 32 and 35 for a 32-token block.

MLX execution is lazy, so force evaluation directly after the operator under test. This turns a delayed compile or GPU execution failure into a failure at the responsible call site:

expected = python_reference(*inputs)
actual = metal_operator(*inputs)
mx.eval(expected, actual)

assert actual.shape == expected.shape
assert actual.dtype == mx.bfloat16
assert mx.allclose(actual, expected, rtol=2e-2, atol=2e-2).item()

When a check fails, inspect the wrapper boundary before the arithmetic. Assert the tensor rank, shape, dtype, and contiguity assumptions in Python or C++, and verify that the encoded buffer indices match the Metal function signature. Then classify the failure:

a pipeline creation error usually means the kernel name, specialization, or Metal compilation is wrong;
an execution or address error usually means a grid, bounds check, stride, or buffer binding is wrong;
a finite but inaccurate result usually means the indexing, reduction, mask, dequantization, or accumulator update is wrong.

For a numerical mismatch, simplify the schedule temporarily. Assign one output to one thread, remove cooperative loads, and compare an intermediate such as a dequantized weight group, a partial dot product, or an online-softmax row. A small debug-only output buffer is usually more useful than printing from every GPU thread. Restore one optimization at a time, rerunning the aligned and tail-shape tests after each change.

Represent Weights With Fewer Bits

Quantization stores each floating-point weight as a value from a small integer codebook plus the parameters needed to reconstruct an approximation. This course uses weight-only 4-bit quantization:

W4 means that each logical weight is represented by a 4-bit code.
A16 means that activations and outputs remain 16-bit floating point.
The resulting path is called W4A16. This course uses BF16 for its activations, scales, biases, and outputs.

The 16 possible codes approximate the original values. In return for some numerical precision, the smaller representation reduces memory traffic.

The kernel does not materialize a dense BF16 weight matrix. It unpacks each 4-bit code, reconstructs the weight in registers, and immediately multiplies it by the corresponding BF16 activation.

Group-Wise Affine Quantization

Rather than applying one scale to an entire weight matrix, divide each row into groups and quantize each group independently. A local scale and bias retain more information about that group’s weight distribution.

For a weight matrix $W$ of shape $(K, N)$ , divide each row into groups of size $G$ . The Qwen3-4B MLX 4-bit checkpoint used in this course has a fixed group size of 128:

Logical weight matrix W: K × N

Group size: G = 128
Number of groups per row = N / G

For each stored group of G consecutive values in a row:
  1. Unpack each unsigned 4-bit code q in [0, 15]
  2. Load the group's stored scale s and bias b
  3. Reconstruct each value as q * s + b

Reconstruct a Stored Group

The checkpoint already contains packed codes and affine parameters. For an unpacked unsigned code $q$ , use the stored scale $s$ and bias $b$ directly:

$\overset{w}{^} = q s + b$

The codes are unsigned, but the stored scale is signed. A positive scale maps code 0 to the lower endpoint and code 15 toward the upper endpoint. A negative scale reverses the mapping: code 0 is the upper endpoint and code 15 moves toward the lower endpoint. Both orientations occur in the shipped Qwen3-4B MLX checkpoint. Use scale and bias as stored instead of reconstructing them from an assumed min/max orientation.

For example, these two stored parameter pairs reconstruct the same endpoint range in opposite code order:

positive orientation: scale =  0.0867, bias = -0.5  => q=0 is -0.5, q=15 is about 0.8
negative orientation: scale = -0.0867, bias =  0.8  => q=0 is  0.8, q=15 is about -0.5

All required quantized-matmul tests use group_size = 128 and BF16 scales, biases, activations, and outputs. Normalize those tensors to BF16 in your solution’s model loader so every later kernel receives one model dtype.

Packed Storage Layout

The 4-bit codes are packed for compact storage and efficient access:

Logical weight matrix: K × N
Dense BF16 storage: K × N bfloat16 (2 bytes each) = 2KN bytes
W4 code storage: K × N int4 (0.5 bytes each) = 0.5KN bytes

Packing: 8 × 4-bit values fit in one uint32 (32 bits)

Packed codes shape: K × (N / 8) uint32
Scales shape: K × (N / G) bfloat16
Biases shape: K × (N / G) bfloat16

Example packing for 8 consecutive 4-bit values [a, b, c, d, e, f, g, h]:

uint32_value = (h << 28) | (g << 24) | (f << 20) | (e << 16) |
               (d << 12) | (c << 8)  | (b << 4)  | a

Unpacking:
  a = (uint32_value >> 0)  & 0xF
  b = (uint32_value >> 4)  & 0xF
  c = (uint32_value >> 8)  & 0xF
  ...
  h = (uint32_value >> 28) & 0xF

Revisit the Decode Roofline

The packed codes are not the entire W4 representation. Each group of 128 weights also stores one BF16 scale and one BF16 bias:

bytes per W4 weight = 0.5 + (2 + 2) / 128 = 0.53125 bytes
streamed W4 bytes   = 4,022,272,000 × 0.53125 = 2.137 GB per token
arithmetic intensity = 8.045 GFLOPs / 2.137 GB = 3.765 FLOPs/byte

Now W4 can be added to the dense comparison:

Weight format	Value bits	Metadata per 128 weights	Effective bytes per weight	Streamed weight bytes per token	Weight arithmetic intensity
FP16	16	None	2	8.045 GB	1.0 FLOP/byte
BF16	16	None	2	8.045 GB	1.0 FLOP/byte
W4	4	One BF16 scale and one BF16 bias	0.53125	2.137 GB	3.765 FLOPs/byte

This representation reduces projection weight traffic by 3.765×. Treat that ratio as a bandwidth ceiling for one-token decode, not as an end-to-end speedup promise.

Theoretical Decode Roofline Across Apple Silicon

Apple publishes unified-memory bandwidth but not a directly comparable BF16 GPU TFLOPS figure. A bandwidth roofline can therefore be calculated without assuming a compute ceiling:

ideal tokens/s = advertised memory bandwidth / streamed weight bytes per token

The table uses the highest-bandwidth configuration of each named chip. GB is decimal, matching Apple’s specifications. The results are theoretical ceilings, not benchmark measurements.

Chip	Bandwidth	FP16/BF16 roofline	W4 roofline
M1 Pro	200 GB/s	24.9 tok/s	93.6 tok/s
M1 Max	400 GB/s	49.7 tok/s	187.2 tok/s
M1 Ultra	800 GB/s	99.4 tok/s	374.4 tok/s
M2 Pro	200 GB/s	24.9 tok/s	93.6 tok/s
M2 Max	400 GB/s	49.7 tok/s	187.2 tok/s
M2 Ultra	800 GB/s	99.4 tok/s	374.4 tok/s
M3 Pro	150 GB/s	18.6 tok/s	70.2 tok/s
M3 Max	400 GB/s	49.7 tok/s	187.2 tok/s
M3 Ultra	819 GB/s	101.8 tok/s	383.3 tok/s
M4 Pro	273 GB/s	33.9 tok/s	127.8 tok/s
M4 Max	546 GB/s	67.9 tok/s	255.5 tok/s

The advertised bandwidths come from Apple’s specifications for M1 Pro and Max, M1 Ultra, M2 Pro and Max, M2 Ultra, M3 Pro and Max, M3 Ultra, and M4 Pro and Max. Apple’s current Mac Studio pairs M4 Max with M3 Ultra, so there is no M4 Ultra row.

These values assume peak advertised bandwidth, one read of every projection weight, and no other traffic or work. A complete model also reads activations and KV, launches other operators, and cannot sustain peak bandwidth continuously, so actual throughput is lower. The performance appendix records measured results separately from this theoretical exercise.

This roofline describes one-token decode, where M = 1 and each streamed weight serves one activation row. Prefill reuses each weight tile across many rows and raises arithmetic intensity, so it needs a matrix schedule. The decode bandwidth ratio does not predict prefill performance.

Quantized Matrix Multiplication

Mathematical Formulation

For standard matrix multiplication $C = A B^{T}$ where:

$A$ : shape $(M, N)$ , bfloat16 (activations)
$B$ : shape $(K, N)$ , quantized to int4 (weights)
$C$ : shape $(M, K)$ , same 16-bit dtype as $A$ (output)

Each element $C [i, k]$ is computed as:

$C [i, k] = j = 0 \sum N - 1 A [i, j] \times B [k, j]$

With quantization, $B [k, j]$ is represented as:

$B [k, j] = B_{quantized} [k, j] \times scale [k, g] + bias [k, g]$

where $g = ⌊ j / G ⌋$ is the group index.

Substituting:

$C [i, k] = g = 0 \sum N / G - 1 j^{'} = 0 \sum G - 1 A [i, g \times G + j^{'}] \times (B_{quantized} [k, g \times G + j^{'}] \times scale [k, g] + bias [k, g])$

Rearranging:

$C [i, k] = g = 0 \sum N / G - 1 scale [k, g] j^{'} = 0 \sum G - 1 A [i, g \times G + j^{'}] \times B_{quantized} [k, g \times G + j^{'}] + bias [k, g] j^{'} = 0 \sum G - 1 A [i, g \times G + j^{'}]$

The scale and bias are constant within a group, so the computation can reuse them across all values in that group.

Computation Flow

Input:
  A: M × N (bfloat16 activations)
  B_quantized: K × (N/8) (uint32, packed weights)
  scales: K × (N/G) (bfloat16)
  biases: K × (N/G) (bfloat16)

Output:
  C: M × K (bfloat16)

For each output element C[i, k]:
  sum = 0  # float accumulator
  for each group g in 0..(N/G - 1):
    scale = scales[k, g]
    bias = biases[k, g]

    # Process G values in the group (G/8 uint32 packs)
    for each pack p in 0..(G/8 - 1):
      packed_value = B_quantized[k, g*(G/8) + p]

      # Unpack 8 × 4-bit values
      for bit_offset in [0, 4, 8, 12, 16, 20, 24, 28]:
        quantized = (packed_value >> bit_offset) & 0xF
        b_value = quantized * scale + bias
        a_value = A[i, g*G + p*8 + bit_offset/4]
        sum = sum + a_value * b_value

  C[i, k] = bfloat16(sum)

Task 1: Implement Quantized Linear and Embedding

src/tiny_llm/quantize.py
src/tiny_llm/embedding.py

Inspect and reuse the starter’s QuantizedWeights.from_mlx_layer packed-weight plumbing. Modify these learner-owned functions:

dequantize_weights and quantized_linear in src/tiny_llm/quantize.py;
QuantizedEmbedding.__call__ and QuantizedEmbedding.as_linear in src/tiny_llm/embedding.py.

QuantizedWeights holds a quantized matrix and its dequantization parameters:

Field	Shape	Description
`weight`	$(K, N /8)$ uint32	Packed quantized weights. Each uint32 stores eight consecutive 4-bit values.
`scales`	$(K, N / G)$ bfloat16	Stored signed per-group scale factors for dequantization. The sign determines which endpoint maps to the low codes.
`biases`	$(K, N / G)$ bfloat16	Stored per-group offsets. Code 0 reconstructs to this value.
`group_size`	int	Number of consecutive values that share the same scale/bias. For the Qwen3 MLX 4-bit weights used here, this is `128`.
`bits`	int	Quantization bit width (typically 4, meaning values are in range $[0, 15]$ )

The supplied from_mlx_layer method extracts these fields from an MLX quantized layer during model loading. Reuse it rather than introducing a second loader.

Then implement quantized_linear as a wrapper around quantized_matmul, using the same input convention as the standard linear function. The next task makes quantized_matmul runnable.

Keep the token embedding table quantized too. Add a QuantizedEmbedding wrapper with two call patterns:

embedding(input_ids) performs a row lookup. Gather the matching packed weights, scales, and biases. Unpack each uint32 with shifts and masks, repeat each group’s scale and bias across its 128 values, and compute q * scale + bias with basic mlx.core array operations. Do not call mx.dequantize. Put this unpacking logic in dequantize_weights(...) so the embedding and its direct tests share one explicit implementation.
embedding.as_linear(h) is the tied output projection. Implement this with quantized_linear(h, embedding_weight) so it uses your quantized matmul path instead of materializing the full vocab_size x hidden_size table. This path starts working once the quantized matmul kernel is implemented in the next tasks.

Task 2: Define the Quantized Matmul Primitive

src/extensions/src/tiny_llm_ext.h
src/extensions/bindings.cpp
src/extensions/src/quantized_matmul.cpp
src/extensions/CMakeLists.txt

The declaration, fail-closed source stub, binding, and build registration are already present. Keep the C++ declarations and definitions in the tiny_llm_ext namespace, and modify these exact functions:

tiny_llm_ext.h — Read the Week 2 Day 3 quantized_matmul(...) declaration and QuantizedMatmul primitive interface; keep its signature in sync with the binding.
bindings.cpp — Verify the existing m.def("quantized_matmul", ...) entry; do not create a second binding.
quantized_matmul.cpp — Replace the body of tiny_llm_ext::quantized_matmul(...) to validate inputs, determine the output shape, return a lazy mx::array, and reject CPU evaluation explicitly in QuantizedMatmul::eval_cpu(...).
CMakeLists.txt — Verify the existing quantized_matmul.cpp source registration; do not add a duplicate.

The extension API lets an mx.array graph node schedule the Metal loop from the next task. MLX owns the array lifetime and command encoder; your primitive supplies the quantized multiplication.

Build now to catch declaration, binding, and registration mismatches. The focused test checks the Task 1 Python wrappers. The primitive becomes runnable after you implement its Metal schedules in Task 3:

pdm run build-ext
pdm run test --week 2 --day 3 -- -k task_1

Task 3: Implement Metal Matrix Products

Before writing the Metal kernels, connect their work to Metal’s four nested execution scopes:

Lane (thread). The smallest unit. Each lane executes the same instruction stream with its own register file. Lanes within a SIMD group can share data through simd_ operations.
SIMD group (warp/subgroup). A fixed-size set of lanes (32 on Apple GPUs) that execute in lockstep. simd_sum, simd_shuffle, and simdgroup_matrix operations work within this scope. A SIMD group cannot directly share registers with another SIMD group in the same threadgroup.
Threadgroup (block). A collection of SIMD groups scheduled together on one GPU core. Threadgroups share threadgroup memory (explicitly allocated with threadgroup address space and synchronized with threadgroup_barrier). The grid is a 1D/2D/3D array of threadgroups.
Grid. The total work dispatched. dispatchThreadgroups launches a grid of threadgroups; the GPU schedules them across available cores. Increasing the grid’s threadgroup count can expose more independent work, but a finer partition can also duplicate reads or require partial-result merging.

Keep two launch knobs separate. Adding SIMD groups within one threadgroup adds threads and can raise register demand; it increases threadgroup-memory use only when the schedule allocates shared storage per group or tile. Either resource can reduce the number of resident threadgroups. Adding threadgroups to the grid changes how output or reduction work is partitioned. Measure both choices; neither guarantees higher throughput.

Begin with the required two-SIMD-group matvec schedule for Qwen. Then benchmark two, four, eight, and sixteen groups per threadgroup as described below. Change the grid partition separately so each measurement isolates one launch knob.

src/extensions/src/quantized_matmul.metal
src/extensions/src/quantized_matmul.cpp

Modify these exact starter functions:

QuantizedMatmul::eval_gpu in quantized_matmul.cpp;
quantized_matmul_vanilla_w4a16_g128 and quantized_matvec_x4_fast_w4a16_g128 in quantized_matmul.metal;
quantized_matmul_vanilla and quantized_matvec_custom in src/tiny_llm/quantize.py for the explicit comparison paths.

Write both Metal kernels, then connect eval_gpu to them. On GPU, the Python quantized_matmul wrapper always dispatches your primitive. The required path never routes through mx.quantized_matmul.

Work in two measured stages. Both implement the same math, with schedules for different shapes:

Vanilla matmul: one Metal thread computes one output element. This is the direct GPU translation of the computation flow above and an inspectable bring-up control.
SIMD matvec: for decode, SIMD lanes cooperate on the reduction for one activation row and calculate several output columns together.

Here, M is the number of activation rows after flattening every leading dimension. Day 3 uses this explicit dispatch:

Activation rows	Kernel	Role at this checkpoint
`M <= 8`	SIMD matvec	Optimized path for decode and other very small matrix inputs.
`M > 8`	Vanilla matmul	Correctness-first prefill path; Day 5 replaces it with a cooperative tiled kernel.

The cutoff does not extend the SIMD kernel to larger M. These are separate schedules: Day 3 optimizes vector-shaped decode and leaves matrix-shaped prefill visible for the later benchmark to select.

Keep the vanilla function callable as quantized_matmul_vanilla so every optimization can be compared directly with its readable control.

Stage 1: Vanilla Matmul

Start with a two-dimensional grid over output row i and output column k. Each thread walks all N input values, unpacks eight int4 weights from each uint32, applies the group scale and bias, and accumulates one C[i, k] in float32. This kernel repeats activation loads and does not share work, but its control flow mirrors the equation and makes it a useful debugging control. The Python mlx.core equation remains the correctness oracle for both Metal schedules.

Keep the vanilla kernel for matrix-shaped prefill in this chapter; Day 5 revisits that workload with cooperative tiling.

Stage 2: SIMD Matvec

Decode normally has M = 1, so an 8×8 matrix tile would leave most rows empty. Instead, let one SIMD group reduce the input dimension and use simd_sum to combine lane-local partial sums. Begin with two output columns per group as an inspectable schedule. For the Qwen3-4B checkpoint, evaluate a four-column path in which each lane loads two adjacent packed words, or 16 activations, and reuses them across the four outputs.

The optimized path also uses the affine identity

$j \sum a_{j} (s q_{j} + b) = s j \sum a_{j} q_{j} + b j \sum a_{j}$

to avoid applying the bias separately to every unpacked value. It also scales the activations once and reads four packed int4 values through a 16-bit mask, avoiding a shift for every weight and output row. This adds live accumulators, so test it as a complete schedule rather than assuming fewer integer instructions must be faster.

Tune the SIMD Schedule

Treat output width, threadgroup size, and shared-memory reuse as benchmark variables. Begin with this Qwen-focused configuration:

flatten all leading activation dimensions into M,
use the custom matvec when M <= 8 and the vanilla matmul when M > 8,
compute four output columns per SIMD group and load two adjacent packed words per lane,
launch two SIMD groups, or eight output rows, per threadgroup.

These thresholds are measured starting points rather than mathematical requirements. Keep them visible in the dispatcher and vary one choice at a time. Compare two, four, and eight output columns per SIMD group. More columns increase activation reuse, but also extend accumulator lifetimes and raise register pressure. Compare two, four, eight, and sixteen SIMD groups per threadgroup. More groups expose additional outputs, but may duplicate activation reads and reduce residency.

Evaluate the affine rearrangement as part of the complete schedule. A lower instruction count helps only if the longer-lived activation sum and output accumulators do not reduce occupancy. Select the schedule with a synchronized whole-model decode benchmark, not with an instruction-count estimate.

Define a row-contiguous Python-to-extension contract for scales, biases, activations, and packed weights. Call mx.contiguous once at that boundary and validate the layout in the C++ primitive before encoding the kernel. Metal receives raw buffers rather than implicit array strides, so layout is a correctness condition as well as a performance condition.

Read activations directly in the kernel. The one-row activation is small and cache-friendly, while staging it in threadgroup memory adds a barrier to every projection. If you test shared staging as an ablation, keep it only when the whole-model result shows that reuse outweighs synchronization.

Kernel Requirements

Implement both required layouts in quantized_matmul.metal:

First, implement the vanilla one-thread-per-output matrix grid.
For M <= 8, assign one SIMD group to an output tile. Cooperatively reduce the input dimension and compute several output columns per group.
For M > 8, dispatch the vanilla matrix grid. Do not loop over rows with the SIMD matvec schedule; Day 5 introduces the tiled prefill schedule.
The required kernel supports bfloat16_t inputs and outputs. The Week 2 checkpoint does not add a second model-storage dtype.
Apply the group-wise dequantization loop defined earlier in this chapter:
- Iterate over groups of 128 values.
- Unpack int4 values from each uint32.
- Dequantize each value with q * scale + bias.
- Accumulate products in float, then cast the result to the kernel dtype.
Add boundary checks (i < M, k < K) before writing output.

The custom kernel only needs to support bits = 4 and group_size = 128. Use the group size to compute groups_per_row and the packed-weight offsets. Instantiate the required Metal kernel for bfloat16_t and select it in eval_gpu. If you retain an optional half specialization, keep it out of the model dispatch in your solution.

GPU Dispatch

Complete eval_gpu in quantized_matmul.cpp with axpby’s GPU dispatch pattern:

Get the Metal device and command encoder from the stream.
Load the quantized matmul kernel matching the output dtype from the Metal library.
Bind the input and output buffers and the dimension constants (M, N, K). The buffer order must match the kernel signature.
Select the matrix-vector layout for M <= 8; otherwise select the vanilla matrix layout. Keep both paths explicit for direct comparisons. Calculate a SIMD-aligned thread-group configuration and tile output columns so packed input values and activations can be reused. Use the four-column, two-packed-word kernel with two SIMD groups.
Dispatch with dispatchThreadgroups.

Run the focused GPU gate:

pdm run build-ext
pdm run test --week 2 --day 3 -- -k gpu

The direct tests cover matvec at M = 1 and M = 8, the vanilla matmul at M = 128, and compare them with an MLX oracle. The oracle checks the result; it is not the implementation under test.

Task 4: Integrate Before Continuing

src/tiny_llm/qwen3_week2.py

Modify Qwen3ModelWeek2.__init__, Qwen3MultiHeadAttention.__call__, Qwen3MLP.__call__, and Qwen3ModelWeek2.__call__. These are the points that load quantized weights, replace dense projections, and keep only the requested logits row.

Now integrate quantized matrix multiplication into the Week 2 Qwen3 model so its linear layers remain quantized throughout inference.

Change the weight type from mx.array to QuantizedWeights for every attention projection (wq, wk, wv, and wo) and MLP projection (w_gate, w_up, and w_down). Replace linear(x, w) with quantized_linear(x, w). In the Week 2 model loader, use QuantizedWeights.from_mlx_layer(...) instead of materializing a 16-bit matrix. Keep the Week 1 model’s boundary intact; its layers still expect plain mx.array weights.

For embeddings, wire the QuantizedEmbedding from Task 1 into the loader. Load embed_tokens with QuantizedWeights.from_mlx_layer(...) and pass it to QuantizedEmbedding. If the model has a separate lm_head, keep that head as QuantizedWeights too and apply it with quantized_linear; lm_head is a projection rather than an embedding lookup.

Normalize each loaded layer’s scales and biases to BF16. Require scales, biases, and activations to match and return BF16. If the output is nan or otherwise invalid, check for a dtype mismatch first.

Preserve the quantized layer’s parameters as well. The model should pass w.group_size and w.bits to the extension, which should validate the course assumptions: group_size = 128 and bits = 4.

Run the complete gate, then the live model checkpoint:

pdm run test --week 2 --day 3

pdm run main --solution tiny_llm --loader week2 \
  --week2-checkpoint quantized-matvec --model qwen3-4b

Measure that same learner solution:

pdm run bench --solution tiny_llm --loader week2 \
  --week2-checkpoint quantized-matvec --model qwen3-4b \
  --num-seqs 1 --min-input-len 128 --max-input-len 128 \
  --min-output-len 65 --max-output-len 65 --warmup 2

Run the same command with --solution tiny_llm_ref to compare it with the reference solution.

Keep the vanilla matrix product callable as an inspectable Metal control. The Python mlx.core equation remains the correctness oracle, while decode integrates only the SIMD matvec.

Verify Quantization in the Complete Model

Before moving on, confirm that model inference actually calls the quantized matvec kernel instead of merely registering and testing it in isolation.

The checkpoint is complete when the model’s projection dispatcher is wired to your custom primitive. Decode-shaped work must route through quantized_linear → quantized_matvec_custom → the extension primitive → the Metal matvec. Matrix-shaped work must route through quantized_linear → quantized_matmul → the extension primitive → its Metal matrix schedule. The supplied tests validate packed model state and the direct operators. Use the live model command to verify that those pieces compose.

Measure the cumulative model and the real projection shapes:

pdm run bench-week2-progression --offline --solution tiny_llm --repeats 2 \
  --variant week2-kv-cache --variant week2-quantized-matvec --variant mlx \
  --model qwen3-4b --input-len 128 --output-len 129 --warmup 2 \
  --prefill-logits last

pdm run profile-week2-kernels --solution tiny_llm --model qwen3-4b \
  --case kv-cache:decode:128 --case quantized-matvec:decode:128 \
  --warmup 4 --iterations 12 \
  --json-output week2-day3-attribution.json

Keep one cumulative model row and one representative real-shape projection comparison. In the checked M4 Pro example, packed W4 reduced attributed projection time by 69.0% and fixed-workload decode rose from 24.38 to 58.90 tokens/s. The re-profile then exposed normalization, position, and activation at 33.5% of attributed time, selecting Day 4. These are bounded observations from one machine and two product samples, not portable timing thresholds. The complete campaign and attribution are in the performance appendix.

If you need to continue without the custom Day 3 kernels, implement the same quantized_linear interface with mx.quantized_matmul and leave the rest of the course model unchanged. This is a local operator off-ramp, not --solution mlx: the latter runs a separate complete model and does not exercise your cache or model wiring.

🚧 Week 2 Day 4: Fused Model Kernels

Day 3 leaves the cached model using packed projections. Day 4 keeps the Week 1 Python equations as readable oracles and completes three separate extension shells already present in the starter:

src/tiny_llm/week2_kernels.py
src/extensions/src/week2_kernels.cpp
src/extensions/src/week2_kernels.metal

Work in checkpoint order: implement and integrate RMSNorm, then RoPE, then SwiGLU. After each operator, run its focused test and the live cumulative checkpoint. This keeps a local operator failure separate from an integration regression before all three are active:

pdm run build-ext
pdm run test --week 2 --day 4 -- -k rms
pdm run test --week 2 --day 4 -- -k rope
pdm run test --week 2 --day 4 -- -k swiglu

RMSNorm, RoPE, and SwiGLU surround the projections in every transformer layer. Week 1 gives you readable Python mlx.core equations; your Week 2 path keeps their interfaces and supplies purpose-built Metal kernels.

Your solution still uses MLX arrays and its extension API. MLX schedules the graph node, owns its buffers, and dispatches the Metal function, but your solution owns the arithmetic inside that function. Your solution does not call mx.fast.rms_norm, mx.fast.rope, or an MLX-provided SiLU implementation.

Why Fusion Helps

Week 1’s Python mlx.core equations already become native GPU work inside the lazy graph. Here, the useful question is how many operations, launches, and memory passes that graph still describes.

For example, RMSNorm expressed as mlx.core operations casts, squares, reduces, takes a reciprocal square root, multiplies, casts again, and applies a learned weight. A compiler may fuse some adjacent element-by-element work, but the row reduction is a boundary. Intermediate values and multiple dispatches remain possible.

A single fused Metal kernel gives you explicit control over the whole operator:

one dispatch replaces several graph operations;
values stay in registers or SIMD-group storage between steps;
float accumulation is used where numerical stability needs it;
inputs are read once when practical, and only the final tensor is written;
the grid matches decode shapes instead of a generic tensor operation.

So compare one purpose-built kernel with a graph of several general-purpose kernels. The source language is not the point; the resulting work is.

Task 1: RMSNorm

Start by replacing the fail-closed RMSNorm bodies: tiny_llm_ext::rms_norm, Week2RMSNorm::eval_cpu, and Week2RMSNorm::eval_gpu in src/extensions/src/week2_kernels.cpp, the week2_rms_norm function in src/extensions/src/week2_kernels.metal, and FastRMSNorm.__call__ in src/tiny_llm/week2_kernels.py. The starter already provides the header, binding, C++/Metal files, and CMake registration, so keep that API rather than adding a parallel one.

Begin with one SIMD group per input row, then benchmark it. A 2,560-element hidden row gives 32 lanes roughly 80 serial elements each; the optimized kernel launches 256 threads, or eight SIMD groups, per row. Each group reduces its portion with simd_sum; lane zero writes eight partial sums to threadgroup memory; the first SIMD group performs the second reduction:

sum_sq = simd_sum(each lane's partial sum)
inverse_rms = rsqrt(sum_sq / hidden_size + epsilon)
output[i] = input[i] * inverse_rms * weight[i]

All 256 lanes then normalize and scale their strided elements. This fuses the reduction and output pass into one dispatch and avoids materializing the squared tensor. Instantiate the required kernel for bfloat16. Keep the reduction, normalization, and weight multiplication in float, then cast the final result once. The Python reference equation rounds once before applying the weight, so compare the two with a tolerance rather than expecting bit-identical results.

The C++ primitive validates shape and dtype, allocates the output through MLX, binds the buffers and scalar constants, allocates eight float partial sums, and launches one 256-thread group per row. Compare this two-level reduction with a single-SIMD-group control to determine whether the extra parallelism offsets the threadgroup reduction on the target machine.

Wire FastRMSNorm into every Week 2 norm as soon as the kernel works. Then run the focused test and record the cumulative model result before touching RoPE:

pdm run build-ext
pdm run test --week 2 --day 4 -- -k rms
pdm run bench --solution tiny_llm --loader week2 \
  --week2-checkpoint rmsnorm --model qwen3-4b

Task 2: RoPE

Next replace tiny_llm_ext::rope, Week2RoPE::eval_cpu, and Week2RoPE::eval_gpu in src/extensions/src/week2_kernels.cpp, the week2_rope function in src/extensions/src/week2_kernels.metal, and FastRoPE.__call__ in src/tiny_llm/week2_kernels.py.

Implement RoPE for the model’s native B, L, H, D layout. A naive element kernel calculates the same angle, sine, and cosine separately for both members of every pair and again for every head. Instead, assign one thread a pair index and a block of four heads. Compute the angle once, then rotate both elements of that pair across the four heads:

angle = (batch_offset + token_position) * base ** (-pair / (dims / 2))
real' = real * cos(angle) - imag * sin(angle)
imag' = imag * cos(angle) + real * sin(angle)

Accept either one scalar offset or one offset per batch row in the Python wrapper. Normalize both cases to an int32 array before dispatch. Supporting per-batch offsets matters once requests at different decode positions share a batch.

Unlike a graph that builds position arrays, gathers sine and cosine values, splits the head, performs several element-by-element operations, and concatenates the result, this kernel reads each input pair and writes each rotated element directly. Reusing trigonometry across four heads is the key optimization. Use Metal’s fast::exp2, fast::sin, and fast::cos for the BF16 path. Normalize a batch’s offsets once in the model call, outside the layer loop, instead of rebuilding the same array in every layer.

Replace the Python mlx.core RoPE in the model you have already optimized. Test and measure that cumulative checkpoint before moving to SwiGLU:

pdm run test --week 2 --day 4 -- -k rope
pdm run bench --solution tiny_llm --loader week2 \
  --week2-checkpoint rope --model qwen3-4b

Task 3: SwiGLU

Finish the operator sequence with tiny_llm_ext::swiglu, Week2SwiGLU::eval_cpu, and Week2SwiGLU::eval_gpu in src/extensions/src/week2_kernels.cpp, the week2_swiglu function in src/extensions/src/week2_kernels.metal, and swiglu in src/tiny_llm/week2_kernels.py.

SwiGLU combines the gate and up branches:

output = (gate / (1 + exp(-gate))) * up

Implement it as one thread per element. That thread loads gate and up, evaluates SiLU with one exponential, multiplies the branches, and performs one output write. The Week 1 form is easier to inspect, but it describes abs, exp, division, selection, and multiplication as separate array operations. The fused kernel removes those intermediate tensors and dispatch boundaries.

Wire the fused expression into the model, then record the third checkpoint:

pdm run test --week 2 --day 4 -- -k swiglu
pdm run bench --solution tiny_llm --loader week2 \
  --week2-checkpoint swiglu --model qwen3-4b

Task 4: Verify the Cumulative Model

Now verify the cumulative switches in Qwen3ModelWeek2.__init__ and the call sites in Qwen3MultiHeadAttention.__call__ and Qwen3MLP.__call__. Task 4 is composition work: it uses the three functions from Tasks 1-3 and adds no new extension function.

Once all three kernels are exposed through C++ MLX primitives, run the complete test file. Keep qwen3_week1.py on its Week 1 Python operators, and leave the Week 2 interfaces reusable by the Week 3 serving model.

pdm run build-ext
pdm run test --week 2 --day 4

Use tolerance-based comparisons with the Python reference equations rather than bit-for-bit equality. Cover both scalar and per-batch RoPE offsets. When timing these lazy operations, call mx.eval inside every measured iteration.

The operator benchmark must also compare the same logical RoPE layout. Your RoPE kernel accepts the model-native B, L, H, D tensor. mx.fast.rope expects B, H, L, D, so transpose into that layout before the MLX call and transpose its result back afterward. Without those transposes, a one-token benchmark accidentally treats the head axis as sequence positions and the timing no longer measures an equivalent operation.

Benchmark Analysis: Decide Whether the Fused Kernels Stay

Measure the three cumulative checkpoints separately so their combined result cannot hide a regression:

pdm run bench-week2-progression --offline --solution tiny_llm --repeats 2 \
  --variant week2-quantized-matvec \
  --variant week2-rmsnorm --variant week2-rope --variant week2-swiglu \
  --variant mlx --model qwen3-4b \
  --input-len 128 --output-len 129 --warmup 2 --prefill-logits last

pdm run profile-week2-kernels --solution tiny_llm --model qwen3-4b \
  --case quantized-matvec:decode:128 --case swiglu:decode:128 \
  --case swiglu:prefill:128 --warmup 4 --iterations 12 \
  --json-output week2-day4-attribution.json

Record one cumulative result per operator, then use the attribution run to choose the next bottleneck. The complete campaign and reference attribution are in the performance appendix.

In the checked M4 Pro example, the fused kernels reduced the attributed normalization/position/activation category by 79.0%. Re-profiling then placed projections at 81.4% of decode attribution and 99.1% of 128-token prefill attribution. That is why the next core chapter is SIMD-Matrix Prefill, not a prescribed attention kernel. Repeat the same measurement on your machine and record the result that would falsify this next-change hypothesis.

If you want to continue without writing one of these kernels, keep its public course interface and delegate only that operator to the corresponding MLX implementation. This local substitution still exercises the cached Week 2 model and the other course-owned operators; --solution mlx does not.

🚧 Week 2 Day 5: SIMD-Matrix Prefill

Day 4 ends with a decision, not a predetermined kernel. Re-profile the fixed 128-token prefill and name the dominant category before changing code. On the checked M4 Pro run, projections accounted for 99.1% of attributed prefill time. That observation selects the matrix-shaped projection path for this chapter.

The swiglu checkpoint still uses Day 3’s correctness-first vanilla W4 matrix kernel when the activation has more than eight rows. You will replace that schedule with a cooperative BF16 SIMD-matrix kernel while preserving the same quantized-linear interface and the last-row-logits product boundary.

The checked numbers in this chapter are one example, not a performance gate. They come from Qwen3-4B on one 20-core M4 Pro running macOS 27 and MLX 0.32.0, with a 128-token prompt, 129 output tokens, two warmups, and two balanced fresh-process samples. Your device and crossover may differ.

Establish the Same-Workload Baseline

Start from the checkpoint you already have. Build the extension and run the focused gate before editing:

pdm run build-ext
pdm run test --week 2 --day 5

Freeze both baselines next. You will repeat these exact commands after the kernel change:

pdm run bench-week2-progression --offline --solution tiny_llm --repeats 2 \
  --variant week2-swiglu --variant week2-simd-matmul --variant mlx \
  --model qwen3-4b --input-len 128 --output-len 129 --warmup 2 \
  --prefill-logits last --json-output week2-day5-product.json

pdm run profile-week2-kernels --solution tiny_llm --model qwen3-4b \
  --case swiglu:prefill:128 --case simd-matmul:prefill:128 \
  --warmup 4 --iterations 12 \
  --json-output week2-day5-attribution.json

Keep the model, phase, token count, prompt rule, prefill-logit mode, warmups, and iteration count identical across the two checkpoints. Do not compare a new prefill kernel at one shape with an old result from another shape.

Task 1: Load One Quantized Tile Cooperatively

Open the three existing extension files; this task stays inside that surface:

src/extensions/src/cooperative_matrix.h
src/extensions/src/quantized_matmul.metal
src/extensions/src/quantized_matmul.cpp

Keep the operation fixed while you change its schedule:

$C = A W^{T},$

where A is BF16 and W is stored as packed W4 codes with one scale and bias per group of 128 values. The mathematical operation does not change. Only the matrix-shaped schedule changes.

Build a 32×32 output tile from 8×8 simdgroup_matrix fragments. SIMD groups cooperate on one 32-value slice of the reduction dimension at a time:

load a contiguous activation tile;
unpack the matching W4 codes and apply their scale and bias;
multiply the BF16 fragments while accumulating in FP32;
advance through the reduction dimension;
store only in-bounds output elements.

Keep the loader and fragment bookkeeping explicit. The course path does not call an MLX or Steel quantized-matmul implementation in place of this exercise. The existing Python equation remains the correctness oracle.

Task 2: Dispatch by Activation Shape

Retain Day 3’s SIMD matvec for M <= 8. Route larger activation matrices to the new tiled kernel and keep the vanilla kernel callable as a bring-up control. Validate dtype, contiguity, group size, bit width, and matrix dimensions at the extension boundary before encoding the GPU command.

The supplied starter dispatch is QuantizedMatmul::eval_gpu in src/extensions/src/quantized_matmul.cpp. Its matrix-shaped Metal entry is quantized_matmul_simdgroup_w4a16_g128 in src/extensions/src/quantized_matmul.metal; an equivalent solution may keep the public dispatch while choosing a different internal kernel name.

The checkpoint feature name is simd-matmul. It includes packed W4 projections and the three fused Day 4 operators. It does not include the optional decode-attention branch from Day 6.

If you want to continue without writing this custom schedule, preserve the course’s quantized_linear interface and route the matrix-shaped projection through mx.quantized_matmul. That is a local operator substitution, not a performance claim and not the separate --solution mlx model.

Task 3: Check Correctness in the Product

Once the new path is connected, get focused feedback before asking the full model to exercise the checkpoint:

pdm run build-ext
pdm run test --week 2 --day 5

pdm run main --solution tiny_llm --loader week2 \
  --week2-checkpoint simd-matmul --model qwen3-4b

An equivalent learner implementation may choose different helper names or a different correct tiling. The observable contract is the quantized-linear result, dtype and shape, checkpoint behavior, fallback behavior, and complete model output—not a private symbol or source-file layout.

Task 4: Re-profile and Decide

Now repeat the exact baseline commands, then close the loop in three sentences:

which operator category dominated the baseline prefill;
whether the candidate changed that category and the matched product phase;
what result would make you revert the candidate or test another schedule.

In the checked run, the SIMD schedule reduced attributed projection time by 86.4% and raised fixed-workload prefill from 106.44 to 721.60 tokens/s. Those large effects justify keeping it for that source tree and workload. They do not establish the same multiplier on another model, Apple GPU, prompt length, or software version.

Day 6 is an optional workload-conditioned operator lab. You may take that branch to study bounded decode attention, or continue directly to Day 7. Day 7 starts from this simd-matmul checkpoint either way.

🚧 Week 2 Day 6 (Optional): Workload-Conditioned Operator Lab

Day 5 restores the matrix-shaped projection path selected by the fixed 128-token prefill profile. Day 6 asks a different question: can a secondary operator earn a place for one named workload?

The supplied worked branch is bounded decode attention. It is useful practice with online softmax, but the checked fixed workload did not identify attention as the next dominant category. Treat this chapter as an optional experiment, not a prerequisite for Day 7 and not evidence of a universal bottleneck.

A successful pass ends with a bounded decision, even when the numbers do not support keeping the branch.

Choose the Workload Before the Operator

Write down the model, checkpoint, phase, prompt or context length, warmups, iterations, and comparison rule before editing code. Start from the Day 5 simd-matmul checkpoint and record the same workload for the candidate:

pdm run profile-week2-kernels --solution tiny_llm --model qwen3-4b \
  --case simd-matmul:decode:128 --case decode-attention:decode:128 \
  --warmup 4 --iterations 12 \
  --json-output week2-day6-attribution.json

The supplied branch uses decode-attention. An equivalent experiment on a different measurement-selected secondary category is valid if it preserves the public checkpoint and decision-record contract. The course grades observable behavior and reasoning, not a private file path, exact Metal symbol, device duration, or schedule choice.

Task 1: Preserve Bounded Decode-Attention Semantics

Use the supplied branch to make that decision concrete. The readable grouped-attention path materializes score and probability rows; for one query row, online softmax can combine the reduction and value-weighted sum without storing the full score row:

m = -infinity
l = 0
o = 0

for each key/value block:
    scores = q @ key_block.T * scale
    block_max = max(scores)
    new_m = max(m, block_max)
    alpha = exp(m - new_m)
    probabilities = exp(scores - new_m)
    l = alpha * l + sum(probabilities)
    o = alpha * o + probabilities @ value_block
    m = new_m

return o / l

Preserve grouped-query head mapping, dense-cache offsets, BF16 inputs and outputs, FP32 online-softmax state, scale, and the existing mask adapter. Keep an exact fallback for shapes outside the tested guard. Do not turn a short-context experiment into a claim about long-context or paged attention.

Task 2: Implement and Verify the Branch

Implement the smallest complete branch: replace only the existing fail-closed Day 6 learner surfaces. Keep the public attention interface stable so Week 3 can reuse it.

The supplied C++ surface is tiny_llm_ext::decode_attention, implemented by Week2DecodeAttention::eval_cpu and Week2DecodeAttention::eval_gpu in src/extensions/src/week2_kernels.cpp; its Metal entry is week2_decode_attention in src/extensions/src/week2_kernels.metal. The product calls it from Qwen3MultiHeadAttention.__call__ through decode_attention_custom. Equivalent internal organization is valid when it preserves this public behavior and fallback.

Build as soon as the branch is wired; run the focused check before the product path:

pdm run build-ext
pdm run test --week 2 --day 6

pdm run main --solution tiny_llm --loader week2 \
  --week2-checkpoint decode-attention --model qwen3-4b

Test supported shapes, grouped heads, offsets, and the exact fallback. A valid solution may use different helper names and internal organization; it must produce the same public attention behavior and preserve the fallback.

If you completed the old Week 2 Day 5 attention exercise before the course was reordered, keep that work. Complete the current Day 5 SIMD checkpoint first, then use this canonical optional Day 6 chapter and its commands to verify your retained attention implementation.

Task 3: Re-measure and Decide

A passing branch establishes correctness. The final decision comes from rerunning the same comparison:

Repeat the frozen workload and compare simd-matmul with decode-attention. Record:

the dominant category before the change;
the category and product effect you actually observed;
the exact context range and fallback you tested;
keep, reject, or inconclusive, plus the next falsifying experiment.

The checked M4 Pro result was equivocal: attributed attention changed from 0.837 ms to 0.831 ms (-0.75%), total attributed time rose 0.97%, and the separate two-sample product control showed decode rising from 74.34 to 76.50 tokens/s (+2.91%). That supports an inconclusive worked example, not a portable speedup claim. Your decision should follow your matched measurement.

Continue to Day 7 from simd-matmul. The split-k checkpoint intentionally excludes this optional attention branch.

🚧 Week 2 Day 7: Conditional Split-K and Final Decision

Day 5 leaves a reusable 32×32×32 SIMD-matrix projection and an exact unsplit fallback. Day 6 is an optional branch and is not inherited here: the split-k checkpoint contains the Day 5 SIMD path plus Split-K, without decode attention.

Begin with an under-filled short shape, then return to the fixed 128×129 product workload. Keep Split-K only for shapes where the same-workload evidence supports it.

Why Split the Reduction Dimension?

For

$C = A W^{T},$

the Day 5 grid spreads work across output rows and columns. When M is small and the Qwen projection width is narrow, it may launch too few independent threadgroups to fill the GPU. Split-K creates parallel work along the reduction dimension:

for each split s:
    partial[s] = A[:, k_start(s):k_end(s)] @ W[:, k_start(s):k_end(s)].T

C = sum(partial, axis=split)

Each split must align to the W4 group size, write to a disjoint partial plane, and accumulate its local dot product in FP32. A second kernel reduces the partial planes in FP32 and casts the final output to BF16.

That extra parallelism also adds a dispatch, a temporary buffer, and another memory pass. Split-K is therefore a shape-conditioned schedule, not an automatic upgrade.

Task 1: Freeze a Short-Shape Control

First verify the inherited Day 5 path and record a 32-token attribution pair:

pdm run build-ext
pdm run test --week 2 --day 7

pdm run profile-week2-kernels --solution tiny_llm --model qwen3-4b \
  --case simd-matmul:prefill:32 --case split-k:prefill:32 \
  --warmup 4 --iterations 12 \
  --json-output week2-day7-short-attribution.json

Capture the exact source, model, phase, token count, prompt rule, software, and device. Do not substitute a 128-token baseline for the 32-token candidate.

Task 2: Reuse the Day 5 Tile for Each Partition

Extend the existing quantized-matmul primitive rather than adding a parallel public operator. Reuse Day 5’s loader, W4 dequantization, and matrix fragments inside each aligned K partition. Validate that:

every split begins and ends on a group-of-128 boundary;
partial planes are disjoint and cover the full reduction exactly once;
edge rows and columns are masked before load or store;
accumulation and reduction remain FP32;
split_k <= 1 dispatches exactly to the Day 5 unsplit kernel.

The tests grade public results, dtype and shape, valid partitioning, and the exact fallback. They do not require a private helper name, Metal symbol, or a particular split-count formula.

Task 3: Make Dispatch Explicit

Expose the split-k checkpoint with an immutable feature set: packed W4, fused pointwise operators, SIMD prefill, no optional decode-attention branch, and Split-K only where its policy selects more than one partition.

Keep that public policy in QuantizedMatmul::eval_gpu; the supplied starter surface is src/extensions/src/quantized_matmul.cpp. Internal helper and Metal kernel names remain implementation choices.

Keep the policy small and inspectable. Static dispatch can demonstrate that a Split-K and reduction kernel exist, but it cannot prove higher occupancy or a product speedup. Those claims require measured evidence.

If you want to continue without Split-K, preserve split_k <= 1 and the Day 5 unsplit result. The chapter’s learning outcome is the conditional decision, not an unconditional custom-kernel win.

Task 4: Re-profile the Short Shape

Rerun the exact 32-token attribution command from Task 1 so the baseline and candidate differ only in schedule. In the checked M4 Pro example, Split-K reduced total attributed time by 4.87% and projection time by 5.01%. Its trace exposed only static Split-K and reduction dispatches; no timeline or counter tree materialized, so no occupancy improvement was inferred.

Write keep, reject, or inconclusive for the 32-token shape, then name the result that would reverse your decision. A sub-percent difference is not a strong conclusion without a larger sample.

Task 5: Close Week 2 at the Fixed Workload

Return to the Day 5 unsplit checkpoint and compare it with Day 7 at the same Qwen3-4B 128×129 product control used throughout the week:

pdm run bench-week2-progression --offline --solution tiny_llm --repeats 2 \
  --variant week2-simd-matmul --variant week2-split-k --variant mlx \
  --model qwen3-4b --input-len 128 --output-len 129 --warmup 2 \
  --prefill-logits last --json-output week2-day7-final.json

pdm run profile-week2-kernels --solution tiny_llm --model qwen3-4b \
  --case simd-matmul:prefill:128 --case split-k:prefill:128 \
  --warmup 4 --iterations 12 \
  --json-output week2-day7-final-attribution.json

On the checked two-sample product control, prefill changed from 721.60 to 718.36 tokens/s (-0.45%) and decode changed by +0.14%. That supports rejecting Split-K for this fixed 128-token product workload while conditionally retaining the short-shape experiment. It does not establish a portable crossover.

Finish with the week’s decision ledger:

Step	Evidence that selected it	Same-workload result	Decision and falsifier
KV cache	Full-prefix recomputation	Matched Week 1 versus cache	Your observation
Packed W4	Cached decode attribution	Repeated decode product and attribution	Your observation
Fused pointwise	Post-W4 re-profile	Repeated decode product and attribution	Your observation
SIMD prefill	Day 4 128-token prefill profile	Repeated prefill product and attribution	Your observation
Optional operator lab	Explicit secondary workload	Before/after/fallback record	`keep`, `reject`, `inconclusive`, or skipped
Split-K	Under-filled 32-token projection	Short control plus fixed 128×129 control	One decision per shape

Close the week with the causal story: what dominated, what changed, what the identical remeasurement showed, and what you chose not to claim.

🚧 Week 3: Build a Mini vLLM

🚧 This overview and chapters carrying the same marker are under review and may change.

Week 3 turns the optimized single-request model into a multi-request serving engine. Students add scheduling, request-owned cache state, shared page pools, and the runtime metadata needed to read noncontiguous K/V directly. The final model uses one page-aware attention interface. A correct direct page-walking implementation may serve every query shape; the completed reference adds a tiled schedule for the supported BF16 long-prefill hot case.

Week 2’s course-owned quantized projections remain the inspectable endpoint of that week’s kernel lessons. Week 3 deliberately switches dense-model projections to mx.quantized_matmul at model construction, while retaining the course-owned normalization, activation, cache, attention, paging, and scheduler paths. This keeps Week 3 focused on serving-system mechanisms rather than carrying the teaching kernel’s projection cost through every benchmark.

What We’ll Cover

Continuous batching and request-slot reuse
Chunked prefill and scheduler fairness
Paged KV storage and page-walking attention
Paged FlashAttention for long prefill
Optional Day 6 Mixture-of-Experts model support
Optional Day 7 speculative decoding over rewindable caches

Day 1 introduces that projection seam and batches independent request states. Day 2 splits long prefills so they cannot monopolize the scheduler. Day 3 replaces a growing dense cache with fixed-size pages while retaining a dense-gather compatibility path. Day 4 removes that gather by teaching attention to walk the page table directly with a correctness-first implementation; one kernel may serve all query shapes. Day 5 then replaces the supported BF16 long-prefill hot case with a tiled schedule. Page translation is therefore introduced before it is optimized.

These five days form the required path in your solution. The final model in your solution runs paged FlashAttention for supported BF16 long prefill and retains correct direct fallbacks for short queries and generic shapes. Every schedule reads the same page pool through the same block-table interface; none rebuilds dense K/V.

Paged attention is not an automatic single-request latency win. The checked trace measures lower KV storage, page reuse, incremental growth, and batching; page-table indirection can make one request slower. It does not establish an admission-capacity gain without a memory-capped sweep. Each chapter ends with a focused measurement, while the performance appendix records the matched chapter-by-chapter results.

Optional Day 6 adds MoE model support independently of the cache and scheduler. Optional Day 7 then adds speculative decoding, whose rejection path needs a precise cache rewind operation and whose multi-token verification needs the page-aware long-query path. Neither extension is required to complete Week 3.

🚧 Week 3 Day 1: Continuous Batching

You begin with the completed Week 2 single-request model: multi-offset RoPE and causal masking already have stable interfaces, and each request can own a dense KV cache. The Day 1 starter leaves four learner-owned slices behind those interfaces:

dense batch assembly and masking in BatchingKvCache;
mlx_quantized_linear plus the per-weight selector and the explicit dispatch_week3_batch_model factory;
selector propagation through Qwen3ModelWeek2; and
Request.try_prefill plus the request-admission/decode loop in batch_generate.

Complete and test those slices in order. The resulting continuous batch keeps several active requests on the device and replaces each request as soon as it finishes. Only quantized projections cross the Week 3 MLX seam; normalization, activation, RoPE, cache state, attention, and scheduling remain course-owned.

So far, each generation loop has processed only one request. That may not provide enough work to use the device efficiently, so Day 1 decodes several requests in each model call.

A static batch could select five prompts and run them together until every request finishes. However, generated sequences have different lengths. If four requests finish quickly while the fifth continues, most of the batch remains idle and queued requests cannot start.

Continuous batching instead sets a maximum number of active decode requests. When one finishes, the scheduler assigns its batch slot and KV-cache entry to a waiting request. This keeps the decode batch populated whenever work is queued.

The scheduler must also interleave prefill and decode work. We will use a simple policy: advance one pending prefill, then decode one token for every active request.

while requests_in_queue_or_in_progress:
    if prefill_request is not None:
        prefill_request.try_prefill()  # Day 1 processes the complete prompt
        if prefill_request.ready:
            if kv_cache.try_add(prefill_request):
                prefill_request = next(requests)
    if active_requests:
        tokens = decode(model, kv_cache)
        for request, token in zip(active_requests, tokens):
            request.append(token)

A complete prompt is admitted in one call on Day 1. This makes the scheduling policy easy to inspect and exposes an important limitation: one long prefill can delay every active request’s next decode step. Day 2 will add a bounded prefill budget to solve that fairness problem.

Task 1: Reuse RoPE and Causal Masking for Batched Requests

src/tiny_llm/week2_kernels.py::FastRoPE  (reuse unchanged)
src/tiny_llm/attention.py::causal_mask   (reuse unchanged)

Continuous batching requires one RoPE offset per batch element and a causal mask whose query and source lengths may differ. Verify those two Week 2 interfaces before adding the scheduler so the serving layer can use one model contract for every request position.

Verify multi-offset RoPE and the rectangular causal mask with:

pdm run test --week 3 --day 1 -- -k task_1

Task 2: Batch KV Cache

src/tiny_llm/kv_cache.py::BatchingKvCache

BatchingKvCache holds one request cache per decode slot. Because requests may have different sequence lengths, it must combine their keys and values into dense tensors and construct a matching B x 1 x L x S mask.

S = max(S_i across active requests)
L = mask_length (input parameter)
request_keys: H, S_i, D
request_values: H, S_i, D
batched_keys: B, H, S, D
batched_values: B, H, S, D
mask: B, 1, L, S

Right-align each active request in the common S dimension. The leading positions remain zero and masked out. Inactive slots remain fully masked.

keys_i, values_i = request_cache[i]
batched_keys[i, :, (S - S_i):S, :] = keys_i
batched_values[i, :, (S - S_i):S, :] = values_i
mask[i, :, 0:L, (S - S_i):S] = causal_mask(L, S_i)

You can verify your solution by running:

pdm run test --week 3 --day 1 -- -k task_2

Task 3: Add the Week 3 Projection Seam

src/tiny_llm/quantize.py::mlx_quantized_linear
src/tiny_llm/qwen3_week2.py::Qwen3ModelWeek2.__init__
src/tiny_llm/models.py::dispatch_week3_batch_model

Week 2 ends with a course-owned quantized matmul so you can inspect its loader, SIMD-matrix operations, and Split-K policy. Week 3 teaches serving mechanisms, so it should not make every cache and scheduler measurement depend on that teaching kernel’s remaining projection overhead.

Add a per-weight use_mlx_quantized_linear selector whose default remains False, preserving every Week 2 checkpoint. When selected, quantized_linear should call mx.quantized_matmul with the same packed weight, scales, biases, group size, bits, and transposed-weight convention. Then implement dispatch_week3_batch_model as the explicit construction seam: it builds the completed dense-cache Week 2 model with that selector enabled.

Call this batch-ready model with several requests, one offset per batch element, and the mask returned by BatchingKvCache. Exercise requests joining and leaving at different positions. Only quantized projections cross the MLX seam; normalization, RoPE, activation, attention, cache state, and scheduling remain in your solution. The model remains request-agnostic; slot ownership and lifecycle belong to the cache and scheduler.

Verify the projection seam and its batch-model factory with:

pdm run test --week 3 --day 1 -- -k task_3

Task 4: Batch Generate

src/tiny_llm/batch.py

First implement Request.try_prefill by prefilling the complete prompt in one call. The visible prefill_max_step input is reserved for Day 2; on Day 1, give this call a budget that covers the complete remaining prompt. Then complete the scheduler in batch_generate: move finished prefills into idle decode slots, collect the next token and offset for each slot, and remove requests that reach EOS or max_seq_len. A prompt longer than max_seq_len must fail before model or cache work, and a generated token that would cross the limit is not emitted.

Results are returned in completion order, because short requests can leave the batch before earlier long requests. Each result’s prompt_idx maps it back to the original input position; do not reorder completed results into input order.

Use the supplied scheduler checkpoint for full prefill, admission and slot reuse, immediate EOS, maximum-length termination, completion-order results, and their original prompt_idx values:

pdm run test --week 3 --day 1 -- -k task_4

Then run the complete scheduler against the real model:

pdm run batch-main --solution tiny_llm --loader week2

By default, batch-main uses Qwen3-0.6B with a batch size of five and a fixed prompt set. Treat it as a product smoke: watch requests enter, decode, finish, and release their slots. Its shuffled prompt order and cumulative wall-clock display are not the source of a decode-gap measurement.

Day 2’s deterministic bench-chunked-prefill runner owns that comparison. Its 512-token budget processes every prompt in the checked 64–512-token trace in one chunk, so that row is the reproducible Day 1 control. The runner records the exact token ids, output budget, seed, process order, and decode-completion gaps before Day 2 changes the prefill budget.

🚧 Week 3 Day 2: Chunked Prefill

🚧 This chapter is under review and may change.

A long prompt can monopolize the device while active decode requests wait for their next token. Chunked prefill gives each scheduler iteration a prompt-token budget, limiting how long decode work can be delayed.

The scheduler policy becomes:

admit at most prefill_max_step prompt tokens
decode one token for every active request
repeat until the queue and active batch are empty

Task 1: Bound Prefill Work

Update Request.try_prefill in src/tiny_llm/batch.py to select one prompt slice, call the model with the slice’s absolute offset, and mark the request ready only after the full prompt has been processed.

for start in range(0, len(prompt_tokens), prefill_max_step):
    chunk = prompt_tokens[start : start + prefill_max_step]
    model(chunk, offset=start, cache=cache)

The final chunk may be smaller than the configured budget. Test prompts shorter than one chunk, exactly one chunk, and one token longer than a chunk.

Task 2: Build Rectangular Causal Masks

When a cache already holds S - L tokens and a chunk contributes L new tokens, the mask is L x S. Every query can attend to the old prefix and to earlier positions in its own chunk.

For a five-token prefix and a three-token chunk, the mask is 3 x 8:

0  0  0  0  0  0  -inf  -inf
0  0  0  0  0  0     0  -inf
0  0  0  0  0  0     0     0

Use the absolute cache offset for RoPE and S - L as the causal diagonal offset. Compare chunked prefill logits with one-shot prefill logits.

Task 3: Materialize Between Chunks

MLX is lazy. Extending an unevaluated cache repeatedly creates a long graph and can grow memory usage. Call each layer cache’s materialize() hook after every chunk so the next scheduler iteration starts from materialized state. A dense cache evaluates its key/value tuple; a paged cache evaluates the page pool storage without first gathering it into a dense tensor.

The hook is part of the cache lifecycle rather than the scheduler’s storage logic. This lets the scheduler use dense and paged caches without inspecting their internal representation.

Task 4: Measure the Fairness Tradeoff

Run the same request trace with several prefill_max_step values. Report total throughput and the longest interval between consecutive decode steps. Smaller chunks usually improve fairness but add scheduler and launch overhead; choose a default from the measured tradeoff rather than treating one chunk size as universal.

pdm run test --week 3 --day 2
pdm run batch-main --solution tiny_llm --loader week2

pdm run bench-chunked-prefill --solution tiny_llm --offline --model qwen3-0.6b \
  --prefill-steps 32 128 512 --num-seqs 8 --batch-size 4 \
  --min-input-len 64 --max-input-len 512 \
  --min-output-len 32 --max-output-len 32 \
  --warmup 1 --repeats 4 --cooldown-seconds 1 \
  --json-output benchmark_results/task367-final-main/raw/learner-week3-chunked-prefill.json

This command measures your tiny_llm solution. The checked reference trace below uses seed 0 and the same 32-token output budget for every request. Each chunk size runs twice in forward order and twice in reverse order in fresh processes. Every row uses the same canonical Week 3 MLX quantized-projection seam and the same course-owned scheduler, dense cache, and attention code; only the prefill budget changes. The JSON stores every prompt token id, the per-request output budget, and their canonical SHA-256 checksum. To reproduce the checked rows separately, rerun the command with --solution ref and --json-output benchmark_results/task367-final-main/raw/week3-chunked-prefill-final-main.json.

A decode-completion gap is the wall-clock interval between two consecutive synchronized decode calls while at least one decode request remains active. It therefore includes intervening prefill and scheduler work; idle time with no decode request is excluded. On the measured M4 Pro, the four-process medians were:

Prefill budget	Output tok/s	Prefill tok/s	Decode tok/s	Requests/s	Decode step p95	Decode gap p95 / max
32	105.23	2,549.62	181.77	3.288	15.82 ms	30.01 / 52.62 ms
128	153.82	4,215.12	242.23	4.807	17.79 ms	45.36 / 53.76 ms
512	170.46	4,769.14	262.01	5.327	17.11 ms	73.56 / 119.90 ms

The 512-token row is the full-prompt Day 1 control for this trace. Relative to that row, the 128-token budget gives up 9.8% output throughput while reducing the p95 completion gap by 38.3% and the maximum gap by 55.2%. The 32-token budget reduces the p95 gap further but gives up substantially more throughput. The course uses 128 as a measured compromise for this workload, not as a universal optimum. The ledger at benchmark_results/task367-final-main/task367-final-main-benchmark-ledger.md keeps this final-main absolute result separate from task #360’s causal projection-seam ablation.

🚧 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

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:

Appending a token usually updates only the current tail page in the pool.
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:

look at that layer cache’s last page
if there is room, append only the new slice into the tail page
otherwise allocate a new page and continue writing
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.

The learner extension already contains the C++ and Metal source files, CMake entries, header declaration, and Python binding. Replace the paged_cache_update stubs in src/extensions/src/paged_attention.cpp and src/extensions/src/paged_attention.metal; do not create or register a second primitive.

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:

compute one-token k and v
check whether the tail page still has space
write into that page if possible
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:

write new K/V into the layer cache’s tail page or newly allocated pages,
gather the layer cache’s pages back into dense K/V,
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:

What page size should this repo use for teaching?
How do we represent the free-page allocator?
How do we prove that paged storage reconstructs the same logical KV as TinyKvFullCache?
How do request cache handles share a layer pool while keeping their own page metadata?
When do we materialize page writes to avoid MLX lazy-graph growth?
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.

Carry forward Day 1’s projection boundary when constructing the dense Qwen3 model: quantized projections use the mx.quantized_matmul seam, while the embedding lookup, normalization, activation, RoPE, cache, and attention paths remain course-owned. Keep use_mlx_quantized_linear=True as the dense Week 3 default and retain the opt-out only as a benchmark/correctness ablation. The optional MoE extension keeps its separately taught router/expert projection contract; do not silently broaden this dense-model seam into that chapter.

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 --batch-decode --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

🚧 Week 3 Day 4: Direct Paged Attention

🚧 This chapter is under review and may change.

In this chapter, you will make direct paged attention handle decode and prefill in both float32 and BF16. The scheduler passes request-local block tables and context lengths to the operator, which reads K/V from the shared layer pool without gathering a dense batch first. One correct page-walking kernel may serve every query shape; separate decode and prefill kernels are an optimization choice. Day 5 replaces the supported BF16 long-prefill hot case with a tiled implementation.

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:

Paged KV cache KV is stored in fixed-size pages.
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:

append the token’s K/V to the current tail page,
allocate a new page only if the tail page is full,
update the current layer cache’s context_len,
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.

Establish the Direct Page-Walking Boundary

First make every supported query shape correct through the same direct paged boundary. You may reuse one kernel for decode and prefill. If you split the workloads, a single tile shape is unlikely to keep the GPU busy for both a one-token query and a long prompt. Use these design rules:

Preserve the Week 2 BF16 model boundary and reuse its internal accumulation policy unchanged.
For short queries, expose parallelism across the cached context. Do not reserve most of a threadgroup for query rows that do not exist.
For prefill, begin with a direct page-walking schedule whose address calculation is easy to validate.
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.

The reference solution uses this optional split. Treat its threshold as a value to verify on your hardware, not part of the public checkpoint:

Shape	Dispatch in your solution	Work decomposition
`L <= 8`	Vector paged decode	One threadgroup per query row; 32 SIMD groups stride over the context and merge partial `(max, sum, output)` states.
`L > 8`	Scalar direct paged prefill	Walk logical K/V tiles through the block table and keep the schedule deliberately inspectable. Day 5 optimizes the supported BF16 hot case.

If you make a shape decision, put it at the extension boundary rather than converting inputs or falling back to dense attention in Python. Keep the model-facing paged_attention API unchanged. Public Day 4 tests grade only its shape, dtype, validation, page addressing, masking, and numerical behavior; they do not require this split or any kernel name.

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 operation with online softmax. It may use one kernel or several internal schedules:

use block_table[b] to find the physical pages for request b,
use context_lens[b] to ignore unused tail capacity,
visit K/V in small tiles instead of materializing dense K/V,
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 5. 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.

For performance, one-token decode benefits from a different work decomposition. A 64-row prefill tile would leave almost every query row idle, so the reference dispatches short queries to a vector-oriented kernel that partitions the context across SIMD groups and merges their partial online-softmax states. This split is an optimization, not a public correctness requirement.

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:

paged storage
block_table / context_lens plumbing
correctness-first page-walking GPU attention
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:

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.
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.
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.
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.
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 the stable starter boundary:

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;
one or more kernels in src/extensions/src/paged_attention.metal that implement the same public behavior. The starter names paged_attention_decode and paged_attention_scalar_f32 mirror the reference design, but they are not required by the tests.

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.

The reference solution implements two correctness-first GPU dispatches:

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.
For longer queries, it assigns query rows to a direct page-walking schedule and resolves 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. Its direct schedule handles both float32 and BF16, using float accumulators for BF16’s dot products, online-softmax state, and output accumulation.

You may instead reuse one correct direct kernel for every query length. Favor inspectable page ownership over the final tiled performance schedule; Day 5 is where the supported BF16 long-prefill region receives a dedicated performance implementation.

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.

Rebuild the extension, then use the BF16 long-prefill checkpoint as the first feedback loop for this task:

pdm run build-ext
pdm run test --week 3 --day 4 -- -k task_2_bfloat16_long_prefill

The focused checkpoint constructs valid page metadata directly. It checks a nine-token prefill and a longer multi-page prefill with noncontiguous physical pages, poisoned unused tail slots, and causal prefix masking. It does not require the quantized embedding, model dispatch, continuous batching, or the Day 5 kernel. Kernel names and implementation structure are not part of the test contract; only the public numerical, metadata, and dtype behavior is.

For the reference 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.

Day 4 must work without a future tiled/cooperative/MMA kernel. A dedicated prefill kernel is optional here: reusing the direct decode implementation for long queries is valid when it preserves the public behavior. Day 5 may replace only the internal supported BF16 long-query region while preserving the same API, generic fallback, and page-table semantics.

Your solution’s boundary

The walkthrough keeps page translation and online softmax in a course-owned Metal operation so you can inspect them. The public checkpoint does not grade an internal helper, tile, symbol, or library choice: a behaviorally equivalent implementation or external low-level library is valid. To preserve the chapter’s systems outcome, the completed operator still consumes page storage and its block table directly rather than rebuilding dense K/V in Python. MLX SDPA remains a useful 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. The reference uses scalar prefill and vector decode schedules, but one correct direct schedule may serve both. Neither choice changes cache dtype or gathers a dense K/V tensor.

This creates the Day 4 routing policy:

every query shape -> correct direct page-walking attention
optional split -> scalar prefill and vector decode

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:

the dense Week 2 attention path in your solution, including any required K/V gather,
your direct paged-attention path,
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 --solution tiny_llm --offline --contexts 128 1024 \
  --page-size 128 --warmup 5 --iterations 60 --repeats 4 \
  --cooldown-seconds 1 \
  --json-output benchmark_results/task367-final-main/raw/learner-week3-attention.json

This command measures your tiny_llm operators. The checked reference values below are medians of four balanced fresh-process medians, with 60 synchronized calls after five warmups per process:

To reproduce those checked rows separately, rerun the command with --solution ref and --json-output benchmark_results/task367-final-main/raw/week3-attention-final-main.json.

Context	Dense + gather	Direct paged	MLX fused
128	201.26 us	228.58 us	188.79 us
1,024	468.39 us	299.14 us	250.04 us

Direct traversal is 13.6% slower than dense-plus-gather at 128 tokens, but 36.1% faster at 1,024 tokens. MLX remains faster at both shapes. The checked BF16 outputs match the readable dense equation within 0.00439453125 at S=128 and 0.001953125 at S=1,024. This operator benchmark contains no model projection, so it isolates the attention paths directly.

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:

assign the lanes of a SIMD group to adjacent elements of a head so K/V loads can be coalesced,
load a page-table entry once and reuse it for all positions in that page,
keep the query and online-softmax state in registers across page tiles,
combine partial dot products with SIMD reductions instead of threadgroup scratch memory and repeated barriers,
specialize the L = 1 decode case so it does not carry prefill control flow,
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 --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

This command compares your Week 2 dense batch reconstruction, Week 3 paged storage with the dense-gather compatibility path, and Week 3 direct paged attention. All three course rows use the same MLX quantized-projection seam; they differ in KV representation and attention path. The runner 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 672.68 prefill tok/s, 46.36 output tok/s, 105.01 decode tok/s, and 0.618 requests/s. Its synchronized decode calls take 28.97/36.78/63.04 ms at median/p95/max; the completion gaps, which include intervening scheduler and prefill work, are 30.16/222.18/239.49 ms.

These are cumulative system results, not an isolated Day 4 kernel speedup. The ledger at benchmark_results/task367-final-main/task367-final-main-benchmark-ledger.md records the fixed trace, balanced process order, and denominator boundary. Reproduce that checked reference trace separately with --solution ref and --json-output benchmark_results/task367-final-main/raw/week3-serving-final-main.json.

pdm run test --week 3 --day 4

🚧 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:

The Week 2 decode-attention lab introduced the online-softmax recurrence.
The Week 2 SIMD-matrix prefill lesson 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:

it resolves each K/V tile through block_table instead of gathering a dense cache;
it keeps only a query tile, one K/V tile, and online-softmax state on chip;
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:

Query shape	Schedule
`L <= 8`	Keep the Day 4 vector paged-decode kernel.
`L > 8`, BF16, `D == 128`	Use the tiled paged FlashAttention kernel.
Every other supported shape	Keep a correct direct page-walking fallback.

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:

compute Q @ Kᵀ with the Week 2 SIMD-matrix fragments;
apply scale and causal bounds;
merge the tile maximum into the running maximum;
rescale the previous sum and output accumulator;
compute exponentials for the current scores and update the running sum;
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:

compare Day 4 page-walking attention with the readable equation written with mlx.core;
compare paged FlashAttention with the Day 4 path;
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:

Storage / attention path	Prefill tok/s	Output tok/s	Decode tok/s	Requests/s	Peak KV	Avoidable KV copy
Dense growth and reconstruction	711.18	35.23	57.59	0.469	1,096 MiB	209,532 MiB
Paged storage + dense gather	725.46	41.64	78.53	0.555	not a total peak	103,445 MiB
Direct paged attention	672.68	46.36	105.01	0.618	576 MiB	504 MiB

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.

8K static checkpoint	Prefill tok/s	Decode tok/s
Week 2 course-owned projections	323.96	17.73
Week 3 seam + course paged path	463.69	27.42
Full MLX	639.73	28.37

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.

🚧 Week 3 Day 6 (Optional): Mixture of Experts

🚧 This chapter is under review and may change. In this chapter, we will implement the feed-forward shape of Mixture of Experts, or MoE, for the Qwen3 family.

This extension is optional. It changes the model’s feed-forward layers but not the scheduler, paged cache, or attention contract, so students can complete the Week 3 serving engine without it.

Start from the learner checkpoint and keep it running as you complete each task:

pdm run test --week 3 --day 6

The focused suite uses small generated weights. It checks the expert mapping, router, sparse block, and a complete mixed dense/sparse model without downloading a model. The 30B commands at the end are optional product smokes after this loop is green.

So far, every transformer block in tiny-llm has used the same dense Qwen3 MLP:

x -> gate_proj
x -> up_proj
SiLU(gate_proj(x)) * up_proj(x) -> down_proj

That is a SwiGLU MLP. Every token visits the same weights.

MoE changes only the feed-forward half of the transformer block. Instead of one dense MLP, the model owns many expert MLPs. A small router chooses which experts each token should use:

token hidden state -> router -> top-k experts -> weighted expert outputs

The attention path does not change. KV cache does not change. The sparse work is inside the MLP half of the block.

Readings

Dense MLP vs MoE MLP

The dense Qwen3 MLP from Week 1 has one set of weights:

w_gate: hidden_dim, dim
w_up:   hidden_dim, dim
w_down: dim, hidden_dim

A Qwen3-MoE sparse block has a bank of those weights:

expert_gate: num_experts, moe_hidden_dim, dim
expert_up:   num_experts, moe_hidden_dim, dim
expert_down: num_experts, dim, moe_hidden_dim

The router produces one score per expert:

router_logits: B, L, num_experts
router_probs:  softmax(router_logits)

Then the model picks num_experts_per_tok experts for each token:

expert_ids:    B, L, num_experts_per_tok
expert_scores: B, L, num_experts_per_tok

For each token, only those selected experts run. Their outputs are weighted and summed:

output[token] = sum(score_i * expert_i(token))

That is the central MoE idea: the model can contain many parameters, but each token activates only a small subset of them.

Qwen3-MoE Shape

Qwen3-MoE keeps the same attention structure as Qwen3, including QK norm, GQA, RoPE, and the same KV cache interface. It replaces some dense MLP layers with a sparse MoE block.

The useful pieces are:

gate: a router linear layer from hidden size to num_experts
switch_mlp: many SwiGLU experts with moe_intermediate_size
num_experts_per_tok: how many experts a token uses
norm_topk_prob: whether selected expert scores are renormalized
decoder_sparse_step and mlp_only_layers: which layers are sparse vs dense

There is no shared expert in the Qwen3-MoE block we are following. The sparse feed-forward output is just the weighted top-k expert mixture.

Grouped Expert Linear

MLX does not give us a single high-level MoE block in mlx.nn. It does have a lower-level primitive, mx.gather_qmm, that performs quantized matrix multiplication while selecting a different matrix for each row. Use that primitive, an equivalent low-level library operation, or your own grouped kernel to implement the public grouped_expert_linear relation.

For MoE, that means:

token rows:  N, D
expert ids:  N
weights:     E, O, D packed as 4-bit QuantizedWeights
output:      N, O

The row with expert_ids[i] = e should multiply by weights[e].

The helper may sort rows by expert id for locality, but it must restore the original token/expert order before returning. The supplied test observes that mapping and the numerical result; it does not grade the library, kernel, or sorting strategy you choose.

Router Step

The router is just a quantized linear layer:

router_logits = quantized_linear(x, w_router)
router_probs = softmax(router_logits, axis=-1)

For a batch of tokens:

x:             B, L, D
router_logits: B, L, E
router_probs:  B, L, E

where E = num_experts.

Qwen3-MoE then uses top-k selection:

expert_ids = argpartition(-router_probs, k)[:k]
expert_scores = take_along_axis(router_probs, expert_ids)

If norm_topk_prob is true, renormalize expert_scores so the selected scores sum to 1 for each token.

Expert Step

Each expert is the same kind of SwiGLU MLP we already know:

expert(x) = down_proj(SiLU(gate_proj(x)) * up_proj(x))

The implementation should build token-expert jobs, group them by expert, and run the expert projections with grouped_expert_linear:

selected expert ids -> expanded token-expert rows
expanded rows -> sort/group by expert id
grouped expert rows -> grouped gate/up projection
SiLU(gate) * up -> grouped down projection
restore original token/top-k order -> weighted sum

The reorder is part of the model implementation. It keeps all token rows for the same expert contiguous so the expert bank can be applied with grouped matrix multiplication.

Task 1: Grouped Expert Linear

src/tiny_llm/moe.py

Implement grouped_expert_linear. It accepts:

x:           ..., D
w_experts:   packed QuantizedWeights for num_experts, output_dim, D
expert_ids:  ...

It returns:

out:         ..., output_dim

Each row uses the expert selected by the matching row in expert_ids:

out[row] = x[row] @ dequantize(w_experts[expert_ids[row]]).T

One direct implementation is:

1. flatten the leading token/expert dimensions,
2. sort rows by expert id and retain the inverse order,
3. call mx.gather_qmm with the matching expert ids,
4. restore the original order and shape.

This task maps to the same idea as QuantizedSwitchLinear in mlx-lm: each token row uses a different packed expert matrix, and the expert ids choose the right matrix. If you want a deeper systems exercise, you may instead add a repository-native grouped C++/Metal operation, binding, and build registration. That native kernel is an optional stretch goal, not part of the Day 6 checkpoint.

Task 2: Router Top-k

src/tiny_llm/moe.py

Modify the existing route_topk function in this file.

Implement route_topk. It accepts hidden states and router weights, then returns:

router probabilities
selected expert ids
selected expert scores

Use quantized_linear and softmax. Use mx.argpartition to select the top num_experts_per_tok experts, then mx.take_along_axis to gather their scores.

Keep norm_topk_prob as an argument because Qwen3-MoE stores this behavior in the model config.

Task 3: Qwen3 Sparse MoE Block

src/tiny_llm/moe.py

Modify Moe.__init__ and Moe.__call__, composing the grouped_expert_linear and route_topk functions from Tasks 1-2.

Implement Moe by composing Task 1 and Task 2:

hidden states -> route_topk
hidden states + expert ids -> grouped gate projection
hidden states + expert ids -> grouped up projection
SiLU(gate) * up -> grouped down projection
weighted sum over num_experts_per_tok

This completes the Qwen3-MoE sparse feed-forward block. There is no shared expert branch in this block.

Task 4: Integrate Qwen3-MoE Layers

src/tiny_llm/qwen3_week3.py
src/tiny_llm/models.py

Modify is_qwen3_moe_sparse_layer and Qwen3ModelWeek3.__init__ in src/tiny_llm/qwen3_week3.py. The existing dispatch_model public path in src/tiny_llm/models.py already recognizes the Qwen3-MoE alias; exercise that path in the checkpoint rather than modifying it.

Add a Qwen3-MoE loader path that reuses the Week 3 Qwen3 attention and paged KV cache behavior, but swaps selected block MLPs for Moe.

The model wrapper should:

keep Qwen3 attention unchanged,
use regular Qwen3MLP for mlp_only_layers,
use Moe for sparse layers selected by decoder_sparse_step,
load router and expert weights as QuantizedWeights from the Qwen3-MoE MLX model,
preserve the same decode call shape:

logits = model(tokens, offset, cache)

No scheduler API change in src/tiny_llm/batch.py is required for correctness.

The focused command from the chapter opening validates this task without a model download. Its final case constructs a two-layer quantized fixture containing one dense mlp_only_layers layer and one sparse layer, dispatches it through the public model alias, runs it with a real KV cache, and compares normalized logits with MLX.

After that passes, you may smoke-test the same public path through the normal generation entrypoints. This optional check downloads a large model and is not required for checkpoint feedback:

hf download Qwen/Qwen3-30B-A3B-MLX-4bit

pdm run main --solution tiny_llm --loader week3 --model qwen3-30b-a3b \
  --prompt "Give me a short introduction to mixture of experts."

pdm run batch-main --solution tiny_llm --loader week3 --model qwen3-30b-a3b \
  --batch-size 2 --prefill-step 16

🚧 Week 3 Day 7 (Optional): Speculative Decoding

🚧 This optional chapter is under review and may change.

Speculative decoding uses a smaller draft model to propose several tokens, then asks the target model to verify them in one call. Accepted draft tokens reduce the number of target-model decode steps without changing the target distribution.

This checkpoint in your solution implements greedy speculative decoding: draft tokens are accepted while they match the target model’s greedy tokens. Extending the same loop to sampling requires the probability-correct acceptance and residual sampling rules; simple token equality is not enough.

Start with the model-free checkpoint. It uses generated token streams and cache objects, so it gives useful feedback without downloading either model:

pdm run test --week 3 --day 7 -- -k "proposal_length or target_only"

Objectives

By the end of this chapter, you should be able to:

generate a bounded proposal with a smaller draft model;
verify several proposed positions in one target-model call;
accept the matching prefix and recover at the first mismatch;
rewind dense and paged caches without corrupting offsets; and
decide whether acceptance rate offsets draft and verification overhead.

Prerequisites

Complete Week 2 cached generation for both the draft and target models.
Complete the Week 3 paged cache and page-aware attention path.
Use compatible tokenizers for the two models. A shared token id must represent the same text in both vocabularies. Validate that contract before either model runs; mismatched prompt encodings, EOS ids, or vocabularies must fail closed.

This extension comes after paged attention for two concrete reasons. Rejected draft tokens must release pages and repair the valid tail length, and verifying several proposed tokens at once is a long-query attention call over the paged prefix. The paged-cache lifecycle and page-aware long-query operator therefore form the stable interface on which speculative decoding is built.

Task 1: Reuse the Cache-Rewind Contract

Day 3 already added rewind(n) to the common KV-cache interface. Verify that prerequisite before building the speculative loop. A dense cache removes the last n logical positions. A paged cache must also return pages that become unused and shorten the valid prefix of the new tail page. Both implementations accept zero through the current logical length and reject other values before changing the cache.

Verify zero-length rewind, a rewind within one page, a rewind across page boundaries, and a full rewind:

pdm run test --week 3 --day 3 -- -k rewind

Task 2: Produce a Bounded Draft

Choose a small proposal length such as four. Starting from the last accepted token, run the draft model one token at a time and retain both the proposed tokens and the draft-cache offset. Stop early at EOS.

Keep the proposal length configurable. A longer proposal reduces target calls only when the acceptance rate remains high enough to repay the extra draft work. Use a default of four and treat zero as an explicit target-only fallback:

speculative_generate(
    draft_model,
    model,
    draft_tokenizer,
    tokenizer,
    prompt,
    proposal_length=4,
    max_tokens=256,
)

max_tokens bounds newly emitted non-EOS tokens. Zero returns without running a model or creating a cache, invalid values fail before either model runs, and EOS may stop generation earlier. Keep the explicit default so existing direct callers remain source-compatible.

After the target-only fallback and bounded proposal work, rerun the focused cases:

pdm run test --week 3 --day 7 -- -k "proposal_length or target_only or budget"

Task 3: Verify in One Target Call

Pass the last accepted token followed by the draft proposal to the target model in one call. Request logits for every supplied position, then compare the target greedy tokens with the aligned draft sequence.

Keep prompt, proposal, and verification token arrays in a supported 32-bit integer dtype. Mixing unsigned and signed 32-bit token arrays can promote a concatenation to 64-bit indices, which quantized embeddings reject.

The first supplied token is already accepted. Starting at the next position, find the longest matching prefix. If every draft token matches, keep the target model’s next token so generation can continue without an extra target call.

Use the one-call and full-acceptance cases as the next checkpoint:

pdm run test --week 3 --day 7 -- -k "verification or full_acceptance"

Task 4: Commit or Rewind

Treat cache offsets as a correctness invariant:

on full acceptance, advance both caches through the accepted proposal and synchronize the draft cache with the target’s extra token;
on a mismatch, emit the target token at that position and rewind every later speculative position from both caches;
after either path, assert that draft offset, target offset, and the logical length of every layer cache agree.

Exercise mismatch at the first, middle, and final proposed token. Also test a fully accepted proposal and EOS inside a proposal. Compare the complete output with ordinary greedy generation from the target model.

pdm run test --week 3 --day 7

Run the integrated path with a small draft model and a larger target model:

pdm run main --solution tiny_llm --loader week3 \
  --draft-model qwen3-0.6b --model qwen3-4b --max-tokens 64

This runs your completed Week 3 solution. To compare the completed reference on the same inputs, rerun it separately with --solution tiny_llm_ref.

The draft-model CLI is greedy-only. It rejects temperature, top-p, and top-k sampling options instead of silently ignoring them. Probability-correct sampled speculation remains a separate future extension.

Design the Measurement

The main command above is a functional smoke test. It does not emit paired target-only and speculative timings, so it is not performance evidence.

For a performance decision, run ordinary cached target generation and speculative generation in balanced fresh processes with the same prompt, tokenizer, --max-tokens value, seed, and synchronization boundary. For example, compare the command above with the same command without --draft-model, keeping --max-tokens 64 on both. Verify identical greedy output, then report proposal length, accepted tokens per proposal, target verification calls, draft-model time, target-model time, cache maintenance time, and end-to-end tokens per second for both paths. Record the raw samples and process order.

Until such a paired artifact exists, this chapter makes no speedup claim. Acceptance rate alone omits draft work, verification, synchronization, and cache maintenance.

🚧 Week 4: Build a Coding Agent

Weeks 1 through 3 ended with a working inference path: render a conversation, run the model, and carry its KV cache into later decoding. Week 4 asks what has to surround that path before model text can act on a project.

The product you are building is one bounded coding-agent run:

task -> model response -> validated action -> workspace observation
     -> approved effect -> receipt -> checkpoint -> compacted view
     -> visible steering -> observable evaluation
     -> two isolated continuations -> explicit selection
     -> bounded retrieval of oversized evidence

Each arrow is a harness decision, not a model privilege. The model proposes one JSON action. Ordinary Python decides whether the action is well formed, enabled, approved, executed, retained, or refused.

What Is Runnable Today

The repository currently ships one cumulative Day 9 declaration scaffold. All Week 4 modules and exports are visible from Day 1, but later-day implementation surfaces remain out of scope until their chapter. Most are TODO stubs; Day 9 explicitly supplies one constructor-validation rule. This is not nine separately materialized starters.

The deterministic learner checkpoint is cumulative within Week 4: pdm run test --week 4 --day N force-refreshes the supplied learner tests for Days 1 through N, then runs those files together. A later checkpoint therefore rechecks every earlier mechanism it builds on.

The real-model pdm run agent command currently exercises the learner loop, workspace, approvals, and receipts from Days 1–3, but its MLX-LM adapter calls mlx_lm.generate directly. It does not exercise the learner-owned generate_response helper or the Week 1–3 course model/cache path. Day 1 now tests that helper directly, and Day 8 reconnects to the course model/cache path in a deterministic test and a manual walkthrough. After all nine days are complete, pdm run week4-capstone composes the deterministic mechanisms in one disposable scenario.

These limits are visible course state, not goals for the learner to repair in the prose-only checkpoint.

The Nine-Day Progression

Day	Product pressure	Learner-owned mechanism	Evidence to inspect
1	Model text is not yet a safe next step.	A validated JSON action protocol and bounded loop.	Parsed events, exact observations, and stop reasons.
2	A fake workspace cannot inspect a project.	Contained directory listing and UTF-8 reads.	Listed paths, returned bytes, and recoverable errors.
3	A read-only agent cannot finish a coding task.	Approval, exact edits and commands, and effect receipts.	Changed bytes, validation status, and receipt facts.
4	A stopped process loses its conversation/model position.	One complete-observation checkpoint and resume boundary.	Saved messages/cache metadata and no effect replay.
5	Completed evidence consumes prompt space.	Receipt-backed deterministic compaction.	Tokens before/after, saved tokens, and unchanged receipts.
6	An operator needs a visible correction point.	Inspect, append one steering message, and resume.	Public status and message ordering.
7	A final sentence is not proof.	A report over declared observable outcomes.	Named file/result/receipt checks.
8	Two continuations should not prefill one identical prefix twice.	Dense token/KV-prefix reuse, isolated effects, and explicit selection.	Prefix offsets, avoided logical prefill, branch-local facts, and reports.
9	A large result should not fill every later prompt.	Content-addressed bytes, bounded previews, and exact range retrieval.	Artifact size/digest, omitted interval, and returned range.

The mechanisms compose in that order. Days 4–9 remain library APIs rather than additions to the real-model agent CLI. The supplied deterministic capstone is the orchestration shell that exercises those completed APIs together; it does not replace the mechanisms you implement here.

Prerequisites and Environment

Complete repository setup and Weeks 1 through 3 first. Day 8 directly uses the course tokenizer, model, and dense KV cache. The other deterministic Week 4 tests use scripted models and temporary workspaces, so they need no model download.

The supported native environment is macOS on Apple Silicon with the project dependencies installed. Real-model sections are manual and nondeterministic. An uncached run also needs network access, free disk space, and enough unified memory for the selected MLX weights. Use only disposable workspaces with no secrets: tool observations become model input, and Day 3 can enable file changes plus one exact allowlisted command.

Work Through One Chapter

For Day N:

Read what the final scaffold already declares and which TODO bodies belong to this day. Ignore future modules even though their declarations are visible.
Predict the named action, count, range, or stop reason before running the focused scenario when the chapter asks for one.
Run the cumulative learner checkpoint:
```
pdm run test --week 4 --day N
```
For Week 4, this command force-refreshes the learner tests for Days 1 through N from the supplied checkpoints, then runs all of them in one pytest invocation. Keep your implementation in src/; do not modify a copied test because the next run replaces it.
Implement only the files and relationships named by that chapter.
Rerun the checkpoint and inspect the artifact that can falsify your prediction: events, files, receipts, checkpoints, reports, cache offsets, or artifact bytes.
After Day 9 is green, run the composed product witness:
```
pdm run week4-capstone
```
Inspect its sorted JSON sections for compaction, both branches, the selected branch, and the externalized artifact range.

Course maintainers can run the corresponding day-local test-refsol command without copying learner tests. Optional model walkthroughs come after the deterministic checkpoint; they are exploration, not correctness evidence.

Read the Metrics as Accounting

Week 4 exposes three kinds of useful counts:

Day 5: transcript tokens before and after compaction;
Day 8: reused prefix tokens, layer offsets, and avoided-prefill tokens; and
Day 9: complete artifact bytes, model-visible bytes, and returned range bytes.

These values prove identity and logical-work accounting inside the teaching mechanisms. They do not establish wall-clock speedup, throughput, model quality, memory-capacity gain, or a universal policy. A manual cached-Qwen run may record the model ID, cache state, device, and observed actions, but its choices remain nondeterministic and non-comparative.

Week Boundary

This is a teaching agent for a trusted operator and disposable local projects. It is not a sandbox, hostile-filesystem defense, process jail, durable transaction system, distributed scheduler, session tree, hidden grader, semantic-perfect memory, network artifact service, or production serving framework. Completed effects are never presented as rewound, and a model’s final prose is never treated as proof by itself.

Day 1: A Validated Agent Loop

Weeks 1 through 3 built a function that turns a conversation into model text. A coding agent needs a small control loop around that function: ask for one response, decide whether it is a final answer or an action, record what happened, and continue when an action produces an observation.

Start with one prediction: if the model first requests a disabled tool and then returns malformed JSON, which requests reach the workspace, which errors enter the next model input, and which configured budget can stop the run first? The event trace at the end of this chapter lets you check every part of that answer.

The model never edits a file directly. It emits text. Ordinary Python validates that text before handing a parsed action to a workspace object. This separation makes the loop deterministic to test even when no model weights are loaded.

The Teaching Boundary

Day 1 teaches only a bounded loop and one JSON action protocol. The supplied test uses a fake workspace with one enabled read-only action. Real project inspection arrives on Day 2; file mutation, command execution, approval, and durable receipts arrive later.

The loop validates and records a model response, but it does not prove that the model solved the task. It is also not a sandbox, background worker, persistent session, or production scheduler.

Files and Public Surface

The repository is a final Day 9 declaration scaffold. Future agent modules and exports are already visible, but their implementation surfaces are not part of Day 1. Most later bodies are TODO stubs; Day 9 also contains one explicitly supplied constructor check. Implement only the following surfaces:

Implement the TODO bodies in these Day 1 starter files:

File	Public names	Responsibility
`src/tiny_llm/agent/generation.py`	`initial_messages`, `generate_response`	Begin a conversation and keep one model-response boundary explicit.
`src/tiny_llm/agent/protocol.py`	`AgentError`, `FinalAction`, `ToolAction`, `parse_action`, `build_system_prompt`	Represent and validate one final answer or one enabled tool request.
`src/tiny_llm/agent/loop.py`	`AgentLimits`, `AgentEvent`, `AgentRun`, `run_agent`	Bound a run, propagate observations, and retain an inspectable trace.

generate_response() remains part of the public Day 1 surface. It renders the messages with the course tokenizer, decodes at most max_tokens with a fresh cache, stops at EOS, and releases every cache in a finally block.

The supplied Day 1 test checks generate_response() with the course tokenizer and model boundary: exact prompt and thinking offsets, EOS stopping, the token limit, a fresh cache for each call, and cache release on normal and exceptional paths. The pdm run agent CLI still uses its own MLX-LM generation adapter, so a successful live CLI run is not evidence for this helper; the cumulative Day 1 checkpoint is.

Run the cumulative learner checkpoint from the repository root:

pdm run test --week 4 --day 1

The command force-refreshes the supplied Day 1 learner test before running it. Before you implement the TODOs, the implementation-dependent cases across ten task groups are expected to fail. No model download is required.

Course maintainers can check the supplied implementation without copying the learner test:

pdm run test-refsol --week 4 --day 1

Task 1: Start the Conversation Deliberately

initial_messages(task, system_prompt) creates the first two messages:

[
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": task},
]

Reject an empty or whitespace-only task. build_system_prompt(workspace) describes only the actions enabled for this run. The prompt is guidance, not enforcement: the protocol and workspace boundary must still reject anything the policy does not allow.

Task 2: Execute One Tool and Finish

run_agent(task, generate, workspace, limits=None) starts from those messages. The test injects a generate callable that returns predetermined strings, so the control flow stays deterministic.

For a valid tool action, call workspace.execute(action), record the action and result, and append both the assistant response and a user observation. When a later response is a valid FinalAction, return a completed AgentRun with the final text.

Task 3: Validate One Structured Decision

A model response is exactly one JSON object. It is either a final answer:

{"final":"I inspected README.md."}

or one tool request:

{"tool":"read_file","path":"README.md"}

parse_action() rejects malformed JSON, non-object values, blank final text, unknown or disabled tools, missing fields, unexpected fields, and fields with the wrong shape. Do not ignore trailing or extra data.

TOOL_FIELDS names the cumulative vocabulary: list_files, read_file, write_file, edit_file, and run_command. Day 1 implements none of those effects. Its fake workspace enables only read_file, which is enough to prove that availability is checked before dispatch.

Malformed or unavailable actions become ordinary error: observations. The model can see the failure and choose another response instead of crashing the Python loop.

Task 4: Stop at the Step Budget

AgentLimits.max_steps bounds how many model decisions one run may attempt. When the loop consumes that budget without a valid final answer, return an incomplete run whose reason is step_limit. The events show exactly how the budget was spent.

Task 5: Return an Inspectable Run

Every interaction becomes an AgentEvent with the step number, raw response, parsed action when one exists, and result or validation error. AgentRun records the completion flag, stop reason, optional final answer, and immutable event tuple.

A run marked completed means only that the model returned a valid final action. Later days add receipts and outcome evaluation; Day 1 keeps the trace small and in memory.

Task 6: Recover from Invalid JSON

After an invalid response, append the raw assistant response and its exact validation error as the next user observation. Reset the identical-action counter, then let the model try again while the invalid-action budget remains.

The focused case sends invalid JSON followed by a valid final response. The first event must retain the recoverable error, and the second must complete the run.

Task 7: Stop Repeated Actions

Serialize each parsed tool name and normalized argument object into a stable signature. Count consecutive identical requests and stop with repeated_action_limit when the count exceeds the configured budget.

This guard matters even when a tool succeeds: repeating the same request can consume the whole run without adding new information.

Task 8: Preserve the Exact Observation

The next model call must receive the complete tool result, not only a marker:

{
    "role": "user",
    "content": "Tool result:\nREADME contents",
}

The normal guard fails if that payload is changed or dropped. It also proves that a known-but-disabled tool such as write_file becomes an error observation and never reaches the fake workspace.

Task 9: Make Every Limit Fail Closed

Require positive values for max_steps, max_context_chars, max_invalid_actions, and max_identical_actions. Zero or negative budgets would disable the intended stopping guarantee and must be rejected.

Before each model call, bound the total message-content characters and stop with context_limit when it is too large. Count invalid actions and stop with invalid_action_limit when that budget is exhausted. Together with the step and repeated-action limits, every run has an explicit terminal reason.

Checkpoint

When Day 1 is green, inspect the focused test rather than only its final pass: confirm the initial system/user pair, one dispatched read_file, the exact observation in the next model input, the completed final event, and each budgeted stop reason. Revisit your opening prediction: a disabled or malformed request must not reach the workspace, and the exact error must remain visible to the following model turn.

This checkpoint proves the scripted protocol and loop, plus the focused generate_response() model/cache boundary. It does not make the separate real-model CLI exercise that helper.

You now have a validated, bounded model → action → observation loop. Continue with Day 2: Inspect a Workspace to replace the fake tool boundary with real contained directory listing and UTF-8 file reads.

🚧 Day 2: Inspect a Workspace

Day 1 built a loop around a fake workspace. The model could request a tool, the loop could return an observation, and the model could finish—but no tool looked at a real project.

Day 2 adds the smallest useful workspace: two read-only tools that let the model list a directory and read a text file. This makes the complete cycle visible without mixing in approval, edits, commands, or durable receipts.

model requests list_files
          |
          v
 Workspace checks the path
          |
          v
 loop returns the listing
          |
          v
 model requests read_file
          |
          v
 loop returns the file text
          |
          v
 model returns a final answer

The Teaching Boundary

Use these tools with a trusted operator and a disposable local repository. Path.resolve() containment and symlink refusal prevent ordinary path mistakes, but they are not a sandbox or a defense against a hostile filesystem. The tool result also becomes model input, so do not point the workspace at files you would not send to the model.

Day 3 adds operator-approved edits, one allowlisted validation command, and a simple receipt log. Keeping those effects out of Day 2 lets you first see the read-only agent loop clearly.

Files and Public Surface

The published starter is the cumulative Day 9 declaration scaffold. Receipt, checkpoint, compaction, and later modules are therefore visible already, but their implementation work belongs to later chapters. Day 2 owns only the workspace relationships below.

Implement the TODO bodies in:

File	Public names	Responsibility
`src/tiny_llm/agent/workspace.py`	`ToolPolicy`, `Workspace`	Bound one directory and expose `list_files` plus `read_file`.
`src/tiny_llm/agent/__init__.py`	Day 1 API plus `ToolPolicy`, `Workspace`	Complete the Day 2 exports within the final scaffold.

The Day 2-owned prefix of ToolPolicy has three fields, in order:

root: Path
max_file_bytes: int = 64 * 1024
max_list_entries: int = 200

Day 2 implements this prefix of Workspace:

available_tools, always {"list_files", "read_file"};
resolve_path(raw, must_exist=True);
list_files(raw=".");
read_file(raw); and
execute(action).

The starter declarations are the contract. Do not implement or depend on the visible write, command, approval, receipt, checkpoint, compaction, steering, evaluation, branching, or evidence-retrieval declarations yet.

Task 1: Normalize One Workspace Root

ToolPolicy receives the directory the agent may inspect. Convert it to a resolved Path, require an existing directory, reject a symlink root, and require both numeric limits to be positive.

Normalizing once keeps every later check relative to the same root:

policy = ToolPolicy(Path("demo-project"))
workspace = Workspace(policy)

Task 2: Resolve Ordinary Relative Paths

resolve_path() accepts a non-empty relative path and returns its resolved location under policy.root. Reject:

absolute paths and .. traversal;
symlinked path components;
.git, .env, .ssh, .aws, and common credential or private-key names;
missing paths when must_exist=True; and
any resolved path outside the workspace.

Return a recoverable AgentError for these cases. The agent loop can turn that error into an observation instead of crashing.

This is deliberately an ordinary local-repository boundary. A different process can still race filesystem checks; defending against a hostile filesystem is outside this course checkpoint.

Task 3: List One Directory

list_files() lists direct children in sorted order. Emit one line per visible regular file or directory:

file README.md
dir src

Omit protected names, symlinks, and special files. Stop after max_list_entries lines. Return (empty directory) when no visible entry remains.

The tool is intentionally not recursive. The model can request another list_files action for a directory it wants to inspect.

Task 4: Read One Text File

read_file() accepts a visible regular file no larger than max_file_bytes. Read its bytes and decode UTF-8. Reject directories, oversized files, invalid UTF-8, protected paths, and symlinks with AgentError.

The size bound keeps a single observation from consuming the whole context window. Later checkpoints can add more deliberate context selection; Day 2 only needs one obvious limit.

Task 5: Turn Failures into Observations

execute() dispatches a parsed ToolAction to list_files or read_file. It returns successful text directly. A recoverable failure becomes a string beginning with error: so run_agent() can append it to the conversation and let the model choose another action.

A known future tool such as write_file is still disabled on Day 2. The system prompt describes only workspace.available_tools, and parse_action() checks that enabled set before dispatch.

Task 6: Run the Read-Only Cycle

The checkpoint uses a scripted model so it is deterministic and does not load weights:

responses = iter([
    '{"tool":"list_files"}',
    '{"tool":"read_file","path":"README.md"}',
    '{"final":"README says hello"}',
])

result = run_agent(
    "inspect the project",
    lambda messages: next(responses),
    workspace,
)

Inspect result.events: the first two events contain the parsed tool actions and exact observations, and the third contains the final answer. The same tool result appears in the next model input as Tool result:\n....

Run the Day 2 Checkpoint

From the repository root, run the cumulative learner checkpoint:

pdm run test --week 4 --day 2

Before you implement the TODOs, the implementation-dependent cases across eight task groups are expected to fail. The command force-refreshes the supplied learner tests for Days 1 and 2, then runs both; completed Day 1 behavior should remain green. During course development, check the supplied implementation with the day-local maintainer command:

pdm run test-refsol --week 4 --day 2

The course-code guard compares the starter and reference public signatures, dataclass fields, package exports, and solution-free method bodies.

Explore with a Real Model

The scripted checkpoint above is the reproducible mechanics proof. This separate manual run is exploratory: a real model chooses the tools, so its wording and exact tool order may vary. Use only a disposable directory that contains no secrets or private source.

The CLI defaults to qwen3-4b (Qwen/Qwen3-4B-MLX-4bit), whose cached weights use about 2 GiB; use that lower-resource option when needed. The recorded exploratory run below used mlx-community/Qwen3-30B-A3B-4bit, whose cached weights use about 16 GiB and require sufficient Apple unified memory. Model behavior and tool order vary with either choice. Both use the local MLX model path already used by this repository. You need macOS on Apple Silicon and the installed MLX dependencies. The first run downloads the selected weights from Hugging Face when they are not cached, so it also needs network access and the corresponding free disk space. If MLX is unavailable or the weights cannot be loaded, the command exits before the agent calls a workspace tool; it does not substitute scripted output.

This CLI uses a supplied MLX-LM generation adapter. It exercises your Day 1 loop and Day 2 workspace, but it does not call the learner-owned generate_response() helper or prove the Weeks 1–3 course model/cache path.

Create a tiny read-only workspace, then give the model one goal:

INSPECT_ROOT="$(mktemp -d)"
mkdir "$INSPECT_ROOT/src"
printf '%s\n' '# Pocket Weather' 'A tiny terminal forecast project.' > "$INSPECT_ROOT/README.md"
printf '%s\n' 'def forecast(city):' '    return f"Sunny in {city}"' > "$INSPECT_ROOT/src/weather.py"

pdm run agent -- --model mlx-community/Qwen3-30B-A3B-4bit --root "$INSPECT_ROOT" \
  "Inspect this workspace and explain its purpose and the behavior implemented in its source file. Use the available workspace tools to gather evidence from both the project overview and the source file. Your first response must be one tool request, every response must contain exactly one JSON object, and you must not finish until you have read the source file."

Watch the printed goal, model responses, parsed actions, and tool observations. There is no predefined action list: the model decides what to list and read before returning a final answer. The policy is read-only, so no approval prompt or receipt is expected. Compare the final answer with the actual disposable files:

find "$INSPECT_ROOT" -type f -print
cat "$INSPECT_ROOT/README.md" "$INSPECT_ROOT/src/weather.py"

Checkpoint

Confirm that the deterministic run listed a directory, read exact UTF-8 file contents, returned recoverable path errors to the model, and completed without enabling any effect tool.

When this checkpoint is green, continue to Day 3: Edit, Validate, and Record.

After all nine days are complete, the supplied deterministic capstone composes this read-only inspection with checkpointing, compaction, steering, evaluation, branching, and bounded evidence. That orchestration is a separate week4-capstone command, not part of the real-model agent CLI.

🚧 Day 3: Edit, Validate, and Record

Day 2 gave the model two read-only tools. It could inspect a disposable project and explain what it found, but it could not fix anything. Day 3 completes one small coding cycle:

read file -> propose exact edit -> operator approves -> recheck bytes
    -> replace file -> record receipt -> run focused check -> record receipt
    -> final answer

The important idea is not broad autonomy. It is an explicit boundary between a model proposal and a local side effect.

Before running the focused cycle, predict three facts: how many approvals and receipts exist after a denied effect, whether stale-read detection changes any file bytes, and which evidence remains after a validation command returns a nonzero status. The workspace, receipt log, and event trace let you check those answers independently of the model’s final sentence.

The Teaching Boundary

Use this checkpoint with one trusted operator, one Python process, and a disposable repository that contains no secrets. Its path checks prevent ordinary mistakes, but they are not a sandbox or a defense against a hostile filesystem. The command tool is not a process jail: an allowed program can access anything the host process can access, spawn children, or use the network.

The receipt file is a simple append-only JSONL teaching record. It detects edited receipt bytes when reopened and handles a repeated call ID in the same process. It is not a transaction log, an fsync protocol, a multi-writer store, or proof that an interrupted effect did or did not happen. Day 3 deliberately stops at this receipt boundary.

Files and Public Surface

Later checkpoint, compaction, steering, evaluation, branching, and evidence modules are already declared in the cumulative Day 9 scaffold. Their later-day implementation work is out of scope. Day 3 owns only the following Day 1–3 surfaces:

Implement the TODO bodies in these starter files:

File	Public names	Responsibility
`src/tiny_llm/agent/workspace.py`	`ToolPolicy`, `Workspace`	Authorize reads, approved edits, and one exact validation command.
`src/tiny_llm/agent/receipts.py`	`EffectReceipt`, `ReceiptStore`	Represent effects and optionally append verified JSONL records.
`src/tiny_llm/agent/__init__.py`	Day 1–3 names within the final scaffold	Export the two receipt types.

ToolPolicy keeps its first three Day 2 fields and adds:

allow_writes: bool = False
allowed_commands: tuple[tuple[str, ...], ...] = ()
max_write_bytes: int = 64 * 1024
command_timeout_seconds: float = 30.0

Writes stay disabled unless allow_writes=True. Commands stay disabled unless their complete argument tuple appears in allowed_commands. There is no shell string or prefix match.

Workspace(policy, confirm_tool=None) creates an in-memory receipt store. Pass a ReceiptStore(path) as the third argument when you want JSONL output. The workspace extends the Day 2 methods with write_file, edit_file, run_command, and modified_files. execute(action, tool_call_id=None) is the model-facing gate. Direct tool methods are useful for focused unit tests; execute performs the approval and receipt steps.

Task 1: Authorize Tools Explicitly

Build available_tools from the policy. Listing and reading are always present. Add both file mutation tools only when writes are enabled, and add run_command only when at least one exact command is configured. Validate all size limits, the timeout, the boolean flag, and every command part.

This configuration permits one focused check:

validation = ("python", "-m", "pytest", "tests/test_math.py", "-q")
policy = ToolPolicy(
    Path("demo-project"),
    allow_writes=True,
    allowed_commands=(validation,),
)

Task 2: Read Before Changing Existing Bytes

Keep Day 2’s path and read rules. When read_file succeeds, remember the SHA-256 digest of the bytes that were returned. Replacing or editing an existing file requires that observation. A new file does not have old bytes to inspect, but its parent directory must already exist.

For edit_file, require a non-empty old string that occurs exactly once. Compute the proposed bytes in memory and enforce max_write_bytes before asking for approval. Whole-file write_file replacements also require a prior read.

Task 3: Ask Once, Default No, Then Recheck

execute preflights the complete action before calling confirm_tool. Missing callbacks, False, and every value other than the boolean True deny the effect. A terminal program can provide a small default-No callback:

def confirm(action):
    answer = input(f"Approve {action.tool} {action.arguments}? [y/N] ")
    return answer.strip().lower() in {"y", "yes"}

After approval, write_file or edit_file reads the destination again and compares its digest with the earlier observation. If another actor changed the bytes while the operator was deciding, return error: file changed since it was read and do not overwrite them.

Task 4: Replace Through the Same Directory

Write the proposed bytes to a temporary file in the destination’s parent, close it, and call os.replace(temporary, destination). Clean up a leftover temporary file after an error. This avoids presenting a partially written destination to ordinary readers.

This small pattern is atomic at the replacement step, but it is not a durable journal and does not close the check-to-replace race against a hostile actor.

Task 5: Run One Exact Validation Command

run_command(argv) accepts a non-empty list of strings only when its tuple is exactly allowlisted. Call subprocess.run without a shell, with the workspace root as cwd, captured text output, and the configured timeout. Bound combined stdout and stderr so one observation cannot consume the whole context window.

Return one observation with the status and captured output:

status: 0
output:
1 passed

A nonzero status and a timeout are ordinary validation results the model can inspect. They are not Python exceptions and do not prove the final answer is correct.

Task 6: Record Simple Effect Receipts

An EffectReceipt has these fields, in order:

tool_call_id: str
tool: str
arguments: dict[str, Any]
exit_state: str
result: str
changed_artifacts: tuple[str, ...] = ()

Its receipt_id is the SHA-256 digest of the canonical JSON payload. A successful write or edit records exactly one normalized workspace-relative artifact. A validation receipt records the exact argv, status and captured output, with no changed artifacts. ReceiptStore(path) loads and verifies an existing JSONL file; ReceiptStore() remains in memory.

The store maps one tool_call_id to one receipt. Repeating the same call ID and action returns the existing result without running the effect again. Reusing the ID for another action is an error. This is proportional duplicate handling for one process, not distributed exactly-once execution.

Task 7: Run the Complete Scripted Cycle

The test uses scripted model responses, so it needs no model weights:

responses = iter([
    '{"tool":"read_file","path":"app.py"}',
    '{"tool":"edit_file","path":"app.py","old":"1","new":"2"}',
    '{"tool":"run_command","argv":["python","-m","pytest","tests/test_math.py","-q"]}',
    '{"final":"changed and validated app.py"}',
])

store = ReceiptStore(Path("demo-project/.agent-receipts.jsonl"))
workspace = Workspace(policy, confirm, store)
result = run_agent("fix app.py", lambda _messages: next(responses), workspace)

Inspect result.events, workspace.modified_files, and the two receipts. The edit receipt names app.py; the validation receipt has an empty artifact tuple. The final answer is still a model statement, so the validation status in the trace is the evidence that matters.

Run the Day 3 Checkpoint

From the repository root, run the cumulative learner checkpoint:

pdm run test --week 4 --day 3

Before you implement the TODOs, the Day 3 cases are expected to fail because the new starter methods return None. The command force-refreshes and runs the supplied learner tests for Days 1–3 together. Keep the Day 3 failures until you solve each task; do not import tiny_llm_ref from the starter.

Course maintainers can check the supplied implementation without copying the learner test:

pdm run test-refsol --week 4 --day 3

The course-code guard checks exact public signatures, dataclass fields, package exports, and TODO-only starter bodies across the final declaration scaffold.

Explore the Full Cycle with a Real Model

Keep the scripted checkpoint as the deterministic proof. This manual exercise lets a real model plan the same Day 3 cycle from one natural-language goal. Its wording and tool order can vary, and completion is not an automated test. Use a fresh disposable directory with no secrets.

The CLI default is qwen3-4b (Qwen/Qwen3-4B-MLX-4bit), whose cached weights use about 2 GiB; use that lower-resource option when needed. The recorded exploratory run below used mlx-community/Qwen3-30B-A3B-4bit, whose cached weights use about 16 GiB and require sufficient Apple unified memory. Model behavior and tool order vary with either choice. Both require macOS on Apple Silicon and the installed MLX dependencies. An uncached first run downloads the selected weights from Hugging Face and needs network access plus the corresponding free disk space. If MLX, network access, disk space, unified memory, or the weights are unavailable, model loading fails before any tool call; do not treat a scripted checkpoint as evidence that this live run occurred.

As on Day 2, this CLI’s supplied MLX-LM adapter bypasses the learner-owned generate_response() helper. It does exercise your loop, workspace, approval, effect, and receipt paths.

Pre-create the workspace, an existing file, and one focused validation fixture:

EDIT_ROOT="$(mktemp -d)"
printf '%s\n' \
  'def greeting(name):' \
  '    return f"Hello, {name}!"' > "$EDIT_ROOT/app.py"
cat > "$EDIT_ROOT/validate.py" <<'PY'
from pathlib import Path
from app import greeting

assert greeting("Ada") == "Welcome, Ada!"
assert Path("NOTES.md").read_text() == "Greeting now says Welcome.\n"
print("validation passed")
PY

Now give the real model one goal. --allow-command names the only command it may run, including its exact arguments. Each proposed write, edit, or command still pauses at a default-No learner approval prompt:

pdm run agent -- --model mlx-community/Qwen3-30B-A3B-4bit --root "$EDIT_ROOT" \
  --allow-writes \
  --allow-command "python validate.py" \
  --receipt-log .agent-receipts.jsonl \
  "Inspect the workspace. Create NOTES.md containing exactly 'Greeting now says Welcome.' followed by a newline, precisely change app.py so greeting says Welcome instead of Hello, run the allowed validation, react to its evidence, and finish with a brief summary. Use workspace tools to gather evidence before any effect. Your first response must be one tool request, every response must contain exactly one JSON object, and do not finish until the requested files and validation evidence have been inspected."

Read every approval payload before answering y. The live trace shows the model response, parsed action, tool observation, validation status/output, and final answer. Afterward, inspect the changed bytes and the durable effect receipts rather than trusting the final prose alone:

cat "$EDIT_ROOT/app.py" "$EDIT_ROOT/NOTES.md"
cat "$EDIT_ROOT/.agent-receipts.jsonl"

The receipt records preserve the approved write_file, edit_file, and run_command arguments and outcomes. The status and captured output in the validation receipt are the evidence to compare with the changed bytes.

Checkpoint

You now have a small end-to-end coding loop: inspect real bytes, propose one precise change, pause for a trusted operator, reject stale observations, replace the file, validate with one exact command, and retain simple evidence of both effects. Continue with Day 4: Checkpoint and Resume to save one complete observation boundary and restore it through a fresh scripted model without turning Day 3 into production infrastructure.

Day 3 is the last mechanism exposed by the real-model CLI. Days 4–9 use deterministic library checkpoints, and the separate supplied Week 4 capstone connects all nine completed days into one runnable deterministic product path.

Day 4: Checkpoint and Resume

Days 1 through 3 build one uninterrupted coding-agent run. The model proposes an action, the harness executes it, and the result becomes the next model observation. But if the process stops, a new model object does not know which conversation prefix it had already processed.

Day 4 makes one boundary visible: save the conversation and the small model snapshot immediately after a complete tool observation, then restore both into a fresh scripted model and continue at the next response. The completed tools stay completed; resume starts after their observations instead of replaying them.

This is a teaching checkpoint for one process and the course’s fake model. It is not a session tree, rewind feature, persistent KV store, transaction log, or exactly-once effect system.

Files and Public Surface

The final scaffold already declares compaction, steering, evaluation, branching, and evidence APIs. Leave those future TODO bodies alone. Day 4 owns only the checkpoint model and the two loop entry points below:

Implement the TODO-only surfaces in:

File	Public names	Responsibility
`src/tiny_llm/agent/checkpoint.py`	`ModelCheckpoint`, `AgentCheckpoint`, `create_checkpoint`	Represent and validate one in-memory conversation/model snapshot.
`src/tiny_llm/agent/loop.py`	`run_to_checkpoint`, `resume_agent`	Stop after a complete observation, then continue with a fresh model.
`src/tiny_llm/agent/__init__.py`	the names above	Complete the Day 4 exports within the final scaffold.

Run the cumulative learner checkpoint from the repository root:

pdm run test --week 4 --day 4

Before you implement the TODOs, all seven Day 4 tasks are expected to fail. The test uses a scripted model, fake cache metadata, a temporary workspace, and one exact Python validation command. It does not load model weights. The command force-refreshes and runs the supplied learner tests for Days 1–4 together.

Course maintainers can check the supplied implementation without copying the learner test:

pdm run test-refsol --week 4 --day 4

One Safe Loop Boundary

A checkpoint is saved only after the harness has appended both halves of a tool interaction:

assistant: {"tool":"edit_file", ...}
user:      Tool result:\nedited app.py
                                      ^ checkpoint here

Saving before the observation would leave the restored model unable to tell whether the tool ran. Day 4 therefore counts completed tool calls and saves at the boundary after _append_tool_result(...) has produced the next complete conversation.

The checkpoint stores the semantic messages, not the AgentEvent history. Day 3 receipts remain separate evidence about the edit or command. They are not copied into the checkpoint and they do not become a replay controller.

Task 1: Represent the Fake Model Snapshot

ModelCheckpoint contains four fields, in order:

conversation_position: int
response_index: int
cached_token_ids: tuple[int, ...]
layer_offsets: tuple[int, ...]

conversation_position is the number of semantic messages at the saved boundary. response_index tells the scripted model which response comes next. cached_token_ids represents the prompt prefix in the fake cache, and every layer_offsets entry must equal its length.

Reject negative positions, invalid token IDs, a missing layer snapshot, or offsets that disagree with the cached prefix. These checks make the fake model state internally coherent without introducing a production cache format.

Task 2: Bind Conversation and Model State

AgentCheckpoint contains:

checkpoint_id: str
task: str
messages: tuple[tuple[str, str], ...]
model: ModelCheckpoint

create_checkpoint(task, messages, model) copies each mutable message into an immutable (role, content) pair. It computes checkpoint_id as the SHA-256 of canonical JSON containing the task, messages, and model fields.

AgentCheckpoint.validate() checks the ordinary resume contract:

the task and messages are structurally valid;
the model’s conversation position equals the saved message count;
the checkpoint ID still matches the content.

This identity check catches an accidental mismatch. It is not a hostile-tamper or authentication scheme.

Task 3: Stop After a Complete Observation

run_to_checkpoint(task, generate, workspace, after_tool_calls=1, limits=None) starts with the same prompt and validation rules as run_agent. It counts a tool call only after execution and observation append. At the requested count, it calls:

model_state = generate.save_checkpoint(messages)

and returns an AgentCheckpoint.

The generator must return a ModelCheckpoint. A missing checkpoint method, a non-positive tool-call count, or a run that finishes before the boundary is a clear error instead of a partial checkpoint.

Task 4: Restore a Fresh Model

resume_agent(checkpoint, fresh_generate, workspace, limits=None) validates the checkpoint, calls:

fresh_generate.restore_checkpoint(checkpoint.model)

rebuilds the semantic message list, and enters the normal loop. The fresh model therefore sees the exact conversation prefix and produces the response at the saved response_index.

No old tool action is submitted again. The first model call after restore sees the already-recorded tool result and chooses the next action or final answer.

Task 5: Make the Fake Cache Visible

The test’s FakeCheckpointModel turns each message content length into one fake token ID. save_checkpoint(messages) records those token IDs, two matching layer offsets, and the next scripted-response index. A new model object restores that state and asserts that its first resumed input has the same token prefix and conversation position.

This deliberately small representation exposes the inference/harness integration: the harness owns semantic conversation state, while the model owns the cache snapshot that accelerates exactly that state.

Task 6: Resume Without Replaying Effects

The end-to-end test scripts:

read app.py
edit app.py
run the exact validation command
checkpoint after the validation observation
construct a fresh scripted model
resume to the final answer

Before resume, the Day 3 store already contains the edit receipt and validation receipt. After resume, the approval log and receipt count are unchanged: neither completed effect ran twice. The checkpoint contains no receipt IDs and makes no exactly-once claim; it simply resumes after the conversation already says those effects completed.

Task 7: Keep the Boundary Small

Day 4 implements only checkpoint.py and two loop entry points. Future modules are already declared in the final scaffold, but do not implement or depend on them here. Do not add session IDs, parent pointers, rewind methods, steering queues, disk cache files, or another checkpoint representation. If the in-memory checkpoint is lost, start a new run.

Checkpoint

You can now stop after a complete tool observation and continue with a fresh scripted model from the same conversation and fake-cache position. Inspect the checkpoint’s messages and model fields, then confirm that the pre-checkpoint edit and command remain single completed effects.

At this point, predict what resume is allowed to do: the next model response may request a new action, but no action already represented by the saved observation should run again. The approval and receipt counts in the focused scenario are the falsifying evidence.

Continue with Day 5: Compact Completed Work to derive a smaller model-visible transcript while keeping the exact effect receipts. The Day 5 checkpoint receives a transcript and receipts directly rather than changing this resume path. After Day 9, the supplied Week 4 capstone composes the two mechanisms through its deterministic orchestration.

Day 5: Compact Completed Work

🚧 Early-review WIP: This chapter is public for early review and may change. Use a disposable workspace when running the agent or enabling writes or commands.

Every tool call adds two model-visible messages: the assistant’s action and the tool observation. A long validation log can eventually crowd out the task and the useful recent steps. Deleting old messages saves space, but it also deletes the evidence behind claims such as “validation passed.”

Day 5 makes one boundary visible: replace an older, completed effect with a small deterministic evidence record while keeping its full EffectReceipt unchanged. The model receives fewer tokens and can continue from that view; the harness still retains the exact action and result.

The Starter Surface

The cumulative Day 9 scaffold already contains later steering, evaluation, branching, and evidence declarations. Leave those future TODO bodies alone. Day 5 owns one small module:

File	Public names	Purpose
`src/tiny_llm/agent/compaction.py`	`CompactionResult`, `compact_completed_interactions`	Derive a smaller model-visible transcript from completed, receipted effects.
`src/tiny_llm/agent/__init__.py`	the names above	Complete the Day 5 exports within the final scaffold.

Run the cumulative learner checkpoint:

pdm run test --week 4 --day 5

Use this command for the supplied implementation:

pdm run test-refsol --week 4 --day 5

Before you implement the TODOs, all six Day 5 tasks are expected to fail. The command force-refreshes and runs the supplied learner tests for Days 1–5 together.

Start From the Existing Transcript

The input is the same list of role/content messages that run_agent gives the model. A completed effect has this shape:

[
    {"role": "assistant", "content": '{"tool":"run_command",...}'},
    {"role": "user", "content": "Tool result:\nstatus: 0\n..."},
]

The compactor does not invent a second event log. It receives this transcript plus the Day 3 EffectReceipt values already produced by the workspace.

Before implementing it, predict which older interaction in the focused fixture is eligible to compact, which recent interaction must stay verbatim, and whether saved_tokens must be positive. The returned messages, exact counter values, and unchanged receipts let you falsify that prediction.

Task 1: Require Exact Receipt Evidence

A pair is eligible only when all three facts match one supplied receipt:

the parsed tool name;
the normalized argument object; and
the complete observation text.

If there is no receipt, or if any of those fields differs, leave both messages verbatim. This deliberately excludes Day 2 reads and listings: the current course receipts effects, not every observation. Day 5 must not pretend that an unreceipted result is durable evidence.

Task 2: Keep a Small, Honest Record

Keep the small assistant action, but replace its large tool-observation message with one bounded evidence record that contains:

the tool and its normalized arguments;
exit_state and changed_artifacts;
a bounded prefix of the result; and
the content-addressed receipt_id.

For example:

Completed tool interaction (compacted evidence):
{"arguments":{"argv":["python","validate.py"]},
 "changed_artifacts":[],"exit_state":"ok",
 "receipt_id":"...","result_preview":"status: 0...",
 "tool":"run_command"}

This is not a model-written summary. It is a deterministic rendering of fields the harness already verified. The bounded preview helps the model explain what happened; receipt.result still contains the full observation.

Task 3: Retain a Recent Tail

keep_recent=1 leaves the newest eligible effect as its original two messages. Older matching effects may compact. The recent tail keeps the next decision grounded in the exact latest interaction without requiring a complicated semantic policy.

The count refers only to receipted effect interactions. Unreceipted reads stay verbatim regardless of this setting.

Task 4: Measure the Real Model Input

The compactor accepts count_tokens(messages) instead of estimating tokens from characters. For the real model, use the same tokenizer and chat template as generation:

def count_tokens(messages):
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    return len(tokenizer.encode(prompt, add_special_tokens=False))

Build the proposed compact view, count it again, and accept it only when the exact counter decreases. CompactionResult reports tokens_before, tokens_after, and saved_tokens. The focused tests inject their own deterministic character counter, which sums the message-content lengths. It is distinct from Day 4’s cached-prefix counter and does not load or download a model.

Task 5: Preserve the Source of Truth

Never mutate the caller’s message list or any receipt. Return copied messages inside CompactionResult. Running the compactor again over its own output is a no-op because a compact evidence record is not a tool action followed by a tool result.

This distinction matters:

original transcript + full receipts   durable evidence owned by the harness
                  |
                  v
       compacted message view          temporary input for the model

If the compact view is lost, derive it again. Do not treat it as a replacement for receipts or checkpoints.

Task 6: Continue With the Existing Model Boundary

CompactionResult.messages contains ordinary role/content mappings. Pass a copied list to the same generation callable used by the loop, then validate the response through the existing protocol:

view = compact_completed_interactions(
    messages,
    receipts,
    count_tokens,
    keep_recent=1,
)
response = generate([dict(message) for message in view.messages])
action = parse_action(response, workspace.available_tools)

The Day 5 test makes the scripted model return a final answer after it sees the compact validation evidence. Nothing about action parsing, tool approval, or workspace execution changes.

Limits of This Teaching Compactor

This checkpoint intentionally does not add semantic-perfect summarization, automatic threshold scheduling inside run_agent, persistent compact views, receipt lookup by summary text, K/V cache editing, session trees, rewind, steering, or exactly-once execution. It compacts only completed effects backed by the receipts the caller supplies.

Checkpoint

You can now make an older validation interaction visibly smaller, inspect the receipt that retains its complete result, and feed the compact view to the next model call. Continue with Day 6: Inspect and Steer a Paused Agent to inspect one checkpoint, add one visible operator message, and resume without replaying completed work.

The Day 6 path resumes the Day 4 transcript; it does not consume CompactionResult.messages. After Day 9, the supplied capstone carries this compacted view into the later control path. Day 5’s token counts and the capstone’s combined report prove only deterministic transcript accounting, not latency, throughput, quality, or memory-capacity improvement.

Day 6: Inspect and Steer a Paused Agent

🚧 Early-review WIP: This chapter is public for early review and may change. Use a disposable workspace when running the agent or enabling writes or commands.

Day 4 gave the harness a safe pause after a complete tool observation. Day 5 showed how to make older completed evidence smaller for the model without discarding the full receipts. Those boundaries also create a useful moment for an operator: inspect what the paused run has actually recorded, add one new instruction, and let a fresh model continue.

The focused Day 6 scenario resumes the original Day 4 transcript rather than CompactionResult.messages. After Day 9, the supplied deterministic capstone connects the compacted view to this later control path.

Day 6 implements only that interaction. It does not inspect hidden reasoning or guess the model’s plan. It reports public facts from the checkpoint, appends one ordinary steering message, and resumes through the existing validated loop.

The Starter Surface

Later evaluation, branching, and bounded-evidence declarations are already visible in the final Day 9 scaffold. Leave those TODO bodies alone. Day 6 owns one module:

File	Public names	Purpose
`src/tiny_llm/agent/steering.py`	`AgentStatus`, `inspect_checkpoint`, `resume_with_steering`	Inspect one complete-observation checkpoint, append one operator message, and resume.
`src/tiny_llm/agent/__init__.py`	the names above	Complete the Day 6 exports within the final scaffold.

Run the cumulative learner checkpoint:

pdm run test --week 4 --day 6

Use this command for the supplied implementation:

pdm run test-refsol --week 4 --day 6

Before you implement the TODOs, all six Day 6 tasks are expected to fail. The command force-refreshes and runs the supplied learner tests for Days 1–6 together.

Start at a Safe Pause

run_to_checkpoint(...) saves after the assistant action and its tool result are both present:

original task
    ...
assistant tool action
tool observation
                  ^ inspect and steer here

This is deliberately not mid-token or mid-tool steering. No process is running in the background. The workspace is quiescent, and the checkpoint binds the original task, the complete message prefix, and the fake-model cache metadata.

Task 1: Derive a Public Status

Implement:

inspect_checkpoint(checkpoint, evidence_chars=160) -> AgentStatus

Validate the checkpoint, then require its final two messages to be a parsed assistant tool action followed by a Tool result:\n... user observation. Return four public facts:

AgentStatus(
    task="fix app.py and validate",
    last_action='{"new":"2","old":"1","path":"app.py","tool":"edit_file"}',
    last_evidence="edited app.py",
    next_step="resume the model after the completed edit_file observation",
)

last_evidence is a bounded prefix of the recorded observation. When it is too long, reserve the final character for an ellipsis so its length never exceeds evidence_chars.

The status function must not receive or call a model or workspace. It reads the validated immutable checkpoint and returns a new frozen value.

Task 2: State Only What the Harness Knows

The checkpoint knows the current task and the last completed action/result. It does not know what the model will decide next. Therefore next_step is a deterministic harness boundary:

resume the model after the completed edit_file observation

Do not replace this with a semantic guess such as “update the tests next.” That would present invented intent as agent state. A richer plan would need its own explicit, model-visible artifact; Day 6 does not add one.

Reject an incomplete boundary rather than producing a misleading card. A checkpoint ending with an ordinary user message, malformed assistant text, or anything other than a complete action/observation pair is not inspectable by this API.

Task 3: Append One Steering Message

Implement:

resume_with_steering(
    checkpoint,
    steering,
    fresh_generate,
    workspace,
    limits=None,
) -> AgentRun

Reject an empty or whitespace-only instruction. For a valid instruction, restore the checkpoint into the fresh generator and derive a mutable copy of the saved transcript. Append exactly one ordinary user message:

{
    "role": "user",
    "content": "Operator steering:\nvalidate before answering",
}

Then pass that list to the existing bounded loop. Do not add a separate queue, control channel, hidden prompt, or special protocol action.

Task 4: Keep the Message in Stable Order

The steering message belongs immediately after the saved checkpoint prefix. If the resumed model chooses another tool, the existing loop appends that assistant action and its observation after the steering:

saved task and evidence
Operator steering: validate before answering
assistant run_command action
validation observation
assistant final answer

Append the steering message once. Because the loop carries its message list forward, every later model call sees the same single message at the same position. Reinserting it on every call would duplicate the instruction and change the conversation.

Task 5: Continue Without Replaying Effects

The focused scenario starts with the original task “fix app.py and validate.” The first model reads and edits app.py, and the harness checkpoints after the complete edit observation. At that point:

app.py already contains the new value;
the edit approval happened once;
the receipt store contains the edit receipt; and
inspection reports the saved edit evidence without touching the workspace.

The operator adds “validate before answering.” A fresh scripted model restores the saved prefix, sees the original task, edit evidence, and steering message, then runs the exact allowed validation command and returns its final answer. The test proves the edit approval remains single, validation executes once, and the edit and command retain separate receipts.

Steering changes the next model input. It does not undo, replay, or rewrite the completed prefix.

Task 6: Keep the Boundary Small

Fail clearly when the checkpoint identity is invalid, the generator cannot restore the saved fake-model state, the evidence limit is not a positive integer, or steering is blank. Reuse AgentError, AgentCheckpoint, AgentLimits, AgentRun, and the existing loop rather than creating parallel versions.

The Day 6 module does not add concurrent interruption, mid-token control, a background worker or status server, a durable steering queue, session trees, branch/rewind, exactly-once reconciliation, or an evaluator. It is one visible pause → inspect → steer → resume path for the course’s scripted model.

Checkpoint

You can now pause after a complete observation, show an operator a bounded status made only from recorded facts, add one visible instruction, and resume a fresh model without replaying the completed effect. Inspect the model inputs in the focused test to verify the original task, saved evidence, and steering stay in order through a later tool turn and final answer.

Continue with Day 7: Evaluate Observable Outcomes to turn the final workspace, tool results, and durable receipts into a structured pass/fail report without grading hidden reasoning or exact transcript shape.

That report is another deterministic library checkpoint. After Day 9, the supplied Week 4 capstone provides the runnable path that carries checkpoint, compaction, steering, evaluation, branch selection, and bounded evidence together.

Day 7: Evaluate Observable Outcomes

🚧 Early-review WIP: This chapter is public for early review and may change. Use a disposable workspace when running the agent or enabling writes or commands.

The first six days built a small coding-agent loop, connected it to a workspace, recorded approved effects, and added checkpoint, compaction, and steering boundaries. The closing question is practical: did one run produce the outcome the task asked for?

Day 7 answers with a small deterministic evaluation harness. It checks declared observable facts: the final answer, exact file contents, tool-result evidence, and named durable receipts. It does not grade hidden reasoning or require one exact transcript shape.

The Starter Surface

The final Day 9 scaffold already exposes branching and bounded-evidence declarations. Leave those future TODO bodies alone. Day 7 owns one module:

File	Public names	Purpose
`src/tiny_llm/agent/evaluation.py`	`FileExpectation`, `ResultExpectation`, `ReceiptExpectation`, `EvaluationCase`, `EvaluationCheck`, `EvaluationReport`, `evaluate_run`	Describe required observable facts and produce a stable pass/fail report.
`src/tiny_llm/agent/__init__.py`	the names above	Complete the Day 7 exports within the final scaffold.

Run the cumulative learner checkpoint:

pdm run test --week 4 --day 7

Use this command for the supplied implementation:

pdm run test-refsol --week 4 --day 7

Before you implement the TODOs, all seven Day 7 tasks are expected to fail. The command force-refreshes and runs the supplied learner tests for Days 1–7 together.

The retired pdm run evaluate-agent launcher and its unconsumed static-grader packages are no longer part of the repository. The supported Day 7 learner checkpoint is evaluate_run through the cumulative test above.

Task 1: Declare the Outcome

An evaluation case names only the facts that matter for this task:

case = EvaluationCase(
    final_contains="validated",
    files=(FileExpectation("app.py", "answer = 2\n"),),
    results=(
        ResultExpectation("run_command", "validation passed"),
    ),
    receipts=(
        ReceiptExpectation(
            "call-1",
            "edit_file",
            "ok",
            "edited app.py",
            ("app.py",),
        ),
        ReceiptExpectation(
            "call-2",
            "run_command",
            "ok",
            "validation passed",
        ),
    ),
)

The two receipt IDs are explicit inputs chosen for this deterministic case. A different harness could discover or correlate effect records another way; Day 7 does not claim that every evaluation needs fixed call IDs.

Reject an invalid specification before evaluating: required strings cannot be blank, file paths must be relative and remain inside the workspace, file paths and receipt IDs must be unique within their groups, and a receipt exit state is either ok or error.

Task 2: Check the Final Answer

Implement:

evaluate_run(run, workspace, receipts, case) -> EvaluationReport

The first check requires a completed run whose public final answer contains the declared substring. This is a small grounding signal, not a prose grader. Do not inspect hidden reasoning, demand exact wording, or ask another model to judge the answer.

Task 3: Check Workspace State

For each FileExpectation, resolve the declared path through the existing Workspace boundary, read it as UTF-8, and compare the exact content. Emit one named check such as file:app.py.

A missing file, directory, unreadable file, or content mismatch is observed evidence that failed. Return a failed check instead of aborting the whole report. That is different from an invalid case definition, which is rejected.

Task 4: Match Result Evidence Without Grading a Trace

Each ResultExpectation requires at least one public AgentEvent with the declared tool and result substring. Search the events as a set of observable facts. Do not require an exact number of turns or an exact event order.

This matters because two useful runs may phrase their final answers differently or place unrelated read-only observations in a different order while producing the same required outcome.

Task 5: Check Named Durable Receipts

Use the public ReceiptStore passed to evaluate_run; do not reach through private workspace state. For every declared call ID, require the expected tool, exit state, result substring, and exact changed-artifact tuple.

Absent receipts, tampered persistent logs, mismatched fields, and lookup errors become failed checks. Evaluation must not append a receipt or rerun an effect.

Task 6: Produce a Stable Report

Return checks in one deterministic order:

final answer;
files in case order;
results in case order; and
receipts in case order.

EvaluationReport.passed is true only when every check passes. Its render() method should produce a compact summary:

evaluation: PASS
- final: PASS (required final observed)
- file:app.py: PASS (content matches)
- result:run_command: PASS (required result observed)
- receipt:call-1: PASS (receipt facts match)
- receipt:call-2: PASS (receipt facts match)

Stable names and ordering make failures easy to inspect without turning the test into an exact transcript comparison.

Task 7: Keep Evaluation Read-Only

The focused scenario asks the existing agent loop to set answer = 2 in app.py, run the exact configured validation command, and finish. The harness then checks the final answer, final file bytes, validation result, edit receipt, and command receipt.

Calling evaluate_run must leave the run, workspace bytes, modified-file list, approval history, and receipt bytes unchanged. It invokes no model, tool, or approval callback. Independent wrong final, file, result, and receipt facts each fail their own named check. An alternate final phrase and event order still pass when the required behavioral evidence is present.

Checkpoint

You can now turn one coding-agent run into a deterministic report over declared observable outcomes. This harness samples the facts a particular case names. It does not prove general task correctness, model quality, security, or production safety, and it is not a hidden grader, benchmark suite, or LLM-as-judge system.

You now have the evidence needed to compare continuations. Continue with Day 8: Fork, Steer, and Select to reuse one real token/KV prefix, steer two isolated branches, and explicitly choose a passing outcome without rewinding completed effects.

This evaluator is exercised as a library boundary. After Day 9, the supplied Week 4 capstone feeds its report into branch selection and then bounded evidence retrieval.

Day 8: Fork, Steer, and Select

🚧 Early-review WIP: Use only pre-created disposable workspaces. A control-state fork does not undo a file edit or any other completed effect.

Days 4 and 6 paused one agent and resumed one continuation. Day 8 asks a new question: after the model has inspected or changed the workspace, can we reuse the same inference prefix, try two explicit directions, and select the branch whose observable result is better?

The answer reconnects Week 4 to the inference system from Weeks 1–3. The course tokenizer renders the checkpoint conversation once. TinyKvFullCache stores the prefix keys and values for every model layer. Each branch gets a fresh control object and cache handles that share only those immutable prefix arrays, then decodes its own suffix. The branch report exposes the reused token count, the layer offsets, and the full-prefix prefill that was avoided.

This is control-state reuse, not effect rollback. Copy the already-modified disposable workspace and its completed receipt log before running either branch. Both copies begin with the same files and evidence; later receipts stay inside their branch.

The Starter Surface

The final scaffold already declares Day 9 evidence APIs. Leave those future TODO bodies alone. Day 8 owns one module and extends the approval result:

File	Public names	Purpose
`src/tiny_llm/agent/workspace.py`	`ApprovalDecision`	Carry an operator’s denial reason back as one ordinary model-visible observation.
`src/tiny_llm/agent/branching.py`	`PrefixReuse`, `KvPrefixGenerator`, `BranchOutcome`, `run_branch`, `select_branch`	Reuse a dense KV prefix, run isolated steered continuations, evaluate them, and make one explicit choice.
`src/tiny_llm/agent/__init__.py`	the names above	Complete the Day 8 exports within the final scaffold.

Run the cumulative learner checkpoint:

pdm run test --week 4 --day 8

Use this command for the supplied implementation:

pdm run test-refsol --week 4 --day 8

Before you implement the TODOs, all five Day 8 tasks are expected to fail. The command force-refreshes and runs the supplied learner tests for Days 1–8 together.

Before running the focused scenario, predict four values or facts: the shared base receipt in both roots, the receipt added only by the validate branch, the reused prefix length, and whether reused_tokens must equal avoided_prefill_tokens. The branch files, receipt bytes, cache offsets, and evaluation reports are the evidence.

Task 1: Return a Reason with a Denial

Add the immutable decision:

ApprovalDecision(approved=False, reason="keep the requested answer at 2")

A structured denial requires a nonblank reason. Workspace.execute returns that reason in its normal error: result so the next model turn can react to the operator’s instruction. It does not execute the effect or append a receipt. Existing callbacks that return plain True or False remain compatible.

The reason is steering, not a secret channel. Keep it short and suitable for the model-visible transcript.

Task 2: Save One Real Token and KV Prefix

KvPrefixGenerator.save_checkpoint(messages) renders the checkpoint messages without a generation prompt, tokenizes them with the course tokenizer, and prefills one TinyKvFullCache per layer. It records the exact token IDs and layer offsets in the existing Day 4 ModelCheckpoint.

The saved prompt must be an exact token prefix of every later steered prompt. Reject a continuation if even a same-length token differs. This binds cache reuse to content, not merely to a position.

fork() creates a fresh generator whose cache handles point at the frozen prefix arrays. When one branch grows, TinyKvFullCache assigns newly concatenated arrays to that branch. The frozen prefix and its sibling remain unchanged. This lesson deliberately uses the dense compatibility path; paged copy-on-write and radix serving are separate scaling topics.

Task 3: Expose What Was Reused

Each continuation reports:

PrefixReuse(
    reused_tokens=prefix_length,
    layer_offsets=(prefix_length, ...),
    avoided_prefill_tokens=prefix_length,
)

The first suffix model call starts at prefix_length; it must not call the model again at offset zero. These numbers make the inference boundary visible: the branch is not cloning only a Python transcript and silently recomputing the whole prompt.

Task 4: Fork Effects and Evidence Explicitly

Suppose the completed prefix changed app.py from answer = 1 to answer = 2 and wrote call-1, the edit receipt. Copy both the post-effect workspace and receipts.jsonl into two roots:

base after checkpoint
├── app.py              answer = 2
└── receipts.jsonl      call-1: edit_file

validate-only/          try-extra-edit/
├── app.py              ├── app.py
└── receipts.jsonl      └── receipts.jsonl

Both receipt files begin byte-identical and contain call-1. The validate-only branch appends call-2 after its exact allowed validation command. The other branch asks to change the answer again; the operator denies it with a reason, so its file and receipt bytes remain unchanged.

Construct each branch with its own Workspace and ReceiptStore, then call:

outcome = run_branch(
    "validate-only",
    "validate without another edit",
    checkpoint,
    prefix_generator.fork(),
    workspace,
    receipts,
    evaluation_case,
)

run_branch composes the Day 6 steered resume with the Day 7 observable-outcome evaluator. It does not copy a directory, infer an evaluation case, or merge effects for you.

Task 5: Select One Passing Branch

Make the choice explicit:

selected = select_branch(outcomes, "validate-only")

The name must identify exactly one outcome, and that outcome must pass its Day 7 report. Reject an absent selected name, a selected name that matches multiple outcomes, or a failing branch. Day 8 does not invent a hidden score or ask another model to judge the traces.

Manual Qwen/MLX Walkthrough

Complete Weeks 1–3 and the Day 8 TODOs first. Use a cached local Qwen model and the same dense compatibility path:

from mlx_lm import load
from tiny_llm import Qwen3ModelWeek3
from tiny_llm.agent import KvPrefixGenerator, create_checkpoint

mlx_model, tokenizer = load("Qwen/Qwen3-0.6B-MLX-4bit")
model = Qwen3ModelWeek3(mlx_model, enable_paged_attention=False)
prefix_generator = KvPrefixGenerator(model, tokenizer, max_tokens=128)

This section is an integration outline, not a complete cached-Qwen program. You still have to construct paused, both copied workspace/receipt roots, the evaluation cases, and each scripted or live continuation. The supplied capstone composes those mechanisms with a deterministic checkpoint model; it does not make this nondeterministic model walkthrough a correctness test.

First create a Day 4 checkpoint named paused after a complete tool observation. Its messages are the control boundary you want to share, while its model field belongs to the generator that created it. Rebind those exact messages to the real tokenizer and dense KV cache before resuming:

messages = [
    {"role": role, "content": content}
    for role, content in paused.messages
]
model_checkpoint = prefix_generator.save_checkpoint(messages)
checkpoint = create_checkpoint(paused.task, messages, model_checkpoint)

Now fork two fresh generators with prefix_generator.fork(). Give them different visible steering messages and the two workspace/receipt copies described above, passing the rebound checkpoint to each run_branch call. Print each outcome.reuse, render both evaluation reports, and select the passing name.

Model responses are nondeterministic, so this walkthrough is manual. Inspect the actual proposed actions, approval reason, final file bytes, receipt logs, and evaluation reports. The deterministic learner test covers the same public boundary with a tiny tokenizer/model and no download.

reused_tokens and avoided_prefill_tokens show exact logical prefix reuse. They are not a wall-clock speedup, throughput measurement, or memory-capacity claim.

Checkpoint

You can now connect a Day 4 control checkpoint to the actual tokenizer and KV cache path, reuse one immutable prefix for two isolated continuations, expose a denial reason to the model without recording an effect, evaluate both branches from declared evidence, and choose one passing result.

Completed effects were copied, not rewound. The two branches do not run concurrently, share a mutable workspace, merge receipts, or provide a session server/tree. Day 8 teaches the boundary visibly before adding any serving-scale machinery.

Continue with Day 9: Bound Tool Evidence to keep oversized results verifiable without placing their complete bytes in each later model prompt.

The Day 9 test starts its own loop rather than consuming the selected Day 8 outcome. After Day 9, the supplied Week 4 capstone exercises that final handoff in one deterministic scenario.

Day 9: Bound Tool Evidence

An agent can read a log that is much larger than the useful part. Appending the whole result to every later model prompt wastes context, but silently slicing it loses the evidence needed to verify what happened.

Day 9 keeps those two concerns separate:

preserve the exact UTF-8 tool-result bytes outside the prompt;
give the model a bounded observation with identity, size, digest, and head/tail previews;
let the model request one explicit byte range and continue in the unchanged agent loop.

This is byte selection, not semantic summarization. The model decides which range to inspect from visible facts.

Files You Implement

This is the final declaration scaffold, so every earlier Week 4 module is already visible. Day 9 owns only the TODO bodies in the evidence adapter and its exports:

File	Public names	Responsibility
`src/tiny_llm/agent/evidence.py`	`ArtifactRef`, `ArtifactStore`, `BoundedEvidenceWorkspace`	Store exact results, render bounded observations, and serve explicit ranges.
`src/tiny_llm/agent/__init__.py`	the names above	Export the completed Day 9 declaration surface.

The protocol, loop, workspace, receipts, and Days 1–8 modules do not change. BoundedEvidenceWorkspace is a small adapter around the existing Workspace.

Run the cumulative learner checkpoint:

pdm run test --week 4 --day 9

Before you implement the TODOs, the implementation-dependent cases across six tasks are expected to fail; the shared constructor-validation cases already pass. The command force-refreshes and runs all nine supplied Week 4 learner tests together.

Use this command for the supplied implementation:

pdm run test-refsol --week 4 --day 9

Task 1: Give Exact Bytes an Identity

ArtifactStore.put(result) encodes the complete result as UTF-8, writes those bytes under its explicit artifact root, and returns:

ArtifactRef(
    artifact_id="artifact-<lowercase SHA-256>",
    byte_count=...,
    sha256="<lowercase SHA-256>",
)

The content-addressed ID and full digest deliberately repeat the same hash in different roles: one is the handle used by the range request; the other is a separately labeled model-visible verification field. The store registers the ID in memory. A different store cannot retrieve it merely because the caller guessed the filename.

Before every range read, verify the stored byte count and digest again. The course store is local and single-process. It does not promise retention, garbage collection, encryption, access control, or a network blob service. It preserves the exact bytes returned by the wrapped tool; earlier tool-level limits, such as Day 3’s command-output cap, still apply before this adapter.

Task 2: Replace Only Oversized Successful Results

Wrap an existing workspace:

from tiny_llm.agent import ArtifactStore, BoundedEvidenceWorkspace

bounded = BoundedEvidenceWorkspace(
    workspace,
    ArtifactStore(artifact_root),
    max_inline_bytes=512,
    preview_bytes=64,
    max_range_bytes=512,
)

Short results and every error: observation remain byte-for-byte unchanged. For a successful result larger than max_inline_bytes, persist the full bytes and return a compact JSON observation containing:

artifact_id, byte_count, and sha256;
valid UTF-8 head and tail previews with their byte ranges;
the omitted half-open byte interval;
one exact read_file range-request example.

The entire compact observation, including metadata and previews, must fit max_inline_bytes. Reduce previews at UTF-8 boundaries when the metadata needs more space. The supplied constructor already requires max_range_bytes >= 4, so the default range can always hold one maximum-width UTF-8 code point. Preserve that supplied rule; the remaining externalization and range behavior is learner-owned. Never split a code point or silently replace one.

Task 3: Reuse the Existing Tool Protocol

Day 9 does not add a new action schema. It reserves one virtual relative-path namespace for the existing read_file action:

.tool-artifacts/<artifact-id>/bytes/<start>-<end>

[start,end) is an exact half-open byte range. The adapter intercepts the reserved prefix before the real workspace sees it. A successful reply names the same artifact, total size, digest, start, end, returned byte count, and the strictly decoded UTF-8 data.

The reply is not sent back through externalization. Its selected data is already limited by max_range_bytes.

Task 4: Fail Closed Without Leaking

Every path beginning with .tool-artifacts/ belongs to the virtual namespace. Malformed paths must not fall through to a learner file of the same name.

Return short ordinary error: observations for:

an invalid or unknown artifact ID;
negative, reversed, out-of-bounds, or oversized ranges;
stored bytes whose size or digest changed;
a range that cuts through a UTF-8 code point.

Do not print the host artifact-root path, enumerate known IDs, or reveal bytes from another store while reporting an error.

Task 5: Continue Through the Same Loop

The deterministic test creates a large ASCII build log whose diagnostic is outside both previews. A scripted model performs three normal steps:

read_file build.log
        |
        v
bounded identity + previews
        |
        v
read_file .tool-artifacts/<id>/bytes/<start>-<end>
        |
        v
exact diagnostic range -> final answer

run_agent is unchanged. Its first event contains only the bounded observation; the second contains only the selected range; the artifact file still matches the complete original result.

Before running it, predict whether the first result stays inline or becomes an artifact, the omitted half-open byte interval, and the exact bytes returned by the model’s range request. The first observation, artifact digest/file, and second observation must agree; a final model sentence is not the evidence.

Task 6: Preserve the Workspace Contract

Delegate policy, available_tools, and modified_files to the wrapped workspace. This lets build_system_prompt, action validation, and the existing event loop operate without knowing about the storage adapter.

The virtual range path is still a normal JSON read_file request, so the learner does not need a second parser or a replacement generation interface.

Manual Cached-Qwen Walkthrough

Complete the Day 9 TODOs first. Create separate disposable workspace and artifact directories, put a large UTF-8 build.log in the workspace, and use the same local-model adapter as the exploratory Week 4 exercise:

import hashlib
from pathlib import Path
from tempfile import TemporaryDirectory

from mlx_lm import generate as mlx_generate, load
from tiny_llm.agent import (
    ArtifactStore,
    BoundedEvidenceWorkspace,
    ToolPolicy,
    Workspace,
    run_agent,
)

workspace_directory = TemporaryDirectory(prefix="tiny-llm-day9-workspace-")
artifact_directory = TemporaryDirectory(prefix="tiny-llm-day9-artifacts-")
workspace_root = Path(workspace_directory.name)
artifact_root = Path(artifact_directory.name)
(workspace_root / "build.log").write_text(
    "build started α\n"
    + "x" * 256
    + "\nERROR code=E42 dependency mismatch\n"
    + "y" * 3_000,
    encoding="utf-8",
)

mlx_model, tokenizer = load("Qwen/Qwen3-0.6B-MLX-4bit")
artifacts = ArtifactStore(artifact_root)
workspace = BoundedEvidenceWorkspace(
    Workspace(ToolPolicy(workspace_root, max_file_bytes=64_000)),
    artifacts,
)

def generate(messages):
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    return mlx_generate(
        mlx_model, tokenizer, prompt, max_tokens=256, verbose=False
    )

run = run_agent(
    "Read build.log. If it is externalized, retrieve one useful byte range.",
    generate,
    workspace,
)

for event in run.events:
    print(event.result)
print(run.final)
for artifact_path in artifact_root.iterdir():
    data = artifact_path.read_bytes()
    print(artifact_path.name, len(data), hashlib.sha256(data).hexdigest())

Save that fragment as /tmp/tiny-llm-day9.py, then run the repository source explicitly so an older editable install cannot be imported by accident:

PYTHONPATH=src pdm run python /tmp/tiny-llm-day9.py

Model choices vary. Inspect the actual first observation, requested artifact ID and range, returned bytes, final answer, and on-disk artifact digest. Do not use a workspace or artifact root containing secrets. After inspection, call workspace_directory.cleanup() and artifact_directory.cleanup().

Invalid or repeated model JSON is an ordinary bounded outcome: the loop may stop at its invalid-action or step budget without externalizing anything. That does not invalidate the deterministic checkpoint, and a successful manual run does not become a correctness or performance result.

Checkpoint

You can now keep a complete large tool result available for verification while placing only bounded facts in the model context. The model can retrieve an explicit range by identity and continue through the same tokenizer and agent loop.

Day 9 does not summarize the result, stream concurrent chunks, retain artifacts for production, or add a network service.

The focused checkpoint starts a standalone bounded-evidence loop. After it is green, run the composed product witness:

pdm run week4-capstone

The command uses the completed learner package in a disposable deterministic scenario. Its sorted JSON reports compaction accounting, both steered branches and their evaluations, the explicitly selected validate-only branch, and the exact E42 artifact range retrieved by identity. The reported token, reuse, and byte counts remain accounting evidence; they do not prove latency, throughput, quality, or memory-capacity improvement.

🚧 Appendix: Performance Evidence Ledger

Status: Experimental, single-machine evidence. See the Week 2 verification matrix before treating a correctness, integration, or performance result as broader proof.

This appendix records the measurements that determined the course order. The numbers are not additive promises: after one bottleneck shrinks, every other operator becomes a larger fraction of model time.

Benchmark Method

The progression runner launches every checkpoint in a fresh process, alternates their order, performs complete-request warmups, synchronizes lazy MLX work inside the timer, and reports the median:

pdm run bench-week2-progression --offline --solution tiny_llm --repeats 2 \
  --model qwen3-4b --input-len 128 --output-len 129 --warmup 2 \
  --prefill-logits last \
  --json-output benchmark_results/task367-final-main/raw/week2-128-final-main.json

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/task367-final-main/raw/week3-serving-final-main.json

--prefill-logits last is a generation-serving workload: both the reference solution and MLX project only the last prompt row into vocabulary logits. Use --prefill-logits all for prompt scoring, but never compare the two modes. Decode throughput excludes the first generated token because that token is produced by prefill.

MLX’s published mlx_lm.benchmark table uses a 2,048-token prompt and 128 generated tokens. That makes 2K/128 a useful static-library comparison point, not a paging acceptance test or a long-context proof. Use a context sweep:

Point	Purpose
128	fixed Week 2 acceptance and short interactive requests
2,048	standard MLX-style static stress comparison
8,192	long-context attention and KV-cache stress
16,384	stress point after the 8K path is healthy

llama-bench commonly uses prompt-processing 512 and token-generation 128 by default, which is another reminder that benchmark lengths are conventions, not universal workloads. Always publish the exact prompt and output lengths.

The measured machine below is an Apple M4 Pro with a 20-core GPU and 64 GB of memory. The current Week 2 control uses two complete warmups and two balanced fresh processes; the continuous-serving rows use one warmup and the median of four balanced fresh processes.

Week 2 Checkpoint Retention Ledger

A polished explanation is not evidence that an optimization belongs in the course. Before retaining a checkpoint, answer six questions: its invariant, why it could be faster, where it wins, where it loses, its fallback, and how the benchmark could mislead us. This ledger records the current answers; links below contain the measurements.

Checkpoint	Required invariant	Performance hypothesis	Retained range and losing shapes	Fallback or control	Main benchmark trap
Dense KV cache	Caller offset equals every layer cache length; K/V append on the sequence axis	Reuse projected prefix K/V instead of recomputing the full model prefix	Wins incremental decode as the prefix grows; repeated `concat` still copies `O(S²)` bytes	Week 1 full-prefix model remains the semantic control; Week 3 pages replace growth copies	Comparing cached MLX with an uncached course model measures different algorithms
Packed quantized matvec	W4, group size 128, BF16 parameters, contiguous packed layout, and the declared transpose convention	Read packed weights once and share unpack/scale work across SIMD lanes	Retained for `M <= 8`; multi-row prefill exposes poor reuse and motivates Day 5	The Python `mlx.core` equation is the correctness oracle; vanilla W4 is an inspectable Metal control; named earlier checkpoints preserve the dense control	Lazy execution or timing post-materialized weights can hide weight traffic
RMSNorm	BF16 I/O with the sum of squares accumulated in FP32	Fuse reduction, normalization, and weight multiply into one dispatch	Retained at Qwen hidden dimensions after both operator and decode gains; unknown dimensions require remeasurement	Python `mlx.core` RMSNorm and the Day 3 checkpoint remain selectable	Adding isolated microseconds as if checkpoint gains were independent
RoPE	One valid offset per batch row; even rotated dimension; tail values preserved	Fuse angle generation and pair rotation without intermediate graphs	Retained for Qwen decode rows; head-count and rotated-dimension changes require remeasurement	Python `mlx.core` RoPE and the RMSNorm-only checkpoint remain selectable	Benchmarking a cached or precomputed angle path against fresh angle construction
SwiGLU	Gate and up tensors have identical shape and dtype	Fuse SiLU and the gate/up product into one elementwise dispatch	Retained for Qwen MLP shapes; tiny tensors and other dtypes are not a performance claim	The Python `mlx.core` SiLU-product and the RoPE checkpoint remain selectable	Accepting an operator win without a repeated complete-model gain
Decode attention (optional lab)	`Hq % Hkv == 0`, `D <= 256`, FP32 online-softmax state, and causal/explicit mask semantics	Avoid score/probability tensors and merge softmax while walking K/V	The checked fixed-workload result is equivocal; retain only for an explicitly measured context and fallback	Python `mlx.core` grouped attention handles unsupported shapes and is the control	Prescribing the lab from chapter order, extrapolating one context, or promoting `n=2` product noise
SIMD-matrix prefill	W4/group-128 layout, BF16 storage, FP32 tile accumulation, and correct partial tiles	Reuse activation and dequantized-weight tiles across prompt rows	Required path for `M > 8`; partial and new model shapes need both correctness and timing sweeps	The Python `mlx.core` matmul is the correctness oracle; Day 3 matvec remains the short-row dispatch and vanilla Metal is a bring-up control	Comparing all-logit course prefill with last-logit MLX serving
Split-K prefill	Partitions align to quantization groups; partial planes are disjoint; final reduction is FP32	Add independent groups only while the ordinary result grid is under-filled	Conditionally retained at the measured 32-token control; rejected at the fixed 128-token product workload	`split_k <= 1` dispatches exactly to the Day 5 unsplit kernel	Static dispatch does not prove occupancy, and a short-shape replay does not prove a fixed-workload gain

This is a retention ledger, not a portability certificate. A new GPU, MLX release, model shape, dtype, or workload reopens the corresponding row.

Long-Context Budget for Week 4

Context length has separate model, memory, and latency limits. For the course Qwen3-4B checkpoint, one token of BF16 K/V state occupies

36 layers * 2 (K and V) * 8 KV heads * 128 values * 2 bytes
    = 147,456 bytes = 144 KiB per token

The checkpoint declares max_position_embeddings = 65,536, but its rope_scaling field is empty. Qwen documents that Qwen3 training covers 32,768 tokens and recommends RoPE scaling for substantially longer inputs. The unmodified course model therefore has a 32,768-token validated limit even though its configuration permits a larger position experiment.

Memory is not the binding limit on the measured 64 GB M4 Pro. MLX reports a 51.84 GiB recommended GPU working set, and the quantized checkpoint occupies 1.99 GiB. Reserving 8 GiB for activations, allocator slack, and outputs gives

floor((51.84 GiB - 1.99 GiB - 8 GiB) / 144 KiB) = 304,738 tokens

That estimate is a capacity calculation, not permission to exceed the model’s trained range. The course limit is the minimum of the limits:

min(32,768 trained, 65,536 configured, 304,738 memory) = 32,768 tokens

Week 4 uses 32,768 total tokens as its hard context budget. It starts compaction before the rendered input exceeds 24,576 tokens, reserving 8,192 tokens for the next model response and a large tool result. The tokenizer must count the complete rendered request, including system instructions and tool schemas.

What Becomes Slow at 300K

FlashAttention removes the quadratic score-matrix allocation; it does not remove the work. Full-attention prefill remains quadratic in context length, so 300K contains about 84 times the attention work of 32K. One-token decode must read a linearly growing K/V history at every layer.

The following synthetic operator sweep uses MLX 0.32.0, one Qwen3-4B-shaped BF16 decode query, three fresh processes, and the median of fifteen synchronized dispatches per process. The final column sums the isolated layer latency across 36 layers and is an optimistic attention-only ceiling; a complete model must also run projections, normalization, sampling, and cache updates.

Context	Full-model BF16 KV	MLX SDPA per layer	Attention-only decode ceiling
2,048	0.28 GiB	0.14 ms	195.33 tok/s
8,192	1.12 GiB	0.29 ms	96.72 tok/s
32,768	4.50 GiB	0.92 ms	30.28 tok/s
65,536	9.00 GiB	1.73 ms	16.08 tok/s
131,072	18.00 GiB	3.65 ms	7.61 tok/s
300,000	41.20 GiB	9.49 ms	2.93 tok/s

The 300K operator allocation runs on this M4 Pro, but an end-to-end 300K run of the course checkpoint would be outside its configured and training ranges, would leave little working-set headroom, and would make initial prefill impractical. It is useful as a kernel stress test, not as a supported course context.

MLX contains several long-context optimizations. Its fused GQA decode path automatically switches to a context-partitioned two-pass reduction; the 0.30.4 release specifically calls out faster long-context vector GQA. Multi-token attention uses a tiled fused path, and MLX-LM chunks prompt evaluation to bound temporary activations. MLX-LM also offers prompt-prefix reuse, a rotating fixed-size cache, and quantized KV storage. Prefix reuse helps repeated prompts; cache rotation changes full-attention semantics; and KV quantization trades numerical precision and sometimes speed for capacity. None makes the first full 300K prefill linear-time.

Reproduce the operator sweep with:

pdm run bench-long-context-attention \
  --json-output benchmark_results/m4-pro-qwen3-4b-long-context-mlx-0.32.0.json

Dependency Upgrade

The project upgraded from MLX 0.29.1 to 0.32.0 and from the mlx-lm 0.28 series to 0.31.3. A matched Qwen3-4B run showed:

Context	Metric	MLX 0.29.1	MLX 0.32.0	Change
128	Prefill tok/s	825.48	828.34	+0.35%
128	Decode tok/s	88.32	88.08	-0.27%
2,048	Prefill tok/s	816.73	820.85	+0.50%
2,048	Decode tok/s	78.42	74.81	-4.60%

The small differences show why the comparison must record exact dependency versions: the MLX denominator is part of the experiment, even when an upgrade does not materially change the result.

Week 2 Performance by Chapter

This section is a checked example of the course’s discover → optimize → re-profile loop. It is not a portability certificate or a set of performance thresholds.

Bound Evidence

The run used exact source add389b747793e910f0506f5720dd0aac373d126 on one Apple M4 Pro with 20 GPU cores and 64 GB unified memory, macOS 27 build 26A428, gpudebug 1.0, Python 3.12.13, MLX 0.32.0, mlx-lm 0.31.3, and Qwen3-4B-MLX-4bit from the local cache.

The fixed product control used a 128-token prompt, 129 output tokens, last-row prefill logits, seed 0, two synchronized warmups, and two balanced fresh-process samples. The attribution cases used four warmups and twelve balanced synchronized iterations. With n=2, product medians can reject a large contradiction; they cannot turn a sub-percent change into a portable claim.

Learners reproduce the method with their own tiny_llm solution:

pdm run bench-week2-progression --offline --solution tiny_llm --repeats 2 \
  --variant week2-kv-cache --variant week2-quantized-matvec \
  --variant week2-swiglu --variant week2-simd-matmul \
  --variant week2-split-k --variant mlx \
  --model qwen3-4b --input-len 128 --output-len 129 --warmup 2 \
  --prefill-logits last --json-output week2-progression.json

pdm run profile-week2-kernels --solution tiny_llm --model qwen3-4b \
  --case kv-cache:decode:128 --case quantized-matvec:decode:128 \
  --case swiglu:decode:128 --case swiglu:prefill:128 \
  --case simd-matmul:prefill:128 --case split-k:prefill:128 \
  --warmup 4 --iterations 12 --json-output week2-attribution.json

The checked compact result is benchmark_results/m4-pro-qwen3-4b-week2-gpudebug-macos27-mlx-0.32.0.json. It records unavailable evidence explicitly. It contains no raw trace, absolute workspace path, screenshot, token output, or portable timing claim.

Day 1: Cache the Prefix

Day 1 changes the generation algorithm: prefill once, retain dense K/V state, and send only the new token through each decode step. The matched Week 1 versus kv-cache product observation measures that algorithmic change before any kernel is replaced. A shader trace is not needed to justify the cache.

Day 2: Discover the First Operator Category

The cached-decode attribution reported 34.527 ms of projection work, or 83.9% of the attributed total. Two BF16 GEMV shaders accounted for 93.60% of the available shader ranking. This selected dense projection weight traffic as the first bounded target.

The evidence-to-next-change decision was: pack W4 weights, change only the projection path, and repeat the identical decode workload. A failure to reduce projection time, or a regression in matched product decode, would falsify the hypothesis.

Day 3: Keep Weights Packed

The packed W4 candidate reduced attributed projection time from 34.527 ms to 10.700 ms (-69.0%) and reduced total attributed time by 56.8%. On the fixed-workload two-sample product control, decode rose from 24.38 to 58.90 tokens/s (+141.6%).

The next re-profile mattered as much as the speedup: normalization, position, and activation work now occupied 5.948 ms, or 33.5% of attributed time. That newly exposed category selected the fused Day 4 operators.

Day 4: Fused Model Kernels

Fused RMSNorm, RoPE, and SwiGLU reduced the selected category from 5.948 ms to 1.251 ms (-79.0%) and total attributed time by 27.3%. The product control improved at every cumulative substep: RMSNorm +10.7%, RoPE +8.7%, and SwiGLU +4.9%.

After the full Day 4 checkpoint, projections again dominated decode at 10.516 ms / 81.4%, while attention was 0.837 ms / 6.5%. At 128-token prefill, the portable attribution put projections at 1,201.306 ms / 99.1%. That prefill result—not a predetermined chapter order—selected SIMD-matrix prefill for Day 5.

Day 5: Restore Matrix-Shaped Prefill

The cooperative W4 SIMD-matrix schedule reduced attributed 128-token projection time to 163.172 ms (-86.4%) and total attributed time by 85.8%. Fixed-workload prefill rose from 106.44 to 721.60 tokens/s (+577.9%). The succeeding capture ranked the SIMD-group W4 matrix shader at 96.86% of available shader cost.

These effects justify retaining the schedule for this source tree and workload. They do not establish the same gain on another Apple GPU, model, prompt length, or dependency version.

Day 6: Keep the Secondary Operator Lab Optional

After Day 4, the checked decode-attention branch changed attributed attention from 0.837 ms to 0.831 ms (-0.75%), while total attributed time rose 0.97%. The separate product control showed a small decode change from 74.34 to 76.50 tokens/s (+2.91%). Those mixed signals support an inconclusive worked branch, not a universal bottleneck or a prerequisite for Day 7.

The capture did confirm that the custom attention shader ran: it accounted for 10.27% of available shader cost while packed projections accounted for 82.83%. That is useful mechanism evidence, but it does not make the optional branch the next dominant optimization.

Day 7: Split K Only Where the Shape Supports It

At 32-token prefill, the unsplit SIMD projection replay exposed an under-filled schedule. Split-K reduced attributed projection time from 48.433 ms to 46.008 ms (-5.01%) and total attributed time by 4.87%. Static inspection found both Split-K and reduction dispatches, but the replay produced no timeline, shader ranking, or counter tree, so no occupancy improvement is inferred.

The fixed 128-token product control rejects a broad claim: prefill changed from 721.60 to 718.36 tokens/s (-0.45%) and decode changed by +0.14%. The checked decision therefore conditionally retains Split-K for the measured short shape and rejects it for the fixed 128-token product workload. Another device or model needs a fresh crossover measurement.

What the Capture Can and Cannot Add

Six of eight checked captures exposed complete shader/counter detail. The pre-SIMD 128-token prefill capture exposed timeline counters but no shader or command ranking. The 32-token Split-K capture exposed only static dispatch. Missing trees remain unavailable; they are not recorded as zero and do not support inferred counters.

The optional macOS 27 profiling lab shows how to create a trace package, hash its files, reduce gpudebug output, record a three-sentence decision, and remove the raw package after preserving compact evidence. The portable benchmark and attribution path remains sufficient for every required checkpoint.

Week 3 Performance by Chapter

Paging adds indirect K/V reads and is not expected to beat contiguous attention for one preallocated static request. Week 3 therefore measures a serving workload with request turnover, incremental unknown-size growth, chunked admission, dense batch reconstruction, and page reuse:

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/task367-final-main/raw/week3-serving-final-main.json

A complete warmup compiles the kernels. The runner then synchronizes and resets every page pool, so the measured paged run starts with zero pages and zero backing capacity.

Ownership and denominators

The projection boundary must be fixed before interpreting any Week 3 table:

Evidence row	Projections	Cache / attention / paging / scheduler	What it establishes
Week 2 SIMD or Split-K	Course-owned zero-Steel W4 kernels, loader, and direct SIMD-matrix helper	Course-owned Week 2 dense cache and operators	Week 2 course implementation versus its explicitly paired full-MLX row.
Week 3 course row	Explicit MLX quantized-projection seam	Course-owned cache, attention, paging, batching, and scheduling	Representative cumulative Week 3 behavior; it does not isolate the seam.
Full `mlx` row	Full MLX model/operator	Full MLX	External denominator, distinct from the hybrid Week 3 course row.
Task #360 seam versus inherited	MLX quantized projections versus inherited Week 2 course projections	Identical course-owned Week 3 mechanisms	Causal projection-seam effect on one measured source tree.

Task #360 and task #367 answer different questions. The former is a causal ablation; the latter is representative final-main absolute evidence. Do not splice one campaign’s absolute values into the other or credit its projection gain to paging, FlashAttention, or scheduling.

The Days 1–2 chunk-size control uses one deterministic Qwen3-0.6B trace with seed 0, eight 64–512-token prompts, a fixed 32-token output budget, and four balanced fresh processes. A gap is measured between synchronized decode-call completions only while a decode request is active. Every row uses the same Week 3 projection seam and course-owned mechanisms; only the budget changes:

Prefill budget	Output tok/s	Prefill tok/s	Decode tok/s	Requests/s	Decode step p95	Decode gap p95 / max
32	105.23	2,549.62	181.77	3.288	15.82 ms	30.01 / 52.62 ms
128	153.82	4,215.12	242.23	4.807	17.79 ms	45.36 / 53.76 ms
512	170.46	4,769.14	262.01	5.327	17.11 ms	73.56 / 119.90 ms

Because 512 covers every prompt in this trace, that row is the full-prompt Day 1 control. Relative to it, 128 gives up 9.8% output throughput while reducing the p95 completion gap by 38.3% and the maximum by 55.2%. The course chooses 128 for this trace, not as a universal chunk-size threshold.

The Day 4 operator control uses B=1, Hq=32, Hkv=8, L=1, D=128, BF16, and 128-token pages. Each row is the median of four balanced fresh-process medians, each containing 60 synchronized calls after five warmups:

Context	Dense + gather	Direct paged	MLX fused
128	201.26 us	228.58 us	188.79 us
1,024	468.39 us	299.14 us	250.04 us

The direct operator is 13.6% slower than dense-plus-gather at 128 tokens and 36.1% faster at 1,024 tokens. MLX remains faster at both shapes. Outputs match the dense BF16 equation within 0.00439453125 and 0.001953125 respectively. This operator contains no model projection and therefore isolates the attention paths directly.

Chapter	Measured checkpoint	Primary result	Change from the preceding comparable path
Day 1	Continuous scheduler	Defines request turnover and active-batch throughput.	Establishes the serving workload.
Day 2	Chunked admission with dense reconstruction	711.18 prefill; 35.23 output; 57.59 decode tok/s	Establishes the dense serving baseline.
Day 3	Paged storage with compatibility gather	725.46 prefill; 41.64 output; 78.53 decode tok/s	+18.2% output; +36.4% decode; -50.6% copy volume.
Day 4	Correct direct paged behavior	105.01 aggregate decode tok/s in the cumulative endpoint	Removes dense K/V reconstruction; this corpus does not isolate Day 4’s scalar prefill.
Day 5	BF16 long-prefill tiled schedule	No isolated scalar-versus-tiled row	The cumulative serving row below includes Day 5 but is not causal evidence for it.

Day 1 introduces scheduling, not a kernel speedup. Day 2 makes the hidden cost measurable: appending one token still reconstructs a padded dense batch. Day 3 makes pages canonical but retains gather_dense() as a compatibility checkpoint. Day 4 then removes that compatibility movement for every query shape. Day 5 changes only the internal schedule for supported BF16 long prefill.

Days 4 and 5 share the final direct-paged process: queries with L <= 8 dispatch to the Day 4 decode schedule, supported BF16 long-prefill calls use the Day 5 tiled schedule, and generic shapes retain a direct scalar fallback. The phase timers report decode and prefill throughput inside the same request trace; they do not isolate the Day 5 schedule.

Every headline number above comes from the same continuous-batch campaign. The cumulative serving endpoints are:

Storage and attention path	Prefill tok/s	Output tok/s	Decode tok/s	Requests/s	Peak KV MiB	Avoidable KV copy MiB
Dense growth and reconstruction	711.18	35.23	57.59	0.469	1,096	209,532
Paged storage plus dense gather	725.46	41.64	78.53	0.555	not a total peak	103,445
Direct paged attention	672.68	46.36	105.01	0.618	576	504

The same raw serving artifact reports synchronized decode-call latency and the completion gaps that include intervening prefill and scheduler work:

Path	Decode step median / p95 / max	Completion gap median / p95 / max
Dense reconstruction	51.03 / 84.49 / 124.52 ms	53.16 / 248.30 / 309.74 ms
Paged + gather	39.80 / 52.79 / 80.09 ms	41.82 / 225.64 / 261.38 ms
Direct paged	28.97 / 36.78 / 63.04 ms	30.16 / 222.18 / 239.49 ms

The compatibility row omits peak storage because an exact peak must include both the page pool and temporary dense staging allocation. Its other counters remain directly comparable.

Direct paged attention is 5.4% lower on prefill, 31.6% higher on output/request throughput, 82.3% higher on decode, and 47.4% lower on measured peak KV storage relative to dense serving. Avoidable logical copy volume falls by 99.76%. Relative to paged storage plus gather, it is 7.3% lower on prefill, 11.3% higher on output/request throughput, 33.7% higher on decode, and removes 99.51% of the remaining copy volume. These cumulative system results do not isolate the Day 5 prefill kernel or prove a short-chunk FlashAttention win.

The 8K static run remains a secondary kernel diagnostic, not a Week 3 headline or acceptance result. At that shape, the Week 3 seam plus course paged path raises prefill from the Week 2 path’s 323.96 to 463.69 tok/s, a 43.1% gain, and reaches 72.5% of the 639.73 tok/s full-MLX row. This does not isolate the projection seam, measure request turnover or admission capacity, or establish long-context support. One-token decode continues to dispatch to the Day 4 vector schedule.

Separate causal projection-seam result

Task #360 holds the Week 3 mechanisms fixed and changes only projection ownership on measured source 170211be3503c0ec0b1fa75bbb3b0c23a86bd3ac:

Causal comparison	MLX seam effect versus inherited Week 2 projections
Chunked prefill, step 512	+10.64% prefill; +11.91% output
Chunked prefill, step 128	+11.74% prefill; +11.76% output
Dense Day 3	+12.17% prefill; +16.82% output; +18.86% decode
Serving	+7.72% prefill; +9.42% output; +13.02% decode

Full MLX remains 17.83% faster than the dense Day 3 seam on prefill (equivalently, the seam is 15.13% below full MLX), because the seam changes projections only. These causal percentages explain the ownership decision; the task #367 tables above provide current absolute values.

The checked-in final-main corpus contains the complete raw samples, exact source commit and tracked-clean flag, host, configuration, execution order, and—where requests are generated—the exact request trace and its checksum:

benchmark_results/task367-final-main/raw/week2-32-final-main.json
benchmark_results/task367-final-main/raw/week2-128-final-main.json
benchmark_results/task367-final-main/raw/week2-2048-final-main.json
benchmark_results/task367-final-main/raw/week2-prefill-operators-final-main.json
benchmark_results/task367-final-main/raw/week3-chunked-prefill-final-main.json
benchmark_results/task367-final-main/raw/week3-attention-final-main.json
benchmark_results/task367-final-main/raw/week3-serving-final-main.json
benchmark_results/task367-final-main/raw/week3-8k-final-main.json

Verify the manifest, all eight raw files, and the evidence ledger with:

(cd benchmark_results/task367-final-main && \
  shasum -a 256 -c task367-final-main-sha256.txt)

Copy counters report logical operation volume, not hardware DRAM traffic. Dense volume includes old K/V copied during each request-cache growth and live K/V copied into a newly padded batch tensor at every decode step. Paged volume includes old physical pages copied only when a layer’s geometric pool grows. Appending a token writes only its page slice, and later requests reuse freed pages.

The raw counters make reuse, fragmentation, logical copy volume, and measured KV headroom visible; static single-request latency cannot. Logical copy volume is not hardware DRAM traffic, and none of these counters establishes admission capacity without a memory-capped sweep.

The workload validates continuous batching, chunked prefill, incremental growth, and page reuse. Prefix sharing and speculative decoding require separate traces with shared prefixes or cache rewind events and are not claimed by this result.

Week 2 Profiling Boundary

The synchronized product benchmark and portable operator-attribution runner are the required evidence path. Metal capture, Xcode visualization, gpudebug, and screenshots remain optional and require macOS 27. The compact checked result records unavailable trees instead of substituting zeros or inferring counters; learners without that toolchain can still complete every checkpoint and reason from the portable artifact.

Optimization Map

Measured bottleneck	Retained change	Chapter
Full-prefix decode recomputation	Dense request KV cache	Week 2 Day 1
Dense projection weight traffic	Packed W4A16 x4 SIMD matvec	Week 2 Day 3
Repeated small graph dispatches	RMSNorm, RoPE, SwiGLU kernels	Week 2 Day 4
Scalar/strided prefill projection loads	Cooperative 32×32×32 quantized matmul	Week 2 Day 5
Explicit secondary workload	Optional online-softmax decode lab or equivalent bounded experiment	Week 2 Day 6
Under-filled short-prefill result grid	Conditional measured split-K dispatch with Day 5 fallback	Week 2 Day 7
Functional whole-cache page updates	Aliasing page-slice write primitive	Week 3 Day 3
Scalar paged final reduction	Compact D=128 SIMD reduction	Week 3 Day 4
Scalar contiguous-page K/V tile loads	Cooperative paged FlashAttention loads	Week 3 Day 5

This is the course progression: optimize one measured cost, benchmark again, then let the evidence choose the next chapter.

Glossary Index

Scaled Dot Product Attention
Multi Head Attention
Linear
Rotary Positional Encoding
Grouped Query Attention
Qwen3 Attention Module
RMSNorm
SiLU
SwiGLU
MLP
Embedding
Qwen3 Transformer Block
Week 1 Qwen3 Model
dequantize_linear
KV Cache
Benchmarking and Profiling
Quantize the Model
Fused Model Kernels
Fused Decode Attention
SIMD-Matrix Prefill
Split-K Prefill
Flash Attention
Paged Attention

Keyboard shortcuts

Tiny-LLM — Build an LLM Serving System with MLX