Programming Blackwell Tensor Cores with CuTe and CUTLASS

Cris Cecka, Mike Rubbelke (NVIDIA GTC | March 21, 2024)

目录

• 什么是 CUTLASS
• NVIDIA Blackwell 架构
• 使用 CuTe 针对 Blackwell 新特性进行编程
• 使用 CUTLASS 适配 Blackwell 新特性
• 结论与未来路线图
• 致谢

什么是 CUTLASS？

CUTLASS 是一个用于在各种规模和尺寸上进行密集计算的 CUDA C++ 模板库。

• 高性能实现: 为深度学习（DL）和高性能计算（HPC）中常见的核心计算核（kernels）提供了高性能的 Tensor Core 实现，例如：通用矩阵乘法（GEMM）、卷积、稀疏矩阵乘法（SPARSE GEMM）、分组矩阵乘法（Grouped GEMM）等。
• 广泛集成: 已被集成到多个深度学习和高性能计算库及框架中，包括 cuBLAS、cuDNN、TensorRT、PyTorch、TensorFlow 和 cuML 等。
• 灵活性与可定制性: 提供了灵活且可定制的组件，使开发者能够轻松编写针对 Tensor Core 的自定义计算核。
• 开源: CUTLASS 是一个开源项目，拥有活跃的社区、详细的文档、开放的贡献和众多的用户。项目地址：https://github.com/NVIDIA/cutlass。
• 跨代支持: 自 2018 年首次推出以来，CUTLASS 为支持 Tensor Core 的各代 NVIDIA 架构提供了支持，包括 Volta、Ampere、Hopper，以及最新的 Blackwell。

NVIDIA Blackwell 架构

Blackwell 硬件新特性概览：

1. Blackwell Tensor Cores - 支持 FP4/FP6 数据类型
2. Tensor Memory (TMEM) - 张量内存
3. 新的调度能力

1. Blackwell Tensor Cores

从 Hopper 到 Blackwell 的演进

Blackwell Tensor Cores 相比 Hopper 实现了显著的性能和功能提升：

• 吞吐量翻倍: 相比 Hopper，Blackwell 在 FP16、BF16、TF32、INT8 和 FP8 等主流数据类型上的吞吐量提升了 2 倍。
• Tensor Memory (TMEM): 使用张量内存（Tensor Memory）替代了寄存器内存（Register Memory）来作为 Tensor Core 的输入和输出。
• 异步 MMA 扩展: 将异步矩阵乘累加（MMA）操作的执行扩展到了 Epilogue（收尾阶段）。

下表对比了 Hopper 和 Blackwell 在 Tensor Core 操作上的主要区别：

阶段	Hopper	Blackwell
操作数加载	从寄存器加载	从张量内存加载
累加	寄存器	张量内存
指令发布	等待组内指令完成	异步执行，与 Epilogue 重叠
指令完成	同步提交	同步提交

扩展 Tensor Core 指令至 2 个 SM

Blackwell 架构将单个 Tensor Core 指令的执行范围扩展到了 2 个流式多处理器（SM）。

• 一个 2x1 的合作线程数组（CTA）集群中的一对 CTA 会被链接起来，并跨越 2 个 SM 执行。

  • 每个 SM 包含 4 个集群，每个集群 4 个 CTR。
  • 1 个集群 + 4 个 CTR 构成一个分区。
  • 2 个集群 + 8 个 CTR 跨 2 个 SM。

• 这对 CTA（CTA 0 和 CTA 1）协同工作，CTA 0 作为“领导者”，CTA 1 作为“链接者”，共同完成一次 MMA 操作。

• A、B 操作数和累加器（Accumulator）被均匀地分配到 2 个 SM 上，每个 SM 使用其本地的共享内存（SMEM）或张量内存（TMEM）。

支持新的块缩放数据格式（Block-scaled Formats）

Blackwell Tensor Core 硬件原生支持新的块缩放数据类型：MXFP8、MXFP6、MXFP4 和 MXINT8。

• MXFP6/MXFP4 支持混合整数输入。

• 性能提升:

  • MXFP8/MXFP6: 吞吐量是 Hopper FP8 的 2 倍。
  • MXFP4: 吞吐量是 Hopper FP8 的 4 倍。

• A 和 B 操作数的缩放因子矩阵需要从张量内存（TMEM）中获取。

下表总结了新的块缩放格式：

格式名称	数据格式	缩放格式	可用于
MXFP8	FP8 E5M2 或 E4M3	FP32 或 None	A, B, C, D
MXFP6	FP6 E3M2	FP32 或 None	A, B, C, D
MXFP4	FP4 E2M1	FP32 或 None	A, B, C, D
MXINT8	INT8	FP32	A, B

2. Blackwell Tensor Memory (TMEM)

TMEM 是每个 SM 上的内联内存，专用于 Tensor Core 的操作数输入和输出。

• 容量: 每个 SM 上的 TMEM 与寄存器文件（Register File）大小相同，为 256 KB。
• Tiled 结构: TMEM 呈分块（Tiled）形状，由 128 个通道（lanes）组成，每个通道 2 KB（512 列，每列 4 字节）。
• 访问限制: 线程只能访问其自身所在分区的 TMEM。

• 分配方式:

  • 编译器会自动管理 TMEM 的分配。
  • 开发者也可以使用 cuda::memcpy_async 显式地分配和管理 TMEM。

• 专用性: TMEM 专用于 Tensor Core (TC) 操作，不支持 SMT（同步多线程）操作。

• 数据移动:
  • 使用 ldmatrix 和 stmatrix 指令与 TMEM 进行数据交换，类似于 Hopper 架构。
  • 支持非线性寻址，数据以预定义的布局（pre-defined layouts）进行移动。

3. 新的调度能力

首选线程块集群 (Preferred Thread Block Clusters)

• Hopper 的局限: Hopper 架构引入了线程块集群（Thread Block Clusters），允许多个 CTA 协同工作。但一个集群内的所有 CTA 必须在同一个 GPC（Graphics Processing Cluster）上，这可能导致部分 SM 闲置。
• Blackwell 的改进: Blackwell 引入了“首选线程块集群”功能，允许以两种不同的线程块集群形状（Shape）来启动计算网格（grid）。
  • 动态选择: 硬件会优先尝试启动首选的、较大的集群配置；如果资源不足，它会自动回退到较小的集群配置来填充剩余的 SM。
  • 编程支持: CuTe 和 CUTLASS 在 Blackwell 上支持此功能。用户可以将两种集群形状作为运行时参数传递给计算核。

运行时持久调度 (Runtime Persistent Scheduling)

静态分块调度 (Static Tile Scheduling)

这是 CUTLASS 在 Hopper 架构上使用的持久化调度方法。

• 机制: 持久化 CTA 的数量等于 SM 的数量。输出分块（output tile）到 SM 的映射在计算核启动时就已固定。
• 问题: 如果某个 SM 因为上下文切换（context switch）而去执行另一个计算核，那么分配给它的输出分块将不会被重新分配。这会导致其他 SM 完成工作后，该 SM 仍有大量任务未完成，形成一个长尾（long instruction tail），影响整体效率。

动态分块调度 (Dynamic Tile Scheduling)

Blackwell 架构引入了新的硬件功能来解决静态调度的问题。
- 机制: Blackwell 允许用户在 SM 上通过 cuda::cluster_arrive_relaxed::fetch_add 指令以编程方式获取新的线程块集群。
- CUTLASS 实现: CUTLASS 利用此功能实现了动态持久化调度器（Dynamic Persistent Scheduler）。
- 输出分块到 SM 的映射是完全动态的，取决于各个 SM 的执行进度。当一个 SM 完成其任务后，它会主动获取新的任务。
- 这是 CUTLASS 在 Blackwell 上的默认调度器。

Blackwell 架构特性总结

• Blackwell Tensor Cores:

  • 主流数据类型吞吐量是 Hopper 的 2 倍。
  • 新增对块缩放格式（MXFP4/6/8）的支持。
  • 扩展 Tensor Core 指令至 2 个 SM，并采用完全异步的编程模型。

• Tensor Memory (TMEM):

  • 每个 SM 上新增了与寄存器文件同样大小的内存。
  • 专用于 Tensor Core 的输入和输出。

