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: Implement custom C++ and Metal kernels to accelerate the model.
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.

Choose a Model for Your Mac

The table below is a conservative starting point for common MacBook unified-memory sizes. 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
16 GB	0.6B / 1.7B	4B / 8B¹	4B / 8B	4B / 8B
32 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 a 16 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.

Week 2 Days 1–2 use the dense Week 1 loader. On a 16 GB Mac, keep using 0.6B until the packed quantized-matvec path is complete on Day 3; the 4B recommendation and 8B maximum apply after that checkpoint. ↩
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

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

On Day 1, we will implement basic attention and multi-head attention. 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

You may use MLX’s softmax; we will revisit lower-level operations in Week 2.

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.

At the end of this task, you should be able to pass the following tests:

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

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.

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)

At the end of the task, you should be able to pass the following tests:

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

You can run all tests for the day with:

pdm run test --week 1 --day 1

Week 1 Day 2: Positional Encodings and RoPE

On Day 2, we will implement the positional encoding used by Qwen3: rotary positional encoding (RoPE). A Transformer needs a way to represent each token’s position in the sequence. Qwen3 applies RoPE to the query and key vectors within its multi-head attention layer.

📚 Readings

Task 1: Implement Traditional Rotary Positional Encoding

You will need to modify the following file:

src/tiny_llm/positional_encoding.py

In traditional RoPE, as described in the readings, positional encoding is applied independently to each head of the query and key vectors. You can precompute the frequencies when initializing the RoPE class.

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.

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.

In practice, D can be even or odd. If it is odd, the final value has no partner and is typically left unchanged. For simplicity, this implementation requires D to be even.

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.

📚 Readings

You can test your implementation by running the following command:

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. 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, and concatenating the results.

📚 Readings

vLLM implementation of RoPE

You can test your implementation by running the following command:

pdm run test --week 1 --day 2 -- -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 2

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.

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. 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, and it applies RMSNorm to each Q and K head before RoPE. We will implement the reusable RMSNorm layer on Day 4, so call mx.fast.rms_norm directly for q_norm and k_norm today. Use non-traditional RoPE.

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

Tests that require an unavailable model will be skipped.

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

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.

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

Run the tests for this task with:

# This task's tests use the 0.6B model and tokenizer.
hf download Qwen/Qwen3-0.6B-MLX-4bit
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.

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:

# Download each model you want to test. Missing models are skipped.
hf download Qwen/Qwen3-0.6B-MLX-4bit
hf download Qwen/Qwen3-1.7B-MLX-4bit
hf download Qwen/Qwen3-4B-MLX-4bit
pdm run test --week 1 --day 5 -- -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 5

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. 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 log probabilities to sampler. 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.

Generate tokens in a loop until the model emits tokenizer.eos_token_id. Append each new token to the token array so that the next model call receives the complete sequence. Feed non-EOS output tokens to tokenizer.detokenizer, and print each new text segment as it becomes available.

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.

You can test your implementation by running the following command:

# Start 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"

# If downloaded, you can also try the larger models.
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"

Each command should produce a reasonable explanation of large language models. Replace --solution tiny_llm with --solution ref to run the reference solution.

Week 1 Day 7: Sampling and Preparing for Week 2

On Day 7, we will implement several sampling strategies and prepare the development environment for Week 2.

Task 1: Sampling

On Day 6, we implemented greedy decoding. In this task, we will add temperature, top-k, and top-p (nucleus) sampling.

src/tiny_llm/sampler.py

📚 mlx-lm sampler implementation

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.

pdm run main --solution tiny_llm --loader week1 --model qwen3-0.6b --sampler-temp 0.5

Top-k Sampling

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

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

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

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, then apply temperature sampling.

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

Task 2: Prepare for Week 2

In Week 2, we will optimize the Qwen3 serving infrastructure with C++ and Metal kernels. You will need Xcode and its command-line tools, including the Metal compiler, to build them.

Install Xcode:

Install Xcode from the Mac App Store or from the Apple Developer website (this may require an Apple Developer account).
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).
Install Xcode Command Line Tools:

Open your Terminal and run:
```
xcode-select --install
```
Set the Default Xcode Path (if needed):

Ensure that your command-line tools are pointing to your newly installed Xcode. You can do this by running:
```
sudo xcode-select --switch /Applications/Xcode.app/Contents/Developer
xcode-select --print-path
```
Adjust the path if Xcode is installed elsewhere.
Accept the Xcode License:

You may also need to accept the Xcode license:
```
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 a missing Metal toolchain, download it and verify the compiler again:
```
xcodebuild -downloadComponent MetalToolchain
xcrun metal --version
```
Install CMake:
```
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 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 continuing. 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. We have implemented all the components required to serve Qwen3. In Week 2, we will optimize the serving infrastructure for Apple silicon.

🚧 Week 2: A Step Closer to vLLM

Status: Experimental. Week 2 is under active development. Each chapter carries its own verification notes; the summary below records what is continuously tested versus what is one-machine research evidence.

Week 2 keeps the Week 1 Python mlx.core model intact and builds a separate optimized Qwen3 path for single-request decoding. It begins with the algorithm change: prefill once, retain a dense KV cache, and decode one new token at a time. Later chapters introduce kernels that address the costs the KV cache exposes.

⏱️ Time commitment. Days 3–7 write and tune custom Metal kernels. Completing the full Week 2 sequence typically takes substantially longer than Week 1 Days 6–7. All seven days are required core material; plan accordingly.

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 are stored as uint32. Numerically sensitive reductions, dot products, and online-softmax state accumulate in FP32 inside Python reference expressions or kernel registers. This contract remains in force for Week 3.

Verification Status

Reference correctness and decode-attention boundaries are continuously tested on ARM64 macOS CI with Qwen3-0.6B. Qwen3-4B performance evidence is one-machine research data measured on an M4 Pro, not a cross-device guarantee. Raw measurements, rejected experiments, and retained dispatch choices live in the performance evidence ledger.

Complete the Core Path

The full reference solution is a substantial performance-engineering project. All seven days are required core material; schedule more than one week if needed:

Required work	Provided infrastructure	Optional work
Days 1–7: cached model integration, matched benchmarking, quantization, fused model kernels, bounded decode-attention, SIMD-matrix prefill, and shape-aware Split-K	Model loading, extension build system, benchmark runners, correctness tests, and Python-reference implementations	The short profiling notice, schedule searches, hardware-specific retuning, and the 80%-of-MLX stretch target

The tests define API and correctness contracts; they do not require a student to rediscover the reference schedule. Metal capture, Xcode visualization, gpudebug, and profiling microbenchmarks are not current requirements and are not acceptance gates.

What We Will Cover

A dense per-request key-value cache for incremental decoding
Synchronized benchmarking and the dense decode roofline
Packed W4 quantization and a SIMD matrix-vector Metal kernel
Fused RMSNorm, RoPE, and SwiGLU Metal kernels
An online-softmax decode-attention kernel
A BF16 SIMD-matrix quantized prefill kernel
A shape-aware split-K schedule for small Qwen prefill matrices
A last-token output interface for generation
An optional stretch target of 80% of MLX prefill and decode throughput on the fixed Week 2 checkpoint

Week 2 does not call MLX-provided implementations of the operators we are learning. Your solution implements quantized matmul, decode attention, RMSNorm, RoPE, and SwiGLU in its own Python, C++, or Metal code. In particular, the completed checkpoint does not use mx.quantized_matmul, mx.dequantize, mx.fast operators, or mx.fast.scaled_dot_product_attention as shortcuts. The Day 1 baseline still uses Week 1’s Python mx.dequantize loading helper; Day 3 replaces that loading path as part of keeping weights packed.

Week 2 uses mlx_lm to load model weights and mlx.core for arrays, graph evaluation, and device synchronization.

Weekly Checkpoints

KV cache: port the Week 1 operators into a Week 2 model, add request-scoped state, and stop recomputing the prefix.
Benchmarking and profiling: measure the cached model against MLX with a matched, synchronized protocol. Profiling is optional and deferred until the macOS 27 tooling is available.
Quantize the model: keep W4 weights packed, implement the matrix-vector Metal path, wire it into the live model, and rerun the Day 2 benchmark.
Fused model kernels: fuse RMSNorm, RoPE, and SwiGLU one operator at a time after packed projections narrow the benchmark gap.
Decode attention: introduce online softmax over its tested short-context range and verify it with a matched workload.
SIMD-matrix prefill: return to the fixed 128-token workload and replace the correctness-first matrix path with cooperative tiles.
Split-K prefill: partition the reduction dimension only for under-filled short projections and fall back to Day 6 at the measured crossover.

Run the Supplied Test Gates

The seven learner days now map one-to-one to the existing supplied selectors:

Course day	Test command selector
Day 1	`--week 2 --day 1`
Day 2	`--week 2 --day 2`
Day 3	`--week 2 --day 3`
Day 4	`--week 2 --day 4`
Day 5	`--week 2 --day 5`
Day 6	`--week 2 --day 6`
Day 7	`--week 2 --day 7`

Run every group assigned to the chapter before continuing. The supplied test filenames and selectors are stable historical machine identities; the chapter headings and navigation now use the same seven-day sequence.

Week 2 to Week 3

The completed Week 2 model decodes one token at a time from a dense KV cache, dispatches separate prefill and decode matrix schedules, and keeps weights quantized throughout. Week 1 continues to use its Python mlx.core full-prefix generation loop.

Week 3 imports these Week 2 interfaces rather than copying or replacing them. It adds page-table translation and combines Week 2’s online softmax, SIMD-matrix tiling, and page walking in one paged FlashAttention operator. That boundary lets each week’s model remain understandable and runnable on its own.

Run pdm run bench-week2-progression to measure each checkpoint against the Week 1 baseline and MLX. Full methodology and cumulative results are in the performance appendix. The default runs reference checkpoints; add --solution tiny_llm to measure your implementation.

🚧 Week 2 Day 1: KV Cache

Status: Experimental. See the Week 2 verification matrix for what is continuously tested, locally measured, and still under review.

In this chapter, we will add a key-value cache to the Qwen3 model. During generation, the cache lets each attention layer 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, not a serving-only Week 3 feature. Without it, every generated token reruns all model layers over an ever-growing prefix, overwhelming the gains from faster individual kernels.

📚 Readings

KV Caching Explained: Optimizing Transformer Inference Efficiency

Recall how Week 1 repeatedly supplied the full sequence to the model:

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. A causal mask also prevents earlier queries from attending to the new token, so their outputs do not change. Recomputing those rows, their softmax values, and their products with V is wasted work. Only the new query row contributes a new output.

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 maintains its own key-value cache. The cache exposes one method, update_and_fetch, which:

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.

In this chapter, the cache passes mask through unchanged and does not use mask_length. Those parameters become important in Week 3 for batching.

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

This is deliberately a simple dense baseline, not a production KV cache. mx.concat allocates a larger buffer and copies the previous K/V contents on every growth step. Over a token-by-token decode of length S, those copies add up to O(S²) bytes even though caching avoids O(S²) prefix recomputation. The reference cache records this traffic as growth_copy_bytes so the profiler can keep it separate from attention. Week 3 replaces this baseline with preallocated pages; do not copy the repeated-concatenation design into a serving cache.

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. Start a 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 in this chapter: the Week 2 model accepts a cache and an offset while Week 1 keeps recomputing the full prefix. This produces the baseline that every later Week 2 chapter will optimize.

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. Do not introduce packed weights or fast kernels yet: measuring one algorithmic change makes the gain attributable. 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 every request gets one cache handle per Transformer layer. Pass the matching layer cache through every block and keep the caller’s offset consistent with the cache’s logical length.

