深入 NVIDIA GPU：高性能矩阵乘法算子解构（三）

本文是文章：Inside NVIDIA GPUs: Anatomy of high performance matmul kernels 的翻译版。本篇文章翻译将分为四个部分，本文是第三部分。

第一、二部分参考文章：
深入 NVIDIA GPU：高性能矩阵乘法算子解构（一）
深入 NVIDIA GPU：高性能矩阵乘法算子解构（二）

设计近乎 SOTA 的同步矩阵乘法内核

在本章中，我们将解构一个在以下限制条件下接近 SOTA 的 fp32 内核：
无 TMA无异步内存指令无张量核心 (Tensor Cores)仅限 fp32（无 bf16）
换句话说，这是在 Volta 架构之前的 GPU 模型下的 SOTA（在 Volta/Ampere 上也接近 SOTA）：
Volta 引入了张量核心。Ampere 引入了异步内存指令。Hopper 引入了 TMA。
我们将学习的技术称为线程束平铺（warp-tiling）。

在深入研究之前，让我们对之前的内核进行微小的修改，看看会发生什么。具体来说，我们将改变 row 和 col 变量的计算方式。

原始版本：

const int row = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE);

修改版本：

const int row = blockIdx.x * BLOCKSIZE + (threadIdx.x % BLOCKSIZE);

换句话说，我们只是交换了 % 和 / 运算符。

交换 row 和 col 是与前一示例相比在逻辑结构上唯一的改变：

图 24：row 和 col 变量的新逻辑组织
以下是修改后的内核现在的表现：

图 25：具有非合并（uncoalesced）GMEM 访问的朴素内核

这个看似无害的微调使得我们的 GMEM 访问变得非合并。

在我的 H100 PCIe 卡上，性能从 3171 GFLOP/s 骤降至仅 243 GFLOP/s——慢了 13 倍。这正是我们之前在 GMEM 章节中看到的惩罚（Stephen Jones 的跨步 GMEM 访问实验）。

从外部看，这只是两个运算符之间微不足道的交换。但如果你没有硬件的认知模型，你永远不会预料到如此戏剧性的影响。

图 26：屋顶线模型（Roofline Model）

观察屋顶线模型，你可以看到我们的内核深陷于图中的内存带宽受限（memory-bandwidth-bound）区域。我们为算力付给 NVIDIA 大笔资金，所以我们理应瞄准计算受限（compute-bound）区域。

📝 屋顶线模型 (Roofline Model)

屋顶线模型在 y 轴上绘制性能 (FLOP/s)，在 x 轴上绘制算术强度 (Arithmetic Intensity, AI)。

算术强度定义为：每从设备内存/GMEM 加载一个字节所执行的浮点运算次数（默认情况下）。

