Inference Zero To Hero

Build your own LLM inference engine from first principles, and understand every byte it moves.

When you send a message to a chatbot, a program reads billions of numbers from memory for every word it writes back. Somebody wrote that program. By the end of this book, you will have written one too. It will be a real inference engine that loads open-weight models from disk, runs them on your GPU (or CPU), streams answers, serves many users at once, and implements the architecture of Qwen3.8-Flash-Next: a 176-billion-parameter hybrid model with linear attention, sparse attention, a 512-expert mixture, hashed n-gram memory and multi-stream residuals.

You will also be able to change models: fine-tune them, attach LoRA adapters, quantize them, look inside them and edit what they compute.

Who this book is for

You are a working software engineer. You can read and write Python, you know what a function, a class and a loop are, and you are comfortable in a terminal. You do not need:

• prior machine-learning experience: we start from “what is a gradient?”;
• GPU or CUDA experience: we start from “what is a thread?”;
• heavy math: high-school algebra plus a willingness to follow small worked examples. When we need a derivative or a matrix identity, we derive it on the page with numbers.

If you already know some of this, skim the early chapters, but do the engine milestones. Later chapters assume your engine works.

What you will build

The book is organized around one project: your engine. Every chapter adds a piece and ends with a milestone test suite you make pass.

The book has 46 chapters (including setup), in eight parts. The reference implementation passes the milestone suites your engine will use; Appendix F records current counts and platform coverage. Its tiny FP32 Flash-Next agrees with the official Hugging Face implementation to about 1e-7. Real-checkpoint and production performance claims require separate validation.

How each chapter works

Every chapter follows the same rhythm, so you always know where you are:

1. Why it matters: the problem, tied to the engine.
2. Concepts: intuition first (a picture, an analogy, a small example worked with real numbers), then the precise version.
3. See it: diagrams and interactive playgrounds you can poke.
4. Code it: the real implementation, in tabs. Python is the main language. C++ and Rust versions sit one click away, and your choice is remembered:

def softmax(x):
    e = [math.exp(v - max(x)) for v in x]
    return [v / sum(e) for v in e]

void softmax(float* x, size_t n) {
    float m = *std::max_element(x, x + n), sum = 0;
    for (size_t i = 0; i < n; ++i) sum += (x[i] = std::exp(x[i] - m));
    for (size_t i = 0; i < n; ++i) x[i] /= sum;
}

#![allow(unused)]
fn main() {
fn softmax(x: &mut [f32]) {
    let max = x.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let sum: f32 = x.iter_mut().map(|v| { *v = (*v - max).exp(); *v }).sum();
    x.iter_mut().for_each(|v| *v /= sum);
}
}

5. Build it: an engine milestone. You implement the chapter’s functions in your engine/ package and run its tests. The book supplies all the scaffolding (file layout, I/O, plotting, test harness), so you only write the ideas.
6. Stretch exercises: graded ★ to ★★★. Each Where note identifies the file to edit or a new experiment to create; Appendix A explains paths, imports and extension tests.
7. Check your understanding: short questions, with answers in Appendix C.
8. Going deeper: exact pointers into the source books, lectures and papers.

Callouts mark things worth stopping for:

Note

Background or a useful aside.

Tip

A practical shortcut or a debugging habit.

Warning

A pitfall that produces plausible-looking wrong answers.

Important

A platform-specific note (for example, DGX Spark), or something that changes how you should read the rest.

Pacing

At 10 to 12 hours a week, plan on roughly one chapter a week. The GPU chapters and the capstone take two. The full eight-part path is about a year, with something working at the end of every week. If you have more time, the milestones are designed to be done in one sitting each.

You can read the whole book without a GPU. The reference milestones have CPU paths, including Triton’s interpreter; CUDA kernels, graph capture and hardware-specific adapters need their target device. A GPU makes the performance chapters far more satisfying, though. The book was developed for an NVIDIA DGX Spark (GB10), and notes for that machine are marked.

Where this book comes from

The book draws on three excellent resources, and borrows from them freely:

• Sebastian Raschka, Build a Large Language Model (From Scratch) (Manning, 2024), cited as BALLM. Parts I and II follow its path from tokens to a trained GPT and reuse its corpus and several of its examples.
• Wen-mei Hwu, David Kirk and Izzat El Hajj, Programming Massively Parallel Processors, 5th edition, cited as PMPP. Part III follows its approach to CUDA, tiling, reductions and attention.
• The GPU Mode lecture series (github.com/gpu-mode/lectures), cited as Ln, for profiling, Triton, FlashAttention, quantization and serving.

You never need them open beside this book, but each chapter ends with exact pointers for a second explanation. Appendix E maps every chapter to its sources.

Start with Setup.

0. Setup: your workspace and engine

The workspace

Everything you do happens in one folder: your copy of the book’s code. It contains:

inference-zero-to-hero-code/
├── engine/      YOUR engine. Every function you'll write is there, with its signature
│                and docstring; the body says  raise NotImplementedError("TODO(Chapter 5) ...")
├── izh/         the complete reference engine: the answer key, same module and function names
├── tests/       one milestone test file per chapter (test_ch05_attention.py, ...)
├── run.py       a runnable demo for each chapter (python run.py --help)
├── data/        training text and small datasets (no downloads needed)
├── rust/  cpp/  the Rust and C++ tracks
└── models/      checkpoints you download later (created when you need it)

engine/ and izh/ have identical structure. That’s deliberate. When you’re stuck, you can open the same file in izh/ and read the reference. When you want to skip ahead, every test and demo can run against the reference instead of your engine:

IZH_IMPL=izh pytest tests/test_ch05_attention.py   # test the reference
python run.py attention --impl engine              # run a demo with YOUR engine

1. Install uv

uv manages Python versions and packages quickly and reproducibly. These commands are for Linux and macOS (Windows users: use WSL2):

curl -LsSf https://astral.sh/uv/install.sh | sh
export PATH="$HOME/.local/bin:$PATH"
uv --version

uv can download Python for you, so you don’t need a separate Python installation.

2. Create your working copy

Download the code archive (or, from a clone of the book’s repository, use docs/inference-zero-to-hero/src/code-examples.zip) and extract it into a fresh folder:

mkdir -p "$HOME/izh" && cd "$HOME/izh"
unzip -n ~/Downloads/code-examples.zip
cd inference-zero-to-hero-code
ls          # engine  izh  tests  run.py  data  rust  cpp ...

unzip -n never overwrites existing files, so re-extracting a newer archive won’t destroy your work. Put the folder under version control right away. Your engine is real code, and you’ll want its history:

git init && git add -A && git commit -m "Starting point"

3. Create the Python environment

uv venv --python 3.12
source .venv/bin/activate          # do this in every new terminal

Now install PyTorch. Choose one route.

# Let uv pick the CUDA build that matches your driver, then verify below.
uv pip install torch numpy pytest --torch-backend=auto

uv pip install torch numpy pytest --torch-backend=cpu

# macOS on M-series: PyTorch's default build includes the MPS backend.
uv pip install torch numpy pytest

Check that it worked, and that a GPU is really doing work if you have one:

python - <<'PY'
import torch
print("torch", torch.__version__, "| CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print(torch.cuda.get_device_name(), torch.cuda.get_device_capability())
    x = torch.ones(4, device="cuda")
    print((x + x).cpu())        # tensor([2., 2., 2., 2.])
PY

Important

DGX Spark / GB10. Spark is an ARM64 machine with a Blackwell GPU of compute capability 12.1 (sm_121). Use a PyTorch wheel built for aarch64 with CUDA 13. --torch-backend=auto normally finds it. CPU and GPU share one pool of LPDDR5x memory (128 GB), which matters for the capacity planning in Chapters 18, 20 and 30.

4. The loop you will repeat every chapter

Run the first milestone’s tests. They fail, because you haven’t written anything yet:

pytest tests/test_ch03_autograd.py

E   NotImplementedError: TODO(Chapter 3): implement __add__ in engine/autograd.py

That message is your to-do list. Each chapter’s Build it section tells you which functions to implement and gives hints. You edit engine/<module>.py, rerun the tests, and repeat until they’re green:

pytest tests/test_ch03_autograd.py      # 5 passed
git commit -am "Chapter 3 milestone"

Then see your code do something:

python run.py autograd --impl engine

Tip

pytest -x stops at the first failure; pytest -k name runs only matching tests; pytest --pdb drops into a debugger where a test fails. The tests are short and readable, so open them: they are the specification.

Try it now with the reference, to confirm your environment works end to end:

IZH_IMPL=izh pytest -q        # ~126 passed, 12 skipped (the GPU-only tests) on a CPU
python run.py tensors

5. Packages you’ll add later

Only PyTorch is needed until Chapter 9. Later chapters ask you to install more, with the command shown at the point of use:

6. Models you’ll download later

The engine never downloads anything by itself. When a chapter needs a real checkpoint, it shows an explicit command like this one, which saves a pinned snapshot into models/:

uv pip install -r optional-requirements.txt
hf download Qwen/Qwen3-0.6B --local-dir models/Qwen3-0.6B

The optional large models are never required to pass a milestone. Every milestone is verified on small configurations, against independent reference implementations.

7. Optional: C++ and Rust

Every core mechanism in the book also exists in C++ and Rust, shown in tabs. The Rust track grows into a CPU engine that loads real Qwen3 weights. To follow along:

# C++17 compiler and CMake (Ubuntu: sudo apt install build-essential cmake)
cmake -S cpp -B build/cpp -DCMAKE_BUILD_TYPE=Release
cmake --build build/cpp -j
build/cpp/izh all          # self-checking demos, one per chapter

# Stable Rust from https://rustup.rs
cd rust
cargo test --release       # unit tests + parity with the Python reference
cargo run --release -- generate --model-dir ../models/Qwen3-0.6B --new-tokens 32

Appendix B describes both tracks and the optional NVIDIA Rust GPU toolchains.

When something goes wrong

Record uv pip freeze and nvidia-smi output alongside any result you want to keep. Next: the big picture.

1. The big picture: what happens when you press Enter

A model in ten lines

Let’s start with the destination and work backwards. Install the libraries that read real checkpoints, and download a small but genuinely capable model, Qwen3-0.6B (about 1.5 GB):

uv pip install -r optional-requirements.txt
hf download Qwen/Qwen3-0.6B --local-dir models/Qwen3-0.6B

Now ask it something. Create experiments/ch01.py in the code root (the directory containing run.py), save the following example there, and run python -m experiments.ch01 from that root. Use this same script for the measurements and stretch exercises below:

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

path = "models/Qwen3-0.6B"
tokenizer = AutoTokenizer.from_pretrained(path)
model = AutoModelForCausalLM.from_pretrained(path, dtype=torch.bfloat16).eval()

messages = [{"role": "user", "content": "Explain what a GPU is in one sentence."}]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
ids = tokenizer(text, return_tensors="pt").input_ids
out = model.generate(ids, max_new_tokens=40, do_sample=False)
print(tokenizer.decode(out[0, ids.shape[1]:], skip_special_tokens=True))

You’ll get an answer along these lines (exact wording varies with library version and hardware):

A GPU (Graphics Processing Unit) is a specialized processor designed to rapidly perform
many calculations in parallel, originally for rendering graphics.

That’s the whole user-facing experience. A library you didn’t write did the work: model.generate. By Chapter 18, that library will be replaced by your engine, and you’ll know what every line of it does and why. The rest of this chapter opens the box.

Tip

No download possible right now? Everything below also works with the tiny random model in the reference package: python run.py cache and python run.py fast exercise the same machinery. The text will be gibberish, but the mechanics are identical.

Step 1: text becomes token IDs

A model never sees characters. A tokenizer chops text into pieces from a fixed vocabulary and replaces each piece with its integer ID:

print(ids[0].tolist())
print([tokenizer.decode([i]) for i in ids[0]])

The output looks like this (middle IDs elided):

[151644, 872, 198, ..., 151645, 198, 151644, 77091, 198]
['<|im_start|>', 'user', '\n', 'Ex', 'plain', ' what', ' a', ' GPU', ' is', ..., '<|im_end|>', '\n', '<|im_start|>', 'assistant', '\n']

Notice three things. Common words are single tokens (' GPU', with its leading space). Rarer words may be split into pieces. And the chat template added special tokens like <|im_start|> (ID 151644) and <|im_end|> (ID 151645) that mark who is speaking. Qwen3’s vocabulary has 151,936 entries. You’ll build a tokenizer like this, trained on a short novel, in Chapter 4.

Step 2: IDs become vectors

The model’s first layer is a lookup table, the embedding, with one row of 1,024 numbers per vocabulary entry. Token 22670 (' GPU') becomes row 22670. The whole prompt becomes a matrix with one row per token: shape [tokens, 1024]. These rows are learned during training so that tokens used in similar ways end up with similar vectors.

Step 3: the vectors flow through layers

Qwen3-0.6B has 28 transformer layers. Each one does two things to every token vector:

1. Attention lets each token gather information from the tokens before it. That’s how ' sentence' comes to “know” it’s part of a request about GPUs. (Chapter 5)
2. The MLP transforms each token’s vector on its own, applying what the model learned. (Chapter 6)

Each layer adds its result to the vector it received (the residual stream), so information accumulates as it flows upward. After 28 layers, each position holds a vector that summarizes everything relevant so far.

Step 4: the last vector becomes a prediction

The final vector, at the last position, is multiplied by a [1024 × 151936] matrix (for Qwen3-0.6B, the same matrix as the embedding table, “tied”). That produces 151,936 scores called logits, one per vocabulary entry: how plausible each token is as the next one.

with torch.no_grad():
    logits = model(ids).logits[0, -1]          # scores for the token after the prompt
probs = torch.softmax(logits.float(), dim=-1)
top = probs.topk(5)
for p, i in zip(top.values, top.indices):
    print(f"{tokenizer.decode([i])!r:12} {p:.3f}")

A typical result (your probabilities will differ a little):

'A'          0.91
'GPU'        0.05
'The'        0.02
...

Softmax turned scores into probabilities that sum to 1. Here the model is quite sure the answer starts with “A”.

Step 5: choose, append, repeat

Sampling picks one token from that distribution. Always taking the most likely one is called greedy decoding. Then comes the trick that makes generation work: append the chosen token to the input and run the model again. Each step produces exactly one new token. A 40-token answer takes 40 trips through all 28 layers. This loop is the autoregressive generation loop, and it’s the heart of inference:

ids_so_far = ids
for _ in range(40):
    with torch.no_grad():
        logits = model(ids_so_far).logits[0, -1]
    next_id = logits.argmax()                       # greedy
    ids_so_far = torch.cat([ids_so_far, next_id.view(1, 1)], dim=1)
    if next_id == tokenizer.convert_tokens_to_ids("<|im_end|>"):
        break

Play with the whole pipeline here:

Two phases with very different costs

The loop above is wasteful. Every iteration re-processes the whole prompt, even though nothing about the earlier tokens has changed. Real engines remember each layer’s intermediate results for earlier tokens (the KV cache, Chapter 16), which splits generation into two phases:

• Prefill. Process all prompt tokens at once, filling the cache. Hundreds or thousands of tokens flow through each layer together, as one big matrix multiplication. GPUs are superb at this, and prefill is limited by arithmetic speed (it’s compute-bound).
• Decode. Produce one token at a time. Each step pushes a single token vector through every layer, reusing the cache. There’s little arithmetic per step, but every weight of the model must still be read from memory. Decode is memory-bound.

The user sees these as two numbers. Time to first token (TTFT) is mostly prefill. Inter-token latency, how fast the words stream, is decode.

The most important number in this book

Here is the back-of-envelope estimate that drives almost every design decision in modern inference engines.

To produce one token, decode multiplies one vector by every weight matrix in the model, so it must read every weight from memory once. Qwen3-0.6B has about 0.6 billion parameters. In BF16, each parameter takes 2 bytes, so each token requires reading about 1.2 GB.

How fast can memory deliver bytes? That’s the memory bandwidth of your device:

$$ \text{tokens per second} ;\lesssim; \frac{\text{memory bandwidth (bytes/s)}}{\text{bytes read per token}} $$

That’s a ceiling: real engines reach 60-90% of it. Compare it with the arithmetic. One token needs about 2 floating-point operations per parameter, roughly 1.2 GFLOP. Even a laptop CPU does that in a few milliseconds, and a GPU in microseconds. During decode, the processor mostly waits for memory.

This one inequality explains a remarkable amount of the field:

• Quantization (Chapter 20) stores weights in 4 bits instead of 16. That’s 4x fewer bytes per token, so decode gets close to 4x faster.
• Batching (Chapter 24) decodes many users’ tokens in the same step. The weights are read once and used for every user, so throughput rises almost for free.
• Speculative decoding (Chapter 26) checks several guessed tokens in one pass over the weights.
• Mixture of experts (Chapter 27) reads only a few experts per token. Flash-Next stores 176B parameters but reads about 6B per token.
• KV-cache size (Chapters 16, 25, 28) adds to the bytes per token, which is why long contexts are slow and why linear attention exists.

What an inference engine does

An inference engine is the program between “a folder of weights” and “a fast, correct stream of tokens for many users”. Its jobs, and where you’ll build each one:

Production engines like vLLM, SGLang and llama.cpp are hundreds of thousands of lines, mostly hardware-specific kernels and model adapters. The core ideas fit in a few thousand lines, and those are the lines you’ll write.

The destination: Qwen3.8-Flash-Next

The capstone model is a deliberate stretch. Qwen3.8-Flash-Next (released August 2026) is a preview of the architecture behind Qwen4. Its 48 layers mix:

• Gated DeltaNet linear attention in 36 layers, which keeps a fixed-size memory instead of a growing KV cache (Chapter 28);
• Qwen Sparse Attention in the other 12 layers, where a small indexer picks which 2,048 past tokens each query reads (Chapter 29);
• a 512-expert mixture in every layer, 10 experts per token plus a shared one (Chapter 27);
• four parallel residual streams with gated reads and writes, and a 51-billion-parameter hashed n-gram memory (Chapter 30).

Every one of those components is a variation on ideas you’ll learn first in their simplest form. By Chapter 30, you’ll write the whole model from scratch and verify it against the official implementation.

Build it

There’s no engine code this week. Instead, build your intuition with a measurement.

1. Time the Hugging Face generation above for 100 new tokens (call torch.cuda.synchronize() before stopping the clock if you’re on a GPU), and compute tokens per second.
2. Look up your device’s memory bandwidth and compute the ceiling from the formula.
3. Compute the ratio. Below 50% is normal for this naive loop. Chapter 19 gets you much closer to the ceiling.
4. Run the reference engine’s demos to see prefill and decode separately:

python run.py cache      # cached vs uncached generation time as output length grows
python run.py fast       # tokens/s of the plain loop vs a decoder without host round trips

Stretch exercises

1. ★ Change do_sample=False to do_sample=True, temperature=1.2. Run it three times. What changed, and what didn’t? Where: copy the example in A model in ten lines into experiments/ch01.py (create it); change its model.generate call.
2. ★★ Measure TTFT separately from decode speed by timing a 1-token generation and a 101-token one with a 500-token prompt. Which phase dominates for short answers? For long ones? Where: the same experiments/ch01.py (create it), around model.generate.
3. ★★ Load the model in FP32 (dtype=torch.float32) and repeat the speed measurement. Predict the slowdown from the bandwidth formula before you run it. Where: the same script’s AutoModelForCausalLM.from_pretrained call.

Check your understanding

1. Why does generating a 40-token answer require at least 40 sequential passes through the model?
2. Prefill of a 1,000-token prompt and decode of one token both read every weight once. Why is prefill not 1,000x slower than one decode step?
3. A 7B-parameter model in BF16 runs on a 1 TB/s GPU. What’s the most tokens per second a single user can get? What if the weights were 4-bit?

Going deeper

• BALLM Chapter 1 (pp. 1-16): what LLMs are and how they’re built and trained.
• PMPP §20.1 (pp. 478-482): the decoder-only transformer from a systems angle, including the generation loop.
• GPU Mode L1 (profiling and integrating CUDA kernels in PyTorch) for a first look at where the time goes.
• Horace He, Making Deep Learning Go Brrrr From First Principles (2022): the compute-bound / memory-bound / overhead-bound framing used throughout this book.

2. Tensors: the data structure of deep learning

Why start here

Every number a language model touches lives in a tensor: the weights you load from disk, the token vectors flowing through layers, the attention scores, the KV cache. Every bug you’ll hit while building an engine is, at bottom, a tensor bug: the wrong axis, the wrong layout, the wrong dtype, or the wrong device. An hour spent getting fluent here saves days later.

A tensor is an array with four properties

Run this in your environment (python inside the activated venv):

import torch
x = torch.arange(24.0).reshape(2, 3, 4)
print(x.shape, x.dtype, x.device, x.stride())

torch.Size([2, 3, 4]) torch.float32 cpu (12, 4, 1)

Four properties describe this tensor completely:

• Shape [2, 3, 4]: two blocks, each with three rows of four numbers. In language models the most common shape is [B, T, D]: a Batch of sequences, each with T tokens, each token a vector of D numbers. Our x could be two sentences of three tokens, each token a 4-number vector.
• Dtype float32: how each number is stored. FP32 takes 4 bytes, BF16 and FP16 take 2, INT8 takes 1. Token IDs are integers (int64 by default). The dtype decides memory size, and as Chapter 1 showed, memory size decides decode speed.
• Device cpu: where the data lives and where operations on it run. x.to("cuda") copies it to GPU memory. Operations between tensors on different devices are errors, not silent copies.
• Strides (12, 4, 1): how to find an element in memory. That needs its own section.

Note

Axis meaning lives in your program, not in the tensor. PyTorch doesn’t know that axis 1 is “tokens”. Writing shapes in comments, like # [B, T, D], is the single most effective habit for avoiding bugs in model code. This book does it everywhere.

Memory is flat: strides

RAM is one long line of bytes. A tensor stores its elements in a flat buffer, and the strides say how far to jump in that buffer to move one step along each axis. For our [2, 3, 4] tensor stored row by row, moving one step along the last axis moves 1 element, along the middle axis 4 elements (one row), and along the first 12 (one whole 3×4 block). The element at index (i, j, k) lives at

$$ \text{offset}(i, j, k) = \text{storage_offset} + 12i + 4j + 1k . $$

So x[1, 2, 3] is at 12 + 8 + 3 = 23, the last element. Strides like these, where the last axis moves fastest, are called contiguous or row-major.

Here’s why strides matter: many operations change only the strides, not the data. Transposing swaps two axes by swapping their sizes and strides. No numbers move:

t = x.transpose(1, 2)
print(t.shape, t.stride(), t.is_contiguous())

torch.Size([2, 4, 3]) (12, 1, 4) False

t is a view: another way of looking at the same memory. Views are free, which is why PyTorch uses them so much. The catch is that code which assumes row-major layout will read a view wrongly. A hand-written GPU kernel that computes row * width + col will happily read the wrong numbers from t, and the shapes will all look right. This exact bug shows up when loading checkpoints (Chapter 9) and writing kernels (Part III).

Some operations need contiguous memory. view refuses to reinterpret a non-contiguous tensor; reshape copies when it has to; contiguous() always produces a compact copy if needed:

t.view(2, 12)            # RuntimeError: view size is not compatible with input tensor's size and stride
t.reshape(2, 12)         # works: silently copies
t.contiguous().stride()  # (12, 3, 1): a fresh row-major copy

A copy costs a full read and write of the tensor. In the performance chapters you’ll learn to spot hidden .contiguous() copies in a profile.

The same idea in C++ and Rust: a tensor struct is just a buffer, a shape and strides, and a transpose swaps two entries:

def contiguous_strides(shape):
    strides, step = [], 1
    for size in reversed(shape):
        strides.append(step)
        step *= size
    return tuple(reversed(strides))

def element_offset(index, strides, storage_offset=0):
    return storage_offset + sum(i * s for i, s in zip(index, strides))

// A tensor is a flat buffer plus shape plus strides; element (i, j, ...) is at
// offset + i*stride[0] + j*stride[1] + ...
struct Tensor {
    std::vector<float> data;
    std::vector<size_t> shape, strides;
    size_t offset = 0;

    static std::vector<size_t> contiguous(const std::vector<size_t>& shape) {
        std::vector<size_t> s(shape.size(), 1);
        for (size_t a = shape.size(); a-- > 1;) s[a - 1] = s[a] * shape[a];
        return s;
    }
    Tensor(std::vector<float> d, std::vector<size_t> sh)
        : data(std::move(d)), shape(sh), strides(contiguous(sh)) {}
    float at(const std::vector<size_t>& index) const {
        size_t flat = offset;
        for (size_t a = 0; a < index.size(); ++a) flat += index[a] * strides[a];
        return data[flat];
    }
    Tensor transpose(size_t a, size_t b) const {    // swap sizes and strides; no data moves
        Tensor t = *this;
        std::swap(t.shape[a], t.shape[b]);
        std::swap(t.strides[a], t.strides[b]);
        return t;
    }
};

#![allow(unused)]
fn main() {
/// A view over a flat f32 buffer. Element (i0, i1, ...) lives at
/// offset + i0*strides[0] + i1*strides[1] + ...
#[derive(Clone, Debug)]
pub struct Tensor {
    pub data: Vec<f32>,
    pub shape: Vec<usize>,
    pub strides: Vec<usize>,
    pub offset: usize,
}

impl Tensor {
    /// Row-major ("contiguous") strides: the last axis moves fastest.
    pub fn contiguous_strides(shape: &[usize]) -> Vec<usize> {
        let mut strides = vec![1; shape.len()];
        for axis in (0..shape.len().saturating_sub(1)).rev() {
            strides[axis] = strides[axis + 1] * shape[axis + 1];
        }
        strides
    }

    pub fn from_vec(data: Vec<f32>, shape: &[usize]) -> Self {
        assert_eq!(data.len(), shape.iter().product::<usize>(), "data length must match the shape");
        Tensor { data, strides: Self::contiguous_strides(shape), shape: shape.to_vec(), offset: 0 }
    }

    pub fn get(&self, index: &[usize]) -> f32 {
        let flat: usize = index.iter().zip(&self.strides).map(|(i, s)| i * s).sum();
        self.data[self.offset + flat]
    }

    /// Swapping two axes swaps their sizes and strides. No data moves.
    pub fn transpose(&self, a: usize, b: usize) -> Self {
        let mut t = self.clone();
        t.shape.swap(a, b);
        t.strides.swap(a, b);
        t
    }

    pub fn is_contiguous(&self) -> bool {
        self.strides == Self::contiguous_strides(&self.shape)
    }
}
}

Elementwise operations and broadcasting

Arithmetic between same-shape tensors works element by element: a + b, a * b, torch.exp(a). More interesting is what happens when shapes differ. Broadcasting lets a smaller tensor be reused across a larger one without copying:

b = torch.tensor([10., 20., 30., 40.])     # shape [4]
print((x + b)[0, 0])                        # the same b added to every 4-vector in x

tensor([10., 21., 32., 43.])

The rule: align shapes from the right. Two sizes are compatible if they’re equal or one of them is 1; a missing axis counts as 1. The result takes the larger size on each axis. So [2, 3, 4] + [4] works (bias added to every token vector), and so does [3, 4] * [3, 1] (each row scaled by its own number):

s = torch.tensor([[1.], [2.], [3.]])       # shape [3, 1]
print(x[0] * s)                             # row i multiplied by s[i]

tensor([[ 0.,  1.,  2.,  3.],
        [ 8., 10., 12., 14.],
        [24., 27., 30., 33.]])

But [2, 3, 4] + [2, 3] fails: aligned from the right, 4 meets 3. If you meant “one number per token”, the scalar tensor must be [2, 3, 1]. Use s.unsqueeze(-1) or s[..., None] to add that axis. You’ll write this rule yourself in the milestone.

Warning

Broadcasting can also succeed when you didn’t intend it. [T, 1] - [T] gives a [T, T] matrix, not a [T] vector. If a loss value or an attention matrix suddenly has an unexpected extra axis, an accidental broadcast is the usual suspect. Assert shapes.

Reductions

A reduction collapses an axis: sum, mean, max, argmax, softmax’s denominator. You name the axis you want to collapse:

print(x.sum(dim=-1))        # sum each 4-vector: [2, 3, 4] -> [2, 3]
print(x.mean(dim=(0, 1)))   # average over batch and tokens: [2, 3, 4] -> [4]

tensor([[ 6., 22., 38.],
        [54., 70., 86.]])
tensor([10., 11., 12., 13.])

keepdim=True keeps the collapsed axis with size 1, which is exactly what you need to broadcast the result back: x - x.mean(-1, keepdim=True) centers every vector. That one line is the first half of LayerNorm (Chapter 6).

Matrix multiplication is many dot products

The dot product of two vectors multiplies them element by element and sums:

$$ [1, 2, 3] \cdot [4, 5, 6] = 1\cdot 4 + 2\cdot 5 + 3\cdot 6 = 32 . $$

A matrix product computes a dot product for every (row of A, column of B) pair:

$$ C_{ij} = \sum_{k} A_{ik} B_{kj}, \qquad [M, K] ;@; [K, N] ;\rightarrow; [M, N]. $$

The shared dimension K is summed away. That’s the shape rule to check first whenever a matmul fails. The work is $M \cdot N \cdot K$ multiply-adds, which we count as $2MNK$ floating-point operations (FLOPs). A transformer spends almost all of its arithmetic here.

def matmul_loops(a, b):
    m, k = a.shape
    k2, n = b.shape
    assert k == k2
    c = torch.zeros(m, n)
    for i in range(m):
        for j in range(n):
            c[i, j] = sum(a[i, p] * b[p, j] for p in range(k))
    return c

// C[M,N] = A[M,K] B[K,N], row-major. The i-k-j order streams rows of B and C sequentially.
inline std::vector<float> matmul(const std::vector<float>& a, const std::vector<float>& b, size_t m, size_t k, size_t n) {
    std::vector<float> c(m * n, 0.f);
    for (size_t i = 0; i < m; ++i)
        for (size_t p = 0; p < k; ++p) {
            float aip = a[i * k + p];
            for (size_t j = 0; j < n; ++j) c[i * n + j] += aip * b[p * n + j];
        }
    return c;
}

#![allow(unused)]
fn main() {
/// C[m][n] = sum_k A[m][k] * B[k][n] for row-major A [M,K] and B [K,N].
/// The i-k-j loop order walks B and C along rows, which keeps memory access sequential.
pub fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec<f32> {
    let mut c = vec![0.0f32; m * n];
    for i in 0..m {
        for p in 0..k {
            let a_ip = a[i * k + p];
            let b_row = &b[p * n..(p + 1) * n];
            let c_row = &mut c[i * n..(i + 1) * n];
            for (c_ij, b_pj) in c_row.iter_mut().zip(b_row) {
                *c_ij += a_ip * b_pj;
            }
        }
    }
    c
}
}

The Python triple loop is about a million times slower than a @ b. The C++ and Rust versions reorder the loops (i, then k, then j) so the inner loop walks memory sequentially. That’s your first glimpse of a theme that runs through Part III: the same arithmetic, arranged to respect memory, runs far faster.

Batched matmul and the linear layer

With more than two axes, @ treats the leading axes as a batch and multiplies the last two:

print((x @ torch.ones(4, 5)).shape)    # [2, 3, 4] @ [4, 5] -> [2, 3, 5]

torch.Size([2, 3, 5])

One weight matrix transforms every token vector in every sequence. That’s a linear layer, the most common operation in a transformer. One convention matters enormously when loading real weights: PyTorch’s nn.Linear(in, out) stores its weight as [out, in] and computes

$$ y = x,W^{\top} + b . $$

Some checkpoints (GPT-2’s original format, Chapter 9) store [in, out] instead. For square matrices both orientations have the same shape, so a missing transpose passes every shape check and produces garbage. You’ll meet this bug in Chapter 9 and will be glad you know about it.

How big is a tensor?

Bytes = number of elements × bytes per element:

print(torch.tensor(1.0).element_size(), torch.tensor(1.0, dtype=torch.bfloat16).element_size())  # 4 2

Qwen3-0.6B’s embedding table is [151936, 1024]: 155.6 million numbers, which is 311 MB in BF16 and 622 MB in FP32. Get in the habit of doing this multiplication for every big tensor you meet. Memory capacity decides what fits; memory bandwidth decides how fast it runs.

Devices and the GPU

x.to("cuda") copies a tensor into GPU memory and returns the copy; operations on it then run on the GPU. Two facts to keep in mind until Chapter 10 explains them properly:

1. GPU operations are asynchronous. y = x @ x on CUDA returns immediately; the GPU works in the background. Anything that needs the value on the CPU (print(y), y.item(), y.tolist()) waits. Naive timing code therefore measures nothing useful.
2. Copies between CPU and GPU are slow compared with GPU memory itself. Keep data on the device; move only small results back.

Important

On DGX Spark the CPU and GPU share physical memory, but PyTorch still treats them as separate devices with separate allocations. x.to("cuda") is still a copy, and .cpu() still synchronizes.

Build it

pytest tests/test_ch02_tensors.py

Tip

For broadcast_shapes, walk the axes from the right with an index i = 1, 2, ... and treat a missing axis as size 1. For linear, one line is enough: the leading axes broadcast through @ on their own.

Stretch exercises

1. ★ Predict, then check: the shapes of torch.randn(8, 1, 6, 1) + torch.randn(7, 1, 5), torch.randn(2, 3).T @ torch.randn(2, 3), and torch.randn(4, 5)[:, None, :] - torch.randn(4, 5)[None, :, :]. Where: experiments/ch02.py (create it).
2. ★★ Write transpose_strides(shape, strides, a, b) and slice_view(shape, strides, offset, axis, start, step) that return the new (shape, strides, offset) without touching data. Verify against x.transpose and x[:, 1::2]. Where: add both helpers to engine/tensors.py.
3. ★★ Time matmul_loops against @ for 64×64 matrices. How many times slower is it? Use torch.utils.benchmark.Timer. Where: experiments/ch02.py (create it), importing engine.tensors.matmul_loops.
4. ★★★ Implement matmul_loops for any batched shapes [..., M, K] @ [..., K, N] with broadcasting on the leading axes, using your broadcast_shapes. Where: extend matmul_loops in engine/tensors.py.

Check your understanding

1. Why does [2, 3, 4] @ [4, 5] produce [2, 3, 5]?
2. A tensor has shape [3, 2] and strides (1, 3). Is it contiguous? Which operation probably produced it?
3. Why can a missing transpose of a square weight matrix pass every shape check?
4. How many bytes does a [28, 2, 8, 4096, 128] BF16 tensor take? (It’s a KV cache you’ll meet in Chapter 16.)

Going deeper

• BALLM Appendix A, §§A.1-A.5 (pp. 251-271): PyTorch tensors and autograd for newcomers.
• PyTorch documentation, Tensor Views and Broadcasting semantics.
• Edward Yang, PyTorch internals (blog, 2019): how strides, storage and dispatch actually work inside PyTorch.

3. How networks learn: gradients, autograd, optimizers

Why an inference book teaches training

You’ll mostly run models in this book. But every weight you load was produced by the process in this chapter, and Parts II and V make you train and fine-tune models yourself. More immediately, the backward-pass machinery you build here is exactly what PyTorch runs when you call loss.backward(). Engine builders who know it can read memory reports, debug fine-tuning, and understand why training needs so much more memory than inference.

Learning is adjusting knobs to reduce a number

A neural network is a function with adjustable numbers inside it, called parameters or weights. Feed it an input, and it produces an output. A loss function turns “how wrong was that output?” into a single number. Learning means changing the parameters to make that number smaller.

The smallest possible example has one parameter w. It predicts w * x, and we want the prediction for x = 3 to be 8. The squared-error loss is

$$ L(w) = (w \cdot x - \text{target})^2 = (3w - 8)^2 . $$

At w = 2, the prediction is 6 and the loss is $(6-8)^2 = 4$. Should we increase or decrease w, and by how much?

The derivative says which way is downhill

The derivative $dL/dw$ is the slope of the loss as a function of w: how much the loss changes per small change in w. For our loss,

$$ \frac{dL}{dw} = 2x(wx - \text{target}) = 2 \cdot 3 \cdot (6 - 8) = -12 . $$

The slope is negative, so increasing w decreases the loss. Gradient descent takes a small step against the slope:

$$ w \leftarrow w - \eta \frac{dL}{dw} $$

where $\eta$ (eta) is the learning rate. With $\eta = 0.01$, the first steps are:

The loss falls, and the steps get smaller as the slope flattens near the minimum at w = 8/3. Try a learning rate twelve times larger ($\eta = 0.12$) and the loss grows: w goes 2.0 → 3.44 → 1.77 → 3.71 and the loss goes 4.0 → 5.38 → 7.24 → 9.75, because every step overshoots. Choosing the learning rate is the first hyperparameter you’ll tune.

Real networks have millions or billions of parameters. The gradient is the vector of all their partial derivatives, $\partial L / \partial w_i$, one per parameter, and gradient descent updates all of them at once. The idea is the same; only the bookkeeping grows.

The chain rule, and computing gradients mechanically

Nobody derives gradients for a billion-parameter model by hand. A network is built from simple operations (add, multiply, exp, …), and each one knows its own local derivative. The chain rule combines them: if $L$ depends on $d$, and $d$ depends on $a$, then

$$ \frac{\partial L}{\partial a} = \frac{\partial L}{\partial d}\cdot\frac{\partial d}{\partial a}. $$

Here’s a small expression worked through by hand. Let $a = 2$, $b = -3$, $c = 10$, $f = -2$, and compute

$$ e = a \cdot b = -6,\qquad d = e + c = 4,\qquad L = d \cdot f = -8 . $$

Work backwards from $L$, multiplying local derivatives as you go:

That procedure is backpropagation: a forward pass computes values and remembers how each was made, then a backward pass visits every operation once, in reverse order, multiplying the incoming gradient by the local derivative and passing it on. The cost of the backward pass is about twice the forward pass, regardless of the number of parameters. That’s why training billion-parameter models is possible at all.

Two details matter when you implement it:

1. Order. A node’s gradient is complete only after every node that uses it has passed its contribution back. Sorting the graph topologically and walking it in reverse guarantees this.
2. Accumulation. If a value is used twice (like x in x * x + x), contributions from each use add up. Gradients are accumulated with +=, never assigned with =.

Build an autograd engine

The Value class below wraps one number, remembers the values it was computed from, and stores a small closure that applies its local chain-rule step. This design follows Andrej Karpathy’s micrograd. Each operation creates the output, then defines how gradient flows from the output to its inputs:

class Value:
    def __init__(self, data, children=(), op=""):
        self.data, self.grad = float(data), 0.0
        self._children, self._op = tuple(children), op
        self._backward = lambda: None

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), "*")
        def backward():
            self.grad += other.data * out.grad     # d(ab)/da = b
            other.grad += self.data * out.grad     # d(ab)/db = a
        out._backward = backward
        return out

    def backward(self):
        order, seen = [], set()
        def visit(node):                           # topological sort
            if id(node) not in seen:
                seen.add(id(node))
                for child in node._children:
                    visit(child)
                order.append(node)
        visit(self)
        self.grad = 1.0
        for node in reversed(order):
            node._backward()

// A scalar autograd node. shared_ptr because one value can feed many later operations.
struct Node {
    double data, grad = 0;
    std::vector<std::shared_ptr<Node>> children;
    std::function<void(Node&)> backward = [](Node&) {};
    explicit Node(double d) : data(d) {}
};
using Value = std::shared_ptr<Node>;

inline Value make(double d) { return std::make_shared<Node>(d); }
inline Value add(Value a, Value b) {
    auto out = make(a->data + b->data);
    out->children = {a, b};
    out->backward = [a, b](Node& self) { a->grad += self.grad; b->grad += self.grad; };
    return out;
}
inline Value mul(Value a, Value b) {
    auto out = make(a->data * b->data);
    out->children = {a, b};
    out->backward = [a, b](Node& self) { a->grad += b->data * self.grad; b->grad += a->data * self.grad; };
    return out;
}
inline Value power(Value a, double n) {
    auto out = make(std::pow(a->data, n));
    out->children = {a};
    out->backward = [a, n](Node& self) { a->grad += n * std::pow(a->data, n - 1) * self.grad; };
    return out;
}
inline void backward(const Value& root) {
    std::vector<Value> order;
    std::set<Node*> seen;
    std::function<void(const Value&)> visit = [&](const Value& v) {
        if (!seen.insert(v.get()).second) return;
        for (auto& c : v->children) visit(c);
        order.push_back(v);
    };
    visit(root);
    root->grad = 1;
    for (auto it = order.rbegin(); it != order.rend(); ++it) (*it)->backward(**it);
}

#![allow(unused)]
fn main() {
#[derive(Clone)]
pub struct Value(Rc<RefCell<Node>>);

struct Node {
    data: f64,
    grad: f64,
    children: Vec<Value>,
    // Given this node's gradient, add the chain-rule contribution to each child.
    backward: Option<Box<dyn Fn(f64, &[Value])>>,
}

impl Value {
    pub fn new(data: f64) -> Self {
        Value(Rc::new(RefCell::new(Node { data, grad: 0.0, children: vec![], backward: None })))
    }
    fn op(data: f64, children: Vec<Value>, backward: impl Fn(f64, &[Value]) + 'static) -> Self {
        Value(Rc::new(RefCell::new(Node { data, grad: 0.0, children, backward: Some(Box::new(backward)) })))
    }
    pub fn data(&self) -> f64 { self.0.borrow().data }
    pub fn grad(&self) -> f64 { self.0.borrow().grad }
    fn add_grad(&self, g: f64) { self.0.borrow_mut().grad += g; }

    pub fn add(&self, other: &Value) -> Value {
        Value::op(self.data() + other.data(), vec![self.clone(), other.clone()], |g, c| {
            c[0].add_grad(g);
            c[1].add_grad(g);
        })
    }
    pub fn mul(&self, other: &Value) -> Value {
        Value::op(self.data() * other.data(), vec![self.clone(), other.clone()], |g, c| {
            let (a, b) = (c[0].data(), c[1].data());
            c[0].add_grad(b * g);
            c[1].add_grad(a * g);
        })
    }
    pub fn powf(&self, n: f64) -> Value {
        Value::op(self.data().powf(n), vec![self.clone()], move |g, c| {
            c[0].add_grad(n * c[0].data().powf(n - 1.0) * g);
        })
    }
    pub fn tanh(&self) -> Value {
        let t = self.data().tanh();
        Value::op(t, vec![self.clone()], move |g, c| c[0].add_grad((1.0 - t * t) * g))
    }

    /// Visit every node after all nodes that use it, then apply each local rule once.
    pub fn backward(&self) {
        let mut order = vec![];
        let mut seen = HashSet::new();
        fn visit(v: &Value, seen: &mut HashSet<usize>, order: &mut Vec<Value>) {
            if seen.insert(Rc::as_ptr(&v.0) as usize) {
                for child in &v.0.borrow().children {
                    visit(child, seen, order);
                }
                order.push(v.clone());
            }
        }
        visit(self, &mut seen, &mut order);
        self.0.borrow_mut().grad = 1.0;
        for v in order.iter().rev() {
            let node = v.0.borrow();
            if let Some(f) = &node.backward {
                f(node.grad, &node.children);
            }
        }
    }
}
}

With it, the worked example above becomes:

from engine.autograd import Value
a, b, c, f = Value(2.0), Value(-3.0), Value(10.0), Value(-2.0)
L = (a * b + c) * f
L.backward()
print(L.data, a.grad, b.grad, c.grad, f.grad)    # -8.0 6.0 -4.0 -2.0 4.0

Neurons, layers and why nonlinearity matters

A neuron computes a weighted sum of its inputs plus a bias, then applies a nonlinear function: $\tanh(w \cdot x + b)$. A layer is many neurons reading the same inputs, which is exactly the linear layer from Chapter 2 followed by an elementwise nonlinearity. A multilayer perceptron (MLP) stacks layers.

Why the nonlinearity? Without it, two linear layers collapse into one: $W_2(W_1 x) = (W_2 W_1)x$, a single matrix. You can stack a hundred linear layers and still only represent linear functions, which can’t even compute XOR. The nonlinearity between layers lets the network bend space. A classic demonstration is XOR, where the output is high when exactly one input is on:

No straight line separates the +1s from the −1s, but a 2-input MLP with two hidden tanh layers of 8 neurons learns it in a few dozen steps. This is the output of python run.py autograd --steps 200 with the reference engine:

{"step": 1, "loss": 5.63926}
{"step": 26, "loss": 2.49858}
{"step": 51, "loss": 0.00013}
{"step": 200, "loss": 0.0}
{"predictions": [-1.0, 1.0, 1.0, -1.0], "targets": [-1.0, 1.0, 1.0, -1.0]}

Transformer MLPs (Chapter 6) are exactly this structure, scaled up to thousands of neurons per layer, with smoother nonlinearities (GELU, SiLU) than tanh.

From scalars to tensors: PyTorch autograd

Your Value engine works on single numbers. PyTorch’s autograd does the same thing on whole tensors, so one node represents a million multiplications, and the local backward rules are themselves tensor operations running on the GPU. The interface is the same idea:

import torch
w = torch.tensor(2.0, requires_grad=True)   # track operations on w
loss = (w * 3.0 - 8.0) ** 2
loss.backward()                             # fills w.grad
print(w.grad)                               # tensor(-12.)

Three switches control it, and they’re easy to confuse:

• requires_grad=True on a tensor (parameters have it by default) means “record operations involving me”.
• torch.no_grad() and torch.inference_mode() stop recording. Every inference path in your engine runs under one of them. Recording costs memory, because each operation keeps its inputs alive for the backward pass.
• model.train() / model.eval() switch layer behavior (dropout on or off). They do not turn gradient recording on or off. That’s a common misconception.

Warning

Gradients accumulate across backward() calls, by design (it lets you sum gradients over several small batches). Forgetting optimizer.zero_grad() makes every step use the sum of all previous gradients. The loss usually explodes after a few steps.

The training loop

Every training script you’ll ever read has the same five lines at its core:

for batch in data:
    loss = loss_fn(model(batch.x), batch.y)   # 1. forward
    optimizer.zero_grad()                     # 2. clear old gradients
    loss.backward()                           # 3. backward: fill .grad for every parameter
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)   # 4. (optional) cap the gradient size
    optimizer.step()                          # 5. update the parameters

The optimizer decides how to turn gradients into parameter updates.

Optimizers: SGD, momentum, Adam, AdamW

SGD (stochastic gradient descent) is the update rule from earlier, applied to gradients estimated on a random batch of data. It works, but one learning rate must suit every parameter, and noisy gradients make it zigzag.

Momentum keeps a running average of recent gradients and steps along it, which smooths the noise:

$$ m \leftarrow \beta m + (1-\beta) g, \qquad w \leftarrow w - \eta, m . $$

Adam also tracks a running average of squared gradients, $v$, and divides each parameter’s step by $\sqrt{v}$. Parameters with consistently large gradients take smaller steps, and rarely updated parameters take relatively larger ones:

$$ m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t,\quad v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2,\quad w \leftarrow w - \eta,\frac{\hat m_t}{\sqrt{\hat v_t} + \epsilon} $$

where $\hat m_t = m_t/(1-\beta_1^t)$ and $\hat v_t = v_t/(1-\beta_2^t)$ correct for starting both averages at zero. Here’s a worked first step with $g_1 = 2$, $\beta_1 = 0.9$, $\beta_2 = 0.999$: $m_1 = 0.2$ and $v_1 = 0.004$, and after correction $\hat m = 2$ and $\hat v = 4$. The step is $\eta \cdot 2/\sqrt 4 = \eta$. Adam’s first step has size equal to the learning rate, whatever the gradient’s scale. (PyTorch agrees: with $\eta = 0.1$, one Adam step moves the weight from 0 to −0.1.)

AdamW adds weight decay, a gentle pull of every weight toward zero, applied separately from the adaptive gradient step: $w \leftarrow (1 - \eta\lambda)w - \eta,\hat m/(\sqrt{\hat v} + \epsilon)$. It’s the default optimizer for transformers.

Two facts matter for engine builders:

1. Optimizer state is big. AdamW stores $m$ and $v$ for every parameter, usually in FP32. Training a 0.6B model needs the weights (1.2 GB in BF16), plus gradients (another 1.2 GB), plus $m$ and $v$ (4.8 GB), plus FP32 master weights in mixed precision (2.4 GB), plus activations. Inference needs only the 1.2 GB. Chapter 22 shows how LoRA avoids most of that.
2. Schedules. Learning rates usually ramp up from near zero for the first few hundred steps (warmup), then decay (often along a cosine curve). Your training loop in Chapter 7 does both.

Build it

pytest tests/test_ch03_autograd.py
python run.py autograd --impl engine      # trains XOR with YOUR engine

Tip

If a gradient is wrong, compare it with the finite-difference estimate $\big(L(w+\epsilon) - L(w-\epsilon)\big)/2\epsilon$ (numerical_gradient in the module). It’s slow but independent of your backward code, which makes it the most useful debugging tool for gradients.

Stretch exercises

1. ★ Add sigmoid to Value. Its derivative is $\sigma(x)(1-\sigma(x))$. Verify it with numerical_gradient. Where: add Value.sigmoid in engine/autograd.py.
2. ★★ Train the XOR MLP with learning rates 0.005, 0.05 and 0.5. Plot the loss curves. Which converges, which is slow, and which oscillates? Where: experiments/ch03.py (create it), adapting run.py’s cmd_autograd.
3. ★★ Implement momentum and Adam for lists of Values and compare the steps-to-convergence on XOR with plain SGD. Where: add optimizer helpers beside sgd_step in engine/autograd.py; compare them in experiments/ch03.py (create it).
4. ★★★ Re-implement the XOR network in PyTorch with nn.Linear and torch.optim.AdamW. Then count the parameters, gradients and optimizer state in bytes for a 1B-parameter model trained this way. Where: experiments/ch03.py (create it).

Check your understanding

1. Why does gradient descent move against the gradient?
2. Why must gradients be accumulated with += in the backward pass?
3. What is lost if you remove every nonlinearity from an MLP?
4. Why is model.eval() not a substitute for torch.no_grad()?
5. Roughly how many bytes does AdamW training need per parameter, compared with BF16 inference?

Going deeper

• BALLM Appendix A §§A.3-A.7 (pp. 258-277): computation graphs, autograd and training loops in PyTorch. Appendix D (pp. 313-321): warmup, cosine decay and gradient clipping.
• Andrej Karpathy, The spelled-out intro to neural networks and backpropagation: building micrograd (video and code), the inspiration for this chapter’s engine.
• Kingma and Ba, Adam: A Method for Stochastic Optimization (2015); Loshchilov and Hutter, Decoupled Weight Decay Regularization (2019), the paper behind AdamW.
• GPU Mode L6 (Jane Xu, Optimizing PyTorch Optimizers): how optimizer steps are fused into a few kernels.

4. Text as numbers: tokenizers and embeddings

A model predicts IDs, not text

A language model is a function from a sequence of integers to a probability distribution over the next integer. A tokenizer is the contract that connects those integers to text. It splits text into units from a fixed vocabulary, maps each unit to an ID, and maps IDs back to text.

The contract is part of the model. A checkpoint’s embedding row 22670 was trained to mean whatever text the tokenizer maps to 22670. Feed it IDs from a different tokenizer and the model receives nonsense, even if every shape is right. When you load real models (Chapters 9 and 18), you’ll always use the checkpoint’s own tokenizer files.

Characters, words or something in between?

Characters are the simplest option: the vocabulary is every character in your training text, and each character gets an ID. It never meets an unknown word. But sequences get long. “tokenization” is 12 tokens, and attention’s cost grows with sequence length (Chapter 5).

Words give short sequences, but the vocabulary explodes (every name, typo, number and inflection is a separate entry) and any word not in it is unrepresentable.

Subwords split the difference: frequent words are single tokens (" the"), rare words are built from pieces (" token" "ization"). A good subword vocabulary of 50k-250k entries covers any text in a few tokens per word. The standard way to build one is byte-pair encoding.

Bytes first

Text in a computer is bytes. UTF-8 encodes ASCII characters as one byte each and other characters as two to four: "é" is the bytes 195 169, and "✓" is 226 156 147. If a tokenizer starts from the 256 possible byte values, every possible string is representable, with no unknown tokens ever. That’s byte-level BPE, used by GPT-2, Llama 3 and Qwen.

Byte-pair encoding, by hand

BPE builds the vocabulary bottom-up from a training text:

1. Start with the 256 single-byte tokens.
2. Count every adjacent pair of tokens in the training text.
3. Merge the most frequent pair into a new token everywhere it occurs, and record the merge.
4. Repeat until the vocabulary is as large as you want.

Here it is on a tiny corpus: “low” five times, “lower” twice, “newest” six times and “widest” three times. Before counting, a pre-tokenizer splits text into chunks (roughly: words with their leading space, punctuation and numbers), and pairs are never counted across chunk boundaries. These are the first eight merges the reference implementation learns:

Now encoding a word it never saw works by applying the learned merges, earliest first:

"lowest"  -> ['low', 'est']        IDs [259, 257]
" newer"  -> [' new', 'e', 'r']     IDs [263, 101, 114]
" lowly"  -> [' low', 'l', 'y']     IDs [260, 108, 121]

Unknown words fall back to smaller pieces, down to single bytes if necessary. The order matters: encoding must apply merges in the order they were learned (lowest new-ID first), because later merges were learned on text where earlier merges had already happened.

The two core helpers are short in every language:

def pair_counts(ids, counts=None):
    counts = Counter() if counts is None else counts
    for pair in zip(ids, ids[1:]):
        counts[pair] += 1
    return counts

def merge_pair(ids, pair, new_id):
    out, i = [], 0
    while i < len(ids):
        if i + 1 < len(ids) and (ids[i], ids[i + 1]) == pair:
            out.append(new_id)
            i += 2
        else:
            out.append(ids[i])
            i += 1
    return out

inline std::vector<int> merge_pair(const std::vector<int>& ids, std::pair<int, int> pair, int new_id) {
    std::vector<int> out;
    for (size_t i = 0; i < ids.size();) {
        if (i + 1 < ids.size() && ids[i] == pair.first && ids[i + 1] == pair.second) { out.push_back(new_id); i += 2; }
        else out.push_back(ids[i++]);
    }
    return out;
}
inline std::map<std::pair<int, int>, int> pair_counts(const std::vector<int>& ids) {
    std::map<std::pair<int, int>, int> counts;
    for (size_t i = 0; i + 1 < ids.size(); ++i) ++counts[{ids[i], ids[i + 1]}];
    return counts;
}

#![allow(unused)]
fn main() {
/// Count adjacent pairs, remembering the order in which each pair first appeared (for ties).
pub fn pair_counts(ids: &[u32], counts: &mut HashMap<(u32, u32), (usize, usize)>, weight: usize) {
    for w in ids.windows(2) {
        let next_rank = counts.len();
        let entry = counts.entry((w[0], w[1])).or_insert((0, next_rank));
        entry.0 += weight;
    }
}

/// Replace non-overlapping occurrences of `pair`, scanning left to right.
pub fn merge_pair(ids: &[u32], pair: (u32, u32), new_id: u32) -> Vec<u32> {
    let mut out = Vec::with_capacity(ids.len());
    let mut i = 0;
    while i < ids.len() {
        if i + 1 < ids.len() && (ids[i], ids[i + 1]) == pair {
            out.push(new_id);
            i += 2;
        } else {
            out.push(ids[i]);
            i += 1;
        }
    }
    out
}

/// Learn `num_merges` merges from whitespace-separated words (a simplified pre-tokenizer).
pub fn train(text: &str, num_merges: usize) -> Vec<((u32, u32), u32)> {
    let mut words: Vec<(Vec<u32>, usize)> = {
        let mut freq: Vec<(String, usize)> = vec![];
        for word in text.split_inclusive(' ') {
            match freq.iter_mut().find(|(w, _)| w == word) {
                Some(entry) => entry.1 += 1,
                None => freq.push((word.to_string(), 1)),
            }
        }
        freq.into_iter().map(|(w, n)| (w.bytes().map(u32::from).collect(), n)).collect()
    };
    let mut merges = vec![];
    for step in 0..num_merges {
        let mut counts = HashMap::new();
        for (ids, n) in &words {
            pair_counts(ids, &mut counts, *n);
        }
        // Most frequent pair; ties go to the pair seen first.
        let Some((&best, _)) = counts.iter().max_by(|a, b| a.1 .0.cmp(&b.1 .0).then(b.1 .1.cmp(&a.1 .1))) else { break };
        let new_id = 256 + step as u32;
        for (ids, _) in words.iter_mut() {
            *ids = merge_pair(ids, best, new_id);
        }
        merges.push((best, new_id));
    }
    merges
}
}

Training on a real text

The book’s training corpus is Edith Wharton’s 1908 short story The Verdict: 20,479 characters and 3,634 words, public domain, and the same text BALLM uses. Training a 400-token vocabulary on it takes under a second (python run.py bpe --vocab 400):

first merges: [' t', 'he', ' a', 'in', ' h', ' s', ' w', ' o', ' the', 'ou', 're', 'it']
last merges:  ['est', 'elf', 'ce', 'qu', ' sh', 'ind', ' Str', ' Stroud']

The first merges are the most common English letter pairs and " the". The last ones already include a character’s name, “Stroud”. With a 1,024-entry vocabulary, the whole story becomes 6,926 tokens instead of 20,479 bytes, about 3 bytes per token. Production tokenizers trained on trillions of tokens of diverse text reach about 4 bytes per token on English with 100-250k entries.

Note

Real tokenizers differ from ours in the pre-tokenizer. GPT-2 and Qwen split text with a regular expression that uses Unicode classes (\p{L} for letters, \p{N} for numbers), which Python’s built-in re module doesn’t support. Ours is a simplified ASCII version. Qwen also splits numbers into single digits, so that arithmetic sees consistent pieces. The merge algorithm itself is the same.

Special tokens and chat templates

Some IDs don’t represent text at all. They mark structure:

• End of text (<|endoftext|>) separates documents during training, so the model learns where text ends.
• Role markers in chat models, such as Qwen’s <|im_start|> and <|im_end|> (the “ChatML” format), mark who is speaking and where a turn ends.

A chat template turns a list of messages into the exact token sequence the model was fine-tuned on. Qwen3’s looks like this:

<|im_start|>user
Explain what a GPU is in one sentence.<|im_end|>
<|im_start|>assistant

The final <|im_start|>assistant\n is the generation prompt: it tells the model that it’s now the assistant’s turn. Omit it and the model may continue the user’s message instead. When the model emits <|im_end|>, its turn is over and your engine should stop. Many “the model rambles forever” bugs are a missing generation prompt or a missing stop token.

Warning

Special tokens must be inserted as IDs, not tokenized as text. If <|im_end|> is typed into a prompt as ordinary characters and tokenized byte by byte, the model sees <, |, im, … instead of its trained marker. Tokenizer libraries handle this with an “allowed special tokens” option; your BPE encode has an allow_special flag.

From a token stream to training examples

A language model learns to predict the next token at every position. From one stream of IDs, a sliding window of length context creates input/target pairs whose targets are the inputs shifted by one:

ids:     10 11 12 13 14 15 16 17 18 19
input:   10 11 12 13        target: 11 12 13 14
input:         12 13 14 15  target:       13 14 15 16      (stride 2: windows overlap)
input:               14 15 16 17  target:   15 16 17 18

One window gives context training examples at once: position 0 learns “after 10 comes 11”, position 1 learns “after 10 11 comes 12”, and so on. The stride controls overlap: stride = context gives disjoint windows, and a smaller stride reuses text.

Warning

Split before you window. If you make overlapping windows first and then split them randomly into training and validation sets, nearly identical windows land in both. Validation loss then measures memorization, not generalization. Split the token stream (or better, whole documents) first, then window each part. split_then_window does this.

Embeddings: from IDs to vectors

The model’s first operation turns each ID into a vector by looking up a row of an embedding table of shape [vocab_size, D]:

emb = torch.nn.Embedding(3, 2)
emb.weight.data = torch.tensor([[1., 0.], [0., 1.], [1., 1.]])
print(emb(torch.tensor([[2, 0]])))       # rows 2 and 0 -> [[[1., 1.], [1., 0.]]]

A lookup is mathematically the same as multiplying a one-hot vector (all zeros except a 1 at the ID) by the table, which is why embeddings train like any other weight. The lookup just skips the multiplication by zeros. Before training, the rows are random; training moves rows of tokens that behave alike closer together.

Two observations matter for engines:

• The table is big. Qwen3-0.6B’s is [151936, 1024]: 156M of the model’s 600M parameters. Many small models tie the output projection to this table (they reuse it to turn the final hidden vector back into vocabulary scores), saving another 156M parameters.
• A lookup reads one row per token, not the whole table. That’s why Flash-Next can afford a 51-billion-parameter hashed n-gram embedding (Chapter 30): each token touches only a few rows of it.

Position

An embedding lookup gives the same vector for " GPU" wherever it appears, so the model can’t tell “dog bites man” from “man bites dog”. GPT-2 adds a second learned table indexed by position: x[t] = token_embedding[id[t]] + position_embedding[t]. That table has a fixed number of rows, which caps the context length. Modern models use rotary position embeddings instead (Chapter 17), which encode position inside attention and generalize better to long sequences.

Build it

pytest tests/test_ch04_text.py
python run.py bpe --vocab 512 --impl engine

Tip

Training is fast if you count each distinct pre-tokenized chunk once, weighted by how often it occurs. “the” appears hundreds of times but needs to be merged only once. In _encode_chunk, the pair to merge next is the one with the lowest merge ID among pairs that have a merge.

Stretch exercises

1. ★ Encode "Hello, world!", "hello world" and "HELLO WORLD" with your 512-token tokenizer. Count the tokens. Why do the counts differ so much? Where: experiments/ch04.py (create it), importing engine.tokenizer.BPETokenizer.
2. ★★ Plot bytes-per-token on The Verdict for vocabulary sizes 300, 500, 1,000, 2,000 and 4,000. Where do the returns diminish? Where: experiments/ch04.py (create it), training on data/the-verdict.txt.
3. ★★ Install tiktoken and compare GPT-2’s tokenization of a paragraph of The Verdict with yours. Which words does GPT-2 keep whole that yours splits? Where: experiments/ch04.py (create it).
4. ★★★ Make _encode_chunk fast: instead of rescanning all pairs after each merge, keep a priority queue of mergeable pairs keyed by merge rank. Measure the speed-up on the whole story. Where: BPETokenizer._encode_chunk in engine/tokenizer.py.

Check your understanding

1. Why can a byte-level BPE tokenizer encode any string, while a word tokenizer cannot?
2. Why must merges be applied in the order they were learned?
3. What goes wrong if a chat prompt omits the final <|im_start|>assistant\n?
4. Why does splitting overlapping windows into train and validation sets inflate validation scores?
5. Why does an embedding table, by itself, carry no information about word order?

Going deeper

• BALLM Chapter 2 (pp. 17-49): tokenizing text, special tokens, byte-pair encoding with tiktoken, sliding-window data loading, token and position embeddings. Our corpus and window examples follow it.
• Sennrich, Haddow and Birch, Neural Machine Translation of Rare Words with Subword Units (2016): the paper that brought BPE to NLP.
• Andrej Karpathy, Let’s build the GPT Tokenizer (video) and the minbpe repository: byte-level BPE in depth, including the GPT-2 and GPT-4 regex pre-tokenizers.
• The Hugging Face tokenizers documentation, for how tokenizer.json stores pre-tokenizers, merges and special tokens.

5. Attention from first principles

The problem: context

After Chapter 4, each token is a vector that describes the token in isolation. But meaning depends on context. In “The animal didn’t cross the street because it was too tired”, the vector for “it” should end up carrying information about “animal”. Predicting the next token needs the whole preceding text, not just the last word.

Before 2017, the standard answer was the recurrent neural network (RNN). It reads tokens one at a time and squeezes everything seen so far into a single fixed-size vector, its hidden state. That has two problems. Information from far back must survive many squeezes and fades. And processing is inherently sequential, so training can’t use a GPU’s parallelism across positions.

Attention takes the opposite approach: when processing token $t$, look directly at every earlier token, decide how relevant each one is, and take a weighted mixture of their information. Nothing is squeezed, and all positions can be computed at once. That one idea, from Attention Is All You Need (Vaswani et al., 2017), is the core of every model in this book. (Recurrence comes back, in a much improved form, in Chapter 28.)

Step 1: attention with no parameters at all

We’ll follow BALLM’s example: six tokens, “Your journey starts with one step”, each already embedded as a 3-dimensional vector:

x = torch.tensor([[0.43, 0.15, 0.89],   # Your
                  [0.55, 0.87, 0.66],   # journey
                  [0.57, 0.85, 0.64],   # starts
                  [0.22, 0.58, 0.33],   # with
                  [0.77, 0.25, 0.10],   # one
                  [0.05, 0.80, 0.55]])  # step

Let’s compute an enriched, context-aware vector for “journey”. The recipe has three steps.

Score every token by its similarity to “journey”. The dot product is a natural similarity measure: it’s large when two vectors point the same way.

scores = x @ x[1] = [0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865]

Normalize the scores into weights that are positive and sum to 1, using softmax:

$$ w_j = \frac{e^{s_j}}{\sum_k e^{s_k}} \quad\Rightarrow\quad w = [0.1385,\ 0.2379,\ 0.2333,\ 0.1240,\ 0.1082,\ 0.1581]. $$

“journey” weighs itself highest, then “starts”, whose vector is nearly identical.

Mix: the new vector is the weighted sum of all token vectors:

$$ z_{\text{journey}} = \sum_j w_j, x_j = [0.4419,\ 0.6515,\ 0.5683]. $$

That’s attention. Doing it for every token at once is two matrix multiplications and a softmax:

$$ \text{weights} = \operatorname{softmax}_{\text{row}}(X X^{\top}), \qquad Z = \text{weights}; X . $$

Row $i$ of the [6, 6] weights matrix says how much token $i$ draws from each token $j$.

Step 2: learned queries, keys and values

Plain dot products only find vectors that are already similar. A model needs to learn what to look for, and that depends on the role a token plays. So attention gives each token three different learned views of itself, each produced by its own weight matrix:

• Query $q_i = x_i W_q$: what token $i$ is looking for.
• Key $k_j = x_j W_k$: what token $j$ offers, which is matched against queries.
• Value $v_j = x_j W_v$: what token $j$ hands over if it’s chosen.

$$ \text{scores} = Q K^{\top}, \qquad Z = \operatorname{softmax}(\text{scores}), V . $$

Separating keys from values is the important design choice. A token can be found for one reason and contribute something else. For “it”, the query might look for “animate noun earlier in the sentence”; the key of “animal” advertises exactly that; and its value passes along information useful for predicting what comes after “it”. These are intuitions, not literal labels. Training simply finds whatever projections lower the loss. The mechanism makes such behavior possible.

Step 3: scale the scores

Dot products of $d$-dimensional vectors with independent, unit-variance components have a standard deviation of about $\sqrt d$. Measured on random vectors:

Large scores make softmax saturate: softmax([1, 2, 3]) is [0.09, 0.24, 0.67], but softmax([8, 16, 24]) is [0.0000, 0.0003, 0.9997]. A saturated softmax puts all its weight on one token, and its gradient is nearly zero, so training stalls. Dividing scores by $\sqrt{d}$ keeps them in a healthy range at any width. This is scaled dot-product attention:

$$ \operatorname{Attention}(Q, K, V) = \operatorname{softmax}!\left(\frac{Q K^{\top}}{\sqrt{d}}\right) V . $$

Step 4: no peeking at the future

A language model learns to predict token $t+1$ from tokens $\le t$. If position $t$ could attend to position $t+1$ during training, it would simply copy the answer. The loss would drop to near zero and the model would learn nothing useful.

The fix is a causal mask: before the softmax, set the scores for future positions to $-\infty$. Since $e^{-\infty} = 0$, they get exactly zero weight, and the remaining weights still sum to 1. For our six tokens (using the plain dot products of Step 1):

             Your  journey starts  with   one   step
Your        1.000  0.000  0.000  0.000  0.000  0.000
journey     0.368  0.632  0.000  0.000  0.000  0.000
starts      0.228  0.389  0.382  0.000  0.000  0.000
with        0.205  0.296  0.292  0.208  0.000  0.000
one         0.175  0.225  0.227  0.157  0.216  0.000
step        0.139  0.218  0.213  0.142  0.099  0.190

The first token can only attend to itself. The last row is unchanged from the unmasked version, because nothing comes after “step”.

Warning

Mask before the softmax. Masking after it (zeroing weights of future tokens) leaves their scores in the denominator, so the remaining weights no longer sum to 1, and information about the future leaks through the normalization.

The best test of a mask doesn’t compare numbers with a library. It tests the rule directly: change only the last input token, and check that every earlier output stays exactly the same. Your milestone tests do this.

Multiple heads

One attention pattern per layer is limiting. A token may need to track its subject, its previous token and the most recent comma all at once. Multi-head attention runs $H$ smaller attentions in parallel, each with its own projections into a $D_h$-dimensional space (typically $D = H \cdot D_h$), then concatenates their outputs and mixes them with an output projection $W_o$.

In code, nobody runs $H$ separate projections. One big projection produces all heads’ queries at once, and a reshape splits them:

x: [B, T, D]
q = x @ Wq.T                     [B, T, H*Dh]
q.reshape(B, T, H, Dh)           [B, T, H, Dh]   split the last axis into heads
 .transpose(1, 2)                [B, H, T, Dh]   heads become a batch axis
scores = q @ k.transpose(-2, -1) [B, H, T, T]    one T×T matrix per head
y = weights @ v                  [B, H, T, Dh]
y.transpose(1, 2)                [B, T, H, Dh]
 .reshape(B, T, H*Dh)            [B, T, D]       heads concatenated per token
out = y @ Wo.T                   [B, T, D]

Warning

The transposes are not optional. reshape(B, H, T, Dh) directly on a [B, T, H*Dh] tensor produces the right shape with the wrong data: it interleaves tokens and heads. Nothing crashes, and the model just computes garbage. Your milestone test checks that head 1 of token 2 really holds features 4-8 of token 2.

Fewer key/value heads: grouped-query attention

Modern models often use fewer key/value heads than query heads. In grouped-query attention (GQA), each KV head is shared by a group of query heads: Qwen3-0.6B has 16 query heads and 8 KV heads, and Flash-Next’s attention layers have 24 query heads and just 2 KV heads. Each query head still computes its own attention pattern, but the K and V tensors are smaller. That matters enormously for the KV cache’s memory (Chapter 16). In code, query head $h$ reads KV head $h \div (H_q/H_{kv})$. The simplest implementation repeats each KV head to match the query heads before the matmul.

Positions, not just “the last T keys”

There’s one more design decision in your causal_attention, and it pays off for the rest of the book. Instead of building a fixed triangular mask, it takes the absolute position of every query and key and applies the rule

$$ \text{key } j \text{ is visible to query } i \iff \text{pos}(k_j) \le \text{pos}(q_i). $$

During ordinary training, queries and keys are both positions $0..T-1$, and this reproduces the triangle. But the same rule also handles, unchanged:

• cached decoding (Chapter 16): 1 new query at position 57 against 58 cached keys;
• chunked prefill: queries 40-47 against keys 0-47, a rectangular mask that a square triangle gets wrong;
• batches of different lengths (Chapters 19 and 24), where each row has its own positions;
• sparse attention (Chapter 29), via the extra allowed mask.

def causal_attention(q, k, v, query_positions=None, key_positions=None, allowed=None, scale=None):
    """Scaled dot-product attention with a causal rule.  (Your engine: Chapter 5)

    query_positions: [T] or [B, T]; defaults to the last T positions of the S keys.
    key_positions:   [S] or [B, S]; defaults to 0..S-1.
    allowed: optional extra boolean mask broadcastable to [B, Hq, T, S] (sliding windows,
             sparse selections). True means "may attend".
    Grouped-query attention: Hq must be a multiple of Hkv; each KV head serves Hq/Hkv query heads.
    """
    batch, q_heads, t, d = q.shape
    kv_heads, s = k.shape[1], k.shape[2]
    if q_heads % kv_heads:
        raise ValueError(f"Query heads ({q_heads}) must be a multiple of KV heads ({kv_heads})")
    if kv_heads != q_heads:
        k = k.repeat_interleave(q_heads // kv_heads, dim=1)
        v = v.repeat_interleave(q_heads // kv_heads, dim=1)
    scale = 1.0 / math.sqrt(d) if scale is None else scale
    scores = (q.float() @ k.float().transpose(-2, -1)) * scale          # [B, Hq, T, S]
    qp = _positions(query_positions, batch, t, s - t, q.device)          # [B, T]
    kp = _positions(key_positions, batch, s, 0, q.device)                # [B, S]
    mask = kp[:, None, None, :] <= qp[:, None, :, None]                  # [B, 1, T, S]
    if allowed is not None:
        mask = mask & allowed
    scores = scores.masked_fill(~mask, float("-inf"))
    weights = torch.softmax(scores, dim=-1)
    return (weights @ v.float()).to(q.dtype)

// q [T, Hq*D] for positions S-T..S-1; k, v [S, Hkv*D]. Query head h reads KV head h/(Hq/Hkv).
inline std::vector<float> causal_attention(const std::vector<float>& q, const std::vector<float>& k, const std::vector<float>& v,
                                           size_t T, size_t S, size_t hq, size_t hkv, size_t d) {
    std::vector<float> out(T * hq * d, 0.f), scores(S);
    float scale = 1.f / std::sqrt(float(d));
    for (size_t i = 0; i < T; ++i)
        for (size_t h = 0; h < hq; ++h) {
            size_t kvh = h / (hq / hkv), visible = S - T + i + 1;
            const float* qv = &q[(i * hq + h) * d];
            for (size_t j = 0; j < visible; ++j) {
                const float* kv = &k[(j * hkv + kvh) * d];
                scores[j] = std::inner_product(qv, qv + d, kv, 0.f) * scale;
            }
            softmax(scores.data(), visible);
            for (size_t j = 0; j < visible; ++j)
                for (size_t e = 0; e < d; ++e) out[(i * hq + h) * d + e] += scores[j] * v[(j * hkv + kvh) * d + e];
        }
    return out;
}

#![allow(unused)]
fn main() {
/// q: [T, Hq*D] for the new tokens at positions start..start+T; k, v: [S, Hkv*D] for every
/// cached position 0..S. Each query reads only keys at positions <= its own (causality).
/// Query head h reads KV head h / (Hq/Hkv): grouped-query attention.
pub fn causal_attention(q: &[f32], k: &[f32], v: &[f32], t: usize, s: usize,
                        q_heads: usize, kv_heads: usize, d: usize) -> Vec<f32> {
    let start = s - t;
    let group = q_heads / kv_heads;
    let scale = 1.0 / (d as f32).sqrt();
    let mut out = vec![0.0; t * q_heads * d];
    let mut scores = vec![0.0; s];
    for i in 0..t {
        let visible = start + i + 1; // keys 0..=start+i
        for h in 0..q_heads {
            let kvh = h / group;
            let qv = &q[(i * q_heads + h) * d..][..d];
            for j in 0..visible {
                let kv = &k[(j * kv_heads + kvh) * d..][..d];
                scores[j] = qv.iter().zip(kv).map(|(a, b)| a * b).sum::<f32>() * scale;
            }
            softmax(&mut scores[..visible]);
            let o = &mut out[(i * q_heads + h) * d..][..d];
            for j in 0..visible {
                let vv = &v[(j * kv_heads + kvh) * d..][..d];
                for (oo, vvv) in o.iter_mut().zip(vv) {
                    *oo += scores[j] * vvv;
                }
            }
        }
    }
    out
}
}

The C++ and Rust versions process one sequence with explicit loops and no batch axis. Read them to see the computation without tensor broadcasting. They implement the same rule: query $i$ of $T$ new tokens sees keys $0 \ldots S-T+i$.

What attention costs

For a sequence of $T$ tokens with model width $D$, the score matrix has $T^2$ entries per head, and computing scores and outputs costs about $4T^2D$ FLOPs per layer. At $T = 8{,}192$ and FP32, one head’s score matrix alone is 256 MiB. This quadratic growth is attention’s central problem at long context. FlashAttention (Chapter 15) removes the memory cost without changing the result. Linear attention (Chapter 28) and sparse attention (Chapter 29) change the computation itself.

Build it

pytest tests/test_ch05_attention.py
python run.py attention --impl engine

Tip

Compute scores in FP32 (q.float() @ k.float().transpose(-2, -1)) and cast the result back to q.dtype. Build the mask by broadcasting positions: key_pos[:, None, None, :] <= query_pos[:, None, :, None] has shape [B, 1, T, S] and broadcasts over heads. Use masked_fill(~mask, float("-inf")) before torch.softmax.

Stretch exercises

1. ★ Recompute this chapter’s six-token causal weight table with your function (use x as q, k and v with a single head, and scale=1.0). Where: experiments/ch05.py (create it), importing engine.attention.causal_attention.
2. ★★ Write a test that catches the “reshape without transpose” bug: build q so that each head’s features are distinguishable and assert that merge_heads(split_heads(x)) equals x, but that a wrong reshape doesn’t. Where: create tests/test_ch05_stretch.py, importing engine.attention.
3. ★★ Implement a sliding-window variant using the allowed argument: each query sees only the last $w$ keys. (You’ll meet this again in Chapter 29.) Where: add a window-mask helper to engine/attention.py and pass its result to causal_attention(allowed=...).
4. ★★★ Measure the runtime and peak memory of causal_attention for $T$ = 512, 1,024, 2,048 and 4,096. Confirm the quadratic growth and estimate when it stops fitting in your device’s memory. Where: experiments/ch05.py (create it), importing engine.attention.causal_attention.

Check your understanding

1. Why does attention normalize over keys (each row) rather than over queries (each column)?
2. Why mask before the softmax?
3. Which inputs can affect the output at position 0 of a causal attention layer?
4. What does dividing by $\sqrt{d}$ fix, and what goes wrong without it?
5. With 16 query heads and 8 KV heads, which KV head does query head 11 read?

Going deeper

• BALLM Chapter 3 (pp. 50-91): attention from simplified weights to causal multi-head attention. This chapter’s worked example comes from §3.3.
• PMPP §20.2-20.3 (pp. 482-488): multi-head attention as matrix operations, and a first CUDA implementation.
• Vaswani et al., Attention Is All You Need (2017); Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (2023).
• Jay Alammar, The Illustrated Transformer: the classic visual walkthrough.

6. The transformer block and a complete GPT

The shape of a GPT

A GPT (“generative pre-trained transformer”) is a stack of identical blocks between an input embedding and an output head:

token IDs [B, T]
  │ token embedding + position embedding
  ▼
x [B, T, D] ──► block 1 ──► block 2 ──► ... ──► block L ──► final LayerNorm ──► head ──► logits [B, T, V]

Every block has the same shape in and out, [B, T, D], so blocks stack like Lego. GPT-2 small has $L = 12$ blocks of width $D = 768$; Qwen3-0.6B has 28 blocks of width 1,024. Once you’ve built one block, you’ve built the model.

The residual stream

Each block contains two branches, attention and an MLP. Each branch reads the current vector, computes something, and adds its result back:

$$ \begin{aligned} a &= x + \operatorname{Attention}(\operatorname{LN}_1(x)) \ y &= a + \operatorname{MLP}(\operatorname{LN}_2(a)) \end{aligned} $$

Picture the vector $x$ as a residual stream that flows from the embedding to the output. Each branch reads from it and writes a small update into it. This has three big consequences:

1. Training works at depth. The gradient of $x + f(x)$ with respect to $x$ is $1 + f’(x)$. Even if a branch’s gradient is tiny, the “1” carries the signal straight back to early layers. Before residual connections (He et al., 2015), very deep networks barely trained.
2. A block can do nothing. If both branches output zero, the block is the identity. A new block starts as a small perturbation of a working network, not a random scrambling of it. Your milestone test checks exactly this: zero the output projections and the block must return its input.
3. Interpretability. Every branch reads and writes the same space, so you can inspect, edit or steer what’s in the stream. Chapter 23 does exactly that.

This “pre-norm” arrangement (normalize the input of each branch) is used by GPT-2 and every modern model. The original 2017 transformer normalized after the addition (post-norm), which is harder to train at depth.

LayerNorm: keep the numbers in range

As updates accumulate in the stream, its vectors can drift to very different scales from token to token. Layer normalization rescales each token vector, independently of every other token and every other example, to mean 0 and variance 1, then applies a learned per-feature scale $\gamma$ and shift $\beta$:

$$ \operatorname{LN}(x)_i = \gamma_i ,\frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta_i, \qquad \mu = \frac1D\sum_i x_i,\quad \sigma^2 = \frac1D\sum_i (x_i-\mu)^2 . $$

For $x = [1, 2, 3]$: $\mu = 2$, $\sigma^2 = 2/3$, and the normalized vector is $[-1.2247, 0, 1.2247]$ (before $\gamma$ and $\beta$). Note that the variance divides by $D$, not $D-1$; the “sample variance” correction from statistics would be a different operation and would break parity with real checkpoints. The small $\epsilon$ (1e-5 in GPT-2) prevents division by zero.

def layernorm(x, gamma, beta, eps=1e-5):          # x [..., D]
    mean = x.mean(-1, keepdim=True)
    var = x.var(-1, keepdim=True, unbiased=False)  # divide by D, not D-1
    return gamma * (x - mean) / torch.sqrt(var + eps) + beta

inline std::vector<float> layernorm(const std::vector<float>& x, const std::vector<float>& g, const std::vector<float>& b, float eps) {
    float mean = std::accumulate(x.begin(), x.end(), 0.f) / x.size(), var = 0;
    for (float v : x) var += (v - mean) * (v - mean);
    float inv = 1.f / std::sqrt(var / x.size() + eps);
    std::vector<float> y(x.size());
    for (size_t i = 0; i < x.size(); ++i) y[i] = (x[i] - mean) * inv * g[i] + b[i];
    return y;
}
inline std::vector<float> rmsnorm(const std::vector<float>& x, const std::vector<float>& w, float eps) {
    float ms = 0;
    for (float v : x) ms += v * v;
    float inv = 1.f / std::sqrt(ms / x.size() + eps);
    std::vector<float> y(x.size());
    for (size_t i = 0; i < x.size(); ++i) y[i] = x[i] * inv * w[i];
    return y;
}

#![allow(unused)]
fn main() {
/// LayerNorm over one token vector: subtract the mean, divide by the standard deviation.
pub fn layernorm(x: &[f32], gamma: &[f32], beta: &[f32], eps: f32) -> Vec<f32> {
    let n = x.len() as f32;
    let mean = x.iter().sum::<f32>() / n;
    let var = x.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / n;
    let inv = 1.0 / (var + eps).sqrt();
    x.iter().zip(gamma).zip(beta).map(|((v, g), b)| (v - mean) * inv * g + b).collect()
}

/// RMSNorm: no mean subtraction, divide by the root mean square.
pub fn rmsnorm(x: &[f32], weight: &[f32], eps: f32) -> Vec<f32> {
    let ms = x.iter().map(|v| v * v).sum::<f32>() / x.len() as f32;
    let inv = 1.0 / (ms + eps).sqrt();
    x.iter().zip(weight).map(|(v, w)| v * inv * w).collect()
}
}

(The C++ and Rust tabs also include RMSNorm, the simpler variant Qwen uses. It’s Chapter 17’s topic.)

The MLP: where most parameters live

Attention moves information between positions. The MLP (also called the feed-forward network) then processes each position on its own: expand to $4D$, apply a nonlinearity, contract back to $D$:

$$ \operatorname{MLP}(x) = W_{\text{down}},\operatorname{GELU}(W_{\text{up}},x + b_{\text{up}}) + b_{\text{down}} . $$

GPT-2 uses GELU, a smooth relative of ReLU. It’s close to 0 for very negative inputs, close to $x$ for positive ones, and smoothly curved in between. GPT-2 uses the tanh approximation:

$$ \operatorname{GELU}(x) \approx \tfrac12 x\left(1 + \tanh!\left(\sqrt{2/\pi},(x + 0.044715x^3)\right)\right). $$

The exact error-function version differs slightly. Using the wrong one is a small but real source of mismatch when loading checkpoints, so match the checkpoint’s convention (approximate="tanh" for GPT-2).

The MLP holds two thirds of each block’s weights ($8D^2$ versus attention’s $4D^2$). A useful mental model, explored in interpretability research, treats it as a key-value memory: the first matrix detects patterns, and the second writes associated information into the stream.

The block in code

Here is the reference block. Attention packs Q, K and V into one linear layer (qkv, width $3D$), which is a storage choice; GPT-2’s checkpoint stores it that way too. Then come the residual additions you just saw:

class GPTAttention(nn.Module):
    """Multi-head causal self-attention with one packed QKV projection."""

    def __init__(self, cfg):
        """Create qkv: width -> 3*width and proj: width -> width, both with bias.  (Your engine: Chapter 6)"""
        super().__init__()
        if cfg.width % cfg.heads:
            raise ValueError("width must be divisible by heads")
        self.heads = cfg.heads
        self.qkv = nn.Linear(cfg.width, 3 * cfg.width)
        self.proj = nn.Linear(cfg.width, cfg.width)

    def forward(self, x, positions, cache=None, layer=0, rows=None):
        """Project, split heads, (cache), attend, merge heads, project.  (Your engine: Chapter 6)"""
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = (split_heads(t, self.heads) for t in (q, k, v))
        key_positions = None
        if cache is not None:
            k, v, key_positions = cache.update(layer, k, v, positions, rows)
        y = causal_attention(q, k, v, positions, key_positions)
        return self.proj(merge_heads(y))

class GPTBlock(nn.Module):
    def __init__(self, cfg):
        """Two LayerNorms, attention, and a 4x-wide GELU MLP.  (Your engine: Chapter 6)"""
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.width, eps=cfg.eps)
        self.attn = GPTAttention(cfg)
        self.ln2 = nn.LayerNorm(cfg.width, eps=cfg.eps)
        self.up = nn.Linear(cfg.width, 4 * cfg.width)
        self.down = nn.Linear(4 * cfg.width, cfg.width)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x, positions, cache=None, layer=0, rows=None):
        """Pre-norm residual block: x + attn(ln1(x)), then + mlp(ln2(.)).  (Your engine: Chapter 6)"""
        x = x + self.drop(self.attn(self.ln1(x), positions, cache, layer, rows))
        hidden = nn.functional.gelu(self.up(self.ln2(x)), approximate="tanh")
        return x + self.drop(self.down(hidden))

The positions, cache and rows arguments let the same block serve generation with a KV cache later (Chapter 16) and batched serving (Chapter 24). For now they’re just passed through. cache is None and positions are $0..T-1$.

Assembling the model

class GPT(nn.Module):
    def __init__(self, cfg):
        """Token and position embeddings, a stack of blocks, a final LayerNorm and a vocabulary head.  (Your engine: Chapter 6)"""
        super().__init__()
        self.cfg = cfg
        self.token = nn.Embedding(cfg.vocab, cfg.width)
        self.position = nn.Embedding(cfg.context, cfg.width)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList(GPTBlock(cfg) for _ in range(cfg.layers))
        self.norm = nn.LayerNorm(cfg.width, eps=cfg.eps)
        self.head = nn.Linear(cfg.width, cfg.vocab, bias=False)
        if cfg.tie_weights:
            self.head.weight = self.token.weight
        self.apply(self._initialize)

    @staticmethod
    def _initialize(module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, std=0.02)
        if isinstance(module, nn.Linear) and module.bias is not None:
            nn.init.zeros_(module.bias)

    def forward(self, ids, cache=None, positions=None, rows=None):
        """ids [B, T] -> logits [B, T, vocab].  (Your engine: Chapter 6)

        Positions default to 0..T-1, or continue after the cache's current length.
        """
        if positions is None:
            start = cache.length if cache is not None else 0
            if start + ids.shape[1] > self.cfg.context:
                raise ValueError(f"Sequence exceeds the model's context of {self.cfg.context}")
            positions = torch.arange(start, start + ids.shape[1], device=ids.device)
        # Explicit position tensors are trusted: checking them would force a GPU->CPU sync.
        x = self.drop(self.token(ids) + self.position(positions))
        for i, block in enumerate(self.blocks):
            x = block(x, positions, cache, i, rows)
        return self.head(self.norm(x))

    @property
    def context_limit(self):
        return self.cfg.context

    def cache_spec(self):
        """(layers, kv_heads, head_dim): what every cache needs to know about this model."""
        return self.cfg.layers, self.cfg.heads, self.cfg.width // self.cfg.heads

    def new_cache(self, batch, capacity):
        if capacity > self.context_limit:
            raise ValueError("Requested cache exceeds the context limit")
        p = self.token.weight
        layers, kv_heads, head_dim = self.cache_spec()
        return KVCache(layers, batch, kv_heads, capacity, head_dim, p.device, p.dtype)

Four details deserve attention:

• Two embeddings are added: token identity plus learned position. The position table has context rows, which hard-limits the sequence length.
• The head produces logits, one score per vocabulary entry at every position: shape [B, T, V]. Position $t$’s logits predict token $t+1$.
• Weight tying: self.head.weight = self.token.weight makes the output head the same tensor as the embedding table, read in reverse. It saves $V \cdot D$ parameters (38.6M for GPT-2) and was used by GPT-2. The parameter count must count it once.
• Initialization: weights start as small random normals (std 0.02, as in GPT-2) and biases at zero. Getting this wrong produces a model that trains slowly or not at all, but it doesn’t matter for loading pretrained weights, which overwrite everything.

Counting parameters exactly

For a block with biases and two LayerNorms:

The whole tied model, with vocabulary $V$ and context $C$:

$$ N = VD + CD + L(12D^2 + 13D) + 2D . $$

For GPT-2 small ($V = 50{,}257$, $C = 1{,}024$, $D = 768$, $L = 12$), that’s 124,439,808, the “124M” in its name. python run.py gpt prints this alongside the count from your model, and your milestone tests that the two agree.

Notice how much of a small model is embedding: 31% of GPT-2 small. As models grow, the $12LD^2$ term dominates, which is why “parameters” and “FLOPs per token” are both close to $12LD^2$ for large models.

A random model already “works”

Before training, your GPT runs end to end and produces logits of the right shape, and the output is noise. That’s useful, not useless: plumbing correctness and learned behavior are separate claims, and you can test the first without the second. The milestone tests check shapes, causality (changing the last token never changes earlier logits), the parameter formula, weight tying, and the identity behavior of zeroed branches. None of that needs training.

Build it

pytest tests/test_ch06_gpt.py
python run.py gpt --impl engine

Tip

Keep the exact attribute names of the reference (token, position, blocks, norm, head, qkv, proj, ln1, ln2, up, down). Chapter 9’s GPT-2 loader maps checkpoint tensors onto them. If you name things differently, update the loader’s mapping.

Stretch exercises

1. ★ Compute by hand the parameter count of GPT-2 medium ($D = 1024$, $L = 24$, 16 heads) and check it with gpt_parameter_formula. Is it really “355M”? Where: paper, then experiments/ch06.py (create it) using engine.gpt.gpt_parameter_formula.
2. ★★ Add a return_hidden option that returns the residual stream after every block. Plot each block’s vector norm for a random input. How does it grow with depth? Where: GPT.forward in engine/gpt.py.
3. ★★ Convert the block to post-norm (normalize after each addition) and verify that zeroed branches no longer give an identity block. Why? Where: GPTBlock.forward in engine/gpt.py; keep a separate post-norm variant for comparison.
4. ★★★ Implement the exact GELU and measure the maximum difference from the tanh version over [−6, 6]. Then estimate how much it would change GPT-2’s logits. Where: experiments/ch06.py (create it) for the numerical comparison; change the GELU in GPTBlock.forward in engine/gpt.py for the logit comparison.

Check your understanding

1. Which axis does the MLP mix, and which does attention mix?
2. Why does a pre-norm block with both branch outputs at zero leave its input unchanged?
3. What changes, in storage and in training, when the head is tied to the token embedding?
4. Why is LayerNorm’s output not always mean-zero and unit-variance?
5. Where do the $12D^2$ parameters of a block come from?

Going deeper

• BALLM Chapter 4 (pp. 92-127): LayerNorm, GELU, shortcut connections and the GPT model, built up in the same order. Its parameter-counting exercise matches this chapter’s formula.
• PMPP §20.1 (pp. 478-482): the decoder block as a sequence of GEMMs and elementwise operations.
• Radford et al., Language Models are Unsupervised Multitask Learners (GPT-2, 2019); He et al., Deep Residual Learning (2015); Ba et al., Layer Normalization (2016).
• Anthropic, A Mathematical Framework for Transformer Circuits (2021): the residual stream view of transformers.

7. Train and evaluate your GPT

The objective: predict the next token

Chapter 4 turned text into windows of inputs and targets shifted by one. For every position $t$ in every window, the model outputs logits $z$, a score per vocabulary entry, and we want the probability of the true next token $y$ to be high. Softmax turns logits into probabilities, and the cross-entropy loss is the negative log of the probability assigned to the correct answer:

$$ p(y) = \frac{e^{z_y}}{\sum_j e^{z_j}}, \qquad L = -\log p(y). $$

The numbers are worth internalizing:

The last row is the most useful number in a training log. An untrained model should start at about $\log V$. If your first loss is 30, the initialization or the loss computation is broken. If it’s 0.5, the model can see its targets.

The gradient has a beautiful form

The derivative of cross-entropy with respect to each logit is

$$ \frac{\partial L}{\partial z_j} = p_j - \mathbb 1[j = y]. $$

If the target has probability 0.2, its logit’s gradient is −0.8, so gradient descent pushes it up. A wrong token with probability 0.3 gets +0.3 and is pushed down. Every logit is nudged in proportion to how wrong its probability is. This gradient then flows back through the head, every block, and the embeddings, via the chain rule from Chapter 3.

Use logits, not probabilities

PyTorch’s F.cross_entropy takes logits and computes a numerically stable log-softmax internally. Passing it probabilities (applying softmax yourself first) silently computes a different, wrong loss. It also wants 2-D logits and 1-D targets, so flatten batch and time together:

def lm_loss(model, x, y):
    """Mean next-token cross-entropy over every position.  (Your engine: Chapter 7)

    logits [B, T, V] and targets [B, T] are flattened to [B*T, V] and [B*T]. cross_entropy
    applies a stable log-softmax itself, so it takes raw logits, never probabilities.
    """
    logits = model(x)
    return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1))

For a validation set, average by tokens, not by batches. If batches have different sizes, a plain average of batch means over-weights the small ones:

@torch.no_grad()
def evaluate(model, x, y, batch_size=32):
    """Token-weighted mean loss over a whole split, in eval mode, without gradients.  (Your engine: Chapter 7)"""
    was_training = model.training
    model.eval()
    total, count = 0.0, 0
    for start in range(0, len(x), batch_size):
        xb, yb = x[start:start + batch_size], y[start:start + batch_size]
        total += lm_loss(model, xb, yb).item() * yb.numel()
        count += yb.numel()
    model.train(was_training)
    return total / count

The training loop

Here’s the full loop you’ll write. It’s the five-line core from Chapter 3, plus the details that make training stable:

def train(model, train_xy, valid_xy, steps, batch_size=16, lr=3e-3, weight_decay=0.1,
          warmup=20, clip=1.0, eval_every=50, seed=0, log=print):
    """AdamW with linear warmup then cosine decay, gradient clipping and periodic evaluation.  (Your engine: Chapter 7)

    Returns a list of {"step", "train_loss", "valid_loss", "lr"} records.
    """
    device = next(model.parameters()).device
    tx, ty = (t.to(device) for t in train_xy)
    vx, vy = (t.to(device) for t in valid_xy)
    generator = torch.Generator().manual_seed(seed)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay, betas=(0.9, 0.95))

    def schedule(step):
        if step < warmup:
            return (step + 1) / warmup
        progress = (step - warmup) / max(1, steps - warmup)
        return 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * progress))

    history = []
    for step in range(steps):
        for group in optimizer.param_groups:
            group["lr"] = lr * schedule(step)
        model.train()
        index = torch.randint(len(tx), (batch_size,), generator=generator).to(device)
        loss = lm_loss(model, tx[index], ty[index])
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        optimizer.step()
        if step % eval_every == 0 or step + 1 == steps:
            record = {"step": step + 1, "train_loss": round(loss.item(), 4),
                      "valid_loss": round(evaluate(model, vx, vy), 4),
                      "lr": round(optimizer.param_groups[0]["lr"], 6)}
            history.append(record)
            if log:
                log(json.dumps(record))
    return history

• AdamW with $\beta_2 = 0.95$ (rather than the default 0.999) is a common choice for transformers. It adapts faster when gradient statistics change.
• Warmup ramps the learning rate up over the first steps. Adam’s early estimates of gradient variance are noisy, so the first updates should be small.
• Cosine decay lowers the learning rate smoothly to 10% of its peak, so the model settles into a minimum instead of bouncing around it.
• Gradient clipping caps the global gradient norm at 1.0, so one unusual batch can’t throw the weights far off course.
• Evaluation runs on held-out data, in eval() mode, without gradients, every eval_every steps.

First, prove the loop can learn

Before training on real data, run two checks that catch most bugs in minutes:

1. Initial loss ≈ log V. It confirms the model and loss are wired correctly.
2. Overfit one batch. Train repeatedly on a single fixed batch. A correct model and loop drive the loss near zero within a couple of hundred steps. If they don’t, there’s a bug: targets not shifted (or shifted twice), the optimizer created before parameters were replaced, zero_grad missing, the head disconnected, or a learning rate far too large.

Both are in your milestone tests. A model that can’t memorize one batch isn’t ready for a dataset.

Experiment 1: a task the model can master

The synthetic colors corpus repeats two phrases, “red green blue” and “blue green red”, 160 times each. With a character tokenizer and a 2-layer model, validation loss falls right along with training loss (python run.py train --corpus colors --steps 200 --context 32 --width 64 --layers 2):

{"step": 121, "train_loss": 0.0928, "valid_loss": 0.0676, ...}
{"step": 161, "train_loss": 0.0688, "valid_loss": 0.0636, ...}
{"step": 200, "train_loss": 0.0687, "valid_loss": 0.053,  ...}

The validation text is new, but its patterns are the same as in training, so what the model learned transfers. This is what successful generalization looks like, on a deliberately easy task.

Experiment 2: a real text, and overfitting

Now train on The Verdict with a 512-token BPE vocabulary, 64-token windows and a 4-layer, 128-wide model of 867,072 parameters (python run.py train --steps 600). The 90/10 split gives 260 training windows and 28 validation windows. On a laptop CPU this takes about 40 seconds:

The run starts at the uniform loss, 6.24 = log 512, as it should. Validation loss is best around step 120-180 and then gets worse, while training loss keeps falling toward zero. This is overfitting. The story is about 7,000 tokens long, and the model has 867k parameters, more than a hundred per training token. After learning the general statistics of English, its cheapest way to keep lowering the training loss is to memorize the training text verbatim. That memorized text is useless, even harmful, on the unseen 10%.

Samples make this concrete. At step 600, prompted with the story’s first words, the model reproduces the opening sentence and then degenerates into fragments of memorized phrases:

I HAD always thought Jack Gisburn rather a cheap genies be sply oweagal note that Emperors of thereerly a

Stopping at step 150, near the validation minimum, gives text that is less memorized but still clearly not fluent English. 7,000 tokens is far too little to learn a language from:

I HAD always thought, and he was the coree to doree of the coree to dorethethetheting, and the coreting

BALLM Chapter 5 shows the same effect on the same text. There’s nothing wrong with the code. The remedy is data, not cleverness: real pretraining uses trillions of tokens, so every token is seen once or a handful of times and memorization is rare. Regularization (dropout, used here at 0.1, and weight decay) and early stopping (keeping the checkpoint with the best validation loss) help at the margin.

Note

Perplexity is $e^{\text{loss}}$: here 512 at the start (uniform over 512 tokens), 83 at the validation minimum, and 1.5 on the memorized training set. Read it as “the model is as uncertain as if it were choosing uniformly among this many tokens”. Perplexities are only comparable between models with the same tokenizer: a 512-token vocabulary and a 150,000-token one measure uncertainty per very different units.

Teacher forcing and free-running generation

During training, every position sees the true previous tokens, even where the model would have predicted something else. This is teacher forcing, and it’s why one forward pass gives $T$ training examples. During generation, the model sees its own previous outputs. One early mistake changes every later input, and the model was never trained on contexts containing its own errors. So low validation loss is necessary but not sufficient for good generations. Always look at samples as well as loss curves.

Saving and loading

A checkpoint for inference needs the weights and the configuration that defines the architecture, plus whatever reconstructs the tokenizer:

def save_checkpoint(path, model, extra=None):
    """Weights + config + anything needed to rebuild the tokenizer. Not an exact-resume record:
    optimizer and RNG state are deliberately omitted."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"config": vars(model.cfg), "state": model.state_dict(), "extra": extra or {}}, path)


def load_checkpoint(path, model_class, config_class, device="cpu"):
    """Only load checkpoints you created: torch.load of untrusted pickles can run code."""
    blob = torch.load(path, map_location=device, weights_only=True)
    model = model_class(config_class(**blob["config"])).to(device)
    model.load_state_dict(blob["state"])
    return model.eval(), blob["extra"]

Exactly resuming training also needs the optimizer state (Adam’s $m$ and $v$), the learning-rate schedule position, and the random-number-generator states. Our checkpoints deliberately omit those.

Warning

torch.load on .pt/.bin files unpickles Python objects, and a malicious pickle can run arbitrary code. Load only files you created, with weights_only=True. Real model distribution has moved to safetensors (Chapter 9), which is plain bytes plus a JSON header.

How big is real pretraining?

A useful rule of thumb: training costs about $6ND$ floating-point operations for $N$ parameters and $D$ training tokens. That’s 2 per parameter per token for the forward pass, and 4 for the backward. For our toy model, $6 \times 867\text{k} \times 600 \text{ steps} \times 16 \times 64$ tokens is about 3.2 TFLOP, a few seconds of a laptop. Qwen3-0.6B was trained on around 36 trillion tokens: about $1.3 \times 10^{23}$ FLOPs, or tens of thousands of GPU-days. This is why you’ll load pretrained weights from Chapter 9 on, and why fine-tuning (Part V) changes a trained model rather than starting over.

Build it

pytest tests/test_ch07_training.py
python run.py train --impl engine --steps 600
python run.py generate --impl engine

Stretch exercises

1. ★ Add early stopping to train: keep a copy of the weights with the lowest validation loss and restore them at the end. How much better do the samples look? Where: train in engine/train.py.
2. ★★ Change one variable at a time (dropout 0.0 vs 0.3, width 64 vs 256, learning rate 1e-3 vs 1e-2) and record the best validation loss of each run. Predict each effect before you run it. Where: experiments/ch07.py (create it), adapting run.py’s cmd_train and its GPTConfig.
3. ★★ Train on a larger public-domain text (any Project Gutenberg novel, or several) with the same model. How does the gap between train and validation loss change with 10x more data? Where: put the new corpus in data/; load it in experiments/ch07.py (create it) using engine.data and engine.train.
4. ★★★ Implement gradient accumulation: compute gradients over $k$ small batches before each optimizer step, weighting by tokens. Verify that $k$ micro-batches of size 8 give the same update as one batch of 8k. Where: train in engine/train.py.

Check your understanding

1. Why does cross-entropy take logits rather than probabilities?
2. What initial loss do you expect for a vocabulary of 50,257, and what does a much larger value suggest?
3. Why can validation loss rise while training loss falls?
4. Why can a model with low teacher-forced loss still produce poor generations?
5. What extra state, beyond the weights, do you need to resume training exactly?

Going deeper

• BALLM Chapter 5 §§5.1-5.2 (pp. 128-150): evaluating generative models, training on The Verdict and observing overfitting. Appendix D (pp. 313-321): warmup, cosine decay, clipping.
• Kaplan et al., Scaling Laws for Neural Language Models (2020) and Hoffmann et al., Training Compute-Optimal Large Language Models (“Chinchilla”, 2022): how parameters, data and compute trade off.
• Andrej Karpathy, Let’s reproduce GPT-2 (124M) (video): a full pretraining run of the model your Chapter 9 loader will import.

8. Generating text: decoding and sampling

The model proposes, the policy decides

The model’s job ends at logits: one score per vocabulary entry for the next position. Turning those scores into one token is a separate decoding policy, and the same model behaves very differently under different policies. Keeping them separate in code is good engineering: you can test the policy on fixed logits, and change it per request without touching the model.

Greedy decoding

Take the highest-scoring token: logits.argmax(-1). No softmax is needed, since the largest logit always has the largest probability. Greedy is deterministic, which makes it ideal for testing (two implementations of the same model must produce identical greedy text). It’s also the best choice for short factual answers and for code.

Its weakness shows on open-ended text. Always taking the locally most likely token tends to fall into loops (“the coree to doree of the coree to doree…”, as the early-stopped model in Chapter 7 did), because the model’s own repeated output makes the repetition ever more likely.

Temperature

To sample instead, convert logits to probabilities and draw from them. Temperature $\tau$ reshapes the distribution first:

$$ p_i = \operatorname{softmax}(z / \tau)_i . $$

For logits $[2, 1, 0]$:

Low temperature sharpens toward greedy, and high temperature flattens toward uniform. $\tau = 1$ is the model’s own distribution. $\tau = 0$ would divide by zero, so implementations treat it as a separate greedy branch.

Truncation: top-k, top-p and min-p

Even at $\tau = 1$, the long tail of a 150,000-token vocabulary holds real probability mass. Thousands of individually unlikely tokens add up, and drawing one of them occasionally derails a generation. Truncation removes the tail before sampling:

• Top-k keeps the $k$ highest-scoring tokens and renormalizes. Simple, but fixed: $k = 40$ is too many when the model is certain and too few when many continuations are reasonable.
• Top-p (nucleus sampling, Holtzman et al. 2020) keeps the smallest set of most-likely tokens whose total probability reaches $p$. For probabilities [0.6, 0.25, 0.1, 0.05] and $p = 0.8$, it keeps the first two tokens (0.6 + 0.25 = 0.85). The set adapts: it’s small when the model is confident and large when it isn’t. Note the boundary rule: keep the token that crosses $p$, not just the tokens before it. Otherwise $p = 0.5$ would keep nothing here.
• Min-p keeps tokens whose probability is at least $p_{\min}$ times the top token’s. For [0.5, 0.3, 0.15, 0.04, 0.01] and $p_{\min} = 0.1$, the threshold is 0.05, so it keeps the first three. It scales its cutoff with the model’s confidence.

The order of operations matters: temperature, then top-k, then top-p, then min-p, then renormalize and draw. Production engines apply them in this order (sometimes with options to reorder).

def sample(logits, temperature=0.0, top_k=None, top_p=None, min_p=None, generator=None):
    """Choose one token per row of logits [B, V]; returns [B, 1].  (Your engine: Chapter 8)

    temperature == 0 is greedy argmax. Otherwise: divide by temperature, keep the top_k
    scores, keep the smallest prefix of the sorted distribution whose mass reaches top_p
    (always keeping the first token that crosses it), drop tokens whose probability is below
    min_p times the top probability, then draw from what is left.
    """
    if temperature < 0:
        raise ValueError("temperature must be >= 0")
    if temperature == 0:
        return logits.argmax(dim=-1, keepdim=True)
    scores = logits.float() / temperature
    if top_k is not None:
        if not 1 <= top_k <= scores.shape[-1]:
            raise ValueError("top_k must be in [1, vocab]")
        threshold = scores.topk(top_k, dim=-1).values[..., -1:]
        scores = scores.masked_fill(scores < threshold, float("-inf"))
    if top_p is not None:
        if not 0 < top_p <= 1:
            raise ValueError("top_p must be in (0, 1]")
        ordered, order = scores.sort(dim=-1, descending=True)
        probs = ordered.softmax(-1)
        # Remove a token when the mass BEFORE it already reaches top_p.
        remove = (probs.cumsum(-1) - probs) >= top_p
        ordered = ordered.masked_fill(remove, float("-inf"))
        scores = torch.full_like(scores, float("-inf")).scatter(-1, order, ordered)
    if min_p is not None:
        probs = scores.softmax(-1)
        scores = scores.masked_fill(probs < min_p * probs.amax(-1, keepdim=True), float("-inf"))
    return torch.multinomial(scores.softmax(-1), 1, generator=generator)

struct Rng {  // xorshift64*: tiny, seeded, reproducible
    uint64_t s;
    explicit Rng(uint64_t seed) : s(seed * 0x9E3779B97F4A7C15ull | 1) {}
    float uniform() { s ^= s >> 12; s ^= s << 25; s ^= s >> 27; return float((s * 0x2545F4914F6CDD1Dull) >> 40) / float(1ull << 24); }
};
inline size_t sample_top_p(std::vector<float> logits, float temperature, float top_p, Rng& rng) {
    std::vector<size_t> order(logits.size());
    std::iota(order.begin(), order.end(), 0);
    if (temperature == 0) return size_t(std::max_element(logits.begin(), logits.end()) - logits.begin());
    std::sort(order.begin(), order.end(), [&](size_t a, size_t b) { return logits[a] > logits[b]; });
    std::vector<float> p(order.size());
    for (size_t r = 0; r < order.size(); ++r) p[r] = logits[order[r]] / temperature;
    softmax(p.data(), p.size());
    float mass = 0;
    size_t keep = 0;
    while (keep < p.size() && mass < top_p) mass += p[keep++];   // keep the token that crosses top_p
    float r = rng.uniform() * mass;
    for (size_t i = 0; i < keep; ++i) { if (r < p[i]) return order[i]; r -= p[i]; }
    return order[keep - 1];
}

#![allow(unused)]
fn main() {
pub struct Policy {
    pub temperature: f32,
    pub top_k: Option<usize>,
    pub top_p: Option<f32>,
}

pub fn sample(logits: &[f32], policy: &Policy, rng: &mut Rng) -> usize {
    let argmax = || logits.iter().enumerate().fold(0, |best, (i, &x)| if x > logits[best] { i } else { best });
    if policy.temperature == 0.0 {
        return argmax();
    }
    // Sort token ids by score, highest first, and keep only the candidates the policy allows.
    let mut order: Vec<usize> = (0..logits.len()).collect();
    order.sort_by(|&a, &b| logits[b].partial_cmp(&logits[a]).unwrap());
    let keep = policy.top_k.unwrap_or(logits.len()).min(logits.len());
    let mut probs: Vec<f32> = order[..keep].iter().map(|&i| logits[i] / policy.temperature).collect();
    softmax(&mut probs);
    if let Some(p) = policy.top_p {
        let (mut mass, mut cut) = (0.0, probs.len());
        for (rank, q) in probs.iter().enumerate() {
            if mass >= p {
                cut = rank; // the token that crossed p (at rank - 1) stays
                break;
            }
            mass += q;
        }
        probs.truncate(cut);
        let total: f32 = probs.iter().sum();
        probs.iter_mut().for_each(|q| *q /= total);
    }
    let mut r = rng.uniform();
    for (rank, q) in probs.iter().enumerate() {
        if r < *q {
            return order[rank];
        }
        r -= q;
    }
    order[probs.len() - 1]
}
}

The Python version handles a batch of rows at once and implements every filter with masking (-inf for removed tokens), so it runs on the GPU without sorting through Python lists. The C++ and Rust versions sort one row explicitly, which makes the top-p boundary rule easy to see.

Repetition penalties

The repetition penalty (from the CTRL paper) discourages tokens already present in the context: positive logits are divided by the penalty, and negative ones multiplied by it. Penalties of 1.05-1.2 reduce loops, but they also discourage legitimately repeated words like names and code identifiers. Related variants subtract a fixed amount per occurrence (the frequency penalty) or once per token seen (the presence penalty). They’re all heuristics layered on the model’s distribution, so use them sparingly.

def repetition_penalty(logits, previous_ids, penalty=1.1):
    """CTRL-style penalty: shrink positive logits and grow negative logits of tokens already seen."""
    if penalty == 1.0:
        return logits
    seen = torch.zeros_like(logits, dtype=torch.bool).scatter(-1, previous_ids, True)
    adjusted = torch.where(logits > 0, logits / penalty, logits * penalty)
    return torch.where(seen, adjusted, logits)

Seeds and reproducibility

Pass an explicit torch.Generator seeded per request. Then the same seed, model, policy and prompt give the same text. In a server, every request needs its own generator: if requests share one, the tokens drawn for one user depend on which other users happened to be in the batch. Reproducibility also depends on the device and kernel versions, so record them (Chapter 18’s manifest does).

The generation loop

Generation repeats: run the model, take the last position’s logits, choose a token, append it, stop if it’s a stop token or the budget is spent:

@torch.inference_mode()
def generate_stream(model, ids, max_new_tokens, cached=True, eos_ids=(), **policy):
    """Yield one new token ID at a time for a single prompt ids [1, P].  (Your engine: Chapter 8)

    Prefill the whole prompt once; its last logits choose the first new token. Each later step
    feeds only the newest token (cached) or the whole sequence again (uncached).
    """
    if ids.ndim != 2 or ids.shape[0] != 1 or ids.shape[1] == 0:
        raise ValueError("generate expects one non-empty prompt of shape [1, P]")
    if ids.shape[1] + max_new_tokens > model.context_limit:
        raise ValueError("Prompt plus new tokens exceed the model's context")
    model.eval()
    if max_new_tokens == 0:
        return
    cache = model.new_cache(1, ids.shape[1] + max_new_tokens) if cached else None
    sequence = ids
    logits = model(ids, cache)
    for step in range(max_new_tokens):
        token = sample(logits[:, -1], **policy)
        yield int(token)                       # .item(): one host sync per token (Chapter 19 removes it)
        if int(token) in eos_ids or step + 1 == max_new_tokens:
            return
        sequence = torch.cat((sequence, token), dim=1)
        logits = model(token, cache) if cached else model(sequence)


def generate(model, ids, max_new_tokens, cached=True, eos_ids=(), **policy):
    """The prompt followed by the generated tokens, as a [1, P+N] tensor."""
    new = list(generate_stream(model, ids, max_new_tokens, cached, eos_ids, **policy))
    return torch.cat((ids, torch.tensor([new], dtype=ids.dtype, device=ids.device)), dim=1)

A few rules, each of which fixes a bug that shows up in real systems:

• The prompt’s last logits choose the first new token. No extra forward pass is needed before the first sample.
• Stop conditions: a stop token (EOS, or <|im_end|> for chat models), a maximum number of new tokens, or the model’s context limit. Decide whether the stop token is included in the output. Here it is; a chat UI would hide it.
• Reject, don’t silently crop. If prompt plus budget exceed the context, raise an error. Silently dropping the start of the prompt changes what the model sees without anyone noticing.
• One host synchronization per token. int(token) copies the token to the CPU, which waits for the GPU to finish. To stream text you have to do that eventually, but Chapter 19 shows how to avoid doing it every step.

The cached flag switches between feeding only the new token (with a KV cache, Chapter 16) and re-running the whole sequence. Both must give identical tokens, and your milestone test checks that. The cache is what makes generation affordable; for now, model.new_cache comes provided.

Watching policies on your trained model

Prompting the model you trained in Chapter 7 (python run.py generate) shows each policy’s character, even on a tiny overfit model:

--- greedy
I HAD always thought Jack Gisburn rather a cheap genies be sply oweagal note that Emperors of thereerly a
--- t=0.8 top_k=20
I HAD always thought Jack Gisburn rat Gisburn wife's biggartw through the hush, why Jvert is, and he said, my enough the mre up the
--- t=1.0 top_p=0.9
I HAD always thought Jack Gisburn rather tap get not feltt."

Beyond sampling

Two other ideas are worth knowing by name. Beam search keeps the $b$ most likely partial sequences instead of one. It’s useful for translation, but it produces bland, repetitive text from LLMs and is rarely used for chat. Constrained decoding masks logits so the output must follow a grammar (valid JSON, a regex, a function signature). It’s a masked_fill before sampling, the same as top-k, just with a smarter mask.

Build it

pytest tests/test_ch08_generation.py
python run.py generate --impl engine --checkpoint runs/model.pt

Tip

Implement top-p on sorted probabilities: remove = (cumsum - probs) >= top_p marks every token whose preceding mass already reached $p$, so the crossing token survives. Then scatter the masked scores back to vocabulary order with scatter.

Stretch exercises

1. ★ Add stop_strings support to generate: stop when the decoded text ends with any given string. Why is this harder than stopping on token IDs? Where: generate / generate_stream in engine/sampling.py; supply a tokenizer for decoded-text matching.
2. ★★ Implement presence and frequency penalties and compare them with the repetition penalty on your Chapter 7 model. Where: add penalty helpers in engine/sampling.py and call them from generate_stream.
3. ★★ Implement beam search with width 4 and compare its output with greedy decoding. Which has the higher total log-probability? Which reads better? Where: add a beam-search helper in engine/sampling.py.
4. ★★★ Implement constrained decoding that forces outputs to be a valid decimal number: precompute, for each vocabulary entry, whether it can continue a partial number, and mask the rest. Where: add a numeric-output mask helper in engine/sampling.py and apply it before sample.

Check your understanding

1. Why is temperature 0 implemented as a separate branch?
2. For probabilities [0.4, 0.3, 0.2, 0.1] and top-p 0.5, which tokens remain?
3. Why should every request in a server have its own random generator?
4. Why does generation not need an extra forward pass before choosing the first new token?

Going deeper

• BALLM §5.3 (pp. 151-158): temperature scaling and top-k sampling.
• Holtzman et al., The Curious Case of Neural Text Degeneration (2020): why greedy and pure sampling fail, and nucleus sampling.
• Keskar et al., CTRL (2019) for the repetition penalty; Nguyen et al., Turning Up the Heat: Min-p Sampling (2024).
• The vLLM and llama.cpp sampler implementations, to see the full set of production options (logit bias, penalties, grammars).

9. Real weights: safetensors and GPT-2

Why load, and what you’re loading

Chapter 7’s estimate put Qwen3-0.6B’s pretraining at about $10^{23}$ FLOPs. Nobody reproduces that to run a model. Open-weight models publish the result of training as files, and an inference engine’s first job is to read those files into exactly the computation that produced them.

A Hugging Face model snapshot is a folder:

models/gpt2/
├── config.json              architecture hyperparameters: layers, widths, vocab size, epsilon, ...
├── tokenizer.json           the tokenizer: vocabulary, merges, special tokens, pre-tokenizer
├── tokenizer_config.json    chat template, special-token roles
├── generation_config.json   default sampling settings and stop tokens
└── model.safetensors        the weights (large models: model-00001-of-00004.safetensors ... + an index)

Download GPT-2 small (548 MB) to follow along:

uv pip install -r optional-requirements.txt
hf download openai-community/gpt2 --local-dir models/gpt2

The safetensors format

Older checkpoints were Python pickles (pytorch_model.bin). Unpickling can execute arbitrary code, so loading an untrusted model was a security hole. They also had to be fully deserialized before you could look at them. Safetensors fixes both with a deliberately simple layout:

┌───────────────┬──────────────────────────────┬────────────────────────────────────────┐
│ 8 bytes: N    │ N bytes: JSON header (UTF-8) │ tensor data: raw little-endian bytes   │
│ (u64, LE)     │                              │ at the offsets the header lists        │
└───────────────┴──────────────────────────────┴────────────────────────────────────────┘

The header maps each tensor name to its dtype, shape and byte range within the data section:

{
  "__metadata__": {"format": "pt"},
  "wte.weight": {"dtype": "F32", "shape": [50257, 768], "data_offsets": [0, 154389504]},
  "h.0.attn.c_attn.weight": {"dtype": "F32", "shape": [768, 2304], "data_offsets": [...]}
}

That’s the entire format. Because data is raw bytes at known offsets, a reader can memory-map the file and read any tensor without touching the others. A loader can inspect a 300 GB checkpoint’s header in milliseconds, which you’ll do in Chapter 30. You can write a valid file by hand:

import json, struct
header = json.dumps({"w": {"dtype": "F32", "shape": [2], "data_offsets": [0, 8]}}).encode()
blob = struct.pack("<Q", len(header)) + header + struct.pack("<2f", 1.5, -2.0)
open("hand.safetensors", "wb").write(blob)        # a real, loadable checkpoint with one tensor

The reader is just as short:

def read_header(path):
    """Return (header dict without __metadata__, metadata dict, byte offset of the data section).  (Your engine: Chapter 9)"""
    with open(path, "rb") as file:
        (length,) = struct.unpack("<Q", file.read(8))
        if length > 100_000_000:
            raise ValueError("Implausibly large safetensors header")
        header = json.loads(file.read(length))
    metadata = header.pop("__metadata__", {}) or {}
    return header, metadata, 8 + length

def load_file(path, device="cpu", names=None):
    """Read tensors into a dict. Memory-maps the file and copies one tensor at a time.  (Your engine: Chapter 9)"""
    header, _, data_start = read_header(path)
    selected = header if names is None else {n: header[n] for n in names}
    tensors = {}
    with open(path, "rb") as file, mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ) as view:
        for name, info in selected.items():
            dtype = DTYPES[info["dtype"]]
            begin, end = info["data_offsets"]
            count = (end - begin) // torch.empty((), dtype=dtype).element_size()
            if count != _numel(info["shape"]):
                raise ValueError(f"{name}: byte range does not match its shape")
            if count == 0:                                   # frombuffer rejects empty reads
                tensors[name] = torch.empty(info["shape"], dtype=dtype, device=device)
                continue
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")          # frombuffer warns that mmap is read-only
                flat = torch.frombuffer(view, dtype=dtype, count=count, offset=data_start + begin)
            tensors[name] = flat.reshape(info["shape"]).to(device, copy=True)
    return tensors

torch.frombuffer creates a tensor that views the mapped bytes with the right dtype, and .to(device, copy=True) makes an owned copy, on the GPU if requested. One tensor at a time keeps peak host memory low. Large models split their weights across shards plus a model.safetensors.index.json that maps each tensor to its file; snapshot_tensors walks them and checks that the index and files agree.

Dtypes, and BF16 for free

Header dtypes are short codes: F32, F16, BF16, I64, U8, F8_E4M3 and so on. BF16 is the most common for modern checkpoints, and it’s worth knowing why it’s so convenient. A BF16 number is exactly the top 16 bits of an FP32: same sign, same 8-bit exponent, mantissa cut from 23 bits to 7. Widening BF16 to FP32 is a 16-bit shift. FP16 has a different exponent width and needs real conversion logic:

# PyTorch does this for you: torch.frombuffer(raw, dtype=torch.bfloat16).float()
# By hand, on the raw 16-bit integers:
def bf16_to_f32(bits16):            # int16 tensor of raw BF16 bit patterns
    return (bits16.to(torch.int32) << 16).view(torch.float32)

// Read just the header: 8-byte little-endian length, then that many bytes of JSON.
inline std::string safetensors_header(const std::string& path) {
    std::ifstream f(path, std::ios::binary);
    if (!f) throw std::runtime_error("cannot open " + path);
    uint64_t n = 0;
    f.read(reinterpret_cast<char*>(&n), 8);           // assumes a little-endian host, like x86 and ARM
    std::string json(n, '\0');
    f.read(json.data(), std::streamsize(n));
    return json;
}
inline float bf16_to_float(uint16_t bits) {
    uint32_t widened = uint32_t(bits) << 16;           // BF16 is the top half of an FP32
    float out;
    std::memcpy(&out, &widened, 4);
    return out;
}

#![allow(unused)]
fn main() {
/// BF16 is the top 16 bits of an f32: widening is a shift.
#[inline]
pub fn bf16_to_f32(bits: u16) -> f32 {
    f32::from_bits((bits as u32) << 16)
}

/// IEEE half precision: 1 sign, 5 exponent, 10 mantissa bits.
pub fn f16_to_f32(bits: u16) -> f32 {
    let sign = if bits >> 15 == 1 { -1.0 } else { 1.0 };
    let exponent = ((bits >> 10) & 0x1f) as i32;
    let mantissa = (bits & 0x3ff) as f32;
    match exponent {
        0 => sign * mantissa * 2f32.powi(-24),                       // subnormal
        31 => if mantissa == 0.0 { sign * f32::INFINITY } else { f32::NAN },
        e => sign * (1.0 + mantissa / 1024.0) * 2f32.powi(e - 15),
    }
}
}

The Rust engine keeps weights in BF16 in memory and widens each value inside the matrix-vector product. That’s half the bytes of FP32, so decode runs about twice as fast, for exactly the reason Chapter 1 gave.

Mapping a checkpoint onto your model

Loading means: for every tensor in the file, find the parameter in your model it belongs to, check its shape, and copy it. Every tensor must be used, and every parameter must be filled. Reject the load if either set has leftovers, because a silently uninitialized parameter produces plausible-looking garbage.

GPT-2’s names map onto your Chapter 6 attribute names like this:

The transpose trap

GPT-2 was written with a Conv1D layer that stores weights as [in, out], while nn.Linear stores [out, in] (Chapter 2). Three of the four matrices change shape when transposed, so forgetting the transpose fails loudly. But c_proj in attention is [768, 768]: square. Copying it untransposed passes the shape check and computes $x W$ instead of $x W^\top$. The model then produces fluent-looking nonsense. This is the canonical example of why shape checks are necessary but not sufficient.

@torch.no_grad()
def load_gpt2(directory, device="cpu", dtype=torch.float32):
    """GPT-2 from Hugging Face's safetensors layout.  (Your engine: Chapter 9)

    GPT-2 stores its projections in a Conv1D layout [in, out]; nn.Linear wants [out, in],
    so those four weights are transposed. The head is tied to the token embedding.
    """
    raw = read_config(directory)
    if raw.get("model_type") != "gpt2" or raw.get("activation_function", "gelu_new") != "gelu_new":
        raise ValueError("Expected a standard GPT-2 checkpoint")
    model = GPT(GPTConfig(vocab=raw["vocab_size"], context=raw["n_positions"], width=raw["n_embd"],
                          layers=raw["n_layer"], heads=raw["n_head"],
                          eps=raw.get("layer_norm_epsilon", 1e-5))).to(device=device, dtype=dtype)
    targets = {"wte.weight": (model.token.weight, False), "wpe.weight": (model.position.weight, False),
               "ln_f.weight": (model.norm.weight, False), "ln_f.bias": (model.norm.bias, False)}
    for i, block in enumerate(model.blocks):
        for source, module, transpose in [("ln_1", block.ln1, False), ("ln_2", block.ln2, False),
                                          ("attn.c_attn", block.attn.qkv, True),
                                          ("attn.c_proj", block.attn.proj, True),
                                          ("mlp.c_fc", block.up, True), ("mlp.c_proj", block.down, True)]:
            targets[f"h.{i}.{source}.weight"] = (module.weight, transpose)
            targets[f"h.{i}.{source}.bias"] = (module.bias, False)
    loaded, head_alias = set(), None
    for name, value in snapshot_tensors(directory):
        name = name.removeprefix("transformer.")
        if name.endswith((".attn.bias", ".attn.masked_bias")):   # old causal-mask buffers, not weights
            continue
        if name == "lm_head.weight":                               # an alias of wte when present
            head_alias = value
            continue
        if name not in targets:
            raise ValueError(f"Unexpected tensor {name}")
        parameter, transpose = targets[name]
        assign(parameter, value.T if transpose else value, name)
        loaded.add(name)
    if set(targets) - loaded:
        raise ValueError(f"Missing tensors: {sorted(set(targets) - loaded)[:5]}")
    if head_alias is not None and not torch.equal(head_alias.to(model.token.weight), model.token.weight):
        raise ValueError("lm_head.weight disagrees with the tied token embedding")
    return model.eval()

The packed c_attn also fixes the order of Q, K and V within its 2,304 outputs (Q first, then K, then V), which matches the chunk(3) in your attention. A checkpoint that packed them as K, Q, V would need reordering, and nothing in the shapes would tell you.

Proving the load is correct

“It loaded without errors” proves little. Correctness has levels, and you should climb them in order:

1. Shapes and names: every tensor used, every parameter filled. The loader enforces this.
2. Logits match an independent implementation on the same input IDs. Hugging Face Transformers is the usual reference. With FP32 on both sides, expect agreement to about 1e-4 or better. Your milestone test builds a small random GPT-2 with Transformers, saves it, loads it with your code, and compares logits to 1e-4.
3. Greedy generations match for several prompts.

When logits disagree, don’t stare at the final output. Find the first intermediate that differs. Run both models with hooks that record each stage, and compare in order: token embeddings, position embeddings, first LayerNorm, Q/K/V of block 0, attention output, MLP output, the stream after block 0, and so on. The first mismatch localizes the bug to one operation, almost always an orientation, an epsilon, an activation variant, or an axis.

Tip

Compare fixed inputs, not generated text. Once two models choose different tokens, everything after is different input, and the comparison is meaningless.

Run GPT-2 in your engine

python run.py gpt2 --model-dir models/gpt2 --prompt "The GPU is a" --new-tokens 30

The tokenizer (from transformers) handles text at the boundary; everything else runs in your code: your safetensors reader, your loader, your GPT, your sampler. GPT-2 small is a 2019 model, so expect grammatical but meandering continuations, which is a big step up from Chapter 7’s model.

Build it

pytest tests/test_ch09_weights.py
python run.py gpt2 --impl engine --model-dir models/gpt2

Tip

torch.frombuffer(view, dtype=..., count=..., offset=...) reads from a memory map without copying. It rejects zero-length reads, so handle empty tensors separately. Silence its “non-writable buffer” warning, then copy (.to(device, copy=True)) so the tensor outlives the mapping.

Stretch exercises

1. ★ Write inventory(dir) output for GPT-2: tensor count and total bytes. Does the total match the parameter formula from Chapter 6 times 4 bytes? Why are there extra tensors? Where: experiments/ch09.py (create it), calling the provided engine.safetensors_io.inventory.
2. ★★ Load GPT-2 in BF16 instead of FP32 and measure the maximum logit difference against the FP32 load. Is greedy text the same for a 50-token generation? Where: experiments/ch09.py (create it), calling engine.loaders.load_gpt2 with each dtype.
3. ★★ Deliberately skip the transpose for c_proj only. Which tests or checks catch it? What does the generated text look like? Where: temporarily change the c_proj assignments in load_gpt2 in engine/loaders.py, then restore them.
4. ★★★ Write a converter that saves your Chapter 7 model as a Hugging Face–compatible GPT-2 snapshot (config, transposed weights, names), then load it with transformers and compare logits. Where: new experiments/export_gpt2.py, using engine.train.load_checkpoint and engine.safetensors_io.save_file.

Check your understanding

1. Why can a safetensors header be read without loading any tensor data?
2. Why is loading a pickle-based checkpoint from an untrusted source dangerous?
3. Why can a missing transpose on a square matrix pass a shape check?
4. A tied checkpoint omits lm_head.weight. Why is that valid, and what must the loader do?
5. Your logits differ from the reference. What’s the first thing to compare?

Going deeper

• BALLM §§5.4-5.5 (pp. 159-168): saving and loading weights, and loading OpenAI’s GPT-2 into a from-scratch model (from TensorFlow checkpoints, with the same transposes).
• The safetensors specification and its security audit.
• Hugging Face’s modeling_gpt2.py, especially Conv1D, the origin of the [in, out] layout.

10. How GPUs run your code, and how to measure it

Why this part of the book exists

Your GPT works. Chapter 1 promised a fast engine, and “fast” means knowing what your hardware can do and making your code do it. The next six chapters teach GPU programming from the ground up: the execution model (this chapter), CUDA (11), matrix multiplication (12), reductions and floating point (13), Triton (14) and FlashAttention (15). Each technique gets used in the engine in Part IV.

Two philosophies of processor design

A CPU core is built to finish one stream of instructions as quickly as possible. Large caches, branch prediction and out-of-order execution all minimize the latency of each instruction. A desktop CPU has 8-32 of these sophisticated cores.

A GPU is built to finish an enormous amount of independent work, without caring how long any single piece takes. It has thousands of simple arithmetic units, small caches per unit, and no speculation. When one group of threads waits for memory, which takes hundreds of cycles, the hardware instantly switches to another group that’s ready. As long as there’s enough independent work, the arithmetic units stay busy. This is throughput-oriented design, and latency hiding through parallelism is the core idea. PMPP Chapter 1 frames the whole field this way.

Matrix multiplication across thousands of tokens and features is about as independent as work gets, which is why deep learning runs on GPUs.

Inside a GPU

An NVIDIA GPU is a collection of streaming multiprocessors (SMs): 132 on an H100, 128 on an RTX 4090. Each SM contains:

• Arithmetic units: FP32/INT32 cores, special-function units for exp and sin, and tensor cores, which perform small matrix multiplications (for example, 16×8×16 in BF16) in a single instruction.
• A large register file (256 KB per SM): the fastest storage, private to each thread.
• Shared memory / L1 cache (up to about 228 KB per SM on Hopper): fast, explicitly managed, shared among the threads of a block.
• Warp schedulers that pick, every cycle, which group of threads issues next.

All SMs share an L2 cache (tens of MB) and the main device memory (HBM or GDDR, tens of GB). Approximate figures for three machines used as examples in this book:

(Vendor datasheet figures; the dense BF16 numbers exclude structured sparsity.)

Threads execute in groups of 32 called warps. All 32 threads of a warp execute the same instruction at the same time on different data, a model NVIDIA calls SIMT (single instruction, multiple threads). Chapter 11 shows how you organize threads into blocks and grids. For now, the takeaway is the memory hierarchy: registers and shared memory are tiny but fast, device memory is huge but relatively slow, and making sure that data loaded from device memory is reused many times before being evicted is most of what “optimizing a kernel” means.

Your Python program submits work; the GPU does it later

When PyTorch runs y = x @ w on a CUDA tensor, the CPU doesn’t multiply anything. It enqueues a kernel on a stream (an ordered queue of GPU work) and returns immediately, usually within 5-20 µs. The GPU executes queued kernels in order, while the CPU races ahead submitting more.

This asynchrony is essential for performance: the GPU never waits for Python if the CPU keeps the queue full. But it has two consequences you must internalize:

1. Anything that needs a GPU value on the CPU waits for the queue to drain. print(y), y.item(), y.tolist(), if y > 0:, .cpu() and .numpy() all synchronize. In a decode loop, one .item() per token means the CPU stops, waits for the GPU, and only then submits the next step’s kernels, leaving the GPU idle while Python runs. Chapter 19 removes these.
2. Naive timing measures submission, not work.

Measuring time correctly

This looks reasonable and is wrong:

start = time.perf_counter()
y = x @ x                     # only *enqueued*
elapsed = time.perf_counter() - start   # ~10 µs, regardless of the matrix size

Two correct options:

• Synchronized wall time. Call torch.cuda.synchronize() before starting the clock (so earlier queued work isn’t counted) and after the operation (so this work is). This measures what a user experiences, including CPU overheads.
• CUDA events. Record an event before and after the work on the GPU’s own timeline and ask for the elapsed time between them. This excludes CPU-side gaps, so it measures kernel time.

def measure_wall(function, samples=10, warmup=3, cuda=None):
    """Median wall-clock milliseconds of function(), synchronizing the GPU before starting and
    after finishing each sample, so queued work is neither excluded nor borrowed.  (Your engine: Chapter 10)"""
    cuda = torch.cuda.is_available() if cuda is None else cuda
    for _ in range(warmup):
        function()
    times = []
    for _ in range(samples):
        if cuda:
            torch.cuda.synchronize()
        start = time.perf_counter()
        function()
        if cuda:
            torch.cuda.synchronize()
        times.append((time.perf_counter() - start) * 1000)
    return {"median_ms": statistics.median(times), "min_ms": min(times), "max_ms": max(times), "samples": samples}


def measure_cuda(function, warmup=5, samples=20, repeats=1):
    """GPU-side milliseconds between two CUDA events recorded on the current stream."""
    for _ in range(warmup):
        function()
    torch.cuda.synchronize()
    times = []
    for _ in range(samples):
        start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        start.record()
        for _ in range(repeats):
            function()
        end.record()
        end.synchronize()
        times.append(start.elapsed_time(end) / repeats)
    return {"median_ms": statistics.median(times), "min_ms": min(times), "max_ms": max(times), "samples": samples}

Three more rules make measurements trustworthy:

• Warm up. The first call pays one-time costs: CUDA context creation, library loading, kernel compilation (Triton, torch.compile) and memory-pool growth. Run a few iterations untimed.
• Report a distribution. Take the median of many samples and keep the min and max. GPU clocks, other processes and thermal throttling all add noise.
• Fix everything except the variable you’re testing. Shapes, dtype, device and input values (some kernels are faster on zeros!) must stay the same across the comparison.

Profiling: why is it slow?

A timer tells you how long; a profiler tells you why. PyTorch’s profiler records every operator and kernel, CPU side and GPU side:

python run.py profile          # writes runs/trace.json and prints the top operators

Open the trace in Perfetto (or chrome://tracing). You’ll see two timelines: the CPU thread submitting operators, and the GPU stream executing kernels. Look for:

• Gaps on the GPU timeline: the GPU is idle waiting for the CPU. The cause is Python overhead, synchronization, or too many tiny kernels. Fixes: fusion, CUDA graphs, batching (Chapters 14, 19, 24).
• Unexpected kernels: a copy_ or contiguous you didn’t mean to cause, or a dtype conversion.
• One kernel dominating: optimize that, and check with the roofline whether it’s already near its limit.

For kernel-level detail, NVIDIA’s Nsight Systems (nsys) shows the whole-program timeline, and Nsight Compute (ncu) shows one kernel’s achieved bandwidth, occupancy and stall reasons. GPU Mode L1 walks through all three tools.

Tip

Profilers add overhead. Use a profile to explain, and a separate un-profiled run to measure.

The roofline model

Before optimizing anything, ask what limits it. Every kernel does some arithmetic (FLOPs) and moves some bytes to and from device memory. Its arithmetic intensity is the ratio:

$$ I = \frac{\text{FLOPs}}{\text{bytes moved}} . $$

A device can do at most $F$ FLOP/s and move at most $W$ bytes/s. So the best achievable throughput is

$$ \text{attainable FLOP/s} = \min(F,; I \cdot W). $$

Plot that against $I$ and you get a “roofline”: a slanted line (memory-bound) that meets a flat line (compute-bound) at the ridge point $I^* = F/W$. For an H100, $I^* \approx 990{,}000/3{,}350 \approx 295$ FLOPs per byte. An operation with lower intensity can’t use the tensor cores fully, no matter how good the kernel is.

Now place real operations on it:

• Vector add in FP32 reads 8 bytes and writes 4 per element, and does 1 FLOP: $I = 1/12$. Hopelessly memory-bound. All you can do is approach the bandwidth limit.
• Matrix × vector (one token of decode through a [4096, 4096] BF16 weight) does $2 \cdot 4096^2$ FLOPs and reads the $2 \cdot 4096^2$-byte matrix: $I \approx 1$. Memory-bound by a factor of about 300 on an H100. This is the Chapter 1 result again, now with a name.
• Batched decode with $B$ tokens per weight read: $I \approx B$. That’s 1 at batch 1, 8 at batch 8, about 120 at batch 128, and about 410 at batch 512 (the input and output reads start to count). Batching is how decode climbs the roofline (Chapter 24).
• Prefill of a 4,096-token prompt through the same weight is a [4096, 4096] @ [4096, 4096] GEMM: $I \approx 1{,}365$, comfortably compute-bound.

This is the most important analytical tool in the book. Before writing any kernel, compute its intensity and its ceiling, so you know when to stop optimizing.

The milestone turns these estimates into functions you can call from any experiment:

def matmul_intensity(m, n, k, bytes_per_element=2):
    """FLOPs per byte of [M,K] @ [K,N] if every operand is read once and C written once.  (Your engine: Chapter 10)"""
    flops = 2 * m * n * k
    traffic = (m * k + k * n + m * n) * bytes_per_element
    return flops / traffic


def attainable_flops(intensity, peak_flops, bandwidth):
    """The roofline: you cannot exceed the compute peak, nor bandwidth x intensity.  (Your engine: Chapter 10)"""
    return min(peak_flops, intensity * bandwidth)


def decode_ceiling(weight_bytes, bandwidth, batch=1, kv_bytes_per_sequence=0):
    """Upper bound on decode tokens/s when every step must read all weights once (shared by the
    batch) plus each sequence's own KV cache.  (Your engine: Chapter 10)"""
    step_bytes = weight_bytes + batch * kv_bytes_per_sequence
    return batch * bandwidth / step_bytes

Amdahl’s law: optimize what matters

If an operation takes a fraction $f$ of the total time and you make it $s$ times faster, the whole program gets

$$ \text{speedup} = \frac{1}{(1-f) + f/s} $$

times faster. Make a 10% operation infinitely fast and the program speeds up by only 1.11x. That’s why you profile the whole engine first (Chapter 19), and then optimize the biggest bar, not the most interesting kernel.

Important

DGX Spark. CPU and GPU share 128 GB of LPDDR5x. Two consequences: the 273 GB/s bandwidth is shared too (a CPU process streaming memory slows your GPU decode), and Linux’s “free” memory figure is misleading. Watch “available” in free -h and use torch.cuda.mem_get_info(). Moving data between cpu and cuda tensors still copies, even though the bytes sit in the same physical memory.

Build it

pytest tests/test_ch10_performance.py
python run.py profile

Stretch exercises

1. ★ Time x @ x for a 4096² FP32 matrix three ways: unsynchronized wall clock, synchronized wall clock, and CUDA events. Explain each number. Where: experiments/ch10.py (create it), using engine.measure.measure_wall and measure_cuda.
2. ★★ Measure the achieved bandwidth of x + y for vector sizes from $2^{10}$ to $2^{28}$ elements. Plot GB/s against size. Where does launch overhead dominate, and what fraction of datasheet bandwidth do you reach at the top? Where: experiments/ch10.py (create it).
3. ★★ Measure the TFLOP/s of a @ b in BF16 for square sizes 256 to 8,192 and plot them on your device’s roofline. At what size do you reach 70% of peak? Where: experiments/ch10.py (create it), using engine.measure for timing and roofline arithmetic.
4. ★★★ Profile one decode step of your GPT-2 with torch.profiler. Count the kernels, and estimate the fraction of step time spent in launch gaps rather than kernels. Where: experiments/ch10.py (create it), wrapping a decode step in torch.profiler.profile.

Check your understanding

1. Why must you synchronize before starting a timer as well as after the operation?
2. Why should profiling and benchmarking be separate runs?
3. What is the arithmetic intensity of decode for a batch of 16 sequences, and is it memory-bound on an H100?
4. A kernel taking 30% of runtime becomes 3x faster. What’s the overall speedup?
5. On DGX Spark, why can a CPU-heavy job slow down GPU decoding?

Going deeper

• PMPP Chapter 1 (heterogeneous computing, latency vs throughput), Chapter 4 §§4.1-4.7 (GPU architecture, warps, scheduling, occupancy), and §22.5 (batching: latency versus throughput).
• GPU Mode L1 (Mark Saroufim, profiling and integrating kernels in PyTorch, with nsys/ncu examples) and L8 (the CUDA performance checklist).
• Williams, Waterman and Patterson, Roofline: An Insightful Visual Performance Model (2009).
• NVIDIA Hopper and Ada architecture whitepapers for exact SM resources.

11. CUDA programming

One function, many threads

A CUDA kernel is a function that runs once per thread, on thousands of threads at the same time. Every thread executes the same code; what differs is its index, which it uses to decide which data to work on. This is the SPMD style (single program, multiple data), and the hardest part of learning it is to stop thinking “loop over elements” and start thinking “I am one element”.

Threads are organized in two levels:

• A block is a group of up to 1,024 threads that run on the same SM. They can cooperate through shared memory and synchronize with barriers (__syncthreads()).
• A grid is all the blocks of one launch. Blocks are independent: the hardware may run them in any order, in parallel or one after another, depending on how many fit on the GPU. A kernel must never assume an order between blocks.

That independence is what makes CUDA programs scale: the same kernel runs on a 48-SM laptop GPU and a 132-SM H100, just with more blocks in flight on the bigger one. PMPP calls this transparent scalability.

Vector addition

The “hello world” of CUDA adds two vectors. Each thread computes its global index from three built-in variables: blockIdx.x (which block), blockDim.x (threads per block) and threadIdx.x (which thread in the block):

$$ i = \texttt{blockIdx.x} \times \texttt{blockDim.x} + \texttt{threadIdx.x} $$

For $N = 1{,}000$ elements and 256 threads per block, we need $\lceil 1000 / 256 \rceil = 4$ blocks, which is 1,024 threads. The last 24 threads have no element to process, so every kernel includes a bounds check. Here’s the same kernel in three languages:

__global__ void add_kernel(const float* a, const float* b, float* out, int64_t n) {
    int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;   // this thread's element
    if (i < n) out[i] = a[i] + b[i];                              // the grid may overshoot n
}

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    """(Your engine: Chapter 14)"""
    pid = tl.program_id(0)                       # which block of BLOCK elements is mine
    offsets = pid * BLOCK + tl.arange(0, BLOCK)  # a vector of BLOCK indices
    mask = offsets < n                           # the last block may run past the end
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)


def vector_add(x, y, block=1024):
    check_device(x, y)
    if x.shape != y.shape or not x.is_contiguous() or not y.is_contiguous():
        raise ValueError("vector_add needs equal-shape contiguous tensors")
    out = torch.empty_like(x)
    n = x.numel()
    add_kernel[(triton.cdiv(n, block),)](x, y, out, n, BLOCK=block)
    return out

#![allow(unused)]
fn main() {
#[cuda_module]
mod kernels {
    use super::*;
    #[kernel]
    pub fn twice(input: &[f32], mut output: DisjointSlice<f32>) {
        let index = thread::index_1d();
        if let Some(value) = output.get_mut(index) {
            *value = input[index.get()] * 2.0;
        }
    }
}
}

The CUDA C++ version is one thread’s view of the work. Triton (Chapter 14) is one block’s view: a program handles BLOCK elements as a vector, and the mask plays the role of the bounds check. NVIDIA’s experimental cuda-oxide compiles Rust to GPU code; its DisjointSlice turns the bounds check into a type-system guarantee that no two threads write the same element.

Launching a kernel

From host code (the CPU side), a launch names the grid and block sizes in triple angle brackets:

int threads = 256;
int blocks = (n + threads - 1) / threads;        // ceiling division
add_kernel<<<blocks, threads, 0, stream>>>(a, b, out, n);

Before that, the data must be in GPU memory. A standalone program allocates device buffers and copies to and from them explicitly (cudaMalloc, cudaMemcpy). The book’s standalone example does this with error checking and CUDA-event timing. Inside PyTorch, tensors already live on the device. You write a small C++ wrapper that checks its inputs, launches the kernel on PyTorch’s current stream, and returns a tensor:

torch::Tensor vector_add(torch::Tensor a, torch::Tensor b) {
    check(a); check(b);                               // CUDA, float32, contiguous
    TORCH_CHECK(a.sizes() == b.sizes(), "shapes differ");
    c10::cuda::CUDAGuard guard(a.device());           // launch on a's GPU
    auto out = torch::empty_like(a);
    int64_t n = a.numel();
    if (n) {                                          // zero-size launches are invalid
        add_kernel<<<(n + 255) / 256, 256, 0, stream()>>>(
            a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), n);
        C10_CUDA_KERNEL_LAUNCH_CHECK();               // surface launch errors immediately
    }
    return out;
}

torch.utils.cpp_extension.load compiles the .cu file the first time you call it and imports the result as a Python module (izh/kernels/cuda.py). The wrapper’s checks are part of correctness. A kernel receives raw pointers and knows nothing about shapes, dtypes or strides; pass it a transposed tensor and it silently reads the wrong elements (Chapter 2).

Warning

Kernel launches are asynchronous, and so are their errors. An out-of-bounds write may only be reported at the next synchronizing call, far from the bug. When something crashes mysteriously, rerun with CUDA_LAUNCH_BLOCKING=1 (launches become synchronous) or under compute-sanitizer (NVIDIA’s memory checker), which pinpoints the offending kernel and thread.

Two-dimensional grids

Grids and blocks can be 2-D or 3-D, which maps naturally onto matrices. For an M × N output with 16×16 blocks:

dim3 block(16, 16);
dim3 grid((N + 15) / 16, (M + 15) / 16);
int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;

A row-major matrix element (row, col) lives at row * N + col. For a width-5 matrix, (2, 3) is at offset 13.

Warps, divergence and occupancy

The hardware runs each block as warps of 32 consecutive threads that execute one instruction at a time, in lockstep. Two performance consequences follow:

• Divergence. If threads in a warp take different branches, the warp runs both paths one after the other, with some threads masked off each time. The bounds check in vector add diverges only in the last warp, which is harmless. Data-dependent branches in a hot loop can halve throughput or worse.
• Occupancy. An SM keeps many warps resident and switches between them to hide memory latency (Chapter 10). How many fit depends on each thread’s registers and each block’s shared memory: use more, and fewer warps fit. Occupancy is resident warps divided by the hardware maximum. It’s a means, not a goal. A kernel with low occupancy but lots of data reuse can beat one with high occupancy and none.

Coalescing: the rule that matters most

When the 32 threads of a warp load from global memory, the hardware combines their requests into as few memory transactions as possible. If thread $t$ reads element $\text{base} + t$, the 32 four-byte reads cover one contiguous 128-byte segment, served by a single transaction. That’s a coalesced access, and it uses the full bandwidth. If thread $t$ reads element $\text{base} + 32t$ (a stride), the warp touches 32 different segments and wastes most of each. Strided access can be 10-30x slower for the same number of useful bytes.

The practical rule: map threadIdx.x, which varies fastest within a warp, to the index that varies fastest in memory (the last axis of a row-major tensor). GPU Mode L8 measures this effect directly (coalesce.cu in the lecture repository).

A naive matrix multiplication

The direct translation of $C_{ij} = \sum_k A_{ik} B_{kj}$ gives each thread one output element:

// One thread per output element. Every thread re-reads a whole row of A and column of B
// from global memory: 2K loads for 2K flops, an arithmetic intensity of ~0.25 flop/byte.
__global__ void naive_matmul_kernel(const float* a, const float* b, float* c,
                                    int M, int K, int N) {
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;   // x varies fastest -> coalesced B and C
    if (row < M && col < N) {
        float total = 0.f;
        for (int k = 0; k < K; ++k) total += a[row * K + k] * b[k * N + col];
        c[row * N + col] = total;
    }
}

Check the address formulas against the shapes: A is [M, K] so its row stride is K, B is [K, N] so its row stride is N. A common bug uses one width for both, which works only for square matrices. Test with non-square shapes like [3, 5] @ [5, 7].

Is it coalesced? Within a warp, col varies (x is the fastest thread index) and row is fixed. Reads of B[k * N + col] are consecutive, which is coalesced. Reads of A[row * K + k] are the same address for every thread, a broadcast, which is fine. Writes to C[row * N + col] are consecutive, also good.

So why is it slow? Apply Chapter 10’s roofline. Each thread loads $2K$ values to do $2K$ FLOPs: an arithmetic intensity of about 0.25 FLOP per byte in FP32. Every element of A is re-read by each of the N threads that use it, and every element of B by M threads. The kernel is starved for memory bandwidth, and its speed is a small fraction of what the arithmetic units can do. Fixing that, by loading tiles into shared memory and reusing them many times, is the subject of the next chapter.

Build it

uv pip install ninja setuptools               # build tools for PyTorch extensions
export CUDA_HOME=/usr/local/cuda              # where your CUDA Toolkit lives
pytest tests/test_ch11_cuda.py -k "vector_add or matmuls"
python run.py kernels --backend cuda --impl engine

uv pip install "nvidia-cuda-nvcc==13.3.*" "nvidia-cuda-runtime==13.3.*" "nvidia-cuda-cccl==13.3.*"
NV=$(python -c "import nvidia, os; print(os.path.dirname(nvidia.__path__[0]))")/nvidia/cu13
$NV/bin/nvcc -std=c++17 -arch=sm_90 -I$NV/include -I$NV/include/cccl cpp/kernels.cu -o build/kernels

Important

DGX Spark compiles for -arch=sm_121 with CUDA 13. PyTorch’s extension builder picks the architecture from the GPU it finds, or from TORCH_CUDA_ARCH_LIST="12.1".

Stretch exercises

1. ★ Make vector add process 4 elements per thread with float4 loads (reinterpret_cast<const float4*>). Measure the bandwidth before and after on 2²⁸ elements. Where: add_kernel in engine/kernels/cuda_ops.cu; handle the tail when the length is not divisible by four.
2. ★★ Write a deliberately uncoalesced copy kernel (thread $t$ reads element $32t \bmod N$) and compare its bandwidth with a coalesced copy. Where: add copy kernels and host launchers in engine/kernels/cuda_ops.cu, then expose them in its PYBIND11_MODULE.
3. ★★ Swap the roles of x and y in the naive matmul (map threadIdx.x to rows). Measure the slowdown and explain it with coalescing. Where: naive_matmul_kernel and its launch geometry in engine/kernels/cuda_ops.cu.
4. ★★★ Write transpose_kernel for a [M, N] FP32 matrix. A naive version must choose between coalesced reads and coalesced writes. Then fix it with a 32×32 shared-memory tile (you’ll need the bank-conflict padding from Chapter 12). Where: add transpose_kernel, a host launcher and a binding in engine/kernels/cuda_ops.cu.

Check your understanding

1. Why launch more threads than elements and then bounds-check, instead of launching exactly N threads?
2. Why can’t one block wait for another block to finish?
3. A warp reads x[threadIdx.x * 16] (FP32). How many 128-byte segments does it touch, and how much of the transferred data is useful?
4. What is the naive matmul’s arithmetic intensity, and what does the roofline predict about its speed?

Going deeper

• PMPP Chapters 2-3 (pp. 23-73): data parallelism, CUDA program structure, multidimensional grids and the naive matmul. Chapter 4 (warps, divergence, occupancy) and §6.1 (coalescing).
• GPU Mode L2 (a recap of PMPP Chapters 1-3), L3 (Jeremy Howard, CUDA for Python programmers, writing kernels inline from a notebook), L4 (compute and memory architecture), and L8 (the performance checklist, with runnable coalesce.cu, divergence.cu and occupancy.cu).
• NVIDIA’s CUDA C++ Programming Guide, chapters on the programming model and memory hierarchy; compute-sanitizer documentation.

12. Fast matrix multiplication

Why matrix multiplication deserves a chapter

In a transformer, nearly all arithmetic is matrix multiplication: the QKV, output and MLP projections of every layer, and the head. Prefill is a sequence of large GEMMs (general matrix-matrix multiplications), and their efficiency sets time-to-first-token. You’ll rarely beat NVIDIA’s cuBLAS or CUTLASS for large GEMMs, and you shouldn’t try. But every kernel you’ll write later (FlashAttention, quantized matmuls, paged attention) is built from the same tiling ideas. This chapter is where you learn them.

Counting the waste

Chapter 11’s naive kernel computes each output with $2K$ loads from global memory. For $M = N = K = 4096$, that’s $2 \times 4096^3 \approx 137$ billion loads to compute 16.8 million outputs, while the matrices themselves contain only $3 \times 4096^2 \approx 50$ million distinct values. Every element of A and B is fetched 4,096 times. The L2 cache catches some of the repeats, but the arithmetic units still mostly wait.

Tiling: load once, use many times

Divide the output into $T \times T$ tiles and give each tile to one thread block. To compute its tile, the block needs a $T$-row strip of A and a $T$-column strip of B. Walk along $K$ in phases. In each phase, the block cooperatively copies one $T\times T$ tile of A and one of B into shared memory, one element per thread, and then every thread accumulates $T$ products from those tiles:

Each loaded value is now used $T$ times instead of once, so global traffic drops by a factor of $T$:

$$ \text{loads}{\text{tiled}} = \frac{2MNK}{T}, \qquad I{\text{phase}} = \frac{2T^3 \text{ FLOPs}}{2T^2 \times 4 \text{ bytes}} = \frac{T}{4} \text{ FLOP/byte (FP32)}. $$

$T = 16$ gives an intensity of 4 instead of 0.25. That’s a 16x higher ceiling on the roofline.

You can verify the traffic arithmetic without a GPU. tiled_matmul performs the same tile-by-tile computation on the CPU and counts every element it “loads into shared memory”:

def tiled_matmul(a, b, tile=16):
    """C = A @ B computed one TILE x TILE output tile at a time, the way a CUDA block does it.
    Returns (C, elements_loaded): every tile of A and B copied into "shared memory" is counted
    as a global-memory load, so you can compare traffic with the naive kernel's 2*M*N*K.  (Your engine: Chapter 12)
    """
    m, k = a.shape
    n = b.shape[1]
    c = torch.zeros(m, n, dtype=torch.float32)
    loads = 0
    for row in range(0, m, tile):                  # one "block" per output tile
        for col in range(0, n, tile):
            acc = torch.zeros(min(tile, m - row), min(tile, n - col))
            for phase in range(0, k, tile):        # march along K one tile at a time
                a_tile = a[row:row + tile, phase:phase + tile].float()   # load into shared memory
                b_tile = b[phase:phase + tile, col:col + tile].float()
                loads += a_tile.numel() + b_tile.numel()
                acc += a_tile @ b_tile                                   # reuse each value `tile` times
            c[row:row + tile, col:col + tile] = acc
    return c, loads

For 64×64 matrices, it loads exactly $1/T$ as many elements as the naive kernel, for every tile size. That’s what the milestone tests check.

The CUDA kernel

constexpr int TILE = 16;
// Each block computes a TILE x TILE output tile. In phase p every thread loads one element of
// A's tile and one of B's into shared memory; then all threads reuse those 2*TILE^2 values
// TILE times each. Global traffic drops by a factor of TILE.
__global__ void tiled_matmul_kernel(const float* a, const float* b, float* c, int M, int K, int N) {
    __shared__ float As[TILE][TILE];
    __shared__ float Bs[TILE][TILE];
    int row = blockIdx.y * TILE + threadIdx.y;
    int col = blockIdx.x * TILE + threadIdx.x;
    float total = 0.f;
    for (int phase = 0; phase < K; phase += TILE) {
        int ak = phase + threadIdx.x, bk = phase + threadIdx.y;
        As[threadIdx.y][threadIdx.x] = (row < M && ak < K) ? a[row * K + ak] : 0.f;   // zero-pad edges
        Bs[threadIdx.y][threadIdx.x] = (bk < K && col < N) ? b[bk * N + col] : 0.f;
        __syncthreads();                         // tile fully loaded before anyone reads it
        for (int k = 0; k < TILE; ++k) total += As[threadIdx.y][k] * Bs[k][threadIdx.x];
        __syncthreads();                         // everyone done reading before the next overwrite
    }
    if (row < M && col < N) c[row * N + col] = total;
}

@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K,
                  stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
                  BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
                  GROUP_M: tl.constexpr):
    """(Your engine: Chapter 14)"""
    # Map the 1-D program id to an output tile. Walking GROUP_M tile-rows at a time keeps the
    # B tiles those rows need hot in L2 cache ("grouped ordering", Triton tutorial 03).
    pid = tl.program_id(0)
    tiles_m, tiles_n = tl.cdiv(M, BLOCK_M), tl.cdiv(N, BLOCK_N)
    group = pid // (GROUP_M * tiles_n)
    first_m = group * GROUP_M
    group_size = min(tiles_m - first_m, GROUP_M)
    pid_m = first_m + (pid % group_size)
    pid_n = (pid % (GROUP_M * tiles_n)) // group_size

    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    rk = tl.arange(0, BLOCK_K)
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)     # accumulate in FP32 whatever the inputs
    for k0 in range(0, K, BLOCK_K):
        a = tl.load(a_ptr + rm[:, None] * stride_am + (k0 + rk)[None, :] * stride_ak,
                    mask=(rm[:, None] < M) & ((k0 + rk)[None, :] < K), other=0.0)
        b = tl.load(b_ptr + (k0 + rk)[:, None] * stride_bk + rn[None, :] * stride_bn,
                    mask=((k0 + rk)[:, None] < K) & (rn[None, :] < N), other=0.0)
        acc += tl.dot(a, b)                                    # tensor cores when dtypes allow
    c = acc.to(c_ptr.dtype.element_ty)
    tl.store(c_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn, c,
             mask=(rm[:, None] < M) & (rn[None, :] < N))


def matmul(a, b, block_m=64, block_n=64, block_k=32, group_m=8):
    check_device(a, b)
    if a.ndim != 2 or b.ndim != 2 or a.shape[1] != b.shape[0]:
        raise ValueError("matmul expects [M, K] @ [K, N]")
    m, k = a.shape
    n = b.shape[1]
    c = torch.empty((m, n), device=a.device, dtype=a.dtype)
    grid = (triton.cdiv(m, block_m) * triton.cdiv(n, block_n),)
    matmul_kernel[grid](a, b, c, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1),
                        c.stride(0), c.stride(1), BLOCK_M=block_m, BLOCK_N=block_n,
                        BLOCK_K=block_k, GROUP_M=group_m)
    return c

Three details make the CUDA version correct:

1. Two barriers per phase. The first __syncthreads() ensures the whole tile is loaded before anyone reads it. The second ensures everyone has finished reading before the next phase overwrites the tile. Remove the second and fast threads corrupt the tile under slow ones. The result is wrong only sometimes, which makes it a nightmare to debug.
2. Zero-padded edges. When $M$, $N$ or $K$ isn’t a multiple of $T$, threads that would load out of bounds load 0 instead. Zeros don’t change the dot product, so edge tiles need no special arithmetic.
3. Every thread reaches every barrier, including threads whose output is outside the matrix. They still help load tiles; they just skip the final write. A thread that returns early would make the barrier wait forever, or worse.

The Triton version (Chapter 14 explains the language) expresses the same algorithm at the level of tiles: tl.load a [BLOCK_M, BLOCK_K] block, tl.dot it with a [BLOCK_K, BLOCK_N] block, and accumulate in FP32. Triton handles shared memory and synchronization itself. It also maps tl.dot onto tensor cores, which our CUDA kernel doesn’t use.

Shared-memory bank conflicts

Shared memory is divided into 32 banks, and consecutive 4-byte words go to consecutive banks. A warp’s 32 accesses complete in one cycle if they hit 32 different banks (or the same word, which is broadcast). If several threads hit different words in the same bank, the accesses are serialized.

Reading a row of a float tile[32][32] across a warp touches 32 consecutive words: 32 banks, no conflict. Reading a column touches words 32 apart, all in the same bank: a 32-way conflict. The classic fix is padding: declare float tile[32][33]. Now advancing one row moves by 33 words, which is one bank over, so a column spans all 32 banks. Our tiled matmul reads As[ty][k] (a broadcast within a warp) and Bs[k][tx] (a row), so it’s conflict-free without padding. The transpose exercise in Chapter 11 needs the padding.

Register tiling and thread coarsening

Tiling into shared memory raised the intensity of global memory traffic. But each FMA in the inner loop still reads two values from shared memory, which has its own bandwidth limit. The next level of reuse moves into registers: each thread computes a small block of outputs, say $4\times4$ or $8\times8$, instead of one. Each value it reads from shared memory then feeds 4 or 8 FMAs held in registers:

for k in tile:
    a_frag[0..7] = As[ty*8 .. ty*8+7][k]      # 8 values of A
    b_frag[0..7] = Bs[k][tx*8 .. tx*8+7]      # 8 values of B
    acc[i][j] += a_frag[i] * b_frag[j]        # 64 FMAs from 16 loads

This is thread coarsening (PMPP §6.5 and Chapter 15): fewer threads, each doing more work, with higher reuse. The costs are more registers per thread (lower occupancy) and more complex code. Well-tuned FP32 SIMT kernels reach roughly 50-70% of FP32 peak this way.

Tensor cores

Modern NVIDIA GPUs have tensor cores: units that compute a small matrix product, such as a 16×8×16 BF16 tile with FP32 accumulation, as a single warp-wide instruction. Their throughput is roughly 8-16x that of the ordinary FP32 cores, which is where the “~990 TFLOP/s” of an H100 comes from. Using them requires BF16, FP16, FP8 or INT8 inputs, data arranged in specific fragment layouts, and enough reuse to feed them, so tiling matters even more.

You can program tensor cores with CUDA’s wmma/mma instructions (GPU Mode L23), with CUTLASS/CuTe templates (L15, L36, L57), or let Triton’s tl.dot do it. For the engine, the right choice is clear: use cuBLAS (torch.matmul) for the big GEMMs, and write kernels only where you can fuse something it can’t. Typical results for a 4096³ BF16 GEMM look like this:

(The fractions are relative to the tensor-core peak, which is why even a well-tuned FP32 SIMT kernel looks small here. Measure your own device in the stretch exercises.)

Decode is a different problem: GEMV

During decode, each projection multiplies a single vector (or a few, with batching) by a weight matrix. That’s a matrix-vector product (GEMV), with an arithmetic intensity of about 1 (Chapter 10). Tiling can’t help, because there’s nothing to reuse: each weight is used exactly once per token. A good GEMV kernel just streams the matrix at full bandwidth:

• coalesced, vectorized loads (each thread reads 16 bytes at a time, consecutive threads read consecutive memory);
• enough parallelism: one or more warps per output row, combining partial sums with warp reductions (Chapter 13);
• no wasted bytes: weights stored compactly, ideally in 4-bit form and dequantized in registers (Chapter 20).

That’s why decode optimization in practice is quantization plus bandwidth-efficient GEMV, and prefill optimization is GEMM. Your Rust engine’s matvec (rows split across CPU threads, BF16 widened in the inner loop) is a CPU GEMV in exactly this spirit.

Build it

pytest tests/test_ch12_matmul.py                         # runs anywhere
pytest tests/test_ch11_cuda.py -k matmuls                # GPU: naive and tiled vs torch

Stretch exercises

1. ★ Run tiled_matmul with tiles 4, 8, 16 and 32 on 256×256 matrices and plot loads against tile size. Then estimate the shared-memory capacity a 64×64 FP32 tile pair would need. Would it fit on your GPU? Where: experiments/ch12.py (create it), calling engine.tiling.tiled_matmul with each tile size.
2. ★★ Add 2×2 register tiling to the CUDA kernel (each thread computes four outputs) and measure the speedup. Where: tiled_matmul_kernel and its launch geometry in engine/kernels/cuda_ops.cu.
3. ★★ Write a GEMV kernel (y = W x, W [N, K] BF16) where each warp computes one output row with a warp-shuffle reduction. Measure its achieved bandwidth against the datasheet. Where: add a GEMV kernel, host launcher and binding in engine/kernels/cuda_ops.cu.
4. ★★★ Use wmma to write a tensor-core BF16 GEMM for multiples of 16, and compare it with cuBLAS. Where: add a WMMA kernel, host launcher and binding in engine/kernels/cuda_ops.cu.

Check your understanding

1. Why does a tiled matmul need two barriers per phase, and what goes wrong without the second?
2. Why can larger tiles be slower, even though they increase reuse?
3. A warp reads column 5 of float s[32][32]. How many bank conflicts occur, and how does s[32][33] fix it?
4. Why doesn’t tiling help a decode-time matrix-vector product?

Going deeper

• PMPP Chapter 5 (pp. 103-130): memory types, tiling, the tiled matmul kernel, boundary checks. §6.4-6.5: bank conflicts and coarsening. Chapter 15 (GEMM): register tiling, software pipelining and tensor-core considerations.
• GPU Mode L5 (Jeremy Howard, tiled matmul in CUDA and Numba from Python), L23 (tensor cores), L15/L36/L57 (CUTLASS and CuTe).
• Simon Boehm, How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance (blog, 2022): ten kernels from naive to 90% of cuBLAS, each measured.

13. Reductions, softmax and floating point

One answer from many values

A reduction combines many values into one: a sum, a maximum, a mean. Transformers are full of them. Softmax needs a maximum and a sum per row, LayerNorm and RMSNorm need a mean of squares per token, and the loss averages over tokens. Unlike elementwise operations, reductions need threads to combine their results, which raises two questions: how to do it in parallel, and how much numerical precision the combination loses.

Parallel reduction: a tree

A serial sum does $N - 1$ additions one after another. A tree adds independent pairs in parallel: $N/2$ additions, then $N/4$, and so on. It’s still $N-1$ additions in total, but only $\log_2 N$ sequential steps. For 8 values, that’s 3 steps instead of 7.

On a GPU, a block reduces its share of the data in two levels:

1. Within a warp, threads exchange register values directly with shuffle instructions. __shfl_down_sync(mask, v, offset) gives each thread the value held by the thread offset lanes above it. Five rounds (offsets 16, 8, 4, 2, 1) leave the warp’s total in lane 0, with no shared memory and no barriers.
2. Across warps, each warp’s lane 0 writes its partial sum to shared memory, the block synchronizes, and the first warp reduces those partials the same way.

// Sum within a warp: each step halves the number of live partial sums, using register
// exchanges instead of shared memory. After 5 steps lane 0 holds the warp's total.
__device__ float warp_sum(float v) {
    for (int offset = 16; offset > 0; offset >>= 1) v += __shfl_down_sync(0xffffffff, v, offset);
    return v;
}

// Block-wide sum: every warp reduces, warp leaders publish to shared memory, warp 0 finishes.
__device__ float block_sum(float v) {
    __shared__ float partial[32];
    int lane = threadIdx.x % 32, warp = threadIdx.x / 32;
    v = warp_sum(v);
    if (lane == 0) partial[warp] = v;
    __syncthreads();
    v = (threadIdx.x < blockDim.x / 32) ? partial[lane] : 0.f;
    if (warp == 0) v = warp_sum(v);
    __syncthreads();                             // partial[] may be reused by the caller
    return v;                                    // valid in thread 0
}

// One block per row; each thread strides over the row (thread coarsening), then reduces.
__global__ void row_sum_kernel(const float* x, float* out, int width) {
    const float* row = x + int64_t(blockIdx.x) * width;
    float v = 0.f;
    for (int i = threadIdx.x; i < width; i += blockDim.x) v += row[i];
    v = block_sum(v);
    if (threadIdx.x == 0) out[blockIdx.x] = v;
}

row_sum_kernel also uses thread coarsening: each thread first sums many elements serially in a register (for (i = tid; i < width; i += blockDim.x)), and only then does the block reduce 256 partials. Reading consecutive elements per iteration keeps the loads coalesced. Two details are worth remembering:

• Identity values. Threads with no data contribute the operation’s identity: 0 for a sum, but $-\infty$ for a max. Contributing 0 to a max of all-negative values would be wrong.
• Barriers across blocks don’t exist. If a reduction spans several blocks, write per-block partial results to global memory and reduce them in a second kernel, or use atomics.

RMSNorm: a reduction plus a broadcast

RMSNorm (Chapter 17) needs one reduction per token (the mean of squares), then the result must reach every thread so each can scale its elements. Thread 0 writes the statistic to shared memory, the block synchronizes, and everyone reads it:

__global__ void rmsnorm_kernel(const float* x, const float* w, float* y, int width, float eps) {
    const float* row = x + int64_t(blockIdx.x) * width;
    float* out = y + int64_t(blockIdx.x) * width;
    __shared__ float inv_rms;
    float sq = 0.f;
    for (int i = threadIdx.x; i < width; i += blockDim.x) sq += row[i] * row[i];
    sq = block_sum(sq);
    if (threadIdx.x == 0) inv_rms = rsqrtf(sq / width + eps);
    __syncthreads();                             // broadcast the statistic to every thread
    for (int i = threadIdx.x; i < width; i += blockDim.x) out[i] = row[i] * inv_rms * w[i];
}

One block per token row, reading the row twice: once for the statistic, once to normalize. A fused Triton version (Chapter 14) also adds the residual connection in the same pass.

Floating point: what your numbers can hold

A floating-point number is $\pm, 1.m \times 2^{e}$: a sign bit, an exponent field, and a mantissa (fraction) field. The exponent bits set the range, and the mantissa bits set the precision, meaning the gap between neighbouring representable numbers. Every format makes a different trade:

BF16 keeps FP32’s 8-bit exponent, so it has the same range and never overflows where FP32 wouldn’t. It pays with precision: only about 2-3 significant decimal digits. That’s why it dominates training and inference. FP16 has more precision but overflows above 65,504, a real problem for activations. FP8 and FP4 are so coarse that they only work with scale factors stored alongside groups of values (Chapter 20).

Rounding makes order matter

In exact arithmetic, $(10^8 + 1) - 10^8 = 1$. In FP32, the gap between representable numbers near $10^8$ is 8, so $10^8 + 1$ rounds back to $10^8$ and the result is 0. Computing $(10^8 - 10^8) + 1$ instead gives 1. Floating-point addition is not associative: a tree reduction, a fused kernel and a serial loop can all give different answers, all legitimately.

Precision of the accumulator matters even more. Summing 0.1 ten thousand times should give 1,000:

BF16 accumulator, sequential:  32.0
FP32 accumulator:              1000.98    (0.1 itself isn't exact in BF16)

The BF16 sum is stuck at 32, where the gap between representable numbers is 0.25: adding 0.1 rounds back down every time. This is why every serious kernel stores in low precision but accumulates in FP32: matmuls on tensor cores, reductions in norms, the softmax denominator. Your attention and norm code converts to FP32 (.float()) for its reductions for exactly this reason.

Stable softmax

$e^{1000}$ overflows every format, so a naive softmax of [1000, 999, 998] returns [nan, nan, nan]. Since softmax is unchanged by subtracting a constant from every input,

$$ \operatorname{softmax}(x)_i = \frac{e^{x_i - m}}{\sum_j e^{x_j - m}}, \qquad m = \max_j x_j , $$

subtracting the row maximum makes the largest exponent $e^0 = 1$ and everything else smaller. The result is $[0.6652, 0.2447, 0.0900]$, exactly as if the inputs had been $[2, 1, 0]$.

def stable_softmax(scores, dim=-1):
    """Softmax that subtracts the row maximum first.  (Your engine: Chapter 13)

    Computes in FP32 and returns the input dtype. Each row needs at least one finite score.
    """
    x = scores.float()
    shifted = x - x.amax(dim=dim, keepdim=True)
    numerators = shifted.exp()
    return (numerators / numerators.sum(dim=dim, keepdim=True)).to(scores.dtype)

inline void softmax(float* x, size_t n) {
    float m = *std::max_element(x, x + n), sum = 0;
    for (size_t i = 0; i < n; ++i) sum += (x[i] = std::exp(x[i] - m));
    for (size_t i = 0; i < n; ++i) x[i] /= sum;
}

#![allow(unused)]
fn main() {
/// Stable softmax in place: subtract the maximum so the largest exponent is exp(0) = 1.
pub fn softmax(x: &mut [f32]) {
    let max = x.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let mut sum = 0.0;
    for v in x.iter_mut() {
        *v = (*v - max).exp();
        sum += *v;
    }
    for v in x.iter_mut() {
        *v /= sum;
    }
}
}

The same trick gives a stable log-sum-exp, $\log\sum_j e^{x_j} = m + \log\sum_j e^{x_j - m}$, which is how cross-entropy is computed without ever forming probabilities (Chapter 7). Keep the softmax’s edge case in mind: if every entry of a row is $-\infty$ (a fully masked attention row), the max is $-\infty$ and the result is NaN. Your mask logic must guarantee at least one visible key, or define what an empty row means.

In Chapter 15, you’ll extend this to compute softmax in tiles, carrying the running maximum forward. That’s the online softmax at the heart of FlashAttention.

Comparing implementations honestly

Because different correct implementations round differently, “equal” means “close enough”, and you must decide how close before you look at the results. The standard test combines absolute and relative tolerances:

$$ |a - r| \le \text{atol} + \text{rtol}\cdot|r| . $$

The absolute term handles values near zero; the relative term handles scale. Good error reports also include aggregate measures, because a maximum alone can hide a systematic drift:

def compare(actual, expected):
    """An error report, not a pass/fail verdict. Choose tolerances per dtype and operation."""
    difference = (actual.float() - expected.float()).abs()
    reference = expected.float()
    return {
        "max_abs": difference.max().item() if difference.numel() else 0.0,
        "mean_abs": difference.mean().item() if difference.numel() else 0.0,
        "relative_l2": (
            torch.linalg.vector_norm(difference)
            / torch.linalg.vector_norm(reference).clamp_min(1e-12)
        ).item(),
    }

Which tolerances? They depend on dtype, operation and magnitude. A BF16 result can’t be closer than about 0.4% relative to anything, so atol=1e-5 on BF16 logits is meaningless. A useful discipline is to make three comparisons:

1. Trusted FP32 vs your FP32: tests your implementation. Expect ~1e-5 or better.
2. Trusted BF16 vs trusted FP32: measures what the precision itself costs.
3. Your BF16 vs trusted BF16: should be about as close as (2), not much worse.

If (3) is much worse than (2), you have a bug or a bad accumulation choice. And whatever the tolerance, rounding never excuses a semantic error: a wrong mask or transpose produces differences far larger than any rounding.

Warning

On NVIDIA GPUs, PyTorch may run “FP32” matmuls with TF32 tensor cores (10-bit mantissa) when torch.backends.cuda.matmul.allow_tf32 is on, giving about 1e-3 relative error. Set the flag explicitly whenever you compare FP32 results.

Note

Some GPU reductions are nondeterministic: using atomic additions or split-K matmuls, the order of additions changes run to run, and so do the last bits of the result. GPU Mode L9 (nondeterminism.py) demonstrates it. Serving engines that promise bitwise-reproducible outputs have to avoid these kernels.

Build it

pytest tests/test_ch13_numerics.py
pytest tests/test_ch11_cuda.py -k reductions     # GPU

Stretch exercises

1. ★ Reproduce the BF16 summation failure, then implement Kahan summation in BF16 (carry a compensation term) and see how close it gets to 1,000. Where: experiments/ch13.py (create it); add a Kahan helper to engine/numerics.py.
2. ★★ Write a block reduction for the maximum and check it on rows of all-negative numbers. What breaks if invalid lanes contribute 0? Where: add max-reduction helpers and a row-max kernel, launcher and binding in engine/kernels/cuda_ops.cu.
3. ★★ Run Qwen3 or your GPT in FP32 and in BF16 on the same input, and report max-abs and relative-L2 error of the logits, and whether greedy tokens agree over 50 steps. Where: experiments/ch13.py (create it), importing engine.numerics.compare.
4. ★★★ Write a single-pass softmax kernel in CUDA: one block per row, a max reduction, a sum reduction, then normalize, with the row held in registers. Compare its bandwidth with torch.softmax. Where: add a row-softmax kernel, launcher and binding in engine/kernels/cuda_ops.cu.

Check your understanding

1. Why does subtracting the row maximum leave softmax unchanged?
2. Why does BF16 have FP32’s range but far less precision?
3. Why should a BF16 kernel accumulate sums in FP32?
4. Why report aggregate error in addition to the maximum error?
5. What identity value should inactive threads contribute to a max reduction?

Going deeper

• PMPP Chapter 10 (reductions: trees, divergence, warp shuffles, coarsening, pp. 225-248) and Appendix A (floating-point considerations).
• GPU Mode L9 (Mark Saroufim, reductions, with runnable kernels and the nondeterminism demo) and L84 (numerics and AI).
• David Goldberg, What Every Computer Scientist Should Know About Floating-Point Arithmetic (1991); Micikevicius et al., FP8 Formats for Deep Learning (2022).

14. Triton and kernel fusion

Why another GPU language

CUDA C++ gives you full control, and you pay for it in detail: thread indices, shared-memory layouts, barriers, bank conflicts, vector loads, tensor-core fragment layouts. Triton (Tillet et al., 2019, now part of PyTorch) raises the level of abstraction by one step. You write the program for one block of data, using operations on whole tiles (tl.load a tile, tl.dot two tiles, tl.sum along an axis), and the compiler decides how threads map onto them, stages data through shared memory, coalesces loads and uses tensor cores. You still choose the tiling, which is the decision that matters most, but the bookkeeping disappears.

The trade-off: Triton reaches 80-100% of hand-written CUDA for most kernels an inference engine needs, in a fraction of the code. PyTorch’s compiler emits Triton, and most new kernels in vLLM and SGLang start as Triton. When it falls short (warp-specialized FlashAttention 3, the very latest hardware features), engines drop to CUDA, CUTLASS or newer tile languages.

The programming model

A Triton kernel is a Python function decorated with @triton.jit. It’s launched over a grid of program instances, and each instance asks which one it is with tl.program_id(axis):

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    """(Your engine: Chapter 14)"""
    pid = tl.program_id(0)                       # which block of BLOCK elements is mine
    offsets = pid * BLOCK + tl.arange(0, BLOCK)  # a vector of BLOCK indices
    mask = offsets < n                           # the last block may run past the end
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)


def vector_add(x, y, block=1024):
    check_device(x, y)
    if x.shape != y.shape or not x.is_contiguous() or not y.is_contiguous():
        raise ValueError("vector_add needs equal-shape contiguous tensors")
    out = torch.empty_like(x)
    n = x.numel()
    add_kernel[(triton.cdiv(n, block),)](x, y, out, n, BLOCK=block)
    return out

Read it as “program pid handles elements [pid*BLOCK, (pid+1)*BLOCK)”:

• tl.arange(0, BLOCK) is a vector of BLOCK indices, not a loop. BLOCK must be a compile-time constant (tl.constexpr) and a power of two.
• Pointers support arithmetic: x_ptr + offsets is a vector of addresses.
• mask disables lanes past the end. Every load and store of a partial block needs it, and masked loads take an other= fill value.
• The launch add_kernel[grid](...) takes a grid tuple, computed with triton.cdiv.

The first call compiles the kernel for the given constants and dtypes (seconds); later calls reuse the compiled binary. Time steady state, not the first call.

Tip

No GPU? Run Triton on the CPU. Set TRITON_INTERPRET=1 before importing triton and pass CPU tensors. The interpreter executes the kernel with NumPy, about 100-1,000x slower but numerically faithful, and you can put print and breakpoints inside it. The book’s tests switch it on automatically when no GPU is present. That’s how every Triton kernel in this book was validated during writing.

Fusion: the optimization that matters most for memory-bound code

Consider softmax over the rows of a [4096, 32000] FP32 matrix (512 MB), written as separate PyTorch operations:

m = x.max(dim=-1, keepdim=True)      # read x (512 MB)            write m (small)
e = torch.exp(x - m)                 # read x, write x - m, read it back, write e (4 passes over 512 MB)
s = e.sum(dim=-1, keepdim=True)      # read e
y = e / s                            # read e, write y

That’s around 8 full passes over memory for an operation whose minimum is 2 (read x once, write y once). Since softmax is memory-bound (Chapter 10), runtime is proportional to passes, so a fused kernel that keeps each row in registers is about 4x faster:

@triton.jit
def softmax_kernel(x_ptr, out_ptr, row_stride, width, BLOCK: tl.constexpr):
    """(Your engine: Chapter 14)"""
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < width
    x = tl.load(x_ptr + row * row_stride + cols, mask=mask, other=-float("inf")).to(tl.float32)
    x = x - tl.max(x, axis=0)                    # stable: largest exponent is exp(0) = 1
    num = tl.exp(x)
    tl.store(out_ptr + row * row_stride + cols, num / tl.sum(num, axis=0), mask=mask)


def softmax(x):
    """Row softmax over the last dim of a contiguous 2-D tensor, one program per row.
    Reads x once and writes once: the unfused version (max, sub, exp, sum, div) makes ~5 passes."""
    check_device(x)
    if x.ndim != 2 or not x.is_contiguous():
        raise ValueError("softmax expects a contiguous matrix")
    out = torch.empty_like(x)
    block = triton.next_power_of_2(x.shape[1])
    if block > 65536:
        raise ValueError("Rows wider than 65536 need a multi-pass (online) kernel")
    softmax_kernel[(x.shape[0],)](x, out, x.stride(0), x.shape[1], BLOCK=block)
    return out

One program per row: load the row (padded with $-\infty$, the identity for max), subtract the max, exponentiate, sum, divide, and store. The intermediates never touch memory. Fusion is how almost every non-GEMM operation in a production engine is implemented: normalization, activations, rotary embeddings, residual additions, sampling.

Fused residual add + RMSNorm

Every transformer layer does x = x + branch_output followed by h = rmsnorm(x). Unfused, that’s a read of both inputs, a write of the sum, a read of the sum for the statistic, and another read to normalize, then a write. Fused, it’s two reads and two writes, all in one kernel:

@triton.jit
def add_rmsnorm_kernel(x_ptr, res_ptr, w_ptr, out_ptr, res_out_ptr, width, eps,
                       HAS_RESIDUAL: tl.constexpr, BLOCK: tl.constexpr):
    """(Your engine: Chapter 14)"""
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < width
    x = tl.load(x_ptr + row * width + cols, mask=mask, other=0.0).to(tl.float32)
    if HAS_RESIDUAL:                             # fuse "x = x + residual" into the same pass
        x += tl.load(res_ptr + row * width + cols, mask=mask, other=0.0).to(tl.float32)
        tl.store(res_out_ptr + row * width + cols, x, mask=mask)
    inv_rms = 1.0 / tl.sqrt(tl.sum(x * x, axis=0) / width + eps)
    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    tl.store(out_ptr + row * width + cols, x * inv_rms * w, mask=mask)


def rmsnorm(x, weight, eps=1e-6, residual=None):
    """RMSNorm(x [+ residual]) * weight over the last dim. With a residual, also returns x + residual,
    so a transformer layer's 'add then normalize' costs one memory pass instead of three."""
    check_device(x, weight)
    shape = x.shape
    x2 = x.reshape(-1, shape[-1]).contiguous()
    out = torch.empty_like(x2)
    res_out = torch.empty_like(x2) if residual is not None else out
    res = residual.reshape(-1, shape[-1]).contiguous() if residual is not None else x2
    add_rmsnorm_kernel[(x2.shape[0],)](x2, res, weight, out, res_out, shape[-1], eps,
                                       HAS_RESIDUAL=residual is not None,
                                       BLOCK=triton.next_power_of_2(shape[-1]))
    if residual is not None:
        return out.view(shape), res_out.view(shape)
    return out.view(shape)

HAS_RESIDUAL is a constexpr, so Triton compiles two specialized versions and the if costs nothing at runtime. You’ll swap this kernel into the engine in Chapter 19.

Matmul in Triton

The tiled algorithm from Chapter 12, at tile level. Each program computes one [BLOCK_M, BLOCK_N] output tile by looping over K in BLOCK_K steps, with tl.dot doing the tile product, on tensor cores when the dtypes allow:

@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K,
                  stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
                  BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
                  GROUP_M: tl.constexpr):
    """(Your engine: Chapter 14)"""
    # Map the 1-D program id to an output tile. Walking GROUP_M tile-rows at a time keeps the
    # B tiles those rows need hot in L2 cache ("grouped ordering", Triton tutorial 03).
    pid = tl.program_id(0)
    tiles_m, tiles_n = tl.cdiv(M, BLOCK_M), tl.cdiv(N, BLOCK_N)
    group = pid // (GROUP_M * tiles_n)
    first_m = group * GROUP_M
    group_size = min(tiles_m - first_m, GROUP_M)
    pid_m = first_m + (pid % group_size)
    pid_n = (pid % (GROUP_M * tiles_n)) // group_size

    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    rk = tl.arange(0, BLOCK_K)
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)     # accumulate in FP32 whatever the inputs
    for k0 in range(0, K, BLOCK_K):
        a = tl.load(a_ptr + rm[:, None] * stride_am + (k0 + rk)[None, :] * stride_ak,
                    mask=(rm[:, None] < M) & ((k0 + rk)[None, :] < K), other=0.0)
        b = tl.load(b_ptr + (k0 + rk)[:, None] * stride_bk + rn[None, :] * stride_bn,
                    mask=((k0 + rk)[:, None] < K) & (rn[None, :] < N), other=0.0)
        acc += tl.dot(a, b)                                    # tensor cores when dtypes allow
    c = acc.to(c_ptr.dtype.element_ty)
    tl.store(c_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn, c,
             mask=(rm[:, None] < M) & (rn[None, :] < N))


def matmul(a, b, block_m=64, block_n=64, block_k=32, group_m=8):
    check_device(a, b)
    if a.ndim != 2 or b.ndim != 2 or a.shape[1] != b.shape[0]:
        raise ValueError("matmul expects [M, K] @ [K, N]")
    m, k = a.shape
    n = b.shape[1]
    c = torch.empty((m, n), device=a.device, dtype=a.dtype)
    grid = (triton.cdiv(m, block_m) * triton.cdiv(n, block_n),)
    matmul_kernel[grid](a, b, c, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1),
                        c.stride(0), c.stride(1), BLOCK_M=block_m, BLOCK_N=block_n,
                        BLOCK_K=block_k, GROUP_M=group_m)
    return c

The grouped ordering at the top deserves a word. The GPU runs programs roughly in pid order. If consecutive programs walked along one row of output tiles, each would need a different column strip of B, and B’s tiles would be evicted from L2 before neighbouring rows reused them. Grouping GROUP_M tile-rows together makes nearby programs share B tiles while they’re still cached. This often gives a 10-20% speedup for free.

Autotuning

The best BLOCK_M, BLOCK_N, BLOCK_K, number of warps and number of pipeline stages depend on the GPU and the matrix shape. Triton can try several and remember the fastest per shape:

@triton.autotune(configs=[
    triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, num_warps=8, num_stages=3),
    triton.Config({"BLOCK_M": 64,  "BLOCK_N": 128, "BLOCK_K": 64}, num_warps=4, num_stages=4),
    triton.Config({"BLOCK_M": 32,  "BLOCK_N": 64,  "BLOCK_K": 64}, num_warps=4, num_stages=4),
], key=["M", "N", "K"])
@triton.jit
def matmul_kernel(...): ...

num_stages controls software pipelining: loading the next K-tile while computing on the current one, which hides memory latency behind arithmetic (PMPP §15.8).

Let the compiler write it: torch.compile

PyTorch’s compiler traces your model, fuses chains of elementwise and reduction operations, and generates Triton kernels. Wrapping a function in torch.compile often gives much of the benefit of hand fusion with no kernel code:

fast_norm = torch.compile(lambda x, r, w: rmsnorm_reference(x + r, w))
# TORCH_LOGS=output_code python script.py   prints the Triton it generated

Reading the generated code is one of the best ways to learn Triton idioms. When does writing your own still pay? When the fusion crosses boundaries the compiler won’t cross (attention, quantized matmuls, paged memory access), when you need a specific algorithm (online softmax), or when you need control over memory layout.

Integrating kernels into an engine

Keep a readable PyTorch reference next to every custom kernel, and select between them in one place:

def rmsnorm(x, w, eps, backend="torch"):
    if backend == "triton" and x.is_cuda and x.shape[-1] <= 65536:
        return triton_basics.rmsnorm(x, w, eps)
    return reference_rmsnorm(x, w, eps)       # always-correct fallback

Then test at two levels. First the kernel against the reference, on awkward shapes (odd widths, one row, the largest supported width). Then the whole model’s logits with the kernel swapped in. A correct kernel can still be called with the wrong axis or a non-contiguous input. Record which backend produced every benchmark number.

Fill a table like this with measurements, not expectations. Decode-sized inputs often benefit less than you’d hope: launch overhead, not memory traffic, dominates tiny kernels (Chapter 19).

Build it

pytest tests/test_ch14_triton.py            # on CPU this uses TRITON_INTERPRET=1 automatically
python run.py kernels --impl engine

Stretch exercises

1. ★ Add a GELU epilogue to the matmul kernel (apply GELU to acc before storing) and test it against F.gelu(a @ b). That’s a fused linear + activation. Where: matmul_kernel in engine/kernels/triton_matmul.py.
2. ★★ Add @triton.autotune to the matmul and plot TFLOP/s against torch.matmul for square BF16 sizes 512-8,192. Where: matmul_kernel and its launcher in engine/kernels/triton_matmul.py.
3. ★★ Write a Triton RoPE kernel that rotates Q and K in place for given positions (Chapter 17 explains RoPE), and test it against the reference apply_rope. Where: new engine/kernels/triton_rope.py; compare with the provided izh.qwen3.apply_rope.
4. ★★★ Write a softmax for rows too long for one block: two passes, or an online version that carries a running max and sum (preview of Chapter 15). Where: softmax_kernel and softmax in engine/kernels/triton_basics.py.

Check your understanding

1. What does one Triton program instance correspond to, and how does it differ from one CUDA thread?
2. Why does fusing a softmax make it several times faster, even though it does the same arithmetic?
3. Why must masked loads in a max-reduction use other=-inf?
4. Why should you keep a PyTorch reference path alongside every custom kernel?

Going deeper

• GPU Mode L14 (Umer Adil, A Practitioner’s Guide to Triton: notebook with debugging tips and the interpreter), L18 (fused kernels), L28 (Liger kernels: fused RMSNorm, RoPE and cross-entropy in production), L29 (Triton internals).
• The official Triton tutorials: vector add, fused softmax, matmul with grouped ordering and autotuning, layer norm.
• Tillet, Kung and Cox, Triton: an intermediate language and compiler for tiled neural network computations (2019).
• PMPP §15.8 (software pipelining) for what num_stages does underneath.

15. FlashAttention

The score matrix is the problem

Standard attention (Chapter 5) computes $S = QK^\top/\sqrt d$, then $P = \operatorname{softmax}(S)$, then $O = PV$. For a sequence of $T$ tokens, $S$ and $P$ are $T \times T$ per head. At $T = 8{,}192$, each is 67 million entries: 256 MiB per head in FP32. With 32 heads and a batch of 8, the intermediate scores alone would need 64 GiB.

Even when the memory fits, the traffic kills performance. Unfused, $S$ is written to HBM, read back for the softmax, $P$ is written, then read again for $PV$: four passes over $T^2$ numbers. All that just to produce an output of size $T \times d$, with $d$ typically 64-256. By Chapter 10’s roofline, unfused attention is memory-bound at long sequence lengths even though its arithmetic is large.

FlashAttention (Dao et al., 2022) computes exactly the same output without ever storing $S$ or $P$ in global memory. It tiles $Q$, $K$ and $V$ into on-chip memory and uses a clever recurrence for the softmax. The output matches standard attention up to rounding; it isn’t an approximation. Sparse and linear attention (Chapters 28-29) are approximations, or different models altogether. This isn’t.

The obstacle: softmax needs the whole row

Tiling a matmul works because a dot product is a sum: process $K$ in chunks and add up partial sums. Softmax is harder. Each weight is $e^{s_j}/\sum_k e^{s_k}$, and the denominator needs every score in the row. The stable version also subtracts the row’s maximum, which you don’t know until you’ve seen every score. It looks like you need the whole row first.

Online softmax: normalize as you go

The trick (Milakov and Gimelshein, 2018) is to keep a running answer that’s correct for the scores seen so far, and to fix it up when new scores arrive. For one query row, keep three running values:

• $m$: the largest score seen so far,
• $\ell$: the sum of $e^{s_j - m}$ over scores seen so far,
• $a$: the sum of $e^{s_j - m}, v_j$ over scores seen so far (a vector of length $d$, not yet normalized).

When a new tile of scores arrives with a larger maximum $m_{\text{new}}$, every old term $e^{s_j - m_{\text{old}}}$ must become $e^{s_j - m_{\text{new}}}$. Since

$$ e^{s_j - m_{\text{new}}} = e^{s_j - m_{\text{old}}} \cdot e^{m_{\text{old}} - m_{\text{new}}}, $$

every old term is fixed by multiplying by the same factor $\alpha = e^{m_{\text{old}} - m_{\text{new}}}$. So the update for a tile with scores $s$ and values $V_{\text{tile}}$ is:

$$ \begin{aligned} m_{\text{new}} &= \max!\big(m_{\text{old}},, \max(s)\big) \ p &= e^{,s - m_{\text{new}}} \ \ell_{\text{new}} &= \alpha,\ell_{\text{old}} + \textstyle\sum p \ a_{\text{new}} &= \alpha, a_{\text{old}} + p, V_{\text{tile}} \end{aligned} \qquad\text{and at the end}\qquad O = a / \ell . $$

Both $\ell$ and $a$ must be rescaled. Forgetting one mixes terms measured against different references, a classic bug.

Worked example

One query, scores $[0, 1, 2, -1]$ against scalar values $[2, 4, 1, 3]$, processed in two tiles of two.

Tile 1, scores $[0, 1]$: $m = 1$, $p = [e^{-1}, e^0] = [0.3679, 1]$, $\ell = 1.3679$, $a = 0.3679\cdot 2 + 1\cdot 4 = 4.7358$.

Tile 2, scores $[2, -1]$: the maximum rises to $m = 2$, so $\alpha = e^{1-2} = 0.3679$. $p = [e^0, e^{-3}] = [1, 0.0498]$.

$$ \ell = 0.3679 \times 1.3679 + 1 + 0.0498 = 1.5530, \qquad a = 0.3679 \times 4.7358 + 1\cdot 1 + 0.0498 \cdot 3 = 2.8915 . $$

$O = 2.8915 / 1.5530 = 1.8619$, exactly what softmax([0,1,2,-1]) @ [2,4,1,3] gives. Your milestone tests check this example.

The algorithm in plain PyTorch

Here’s the recurrence over key tiles, for all queries and heads at once, with causal masking by position:

def online_attention(q, k, v, query_positions=None, key_positions=None, tile=16):
    """The same result as causal_attention, computed one key tile at a time without ever
    holding the full [T, S] score matrix. This is FlashAttention's recurrence in plain
    PyTorch.  (Your engine: Chapter 15)

    Running state per query row: m (max score so far), l (sum of exp(score - m)),
    acc (sum of exp(score - m) * value). A new tile with a larger maximum rescales the old
    l and acc by exp(m_old - m_new) so every term shares one exponent reference.
    """
    if tile < 1:
        raise ValueError("tile must be positive")
    batch, q_heads, t, d = q.shape
    kv_heads, s = k.shape[1], k.shape[2]
    if q_heads % kv_heads:
        raise ValueError("Query heads must be a multiple of KV heads")
    group = q_heads // kv_heads
    k = k.repeat_interleave(group, dim=1) if group > 1 else k
    v = v.repeat_interleave(group, dim=1) if group > 1 else v
    qp = _positions(query_positions, batch, t, s - t, q.device)
    kp = _positions(key_positions, batch, s, 0, q.device)
    q32 = q.float() / math.sqrt(d)
    m = torch.full((batch, q_heads, t, 1), float("-inf"), device=q.device)
    l = torch.zeros((batch, q_heads, t, 1), device=q.device)
    acc = torch.zeros((batch, q_heads, t, v.shape[-1]), device=q.device)
    for start in range(0, s, tile):
        stop = min(start + tile, s)
        scores = q32 @ k[:, :, start:stop].float().transpose(-2, -1)
        visible = kp[:, None, None, start:stop] <= qp[:, None, :, None]
        scores = scores.masked_fill(~visible, float("-inf"))
        m_new = torch.maximum(m, scores.amax(-1, keepdim=True))
        # A row that has seen no visible key yet keeps m = -inf; use 0 as a safe reference.
        reference = torch.where(torch.isfinite(m_new), m_new, torch.zeros_like(m_new))
        alpha = torch.exp(m - reference)               # exp(-inf) = 0 for the first visible tile
        p = torch.exp(scores - reference)
        l = alpha * l + p.sum(-1, keepdim=True)
        acc = alpha * acc + p @ v[:, :, start:stop].float()
        m = m_new
    if bool((l == 0).any()):
        raise ValueError("A query row has no visible key")
    return (acc / l).to(q.dtype)

Two edge cases need care. A query row may see no visible key in an early tile (with causal masking, the first tiles of a long query block). Then $m = -\infty$, and $e^{-\infty - (-\infty)}$ is NaN. The code uses a safe reference of 0 for rows that haven’t seen any score yet. And if a row never sees a key at all, the result is undefined, so the function raises.

The same recurrence for one query and one head, in C++ and Rust:

inline std::vector<float> online_attention(const std::vector<float>& q, const std::vector<float>& k,
                                           const std::vector<float>& v, size_t S, size_t d, size_t tile) {
    float m = -INFINITY, l = 0, scale = 1.f / std::sqrt(float(d));
    std::vector<float> acc(d, 0.f);
    for (size_t start = 0; start < S; start += tile) {
        size_t end = std::min(S, start + tile);
        std::vector<float> s(end - start);
        float m_new = m;
        for (size_t j = start; j < end; ++j) {
            s[j - start] = std::inner_product(q.begin(), q.end(), k.begin() + j * d, 0.f) * scale;
            m_new = std::max(m_new, s[j - start]);
        }
        float alpha = std::exp(m - m_new);      // rescale what was accumulated under the old max
        l *= alpha;
        for (float& a : acc) a *= alpha;
        for (size_t j = start; j < end; ++j) {
            float p = std::exp(s[j - start] - m_new);
            l += p;
            for (size_t e = 0; e < d; ++e) acc[e] += p * v[j * d + e];
        }
        m = m_new;
    }
    for (float& a : acc) a /= l;
    return acc;
}

#![allow(unused)]
fn main() {
/// One query, one head, keys processed in tiles of `tile` with a running max m, running
/// denominator l and an unnormalized accumulator acc (Chapter 15). Equals softmax(qK^T)V.
pub fn online_attention(q: &[f32], k: &[f32], v: &[f32], s: usize, d: usize, tile: usize) -> Vec<f32> {
    let scale = 1.0 / (d as f32).sqrt();
    let (mut m, mut l) = (f32::NEG_INFINITY, 0.0f32);
    let mut acc = vec![0.0f32; d];
    for start in (0..s).step_by(tile) {
        let end = (start + tile).min(s);
        let scores: Vec<f32> = (start..end)
            .map(|j| q.iter().zip(&k[j * d..][..d]).map(|(a, b)| a * b).sum::<f32>() * scale)
            .collect();
        let m_new = scores.iter().cloned().fold(m, f32::max);
        let alpha = (m - m_new).exp(); // rescales everything accumulated so far
        l *= alpha;
        acc.iter_mut().for_each(|a| *a *= alpha);
        for (j, sc) in (start..end).zip(&scores) {
            let p = (sc - m_new).exp();
            l += p;
            for (a, vv) in acc.iter_mut().zip(&v[j * d..][..d]) {
                *a += p * vv;
            }
        }
        m = m_new;
    }
    acc.iter().map(|a| a / l).collect()
}
}

This PyTorch version saves memory but not time: each tile is still several separate kernels launched from a Python loop. The speed comes from doing a whole tile’s work inside one kernel, with the tile in on-chip memory.

The kernel

A FlashAttention forward kernel assigns each program a block of BLOCK_M queries for one (batch, head). The program loads its queries once, then streams BLOCK_N keys and values at a time through on-chip memory, applying the recurrence:

@triton.jit
def flash_fwd_kernel(q_ptr, k_ptr, v_ptr, o_ptr,
                     sqb, sqh, sqt, skb, skh, sks, svb, svh, svs, sob, soh, sot,
                     T, S, offset, heads, group, scale,
                     D: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    """(Your engine: Chapter 15)"""
    pid_m = tl.program_id(0)                 # which tile of queries
    bh = tl.program_id(1)                    # which (batch, query head)
    b = bh // heads
    h = bh % heads
    kvh = h // group                         # GQA: several query heads read one KV head
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rd = tl.arange(0, D)
    q = tl.load(q_ptr + b * sqb + h * sqh + rm[:, None] * sqt + rd[None, :], mask=rm[:, None] < T, other=0.0)
    q_pos = rm + offset                      # absolute position of each query row
    m = tl.full((BLOCK_M,), -float("inf"), tl.float32)
    l = tl.zeros((BLOCK_M,), tl.float32)
    acc = tl.zeros((BLOCK_M, D), tl.float32)
    # Causality: no key beyond the last query's position is ever needed.
    last_key = tl.minimum(S, offset + (pid_m + 1) * BLOCK_M)
    for start in range(0, last_key, BLOCK_N):
        rn = start + tl.arange(0, BLOCK_N)
        k = tl.load(k_ptr + b * skb + kvh * skh + rn[:, None] * sks + rd[None, :], mask=rn[:, None] < S, other=0.0)
        v = tl.load(v_ptr + b * svb + kvh * svh + rn[:, None] * svs + rd[None, :], mask=rn[:, None] < S, other=0.0)
        # input_precision="ieee": FP32 inputs get true FP32 math, not TF32 (BF16/FP16 are unaffected).
        s = tl.dot(q, tl.trans(k), input_precision="ieee").to(tl.float32) * scale   # [BLOCK_M, BLOCK_N]
        visible = (rn[None, :] <= q_pos[:, None]) & (rn[None, :] < S)
        s = tl.where(visible, s, -float("inf"))
        m_new = tl.maximum(m, tl.max(s, axis=1))
        m_safe = tl.where(m_new == -float("inf"), 0.0, m_new)           # rows with nothing visible yet
        alpha = tl.exp(m - m_safe)                                        # rescale the old partial sums
        p = tl.exp(s - m_safe[:, None])
        l = alpha * l + tl.sum(p, axis=1)
        acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v, input_precision="ieee").to(tl.float32)
        m = m_new
    out = acc / l[:, None]
    tl.store(o_ptr + b * sob + h * soh + rm[:, None] * sot + rd[None, :], out.to(o_ptr.dtype.element_ty),
             mask=rm[:, None] < T)


def flash_attention(q, k, v, query_offset=None, block_m=32, block_n=32):
    """q [B, Hq, T, D], k/v [B, Hkv, S, D]; query i sits at absolute position offset + i
    (default offset = S - T, i.e. the queries are the newest tokens). D must be a power of two >= 16."""
    check_device(q, k, v)
    batch, heads, t, d = q.shape
    kv_heads, s = k.shape[1], k.shape[2]
    if heads % kv_heads or d & (d - 1) or d < 16:
        raise ValueError("Need Hq % Hkv == 0 and a power-of-two head dim >= 16")
    offset = s - t if query_offset is None else query_offset
    q, k, v = (x.contiguous() for x in (q, k, v))
    o = torch.empty_like(q)
    grid = (triton.cdiv(t, block_m), batch * heads)
    flash_fwd_kernel[grid](q, k, v, o, *q.stride()[:3], *k.stride()[:3], *v.stride()[:3], *o.stride()[:3],
                           t, s, offset, heads, heads // kv_heads, 1.0 / math.sqrt(d),
                           D=d, BLOCK_M=block_m, BLOCK_N=block_n)
    return o

Details worth noticing:

• Causal block skipping. No key past the last query’s position is needed, so the loop stops at offset + (pid_m + 1) * BLOCK_M. That halves the work for causal attention.
• Query offset. Query row i sits at absolute position offset + i. With offset = S - T, one kernel serves full prefill (T = S), chunked prefill (T < S) and decode (T = 1), exactly like your Chapter 5 position rule.
• Grouped-query attention for free. Query head h reads KV head h // group. No repeated K/V tensors are materialized.
• FP32 statistics, low-precision inputs. m, l and acc stay in FP32 registers. p is cast to the value dtype for the second tl.dot, so both matmuls run on tensor cores.
• input_precision="ieee". On NVIDIA GPUs tl.dot runs FP32 inputs on TF32 tensor cores by default (Chapter 13’s 10-bit mantissa, about 1e-3 error), so an FP32 test against causal_attention fails at a 1e-4 tolerance even though the interpreter passes. "ieee" asks for true FP32 math. It changes only FP32 inputs: BF16 and FP16 still run on tensor cores.

How much memory traffic does this save? Each query block streams the visible K and V once, so the kernel reads about $2T^2 d/\text{BLOCK_M}$ elements of K and V. With $d = 128$ and BLOCK_M = 128, that’s about $2T^2$, against roughly $4T^2$ score elements moved by the unfused version (plus its own K/V reads). The bigger wins are elsewhere: nothing of size $T^2$ is ever allocated or written, and the time moves out of memory-bound elementwise kernels into two compute-bound matmuls per tile.

Note

Production FlashAttention adds much more. FlashAttention-2 reorders loops to parallelize over the sequence and cut non-matmul work. FlashAttention-3 uses Hopper’s asynchronous copies (TMA) and warp specialization (some warps load, others compute) and supports FP8. FlashInfer and vLLM’s kernels add paged KV caches (Chapter 25) and many variants. GPU Mode L12 (Thomas Viehmann’s FlashAttention lecture, with a from-scratch CUDA version) and L36 (FlashAttention-3 with CUTLASS) cover them.

Decode is different: flash-decoding

During decode, there’s one query per sequence and a long KV cache. One program per (batch, head) would leave most of the GPU idle: with batch 1 and 8 KV heads, that’s 8 programs for 132 SMs. Flash-decoding splits the KV sequence into chunks processed by different programs in parallel. Each produces a partial result $(m_i, \ell_i, a_i)$ for its chunk, and a second small kernel merges them with the same rescaling rule:

$$ m = \max_i m_i,\qquad \ell = \sum_i e^{m_i - m},\ell_i,\qquad a = \sum_i e^{m_i - m},a_i,\qquad O = a/\ell . $$

This merge rule is associative: any number of partial softmaxes can be combined in any grouping. That same property powers ring attention across multiple GPUs (GPU Mode L13) and the paged decode kernel in Chapter 25.

What FlashAttention does and doesn’t change

• Changes: memory for intermediates goes from $O(T^2)$ to $O(T)$, and HBM traffic drops several-fold, so long-sequence attention becomes compute-bound and fast.
• Doesn’t change: the result, up to rounding, and the arithmetic. Attention still costs $O(T^2 d)$ FLOPs. At very long contexts that quadratic arithmetic itself becomes the bottleneck, which is the motivation for Chapters 28 and 29.

Build it

pytest tests/test_ch15_flash.py

Stretch exercises

1. ★ Instrument causal_attention and online_attention with torch.cuda.max_memory_allocated() for T = 4,096. How much memory does the online version save? Where: experiments/ch15.py (create it), importing both functions from engine.attention.
2. ★★ Implement the flash-decoding merge: split online_attention’s key range into 4 chunks, compute each chunk’s $(m, \ell, a)$ separately, and merge. Verify against the unsplit result. Where: add a split/merge attention helper beside online_attention in engine/attention.py.
3. ★★ Add a sliding-window option to the Triton kernel: skip tiles entirely before q_pos - window, and mask within the boundary tile. Where: flash_fwd_kernel and flash_attention in engine/kernels/triton_flash.py.
4. ★★★ Write the backward pass of attention in PyTorch, using FlashAttention’s trick: recompute $P$ from the saved log-sum-exp per row instead of storing it. Check it against autograd. Where: add a backward helper in engine/attention.py; compare with autograd in tests/test_ch15_stretch.py.

Check your understanding

1. Why is the accumulator $a$ kept unnormalized until the very end?
2. Why is applying an ordinary softmax to each tile independently incorrect?
3. Which complexity does FlashAttention reduce: arithmetic, intermediate memory, or both?
4. Why does decode need a different parallelization (flash-decoding) than prefill?

Going deeper

• PMPP §20.5 (pp. 492-503): FlashAttention, derived and implemented in CUDA; §20.6 (KV-cache arithmetic intensity) previews Chapter 16.
• GPU Mode L12 (FlashAttention, with flash_attention.cu and a notebook), L13 (ring attention and the log-sum-exp merge, with howto_log_sum_exp.ipynb), L36 (FlashAttention-3).
• Dao et al., FlashAttention (2022) and FlashAttention-2 (2023); Shah et al., FlashAttention-3 (2024); Milakov and Gimelshein, Online normalizer calculation for softmax (2018).
• The Triton tutorial Fused Attention, which this chapter’s kernel simplifies.

16. Prefill, decode and the KV cache

Generation without a cache repeats itself

Chapter 8’s loop with cached=False re-runs the entire sequence for every new token. With a prompt of $P$ tokens, generating $N$ tokens processes $P$, then $P+1$, …, up to $P+N-1$ positions. Each of those passes recomputes the same projections, MLPs and attention for every earlier position, getting exactly the same numbers as last time. The total work grows like $N \cdot P + N^2/2$, so it’s quadratic in the output length.

Measured on the reference engine (a tiny 4-layer Qwen3 on a laptop CPU, python run.py cache):

The uncached time grows quadratically and the cached time roughly linearly. Real models with thousand-token prompts make the gap enormous.

The observation: the past doesn’t change

In a causal model, position $t$’s hidden states depend only on tokens $\le t$. Appending a token can’t change anything computed for earlier positions. So their keys and values, the only things later tokens read from them, are fixed once computed.

A new token needs, at every layer: its own query, and the keys and values of all positions so far. So store each layer’s K and V as they’re computed, and on the next step process only the new token: compute its Q, K and V, append K and V to the cache, and attend over the whole cache. That’s the KV cache. Two things it deliberately does not store:

• Old queries. A query is used only once, by its own position.
• Old hidden states or MLP results. They’re never read again; only K and V are.

Prefill and decode

Generation now has two phases, as Chapter 1 promised:

prefill:   input IDs at positions 0,1,2       write K/V at 0,1,2      last logits -> token 3
decode 1:  input token 3 at position 3        write K/V at 3          logits -> token 4
decode 2:  input token 4 at position 4        write K/V at 4          logits -> token 5

Two off-by-one facts trip people up:

1. The first new token comes from prefill. The prompt’s last position already predicts it. No separate decode step is needed before the first sample.
2. The last generated token is never fed back unless you want another prediction. After $N$ new tokens, decode has run $N-1$ times.

For a prompt of 3 tokens and 3 generated tokens, the uncached loop projects $3 + 4 + 5 = 12$ token positions; the cached one projects $3 + 1 + 1 = 5$.

Layout and validity

For each layer, store keys and values as [B, H_kv, capacity, D_h], preallocated to the maximum length you’ll need, plus a count of valid positions:

class KVCache:
    """Contiguous preallocated cache; every row in the batch has the same length.  (Your engine: Chapter 16)"""

    def __init__(self, layers, batch, kv_heads, capacity, head_dim, device="cpu", dtype=torch.float32):
        self.capacity = capacity
        shape = (batch, kv_heads, capacity, head_dim)
        self.keys = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(layers)]
        self.values = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(layers)]
        self.lengths = [0] * layers

    @property
    def length(self):
        """Tokens already processed into every layer."""
        if len(set(self.lengths)) != 1:
            raise RuntimeError("A forward pass stopped part-way; reset or truncate the cache")
        return self.lengths[0]

    def update(self, layer, k, v, positions=None, rows=None):
        """Write k/v after the existing entries and return views of the whole valid history.  (Your engine: Chapter 16)"""
        if rows is not None:
            raise ValueError("KVCache rows always advance together; use StaticKVCache for slots")
        start = self.lengths[layer]
        end = start + k.shape[2]
        if end > self.capacity:
            raise ValueError(f"KV capacity {self.capacity} exceeded")
        if k.shape[:2] != self.keys[layer].shape[:2] or k.shape[3] != self.keys[layer].shape[3]:
            raise ValueError("Cache batch/head/width shapes differ")
        self.keys[layer][:, :, start:end].copy_(k)
        self.values[layer][:, :, start:end].copy_(v)
        self.lengths[layer] = end
        return self.keys[layer][:, :, :end], self.values[layer][:, :, :end], None

    def truncate(self, length):
        """Forget everything after `length` tokens (speculative-decoding rollback).  (Your engine: Chapter 16)"""
        if not 0 <= length <= self.length:
            raise ValueError("Invalid truncation length")
        self.lengths = [length] * len(self.lengths)

    def reset(self):
        self.truncate(0)

    @property
    def bytes(self):
        return sum(t.numel() * t.element_size() for t in self.keys + self.values)

struct KvCache {
    std::vector<std::vector<float>> keys, values;  // per layer: [positions, kv_heads*head_dim]
    size_t width, len = 0;
    KvCache(size_t layers, size_t width_, size_t capacity) : keys(layers), values(layers), width(width_) {
        for (size_t l = 0; l < layers; ++l) { keys[l].reserve(capacity * width); values[l].reserve(capacity * width); }
    }
    void append(size_t layer, const std::vector<float>& k, const std::vector<float>& v) {
        keys[layer].insert(keys[layer].end(), k.begin(), k.end());
        values[layer].insert(values[layer].end(), v.begin(), v.end());
    }
    void truncate(size_t n) {
        for (size_t l = 0; l < keys.size(); ++l) { keys[l].resize(n * width); values[l].resize(n * width); }
        len = n;
    }
};

#![allow(unused)]
fn main() {
/// Rows are positions; each row holds kv_heads * head_dim values. Capacity is reserved up
/// front so appending never reallocates (and never copies the history).
pub struct KvCache {
    pub keys: Vec<Vec<f32>>,   // one buffer per layer
    pub values: Vec<Vec<f32>>,
    pub width: usize,          // kv_heads * head_dim
    pub len: usize,            // positions already processed
}

impl KvCache {
    pub fn new(layers: usize, width: usize, capacity: usize) -> Self {
        KvCache {
            keys: (0..layers).map(|_| Vec::with_capacity(capacity * width)).collect(),
            values: (0..layers).map(|_| Vec::with_capacity(capacity * width)).collect(),
            width,
            len: 0,
        }
    }

    /// Append one position's K and V for one layer; return the whole history for that layer.
    pub fn append(&mut self, layer: usize, k: &[f32], v: &[f32]) -> (&[f32], &[f32]) {
        self.keys[layer].extend_from_slice(k);
        self.values[layer].extend_from_slice(v);
        (&self.keys[layer], &self.values[layer])
    }

    /// Roll back to `len` positions (speculative decoding).
    pub fn truncate(&mut self, len: usize) {
        for layer in 0..self.keys.len() {
            self.keys[layer].truncate(len * self.width);
            self.values[layer].truncate(len * self.width);
        }
        self.len = len;
    }

    pub fn bytes(&self) -> usize {
        2 * self.keys.len() * self.len * self.width * 4
    }
}
}

Design choices worth noticing:

• Preallocate; don’t concatenate. torch.cat on every step reallocates and copies the whole history, which is $O(T)$ work per token and $O(T^2)$ overall, the very cost the cache was meant to remove. Writing into a preallocated slice is $O(1)$ per token. (Chapter 25’s paged cache gets the same benefit without reserving the maximum up front.)
• Capacity versus valid length. Slots beyond length hold zeros or stale data. update returns a view of just the valid prefix, so attention never sees them.
• Per-layer lengths. Every layer appends during a forward pass. If an exception interrupts a pass halfway, some layers are one token ahead of others. The length property detects this and refuses to continue, rather than silently misaligning positions.
• Truncate for rollback. truncate(n) forgets everything after $n$ tokens without copying. Speculative decoding (Chapter 26) depends on this.

The cache follows a one-method protocol: update(layer, k, v, positions, rows) writes new entries and returns everything attention may read. Your models never touch cache internals. Later caches (static for CUDA graphs in Chapter 19, slot-based for batching in Chapter 24, paged in Chapter 25) implement the same method and plug into the same model code unchanged.

Positions continue, and masks become rectangles

When a decode step feeds token 57, it must be treated as position 57, not position 0. GPT adds position_embedding[57], and Qwen rotates its query and key by angle 57 (Chapter 17). Your models compute positions as cache.length + arange(T).

The mask changes too. Without a cache, queries and keys are the same $T$ positions, and the mask is a square lower triangle. With a cache, $T$ new queries attend to $S$ keys ($S > T$): a rectangle. For a cached prefix of 3 tokens and a chunk of 2 new tokens:

             key 0  1  2  3  4
query pos 3      1  1  1  1  0
query pos 4      1  1  1  1  1

A square-triangle shortcut applied to this 2-row chunk would let query 3 see only key 0 and query 4 only keys 0-1, hiding the most recent history. This is exactly why your causal_attention takes positions (Chapter 5): the rule key_pos <= query_pos produces the right rectangle automatically. (Watch out for PyTorch’s scaled_dot_product_attention(is_causal=True): with $T \ne S$ it uses a top-left-aligned triangle, which is the wrong rectangle for cached decoding.)

Prove equivalence

The cache is an optimization, so it must change nothing except speed. The strongest test feeds one fixed sequence three ways:

1. the whole sequence at once, with no cache;
2. a prefix, then a multi-token chunk, then single tokens, all through one cache;
3. compare all logits position by position.

They must agree to floating-point tolerance (about 1e-5 in FP32). Use fixed token IDs, not generated ones: once two runs pick different tokens, they’re processing different inputs and the comparison is meaningless.

How big is the cache?

Each layer stores K and V for every position, so:

$$ \text{KV bytes} = 2 \times \text{layers} \times \text{batch} \times \text{tokens} \times H_{kv} \times D_h \times \text{bytes per value}. $$

For Qwen3-0.6B (28 layers, 8 KV heads of dimension 128, BF16):

• 112 KiB per token;
• 224 MiB for a 2,048-token conversation;
• 3.5 GiB for a 32k-token context, about three times the model’s own weights.

def kv_cache_bytes(layers, batch, tokens, kv_heads, head_dim, bytes_per_value=2):
    """2 (K and V) x layers x batch x tokens x kv_heads x head_dim x bytes."""
    return 2 * layers * batch * tokens * kv_heads * head_dim * bytes_per_value

Why the cache shapes model design

Two consequences follow from this formula, and they explain several architecture choices you’ll meet:

• Capacity. In serving, the cache limits how many conversations fit on a GPU at once (Chapters 24-25). Note that the formula uses KV heads, not query heads. Grouped-query attention with 8 KV heads instead of 32 shrinks the cache 4x: for an 8B-class model (32 layers, head dimension 128) at 8k tokens, that’s 1 GiB per sequence instead of 4 GiB. Multi-query attention (one KV head) and DeepSeek’s multi-head latent attention push further.
• Bandwidth. Each decode step reads the whole cache of each sequence once, in addition to the weights. Attention during decode has an arithmetic intensity of about 1 FLOP per byte per query head (PMPP §20.6), as memory-bound as the weight reads. At long contexts, the KV reads dominate the step time. That’s what motivates KV-cache quantization (Chapter 20), sparse attention that reads only part of the cache (Chapter 29), and linear attention with a fixed-size state instead of a cache (Chapter 28).

Build it

pytest tests/test_ch16_kv_cache.py
python run.py cache --impl engine

Stretch exercises

1. ★ Replace the preallocated write with torch.cat and measure generation time for 64, 256 and 1,024 new tokens. Where does the quadratic copy cost become visible? Where: a separate concatenating-cache variant of KVCache.update in engine/kv_cache.py; compare it in experiments/ch16.py (create it).
2. ★★ Add a layer_bytes() report and print how much of a 32k-token Qwen3-0.6B decode step’s traffic is KV cache versus weights. Where: add KVCache.layer_bytes in engine/kv_cache.py.
3. ★★ Implement a sliding-window cache that keeps only the last $w$ positions in a ring buffer. What must change in the positions passed to attention? Where: add a ring-buffer cache class in engine/kv_cache.py; wire its absolute positions into Qwen3Attention.forward in engine/qwen3.py.
4. ★★★ Store the cache in FP8 (E4M3, with one scale per head) and dequantize on read. Measure the logit error against a BF16 cache for 1,000 decode steps. Where: add an FP8 cache variant in engine/kv_cache.py; integrate its dequantized reads in engine/qwen3.py.

Check your understanding

1. Why cache keys and values, but not queries?
2. Which forward pass produces the logits for the first generated token?
3. Why can a lower-triangular mask be wrong for a 2-token chunk after a 3-token prefix?
4. Why does grouped-query attention reduce cache size, but not the number of attention score computations?
5. Why does decode at long context become dominated by cache reads?

Going deeper

• PMPP §20.4 (pp. 488-492): KV caching; §20.6 (pp. 504-508): the arithmetic intensity and memory requirement of the KV cache; §20.7: MQA and GQA.
• Pope et al., Efficiently Scaling Transformer Inference (2022): the analysis of cache memory and bandwidth that much of modern serving builds on.
• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need (MQA, 2019); DeepSeek-V2 (2024) for multi-head latent attention.

17. A modern architecture: Qwen3

From GPT-2 to a 2025 model

The transformer block hasn’t changed shape since GPT-2: embed, then $L$ blocks of attention and MLP on a residual stream, then norm and head. What changed is the ingredients. Almost every open model since Llama (2023), including Qwen, Mistral and DeepSeek, made the same set of substitutions:

Dense Qwen3 is the first real model your engine will run. Here’s the contract for the 0.6B checkpoint, from its official config.json:

Notice that the query projection has width $16 \times 128 = 2{,}048$, which is twice the hidden size. Code that assumes head_dim = hidden_size / heads reshapes the projections wrongly and crashes (or worse, doesn’t). The book’s tiny default config deliberately has the same mismatch, so the bug surfaces in the tests.

RMSNorm

LayerNorm subtracts the mean and divides by the standard deviation. RMSNorm (Zhang and Sennrich, 2019) skips the mean and divides by the root mean square:

$$ \operatorname{RMSNorm}(x)_i = \gamma_i, \frac{x_i}{\sqrt{\frac1D\sum_j x_j^2 + \epsilon}} . $$

For $x = [3, 4]$: the mean square is $(9 + 16)/2 = 12.5$, the RMS is 3.5355, and the output is $[0.8485, 1.1314]$ before $\gamma$. The vector’s direction is unchanged, and its RMS becomes 1. It’s one reduction instead of two, with no bias, and models train just as well.

class RMSNorm(nn.Module):
    def __init__(self, width, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(width))
        self.eps = eps

    def forward(self, x):
        """x / sqrt(mean(x^2) + eps) * weight, reduced in FP32.  (Your engine: Chapter 17)"""
        x32 = x.float()
        normalized = x32 * torch.rsqrt(x32.square().mean(-1, keepdim=True) + self.eps)
        return self.weight * normalized.to(x.dtype)

Two parity details matter when loading real weights:

1. Compute the mean of squares in FP32 even for BF16 inputs (Chapter 13).
2. The cast order is part of the contract. Qwen3 normalizes in FP32, casts back to the input dtype, then multiplies by $\gamma$. Gemma and the Qwen3.5 family use $(1+\gamma)$ with $\gamma$ initialized at zero, multiplying before the cast (you’ll meet that “zero-centered” variant in Chapter 29). In BF16, the difference changes the last bits of every layer’s output.

Rotary position embeddings

GPT-2 added a learned position vector to each token. RoPE (Su et al., 2021) instead rotates the query and key vectors by an angle proportional to their position, right before the attention scores are computed.

Take one pair of coordinates and rotate it by angle $\phi$:

$$ R(\phi) = \begin{pmatrix} \cos\phi & -\sin\phi \ \sin\phi & \cos\phi \end{pmatrix}. $$

Rotate a query at position $p$ by $p,\omega$ and a key at position $s$ by $s,\omega$. Their dot product is

$$ \big(R(p\omega),q\big)\cdot\big(R(s\omega),k\big) = q^{\top} R\big((s-p),\omega\big), k , $$

because rotations compose and $R(a)^\top = R(-a)$. The score depends only on the distance $s - p$, not on absolute positions. Shift both tokens by 10 and nothing changes. For example, with $q = k = [1, 0]$ and $\omega = 0.35$, two tokens 2 apart score $\cos(0.7) = 0.7648$, at positions (0, 2) or (10, 12) alike.

A head of dimension $d$ has $d/2$ coordinate pairs, each rotating at its own frequency:

$$ \omega_i = \theta^{-2i/d}, \qquad i = 0, \ldots, d/2 - 1 . $$

With $\theta = 10^6$ and $d = 128$, the fastest pair turns 1 radian per token and the slowest about $1.2 \times 10^{-6}$ radians, a wavelength of about 5 million tokens. Fast pairs resolve nearby order; slow pairs let distant tokens still be distinguished. Raising $\theta$ (Qwen3 uses $10^6$, the original paper $10^4$) stretches every wavelength, which helps long contexts.

The pairing convention is part of the contract

Which coordinates form a pair? The original paper paired adjacent coordinates $(0,1), (2,3), \ldots$ Llama, Qwen and most Hugging Face models pair coordinate $i$ with $i + d/2$, the split-half (“rotate half”) convention. Both are valid RoPEs, and a model is trained with exactly one. Using the other produces correct shapes and fluent-looking garbage. Our implementation uses split-half:

def rope_cos_sin(positions, rotary_dim, theta, dtype=torch.float32):
    """cos/sin tables for absolute positions [T] or [B, T] -> [B or 1, 1, T, rotary_dim].  (Your engine: Chapter 17)

    Frequencies theta^(-2i/rotary_dim) for i in [0, rotary_dim/2), repeated for the two halves.
    """
    if rotary_dim % 2:
        raise ValueError("RoPE needs an even number of rotated dimensions")
    inv_freq = theta ** (-torch.arange(0, rotary_dim, 2, device=positions.device, dtype=torch.float32) / rotary_dim)
    pos = positions.float() if positions.ndim == 2 else positions.float()[None]
    angles = pos[..., None] * inv_freq                    # [B, T, rotary_dim/2]
    angles = torch.cat((angles, angles), dim=-1)[:, None]  # [B, 1, T, rotary_dim]
    return angles.cos().to(dtype), angles.sin().to(dtype)


def rotate_half(x):
    a, b = x.chunk(2, dim=-1)
    return torch.cat((-b, a), dim=-1)


def apply_rope(x, cos, sin):
    """Rotate the first cos.shape[-1] features of x (split-half pairing); pass the rest through.  (Your engine: Chapter 17)

    Pairs are (i, i + rotary_dim/2), NOT adjacent features (2i, 2i+1): that is the checkpoint's
    convention and getting it wrong still produces plausible-looking shapes.
    """
    rotary = cos.shape[-1]
    x_rot, x_pass = x[..., :rotary], x[..., rotary:]
    rotated = x_rot * cos + rotate_half(x_rot) * sin
    return torch.cat((rotated, x_pass), dim=-1) if x_pass.shape[-1] else rotated

// Split-half pairing: feature i rotates with feature i + d/2.
inline void rope(float* x, size_t d, size_t pos, float theta) {
    for (size_t i = 0; i < d / 2; ++i) {
        float freq = std::pow(theta, -2.f * i / d), s = std::sin(pos * freq), c = std::cos(pos * freq);
        float a = x[i], b = x[i + d / 2];
        x[i] = a * c - b * s;
        x[i + d / 2] = a * s + b * c;
    }
}

#![allow(unused)]
fn main() {
/// Rotate one head vector in place for absolute position `pos`, pairing feature i with
/// feature i + d/2 (the split-half convention used by Qwen, Llama and the book's Python code).
pub fn rope(x: &mut [f32], pos: usize, theta: f32) {
    let half = x.len() / 2;
    for i in 0..half {
        let freq = theta.powf(-2.0 * i as f32 / x.len() as f32);
        let (sin, cos) = (pos as f32 * freq).sin_cos();
        let (a, b) = (x[i], x[i + half]);
        x[i] = a * cos - b * sin;
        x[i + half] = a * sin + b * cos;
    }
}
}

apply_rope rotates only the first cos.shape[-1] features and passes the rest through. Qwen3 rotates all 128, but Flash-Next’s attention rotates only the first 64 of 256 (partial RoPE, Chapter 29). The same function handles both.

Where RoPE sits in the computation matters for caching: Q and K are rotated before K is cached. A cached key already contains its position. Never rotate cached keys again; each new token rotates only its own Q and K, at its own position.

Attention with grouped KV heads and Q/K normalization

Qwen3’s attention is your Chapter 5 attention with three additions:

1. Separate Q, K, V projections, no biases, with output widths $H_q D_h$, $H_{kv} D_h$, $H_{kv} D_h$.
2. Q/K norm: an RMSNorm over each head’s 128 features, with a learned $\gamma$ shared across heads, applied to Q and K before RoPE. This keeps attention logits from growing without bound during training, a known source of instability.
3. GQA: 16 query heads share 8 KV heads; your causal_attention handles the mapping.

class Qwen3Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        if cfg.num_attention_heads % cfg.num_key_value_heads:
            raise ValueError("Query heads must be a multiple of KV heads")
        self.cfg = cfg
        d, hq, hkv = cfg.head_dim, cfg.num_attention_heads, cfg.num_key_value_heads
        self.q_proj = nn.Linear(cfg.hidden_size, hq * d, bias=False)
        self.k_proj = nn.Linear(cfg.hidden_size, hkv * d, bias=False)
        self.v_proj = nn.Linear(cfg.hidden_size, hkv * d, bias=False)
        self.o_proj = nn.Linear(hq * d, cfg.hidden_size, bias=False)
        self.q_norm = RMSNorm(d, cfg.rms_norm_eps)
        self.k_norm = RMSNorm(d, cfg.rms_norm_eps)

    def forward(self, x, positions, rope, cache=None, layer=0, rows=None):
        """Project, per-head RMSNorm on q and k, RoPE, cache, GQA attention, output projection.  (Your engine: Chapter 17)"""
        b, t, _ = x.shape
        c = self.cfg
        q = self.q_proj(x).view(b, t, c.num_attention_heads, c.head_dim)
        k = self.k_proj(x).view(b, t, c.num_key_value_heads, c.head_dim)
        v = self.v_proj(x).view(b, t, c.num_key_value_heads, c.head_dim)
        q = self.q_norm(q).transpose(1, 2)     # norm over head_dim, before RoPE
        k = self.k_norm(k).transpose(1, 2)
        v = v.transpose(1, 2)
        cos, sin = rope
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        key_positions = None
        if cache is not None:
            k, v, key_positions = cache.update(layer, k, v, positions, rows)   # keys are cached already rotated
        y = causal_attention(q, k, v, positions, key_positions)
        return self.o_proj(y.transpose(1, 2).reshape(b, t, -1))

SwiGLU

The GPT-2 MLP was expand → GELU → contract. Qwen3’s MLP has two expansions, one of which gates the other (Shazeer, 2020):

$$ \operatorname{SwiGLU}(x) = W_{\text{down}}\big(\operatorname{SiLU}(W_{\text{gate}},x) \odot W_{\text{up}},x\big), \qquad \operatorname{SiLU}(z) = z,\sigma(z). $$

The gate decides, feature by feature, how much of the “up” signal passes. If a gate pre-activation is 0, $\operatorname{SiLU}(0) = 0$ and the feature is shut off. If it’s 2, $\operatorname{SiLU}(2) = 1.7616$, and with an up value of 3 the feature passes $5.2848$. Gated MLPs consistently outperform ungated ones at the same parameter count. To keep that count comparable with a $4D$ GELU MLP, the intermediate size is usually about $\frac{8}{3}D$ (Qwen3-0.6B uses $3D$).

class SwiGLU(nn.Module):
    def __init__(self, width, hidden):
        super().__init__()
        self.gate_proj = nn.Linear(width, hidden, bias=False)
        self.up_proj = nn.Linear(width, hidden, bias=False)
        self.down_proj = nn.Linear(hidden, width, bias=False)

    def forward(self, x):
        """down(silu(gate(x)) * up(x)).  (Your engine: Chapter 17)"""
        return self.down_proj(nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))

The whole model

class Qwen3Layer(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.input_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.self_attn = Qwen3Attention(cfg)
        self.post_attention_layernorm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
        self.mlp = SwiGLU(cfg.hidden_size, cfg.intermediate_size)

    def forward(self, x, positions, rope, cache=None, layer=0, rows=None):
        x = x + self.self_attn(self.input_layernorm(x), positions, rope, cache, layer, rows)
        return x + self.mlp(self.post_attention_layernorm(x))

class Qwen3(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.model = Qwen3Backbone(cfg)
        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
        if cfg.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

    @property
    def context_limit(self):
        return self.cfg.max_position_embeddings

    def forward(self, ids, cache=None, positions=None, rows=None, return_hidden=False):
        """ids [B, T] -> logits [B, T, vocab]. Same cache/position contract as GPT.forward.  (Your engine: Chapter 17)"""
        if positions is None:
            start = cache.length if cache is not None else 0
            if start + ids.shape[1] > self.context_limit:
                raise ValueError("Sequence exceeds max_position_embeddings")
            positions = torch.arange(start, start + ids.shape[1], device=ids.device)
        x = self.model.embed_tokens(ids)
        rope = rope_cos_sin(positions, self.cfg.head_dim, self.cfg.rope_theta, x.dtype)
        for i, layer in enumerate(self.model.layers):
            x = layer(x, positions, rope, cache, i, rows)
        hidden = self.model.norm(x)
        return hidden if return_hidden else self.lm_head(hidden)

    def cache_spec(self):
        c = self.cfg
        return c.num_hidden_layers, c.num_key_value_heads, c.head_dim

    def new_cache(self, batch, capacity):
        if capacity > self.context_limit:
            raise ValueError("Requested cache exceeds the context limit")
        c, p = self.cfg, self.model.embed_tokens.weight
        return KVCache(c.num_hidden_layers, batch, c.num_key_value_heads, capacity, c.head_dim, p.device, p.dtype)

The parameter budget of Qwen3-0.6B, counted from this code:

A quarter of this small model is the embedding table. In bigger models, the 28 × 15.7M of layers dominates.

Why parity tests matter now

From now on, your models must match someone else’s trained weights exactly, and small discrepancies don’t announce themselves. A wrong RoPE pairing, an epsilon of 1e-5 instead of 1e-6, LayerNorm instead of RMSNorm, or norm-then-RoPE swapped to RoPE-then-norm all produce a model that runs, generates English-like text after loading, and is subtly or badly wrong. The milestone tests therefore compare your Qwen3 against Hugging Face’s on small random configurations. Random weights exercise every code path, and a 1e-5 match on random weights is strong evidence of a correct implementation.

Build it

uv pip install -r optional-requirements.txt     # for the comparison with Transformers
pytest tests/test_ch17_qwen3.py
python run.py cache --impl engine                # your Qwen3, cached vs uncached

Stretch exercises

1. ★ Verify numerically that RoPE preserves vector norms, and that the score at distance 0 is just the unrotated dot product. Where: experiments/ch17.py (create it), importing engine.qwen3.rope_cos_sin and apply_rope.
2. ★★ Implement the adjacent-pair RoPE convention and show that loading Qwen3 weights with it changes the logits (by how much?). Where: add an adjacent-pair variant of apply_rope in engine/qwen3.py and select it in Qwen3Attention.forward.
3. ★★ Implement YaRN or linear RoPE scaling (multiply positions by a factor, or interpolate frequencies) and explain what it does to the wavelength table. Where: rope_cos_sin in engine/qwen3.py.
4. ★★★ Write a fused Triton kernel that applies Q/K-norm and RoPE to the Q and K projections in one pass, and plug it into Qwen3Attention behind a backend switch. Where: new engine/kernels/triton_rope.py, selected by Qwen3Attention.forward in engine/qwen3.py.

Check your understanding

1. Why can a correct-looking reshape still be wrong for Qwen3’s attention projections?
2. What happens to a key, in order, before it enters the cache?
3. What property of rotations makes a shared position shift cancel out in a RoPE dot product?
4. Why does SwiGLU have three weight matrices where GPT-2’s MLP has two?
5. Why are random-weight parity tests enough to establish the architecture, before you ever load the real checkpoint?

Going deeper

• BALLM has no Qwen3 chapter, but Raschka’s companion repository does: ch05/11_qwen3 builds Qwen3 from scratch in the book’s style, and ch05/07_gpt_to_llama walks the GPT-2 → Llama changes one by one.
• Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding (2021); Zhang and Sennrich, Root Mean Square Layer Normalization (2019); Shazeer, GLU Variants Improve Transformer (2020).
• Qwen Team, Qwen3 Technical Report (2025); the official modeling_qwen3.py in Transformers.

18. Engine v1: load Qwen3 and chat

From model to engine

You have a model (Chapter 17), a sampler (8) and a cache (16). An engine joins them under a contract:

• Input: token IDs plus generation settings (temperature, top-p, max new tokens, stop tokens, seed).
• Output: generated token IDs, a finish reason (stop or length), and measurements.
• Owned state: the KV cache for the request, created when the request starts and released when it ends, even if it ends with an error.

Text appears only at the boundary. Turning a chat into token IDs and IDs back into text is the tokenizer’s job, and keeping it outside the engine core has practical benefits. The engine is testable with integer inputs alone, the same engine serves different tokenizers, and token-level outputs make parity checks exact.

Loading Qwen3-0.6B

Download the checkpoint (if you didn’t in Chapter 1):

uv pip install -r optional-requirements.txt
hf download Qwen/Qwen3-0.6B --local-dir models/Qwen3-0.6B

Because your parameter names match the checkpoint’s exactly (Chapter 17), loading is a checked copy:

@torch.no_grad()
def load_qwen3(directory, device="cpu", dtype=torch.bfloat16):
    """Dense Qwen3: our parameter names equal the checkpoint's, so this is a checked copy.  (Your engine: Chapter 18)"""
    cfg = Qwen3Config.from_hf(read_config(directory))
    with torch.device("meta"):
        model = Qwen3(cfg)                      # no memory yet: parameters are shapes only
    model = model.to_empty(device=device).to(dtype)
    if cfg.tie_word_embeddings:                 # to_empty breaks the tie; restore it
        model.lm_head.weight = model.model.embed_tokens.weight
    destinations = dict(model.named_parameters())
    loaded, head_alias = set(), None
    for name, value in snapshot_tensors(directory):
        if name == "lm_head.weight" and cfg.tie_word_embeddings:
            head_alias = value
            continue
        if name not in destinations:
            raise ValueError(f"Unexpected tensor {name}")
        assign(destinations[name], value, name)
        loaded.add(name)
    missing = set(destinations) - loaded
    if missing:
        raise ValueError(f"Missing tensors: {sorted(missing)[:5]}")
    if head_alias is not None and not torch.equal(head_alias.to(model.lm_head.weight), model.lm_head.weight):
        raise ValueError("Tied lm_head and embed_tokens disagree")
    return model.eval()

Three engineering points:

• Build on the meta device first. Qwen3(cfg) on the CPU would allocate and randomly initialize 596M parameters, only to overwrite them. Built under torch.device("meta"), parameters are shapes without storage. to_empty(device=...) then allocates uninitialized memory directly on the target device, and the loader fills it shard by shard. Peak memory is about one model plus one shard, not two models.
• Restore weight tying. to_empty creates fresh storage for every parameter, which silently unties the head from the embedding. Re-tie them, and if the checkpoint also contains lm_head.weight, verify it equals the embedding rather than letting file order decide which copy wins.
• Refuse what you don’t implement. Qwen3Config.from_hf rejects other model types, scaled RoPE and sliding windows. A loader that “mostly works” on an unsupported variant is worse than one that refuses.

The text boundary

The tokenizer comes from the checkpoint’s own files. Your engine uses Hugging Face’s implementation through AutoTokenizer, loaded with local_files_only=True so nothing is fetched or executed remotely. (Chapter 4’s BPE explains what it does; reading Qwen’s tokenizer.json with your own BPE is a stretch exercise.)

Three details decide whether chat works at all:

1. The chat template wraps messages in <|im_start|>role ... <|im_end|> markers and ends with the generation prompt <|im_start|>assistant\n (Chapter 4).
2. Thinking mode. Qwen3 can “think” in a <think>...</think> block before answering. enable_thinking=False in the template turns it off for short, direct answers; True gives better reasoning at the cost of many more tokens.
3. Stop tokens. The assistant ends its turn with <|im_end|> (ID 151645); <|endoftext|> (151643) ends a document. Stop on both, or the model will continue with an imaginary next user turn.

The engine

class LLM:
    def __init__(self, model, tokenizer=None, manifest=None):
        self.model = model.eval()
        self.tokenizer = tokenizer
        self.device = next(model.parameters()).device
        self.manifest = manifest or {}

    @classmethod
    def from_pretrained(cls, directory, device=None, dtype="bf16", **load_options):
        device = device or default_device()
        model = load_model(directory, device, parse_dtype(dtype), **load_options)
        tokenizer = None
        try:   # The tokenizer is the only third-party piece, used at the text boundary.
            from transformers import AutoTokenizer
            tokenizer = AutoTokenizer.from_pretrained(directory, local_files_only=True)
        except Exception:
            pass
        config = read_config(directory)
        manifest = {"model_dir": str(directory), "model_type": config.get("model_type"),
                    "dtype": str(parse_dtype(dtype)), "device": str(device),
                    "torch": torch.__version__, "python": platform.python_version(),
                    "gpu": torch.cuda.get_device_name() if str(device).startswith("cuda") else None}
        return cls(model, tokenizer, manifest)

    def encode(self, text, chat=False, thinking=False):
        if self.tokenizer is None:
            raise RuntimeError("No tokenizer: pass token IDs directly")
        if chat:
            text = self.tokenizer.apply_chat_template([{"role": "user", "content": text}], tokenize=False,
                                                      add_generation_prompt=True, enable_thinking=thinking)
        return self.tokenizer(text, add_special_tokens=False).input_ids

    def default_stop_ids(self):
        if self.tokenizer is None or self.tokenizer.eos_token_id is None:
            return ()
        ids = {self.tokenizer.eos_token_id}
        for token in ("<|im_end|>", "<|endoftext|>"):
            index = self.tokenizer.convert_tokens_to_ids(token)
            if isinstance(index, int) and index >= 0 and index != self.tokenizer.unk_token_id:
                ids.add(index)
        return tuple(ids)

    @torch.inference_mode()
    def stream(self, prompt_ids, config=GenerationConfig()):
        """Yield token IDs as they are produced; afterwards self.last_result holds the summary.  (Your engine: Chapter 18)

        Measures time to first token (prefill + first sample) and the mean time per later token.
        """
        ids = torch.tensor([list(prompt_ids)], device=self.device)
        total = ids.shape[1] + config.max_new_tokens
        if total > self.model.context_limit:
            raise ValueError("Prompt plus max_new_tokens exceed the model's context")
        generator = torch.Generator(device=self.device).manual_seed(config.seed)
        policy = dict(temperature=config.temperature, top_k=config.top_k, top_p=config.top_p,
                      min_p=config.min_p, generator=generator)
        cache = self.model.new_cache(1, total)
        synchronize(self.device)
        start = time.perf_counter()
        logits = self.model(ids, cache)
        produced, finish, first_token_time = [], "length", None
        try:
            for step in range(config.max_new_tokens):
                token = sample(logits[:, -1], **policy)
                value = int(token)                    # host sync: we must know it to stream it
                if first_token_time is None:
                    first_token_time = time.perf_counter()
                produced.append(value)
                yield value
                if value in config.stop_ids:
                    finish = "stop"
                    break
                if step + 1 < config.max_new_tokens:
                    logits = self.model(token, cache)
        finally:
            end = time.perf_counter()
            decode_steps = max(len(produced) - 1, 0)
            self.last_result = GenerationResult(produced, finish, metrics={
                "prompt_tokens": ids.shape[1], "new_tokens": len(produced),
                "ttft_ms": round(1000 * ((first_token_time or end) - start), 3),
                "ms_per_output_token": round(1000 * (end - (first_token_time or end)) / decode_steps, 3) if decode_steps else None,
                "tokens_per_s": round(len(produced) / (end - start), 2) if produced else 0.0})

    def generate(self, prompt, config=None, chat=False):
        """prompt is text (needs a tokenizer) or a list of token IDs."""
        config = config or GenerationConfig(stop_ids=self.default_stop_ids())
        prompt_ids = self.encode(prompt, chat=chat) if isinstance(prompt, str) else list(prompt)
        for _ in self.stream(prompt_ids, config):
            pass
        result = self.last_result
        if self.tokenizer is not None:
            result.text = self.tokenizer.decode(result.token_ids, skip_special_tokens=True)
        return result

    def run_manifest(self, prompt_ids, config):
        """Everything needed to reproduce (and compare) a run."""
        return {**self.manifest, "prompt_token_ids": list(prompt_ids), "generation": asdict(config)}

stream is the heart: prefill the prompt, sample, then decode one token at a time, yielding each token as soon as it exists so a UI can display it. The try/finally records the result even if the consumer stops reading early (a user pressing “stop” is a normal event, not an error). It measures two numbers users feel:

• TTFT (time to first token): from the start of the request to the first sampled token. That’s prefill plus one sample. It grows with prompt length.
• Time per output token (TPOT): the average time between later tokens. That’s decode. It’s nearly constant, rising slowly as the KV cache grows.

run_manifest records everything needed to reproduce a run: model directory and type, dtype, device, library versions, prompt token IDs and generation settings. Save one next to every result you want to keep or compare. Later, “it got faster” or “it got worse” questions become answerable.

Chat with it

python run.py chat --model-dir models/Qwen3-0.6B --prompt "Why is the sky blue? Answer in two sentences."

The text streams token by token, followed by the measurements, like this (exact text and timings will vary):

The sky appears blue because molecules in the atmosphere scatter shorter (blue) wavelengths of sunlight
much more strongly than longer (red) ones. ...
{"prompt_tokens": 23, "new_tokens": 41, "ttft_ms": ..., "ms_per_output_token": ..., "tokens_per_s": ..., "finish": "stop"}

Tip

If the model rambles, repeats the question or starts a fake “user:” turn, check the boundary before suspecting the model: was the generation prompt added, and are both stop IDs in stop_ids? llm.encode(text, chat=True) and tokenizer.convert_ids_to_tokens(ids) show exactly what the model sees.

A ladder of evidence

“It produces English” is weak evidence. A model with a subtly wrong RoPE still writes fluent sentences. Climb this ladder instead, and stop at the first rung that fails:

1. Unit operations against small hand-checked cases: Chapter 17’s tests.
2. Internal consistency: full, chunked and token-by-token forwards agree (Chapter 16’s tests). This catches position and cache bugs.
3. Independent parity on the real checkpoint: load the same snapshot into Hugging Face Transformers and into your engine, in FP32, and compare logits on fixed token IDs. Expect a maximum absolute difference around 1e-4 to 1e-3.
4. Layer-by-layer localization when rung 3 fails: register forward hooks on both models’ decoder layers, and compare the residual stream after each layer. The first layer that disagrees contains the bug.
5. Behavior: identical greedy text over several prompts, in FP32.

import torch
from transformers import AutoModelForCausalLM
from engine.loaders import load_qwen3

path = "models/Qwen3-0.6B"
ref = AutoModelForCausalLM.from_pretrained(path, dtype=torch.float32, attn_implementation="eager").eval()
mine = load_qwen3(path, device="cpu", dtype=torch.float32)
ids = torch.randint(0, 151936, (1, 33))
with torch.no_grad():
    diff = (mine(ids) - ref(ids).logits).abs()
print("max abs logit difference:", diff.max().item())

For BF16, compare like with like: your BF16 against Transformers’ BF16, both judged against FP32 (the three-comparison rule from Chapter 13). Greedy tokens can legitimately diverge after a few dozen steps in BF16, because a near-tie between two tokens flips on rounding. Compare logits, not long generations.

Memory budget

For Qwen3-0.6B in BF16 on a single request:

The logits row is a surprise to most people. Prefill only needs the last position’s logits, so computing all 4,096 × 151,936 of them wastes memory and a large matmul. Production engines slice the hidden states to the last position before the head. That’s a one-line stretch exercise below.

Measure against the ceiling

Use your Chapter 10 tools. Predict, then measure:

• Decode ceiling: decode_ceiling(1.19e9, bandwidth), about 230 tokens/s on DGX Spark and about 840 on an RTX 4090.
• Achieved: ms_per_output_token from the engine.

Engine v1 typically reaches a small fraction of the ceiling on a GPU. A 0.6B model’s decode step is hundreds of small kernels, each waiting on Python, plus one int(token) synchronization per step, so the GPU idles between kernels. A profile (torch.profiler) shows the gaps clearly. Chapter 19 removes them.

The Rust engine does the same, on the CPU

The Rust track loads the same snapshot and generates on the CPU. Its weights stay in BF16 and are widened inside a multi-threaded mat-vec, which is a direct application of the memory-bound decode model:

#![allow(unused)]
fn main() {
    /// Process one token at position cache.len and return the next-token logits.
    pub fn forward_token(&self, token: u32, cache: &mut KvCache) -> Vec<f32> {
        let c = &self.cfg;
        let pos = cache.len;
        let d = c.head_dim;
        let mut x = self.embed.row(token as usize);
        for (li, layer) in self.layers.iter().enumerate() {
            let h = rmsnorm(&x, &layer.input_norm, c.eps);
            let (mut q, mut k, v) = (layer.q.matvec(&h), layer.k.matvec(&h), layer.v.matvec(&h));
            for head in q.chunks_mut(d) {                    // per-head RMSNorm, then rotate
                head.copy_from_slice(&rmsnorm(head, &layer.q_norm, c.eps));
                rope(head, pos, c.theta);
            }
            for head in k.chunks_mut(d) {
                head.copy_from_slice(&rmsnorm(head, &layer.k_norm, c.eps));
                rope(head, pos, c.theta);
            }
            let (keys, values) = cache.append(li, &k, &v);
            let attended = causal_attention(&q, keys, values, 1, pos + 1, c.heads, c.kv_heads, d);
            for (xi, oi) in x.iter_mut().zip(layer.o.matvec(&attended)) {
                *xi += oi;
            }
            let h = rmsnorm(&x, &layer.post_norm, c.eps);
            let gated: Vec<f32> = layer.gate.matvec(&h).iter().zip(layer.up.matvec(&h)).map(|(g, u)| silu(*g) * u).collect();
            for (xi, mi) in x.iter_mut().zip(layer.down.matvec(&gated)) {
                *xi += mi;
            }
        }
        cache.len += 1;
        let h = rmsnorm(&x, &self.norm, c.eps);
        self.head.as_ref().unwrap_or(&self.embed).matvec(&h)
    }
}

#![allow(unused)]
fn main() {
/// Prefill the prompt token by token (simple, not fast), then decode `new_tokens` tokens.
pub fn generate(model: &Qwen3, prompt: &[u32], new_tokens: usize, policy: &crate::sampling::Policy,
                rng: &mut crate::rng::Rng, stop: &[u32]) -> Vec<u32> {
    let mut cache = model.new_cache(prompt.len() + new_tokens);
    let mut logits = vec![];
    for &t in prompt {
        logits = model.forward_token(t, &mut cache);
    }
    let mut out = vec![];
    for _ in 0..new_tokens {
        let next = crate::sampling::sample(&logits, policy, rng) as u32;
        out.push(next);
        if stop.contains(&next) || out.len() == new_tokens {
            break;
        }
        logits = model.forward_token(next, &mut cache);
    }
    out
}
}

cd rust
IDS=$(python tokenize_ids.py encode --model-dir ../models/Qwen3-0.6B "Why is the sky blue?")
OUT=$(cargo run --release -- generate --model-dir ../models/Qwen3-0.6B --ids "$IDS" --new-tokens 48)
python tokenize_ids.py decode --model-dir ../models/Qwen3-0.6B "$OUT"

Its cargo test reproduces the Python reference’s logits at every position to 1e-4, for both F32 and BF16 checkpoints. Compare its CPU tokens/s with your machine’s memory bandwidth divided by 1.19 GB.

Build it

pytest tests/test_ch18_engine.py
python run.py chat --impl engine --model-dir models/Qwen3-0.6B

Stretch exercises

1. ★ Compute logits only for the last position during prefill (return_hidden=True, slice, then lm_head). Measure the memory saved on a 4,096-token prompt. Where: the prefill model call in LLM.stream in engine/engine.py.
2. ★★ Implement rung 4: a function that registers hooks on every decoder layer of both models and prints the first layer whose output differs by more than a tolerance. Where: experiments/ch18.py (create it), using engine.qwen3 and the corresponding Transformers model.
3. ★★ Add a --thinking flag to run.py chat and compare the answer quality and token count of thinking and non-thinking modes on five reasoning questions. Where: the argument parser and cmd_chat in run.py; pass it as thinking=... to LLM.encode (which forwards it to the chat template).
4. ★★★ Load Qwen’s tokenizer.json with your own BPE implementation: parse the vocabulary and merges, implement the byte-to-unicode mapping and the pre-tokenizer regex (use the third-party regex package), and check that you reproduce AutoTokenizer’s IDs on 1,000 lines of text. Where: add a checkpoint-tokenizer loader in engine/tokenizer.py; compare it in experiments/ch18.py (create it).

Check your understanding

1. Why should you compare intermediate layer outputs before generated text?
2. Why can a model that fits in memory still fail to allocate for a long prompt?
3. Why are “tokens emitted” and “tokens processed into the cache” different counts?
4. Why does to_empty require you to re-tie the embedding and head?
5. What two stop tokens does a Qwen3 chat need, and what happens with only one?

Going deeper

• BALLM Chapter 5 §§5.4-5.5 for the GPT-2 version of loading pretrained weights; Raschka’s ch05/11_qwen3 notebook loads real Qwen3 weights into a from-scratch model.
• PMPP §20.4 and §20.6 for cache sizing and decode intensity.
• The Qwen3 model card (chat template, thinking mode, recommended sampling settings: temperature 0.7, top-p 0.8 for non-thinking mode).
• vLLM’s LLM class and SGLang’s Engine: the production versions of this chapter’s interface.

19. Fast decode: syncs, CUDA graphs and compilation

Where the time goes

Chapter 18 ended with a puzzle. The decode ceiling for Qwen3-0.6B in BF16 on an RTX 4090 is about 840 tokens/s (1.19 GB of weights read per token at 1,008 GB/s), but engine v1 runs far below it. Let’s count what one decode step asks the CPU to do.

Each of the 28 layers runs about 20 PyTorch operations: two RMSNorms (several kernels each without fusion), three projections, two QK-norms, RoPE on Q and K, the cache write, attention, the output projection, a residual add, the gate and up projections, SiLU, a multiply, the down projection, another residual add. That’s roughly 600 kernel launches per token. Each one costs the CPU several microseconds of Python and PyTorch dispatch, plus a few microseconds for the CUDA driver to launch it.

The GPU finishes each small kernel faster than the CPU can issue the next one, so it idles between them. A decode step of a small model is launch-bound: its speed is set by CPU overhead, not by memory or arithmetic. Larger models move toward the bandwidth ceiling, because each kernel does more work while the number of launches stays similar, but even an 8B model loses a noticeable fraction to overhead.

There are three fixes, in order of how much they help:

1. Remove host synchronizations, so the CPU can run ahead and queue work while the GPU executes.
2. Replay the whole step as a CUDA graph, so ~600 launches become one.
3. Fuse kernels, so there are fewer of them and less memory traffic between them (Chapter 14).

Host synchronizations

CUDA kernel launches are asynchronous. When Python calls torch.matmul on CUDA tensors, the CPU enqueues the kernel and returns immediately, typically long before the kernel runs. While the queue is non-empty, the GPU never waits for the CPU: launch overhead overlaps with GPU execution.

Some operations break this. Anything that needs a value from the GPU on the CPU must wait for every queued kernel to finish:

• int(token), token.item(), tensor.tolist(), print(tensor);
• if (tensor > 0).any(): (the bool conversion is a sync);
• torch.nonzero, boolean-mask indexing like x[mask], torch.unique (output shape depends on data);
• copying a CPU tensor to the GPU from non-pinned memory, and torch.cuda.synchronize() itself.

Engine v1 calls int(token) every step, to check stop tokens and to yield the token. After that call, the GPU queue is empty, and the next step starts from scratch: the GPU waits for the first launch, the second, and so on. The CPU never gets ahead.

The fix is to keep everything the loop needs on the device:

• The next input token, its position and the output buffer are device tensors that the step updates in place.
• Sampling happens on the device and produces a device tensor.
• Stop tokens are checked every check_every steps, with one sync per check. A few tokens generated after a stop are trimmed afterwards. That’s wasted work of at most check_every - 1 steps, in exchange for removing almost all syncs.

Tip

torch.cuda.set_sync_debug_mode("warn") makes PyTorch print a warning at every synchronizing operation. Run one decode step with it on to find the syncs you didn’t know about.

Sampling without the CPU

Chapter 8’s sampler used torch.multinomial, which is fine, but top-p and min-p filters involve sorting and data-dependent cutoffs. For the graph-friendly path, the decoder uses the Gumbel-max trick, which draws an exact sample from $\operatorname{softmax}(z)$ with only elementwise operations and an argmax:

$$ \text{if } g_i = -\log(-\log u_i),\ u_i \sim \text{Uniform}(0,1) \text{ independently, then } \arg\max_i (z_i + g_i) \sim \operatorname{softmax}(z). $$

Why it works, in one line: $z_i + g_i$ is a Gumbel random variable with location $z_i$, and the probability that the $i$-th of several independent Gumbels is the largest is $e^{z_i}/\sum_j e^{z_j}$. Dividing $z$ by the temperature first gives temperature sampling.

def sample_on_device(logits, temperature):
    """Greedy when temperature == 0; otherwise the Gumbel-max trick, which draws exactly from
    softmax(logits / temperature) using only elementwise ops and argmax (no host sync).  (Your engine: Chapter 19)"""
    if temperature == 0:
        return logits.argmax(-1, keepdim=True)
    uniform = torch.rand_like(logits, dtype=torch.float32).clamp_(1e-10, 1.0)
    gumbel = -torch.log(-torch.log(uniform))
    return (logits.float() / temperature + gumbel).argmax(-1, keepdim=True)

The milestone test draws 20,000 samples and checks their frequencies against the softmax. Top-k and top-p filters can be added before the argmax (set filtered logits to $-\infty$); top-k with a fixed k keeps shapes static, and so does top-p implemented with a sort and a mask rather than slicing.

Static shapes and a static cache

A CUDA graph records exact kernels with exact pointer arguments. Replaying it re-runs them on whatever data is at those addresses now. So everything a step reads and writes must live at fixed addresses with fixed shapes:

• The input token and position: one-element tensors updated in place with copy_ and add_.
• The KV cache: Chapter 16’s KVCache returns a view of the valid prefix, whose shape grows every step. That’s useless for a graph.

StaticKVCache solves the cache problem with the attention rule you’ve used since Chapter 5. It writes keys at their absolute positions in a preallocated buffer and always returns the full capacity, together with the position of each slot, key_positions = arange(capacity). Slots that haven’t been written yet sit at positions greater than every live query, so key_pos <= query_pos masks them, with no extra bookkeeping:

class StaticKVCache:
    """Fixed-shape cache indexed by absolute position. Each row is a slot that a request can own.

    Writes go to explicit positions, so rows may hold different lengths. Reads return the full
    capacity; unwritten or stale slots sit at positions greater than every live query, so the
    causal rule (key_pos <= query_pos) hides them without any extra mask. Shapes never change,
    which is what CUDA graphs and torch.compile need.  (Your engine: Chapter 19)
    """

    def __init__(self, layers, slots, kv_heads, capacity, head_dim, device="cpu", dtype=torch.float32):
        self.capacity = capacity
        shape = (slots, kv_heads, capacity, head_dim)
        # zeros, not empty: masked weights are exactly 0, and 0 * garbage could still be NaN.
        self.keys = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(layers)]
        self.values = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(layers)]
        self.key_positions = torch.arange(capacity, device=device)

    def update(self, layer, k, v, positions, rows=None):
        """(Your engine: Chapter 19)"""
        if positions.ndim == 1:
            positions = positions.expand(k.shape[0], -1)
        keys, values = self.keys[layer], self.values[layer]
        if rows is None:
            rows = torch.arange(k.shape[0], device=k.device)
        # One scatter per tensor: row r, head h, position positions[r, t] <- k[r, h, t].
        b, h, t, d = k.shape
        row_index = rows[:, None].expand(b, t)
        keys[row_index, :, positions] = k.transpose(1, 2)
        values[row_index, :, positions] = v.transpose(1, 2)
        return keys[rows], values[rows], self.key_positions

    @property
    def bytes(self):
        return sum(t.numel() * t.element_size() for t in self.keys + self.values)

Two details matter:

• Zeros, not empty. A masked attention weight is exactly 0, but $0 \times \text{NaN} = \text{NaN}$. Uninitialized memory can contain NaN bit patterns, so the buffers are zero-filled once.
• Rows are slots. Each row of the cache can hold a different request at a different length. Writes go to (row, position) pairs with one scatter. That’s exactly what continuous batching needs in Chapter 24, so the same class serves both.

The cost of a static cache: attention always reads capacity keys, even when only 50 are valid. For a short conversation in a large buffer, that wastes bandwidth. Production engines pass the valid length as a device tensor to a custom decode kernel that stops early (Chapter 25’s paged decode kernel does this). Choose a capacity near the longest request you expect, or capture one graph per capacity bucket.

The fast decoder

class FastDecoder:
    """Batch-1 decoder over static buffers.  (Your engine: Chapter 19)"""

    def __init__(self, model, capacity, mode="eager", temperature=0.0):
        if mode not in ("eager", "graph", "compile"):
            raise ValueError("mode must be eager, graph or compile")
        if mode == "graph" and not torch.cuda.is_available():
            raise RuntimeError("CUDA graphs need a CUDA device")
        self.model, self.mode, self.temperature = model.eval(), mode, temperature
        p = next(model.parameters())
        layers, kv_heads, head_dim = model.cache_spec()
        self.cache = StaticKVCache(layers, 1, kv_heads, capacity, head_dim, p.device, p.dtype)
        self.capacity = capacity
        self.token = torch.zeros(1, 1, dtype=torch.long, device=p.device)      # next input token
        self.position = torch.zeros(1, 1, dtype=torch.long, device=p.device)   # its absolute position
        self.outputs = torch.zeros(1, capacity, dtype=torch.long, device=p.device)
        self.step_fn = self._step
        self.graph = None
        if mode == "compile":
            self.step_fn = torch.compile(self._step, mode="reduce-overhead", fullgraph=False)

    def _step(self):
        """One decode step that touches only static tensors: run, sample, record, advance.  (Your engine: Chapter 19)"""
        logits = self.model(self.token, self.cache, positions=self.position)[:, -1]
        nxt = sample_on_device(logits, self.temperature)
        self.outputs.index_copy_(1, self.position[0], nxt)    # output i is stored at slot of its input
        self.token.copy_(nxt)
        self.position.add_(1)
        return nxt

    def _capture(self):
        """Record one step as a CUDA graph. Warm up on a side stream first, as PyTorch requires."""
        stream = torch.cuda.Stream()
        stream.wait_stream(torch.cuda.current_stream())
        saved = (self.token.clone(), self.position.clone(), self.outputs.clone())
        with torch.cuda.stream(stream):
            for _ in range(2):
                self._step()
        torch.cuda.current_stream().wait_stream(stream)
        self.graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self.graph):
            self._step()
        # Warmup and capture advanced the state; put it back. Cache slots they wrote lie at
        # positions the real run will overwrite before reading.
        self.token.copy_(saved[0]); self.position.copy_(saved[1]); self.outputs.copy_(saved[2])

    @torch.inference_mode()
    def generate(self, prompt_ids, max_new_tokens, stop_ids=(), check_every=16):
        """Return the generated token IDs (a Python list) for one prompt.  (Your engine: Chapter 19)"""
        ids = torch.as_tensor(prompt_ids, device=self.token.device).view(1, -1)
        prompt = ids.shape[1]
        if prompt + max_new_tokens > self.capacity:
            raise ValueError("Prompt plus new tokens exceed the decoder's capacity")
        positions = torch.arange(prompt, device=ids.device)[None]
        logits = self.model(ids, self.cache, positions=positions)[:, -1]       # prefill (eager)
        first = sample_on_device(logits, self.temperature)
        self.token.copy_(first)
        self.position.fill_(prompt)
        self.outputs.zero_()
        if self.mode == "graph" and self.graph is None:
            self._capture()
        stop = torch.tensor(list(stop_ids) or [-1], device=ids.device)
        produced = 1
        while produced < max_new_tokens:
            steps = min(check_every, max_new_tokens - produced)
            for _ in range(steps):
                if self.graph is not None:
                    self.graph.replay()
                else:
                    self.step_fn()
            produced += steps
            window = torch.cat((first, self.outputs[:, prompt:prompt + produced - 1]), dim=1)
            if stop_ids and bool(torch.isin(window, stop).any()):     # one sync per check
                break
        tokens = torch.cat((first, self.outputs[:, prompt:prompt + produced - 1]), dim=1)[0].tolist()
        for i, token in enumerate(tokens):                            # trim anything after a stop token
            if token in stop_ids:
                return tokens[:i + 1]
        return tokens[:max_new_tokens]

Read _step first. It’s the whole decode step, and it touches only static tensors: run the model on self.token at self.position, sample on the device, store the token in self.outputs, copy it into self.token, and advance the position. No Python value depends on the GPU’s results, so nothing waits.

generate runs prefill eagerly (the prompt length varies, so it can’t be graphed without bucketing), samples the first token, then runs _step in groups of check_every, checking for stop tokens once per group.

CUDA graphs

A CUDA graph is a recording of a stream of GPU work. You capture it once and replay it as a single launch:

graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):     # kernels are recorded, not executed
    step()
graph.replay()                    # re-run every recorded kernel, one launch from the CPU

During capture, PyTorch records the kernels and their arguments; the memory they allocate comes from a private pool that stays reserved for the graph’s lifetime, so replays reuse the same addresses. The rules follow from “replay re-runs the recorded kernels on the same addresses”:

1. No host syncs during capture. A capture can’t wait for results that don’t exist yet; PyTorch raises an error if you try.
2. No data-dependent Python control flow. An if on a tensor’s value is evaluated once, at capture time, and its branch is baked in.
3. Inputs must be updated in place. Assigning self.token = new_tensor gives a new address that the graph never sees. Use self.token.copy_(new).
4. Warm up first, on a side stream. The first call to many operations triggers lazy initialization (cuBLAS handles, autotuning, Triton compilation), which must not happen during capture. _capture runs the step twice on a side stream before recording.
5. Random numbers work, because PyTorch registers the generator’s state with the graph and advances its offset each replay. Each replay draws fresh Gumbel noise.

The warmup and capture each executed a step, which advanced the token, the position and the outputs. _capture restores them afterwards. They also wrote K and V into one or two cache slots beyond the prompt; those slots sit at positions the real run will overwrite before any query can see them, so they’re harmless. Thinking through “what state did capture modify?” is a habit worth forming: it’s the most common source of graph bugs.

torch.compile

torch.compile(step, mode="reduce-overhead") does two jobs at once: it fuses chains of small operations into generated Triton kernels (fewer launches, less traffic), and it captures CUDA graphs automatically (“reduce-overhead” means “use graphs”). mode="max-autotune" additionally benchmarks matmul configurations. The first calls are slow (compilation can take a minute); measure only after warmup.

Compilation needs the same discipline as graphs: static shapes (or marked dynamic dimensions), no syncs inside the compiled region, and in-place state updates. Code written for mode="graph" is already compile-friendly, which is why the decoder supports both.

Measure it

python run.py fast --new-tokens 128 --device cuda

The command measures the plain generate() loop from Chapter 8, then the fast decoder in eager mode, then (on CUDA) in graph mode, and checks that greedy tokens agree. On a laptop CPU with the tiny test model, it printed:

{"mode": "generate()", "tokens_per_s": 384.0}
{"mode": "eager", "tokens_per_s": 359.6, "same_tokens": true}

No improvement, and that’s the right result: on a CPU, operations run synchronously, so there’s no queue to keep full and nothing for sync removal to win. On a GPU the picture changes completely. The shape of what you should expect for Qwen3-0.6B on a consumer GPU (illustrative, not a measurement from this book’s validation; run it and record yours):

The remaining gap comes from kernels that don’t reach peak bandwidth at batch 1 (a 1,024 × 3,072 mat-vec is too small to saturate the memory system), from attention reading the whole static capacity, and from unfused elementwise operations.

Read a profile

Guessing where time goes is unreliable. Profile one decode step:

from torch.profiler import profile, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    decoder.generate(prompt, 32)
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))
prof.export_chrome_trace("decode.json")      # open in https://ui.perfetto.dev

In the trace, look at the GPU stream row:

• Gaps between kernels mean launch-bound. Fix with graphs.
• Back-to-back kernels, mat-vecs taking most of the time means memory-bound, which is the goal for decode. Check achieved bandwidth: bytes read ÷ kernel time.
• A cudaStreamSynchronize or aten::item on the CPU row in the middle of a step is a sync you missed.

NVIDIA Nsight Systems (nsys profile python run.py fast ...) shows the same picture for the whole process, including Python, and is the tool GPU Mode lectures use most.

Build it

pytest tests/test_ch19_fast.py                 # graph and compile tests run only on CUDA
python run.py fast --impl engine --device cuda

Stretch exercises

1. ★ Run one step of engine v1 under torch.cuda.set_sync_debug_mode("warn") and list every sync it reports. Then do the same for the fast decoder. Where: experiments/ch19.py (create it), comparing engine.engine.LLM with engine.fast.FastDecoder.
2. ★★ Fuse Q, K and V into one projection (concatenate the three weight matrices at load time) and gate and up into another. Count the kernels per step before and after with the profiler. Where: Qwen3Attention and SwiGLU in engine/qwen3.py, with weight conversion in engine/loaders.py.
3. ★★ Swap Chapter 14’s fused add+RMSNorm Triton kernel into Qwen3Layer, keeping the PyTorch path as a fallback. Verify logits, then measure graph-mode tokens/s. Where: Qwen3Layer.forward in engine/qwen3.py, calling engine.kernels.triton_basics.rmsnorm.
4. ★★ Graph the prefill too: capture one graph per prompt-length bucket (64, 128, 256, …) and pad prompts up to the next bucket. What must the padding positions be so that they don’t affect the real tokens? Where: FastDecoder._capture / generate in engine/fast.py, with padding masks passed through engine/qwen3.py.
5. ★★★ Add top-k and top-p filtering to sample_on_device without any operation whose output shape depends on data, and check the sample distribution against Chapter 8’s sampler. Where: sample_on_device in engine/fast.py.

Check your understanding

1. Why does removing int(token) from the loop make the GPU faster, even though the GPU does the same work?
2. Why can’t a CUDA graph contain if logits.argmax() == stop_id:?
3. How does StaticKVCache hide slots that haven’t been written yet, without a separate mask?
4. Why does the fast decoder show no speedup on a CPU?
5. What happens if you write self.token = nxt instead of self.token.copy_(nxt) inside a captured step?

Going deeper

• GPU Mode L1 (Profiling and integrating CUDA kernels in PyTorch) and L16 (Hands-on profiling) for the profiler and Nsight; L6 (Optimizing PyTorch optimizers) for horizontal fusion and why many tiny launches are slow; L35 (SGLang performance optimization) for CUDA graphs and overhead removal in a production engine.
• The PyTorch blog Accelerating Generative AI with PyTorch II: GPT, Fast and the gpt-fast repository: static KV cache, torch.compile(mode="reduce-overhead"), and int8/int4 weight-only quantization in under 1,000 lines. This chapter follows its approach.
• PyTorch documentation: CUDA Graphs (capture rules, memory pools, make_graphed_callables) and torch.compile troubleshooting (graph breaks, recompilations).
• Gumbel (1954) and Maddison, Tarlow and Minka, A Sampling* (2014) for the Gumbel-max trick.
• vLLM’s and SGLang’s CUDA-graph runners, which capture one graph per batch size bucket.

20. Quantization

Fewer bytes, faster tokens

Decode reads every weight once per token, so its ceiling is bandwidth ÷ bytes per token (Chapter 10). Halve the bytes and you double the ceiling. For an 8B-parameter model:

Quantization also decides what fits. A 32B model in BF16 needs 64 GB, more than an RTX 4090’s 24 GB; in INT4 it needs about 17 GB. For the sparse models of Part VII, with hundreds of billions of parameters, low-bit weights are the difference between running and not running.

The price is error: every weight becomes the nearest value on a coarse grid. The whole subject is about spending a fixed bit budget where it matters most.

Symmetric quantization, from scratch

Take a group of weights $w_1, \dots, w_n$ and a bit width $b$. A signed $b$-bit integer can hold $-2^{b-1} \dots 2^{b-1}-1$; we use the symmetric range $[-Q, Q]$ with $Q = 2^{b-1}-1$ (127 for INT8, 7 for INT4) so that zero is exact and the grid is symmetric. Map the largest magnitude $a = \max_i |w_i|$ to $Q$:

$$ s = \frac{a}{Q}, \qquad q_i = \operatorname{clamp}!\Big(\operatorname{round}\big(w_i / s\big), -Q, Q\Big), \qquad \hat w_i = s, q_i . $$

The integers $q_i$ (the codes) are stored, along with one floating-point scale $s$ per group. Without clamping (the range covers every weight), rounding moves each value by at most half a step:

$$ |\hat w_i - w_i| \le \tfrac{s}{2} = \frac{a}{2Q}. $$

Worked example. $w = [-2, -1, 0, 1, 2]$ at 4 bits: $Q = 7$, $s = 2/7 \approx 0.2857$. Dividing gives $[-7, -3.5, 0, 3.5, 7]$. PyTorch rounds halves to even, so the codes are $[-7, -4, 0, 4, 7]$ and $\hat w = [-2, -1.143, 0, 1.143, 2]$. The error on $\pm 1$ is $0.143 = s/2$, exactly at the bound. (Rounding conventions differ between libraries, which is one reason to test against the specific runtime you’ll use.)

Granularity: who shares a scale

The bound says the error is proportional to the largest magnitude in the group. One large weight coarsens the grid for everything that shares its scale. So the choice of groups matters as much as the bit width:

• Per tensor: one scale for the whole matrix. A single outlier anywhere ruins it. Fine for INT8 activations with care, poor for low-bit weights.
• Per row (per output channel): one scale per row of the [out, in] weight. Standard for INT8 weights.
• Per group: one scale per g consecutive inputs within a row, typically $g = 32$-$128$. Standard for INT4.

def quantize_int8_rows(weight):
    """One scale per output row: scale = max|w| / 127, q = round(w / scale).  (Your engine: Chapter 20)"""
    maximum = weight.float().abs().amax(dim=-1, keepdim=True)
    scale = torch.where(maximum > 0, maximum / 127.0, torch.ones_like(maximum))
    return (weight.float() / scale).round().clamp(-127, 127).to(torch.int8), scale


def dequantize_int8_rows(q, scale):
    return q.float() * scale

def quantize_groupwise(weight, bits=4, group_size=128, clip=1.0):
    """Symmetric signed codes with one scale per (row, group of group_size inputs).  (Your engine: Chapter 20)

    weight [O, I] -> codes int8 [O, I] in [-(2^(b-1)-1), 2^(b-1)-1], scales [O, ceil(I/g)].
    clip < 1 shrinks the range: large outliers saturate, everything else gets finer steps.
    """
    if weight.ndim != 2 or bits not in (2, 3, 4, 8) or group_size < 1 or not 0 < clip <= 1:
        raise ValueError("Expected a matrix, 2/3/4/8 bits, positive group size and 0 < clip <= 1")
    out_features, in_features = weight.shape
    groups = -(-in_features // group_size)
    padded = F.pad(weight.float(), (0, groups * group_size - in_features))
    blocks = padded.view(out_features, groups, group_size)
    limit = 2 ** (bits - 1) - 1
    maximum = blocks.abs().amax(-1, keepdim=True) * clip
    scales = torch.where(maximum > 0, maximum / limit, torch.ones_like(maximum))
    codes = (blocks / scales).round().clamp(-limit, limit).to(torch.int8)
    return codes.view(out_features, -1)[:, :in_features], scales.squeeze(-1)


def dequantize_groupwise(codes, scales, group_size):
    out_features, in_features = codes.shape
    expanded = scales.repeat_interleave(group_size, dim=1)[:, :in_features]
    return codes.float() * expanded

Two edge cases the code handles: a group of all zeros gets scale 1 (any scale works, and 0/0 doesn’t); the last group of a row may be short when in_features isn’t a multiple of g, so the row is padded for the computation and the padding is dropped.

The metadata the name leaves out

“4-bit” counts the codes only. With groups of $g$ inputs and a 16-bit scale per group, the true cost is $4 + 16/g$ bits per weight: 4.125 for $g = 128$, 4.5 for $g = 32$. Asymmetric schemes add a zero point per group. Layers usually left in high precision (embeddings, the LM head, MoE routers, norms) add more. For Qwen3-0.6B, whose tied embedding is 26% of the parameters, “4-bit” weights come to well over 4 bits on average. Always report measured bytes.

Packing

PyTorch has no 4-bit dtype, so an int8 tensor holding values in $[-7, 7]$ still costs 8 bits each. Pack two codes per byte: the first in the low nibble, the second in the high nibble, each in 4-bit two’s complement (values 8-15 of a nibble mean -8 to -1):

def pack_int4(codes):
    """Two signed 4-bit values per byte, two's complement, first value in the low nibble."""
    values = codes.to(torch.int16)
    if bool(((values < -8) | (values > 7)).any()):
        raise ValueError("INT4 values must lie in [-8, 7]")
    flat = values.flatten()
    if flat.numel() % 2:
        flat = torch.cat((flat, flat.new_zeros(1)))
    nibbles = flat & 0xF
    return (nibbles[0::2] | (nibbles[1::2] << 4)).to(torch.uint8)


def unpack_int4(packed, shape):
    """Inverse of pack_int4: nibbles 8..15 decode to -8..-1."""
    nibbles = torch.stack((packed & 0xF, packed >> 4), dim=-1).flatten().to(torch.int16)
    signed = torch.where(nibbles >= 8, nibbles - 16, nibbles)
    count = 1
    for size in shape:
        count *= size
    return signed[:count].to(torch.int8).reshape(shape)

inline std::pair<std::vector<int8_t>, std::vector<float>> quantize_row(const std::vector<float>& w, int bits, size_t group) {
    float limit = float((1 << (bits - 1)) - 1);
    std::vector<int8_t> codes;
    std::vector<float> scales;
    for (size_t g = 0; g < w.size(); g += group) {
        size_t end = std::min(w.size(), g + group);
        float m = 0;
        for (size_t i = g; i < end; ++i) m = std::max(m, std::fabs(w[i]));
        float scale = m > 0 ? m / limit : 1.f;
        scales.push_back(scale);
        for (size_t i = g; i < end; ++i) codes.push_back(int8_t(std::clamp(std::round(w[i] / scale), -limit, limit)));
    }
    return {codes, scales};
}
inline std::vector<uint8_t> pack_int4(const std::vector<int8_t>& c) {  // first value in the low nibble
    std::vector<uint8_t> out((c.size() + 1) / 2, 0);
    for (size_t i = 0; i < c.size(); ++i) out[i / 2] |= uint8_t((c[i] & 0x0f) << (4 * (i % 2)));
    return out;
}
inline int8_t unpack_nibble(uint8_t byte, int high) {
    int n = (byte >> (4 * high)) & 0x0f;
    return int8_t(n >= 8 ? n - 16 : n);                 // two's complement
}

#![allow(unused)]
fn main() {
/// Quantize one row in groups of `group` values: scale = max|w| / (2^(bits-1) - 1).
pub fn quantize_row(w: &[f32], bits: u32, group: usize) -> (Vec<i8>, Vec<f32>) {
    let limit = ((1 << (bits - 1)) - 1) as f32;
    let mut codes = Vec::with_capacity(w.len());
    let mut scales = vec![];
    for chunk in w.chunks(group) {
        let max = chunk.iter().fold(0.0f32, |m, v| m.max(v.abs()));
        let scale = if max > 0.0 { max / limit } else { 1.0 };
        scales.push(scale);
        codes.extend(chunk.iter().map(|v| (v / scale).round().clamp(-limit, limit) as i8));
    }
    (codes, scales)
}

/// Two signed 4-bit codes per byte, first code in the low nibble (two's complement).
pub fn pack_int4(codes: &[i8]) -> Vec<u8> {
    codes.chunks(2).map(|p| {
        let lo = (p[0] as u8) & 0x0f;
        let hi = (*p.get(1).unwrap_or(&0) as u8) & 0x0f;
        lo | (hi << 4)
    }).collect()
}

pub fn unpack_int4(packed: &[u8], count: usize) -> Vec<i8> {
    let nibble = |n: u8| if n >= 8 { n as i8 - 16 } else { n as i8 };
    packed.iter().flat_map(|b| [nibble(b & 0x0f), nibble(b >> 4)]).take(count).collect()
}
}

Packed layouts are a contract between the quantizer and the kernel, and every library chooses differently: nibble order, interleaving for vectorized loads, and transposition for tensor-core fragment layouts. A GPTQ checkpoint packs eight 4-bit values into an int32; Marlin kernels reorder them again for fast loading. A checkpoint is only usable by a runtime that implements its exact format.

Calibration: optimize what you care about

Clipping (shrinking $a$ by a factor clip < 1) saturates the largest weights but gives everything else a finer grid. Whether that helps depends on how the weights are used. A large weight attached to an input feature that’s almost always near zero barely affects the output, so it’s worth clipping; a large weight on an active feature isn’t.

So choose quantization parameters by output error on representative inputs (calibration data), not by weight error:

def choose_clip(weight, calibration_inputs, bits=4, group_size=128, candidates=(1.0, 0.95, 0.9, 0.85, 0.8, 0.7)):
    """Pick the clip that minimizes OUTPUT error on calibration activations, not weight error."""
    reference = calibration_inputs.float() @ weight.float().T
    best = None
    for clip in candidates:
        codes, scales = quantize_groupwise(weight, bits, group_size, clip)
        approx = calibration_inputs.float() @ dequantize_groupwise(codes, scales, group_size).T
        error = ((approx - reference).norm() / reference.norm().clamp_min(1e-12)).item()
        if best is None or error < best[0]:
            best = (error, clip)
    return best[1], best[0]

The lab plants exactly that situation: a 16 × 128 weight whose input feature 0 has weights 12× larger than the rest, while that feature is nearly silent in the activations. It picks the clip on 256 calibration rows and reports error on 256 different held-out rows:

python model_workflows.py quantize

{"bits": 8, "group": 128, "clip": 0.95, "weight_error": 0.03627, "held_out_output_error": 0.02119, ...}
{"bits": 8, "group": 32,  "clip": 1.0,  "weight_error": 0.00874, "held_out_output_error": 0.01161, ...}
{"bits": 4, "group": 128, "clip": 0.7,  "weight_error": 0.28655, "held_out_output_error": 0.28023, ...}
{"bits": 4, "group": 64,  "clip": 0.7,  "weight_error": 0.25714, "held_out_output_error": 0.22791, ...}
{"bits": 4, "group": 32,  "clip": 0.85, "weight_error": 0.17493, "held_out_output_error": 0.19074, ...}

Three lessons: 4-bit chose aggressive clipping because the outlier is harmless to clip; smaller groups help most at low bit widths; and at group 32 the output error is larger than the weight error, while at group 64 it’s smaller. Weight error is not the objective. Keep calibration and evaluation data separate, or you’ll fool yourself.

A quantized linear layer

class QuantLinear(nn.Module):
    """Weight-only quantized linear layer: integer weights in memory, higher-precision math.  (Your engine: Chapter 20)

    The reference forward dequantizes the whole matrix each call, which saves memory but not
    time. A real W4A16 kernel (kernels/triton_quant.py) dequantizes tiles in registers so the
    weights cross the memory bus as 4-bit values.
    """

    def __init__(self, linear, bits=4, group_size=128, clip=1.0):
        super().__init__()
        self.in_features, self.out_features = linear.in_features, linear.out_features
        self.bits, self.group_size = bits, group_size
        codes, scales = quantize_groupwise(linear.weight.data, bits, group_size, clip)
        if bits == 4:
            self.register_buffer("packed", pack_int4(codes))
        else:
            self.register_buffer("packed", codes)
        self.register_buffer("scales", scales.to(linear.weight.dtype))
        self.bias = None if linear.bias is None else nn.Parameter(linear.bias.data.clone(), requires_grad=False)
        self.compute_dtype = linear.weight.dtype

    def codes(self):
        shape = (self.out_features, self.in_features)
        return unpack_int4(self.packed, shape) if self.bits == 4 else self.packed

    def dequantized_weight(self):
        """(Your engine: Chapter 20)"""
        return dequantize_groupwise(self.codes(), self.scales.float(), self.group_size).to(self.compute_dtype)

    def forward(self, x):
        """(Your engine: Chapter 20)"""
        return F.linear(x, self.dequantized_weight(), self.bias)

    def storage_bytes(self):
        return sum(t.numel() * t.element_size() for t in (self.packed, self.scales))

def quantize_model(model, bits=4, group_size=128, skip=("lm_head", "head", "gate", "router", "shared_expert_gate")):
    """Replace every nn.Linear with a QuantLinear, in place, unless a component of its name is in
    `skip` (output heads and MoE routers are small and sensitive, so they stay in high precision)."""
    replaced = 0
    for name, module in list(model.named_modules()):
        for child_name, child in list(module.named_children()):
            full = f"{name}.{child_name}" if name else child_name
            if isinstance(child, nn.Linear) and not set(full.split(".")) & set(skip):
                setattr(module, child_name, QuantLinear(child, bits, group_size))
                replaced += 1
    return replaced

quantize_model replaces every nn.Linear in place, except layers whose names mark them as sensitive: the LM head (its errors land directly on the logits) and MoE routers (small, and a flipped expert choice changes the computation discretely). On the tiny Qwen3 test model in FP32:

python run.py quant --bits 4

{"bits": 4, "bytes_before": 15739904, "bytes_after": 3124224,
 "logit_error": {"max_abs": 0.1275, "mean_abs": 0.0178, "relative_l2": 0.0396}, "same_top1": 0.953}
{"bits": 8, "bytes_before": 15739904, "bytes_after": 4959232,
 "logit_error": {"max_abs": 0.0075, "mean_abs": 0.00098, "relative_l2": 0.0022}, "same_top1": 1.0}

(The 8-bit line comes from --bits 8; both runs are from a laptop CPU.) The 8-bit model’s size isn’t a quarter of FP32’s because the embedding stays unquantized. At 4 bits, 95% of positions keep the same top token, with 4% relative logit error. Real models are more robust than this tiny random one, but measure, don’t assume.

Speed needs a kernel

QuantLinear.forward dequantizes the whole matrix to BF16 and calls an ordinary matmul. That saves memory but not time: the full-precision weights are materialized and read again every call, so decode is slower than BF16. The speedup only appears when a kernel reads the packed bytes and dequantizes them in registers, so weights cross the memory bus at 4 bits. That’s a W4A16 kernel: 4-bit weights, 16-bit activations, 16/32-bit math.

@triton.jit
def w4a16_kernel(x_ptr, packed_ptr, scales_ptr, y_ptr, M, N, K, groups,
                 GROUP: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    """(Your engine: Chapter 20)"""
    pid_m, pid_n = tl.program_id(0), tl.program_id(1)
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    half = tl.arange(0, BLOCK_K // 2)                       # byte index inside the K tile
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k0 in range(0, K, BLOCK_K):
        rk = k0 + tl.arange(0, BLOCK_K)
        x = tl.load(x_ptr + rm[:, None] * K + rk[None, :], mask=(rm[:, None] < M) & (rk[None, :] < K), other=0.0)
        # Packed row n holds K/2 bytes; byte j carries code 2j (low nibble) and 2j+1 (high nibble).
        byte_cols = k0 // 2 + half
        packed = tl.load(packed_ptr + rn[:, None] * (K // 2) + byte_cols[None, :],
                         mask=(rn[:, None] < N) & (byte_cols[None, :] < K // 2), other=0).to(tl.int32)
        low, high = packed & 0xF, (packed >> 4) & 0xF
        codes = tl.reshape(tl.join(low, high), (BLOCK_N, BLOCK_K))     # interleave back to k order
        codes = tl.where(codes >= 8, codes - 16, codes)                  # two's complement nibble
        scale = tl.load(scales_ptr + rn * groups + k0 // GROUP, mask=rn < N, other=0.0).to(tl.float32)
        w = codes.to(tl.float32) * scale[:, None]                        # [BLOCK_N, BLOCK_K]
        acc += tl.dot(x.to(tl.float32), tl.trans(w), input_precision="ieee")   # FP32, not TF32
    tl.store(y_ptr + rm[:, None] * N + rn[None, :], acc.to(y_ptr.dtype.element_ty),
             mask=(rm[:, None] < M) & (rn[None, :] < N))


def w4a16_matmul(x, packed, scales, group_size, out_features, block_m=16, block_n=32, block_k=32):
    """x [M, K]; packed uint8 [N * K / 2] (row-major, from quant.pack_int4); scales [N, K / group]."""
    check_device(x, packed, scales)
    m, k = x.shape
    if k % 2 or group_size % block_k or k % block_k:
        raise ValueError("K must be a multiple of BLOCK_K, and BLOCK_K must divide the group size")
    y = torch.empty((m, out_features), device=x.device, dtype=x.dtype)
    grid = (triton.cdiv(m, block_m), triton.cdiv(out_features, block_n))
    w4a16_kernel[grid](x.contiguous(), packed, scales.contiguous(), y, m, out_features, k, scales.shape[1],
                       GROUP=group_size, BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k)
    return y

It’s the Chapter 14 matmul with the B-tile loaded differently. Each program loads BLOCK_K / 2 bytes per weight row, splits each byte into its two nibbles, interleaves them back into $k$ order with tl.join and a reshape, sign-extends, and multiplies by the group’s scale. Requiring BLOCK_K to divide the group size means one scale per row per K-tile. As in Chapter 15, input_precision="ieee" keeps the FP32 tl.dot exact on the GPU instead of TF32.

What production W4A16 kernels (Marlin, Machete, AWQ’s GEMV, ExLlama) add: weight layouts pre-shuffled so a 128-bit load fills a tensor-core fragment directly, dequantization via bit tricks instead of shifts and multiplies, and split-K for the batch-1 decode shape. At batch 1 a good INT4 kernel approaches 3.5-3.9× BF16 decode speed; at large batches the kernel becomes compute-bound and the advantage shrinks, because the dequantization arithmetic is extra work.

Beyond round-to-nearest

Round-to-nearest with good groups is a strong baseline at 8 bits and a reasonable one at 4. The research methods address different parts of the problem:

The scaling identity behind SmoothQuant and AWQ is worth remembering: for a diagonal $S$, $XW^\top = (XS^{-1})(WS)^\top$ exactly. It moves magnitude between activations and weights without changing the function, so you can choose where the outliers live before you round.

Quantizing the KV cache

At long contexts, the KV cache dominates memory traffic (Chapter 16). Quantizing it to FP8 or INT8, with a scale per head or per token, halves the traffic of attention during decode. Keys are more sensitive than values (errors pass through the softmax’s exponential), and keys often have outlier channels, so KV quantization usually uses per-channel scales for K and per-token scales for V (KIVI, 2024). Chapter 16’s last stretch exercise implements a simple version.

Quantize a real checkpoint

For Qwen3-0.6B, the companion script compares BF16 against bitsandbytes NF4 or LLM.int8() on held-out instruction data, then saves and reloads the quantized model to confirm the output survives the round trip:

uv pip install -r workflow-requirements.txt      # transformers, safetensors, peft
uv pip install 'bitsandbytes>=0.48'               # check platform support first
python quantize_checkpoint.py --model-dir models/Qwen3-0.6B --bits 4 \
    --valid-file data/sft-valid.jsonl --output runs/qwen3-nf4

It reports the response loss of both models, the KL divergence between their next-token distributions, and parameter bytes. These are what to look at when judging a quantization: loss and KL on your own data first, then task scores and some generations read by a human.

Warning

bitsandbytes ships compiled CUDA binaries. If its CUDA version doesn’t match PyTorch’s, import fails or 4-bit layers silently fall back. The bitsandbytes installation guide explains BNB_CUDA_VERSION; Chapter 22 shows how to diagnose it. The exported checkpoint runs in Transformers + bitsandbytes, not in your engine’s loader, which expects ordinary floating-point tensors.

An evaluation protocol

1. Fix the evaluation set, token IDs and decoding settings before quantizing anything.
2. Report held-out loss (or perplexity) and the KL from the original model, then task metrics.
3. Measure memory (weights, KV cache, peak) and speed separately for prefill and decode, with batch size, lengths, GPU, backend and warmup recorded.
4. Re-evaluate after every transformation that follows (merging adapters, editing, exporting): quantization doesn’t commute with them.

Build it

pytest tests/test_ch20_quant.py
python run.py quant --impl engine --bits 4

Stretch exercises

1. ★ Plot held-out output error against bits per weight (including scales) for group sizes 16-256 on a real Qwen3-0.6B layer. Where: experiments/ch20.py (create it), using engine.quant.quantize_groupwise and dequantize_groupwise.
2. ★★ Make QuantLinear call your W4A16 kernel on CUDA, and measure decode tokens/s for Qwen3-0.6B against BF16 in graph mode (Chapter 19). Where: QuantLinear.forward in engine/quant.py, calling engine.kernels.triton_quant.w4a16_matmul.
3. ★★ Implement AWQ’s core: for one linear layer, search a per-input-channel scale $s_j = \bar{|x_j|}^{\alpha}$ over $\alpha \in [0, 1]$ that minimizes calibration output error after 4-bit quantization of $W \operatorname{diag}(s)$. Where: add a calibration/search helper in engine/quant.py; Chapter 39’s engine/formats/awq.py provides the later integration.
4. ★★★ Implement GPTQ for one layer: accumulate $H = 2X^\top X$ from calibration inputs, then quantize column by column, spreading each column’s error over the remaining columns with the Cholesky factor of $H^{-1}$. Compare with round-to-nearest at 3 and 4 bits. Where: add a one-layer GPTQ helper in engine/quant.py; Chapter 39’s engine/formats/gptq.py provides the later integration.

Check your understanding

1. Why does halving the bytes per weight roughly double the decode ceiling, but barely change large-batch prefill speed?
2. Why do smaller groups reduce error, and what do they cost?
3. Why can a clip below 1.0 reduce output error while increasing weight error?
4. Why doesn’t dequantize-then-matmul make decode faster?
5. Why are the LM head and MoE routers often left unquantized?

Going deeper

• GPU Mode L7 (Advanced quantization: weight-only and dynamic quantization in torchao, with Triton kernels), L33 (BitBLAS: mixed-precision kernels for arbitrary bit widths), L30 (Quantized training).
• PMPP Appendix A (Numerical considerations): the low-precision floating-point formats (FP16, BF16, FP8, block-scaled formats) and how GPUs implement them.
• Papers: LLM.int8() (2208.07339), SmoothQuant (2211.10438), GPTQ (2210.17323), AWQ (2306.00978), QLoRA/NF4 (2305.14314), KIVI (2402.02750); Micikevicius et al., FP8 Formats for Deep Learning (2022).
• gpt-fast’s quantize.py (int8 and int4 weight-only with GPTQ in a few hundred lines); the Marlin kernel write-up; Maarten Grootendorst’s A Visual Guide to Quantization.

21. Fine-tuning: classifiers and instruction following

What fine-tuning changes

A pretrained model has learned language, facts and patterns from trillions of tokens, but its only skill is continuing text. Ask a base model “What is the capital of France?” and it may continue with three more quiz questions. Fine-tuning keeps training the same weights on a smaller, targeted dataset so the model does a specific job.

The important intuition: fine-tuning mostly teaches format and behavior, and draws on knowledge already in the weights. A thousand examples can teach a model to answer in one sentence and stop. They can’t teach it chemistry. You’ll see this directly below, when a model trained from scratch learns the answer format perfectly and the facts not at all.

BALLM Chapters 6 and 7 build both classification and instruction tuning on GPT-2; this chapter follows them closely, then moves the same ideas to Qwen3.

Classification: replace the head

A language model ends in a head that maps the final hidden state to vocabulary scores. For a classifier, replace it with a head that maps to class scores: for spam detection, nn.Linear(width, 2) instead of nn.Linear(width, 50257).

Which position’s hidden state should be classified? Under a causal mask, position $t$ has seen tokens $0 \dots t$ only. The last token is the only one that has seen the whole message. With right-padded batches, the last real token is at length - 1, not at -1:

def replace_head(model, num_classes, train_last_blocks=1):
    """Freeze the model, swap the vocabulary head for a num_classes head, and unfreeze the
    last few blocks and the final norm (BALLM Chapter 6's recipe)."""
    for parameter in model.parameters():
        parameter.requires_grad_(False)
    width = model.head.in_features
    p = model.head.weight
    model.head = nn.Linear(width, num_classes, device=p.device, dtype=p.dtype)   # new, trainable
    for block in list(model.blocks)[len(model.blocks) - train_last_blocks:]:
        block.requires_grad_(True)
    model.norm.requires_grad_(True)
    return model


def last_token_logits(model, ids, lengths):
    """Class scores read at each sequence's last REAL token.  (Your engine: Chapter 21)

    Under a causal mask only the last position has seen the whole text, and with right padding
    that position is lengths - 1, not -1.
    """
    logits = model(ids)
    return logits[torch.arange(ids.shape[0], device=ids.device), lengths - 1]


def pad_sequences(sequences, pad_id, device="cpu"):
    width = max(len(s) for s in sequences)
    ids = torch.full((len(sequences), width), pad_id, dtype=torch.long)
    for row, seq in enumerate(sequences):
        ids[row, :len(seq)] = torch.tensor(seq)
    return ids.to(device), torch.tensor([len(s) for s in sequences], device=device)


@torch.no_grad()
def accuracy(model, sequences, labels, pad_id, batch_size=32):
    model.eval()
    device = next(model.parameters()).device
    correct = 0
    for i in range(0, len(sequences), batch_size):
        ids, lengths = pad_sequences(sequences[i:i + batch_size], pad_id, device)
        predicted = last_token_logits(model, ids, lengths).argmax(-1).cpu()
        correct += int((predicted == torch.tensor(labels[i:i + batch_size])).sum())
    return correct / len(sequences)


def train_classifier(model, sequences, labels, pad_id, steps, batch_size=16, lr=5e-4, seed=0):
    device = next(model.parameters()).device
    rng = random.Random(seed)
    trainable = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(trainable, lr=lr, weight_decay=0.1)
    for _ in range(steps):
        model.train()
        index = rng.sample(range(len(sequences)), min(batch_size, len(sequences)))
        ids, lengths = pad_sequences([sequences[i] for i in index], pad_id, device)
        target = torch.tensor([labels[i] for i in index], device=device)
        loss = F.cross_entropy(last_token_logits(model, ids, lengths), target)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
    return loss.item()

replace_head follows BALLM’s recipe: freeze everything, add a fresh trainable head, and unfreeze the last transformer block and the final norm. Training only the top of the network is fast, needs little data, and keeps the pretrained features intact. Train more blocks when you have more data and the task differs more from pretraining.

Run it

Without any download, the demo trains a small GPT on a synthetic task that requires reading the whole sequence (does the token 7 appear anywhere?):

python run.py classify --steps 300

{"before": 0.49166666666666664}
{"train_accuracy": 1.0, "test_accuracy": 1.0}

That took 12 seconds on a laptop CPU. For a real task, download the SMS Spam Collection (5,574 messages, tab-separated label<TAB>text) and pass --data SMSSpamCollection. With pretrained GPT-2 small, frozen except the last block, BALLM reports 95.7% test accuracy after five epochs (BALLM §6.7); balance the classes first, as the book does, so that “always say ham” doesn’t score 87%.

Note

The demo trains its GPT from scratch with all blocks unfrozen (train_last_blocks=layers), because there are no pretrained features to preserve. With GPT-2 weights, keep the default of one block.

Instruction tuning

Instruction tuning is still next-token prediction. What changes is the data: each example is a prompt and a desired response, joined by a template, and the loss covers only the response.

A template

The template marks where the instruction ends and the answer begins. BALLM uses the Alpaca format:

Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
Rewrite the sentence using a simile.

### Input:
The car is very fast.

### Response:
The car is as fast as lightning.<|endoftext|>

Chat models use their own chat template with special tokens instead (Qwen’s <|im_start|> and <|im_end|>, Chapter 4). Two rules matter more than the template you choose: training and inference must use exactly the same format, and the response must end with the stop token the engine stops on. A model never trained to emit <|endoftext|> never stops.

def format_prompt(entry):
    """The Alpaca-style template from BALLM Chapter 7. The response follows the final header."""
    text = ("Below is an instruction that describes a task. "
            "Write a response that appropriately completes the request."
            f"\n\n### Instruction:\n{entry['instruction']}")
    if entry.get("input"):
        text += f"\n\n### Input:\n{entry['input']}"
    return text + "\n\n### Response:\n"


def load_instructions(path):
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def split_entries(entries, valid_fraction=0.1, test_fraction=0.1, seed=0):
    """Shuffle once, then split by item: every version of an item stays in one split."""
    order = list(range(len(entries)))
    random.Random(seed).shuffle(order)
    n_test, n_valid = int(len(entries) * test_fraction), int(len(entries) * valid_fraction)
    pick = lambda idx: [entries[i] for i in idx]
    return pick(order[n_test + n_valid:]), pick(order[n_test:n_test + n_valid]), pick(order[:n_test])


def encode_examples(entries, encode, eos_id, max_length):
    """(prompt_ids, response_ids + [eos]) pairs. Over-long examples are dropped and counted, never
    truncated: truncation would silently cut the response, the only part that is learned."""
    examples, dropped = [], 0
    for entry in entries:
        prompt, response = encode(format_prompt(entry)), encode(entry["output"]) + [eos_id]
        if len(prompt) + len(response) > max_length + 1:      # inputs are one shorter than the pair
            dropped += 1
            continue
        examples.append((prompt, response))
    return examples, dropped

The dataset is BALLM’s 1,100 instruction-response pairs, in data/instruction-data.json. split_entries splits it into 935 training, 55 validation and 110 test examples before any processing, so no item leaks between splits. encode_examples drops over-long examples instead of truncating them; truncation would cut off the response, which is the only part that’s learned.

Response-only loss and the label shift

Should the model learn to predict the prompt tokens too? Usually not: the prompt is given at inference time, and learning to generate instructions wastes capacity. So prompt positions get the target $-100$, which F.cross_entropy ignores by default. Padding positions get $-100$ too.

The subtle part is the shift. For a prompt [10, 11] and a response [20, 21, EOS], the complete sequence is [10, 11, 20, 21, EOS]. Inputs are all tokens but the last; targets are all tokens but the first:

inputs     10    11    20    21
targets  -100    20    21   EOS

The input 11, the last prompt token, must predict 20, the first response token. That’s the most important prediction in the example, and the one most often lost by masking on the wrong side of the shift.

def instruction_batch(examples, pad_id, device="cpu"):
    """examples: list of (prompt_ids, response_ids_ending_in_eos).  (Your engine: Chapter 21)

    Returns inputs x and targets y, both [B, L-1], already shifted for next-token loss.
    Prompt and padding targets are -100 so cross_entropy ignores them; only response tokens
    (including the final EOS) are learned.
    """
    if not examples or any(not p or not r for p, r in examples):
        raise ValueError("Each example needs a non-empty prompt and response")
    width = max(len(p) + len(r) for p, r in examples)
    ids = torch.full((len(examples), width), pad_id, dtype=torch.long)
    labels = torch.full_like(ids, -100)
    for row, (prompt, response) in enumerate(examples):
        complete = list(prompt) + list(response)
        ids[row, :len(complete)] = torch.tensor(complete)
        labels[row, len(prompt):len(complete)] = torch.tensor(list(response))
    return ids[:, :-1].to(device), labels[:, 1:].to(device)

Warning

Hugging Face models shift internally: you pass unshifted input_ids and labels of the same length, and the model drops the first label. If you shift in your data pipeline and pass the result to a Hugging Face model, every target moves one token too far and the model learns to predict two tokens ahead. Your own GPT expects shifted targets; finetune.py for Qwen3 passes unshifted ones. Know which convention each model uses.

When the pad token equals the EOS token (common: GPT-2 has no pad token), mask padding by position, not by ID. Masking every occurrence of the EOS ID would also mask the real end of the response, and the model would never learn to stop.

The training loop

def response_loss(model, x, y):
    """Mean cross-entropy over response targets only; -100 marks prompt and padding."""
    logits = model(x)
    return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1), ignore_index=-100)


@torch.no_grad()
def evaluate_responses(model, examples, pad_id, batch_size=8):
    """Average loss per response token (not per batch), so long and short answers weigh fairly."""
    model.eval()
    device = next(model.parameters()).device
    total, count = 0.0, 0
    for i in range(0, len(examples), batch_size):
        x, y = instruction_batch(examples[i:i + batch_size], pad_id, device)
        logits = model(x)
        total += F.cross_entropy(logits.reshape(-1, logits.shape[-1]), y.reshape(-1),
                                 ignore_index=-100, reduction="sum").item()
        count += int((y != -100).sum())
    return total / max(count, 1)


def finetune(model, train_examples, valid_examples, pad_id, steps, batch_size=8, lr=5e-5,
             weight_decay=0.1, warmup=10, eval_every=25, seed=0, log=print):
    """AdamW on response-only loss, each batch padded only to its own longest example."""
    device = next(model.parameters()).device
    rng = random.Random(seed)
    trainable = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(trainable, lr=lr, weight_decay=weight_decay)
    history = []
    for step in range(steps):
        scale = (step + 1) / warmup if step < warmup else 0.5 * (1 + math.cos(math.pi * (step - warmup) / max(1, steps - warmup)))
        for group in optimizer.param_groups:
            group["lr"] = lr * scale
        model.train()
        x, y = instruction_batch(rng.sample(train_examples, min(batch_size, len(train_examples))), pad_id, device)
        loss = response_loss(model, x, y)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(trainable, 1.0)
        optimizer.step()
        if step % eval_every == 0 or step + 1 == steps:
            record = {"step": step + 1, "train_loss": round(loss.item(), 4),
                      "valid_loss": round(evaluate_responses(model, valid_examples, pad_id), 4)}
            history.append(record)
            if log:
                log(json.dumps(record))
    return history

Two details differ from Chapter 7’s loop:

• Dynamic padding. Each batch is padded to its own longest example, not to the dataset’s. Instruction lengths vary a lot, and padding to the maximum wastes compute.
• Token-weighted evaluation. Validation loss is the sum over response tokens divided by their count, not an average of batch averages, which would weight a three-word answer the same as a three-sentence one.

Run it

Without pretrained weights, run.py sft trains its own BPE and a 4-layer GPT from scratch on the 935 training examples only:

python run.py sft --steps 600 --width 192 --layers 4

On a laptop CPU, in 55 seconds:

{"train": 935, "valid": 55, "test": 110, "dropped_overlong": 0}
{"step": 1, "train_loss": 6.9229, "valid_loss": 6.8699}
{"step": 301, "train_loss": 3.9705, "valid_loss": 4.2244}
{"step": 600, "train_loss": 2.8521, "valid_loss": 3.8971}
### What is the contraction for 'it is'?
  expected: The contraction for 'it is' is 'it's.'
  model:    The gas' is 'bua'.
### Convert the number 5 from decimal to binary.
  expected: The binary equivalent of the decimal number 5 is 101.
  model:    The approximately 14 is 24.
### Rewrite the sentence using a simile. The baby is very cute.
  expected: The baby is as cute as a button.
  model:    The ma is a pa.

The model has learned the shape of an answer (“The X is Y.”, then stop) and nothing else. It has never read enough text to know words, let alone facts. That’s fine-tuning without pretraining.

Now the real thing. Load GPT-2 (Chapter 9) and run the same command with --model-dir models/gpt2-medium (a GPU makes this a couple of minutes; BALLM reports 15.8 minutes for GPT-2 medium on an M3 MacBook Air CPU). BALLM’s GPT-2 medium (355M), after two epochs on this dataset, answered (BALLM §7.7):

### Rewrite the sentence using a simile. The car is very fast.
  expected: The car is as fast as lightning.
  model:    The car is as fast as a bullet.
### What type of cloud is typically associated with thunderstorms?
  expected: The type of cloud typically associated with thunderstorms is cumulonimbus.
  model:    The type of cloud associated with thunderstorms is a cumulus cloud.

Same data, same loop; the difference is entirely pretraining. Note the second answer: fluent, confident and wrong. Evaluating instruction-tuned models is hard.

Fine-tuning Qwen3-0.6B

For the real checkpoint, finetune.py uses Hugging Face Transformers for the model, so you can focus on the data and the measurements. The same instruction data is in chat format in data/instruction-{train,valid,test}.jsonl:

uv pip install -r optional-requirements.txt
python finetune.py train --mode full --model-dir models/Qwen3-0.6B \
    --train-file data/instruction-train.jsonl --valid-file data/instruction-valid.jsonl \
    --output runs/qwen3-sft --steps 200 --lr 1e-5 --batch-size 4 --accumulation 4
python finetune.py evaluate --model-dir models/Qwen3-0.6B \
    --valid-file data/instruction-valid.jsonl --output runs/qwen3-sft

What the script does, each of which you’ve now seen from first principles:

• Formats every example with Qwen’s own chat template (enable_thinking=False), and checks that the prompt-with-generation-prompt is an exact prefix of the full conversation, so the mask boundary is right.
• Masks the prompt and padding, keeps the <|im_end|> that ends the response, and refuses over-long examples.
• Accumulates gradients over --accumulation micro-batches, dividing by the total number of response tokens in the update (not by the number of micro-batches).
• Saves the model with a manifest.json recording the base config’s checksum, dataset checksums, arguments and library versions, and evaluate reloads it in a fresh process.

Qwen3-0.6B is already instruction-tuned, so this run teaches it the dataset’s terse answer style. To see the full effect of SFT, start from a base checkpoint such as Qwen/Qwen3-0.6B-Base.

The memory of full fine-tuning

Inference needs the weights. Full fine-tuning with AdamW in mixed precision needs, per parameter:

For 0.6B parameters, that’s about 11 GB before activations; for an 8B model, 144 GB. Activations add more, growing with batch size × sequence length × layers (gradient checkpointing trades recomputation for most of it). That’s why the next chapter’s methods exist.

Evaluating a fine-tuned model

A falling training loss proves only that the model memorizes the training set. In increasing order of cost and value:

1. Held-out loss on response tokens, computed identically before and after.
2. Exact checks where answers are checkable: classification accuracy, arithmetic, format compliance (“did it stop?”, “is it valid JSON?”).
3. Regression checks on prompts unrelated to the fine-tuning data: did general ability survive? Small datasets and high learning rates cause catastrophic forgetting.
4. Model-as-judge: a stronger model scores responses against references. BALLM §7.8 does this with Llama 3 8B through Ollama, scoring each response 0-100. Useful, but biased toward long and confident answers; spot-check it.
5. Human review of a sample, especially of failures.

Build it

pytest tests/test_ch21_finetune.py
python run.py classify --impl engine
python run.py sft --impl engine --steps 600 --width 192

Stretch exercises

1. ★ Train on all tokens (prompt included) instead of response-only, with the same budget. Compare held-out response loss and the test answers. Where: the response mask built by instruction_batch in engine/lora.py, consumed by engine.sft.response_loss.
2. ★★ Classify SMS spam with pretrained GPT-2, comparing three settings: head only, head + last block, all layers. Plot test accuracy against trainable parameters. Where: experiments/ch21.py (create it), adapting run.py’s cmd_classify and calling engine.sft.train_classifier.
3. ★★ Implement sequence packing: concatenate several examples into one row, with a block-diagonal causal mask and positions restarting at 0 for each example. Your causal_attention already takes an allowed mask (Chapter 5). Measure the speedup over dynamic padding. Where: add a packed batch builder in engine/lora.py; pass its masks/positions through engine/sft.py and engine/gpt.py to engine.attention.causal_attention.
4. ★★★ Implement DPO (Rafailov et al., 2023) on BALLM’s preference dataset (ch07/04_preference-tuning-with-dpo in Raschka’s repository): the loss compares the policy’s and a frozen reference model’s log-probabilities of the chosen and rejected responses. Where: add a DPO loss/training helper in engine/sft.py; load chosen/rejected pairs in experiments/ch21.py (create it).

Check your understanding

1. Why does a classifier read the last token’s hidden state, not the first?
2. In the shifted inputs and targets, which input position predicts the first response token?
3. Why must padding be masked by position when the pad ID equals the EOS ID?
4. Why does full fine-tuning need many times the memory of inference?
5. Why did the from-scratch model learn the answer format but none of the facts?

Going deeper

• BALLM Chapter 6 (classification fine-tuning: §6.5 adding the head, §6.6 the last token, §6.7 training) and Chapter 7 (instruction fine-tuning: §7.3 batching and masking, §7.6 training, §7.8 evaluation with a judge model). Raschka’s repository has bonus material on DPO, LoRA variants and larger models.
• Ouyang et al., Training language models to follow instructions with human feedback (InstructGPT, 2022); Taori et al., Alpaca (2023); Zhou et al., LIMA: Less Is More for Alignment (2023), on how little data format-teaching needs.
• Rafailov et al., Direct Preference Optimization (2023).
• Hugging Face TRL’s SFTTrainer documentation, the production version of this chapter’s loop.

22. LoRA and QLoRA

The cost of full fine-tuning

Chapter 21 ended with an uncomfortable table: full fine-tuning with AdamW needs about 18 bytes per parameter, 144 GB for an 8B model before activations. And every fine-tuned variant is a complete copy of the model: ten customers with ten fine-tunes of an 8B model means 160 GB of checkpoints.

LoRA (Low-Rank Adaptation, Hu et al., 2021) fixes both. It freezes every pretrained weight and learns a small correction for selected matrices. Only the corrections get gradients and optimizer state, and only they are saved.

The low-rank idea

Fine-tuning changes a weight matrix $W \in \mathbb{R}^{d_\text{out} \times d_\text{in}}$ to $W + \Delta W$. Empirically, the useful $\Delta W$ for adapting a pretrained model has low intrinsic rank (Aghajanyan et al., 2020): it can be well approximated by a product of two thin matrices. So LoRA parameterizes it that way:

$$ W’ = W + \frac{\alpha}{r} B A, \qquad A \in \mathbb{R}^{r \times d_\text{in}},\ B \in \mathbb{R}^{d_\text{out} \times r},\ r \ll \min(d_\text{in}, d_\text{out}). $$

For a 4,096 × 4,096 projection, full fine-tuning trains 16.8M numbers; LoRA with $r = 8$ trains $8 \times (4096 + 4096) = 65{,}536$, which is 0.39%.

The forward pass never forms $BA$. It computes the frozen path and the low-rank path separately and adds them:

$$ y = W x + \frac{\alpha}{r}, B (A x). $$

$Ax$ is an $r$-vector, cheap to compute. The extra FLOPs are $2r(d_\text{in} + d_\text{out})$ per token, under 1% of the layer.

class LoRALinear(nn.Module):
    """y = base(x) + (alpha / r) * B(A(x)), with the base frozen.  (Your engine: Chapter 22)

    A [r, in] starts random and B [out, r] starts at zero, so the adapted layer initially
    equals the base exactly; B receives a gradient on the first step, A on the next.
    """

    def __init__(self, base, rank=8, alpha=16):
        super().__init__()
        if rank < 1:
            raise ValueError("rank must be positive")
        self.base = base
        for parameter in base.parameters():
            parameter.requires_grad_(False)
        self.scale = alpha / rank
        self.A = nn.Parameter(base.weight.new_empty(rank, base.in_features))
        self.B = nn.Parameter(base.weight.new_zeros(base.out_features, rank))
        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))

    def forward(self, x):
        """(Your engine: Chapter 22)"""
        return self.base(x) + F.linear(F.linear(x, self.A), self.B) * self.scale

    @torch.no_grad()
    def merged(self):
        """A plain nn.Linear with W + scale * B @ A folded in. Do not also keep the adapter active.  (Your engine: Chapter 22)"""
        out = nn.Linear(self.base.in_features, self.base.out_features, bias=self.base.bias is not None,
                        device=self.base.weight.device, dtype=self.base.weight.dtype)
        out.weight.copy_(self.base.weight + self.scale * (self.B @ self.A))
        if self.base.bias is not None:
            out.bias.copy_(self.base.bias)
        return out

Initialization and the scale

• $B = 0$, $A$ random. The adapted model starts exactly equal to the base, so training begins from the pretrained behavior. On the first step, $B$ gets a gradient ($\partial L / \partial B = \frac{\alpha}{r} g, (Ax)^\top$, which is non-zero) while $A$’s gradient is zero ($\partial L/\partial A = \frac{\alpha}{r} B^\top g, x^\top$, and $B = 0$). After one update, $B \ne 0$ and both train. The milestone test checks this.
• $\alpha/r$. Scaling the update by $\alpha/r$ keeps its size roughly stable when you change $r$, so a learning rate tuned at one rank still works at another. A common choice is $\alpha = 2r$. (rsLoRA, Kalajdzievski 2023, argues for $\alpha/\sqrt{r}$ at large ranks.)

Which layers to adapt

def add_lora(model, targets=("q_proj", "v_proj"), rank=8, alpha=16):
    """Freeze the model, then wrap every Linear whose attribute name is in targets."""
    for parameter in model.parameters():
        parameter.requires_grad_(False)
    wrapped = []
    for name, module in list(model.named_modules()):
        for child_name, child in list(module.named_children()):
            if child_name in targets and isinstance(child, nn.Linear):
                setattr(module, child_name, LoRALinear(child, rank, alpha))
                wrapped.append(f"{name}.{child_name}")
    return wrapped


def merge_lora(model):
    """Replace every LoRALinear with its merged Linear, in place."""
    for module in list(model.modules()):
        for child_name, child in list(module.named_children()):
            if isinstance(child, LoRALinear):
                setattr(module, child_name, child.merged())
    return model

The original paper adapted only the attention query and value projections. The QLoRA paper found that adapting all linear layers (attention and MLP) matters more than the rank: with all layers adapted, ranks from 8 to 64 performed similarly. Embeddings and the LM head are usually left alone; for a model with tied embeddings, adapting the head would also change the input embedding through the shared tensor.

Run it: a new style with 5% of the parameters

Chapter 21’s run.py sft saved its instruction-tuned GPT to runs/sft.pt. Now teach it a new behavior, answering in capital letters, by training LoRA adapters on every linear layer with the same instructions and upper-cased responses:

python run.py lora --steps 300

On a laptop CPU, in 33 seconds:

{"before_lora": "The gas' is 'bua'.", "valid_loss_on_caps": 6.5372, "dropped_overlong": 0}
{"wrapped_layers": 16, "trainable": 98304, "total": 2025600, "trainable_fraction": 0.0485}
{"step": 1, "train_loss": 6.5902, "valid_loss": 6.4817}
{"step": 151, "train_loss": 3.6265, "valid_loss": 3.6524}
{"step": 300, "train_loss": 3.2451, "valid_loss": 3.5377}
{"instruction": "What is the contraction for 'it is'?", "after_lora": "THE OMES 'TES 'S 'S 'TETES 'TETETETETESTETESSTESTESTE"}
{"instruction": "Convert the number 5 from decimal to binary.", "after_lora": "15 MES: 14"}
{"merged_max_abs_difference": 4.410743713378906e-06, "parameters_after_merge": 2025600}

The style moved completely; the content was nonsense before and remains so (this model never pretrained), and the first answer degenerates into repetition, Chapter 8’s classic greedy failure. The fraction is high here only because the model is tiny. For Qwen3-0.6B, rank 8 on q_proj and v_proj trains 1.15M parameters, 0.19% of the model. On pretrained GPT-2 small, BALLM Appendix E’s LoRA (rank 16 on every linear layer, 2.7M trainable parameters, 2% of the model) reaches 98.0% test accuracy on spam classification, slightly above full fine-tuning of the last block in Chapter 6.

The last line is the merge check, discussed next.

Merge, or keep adapters separate

After training, an adapter can be merged: compute $W + \frac{\alpha}{r}BA$ once and replace $W$. The merged model is an ordinary model with zero extra inference cost. The demo’s merge changed the logits by $4 \times 10^{-6}$, floating-point rounding from adding in a different order.

Or keep adapters separate. One base model in GPU memory can then serve hundreds of fine-tunes, each a few megabytes, applying each request’s adapter in its forward pass. Batching requests with different adapters needs a gathered low-rank matmul (S-LoRA and Punica’s BGMV kernels; vLLM and SGLang support this). That’s how serving providers host many customer fine-tunes cheaply.

Rules that prevent subtle bugs:

• Never apply an adapter to a model it has already been merged into: the update would count twice.
• A merged model that you then quantize is a different model from the adapter-on-quantized-base you trained with QLoRA. Evaluate the artifact you’ll ship.
• An adapter is only valid for the exact base weights it was trained on. Record the base’s revision with the adapter (finetune.py writes it to its manifest).

QLoRA: a 4-bit frozen base

LoRA removes gradients and optimizer state for the base, but the base weights themselves still sit in BF16. QLoRA (Dettmers et al., 2023) stores the frozen base in 4 bits and dequantizes on the fly in the forward and backward passes, while the adapters train in BF16. It combines three ideas:

• NF4, a 4-bit data type whose 16 levels are quantiles of a normal distribution, matching how pretrained weights are distributed (Chapter 20). Unlike Chapter 20’s uniform INT4 grid, dequantizing is a table lookup.
• Double quantization: the per-block scales (one per 64 weights) are themselves quantized to 8 bits, saving about 0.37 bits per parameter.
• Paged optimizers: optimizer state that can spill to CPU memory during memory spikes.

Memory per parameter of the base, with activations and the small adapters aside:

That’s why QLoRA fine-tunes an 8B model on a single consumer GPU, and why the QLoRA paper could fine-tune a 65B model on one 48 GB GPU. The cost is speed: dequantizing every weight in every forward and backward pass makes each step slower than BF16 LoRA.

LoRA and QLoRA on Qwen3

finetune.py uses Hugging Face PEFT for the adapters, with the same data and loss handling as Chapter 21:

uv pip install -r workflow-requirements.txt
python finetune.py train --mode lora --rank 8 --model-dir models/Qwen3-0.6B \
    --train-file data/instruction-train.jsonl --valid-file data/instruction-valid.jsonl \
    --output runs/qwen3-lora --steps 200 --lr 2e-4 --batch-size 4 --accumulation 4
python finetune.py evaluate --model-dir models/Qwen3-0.6B \
    --valid-file data/instruction-valid.jsonl --output runs/qwen3-lora

The saved directory holds only the adapter (a few MB) and the tokenizer. Note the learning rate: LoRA typically needs 10-20× the learning rate of full fine-tuning, because the update starts at zero and lives in a small subspace.

For QLoRA, install bitsandbytes and pass --mode qlora. The script loads the base with NF4, double quantization and BF16 compute, prepares it for k-bit training (casting norms to FP32, enabling gradient checkpointing), and adds the same adapters.

When bitsandbytes’ CUDA build doesn’t match PyTorch’s

bitsandbytes ships compiled CUDA libraries for specific CUDA versions. If PyTorch uses a newer CUDA than any bundled library, bitsandbytes looks for a file that doesn’t exist (for example libbitsandbytes_cuda132.so with PyTorch on CUDA 13.2) and fails or falls back. Diagnose first:

python -m bitsandbytes            # prints the CUDA version it detected and the library it loaded

If a library for an older compatible CUDA is bundled and that toolkit is installed, the documented override selects it, as in this example for a CUDA 13.0 build:

export BNB_CUDA_VERSION=130
export LD_LIBRARY_PATH="/usr/local/cuda-13.0/lib64${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
python -m bitsandbytes            # must now report success before you train

This is an environment-specific workaround, not a fix for an unsupported GPU or platform (check the installation guide for ARM64 and newer architectures). If QLoRA can’t run, use plain LoRA; don’t call the result QLoRA.

Build it

pytest tests/test_ch22_lora.py
python run.py lora --impl engine --steps 300

Stretch exercises

1. ★ Run run.py lora with ranks 1, 2, 8 and 32, and plot final validation loss against trainable parameters. Where does a larger rank stop helping? Where: terminal: python run.py lora --impl engine --rank R for each rank; record results in experiments/ch22.py (create it).
2. ★★ Compare full SFT (Chapter 21) and LoRA on GPT-2 with the same data and number of steps: held-out loss, peak memory (torch.cuda.max_memory_allocated) and seconds per step. Where: experiments/ch22.py (create it), adapting run.py’s cmd_sft and cmd_lora.
3. ★★ Serve two adapters at once: for a batch where row 0 uses adapter 1 and row 1 uses adapter 2, compute $B_{i}(A_{i} x)$ per row with gathered weights, and verify against running each row separately. Where: add a gathered-adapter helper in engine/lora.py; Chapter 43 integrates it in engine/serve/adapters.py.
4. ★★★ Implement DoRA (Liu et al., 2024): decompose $W$ into magnitude and direction, $W’ = m \cdot \frac{W + BA}{\lVert W + BA \rVert_c}$ with a trainable per-column magnitude $m$, and compare with LoRA at equal parameter count. Where: add a DoRA layer beside LoRALinear in engine/lora.py.

Check your understanding

1. Why does initializing $B$ to zero make the adapted model equal to the base at the start?
2. Why is $A$’s gradient zero on the first step, and why doesn’t that stop training?
3. What does merging cost at inference time, and what does keeping adapters separate buy you?
4. Why does QLoRA reduce memory but slow each step down?
5. Why can’t an adapter trained on one base model be applied to another?

Going deeper

• BALLM Appendix E (Parameter-efficient fine-tuning with LoRA): LoRA from scratch on GPT-2 for spam classification, the source of this chapter’s comparison.
• Hu et al., LoRA (2021); Dettmers et al., QLoRA (2023); Aghajanyan et al., Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning (2020); Liu et al., DoRA (2024); Sheng et al., S-LoRA (2023) and Chen et al., Punica (2023) for multi-adapter serving.
• GPU Mode L32 (Unsloth: LLM systems engineering): fused Triton kernels for LoRA and QLoRA training.
• The Hugging Face PEFT documentation (LoRA configuration, merging, quantized bases) and Sebastian Raschka’s article Practical Tips for Finetuning LLMs Using LoRA.

23. Inside the model: hooks, the logit lens and steering

The residual stream

Look at Chapter 6’s block again: x = x + attention(norm(x)), then x = x + mlp(norm(x)). Every layer reads the current $x$ and adds its result back. Nothing ever overwrites $x$. So the hidden state after layer $L$ is a sum:

$$ x_L = \underbrace{x_0}{\text{embedding}} + \sum{\ell=1}^{L} \big(\text{attn}\ell + \text{mlp}\ell\big). $$

This residual stream is a shared workspace: a $d$-dimensional vector per position that each component reads from and writes to, and that the final norm and head decode into the next-token prediction. Two consequences make this chapter possible:

1. Because the head reads the stream through one linear map, you can apply it to the stream at any layer and see what’s there. That’s the logit lens.
2. Because components communicate by adding vectors, a feature the model uses is often represented as a direction in this space (the linear representation hypothesis). Adding or removing a direction can change behavior directly.

Hooks: observing without changing code

PyTorch’s forward hooks run a function after a module’s forward pass, with access to its inputs and output. If the function returns a value, it replaces the output. That’s all you need to read and edit activations in any model, yours or a library’s, without changing its source:

@contextmanager
def capture(module, store, name):
    """Record a module's output into store[name] for every forward call inside the block."""
    def hook(_module, _inputs, output):
        store.setdefault(name, []).append((output[0] if isinstance(output, tuple) else output).detach())
    handle = module.register_forward_hook(hook)
    try:
        yield store
    finally:
        handle.remove()                       # always detach, even when the forward raises


@torch.no_grad()
def residuals(model, ids, layers=None):
    """Hidden state after each requested decoder layer, as {layer_index: [B, T, D]}."""
    blocks = list(_blocks(model))
    layers = range(len(blocks)) if layers is None else layers
    store, handles = {}, []
    for i in layers:
        handles.append(blocks[i].register_forward_hook(
            lambda _m, _inp, out, i=i: store.__setitem__(i, (out[0] if isinstance(out, tuple) else out).detach())))
    try:
        model(ids)
    finally:
        for handle in handles:
            handle.remove()
    return store

Two habits matter. Always remove hooks in a finally block: a forgotten hook keeps editing every later forward pass, and an exception inside a forward is exactly when you’d forget. And detach() anything you store, or you’ll keep the whole autograd graph alive.

The logit lens

Apply the final norm and the head to the residual stream after each layer, and look at the top tokens:

@torch.no_grad()
def logit_lens(model, ids, top=5):
    """Decode every layer's residual through the final norm and head: what would the model
    predict if it stopped here? Returns {layer: top token IDs at the last position}."""
    norm = model.norm if hasattr(model, "norm") else model.model.norm
    head = model.head if hasattr(model, "head") else model.lm_head
    return {layer: head(norm(h[:, -1])).topk(top, dim=-1).indices[0].tolist()
            for layer, h in residuals(model, ids).items()}

On a real model the picture is striking (nostalgebraist, 2020). Early layers predict tokens related to the current input; middle layers converge toward plausible continuations; the last few layers sharpen the final answer. For GPT-2 on “The Eiffel Tower is in the city of”, “ Paris“ typically rises to the top only in the later layers. Different models use their early layers differently, and the plain logit lens can be misleading for some; the tuned lens (Belrose et al., 2023) learns a small linear translator per layer to fix this.

Your 4-layer instruction-tuned GPT from Chapter 21, teacher-forced through “The contraction for”:

python run.py steer

{"text_end": " Response:\nThe contraction for", "layer_0": [" '", " \"", "'"], "layer_1": [" '", " \"", "'"],
 "layer_2": [" '", " \"", " the"], "layer_3": [" '", " \"", " the"]}

Every layer already expects a quote, the dataset’s style for this kind of answer (“The contraction for ‘it is’ is …”). A small model decides early. The last layer’s lens equals the real prediction exactly, which the milestone test checks.

Directions from contrasting examples

How do you find the direction for a concept? The simplest method that works surprisingly well is a difference of means. Collect activations at one layer for examples with the property (positive) and without it (negative), and take the normalized difference of their averages:

$$ d = \frac{\bar{h}\text{pos} - \bar{h}\text{neg}}{\lVert \bar{h}\text{pos} - \bar{h}\text{neg} \rVert}. $$

The pairs should differ only in the property, so everything else averages out: “Explain recursion in a cheerful conversational tone” versus “… in a formal technical tone”, for 32 topics (data/style-positive.jsonl and style-negative.jsonl).

Two edits use a direction:

• Steering (adding): $h \leftarrow h + \lambda d$ at one layer, at every position. Positive $\lambda$ pushes the model toward the property.
• Directional ablation (projecting out): $h \leftarrow h - (h \cdot d), d$. The component along $d$ becomes exactly zero, so downstream layers can no longer read the property from it.

def direction_from_means(positive, negative):
    """Unit vector from the mean of negative examples toward the mean of positive ones.  (Your engine: Chapter 23)"""
    delta = positive.float().mean(0) - negative.float().mean(0)
    norm = delta.norm()
    if not torch.isfinite(norm) or norm < 1e-8:
        raise ValueError("The two groups have (almost) the same mean: no usable direction")
    return delta / norm


def project_out(hidden, direction, strength=1.0):
    """h - strength * (h . d) d. With strength 1 the d-component becomes exactly zero.  (Your engine: Chapter 23)"""
    h = hidden.float()
    return (h - strength * (h @ direction.float()).unsqueeze(-1) * direction.float()).to(hidden.dtype)

@contextmanager
def steer(model, layer, direction, mode="add", strength=4.0):
    """Temporarily edit one block's output: 'add' pushes along the direction, 'ablate' removes it."""
    block = _blocks(model)[layer]

    def hook(_module, _inputs, output):
        hidden = output[0] if isinstance(output, tuple) else output
        if mode == "add":
            edited = hidden + strength * direction.to(hidden)
        elif mode == "ablate":
            edited = project_out(hidden, direction.to(hidden.device), strength)
        else:
            raise ValueError("mode must be 'add' or 'ablate'")
        return (edited, *output[1:]) if isinstance(output, tuple) else edited
    handle = block.register_forward_hook(hook)
    try:
        yield model
    finally:
        handle.remove()

Run it

run.py steer builds a direction at layer 2 of your Chapter 21 model from 300 training answers, written normally (negative) and in capitals (positive). Then it adds the direction, scaled by multiples of the gap between the two means, and measures the probability that the first response token is in capitals, over 60 test prompts:

{"layer": 2, "added": "0.0 x gap (0.0)",  "first_token_all_caps_probability": 0.2227}
{"layer": 2, "added": "0.5 x gap (5.4)",  "first_token_all_caps_probability": 0.2877}
{"layer": 2, "added": "1.0 x gap (10.8)", "first_token_all_caps_probability": 0.3699}
{"layer": 2, "added": "1.5 x gap (16.1)", "first_token_all_caps_probability": 0.4431}

One vector, no training, and the probability doubles. (The baseline is 22% because single capital-letter tokens like “A” count.) In a model this small, though, the direction is crude: push harder, or generate long texts, and fluency collapses before the style changes cleanly. Steering works far better in large models, whose representations are more linear and more robust. Turner et al. (2023) steered GPT-2 XL’s topic and sentiment with a single prompt-pair difference, and Arditi et al. (2024) found that refusal in 13 open chat models is mediated by one direction.

Making an edit permanent

A hook works at runtime. To ship an edit, change the weights. Every write into the residual stream comes from a matrix: the embedding, each attention output projection and each MLP down projection. If each writer’s output has no component along $d$, the stream never gains one. For a linear layer $y = Wx + b$, replace

$$ W \leftarrow (I - d d^\top) W, \qquad b \leftarrow (I - d d^\top), b , $$

so that $d^\top y = 0$ for every input:

@torch.no_grad()
def orthogonalize_output(linear, direction):
    """Make a linear layer unable to write along `direction`: W <- (I - d d^T) W, b <- (I - d d^T) b.
    Applied to every module that writes into the residual stream, this makes ablation permanent."""
    d = direction.to(linear.weight)
    linear.weight -= torch.outer(d, d @ linear.weight)
    if linear.bias is not None:
        linear.bias -= d * (d @ linear.bias)

Applied to all writers (and the embedding’s rows), this is weight orthogonalization, popularized as “abliteration” after Arditi et al.’s refusal paper. The milestone test checks that the edited layer equals the runtime projection exactly.

Warning

Directional ablation of a refusal direction removes a model’s safety behavior. It’s a well-documented interpretability result, and it’s why open-weight safety can’t rely on refusals alone. This book uses tone and capitalization as its examples; if you study refusal, do it on models and in settings where removing safeguards is appropriate, and don’t distribute such edits.

Editing a real Qwen3

edit_checkpoint.py runs the whole procedure on a Hugging Face Qwen3 checkpoint: it fits a direction at one block’s output from the last prompt token of the positive and negative prompts, ablates it with a hook on every position, and compares the edited and original models on held-out data:

python edit_checkpoint.py --model-dir models/Qwen3-0.6B --layer 14 \
    --positive-file data/style-positive.jsonl --negative-file data/style-negative.jsonl \
    --valid-file data/instruction-valid.jsonl --output runs/tone-direction.pt

It reports response loss before and after, the KL divergence from the original model on held-out responses, one greedy generation each way, and saves the direction with a manifest. No weights are changed; apply orthogonalize_output to make it permanent once you’re convinced.

Evaluate edits like any other change

• Did it do what you wanted? Measure the targeted behavior on held-out prompts, not the ones used to fit the direction.
• What else did it change? Report KL from the original on unrelated data, and held-out loss. A direction that changes everything is not a “tone” direction.
• Which layer? Directions work best in the middle layers; sweep a few and pick by held-out measurements, not by eye.
• Combine carefully. Ablation, LoRA (Chapter 22) and quantization (Chapter 20) don’t commute. Evaluate the artifact you ship.

Build it

pytest tests/test_ch23_steering.py
python run.py steer --impl engine

Stretch exercises

1. ★ Run run.py steer --layer L for every layer, and plot the all-caps probability at 1.0 × gap against the layer. Where: terminal: python run.py steer --impl engine --layer L; plot results in experiments/ch23.py (create it).
2. ★★ Load GPT-2 (Chapter 9) and print the logit lens for “The Eiffel Tower is in the city of”. At which layer does “ Paris“ enter the top 5? Where: experiments/ch23.py (create it), calling engine.loaders.load_gpt2 and engine.steering.logit_lens.
3. ★★ Steer Qwen3-0.6B’s tone with the style prompts: add $\lambda d$ at the best layer during generation and compare five answers at $\lambda = 0$, 4 and 8. Where: experiments/ch23.py (create it), using engine.steering.steer with engine.engine.LLM.
4. ★★★ Orthogonalize every residual writer of your Qwen3 engine against a direction, save the edited checkpoint, and verify its logits equal the hooked model’s. Where: extend orthogonalize_output in engine/steering.py for Qwen3’s residual writers; save with engine.safetensors_io.save_file.

Check your understanding

1. Why can the final norm and head be applied to the residual stream after any layer?
2. Why should positive and negative examples differ only in the property you want?
3. What’s the difference between adding a direction and projecting it out?
4. Why does orthogonalizing every residual writer make ablation permanent?
5. How would you show that an edit changed only what you intended?

Going deeper

• nostalgebraist, interpreting GPT: the logit lens (2020); Belrose et al., Eliciting Latent Predictions from Transformers with the Tuned Lens (2023).
• Elhage et al., A Mathematical Framework for Transformer Circuits (Anthropic, 2021): the residual-stream view this chapter uses.
• Turner et al., Activation Addition: Steering Language Models Without Optimization (2023); Zou et al., Representation Engineering (2023); Arditi et al., Refusal in Language Models Is Mediated by a Single Direction (2024).
• Sparse autoencoders, which find thousands of interpretable directions at once: Bricken et al., Towards Monosemanticity (2023) and Templeton et al., Scaling Monosemanticity (2024); Gemma Scope and the TransformerLens and nnsight libraries for doing this on real models.

24. Serving many users: continuous batching

One user leaves the GPU idle

Decode at batch 1 is memory-bound: each step reads every weight once to produce one token (Chapter 10). A 1,024 × 4,096 BF16 linear layer does 2 FLOPs per weight read, for an arithmetic intensity of about 1 FLOP per byte. An H100 can do around 300 FLOPs in the time it reads one byte. At batch 1, more than 99% of its arithmetic capability idles.

Now run 64 users’ decode steps in one forward pass. The weights are read once and multiplied by 64 vectors: 64 times the work for the same weight traffic. PMPP §20.6 works the numbers for that layer: arithmetic intensity grows from 1 FLOP/B at batch 1 to 315 FLOP/B at batch 512. Until the batch is large enough to become compute-bound, each extra user’s decode is nearly free.

There’s an important exception. Each user has their own KV cache, so attention reads grow linearly with the batch:

$$ \text{bytes per step} \approx \underbrace{\text{weight bytes}}{\text{shared}} + \sum{\text{requests}} \underbrace{\text{KV bytes}(\text{context}i)}{\text{per request}} . $$

For Qwen3-0.6B (1.19 GB of weights, 112 KiB of KV per token), 64 users at 2,000 tokens each read about 15 GB of cache per step, twelve times the weights. Batching helps until the KV cache dominates, and the KV cache also limits how many users fit in memory. Those two facts drive the next two chapters: Chapter 25 packs caches without waste, and Chapter 26 gets more tokens out of each cache read.

Static batching and its waste

The simplest batching collects $B$ requests, pads their prompts to the same length, and generates until all of them finish:

step:         1  2  3  4  5  6  7  8  9  10
request A:    ■  ■  ■  ✓  ·  ·  ·  ·  ·  ·      finished at step 4, slot idles
request B:    ■  ■  ■  ■  ■  ■  ■  ■  ■  ✓
request C:    ■  ■  ✓  ·  ·  ·  ·  ·  ·  ·      new request D waits for the whole batch

Outputs vary from a few tokens to thousands, so most slots idle most of the time, and newly arrived requests wait for the slowest member of the current batch (head-of-line blocking).

Continuous batching (Orca, Yu et al., 2022) makes the decision at every iteration instead: after each step, finished requests leave and waiting requests join. The batch is a set of independent requests that happen to share a forward pass.

Prefill and decode in one step

Joining requests need their prompts prefilled, which is a very different workload from decode (Chapter 16): hundreds of tokens at once, compute-bound. If a 4,000-token prompt is prefilled in one step, every running user’s next token waits for it, and their time per output token jumps.

The fix is a token budget per step and chunked prefill (Sarathi-Serve, Agrawal et al., 2024):

1. Admit waiting requests while cache slots are free.
2. Decode first: every request that has finished its prefill gets one token in this step.
3. Spend the rest of the budget on prefill chunks of at most prefill_chunk tokens, oldest request first.

A long prompt is then spread over several steps, and decode latency stays steady. The budget is the knob: a large budget finishes prefills sooner (lower time to first token) but makes each step slower (higher time per output token).

class Scheduler:
    """Decode-first, token-budgeted, chunked-prefill planner.  (Your engine: Chapter 24)

    Each step: admit waiting requests while slots are free; give every running request that
    has finished prefill one decode token; spend what is left of the budget on prompt chunks
    of at most prefill_chunk tokens, oldest request first.
    """

    def __init__(self, max_slots=8, token_budget=64, prefill_chunk=32):
        if min(max_slots, token_budget, prefill_chunk) < 1:
            raise ValueError("Scheduler limits must be positive")
        self.max_slots, self.token_budget, self.prefill_chunk = max_slots, token_budget, prefill_chunk
        self.free_slots = deque(range(max_slots))
        self.waiting, self.running = deque(), []

    def submit(self, request):
        if not request.prompt or request.max_new_tokens < 1:
            raise ValueError("A request needs a prompt and at least one new token")
        self.waiting.append(request)

    def plan(self):
        """(Your engine: Chapter 24)"""
        while self.waiting and self.free_slots:
            request = self.waiting.popleft()
            request.slot = self.free_slots.popleft()
            self.running.append(request)
        budget = self.token_budget
        decode = [r for r in self.running if r.prefilled == len(r.prompt)][:budget]
        budget -= len(decode)
        prefill = []
        for request in self.running:
            remaining = len(request.prompt) - request.prefilled
            if remaining and budget:
                length = min(remaining, self.prefill_chunk, budget)
                prefill.append((request, request.prefilled, length))
                budget -= length
        return Plan(prefill, decode)

    def release(self, request):
        """Return a finished or cancelled request's slot."""
        self.running.remove(request)
        self.free_slots.append(request.slot)

    @property
    def idle(self):
        return not self.waiting and not self.running

The scheduler is pure policy: it moves no tensors. Keeping it separate from execution makes it testable with plain Python lists, and it’s where production engines differ most (priorities, preemption, fairness, prefix-aware ordering).

The engine loop

class ContinuousBatchingEngine:
    """Runs Scheduler plans on a model that follows the book's cache protocol.  (Your engine: Chapter 24)"""

    def __init__(self, model, max_slots=8, capacity=512, token_budget=64, prefill_chunk=32):
        self.model = model.eval()
        p = next(model.parameters())
        layers, kv_heads, head_dim = model.cache_spec()
        self.cache = StaticKVCache(layers, max_slots, kv_heads, capacity, head_dim, p.device, p.dtype)
        self.capacity, self.device = capacity, p.device
        self.scheduler = Scheduler(max_slots, token_budget, prefill_chunk)
        self.generators = {}
        self.steps = 0

    def submit(self, request):
        if len(request.prompt) + request.max_new_tokens > self.capacity:
            raise ValueError(f"{request.rid}: prompt plus output exceed slot capacity")
        request.arrival = time.perf_counter()
        self.generators[request.rid] = torch.Generator(device=self.device).manual_seed(request.seed)
        self.scheduler.submit(request)

    def _emit(self, request, logits_row):
        token = int(sample(logits_row[None], request.temperature, generator=self.generators[request.rid]))
        if request.first_token_time is None:
            request.first_token_time = time.perf_counter()
        request.output.append(token)
        if token in request.stop_ids:
            request.finish_reason = "stop"
        elif len(request.output) >= request.max_new_tokens:
            request.finish_reason = "length"

    @torch.inference_mode()
    def step(self):
        """One engine iteration. Returns the requests that finished during it.  (Your engine: Chapter 24)"""
        plan = self.scheduler.plan()
        for request, start, length in plan.prefill:            # prompt chunks, one request at a time
            ids = torch.tensor([request.prompt[start:start + length]], device=self.device)
            positions = torch.arange(start, start + length, device=self.device)[None]
            rows = torch.tensor([request.slot], device=self.device)
            logits = self.model(ids, self.cache, positions=positions, rows=rows)
            request.prefilled += length
            if request.prefilled == len(request.prompt):        # the last prompt position predicts token 1
                self._emit(request, logits[0, -1])
        if plan.decode:                                         # every decoding request, one batched forward
            ids = torch.tensor([[r.output[-1]] for r in plan.decode], device=self.device)
            positions = torch.tensor([[len(r.prompt) + len(r.output) - 1] for r in plan.decode], device=self.device)
            rows = torch.tensor([r.slot for r in plan.decode], device=self.device)
            logits = self.model(ids, self.cache, positions=positions, rows=rows)
            for i, request in enumerate(plan.decode):
                self._emit(request, logits[i, -1])
        finished = [r for r in self.scheduler.running if r.done]
        for request in finished:
            request.finish_time = time.perf_counter()
            self.scheduler.release(request)
            self.generators.pop(request.rid, None)
        self.steps += 1
        return finished

    def run(self, requests):
        """Submit everything, step until idle, return {rid: request}."""
        for request in requests:
            self.submit(request)
        results = {}
        while not self.scheduler.idle:
            for request in self.step():
                results[request.rid] = request
        return results

Each request owns one row of a StaticKVCache (Chapter 19), its slot. Rows hold different lengths; writes go to (row, position) pairs and the causal rule hides everything a row hasn’t written. That’s why one batched forward can decode requests at positions 7, 58 and 1,203 at once: the model receives positions of shape [B, 1] and rows of shape [B], and the cache protocol does the rest. No padding of prompts, no attention mask beyond key_pos <= query_pos.

Per-request state lives on the Request: prompt, output, how much has been prefilled, stop tokens, finish reason, and its own random generator, seeded per request so that a request’s samples don’t depend on who else shares the batch.

The correctness test

Batching is an optimization, so it must not change any request’s output. The milestone test runs four requests with different prompt and output lengths through 3 slots, with chunked prefill and a small budget, and checks that each one’s greedy output equals the output of generating it alone. This catches wrong positions, wrong rows, a stale cache slot from a previous request and off-by-one errors in the prefill-to-decode handoff. (With sampling, outputs match only if each request’s random stream is independent of the batch, which is why generators are per request.)

Note

Even with correct code, batched and solo outputs can differ in the last bits in BF16 on a GPU: a matmul over a batch of 8 may use a different kernel or reduction order than a batch of 1. Greedy tokens can then diverge at a near-tie. Production engines that need batch-invariant results use kernels designed for it (Thinking Machines, Defeating Nondeterminism in LLM Inference, 2025).

Run it

python run.py batch --requests 12 --slots 8

Twelve requests with prompts of 8-64 tokens and outputs of 8-48 tokens, through the 2-layer test model on a laptop CPU, first one at a time, then continuously batched:

{"requests": 12, "output_tokens": 322, "serial_tok_s": 398.4, "batched_tok_s": 1037.6, "engine_steps": 55}
{"rid": "5", "ttft_ms": 25.06, "tpot_ms": 8.15, "tokens": 9, "finish": "length"}
{"rid": "6", "ttft_ms": 29.77, "tpot_ms": 7.69, "tokens": 10, "finish": "length"}
{"rid": "2", "ttft_ms": 9.11, "tpot_ms": 7.53, "tokens": 21, "finish": "length"}

2.6× the throughput, in 55 engine steps instead of 322 sequential ones. With --slots 1 the engine does no batching and reached 549 tokens/s against the serial loop’s 640 on the same run: the scheduler’s bookkeeping costs a little, and batching is where the win comes from. On a GPU, where a batch of 8 costs about the same as a batch of 1, the gain approaches the slot count.

Serving metrics

Report percentiles (p50, p90, p99), not averages: a scheduler that’s great on average and terrible for 1% of users is a bad scheduler. And benchmark with a realistic arrival process (requests arriving over time, for example Poisson), not everything submitted at once, which hides queueing.

Production engines add more scheduling tools: preemption (pause a request and free its cache when memory runs out, then recompute or swap it back in), priorities, prefix-aware ordering (Chapter 25), and disaggregation, running prefill and decode on separate GPUs so the two workloads stop interfering (DistServe, Splitwise, 2024).

Batched prefill in production

This engine prefills one request’s chunk per forward call, which is simple and correct but launches many small forward passes. Production engines flatten all of a step’s tokens, prefill chunks and decode tokens of every request, into one long sequence of $N$ tokens, with per-token positions and request IDs. The linear layers see one [N, d] matrix, ideal for tensor cores. Attention uses a “varlen” kernel that receives the boundaries between requests (cu_seqlens in FlashAttention) and each request’s cache location. Your model’s positions-based interface already supports this; what’s missing is an attention kernel that takes ragged batches (Chapter 25’s paged decode kernel is a step in that direction).

Build it

pytest tests/test_ch24_batching.py
python run.py batch --impl engine

Stretch exercises

1. ★ Sweep --slots from 1 to 32 on a GPU with Qwen3-0.6B, and plot throughput and mean TPOT. Where does throughput stop growing? Where: create experiments/ch24.py from run.py’s cmd_batch, replacing its tiny model with engine.loaders.load_qwen3 and collecting per-request timings. Sweep max_slots in that script; run.py batch --slots N is the tiny-model baseline.
2. ★★ Add Poisson arrivals: submit requests at random times while the engine runs, and report p50 and p99 TTFT for token budgets 64, 256 and 1,024. Where: experiments/ch24.py (create it), adapting run.py’s cmd_batch and submitting requests to engine.scheduler.ContinuousBatchingEngine.
3. ★★ Serve your engine over HTTP with an OpenAI-compatible /v1/completions endpoint (FastAPI and server-sent events), running the engine loop in a background thread and streaming each request’s tokens as they’re emitted. Where: new experiments/batching_server.py, wrapping engine.scheduler.ContinuousBatchingEngine; Chapter 36 supplies the later unified server.
4. ★★★ Implement preemption: when a new high-priority request arrives and no slot is free, evict the request with the most remaining work, and recompute its cache from prompt + output when it’s readmitted. Verify outputs are still identical to solo runs. Where: Scheduler and ContinuousBatchingEngine in engine/scheduler.py.

Check your understanding

1. Why does batching speed up the weight reads of decode but not its KV-cache reads?
2. What does continuous batching fix that static batching can’t?
3. Why does the scheduler give decode tokens priority over prefill chunks?
4. How does one batched forward serve requests at different positions without padding?
5. Why should each request have its own random generator?

Going deeper

• PMPP §20.6 (KV cache arithmetic intensity and memory requirement, pp. 504-508): batching’s effect on intensity and the batch-size versus context-length trade-off.
• Yu et al., Orca: A Distributed Serving System for Transformer-Based Generative Models (OSDI 2022); Agrawal et al., Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve (OSDI 2024); Zhong et al., DistServe (2024).
• GPU Mode L35 (SGLang performance optimization), and the vLLM and SGLang schedulers (vllm/v1/core/sched/scheduler.py, sglang/srt/managers/scheduler.py), which follow the plan-then-execute structure used here.
• Anyscale’s blog post How continuous batching enables 23x throughput in LLM inference for measurements on real GPUs.

25. Paged attention and prefix caching

Reserved memory is wasted memory

Chapter 24’s engine gives each request a slot of fixed capacity, say 1,024 tokens. A request that uses 50 tokens still reserves 1,024. With requests of varied lengths, most of the reserved cache is empty:

python run.py paged

{"requests": 32, "tokens_stored": 9924, "static_slots_reserved": 32768, "paged_slots_reserved": 10144,
 "static_utilization": 0.303, "paged_utilization": 0.978}

Thirty-two live requests of 20-500 tokens fill 30% of their reserved slots. The vLLM paper (Kwon et al., 2023) measured existing systems using only 20-38% of their KV memory for actual tokens. Since the cache limits how many requests fit (Chapter 24), that waste directly cuts throughput.

Allocating exactly the right amount isn’t possible either: you don’t know how long an answer will be until it’s done. And growing a contiguous buffer means copying it, or fragmenting memory into holes that no request fits.

Paging

Operating systems solved the same problem for process memory decades ago. PagedAttention applies the solution to the KV cache:

• Split every sequence’s cache into fixed-size blocks of $b$ tokens (16 is typical).
• Keep all blocks of all sequences in one pool: per layer, tensors of shape [num_blocks, kv_heads, b, head_dim].
• Give each sequence a block table mapping its logical blocks to physical blocks in the pool, which need not be contiguous or in order.

Logical position $p$ lives in physical block table[p // b] at offset p % b. For $b = 4$ and table [5, 1, 7], position 9 is in logical block 2, physical block 7, offset 1.

print(table[9 // 4], 9 % 4)      # (7, 1)

inline std::pair<size_t, size_t> physical_slot(const std::vector<size_t>& table, size_t pos, size_t block) {
    return {table[pos / block], pos % block};
}

#![allow(unused)]
fn main() {
pub fn physical_slot(table: &[usize], position: usize, block_size: usize) -> (usize, usize) {
    (table[position / block_size], position % block_size)
}
}

A sequence allocates a new block only when it fills its last one, so it wastes at most $b - 1$ slots, under 4% in the demo above. Blocks come from a free list and go back when the request finishes.

The allocator

class BlockAllocator:
    """Free list plus reference counts.  (Your engine: Chapter 25)"""

    def __init__(self, num_blocks):
        self.free = deque(range(num_blocks))
        self.refs = [0] * num_blocks

    def allocate(self):
        """(Your engine: Chapter 25)"""
        if not self.free:
            raise MemoryError("KV block pool exhausted")
        block = self.free.popleft()
        self.refs[block] = 1
        return block

    def share(self, block):
        self.refs[block] += 1

    def release(self, block):
        """(Your engine: Chapter 25)"""
        self.refs[block] -= 1
        if self.refs[block] == 0:
            self.free.append(block)
        elif self.refs[block] < 0:
            raise RuntimeError(f"Block {block} released more times than it was referenced")

    @property
    def num_free(self):
        return len(self.free)

Every block carries a reference count: the number of block tables (and caches, below) that point to it. release decrements it, and only a count of zero returns the block to the free list.

The paged cache

class PagedKVCache:
    """Implements the book's cache protocol over a block pool.  (Your engine: Chapter 25)

    Before a forward pass the caller names the sequences in the batch (set_batch) and
    reserves room for the new tokens (reserve); allocation failures therefore happen
    before any GPU work is launched.
    """

    def __init__(self, layers, num_blocks, block_size, kv_heads, head_dim, device="cpu", dtype=torch.float32):
        shape = (num_blocks, kv_heads, block_size, head_dim)
        self.k_pool = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(layers)]
        self.v_pool = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(layers)]
        self.block_size, self.device = block_size, device
        self.allocator = BlockAllocator(num_blocks)
        self.tables, self.lengths = {}, {}
        self.batch = []

    def blocks_needed(self, tokens):
        return -(-tokens // self.block_size)

    def create(self, seq):
        if seq in self.tables:
            raise ValueError(f"Sequence {seq} already exists")
        self.tables[seq], self.lengths[seq] = [], 0

    def reserve(self, seq, new_tokens):
        """Make blocks exist for positions [length, length + new_tokens), copying any shared
        block that is about to be written (copy-on-write).  (Your engine: Chapter 25)"""
        table, start = self.tables[seq], self.lengths[seq]
        end = start + new_tokens
        first_written = start // self.block_size
        for index in range(first_written, min(len(table), self.blocks_needed(end))):
            block = table[index]
            if self.allocator.refs[block] > 1:                 # shared: copy before writing
                fresh = self.allocator.allocate()
                for pool in self.k_pool + self.v_pool:
                    pool[fresh].copy_(pool[block])
                self.allocator.release(block)
                table[index] = fresh
        while len(table) < self.blocks_needed(end):
            table.append(self.allocator.allocate())

    def set_batch(self, seqs):
        self.batch = list(seqs)
        width = max(len(self.tables[s]) for s in self.batch)
        rows = [self.tables[s] + [0] * (width - len(self.tables[s])) for s in self.batch]
        self.table_tensor = torch.tensor(rows, device=self.device, dtype=torch.long)

    def update(self, layer, k, v, positions, rows=None):
        """Scatter new K/V into their blocks, then gather each sequence's history (reference path).  (Your engine: Chapter 25)"""
        if positions.ndim == 1:
            positions = positions.expand(k.shape[0], -1)
        batch_index = torch.arange(k.shape[0], device=k.device)[:, None]
        blocks = self.table_tensor[batch_index, positions // self.block_size]          # [B, T]
        offsets = positions % self.block_size
        self.k_pool[layer][blocks, :, offsets] = k.transpose(1, 2)
        self.v_pool[layer][blocks, :, offsets] = v.transpose(1, 2)
        if layer == len(self.k_pool) - 1:                       # the last layer commits the lengths
            for b, seq in enumerate(self.batch):
                self.lengths[seq] = max(self.lengths[seq], int(positions[b].max()) + 1)
        span = int(positions.max()) + 1
        logical = torch.arange(span, device=k.device)
        gather_blocks = self.table_tensor[:, logical // self.block_size]               # [B, span]
        keys = self.k_pool[layer][gather_blocks, :, logical % self.block_size]         # [B, span, H, D]
        values = self.v_pool[layer][gather_blocks, :, logical % self.block_size]
        return keys.transpose(1, 2), values.transpose(1, 2), logical

    def fork(self, parent, child):
        """child starts as an alias of parent's blocks (beam search, n>1 sampling, shared prompts).  (Your engine: Chapter 25)"""
        self.create(child)
        self.tables[child] = list(self.tables[parent])
        self.lengths[child] = self.lengths[parent]
        for block in self.tables[child]:
            self.allocator.share(block)

    def free(self, seq):
        for block in self.tables.pop(seq):
            self.allocator.release(block)
        self.lengths.pop(seq)

    def truncate(self, seq, length):
        """Roll back to `length` tokens, returning now-unused blocks."""
        keep = self.blocks_needed(length)
        for block in self.tables[seq][keep:]:
            self.allocator.release(block)
        del self.tables[seq][keep:]
        self.lengths[seq] = length

The cache implements the same update protocol as every cache since Chapter 16, so your Qwen3 runs on it unchanged. The flow per step:

1. reserve(seq, n) before the forward pass makes sure blocks exist for the $n$ new positions. Allocation failures happen here, before any GPU work is launched, where the scheduler can handle them (by waiting, or preempting another request).
2. set_batch(seqs) builds the block-table tensor for this batch, padded to the longest table.
3. update scatters the new keys and values to (table[p // b], p % b) and returns each sequence’s history. This reference version gathers the history into a contiguous tensor for the existing attention code, which costs a copy; the kernel at the end of the chapter removes it.

Sharing blocks: forks and copy-on-write

Many requests share content. Parallel sampling (n = 4 answers to one prompt) and beam search need several sequences with the same prompt. With block tables, a fork copies the table, not the data, and increments each block’s reference count.

The sequences can then diverge. When one of them is about to write into a block whose count is above 1, reserve first copies that block to a fresh one and points this sequence’s table at the copy: copy-on-write. Full blocks of the prompt stay shared forever; only the partially filled last block gets copied.

{"parent_blocks": 3, "new_blocks_after_4_forks": 0, "new_blocks_after_each_sample_writes_a_token": 4}

A 40-token prompt takes 3 blocks of 16 (the last one half full). Four forks cost no new blocks. When each sample writes its first token, each copies the shared partial block: 4 new blocks in total, instead of 12 for four independent copies of the prompt.

The milestone test checks the subtle part: after a fork, the parent and the child each continue with a different token, and both must produce exactly the logits of running their full sequence from scratch. If copy-on-write is missing, one sequence’s write corrupts the other’s history.

Prefix caching

Most production traffic shares long prefixes: the same system prompt, the same few-shot examples, the same document with different questions, the growing history of a multi-turn chat. Their KV cache is identical, so it can be computed once.

class PrefixCache:
    """Reuse full blocks whose token contents (and everything before them) match exactly.

    Keys are the whole token prefix up to the end of a block, so two blocks with the same
    text but different histories never collide. The cache holds its own reference on each
    block; least-recently-used entries are evicted when the pool runs low.
    """

    def __init__(self, cache):
        self.cache = cache
        self.entries = OrderedDict()    # tuple(prefix tokens) -> block

    def attach(self, seq, prompt):
        """Create seq with the longest cached prefix of prompt. Returns tokens reused."""
        self.cache.create(seq)
        size, reused = self.cache.block_size, 0
        # Keep at least one prompt token to prefill: its logits predict the first output.
        while (reused + size) <= len(prompt) - 1:
            key = tuple(prompt[:reused + size])
            block = self.entries.get(key)
            if block is None:
                break
            self.entries.move_to_end(key)
            self.cache.allocator.share(block)
            self.cache.tables[seq].append(block)
            reused += size
        self.cache.lengths[seq] = reused
        return reused

    def remember(self, seq, tokens):
        """After prefill, publish seq's full blocks for later requests."""
        size = self.cache.block_size
        for index in range(len(tokens) // size):
            key = tuple(tokens[:(index + 1) * size])
            if key not in self.entries:
                block = self.cache.tables[seq][index]
                self.cache.allocator.share(block)
                self.entries[key] = block

    def evict(self, blocks_wanted):
        """Drop LRU entries that nobody else is using until enough blocks are free."""
        for key in list(self.entries):
            if self.cache.allocator.num_free >= blocks_wanted:
                break
            block = self.entries[key]
            if self.cache.allocator.refs[block] == 1:
                del self.entries[key]
                self.cache.allocator.release(block)

The key for a cached block is the entire token prefix up to the end of that block, not just the block’s own tokens. A block’s keys and values depend on every earlier token, so two blocks with identical text but different histories must never match. (Production engines hash the chain of block contents instead of storing the tuples, and the key must also include anything else that changes activations: the model revision, LoRA adapter, and multimodal inputs.)

attach walks the prompt one full block at a time while the prefix is cached, shares those blocks, and leaves at least one token to prefill, because the prompt’s last position must run to produce the first output’s logits. remember publishes a sequence’s full blocks after prefill. The cache holds its own reference to each block, so cached blocks survive their requests, and evict drops least-recently-used entries nobody else is using when the pool runs low.

Eight requests sharing a 96-token system prompt, each with 8 tokens of its own:

{"prompt_tokens": 832, "tokens_prefilled": 160, "blocks_in_use": 14, "blocks_without_sharing": 56}

The first request prefills all 104 tokens; the other seven prefill only their own 8. That’s 81% less prefill compute, which lowers time to first token directly, and a quarter of the memory. SGLang’s RadixAttention generalizes this to a radix tree over all cached prefixes and schedules requests to maximize hits.

Attention that reads through the table

The gather in update builds a contiguous copy of every sequence’s history at every layer, every step. A paged attention kernel reads K and V directly from the pool instead. Here’s a decode kernel: one program per (sequence, query head), walking the sequence’s block table and applying Chapter 15’s online softmax one block at a time:

@triton.jit
def paged_decode_kernel(q_ptr, k_pool, v_pool, tables_ptr, lengths_ptr, out_ptr,
                        max_blocks, heads, group, scale,
                        s_pool_block, s_pool_head, s_pool_slot,
                        D: tl.constexpr, BLOCK: tl.constexpr):
    """(Your engine: Chapter 25)"""
    seq, h = tl.program_id(0), tl.program_id(1)
    kvh = h // group
    rd = tl.arange(0, D)
    slots = tl.arange(0, BLOCK)
    q = tl.load(q_ptr + (seq * heads + h) * D + rd).to(tl.float32)
    length = tl.load(lengths_ptr + seq)
    m = tl.full((), -float("inf"), tl.float32)
    l = tl.zeros((), tl.float32)
    acc = tl.zeros((D,), tl.float32)
    for i in range(0, tl.cdiv(length, BLOCK)):
        block = tl.load(tables_ptr + seq * max_blocks + i)            # logical block i -> physical block
        base = block * s_pool_block + kvh * s_pool_head
        valid = i * BLOCK + slots < length
        k = tl.load(k_pool + base + slots[:, None] * s_pool_slot + rd[None, :], mask=valid[:, None], other=0.0)
        v = tl.load(v_pool + base + slots[:, None] * s_pool_slot + rd[None, :], mask=valid[:, None], other=0.0)
        s = tl.sum(k.to(tl.float32) * q[None, :], axis=1) * scale
        s = tl.where(valid, s, -float("inf"))
        m_new = tl.maximum(m, tl.max(s, axis=0))
        alpha = tl.exp(m - m_new)
        p = tl.exp(s - m_new)
        l = l * alpha + tl.sum(p, axis=0)
        acc = acc * alpha + tl.sum(p[:, None] * v.to(tl.float32), axis=0)
        m = m_new
    tl.store(out_ptr + (seq * heads + h) * D + rd, (acc / l).to(out_ptr.dtype.element_ty))


def paged_decode_attention(q, k_pool, v_pool, block_tables, lengths, block_size):
    """q [B, Hq, D]; pools [N, Hkv, block_size, D]; block_tables int [B, max_blocks]; lengths int [B] (>= 1)."""
    check_device(q, k_pool, v_pool, block_tables, lengths)
    batch, heads, d = q.shape
    out = torch.empty_like(q)
    k_pool, v_pool = k_pool.contiguous(), v_pool.contiguous()
    paged_decode_kernel[(batch, heads)](q.contiguous(), k_pool, v_pool, block_tables.contiguous(),
                                        lengths.contiguous(), out, block_tables.shape[1], heads,
                                        heads // k_pool.shape[1], 1.0 / math.sqrt(d),
                                        k_pool.stride(0), k_pool.stride(1), k_pool.stride(2),
                                        D=d, BLOCK=block_size)
    return out

Compared with the flash kernel, the only new thing is the indirection: tl.load(tables_ptr + seq * max_blocks + i) turns logical block i into a physical block number, and the K/V addresses come from that. The lengths tensor masks the tail of the last block. Grouped-query attention works the same way as before: query head h reads KV head h // group.

The reference paged_decode_attention in paged.py gathers and computes the same result; the test checks the kernel against it with scrambled block tables.

What production paged kernels add: splitting long sequences across programs with a final merge (flash-decoding, Chapter 15), processing all query heads of a KV group together so K/V is loaded once, and FP8 KV caches. FlashInfer and vLLM’s attention backends implement these for prefill, decode and mixed batches.

Choosing the block size

16 is vLLM’s default. Some FlashAttention-3-based paged kernels prefer larger blocks.

Build it

pytest tests/test_ch25_paged.py
python run.py paged --impl engine

Stretch exercises

1. ★ Plot paged_utilization from run.py paged against block sizes 1-128 for the same request lengths. Where: experiments/ch25.py (create it), adapting run.py’s cmd_paged with different block sizes.
2. ★★ Replace StaticKVCache in Chapter 24’s engine with PagedKVCache and a PrefixCache, calling reserve for every planned chunk. When reserve raises MemoryError, try evict, then preempt the newest request. Verify outputs still match solo runs. Where: ContinuousBatchingEngine.__init__ / step in engine/scheduler.py, using engine.paged.
3. ★★ Implement $n$-way parallel sampling in the batching engine with fork: one prefill, then $n$ independently seeded samples. Where: request/slot handling in engine/scheduler.py and PagedKVCache.fork in engine/paged.py.
4. ★★★ Write a paged prefill kernel: a block of queries from one sequence attends causally to that sequence’s paged history plus the new tokens. Where: add a prefill kernel and launcher beside paged_decode_kernel in engine/kernels/triton_paged.py.

Check your understanding

1. Why can’t an engine simply allocate exactly the cache each request will need?
2. Where does the logical position $p$ of a sequence live in the pool?
3. Why is the key of a cached prefix block the whole prefix, not just the block’s tokens?
4. When does copy-on-write copy a block, and which block is it usually?
5. Why does attach always leave at least one prompt token to prefill?

Going deeper

• Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention (SOSP 2023), the vLLM paper; Zheng et al., SGLang: Efficient Execution of Structured Language Model Programs (2024) for RadixAttention.
• GPU Mode L40 (FlashInfer): paged and ragged attention kernels in production; L35 (SGLang).
• PMPP §20.7 (Alleviating the memory requirements of the attention mechanism).
• vLLM’s vllm/v1/core/block_pool.py and kv_cache_manager.py, and its paged attention kernels, which follow the structure of this chapter.

26. Speculative decoding

Verification is almost free

At batch 1, a decode step reads every weight to produce one token’s logits (Chapter 10). Feeding the model five tokens instead of one reads the same weights and returns five positions’ logits: the step is still memory-bound and takes nearly the same time. Chapter 24 used this by batching different users. Speculative decoding uses it along the sequence of one user.

The catch: to feed five tokens, you need to know them, and you only know the next token after computing it. Unless you guess. If a cheap draft model guesses the next $k$ tokens, the expensive target model can check all of them in one forward pass. Where its own predictions agree, those tokens are done. At the first disagreement, the target’s own prediction replaces the guess, so every round produces at least one token, and up to $k+1$.

Greedy speculation

With greedy decoding, the rule is simple. Feed the target [last, d1, d2, d3, d4]. Its logits at each position give its own choice for the next token:

target input:       last    d1     d2     d3     d4
target's argmax:     t1     t2     t3     t4     t5
draft proposed:      d1     d2     d3     d4

Accept $d_1$ if $t_1 = d_1$, then $d_2$ if $t_2 = d_2$, and so on. At the first mismatch, take the target’s token $t_i$ instead and stop. If all four match, $t_5$ is a free bonus token. Every committed token is exactly what greedy decoding of the target alone would have produced, so the output is identical; only the number of target calls changes.

Sampling: rejection sampling keeps the distribution exact

With temperature sampling, “agree with the argmax” is too strict and would also change the output distribution. Leviathan et al. (2023) and Chen et al. (2023) showed how to accept draft tokens so that every committed token is distributed exactly as if sampled from the target.

Let $q$ be the draft’s distribution at a position, $p$ the target’s, and $x \sim q$ the draft’s token:

1. Accept $x$ with probability $\min!\big(1, p(x)/q(x)\big)$.
2. Otherwise sample a replacement from the residual distribution $r(y) = \dfrac{\max(0,\ p(y) - q(y))}{\sum_z \max(0,\ p(z) - q(z))}$, and stop the round.

Why this produces exactly $p$: the probability of outputting $y$ through acceptance is $q(y) \cdot \min(1, p(y)/q(y)) = \min(q(y), p(y))$. The probability of rejecting is $1 - \sum_z \min(q(z), p(z)) = \sum_z \max(0, p(z) - q(z))$, which is exactly the residual’s normalizer. So rejection contributes $\max(0, p(y) - q(y))$. Adding the two:

$$ \min(p(y), q(y)) + \max(0,\ p(y) - q(y)) = p(y). \quad\checkmark $$

Worked example. Vocabulary of three tokens, $p = [0.6, 0.3, 0.1]$, $q = [0.2, 0.3, 0.5]$. The draft proposes token 2 half the time, but the target accepts it with probability $0.1/0.5 = 0.2$. Token 0 is accepted always ($0.6/0.2 > 1$). The residual is $\max(0, p - q) = [0.4, 0, 0]$: every rejection becomes token 0. Total probability of token 0: $0.2$ (proposed and accepted) $+ (0.5 \times 0.8)$ (token 2 rejected) $= 0.6$. The milestone test draws 20,000 samples and checks all three frequencies.

def accept_or_correct(p, q, token, generator=None):
    """One verification step for draft token `token` drawn from q.  (Your engine: Chapter 26)

    Accept with probability min(1, p[token] / q[token]). On rejection, draw from the
    residual max(p - q, 0) / sum(max(p - q, 0)). Returns (token, accepted).
    The two cases together produce exactly p: accepted mass min(p, q) plus residual mass
    max(p - q, 0) add up to p at every token.
    """
    ratio = (p[token] / q[token].clamp_min(1e-20)).clamp(max=1.0)
    if torch.rand((), generator=generator, device=p.device) < ratio:
        return token, True
    residual = (p - q).clamp_min(0)
    if residual.sum() <= 0:          # p == q numerically: nothing to correct
        residual = p
    return int(torch.multinomial(residual / residual.sum(), 1, generator=generator)), False

inline std::pair<size_t, bool> accept_or_correct(const std::vector<double>& p, const std::vector<double>& q, size_t x, Rng& rng) {
    if (rng.uniform() < std::min(1.0, p[x] / q[x])) return {x, true};
    std::vector<double> residual(p.size());
    double total = 0;
    for (size_t i = 0; i < p.size(); ++i) total += (residual[i] = std::max(0.0, p[i] - q[i]));
    double r = rng.uniform() * total;
    for (size_t i = 0; i < p.size(); ++i) { if (r < residual[i]) return {i, false}; r -= residual[i]; }
    return {x, false};
}

#![allow(unused)]
fn main() {
/// Draft token `x` came from q. Accept with probability min(1, p[x]/q[x]); otherwise draw
/// from the leftover distribution max(p - q, 0), renormalized. The result is distributed as p.
pub fn accept_or_correct(p: &[f64], q: &[f64], x: usize, rng: &mut Rng) -> (usize, bool) {
    if (rng.uniform() as f64) < (p[x] / q[x]).min(1.0) {
        return (x, true);
    }
    let residual: Vec<f64> = p.iter().zip(q).map(|(a, b)| (a - b).max(0.0)).collect();
    let total: f64 = residual.iter().sum();
    let mut r = rng.uniform() as f64 * total;
    for (i, mass) in residual.iter().enumerate() {
        if r < *mass {
            return (i, false);
        }
        r -= mass;
    }
    (residual.iter().rposition(|m| *m > 0.0).unwrap_or(x), false)
}
}

How much faster?

Suppose each draft token is accepted independently with probability $\alpha$ (a simplification, but a useful one). A round accepts $i$ tokens with probability $\alpha^i(1-\alpha)$ and always adds one target token, so the expected tokens per round is

$$ E[\text{tokens}] = 1 + \alpha + \alpha^2 + \dots + \alpha^k = \frac{1 - \alpha^{k+1}}{1 - \alpha}. $$

A round costs one target forward plus $k$ draft forwards. If a draft step costs a fraction $c$ of a target step, the speedup over plain decoding is

$$ \text{speedup} = \frac{1 - \alpha^{k+1}}{(1 - \alpha)(1 + k c)} . $$

With $\alpha = 0.8$, $k = 4$ and $c = 0.05$ (a draft about 20× smaller), that’s $3.36 / 1.2 = 2.8\times$. With $\alpha = 0.5$, it’s $1.94 / 1.2 = 1.6\times$. Larger $k$ helps only while $\alpha$ is high: the chance of reaching the $k$-th guess shrinks geometrically, while its cost doesn’t.

def expected_tokens_per_round(alpha, k):
    """With per-token acceptance probability alpha, a round yields (1 - alpha^(k+1)) / (1 - alpha) tokens."""
    if alpha >= 1:
        return k + 1
    return (1 - alpha ** (k + 1)) / (1 - alpha)

The algorithm, with cache rollback

@torch.inference_mode()
def speculative_generate(target, draft, prompt_ids, max_new_tokens, k=4, temperature=0.0,
                         stop_ids=(), generator=None):
    """Generate with a draft model. Both models must share the tokenizer.  (Your engine: Chapter 26)

    Invariant at the top of each round: both caches hold every committed token except the
    newest one, `last`, which has been chosen but not yet fed to either model.
    Returns (new_token_ids, stats).
    """
    device = next(target.parameters()).device
    ids = torch.as_tensor(prompt_ids, device=device).view(1, -1)
    capacity = ids.shape[1] + max_new_tokens + k + 1
    t_cache, d_cache = target.new_cache(1, capacity), draft.new_cache(1, capacity)
    target(ids[:, :-1], t_cache) if ids.shape[1] > 1 else None
    draft(ids[:, :-1], d_cache) if ids.shape[1] > 1 else None
    last = ids[:, -1:]
    output, rounds, accepted_total, proposed_total = [], 0, 0, 0
    while len(output) < max_new_tokens:
        rounds += 1
        base = t_cache.length
        # 1. Draft proposes k tokens, one cheap forward each.
        proposals, q_dists, token = [], [], last
        for _ in range(k):
            logits = draft(token, d_cache)[:, -1]
            if temperature == 0:
                token = logits.argmax(-1, keepdim=True)
            else:
                q = _probs(logits, temperature)[0]
                q_dists.append(q)
                token = torch.multinomial(q, 1, generator=generator)[None]
            proposals.append(token)
        # 2. Target scores last + all proposals in ONE forward: k+1 next-token distributions.
        block = torch.cat([last] + proposals, dim=1)
        target_logits = target(block, t_cache)[0]                       # [k+1, V]
        # 3. Accept a prefix of the proposals, then add one corrected or bonus token.
        new_tokens, all_accepted = [], True
        for i, proposal in enumerate(proposals):
            token_id = int(proposal)
            if temperature == 0:
                best = int(target_logits[i].argmax())
                ok, chosen = best == token_id, best
            else:
                chosen, ok = accept_or_correct(_probs(target_logits[i], temperature), q_dists[i], token_id, generator)
            new_tokens.append(chosen)
            if not ok:
                all_accepted = False
                break
        accepted = len(new_tokens) if all_accepted else len(new_tokens) - 1
        if all_accepted:                                                # all accepted: free bonus token
            bonus = target_logits[k]
            chosen = int(bonus.argmax()) if temperature == 0 else int(
                torch.multinomial(_probs(bonus, temperature), 1, generator=generator))
            new_tokens.append(chosen)
        accepted_total += accepted
        proposed_total += k
        # 4. Roll both caches back to "everything committed except the newest token".
        t_cache.truncate(base + accepted + 1)
        if accepted == k:
            draft(proposals[-1], d_cache)          # the draft never fed its own last proposal
        d_cache.truncate(base + accepted + 1)
        for token_id in new_tokens:
            output.append(token_id)
            if token_id in stop_ids or len(output) == max_new_tokens:
                return output, {"rounds": rounds, "acceptance_rate": accepted_total / proposed_total,
                                "tokens_per_target_call": len(output) / rounds}
        last = torch.tensor([[new_tokens[-1]]], device=device)
    return output, {"rounds": rounds, "acceptance_rate": accepted_total / max(proposed_total, 1),
                    "tokens_per_target_call": len(output) / max(rounds, 1)}

The bookkeeping that makes it correct is one invariant: at the top of each round, both caches hold every committed token except the newest one, last, which has been chosen but not yet fed to either model. Then:

1. The draft feeds last and its own proposals one by one, writing $k$ positions to its cache.
2. The target feeds last and all $k$ proposals at once, writing $k+1$ positions to its cache and returning $k+1$ distributions.
3. Some prefix of the proposals is accepted, plus one corrected or bonus token.
4. Both caches are truncated back to the base length plus the accepted tokens plus last, the new invariant. Rejected proposals’ keys and values are simply forgotten. That’s why Chapter 16’s cache has truncate: rollback costs nothing.

One edge case: if all $k$ proposals are accepted, the draft never fed its own last proposal, so it’s fed once more before truncating.

Run it

python run.py speculate --k 4

Both models are random here (a 6-layer tiny Qwen3 as target), so the numbers show mechanics, not real-world speed:

{"draft": "target itself (perfect draft)", "rounds": 7, "acceptance_rate": 1.0, "tokens_per_target_call": 4.57}
{"draft": "first 2 of 6 layers", "identical_to_target_greedy": true, "rounds": 14, "acceptance_rate": 0.339, "tokens_per_target_call": 2.29}

With a perfect draft, every round yields $k+1 = 5$ tokens (the last round is cut short at 32 tokens). The “early exit” draft, the target’s own first two layers followed by its final norm and head, agrees with the full model a third of the time, still halving the target calls. And the output is identical to the target’s greedy output, as it must be. With --k 2 the early-exit draft’s acceptance rose to 44% but tokens per call fell to 1.78; with --k 8 acceptance fell to 21%.

On real models, a well-matched pair does much better: drafts from the same family and training data agree on most “easy” tokens (punctuation, the rest of a word, boilerplate code). Typical acceptance rates are 0.6-0.8 for chat and higher for code.

Where drafts come from

Tree verification checks several candidate continuations at once: the drafts form a tree, and one target forward pass evaluates all branches using an attention mask where each node sees only its ancestors. Your position-based attention (Chapter 5) supports this directly: give each node its depth as its position and an allowed mask from the tree.

When it doesn’t help

Speculation turns spare memory bandwidth into tokens. At large batch sizes, decode is already compute-bound (Chapter 24), so verifying $k+1$ tokens per request costs $k+1$ times as much, and rejected tokens are pure waste. Production engines reduce $k$ or turn speculation off as the batch grows. It also helps less when acceptance is low: high-temperature creative writing, or a draft trained on different data.

Build it

pytest tests/test_ch26_speculative.py
python run.py speculate --impl engine --k 4

Stretch exercises

1. ★ Measure acceptance rate and wall-clock speedup for Qwen3-1.7B (or 8B) as target and Qwen3-0.6B as draft on a GPU, for $k = 2, 4, 6$. Compare with the formula. Where: experiments/ch26.py (create it), calling engine.speculative.speculative_generate.
2. ★★ Implement prompt lookup decoding: find the most recent earlier occurrence of the last 3 tokens in the context and propose the $k$ tokens that followed. Measure acceptance on a “rewrite this paragraph” task. Where: add a prompt-lookup draft helper in engine/speculative.py.
3. ★★ Make $k$ adaptive: track a running acceptance rate and choose the $k$ that maximizes the expected speedup each round. Where: the round loop in speculative_generate in engine/speculative.py.
4. ★★★ Implement tree verification for greedy decoding: the draft proposes its top-2 tokens at each of 3 depths (a tree of 14 nodes), verify all branches in one target call with a tree attention mask, and commit the longest accepted path. Where: add tree drafting/verification in engine/speculative.py; pass tree masks through the target model to engine.attention.causal_attention.

Check your understanding

1. Why does verifying five tokens take about as long as generating one, at batch 1?
2. Show that the accept-or-resample rule outputs token $y$ with probability exactly $p(y)$.
3. Why does every round produce at least one token?
4. Why must both caches be truncated after a round, and to what length?
5. Why can speculation make a large-batch server slower?

Going deeper

• Leviathan, Kalman and Matias, Fast Inference from Transformers via Speculative Decoding (2023); Chen et al., Accelerating Large Language Model Decoding with Speculative Sampling (DeepMind, 2023).
• GPU Mode L22 (Hacker’s guide to speculative decoding in vLLM, Cade Daniels).
• PMPP §20.6 names batching and speculative decoding as the two main ways to raise the arithmetic intensity of the generation phase.
• Cai et al., Medusa (2024); Li et al., EAGLE, EAGLE-2, EAGLE-3 (2024-25); Elhoushi et al., LayerSkip (2024); Gloeckle et al., Better & Faster Large Language Models via Multi-token Prediction (2024); DeepSeek-V3 technical report (MTP).

27. Mixture of experts

Size without compute

In a dense transformer, every parameter takes part in every token. Doubling the parameters doubles the compute per token and the bytes read per decode step. A mixture of experts (MoE) breaks that link. Each MLP is replaced by $E$ smaller MLPs, the experts, and a router that sends each token to only $k$ of them.

dense layer:   x ──► MLP ──► y                               every token uses all of it
MoE layer:     x ──► router ──► top-k of E experts ──► weighted sum ──► y

Qwen3-30B-A3B has 128 experts per layer and uses 8 per token. It stores 30.5B parameters but each token is processed by about 3.3B active ones. The model knows roughly what a 30B model knows, and computes like a 3B one.

python run.py moe

{"model": "Qwen3-30B-A3B",      "expert_params_stored_B": 29.0,   "expert_params_active_B": 1.82,  "active_fraction": 0.063}
{"model": "Qwen3-235B-A22B",    "expert_params_stored_B": 227.15, "expert_params_active_B": 14.24, "active_fraction": 0.063}
{"model": "Qwen3.8-Flash-Next", "expert_params_stored_B": 121.09, "expert_params_active_B": 2.66,  "active_fraction": 0.022}

(Expert parameters only; attention and embeddings add the rest.) Flash-Next, the capstone’s target, routes each token to 10 of 512 experts plus a shared expert: it stores 121B parameters of experts and uses 2.2% of them per token.

The router

The router is one linear layer from the hidden state to $E$ scores. Qwen-family models then use:

1. Softmax over all experts, in FP32: $\pi = \operatorname{softmax}(W_r x)$.
2. Top-k: keep the $k$ largest probabilities and their expert indices.
3. Renormalize (when norm_topk_prob is set): divide the kept weights by their sum so they add to 1.

The layer’s output is the weighted sum of the chosen experts’ outputs: $y = \sum_{i \in \text{top-}k} w_i, \text{expert}_i(x)$.

Example. Four experts, $k = 2$, router probabilities $[0.1, 0.5, 0.3, 0.1]$. Experts 1 and 2 are chosen with weights $0.5$ and $0.3$; renormalized, they become $0.625$ and $0.375$.

class TopKRouter(nn.Module):
    def __init__(self, hidden, experts, top_k, normalize=True):
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(experts, hidden))
        self.top_k, self.normalize = top_k, normalize

    def forward(self, x):
        """x [N, D] -> (logits [N, E], weights [N, k], experts [N, k]).  (Your engine: Chapter 27)

        Softmax over all experts in FP32, keep the k largest, and (if normalize) rescale the
        kept probabilities to sum to 1.
        """
        logits = F.linear(x, self.weight)
        probs = torch.softmax(logits, dim=-1, dtype=torch.float32)
        weights, experts = probs.topk(self.top_k, dim=-1)
        if self.normalize:
            weights = weights / weights.sum(-1, keepdim=True)
        return logits, weights.to(x.dtype), experts

Two details matter for exact parity. The softmax must be in FP32: the router’s decision is discrete, so a BF16 rounding difference can swap two near-tied experts and change the output completely. And ties in topk are broken by index, which matters when comparing against another implementation (Chapter 29 meets this problem in earnest).

Other routers exist: DeepSeek-V3 uses a sigmoid score per expert with a learned bias for load balancing, and some models route groups of experts first. Each checkpoint documents its own; from_hf refuses configurations it doesn’t implement.

The shared expert

Some models add one shared expert that every token uses, alongside the routed ones. In Qwen3-Next and Flash-Next, its output is scaled by a learned gate, $\sigma(w_s^\top x)$, so each token decides how much of the shared knowledge to mix in:

$$ y = \sum_{i \in \text{top-}k} w_i, \text{expert}_i(x) + \sigma(w_s^\top x), \text{shared}(x). $$

class SparseMoeBlock(nn.Module):
    """Router + routed experts (+ an optional shared expert whose output is scaled by
    sigmoid(shared_expert_gate(x)), as in Qwen3-Next and Flash-Next)."""

    def __init__(self, hidden, experts, top_k, intermediate, shared_intermediate=0, normalize=True):
        super().__init__()
        self.gate = TopKRouter(hidden, experts, top_k, normalize)
        self.experts = Experts(experts, hidden, intermediate)
        self.shared_expert = SwiGLU(hidden, shared_intermediate) if shared_intermediate else None
        self.shared_expert_gate = nn.Linear(hidden, 1, bias=False) if shared_intermediate else None
        self.dispatch = "grouped"

    def forward(self, x):
        shape = x.shape
        flat = x.reshape(-1, shape[-1])
        _, weights, experts = self.gate(flat)
        run = self.experts.forward_grouped if self.dispatch == "grouped" else self.experts.forward_loop
        out = run(flat, weights, experts)
        if self.shared_expert is not None:
            out = out + torch.sigmoid(self.shared_expert_gate(flat)) * self.shared_expert(flat)
        return out.reshape(shape)

Dispatch: getting tokens to their experts

Experts are stored stacked: one tensor gate_up_proj [E, 2I, D] holding every expert’s gate and up projections, and one down_proj [E, D, I]. Checkpoints usually store each expert’s matrices separately (mlp.experts.17.gate_proj.weight), and the loader stacks them.

class Experts(nn.Module):
    """E SwiGLU experts in two stacked tensors: gate_up_proj [E, 2I, D] (gate rows, then up rows)
    and down_proj [E, D, I]. One tensor per kind keeps the layout friendly to grouped kernels."""

    def __init__(self, experts, hidden, intermediate):
        super().__init__()
        self.gate_up_proj = nn.Parameter(torch.empty(experts, 2 * intermediate, hidden))
        self.down_proj = nn.Parameter(torch.empty(experts, hidden, intermediate))
        nn.init.normal_(self.gate_up_proj, std=0.02)
        nn.init.normal_(self.down_proj, std=0.02)

    @property
    def num_experts(self):
        return self.gate_up_proj.shape[0]

    def expert(self, e, x):
        gate, up = F.linear(x, self.gate_up_proj[e]).chunk(2, dim=-1)
        return F.linear(F.silu(gate) * up, self.down_proj[e])

    def forward_loop(self, x, weights, experts):
        """Readable reference: for each expert, gather its tokens, run it, scatter-add back."""
        out = torch.zeros_like(x)
        for e in experts.unique().tolist():
            token, slot = torch.where(experts == e)
            out.index_add_(0, token, self.expert(e, x[token]) * weights[token, slot, None])
        return out

    def forward_grouped(self, x, weights, experts):
        """Sort the N*k assignments by expert so each expert's rows are contiguous: this is the
        layout a grouped GEMM consumes (one launch, many independent matmuls).  (Your engine: Chapter 27)"""
        n, k = experts.shape
        flat_expert = experts.reshape(-1)
        order = flat_expert.argsort(stable=True)
        token_of = order // k                                  # which token each sorted assignment came from
        counts = torch.bincount(flat_expert, minlength=self.num_experts).tolist()
        rows = x[token_of]                                      # [N*k, D] grouped by expert
        results, start = [], 0
        for e, count in enumerate(counts):                     # a grouped GEMM does these as one kernel
            if count:
                results.append(self.expert(e, rows[start:start + count]))
            start += count
        expert_out = torch.cat(results) * weights.reshape(-1)[order, None]
        return torch.zeros_like(x).index_add_(0, token_of, expert_out)

The loop dispatch is the readable reference: for each expert used in the batch, find its tokens, run the expert on them, and add the weighted result back to those tokens’ rows. Each expert is a separate small matmul, and the loop body runs once per active expert, with a host sync from unique().

The grouped dispatch prepares the layout a fast kernel needs:

1. Flatten the $N \times k$ assignments and sort them by expert (argsort, stable). Now each expert’s rows are contiguous.
2. Permute the token rows into that order (x[token_of]), duplicating each token $k$ times.
3. Run each expert on its contiguous slice. A grouped GEMM kernel does all of these as one launch: many independent matmuls of different sizes.
4. Multiply by the routing weights and unpermute with index_add_, summing each token’s $k$ contributions.

On a CPU, with 512 tokens, 32 experts and $k = 4$:

{"tokens": 512, "top_k": 4, "assignments": 2048, "busiest_expert": 80, "idlest_expert": 52, "ideal_per_expert": 64}
{"dispatch": "loop",    "ms": 16.1, "checksum": 1027.975}
{"dispatch": "grouped", "ms": 9.3,  "checksum": 1027.975}

Same result, and even the Python-level grouped version is faster. On a GPU, production MoE layers fuse the permutation, the grouped GEMMs and the unpermutation into a few kernels (vLLM’s fused_moe Triton kernel, MegaBlocks, DeepGEMM).

The first line shows load imbalance: with a random router, one expert got 80 tokens and another 52, against an ideal of 64. Trained routers are pushed toward balance during training, with an auxiliary loss (Switch Transformer: $E \sum_i f_i P_i$, the fraction of tokens routed to expert $i$ times its mean probability) or a bias adjusted on the fly (DeepSeek-V3’s auxiliary-loss-free balancing). At inference, imbalance means the busiest expert sets the step time.

The economics of MoE inference

Batch-1 decode is fast. A decode step reads only the experts its token uses. For Qwen3-30B-A3B in BF16, that’s about 6.6 GB per token (3.3B active parameters × 2 bytes), not 61 GB: decode runs at the speed of a 3B dense model.

Memory is the price. All experts must be resident, because the next token may need any of them. A 30B MoE needs 30B parameters of memory, the same as a 30B dense model. That’s why MoE models pair so well with quantization (experts are most of the bytes, Chapter 20) and with large-memory machines: the 128 GB of unified memory in a DGX Spark holds Flash-Next with 4-bit experts.

Larger batches touch more experts. With $B$ tokens each choosing $k$ of $E$ experts roughly uniformly, a layer touches about $E,\big(1 - (1 - k/E)^B\big)$ distinct experts:

By batch 32, a decode step of Qwen3-30B-A3B reads almost every expert, the full 61 GB, while each expert does only a little work (about $Bk/E = 2$ tokens each). MoE decode at moderate batch is therefore more memory-bound per token than a dense model with the same active parameters. Throughput still grows with batch until each expert gets enough tokens to be compute-bound, which needs batches in the hundreds or thousands. That’s why large MoE deployments use expert parallelism: experts spread over many GPUs, tokens exchanged with all-to-all communication, so that each GPU holds few experts and gets many tokens for each (Chapter 41).

Prefill is easy. Thousands of prompt tokens give every expert plenty of rows; grouped GEMMs run near peak.

Qwen3-MoE

Qwen3-MoE is Qwen3 with every MLP replaced by a SparseMoeBlock, and nothing else changed. Your Chapter 17 code gives you everything but the block:

@dataclass
class Qwen3MoeConfig(Qwen3Config):
    num_experts: int = 8
    num_experts_per_tok: int = 2
    moe_intermediate_size: int = 64
    norm_topk_prob: bool = True

    @classmethod
    def from_hf(cls, raw):
        if raw.get("model_type") != "qwen3_moe":
            raise ValueError("Expected model_type 'qwen3_moe'")
        if raw.get("mlp_only_layers") or raw.get("decoder_sparse_step", 1) != 1:
            raise ValueError("Dense layers inside a Qwen3-MoE stack are not implemented")
        raw = dict(raw)
        if "num_experts" not in raw and "num_local_experts" in raw:      # the name Transformers 5 writes
            raw["num_experts"] = raw["num_local_experts"]
        dense = Qwen3Config.from_hf({**raw, "model_type": "qwen3"})
        own = [f.name for f in fields(cls) if not hasattr(Qwen3Config, f.name)]
        missing = [name for name in own if name not in raw]
        if missing:                                  # never fall back to defaults for the model's shape
            raise ValueError(f"config.json lacks {missing}")
        return cls(**vars(dense), **{name: raw[name] for name in own})


class Qwen3MoeLayer(Qwen3Layer):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.mlp = SparseMoeBlock(cfg.hidden_size, cfg.num_experts, cfg.num_experts_per_tok,
                                  cfg.moe_intermediate_size, 0, cfg.norm_topk_prob)


class Qwen3Moe(Qwen3):
    """Qwen3 with every MLP replaced by a routed-expert block."""
    def __init__(self, cfg):
        super().__init__(cfg)
        self.model.layers = nn.ModuleList(Qwen3MoeLayer(cfg) for _ in range(cfg.num_hidden_layers))


@torch.no_grad()
def load_qwen3_moe(directory, device="cpu", dtype=torch.bfloat16):
    cfg = Qwen3MoeConfig.from_hf(read_config(directory))
    model = Qwen3Moe(cfg).to(device=device, dtype=dtype)
    params = dict(model.named_parameters())
    pending = {}
    for name, value in snapshot_tensors(directory):
        parts = name.split(".")
        if ".mlp.experts." in name:                        # model.layers.L.mlp.experts.E.kind.weight
            layer, e, kind = int(parts[2]), int(parts[5]), parts[6]
            pending.setdefault(layer, {})[(e, kind)] = value
            continue
        if name == "lm_head.weight" and cfg.tie_word_embeddings:
            continue
        if name not in params:
            raise ValueError(f"Unexpected tensor {name}")
        assign(params[name], value, name)
    for layer, tensors in pending.items():
        gate_up, down = stack_expert_tensors(tensors, cfg.num_experts, cfg.moe_intermediate_size)
        block = model.model.layers[layer].mlp.experts
        assign(block.gate_up_proj, gate_up, f"layer {layer} gate_up_proj")
        assign(block.down_proj, down, f"layer {layer} down_proj")
    return model.eval()

The loader collects the per-expert tensors as it streams the safetensors shards, then stacks them per layer. Qwen3-30B-A3B is 61 GB in BF16, so this chapter’s tests build tiny random Qwen3-MoE checkpoints, save them in Hugging Face’s format, and compare your model’s logits with transformers’ Qwen3MoeForCausalLM, with and without renormalization. The largest logit difference is about $10^{-7}$ in FP32, on logits of magnitude 0.5.

Note

Writing that test exposed a real bug. Transformers 5 saves the expert count as num_local_experts, while Qwen’s published configs say num_experts. The first version of from_hf silently fell back to its default of 8 experts and passed every test that happened to use 8. Now it accepts both names and refuses a config that lacks the field. Two lessons: never let a model’s shape come from a default, and make test models use non-default sizes everywhere.

def moe_parameter_split(hidden, layers, experts, top_k, intermediate, shared_intermediate=0):
    """(stored expert parameters, active expert parameters per token) for the MLP part only."""
    per_expert = 3 * hidden * intermediate
    shared = 3 * hidden * shared_intermediate
    stored = layers * (experts * per_expert + shared + experts * hidden)
    active = layers * (top_k * per_expert + shared + experts * hidden)
    return stored, active

Build it

pytest tests/test_ch27_moe.py
python run.py moe --impl engine

Stretch exercises

1. ★ Run a tiny Qwen3-MoE on 1,000 tokens of text and record every layer’s expert counts. Is the load more balanced in early or late layers? (Use a real checkpoint if you can; random routers are uninformative.) Where: experiments/ch27.py (create it), registering hooks on TopKRouter in engine/moe.py.
2. ★★ Write a Triton grouped GEMM: one program per (expert, output tile), with a prefix-sum array of row offsets telling each program where its expert’s rows start. Where: new engine/kernels/triton_grouped.py, called by Experts.forward_grouped in engine/moe.py.
3. ★★ Quantize only the experts to 4 bits with Chapter 20’s QuantLinear logic adapted to stacked tensors, leaving attention and the router in BF16. Measure memory and logit error. Where: add quantized stacked-expert storage/dispatch in engine/moe.py, using engine.quant.
4. ★★★ Expert offloading: keep experts in CPU memory, keep a GPU cache of the most recently used ones, and copy missing experts on demand. Measure cache hit rates over a conversation. Where: add an expert-cache variant in engine/moe.py; Chapter 40’s engine/offload.py provides the later integration.

Check your understanding

1. Why does an MoE’s decode speed at batch 1 depend on active parameters, while its memory depends on total parameters?
2. Why must the router’s softmax run in FP32?
3. What does sorting the assignments by expert buy you?
4. Why does a moderate batch make MoE decode read nearly all expert weights?
5. What problem does a load-balancing loss solve, and why does imbalance matter at inference too?

Going deeper

• Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer (2017); Fedus, Zoph and Shazeer, Switch Transformers (2021); Jiang et al., Mixtral of Experts (2024); DeepSeek-AI, DeepSeekMoE (2024) and the DeepSeek-V3 technical report (fine-grained and shared experts, auxiliary-loss-free balancing).
• Gale et al., MegaBlocks: Efficient Sparse Training with Mixture-of-Experts (2022), on grouped and block-sparse kernels.
• The Qwen3 technical report (2025) for Qwen3-MoE, and the Qwen3-Next model card for the gated shared expert.
• vLLM’s fused_moe Triton kernel and SGLang’s MoE runner; GPU Mode L11 (Sparsity) for the broader picture of sparse computation on GPUs.

28. Linear attention and Gated DeltaNet

The cost of remembering everything

Softmax attention keeps every past key and value. At decode, each new token reads the whole KV cache: time and memory per token grow linearly with context, and prefill grows quadratically. Chapter 16’s formula put Qwen3-0.6B’s cache at 3.5 GiB for a 32k context, three times the model. At the million-token contexts that agents and long documents want, the cache, not the weights, is the model.

Could a layer instead keep a fixed-size summary of the past, updated once per token, like an RNN? That’s linear attention. Flash-Next uses it in 36 of its 48 layers.

Removing the softmax

Attention for query $t$ is $o_t = \sum_{i \le t} \frac{\exp(q_t \cdot k_i)}{Z}, v_i$. The exponential couples $q_t$ and $k_i$ inside each term, so nothing can be precomputed. Replace $\exp(q \cdot k)$ by a plain dot product (or by $\phi(q) \cdot \phi(k)$ for some feature map $\phi$), and the sum factorizes (Katharopoulos et al., 2020, Transformers are RNNs):

$$ o_t = \sum_{i \le t} (q_t^\top k_i), v_i = q_t^\top \underbrace{\Big(\sum_{i \le t} k_i v_i^\top\Big)}_{S_t ,\in, \mathbb{R}^{d_k \times d_v}} . $$

$S_t$ is a $d_k \times d_v$ matrix, the state, and it updates with one outer product per token:

$$ S_t = S_{t-1} + k_t v_t^\top, \qquad o_t = S_t^\top q_t . $$

Per-token cost and memory are now constant, whatever the context. Prefill is linear in $T$ instead of quadratic.

The state is an associative memory

Think of $S$ as a memory that stores value $v$ under key $k$. Reading with a query equal to a stored key returns $S^\top k_j = \sum_i (k_i \cdot k_j), v_i$: the right value, if keys are orthonormal, plus interference from every other stored value in proportion to how similar its key is. A $d_k$-dimensional state can hold at most $d_k$ orthogonal keys. Writing more than that, and real sequences are thousands of tokens, the memories blur together.

python run.py linear

The first part stores $n$ random key-value pairs in a 64 × 64 state and reads each back with its key (unit-norm keys, relative error of the retrieved value):

{"pairs": 16,  "key_dim": 64, "additive_error_all": 0.45,  "delta_error_all": 0.303, "delta_error_last_8": 0.224}
{"pairs": 64,  "key_dim": 64, "additive_error_all": 0.967, "delta_error_all": 0.714, "delta_error_last_8": 0.23}
{"pairs": 256, "key_dim": 64, "additive_error_all": 2.013, "delta_error_all": 1.199, "delta_error_last_8": 0.282}

Plain additive writes degrade steadily, and past 64 pairs the “retrieved” value is mostly noise. Two ideas fix this, and Gated DeltaNet uses both.

Forget: decay

Multiply the state by a decay $\alpha_t \in (0, 1)$ before each write: $S_t = \alpha_t S_{t-1} + k_t v_t^\top$. Old memories fade, making room for new ones. If $\alpha_t$ is computed from the input (a gate), the model can decide per token how much to forget: keep everything inside a sentence, wipe the slate at a document boundary. This is the idea behind RetNet, GLA and Mamba-2.

Correct: the delta rule

Instead of adding $v_t$ blindly, first ask what the memory currently returns for $k_t$, and write only the error:

$$ S_t = S_{t-1} + \beta_t, k_t \big(v_t - S_{t-1}^\top k_t\big)^\top . $$

This is the delta rule (Widrow-Hoff, 1960; DeltaNet, Schlag et al., 2021): one step of gradient descent on $\tfrac12 \lVert S^\top k_t - v_t \rVert^2$ with learning rate $\beta_t$. With $\beta = 1$ and a unit key, reading $k_t$ afterwards returns exactly $v_t$: the old association along $k_t$ is replaced, not piled on. That’s the “delta” columns above: the most recent pairs are retrieved well even at 256 pairs, and overall error is much lower.

Worked example (the milestone test). $k = q = [1, 0]$, $v = [2, 4]$, empty state. With $\beta = 0.5$: the prediction is $[0, 0]$, the error is $[2, 4]$, and $S$ gains $0.5 \cdot [1, 0]^\top [2, 4]$, so its first row is $[1, 2]$ and the output is $[1, 2]$. Repeat with $\beta = 1$: the prediction is now $[1, 2]$, the error $[1, 2]$, and the first row becomes $[2, 4]$, exactly $v$.

Gated DeltaNet

Combine both (Yang, Kautz and Hatamizadeh, 2024):

$$ S_t = \alpha_t S_{t-1} + \beta_t, k_t \big(v_t - \alpha_t S_{t-1}^\top k_t\big)^\top, \qquad o_t = S_t^\top q_t , $$

with $\alpha_t = e^{g_t}$ for a learned log-decay $g_t \le 0$, and $q, k$ L2-normalized (so that $\beta \le 1$ keeps updates stable).

The recurrent form

Decode processes one token at a time, so the recurrence is exactly what it needs:

def recurrent_gated_delta_rule(q, k, v, g, beta, state=None):
    """One token at a time.  (Your engine: Chapter 28)

    q, k [B, H, T, Dk] (already L2-normalized; q also scaled by Dk^-0.5); v [B, H, T, Dv];
    g, beta [B, H, T]; state [B, H, Dk, Dv] or None. Returns (o [B, H, T, Dv], final state).
    """
    b, h, t, dk = k.shape
    dv = v.shape[-1]
    S = torch.zeros(b, h, dk, dv, device=q.device, dtype=torch.float32) if state is None else state.float()
    q, k, v, g, beta = (x.float() for x in (q, k, v, g, beta))
    out = torch.empty(b, h, t, dv, device=q.device, dtype=torch.float32)
    for i in range(t):
        S = S * g[:, :, i, None, None].exp()
        prediction = torch.einsum("bhk,bhkv->bhv", k[:, :, i], S)
        delta = (v[:, :, i] - prediction) * beta[:, :, i, None]
        S = S + torch.einsum("bhk,bhv->bhkv", k[:, :, i], delta)
        out[:, :, i] = torch.einsum("bhk,bhkv->bhv", q[:, :, i], S)
    return out, S

// S is [dk][dv] row-major: decay, predict v from k, write beta * error along k, read with q.
inline std::vector<float> delta_step(std::vector<float>& S, const std::vector<float>& q, const std::vector<float>& k,
                                     const std::vector<float>& v, float decay, float beta) {
    size_t dv = v.size();
    for (float& s : S) s *= decay;
    std::vector<float> pred(dv, 0.f), out(dv, 0.f);
    for (size_t i = 0; i < k.size(); ++i) for (size_t j = 0; j < dv; ++j) pred[j] += k[i] * S[i * dv + j];
    for (size_t i = 0; i < k.size(); ++i) for (size_t j = 0; j < dv; ++j) S[i * dv + j] += beta * k[i] * (v[j] - pred[j]);
    for (size_t i = 0; i < q.size(); ++i) for (size_t j = 0; j < dv; ++j) out[j] += q[i] * S[i * dv + j];
    return out;
}

#![allow(unused)]
fn main() {
/// S is [dk][dv] stored row-major. Decay, predict v from k, write the scaled error, read with q.
pub fn delta_step(s: &mut [f32], q: &[f32], k: &[f32], v: &[f32], decay: f32, beta: f32) -> Vec<f32> {
    let dv = v.len();
    s.iter_mut().for_each(|x| *x *= decay);
    let mut prediction = vec![0.0; dv];
    for (i, ki) in k.iter().enumerate() {
        for j in 0..dv {
            prediction[j] += ki * s[i * dv + j];
        }
    }
    for (i, ki) in k.iter().enumerate() {
        for j in 0..dv {
            s[i * dv + j] += beta * ki * (v[j] - prediction[j]);
        }
    }
    let mut out = vec![0.0; dv];
    for (i, qi) in q.iter().enumerate() {
        for j in 0..dv {
            out[j] += qi * s[i * dv + j];
        }
    }
    out
}
}

Each step is a few $d_k \times d_v$ elementwise operations and mat-vecs per head: memory-bound on the state, but the state is small (16 KiB per head in FP32 for $d_k = d_v = 64$) and doesn’t grow. The state stays in FP32 even when the model runs in BF16, because it accumulates over thousands of steps (Chapter 13).

The chunked form: recurrence as matrix multiplication

For prefill, a loop of $T$ sequential steps wastes a GPU. The chunked algorithm processes the sequence in chunks of $C$ tokens (64 is typical). Within a chunk, it unrolls the recurrence algebraically into matrix products; across chunks, it carries the state.

Let $G_i = \sum_{j \le i} g_j$ be the cumulative log-decay within the chunk, and $D_{ij} = e^{G_i - G_j}$ for $j \le i$ (the decay from position $j$ to $i$). Unrolling shows that the values actually written, $U$ (each $v_i$ minus what the memory predicted for $k_i$, scaled by $\beta_i$), satisfy a unit lower-triangular linear system: each position’s write depends on the writes before it in the chunk. Solving that system (the “UT transform”) turns the sequential dependency into one triangular solve per chunk. Then the outputs and the next chunk’s state are matrix products:

$$ \begin{aligned} \big(I + \operatorname{strict_tril}(\beta_i, k_i \cdot k_j, D_{ij})\big), U &= \beta \odot V - \beta \odot e^{G} \odot (K S_0) \ O &= e^{G} \odot (Q S_0) + \operatorname{tril}!\big(QK^\top \odot D\big), U \ S_\text{end} &= e^{G_C} S_0 + \textstyle\sum_j e^{G_C - G_j}, k_j U_j^\top \end{aligned} $$

def chunk_gated_delta_rule(q, k, v, g, beta, state=None, chunk=64):
    """Same contract and result as the recurrent form, computed chunk by chunk.  (Your engine: Chapter 28)

    Inside a chunk, let G_i be the cumulative log-decay up to position i and
    D_ij = exp(G_i - G_j) for j <= i (0 above the diagonal). Unrolling the recurrence shows the
    values actually written, U, solve a unit-lower-triangular system:
        (I + strict_tril(beta_i k_i.k_j D_ij)) U = beta * v - beta * exp(G) * (K S_0)
    (the "UT transform"); the chunk's outputs and final state then follow from matmuls:
        O = exp(G) * (Q S_0) + (tril(Q K^T * D)) U
        S_end = exp(G_last) S_0 + sum_j exp(G_last - G_j) k_j U_j^T
    """
    b, h, t, dk = k.shape
    dv = v.shape[-1]
    q, k, v, g, beta = (x.float() for x in (q, k, v, g, beta))
    pad = (-t) % chunk
    if pad:
        q, k, v = (F.pad(x, (0, 0, 0, pad)) for x in (q, k, v))
        g, beta = (F.pad(x, (0, pad)) for x in (g, beta))
    n = q.shape[2] // chunk
    q, k, v = (x.view(b, h, n, chunk, -1) for x in (q, k, v))
    g, beta = g.view(b, h, n, chunk), beta.view(b, h, n, chunk)
    G = g.cumsum(-1)                                                       # [b,h,n,C]
    lower = torch.ones(chunk, chunk, dtype=torch.bool, device=q.device).tril()
    decay = (G[..., :, None] - G[..., None, :]).masked_fill(~lower, float("-inf")).exp()   # D_ij
    k_beta, v_beta = k * beta[..., None], v * beta[..., None]
    system = (k_beta @ k.transpose(-1, -2)) * decay                        # beta_i k_i.k_j D_ij
    system = system.tril(-1) + torch.eye(chunk, device=q.device)           # I + strictly-lower part
    # Solve once for both right-hand sides: the part that depends on S_0 and the part that does not.
    w = torch.linalg.solve_triangular(system, k_beta * G[..., None].exp(), upper=False, unitriangular=True)
    u = torch.linalg.solve_triangular(system, v_beta, upper=False, unitriangular=True)
    attn = (q @ k.transpose(-1, -2)) * decay                               # causal, decayed Q K^T
    S = torch.zeros(b, h, dk, dv, device=q.device) if state is None else state.float()
    out = torch.empty_like(v)
    for c in range(n):
        U = u[:, :, c] - w[:, :, c] @ S                                    # values actually written
        out[:, :, c] = (q[:, :, c] * G[:, :, c, :, None].exp()) @ S + attn[:, :, c] @ U
        tail = (G[:, :, c, -1:] - G[:, :, c]).exp()                       # exp(G_last - G_j)
        S = S * G[:, :, c, -1, None, None].exp() + (k[:, :, c] * tail[..., None]).transpose(-1, -2) @ U
    return out.view(b, h, n * chunk, dv)[:, :, :t], S

The milestone test runs chunk sizes 1, 5, 8 and 64 against the recurrent form on a sequence of 37 tokens (so the last chunk is partial) with a non-zero initial state. They all agree to $10^{-5}$.

Note

The first version of this function kept the wrong triangle of the system matrix, a transposition that’s easy to make when deriving it on paper. Outputs looked plausible, and the chunk-size-1 test even passed, because a 1 × 1 chunk has no off-diagonal part. Testing several chunk sizes, including ones that don’t divide the length, caught it. Test a fast algorithm against the slow one over every boundary case you can construct.

On a laptop CPU, for 4 heads and 1,024 tokens:

{"form": "recurrent", "tokens": 1024, "ms": 132.0}
{"form": "chunked",   "tokens": 1024, "ms": 12.2}
{"max_difference": 1.043081283569336e-07, "state_bytes_per_head": 16384, "kv_bytes_per_head_at_1024_tokens": 262144}

Eleven times faster, numerically identical, and the state is 16× smaller than one head’s KV cache at only 1,024 tokens (in BF16). The production kernels (the flash-linear-attention library’s Triton kernels, which Qwen’s models use) fuse these steps per chunk and run at near-matmul speed.

The full layer

The Gated DeltaNet layer in Qwen3.5 and Flash-Next wraps the rule with projections, gates and a short convolution. Parameter names match the checkpoints:

class GatedDeltaNet(nn.Module):
    """A Qwen3.5 / Flash-Next linear-attention layer.  (Your engine: Chapter 28)"""

    def __init__(self, hidden, key_heads, value_heads, key_dim, value_dim, conv_kernel=4, eps=1e-6,
                 gate_activation="silu"):
        super().__init__()
        if value_heads % key_heads:
            raise ValueError("value heads must be a multiple of key heads")
        self.kh, self.vh, self.dk, self.dv = key_heads, value_heads, key_dim, value_dim
        self.key_size, self.value_size = key_heads * key_dim, value_heads * value_dim
        conv_dim = 2 * self.key_size + self.value_size
        self.in_proj_qkv = nn.Linear(hidden, conv_dim, bias=False)
        self.in_proj_z = nn.Linear(hidden, self.value_size, bias=False)
        self.in_proj_b = nn.Linear(hidden, value_heads, bias=False)
        self.in_proj_a = nn.Linear(hidden, value_heads, bias=False)
        self.conv1d = nn.Conv1d(conv_dim, conv_dim, conv_kernel, groups=conv_dim, bias=False)
        self.dt_bias = nn.Parameter(torch.ones(value_heads))
        self.A_log = nn.Parameter(torch.log(torch.empty(value_heads).uniform_(0.01, 16)))
        self.norm = RMSNormGated(value_dim, eps, gate_activation)
        self.out_proj = nn.Linear(self.value_size, hidden, bias=False)

    def forward(self, x, state=None, chunk=64):
        """x [B, T, D] -> (y [B, T, D], new LinearAttentionState).  (Your engine: Chapter 28)"""
        b, t, _ = x.shape
        state = state or LinearAttentionState()
        mixed, conv_state = causal_conv1d(self.in_proj_qkv(x).transpose(1, 2), self.conv1d.weight[:, 0], state.conv)
        q, k, v = mixed.transpose(1, 2).split([self.key_size, self.key_size, self.value_size], dim=-1)
        q = q.reshape(b, t, self.kh, self.dk).transpose(1, 2)
        k = k.reshape(b, t, self.kh, self.dk).transpose(1, 2)
        v = v.reshape(b, t, self.vh, self.dv).transpose(1, 2)
        repeat = self.vh // self.kh                                  # several value heads share a key head
        q, k = q.repeat_interleave(repeat, 1), k.repeat_interleave(repeat, 1)
        q = l2norm(q.float()) * self.dk ** -0.5
        k = l2norm(k.float())
        beta = torch.sigmoid(self.in_proj_b(x)).transpose(1, 2)               # write strength in (0,1)
        g = (-self.A_log.float().exp() * F.softplus(self.in_proj_a(x).float() + self.dt_bias)).transpose(1, 2)
        rule = recurrent_gated_delta_rule if t == 1 else chunk_gated_delta_rule
        kwargs = {} if t == 1 else {"chunk": chunk}
        o, recurrent = rule(q, k, v, g, beta, state.recurrent, **kwargs)
        z = self.in_proj_z(x).reshape(b, t, self.vh, self.dv)
        o = self.norm(o.transpose(1, 2).to(x.dtype), z)
        return self.out_proj(o.reshape(b, t, -1)), LinearAttentionState(conv_state, recurrent)

Step by step:

1. Project $x$ to $q, k, v$ in one matrix (in_proj_qkv), plus an output gate $z$ (in_proj_z), a write strength $b$ and a decay input $a$ (one scalar per value head each).
2. Short causal convolution over time on $q, k, v$ (kernel 4, one filter per channel, then SiLU). It lets each token mix in its three predecessors before the state sees it: cheap local context. Its own state is the last 3 inputs.
3. Heads: there are more value heads than key heads (Flash-Next: 16 key heads, 48 value heads), and each key head is shared by several value heads, like grouped-query attention in reverse.
4. Normalize: $q$ and $k$ L2-normalized, $q$ scaled by $d_k^{-1/2}$.
5. Gates: $\beta = \sigma(b)$, and the log-decay $g = -e^{A_\text{log}} \cdot \operatorname{softplus}(a + \text{dt_bias})$, the same parameterization as Mamba’s $\Delta$: always negative, so $\alpha = e^{g} \in (0, 1)$.
6. The rule: recurrent for a single token, chunked otherwise.
7. Gated RMSNorm: normalize each head’s output, multiply by $\operatorname{SiLU}(z)$, then project back to the hidden size.

def causal_conv1d(x, weight, state=None, activation=True):
    """Depthwise causal convolution over time with carried history.  (Your engine: Chapter 28)

    x [B, C, T]; weight [C, K]; state [B, C, K-1] = the previous K-1 inputs (zeros at start).
    Returns (y [B, C, T], new_state). Output t sees inputs t-K+1 .. t only.
    """
    kernel = weight.shape[-1]
    if state is None:
        state = x.new_zeros(x.shape[0], x.shape[1], kernel - 1)
    joined = torch.cat((state.to(x.dtype), x), dim=-1)
    y = F.conv1d(joined, weight[:, None, :].to(x.dtype), groups=x.shape[1])
    return (F.silu(y) if activation else y), joined[..., -(kernel - 1):]

class RMSNormGated(nn.Module):
    """RMSNorm(x) * weight * silu(z): normalize the read-out, then let a gate decide how much passes."""
    def __init__(self, width, eps=1e-6, activation="silu"):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(width))
        self.eps, self.activation = eps, activation

    def forward(self, x, z):
        x32 = x.float()
        normed = x32 * torch.rsqrt(x32.square().mean(-1, keepdim=True) + self.eps)
        gate = F.silu(z.float()) if self.activation == "silu" else torch.sigmoid(z.float())
        return (self.weight * normed.to(x.dtype) * gate).to(x.dtype)

The state a request carries

A Gated DeltaNet layer’s per-request state is the convolution history plus the recurrent matrices, fixed size whatever the context. For Flash-Next: 48 value heads × 128 × 128 FP32 values per layer, 3 MiB, times 36 linear layers, 108 MiB per sequence. Its 12 attention layers’ KV cache at 32k tokens is 0.75 GiB; at 262k tokens, 6 GiB. An all-attention model of the same shape would need four times as much.

The fixed state changes the engine in three ways:

• No truncation. A KV cache can be rolled back by forgetting positions (Chapter 26’s speculative decoding). A recurrent state can’t: the rejected tokens have been mixed into it. Engines snapshot the state before speculating and restore it on rejection (Chapter 30’s HybridState).
• Prefix caching needs snapshots too, stored at block boundaries: the state after a shared prefix, not a list of per-position entries.
• Batching is simple: every request’s state has the same shape.

Why hybrid?

Linear attention’s fixed memory is also its weakness. Tasks that need exact recall of an arbitrary earlier token (copy this ID, find the needle in the haystack) are hard for any fixed-size state once the context is long, as the recall demo showed. A few full-attention layers restore precise lookup, and models with a 3:1 ratio of linear to attention layers (Qwen3-Next, Qwen3.5, Flash-Next) match or beat pure attention models on long-context benchmarks at a fraction of the cache. Flash-Next goes one step further: its attention layers are themselves sparse (Chapter 29).

Build it

pytest tests/test_ch28_gdn.py
python run.py linear --impl engine

Stretch exercises

1. ★ Extend the recall demo with decay: at $\alpha = 0.95$, how does error depend on how long ago a pair was written? Where: experiments/ch28.py (create it), adapting the recall example in run.py’s cmd_linear.
2. ★★ Measure the chunked form’s time for chunk sizes 16, 32, 64, 128 and $T$ = 4,096. Explain the optimum in terms of the triangular solve’s cost ($O(C^2)$ per position) against the number of sequential chunk steps ($T/C$). Where: experiments/ch28.py (create it), calling engine.gdn.chunk_gated_delta_rule.
3. ★★ Write a Triton kernel for the recurrent step at decode: one program per (sequence, head), the state tile in registers, all of decay, predict, correct and read fused. Where: new engine/kernels/triton_gdn.py, selected by GatedDeltaNet.forward in engine/gdn.py.
4. ★★★ Implement Mamba-2’s selective state-space recurrence ($S_t = \alpha_t S_{t-1} + k_t v_t^\top$ with scalar $\alpha_t$ per head) in the same chunked style, and show it’s Gated DeltaNet with $\beta$’s correction term removed. Where: add a selective-state-space helper beside chunk_gated_delta_rule in engine/gdn.py.

Check your understanding

1. Why does replacing $\exp(q \cdot k)$ by $q \cdot k$ make attention a recurrence?
2. Why do additive writes interfere, and what limits how many associations a $d_k \times d_v$ state can hold?
3. What does the delta rule write instead of $v_t$, and why does $\beta = 1$ replace an association exactly?
4. Why does the chunked form need a triangular solve?
5. Why can’t a recurrent state be truncated like a KV cache, and what do engines do instead?

Going deeper

• Katharopoulos et al., Transformers are RNNs (2020); Schlag, Irie and Schmidhuber, Linear Transformers Are Secretly Fast Weight Programmers (2021); Yang et al., Gated Linear Attention (2023), Parallelizing Linear Transformers with the Delta Rule over Sequence Length (2024) and Gated Delta Networks (2024).
• Gu and Dao, Mamba (2023) and Dao and Gu, Transformers are SSMs (Mamba-2, 2024), for the state-space view of the same family.
• Songlin Yang’s blog series DeltaNet Explained (three parts), the clearest derivation of the chunked algorithm; the flash-linear-attention repository for production kernels.
• GPU Mode L20, L21 and L24 (Scan), the parallel-prefix algorithm behind Mamba’s selective scan; PMPP Chapter 11 (Scan: Kogge-Stone and Brent-Kung parallel prefix sums).
• The Qwen3-Next and Qwen3.5 model cards for the hybrid 3:1 layout.

29. Sparse attention: windows, top-k and learned indexers

Most attention weights are tiny

A trained model’s attention is usually peaked: for a given query, a handful of keys carry most of the weight and thousands share the rest. If the weight on a key is $10^{-6}$, skipping it changes the output by about $10^{-6}$ of that value. Sparse attention bets on this: read only the keys that matter, and save the memory traffic of the rest.

The bet doesn’t always pay. run.py sparse compares exact top-$k$ attention (keep each query’s $k$ highest-scoring keys, renormalize) with dense attention over 1,024 keys, for flat and for peaked score distributions:

python run.py sparse

{"context": 1024, "score_scale": 1.0, "relative_error": {"budget_16": 4.4,   "budget_64": 1.892, "budget_256": 0.595}}
{"context": 1024, "score_scale": 4.0, "relative_error": {"budget_16": 0.123, "budget_64": 0.026, "budget_256": 0.002}}

With flat scores (random queries and keys), the output is an average over many values, and keeping 256 of 1,024 keys still leaves a 60% error. With peaked scores, 64 keys give 2.6% error and 256 give 0.2%. Real models sit closer to the second case for most heads and most layers, which is why sparse attention works in practice, and why it’s trained into the model rather than applied after the fact.

Fixed patterns: windows and sinks

The simplest selection ignores content. A sliding window lets each query see only the last $w$ keys: memory is $O(w)$ per sequence, and the KV cache becomes a ring buffer. Mistral 7B used $w = 4{,}096$; Gemma alternates windowed and global layers.

Xiao et al. (2023) found a twist: windowed attention collapses when the first tokens leave the window, because models learn to dump excess attention weight on the very first positions (“attention sinks”). Keeping the first few tokens visible forever fixes it.

def sliding_window_allowed(query_positions, key_positions, window, sinks=0):
    """Keys within `window` positions of the query, plus the first `sinks` tokens of the sequence."""
    qp = query_positions[..., :, None]
    kp = key_positions[..., None, :]
    return ((qp - kp) < window) | (kp < sinks)

Your causal_attention takes an allowed mask (Chapter 5), so a window is one more boolean term. The milestone test checks a window of 2 with 1 sink: position 5 sees keys 0, 4 and 5.

Content-based selection: top-k

Windows can’t recall something important from 50,000 tokens ago. Top-k attention selects by score instead:

def topk_attention(q, k, v, budget, query_positions=None):
    """Exact scores, but each query keeps only its `budget` highest-scoring visible keys.
    Teaches selection; it still computes every score, so it saves nothing by itself."""
    t, s = q.shape[-2], k.shape[-2]
    qp = torch.arange(s - t, s, device=q.device) if query_positions is None else query_positions
    scores = (q.float() @ k.float().transpose(-2, -1)) / math.sqrt(q.shape[-1])
    visible = torch.arange(s, device=q.device)[None, :] <= qp[:, None]
    scores = scores.masked_fill(~visible, float("-inf"))
    keep = scores.topk(min(budget, s), dim=-1).indices
    allowed = torch.zeros_like(scores, dtype=torch.bool).scatter(-1, keep, True) & visible
    return causal_attention(q, k, v, qp, None, allowed)

This version is useful for understanding and for measuring error, but it saves nothing: to find the top $k$ scores it computes all of them, reading every key. The savings come only when something cheaper than attention decides which keys to read. That’s an indexer.

Qwen Sparse Attention

Flash-Next’s 12 attention layers use QSA: a small learned indexer picks, for each query, a budget of 2,048 keys out of the whole context, and full attention reads only those. Three ideas make the indexer cheap:

1. Blocks, not tokens. The indexer’s keys are averaged over consecutive blocks of 4 tokens (the compression ratio). Selection happens per block: one decision per 4 tokens, and contiguous memory reads.
2. Small and shared. The indexer has its own tiny projection: 4 query heads of dimension 128, and one key per block shared by all of them, against attention’s 2 KV heads of dimension 256 per token.
3. A simple score. For query $q$ (with heads $h$) and pooled block key $\bar k_b$: $$ \text{score}(q, b) = \frac{1}{\sqrt{d}} \sum_h \operatorname{ReLU}\big(q_h \cdot \bar k_b\big). $$ The ReLU lets each indexer head vote only for blocks, never against.

Then each query keeps its top $\text{budget}/\text{ratio} = 512$ complete visible blocks, plus the open block, the partially filled block it’s in, which is always readable so recent context is never lost. Both the indexer’s queries and the pooled block keys get RoPE (the block key at the block’s first position), so the indexer is position-aware.

class QSAIndexer(nn.Module):
    """Chooses, for every query, which keys attention may read.  (Your engine: Chapter 29)

    Keys are averaged over consecutive blocks of `ratio` tokens; one shared indexer key per
    block (RMSNorm, then RoPE at the block's first position) is scored against several
    small indexer query heads: score = sum over heads of ReLU(q . k) / sqrt(d). Each query
    keeps its top budget/ratio complete visible blocks plus the incomplete trailing block.
    """

    def __init__(self, hidden, heads, head_dim, budget, ratio, eps=1e-6):
        super().__init__()
        if budget % ratio:
            raise ValueError("budget must be a multiple of the compression ratio")
        self.heads, self.head_dim, self.ratio, self.block_topk = heads, head_dim, ratio, budget // ratio
        self.index_qk_proj = nn.Linear(hidden, (heads + 1) * head_dim, bias=False)
        self.q_layernorm = ZeroCenteredRMSNorm(head_dim, eps)
        self.k_layernorm = ZeroCenteredRMSNorm(head_dim, eps)

    def forward(self, x, positions, rotary_dim, theta, cached_keys=None):
        """x [B, T, D] at absolute positions [T] -> (allowed [B, 1, T, S], all raw keys [B, S, d]).  (Your engine: Chapter 29)"""
        b, t, _ = x.shape
        q, raw = self.index_qk_proj(x).split([self.heads * self.head_dim, self.head_dim], dim=-1)
        q = self.q_layernorm(q.view(b, t, self.heads, self.head_dim))
        cos, sin = rope_cos_sin(positions, rotary_dim, theta, q.dtype)       # [1, 1, T, r]
        q = apply_rope(q, cos.transpose(1, 2), sin.transpose(1, 2))          # broadcast over heads
        keys = raw if cached_keys is None else torch.cat((cached_keys, raw), dim=1)
        s = keys.shape[1]
        n_blocks = s // self.ratio
        key_pos = torch.arange(s, device=x.device)
        visible_tail_start = ((positions + 1) // self.ratio) * self.ratio    # first token of the open block
        tail = (key_pos[None, :] >= visible_tail_start[:, None]) & (key_pos[None, :] <= positions[:, None])
        if n_blocks == 0:
            return tail[None, None].expand(b, 1, t, s), keys
        pooled = keys[:, :n_blocks * self.ratio].view(b, n_blocks, self.ratio, -1).float().mean(2).to(keys.dtype)
        pooled = self.k_layernorm(pooled)
        starts = torch.arange(n_blocks, device=x.device) * self.ratio
        bcos, bsin = rope_cos_sin(starts, rotary_dim, theta, pooled.dtype)   # [1, 1, n, r]
        pooled = apply_rope(pooled, bcos[:, 0], bsin[:, 0])
        scores = torch.relu(torch.einsum("bthd,bnd->btnh", q.float(), pooled.float())).sum(-1) / math.sqrt(self.head_dim)
        # Queries that see the same number of complete blocks run topk over exactly those blocks.
        # (topk breaks ties differently for different vector lengths; ReLU makes ties at 0 common,
        # so slicing rather than masking with -inf is what keeps us identical to the reference.)
        visible_blocks = (positions + 1) // self.ratio                         # [T]
        chosen = torch.zeros(b, t, n_blocks, dtype=torch.bool, device=x.device)
        for count in visible_blocks.unique().tolist():
            if count == 0:
                continue
            rows = visible_blocks == count
            top = scores[:, rows, :count].topk(min(self.block_topk, count), dim=-1).indices
            chosen[:, rows] = chosen[:, rows].scatter(-1, top, True)
        block_of_key = (key_pos // self.ratio).clamp(max=n_blocks - 1)        # tail keys are masked below
        from_blocks = chosen.gather(-1, block_of_key.expand(b, t, s))
        from_blocks = from_blocks & (key_pos < n_blocks * self.ratio)
        return (from_blocks | tail[None])[:, None], keys

The indexer’s output is just an allowed mask [B, 1, T, S], and attention uses it exactly like a causal mask. A real kernel instead gathers the selected blocks’ K and V (Chapter 25’s block tables are the natural fit) and never reads the others.

What it saves

The arithmetic for one Flash-Next attention layer at decode, in BF16 (the second half of run.py sparse):

{"context": 8192,   "dense_kv_MB_per_layer_step": 16.8,  "qsa_MB_per_layer_step": 4.7,  "reduction": 3.5}
{"context": 32768,  "dense_kv_MB_per_layer_step": 67.1,  "qsa_MB_per_layer_step": 6.3,  "reduction": 10.7}
{"context": 262144, "dense_kv_MB_per_layer_step": 536.9, "qsa_MB_per_layer_step": 21.0, "reduction": 25.6}

QSA reads the K and V of 2,048 selected tokens plus every block’s 128-dimensional indexer key. The indexer’s reads still grow with context, but 16 times more slowly than the full KV cache (one 256-byte key per 4 tokens, against 4 KiB of K and V per token). The cache itself still stores everything: sparse attention saves bandwidth and compute, not memory. That’s why Flash-Next combines it with linear attention (Chapter 28), which saves memory.

A subtlety: a query may not see itself

A query at position 11, with blocks of 4, completes block 2. Its open block is now empty, and block 2 is a complete block that must compete in the top-$k$ with blocks 0 and 1. With a budget of two blocks and a low score for block 2, the query attends to blocks 0 and 1 and not to itself. The milestone test asserts that this can happen, rather than assuming that every query sees its own token. It always sees some keys (at least one complete block wins the top-$k$), just not necessarily its own.

Gated attention with partial RoPE

Flash-Next’s attention layer (inherited from Qwen3-Next and Qwen3.5) changes three more things from Chapter 17’s Qwen3:

• An output gate. q_proj produces twice the query width: a query and a gate per head. The attention output is multiplied by $\sigma(\text{gate})$ before o_proj. The gate lets a head output nothing when it has nothing useful to add, which also removes the need for attention sinks (Qiu et al., 2025, Gated Attention for Large Language Models).
• Partial RoPE. Only the first 25% of each head’s 256 dimensions are rotated (partial_rotary_factor); the remaining 192 carry content without position. Your apply_rope from Chapter 17 rotates the first rotary_dim features.
• Zero-centered RMSNorm for the Q and K norms: the scale is $1 + w$ with $w$ initialized at 0, so weight decay pulls toward the identity instead of toward zero.

class ZeroCenteredRMSNorm(nn.Module):
    """RMSNorm with scale (1 + weight), weight initialized at 0. Optionally normalizes each
    group of `group_size` features separately (used on the 4 residual streams in Chapter 30)."""
    def __init__(self, width, eps=1e-6, group_size=None):
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(width))
        self.eps, self.group_size = eps, group_size

    def forward(self, x):
        x32 = x.float()
        if self.group_size:
            x32 = x32.unflatten(-1, (-1, self.group_size))
        x32 = x32 * torch.rsqrt(x32.square().mean(-1, keepdim=True) + self.eps)
        if self.group_size:
            x32 = x32.flatten(-2)
        return (x32 * (1.0 + self.weight.float())).type_as(x)

class GatedAttention(nn.Module):
    """Qwen3.5/Flash-Next attention: zero-centered Q/K norms, partial RoPE, GQA, an optional QSA
    indexer, and a sigmoid output gate produced by the same q_proj.  (Your engine: Chapter 29)"""

    def __init__(self, hidden, heads, kv_heads, head_dim, rotary_dim, theta, eps=1e-6, indexer=None):
        super().__init__()
        self.heads, self.kv_heads, self.head_dim = heads, kv_heads, head_dim
        self.rotary_dim, self.theta = rotary_dim, theta
        self.q_proj = nn.Linear(hidden, heads * head_dim * 2, bias=False)    # [query | gate] per head
        self.k_proj = nn.Linear(hidden, kv_heads * head_dim, bias=False)
        self.v_proj = nn.Linear(hidden, kv_heads * head_dim, bias=False)
        self.o_proj = nn.Linear(heads * head_dim, hidden, bias=False)
        self.q_norm = ZeroCenteredRMSNorm(head_dim, eps)
        self.k_norm = ZeroCenteredRMSNorm(head_dim, eps)
        self.indexer = indexer

    def forward(self, x, positions, state=None):
        """x [B, T, D], positions [T] -> (y [B, T, D], new AttentionState).  (Your engine: Chapter 29)"""
        b, t, _ = x.shape
        state = state or AttentionState()
        query, gate = self.q_proj(x).view(b, t, self.heads, 2 * self.head_dim).chunk(2, dim=-1)
        q = self.q_norm(query).transpose(1, 2)
        k = self.k_norm(self.k_proj(x).view(b, t, self.kv_heads, self.head_dim)).transpose(1, 2)
        v = self.v_proj(x).view(b, t, self.kv_heads, self.head_dim).transpose(1, 2)
        cos, sin = rope_cos_sin(positions, self.rotary_dim, self.theta, q.dtype)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)      # only the first rotary_dim features
        allowed, index_keys = None, None
        if self.indexer is not None:
            allowed, index_keys = self.indexer(x, positions, self.rotary_dim, self.theta, state.index_keys)
        state = state.append(k, v)
        state.index_keys = index_keys                                # all raw indexer keys so far
        y = causal_attention(q, state.k, state.v, positions, None, allowed)
        y = y.transpose(1, 2).reshape(b, t, -1) * torch.sigmoid(gate.reshape(b, t, -1))
        return self.o_proj(y), state

AttentionState carries the layer’s per-request cache: K, V and all raw indexer keys (pooling happens on the fly, so a decode step that completes a block sees it immediately). This reference concatenates; in Chapter 30’s engine this is the state that a paged or static cache would replace.

Consistency: the hard requirement

Selection is discrete: a block is in or out. If a full forward and a token-by-token decode compute scores with tiny differences, they can select different blocks and produce very different outputs, and the cache-equivalence test from Chapter 16 fails in a way that looks random. Two rules keep them identical:

1. Score each query against exactly the blocks it can see, the same way in both paths. A decode step at position $p$ must see the same pooled keys a full forward does for row $p$.
2. Break ties the same way. ReLU makes exact ties at 0 common (many blocks score 0). PyTorch’s topk breaks ties differently depending on the length of the vector it searches. Masking unseen blocks with $-\infty$ in a long vector and slicing a short vector to just the visible blocks can choose different zero-scored blocks. The reference therefore groups queries by their number of visible blocks and runs topk over exactly that slice, for every path.

The second rule was found the hard way: the first version of the whole Flash-Next model differed from Hugging Face’s on a few positions, and the cause was exactly this tie-breaking. The milestone test runs prefill of 11 tokens followed by 14 single-token steps against one full forward, with a budget small enough that selection matters.

Build it

pytest tests/test_ch29_sparse.py
python run.py sparse --impl engine

Stretch exercises

1. ★ Add an attention-sink option to topk_attention (always keep the first 4 keys) and measure the error change with flat and peaked scores. Where: topk_attention in engine/sparse.py.
2. ★★ Implement a ring-buffer KV cache for sliding-window layers, storing only the last $w$ positions, and verify a windowed model’s logits against the full-cache version. Where: AttentionState in engine/sparse.py, with position handling in GatedAttention.forward.
3. ★★ Make the indexer’s selection explicit: return the selected block indices [B, T, budget/ratio] instead of a mask, and write attention that gathers only those blocks’ K and V. Check it against the mask version. Where: QSAIndexer.forward and GatedAttention.forward in engine/sparse.py.
4. ★★★ Write a Triton decode kernel for QSA: one program per (sequence, head) that reads the selected block list and streams those blocks’ K/V with the online softmax (Chapter 25’s paged kernel is most of it). Where: new engine/kernels/triton_sparse.py, called by GatedAttention.forward in engine/sparse.py.

Check your understanding

1. Why does top-k attention with exact scores save no computation?
2. Why does Flash-Next’s indexer select blocks rather than tokens?
3. What is the open block, and why is it always readable?
4. Why does sparse attention save bandwidth but not cache memory?
5. Why must full and incremental forwards break top-k ties identically?

Going deeper

• Beltagy et al., Longformer (2020) and Zaheer et al., BigBird (2020) for fixed sparse patterns; Xiao et al., Efficient Streaming Language Models with Attention Sinks (2023).
• DeepSeek-AI, Native Sparse Attention (2025) and the DeepSeek-V3.2 report (DeepSeek Sparse Attention, a lightning indexer with ReLU scores): the closest published relatives of QSA.
• Qiu et al., Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free (2025), the gate used here.
• The Transformers qwen4_exp modeling file, the reference implementation this chapter’s code is checked against.

30. Capstone: Qwen3.8-Flash-Next from scratch

The target

Qwen3.8-Flash-Next is a multimodal model whose language part combines nearly every idea in Part VII. Its model card describes a 125B-parameter backbone with 6B parameters active per token, plus 51B parameters of n-gram memory, 4B parameters of multi-token-prediction heads, and a vision encoder. The text configuration (qwen4_exp_text in Transformers 5.18 and later):

Twenty-eight chapters of preparation cover all but two rows. Before writing any code, the professional habit is an architecture ledger: for every component, the equations, tensor names and shapes, and the reference you’ll check against. The official implementation (transformers/models/qwen4_exp/modeling_qwen4_exp.py) is the ground truth; this chapter’s code was written from it and is tested against it.

What’s already built

• Gated DeltaNet layers with their convolution and recurrent state, exactly Chapter 28’s GatedDeltaNet, with the same parameter names.
• Gated attention with zero-centered Q/K norms, partial RoPE, two KV heads and the QSA indexer: Chapter 29’s GatedAttention.
• MoE with 512 experts, top-10, renormalized softmax router and a sigmoid-gated shared expert: Chapter 27’s SparseMoeBlock.
• Loading sharded safetensors, meta-device construction and expert stacking: Chapters 9, 18 and 27.

Two things are new.

Hyper-connections: four residual streams

Every model so far had one residual stream: x = x + sublayer(norm(x)) (Chapter 23). Flash-Next keeps four parallel streams of width 2,560, and each sublayer learns how to read from them and how to write back (hyper-connections, Zhu et al., 2024). The embedding initializes all four streams to the same vector, and the streams drift apart as layers write to them differently.

Around each sublayer (attention or linear attention, then the MoE), a GatedResidual module:

1. Normalizes each stream separately: a zero-centered RMSNorm over groups of 2,560 features. This replaces the usual pre-norm; there’s no other norm before the sublayer.
2. Reads: computes per-feature weights for each stream through a low-rank bottleneck, $w = \sigma\big(W_\text{up}, \operatorname{SiLU}(W_\text{down}, \hat s / 4)\big)$, with rank 320. The sublayer’s input is the mean over streams of $w \odot \hat s$.
3. Writes: computes one scalar per stream, $\gamma = 2\sigma(W_\text{inject}, \hat s / 4) \in (0, 2)$, and adds $\gamma_i \cdot \text{output}$ to stream $i$.

At the end, a final GatedResidual without the write part (combine=False) reads one vector from the four streams for the LM head. There’s no final norm.

class GatedResidual(nn.Module):
    """Hyper-connection around one sublayer.  (Your engine: Chapter 30)

    streams [B, T, hc*D] -> normalize each stream (grouped zero-centered RMSNorm);
    read weights  w = sigmoid(up(silu(down(normed) / hc)))      [B, T, hc, D]
    sublayer input = mean over streams of w * normed            [B, T, D]
    write weights  = 2 * sigmoid(inject(normed) / hc)           [B, T, hc]
    After the sublayer: streams + write_weights[..., None] * output[..., None, :].
    """

    def __init__(self, hidden, hc, rank, eps=1e-6, combine=True):
        super().__init__()
        self.hc, self.hidden = hc, hidden
        self.hc_norm = ZeroCenteredRMSNorm(hc * hidden, eps, group_size=hidden)
        self.input_mix_weight_down = nn.Linear(hc * hidden, rank, bias=False)
        self.input_mix_weight_up = nn.Linear(rank, hc * hidden, bias=False)
        self.block_inject_weight = nn.Linear(hc * hidden, hc, bias=False) if combine else None

    def read(self, streams):
        """(Your engine: Chapter 30)"""
        normed = self.hc_norm(streams)
        weights = torch.sigmoid(self.input_mix_weight_up(F.silu(self.input_mix_weight_down(normed) / self.hc)))
        mixed = (weights.unflatten(-1, (self.hc, self.hidden)) * normed.unflatten(-1, (self.hc, self.hidden))).mean(-2)
        if self.block_inject_weight is None:
            return mixed, None
        return mixed, 2 * torch.sigmoid(self.block_inject_weight(normed) / self.hc)

    @staticmethod
    def write(streams, output, inject):
        return streams + (output.unsqueeze(-2) * inject.unsqueeze(-1)).flatten(-2)

Why bother? With one stream, each layer’s output is added with weight 1, and depth works through a single shared channel. Multiple streams with learned read and write weights let the network keep some information untouched by later layers and vary how strongly each layer contributes, which in the hyper-connection papers improves training stability and quality at a cost of a few small matrices per layer. For inference, it means the hidden state is four times wider between layers: 20 KiB per token instead of 5 KiB in BF16, which matters for activation memory during prefill and not at all for the KV cache.

The n-gram memory

The second new piece is a huge lookup table indexed by the last few tokens. The idea: many next-token facts depend only on the immediately preceding tokens (“New York” → “City”), and a model shouldn’t spend attention and MLP compute rediscovering them. A hashed table can store an embedding for every frequent bigram and trigram, and the model looks it up in $O(1)$.

Hashing n-grams

For position $t$, the bigram is $(x_{t-1}, x_t)$ and the trigram $(x_{t-2}, x_{t-1}, x_t)$. There are $248{,}320^3$ possible trigrams, far too many to store, so each n-gram is hashed into a table of about 20 million rows:

$$ h = \Big(\bigoplus_{i=0}^{n-1} x_{t-i} \cdot m_i\Big) \bmod p , $$

with odd multipliers $m_i$ derived from a seed with splitmix64, XOR ($\oplus$) to combine, and a prime table size $p$. Hash collisions are inevitable, so each n-gram order uses 8 independent hash heads, each with its own prime size (the 16 smallest primes above 20 million), each returning a 160-dimensional row. Concatenated, the 16 rows form one 2,560-dimensional embedding. Two n-grams colliding in one head are very unlikely to collide in all eight.

MASK64 = (1 << 64) - 1


def splitmix64(value):
    value = (value + 0x9E3779B97F4A7C15) & MASK64
    value = ((value ^ (value >> 30)) * 0xBF58476D1CE4E5B9) & MASK64
    value = ((value ^ (value >> 27)) * 0x94D049BB133111EB) & MASK64
    return (value ^ (value >> 31)) & MASK64


def layer_multipliers(vocab, ngram_size, ple_index, seed):
    """Odd multipliers, small enough that token_id * multiplier never overflows int64."""
    half = max(1, (((1 << 63) - 1) // max(vocab, 1)) // 2)
    base = seed + 10007 * ple_index
    return [2 * (splitmix64((base + 0x9E3779B97F4A7C15 * (i + 1)) & MASK64) % half) + 1 for i in range(ngram_size)]


def next_primes(start, count):
    """The `count` smallest primes greater than `start` (one distinct table size per hash head)."""
    def is_prime(n):
        if n < 2 or (n % 2 == 0 and n != 2):
            return n == 2
        return all(n % d for d in range(3, math.isqrt(n) + 1, 2))
    primes, n = [], start
    while len(primes) < count:
        n += 1
        if is_prime(n):
            primes.append(n)
    return primes

Two details that a careless implementation gets wrong:

• Document boundaries. An n-gram must not span an end-of-sequence token: after an EOS, the “previous token” positions read as EOS. shifted computes, for every position, how far back the current document starts, and substitutes EOS beyond it.
• Carried history. In decode, the n-gram of the new token needs the previous two tokens, which belong to earlier calls. The layer’s state carries them.

class NGramEmbedding(nn.Module):
    """Hashed bigram and trigram embeddings, heads_per_ngram independent hashes per order.  (Your engine: Chapter 30)

    The n-gram ending at token t is (x[t-n+1], ..., x[t]); its id for one head is
    (XOR over i of x[t-i] * m_i) mod p_head, offset into one big table. N-grams never cross an
    EOS: positions before the most recent EOS read as EOS.
    """

    def __init__(self, cfg, ple_index):
        super().__init__()
        self.n, self.heads_per = cfg.ngram_size, cfg.heads_per_ngram
        self.heads = (cfg.ngram_size - 1) * cfg.heads_per_ngram
        self.eos = cfg.eos_token_id
        sizes = next_primes(cfg.ngram_vocab_size_base - 1, (ple_index + 1) * self.heads)[ple_index * self.heads:]
        self._sizes = sizes
        offsets = [sum(sizes[:i]) for i in range(len(sizes))]
        padded = -(-sum(sizes) // cfg.make_ngram_vocab_size_divisible_by) * cfg.make_ngram_vocab_size_divisible_by
        self.register_buffer("layer_multipliers", torch.tensor(layer_multipliers(cfg.vocab_size, self.n, ple_index, cfg.seed)))
        self.register_buffer("ngram_heads_vocab_sizes", torch.tensor(sizes))
        self.register_buffer("ngram_heads_offsets", torch.tensor(offsets))
        self.ngram_embedding = nn.Embedding(padded, cfg.ple_embed_dim // self.heads)

    def shifted(self, history, shift):
        """history[t - shift], or EOS when that position is before the current document.  (Your engine: Chapter 30)"""
        if shift == 0:
            return history
        b, length = history.shape
        pos = torch.arange(length, device=history.device)
        eos_at = torch.where(history == self.eos, pos, -1)
        last_eos_before = torch.cat((eos_at.new_full((b, 1), -1), eos_at.cummax(1).values[:, :-1]), dim=1)
        in_document = pos - (last_eos_before + 1)
        source = (pos - shift).clamp_min(0).expand(b, -1)
        valid = (in_document >= shift) & ((pos - shift) >= 0)
        return torch.where(valid, history.gather(1, source), torch.full_like(history, self.eos))

    def forward(self, ids, context=None):
        """ids [B, T]; context = the previous n-1 token ids (EOS at the start). Returns (emb [B, T, E], new context).  (Your engine: Chapter 30)"""
        if context is None:
            context = ids.new_full((ids.shape[0], self.n - 1), self.eos)
        history = torch.cat((context, ids.long()), dim=1)
        shifted = [self.shifted(history, s) for s in range(self.n)]
        blocks = []
        for order in range(2, self.n + 1):
            mixed = shifted[0] * self.layer_multipliers[0]
            for position in range(1, order):
                mixed = mixed ^ (shifted[position] * self.layer_multipliers[position])
            first = (order - 2) * self.heads_per
            sizes = self.ngram_heads_vocab_sizes[first:first + self.heads_per]
            blocks.append(mixed[..., None] % sizes + self.ngram_heads_offsets[first:first + self.heads_per])
        index = torch.cat(blocks, dim=-1)[:, -ids.shape[1]:]
        table = self.ngram_embedding.weight
        rows = table[index.to(table.device)].to(ids.device)        # the table may live on the host
        return rows.flatten(-2), history[:, -(self.n - 1):]

Injecting it: per-layer embeddings

The lookup result is injected into the residual streams at layer 2 by a PLELayer (“per-layer embedding”):

1. A key projection of the n-gram embedding (one per stream) and a query from the normalized streams give a gate per stream: their dot product over $\sqrt{d}$, passed through a signed square root (to tame large values while keeping the sign) and a sigmoid.
2. A value projection of the n-gram embedding, scaled by each stream’s gate, is the addition to that stream.
3. A dilated causal convolution (kernel 4, dilation 3, depthwise) over recent gated values adds local context, with its own carried state of the last 9 positions.

class PLELayer(nn.Module):
    """Per-layer embedding: gate the n-gram value into each residual stream, then add a dilated
    causal depthwise convolution over recent gated values.  (Your engine: Chapter 30)"""

    def __init__(self, cfg, ple_index):
        super().__init__()
        d, hc, e = cfg.hidden_size, cfg.hc_count, cfg.ple_embed_dim
        self.hidden, self.hc, self.dilation = d, hc, cfg.ngram_size
        self.state_len = (cfg.ple_conv_kernel_size - 1) * cfg.ngram_size
        self.ple_embedding = NGramEmbedding(cfg, ple_index)
        self.key_proj = nn.Linear(e, hc * d, bias=False)
        self.value_proj = nn.Linear(e, d, bias=False)
        self.norm_key = ZeroCenteredRMSNorm(hc * d, cfg.rms_norm_eps, group_size=d)
        self.norm_query = ZeroCenteredRMSNorm(hc * d, cfg.rms_norm_eps, group_size=d)
        self.norm_conv = ZeroCenteredRMSNorm(hc * d, cfg.rms_norm_eps, group_size=d)
        self.conv1d = nn.Conv1d(hc * d, hc * d, cfg.ple_conv_kernel_size, groups=hc * d,
                                dilation=cfg.ngram_size, bias=False)

    def forward(self, streams, ids, state=None):
        """state = (ngram context ids, conv history). Returns (addition to streams, new state).  (Your engine: Chapter 30)"""
        context, conv_state = state if state is not None else (None, None)
        embedding, context = self.ple_embedding(ids, context)
        key = self.norm_key(self.key_proj(embedding)).unflatten(-1, (self.hc, self.hidden))
        value = self.value_proj(embedding)
        query = self.norm_query(streams).unflatten(-1, (self.hc, self.hidden))
        gate = (key * query).sum(-1, keepdim=True) / math.sqrt(self.hidden)
        gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign()              # signed square root
        gated = (torch.sigmoid(gate) * value.unsqueeze(-2)).flatten(-2)      # [B, T, hc*D]
        x = self.norm_conv(gated).transpose(1, 2)
        if conv_state is None:
            conv_state = x.new_zeros(x.shape[0], x.shape[1], self.state_len)
        joined = torch.cat((conv_state, x), dim=-1)
        conv = F.silu(self.conv1d(joined)).transpose(1, 2)
        return gated + conv, (context, joined[..., -self.state_len:])

Why it’s cheap at inference: each token reads 16 rows of 160 BF16 values, about 5 KB, from a table of 51 billion parameters. The table’s size is a storage problem, not a bandwidth problem, so it can live in host memory or even on an SSD. Closely related published work: DeepSeek’s Engram conditional memory (2026).

The layer and the model

class FlashNextLayer(nn.Module):
    def __init__(self, cfg, index):
        super().__init__()
        self.kind = cfg.layer_types[index]
        if self.kind == "linear_attention":
            self.linear_attn = GatedDeltaNet(cfg.hidden_size, cfg.linear_num_key_heads, cfg.linear_num_value_heads,
                                             cfg.linear_key_head_dim, cfg.linear_value_head_dim,
                                             cfg.linear_conv_kernel_dim, cfg.rms_norm_eps, cfg.output_gate_type)
        else:
            indexer = None
            if cfg.indexer_n_heads:
                indexer = QSAIndexer(cfg.hidden_size, cfg.indexer_n_heads, cfg.indexer_head_dim,
                                     cfg.indexer_budget, cfg.indexer_compress_ratio, cfg.rms_norm_eps)
            self.self_attn = GatedAttention(cfg.hidden_size, cfg.num_attention_heads, cfg.num_key_value_heads,
                                            cfg.head_dim, cfg.rotary_dim, cfg.rope_theta, cfg.rms_norm_eps, indexer)
        self.mlp = SparseMoeBlock(cfg.hidden_size, cfg.num_experts, cfg.num_experts_per_tok, cfg.moe_intermediate_size,
                                  cfg.shared_expert_intermediate_size, cfg.norm_topk_prob)
        ple_index = cfg.ple_layer_ids.index(index + 1) if (index + 1) in cfg.ple_layer_ids else None
        self.ple = PLELayer(cfg, ple_index) if ple_index is not None else None
        self.attn_hyper_connection = GatedResidual(cfg.hidden_size, cfg.hc_count, cfg.hc_lowrank, cfg.rms_norm_eps)
        self.mlp_hyper_connection = GatedResidual(cfg.hidden_size, cfg.hc_count, cfg.hc_lowrank, cfg.rms_norm_eps)

    def forward(self, streams, ids, positions, state):
        """state = {"mixer": layer state, "ple": PLE state}. Returns (streams, new state).  (Your engine: Chapter 30)"""
        new_state = {}
        if self.ple is not None:
            addition, new_state["ple"] = self.ple(streams, ids, state.get("ple"))
            streams = streams + addition
        x, inject = self.attn_hyper_connection.read(streams)
        if self.kind == "linear_attention":
            y, new_state["mixer"] = self.linear_attn(x, state.get("mixer"))
        else:
            y, new_state["mixer"] = self.self_attn(x, positions, state.get("mixer"))
        streams = GatedResidual.write(streams, y, inject)
        x, inject = self.mlp_hyper_connection.read(streams)
        streams = GatedResidual.write(streams, self.mlp(x), inject)
        return streams, new_state

Each layer: inject the n-gram memory (layer 2 only), read from the streams, run the mixer (Gated DeltaNet or gated sparse attention), write, read again, run the MoE, write.

class HybridState:
    """Everything one sequence carries between calls. Updates return a new object, so keeping the
    old one is a free snapshot: speculative decoding rolls back by simply not adopting the new state."""
    def __init__(self, layers=None, length=0):
        self.layers = layers or []
        self.length = length


class HybridCache:
    """Adapts the functional HybridState to the engines' mutable-cache convention (Chapters 18-19):
    model(ids, cache) updates cache.state in place and returns logits only."""
    def __init__(self):
        self.state = None

    @property
    def length(self):
        return 0 if self.state is None else self.state.length

    def snapshot(self):
        return self.state                     # states are never modified in place: this is a free copy

    def restore(self, state):
        self.state = state

    def truncate(self, length):
        raise NotImplementedError("A recurrent state cannot be truncated: restore a snapshot instead")


class FlashNextBackbone(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
        self.layers = nn.ModuleList(FlashNextLayer(cfg, i) for i in range(cfg.num_hidden_layers))
        self.hyper_connection_mixer = GatedResidual(cfg.hidden_size, cfg.hc_count, cfg.hc_lowrank,
                                                    cfg.rms_norm_eps, combine=False)


class FlashNext(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.model = FlashNextBackbone(cfg)
        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)

    @property
    def context_limit(self):
        return self.cfg.max_position_embeddings

    def new_cache(self, batch, capacity):
        if capacity > self.context_limit:
            raise ValueError("Requested cache exceeds the context limit")
        return HybridCache()

    def forward(self, ids, state=None):
        """ids [B, T] continuing `state` -> (logits, new HybridState); with a HybridCache, logits only."""
        if isinstance(state, HybridCache):
            logits, state.state = self.step(ids, state.state)
            return logits
        return self.step(ids, state)

    def step(self, ids, state=None):
        """ids [B, T] continuing `state` -> (logits [B, T, V], new HybridState).  (Your engine: Chapter 30)"""
        state = state or HybridState([{} for _ in self.model.layers])
        positions = torch.arange(state.length, state.length + ids.shape[1], device=ids.device)
        streams = self.model.embed_tokens(ids).repeat(1, 1, self.cfg.hc_count)    # every stream starts equal
        new_layers = []
        for layer, layer_state in zip(self.model.layers, state.layers):
            streams, layer_state = layer(streams, ids, positions, layer_state)
            new_layers.append(layer_state)
        hidden, _ = self.model.hyper_connection_mixer.read(streams)               # collapse streams
        return self.lm_head(hidden), HybridState(new_layers, state.length + ids.shape[1])

The state of a request

One Flash-Next request carries four kinds of state, all owned by HybridState:

step never modifies a state in place: it returns a new HybridState. Keeping the old one is therefore a free snapshot, which is exactly what speculative decoding needs, since the recurrent state can’t be truncated (Chapter 28). The milestone test runs a branch from a state, discards it, runs another from the same state, and checks they’re identical. For the engines of Chapters 18-19, new_cache returns a HybridCache, a small mutable wrapper with snapshot and restore, so LLM.stream serves Flash-Next unchanged.

Exactly equal to the official implementation

The tests build a tiny random Flash-Next with Transformers’ Qwen4ExpForCausalLM: 8 layers, 2 of them attention with a QSA budget small enough to matter, 8 experts, two PLE layers with sharded n-gram tables. Every norm and convolution weight is randomized, since zero-initialized ones would hide bugs. The tests save it as a checkpoint and load it with your loader:

max abs logit difference: 1.49e-08   (logits of magnitude 0.12; FP32)

and check that:

• a full forward over 2 × 41 tokens with an EOS in the middle matches the reference;
• prefill of 13 tokens, a chunk of 16, then 11 single-token steps equals one full forward;
• the state is a free snapshot;
• the real configuration’s memory plan has the published sizes;
• the n-gram tables memory-mapped from disk give bit-identical logits, and 4-bit experts stay close;
• your Chapter 18 LLM engine generates the same tokens as the model’s own greedy loop.

Getting there surfaced real bugs, each now a lesson in an earlier chapter: a transposed triangle in the chunked delta rule (28) and topk tie-breaking under ReLU zeros (29). A third, a config field renamed between library versions (27), would have broken any loader that trusted defaults.

Loading the real checkpoint

@torch.no_grad()
def load_flashnext(directory, device="cpu", dtype=torch.bfloat16, ngram_device="cpu", expert_bits=None,
                   ngram_mmap=False, group_size=128):
    """Load the text model from a local snapshot of the official checkpoint.  (Your engine: Chapter 30)

    Renames multimodal prefixes, skips the vision tower and MTP head, and stacks per-expert
    tensors. expert_bits=4 or 8 quantizes each layer's experts as soon as they are read;
    ngram_mmap=True leaves the n-gram tables on disk (MappedRows); otherwise they are loaded to
    `ngram_device` (host memory by default: lookups touch a few rows per token).
    """
    raw = json.loads((Path(directory) / "config.json").read_text())
    cfg = FlashNextConfig.from_hf(raw)
    with torch.device("meta"):
        model = FlashNext(cfg)
    for layer in model.model.layers:                    # never allocate what will be stored differently
        if expert_bits:
            layer.mlp.experts = nn.Module()
        if ngram_mmap and layer.ple is not None:
            del layer.ple.ple_embedding.ngram_embedding
    model = model.to_empty(device=device).to(dtype)
    for index, layer in enumerate(model.model.layers):
        if layer.ple is not None:                       # buffers are derived, not learned: recompute them
            fresh = _ngram_buffers(cfg, cfg.ple_layer_ids.index(index + 1))
            for name, value in fresh.items():
                setattr(layer.ple.ple_embedding, name, value.to(device))
            if not ngram_mmap:
                layer.ple.ple_embedding.ngram_embedding.to(ngram_device)
    params = dict(model.named_parameters())
    rename = lambda name: name.replace("model.language_model.", "model.")
    is_ngram = lambda name: ".ngram_embedding." in name
    experts, ngram_shards = {}, {}
    skip = (lambda name: is_ngram(name)) if ngram_mmap else None
    for name, value in snapshot_tensors(directory, skip=skip):
        name = rename(name)
        if name.startswith(("model.visual.", "mtp.", "model.mtp")) or name.endswith(
                ("layer_multipliers", "ngram_heads_vocab_sizes", "ngram_heads_offsets")):
            continue
        if ".mlp.experts." in name:
            parts = name.split(".")
            layer = int(parts[2])
            experts.setdefault(layer, {})[(int(parts[5]), parts[6])] = value
            if len(experts[layer]) == 3 * cfg.num_experts:          # a whole layer has arrived
                _install_experts(model, cfg, layer, experts.pop(layer), expert_bits, group_size, device, dtype)
            continue
        if is_ngram(name):
            prefix, _, shard = name.partition(".shard_")
            ngram_shards.setdefault(prefix.removesuffix(".weight"), {})[int(shard.split(".")[0]) if shard else 0] = value
            continue
        if name not in params:
            raise ValueError(f"Unexpected tensor {name}")
        assign(params[name], value, name)
    if experts:
        raise ValueError(f"Incomplete experts for layers {sorted(experts)}")
    if ngram_mmap:
        for original, path in tensor_files(directory).items():
            name = rename(original)
            if is_ngram(name):
                prefix, _, shard = name.partition(".shard_")
                ngram_shards.setdefault(prefix.removesuffix(".weight"), {})[
                    int(shard.split(".")[0]) if shard else 0] = map_tensor(path, original)
    for prefix, shards in ngram_shards.items():
        ordered = [shards[i] for i in sorted(shards)]
        owner = model.get_submodule(prefix.rsplit(".", 1)[0])
        if ngram_mmap:
            owner.ngram_embedding = MappedRows(ordered, dtype)
        else:
            assign(params[prefix + ".weight"], torch.cat(ordered, dim=0), prefix)
    return model.eval()


def _ngram_buffers(cfg, ple_index):
    with torch.device("meta"):
        probe = NGramEmbedding(cfg, ple_index)          # meta: the big table is not allocated
    return {"layer_multipliers": torch.tensor(layer_multipliers(cfg.vocab_size, cfg.ngram_size, ple_index, cfg.seed)),
            "ngram_heads_vocab_sizes": torch.tensor(probe._sizes),
            "ngram_heads_offsets": torch.tensor([sum(probe._sizes[:i]) for i in range(len(probe._sizes))])}


def _install_experts(model, cfg, layer, tensors, bits, group_size, device, dtype):
    gate_up, down = stack_expert_tensors(tensors, cfg.num_experts, cfg.moe_intermediate_size)
    block = model.model.layers[layer].mlp
    if bits:
        block.experts = QuantizedExperts(gate_up.to(device, dtype), down.to(device, dtype), bits, group_size)
    else:
        assign(block.experts.gate_up_proj, gate_up, "gate_up_proj")
        assign(block.experts.down_proj, down, "down_proj")

The loader streams the safetensors shards one at a time and handles the checkpoint’s quirks: the multimodal prefix model.language_model. is renamed, the vision tower and MTP weights are skipped, per-expert tensors are stacked, and n-gram tables stored as shards are concatenated. Two options make it fit on one machine:

class QuantizedExperts(nn.Module):
    """Routed experts stored as groupwise INT4/INT8 codes (Chapter 20). An expert is dequantized
    only when a token is routed to it, so a BF16 copy of all experts never exists."""

    def __init__(self, gate_up, down, bits=4, group_size=128):
        super().__init__()
        self.bits, self.dtype, self.shapes, self.groups = bits, gate_up.dtype, {}, {}
        for name, w in (("gate_up", gate_up), ("down", down)):
            experts, rows, cols = w.shape
            codes, scales = quantize_groupwise(w.reshape(experts * rows, cols), bits, group_size)
            codes = pack_int4(codes).view(experts, -1) if bits == 4 else codes.view(experts, rows, cols)
            self.register_buffer(f"{name}_codes", codes)
            self.register_buffer(f"{name}_scales", scales.view(experts, rows, -1).to(torch.float16))
            self.shapes[name], self.groups[name] = (rows, cols), group_size
        self.num_experts = gate_up.shape[0]

    def weight(self, name, e):
        codes = getattr(self, f"{name}_codes")[e]
        if self.bits == 4:
            codes = unpack_int4(codes, self.shapes[name])
        scales = getattr(self, f"{name}_scales")[e].float()
        return dequantize_groupwise(codes, scales, self.groups[name]).to(self.dtype)

    def expert(self, e, x):
        gate, up = F.linear(x, self.weight("gate_up", e)).chunk(2, dim=-1)
        return F.linear(F.silu(gate) * up, self.weight("down", e))

    forward_loop = Experts.forward_loop
    forward_grouped = Experts.forward_grouped


class MappedRows:
    """An embedding table left on disk: rows are read through memory maps when looked up, and the
    operating system's page cache keeps the hot ones in RAM. Lookups touch a few rows per token,
    so the 95 GiB n-gram table never has to fit in memory."""

    def __init__(self, shards, dtype):
        self.shards, self.dtype, self.device = shards, dtype, torch.device("cpu")
        sizes = torch.tensor([0] + [t.shape[0] for t in shards])
        self.starts = sizes.cumsum(0)

    @property
    def weight(self):
        return self

    def __getitem__(self, index):
        flat = index.reshape(-1).cpu()
        shard = torch.searchsorted(self.starts, flat, right=True) - 1
        out = torch.empty(flat.numel(), self.shards[0].shape[1], dtype=self.dtype)
        for s in shard.unique().tolist():
            hit = shard == s
            out[hit] = self.shards[s][flat[hit] - self.starts[s]].to(self.dtype)
        return out.view(*index.shape, -1)

• expert_bits=4 replaces each layer’s experts with QuantizedExperts as soon as that layer’s 1,536 expert tensors have arrived: the BF16 experts of more than one layer never exist at once, and the model is built on the meta device so their full-size placeholders are never allocated. An expert is dequantized only when a token is routed to it (a W4A16 grouped kernel, Chapters 20 and 27, is the fast version).
• ngram_mmap=True never reads the n-gram shards into memory. MappedRows wraps memory maps of the safetensors files; a lookup reads a few pages and the operating system caches the hot ones.

The memory plan

def memory_plan(cfg, weight_bits=16, expert_bits=4, ngram_bits=16, context=32768, kv_bits=16):
    """Rough byte budget for the text model: what must be resident, and where."""
    d, l = cfg.hidden_size, cfg.num_hidden_layers
    n_attn = sum(t != "linear_attention" for t in cfg.layer_types)
    n_lin = l - n_attn
    expert = cfg.num_experts * 3 * d * cfg.moe_intermediate_size * l
    shared = 3 * d * cfg.shared_expert_intermediate_size * l + l * (cfg.num_experts + 1) * d
    kd, vd = cfg.linear_num_key_heads * cfg.linear_key_head_dim, cfg.linear_num_value_heads * cfg.linear_value_head_dim
    linear = n_lin * (d * (2 * kd + vd) + d * vd + vd * d + 2 * d * cfg.linear_num_value_heads)
    attn = n_attn * (d * cfg.num_attention_heads * cfg.head_dim * 3 + 2 * d * cfg.num_key_value_heads * cfg.head_dim)
    hc = (2 * l + 1) * 2 * cfg.hc_count * d * cfg.hc_lowrank
    embed = 2 * cfg.vocab_size * d
    ngram_rows = sum(next_primes(cfg.ngram_vocab_size_base - 1, (cfg.ngram_size - 1) * cfg.heads_per_ngram * len(cfg.ple_layer_ids)))
    ngram = ngram_rows * cfg.ple_embed_dim // ((cfg.ngram_size - 1) * cfg.heads_per_ngram)
    kv = 2 * n_attn * context * cfg.num_key_value_heads * cfg.head_dim * kv_bits // 8
    recurrent = n_lin * cfg.linear_num_value_heads * cfg.linear_key_head_dim * cfg.linear_value_head_dim * 4
    gib = 1024 ** 3
    return {
        "routed_experts_GiB": expert * expert_bits / 8 / gib,
        "dense_weights_GiB": (shared + linear + attn + hc + embed) * weight_bits / 8 / gib,
        "ngram_tables_GiB": ngram * ngram_bits / 8 / gib,
        "kv_cache_GiB_per_sequence": kv / gib,
        "recurrent_state_MiB_per_sequence": recurrent / 1024 ** 2,
        "active_params_per_token_B": (cfg.num_experts_per_tok * 3 * d * cfg.moe_intermediate_size * l
                                      + shared + linear + attn + hc) / 1e9,
    }

python run.py flashnext

{"expert_bits": 16, "context": 32768,  "routed_experts_GiB": 225.0,  "dense_weights_GiB": 9.11, "ngram_tables_GiB": 95.37, "kv_cache_GiB_per_sequence": 0.75, "recurrent_state_MiB_per_sequence": 108.0, "active_params_per_token_B": 5.98}
{"expert_bits": 8,  "context": 32768,  "routed_experts_GiB": 112.5, ...}
{"expert_bits": 4,  "context": 32768,  "routed_experts_GiB": 56.25, ...}
{"expert_bits": 4,  "context": 262144, "routed_experts_GiB": 56.25, ..., "kv_cache_GiB_per_sequence": 6.0, ...}

Putting it together for three machines (estimates from the plan, not measurements):

The ceilings use Chapter 10’s rule with the bytes read per token: 2.36B active expert parameters at 4 bits (1.2 GB) plus 4.25B other active parameters including the LM head at 2 bytes (8.5 GB), about 9.7 GB per token in BF16, or 5.4 GB with INT8 dense weights. The dense part, not the experts, dominates decode traffic once the experts are 4-bit: quantize it next.

Run it

hf download Qwen/Qwen3.8-Flash-Next --local-dir models/Qwen3.8-Flash-Next      # ~360 GB on disk
python run.py chat --model-dir models/Qwen3.8-Flash-Next --expert-bits 4 --ngram-mmap \
    --prompt "Explain hyper-connections in two sentences." --new-tokens 128

Important

This book’s code was validated against the official implementation on tiny random checkpoints with the official architecture (Appendix F); the full checkpoint was not run in the validation environment. On a real machine, climb Chapter 18’s ladder of evidence again: compare your logits with Transformers’ on a few hundred tokens in BF16, layer by layer if they differ, before trusting generations.

Multi-token prediction

The checkpoint also ships about 4B parameters of multi-token prediction (MTP) heads, which both your loader and Transformers skip. In the DeepSeek-V3 style, an MTP module takes the final hidden state at position $t$ and the embedding of token $t+1$ and predicts token $t+2$, reusing the model’s embedding and head. At inference, it’s a built-in draft model for speculative decoding (Chapter 26): one cheap extra module proposes the next token, and the main model verifies. With HybridState snapshots, your speculative loop needs only one change: instead of truncating caches after a rejection, restore the snapshot and re-run the accepted tokens. That’s the first stretch exercise.

Build it

pytest tests/test_ch30_flashnext.py
python run.py flashnext --impl engine

Stretch exercises

1. ★★ Speculative decoding for Flash-Next: change speculative_generate to snapshot and restore HybridCache instead of truncating, and use a smaller model with the same tokenizer as the draft. Verify greedy outputs are unchanged. Where: speculative_generate in engine/speculative.py, using HybridCache.snapshot / restore in engine/flashnext.py.
2. ★★ Quantize the dense weights to INT8 (attention, linear attention, shared experts, LM head) and measure the logit error and decode speed against the plan’s prediction. Where: weight installation in load_flashnext in engine/flashnext.py, using engine.quant.
3. ★★★ Load and use the MTP head as a draft: read its weights (prefix mtp.) and implement its forward from the configuration and the DeepSeek-V3 report’s description, then measure acceptance rates. Where: add an MTP module and load its weights in engine/flashnext.py; call it from engine/speculative.py.
4. ★★★ Expert offloading for a 24 GB GPU: keep QuantizedExperts codes in pinned host memory and copy only the routed experts per layer, overlapping the copy for layer $\ell+1$ with compute for layer $\ell$ using CUDA streams. Report tokens/s against the PCIe bound. Where: QuantizedExperts and expert installation in engine/flashnext.py.
5. ★★★ Continuous batching for a hybrid model: give each request a HybridState, batch the linear-attention layers (every state has the same shape) and the attention layers (ragged KV), and verify solo-equivalence as in Chapter 24. Where: new experiments/hybrid_batching.py, adapting engine.scheduler.ContinuousBatchingEngine for engine.flashnext.HybridState.

Check your understanding

1. Which Flash-Next components come from which earlier chapters, and which two are new?
2. How do the read and write weights of a hyper-connection differ in shape and range?
3. Why does the n-gram memory use 8 hash heads with different prime sizes?
4. Why can a 51B-parameter table live on disk without slowing decode much?
5. Why is a functional (never modified in place) state convenient for speculative decoding?
6. After quantizing the experts to 4 bits, what dominates the bytes read per decode token?

Going deeper

• The Qwen3.8-Flash-Next model card and configuration, and modeling_qwen4_exp.py in Transformers 5.18+.
• Zhu et al., Hyper-Connections (2024), and the follow-ups on manifold-constrained hyper-connections; DeepSeek-AI’s Engram (2026) for hashed n-gram memory; the DeepSeek-V3 technical report for MTP.
• The Qwen3-Next and Qwen3.5 model cards, the direct ancestors of this architecture (3:1 Gated DeltaNet / gated attention, sigmoid-gated shared expert).
• vLLM and SGLang’s model files for Qwen3-Next, to see how production engines batch a hybrid model’s states.

31. Engine v2: one core for every request

Five good pieces that don’t fit together

Look at what Parts IV-VI gave you. Chapter 19’s FastDecoder removes launch overhead, but serves one request. Chapter 24’s ContinuousBatchingEngine serves many, but gives each request a fixed-size row of a static cache and prefills one request’s chunk per forward call. Chapter 25’s PagedKVCache packs caches without waste and shares prefixes, but nothing schedules requests onto it. Chapter 26’s speculative decoding handles one request with a contiguous cache. Each piece passed its tests. None of them can be combined with the others without rewriting the seams between them.

That’s the situation every serving engine reaches. vLLM’s first version grew one feature at a time and was rewritten in 2025 as “V1” around a simpler core; SGLang, TensorRT-LLM and LMDeploy converged on similar designs. This chapter builds that core. The next twelve chapters plug the rest of a production server into it: fast kernels, CUDA graphs, a full sampler, constrained decoding, a tokenizer, an HTTP API, speculative decoding, real checkpoint formats, other hardware, several GPUs and more model families.

The design rests on three decisions:

1. One kind of work. A request’s state is its token list (prompt, then output) and one number, num_computed_tokens: how many of those tokens already have keys and values in the cache. The scheduler’s only output is a list of (request, n): compute the next n tokens of this request.
2. One batch layout. All scheduled tokens of all requests are concatenated into one flat sequence of $N$ tokens. Matmuls see [N, D]; attention gets a small metadata object that says where each request’s rows start and where its blocks live.
3. One memory pool. Every layer’s K and V live in a pool of fixed-size blocks shared by all requests, with reference counts and a prefix cache keyed by hash chains.

One kind of work

With num_computed_tokens as the only state, every situation an engine meets is a value of n: