Stream-K: Work-centric Parallel Decomposition for Dense Matrix-Matrix Multiplication on the GPU

在 NVIDIA GTC22 秋季会议上，CUTLASS: Python API, Enhancements, and NVIDIA Hopper 介绍了 CUTLASS~2.11 中引入的 Stream-K 分解：
几个月后公开的 Stream-K: Work-centric Parallel Decomposition for Dense Matrix-Matrix Multiplication on the GPU 论文对其进行了更详细的介绍。CUTLASS 的 GEMM 实现由三大方法组成：

• SlicedK，即论文中的 Data-parallel；
• SplitK，即论文中的 Fixed-split；
• Stream-K。

Stream-K 是一种以任务为中心的并行工作负载分解技术，用于在 GPU 等宽架构上调度通用矩阵乘法（GEMM）及类似计算。与当前主要基于图块的分解不同，Stream-K 通过在 SM 之间平均（误差在一以内）分配的聚合内部循环迭代来运行。MAC 循环迭代是其跨处理器核心的工作负载量化单位。无论给定问题的输出分块在底层处理单元之间的量化效率如何，都能近乎完美地利用计算资源。这提供了极佳的强扩展性和工作负载平衡，因为其成本：

• 相对于问题形状是常数；
• 显著小于整个输出图块的成本。

此外，Stream-K 产生的分割缝数量为 $O(p)$，受处理器核心数量限制。因此，临时存储开销与处理器宽度成正比，而非问题规模。 这对于许多应用程序来说是一个值得欢迎的特性，因其无法负担分配相当于问题输出的大量临时存储。

作者在广泛的 GEMM 形状和大小范围内评估了所提 Stream-K 方法。证明了通过每种浮点类型的单一分块配置，Stream-K 能够：

• 达到与 NVIDIA 的 cuBLAS 库相媲美甚至超越的绝对性能水平，即使后者的处理器利用率接近峰值；
• 同时具有更高的性能一致性。

最后，Stream-K 对于库的构建和维护是一个有吸引力的选择，因为借助它可以在损害性能的情况下将分发大小减少一个数量级，并且无需复杂的手工编码启发式方法或机器学习模型来选择内核。

注意： 论文中是一种非常理想化的描述。Stream-K 拆分导致 IO 成本过高，性能并不能超过其他方法。改良的混合调度为 SlicedK + Stream-K 或者 SplitK + Stream-K，主体为原来的方法。所以，与其说 Stream-K 是一种以任务为中心的并行工作负载分解技术，不如说是一种尾数处理手段。

Introduction

经典的 GEMM 实现通过对输出矩阵进行数据并行切块来分块计算，将独立生成输出图块的任务分配给并行的线程（或线程组）[KBLAS, CUTLASS, 14]。每个输出图块的工作是规则的，图块生产往往以“波浪”形式在空闲的物理核心之间进行调度。当有许多波浪时，即输出切块的数量大大超过核心的数量时，整体工作负载平衡良好，处理器利用率最高。

然而，随着处理器规模的扩大，这种超额订阅已经大幅减少：

• 增加的核心数量将需要更少的波浪来产生给定的切块数量；
• 更大的核心将迫使矩阵的分块因子更大，导致更大的切块、更少的波浪。

一般而言，波浪更少的执行调度更容易出现量化效率低下，即当输出切块的数量不是处理器核心数量的整数倍时发生的处理器利用率不足。由于“黄式刀法”，NVIDIA GPU 的 SM 数量几乎都不是2的幂。实际中难以被任务整除。当最后一波不满载时，未使用的核心必须等待剩余的线程执行数百万（如果不是数十亿）的乘加（MAC）指令，然后才能执行任何依赖的工作。

图 1a 在一个假设的拥有流式多处理器核心 (SM) 的 GPU 上展示了这样一个场景。如果我们将一个 $384\times 384\times 128$ 的 GEMM 计算划分为九个 $128\times 128$ 的输出块，数据并行分解实现无法超过处理器额定吞吐量的75% 。通过将图块大小减半，可以将理论利用率上限提高到 90%，如图 1b 所示。然而，更细粒度的分块因子将会降低缓存和暂存器的效率，并可能阻碍任何实际的性能改进。

对于越来越宽的处理器（例如 GPU，其中每个核的 ALU 和每个处理器的核数目前都在数百个），量化效率低下是一个令人担忧的问题。因此，许多常见的类似 GEMM 的工作负载现在呈现出最后一个部分完整的波浪，这构成了总计算时间的很大一部分。

目前基于 GPU 的数学和深度学习库采用的解决办法是部署一系列的切块配置。当理想的分块因子不能很好地量化时，库会在具有较小并发工作量的分块替代方案中进行选择，例如图 1b 和图 2a 中所示的方案。

然而，基于分块的集成方法为数学库提供了性能和组织管理方面的挑战，这些数学库试图在不同的问题大小和形状中提供最佳的性能。对于大型集合，可分发代码大小可能成问题。例如，NVIDIA 的 cuBLAS 库有数百兆字节，通常为给定的 API 入口点提供超过二十种针对每种架构的预编译内核特化。大型集成还需要复杂的选择启发法。在实际当中，这些启发式方法可能难以一致地识别任意问题的最优配置。

Background

通用矩阵乘法（General Matrix Multiplication，GEMM）定义为乘积 $\text{C}=\alpha \text{AB}+\beta \text{C}$，其中 $\alpha$ 和 $\beta$ 是标量值，$\text{A}$、$\text{B}$ 和 $\text{C}$ 是矩阵。（为简单起见，本文假设 $\alpha =1$，$\beta =0$。）作者通过 GEMM 问题计算的体积范围来指称 GEMM 问题的形状。例如，一个 $m\times n\times k$ 的 GEMM 问题消耗 $m\times k$ 的输入矩阵 $\text{A}$ 和 $k\times n$ 的输入矩阵 $\text{B}$，执行 $m\times n\times k$ 的乘累加运算，并生成一个 $m\times n$ 的输出矩阵 $\text{C}$。

Larsen 和 McAllister 关于 GPU 矩阵-矩阵乘法的早期工作将计算框架定为多纹理乘法和混合操作[11]。后续 GPU 架构提供的用户可编程共享内存使得性能更高的数据并行方案成为可能，这些方案采用了两级分块（共享内存和寄存器），并通过广泛的微基准测试分析[2, 14, 17, Triton]以及自动调优[5, 7, 12]来确定图块大小。

MAGMA GPU 数学库可能是第一个针对多样化的 GEMM 问题形状进行优化的库[9]。他们的解决方案是将一组带约束的拼块参数应用于模板化的 CUDA C++ 代码模板，为每个 API 原语（例如，用于半精度转置-转置 GEMM 的 hgemm_tt()）生成数百个数据并行变体。他们评估了这些变体，以提取出一个通常由三到五个内核组成的小型集合，这些内核在各种问题形状上集体表现良好。针对给定问题的内核选择和分派由通过简单手写规则表达的尺寸阈值控制。

随后的 GPU 数学库采用了更复杂的代码生成和内核选择组件。 例如，ISAAC 项目使用机器学习技术来预测给定 GEMM 形状的最佳拼块和（或）切分参数化，然后可以通过 PTX 级别的代码生成器在线或离线实例化。

NVIDIA 的 cuBLAS 库提供了一个扩展的cublasGemmEx接口，允许调用者从 24 种不同的 GEMM “算法”中进行选择。当使用默认接口时，经过精心训练的启发式算法会在这些大量的替代方案中做出选择。这些算法实现了多种不同的数据并行和固定分割变体，cuBLAS 通常会将每个变体组装成其自身特定架构的内核程序，以便进行代码优化。GEMM API 功能、策略变体和微体系结构的交叉积导致发行版变得越来越庞大，可执行代码超过数百兆字节。

