深度长文：“深入NVIDIA GPU：高性能矩阵乘法（matmul）内核剖析”

Bosch

转自微博用户@蚁工厂

————————————————————————————

深度长文：“深入NVIDIA GPU：高性能矩阵乘法（matmul）内核剖析”

网页链接

作者是前微软和DeepMind软件/机器学习工程师Aleksa Gordić.
下面是作者写的推荐语：
如果你想深入了解如何在CUDA中编写最顶尖的（state-of-the-art）矩阵乘法内核，请继续读下去。
（请记住，matmul 是Transformer模型在训练和推理过程中执行的唯一且最重要的操作。NVIDIA绝大多数的算力都消耗在这上面。效率上每提升1%，就能节省下相当于好几座核反应堆的能源 ）

我再次发现，我低估了写这篇文章所需的工作量。这简直又是一本小册子了（笑）。足足有47张图！

我在这篇文章中涵盖了以下内容：

1️⃣GPU架构的基础知识，重点是内存层次结构。我为你建立了关于GMEM、SMEM和L1/L2缓存的心智模型，并将其与CUDA编程模型联系起来。在此过程中，我们还探讨了GPU的“光速”（理论性能上限）以及它如何受到功耗的限制——让硬件的真实情况融入到我们的模型中。

2️⃣PTX/SASS（GPU汇编），以及如何引导编译器生成我们真正想要的代码（比如，循环展开了吗？我们用上像 LDG.128 这样的向量化加载指令了吗？等等）。我为一个简单的 matmul 内核，极其详尽地注释了一个PTX/SASS示例。即使你是编译器新手，也应该会觉得这部分很有用。
    （我其实在两个编译器里都发现了不少效率低下的地方——真有趣！）

3️⃣许多核心概念，例如分块/波次量化（tile/wave quantization）、占用率（occupancy）、指令级并行（ILP）、Roofline模型等。同时，还帮助你建立对一些基本等价关系的直觉：比如点积可以看作是部分外积的和，以及为什么方形的分块（tile）是实现高算术强度的理想形状等。

4️⃣Warp Tiling（线程束平铺）方法——在假设你不能使用Tensor Cores、TMA、异步内存指令和bf16数据类型的前提下，这是一种近乎SOTA的技术。它仅仅依靠CUDA Cores、寄存器和共享内存，就将GPU的性能发挥到极致。

5️⃣最后，我们进入了Hopper（H100）架构的世界：TMA（张量内存加速器）、Swizzling（地址交错技术）、Tensor Cores和wgmma指令、异步加载/存储流水线、像希尔伯特曲线这样的调度策略、利用TMA多播的集群（Clusters）、更快的PTX屏障（barriers）等等。

和往常一样，文章里有大量的示例和丰富的图解。这是我第一次看到 warp tiling 内核时，能脱口而出“哦，我完全懂了”的时刻。我只是需要把我脑海中的图像变成一张真正的图片。

几年前，我深受Si_Boehm 的一篇关于 matmul 工作原理的精彩博客的启发，但我也发现其中存在一些错误和不清晰的解释，而且内容也有些过时了。在pranjalssh 的杰出工作（他为H100构建了SOTA内核）和我自己的研究基础上，我完成了这篇文章。