CUDA Tensor Core编程实战指南

可视化演示

硬件架构原理

食堂后厨里，最忙的永远是那位“配餐大师”。普通厨师炒一份菜要起锅烧油、颠勺调味，一套动作十几步；而配餐大师面前有十台标准模具，只要把食材按格子摆好，一压一推，十份套餐同时出炉。Tensor Core 就是 GPU 里的这位配餐大师——它不处理灵活的单条指令，而是在一个周期内完成一次完整的矩阵乘加运算，把 $D = A \times B + C$ 这个深度学习中最常见的公式变成了“肌肉记忆”。

在 Volta 架构出现之前，GPU 只能靠成百上千个 CUDA Core 逐个计算乘法和加法，就像让十位厨师一人炒一道菜再拼盘，速度慢且耗能。Tensor Core 的出现改变了局面：每个 Tensor Core 内部有专门的高速乘法器和加法树，输入两个小矩阵，直接输出乘加结果。以 A100 为例，每个 SM 中有四个 Tensor Core，FP16 峰值算力可达 312 TFLOPS，是同卡 FP32 CUDA Core 算力的数倍。

这种“以空间换时间”的设计，本质上是把硬件单元从“通用”推向“专用”。通用 CUDA Core 负责灵活的控制流和标量运算，Tensor Core 则专注于矩阵乘加这条主干道。两者各司其职，共同支撑起深度学习训练与推理的巨大算力需求。

$$D = A \times B + C$$

1. 符号说明：$D$ 是输出矩阵，$A$ 和 $B$ 是两个输入矩阵，$C$ 是累加矩阵（通常初始为 0 或上一次的偏置结果）。

$$D = A \times B + C$$

2. 乘法优先：先计算 $A$ 与 $B$ 的矩阵乘法，这要求 $A$ 的列数等于 $B$ 的行数，结果矩阵的尺寸由 $A$ 的行数和 $B$ 的列数决定。

$$A_{(M \times K)} \times B_{(K \times N)} = R_{(M \times N)}$$

3. 累加融合：将乘法结果 $R$ 与矩阵 $C$ 逐元素相加，这一步在 Tensor Core 内部的加法树中与乘法同步完成，没有额外的访存开销。

$$D_{ij} = R_{ij} + C_{ij}$$

食堂里配餐大师的模具尺寸是固定的，Tensor Core 也有固定形状偏好（如 16×16×16）。接下来我们看看如何通过 wmma API 与这位大师沟通。

wmma编程模型

配餐大师不会说普通话，他只听“套餐模板”——一种印好格子的标准订单。你不需要告诉他每一份菜怎么炒，只需要把食材按格子放好，在订单上打勾。CUDA 提供的 wmma API（Warp Matrix Multiply-Accumulate）正是这张标准订单：它把矩阵运算包装成三个固定步骤，让程序员无需触碰底层硬件细节。

wmma 的核心是 fragment 这个“分片托盘”。矩阵太大搬不动，fragment 负责把一块 16×16 的小矩阵从全局内存或共享内存“端”到每个线程的寄存器里。load_matrix_sync 函数同步加载数据到 fragment，mma_sync 执行乘加，store_matrix_sync 把结果写回内存。这三个函数调用，就对应了配餐流程中的“备菜 → 压模 → 出餐”。

之所以叫 sync（同步），是因为一个 warp（32 个线程）需要集体协作完成这些操作。每个线程只持有 fragment 中的一部分元素，用“大家同时伸手拿自己那份”来理解最贴切。所有线程步调一致，缺一不可——这正是 wmma 与普通 CUDA 编程最大的思维差异：你不是在写“我的线程做什么”，而是在写“整个 warp 共同完成什么”。

$$\text{C\_frag} += \text{A\_frag} \times \text{B\_frag}$$

1. fragment 初始化：声明 fragment 时指定矩阵类型（matrix_a / matrix_b / accumulator）、精度（half / float）和形状（如 16×16×16），相当于告诉食堂“我要三人份标准套餐”。

$$\text{fragment}<\text{matrix\_a}, \text{half}, 16, 16, 16> \ \text{a\_frag}$$

2. 同步加载：load_matrix_sync 把共享内存或全局内存中的一块数据按固定布局填入 fragment，每个线程拿到属于自己的那几份“食材”。

$$\text{load\_matrix\_sync}(\text{a\_frag}, \text{ptr}, \text{ldm})$$

3. 乘加执行：mma_sync 让整个 warp 的 fragment 进行一次矩阵乘加，硬件内部完成所有乘法和加法，结果累积到累加器 fragment 中。

