想象你是“福满楼”的行政总厨,接下一场一千二百人的婚宴。几十道菜、上百种食材,靠一位厨师一道菜一道菜地做,天亮也上不完。于是你购进了一台“超级料理机”:只要把食材按配方放进去,它一次就能把一整批菜品混合、调味、装盘全部搞定。这台机器就是 GPU 里的 Tensor Core——专为矩阵乘加打造的硬件引擎。可光有机器还不够,你手下有 32 位帮厨(一个 Warp),必须精密配合:谁来递食材、每人手里端哪几样、按什么顺序下锅、出锅的菜怎么摆盘。这套“32 人协作驾驭超级料理机”的标准化操作手册,就是 NVIDIA 的 WMMA API——Warp 级矩阵乘加接口。它规定了如何把矩阵切成 Fragment(食材碟)、何时调用 mma_sync(下锅)、又如何取回结果。本文就带你读懂这份手册,从硬件原理到代码实战,一步到位。
一句话价值:掌握 WMMA API 是解锁 Tensor Core 算力的钥匙,让你在推理引擎、大模型部署与算子优化岗位上,写出性能翻倍的矩阵运算核心。
可视化演示
展示Tensor Core的WMMA API三阶段流程:矩阵分片加载、MMA乘加计算与结果存储,步进切换查看每个阶段的核心机制。
本节配有交互动效(进入视口后自动加载,约 3KB)。
加载阶段:矩阵分片载入寄存器 load_matrix_sync从共享内存读取矩阵A和B分片 Warp中32线程协作分配矩阵片段到寄存器 load_matrix_sync(A, B, C) 支持fp16/bf16/tf32数据类型
硬件算力基石
“福满楼”的后厨新进了一台超级料理机。它不是普通的锅碗瓢盆,而是一台能同时把整批食材打成浆、拌好料、装好盘的专用机器。这台机器就是 Tensor Core——GPU 里专门为矩阵乘加打造的硬件单元。它不干别的,只做一件事:把两个矩阵相乘、再加一个矩阵,而且一条指令就能完成一小块矩阵的全部乘加操作。
在 Tensor Core 出现之前,GPU 靠成百上千个通用 CUDA Core 一个个地算乘加;Tensor Core 出现后(Volta 架构起),一块 4×4 的矩阵乘加可以在一条指令里完成。对 AI 来说,神经网络里 90% 以上的算力都花在矩阵乘加(GEMM)上,所以这台“料理机”直接决定了 AI 训练和推理的速度上限。WMMA 就是使用这台料理机的接口,全称 Warp-level Matrix Multiply-Accumulate,核心计算可以写成:把输入矩阵 A、B 相乘,再与累加矩阵 C 相加,得到输出 D。每一个 WMMA 调用,本质上都是让 32 个线程协作,把一块小矩阵的乘加任务交给 Tensor Core 去完成。
理解了这台料理机能做什么,下一个问题就是:32 位帮厨怎么把食材递给它,才能让机器一刻不空转。
- 输入矩阵 A形状为 M×K,记录每桌客人需要的 K 种食材数量,是左侧的“食材总表”。
- 输入矩阵 B形状为 K×N,是“食材转调料”的换算表,把每种食材换算成 N 种调料用量。
- 累加矩阵 C相当于下锅前就铺好的“底味”,可为全零矩阵,也可携带偏置或上次累加的中间结果。
- 逐项乘加求和输出 D 的每个元素都是 A 的一行与 B 的一列逐项相乘再累加,最后叠上 C 的底味。
- 输出矩阵 D最终每桌客人的成品菜,形状与 C 完全一致。
WMMA 基本概念
有了料理机还不够,福满楼的 32 位帮厨得先把食材切成小份,每人端一碟。在 WMMA 里,这一碟“食材”就是 Fragment——矩阵被切成许多小块后,每个线程手里持有的一小撮数据。一个 Warp 有 32 个线程,它们集体处理一个输出小块,谁也离不开谁。
WMMA 组织一次计算的方式很固定:先把 A、B 矩阵的对应切片加载成 a_fragment 和 b_fragment,再把累加器 c_fragment 一起交给 mma_sync。一次 mma_sync 等于让料理机按配方快速颠勺一次,把一小块矩阵乘加做完。这里的 Fragment 不是普通数组,而是编译器管理的寄存器数据,每个线程只能看见自己手里的那几份。
累加器为什么重要?因为矩阵乘法本质上是大量乘积的累加。把 C 放进乘加公式,意味着每次 mma 都能接着上次的结果继续加,避免把中间结果反复写回显存。这正是 WMMA 又快又稳的关键——就像炒一大锅菜时,厨师不会每加一味料就把锅端离火灶,而是不断翻搅累积味道。
现在明白了“每人端一碟”,接下来要回答:碟子里的食材到底怎么摆,料理机才认得出?
- 行数 M 与列数 N输出矩阵的规模,对应婚宴的桌数和菜数,两者相乘就是输出矩阵的元素总量。
- 总元素数以最常用的 16×16 输出块为例,共有 256 个输出元素需要计算。
- 除以 32一个 Warp 固定有 32 个线程,大家把计算责任均匀分摊。
- 每线程 8 个元素每个线程的累加器 Fragment 恰好持有 8 个输出元素的累加任务。
形状与数据布局
料理机对食材摆放有严格要求:块切得太大塞不进去,切得太小又浪费产能。WMMA 规定了一次计算的标准尺寸,其中最常用的是 16×16×16。这三个数字分别对应输出矩阵的行数 m、列数 n,以及一次乘加的累加深度 k。
m16n16k16 的含义是:32 个线程合作,一次算出一个 16×16 的输出块,每个输出元素都对 16 个乘积求和。除了这个组合,硬件还支持 16×8×8、8×16×8 等形状。所有维度必须严格匹配,这是硬件排线和寄存器布局共同决定的,程序员不能自由发挥——好比料理机的模具是固定尺寸,想换花样只能换模具,不能硬塞。
数据的摆放方式同样重要:row_major(行主序)表示矩阵在内存里按行连续存放,col_major(列主序)则按列连续存放。加载 Fragment 时必须告诉 API 实际布局,否则数据就像菜谱拿反了一样,怎么看都对不上。此外,每个线程具体拿到哪几个元素,由硬件固定的映射规则决定,程序员无权指定——这是为了确保 32 位帮厨的碟子内容互不重叠、一步到位。
食材摆好了,下面看完整的操作流水线:怎么备菜、怎么下锅、怎么出锅。
- 输出行数 mm16n16k16 中的第一个 16,代表一次算出的输出行数,类似一桌摆几道菜。
- 输出列数 n第二个 16,代表输出列数,类似同时服务的桌数。
- 累加深度 k第三个 16,代表一次乘加要累加多少层,类似每道菜混合多少种食材。
- 其他可选模具硬件还支持更小的组合,供不同矩阵形状选择。
API 核心流程
福满楼做一道菜的完整流程是:帮厨备菜(把食材装碟)→ 递进料理机 → 料理机颠勺 → 端出成品摆盘。WMMA 的 API 流程与这四步严丝合缝。第一步 load_matrix_sync 把内存里的矩阵切片装进 Fragment;第二步 fill_fragment 把累加器清零(或预置偏置);第三步 mma_sync 是核心计算指令,所有 32 个线程同步执行,把乘加结果写进累加器 Fragment;第四步 store_matrix_sync 把累加器写回内存。
这四步都发生在核函数内部,而且是 Warp 级别的集合操作——32 个线程必须一起调用同一个 API,步调一致,不能各干各的。这就像帮厨们必须同时把碟子递上操作台,如果谁慢了半拍,流水线就全乱了。正因如此,WMMA 代码里你几乎看不到 if(lane==0) 之类的分支,所有线程走同一条路径。
开发者常犯的一个直觉误区是“加载 → 计算 → 存储”三步走,其实还有隐藏的第四步:fill_fragment 初始化累加器。如果跳过它,累加器里残留的是上次的垃圾数据,结果可想而知——好比厨师没洗锅就直接炒下一道菜,整锅菜都串了味。
流程清楚了,还得选对食材等级——精度。不同精度的食材,出菜速度和味道天差地别。
- 加载 A 碎片
load_matrix_sync把 A 矩阵的切片从全局或共享内存搬进每个线程的寄存器。 - 加载 B 碎片B 矩阵的切片以同样方式装入,两个输入都就位。
- 同步乘加
mma_sync让 Tensor Core 一次性完成 k 层乘加,结果累加进 acc,全程无内存读写。 - 存储结果
store_matrix_sync把累加器里的成品写回内存,供上层算法继续使用。
精度类型选择
婚宴规模不同,食材档次也不同。FP16 像普通速冻食材,便宜量大,但保存条件苛刻;BF16 像耐放的干货,不挑存储但味道清淡;INT8 则是大路货,必须提前腌制(量化)才能下锅。Tensor Core 对这些精度的支持,决定了 AI 计算的速度和准确度。
FP16(半精度)在 Tensor Core 上吞吐最高,但数值范围窄、尾数精度有限,最大值约 65504,梯度稍大就会溢出。BF16 舍弃了尾数位,保留与 FP32 相同的 8 位指数,因此数值范围很大,适合训练场景中跨度巨大的中间结果。INT8 需要先把浮点模型量化成 8 位整数,吞吐进一步翻倍,常用于推理。新型号 GPU 还支持 TF32,能直接接受 FP32 输入,省去转换开销。
选型的核心在指数位和尾数位:指数位决定“能表示多大的数”,尾数位决定“数有多精确”。FP16 是 1 位符号 + 5 位指数 + 10 位尾数,BF16 是 1 + 8 + 7,FP32 是 1 + 8 + 23。值得强调的是:WMMA 计算时输入可以是半精度,但累加器通常用 FP32,保证长期累加不漂移——就像炖汤时虽然食材是速冻的,但汤底一定要用大火收浓,味道才稳。
选好了食材,最后看效率:怎样让料理机一刻都不闲着,出菜又快又多。
- 符号位最高位表示正负,三种格式都是 1 位,规则完全相同。
- 指数位决定可表示数的动态范围,BF16 与 FP32 同为 8 位,因此范围一样大。
- 尾数位决定数值精度,尾数越多,小数部分刻画越精细。
- 取值公式浮点数实际值由符号、尾数、指数共同还原。
性能优化策略
福满楼生意火爆时,主厨不会等一道菜完全做好才开始下一道,而是这边料理机在颠勺,那边帮厨已经把下一批食材摆上了台面。这就是流水线。WMMA 的性能优化,本质上就是让“备菜”和“做菜”这两件事重叠起来。
实际矩阵的 K 维通常远超 16,比如 K = 1024。一次 mma_sync 只算 16 层累加,必须循环调用 64 次,每次都把结果累加到同一个累加器里。这个循环就是 K 循环累加,它让 Tensor Core 像接力赛一样,一棒接一棒不停跑。优化手段主要有三:一是让数据加载与 mma_sync 计算重叠,用双缓冲隐藏访存延迟;二是把矩阵分块放进共享内存,减少全局内存访问;三是保证指针 16 或 32 字节对齐,让显存带宽跑满。记住一句话:访存永远是瓶颈,计算单元吃得饱不饱,全看“投喂”效率高不高。
优化手段千千万,最后都要落在工程细节上。下一节看看最容易翻车的几个坑。
- 切分 K 维把大 K 分成若干段,每段恰好 16 列,便于交给 mma_sync。
- 逐段乘加每一轮取一块 A 和一块 B 做 MMA,结果累加到同一个累加器。
- 循环范围从 k=0 到最后一个分块,索引必须写清楚,避免越界。
- 流水线重叠第 k+1 段的加载与第 k 段的计算同时进行,隐藏访存延迟。
工程实践与陷阱
再熟练的帮厨也有打翻盘子的时刻。WMMA 工程实践里最常见的翻车点,就是维度对不齐。比如 K = 40,而 mma_sync 一次只吃 16 层,40 不是 16 的倍数,硬件不会好心帮你处理剩余的 8 层——结果就是数据错位、输出全乱。
正确做法是把 K 补齐到 48(补零填充),或者让上层循环跳过多余部分。M、N 不满足 16 的倍数同理,都要补零凑整。其他常见坑还包括:ldm(leading dimension,主维度跨度)写错,导致数据错位读取;Fragment 未用 fill_fragment 初始化;矩阵布局声明与数据实际存储方式不一致。调试时建议先用小尺寸矩阵加 CPU 参考实现对比验证,再用 compute-sanitizer 检查越界访问。
掌握这些细节,你就从“听说过有料理机”进化到“能指挥料理机高效出菜”的厨房主管了。剩下的就是多颠几次勺,把手感练出来——毕竟,纸上谈兵永远不如亲手跑一个正确的 WMMA Kernel 来得踏实。
- M 对齐输出行数必须是 16 的整数倍,不满足时补零。
- N 对齐输出列数必须是 16 的整数倍。
- K 对齐累加深度必须是 16 的整数倍,这是最容易被忽略的一条。
- 零填充扩展出来的区域用 0 填充,不影响乘加结果。
AI 应用场景
大模型推理引擎的 GEMM 算子重写
阶段:模型部署与推理优化
解决的问题:传统方式直接调用 cuBLAS 库,遇到 decode 阶段 M 很小、且需要把 bias 融合进算子时,通用库无法定制,Kernel 启动开销和显存带宽浪费严重。AI 解法是用 WMMA API 手写 GEMM 核心,把矩阵乘、偏置加、激活函数融合进一个核函数,让 Tensor Core 全速运转。
具体案例:某团队在自研 LLM 推理引擎中,用 WMMA 重写 QKV 投影的批量矩阵乘,采用 m16n16k16 形状和 K 循环累加,配合双缓冲流水线,把 decode 阶段单 Token 生成延迟降低了约 30%,显存占用同步下降。
混合精度训练框架的矩阵乘路径
阶段:模型训练加速
解决的问题:全 FP32 训练耗时巨大;PyTorch 自动混合精度(AMP)把 Linear 层的输入转成 FP16,底层需要高效的 FP16 矩阵乘——这正是 Tensor Core 和 WMMA 的主战场。它把“食材快速下锅”的能力交给硬件,同时用 FP32 累加器保住训练精度。
具体案例:DeepSpeed 和 Megatron 的 Transformer 模块中,Linear 层通过 cuTLASS 的 WMMA 实现(与 nvcuda::wmma 同源),在 FP16 输入、FP32 累加的组合下,大规模 GPT 训练的矩阵乘吞吐提升数倍。
Python 演示
import numpy as np def wmma_simulate(shape=(16, 16, 16)): """用 32 个线程模拟一次 16x16x16 的 WMMA 计算.""" M, N, K = shape THREADS = 32 rng = np.random.default_rng(42) A = rng.standard_normal((M, K)).astype(np.float32) B = rng.standard_normal((K, N)).astype(np.float32) C0 = rng.standard_normal((M, N)).astype(np.float32) expected = A @ B + C0 def owner_of(row, col): # 简化版线程映射: 每个线程负责半行 8 个输出元素 return row * 2 + col // 8 D_sim = np.zeros((M, N), dtype=np.float32) holders = {} for row in range(M): for col in range(N): t = owner_of(row, col) acc = C0[row, col] # 内层 K 循环对应 mma_sync 每次完成的 16 层乘加 for k in range(K): acc += A[row, k] * B[k, col] D_sim[row, col] = acc holders.setdefault(t, []).append((row, col)) np.testing.assert_allclose(D_sim, expected, rtol=1e-4) print("WMMA 模拟验证通过: 32 线程各持 8 个元素") print("线程 0 负责的前 4 个坐标:", holders[0][:4]) print("D[0][0] =", D_sim[0][0]) return D_sim if __name__ == "__main__": wmma_simulate()
def wmma_simulate(shape=(16, 16, 16)):定义模拟函数,默认复刻 m16n16k16 标准形状。A = rng.standard_normal((M, K)).astype(np.float32)生成输入矩阵 A,对应“食材总表”。expected = A @ B + C0用 NumPy 的完整矩阵乘法算出参考结果。def owner_of(row, col):定义线程映射函数,决定每个输出元素由哪位“帮厨”负责。return row * 2 + col // 8简化映射:每个线程负责半行 8 个输出元素。acc = C0[row, col]累加器从 C 矩阵的底味出发,对应fill_fragment后的初始状态。for k in range(K):内层 K 循环,对应mma_sync每次完成的累加深度。np.testing.assert_allclose(D_sim, expected, rtol=1e-4)一致性校验,确保碎片化计算与整体矩阵乘结果吻合。holders.setdefault(t, []).append((row, col))记录每个线程负责的坐标,验证 32 人分工没有遗漏。
常见误区
误区我平时用 Python 写 AI,以为 WMMA 是某个 Python 库,`import wmma` 之后就能像调用 `torch.matmul` 一样使用 Tensor Core。
正解WMMA 是 CUDA C++ 的 Warp 级 API,位于 nvcuda::wmma 命名空间,必须写在核函数里并由 32 个线程同步调用。PyTorch 这类框架底层(cuBLAS/cuTLASS)才封装了等价逻辑。把 WMMA 当 Python 函数调用,等于想用电话遥控后厨的料理机——它根本不在你手边。
为何常见:很多 AI 开发者日常只用 Python,第一次接触 CUDA 时不熟悉“框架 → 编程库 → 硬件指令”的分层结构,误以为所有加速 API 都是 Python 可导入的。
误区Fragment 名字里带“片段”,我觉得它就是矩阵的一个切片,应该是一个可以随意下标访问的数组,放在显存或者共享内存里。
正解Fragment 是编译器管理的寄存器片段,每个线程只持有其中一部分元素,只能访问属于自己的元素,并没有统一的“数组地址”。load_matrix_sync 只是把数据从内存搬运进这些寄存器。想直接按数组下标访问整个矩阵,就像让一个帮厨去拿别人手里的碟子,后厨立刻乱套。
为何常见:直觉上“数据片段”像数组切片,而数组默认存放在内存中,容易混淆存储层级。
误区我习惯了 `torch.matmul` 支持任意形状,以为 `mma_sync` 也会自动把大矩阵切成合适的块,我只需要把整个矩阵传进去就行。
正解WMMA 只支持固定形状组合(16×16×16、16×8×8、8×16×8 等)。尺寸不匹配时,编译可能通过但运行结果错误,甚至触发非法指令。程序员必须自己把大矩阵切成小块,并保证 K 对齐。机器不会替你把菜谱改小,它只认固定模具。
为何常见:高级框架 API 用多了,容易忽略底层硬件“固定模具”的约束,低估了算子开发中形状管理的复杂度。
误区半精度能让 Tensor Core 提速翻倍,那我把所有数据都转成 FP16 不就万事大吉?
正解FP16 吞吐高的前提是数据范围合适。FP16 最大值约 65504,尾数仅 10 位,梯度或中间结果过大时会溢出或严重舍入。跨度过大的场景该用 BF16,或采用混合精度(输入半精度、累加 FP32)。无脑转 FP16 就像把高档食材全部剁成碎末,速度上去了味道却毁了。
为何常见:GPU 厂商的宣传材料常说“半精度算力翻倍”,让人忽略精度适配的前提条件。
误区Fragment 既然是个“片段”,那每个线程应该能像访问数组一样访问 fragment 的任意下标,拿到完整矩阵。
正解Fragment 的索引是线程局部的——每个线程只能按自己的局部下标访问自己持有的元素。试图访问其他线程持有的元素会把硬件布局搞乱,结果是错得莫名其妙的矩阵。正确姿势是只操作自己“碟子”里的数据,跨线程交换必须走 shuffle 指令或共享内存中转。
为何常见:习惯了普通数组“全局可寻址”的思维,不理解 Fragment 的私有分布本质是 WMMA 高性能的来源之一。
面试题
Q1. 一个 Warp 在 CUDA 中固定包含多少个线程?
- A.8 个线程
- B.16 个线程
- C.32 个线程← 正确答案
- D.64 个线程
答案:C。Warp 是 GPU 调度的最小单位,固定 32 个线程,就像福满楼的帮厨组固定 32 人。WMMA 的"W"就是指 Warp,所有 Fragment 按 32 份分发。
Q2. 下列哪一个是 WMMA API 支持的标准计算形状?
- A.8×8×8 矩阵乘加
- B.16×16×16 乘加← 正确答案
- C.32×32×32 乘加
- D.64×64×64 乘加
答案:B。WMMA 最常用的形状是 m16n16k16,还有 16×8×8、8×16×8。8×8×8、32×32×32、64×64×64 都不在支持清单内,硬件模具没有这些尺寸。
Q3. 16×16 输出矩阵由 32 个线程分摊,每个线程负责几个输出元素?
- A.2 个输出元素
- B.4 个输出元素
- C.8 个输出元素← 正确答案
- D.16 个输出元素
答案:C。16×16 共 256 个元素,除以 32 个线程,每个线程负责 8 个。这就是“总菜量除以帮厨人数”的均分逻辑,也是 Fragment 大小的来源。
Q4. WMMA 的 Fragment 数据实际存放在哪里?
- A.全局内存显存
- B.共享内存区块
- C.线程私有寄存器← 正确答案
- D.L2 高速缓存中
答案:C。Fragment 是编译器管理的寄存器片段,每个线程只持有其中部分元素,访问速度最快。它不在全局内存,也不在共享内存,更不能被其他线程直接访问。
Q5. load_matrix_sync 的作用是什么?
- A.从内存加载到片段← 正确答案
- B.把片段存回内存
- C.执行矩阵乘加计算
- D.清空所有累加器
答案:A。load_matrix_sync 负责把矩阵切片从全局或共享内存搬进 Fragment,对应“帮厨备菜装碟”。存回内存是 store_matrix_sync,执行乘加是 mma_sync。
Q6. 哪种数据类型的指数位与 FP32 相同,拥有更广的数值范围?
- A.FP16 半精度
- B.BF16 大脑浮点← 正确答案
- C.INT8 整型类型
- D.TF32 截断浮点
答案:B。BF16 的配置是 1 符号位 + 8 指数位 + 7 尾数位,指数位与 FP32 相同,因此动态范围大、精度低,适合大跨度训练场景。
Q7. 当 K 维度为 40 时,直接用 m16n16k16 计算会发生什么?
- A.自动补齐到 48
- B.程序直接崩溃
- C.结果错误或未定义← 正确答案
- D.编译阶段会报错
答案:C。40 不是 16 的倍数,硬件不会自动补零,剩余 8 层无人处理,结果错误或行为未定义。程序员必须手动把 K 补齐到 48,并配合零填充。
Q8. 下列哪些精度类型可以被 WMMA/Tensor Core 使用?
- A.FP16 半精度← 正确答案
- B.BF16 大脑浮点← 正确答案
- C.FP64 双精度
- D.INT8 整型量化← 正确答案
答案:ABD。Tensor Core 支持 FP16、BF16、INT8,以及新架构的 TF32。FP64 是通用 CUDA Core 的领域,Tensor Core 并不直接用于双精度矩阵乘加。
Q9. 以下哪些属于 WMMA 性能优化策略?
- A.K 维度循环累加← 正确答案
- B.双缓冲流水线← 正确答案
- C.数据 32B 对齐← 正确答案
- D.使用 Python 编写内核
答案:ABC。K 循环累加让 Tensor Core 接力运算,双缓冲让加载与计算重叠,32B 对齐跑满显存带宽。Python 无法编写 CUDA 内核,更谈不上优化 WMMA。
Q10. 在 m16n16k16 的 FP16 WMMA 计算中,累加器 fragment 的默认数据类型是?
- A.FP16 半精度类型
- B.BF16 大脑浮点类型
- C.FP32 单精度浮点← 正确答案
- D.INT32 整型累加类型
答案:C。WMMA 的设计精髓在于输入用 FP16、累加用 FP32。即使每轮输入都是半精度,累加器仍以 FP32 保存中间结果,避免逐轮舍入误差累积。这就像炒菜时厨师用大汤勺不断翻匀,调料误差不会越积越多。