USE CUTLASS TO FUSE MULTIPLE GEMMS TO EXTREME PERFORMANCE

BING LIU, DEVTECH

议程 (AGENDA)

• GEMM 性能分析
• 将2个GEMM融合成1个核函数
• 将多个GEMM融合成1个核函数
• 融合2个GEMM，一个更通用的实现
• 结论

背景 (BACKGROUND)

痛点

• 尽管GPU的峰值计算能力非常强大，但在模型中的小型GEMM（通用矩阵乘法）操作中很难发挥出较好的性能。
• 在TensorFlow等框架中，碎片化的核函数启动会带来严重的核函数启动开销。

GEMM性能分析 (GEMM PERFORMANCE ANALYSIS)

GEMM 相关基本概念

下图展示了GEMM操作在GPU上执行时的数据层次和划分方式，从全局内存（Global Memory）到共享内存（Shared Memory），再到寄存器文件（Register File），最终在CUDA核心或张量核心（Tensor Cores）上进行计算。

• Blocked GEMM: 整个GEMM操作被划分为多个块。
• Thread Block Tile: 每个线程块（Thread Block）处理一个数据块（Tile）。数据从全局内存加载到共享内存。
• Warp Tile: 线程块中的每个Warp处理共享内存中的一个更小的数据块。
• Thread Tile: 每个线程处理寄存器中的片段（Fragment），并由CUDA核心或张量核心执行计算。

CUTLASS 高效GEMM实现

CUTLASS的GEMM实现流程在标准的计算流程后增加了一个Epilogue（结尾）阶段。在主循环的矩阵乘累加（MMA）操作完成后，Epilogue阶段负责处理结果，例如进行激活函数、偏置加法等操作，并将最终结果写回全局内存。整个CTA（Cooperative Thread Array）的执行时间主要由K维度的迭代计算（K Iteration MMA）和Epilogue构成。

CUTLASS 主循环流水线实现

下图详细展示了GEMM/Conv操作的流水线，分为三个主要阶段：prologue（序言）、main loop（主循环）和epilogue（结尾）。

• Prologue（序言）

  • 准备MatA和MatB的数据切片。
  • 指令：LDG->STS->BAR->LDS。
  • 必须等待全局延迟。

• Main loop（主循环）

  • 从全局内存预取下一个MatA和MatB切片。
  • 从共享内存中获取片段A和B。
  • 执行FMA或MMA指令。
  • 重复此过程直到K维度完成。
  • 大多数全局内存延迟和共享内存延迟都可以被很好地隐藏。
  • 主循环时间与K的大小成正比。

• Epilogue（结尾）

  • 执行激活函数和加偏置。
  • 将累加器（accumulators）存储到共享内存并重组布局以最大化全局带宽。
  • 将MatC tile存储回全局内存。
  • 大部分时间花费在全局写回上，时间与MatC tile的大小成正比。

实现峰值TFlops的关键

• 实现峰值TFlops的公式为：Peak TFlops = SM_num * Clk_freq * IPC_per_sm。
• 换言之，我们需要在整个核函数运行期间让所有SM（Streaming Multiprocessor）单元保持繁忙。
• 尽管序言和结尾部分的时间会影响我们达到峰值性能，但它们是必不可少的。
• 我们希望最小化 "wave quant" 的影响，让所有SM都处于繁忙状态。
• 如下图所示，当K维度较小时（small k），主循环时间短，启动、加载和存储的开销占比高。当K维度较大时（large k），主循环占据了绝大部分时间，计算效率更高。

• 小规模的GEMM问题通常很棘手，有限的线程块数量和较小的K值会导致数学计算吞吐量低，GPU资源利用率低下。
• 如果应用中有大量的小规模GEMM，将导致GPU资源的利用率非常低。

张量操作融合（Tensor Op Fusion）

• 从CUTLASS 2.5开始，引入了张量操作融合的示例。
• 核心思想是：第一个GEMM的输出无需写回全局内存，可以直接被第二个GEMM重用。
• 如下图所示，非融合（Non-fused）方案需要两次完整的核函数启动、加载和存储。而融合（Fused）方案将两次计算合并在一个核函数中，减少了内存访问和启动开销。