$$\text{mma\_sync}(\text{c\_frag}, \text{a\_frag}, \text{b\_frag}, \text{c\_frag})$$

现在你有了“标准订单”，但整张菜单可能有几千个菜。下一节我们看看如何把大矩阵“切碎”，分批喂给这位配餐大师。

矩阵分块策略

食材太大，模具放不下怎么办？答案很简单：切。一整头牛不能直接塞进模具，但切成 16×16 的小肉块就刚刚好。矩阵分块（Tiling）就是干这个的——传统的矩阵乘法需要 $M \times N \times K$ 次乘加，如果直接把整个矩阵交给 Tensor Core，硬件模具根本装不下。聪明的做法是把大矩阵切成若干个 16×16（或 16×8×8）的小 tile，逐个送入 Tensor Core 计算。

在 CUDA 程序的视角里，分块策略是一场“三层嵌套的配餐循环”：外层循环遍历输出分块的行和列，内层循环遍历 K 维度上的累加步数。每一轮内层循环，我们都从全局内存取一块 A 的 tile 和一块 B 的 tile，放进共享内存这个“备菜台”，再让 warp 们从备菜台取食材送入 Tensor Core。

共享内存在这里扮演了“冰箱”的角色——全局内存太慢，直接把每个 tile 送到 Tensor Core 会饿着它。先把当前循环需要的 A 和 B 分块搬进共享内存，warp 再从共享内存快速取用，访存速度能提升一个数量级。而 tile 尺寸的选择也很有讲究：16×16 能完美贴合 Tensor Core 的模具，但更大的 tile（如 32×32）能减少循环次数、提高数据复用率，代价是共享内存占用更大，可能降低 SM 上同时驻留的线程块数量。

$$\text{C}_{\text{tile}} = \sum_{k=0}^{K/16-1} \text{A}_{\text{tile}}^{(k)} \times \text{B}_{\text{tile}}^{(k)}$$

1. 外层行块循环：遍历输出矩阵的行方向，每次取 $M_{\text{tile}}$ 行，对应 $\text{C}_{\text{tile}}$ 的第一维。

$$\text{for } i \text{ in range}(0, M, 16)$$

2. 外层列块循环：遍历输出矩阵的列方向，与行块组合确定一个输出分块。

$$\text{for } j \text{ in range}(0, N, 16)$$

3. 内层 K 循环：沿 K 维逐步累加，每次取 A 的 $16 \times 16$ 块和 B 的 $16 \times 16$ 块做乘加。

$$\text{C}_{ij} += \text{A}_{ik} \times \text{B}_{kj}, \quad k = 0, 16, 32, \dots$$

分块是空间上的切配艺术，而混合精度则是“食材用量”上的巧妙计算。接下来我们看看如何用“半份食材”做出同样美味的菜。

混合精度实战

食堂要控制成本，食材贵怎么办？聪明的做法是：排骨切成薄片（FP16）来配菜，但记账时仍然用精确到分的计算器（FP32）。混合精度训练的哲学与此如出一辙——用 FP16 存储和搬运矩阵数据，节省一半内存带宽，但在 Tensor Core 内部用 FP32 做累加，保证计算精度不“缩水”。

为什么要这样做？FP16 只用 16 位存储，数值范围小、精度低，但它的存储体积是 FP32 的一半，搬运速度翻倍。Tensor Core 的硬件设计恰好迎合了这种“不对称”：输入两个 FP16 矩阵，内部用高精度乘法器和 FP32 累加器，最终输出 FP32 结果。就像配餐时食材是半份，但每份菜的分量都由电子秤精确称量。

实际使用中还有一个经典坑：FP16 的动态范围只有约 5.96e-8 到 65504，梯度稍微小一点就会“下溢”变成零，导致训练不收敛。损失缩放（Loss Scaling）技术应运而生——训练时把损失值放大若干倍（如 1024），等梯度计算完毕再缩小回来，相当于把整道菜先放大再摆盘，确保微小食材也能被看见。

$$D_{\text{fp32}} = \text{fp32}(A_{\text{fp16}}) \times \text{fp32}(B_{\text{fp16}}) + C_{\text{fp32}}$$

1. 存储格式：$A$ 和 $B$ 以 FP16 格式存放在内存中，内存占用减半、带宽需求减半。

$$A_{\text{fp16}}, \ B_{\text{fp16}}$$

2. 内部提升：进入 Tensor Core 后，FP16 值被转换为 FP32 参与乘法，避免精度损失。

$$\text{fp32}(A_{\text{fp16}}) \times \text{fp32}(B_{\text{fp16}})$$

