原著：Jacob Austin 等，Google DeepMind，《How to Scale Your Model》。本章原文。原文固定版本：173fb0851d7e，获取日期：2026-10-26，MIT 许可。
本章为完整范围的机器辅助中文初译

简介：本章介绍 TPU 如何工作、如何组成网络以支持多芯片训练和推理，以及这些特点如何影响常见算法的性能。GPU 使用者也能从中获得有用的认识。

Scaling Book 中文导读 · 原著目录

你可能还会喜欢阅读有关 NVIDIA GPU 的新 第 12 节！

什么是TPU？

**TPU 基本上是一个专门从事矩阵乘法的计算核心（称为 TensorCore），附加到快速内存堆栈（称为高带宽内存或 HBM）

图： TPU芯片的基本组成。 TensorCore 是左侧的灰色框，包含矩阵乘法单元 (MXU)、向量单元 (VPU) 和向量内存 (VMEM)。

你可以将 TensorCore 视为基本上只是一个非常好的矩阵乘法机，但它还有一些其他值得注意的功能。 TensorCore 具有三个关键单元：

• MXU（矩阵乘法单元）是 TensorCore 的核心。在 TPU7x 上，它使用脉动阵列在每个 4 周期执行一次 bf16[8,256] @ bf16[256,256] -> f32[8,256] 矩阵乘法（有关详细信息，请参阅 附录 B）。

  • 这大约是 TPU7x 上每个 MXU 在 2.2GHz 下的 5.8e14 bf16 FLOPs/s。每个 TPU7x 芯片有 2 个 TensorCore，每个 TensorCore 有 2 个 MXU，因此 TPU7x 的总 bf16 FLOPs/s 约为 2.3e15。
  • TPU 还支持具有较高吞吐量的较低精度 matmul（例如，每个 TPU7x 芯片可以执行 4.6e15 fp8 OPs/s）。

• VPU（向量处理单元）执行一般数学运算，例如 ReLU 激活或向量之间的逐点加法或乘法。此处还执行归约（求和）。 附录 A 提供了更多详细信息。

• VMEM（矢量内存）是位于 TensorCore 中靠近计算单元的片上暂存器。它比 HBM 小得多（例如，TPU7x 上每个芯片 128 MiB，而 HBM 为 192 GiB），但 MXU 的带宽要高得多。 VMEM 的运行方式有点像 CPU 上的 L1/L2 缓存，但要大得多并且由程序员控制。 HBM 中的数据需要先复制到 VMEM 中，然后 TensorCore 才能使用它进行任何计算。

TPU 在矩阵乘法方面非常非常快。 TPU7x 是迄今为止最强大的 TPU 之一，可以执行 1.15e15 bf16 FLOPs/秒/核心或 2.3e15 bf16 FLOPs/秒/芯片。单个 9216 芯片 Pod 可以执行 21 bf16 exaFLOPs/s。（译注：TPU，特别是它们的脉动阵列，是非常强大的硬件加速器，因为矩阵乘法是使用 $O(n^3)$ 计算 $O(n^2)$ 字节的少数算法之一。这使得普通 ALU 很容易成为计算瓶颈，而不是内存带宽。）

上图还包括一些其他组件，例如 SMEM 和标量单元，它们用于控制流处理，并在 附录 A 中进行了简要讨论，但对于理解它们并不重要。另一方面，HBM 很重要而且相当简单：

• HBM（高带宽内存）是一大块快速内存，用于存储供 TensorCore 使用的张量。 HBM 容量通常为数十至数百 GB 量级（例如 TPU7x 的 HBM 容量为 192GiB）。

  • 当需要计算时，张量通过 VMEM（见下文）从 HBM 流式传输到 MXU，并将结果从 VMEM 写回 HBM。

  • HBM 和 TensorCore 之间的带宽（通过 VMEM）被称为“HBM 带宽”（TPU7x 约为 7TB/秒），它限制了内存密集型工作负载中的计算速度。

一般来说，所有 TPU 操作都是流水线和重叠的。 要执行 matmul $X \cdot A \to Y$，TPU 首先需要将矩阵块 $A$ 和 $X$ 从 HBM 复制到 VMEM 中，然后将它们加载到 MXU 中，MXU 将 8x256 块（对于 $X$）相乘， 256x256（对于 $A$），然后将结果逐块复制回 HBM。为了有效地做到这一点，matmul 是管道化的，因此进出 VMEM 的副本与 MXU 工作重叠。这允许 MXU 继续工作，而不是等待内存传输，从而保持 matmuls 受计算受限，而不是受内存限制。