鉴于当代深度学习的快速发展和不断变化的特性，最近的研究侧重于简化表达的编程模型和构造改变或补充 GEMM 计算的高性能内核。CUTLASS C++库提供了数据移动和乘法累加类，用于在 GPU 线程层次结构的各个级别组合自定义的类 GEMM 计算。Triton 是一种用于张量编程的领域特定语言，以块（图块）概念的表达、转换和优化为中心。其他领域特定编程语言，如 Halide 和 TVM，将逐点运算符的表达式与循环调度表达式分开。Fireiron 进一步将数据移动构造添加到调度语法中。

Existing Work Decomposition Strategies

现代处理器通常将 A、B 和 C 存储在大容量、速度慢、距离远的内存中，并且可以访问小容量、速度快的便签存储器或高速缓存。任何 GEMM 实现的一个主要目标是利用这些本地存储资源，使得最终实现是计算受限的。

Sequential Cache-Blocked

经典的 GEMM 缓存分块公式将其计算体积划分成块，并选择一种遍历顺序以展现内存局部性。算法 1 给出了一个包含6个循环的简化实现。最内层的三个循环在分块因子 BLK_M、BLK_N 和 BLK_K 内进行迭代，而最外层的三个循环则在它们之间迭代。如果缓存可以从三个矩阵中各捕获一个块，这些元素之间的数据重用将显著减少对最后一级内存的访问次数[10]。

Data-parallel

如算法 2 所示，GEMM 的数据并行 GPU 表述在并行线程块网格中分解，后者称为协作线程数组（cooperative thread array，CTA）。GPU 线程块在 CTA 中协同调度，这些 CTA 虚拟化了硬件的流式多处理器核心（streaming multiprocessor core，SM）。网格的大小设置为每个 CTA 生成自己的（BLK_M $\times$ BLK_N）输出图块。

为了便于说明，算法 3 中的MacLoop()子程序封装了计算 CTA 输出图块值的乘累加工作负载。它在累积域中执行一系列 MAC-loop 迭代，例如，GEMM 中的 k 轴。每个 MAC-loop 迭代包含每个线程（BLK_M $\times$ BLK_N $\times$ BLK_K) $/$ CTA_THREADS 的 MAC 操作量。随着计算的进行，输入矩阵的片段依次到达 SM 的共享内存，以便在各个线程之间进行局部重用。

尽管这个特定的MacLoop()展示为每个输出图块元素调用一个线程，但在 CUTLASS 和 cuBLAS 中的复杂实现将会：

• 完全展开每个线程的 MAC-loop 迭代；
• 在 warp 和（或）线程级别实施额外的分块；
• 在 MAC - loop 迭代中编排共享内存数据移动的软件流水。

遗憾的是，这种经典的数据并行分解在现代 GPU 上容易出现量化效率低下的问题，如图 1 所示。尽管多样化的分块因子集合可能会发现提高处理器利用率的机会，但它不太可能为任意问题大小提供完美的量化。此外，较小的分块因子有两个缺点：

• 在流水线实现中，每个 MAC-loop 迭代的指令较少，无法覆盖全局和共享内存传输的延迟；
• 相对于 MAC 指令，内存操作的比例更高，这可能阻止它们成为计算受限。

Fixed-split

或者，可以通过在累加维度上进行并行化来减少分配给每个 CTA 的工作粒度。 对于给定的输出图块，加法的结合律允许迭代域在多个并发 CTA 之间进行拆分，然后通过一个依赖的“修正”步骤来归约每个 CTA 计算的部分和。算法 4 中展示了这种固定分割方法，其中每个输出图块由 $s$ 个 CTA 协作产生。值得注意的是，当分割因子 $s=1$ 时，它的功能与数据并行分解完全相同。

固定分割分解也是 CUTLASS 和 cuBLAS 中的特性。分割因子作为运行时参数实现，允许单个内核可执行文件支持多种工作量，同时保留理想的分块因子以实现最佳数据共享和延迟隐藏。然而，正如图 2a 所示，通过均匀图块切分实现完美的量化的的前景并不乐观。此外，通信和同步的额外开销与整个问题的大小和分割系数的成正比。

Our Stream-K Decomposition

Stream-K 分解是一种图块切分并行化，其中切分线与拼块结构本身完全脱离。虽然我们采用熟悉的分块和拼贴策略来实现数据重用，但我们将 GEMM 计算量化为 MAC 循环迭代，即小量的 CTA 范围的 BLK_M $\times$ BLK_N$\times$ BLK_K 工作。如算法 5 所示，Stream-K 在 $g$ 个 CTA 构成的固定大小网格上均匀地划分 GEMM 的 MAC 循环迭代的总体工作负载。将每个 CTA 的 MAC 循环迭代范围连续映射到 GEMM 形状的 $m\to n\to k$ 线性化中，可能会跨越输出图块边界。

如果给定的 CTA 的开始和（或）结束迭代不与图块边界重合（这是预期的常见情况），它必须将其部分结果与处理同一图块的其他 CTA 的结果进行合并。在这个基本实现中，$\mathbf{C}$ 中的每个输出图块都由执行该图块的 $k=0$ MAC 循环迭代的 CTA 写入。然而，在此之前，它必须在临时全局存储中累积从其他 CTA 共享的所有部分和。值得注意的是，Stream-K 的通信、同步和全局存储开销与问题大小无关，而是与 CTA 的数量 $g$ 成比例增长。

Stream-K 的第二个好处是，当输出图块的数量大于 CTA 的数量时，同步等待可能几乎可以忽略不计。在这种情况下，每个输出图块最多由两个 CTA 处理，图块处理的偏差确保累加 CTA 在其协作者完成产生这些贡献之后很久才需要它们的贡献。

继续我们之前的例子，图 2b 展示了在一个假设的四 SM GPU 上 $384\times 384\times 128$ GEMM 问题的基本 Stream-K 执行调度。为了充分占用 GPU，我们启动了 $g=4$ 个 CTA。假设 BLK_M $=128$，BLK_N $=128$，并且 BLK_K $=4$，每个 CTA 负责一个 $128\times 128\times 288$ 的工作量，包含72次 MAC 循环迭代。这产生了100% 的量化效率，因为所有四个 SM 将执行相同数量的 MAC 指令。

此外，单次 MAC 循环迭代的工作量比整个输出图块的工作量小32倍。因此，32路固定分割分解也能提供100%的量化效率，但代价是8倍大的“fixup”开销。而且，由于共享部分和时写入器和读取器之间的时间偏差，Stream-K 能够更好地隐藏 CTA 间同步的延迟。

Stream-K 同样适用于固定分割和数据并行分解：

• 当网格大小 $g$ 是输出图块数量的偶数倍时，Stream-K 的功能与固定分割分解完全相同；
• 类似地，当 $g$ 等于输出图块的数量时，Stream-K 的行为与数据并行分解一致。

我们利用这种泛化，在后续章节 (5.2) 中创建了 Stream-K 分解的优化混合体。

Implementation Details

上一节中介绍的工作分解可以通过多种不同的方式实例化，以适应不同硬件架构和软件库设计的需求。本文的实现针对 NVIDIA GPU，并旨在集成到现有库如 cuBLAS 和 CUTLASS 中。本节描述了如何配置启动的内核，并引入了一种混合方案，该方案有助于确保用户在尽可能广泛的问题形状范围内实现最大的 GEMM 性能。

需要强调的是，这些确实是内部实现细节。它们对使用类似 BLAS 库的用户是完全透明的，并且不会改变库的接口。唯一可观察到的影响是文章在第6节中分析的性能特性的提升。