• 新的调度能力:

  • 首选线程块集群: 允许计算网格（Grid）以两种集群配置启动，提高硬件利用率。
  • 运行时持久调度: 实现输出分块到 SM 的动态映射，消除长尾效应。

使用 CuTe 针对 Blackwell 新特性进行编程

本章介绍如何使用 CuTe 库来为 NVIDIA Blackwell 架构的新特性进行编程，主要内容包括：

• 使用 Blackwell MMA 编程
• 使用 TMA
• 使用 MMA.2SM 和 TMA.2SM
• TMEM 累加器与 Epilogue

使用 Blackwell Tensor Core 编程 (MMA)

Blackwell Tensor Cores: CuTe Atoms

CuTe 的 MMA 原子操作 (atom) 为使用 MMA (Matrix Multiply-Accumulate) 提供了 PTX (Parallel Thread Execution) 和元数据。

• MMA_Op：一个 PTX tile

  • A、B 和 C 的矩阵描述符
  • 延迟相关信息
  • 针对源数据类型和目标数据类型的指令描述符

• Make_Traits：MMA 的模板

  • 定义 MMA_Atom 的默认 Dispatch
  • 定义片段布局
  • 分区 A、B 和 C 的模式
  • 定义指令描述符更新
  • 其他默认参数

Hopper 与 Blackwell 的区别

• Hopper tiled_mma 使用线程ID（tid）在 CTA（Cooperative Thread Array）内部进行分区。
• Blackwell tiled_mma 使用 CTA ID（cta_id）在 CTAs 之间进行分区。

---

Blackwell TiledMMA: CuTe Atoms

用户选择一个 MMA_Op 来创建一个 TiledMMA。CuTe 为每个 M100 Tensor Core 指令都提供了相应的 TiledMMA。

下图展示了如何从一个 MMA_Op 构建 TiledMMA，它描述了 CTA 级别的计算分块。右侧的代码示例演示了如何用指定的 MMA_Atom 和操作数布局来创建一个 TiledMMA。

---

Blackwell MMA GEMM 示例

这是一个完整的通用矩阵乘法（GEMM）示例，展示了从数据布局到计算执行的完整流程。

1. 在全局内存中表示张量

首先，定义输入矩阵 A、B 和输出/累加矩阵 C 在全局内存（Global Memory）中的布局。
- 矩阵 A: (M, K)
- 矩阵 B: (K, N)
- 矩阵 C: (M, N)

2. 创建分区视图

使用 Tile 来创建全局内存张量的分区视图。在 Blackwell 架构中，我们使用 mma_tiler 而不是 Hopper 中常用的 cta_tiler。这个 tiler 定义了计算任务如何在 CTA 网格上进行划分。

3. 切分张量

根据 global_coord（而不是 cta_id）从分区视图中为每个 CTA 切分出对应的数据块。这一步确定了每个 CTA 需要处理的 A、B 和 C 矩阵的具体部分。

4. 在 CTA 间划分 MMA Tiles

使用 TiledMma 对象来将整个 MMA 计算任务在不同的 CTA 之间进行划分。这决定了每个 CTA 内部线程块的计算范围和数据分工。

5. 创建 SMEM 布局并分配共享内存

为 CTA 本地的数据（A 和 B 的分块）创建共享内存（Shared Memory, SMEM）布局，并分配相应的内存空间。数据将从全局内存加载到共享内存，以供 Tensor Core 高效访问。

6. 创建 MMA 片段 (Fragments)

为 MMA 操作创建寄存器级别的 "Fragments"。这些 Fragments 是 MMA 指令直接消耗的数据单元，代表了将从共享内存加载到寄存器的数据。

7. 协同拷贝 (GMEM -> SMEM)

使用 cooperative_copy 将数据从全局内存（GMEM）异步拷贝到共享内存（SMEM）。这是一个协同操作，由 CTA 内的所有线程共同完成。

8. 执行 MMA (SMEM -> REG)

最后，执行 gemm 操作。数据从共享内存加载到寄存器（Fragments），然后由 Tensor Core 执行矩阵乘加运算，结果累加到 C 的 Fragments 中。

