the path · chapter 08 of 15 · part i, the descent

TPU chip

Two compute units, one scratchpad, one far memory. Every performance question resolves here.

mastery work · this chapter0/2
  1. go →
manual items are your word; auto items complete from your streaks, labs, and can-you ticks · stored in your browser only
the layer

The four numbers

A TPU v5e: 1.97e14 bf16 FLOPs per second in the MXU, 8.2e11 bytes per second from HBM, roughly 128 MiB of VMEM scratchpad, 16 GB of HBM. Every kernel decision in this track is arithmetic over these four numbers. The ratio of the first two is the ridge (about 240 FLOPs per byte): ops below it are memory-bound, ops above it compute-bound, and no cleverness moves an op across the ridge except changing its algorithm.

The MXU is a 128x128 systolic array (256x256 on v6e): matmuls only. The VPU handles everything elementwise in (8, 128) lanes, which is where the tiling lattice that haunts every layer above actually comes from. VMEM is software-managed: nothing is cached for you; what is resident is what your schedule staged. Stage 0's EX·06 draws the hierarchy to these real widths, and EX·01 lets you drag op shapes across the roofline.

The habit the whole track builds on this layer: predict from the constants before measuring, then explain every miss. The constants are public, the arithmetic is short, and an engineer who cannot predict a latency floor from them is guessing everywhere above.

later on the path Stage 0 lives on this layer (chapter 10). Keep descending; the path arrives there in order.