3. 高精度累加：乘法结果与 FP32 累加矩阵 $C$ 相加，全程保持单精度累积，只在最终写回时选择是否转换为低精度。

$$D_{ij} \leftarrow (\text{乘积})_{ij} + C_{ij}$$

食材备好了，模板也有了，但如何让流水线持续高速运转不空转？这就进入了性能优化的深水区。

性能优化途径

食堂开业后你发现：配餐大师经常“等菜”。食材没送到，模具空转；模具压完了，下一个分块还没切好。性能优化的本质，就是消除这些“等待”——让食材供应链、切配台和压模机三者之间永远没有空闲。

第一大优化是内存对齐。Tensor Core 要求矩阵的行主序步长（ldm）是 16 字节的倍数，相当于托盘上的每道菜必须按固定间距摆放，否则取餐机器人会拿错。如果你用行主序存储一个 17 列的矩阵，每行末尾要补上 padding 字节，确保每行首地址对齐。

第二大优化是规避共享内存的 bank 冲突。共享内存被划分成 32 个 bank，同一时钟周期内多个线程访问同一个 bank 就会排队。当多个线程同时读取同一行相邻地址时，如果地址恰好映射到同一个 bank，性能骤降。解决办法是在共享内存数组的每行末尾加一个元素的 padding，让访问“错开”。

第三大优化是双缓冲流水线。原理和食堂的“双灶台”一样：一个灶台在做菜，另一个灶台已经在预热。用两个共享内存缓冲区交替加载和计算，cp.async 异步拷贝让数据搬运与矩阵乘加重叠执行，把访存延迟完全隐藏在计算背后。这样流水线的吞吐公式就变成了两个阶段中较慢者的速度，而不是两者之和。

$$T_{\text{总}} \approx \max(T_{\text{加载}}, T_{\text{计算}})$$

1. 加载阶段：共享内存缓冲区 $B_0$ 正在从全局内存装载下一块数据。

$$T_{\text{加载}} = \text{全局内存传输时间}$$

2. 计算阶段：共享内存缓冲区 $B_1$ 中的数据正被 Tensor Core 消耗计算。

$$T_{\text{计算}} = \text{wmma 乘加执行时间}$$

3. 重叠效果：双缓冲让两个阶段并行，总时间由慢者决定。

$$T_{\text{总}} = \max(T_{\text{加载}}, T_{\text{计算}})$$

优化之后，如何确认食堂真的变快了？这就轮到质检员登场了。

调试性能分析

食堂老板想知道“出餐速度到底卡在哪个环节”，光靠感觉不行，得用数据说话。GPU 性能分析工具 Nsight Compute（ncu）就是这位质检员：它会详细记录你的 kernel 在每个硬件单元上的耗时、占用率和吞吐量，告诉你 Tensor Core 的利用率是 90% 还是只有 30%。

使用 ncu 的基本流程是“三问”：第一问，计算吞吐够不够？看 tensor 指令的占用率，如果 Tensor Core 只有 20% 活跃，说明数据供给跟不上，要优化访存或分块策略。第二问，访存效率高不高？看 Memory Throughput 和 L2 Cache Hit Rate，如果全局内存访问频繁且命中率低，说明共享内存复用不够。第三问，SM 占用率合理吗？看 Achieved Occupancy，如果线程块太大或寄存器用太多，SM 上驻留的 warp 数量会减少，延迟隐藏能力变弱。

瓶颈定位有一个经典的“剥洋葱”思路：先用 ncu 的全量报告看大方向（计算密集还是访存密集？），再逐项深入——如果是访存瓶颈，检查是否有未合并访问、bank 冲突和低缓存命中；如果是计算瓶颈，检查 Tensor Core 是否有空转周期。每一步都用数据佐证，而不是凭直觉改代码。

$$\text{利用率} = \frac{\text{实际活跃周期}}{\text{理论最大周期}} \times 100\%$$

1. 实际活跃周期：ncu 统计出的 Tensor Core 真正执行乘加指令的时钟数。

$$T_{\text{active}}$$

2. 理论最大周期：kernel 运行总时长内 Tensor Core 能执行指令的最大时钟数。

$$T_{\text{total}}$$

3. 利用率计算：两者相除得百分数，低于 50% 说明流水线存在明显“饿肚子”。

$$\text{Utilization} = \frac{T_{\text{active}}}{T_{\text{total}}}$$

质检合格后，终于到了出正式菜单的时候。最后我们把所有技巧打包，看看真实工程中如何应用。

工程应用实践