以下是如何从 HBM 执行元素乘积的示例：

图： 动画显示在 TPU 上执行的逐点乘积，其中字节从 HBM 加载。请注意字节如何以块的形式从内存中流出，以及部分结果如何通过管道传输回来，而无需等待完整的数组被具体化。

matmul 看起来几乎相同，只是它会加载到 MXU 而不是 VPU/Vector 单元，并且加载和存储将以不同的顺序发生，因为相同的权重块用于多个激活块。你可以看到数据块流入 VMEM，然后流入 VREG（向量寄存器），然后流入向量单元，然后返回 VMEM 和 HBM。正如我们即将看到的，如果从 HBM 到 VMEM 的负载比向量单元（或 MXU）中的 FLOP 慢，我们就会受到“带宽受限”，因为我们导致 VPU 或 MXU 的工作量不足。

要点： TPU 非常简单。它们将权重从 HBM 加载到 VMEM，然后从 VMEM 加载到每秒可执行约 290 万亿次乘法加法的脉动阵列。 HBM $\leftrightarrow$ VMEM 和 VMEM $\leftrightarrow$ 脉动阵列带宽对 TPU 可以高效执行的计算设置了基本限制。

VMEM 和算术强度： VMEM 比 HBM 小得多，但它对 MXU 的带宽要高得多。正如我们在第 1 节 中看到的，这意味着如果算法可以在 VMEM 中容纳其所有输入/输出，那么遇到通信瓶颈的可能性就会小得多。当计算的算术强度较差时，这尤其有用：VMEM 带宽比 TPU7x 的 HBM 带宽高约 7 倍，这意味着从 VMEM 读取/写入的 MXU 操作所需的算术强度比使用 HBM 的操作低约 7 倍，以实现峰值 FLOP 利用率。这意味着如果我们可以将权重放入 VMEM 而不是 HBM，那么我们的矩阵乘法就可以在更小的批量大小下实现 FLOP。这意味着从根本上来说算术强度较低的算法仍然是高效的。 VMEM 太小了，这通常是一个挑战。（原注：我们有时会谈论 VMEM 预取，指的是在 VMEM 中提前加载权重，这样我们就可以掩盖 matmul 的加载成本。例如，在普通的 Transformer 中，我们有时可以在注意力期间将大前馈权重加载到 VMEM 中，如果我们受到内存带宽受限，这可以隐藏权重加载的成本。这要求我们的权重足够小或足够分片，以便将单个层放入 VMEM 中，并留出空间。）

TPU 芯片通常（但并非总是）由两个共享内存的 TPU 核心组成，可被视为一个大型加速器，具有两倍的 FLOP（称为“兆核”配置）。对于 v4、v5 和 v6 TPU 来说都是如此（TPU v7 删除了超级核心，而是在两个核心之间建立了高带宽链接）。较旧的 TPU 芯片具有独立的内存，并被视为两个独立的加速器（TPU v3 及更早版本）。 TPU v5e 等推理优化芯片每个芯片只有一个 TPU 核心。

芯片以 **4 组的形式排列在通过 PCIe 网络连接到 CPU 主机的“托盘”上。 这是大多数读者熟悉的格式，4 个芯片（TPU7x 上有 8 个核心，因为每个芯片有两个 TensorCore，尽管一些老一代将这些组合成 4 个逻辑兆核）通过 Colab 或单个 TPU-VM 公开。对于像 TPU v6e 这样的推理芯片，我们每个主机有 2 个托盘，而不是 1 个，但每个芯片只有 1 个核心，因此 8 个芯片 = 8 个核心。（原注：在 Cloud TPU VM 上，每个托盘作为单独 VM 的一部分公开，因此再次有 4 个可见核心。）

PCIe 带宽有限： 与 HBM $\leftrightarrow$ VMEM 链接一样，CPU $\leftrightarrow$ HBM PCIe 连接具有特定带宽，限制从主机内存加载到 HBM 的速度，反之亦然。例如，TPU v4 的 PCIe 带宽为单路 16GB/秒，比 HBM 慢近 100 倍。我们可以将数据加载/卸载到主机 (CPU) RAM 中，但速度不是很快。

TPU网络

译者补充：host、TPU VM、ICI 与 slice（切片）分别是什么？

这里需要区分“CPU 主机怎样管理 TPU”和“TPU 芯片怎样互传数据”两种连接。

