在 H100 上超越 cuBLAS：一份工作日志

CUDA 矩阵乘法内核 - 从零开始

作者：Pranjal Shankhdhar | 2024年11月29日

来源： cudaforfun.substack.com

---

在这篇文章中，我们将在最新一代 NVIDIA 硬件 H100 上，迭代地实现一个 CUDA 矩阵乘法内核。

我们将深入理解 H100 的架构，并一步步展示这些优化。最终内核在 N=4096 时比 cuBLAS 快 7%。它可以装进一个 C++ 文件中，没有任何依赖。

这篇文章可以看作是 Simon 的经典博客的续篇，该博客展示了在 A6000 GPU 上的类似优化。然而 H100 GPU 是完全不同的怪兽，需要完全不同的算法。举个例子，Simon 博客中的算法在 H100 上只能达到 cuBLAS 性能的 4%。在这篇文章中，我们将从 Simon 的博客出发，迭代地达到 cuBLAS 的 107%。

我所有的代码都可以在 Github 上获得。

我们的目标不是成为 cuBLAS 的替代品，而是设计一个略快但更简洁的矩阵乘法内核，适用于一般较大的矩阵。cuBLAS 在各种矩阵尺寸上表现良好，比如 k 维度很大的小矩阵或矩阵-向量乘法（与 LLM 推理相关）。

Simon 博客的快速回顾

让我们回顾一下 Simon 介绍的矩阵乘法算法的基本结构。我们计算 C[m, n] = A[m, k] x B[k, n]，如下图所示：

我们将大矩阵乘法分解为计算多个输出分块（tile）。每个分块大小为 BM x BN —— 代表输出矩阵 C 的一部分。我们分配一个线程块来计算这个分块 —— 该线程块有多达 1024 个线程协同工作。

为了计算这个分块中的所有输出，我们需要从 A 中读取 BM x K 的行块（蓝色）和从 B 中读取 K x BN 的列块（绿色）。这些值被多次访问，所以我们需要将它们存储在 SMEM 中以提高性能。

然而，这些块太大了，无法存储在 SMEM 中。所以我们将它们按 BK 大小分块存储。对于每个块，我们可以将 BM x BK 和 BK x BN 矩阵相乘，得到输出分块的 BM x BN 矩阵。回忆一下朴素矩阵乘法：

由于这些块是在 k 维度上的，我们只需将所有这些矩阵求和就能计算出 BM x BN 输出分块的最终值。输出分块中的所有值都存储在寄存器中，所以累加很容易。

我们的 H100 矩阵乘法内核将遵循这种通过乘以较小矩阵块来计算输出分块的结构。Simon 的博客进一步充分利用寄存器空间，将块的部分从 SMEM 移动到寄存器中。我们不会在这里详述这些细节，因为我们的内核不使用它们。

实验设置

在博客的其余部分，我们将考虑方阵大小为 (M=N=K=4096) 的 bfloat16 类型大矩阵。bfloat16 是一种专门的 16 位数据类型，用于近期的深度学习应用。对于我们的内核性能而言，它与普通的 fp16 没有什么不同。矩阵 B 和 C 以列主序存储，而 A 以行主序存储。这是矩阵乘法基准测试的常见设置。

我们使用 mean = 0 和 std_dev = 1 的正态分布来初始化矩阵。事实证明，这是性能测量的最佳分布。

对于测量 FLOPS，我们取 8 次运行的平均时间（忽略第一次预热运行）。然后 FLOPS 通过 2 * m * n * k / time 计算。

所有基准测试都在 H100 SXM 上运行，使用 CUDA 工具包 12.6，V12.6.68

H100 的内部世界

让我们了解一些 H100 的规格参数，以理解这款 GPU 的新特性。
H100 有两个变体：PCIe 和 SXM。它们非常相似，只是 SXM 变体略快。我的机器使用 H100 SXM，其规格如下：

• 132 个流式多处理器（SM）
• 每个 SM 1024 个线程
• 每个 SM 4 个张量核心
• 80GB 高带宽内存（3.35TB/s）
• 每个 SM 256KB 共享内存 + L1 缓存组合
• 每个 SM 65,536 个寄存器
• 50MB L2 缓存，所有 SM 共享

