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.
4/5 gates passed · your progress 0/15 chapters
03/06step 2 · one pass over K1, V1: rescale (m, l, acc), never materializing scores to HBM
JAX / PyTorch
Math on whole arrays. No notion of hardware.
jaxpr
The traced program: a flat list of primitive ops.
XLA
Fuses along dataflow edges. Cannot restructure an algorithm.
the gap
XLA stops at fusion. Pallas asks for everything by hand. Between them: nothing. Stage 3 teaches you to cross on foot.
Pallas
Hand choreography: BlockSpecs, grids, VMEM residency, DMA.
Mosaic → LLO
The TPU backend. Kernels compile here; you never touch it.
The machines
MXU, VPU, VMEM, HBM, and the profiler's own account of them. The GPU taught alongside, for contrast.
The workshop
Below the stack sits the equipment: twelve instruments, eight gym stations, eight museum exhibits, eighteen labs. Every one is assigned to a chapter above; press ⌘K to jump anywhere.
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.