host 是 CPU 主机；TPU VM 是运行在 CPU 主机上的虚拟机环境，你在其中运行 Python、JAX 等程序。TPU 芯片则是连接到主机的加速器。所以，“4 块芯片对应一个 VM”是某些机型的资源配置，不是说 4 块 TPU 本身构成一台虚拟机。每台 VM／主机对应多少芯片，要看 TPU 代际和具体配置，不能一律按 4 块计算。

slice 是从同一个 TPU Pod 中分配出来、按指定拓扑通过 ICI 互连的一组 TPU 芯片。它可以由一个 host 管理，也可以跨多个 host。切片的边界不等于 VM 的边界；跨 host 不一定就跨切片。拓扑中的 2x2x2 表示三个方向的芯片数量，乘积是 8 块芯片，不是 8 台 VM。

以每台 host／VM 对应 4 块芯片的 TPU v5p 配置为例，下表每一行都表示申请了一个相应拓扑的切片：

申请的切片拓扑	TPU 芯片数	host／VM 数量	切片类型
2x2x1	4	1	单 host 切片
2x2x2	8	2	跨两个 host 的单切片
4x4x4	64	16	跨十六个 host 的单切片

因此，申请一个 2x2x2 切片后，两台 VM 各管理 4 块 TPU，这 8 块芯片通过 ICI 组成一个整体。ICI 连接的是 TPU 芯片，不是两台 VM 的 CPU 网络接口。即使发送方和接收方 TPU 分属不同 host，切片内的 TPU 张量也可以沿 ICI 传输，不必先复制到 CPU 内存再经过主机网络。单 host 的 4 块芯片之间同样可以使用 ICI，ICI 并不是从第二台 host 才开始出现。

反过来，分别申请两个独立的 2x2x1 切片，再让两台 VM 通过普通网络通信，并不会自动得到一个 2x2x2 ICI 切片。判断时应查看申请或分配的切片资源及其拓扑，而不能仅凭芯片总数、VM 数量或它们是否在同一个 Pod 中推断。多个切片协同运行称为 Multislice：切片内部的 TPU 通信使用 ICI，切片之间的数据传输使用 DCN，经由 host 完成；两个切片也可能位于同一个 Pod 中。

另外，是否具有首尾环绕连接，不是判断切片的标准。v5p 的 2x2x2 虽然没有完整环面拓扑的环绕链路，仍然是一个 ICI 切片。环绕连接影响通信路径和性能，不决定它是不是切片。

本补充根据 Google Cloud 的 TPU v5p 配置表、TPU 系统架构和 Multislice 说明核对，日期为 2026-10-27。

芯片通过 Pod 中的 ICI 网络相互连接。在 Trillium (TPU v6e) 中，ICI（“芯片间互连”）将每个芯片连接到其 4 个最近的邻居（通过边缘链接形成 2D 环面），这对于推理优化芯片来说是典型的。另一方面，TPU7x 中的每个芯片都连接到最近的 6 个邻居（形成 3D 环面），为训练工作负载提供更多 ICI 带宽。请注意，这些连接不通过其主机，它们是芯片之间的直接链接。

环形结构将任意两个节点之间的最大距离从$N$减少到$N / 2$，使得通信速度更快。 TPU 还具有“扭曲环面”配置，将环面包裹在莫比乌斯带状拓扑中，以进一步减少节点之间的平均距离。

TPU pod（通过 ICI 连接）可以变得非常大： TPU7x 的最大 pod 尺寸（称为 superpod）由 9216 个芯片组成，这些芯片由 4x4x4 芯片的可重新配置立方体组成（原注：可重新配置，我们的意思是它们可以使用共享光路开关（OCS）动态地重新整形为 4nx4mx4k 的任何 3D 拓扑。 OCS 只是具有相同 ICI 带宽的可重新配置连接。它只是让我们连接立方体，同时保留环绕链接。）通过光学环绕链接连接，我们可以重新配置以连接非常大的拓扑。

也可以请求更小的拓扑（例如 2x2x1、2x2x2），尽管没有环绕。这是一个重要的警告，因为它通常会使大多数通信的时间加倍。任何完整立方体的倍数（例如 4x4x4 或 4x4x8）都将具有由光学开关提供的环绕。（原注：请注意，2x2x4 不会有任何环绕，因为它们是由仅在完整立方体上可用的光学开关提供的。然而，TPU v6e 8x16 将在长轴上有一个环绕，因为它不使用可重新配置的光纤网络。）

