想象你是“福满楼”的行政总厨,接下一场一千二百人的婚宴。几十道菜、上百种食材,靠一位厨师一道菜一道菜地做,天亮也上不完。于是你购进了一台“超级料理机”:只要把食材按配方放进去,它一次就能把一整批菜品混合、调味、装盘全部搞定。这台机器就是 GPU 里的 Tensor Core——专为矩阵乘加打造的硬件引擎。可光有机器还不够,你手下有 32 位帮厨(一个 Warp),必须精密配合:谁来递食材、每人手里端哪几样、按什么顺序下锅、出锅的菜怎么摆盘。这套“32 人协作驾驭超级料理机”的标准化操作手册,就是 NVIDIA 的 WMMA API——Warp 级矩阵乘加接口。它规定了如何把矩阵切成 Fragment(食材碟)、何时调用 mma_sync(下锅)、又如何取回结果。本文就带你读懂这份手册,从硬件原理到代码实战,一步到位。

岗位关联

一句话价值:掌握 WMMA API 是解锁 Tensor Core 算力的钥匙,让你在推理引擎、大模型部署与算子优化岗位上,写出性能翻倍的矩阵运算核心。

目录
  1. 可视化演示
  2. 硬件算力基石
  3. WMMA 基本概念
  4. 形状与数据布局
  5. API 核心流程
  6. 精度类型选择
  7. 性能优化策略
  8. 工程实践与陷阱
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

展示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 位帮厨怎么把食材递给它,才能让机器一刻不空转。

D=A×B+CD = A \times B + C
这个公式是 Tensor Core 的灵魂:一次调用同时完成“乘”和“加”,正是神经网络全连接层、卷积、注意力机制最常用的运算。
  1. 输入矩阵 A形状为 M×K,记录每桌客人需要的 K 种食材数量,是左侧的“食材总表”。
    A∈RM×KA \in \mathbb{R}^{M \times K}
  2. 输入矩阵 B形状为 K×N,是“食材转调料”的换算表,把每种食材换算成 N 种调料用量。
    B∈RK×NB \in \mathbb{R}^{K \times N}
  3. 累加矩阵 C相当于下锅前就铺好的“底味”,可为全零矩阵,也可携带偏置或上次累加的中间结果。
    C∈RM×NC \in \mathbb{R}^{M \times N}
  4. 逐项乘加求和输出 D 的每个元素都是 A 的一行与 B 的一列逐项相乘再累加,最后叠上 C 的底味。
    Dij=∑k=1KAik×Bkj+CijD_{ij} = \sum_{k=1}^{K} A_{ik} \times B_{kj} + C_{ij}
  5. 输出矩阵 D最终每桌客人的成品菜,形状与 C 完全一致。
    D∈RM×ND \in \mathbb{R}^{M \times N}

WMMA 基本概念

有了料理机还不够,福满楼的 32 位帮厨得先把食材切成小份,每人端一碟。在 WMMA 里,这一碟“食材”就是 Fragment——矩阵被切成许多小块后,每个线程手里持有的一小撮数据。一个 Warp 有 32 个线程,它们集体处理一个输出小块,谁也离不开谁。

WMMA 组织一次计算的方式很固定:先把 A、B 矩阵的对应切片加载成 a_fragment 和 b_fragment,再把累加器 c_fragment 一起交给 mma_sync。一次 mma_sync 等于让料理机按配方快速颠勺一次,把一小块矩阵乘加做完。这里的 Fragment 不是普通数组,而是编译器管理的寄存器数据,每个线程只能看见自己手里的那几份。

累加器为什么重要?因为矩阵乘法本质上是大量乘积的累加。把 C 放进乘加公式,意味着每次 mma 都能接着上次的结果继续加,避免把中间结果反复写回显存。这正是 WMMA 又快又稳的关键——就像炒一大锅菜时,厨师不会每加一味料就把锅端离火灶,而是不断翻搅累积味道。

现在明白了“每人端一碟”,接下来要回答:碟子里的食材到底怎么摆,料理机才认得出?