使用 TMA (Tensor Memory Accelerator)

TMA 是用于高效数据传输的硬件单元。

TMA Atom

CuTe 的 TMA 原子操作为使用 TMA 提供了 PTX 和元数据支持。

• TMA_Op - PTX for TMA:

  • COPY: 基本拷贝
  • COPY_MULTICAST: 组播拷贝
  • IM2COL: Im2col 转换
  • FILL: 填充操作

• TMA atom:

  • 选择 TMA Op
  • 提供 TMA 布局
  • 提供 Tile 以获取 TMA Tiling 信息

以下是 TMA_Op 和 Copy_Atom 的代码示例：

Blackwell MMA+TMA GEMM: 使用 TMA

为了在 Blackwell 架构上利用张量内存加速器（TMA）进行通用矩阵乘法（GEMM），我们需要对现有的基于全局内存（GMEM）和共享内存（SMEM）的 GEMM 核函数进行修改。

1. 使用 TMA 感知的全局内存张量

第一步是修改全局内存张量的定义方式。标准的 cute::make_tensor 用于创建通用张量视图。为了让 TMA 能够识别和处理这些张量，我们需要使用一个专门的构造函数 cute::make_tma_tensor。此函数位于 cute/tensor_map.hpp 中，它会为张量附加必要的信息，以便 TMA 硬件能够正确地执行加载操作。

2. TMA 分区与同步

定义了 TMA 张量后，下一步是规划数据加载。

• TMA 分区：需要应用 TMA 分区（TMA partitioning）来确定每个 tile 需要一个或多个 TMA 操作来完成数据加载。这通过 make_tma_copy 和 partition_S 来实现。
• TMA 同步：TMA 操作是异步的。为了确保在数据被使用前已经从全局内存加载到共享内存，必须使用 TMA 同步屏障（TMA barrier）。通过 cute::TMA_Barrier 创建一个屏障对象，并在主循环中使用 tma_arrive_and_wait(tma_bar, ...) 等待 TMA 操作完成。

通过这些修改，原有的 copy(gA, tAgA, ...) 操作被替换为 copy(tma_load_a, ...)，并配合同步原语，从而将数据加载任务卸载到 TMA 硬件上。

---

使用 MMA.2SM 和 TMA.2SM

Blackwell 架构引入了 .2sm 后缀的指令，允许单个 MMA（矩阵乘法累加）或 TMA 操作跨越两个流式多处理器（Streaming Multiprocessors, SMs）执行。这使得两个协作的 CTA（Cooperative Thread Array）可以共同处理一个更大的计算任务。

CuTe MODE 原子操作

CuTe 通过其“原子操作”（Atoms）抽象来支持这些新的硬件特性。

• MMA.2SM:
  • 分发的 PTX 指令为 mma.2sm。
  • 为了在 CuTe 中定义一个 2-SM 的 MMA 操作，我们需要指定一个 2x1 的 SM 拓扑。这通过在 MMA Atom 的定义中设置 SmShape = Shape<_2,_1> 来实现。这表示操作将在两个 SM 组成的集群中执行，布局为 2 行 1 列。

• TMA.2SM:

  • 分发的 PTX 指令为 tma.2sm。

  • .2sm 版本的 TMA copy 操作需要额外的调度信息：

    • CtaSchd (cute::CtaSchd)：用于在 SM 间进行分区。
    • Multicast (cute::Multicast)：一个掩码，用于确定到 SM 的映射。

  • 在 CuTe 中，通过向 make_tma_copy 函数传递这些额外的调度参数来构造一个 2-SM 的 TMA 操作。

内核修改

为了在 GEMM 内核中启用 2-SM 操作，需要进行以下修改：

1. 使用 blockIdx.z 进行 SM 映射:
   CUDA grid 的 Z 维度（blockIdx.z）现在被用作 SM 集群内的索引。例如，对于一个 2-SM 操作，cta_id_z 为 0 的 CTA 运行在一个 SM 上，cta_id_z 为 1 的 CTA 运行在另一个协作的 SM 上。