To verify correctness, run the following test, which is similar to the Week 1 model test:

pdm run test --week 2 --day 1

Task 4: Connect the Serving Loop

src/tiny_llm/generate.py

The first model call prefills the cache with the complete prompt. Each later call passes only the token produced by the preceding step, together with the number of tokens already cached. The same lifecycle will be owned by the continuous-batching scheduler in Week 3.

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

Run the cached Week 1 checkpoint end to end before changing any operator:

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

Record this number in your optimization ledger. The next chapter teaches how to compare it fairly with Week 1 and MLX; every later command changes exactly one cumulative checkpoint.

Day 1 is an algorithmic checkpoint, so it does not invent a shader-level limiter from a GPU trace. The checkpoint removes full-prefix recomputation; use the end-to-end benchmark to measure that algorithmic change. Day 2 measures this model and identifies the projection-weight bandwidth bottleneck. Day 3 introduces 4-bit quantization and implements the SIMD matvec kernel that operates on packed weights directly.

🚧 Week 2 Day 2: Benchmarking and Profiling

Status: Experimental. See the Week 2 verification matrix for what is continuously tested, locally measured, and still under review.

Day 1 gave us a cached model. Day 2 establishes a trustworthy dense BF16 baseline: how fast are prefill and decode under one matched protocol, and what architectural cost should the next chapter attack? Benchmarking is required. Profiling is optional and is not a prerequisite or acceptance gate.

Benchmark the Cached Model

Optimization starts with a trustworthy comparison. Prefill processes many prompt tokens at once; decode usually processes one token per request and is dominated by repeatedly reading dense BF16 projection weights at this checkpoint. A change can improve one phase while hurting the other, so benches/bench.py reports both:

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 belongs to prefill. Excluding it from decode prevents prompt length from distorting the decode number.

Choose the prefill workload before comparing implementations. Prompt scoring needs logits for every position, while serving needs only the final prompt logit. Use --prefill-logits all for the former and --prefill-logits last for the latter. The runner applies the choice to your solution and MLX alike. Never compare a final-row run from your solution with an all-row MLX run.

Both sides of the Week 2 comparison use a KV cache: prefill the prompt once, then pass only the newly generated token on each decode step. Comparing a cached MLX baseline with your solution recomputing the full prefix would measure two different algorithms and make the next optimization target meaningless.

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 4 \
  --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 let the machine return to a stable temperature before comparing runs. Run each command several times, report the median, and include the hardware, MLX and mlx-lm versions, prefill-logit mode, and exact model with the result. A dependency upgrade changes the comparison baseline, so remeasure MLX rather than carrying an old denominator forward.

Synchronize Lazy Work

MLX builds lazy computation graphs. Timing only the Python call measures graph construction, not GPU execution. Every timed iteration must evaluate the 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 so cache implementations with owned or shared resources can return them:

pdm run test --week 2 --day 2

Optional Profiling Boundary

The required Day 2 work ends with the synchronized benchmark JSON. Metal capture, Xcode visualization, gpudebug, and related profiling microbenchmarks are not part of the current course requirements. They require the macOS 27 tooling release and will return as optional material after that release is available.

The optional profiling notice records this boundary. You may skip it and continue directly to Day 3. No profiling tool, trace, screenshot, or microbenchmark is a prerequisite or acceptance gate.

Why Quantize: The Decode Roofline

The decode phase of LLM inference is typically memory-bandwidth bound: each token requires reading the model’s weights but performs 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

The tied embedding matrix is counted once as the vocabulary projection. The single-row embedding lookup, normalization weights, activations, KV reads, and attention work are omitted. This makes the result an upper bound for linear layers, not 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: Metal Profiling

Week 2 requires reproducible benchmarks, not a GPU capture. The course’s full Metal profiling workflow depends on capture and visualization tooling that will ship with macOS 27, so those instructions are intentionally deferred until that release is available.

These materials are not required: gpudebug, an Xcode GPU capture, a .gputrace, screenshots, or profiling microbenchmarks to complete any current checkpoint. Continue from Day 2’s matched benchmark directly to Day 3’s quantized live-model integration.

When the macOS 27 workflow is added back, it will remain optional and will supplement—not replace—the ordinary synchronized benchmark and correctness gates.

🚧 Week 2 Day 3: Quantize the Model

Status: Experimental. See the Week 2 verification matrix for what is continuously tested, locally measured, and still under review.

Day 2 established a synchronized dense BF16 baseline. Day 3 reduces projection weight traffic with packed W4A16 weights, implements the Metal operators that consume them directly, wires those operators into the live model, and reruns the same benchmark. Packed storage or an isolated fast kernel is not completion: the cached model must use the quantized path.

📚 Readings

Debug Metal Without a CPU Twin

A C++ CPU version is possible but not required. Use this three-level validation ladder instead:

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. Do not debug an optimized kernel by comparing only full-model text output.

Make Failures Small and Synchronous

Start with 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()

Check the wrapper boundary before inspecting 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, temporarily simplify the schedule. 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 often more useful than printing from every GPU thread. Restore one optimization at a time and rerun both the aligned and tail-shape tests after each change.

Represent Weights With Fewer Bits

Quantization represents floating-point weights with values from a small integer codebook plus the parameters needed to approximately reconstruct the original values. 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.

With only 16 possible codes, the reconstructed weights approximate the original values. The smaller representation trades some numerical precision for less 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

Instead of applying one scale to an entire weight matrix, we divide each row into groups and quantize every group independently. Local scales and biases preserve more information about each 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 the packed codes and their 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 that orientation: code 0 is the upper endpoint and code 15 moves toward the lower endpoint. Both orientations occur in the shipped Qwen3-4B MLX checkpoint, so do not recompute scale and bias 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

The smaller representation reduces the projection weight traffic by 3.765×. That ratio is a bandwidth ceiling for one-token decode, not a promise of the same end-to-end speedup.

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. These are theoretical ceilings, not benchmark results.

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. Actual throughput is lower because the complete model also reads activations and KV, launches other operators, and does not sustain peak bandwidth continuously. 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, increasing arithmetic intensity. It therefore needs a matrix schedule; the decode bandwidth ratio should not be treated as a prefill prediction.

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

Modify these exact starter functions:

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

The starter code provides QuantizedWeights, a container for 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]$ )

Its from_mlx_layer method extracts these fields from an MLX quantized layer when loading the model.

Next, implement quantized_linear, a wrapper around quantized_matmul with the same input convention as the standard linear function. You will implement quantized_matmul in the next task.

Keep the token embedding table quantized as well. 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 starter already contains the declaration, fail-closed source stub, binding, and build registration. 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 is infrastructure: it lets an mx.array graph node schedule the Metal loop you write in the next task. MLX owns the array lifetime and command encoder, but it does not supply the quantized multiplication.

Build the extension to catch declaration, binding, and registration mismatches. The focused test below 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 your first Metal kernel, understand the execution model. Metal organizes GPU work in four nested 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. More SIMD groups within one threadgroup add threads and can raise register demand; they increase threadgroup-memory use only when the schedule allocates shared storage per group or tile. Either resource can reduce the number of resident threadgroups. More threadgroups in the grid change how the output or reduction work is partitioned. Neither change guarantees higher throughput.

Use the required two-SIMD-group matvec schedule as the Qwen starting point, then benchmark two, four, eight, and sixteen groups per threadgroup as described below. Change the grid partition separately so each measurement answers which launch knob helped.

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 the Metal kernels and connect eval_gpu to them. The Python quantized_matmul wrapper always dispatches the primitive you implement on GPU; the required path in your solution never routes through mx.quantized_matmul.

Do this in two measured stages. They expose the same math but schedule different shapes differently:

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 6 replaces it with a cooperative tiled kernel.

The cutoff does not mean the SIMD kernel expands to cover larger M. The two paths are separate schedules: Day 3 optimizes the vector-shaped decode bottleneck and leaves matrix-shaped prefill visible for the later benchmark to select.

Keep the vanilla function callable as quantized_matmul_vanilla. An optimization is much easier to trust when it can be compared directly with the implementation it replaces.

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 6 revisits that workload with cooperative tiling.

Stage 2: SIMD Matvec

Decode normally has M = 1; an 8×8 matrix tile would leave most rows empty. Instead, one SIMD group reduces the input dimension and uses simd_sum to combine lane-local partial sums. Start with two output columns per group as an inspectable schedule. For the Qwen3-4B checkpoint, then 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. Use this Qwen-focused starting point:

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, not mathematical requirements. Keep them visible in the dispatcher, then 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. Its lower instruction count is useful 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 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.

Use direct activation reads for your 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, report the whole-model result and keep it only when reuse outweighs synchronization.

Kernel Requirements

Implement both required kernel 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 6 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 by following 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.

You can test your solution by running:

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__ in this task. These are the exact points that load quantized weights, replace dense projections, and keep only the requested logits row.

Integrate quantized matrix multiplication into the Week 2 Qwen3 model so that the 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, not 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.

You can test your solution by running:

pdm run test --week 2 --day 3

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

You can also benchmark your 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.

The vanilla matrix product remains callable as an inspectable Metal control, but the Python mlx.core equation is the correctness oracle and only the SIMD matvec is integrated into decode.

Verify Quantization in the Complete Model

Before moving on, confirm that the quantized matvec kernel is actually called during model inference, not just registered and tested in isolation.

🚧 Acceptance criterion. Your checkpoint is incomplete until 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. Use a source trace through those branches in your completed dispatcher and model wiring. The supplied tests validate packed model state and the direct operators, while the matched benchmark reports complete-model throughput; neither proves the live Metal pipeline identity by itself. Use a direct source trace of the dispatch branches, then treat the throughput comparison as a separate result.

Measure the cumulative model and the real projection shapes:

pdm run bench-week2-progression --offline --solution tiny_llm --repeats 4 \
  --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 bench-week2-operators --solution tiny_llm --model qwen3-4b \
  --section decode-projections --context 128

Attach the complete-model before/after rows, the per-projection latency table, and the direct dispatch trace. First require a clear decode gain over kv-cache. Then compare each projection with MLX at the identical shape. Projections may remain the largest absolute category because the model performs them in every layer; once their operator latency is close to MLX, that bar is no longer the largest removable gap.

Continue to Day 4 only after the correctness tests pass, the source trace proves that the live model selects the intended matrix and matvec branches, the matched model run improves decode over kv-cache, and the projection table is close to MLX at the same shapes. If the projection comparison is still far behind, keep tuning the matvec instead. Once that gap shrinks, Day 4 turns to the recurring normalization, position, and activation work around those projections.

Optional profiling evidence. A kernel-group replay or operator attribution can corroborate that transition, but neither gates progress. The reference checkpoint includes both alongside the model and projection measurements above.

🚧 Week 2 Day 4: Fused Model Kernels

Status: Experimental. See the Week 2 verification matrix for what is continuously tested, locally measured, and still under review.

Day 3 removed the largest projection gap. Day 4 now targets RMSNorm, RoPE, and SwiGLU, which recur around those projections in every transformer layer. Week 1 expresses them as Python mlx.core equations. Week 2 keeps those implementations intact and asks you to write three Metal kernels behind a separate interface:

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

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.

Optional profiling evidence. The Day 3 kernel-group replay and the reference-solution attribution show the pointwise cluster behind the optimized projections. They explain the chapter order but are not prerequisites or acceptance gates.

Why Fusion Helps

Week 1’s Python mlx.core equations already run as native GPU kernels inside the lazy graph. The important difference is how many operations and memory passes the graph 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.

