深入解析 CUTLASS Ping-Pong GEMM 内核 – PyTorch

Figure 1. FP8 GEMM Throughput Comparison CUTLASS vs Triton

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

总结

在本篇文章中，我们将概述 CUTLASS Ping-Pong GEMM 内核，并提供相关的 FP8 推理内核基准测试数据。

Ping-Pong 是目前 Hopper GPU 架构下速度最快的矩阵乘法 (GEMM) 内核架构之一。Ping-Pong 属于“Warp Group Specialized Persistent Kernels”（线程束组专用持久内核）家族，该家族包含 Cooperative（协作）和 Ping-Pong 两种变体。与以往的 GPU 相比，Hopper 强大的张量核心（Tensor Core）计算能力要求极高的异步软件流水线技术才能达到峰值性能。

Ping-Pong 和 Cooperative 内核正是这一范式的典范。其核心设计模式是利用持久内核来摊销启动和序言开销，并通过“一切皆异步”的理念，配合具备“两消费者、一生产者”结构的专用线程束组，构建出一个高度重叠的处理流水线，从而能够持续地为张量核心提供数据。

当 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 Specialization）。与传统的同构内核不同，这里由“线程束组”（Warp Groups）承担专门的角色。请注意，一个线程束组由 4 个线程束组成，每个线程束包含 32 个线程，总计 128 个线程。

在旧架构上，延迟通常通过在每个 SM 上运行多个线程块来掩盖。然而，Hopper 的张量核心吞吐量极高，必须采用更深层的流水线，而这又会阻碍在每个 SM 上运行多个线程块。因此，持久化线程块现在需跨多个瓦片（tile）和多个线程束组执行集体主循环。线程块集群（Thread block clusters）根据总 SM 数量进行分配。

对于 Ping-Pong 而言，每个线程束组都承担特定的角色：数据生产者（Producer）或数据消费者（Consumer）。

生产者线程束组专注于数据移动，负责填充共享内存缓冲区（通过 TMA）。另外两个线程束组作为专用消费者，负责使用张量核心处理数学运算（MMA），完成后续工作，并将结果写回全局内存（尾声/epilogue）。

生产者线程束组与 TMA（张量内存加速器）协同工作，并被刻意设计得尽可能轻量化。事实上，在 Ping-Pong 中，它们有意减少寄存器资源以提高占用率。生产者的最大寄存器数量比消费者减少了 40 个，而消费者的最大寄存器数量增加了 232 个，这一效果可以在 CUTLASS 源代码和相应的 SASS 代码中观察到。

Ping-Pong 的独特之处在于，每个消费者处理独立的 C 输出瓦片。（作为参考，Cooperative 内核与 Ping-Pong 在很大程度上等价，但两个消费者组在同一 C 输出瓦片上工作）。此外，两个消费者线程束组将工作在主循环 MMA 和尾声（epilogue）之间进行分配。

如下图所示

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

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

拥有两个消费者意味着一个消费者可以使用张量核心进行 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());

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

生产者线程束专注于数据移动——具体来说，它们保持尽可能轻量化，甚至将部分寄存器空间让给消费者线程束（仅保留 40 个寄存器，而消费者获得 232 个）。它们的主要任务是在共享内存缓冲区被标记为空闲时，发出 TMA（张量内存加速器）指令，将数据从全局内存移动到共享内存。

关于 TMA（张量内存加速器），它是 H100 中引入的硬件组件，用于异步处理从 HBM（全局内存）到共享内存的数据传输。通过拥有专门的硬件单元进行数据移动，工作线程从计算和管理数据移动的任务中解放出来。TMA 不仅处理数据本身的移动，还能计算所需的共享内存地址，对数据应用任何转换（如缩减等），并能处理布局转换，以“交错（swizzled）”模式将数据传送到共享内存，从而确保使用时不会产生存储体冲突（bank conflicts）。最后，如果需要，它还可以将相同数据组播（multicast）到同一线程集群中的其他 SM。数据交付完成后，TMA 会向相关消费者发出信号，通知数据已就绪。

CUTLASS 异步流水线类

生产者和消费者之间的信号传输是通过新的异步流水线类（Asynchronous Pipeline Class）进行协调的，CUTLASS 对此描述如下：

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

这种复杂性对于程序员手动管理来说过于繁琐。

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

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_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) 线程束组获取共享内存缓冲区的锁。

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

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

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

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

E – 相关的消费者线程束组开始工作，发出多个 wgmma.mma_async 命令，这些命令在 wgmma.mma_async 矩阵乘法操作中将数据从 smem 读取到张量核心。

F – 当瓦片完成时，MMA 累加器值被写入寄存器内存。

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

H – 生产者线程束组开始工作，发出下一条 tma 指令以重新填充现已空闲的 smem 缓冲区。

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

J – 消费者线程束发出 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 中启用了线程束专用化等功能，我们计划进行另一次深入探讨，分析诸如 FP8 GEMM 和 FlashAttention 之类的 Triton 内核如何利用像 Ping-Pong 这样的内核设计在 Hopper GPU 上实现加速。

深入探讨 CUTLASS 乒乓球 GEMM 核

总结