跳转至

TPU 执行与数据流⚓︎

理解 TPU 时，不能只看峰值 FLOPs，还要看算子落在哪个执行单元、数据经过哪些路径，以及这些路径上的带宽是否足够支撑计算吞吐。对 LLM 而言，TPU 的性能分析核心不是“有没有很多算力”，而是“矩阵相关算子是否能够以足够高的复用率把数据送入 MXU，以及跨芯片通信是否会反过来卡住计算”。

TPU 计算视角⚓︎

从最简化的角度看，TPU 可以理解为一个专门面向矩阵乘法的计算核心，连接着一组高带宽但容量有限的片上资源和一组容量更大的外部存储。对主要部件可以先建立以下分工：

  • MXU：负责矩阵乘法，是 TPU 上最核心的吞吐来源。
  • VPU：负责逐元素运算、向量算术和部分归约操作。
  • VMEM：靠近计算单元的片上 scratchpad，用于承接当前计算阶段所需的数据块。
  • HBM：容量更大的外部高带宽内存，用于存放权重、激活和结果张量。

其中最关键的事实是：TPU 并不会直接从 HBM 把数据送进 MXU 做计算。数据通常先从 HBM 进入 VMEM，再从 VMEM 进入寄存器或计算阵列，最后结果再写回 VMEM 与 HBM。因此，一个算子的性能不仅取决于 FLOPs 总量，还取决于它是否落在 MXU 上、是否能够持续喂满 MXU，以及这条数据路径上的每一级带宽是否足够。

从这个意义上说，TPU 性能问题往往不是“算力够不够”，而是“数据能否按正确路径、以足够高的复用率送到正确计算单元”。只有把计算单元和数据路径一起看，roofline 才有解释力。

片上存储与数据流⚓︎

TPU 文档里常出现 VMEM、SMEM、TMEM 等片上存储名称。理解这些名词时，更重要的是区分它们在数据流中的职责，而不是机械记忆缩写。

从性能视角看，HBM 提供大容量但相对更慢的外部带宽，片上存储则负责承接当前 tile 所需的数据，并在计算阶段尽可能复用。只要某个输入块或中间结果能够留在片上，就能减少一次 HBM 往返；只要某个 tile 因容量不足、调度方式不佳或切分粒度不合理而被迫反复搬运，就会直接降低有效带宽利用率。

一个重要的数量级关系是，VMEM 到 MXU 的带宽通常显著高于 HBM 到 VMEM 的带宽。对 TPU v5e，这个比例可接近二十倍量级。其直接含义是：

  • 如果一个算子的输入、权重和输出都能主要在 VMEM 内部周转，那么达到峰值 FLOPs 所需的算术强度会明显下降。
  • 如果权重能够提前 prefetched 到 VMEM，那么一部分原本来自 HBM 的搬运成本可以被隐藏在其他算子的计算阶段中。
  • 如果权重太大、切分方式不合适或 VMEM 空间不足，算子就仍会被 HBM 带宽主导。

因此，TPU 上很多“算力没有吃满”的现象，根源并不在公式本身，而在于数据是否以合适粒度留在片上，以及不同存储层之间的搬运是否成为主瓶颈。

矩阵乘法在 TPU 上的执行⚓︎

TPU 的矩阵乘法不是把一个大矩阵整体送入 MXU，而是先把矩阵切成多个块，再逐块执行乘法与局部累加。其根本原因不是算法变化，而是单次可放入片上存储和计算阵列的数据规模有限。

以常见 TPU 代际为例，MXU 的核心是一个 128 x 128 的 systolic array；TPU v6e 则提升到 256 x 256。对前者而言，当阵列充分饱和时，可在数个 cycle 内完成一次形如：

bfloat16[8, 128] @ bfloat16[128, 128] \rightarrow float32[8, 128]

的矩阵乘法片段。需要注意两点：

  • 激活块和权重块不是同时一次性装入，而是以流式方式送入阵列。
  • 输出也不是等整个大矩阵完成后才统一回写，而是随着局部累加逐步流出。

因此，一次大 matmul 的真实执行过程更接近：加载一个 tile、执行局部乘加、保留部分和、继续加载下一组 tile、重复累加，直到完成整块输出。如果输入块或权重块能够在多个步骤中复用，实际搬运成本就会下降；如果同一个块因为 tile 设计不当被多次回读，理论上的高算术强度就无法完全兑现为实际吞吐。

