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

A1 主要贡献

本文介绍了一种名为 Stream-K 的新颖并行分解方法，该方法以工作为中心，用于稠密线性代数中的矩阵乘法（GEMM）及相关计算。与当前主要基于分块（tile-based）的分解方法不同，Stream-K 通过将总内循环迭代次数均匀地分配给物理处理单元来实现并行化。这种方法的核心优势在于，无论给定问题的输出分块如何有效地在底层处理单元上进行量化，它都能提供近乎完美的计算资源利用率。

核心问题： 传统的基于数据并行的 GEMM 实现通过对输出矩阵进行分块，并将独立的输出块分配给并发的线程组。当输出块的数量远大于处理器核心数时，这种方法能实现良好的负载均衡。然而，随着处理器核心数量的增加，这种超额订阅（oversubscription）的程度已大幅缩小。当输出块的数量不能被核心数整除时，就会出现量化低效（quantization inefficiency）问题，导致最后一波（wave）计算任务只有部分核心在工作，而其他核心处于空闲等待状态，从而严重降低了处理器利用率。如下图1a所示，在4个SM的GPU上计算384×384×128的GEMM，若划分为9个128x128的输出块，利用率上限仅为75%。尽管可以通过减小块大小（图1b）来提高利用率，但这会降低缓存效率，可能无法带来实际性能提升。

(a) 数据并行分解，网格大小g=9个CTA，采用大的128 × 128 × 128 CTA工作量，处理器利用率上限为75%


(b) 数据并行分解，网格大小g=18个CTA，采用较小的128 × 64 × 128 CTA工作量，处理器利用率上限为90%

图1. 在一个假设的四SM GPU上，对384 × 384 × 128 GEMM的数据并行执行调度。

当前数学库（如cuBLAS）的解决方案是部署一个包含多种分块配置的内核集合（ensemble），并使用复杂的启发式方法在运行时选择最佳内核。这种方法不仅导致库的体积庞大（cuBLAS达数百兆），而且启发式选择也未必总能找到最优配置。图2a展示了另一种方法，即固定分块（fixed-split），但同样难以实现完美的量化。

(a) 固定分块分解，分裂因子s=2，网格大小g=18个CTA，采用较小的128 × 128 × 64 CTA工作量，量化效率为90%


(b) 基本的Stream-K分解，网格大小g=4个CTA，采用较大的128 × 128 × 288 CTA工作量，量化效率接近100%
图2. 在一个假设的四SM GPU上，对384 × 384 × 128 GEMM的分块分裂执行调度。

研究目标与创新点： 本文提出的Stream-K分解方法旨在解决上述量化低效问题。其核心创新点在于：
1. 工作中心分解（Work-centric Decomposition）：Stream-K 将 GEMM 计算的总乘加累积（MAC）循环迭代次数在流式多处理器（SM）之间进行均匀分配（误差在一个迭代以内）。由于单个MAC循环迭代的工作量远小于整个输出块，核心工作负载的差异几乎可以忽略不计，从而实现了近乎完美的负载均衡和处理器利用率（如图2b所示）。
2. 单一内核配置：Stream-K 无论问题形状如何，都使用理想的分块因子，并且其通信开销与处理器宽度成比例，而非输出块数量。这使得它可以用单个内核来应对各种问题，避免了大型内核集合和复杂的选择启发式。
3. 高性能与高一致性：在NVIDIA A100 GPU上对32,824个GEMM问题形状的评估表明，与最先进的数学库CUTLASS和cuBLAS相比，Stream-K不仅性能响应更高、更一致，而且峰值加速比分别高达14倍和6.7倍。

A3 背景知识与现有工作分解策略

背景知识

通用矩阵乘法（GEMM）的定义与重要性。通用矩阵乘法（GEMM）定义为 $C = \alpha AB + \beta C$，其中 $\alpha$ 和 $\beta$ 是标量，A、B、C是矩阵。为简化讨论，本文假设 $\alpha=1, \beta=0$。GEMM问题的形状由其计算的体积范围定义，例如，一个 $m \times n \times k$ 的GEMM消耗一个 $m \times k$ 的输入矩阵A和一个 $k \times n$ 的输入矩阵B，执行 $m \times n \times k$ 次乘加累积操作，并产生一个 $m \times n$ 的输出矩阵C。GEMM是许多大规模工程和科学应用中性能关键的子程序，在LU、QR和Cholesky等矩阵分解方法中扮演重要角色，这些方法被广泛应用于工程、气候模拟、宇宙学、量子化学等高性能建模与仿真领域。

GEMM在深度学习中的核心地位。矩阵乘法也是现代深度学习（DL）方法的基础构建模块。深度神经网络（DNN）的训练通常在大型分布式系统上使用海量数据集进行【13, MLPerf Training Benchmark, 2020】。许多DL训练和推理操作都被转换为矩阵乘法。例如，依赖卷积的图像识别和计算机视觉模型，其卷积操作可以直接实现为滤波器和图像数据集的乘积【4, cuDNN: Efficient Primitives for Deep Learning, 2014】。主导自然语言处理等应用的Transformer架构，其性能几乎完全受限于大型矩阵乘法的性能。

GPU上GEMM实现的历史演进。Larsen和McAllister的早期工作将GPU矩阵乘法框架化为多纹理乘法和混合操作【11, Fast Matrix Multiplies using Graphics Hardware, 2001】。后续GPU架构提供的用户可编程共享内存使得更高性能的数据并行方案成为可能，这些方案采用两级分块（共享内存和寄存器），块大小通过广泛的微基准测试分析【2, 14, 17, 19】和自动调优【5, 7, 12】来确定。

针对不同GEMM问题形状的优化库。MAGMA GPU数学库可能是最早为多样化GEMM问题形状进行优化的库之一【9, Autotuning GEMM Kernels for the Fermi GPU, 2012】。其解决方案是将一组受限的分块参数应用于模板化的CUDA C++代码模板，为每个API原语（如hgemm_tt()）生成数百个数据并行变体。通过评估这些变体，提炼出一个由三到五个内核组成的小型集合，该集合在各种问题形状上均表现良好。给定问题的内核选择和分发由简单的手写规则表示的大小阈值控制。

现代GPU数学库的复杂性。后续的GPU数学库采用了更复杂的代码生成和内核选择组件。例如，ISAAC项目使用机器学习技术来预测给定GEMM形状的最佳分块和/或分裂参数化，然后可以通过PTX级代码生成器在线或离线实例化【18, Input-Aware Auto-Tuning of Compute-Bound HPC Kernels, 2017】。

NVIDIA cuBLAS的实现策略。NVIDIA的cuBLAS库【15, CUDA cuBLAS Library (v9.2), 2020】提供了一个扩展的cublasGemmEx接口，允许调用者从24种不同的GEMM“算法”中选择。在使用默认接口时，经过精心训练的启发式方法会从这个巨大的备选空间中进行选择。这些算法实现了各种不同的数据并行和固定分裂变体，并且cuBLAS通常为了代码优化而将每个变体组装成其自己的特定于架构的内核程序。GEMM API功能、策略性变体和微架构的交叉组合导致了其发行版越来越庞大，可执行代码超过数百兆字节。

新兴的编程模型和领域特定语言。鉴于当代深度学习的快节奏和快速变化的特性，最近的工作集中在简化表达和构建高性能内核的编程模型上，这些内核可以修改或补充GEMM计算。CUTLASS C++库【8, CUTLASS: Fast Linear Algebra in CUDA C++, 2017】为在GPU线程层次结构的所有级别上组合自定义的类GEMM计算提供了数据移动和乘加累积类。Triton【19, Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, 2019】是一种用于张量编程的领域特定语言，其核心是块/瓦片概念的表达、转换和优化。其他领域特定编程语言如Halide【16, Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines, 2013】和TVM【3, TVM: End-to-End Optimization Stack for Deep Learning, 2018】将逐点运算符的表达与循环调度分离开。Fireiron【6, Fireiron: A Data-Movement-Aware Scheduling Language for GPUs, 2020】进一步将数据移动构造添加到调度语法中。

现有工作分解策略

