上次为大家带来了Hopper上的新特性之TMA，这次我们来一起看看Hopper上的新矩阵乘法操作WGMMA。
引子
如果一个 CUDA 教程没有关于通用矩阵乘法（GEMM）的章节，那么就是不完整的。可以说，GEMM 是现代 GPU 上最重要的例程，它在神经网络、大型语言模型和许多图形应用程序中构成了大部分计算。尽管 GEMM 无处不在，但它以难以有效实现而闻名。
这个由三部分组成的教程系列旨在让读者全面了解如何使用 CUTLASS 库在 NVIDIA Hopper GPU 上编写高效的 GEMM 内核。
[第 1 部分，即本部分] 讨论了 warp 组矩阵乘法累加（WGMMA）指令。这些是针对基于 Hopper 架构的 NVIDIA GPU 的 Tensor Core 的原始指令。 [第 2 部分] 将讨论高效 GEMM 内核的整体设计，包括 CUTLASS 内核中使用的高级技术，如 warp 特化和乒乓调度。 [第 3 部分] 将讨论持久内核和 Stream-K，这是一种针对 GEMM 的负载平衡策略，可在大量问题几何形状中实现最先进的效率。
本系列的三个部分大致遵循通用矩阵乘（GEMM）内核的整个开发过程，但采用“由内向外”的方式。首先，我们有按分块进行的 GEMM 基本操作，它调用张量核心（Tensor Cores）来最终进行计算。其次，我们有从每个线程束协同线程组（CTA）角度看到的 GEMM 内核设计——由序言、主循环和尾声组成——其中主要挑战是避免内存加载成为快速张量核心的瓶颈。最后，我们在最外层网格级别对 CTA 进行调度，此时负载平衡考虑因素成为首要问题。
我们希望在阅读完本系列后，读者将成为 GEMM 算法的专家，并能够利用该算法中的一些优秀理念来设计和实现他们自己工作中的其他内核。
Asynchronous Warpgroup MMA (WGMMA)
Hopper 引入了异步线程束组级矩阵乘法累加运算（WGMMA）。一个线程束组由四个连续的线程束组成，即 128 个连续的线程，其中第一个线程束的线程束编号是 4 的倍数。wgmma.mma_async指令由线程束组中的所有 128 个线程共同执行。此操作通常采用以下形式之一，其中矩阵C用作累加器：
C = A * B + CC = A * B, where the input from accumulatorCis disabled.
WGMMA 的一个显著要求是操作数B必须始终存储在共享内存（SMEM）中。相比之下，操作数A可以位于共享内存或寄存器内存（RMEM）中，并且累加器C始终保存在 RMEM 中。
这篇博客文章的结构如下。首先，我们讨论在 CUTLASS 中调用wgmma.mma_async指令的要点。这涉及构建相关的TiledMMA，以及创建和划分 SMEM 张量以与 WGMMA 兼容。其次，我们讨论确保 WGMMA 正确性所需的同步机制。最后，我们更详细地讨论 WGMMA 中使用的布局，包括来自 SMEM 的操作数的核心矩阵和矩阵描述符的概念。
在整个过程中，为了简洁起见，我们将wgmma.mma_async缩写为wgmma。
CUTLASS kernel 中的WGMMA
在本教程中，我们的主要目标是解释用于调用 Hopper Tensor Cores 进行基于分块的 GEMM 的wgmma原语，以及如何将其作为cute::gemm调用的一部分进行调用。为了做好准备，考虑一个标准的 GEMM 内核，它接收维度为MxNxK的输入矩阵A和B，并计算C = A*B。为了并行化计算，内核固定静态分块大小bM、bN和bK，并启动一个由⌈M/bM⌉x⌈N/bN⌉多个线程块（CTAs），每个 CTA 计算输出矩阵的一个bMxbN瓦片rC。这将在被写回到全局C矩阵之前保存在 CTA 的本地内存（RMEM）中。
根据CTA，我们就有了内核的主循环。通过多次迭代，我们循环内部维度，并将A和B的bMxbK和bNxbK块依次从全局加载到共享内存中，作为sA和sB；请注意，在CUTLASS中，我们将sB的形状固定为数学上的转置。（事实上，反映了常见的做法，我们将A和B的块加载到循环ME M缓冲区中，其中的级数由编译时整数给出，例如2或3。然后sA和sB的形状元组的最后一种模式由该阶段计数给出。）cute::gemm调用然后计算sA和sB的（分阶段切片）的乘积，并将值连续累加到rC中。主循环完成后，最后将rC写入全局内存。
现在，我们希望解释以下cute::gemm调用及其参数。
template <class TiledMMA, ... >
__global__ device_gemm(TiledMMA tiled_mma, ...) {
// PROLOGUE
// ...
// Define A/B partitioning and C accumulators
ThrMMA thr_mma = tiled_mma.get_thread_slice(threadIdx.x);
Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
Tensor tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N)
// Allocate accumulators and clear them
Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N)
clear(tCrC);
// Allocate"fragments"
Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)
Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE)
// PIPELINED MAIN LOOP
while(k_tile_count > -K_PIPE_MAX) {
// ...
// MMAs to cover 1 K_TILE
cute::warpgroup_arrive();
// (V,M,K) x (V,N,K) => (V,M,N)
cute::gemm(tiled_mma, tCrA(_,_,_,read_pipe), tCrB(_,_,_,read_pipe), tCrC);
cute::warpgroup_commit_batch();
// Waitforall MMAsina K_TILE to complete
cute::warpgroup_wait<0>();
// ...
}
// EPILOGUE
// ...
}
在 CUTLASS 的 MMA（矩阵乘法累加）范式中，“MMA范式”里的cute::gemm方法旨在通过统一的接口展示特定架构的 MMA 指令。（实际上，如果你查看SM80 教程的 GEMM 内核，你会看到那里的cute::gemm调用在语法上与上述相同。）然而，cute::gemm调用所涉及的参数定义包含许多 WGMMA 特定的方面：
TiledMMA对象tiled_mma的定义封装了cute::gemm分派到特定wgmmaPTX指令所需的信息。 必须定义SMEM张量sA和sB的布局以与wgmma兼容。 片段 tCrA、tCrB和tCrC使用TiledMMA对象构建为数据的线程级分区，因此具有程序员应该注意的WGMMA特定布局。片段 tCrA（如果从SMEM获取操作数A）和tCrB不是寄存器支持的张量，其值是从SMEM复制的，而是在SMEM之上构建的矩阵描述符。
最后，当然，在cute::gemm调用周围有线程组同步原语。我们将依次解释所有这些概念。
WGMMA中的TiledMMA对象
以下内容中，假设数据类型为 FP16，且A和B是MN，所以在 BLAS 表示法中，我们正在计算一个NT gemm。我们使用cute::make_tiled_mma方法在主机上构造TiledMMA对象，如下所示：
TiledMMA tiled_mma = cute::make_tiled_mma(
SM90_64x64x16_F16F16F16_SS<GMMA::Major::MN,GMMA::Major::MN>{});
虽然cute::make_tiled_mma也有一些可选参数，但让我们专注于当前的这个参数——矩阵乘法累加原子（MMA Atom）。这是一个结构体，它封装了一个底层的 PTX 调用，在这种情况下是：
wgmma.mma_async.sync.aligned.m64n64k16.f16.f16.f16
CUTLASS 符号系统使得人们可以立即读出包装后的 PTX 指令与 MMA 原子之间的关系。首先，SM90 是 Hopper 架构的另一个名称。然后，SM90 MMA 原子被标记为SM90_MxNxK_XYZ_SS或SM90_MxNxK_XYZ_RS，其中有两个模板参数可以是GMMA::Major::MN或GMMA::Major::K。它们的含义如下：
X和Y是操作数的数据类型。 Z是累加器的数据类型。 MxNxK是wgmma指令计算的大小-“wgmma atom”。并非MxNxK的所有值都是可进行mma操作的的。这是允许形状的列表：M始终为64，N是从8到256的8倍数，对于16位操作数数据类型，K为16（更一般地说，K固定为32字节）。 后缀RS或SS指示操作数A是来自寄存器（R）还是来自共享内存（S）。操作数B始终来自共享内存，因此S。 这两个模板参数指示操作数A和B在MN模式或K模式下是内存连续的。例如，在BLAS表示法中，两个操作数都是K-Major将对应于一个TN gemm（参见此表）。请注意，对于16位操作数数据类型，可以灵活地将内存布局设置为MN-Major或K-Major。但是，对于非16位操作数数据类型，布局必须始终为K-Major。
这就是你需要了解的 MMA Atom 的语法！现在，我们已经强调过 WGMMA 是一个全线程组指令。在代码中，你可以使用其大小来检索参与由 TiledMMA 对象定义的 MMA 操作的线程数量。例如，以下主机代码。
dim3 dimBlock(cute::size(tiled_mma));
规定内核中的每个 CTA 以 1 个包含 128 个线程的线程束组启动。假设我们想要2 个线程束组来执行 WGMMA，由不同的线程束组独立计算输出块的一半（并且每个线程束组发出各自的wgmma指令）。为此，我们可以将一个非平凡的布局（AtomLayoutMNK）作为第二个参数传递给make_tiled_mma方法。例如，以下代码。
TiledMMA tiled_mma = make_tiled_mma(
SM90_64x64x16_F16F16F16_SS{},
Layout<Shape<_2,_1,_1>>{});
定义了一个 WGMMA 操作，其中 warp 组 1 和 2 分别计算输出瓦片的上半部分和下半部分，沿M模式划分（现在假设bM是 128 的倍数）。此外，size(tiled_mma)将等于 256。
一般来说，make_tiled_mma的两个可选布局参数——AtomLayoutMNK和PermutationMNK——对于任何 MMA 原子都同样适用。
共享内存的布局约束了WGMMA
接下来，我们解释在给定 MMA 原子选择的情况下，共享内存中操作数矩阵的瓦片大小和布局的约束。首先，对于任何 MMA 指令，MMA 原子的MxNxK需要能够整除操作数和累加器Tile的大小。在我们的例子中，这意味着bM应该是 64 的倍数，bN是 64 的倍数，bK是 16 的倍数。
其次，WGMMA 对sA和sB的共享内存布局（包括形状和跨度）施加了一个额外的约束，并且这个约束会随着所选的交错模式而变化。特别是，（分阶段切片的）sA的布局通常不是简单的(bM,bK)
1,bM)或(bM,bK)bK,1)，sB也是如此。
为了深入理解这些要求，我们需要“核心矩阵”的概念，我们将在下面介绍。然而，实际上，我们总是可以使用 CUTLASS 提供的某些预定义布局原子，然后使用cute::tile_to_shape方法构建保证与
auto bM = Int<128>{};
auto bN = Int<128>{};
auto bK = Int< 64>{};
auto bP = Int< 3>{}; // Pipeline
auto sA = cute::tile_to_shape(
GMMA:
ayout_MN_SW128_Atom<T>{},
cute::make_shape(bM, bK, bP)
);
auto sB = cute::tile_to_shape(
GMMA:ayout_MN_SW128_Atom<T>{},
cute::make_shape(bN, bK, bP)
);
在这里，MN表示布局原子适用于MN主操作数，而SW128是 128 字节的交错模式。输出sA或sB会显示。
Sw<3,4,3> o smem_ptr[16b](unset) o ((_64,_2),(_8,_8),_3)(_1,_512),(_64,_1024),_8192)
这个布局是从哪里来的？cute::tile_to_shape采用一个布局（同名的tile）并复制它以平铺在更大的形状上（类似于numpy.tile）。抛开swizzle函数Sw<3,4,3>，我们知道布局原子由(64,8)1,64)给出，并以列主要方式平铺在形状(128, 64, 3)上，因此对于MxK形状，512的较小外步幅位于M模式，而1024的较大外步幅位于K模式。（8192的最大步幅在于阶段计数P模式，这是有道理的，因为sA或sB的不同阶段切片不应该在内存中混合。）
请注意，64乘以sizeof(half_t)等于128字节，这是swizzle模式的名称。这是设计：由于核心矩阵的工作方式，我们总是在连续方向上安排布局原子的长度以等于swizzle字节数-对于无swizzle，可以是16，或者32、64或128之一。
相对的，如果我们考虑：
auto sA = cute::tile_to_shape(
GMMA:ayout_K_SW128_Atom<T>{},
cute::make_shape(bM,bK,bP)
);
auto sB = cute::tile_to_shape(
GMMA:ayout_K_SW128_Atom<T>{},
cute::make_shape(bN,bK,bP)
);
打印sA会得到我们预期的结果。
Sw<3,4,3> o smem_ptr[16b](unset) o (_128,_64,_3)_64,_1,_8192)
由于我们改为在(8,64)64,1)上平铺(8,64)64,1)。（请注意，布局((_8,_16),(_64,_1),_3)(_64,_512),(_1,_0),_8192)合并为(_128,_64,_3)_64,_1,_8192)）。
一般来说，我们可以在8种布局原子的可能性中进行选择，它们对应于MN或K为主以及四种混洗模式之一：
无交错：无交错。隐含 16 字节边界。 32 字节交错：交错 2 个连续的 16 字节段。 64 字节交错：交错 4 个连续的 16 字节段。 128 字节交错：交错 8 个连续的 16 字节段。
GMMA:ayout_MN_INTER_Atom<T>
GMMA:ayout_MN_SW32_Atom<T>
GMMA:ayout_MN_SW64_Atom<T>
GMMA:ayout_MN_SW128_Atom<T>
GMMA:ayout_K_INTER_Atom<T>
GMMA:ayout_K_SW32_Atom<T>
GMMA::Layout_K_SW64_Atom<T>
GMMA::Layout_K_SW128_Atom<T>
然后，必须将这些布局原子传入tile_to_shape，其中sA和sB的共享内存（SMEM）形状由make_shape(bM,bK,bP)或make_shape(bN,bK,bP)给出，形状的模式按照该顺序给出，使得布局原子的分块大小能够整除较大的 SMEM 形状的分块大小。这最终是由混洗模式的选择对 SMEM 形状造成的约束，并且与由矩阵乘法累加（MMA）原子形状施加的另一个约束分开。
WGMMA 片段和描述符
我们创建了TiledMMA对象，并在主机上相应地准备了共享内存（SMEM）布局。现在，在设备上，我们可以使用TiledMMA对象tiled_mma来构建适当的分区张量，以便传递到cute::gemm调用中。首先，我们通过在tiled_mma上调用带有线程索引的get_thread_slice方法来创建一个名为thr_mma的ThrMMA对象。在我们的例子中，线程索引从0到127。
接着，参考上面的内核代码片段，打印张量tCsA和tCsB对于任何线程索引，显示如下：
tCsA: Sw<3,4,3>_smem_ptr[16b](0x7f8800000400) o
((_64,(_8,_2)),_2,_4,_3)(_1,(_64,_1024)),_512,_2048,_8192)
tCsB: Sw<3,4,3>_smem_ptr[16b](0x7f880000c400) o
((_64,(_8,_2)),_2,_4,_3):((_1,(_64,_1024)),_512,_2048,_8192)
根据注释，tCsA的形状应被视为(MMA,MMA_M,MMA_K,PIPE)：
MMA是MMA Atom的NxK形状。 MMA_M和MMA_K是它在sA的M和K模式上平铺的范围（因此MMA_M=bM/64=2和MMA_K=bK/16=4）。 PIPE是stage数。
tCrA: GMMA:
escriptorIterator o (_1,_2,_4,_3):(_0,_64,_256,_1024)
tCrB: GMMA:escriptorIterator o (_1,_2,_4,_3):(_0,_64,_256,_1024)
在内部，CUTLASS 构造一个“矩阵描述符”，这是一个保存在寄存器中的 64 位值，以一种适合wgmma指令使用的方式描述共享内存（SMEM）。对于程序员来说，最重要的是要记住，共享内存的值不会被复制到寄存器内存（RMEM）中；相反，访问 tCrA 和 tCrB 的值实际上是访问这些 64 位描述符。此外，这些张量作为“迭代器”意味着在任何时候，对于给定的wgmma指令，只有一个 64 位描述符保存在寄存器中（例如，与全部 24 个不同）。
与操作数相比，累加器张量以更标准的方式定义。打印线程 0 的tCgC和tCrC显示：
tCgC: gmem_ptr[16b](0x7f877a780000) o ((_2,_2,_8),_2,_2):((512,_8,4096),_64,32768)
tCrC: ptr[16b](0x7feee1fffbe0) o ((_2,_2,_8),_2,_2):((_1,_2,_4),_32,_64)
tCgC是输出 GMEM 张量的一部分，我们希望在尾声中将累加器的值复制到该部分，而tCrC是为了在主循环中计算这些值时保存这些值而创建的基于寄存器的张量。这些张量的(MMA,MMA_M,MMA_N)形状可以如下解释：在 MMA 原子的MxN=64x64输出块中，128 个线程中的每个线程都持有32=2*2*8个值，并且MMA_M=MMA_N=2与tCsA和tCsB相同。
每个线程以一种需要将 32 分解为(2,2,8)形状的方式持有原子的 32 个值，以便能够为tCgC的布局定义相应的步长。具体的分区模式可以从取自 PTX 文档的这张图片中读出：