TPU v5e 和 Trillium (TPU v6e) Pod 由单个 16x16 2D 环面组成，沿任何尺寸为 16 的轴都有环绕（这意味着 8x16 在长轴上有环绕）。对于这些代，拓扑无法扩展到超过 16x16 环面，但 Pod 仍然可以通过标准数据中心网络 (DCN) 相互通信，该网络将 TPU 主机相互连接。同样，可以请求更小的拓扑，而无需包装暗调 $<16$。

这种最近邻连接是 TPU 和 GPU 之间的关键区别。 GPU 通过层次结构的交换机连接，这些交换机近似于每个 GPU 之间的点对点连接，而不是像 TPU 那样使用本地连接。通常，节点内的 GPU（H100 为 8 个 GPU，B200 NVL72 为多达 72 个 GPU）直接连接，而较大的拓扑需要每个 GPU 之间的 O(log(N)) 跳。一方面，这意味着 GPU 可以在少量跳数内发送任意数据。另一方面，TPU 的价格要便宜得多（因为 NVLink 交换机价格昂贵），连接起来更简单，并且可以扩展到更大的拓扑，因为每个设备的链路数量和每个设备的带宽是恒定的。阅读更多信息此处。

ICI 相对于 DCN 来说非常快，但仍然比 HBM 带宽慢。 例如，TPU7x 具有：

• 每个芯片的 HBM 带宽为 7.4e12 字节/秒 (7380 GB/秒)。
• 每个芯片的双向 ICI 带宽为 1.2e12 字节/秒（1200 GB/秒），或每轴 4e11 字节/秒（400 GB/秒）。（原注：根据所执行的操作，TPU ICI 链路的带宽略有不同。你通常可以放心使用本文档中的数字。）
• 每个芯片的 DCN（出口）带宽为 1.25e10 字节/秒（12.5 GB/秒）。

这意味着当我们将模型拆分到多个芯片上时，我们需要小心避免跨设备通信速度较慢而成为 MXU 的瓶颈。

多切片训练：一组ICI连接的TPU称为切片。不同的切片可以使用 DCN 相互连接，例如链接不同 Pod 上的切片。由于 DCN 的连接速度比 ICI 慢得多，因此我们应该尝试限制计算等待来自 DCN 的数据的时间。 DCN 是主机到主机的，因此要通过 DCN 将缓冲区从 TPU 传输到 TPU，我们首先需要通过 PCIe 传输到主机，然后通过网络传出，然后通过目标主机网络传入，然后通过 PCIe 传输到 HBM。

译者补充：为什么使用多个切片，而不申请一个更大的 ICI 切片？

如果单个切片的规模足够、资源也能及时获得，通常优先考虑单切片：TPU 张量通信可以留在高速 ICI 网络内。采用多个切片，往往是为了应对下面几种情况：

• 需要的规模超过单个切片支持的上限。单切片能容纳多少芯片，取决于 TPU 型号和支持的配置。超过这个范围，就需要通过 DCN 组合多个切片，继续扩大训练规模。
• 暂时申请不到足够大的拓扑资源。集群虽然还有空闲芯片，却未必能凑出满足所需拓扑的一整个大切片。多个小切片可能更容易获得，让任务更早开始。这里的“连续”指满足互连拓扑要求，不是简单地说芯片摆在相邻位置。
• 需要利用不同 Pod 中的资源。一个切片位于同一个 Pod 内；如果训练需要组合多个 Pod 中的 TPU，就要使用多个切片，通过 DCN 传输跨切片数据。

多切片也可以位于同一个 Pod 中，并不一定跨 Pod。关键是任务使用了几个独立分配的切片：每个切片内部走 ICI，切片之间走 DCN。多切片的取舍，是扩大可用规模或缩短资源等待时间，同时承担更慢的跨切片通信；它并不保证每一步训练都比同等芯片数量的单切片更快。

以上场景与通信方式参见 Google Cloud Multislice 官方说明（2026-10-27 核对）。

要点

• TPU 很简单，在大多数情况下可以被视为连接到内存（超快）、ICI 上的其他芯片（相当快）以及 DCN 上的数据中心其余部分（稍快）的矩阵乘法单元。

• 通信受到我们各种网络带宽的限制（按速度顺序排列）：

  • HBM 带宽：TensorCore 与其关联的 HBM 之间。
  • ICI 带宽：TPU 芯片与其最近的 4 或 6 个邻居之间的带宽。
  • PCIe 带宽：CPU 主机与其关联的芯片托盘之间。
  • DCN 带宽：多个 CPU 主机之间，通常是不通过 ICI 连接的主机。

