Google TPU 8:训练扩到9600颗,推理为什么要换一套架构?

潘驴邓晓闲缺一
08-26

Google 已经迭代了八代 TPU,第八代第一次明确采用两条硬件路线。TPU 8t 面向大规模预训练,TPU 8i 面向后训练、推理和高并发模型服务。两颗芯片共享 Google 的软件体系和基础设施,但内存配置、专用计算单元和网络拓扑已经明显分开。

参数上的差异非常直接。TPU 8t 单芯片拥有 216GB HBM、6.528TB/s HBM 带宽、128MB 片上 SRAM,FP4 峰值计算能力达到 12.6 PFLOPS;TPU 8i 的 FP4 峰值算力略低,为 10.1 PFLOPS,HBM 却增加到 288GB,带宽达到 8.601TB/s,片上 SRAM 更达到 384MB。两款芯片的双向芯片间互连(ICI)带宽均为 19.2Tb/s。

推理芯片没有追求更高的峰值 FLOPS,却拿到了更多内存、更高内存带宽和三倍于 8t 的 SRAM。这组资源分配已经透露出 TPU 8 的设计方向:训练继续向更大的计算域、更高吞吐和更稳定的长时间运行扩张;推理则越来越受模型驻留、KV Cache、MoE 专家路由以及集合通信延迟约束。

训练和推理正在要求两种不同的 AI 系统。

Google TPU 8t and TPU 8i official product image

Google 第八代 TPU 同时推出训练型 TPU 8t 与推理型 TPU 8i。

一、TPU 8 的分界线:训练追求计算域,推理压缩数据路径

大模型训练天然偏向吞吐。一次训练迭代需要执行大量矩阵运算,同时维护参数、梯度和优化器状态,并在大量加速器之间同步数据。模型规模继续上升以后,张量并行、流水线并行、专家并行等策略会把越来越多通信引入训练过程。系统只要能够保持较高的扩展效率,更大的计算域就仍然能够换取更高的总体训练吞吐。

在线推理面对的是另一组约束。模型权重需要长期驻留在加速器内存。上下文窗口拉长以后,每个会话还会不断积累 KV Cache。对于 MoE 模型,一个 Token 经过路由器后可能被送往其他芯片上的专家模块,再将计算结果汇总回来。自动回归生成的串行特征又意味着,部分等待时间会重复进入每一个 Token 的生成过程。

于是,芯片面积和功耗预算开始出现不同的分配方式。TPU 8t 的 FP4 峰值算力为 12.6 PFLOPS,高于 TPU 8i 的 10.1 PFLOPS;TPU 8i 则把 HBM 容量提高到 288GB,HBM 带宽较 8t 高约 32%,片上 SRAM 达到 384MB。Google 明确将更大的 SRAM 与长上下文解码中的 KV Cache 联系起来,用于降低计算核心等待数据的时间。

TPU 8t and TPU 8i core specification comparison

TPU 8i 的峰值计算能力低于 8t,但 HBM、HBM 带宽和片上 SRAM 均明显增加。

网络设计也随之分叉。TPU 8t 延续 Google 已使用多代的 3D Torus(三维环面)拓扑,专用单元包括 SparseCore 和 LLM Decoder Engine;TPU 8i 换用新的 Boardfly 拓扑,并加入 CAE(Collectives Acceleration Engine,集合通信加速引擎)

这一代 TPU 已经很难再用一套统一的“AI 加速器参数表”来解释。8t 优先解决大规模训练中的计算利用率,8i 则明显增加了用于保存数据和减少通信等待的硬件资源。

二、TPU 8t:把 9,600 颗芯片维持在一个训练域里

TPU 8t 的基本系统单位已经不是一颗芯片,而是一个包含 9,600 颗 TPU 的超级计算单元(Superpod)。Google 公布的单个 TPU 8t Superpod 拥有约 2PB 共享 HBM,FP4 总计算能力达到 121 ExaFLOPS。9,600 颗芯片乘以单颗 12.6 PFLOPS,对应约 120.96 ExaFLOPS,与官方 121 ExaFLOPS 的系统规格基本一致。

单颗 8t 拥有 216GB HBM、6.528TB/s HBM 带宽和 128MB SRAM,并首次在训练型 TPU 中原生加入 FP4。Google 的架构思路是降低矩阵计算过程中的数据搬运压力,同时提高矩阵乘法单元(MXU)吞吐;向量处理单元(VPU)与 MXU 之间的执行重叠也被强化,用于减少量化、Softmax、LayerNorm 等向量运算造成的等待。

TPU 8t ASIC block diagram

