如何让计算受限问题真正达到计算受限：CUDA 核心上的矩阵乘法优化

免责声明：本博客内容反映了我在业余时间学习 GPU 编程时的个人经验和观点。所有呈现的信息均为公开信息，不代表 NVIDIA Corporation 或其任何关联公司的观点或立场。

0. 引言

当计算机架构师想要加速一个应用时，他们首先要做的事情之一就是判断该应用是计算受限还是内存受限。为此，他们通常会使用roofline 模型。

如果应用被认为是内存受限的，那么架构师会专注于以下方法：

• 首先为基线实现达到峰值内存带宽。
• 减少内存流量：例如量化、稀疏化、压缩、算子融合等。
• 提高内存带宽：例如增加更多 HBM 站点、拥有更多 SRAM bank 等。

只要保持足够的内存级并行（MLP），为基线实现达到峰值内存带宽并不难。内存系统会做合理的事情来使自己饱和。并且有大量的研究提案来缓解内存瓶颈。

如果应用被认为是计算受限的，类似地，架构师会专注于以下方法：

• 首先为基线实现达到峰值计算吞吐量。
• 减少计算量：例如稀疏化，或使用计算效率更高的算法等。
• 提高计算吞吐量：例如提高频率、增加更多核心等。

我认为如果不仔细思考，对于计算受限的应用，尤其是在 GPU 上，达到峰值计算吞吐量并不容易。很多时候你甚至无法达到计算 roofline。原因在于我们在 roofline 模型中常用来推理性能的简单机器模型信息量不够。它无法为我们应该做哪些优化提供清晰、系统的指导。有了更真实的机器模型，我们就可以从第一性原理出发，推理我们为达到峰值计算吞吐量而应用的各种优化。

在这篇博客中，我将以 CUDA 核心上的矩阵乘法为例，说明如何达到理论峰值计算吞吐量。简而言之，内存系统需要以与计算单元消耗速率相匹配的速率将操作数传递给计算单元，并在内存层次的每一级进行仔细的分块，以实现带宽放大。然后我将简要提及一些实现这种带宽放大的替代方法。

这篇博客将围绕对于 SIMT 程序，如何实现吞吐量匹配，而不是跨流水线阶段的延迟隐藏。尽管我们（包括我自己）所接受的传统 GPU SIMT 编程（非 tensor core）教育告诉我们，目标是实现完全的（内存）延迟隐藏。我认为这不是最好的思考方式。特别是内存延迟是内存吞吐量的函数（即 $\text{latency} = \text{static_latency} + \frac{\text{tile_size}}{\text{memory_BW}}$）。因此吞吐量是第一公民，而延迟是副产品。所以首要任务是吞吐量匹配。

1. 示例问题：CUDA 核心上的矩阵乘法

1.1 动机

许多人尝试编写能与 cuBLAS 性能匹敌的 CUDA 核心 fp32 gemm 内核。Simon Boehm 的博客是我参考的那篇。尽管他的博客信息量极大，但我觉得最重要的优化——能让你完成大部分工作的那个——被埋没在许多不那么重要的优化之中。我自己尝试在 Nvidia V100 GPU 上编写了一个 CUDA 核心 gemm，仅通过两个超越朴素实现的优化，就达到了 cuBLAS 性能的 88%：

1. 共享内存（smem）分块。
2. 寄存器文件（RF）分块。

你可能已经注意到，这两个优化都与分块有关，但处于内存层次结构的不同层级。希望通过这篇博客的其余部分，我能向你展示为什么分块本身就足以达到（几乎）峰值计算吞吐量，以及你应该如何选择分块大小。我所有的推理都将是纯分析性的，不涉及任何芯片性能剖析，而这已足以让我们达到峰值计算吞吐量。

1.2 背景

首先，让我们定义 gemm 问题以及一些术语。我们有两个矩阵 A 和 B，我们想计算乘积 C = A * B。矩阵的大小为 [M, K] 和 [K, N]，我们想计算大小为 [M, N] 的乘积 C。为简单起见，在本博客中我们以 M=N=K=4096 作为示例问题规模。

由于 GPU 有许多 SM（流式多处理器），我们需要以某种方式将 gemm 问题并行化到所有 SM 上。我们并行化的方式属于 dataflow 的一部分（或计算调度，参考文献1，参考文献2，参考文献3）。人们在将 gemm 映射到 SM 时最常用的数据流称为 output-stationary (OS) dataflow（或内积（IP）数据流）。下图展示了 OS 数据流以及 1 个 SM 的工作。这意味着每个 SM 负责生成输出矩阵的一个不相交分块（C(0,0)）。输出分块在 SM 中是驻留的。输出在 SM 中获得最大重用。高亮的 SM 负责执行 [BLOCK_M, K] (A0) x [K, BLOCK_N] (B0) -> [BLOCK_M, BLOCK_N] (C(0,0))。

