第八代 TPU 架构深度解析

# TPU 8t 与 TPU 8i 技术深度解析 来源：https://cloud.google.com/blog/products/compute/tpu-8t-and-tpu-8i-technical-deep-dive 在 Google，我们的 TPU 设计哲学始终围绕三大支柱：可扩展性、可靠性、效率。随着 AI 模型从稠密大语言模型（LLM）走向超大规模混合专家（MoE）与重推理架构，硬件不再只是堆加每秒浮点运算（FLOPS），而必须针对最新工作负载的特定运算强度持续演进。 代理式 AI 的兴起，需要基础设施能够承载长上下文窗口与复杂顺序逻辑；同时，世界模型（world models）从“预测下一 token”走向“模拟未来场景、预判后果、靠‘想象力’而非高风险试错学习”，这要求新一代硬件提供完全不同的支持。第八代 TPU（TPU 8t 与 TPU 8i）正是为此而生——确保从训练的第一个 token 到多轮推理链的最后一步，都能跑在最高效路径上。它们专为高效训练并服务 Google DeepMind 的 Genie 3 等世界模型而打造，让数百万智能体可在多样仿真环境中并行练习、迭代推理。 ### **TPU 8：为专用而生** 我们意识到预训练、后训练与实时推理的基础设施需求早已分化，因此第八代 TPU 推出两套独立系统：TPU 8t 与 TPU 8i。二者均是 Google Cloud AI Hypercomputer 超级计算架构的核心组件，该架构将硬件、软件与网络整合，贯通 AI 全生命周期。虽然共享 Google AI 栈的核心 DNA 并支持全周期工作负载，但两者分别瞄准不同瓶颈，为关键阶段极致提效。此外，第八代全系引入基于 Arm 的 Axion CPU 头节点，消除数据准备延迟导致的主机瓶颈；Axion 提供充足算力完成复杂预处理与编排，让 TPU 持续“吃饱”不空转。 ### TPU 8t：预训练巨无霸 专为超大规模预训练与重度嵌入场景优化，TPU 8t 在单 superpod 内将 proven 的 3D 环面网络拓扑扩展到 9,600 芯片，跨数百 superpod 提供最大吞吐，确保训练按时完成。 相比前代的关键提升： - **SparseCore 优势**：TPU 8t 核心在于 SparseCore——专用加速器处理嵌入查表的不规则访存。矩阵乘单元（MXU）负责矩阵运算，SparseCore 则卸载数据相关的 all-gather 等集合通信，避免通用芯片常见的“零操作”瓶颈。 - **VPU/MXU 重叠与均衡扩展**：通过更均衡的向量单元（VPU）规模设计，最大化已配 FLOPS 利用率，把量化、softmax、layernorm 等与 MXU 矩阵乘更好重叠，减少向量任务空泡。 - **原生 FP4**：引入原生 4 位浮点（FP4），在保持大模型精度的同时把 MXU 吞吐翻倍，降低能耗巨大的数据搬移，让更大层模型常驻本地缓存，实现峰值算力利用。 https://storage.googleapis.com/gweb-cloudblog-publish/images/1_TPU_8t_ASIC_block_diagram.max-2000x2000.png 图 1：TPU 8t ASIC 框图 - **Virgo 网络拓扑 + 4× 数据中心网络带宽**：面对 TPU 8t 的庞大数据需求，我们推出 Virgo 网络（https://cloud.google.com/blog/products/networking/introducing-virgo-megascale-data-center-fabric）。新网络架构让 TPU 8t 训练侧数据中心网络（DCN）带宽提升最高 4 倍。Virgo 采用高基数交换机，减少网络层级，构建扁平两层无阻塞拓扑，显著降低跳数与延迟；多平面独立控制域连接 TPU 8t 芯片，同时通过 Jupiter 南北向骨干访问计算与存储，提供超大对分带宽与确定性低延迟，支撑全球最大训练集群高可用运行。 芯片内互连（ICI）scale-up 带宽翻倍，DCN scale-out 裸带宽提升 4 倍，大幅缓解数据瓶颈。再借助 JAX（https://docs.jax.dev/en/latest/index.html）与 Pathways（https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro），**我们已可在单一训练集群内扩展至百万级 TPU 芯片**。Virgo 单 fabric 可互联超 134,000 颗 TPU 8t 芯片，提供高达 47 Pbit/s 的无阻塞对分带宽，实现 160 万 ExaFlops 近线性扩展性能。 https://storage.googleapis.com/gweb-cloudblog-publish/images/2_TPU_8t_rack_level_connectivity_to_Virgo_.max-2000x2000.png 图 2：TPU 8t 机架级与 Virgo fabric 的连接 - **更快存储访问**：TPU 8t 引入 **TPUDirect RDMA** 与 **TPUDirect Storage**。TPUDirect RDMA 让 TPU 内存（HBM）与网卡直接互传数据，绕过主机 CPU 与 DRAM，降低延迟、提升 TPU 间有效带宽；TPUDirect Storage 同样绕过 CPU，让 TPU 直接访问 10T Lustre 等高速托管存储，大文件传输有效带宽翻倍，确保 MXU 在处理百 PB 级多模态数据集时始终满载。 托管 Lustre 10T（https://cloud.google.com/blog/topics/hpc/google-cloud-ranks-on-io500-benchmark-with-lustre）与 TPUDirect Storage 把百 PB 数据集直送硅片，消除数据摄取瓶颈，相比第七代 Ironwood TPU 训练存储访问提速 10 倍。 https://storage.googleapis.com/gweb-cloudblog-publish/images/3_rq0yjyX.max-2000x2000.png 图 3：上图无 TPUDirect Storage，下图展示 TPU 8t 通过 TPUDirect Storage 与托管 10T Lustre 直连的数据路径 ### TPU 8i：采样与推理专家 专为后训练与高并发推理优化，TPU 8i 拥有最大片上 SRAM、全新集合通信加速引擎（CAE）以及面向推理的网络拓扑 Boardfly。 - **大容量片上 SRAM**：片上 SRAM 提升 3 倍，可将更大 KV Cache 完全放在片上，显著降低长上下文解码时的核心空闲时间。 https://storage.googleapis.com/gweb-cloudblog-publish/images/4_TPU_8i_ASIC_block_diagram.max-2000x2000.png 图 4：TPU 8i ASIC 框图 - **集合通信加速引擎（CAE）**：针对采样瓶颈，TPU 8i 引入 CAE，以近零延迟跨核聚合结果，专门加速自回归解码与“思维链”中的规约同步。每颗 TPU 8i 芯片含 2 个 Tensor Core（TC）与 1 个 CAE（位于 chiplet），替换了前代 Ironwood TPU 核心上的 4 个 SparseCore。专用 CAE 将片上集合延迟再降 5 倍，更少等待、更高吞吐，支撑百万级智能体并发。 - **Boardfly ICI 拓扑**：3D 环面可连数千芯片，但大网格带来多跳与更高 all-to-all 延迟。8i 改用“板级全互连→组级全互连”的高基数设计，最多 1,152 芯片直连，将网络直径与跨系统跳数减半，all-to-all 通信延迟降低 50%，直击 MoE 与推理模型通信痛点。 https://storage.googleapis.com/gweb-cloudblog-publish/images/5_I1mUzjb.max-1300x1300.png 图 5：TPU 8i 层级 Boardfly 拓扑——四芯片全连构成 Building Block，八板全连成 Group，36 Group 经 OCS 全连构成单 Pod Boardfly 层级结构： - **Building Block（BB）**：每托盘四芯片环网，内部 ICI 互联，对外提供 16 条链路。 - **Group（G）**：八块板通过铜缆全互连，组成局部 Group，用 11 条外链做组内通信。 - **Pod**：36 Group（最多 1,024 活跃芯片）通过光路交换机（OCS）互联，任意芯片间最多七跳。 ### 深度对比：Boardfly vs 环面数学 为何弃用环面？关键在于**网络直径**。 3D 环面 8×8×16（1,024 芯片）到最远芯片需跨越各环一半： 3D 环面 = 8/2(X) + 8/2(Y) + 16/2(Z) = 16 跳 环面对邻居通信极佳，却给 all-to-all 模式带来“延迟税”。推理与 MoE 时代，任意芯片都可能需即时交换 token，跳数决定成败。 Boardfly 高基数拓扑借鉴 Dragonfly（https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/34926.pdf）思想，通过增加组间长距光链路“压扁”网络，同样 1,024 芯片 Pod 直径从 16 跳降至 7 跳，降幅 56%，尾部延迟更低，CAE 不再空等。 https://storage.googleapis.com/gweb-cloudblog-publish/images/6_Qu7H2lI.max-1300x1300.png 图 6：TPU 8i Pod 经 OCS 实现任意芯片间最大七跳 ICI 网络直径示意 ### TPU 8t 与 TPU 8i 一览 | 特性 | TPU 8t | TPU 8i | |---|---|---| | 主要场景 | 大规模预训练 | 采样、推理与思维链 | | 网络拓扑 | 3D 环面 | Boardfly | | 专用单元 | SparseCore（嵌入）+ LLM 解码引擎 | CAE（集合通信加速引擎） | | HBM 容量 | 216 GB | 288 GB | | 片上 SRAM | 128 MB | 384 MB | | 峰值 FP4 | 12.6 PFLOPs | 10.1 PFLOPs | | HBM 带宽 | 6,528 GB/s | 8,601 GB/s（≈1.3×） | | CPU 头节点 | Arm Axion | Arm Axion | ### 软件加持：性能优先的 AI 栈 硬件再强，也需软件驱动。第八代 TPU 沿用第七代 Ironwood 开创的“性能优先”栈，让自定义核开发触手可及，又不失高级框架抽象： - **Pallas 与 Mosaic**：一等公民支持 Pallas（https://docs.jax.dev/en/latest/pallas/tpu/），用 Python 写硬件感知核，榨干 TPU 8i CAE 与 TPU 8t SparseCore 每滴性能。 - **原生 PyTorch 体验**：**TPU 原生 PyTorch 支持**（https://developers.googleblog.com/torchtpu-running-pytorch-natively-on-tpus-at-google-scale/）已开启预览。现有 PyTorch 模型无需改动即可在 TPU 上运行，完整支持 Eager 等原生特性。 - **可移植性**：同一套 JAX、PyTorch 或 Keras 代码从 Ironwood 无缝扩展到本代，XLA 在后台自动搞定 Broadly 拓扑与 CAE 同步，你只关心模型本身。 ### 代际飞跃：性能跃升 软硬件协同设计持续兑现红利，相较第七代 Ironwood TPU，第八代带来： - **训练性价比**：TPU 8t 提升高达 2.7 倍性能/美元。 - **推理性价比**：TPU 8i 提升高达 80% 性能/美元，尤其在大 MoE 低延迟场景。 - **能效**：两代芯片均实现**2 倍性能/瓦**，为 AI 可持续扩展奠基。 ### 展望未来 为了让 Google Cloud 客户率先拥抱下一波创新，我们把 TPU 8t 与 TPU 8i 设计为两套专用系统，分别精准匹配 AI 生命周期不同阶段的多样化需求。二者均深度集成 AI Hypercomputer 软件栈：JAX、PyTorch、vLLM、XLA、Pathways，并与 Google DeepMind 联合从头重构，实现极致性价比与能效。 第八代架构的模块化设计为未来提供清晰而独特的路线图。正如历次计算范式跃迁都需基础设施突破，代理时代亦然。能够“规划-执行-学习”闭环的推理代理，无法在仅为传统训练或交易优化的硬件上高效运行。