将2个GEMM融合成1个核函数 (Fuse 2 GEMM into 1 kernel)

融合需求

将两个连续的GEMM操作融合到一个核函数中，需要满足以下要求：
1. 第一个和第二个GEMM可能有不同的CTA（线程块）配置。
2. 第二个GEMM的计算依赖于第一个GEMM的输出（Mat C0）。

可能的方法1 - 通过全局内存传递Mat C0

这种方法通过全局内存（希望命中L2缓存）来传递两个GEMM之间的中间结果。

• 方法：
  1. 使用协作组（cooperative group）启动。
  2. 将第一个GEMM的结果Mat C0写回全局内存。
  3. 同步该组内的所有线程块。
  4. 第二个GEMM加载Mat C0并继续计算。
• 限制：
  1. 只能启动有限数量的线程块，以避免死锁。
  2. Mat C0缓存在L2中。

可能的方法2 - 通过寄存器传递Mat C0 (CUTLASS)

这种方法在CTA内部通过寄存器直接传递中间结果，避免了全局内存的读写。

• 方法：
  1. 在一个CTA内安排4x1个warps。
  2. Warps之间无需同步。
  3. 中间结果的片段（Fragments）保存在寄存器中。
• 限制：
  1. N0_tile == N0 且 N1_tile == N1。
  2. 难以处理较大的N值，例如N=512。
  3. 寄存器压力会更高。
  4. 线程块数量为 M / M_tile。

方法2实现细节

下图展示了通过寄存器传递数据的具体流程。Warp 0,1,2,3并行计算。第一个GEMM的输入片段（Frag A0, Frag B0）计算得到FP32的累加结果（Accu C0），经过偏置、激活和类型转换后，作为第二个GEMM的输入片段（Frag A1），与Frag B1计算得到最终的累加结果（Accu C1），最后输出FP16的片段。

• 该流程展示了F16 * F16 + F32的操作，以及16-by-8-by-8的计算模式。
• 第二代和第三代张量核心的输入和输出布局相同，这使得片段的转换更加容易。

该方法将两个连续的通用矩阵乘法（GEMM）操作融合到一个CUDA内核中，避免了将中间结果写回全局内存。

数据流与计算过程：
1. 第一个GEMM:
* Warp 0, 1, 2, 3 从全局内存中获取矩阵 A0 (Frag A0) 和 B0 (Frag B0) 的分片。
* 执行矩阵乘累加操作，得到FP32精度的累加结果 C0 (Accu C0)。
2. 中间处理:
* 对累加结果 C0 应用 Bias 和激活函数（Activation）。
* 将结果转换（Cast）为FP16精度，作为下一个GEMM的输入矩阵 A1 (Frag A1)。这个中间结果 A1 保留在寄存器中，不写入全局内存。
3. 第二个GEMM:
* Warp 0, 1, 2, 3 获取第二个权重矩阵 B1 (Frag B1) 的分片。
* 使用寄存器中的 A1 和获取的 B1 执行矩阵乘累加，得到FP32精度的累加结果 C1 (Accu C1)。
4. Epilogue:
* 对最终的累加结果 C1 应用 Bias 和激活函数。
* 将最终结果（Output Frag FP16）写回全局内存。

伪代码逻辑:
* 代码部分:
* Prologue 0: 获取矩阵A0、B0的分片，存入共享内存。
* Main loop 0:
* 获取下一个A0、B0的分片。
* 获取当前的mma fragments A0和B0。
* 执行mma指令。
* Epilogue 0 & Prologue 1:
* 获取矩阵B1的分片。
* 添加Bias。
* 执行激活函数。
* 将B1存入共享内存。
* Cast: 将累加器转换为fragment A1。
* Main loop 1:
* 获取下一个B1分片。
* 获取当前的mma fragments B1。
* 执行mma指令。
* Epilogue 1: 添加Bias、执行激活函数，将结果存储到全局内存。

将多个GEMM融合成1个核函数

这个概念可以从两个GEMM扩展到融合任意N个GEMM。

如上图所示，计算链可以持续进行，前一个GEMM的输出（经过Bias、激活和类型转换后）直接作为下一个GEMM的输入，直到最后一个GEMM计算完成，其结果才被写回全局内存。