Kernel Configuration

用于分块 GEMM 计算的块尺寸当然是一个关键参数，控制着 GEMM 内核的性能。对于现代 NVIDIA GPU，合适的块大小由 GPU 的 Tensor Core 支持的矩阵形状决定。基于广泛的实证经验，作者选择了最小的 CTA 级别的块尺寸，在每个支持的精度下，对于非常大的 GEMM，其能够达到 GPU 峰值 TFLOP/s 的99% 。在我们实验中使用的 NVIDIA A100 GPU，这些尺寸分别是 FP64 问题的 64$\times$ 64 $\times$ 16，以及 FP16 $\to$ 32 问题的 128 $\times$ 128 $\times$ 32。

为了从 Stream-K 并行化中获得最大的 GEMM 性能，还需要某种程度的针对特定问题的动态配置。在启动内核之前，选择一个可能在当前特定问题形状上提供最佳性能的网格大小。 这与基于集成的方法不同，后者基于工作负载分解和分块因子静态生成许多核变体来适应不同的问题形状。

论文中的网格大小选择启发式算法基于一个简单的分析模型，该模型最大限度地降低了读取、写入和累加部分和的开销，同时平均分配每个 CTA 的 MAC 循环迭代。该分析模型的详细信息见补充材料（附录 A.1）。模型的参数是通过经验测量简单地选择的，并且每个目标架构只需要进行一次。然后将得到的参数静态编译到库中。再次强调，这与基于集成的方法形成对比，后者在运行时依赖潜在的复杂启发式算法和机器学习模型来选择内核。

Data-parallel Hybridization

在某些情况下，基本的 Stream-K 分解可能会表现出切片处理偏斜，从而对缓存性能产生潜在的不利影响。当输出块的数量 $t$ 不是网格大小 $g$ 的整数倍时，每个 CTA 中第一个 MAC 循环迭代的起始 $k$ 偏移将会不同。根据输入矩阵的大小和形状以及分块因子，这种偏差可能阻止 CTA 之间在 GPU 的缓存结构中重用这些片段。例如，在图 3a 中，四个 CTA 的初始 $k$ 轴片段偏移分别为 $k=0$、$k=32$、$k=64$ 和 $k=96$。此外，CTA 之间的这 32 元素偏差将在整个 GEMM 计算过程中持续存在。

图块处理偏差是 Stream-K 工作负载均衡策略的直接后果。然而，我们可以采取措施通过将 Stream-K 的迭代平衡应用于总迭代域的一个较小的、图块对齐的区域来限制其持续时间，使得剩余的图块可以在完全的、时间对齐的波中生成。

最简单的混合方案是图 3b 中所示的“数据并行 + 单图块 Stream-K”调度。它仅在最后一个部分填满的数据并行波中剩余的图块之间应用迭代平衡。完整波的总数为 $w=\lfloor t/p\rfloor$，其中 $t$ 是输出图块的数量，$p$ 是 GPU 中的 SM 核心数。因此，每个 Stream-K CTA 收到的迭代份额相等，且少于一个图块的量。

不幸的是，在三个或更多的 CTA 处理同一瓦片时，该策略几乎无法隐藏用于交换部分和的同步延迟。在这些场景中，累积的 CTA 可能被迫等待其他 CTA 的贡献变得可见，因为除了最后一个 CTA 外，其他所有 CTA 都将在大致相同的时间完成其最终迭代。此外，基础的 Stream-K 方案中在单个 CTA 内部串行聚合部分和，因此当每个图块的贡献 CTA 数量很大时，可能会导致 SM 工作负载不均衡。

