AI Infra 硬件体系与编程模型：5. Tensor Core 解析

AI Infra 硬核拆解：SM 流处理器的 Tensor Core 工作原理与性能魔法

在我们的 SM 架构系列文章中，我们已经深入解析了 SM 的整体架构和 CUDA Core 的工作原理。今天，我们终于要揭开 AI 时代 GPU 最核心、最神秘也最强大的组件的面纱：Tensor Core（张量核心）。

毫不夸张地说，没有 Tensor Core，就没有今天的大语言模型。正是 Tensor Core 带来的数量级性能提升，让训练万亿参数模型成为可能。但同时，Tensor Core 也是最容易被误解的硬件单元——很多人知道它快，但不知道它为什么快，更不知道如何才能让它真正跑满。

今天，我们就从硬件底层到软件执行，彻底搞清楚 Tensor Core 到底是怎么工作的，以及它是如何实现"性能魔法"的。

一、为什么需要 Tensor Core？CUDA Core 的极限

在深入 Tensor Core 之前，我们先思考一个问题：为什么 NVIDIA 要在 GPU 中加入一个全新的计算单元？CUDA Core 不够用吗？

答案很简单：CUDA Core 的架构设计决定了它在矩阵乘法上的效率有天然的极限。

我们来算一笔账：一个 16x16x16 的矩阵乘法（C = A × B）需要多少次运算？

• 每个元素需要 16 次乘法和 15 次加法
• 总共 16×16×16 = 4096 次乘法和 16×16×15 = 3840 次加法
• 总计约 8000 次运算

如果用 CUDA Core 来做这个运算：

• 每个 CUDA Core 每周期只能执行 1 次 FMA（乘加）运算
• 一个 SM 有 128 个 CUDA Core（H100）
• 理论上需要 8000 ÷ 128 ≈ 62.5 个周期才能完成

但实际上，由于数据依赖、内存访问和指令调度的开销，实际需要的周期数会远远超过这个数字。

而 Tensor Core 呢？一个 H100 的 Tensor Core 只需要 1 个周期就能完成这个 16x16x16 的矩阵乘法！

CUDA 核心一次处理一个融合的乘法和加法运算，也就是说，每条指令都将矩阵中的一个元素A乘以矩阵中的另一个元素B，并将结果累加起来C。而 Tensor 核心则在一条指令内完成 MMA 运算（例如，一个 64×128 的图块）。

因此，性能提升也就不足为奇了。对于维度为 n 的大规模矩阵MxN，我们可以通过将矩阵乘法分解成更小的矩阵乘法（一种称为分块的技术，将在后续博文中讨论）来大幅提升速度，并使用更少的指令执行计算。最终，我们获得了显著的速度提升。

每个SM的Tensor核心数量（如果你好奇的话，是4个）并不重要。这是因为当我们执行MMA指令时，硬件会自动负责MMA在各个Tensor核心之间的分区/分配。

这就是 60 多倍的性能差距。这不是简单的优化，这是计算范式的根本转变：从"一次算一个数"变成"一次算一个矩阵"。

二、Tensor Core 的本质：硬件化的矩阵乘加单元

Tensor Core 不是一个更快的 CUDA Core，它是一个完全不同的计算单元。

CUDA Core 是一个通用的算术逻辑单元（ALU），可以执行各种标量运算。而 Tensor Core 是一个专用的矩阵乘加（Matrix Multiply-Accumulate, MMA）单元，它唯一擅长的就是执行 D = A × B + C 这个操作。

正确的比喻：如果 CUDA Core 是一个普通计算器，一次只能算一个加减乘除；那么 Tensor Core 就是一个矩阵计算器，一次就能算完一整个矩阵乘法。

2.1 基本 MMA 操作

Tensor Core 的所有能力都建立在一个基本操作之上：矩阵乘加（MMA）。

D = A × B + C

其中：

• A 和 B 是输入矩阵，维度为 M×K 和 K×N
• C 是累加矩阵，维度为 M×N
• D 是输出矩阵，维度为 M×N

这个操作看起来简单，但它是神经网络中 90% 以上运算的基础。无论是全连接层、卷积层还是注意力机制，最终都可以分解为一系列的矩阵乘法。

2.2 操作维度的演进

随着架构的演进，Tensor Core 支持的基本 MMA 操作维度也在不断增大：

关键趋势：每一代架构都在翻倍 Tensor Core 的基本操作维度，从而在不增加核心数量的情况下翻倍理论峰值性能。

三、Tensor Core 的硬件结构与工作流程

现在我们来深入 Tensor Core 的内部，看看它是如何在一个周期内完成数千次运算的。

3.1 内部硬件结构

一个现代的 Tensor Core（以 Hopper 为例）主要由以下几个部分组成：