融合的优势

• 以融合3个GEMM为例。
• 通过融合，可以节省中间结果矩阵C0和C1的读写时间。

• 融合前 (Before): 每个GEMM都需要一次独立的kernel launch。每次启动都包含加载激活（Load Activation）、加载权重（Load Weights）、主循环（Main Loop）和存储激活（Store Activation）的完整过程。这导致了多次对全局内存的读写以及kernel launch的开销。
• 融合后 (After): 只需要一次kernel launch。第一个GEMM加载激活和权重，其后每个GEMM只需加载权重，中间的激活值在寄存器或共享内存中传递。只有最后一个GEMM完成后才将最终结果存储到全局内存。这显著减少了内存I/O和内核启动开销，从而带来潜在的性能提升（Potential Speed Up）。

自动化代码生成

为了避免针对不同输入形状、不同融合层和不同激活函数进行大量重复性工作，开发了一个Python脚本，该脚本基于cutlass自动生成相应的代码。

左侧是用于定义融合GEMM信息的Python配置（fuse_gemm_info），右侧是代码生成工具（40_multi_gemm_ir_and_codegen）的文件目录结构，展示了用于生成不同组件（如epilogue、kernel、device代码等）的模块。

配置脚本

特性 (Features):
* GEMM的M维度可以是动态的。
* 融合层的数量可以是动态的。

限制 (Restrictions):
* GEMM的N和K维度必须是确定的。
* N不应大于256。
* 矩阵A需要行主序（Row Major）布局。
* 矩阵B需要列主序（Column Major）布局。
* 矩阵C和D需要行主序布局。
* 当前原型仅支持FP16精度。
* 每一层的epilogue（收尾操作）需要确定。

未来工作 (Future Works):
* 支持N维度最高到512。
* 支持更多精度。
* 支持自定义epilogue代码生成。
* 可能会通过官方CUTLASS示例发布。

性能测试结果

以下图表展示了在不同GPU（T4, A100, A30, A10）上，使用不同矩阵尺寸配置的CUTLASS融合内核相对于未融合版本的加速比。

配置: Mx256x32; Mx64x256; Mx1x64

配置: Mx128x512; Mx64x128; Mx1x64

配置: Mx256x512; Mx128x256; Mx64x128; Mx32x64

配置: Mx128x128; Mx128x128; Mx128x128; Mx128x128;

将2个GEMM融合：一种更通用的实现

对上述实现的优缺点分析

优点 (Pros):
* 对融合层的数量没有限制。

缺点 (Cons):
* 所有的内核都必须被预定义。
* 由于N_tile和N必须相等，限制了适用性。
* 对具有较大N值问题的支持有限。

接下来将探讨如何改进这两个缺点。

更通用的方法：Split K

对于小规模问题，线程块（threadblocks）数量太少，无法有效占满整个GPU的计算资源。

如上图所示，在标准GEMM中，一个CTA（Cooperative Thread Array，等同于线程块）负责计算输出矩阵C0的一个分块。如果问题规模小，总的CTA数量不足，会导致GPU利用率低下。

Split K 原理

通过在K维度上进行拆分，可以增加并行度。

将K维度一分为二（K/2），启动两倍的CTA。例如，blockIdx.z = 0的CTA组处理前K/2的数据，blockIdx.z = 1的CTA组处理后K/2的数据。它们分别计算出一个部分的子矩阵C0（Sub Mat C0），从而并行地完成原有的归约（Reduction）任务。

最后，将不同CTA组计算出的子矩阵结果（Sub Mat C0 和 Sub Mat C1）相加，然后应用Bias和激活函数，得到最终的输出矩阵C。这种方法通过增加并行任务数量，提高了GPU在处理小规模或特定尺寸问题时的利用率。

该方法将两个连续的矩阵乘法（GEMM）操作融合成一个单一的GPU内核。其基本流程如下：
1. 第一个GEMM: 计算 Mat C0 = Mat A * Mat B0。这是一个常规的GEMM操作。
2. 第二个GEMM: 使用第一个GEMM的结果 Mat C0 作为输入，计算 Mat C1 = Mat C0 * Mat B1。这个阶段采用 "Split K" 实现，其中 Mat C0 的K维度被分割处理。
3. 最终处理: 将第二个GEMM产生的多个部分和（C0-C3）相加，然后加上偏置（Bias）并应用激活函数（Act），得到最终的输出矩阵 Mat C。

