Developing Optimal CUDA Kernels on Hopper Tensor Cores

Pradeep Ramani, Cris Cecka | March 22, 2023

目录

• CUTLASS 简介
• NVIDIA Hopper 上的 Tensor Core 性能
• 发展路线图 (Roadmap
• 议程 (Agenda
• NVIDIA Hopper 架构
• CuTe
• CUTLASS 3.0
• CUTLASS Python
• 结论
• 致谢 (Acknowledgements
• 参考文献

CUTLASS 简介

CUTLASS (CUDA C++ Template Library for Deep Learning and High Performance Computing) 是一个用于在各种范围和规模上进行矩阵计算的优化 CUDA C++ 模板库。

关键信息:
- 开源: https://github.com/NVIDIA/cutlass (new BSD license)
- 最新版本: CUTLASS 3.0
- 文档: https://github.com/NVIDIA/cutlass#documentation
- 功能: https://github.com/NVIDIA/cutlass/blob/master/media/docs/functionality.md
- 历届 GTC 演讲: GTC'18, GTC'19, GTC'20, GTC'21, GTC'22

NVIDIA Hopper 上的 Tensor Core 性能

使用 CUTLASS 3.0 和 CUDA 12.0 Toolkit 在 NVIDIA H100 上的性能表现。下图展示了在不同 GPU 架构（A100, A40, H100, L40）上，各种 GEMM (通用矩阵乘法) 配置的相对峰值性能。结果显示，Hopper 架构（H100, L40）在多种精度和数据类型下均能达到接近理论峰值的性能。

发展路线图 (Roadmap)

以下是 CUTLASS 在 2023 年的开发路线图，可能会有变动。

议程 (Agenda)

本次演讲将涵盖以下主题：
- Hopper 架构 (Hopper Architecture)
- CuTe
- CUTLASS 3.0
- CUTLASS Python
- 结论 (Conclusion)

NVIDIA Hopper 架构

NVIDIA H100 引入了多项架构改进，以提升性能。

主要特性

• 更快的新 Tensor Core 指令:

  • 16 位浮点运算：比 F32 CUDA Cores 快 16 倍和 32 倍。
  • 8 位浮点运算：比 F32 CUDA Cores 快 32 倍和 64 倍。
  • 改进的整数运算：比 F32 CUDA Cores 快 32 倍和 64 倍。
  • 改进的 32 位浮点运算：比 F32 CUDA Cores 快 8 倍和 16 倍。
  • 改进的 IEEE 双精度运算：比 F64 CUDA Cores 快 2 倍。

• 新的线程组织层次 - 线程块集群 (Thread Block Clusters): 帮助在线程块之间实现数据的优化共享。

• 额外的数据类型: 支持 8 位浮点类型 E5M2 和 E4M3。
• 使用 TMA 的异步复制 (Asynchronous Copy):
  • TMA (Tensor Memory Accelerator) 可以为全局加载和经过变换（swizzled）的共享内存存储执行地址计算，并进行边界检查。
  • 能够在线程块集群中广播数据。

更多细节请参阅 "NVIDIA H100 Tensor Core GPU Architecture" 白皮书。

Hopper 架构 - Tensor Core 运算性能

下表比较了 Hopper (H100)、Ampere (A100) 和 Volta (V100) 架构上 Tensor Core 运算的理论峰值 TFLOPS。Hopper 架构在所有数据类型上都展现出显著的性能飞跃。

Tensor Core 运算: 基本形态

Tensor Core 的核心是矩阵乘加运算：D = op(A, B) + D。

• M-by-N-by-K 矩阵运算:
  • 这是一个 Warp-Group 范围的异步集体操作。
  • Warp-Group 内的 128 个线程共同持有 D 操作数。
  • 操作数 A 和 B 可从寄存器内存加载，但使用描述符直接从共享内存加载是最高效的方式。
  • M、N、K 的有效值取决于数据类型，例如 M=64; N ∈ [8, 256]; K ∈ [8, 16, 32, 256]。

F16 * F16 + F32 运算示例

下图展示了一个 64-by-N-by-16 形态的 F16 乘法与 F32 累加操作的示例，并给出了对应的 wgmma.mma_async 指令的汇编代码。

• wgmma.mma_async 指令会发起一个 Warp-Group 范围的异步 MxNxK 矩阵乘加操作。
• 矩阵 A 和 B 使用描述符直接从共享内存加载。
• 该指令本身支持对输入进行转置和缩放。

新的数据移动方式：使用 TMA 的异步复制

Hopper 架构引入了 Tensor Memory Accelerator (TMA) 来实现一种新的高效数据移动方式。

• 新机制: 实现了从全局内存到共享内存或从共享内存到全局内存的异步复制。
• 高效: 避免了通过寄存器文件的额外往返，直接读写共享内存。
• 多播 (Multicast): 可以在集群内的其他线程块间多播数据并更新异步屏障 (async barriers)。
• 安全: 执行数据边界检查，并在越界时自动向共享内存填充零。
• 灵活: 能够在共享内存中以多种不同的变换（swizzled）格式读写数据。

• Hopper 支持的类型:

  • 全局内存与共享内存之间的 1D-5D 张量复制。
  • 两种模式：TILED 和 IM2COL（对卷积有用）。
  • 支持归约（reductions）和简化批量复制的变体。

• 描述符（Descriptors）通过驱动 API 创建，用于传递待复制张量的信息。

TMA 多播 (Multicast) 可视化

TMA 支持将数据从全局内存高效地多播到块集群 (Block Cluster) 内的多个 SM (Streaming Multiprocessor)。一个 SM 的 TMA 从全局内存读取数据后，可以将数据和屏障更新直接多播给同一集群中的其他 SM，从而避免了重复的全局内存读取。

copy.async.bulk

copy.async.bulk 是一系列指令，用于发出一个 warp 统一的异步复制操作。它支持 2D、平铺模式（tiled mode）和多播（multicast）。

• 功能:

  • copy.async.bulk 系列指令发出一个 warp 统一的异步复制操作。
  • 复制可以在共享内存 (Shared) 到全局内存 (Global) 或全局内存到共享内存之间进行。

• 多播:

  • 多播掩码 (Multicast mask) 指定了集群 (Cluster) 中的哪些 CTA (Cooperative Thread Arrays) 也需要接收这份数据。
  • 此功能仅适用于从全局内存到共享内存的复制。

• 描述符 (Descriptors):

  • 描述符在主机端使用 CUDA 驱动 API cuTensorMapEncode* 创建，并通过 grid private 常量传递到设备端。

• 汇编示例与参考:

  • 右侧代码展示了 cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster 指令的用法。
  • 参考: cute::SM90_TMA_LOAD_2D_MULTICAST

CuTe

什么是 CuTe？

CuTe 是一个用于 CUDA Tensors 的库。

• CuTe 提供 Layout 和 Tensor 类型

  • 紧凑地封装了数据的类型 (type)、形状 (shape)、内存空间 (memory space) 和 布局 (layout)。
  • 为用户执行复杂的索引计算。
  • 提供用于定义和操作分层多维布局的抽象。
  • 形状 (Shapes) 和步幅 (strides) 可以是完全静态的、动态的或混合的。

• 一个关于 Layout 的形式代数

  • 布局可以被组合和操作。
  • 可以从简单的、不变的原始布局构建复杂的布局。
  • 跨其他布局进行分区，以在每个级别上保持逻辑一致性。

为什么需要 CuTe？

因为布局 (Layouts) 和张量 (Tensors) 无处不在。

• Layout 包含了 CUTLASS-2 迭代器的功能

  • 将复杂性凝聚到单个实现中。

• 用于操作 Layout 的形式代数

  • 包括组合 (composition)、补 (complement)、右逆/左逆 (right_inverse, left_inverse)、"积" ("product") 和 "除" ("divide") 等操作。

• 为线程和数据提供统一的 Layout

从 CUTLASS 2.x 到 3.x 的演进，体现了从大量特定的、硬编码的布局类型（如 RowMajor, ColumnMajor, TensorNCHW 等）到一个统一的、可组合的 Layout<Shape, Stride> 抽象的转变，极大地简化了代码。

CuTe in CUTLASS-3.0

在 CUTLASS 3.0 中，CuTe 的应用主要体现在以下几个方面：

• 布局 (Layouts)
• 代数 (Algebra)
• 张量核心 (Tensor Cores)

布局 (Layouts) 将坐标映射到存储

布局定义了一个函数，它将逻辑上的多维坐标映射到线性的物理存储地址。下图展示了对于一个形状为 (4,3) 的逻辑张量，不同的布局函数（行主序、列主序、带填充、混合模式）如何将其映射到一维内存空间。

• Row-major (行主序): f: (int, int) => int
• Col-major (列主序): g: (int, int) => int
• Padded (带填充): h: (int, int) => int
• Mixed (混合): a: (int, int) => int

布局表示 (Layout Representation)

布局可以通过形状 (Shapes) 和步幅 (Strides) 来表示。内存中的偏移量可以通过坐标与步幅的内积计算得出：f(coord) = inner_product(coord, stride)。

• 示例 1: Shape: (2,3), Stride: (1,2)
• 示例 2: Shape: (2,3), Stride: (3,1)
• 示例 3 (多维): Shape: (2,2,2), Stride: (4,1,2)

下图展示了不同的形状和步幅组合如何产生不同的内存布局。

分层布局 (Hierarchical Layouts)

张量可以具有分层的形状和折叠的模式。一个逻辑上具有分层形状的张量可以被“折叠”并视为一个简单的矩阵，这有助于简化操作和理解。

• 一个 Shape: (2,2,2)、Stride: (4,1,2) 的张量。
• 可以被折叠并视为一个 Shape: (2,4)、Stride: (4,1) 的矩阵。
• 坐标到偏移量的映射函数保持为：f(coord) = inner_product(coord, stride)。

• 同样地，上述张量也可以被折叠并视为一个 Shape: (4,2)、Stride: (?,1) 的矩阵。

• 或者，可以将其视为一个嵌套的形状，例如 Shape: ((2,2),2)，其对应的步幅为 Stride: ((4,2),1)。这种表示法清晰地揭示了数据的层次结构。

核心设计要素 (Key design ingredients)

CuTe Tensors 的核心设计要素包括：

• Shape: IntTuple 或 tuple of IntTuples 的概念。

  • 例如：N, (N,M), (N,M,P), ((N1,N2),N3), ((Na,Nb),(N2,(N3,N4,N5)))

• Stride: DTuple 或 tuple of DTuples 的概念。

• Layout<Shape, Stride>:

  • 一个从 n-D 逻辑坐标到 1-D 索引坐标的映射，以及其逆映射（如果可用）。
  • (I) ⇔ (i, j) ⇔ (i, (j1, j2)) ⇔ [k]

• Tensor<Ptr, Layout>:

  • Layout 与底层随机访问迭代器（如 T[], T*, smem_ptr<T>, gmem_ptr<T>）的组合。

• 布局的代数 (Algebra of Layouts):

  • concatenation(Layouts...) -> Layout
  • composition(LayoutA, LayoutB) -> Layout
  • complement(Layout, M) -> Layout
  • right_inverse(Layout) -> Layout

布局映射 (Layout Mappings)

布局定义了从逻辑坐标到线性索引的完整映射流程。

1. 逻辑 1-D 坐标 (e.g., A(I))
2. 通过坐标映射 (Coordinate Mapping) 转换为 逻辑 n-D 坐标 (e.g., A(i,j))
3. 通过坐标映射 (Coordinate Mapping) 转换为 逻辑 h-D (分层) 坐标 (e.g., A(i,(j1,j2)))

4. 通过索引映射 (Index Mapping) 转换为 线性 1-D 存储索引 (e.g., A[k])

5. Shape 定义了坐标映射: (I) ⇔ (i,j) ⇔ (i,(j1,j2))

6. Stride 定义了索引映射: (i,(j1,j2)) ⇔ [k]

布局示例 (Example Layout)

以下是一个使用 CuTe 构建复杂布局的示例。

• 代码: 通过 make_layout, blocked_product 和 make_tensor 等操作创建了一个分块的 Morton 序（Z-order curve）布局。
• Shape: shape(A) = ((2,2,2),(2,(2,2)))

• 逻辑坐标:

  • 元素可以被 1D、2D 或更高维度的逻辑坐标索引。
  • 例如，A(37)、A(5,4)、A(1,2),(0,2) 等都指向同一个物理地址，对应的值为 49。

• 切片 (Slicing):

  • 可以沿着逻辑子边界进行切片操作。
  • A(_,2) 和 A(_,1),(_,2) 展示了如何提取数据的子集。

布局转写示例 (Layout Transcription Examples)

CuTe 的 Layout<Shape, Stride> 抽象能够表示许多传统和复杂的内存布局。

• ColumnMajor + Padding: Layout<Shape<_3,_5>, Stride<_1,_4>>
• PitchLinear<>: Layout<Shape<_4,_3>, Stride<_3,_1>> (注意，这里的形状和步幅是基于 ((4,3)) 和 ((3,1)) 的元组)
• ColumnMajorInterleaved<4>: Layout<Shape<_4,Shape<_4,_2>>, Stride<_4,Stride<_1,_16>>>

Swizzle 布局示例 (Swizzle Layout Examples)

Swizzle 布局通过重排数据来优化内存访问模式，以减少缓存冲突和提高带宽利用率，这在 GPU 编程中非常常见。CuTe 同样可以简洁地表示这些复杂的布局。

• Swizzle<2,2,2>: Layout<Shape<_4,Shape<_4>>, Stride<_4,Stride<_1,_16>>>
• Swizzle<2,0,2>: Layout<Shape<_4,_16>, Stride<_1,_4>>
• Swizzle<3,0,0>: Layout<Shape<_8,_8>, Stride<_8,_1>>

组合能力 (Composition Power) 示例

幻灯片通过一个逐步构建的示例来阐述“组合能力”的概念。

首先，我们有一个一维的逻辑数据数组。

接着，对这个数组进行分区，定义了值的布局。此布局由 ((2,3)) 和 ((1,4)) 这样的形状元组以及一个具体的值索引表 Values 来描述。

然后，引入了第二层分区，将不同的值集（以不同颜色表示）分配给不同的线程。

这个多层分区是通过线程布局来定义的。Threads 布局由 ((2, 2)) 和 ((2, 12)) 等元组描述。这建立了一个从线程ID（tid）和值ID（vid）到最终数据坐标（coord c）的映射函数。

这个过程的核心思想是函数组合。线程布局（Threads）和值布局（Values）可以被看作两个独立的函数。将它们组合起来，就可以得到一个将数据划分到不同线程的最终布局。

在代码实现中，这个过程表现为：
1. 创建一个输入张量（make_tensor）。
2. 将输入张量与一个线程-值（thr_val）布局进行组合（composition），生成一个统一的线程-值视图（input_TV）。
3. 通过线程ID（tid）对这个组合视图进行切片，从而得到每个线程负责的数据子集（thr_input）。

核心思想总结：给定一个从（线程，值）到坐标的映射，分区（Partitioning）本质上就是函数组合（Functional Composition）后进行切片（Slicing）。

MMA 特征 (MMA Traits)

这部分内容定义了不同NVIDIA GPU架构上矩阵乘累加（MMA）操作的元数据和内存布局。这些定义（Traits）封装了特定硬件指令的细节。

Volta FP16 8x8x4 元数据

MMA_Traits 结构体为Volta架构（SM70）上的 8x8x4 FP16 MMA操作定义了元数据。
- MMA形状 (Shape_MNK): 8x8x4
- 线程网格 (ThrID): 4x2
- 数据类型: 输入为 half_t，累加器为 float。
- 布局 (Layouts): 为矩阵A、B、C定义了数据在线程和寄存器中的具体布局。例如，ALayout 将 (T8, V4) 映射到 (M8, K4)，意味着8个线程和4个值构成了矩阵A的一个8x4的块。

Ampere FP64 8x8x4 元数据

为Ampere架构（SM80）上的 8x8x4 FP64 MMA操作定义了元数据。
- MMA形状 (Shape_MNK): 8x8x4
- 线程数 (ThrID): 32
- 数据类型: 输入和累加器均为 double。
- 布局 (Layouts): ALayout 将 (T32, V1) 映射到 (M8, K4)。

Ampere FP16 16x8x8 元数据

为Ampere架构（SM80）上的 16x8x8 FP16 MMA操作定义了元数据。
- MMA形状 (Shape_MNK): 16x8x8
- 线程数 (ThrID): 32
- 数据类型: 输入为 half_t，累加器为 float。
- 布局 (Layouts): ALayout 将 (T32, V4) 映射到 (M16, K8)。

Hopper FP16 64x16x16 元数据

为Hopper架构（SM90）上的 64x16x16 FP16 MMA操作定义了元数据。此定义使用了更通用的 GMMA (可能指代 Generic MMA) 模板。
- MMA形状 (Shape_MNK): 64x16x16
- 线程数 (ThrID): 128
- 数据类型: 输入为 half_t，累加器为 float。
- 布局 (Layouts): 使用 GMMA::ABLayout 等更高级的抽象来定义，以适应Hopper架构的特性。

布局代数 (Layout Algebra)

Layout Algebra 定义了一套在布局上进行操作的代数法则，允许以声明式的方式构建和变换复杂的内存布局。

逻辑积 (Logical Product)

公式: f_A ⊗ g_B = (f_A ∘ g_B) → (f_A, h_B')

描述: "生成一个布局，其中布局B的每个元素都是一个布局A。" 这是一种创建分块或瓦片式布局的操作。常见的操作包括 logical_product、blocked_product、raked_product 和 tile_to_shape。

逻辑除 (Logical Divide)

公式: f_A ⊘ g_B = f_A ∘ (g_B, g_B*) → (h_B', l_C)

描述: "将布局A拆分为由布局B指向的元素和其他剩余部分。" 这是一种对布局进行解构或分区的操作，与逻辑积互逆。常见的操作包括 logical_divide、zipped_divide 和 tiled_divide。

更多关于CuTe库的讨论和示例，请参阅其官方文档：
https://github.com/NVIDIA/cutlass/tree/master/media/docs/cute

CUTLASS 3.0

CUTLASS 3.0 的新特性

CUTLASS 3.1：
- 针对 TF32 的寄存器支持的 WGMMA (Warp Group Matrix Multiply Accumulate) 内核。
- 用于 CUTLASS 的全新 Pythonic 接口。
- 具有融合功能的高效 Epilogue。

CUTLASS 3.0：
- 利用 CuTe 后端进行了一次重大的重构。
- 高效的 Hopper Tensor Core 指令与使用 TMA (Tensor Memory Accelerator) 的异步拷贝。
- Warp Specialized（Warp 特化）和 Persistent（持久化）内核实现。
- Collective Builders、文档、性能分析器支持、PyCUTLASS 集成、SDK 示例等。

CUTLASS 2.11：
- 针对 Ampere 内核的 Fused MHA（多头注意力）。
- Stream-K - 一个新的通用 Split-K 实现。
- 支持新 Hopper 双精度指令的 BLAS3 功能。

NVIDIA Hopper 上的 Tensor Core 性能

以下图表展示了使用 CUTLASS 3.0、CUDA 12.0 Toolkit 在 NVIDIA H100 上的 Tensor Core 性能。实验中 m=2048, n=8848。

• 混合精度浮点 (Mixed Precision Floating Point):

  • Tensor Core BF16, F16 相较于 FP32 SIMT GEMM 实现了高达 19.1倍 的加速。
  • Tensor Core TF32 相较于 FP32 SIMT GEMM 实现了高达 10.3倍 的加速。

• 双精度浮点 (Double Precision Floating Point):

  • Tensor Core F64 相较于 FP64 SIMT GEMM 实现了高达 2.2倍 的加速。

• 混合精度整数 (Mixed Precision Integer):

  • Tensor Core INT8 相较于 INT8 SIMT GEMM 实现了高达 14.4倍 的加速。

分块 GEMM (Blocked GEMM) 回顾

传统的通用矩阵乘法（GEMM）采用分块、分层的模型，在共享内存（Shared Memory）和寄存器（Registers）中复用数据。数据流如下：
全局内存 (Global Memory) -> 共享内存 (Shared Memory) -> 寄存器文件 (Register File) -> CUDA/Tensor Cores -> SMEM -> CUDA Cores -> 全局内存 (Global Memory)。

更多关于此模型的详细信息，请参见 CUTLASS GTC 2018 和 2020 的演讲。

Hopper 架构带来的变化

在 Hopper 架构中，数据流发生了变化。它采用了线程块集群分块（Block Cluster Tiled），实现了全局内存的大块数据拷贝，并直接从共享内存中复用数据。
新的数据流如下：
全局内存 (Global Memory) -> 共享内存 (Shared Memory) -> Tensor Cores -> CUDA Cores -> SMEM -> 全局内存 (Global Memory)。

CUTLASS 3 的概念性 GEMM 层次结构

CUTLASS 3 的层次结构不再以硬件层次为中心。

• 设备层 (Device layer): 主机端（host side）的设置和接口。

• 内核层 (Kernel layer): 启动 API、网格规划逻辑、负载均衡调度和内核线程调度。

  • 概念上：网格中所有线程块/集群的集合。

• 集合层 (Collective layer): 主循环，用于协调具有特定架构同步功能的拷贝/数学微内核。

  • 概念上：动态调用拷贝/数学微内核，以计算外积的累积内积。

• 分块 MMA/拷贝层 (Tiled MMA/Copy): GPU 微内核接口。

  • 概念上：构建 GPU 微内核所跨越的最大线程集合，用于数学/拷贝操作。

• 原子层 (Atom layer): 架构指令及其相关的元信息。

  • 概念上：必须参与架构加速的指定数学/拷贝操作的最少线程数和值。

上图展示了 CUTLASS 3 的概念性 GEMM 层次结构：

• Device Layer (设备层): Gemm Universal, Conv, etc.
• Kernel Layer (内核层): Multistage, Specialized, Specialized Persistent, ...
• Collective Layer (集合层): SM70, SM90 Multistage, SM90 TMA, SM90 TMA Specialized, ...
• Tiled Copy & Tiled MMA (分块拷贝与分块 MMA 层)
• Atom (原子层): LDMATRIX, CP.ASYNC, TMA, MMA SM80, MMA SM90

CUTLASS 3.x API 入口点

CUTLASS 3.x 减少了 API 的表面积。

• 设备层 (Device layer): cutlass::gemm::device::GemmUniversalAdapter<>

  • 可用于 2.x 或 3.x API 内核。
  • 是内核类型的无状态包装器。

• 内核层 (Kernel layer): cutlass::gemm::kernel::GemmUniversal<>

  • 将 GEMM 视为集合主循环 (collective mainloop) 和集合结尾 (collective epilogue) 的组合。
  • 每个新的内核调度都是针对调度标签 (schedule tags) 分派的特化版本。

• 集合层 (Collective layer): cutlass::gemm::collective::CollectiveMma<>

  • 根据定义了它们可以组合的内核调度集的策略进行分派。

• 微内核层 (Microkernel layer): cute::TiledMma<> 和 cute::TiledCopy<>

  • 在广泛的 GPU 架构中提供稳健的表示。

• 在各处使用静态断言（static asserts）来防止无效的组合或不正确的布局。

CUTLASS 3: 内核 API

cutlass::gemm::kernel::GemmUniversal<>

• 每个 GEMM 都是一个主循环 (mainloop) 和一个结尾 (epilogue) 的融合。
• 负责线程块/集群范围的交换（swizzling）、网格规划逻辑、负载均衡调度。
• 通过 Warp 特化的集合进行线程调度。
• 通过分派策略的 Schedule 标签进行选择。
  • 经验法则：
    • 每个主循环变体都有其自己的策略类型用于分派。
    • 它可以与一组它指定并检查的内核调度进行组合。

以下调度策略示例展示了 Warp 特化的主循环如何与持久化和非持久化内核调度组合：

template<
  int Stages_,
  class ClusterShape_ = Shape<_1, _1, _1>,
  class KernelSchedule = KernelTmaWarpSpecialized // or KernelTmaWarpSpecializedPersistent
>
struct MainloopSm90TmaWarpSpecialized {
  constexpr static int Stages = Stages_;
  using ClusterShape = ClusterShape_;
  using ArchTag = arch::Sm90;
  using Schedule = KernelSchedule;
};

CUTLASS 3: 集合 API

cutlass::gemm::collective::Gemm<>
cutlass::epilogue::collective::Epilogue<>

• 拥有线程块集群范围的操作和工作所有权。
• 调度策略允许在提供保护措施的同时，与任何内核调度自由组合。
• 以下代码片段展示了两种不同的主循环实现：一种是 n-buffer in smem, pipelined with Hopper GMMA and TMA；另一种是 n-buffer in smem, pipelined with Hopper GMMA and TMA, warp-specialized。

// n-buffer in smem, pipelined with Hopper GMMA and TMA
template<
  int Stages_,
  class ClusterShape_ = Shape<_1, _1, _1>,
  int PipelineAsyncMmaStages_ = 1
>
struct MainloopSm90TmaGmma {
  constexpr static int Stages = Stages_;
  using ClusterShape = ClusterShape_;
  constexpr static int PipelineAsyncMmaStages = PipelineAsyncMmaStages_;
  using ArchTag = arch::Sm90;
  using Schedule = KernelTma;
};

// n-buffer in smem, pipelined with Hopper GMMA and TMA, warp-specialized
template<
  int Stages_,
  class ClusterShape_ = Shape<_1, _1, _1>,
  class KernelSchedule = KernelTmaWarpSpecialized
>
struct MainloopSm90TmaGmmaWarpSpecialized {
  constexpr static int Stages = Stages_;
  using ClusterShape = ClusterShape_;
  using ArchTag = arch::Sm90;
  using Schedule = KernelSchedule;
};

CUTLASS 3: 集合构建器 (Collective Builders)

cutlass::gemm::collective::CollectiveBuilder<>

• 这是一个在主集合 API 之上的便捷接口。
• 元程序 (Meta-programs) 映射常见的 2.x 风格参数以生成 3.x 类型。
• 替代了 2.x 中的 DefaultXConfiguration 特化。

示例 1: "我只想要一个 Hopper 主循环"

using CollectiveOp = typename cutlass::gemm::collective::CollectiveBuilder<
  arch::Sm90, arch::OpClassTensorOp,
  half_t, LayoutA, 8,
  half_t, LayoutB, 8,
  float,
  Shapes<_128, _128, _64>, Shapes<_1, _2, _1>,
  gemm::collective::StageCountAuto,
  gemm::collective::KernelScheduleAuto
>::CollectiveOp;

示例 2: "我想要一个 Hopper 主循环，但使用持久化调度和 5 个阶段"

using CollectiveOp = typename cutlass::gemm::collective::CollectiveBuilder<
  arch::Sm90, arch::OpClassTensorOp,
  half_t, LayoutA, 8,
  half_t, LayoutB, 8,
  float,
  Shapes<_128, _128, _64>, Shapes<_1, _2, _1>,
  gemm::collective::StageCount<5>,
  gemm::KernelTmaWarpSpecializedPersistent
>::CollectiveOp;

实现峰值性能的秘诀

异步机器的攻击

• 使用 TMA 的全局内存访问:

  • 需要通过编程方式进行多播，以实现 L2 带宽效率。
  • 保存谓词和地址算术操作。
  • 发出 TMA 指令的线程只需发出 TMA 指令，以最小化延迟覆盖需求。

• MMA 操作:

  • 同时发出多个异步 MMA 指令，并与 TMA 屏障同步。
  • 发出 MMA 指令的线程只需发出 MMA，以达到峰值速率。

• 为峰值数学吞吐量隐藏数据移动延迟:

  • 通过共享内存中的缓冲区，对 gmem->smem 进行软件流水线化。
  • 高效同步异步操作，以达到峰值性能。

• 还必须隐藏软件流水线头部和尾部引起的空泡:

  • MMA 指令具有很强的缩放 A lot。
  • 持久化 GEMM 是不可避免的，并且还需要 CTA（协作线程数组）的重新配置。

• 仍然需要在 smem 中进行数据交换 (swizzle):

  • 现在由 MMA 规定并由 TMA 实现。

异步流水线 (Async. Pipelines)

管理到达等待 (Arrive Wait) 和事务屏障 (Transaction Barriers)

• Hopper H100 严重依赖深度异步软件流水线以实现峰值性能。
• 对于十几个循环缓冲阶段，直接通过操作多个生产者和消费者的屏障来管理异步流水线是繁琐的。
• CUTLASS 增加了对 Async Pipeline 类的支持，这些类提供了一个函数式抽象 API，利用底层硬件特性来实现同步。

PipelineTmaAsync 模板类接口示例:

template <int Stages, class ClusterShape>
class PipelineTmaAsync {
  // Acquire a stage in Smem before writing to it
  void producer_acquire(PipelineStage<Stages> state);

  // Commit a stage after writing to Smem (optional)
  void producer_commit(PipelineStage<Stages> state);

  // Wait for Commit before consuming a stage in Smem
  void consumer_wait(PipelineStage<Stages> state);

  // Notify end of consumption of Smem stage
  void consumer_release(PipelineStage<Stages> state);
};

异步 Warp 特化/专用内核 (Async. Warp Specialized Kernel)

共享多处理器 (Shared Multiprocessor) 的逻辑视图

该内核设计将 SM (Shared Multiprocessor) 内的 Warp 分为两类：
- MMA Warps: 负责执行矩阵乘加 (Matrix Multiply-Accumulate) 计算。
- DMA Warps: 负责通过张量内存加速器 (Tensor Memory Accelerator, TMA) 进行数据移动 (Direct Memory Access)。

数据加载流程 (Producer):

1. 发起异步拷贝: DMA Warps 发起 cp.async.bulk 指令，通过 TMA 从全局内存 (Global Memory) 异步加载数据。同时，它们会在共享内存 (Shared Memory) 上获取一个异步屏障 (Async. Barrier) 的所有权，为即将写入的数据块做准备。

   

2. 数据到达: TMA 完成从全局内存的数据读取，并将数据写入共享内存。同时，它会更新对应的屏障状态，以通知等待该数据的消费者。此操作支持将数据和屏障更新多播 (Multicast) 到另一个线程块。

   

计算流程 (Consumer):

1. 等待数据: MMA Warps 在共享内存的异步屏障上执行 Wait 操作，暂停执行，直到 DMA Warps 加载的数据准备就绪。

   

2. 执行计算: 一旦屏障被满足，MMA Warps 被唤醒。它们发起 wgmma.mma_async (Warp Group MMA Asynchronous) 指令，让张量核心 (Tensor Cores) 从共享内存中读取操作数并开始计算。

   

3. 释放屏障: 计算结果被写入寄存器内存 (Register Memory)。wgmma 操作完成后，MMA Warps 会释放 (Release) 异步屏障，表明它们已经消耗完共享内存中的数据。这块共享内存区域现在可以被 DMA Warps 用于加载下一批数据，从而实现流水线操作。

   

数据写回流程 (Epilogue):

1. 存储到共享内存: 计算结果从寄存器内存通过 stmatrix 指令存储回共享内存。

   

2. 异步写回全局内存: MMA Warps 发起 copy.bulk.async 指令，通过 TMA 将共享内存中的最终结果异步写回到全局内存。

   

完整流程总结:
下图展示了使用异步屏障的完整生产者-消费者模型。DMA Warps (生产者) 负责从全局内存加载数据到共享内存，并通过屏障通知 MMA Warps (消费者)。MMA Warps 等待数据就绪，然后使用张量核心进行计算，并将结果写回。整个过程是异步流水线化的，以最大化计算和数据传输的重叠。

紧凑的主循环表示 (Compact Mainloop Representation)

下图展示了 DMA 和 MMA 主循环的伪代码表示，体现了生产者-消费者模式。

• DMA 主循环 (生产者):

  • pipeline.producer_acquire: 锁定共享内存管道以进行写入。
  • copy(tma_load_bwith...): 使用 TMA 从全局内存加载数据块到共享内存的指定阶段。
  • ++smem_pipe_write: 推进写入指针到下一个阶段。

• MMA 主循环 (消费者):

  • pipeline.consumer_wait: 等待共享内存管道中的数据可用。
  • cute::gemm: 执行 GEMM 计算。
  • warpgroup_wait: 等待 GMMA (Group MMA) 操作完成。
  • pipeline.consumer_release: 释放共享内存管道，表示数据已消耗完毕。
  • ++smem_pipe_read, ++smem_pipe_release: 推进读取和释放指针。

持久化 Warp 专门化 (Persistent Warp Specialization)

这是一种隐藏非张量核心操作开销的方法。

• 背景: GEMM 内核的生命周期包括 Prologue (前序)、Tensor Core Operations (张量核心操作) 和 Epilogue (后序)。前序和后序部分包含非张量核心操作，通常会受延迟和/或带宽的限制。
• 传统方法: 过去架构中隐藏这些开销的典型方案是在每个 SM 上运行多个线程块。
• 挑战: 随着张量核心吞吐量的增加，软件流水线变得更深，这限制了在每个 SM 上运行多个线程块的能力。
• Hopper 解决方案: 引入一种新方法，通过持久化线程块来解决这个问题。这些持久化线程块跨多个输出图块 (output tiles) 和多个 Warp 组 (warp-groups, WGs) 发起集合性主循环 (collective Mainloops)，并为 Warp 组分配专门的功能。
• 专门化: 每个 WG 被分配一个特定的角色，例如数据生产者 (Data Producer) 或数据消费者 (Data Consumer)。

Warp 专门化和持久化内核

• Warp 专门化: 内核将线程组织到完全独立的执行路径中。

  • 发起 TMA 的线程只执行拷贝操作。
  • 发起 MMA 的线程只执行数学运算和后序操作。

• 持久化网格 (Persistent grids): 通过摊销内核启动和前序成本，并重叠 MMA 主循环和后序的执行，来提升效率。

下图展示了这种模式：一个生产者 (Producer, DMA Warps) 持续通过 TMA 准备数据，而多个消费者 (Consumers, MMA Warps) 并行地进行张量核心计算和后序处理，实现了高度流水化的执行。

CUTLASS Python

CUTLASS Python 接口

• 发布计划: 实验性版本将在 CUTLASS 3.1 中发布。

• 目标:

  • 简化 CUTLASS 内核的声明、生成和编译。
  • 在 Python 层面捕获常见的编译和运行时错误，以简化调试。
  • 将 CUTLASS 内核轻松集成到深度学习框架（例如 PyTorch）中。

目标：简化 CUTLASS 内核的声明、生成和编译

下图对比了 C++ 和 Python 接口的易用性。C++ 接口使用复杂的模板元编程，代码冗长且难以理解。而 Python 接口提供了一种更简洁、更具声明性的方式来构建和配置内核，例如可以轻松地更改 swizzling_functor 或添加 relu 激活函数。

目标：在 Python 中捕获常见的编译和运行时错误

下图展示了错误处理方面的改进。C++ 模板编译错误通常非常冗长且晦涩难懂，给调试带来巨大困难。而 Python 接口能在运行时抛出清晰、易于理解的异常，明确指出不支持的操作、数据类型或布局组合，极大地改善了开发体验。

与深度学习框架（如 PyTorch）的轻松集成

目标是实现与深度学习框架（如 PyTorch）的轻松集成。通过 CUTLASS Python 接口，可以定义一个计算计划（plan），然后生成相应的 PyTorch 扩展。

具体流程如下：

1. 定义计划：使用 cutlass.GroupedGemm 等类来定义所需的计算操作，例如指定元素类型为 torch.float16 和布局为行主序。
2. 代码生成：调用 cutlass.emit.pytorch 函数，传入构建好的计划，以生成 PyTorch 扩展所需的源代码文件。这会创建 Python (setup.py)、C++ (grouped_gemm.cpp) 和 CUDA (grouped_gemm.cu) 文件。
3. 编译安装：在命令行中运行 python setup.py install 来编译和安装生成的扩展。
4. 在 Python 中使用：安装完成后，可以直接在 Python 脚本中 import 该模块，并像调用普通 Python 函数一样运行高性能的 CUTLASS 内核。

结论

本报告对 CUTLASS 及其最新进展进行了总结。

CUTLASS 路线图
- CUTLASS 2.11 于 2022 年 11 月发布，将作为最后一个 CUTLASS 2.x 系列的版本。
- CUTLASS 3.0 于 2023 年 1 月发布。
- CUTLASS 3.1 预计于 2023 年 4 月可用。

CUTE
- CUTE 是一种思考张量 (Tensor) 和布局 (Layout) 的新方式。
- 它极大地简化了地址计算逻辑。
- CUTE 是 CUTLASS 3 的后端。

CUTLASS 3.0
- 为用户提供了灵活的抽象，用于组合自定义的内核和集合操作 (collectives)。
- 使用第四代张量核心 (Tensor Cores) 实现最优计算。
- 支持异步持久化 (async Persistent) 的生产者-消费者同步模型。
- 开源代码库：https://github.com/NVIDIA/cutlass

致谢 (Acknowledgements)

• CUTLASS GitHub 社区: 拥有 2.5K 星标，每月 1.6M 次克隆，70 多位贡献者和众多活跃用户。已集成到 PyTorch、Onflow、TVM、PaddlePaddle、AITemplate、PyG 等 300 多个 GitHub 项目中。
• 开发团队: 特别感谢 CUTLASS 开发者、CuTe 开发者、CUTLASS 产品管理团队以及其他所有贡献者。

参考文献

NVIDIA Hopper 架构与 CUDA
- "Inside the NVIDIA Hopper Architecture" (GTC 2022)
- "CUDA New Features and Beyond" (GTC 2022, GTC 2023)
- "Optimizing Applications for Hopper Architecture" (GTC 2023)
- "NVIDIA Hopper Architecture In-Depth" (博客文章)

PTX ISA
- 使用并行线程执行 (Parallel Thread Execution) 和指令集架构的编程指南 (CUDA 文档)。

CUTLASS
- https://github.com/NVIDIA/cutlass (开源软件，New BSD 许可证)
- CUTLASS Parallel For All 博客文章
- 往届 GTC CUTLASS 演讲：GTC'18, GTC'19, GTC'20, GTC'21, GTC'22, GTC'22