1. 输入寄存器堆：存储输入矩阵 A 和 B 的数据
2. 乘法阵列：由大量乘法器组成的二维阵列，并行执行所有乘法操作
3. 加法树：将乘法结果按行累加，得到中间结果
4. 累加寄存器堆：存储累加矩阵 C 和最终结果 D
5. 控制逻辑：控制数据的流动和运算的执行

核心设计：Tensor Core 采用了空间并行的架构。它不是在时间上串行执行运算，而是在空间上布置了数千个运算单元，让它们同时工作。

3.2 一个周期内的完整工作流程

我们以 H100 的 16×16×16 FP16 MMA 操作为例，看看一个 Tensor Core 在一个时钟周期内都做了什么：

1. 数据加载：从 SM 的寄存器文件中加载 16×16 的矩阵 A 和 16×16 的矩阵 B 到 Tensor Core 的输入寄存器
2. 乘法阶段：256 个乘法器同时工作，计算 A 的每一行与 B 的每一列的元素乘积，得到 256 个中间结果
3. 累加阶段：16 个加法树同时工作，每个加法树将 16 个乘积结果相加，得到 16 个中间和
4. 最终累加：将中间和与累加寄存器中的 C 矩阵对应元素相加
5. 结果写回：将最终结果 D 写回 SM 的寄存器文件

整个过程在一个时钟周期内完成！这就是 Tensor Core 性能魔法的核心——极致的空间并行性。

3.3 累加精度的重要性

你可能已经注意到了，Tensor Core 的输入和输出通常使用不同的精度。例如，FP16 输入，FP32 累加。

为什么要这样做？因为如果全程使用 FP16，累加过程中会产生严重的精度损失，导致模型无法收敛。而使用更高精度的累加器，可以在保持高性能的同时，保证计算精度。

各代架构的累加精度：

• Volta：FP16 输入，FP32 累加
• Ampere：BF16/TF32 输入，FP32 累加
• Hopper：FP8 输入，FP16/FP32 累加
• Blackwell：NVFP4 输入，FP8/FP16/FP32 累加

四、Tensor Core 的调度与执行模型

理解了 Tensor Core 的硬件结构，我们再来看看它是如何与 SM 的整体调度系统协同工作的。

4.1 从 Warp 到 Warp Group

在 CUDA Core 中，SM 以 Warp（32 个线程）为基本调度单位。但在 Tensor Core 中，这个基本单位发生了变化。

• Volta/Ampere：一个 Warp（32 个线程）协同执行一个 MMA 指令
• Hopper/Blackwell：一个 Warp Group（4 个 Warp，共 128 个线程）协同执行一个 MMA 指令

为什么要引入 Warp Group？因为随着 Tensor Core 操作维度的增大，单个 Warp 已经无法提供足够的线程来加载和存储数据。需要更多的线程协同工作，才能喂饱越来越强大的 Tensor Core。

4.2 完整的 MMA 指令执行流程

我们以 H100 的 Warp Group MMA 为例，看看一个完整的 Tensor Core 指令是如何执行的：

1. 数据准备：Warp Group 中的 128 个线程协同工作，将全局内存中的矩阵数据加载到共享内存
2. 数据分发：每个线程负责加载自己的那部分数据到寄存器
3. Warp Group 同步：所有线程到达同步点，确保数据准备完毕
4. MMA 指令发射：Warp 调度器向 Tensor Core 发射 MMA 指令
5. Tensor Core 执行：Tensor Core 在一个周期内完成矩阵乘加运算
6. 结果写回：运算结果写回寄存器文件
7. 后续处理：线程继续执行后续的指令，如激活函数、归一化等

关键优化点：整个过程是高度流水线化的。当一个 MMA 指令在 Tensor Core 上执行时，线程可以同时准备下一个 MMA 指令所需的数据，从而实现计算和数据传输的重叠。

4.3 多个 Tensor Core 的并行工作

每个 SM 包含 4 个独立的 Tensor Core，它们可以并行执行不同的 MMA 指令。

这意味着，一个 H100 的 SM 每周期可以执行：

• 4 个 Tensor Core × 1024 FMA/周期 = 4096 FP16 FMA/周期

而整个 H100 GPU 有 144 个 SM，所以理论峰值性能为：

• 144 SM × 4096 FMA/周期 × 1.83 GHz ≈ 1080 TFLOPS FP16

这就是我们常说的 GPU 理论峰值算力的由来。

五、Tensor Core 支持的精度格式详解

Tensor Core 的性能与它支持的精度格式密切相关。不同的精度格式有不同的吞吐量，选择合适的精度格式是性能优化的关键。

5.1 各代架构支持的精度对比

