TPU 的 4x4x4 到底在乘什么？一篇看懂 TPU Slice 与 GPU 设计哲学

刚开始学习 TPU 架构时，规格中 4x4x4、8x16、16x16 这类 slice（切片）尺寸让我十分困惑。当下冒出一连串疑问：为什么是“几乘几”？这些数字是矩阵乘法的维度吗？为什么 GPU 的世界好像不用管这些，TPU 却要把拓扑搞得这么复杂？

之前在 AI 有意识吗？为什么 ChatGPT 这么好用却常一本正经地出错 ¹ 中拆解过 LLM 的本质：一台用矩阵计算进行概率预测的机器。这篇文章想继续深入硬件层：当矩阵大到一颗芯片装不下时，Google 是怎么把一个数学问题，变成一座机房的接线问题？

结论在前：TPU slice 的“几乘几”（例如 4x4x4 或 8x16）指的是芯片在数据中心网络中的物理排列拓扑，不是矩阵乘法在数学上的切分方式。这两者是两个不同层次的问题，中间由 XLA 编译器负责对应。而 TPU 之所以选择 2D/3D 环面（torus）拓扑，而不是像 GPU 集群那样构建一套昂贵的交换网络，本质上是一个成本与复杂度的取舍：Google 用软件与编译器的复杂度，换取硬件规模化时的极致性价比。

TL;DR：四个核心判断

作为我个人的阅读与消化笔记，本篇讨论以下内容：

• 深度学习的本质是矩阵乘法：神经网络在前向与反向传播时，多数计算都可化为矩阵计算，这也正是 TPU 把整颗芯片押注在“脉动阵列”上的原因。
• Slice 的“几乘几”是芯片接线图，不是矩阵切分：4x4x4 描述的是 64 颗芯片的 3D 网格排列；矩阵怎么切、切给谁算，是 XLA 编译器把逻辑网格“贴”到物理拓扑上的结果。
• TPU 不是“完全没有交换机”：机柜内的 cube 用铜缆直连，cube 之间其实靠 OCS（光学电路交换）串接。真正省下的是昂贵的电子数据包交换机，OCS 与光学元件只占整体系统成本不到 5%。
• GPU 的复杂拓扑没有消失：GPU 集群把复杂度藏进硬件账单（NVSwitch、InfiniBand 胖树等拓扑结构）里，而这些硬件元件在数据中心往往成本高昂。TPU 则省下这部分硬件开销，把复杂度交给 XLA 编译器与工程师。

为什么深度学习离不开矩阵乘法？

在以神经网络构建的深度学习中，每一层本质上是“一组输入乘上一组权重，再加总”的线性变换，再接一个非线性激活函数：每个神经元把上一层的输出各自乘上权重、相加、过一个非线性函数。当你把一整层神经元、整个批次（batch）数据同时考虑进来，这些成千上万个“乘法加总”就能自然地打包成一个矩阵乘法。

---
config:
  layout: dagre
  flowchart:
    curve: linear
    nodeSpacing: 15
    rankSpacing: 110
---
flowchart LR
    subgraph IN["输入层（上一层的输出）"]
        x1(("x1"))
        x2(("x2"))
        x3(("x3"))
    end
    subgraph HID["隐藏层神经元"]
        h1(("h1"))
        h2(("h2"))
        h3(("h3"))
        h4(("h4"))
    end
    subgraph OUT["输出层（传给下一层）"]
        y1(("y1"))
        y2(("y2"))
    end
    x1 --> h1 & h2 & h3 & h4
    x2 --> h1 & h2 & h3 & h4
    x3 --> h1 & h2 & h3 & h4
    h1 --> y1 & y2
    h2 --> y1 & y2
    h3 --> y1 & y2
    h4 --> y1 & y2

图中每一条箭头都代表一次“乘上权重”：一个神经元的输出，就是所有进入箭头的加权总和再过一个非线性函数。把一整层的箭头合起来看，就是输入向量乘上权重矩阵；再把整个批次数据叠起来，就是一个完整的矩阵乘法 X × W。

打包成矩阵有两个巨大的工程好处：

1. 并行化：矩阵乘法里的大量乘加计算彼此独立，可以交给成千上万个计算单元同时做，完美匹配 GPU / TPU 这类高度并行计算的硬件。
2. 数据重用：矩阵乘法中同一笔数据会被多次使用，硬件可以让数据在芯片内部流动、重复利用，减少昂贵的内存存取。在现代芯片中，搬动数据的能耗往往远高于计算本身。

从 Transformer 到今天的 LLM，底层跑的几乎都是同一件事：把 attention、前馈网络全部化约成一连串巨大的矩阵乘法。

flowchart LR
    A["神经网络的一层<br>数千个神经元 × 数千维输入"] --> B["打包成矩阵乘法<br>输入矩阵 X × 权重矩阵 W"]
    B --> C["高度并行的硬件<br>GPU / TPU 一次完成数万个乘加运算"]
    C --> D["数据在芯片内流动重用<br>减少内存搬运与能耗"]

TPU 的核心：脉动阵列 MXU

TPU（Tensor Processing Unit）是 Google 为此打造的专用芯片，主要任务是执行矩阵计算，也就是大量的乘法与累加。它内建数千甚至上万个彼此直连的乘加器，组成一个大型的实体计算阵列。

它的核心是 MXU（Matrix Multiply Unit，矩阵乘法单元），采用一种叫“脉动阵列”（systolic array）的设计 ²。

可以把脉动阵列想象成一条精心编排的工厂流水线：权重先被“预加载”到一个二维计算单元网格中并保持不动，输入数据则像水流一样从侧面一波波流入。每个计算单元完成“乘法 + 累加”后，再把结果向下传递。数据在网格中规律地“脉动”，无须反复读写内存——这正是它比通用处理器更加节能、高效的原因。

TPU 主机会将数据串流至输入（infeed）队列。TPU 再从输入队列载入数据，并将其储存到高带宽内存（HBM）。计算完成后，TPU 会把结果写入输出（outfeed）队列；接着，TPU 主机从输出队列读取结果，并储存在主机内存中。若要执行矩阵计算，TPU 会将参数从 HBM 载入矩阵乘法单元（MXU）。

各世代 TPU 的 MXU 规格与拓扑如下 ^{2 3 4 5 6 7 8 9}：

这张表有两个值得注意的细节：

• MXU 的边长从 v2 到 v5p 一直维持 128x128，到了第六代 Trillium（v6e）才扩大为 256x256，并由第七代 Ironwood（TPU7x）延续。边长翻倍，意味着每个时钟周期的乘加计算量变成四倍 ^{2 7 10}。
• 这个数字对开发者有实际意义：要充分利用 MXU，矩阵维度最好是 128 的倍数（v6e / TPU7x 上是 256），否则硬件会自动补零（padding），等于花同样的时间算一堆零 ^{10 11}。

再看最新两代的走向：

• 第七代 Ironwood（TPU7x，2025）延续 v5p 的 3D torus + OCS 路线，单一 superpod 直上 9,216 颗芯片、FP8 峰值算力 42.5 ExaFLOPS，更大规模则交给数据中心网络（DCN）串接多个 superpod ^{8 12 13}。
• 2026 年 Cloud Next 发布的第八代则首次拆分为两颗专用芯片：用于训练的 TPU 8t 与用于低延迟推理的 TPU 8i。其中，8i 针对 AI 智能体工作负载引入新的“Boardfly”拓扑，通过直连光路把芯片间通信的最大跳数从 16 跳降到 7 跳 ^{9 14}。

训练与推理的通信模式不同，最佳拓扑也开始分道扬镳。拓扑不是教条，而是跟随工作负载演进的成本工程，这正是本文主线最好的注脚。

Slice 是什么？芯片的“座位表”，不是矩阵的切法

理解了单颗芯片，接下来就是本文最容易搞混的地方。当你在 Google Cloud 租用一个 TPU slice，例如 v5e 的 8x16，这个“8x16”的意思是：128 颗芯片在机房里排成 8 行 16 列的二维网格，每颗芯片只和上下左右的邻居用高速通道（ICI，芯片间高速互连）直连 ⁵。

（实际工作负载示例可参考 GKE AI 系列：在 TPU 上用 JAX 跑第一个训练任务 ¹⁵。）

2D Torus 拓扑示意：以 4×4 slice 为例，每个方块是一颗 TPU 芯片；实线是和上下左右邻居的 ICI 直连，虚线是每行每列头尾相接的“环绕”连线（8×16 同理，只是网格更大）

上图用 4x4 缩小示意：每颗芯片只和上下左右的邻居直连，而每一行、每一列的头尾再用环绕连线相接，这就是下一段要讲的“环面”（torus）。

这条芯片间的高速通道叫 ICI（Inter-Chip Interconnect）。3D 拓扑的芯片（v4、v5p）有 6 个 ICI 端口，朝 ±X、±Y、±Z 六个方向各接一个邻居；2D 拓扑的芯片（v5e、v6e）则有 4 个，接上下左右 ^{4 11}。网格的头尾再绕回来相接，就成了“环面”（torus），像把一张棋盘卷成甜甜圈，让最边缘的芯片也不会变成孤岛。

这样设计的首要考量是降低延迟，但也连带几个好处：

• 每颗芯片的“出入口”固定，布线成本可控：每颗芯片只需要几个 ICI 端口（2D 是 4 个、3D 是 6 个），硬件设计可以标准化量产；扩展到更大规模时，不必把每颗芯片都接到昂贵的交换机。例如，NVIDIA 使用 NVSwitch 与专用 NVLink 链路实现 GPU 之间的高速通信。
• 通信距离随维度增加而缩短：在环面上，任意两点之间的最远距离大约等于各维度长度的一半之和。将狭长的 1D 拓扑改为 2D 或 3D 后，最坏情况下的跳数（hop）会显著下降，集合通信（collectives）也更容易获得稳定的带宽和低延迟。
• 集合通信有“天然路径”：All-Reduce、All-Gather 这类操作可以沿环面执行 ring / mesh 算法，让数据像接力一样绕行一圈或多圈即可完成同步，非常适合通信模式规律的深度学习工作负载。

但这样的设计也意味着软件层通常需要管理 TPU 复杂的拓扑网络。在 GPU 胖树（Fat-Tree）或全连接网络中，开发者通常不必过多关注哪些 GPU 彼此相邻。但在 TPU slice 的 torus 上，并行切分必须与芯片的物理邻接关系对齐。一旦没有对齐，大量通信就会被挤入长距离路径，导致任务无法同步，甚至出现死锁（deadlock）；最直接的表现就是任务卡住。

按照规模从小到大，整个层级如下：

flowchart TD
    A["TPU 芯片<br>内含 TensorCore 与 MXU"] --> B["Tray 主板<br>4 颗芯片配 1 台 CPU 主机"]
    B --> C["机柜 Rack<br>4x4x4 Cube = 64 颗芯片<br>铜缆直连成 3D 网格"]
    C --> D["OCS 光学电路交换<br>把多个 Cube 接成任意大小的 Slice"]
    D --> E["Pod<br>v4: 4096 颗 / v5p: 8960 颗 / v7: 9216 颗"]
    E --> F["跨 Pod：数据中心网络 DCN<br>Multislice 训练"]

以 TPU v4 为例：基本单位是一个机柜里的 4x4x4 cube（64 颗芯片），cube 内部全部使用成本较低的铜缆直连；需要组成更大的 slice 时，再由 OCS 通过光纤连接多个 cube。slice 的形状可以采用各维度为 4 的倍数的多种组合，例如 4x4x12、4x8x8，甚至不必是 2 的幂 ^{4 16}。

顺带一提，v4 的命名如 v4-128 指的是 128 个 TensorCore（每颗芯片有 2 个），实际上就是 64 颗芯片、正好一个 cube。

TPU v4 常见的 slice 拓扑形状：从 2×2×1（4 颗芯片）到一整个 4×4×4 cube（64 颗芯片）

所以，回到开头的问题：slice 的“几乘几”是实际矩阵乘法的分布吗？

不是。它比较像是芯片的物理座位表。

逻辑网格 vs 物理拓扑：XLA 如何把矩阵映射到芯片上

训练大模型时，我们会在逻辑层把计算切分成不同的并行维度，最常见的两种：

• 数据并行（Data Parallelism）：不同芯片处理不同批次的数据，每轮结束后同步梯度。
• 模型／张量并行（Tensor Parallelism）：单一权重矩阵太大，切成好几块分给不同芯片，各算一部分。

在 JAX 中，这通常表示为一个逻辑网格（mesh）。例如，可以把 128 颗芯片声明为 Mesh(('data', 'model'), (8, 16))：8 路数据并行 × 16 路张量并行 ¹⁷。

关键在于：逻辑网格和物理拓扑是两张不同的图，由 XLA 编译器负责把前者“映射”到后者上。因为 ICI 只连接相邻芯片，距离较远的芯片必须通过中间芯片逐跳通信 ¹¹。编译器会尽量把“通信最频繁的逻辑维度”映射到“物理上相邻的芯片轴”：

flowchart LR
    subgraph L["逻辑层（数学问题）"]
        A["巨大的矩阵乘法"] --> B["切分成并行维度<br>data = 8, model = 16"]
    end
    subgraph P["物理层（机房接线）"]
        C["8x16 芯片网格<br>ICI 只连上下左右邻居"]
    end
    B -->|XLA 编译器对应| C
    C --> D["每颗芯片内部<br>仍由 128x128 或 256x256 的 MXU 计算"]

以 8x16 的 v5e slice 配合 data = 8, model = 16 的逻辑网格为例：编译器会把 8 路数据并行映射到物理网格的一条轴，16 路张量并行映射到另一条轴。这样一来，张量并行所需的高带宽通信（每层前向和反向传播都要交换部分结果）会发生在物理相邻的芯片之间，而数据并行中频率相对较低的梯度同步（All-Reduce）也能沿环面高效执行。

总结这一段：

• Slice 的尺寸（如 8x16）：芯片网络的物理互连结构。
• 矩阵的切分：数学上的张量分片（sharding）维度。
• 实际运行：XLA 把数学切分映射到物理拓扑上，尽量缩短通信路径；无论外层如何切分，每颗芯片内部最终仍由 MXU 的脉动阵列完成计算。

为什么不排成一条线就好？1D 拓扑的物理极限

看到这里你可能会想：2D、3D 拓扑还要求工程师考虑芯片的“座位表”，未免太复杂，为什么不简化成一条线（1D）？

答案很直接：一维拓扑在物理上无法高效支撑数千颗芯片。设想把 4096 颗芯片排成一条直线：头尾两颗芯片通信时，数据必须穿过中间几千颗芯片；即使首尾相连成环，最远仍要经过 2048 跳。如果改排成 16x16x16 的 3D 环面，任意两颗芯片之间的最远距离会立刻缩短到 24 跳（每个维度最多 8 跳）。维度越高，芯片间的平均通信距离越短。TPU 正是利用这一几何特性，在计算性能与成本之间取得更好的平衡。

那为什么 GPU 集群“感觉”不用管拓扑？因为 GPU 也不是 1D，而是把拓扑藏了起来。

以常见的 NVIDIA DGX 方案为例，服务器内部使用 NVSwitch 连接 8 颗 GPU，让它们在逻辑上如同一颗巨型 GPU；服务器之间则通过 InfiniBand 交换机构建胖树（Fat-Tree）拓扑。例如，DGX SuperPOD 参考架构使用 Quantum-2 InfiniBand 组成三层胖树 ¹⁸。对软件而言，所有 GPU 仿佛都连接在同一个大型交换网络中，逻辑拓扑相对“扁平”。

GPU 集群网络简化示意：每台 DGX 服务器内部由 NVSwitch 连接 8 颗 GPU；服务器之间则由 Leaf 和 Spine 层的 InfiniBand 交换机组成胖树（Fat-Tree）拓扑。实际的 SuperPOD 使用三层胖树。

