GPU GEMM 优化完全拆解（三）：Tensor Core、MMA/wgmma 指令与异步流水线

FFMA 的天花板：SIMT 架构的算术强度瓶颈#

先从第二篇的终点往回看。kernel 10 的指令流几乎全是 FFMA——每个 SM 的 FP32 单元每周期只能执行有限次乘加。以 V100 为例：80 个 SM，每个 SM 64 个 FP32 单元，boost 频率 1.53 GHz，每个单元每周期一次 FMA（算 2 FLOP）：

15.7 TFLOPS 是 2017 年最强数据中心 GPU 的 FP32 算力。而深度学习对矩阵乘算力的胃口远超这个数——2017 年 CNN 已经靠卷积打遍视觉任务，同年 Transformer 也诞生了，两者底层都是稠密矩阵运算。问题出在哪里？

FFMA 指令的形态是”一次乘加、三个操作数”：c = a * b + c。要算一个 $16 \times 16 \times 16$ 的矩阵乘，需要 $16^3 = 4096$ 次 FMA；每次 FMA 都要从寄存器取两个操作数。硬件上，这就是”每个线程每周期能做一次乘加”，算力提升只能靠堆 SM、堆频率，而功耗墙很快顶住。更根本的问题是：GEMM 是大量可并行的乘加，但 SIMT 架构把它拆成了逐条 FMA 指令流，指令发射、操作数搬运都成了开销。

换个思路：与其让一个线程算一个输出元素（内积视角），不如让硬件直接做一个外积（outer product）——把 $A$ 的一列（16 个元素）和 $B$ 的一行（16 个元素）两两相乘，一次产生 16×16=256 个乘加结果。这 256 个乘加互不依赖、可以同时执行，而且输入只有 32 个操作数——每个操作数被复用了 16 次。这就是 Tensor Core 的出发点。

V100 的答案：2017 年的 4×4×4 矩阵单元#

2017 年 5 月，NVIDIA 在 GTC 上发布了 Volta 架构的 V100，首次引入 Tensor Core（张量核）。它的规格很简单：每个 Tensor Core 每周期完成一次 $4 \times 4 \times 4$ 的矩阵乘加（FP16 输入、FP32 累加），等价于每周期 64 次 FMA、128 FLOP。V100 每个 SM 放了 8 个 Tensor Core：

125 vs 15.7——同样的频率、同样的 SM 数，Tensor Core 把矩阵乘算力翻了 8 倍。 而且功耗几乎没增加：矩阵单元内部的操作数不需要经过指令流水线，硬连线直接相乘累加。

Tensor Core 的运算形态用一张图概括：两个 FP16 矩阵相乘，累加器可以是 FP16 或 FP32（原图见 NVIDIA 2017 年发布的 Programming Tensor Cores in CUDA 9 博客）：

A、B 是两个矩阵而非两个标量——这是 Tensor Core 与 CUDA Core 的本质区别：一条指令做整个矩阵块的乘加，而不是一个元素的乘加。

同一时期 Google 也在做同样的事，只是方向不同：TPU v1 用 15 个月从立项到流片，走的是脉冲阵列 + CISC 指令的路线；V100 则把矩阵单元嵌进 SM，保持 CUDA 编程模型不动（TPU v1 的完整故事见 TPU v1 拆解）。对用户来说，Tensor Core 最大的优势是不用学新编程模型——还是 CUDA，只是在现有 warp 里多了一类矩阵指令。cuBLAS 率先用上，官方博客给出了一张基准图：

同为 cuBLAS 库函数，P100（无 Tensor Core）与 V100（有 Tensor Core）的 FP16 GEMM 性能对比。Tensor Core 的收益完全在库内部实现，用户代码不用改一行。

五代 Tensor Core 演进总表#

从 2017 年到现在，Tensor Core 已经过了五代。下面这张表是本文后续所有讨论的地图，先通览一遍：

A100、H100、B200 数据取 SXM/数据中心版官方规格，均为无稀疏的密集算力；各代官方宣传口径常按 2:4 稀疏给 2 倍数值。T4 的 65 TFLOPS 是 FP16 密集值。

表中每一行都对应一次”解耦”的深化，后面四节逐一展开。先记住两个核心变量：指令的矩阵形态（一次算多大）和操作数从哪来（寄存器、共享内存还是专用存储）。五代演进本质上是把”矩阵乘加”这个操作逐步从通用线程模型里剥离出去，让 warp 越来越像”搬运工”而不是”计算器”。

从 wmma 到 mma.sync：两级接口#

硬件有了矩阵单元，软件怎么用？NVIDIA 给了两级接口：

第一级：wmma（Warp Matrix Multiply-Accumulate）C++ API，CUDA 9 随 Volta 引入。它把”矩阵分块、线程排布”这些细节封装起来，程序员只需要用 wmma::fragment 声明一个”片段”，用 load_matrix_sync / store_matrix_sync 把矩阵搬进搬出，再用 mma_sync 做乘加。下图是官方文档里 16×16×16 操作的示意——一个 warp 的 32 个线程合作完成 16×16 的输出：

wmma 的 16×16×16 形状：输出 D 是 16×16，A 是 16×16，B 是 16×16，K 维为 16。一个 warp（32 线程）持有全部数据，每个线程只拿其中一小片。

wmma 的问题在于”自动布局”是一把双刃剑：它隐藏了线程↔数据映射，程序员无法控制共享内存的 bank 分布，性能上限被卡死。NVIDIA 的官方建议很直白：追求性能请用底层的 mma.sync PTX 指令（在 CUTLASS 里全是这么干的）。

第二级：mma.sync PTX 指令，这是本节的解剖对象。它暴露了完整的线程级数据排布（fragment 布局）——每个线程的每个寄存器里放矩阵的哪个元素，由 PTX ISA 规范硬性规定，编译器不帮你排。理解这套布局，是手写 Tensor Core GEMM 的第一道门槛，也是理解后续 wgmma、tcgen05 布局演进的钥匙。

mma.sync.m16n8k16：一纸”座位表”#

以最常用的 mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 为例（FP16 输入、FP32 累加）：一次指令计算

由整个 warp 的 32 个线程集体执行（.aligned 要求 32 个线程同时到达）。每个线程的寄存器里持有矩阵的一小片，这一小片就叫 fragment（片段）。各 fragment 的寄存器开销是固定的：

验证一下账：A 有 256 个元素 ÷ 32 线程 = 8 个/线程 ✓；B 有 128 个元素 ÷ 32 = 4 个/线程 ✓；C/D 有 128 个元素 ÷ 32 = 4 个/线程 ✓。每线程总计 10 个寄存器装数据，换来 16×8×16 = 2048 次乘加——每个”数据寄存器”摊到约 205 次 FMA，而 kernel 3 的每个数据寄存器只摊到 1 次。这就是 Tensor Core 的算术强度来源，也是第二篇”寄存器分块”思想的硬件终极形态。

元素到线程的映射由两个量决定：

32 个线程被分成 8 组 × 4 线程。以累加器 C/D 为例，线程 (groupID, threadID_in_group) 持有 4 个元素，位置是：

• 第 groupID 行、第 2t 和 2t+1 列；
• 第 groupID + 8 行、第 2t 和 2t+1 列。

也就是说：行方向上 8 组线程各占一行（以及对称的 +8 行），列方向上每组 4 个线程把 8 列切成 4 段、每段 2 列。A、B 的排布类似（A 的 16 行按”上下两半”处理，B 的 8 行按 2×4 分组），细节以 PTX ISA 的 Matrix Fragments 一节 为准。

这个”8 组 × 4 线程”的结构从哪来？硬件矩阵单元的内部组织是 16 行 × 8 列的输出，配上足够宽的乘加阵列；32 个线程、每个线程 4 个元素，正好把 16×8 的输出摊平。而 A 用 half2 打包（一个 32 位寄存器装两个 FP16）是为了让”线程到元素”的映射和 32 位寄存器粒度对齐——mma.sync 的 A 操作数寄存器类别是 .b32 而不是 .f16，这是新手最容易踩的坑：拿 half 类型的寄存器去喂 mma.sync 是编译不过的。

ldmatrix：为 fragment 定制的搬运指令#

布局是死的，怎么把共享内存里的数据搬进这些指定位置？逐元素 ld.shared 当然可以，但 32 个线程各自算地址、各自访问，既慢又容易踩 bank conflict。Turing 开始硬件提供了一条专用指令：ldmatrix.sync.aligned.m8n8.xN.shared.b16——一次让一个 warp 从共享内存加载若干个 8×8 的 FP16 矩阵，硬件自动把元素分发到 32 个线程的指定寄存器位置（.x1/.x2/.x4 表示一次加载 1/2/4 个 8×8 矩阵，.trans 表示同时转置）。