2. Leader CTA:
   集群中的一个 CTA（通常是 cta_id_z == 0 的那个）被指定为 "leader"。is_leader 变量用于识别 leader CTA。

   

3. 调度与同步:

   • 核心的 cute::gemm 计算循环仅由 leader CTA 执行。
   • 必须在关键节点插入 cute::cluster_sync()，这是一个新的同步原语，用于同步集群内所有协作的 CTAs。

   

---

TMEM 累加器与 Epilogue

Blackwell 架构引入了张量内存（Tensor Memory, TMEM），这是一个由 warp 寻址的、显式管理的内存空间，旨在作为寄存器和共享/全局内存之间的高效数据交换媒介，特别适用于 GEMM 的 Epilogue（收尾）阶段。

TMEM 硬件与 PTX

• 寻址: TMEM 是 warp 寻址的，地址由 128 个“level”和 512 个“column”组成。一个 warp 组内的每个 entry 只能访问 32 个 lane。
• 管理: TMEM 必须通过 PTX 指令显式分配和释放。
  • 分配: tex.mbarrier.init.shared.b64，分配粒度为 512B (16B rows x 32 columns)，返回一个 64-bit 的本地 TMEM 地址。
  • 释放: tex.mbarrier.invalidate.shared.b64。

TMEM 软件接口与操作

• 间接访问: TMEM 地址不能被直接解引用。数据移动必须通过 warp-group 范围内的状态指令集体完成，如 tmem.load.global 和 tmem.store.global。
• 固定模式: 这些指令的访存模式由预定义的队列（queues）决定，通过描述符（如 SM100_TMEM_LOAD_L2_ENABLED_DESC_A）来指定。不同的描述符对应不同的数据布局（行主序、列主序、转置等）。

CuTe 对 TMEM 的抽象

CuTe 为 TMEM 提供了高层抽象，简化了其使用：
- Copy_Atom 和 TiledCopy: CuTe 提供了用于 TMEM 的原子操作和 TiledCopy 特性，封装了 TMEM 的布局和复制逻辑。
- TmemTensor 创建: 可以使用 make_fragment_A, _B, _C 等函数创建与 MMA 操作布局一致的 TmemTensor。
- TiledCopy 创建: make_tiled_copy 函数可以基于 TmemTensor 创建一个 TiledCopy 对象，用于执行实际的数据传输。

CuTe TMEM Epilogue 示例

以下步骤展示了如何在 GEMM epilogue 中使用 TMEM 将累加器中的结果（矩阵 C）写回到全局内存（矩阵 D），同时可能进行量化等操作。

1. 准备阶段:

   • 首先，像往常一样计算出线程块的坐标和线程索引。

   

2. 分区与张量创建:

   • 获取输出张量的局部视图: 使用 local_tile 从全局输出张量 D_Global 中划分出当前 CTA 负责的局部视图 D_local。
   • 创建 TMEM 张量: 使用 make_fragment_C 将存储在寄存器中的累加器 rC 转换为一个 TMEM 张量 tC。这一步是逻辑上的转换，它为寄存器数据赋予了 TMEM 的布局信息，为后续的物理拷贝做准备。

   

接下来的步骤将是使用 cute::copy 指令，通过先前定义的 TiledCopy 对象，将寄存器中的 rC 复制到 TMEM，然后再从 TMEM 复制到全局内存的 D_local 中。TMEM 在此过程中充当了一个高性能的暂存区。

使用 CUTLASS 适配 Blackwell 新特性

本章介绍如何利用 CUTLASS 库来充分发挥 Blackwell 架构的新特性。

CuTe 与 CUTLASS 3.x 概念层级

CUTLASS 3.x 建立在一个分层的概念体系之上，底层是 CuTe 库，提供了对硬件指令的精细控制，顶层则提供了高生产力的编程接口。

• CuTe Atoms (底层): 封装了硬件指令（如 LDG, STS, MMA, TMA）及其元数据，提供最大程度的控制。
• CuTe Tensors and Tile views: 在 CuTe Atoms 基础上进行通用化的分块（Tiling），覆盖 CTA（线程块）和 Cluster（线程块簇）级别。
• Collective Mainloop/Epilogue: 组合了 Tile 迭代器和核心卡片（Core Cards），用于调度和执行数据流水线。Collective Builder 帮助实例化最优的 collectives。
• Kernel Layer: 围绕 collectives 进行动态分发，负责启动网格、Tile 调度和 Warp 专业化。
• Device Layer (顶层): 主机端的 Kernel 和接口，面向生产力。

