How to Think About TPUs | How To Scale Your Model

[jacobaustin123]

Discussion about TPUs!

[mitchellgoffpc]

In the solution for question 5, it looks like bytes transferred should be 1.7e7 rather than 1.7e10, and transfer time should be 170µs rather than 170ms

[jacobaustin123]

Very good catch, several parts of this answer are wrong. I've updated the repo with the fix, will take a few minutes to update the website.

[kishorepv]

Multiplication is more expensive than an addition right? In multiplying a matrix when we say the number of operations is ~ 2 x B x D x F (BDF for multiplication, and B(D-1)F for addition)?, do we consider multiplication and addition as equally expensive ?

[jacobaustin123]

The hardware typically treats them the same (each cycle it does an equal number of multiplies and adds to perform a matmul). They may require different numbers of transistors but the overall speed is the same.

With that said, the TPU MXU does matrix multiplications and not, say, pure sums/reductions. You can use the MXU speed to calculate how long matmuls will take, but you can just say "hey I want to do 1e10 adds, let me just divide that by the TPU/GPU bfloat16 FLOPs figures".

[kishorepv]

Got it. Thanks.
If the hardware treats multiplication and additions as the same, then do techniques like Karatsuba / Toom-Cook algorithm (which reduces the number of multiplications at the cost of increased additions / subtractions) have any use in speeding up matmul on a GPU/TPU?

[kerrickstaley]

@kishorepv: The Karatsuba algorithm is for multiplying integers. Did you mean the Strassen algorithm which is a fast algorithm for multiplying matrices?

[fedelebron]

You can implement Strassen, but trading what the TPU is excellent at (matrix multiplications) for something it's merely OK at (pointwise operations) might not be a net win, on matrix sizes you care about (i.e. before the asymptotics totally dominate). Sharding (the next chapter) is another way to speed up large matrix multiplications that scales with better constants, at the cost of, well, more hardware :)

[kerrickstaley]

In question 5, the answer mentions v4e; I think it should be v5e. Also, I don't understand why the total transfer time is multiplied by the number of hops. Streaming from one core to another should take num_bytes / bandwidth + latency_of_first_byte, which is approximately num_bytes / bandwidth when num_bytes is large. This means it should take about num_bytes / bandwidth = 170 us for the whole transfer.

[jacobaustin123]

Thanks for catching this. This got changed around a few times and we missed this. Fixed.

[kishorepv]

There is a typo in the line "In each frame above, we multiply all the overlapped green and blue units, sum the result with any residual passed in from above, and then pass the result in turn down one unit...". It should be "In each frame below" instead of "In each frame above"

[burichh]

