谷歌 TPU v4 与光路交换：AISystem 视角下的 AI 芯片架构与超级互联解析

本文是 AISystem 仓库《国外 AI 芯片》系列课程的 TPU 正代芯片终章，承接 04TPUIntrol.md 的 TPU 历史演进与 07TPU3.md 的 TPU v3 Pod 形态。全文以 TPU v4 芯片架构、Sparse Core 稀疏计算核、3D Torus 光互联与 Palomar 光路交换机为主线，完整覆盖芯片微架构、产品规格、集群拓扑与光交换原理，并补充与英伟达 DGX SuperPod 的集群对比与优缺点权衡。读完本文，你将掌握 TPU v4 从单芯片到 4096 节点超级计算机的完整技术栈，理解稀疏计算硬件加速、光路交换与 3D Torus 拓扑在大规模 AI 训练集群中的设计逻辑。

为什么 TPU v4 值得单独一章

TPU v3（2018 年）到 TPU v4 之间相隔了四年。这四年里谷歌并没有停下脚步——它把更多精力放在了面向大众的芯片上，例如 Pixel 手机系列里的 Pixel Visual Core、Pixel Neural Core 与谷歌 Tensor，但 TPU 正代芯片一直没有更新。

而这四年间，业界环境发生了三个关键变化：

• 英伟达迭代了三代架构：从 Volta（V100）到 Ampere（A100）再到 Hopper（H100），GPU 在 AI 训练市场的统治力持续增强；
• 训练框架易主：PyTorch 取代 TensorFlow 成为业界的首选训练框架；
• 超大算力需求爆发：大模型的出现对单卡算力之外的大规模并行计算能力提出了前所未有的要求。

TPU v4 正是谷歌针对上述问题给出的答卷。它围绕一个核心命题展开设计：大模型的训练瓶颈已经从单芯片算力转移到集群互联与规模化能力。这也是 TPU v4 首次在 Pod 层面引入光交换（OCS）与 3D Torus 拓扑的根本原因。

TPU v4 芯片架构总览

每个 TPU v4 包含两个 Tensor Core，每个 Tensor Core 内部由六个单元组成：

• 四个 MXU：TPU 最核心的脉动阵列（Systolic Array）矩阵乘法单元，延续了自 TPU v1 以来的脉动阵列设计（详见 05TPU1.md）；
• 一个 Scalar Unit：负责标量计算；
• 一个 Vector Unit：负责向量计算。

此外，芯片上还有两个 HBM 内存模块，分别放置在两个 Tensor Core 的左右两侧。这种布局并非随意为之——将 HBM 紧贴计算核心两侧，是为了尽量缩短数据在片内传输的物理路径，从而降低电缆与走线的时延，这在超大算力集群中直接影响整体性能。

产品形态对比

TPU v4 相比前代在产品形态上变化巨大，主要体现在四个方面：

• 制造工艺：从 TPU v3 的 16 nm 提升到 7 nm 工艺，是谷歌在 TPU 制程工艺上最大的一次更新；
• 硬件规模：MXU 数量相比 TPU v3 再翻一倍；片上缓存增加 9 倍，达到 244 MB（不同资料来源存在 244 MB 与 288 MB 两种表述，本文以本节文档为准）；HBM 容量保持 32 GB 不变，但内存带宽提升到 1.2 TB/s，相比前代提升了 33%；
• Sparse Core：硬件层面新增稀疏计算核，专门利好稀疏计算场景，并基于 TPU v4 改良了自有 Transformer 模型架构；
• 3D Torus 互联：首次亮相 3D Torus 互联方式，紧密耦合 4096 个 TPU v4 引擎，使 TPU v4 Pod 总计提供 1.126 Exaflops 的 BF16 峰值算力。

其中 Sparse Core 与 3D Torus 是 TPU v4 区别于前代的两大核心技术，下面分别展开。

Sparse Core：为稀疏计算而生的专用硬件

Embedding 层的计算本质

先看一个计算范式问题。Embedding 处理的是离散型分类特征（Categorical Features），这是稀疏化的典型计算范式。NLP、搜索推荐算法只支持字符串形式输入，这些输入被标示为离散的稀疏向量特征。

这些特征本质上更像一张哈希表，而不是适合映射到硬件矩阵乘法单元进行张量计算的稠密数据。原因在于：

