Preface

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 three weeks. We will serve Qwen3 MLX models and optimize the serving path throughout the course.

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.

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.

Setting Up the Environment

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

Install PDM

Follow the official installation guide to install PDM.

Clone the Repository

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

The repository is organized as follows:

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

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

Install Dependencies

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

Check the Installation

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

Run Unit Tests

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

pdm run test

Download the Model Parameters

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

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

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

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

Then, you can run:

pdm run main --solution ref --loader week1

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

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

Week 1: From Matmul to Text

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

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

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

What We Will Cover

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

What We Will Not Cover

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

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

Basic Matrix APIs

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

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

Qwen3 Models

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

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

📚 Readings

Week 1 Day 1: Attention and Multi-Head Attention

On Day 1, we will implement basic attention and multi-head attention. An attention layer processes an input sequence and weighs the relevance of its different positions when producing each output. Attention is a key building block of Transformer models.

📚 Reading: Transformer Architecture

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

📚 Reading: LLM Inference, the Decode Phase

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

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

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

Task 1: Implement `scaled_dot_product_attention_simple`

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

src/tiny_llm/attention.py

📚 Readings

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

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

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

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

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

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

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

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

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

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

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

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

Task 2: Implement `SimpleMultiHeadAttention`

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

src/tiny_llm/attention.py

📚 Readings

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

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

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

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

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

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

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

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

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

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

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

You can run all tests for the day with:

pdm run test --week 1 --day 1

Week 1 Day 2: Positional Encodings and RoPE

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

📚 Readings

Task 1: Implement Traditional Rotary Positional Encoding

You will need to modify the following file:

src/tiny_llm/positional_encoding.py

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

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

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

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

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

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

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