• 在切片内，TPU 仅通过 ICI 连接到最近的邻居。这意味着切片中远程芯片之间通过 ICI 进行的通信需要首先跳过中间芯片。

• 权重矩阵需要在两个维度上填充到可被 256 整除的大小（TPU v6e 之前的各代为 128），以填充 MXU（事实上，较小的轴会填充到 256）。

• 较低精度的矩阵乘法往往会更快。 TPU v6e 执行 int8/int4 矩阵乘法的速度比 bfloat16 快 2 倍/4 倍，而 TPU7x 执行 fp8 矩阵乘法的速度比 bfloat16 多 2 倍。 VPU 操作仍然在 fp32 或 bf16 中执行（尽管 TPU v6e 之前的各代仅支持 fp32 VPU 操作）。

• 为了避免 TPU 计算单元出现瓶颈，我们需要确保每个通道的通信量与其速度成正比。

工作中遇到的问题

问题 3 [PCIe 操作强度]： 假设我们被迫在主机 DRAM 中存储一个 $\text{bf16}[D, F]$ 类型的大权重矩阵 $A$ 和一批 $\text{bf16}[B, D]$ 类型的激活 $x$ ，并希望对它们进行矩阵乘法。它在单个主机上运行，​​我们使用连接到它的单个 TPU v6e 芯片。你可以假设 $B \ll D$ 和 $F = 4D$ （我们将在以后的章节中看到为什么这些是合理的假设）。我们需要保持 PCIe 上绑定的 FLOP 所需的最小批量大小 $B$ 是多少？假设 PCIe 带宽为 1.6e10 字节/秒。

答案：

答案： 我们要执行$2BDF$浮点运算，每个芯片每秒可以执行9.2e14浮点运算。这需要 $2BDF / 9.2e14$ 秒才能执行。我们必须从 DRAM 加载 $2DF + 2BD$ 字节，并将 $2BF$ 字节写回其中。我们受到 PCIe 传输速度的瓶颈，因此我们需要 $2 \cdot (BD + DF + BF) / 1.6e10$ 秒来与 TPU 传输数据。由于我们希望计算比重量加载花费更长的时间，假设我们可以将所有重量加载与计算重叠，我们需要 $2BDF / 9.2e14 > 2 \cdot (BD + DF + BF) / 1.6e10$。我们可以使用我们的假设 $B \ll D$ 和 $F = 4D$ 来简化它，得到

$$\frac{8BD^2}{9.2 \times 10^{14}} > \frac{8D^2}{1.6 \times 10^{10}}$$

或者

$$B > \frac{9.2 \times 10^{14}}{1.6 \times 10^{10}} \simeq 57{,}500$$

问题 4 [一般 matmul 延迟]： 假设我们想要将权重矩阵 int8[16384, 4096] 乘以大小为 int8[B, 4096] 的激活矩阵，其中 B 是某个未知的批量大小。假设我们从 1 个 TPU v5e 开始。

1. 作为 B 的函数，这个乘法需要多长时间？ 提示：它可能有助于计算从 HBM 加载数组需要多长时间以及乘法实际需要多长时间。哪个是你的瓶颈？
2. 如果我们想在 VMEM 之外运行此操作怎么办？作为 B 的函数需要多长时间？

答案：

答案： （1）我们需要执行的操作数是$2 \cdot 4096 \cdot 16384 \cdot B = 1.3 \times 10^{8} \cdot B$。所以 $T_{\text{math}} = (1.3 \times 10^{8} \cdot B) / 3.94 \times 10^{14}$ 秒。我们需要将 $16384 \cdot 4096 + 4096 \cdot B$ 字节从 HBM 加载到 VMEM，并将 $16384 \cdot B$ 字节从 VMEM 写回到 HBM。这意味着 $T_{\text{comms}} = (6.7 \times 10^{7} + 2 \times 10^{4} \cdot B) / 8.2 \times 10^{11}$ 秒。假设通信和计算尽可能多地重叠，整个乘法将大约花费

$$\max{T_{\text{math}}, T_{\text{comms}}} = \max\left{\frac{1.3 \times 10^{8} \cdot B}{3.94 \times 10^{14}}, \frac{6.7 \times 10^{7} + 2 \times 10^{4} \cdot B}{8.2 \times 10^{11}}\right}$$