第二个GEMM的计算过程

第二个GEMM的计算由一个Cooperative Thread Array (CTA) 内的多个Warp协同完成。一个CTA负责计算输出矩阵的一个M_tile x N1_tile的块。这个计算过程通过一个循环（N_iter）来遍历N1维度上的所有tile。

• 迭代 0 (N_iter = 0): CTA处理N1维度的第一个tile。
  

• 迭代 1 (N_iter = 1): CTA处理N1维度的第二个tile。
  

• 迭代 2 (N_iter = 2): CTA处理N1维度的第三个tile。
  

• 迭代 3 (N_iter = 3): CTA处理N1维度的第四个tile。
  

融合Kernel的伪代码结构

整个融合Kernel的执行流程可以概括为两部分GEMM的串行执行：
- 第一个GEMM部分:
- Prologue 0
- Main loop 0
- Epilogue 0
- 数据类型转换 (Cast)
- 第二个GEMM部分:
- for iter = 0 to N_iter:
- Main loop 1
- Epilogue 1

其中，N_iter 的计算方式为 (N1 + N1_tile - 1) / N1_tile。

Split K 实现方法一：自旋锁 / 串行归约 (Spin Lock / Serial Reduction)

为了对第一个GEMM产生的中间结果（在K维度上分割）进行求和，可以采用串行归约的方式。这种方法使用自旋锁（Semaphore）来保证对同一输出内存位置的原子更新。

执行步骤:
1. 每个线程块（block）等待信号量（获取锁）。
2. 加载前一个CTA计算的结果作为当前矩阵C的初始值。
3. 执行Epilogue操作（例如，加上偏置，应用激活函数）。
4. 增加信号量的值并释放锁，以便下一个线程块可以继续。

代码示例 (CUTLASS):
代码片段展示了如何使用semaphore.wait()和semaphore.release()来同步线程块，确保对输出张量的更新是串行化的。

时间线分析

下图展示了使用自旋锁时多个CTA的执行时间线。每个CTA可以并行执行其主循环（Mainloop），但在Epilogue阶段，由于需要等待锁，导致对同一输出位置的更新操作被串行化。

性能建模

• 核心思想: 每个M_tile x N1大小的输出矩阵D块共享同一个锁。
• 性能公式 (版本1):
• M_blocks = (M + M_tile - 1) / M_tile
• N0_blocks = (N0 + N0_tile - 1) / N0_tile
• N1_iter = (N1 + N1_tile - 1) / N1_tile
• Time = (Mainloop_time + Epilogue_time) * (2 * N1_iter - 1) + (N0_blocks - 1) * Epilgoue_time
• Workspace_required = M_blocks * 1 * sizeof(int)
• Epilogue Read Bytes = MatrixC + (N0_blocks - 1) * MatrixD
• Epilogue WriteBytes = N0_blocks * MatrixD

• 性能公式 (版本2):
• Time = (Mainloop_time + Epilogue_time) * N1_iter + (N0_blocks - 1) * Epilgoue_time
• Workspace_required = M_blocks * N1_iter * sizeof(int)

性能测试结果 (T4 GPU)

以下图表展示了自旋锁方法在T4 GPU上的性能。

• 测试 1 (4096xN0xK0x128): 当N0维度增加时，由于锁竞争加剧，性能呈下降趋势。
  

• 测试 2 (4096x256xK0xN1): 当N1维度增加时，性能同样表现出下降趋势。
  

Split K 实现方法二：并行 + 归约 (Parallel + Reduction)

为了避免自旋锁导致的串行化瓶颈，可以采用并行方法。

执行步骤:
1. 每个线程块将其计算出的累加器结果（部分和）存储到一个独立的全局缓冲区中。
2. 所有线程块完成计算后，启动一个额外的归约内核（Reduction Kernel）来将全局缓冲区中的所有部分和相加，得到最终结果。