这些术语大多在之前的博客中已经熟悉了。与前几代相比，H100 GPU 有更多的 SM、更快的全局内存、更快的时钟速度、更多的共享内存以及更大更快的 L2 缓存。矩阵乘法内核使用所有这些特性 —— 所以我们可以预期旧一代的算法在 H100 上自然会运行得更快。我们确实看到了从 A100 上的 21 TFLOPs → H100 上的 32 TFLOPs 的跳跃。

我们距离接近 cuBLAS（716 TFLOPs）显然还有很长的路要走。关键在于一个我们以前没有见过的新规格：

张量核心

张量核心是 GPU 中的一个特殊硬件单元，能在单条硬件指令中执行小型矩阵-矩阵乘法。它有多种形式 —— mma、wmma 和 wgmma 指令。在这篇博客中，我们将研究 Hopper 架构引入的 wgmma 指令。

不幸的是，CUDA C++ 指南中没有这些指令的文档，我们必须查看 PTX 指南。让我们看一个示例指令：

wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16

这执行一个矩阵乘法操作 C = A*B + C，其中 m=64, n=16, k=16。

• A：存储在共享内存中的 mxk bfloat16 类型矩阵。
• B：存储在共享内存中的 kxn bfloat16 类型矩阵。
• C：存储在寄存器中的 mxn 32 位浮点类型矩阵。

存储 A 和 B 需要 (64*16 + 16*16) * 2 字节 = 2.5KB 共享内存。

存储 C 需要 64*16 = 1024 个寄存器。

注意，单个 GPU 线程最多只能存储 256 个寄存器 —— 因此张量核心指令需要将 C 分布存储在 SM 中的 128 个线程上！注意一个 warp = 32 个线程，所以 128 个线程将包含 4 个 warp。在 Hopper 架构中，一组 4 个 warp 称为一个 warp-group（线程组）。当我们将 C 分布在一个 warp-group 上时，每个线程需要 1024/128 = 8 个寄存器 —— 这是一个合理得多的数字。注意这个指令叫做 wgmma，代表 warp-group-matrix-multiply-add。

异步性

注意张量核心指令中的 mma_async 术语。这些指令在每个 SM 的 4 个张量核心上异步运行。连续的张量核心指令可以批量组合并发送到张量核心，并行运行。这对充分利用所有张量核心至关重要。

指令尺寸

H100 提供多种不同大小的矩阵乘法指令。来自 PTX 指南：

.shape   = {.m64n8k16, .m64n16k16, .m64n24k16, .m64n32k16,
                                 ...
                                 ...
            .m64n232k16, .m64n240k16, .m64n248k16, .m64n256k16};

在所有这些指令中，m=64 和 k=16 保持不变。n 可以从 8 变化到 256。根据我的经验，使用单个较大 n 的指令比使用多个较小 n 的指令更快。然而，请注意较大的 n 使用更多资源。n=256 需要惊人的 40KB SMEM 和每线程 128 个寄存器！

共享内存布局（Swizzling）

张量核心指令需要 SMEM 中非常特定的块布局 —— 不是简单的行主序或列主序。这种布局经过大量 swizzling，太复杂而无法手工加载。Nvidia 实现了 swizzling 以避免共享内存 bank 冲突。此外，内存布局的文档记录是不正确的。

内核演进

内核 1：Simon 的博客

Simon 的算法是为 FP32 类型设计的。将其适配到 bfloat16 类型给我们 32 TFLOPs。

注意这不是一个公平的比较，因为 cuBLAS 利用了 bfloat16 的张量核心操作，而这对 FP32 是不可用的。在博客文章中，Simon 声称这可以将性能提高 3.5 倍，但没能实现。我们将从他停下的地方继续，开始使用它们。

---

内核 2：使用张量核心指令

我们现在准备编写一个使用张量核心指令计算输出分块的简单内核。这一节比我预想的要长一些，但它建立了我们在后续内核中需要的核心概念。