In the first gif in "Appendix B: How does a systolic array work?", the last image (can't embed into this message:( ) shows an incorrect calculation of the bottom left output value (in purple). It shows $W_{00} \cdot X_{01} + W_{10} \cdot X_{00}$ while actually it should be $W_{01} \cdot X_{01} + W_{00} \cdot X_{00}$.

It's a bit strange (but maybe only for me) that for both $W$ and $X$ it is the second index that runs across the multiplication, while in general the notation is such that the second and first indices are running:

$$(W\cdot X)_{nm} = \sum_i W_{ni} \cdot X_{im}$$

But that's actually just a question of notation.

[burichh]

Ups, sorry, I wanted to add a new comment, not reply to yours. Upsie!

[manavgarg]

I actually found this to be more intuitive when understanding systolic arrays.

[Shua1]

H100 spec says:

INT8 Tensor Core* 3,958 TOPS
FP16 Tensor Core* 1,979 teraFLOPS
, meaining INT8 is 2x of FP16 flops. Simplistically, understandable since FP16 is 2x the size of INT8. Though I expect integer operation is faster than FP operations.

But the text sasy:

This is about 5e13 bf16 FLOPs/s per MXU at 1.5GHz on TPU v5e.
TPUs also support lower precision matmuls with higher throughput (e.g. each TPU v5e MXU can do 4e14 int8 OPs/s)

int8 is about 8x faster than FP16. Any idea why the difference between 2x on H100 and 8x on TPU v5e? Thanks!

[jacobaustin123]

This was supposed to say "each TPU v5e chip can do 4e14 FLOPs", i.e. it's only 2x the FLOPs of bfloat16. Fixed.

[Kumawat-Akhilesh]

I am still grappling with the idea of when I should use TPU (say V6e) versus an Nvidia B200. The computation FLOPs/s is higher for B200 (4500 TFLOPs v 920 TFLOPs, bf16), HBM bandwidth is higher for B200 (8TB/s v 1.6TB/s), HBM size is higher for B200 (192GB v 32GB), interconnect is faster for B200 (NVlink at 1800GB/s v ICI BW at 180GB/s)...
I understand the idea of larger pods for TPUs but even if B200 connects to say 576 GPUs pod, beyond that they connect at 800GB/s InfiniBand or Ethernet and even this DCN is faster than the ICI speed of TPU (180GB/s)...

Hence, I am a bit confused on the tradeoff. Isn't B200 outperforming on all fronts (except costs perhaps due to switch costs)?
In fact, even H100 outperforms any TPU on all these metrics.. so I am not understanding when do TPUs outperform these GPUs.. for what workloads?

[jacobaustin123]

I think the key is that nothing matters except performance / dollar. FLOPs / dollar. ICI bandwidth / dollar. It's relatively easy to package more FLOPs or bandwidth in a package. What's hard is making a product that does this cheaply. That's partly why Google's inference chips (e.g. TPU v5e) have worse specs than an H100 in almost every possible way (and have worse specs than many other TPUs) but are still the best inference solution on the market. Because they're incredibly cheap for what they do.

Even on Google Cloud, which has a markup over real costs, we have 918e12 bfloat16 FLOPs/s at $2.7 / hour, compared to about 1000e12 FLOPs/s from an H100 for $12 / hour. B100s aren't available on AWS yet but even a TPU v5e has 200e12 FLOPs/s at $1.2 / hour, so 1/5 the FLOPs for 1/10th the price.

[s2bk]

Answer 5. "the first byte will arrive in about 6us and the total transfer will take 188us." Wouldn't the total transfer time will be 188 + 6us = 194us?

[samuela]

What's the difference between a "slice" and a "pod"? The article simultaneously states:

• "Chips are connected to each other through the ICI network in a Pod" when introducing a "pod"
• "A set of ICI-connected TPUs is called a slice." when defining a "slice"

[jacobaustin123]

A pod is the maximum slice for a given topology. So e.g. 16x16x16 for TPU v4, or 16x16 for TPU v5e

[sshkhr]

@jacobaustin123 Then what is the difference between a pod and a super pod? From the text:

the maximum pod size (called a superpod) is 16x16x16 for TPU v4 and 16x20x28 for TPU v5p

This is what I found from a different reference (TPU architecture: https://cloud.google.com/tpu/docs/system-architecture-tpu-vm):

TPU Pod
A TPU Pod is a contiguous set of TPUs grouped together over a specialized network. The number of TPU chips in a TPU Pod is dependent on the TPU version.

Slice
A slice is a collection of chips all located inside the same TPU Pod connected by high-speed inter chip interconnects (ICI). Slices are described in terms of chips or TensorCores, depending on the TPU version.

What is this 'specialized network' that pods are grouped over: is it just ICI (I assume not, since that would be a slice)? Or does TPU pods include both the TPUs and host devices (hence PCIe is also included perhaps)?

[jacobaustin123]

I think Google is pretty inconsistent with this terminology. I think basically in my view pod and superpod should be interchangeable, and refer to a slice of the maximum size.

[Kumawat-Akhilesh]

Although ICI connects the TPUs within a pod (e.g., 8,076 TPUs), it would not be sufficient to directly connect a much larger number of TPUs—say 30,000 or even 100,000—as required for extremely large-scale training. In such scenarios, how are these TPUs interconnected?

Are all the TPUs connected individually to the DCN (e.g., via Ethernet) using NICs, similar to how NVIDIA connects GPUs to a scale-out network despite having an NVLink network? Or do the OCS switches connecting the 8,076-TPU pod also interface with DCN switches, forming a hierarchical structure? This hierarchical approach would mean there isn’t a completely separate scale-up (intra-pod) and scale-out (inter-pod) network, as in NVIDIA’s systems, but rather a unified network with a layered design. Could you clarify?

[jacobaustin123]

Beyond the ICI network, TPUs are connected to their hosts via PCIe and the hosts are connected to NICs which connect to DCN. So I think similar to NVIDIA GPUs beyond the NVLink layer.

With that said, JAX/XLA can abstract DCN as yet another axis of a multi-TPU network, so from a software standpoint it's a minimal change from scaling over ICI.

[Kumawat-Akhilesh]

Thank you for the clarification. Does the each TPU get connected to one NIC and hence one port of the DCN switch or because of ICI, we don't need one DCN switch port or one NIC per GPU?

[jacobaustin123]

The details differ a bit generation by generation, but generally there is one NIC per host (so per 4 or 8 TPUs). When it's possible to use ICI, you generally prefer to, but then you can do direct host-to-host transfers over DCN between any TPUs.

[idnm]

I have several questions trying to connect a high-level picture (FLOPs, Bandwidths, etc) with the details of how systolic arrays work.

1. I don't think I understand all the details of how systolic arrays perform matmuls. If there are any additional refs with thorough explanations, that would be highly appreciated!

2. In particular, you state that

When fully saturated the systolic array can perform one bfloat16[8,128] @ bf16[128x128] -> f32[8,128]9 multiplication per 8 clock cycles.

Where does the "8" come from? I like to think that systolic arrays perform a single matrix-vector multiplication per clock cycle, is that a bad mental model?

3. More importantly, is the critical batch size $B\approx 240$ directly related to the MXU geometry, i.e. it's height? Intuitively, it should be, because for smaller $B$ loading weights would take longer than loading activations, and they could not be perfecly overlapped. The number $B\approx 240$ is presented everywhere as the ratio of FLOPs/bandwidth, but I wonder if the bandwidth was adjusted specifically to match the geometry-constrained batch size?

4. Along a similar lines. A general argument shows that for large matmuls compute ($N^3$) starts to dominate the memory ($N^2$) and we should be able to get to the compute-bound regime by simply scaling the matrix size. However, systolic arrays effectively perform $N$ operations per clock cycle in parallel, to their actual compute time also scales as $N^2$ instead of $N^3$. Then, we generally can't compensate insufficient bandwidth by increasing the chip size. This is also clear from the systolic array dynamics. If the weights/activations can be feed at the same rate as they propagate through a systolic array (around $1.5GHz$) we won't be memory-bound.
   So, can the ideal memory BW be alternatively computed from the frequency and geometry of the chip?

p.s. Thanks for the fantastic guide!

[jacobaustin123]

A couple quick answers to these questions:

1. someone linked https://ecelabs.njit.edu/ece459/lab3.php below which is maybe better? You kind of just have to stare at some animations for a while.
2. To some extent it's just a magic number. You can think of the activations flowing through the systolic array at some rate, which is one batch of size 8 every 8 clock cycles. The mental model you suggest isn't really correct since it isn't doing them one at a time, but it's reasonable if it helps you remember.
3. It's related to the size of the MXU in the sense that that determines how many FLOPs/s it can do. If you made the MXU bigger without increasing your HBM bandwidth (or added another MXU), you'd increase this number. I think this ratio is something the designers picked based on their sense of what common ML workloads look like.
4. I think it's rather the opposite. If we reduced the size of the MXU, we'd again reach a compute-bound regime, since we'd reduce the number of FLOPs the machine can do per-byte. That would be wasteful though for many workloads.

Thanks :)

[idnm]

@jacobaustin123 Thanks, that clarifies a lot! Please let me follow up on a particular point, though.

I can't help thinking that batch size $B\approx 240$ being so similar to the size of a TPU $256x256$ is not a coincidence, but a very direct relation.

First, by looking at the pictures with sliding parallelograms, it seems that the width (batch size) of the parallelogram sliding from the left must be at least the size of TPU, otherwise it will be too short and cause bubbles and underutilization. So I'd think that we must have a relation batch size >= TPU width. Is that misguided?

From a slightly different angle. Say, for TPUv5e we have four MXU, which effectively gives a 256x256 systolic array operating at a frequency 1.5 GHz. The HBM bandwidth is 820 GB/s. If we compute the number of bytes per cycle per input (assume we only slide activations from the left, so that there are 256 inputs) we get 820 BG/s / 1.5 GHz / 256 = 2.14 bytes, which is very close to 2 bytes per cycle per input that need to be fed to the systolic array. Is that a coincidence?

Or, framing this another way, if we were to increase the HMB bandwidth by 10x, would the critical batch size drop by 10x? I don't think it will, because the systolic array frequency remains fixed and already saturated, so the ability to feed the inputs faster won't change the processing speed.

[FL33TW00D]

I might be missing something obvious here, forgive me.

The link you provided for the v5p lists the Interchip Interconnect BW (ICI) as 4800Gbps (600GB/s): https://cloud.google.com/tpu/docs/v5p#system_architecture

If it has 6 way interconnect, wouldn't that be 100GB/s? Where does 90GB/s come from?

[jacobaustin123]

It's less than ideal but all TPU docs report slightly different numbers for these bandwidths. Sometimes they're rounded for simplicity, other times there are bottlenecks that limit the peak bandwidth for different operations. I've added a short note.

[tianshub]

This is really insightful. Minor nit: the equation in the answer of Q4 has lhs max{T_math, T_comm} but rhs max{T_comm, T_math}

[mrinal-essential]

Could you give a brief description of how the Optical Switches (OCS) compare & contrast against the Electrical Switches that are the standard in datacenters today? Is it mostly savings for Google or are there advantages that ML practitioners can benefit from?

[jacobaustin123]

https://arxiv.org/pdf/2208.10041 I think provides some good answers. My general sense is that the main advantage is reconfigurability, but I'm not an expert

[Kumawat-Akhilesh]

I see above that ICI bidirectional bandwidth per link for TPU v6e (Trillium) is 180 GB/s. For the newly introduced Ironwood, the ICI bidirectional bandwidth per link is 1.2Tbps or 1200/8= 150 GB/s. But the Ironwood announcement says the ICI for Ironwood is 1.5x of Trillium. How shall I reconcile or I am mis understanding here.
Thank you.

[jacobaustin123]

Ironwood is connected in a 3D torus, so we have an addition 3/2 total bandwidth.

[Kumawat-Akhilesh]

But then it means Ironwood is 1.25x better... Ironwood has 150 GBps while Trillum has 180 GBps.. So Ironwood would be (150/180)*(3/2) = 1.25x times more bandwidth and not 1.5x?

Also this assumes there is one link connecting one TPU to another. So Trillium with 2 d torus would be 4 links and Ironwood with 3 d torus would be 6 links.

At the same time, if I visit Google's Trillium page (https://cloud.google.com/tpu/docs/v6e), Trillum has 3584 Gbps or 448 GBps. With 4 links, I get 112GBps for Trillum. This is much lower than above in the table stated 180 GBps...

[jacobaustin123]

Ironwood also has 180 GBps depending on the workload. As noted in a comment above, these bandwidths differ slightly depending on the workload.

[Kumawat-Akhilesh]

Thank you, Jacob!

[domluna]

I'm a bit confused by the ICI calculation in Q6. Does it assume 15GB is on adjacent TPUs, so 7.5GB can travel on each link in one hop? The way I reasoned about this was that the upper bound would be transferring 1GB from one corner to the other, which would be 6 hops, so 1e9 / 9e10 ~= 11ms * 6 ~= 66ms to get the final GB to the destination TPU. This doesn't factor in latency calculations.

[domluna]

I can't edit above: 66ms < 166ms (the lower bound in the answer). But would latency of the hops be greater than 100ms?

[jacobaustin123]

No, there are many more than one hops. But this is the absolute minimum amount of time from a throughput standpoint. Somehow, 15GB needs to pass through those two links in one direction. Each one has 4.5e10 bytes / s of bandwidth. 15e9 / (4.5e10 * 2) = 166ms. The per-hop reasoning ignores that we're throughput and not latency-bound.

[FrankLong1]

This is phenomenal reading thank you so much for writing this!!

[vorushin]

CS336 (Spring 2025) has a great video about the modern GPU architecture (a nice addendum to the Appendix A): https://www.youtube.com/watch?v=6OBtO9niT00

[RissyRan]

Thanks for the great doc!

I am curious about For most TPU generations, it performs one bfloat16[8,128] @ bf16[128,128] -> f32[8,128] matrix multiply. It mentions LHS (activation) use 8x128 and RHS (weight) uses 128x128. Wondering if we should put opposite when sequence length is increasing? i.e. activation is (batch, seq, model_dim) and weight is (model_dim, intermediate_dim). In recent DeepSeek v3 case, seq length could be up to 128K, and intermediate_dim for moe layer is 2k.

If we could opposite the padding size, as a JAX user, how could we update our code accordingly? Thank you!

[jacobaustin123]

Sorry for missing this comment! I think generally speaking the compiler is smart enough to transpose matrices as needed to get the best performance out of the matmuls in cases like this. So it can choose what goes in the LHS and RHS slots.

[kiankyars]

In exercise three how come We don't additionally have to compute the latency to go from HBM to the VMEM. Is this because it's much smaller than the HBM latency that it doesn't matter?

[kiankyars]

question four*

[jacobaustin123]

If we were asking about latency we would have to think about it, but as a pipeline, slower link -> faster link means we can totally ignore the faster link from a throughput standpoint.

[kiankyars]

So when you say latency, what is that in contrast to? The question asks how long does it take, which I interpret as latency.

[kiankyars]

How do we have a bidirectional speed w/o wraparound in question 6. It says in the article that the latter is necessary for the former.

[jacobaustin123]

We don't actually have bidirectional speed, I mention that speed but in practice we can only send in 2 directions, so we only get 1/4 of that, really

[kiankyars]

That makes sense except for the 1/4 part. In the question we do 4.5e10 * 2 (because we can send in two directions) which is equal to the bidirectional speed of 9e10.

Perhaps you can specify what "that" is referring to in "so we only get 1/4 of that, really"

[jacobaustin123]

I mean if you just look at the target TPU{0,0}, it has two incoming links from below and from the right, and each can send data at 4.5e10 bytes / second. So even though it adds up to the bidirectional speed, it's the unidirectional speed over two links.

[jh-michael-shin]

Hi!

What a great post! I really enjoyed reading them!

I believe it would be a great resource for the AI community here in Korea. With your team's permission, I would love to translate it and share it with them.
I occasionally do translations, and here is an example of my previous work.

• Reasoning in LLMs: https://tulip-phalange-a1e.notion.site/Reasoning-LLMs-190c32470be2806d834ee0ad98aaa0b6
• AI Domain-Specific Architectures: https://tulip-phalange-a1e.notion.site/AI-Domain-Specific-Architectures-23cc32470be28025bbffe42e15e58d99

Of course, I would ensure that full credit and a link back to your original article are prominently featured.

Thank you for your consideration.

Best regards,
Micael

[jacobaustin123]

Hi Micael. Feel free to translate it, happy that you found it useful!