性能建模:
- Time = (Mainloop_time + Epilogue_time) * N1_iter + Reduction Kernel
- Workspace_required = N0_blocks * MatrixD
- 读写字节数需要额外包含归约内核的开销。

这种方法以额外的内存空间（Workspace）和一次额外的内核启动为代价，换取了更高的并行度。

实现细节与优化

• 在CUTLASS实现中，Epilogue阶段的每个warp倾向于将其整个Ntile元素存储到全局内存。
• Epilogue的伪代码通常包含 BAR.SYNC（线程块内同步障），以确保数据写入的一致性。
• 优化: 如果块内的warps配置为4x1的形状，它们在写操作时不会有冲突，此时可以移除BAR.SYNC，从而减少同步开销，提升性能。

性能对比 (Spin Lock vs Parallel Split K)

在T4 GPU上，针对4096xN0x128x128问题的性能对比结果如下：
- SpinLock: 在N0较小（锁竞争不激烈）时性能更优。
- Parallel Split K: 当N0增大时，其可扩展性优势显现，性能超越SpinLock。
- Parallel Split K no BAR: 移除同步障带来了稳定的小幅性能提升。

结论：自旋锁方法适用于低竞争场景，而并行归约方法在高竞争场景下更具优势。

在 T4 GPU 上，针对 MxN0x0xK0xN1 (4096x256x128xN1) 问题的 2GEMM 融合性能比较，对比了 Spin Lock、Parallel Split K 和 Parallel Split K no BAR 三种方法。结果显示，Parallel Split K no BAR 方法在各种维度下都表现出最优的性能。

Split K 实现方法三：原子操作 + 归约 (Atomic + Reduction)

此方法利用原子操作将累加器存储到带填充的全局缓冲区（padded global buffer），这需要一个额外的归约内核（reduction kernel）来完成最终计算。

其计算开销和所需资源如下：
1. 使用原子操作将累加器存储到带填充的全局缓冲区，需要一个额外的归约内核。

-   M_blocks = (M + M_tile - 1) / M_tile
-   N0_blocks = (N0 + N0_tile - 1) / N0_tile
-   N1_iter = (N1 + N1_tile - 1) / N1_tile
-   Time = (Mainloop_time + Epilogue_time) * N1_iter + Reduction Kernel
-   Workspace_required = padded MatrixD
-   Epilogue Read Bytes = MatrixC + MatrixD (包含归约内核)
-   Epilogue WriteBytes = (N0_blocks + 1) * MatrixD (包含归约内核)

性能对比

以下是在 T4 GPU 上，不同问题配置下的 2GEMM 融合性能对比图。

• Spin Lock vs Parallel Split K vs Atomic (MxN0xK0xN1, 4096xN0xK0x128)
  在此配置下，Atomic 方法的性能与 Parallel Split K no BAR 方法相当，两者均显著优于 Spin Lock 基线。

  

• Spin Lock vs Parallel Split K no BAR vs Atomic (MxN0xK0xN1, 4096x256x128xN1)
  在此配置下，结果与前一图类似，Atomic 和 Parallel Split K no BAR 方法性能接近，且远超 Spin Lock。

  

• Spin Lock vs Atomic Fair (MxN0x0xK0xN1)
  "Fair" 模式意味着将结果直接加到输出中，而不是写入带填充的中间工作区。与 Spin Lock 相比，Atomic Fair 方法在 N0 和 N1 维度上均表现出更好的性能。

  

• 2GEMM 融合：Atomic 性能
  下图展示了在 T4 和 A10 GPU 上，针对不同问题规模 (MxN0x0xK0xN1)，Atomic 方法的性能表现。

  

未来工作

未来的工作方向包括：
- 移除 memset 和归约内核（reduction kernel）。
- 实现混合原子操作（hybrid atomic）和自旋锁（spin lock）的实现。

结论

1. 我们介绍了使用 cutlass 来融合多层小规模 GEMM，以实现显著的性能提升。
2. 讨论了更灵活的两层 GEMM 融合方法，并证明其在某些情况下可以提升性能。
3. 我们也可以在 TensorFlow 等框架中使用这种方法，通过减少内核启动开销来获得额外的性能提升。