1. 输入矩阵中存在极大量的 0，这些 0 本身不需要参与计算，或者可以大幅简化计算；
2. 若按照传统矩阵乘法逐项相乘，会反复执行大量"0 × 权重"的无意义乘法，白白浪费计算资源；
3. 将稀疏变量放到 Tensor Core 上计算，往往涉及小规模的内存访问动作，很容易触及 CPU 与 DRAM 的性能瓶颈——尤其是在 TPU 与 CPU 比例为 4:1 的 TPU 集群中，网络中心的延迟还会进一步放大这种负面影响。

因此，TPU 需要一种能够高效处理稀疏变量的计算范式。深度学习神经网络通常在稠密 Tensor 上计算性能更优，所以需要先对稀疏矩阵进行信息压缩，通过特定的 Embedding 层生成稠密的原始数据表达——这通常就是 NLP 或搜推算法的第一层。

上图是搜索推荐场景最经典的模型框架 Deep & Wide Model。Wide & Deep 本身是一个框架：Wide 部分可以是任意广义线性模型，而 Deep 部分则是 Embedding 加全连接神经网络。图中 "Deep Models" 部分清晰展示了用户稀疏特征先被处理成稠密 Embeddings，再进入中间的隐藏层，协同 "Wide Models" 输出最终单元。

TPU v4 的 Sparse Core 在硬件中原生支持 Embedding 计算的模型并行与数据并行，提供了极大的并行灵活度，让大规模 Embedding 计算能够在超大规模集群内被妥善处理。当 NLP 或推荐的语料规模非常庞大时，就需要在 Embedding 层进行切分、并行，TPU v4 直接把这些特性固化到了硬件里。

Sparse Core 核心架构

Sparse Core（SC）中最通用的单元是 16 个 tiles（计算瓦片），每个瓦片都关联一个 HBM 通道，并支持多个内存访问。每个瓦片内部包含三个关键部件：

部件	职责
Fetch Unit	从 HBM 读取激活函数和参数，加载到瓦片内的 2.5 MiB 稀疏向量内存（Spmem） 切片中
scVPU	可编程的 8 路 SIMD 向量处理单元（注意与 Tensor Core 中的 Vector Unit / VPU 区分），是稀疏数据快速计算的核心
Flush Unit	在反向传播过程中，将更新后的参数写回 HBM

除了 16 个瓦片（图中的深蓝色方框），还有五个跨通道单元（图中的金色方框），在 16 个 Spmem 上执行与其名称相对应的嵌入操作（例如稀疏归约、排序、数据搬运等，具体见 srt/08.srt 中对应讲解）。与 TPU v1 一样，这些单元执行类似 CISC 的指令，处理可变长度的输入——每条指令的运行时间取决于实际数据。这正是稀疏计算与传统稠密矩阵计算在指令模型上的根本差异：数据长度不固定，指令执行时间随数据动态变化。

TPU v4 Pod：1.126 Exaflops 的超级计算机

TPU v4 Pod 能提供 Exaflops 级别的算力，核心原因在于其强大的芯片间互联能力。谷歌的构建方式分两步：

1. 将 4 × 4 × 4 = 64 个 TPU v4 芯片互联在一起，形成一个立方体结构（Cube）；
2. 再把多个 4 × 4 × 4 Cube 用光互联连在一起，形成一个总共包含 4096 个 TPU v4 芯片的超级计算机。

光互联（OCS）

在如此规模的超级计算机中，芯片间互联在很大程度上决定整体计算效率。如果数据互联效率不够高，很多时候芯片都在等待来自其他芯片的数据到达才开始计算——计算单元空转，整体利用率上不去。因此必须确保芯片之间高带宽、低延迟的互联，而光互连对于物理距离较远的芯片就成为首选。

谷歌在 TPU v4 Pod 中应用了光交换器（OCS，Optical Circuit Switching） 来避免计算与通信的相互等待：

• 64（4 × 4 × 4）颗 TPU 构成一组，负责 Slice 间的互连，实现 Pod 内 Slice 间全光互连（4096 TPUs）；
• OCS 同样可用于 Pod 之间的互连，而不仅局限于芯片之间。

