Update TPU chapter for latest generation (Ironwood/TPU7x) by mwdaub-google · Pull Request #124 · jax-ml/scaling-book

mwdaub-google · 2026-06-22T21:36:20Z

Update primary focus of "All About TPUs" chapter to Ironwood/TPU7x (and Trillium/TPU v6e for highlighting differences between training and inference chips).
Some sections and exercises still discuss TPU v5e and TPU v5p, but the main thrust of the chapter was simplified to the latest generations.
Also use specs from public TPU documentation in more places.

jacobaustin123

Thank you! A few general thoughts (I only got through about 2/3 of this).

In general my instinct is to avoid changing worked problems that people may have asked questions about in the comments or solved themselves, so it stays a reference people trust won't change under them. Would you be open to just adding some new questions that are v7-specific, maybe highlighting areas where v7 is different from previous chips?
I'd love to keep this comprehensive to include older TPU gens even though v7 is the current one. When v8 comes out, it would be great to not feel like we need to change it again, but merely add some new details about the latest gen. I'd love to tweak a few of the examples in the text to emphasize v7, but especially showing how it differs from previous gens.

jacobaustin123 · 2026-06-22T22:07:27Z

 The diagram above also includes a few other components like SMEM and the scalar unit, which are used for control flow handling and are discussed briefly in <a href="#appendix-a-more-on-tpu-internals">Appendix A</a>, but aren't crucial to understand. On the other hand, HBM is important and fairly simple:

-* **HBM** (High Bandwidth Memory) is a big chunk of fast memory that stores tensors for use by the TensorCore. HBM usually has capacity on the order of tens of gigabytes (for example, [TPU v5e has 16GiB of HBM](https://cloud.google.com/tpu/docs/v5e#system_architecture)).
+* **HBM** (High Bandwidth Memory) is a big chunk of fast memory that stores tensors for use by the TensorCore. The HBM capacity of TPU7x is [192GiB](https://cloud.google.com/tpu/docs/tpu7x#system_architecture).


Can we change this to tens to hundreds of GB so it stays comprehensive across generations.

jacobaustin123 · 2026-06-22T22:09:19Z

-{% include figure.liquid path="assets/img/cores.png" class="img-fluid img-small" %}
-
-**Chips** are arranged in **sets of 4 on a 'tray'** connected to a **CPU host via PCIe network.** This is the format most readers will be familiar with, 4 chips (8 cores, though usually treated as 4 logical megacores) exposed through Colab or a single TPU-VM. For inference chips like the TPU v5e, we have 2 trays per host, instead of 1, but also only 1 core per chip, giving us 8 chips = 8 cores.<d-footnote>On Cloud TPU VMs, each tray is exposed as part of a separate VM, so there are once again 4 cores visible.</d-footnote>
+**Chips** are arranged in **sets of 4 on a 'tray'** connected to a **CPU host via PCIe network.** This is the format most readers will be familiar with, 4 chips (8 cores, since TPU7x has two TensorCores per chip) exposed through Colab or a single TPU-VM. For inference chips like TPU v6e, we have 2 trays per host, instead of 1, but also only 1 core per chip, giving us 8 chips = 8 cores.<d-footnote>On Cloud TPU VMs, each tray is exposed as part of a separate VM, so there are once again 4 cores visible.</d-footnote>


likewise can we keep this more general to mention both megacore for older gens, i.e. 4 chips (8 cores, though some older generations combine these into a single 'megacore').

jacobaustin123 · 2026-06-22T22:10:08Z

 ## TPU Networking

-**Chips are connected to each other through the ICI network in a Pod**. In older generations (TPU v2 and TPU v3), inference chips (e.g., TPU v5e), and Trillium (TPU v6e), ICI ("inter-chip interconnects") connects the 4 nearest neighbors (with edge links to form a 2D torus). TPU v4 and TPU v5p are connected to the nearest 6 neighbors (forming a 3D torus). Note these connections do **not** go through their hosts, they are direct links between chips.
+**Chips are connected to each other through the ICI network in a Pod**. In inference-optimzed Trillium (TPU v6e), ICI ("inter-chip interconnects") connects each chip to its 4 nearest neighbors (with edge links to form a 2D torus). On the other hand, each chip in training-optimized TPU7x is connected to the nearest 6 neighbors (forming a 3D torus). Note these connections do **not** go through their hosts, they are direct links between chips.


Can we drop the phrase "training-optimized" since it reads a bit like branding, esp since there's also a 3D inference optimized version. Maybe just add v7x to this list.

The phrase "inference-optimized" was already present in this chapter, and I don't think training-optimized is any different. Regardless, I rephrased so it isn't directly applied to "Trillium".

I think it's more true that v5e is inference optimized than that v7x is training optimized (since there's an explicit inference optimized version of the v7 chip)

Fair enough. I think the new phrasing doesn't contradict this, though happy to hear suggestions for improvements.

jacobaustin123 · 2026-06-22T22:11:35Z

-* `9e10` bytes/s (90 GB/s) of ICI bandwidth per axis, with 3 axes per chip.<d-footnote>The page above lists 100 GB/s of bandwidth, which is slightly different from what's listed here. TPU ICI links have slightly different bandwidths depending on the operation being performed. You can generally use the numbers in this doc without worry.</d-footnote>
-* `6.25e9` bytes/s (6.25 GB/s) of DCN (egress) bandwidth per TPU (via 1-2 NICs on each host).<d-footnote>TPU v6e and TPU7x have 12.5e9 bytes/s and v5e has 3.125e9 bytes/s.</d-footnote>
+* `7.92e12` bytes/s (7380 TiB/s) of HBM bandwidth per chip.
+* `1.2e12` bytes/s (1200 GB/s) of bidirectional ICI bandwidth per chip.


Can we keep this per-axis and keep the foot-note? that's very important

jacobaustin123 · 2026-06-22T22:11:46Z

 * **Within a slice, TPUs are only connected to their nearest neighbors via ICI.** This means communication over ICI between distant chips in a slice needs to hop over the intervening chips first.

-* **Weight matrices need to be padded to at least size 128** (256 on TPU v6e) in both dimensions to fill up the MXU (in fact, smaller axes are padded to 128).
+* **Weight matrices need to be padded to size divisible by 256** in both dimensions to fill up the MXU (in fact, smaller axes are padded to 256).


Can we keep this covering v5e/v6e?

jacobaustin123 · 2026-06-22T22:12:16Z

+Host size refers to the topology of TPUs connected to a single host (e.g. TPU v5e has a single CPU host connected to 8 TPUs in a 4x2 topology). Here are interconnect figures:

-We include both one-way (unidirectional) bandwidth and bidi (bidirectional) bandwidth since unidirectional bandwidth is more true to the hardware but bidirectional bandwidth occurs more often in equations involving a full ring.<d-footnote>By bidi (bidirectional) bandwidth we mean the total bytes that can be sent along a single link in both directions, or equally, the total number of outgoing bytes from a single TPU along a particular axis, assuming we can use both links efficiently. This is true when we have a functioning ring, AKA when we have a wraparound connection on the particular axis. This occurs on inference chips when we have a full 16 axis, or on training chips (v*p) when we have an axis which is a multiple of 4. We prefer to use the bidirectional bandwidth because it appears frequently in calculations involving bidirectional comms.</d-footnote>
+| Model       | ICI BW/link (bidi, bytes/s) |


These numbers aren't quite right. The original values are correct I think

jacobaustin123 · 2026-06-22T22:12:38Z

-**Question 1 [bounding LLM latency]:** Say you want to sample from a 200B parameter model in bf16 that's split across 32 TPU v4p. How long would it take to load all the parameters from HBM into the systolic array? *Hint: use the numbers above.*
+**Question 1 [bounding LLM latency]:** Say you want to sample from a 200B parameter model in bf16 that's split across 32 TPU v6e chips. How long would it take to load all the parameters from HBM into the systolic array? *Hint: use the numbers above.*

 {% details Click here for the answer. %}


I worry about changing problems post-hoc, could we keep these unchanged and just add new problems for TPU7?

jacobaustin123 · 2026-06-22T22:12:47Z

 {% details Click here for the answer. %}

-**Answer:** For TPU v5e, each pod is `16x16` and each host is a 4x2 slice, so we have `16*16 / 8 = 32` hosts. For TPU v5e, each TPU has only one core, so we have 256 TensorCores. The total FLOPs/s is `16*16*2e14 = 5.1e16` in bfloat16. Each chip has 16GB of HBM, so that's `256 * 16 = 4TB` of memory.
+**Answer:** For TPU v6e, each pod is 256 chips and each host is a 4x2 slice, so we have `256 / 8 = 32` hosts. For TPU v6e, each TPU has only one core, so we have 256 TensorCores. The total FLOPs/s is `256*9.2e14 = 2.4e17` in bfloat16. Each chip has 32GiB of HBM, so that's `256 * 32 = 8TiB` of memory.


I think 16x16 is good since that is the topology.

jacobaustin123 · 2026-06-22T22:16:22Z

Oh also a recent PR adds v7 specs, so maybe worth reconciling those (there's a small conflict).
Also a bunch of the TPU specs here are possibly not quite right, there are a lot of inconsistent numbers but I think the existing ones are close to the true values.

Update TPU chapter for latest generation (Ironwood/TPU7x)

2b2fadc

jacobaustin123 reviewed Jun 22, 2026

View reviewed changes

mwdaub-google added 2 commits June 26, 2026 07:33

Merge branch 'main' into all-about-tpus-update-7x

5890275

First round of comments

bed433c

Uh oh!

Conversation

mwdaub-google commented Jun 22, 2026

Uh oh!

jacobaustin123 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jacobaustin123 commented Jun 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants