Tensor Core 与新指令集：从 WMMA 到 wgmma、TMA、tcgen05

① 先看全貌：每一代都在把活从线程手里拿走

为什么要专门的矩阵指令：C4 的 v2 里，一个线程每条 FMA 指令做 1 次乘加。Ampere 的 mma.sync.m16n8k16 是一个 warp 一起发的一条指令，一次做 16 × 8 × 16 = 2048 次乘加。指令数少了两个数量级，取指、调度、寄存器读写的开销都随之摊薄，这就是 Tensor Core 快的根本原因。本仓库的实测：RTX 5090 上 FP32 CUDA core 的 cuBLAS 是 65 TFLOP/s（C4），FP16 Tensor Core 的 Triton GEMM 是 218 TFLOP/s（T4）。

但 Tensor Core 算得越快，喂数据就越成问题。每一代新增的指令，几乎都是为了让数据跟上：

② 第一步：WMMA，用 Tensor Core 改写 C4

nvcuda::wmma（#include <mma.h>）是 CUDA 9 起的 C++ 接口，最容易上手。它的核心是 fragment：一个 warp 共同持有的一小块矩阵，具体每个线程拿哪几个元素是不公开的（opaque）。

cuda-samples 里最简单的版本 cudasamplessimple_wmma_gemm，核心就是这几行：

wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> a_frag;
wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::col_major> b_frag;
wmma::fragment<wmma::accumulator, 16, 16, 16, float> acc_frag;
wmma::fill_fragment(acc_frag, 0.0f);
for (int k = 0; k < K; k += 16) {
  wmma::load_matrix_sync(a_frag, a + k + aRow * lda, lda);
  wmma::load_matrix_sync(b_frag, b + k + bCol * ldb, ldb);
  wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);
}
wmma::store_matrix_sync(c + cCol + cRow * ldc, acc_frag, ldc, wmma::mem_row_major);

改写 C4 的路线（面试题里的第一步）：

1. 先照抄上面这个版本：直接从全局内存 load fragment，没有 shared memory。它能工作，但每个 fragment 都从显存读，复用很差。
2. 把 C4 v1 的 shared memory tiling 加回来：一个 block 把 A、B 的 tile 读进 shared memory，多个 warp 各自从 shared memory load fragment。cuda-samples 的完整版 cudasamplescompute_gemm 就是这样，cudasamples每个 warp 持有多个 a、b fragment，相当于 C4 v2 的“寄存器 tile”。
3. 和 cuBLAS 的 FP16 GEMM 对比。参考上限：本仓库 T4 在 5090 上实测 cuBLAS FP16 约 218 TFLOP/s。

WMMA 的局限：fragment 的寄存器布局不公开，所以你没法自己决定每个线程持有哪些元素，也就没法和 softmax、量化之类的计算精细融合；shared memory 到 fragment 的加载方式也是固定的。这正是下一代写法（直接用 mma.sync PTX，或者用 CuTe）要解决的问题。文档：Programming Guide · WMMA。

③ Ampere：mma.sync、ldmatrix、cp.async，以及 CuTe 的 MMA Atom

三条新指令

MMA Atom：把一条指令的布局写成 layout

mma.sync 的难点在于：32 个线程各自持有 A、B、C 的哪几个元素，是硬件规定好的。CuTe 用TV layout（(线程号, 值号) → (行, 列)）把这个规定写下来：cutlassMMA_Traits SM80_16x8x16 里的 ALayout、BLayout、CLayout，例如 C 的布局是 cutlassSM80_16x8_Row。这就是 CuTe 专题 第 ⑦ 节说的“硬件也用 Layout 说话”，也是 G1 里你写的“lane 持有哪几个元素”的计算。

TiledMMA：用小印章拼大印章

一条 mma 只算 16×8 的 C。make_tiled_mma 把这个 Atom 在多个 warp 上、在 M / N 方向上重复，拼成一个大的 tile：cutlassmake_tiled_mma(SM80_16x8x16_F16F16F16F16_TN, …)。之后的代码只调用 gemm(mma, A, B, C)（cutlass298 行），每个线程该算哪些、该持有哪些，全部由 layout 推出来。同一份 kernel 换成 cutlassUniversalFMA（普通 FMA）也能跑，只是换了一个 Atom。

多级流水线

cutlassK_PIPE_MAX 个 stage：先发出前几个 stage 的 cp.async，主循环里 cutlass等最老的一个 stage 到齐（cp_async_wait<K_PIPE_MAX−2>），算它，同时发出下一个 stage 的拷贝。

示意图，不是实测时间。有了异步拷贝，第 k 块的计算和第 k+1、k+2 块的搬运同时进行，总时间从“搬 + 算”的和变成接近两者中较大的那个。

④ Hopper（sm_90a）：TMA、wgmma、mbarrier 与 warp specialization

TMA：一个线程搬一整个 tile

Ampere 的 cp.async 仍然是每个线程算自己那几个元素的地址、各发各的拷贝。Hopper 的 TMA（Tensor Memory Accelerator）是一个专用的拷贝引擎：

• host 端先建一个 tensor map（描述整个张量的形状、stride、tile 大小、swizzle 方式）：cutlassmake_tma_atom(SM90_TMA_LOAD, mA, sA, tile_shape)。
• kernel 里一个线程发一条 cp.async.bulk.tensor，整个 tile（比如 128×64）就被搬进 shared memory：cutlassSM90_TMA_LOAD_2D。地址计算、越界处理、swizzle 全部由硬件完成。
• 完成通知走 mbarrier，按字节数计：cutlassarrive_and_expect_tx(barrier, 字节数) 之后发 TMA，消费者 cutlasswait(barrier, phase)，等到这么多字节都到了才继续。
• multicast：一个 cluster 里的几个 CTA 要同一块数据时，TMA 可以一次读、同时写进多个 CTA 的 shared memory：cutlassSM90_TMA_LOAD_MULTICAST_2D。cluster 是 Hopper 新加的一层：几个 CTA 保证同时运行在相邻的 SM 上，还能读写彼此的 shared memory（distributed shared memory）。

wgmma：四个 warp 合用一台更大的机器，操作数直接从 shared memory 读

• warpgroup = 4 个连续的 warp（128 个线程）一起发一条 wgmma.mma_async，M 固定为 64，N 从 8 到 256，FP16 时 K = 16：cutlassMMA_64x8x16 … m64n8k16（一直到 m64n256k16）。
• 操作数来自 shared memory（通过一个 64 位的描述符，A 也可以放在寄存器里），不再需要 ldmatrix 把数据先搬进寄存器。
• 异步：发出之后线程可以继续做别的事，之后再等它完成：cutlasswgmma.fence / commit_group / wait_group，在 CuTe 里是 cutlasswarpgroup_arrive → gemm → warpgroup_commit_batch → warpgroup_wait。TiledMMA 的构造和 Ampere 一样：cutlassmake_tiled_mma(SM90_64x64x16_F16F16F16_SS)。

warp specialization：分工，而不是每个人什么都干

有了异步的 TMA 和异步的 wgmma，最自然的写法是分工：

• setmaxnreg：生产者几乎不需要寄存器（只是发 TMA），消费者需要大量寄存器存累加器。Hopper 允许 warpgroup 在运行中交换寄存器配额：FA3 里生产者 fawarpgroup_reg_dealloc、消费者 fawarpgroup_reg_alloc；两个消费者 warpgroup、用 TMA 读 K/V 时，生产者降到每线程 24 个寄存器，消费者升到 240 个，其他配置的数字在 faLoadRegisterRequirement / MmaRegisterRequirement。文档：PTX · setmaxnreg。

读 FlashAttention-3：生产者 + 两个消费者的乒乓

FA3 的 forward kernel 把上面这些全部用上：

• 角色：fawarp_group_idx == 0 是生产者（发 Q、K、V 的 TMA），fa其余是消费者（做 QKᵀ、softmax、PV）；消费者的个数是 faNumMmaWarpGroups，1–3 个。
• 主循环：生产者的 faload() 和消费者的 famma()。
• 乒乓调度：softmax 用的是普通 ALU（exp、max、sum），GEMM 用的是 Tensor Core，两者可以同时进行。两个消费者 warpgroup 用 fawarp_scheduler_barrier_sync / arrive 轮流使用 Tensor Core：一个在做 GEMM 时，另一个在做 softmax。另外还有 warpgroup 内部的重叠（faIntraWGOverlap）：发出下一块的 QKᵀ 之后，再做上一块的 softmax。

Tensor Core（绿色）几乎一直有活干：总有一个 warpgroup 在做 GEMM。softmax（紫色）的指数运算吞吐比 GEMM 低得多，如果不和 GEMM 重叠，它会成为瓶颈。

读 DeepGEMM（SM90）：FP8 GEMM 的精度补偿

DeepGEMM 的 SM90 FP8 GEMM 结构完全一样：deepgemmkNumTMAThreads = 128（一个 warpgroup）+ 若干 math warpgroup，deepgemmTMA 侧 reg_dealloc、deepgemmmath 侧 reg_alloc，deepgemmfull / empty barrier。

它多做了一件事，和 M3 直接相关：每个 K 块做完 deepgemmFP8 的 wgmma 之后，把结果乘上这一块的 scale，deepgemm加到 FP32 的 final_accum 里（在 CUDA core 上做）。这就是 M3 里 acc += tl.dot(a, b) * (sa * sb) 的硬件级版本，同时也补偿了 Hopper FP8 Tensor Core 内部累加精度有限的问题。

⑤ Blackwell（sm_100a）：tcgen05、Tensor Memory 与 2-CTA MMA

tcgen05.mma：一个线程发起，累加器放在 Tensor Memory

