The physical picture
Chips in a TPU pod connect to their neighbors over ICI links (4.5e10 bytes per second each way per link on v5e), and slices along a mesh axis form physical rings. Every collective you have ever called resolves to movements over these links: an all-gather is N−1 hops around the ring, each chip pushing one shard to its neighbor; a reduce-scatter is the same ring with an add at each hop.
The TPU's native distributed operation is the async remote DMAAn asynchronous copy between memories that runs while compute continues; the grid pipeline is DMAs the runtime writes for you.taught in /l/pallas →: a chip pushes a buffer directly into a neighbor's VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → and signals a semaphore, while its compute units keep working. That last clause is the entire economics of distributed kernels: transfers that hide behind compute are free, and transfers that do not are the dominant tax at scale. The craft is arranging the overlap.
The deep continuity with everything above: sharding decides where data lives, and a kernel is always a per-device program on its local shard. The streaming-attention algebra from stage 3 does not care whether a KV block arrived from local HBMThe chip’s main memory: large, far, and the resource memory-bound ops spend. 8.2e11 bytes per second on v5e, 1.6e12 on v6e.taught in /l/tpu → or from the neighbor chip, which is why ring attention (stage 4) is a composition, not a new derivation. EX·04 shows it: the same schedule as EX·03, longer arrows.
Lessons
- 01The TPU fabricTPUs scale by wiring each chip to its neighbors and to nothing else. Everything about distributed kernels follows from that one wiring decision. ·
- 02The vocabulary, as ring movementsAll-gather, reduce-scatter, all-reduce: each one is the same ring step repeated, shards moving neighbor to neighbor by remote DMA. ·
- 03Cost formulas you can do in your headRing collective cost is desk arithmetic, not measurement: bytes, chips, and one link rate are the whole formula. ·
- 04The collective is a kernelNothing about a collective is special hardware. It is the same async remote DMA and semaphore pairing, written as a kernel you can read. ·
- 05The GPU fabricInside a node, every GPU is one hop from every other. Cross the node boundary and you are on a different network with different numbers. ·
- 06The optical patch panelA TPU v4 pod has almost no long fixed cables. Forty-eight optical switches aim mirrors at fibers, and the shape of your slice is a routing table they load when the job starts. ·