深入剖析 CUTLASS Ping-Pong GEMM 内核

作者：Less Wright, Adnan Hoque

Figure 1. FP8 GEMM Throughput Comparison CUTLASS vs Triton

图 1. FP8 GEMM 吞吐量比较：CUTLASS vs Triton

摘要

本文概述了 CUTLASS Ping-Pong GEMM 内核，并提供了相关的 FP8 推理内核基准测试。

Ping-Pong 是 Hopper GPU 架构上最快的矩阵乘法 (GEMM) 内核架构之一。Ping-Pong 是 Warp Group Specialized Persistent Kernels 系列的成员，该系列包括 Cooperative 和 Ping-Pong 变体。与之前的 GPU 相比，Hopper 巨大的张量核心计算能力需要深度的异步软件流水线才能实现峰值性能。

Ping-Pong 和 Cooperative 内核是这种范例的典范，其关键设计模式是持久化内核以分摊启动和序言开销，以及“一切皆异步”，通过具有两个消费者和一个生产者的专用 Warp Group 创建一个高度重叠的处理流水线，能够持续向张量核心提供数据。

当 H100 (Hopper) GPU 发布时，Nvidia 将其宣传为第一个真正的异步 GPU。这一声明强调了 H100 特定的内核架构也必须是异步的，才能最大限度地提高计算/GEMM 吞吐量。

CUTLASS 3.x 中引入的 Ping-Pong GEMM 就是这方面的例证，它将内核的所有方面都转移到了“完全异步”的处理范例中。在这篇博客中，我们将展示 Ping-Pong 内核设计的核心特性，以及它在推理工作负载中与 cublas 和 triton split-k 内核的性能对比。

Ping-Pong 内核设计

Ping-Pong（或者更准确地说，“sm90_gemm_tma_warpspecialized_pingpong”）通过异步流水线运行，利用了 warp 特化。与更经典的同质内核不同，“warp group”承担着特定的角色。请注意，一个 warp group 由 4 个 warps 组成，每个 warp 有 32 个线程，总共 128 个线程。

在较早的架构上，延迟通常通过每个 SM 运行多个线程块来隐藏。然而，对于 Hopper，张量核心吞吐量非常高，以至于需要转向更深的流水线。这些更深的流水线反而阻碍了每个 SM 运行多个线程块。因此，持久化线程块现在跨多个 Tile 和多个 Warp Group 发出集体主循环。线程块集群根据总 SM 计数进行分配。

对于 Ping-Pong，每个 warp group 承担一个特定的角色，要么是数据生产者，要么是数据消费者。

生产者 warp group 专注于产生数据移动，以填充共享内存缓冲区（通过 TMA）。另外两个 warp group 是专用的消费者，它们处理数学（MMA）部分，使用张量核心，然后执行后续工作并将结果写回全局内存（epilogue）。

生产者 warp group 与 TMA (Tensor Memory Accelerator) 协同工作，并且被有意保持尽可能轻量。事实上，在 Ping-Pong 中，它们有意减少寄存器资源以提高占用率。生产者将其最大寄存器计数减少 40，而消费者将其最大寄存器计数增加 232，这种效果可以在 CUTLASS 源代码和相应的 SASS 中看到。

source code

Ping-Pong 的独特之处在于，每个消费者都在不同的 C 输出 Tile 上工作。（作为参考，Cooperative 内核与 Ping-Pong 基本等同，但两个消费者组都在同一个 C 输出 Tile 上工作）。此外，两个消费者 warp group 然后将它们的工作分配给主循环 MMA 和 epilogue。

这显示在下图所示：

Figure 2: An overview of the Ping-Pong Kernel pipeline. Time moves left to right.

图 2：Ping-Pong 内核流水线概述。时间从左向右移动。

通过拥有两个消费者，意味着一个可以使用张量核心进行 MMA，而另一个执行 epilogue，然后再反过来。这最大化了每个 SM 上张量核心的“持续使用”，是实现最大吞吐量的关键原因之一。张量核心可以持续获取数据以实现其（接近）最大计算能力。（参见上面的图 2 插图底部部分）。

与生产者线程专注于数据移动类似，MMA 线程只发出 MMA 指令以实现峰值发射率。MMA 线程必须发出多个 MMA 指令并保持它们在 TMA 等待障碍之外执行。

以下是内核代码的摘录，以强调特化方面：

// Two types of warp group 'roles' 
enum class WarpGroupRole {
      Producer = 0,
      Consumer0 = 1,
      Consumer1 = 2
    };

//warp group role assignment
auto warp_group_role = WarpGroupRole(canonical_warp_group_idx());

使用生产者和张量内存加速器进行数据移动

生产者 warps 完全专注于数据移动——特别地，它们被保持得尽可能轻量，甚至将一部分寄存器空间让给消费者 warps（只保留 40 个寄存器，而消费者将获得 232 个）。它们的主要任务是发出 TMA (Tensor Memory Accelerator) 命令，在共享内存缓冲区被标记为空时立即将数据从全局内存移动到共享内存。

进一步解释 TMA，或 Tensor Memory Accelerator，TMA 是 H100 中引入的一种硬件组件，它异步处理内存从 HBM（全局内存）到共享内存的传输。通过拥有一个专用的硬件单元进行内存移动，工作线程得以解放，可以从事其他工作，而不是计算和管理数据移动。TMA 不仅处理数据本身的移动，还计算所需的目标内存地址，可以对数据应用任何转换（归约等），并且可以处理布局转换，以“交织”模式将数据传输到共享内存，使其无需任何 bank conflict 即可使用。最后，如果需要，它还可以将相同的数据多播到同一线程集群的其他 SM。数据传输完成后，TMA 将通知相关的消费者数据已准备好。