从单道菜的烹饪技巧到一家餐厅的完整菜单，中间还差一步：标准化。CUTLASS 是 NVIDIA 开源的 GEMM 算子库，它把分块、流水线、混合精度这些技巧固化成模板代码，你只需要声明“我要 FP16 输入、FP32 累加、16×16 tile”，就能自动生成一个高性能矩阵乘算子。就像食堂把配餐流程写成标准作业手册，新员工照着做就能出菜。

在 Transformer 架构中，Tensor Core 的身影无处不在：QKV 投影矩阵乘、注意力分数计算、FFN 两层线性变换，每一步都是 GEMM。以 GPT-3 为例，一次前向推理要执行上千次矩阵乘法，这些运算约占整体计算量的 90% 以上。用 Tensor Core 加速这些 GEMM，推理速度能提升数倍，这也是为什么大模型推理框架（如 TensorRT-LLM、vLLM）深度依赖 Tensor Core。

工程中还有两个实战心得：第一，算子融合——把“矩阵乘 + 偏置 + 激活函数”打包成一个 kernel，减少中间结果的显存读写，相当于在食堂里把“配餐+热菜+装盘”合并成一个工位；第二，版本兼容——不同 GPU 架构（Volta/Turing/Ampere/Hopper）的 Tensor Core 指令集和 wmma API 有差异，编写算子时要用 __CUDA_ARCH__ 宏做分支，确保老卡也能运行。

$$\text{GEMM}: \quad C_{(M \times N)} = \sum_{k=1}^{K} A_{(M \times K)} \cdot B_{(K \times N)}$$

1. 输出矩阵：$C$ 的每个元素是 $A$ 的第 $i$ 行与 $B$ 的第 $j$ 列的内积。

$$C_{ij} = \sum_{k=1}^{K} A_{ik} B_{kj}$$

2. 计算规模：总乘加次数为 $M \times N \times K$，当 $M=N=K=4096$ 时超 680 亿次运算。

$$\text{FLOPs} = 2 \times M \times N \times K$$

3. 并行分解：所有 $C_{ij}$ 互不依赖，天然可并行——这正是分块交给不同 SM 执行的硬件基础。

$$C = \bigcup_{\text{tiles}} C_{\text{tile}}$$

从硬件原理到工程实践，Tensor Core 编程的完整链路已经打通。就像一家食堂从采购、切配、压模、出餐到质检的全流程管理，每一步都有章可循。

AI 应用场景

Python 演示

import numpy as np  def wmma_simulate(A, B, C, tile=16):    """用分块三重循环模拟 Tensor Core 的 wmma 乘加流程."""    M, K = A.shape    K2, N = B.shape    assert K == K2, "A 的列数必须等于 B 的行数"    for i in range(0, M, tile):        for j in range(0, N, tile):            acc = C[i:i+tile, j:j+tile].astype(np.float32)            for k in range(0, K, tile):                a_tile = A[i:i+tile, k:k+tile].astype(np.float16)                b_tile = B[k:k+tile, j:j+tile].astype(np.float16)                acc += a_tile.astype(np.float32) @ b_tile.astype(np.float32)            C[i:i+tile, j:j+tile] = acc    return C  def main():    M, N, K = 64, 64, 64    rng = np.random.default_rng(42)    A = rng.standard_normal((M, K)).astype(np.float32)    B = rng.standard_normal((K, N)).astype(np.float32)    C = np.zeros((M, N), dtype=np.float32)     C_out = wmma_simulate(A, B, C, tile=16)     C_ref = A.astype(np.float32) @ B.astype(np.float32)    err = np.max(np.abs(C_out.astype(np.float32) - C_ref))    print("最大绝对误差:", err)    print("分块结果与参考一致:", err < 0.5)  if __name__ == "__main__":    main()

1. for i in range(0, M, tile):外层循环按 tile 步长遍历输出矩阵行方向，对应 CUDA 中线程块在网格中的行划分。
2. acc = C[i:i+tile, j:j+tile].astype(np.float32)取出输出分块并转为 FP32，模拟 Tensor Core 中累加器 fragment 的精度。
3. a_tile = A[i:i+tile, k:k+tile].astype(np.float16)将 A 的分块转为 FP16，模拟 wmma 中以半精度加载输入矩阵。
4. acc += a_tile.astype(np.float32) @ b_tile.astype(np.float32)核心乘加语句：FP16 转 FP32 后相乘累加，模拟 mma_sync 的行为。
5. C[i:i+tile, j:j+tile] = acc将累加结果写回输出矩阵，对应 store_matrix_sync 函数。

常见误区

面试题