支持众多 Kernel 变体

CUTLASS 的设计目标之一是支持大规模的 Kernel 组合，以应对不同的算法、数据类型、数据布局和性能优化选项。这导致了数百种不同的 Kernel 变体。

• 操作类型: 包括混合精度、逐通道/张量 GEMM、分组 GEMM、稀疏 GEMM 等。
• 数据类型与精度: 支持从 int4/fp8 到 fp64 以及复数等多种类型。
• 数据布局: 支持 NCHW、NHWC 等多种布局。
• 性能选项: 支持不同的 Tile 尺寸、Split-K 策略、用户自定义的 Main Loop 和 Epilogue 等。

CUTLASS Kernel 示例：从 Hopper 到 Blackwell

构建一个 CUTLASS Kernel 通常遵循以下三个步骤：
1. 选择 Mainloop: 定义核心的 MMA 计算以及输入数据的流入方式。这通过 CollectiveMainloop 来实现。
2. 选择 Epilogue: 定义如何对 MMA 的输出进行后处理。这通过 CollectiveEpilogue 来实现。
3. 组合成 Kernel: 使用调度策略将 Mainloop 和 Epilogue 组合在一起，形成 内核层 (Kernel Layer)。

将一个为 Hopper (SM90) 架构编写的 CUTLASS Kernel 迁移到 Blackwell (SM100) 架构非常直接，主要涉及以下修改：

1. 更改架构标签: 将 ArchTag 从 cutlass::arch::Sm90 修改为 cutlass::arch::Sm100_...。
2. 更新调度策略: 将 DispatchPolicy 从 Hopper 的 ...Sm90TmaGmma... 修改为 Blackwell 的 ...Sm100TmaWgmma...，以使用 Blackwell 引入的 Warp-Group MMA (WGMMA) 指令。
3. 调整 Tile 形状: TileShape 可能需要从基于 CTA 的定义调整为基于 MMA 的定义，以更好地匹配新硬件的特性。

如代码高亮所示，迁移工作主要集中在修改几个关键的类型别名（using 声明），而 Kernel 的主体逻辑保持不变。

Collective Mainloop 与 Builder

• Collective Mainloop: 它是对数据流和 MMA 计算的封装，通过 cute::TiledMma 和 cute::TiledCopy 实现。它内部实现了流水线机制（如双缓冲），以重叠数据传输和计算，并通过定义同步点（如 tma_load, mma_prologue）来协调操作。
• Collective Builder: 这是一个高级辅助工具，可根据架构、Tile 形状、数据加载策略等高层级参数，自动推断并构建出最优的 CollectiveMainloop 类型。这极大地简化了 Kernel 的定义过程。

// 使用 CollectiveBuilder 简化 Mainloop 的定义
using Mainloop = typename CollectiveBuilder<
    cutlass::arch::Sm100,
    cutlass::gemm::MainloopSm100TmaWgmmaFp16Crosswise,
    Shape<_128, _128, _32>, // TileShape_MNK
    /* ... other params ... */
>::CollectiveType;

为 Blackwell 设计的新集合 (New Collectives for Blackwell)

为支持新的硬件特性，CUTLASS 提供了新的 Collectives 和 Mainloop 实现。

Blackwell 架构引入了新的 WGMMA 指令，所有的 Mainloop 都基于 TmaWGMMA。CUTLASS 为此提供了新的 Collectives，以支持不同规模的 MMA 操作。
- 非线程块级 MMA (Per-Warp): 适用于 MMA 计算在单个 Warp 内部完成的场景。这通过使用 MainloopSm100TmaWgmma*PerWarp* 调度策略来实现。
- 线程块级 MMA (Block-wide/Crosswise): 适用于需要整个线程块协作完成的更大规模 MMA 计算。这通过使用 MainloopSm100TmaWgmma*Crosswise* 调度策略来实现。