The useful comparison is not “Metal versus Python arithmetic,” but one purpose-built kernel versus a graph of several general-purpose kernels.

Task 1: RMSNorm

Modify 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 header, binding, C++/Metal files, and CMake registration already exist for this checkpoint; replace the fail-closed bodies instead of adding parallel APIs.

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.

Integrate FastRMSNorm into every Week 2 norm immediately, run the RMSNorm tests, and record the cumulative model result before writing 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

Modify 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 already optimized model, then test and measure that cumulative checkpoint before implementing 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

Modify 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.

Integrate the fused expression immediately and 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

Verify the cumulative switches in Qwen3ModelWeek2.__init__ and the call sites in Qwen3MultiHeadAttention.__call__ and Qwen3MLP.__call__. Task 4 should not introduce another extension function; it composes the three functions from Tasks 1-3.

After exposing all three kernels through C++ MLX primitives, run the complete test file to verify their composition. Keep qwen3_week1.py on its Week 1 Python operators, and make the Week 2 interfaces reusable by the Week 3 serving model.

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

Compare against the Python reference equations with tolerances rather than bit-for-bit equality. Test RoPE with scalar and per-batch offsets. Always call mx.eval inside a timed iteration when measuring these lazy operations.

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

Keep the three cumulative checkpoints separate so a regression cannot hide inside their combined gain:

pdm run bench-week2-progression --offline --solution tiny_llm --repeats 4 \
  --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 bench-week2-operators --solution tiny_llm --model qwen3-4b \
  --section model-kernels --context 128

Attach each cumulative model row, the three Python-reference/optimized/MLX operator rows, and the direct dispatch trace. Let the matched benchmark results decide whether the kernels stay.

Continue to Day 5 when all three correctness gates pass, the direct fused-dispatch source trace reaches the intended kernels, the cumulative rows retain the gain, and the three operator comparisons justify keeping the fused implementations. Day 5 then tests whether attention is the next removable gap by sweeping cached context and query length before setting a dispatch guard.

Optional profiling evidence. The reference checkpoint pairs the cumulative and operator measurements with an updated attribution. That attribution can explain the transition, but it does not replace the checkpoint evidence above.

🚧 Week 2 Day 5: Fused Decode Attention

Status: Experimental. See the Week 2 verification matrix for what is continuously tested, locally measured, and still under review.

This chapter starts only after the Day 4 evidence has verified that the fused model kernels reduced the repeated pointwise cluster. Linear projections remain important, but their operator latency is already close to the external denominator, while attention is the next measured removable gap through cached context S <= 256. Longer-context measurements use the Week 1 Python fallback for caches beyond 256; every checked context through 256 uses the optimized path. During single-request decode, query length is normally one while the cached key/value sequence grows by one token at a time. Week 1 expresses attention as matrix multiplication, masking, softmax, and another matrix multiplication. That is expressed with mlx.core, but it materializes the complete score and probability rows.

First write a Python mlx.core composition to preserve the equation, then replace its matmuls and softmax with an online-softmax Metal kernel in your solution. Measure the complete model before deciding whether to retain the dispatch. The kernel does not call mx.matmul or an MLX-provided scaled-dot-product-attention implementation; MLX still provides arrays, streams, buffers, and extension dispatch.

Task 1: Preserve the Interface

Modify scaled_dot_product_attention in src/tiny_llm/week2_kernels.py. Keep this readable function as the oracle and fallback; Task 2 modifies the separate decode_attention_custom entry point.

Implement scaled_dot_product_attention in week2_kernels.py with these model-facing shapes:

query: B, H_q,  L, D
key:   B, H_kv, S, D
value: B, H_kv, S, D
out:   B, H_q,  L, D

Validate that H_q is divisible by H_kv. Flatten batch and head dimensions for the extension and map each query head to its shared KV head with:

kv_head = query_head / (H_q / H_kv)

Normalize explicit masks to B * H_q, L, S. Also pass a causal flag so the kernel can skip future positions without constructing a causal-mask tensor.

As a Python intermediate step, reshape query heads into H_kv groups and a repeat dimension. Broadcasting then pairs several query heads with one KV head without physically repeating the key and value tensors. Express scaled scores, softmax, and the weighted-value product explicitly. Use this form as a correctness oracle and ablation, not as the completed optimized path: its matmuls are MLX-provided operator implementations.

Task 2: Implement Online Softmax in Metal

Modify tiny_llm_ext::decode_attention, Week2DecodeAttention::eval_cpu, and Week2DecodeAttention::eval_gpu in src/extensions/src/week2_kernels.cpp, the week2_decode_attention function in src/extensions/src/week2_kernels.metal, and decode_attention_custom in src/tiny_llm/week2_kernels.py. The starter declaration, binding, source stub, Metal file, and CMake registration are already present and labeled Week 2 Day 5; replace those fail-closed bodies rather than adding new names.

Expose decode_attention_custom for the Metal implementation. Cache the scaled query fragment in registers before walking the cache; loading it again for every key position is avoidable. Assign 32 32-lane SIMD groups to each query row on the 128-192 token benchmark. Each group visits every 32nd cached position; within a group:

Each lane multiplies a regularly spaced subset of query and key values.
simd_sum combines those partial dot products into one score.
Apply the scale, optional mask, and causal check.
Update a running maximum, softmax denominator, and weighted value accumulator.

The online update is:

new_max = max(running_max, score)
old_factor = exp(running_max - new_max)
score_factor = exp(score - new_max)
denominator = denominator * old_factor + score_factor
accumulator = accumulator * old_factor + score_factor * value

After its last cached position, each group writes its partial maximum, denominator, and value accumulator to threadgroup memory. The first SIMD group computes the common maximum and rescale factors. One thread computes the final denominator, then the first D threads each combine one output dimension. This keeps the final value reduction parallel across the head dimension. Subtracting the maxima gives stable softmax without storing all S scores or probabilities.

This removes two large intermediates and several dispatch boundaries from the Week 1 graph. It is especially relevant as context grows: the avoided score and probability tensors are proportional to L * S, while decode needs only the final D-element result for each query head.

Load and store BF16 directly, but accumulate dot products, softmax state, and weighted values in float32. Casting whole Q, K, and V tensors outside the kernel creates extra dispatches and memory traffic; doing the conversion in registers avoids that cost.

Use fast::exp for the rescale factors and compute each factor once before applying it to the denominator and all value dimensions. These ideas also appear in production vector-attention kernels, including MLX’s SDPA sources. Your kernel reimplements the algorithm and scheduling in its own Metal code; it does not include or instantiate the MLX kernel.

Scheduling Experiment

Compare eight, sixteen, and thirty-two SIMD groups with Qwen3-4B while holding the context fixed. The number of groups is a workload parameter, not a universal constant: more groups expose parallel score work but consume more threads and threadgroup memory. Record the synchronized operator and complete-model result for each schedule, then repeat the experiment when context length changes.

Task 3: Integrate and Measure

Modify Qwen3MultiHeadAttention.__call__ in src/tiny_llm/qwen3_week2.py to apply the measured dispatch guard. Keep scaled_dot_product_attention as the explicit fallback and call decode_attention_custom only inside the supported region.

Route short-query, short-context Week 2 attention through the Metal implementation. Dispatch back to the Python mlx.core composition when the cached context exceeds the measured crossover; a schedule that wins at 128 tokens should not be forced onto 2,048 tokens. Retain the Python composition for tests and ablations. Week 3 later combines this recurrence with paged K/V and SIMD-matrix tiles for FlashAttention; prefill is a different workload where both query and context lengths are large.

Set a concrete dispatch guard: use your Metal kernel only when query length is at most two and cached context length is at most 256. Otherwise use the Python grouped-attention path. Keep this condition at the model call site so the benchmarked operating range remains reviewable instead of becoming a hidden performance policy inside the Metal kernel.

Keep arbitrary dense, per-request masks on the Python model path. The primitive still accepts explicit masks so its arithmetic contract can be tested, but the Week 2 dispatch guard selects the custom kernel only for None or "causal". Explicit masks appear in the first continuous-batching exercise, while normal single-request decode uses no mask. Week 3 replaces dense batch masks with paged-attention metadata instead of making them a hidden performance policy in this focused model path.

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

Test grouped-query head mapping, output shape, causal behavior, and explicit masks against the Python Week 1 implementation. The reference suite uses Qwen’s D = 128, query lengths 1 and 8, GQA ratios 1 and 4, and cached contexts 1, 31, 32, 127, 128, 129, 255, 256. It also checks both sides of the model’s L <= 2 and S <= 256 dispatch guard. Use a tolerance because online softmax changes the floating-point reduction order.

Correctness over that grid does not prove that a fixed 32-SIMD-group schedule is efficient. At contexts 1, 8, and 31, many of its 1,024 threads have no score position to process. Run the same real-shape operator sweep on each target machine before retaining the schedule:

for context in 1 31 32 127 128 129 255 256; do
  pdm run bench-week2-operators --solution tiny_llm --model qwen3-4b \
    --section attention --context "${context}" \
    --query-length 1 --gqa-ratio 4 --attention-mask none
done

for context in 8 31 32 127 128 129 255 256; do
  pdm run bench-week2-operators --solution tiny_llm --model qwen3-4b \
    --section attention --context "${context}" \
    --query-length 8 --gqa-ratio 4 --attention-mask causal
done

Repeat representative points with --gqa-ratio 1 and --attention-mask explicit. Keep M1 and M4 results as separate records; a correctness run on the M1 CI runner is not evidence that the M4 crossover applies there.

Run the preceding checkpoint and your solution with the new dispatch under otherwise identical settings:

pdm run bench --solution tiny_llm --loader week2 \
  --week2-checkpoint swiglu --model qwen3-4b \
  --num-seqs 1 --min-input-len 32 --max-input-len 32 \
  --min-output-len 97 --max-output-len 97 --warmup 2 \
  --prefill-logits last

pdm run bench --solution tiny_llm --loader week2 \
  --week2-checkpoint decode-attention --model qwen3-4b \
  --num-seqs 1 --min-input-len 32 --max-input-len 32 \
  --min-output-len 97 --max-output-len 97 --warmup 2 \
  --prefill-logits last

Prefill produces the first token, so the 96 timed decode calls grow the cache from S=33 through S=128. Every one is inside the custom dispatch guard. Your solution falls back to the exact Python Week 1 composition outside that validated range.

Benchmark Analysis: Verify Prefill Projections Are the Next Bottleneck

Measure the attention operator and the cumulative checkpoint separately. The first progression is the matched short-context acceptance test for this bounded kernel. The second keeps the fixed Week 2 denominator: its 128-token prefill remains outside the query-length guard, while timed one-token decode steps with S=129 through S=256 enter the current context guard:

pdm run bench-week2-operators --solution tiny_llm --model qwen3-4b \
  --section attention --context 32 --context 128 --context 160 \
  --context 192 --context 256 --context-repeats 6 \
  --json-output benchmark_results/week2-attention-context-sweep.json

pdm run bench-week2-progression --offline --solution tiny_llm --repeats 4 \
  --variant week2-swiglu --variant week2-decode-attention --variant mlx \
  --model qwen3-4b --input-len 32 --output-len 97 --warmup 2 \
  --prefill-logits last

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

Repeat the attention microbenchmark at contexts 32, 128, 160, 192, and 256, and attach that context sweep beside the short-context swiglu/decode-attention model rows. The intermediate points reveal whether the custom kernel has a useful measured crossover rather than assuming that an endpoint applies to every context. Reject the custom dispatch if repeated fresh-process short-context runs do not improve, even when the isolated kernel looks faster. If the operator wins only over a limited context range, encode that measured crossover in the dispatch guard.

