在生成式 AI 时代，充分发挥 GPU 的潜力对于训练更好的模型和大规模服务用户至关重要。通常，这些模型的层由于细微的修改而无法表示为现成的库运算，而 DL 编译器通常会放弃最后几个百分点的优化，以实现其部署。

为了向 NVIDIA CUDA 开发者提供更大限度地提高 DL 和 HPC 内核性能所需的功率和控制水平，自 2017 年以来，我们一直在 CUTLASS 上构建和迭代。

借助新的 Python 接口，它现已进入下一阶段的开发阶段。重新设计 CUTLASS 3.x 时引入的基本抽象概念在 Python 和 CUTLASS 4.0 中直接公开。在本文中，我们讨论了 CUTLASS 3.x 背后的设计原则、其核心后端库、CUDA 张量和空间微核 (CuTe) ，以及利用 CuTe 关键功能的优化示例。

来自 CUTLASS 3.x 的亮点

CUTLASS 3 引入了 CuTe，这是一个基于布局概念的新库，作为描述和操作线程和数据的统一且可组合的抽象。通过将布局提升为编程模型的一级公民，CuTe 的使用大大简化了线程数据的组织。CuTe 以易于理解和静态可检查的方式向开发者展示索引逻辑，同时保持与 CUTLASS 2.x 中相同的高性能和 Tensor Core 运算覆盖率。

除了这种更有意义的布局方法之外，CUTLASS 3 的目标与所有之前版本的 CUTLASS 相同，即通过围绕最新硬件功能开发直观的编程模型，帮助 CUDA 开发者编写高性能 GPU 线性代数核函数。在这一新的主要迭代中，我们强调了以下几点：

• 能够在库设计中自定义任何层，同时保持与其他层的可组合性，从而提高开发者的工作效率并更清晰地分离运动部件
• 编译时检查，以确保内核结构的正确性。这可以保证，如果编译，它将正确运行，否则将使用可操作的静态 assert 消息。
• 通过更少的命名类型和更平滑的学习曲线，以及自定义 Hook 的单入口点来减少 API 表面积。
• NVIDIA Hopper H100 和 NVIDIA Blackwell B200 性能出色，可使用 WGMMA (适用于 Hopper) 或 UMMA (适用于 Blackwell) 、Tensor Memory Accelerator for Hopper (TMA) 和线程块集群等功能。

CuTe

CUTLASS 3.x 的核心是 CuTe，这是一个用于描述和操作线程和数据张量的新库。CuTe 由两部分组成：强大的布局表示和作用于这些布局的运算代数。

CuTe 的布局表示采用原生分层式，自然支持静态和动态信息，并用于表示多维张量。相同的布局表示用于描述数据张量和线程张量。在多个独立资源中使用相同的词汇类型显示了 CuTe 布局概念的广泛适用性。

基于这种表征能力，CuTe 提供了形式化的布局代数，使用户能够根据简单的已知布局构建复杂的布局，或将一个布局分割到另一个布局。这使得程序员可以专注于其算法的逻辑描述，而 CuTe 可以为其进行机械记帐。借助这些工具，用户可以快速设计、实施和修改密集线性代数算法。

与之前的任何 GPU 编程模型不同，线程和数据张量的功能组合消除了 GPU 编程中最复杂的障碍之一，即始终将大量线程映射到其所运行的数据。一旦描述了独立于将要操作的数据布局的线程布局，CuTe 的布局代数就可以跨线程对数据进行分区，而无需手动实施复杂的后分区迭代方案。

CuTe 布局和张量

有关布局和张量的更多 CuTe 文档可在其专用文档目录中找到。

CuTe 提供 Layout 和 Tensor objects，可紧凑地封装数据的类型、形状、内存空间和布局，同时为用户执行复杂的索引。

• Layout<Shape,Stride> 提供 Shape 内逻辑坐标与使用 Stride 计算出的索引之间的映射。(请参见图 1 示例) Shape 定义了一个或多个坐标空间以及它们之间的映射。 Stride 定义了将坐标转换为索引的索引图。
• Tensor<Engine,Layout> 通过迭代器提供 Layout 的合成。迭代器可以是指向全局内存、共享内存、寄存器内存或任何其他提供随机访问偏移和解引用的数据的指针。