重要结论：

• 每降低一半的位宽，Tensor Core 的吞吐量通常会翻倍
• FP8 的吞吐量是 FP16 的 2 倍，NVFP4 的吞吐量是 FP8 的 2 倍
• TF32 的吞吐量是 FP16 的一半，但精度与 FP32 相当

5.2 关键精度格式详解

TF32（TensorFloat-32）

• 由 Ampere 架构引入
• 19 位表示：1 位符号，8 位指数，10 位尾数
• 精度与 FP32 相当，吞吐量与 FP16 相当
• 最大优势：不需要修改代码，只需打开一个开关，就能自动利用 Tensor Core 加速 FP32 训练

BF16（BFloat16）

• 由 Google 提出，Ampere 架构开始支持
• 16 位表示：1 位符号，8 位指数，7 位尾数
• 指数范围与 FP32 相同，尾数精度比 FP16 低
• 特别适合深度学习训练，因为神经网络对指数范围的要求远高于尾数精度

FP8（8位浮点数）

• 由 Hopper 架构引入
• 有两种格式：E4M3（4 位指数，3 位尾数）和 E5M2（5 位指数，2 位尾数）
• E4M3 适合前向传播，E5M2 适合反向传播
• 吞吐量是 FP16 的 2 倍，是目前大模型训练的主流精度

NVFP4（NVIDIA 4位浮点数）

• 由 Blackwell 架构引入
• 4 位表示，结合双级缩放机制
• 在保持精度的同时，吞吐量是 FP8 的 2 倍
• 专门为大模型推理优化，能带来 2-4 倍的推理性能提升

六、Tensor Core 与 CUDA Core 的协同工作

在现代 AI 系统中，Tensor Core 和 CUDA Core 不是竞争关系，而是互补关系。它们各自擅长不同的运算，协同工作才能实现最佳性能。

6.1 典型 Transformer 层的运算分布

我们以一个典型的 7B 参数 Transformer 层为例，看看 Tensor Core 和 CUDA Core 分别负责哪些运算：

可以看到：虽然矩阵乘法在理论上占了 80% 以上的运算量，但在实际运行中，由于 Tensor Core 速度极快，CUDA Core 负责的非矩阵运算往往会成为瓶颈。

这就是为什么 FlashAttention 能带来如此巨大的性能提升——它不仅优化了矩阵乘法，更重要的是它将多个非矩阵运算融合到了同一个 Kernel 中，大大减少了 CUDA Core 的开销。

6.2 理想的协同工作模式

在一个优化良好的 AI 系统中，Tensor Core 和 CUDA Core 应该是并行工作的，而不是串行工作的。

理想的流水线：

1. Tensor Core 执行第 N 层的矩阵乘法
2. 同时，CUDA Core 处理第 N-1 层的激活函数和归一化
3. 同时，DMA 控制器加载第 N+1 层的数据

这样，所有硬件单元都在同时工作，没有空闲时间，从而实现最高的整体吞吐量。

挑战：如何精确平衡 Tensor Core 和 CUDA Core 的工作负载，避免其中一个成为瓶颈。这需要对运算的计算量、内存访问量和延迟进行精确的建模和分析。

七、Tensor Core 性能优化的核心要点

理解了 Tensor Core 的工作原理，我们就可以针对性地进行性能优化。以下是每个 AI Infra 工程师都必须掌握的核心优化原则：

7.1 确保矩阵维度对齐

这是最基本也是最重要的一点。Tensor Core 对矩阵维度有严格的要求，只有对齐的矩阵才能获得最佳性能。

对齐要求：

• M、N、K 维度必须是 16 的倍数（对于 FP16/BF16）
• 对于 FP8，必须是 32 的倍数
• 对于 NVFP4，必须是 64 的倍数

如果维度不对齐会发生什么？Tensor Core 仍然可以工作，但会有部分运算单元闲置，导致性能下降。例如，如果 K 维度是 17，那么 Tensor Core 会执行两次 16 的运算，其中第二次只有 1 个元素是有效的，性能下降 50%。

最佳实践：在设计模型时，尽量让所有矩阵维度都是 128 的倍数。如果必须使用非对齐维度，可以在矩阵末尾填充零来对齐。

7.2 使用正确的数据布局

Tensor Core 对数据布局非常敏感。不同的数据布局会导致完全不同的性能。

推荐的数据布局：

• 对于矩阵 A：行优先布局
• 对于矩阵 B：列优先布局

为什么？因为这样可以让 Tensor Core 以连续的方式访问内存，实现最高的内存带宽利用率。

常见的错误：使用 NHWC 布局而不是 NCHW 布局。NHWC 布局会导致 Tensor Core 访问非连续的内存地址，性能下降 2-3 倍。