Optional profiling evidence. Decode and prefill kernel-group results can explain how the workload divides its time, but they are reference evidence, not required output for this checkpoint.

The checked Qwen3-4B sweep on an M4 Pro used six forward/reverse context passes, rotated every implementation order, and recorded all 60 samples per implementation and pass:

Context	Python reference	Metal	MLX	Metal speedup
32	143.0 us	125.7 us	116.3 us	1.138x
128	149.3 us	136.3 us	120.6 us	1.095x
160	151.2 us	140.1 us	120.9 us	1.079x
192	154.0 us	143.9 us	121.9 us	1.071x
256	158.0 us	150.7 us	122.8 us	1.048x

The Metal path wins at every measured point through 256, so 256 is the largest evidenced context guard. The Python mlx.core path remains the policy beyond that range; do not extrapolate the final 4.8% operator win to longer caches. The raw record, including exact source SHA, model configuration, MLX and mlx-lm versions, Metal compiler version, device information, execution order, samples, and medians, is benchmark_results/m4-pro-qwen3-4b-week2-attention-context-sweep-mlx-0.32.0.json.

The production-boundary sweep held context at 128, selected Qwen3-4B’s 4:1 GQA ratio, and balanced L1/L2/L4/L8 order over six passes. It used the causal form for every query length: at L=1 that mask permits the entire existing cache and is equivalent to unmasked one-token decode, while L>1 measures causal multi-token chunks. Each pass also rotated the three implementation orders and retained every sample:

Query length	Python reference	Metal	MLX	Metal speedup	Pass wins
1	244.4 us	213.1 us	155.9 us	1.147x	6/6
2	341.4 us	258.8 us	185.3 us	1.319x	6/6
4	322.7 us	297.3 us	197.4 us	1.085x	4/6
8	377.7 us	491.5 us	290.6 us	0.768x	0/6

L4’s aggregate median improved, but it lost two of six balanced passes. L2 is the largest repeat-consistent win, so the dispatch guard remains conservative at L <= 2; L4 and L8 use the Python path. Reproduce the recorded sweep with:

pdm run bench-week2-operators --solution tiny_llm_ref --model qwen3-4b \
  --section attention --context 128 \
  --query-length 1 --query-length 2 --query-length 4 --query-length 8 \
  --gqa-ratio 4 --attention-mask causal --context-repeats 6 \
  --warmup 12 --iterations 60 \
  --json-output benchmark_results/week2-attention-query-sweep.json

The checked raw record is benchmark_results/m4-pro-qwen3-4b-week2-attention-query-sweep-mlx-0.32.0.json.

In the fixed 128/129 workload, prefill has L=128 and uses the Python path. The first timed decode call appends the new token before the guard sees S=129; the one-token decode calls through S=256 therefore use the custom path. Keep the fixed workload separate from the short-context acceptance run. Continue to Day 6 after the correctness tests pass, the direct source trace proves the bounded attention dispatch and its fallback, repeated short-context runs retain the gain, and the fixed 128/129 control confirms that prefill is unchanged and still routes through Day 3’s matrix-shaped projection path.

Optional profiling evidence. The reference checkpoint pairs the context sweep, short-context model delta, and fixed-workload control with a separate prefill attribution. The attribution explains why the course targets matrix-shaped projections next; it is not a prerequisite for Day 6.

🚧 Week 2 Day 6: SIMD-Matrix Prefill

Status: Experimental. See the Week 2 verification matrix for what is continuously tested, locally measured, and still under review.

Day 5 ends by switching the benchmark from one-token decode to multi-token prefill. Its source trace shows that the 128-token prefill still uses Day 3’s correctness-first vanilla quantized matrix path for M > 8. Day 6 replaces that inherited multi-row schedule, then measures the complete model and the real projection shapes to decide whether the new path stays.

MLX remains an external performance denominator; the SIMD-matrix path in your solution continues to call the C++/Metal primitive you implement for every projection.

Optional profiling evidence. The checked dependency-aware attribution and the reference-solution attribution explain why projections are the reference solution’s next target. They are not required learner output and do not gate this chapter.

The implementation remains deliberately narrow:

W4A16 weights with four bits and group size 128;
BF16 activations, quantization parameters, and output;
Qwen3-4B projection dimensions;
FP32 matrix accumulators;
the Day 3 SIMD matvec remains in use for M <= 8.

From a Matvec to a Cooperative Tile

The vanilla one-thread dot product and a single-group 8×8 tile are useful Metal bring-up controls, but neither provides enough cooperative reuse for multi-row prefill. Compare both with the Python MLX correctness oracle. The performance schedule must share both activations and dequantized weights across a larger result tile.

The optimized kernel assigns four SIMD groups, or 128 threads, to one 32×32×32 tile:

                  32 output columns
               +--------------------+
32 prompt rows |  four 16x16 SIMD   |
               |  output quadrants  |
               +--------------------+
                         ^
                         |
             shared 32-value K step

For each 32-value reduction step, the threadgroup:

loads one 32×32 activation tile into padded threadgroup memory;
unpacks and dequantizes one 32×32 weight tile there;
lets four SIMD groups reuse both tiles;
accumulates four 16×16 quadrants from Metal 8×8 matrix fragments;
advances to the next reduction tile.

The 40-element shared-memory stride pads the 32-value rows to avoid an unhelpful bank-access pattern. Tail rows and columns are zero-filled or guarded at the final store.

Your Metal kernel may use MLX’s low-level Steel BlockLoader and BlockMMA headers as building blocks. Those helpers provide cooperative loads and matrix-fragment bookkeeping. Your solution still owns the W4A16 unpacking, dequantization, tile layout, primitive, dispatch, split policy, and reduction; it does not call MLX’s quantized-matmul operator.

Task 1: Preserve the Workload Dispatch

Modify QuantizedMatmul::eval_gpu in src/extensions/src/quantized_matmul.cpp and quantized_matmul_simdgroup_w4a16_g128 in src/extensions/src/quantized_matmul.metal. Keep the Day 3 quantized_matvec_x4_fast_w4a16_g128 function intact for M <= 8.

Keep the Day 3 decode schedule and add the matrix schedule behind the same quantized-linear interface:

M <= 8  -> quantized SIMD matvec
M > 8   -> 32x32x32 quantized SIMD-matrix kernel

Expose the new path through the cumulative simd-matmul checkpoint. Test the vanilla, tiled, and MLX results on an aligned shape and on partial row and column tiles. The result must retain the model-facing 16-bit dtype.

Task 2: Make Device Loads Contiguous

Continue modifying quantized_matmul_simdgroup_w4a16_g128 (and its private Metal helper, if you factor one) in src/extensions/src/quantized_matmul.metal; do not change the public quantized_matmul binding.

Use a cooperative block loader so adjacent threads and each thread’s local reads form contiguous transactions. This is a requirement of the schedule, not a cosmetic detail. Benchmark Q, K/V, gate/up, and down projections separately at their Qwen3-4B dimensions so both wide and narrow output grids are covered.

Task 3: Hoist Quantization Parameters

Continue modifying quantized_matmul_simdgroup_w4a16_g128 in src/extensions/src/quantized_matmul.metal. This task changes the tiled kernel’s load/reuse strategy, not its C++ or Python signature.

One scale and bias apply to 128 reduction values. Loading them for every 32-value tile repeats the same device access four times. Have one thread load the scale and bias for each of the 32 output columns into threadgroup memory, then let the four weight-unpack threads for that column reuse them for the next four reduction tiles.

Keep the scale, bias, and unpacked operands in BF16 storage, while the matrix accumulator remains FP32. Cast once when writing the final model output.

Task 4: Project Only Required Logits

Modify Qwen3ModelWeek2.__call__ in src/tiny_llm/qwen3_week2.py so logits_to_keep=1 slices before the vocabulary projection. Do not add a new extension function for this model-level optimization.

Generation needs only the final prompt row to produce the first sampled token. Accept logits_to_keep=1 and apply the vocabulary projection only to that row. The benchmark applies the same last-logit workload to MLX, while prompt-scoring callers can still request every logit row.

Task 5: Verify, Benchmark, and Name the Next Bottleneck

Task 5 adds no function. Verify the cumulative QuantizedMatmul::eval_gpu/quantized_matmul_simdgroup_w4a16_g128 path and the Qwen3ModelWeek2.__call__ projection boundary from Tasks 1-4.

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

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

pdm run bench-week2-progression --offline --solution tiny_llm --repeats 4 \
  --variant week2-decode-attention --variant week2-simd-matmul --variant mlx \
  --model qwen3-4b --input-len 32 --output-len 33 --warmup 2 \
  --prefill-logits last

Inspect the projection sweep as well as complete-model throughput. Continue to Day 7 when the long-M projections are healthy but short, narrow K/V projections launch too few 32×32 result tiles to fill the GPU. If the same kernel remains slow at large M, improve its loads or matrix schedule before adding reduction partitions.

At long M, the two-dimensional tile grid is already large. Do not force the next optimization there: additional reduction partitions would only add a temporary buffer and another launch.

Benchmark Analysis: Identify Under-Filled Prefill Shapes

Compare the matrix kernel at both an occupied control shape and the short K/V shape, then benchmark the latter without enabling Split-K:

for context in 32 128 2048; do
  for projection in q k v o gate up down; do
    pdm run bench-week2-operators --solution tiny_llm --model qwen3-4b \
      --section prefill-projections --context "${context}" \
      --prefill-projection "${projection}"
  done
done

The dispatch formula gives the unsplit 32-row K projection 32 independent threadgroups.

Attach the complete-model prefill delta and per-projection tables at 32, 128, and 2,048 rows. Do not select Split-K merely because projections still occupy most of prefill. First require the long or wide controls to approach MLX while the short, narrow projection remains disproportionately slow.

Use the dispatch calculation and short-shape operator sweep to establish that the unsplit result grid has too few independent threadgroups. Use the matched long-shape control to rule out costly work inside each tile; if it exposes such a cost, repair Day 6 before multiplying the grid. The reference checkpoint pairs the prefill gain with long and short operator controls and the dispatch geometry that motivates Split-K. A remaining arithmetic hot spot would send you back to Day 6 instead.

Optional profiling evidence. A 32/128-row attribution can corroborate the shape analysis, but it does not replace the matched complete-model delta, projection controls, and dispatch calculation above.

🚧 Week 2 Day 7: Split-K Prefill

Status: Experimental. See the Week 2 verification matrix for what is continuously tested, locally measured, and still under review.

Day 6’s cooperative loads brought long-row prefill near MLX. Its follow-up sweep shows a different problem at short prefill: Qwen’s narrow K/V projections do not launch enough independent result tiles to occupy the GPU. Today we split the reduction dimension only until that grid is large enough.

This chapter is not a general split-K library. It optimizes the model shapes we actually run:

Model	Reduction `N`	Q output `K`	K/V output `K`
Qwen3-4B	2,560	4,096	1,024

Why Split the Reduction Dimension?

For C = A @ W.T, Day 6 launches:

ceil(M / 32) * ceil(K / 32) threadgroups

Split-K adds a partition grid dimension:

partial[p, :, :] = A[:, N_start[p]:N_end[p]]
                    @ W[:, N_start[p]:N_end[p]].T
C = reduce(partial, partition axis)

This exposes more independent work, but rereads part of A, allocates a temporary tensor, and launches a reduction kernel. It is useful only while the original two-dimensional grid is under-filled.

Task 1: Reproduce the Under-Filled Grid

Task 1 changes no function. Benchmark the existing quantized_matmul_simdgroup_w4a16_g128 Day 6 kernel before editing the Split-K stubs.