值得强调的是，CuTe 中的布局是分层的，并受张量代数中折叠张量运算的启发。如图所示，分层形状和步长可实现远超简单行主和列主的布局表示。同时，分层布局仍然可以像正常张量一样访问 (例如，所示的逻辑 2D 坐标) ，因此这些更高级的数据布局在算法开发中被抽象化。

CUTLASS 3.x 中的 CuTe

CUTLASS 3.x 使用单一词汇表类型 (cute::Layout) ，从而实现简化、形式化和统一的布局表示，帮助用户轻松编写超快的内核。

用于转换和分区的 CuTe 布局

CuTe 布局支持将功能合成作为核心运算。功能合成可用于转换另一个布局的形状和顺序。如果我们有一个具有坐标 (m,n) 的数据布局，而我们想改用坐标 (tg_ 15) ，则我们会使用描述映射的布局来编写数据布局 (tg_ 16) – > (tg_ 17) 。

结果是具有坐标 (thread_idx,value_idx) 的数据布局，我们可以使用该布局轻松访问每个线程的每个值！

例如，考虑 4 × 8 的数据布局。此外，假设我们要为 4 × 8 数据的每个坐标分配线程和值。我们编写一个“电视布局”，记录特定的分区模式，然后在数据布局和电视布局之间执行功能合成。

如图所示，合成会对数据进行排列和重塑，以便在结果的每一行中排列每个线程的值。只需使用线程索引对结果进行切片，即可完成分区。

更直观的分区模式视图是电视布局的反向视图。

此布局显示了从 4 × 8 数据布局中的每个坐标到线程和值的映射。可以记录任意分区模式，并将其应用于任意数据布局。

有关 CuTe 布局代数的其他文档可在 GitHub 上找到。

CuTe 矩阵乘积累加原子

原子是必须协同参与执行硬件加速数学运算或复制运算的最小线程和数据集合。

Atom 将 PTX 指令与有关线程形状和安排的元数据以及必须参与该指令的值相结合。此元数据表示为 CuTe TV 布局，然后可用于划分输入和输出数据的任意张量。用户通常不必扩展此层，因为我们将为新架构提供 CuTe 原子的实现。

上图显示了 SM70_8x8x4_F32F16F16F32_NT 指令及其关联的 MMA_Traits 元数据。在左侧，映射 (thread_id,value_id) -> coord 的 TV 布局记录在特征中，在右侧，通过 tg_ 24 映射可视化特征。可以使用以下命令生成右侧图像

print_latex(make_tiled_mma(cute::SM70_8x8x4_F32F16F16F32_NT{}))

有关矩阵乘积累加 (MMA) 原子的其他 CuTe 文档位于 GitHub 上。

CuTe 平铺 MMA

平铺式 MMA 和平铺式文案分别是 MMA 原子和拷贝原子的平铺式。我们将此级别称为“平铺”，因为它在原子之上构建更大的运算，就像将单个图块组合在一起以构建可重复使用的马赛克组件一样。这些平铺在线程和数据之间重现原子，并且可能存在原子的排列和交错。

此层最类似于 CUTLASS 2.x 中 MMA 指令的线程束级平铺；但是，它会从参与操作的所有线程的角度查看平铺，并将此概念推广到复制操作。此层的目的是基于大量硬件加速的数学和数据移动操作构建可组合的 GPU 微内核，每个操作都可能具有自己的线程和数据内部布局。平铺的 MMA 和平铺的 Copy 类型通过一个统一的 API 来划分数据，从而呈现所有这些硬件加速的 CuTe 原子。

例如，CuTe 可能会提供一个 MMA 原子，用户可以针对固定的 M、N 和 K 维度在单个线程束中调用该原子。然后，我们可以使用 CuTe 运算 make_tiled_mma 将此原子转换为适用于整个线程块的运算，以处理更大的 M、N 和 K 维度。在上一节中，我们已经看到了 Tiled MMA 的一个示例，即 SM70_8x8x4_F32F16F16F32_NT 的 1x1x1 平铺。

下图显示了另外两个使用相同 SM70_8x8x4_F32F16F16F32_NT 原子的平铺 MMA。在左侧，其中四个原子组合成 2 × 2 行的主要布局，以生成 16x16x4 的单经 MMA。在右侧，其中四个原子是 2 × 2 行的主要布局，以产生 16x16x4 的单曲面 MMA，然后行 (M) 和列 (N) 被排列以交错这些原子。这两种方法都会产生可应用于任何数据布局的分区模式，如下节所示。

CuTe GEMM 和主回路

借助与架构无关的平铺 API，用户可以构建通往 GEMM 外部循环的一致接口，其中包含来自原子层的内部循环。

Tensor gA = . . . // Tile of 64x16 gmem for A
Tensor gB = . . . // Tile of 96x16 gmem for B
Tensor gC = . . . // Tile of 64x96 gmem for C

// 64x16 static-layout padded row-major smem for A
Tensor sA = make_tensor(make_smem_ptr<TA>(smemAptr),
                        Layout<Shape <    _64,_16>,
                               Stride<Int<17>, _1>>{});
// 96x16 static-layout interleaved col-major smem for B
Tensor sB = make_tensor(make_smem_ptr<TB>(smemBptr),
                        Layout<Shape <Shape <_32,  _3>,_16>,
                               Stride<Stride< _1,_512>,_32>>{});

// Partition tensors across threads according to the TiledMMA
ThrMMA thr_mma = tiled_mma.get_slice(thread_idx);
Tensor tCsA = thr_mma.partition_A(sA);        // (MMA, MMA_M, MMA_K) smem
Tensor tCsB = thr_mma.partition_B(sB);        // (MMA, MMA_N, MMA_K) smem
Tensor tCgC = thr_mma.partition_C(gC);        // (MMA, MMA_M, MMA_N) gmem

// Make register tensors the same shape/layout as above
Tensor tCrA = thr_mma.make_fragment_A(tCsA);  // (MMA, MMA_M, MMA_K) rmem
Tensor tCrB = thr_mma.make_fragment_B(tCsB);  // (MMA, MMA_N, MMA_K) rmem
Tensor tCrC = thr_mma.make_fragment_C(tCgC);  // (MMA, MMA_M, MMA_N) rmem

// COPY from smem to rmem thread-level partitions
cute::copy(tCsA, tCrA);
cute::copy(tCsB, tCrB);
// CLEAR rmem thread-level partition (accumulators)
cute::clear(tCrC);

// GEMM on rmem: (V,M,K) x (V,N,K) => (V,M,N)
cute::gemm(tiled_mma, tCrA, tCrB, tCrC);
// Equivalent to
// for(int k = 0; k < size<2>(tCrA); ++k)
//   for(int m = 0; m < size<1>(tCrC); ++m)
//     for(int n = 0; n < size<2>(tCrC); ++n)
//       tiled_mma.call(tCrA(_,m,k), tCrB(_,n,k), tCrC(_,m,n));

// AXPBY from rmem to gmem thread-level partitions
cute::axpby(alpha, tCrC, beta, tCgC);
// Equivalent to
// for(int i = 0; i < size(tCrC); ++i)
//   tCgC(i) = alpha * tCrC(i) + beta * tCgC(i)

对于上述代码，现在有许多关于计算和复制指令的时间交错的决策需要做出

• 仅将 rmem 分配为 A: (MMA,MMA_M) 以及 tg_ 33 和 tg_ 34Tensors，并在每次 k-block 迭代时复制到其中。
• 考虑 gmem 的多个 K 图块，并在每次 K 图块迭代中复制到 smem。
• 将上述复制阶段与计算阶段异步重叠。
• 通过寻找更好的 smem 布局来优化，从而改进 smem – > rmem 文案的访问模式。
• 通过为 gmem – > smem 复制找到高效的 TiledCopy 分区模式来进行优化。

这些问题被视为“时间微核”的一部分，而非 CuTe 提供的“空间微核”。通常，有关管线和 CuTe 张量指令执行的决策将由 CUTLASS 级别做出，并将在本系列的下一部分中进行讨论。

总结

总之，CuTe 通过抽象出张量布局和线程映射的低级细节，并为现代 NVIDIA GPU 上的密集线性代数提供统一的代数接口，使开发者能够编写更具可读性、可维护性和高性能的 CUDA 代码。

有关更多信息，您可以在 GitHub 上下载软件，阅读我们的文档，或加入我们的开发者论坛进行更深入的讨论。

致谢

感谢 Jack Kosaian、Mark Hoemmen、Haicheng Wu 和 Matt Nicely 为本文做出的贡献。特别感谢 Jay Shah、Paul VanKoughnett 和 Rya Asai 的 Colfax International 团队。