设备软件

线程束组Warpgroup

Warpgroup（线程束组）是指四个连续线程束 (warps) 的集合，且该组中第一个线程束的秩（warp-rank）必须是 4 的倍数。

在分发 Warpgroup 级别的指令时，系统会协同 128 个线程 —— 即每个 Warpgroup 包含 4 个线程束，每个线程束包含 32 个线程。以更大的粒度进行操作，消除了显式线程束间同步的需求，并允许单条指令处理更大规模的问题，特别是更大的矩阵乘法。更大规模的矩阵乘法能够更容易地使近期数据中心 GPU 中 Tensor Core (张量核心) 的海量算术带宽达到饱和。

Warpgroup 引入于 NVIDIA 的 Hopper 流式多处理器（SM）架构中，用于支持 Warpgroup 级别的矩阵乘法，例如 wgmma.mma_async 指令。欲深入了解，请参阅 Colfax 的这篇博客文章。Warpgroup 在高性能 Hopper 和 Blackwell 内核（kernels）（如 Flash Attention 4）的流水线组件组织中占据重要地位。

在并行线程执行 (PTX) 中间表示（IR）中，一个线程束的秩（warp-rank）定义为：

int linearIdx = (%tid.x + %tid.y * %ntid.x  + %tid.z * %ntid.x * %ntid.y);
int warpRank = linearIdx / 32;

其中 tid 是线程索引，通过特殊的 PTX 寄存器进行访问。

因此，对于包含 8 个线程束的调度，有效的 Warpgroup 分别为：

Warpgroup 0: 秩为 0, 1, 2 和 3 的线程束
Warpgroup 1: 秩为 4, 5, 6 和 7 的线程束

据我们需要，关于这种线程束秩（warp-rank）对齐限制的目的尚无官方文档说明。但近期数据中心 GPU 的流式多处理器（SM）似乎包含四个（未命名的）子单元，每个子单元都拥有独立的线程束调度器和 Tensor Core。

英文原文(Modal GPU Glossary)

A warpgroup is a set of four contiguous warps such that the warp-rank of the first warp is a multiple of 4.

Upon dispatching a warpgroup-level instruction, we coordinate 128 threads -- 4 warps per warpgroup × 32 threads per warp. Operating at a larger granularity removes the need for explicit inter-warp synchronization and allows work to be performed on larger problem sizes per instruction, especially larger matrix multiplications. Larger matrix multiplications more readily saturate the massive arithmetic bandwidth of the Tensor Cores of recent data center GPUs.

Warpgroups were introduced in NVIDIA's Hopper Streaming Multiprocessor architecture, where they are used to support warpgroup-level matrix multiplication, like wgmma.mma_async. See this blog post from Colfax for a deep dive. Warpgroups feature prominently in the organization of pipeline components of high-performance Hopper and Blackwell kernels, like Flash Attention 4.

In Parallel Thread Execution (PTX) IR, the warp-rank of a warp is:

int linearIdx = (%tid.x + %tid.y * %ntid.x  + %tid.z * %ntid.x * %ntid.y);
int warpRank = linearIdx / 32;

where tid is the thread index, accessed via special PTX registers.

So the valid warpgroups for an 8-warp dispatch are:

Warpgroup 0: warp-ranks 0, 1, 2, and 3
Warpgroup 1: warp-ranks 4, 5, 6, and 7.

To our knowledge, the purpose of the warp-rank alignment restriction is not documented. But Streaming Multiprocessors for recent data center GPUs appear to contain four (unnamed) subunits, each with their own Warp Scheduler and Tensor Core.

相关词条