4. TPU架构总览

Google TPU（Tensor Processing Unit）是专为神经网络推理与训练设计的领域专用加速器（DSA）。它并非通用GPU的替代品，而是围绕张量运算这一核心负载，从指令集、数据通路到系统互连进行垂直优化的产物。理解TPU，需要同时看清三条线索：代际演进带来的能力边界变化、设计哲学决定的取舍逻辑，以及核心组件与系统层级构成的完整执行栈。

一、设计哲学：以确定性吞吐换通用性

TPU的第一性原理是：神经网络的绝大部分计算可归结为矩阵乘加与激活函数。与其让大量通用ALU去猜测下一条指令，不如用脉动阵列（Systolic Array）让数据在计算单元间有节奏地流动，实现高复用、低功耗的乘加运算。其核心取舍体现在三点：

• 确定性执行：TPU采用静态调度，编译器提前编排数据搬运与计算，避免运行时动态调度开销，适合批处理场景。
• 大容量片上存储：以软件管理的Scratchpad Memory替代复杂缓存层级，减少不可预测的访存延迟。
• 量化优先：从INT8推理到BF16训练，TPU在数值格式上主动裁剪，以更小位宽换取更高吞吐与能效。

二、代际演进：从推理专用到训练集群

TPU的演进并非简单堆算力，而是逐步补齐训练、互连与生态短板。下表梳理了主要代际的关键特征：

可以看到，v1到v4的主线是从单卡推理走向Pod级训练：算力提升之外，互连拓扑、光交换与编译器协同成为决定实际性能的关键。v5系列则进一步细分市场，用不同规格覆盖成本敏感与极致性能两类需求。

三、核心组件：一张TPU芯片里有什么

单颗TPU的核心组件可归纳为五类：矩阵乘法单元（MXU）负责脉动阵列乘加；向量处理单元（VPU）处理激活、归一化等逐元素运算；标量单元执行控制流与地址计算；片上存储包括累加器、权重与激活缓冲区；高带宽互连接口连接HBM与ICI。下图展示了这些组件在数据流中的位置关系：

四、系统层级：从芯片到Pod的四级结构

TPU的性能释放依赖清晰的层级划分。第一级是芯片级，MXU与VPU协同完成单卡计算；第二级是板卡级，多颗芯片通过ICI组成Torus拓扑；第三级是Pod级，成百上千芯片通过光交换或高速链路互联，形成统一地址空间；第四级是集群级，多个Pod通过数据中心网络连接，配合XLA与JAX/PyTorch前端完成分布式训练。

五、案例场景：一次典型训练任务的执行路径

以JAX在TPU v4 Pod上训练一个Transformer为例，执行路径大致如下：

1. 前端将模型计算图交给XLA，XLA进行算子融合与布局优化。
2. 编译器将矩阵乘映射到MXU，将LayerNorm、GELU等映射到VPU。
3. DMA引擎按静态调度将权重与激活从HBM搬入片上存储。
4. MXU完成乘加后，结果经VPU后处理，再写回累加器或HBM。
5. 梯度同步通过ICI在Pod内All-Reduce，跨Pod则走数据中心网络。

# 典型XLA编译日志片段（示意）
XLA: Compiling computation...
  Fusion: matmul + bias_add + relu -> MXU kernel
  Layout: weight [4096,4096] -> tiled for systolic array
  Memory: activation buffer 32MB, weight buffer 64MB
  Schedule: static, 128 steps, double buffering enabled
  ICI: all-reduce ring across 256 chips

六、常见误区与复盘