这条指令的意义有两个：一是把”布局转换”从软件搬进硬件——ldmatrix 内部按无 bank conflict 的方式组织访问（配合共享内存的 swizzle 布局，这正是第一篇预告的”现代 GEMM 普遍用 swizzle 方案”的落点）；二是一条指令代替几十条 LDS，指令发射开销大减。m16n8k16 的 A 是 16×16，正好两个 8×8 块；B 是 16×8，也是两个 8×8 块（取其中 8 列）。所以一次 mma 的典型前置是：24 条 ldmatrix + 01 条 mma.sync。

组装一个 Tensor Core GEMM kernel#

把前几节拼起来，一个最小可用的 Tensor Core GEMM 主循环长这样（忽略边界、假设尺寸对齐）：

1
#include <mma.h>
2
using namespace nvcuda;
3

4
// 每个 warp 算一个 16×16 的输出块（2 个 mma.m16n8k16 并排）
5
// 每线程持有 2×4=8 个 FP32 累加器
6
wmma::fragment<wmma::accumulator, 16, 16, 16, float> acc;
7
wmma::fragment<wmma::matrix_a,    16, 16, 16, half> a_frag;
8
wmma::fragment<wmma::matrix_b,    16, 16, 16, half> b_frag;
9
wmma::fill_fragment(acc, 0.0f);
10

11
for (int k = 0; k < K; k += 16) {
12
  // 从共享内存装载 fragment（内部就是 ldmatrix 序列）
13
  wmma::load_matrix_sync(a_frag, As + offsetA, lda);   // A 的 16×16 子块
14
  wmma::load_matrix_sync(b_frag, Bs + offsetB, ldb);   // B 的 16×16 子块
15
  // 一次乘加：D = A×B + C
16
  wmma::mma_sync(acc, a_frag, b_frag, acc);
17
}
18
// epilogue：把累加器写回全局内存
19
wmma::store_matrix_sync(C + offsetC, acc, ldc, wmma::mem_row_major);

用 mma.sync PTX 手写时，load_matrix_sync 换成 ldmatrix + 显式 fragment 变量，逻辑一致。注意这个循环结构和第二篇 kernel 10 的”外层 K 循环 + 内层 mma 序列”完全同构——warp tiling 时我们手动把 64×64 的 warp 分块切成 8×4 的线程小块，那个 8×4 小块，现在就是由硬件 mma 指令直接承接的 fragment。kernel 10 的结构是为 Tensor Core 量身预热的：warp 分块映射到 warp 级 mma 序列，共享内存的 swizzle 布局为 ldmatrix 服务，双缓冲流水线为异步搬运服务。

数值精度：为什么 FP32 的 cuBLAS 不用 Tensor Core#

这里有个必须澄清的细节：mma.sync 的输入精度是 FP16/BF16/TF32/INT8 等，没有原生 FP32 输入（V100 的 4×4×4 单元只有 FP16 输入）。所以 A6000 上 cuBLAS 的 FP32 SGEMM 全程用 FFMA——不是 NVIDIA 偷懒，而是 FP32 输入没有对应的矩阵指令（第二篇的结论”cuBLAS 的 FP32 也没用 Tensor Core”在此得到解释）。

FP32 用户怎么办？三条路：

1. TF32（Ampere 起）：把 FP32 操作数截断成 19 位（8 位指数 + 10 位尾数，比 FP16 的 10 位尾数少 1 位但指数范围大），Tensor Core 原生支持，A100 上 156 TFLOPS。代价是精度损失：单次 TF32 乘法的相对误差约 $2^{-11}$，明显大于 FP32 的 $2^{-24}$；
2. 3×TF32（cuBLAS 提供）：把 FP32 拆成”高位 TF32 + 低位 TF32”两部分，三次 TF32 乘法叠加，精度接近 FP32。代价是约 3 倍的 TF32 指令量，实测速度大致为纯 FFMA 的 1.5~2.5 倍（取决于拆分与融合的开销）。这是 Ampere 上”FP32 GEMM 加速”的主流工程方案；
3. BF16×9 模拟（Hopper/Blackwell 的 cuBLAS 与 CUTLASS）：用 9 次 BF16 乘法模拟一次 FP32 乘法，B200 上 FP32 GEMM 也能吃到 Tensor Core 红利。

推理场景更干脆：模型权重直接量化到 FP16/INT8/FP4，用低精度换 4~16 倍张量算力，再用量化感知训练或校准补偿精度损失——这是 XFP 码本量化 那一类工作的前提。所以”GEMM 快不快”在现代 GPU 上，实际是”低精度 GEMM 快不快”。