利用局部存储以实现计算密集。现代处理器通常将矩阵A、B和C存储在较大、较慢、较远的内存中，并可以访问一个较小、较快的高速缓存或暂存器内存。任何GEMM实现的主要目标都是利用这些本地存储资源，使得最终的实现是计算受限的。

3.1 顺序缓存分块

经典六重循环结构。GEMM的经典缓存分块（cache-blocked）公式将其计算量划分为块，并选择一种能够暴露内存局部性的遍历顺序。算法1展示了一个简化的实现，包含六个循环。最内层的三个循环在分块因子BLK_M、BLK_N和BLK_K内迭代，而最外层的三个循环则跨越这些块进行迭代。如果缓存能够容纳来自三个矩阵的各一个块，那么这些元素之间的数据重用将显著减少对最后一级内存的访问次数【10, The Cache Performance and Optimizations of Blocked Algorithms, 1991】。

// 算法 1 顺序缓存分块GEMM
1: ▷ tile-processing outer loops   
2: for mm ← 0 to m step BLK_M do   
3:   for nn ← 0 to n step BLK_N do   
4:     ▷ zero-initialize output tile   
5:     for mmm ← mm to (mm + BLK_M) do   
6:       for nnn ← nn to (nn + BLK_N) do   
7:         C[mmm,nnn] ← 0   
8:       end for   
9:     end for   
10:    ▷ perform the MAC iterations for this tile   
11:    for kk ← 0 to k step BLK_K do   
12:      ▷ MAC iteration (fully unrolled)   
13:      for mmm ← mm to (mm + BLK_M) do   
14:        for nnn ← nn to (nn + BLK_N) do   
15:          for kkk ← kk to (kk + BLK_K) do   
16:            C[mmm,nnn] ← C[mmm, nnn] +   
17:              (A[mmm,kkk] × B[kkk,nnn])   
18:          end for   
19:        end for   
20:      end for   
21:    end for   
22:  end for
23: end for

3.2 数据并行

GPU上的并行分解。如算法2所示，GEMM的数据并行GPU公式被分解到一个由并行线程块（或称为协同线程阵列，CTA）组成的网格中。网格的大小被设定为每个CTA产生其自己的(BLK_M × BLK_N)输出块。为便于说明，算法3中的MacLoop()子程序封装了计算CTA输出块值的乘加累积工作负载。它在累积域（例如GEMM的k轴）执行一系列MAC循环迭代。每次MAC循环迭代包含一个线程级别的(BLK_M × BLK_N × BLK_K) / CTA_THREADS的MAC操作量。在计算过程中，输入矩阵的片段被暂存到SM的共享内存中，以供单个线程局部重用。

MacLoop()子程序的复杂实现。尽管算法3中MacLoop()的展示方式是每个输出块元素部署一个线程，但CUTLASS【8, CUTLASS: Fast Linear Algebra in CUDA C++, 2017】和cuBLAS中的复杂实现会：(1) 完全展开每个线程的MAC循环迭代；(2) 在warp和/或线程级别实现额外的分块；(3) 在MAC循环迭代之间精心安排一个共享内存数据移动的软件流水线。

数据并行分解的局限性。不幸的是，这种经典的数据并行分解在现代GPU上容易遭受量化低效的影响，如图1所示。尽管一个包含多种分块因子的集合可能会发现提高处理器利用率的机会，但对于任意问题大小，它不太可能实现完美的量化。此外，较小的分块因子有两个缺点：(1) 在流水线实现中，每次MAC循环迭代的指令数较少，难以覆盖全局和共享内存传输的延迟；(2) 内存操作相对于MAC指令的比例更高，这可能导致它们无法成为计算受限的。

// 算法 2 数据并行GPU GEMM
1: _shared_ accum[BLK_M,BLK_N]   
2: iters_per_tile ← ⌈k/BLK_K⌉   
3: ▷ instantiate one CTA per output tile   
4: fork CTA[x] in [ ⌈m/BLK_M⌉ × ⌈n/BLK_N⌉ ] do   
5:   ▷ perform the MAC iterations for this tile   
6:   accum ← MacLoop(x, 0, iters_per_tile)   
7:   ▷ store accumulators to output tile   
8:   StoreTile(C, x, accum)   
9: join