作者通过“两图块 Stream-K + 数据并行”混合调度来解决这些问题，如图 3c 所示。它执行的完整数据并行波少一个，以换取每个 Stream-K CTA 接收超过一个图块（但少于两个）的迭代量。当 $w\ge 2$ 时，这种方式可以更好地隐藏延迟，并且每个累积 CTA 只需要从另一个贡献的 CTA 接收部分和。除此之外，其行为与“DP + 单图块 SK”调度完全相同。这种混合方法在改善内存访问模式和隐藏延迟方面都有所提升。它还展示了通用 *Stream-K 循环结构的多功能性，能够在同一个内核实例中实现不同的调度策略。

Performance Evaluation

作者使用 NVIDIA 的 CUTLASS 库实现了 Stream-K 分解，该库提供了用于编写 GEMM 类计算的 CUDA C++模板抽象。CUTLASS 提供了算法 3中 CTA 范围的MacLoop()子程序的优化等效实现，该子程序执行类似于闭源的 cuBLAS 和 cuDNN 实现的分块、拼贴和软件流水线数据传输。我们的评估包括：

• 双精度 FP64 GEMM；
• 以及混合精度 FP16 $\to$ 32 GEMM。

对于后者，输入矩阵 $\text{A}$ 和 $\text{B}$ 由半精度 FP16值组成，但内部累加以及输出矩阵 $\text{C}$ 的值为单精度 FP32。

Hardware environment

测试 GPU 为 NVIDIA A100，包含108个 SM 核心。为了测量稳定性，作者将功率限制锁定在400W，SM 时钟锁定在1005MHz（约为其动态峰值的 71%）。这使得 FP64 张量核心峰值吞吐量为13.9TFLOP/s，混合 FP16 $\to$ 32 张量核心峰值吞吐量为222.3 TFLOP/s。

Dataset

测试资料库旨在近似 GPU 数学内核库设计用于适应的设备范围 GEMM问题的巨大广度和范围。如图 4所示，作者评估了32,824种不同的问题大小和形状，这些问题在 $m$、$n$ 和 $k$ 矩阵维度的域内随机对数采样，其体积跨越六个数量级。

Methodology

对于这两种 GEMM 精度，作者根据第5节中的指南分别构建了一个专门的单一 Stream-K 内核。此外，这些内核实现了所提出的“双块 Stream-K + 数据并行”混合分解。评估实验将每个 Stream-K 内核与以下项进行比较：

• 具有相同分块因子的默认数据并行 CUTLASS 内核；
• 该精度的 cuBLAS 集成（CUDA 11.6）；
• 一个理想化的预测器，它总是选择性能最高的数据并行 CUTLASS 分块因子来执行给定的 GEMM 实例。

对于 FP64 问题，该预测器在以下分块因子的专门配置中进行选择：

• (32 $\times$ 32 $\times$ 16)
• (32 $\times$ 64 $\times$ 16)
• (64 $\times$ 64 $\times$ 16)
• (128 $\times$ 128 $\times$ 16)

对于 FP16 $\to$ 32 问题，它从以下分块因子的集合中选择：

• (64 $\times$ 64 $\times$ 64)
• (64$\times$ 128 $\times$ 32)
• (128 $\times$ 128 $\times$ 32)
• (128 $\times$ 256 $\times$ 32)

这些特定的专门配置是相应 cuBLAS GEMM 内核集合的开源严格子集替代方案。

图 6a 和图 5a 的“屋顶线”图突出显示了单一数据并行 CUTLASS 内核产生的性能分布。它们将 FP64和 FP16 $\to$ 32的处理器利用率百分比作为计算强度的函数进行绘制。理想情况下，GEMM 实现的性能响应将表现为紧密贴合机器带宽限制和计算限制性能上限的窄带。在这里，数据并行内核在任何给定的算术强度区制下都表现出相当大的动态范围。相比之下，图 6d 和图 5d 中相应的 Stream-K 内核的性能响应要紧密得多。这些观察结果得到了表 1 和表 2 的证实，表格显示 FP64和 FP16 $\to$ 32 类型 Stream-K 内核的性能平均值分别达到其数据并行等效内核的1.23倍和1.63倍。对于 $m\times n$ 较小而 $k$ 较大的极端强缩放场景，所提出的 Stream-K 内核分别取得了高达5.63倍和14.7倍的加速。