问题重述：算力涨了 8 倍，带宽没跟上#

Tensor Core 把算力翻了 8 倍，但显存带宽没翻 8 倍：V100 的 HBM2 带宽 900 GB/s，H100 的 HBM3 带宽 3.35 TB/s——五代之间算力涨了 72 倍（125 → 9000 TFLOPS），带宽只涨了 8.9 倍（0.9 → 8 TB/s）。结果是：GEMM 的算术强度要求越来越高，数据搬运越来越成为瓶颈。

前两篇的解法是”占用率隐藏延迟”：让很多 warp 排队，一个 warp 等数据时换另一个上。但 Tensor Core 时代这招失效了——因为 Tensor Core 的指令是 warp 级的，一个 warp 发出一条 mma 指令后，整个 warp 的 32 个线程都在等这一个矩阵乘的结果；等数据时，这个 warp 占着的寄存器、调度槽位全在空转。更糟的是，GEMM 的每个 warp 都在等同样的事：加载下一块 tile。当所有 warp 一起等 HBM 时，占用率再高也没用。

方向变了：不再靠”多 warp 轮换”隐藏延迟，而是让”搬运”和”计算”真正同时发生——软件流水线（software pipelining）。Colfax 的 CUTLASS 教程把这个目标概括为一句很形象的话：feeding the beast（喂饱这头野兽），野兽就是 Tensor Core，它每周期能吞 2048 次乘加，绝不能让它在等数据。

cp.async：Ampere 的异步拷贝指令#

流水线的地基是异步拷贝。Ampere（sm_80）起，PTX 提供了 cp.async 指令：从全局内存直接拷贝到共享内存，完全绕过寄存器和指令流水线，由 SM 上的异步拷贝单元（MIO 的一部分）执行。线程发起拷贝后立即继续执行，不等待数据到达。

1
// 每个线程发起一次 16 字节的异步拷贝：GMEM → SMEM
2
__pipeline_memcpy_async(&As[stage][idx], &A[offset], 16);
3
// 把"本阶段"的拷贝登记为一个组
4
__pipeline_commit();
5
// 等到只剩 N 组拷贝未完成（N=2 表示允许 2 个阶段在途）
6
__pipeline_wait_prior(2);
7
__syncthreads();  // 保证本阶段数据就绪后再计算

三段式 API 对应硬件机制：cp.async 拷贝由硬件异步执行；commit 把当前批次的拷贝登记成一个组（group）；wait_prior(N) 等待最多 N 个组还在路上。程序员用”组”管理在途拷贝的数量，从而控制流水线的深度。

配合**多缓冲（multi-buffering）**就得到了流水线。共享内存里给 A、B 各分配 $S$ 份缓冲（$S$ 叫 stage 数，CUTLASS 里就是 num_stages）：

• $S=1$（无缓冲）：加载 tile → 同步 → 计算 → 加载下一块。加载和计算完全串行，总时间 $T = T_{\text{load}} + T_{\text{compute}}$；
• $S=2$（双缓冲）：计算第 $i$ 块的同时，异步加载第 $i+1$ 块。理想情况下加载完全被计算掩盖，$T \approx \max(T_{\text{load}}, T_{\text{compute}})$；
• $S=3$（三级流水）：再多一级缓冲，让加载的波动也有地方缓冲。

一般化地，$S$ 级流水在稳态下，每个阶段暴露的加载延迟大约是总加载时间的 $1/S$（只有”轮到自己消费的那一级”是必须等待的，其余 $S-1$ 级都在后台预取），所以直觉模型：

其中 $T_{\text{compute}}$ 是计算一块 tile 的时间，$T_{\text{load}}$ 是加载一块 tile 的时间，$S$ 是流水线级数。从这个公式能读出三条设计准则：$T_{\text{load}}$ 远大于 $T_{\text{compute}}$ 时（算力过剩、带宽紧张），增大 $S$ 显著有效；$S$ 从 1 变 2 收益最大；$S$ 不能无限大——每级缓冲都要占共享内存（约束：(BM×BK + BN×BK) × 4 字节 × S ≤ 共享内存上限，H100 上约 228 KB），stage 太多会挤占共享内存导致编译失败或占用率崩掉。Ampere 时代 cp.async 的多级流水是最优解；但它有个弱点——每个线程都要自己算地址、自己发拷贝指令，32 个线程 × 每块 tile 的地址计算和指令发射开销仍然不小，而且同步用的是软件轮询。