CUTLASS 3.6 预调优了以下集合：
- 针对非块缩放 MMA 的带有 TMA 加载的密集 GEMM
- MainloopSm90TmaGmmaRmemAsync

• 针对块缩放 MMA 的带有 TMA 加载的密集 GEMM

  • MainloopSm100TmaGmmaRmemAsync

• 密集分组 GEMM

  • MainloopSm90TmaGmmaRmemAsyncGrouped

• 带有 TMA 加载的 h1688gemm 卷积

  • MainloopSm90TmaGmmaRmemAsyncImplicitGemm

• 快速模拟 FP32 密集 GEMM

  • MainloopSm90TmaGmmaSimtRmemAsync

• 带有软件块缩放的密集 GEMM

  • MainloopSm90TmaGmmaRmemAsyncSoftwareScaling

新的主循环构建器特性 (New Mainloop Builder Features)

运行时 - 首选线程块集群 (Runtime - Preferred Thread Block Clusters)

• 支持将首选线程块集群作为运行时参数传递给内核。
• 无需为不同的线程块集群大小编译单独的内核。

运行时数据类型 (Runtime Data Types)
- 支持将数据类型作为运行时参数指定给内核。
- 无需为具有相同位宽但类型不同的参数（例如 s4 和 u4）编译单独的内核，两者都由 cute::uint_sub_byte_t 分派。

Collective Epilogue 与 Builder

Collective Epilogue
- 使用 cute::Tensor 进行灵活的输出后处理。
- Hopper 支持基于 TMA 的收尾（Sm90Tma）。
- Blackwell 完全支持 bfloat 类型和基于 TMT 的收尾。

Collective Builder
- 基于问题大小、设备、切片形状和数据并行策略，在运行时选择最优的集合。
- 调度内部的启发式方法 (Heuristics)。

为 Blackwell 设计的新的预调优集合收尾

为 Blackwell 架构引入了新的收尾（epilogue）实现，这些实现从 TMA 和 adaptor::CUT 进行了重构。

• 支持 sm90/sm100 的 TMA 存储收尾及其对应的 ema/grouped GEMM

  • Sm100TmaSpecialized
  • Sm100EmaTmaFusionSpecialized
  • Sm100TmaGroupedSpecialized
  • 所有实现都支持完整的 CUT 融合。

• 不使用任何共享内存进行 D 输出存储的 TMT (Direct) 存储收尾

  • Sm100TmtSpecialized
  • Sm100EmaTmtFusionSpecialized
  • 所有实现都支持完整的 CUT 融合。

Blackwell TMEM Epilogue (收尾阶段)

Blackwell 架构引入了张量内存 (TMEM) 来优化 Epilogue 阶段的操作。该流程旨在将累加器中的数据高效地写回全局内存 (GMEM)。

1. 累加器作为起点

Epilogue 流程的输入是 MMA 计算完成后存储在寄存器中的累加器 tC(D)。

// 为累加器定义一个 Gmem 布局
auto tC_gmem = make_tensor(make_gmem_ptr(C), LayoutC{});

2. 创建 TMEM 拷贝操作

为了将累加器的数据输出到 GMEM，需要创建一个 TmaCopy 对象来管理通过 TMEM 进行的数据传输。

// 1. 创建一个 TmaCopy 对象
auto gmem_tma = make_tma_copy(
    TmaMode{},
    tC_gmem);

// 2. 获取一个用于存放 TmaCopy 输入数据的 Fragment
auto tC_frg = make_fragment(tC_acc, tile_shape_MN);

3. TMEM 目标分区

在执行拷贝之前，需要根据线程索引对目标 GMEM 张量进行分区，以确保每个线程正确写入其负责的数据块。

// 3. 对 Gmem 张量进行分区
auto tC_gmem_part = partition_C(gmem_tma, thread_idx);

4. 执行到 GMEM 的拷贝

最后，调用 copy 函数，将累加器片段中的数据通过 TMEM 异步拷贝到已分区的 GMEM 目标位置。

// 4. 将累加器片段拷贝到已分区的 Gmem 目标
copy(gmem_tma, tC_gmem_part, tC_frg);