我们将为每个输出分块分配一个线程块，其中有 128 个线程协作执行张量核心指令。这与 Simon 博客中的 Kernel 5 非常相似。我们将简单地用张量核心指令替换手写的 blocktile 矩阵乘法。

让我们使用 WGMMA_M=64, WGMMA_N=64, WGMMA_K=16 的符号来表示 wgmma 操作的大小。为简单起见，让我们将块大小与 wgmma 大小匹配，使内核更简单。以下是整体内核结构：

注意我们遵循之前讨论的相同内核结构：

• 我们沿 K 维度按大小为 BK 的块循环。对于每个块：
  • 我们将 A 和 B 的对应块加载到 SMEM 中
  • 我们使用张量核心指令将这些块相乘并将结果存储在寄存器中
• 所有块处理完毕后，我们将寄存器中的值写入 C 中对应的分块。

注意我们跳过了加载和存储块的代码。让我们先看加载部分。

张量核心指令需要 SMEM 中非常特定的块布局 —— 不是简单的行主序或列主序。这种布局经过大量 swizzling，太复杂而无法手工加载。Nvidia 实现了 swizzling 以避免共享内存 bank 冲突。此外，内存布局的文档记录是不正确的。幸运的是，Nvidia 提供了一种开箱即用的方式来加载分块，而不用担心这些布局：张量内存加速器（TMA）。

使用张量内存加速器（TMA）加载

TMA 是 Hopper 架构引入的新硬件。它是一种更快的方式来在 GMEM 和 SMEM 之间加载多维矩阵的分块。这在一个独立的硬件单元中实现，使其比自定义加载快得多。TMA 加载直接支持张量核心所需的 swizzling 模式。

TMA 接受矩阵的分块配置，可以将任何请求的分块加载到 SMEM 中。

TMA 加载与之前的不同之处在于，它需要从单个线程调用。以前，多个 CUDA 线程协作加载一块内存。使用 TMA，单个线程可以发出 TMA 调用，所有线程等待它完成。以下示例取自 CUDA 编程指南，可用于将 A 的一个分块加载到 SMEM 中。它使用 cuda 屏障来等待加载完成。

存储输出分块

输出分块的值分布存储在线程块中的 128 个线程上。可以计算从线程 id、寄存器索引 → 对应全局内存地址的映射。

一旦我们有了映射，就可以将所有寄存器值存储到 GMEM。映射函数没有什么特别的，我不建议读者深入研究它。我们只需要知道它是一个简单的算术映射，可以在需要时计算。

性能

我们达到了 317 TFLOPs 的吞吐量，比上一个内核的 32 TFLOPs 有了巨大的飞跃。我们在这一节中还引入了几个新特性：张量核心、TMA 和 CUDA 屏障。所有这些协同工作，给我们带来了 10 倍的性能提升！

张量核心确实蕴含着强大的力量。注意 A6000 GPU 也有张量核心，但可以在不使用它们的情况下达到 92% 的吞吐量。这在 H100 上不成立，在 H100 上张量核心对于高吞吐量是必需的。

---

内核 3：处理更大的输出分块

我们可以将分块大小增加到超出 wgmma 指令维度。我们不必将 BM、BN 匹配到 wgmma 维度，而是可以将更大的分块分解为多个较小的 wgmma 操作。我们简单地在 M、N、K 维度上循环，执行适当大小块的常规矩阵乘法。分块变大后，从全局内存搬进共享内存的矩阵数据就能被参与到更多的交叉相乘中，从而摊薄了内存读取的时间成本。

下面这张图中A, B都是BM * BK和BN * BK的模式进行排列的。

由于wgmma一次性会取多行，例如16*16，所以需要swizzling

性能

使用 BM=128, BN=128, BK=64 配合 m64n128k16 wgmma 指令，达到 423 TFLOPs。一个关键发现：选择最大的可用指令并设置 BN = WGMMA_N 始终比其他方案性能更好。

---

性能分析

让我们对内核进行性能分析，了解时间花在了哪里。