TMA：Hopper 的批量搬运引擎#

Hopper（sm_90）把异步搬运升级成了硬件引擎：TMA（Tensor Memory Accelerator，张量内存加速器）。TMA 是每个 SM 里的独立 DMA 单元，专门做”全局内存 ↔ 共享内存”的多维张量搬运。

用法和 cp.async 完全不同：整个搬运只需要一个线程发起。先在主机端用 cuTensorMapEncodeTiled 把搬运的”地图”编码成一个 tensor map 描述符（包含基地址、各维 stride、tile 形状、swizzle 模式），然后任意一个线程执行一条 cp.async.bulk.tensor 指令，TMA 引擎自己完成全部地址计算和搬运。它的优势：

• 地址计算和指令发射从 32 个线程减到 1 个线程 1 条指令——warp 完全从搬运工作中解放出来；
• TMA 支持 2D/3D/4D tile 搬运、自动 swizzle（搬运时就把共享内存布局重排好，配合 ldmatrix/wgmma 无冲突读取）、集群级 multicast（一个 block 簇内的多个 SM 共享同一次 HBM 读取）；
• 完成信号由硬件直接写入mbarrier（内存屏障）：发起前用 expect_tx(字节数) 告诉屏障”将要有这么多字节写入”，TMA 每写完一部分，硬件自动 arrive_tx，字节数到齐屏障自动翻转相位。生产者（TMA）和消费者（计算 warp）通过屏障握手，全程不需要软件轮询。

CUTLASS 的 PipelineTmaAsync 里，producer 的”提交”动作甚至是个 no-op——TMA 拷贝完成本身就会触发屏障。下面的图是 CUTLASS 官方文档里的软件流水线示意，展示了一个 3 级流水主循环（prologue → mainloop → epilogue）里，TMA 加载、mma 计算、屏障等待如何交替：

图上方是流水线的三个阶段（prologue 预取、mainloop 计算、epilogue 写回），下方是共享内存里多个 stage 缓冲的轮转。每个计算步骤都”超前”于数据就绪状态——这就是”让计算永远不等加载”的含义。

TMA 还有一个反直觉的事实：TMA 拷贝本身并不比 cp.async 快——它的收益完全来自”单线程发起 + 硬件调度 + 自动屏障”，把 warp 的发射带宽和寄存器让给了计算。Colfax 的教程测过：紧跟着等待 TMA 拷贝的 kernel 性能只有 0.9~1.0×（和 cp.async 无差别），而正确流水化后 2 级流水在低算术强度负载上能到 1.7×。异步的价值不在指令，在重叠。

warp specialization：生产者-消费者分工#

有了 TMA 和 wgmma（下一节），Hopper 上出现了更激进的组织方式：warp specialization（warp 特化）。把 block 里的 warp 分成两类，各自专职：

• producer warpgroup（生产者）：只发 TMA 拷贝，把数据从 HBM 灌进共享内存；
• consumer warpgroup（消费者）：只做 mma/wgmma 计算和 epilogue。

两类 warp 通过 mbarrier 握手，互不等候。和”所有 warp 既搬运又计算”的多级流水相比，warp specialization 是空间分工（不同 warp 干不同活）而非时间重叠（同一 warp 轮转干活）。它在新硬件上成立的条件：TMA 让生产者几乎不占寄存器；wgmma 让消费者直接从共享内存读操作数；再加上 Hopper 的 setmaxnreg 指令——硬件支持按 warpgroup 动态重新分配寄存器，生产者 warpgroup 只分 2440 个寄存器，消费者 warpgroup 分到 240256 个。没有这个能力，消费者端的 FP32 累加器会因寄存器不足而灾难性溢出（Colfax 实测：不做寄存器再分配时，FP32 累加 kernel 性能从 ~500 TFLOPS 掉到 ~21 TFLOPS）。

FlashAttention-3（FA3）是这套技术的教科书案例。2024 年 7 月，Colfax 与 Tri Dao 团队把 FA3 论文 发布在 arXiv 上，用生产者-消费者 + 乒乓调度重写了注意力 kernel：producer warpgroup 用 TMA 预取 Q、K、V，consumer warpgroup 用 wgmma 做 $QK^{\top}$ 与 $PV$ 两个矩阵乘。FA2 在 H100 上只能达到 35% 的利用率（它的 kernel 是为 Ampere 的同步模型设计的），FA3 用异步流水线把它推到 FP16 下最高 740 TFLOPS、75% 利用率，FP8 下接近 1.2 PFLOPS——对 FA2 是 1.5~2.0 倍加速。

