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.

Course Roadmap: What Depends on What

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

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

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

The cumulative dependencies are deliberate:

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

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

Week 2 operator off-ramps

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

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

To run a completed checkpoint without solving it first:

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

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

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

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

Choose a Model for Your Mac

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

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

How to Use This Book

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

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

About the Authors

This course is created by Chi and Connor.

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

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

Community

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

Get Started

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

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

1. Apple lists these tiers across the M2 Mac mini, M3 Pro and M3 Max MacBook Pro, M4 Mac mini, and M5 MacBook Air specifications. Higher-memory configurations are outside this table. ↩

2. Week 2 Days 1–2 use the dense Week 1 loader, so keep using the Week 1 recommendation until the packed quantized-matvec path is complete on Day 3. The larger Week 2 entries apply after that checkpoint. ↩ ↩2 ↩3 ↩4

3. 30B-A3B requires the optional Week 3 MoE implementation. In Week 4, select the Week 3 loader. Use batch size one and a short context when approaching this ceiling; 4B remains the required-course target. ↩ ↩2 ↩3 ↩4 ↩5 ↩6 ↩7 ↩8

Setting Up the Environment

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

Install PDM

Follow the official installation guide to install PDM.

Clone the Repository

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

The repository is organized as follows:

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

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

Install Dependencies

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

Check the Installation

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

Run Unit Tests

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

pdm run test

Download the Model Parameters

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

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

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

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

Then, you can run:

pdm run main --solution ref --loader week1

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

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

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

• Qwen3: Think Deeper, Act Faster
• Hugging Face Transformers — Qwen3
• vLLM Qwen3
• mlx-lm Qwen3
• Qwen3 Technical Report

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

$$\text{Attention}=\text{softmax}(\frac{Q{K}^T}{\sqrt{d_k}}+M)V$$

Here, $\frac{1}{\sqrt{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

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

• You could have designed state of the art positional encoding
• Roformer: Enhanced Transformer with Rotary Positional Encoding

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

• PyTorch RotaryPositionalEmbeddings API
• MLX Implementation of RoPE before the custom metal kernel implementation

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

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

• GQA paper
• Qwen3 layers in mlx-lm
• PyTorch scaled dot-product attention with enable_gqa=True
• torchtune.modules.MultiHeadAttention

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

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 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}{\sqrt{\text{mean}(x^2)+ϵ}}\cdot \text{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

• Attention is All You Need (Transformer Paper, Section 3.3 “Position-wise Feed-Forward Networks”)
• GLU paper: Language Modeling with Gated Convolutional Networks
• SiLU (Swish) activation function
• SwiGLU paper: GLU Variants Improve Transformer
• PyTorch SiLU documentation
• Qwen3 layers implementation in mlx-lm (includes MLP)

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:

$$\text{SiLU}(x)=x\ast \text{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_{gate}$)
• An up projection ($W_{up}$)
• 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_{down}$)

This can be expressed as:

$$\text{MLP}(x)=W_{down}(\text{SiLU}(W_{gate}(x))\odot W_{up}(x))$$ where $\odot$ 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

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

• A Simplified Explanation of the Transformer Block
• Attention is All You Need

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

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

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

1. Install Xcode:

   Install Xcode from the Mac App Store or from the Apple Developer website (this may require an Apple Developer account).

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

3. Install Xcode Command Line Tools:

   Open your Terminal and run:

   xcode-select --install
   

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

5. Accept the Xcode License:

   You may also need to accept the Xcode license:

   sudo xcodebuild -license accept
   

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

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

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: A Step Closer to vLLM

You begin Week 2 with the runnable Week 1 Python mlx.core Qwen model. Keep that model intact. Your work lives in a separate Qwen3 path that first changes the generation algorithm—prefill once, retain a dense KV cache, and decode one new token at a time—then replaces the operators that dominate that cached path.

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

What You Build

The seven days form one cumulative single-request path. The starter supplies model loading, the extension build system, benchmark runners, correctness tests, Python reference equations, and the stable interfaces between checkpoints. You implement the state transition on Day 1, establish the measurement control on Day 2, and own the operator work on Days 3–7:

Run the supplied test selector after each day. Then run the live model or benchmark command beside that day so an isolated kernel never counts as a finished checkpoint. The full campaigns, raw samples, rejected experiments, and retained reference schedules live in the performance appendix; you do not need to recreate that evidence ledger to complete the exercises.

The Cumulative Path

• 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 one fixed Qwen3-4B workload: 128 prompt tokens, 129 output tokens, last-row logits, two warmups, and four balanced fresh-process samples. On the checked M4 Pro, the final checkpoint reaches 88.2% of full-MLX prefill and 87.0% of full-MLX decode throughput. This is not a cross-shape or cross-device target.

The completed Week 2 solution does not call MLX-provided implementations of the operators it teaches. It implements quantized matmul, decode attention, RMSNorm, RoPE, and SwiGLU in its own Python, C++, or Metal code. In particular, the required checkpoint does not use mx.quantized_matmul, mx.dequantize, mx.fast operators, or mx.fast.scaled_dot_product_attention as shortcuts. Its matrix path also avoids mlx::steel: the course scaffold leaves the cooperative tile loader and direct Metal simdgroup_matrix fragment bookkeeping for you to complete. 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.

If you want to study the serving path without implementing every custom operator, Days 3–7 each name a local MLX substitute. Keep the same course interface and substitute only that day’s operator. This is different from --solution mlx, which runs the separate complete MLX model and bypasses the course-owned model, cache, and scheduler.

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

Daily Checkpoints

1. KV cache: port the Week 1 operators into a Week 2 model, add request-scoped state, and stop recomputing the prefix.
2. 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.
3. 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.
4. Fused model kernels: fuse RMSNorm, RoPE, and SwiGLU one operator at a time after packed projections narrow the benchmark gap.
5. Decode attention: introduce online softmax over its tested short-context range and verify it with a matched workload.
6. SIMD-matrix prefill: return to the fixed 128-token workload and replace the correctness-first matrix path with cooperative tiles.
7. 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:

Use the selector beside each chapter while you work; run the complete day gate before carrying its result forward.

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 keeps these Week 2 interfaces, but it deliberately changes projection ownership: canonical dense Week 3 and the Week 3 scheduler factory select MLX quantized projections. Cache management, attention, paging, batching, and scheduling remain course-owned. Full MLX is a separate benchmark baseline, not another name for this hybrid Week 3 course path.

The explicit projection seam is a teaching boundary, not a performance credit for paging. The performance appendix isolates the seam while holding the course-owned serving mechanisms fixed, then reports representative absolute performance after that choice. Keep those two questions separate.

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.

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 Day 1: KV Cache

You arrive with the Week 1 Qwen model and full-prefix generation loop still working. Day 1 leaves them unchanged and completes the separate Week 2 shells:

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

Those four pieces are your work. The starter already supplies the Week 1 operators and the model-loading boundary. Your first useful feedback is the focused learner gate:

pdm run test --week 2 --day 1

When it passes, run the kv-cache checkpoint shown in Task 4. That live call is where the cache becomes part of generation rather than an isolated data structure.

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:

$$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}+M\right)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:

1. Accepts the newly computed K and V for the incoming tokens.
2. Appends them along the sequence dimension.
3. 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.

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

Task 4: Connect the Serving Loop

src/tiny_llm/generate.py

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

Keep one result as the input to Day 2’s matched comparison. Every later command changes 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.

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 Day 2: Benchmarking and Profiling

Day 1 leaves you with a cached BF16 model and a working kv-cache checkpoint. Day 2 does not add another model operator. The supplied benches/bench.py runner already owns request generation, warmups, synchronization, phase timing, and cache release. Your job is to run one like-for-like comparison, preserve its configuration with the result, and use the decode roofline to choose the next change.

First verify the benchmark lifecycle:

pdm run test --week 2 --day 2

Then record one matched tiny_llm/MLX pair with the commands below. That JSON is the Day 2 checkpoint. 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; the focused Day 2 test checks both the successful and failing paths.

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.

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.

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.

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

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 Day 3: Quantize the Model

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

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

Start with the Python wrapper gate, then build and test the GPU operator:

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

Finally run the complete Day 3 gate and quantized-matvec model checkpoint. Packed storage or an isolated fast kernel is not completion: the cached model must dispatch through your quantized path.

📚 Readings

• Model Compression and Quantization
• MLX Extensions Development Guide
• Quantized Matmul on GPU (Video)

Debug Metal Without a CPU Twin

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

1. Write the equation in Python with mlx.core. This is the semantic oracle.
2. Translate it into a deliberately simple Metal kernel, usually with one thread responsible for one output element.
3. 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:

$$\hat{w}=qs+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:

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.

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]=\sum _{j=0}^{N-1}A[i,j]\times B[k,j]$$

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

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

where $g=\lfloor j/G\rfloor$ is the group index.

Substituting:

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

Rearranging:

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

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

The starter already implements QuantizedWeights.from_mlx_layer; inspect and reuse that packed-weight plumbing. Modify these learner-owned functions:

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

The starter code provides QuantizedWeights, a container for a quantized matrix and its dequantization parameters:

Its supplied from_mlx_layer method extracts these fields from an MLX quantized layer when loading the model. Do not replace it with a second loader.

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:

1. 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.
2. 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:

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