内核执行分为三个阶段：加载、张量核心操作和存储。

加载和张量核心操作在跨越 k 维度的循环中执行。所有值在计算完成后存储到输出矩阵。

每个线程的时钟周期测量：

加载: 1415   张量核心: 703   存储: 4572

张量核心操作比加载快 2 倍。存储操作比加载慢 6.4 倍，但只运行一次；加载/计算循环运行 128 次。这些指标根据分块配置而波动，但量化它们可以揭示瓶颈并指导后续优化。

---

内核 4：隐藏加载延迟

注意加载和张量核心操作是独立的，可以并行执行！这是生产者-消费者问题，其中生产者（加载）和消费者（张量核心）可以独立运行。

我们将使用 CUDA 编程指南中的 "Warp 特化" 概念。我们在单个线程块内启动 2 个 warp-group。第一个 warp-group 将作为生产者，将块加载到共享内存中。第二个 warp-group 将作为消费者，对加载的块使用张量核心指令。它们将使用屏障和循环缓冲区作为共享队列进行通信。

我们可以如下设置队列所需的共享数据结构。生产者在加载分块之前，调用 empty[i].wait() 来验证缓冲区槽是否可用。加载后，full[i].arrive() 发出数据就绪信号。消费者镜像此操作：full[i].wait() 确认数据存在，然后 empty[i].arrive() 发出槽位可以重用的信号。

生产者初始化：

初始屏障状态假设共享缓冲区开始时是空的，使生产者能够立即开始。

生产者循环：

消费者循环：

循环缓冲区存储多个分块。生产者持续加载到位置 0、1、2 等，然后环绕。消费者按顺序从相同位置读取。当队列有缓冲项时，两个操作都可以并行进行。

QSIZE=2 时的时序图：

性能

使用 128 x 128 分块大小和 QSIZE=5，达到 498 TFLOPs。

---

内核 5：挑战分块大小极限

分块大小越大，SMEM复用率越高，arithmetic intensity越高。

让我们尝试将分块大小增加到 128 x 256。我们遇到了编译器警告：

ptxas info : (C7511) Potential Performance Loss: wgmma.mma_async
instructions are serialized due to insufficient register resources

性能下降了 5x，降到 123 TFLOPs！

输出分块需要 128 x 256 = 32768 个寄存器分布在线程块中 —— 只有 SM 总寄存器的 50%。然而，检查每线程使用量揭示了实际问题：128 个线程组成的 warpgroup 中，每个线程需要 256 个输出寄存器，这是 H100 每线程的最大硬件限制（于是每个thread就没有其他寄存器去做基本运算了）。

当线程耗尽寄存器时，编译器执行"寄存器溢出" —— 临时将寄存器存储到内存中。这个操作发生在张量核心指令之间，导致序列化，阻止了正常情况下的批处理优化。

两个消费者 warp-group

该策略保持总体寄存器使用量不变，但将工作分配到更多线程上。不是一个 warpgroup（128 个线程）处理 128 x 256 分块，而是两个 warpgroup（总共 256 个线程）分工：每个处理一个 64 x 256 分块。

每线程寄存器使用量从 256 降到 128，同时保持相同的总消耗。内核现在用 128 * 3 个线程启动：一个生产者 warpgroup 和两个消费者 warpgroup。

两个消费者处理不同的分块部分，但在具有更高令牌数的相同屏障上同步。

性能

这产生了 610 TFLOPs。更大的分块大小增加了共享内存需求，迫使 QSIZE 从 5 减少到 3，但整体吞吐量提高了。

每 warpgroup 寄存器规格

性能分析显示每个线程使用 168 个寄存器。块总使用量：64512 个寄存器，刚好低于 65536 的限制。然而，生产者线程不需要张量操作，所以它们需要的寄存器比消费者少。

使用 PTX 级别的寄存器规格允许为每个 warpgroup 分配不同的数量。我们将消费者设置为每线程使用 240 个寄存器，生产者为每线程 24 个寄存器，总计相同的 64512，但优化了分配。

这个优化将性能提升到 631 TFLOPs。假设是消费者中更大的寄存器数量减少了寄存器 bank 冲突，不过对其底层机制还不确定。

---

内核 6：隐藏存储延迟

我们通过分离生产者和消费者成功地隐藏了加载延迟。让我们看看如何隐藏存储延迟。

一个 SM 在内核的整个生命周期中处理多个输出分块。对于第一个分块，我们看到加载和张量核心操作是并行的。最后，我们将所有计算值存储到 C 矩阵。在这段时间里，我们也可以开始为下一个输出分块加载块。注意存储和加载操作不使用任何共同资源。加载存储到 SMEM 中，而存储是从 RMEM 到 GMEM。

根据我们的性能分析，大约 4572 个周期花在将值存储到 GMEM 上。如果我们在这段时间里开始为下一个线程块加载块，我们可以加载 4572 / 1415 = 3.2 个块。这意味着，下一个线程块的消费者可以在当前线程块完成后立即开始运行！

为了实现这一点，我们以与 SM 数量相同的线程块启动内核 —— H100 为 132。

现在，我们需要决定哪些分块分配给哪个 SM。以前，我们为每个线程块分配 1 个分块，让 GPU 将它们调度到不同的 SM 上。现在，我们需要自己做这个调度。让我们遵循一个简单的调度逻辑，将连续的输出分块分配给一个 SM：

我们不需要太多额外逻辑来跨分块重叠存储和加载。处理新分块时，我们将重用屏障和共享队列，而不是重新初始化它们。一旦生产者完成一个分块的块加载，它将立即开始为下一个分块加载块。消费者也会知道它们何时完成了一个分块的处理，可以开始从共享队列的下一个位置读取下一个分块。

性能

使用这个策略我们看到 400 TFLOPs，这比之前的 640 TFLOPs 是一个退步！这没有按我们计划的那样好。我们的重叠存储逻辑相当合理，让我们看看是不是调度逻辑出了问题。

调度与 L2 缓存

与其看单个 SM 处理的分块，不如看看各个 SM 处理的第一个分块。这些分块将在同一时间被处理。

我们看到 SM 同时处理距离很远的分块。这意味着同时从 A 和 B 矩阵加载非常不同的内存块。如果我们能同时调度附近的分块，那么它们加载的值将有大量 A 和 B 矩阵的共同部分。这些共同部分将被缓存在 GPU 的 L2 缓存中 —— 这意味着我们不必总是从 GMEM 加载分块！让我们看看这种调度是什么样的：

对比更好的调度：

图示不按比例，4 x 4 区域对应 128 个分块（16 x 8）

注意相同颜色的分块在同一时间被调度。这意味着同时对 A/B 矩阵有大量共同访问 —— 全部由 L2 缓存服务！

注意我们在图中只使用了 128 个 SM，因为它很容易分成 16 x 8 的配置。让一切都是 2 的幂使我们的调度逻辑简单得多。

性能

660 TFLOPs。我们达到了 83% 的 L2 缓存命中率，而 cuBLAS 只有 70% 的 L2 缓存命中率。

修改逻辑使用全部 132 个 SM 并不困难。我们仍然可以保持这个配置，但将下一组中的一些分块分配给剩余的 SM。我们一直这样做直到遍历所有分块组。此外，我们的分块配置不必是 16x8，也可以像 2x2 这样小。

在尝试了几种分块组配置后，我发现使用 132 个 SM 比 128 个 SM 慢 (655 TFLOPs)。这是因为我们的分块数量可以整齐地划分为 16 x 8 的区域 —— 如果我们使用 128 个 SM，可以获得更好的 L2 缓存命中。

---

内核 7：更快的屏障

注意我们当前的屏障实现是 CUDA 编程指南推荐的。事实证明，存在一种更快的屏障实现，可以显著加速我们的内核。这种实现只在 PTX 指南中被引用，没有任何 CUDA API。至于哪个更好，就留给读者作为练习吧 :)

让我们列出两种屏障 API，并开始使用新的那个！

CUDA 屏障 API：

PTX 屏障 API：

API 有 2 个区别：

