TPU lineage and roofline thinking

This chapter does two things. It walks the ten TPU generations as a sequence of answers to specific bottlenecks, and it sets up the one model the rest of the course reasons with: the roofline. The two are the same story. Every generation is Google moving a ridge point.

Ten generations in one table

Peaks are BF16 unless stated. "2 × 4 × 128×128" means two TensorCores, each with four 128×128 matrix units. Google publishes only FP4 peaks for TPU 8. The JAX hardware table lists BF16 at 996 TFLOP/s for 8t and 1,100 TFLOP/s for 8i provisional [J-hw]. Neither chip had reached general availability as of 2026-10-23.

Each generation answers a bottleneck

TPUv1 (2015): inference throughput per watt

Google projected in 2013 that speech recognition alone could double its datacenter compute [P-v1, §1]. The answer was a PCIe coprocessor with one 256×256 int8 systolic array, 24 MiB of on-chip buffer and 8 GiB of DDR3 at 34 GB/s. It ran 15–30× faster than the K80 GPU and Haswell CPU it was measured against, at 30–80× better TOPS/W. Its own roofline exposed its limit: the ridge sat at 1,350 operations per weight byte, so most production models were memory-bound, and the paper estimates GDDR5 would have tripled performance [P-v1, §4, §7].

TPUv2 (2017): training

Training needs floating point, more memory bandwidth and many chips. v2 replaced DDR3 with HBM (700 GB/s, about 20× v1), added a programmable vector unit with bf16 multiply and fp32 accumulate, and added the inter-chip interconnect, ICI. Chips connect directly in a 2D torus with routers on the chip and no external switch [P-Norrie] [P-v23]. The array shrank to 128×128, because a 256×256 array wastes more of its cells on the matrix shapes real models produce [HC-v23, Why 128x128].

TPUv3 (2018): more of the same, cooled by liquid

Two matrix units per core, 940 MHz instead of 700, twice the HBM, and pods of 1,024 chips. Liquid cooling allowed the extra power [P-v23].

TPU v4 (2020): scale, failures and embeddings

v4 has three new pieces [P-v4]:

• A 3D torus of 4×4×4 cubes, joined by optical circuit switches (OCS). The switches let a 4,096-chip pod be re-cut into slices, route around failed racks, and use twisted torus shapes. They cost under 5% of the system.
• SparseCores for embedding lookups. They take about 5% of die area and run embedding-heavy models 5–7× faster.
• 128 MiB of shared CMEM on the chip.

TPU v5e and v5p (2023): one line for cost, one for scale

v5e is a single-core chip in 256-chip 2D pods, priced for serving and small training [D-v5e]. v5p is the training flagship: 95 GiB of HBM, 2,765 GB/s, and 8,960-chip pods [D-v5p].

TPU v6e Trillium (2024): a bigger matrix unit

v6e moved to 256×256 arrays, giving 4.7× v5e's peak per chip [G-v6e] [D-arch]. v1 had used a 256×256 int8 array; v2 through v5 used 128×128. By 2024 the matrix dimensions in production models were large enough to fill 256×256 arrays.

TPU7x Ironwood (2025): inference at training scale

Two compute chiplets per package, each seen by JAX as its own device. 192 GiB of HBM at 7,380 GB/s, native FP8 at 4,614 TFLOP/s, and 9,216-chip pods at about 10 MW [D-7x] [G-7x]. Google's launch post calls it "the first Google TPU for the age of inference". The HBM capacity is for KV caches and the bandwidth is for decode.

TPU 8t and 8i (announced April 2026): two chips

The first generation split by workload [G-8dd]:

• 8t keeps the 3D torus, scales the pod to 9,600 chips and doubles ICI bandwidth relative to Ironwood.
• 8i has 288 GB of HBM at 8,601 GB/s and 384 MB of on-chip SRAM, three times Ironwood's, sized to keep KV cache on chip. It replaces the torus with a lower-diameter "Boardfly" network of at most 7 hops, and replaces SparseCores with a collectives engine.

The rest of this chapter explains why decode traffic pushed Google there.

FLOPs outgrew bandwidth

From TPUv2 to Ironwood, per-chip BF16 FLOP/s grew 50× and HBM bandwidth grew 10.5×. Google's own retrospective gives the same shape: about 100× peak and 10× HBM over five generations, counting Ironwood's FP8 peak [P-Micro26]. The ratio of the two is the number this chapter is about.

Two times bound every kernel

A kernel spends time in two places: arithmetic and data movement. The Scaling Book models both as a count divided by a rate [SB1].

Time on arithmetic is work over peak FLOP/s, C. Time on data movement is bytes over the bandwidth W of whichever link carries them: HBM, ICI, PCIe or DCN.

With perfect overlap the kernel takes $\max(T_\text{math}, T_\text{comms})$. With none it takes the sum, at most 2× worse. The course reasons from the max.

Arithmetic intensity is the ratio of the two counts: FLOPs per byte moved. A chip has an intensity too, the ratio of its two rates. A kernel is compute-bound when its intensity exceeds the chip's:

A kernel keeps the matrix units busy only if it does more arithmetic per byte than the chip can do per byte of bandwidth.

The right-hand side is the chip's ridge point. TPUv2's was 66 FLOP/B and Ironwood's is 313. For the chips used in this course, against HBM:

Values come from Cloud TPU docs, the JAX hardware table for TPU 8's BF16, FP8 and INT8 rows, and NVIDIA product pages with sparse figures halved [D-7x] [J-hw] [G-8dd] [N-h100]. The Scaling Book uses 240 for v5e; the docs-based figure here is 241. v5p's ridge, 166, is lower than v5e's because v5p carries 3.4× the bandwidth for 2.3× the FLOPs. v5p and v6e list FP8 at the same rate as BF16, so FP8 on those chips saves bytes but not time [D-v5p] [J-hw].

A matmul's intensity is its batch

Take $X[B, D] \cdot W[D, F]$ in bf16. It does $2BDF$ FLOPs and moves $2BD + 2DF + 2BF$ bytes:

When the token batch is small next to the weight dimensions, weight bytes dominate, and every weight byte is used once per token. Intensity equals the number of tokens.

So a bf16 matmul on Ironwood is compute-bound once the token batch per replica exceeds about 313. On v5e the figure is 241 and on an H100 it is 295 [SB1]. The batch counts tokens, not sequences. Model sharding does not change it, because sharding divides FLOPs and bandwidth by the same factor.

Precision moves the ridge

Store the weights in int8 or fp8 and keep the math in bf16. FLOPs stay at $2BDF$, but weight bytes halve to $DF$, so intensity becomes about $2B$:

Each weight byte read supports 2 / (bytes per weight) FLOPs per token, so intensity is about 2B divided by bytes per weight. C is the peak for the precision the math runs in. Halving bytes per weight halves the batch needed; doubling the math rate doubles it back.

The Scaling Book writes this as $\beta \cdot C/W_\text{hbm}$, with $\beta$ the ratio of weight bits to activation bits and $C$ the BF16 peak [SB7]. For bf16 math the two forms agree. Formula 1.4 also covers low-precision math directly.

On v5e, int8 weights with bf16 math bring the compute-bound batch from 240 to 120 [SB1, Q2]. On Ironwood the example shape goes from 329 to 164.

Moving the math to low precision as well doubles $C$. On Ironwood, fp8 math runs at 4,614 TFLOP/s [D-7x], so the ridge doubles to 625 and the compute-bound batch returns to 329. The step still gets faster: at batch 64 it drops from 64 µs to 32 µs, because the bytes halved. The batch at which the matrix units saturate is what stays put.

The provisional TPU 8i figures push this further. Its fp8 ridge is about 1,024 FLOP/B provisional. With fp8 weights and fp8 math, the example shape needs more than 557 tokens per chip to be compute-bound. Decode batches that large are rare, especially with long contexts, as Chapter 7 shows. Google's design answer for 8i was bandwidth and on-chip SRAM, not FLOPs: 8,601 GB/s of HBM and 384 MB of SRAM [G-8dd].

The same roofline for a network link

Replace HBM with an inter-chip link and the model is unchanged. The Scaling Book's example splits the contracting dimension $D$ of $X[B, D] \cdot W[D, F]$ across two chips [SB1]. Each chip multiplies its half, $BDF$ FLOPs, then the chips exchange $2BF$ bytes of bf16 partial sums:

Splitting the contracting dimension is worthwhile only if each chip's half of the matmul takes longer than shipping the partial result to the other chip.

The threshold now depends on $D$, not on the batch. On v5e, with 45 GB/s one way per link, it is $D > 8{,}755$ [SB1]. On Ironwood, the Cloud docs give 200 GB/s per link counting both directions, so 100 GB/s one way [D-7x]. The threshold becomes $D > 46{,}140$ derived: no current model has a hidden dimension that large. Ironwood's FLOPs grew 11.7× over v5e while its per-link bandwidth grew 2.2×. Chapters 5 and 6 work through what that does to tensor parallelism.

What the HBM roofline leaves out

Three things the model above does not capture, each of which shows up in Chapter 10's profiles.

Vector-unit work has its own roofline. Elementwise ops, reductions and softmax run on the VPU, not the matrix units. On v5p the VPU peaks at about 7 TFLOP/s per core, a critical intensity of about 3 [SB1]. A dot product, at 0.5 FLOP/B, is memory-bound on any chip. These ops are cheap only when XLA fuses them into the matmul that produces or consumes their input.

Small transfers are latency-bound. Each ICI hop costs about 1 µs regardless of size. On v5e, any per-hop chunk under about 45 kB spends more time on latency than bandwidth [SB3]. That is the normal regime for decode-time collectives (Chapter 7).

Peak is not achievable peak. The Scaling Book states that H100s and B200s usually reach 80–85% of claimed peak FLOP/s, while TPUs get closer to 95% [SB1]. The book does not publish the measurements behind those figures, so treat them as practitioner estimates.