Tensor Core 与 MMA：wmma、mma.sync 与数据布局

GEMM 篇 把手写 FP32 GEMM 优化到了 6.4 TFLOP/s，约 FP32 峰值的 39%。但深度学习训练和推理的算力其实大部分不跑在 FP32 CUDA core 上，而跑在 Tensor Core 这种专用矩阵单元上。这一篇讲 Tensor Core 是什么、怎么用、有什么约束，以及为什么用了它常常还是快不起来。

一、Tensor Core：把矩阵乘做进硬件

普通 CUDA core 一次做一个标量 FMA。Tensor Core 不一样——它一条指令直接算一个小矩阵的乘加：

\[ D = A \times B + C \]

其中 \(A\)、\(B\)、\(C\)、\(D\) 是小矩阵（如 \(16\times16\)）。这条指令叫 MMA（Matrix Multiply-Accumulate），由一个 warp 的 32 个线程协作执行，硬件在内部并行完成所有乘加。把 \(O(n^3)\) 的乘加打包进专用电路，单位时间的吞吐远高于标量 FMA。

Tensor Core 从 Volta（第一代）引入，本系列测试卡 RTX 3060 Ti 是 Ampere 消费级（GA10x），第三代 Tensor Core，每个 SM 4 个，共 152 个。它支持 FP16、BF16、TF32 输入，FP32 累加；不支持 FP8（FP8 Tensor Core 从 Hopper sm_90 才有，见 第 21 篇）。

二、实测：Tensor Core 比 FP32 快多少

用一个微基准量化：每个 warp 反复执行 wmma 的 \(16\times16\times16\) MMA（FP16 输入、FP32 累加），用多个独立累加器填满流水线，统计达到的浮点吞吐。RTX 3060 Ti 实测：

Tensor Core 的 FP16 吞吐达 72.8 TFLOP/s，是 FP32 峰值的约 4.5 倍、是手写 FP32 GEMM 的约 11 倍。这就是为什么混合精度训练和低精度推理是标配——不是为了省显存那么简单，而是 Tensor Core 的算力本来就远超 FP32 通路。

需要说明：这是消费级 Ampere（GA10x）的数字。数据中心卡（A100/H100）的 Tensor Core 吞吐更高；消费级 GA10x 的 FP16 输入 / FP32 累加这条通路相比数据中心 GA100 被削减（FP32 累加速率受限），但这只是「跨卡对比」的相对结论——本卡实测的 72.8 TFLOP/s 正是在这条受限通路上跑出来的真实值，二者并不矛盾，一个是本卡实测、一个是与 GA100 的横向对比。不同代际、不同精度（FP16/BF16/TF32）的峰值差异很大，做 Roofline 分析时要用对应精度的峰值。

三、wmma API：fragment 与三步操作

CUDA 提供 nvcuda::wmma（warp matrix multiply-accumulate）这套 C++ API 来用 Tensor Core。它的核心抽象是 fragment——一个 warp 协作持有的小矩阵分片，数据分布在 32 个线程的寄存器里，布局对程序员不透明（由硬件决定）。

用法三步：load → mma → store。

#include <mma.h>
using namespace nvcuda;

wmma::fragment<wmma::matrix_a, 16,16,16, half, wmma::row_major> a;
wmma::fragment<wmma::matrix_b, 16,16,16, half, wmma::col_major> b;
wmma::fragment<wmma::accumulator, 16,16,16, float> c;

wmma::fill_fragment(c, 0.0f);
wmma::load_matrix_sync(a, ptrA, lda);   // 从 global/shared 加载分片
wmma::load_matrix_sync(b, ptrB, ldb);
wmma::mma_sync(c, a, b, c);             // c += a * b，整个 warp 协作
wmma::store_matrix_sync(ptrC, c, ldc, wmma::mem_row_major);

几个要点：