• 阶段变量： 我们手动跟踪阶段变量，它是我们在屏障上调用 wait 次数的奇偶性。阶段变量没有其他意义。底层 API 要求我们手动跟踪并在 API 中传递它。这是一个抽象泄漏，可能是性能原因所需。

  • 注意，一旦 wait 调用完成，我们不需要重新初始化屏障。我们可以简单地重用它，就好像它用之前的值重新初始化了一样。屏障通常被重用数百次，因为我们在共享队列中加载数百个分块。

• 令牌： 另一个区别是这个 API 在 arrive 和 wait 调用中不使用任何令牌。这使实现更简洁，并允许我们进一步优化同步。这意味着并非所有执行 wait 的线程都需要先调用 arrive。我们可以将令牌同步的数量从 257 减少到 3（每个生产者和消费者各一个）。使用更少的同步使代码更快：

注意新的 API 需要在 PTX 中实现。这些是 PTX 代码上的简单 CUDA 包装器，在 github 代码中有高亮。

性能

新的屏障 API 带来了 10% 的性能提升，使我们达到 704 TFLOPs。我们现在已经达到了 cuBLAS 性能的 98%。剩余的优化将带来更小的回报，但会慢慢将我们提升到 cuBLAS 及以上。

---

内核 8：线程块集群

集群（Cluster）是 Hopper 的新特性，它将多个线程块分组，同时在多个 SM 上并发运行。集群中的多个 SM 可以同步，并协作地获取和交换数据。

要使用这个特性，我们需要在内核函数定义中声明：

// 这将启动一个每个集群有 2 个 SM 的内核。

__global__ void __cluster_dims__(2, 1, 1) kernel(...) {
  // ... 内核代码
}

TMA 多播

集群中的多个 SM 可以使用 TMA 多播操作加载相同的分块。这比从 L2 缓存加载分块两次更快。附近的分块读取输入矩阵的相同块，使这个特性非常有用。

上图显示了当 2 个垂直相邻的分块在同一集群的不同 SM 上运行时的情况。它们需要从 A 加载 2 个不同的块，但从 B 加载相同的块。B 的这个块可以多播到集群中的 SM。TMA 支持这个功能。

TMA 多播操作是一条 PTX 指令。和其他情况一样，这并不复杂，但缺少 CUDA 中的包装函数。

cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes.multicast::cluster

为了使用这个，我们还需要跨集群中不同 SM 同步屏障。PTX 屏障通过在 arrive 函数后附加 cluster 关键字来提供这个功能。我们在 github 代码中提供了两种方法的包装器。

性能

这使我们达到 734 TFLOPs。我们现在比 cuBLAS 做得稍好，达到 cuBLAS 性能的 102%。

注意可以以不同方式分组集群（水平分块），甚至使用大小为 4 的集群，将 2x2 的分块聚在一起。我们的实现支持所有集群形状，但我们发现垂直聚类 2 个分块是最快的。这弥补了我们不均匀的分块大小（128 x 256）。使用更大的集群大小要慢得多，可能是由于昂贵的跨 SM 同步。

---

内核 9：微优化

这个内核包含一系列小优化。

• 重排存储：
  • 我们将多个寄存器的值写入 GMEM。我们可以按一定顺序排列，使连续写入映射到附近的内存位置。这带来了略好的性能。

• 写入时跳过 L1/L2 缓存：
  • 我们可以使用缓存提示直接将值存储到 GMEM，跳过 L1/L2 缓存 —— 为 A、B 矩阵释放略多的空间。使用 CUDA 提供的 __stwt() 方法。

• 跳过将寄存器重置为 0：
  • 记住张量核心操作会累加值，因此我们需要在处理不同分块之间将寄存器重置为 0。
  • 如果我们查看张量核心规格，可以设置一个标志来控制张量核心操作是否进行累加。这将张量核心操作在 C = A*B 和 C = A*B+C 之间切换。
  • 我们可以在第一次为输出分块使用张量核心指令时设置这个标志。这帮助我们避免每次处理输出分块时都将寄存器重置为 0。

性能