当 $\frac{6.7 \times 10^{7} + 2 \times 10^{4} \cdot B}{8.2 \times 10^{11}} < \frac{1.3 \times 10^{8} \cdot B}{3.94 \times 10^{14}}$ 或等效的 $B > 267$ 时，我们将受到 FLOP 限制。这比我们在第 1 节 中得出的 240 数字稍大，因为我们考虑了 $D$ 和 $F$ 的全部影响。

(2) 如果我们从 VMEM 加载，则将 MXU 的 VMEM 带宽视为 HBM $\leftrightarrow$ VMEM 带宽的 22 倍。这将我们的数据加载分母从 8.2e11 变为 1.80e13，我们得到 $B > 11$。请注意，实际上，我们无法将所有 VMEM 带宽用于加载权重矩阵，因此实际上它会更接近 20。

问题 5 [ICI 带宽]： 假设我们有一个 TPU v5e 4x4 切片。假设我们想要将 bf16[8, 128, 8192] 类型的数组从 TPU{0,0} 发送到 TPU{3, 3}。假设 TPU v5e 的每跳延迟为 $1\mu s$。

1. 第一个字节多久会到达目的地？
2. 总共转账需要多长时间？

答案：

答案： 在 TPU v5e 中，我们有 2D 连接。因为我们只有 4x4 切片（没有大小为 16 的轴），所以我们没有环绕连接。因此，我们的目标芯片可以从两个端口接收数据，同样，我们的源芯片可以从两个端口发送数据。我们要传输的数据量是 2 * 8 * 128 * 8192 = 1.7e7 字节。我们可以同时从两个端口进行传输（即向右发送一半数组，向下发送一半数组），因此每秒传输的字节数为 2 * 4.5e10 = 9e10，这意味着传输整个数组大约需要 1.7e7 / 9e10 = 188us（假设我们受到带宽受限）。在 4x4 切片中，我们在芯片 $(0, 0)$ 和 $(3, 3)$ 之间有 6 个跃点，因为对于少于 16 个芯片的轴没有环绕链接。由于每一跳的延迟约为 $1\mu s$，第一个字节将在大约 6us 内到达，并且总传输将花费大约 188 + 6 = 194us，因为最后一个字节在离开源后必须同样遍历六个跳（一般来说，延迟和带宽项是相加的，尽管这里延迟是一个小的修正）。

问题 6 [将所有内容整合在一起，很难]： 假设你有一个大矩阵 A：int8[128 * 1024, 128 * 1024] 在 TPU v5e 4x4 切片上均匀分片，但卸载到每个芯片上的主机 DRAM。假设你要将整个数组复制到 TPU{0, 0} 并将其乘以向量 bf16[8, 128 * 1024]。这需要多长时间？ 提示：使用上面的数字。

答案：

答案： 让我们首先概述我们必须执行的操作。我们的阵列大约有 16GB。从上表中可以看出，TPU v5e 主机具有 4x2 拓扑，因此 4x4 有 2 个主机。因此，由于我们的阵列是均匀分片的，因此每个主机实际上包含阵列的 1/2 块，即 8GB。我们需要将这些块全部复制到 TPU{0,0}，这给了我们两个选择：

1. 我们可以通过 DCN 进行复制，然后通过 PCIe 将整个非分片阵列加载到 HBM 中。
2. 我们可以将分片数组加载到相应的 TPU 上，然后通过 ICI 执行收集，然后在 TPU{0,0} 上执行 matmul。

应该清楚的是，选项（2）更好。与 ICI 相比，DCN 速度较慢，我们更愿意通过许多 PCIe 链路加载大型阵列，而不是仅加载几个（主机 0 上的 8 个）。这是系统部分的图。如上所述，请注意，TPU 通过 ICI 连接到其邻居（甚至跨主机），所有 TPU 都连接到其主机 CPU（通过 PCIe），主机通过 DCN 连接。

每个芯片实际上都有自己的 PCIe 链路连接到其主机，但为了清楚起见，此处仅显示了一个。

现在让我们计算一下每部分需要多长时间：

1. PCIe 负载：我们通过 16 个 PCIe 链路加载 16GB 的块，每个链路的带宽为 1.6e10 字节/秒。因此这将花费大约 63 毫秒。

2. ICI 副本： 每个 TPU 现在有 16GB / 16 = 1GB 的阵列。我们的 ICI 带宽为每条链路 9e10 字节/秒双向，你会从上图中注意到，TPU v5e 上的 4 个 ICI 链路中只有 2 个在 TPU{0,0} 的拓扑中使用。由于 TPU{0,0} 需要以 4.5e10 字节/秒/链接沿 2 个轴接收总共 15GB 的数据，因此我们可以将时间下限 15e9 / (4.5e10 * 2) = 167ms。实际上，这可能无法实现，因为负载非常不均匀，但可能在 2 倍以内。正如你将在第 3 节中看到的，执行完整的 AllGather 也大约需要 16e9 / (4.5e10 * 2)，因此这接近最佳值。