// 算法 3 MacLoop子程序
1: procedure MacLoop(tile_idx, iter_begin, iter_end)   
2:   _shared_ accum[BLK_M,BLK_N]   
3:   _shared_ frag_a[BLK_M,BLK_K]   
4:   _shared_ frag_b[BLK_K,BLK_N]   
5:   ▷ determine output tile coordinates   
6:   mm ← BLK_M × (tile_idx / ⌈m/BLK_M⌉)   
7:   nn ← BLK_N × (tile_idx % ⌈m/BLK_M⌉)   
8:   ▷ zero-initialize local accumulator storage   
9:   accum ← 0   
10:  ▷ perform the specified range of MAC iters for this tile   
11:  for iter ← iter_begin to iter_end do   
12:    kk ← iter × BLK_K   
13:    ▷ copy global matrix fragments to local storage   
14:    frag_a ← LoadFragment(A, mm, kk)   
15:    frag_b ← LoadFragment(B, kk, nn)   
16:    fork THREAD[mmm,nnn] in [BLK_M, BLK_N] do   
17:      ▷ MAC iteration per thread (fully unrolled)   
18:      for kkk ← 0 to BLK_K do   
19:        accum[mmm, nnn] ← accum[mmm,nnn] +   
20:          (frag_a[mmm,kkk] × frag_b[kkk,nnn])   
21:      end for   
22:    join   
23:  end for   
24:  return accum   
25: end procedure

3.3 固定分块 (Fixed-split)

沿累积维度并行化。另一种方法是通过在累积维度上进行并行化来减小分配给每个CTA的工作粒度。对于给定的输出块，加法的结合律允许将迭代域在多个并发的CTA之间进行划分，然后通过一个依赖的“修复”（fixup）步骤来归约每个CTA计算出的部分和。我们在算法4中突出了这种固定分块方法，其中每个输出块由s个CTA协同产生。值得注意的是，当分裂因子s=1时，它的功能与数据并行分解完全相同。

固定分块的实现与缺点。固定分块分解也出现在CUTLASS和cuBLAS中。分裂因子被实现为一个运行时参数，允许单个内核可执行文件支持多种工作量，同时保留理想的分块因子以实现最佳的数据共享和延迟隐藏。然而，如图2a所示，通过均匀地分块来达到完美量化的可能性很小。此外，通信和同步的额外开销会随着整体问题大小和分裂因子的增加而增加。

// 算法 4 使用分裂因子s的固定分块GPU GEMM
1: _shared_ accum[BLK_M,BLK_N]   
2: iters_per_tile ← ⌈k/BLK_K⌉   
3: iters_per_split ← ⌈iters_per_tile/s⌉   
4: ▷ instantiate s CTAs per output tile   
5: fork CTA[x,y] in [ ⌈m/BLK_M⌉ × ⌈n/BLK_N⌉, s] do   
6:   ▷ perform the range of MAC iterations for this split   
7:   iter ← y × iters_per_split   
8:   iter_end ← min(iters_per_tile, iter + iters_per_split)   
9:   accum ← MacLoop(x, iter, iter_end)   
10:  ▷ consolidate partial-sums across CTAs   
11:  if y ≠ 0 then   
12:    ▷ store accumulators to temporary global storage   
13:    StorePartials(partials[x,y], accum)   
14:    Signal(flags[x,y])   
15:  else   
16:    ▷ accumulate partial sums from other CTAs contributing to this tile   
17:    for cta ← 1 to s do   
18:      Wait(flags[x,cta])   
19:      accum ← accum + LoadPartials(partials[x,cta])   
20:    end for   
21:    ▷ store accumulators to output tile   
22:    StoreTile(C, tile_id, accum)   
23:  end if   
24: join

A2 方法细节

4. Stream-K 分解