• 模板参数 16,16,16 是 MMA 的形状 \(M,N,K\)（Ampere FP16 支持 16×16×16 等几种）。
• matrix_a/matrix_b 要指定行主序还是列主序；累加器是 FP32。
• 所有 *_sync 操作是 warp 级的，warp 内 32 线程必须一起调用，不能在发散的分支里调。
• fragment 的内部布局不透明，所以不能手动索引 fragment 元素，只能整体操作。

四、wmma 之下：mma.sync PTX

wmma 是较高层的封装，简单但不够灵活——fragment 布局不透明，限制了和 shared memory 的精细配合。高性能库（CUTLASS）直接用更底层的 mma.sync PTX 指令，它暴露明确的线程-数据映射（哪个 lane 持有矩阵的哪些元素），让库能精确控制数据如何从 shared 进入 fragment，配合 swizzle 布局消除 bank conflict。

mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32  {d0,d1,d2,d3}, {a0,a1,a2,a3}, {b0,b1}, {c0,c1,c2,c3};

手写 mma.sync 要管理寄存器到矩阵元素的映射，非常繁琐，所以一般不直接手写——要么用 wmma，要么用 CUTLASS 的 CuTe 抽象（第 12 篇）来管理这些布局。

五、真正的瓶颈：喂数据

一个反复出现的现象：换上 Tensor Core，性能却没涨多少。原因几乎总是数据喂不上。Tensor Core 算得太快，如果数据从 global/shared 进入 fragment 的速度跟不上，它就大量空转。

回到 Roofline：Tensor Core 把算力屋顶大幅抬高（本卡 FP16 约 73 TFLOP/s vs FP32 16 TFLOP/s），脊点的算术强度要求随之提高——原本 compute-bound 的 GEMM 在 Tensor Core 下可能重新变成 memory-bound。所以用 Tensor Core 的 GEMM 必须把访存做到极致：

• shared memory staging：把 tile 搬进 shared，从 shared 喂 fragment，避免反复打 global。
• 异步拷贝：Ampere 引入 cp.async，让数据从 global 直接异步搬进 shared、不占寄存器、不阻塞计算，实现 load 和 mma 的重叠（通信重叠篇 会再提）。
• swizzle 布局：精心安排 shared 里的数据布局，让 fragment 加载无 bank conflict。
• double buffering：算当前 tile 时预取下一个 tile。

这些正是 CUTLASS 在做的事——把 Tensor Core 喂饱，让它逼近峰值。手写一个能跑满 Tensor Core 的 GEMM 工作量极大，所以下一篇转向 CUTLASS。

六、什么时候直接用、什么时候交给库

• 标准 GEMM/卷积：直接用 cuBLAS/cuDNN，它们的 Tensor Core 实现接近峰值，没有手写的理由。
• 自定义融合算子里的矩阵乘（如 FlashAttention 内部的 \(QK^\top\) 和 \(PV\)）：需要把 MMA 嵌进自己的 kernel，这时用 CUTLASS 的 collective/CuTe，或在能接受的性能下用 wmma。
• 学习与原型：wmma 是理解 Tensor Core 编程模型的最佳入口，但别指望它达到库的性能。

七、小结与下一步

• Tensor Core 用一条 MMA 指令算一个小矩阵乘加，由 warp 协作执行，吞吐远超标量 FMA。
• 本卡实测 FP16 Tensor 吞吐 72.8 TFLOP/s，约 FP32 峰值 4.5 倍、手写 FP32 GEMM 的 11 倍；支持 FP16/BF16/TF32，不支持 FP8。
• wmma 提供 fragment 抽象（load/mma/store，warp 级），底层是 mma.sync PTX；高性能库用后者精细控制布局。
• 用上 Tensor Core 后瓶颈常转回访存，必须靠 shared staging、cp.async、swizzle、double buffering 把它喂饱。

把 Tensor Core 喂到接近峰值是一项系统工程。下一篇看工业界的答案——CUTLASS 与 CuTe：模板化 GEMM 与布局代数。