CUTLASS 内核层 (CUTLASS Kernel Layer)

内核层将所有组件（主循环、收尾）整合在一起。

• 切片调度 (Tile scheduling)：内核结构是 Warp-specialized (Warp 专用)。

  • 在集合之外的最外层循环。
  • 默认使用动态持久化。
  • 支持 StreamK 和 Data Parallelism (数据并行)。

• 内核结构 (Kernel Structure)：

  • 新的 Main 专业化内核用于 Blackwell，以处理 TMA 和 Epilogue。
  • KernelTmaWarpSpecialized：当不需要后处理 MMA 结果时使用（例如，仅进行 GEMM）。
  • KernelTmaWarpSpecializedCooperative：当需要后处理 MMA 结果时使用（例如，在 Epilogue 中进行融合操作）。

Warp 专用持久化内核 - Hopper vs Blackwell 对比

该图对比了 Hopper 和 Blackwell 架构在 Warp 专用持久化内核上的执行模型。

• Hopper: Ping-pong Mainloop 优化

  • 在 Hopper 中，Epilogue 和 MMA 在相同的线程上执行，通过“乒乓”机制在两组 MMA 之间交替执行，以重叠开销。

• Blackwell: 真正并发 (True Concurrency)

  • 在 Blackwell 中，Epilogue 和 MMA 在独立的线程上执行。这实现了真正的并发，不需要两个潜在的线程来重叠 Epilogue 的执行。

结论与未来路线图

CUTLASS 3.8 新特性

为 Blackwell 提供了全面的高性能支持。

特性 (Features):

• 全面支持所有 Blackwell Tensor Core 媒体类型。
• 全面支持 TMA 和所有新的 Blackwell 拷贝指令：TMMA, TMT。
• 新的 CUTLASS 流水线，用于 Blackwell 上扩展的异步执行专业化：JCTM, CLT。
• 用于 Blackwell 异步专业化的新 CUTLASS 集合和内核。

内核 (Kernels):
- 适用于所有 Hopper 类型的密集 GEMM 内核（FP32、TF32、FP16、BF16、FP8），性能比 Hopper 提升 2 倍。
- 适用于新 Hopper 数据类型的密集 GEMM 内核（benchmarking a la fp6/fp4），性能比 Hopper 提升高达 4 倍。
- 适用于所有 Hopper 类型的密集隐式 GEMM 卷积内核（FP32、TF32、FP16、BF16、FP8），性能比 Hopper 提升 2 倍。
- 使用 Blackwell Tensor Cores 的分组 GEMM 内核，支持块缩放类型。
- 使用 M16 矩阵的模拟 FP32 内核，利用 Blackwell Tensor Cores。

2025 年规划

CUTLASS 4.x - 将 CUTLASS 设备级与 Python 互操作性结合

CUTLASS C++
- 关键特性
- 使用 Blackwell Tensor Cores 的 Blackwell 专用稠密 GEMM 内核。
- Blackwell fp4/fp6 支持，完全与 cuBLASLt 早期版本对齐。
- Blackwell INT2/INT1 支持，带有新的 MINT 块缩放 GEMM。
- 稀疏性支持：建立在 fp4/fp6 块缩放 GEMM、INT2/INT1 和 cuSPARSELt 之上。

• 其他近期（6-12个月内）特性
  • Blackwell 上的 GEMM 支持。
  • BF16 分组 GEMM 内核。
  • INT8/FP8 复杂 GEMM 内核。
  • INT8/FP8 Transposed/Non-Transposed GEMM 的分组 GEMM 支持。
  • 分布式 GEMM 和 cuBLASLt。
  • INT8 SD-A。
  • ... 以及更多。

致谢

向以下社区和开发者致谢：

• CUTLASS GitHub 社区: 9k+ 星标, 6.4M+ 下载, 1.8k+ 复刻, 以及众多活跃用户。
• CUTLASS 开发者: 列出了众多核心开发人员。
• CuTe 开发者: Che-Hsien Lin, Vijay Thakkar。
• CUTLASS Product Manager: Matthew Nicely。
• 贡献者与致谢: 列出了来自学术界和工业界的众多贡献者。