[论文解读] Stream-K: Work-centric Parallel Decomposition for Dense Matrix-Matrix Multiplication on the GPU

Muhammad Osama, Duane Merrill+3 more|arXiv (Cornell University)|Jan 9, 2023

Parallel Computing and Optimization Techniques被引用 5

一句话总结

Stream-K 在 GPU 上引入了一种以工作为中心的并行分解方法，用于密集矩阵-矩阵乘法（GEMM），将乘加（MAC）循环迭代均匀分配给流式多处理器（SMs），而非对输出矩阵进行分块。这种方法在各种问题形状下实现了接近 100% 的处理器利用率，相较于 cuBLAS 和 CUTLASS，最高可实现 14× 的峰值加速和 6.7× 的平均性能提升，且每个精度仅使用一个内核，而非复杂的内核集合。

ABSTRACT

We introduce Stream-K, a work-centric parallelization of matrix multiplication (GEMM) and related computations in dense linear algebra. Whereas contemporary decompositions are primarily tile-based, our method operates by partitioning an even share of the aggregate inner loop iterations among physical processing elements. This provides a near-perfect utilization of computing resources, regardless of how efficiently the output tiling for any given problem quantizes across the underlying processing elements. On GPU processors, our Stream-K parallelization of GEMM produces a peak speedup of up to 14$ imes$ and 6.7$ imes$, and an average performance response that is both higher and more consistent across 32,824 GEMM problem geometries than state-of-the-art math libraries such as CUTLASS and cuBLAS. Furthermore, we achieve this performance from a single tile size configuration per floating-point precision, whereas today's math libraries employ complex kernel-selection heuristics to select from a large ensemble of kernel variants.

研究动机与目标

• 为解决在宽 GPU 上基于分块的 GEMM 工作负载中因量化效率低下导致的性能下降问题。
• 通过在各种问题形状下实现一致且高性能的执行，消除数学库中复杂的内核选择启发式方法的需求。
• 通过在乘加循环级别重新定义工作量化，无论问题大小或形状如何，均实现接近 100% 的处理器利用率。
• 通过用每个精度仅一个高度可移植的内核替代大型内核集合，减少库的膨胀和维护复杂性。

提出的方法

• Stream-K 将乘加（MAC）循环迭代的总数在 GPU 流式多处理器（SMs）之间均匀分配，确保工作负载分布均衡。
• 它将 MAC 循环迭代作为工作量量化的基本单位，其成本恒定且最低，相较于整个输出分块而言。
• 该方法动态创建 O(p) 个分割缝（p 为 SM 数量），按处理器宽度而非问题大小扩展工作量分配。
• 它为每种浮点精度使用单一固定分块因子，避免了对多个内核变体或复杂选择启发式方法的需求。
• 该方法最小化了通信和同步开销，这些开销随 SM 数量增长，而非随问题大小增长。
• 该方法已集成至 CUTLASS 2.11，支持 FP64、FP16→FP32 及其他精度格式，并保持一致的高性能表现。

(a) Data parallel decomposition with grid size $g$ =9 CTAs, large $128\times 128\times 128$ CTA work volumes, and 75% processor utilization ceiling

实验结果

研究问题

• RQ1 以工作为中心的分解方法，若在 MAC 循环级别进行量化，是否能在现代 GPU 上实现比基于分块的 GEMM 更高且更一致的性能？
• RQ2 在不牺牲性能的前提下，每个精度仅用一个内核是否足以替代 cuBLAS 等数学库中庞大的内核集合？
• RQ3 在各种不同的问题几何形状下，Stream-K 的工作量平衡与基于分块的方法相比，在量化效率方面表现如何？
• RQ4 与复杂的内核选择启发式方法相比，采用统一的、以工作为中心的分解方法，其性能影响如何？

主要发现

• 在 32,824 种 GEMM 问题几何形状中，Stream-K 相较于 cuBLAS 实现了最高 14× 的峰值加速和 6.7× 的平均性能提升。
• 对于 FP64，Stream-K 的吞吐量比 cuBLAS 高 6%；对于 FP16→FP32，其吞吐量高出 13%。
• Stream-K 的性能分布范围比 cuBLAS 和理想化的 CUTLASS 模拟器（oracle）更窄，表明其具备更优的一致性和资源利用率。
• 在计算密集型场景下（FP64 每字节操作数 >150，FP16→FP32 每字节操作数 >400），Stream-K 持续优于 cuBLAS 内核集合。
• 与 cuBLAS 相比，Stream-K 将分发大小减少了最多 20×，仅用一个内核每精度即可实现峰值性能，而非数百个专用内核。
• 该方法无论问题形状如何，均实现了接近 100% 的处理器利用率，消除了基于分块方法因量化效率低下而产生的 75–90% 理论上限。

(b) Data parallel decomposition with grid size $g$ =18 CTAs, smaller $128\times 64\times 128$ CTA work volumes, and 90% processor utilization ceiling