You can implement this operation by reshaping x to (N, L, H, D // 2, 2) and applying the formula to each pair.

📚 Readings

You can test your implementation by running the following command:

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

Task 2: Implement Non-Traditional `RoPE`

Qwen3 uses a non-traditional arrangement of RoPE pairs. Split the head dimension into two halves, then pair corresponding values from the halves. Let x1 = x[..., :HALF_DIM] and x2 = x[..., HALF_DIM:].

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

Implement this form by selecting the first and second halves of x directly, applying the rotations, and concatenating the results.

📚 Readings

vLLM implementation of RoPE

You can test your implementation by running the following command:

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

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

pdm run test --week 1 --day 2

Week 1 Day 3: Grouped Query Attention (GQA)

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

Readings

Task 1: Implement `scaled_dot_product_attention_grouped`

You will need to modify the following file:

src/tiny_llm/attention.py

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

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

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

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

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

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

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

Finally, reshape the result to the expected output shape.

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

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

In addition to grouped heads, this function supports different query and key/value sequence lengths: Q uses length L, while K and V use length S.

You can test your implementation by running the following command:

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

Task 2: Causal Masking

Readings

Writing an LLM from scratch, part 9 – causal attention

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

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

The additive causal mask has shape (L, S), where L is the query sequence length and S is the key/value sequence length. Allowed positions contain 0, and masked positions contain -inf. When S is greater than L, shift the diagonal by S - L so that the queries correspond to the final L positions in the key/value sequence. For example, if L = 3 and S = 5, the mask is:

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

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

You can test your implementation by running the following command:

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

Task 3: Qwen3 Grouped Query Attention

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

src/tiny_llm/qwen3_week1.py

Qwen3MultiHeadAttention implements attention for Qwen3. Follow this pseudocode:

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

Qwen3 attention has no Q/K/V projection biases, and it applies RMSNorm to each Q and K head before RoPE. We will implement the reusable RMSNorm layer on Day 4, so call mx.fast.rms_norm directly for q_norm and k_norm today. Use non-traditional RoPE.

You can test your implementation by running the following command:

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

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

pdm run test --week 1 --day 3

Week 1 Day 4: RMSNorm and the Multilayer Perceptron

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

Task 1: Implement `RMSNorm`

In this task, we will implement the RMSNorm layer.

src/tiny_llm/layer_norm.py

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

📚 Readings

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

RMSNorm is defined as:

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

where:

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

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

D is the embedding dimension.

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

You can test your implementation by running:

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

Task 2: Implement the MLP Block

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

src/tiny_llm/qwen3_week1.py

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

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

A plain FFN can be abstracted as:

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

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

SwiGLU is the GLU variant used by Qwen3:

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

📚 Readings

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

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

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

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

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

The negative branch is algebraically equivalent to the direct sigmoid formula, but it prevents exp(-x) from becoming exp(large positive) when x is a large negative value. In vector code, compute the direct branch with abs(x), then use 1 - y for negative inputs. This approach closely matches MLX’s low-precision GPU path.

Then implement Qwen3MLP. Qwen3’s MLP contains:

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

This can be expressed as:

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

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

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

You can test your implementation by running:

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

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

pdm run test --week 1 --day 4

Week 1 Day 5: The Qwen3 Model

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

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

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

Tests that require an unavailable model will be skipped.

Task 1: Implement `Qwen3TransformerBlock`

src/tiny_llm/qwen3_week1.py

📚 Readings

Qwen3 uses the following Transformer block structure:

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

Run the tests for this task with:

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

Task 2: Implement `Embedding`

src/tiny_llm/embedding.py

📚 Readings

LLM Embeddings Explained: A Visual and Intuitive Guide

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

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

This can be implemented with array indexing.

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

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

Run the tests for this task with:

# This task's tests use the 0.6B model and tokenizer.
hf download Qwen/Qwen3-0.6B-MLX-4bit
pdm run test --week 1 --day 5 -- -k task_2

Task 3: Implement `Qwen3ModelWeek1`

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

src/tiny_llm/qwen3_week1.py

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

The Qwen3 model has the following layers:

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

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

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

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

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

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

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

Run the tests for this task with:

# Download each model you want to test. Missing models are skipped.
hf download Qwen/Qwen3-0.6B-MLX-4bit
hf download Qwen/Qwen3-1.7B-MLX-4bit
hf download Qwen/Qwen3-4B-MLX-4bit
pdm run test --week 1 --day 5 -- -k task_3

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

pdm run test --week 1 --day 5

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

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

Task 1: Implement `simple_generate`

src/tiny_llm/generate.py

simple_generate takes a model, tokenizer, prompt, and optional sampler, then streams the generated response to standard output. Generation has two phases: prefill and decode.

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

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

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

logits = output_logits[:, -1, :]

You may normalize these logits into log probabilities with the log-sum-exp trick. This normalization does not change the result of argmax, but the sampler introduced on Day 7 expects log probabilities. If sampler is None, use mx.argmax along the final, vocabulary dimension. Otherwise, pass the log probabilities to sampler. Selecting the highest-scoring token at every step is called greedy decoding.

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

With _step complete, implement the rest of simple_generate. Begin by encoding the prompt into a one-dimensional token array with tokenizer.encode.

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

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

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

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

You can test your implementation by running the following command:

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

# If downloaded, you can also try the larger models.
pdm run main --solution tiny_llm --loader week1 --model qwen3-1.7b \
  --prompt "Give me a short introduction to large language model"
pdm run main --solution tiny_llm --loader week1 --model qwen3-4b \
  --prompt "Give me a short introduction to large language model"

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

Week 1 Day 7: Sampling and Preparing for Week 2

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

Task 1: Sampling

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

src/tiny_llm/sampler.py

📚 mlx-lm sampler implementation

Temperature Sampling

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

To implement temperature sampling, divide the log probabilities by the temperature and pass them to mx.random.categorical.

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

Top-k Sampling

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

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

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

Top-p (Nucleus) Sampling

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

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

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

Task 2: Prepare for Week 2

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

Install Xcode:

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

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

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

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

You may also need to accept the Xcode license:
```
sudo xcodebuild -license accept
```
Install CMake:
```
brew install cmake
```

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

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

That completes Week 1. We have implemented all the components required to serve Qwen3. In Week 2, we will optimize the serving infrastructure for Apple silicon.

🚧 Week 2: A Step Closer to vLLM

Week 2 keeps the readable Week 1 model intact and builds a separate optimized Qwen3 path for single-request decoding. It begins by changing the algorithm: prefill once, retain a dense KV cache, and decode one new token at a time. Only then does a matrix-vector kernel describe the workload we are optimizing.

Every later chapter starts from the runnable checkpoint produced by the previous chapter. Implement one replacement, integrate it into the Week 2 model immediately, verify correctness, and measure the new end-to-end decode rate before continuing. There is no final chapter where a pile of isolated operators suddenly becomes a model.

Week 2 inherits Week 1’s BF16 model-storage contract. Dense and quantized weights, activations, projections, KV-cache entries, and model-facing kernel outputs are BF16. Numerically sensitive reductions, dot products, and online-softmax state accumulate in FP32 inside readable expressions or kernel registers. Week 2 extensions are GPU-only: readable Python/MLX equations and vanilla Metal kernels provide correctness references without requiring CPU BF16 support. This contract remains in force for Week 3, so later chapters only describe new storage and scheduling behavior.

What We Will Cover

A dense per-request key-value cache for incremental decoding
Synchronized benchmarking of the cached baseline against MLX
A readable quantized matrix product and a SIMD matrix-vector decode kernel
The course-owned decode-attention primitive
Fast RMSNorm, RoPE, and SwiGLU operations
A BF16 SIMD-matrix quantized prefill kernel
A last-token output interface for generation
An acceptance target of 70% of MLX decode throughput

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

We are still building on MLX as infrastructure. mlx_lm loads the official Qwen3 4-bit checkpoint and tokenizer. mlx.core supplies arrays, lazy graph evaluation, memory management, device streams, and synchronization. The MLX extension API registers our C++ primitive and dispatches our Metal kernels. Those facilities are the platform on which the course implementation runs; they are not substitutes for the operator implementations themselves.

The order is intentional:

KV cache: copy the readable Week 1 operators into a Week 2 model, add request-scoped state, and stop recomputing the prefix.
Benchmark: measure that cached model and the matched cached MLX baseline.
Quantized matvec: keep weights packed, integrate the SIMD decode kernel, and measure the first operator replacement.
Decode attention: replace the exact Week 1 float32 attention composition with the course-owned online-softmax kernel and measure the whole model.
Fast kernels: replace RMSNorm, then RoPE, then SwiGLU. Each replacement is integrated and benchmarked before the next one begins.
SIMD-matrix prefill: optimize the larger matrix shape, introduce 8×8 matrix fragments, and retain FP32 accumulators behind BF16 inputs/outputs.

A later chapter never becomes an undeclared prerequisite for an earlier one.

Unlike Week 1, the completed Week 2 model prefills a dense KV cache once, passes only the new token during decode, keeps its linear and embedding weights quantized, dispatches separate decode and prefill matrix schedules, and imports optimized operations from week2_kernels.py. Week 1 continues to use its readable full-prefix generation loop and Python RMSNorm, RoPE, attention, and MLP implementations.

Week 3 imports these Week 2 interfaces rather than copying or replacing them. Its paged-attention chapters combine Day 4 online softmax and Day 6 matrix fragments only after page-table translation has been introduced. That boundary lets each week’s model remain understandable and runnable on its own.

The cumulative ladder is executable at any time. The performance appendix records the matched results:

pdm run bench-week2-progression --offline --repeats 3 \
  --model qwen3-0.6b --input-len 128 --output-len 65 --warmup 2

The runner executes each checkpoint in a fresh process and reports its median against Week 1 and MLX. The performance appendix records the cumulative percentages in one place. They are not additive promises: replacing one bottleneck changes how much every later replacement matters.

The default runs the reference checkpoints. After implementing the cumulative selector in your model, add --solution tiny_llm to measure your own complete ladder. Preserve the named checkpoints as you work; a later implementation should add a new branch without changing what an earlier checkpoint executes.

🚧 Week 2 Day 1: KV Cache

In this chapter, we will add a key-value cache to the Qwen3 model. During generation, the cache lets each attention layer reuse the keys and values from previous tokens instead of recomputing the entire prefix at every step.

This is the foundation of Week 2 decode optimization, not a serving-only Week 3 feature. Without it, every generated token reruns all model layers over an ever-growing prefix, overwhelming the gains from faster individual kernels.

📚 Readings

KV Caching Explained: Optimizing Transformer Inference Efficiency

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

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

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

The attention mechanism is computed as:

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

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

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

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

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

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

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

[a b c d] represent cached values

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

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

Task 1: Implement the Key-Value Cache

src/tiny_llm/kv_cache.py

Each Transformer layer maintains its own key-value cache. The cache exposes one method, update_and_fetch, which:

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

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

You may implement this in kv_cache.py as TinyKvFullCache:

L_new = number of incoming tokens

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

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

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

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

return key, value, self.offset, mask

Task 2: Preserve the Week 1 Boundary

src/tiny_llm/qwen3_week2.py

Keep the readable Week 1 model and its full-prefix generation loop unchanged. Start a separate qwen3_week2.py model with the same dense weights and readable 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 readable Week 1 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 readable 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; “readable” describes the implementation, not a return to an FP32 model.

Task 3: Create Request-Scoped Caches

src/tiny_llm/qwen3_week2.py

Implement create_kv_cache so every request gets one cache handle per Transformer layer. Pass the matching layer cache through every block and keep the caller’s offset consistent with the cache’s logical length.

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

pdm run test --week 2 --day 1

Task 4: Connect the Serving Loop

src/tiny_llm/generate.py

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

For example:

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

You can test your implementation with:

pdm run main --solution tiny_llm --loader week2 \
  --week2-checkpoint kv-cache --model qwen3-0.6b
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-0.6b

Integrate and Measure

Run the cached readable checkpoint end to end before changing any operator:

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

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

🚧 Week 2 Day 2: Benchmark Decode

Optimization starts with a trustworthy comparison. In this chapter, we measure prefill and decode separately, synchronize MLX’s lazy execution inside every timed iteration, and record a baseline that later changes must beat.

Prefill and Decode Are Different Workloads

Prefill processes many prompt tokens at once, so its matrix multiplications have a larger row dimension. Decode usually processes one token per request and is dominated by repeatedly reading quantized weights. A change can improve one phase while hurting the other, so 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.

For a matched prefill comparison, all implementations compute logits for every prompt position. The generation path may request only the final logit row, but using that shortcut for the course model while MLX projects all rows would make the prefill columns incomparable. Cached decode has L = 1, so logits_to_keep=1 removes no decode work.

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 a course model that recomputes the full prefix would measure two different algorithms and make the kernel target meaningless. Day 1 already produced the cached readable model used as this week’s starting point.

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 release request-owned caches after warmups and timed runs so a later sample does not inherit allocator state:

pdm run test --week 2 --day 2

Debug Metal Without a CPU Twin

From Day 3 onward, the course extensions are GPU-only. A second C++ CPU implementation would duplicate the equation without exercising the dispatch, memory, or synchronization behavior that makes a Metal kernel fail. Use this three-level validation ladder instead:

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

Compare each level with the one immediately above it. Do not debug an optimized kernel by comparing only full-model text output.

Make Failures Small and Synchronous

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

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

expected = readable_operator(*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.

Metal API Validation and an Xcode GPU capture can help diagnose dispatch and resource problems, but they supplement this ladder rather than replace its small deterministic comparisons. Only profile after the vanilla and optimized kernels agree with the readable oracle.

The isolated benchmarks in benches/ use the same rule. Evaluate input setup before invoking the benchmark fixture so setup does not leak into the result. The Week 2 operator ladder compares the readable implementation, the optimized course implementation, and MLX at the selected model’s real tensor shapes:

pdm run bench-week2-operators --model qwen3-0.6b --context 128

For the measurements quoted later in this week, the machine was an M4 Pro with a 20-core GPU and 64 GB of memory. Each operator used 20 warmup iterations and 100 synchronized timed iterations in each of three fresh processes. The reported result is the median process-level speedup. The matched end-to-end commands used two complete warmups and three fresh measured runs.

Record a Matched Baseline

Use the same model, prompt length, output length, device, and warmup count for your implementation and the MLX run. Replace tiny_llm with tiny_llm_ref to compare against the course reference:

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

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

Or run the complete cumulative ladder in fresh processes. At this point, only the Week 1, KV-cache, and MLX rows are course prerequisites; later rows become meaningful as you complete their chapters.

pdm run bench-week2-progression --offline --repeats 3 \
  --model qwen3-0.6b --input-len 128 --output-len 65 --warmup 2

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 with the result.

Acceptance Target

The Week 2 target is:

reference decode throughput / MLX decode throughput >= 0.70

Reaching 70% is the acceptance threshold, not a promise that every educational kernel individually matches its MLX counterpart. MLX is the comparison baseline; the Week 2 solution must reach the target with course-owned operator implementations.

🚧 Week 2 Day 3: Quantized Matvec

In this chapter, we will study and implement quantized matrix multiplication. Quantizing weights from 16-bit floating point to 4-bit integers reduces both model size and the memory traffic required for each generated token.

📚 Readings

Why Quantization?

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. The following rough calculation illustrates this for Qwen3-0.6B:

Per-token linear layers in decode phase:
- Input: 1 token × 1024 dimensions = 1024 bfloat16 values = 2 KB
- MLP weights: 1024 × 3072 × 3 matrices × 2 bytes = ~19 MB per layer
- Attention weights:
  - q_proj / o_proj: 1024 × 2048 × 2 matrices × 2 bytes = ~8 MB per layer
  - k_proj / v_proj: 1024 × 1024 × 2 matrices × 2 bytes = ~4 MB per layer
- Total weights per layer: ~31 MB
- Total for 28 layers: ~880 MB

FLOPs (2 per multiply-accumulate):
- MLP per layer: 2 × 3 × 1024 × 3072 ≈ 19M
- Attention projections per layer: 2 × (1024 × 2048 × 2 + 1024 × 1024 × 2) ≈ 13M
- 28 layers: ~880 million per token

Memory access: ~880 MB
Arithmetic intensity: 880M FLOPs / 880 MB ≈ 1.0 FLOPs/Byte

With M3 Max’s 400 GB/s memory bandwidth and ~10 TFLOPS compute:

Memory-bound throughput: 400 GB/s × 1.0 FLOPs/Byte = 400 GFLOPS
Compute-bound throughput: 10 TFLOPS

This workload can use only about 4% of the available compute before exhausting memory bandwidth.

The Solution: Quantization

Compressing BF16 weights to 4-bit integers (int4) can:

Reduce memory traffic by 4×: 880 MB → ~220 MB per token
Improve arithmetic intensity by 4×: 1.0 → ~4.0 FLOPs/Byte
Increase throughput by ~4×: 400 GFLOPS → ~1.6 TFLOPS

The tradeoff is a small loss in model accuracy when the weights are quantized carefully.

Group-wise 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 MLX 4-bit checkpoints used in this course have a fixed group size of 128:

Original weight matrix W: K × N (bfloat16)

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

For each group of G consecutive values in a row:
  1. Find min and max values
  2. Compute scale and bias to map [min, max] → [0, 15] (4-bit range)
  3. Quantize each value using: quantized = round((value - bias) / scale)

All required quantized-matmul tests use group_size = 128 and BF16 scales, biases, activations, and outputs. Normalize those tensors to BF16 when loading the course model so every later kernel receives one model dtype.

Affine Quantization

We use affine (asymmetric) quantization, which maps a floating-point range onto the full integer range:

$quantized = round (\frac{value - bias}{scale})$

$dequantized = quantized \times scale + bias$

For 4-bit quantization, the quantized values are in the range $[0, 15]$ .

Given a group with minimum value $v_{min}$ and maximum value $v_{ma x}$ :

$scale = \frac{v _{ma x} - v _{min}}{2 ^{bits} - 1} = \frac{v _{ma x} - v _{min}}{15}$

$bias = v_{min}$

Example:

Group values: [-0.5, -0.3, 0.1, 0.4, 0.8]
min = -0.5, max = 0.8

scale = (0.8 - (-0.5)) / 15 = 1.3 / 15 ≈ 0.0867
bias = -0.5

Quantization:
  -0.5 → round((-0.5 - (-0.5)) / 0.0867) = 0
  -0.3 → round((-0.3 - (-0.5)) / 0.0867) = 2
   0.1 → round((0.1 - (-0.5)) / 0.0867) = 7
   0.4 → round((0.4 - (-0.5)) / 0.0867) = 10
   0.8 → round((0.8 - (-0.5)) / 0.0867) = 15

Quantized: [0, 2, 7, 10, 15] (4 bits each)

Storage Format

The quantized values are packed for compact storage and efficient access:

Original: K × N bfloat16 (2 bytes each) = 2KN bytes
Quantized: K × N int4 (0.5 bytes each) = 0.5KN bytes

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

Weight matrix shape: K × N
Quantized storage 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

Quantized Matrix Multiplication

Mathematical Formulation

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

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

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

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

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

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

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

Substituting:

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

Rearranging:

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

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

Computation Flow

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

Output:
  C: M × K (bfloat16)

For each output element C[i, k]:
  sum = 0  # float accumulator
  for each group g in 0..(N/G - 1):
    scale = scales[k, g]
    bias = biases[k, g]
    
    # Process G values in the group (G/8 uint32 packs)
    for each pack p in 0..(G/8 - 1):
      packed_value = B_quantized[k, g*(G/8) + p]
      
      # Unpack 8 × 4-bit values
      for bit_offset in [0, 4, 8, 12, 16, 20, 24, 28]:
        quantized = (packed_value >> bit_offset) & 0xF
        b_value = quantized * scale + bias
        a_value = A[i, g*G + p*8 + bit_offset/4]
        sum = sum + a_value * b_value
  
  C[i, k] = bfloat16(sum)

Task 1: Implement Quantized Linear and Embedding

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

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

Field	Shape	Description
`weight`	$(K, N /8)$ uint32	Packed quantized weights. Each uint32 stores eight consecutive 4-bit values.
`scales`	$(K, N / G)$ bfloat16	Per-group scale factors for dequantization. Each group of $G$ consecutive values shares one scale. Recall: $scale = (v_{ma x} - v_{min}) /15$
`biases`	$(K, N / G)$ bfloat16	Per-group bias (offset) for dequantization. Recall: $bias = v_{min}$
`group_size`	int	Number of consecutive values that share the same scale/bias. For the Qwen3 MLX 4-bit weights used here, this is `128`.
`bits`	int	Quantization bit width (typically 4, meaning values are in range $[0, 15]$ )

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

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

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

embedding(input_ids) performs a row lookup. Gather the matching packed weights, scales, and biases. Unpack each uint32 with shifts and masks, repeat each group’s scale and bias across its 128 values, and compute q * scale + bias with basic mlx.core array operations. Do not call mx.dequantize. Put this readable 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: Register a GPU-Only Primitive

Register quantized matrix multiplication as an MLX C++ extension. Follow the existing axpby example for array validation, lazy primitive construction, bindings, and Metal dispatch. The course implementation is GPU-only; its eval_cpu method should raise a clear unsupported-device error.

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

You will update four files. Keep the C++ declarations and definitions in the tiny_llm_ext namespace:

tiny_llm_ext.h — Declare the quantized_matmul(...) function signature and define a QuantizedMatmul primitive class (inheriting mx::Primitive). Store group_size and bits as private members.
bindings.cpp — Add an m.def(...) call to expose the function to Python.
quantized_matmul.cpp — Implement quantized_matmul(...) to validate inputs, determine the output shape, return a lazy mx::array, and reject CPU evaluation explicitly.
CMakeLists.txt — Add the new C++ source to the extension target.

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 and test the extension:

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

Task 3: Implement Metal Matrix Products

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

Write the Metal kernels and connect eval_gpu to them. The Python quantized_matmul wrapper always dispatches this course-owned primitive on GPU; the required path never routes through mx.quantized_matmul.

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

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

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 is a useful debugging oracle.

Prefill Tiling Comes on Day 6

Prefill has many activation rows and benefits from a different matrix-matrix schedule. Keep the vanilla kernel for that shape today so Day 3 stays focused on decode. Day 6 replaces it with 8×8 simdgroup_matrix tiles after the course has established the packed format, dispatcher, and synchronized benchmark.

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. The regular projection kernel calculates two output columns per SIMD group. The much wider tied vocabulary projection uses eight columns so each activation load is reused across more weights.

The wide path also uses the affine identity

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

to avoid applying the bias separately to every unpacked value. This adds an activation-sum accumulator, so test it separately for ordinary and wide output tiles in the scheduling experiment below.

Tune the SIMD Schedule

Treat output width, threadgroup size, and shared-memory reuse as benchmark variables. Start with this host dispatch plan:

flatten all leading activation dimensions into M,
use the custom matvec when M <= 8,
compute two output columns per SIMD group for ordinary projections,
switch to eight output columns when K >= 8192 for the wide vocabulary head,
launch eight SIMD groups per threadgroup.

These thresholds are measured starting points, not mathematical requirements. Keep them visible in the dispatcher, then vary one choice at a time. The reference measurements below are one-factor scheduling checks from an M1 Pro; use them to form hypotheses, but report the results from your own hardware.

For output tiling, four columns in an ordinary projection reduced full-model decode from about 249 to 232 tok/s because the extra accumulators increased register pressure. Eight columns helped the unusually wide vocabulary projection because each activation load was reused across more weights.

Apply the affine rearrangement selectively as well. It helps the wide path, but applying it to the two-output projection reduced the measured full-model result from about 249 to 244.5 tok/s. Fewer arithmetic operations do not imply a faster kernel when they extend register lifetimes or complicate scheduling.

At the Python-to-extension boundary, make scales, biases, activations, and packed weights row-contiguous once with mx.contiguous. The C++ primitive validates that contract before encoding the kernel. A Metal kernel receives raw buffers and strides are not implicit, so silently accepting a noncontiguous view would produce either wrong addressing or a slower hidden copy in a less explicit layer.

Do not copy the 2 KB activation vector into threadgroup memory by default. On the reference machine, rereading this cache-hot vector avoided a barrier and raised decode from roughly 238.6 to 246-249 tok/s. Verify this result by adding the shared-memory variant as an ablation; it is a useful demonstration that reuse helps only when it costs less than synchronization.

Finally, compare four, eight, and sixteen SIMD groups per threadgroup. Four groups measured about 248.5 tok/s and did not improve on eight. Using sixteen groups only for the vocabulary path reduced 250.6 to 247.2 tok/s. Start with eight groups for both shapes, then retune on a different GPU rather than assuming that either smaller scheduling units or fewer threadgroups must win.

Direct Quantized Embedding Comes on Day 6

Week 2 performs row lookup and dequantization with basic mlx.core array operations at this checkpoint. Day 6 optionally fuses row gather, int4 unpacking, and affine dequantization into one Metal dispatch.

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.
The required kernel supports bfloat16_t inputs and outputs. The course 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 course-model dispatch.

GPU Dispatch

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

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

You can test your implementation 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

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 readable boundary; 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 implementation by running:

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

You can also benchmark throughput and compare your implementation with the reference solution:

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

Compare this result with the Day 1 kv-cache row. Do not start the decode attention chapter until the complete model uses packed weights and the end-to-end number has been recorded. The vanilla matrix product remains callable as a correctness oracle, but only the SIMD matvec is integrated into decode.

🚧 Week 2 Day 4: Decode Attention

This chapter starts from the quantized-matvec checkpoint. During single-request decode, query length is normally one while the cached key/value sequence grows by one token at a time. Week 1 expresses attention as matrix multiplication, masking, softmax, and another matrix multiplication. That is readable, but it materializes the complete score and probability rows.

First write a readable composition to preserve the equation, then replace its matmuls and softmax with a course-owned online-softmax Metal kernel. 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

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

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

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

The online update is:

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

After its last cached position, each group writes its partial maximum, denominator, and value accumulator to threadgroup memory. The first SIMD group computes the common maximum and rescale factors. One thread computes the final denominator, then the first D threads combine one output dimension each. This parallel final reduction was faster than making the first 32 lanes each reduce four dimensions. 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. The course 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 while holding the model and 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. On the reference M1 Pro run, the three schedules reached about 215, 232, and 238-239 decode tok/s respectively. Record your own result and repeat the experiment when context length or head dimension changes.

Task 3: Integrate and Measure

Route short-query, short-context Week 2 attention through the Metal implementation. Dispatch back to the readable 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 readable 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 the course kernel only when query length is at most eight and cached context length is at most 256. Otherwise use the readable 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 readable compatibility path. They appear in the first continuous-batching exercise, while the normal Week 2 decode path uses no explicit mask. Week 3 replaces dense batch masks with paged attention metadata instead of complicating this focused decode kernel.

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

Test grouped-query head mapping, output shape, causal behavior, and explicit masks against the readable Week 1 implementation. Use a tolerance because the online softmax changes the floating-point reduction order.

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

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

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

The model dispatches short-query contexts through the course kernel and falls back to the exact readable Week 1 composition outside the measured range. This is the same evidence-driven decision used for tile sizes, barriers, and threadgroup layouts elsewhere in the course.

🚧 Week 2 Day 5: Fast Kernels

Week 1 expresses RMSNorm, RoPE, and SiLU as readable mlx.core equations. Week 2 keeps those implementations intact and writes three course-owned Metal kernels behind a separate interface:

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

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

Why a Metal Kernel Helps

Calling the Week 1 code “Python” does not mean Python visits every tensor element. Python builds a lazy graph whose individual array operations already run as native kernels. The important difference is how many operations and memory passes the graph describes.

For example, readable RMSNorm 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 course-owned Metal kernel gives us explicit control over the whole model 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.

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

Task 1: RMSNorm

Begin with one SIMD group per input row, then profile it. A 1024-element hidden row gives 32 lanes too much serial work. 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 readable Week 1 equation rounds once before applying the weight, so compare the two with a tolerance rather than expecting bit-identical results. The single final cast also tracks the MLX model more closely in the Qwen3-4B end-to-end correctness test.

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. The two-level reduction was the largest small-operator improvement in the measured stack; a single SIMD group left too little parallel work available.

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 5 -- -k rms
pdm run bench --solution tiny_llm --loader week2 \
  --week2-checkpoint rmsnorm --model qwen3-0.6b

Task 2: RoPE

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 readable RoPE in the already optimized model, then test and measure that cumulative checkpoint before implementing SwiGLU:

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

Task 3: SwiGLU

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 5 -- -k swiglu
pdm run bench --solution tiny_llm --loader week2 \
  --week2-checkpoint swiglu --model qwen3-0.6b

Task 4: Verify the Cumulative Model

At this point all three kernels have already been exposed through C++ MLX primitives, integrated, and measured. Run the complete test file now to verify their composition. qwen3_week1.py must still use its readable operators, and Week 3 should import the Week 2 interfaces so the serving model does not regress.

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

Day 6 changes workload shape rather than replacing another element-wise or reduction operator: it introduces SIMD-matrix fragments for quantized prefill. Keep today’s swiglu checkpoint intact so that prefill tiling has a clean before-and-after comparison.

Compare against the readable 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. The course 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.

🚧 Week 2 Day 6: SIMD-Matrix Prefill

Day 3 optimized the one-to-eight-row matrix-vector shape used during decode. Prefill is a different workload: it projects tens or hundreds of token rows at once. Today we add a second quantized-linear schedule for that larger M dimension and introduce Metal’s 8×8 SIMD-matrix fragments. Week 3 will reuse the same fragment API after paged KV addressing is established.

This chapter assumes the packed W4A16 format and dispatch interface from Day 3, the synchronized prefill benchmark from Day 2, and the BF16 model contract from the Week 2 overview. No quantization format changes today.

From SIMD Reduction to SIMD-Matrix Fragments

The Day 3 decode kernel gives a SIMD group several output columns and reduces each dot product across its 32 lanes. That schedule fits a very small number of activation rows. During prefill, the same weight row is reused across many activation rows, so it is better to tile both output dimensions.

Metal exposes an 8×8 simdgroup_matrix fragment. A fragment is a register-held matrix tile distributed across the 32 lanes of one SIMD group; it is not an ordinary thread-local array. The public matrix-multiply-accumulate operation computes:

C[8, 8] += A[8, 8] × B[8, 8]

Each SIMD group owns one 8×8 output tile. It advances through the reduction dimension eight values at a time, loading one activation fragment and one dequantized weight fragment before issuing the multiply-accumulate.

activation rows M
        |
        v
  +-----------+     reduction N in 8-value steps
  | A 8 x 8   |  ×  | B 8 x 8 |  --->  C 8 x 8
  +-----------+                         one SIMD group
        |
        +---- next 8 activation rows

output columns K are covered by independent 8-column tiles

This is the first use of SIMD-matrix fragments in the course. Keep the scalar Day 3 kernel as the correctness oracle while bringing up the tiled path.

Mixed-Precision Boundary

The completed course model stores activations, scales, biases, and outputs in BF16. Unpack each 4-bit weight and combine it with its BF16 scale and bias, but accumulate the matrix product in FP32 fragments. Cast only once when storing the final BF16 output tile.

This distinction applies for the rest of the course:

Location	Dtype
Model activations and KV storage	BF16
Quantization scale and bias storage	BF16
Matrix operands loaded by the kernel	BF16
Dot-product and online-softmax accumulators	FP32
Model-facing kernel output	BF16

BF16 is the model-storage contract; FP32 is an internal arithmetic choice. A correct result with an FP32 output still violates the model contract.

Task 1: Add Shape-Based Dispatch

Keep the Day 3 matrix-vector path for M <= 8. For larger M, dispatch a new SIMD-matrix kernel:

M <= 8  -> quantized SIMD matvec
M > 8   -> quantized SIMD-matrix matmul

Expose the prefill path through QuantizedWeights.use_simdgroup_matmul while you compare checkpoints. The dispatch must preserve the same function signature, packed weights, output shape, and BF16 dtype as Day 3.

Task 2: Implement the 8×8 Kernel

Use one SIMD group per 8×8 output tile. For every eight-wide reduction step:

Load the activation fragment cooperatively.
Unpack the matching 4-bit weights.
Apply the group scale and bias while forming the weight fragment.
Issue simdgroup_multiply_accumulate into an FP32 accumulator fragment.
After the complete reduction, store only valid rows and columns as BF16.

Zero-fill partial input fragments. Guard partial output tiles at the final store. A prefill length such as ten tokens must exercise both a complete tile and a two-row tail without changing the accumulation dtype.

Task 3: Hoist Quantization Parameters

One scale and bias cover 128 reduction elements, or sixteen consecutive eight-wide matrix steps. Loading them inside every step repeats address calculation and device reads.

Loop over quantization groups first. Load the scale and bias values needed by the output fragment into registers, then reuse them for all sixteen steps:

for each 128-value quantization group:
    load scale and bias for this output fragment
    for each of its sixteen 8-value reduction tiles:
        unpack and dequantize weights
        matrix-multiply-accumulate into FP32

Measure this change independently. More reuse is only a hypothesis until a synchronized benchmark shows that its register cost is worthwhile.

Task 4: Fuse Quantized Embedding Lookup

Prompt tokens do not need a matrix multiplication. They select rows from the quantized embedding table and dequantize only those rows. The readable Day 3 path expresses gather, int4 unpacking, scale, and bias as separate array operations. Add a direct Metal kernel that fuses those steps into one dispatch.

Map one thread to one requested embedding element. Accept both int32 prompt IDs and uint32 sampled IDs so generation does not insert a token-cast graph node. This kernel reuses the W4A16 unpacking equation from Day 3; the only new idea is fusing row selection with dequantization.

Task 5: Keep Only Needed Logits

Prefill computes hidden states for the whole prompt, but generation samples only from the final position. Preserve an optional logits_to_keep model argument and slice hidden states before the final norm and vocabulary projection:

if logits_to_keep is not None:
    h = h[:, -logits_to_keep:, :]

Generation requests one row. Correctness tests and prompt-scoring callers can still pass None to obtain every position. This is a model-interface optimization, not a change to attention or sampling.

Task 6: Integrate the Prefill Checkpoint

Add a cumulative simd-matmul checkpoint after swiglu. It enables use_simdgroup_matmul for every quantized projection and the direct embedding kernel while retaining the Day 3 matvec for decode. Update generation to request only the final logit row. Run:

pdm run test --week 2 --day 6

pdm run bench-week2-progression --offline --repeats 3 \
  --model qwen3-0.6b --input-len 128 --output-len 65 --warmup 2

Report prefill and decode separately. The new schedule should improve prefill; it should not change the one-token decode path.

Prepare for Paged FlashAttention

Day 4 introduced online softmax, and this chapter introduced BF16 SIMD-matrix fragments. Week 3 first adds page-table translation, then combines these established ideas in paged FlashAttention: one SIMD-matrix product forms score tiles, online softmax updates the running state, and a second SIMD-matrix product applies probabilities to value tiles.

🚧 Week 3: Build a Mini vLLM

Week 3 takes the optimized single-request model from Week 2 and builds a serving engine around it. The Week 2 operator interfaces remain intact; this week adds multi-request cache ownership, scheduling, and runtime metadata. Week 2 has already introduced BF16 model storage, online softmax, and SIMD-matrix fragments. Week 3 adds page translation before combining those pieces into paged FlashAttention.

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 serving and scheduling experiments
Optional MoE and speculative-decoding extensions

The ordering is intentional. Day 1 batches independent request states. Day 2 splits long prefills so they cannot monopolize the scheduler. Day 3 replaces a growing dense cache with fixed-size pages while retaining a dense-gather compatibility path. Day 4 removes that gather by teaching attention to walk the page table directly with a correctness-first schedule. Day 5 then tiles that same page-walking operation with Week 2’s matrix fragments. Page translation is therefore introduced before it is optimized.

The final model 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. It primarily improves serving capacity, cache reuse, and batching; page-table indirection can make one request slower. This week measures those tradeoffs rather than assuming an algorithm with a production name is already production-fast. The performance appendix keeps kernel throughput and serving usability as separate measurements.

Week 4 owns application concerns such as RAG and tool calling. This separation keeps Week 3 focused on the reusable serving substrate.

🚧 Week 3 Day 1: Continuous Batching

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

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

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

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

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

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

Day 2 will refine this scheduler with chunked prefill. A long prompt can make one prefill step much slower than a decode step, delaying every active request’s next token. Splitting the prompt into smaller chunks bounds the amount of prefill work in each scheduler iteration.

Each chunk adds another range of prompt tokens to the request’s KV cache:

# prompt_tokens contains 400 tokens; the chunk size is 128
_step(model, prompt_tokens[0:128], offset=0, kv_cache)
_step(model, prompt_tokens[128:256], offset=128, kv_cache)
_step(model, prompt_tokens[256:384], offset=256, kv_cache)
_step(model, prompt_tokens[384:400], offset=384, kv_cache)

The causal mask for each chunk has shape L x S, where L is the chunk length and S is the total sequence length after appending the chunk. For example, if the cache already contains five tokens and the next chunk contains three, the mask has shape 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

Each row can attend to all five cached tokens, itself, and any earlier token in the same chunk.

Task 1: Verify the Week 2 Batch Contract

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

Week 3 begins by exercising interfaces established earlier rather than editing them retroactively. Confirm that Week 2 FastRoPE accepts one integer offset per batch element and that the existing causal-mask helper handles L != S, as required by chunked prefill. If either contract is missing, return to the corresponding earlier-week task and complete it there; do not create a second incompatible implementation in Week 3.

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 implementation by running:

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

Task 3: Exercise the Batch-Ready Model

src/tiny_llm/qwen3_week2.py  (reuse unchanged)

The Week 2 model already accepts multiple requests, a separate offset for each batch element, and the mask returned by BatchingKvCache. Exercise that contract with several requests joining and leaving at different positions. The new Week 3 work belongs in the cache and scheduler; do not modify the Week 2 model to make this test pass.

You should pass all of the tests by running:

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

Task 4: Batch Generate

src/tiny_llm/batch.py

First implement Request.try_prefill by prefilling the complete prompt in one call. Then complete the scheduler in batch_generate: move finished prefills into idle decode slots, collect the next token and offset for each slot, and remove requests that reach EOS or max_seq_len.

Day 2 Preview: Chunked Prefill

src/tiny_llm/batch.py

On Day 2, modify Request.try_prefill to process at most prefill_max_step prompt tokens per call.

Materialize the KV cache between chunks. MLX evaluates lazily, so repeatedly extending an unevaluated cache creates an increasingly long computation graph and allows memory usage to grow. Calling mx.eval on every layer’s key and value tensors after each chunk stores the current cache and truncates that graph.

You can test your implementation by running:

pdm run batch-main

By default, this command uses Qwen3-0.6B with a batch size of five and a fixed set of prompts.

🚧 Week 3 Day 2: Chunked Prefill

A long prompt can monopolize the device while active decode requests wait for their next token. Chunked prefill limits the number of prompt tokens admitted in one scheduler step.

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)

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.

Materialize Between Chunks

MLX is lazy. Extending an unevaluated cache repeatedly creates a long graph and can grow memory usage. Evaluate every layer’s key and value tensors after each chunk so the next scheduler iteration starts from materialized state.

Task: Bound Prefill Work

Update Request.try_prefill in src/tiny_llm/batch.py to process at most prefill_max_step tokens, advance its offset, materialize the cache, and mark the request ready only when the full prompt is complete.

pdm run test --week 3 --day 2
pdm run batch-main

Compare time-to-next-token for active decode requests with a small and a large prefill step. Smaller chunks improve fairness but add scheduling overhead.

🚧 Week 3 Day 3: Paged Attention, Part 1

In this chapter, we will design the paged KV cache. This is the storage abstraction behind paged attention.

By the end of Week 3 Day 2, our serving stack already supports:

per-request KV cache
chunked prefill
continuous batching
the Week 2 SIMD-matrix prefill operator

That gives us a working miniature serving engine, but the memory layout is still too simple. KV for each request is treated as one growing dense tensor, and batching rebuilds dense K/V for all active requests. That approach is easy to teach, but it does not scale well once requests become long and numerous.

Paged attention starts by fixing the storage layout.

📚 Readings

Why the Dense KV Layout Becomes Expensive

Right now, the mental model looks like this:

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

Before attention, the runtime repacks them into:

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

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

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

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

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

The Page Abstraction

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

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

Each layer cache keeps a small page table:

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

That means:

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

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

In our Part 1 teaching implementation, 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.

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

Why Fixed-Size Pages Help

The page abstraction gives us two immediate wins:

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

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

Data Structures We Need

1. `PagePool`

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

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

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

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

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

In the reference solution, this becomes 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

In the reference solution, this becomes TinyKvPagedCache. It is created with a pool from the model. It should not allocate its own pool, because that would isolate one request from the shared page allocator.

The reference solution creates one TinyKvPagedCache per transformer layer. Those caches share the pool, but they do not share metadata: each layer cache owns its own page_ids, page_lens, and offset.

3. Tail-Append Logic

When new K/V arrives for one layer:

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

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

Prefill with Pages

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

chunk = [t0 t1 t2 t3 t4 t5]

One possible layout is:

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

That layer cache’s metadata becomes:

page_ids = [7, 2]
context_len = 6

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

Decode with Pages

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

With paged storage:

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

So if page_size = 4 and context_len = 9:

page_ids = [12, 5, 3]

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

Stage A: Keep 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 is not the final paged attention runtime yet, but it is a very useful intermediate step:

small surface-area change
easier debugging
direct correctness comparison against TinyKvFullCache

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 key Part 1 behavior is:

write new K/V into the layer cache’s tail page or newly allocated pages,
gather the layer cache’s pages back into dense K/V,
feed that dense K/V into the old attention path.

So Part 1 changes the storage model first, not the attention kernel yet.

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

Part 1 also keeps a small rewind(n) lifecycle hook. Rewind is useful for speculative decoding: if some drafted tokens are rejected, the cache must forget their K/V. In the paged cache, rewind frees whole pages that are no longer needed and shortens the valid length of the final remaining page.

Design Questions for Part 1

Before implementing, make sure the following are clear:

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

Task 1: Design `PagePool`

src/tiny_llm/paged_kv_cache.py

Design layer-owned page pools that:

own one free-page allocator per layer,
stores flat fixed-size K/V pages,
allocates and frees page ids,
supports writing a chunk into page storage,
grows backing capacity geometrically,
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

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

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 Buys on a Mac

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 capacity and lifecycle wins. They should be measured with live-request count, allocated bytes, fragmentation, and scheduler throughput—not inferred from one request’s token latency.

pdm run test --week 3 --day 3

🚧 Week 3 Day 4: Paged Attention, Part 2

In this chapter, we move from paged KV storage to the runtime metadata and execution path needed for real paged attention.

Part 1 introduced fixed-size pages, one model-owned physical pool per layer, and request-local page metadata. That change already improves the storage abstraction, but it does not yet remove the dense gather before attention. To get the full benefit, the attention path itself must understand how to read from pages.

Paged KV Cache vs Paged Attention

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

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

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

The Metadata a Paged Runtime Needs

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

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

For the current layer being executed:

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

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

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

Why `block_table` Matters

Suppose one layer cache for request A has:

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

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

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

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

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

That is exactly what block_table and context_lens encode.

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

Compared with the earlier dense path, the important difference is that the source length is no longer represented as one contiguous tensor dimension. It is reconstructed logically from the page table.

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

Prefill Metadata

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

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

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

Decode Metadata

During decode, each active request typically writes one token.

The runtime should be able to:

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

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

Choose a Schedule for Each Query Shape

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

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

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

Shape	Course-owned dispatch	Work decomposition
`L <= 8`	Vector paged decode	One threadgroup per query row; 32 SIMD groups stride over the context and merge partial `(max, sum, output)` states.
`L > 8`	Direct paged prefill	Walk logical K/V tiles through the block table and keep the schedule deliberately inspectable. Day 5 optimizes it.

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

How This Maps to `tiny-llm`

`src/tiny_llm/attention.py`

Add a new function:

def paged_attention(...):
    ...

For this course implementation, make it a correctness-first page-walking Metal kernel with online softmax:

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

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

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

After that lookup, the online-softmax update is the same recurrence as Week 2 Day 4. Keep the first page-walking schedule simple enough that block-table and tail-page bugs are visible. Day 5 will replace its inner matrix work with SIMD-matrix tiles 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.

Recommended Incremental Rollout

The safest implementation order is:

paged storage
block_table / context_lens plumbing
correctness-first page-walking GPU attention
model and batch dispatch

This order matters because it gives us a clean correctness baseline at each step.

Correctness Invariants

These are the invariants worth checking in tests:

context_len always equals the number of written logical token positions.
block_table reconstructs the same logical KV order as the dense baseline.
the allocator never hands the same page to two live cache handles unless explicit sharing is implemented.
releasing a request returns all pages owned by all of its layer caches exactly once.
decode allocates a new page only when the tail page overflows.

Task 1: Add Batch Metadata

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

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

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.

The reference solution walks every request’s block table and keeps online softmax state:

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

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

Implement two correctness-first GPU dispatches:

For L <= 8, partition logical context positions across SIMD groups and merge their partial (max, sum, output) states in threadgroup memory. The reference schedule uses 32 SIMD groups per query so group g visits positions g, g + 32, g + 64, and so on.
For longer queries, assign query rows to a direct page-walking schedule and resolve every K/V tile through block_table. Favor inspectable ownership over the final tiled performance schedule.

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

Course implementation 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 course-owned: 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/performance reference.

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

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

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.

Start by recording three operator baselines on the same machine:

the dense course attention path, including any required K/V gather,
your direct paged-attention path,
the MLX attention path as a production-library reference.

These M4 Pro measurements are reference points, not target values to copy into your report:

Batch / context	Dense course attention	Course paged attention	MLX attention
1 / 128	226 µs	219 µs	170 µs
1 / 512	300 µs	309 µs	192 µs
1 / 2048	729 µs	726 µs	263 µs
8 / 512	800 µs	626 µs	225 µs
8 / 2048	1,347 µs	1,322 µs	460 µs

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.

A correct first version may be slower than dense attention. At this checkpoint, the useful result is a trustworthy baseline and a working runtime interface.

Checkpoint 2: Design the Decode Schedule

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

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

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.

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 course path 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 can admit more concurrent requests. Day 5 optimizes its long-prefill schedule rather than routing around the page-table contract.

pdm run test --week 3 --day 4

🚧 Week 3 Day 5: Paged FlashAttention

Day 4 made attention understand paged KV storage. Its first job was correctness: translate logical token positions through block_table, visit every visible page, and merge the result with online softmax. Today we keep that exact API and storage layout but replace the long-prefill schedule with tiled FlashAttention.

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 ideas already introduced earlier:

Week 2 Day 4 introduced the online-softmax recurrence.
Week 2 Day 6 introduced BF16 8×8 SIMD-matrix fragments and partial tiles.
Week 3 Day 3 introduced physical pages and block tables.
Week 3 Day 4 introduced direct page-walking attention and the decode schedule.

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

Why Optimize the Paged Path

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

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

The algorithm is still exact attention. Only the order of loads and reductions changes.

Keep the Day 4 Interface

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

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

Put the shape dispatch inside the extension:

Query shape	Schedule
`L <= 8`	Keep the Day 4 vector paged-decode kernel.
`L > 8`, BF16, `D == 128`	Use the tiled paged FlashAttention kernel.

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

Task 1: Tile Queries and Paged K/V

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.

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

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

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

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

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

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

compare Day 4 page-walking attention with the readable MLX equation;
compare paged FlashAttention with the Day 4 path;
only then benchmark the tiled kernel.

Required fixtures include:

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

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

pdm run test --week 3 --day 5

Task 4: Integrate and Measure

The Week 3 model should use the tiled paged path automatically for supported long prefills. Short queries continue through the vector paged-decode schedule. Neither path gathers a dense K/V tensor.

Measure long prompts as well as the short course checkpoint. Report prompt length, page size, batch size, hardware, and both prefill and decode throughput:

pdm run bench --solution tiny_llm --loader week3 --batch-decode \
  --num-seqs 16 --batch-size 4 --prefill-step 128 \
  --min-input-len 512 --max-input-len 512

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.

The serving performance lab next varies page size, chunk size, batch size, and request mix without changing this completed operator contract.

Week 3 Day 6 (Optional): Mixture of Experts

In this chapter, we will implement the feed-forward shape of Mixture of Experts, or MoE, for the Qwen3 family.

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

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

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

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

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

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

Readings

Dense MLP vs MoE MLP

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

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

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

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

The router produces one score per expert:

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

Then the model picks num_experts_per_tok experts for each token:

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

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

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

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

Qwen3-MoE Shape

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

The useful pieces are:

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

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

Grouped Quantized Matmul

MLX does not give us a single high-level MoE block in mlx.nn. It does have a lower-level primitive, mx.gather_qmm, that performs quantized matrix multiplication while selecting a different matrix for each row. In this chapter, we will build a narrow teaching version of that idea: grouped_quantized_matmul.

For MoE, that means:

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

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

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

Router Step

The router is just a quantized linear layer:

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

For a batch of tokens:

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

where E = num_experts.

Qwen3-MoE then uses top-k selection:

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

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

Expert Step

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

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

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

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

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

Task 1: Grouped Quantized Matmul

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

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

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

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

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

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

The model wrapper should:

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

logits = model(tokens, offset, cache)

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

Run the focused tests with:

pdm run test --week 3 --day 6

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

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

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

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

🚧 Week 3 Day 7 (Optional): Serving Performance Lab

Week 2 optimized the model operators. Week 3 added request scheduling and paged KV ownership. This lab measures the serving system as a system: batch size, prefill chunk size, page size, request mix, allocator reuse, and fairness. It does not introduce another matrix kernel or change the model dtype.

Start from the Completed Interfaces

The following boundaries are prerequisites, not tasks for this lab:

Week 2 quantized linear dispatches SIMD matvec for decode and SIMD-matrix matmul for prefill.
Week 3 paged FlashAttention handles long prefill.
Generation requests only the final logit row.
Week 3 continuous batching owns request admission and slot reuse.
Week 3 paged caches own physical K/V pages and expose page metadata to paged attention.

If one of those paths is missing, return to its chapter. Do not hide an incomplete prerequisite behind a benchmark-only workaround.

Choose Serving Metrics

One number cannot describe a serving engine. Record at least:

Metric	What it reveals
Time to first token	prompt admission and prefill delay
Time between tokens	decode latency and scheduler stalls
Aggregate decode tok/s	useful work across active requests
Request throughput	admission, completion, and slot reuse
Peak KV pages	cache capacity and fragmentation
P50/P95 latency	average behavior and tail fairness

Keep the workload fixed while varying one scheduling choice. Record model, MLX version, hardware, prompt/output distribution, warmup, and repeat count.

Experiment 1: Prefill Chunk Size

Run the same mixture of long prompts and active decoders with several prefill_max_step values. A smaller chunk bounds scheduler stalls but adds more model calls and cache updates. Compare aggregate throughput and P95 time-between-tokens; neither metric alone is sufficient.

pdm run bench --solution tiny_llm_ref --loader week3 --batch-decode \
  --num-seqs 16 --batch-size 4 --prefill-step 32 --model qwen3-0.6b

pdm run bench --solution tiny_llm_ref --loader week3 --batch-decode \
  --num-seqs 16 --batch-size 4 --prefill-step 128 --model qwen3-0.6b

Experiment 2: Batch Size and Request Turnover

Increase active batch size while keeping the queued request set fixed. Verify that finished requests release their pages and that newly admitted requests reuse capacity. Report aggregate throughput together with per-request latency; a larger batch can improve the former while worsening the latter.

Experiment 3: Page Size

Compare several page sizes with the same sequence-length distribution. Small pages reduce unused tail capacity but lengthen block tables and increase page management. Large pages shorten metadata but waste more space at request tails.

Measure page count, unused tail slots, and attention latency. Do not call a page size better merely because one isolated kernel becomes faster.

Experiment 4: Dense and Paged Attention

Use the dense-gather compatibility path from Day 3 as an ablation against the direct paged path from Day 4. They must run the same requests and reuse the same Week 2 model operators. Attribute differences to gather/repack work, page-table indirection, and cache reuse rather than to unrelated prefill kernels.

Long prefill should run Week 3 paged FlashAttention directly from page storage. A later chunk with cached history uses paged prefill. Short decode uses the paged vector schedule. Measure these shapes separately before combining them in one serving workload.

Preserve an Optimization Ledger

For each experiment record:

the hypothesis;
the single changed variable;
correctness status;
synchronized measurements;
the keep or reject decision and its tradeoff.

Fewer graph nodes, fewer page-table entries, or more active requests are hypotheses. Only matched measurements reveal their effect on latency, throughput, memory, and fairness.

Validate Before Measuring

Run the completed serving tests before collecting a performance result:

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

The performance appendix contains representative results and the retained optimization ledger. Use it as a reasonableness check, not as a substitute for measuring your machine.

🚧 Week 3 Optional Extension: Speculative Decoding

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 course checkpoint 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 Week 3 paged-cache lifecycle support before using the Week 3 loader.
Use compatible tokenizers for the two models. A shared token id must represent the same text in both vocabularies.

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.

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.

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.

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

Measure the Decision

Report proposal length, accepted tokens per proposal, target verification calls, draft-model time, target-model time, and end-to-end tokens per second. Compare against ordinary cached target generation on the same prompt. Keep speculative decoding optional when draft quality, cache rewind, or verification overhead makes it slower.

Do not infer a speedup from acceptance rate alone: a successful serving result must include both models, verification, synchronization, and cache maintenance.

🚧 Week 4: Build a Coding Agent

Weeks 1 through 3 turned tokens into text, made decoding efficient, and introduced serving techniques. Week 4 adds the next layer: an agent loop that lets the model observe a workspace, choose a tool, see the result, and continue until a coding task is complete.

The goal is not to reproduce a production coding agent. It is to understand the small mechanism underneath one and identify where reliability, efficiency, and safety come from. By the end of the week, you will have a local CLI agent powered by the inference stack you already built.

What You Will Build

The finished agent can:

inspect a repository without loading every file into the prompt;
read files in bounded chunks and make exact, reviewable edits;
run a narrowly scoped test command and use its output as feedback;
save and resume a session;
compact an overlong context while retaining task state;
checkpoint and undo its own file mutations;
accept steering messages and interrupt long-running work; and
solve a small repository task graded by held-out tests.

This is a deliberately small target. Features such as multi-agent delegation, remote execution, MCP integrations, and long-term user memory remain extensions rather than prerequisites.

The Core Loop

Every chapter builds on the same loop:

Render the task, project instructions, recent events, and tool descriptions.
Decode one structured action using the model and KV cache.
Parse and validate the action before it reaches the operating system.
Run one workspace tool and append its observation to the session.
Repeat until the model returns a final answer or reaches a budget.

The model does not edit files directly. It proposes an action; ordinary code decides whether that action is valid and performs it. This boundary makes agent behavior easy to inspect and test.

task + session events
        |
        v
  context builder ---> model ---> action validator
        ^                            |
        |                            v
   tool result <--- tool runner <--- validated action
                         |
                         v
                      workspace

A Small Tool Surface

The target agent uses four tools inspired by small coding-agent harnesses:

read(path, offset?, limit?)
edit(path, old_text, new_text)
write(path, content)
bash(command, timeout?)

read and edit are preferable to shell equivalents because they can enforce consistent bounds and return structured errors. bash supplies repository search, file discovery, and test execution without requiring a separate tool for every command-line program.

A shell working directory is not a security sandbox. During this course, run the agent only in a disposable exercise workspace. A production agent would need a container, virtual machine, or similarly strong isolation boundary.

The week then extends that loop with retrieval-augmented generation and a serving-oriented tool API. The WIP milestone is to move these application layers onto the Week 3 engine without changing the model kernels; the initial demo still calls the model directly.

Seven-Day Plan

Day	Topic	Working milestone
1	Agent loop	The model alternates between actions and observations.
2	Tools	The agent can read, edit, write, and run a bounded command.
3	Safety and validation	Mutations are confined, reviewable, and followed by validation.
4	Sessions	Work survives process exit and can be resumed.
5	Compaction	Long sessions retain a small, useful working context.
6	Control and recovery	The user can steer, interrupt, checkpoint, and undo.
7	Evaluation	The agent fixes a small bug and passes held-out tests.

Run the Starting Demo

The repository contains a minimal demonstration that uses the reference model implementation:

pdm run agent "inspect this project and summarize its files"

Pass --solution tiny_llm to use your implementation, or --solution mlx to use MLX-LM’s optimized executor. The starting program is intentionally smaller than the final agent: each day replaces one shortcut with an explicit component that can be inspected and tested.

The default Qwen3 4B model follows the structured action protocol more reliably. Use --model qwen3-0.6b on memory-constrained machines and expect to spend more time on malformed-action recovery.

Milestones

Minimal: the model can inspect a workspace and produce one valid action.
Useful: the agent can make a precise change and run its test.
Recoverable: the session can resume, compact, and undo its own changes.
Controllable: budgets bind and the user can steer or interrupt work.
Measurable: a repeatable task suite distinguishes progress from anecdotes.

🚧 Appendix: Performance by Chapter

This appendix connects each chapter to a measurable performance or serving effect. It distinguishes latency and throughput from usability. Paged allocation, chunked prefill, and continuous batching can make a server more usable under load even when a teaching kernel is slower for one request.

The percentages below are not additive. Replacing one bottleneck changes the fraction of time spent in every other operator. Use synchronized end-to-end measurements to judge a checkpoint, and use one-factor ablations only to explain why it moved.

Reproduce the Checkpoint Comparison

Run all checkpoints in fresh, sequential processes and report the median:

pdm run bench-course-progression --offline --repeats 3 \
  --model qwen3-0.6b --input-len 128 --output-len 65 --warmup 2

For the chapter-by-chapter Week 2 ladder, run:

pdm run bench-week2-progression --offline --repeats 3 \
  --model qwen3-0.6b --input-len 128 --output-len 65 --warmup 2

The runner prints prefill and decode throughput, speed relative to Week 1, and the remaining gap to MLX. It alternates checkpoint order across repeats to reduce systematic thermal bias. It cannot isolate the machine for you: stop other CPU/GPU workloads, keep power mode fixed, and use --cooldown-seconds when the system does not return to a stable temperature between processes. Add --json-output course-performance.json to retain every sample, the medians, seed, selected variants, MLX version, and non-identifying hardware details.

The prefill timer computes logits for every prompt position on Week 1, Week 2, Week 3, and MLX. Decode requests only the final row, but its input already has one token, so that shortcut does not change the compared decode work.

The default uses the reference solution. Add --solution tiny_llm to run the same checkpoint sequence against your implementation.

For operator attribution and serving throughput, also use:

pdm run bench-week2-operators --model qwen3-0.6b --context 128

pdm run bench --solution tiny_llm_ref --loader week3 --batch-decode \
  --num-seqs 16 --batch-size 4 --prefill-step 128 \
  --min-input-len 128 --max-input-len 128 \
  --min-output-len 65 --max-output-len 65 --warmup 1 \
  --model qwen3-0.6b

Always record the model, MLX version, hardware, prompt/output lengths, warmup, repeat count, and whether the number is a median or one representative run.

Course Optimization Map

Use this table to navigate from each optimization concept to its implementation boundary and teaching chapter. Complete the rows in chapter order; use the Week 3 serving lab to record scheduler and cache experiments that you decide not to keep.

Course step	Required implementation boundary	Where it is explained
Synchronized, matched measurement	`bench.py` evaluates timed work, computes full prompt logits for matched prefill, releases caches, and separates prefill from decode. `bench_course_progression.py` uses fresh processes, alternating order, medians, and optional pauses between runs.	Week 2 Day 2 and this appendix
KV cache before decode tuning	`TinyKvFullCache`, `create_kv_cache`, and cached generation prefill once and then pass one new token.	Week 2 Day 1
Packed 4-bit weights	`QuantizedWeights` preserves weight, scale, bias, group size, and bit width instead of materializing dense 16-bit matrices.	Week 2 Day 3
Shape-dispatched SIMD matvec	`quantized_matvec_x2` handles ordinary projections and `quantized_matvec_x8` handles `K >= 8192`; eight SIMD groups share activation loads and reduce with `simd_sum`.	Week 2 Day 3
Matvec scheduling cleanup	Reread the cache-hot activation vector instead of copying it to threadgroup memory and waiting at a barrier.	Week 2 Day 3 scheduling experiments
Online decode attention	`decode_attention_custom` keeps the query in registers, uses 32 SIMD groups, FP32 online softmax, and a parallel final reduction for `L <= 8` and context at most 256.	Week 2 Day 4
Fused RMSNorm	One 256-thread group performs a two-level reduction, normalization, and learned scaling with FP32 accumulation.	Week 2 Day 5
Fused RoPE	One thread computes an angle once for a pair and reuses it across four heads; batch offsets are normalized once per model call.	Week 2 Day 5
Fused SwiGLU	One Metal dispatch computes SiLU and multiplies the up branch without graph intermediates.	Week 2 Day 5
SIMD-matrix quantized prefill	`quantized_matmul_simdgroup` uses 8-by-8 matrix fragments, unpacking weights directly and accumulating in FP32.	Week 2 Day 6
Direct quantized embedding	The embedding kernel fuses row gather, int4 unpacking, scale, and bias for both signed and unsigned token IDs.	Week 2 Day 6
Decode graph cleanup	Generation requests `logits_to_keep=1`, and single-token decode omits the causal-mask graph.	Week 2 Day 6
Paged FlashAttention	The BF16 SIMD-matrix kernel streams K/V from page storage, uses online softmax, bounds causal work, and avoids both a dense gather and the quadratic score tensor.	Week 3 Day 5
Serving and cache scheduling	Continuous batching reuses active slots, chunked prefill bounds scheduler stalls, and paged attention separates logical context from reusable physical pages.	Week 3 Days 1–4

The model-loading boundary is not an optimization shortcut. MLX still supplies model files, arrays, streams, allocation, and extension dispatch; the Week 2 learned operators above execute course-owned Python, C++, or Metal code.

Reference End-to-End Checkpoints

Use the following Qwen3-0.6B snapshot as a reasonableness check after running the commands above. It was measured on an Apple M4 Pro with a 20-core GPU and 64 GB of memory, MLX 0.29.1, and Python 3.12.13. Every row uses a 128-token prompt, 65 generated tokens, two complete warmups, and the median of three fresh processes. The runner alternated checkpoint order across repeats.

Checkpoint	Prefill tok/s	Versus Week 1	Gap to MLX	Decode tok/s	Versus Week 1	Gap to MLX
Week 1 readable model	3,310.10	baseline	24.5% slower	19.44	baseline	93.9% slower
Week 2 completed SIMD-matrix model	1,928.49	41.7% slower	56.0% slower	244.87	12.60x	23.5% slower
Week 3 completed paged FlashAttention model	1,859.97	43.8% slower	57.6% slower	207.40	10.67x	35.2% slower
MLX	4,383.47	1.32x	baseline	320.06	16.46x	baseline

Week 1 prefill is fast because it materializes dense 16-bit weights and uses efficient dense matrix multiplication. Week 2 deliberately optimizes the memory-bound one-token decode shape first; its scalar quantized prefill path therefore regresses until Week 2 Day 6 introduces SIMD-matrix fragments. Week 3’s page-table and serving-runtime work is not free: the single-request checkpoint is 15.3% slower than Week 2 for decode and 3.6% slower for prefill at this short shape. The serving measurement below shows the separate benefit when four requests share the decode batch.

Matched Week 2 Chapter Ladder

This separate three-process campaign uses the same model, prompt/output shape, warmups, hardware, and alternating order. Its matched denominators are Week 1 at 3,313.05 prefill and 19.34 decode tok/s, and MLX at 4,373.56 prefill and 316.35 decode tok/s. End-to-end rows are cumulative; percentages are changes from the preceding executable checkpoint unless stated otherwise.

Week 1: Establish the Readable Baseline

Chapter	Performance or usability effect	Comparison
1.1 Attention	Materializes scores and probabilities so the algorithm is inspectable.	No speed claim; this is the readable attention baseline.
1.2 RoPE	Builds positions, sine, and cosine with array operations.	No isolated checkpoint gain; Week 2 later fuses this work.
1.3 GQA	Shares K/V heads, reducing KV storage versus full multi-head attention.	The Qwen3 architecture already fixes the head ratio, so the course has no MHA-to-GQA timing ablation.
1.4 RMSNorm and MLP	Uses readable array expressions and exposes intermediate tensors.	No speed claim; Week 2 measures fused replacements against these equations.
1.5 Qwen3 Model	Produces the first complete model and the baseline prefill graph.	The completed Week 1 checkpoint measured 3,313.05 prefill tok/s, 24.2% below MLX.
1.6 Generation	Recomputes the growing prefix for every output token.	The matched checkpoint reached 19.34 decode tok/s at context 128, 93.9% below MLX.
1.7 Sampling	Adds temperature/top-p/top-k policy after logits.	Not a model-throughput optimization; deterministic argmax is used by the course benchmark.

Week 2: Build the Fast Decode Checkpoint

Chapter	Change	Measured estimate	How to interpret it
2.1 KV Cache	Prefill once, retain K/V, and process one new query per decode step while preserving the readable Week 1 operators.	Prefill 3,296.76; decode 98.47 tok/s, or 5.09x Week 1 and 68.9% below MLX.	This is an algorithmic gain, not a faster operator. Its size depends strongly on context and output length and grows as full-prefix recomputation becomes more expensive.
2.2 Benchmark	Synchronize lazy work, separate prefill from decode, and establish the cached and MLX baselines.	0% direct speedup.	Measurement prevents changes in inputs, synchronization, or machine state from being mistaken for optimization.
2.3 Quantized Matvec	Keep weights packed and use the course SIMD matvec for every decode projection.	Prefill 747.08; decode 130.19 tok/s, +32.2% from Day 1 and 58.8% below MLX.	Packed weights reduce decode memory traffic but make the still-scalar prefill path much slower. The isolated table below separates kernel latency from this end-to-end result.
2.4 Decode Attention	Replace the readable score/softmax/value composition with a course-owned online-softmax kernel in its measured short-context range.	Prefill 732.02; decode 137.71 tok/s, +5.8%.	Quantized weight reads still dominate short decode. Report dtype, context, schedule, and end-to-end throughput rather than treating avoided intermediates as a speedup by themselves.
2.5a RMSNorm	Fuse square, reduction, normalization, and scaling in a course Metal kernel.	Prefill 731.85; decode 185.55 tok/s, +34.7%.	The cumulative delta includes all prior chapters; the isolated kernel is 1.31x faster than the readable equation.
2.5b RoPE	Rotate pairs directly and reuse each angle across four heads.	Prefill 766.62; decode 214.61 tok/s, +15.7%.	The corrected-layout isolated comparison is 1.52x faster than readable RoPE.
2.5c SwiGLU	Fuse SiLU and multiplication into one element-by-element kernel.	Prefill 777.80; decode 234.78 tok/s, +9.4%.	The isolated fused kernel is 1.38x faster than the readable expression.
2.6 SIMD-Matrix Prefill	Add the tiled quantized prefill schedule, direct quantized embedding, and last-token projection.	Prefill 1,932.94 tok/s, +148.5%; decode 240.98 tok/s, +2.6%.	This is the workload-shape handoff: matrix fragments repair prefill while leaving the one-token SIMD matvec path intact.
Week 2 checkpoint	Combine the dense KV cache, SIMD matvec, decode attention, fused element-wise kernels, and SIMD-matrix prefill.	12.46x Week 1 decode; 23.8% below MLX decode and 55.8% below MLX prefill.	All figures now come from one matched campaign.

Week 3: Improve Serving Structure

Chapter	What improves	Performance reality
3.1 Continuous Batching	Reuses decode slots as requests finish and keeps multiple requests active.	Improves aggregate utilization, not single-request latency. The fixed serving snapshot below reports both wall-clock output and timed decode throughput.
3.2 Chunked Prefill	Bounds how long a prompt can delay active decoders.	Primarily improves fairness and time-to-next-token. Smaller chunks add launches and may reduce aggregate throughput; report both latency and throughput.
3.3 Paged Attention, Part 1	Allocates a paged KV cache with fixed-size reusable pages and separates logical context from physical storage.	Per-layer pools grow geometrically and share storage across requests in that layer. This improves capacity, reuse, removal, and fragmentation without serializing every layer through one backing tensor.
3.4 Paged Attention, Part 2	Reads K/V through page tables without rebuilding dense per-request K/V.	The completed single-request Week 3 checkpoint reached 207.40 decode tok/s. With four active slots, the paged serving path reached a 327.62 tok/s median during timed decode sections.
3.5 Paged FlashAttention	Tiles the direct Day 4 page walk with Week 2 SIMD-matrix fragments and online softmax.	Its opportunity is long prefill; decode retains the Day 4 vector schedule. The completed checkpoint reached 1,859.97 prefill tok/s at the short course shape.
3.6 Optional MoE	Routes tokens through selected experts and supports sparse Qwen3 variants.	Expands model coverage and can reduce active parameter work, but the course has no controlled dense-versus-MoE speed claim.
3.7 Optional Serving Lab	Varies chunk size, batch size, page size, request mix, and dense-versus-paged policy without adding a model kernel.	Report latency, throughput, memory, fairness, and allocator behavior together.
Optional Speculative Decoding	Drafts several tokens and verifies them with the target model.	Work in progress. Speed depends on acceptance rate and cache-rewind cost; no result should be claimed yet.

Current Isolated Operator Snapshot

The operator runner uses synchronized median latency over 50 iterations after 10 warmups at the Qwen3-0.6B shapes. Lower latency is better. The speedup column compares the course kernel with its readable or vanilla course implementation, not with MLX.

Operator	Readable/vanilla µs	Course µs	MLX µs	Course speedup
Quantized embedding	150.7	177.4	137.6	0.85x
Q projection	190.8	135.1	106.8	1.41x
K projection	154.3	109.7	106.5	1.41x
V projection	156.5	112.9	109.7	1.39x
O projection	195.8	123.9	109.6	1.58x
Gate projection	229.2	124.1	113.8	1.85x
Up projection	230.1	123.7	114.6	1.86x
Down projection	235.3	126.1	120.6	1.87x
Tied LM head	4,547.7	439.2	427.7	10.35x
Prefill Q matmul	729.6	355.7	232.7	2.05x
RMSNorm	167.9	128.2	126.3	1.31x
RoPE	168.9	111.5	108.6	1.52x
SwiGLU	147.5	106.6	108.1	1.38x
Decode attention	138.8	126.6	107.6	1.10x

The direct embedding kernel loses this isolated one-token comparison, so it should not be described as a standalone speedup. Day 6’s end-to-end prefill gain comes from the SIMD-matrix products and last-token output boundary. The tied output head remains the largest isolated matvec win because it reads the widest packed matrix.

Current Four-Slot Serving Snapshot

This serving comparison fixes all 16 requests at 128 prompt and 65 output tokens, uses a chunk size of 128, warms up once, alternates paged and dense-gather processes, and reports the median of three runs per path.

Week 3 path	Output tok/s	Total tok/s	Prefill tok/s	Decode tok/s
Day 4 dense-gather ablation	175.02	519.66	2,423.55	207.64
Completed paged path	248.50	737.87	2,315.78	327.62

Paging is 4.4% slower during prefill in this workload, but removes enough repacking during decode to raise timed decode throughput by 57.8% and wall-clock output throughput by 42.0%. This is a serving-system result, not a claim that indirect page reads beat contiguous reads in every isolated kernel.

Week 4: Measure Application Quality Separately

Chapter	Appropriate metric
4.1 Agent Loop	Valid-action rate, steps, generated tokens, and wall-clock latency. The loop adds work; it does not improve model tok/s.
4.2 Tools	Successful bounded reads, exact-edit success, command failures, and observation size.
4.3 Safety and Validation	Rejected unsafe actions, false rejections, mutation scope, and post-edit test results.
4.4 Sessions	Resume fidelity, serialized state size, and time to restore useful context.
4.5 Compaction	Token reduction, retained-task-state accuracy, and task completion before and after compaction.
4.6 Control and Recovery	Steering latency, interrupt latency, and successful checkpoint/undo recovery.
4.7 Evaluation	Held-out task-completion rate, test pass rate, steps, tokens, and end-to-end latency.

The Week 4 chapters remain works in progress. Their success criterion is useful, safe application behavior on top of the Week 2/3 inference interfaces, not a higher isolated token rate.

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
Quantized Matmul
Fast RMSNorm, RoPE, and SwiGLU
KV Cache
Decode Attention
SIMD-Matrix Prefill
Flash Attention
Paged Attention

Keyboard shortcuts

Tiny LLM - LLM Serving in a Week