elements per thread=M×N32\text{elements per thread} = \frac{M \times N}{32}
这是理解 Fragment 最直观的公式:输出矩阵总元素数除以 32 个线程,得到每个线程手里的“食材份额”。
  1. 行数 M 与列数 N输出矩阵的规模,对应婚宴的桌数和菜数,两者相乘就是输出矩阵的元素总量。
    M×NM \times N
  2. 总元素数以最常用的 16×16 输出块为例,共有 256 个输出元素需要计算。
    M×N=16×16=256M \times N = 16 \times 16 = 256
  3. 除以 32一个 Warp 固定有 32 个线程,大家把计算责任均匀分摊。
    25632=8\frac{256}{32} = 8
  4. 每线程 8 个元素每个线程的累加器 Fragment 恰好持有 8 个输出元素的累加任务。
    frag_size=8\text{frag\_size} = 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 位帮厨的碟子内容互不重叠、一步到位。

食材摆好了,下面看完整的操作流水线:怎么备菜、怎么下锅、怎么出锅。

MMA⁡(m,n,k):(m,n,k)∈{(16,16,16),(16,8,8),(8,16,8)}\operatorname{MMA}_{(m,n,k)}: \quad (m,n,k) \in \{(16,16,16), (16,8,8), (8,16,8)\}
这是 WMMA 的“模具清单”:硬件只认这几档尺寸,选型时以目标架构手册为准。
  1. 输出行数 mm16n16k16 中的第一个 16,代表一次算出的输出行数,类似一桌摆几道菜。
    m=16m = 16
  2. 输出列数 n第二个 16,代表输出列数,类似同时服务的桌数。
    n=16n = 16
  3. 累加深度 k第三个 16,代表一次乘加要累加多少层,类似每道菜混合多少种食材。
    k=16k = 16
  4. 其他可选模具硬件还支持更小的组合,供不同矩阵形状选择。
    (16,8,8), (8,16,8)(16,8,8), \ (8,16,8)

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 初始化累加器。如果跳过它,累加器里残留的是上次的垃圾数据,结果可想而知——好比厨师没洗锅就直接炒下一道菜,整锅菜都串了味。

流程清楚了,还得选对食材等级——精度。不同精度的食材,出菜速度和味道天差地别。

accfrag=MMA⁡(Load⁡(Afrag),Load⁡(Bfrag),accfrag)\text{acc}_{\text{frag}} = \operatorname{MMA}\big(\operatorname{Load}(A_{\text{frag}}), \operatorname{Load}(B_{\text{frag}}), \text{acc}_{\text{frag}}\big)
这行公式浓缩了 WMMA 的标准流水线:加载两个输入碎片,乘以累加器,再写回累加器。
  1. 加载 A 碎片load_matrix_sync 把 A 矩阵的切片从全局或共享内存搬进每个线程的寄存器。
    fragA=load_matrix_sync(Aslice)\text{fragA} = \text{load\_matrix\_sync}(A_{\text{slice}})
  2. 加载 B 碎片B 矩阵的切片以同样方式装入,两个输入都就位。
    fragB=load_matrix_sync(Bslice)\text{fragB} = \text{load\_matrix\_sync}(B_{\text{slice}})
  3. 同步乘加mma_sync 让 Tensor Core 一次性完成 k 层乘加,结果累加进 acc,全程无内存读写。
    accfrag=mma_sync(fragA,fragB,accfrag)\text{acc}_{\text{frag}} = \text{mma\_sync}(\text{fragA}, \text{fragB}, \text{acc}_{\text{frag}})
  4. 存储结果store_matrix_sync 把累加器里的成品写回内存,供上层算法继续使用。
    store_matrix_sync(ptr,accfrag)\text{store\_matrix\_sync}(\text{ptr}, \text{acc}_{\text{frag}})

精度类型选择

婚宴规模不同,食材档次也不同。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,保证长期累加不漂移——就像炖汤时虽然食材是速冻的,但汤底一定要用大火收浓,味道才稳。

选好了食材,最后看效率:怎样让料理机一刻都不闲着,出菜又快又多。

FP16: [1,5,10],BF16: [1,8,7],FP32: [1,8,23]\text{FP16: } [1,5,10], \quad \text{BF16: } [1,8,7], \quad \text{FP32: } [1,8,23]
三个数字分别代表符号位、指数位、尾数位。这一组配置直接决定了每个精度能表示的范围和精度。
  1. 符号位最高位表示正负,三种格式都是 1 位,规则完全相同。
    s∈{0,1}s \in \{0, 1\}
  2. 指数位决定可表示数的动态范围,BF16 与 FP32 同为 8 位,因此范围一样大。
    eFP16=5,eBF16=8,eFP32=8e_{\text{FP16}} = 5, \quad e_{\text{BF16}} = 8, \quad e_{\text{FP32}} = 8
  3. 尾数位决定数值精度,尾数越多,小数部分刻画越精细。
    fFP16=10,fBF16=7,fFP32=23f_{\text{FP16}} = 10, \quad f_{\text{BF16}} = 7, \quad f_{\text{FP32}} = 23
  4. 取值公式浮点数实际值由符号、尾数、指数共同还原。
    (−1)s×1.f×2e−bias(-1)^s \times 1.f \times 2^{e - \text{bias}}

