Abstract
This article delves into the hardware architecture and design philosophy of Google’s Tensor Processing Unit (TPU). It covers the components of a single TPU chip (including matrix multiplication units, vector processing units, and memory hierarchy) to the interconnection topology of multiple chips (such as 2D/3D toroidal connections and optical switching technology). The article discusses in detail how TPUs achieve extremely high matrix computation throughput through pulsed arrays, optimize energy efficiency with ahead-of-time compilation, and utilize reconfigurable interconnect topologies to support flexible scaling from a few chips to thousands. Additionally, it analyzes the performance impacts of different parallel strategies (data parallelism, pipeline parallelism, tensor parallelism) on various TPU slice topologies, supplemented by real deployment cases (such as PaLM training) to illustrate their large-scale applications.
Table of Contents
- 1. Background Introduction
- 2. TPU Single-Chip Architecture
- 3. TPU Design Philosophy
- • Pulsed Arrays and Pipelines
- • Ahead-of-Time Compilation and Caching Strategies
- • Board-Level Interconnection (4 Chips)
- • Rack-Level Interconnection (64 Chips)
- • Full Pod-Level Interconnection (4096 Chips)
- • Multi-Pod Level Interconnection
Main Text
1 Background Introduction
The TPU (Tensor Processing Unit) is a dedicated integrated circuit (ASIC) designed by Google specifically for machine learning workloads. Its core design goals focus on two key factors: extremely high matrix multiplication throughput and excellent energy efficiency.
The origin of the TPU dates back to 2006 when Google was evaluating whether to use GPUs, FPGAs, or custom ASICs to meet its computational needs. Initially, only a few applications required dedicated hardware, which could be satisfied by the surplus CPU resources in data centers. This situation changed in 2013 when Google’s voice search began to rely on neural networks, and internal predictions indicated that if this feature became widespread, computational demands would increase significantly.
Today, TPUs support most of Google’s AI services, including the training and inference of models like Gemini and Veo, as well as the deployment of recommendation models (such as DLRM).
2 TPU Single-Chip Architecture
The following analysis focuses on TPUv4, but its architectural layout is also applicable to the latest generation of TPUs (such as TPUv7 “Ironwood,” which has not yet released detailed information as of June 2025).
The layout of a single TPUv4 chip is as follows:
Each chip contains two TPU TensorCores responsible for computation tasks (note: the TPU dedicated to inference has only one TensorCore). The two TensorCores share memory units: CMEM (128MiB) and HBM (32GiB).
Each TensorCore internally contains computation units and small memory buffers:
- 1. Matrix Multiplication Unit (MXU)
- • This is the core component of the TensorCore, consisting of a 128×128 pulsed array.
- • Details of the pulsed array will be elaborated later.
- • Handles general element-wise operations (such as ReLU, element-wise addition/multiplication, reduction operations).
- • Memory buffer. Data must be copied from HBM to VMEM before the TensorCore can perform computations.
- • Controls the operations of the VPU and MXU.
- • Manages control flow, scalar operations, and memory address generation.
For readers familiar with NVIDIA GPUs, the following observations may differ:
- 1. The on-chip memory units (CMEM, VMEM, SMEM) on the TPU are significantly larger than the L1 and L2 caches of GPUs (H100’s L1 and L2 caches are 256KB and 50MB, respectively).
- 2. The HBM capacity of the TPU is much smaller than that of GPUs (H100 has 80GB).
- 3. The number of cores responsible for computation is noticeably fewer.
This contrasts with GPU designs: GPUs have smaller L1 and L2 caches, larger HBM, and tens of thousands of cores.
It is important to note that TPUs can also achieve extremely high throughput. TPU v5p can reach 500 TFLOPs/s per chip, and a complete Pod containing 8960 chips can provide approximately 4.45 ExaFLOPs/s. The latest “Ironwood” TPUv7 is reportedly capable of achieving 42.5 ExaFLOPS/s per Pod (9216 chips).
Understanding how TPUs achieve this performance requires a deep dive into their design philosophy.
3 TPU Design Philosophy
TPUs achieve exceptional throughput and energy efficiency through two main pillars and one key assumption: pulsed arrays and pipeline technology, ahead-of-time (AoT) compilation, and the assumption that most operations can be mapped to pulsed arrays. In modern deep learning, matrix multiplication constitutes the bulk of computations, making it ideally suited for processing by pulsed arrays.
3.1 Pulsed Arrays and Pipelines
Pulsed Arrays are a hardware design architecture composed of a grid of interconnected processing elements (PEs). Each PE performs small computations (such as multiply-accumulate operations) and passes the results to adjacent PEs.
The advantage of this design is that once data enters the pulsed array, no additional control logic is needed to guide data processing. Moreover, when the pulsed array is sufficiently large, there are no other memory read/write operations apart from input and output.
Due to its rigid organizational structure, the pulsed array can only handle operations with fixed data flow patterns, but matrix multiplication and convolution fit this pattern perfectly.
Additionally, pipeline technology can clearly overlap computation with data movement. The following diagram illustrates the pipelined element-wise operations on the TPU.
Limitations of Pulsed Arrays: Sparsity Pulsed arrays excel at handling dense matrices (i.e., where each PE is active in nearly every cycle). However, their drawback is that they do not provide performance improvements for sparse matrices of the same size: even for zero-value elements, PEs still need to work, executing the same number of cycles.
Addressing this systemic limitation of pulsed arrays will become increasingly important, especially in the context of the deep learning community’s preference for irregular sparsity (such as mixture of experts models, MoE).
3.2 Ahead-of-Time Compilation and Reduced Cache Dependency
This section explains how TPUs avoid using caches through hardware-software co-design (TPU + XLA compiler) to achieve high energy efficiency.
Traditional cache designs are intended to handle unpredictable memory access patterns. Different applications may have vastly different memory access patterns. Essentially, caches allow hardware to flexibly adapt to a wide range of applications. This is a significant reason why GPUs (compared to TPUs) are very flexible.
However, cache access (and general memory access) consumes a lot of energy. The following diagram provides a rough estimate of the energy costs of operations on the chip (45nm, 0.9V; [18]). The key point is that memory access and control consume most of the energy, while arithmetic operations consume significantly less energy.
But what if the application is highly specific, and its computation/memory access patterns are highly predictable?
As an extreme example, if the compiler can determine all required memory accesses in advance, the hardware can rely solely on register memory as a buffer, without needing caches.
This is precisely what the TPU design philosophy aims for, and it is why the TPU is co-designed with the XLA compiler to achieve this goal. The XLA compiler generates optimized programs by analyzing the computation graph ahead of time.
Collaboration of JAX and TPU JAX+XLA operates in a hybrid space of JIT (just-in-time compilation) and AOT (ahead-of-time compilation) on TPUs, which can sometimes be confusing. When a jitted function in JAX is called for the first time, JAX tracks it to create a static computation graph. This graph is passed to the XLA compiler, which converts it into a fully static binary file specific to the TPU. It is at this final conversion stage that TPU-specific optimizations (such as minimizing memory access) are performed to customize the process.
However, it is important to note: if the jitted function is run with different input shapes, it needs to be recompiled and cached. This is why JAX does not perform well when handling any dynamic padding or loops of varying lengths that depend on input.
Of course, this approach sounds good, but it also has its inconveniences, namely a lack of flexibility, and heavily relying on the compiler is a double-edged sword.
TPU and Energy Efficiency (TPUv4) The previous energy usage graph does not accurately represent the TPU; below is a breakdown of energy usage for TPUv4. Note that TPUv4 uses a 7nm process, and the graph includes 45nm data for comparison ([3], [16]).
Energy per Operation (TPUv4, 7 nm)
The left bar chart visually displays the values, but it is important to note that modern chips use HBM3, which consumes far less energy than the DDR3/4 DRAM shown here. Nevertheless, the chart still indicates that memory operations consume energy at several orders of magnitude higher.
This aligns well with the modern scaling law: we are very willing to increase FLOPS in exchange for reducing memory operations. Therefore, reducing memory operations has dual optimization benefits, as they not only make programs faster but also consume significantly less energy.
4 TPU Multi-Chip Architecture
4.1 Board-Level (also known as “Board Card”; 4 Chips)
A single TPU board contains 4 TPU chips or 8 TensorCores (referred to as “cores”). Each board has its own CPU host (note: for inference TPUs, one host accesses 2 boards since each chip has only 1 core).
The host ⇔ chip connection is via PCIe, while the chip ⇔ chip connection is through inter-core interconnect (ICI), which has higher bandwidth.
However, ICI connections can be further extended to multiple board cards. For this, we need to scale up to the rack level.
4.2 Rack-Level (4x4x4 Chips)
The particularly exciting aspect of TPU is its scalability, which becomes evident at the rack level.
A TPU rack consists of 64 TPUs connected in a 4x4x4 3D toroidal manner. If you have seen Google’s promotional materials for TPUs, this image shows 8 TPU racks.
Before delving deeper into the rack, it is important to clarify some easily confused terms: Rack, Pod, and Slice.
Differences Between Rack, Pod, and Slice Different Google sources use them slightly differently, and sometimes “TPU Pod” and “TPU Slice” can be used interchangeably. However, in this article, we will follow the definitions used in Google’s TPU papers and GCP TPU documentation ([3], [7], [9]).
- 1. TPU Rack:
- • A physical unit containing 64 chips. Also referred to as a “cube.”
- • The largest TPU unit connected via ICI and fiber optics.
- • Also known as “SuperPod” or “Full Pod.” For example, a TPUv4 Pod consists of 4096 chips or 64 TPU racks.
- • Any TPU configuration between 4 chips and SuperPod size.
The key distinction is that TPU racks and TPU Pods are physical measurement units, while TPU slices are an abstract unit. Of course, setting up TPU slices has important physical properties, but we will temporarily abstract that point.
Currently, we will use physical measurement units: TPU racks and TPU Pods. Understanding how TPU systems are physically connected will help us better understand the TPU design philosophy.
Returning to the TPU rack (TPUv4): A single TPU rack consists of 64 chips connected via ICI and optical circuit switching (OCS). Essentially, we connect multiple board cards together to simulate a 64-chip system. This theme of combining small components into a supercomputer will continue to appear later.
The following diagram is a schematic of a single TPUv4 rack. It is a 4x4x4 3D toroidal structure, where each node is a chip, the blue arrows represent ICI, and the lines on the surface represent OCS (redrawn according to [7]).
However, this diagram raises several questions.Why is OCS only used on the surface? In other words,what are the benefits of using OCS? There are three major benefits, and we will introduce two more below.
Benefit of OCS #1: Surrounding Connections Faster communication between nodes is achieved through surrounding connections. OCS also serves as the surrounding connection for a given TPU configuration. This reduces the worst-case number of hops between two nodes from (N-1) hops per axis to (N-1)/2 hops, as each axis becomes a ring (1D toroidal). As the scale increases, this effect becomes more significant, as reducing communication latency between chips is crucial for high parallelization.
Not all TPUs have 3D toroidal topology Note: Older TPU generations (such as TPUv2, v3) and inference TPUs (such as TPUv5e, TPUv6e) have a 2D toroidal topology instead of the 3D toroidal shown below. However, TPUv7 “Ironwood” appears to be 3D toroidal, although it is marketed as an inference chip (note: this judgment is based solely on its promotional materials).
4.3 Full Pod Level (also known as “SuperPod”; TPUv4 has 4096 chips)
Just as we connect multiple chips to form a TPU rack, we can connect multiple racks to form a large SuperPod. A SuperPod also refers to the maximum configuration of interconnected chips that TPU can achieve (using only ICI and OCS). There is also a multi-Pod level, but this must be achieved through slower interconnections, which we will discuss later. This varies by generation, but for TPUv4, it is 4096 chips (i.e., 64 racks, each with 4x4x4 chips). For the latest TPUv7 “Ironwood,” this is 9216 chips. The following diagram shows a SuperPod of TPUv4.
Note how each cube (i.e., a TPU rack) is interconnected via OCS. This also allows us to obtain TPU slices within a Pod.
Using OCS for TPU Slices We can request a subset of TPUs within a Pod, which are TPU slices. However, even if you want N chips, there are various topologies to choose from. For example, suppose you need a total of 512 chips. You could request a cube (8x8x8), a cigar shape (4x4x32), or a rectangular shape (4x8x16). Choosing the topology of the slice itself is a hyperparameter. The topology you choose will affect the communication bandwidth between nodes. This directly impacts the performance of different parallel methods. For instance, a cube (e.g., 8x8x8) is better suited for all-to-all communications, such as data parallelism or tensor parallelism, because it has the highest bipartite bandwidth. However, a cigar shape (e.g., 4x4x32) is more suitable for pipeline parallelism because it can communicate faster with sequential layers (assuming one layer fits a 4×4 chip sub-slice).
Of course, the optimal topology depends on the model, and finding it is a task in itself. The TPUv4 paper [9] also measured this to show how topology variations can accelerate throughput (note: I am not sure which LLM architecture the first line refers to, as it is not specified).
We have introduced TPU slices, but there is an important feature that contributes to the high operational stability of TPUs. Due to the presence of OCS, these slices do not have to be contiguous racks. This is the second benefit of using OCS—possibly the greatest—that we have not mentioned before.
Benefit of OCS #2: (Reconfigurable) Non-contiguous Multi-node Slices Note that this is different from hard-wiring multiple nodes to simulate non-contiguous slices. Because OCS is a switch rather than a hard connection, there are far fewer physical lines between nodes, allowing for greater scalability (i.e., larger TPU Pod sizes). This allows for large-scale flexible node configurations. For example, suppose we want to run three jobs on a Pod. While naive scheduling would not allow this, OCS connections enable us to abstract away the location of nodes and treat the entire Pod as merely a “bag of nodes” (redrawn according to [6]).
This improves Pod utilization and may make maintenance easier when nodes fail. Google describes this as “small dead node explosion radius.” However, I am unsure how its liquid cooling would be affected when only certain nodes need to be shut down. Finally, this flexible OCS has an interesting extension: we can also change the topology of TPU slices, for example, from a regular toroidal to a twisted toroidal.
Benefit of OCS #3: Twisted TPU Topology Previously, we saw how to achieve different TPU slice topologies by changing the (x,y,z) dimensions of a fixed number of chips. This time, however, we will deal with fixed (x,y,z) dimensions but change their connectivity to achieve different topologies. A notable example is transforming a cigar-shaped regular toroidal into a twisted cigar toroidal, as shown below.
The twisted toroidal case allows for faster communication between chips across the twisted 2D plane. This is particularly useful for accelerating all-to-all communications. Let’s delve a bit deeper and imagine a specific scenario where this would be beneficial.
Accelerating Training with Twisted Toroidal Theoretically, the twisted toroidal will bring maximum benefits for tensor parallelism (TP) since each layer has multiple all-gather and reduce-scatter operations. It may also provide moderate benefits for data parallelism (DP) since each training step also has an all-reduce, but this will be less frequent. Suppose we are training a standard decoder-only Transformer and we want to adopt a large amount of parallelism to accelerate training. We will see two scenarios below.
Scenario #1: 4x4x16 Topology (TP + PP; Total 256 Chips) Our z-axis will be our pipeline parallel (PP) dimension, and our 2D TP dimension will be 4×4. Essentially, assume each layer k is located at z=k, and each layer is sliced across 16 chips. Assume standard OCS connections are used without explicit drawing (i.e., nearest neighbors).
We will twist the 2D toroidal at each z=k, which makes communication between chips within each TP layer faster. Twisting along our PP dimension is unnecessary since they primarily rely on point-to-point communication.Note: In practice, the twisted toroidal will be beneficial when the number of chips exceeds 4×4. We use 4×4 only for visualization purposes.
Scenario #2: 16x4x16 Topology (DP + TP + PP; Total 1024 Chips) As an extension, we will add a size 4 DP dimension to the previous scenario. This means there are 4 instances of Scenario #1 along the x-axis.
Note that the twisted toroidal is limited to each TP dimension of each DP model (i.e., for each z=k, given k=1…16, a 4×4 2D plane). The DP dimension has only one surrounding, so each row becomes a horizontal ring of size 16. You may have noticed another topology of 8x8x16 (i.e., 2×2 DP dimension), but this becomes more complex as we mix DP and TP dimensions together. Specifically, it is unclear how we should build the OCS surrounding for the y-axis while accommodating the twisted toroidal for each TP dimension.
4.4 Multi-Pod Level (also known as “Multi-Slice”; TPUv4 has 4096+ chips)
The final level of the TPU hierarchy is the multi-Pod level. Here, you can think of multiple Pods as a large machine. However, communication between Pods occurs through data center networks (DCN), which have lower bandwidth than ICI.
This diagram shows how to configure multi-Pod training. PaLM was trained this way. It took 56 days to train on 6144 TPUv4 (2 Pods). Below you can see the TPU job allocation across 6 Pods: green for PaLM, red for unallocated, and the rest for other jobs. Note that each square is a 4x4x4 TPU cube.
Achieving this is challenging in itself, but what is even more impressive is the focus on developer experience. Specifically, the emphasis is on the question, “How can we abstract the system/HW parts of model scaling as much as possible?” Google’s answer is to let the XLA compiler handle coordinating communication across large-scale chips. By providing the correct flags (i.e., DP, FSDP, TP parallel dimensions, number of slices, etc.) to researchers, the XLA compiler inserts the correct hierarchical collective operations for the TPU topology at hand (Xu et al, 2021: GSPMD [2]). The goal is to enable large-scale training with as few code changes as possible.
For example, the following is a breakdown of the all-reduce operation across multiple slices from Google’s blog [1].
This indicates that the XLA compiler is responsible for handling collective communication operations between and within slices. To give a specific example, the TPU topology for training a model might look like this. Activation communication occurs within slices via ICI, while gradient communication occurs across slices via DCN (i.e., across DCN DP dimensions) [1].
5 Actual System Correspondence
I find it helpful to place charts in the background when there are actual photos of the hardware. Below is a summary. If you have seen images from Google’s TPU promotional materials, you may have encountered the following image.
This shows 8 TPU Pods, where each unit is a 4x4x4 3D toroidal structure as we saw above. Each row in a Pod has 2 board cards, meaning each row has 8 TPU chips. Here is a single TPUv4 board card:
Note that this diagram simplifies to only one PCIe port, but in the actual board card, there are 4 PCIe ports (on the left)—one for each TPU. Below is a single chip:
The center part is the ASIC, surrounded by 4 HBM stacks. This is the TPU v4 we see, so there are 2 TensorCores inside, resulting in a total of 4 HBM stacks. I could not find a layout diagram for TPUv4, so here is a layout diagram for TPUv4i, which is similar, except that it has only 1 TensorCore because it is an inference chip [3].
Note that CMEM occupies a considerable amount of space in the layout of TPUv4i.
6 Acknowledgments
Thanks to Google TPU Research Cloud (TRC) for providing TPU support!
7 References
[1] Google Blog: TPU Multi-Slice Training[1][2] Xu, et al. “GSPMD: General and Scalable Parallelization for ML Computation Graphs”[2][3] Jouppi et al. “Ten Lessons From Three Generations Shaped Google’s TPUv4i”[3][4] How to Scale Your Model – TPUs[4][5] Domain Specific Architectures for AI Inference – TPUs[5][6] HotChips 2023: TPUv4[6][7] Google Cloud Docs: TPUv4[7][8] Jouppi et al. “In-Datacenter Performance Analysis of a Tensor Processing Unit” — TPU origins paper[8][9] Jouppi et al. “TPU v4”– TPUv4 paper[9][10] PaLM training video[10][11] HotChips 2021: “Challenges in large scale training of Giant Transformers on Google TPU machines”[11][12] HotChips 2020: “Exploring Limits of ML Training on Google TPUs”[12][13] Google Blog: Ironwood[13][14] HotChips 2019: “Cloud TPU: Codesigning Architecture and Infrastructure”[14][15] ETH Zurich’s Comp Arch Lecture 28: Systolic Array Architectures[15][16] Patterson presentation: “A Decade of Machine Learning Accelerators: Lessons Learned and Carbon Footprint”[16][17] Camara et al. “Twisted Torus Topologies for Enhanced Interconnection Networks.” [17][18] Horowitz article: “Computing’s Energy Problem(and what we can do about it)”[18]
This article is translated from TPU Deep Dive[19]
Reference Links
<span>[1]</span> Google Blog: TPU Multi-Slice Training: https://cloud.google.com/blog/products/compute/using-cloud-tpu-multislice-to-scale-ai-workloads<span>[2]</span> Xu, et al. “GSPMD: General and Scalable Parallelization for ML Computation Graphs”: https://arxiv.org/pdf/2105.04663<span>[3]</span> Jouppi et al. “Ten Lessons From Three Generations Shaped Google’s TPUv4i”: https://gwern.net/doc/ai/scaling/hardware/2021-jouppi.pdf<span>[4]</span> How to Scale Your Model – TPUs: https://jax-ml.github.io/scaling-book/tpus/<span>[5]</span> Domain Specific Architectures for AI Inference – TPUs: https://fleetwood.dev/posts/domain-specific-architectures#google-tpu<span>[6]</span> HotChips 2023: TPUv4: https://hc2023.hotchips.org/assets/program/conference/day2/ML%20training/HC2023.Session5.ML_Training.Google.Norm_Jouppi.Andy_Swing.Final_2023-08-25.pdf<span>[7]</span> Google Cloud Docs: TPUv4: https://cloud.google.com/tpu/docs/v4<span>[8]</span> Jouppi et al. “In-Datacenter Performance Analysis of a Tensor Processing Unit” — TPU origins paper: https://arxiv.org/abs/1704.04760<span>[9]</span> Jouppi et al. “TPU v4”– TPUv4 paper: https://arxiv.org/abs/2304.01433<span>[10]</span> PaLM training video: https://www.youtube.com/watch?v=0yPFBxkOKRY<span>[11]</span> HotChips 2021: “Challenges in large scale training of Giant Transformers on Google TPU machines”: https://hc33.hotchips.org/assets/program/tutorials/HC2021.Google.Sameer%20Kumar.pdf<span>[12]</span> HotChips 2020: “Exploring Limits of ML Training on Google TPUs”: https://hc32.hotchips.org/assets/program/tutorials/HC2020.Google.SameerKumarDehaoChen.v02.pdf<span>[13]</span> Google Blog: Ironwood: https://blog.google/products/google-cloud/ironwood-tpu-age-of-inference/<span>[14]</span> HotChips 2019: “Cloud TPU: Codesigning Architecture and Infrastructure”: https://old.hotchips.org/hc31/HC31_T3_Cloud_TPU_Codesign.pdf<span>[15]</span> ETH Zurich’s Comp Arch Lecture 28: Systolic Array Architectures: https://www.youtube.com/watch?v=XkgtANeDrm8<span>[16]</span> Patterson presentation: “A Decade of Machine Learning Accelerators: Lessons Learned and Carbon Footprint”: https://www.cs.ucla.edu/wp-content/uploads/cs/PATTERSON-10-Lessons-4-TPU-gens-CO2e-45-minutes.pdf<span>[17]</span> Camara et al. “Twisted Torus Topologies for Enhanced Interconnection Networks.”: https://personales.unican.es/vallejoe/Publications/C%C3%A1mara%20-%20TPDS’10%20-%20Twisted%20Torus%20Topologies%20for%20Enhanced%20Interconnection%20Networks.pdf<span>[18]</span> Horowitz article: “Computing’s Energy Problem(and what we can do about it)”: https://gwern.net/doc/cs/hardware/2014-horowitz-2.pdf<span>[19]</span> TPU Deep Dive: https://henryhmko.github.io/posts/tpu/tpu.html