Stream-K的核心思想。我们的Stream-K分解是一种分块分裂（tile-splitting）的并行化方法，其中分裂的边界与分块结构本身完全分离。尽管我们采用熟悉的分块和切片策略来实现数据重用，但我们将GEMM计算量化为MAC循环迭代，即CTA范围内的BLK_M × BLK_N × BLK_K的小工作单元。如算法5所示，Stream-K将GEMM的总MAC循环迭代工作负载均匀地划分给一个固定大小为g的CTA网格。每个CTA的MAC循环迭代范围被连续地映射到GEMM形状的 $m \to n \to k$ 线性化空间中，并可能跨越输出块的边界。

部分和的合并机制。如果某个CTA的起始和/或结束迭代与块边界不重合（这预计是常见情况），它必须将其部分结果与覆盖同一块的其他CTA的结果进行合并。在这个基本实现中，C中的每个输出块由执行该块k=0 MAC循环迭代的CTA写入。然而，在写入之前，它必须累积从其他CTA共享到临时全局存储中的任何部分和。值得注意的是，Stream-K的通信、同步和全局存储开销与问题大小无关，而是与CTA的数量g成比例。

同步等待的优化。Stream-K的第二个好处是，当输出块的数量大于CTA的数量时，同步等待可能可以忽略不计。在这种情况下，每个输出块最多被两个CTA覆盖，并且块处理的偏斜（skew）确保了进行累加的CTA直到其协作者完成生成贡献值很久之后才需要这些值。

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

与固定分块的比较。此外，单个MAC循环迭代的工作量比整个输出块小32倍。因此，一个32路（32-way）的固定分块分解也能提供100%的量化效率，但代价是8倍大的“修复”开销。此外，由于在共享部分和时，写入者和读取者之间存在时间上的偏斜，Stream-K能够更好地隐藏CTA间的同步延迟。

Stream-K的通用性。Stream-K也可以泛化为固定分块和数据并行分解。当网格大小g是输出块数量的偶数倍时，Stream-K的行为与固定分块分解完全相同。同样，当g等于输出块数量时，Stream-K的行为与数据并行分解完全相同。我们利用这种通用性，在第5.2节中创建了一个Stream-K分解的优化混合版本。

// 算法 5 Stream-K GPU GEMM
1: _shared_ accum[BLK_M,BLK_N]   
2: iters_per_tile ← ⌈k/BLK_K⌉   
3: total_iters ← ⌈m/BLK_M⌉ × ⌈n/ BLK_N⌉ × iters_per_tile   
4: iters_per_cta ← ⌈total_iters / g⌉   
5: ▷ instantiate g CTAs   
6: fork CTA[x] in [g] do   
7:   iter ← x × iters_per_cta   
8:   iter_end ← iter + iters_per_cta   
9:   ▷ iteration-processing outer loop   
10:  while iter < iter_end do   
11:    tile_idx ← iter / iters_per_tile   
12:    tile_iter ← tile_idx × iters_per_tile   
13:    tile_iter_end ← tile_iter + iters_per_tile   
14:    ▷ perform the range of MAC iterations for this tile   
15:    local_iter ← iter - tile_iter   
16:    local_iter_end ← min(iter_end, tile_iter_end) - tile_iter   
17:    accum ← MacLoop(tile_id, local_iter, local_iter_end)   
18:    ▷ consolidate partial-sums across CTAs   
19:    tile_started ← iter = tile_iter   
20:    tile_ended ← (iter_end ≥ tile_iter_end)   
21:    if ¬tile_started then   
22:      ▷ store accum to temporary global storage   
23:      StorePartials(partials[x], accum)   
24:      Signal(flags[x])   
25:    else   
26:      ▷ store accumulators to output tile   
27:      if ¬tile_ended then   
28:        ▷ accumulate partial sums from other CTA contributing to this tile   
29:        cta_end ← tile_iter_end / iters_per_tile   
30:        for cta ← (x+1) in cta_end do   
31:          Wait(flags[cta])   
32:          accum ← accum + LoadPartials(partials[cta])   
33:        end for   
34:      end if   
35:      StoreTile(C, tile_id, accum)   
36:    end if   
37:    iter ← tile_iter_end   
38:  end while   
39: join

5. 实现细节