7.3 充分利用共享内存

Tensor Core 的数据来自两个地方：寄存器文件和共享内存。共享内存是 Tensor Core 性能的关键。

为什么需要共享内存？

• 全局内存的延迟太高（约 400-800 个周期）
• 寄存器文件的容量有限，无法存储大的矩阵块
• 共享内存的延迟只有约 20-30 个周期，带宽是全局内存的 10 倍以上

优化方法：

• 将大矩阵分成小块（Tile），每次只加载一个小块到共享内存
• 使用异步内存拷贝（Async Copy）将数据从全局内存加载到共享内存
• 使用预取技术，提前加载下一个小块的数据

7.4 避免共享内存 Bank 冲突

共享内存被分成 32 个 Bank，每个 Bank 每个周期只能被访问一次。如果多个线程同时访问同一个 Bank，就会发生 Bank 冲突，导致性能下降。

Tensor Core 特有的 Bank 冲突问题：由于 Tensor Core 以 16×16 的块访问共享内存，如果数据布局不正确，很容易导致严重的 Bank 冲突。

解决方法：

• 在共享内存中添加填充（Padding）
• 使用转置操作改变数据布局
• 使用 NVIDIA 推荐的优化数据布局

7.5 不要自己写 Tensor Core 代码

除非你是 NVIDIA 的工程师，否则不要自己写 Tensor Core 代码。

Tensor Core 的编程非常复杂，需要考虑维度对齐、数据布局、Bank 冲突、流水线等众多因素。即使是经验丰富的工程师，也很难写出比 cuBLAS 和 CUTLASS 更好的代码。

最佳实践：

• 对于标准的矩阵乘法，直接使用 cuBLAS
• 对于自定义的融合算子，使用 CUTLASS 库
• 只有在万不得已的情况下，才使用 WMMA API 直接编程 Tensor Core

八、常见误区与最佳实践

误区 1：只要用了 Tensor Core 就一定快

真相：Tensor Core 只有在处理大矩阵乘法时才能发挥出最佳性能。对于小矩阵乘法，CUDA Core 可能比 Tensor Core 更快。

例子：一个 4×4×4 的矩阵乘法，用 CUDA Core 只需要几个周期，而用 Tensor Core 会有很大的调度开销，反而更慢。

最佳实践：对于 M、N、K 小于 64 的矩阵乘法，使用 CUDA Core 实现。

误区 2：Tensor Core 的利用率越高越好

真相：高 Tensor Core 利用率并不一定意味着高整体性能。有时候，为了提高 Tensor Core 利用率，会增加更多的内存访问开销，导致整体性能下降。

例子：FlashAttention 的 Tensor Core 利用率通常只有 60-70%，但它的整体性能比 cuBLAS 实现高 2-3 倍，因为它大大减少了内存访问。

最佳实践：关注整体吞吐量，而不是单一硬件单元的利用率。

误区 3：所有精度的 Tensor Core 性能都一样

真相：不同精度的 Tensor Core 吞吐量差异很大。FP4 的吞吐量是 FP16 的 4 倍，但很多人仍然在使用 FP16 进行推理。

最佳实践：对于推理工作负载，尽量使用最低的可行精度。从 FP16 迁移到 FP8 可以获得 2 倍的性能提升，从 FP8 迁移到 NVFP4 可以再获得 2 倍的性能提升。

九、总结与学习建议

Tensor Core 是 AI 时代 GPU 最伟大的创新之一。它通过专用化的硬件设计，实现了数量级的性能提升，为大语言模型的发展奠定了基础。

核心要点回顾：

1. Tensor Core 是硬件化的矩阵乘加单元，一次能完成一个完整的矩阵乘法
2. 它采用空间并行架构，在一个周期内可以执行数千次运算
3. 不同精度格式的吞吐量差异很大，选择合适的精度是性能优化的关键
4. Tensor Core 和 CUDA Core 是互补关系，协同工作才能实现最佳性能
5. 性能优化的核心是：维度对齐、正确的数据布局、充分利用共享内存

学习建议：

1. 阅读 NVIDIA 的《Tensor Core 编程指南》，了解 WMMA API 的基本使用
2. 学习 CUTLASS 库的源码，看看它是如何优化 Tensor Core 性能的
3. 深入研究 FlashAttention 的实现，学习如何融合算子和优化内存访问
4. 使用 NVIDIA Nsight Compute 工具分析 Kernel 的执行情况，查看 Tensor Core 利用率和内存带宽利用率

理解 Tensor Core 的工作原理，是成为一名优秀 AI Infra 工程师的必经之路。它能让你透过现象看本质，真正理解 AI 系统的性能瓶颈在哪里，从而写出真正高效的代码。