这些优化总共帮助我们从 734 TFLOPs 达到 747 TFLOPs。我们已经开始看到优化的收益递减，但这不会阻止我们。

---

内核 10：异步存储

我们已经花了一些时间优化存储操作的性能，但还有另一种方式可以达到类似的效果。我们可以将寄存器值存储到 SMEM，然后使用 TMA 异步地将这些值存储到 GMEM！

唯一的注意事项是我们的共享队列可用的 SMEM 空间会减少。很难判断这是否更好，让我们试试看吧！

性能

这使我们达到 758 TFLOPs，又提升了 2%。此时，我们的想法快用完了，所以让我们请出一些大招。

---

内核 11：希尔伯特曲线

让我们重新审视 SM 上输出分块的调度。我们将相同颜色的分块同时调度到 SM 上。这次我们按分块在 SM 上运行的顺序编号。注意我们不会显式等待所有 SM 处理完分配的分块才调度下一组。这自然发生，因为我们假设 SM 处理分块所需的时间相似。

注意，虽然我们在同一分块组内看到大量的 L2 缓存命中，但我们的调度在分块组之间并不是最优的。我们按顺序运行分块：蓝、绿、灰、红。绿和灰分块不会共享 A/B 的任何共同块。我们可以通过交换红和灰分块的调度来修复这一点。

对大矩阵实现这一点会变得非常复杂。幸运的是，按空间顺序填充矩阵是一个研究成熟的问题 —— 答案就是希尔伯特曲线。

希尔伯特曲线是一种空间填充曲线，它覆盖矩阵的所有单元格，同时确保"附近"的单元格被一起访问。如果我们取它的任何一段，我们会发现覆盖的所有单元格在空间上都是接近的。这给了我们一个新的调度算法。在 [M/BM, N/BN] 矩阵上创建希尔伯特曲线，并使用这个顺序来调度分块。连续的分块将在同一时间被调度。

性能

这给我们带来了 1% 的提升，达到 764 TFLOPs。我们已经走了很远，达到了 cuBLAS 性能的 107%。这是一个停下来总结我们想法的好时机。

---

结论

我们的内核在不同的 N 值下性能各异：

• 对于 N=512 快 2%
• 对于 N=1024 快 17%
• 对于 N=2048,4096 快 7-8%
• 对于 N=8192 快 1.5%

对于小 N，矩阵乘法内核受内存带宽限制。这导致改进空间很小。

对于非常大的 N 值，矩阵乘法内核变成功耗受限！H100 GPU 的最大功率上限为 700W —— 不足以同时使用所有张量核心。这导致非常大的 N 值出现收益递减。

注意我们并非在所有 N 值上都更快 —— 在不同的 N 值上，有时更慢，有时更快，是混合的。然而，我相信通过广泛地自动调优内核参数，可以达到持平。

还可以通过调整 GPU 设置来提高性能，将功率从 L2 缓存转移到张量核心。这应该会在 cuBLAS 和我们的内核上都带来性能提升。

我所有的代码都可以在 Github 上获得。

我最近开始将编写 GPU 内核作为一项爱好 —— 希望能做更多 :)

参考资料

以下是一些帮助我学习 GPU 编程的资源：

• 《大规模并行处理器编程》视频讲座
• Simon 的从零开始的矩阵乘法博客。
• GPU Mode Discord 群组
• Nvidia 的 Hopper 白皮书，包含新 Hopper 架构的详细信息。
• CUTLASS 文档，用于高效矩阵乘法
• Flash Attention 3 论文：突出了几种 Hopper 特定技术

---

附录

我们将讨论之前为了简洁而省略的一些细节。

张量核心操作

以下是 m64n16k16 张量核心操作的 PTX 实现：

它接受存储 A 和 B 的共享内存描述符，以及存储输出的寄存器。这个操作每线程使用 8 个寄存器。更大的指令需要更多的寄存器作为参数。

批量 WGMMA 操作

由于 WGMMA 操作在每个 SM 的 4 个张量核心上异步执行，我们可以批量多个张量调用并并行执行它们。