实现目标与透明性。我们在上一节中介绍的工作分解可以以多种不同的方式实例化，以适应不同硬件架构和软件库设计的需求。我们的实现针对NVIDIA GPU，并设计用于集成到现有的库如cuBLAS和CUTLASS中。在本节中，我们描述了如何配置我们启动的内核，并引入了一种混合方案，以帮助确保用户在尽可能广泛的问题形状范围内实现最大的GEMM性能。我们还强调，这些都是真正的内部实现细节。它们对类似BLAS库的用户是完全透明的，并且不会改变库的接口。唯一可观察到的影响是我们在第6节中分析的改进的性能特性。

5.1 内核配置

块大小的选择。为GEMM计算选择的块大小无疑是控制GEMM内核性能的关键参数。对于现代NVIDIA GPU，合适的块大小由GPU张量核心（Tensor Cores）支持的矩阵形状决定。基于广泛的经验，我们为每种支持的精度选择了能够为非常大的GEMM计算量达到GPU峰值TFLOP/s 99%的最小CTA级块大小。对于我们实验中使用的NVIDIA A100 GPU，这些尺寸对于FP64问题是64×64×16，对于FP16→32问题是128×128×32。

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

网格大小选择的启发式模型。我们的网格大小选择启发式基于一个简单的分析模型，该模型在均匀分配每个CTA的MAC循环迭代的同时，最小化读取、写入和累加部分和的成本。此分析模型的细节在补充材料（附录A.1）中提供。模型的参数可以通过经验测量轻松选择，并且每个目标架构只需进行一次。然后，得到的参数可以静态编译到库中。这再次与基于集合的方法形成对比，后者依赖于可能复杂的启发式和机器学习模型在运行时进行内核选择。

5.2 数据并行混合

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

混合方案的动机。块处理偏斜是Stream-K工作负载平衡策略的直接后果。然而，我们可以采取措施限制其持续时间，方法是将Stream-K的迭代平衡应用于总迭代域中一个较小的、与块对齐的区域，以便剩余的块可以以完整的、时间上对齐的波次产生。

“数据并行 + 单块Stream-K”方案。最简单的混合方案是图3b所示的“数据并行 + 单块Stream-K”调度。它仅在为最后一个部分填充的数据并行波次而剩余的块之间应用迭代平衡。完整波次的总数是 $w = \lfloor t/p \rfloor$，其中t是输出块的数量，p是GPU中的SM核心数。因此，每个Stream-K CTA接收到的迭代份额均匀，但少于一个块的量。不幸的是，当三个或更多CTA覆盖同一块时，此策略在隐藏交换部分和的同步延迟方面能力有限。在这些情况下，累加的CTA可能被迫等待其他CTA的贡献变得可见，因为除了最后一个CTA外，所有CTA都将大致在同一时间完成其最终迭代。此外，我们用于聚合部分和的基本方案在单个CTA内是串行化的，因此当每个块的贡献CTA数量很大时，可能会导致SM工作负载不平衡。


图3. 在一个假设的四SM GPU上，对896 × 384 × 128 GEMM的基本Stream-K与混合执行调度的对比。

“两块Stream-K + 数据并行”方案。我们用图3c所示的“两块Stream-K + 数据并行”混合调度来解决这些问题。它通过减少一个完整的数据并行波次，换取每个Stream-K CTA接收超过一个块但少于两个块的迭代量。当 $w \ge 2$ 时，这提供了更好的延迟隐藏，并且每个累加的CTA只需要从另一个贡献CTA接收部分和。否则，它的行为与“DP + 单块SK”调度完全相同。这种混合方法改善了内存访问模式和延迟隐藏。它还展示了通用Stream-K循环结构在同一内核实例内实现不同调度策略的多功能性。

A4 实验

实验环境


图4. 用于性能评估的32,824个GEMM问题形状和大小的测试域。{m} = {128 ... 8192}, {n} = {128 ... 8192}, {k} = {128 ... 8192}

实验结果