“脊点”（ridge point）出现在：峰值性能 / GMEM 带宽。对于我的 H100 PCIe，这个数值大约是 410。只有当算术强度超过这个值时，内核才能进入计算受限状态。
在继续之前，让我们重新审视一下串行矩阵乘法代码。供参考：for (int m = 0; m < M; m++) {我想在这里强调的关键点是：语义对循环顺序是不变量。换句话说，我们可以将这三个嵌套循环以 3! = 6 种方式中的任何一种进行置换，结果仍然是一个正确的矩阵乘法。

在这六种置换中，最有趣的是将 K 放在最外层的顺序。（m 和 n 的相对顺序较不重要，所以让我们假设“规范的” m-n 顺序）：for (int k = 0; k < K; k++) {如果这些加载来自 GMEM，通过将 A 的加载次数从N^3 减少到 N^2，我们刚刚节省了大约 2x 的带宽。

但更重要的洞察是算法层面的：这个版本将矩阵乘法计算为外积的偏部分和（partial sum of outer products）。这种视角对于理解我们接下来要深入探讨的线程束平铺（warp-tiling）方法至关重要。

图 27：矩阵乘法作为部分外积之和

这可能显而易见，但值得强调：一个点积等同于多个部分点积之和：

图 28：点积等同于部分点积之和

这很重要，因为它允许我们将计算分解为一系列块矩阵乘法（block matmuls）（每个块产生部分点积）。通过在执行计算之前将这些块移动到 SMEM 中，我们可以减少 GMEM 流量并显著提高速度。

如果不进行分块（chunking），我们根本无法将其放入 SMEM 内部。

还请回想一下，我们最初的内核算术强度非常低——它们在加载每个字节时完成的工作很少。为了改进这一点，我们需要：

1. 每个线程计算多个输出元素。2. 使输出分块（tiles）尽可能接近正方形。这里有一个视觉直觉，解释了为什么这很重要：

图 29：当每个线程计算多个输出且分块接近正方形时，算术强度会提高

至此，我们已经收集了理解线程束平铺（warp-tiling）所需的大部分拼图。让我们把它们拼在一起。

我们知道两件关键的事：输出分块应该是正方形的（以最大化算术强度）。计算应该分解为子步骤，以便中间块可以放入 SMEM。
考虑到这一点，算法的高层结构如下所示：

图 30：线程束平铺算法的高层结构，也称为块平铺（block tiling）

参考代码在这里：
https://github.com/siboehm/SGEMM_CUDA/blob/master/src/kernels/10_kernel_warptiling.cuh。我建议先从我的图表开始，然后打开代码将所有要点连接起来。

📝 注意：
我将使用与 Simon 博客文章中相同的分块大小（未针对我的 H100 进行自动调优）：
Bm = Bn = 128, Bk = 16
由于每个块的计算是独立的——而且我们已经确信部分点积可以累加为完整的点积——我们只需要关注单个块的单个步骤。其余部分（另外 1023 个块，4096/128 4096/128 = 3232 = 1024 总计）将遵循相同的逻辑。

📝 给自己的笔记：
出于某种原因，我很难忽略其他的块。所以，念咒语时间：“其他一切都是正确的；我只需要专注于下一步。局部正确性导致全局正确性。” :)

带着这种心态，让我们放大到蓝色块的第一步（红箭头转换前的计算），它对应于输出分块 C[0,0]（注意是分块，而不是单个元素）。

矩阵 A的分块维度为 Bm × Bk，矩阵 $B$ 的分块维度为 Bk × Bn。这些数据被加载到 SMEM 缓冲区 As 和 Bs 中。

加载/存储 B \to Bs 是很直接的，因为 Bs 没有经过转置。4 个线程束中的每一个都从 GMEM 抓取 B 的一行，每个线程发布一次向量化加载（LDG.128），随后执行一次向量化存储（STS.128）。每个线程束循环 4 次，步长为 4 行。

对应代码（我增加了注释并删除了 Simon 注释掉的代码）：

for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {

图 31：将 B 的分块（GMEM）加载到 Bs（SMEM）中

加载 A \to As。这一步更棘手，因为 As 是经过转置的。转置的原因是它允许在随后的计算阶段进行向量化加载（LDS.128）。

权衡之处在于存储无法向量化：从 A 的一行中提取的 4 个浮点数现在必须离散地存入 As 的一列中，而这一列映射到了同一个内存银行（bank）。这是可以接受的，因为我们优先考虑快速加载——在计算过程中，As 的每个元素会被多次访问，而存储仅发生一次。

图中的 innerRowX 和 innerColX 注解准确展示了每个线程负责的工作。

对应代码：

for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {

图 32：将 A 的分块（GMEM）加载到 As（SMEM）中(1)(2)

加载完成后，我们同步线程块（__syncthreads()），以确保所有数据在 As 和 Bs 中均已就绪。

现在进入计算阶段。

对应代码（建议扫视一下代码，并在代码与绘图之间进行几次对照阅读）：

for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {

图 33：将 As 和 Bs 之间的矩阵乘法执行为一系列线程级外积（线程束平铺 + 线程平铺）

一旦分块处理完毕，我们再次同步。这可以防止竞争条件——如果没有它，一些线程可能会开始将下一个分块写入 As 和 Bs，而其他线程仍在处理当前分块。

同步后，我们将 A 和 B 的指针推进 Bk 距离，算法重复执行直到所有分块处理完毕。

A += BK;     // 将 BK 列向右移动

最后，一旦循环完成，128 个线程将其私有的 threadResults 寄存器刷入矩阵 C 对应的输出分块中（此时该分块已包含完整的点积结果！）。

在实践中，你会针对特定的 GPU 对该算法的参数进行自动调优。但正如前面指出的，这种风格的内核已不再是首选方法——现代 GPU 拥有异步内存机制和张量核心（Tensor Cores），能将性能推向远超单靠线程束平铺所能达到的水平。

接下来，让我们转向 Hopper 上的真正 SOTA。

📝 下一章的补充阅读：
我强烈推荐 Pranjal 的优秀博文 [15]，它读起来更像是一份工作日志。在本章中，我将遵循他日志中的内核。与 Simon 的工作一样，大部分代码似乎也受到了 CUTLASS 的启发（例如这些帖子：CUTLASS ping pong 内核 [16] 和高效 GEMM）。

值得注意的是，细节决定成败，Pranjal 成功超越了 cuBLAS SOTA——在一些目标矩阵维度上达到了 cuBLAS 性能的约 107%。

达坦科技始终致力于打造高性能AI+Cloud基础设施平台，积极推动AI应用的落地。达坦科技通过软硬件深度融合的方式，提供AI推理引擎和高性能网络，为AI应用提供弹性、便利、经济的基础设施服务，以此满足不同行业客户对AI+Cloud的需求。

公众号：达坦科技DatenLord
DatenLord官网：https://datenlord.github.io/zh-cn/
知乎账号：https://www.zhihu.com/org/da-tan-ke-ji
B站：https://space.bilibili.com/2017027518
邮箱：info@datenlord.com

如果您有兴趣加入达坦科技Rust前沿技术交流群、硬件敏捷开发和验证方法学讨论群或AI Infra 交流群，请添加小助手微信：DatenLord_Tech