Begin with the narrow K projection at M=32 before changing dispatch. This is the smallest baseline needed to reproduce the under-filled Day 6 grid; Task 4 runs the full all-projection, all-row sweep after Split-K exists:

pdm run bench-week2-operators --solution tiny_llm --model qwen3-4b \
  --section prefill-projections --context 32 --prefill-projection k \
  --warmup 5 --iterations 30

Record synchronized Day 6 and MLX latency before implementing Split-K. The narrow K/V shape is the clearest small-grid case. Large output widths or prompt lengths may already have enough row-by-column tiles and should become controls.

Task 2: Reuse the Day 6 Kernel for Each Partition

Implement quantized_matmul_simdgroup_splitk_w4a16_g128 in src/extensions/src/quantized_matmul.metal, reusing the Day 6 tiled helper behind quantized_matmul_simdgroup_w4a16_g128.

Add group_id.z as the partition index. Every partition must:

have the same reduction length;
start and end on a 128-value quantization-group boundary;
reuse the validated Day 6 loader, dequantizer, and 32×32 tile;
write to its own [M, K] plane without atomics.

Store partial planes in BF16 to keep the temporary small and perform the final sum in FP32 before the output cast. This introduces one extra BF16 rounding boundary compared with the unsplit FP32 accumulator, so tests use a BF16-appropriate tolerance. An FP32 temporary is a useful bring-up oracle, but it doubles the partial-buffer traffic.

Task 3: Choose Partitions From Occupancy

Modify QuantizedMatmul::eval_gpu in src/extensions/src/quantized_matmul.cpp to select the partition count and dispatch the Split-K kernel. Keep tiny_llm_ext::quantized_matmul and its Python binding unchanged; the existing use_split_k argument carries this cumulative checkpoint.

Use a small explicit policy:

base_groups = ceil(M / 32) * ceil(K / 32)
split_k = min(16, floor(320 / base_groups), N / 128)
decrease split_k until N is divisible by split_k * 128
use Day 6 unchanged when split_k <= 1

For the Qwen3-4B target, use roughly 320 threadgroups and a cap of 16 as explicit tuning parameters. They are not universal GPU properties. Unlike a hard-coded prompt-length cutoff, the grid calculation naturally stops splitting a narrow projection once more row tiles are present, and stops immediately for already wide grids.

For Qwen3-4B, the policy selects these schedules:

Projection	Base groups at `M=32`	Selected split at `M=32`	Selected split at `M=128`
Q, `2560 -> 4096`	128	2	1
K/V, `2560 -> 1024`	32	10	2
O, `4096 -> 2560`	80	4	1
MLP gate/up, `2560 -> 9728`	304	1	1
MLP down, `9728 -> 2560`	80	4	1

A split of one means the dispatcher uses the Day 6 kernel unchanged. At the 128-token acceptance shape only the narrow K/V projections remain eligible, with a two-way split; the other major projections already expose enough output tiles. At 2,048 tokens every projection uses the unsplit kernel.

Expose the policy through a cumulative split-k checkpoint. Keep Day 6 selectable so the benchmark always has an unsplit control.

Task 4: Reduce and Verify

Implement quantized_matmul_splitk_reduce in src/extensions/src/quantized_matmul.metal and complete the corresponding reduction dispatch in QuantizedMatmul::eval_gpu. Do not add a second public matmul function.

Launch one reduction thread per output element. Sum all partition values in FP32 and cast once to the model dtype. Test:

Qwen3-4B’s 2560 -> 1024 K/V projection;
a partial 32-column output tile;
a shape whose base grid already reaches 320 groups and therefore falls back exactly to Day 6.

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

for context in 16 32 64 128 2048; do
  for projection in q k v o gate up down; do
    pdm run bench-week2-operators --solution tiny_llm --model qwen3-4b \
      --section prefill-projections --context "${context}" \
      --prefill-projection "${projection}" --include-split-k
  done
done

Benchmark Analysis: Complete Week 2

Compare Day 6, Day 7, and MLX at short, acceptance, and long prompt lengths. Split-K should help only while the unsplit output grid is under-filled. Verify that one-token decode remains unchanged because it still dispatches to Day 3’s matvec, and that sufficiently large prefill shapes select the unsplit Day 6 kernel instead of paying for partial storage and reduction.

Keep a short complete-model control beside the under-filled shape sweep, then run the fixed Week 2 acceptance workload from Day 3. The performance appendix is the single place for the measured hardware, dependency versions, checkpoint table, and final MLX ratios.

pdm run bench-week2-progression --offline --solution tiny_llm --repeats 4 \
  --variant week2-simd-matmul --variant week2-split-k --variant mlx \
  --model qwen3-4b --input-len 32 --output-len 33 --warmup 2 \
  --prefill-logits last

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

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

Repeat the operator comparison at the 128-token acceptance shape and at a long control such as 2,048 tokens. Attach the three end-to-end comparisons and the per-projection SIMD/Split-K/MLX table at each crossover candidate. Retain Split-K only below the measured crossover: it must improve the under-filled projection, preserve one-token decode, and fall back exactly to Day 6 when the ordinary result grid is already occupied. Record the accumulation and reduction dispatches beside the calculated partition policy and operator table. The final stretch-goal acceptance run must still reach 80% of MLX in both phases.

The reference checkpoint pairs the short-shape operator gains with the end-to-end result and keeps the neutral acceptance and long controls separate. Verify directly that the short shape executes the accumulation and merge pipelines, while the calculated policy names the partitions and the shape sweep prevents their overhead from leaking into occupied controls. Week 3 then changes the benchmark itself: request turnover and dense KV reconstruction, rather than another static projection, become the measured serving bottleneck.

The Week 2 loop is now complete:

optimize matvec -> benchmark decode -> optimize model kernels -> benchmark decode
-> optimize attention -> benchmark prefill -> optimize cooperative matmul
-> measure tile occupancy -> optimize split-K -> benchmark the complete checkpoint

Week 3 inherits these projection schedules. Paging is evaluated separately on cache writes, direct page reads, attention time, and end-to-end throughput; it does not receive credit for the Day 7 projection gain.

🚧 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 with separate schedules for one-token decode and multi-token prefill.

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 speculative decoding over rewindable caches
Optional Mixture-of-Experts model support

Day 1 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 schedule. Day 5 then tiles that same page-walking operation with Week 2’s matrix fragments. 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 long prefill and the paged vector kernel for short queries. Both schedules read the same page pool through the same block-table interface; neither 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.

Speculative decoding follows the paged-attention chapters because rejection needs a precise cache rewind operation, and multi-token verification needs the page-aware long-query path. MoE is independent of the cache and scheduler, so it remains an optional model extension and is not required to complete Week 3.

🚧 Week 3 Day 1: Continuous Batching

🚧 This chapter is under review and may change.

In this chapter, we will implement continuous batching, which keeps a batch of active requests on the device and replaces each request as soon as it finishes.

So far, each generation loop has processed only one request. That may not provide enough work to use the device efficiently, so we will decode 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 both attention paths 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: Exercise the Batch-Ready Model

src/tiny_llm/qwen3_week2.py  (reuse unchanged)

Call the Week 2 model with several requests, one offset per batch element, and the mask returned by BatchingKvCache. Exercise requests joining and leaving at different positions. The model remains request-agnostic; slot ownership and lifecycle belong to the cache and scheduler.

You should pass all of the tests by running:

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. 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.

Run the complete scheduler with:

pdm run batch-main

By default, this command uses Qwen3-0.6B with a batch size of five and a fixed set of prompts. Record the longest interval between consecutive decode steps when one queued request has a much longer prompt. That interval is the baseline for Day 2. Use Day 2’s bench-chunked-prefill runner for a publishable 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 rather than relying on the interactive prompts above.

🚧 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

pdm run bench-chunked-prefill --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/m4-pro-qwen3-0.6b-week3-chunked-prefill-mlx-0.32.0.json

The checked trace 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. The JSON stores every prompt token id, the per-request output budget, and their canonical SHA-256 checksum.

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	Requests/s	Decode step p95	Decode gap p95 / max
32	105.47	3.296	17.52 ms	30.39 / 32.47 ms
128	144.91	4.528	18.78 ms	46.52 / 48.80 ms
512	157.00	4.906	19.57 ms	76.04 / 122.16 ms

The 512-token row is the full-prompt Day 1 control for this trace. Reducing the budget makes the p95 completion gap monotonically smaller, while the 32-token budget gives up substantial throughput. The course uses 128 as a measured compromise for this workload, not as a universal optimum.

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

Complete every learner-extension integration point before rebuilding:

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

Then rebuild:

pdm run build-ext

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

Prefill with Pages

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

chunk = [t0 t1 t2 t3 t4 t5]

One possible layout is:

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

That layer cache’s metadata becomes:

page_ids = [7, 2]
context_len = 6

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

Decode with Pages

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

With paged storage:

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.

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

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

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

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

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

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

What Paging Changes

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

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

pdm run test --week 3 --day 3

🚧 Week 3 Day 4: Direct Paged Attention

🚧 This chapter is under review and may change.

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

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

Paged KV Cache vs Paged Attention

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

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.