评估方法：我们为每种精度构建了一个单一的Stream-K内核，该内核实现了“两块Stream-K + 数据并行”混合分解。将其与以下对象进行比较：
1. CUTLASS数据并行：使用相同分块因子的默认数据并行CUTLASS内核。
2. cuBLAS：对应精度的cuBLAS内核集合（CUDA 11.6）。
3. CUTLASS Oracle：一个理想化的预言机，总能为给定GEMM实例选择性能最高的数据并行CUTLASS分块因子。
- FP64 Oracle集合：{(32×32×16), (32×64×16), (64×64×16), (64×128×16), (128×128×16)}
- FP16→32 Oracle集合：{(64×64×64), (64×128×32), (128×128×32), (128×256×32)}

性能分析：

  1. 与CUTLASS数据并行的比较：

    • 性能一致性：图5a和6a的Roofline图显示，单一的数据并行内核性能分布非常分散。相比之下，图5d和6d中的Stream-K内核性能响应要紧凑得多，更贴近机器的性能上限。
    • 平均与峰值加速比：如表1和表2所示，Stream-K内核在FP64和FP16→32上分别比其数据并行对应版本平均快1.23倍和1.63倍。在m×n小而k大的强扩展场景中，峰值加速比分别高达5.63倍和14.7倍。

    表1. Stream-K FP64 相对性能

    表2. Stream-K FP16→32 相对性能

  2. 与cuBLAS的比较：

    • 平均与峰值性能提升：表1和表2的第二列显示，我们的FP64和FP16→32 Stream-K内核分别比对应的cuBLAS集合平均吞吐量高6%和13%，峰值性能提升分别达到2.55倍和6.74倍。这是一个显著的改进，尤其考虑到Stream-K仅用单个内核（可执行代码量减少20倍）就实现了这一点。

  3. 与CUTLASS Oracle的比较：

    • 启发式选择的挑战：通过对比cuBLAS（图5b, 6b）和理想化的CUTLASS Oracle（图5c, 6c）的性能，可以看出设计一个在所有情况下都能选出最优内核的启发式算法是非常困难的。cuBLAS的性能分布范围比Oracle大得多。
    • Stream-K的优越性：Stream-K的性能分布比理想化的Oracle还要窄，峰值性能甚至比Oracle高出4.6倍，这突显了它能够达到传统以分块为中心的工作分解方法无法企及的利用率水平。


图5. 在NVIDIA A100上，针对32K个GEMM问题形状和大小的FP16→FP32 GEMM“roofline”性能利用率图。


图6. 在NVIDIA A100上，针对32K个问题形状和大小的FP64 GEMM“roofline”性能利用率图。

  1. 计算密集型与访存密集型问题分析：
    • 访存密集型场景：我们观察到，在一些小的、受带宽限制的问题形状上，我们较大的分块因子与cuBLAS相比不具优势。这是因为Stream-K通过增加内存工作负载来尝试加速内存受限的计算。
    • 计算密集型场景：然而，如果我们将范围限制在计算密集型问题（FP64问题计算强度 > 150 ops/byte，FP16→32问题 > 400 ops/byte），图7a和7b显示，我们的单一Stream-K内核实现了比cuBLAS集合一致更高的性能。
    • 未来工作启示：这表明未来的工作可以探索为访存密集型场景建立独立的成本模型，或者将一个具有较小块大小的第二个Stream-K内核捆绑成一个双内核集合。


图7. Stream-K相对于cuBLAS的加速比。

A5 结论

本文提出了一种名为Stream-K的新型并行工作负载分解技术，用于在GPU等宽架构上调度通用矩阵乘法（GEMM）及类似计算。与其他分块分裂技术不同，Stream-K以MAC循环迭代作为跨处理器核心的工作负载量化单位。这带来了卓越的强扩展性和工作负载平衡能力，因为其成本（1）相对于问题形状是恒定的，（2）远小于整个输出块的成本。

此外，Stream-K产生的分割边界数量为 $O(p)$，受处理器核心数量限制。因此，强扩展和工作负载平衡的开销与处理器宽度成比例，而不是问题大小。这对于许多无法承担分配与问题输出等量的大量临时存储的应用来说是一个受欢迎的特性。