systolic array 的执行方式也解释了为什么 TPU 特别适合矩阵乘法。矩阵乘法具有较高的数据复用率，同一批权重或激活在阵列内部可以支撑大量乘加，从而把计算密度维持在较高水平。这也是 MXU 的峰值吞吐远高于 VPU 的主要原因之一。

VPU 与非矩阵算子⚓︎

与 MXU 相比，VPU 更适合逐元素运算、向量加减、激活函数和部分归约。其峰值吞吐通常显著低于 MXU，可能低一个数量级左右。这意味着：

  • 一个模型即使主要计算量来自 matmul，仍可能在某些阶段被 VPU 路径或归约路径拖慢。
  • 某些低算术强度或跨 lane 归约较重的操作，即使 FLOPs 总量不大，也不容易逼近 MXU 那样的利用率。

因此，在 TPU 上分析算子瓶颈时，第一步常常不是计算 FLOPs，而是先判断它到底落在哪个执行单元上。

tile 与带宽约束⚓︎

tile 大小决定了三个关键量：单次进入片上的数据块规模、块内可实现的复用程度，以及为了完成整次 matmul 需要重复多少轮 HBM 搬运。tile 太小，会导致边界开销和重复加载占比上升；tile 太大，又可能超出片上存储或调度友好的范围。

这意味着理论上由 FLOPs 与 HBM 字节数直接算出的算术强度，通常只是理想上界。真实实现里，还要扣除以下因素：

  • tile 切分导致的重复读取
  • 片上缓存容量与寄存器资源限制
  • 数据在 HBM、VMEM、寄存器和计算阵列之间搬运的额外成本
  • 初始 pipeline bubble 与不同 tile 之间的调度开销

因此，同一个矩阵形状在不同 tile 配置下，可能表现出不同的有效 roofline 位置。对大型 matmul 而言，理想情况下的算术强度可近似由 tile 维度决定；而在 TPU 执行视角下，这个结论可以进一步理解为：tile 设计本质上决定了一个块在进入 MXU 之前和之后要为复用付出多少搬运代价。

另一个直接影响是张量尺寸必须足够大，才能把 MXU 填满。对常见 128 x 128 的阵列而言，若矩阵维度明显小于 128，则硬件会出现 padding 或利用率不足的问题；在有多个 MXU 的代际中，若 tiling 规模仍偏小，也很难充分吃满整体吞吐。TPU v6e 由于阵列扩大到 256 x 256，这一约束会进一步增强：张量维度需要更大，才能维持同等利用率。

TPU 与 roofline 的关系⚓︎

把 TPU 的执行路径放回 roofline 框架后，可以把问题分成两层：

  • 第一层是算子理论上有多少 FLOPs、多少字节搬运。
  • 第二层是这些字节究竟来自 HBM、片上存储重载，还是更慢的跨设备或跨主机通信。

对 MXU 上的大型 matmul，roofline 往往更接近“是否能把 tile 以足够高复用率送入计算阵列”的问题；对低复用算子、控制型算子或 VPU 主导算子，则更容易受带宽与调度路径限制。也就是说，TPU 上的 roofline 不能只看静态公式。一个算子即使纸面上有很高算术强度，只要 tile 设计导致同一数据块需要反复从 HBM 读取，实际表现仍可能偏向 memory-bound。反过来，如果数据路径顺畅、片上复用充分，算子就更容易逼近 MXU 的峰值吞吐。

一个重要经验是，如果一个矩阵乘法能主要从 VMEM 而不是 HBM 读写，其达到 peak FLOPs 利用率所需的算术强度会大幅降低。这意味着同样的矩阵形状，在“权重每次都从 HBM 读”和“权重已预取到 VMEM”两种实现下，roofline 位置可能完全不同。

用算例理解执行路径的影响⚓︎

下面用几个算例归纳执行路径如何改变瓶颈判断。

从 HBM 加载大模型参数给出单步延迟下界⚓︎

若一个 200B 参数、bfloat16 存储的模型被切到 32 张 TPU v4p 上，则总参数量约为：

400 \times 10^9 \text{ bytes}

平均到每张卡约为 12.5e9 bytes。若每卡 HBM 带宽约为 1.23e12 bytes/s，则仅把参数从 HBM 送入计算路径的时间下界约为 10ms。这个估计的重要性不在于精确到单毫秒，而在于它说明：单步推理延迟不可能低于参数加载时间的下界。

PCIe 可以成为比 HBM 更强的外层瓶颈⚓︎