TPU 8t 芯片框图:Tensor Core、HBM 与高速 ICI 构成主要数据通路。

芯片计算吞吐只是第一层。TPU 8t 的双向 ICI 带宽达到 19.2Tb/s,是 Ironwood 的两倍。多个 TPU 计算单元继续通过 OCS(Optical Circuit Switch,光路交换机)连接,系统可以动态重新配置光路,在不同训练任务之间组织计算资源,同时绕过部分失效的 ICI 光链路。

数千颗加速器进入同一个训练任务以后,设备故障已经成为统计上的常态。芯片、光模块、链路、交换设备任何一个环节出现异常,都可能拖慢同步训练。Google 因此把 有效计算占比(Goodput)超过 97%列为 TPU 8t 的系统目标,并加入跨数万颗芯片的实时遥测、ICI 故障链路自动检测与绕行以及 OCS 自动重构。

TPU 8t Superpod and OCS compute domain illustration

TPU 8t 通过 OCS 动态重构光路,用于计算域组织、任务切片与故障绕行。示意图;资料依据:Google Cloud、Hot Chips 2026

Hot Chips 2026 进一步披露了运行期间的硬件检测机制,芯片可以利用部分空闲周期检查内部单元状态。大型 AI 系统的工程前提已经发生变化:系统架构必须允许局部硬件持续发生故障,同时尽量避免整个训练任务停下来等待维修。

9,600 颗芯片仍然只是一个 Superpod 的边界。继续向外扩展时,网络进入 Virgo 这一层。TPU 8t 每颗芯片的跨系统扩展网络带宽达到 400Gb/s,Ironwood 为 100Gb/s。Virgo 采用高基数交换机、扁平两层非阻塞架构和多平面设计,一个 Virgo 网络可以连接约 13.4 万颗 TPU 8t,提供最高 47Pbit/s 非阻塞二分带宽,对应超过 1,700 ExaFLOPS 的计算规模。

TPU 8t rack level connectivity to Virgo fabric

TPU 8t 从 Superpod 内部的 ICI / OCS 继续通过 Virgo 向数据中心级计算集群扩展。

这里存在三个不同的系统尺度:9,600 颗 TPU对应单个 TPU 8t Superpod;约 13.4 万颗 TPU对应单个 Virgo 网络能够连接的规模;JAX 和 Pathways 进一步把多个计算域组织起来后,Google 宣称单个逻辑训练集群可以扩展到 100 万颗以上 TPU

TPU 8t 这一代的主要增量因此不仅存在于芯片内部。ICI 带宽翻倍,跨系统网络带宽提高到上一代四倍,OCS 负责计算域重构,Virgo 继续扩大数据中心级网络规模。训练系统的竞争焦点已经延伸到如何让成千上万颗计算芯片持续保持高利用率。

三、TPU 8i:推理芯片为什么反而配置更多内存

TPU 8i 的设计取舍可以从三组数据直接看出:HBM 容量从 8t 的 216GB 增加到 288GB;HBM 带宽从 6.528TB/s 提高到 8.601TB/s;片上 SRAM 从 128MB 提高至 384MB。与此同时,FP4 峰值算力却从 12.6 PFLOPS 下降到 10.1 PFLOPS。

TPU 8i ASIC block diagram

TPU 8i 芯片框图:两个 Tensor Core 与位于小芯片裸片侧的 CAE 共同构成推理侧专用架构。

大型模型在线推理正在越来越明显地撞上“内存墙”。第一层压力来自模型权重。参数规模越大,单颗芯片能够容纳的模型比例越低,模型切分和跨芯片通信随之增加。第二层压力来自 KV Cache。上下文窗口拉长、并发会话增加后,每个请求都在持续占用内存。第三层来自自动回归解码。模型每生成一个 Token,都需要再次读取权重和上下文状态。理论算力再高,如果数据无法及时进入计算核心,MXU 依然会出现空闲。

Google 因此把 TPU 8i 的片上 SRAM 增加到 384MB。官方称,这是其 TPU 中最高的片上 SRAM 配置,并明确指向更大的片上 KV Cache 以及长上下文解码中的核心利用率。

TPU 8t versus TPU 8i memory resources

TPU 8i 将更多硅资源交给模型驻留、KV Cache 与数据供给。

MoE 又在内存问题之上叠加了一层网络问题。MoE 只在每次前向计算中激活部分专家模块,可以扩大总参数规模而不同比例增加计算量,但不同专家往往分布在不同加速器。Token 完成路由之后,需要跨芯片发送到目标专家,随后重新聚合。

