The TensorCore complex

Google publishes two kinds of document about a TPU. Product pages give you a peak number, a memory size, and a link to pricing. ISCA papers give you the parts. The papers stop at v4, which is why v4 is the specimen worth opening: one paragraph of the 2023 paper names every compute block on the die, and Table 4 of the same paper prints the clock those blocks run at. After v4 there are product pages and launch posts, and neither kind of document lists parts.

each TPU v4 contains two TensorCores (TC). Each TC contains four 128x128 Matrix Multiply Units (MXUs) and a Vector Processing Unit (VPU) with 128 lanes (16 ALUs per lane) and a 16 MiB Vector Memory (VMEM). The two TCs share a 128 MiB Common Memory (CMEM).

Two things in that sentence never reach a product page. CMEM is 128 MiB of memory the two TensorCores share, and no later generation is documented as keeping it or dropping it. The VPUThe vector unit for elementwise work, organized as (8, 128) lanes; the origin of the tiling lattice every layer above obeys.taught in /l/tpu → is given as 128 lanes with 16 ALUs each, a count you should always carry its generation with: the scaling book describes a v5 VPU as a (8, 128) grid with 4 ALUs per lane-sublane pair, which is a different chip counted a different way. Neither source is wrong. Quoting one shape without its generation is.

Two cores sharing one pool of memory raises a question the paper does not answer: does software see one accelerator or two? It never uses the word MegaCore for the arrangement either. Google Cloud’s docs and the scaling book do the answering, describing a chip whose two cores share memory and can be treated as one large accelerator. So the fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → is a software-visible arrangement resting on a hardware fact, and the hardware fact is CMEM. Which core jax.devices() hands you is the generations lesson at /l/tpu/generations.

The chip lesson at /l/tpu/tpu-chip already said what these units do, and what 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 → is instead of a cache. This one counts them. Counting is what turns a spec sheet from a claim into arithmetic you can check.

A peak FLOPs figure is three numbers multiplied together: how many multiply-accumulate cells the chip has, two operations per cell per cycle, and the clock. Start with v1, where all three sit in one paper. Its matrix unit is a single 256x256 grid, 65,536 cells, running at 700 MHz. That gives 91.75 TOPS against a published 92.

Now v4, where the cells are spread across eight smaller arrays instead of one big one. Two TensorCores, four MXUs each, 128x128 per MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu →: 131,072 cells, exactly twice what v1 carried in its single array. At 1050 MHz that is 275.25 TFLOPS against a published 275.

Agreement at three digits is worth pausing on, because of what it tells you about the number being agreed with. There is no derating in it. Nothing accounts for the cycles the systolic array spends filling and draining, nothing accounts for a kernel that cannot keep the array fed. The vendor did the same multiplication you just did. A peak figure is geometry times a clock, and every real kernel lives underneath it.

The same arithmetic survives one more zoom. 275 TFLOPS across a full 4096-chip v4 pod is 1.126 exaflops, and Google’s v4 page advertises 1.1 exaflops. Peak numbers at pod scale are the chip number times the chip count, nothing else.

v1   256 x 256 MACs         =  65,536         ISCA 2017, section 2
     x 2 ops per MAC        = 131,072         one multiply, one add
     x 700 MHz             =  91.75 TOPS     ISCA 2017, Table 2
     Google publishes         92 TOPS        ISCA 2017, Table 2

v4   2 TC x 4 MXU x 128x128 = 131,072 MACs   ISCA 2023, section 2
     x 2 ops per MAC        = 262,144
     x 1050 MHz            = 275.25 TFLOPS   ISCA 2023, Table 4
     Google publishes        275 TFLOPS      ISCA 2023, Table 4

v1   256 x 256 MACs         =  65,536         ISCA 2017, section 2
     x 2 ops per MAC        = 131,072         one multiply, one add
     x 700 MHz             =  91.75 TOPS     ISCA 2017, Table 2
     Google publishes         92 TOPS        ISCA 2017, Table 2

v4   2 TC x 4 MXU x 128x128 = 131,072 MACs   ISCA 2023, section 2
     x 2 ops per MAC        = 262,144
     x 1050 MHz            = 275.25 TFLOPS   ISCA 2023, Table 4
     Google publishes        275 TFLOPS      ISCA 2023, Table 4

v5p  2 TC x 4 MXU x 128x128 = 131,072 MACs   Google v5p page (2 TC);
     x 2 ops per MAC        = 262,144        4 MXUs/TC read across from v4 and v5e
     x about 1.75 GHz      = 458.8 TFLOPS    scaling book, SECONDARY, no vendor clock
     Google publishes        459 TFLOPS      Google v5p page

