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.
3/6 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. Where all the time is actually spent; the GPU taught alongside, for contrast.
The workshop
Below the stack sits the equipment: twelve instruments, six gym stations, eight museum exhibits, fifteen 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.