跳到主要内容

Tensor Core

Tensor Core 是 NVIDIA 硬件中用于矩阵乘累加(matrix multiply-accumulate)运算的专用执行单元。NVIDIA 在 Volta 架构中引入了它们。一个 warp 或一组 warp 可以使用一条矩阵指令来更新一小块输出分块,而不是发出长长一串标量乘加指令。

Transformer 的大部分时间都花在矩阵乘法上,包括注意力投影和前馈层。因此,只要内核满足所需的精度、分块和布局约束,Tensor Core 就能提供比通用算术单元高得多的吞吐量。

矩阵指令与分块​

在概念层面,一条 Tensor Core 指令执行如下形式的分块运算:

D=A×B+CD = A \times B + C

这四个字母是单次矩阵乘累加的操作数。每个都是小块的、分块大小的矩阵:

  • A 和 B 是被相乘的两个输入矩阵,A × B 是矩阵乘积。
  • C 是累加器输入,即已经累加的值。
  • D 是输出。

应用程序暴露的矩阵可能远大于硬件指令分块。内核会把它们分割成多个分块,将分块片段分发到各线程,发出矩阵乘累加指令,再合并部分结果。

编程接口取决于抽象层级。例如:

  • cuBLAS、cuBLASLt 等库会在内部选择 Tensor Core 内核
  • Triton 等编译器可以从分块张量程序生成 Tensor Core 指令

更高层的接口减少实现工作量。更底层的接口可以对数据搬运和调度解锁更多控制,但也会让内核与单一架构绑定得更紧。

支持的精度​

Tensor Core 并不是在每一代 GPU 上都接受所有类型。支持范围随时间不断扩大:

  • Volta 引入了 FP16 矩阵输入,并支持更高精度的累加。
  • Ampere 增加了 BF16、TF32 等格式和模式,并支持更广泛的整数类型。
  • Hopper 增加了 FP8 Tensor Core 路径和新的矩阵指令。
  • Blackwell 通过 FP4 和 FP6 格式以及第二代 Transformer Engine 进一步扩展了 Tensor Core。

注意,输入精度和累加器精度可以不同。例如,内核可以以 FP16 或 BF16 输入做乘法,同时累加到 FP32 以减少数值误差(求和这一步正是数值误差容易累积的地方)。16 位的累加和可能溢出有限的范围,或者在把较小的值加到很大的总和上时丢失这些较小的值。

输入类型、累加器类型、分块形状和吞吐量的确切组合取决于计算能力。编写代码时应面向具体硬件,而不是假设某一种 Tensor Core 模式在所有地方都有效。

分块与布局要求​

只有周边内核高效地喂数据,Tensor Core 才能带来高吞吐量。重要因素包括:

  • 形状:矩阵维度需要有足够多的完整指令分块。边缘分块可能需要填充、掩码或回退路径。
  • 对齐:指针和主维度通常需要合适的对齐,以实现高效的向量化加载和矩阵指令。
  • 布局:线程期望矩阵片段按架构专属的寄存器布局排列。共享内存布局可能需要 swizzle 或填充,以避免存储体冲突。
  • 精度:输入必须使用受支持的类型和累加模式。
  • 数据供给:HBM 和共享内存的传输必须跟上 Tensor Core 流水线的速度。

内核即使发出了 Tensor Core 指令,也可能仍然运行得很差。如果分块加载主导了执行,或矩阵维度太小,理论上的 Tensor 吞吐量就不会转化为应用的吞吐量。

为什么 Tensor Core 对 LLM 推理很重要​

LLM 推理在很大程度上是一系列大型矩阵乘法(matmul):QKV 投影和输出投影、前馈(MLP)层,以及注意力内部的 matmul。它们正是 Tensor Core 为加速而生的形状,这就是为什么 Tensor Core 吞吐量(以 TFLOPS 或 TOPS 衡量)在很大程度上决定了推理中计算密集部分的运行速度。

如上所述,Tensor Core 也支持低精度输入:FP16、BF16、TF32、INT8、FP8 和 FP4。这使它们直接与量化相关联。把权重和激活应用到更低精度的格式,既能缩小内存占用,又能让它们走上更快的 Tensor Core 路径,因此这两种优化相互促进。

具体来说,对于 LLM 推理的两个阶段:

  • 在预填充期间,许多 token 会同时被处理,matmul 很大,算术强度很高。Tensor Core 得到充分供数,这一阶段通常是计算受限的,所以 Tensor Core 吞吐量才是关键。
  • 在解码期间,瓶颈转向内存带宽(流式读取权重和 KV 缓存),因此原始的 Tensor Core FLOPs 不再成为限制,Tensor Core 部分闲置。对多个序列进行批处理会恢复更大的矩阵维度,并重新拉高 Tensor Core 利用率,这是推理服务器积极批处理的一个主要原因。

因此,Tensor Core 对推理性能至关重要,但并非在所有场景下都同样重要。它们常常为计算密集的预填充和大批量的解码设定上限。小批量的解码通常仍然是带宽受限的,所以单纯增加 Tensor Core 吞吐量可能没有帮助。与占用率一样,它们告诉你计算上限在哪里。

常见问题​

CUDA Core 和 Tensor Core 有什么区别?​

CUDA Core 是用于标准浮点和整数运算的通用执行单元。Tensor Core 是专门为小型分块(例如 4×4 或 16×16,取决于架构和数据类型)上的矩阵乘累加而设计的专用单元。对于像 Transformer 层这样矩阵密集的工作负载,Tensor Core 提供的吞吐量显著高于 CUDA Core。

使用 FP16 或 FP8 会自动激活 Tensor Core 吗?​

不会。运算、维度、布局、对齐和生成的指令都必须匹配受支持的 Tensor Core 路径。务必测试你的代码,确认内核是否真的发出了矩阵指令,以及这些流水线的利用率如何。

扩展阅读​