一个值得关注的数据点：在使用可配置光互连（以及光路开关）时，假设芯片可靠率为 99%，其整体系统的平均性能提升比不使用 OCS 可高达 6 倍。光互连开关的重要性由此可见一斑。

3D Torus 拓扑结构

TPU v4 Pod 采用 3D Torus 拓扑结构实现芯片间互联。在 3D Torus 网络中，每个节点通过三个维度的连接与相邻节点相连，形成一个三维网格结构：

• 节点在 X-Y-Z 三个维度上形成一个连续循环（环面），因此最边上的节点会连接到相对面上的节点；
• 每个节点与上下、左右、前后共 6 个节点互联，形成高度互联的闭合结构；
• 这种结构提供高带宽和低延迟的通信能力，非常适合高性能计算和大规模并行处理。

TPU v4 通过该拓扑紧密耦合 4096 个 TPU v4 引擎，实现总计 1.126 Exaflops 的 BF16 计算能力。每个 TPU v4 端口对应一个 TPU v4 芯片，端口连接交换机提供 6 Tb/sec 的带宽，作为网络接口卡和 3D Torus 网络的基础。

Cube 与 OCS 的链路计算

进一步拆解这个 Cube：为了实现 6 面连接，每个面需要 16 条链路，因此每个 Cube 总共有 96 条光链路连接到 OCS 上。为了让整个 3D Torus 连接起来，相对的两个侧面需要连接到同一个 OCS 上（这正是图中不同色块与 Cube 之间长环的含义）。于是：

$$每个 Cube 连接的 OCS 数量 = 6 \times 16 \div 2 = 48$$

这 48 个 OCS 一共连接来自 64 个 Cube 的 48 对光缆，总计并联 4096 个 TPU v4 芯片。

由此可以推算搭建一个 TPU v4 集群的基本清单：4096 片 TPU v4 芯片 + 48 个 OCS 光互联交换器，成本相当之高。值得注意的是，在"每一家公司和投资人都愿意无上限为模型花钱"的时代，传统大厂在成本投入上的优势相对于新兴 AI 公司（例如 OpenAI 与 Anthropic）正在缩小，芯片产能与电力资源这类物理限制，反而成了各家迭代模型的最大瓶颈。

与英伟达 DGX SuperPod 的集群对比

当时的竞品状态如何？英伟达 DGX SuperPod 搭配第四代 NVLink/NVSwitch：

• 最多连接 32 个 node、总计 256 颗 H100 芯片；
• 每颗 GPU 实现 900 G/s 的互连带宽；
• NV 每机架 4 台 DGX（共 32 颗 H100 GPU），机架内/外需要光纤连接；
• NV 的每机架算力密度相对更小/更窄，需要更多的收发激光器和光纤线材，网络成本高；
• 若要部署 4096 颗 GPU 集群，必须切分成更多个 SuperPod 并独立规划互连网络层，中间完成多层交换，集群内总计需要采购大约 568 个 InfiniBand Switch。

反观 TPU v4 Pod：与超级计算机一样，工作负载由不同规模的算力承担，称为切片（Slice），例如 64 芯片、128 芯片、256 芯片等。与 InfiniBand 相比：

• OCS 的成本更低、功耗更低、速度更快：成本不到系统成本的 5%，功率不到系统功率的 3%；
• 每个 TPU v4 都包含 SparseCores 数据流处理器，可将依赖嵌入的模型加速 5 至 7 倍，但仅使用 5% 的裸片面积和功耗。

Palomar：基于 MEMS 的光路交换机

光路开关芯片

上文提到的光路开关芯片名为 Palomar，采用基于 MEMS 反射镜阵列的技术。其基本原理是：使用一个 2D MEMS 反射镜阵列，通过控制反射镜的位置来调整光路，从而实现光路切换。相比传统光电转换交换，MEMS 光路开关芯片具备低损耗、低切换延迟（毫秒级别）、低功耗、低成本四大优势。

2D MEMS 阵列

每个陶瓷封装内部是一个大型芯粒，芯粒内有 176 个可单独控制的微反射镜。从反射镜阵列的热成像图可以清楚看到：每个微反射镜四周都有四个梳状驱动区域，用于在两个方向上旋转反射镜，使光可以进行光路切换。

Palomar 工作原理

Palomar OCS 的工作原理可以从信号路径来理解：