下图是论文 Figure 2：2-stage WGMMA-softmax 流水——把一个 warpgroup 内部的 $QK^{\top}$ 矩阵乘和 softmax 重叠起来：

蓝色块是 wgmma（矩阵乘），橙色块是 softmax 相关计算。2 级流水把”矩阵乘 → softmax → 矩阵乘”的串行链拆成两条交替的流水线。

FA3 论文里还给出了 3 级流水（把 softmax 再拆成指数/归一化两步）与 ping-pong 调度（两个 consumer warpgroup 轮流消费两个独立的 $QK^{\top}/PV$ 链）。ping-pong 的思想值得单独看图——它用两个独立计算链互相掩盖对方的启动与收尾：

乒乓调度的核心：当一条链在等待屏障/同步时，另一条链的 wgmma 还在满负荷运转。图上方和下方的两条时间线交替覆盖彼此的间隙。

FA3 论文对 2 级、3 级流水与 ping-pong 调度都做了消融测量（论文 Table 2）——3 级和 ping-pong 相对 2 级还有提升，但收益明显递减，印证了流水线的普遍规律：第一级异步（从无到 2 级流水）吃掉绝大部分收益，后面的工程是锦上添花。

wgmma：warpgroup 级的异步矩阵乘#

mma.sync 是 warp（32 线程）级、同步的：发出一条指令，整个 warp 被矩阵单元占用，等结果写回寄存器才能干别的。Hopper 的 wgmma.mma_async（WGMMA，warpgroup MMA）把粒度放大到 warpgroup（4 个 warp = 128 线程），并且是异步的：

• 指令形状固定 M=64：wgmma.m64nNk16，FP16/BF16 时 K=16，N 可取 8 到 256 的多种值（常用 64/128/256）；FP8 时 K=32；
• 操作数直接从共享内存读取，不再经过寄存器。A、B 的地址用一个 64 位的 **matrix descriptor（矩阵描述符）**描述：基地址、leading dimension 字节偏移（行间步长）、stride 字节偏移（列间步长）、swizzle 模式。wgmma 有 SS 形式（A、B 都在共享内存）和 RS 形式（A 在寄存器、B 在共享内存，用于广播类操作）；
• 指令发出后立即返回，warpgroup 可以继续做别的；完成状态由 wgmma.commit_group / wgmma.wait_group 配合 mbarrier 追踪；
• 累加器仍在寄存器：m64n128k16 的累加器 64×128 = 8192 个 FP32，由 128 个线程每人持有 64 个（m64n256 时每人 128 个）——排布方式和 mma.sync 的”8 组 × 4 线程”同构，只是行、列规模放大。FA3 论文 Figure 3 专门画了这张累加器座位表（行 0/8 与线程 0-3 的对应关系）。

wgmma 的意义：mma.sync 时代”共享内存 → ldmatrix → 寄存器 → 矩阵单元”的数据链被砍成”共享内存 → 矩阵单元”——寄存器不再充当操作数的中转站，128 个线程的寄存器预算几乎全部让给累加器。配合 TMA 的异步搬运，Hopper 的 GEMM 主循环可以做到：生产者 warpgroup 发 TMA 灌数据，消费者 warpgroup 发 wgmma 算矩阵，两边各干各的，只有 mbarrier 上的一点点握手。

tcgen05：Blackwell 把累加器也搬走了#

Blackwell（sm_100）的 tcgen05.mma（CUTLASS 里叫 UMMA，Ultra-MMA）把解耦推到终点，同时解决了 wgmma 的最后两个牵绊：

1. 累加器从寄存器搬进 TMEM。Blackwell 为矩阵单元新增了专用片上存储 TMEM（Tensor Memory）：每 SM 256 KB（128 列 × 512 行 × 32 位），只能由矩阵单元读写。tcgen05.mma 的累加器（C/D）就在 TMEM 里，算完用 tcgen05.ld 把结果拷回寄存器做 epilogue。寄存器压力被彻底清零——SemiAnalysis 的分析指出，TMEM 之所以值得单独建一块存储，是因为累加器 D 的访问频率比操作数 A、B 高得多，把高频访问的 D 放在能效更高的专用存储里，比放在通用寄存器文件里更划算；
2. 单线程发起（single-thread issue）。tcgen05.mma 是标量指令：一个线程发出一条指令，整个矩阵乘由硬件独立执行，不占用 warp 的发射槽，也不要求 128 个线程同步到达。执行期间这个线程可以继续干别的——矩阵单元像一台挂载在 SM 上的独立 DMA 引擎。配置通过指令描述符（idesc）携带：数据类型、稀疏模式、转置、符号翻转等元数据一次打包。