表 1 和表 2 的第二列将 Stream-K 的性能与 cuBLAS 进行了比较。平均而言，FP64和 FP16 $\to$ 32类型 Stream-K GEMM 内核分别提供了比相应的 cuBLAS 集成高出6% 和13%的吞吐量，峰值改进分别达到2.55倍和6.74倍。这是在 32K GEMM 问题形状和大小范围内的一项显著改进，与 NVIDIA 的供应商 GEMM 库 cuBLAS 相比，Stream-K 的执行代码仅有后者的1/20（每种精度一个内核）。

此外，FP64 和 FP16 $\to$ 32 cuBLAS 性能响应（图 6b 和图 5b）与我们假设的 CUTLASS 理想化集合（图 6c 和图 5c）的对比，揭示了设计内核选择启发式方法以提供始终如一的良好性能的困难。 尽管可以使用相同的分块因子特化，但 cuBLAS 的动态范围远远大于 CUTLASS 理想化集合。
Stream-K 内核的性能分布范围更窄，最多可达到理想预言机性能的4.6倍，并强调了它们达到以图块为中心的工作分解所无法实现的利用率水平的能力。

最后，我们观察到在小型、带宽受限的问题形状的区域，我们的较大的分块因子在与 cuBLAS 竞争时表现不佳。然而，如果我们将范围限制在计算受限问题上（即，计算强度大于150 ops/byte 的 FP64问题和大于400 ops/byte 的 FP16 $\to$ 32 问题），图7a 和图 7b 展示了所提单例 Stream-K 内核在性能上全面超越 cuBLAS 集合。

由于 Stream-K 试图通过增加内存工作量来提高内存约束计算的运行速度，因此在低于这些阈值的情况下出现 “不稳定 ”的相对性能并不奇怪。这表明了未来工作的几个方向，即

• 对内存受限区域进行单独的成本建模；
• 将具有更小图块尺寸的第二个 Stream-K 内核打包到双内核集成中。

Supplementary Material

Analytical Modeling for Stream-K Configuration

在实践中，使用尽可能多的可在 GPU 上活跃驻留的 CTA 来调用 Stream-K 分解并不总是有利的。因为这种方法是一种图块分割方法，它会产生比简单的数据并行分解更高的修复成本。因此，基本的命题是强扩展性：在额外的开销导致投资回报为负之前，可以表达多少额外的并行性。根据问题形状的不同，最优的分割可能足以填满整个处理器（即 $g\leftarrow p$），也可能完全不分割（即 $g\leftarrow t$），或者介于两者之间。

为了预测这一拐点，文中提出了一种简单的方法，将 Stream-K 的运行时间建模为网格大小 $g$ 的函数。在 GPU 上没有其他工作的情况下，整个 Stream-K 调度的运行时间将与其中一个输出图块 CTA 的运行时间相同，将其公式化如下：

$$\begin{array}{rl}tim{e}_{CTA}(g)\leftarrow & a+b(FixupPeers(g)>1)\\ & +c(ItersPerCta(g))+d(FixupPeers(g)-1)\end{array}$$
其中：
$$\begin{array}{rl}ItersPerCta(g)\leftarrow & \lceil \frac{\lceil \frac{m}{\text{BLK\_M}}\rceil \times \lceil \frac{n}{\text{BLK\_N}}\rceil \times \lceil \frac{k}{\text{BLK\_K}}\rceil }{g}\rceil \\ FixupPeers(g)\leftarrow & \lceil \frac{\lceil \frac{k}{\text{BLK\_K}}\rceil }{IterationsPerCta(g)}\rceil \end{array}$$

