AI Infra 硬核拆解SM 流处理器的 Tensor Core 工作原理与性能魔法在我们的 SM 架构系列文章中我们已经深入解析了 SM 的整体架构和 CUDA Core 的工作原理。今天我们终于要揭开 AI 时代 GPU 最核心、最神秘也最强大的组件的面纱Tensor Core张量核心。毫不夸张地说没有 Tensor Core就没有今天的大语言模型。正是 Tensor Core 带来的数量级性能提升让训练万亿参数模型成为可能。但同时Tensor Core 也是最容易被误解的硬件单元——很多人知道它快但不知道它为什么快更不知道如何才能让它真正跑满。今天我们就从硬件底层到软件执行彻底搞清楚 Tensor Core 到底是怎么工作的以及它是如何实现性能魔法的。一、为什么需要 Tensor CoreCUDA 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 CoreH100理论上需要 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×NC 是累加矩阵维度为 M×ND 是输出矩阵维度为 M×N这个操作看起来简单但它是神经网络中 90% 以上运算的基础。无论是全连接层、卷积层还是注意力机制最终都可以分解为一系列的矩阵乘法。2.2 操作维度的演进随着架构的演进Tensor Core 支持的基本 MMA 操作维度也在不断增大架构基本 MMA 维度每周期运算量FP16Volta4×4×4128 FMA/周期Turing8×8×4256 FMA/周期Ampere16×8×8512 FMA/周期Hopper16×16×161024 FMA/周期Blackwell16×16×322048 FMA/周期关键趋势每一代架构都在翻倍 Tensor Core 的基本操作维度从而在不增加核心数量的情况下翻倍理论峰值性能。三、Tensor Core 的硬件结构与工作流程现在我们来深入 Tensor Core 的内部看看它是如何在一个周期内完成数千次运算的。3.1 内部硬件结构一个现代的 Tensor Core以 Hopper 为例主要由以下几个部分组成输入寄存器堆存储输入矩阵 A 和 B 的数据乘法阵列由大量乘法器组成的二维阵列并行执行所有乘法操作加法树将乘法结果按行累加得到中间结果累加寄存器堆存储累加矩阵 C 和最终结果 D控制逻辑控制数据的流动和运算的执行核心设计Tensor Core 采用了空间并行的架构。它不是在时间上串行执行运算而是在空间上布置了数千个运算单元让它们同时工作。3.2 一个周期内的完整工作流程我们以 H100 的 16×16×16 FP16 MMA 操作为例看看一个 Tensor Core 在一个时钟周期内都做了什么数据加载从 SM 的寄存器文件中加载 16×16 的矩阵 A 和 16×16 的矩阵 B 到 Tensor Core 的输入寄存器乘法阶段256 个乘法器同时工作计算 A 的每一行与 B 的每一列的元素乘积得到 256 个中间结果累加阶段16 个加法树同时工作每个加法树将 16 个乘积结果相加得到 16 个中间和最终累加将中间和与累加寄存器中的 C 矩阵对应元素相加结果写回将最终结果 D 写回 SM 的寄存器文件整个过程在一个时钟周期内完成这就是 Tensor Core 性能魔法的核心——极致的空间并行性。3.3 累加精度的重要性你可能已经注意到了Tensor Core 的输入和输出通常使用不同的精度。例如FP16 输入FP32 累加。为什么要这样做因为如果全程使用 FP16累加过程中会产生严重的精度损失导致模型无法收敛。而使用更高精度的累加器可以在保持高性能的同时保证计算精度。各代架构的累加精度VoltaFP16 输入FP32 累加AmpereBF16/TF32 输入FP32 累加HopperFP8 输入FP16/FP32 累加BlackwellNVFP4 输入FP8/FP16/FP32 累加四、Tensor Core 的调度与执行模型理解了 Tensor Core 的硬件结构我们再来看看它是如何与 SM 的整体调度系统协同工作的。4.1 从 Warp 到 Warp Group在 CUDA Core 中SM 以 Warp32 个线程为基本调度单位。但在 Tensor Core 中这个基本单位发生了变化。Volta/Ampere一个 Warp32 个线程协同执行一个 MMA 指令Hopper/Blackwell一个 Warp Group4 个 Warp共 128 个线程协同执行一个 MMA 指令为什么要引入 Warp Group因为随着 Tensor Core 操作维度的增大单个 Warp 已经无法提供足够的线程来加载和存储数据。需要更多的线程协同工作才能喂饱越来越强大的 Tensor Core。4.2 完整的 MMA 指令执行流程我们以 H100 的 Warp Group MMA 为例看看一个完整的 Tensor Core 指令是如何执行的数据准备Warp Group 中的 128 个线程协同工作将全局内存中的矩阵数据加载到共享内存数据分发每个线程负责加载自己的那部分数据到寄存器Warp Group 同步所有线程到达同步点确保数据准备完毕MMA 指令发射Warp 调度器向 Tensor Core 发射 MMA 指令Tensor Core 执行Tensor Core 在一个周期内完成矩阵乘加运算结果写回运算结果写回寄存器文件后续处理线程继续执行后续的指令如激活函数、归一化等关键优化点整个过程是高度流水线化的。当一个 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 各代架构支持的精度对比架构支持的精度格式相对吞吐量FP161典型应用VoltaFP161早期深度学习训练TuringFP16, INT8, INT41, 2, 4推理加速AmpereFP16, BF16, TF32, INT81, 1, 0.5, 2主流训练和推理HopperFP16, BF16, TF32, FP8, INT81, 1, 0.5, 2, 2大模型训练和推理BlackwellFP16, BF16, TF32, FP8, INT8, NVFP41, 1, 0.5, 2, 2, 4极致推理性能重要结论每降低一半的位宽Tensor Core 的吞吐量通常会翻倍FP8 的吞吐量是 FP16 的 2 倍NVFP4 的吞吐量是 FP8 的 2 倍TF32 的吞吐量是 FP16 的一半但精度与 FP32 相当5.2 关键精度格式详解TF32TensorFloat-32由 Ampere 架构引入19 位表示1 位符号8 位指数10 位尾数精度与 FP32 相当吞吐量与 FP16 相当最大优势不需要修改代码只需打开一个开关就能自动利用 Tensor Core 加速 FP32 训练BF16BFloat16由 Google 提出Ampere 架构开始支持16 位表示1 位符号8 位指数7 位尾数指数范围与 FP32 相同尾数精度比 FP16 低特别适合深度学习训练因为神经网络对指数范围的要求远高于尾数精度FP88位浮点数由 Hopper 架构引入有两种格式E4M34 位指数3 位尾数和 E5M25 位指数2 位尾数E4M3 适合前向传播E5M2 适合反向传播吞吐量是 FP16 的 2 倍是目前大模型训练的主流精度NVFP4NVIDIA 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 分别负责哪些运算运算类型理论占比实际占比执行单元QKV 投影矩阵乘法35%40%Tensor Core注意力分数矩阵乘法25%30%Tensor Core输出投影矩阵乘法20%20%Tensor Core激活函数SiLU8%5%CUDA Core SFU归一化RMSNorm7%3%CUDA CoreSoftmax4%1%CUDA Core其他残差、Dropout1%1%CUDA Core可以看到虽然矩阵乘法在理论上占了 80% 以上的运算量但在实际运行中由于 Tensor Core 速度极快CUDA Core 负责的非矩阵运算往往会成为瓶颈。这就是为什么 FlashAttention 能带来如此巨大的性能提升——它不仅优化了矩阵乘法更重要的是它将多个非矩阵运算融合到了同一个 Kernel 中大大减少了 CUDA Core 的开销。6.2 理想的协同工作模式在一个优化良好的 AI 系统中Tensor Core 和 CUDA Core 应该是并行工作的而不是串行工作的。理想的流水线Tensor Core 执行第 N 层的矩阵乘法同时CUDA Core 处理第 N-1 层的激活函数和归一化同时DMA 控制器加载第 N1 层的数据这样所有硬件单元都在同时工作没有空闲时间从而实现最高的整体吞吐量。挑战如何精确平衡 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 实现。误区 2Tensor 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 最伟大的创新之一。它通过专用化的硬件设计实现了数量级的性能提升为大语言模型的发展奠定了基础。核心要点回顾Tensor Core 是硬件化的矩阵乘加单元一次能完成一个完整的矩阵乘法它采用空间并行架构在一个周期内可以执行数千次运算不同精度格式的吞吐量差异很大选择合适的精度是性能优化的关键Tensor Core 和 CUDA Core 是互补关系协同工作才能实现最佳性能性能优化的核心是维度对齐、正确的数据布局、充分利用共享内存学习建议阅读 NVIDIA 的《Tensor Core 编程指南》了解 WMMA API 的基本使用学习 CUTLASS 库的源码看看它是如何优化 Tensor Core 性能的深入研究 FlashAttention 的实现学习如何融合算子和优化内存访问使用 NVIDIA Nsight Compute 工具分析 Kernel 的执行情况查看 Tensor Core 利用率和内存带宽利用率理解 Tensor Core 的工作原理是成为一名优秀 AI Infra 工程师的必经之路。它能让你透过现象看本质真正理解 AI 系统的性能瓶颈在哪里从而写出真正高效的代码。