Choose a Schedule for Each Query Shape

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

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.

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

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`	Direct paged prefill	Walk logical K/V tiles through the block table and keep the schedule deliberately inspectable. Day 5 optimizes it.

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

How This Maps to `tiny-llm`

`src/tiny_llm/attention.py`

Add a new function:

def paged_attention(...):
    ...

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

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 4. Keep the page-walking schedule simple enough that block-table and tail-page boundary errors are visible. Day 5 will tile its inner matrix work while preserving this address calculation.

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

The page pool should therefore expose contiguous physical storage:

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

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

`src/tiny_llm/qwen3_week3.py`

The attention module should call the paged runtime directly:

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

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

`src/tiny_llm/batch.py`

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

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

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

Implementation Order

Use this implementation order:

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 these exact starter functions:

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

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

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

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

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

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

Implement two correctness-first GPU dispatches:

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, assign query rows to a direct page-walking schedule and resolve every K/V tile through block_table. When a tile is aligned and cannot cross a page boundary, share its one physical page id across the whole tile. Favor inspectable ownership over the final tiled performance schedule.

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

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

Your solution’s boundary

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

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

Task 3: Dispatch from the Model

src/tiny_llm/qwen3_week3.py

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

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

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

This creates the Day 4 routing policy:

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

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

Task 4: Connect It to Continuous Batching

src/tiny_llm/batch.py

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

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

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

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

Measure the Direct Page Walk

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

Record three operator baselines on the same machine:

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 --offline --contexts 128 1024 \
  --page-size 128 --warmup 5 --iterations 60 --repeats 4 \
  --cooldown-seconds 1 \
  --json-output benchmark_results/m4-pro-qwen3-4b-week3-attention-mlx-0.32.0.json

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

Context	Dense + gather	Direct paged	MLX fused
128	184.01 us	187.55 us	153.59 us
1,024	420.88 us	249.79 us	207.18 us

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

Checkpoint 1: Establish a Correct Direct Path

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

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

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

Checkpoint 2: Design the Decode Schedule

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

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 --offline --repeats 4 \
  --model qwen3-4b --num-seqs 16 --batch-size 4 \
  --min-input-len 128 --max-input-len 1024 \
  --min-output-len 32 --max-output-len 128 --prefill-step 128 \
  --warmup 1 --cooldown-seconds 1 \
  --json-output benchmark_results/m4-pro-qwen3-4b-week3-serving-mlx-0.32.0.json

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

pdm run test --week 3 --day 4

🚧 Week 3 Day 5: Paged FlashAttention

🚧 This chapter is under review and may change.

In this chapter, we will tile page-aware attention for multi-token queries. The operator translates logical K/V positions through block_table, stages page-backed tiles on chip, and combines them with online softmax. Short queries continue to use the vector decode schedule from Day 4; long prefill chunks use the tiled schedule developed here.

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

Prerequisites

This chapter combines four prerequisites:

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

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

Why Optimize the Paged Path

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

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 is still exact attention. 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",
)

Put the 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.

The completed Week 3 model therefore has one paged-attention contract and two workload-specific GPU schedules.

Task 1: Tile Queries and Paged K/V

Begin paged_attention_mma_bf16_d128 in src/extensions/src/paged_attention.metal. Keep paged_attention_decode and paged_attention_scalar_f32 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. Your Metal kernel may use MLX’s low-level Steel block-loader header for this load primitive, while your solution owns the page translation, tile schedule, online softmax, primitive, and dispatch. It does not instantiate MLX attention.

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;
only then benchmark the tiled kernel.

Required fixtures include:

a context contained in one page;
a tile that crosses a page boundary;
non-consecutive physical page ids;
L = 65 and a context whose length is not a tile multiple;
causal decode after the paged prefill;
GQA where multiple query heads map to one K/V head;
output dtype remains BF16.

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

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.

Measure the completed operator in the continuous-serving trace. Report prompt range, page size, batch size, hardware, prefill throughput, decode throughput, request throughput, peak KV storage, and logical KV copy volume:

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

FlashAttention is expected to matter more as prefill grows. It should not replace the Day 4 decode schedule: a one-token query has no query-tile reuse.

On the checked M4 Pro trace, the complete direct-paged path reaches 679.56 prefill tok/s, 41.88 output tok/s, 82.11 decode tok/s, and 0.558 requests/s. Relative to dense serving on the same trace, output and request throughput are 28.7% higher, decode throughput is 62.8% higher, and peak KV storage is 47.4% lower. These are cumulative Week 3 path results; the serving trace does not isolate the Day 5 prefill schedule from paging, direct decode, or scheduling.

Use a separate 8K static sweep as a kernel diagnostic after the serving trace. It shows when query tiling begins to offset page-table overhead, but it does not 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 --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/m4-pro-qwen3-4b-week3-8k-mlx-0.32.0.json

8K static checkpoint	Prefill tok/s	Decode tok/s
Week 2	323.26	17.62
Complete Week 3	424.14	24.96
MLX	594.21	25.24

The complete Week 3 prefill path is 31.2% faster than Week 2 and reaches 71.4% of MLX at this shape. Its decode row is the Day 4 vector schedule, not evidence for the tiled prefill kernel.

🚧 Week 3 Optional Extension: 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.

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: Make Cache Rewind a Contract

Add rewind(n) to the common KV-cache interface. 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.

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,
)

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.

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_ref --loader week3 \
  --draft-model qwen3-0.6b --model qwen3-4b

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, output budget, seed, and synchronization boundary. 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 3 Optional Extension: 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.

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 Quantized Matmul

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. In this chapter, we will build a narrow teaching version of that idea: grouped_quantized_matmul.

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].

Task 1 will assume the rows are already sorted by expert id. The MoE helper will keep the inverse order from the sort so the result can be restored to the original token order.

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

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 Quantized Matmul

src/extensions/src/quantized_matmul.cpp
src/extensions/src/quantized_matmul.metal
src/tiny_llm/quantize.py
src/tiny_llm/moe.py

Implement grouped_quantized_matmul, then use it from grouped_expert_linear. This is the quantized grouped-matmul core of MoE.

This optional interface is intentionally not predeclared in src/extensions/src/tiny_llm_ext.h, the bindings, or the core CMake target. The required Week 2/3 interfaces are scaffolded from setup, but this optional chapter is a staged reveal: if you choose the extension variant, add the new tiny_llm_ext::grouped_quantized_matmul declaration, binding, C++ source function, grouped_quantized_matmul Metal kernel, and build registration here. Then modify the existing grouped_expert_linear function in src/tiny_llm/moe.py to call it. Keeping it out of the core starter prevents an optional future interface from appearing to be required by earlier checkpoints.

grouped_quantized_matmul accepts:

a:           R, D
w_experts:   packed QuantizedWeights for num_experts, output_dim, D
expert_ids:  R, sorted by expert id

It returns:

out:         R, output_dim

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

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

The implementation should:

1. add a Python wrapper for grouped_quantized_matmul,
2. extend the quantized matmul extension with a grouped entrypoint,
3. read expert_ids[row] inside the kernel,
4. use that expert id to choose the expert weight, scale, and bias row.

After that, implement grouped_expert_linear in src/tiny_llm/moe.py:

1. flatten token rows and expert ids,
2. sort rows by expert id,
3. call grouped_quantized_matmul,
4. restore the original order.

The call should look like:

out = grouped_quantized_matmul(
    w_experts.scales,
    w_experts.biases,
    group_size=w_experts.group_size,
    bits=w_experts.bits,
    a=grouped_rows,
    b=w_experts.weight,
    expert_ids=grouped_expert_ids,
    transpose_b=True,
)

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.

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, plus dispatch_model in src/tiny_llm/models.py.

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.

Run the focused tests with:

pdm run test --week 3 --day 6

Run this task through the normal generation entrypoints instead of adding a separate serving entrypoint. For example:

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 4: Build a Coding Agent

Course status: Week 4 is being published one checkpoint at a time. Days 1 through 8 are ready to learn and review. Additional capabilities will appear only after their implementation, starter, and reviews are ready.

Weeks 1 through 3 turn tokens into text and make serving that text efficient. Week 4 starts a different kind of program: an agent repeatedly asks the model for one structured action, records the result, and gives that result back to the model. Day 1 builds the bounded deterministic loop. Day 2 gives that loop a small read-only workspace for listing and reading project files. Day 3 adds an approval boundary for edits, one exact validation command, and simple receipts. Day 4 saves a complete conversation boundary and the fake model’s cache metadata, then restores both into a fresh model without replaying effects. Day 5 derives a smaller model-visible transcript from older completed effects while their exact receipts remain unchanged. Day 6 inspects one safe checkpoint, adds one visible operator steering message, and resumes a fresh model without replaying the completed effect. Day 7 evaluates one completed run from declared final, file, result, and receipt facts without grading hidden reasoning or exact transcript shape. Day 8 reconnects the agent to the inference system from Weeks 1–3: it saves one real tokenizer/KV prefix, forks two isolated steered continuations without rewinding completed effects, evaluates both, and makes one explicit selection.

What Day 1 Builds

Day 1 establishes the protocol and control flow that later checkpoints will extend:

task + system instruction
          |
          v
     model response
          |
          v
  validate one JSON action
      |             |
      | tool        | final answer
      v             v
 execute through  stop the run
 a supplied workspace
      |
      v
 append an observation and continue

The workspace in the Day 1 tests is a small fake object with one enabled read-only action. This keeps the learning target focused: validate the model’s text, bound the loop, and preserve the conversation. Day 2 replaces the fake with real read-only tools without changing the Day 1 loop contract.

Day 1 Checkpoint

The Day 1 starter exposes exactly these public names:

File	Public names	Why they are here
`src/tiny_llm/agent/generation.py`	`initial_messages`, `generate_response`	Begin a conversation and keep the one-response model boundary explicit.
`src/tiny_llm/agent/protocol.py`	`AgentError`, `FinalAction`, `ToolAction`, `parse_action`, `build_system_prompt`	Represent and validate one final answer or one tool request.
`src/tiny_llm/agent/loop.py`	`AgentLimits`, `AgentEvent`, `AgentRun`, `run_agent`	Bound the run and retain an inspectable trace.

Implement the Day 1 exercise, then run:

pdm run test --week 4 --day 1

The learner test uses scripted responses and a fake workspace, so it does not download or load a model. To check the supplied implementation without copying the learner test, run:

pdm run test-refsol --week 4 --day 1

Publication Boundary

After Day 1 passes, continue with Day 2: Inspect a Workspace. The cumulative Day 2 command is:

pdm run test --week 4 --day 2

After Day 2 passes, continue with Day 3: Edit, Validate, and Record. Its cumulative command is:

pdm run test --week 4 --day 3

After Day 3 passes, continue with Day 4: Checkpoint and Resume. Its cumulative command is:

pdm run test --week 4 --day 4

After Day 4 passes, continue with Day 5: Compact Completed Work. Its cumulative command is:

pdm run test --week 4 --day 5

After Day 5 passes, continue with Day 6: Inspect and Steer a Paused Agent. Its cumulative command is:

pdm run test --week 4 --day 6

After Day 6 passes, continue with Day 7: Evaluate Observable Outcomes. Its cumulative command is:

pdm run test --week 4 --day 7

After Day 7 passes, continue with Day 8: Fork, Steer, and Select. Its cumulative command is:

pdm run test --week 4 --day 8

Only the Day 1 through Day 8 starter modules are published. Do not add session trees, effect rewind, reconciliation, an LLM judge, radix serving, or other later public APIs to your solution.

Day 1: A Validated Agent Loop

Day 1 scope: This chapter teaches a bounded loop and a JSON action protocol. The supplied test uses a fake read-only workspace. File mutation and command execution are not Day 1 capabilities.

A text generator returns one response and stops. A coding agent needs a small control loop: it asks for one response, decides whether that response is a final answer or an action, records what happened, and repeats when an action produces an observation.

The model never edits a file directly. It emits text. Ordinary Python code validates that text before handing a parsed action to the workspace object. That separation is what makes the loop testable without a model.

Files and Commands

Implement these Day 1 starter functions:

File	Function or type	Your responsibility
`src/tiny_llm/agent/generation.py`	`initial_messages()`, `generate_response()`	Reject a blank task, create the first messages, and decode one response with a fresh cache.
`src/tiny_llm/agent/protocol.py`	`AgentError`, `FinalAction`, `ToolAction`, `parse_action()`, `build_system_prompt()`	Define the exact action vocabulary, validate one JSON object, and describe only enabled actions.
`src/tiny_llm/agent/loop.py`	`AgentLimits`, `AgentEvent`, `AgentRun`, `run_agent()`	Bound the loop, append observations, and return an auditable result.

Run the focused learner check:

pdm run test --week 4 --day 1

It should pass without loading a model. For the supplied reference check:

pdm run test-refsol --week 4 --day 1

generate_response() is still a Day 1 public boundary even though the focused test deliberately avoids model weights. Render the messages with the tokenizer, decode at most the requested token count using a fresh cache, stop at EOS, and release every cache in a finally block. The scripted loop tests are the fast way to verify the control flow; a real model is not required for this checkpoint.

One Response, One Structured Decision

Use JSON because it makes the protocol visible in a trace. A response is either a final answer:

{"final":"I inspected README.md."}

or a tool request:

{"tool":"read_file","path":"README.md"}

parse_action() must accept exactly one JSON object. It rejects malformed JSON, non-object values, an empty final answer, an unknown tool, disabled tools, missing required fields, unexpected fields, and fields with the wrong shape. Do not quietly ignore trailing or extra data.

TOOL_FIELDS names the complete future vocabulary: list_files, read_file, write_file, edit_file, and run_command. Day 1 does not implement those effects. Its fake workspace enables only read_file, which is enough to prove that the loop validates availability before dispatching an action.

Start a 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. A clear first message lets later turns grow from a known history instead of assembling prompt fragments ad hoc.

build_system_prompt(workspace) describes the enabled action set for this run. The prompt is guidance, not enforcement: parse_action() and the workspace boundary must independently reject anything the policy does not allow.

The Bounded Loop

run_agent() receives a task, a generate callable, and a workspace. The tests substitute a callable that returns predetermined strings, so the loop’s behavior stays deterministic.

messages = initial_messages(task, build_system_prompt(workspace))
for step in range(1, limits.max_steps + 1):
    response = generate(messages)
    action = parse_action(response, workspace.available_tools)

    if action is a final answer:
        record it and stop

    result = workspace.execute(action)
    record the action and result
    messages = append the assistant response and tool observation

The real implementation also turns a parse failure into an observation such as error: response is not valid JSON: ..., then lets the model try again. This is more useful than crashing the agent for one malformed answer.

Every interaction becomes an AgentEvent containing the step number, raw response, parsed action when one exists, and result or validation error. The returned AgentRun records whether a valid final answer completed the protocol. It does not prove that a task was solved; task grading is a later course concern.

Stop Conditions Are Part of Correctness

AgentLimits requires positive values. Implement all of these terminal cases:

a valid final answer returns completed;
reaching max_steps returns step_limit;
too many invalid actions returns invalid_action_limit;
an overlong conversation returns context_limit; and
too many identical tool requests returns repeated_action_limit.

The repeated-action check matters even when a tool succeeds. Repeating the same request can burn the whole budget while adding no new information.

Exercise Checklist

Before considering Day 1 complete, make the focused test demonstrate all of these behaviors:

A task starts with a system message and a user message.
A read_file request reaches the fake workspace, its result becomes an observation, and a later final answer stops the run.
Invalid JSON and an unavailable tool become recoverable error observations.
The loop stops at the step budget.
Repeated identical actions stop before the general step budget is spent.

Keep the solution inside the Day 1 starter files.

🚧 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

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`	Make the cumulative checkpoint importable.

