深入探讨 CUTLASS 乒乓 GEMM 内核 – PyTorch

Figure 1. FP8 GEMM Throughput Comparison CUTLASS vs Triton

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

总结

在这篇文章中，我们将概述 CUTLASS 乒乓 GEMM 内核，并提供相关的 FP8 推理内核基准测试。

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

乒乓和协同内核是这种范例的典范，因为关键设计模式是持久内核以分摊启动和前言开销，以及“一切皆异步”，并使用具有两个消费者和一个生产者的专用 warp 组，以创建高度重叠的处理流水线，能够持续为张量核心提供数据。

当 H100 (Hopper) GPU 发布时，Nvidia 将其称为第一个真正的异步 GPU。这一声明强调了 H100 专用内核架构也需要异步才能充分最大化计算/GEMM 吞吐量。

CUTLASS 3.x 中引入的乒乓 GEMM 通过将内核的所有方面都移动到“完全异步”处理范例中，从而体现了这一点。在这个博客中，我们将展示乒乓内核设计的核心功能，并展示其在推理工作负载上与 cublas 和 triton split-k 内核相比的性能。

乒乓内核设计

乒乓（或技术上称为“sm90_gemm_tma_warpspecialized_pingpong”）以异步流水线运行，利用了 warp 专业化。与更经典的同构内核不同，“warp 组”承担了专门的角色。请注意，一个 warp 组由 4 个 warps 组成，每个 warp 32 个线程，总共 128 个线程。

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

对于乒乓，每个 warp 组都承担了数据生产者或数据消费者的专门角色。

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

生产者 warp 组与 TMA（张量内存加速器）协同工作，并被刻意保持尽可能轻量级。事实上，在乒乓中，它们刻意减少其寄存器资源以提高占用率。生产者将最大寄存器计数减少 40，而消费者将最大寄存器计数增加 232，这种效果我们可以在 CUTLASS 源代码和相应的 SASS 中看到。

乒乓的独特之处在于，每个消费者都在单独的 C 输出瓦片上工作。（作为参考，协同内核与乒乓大致相同，但两个消费者组都在同一个 C 输出瓦片上工作）。此外，两个消费者 warp 组随后在主循环 MMA 和尾声之间分配其工作。

这在下图中显示

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

图 2：乒乓内核流水线的概述。时间从左向右移动。

通过拥有两个消费者，意味着一个可以使用张量核心进行 MMA，而另一个执行尾声，反之亦然。这最大化了每个 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（张量内存加速器）命令，一旦共享内存缓冲区被标记为空，就立即将数据从全局内存移动到共享内存。

更详细地解释 TMA 或张量内存加速器，TMA 是 H100 中引入的硬件组件，它异步处理从 HBM（全局内存）到共享内存的内存传输。通过拥有专用的内存移动硬件单元，工作线程可以自由地从事其他工作，而不是计算和管理数据移动。TMA 不仅处理数据本身的移动，还计算所需的目的地内存地址，可以对数据应用任何转换（缩减等），并且可以处理布局转换，以“交错”模式将数据传送到共享内存，以便无需任何 bank 冲突即可使用。最后，如果需要，它还可以将相同的数据多播到属于同一线程集群的其他 SM。数据传输完成后，TMA 将向感兴趣的消费者发出数据已准备就绪的信号。

CUTLASS 异步流水线类

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

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

“这种复杂性对于人类程序员来说太难手动管理了。”

“因此，我们开发了 [CUTLASS 流水线异步类]…”

乒乓异步流水线中的屏障和同步

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

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

数据移动完成后，生产者发出“producer_commit”方法，向消费者线程发出数据已准备就绪的信号。
然而，对于乒乓来说，这实际上是一个空操作指令，因为基于 TMA 的生产者屏障在写入完成后会自动由 TMA 更新。

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

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

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

提供伪代码概述

//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：乒乓全异步流水线的概述

乒乓计算循环的逐步分解

最后，乒乓处理循环更详细的逻辑分解

A – 生产者 (DMA) warp 组获得共享内存缓冲区上的锁。

B – 这允许它向 TMA 芯片发起 tma cp_async.bulk 请求（通过单个线程）。

C – TMA 计算所需的实际共享内存寻址，并将数据移动到共享内存。作为其中一部分，执行了交错以在 smem 中布局数据以实现最快（无 bank 冲突）访问。

C1 – 潜在地，数据也可以多播到其他 SM，并且/或者它可能需要等待来自其他 TMA 多播的数据完成加载。（线程块集群现在跨多个 SM 共享共享内存！）

D – 此时，屏障被更新以指示数据到达 smem。

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

F – MMA 累加器值在瓦片完成后写入寄存器内存。

G – 消费者 warp 组释放共享内存上的屏障。

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

I – 消费者 warp 组同时对累加器应用任何尾声操作，然后将数据从寄存器移动到不同的 smem 缓冲区。

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

循环重复，直到工作完成。希望这能让您对乒乓令人印象深刻的性能背后的核心概念有一个实际的理解。

微基准测试

为了展示乒乓的一些性能，下面是我们关于设计快速推理内核工作的一些比较图表。

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

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

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

并将其转化为乒乓与 cuBLAS 和 Triton 的相对加速图

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

图 5，上方：乒乓与两个最接近的内核的相对加速。

乒乓内核的完整源代码在此（619 行深度模板化的 CUTLASS 代码，或者用著名的乌龟梗来解释——“全是模板……一路往下！”）

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

此外，我们已将乒乓实现为 CPP 扩展，以便于与 PyTorch 集成使用（以及一个简单的测试脚本，显示其用法）

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

最后，为了持续学习，Nvidia 有两个 GTC 视频深入探讨了 CUTLASS 的内核设计

未来工作

数据移动通常是任何内核实现最高性能的最大障碍，因此，对 Hopper 上的 TMA（张量内存加速器）有一个最佳的策略理解至关重要。我们之前发表过关于Triton 中 TMA 使用的工作。一旦在 Triton 中启用 warp 专门化等功能，我们计划再进行一次深入探讨，了解 Triton 内核（如 FP8 GEMM 和 FlashAttention）如何利用乒乓等内核设计在 Hopper GPU 上进行加速。

深入探讨 CUTLASS 乒乓球 GEMM 核

总结