该 CTA 运行时模型由四个部分组成：

• $a$ 工作负载包括每个 CTA 产生的一次性、固定大小的成本，如网格启动延迟、初始强制缓存缺失、将最终输出瓦片写入 $\text{C}$ 的成本等。

• 第二部分 $b$ 包含了输出临时部分和的条件成本，适用于输出图块数量无法在整个处理器上完美量化的情况。

• 第三部分——每次迭代的负载 $c$ ——代表了每个 MAC 迭代的指令和停滞工作负载。

• 最后，每个协作者的负载 $d$ 是从处理相同图块的另一个 CTA 读取并累积部分和的成本。

工作负载常数集合 {$a$, $b$, $c$, $d$} 对于分块因子、矩阵数据类型和GPU 微架构的每种组合都是独一无二的，并且可以通过微基准测试来确定。

图 8 展示了论文中网格尺寸选择模型在 NVIDIA A100 GPU 上使用分块因子 BLK_M $=128$、BLK_N $=128$ 和 BLK_K $=32$ 参数化 fp16精度 GEMM 的行为。具体来说，作者强调了三种强缩放 GEMM 场景，在这些场景中，输出图块的数量不足以在处理器的 108 个 SM 核心上产生单个完整波。

第一个 GEMM 形状通过一个大尺寸的 $k$ 维累积，产生一个短而宽的输出矩阵。在这种情况下，与接缝修复成本的增加相比，MAC 循环时间的减少是单调改善的。因此，最优网格尺寸与 $g=108$ CTA 的最大并行度相吻合。

第二个形状通过中等尺寸的 $k$ 维累积，产生一个具有 64 个输出图块的方形矩阵。在这种情况下，$b$ 和 $d$ 的修复成本超过了 MAC 循环迭代次数的任何减少，正如 $g=64$ CTA 的全局最小值"凹陷"所见。

第三个形状通过一个巨大的 $k$ 维累积后产生一个单一的输出图块，类似于图 9 中的执行调度。虽然强扩展的机会相当大，但串行归约的每个对等点成本完全由单个 CTA 承担。当网格大小 $g>8$ 时，这些累积成本开始超过迭代次数的进一步减少。

参考资料：

• Stream-k matmul implementation very slow mostly because of if/else inside for loop #1393
• CUTLASS: Python API, Enhancements, and NVIDIA Hopper
• Z-order curve
• Morton/Z-order indexing
• Low utilization of warp per scheduler for bit pack operation (Morton Code)
• Thinking Parallel, Part III: Tree Construction on the GPU
• [翻译]Thinking Parallel, Part III: Tree Construction on the GPU 并行思考，第三部分：基于GPU的树状结构构建
• Multi-level Optimization of Matrix Multiplication for GPU-equipped Systems
• Extended Morton Codes for High Performance Bounding Volume Hierarchy Construction
• [QST] Why use blockIdx.z as the indexing method for batch #1050
• 2019-04-16-ModernCpp
• Optimizing Memory Access on GPUs using Morton Order Indexing
• Morton Order - Introduction
• 一种基于随机采样的SPM管理机制
• GPU基础：Occupancy、wave and tail effect
• 如何：使用过度订阅偏移延迟
• 优化 2 ：在 CPU – GPU 之间进行数据分区的直接内存访问
• 提高 GPU 内存超额订阅性能
• 数据中心网络的研究进展与趋势
• 基于分布式优化的数据中心网络混流调度机制
• 数据中心网络架构
• 数据中心网络架构浅谈（一）
• 【AI System】第7章：异构计算集群调度与资源管理系统
• Fast implementation of DGEMM on Fermi GPU论文笔记
• PPT - Automatic Tuning Matrix Multiplication Performance on Graphics Hardware
• git rebase 命令介绍
• [QST] Use stream-k feature in cutlass profiler #1004
• cute 之 简单GEMM实现