源码：cutlassSM100_MMA_F16BF16_SS → tcgen05.mma.cta_group::1.kind::f16；TMEM 的分配 cutlasstcgen05.alloc；2-CTA 版本 cutlassSM100_MMA_F16BF16_2x1SM_SS → cta_group::2；block-scaled 版本 cutlassSM100_MMA_MXF8F6F4_SS → kind::mxf8f6f4.block_scale。

CuTe 的 Blackwell 教程按难度排好了：cutlass01：tcgen05.mma + tcgen05.ld（cutlass分配 TMEM、cutlassgemm 之后 umma_arrive 通知、cutlassmake_tmem_copy 把结果读回寄存器）→ 02：加 TMA → 03：TMA multicast → cutlass04：2-CTA 的 MMA（cutlassSM100_TMA_2SM_LOAD_MULTICAST）→ 05：TMA epilogue。

读 FlashAttention-4：用 CuTe DSL 写成，角色分得更细

FA4（faflash_fwd_sm100.py，用 Python 的 CuTe DSL 写）的角色分配一目了然：fasoftmax0、softmax1 各 4 个 warp，correction 4 个 warp，MMA 1 个 warp，epilogue 1 个 warp，load 1 个 warp。

• S 和 O 都在 TMEM 里：fatmem_s_offset、tmem_o_offset，QKᵀ 的结果 S 和输出累加器 O 各占 TMEM 的一段，两个 Q tile（softmax0 和 softmax1）乒乓。
• correction warp：online softmax 里“max 变了就要把 O 乘一个系数”（T5 讲过的 alpha）由专门的 warp 去做，不挡 softmax 和 MMA 的路。
• exp2 软件模拟：Blackwell 的 Tensor Core 太快了，softmax 里的指数运算（硬件的 MUFU 单元）反而成了瓶颈，FA4 用多项式在 FMA 单元上模拟一部分 exp2：faex2_emu 的配置。这正是第 ① 节那条主线的延续：Tensor Core 越快，非矩阵的部分越要想办法挤。
• 寄存器配额同样在运行中调整：fasetmaxregister_decrease。

读 DeepGEMM（SM100）：一个 warp 一个角色

deepgemmwarp 0：TMA load，deepgemmwarp 1：MMA issue（只在 leader CTA 上），deepgemm其余 warp group：epilogue。TMEM 分配按是否 2-CTA 选择 deepgemmAllocator1Sm 或 Allocator2Sm，MMA 的形状见 deepgemmUMMA_M / UMMA_N / UMMA_K。和 SM90 版本相比，FP8 的 scale 由 Tensor Core 的 block scaling 在硬件里乘，不再需要 CUDA core 上的 promotion。

⑥ 哪块 GPU 能跑哪些

sm_90a、sm_100a 里的 a 表示“架构专用特性”：编译出来的代码只能在这一代上运行，不向后兼容。这也是 FlashAttention、DeepGEMM 为每一代 GPU 单独写一套主循环的原因。

⑦ 推荐的阅读和动手路线

1. 在 5090 上就能做：用 WMMA 改写 C4（第 ②节），然后读 cutlasssgemm_sm80 和 cutlass0t_mma_atom.md，把 MMA Atom 的 TV layout 和 G1 / G2 对照起来。
2. 有 H100 时：跑 cutlasswgmma_tma_sm90.cu，再读 FA3 的 fa生产者 / 消费者分支 和 DeepGEMM SM90 的主循环。
3. 有 B200 时：按顺序跑 cutlassblackwell/01–05，再读 FA4 和 DeepGEMM SM100。

源码与资料索引

本仓库：C4 讲解 · T4 讲解 · T5 讲解 · M3 讲解 · CuTe 专题 · G1 讲解

WMMA：cudasamplessimple_wmma_gemm · cudasamplescompute_gemm · Programming Guide · WMMA

CuTe（Ampere）：cutlassSM80_16x8x16 · cutlassMMA_Traits · cutlassldmatrix · cutlasscp.async · cutlasssgemm_sm80 · cutlass0t_mma_atom.md

Hopper：cutlassTMA · cutlasswgmma · cutlasswgmma_tma_sm90.cu · faFA3 kernel · faFA3 主循环 · deepgemmDeepGEMM SM90 · flashinferFlashInfer 的 Hopper prefill（同样的生产者 / 消费者）

Blackwell：cutlasstcgen05.mma · cutlassTMEM 分配 · cutlassblackwell/01–05 · faFA4 · deepgemmDeepGEMM SM100

PTX 文档：mma（m16n8k16 的寄存器布局） · ldmatrix · cp.async · cp.async.bulk.tensor（TMA） · mbarrier · wgmma · setmaxnreg · tcgen05.mma · Tensor Memory 分配 · CTA pair · block scaling

CUDA Programming Guide：TMA · Thread Block Clusters

论文：FlashAttention-3 · DeepSeek-V3（FP8 训练与 DeepGEMM 的背景）