3. HBM $\rightarrow$ MXU 加载： 为了执行最终的 matmul，我们需要通过 HBM 带宽将这 16e9 字节加上 bf16[8, 128 * 1024] 数组（另外 2MB，因此可以忽略不计）加载到 MXU 中，这将占用 16e9 / 8.2e11 = 20ms。

4. FLOPs： 我们总共执行了 $$2 \cdot 8 \cdot 128 \cdot 1024 \cdot 128 \cdot 1024 = 2.7 \times 10^{11}$$ FLOPs，并且由于我们可以执行 1.97e14 bf16 FLOPs/s，因此我们得到 1.4ms。

总时间的上限是所有这些时间的总和，但由于 TPU 通常可以重叠这些操作，因此我们可以将其视为因最慢部分而成为瓶颈的流水线问题。假设这是真的，那么答案至少是 167 毫秒，在不完美重叠的情况下可能接近 200 毫秒。

第 2 章到此结束！第 3 章将讨论数组分片与跨 TPU 通信。

附录

附录 A：有关 TPU 内部结构的更多信息

在这里，我们将更深入地了解 TPU 的内部操作。除非另有说明，我们将提供 TPU7x 的规格。

视觉处理单元

VPU是TPU的矢量运算核心。 VPU 由执行 vadd（向量加法）或 vmax（元素最大值）等元素算术运算的二维 SIMD 向量机（VPU）和一组称为 VREG 的向量寄存器组成，这些寄存器保存 VPU 和 MXU 的数据。

VREG： 每个 TPU7x 核心都有 64 个 32 位 VREG（TPU v4 中有 32 个），每个核心总共提供约 64 * 8 * 128 * 4 = 256kB 的 VREG 内存（或者整个芯片的 2 倍，因为我们有两个核心）。 TPU7x 每个周期可以从 VMEM 加载 3 个寄存器，每个周期向 VMEM 写入 2 个寄存器，但总读取/写入次数限制为 4 次，因此可以是 3 次读取 + 1 次写入或 2 次读取 + 2 次写入。

对于在 float32 输入上运行的程序，VREG 的每个 32 位字包含一个 float32 元素。使用较低精度的数据类型时，可以将多个元素打包到单个 32 位字中，以更有效地使用 VMEM 和 VREG，并且某些操作希望其操作数以这种方式打包。

VPU： VPU 是形状为 (8, 128) 的 2D 向量运算单元，其中 128 维称为车道轴，8 维称为子车道轴。 TPU7x 上的每个（通道、子通道）对包含 4 个彼此独立的标准浮点 ALU。 VPU 在每个 ALU 的一个周期内执行大多数算术指令（例如 vadd 或向量加法），延迟为 2 个周期，因此在 TPU7x 中，你可以在每个周期中将 VREG 中的 4 对 f32 值相加。典型的 VPU 指令可能类似于 {v2 = vadd.8x128.f32 v0, v1}，其中 v0 和 v1 是输入 VREG，v2 是输出 VREG。所有 TPU 代的 ALU 还支持一组元素整数、按位和逻辑运算，并且 TPU v6e 和 TPU7x 还可以在打包 bfloat16 VREG 上执行元素算术运算。

所有通道和子通道在每个周期以纯SIMD方式执行相同的程序，但每个ALU可以执行不同的操作。所以我们可以例如在单个周期中处理 1 个 vadd 和 1 个 vsub，每个过程都在两个完整的 VREG 上运行，并将输出写入第三个 VREG。

小测验[计算 VPU 吞吐量]： 使用上述信息，计算 TPU7x 可以执行多少矢量 FLOP/s。 TPU7x 的时钟速度约为 2.2GHz。

点击此处查看答案。

答案：每个周期，每个内核可以在 8 * 128 ALU 上执行 4 条向量指令。这为我们提供了每个核心的 8 * 128 * 4 FLOPs/cycle，或 8 * 128 * 4 * 2.2e9 = 9e12 FLOPs/s（1.8e13 bfloat16 FLOPs/s，因为打包的 bfloat16 VREG 包含的元素数量是 float32 VREG 的两倍）。请注意，这比每核大约 1.2e15 的 MXU FLOPs/s（大约是 bfloat16 的 64 倍）小得多。