tcgen05 的能力上限：单次指令最大 m256n256k16（FP16 时 K=16，FP8/FP4 时 K=32），是 wgmma 最大原子（m64n256）的两倍以上；支持原生 FP4/FP6（配合 MX 块缩放，scale factor 也存在 TMEM 里）；两个相邻 CTA 可以组成 CTA-pair 协作完成同一个矩阵乘。代价是兼容性断裂：wgmma 在 Blackwell 上被标记为 deprecated（仍可运行），但 NVIDIA 的工程建议是——MX-FP4 负载上继续用 wgmma 会留下 30~40% 的峰值性能在桌上。

三代指令的对比，一张表收束：

演进主线：指令粒度从”32 线程集体执行”变成”1 个线程发起”，操作数从”寄存器中转”变成”直读共享内存”再到”累加器搬进专用存储”。每一次演进，通用线程离计算本身都更远一步——warp 越来越像给矩阵单元喂料的搬运工。

对 kernel 设计的影响：Blackwell 的 FA4 为什么长这样#

这套演进不是抽象概念——本系列之前拆解的 FlashAttention 完全拆解（四）（FA-4，2025 年）就是深度使用 tcgen05 的生产 kernel：它的主循环里，consumer 端是”一个线程发 tcgen05.mma + tcgen05.ld”，producer 端是 TMA 灌数据，两个 warpgroup 之间的同步完全靠 mbarrier。对比 Ampere 时代的 FA2 kernel（mma.sync + cp.async + __syncthreads），代码形态判若两个世界——这不是风格差异，是指令集把软件能表达的最高性能形态整体抬高了。

从 kernel 到库：层次结构#

前面几节的手写 kernel 把原理讲清楚了，但生产级 GEMM 的复杂度远不止于此：要支持任意矩阵尺寸、任意数据类型、任意 GPU 世代，还要在每种组合下都接近峰值。NVIDIA 的开源答案是 CUTLASS——一个模板化的线性代数库，2017 年随 Volta 发布，官方博客 介绍了它的设计初衷：把高性能 GEMM 的工程经验组织成可复用、可组合的组件（它与 cuBLAS 遵循同一套设计原则，被 NVIDIA 内部和社区广泛用于构建各类高性能 kernel）。

CUTLASS 对 GEMM 的建模是三层嵌套分块，efficient GEMM 文档 给出了权威的循环结构（略去循环体）：