• 绿色线条表示带内光信号路径；
• 与带内信号路径叠加的是一个 850 纳米波长的监控通道（红色箭头），用于校准镜面；
• 输入/输出的光学信号通过二维（2D）光纤准直器阵列（Fiber collimator array） 进入光学核心，每个准直器阵列由一个 NxN 光纤阵列和 2D 透镜阵列组成；
• 核心是其中的两组 2D MEMS 阵列：MEMS 的镜面在监控通道的校准下被驱动、倾斜，将信号切换到相应的输入/输出准直光纤。

为什么要用 OCS 而不是传统交换机？传统交换机传进来的是光缆，需要先把光转化成电、再把电转成光输出，这个光-电-光转换过程会造成大量性能损耗。而 OCS 提供的是光到光的直接转换，从而节省了非常多的电能，更重要的是解决了大量时延问题——这正是大规模 AI 集群中互联延迟敏感场景所迫切需要的。

TPU v4 优缺点分析

TPU v4 Pod 在成本、功耗和速度方面展示了明显优势。但是，尽管 3D Torus 拓扑带来了低延迟和低网络成本的益处，也存在一些挑战，包括系统成熟度、拓扑僵硬性以及负载均衡问题。下面从正反两面完整梳理。

优点

• 低时延：3D Torus 的相邻节点之间是短而直接的连线，因此延迟更低；尤其当节点间需要运行密集 I/O、紧耦合的并行任务时特别有用；
• 低网络代价：对于相同数量的节点，3D Torus 拓扑的网络直径低于 Clos 拓扑，前者的交换机/线材/连接器保有量更低，网络层次更少，节省硬件成本；
• 路由可重配：谷歌 OCS 网络支持动态可重配路由，Slice 集群在部署之后可以立即投入生产，无需等待整个网络收敛；并且这种特性更容易隔离/下线故障节点；
• 更好的集群布局：集群布局让物理上相邻的节点在逻辑上也临近，让密集 I/O 通信、data-flow 发生在局部流域，换来更低的通信开销，同时优化延迟和功耗。这就是 3D Torus 将大集群逻辑切割成紧耦合的局部域，局部互连并共享作业的设计思想。

缺点

• 系统成熟度低：Clos 拓扑本身具备非阻塞特点，性能始终保持一致且可以预测，所有输入/输出都是全带宽同时连接，无冲突无阻塞——这在 3D Torus 拓扑中无法保证；
• 拓扑僵硬：在 Clos 这种 Spine-Leaf 脊叶拓扑中，扩容新的叶交换机相对简单，无需更改当前架构；相比之下，扩缩 3D Torus 结构比较复杂和耗时，可能需要重新配置整个拓扑；
• 负载均衡问题：Clos 网络在任意两个节点之间提供更多路径，从而实现负载均衡和冗余；虽然 3D Torus 结构也提供多路径冗余，但 Clos 的替代路径数量更多，具体取决于网络的配置。

小结与思考

回顾整章内容，可以提炼出三条主线：

1. 模型的演进总是快于芯片的设计迭代。随着大模型的涌现，训练越来越依赖 AI 计算集群的规模化能力而非单芯片算力，TPU v4 的设计重心正是对高效互联和规模化的支持；
2. 算法-芯片协同设计是 TPU v4 的灵魂。AI 芯片与算法之间的结合继续保持紧密关系，具体体现在对新数据格式（HF32/BF16）的支持、对稀疏计算的原生支持、以及对模型关键算法专用加速器的引入（如 Sparse Core）；
3. Benchmark 之外的很多生产环境性能指标无法直观评测。XLA 编译器的优化效益也难以直观评测，TPU v4 是否如网上披露的信息那般强悍，仍有待实际生产环境的验证。

作为历史上最有远见的科技公司之一，谷歌能否在大模型时代继续保持曾经的远见，才是决定它能否在这个时代乃至下一个时代引领世界的关键。

本文内容基于 08TPU4.md 整理扩展，配套 PPT 与视频资源位于 02Hardware/05Abroad 目录（含 08TPU4.pdf 与字幕 srt/08.srt）。建议按顺序阅读 04TPUIntrol.md → 05TPU1.md → 06TPU2.md → 07TPU3.md → 本文，以形成 TPU 全系列的知识闭环。