ToolPolicy has three fields, in order:

root: Path
max_file_bytes: int = 64 * 1024
max_list_entries: int = 200

Workspace exposes:

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 add write, command, approval, receipt, session, checkpoint, or rewind APIs 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 Cumulative Checkpoint

From the repository root, run:

pdm run test --week 4 --day 2

This copies the cumulative Day 2 learner test into tests/ and runs it against tiny_llm. During course development, check the supplied implementation without copying the learner test:

pdm run test-refsol --week 4 --day 2

The cumulative 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.

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"

When this checkpoint is green, continue to Day 3: Edit, Validate, and Record.

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

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

Implement the TODO bodies in these cumulative 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`	cumulative Day 1–3 API	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 Cumulative Checkpoint

From the repository root, copy and run the learner checkpoint:

pdm run test --week 4 --day 3

Before you implement the TODOs, the copied test is expected to fail because the new starter methods return None. Keep those 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 cumulative course-code guard checks exact public signatures, dataclass fields, package exports, TODO-only starter bodies, and absence of future APIs.

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.

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 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.

Check the Chapter

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	Export the cumulative Day 4 API.

Run the 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.

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

The Day 4 starter adds only checkpoint.py and two loop entry points. Do not add session IDs, parent pointers, branches, rewind methods, compaction summaries, steering queues, disk cache files, or later-day modules. 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.

Continue with Day 5: Compact Completed Work to derive a smaller model-visible transcript while keeping the exact effect receipts.

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

Day 5 adds 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	Export the cumulative Day 5 API.

Copy the learner test, then run it:

pdm run copy-test --week 4 --day 5
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.

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.

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 the Day 4 fake model’s simple counter, so they do 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.

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.

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.

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

Day 6 adds 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	Export the cumulative Day 6 API.

Copy and run the six learner tasks:

pdm run copy-test --week 4 --day 6
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.

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.

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

Day 7 adds 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	Export the cumulative Day 7 API.

Copy and run the seven learner tasks:

pdm run copy-test --week 4 --day 7
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.

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.

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

Day 8 adds 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	Export the cumulative Day 8 API.

Copy and run the five learner tasks:

pdm run copy-test --week 4 --day 8
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.

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)

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.

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.

🚧 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 --repeats 4 --cooldown-seconds 1 \
  --model qwen3-4b --input-len 128 --output-len 129 --warmup 2 \
  --prefill-logits last --json-output week2-128.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/m4-pro-qwen3-4b-week3-serving-mlx-0.32.0.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. Static Week 2 rows use two complete warmups and the median of four 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 6	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	`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	Model dispatch is `L <= 2`, `S <= 256`, and no explicit array mask; the context sweep wins 6/6 passes through 256, while the query sweep is repeat-consistent only through `L=2`	Python `mlx.core` grouped attention handles longer queries, longer contexts, and explicit array masks	Fixed implementation order, GPU performance-state drift, extrapolating beyond 256, or treating correctness at `S=1` as schedule efficiency
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	Helps short narrow Qwen projections, is neutral around the 128-token acceptance shape, and loses once the base grid is occupied	`split_k <= 1` dispatches exactly to the Day 6 unsplit kernel	Profiling independent layers can hide under-occupancy that appears in the dependency-ordered model

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

Week 2 has one fixed acceptance shape: Qwen3-4B, a 128-token prompt, 128 timed decode steps, last-row logits, two complete warmups, and the median of four fresh processes. Two passes use forward checkpoint order and two use reverse order. The output length is 129 because prefill produces the first generated token.

Each row is cumulative. Day 2 retains the Day 1 checkpoint while it establishes the synchronized benchmark. Day 3 then completes the packed quantized-matvec checkpoint.

Chapter	Cumulative checkpoint	Prefill tok/s	Decode tok/s	Output tok/s	Change selected by the preceding evidence
Day 1	Dense request KV cache	730.43	24.63	24.01	Stop full-prefix decode recomputation.
Day 2	Benchmark baseline	730.43	24.63	24.01	Measure dense projection weight traffic.
Day 3	Quantized matvec	105.00	58.71	37.95	Keep weights packed and add the x4 decode kernel.
Day 4	Fused model kernels	105.97	75.21	44.33	Remove the newly exposed pointwise graph launches.
Day 5	Bounded decode attention	105.99	75.75	44.50	Historical row from before the guard extended through `S=256`; do not use it as the current checkpoint delta.
Day 6	SIMD-matrix prefill	797.45	75.12	69.17	Fix the quantized matrix path exposed by Day 3.
Day 7	Split-K prefill	792.55	75.41	69.37	Fill the GPU only for under-occupied short projections.
Baseline	MLX 0.32.0	830.49	89.37	81.30	External denominator.

The checked-in progression file predates the current L <= 2, S <= 256 guard. With the current implementation, prefill has L=128 and stays on the Python mlx.core path, while timed one-token decode steps see S=129 through S=256 and enter the custom path. The historical Day 4-to-Day 5 difference is not a current end-to-end measurement; the balanced context and query sweeps below are the checked evidence for the production guard.

Checked Operator Attribution That Selects Each Chapter

The checked reference-solution attribution does not replace an operator with an MLX operator. It calls the projection, attention, pointwise, and cache paths from tiny_llm_ref at Qwen3-4B shapes and replays each group at the model’s real dispatch count. The projection replay preserves the transformer dependency order so work from a later MLP cannot hide an under-filled attention projection. Each round rotates the category order, synchronizes every category once, and the median follows four warmups and twelve samples. This historical evidence is checked in for readers; reproducing it is not a learner requirement.

The bar widths below are normalized within a checkpoint. The time at the right is the sum of the synchronized category medians, not a throughput measurement. Forcing category boundaries prevents some whole-graph fusion, so use the shares to rank work and the fresh-process checkpoint table above to accept or reject a change.

This is an operator-attribution chart, not a Metal flame graph. It ranks model operator families and explains why the course tackles the kernels in this order.

The profile makes the progression concrete:

Cached decode spends 81.5% of attributed time in dense projections. Day 3 therefore changes weight storage and the decode projection schedule first.
After packed matvec, the pointwise group is 35.8% while attention is only 4.5% at the 128-token acceptance context. Day 4 therefore removes the measured normalization, position, and activation overhead first.
After the Day 4 pointwise kernels, the balanced operator sweeps isolate a removable attention gap through S=256 and a repeat-consistent query-length win through L=2. Day 5 tests online softmax inside those bounds.
At the fixed workload, 128-token prefill remains outside the query-length guard. Its profile makes the vanilla quantized projection path 99.0% of attributed prefill time, which selects the cooperative matrix kernel in Day 6; one-token decode uses the bounded Day 5 path.
After Day 6, projections remain most of the inherent prefill work, but the long-shape operator comparison is already close to MLX. The 32-token shape sweep then isolates under-occupied Qwen projections and selects Split-K only below their measured crossover.

The checked-in raw profile is benchmark_results/m4-pro-qwen3-4b-week2-kernel-profile-mlx-0.32.0.json. The balanced fresh-process samples are benchmark_results/m4-pro-qwen3-4b-week2-progression-mlx-0.32.0.json.

The operator tables below use bench-week2-operators with twelve warmup rounds and sixty measured rounds. Each round synchronizes every implementation, and the runner rotates through every execution order so GPU performance-state drift does not consistently favor Python reference code, the course kernel, or MLX. These latencies are microbenchmarks; only the fresh-process table above accepts an end-to-end checkpoint.

Day 1: Cache the Prefix

The dense cache makes prefill a one-time cost, but every decode projection still reads dense weights. Day 1 therefore starts with respectable prefill and only 24.63 decode tok/s. The result gives Day 2 a real cached baseline to measure.

Day 2: Measure Before Optimizing

Day 2 changes the measurement discipline rather than the model. The end-to-end row and synchronized attribution answer different parts of the handoff:

Evidence	Result	Decision
Complete-model decode	24.63 tok/s; MLX 89.37 tok/s	A large decode gap remains.
Dense projections	33.66 ms, 81.5% of attributed time	Optimize projection weight traffic first.
Pointwise operators	6.45 ms, 15.6%	Defer until projections shrink.
Attention	0.85 ms, 2.1%	Do not select attention from this workload.
KV growth	0.33 ms, 0.8%	The dense cache already removed prefix recomputation.

The operator-family result is sufficient to select the quantized-matvec work for Day 3. The isolated packed-W4 control is not the Day 2 model’s dense projection; it remains a readable schedule comparison without pretending that one shader ranked the complete model.

Day 3: Keep Weights Packed

The x4 W4A16 matvec raises complete-model decode from 24.63 to 58.71 tok/s, a 138.4% gain. Prefill falls from 730.43 to 105.00 tok/s because matrix-shaped inputs still use the vanilla Metal quantized kernel. The operator microbenchmark checks whether the decode gain came from the intended projection schedule:

Qwen3-4B projection, `M=1`	Vanilla Metal	Packed matvec	MLX
Q	750.3 us	187.6 us	183.4 us
K	239.5 us	145.1 us	147.8 us
V	244.8 us	147.0 us	138.9 us
O	590.3 us	163.7 us	160.2 us
MLP gate	908.8 us	182.5 us	177.2 us
MLP up	948.0 us	185.6 us	182.9 us
MLP down	1,243.3 us	188.3 us	181.6 us
Vocabulary head	11,086.1 us	1,030.2 us	1,029.3 us

The packed operator is close to MLX at every listed shape. Projections still occupy 57.9% of the synchronized model replay because every layer inherently uses them, but normalization, position, and activation now occupy 35.8% and are the larger removable gap. That combination, rather than the absolute height of the projection bar, selects Day 4.

Day 4: Fused Model Kernels

The cumulative model and operator results agree on all three retained changes:

Checkpoint	Decode tok/s	Python reference	Fused operator	MLX operator
Day 3 packed matvec	58.71	–	–	–
Fast RMSNorm	65.94	210.0 us	168.2 us	147.1 us
Fast RoPE	71.16	180.9 us	144.8 us	118.7 us
Fused SwiGLU	75.21	189.4 us	125.7 us	137.2 us

