a mastery course · one path · down the stack

Learn to make TPUs fast.

This page is the stack. Scroll and you descend it: four layers of compiler land on paper, the gap, then the shaded depths where kernels live. Each layer is a chapter with assigned mastery work; the stages are pinned where they belong; your progress rides on every entry.

start the path → chapter 1 the repo

0/6 gates passed · your progress 0/15 chapters

01 Read any IR on sight · narrate unfamiliar dumps, decode any dot_general
02 Write kernels that beat XLA · derived, differential-tested, benchmarked
03 Treat the cluster as one kernel · communication hidden behind compute
EX·03 streaming attention · a chapter 13 instrument, live
HBMQQKK0K1K2K3VV0V1V2V3VMEMqQ (resident)kAK0vAV0kBK1vBV1CARRIED STATEm : max(m, rowmax S1)l : l·e^(m₋−m) + rowsum e^(S1−m)acc : acc·e^(m₋−m) + e^(S1−m)·V1MXU · S = Q·K1ᵀ, then acc

03/06step 2 · one pass over K1, V1: rescale (m, l, acc), never materializing scores to HBM
scrub it: the carried (m, l, acc) state rescales exactly; the score matrix never exists
↓ the descent · nine layers · six stages pinned along the way
ch 01

JAX / PyTorch

Math on whole arrays. No notion of hardware.

ch 02

jaxpr

The traced program: a flat list of primitive ops.

ch 03

StableHLO

The portable tensor IR every framework lowers into.

ch 12 The IR stack
ch 04

XLA

Fuses along dataflow edges. Cannot restructure an algorithm.

ch 05

the gap

XLA stops at fusion. Pallas asks for everything by hand. Between them: nothing. Stage 3 teaches you to cross on foot.

ch 13 The priesthood kernels
ch 06

Pallas

Hand choreography: BlockSpecs, grids, VMEM residency, DMA.

ch 11 Pallas fundamentals ch 15 Capstone: shipped upstream
ch 07

Mosaic → LLO

The TPU backend. Kernels compile here; you never touch it.

ch 08

TPU chip

MXU, VPU, VMEM, HBM. Where all the time is actually spent.

ch 10 The machine
ch 09

ICI ⇄ chips

The inter-chip links. Collectives are kernels here too.

ch 14 Distributed Pallas
you have reached the floor

The workshop

Below the stack sits the equipment: twelve instruments, four drills, six museum exhibits, thirteen labs. Every one is assigned to a chapter above; press ⌘K to jump anywhere.

EX·01 the roofline playground drag shapes across the ridge
EX·09 the BlockSpec sandbox legal and illegal schedules, live
EX·05 the IR x-ray one program, three layers, hover-synced
EX·03 streaming attention the score matrix that never exists
EX·10 the semaphore timeline step into a deadlock on purpose
EX·11 the mesh, hop by hop collectives over real link budgets
the honesty rule

Gates pass only on real hardware, and failures stay published. Every value on this site states what it is: a roofline prediction, an off-chip verification, or a measured run with chip, dtype, shapes, and date attached.

roofline-predicted floors · v5e · hardware rows replace these
matmul NxNxN697.7 µs
matmul skinny 8xNxN41.1 µs
elementwise add122.8 µs
the full bench →