性能优化策略

福满楼生意火爆时,主厨不会等一道菜完全做好才开始下一道,而是这边料理机在颠勺,那边帮厨已经把下一批食材摆上了台面。这就是流水线。WMMA 的性能优化,本质上就是让“备菜”和“做菜”这两件事重叠起来。

实际矩阵的 K 维通常远超 16,比如 K = 1024。一次 mma_sync 只算 16 层累加,必须循环调用 64 次,每次都把结果累加到同一个累加器里。这个循环就是 K 循环累加,它让 Tensor Core 像接力赛一样,一棒接一棒不停跑。优化手段主要有三:一是让数据加载与 mma_sync 计算重叠,用双缓冲隐藏访存延迟;二是把矩阵分块放进共享内存,减少全局内存访问;三是保证指针 16 或 32 字节对齐,让显存带宽跑满。记住一句话:访存永远是瓶颈,计算单元吃得饱不饱,全看“投喂”效率高不高。

优化手段千千万,最后都要落在工程细节上。下一节看看最容易翻车的几个坑。

Cacc=∑k=0K/16−1A[:, 16k:16k+16]×B[16k:16k+16, :]C_{\text{acc}} = \sum_{k=0}^{K/16-1} A_{[:, \, 16k : 16k+16]} \times B_{[16k : 16k+16, \, :]}
大 K 被切成许多 16 层的小段,每一段交给一次 mma_sync,所有段的结果累加进同一个累加器。
  1. 切分 K 维把大 K 分成若干段,每段恰好 16 列,便于交给 mma_sync。
    Ktiles=K/16K_{\text{tiles}} = K / 16
  2. 逐段乘加每一轮取一块 A 和一块 B 做 MMA,结果累加到同一个累加器。
    Cacc+=Ak×BkC_{\text{acc}} \mathrel{+}= A_k \times B_k
  3. 循环范围从 k=0 到最后一个分块,索引必须写清楚,避免越界。
    ∑k=0Ktiles−1\sum_{k=0}^{K_{\text{tiles}}-1}
  4. 流水线重叠第 k+1 段的加载与第 k 段的计算同时进行,隐藏访存延迟。
    load(k+1)∥compute(k)\text{load}(k+1) \parallel \text{compute}(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 mod 16=0,N mod 16=0,K mod 16=0M \bmod 16 = 0, \quad N \bmod 16 = 0, \quad K \bmod 16 = 0
这是 WMMA 的“对齐三连”:行、列、累加深度都必须能被 16 整除,否则补零。
  1. M 对齐输出行数必须是 16 的整数倍,不满足时补零。
    M′=M+(16−M mod 16) mod 16M' = M + (16 - M \bmod 16) \bmod 16
  2. N 对齐输出列数必须是 16 的整数倍。
    N′=N+(16−N mod 16) mod 16N' = N + (16 - N \bmod 16) \bmod 16
  3. K 对齐累加深度必须是 16 的整数倍,这是最容易被忽略的一条。
    K′=K+(16−K mod 16) mod 16K' = K + (16 - K \bmod 16) \bmod 16
  4. 零填充扩展出来的区域用 0 填充,不影响乘加结果。
    A′=[A∣0]A' = [A \mid 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 演示

python40 行
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()
目的:用纯 Python + NumPy 模拟 WMMA 的“32 线程分摊 16×16 输出块”核心思想,验证 Fragment 分配后的计算结果与完整矩阵乘法一致。
  1. def wmma_simulate(shape=(16, 16, 16)):定义模拟函数,默认复刻 m16n16k16 标准形状。
  2. A = rng.standard_normal((M, K)).astype(np.float32)生成输入矩阵 A,对应“食材总表”。
  3. expected = A @ B + C0用 NumPy 的完整矩阵乘法算出参考结果。
  4. def owner_of(row, col):定义线程映射函数,决定每个输出元素由哪位“帮厨”负责。
  5. return row * 2 + col // 8简化映射:每个线程负责半行 8 个输出元素。
  6. acc = C0[row, col]累加器从 C 矩阵的底味出发,对应 fill_fragment 后的初始状态。
  7. for k in range(K):内层 K 循环,对应 mma_sync 每次完成的累加深度。
  8. np.testing.assert_allclose(D_sim, expected, rtol=1e-4)一致性校验,确保碎片化计算与整体矩阵乘结果吻合。
  9. 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 中固定包含多少个线程?

单选题 · easy

  • A.8 个线程
  • B.16 个线程
  • C.32 个线程← 正确答案
  • D.64 个线程

答案:C。Warp 是 GPU 调度的最小单位,固定 32 个线程,就像福满楼的帮厨组固定 32 人。WMMA 的"W"就是指 Warp,所有 Fragment 按 32 份分发。

Q2. 下列哪一个是 WMMA API 支持的标准计算形状?

单选题 · easy

  • 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 个线程分摊,每个线程负责几个输出元素?

单选题 · easy

  • A.2 个输出元素
  • B.4 个输出元素
  • C.8 个输出元素← 正确答案
  • D.16 个输出元素

答案:C。16×16 共 256 个元素,除以 32 个线程,每个线程负责 8 个。这就是“总菜量除以帮厨人数”的均分逻辑,也是 Fragment 大小的来源。

Q4. WMMA 的 Fragment 数据实际存放在哪里?

单选题 · mid

  • A.全局内存显存
  • B.共享内存区块
  • C.线程私有寄存器← 正确答案
  • D.L2 高速缓存中

答案:C。Fragment 是编译器管理的寄存器片段,每个线程只持有其中部分元素,访问速度最快。它不在全局内存,也不在共享内存,更不能被其他线程直接访问。

Q5. load_matrix_sync 的作用是什么?

单选题 · mid

  • A.从内存加载到片段← 正确答案
  • B.把片段存回内存
  • C.执行矩阵乘加计算
  • D.清空所有累加器

答案:A。load_matrix_sync 负责把矩阵切片从全局或共享内存搬进 Fragment,对应“帮厨备菜装碟”。存回内存是 store_matrix_sync,执行乘加是 mma_sync。

Q6. 哪种数据类型的指数位与 FP32 相同,拥有更广的数值范围?

单选题 · mid

  • A.FP16 半精度
  • B.BF16 大脑浮点← 正确答案
  • C.INT8 整型类型
  • D.TF32 截断浮点

答案:B。BF16 的配置是 1 符号位 + 8 指数位 + 7 尾数位,指数位与 FP32 相同,因此动态范围大、精度低,适合大跨度训练场景。

Q7. 当 K 维度为 40 时,直接用 m16n16k16 计算会发生什么?

单选题 · hard

  • A.自动补齐到 48
  • B.程序直接崩溃
  • C.结果错误或未定义← 正确答案
  • D.编译阶段会报错

答案:C。40 不是 16 的倍数,硬件不会自动补零,剩余 8 层无人处理,结果错误或行为未定义。程序员必须手动把 K 补齐到 48,并配合零填充。

Q8. 下列哪些精度类型可以被 WMMA/Tensor Core 使用?

多选题 · mid

  • A.FP16 半精度← 正确答案
  • B.BF16 大脑浮点← 正确答案
  • C.FP64 双精度
  • D.INT8 整型量化← 正确答案

答案:ABD。Tensor Core 支持 FP16、BF16、INT8,以及新架构的 TF32。FP64 是通用 CUDA Core 的领域,Tensor Core 并不直接用于双精度矩阵乘加。

Q9. 以下哪些属于 WMMA 性能优化策略?

多选题 · hard

  • A.K 维度循环累加← 正确答案
  • B.双缓冲流水线← 正确答案
  • C.数据 32B 对齐← 正确答案
  • D.使用 Python 编写内核

答案:ABC。K 循环累加让 Tensor Core 接力运算,双缓冲让加载与计算重叠,32B 对齐跑满显存带宽。Python 无法编写 CUDA 内核,更谈不上优化 WMMA。

Q10. 在 m16n16k16 的 FP16 WMMA 计算中,累加器 fragment 的默认数据类型是?

单选题 · extreme

  • A.FP16 半精度类型
  • B.BF16 大脑浮点类型
  • C.FP32 单精度浮点← 正确答案
  • D.INT32 整型累加类型

答案:C。WMMA 的设计精髓在于输入用 FP16、累加用 FP32。即使每轮输入都是半精度,累加器仍以 FP32 保存中间结果,避免逐轮舍入误差累积。这就像炒菜时厨师用大汤勺不断翻匀,调料误差不会越积越多。