这使推理性能越来越依赖两项指标:一颗芯片能够容纳多少模型和缓存,以及一个 Token 需要经过多长的数据路径。TPU 8i 提高 HBM 和 SRAM 解决第一项,Boardfly 解决第二项。

四、Boardfly:Google 为什么把 16 跳压到 7 跳

3D Torus 长期被 Google 用于 TPU 训练系统。这种拓扑让芯片通过多个维度与相邻节点连接,链路数量可控,特别适合训练中大量规律性的邻近通信和集合操作。但当通信模式逐渐转向 All-to-All(全互联通信),网络直径开始成为延迟来源。

MoE 正是典型场景。一个 Token 可能需要被发送到系统中距离很远的专家模块。随着计算域扩大,如果网络仍然依赖多跳转发,数据在不同芯片之间移动的时间就会持续增加。

Boardfly 采用分层、高基数拓扑。最底层每个计算托盘由 4 颗 TPU 组成一个基本单元,并提供 16 条外部 ICI 连接;8 块板通过铜缆组成一个组,其中 11 条外部连接用于组内通信;最终 36 个组通过 OCS 连接成 TPU 8i Pod。

TPU 8i hierarchical Boardfly topology

Boardfly 从 4 颗 TPU 基本单元逐级扩展到 8 块板组成的组,并最终连接 36 个组。

Google 目前公开资料同时存在两个系统规模数字。产品材料称 TPU 8i 单 Pod 最多连接 1,152 颗 TPU;Boardfly 架构说明中则写明 36 个组对应 最多 1,024 颗活跃芯片(active chips)。Google 没有进一步公开两者之间的资源映射,因此不应自行把 128 颗差额解释为备用芯片。

Google 使用 1,024 颗芯片进行了直接对比。如果采用 8×8×16 的 3D Torus,最远两个节点之间的距离为 16 hops;TPU 8i 采用 Boardfly 后,同样规模下的最大网络直径降至 7 hops,下降约 56%。Google 测试显示,对于通信密集型工作负载,Boardfly 能够带来最高约 50% 的延迟改善。

Boardfly maximum seven-hop ICI network diameter

同为 1,024 颗芯片规模,Boardfly 将最大网络直径从 16 跳降至 7 跳。来源:Google Cloud

Boardfly 的价值主要落在低延迟通信。在线推理尤其是 MoE 和复杂推理型模型中,一个请求可能反复触发专家路由、Reduce 和其他集合操作。单次网络传输节省的延迟并不大,但这些等待会反复进入 Token 生成链路,并进一步影响尾延迟。8i 重新设计的已经不只是“网络带宽”,而是整个推理计算域内部的数据距离。

五、CAE:集合通信开始占用专门的芯片面积

Boardfly 降低芯片之间的网络距离,CAE 处理芯片内部的集合操作。TPU 8i 每颗芯片包含两个位于核心计算裸片上的 Tensor Core,以及一个位于小芯片裸片上的 CAE。Ironwood 此前在核心计算裸片上配置的四个 SparseCore,在 8i 中被 CAE 取代。

CAE 用于加速自动回归解码过程中频繁出现的 Reduce、Synchronization 等集合通信操作。Google 给出的数据是,专用 CAE 可以将 片上集合通信延迟最高降低 5 倍

这里反映出 AI 加速器内部资源分配的另一项变化。早期专用 AI 芯片的核心竞争主要围绕矩阵计算单元、数值格式和 HBM;计算域扩大以后,通信已经占据足够高的系统成本,开始获得专门的硅资源。

Boardfly 减少数据跨芯片需要走多少跳,CAE 缩短集合操作在芯片内部需要等待多久。两项设计都围绕同一个目标——降低推理过程中由通信产生的空闲时间。

NVIDIA 在 GPU 系统中通过 NVLink 与 NVSwitch 提高机架级互连能力,Google 则形成 ICI、OCS、Boardfly、CAE 和 Virgo 等不同层级的自研体系。路径并不相同,但整个行业正在面对相同约束:单颗加速器的计算能力上升以后,系统性能越来越取决于数据能否及时抵达计算单元。

六、从 Axion 到 TPUDirect,TPU 已经变成一套完整机器

TPU 8 还有一项容易被芯片参数掩盖的变化:Google 自己的 Arm CPU 开始更深地进入 TPU 系统。TPU 8t 和 TPU 8i 都使用 Axion CPU 主机。Google 把这项变化直接与数据预处理和系统编排联系起来,希望减少 CPU 侧准备数据造成的 TPU 等待。TPU 8i 进一步把每台服务器的物理 CPU 主机数量提高一倍,并通过 NUMA 进行隔离。

