深入剖析 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 架构可用的最快的 matmul (GEMM) 内核架构之一。Ping-Pong 是 Warp Group Specialized Persistent Kernels 系列的成员，该系列包括 Cooperative 和 Ping-Pong 变体。相对于之前的 GPU，Hopper 强大的张量核心计算能力需要深入的异步软件流水线才能实现峰值性能。

Ping-Pong 和 Cooperative 内核例证了这种范例，因为关键设计模式是持久内核，用于分摊启动和 prologue 开销，以及“异步一切”，使用具有两个消费者和一个生产者的专用 warp 组，以创建高度重叠的处理流水线，该流水线能够持续向张量核心提供数据。

当 H100 (Hopper) GPU 发布时，Nvidia 称其为首款真正的异步 GPU。该声明强调了 H100 特定内核架构也需要是异步的，以便充分最大化计算/GEMM 吞吐量。

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

Ping-Pong 内核设计

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

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

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

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

生产者 warp 组与 TMA（张量内存加速器）一起工作，并有意保持尽可能轻量化。事实上，在 Ping-Pong 中，他们有意减少其寄存器资源以提高占用率。生产者将最大寄存器计数减少 40，而消费者将最大寄存器计数增加 232，这是我们可以在 CUTLASS 源代码和相应的 SASS 中看到的效果

source code

Ping-Pong 独有的特点是，每个消费者都在单独的 C 输出 tile 上工作。（作为参考，协同内核在很大程度上等同于 Ping-Pong，但两个消费者组都在同一个 C 输出 tile 上工作）。此外，两个消费者 warp 组然后将它们的工作在主循环 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());

生产者和张量内存加速器的数据移动

生产者 warp 专注于数据移动 - 特别是它们被保持尽可能轻量化，并且实际上将它们的一些寄存器空间让给消费者 warp（仅保留 40 个寄存器，而消费者将获得 232 个）。它们的主要任务是发出 TMA（张量内存加速器）命令，以便在共享内存缓冲区被标记为空时立即将数据从全局内存移动到共享内存。

为了扩展 TMA 或张量内存加速器，TMA 是 H100 引入的硬件组件，它可以异步处理从 HBM（全局内存）到共享内存的内存传输。通过拥有专用的硬件单元进行内存移动，工作线程可以自由地从事其他工作，而不是计算和管理数据移动。TMA 不仅处理数据本身的移动，还计算所需的目的地内存地址，可以对数据应用任何转换（缩减等），并且可以处理布局转换，以便以“swizzled”模式将数据传送到共享内存中，从而使其可以立即使用，而不会发生任何 bank 冲突。最后，如果需要，它还可以将相同的数据多播到同一线程集群的成员的其他 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，这实际上是一个 noop 指令，因为基于 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：Ping-Pong 的完整异步流水线概述

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

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

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

B - 这使其可以启动对 tma 芯片的 tma cp_async.bulk 请求（通过单个线程）。

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

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

D - 此时，屏障已更新，以向 smem 发出数据到达的信号。

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

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

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

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

I - 消费者 warp 组同时对累加器应用任何 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（张量内存加速器）有最佳策略理解至关重要。我们之前发表了关于 Triton 中 TMA 用法的工作。一旦在 Triton 中启用 warp 专精化等功能，我们计划再次深入研究 Triton 内核（如 FP8 GEMM 和 FlashAttention）如何利用 Ping-Pong 等内核设计来加速 Hopper GPU。

摘要