减少：通常，跨子车道维度的通信或减少比跨车道维度更容易。例如，VPU支持通道内洗牌操作，该操作可以在大约一个周期内沿尺寸为8的轴滚动。这可用于沿子通道维度执行有效的缩减（只需按 4、2 和 1 进行混洗，并执行 3 对元素求和）。

跨车道减少要困难得多，并且涉及称为 XLU 或“跨车道单元”的单独硬件单元，该单元速度缓慢且相当昂贵。

与 GPU 的比较： 对于熟悉 NVIDIA GPU 的人来说，VPU 中的每个 ALU 类似于 CUDA 核心，单个 VPU 通道类似于“Warp Scheduler”，即执行 SIMD 算术的通常 32 个 CUDA 核心的集合。车道内减速非常容易，但如果我们需要穿越车道，我们至少需要换乘VMEM/XLU/SMEM，这要慢得多。有关更多详细信息，请参阅GPU 部分。

标量核心

标量核心是TPU的控制单元。它获取并分派所有指令并执行从 HBM 到 VMEM 的传输，并且可以编程来执行标量元数据工作。由于标量核心是单线程的，因此其副作用是 TPU 的每个核心每个周期只能创建一个 DMA 请求。

具体来说，单个标量核心控制一个 VPU（由 4096 个 ALU 组成）、2 个 MXU、2 个 XLU 和多个 DMA 引擎。每单位计算控制的高度倾斜性质是硬件效率的一个来源，但也限制了以任何有趣的方式进行数据相关矢量化的能力。

附录 B：脉动阵列如何工作？

TPU MXU 的核心是 256x256 脉动阵列（TPU v6e 之前的 128x128）。当完全饱和时，脉动阵列可以执行一个bf16[8,256] @ bf16[256,256] -> f32[8,256]（原注：如果你不熟悉这种表示法，它的意思是：将具有bfloat16元素的8x256矩阵乘以具有bfloat16元素的256x256矩阵，并将结果存储在具有float32元素的8x256矩阵中。）每4个时钟周期（原注：MXU可以每4个周期加载8行f32输入并将其转换为bf16，或者每8个周期加载16行bf16输入。无论哪种情况，每 4 个周期相当于 8 行 bf16 输入。（TPUv6e 之前为 8 行）。

• 脉动阵列的核心是 ALU 的 2D 256x256 (=65,536) 网格，每个 ALU 都能够执行乘法和加法运算。
• 权重（W，256x256 输入）从上方（称为 RHS）向下传递，而输入（X，8x256 输入）从左侧（称为 LHS）传递。

这是一组权重（蓝色）与一组激活（绿色）相乘的简化动画。你会注意到，权重 (RHS) 首先沿对角线方向部分加载，然后激活也沿对角线方向输入。在下面的每一帧中，我们将所有重叠的绿色和蓝色单位相乘，将结果与从上面传入的任何残差相加，然后将结果依次向下传递一个单位。

这是该动画的更通用版本，显示了从计算中流出的输出：

下面的图表显示了如何跨多个 RHS 和 LHS 阵列进行管道传输：

加载权重 (RHS) 和激活 (LHS) 时会出现初始管道气泡。在初始气泡之后，可以加载新的输入和权重，而无需额外的气泡。

这是 bf16[2, 3] x bf16[3, 3] 矩阵乘法的糟糕动画，你可以将其想象为 2x3 权重矩阵的 matmul，输入激活为批次 1，大小为 3。与之前的幻灯片相比，这是旋转的，输入向右而不是向下流动，但你可以大致看到结构。

我们可以有效地将其管道化以乘以大型矩阵，而不会出现太大的管道气泡。话虽如此，重要的是我们的矩阵形状必须大于 MXU 的边尺寸（128x128 或 256x256）。某些 TPU（自 TPU v3 起）具有多个 MXU，TPU v3、TPU v6e 和 TPU7x 为 2 个，TPU v4/5 为 4 个，因此我们需要确保平铺尺寸大于 256（对于旧代为 128）* MXU 数量。 这里 一个很好的动画。

这篇博文 还有另一个关于固定权重矩阵的脉动数组乘法的优秀动画。

---

译文与许可：本章是基于固定版本 173fb0851d7e654872f0327aff2e17607cf24106 的机器辅助中文初译，并非逐句人工精校。原文与配图基于 MIT License 发布。Copyright (c) 2022 Maruan Al-Shedivat；Copyright (c) 2025 Google LLC。完整许可见 MIT License。