TPU (Tensor Processing Unit)

张量处理单元 (TPU) 是 Google 设计的一种专用集成电路 (ASIC)，用于加速神经网络的训练和推理。TPU 以大型矩阵乘法单元为核心构建，因为矩阵乘法占据了深度学习的大部分计算量。自 2015 年以来，Google 一直在其数据中心运行 TPU，并使用 TPU 训练 Gemini 模型，还通过 Cloud TPU向其他公司出租 TPU。2026 年 4 月公布的第八代产品分为两款芯片：用于训练的 TPU 8t 和用于推理的 TPU 8i。

TPU 的工作原理#

每个 TPU 的核心是矩阵乘法单元（MXU），它以脉动阵列形式排列，由相互传递数据的乘加单元构成。权重加载到单元中一次后，输入每个时钟周期在网格中前进一格，部分和则沿列向下传递，因此从内存中读取的每个值都能被多个单元重复使用。读取内存所需的时间和能量远高于执行算术运算，这正是该设计高效的原因。在 Cloud TPU v2 和 v3 中，每个 MXU 都是 128×128 阵列，使用 bfloat16执行乘法，并以 FP32累加。下面用一个简短模拟展示 4×4 网格上的数据流：

import numpy as np

rng = np.random.default_rng(0)
W = rng.integers(-5, 5, (4, 4))  # weights, preloaded one per cell of a 4x4 grid
X = rng.integers(-5, 5, (6, 4))  # six input rows stream through the grid
Y = np.zeros((6, 4), dtype=int)

# Cell (k, j) holds W[k, j]. Input X[i, k] enters row k from the left and reaches
# cell (k, j) at cycle i + k + j, where it is multiplied by the stored weight and
# added to the partial sum for Y[i, j] passing down column j
cycles = 6 + 4 + 4 - 2
for t in range(cycles):
    for k in range(4):
        for j in range(4):
            i = t - k - j
            if 0 <= i < 6:
                Y[i, j] += X[i, k] * W[k, j]

print(f"{X.size * 4} multiply-adds in {cycles} cycles, equal to X @ W: {np.array_equal(Y, X @ W)}")

它会打印 96 multiply-adds in 12 cycles, equal to X @ W: True：16 个单元会计算出精确的矩阵乘积，每个权重只从内存中读取一次，随后供全部六行输入重复使用。第一代 TPU在 256×256 的 8 位单元网格上执行这种计算模式，每个周期可完成 65,536 次乘加运算，峰值性能为每秒 92 万亿次运算。其简单、确定性的执行方式符合 Google 的目标，即满足推理任务第 99 百分位数的响应时间目标。

TPU 运行由 XLA编译的代码。XLA 会获取来自 JAX、TensorFlow或 PyTorch的整个计算图，并将其映射到矩阵单元上。芯片通过专用互连组成包含数千个芯片的 Pod，这使得最大的基础模型能够同时在多个 TPU 上训练。

TPU、GPU 与 CPU 的比较#

• CPU：由少量强大的核心构成，擅长处理分支和通用代码；在大多数训练流水线中，CPU 负责数据加载和预处理。
• GPU：拥有数千个可编程核心，以及用于矩阵运算的 Tensor Core。GPU 能够运行几乎任何模型和框架，并由许多供应商和云服务商销售。
• TPU：牺牲灵活性，换取大型稠密矩阵运算的效率。静态张量形状的工作负载效果最好，因为形状变化会导致 XLA 重新编译；自定义运算、大量分支或高精度算术则不太适合 TPU。
• NPU：手机、笔记本电脑和边缘芯片中神经网络加速器的统称，Edge TPU 也属于这一类别。

Edge TPU#

Edge TPU是一款仅用于推理的小型 TPU，也是 Google Coral USB 加速器、M.2 和 PCIe 模块以及开发板中的芯片。它以 2 瓦功耗每秒执行 4 万亿次运算，并且只运行 8 位整数运算，因此网络必须通过模型量化完全转换为 INT8，然后使用 Google 的 Edge TPU Compiler 编译；编译器会将受支持的运算放到芯片上执行，其余运算则留给主机 CPU。Ultralytics YOLO 通过 format="edgetpu"自动完成这两个步骤：该参数会启用 INT8 量化，并在 x86-64 Linux 计算机上运行编译器，生成可在配备 Coral USB 加速器的 Raspberry Pi等开发板上运行的模型：

from ultralytics import YOLO

# Load the smallest YOLO26 detection model
model = YOLO("yolo26n.pt")

# Quantize to INT8 and compile for the Edge TPU (x86-64 Linux only)
path = model.export(format="edgetpu")
print(f"Edge TPU model: {path}")

Google 的 Coral 平台目前以 Coral NPU为核心。这是一种于 2025 年 10 月发布、面向可穿戴设备的开源 RISC-V 加速器设计，目标是在数毫瓦功耗下实现每秒 5120 亿次运算。