人们倾向于使用 OS 数据流的原因是 SM 间通信开销很大。假设我们不使用 OS 数据流，这意味着多个 SM 协作生成 1 个输出分块。它们必须通过全局内存（缓存在 L2 中）进行通信和同步，这极其缓慢。因此，使用 OS 数据流来避免昂贵的 SM 间通信是一种常见做法。

随着 Hopper 中引入分布式共享内存，这种权衡可能会改变，它在线程块集群内提供了更快的 SM 间通信。我将在未来的博客中探讨这一点。

2. 简单机器模型信息量不足

现在我已经定义了 gemm 问题，让我们看看如何使用 roofline 模型分析它在 Nvidia A100 80GB SXM GPU 上是计算受限还是内存受限问题。

经典的 roofline 模型假设一个简单的机器模型，其中我们有一个片外内存（DRAM）、一个容量无限的片上存储（SRAM）以及计算单元（SM CUDA 核心）。片外到片上内存的带宽有限。片上内存到计算单元的带宽无限。计算单元具有固定的吞吐量。代入 A100 的数据，我们得到如上所示的简单机器模型。

注意，这是 A100 GPU 的每 SM 切片，意味着 DRAM 带宽是每个 SM（108 个 SM）能从总 DRAM 带宽（2039 GB/s）中获得的带宽。$\frac{2039GB/s}{108 \times 1.41GHz} = 13.3B/cycle$。而计算单元吞吐量是针对 1 个 SM 的（即 64 FMA/cycle）。或者，你可以通过聚合所有 SM 的吞吐量来为整个 GPU 创建机器模型。这基本上等同于将我们机器模型中的所有数字乘以 108（SM 数量），并且完全不会影响 roofline 分析。

现在我们用简单机器模型计算 4096x4096x4096 (M, N, K) gemm 问题的算术强度。强制流量（即无限片上内存）为 $4096 \times 4096 \times 3 = 48MB$。FMA（融合乘加）的数量为 $4096 \times 4096 \times 4096 = 64G$。因此 gemm 问题的算术强度为 $\frac{64GFMA}{48MB} = 1365FMA/B$。

A100 80GB SXM GPU 的算术强度为 $\frac{64FMA/cycle}{13.3B/cycle} = 4.8FMA/B$。由于 1365 > 4.8，我们可以得出结论：此 gemm 在 A100 上运行时将是计算受限的，即 CUDA 内核的目标是饱和计算吞吐量。但请记住，这个结论基于至少两个理想化假设：

1. 无限片上存储。
2. 无限片上存储到计算单元的带宽。

在实际的 A100 规格中，这两种资源并非无限，我们仍然计算受限吗？看起来是的，因为 cuBLAS 达到了计算 roofline。为什么我们的朴素实现只有 cuBLAS 速度的 16\%？不可知。我们应该如何编写内核，使其像 cuBLAS 一样充分利用计算单元？不可知。简单机器模型的信息量不足以指导我们编写最优化的内核来实现计算受限目标。

3. 更现实的机器模型

这就是更现实的机器模型发挥作用的地方。它在抽象和细节之间取得平衡。它忠实地对 GPU 内存层次结构建模，同时足够高层，不会让我们被硬件实现细节分心。更现实的机器模型为我们提供了清晰的指导，说明如何编写高性能 CUDA 核心 gemm 内核，以实现计算单元的完全利用。

左侧我展示了更现实的机器模型。注意与简单模型的区别：

1. 我明确展示了 3 级内存层次结构（DRAM->smem->RF）。
2. 每一级都有有限容量（假设我们总能将所有内容放入 DRAM）。
3. 每一级到下一级的带宽有限。

注意从 RF 到 CUDA 核心的带宽是 512 B/cycle。这是完全饱和 CUDA 核心所需的最小输入数据带宽。我们有 64 个 CUDA 核心，每个每周期需要两个 4B 输入。总共，内存层次结构需要向 CUDA 核心提供 $64 \times 2 \times 4B/cycle = 512B/cycle$ 输入数据才能完全饱和它。实际的 RF 到 CUDA 核心带宽可能更高，但我们不需要超过 512 B/cycle 的输入带宽。所以我们使用这两个数字中的最小值。

为什么我们只关注输入数据带宽需求，而不考虑输入和输出两者？这是因为我们对内核实现做了一个假设，即在内存层次结构的所有层级都使用output-stationary (OS) dataflow。只有当所有部分结果在当前层级完全累积后，输出才会写入下一层级。它获得了完全重用。这意味着与输入带宽需求相比，输出带宽需求极小，因为我们需要多次读取相同的输入。因此，为了简单起见，我们可以忽略输出。如果我们选择不在所有地方使用OS，这个假设就不成立。我们就需要考虑输入和输出带宽需求。

4. 内核优化目标

考虑到更现实的机器模型，我们已经确定CUDA核心需要从内存层次结构获得512 B/周期的输入数据才能实现完全计算利用率，这是我们的最终目标。所有输入数据都来自DRAM。通过我们的内核优化，我们需要制造一种假象，即CUDA核心以512 B/周期的带宽直接读取DRAM。

从现实的机器模型来看，输入数据传递流水线是DRAM->smem->RF->CUDA core。我们希望这个流水线的吞吐量为512 B/周期。这意味着每个流水线阶段都需要达到512 B/周期。

不幸的是，在没有任何优化的情况下，DRAM->smem带宽仅为13.3 B/周期，这限制了整体数据传递流水线的吞吐量。我们需要一种方法来放大带宽，以匹配整个流水线的数据传递吞吐量要求。同样，smem->RF也不够大（128 B/周期），无法匹配整体流水线吞吐量要求（512 B/周期）。我们也需要一种方法来放大这个带宽。总之，总体目标是放大所有阶段之间的带宽，以匹配吞吐量要求（即CUDA核心数据传递带宽）。

正如我已经暗示的，所有阶段的带宽放大问题可以解耦。在我们的示例问题和机器中，我们可以将带宽放大问题分解为两个独立的子问题：

1. DRAM->smem带宽放大（红色虚线框）。
2. smem->RF带宽放大（蓝色虚线框）。

在每个子问题中，带宽放大后，我们希望在每个阶段传递512 B/周期。如果两个子问题都成功，我们就可以在整个流水线的所有阶段之间传递512 B/周期的输入数据。从而充分利用CUDA核心。

5. 如何实现带宽放大

有许多方法可以在内存层次结构中实现带宽放大。我将讨论两种流行技术：分块（第5.1节）和多播（第5.2节）。并简要提及其他技术（第5.3节）。这些技术应该与我们所在的任何内存层次结构层级无关，这样我们可以在不同层级使用不同的技术。我们也可以在同一层级组合多种技术。

5.1 分块

分块的一般前提是它利用数据重用来实现带宽放大。在本节中，我将展示如何在DRAM->smem阶段（也称为smem分块（第5.1.1节））和smem->RF阶段（也称为RF分块（第5.1.2节））应用分块来实现带宽放大。我还将描述如何选择分块大小，以便我们在每个阶段有足够的带宽放大因子来实现完全计算利用率。

5.1.1 共享内存（smem）分块

共享内存（smem）分块试图在DRAM->smem阶段实现带宽放大。下图展示了该场景。例如，在左图中，我们将一块输入数据从 DRAM 传输到 smem。我们期望该块中的每个元素能被下一阶段仅从 smem 读取两次（不接触 DRAM）。换句话说，输入元素的复用次数为 2。这等价于下一阶段从 DRAM 读取同一元素两次（绕过 smem）（如右图所示）。因此我们有效地实现了 2 倍的带宽放大系数。

代入真实的 A100 数值，我们希望放大后的带宽为 512 B/cycle，而目前DRAM->smem带宽为 13.3 B/cycle。因此 smem 分块应至少给我们带来 $\frac{512}{13.3}=39$ 的放大系数。这意味着，平均而言，smem 中的每个输入元素应至少被复用 39 次。如果我们实现 39 的复用系数，那么 smem 就能给 CUDA 核心（以及整个流水线）一种错觉，即它可以直接从 DRAM 提供 512 B/cycle。

现在我将解释如何对 gemm 进行分块，使得每个输入元素被复用 39 次。再次，我们假设使用output-stationary (OS) dataflow。下图展示了内核中 1 个线程块（SM）执行的计算。

现在让我们计算每个输入元素的复用系数。总计算量为BLOCK_M * BLOCK_N * K FMA。我们从 DRAM 读取一次大小为BLOCK_M * K的 A 块，并从 DRAM 读取一次大小为K * BLOCK_N的 B 块。因此，每个 A 元素的复用系数为 $\frac{BLOCK_M * BLOCK_N * K}{BLOCK_M * K} = BLOCK_N$。每个 B 元素的复用系数为 $\frac{BLOCK_M * BLOCK_N * K}{K * BLOCK_N} = BLOCK_M$。我们需要平均输入复用系数为 39。因此BLOCK_M = BLOCK_N = 39就足以在DRAM->smem阶段实现完全的带宽放大。