CUTLASS 异步流水线类

生产者和消费者之间的这种信令通过新的异步流水线类进行协调，CUTLASS 对其描述如下：

“实现一个持久化 GEMM 算法需要管理数十种不同类型的异步执行操作，这些操作使用多个组织成循环列表的障碍进行同步。

这种复杂性对于人工编程来说太难管理了。

因此，我们开发了 [CUTLASS Pipeline Async Class]…”

Ping-Pong 异步流水线中的障碍和同步

生产者必须通过“producer_acquire”来“获取”给定的 smem 缓冲区。一开始，流水线是空的，这意味着生产者线程可以立即获取障碍并开始移动数据。

PipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state<MainloopPipeline>();

数据移动完成后，生产者发出“producer_commit”方法，以通知消费者线程数据已准备好。
然而，对于 Ping-Pong，这实际上是一个无操作指令，因为基于 TMA 的生产者障碍会在写入完成后由 TMA 自动更新。

consumer_wait - 等待生产者线程的数据（阻塞）。

consumer_release - 通知等待的生产者线程，它们已完成从给定 smem 缓冲区消费数据。换句话说，允许生产者开始工作，用新数据重新填充该缓冲区。

从那里开始，同步将真正启动，生产者将通过阻塞的 producer acquire 进行等待，直到它们能够获取锁，此时它们的数据移动工作将重复进行。这会一直持续到工作完成。

提供伪代码概览：

//producer
While (work_tile_info.is_valid_tile) {

	collective_mainloop.dma() // fetch data with TMA
	scheduler.advance_to_next_work()
	Work_tile_info = scheduler.get_current_work()

}

// Consumer 1, Consumer 2
While (work_tile_info.is_valid_tile()) {

	collective_mainloop.mma()
	scheduler.advance_to_next_work()
	Work_tile_info = scheduler.get_current_work()

}

以及将所有部分与底层硬件结合在一起的视觉鸟瞰图：

Figure 3: An overview of the full async pipeline for Ping-Pong

图 3：Ping-Pong 的完整异步流水线概览

Ping-Pong 计算循环的逐步分解

最后，Ping-Pong 处理循环更详细的逻辑分解：

A - 生产者 (DMA) warp group 获取共享内存缓冲区上的锁。

B - 这使得它能够（通过单个线程）向 tma 芯片发起 tma cp_async.bulk 请求。

C - TMA 计算实际所需的共享内存地址，并将数据移动到共享内存。作为此过程的一部分，进行交织以在 smem 中布局数据，以实现最快（无 bank conflict）的访问。

C1 - 可能，数据也可以多播到其他 SM，或者可能需要等待其他 tma 多播的数据完成加载。（线程块集群现在跨多个 SM 共享共享内存！）

D - 在此点，障碍被更新以发出数据到达 smem 的信号。

E - 相关的消费者 warpgroup 现在开始工作，发出多个 wgmma.mma_async 命令，然后作为其 wgmma.mma_async 矩阵乘法操作的一部分，将数据从 smem 读取到张量核心。

F - MMA 累加器值在 Tile 完成后写入寄存器内存。

G - 消费者 warp group 释放共享内存上的障碍。

H - 生产者 warp group 开始工作，发出下一个 tma 指令以重新填充现在空闲的 smem 缓冲区。

I - 消费者 warp group 同时将任何 epilogue 操作应用于累加器，然后将数据从寄存器移动到另一个 smem 缓冲区。

J - 消费者 warp 发出 cp_async 命令，将数据从 smem 移动到全局内存。

这个循环重复，直到工作完成。希望这能为您提供对驱动 Ping-Pong 令人印象深刻的性能的核心概念的工作理解。

微基准测试

为了展示 Ping-Pong 的部分性能，以下是与我们设计快速推理内核相关的一些比较图表。

首先是目前最快的三个内核的通用基准测试（越低越好）：\

Figure 4, above: Benchmark timings of FP8 GEMMs, lower is better (faster)

图 4，上图：FP8 GEMM 基准测试时间，越低越好（越快）

并将其转化为 Ping-Pong 相对于 cuBLAS 和 Triton 的相对加速图表：

Figure 5, above: Relative speedup of Ping-Pong vs the two closest kernels.

图 5，上图：Ping-Pong 相对于两个最接近的内核的相对加速。

Ping-Pong 内核的完整源代码在这里（619 行深度模板化的 CUTLASS 代码，或者套用著名的乌龟梗——“全是模板……一直向下！”）：

https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp

此外，我们已将 PingPong 实现为 CPP 扩展，以便轻松集成到 PyTorch 中使用（以及展示其用法的一个简单测试脚本）：

https://github.com/pytorch-labs/applied-ai/tree/main/kernels/cuda/cutlass_gemm

最后，为了进一步学习，Nvidia 有两个 GTC 视频深入探讨了使用 CUTLASS 设计内核：

未来工作

数据移动通常是任何内核实现最高性能的最大障碍，因此对 Hopper 上的 TMA (Tensor Memory Accelerator) 有一个最优策略理解至关重要。我们之前发表了关于在 Triton 中使用 TMA 的工作。一旦 Triton 中启用了 Warp 特化等功能，我们计划再次深入探讨 Triton 内核（如 FP8 GEMM 和 FlashAttention）如何利用 Ping-Pong 等内核设计在 Hopper GPU 上实现加速。

摘要