v6e  1 TC x 2 MXU x 256x256 = 131,072 MACs   Google v6e page + architecture page
     x 2 ops per MAC        = 262,144
     Google publishes        918 TFLOPS      Google v6e page
     implied clock          = 918e12 / 262,144 = 3.50 GHz
     no v6e clock is published, and 3.50 GHz is not a plausible TPU clock

v5p is where the derivation starts needing a footnote. Google’s v5p page gives two TensorCores per chip and a peak of 459 TFLOPS bf16. It does not print an MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → count per TensorCore, though v4 and v5e are both documented at four, and the v5p page is written to the same shape. Take that layout and the MAC count comes out at 131,072, exactly what v4 had. Same cells, 1.67x the FLOPs. The extra has to be clock, and Google publishes no TPU clock after v4.

One public figure fills the gap, and it is secondary: the scaling book states that a TPU v5p runs at about 1.75 GHz. Run it forward. 262,144 FLOPs per clock times 1.75e9 is 458.8 TFLOPS against a published 459. A secondary number that reproduces a vendor number to three digits has earned a place in the lesson, with its label still attached. If Google ever prints a v5p clock, that is the citation this paragraph should carry instead.

Notice what the closing derivation buys you: it runs backwards. Whenever a vendor publishes a peak and enough geometry, you can solve for whatever constant they withheld, and then judge a third-party claim by whether it lands where the arithmetic says it must. What that clock did to the ratio between compute and bandwidth is the generations lesson at /l/tpu/generations, which owns the ridge.

Trillium, the v6e generation, breaks the pattern. Google’s v6e page gives one TensorCore per chip with two matrix-multiply units on it, and a peak of 918 TFLOPS bf16. The architecture page gives the v6e MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → as 256x256. Multiply those together and you get 131,072 MACs again, 262,144 FLOPs per clock, and an implied clock of 3.50 GHz.

That clock is not credible. The fastest TPU Google itself has printed is v4 at 1050 MHz, and the fastest figure in the public record of any kind is the scaling book’s 1.75 GHz for v5p. So at least one of the three inputs is wrong, and no published clock exists to arbitrate: the Trillium launch post says the team expanded the MXUs and raised the clock speed, and attaches no number to either.

Google’s own architecture page also undercuts one of the inputs. It gives the MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → size for v6e and TPU7x as 256 x 256 multiply-accumulators, and elsewhere in the same passage it puts an MXU at 16K multiply-accumulate operations per cycle. 16K is 16,384, which is 128 x 128. Whichever of those two describes v6e, they do not describe it together, and anyone quoting both in one breath is publishing a contradiction.

You can make the numbers close by assuming four 256x256 MXUs instead of two, which puts the implied clock at 1.75 GHz and matches v5p exactly. That is arithmetic, not evidence. Google’s page says two. So the honest result is an open derivation, written down as open: the published v6e MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → shape and the published 918 TFLOPS cannot both hold at any plausible clock, and the lesson stops there rather than inventing the number that would rescue it.

Every figure in this lesson is a count of cells, never a count of cores, and that is deliberate. Google’s TensorCore is a whole core: four MXUs, a vector unit, and a scalar unit under one instruction stream, two per chip on v4 and 8192 across a full 4096-chip pod. NVIDIA’s Tensor Core is a functional unit inside one of an SM’s four processing blocks, 528 of them on an H100 SXM5. The GPU chapter works that name collision through properly at /s/machine/gpu-chip; what it does to arithmetic is this lesson’s problem.

A core count cannot be multiplied by anything. One chip has two of them and the other has 528, and neither figure enters a derivation, because the two vendors are counting objects at different scales. MACs per clock per chip is the quantity both sides can produce, which is why every derivation above turns a TensorCore count into 131,072 cells in the same breath.

The count also moves between generations for reasons that have little to do with capability. v4 and v5p carry two TensorCores, v5e and v6e carry one. v6e drops to a single core with two 256x256 MXUs, and on those published numbers the cell count lands exactly where it has sat since v4, at 131,072. That is the same shape whose peak refuses to close, so hold it loosely; the point stands either way, which is that the core count is packaging and the cell count is the machine.

Past v4 the parts list thins out fast, and the gaps are worth knowing by name rather than discovering mid-argument. No clock is published for v5e, v5p, or v6e. 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 → is public for v3 and v4 and, secondhand through the scaling book, for v5e; for v5p and v6e it is not public at all. SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → gets named by both Google and the scaling book and sized by neither, at any generation. Whether CMEM survives past v4 is not stated anywhere. Die size, transistor count, and process node stop at v4 with the ISCA papers.

What to do with a gap is the part worth practising. The temptation is to divide two published numbers and print the quotient as a spec, which is exactly the move that yields a 3.50 GHz TPU. Write the derivation instead, with its inputs and their sources, and label the result: closed against vendor numbers, closed only against a secondary source, or open. A lesson that says which figures are not public is more useful than one that quietly fills them in, because the reader can tell what to trust.