如果BLOCK_M != BLOCK_N，则每个输入元素的平均复用系数为 $\frac{2 * BLOCK_M * BLOCK_N * K}{BLOCK_M * K + K * BLOCK_N} = 2 * \frac{BLOCK_M * BLOCK_N}{BLOCK_M + BLOCK_N}$（因为总读取量为 $2 * BLOCK_M * BLOCK_N * K$）。我们只需确保该平均复用系数至少为 39。

更现实地说，2 的幂次大小的块对索引计算和内存系统更友好。因此我们为 smem 块大小选择BLOCK_M = BLOCK_N = 64。

5.1.2 寄存器文件（RF）分块

类似地，寄存器文件（RF）分块试图在smem->RF阶段实现带宽放大。上图展示了该场景。RF 每周期只能从 smem 读取 128 B，但 CUDA 核心每周期可以从 RF 消耗 512 B。因此每周期，CUDA 核心需要 512B 数据，但 smem 只能提供 128B。这意味着从 smem 读取的每个输入元素将被 CUDA 核心读取 4 次，即在smem->RF阶段需要 4 的复用系数。

如何在smem->RF阶段实现 4 的复用系数？我们在smem->RF阶段使用相同的 OS 数据流，意味着 RF 持有一个大小为BLOCK_M * K的 A 块和一个大小为K * BLOCK_N的 B 块。A 输入复用系数为 4 意味着 RF 块大小为BLOCK_M = BLOCK_N = 4。那么smem->RF阶段也给 CUDA 核心一种错觉，即它可以直接从 smem 提供 512 B/cycle。

5.1.3 整合所有内容

将 smem 和 RF 分块放在一起，如上所示。整个流水线为DRAM->smem->RF->CUDA core：

• 通过 smem 分块（块大小为 64），我们可以确保DRAM->smem阶段向流水线的其余部分提供 512 B/cycle，即通过分块，它表现得就像我们可以直接以 512 B/cycle 的吞吐量读取 DRAM。
• 使用 RF 分块（tile 大小为 4），我们可以确保 smem->RF 阶段向流水线提供 512 B/cycle，也就是说，通过分块，它表现得就像我们可以直接以 512 B/cycle 的吞吐量读取 smem。

既然我们已经知道 RF->CUDA core 具有足够的 512 B/cycle 吞吐量。我们可以得出结论：整个流水线具有 512 B/cycle 的吞吐量，将输入数据从 DRAM 传递到 CUDA 核心。因此，我们实现了完整的计算利用率。

分块实现了数据重用，进而实现了带宽放大。有了足够的带宽放大，我们就能够实现完整的计算利用率。这就是为什么仅靠分块就能让我们获得 cuBLAS 的大部分性能，因为它足以满足所有带宽放大需求。

5.1.4 局部与全局重用因子

如果你不想被搞糊涂，就不要读这一节。如果你好奇并且和我有同样的问题，请继续。

有一个微妙之处我想澄清，它在我写这篇博客时最初让我非常困惑。问题是，对于 DRAM->smem 阶段，为什么我们需要将带宽放大到 512 B/cycle 而不是 128 B/cycle？smem 的消费者，即 RF，无论如何也只能从 smem 读取 128 B/cycle。那么为什么我们不直接将带宽放大到 128 B/cycle，即我们只需要在 DRAM->smem 阶段有 $\frac{128}{13.3}=9.6$ 的输入重用因子（例如 tile 大小为 16 就足够了）？

答案是数字 16 是 DRAM->smem 级别的局部重用因子。这并不意味着 smem tile 大小为 16。smem tile 大小可以更大（例如 64）。这可以用 gemm 内核的实际循环嵌套来更好地解释。