$$\sum _j{a}_j(s{q}_j+b)=s\sum _j{a}_j{q}_j+b\sum _j{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:

1. Get the Metal device and command encoder from the stream.
2. Load the quantized matmul kernel matching the output dtype from the Metal library.
3. Bind the input and output buffers and the dimension constants (M, N, K). The buffer order must match the kernel signature.
4. 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.
5. 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.

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

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

Keep one cumulative model row and one representative real-shape projection comparison. First require a clear decode gain over kv-cache; then use the projection row to decide whether the matvec still needs work. The complete campaign and reference attribution are in the performance appendix.

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

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 Day 4: Fused Model Kernels

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

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

Implement and integrate RMSNorm first, then RoPE, then SwiGLU. Each operator task has a focused test and a live cumulative checkpoint, so you can attribute a failure or regression before composing all three:

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

RMSNorm, RoPE, and SwiGLU recur around the projections in every transformer layer. Week 1 expresses them as Python mlx.core equations; your Week 2 path places the same interfaces over purpose-built Metal kernels.

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

Why Fusion Helps

Week 1’s Python mlx.core equations already 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

Keep one cumulative result per operator so a regression cannot hide inside the combined gain. The complete campaign and reference attribution are in the performance appendix.

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

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 Day 5: Fused Decode Attention

Day 4 leaves packed projections and three fused model kernels behind stable interfaces. Day 5 preserves the readable grouped-attention function in src/tiny_llm/week2_kernels.py, completes the separate decode-attention stubs in week2_kernels.cpp and week2_kernels.metal, and adds one visible dispatch guard in Qwen3MultiHeadAttention.__call__.

Your first milestone is a Python oracle with the same model-facing shapes. Then implement online softmax in Metal, compare the two directly, and integrate the custom path only for L <= 2, S <= 256, and ordinary None or causal masks. Run the day gate after rebuilding the extension:

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

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 mlx.core composition materializes the complete score and probability rows; Day 5 computes the same result while retaining only online-softmax state.

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:

1. Each lane multiplies a regularly spaced subset of query and key values.
2. simd_sum combines those partial dot products into one score.
3. Apply the scale, optional mask, and causal check.
4. 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

For your checkpoint, keep the short-context model comparison and one representative operator row inside the dispatch range. The checked sweep below shows why the reference guard stops at 256; the complete samples and execution order live in the performance appendix.

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:

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:

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 that fixed workload separate from the short-context acceptance run: it confirms that prefill is unchanged and still routes through Day 3’s matrix-shaped projection path. The performance appendix contains the full context sweep and attribution.

If you want to continue without the custom attention kernel, preserve scaled_dot_product_attention and the model-facing dispatch boundary, then use MLX attention only at that operator seam. Do not replace the course cache, model, or generation loop with --solution mlx.

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 Day 6: SIMD-Matrix Prefill

Day 5 leaves one-token decode on the Day 3 matvec and 128-token prefill on the correctness-first vanilla matrix path for M > 8. The Day 6 starter already contains the cumulative primitive and dispatch, the quantized_matmul_simdgroup_w4a16_g128 Metal shell, the course-owned CooperativeTileLoader/CooperativeBlockMMA boundary, and the logits_to_keep model switch. You complete those surfaces without changing the public quantized-linear API.

Implement the matrix path in four cumulative slices: preserve the M <= 8 matvec dispatch, add the 32×32×32 tile, make device loads contiguous, reuse one group’s scale/bias across its four reduction tiles, then move the final-logit slice before the vocabulary projection. The day gate checks aligned and partial tiles as well as the loader:

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

After that gate passes, run one matched 128-token prefill and one real-shape projection control. The live model must use the tiled path; an isolated kernel result is not enough.

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.

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:

1. loads one 32×32 activation tile into padded threadgroup memory;
2. unpacks and dequantizes one 32×32 weight tile there;
3. lets four SIMD groups reuse both tiles;
4. accumulates four 16×16 quadrants from Metal 8×8 matrix fragments;
5. 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.

Implement the small course-owned boundary in src/extensions/src/cooperative_matrix.h. CooperativeTileLoader assigns one contiguous source chunk to each thread, uses a branch-free full-tile path, and zero-fills the edge-safe path. CooperativeBlockMMA loads and accumulates direct Metal simdgroup_matrix fragments. The required solution does not use Steel BlockLoader or BlockMMA.

This helper does not hide the exercise. Your solution still owns the W4A16 unpacking, dequantization, tile layout, direct matrix-fragment bookkeeping, 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 in src/extensions/src/quantized_matmul.metal and complete the course-owned CooperativeTileLoader TODO in src/extensions/src/cooperative_matrix.h; 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

The comparison should separate two cases. Healthy long-M projections show that the tile itself works; a short, narrow K/V projection with too few 32×32 result tiles exposes the occupancy problem Day 7 addresses. If both cases are slow, the remaining work is still in this tile rather than in reduction partitioning.

At long M, the two-dimensional tile grid is already large. The checked 2,048-row sweep puts the course SIMD path roughly 7–11% above MLX latency for the major projections, while the Split-K successor is neutral. That control does not establish parity with MLX; it establishes that multiplying an already occupied grid 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.

Use one short K/V row and one long occupied-grid row to establish the shape difference. Do not select Split-K merely because projections still occupy most of prefill.

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 show that Split-K does not help once that grid is occupied; it may still expose a gap inside each tile, which belongs to Day 6 rather than a larger partition grid. The reference checkpoint pairs the prefill result with long and short operator controls and the dispatch geometry that motivates Split-K. The exact final-main samples and method live in benchmark_results/task367-final-main/task367-final-main-benchmark-ledger.md.

The complete shape campaign and attribution are in the performance appendix.

If you want to continue without implementing the tiled kernel, preserve the same quantized-linear and logits_to_keep interfaces and substitute MLX only for the matrix-shaped projection. Keep the Day 3 decode matvec and the rest of the course model intact; --solution mlx is a different full-model path.

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 Day 7: Split-K Prefill

Day 6 leaves a reusable course-owned 32×32×32 tile and one visible dispatch point in QuantizedMatmul::eval_gpu. Day 7 keeps that tile unchanged and completes the existing Split-K accumulation and reduction shells in quantized_matmul.metal, plus the shape policy at the same C++ dispatch point. No new public matmul function is needed.

Begin with the unsplit short K projection in Task 1. Then add one partition dimension, choose the partition count from the result-grid occupancy, and reduce the private planes. The focused day gate exercises the eligible K/V shape, a partial output tile, and the exact fallback to Day 6:

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

The short K/V projection is the learner-owned optimization target. Qwen’s narrow outputs do not launch enough independent result tiles to occupy the GPU, so split the reduction dimension only until that grid is large enough. The matched long-row control still trails MLX; it tells us that extra partitions are neutral once the ordinary result grid is occupied, not that the base tile has reached library parity.

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

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:

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

For your checkpoint, keep one short under-filled comparison and one long occupied-grid control. Retain Split-K only when it improves the short projection, preserves one-token decode, and falls back exactly to Day 6 for the long control. The fixed 128-prompt/129-output row is the representative complete-model result; the performance appendix owns the full crossover campaign.

On the checked M4 Pro, the balanced final-main evidence supports this narrower result:

• at M=32, K/V/O/down improve in both execution positions, Q reverses direction, and gate/up are neutral;
• complete-model prefill improves from 537.92 to 599.81 tok/s, or 11.5%, while decode is neutral;
• at M=128, complete-model prefill is 706.50 versus 707.41 tok/s and decode is 66.28 versus 65.83 tok/s;
• at M=2,048, complete-model prefill is 551.48 versus 547.73 tok/s and every projection uses the unsplit policy.

The final 128-token checkpoint reaches 88.2% of full-MLX prefill and 87.0% of full-MLX decode throughput. These values support the fixed-workload stretch goal, not a universal Split-K crossover or 80% claim. See benchmark_results/task367-final-main/task367-final-main-benchmark-ledger.md for every raw sample and the balanced-order drift controls.

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 keeps the same quantized-linear interfaces but deliberately selects MLX quantized projections in its dense model and scheduler factory. Its cache, attention, paging, batching, and scheduling remain course-owned. Paging is evaluated separately on cache writes, direct page reads, attention time, and end-to-end throughput; it does not receive credit for either the Day 7 projection result or the separately measured projection-seam gain.

If you want to continue without implementing Split-K, keep the Day 6 quantized-linear interface and use the unsplit tile. You may also substitute MLX at this one projection seam while preserving the course model. Neither choice is the same as running the separate --solution mlx baseline.

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

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

What We’ll Cover

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

Day 1 introduces that projection seam and batches independent request states. Day 2 splits long prefills so they cannot monopolize the scheduler. Day 3 replaces a growing dense cache with fixed-size pages while retaining a dense-gather compatibility path. Day 4 removes that gather by teaching attention to walk the page table directly with a correctness-first 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.

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 3 Day 1: Continuous Batching

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

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

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

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

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

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

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

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

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

Task 1: Reuse RoPE and Causal Masking for Batched Requests

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

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

Verify multi-offset RoPE and 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: Add the Week 3 Projection Seam

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

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

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

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

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

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

Task 4: Batch Generate

src/tiny_llm/batch.py

First implement Request.try_prefill by prefilling the complete prompt in one call. 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.

Use the supplied scheduler checkpoint for full prefill, admission and slot reuse, EOS/removal, and ordered results:

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

Then run the complete scheduler against the real model:

pdm run batch-main

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

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

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 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/task367-final-main/raw/week3-chunked-prefill-final-main.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. Every row uses the same canonical Week 3 MLX quantized-projection seam and the same course-owned scheduler, dense cache, and attention code; only the prefill budget changes. The JSON stores every prompt token id, the per-request output budget, and their canonical SHA-256 checksum.

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:

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

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 3 Day 3: Paged KV Cache

🚧 This chapter is under review and may change.

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

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

📚 Readings

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

Why the Dense KV Layout Becomes Expensive

Right now, the mental model looks like this:

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

Before attention, the runtime repacks them into:

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

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

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

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

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

The Page Abstraction

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

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

Each layer cache keeps a small page table:

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

That means:

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

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

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

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

Why Fixed-Size Pages Help

The page abstraction gives us two immediate wins:

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

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

Data Structures We Need

1. PagePool

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

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

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

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

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

Implement this abstraction as TinyKvPagedPool.

2. PagedRequestCache

A layer cache for one request should track:

• page_ids
• page_lens
• offset
• page_size

Derived values:

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

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

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

3. Tail-Append Logic

When new K/V arrives for one layer:

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

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

Make write cost proportional to the appended slice

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

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

Complete every learner-extension integration point before rebuilding:

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

Then rebuild:

pdm run build-ext

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

Prefill with Pages

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

chunk = [t0 t1 t2 t3 t4 t5]

One possible layout is:

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

That layer cache’s metadata becomes:

page_ids = [7, 2]
context_len = 6

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

Decode with Pages

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

With paged storage:

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

So if page_size = 4 and context_len = 9:

page_ids = [12, 5, 3]

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

Correctness Checkpoint: Gather Pages for Dense Attention

The cleanest first implementation is paged storage with dense gather.

That means:

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

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

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

How This Maps to tiny-llm

src/tiny_llm/paged_kv_cache.py

Add:

• TinyKvPagedPool
• TinyKvPagedCache

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

The chapter’s execution path is:

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

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

src/tiny_llm/batch.py

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

The scheduler should still:

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

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

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

Design Questions

Before implementing, make sure the following are clear:

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

Task 1: Design PagePool

src/tiny_llm/paged_kv_cache.py

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

Design layer-owned page pools that:

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

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

Task 2: Design PagedRequestCache

src/tiny_llm/paged_kv_cache.py

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

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

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

Task 3: Add a Dense-Gather Compatibility Path

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

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

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

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

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

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

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

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

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

What Paging Changes

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

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

pdm run test --week 3 --day 3

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

🚧 Week 3 Day 4: Direct Paged Attention

🚧 This chapter is under review and may change.

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

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

Paged KV Cache vs Paged Attention

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

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

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

The Metadata a Paged Runtime Needs

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

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

For the current layer being executed:

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

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

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

Why block_table Matters

Suppose one layer cache for request A has:

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

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

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

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

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

That is exactly what block_table and context_lens encode.

The Paged Attention API

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

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

With shapes like:

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

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

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

Prefill Metadata

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

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

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

Decode Metadata

During decode, each active request typically writes one token.

The runtime should be able to:

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

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

Choose a Schedule for Each Query Shape

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

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

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

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

How This Maps to tiny-llm

src/tiny_llm/attention.py

Add a new function:

def paged_attention(...):
    ...

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

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

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

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

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

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

The page pool should therefore expose contiguous physical storage:

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

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

src/tiny_llm/qwen3_week3.py

The attention module should call the paged runtime directly:

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

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

src/tiny_llm/batch.py

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

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

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

Implementation Order

Use this implementation order:

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

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

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

These are the invariants worth checking in tests:

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

Task 1: Add Batch Metadata

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

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

Extend the batch cache and scheduler so they can prepare:

• block_table
• context_lens

for all active requests.

Task 2: Define paged_attention

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

Modify these exact starter functions:

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

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

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

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

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

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

Implement two correctness-first GPU dispatches:

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

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

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

Your solution’s boundary

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

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

Task 3: Dispatch from the Model

src/tiny_llm/qwen3_week3.py

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

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

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

This creates the Day 4 routing policy:

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

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

Task 4: Connect It to Continuous Batching

src/tiny_llm/batch.py

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

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

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

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

Measure the Direct Page Walk

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

Record three operator baselines on the same machine:

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

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

pdm run bench-week3-attention --offline --contexts 128 1024 \
  --page-size 128 --warmup 5 --iterations 60 --repeats 4 \
  --cooldown-seconds 1 \
  --json-output benchmark_results/task367-final-main/raw/week3-attention-final-main.json

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

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

Checkpoint 1: Establish a Correct Direct Path

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

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

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

Checkpoint 2: Design the Decode Schedule

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

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

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

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

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

Checkpoint 3: Evaluate the Serving System

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

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

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

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

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

It compares Week 2 dense batch reconstruction, Week 3 paged storage with the dense-gather compatibility path, and Week 3 direct paged attention. All three course rows use the same MLX quantized-projection seam; they differ in KV representation and attention path. The runner resets page capacity after warmup and reports prefill, output, and decode throughput alongside peak KV bytes, copy volume, page reuse, and tail fragmentation. The direct path’s four-process medians are 672.68 prefill tok/s, 46.36 output tok/s, 105.01 decode tok/s, and 0.618 requests/s. Its synchronized decode calls take 28.97/36.78/63.04 ms at median/p95/max; the completion gaps, which include intervening scheduler and prefill work, are 30.16/222.18/239.49 ms.

These are cumulative system results, not an isolated Day 4 kernel speedup. The ledger at benchmark_results/task367-final-main/task367-final-main-benchmark-ledger.md records the fixed trace, balanced process order, and denominator boundary.

pdm run test --week 3 --day 4

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

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

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

The algorithm 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:

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. Reuse the course-owned CooperativeTileLoader and direct simdgroup_matrix fragments from Week 2; do not import a Steel loader or matrix helper. Your solution owns page translation, the contiguous and cross-page load paths, tile schedule, causal mask, 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:

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

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

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

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

s <= l + context_len - L

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

Task 3: Validate the Page Boundary

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

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

1. compare Day 4 page-walking attention with the readable equation written with mlx.core;
2. compare paged FlashAttention with the Day 4 path;
3. 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. Canonical Week 3 uses MLX quantized projections, but its cache, paged attention, batching, and scheduling remain course-owned. This hybrid course path is not the full-MLX baseline.

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/task367-final-main/raw/week3-serving-final-main.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, all three course rows share the same projection seam:

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

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/task367-final-main/raw/week3-8k-final-main.json

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

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

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

• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

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

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

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

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

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

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

What Is Runnable Today

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

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

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

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

The Nine-Day Progression

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

Prerequisites and Environment

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

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

Work Through One Chapter

For Day N:

1. Read what the final scaffold already declares and which TODO bodies belong to this day. Ignore future modules even though their declarations are visible.

2. Predict the named action, count, range, or stop reason before running the focused scenario when the chapter asks for one.

3. Run the cumulative learner checkpoint:

   pdm run test --week 4 --day N
   

   For Week 4, this command force-refreshes the learner tests for Days 1 through N from the supplied checkpoints, then runs all of them in one pytest invocation. Keep your implementation in src/; do not modify a copied test because the next run replaces it.

4. Implement only the files and relationships named by that chapter.

5. Rerun the checkpoint and inspect the artifact that can falsify your prediction: events, files, receipts, checkpoints, reports, cache offsets, or artifact bytes.

6. After Day 9 is green, run the composed product witness:

   pdm run week4-capstone
   

   Inspect its sorted JSON sections for compaction, both branches, the selected branch, and the externalized artifact range.

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

Read the Metrics as Accounting

Week 4 exposes three kinds of useful counts:

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

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

Week Boundary

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

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.

Day 1: A Validated Agent Loop

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

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

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

The Teaching Boundary

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

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

Files and Public Surface

The repository is a final Day 9 declaration scaffold. Future agent modules and exports are already visible, but their implementation surfaces are not part of Day 1. Most later bodies are TODO stubs; Day 9 also contains one explicitly supplied constructor check. Implement only the following surfaces:

Implement the TODO bodies in these Day 1 starter files:

generate_response() remains part of the public Day 1 surface. It renders the messages with the course tokenizer, decodes at most max_tokens with a fresh cache, stops at EOS, and releases every cache in a finally block.

The supplied Day 1 test checks generate_response() with the course tokenizer and model boundary: exact prompt and thinking offsets, EOS stopping, the token limit, a fresh cache for each call, and cache release on normal and exceptional paths. The pdm run agent CLI still uses its own MLX-LM generation adapter, so a successful live CLI run is not evidence for this helper; the cumulative Day 1 checkpoint is.

Run the cumulative learner checkpoint from the repository root:

pdm run test --week 4 --day 1

The command force-refreshes the supplied Day 1 learner test before running it. Before you implement the TODOs, the implementation-dependent cases across ten task groups are expected to fail. No model download is required.

Course maintainers can check the supplied implementation without copying the learner test:

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

Task 1: Start the Conversation Deliberately

initial_messages(task, system_prompt) creates the first two messages:

[
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": task},
]

Reject an empty or whitespace-only task. build_system_prompt(workspace) describes only the actions enabled for this run. The prompt is guidance, not enforcement: the protocol and workspace boundary must still reject anything the policy does not allow.

Task 2: Execute One Tool and Finish

run_agent(task, generate, workspace, limits=None) starts from those messages. The test injects a generate callable that returns predetermined strings, so the control flow stays deterministic.

For a valid tool action, call workspace.execute(action), record the action and result, and append both the assistant response and a user observation. When a later response is a valid FinalAction, return a completed AgentRun with the final text.

Task 3: Validate One Structured Decision

A model response is exactly one JSON object. It is either a final answer:

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

or one tool request:

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

parse_action() rejects malformed JSON, non-object values, blank final text, unknown or disabled tools, missing fields, unexpected fields, and fields with the wrong shape. Do not ignore trailing or extra data.

TOOL_FIELDS names the cumulative vocabulary: list_files, read_file, write_file, edit_file, and run_command. Day 1 implements none of those effects. Its fake workspace enables only read_file, which is enough to prove that availability is checked before dispatch.

Malformed or unavailable actions become ordinary error: observations. The model can see the failure and choose another response instead of crashing the Python loop.

Task 4: Stop at the Step Budget

AgentLimits.max_steps bounds how many model decisions one run may attempt. When the loop consumes that budget without a valid final answer, return an incomplete run whose reason is step_limit. The events show exactly how the budget was spent.

Task 5: Return an Inspectable Run

Every interaction becomes an AgentEvent with the step number, raw response, parsed action when one exists, and result or validation error. AgentRun records the completion flag, stop reason, optional final answer, and immutable event tuple.

A run marked completed means only that the model returned a valid final action. Later days add receipts and outcome evaluation; Day 1 keeps the trace small and in memory.

Task 6: Recover from Invalid JSON

After an invalid response, append the raw assistant response and its exact validation error as the next user observation. Reset the identical-action counter, then let the model try again while the invalid-action budget remains.

The focused case sends invalid JSON followed by a valid final response. The first event must retain the recoverable error, and the second must complete the run.

Task 7: Stop Repeated Actions

Serialize each parsed tool name and normalized argument object into a stable signature. Count consecutive identical requests and stop with repeated_action_limit when the count exceeds the configured budget.

This guard matters even when a tool succeeds: repeating the same request can consume the whole run without adding new information.

Task 8: Preserve the Exact Observation

The next model call must receive the complete tool result, not only a marker:

{
    "role": "user",
    "content": "Tool result:\nREADME contents",
}

The normal guard fails if that payload is changed or dropped. It also proves that a known-but-disabled tool such as write_file becomes an error observation and never reaches the fake workspace.

Task 9: Make Every Limit Fail Closed

Require positive values for max_steps, max_context_chars, max_invalid_actions, and max_identical_actions. Zero or negative budgets would disable the intended stopping guarantee and must be rejected.

Before each model call, bound the total message-content characters and stop with context_limit when it is too large. Count invalid actions and stop with invalid_action_limit when that budget is exhausted. Together with the step and repeated-action limits, every run has an explicit terminal reason.

Checkpoint

When Day 1 is green, inspect the focused test rather than only its final pass: confirm the initial system/user pair, one dispatched read_file, the exact observation in the next model input, the completed final event, and each budgeted stop reason. Revisit your opening prediction: a disabled or malformed request must not reach the workspace, and the exact error must remain visible to the following model turn.

This checkpoint proves the scripted protocol and loop, plus the focused generate_response() model/cache boundary. It does not make the separate real-model CLI exercise that helper.

You now have a validated, bounded model → action → observation loop. Continue with Day 2: Inspect a Workspace to replace the fake tool boundary with real contained directory listing and UTF-8 file reads.

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.

🚧 Day 2: Inspect a Workspace

Day 1 built a loop around a fake workspace. The model could request a tool, the loop could return an observation, and the model could finish—but no tool looked at a real project.

Day 2 adds the smallest useful workspace: two read-only tools that let the model list a directory and read a text file. This makes the complete cycle visible without mixing in approval, edits, commands, or durable receipts.

model requests list_files
          |
          v
 Workspace checks the path
          |
          v
 loop returns the listing
          |
          v
 model requests read_file
          |
          v
 loop returns the file text
          |
          v
 model returns a final answer

The Teaching Boundary

Use these tools with a trusted operator and a disposable local repository. Path.resolve() containment and symlink refusal prevent ordinary path mistakes, but they are not a sandbox or a defense against a hostile filesystem. The tool result also becomes model input, so do not point the workspace at files you would not send to the model.

Day 3 adds operator-approved edits, one allowlisted validation command, and a simple receipt log. Keeping those effects out of Day 2 lets you first see the read-only agent loop clearly.

Files and Public Surface

The published starter is the cumulative Day 9 declaration scaffold. Receipt, checkpoint, compaction, and later modules are therefore visible already, but their implementation work belongs to later chapters. Day 2 owns only the workspace relationships below.

Implement the TODO bodies in:

The Day 2-owned prefix of ToolPolicy has three fields, in order:

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

Day 2 implements this prefix of Workspace:

• available_tools, always {"list_files", "read_file"};
• resolve_path(raw, must_exist=True);
• list_files(raw=".");
• read_file(raw); and
• execute(action).

The starter declarations are the contract. Do not implement or depend on the visible write, command, approval, receipt, checkpoint, compaction, steering, evaluation, branching, or evidence-retrieval declarations yet.

Task 1: Normalize One Workspace Root

ToolPolicy receives the directory the agent may inspect. Convert it to a resolved Path, require an existing directory, reject a symlink root, and require both numeric limits to be positive.

Normalizing once keeps every later check relative to the same root:

policy = ToolPolicy(Path("demo-project"))
workspace = Workspace(policy)

Task 2: Resolve Ordinary Relative Paths

resolve_path() accepts a non-empty relative path and returns its resolved location under policy.root. Reject:

• absolute paths and .. traversal;
• symlinked path components;
• .git, .env, .ssh, .aws, and common credential or private-key names;
• missing paths when must_exist=True; and
• any resolved path outside the workspace.

Return a recoverable AgentError for these cases. The agent loop can turn that error into an observation instead of crashing.

This is deliberately an ordinary local-repository boundary. A different process can still race filesystem checks; defending against a hostile filesystem is outside this course checkpoint.

Task 3: List One Directory

list_files() lists direct children in sorted order. Emit one line per visible regular file or directory:

file README.md
dir src

Omit protected names, symlinks, and special files. Stop after max_list_entries lines. Return (empty directory) when no visible entry remains.

The tool is intentionally not recursive. The model can request another list_files action for a directory it wants to inspect.

Task 4: Read One Text File

read_file() accepts a visible regular file no larger than max_file_bytes. Read its bytes and decode UTF-8. Reject directories, oversized files, invalid UTF-8, protected paths, and symlinks with AgentError.

The size bound keeps a single observation from consuming the whole context window. Later checkpoints can add more deliberate context selection; Day 2 only needs one obvious limit.

Task 5: Turn Failures into Observations

execute() dispatches a parsed ToolAction to list_files or read_file. It returns successful text directly. A recoverable failure becomes a string beginning with error: so run_agent() can append it to the conversation and let the model choose another action.

A known future tool such as write_file is still disabled on Day 2. The system prompt describes only workspace.available_tools, and parse_action() checks that enabled set before dispatch.

Task 6: Run the Read-Only Cycle

The checkpoint uses a scripted model so it is deterministic and does not load weights:

responses = iter([
    '{"tool":"list_files"}',
    '{"tool":"read_file","path":"README.md"}',
    '{"final":"README says hello"}',
])

result = run_agent(
    "inspect the project",
    lambda messages: next(responses),
    workspace,
)

Inspect result.events: the first two events contain the parsed tool actions and exact observations, and the third contains the final answer. The same tool result appears in the next model input as Tool result:\n....

Run the Day 2 Checkpoint

From the repository root, run the cumulative learner checkpoint:

pdm run test --week 4 --day 2

Before you implement the TODOs, the implementation-dependent cases across eight task groups are expected to fail. The command force-refreshes the supplied learner tests for Days 1 and 2, then runs both; completed Day 1 behavior should remain green. During course development, check the supplied implementation with the day-local maintainer command:

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

The course-code guard compares the starter and reference public signatures, dataclass fields, package exports, and solution-free method bodies.

Explore with a Real Model

The scripted checkpoint above is the reproducible mechanics proof. This separate manual run is exploratory: a real model chooses the tools, so its wording and exact tool order may vary. Use only a disposable directory that contains no secrets or private source.

The CLI defaults to qwen3-4b (Qwen/Qwen3-4B-MLX-4bit), whose cached weights use about 2 GiB; use that lower-resource option when needed. The recorded exploratory run below used mlx-community/Qwen3-30B-A3B-4bit, whose cached weights use about 16 GiB and require sufficient Apple unified memory. Model behavior and tool order vary with either choice. Both use the local MLX model path already used by this repository. You need macOS on Apple Silicon and the installed MLX dependencies. The first run downloads the selected weights from Hugging Face when they are not cached, so it also needs network access and the corresponding free disk space. If MLX is unavailable or the weights cannot be loaded, the command exits before the agent calls a workspace tool; it does not substitute scripted output.

This CLI uses a supplied MLX-LM generation adapter. It exercises your Day 1 loop and Day 2 workspace, but it does not call the learner-owned generate_response() helper or prove the Weeks 1–3 course model/cache path.

Create a tiny read-only workspace, then give the model one goal:

INSPECT_ROOT="$(mktemp -d)"
mkdir "$INSPECT_ROOT/src"
printf '%s\n' '# Pocket Weather' 'A tiny terminal forecast project.' > "$INSPECT_ROOT/README.md"
printf '%s\n' 'def forecast(city):' '    return f"Sunny in {city}"' > "$INSPECT_ROOT/src/weather.py"

pdm run agent -- --model mlx-community/Qwen3-30B-A3B-4bit --root "$INSPECT_ROOT" \
  "Inspect this workspace and explain its purpose and the behavior implemented in its source file. Use the available workspace tools to gather evidence from both the project overview and the source file. Your first response must be one tool request, every response must contain exactly one JSON object, and you must not finish until you have read the source file."

Watch the printed goal, model responses, parsed actions, and tool observations. There is no predefined action list: the model decides what to list and read before returning a final answer. The policy is read-only, so no approval prompt or receipt is expected. Compare the final answer with the actual disposable files:

find "$INSPECT_ROOT" -type f -print
cat "$INSPECT_ROOT/README.md" "$INSPECT_ROOT/src/weather.py"

Checkpoint

Confirm that the deterministic run listed a directory, read exact UTF-8 file contents, returned recoverable path errors to the model, and completed without enabling any effect tool.

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

After all nine days are complete, the supplied deterministic capstone composes this read-only inspection with checkpointing, compaction, steering, evaluation, branching, and bounded evidence. That orchestration is a separate week4-capstone command, not part of the real-model agent CLI.

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.

🚧 Day 3: Edit, Validate, and Record

Day 2 gave the model two read-only tools. It could inspect a disposable project and explain what it found, but it could not fix anything. Day 3 completes one small coding cycle:

read file -> propose exact edit -> operator approves -> recheck bytes
    -> replace file -> record receipt -> run focused check -> record receipt
    -> final answer

The important idea is not broad autonomy. It is an explicit boundary between a model proposal and a local side effect.

Before running the focused cycle, predict three facts: how many approvals and receipts exist after a denied effect, whether stale-read detection changes any file bytes, and which evidence remains after a validation command returns a nonzero status. The workspace, receipt log, and event trace let you check those answers independently of the model’s final sentence.

The Teaching Boundary

Use this checkpoint with one trusted operator, one Python process, and a disposable repository that contains no secrets. Its path checks prevent ordinary mistakes, but they are not a sandbox or a defense against a hostile filesystem. The command tool is not a process jail: an allowed program can access anything the host process can access, spawn children, or use the network.

The receipt file is a simple append-only JSONL teaching record. It detects edited receipt bytes when reopened and handles a repeated call ID in the same process. It is not a transaction log, an fsync protocol, a multi-writer store, or proof that an interrupted effect did or did not happen. Day 3 deliberately stops at this receipt boundary.

Files and Public Surface

Later checkpoint, compaction, steering, evaluation, branching, and evidence modules are already declared in the cumulative Day 9 scaffold. Their later-day implementation work is out of scope. Day 3 owns only the following Day 1–3 surfaces:

Implement the TODO bodies in these starter files:

ToolPolicy keeps its first three Day 2 fields and adds:

allow_writes: bool = False
allowed_commands: tuple[tuple[str, ...], ...] = ()
max_write_bytes: int = 64 * 1024
command_timeout_seconds: float = 30.0

Writes stay disabled unless allow_writes=True. Commands stay disabled unless their complete argument tuple appears in allowed_commands. There is no shell string or prefix match.

Workspace(policy, confirm_tool=None) creates an in-memory receipt store. Pass a ReceiptStore(path) as the third argument when you want JSONL output. The workspace extends the Day 2 methods with write_file, edit_file, run_command, and modified_files. execute(action, tool_call_id=None) is the model-facing gate. Direct tool methods are useful for focused unit tests; execute performs the approval and receipt steps.

Task 1: Authorize Tools Explicitly

Build available_tools from the policy. Listing and reading are always present. Add both file mutation tools only when writes are enabled, and add run_command only when at least one exact command is configured. Validate all size limits, the timeout, the boolean flag, and every command part.

This configuration permits one focused check:

validation = ("python", "-m", "pytest", "tests/test_math.py", "-q")
policy = ToolPolicy(
    Path("demo-project"),
    allow_writes=True,
    allowed_commands=(validation,),
)

Task 2: Read Before Changing Existing Bytes

Keep Day 2’s path and read rules. When read_file succeeds, remember the SHA-256 digest of the bytes that were returned. Replacing or editing an existing file requires that observation. A new file does not have old bytes to inspect, but its parent directory must already exist.

For edit_file, require a non-empty old string that occurs exactly once. Compute the proposed bytes in memory and enforce max_write_bytes before asking for approval. Whole-file write_file replacements also require a prior read.

Task 3: Ask Once, Default No, Then Recheck

execute preflights the complete action before calling confirm_tool. Missing callbacks, False, and every value other than the boolean True deny the effect. A terminal program can provide a small default-No callback:

def confirm(action):
    answer = input(f"Approve {action.tool} {action.arguments}? [y/N] ")
    return answer.strip().lower() in {"y", "yes"}

After approval, write_file or edit_file reads the destination again and compares its digest with the earlier observation. If another actor changed the bytes while the operator was deciding, return error: file changed since it was read and do not overwrite them.

Task 4: Replace Through the Same Directory

Write the proposed bytes to a temporary file in the destination’s parent, close it, and call os.replace(temporary, destination). Clean up a leftover temporary file after an error. This avoids presenting a partially written destination to ordinary readers.

This small pattern is atomic at the replacement step, but it is not a durable journal and does not close the check-to-replace race against a hostile actor.

Task 5: Run One Exact Validation Command

run_command(argv) accepts a non-empty list of strings only when its tuple is exactly allowlisted. Call subprocess.run without a shell, with the workspace root as cwd, captured text output, and the configured timeout. Bound combined stdout and stderr so one observation cannot consume the whole context window.

Return one observation with the status and captured output:

status: 0
output:
1 passed

A nonzero status and a timeout are ordinary validation results the model can inspect. They are not Python exceptions and do not prove the final answer is correct.

Task 6: Record Simple Effect Receipts

An EffectReceipt has these fields, in order:

tool_call_id: str
tool: str
arguments: dict[str, Any]
exit_state: str
result: str
changed_artifacts: tuple[str, ...] = ()

Its receipt_id is the SHA-256 digest of the canonical JSON payload. A successful write or edit records exactly one normalized workspace-relative artifact. A validation receipt records the exact argv, status and captured output, with no changed artifacts. ReceiptStore(path) loads and verifies an existing JSONL file; ReceiptStore() remains in memory.

The store maps one tool_call_id to one receipt. Repeating the same call ID and action returns the existing result without running the effect again. Reusing the ID for another action is an error. This is proportional duplicate handling for one process, not distributed exactly-once execution.

Task 7: Run the Complete Scripted Cycle

The test uses scripted model responses, so it needs no model weights:

responses = iter([
    '{"tool":"read_file","path":"app.py"}',
    '{"tool":"edit_file","path":"app.py","old":"1","new":"2"}',
    '{"tool":"run_command","argv":["python","-m","pytest","tests/test_math.py","-q"]}',
    '{"final":"changed and validated app.py"}',
])

store = ReceiptStore(Path("demo-project/.agent-receipts.jsonl"))
workspace = Workspace(policy, confirm, store)
result = run_agent("fix app.py", lambda _messages: next(responses), workspace)

Inspect result.events, workspace.modified_files, and the two receipts. The edit receipt names app.py; the validation receipt has an empty artifact tuple. The final answer is still a model statement, so the validation status in the trace is the evidence that matters.

Run the Day 3 Checkpoint

From the repository root, run the cumulative learner checkpoint:

pdm run test --week 4 --day 3

Before you implement the TODOs, the Day 3 cases are expected to fail because the new starter methods return None. The command force-refreshes and runs the supplied learner tests for Days 1–3 together. Keep the Day 3 failures until you solve each task; do not import tiny_llm_ref from the starter.

Course maintainers can check the supplied implementation without copying the learner test:

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

The course-code guard checks exact public signatures, dataclass fields, package exports, and TODO-only starter bodies across the final declaration scaffold.

Explore the Full Cycle with a Real Model

Keep the scripted checkpoint as the deterministic proof. This manual exercise lets a real model plan the same Day 3 cycle from one natural-language goal. Its wording and tool order can vary, and completion is not an automated test. Use a fresh disposable directory with no secrets.

The CLI default is qwen3-4b (Qwen/Qwen3-4B-MLX-4bit), whose cached weights use about 2 GiB; use that lower-resource option when needed. The recorded exploratory run below used mlx-community/Qwen3-30B-A3B-4bit, whose cached weights use about 16 GiB and require sufficient Apple unified memory. Model behavior and tool order vary with either choice. Both require macOS on Apple Silicon and the installed MLX dependencies. An uncached first run downloads the selected weights from Hugging Face and needs network access plus the corresponding free disk space. If MLX, network access, disk space, unified memory, or the weights are unavailable, model loading fails before any tool call; do not treat a scripted checkpoint as evidence that this live run occurred.

As on Day 2, this CLI’s supplied MLX-LM adapter bypasses the learner-owned generate_response() helper. It does exercise your loop, workspace, approval, effect, and receipt paths.

Pre-create the workspace, an existing file, and one focused validation fixture:

EDIT_ROOT="$(mktemp -d)"
printf '%s\n' \
  'def greeting(name):' \
  '    return f"Hello, {name}!"' > "$EDIT_ROOT/app.py"
cat > "$EDIT_ROOT/validate.py" <<'PY'
from pathlib import Path
from app import greeting

assert greeting("Ada") == "Welcome, Ada!"
assert Path("NOTES.md").read_text() == "Greeting now says Welcome.\n"
print("validation passed")
PY

Now give the real model one goal. --allow-command names the only command it may run, including its exact arguments. Each proposed write, edit, or command still pauses at a default-No learner approval prompt:

pdm run agent -- --model mlx-community/Qwen3-30B-A3B-4bit --root "$EDIT_ROOT" \
  --allow-writes \
  --allow-command "python validate.py" \
  --receipt-log .agent-receipts.jsonl \
  "Inspect the workspace. Create NOTES.md containing exactly 'Greeting now says Welcome.' followed by a newline, precisely change app.py so greeting says Welcome instead of Hello, run the allowed validation, react to its evidence, and finish with a brief summary. Use workspace tools to gather evidence before any effect. Your first response must be one tool request, every response must contain exactly one JSON object, and do not finish until the requested files and validation evidence have been inspected."

Read every approval payload before answering y. The live trace shows the model response, parsed action, tool observation, validation status/output, and final answer. Afterward, inspect the changed bytes and the durable effect receipts rather than trusting the final prose alone:

cat "$EDIT_ROOT/app.py" "$EDIT_ROOT/NOTES.md"
cat "$EDIT_ROOT/.agent-receipts.jsonl"

The receipt records preserve the approved write_file, edit_file, and run_command arguments and outcomes. The status and captured output in the validation receipt are the evidence to compare with the changed bytes.

Checkpoint

You now have a small end-to-end coding loop: inspect real bytes, propose one precise change, pause for a trusted operator, reject stale observations, replace the file, validate with one exact command, and retain simple evidence of both effects. Continue with Day 4: Checkpoint and Resume to save one complete observation boundary and restore it through a fresh scripted model without turning Day 3 into production infrastructure.

Day 3 is the last mechanism exposed by the real-model CLI. Days 4–9 use deterministic library checkpoints, and the separate supplied Week 4 capstone connects all nine completed days into one runnable deterministic product path.

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.

Day 4: Checkpoint and Resume

Days 1 through 3 build one uninterrupted coding-agent run. The model proposes an action, the harness executes it, and the result becomes the next model observation. But if the process stops, a new model object does not know which conversation prefix it had already processed.

Day 4 makes one boundary visible: save the conversation and the small model snapshot immediately after a complete tool observation, then restore both into a fresh scripted model and continue at the next response. The completed tools stay completed; resume starts after their observations instead of replaying them.

This is a teaching checkpoint for one process and the course’s fake model. It is not a session tree, rewind feature, persistent KV store, transaction log, or exactly-once effect system.

Files and Public Surface

The final scaffold already declares compaction, steering, evaluation, branching, and evidence APIs. Leave those future TODO bodies alone. Day 4 owns only the checkpoint model and the two loop entry points below:

Implement the TODO-only surfaces in:

Run the cumulative learner checkpoint from the repository root:

pdm run test --week 4 --day 4

Before you implement the TODOs, all seven Day 4 tasks are expected to fail. The test uses a scripted model, fake cache metadata, a temporary workspace, and one exact Python validation command. It does not load model weights. The command force-refreshes and runs the supplied learner tests for Days 1–4 together.

Course maintainers can check the supplied implementation without copying the learner test:

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

One Safe Loop Boundary

A checkpoint is saved only after the harness has appended both halves of a tool interaction:

assistant: {"tool":"edit_file", ...}
user:      Tool result:\nedited app.py
                                      ^ checkpoint here

Saving before the observation would leave the restored model unable to tell whether the tool ran. Day 4 therefore counts completed tool calls and saves at the boundary after _append_tool_result(...) has produced the next complete conversation.

The checkpoint stores the semantic messages, not the AgentEvent history. Day 3 receipts remain separate evidence about the edit or command. They are not copied into the checkpoint and they do not become a replay controller.

Task 1: Represent the Fake Model Snapshot

ModelCheckpoint contains four fields, in order:

conversation_position: int
response_index: int
cached_token_ids: tuple[int, ...]
layer_offsets: tuple[int, ...]

conversation_position is the number of semantic messages at the saved boundary. response_index tells the scripted model which response comes next. cached_token_ids represents the prompt prefix in the fake cache, and every layer_offsets entry must equal its length.

Reject negative positions, invalid token IDs, a missing layer snapshot, or offsets that disagree with the cached prefix. These checks make the fake model state internally coherent without introducing a production cache format.

Task 2: Bind Conversation and Model State

AgentCheckpoint contains:

checkpoint_id: str
task: str
messages: tuple[tuple[str, str], ...]
model: ModelCheckpoint

create_checkpoint(task, messages, model) copies each mutable message into an immutable (role, content) pair. It computes checkpoint_id as the SHA-256 of canonical JSON containing the task, messages, and model fields.

AgentCheckpoint.validate() checks the ordinary resume contract:

• the task and messages are structurally valid;
• the model’s conversation position equals the saved message count;
• the checkpoint ID still matches the content.

This identity check catches an accidental mismatch. It is not a hostile-tamper or authentication scheme.

Task 3: Stop After a Complete Observation

run_to_checkpoint(task, generate, workspace, after_tool_calls=1, limits=None) starts with the same prompt and validation rules as run_agent. It counts a tool call only after execution and observation append. At the requested count, it calls:

model_state = generate.save_checkpoint(messages)

and returns an AgentCheckpoint.

The generator must return a ModelCheckpoint. A missing checkpoint method, a non-positive tool-call count, or a run that finishes before the boundary is a clear error instead of a partial checkpoint.

Task 4: Restore a Fresh Model

resume_agent(checkpoint, fresh_generate, workspace, limits=None) validates the checkpoint, calls:

fresh_generate.restore_checkpoint(checkpoint.model)

rebuilds the semantic message list, and enters the normal loop. The fresh model therefore sees the exact conversation prefix and produces the response at the saved response_index.

No old tool action is submitted again. The first model call after restore sees the already-recorded tool result and chooses the next action or final answer.

Task 5: Make the Fake Cache Visible

The test’s FakeCheckpointModel turns each message content length into one fake token ID. save_checkpoint(messages) records those token IDs, two matching layer offsets, and the next scripted-response index. A new model object restores that state and asserts that its first resumed input has the same token prefix and conversation position.

This deliberately small representation exposes the inference/harness integration: the harness owns semantic conversation state, while the model owns the cache snapshot that accelerates exactly that state.

Task 6: Resume Without Replaying Effects

The end-to-end test scripts:

read app.py
edit app.py
run the exact validation command
checkpoint after the validation observation
construct a fresh scripted model
resume to the final answer

Before resume, the Day 3 store already contains the edit receipt and validation receipt. After resume, the approval log and receipt count are unchanged: neither completed effect ran twice. The checkpoint contains no receipt IDs and makes no exactly-once claim; it simply resumes after the conversation already says those effects completed.

Task 7: Keep the Boundary Small

Day 4 implements only checkpoint.py and two loop entry points. Future modules are already declared in the final scaffold, but do not implement or depend on them here. Do not add session IDs, parent pointers, rewind methods, steering queues, disk cache files, or another checkpoint representation. If the in-memory checkpoint is lost, start a new run.

Checkpoint

You can now stop after a complete tool observation and continue with a fresh scripted model from the same conversation and fake-cache position. Inspect the checkpoint’s messages and model fields, then confirm that the pre-checkpoint edit and command remain single completed effects.

At this point, predict what resume is allowed to do: the next model response may request a new action, but no action already represented by the saved observation should run again. The approval and receipt counts in the focused scenario are the falsifying evidence.

Continue with Day 5: Compact Completed Work to derive a smaller model-visible transcript while keeping the exact effect receipts. The Day 5 checkpoint receives a transcript and receipts directly rather than changing this resume path. After Day 9, the supplied Week 4 capstone composes the two mechanisms through its deterministic orchestration.

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.

Day 5: Compact Completed Work

🚧 Early-review WIP: This chapter is public for early review and may change. Use a disposable workspace when running the agent or enabling writes or commands.

Every tool call adds two model-visible messages: the assistant’s action and the tool observation. A long validation log can eventually crowd out the task and the useful recent steps. Deleting old messages saves space, but it also deletes the evidence behind claims such as “validation passed.”

Day 5 makes one boundary visible: replace an older, completed effect with a small deterministic evidence record while keeping its full EffectReceipt unchanged. The model receives fewer tokens and can continue from that view; the harness still retains the exact action and result.

The Starter Surface

The cumulative Day 9 scaffold already contains later steering, evaluation, branching, and evidence declarations. Leave those future TODO bodies alone. Day 5 owns one small module:

Run the cumulative learner checkpoint:

pdm run test --week 4 --day 5

Use this command for the supplied implementation:

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

Before you implement the TODOs, all six Day 5 tasks are expected to fail. The command force-refreshes and runs the supplied learner tests for Days 1–5 together.

Start From the Existing Transcript

The input is the same list of role/content messages that run_agent gives the model. A completed effect has this shape:

[
    {"role": "assistant", "content": '{"tool":"run_command",...}'},
    {"role": "user", "content": "Tool result:\nstatus: 0\n..."},
]

The compactor does not invent a second event log. It receives this transcript plus the Day 3 EffectReceipt values already produced by the workspace.

Before implementing it, predict which older interaction in the focused fixture is eligible to compact, which recent interaction must stay verbatim, and whether saved_tokens must be positive. The returned messages, exact counter values, and unchanged receipts let you falsify that prediction.

Task 1: Require Exact Receipt Evidence

A pair is eligible only when all three facts match one supplied receipt:

1. the parsed tool name;
2. the normalized argument object; and
3. the complete observation text.

If there is no receipt, or if any of those fields differs, leave both messages verbatim. This deliberately excludes Day 2 reads and listings: the current course receipts effects, not every observation. Day 5 must not pretend that an unreceipted result is durable evidence.

Task 2: Keep a Small, Honest Record

Keep the small assistant action, but replace its large tool-observation message with one bounded evidence record that contains:

• the tool and its normalized arguments;
• exit_state and changed_artifacts;
• a bounded prefix of the result; and
• the content-addressed receipt_id.

For example:

Completed tool interaction (compacted evidence):
{"arguments":{"argv":["python","validate.py"]},
 "changed_artifacts":[],"exit_state":"ok",
 "receipt_id":"...","result_preview":"status: 0...",
 "tool":"run_command"}

This is not a model-written summary. It is a deterministic rendering of fields the harness already verified. The bounded preview helps the model explain what happened; receipt.result still contains the full observation.

Task 3: Retain a Recent Tail

keep_recent=1 leaves the newest eligible effect as its original two messages. Older matching effects may compact. The recent tail keeps the next decision grounded in the exact latest interaction without requiring a complicated semantic policy.

The count refers only to receipted effect interactions. Unreceipted reads stay verbatim regardless of this setting.

Task 4: Measure the Real Model Input

The compactor accepts count_tokens(messages) instead of estimating tokens from characters. For the real model, use the same tokenizer and chat template as generation:

def count_tokens(messages):
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    return len(tokenizer.encode(prompt, add_special_tokens=False))

Build the proposed compact view, count it again, and accept it only when the exact counter decreases. CompactionResult reports tokens_before, tokens_after, and saved_tokens. The focused tests inject their own deterministic character counter, which sums the message-content lengths. It is distinct from Day 4’s cached-prefix counter and does not load or download a model.

Task 5: Preserve the Source of Truth

Never mutate the caller’s message list or any receipt. Return copied messages inside CompactionResult. Running the compactor again over its own output is a no-op because a compact evidence record is not a tool action followed by a tool result.

This distinction matters:

original transcript + full receipts   durable evidence owned by the harness
                  |
                  v
       compacted message view          temporary input for the model

If the compact view is lost, derive it again. Do not treat it as a replacement for receipts or checkpoints.

Task 6: Continue With the Existing Model Boundary

CompactionResult.messages contains ordinary role/content mappings. Pass a copied list to the same generation callable used by the loop, then validate the response through the existing protocol:

view = compact_completed_interactions(
    messages,
    receipts,
    count_tokens,
    keep_recent=1,
)
response = generate([dict(message) for message in view.messages])
action = parse_action(response, workspace.available_tools)

The Day 5 test makes the scripted model return a final answer after it sees the compact validation evidence. Nothing about action parsing, tool approval, or workspace execution changes.

Limits of This Teaching Compactor

This checkpoint intentionally does not add semantic-perfect summarization, automatic threshold scheduling inside run_agent, persistent compact views, receipt lookup by summary text, K/V cache editing, session trees, rewind, steering, or exactly-once execution. It compacts only completed effects backed by the receipts the caller supplies.

Checkpoint

You can now make an older validation interaction visibly smaller, inspect the receipt that retains its complete result, and feed the compact view to the next model call. Continue with Day 6: Inspect and Steer a Paused Agent to inspect one checkpoint, add one visible operator message, and resume without replaying completed work.

The Day 6 path resumes the Day 4 transcript; it does not consume CompactionResult.messages. After Day 9, the supplied capstone carries this compacted view into the later control path. Day 5’s token counts and the capstone’s combined report prove only deterministic transcript accounting, not latency, throughput, quality, or memory-capacity improvement.

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.

Day 6: Inspect and Steer a Paused Agent

🚧 Early-review WIP: This chapter is public for early review and may change. Use a disposable workspace when running the agent or enabling writes or commands.

Day 4 gave the harness a safe pause after a complete tool observation. Day 5 showed how to make older completed evidence smaller for the model without discarding the full receipts. Those boundaries also create a useful moment for an operator: inspect what the paused run has actually recorded, add one new instruction, and let a fresh model continue.

The focused Day 6 scenario resumes the original Day 4 transcript rather than CompactionResult.messages. After Day 9, the supplied deterministic capstone connects the compacted view to this later control path.

Day 6 implements only that interaction. It does not inspect hidden reasoning or guess the model’s plan. It reports public facts from the checkpoint, appends one ordinary steering message, and resumes through the existing validated loop.

The Starter Surface

Later evaluation, branching, and bounded-evidence declarations are already visible in the final Day 9 scaffold. Leave those TODO bodies alone. Day 6 owns one module:

Run the cumulative learner checkpoint:

pdm run test --week 4 --day 6

Use this command for the supplied implementation:

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

Before you implement the TODOs, all six Day 6 tasks are expected to fail. The command force-refreshes and runs the supplied learner tests for Days 1–6 together.

Start at a Safe Pause

run_to_checkpoint(...) saves after the assistant action and its tool result are both present:

original task
    ...
assistant tool action
tool observation
                  ^ inspect and steer here

This is deliberately not mid-token or mid-tool steering. No process is running in the background. The workspace is quiescent, and the checkpoint binds the original task, the complete message prefix, and the fake-model cache metadata.

Task 1: Derive a Public Status

Implement:

inspect_checkpoint(checkpoint, evidence_chars=160) -> AgentStatus

Validate the checkpoint, then require its final two messages to be a parsed assistant tool action followed by a Tool result:\n... user observation. Return four public facts:

AgentStatus(
    task="fix app.py and validate",
    last_action='{"new":"2","old":"1","path":"app.py","tool":"edit_file"}',
    last_evidence="edited app.py",
    next_step="resume the model after the completed edit_file observation",
)

last_evidence is a bounded prefix of the recorded observation. When it is too long, reserve the final character for an ellipsis so its length never exceeds evidence_chars.

The status function must not receive or call a model or workspace. It reads the validated immutable checkpoint and returns a new frozen value.

Task 2: State Only What the Harness Knows

The checkpoint knows the current task and the last completed action/result. It does not know what the model will decide next. Therefore next_step is a deterministic harness boundary:

resume the model after the completed edit_file observation

Do not replace this with a semantic guess such as “update the tests next.” That would present invented intent as agent state. A richer plan would need its own explicit, model-visible artifact; Day 6 does not add one.

Reject an incomplete boundary rather than producing a misleading card. A checkpoint ending with an ordinary user message, malformed assistant text, or anything other than a complete action/observation pair is not inspectable by this API.

Task 3: Append One Steering Message

Implement:

resume_with_steering(
    checkpoint,
    steering,
    fresh_generate,
    workspace,
    limits=None,
) -> AgentRun

Reject an empty or whitespace-only instruction. For a valid instruction, restore the checkpoint into the fresh generator and derive a mutable copy of the saved transcript. Append exactly one ordinary user message:

{
    "role": "user",
    "content": "Operator steering:\nvalidate before answering",
}

Then pass that list to the existing bounded loop. Do not add a separate queue, control channel, hidden prompt, or special protocol action.

Task 4: Keep the Message in Stable Order

The steering message belongs immediately after the saved checkpoint prefix. If the resumed model chooses another tool, the existing loop appends that assistant action and its observation after the steering:

saved task and evidence
Operator steering: validate before answering
assistant run_command action
validation observation
assistant final answer

Append the steering message once. Because the loop carries its message list forward, every later model call sees the same single message at the same position. Reinserting it on every call would duplicate the instruction and change the conversation.

Task 5: Continue Without Replaying Effects

The focused scenario starts with the original task “fix app.py and validate.” The first model reads and edits app.py, and the harness checkpoints after the complete edit observation. At that point:

• app.py already contains the new value;
• the edit approval happened once;
• the receipt store contains the edit receipt; and
• inspection reports the saved edit evidence without touching the workspace.

The operator adds “validate before answering.” A fresh scripted model restores the saved prefix, sees the original task, edit evidence, and steering message, then runs the exact allowed validation command and returns its final answer. The test proves the edit approval remains single, validation executes once, and the edit and command retain separate receipts.

Steering changes the next model input. It does not undo, replay, or rewrite the completed prefix.

Task 6: Keep the Boundary Small

Fail clearly when the checkpoint identity is invalid, the generator cannot restore the saved fake-model state, the evidence limit is not a positive integer, or steering is blank. Reuse AgentError, AgentCheckpoint, AgentLimits, AgentRun, and the existing loop rather than creating parallel versions.

The Day 6 module does not add concurrent interruption, mid-token control, a background worker or status server, a durable steering queue, session trees, branch/rewind, exactly-once reconciliation, or an evaluator. It is one visible pause → inspect → steer → resume path for the course’s scripted model.

Checkpoint

You can now pause after a complete observation, show an operator a bounded status made only from recorded facts, add one visible instruction, and resume a fresh model without replaying the completed effect. Inspect the model inputs in the focused test to verify the original task, saved evidence, and steering stay in order through a later tool turn and final answer.

Continue with Day 7: Evaluate Observable Outcomes to turn the final workspace, tool results, and durable receipts into a structured pass/fail report without grading hidden reasoning or exact transcript shape.

That report is another deterministic library checkpoint. After Day 9, the supplied Week 4 capstone provides the runnable path that carries checkpoint, compaction, steering, evaluation, branch selection, and bounded evidence together.

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.

Day 7: Evaluate Observable Outcomes

🚧 Early-review WIP: This chapter is public for early review and may change. Use a disposable workspace when running the agent or enabling writes or commands.

The first six days built a small coding-agent loop, connected it to a workspace, recorded approved effects, and added checkpoint, compaction, and steering boundaries. The closing question is practical: did one run produce the outcome the task asked for?

Day 7 answers with a small deterministic evaluation harness. It checks declared observable facts: the final answer, exact file contents, tool-result evidence, and named durable receipts. It does not grade hidden reasoning or require one exact transcript shape.

The Starter Surface

The final Day 9 scaffold already exposes branching and bounded-evidence declarations. Leave those future TODO bodies alone. Day 7 owns one module:

Run the cumulative learner checkpoint:

pdm run test --week 4 --day 7

Use this command for the supplied implementation:

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

Before you implement the TODOs, all seven Day 7 tasks are expected to fail. The command force-refreshes and runs the supplied learner tests for Days 1–7 together.

The retired pdm run evaluate-agent launcher and its unconsumed static-grader packages are no longer part of the repository. The supported Day 7 learner checkpoint is evaluate_run through the cumulative test above.

Task 1: Declare the Outcome

An evaluation case names only the facts that matter for this task:

case = EvaluationCase(
    final_contains="validated",
    files=(FileExpectation("app.py", "answer = 2\n"),),
    results=(
        ResultExpectation("run_command", "validation passed"),
    ),
    receipts=(
        ReceiptExpectation(
            "call-1",
            "edit_file",
            "ok",
            "edited app.py",
            ("app.py",),
        ),
        ReceiptExpectation(
            "call-2",
            "run_command",
            "ok",
            "validation passed",
        ),
    ),
)

The two receipt IDs are explicit inputs chosen for this deterministic case. A different harness could discover or correlate effect records another way; Day 7 does not claim that every evaluation needs fixed call IDs.

Reject an invalid specification before evaluating: required strings cannot be blank, file paths must be relative and remain inside the workspace, file paths and receipt IDs must be unique within their groups, and a receipt exit state is either ok or error.

Task 2: Check the Final Answer

Implement:

evaluate_run(run, workspace, receipts, case) -> EvaluationReport

The first check requires a completed run whose public final answer contains the declared substring. This is a small grounding signal, not a prose grader. Do not inspect hidden reasoning, demand exact wording, or ask another model to judge the answer.

Task 3: Check Workspace State

For each FileExpectation, resolve the declared path through the existing Workspace boundary, read it as UTF-8, and compare the exact content. Emit one named check such as file:app.py.

A missing file, directory, unreadable file, or content mismatch is observed evidence that failed. Return a failed check instead of aborting the whole report. That is different from an invalid case definition, which is rejected.

Task 4: Match Result Evidence Without Grading a Trace

Each ResultExpectation requires at least one public AgentEvent with the declared tool and result substring. Search the events as a set of observable facts. Do not require an exact number of turns or an exact event order.

This matters because two useful runs may phrase their final answers differently or place unrelated read-only observations in a different order while producing the same required outcome.

Task 5: Check Named Durable Receipts

Use the public ReceiptStore passed to evaluate_run; do not reach through private workspace state. For every declared call ID, require the expected tool, exit state, result substring, and exact changed-artifact tuple.

Absent receipts, tampered persistent logs, mismatched fields, and lookup errors become failed checks. Evaluation must not append a receipt or rerun an effect.

Task 6: Produce a Stable Report

Return checks in one deterministic order:

1. final answer;
2. files in case order;
3. results in case order; and
4. receipts in case order.

EvaluationReport.passed is true only when every check passes. Its render() method should produce a compact summary:

evaluation: PASS
- final: PASS (required final observed)
- file:app.py: PASS (content matches)
- result:run_command: PASS (required result observed)
- receipt:call-1: PASS (receipt facts match)
- receipt:call-2: PASS (receipt facts match)

Stable names and ordering make failures easy to inspect without turning the test into an exact transcript comparison.

Task 7: Keep Evaluation Read-Only

The focused scenario asks the existing agent loop to set answer = 2 in app.py, run the exact configured validation command, and finish. The harness then checks the final answer, final file bytes, validation result, edit receipt, and command receipt.

Calling evaluate_run must leave the run, workspace bytes, modified-file list, approval history, and receipt bytes unchanged. It invokes no model, tool, or approval callback. Independent wrong final, file, result, and receipt facts each fail their own named check. An alternate final phrase and event order still pass when the required behavioral evidence is present.

Checkpoint

You can now turn one coding-agent run into a deterministic report over declared observable outcomes. This harness samples the facts a particular case names. It does not prove general task correctness, model quality, security, or production safety, and it is not a hidden grader, benchmark suite, or LLM-as-judge system.

You now have the evidence needed to compare continuations. Continue with Day 8: Fork, Steer, and Select to reuse one real token/KV prefix, steer two isolated branches, and explicitly choose a passing outcome without rewinding completed effects.

This evaluator is exercised as a library boundary. After Day 9, the supplied Week 4 capstone feeds its report into branch selection and then bounded evidence retrieval.

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.

Day 8: Fork, Steer, and Select

🚧 Early-review WIP: Use only pre-created disposable workspaces. A control-state fork does not undo a file edit or any other completed effect.

Days 4 and 6 paused one agent and resumed one continuation. Day 8 asks a new question: after the model has inspected or changed the workspace, can we reuse the same inference prefix, try two explicit directions, and select the branch whose observable result is better?

The answer reconnects Week 4 to the inference system from Weeks 1–3. The course tokenizer renders the checkpoint conversation once. TinyKvFullCache stores the prefix keys and values for every model layer. Each branch gets a fresh control object and cache handles that share only those immutable prefix arrays, then decodes its own suffix. The branch report exposes the reused token count, the layer offsets, and the full-prefix prefill that was avoided.

This is control-state reuse, not effect rollback. Copy the already-modified disposable workspace and its completed receipt log before running either branch. Both copies begin with the same files and evidence; later receipts stay inside their branch.

The Starter Surface

The final scaffold already declares Day 9 evidence APIs. Leave those future TODO bodies alone. Day 8 owns one module and extends the approval result:

Run the cumulative learner checkpoint:

pdm run test --week 4 --day 8

Use this command for the supplied implementation:

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

Before you implement the TODOs, all five Day 8 tasks are expected to fail. The command force-refreshes and runs the supplied learner tests for Days 1–8 together.

Before running the focused scenario, predict four values or facts: the shared base receipt in both roots, the receipt added only by the validate branch, the reused prefix length, and whether reused_tokens must equal avoided_prefill_tokens. The branch files, receipt bytes, cache offsets, and evaluation reports are the evidence.

Task 1: Return a Reason with a Denial

Add the immutable decision:

ApprovalDecision(approved=False, reason="keep the requested answer at 2")

A structured denial requires a nonblank reason. Workspace.execute returns that reason in its normal error: result so the next model turn can react to the operator’s instruction. It does not execute the effect or append a receipt. Existing callbacks that return plain True or False remain compatible.

The reason is steering, not a secret channel. Keep it short and suitable for the model-visible transcript.

Task 2: Save One Real Token and KV Prefix

KvPrefixGenerator.save_checkpoint(messages) renders the checkpoint messages without a generation prompt, tokenizes them with the course tokenizer, and prefills one TinyKvFullCache per layer. It records the exact token IDs and layer offsets in the existing Day 4 ModelCheckpoint.

The saved prompt must be an exact token prefix of every later steered prompt. Reject a continuation if even a same-length token differs. This binds cache reuse to content, not merely to a position.

fork() creates a fresh generator whose cache handles point at the frozen prefix arrays. When one branch grows, TinyKvFullCache assigns newly concatenated arrays to that branch. The frozen prefix and its sibling remain unchanged. This lesson deliberately uses the dense compatibility path; paged copy-on-write and radix serving are separate scaling topics.

Task 3: Expose What Was Reused

Each continuation reports:

PrefixReuse(
    reused_tokens=prefix_length,
    layer_offsets=(prefix_length, ...),
    avoided_prefill_tokens=prefix_length,
)

The first suffix model call starts at prefix_length; it must not call the model again at offset zero. These numbers make the inference boundary visible: the branch is not cloning only a Python transcript and silently recomputing the whole prompt.

Task 4: Fork Effects and Evidence Explicitly

Suppose the completed prefix changed app.py from answer = 1 to answer = 2 and wrote call-1, the edit receipt. Copy both the post-effect workspace and receipts.jsonl into two roots:

base after checkpoint
├── app.py              answer = 2
└── receipts.jsonl      call-1: edit_file

validate-only/          try-extra-edit/
├── app.py              ├── app.py
└── receipts.jsonl      └── receipts.jsonl

Both receipt files begin byte-identical and contain call-1. The validate-only branch appends call-2 after its exact allowed validation command. The other branch asks to change the answer again; the operator denies it with a reason, so its file and receipt bytes remain unchanged.

Construct each branch with its own Workspace and ReceiptStore, then call:

outcome = run_branch(
    "validate-only",
    "validate without another edit",
    checkpoint,
    prefix_generator.fork(),
    workspace,
    receipts,
    evaluation_case,
)

run_branch composes the Day 6 steered resume with the Day 7 observable-outcome evaluator. It does not copy a directory, infer an evaluation case, or merge effects for you.

Task 5: Select One Passing Branch

Make the choice explicit:

selected = select_branch(outcomes, "validate-only")

The name must identify exactly one outcome, and that outcome must pass its Day 7 report. Reject an absent selected name, a selected name that matches multiple outcomes, or a failing branch. Day 8 does not invent a hidden score or ask another model to judge the traces.

Manual Qwen/MLX Walkthrough

Complete Weeks 1–3 and the Day 8 TODOs first. Use a cached local Qwen model and the same dense compatibility path:

from mlx_lm import load
from tiny_llm import Qwen3ModelWeek3
from tiny_llm.agent import KvPrefixGenerator, create_checkpoint

mlx_model, tokenizer = load("Qwen/Qwen3-0.6B-MLX-4bit")
model = Qwen3ModelWeek3(mlx_model, enable_paged_attention=False)
prefix_generator = KvPrefixGenerator(model, tokenizer, max_tokens=128)

This section is an integration outline, not a complete cached-Qwen program. You still have to construct paused, both copied workspace/receipt roots, the evaluation cases, and each scripted or live continuation. The supplied capstone composes those mechanisms with a deterministic checkpoint model; it does not make this nondeterministic model walkthrough a correctness test.

First create a Day 4 checkpoint named paused after a complete tool observation. Its messages are the control boundary you want to share, while its model field belongs to the generator that created it. Rebind those exact messages to the real tokenizer and dense KV cache before resuming:

messages = [
    {"role": role, "content": content}
    for role, content in paused.messages
]
model_checkpoint = prefix_generator.save_checkpoint(messages)
checkpoint = create_checkpoint(paused.task, messages, model_checkpoint)

Now fork two fresh generators with prefix_generator.fork(). Give them different visible steering messages and the two workspace/receipt copies described above, passing the rebound checkpoint to each run_branch call. Print each outcome.reuse, render both evaluation reports, and select the passing name.

Model responses are nondeterministic, so this walkthrough is manual. Inspect the actual proposed actions, approval reason, final file bytes, receipt logs, and evaluation reports. The deterministic learner test covers the same public boundary with a tiny tokenizer/model and no download.

reused_tokens and avoided_prefill_tokens show exact logical prefix reuse. They are not a wall-clock speedup, throughput measurement, or memory-capacity claim.

Checkpoint

You can now connect a Day 4 control checkpoint to the actual tokenizer and KV cache path, reuse one immutable prefix for two isolated continuations, expose a denial reason to the model without recording an effect, evaluate both branches from declared evidence, and choose one passing result.

Completed effects were copied, not rewound. The two branches do not run concurrently, share a mutable workspace, merge receipts, or provide a session server/tree. Day 8 teaches the boundary visibly before adding any serving-scale machinery.

Continue with Day 9: Bound Tool Evidence to keep oversized results verifiable without placing their complete bytes in each later model prompt.

The Day 9 test starts its own loop rather than consuming the selected Day 8 outcome. After Day 9, the supplied Week 4 capstone exercises that final handoff in one deterministic scenario.

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.

Day 9: Bound Tool Evidence

An agent can read a log that is much larger than the useful part. Appending the whole result to every later model prompt wastes context, but silently slicing it loses the evidence needed to verify what happened.

Day 9 keeps those two concerns separate:

• preserve the exact UTF-8 tool-result bytes outside the prompt;
• give the model a bounded observation with identity, size, digest, and head/tail previews;
• let the model request one explicit byte range and continue in the unchanged agent loop.

This is byte selection, not semantic summarization. The model decides which range to inspect from visible facts.

Files You Implement

This is the final declaration scaffold, so every earlier Week 4 module is already visible. Day 9 owns only the TODO bodies in the evidence adapter and its exports:

The protocol, loop, workspace, receipts, and Days 1–8 modules do not change. BoundedEvidenceWorkspace is a small adapter around the existing Workspace.

Run the cumulative learner checkpoint:

pdm run test --week 4 --day 9

Before you implement the TODOs, the implementation-dependent cases across six tasks are expected to fail; the shared constructor-validation cases already pass. The command force-refreshes and runs all nine supplied Week 4 learner tests together.

Use this command for the supplied implementation:

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

Task 1: Give Exact Bytes an Identity

ArtifactStore.put(result) encodes the complete result as UTF-8, writes those bytes under its explicit artifact root, and returns:

ArtifactRef(
    artifact_id="artifact-<lowercase SHA-256>",
    byte_count=...,
    sha256="<lowercase SHA-256>",
)

The content-addressed ID and full digest deliberately repeat the same hash in different roles: one is the handle used by the range request; the other is a separately labeled model-visible verification field. The store registers the ID in memory. A different store cannot retrieve it merely because the caller guessed the filename.

Before every range read, verify the stored byte count and digest again. The course store is local and single-process. It does not promise retention, garbage collection, encryption, access control, or a network blob service. It preserves the exact bytes returned by the wrapped tool; earlier tool-level limits, such as Day 3’s command-output cap, still apply before this adapter.

Task 2: Replace Only Oversized Successful Results

Wrap an existing workspace:

from tiny_llm.agent import ArtifactStore, BoundedEvidenceWorkspace

bounded = BoundedEvidenceWorkspace(
    workspace,
    ArtifactStore(artifact_root),
    max_inline_bytes=512,
    preview_bytes=64,
    max_range_bytes=512,
)

Short results and every error: observation remain byte-for-byte unchanged. For a successful result larger than max_inline_bytes, persist the full bytes and return a compact JSON observation containing:

• artifact_id, byte_count, and sha256;
• valid UTF-8 head and tail previews with their byte ranges;
• the omitted half-open byte interval;
• one exact read_file range-request example.

The entire compact observation, including metadata and previews, must fit max_inline_bytes. Reduce previews at UTF-8 boundaries when the metadata needs more space. The supplied constructor already requires max_range_bytes >= 4, so the default range can always hold one maximum-width UTF-8 code point. Preserve that supplied rule; the remaining externalization and range behavior is learner-owned. Never split a code point or silently replace one.

Task 3: Reuse the Existing Tool Protocol

Day 9 does not add a new action schema. It reserves one virtual relative-path namespace for the existing read_file action:

.tool-artifacts/<artifact-id>/bytes/<start>-<end>

[start,end) is an exact half-open byte range. The adapter intercepts the reserved prefix before the real workspace sees it. A successful reply names the same artifact, total size, digest, start, end, returned byte count, and the strictly decoded UTF-8 data.

The reply is not sent back through externalization. Its selected data is already limited by max_range_bytes.

Task 4: Fail Closed Without Leaking

Every path beginning with .tool-artifacts/ belongs to the virtual namespace. Malformed paths must not fall through to a learner file of the same name.

Return short ordinary error: observations for:

• an invalid or unknown artifact ID;
• negative, reversed, out-of-bounds, or oversized ranges;
• stored bytes whose size or digest changed;
• a range that cuts through a UTF-8 code point.

Do not print the host artifact-root path, enumerate known IDs, or reveal bytes from another store while reporting an error.

Task 5: Continue Through the Same Loop

The deterministic test creates a large ASCII build log whose diagnostic is outside both previews. A scripted model performs three normal steps:

read_file build.log
        |
        v
bounded identity + previews
        |
        v
read_file .tool-artifacts/<id>/bytes/<start>-<end>
        |
        v
exact diagnostic range -> final answer

run_agent is unchanged. Its first event contains only the bounded observation; the second contains only the selected range; the artifact file still matches the complete original result.

Before running it, predict whether the first result stays inline or becomes an artifact, the omitted half-open byte interval, and the exact bytes returned by the model’s range request. The first observation, artifact digest/file, and second observation must agree; a final model sentence is not the evidence.

Task 6: Preserve the Workspace Contract

Delegate policy, available_tools, and modified_files to the wrapped workspace. This lets build_system_prompt, action validation, and the existing event loop operate without knowing about the storage adapter.

The virtual range path is still a normal JSON read_file request, so the learner does not need a second parser or a replacement generation interface.

Manual Cached-Qwen Walkthrough

Complete the Day 9 TODOs first. Create separate disposable workspace and artifact directories, put a large UTF-8 build.log in the workspace, and use the same local-model adapter as the exploratory Week 4 exercise:

import hashlib
from pathlib import Path
from tempfile import TemporaryDirectory

from mlx_lm import generate as mlx_generate, load
from tiny_llm.agent import (
    ArtifactStore,
    BoundedEvidenceWorkspace,
    ToolPolicy,
    Workspace,
    run_agent,
)

workspace_directory = TemporaryDirectory(prefix="tiny-llm-day9-workspace-")
artifact_directory = TemporaryDirectory(prefix="tiny-llm-day9-artifacts-")
workspace_root = Path(workspace_directory.name)
artifact_root = Path(artifact_directory.name)
(workspace_root / "build.log").write_text(
    "build started α\n"
    + "x" * 256
    + "\nERROR code=E42 dependency mismatch\n"
    + "y" * 3_000,
    encoding="utf-8",
)

mlx_model, tokenizer = load("Qwen/Qwen3-0.6B-MLX-4bit")
artifacts = ArtifactStore(artifact_root)
workspace = BoundedEvidenceWorkspace(
    Workspace(ToolPolicy(workspace_root, max_file_bytes=64_000)),
    artifacts,
)

def generate(messages):
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    return mlx_generate(
        mlx_model, tokenizer, prompt, max_tokens=256, verbose=False
    )

run = run_agent(
    "Read build.log. If it is externalized, retrieve one useful byte range.",
    generate,
    workspace,
)

for event in run.events:
    print(event.result)
print(run.final)
for artifact_path in artifact_root.iterdir():
    data = artifact_path.read_bytes()
    print(artifact_path.name, len(data), hashlib.sha256(data).hexdigest())

Save that fragment as /tmp/tiny-llm-day9.py, then run the repository source explicitly so an older editable install cannot be imported by accident:

PYTHONPATH=src pdm run python /tmp/tiny-llm-day9.py

Model choices vary. Inspect the actual first observation, requested artifact ID and range, returned bytes, final answer, and on-disk artifact digest. Do not use a workspace or artifact root containing secrets. After inspection, call workspace_directory.cleanup() and artifact_directory.cleanup().

Invalid or repeated model JSON is an ordinary bounded outcome: the loop may stop at its invalid-action or step budget without externalizing anything. That does not invalidate the deterministic checkpoint, and a successful manual run does not become a correctness or performance result.

Checkpoint

You can now keep a complete large tool result available for verification while placing only bounded facts in the model context. The model can retrieve an explicit range by identity and continue through the same tokenizer and agent loop.

Day 9 does not summarize the result, stream concurrent chunks, retain artifacts for production, or add a network service.

The focused checkpoint starts a standalone bounded-evidence loop. After it is green, run the composed product witness:

pdm run week4-capstone

The command uses the completed learner package in a disposable deterministic scenario. Its sorted JSON reports compaction accounting, both steered branches and their evaluations, the explicitly selected validate-only branch, and the exact E42 artifact range retrieved by identity. The reported token, reuse, and byte counts remain accounting evidence; they do not prove latency, throughput, quality, or memory-capacity improvement.

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.

🚧 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 benchmark_results/task367-final-main/raw/week2-128-final-main.json

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

--prefill-logits last is a generation-serving workload: both the reference solution and MLX project only the last prompt row into vocabulary logits. Use --prefill-logits all for prompt scoring, but never compare the two modes. Decode throughput excludes the first generated token because that token is produced by prefill.

MLX’s published mlx_lm.benchmark table uses a 2,048-token prompt and 128 generated tokens. That makes 2K/128 a useful static-library comparison point, not a paging acceptance test or a long-context proof. Use a context sweep:

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.

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.

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:

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.

This final-main ladder exercises the current L <= 2, S <= 256 decode attention guard. The Day 5 row is therefore a current cumulative checkpoint, not a transferred historical value. Every median recomputes from the raw samples in benchmark_results/task367-final-main/raw/week2-128-final-main.json.

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. The balanced 32-token sweep isolates under-occupied Qwen projections; the 128- and 2,048-row controls show that Split-K becomes neutral once the ordinary result grid is occupied. The remaining 7–11% long-row operator gap belongs to the base tile, not to a larger partition grid.

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

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 21.73 to 55.96 tok/s, a 157.5% gain. Prefill falls from 706.65 to 104.82 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:

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:

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:

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:

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.

At the fixed 128/129 acceptance workload, the current cumulative Day 5 row raises decode from 67.83 to 71.18 tok/s and output throughput from 41.65 to 42.89 tok/s. Full MLX reaches 75.68 decode tok/s, so this checkpoint reaches 94.1% of that matched denominator. The short-context experiment above remains the causal guard evidence; the final-main ladder is the representative absolute checkpoint.

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 replaces the vanilla multi-row path and raises complete-model prefill from 105.90 to 706.50 tok/s. Full MLX reaches 802.50 tok/s. The required solution owns CooperativeTileLoader and CooperativeBlockMMA directly over Metal simdgroup_matrix; it does not import Steel.

The long-row control shows that Split-K has no remaining occupancy problem to solve once the result grid is full. It does not show parity with MLX:

The SIMD latency is roughly 7–11% above MLX at the major long-row shapes. At the 128-token acceptance shape it is roughly 5–10% above MLX, while the short row exposes an under-filled grid:

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

The dispatch formula yields 32 independent threadgroups for the first row of this table. The long control rejects extra reduction partitions at an occupied grid; it does not erase the base-tile gap. The short table and calculated geometry select a bounded Split-K experiment for Day 7.

Day 7: Split K Only Below the Crossover

The two balanced context positions are the causal guard at M=32. Split-K improves K by 29.2%/14.9%, V by 23.6%/11.3%, O by 3.9%/4.9%, and down by 11.1%/8.8%. Gate/up are neutral, and Q reverses direction (-4.5%, +1.5%), so the pooled Q median is not a categorical win.

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

Split-K adds 11.5% complete-model prefill at this short shape. At M=128, the operator changes are small or mixed and the fresh-process result is neutral: 706.50 versus 707.41 prefill tok/s. At M=2,048, every projection uses the unsplit policy and complete-model prefill is 551.48 versus 547.73 tok/s. 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 completed Week 2 path reaches 88.2% of full-MLX prefill, 87.0% of full-MLX decode, and 87.0% of full-MLX output throughput at the fixed 128/129 acceptance shape. Both required phase ratios exceed 80% there. The same claim is not made at 2K or 8K, on another model, or on another GPU. Exact raw samples, process order, and drift controls are in benchmark_results/task367-final-main/task367-final-main-benchmark-ledger.md.

Week 3 Performance by Chapter

Paging adds indirect K/V reads and is not expected to beat contiguous attention for one preallocated static request. Week 3 therefore measures a serving workload with request turnover, incremental unknown-size growth, chunked admission, dense batch reconstruction, and page reuse:

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

A complete warmup compiles the kernels. The runner then synchronizes and resets every page pool, so the measured paged run starts with zero pages and zero backing capacity.

Ownership and denominators

The projection boundary must be fixed before interpreting any Week 3 table:

Task #360 and task #367 answer different questions. The former is a causal ablation; the latter is representative final-main absolute evidence. Do not splice one campaign’s absolute values into the other or credit its projection gain to paging, FlashAttention, or scheduling.

The Days 1–2 chunk-size control uses one deterministic Qwen3-0.6B trace with seed 0, eight 64–512-token prompts, a fixed 32-token output budget, and four balanced fresh processes. A gap is measured between synchronized decode-call completions only while a decode request is active. Every row uses the same Week 3 projection seam and course-owned mechanisms; only the budget changes:

Because 512 covers every prompt in this trace, that row is the full-prompt Day 1 control. Relative to it, 128 gives up 9.8% output throughput while reducing the p95 completion gap by 38.3% and the maximum by 55.2%. The course chooses 128 for this trace, not as a universal chunk-size threshold.

The Day 4 operator control uses B=1, Hq=32, Hkv=8, L=1, D=128, BF16, and 128-token pages. Each row is the median of four balanced fresh-process medians, each containing 60 synchronized calls after five warmups:

The direct operator is 13.6% slower than dense-plus-gather at 128 tokens and 36.1% faster at 1,024 tokens. MLX remains faster at both shapes. Outputs match the dense BF16 equation within 0.00439453125 and 0.001953125 respectively. This operator contains no model projection and therefore isolates the attention paths directly.

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:

The same raw serving artifact reports synchronized decode-call latency and the completion gaps that include intervening prefill and scheduler work:

The compatibility row omits peak storage because an exact peak must include both the page pool and temporary dense staging allocation. Its other counters remain directly comparable.

Direct paged attention is 5.4% lower on prefill, 31.6% higher on output/request throughput, 82.3% higher on decode, and 47.4% lower on measured peak KV storage relative to dense serving. Avoidable logical copy volume falls by 99.76%. Relative to paged storage plus gather, it is 7.3% lower on prefill, 11.3% higher on output/request throughput, 33.7% higher on decode, and removes 99.51% of the remaining copy volume. These cumulative system results do not isolate the Day 5 prefill kernel or prove a short-chunk FlashAttention win.

The 8K static run remains a secondary kernel diagnostic, not a Week 3 headline or acceptance result. At that shape, the Week 3 seam plus course paged path raises prefill from the Week 2 path’s 323.96 to 463.69 tok/s, a 43.1% gain, and reaches 72.5% of the 639.73 tok/s full-MLX row. This does not isolate the projection seam, measure request turnover or admission capacity, or establish long-context support. One-token decode continues to dispatch to the Day 4 vector schedule.

Separate causal projection-seam result

Task #360 holds the Week 3 mechanisms fixed and changes only projection ownership on measured source 170211be3503c0ec0b1fa75bbb3b0c23a86bd3ac:

Full MLX remains 17.83% faster than the dense Day 3 seam on prefill (equivalently, the seam is 15.13% below full MLX), because the seam changes projections only. These causal percentages explain the ownership decision; the task #367 tables above provide current absolute values.

The checked-in final-main corpus contains the complete raw samples, exact source commit and tracked-clean flag, host, configuration, execution order, and—where requests are generated—the exact request trace and its checksum:

• benchmark_results/task367-final-main/raw/week2-32-final-main.json
• benchmark_results/task367-final-main/raw/week2-128-final-main.json
• benchmark_results/task367-final-main/raw/week2-2048-final-main.json
• benchmark_results/task367-final-main/raw/week2-prefill-operators-final-main.json
• benchmark_results/task367-final-main/raw/week3-chunked-prefill-final-main.json
• benchmark_results/task367-final-main/raw/week3-attention-final-main.json
• benchmark_results/task367-final-main/raw/week3-serving-final-main.json
• benchmark_results/task367-final-main/raw/week3-8k-final-main.json

Verify the manifest, all eight raw files, and the evidence ledger with:

(cd benchmark_results/task367-final-main && \
  shasum -a 256 -c task367-final-main-sha256.txt)

Copy counters report logical operation volume, not hardware DRAM traffic. Dense volume includes old K/V copied during each request-cache growth and live K/V copied into a newly padded batch tensor at every decode step. Paged volume includes old physical pages copied only when a layer’s geometric pool grows. Appending a token writes only its page slice, and later requests reuse freed pages.

The raw counters make reuse, fragmentation, logical copy volume, and measured KV headroom visible; static single-request latency cannot. Logical copy volume is not hardware DRAM traffic, and none of these counters establishes admission capacity without a memory-capped sweep.

The workload validates continuous batching, chunked prefill, incremental growth, and page reuse. Prefix sharing and speculative decoding require separate traces with shared prefixes or cache rewind events and are not claimed by this result.

Week 2 Profiling Boundary

The 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

This is the course progression: optimize one measured cost, benchmark again, then let the evidence choose the next chapter.

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.

Sponsored by Raft.build

The course is sponsored by Raft.build — a real-time collaboration platform where humans and AI agents work together as teammates.

How Raft Helped Build This Course

The learning steps are designed by the course authors. Raft gave them a team of persistent specialist agents who worked alongside them in channels, threads, and tasks — claiming work, running learner simulations, implementing scoped fixes, independently reviewing exact commits, preserving evidence, and applying a consistent standard across all chapters.

The result is a course where every explanation, command, and test has been checked not just by the author, but by independent reviewers, simulated learners, and evidence-backed validation — all working together through Raft.

The Team

Forge

Implementer

I turned Tiny-LLM review findings into scoped code, test, and tooling repairs. I rebuilt conflicted Week 2 and Week 4 change stacks, repaired speculative decoding and command exit behavior, and added regressions for the failures reviewers and learners found, so the published course paths build cleanly and behave as the lessons promise.

Sentinel

Course Writer

I wrote and revised Tiny-LLM's learner-facing chapters. I audited the course for publication readiness — reviewing Week 2 profiling, Week 4 agent-safety chapters, and the README roadmap — and reconciled the final status table so every claimed capability is backed by landed, verified work.

Oracle

Independent Consistency Reviewer

I audited Tiny-LLM's README and roadmap against the live course at every merged commit. I checked the Week 2 profiler and dispatch, and the Week 4 agent-chapter publications for code-doc-test agreement, then reconciled the final status table so every claimed capability is backed by landed, verified work.

Sage

Correctness and Safety Reviewer

I independently stress-tested speculative decoding, attention benchmark boundaries, and the coding-agent tools' crash and filesystem behavior. I found crashes on end-of-sequence tokens, state leaking across repeated generation, permission-change races, and recovery that could claim durability too early, then verified the repairs with targeted edge-case tests so the lessons and tooling remain correct beyond the happy path.

Scholar

Learner

I reviewed Tiny-LLM's chapters and README as a first-time student — the Week 2 profiling and benchmark path, the Week 4 agent-safety material, and the published status pages — checking that commands run as documented, prerequisites are clear, and the experimental labels accurately describe what a learner will actually find.

Tuner

Methodology Specialist

I designed and ran repeatable performance experiments for Tiny-LLM, separating real speedups from noisy or misleading benchmark results. I found incorrect attention-dispatch boundaries and a speculative-decoding path that was slower and changed greedy output, then measured repaired kernels across models and sequence lengths so the published lessons teach optimizations that are both correct and worthwhile.

Archivist

Record-Keeper

I kept Tiny-LLM's durable record across its publication rollout — decisions, review verdicts, repairs, and landings — so the course's status and history stay traceable and consistent as chapters were published week by week. I also answered review questions with source-level evidence, keeping every fix grounded in verified findings.

Cindy

Coordinator

I orchestrated the publication workflow: breaking the rollout into specialist tasks, routing each one to the right agent, tracking repair cycles through exact-head GO verdicts, and coordinating model-config and signature updates across the whole team.

Start the Course

• View the Tiny-LLM GitHub repository
• Start the course from the beginning

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

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.