1
for (int cta_n ...) { for (int cta_m ...) {        // threadblock 级：GMEM → SMEM
2
  for (int cta_k ...) {                            // GEMM mainloop，每轮一个 stage
3
    for (int warp_n ...) { for (int warp_m ...) {  // warp 级：SMEM → 寄存器
4
      for (int warp_k ...) {                       // 完全展开，一个 k-group
5
        for (int mma_k ...) { for (int mma_n ...) { for (int mma_m ...) {
6
          mma_instruction(d, a, b, c);             // 指令级：Tensor Core
7
        } }
8
      }
9
    } }
10
  }
11
} }

每一层循环对应硬件的一个并行层次和一个存储层级：threadblock 层（不同 block 跑在不同 SM 上，操作全局内存）、warp 层（同一 SM 内不同 warp 跑在不同调度器上，操作共享内存）、mma 层（同一 warp 内连续发射的指令，操作寄存器/TMEM）。这个结构和第二篇 kernel 10 的三级分块一一对应——CUTLASS 只是把它模板化、参数化了。下图是文档中的全景图：

从左到右依次是 threadblock tile（加载进共享内存）、warp tile（加载进寄存器）、mma 指令（喂给 Tensor Core），右侧是 epilogue（写回全局内存）。每一层平铺都对应一个可以独立调优的模板参数。

mainloop 与 epilogue：把”算”和”写”分开#

CUTLASS 把 kernel 分成两个阶段。mainloop 只做矩阵乘累加（上文的三层循环），输出留在寄存器的累加器里——注意，累加器里 C 元素的排布是为矩阵乘最优的，直接写回全局内存会完全不合并（这个坑在第二篇 kernel 10 的写回里已经埋过）。所以有 epilogue 阶段：先让线程通过共享内存把累加器数据转置重排成全局内存友好的连续布局，再按合并访问写回；顺便在写回前把 $\alpha C + \beta C$、ReLU、量化缩放这类逐元素操作一并做掉（FP16 累加器的转置重排 + dequant 是量化推理 kernel 的 epilogue 标配）。

双缓冲的两个 scope、rasterization 与 Split-K#

CUTLASS 的流水线在两个 scope 同时双缓冲（文档原话）：

• threadblock 级：共享内存里分配两份（或多份）tile，一份供当前 mainloop 迭代计算，一份接收下一轮从全局内存异步加载的数据（cp.async/TMA 完成）；
• warp 级：寄存器里分配两份 fragment，一份正在喂给 Tensor Core，一份接收下一轮从共享内存装载的结果（ldmatrix 完成）。

两级缓冲合起来，就是前文”软件流水线”图里的完整结构——加载、装载、计算、写回四个动作在各存储层级上流水推进。

除此之外，生产级 GEMM 还要处理三个前面没遇到的设计点：

Threadblock rasterization（栅格化）：block 到输出 tile 的映射不是简单的行优先扫描，而是把连续启动的 block 映射到二维邻近区域，让同时运行的 block 尽量访问全局内存中相邻的 tile，提高 L2 命中率（第二篇 kernel 10 的 thread swizzling 负结果在 CUTLASS 里是正式功能——因为 CUTLASS 的 block 调度是大规模 grid，L2 压力比单 kernel 实验大得多）。

Split-K 与 Sliced-K：矩阵太小时，单 block 的计算量喂不饱整张 GPU。Split-K 把 K 维度切给多个 block 并行算部分和，再用第二个归约 kernel 合并（第二篇看到 cuBLAS 对 256 尺寸就是这么干的）；Sliced-K 是在 block 内部把 K 切片给多个 warp，代价是最后要一次 warp 间归约。

Hopper 时代的三种 kernel 设计#

CUTLASS 3.0 为 Hopper 提供了三种 warp-specialized kernel 设计，是”生产者-消费者”思想在库里的完整形态：

1. Warp-Specialized（WS）：一个 producer warpgroup 专职 TMA 搬运，一个/两个 consumer warpgroup 专职 wgmma，通过 mbarrier 的空/满屏障握手；
2. Warp-Specialized Persistent Cooperative：block 数量 = SM 数（persistent，常驻），每个 block 通过 Tile Scheduler 领取多个输出 tile，摊薄启动开销；两个 consumer warpgroup 合作切分同一个大 tile 的 M 维，降低各自寄存器压力；
3. Warp-Specialized Persistent Ping-Pong：在合作式基础上让两个 consumer warpgroup 各领一个输出 tile，一个做 epilogue 时另一个做矩阵乘——epilogue 的时间被另一条链的计算掩盖。

第三种正是 FA3 的 ping-pong 思想在 GEMM 库中的版本。CUTLASS 3.x 的工程底座是自研的 CuTe 模板库（类型化的 tile 与布局系统），用户通过 CollectiveBuilder 声明式组合主循环（选 MMA 指令、流水线类型、共享内存 swizzle），其余由库生成——这也是 Triton 等编译器后端的对标对象。

性能：CUTLASS 的 Hopper 水平#

Colfax 的 CUTLASS 教程 给了一组 H100 上的实测，可以作为”这套工程做到什么水平”的标尺（H100 SXM5，FP16 GEMM，峰值 989.5 TFLOPS）：

多级流水和 warp specialization 都能把教学级实现推到 ~65% 峰值——第一级异步吃掉了大部分收益。注意这是 Colfax 教程的简化框架下测的，经过深度调优的生产级 kernel（CUTLASS 社区最优实现、cuBLAS 的高性能路径）可以达到更高。

对比第二篇的 A6000 数据（FP32 FFMA 路线 cuBLAS 93.7%），会发现一个微妙的反差：FFMA 路线能逼近峰值 90%+，Tensor Core 路线的常见实现却停在 65~75%。原因在于两种路线的”峰值”定义不同：FFMA 的峰值只算纯计算，而 Tensor Core 的峰值假设矩阵单元满负荷——喂满它需要的搬运带宽远超 HBM 实际供给，所以 65% 的”缺口”里很大一部分是数据供给的天花板，而不是实现不行。这也是为什么 H100 的 FP16 GEMM 实测算术强度远低于理论值，以及为什么 FlashAttention 这类 kernel 拼命把 HBM 流量降到最低。