L1= 16 # DRAM->smem reuse factor
L2= 4 # smem->RF reuse factor
-----------------------------------------------------------
gA = [M, K]
gB = [K, N]
for m0 in range(M // L1 // L2):
  for n0 in range(N // L1 // L2):                   DRAM
-----------------------------------------------------------
    sA = gA[m0 * L1 * L2 : (m0 + 1) * L1 * L2, :] # [L1 * L2, K]
    sB = gB[:, n0 * L1 * L2 : (n0 + 1) * L1 * L2] # [K, L1 * L2]
    for m1 in range(L1):
      for n1 in range(L1):                          smem
-----------------------------------------------------------
        rA = sA[m1 * L2 : (m1 + 1) * L2, :] # [L2, K]
        rB = sB[:, n1 * L2 : (n1 + 1) * L2] # [K, L2]
        for m2 in range(L2):
          for n2 in range(L2):                        RF
-----------------------------------------------------------
            for k in range(K):
              idx_m = m0 * L1 * L2 + m1 * L2 + m2
              idx_n = n0 * L1 * L2 + n1 * L2 + n2
              C[idx_m, idx_n] += rA[m2, k] * rB[k, n2]

从循环嵌套中可以清楚地看出，数字 16 是 DRAM->smem 级别的局部重用因子，即 L1 = 16。这并不意味着 smem tile 大小为 16。因为在 smem->RF 级别还可以有进一步的重用，例如 L2 = 4。所以全局 smem 输入重用因子是 L1 * L2 = 64，即一个元素被取入 smem 一次，然后被所有较低级别总共消耗 64 次。在这 64 次中，RF 将读取 smem 16（L1）次。每次读取，都会触发从 CUDA 核心进一步读取 4（L2）次。因此，为了实现 64 的 smem 重用因子，我们需要 64 的 smem tile 大小。所以这里的数字 16 意味着在 DRAM->smem 级别，我们以 16 的因子重用 RF tile。但 RF tile 中的每个元素在较低级别有 4 的重用因子。所以对于每个单独的元素，它被重用 16 * 4 = 64 次。我们经常提到的 smem tile 大小是在单个元素级别。所以 smem 级别的单个元素 tile 大小是 64。

类似地，在 RF 级别，局部重用因子是 L2 = 4。由于 RF 之下没有进一步的重用，对于每个元素，全局 RF 输入重用因子确实是 4。为了实现这个 4 的重用因子，我们需要 4 的 RF tile 大小。

TL DR：我认为这种局部/全局重用因子的区分确实让人困惑，我们不应该考虑它。为简单起见，我们应该只考虑全局重用因子，即我们希望在每个阶段都有带宽放大，以匹配整个流水线的 512 B/cycle 要求。而带宽放大对每个阶段的输入元素有（全局）重用因子的要求。然后这个重用因子要求将转化为每个阶段的 tile 大小要求。

5.2 多播

平铺并不是实现带宽放大的唯一方法。另一种常用技术是多播。多播的工作原理是，如果多个单元请求相同的数据，内存系统只需生成 1 个请求，并且通常只得到 1 个响应。当响应到达内存系统中更靠近单元的位置时，它会被复制到所有单元。这样，响应带宽就高于请求带宽，即得到了放大。

要使多播正常工作，它依赖于以下条件：

1. 系统中有多个计算单元。
2. 它们请求相同的输入数据。

我们如何利用多播来实现带宽放大？我们将以 DRAM->smem 阶段为例，并假设 DRAM->smem NoC 中存在某种带宽放大机制。

在上面的例子中，当我们进行 smem 平铺时，我们得出结论：在 DRAM->smem 阶段，我们至少需要 39 的输入重用因子才能实现全带宽放大，这转化为最小 smem 平铺尺寸为 39。假设我们的 smem 非常小，无法容纳 39 的平铺尺寸。它只够容纳 32 的平铺尺寸。看起来我们无法在 DRAM->smem 阶段达到完整的 512 B/周期带宽要求。我们注定失败了吗？不。我们仍然可以通过利用多播来弥补差距，从而实现完整的带宽要求。

我们展示一个例子：4 个线程块（在 4 个 SM 上），每个都有一个 BLOCK_M=BLOCK_N=32 smem 平铺尺寸，带有 OS dataflow。根据上述分析，平铺尺寸 32 不足以在 DRAM->smem 阶段实现全带宽放大。总共，4 个 SM 将从 DRAM 读取 4 个 A 平铺（4 * BLOCK_M * K）和 4 个 B 平铺（4 * K * BLOCK_N），即 4 * (32 + 32) * K * 4B = 1024*K Byte。将其除以 4 个 SM 的 DRAM 带宽（4 * 13.3 B/周期），我们需要 19.2 * K 个周期来读取所有数据。而 CUDA 核心将执行 4 * BLOCK_M * BLOCK_N * K 次 FMA，消耗 4 * BLOCK_M * BLOCK_N * K * 2 个 A 和 B 元素（速率为 512 B/周期）。因此计算周期为 4 * 32 * 32 * K * 2 / 512 = 16 * K 个周期。由于 19.2 * K > 16 * K，我们受限于 DRAM->smem 带宽。这只是另一种说法，说明为什么平铺尺寸 32 不够。

但多播将解决这个问题。注意，4 个 SM 并不是在读取完全不同的数据。事实上，它们只读取 2 个不同的 A 平铺（2 * BLOCK_M * K）和 2 个不同的 B 平铺（2 * K * BLOCK_N）。而且我们知道有一种 NoC 魔法可以免费将 2 个 A 平铺和 2 个 B 平铺多播到所有 4 个 SM。因此读取所有数据所需的周期数为 19.2 * K / 2 = 9.6 * K 个周期，是原来无多播时 19.2 * K 个周期的一半。而计算周期为 16 * K 个周期。所以我们不再受限于 DRAM->smem 带宽，而是受限于计算。

5.2.1 多播的一种视角

一种思考方式是，多播基本上使 DRAM 带宽翻倍，因为在 4 个 SM 之间发生了 2 倍的请求/响应合并。然后高层次的结论是，即使从单个 SM 的角度来看，平铺尺寸 32 不够，跨 SM 的多播可以在 DRAM->smem 阶段再提供 2 倍的带宽。因此，39 的重用因子要求可以分解为两部分：

• SM 内的 smem 平铺重用（例如重用因子为 L1=32）
• 跨 SM 的多播（例如重用因子为 L2=2）

只要 L1 * L2 >= 39，我们就能在 DRAM->smem 阶段实现全带宽放大。

5.2.2 多播的另一种视角

另一种等价的思考多播的方式是，它实际上增大了共享内存块的大小。如果我们把 4 个线程块（4 个 SM）打包成一个集群（即一个 线程块集群），那么对于这个集群来说，我们实际上拥有 BLOCK_M=BLOCK_N=64。因此输入复用因子为 64，高于我们所需的 39。所以即使每个 SM 的共享内存块大小只有 32，跨 SM 的多播实际上给了我们 64 的共享内存块大小。

这正是线程块集群的用例。更正式地说，这里我们创建了一个 2x2 线程块集群，其线程块块大小为 BLOCK_M=BLOCK_N=32。集群块大小为 BLOCK_M=BLOCK_N=64。这样做的好处是，如果由于某些原因（例如共享内存容量限制）你无法拥有较大的线程块块大小，从而受到 DRAM->smem 带宽限制，你可以使用线程块集群来有效增大共享内存块大小，并实现完整的带宽放大。

5.2.3 线程块光栅化

重要的是，要让多播生效，你需要多个 SM 请求相同的数据。这通常通过改变内核的数据流/映射来实现。这也被称为 线程块光栅化。它基本上是在说 SM 执行所有线程块的顺序是什么。

上面我们展示了在一个 4x4 线程块 gemm 和 4 个 SM 下的三种不同光栅化顺序。相同的颜色表示同一波计算。总计算量为 4 波。假设 BLOCK_M=BLOCK_N=32。左侧是一种朴素的光栅化顺序，我们按行优先顺序执行线程块。如你所见，在这种行优先顺序下，所有 4 个 SM 请求的唯一输入数据量为 1 * BLOCK_M * K + 4 * K * BLOCK_N = 160 * K。而请求的输入数据总量为 4 * (BLOCK_M * K + K * BLOCK_N) = 256 * K。在这种光栅化顺序下，我们的多播因子为 256 / 160 = 1.6。

中间是我们在上面的示例中用来利用多播的光栅化。在这种光栅化顺序下，所有 4 个 SM 请求的唯一输入数据量为 2 * BLOCK_M * K + 2 * K * BLOCK_N = 128 * K。而请求的输入数据总量为 2 * (BLOCK_M * K + K * BLOCK_N) = 256 * K。在这种光栅化顺序下，我们的多播因子为 256 / 128 = 2，优于行优先顺序的 1.6！

右侧是一种故意构造的糟糕光栅化顺序，它将没有多播。如你所见，4 个 SM 之间没有输入数据共享（局部性差）。因此多播因子为 1。但有趣的是，这正是我们进行共享内存 swizzling 以避免 bank 冲突所需的模式。我们改天再探讨。

因此，通过改变光栅化顺序，我们可以提高多播因子。这通常由用户在软件中完成。两个例子是 cutlass gemm 光栅化 和 triton gemm 光栅化（在教程中称为 L2 缓存优化）。

5.2.4 硬件中的多播

你可能已经想到，多播需要一些硬件支持。在 Nvidia GPU（Hopper+）中，跨内存层次至少有 3 种多播：

• 线程块集群：用户指定哪些线程块组成一个集群（例如上面示例中的 2x2）。硬件将负责集群内的多播。
• L2 请求合并器（LRC）：用户不指定哪些线程块组合在一起。硬件将机会性地找到重复的输入请求并将它们组合在一起进行多播。
• smem：在一个 warp 内，如果没有 bank conflict，每个周期全部 32 个线程可以从 smem 读取 128B 数据。如果每个线程读取超过 4B（例如使用 LDS.64 读取 8B、使用 LDS.128 读取 16B），并且全部 32 个线程都从某些重叠的地址范围读取数据，那么就会发生从 smem 到 32 个线程的某种多播，因此有效 smem 带宽会高于 128B/cycle 的 smem 带宽。更多细节请参考我朋友 Axel Feldmann 的博客。

如果你遇到分块不足以实现完整带宽放大的情况，可以通过在多个计算单元之间智能地共享数据，使用多播来弥补差距。

5.3 其他技术

还有其他实现带宽放大的方法。我在这里简要列出：

• 更高的内存带宽：如果你没有足够的带宽，那就直接增加它 :)。这是一种需要更改硬件的暴力解决方案。
• 压缩：在内存层级中进行数据传输时压缩数据，将提高有效带宽。
• 稀疏性：如果数据是稀疏的，你只传输非零值以及一些指示非零值位置的元数据。这类似于压缩，可以提高有效带宽。
• 融合：如果你有多个 kernel 读取相同的数据，你可以将它们融合为一个 kernel。这样数据只被读取一次，然后被所有 kernel 消费。这类似于多播，但发生在 kernel 层面。Batched gemm（A 在各 batch 间共享）就是融合的一个例子。

6. 从理论到实践

这篇博客主要聚焦于分析性分析，探讨如何通过内存层级间的吞吐量匹配来实现完整的计算利用率。在实践中，实现理论分析还需要付出一些努力。此外，内存吞吐量匹配并不是唯一需要考虑的因素（尽管可能是最重要的因素）。以下是一些其他需要考虑的因素：

• 功耗限制：正如 Horace He 在他的博客中指出的，如今的大方阵 gemm 会受到功耗限制。因此，在固定功耗上限下优化 perf/W 变得更加重要。因为数据移动会消耗功耗，数据移动消耗的功耗越少，我们就能把越多功耗投入到 tensor core 上。因此，在计算受限且受功耗限制的 kernel 中减少内存流量出乎意料地重要。减少内存流量的方法是利用更多的数据复用，即更大的 tile 尺寸或更大的多播因子。因此，选择一个刚好满足带宽放大因子最低要求的 tile 尺寸是不够的。我们希望它尽可能大。
• （静态）内存延迟：静态内存延迟在历代 HBM 中不断增长。这意味着如果我们不提供足够的内存级并行（MLP），我们可能会受限于延迟而非带宽。我的看法是，如果你受限于延迟，那只是意味着你实际达到的带宽甚至低于峰值带宽（13.3 B/cycle/SM），因此这仍然是一个带宽不足的问题。这就是为什么 cutlass 和 triton（triton.config 中的 num_stages）都使用多个 DMA 阶段来隐藏内存延迟。
• DRAM Bank Camping：取决于数据访问模式，某些 DRAM bank 收到的请求会比其他 bank 少。因此这些 bank 可能无法达到峰值带宽。这样我们实际达到的带宽就会低于峰值带宽，进而需要更多的带宽放大。
• Bank Conflict：如果你不仔细考虑自己的 smem 访问模式，就会遇到 smem bank conflict，这会降低有效的 smem 带宽。你甚至无法达到 128 B/cycle 的带宽。有很多方法可以避免 bank conflict 并利用 smem multicasting。更多细节请参考我的朋友 Axel Feldmann 的博客。

显然，本博客中的所有分析都适用于 tensor core gemm（或任何应用）。但正因为上述因素，要饱和 tensor core 吞吐量更加困难。在我们简单的 CUDA core gemm 示例中，我们并未受到功耗限制。而且由于 CUDA core 吞吐量低于 tensor core，操作数交付吞吐量需求也更低。所以即使我们不做多级 DMA 来隐藏内存延迟或摊销 bank camping，也没问题。但对于 tensor core gemm，我们需要考虑上述所有因素。

最后一点想法是，所有分析都假设采用 output-stationary（OS）数据流。在不同的硬件配置下，最佳数据流可能会改变。而如果你改变数据流，吞吐量匹配要求也会改变。但跨内存层级进行吞吐量匹配的总体原则仍然适用。

7. 附加：一个 GEMV 示例

到目前为止，我们只涵盖了一个示例，即一个 gemm kernel。让我们看另一个示例，一个 fp32 gemv kernel，以证明上述分析具有良好的泛化能力（甚至对 memory-bound kernel 也适用）。

7.1 使用真实机器模型估算 FLOPS

因为这是一个 gemv 案例，N=1。假设我们采用常见的 output-stationary（OS）数据流，其中每个 SM 负责生成一个 BLOCK_M=BLOCK_N=1 的输出 tile。并且我们直接从 DRAM 加载值到 RF（绕过 smem）。因此新的机器模型如下所示：

对于这个计算，每个 SM 都在进行一个向量（长度为 K）点积。给定这个 tiling 因子，RF 级数据复用因子为 $\frac{2 * BLOCK_M * BLOCK_N * K}{BLOCK_M * K + K * BLOCK_N} = 2 * \frac{BLOCK_M * BLOCK_N}{BLOCK_M + BLOCK_N} = 1$。这意味着没有 RF 级数据复用/带宽放大。增大 BLOCK_M 会有帮助。

现在让我们用这个映射计算可达到的 flops。由于没有 BW 放大，RF 以与 DRAM 供给 RF 相同的 13.3 B/cycle 速率供给 CUDA core。每次 FMA 需要 8 B/cycle（2 个 4B 操作数）。所以可达到的 flops 为 13.3 / 8 = 1.64 FMA/cycle/SM。

7.2 使用 Roofline 模型估算 FLOPS

gemv 的算术强度为 $\frac{M * N * K}{4 * (M * K + N * K)} = \frac{M * N}{4 * (M + N)} \sim \frac{1}{4} \text{FMA}/B$。而 A100 的算术强度拐点为 $\frac{64 \text{FMA}/cycle}{13.3 B/cycle} = 4.8 \text{FMA}/B$。所以 gemv 是 memory-bound 的。那么每个 SM 的 flops 为 arithmetic intensity * BW = 1/4 FMA/B * 13.3 B/cycle = 3.33 FMA/cycle/SM。

这比真实机器模型下可达到的 flops 高 2 倍！这意味着真实的 gemv kernel 甚至还没有达到内存 roofline。

7.3 哪里出了问题？

这是因为真实机器模型还不够真实，哈哈。如果每个 SM 都单独连接到 DRAM，那么 1.64 FMA/cycle 确实是可达到的 flops。所有 SM 都通过 L2 连接，L2 通过 LRC 提供 multicasting 能力（如 第 5.2.4 节所述）。

在我们的 gemv kernel 语境下，这意味着什么？这意味着我们从 DRAM 到 RF 获得了带宽放大，因为所有 SM 都读取同一个向量 B in。所以在所有 SM 上，我们的输入复用因子是 $\frac{2 * M * N * K}{M * K + K * N} \sim 2$。这意味着放大后的 DRAM->SM 带宽是 2 * 13.3 B/cycle = 26.6 B/cycle。因此，可达到的 flops 是 26.6 B/cycle / 8 B/cycle = 3.33 FMA/cycle/SM，这与 roofline 预测相符。借助 LRC，gemv kernel 可以达到内存 roofline。

LRC 带宽放大的效果如上所示。我们将每个 SM 的有效 DRAM 带宽翻倍。另一种理解方式是，如果没有 LRC，只有一半的 DRAM 带宽用于获取矩阵 A，另一半用于冗余地获取同一个向量 B。有了 LRC，我们可以只获取一次 B 并将其多播到所有 SM，因此大部分 DRAM 带宽用于获取 A，从而达到内存 roofline。

8. 总结

• 即使 roofline 分析告诉你受计算限制，要实现完整的计算利用率也并不简单。我们需要一个更真实的机器模型来暴露内存层次结构，以指导我们实现完整的计算利用率。
• 为了实现完整的计算利用率，跨内存层次结构的操作数交付吞吐量需要与 CUDA core 的计算吞吐量相匹配。内存层次结构每一层的带宽放大是实现吞吐量匹配的关键。
• Tiling 和 multicasting 是两种通过利用数据复用概念来实现带宽放大的常用技术。它们可以用于内存层次结构的不同层级，并且可以组合使用。
• 在实现操作数交付吞吐量匹配之后，还有其他因素需要考虑以实现完整性能，包括功耗节流、内存延迟和 bank conflict 等。

感谢我的朋友 Axel Feldmann 对这个主题的启发，以及 Jenny Huang 的校对和有益讨论。

9. 附加参考

• 关于 threadblock rasterization 的更多内容
• GTC 2025 How to Ace a Finance Developer Interview: A Deep Dive into GPU Matrix Optimization [S73619]