当然，每种架构都有取舍。这类设计最直接的代价就是成本。胖树拓扑若要维持任意两点之间的无阻塞带宽，所需交换机、光模块和线缆的数量会随规模快速增加。据估计，网络设备往往占 GPU 集群总成本的 15–20%，而且规模越大，这一比例越值得关注 ¹⁹。

TPU 的省钱哲学：ICI + OCS

如果你以为 TPU 完全不用交换机，只要把芯片绑在一起就能扩充，那并不完全正确。更精确的说法是 ⁴：

1. Cube 内部（64 颗芯片）：芯片通过 ICI 和无源铜缆直连，组成 4x4x4 的 3D 网格；这一层确实没有交换机，也就是所谓的 glueless 设计。
2. Cube 之间：靠 OCS（Optical Circuit Switch，光学电路交换）串接；一个 v4 Pod 会用 48 台 OCS 把 64 个 cube（共 4096 颗芯片）连成完整系统。

OCS 与 InfiniBand 这类“电子数据包交换机”有本质区别。普通交换机就像邮件分拣中心：每个经过的数据包都要检查目标地址，才能决定转发方向。这需要交换芯片参与处理，也会增加延迟与功耗。

OCS 则简单得多：其内部是一组可以调整角度的微型镜面阵列（MEMS），工作原理就是把光反射到另一条光纤上，类似老式电话交换台将 A 端线路直接接到 B 端。光路一旦建立，数据穿过 OCS 时就像沿一条直通光纤传输，全程不解析数据包，也不检查内容。这带来三个关键优势 ⁴：

• 便宜省电：OCS 与底层光学元件占整体系统成本不到 5%、功耗不到 3%。
• 拓扑可重组：无须重新布线，只要重新配置光路，就能将 slice 组成不同形状，甚至支持“twisted torus”这类变形拓扑。通过“扭转”环面的连接方式，可减少远距离芯片之间的跳数，更适合 all-to-all 通信；用户需要 4x4x12，就可以配置为 4x4x12。
• 容错与调度：某个 cube 故障时，光路可以绕过它并连接备用 cube。调度器也不必寻找物理位置连续的空闲芯片，只需选择若干可用 cube，即可组合成一个 slice。

把两种设计哲学放在一起看：

所以，TPU 将拓扑直接呈现在开发者面前，并非刻意增加工程难度，而是一项明确的取舍：以软件和编译器的复杂度，换取硬件规模化时的性价比与灵活性。

这项取舍之所以成立，是因为深度学习训练的通信模式高度规律：每一轮计算中，芯片通常都与固定的相邻芯片交换数据，例如同步梯度，像工厂流水线一样重复相同动作。环面拓扑最擅长的正是这种邻居间通信。反过来说，TPU 的设计也因此高度针对深度学习训练进行了优化；如果用于通用计算，通信双方和通信时机都难以预测，数据往往需要在网格中绕行较远的路径，这套设计就未必经济。

结语

回到开头的问题：TPU slice 为什么是 4x4x4？一路分析下来会发现，答案不在数学课本里，而在机房的布线方式与建设成本。本文的重点可以归纳为以下几点：

• “几乘几”是芯片的座位表，不是矩阵的切法。4x4x4 说的是 64 颗芯片在机柜里如何排列、如何连接；矩阵如何切分、每颗芯片负责哪一部分，则由 XLA 编译器把逻辑网格映射到这张座位表上。
• TPU 节省的是交换基础设施成本。cube 内部用铜缆直连，cube 之间由 OCS 切换光路。以 v4 为例，OCS 加上光模块、光纤等整套光学组件，只占系统成本不到 5%。据估计，GPU 集群会将 15–20% 的预算投入 InfiniBand 胖树等网络设备（NVSwitch 则计入服务器本身的成本），以硬件投入隐藏拓扑复杂度，降低软件开发难度。
• 复杂度不会消失，只会转移。Google 节省的硬件成本，转化为编译器与工程师需要处理的问题：矩阵维度应对齐 128（或 256），并行切分也应贴合物理拓扑。否则，轻则 MXU 利用率下降，重则任务直接卡死。

希望这篇文章能帮助你理解 TPU 架构的设计取舍，并在评估 AI 基础设施或阅读规格表时，做出更务实的判断。

参考资源