STH 根据 Hot Chips 演讲给出的 TPU 8i 系统配置约为 2 颗 TPU 对应 1 颗 Axion CPU。这一比例目前可以视为 Google 系统设计信息,但尚未成为 Google Cloud 正式实例规格,因此不进一步推导具体虚拟机配置。

对于在线推理,CPU 承担请求调度、数据预处理、网络 I/O 和运行时编排等任务。TPU 算力继续增加以后,CPU 端准备数据的速度同样需要提高,否则主机侧会成为新的等待来源。

TPU 8t 则进一步减少 CPU 参与数据搬运。TPUDirect RDMA允许 TPU HBM 与网络接口卡直接交换数据,绕过主机 CPU 和系统 DRAM;TPUDirect Storage把这一路径延伸到高速存储。Google 称,配合 Managed Lustre 10T 后,TPU 8t 的存储访问速度可达到 Ironwood 训练系统的约 10 倍

TPUDirect RDMA and Storage data path

TPUDirect 减少主机 CPU 和系统内存参与的数据搬运环节,使网络与存储更直接地访问 TPU 内存。来源:Google Cloud

这些设计和 ICI、Virgo 处理的是同一个问题:数据进入 Tensor Core 之前,需要尽可能减少不产生有效计算的中间环节。散热也被纳入整套机器设计。TPU 8t 和 TPU 8i 都支持 Google 第四代液冷系统,Google 给出的系统级能效提升最高达到 Ironwood 的两倍。

软件承担着另一层抽象。第八代 TPU 继续使用 XLA 屏蔽部分底层硬件差异;JAX、PyTorch 和 Keras 代码可以向新一代 TPU 迁移,Pallas 与 Mosaic 则用于开发更贴近 TPU 硬件的定制内核。Google 的原生 PyTorch TPU 支持目前也已经进入预览阶段。

硬件越专用,编译器、运行时和调度系统需要承担的工作就越多。Boardfly、CAE 和 SparseCore 如果全部暴露给模型开发者手工管理,硬件优势很难转化成实际使用效率。Google 试图解决的问题,是让这些复杂性留在 AI 超级计算系统(AI Hypercomputer)内部。

结语:训练扩大计算域,推理缩短 Token 的数据路径

截至 2026 年 8 月 26 日,Google Cloud 仍然将 TPU 8t 和 TPU 8i 标记为 Coming soon。两款产品已经正式发布并公开主要系统架构,但尚未进入 Google Cloud 全面商用阶段。目前处于 GA(正式可用)状态的最新 TPU 仍是第七代 Ironwood。

TPU 8 已经把未来几代 AI 基础设施的两条技术路线摆得相当清楚。训练侧继续扩大计算域:TPU 8t 从单颗 12.6 PFLOPS FP4 扩展到 9,600 颗芯片、121 ExaFLOPS 和约 2PB HBM,再通过 Virgo 把物理网络扩大到十万颗芯片级别。ICI、OCS、Virgo、TPUDirect 和可靠性设计的作用都指向同一项系统指标:让昂贵的计算单元尽可能长时间保持有效工作。

推理侧则开始压缩数据路径。TPU 8i 的峰值算力低于 8t,却拥有 288GB HBM、8.601TB/s HBM 带宽和 384MB SRAM;Boardfly 把 1,024 颗芯片系统的最大网络直径从 16 跳降至 7 跳,CAE 继续压低芯片内部的集合通信延迟。

训练在扩大“能同时工作的芯片数量”,推理在缩短“一个 Token 需要等待的距离和时间”。

过去评价 AI 加速器,FLOPS、HBM 容量和 HBM 带宽往往已经足够抓住性能的大致轮廓。进入万卡训练和大规模 MoE 推理以后,这些数字仍然重要,却越来越难单独解释系统表现。训练一套前沿模型,需要计算数千颗芯片到底有多少时间处于有效工作状态;服务一个大型 MoE 模型,则要计算 Token 经过多少次路由、多少个网络跳点以及多少轮集合通信。

TPU 8 透露出的变化非常具体:AI 基础设施的性能竞争,正在从“芯片能算多快”,延伸到“数据能多快抵达芯片”。

主要资料:Google Cloud《Inside the eighth-generation TPU: An architecture deep dive》、Google Cloud TPU 产品页、Hot Chips 2026 Google TPU 演讲、ServeTheHome 现场报道。

官方架构图版权归 Google 所有;中文数据图与 OCS 示意图依据公开资料整理。仅供参考学习,不构成投资建议。

免责声明:上述内容仅代表发帖人个人观点,不构成本平台的任何投资建议。

精彩评论

我们需要你的真知灼见来填补这片空白
发表看法