若权重和激活不在 HBM，而在 host DRAM，需要经由 PCIe 送入 TPU，则瓶颈会立即外移。因为 PCIe 带宽通常比 HBM 再低一个数量级以上，所以一个本来在片内接近 compute-bound 的算子，迁移到 host-to-device 路径后，可能完全变成 PCIe-bound。在这种场景下，为了保持 FLOPs-bound，所需 batch 往往要显著更大。

VMEM 可显著降低临界 batch⚓︎

以 int8[16384, 4096] 权重矩阵与 int8[B, 4096] 激活矩阵为例，在 TPU v5e 上，如果主要从 HBM 读取数据，则进入 FLOPs-bound 的临界 batch 大约在几百量级；若主要从 VMEM 读取，由于带宽提升约二十倍，临界 batch 可以降到十几到几十的量级。这一结论再次表明：同一个矩阵形状，其瓶颈并不只由 FLOPs 决定，而高度依赖数据从哪里来。

TPU 网络与跨设备通信⚓︎

单芯片内部之外，TPU 还要面对跨芯片网络约束。TPU 芯片在 pod 内通常通过 ICI 相连，而不同 slice 或不同 host 之间则通过 DCN 连接。与 GPU 倾向于交换机层级互联不同，TPU 的 ICI 更接近最近邻 torus 网络，这意味着通信成本强烈依赖拓扑形状、是否存在 wraparound，以及数据需要跨越多少 hop。

从带宽层次看，可以按以下顺序理解：

  • HBM 带宽最快，决定片内大规模数据供给能力。
  • ICI 带宽次之，决定 slice 内跨 TPU 通信效率。
  • PCIe 更慢，决定 host 与 TPU 之间的数据装载速度。
  • DCN 更慢，决定跨 slice 或跨 pod 的外层通信能力。

这带来一个非常直接的工程原则：每一层链路上的通信量，都需要与该层链路的速度相匹配。如果模型切分方式导致大量数据必须跨 DCN 或 PCIe 往返，那么再高的 MXU 峰值也很难转化为端到端吞吐。

网络算例的含义⚓︎

以 4 x 4 TPU v5e slice 上的数组传输为例，哪怕数组规模本身不算极端，通信完成时间也由两部分决定：

  • 首字节到达时间由 hop 数与每 hop 延迟决定。
  • 全量传输时间由有效带宽与可并行使用的链路数决定。

这说明 TPU 通信不能只看带宽，也要看拓扑。特别是在没有 wraparound 的小拓扑中，远端芯片之间的通信 hop 数会变长，首字节延迟和链路压力都会增加。

更复杂的综合场景也说明，如果一个大矩阵先分散保存在 host DRAM 或多个 TPU 上，最终又必须汇聚到单个 TPU 执行 matmul，那么端到端时间往往由最慢的一段决定，而不是由 FLOPs 决定。在这种场景里，ICI gather 往往比最终的 HBM 读取和 MXU 计算更慢，因此整体时延主要受跨卡搬运控制。

工程判断⚓︎

分析 TPU 上的性能问题时，可以按以下顺序判断：

  • 先看算子主要落在哪个计算单元，是 MXU 还是 VPU。
  • 再看数据主要来自 HBM、VMEM 复用，还是 host / 跨芯片路径。
  • 再看 tile 设计是否造成重复加载、padding 浪费或片上资源不足。
  • 最后再把这些量映射到 roofline 框架里，判断瓶颈位于算力、片内带宽还是跨设备通信。

对矩阵乘法类算子，判断重点通常不是“有没有足够多的 FLOPs”，而是“是否把足够大的 tile 以足够高的复用率送入 MXU”。对逐元素或低复用算子，则更容易直接受带宽限制。对跨设备场景，还需要额外判断通信是否主要走 ICI、PCIe 还是 DCN，因为这三者的量级差异足以彻底改变瓶颈位置。

可以把 TPU 的几个稳定结论归纳为：

  • TPU 最擅长高复用、高密度的矩阵乘法。
  • VMEM 的价值在于把原本受 HBM 限制的算子拉向更高利用率区域。
  • tile 设计决定理论算术强度能否兑现为实际吞吐。
  • 小矩阵、低复用算子或 VPU 主导算子更容易被带宽拖住。
  • 多卡场景中，ICI、PCIe 和 DCN 往往比 MXU 峰值更早成为瓶颈。

因此，理解 TPU 的关键不是背规格表，而是把执行单元、数据路径、片上复用与 roofline 放进同一个分析框架里。只有这样，才能解释为什么同一台 TPU 上不同算子的性能特征会完全不同。

评论