The pointwise group falls from 35.8% after Day 3 to 10.5%. Projections are now 80.5% of attributed decode time but are already close to their MLX operator latencies. A direct dispatch trace can verify that the RMSNorm, RoPE, and SwiGLU pipelines all ran. The balanced S=32,128,160,192,256 sweep then isolates an attention opportunity through the largest measured context; the query-length sweep supplies the other dispatch boundary.

Day 5: Fused Decode Attention

The matched short-context model checkpoint uses a 32-token prompt and an output length of 97. Prefill produces the first token, so all 96 timed decode calls grow the cache from S=33 through S=128 and enter the custom guard. Under that workload, fused attention raises median decode from 59.90 to 61.78 tok/s (+3.1%) and output throughput from 48.52 to 49.54 tok/s (+2.1%). MLX reaches 68.86 decode tok/s, so the bounded checkpoint reaches 89.7% of that matched denominator. The raw samples are checked in at benchmark_results/m4-pro-qwen3-4b-week2-short-context-mlx-0.32.0.json.

The current context sweep includes the FP32 promotion and output cast used by the Python mlx.core fallback. It uses six forward/reverse context passes, rotates every implementation order, and retains 60 samples per implementation and pass:

Cached context	Python reference	Fused	MLX	Fused vs Python	Pass wins
32	143.0 us	125.7 us	116.3 us	1.138x	6/6
128	149.3 us	136.3 us	120.6 us	1.095x	6/6
160	151.2 us	140.1 us	120.9 us	1.079x	6/6
192	154.0 us	143.9 us	121.9 us	1.071x	6/6
256	158.0 us	150.7 us	122.8 us	1.048x	6/6

The query-length sweep holds S=128, Qwen3-4B’s 4:1 GQA ratio, and the causal form while balancing L1/L2/L4/L8 order over six passes:

Query length	Python reference	Fused	MLX	Fused vs Python	Pass wins
1	244.4 us	213.1 us	155.9 us	1.147x	6/6
2	341.4 us	258.8 us	185.3 us	1.319x	6/6
4	322.7 us	297.3 us	197.4 us	1.085x	4/6
8	377.7 us	491.5 us	290.6 us	0.768x	0/6

At L=1, the causal mask permits the entire existing cache and is equivalent to unmasked one-token decode; longer rows measure causal multi-token chunks. The context sweep supports S <= 256, while L=2 is the largest repeat-consistent query-length win. Those results define the current L <= 2, S <= 256 guard. The checked raw records are benchmark_results/m4-pro-qwen3-4b-week2-attention-context-sweep-mlx-0.32.0.json and benchmark_results/m4-pro-qwen3-4b-week2-attention-query-sweep-mlx-0.32.0.json.

In the fixed 128-token workload, prefill remains outside the query-length guard and attributes 1,196.34 ms of 1,208.78 ms, or 99.0%, to quantized projections; attention accounts for 6.08 ms and the pointwise group for 6.35 ms. That prefill bottleneck selects the matrix-shaped projection kernel in Day 6.

Day 6: Use Cooperative Loads for Quantized Prefill

At the fixed-workload prefill checkpoint, quantized projections account for 1,196.34 ms of the 1,208.78 ms attributed profile, or 99.0%. The cooperative matrix schedule reduces attributed projection time to 147.63 ms and raises complete-model prefill from 105.99 to 797.45 tok/s. MLX reaches 830.49 tok/s.

The long-row controls show that the tile arithmetic and cooperative loads are healthy once the result grid is occupied:

Projection at `M=2048`	Day 6	MLX
Q, `2560 -> 4096`	6.64 ms	6.65 ms
K, `2560 -> 1024`	1.78 ms	1.79 ms
O, `4096 -> 2560`	6.82 ms	6.66 ms
MLP gate, `2560 -> 9728`	15.65 ms	15.75 ms
MLP down, `9728 -> 2560`	16.25 ms	15.90 ms

At the 128-token acceptance shape, the same projections are within 3.1% of MLX. The short-row control exposes a different pattern:

Projection at `M=32`	Day 6 SIMD	MLX	Gap
Q	733.0 us	643.9 us	13.8%
K	221.0 us	205.1 us	7.8%
O	256.3 us	241.4 us	6.2%
MLP gate	410.0 us	406.6 us	0.8%
MLP down	440.1 us	391.7 us	12.4%

The operator gaps correlate with result-grid size rather than reduction width or arithmetic. For the narrow K projection, the unsplit launch geometry is:

Prompt rows	Row tiles	Output tiles	Independent threadgroups
32	1	32	32
128	4	32	128
2,048	64	32	2,048

The dispatch formula yields 32 independent threadgroups for the first row of this table. The long controls rule out a generally slow schedule, while the short operator table and calculated dispatch geometry select Split-K for Day 7.

Day 7: Split K Only Below the Crossover

The per-projection microbenchmark tests the proposed occupancy fix directly at M=32:

Projection	Day 6 SIMD	Split-K	MLX	Split-K effect
Q	733.0 us	612.3 us	643.9 us	1.20x faster
K	221.0 us	201.3 us	205.1 us	1.10x faster
O	256.3 us	243.8 us	241.4 us	1.05x faster
MLP gate	410.0 us	414.2 us	406.6 us	Falls back; within noise
MLP down	440.1 us	395.1 us	391.7 us	1.11x faster

The complete 32-token model confirms that the useful projection changes survive composition:

Checkpoint	Prefill tok/s	Decode tok/s	Prefill / MLX
Day 6 cooperative matmul	607.36	83.53	82.3%
Day 7 split-K	679.50	83.53	92.1%
MLX 0.32.0	737.55	90.52	100%

Split-K adds 11.9% complete-model prefill at this short shape. At M=128, the operator sweep is neutral: Q, O, gate, and down dispatch unchanged, while the narrow K split measures 221.2 us versus 222.8 us unsplit. The fresh-process acceptance result is likewise neutral at 792.55 versus 797.45 prefill tok/s. At M=2048, every projection falls back exactly to Day 6. Because the dispatch geometry, operator table, and end-to-end result agree. The direct dispatch trace must show the accumulation and merge pipelines, while the calculated policy supplies the partition count and the shape sweep decides where those costs are worthwhile.

The fresh-process samples for the short control point are checked in at benchmark_results/m4-pro-qwen3-4b-week2-32-mlx-0.32.0.json. The completed Week 2 path reaches 95.4% of MLX prefill, 84.4% of MLX decode, and 85.3% of MLX end-to-end output throughput at the 128-token acceptance shape. Both required phase ratios exceed 80%. Longer static sweeps remain attention diagnostics; they do not test the memory-management reasons for paging.

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/m4-pro-qwen3-4b-week3-serving-mlx-0.32.0.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.

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:

Prefill budget	Output tok/s	Requests/s	Decode step p95	Decode gap p95 / max
32	105.47	3.296	17.52 ms	30.39 / 32.47 ms
128	144.91	4.528	18.78 ms	46.52 / 48.80 ms
512	157.00	4.906	19.57 ms	76.04 / 122.16 ms

Because 512 covers every prompt in this trace, that row is the full-prompt Day 1 control. The monotonically smaller p95 gap at smaller budgets comes with lower throughput; 128 is the measured course compromise.

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	184.01 us	187.55 us	153.59 us
1,024	420.88 us	249.79 us	207.18 us

The direct operator is 1.9% slower than dense-plus-gather at 128 tokens and 40.7% faster at 1,024 tokens. MLX remains faster at both shapes. All three paths pass the checked BF16 correctness tolerance before timing.

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	718.30 prefill; 32.54 output; 50.42 decode tok/s	Establishes the dense serving baseline.
Day 3	Paged storage with compatibility gather	730.69 prefill; 38.44 output; 65.88 decode tok/s	+18.1% output; +30.6% decode; -50.6% copy volume.
Day 4	Direct paged decode schedule	82.11 aggregate decode tok/s	+24.6% decode over the compatibility gather path.
Day 5	Complete direct paged path	679.56 prefill; 41.88 output; 82.11 decode tok/s	+28.7% output and request throughput over dense serving.

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. Days 4 and 5 then remove that compatibility movement for decode and long-query prefill respectively.

Days 4 and 5 share the final direct-paged process: queries with L <= 8 dispatch to the Day 4 decode schedule, while longer chunks dispatch to the Day 5 tiled schedule. The phase timers report their decode and prefill throughput inside the same request trace; they are not results from different workloads.

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	718.30	32.54	50.42	0.433	1,096	209,532
Paged storage plus dense gather	730.69	38.44	65.88	0.512	—	103,445
Direct paged attention	679.56	41.88	82.11	0.558	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	58.95 / 90.60 / 169.01 ms	62.61 / 241.29 / 344.04 ms
Paged + gather	48.86 / 53.64 / 57.82 ms	50.65 / 221.90 / 233.89 ms
Direct paged	38.27 / 39.83 / 43.46 ms	39.13 / 224.70 / 240.99 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 improves output and request throughput by 28.7%, aggregate decode by 62.8%, and peak KV storage by 47.4% relative to dense serving. Avoidable logical copy volume falls by 99.8%. Relative to paged storage plus gather, the direct operator adds 9.0% output throughput, 24.6% decode throughput, and removes 99.5% of the remaining copy volume. Prefill is 5.4% below dense and 7.0% below gather at the 128-token serving chunk, so the chapter does not claim a short-chunk FlashAttention speedup.

The 8K static run remains a secondary kernel diagnostic, not a Week 3 headline or acceptance result. At that shape, paged FlashAttention raises prefill from 323.26 to 424.14 tok/s. MLX reaches 594.21 tok/s, so the complete Week 3 path in the reference solution reaches 71.4%. This shows where the cumulative paged prefill path begins to help without mixing a static denominator into the serving progression. One-token decode continues to dispatch to the Day 4 vector schedule.

The checked-in Week 3 files contain 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/m4-pro-qwen3-0.6b-week3-chunked-prefill-mlx-0.32.0.json
benchmark_results/m4-pro-qwen3-4b-week3-attention-mlx-0.32.0.json
benchmark_results/m4-pro-qwen3-4b-week3-8k-mlx-0.32.0.json
benchmark_results/m4-pro-qwen3-4b-week3-serving-mlx-0.32.0.json

Verify all four file hashes from the repository root with:

shasum -a 256 -c benchmark_results/m4-pro-week3-evidence-mlx-0.32.0.sha256

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 direct-paged median reaches 1,116 live pages out of 1,152 reserved pages, reuses 2,196 page allocations, and records 15,840 unused tail slots across layer caches. At the same peak-tail-waste snapshot, all live pages contain 133,632 token slots, so tail waste is 11.9%, or 61.9 MiB of KV storage. This denominator excludes unused reserved pool capacity. The run grows the layer pools 144 times because it starts empty. These counters make reuse, fragmentation, and measured KV headroom visible; static single-request latency cannot. They do not establish 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 balanced JSON tables and SVG above are the checked-in evidence for the current course. Learners are not required to generate Metal captures, Xcode visualizations, gpudebug reports, profiling microbenchmarks, or screenshots. The full profiling workflow will return when the macOS 27 tooling is available; until then, matched synchronized benchmarks are the acceptance evidence.

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
Growing short-context attention	Online-softmax decode kernel	Week 2 Day 5
Scalar/strided prefill projection loads	Cooperative 32×32×32 quantized matmul	Week 2 Day 6
Under-filled short-prefill result grid	Measured split-K dispatch	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