最后，我们对Stream-K方法在广泛的GEMM形状和大小范围内进行了评估。我们证明了，单一分块配置的Stream-K可以（1）达到与NVIDIA cuBLAS库相当甚至超过其的绝对性能水平，即使后者在接近处理器峰值利用率下运行；（2）并且性能一致性要高得多。此外，Stream-K对于库的构建和维护是一个有吸引力的选择，因为它提供了将发行版大小减少一个数量级的机会，并且无需复杂的硬编码启发式或机器学习模型进行内核选择，同时不牺牲性能。Stream-K已在CUTLASS 2.11中开源，本文展示的性能可以在使用CUDA 11.8编译时复现。

对于未来的工作，我们确定了缓存感知的块访问模式（如莫顿序）作为一个优化方向。我们还相信，Stream-K分解可以为其他同样受量化低效问题困扰的类GEMM工作负载提供类似的性能提升。

A6 附录

A.1 用于Stream-K配置的分析模型

配置Stream-K的挑战。在实践中，并非总是以GPU上可驻留的最大CTA数量来调用Stream-K分解就是最优的。因为它是一种分块分裂方法，会产生超出简单数据并行分解的修复成本。因此，基本问题是强扩展性问题：在额外开销导致负投资回报之前，可以表达多少额外的并行性。根据问题形状，最佳分裂可能足以填满整个处理器（即 $g \leftarrow p$），或者完全不分裂（即 $g \leftarrow t$），或者介于两者之间。

运行时分析模型。为了预测这个拐点，我们提出了一种简单的方法来模拟Stream-K的运行时作为网格大小g的函数。在GPU上没有其他工作的情况下，整个Stream-K调度的运行时将与其一个输出块的CTA的运行时相同，我们将其表示如下：

其中：

模型组件说明。这个CTA运行时模型包含四个组成部分。
- a (固定成本)：a工作负载涵盖了每个CTA产生的一次性、固定大小的成本，例如网格启动延迟、初始的强制性缓存未命中、将最终输出块写入C的成本等。
- b (部分和输出成本)：第二个组成部分b包含了在输出块数量不能在处理器上完美量化的情况下，输出临时部分和的条件成本。
- c (每迭代成本)：第三个部分——每迭代工作负载c——代表了每次MAC迭代的指令和停顿工作负载。
- d (协作成本)：最后一个部分，每个协作者的工作负载d是读取和累加来自覆盖同一块的另一个CTA的部分和的成本。
工作负载常数集合 {a, b, c, d} 对于分块因子、矩阵数据类型和GPU微架构的每种组合都是唯一的，并且可以通过微基准测试凭经验确定。

模型行为示例。图8展示了我们的网格大小选择模型的行为，该模型针对NVIDIA A100 GPU上的fp16精度GEMM进行了参数化，使用的分块因子为 BLK_M = 128, BLK_N = 128, BLK_K = 32。具体来说，我们突出了三种强扩展GEMM场景，其中输出块的数量不足以在处理器的108个SM核心上产生一个完整的波次。


图8. 在NVIDIA A100（108个SM）上，针对BLK_M=128, BLK_N=128, BLK_K=32的Stream-K性能模型。

  1. 大k维，短宽输出：第一个GEMM形状通过一个大尺寸的k维度进行累加，以产生一个短而宽的输出矩阵。在这种情况下，MAC循环时间的减少相对于修复边界成本的增加是单调改善的。因此，最佳网格大小与最大并行度一致，即 $g = 108$ 个CTA。
  2. 中等k维，方形输出：第二个形状通过一个中等大小的k维度进行累加，产生一个有64个输出块的方形矩阵。在这种情况下，b和d的修复成本超过了MAC循环迭代次数的任何减少，如在 $g = 64$ 个CTA处的全局最小值“凹陷”所示。
  3. 单输出块，巨大k维：第三个形状在通过一个巨大的k维度进行累加后产生单个输出块，类似于图9中的执行调度。尽管强扩展的机会非常大，但每个对等的串行归约成本完全由单个CTA承担。当网格大小 $g > 8$ 时，这些累加成本开始超过迭代次数的任何进一步减少。


图9. 在一个假设的四SM GPU上，对128 × 128 × 384 GEMM的数据并行和Stream-K执行调度的强扩展比较。数据并行导致巨大的k维度在单个CTA内顺序处理，而Stream-K能够利用k维度上可用的并行性。