TPU(张量处理单元)
也称为 TPU、Tensor Processing Unit、Google TPU、Cloud TPU
先用大白话
TPU 是 Google 为自家 AI 负载设计的加速器,现在也开始对外出售或出租,用于其他公司的推理服务。
技术定义
TPU(Tensor Processing Unit)是 Google 自研的加速器,核心由大尺寸 systolic 矩阵单元、片上向量存储(VMEM)、处理不规则计算的 SparseCore,以及名为 ICI 的芯片间互连组成。
工程细节
Google 内部的搜索、广告、YouTube 和每一代 Gemini 都运行在 TPU 上。每颗芯片包含负责矩阵乘法的 TensorCore(内部是 MXU systolic 阵列)和负责 embedding 查表与数据搬运的 SparseCore。XLA 编译器在这些单元之间调度计算,编译器切分不好的算子则由 Pallas kernel 补齐。芯片之间通过 ICI 直接互连,组成数千颗芯片的 pod,数据不经过主机 CPU。Ironwood(TPUv7)是 Google 第一代面向外部推理客户的产品,TPUv8 则拆分为训练用的 8t 和推理用的 8i。
为什么重要
这种设计用单芯片峰值换取系统级成本优势。宽矩阵单元、低延迟 torus 网络和编译器驱动的调度,让 Google 的每 token 成本低于规格表的推算;但 tile 几何会惩罚填不满阵列的模型形状。TPU 能否被外部广泛采用,取决于 PyTorch 原生的 TorchTPU 栈能否在 vLLM 和 SGLang 中追平 CUDA 生态。
如何在 InferenceX 中解读
InferenceX Official Preview 发布了首批第三方 TPUv7 结果,在 FP8 聚合服务下与 B200、B300 对比,Ironwood 的每美元性能最高高出 50%。TorchTPU 栈开源后,AgentX 和 TPU 分离式服务的结果也会陆续发布。