CUTLASS: Python API, Enhancements, and CUTLASS 3.0 Preview

Andrew Kerr, Cris Cecka | GTC Fall 2022

目录

致谢

CUTLASS GitHub 社区
- 2.1K 星标，250K 克隆/月，50 位贡献者，以及许多活跃用户。
- 来自 NVIDIA 外部的许多贡献和 PR。
- 已集成到 TVM、PyG 等项目中。

CUTLASS 工程师
- Andrew Kerr, Haicheng Wu, Cris Cecka, Pradeep Ramani, Aniket Shivam, Vijay Thakkar, Jin Wang, Honghao Lu, Ethan Yan, Shang Zhang, Jack Chen, Petrick Liu, Zhaodong Chen, Yujia Zhai, Jack Kosaian, Dustyn Blasig, Duane Merrill

CUTLASS 产品管理
- Matthew Nicely, Timothy Costa

贡献者
- Vedaanta Agarwalla, Roman Anders, Maximilien Breughe, Naila Farooqui, Manish Gupta, Markus Hohnerbach, Gautam Jain, Alan Kaatz, Wei Liu, Piotr Majcher, Dhiraj Reddy Nallapa, Kyrylo Perelygin, Paul Springer, Pawel Tabaszewski, Chinmay Talegaonkar, John Tran, Yang Xu, Scott Yokim

致谢
- Bing Xu, Leyuan Wang, Masahiro Masuda, Hao Lu, Olivier Giroux, Mostafa Hagog, Bryce Lelbach, Julien Demouth, Joel McCormack, Aartem Belewich, Peter Han, Timmy Liu, Yang Wang, Nich Zhao, Jack Yang, Vicki Wang, Junkai Wu, Ivan Yin, Aditya Alturi, Takuma Yamaguchi, Stephen Jones, Luke Durant, Harun Bayraktar

CUTLASS 简介

CUTLASS 是一个用于深度学习和高性能计算的 CUDA C++ 模板库。它为各种范围和规模的矩阵计算提供最优的 CUDA C++ 模板。

范围 描述
Device { GEMM, 卷积, 归约 } x { 所有数据类型 } x { SIMT, Tensor Cores } x { 所有架构 }
Kernel GEMM, Batched GEMM, 卷积, 归约, 融合输出操作, 融合输入操作
Collective 流水线矩阵乘法, Epilogue, 对张量的集体访问, 卷积矩阵访问
Atom Tensor Core 乘加操作, 对置换张量布局的高效访问
Thread 数值转换, 数组上的 <functional> 操作符, 快速数学算法
Architecture intrinsic 封装架构特定 PTX 指令的模板 (例如 mma, cp.async, ldmatrix, cvt)

CUTLASS 与 CUDA 编译器

  • CUTLASS 团队与 NVCC 团队和 NVResearch 团队紧密合作，以优化不同深度学习/高性能计算（DL/HPC）核心的性能。
  • 特别感谢 Duane Merrill, Malay Sanghi, Howard Chen, Rishkul Kulkarni, Hyun Dok Cho, Fei Peng, Balaji Atukuri, Eddie Gornish, Justin Holewinski, Kartik Haria, Christos Angelopoulos, Xingxing Pan, Jayashree Venkatesh, Xiaohua Zhang, Jerry Zheng, Shekhar Divekar, Cody Addison, Hari Sandanagobalane, Gautam Chakrabarti, Chu-Cheow Lim, Brian Deitrich, Dibyapran Sanyal, Vatsa Santhanam, 以及许多其他编译器团队成员。
  • 不同版本的 NVCC 为不同类型的核心提供了优化：
    • 11.3 - Tensor Core GEMMs
    • 11.4 - Tensor Core Implicit GEMMs (Fprop, Dgrad, Wgrad)
    • 11.5 - Sparse Tensor Core GEMMs
    • 11.6 - FP32 Tensor Core Emulation Kernels
    • 11.7 - SIMT kernels; Tensor Core kernels 的性能变化
    • 11.8 - CUTLASS kernels 的最佳性能，特别是 Grouped GEMM

议程

  • 路线图
  • CUTLASS Python
  • CUTLASS 2.x 增强功能
  • CUTLASS 3.0 预览

CUTLASS 路线图 (CUTLASS 2.x 和 3.x)

Page 6, CUTLASS 发展路线图，展示了从 CUTLASS 2.10 到 3.0 预览版的发布计划和主要特性。
Page 6, CUTLASS 发展路线图，展示了从 CUTLASS 2.10 到 3.0 预览版的发布计划和主要特性。

2022 年

  • 八月 (CUDA 11.7)

    • CUTLASS 2.10
    • CUTLASS Python
    • 增强功能:
      • Grouped GEMM 优化
      • Fused Softmax, Layernorm, 和 Multihead Attention
      • Grouped 和 Depthwise separable convolution

  • 十月 (CUDA 11.8)

    • CUTLASS 2.11

    • NVIDIA Hopper

      • 2x FP64 Tensor Cores
      • FP8 数值类型

    • Stream-K 动态调度算法

  • 十一月 (CUDA 12.0)

    • CUTLASS 3.0 预览版
    • CUDA 中的第四代 Tensor Cores
    • CuTe 编程模型

CUTLASS Python

本节介绍基于 Python 的 API、核函数融合、JIT 编译和缓存，以及示例与性能。

为什么需要 CUTLASS PYTHON？

目标：在 CUTLASS C++ API 和 Python 环境之间架起桥梁。

Page 8, 该图展示了 CUTLASS Python 如何作为桥梁，连接底层的 CUTLASS CUDA C++ 模板与上层的 Python 用户环境（如 PyTorch, TensorFlow, NumPy 等）。
Page 8, 该图展示了 CUTLASS Python 如何作为桥梁，连接底层的 CUTLASS CUDA C++ 模板与上层的 Python 用户环境（如 PyTorch, TensorFlow, NumPy 等）。

CUTLASS Python 解决了以下问题：

  • 如何构造模板实例？
  • 如何实现一个分发层，根据数据类型、张量布局、对齐方式和问题规模来选择最佳的核函数？
  • 如何将用户环境的参数传递给 CUTLASS 核函数？
  • 以及其他相关问题...

CUTLASS Python 架构

Page 9, CUTLASS Python 架构图，分为编译时和运行时两个部分。
Page 9, CUTLASS Python 架构图，分为编译时和运行时两个部分。

该架构分为编译时和运行时两个阶段：

编译时 (Compile Time):
- 客户端定义的配方 (Client-defined Recipe): 用户定义数据类型、布局、对齐、分块大小、交换函子、Epilogue 等。
- 操作描述 (Operation Description): 将用户配方转化为具体的操作描述。
- 代码发射器 (Code Emitter) 和 主机与设备编译器 (Host & Device Compiler): 基于操作描述生成并编译 CUDA C++ 代码，底层依赖于 CUTLASS C++ 库。

运行时 (Runtime):
- 应用程序提供的张量 (Application-provide Tensors): 如 torch.Tensor, np.ndarray, cp.ndarray 等。
- 参数包装器 (Argument Wrappers): 包装用户提供的张量。
- 运行时 (Runtime): 管理执行流程，包含一个 内存池管理器 (Memory Pool Manager)。
- 已编译构件管理器 (Compiled Artifact Manager): 管理 JIT 编译生成的核函数，并在运行时调用。

CUTLASS Python 示例

Page 10, 展示了一个完整的 CUTLASS Python 代码示例，从定义操作到执行。
Page 10, 展示了一个完整的 CUTLASS Python 代码示例，从定义操作到执行。

该示例代码展示了如何使用 CUTLASS Python API 来定义、编译和执行一个 GEMM 操作。
- 左侧代码块: 定义了操作的各个组成部分，包括数学指令（Math Instruction）、分块描述（Tile Description）、张量操作数（Tensor Operands）和 Epilogue 函子（Epilogue Functor）。
- 右上代码块 (Operation Description): 将上述组件组合成一个完整的操作描述。
- 右上代码块 (JIT Compilation): 使用 pycutlass.compiler.add_module 对操作进行即时编译。
- 右下代码块 (Launch):
- 提供用户张量（例如来自 PyTorch、NumPy）和问题规模。
- 准备参数。
- 运行操作并同步。

CUTLASS Python 与 CUTLASS C++ 性能对比

Conv2D 性能

Page 11, 对比了 CUTLASS Python 和 CUTLASS C++ 在 ResNet50 的 Conv2D 层上的性能。
Page 11, 对比了 CUTLASS Python 和 CUTLASS C++ 在 ResNet50 的 Conv2D 层上的性能。
  • 基准测试: RESNET50 中的各个卷积层。
  • 硬件: NVIDIA A100。
  • 配置: 混合精度训练 (F16<-F16*F16 + F32)。
  • 结果: CUTLASS Python 的执行时间与 CUTLASS C++ 非常接近，相对性能在 98% 到 101% 之间，表明 Python API 的开销极小。

GEMM 性能

Page 12, 对比了 CUTLASS Python 和 CUTLASS C++ 在 BERT Large 模型中不同 GEMM 操作上的性能。
- 基准测试: BERT Large 模型中的 GEMM 操作。
- 硬件: NVIDIA A100。
- 配置: 半精度推理 (F16<-F16F16 + F16)。
- 结果*: 对于 BERT 中的各种 GEMM 计算（如 Compute QKV, QK^T, AV, SelfOutput 等），CUTLASS Python 的性能几乎与 C++ 版本持平，达到了接近 100% 的相对性能。

CUTLASS Python 核函数融合

将 CUTLASS 中丰富的融合配方暴露给 Python

Page 13, 该表格详细列出了 CUTLASS Python 支持的核函数融合模式及其用例。
Page 13, 该表格详细列出了 CUTLASS Python 支持的核函数融合模式及其用例。

CUTLASS 支持在计算的不同阶段（Mainloop, Epilogue）进行多种模式的融合，以减少访存和核函数启动开销。

阶段 模式 用例
Mainloop 元素级计算 转换、缩放、掩码
A 或 B 的归约 GEMM + bias grad
广播 加载向量并跨通道广播
Epilogue 多个张量操作数 GEMM + GELU with Aux tensor; GEMM + RELU with output bitmask; GEMM + dRELU loading bitmask
多个向量操作数 Alpha 和 beta 缩放为向量
跨列广播 Bias add
跨列归约 GEMM + Bias Grad
Layer norm
跨行归约 GEMM + Softmax
元素级计算 算术、转换、激活函数
Composition 背靠背 GEMM
背靠背 CONV

CUTLASS 2.x 增强功能

本节介绍 CUTLASS 2.x 版本中的重要增强功能，包括：
- GEMM 融合与置换
- 卷积: 分组与深度可分离
- Stream-K
- NVIDIA Hopper 架构支持

GEMM 融合与置换

GEMM Permute (置换)

CUTLASS 实现和布局函数

  • 在深度学习应用中，GEMM 操作后通常会跟随布局转换（Layout transformations）。
# BERT self attention
context_layer = torch.matmul(attention_probs, value_layer)
context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
  • 在 CUTLASS 中，GEMM 的 epilogue 通过共享内存（Shared Memory）交换数据，然后使用高效的条带化访问模式（striped access patterns）协同访问全局内存（Global Memory）。
  • 在 GEMM epilogue 中，我们提供了 Layout plugin (include/cutlass/layout/permute.h)，用于对 CUTLASS 布局进行置换。
  • 该插件可以与现有的全局内存地址计算相结合，以在内存中重排数据，从而避免了额外的核函数调用。
Page 15, 该图展示了 GEMM epilogue 中如何通过一个布局函数（Layout function）在写入全局内存时直接进行数据重排。
Page 15, 该图展示了 GEMM epilogue 中如何通过一个布局函数（Layout function）在写入全局内存时直接进行数据重排。

GEMM Permute 性能

本页展示了在NVIDIA A100上进行的实验，旨在评估GEMM（通用矩阵乘法）与Permute（置换）操作融合的性能。

  • 实验设置:

    • GPU: NVIDIA A100
    • BatchedGEMM 尺寸: batch-count=1536, m=128, n=64, k=128
    • 输出张量形状: [128*12, 128, 64]
    • 重塑为: [128, 12, 128, 64]
    • 最终置换为: [128, 128, 12, 64]

  • 性能对比:

    • 融合 (Fused): BMM + Permute 总计 0.09789 ms
    • 非融合 (Unfused): BMM (0.09569 ms) + Permute (0.09159 ms) = 0.1873 ms
    • 加速比: 1.91x

如下图所示，融合操作将执行时间缩短了近一半，显著提升了性能。

GEMM Permute 性能图表
GEMM Permute 性能图表

融合层归一化 (Fused Layer Normalization)

层归一化（Layer Normalization）可以与前后两个GEMM操作进行融合，形成 GEMM₀ → Normalization → GEMM₁ 的流水线。

  • 层归一化公式:
    z̃ = (z - μ) / σ * γ + β
    • z̃: 输入/输出向量
    • μ: 向量均值
    • σ: 标准差
    • γ: 缩放向量
    • β: 偏置向量

归一化操作可以被分解并与GEMM层融合，过程如下：

  1. GEMM₀ 附带结尾融合 (epilogue fusion): 在CTA（Cooperative Thread Array）块内计算部分和，用于后续计算均值 μ 和方差 σ²。归约工作区的大小是 Z 的 1/CtaTile_M。
  2. 最终归约和标量计算: 进行轻量级的完全归约，计算出最终的均值 μ 和标准差的倒数 σ⁻¹。
  3. GEMM₁ 附带主循环融合 (mainloop fusion): 在主循环中进行元素级运算，计算最终归一化结果 z̃。
融合层归一化流程图
融合层归一化流程图

CUTLASS 全融合层归一化示例基准测试

该基准测试展示了在不同问题规模下，CUTLASS中完全融合的层归一化与非融合实现的性能对比。

  • 测试设置:
    • GEMM0: (768 x n0 x 768)
    • GEMM1: (3072 x n0 x 768)
    • n0 取值分别为 1024, 4096, 6144, 8192

如下图所示，融合（fused）版本的执行时间始终优于非融合（unfused）版本。非融合版本的相对执行时间比融合版本高出约13%-15%。

CUTLASS 全融合层归一化基准测试图表
CUTLASS 全融合层归一化基准测试图表

融合 Softmax (Fused Softmax)

Softmax操作可以与前序的GEMM操作进行融合，以减少数据移动和内核启动开销。

  • 带max技巧的Softmax公式:
    σ(z)ᵢ = e^(zᵢ - max) / Σₖ e^(zₖ - max)

    • σ(z)ᵢ: 第 i 个输出softmax向量
    • z: 输入向量
    • eᶻᵢ: 对输入向量的标准指数函数
    • max: 输入向量中的最大值

  • 融合过程:
    Softmax操作包含两次归约（求最大值、求和）和两次元素级运算（指数、缩放）。在融合实现中，来自GEMM的输出矩阵直接作为输入。线程内归约和warp内归约（通过shuffle指令）被融合在GEMM的epilogue（结尾操作）中，计算出max和sum值。这些值随后被用于元素级运算，得到最终的Softmax结果。

融合 Softmax 流程图
融合 Softmax 流程图

CUTLASS Softmax 带批量支持示例基准测试

此基准测试比较了CUTLASS中的融合Softmax与两遍式基线（Two-pass baseline）实现的性能。测试涵盖了16批次和1批次的场景，以及不同的GEMM形状 {m, n, k}。

结果显示，在所有测试用例中，融合Softmax（Fused Softmax）的性能均优于基线实现。基线版本的相对执行时间比融合版本高出20%至60%不等。

CUTLASS Softmax 基准测试图表
CUTLASS Softmax 基准测试图表

融合多头注意力 (Fused Multihead Attention, MHA)

多头注意力机制的核心计算可以被高效地融合。

  • 多头注意力计算:

    1. P = scalar * Q * K
    2. P = softmax(P)
    3. O = P * V

  • CUTLASS融合MHA示例（使用Grouped GEMM）:
    该问题可以看作是(BS x HN)个具有可变序列长度的MHA问题。

    1. Grouped GEMM0 附带结尾融合: 计算 Pᵢ = Qᵢ * Kᵢ。在结尾部分融合了用于Softmax的部分归约操作。
    2. Sum Max: 进行轻量级的完全归约，计算Softmax所需的最终 sum 和 max。
    3. Grouped GEMM1 附带主循环融合: 计算 Oᵢ = Pᵢ * Vᵢ。在主循环中融合了Softmax的元素级运算。
融合多头注意力流程图
融合多头注意力流程图

CUTLASS 融合多头注意力示例基准测试

该基准测试展示了CUTLASS实现的融合MHA相比于PyTorch的性能加速比。

  • 结论: 平均比PyTorch快约 6.25倍。

  • 测试场景:

    • CUTLASS v.s. PyTorch + Mask
    • CUTLASS v.s. PyTorch + No Mask

  • 测试范围: 最大序列长度从64到1024。

测试结果表明，在所有序列长度下，CUTLASS均展现出显著的性能优势，加速比在3.5倍到9.5倍以上。

CUTLASS 融合 MHA 基准测试图表
CUTLASS 融合 MHA 基准测试图表

更多细节请参考：https://github.com/NVIDIA/cutlass/tree/master/examples/41_multi_head_attention/fused_multihead_attention.cu

卷积: 分组与深度可分离

分组卷积 (Grouped Convolution)

分组卷积通过对滤波器进行分组来降低卷积层的计算复杂度。

  • 核心思想: 每个滤波器组仅与图像通道的一个子集进行卷积。
  • 关系式: filter.channels = image.channels / groups
  • 图示: 下图以一个 groups = 4 的例子展示了分组卷积的过程。输入张量的通道被分成4组，每组分别与对应的滤波器组进行卷积，最后将结果拼接起来。
分组卷积示例
分组卷积示例

CUTLASS 分组卷积实现

CUTLASS 使用隐式GEMM（Implicit GEMM）的方式来实现分组卷积。

  • 实现要点:

    • 在隐式GEMM的表示中引入了块稀疏结构。
    • C（通道）维度根据分组数量进行条带化划分。
    • 单个CTA（Cooperative Thread Array）内的工作负载根据分组数量进行划分。

  • 图示: 下图展示了激活（Activations）、滤波器组（Filter Groups）和层输出（Layer Output）在内存和计算单元中的映射关系。滤波器矩阵呈现出块对角线的稀疏结构。

CUTLASS 分组卷积实现示意图
CUTLASS 分组卷积实现示意图

CUTLASS 分组卷积实现: kSingleGroup 模式

cutlass::conv::kSingleGroup 模式是分组卷积的一种实现方式。

  • 工作方式: 每个CTA计算一个分组。例如，CTA 0 负责处理 group 0，CTA 2 负责处理 group 1。
  • 优势: 这种方式可以将主循环的迭代次数减少4倍（在groups=4的例子中）。
  • 图示: 下图详细展示了在CTA层级上，数据是如何被不同CTA处理的。
kSingleGroup 模式图示
kSingleGroup 模式图示

CUTLASS 分组卷积实现: kMultipleGroup 模式

cutlass::conv::kMultipleGroup 模式是分组卷积的另一种实现方式。

  • 适用场景: 当 filters_per_group < cta_tile_N 时，即每个分组的滤波器数量小于CTA块的N维度大小时。
  • 工作方式: 每个CTA可以计算多个分组。例如，CTA 0 同时负责处理 group 0 和 group 1。
  • 优势: 这种方式可以将主循环的迭代次数减少2倍（在groups=4的例子中）。
  • 图示: 下图展示了单个CTA如何处理多个分组的数据。
kMultipleGroup 模式图示
kMultipleGroup 模式图示

分组卷积代码片段 (CUTLASS 2.11)

以下代码片段展示了如何在CUTLASS 2.11中使用分组卷积。

  • 编译时 (C++ 模板实例化): 通过 DefaultConv2dGroupFprop 定义一个分组卷积核，并将 GroupMode 设置为 kSingleGroup。
  • 运行时 (隐式GEMM参数结构): 在定义问题尺寸时，通过 groups 参数指定分组数量，例如 groups = 2。
分组卷积代码示例
分组卷积代码示例

深度可分离卷积 (Depthwise Separable Convolution)

深度可分离卷积是一种特殊的分组卷积，其滤波器分组数等于通道数，且每个滤波器的通道数filter.channels = 1。

  • 前向传播: 每个输入通道都由一个独立的滤波器进行卷积。
  • 图示: 下图展示了一个深度可分离卷积的例子。一个3通道的激活张量（Activation Tensor）与一个包含3个单通道滤波器的滤波器张量（Filter Tensor）进行卷积，生成一个3通道的输出张量（Output Tensor）。
深度可分离卷积图示
深度可分离卷积图示

深度可分离卷积 (隐式GEMM版本)

深度可分离卷积的隐式GEMM实现具有以下特点：

  • 在每个平铺迭代中:
    • 激活矩阵（Activation matrix）的构造方式与常规的隐式GEMM卷积相同。
    • 用于计算的滤波器矩阵（Filter matrix）是类对角矩阵。
    • SMEM（共享内存）存储是紧凑的。只从GMEM（全局内存）加载和存储实际的滤波器元素。
    • 减少了LDS（本地数据共享，即共享内存）的使用量。SMEM中只加载需要使用的元素（图中的彩色部分）。

下图通过四次迭代展示了如何使用紧凑的滤波器表示来处理激活矩阵的不同部分，从而减少内存占用。

深度可分离卷积的隐式GEMM实现
深度可分离卷积的隐式GEMM实现

深度可分离卷积代码片段

在 CUTLASS 2.11 中，可通过特定的 Device-level DepthwiseFpropKernel Instance 模板实例化来启用深度可分离卷积的代码路径。以下代码片段对比了常规的 Conv2d 前向传播（fprop）核函数与深度可分离 Conv2d 前向传播核函数的实例化。

关键区别在于：

  • 操作类别 (OpClass)：深度可分离版本使用 cutlass::arch::OpClassSimt，这表示由于核函数受内存限制，它将使用 CUDA Cores 而不是 Tensor Cores。
  • Epilogue: 深度可分离版本使用专门的 cutlass::epilogue::thread::LinearCombinationDepthwise。
Page 31 - 深度可分离卷积代码片段
Page 31 - 深度可分离卷积代码片段

Stream-K

"经典" GEMM/CONV 并行化方法

传统的通用矩阵乘法（GEMM）和卷积（CONV）并行化方法是输出导向的，采用数据并行分解，即为每个输出瓦片（tile）分配一个协作线程数组（CTA）。

这种方法存在以下问题：
- 量化效率低下：CTA 以“波次”（waves）的形式在芯片上分派。最后一波可能只有部分被填充，导致 GPU 资源未被充分利用。
- 多种瓦片尺寸策略：需要选择启发式算法来确定瓦片大小，但可能仍无法与 GPU 的占用率完全匹配。

示例：如下图所示，9个 CTA 被调度到 4 个流式多处理器（SM）上，共分 3 个波次来生成九个 128x128 的元素瓦片。在第三个波次中，只有 SM₀ 在工作，而 SM₁, SM₂, SM₃ 处于空闲状态，导致利用率上限仅为 75%。

Page 33 - 经典并行化方法的资源利用率问题
Page 33 - 经典并行化方法的资源利用率问题

Stream-K 分解 (STK)

Stream-K (STK) 是一种新的分解方法。它不直接对输出矩阵进行瓦片化，而是考虑整个 GEMM 计算的聚合累加工作，即主循环的迭代总次数。

总迭代次数的量级为 O(m * n * k)，STK 将这些迭代视为一个线性的工作序列。

Page 34 - Stream-K 将计算分解为主循环迭代
Page 34 - Stream-K 将计算分解为主循环迭代

STK 的核心思想是：
- 假设有 p 个 SM，在每个 SM 上启动恒定数量的 CTA。
- 为每个 CTA 分配均等份额的聚合主循环迭代。
- 引入一个 “修正”（Fixup）阶段，用于累加来自每个 CTA 的“结转”部分和（carry-out partial-sums）。这在计算单个输出瓦片的迭代跨越多个 CTA 时是必需的。

示例：36 次迭代被调度到四个 128x128 的 CTA 上。修正阶段将完成那些其计算迭代跨越了多个 CTA 的瓦片。

Page 35 - Stream-K 将迭代均匀分配给 CTA
Page 35 - Stream-K 将迭代均匀分配给 CTA

下图中的蓝色竖线标示了需要进行“修正”操作的位置，即 CTA 之间工作负载的边界。

Page 36 - Stream-K 的修正阶段示意图
Page 36 - Stream-K 的修正阶段示意图

"Stream-K" GEMM/CONV 并行化

Stream-K 是一种以迭代为中心的并行化方法，具有以下优点：
- 工作负载均衡：每个 SM 接收均等份额的聚合主循环迭代，从而提高了 SM 的利用率。
- 低“修正”开销：跨 SM 聚合部分和所需的工作量和存储空间开销为 O(p)。
- 通用性：以迭代为中心的设计可以模拟“数据并行”和“split-k”等分解方法。

如下图所示，与经典方法不同，所有 SM 从始至终都保持活跃状态，避免了因波次调度而产生的空闲时间，从而实现了更高的硬件利用率。

Page 37 - Stream-K 的并行化模型与资源利用率
Page 37 - Stream-K 的并行化模型与资源利用率

添加 Stream-K 外层循环

在现有抽象上添加 Stream-K 只需要最小的额外结构。其实现包含一个 __global__ 核函数，其中有一个持久化的工作处理循环。

代码结构如下：
1. 动态线程块调度逻辑：在 while 循环中，每个 CTA 动态获取其要处理的瓦片ID（tile_id）和迭代范围（tile_first_itr, tile_stop_itr）。
2. 标准 GEMM 主循环：调用标准的基于瓦片的 GEMM 计算。
3. “修正”逻辑：
- 如果一个 CTA 完成了某个瓦片的一部分工作，它会存储部分和并发出信号。
- 如果另一个 CTA 开始处理同一个瓦片的剩余部分，它会等待前一个 CTA 完成，然后累加之前存储的部分和。

  1. Epilogue：生成最终的输出瓦片。
Page 38 - Stream-K 内核的伪代码结构
Page 38 - Stream-K 内核的伪代码结构

性能提升

通过在一个 NVIDIA A100 GPU 上对 4,096 个 GEMM 问题进行的抽样研究，对比了 cuBLAS（使用约20个核函数的集成）和 Stream-K（使用单个核函数）的性能。

结果：
- 性能：
- 平均加速比 1.1x
- 最大加速比 5.4x
- 最小加速比 0.6x

  • 一致性：Stream-K 提供了更强的性能一致性。
  • 库大小：需要的核函数特化更少，减小了库的体积。

下面的图表显示了 Tensor Core 利用率与计算强度（每字节操作数）的关系。与 cuBLAS 相比，Stream-K 在各种计算强度下都能实现更高且更稳定的 Tensor Core 利用率。

Page 39 - cuBLAS 与 Stream-K 性能对比图
Page 39 - cuBLAS 与 Stream-K 性能对比图

在 CUTLASS 中使用 Stream-K

在 CUTLASS 2.11 中，Stream-K 通过一个新的线程块光栅化函数在 CUTLASS GEMM 计算中启用，扩展了现有的接口。

使用方法：
1. 定义核函数：使用 cutlass::gemm::kernel::DefaultGemmUniversal 定义一个启用了 Stream-K 的 GemmKernel。
2. 启动核函数：
- 在参数结构中，将 GemmUniversalMode 设置为 kStreamK。
- 可以通过 streamk_blocks 等参数控制参与协作处理的 CTA 数量。
- 当使用默认值时，调度器会尝试进行负载均衡。

Page 40 - 在 CUTLASS 中使用 Stream-K 的代码示例
Page 40 - 在 CUTLASS 中使用 Stream-K 的代码示例

NVIDIA Hopper 架构支持

NVIDIA H100 GPU 引入了多项架构级加速。

全新更快的第四代 Tensor Core：

  • 混合精度浮点 Tensor Core 操作速度是 NVIDIA A100 的 2倍。
  • IEEE 双精度 Tensor Core 速度是 NVIDIA A100 的 2倍。

FP8 数据类型和模式：
- 支持 FP8 浮点数据类型。
- FP8 Tensor Core 的速度是 NVIDIA A100 上 FP16 Tensor Core 的 4倍。

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

Page 42 - NVIDIA Hopper 架构特性
Page 42 - NVIDIA Hopper 架构特性

这些新特性可通过 CUDA 11.8 工具包和 CUTLASS 2.11 访问。
- 架构性加速（混合精度、双精度）可通过 CUDA 11.8 工具包使用。
- 新增的数据类型（如 FP8）由 CUTLASS 2.11 提供支持。

Page 43 - Hopper 特性与软件支持的对应关系
Page 43 - Hopper 特性与软件支持的对应关系

双精度：F64 * F64 + F64

Hopper 架构支持 16-by-8-by-4 的双精度（F64）矩阵乘加（MMA）指令。该指令通过内联 PTX mma.sync.aligned 调用。

下图展示了该操作的数据布局和一个 PTX 代码示例，其中 mma.sync.aligned.m16n8k4.row.col.f64.f64.f64.f64 指令执行了 F64 的矩阵乘加。

Page 44 - Hopper 上的双精度 MMA 操作
Page 44 - Hopper 上的双精度 MMA 操作

NVIDIA Hopper 上 2倍速的双精度 Tensor Cores

在 CUTLASS 2.11 中，开发者可以通过调整 InstructionShape 参数，在 NVIDIA Ampere 和 NVIDIA Hopper 架构上实现最佳的双精度性能。

如下代码所示：
- NVIDIA Ampere (sm80)：InstructionShape 设置为 cutlass::gemm::GemmShape<8, 8, 4>。
- NVIDIA Hopper (sm90)：InstructionShape 设置为 cutlass::gemm::GemmShape<16, 8, 4>，以利用其新的 MMA 指令。

这展示了 CUTLASS API 的灵活性，允许为不同架构优化性能，而无需更改高层代码结构。

Page 45 - 在 CUTLASS 中为 Ampere 和 Hopper 配置双精度 GEMM
Page 45 - 在 CUTLASS 中为 Ampere 和 Hopper 配置双精度 GEMM

FP8 浮点数类型

Hopper 架构引入了新的 FP8 数据类型，主要有两种格式：E4M3 和 E5M2。

  • E4M3: 1个符号位 (s)，4个指数位 (e3 e2 e1 e0)，3个尾数位 (m2 m1 m0)。
  • E5M2: 1个符号位 (s)，5个指数位 (e4 e3 e2 e1 e0)，2个尾数位 (m1 m0)。

软件支持:
- CUDA 11.8: 启用了硬件加速的转换和打包操作，例如：
- 4 x FP8 ↔ 4 x F16
- 4 x F32 → 4 x FP8

  • CUTLASS 2.11: 实现了 "fast math" 内核，其计算流程为：FP8 输入 → F16 * F16 + F32 计算 → FP8 输出。

下图展示了 FP8 在 CUTLASS 中的数据流：数据从全局内存 (FP8) 加载到共享内存 (FP8)，然后转换到寄存器文件 (F16) 中，送入 Tensor Cores 进行 F16*F16+F32 的混合精度计算。计算结果 (F32) 暂存至共享内存 (SMEM)，再由 CUDA Cores 进行 Epilogue 操作，最终写回全局内存 (FP8)。

FP8 数据类型和计算流程图
FP8 数据类型和计算流程图

CUTLASS 3.0 预览

接下来的内容将预览 CUTLASS 3.0 的主要特性，包括 CuTe 布局函数、Tensor Core 编程模型和 CUTLASS 3.0 API。

CuTe 与分层布局

分层布局 (Hierarchical Layouts)

分层布局通过坐标 (Coordinates) 和索引 (Indices) 来描述数据结构。

以一个 4x4 矩阵为例，其布局可以定义为：
- Shape: (4, (2, 2))
- Stride: (2, (1, 8))

数据在不同抽象层级下的映射关系如下：
1. 逻辑一维坐标 A(I): 将矩阵元素视为一维数组。
2. 坐标映射 (Coordinate Mapping): 将一维坐标 I 映射到逻辑 n 维坐标 (i, j)。
3. 坐标映射 (Coordinate Mapping): 进一步将 n 维坐标 (i, j) 映射到分层的 h 维坐标 (i, (j1, j2))。
4. 索引映射 (Index Mapping): 将分层坐标 (i, (j1, j2)) 映射到最终的线性一维存储索引 k。

核心思想:
- Shape 定义了不同维度坐标之间的映射关系。
- Stride 定义了从高维坐标到一维线性内存索引的映射关系。

分层布局的坐标和索引映射示意图
分层布局的坐标和索引映射示意图

布局示例

本页展示了一个更复杂布局的例子及其操作。

通过 make_layout, blocked_product 和 make_tensor 等函数可以构建复杂的数据张量布局。

布局示例代码与可视化
布局示例代码与可视化

  • Shape 与 Size:

    • size(A) = 64
    • size<0>(A) = 8
    • size<1>(A) = 8

  • 逻辑坐标:

    • 可以使用1D、2D或更高维度的逻辑坐标来访问同一个线性存储位置。例如，线性索引为 49 的元素可以通过 A(37)、A(5,4)、A(1,2),(0,2) 或 A((1,0,1)),(0,(0,1))) 等多种方式访问。

  • 切片 (Slice):

    • 可以沿着逻辑子边界进行切片操作，例如 A(_,2) 或 A(_,(_,2))。

布局与张量：为何选择 CuTe？

CuTe (CUDA Templates for Tensors) 的设计动机是为了简化和统一数据布局的表示。

  • 布局的代数形式化: CuTe 对布局 (Layout) 进行了形式化的代数定义，支持一系列操作，如：

    • composition (复合)
    • right_inverse / left_inverse (左右逆)
    • complement (补)
    • "product" (积)
    • "divide" (除)
    • "tiling" (分块)
    • "partitioning" (分区)

  • 统一数据与线程的布局: CuTe 使用统一的布局概念来描述数据和线程的组织方式。

对比:
- 传统 CUTLASS: 需要为各种内存布局定义大量特定的类型，如 ColumnMajor, RowMajorInterleaved, VoltaTensorOpMultiplicandCongruous 等，种类繁多且复杂。
- CuTe: 将所有这些复杂的布局统一抽象为 Layout<Shape, Stride> 的形式，极大地简化了编程模型。

CuTe 布局与传统 CUTLASS 布局对比
CuTe 布局与传统 CUTLASS 布局对比

Tensor Core 编程模型

MMA_Op 和 MMA_Traits

这是 Tensor Core 编程模型的底层抽象，涉及 PTX (Parallel Thread eXecution) 指令和其元信息。

  • MMA_Op (Raw PTX): 代表原始的 PTX MMA (Matrix Multiply-Accumulate) 指令。图中示例为 SM90_16x8x4_F64F64F64_TN，它封装了底层的 mma.sync.aligned 汇编指令。
  • MMA_Traits (PTX meta-info): 提供了关于 MMA 指令的元信息，是一个模板结构体。它定义了操作数的数据类型、形状 (Shape_MNK)，参与计算的线程数 (32 threads, one warp)，以及操作数 A、B、C 在寄存器中的布局 (ALayout, BLayout, CLayout)。
MMA_Op 和 MMA_Traits 的代码示例
MMA_Op 和 MMA_Traits 的代码示例

MMA_Atom

MMA_Atom 是基于 MMA_Op 和 MMA_Traits 构建的基本计算单元。

  • 作用: 它将底层的 PTX 指令和其元信息结合起来，提供了经过检查的调用接口 (checked call interfaces) 和片段生成 (fragment generation) 功能。
  • 基本构建块: MMA_Atom 是构建更复杂矩阵运算的基础。
  • 下图右侧可视化了单个 MMA_Atom 操作（例如 16x8 矩阵乘法）中，32个线程 (T0-T31) 如何分布以及它们各自持有的数据片段 (V0, V1, ...)。
MMA_Atom 的构建与可视化
MMA_Atom 的构建与可视化

Tiled_MMA：构建更大规模的操作

通过组合 MMA_Atom，可以构建更大规模的矩阵运算，这就是 Tiled_MMA 的作用。

  • 构建: 使用 make_tiled_mma 函数，传入一个 MMA_Atom 和一个布局 (Layout)，即可创建一个由 MMA_Atom 平铺而成的更大操作。例如，Layout<Shape<_2,_2>>() 表示创建一个 2x2 的 MMA_Atom 网格。
  • 功能: Tiled_MMA 负责管理 MMA_Atom 的布局，并提供分区工具 (partition utilities)。
  • 下图右侧展示了一个由多个 MMA_Atom 块平铺而成的更大矩阵的可视化表示。
Tiled_MMA 的构建与可视化
Tiled_MMA 的构建与可视化

Tiled_MMA 的线程切片 (Thread slice)

Tiled_MMA 定义了整个 warp 级别或线程块级别的操作，而每个单独的线程只负责其中的一部分。ThrMMA 代表了单个线程的视角。

  • 获取线程视图: 通过 tiled_mma.get_slice(threadIdx.x)，可以获得当前线程负责处理的 Tiled_MMA 的一部分，称为 ThrMMA。
  • 功能: ThrMMA 继承了分区和切片的功能，使得每个线程可以独立地管理自己的数据片段。
  • 编程流程:
    1. 定义全局张量 A, B, C。
    2. 获取当前线程的 ThrMMA 切片。
    3. 将 ThrMMA 分区，得到该线程负责的 A, B, C 的逻辑片段 thr_A, thr_B, thr_C。
    4. 在寄存器中为这些片段创建张量 rA, rB, rC。
    5. 执行数据拷贝（如从共享内存到寄存器）。
    6. 在 k 维度上循环，调用 mma.call 执行矩阵乘法累加。
    7. 将计算结果从寄存器写回。

下图展示了完整的 GEMM (General Matrix Multiply) 内核的线程级代码逻辑。

线程级 GEMM 代码示例
线程级 GEMM 代码示例

CUTE 与拷贝 (Copy) 操作

CUTE 的分层抽象设计不仅适用于计算操作（如 MMA），也同样适用于数据拷贝操作。

  • 拷贝操作的层次结构: 与 MMA 类似，拷贝操作也有一套从底层到高层的抽象：

    • Copy_Op (原始 PTX 指令) 和 Copy_Traits (元信息)
    • Copy_Atom (基本拷贝单元)
    • TiledCopy (由 Copy_Atom 平铺而成的大规模拷贝操作)
    • ThrCopy (单个线程的拷贝任务切片)

  • 代码示例: 代码展示了如何创建 Copy_Atom 和 TiledCopy，然后获取线程的 ThrCopy 切片，并最终调用 copy 函数执行数据搬运。

  • 这种统一的编程模型简化了为不同硬件和数据布局编写高效内核的复杂性。
CUTE 拷贝操作的层次结构和代码示例
CUTE 拷贝操作的层次结构和代码示例

CUTLASS 3.0 API

CUTLASS 3.0: 设备、内核、主循环 API

CUDA C++ 深度学习与高性能计算模板库

CUTLASS 3.0 计算层级结构:

层级 描述
Device (设备) 通用的、内核无关的主机接口，用于参数构造和内核启动。
Kernel (内核) CTA (Cooperative Thread Array) 的入口点，这些 CTA 可能会也可能不会组织成一个集群。用于融合背靠背 GEMM、epilogue 等的组合点。
Collective (集体) 协同工作的线程数量。流水线矩阵乘法、epilogue、加速同步、线程块集群。主循环融合、epilogue 偏置融合等的组合点。
TiledMMA / TiledCopy 从集体中复制或数学原子的布局。
Atom (原子) 由一个或多个线程管理的最小操作（数学/复制）。FFMA、LDS、mma.sync、LDSM、3xTF32、复杂 MMA 等。组合的 PTX 指令和 PTX 元信息。

  • 定制化: 可以在不破坏与其他层可组合性的前提下，定制层次结构中的任何层。

    • 极端关注开发者生产力。

  • 静态检查: 在每一层进行静态检查，以确保布局兼容性。

    • 如果能编译，它就是正确的——否则会提供可操作的断言信息。

  • 通过精心布局进行优化:

    • CUTLASS 3.0 将为常见用例提供默认布局。
Page 62
Page 62

CUTLASS 3.0: 2.x 设备兼容性 API

  • 与 CUTLASS 2.x 向后兼容的模板接口。

  • 默认的 GEMM 配置将基本的 GEMM 内核配置适配到 3.0 API。

    • 同时提供 3.0 内核配方 (kernel recipes) 的示例。

  • 对于简单的用例，提供更熟悉的开箱即用 (OOTB) 内核可用性。

    • 仅遵循数据类型和布局。
    • 指令形状 (Instruction shape)、扭曲形状 (warp shape)、切片形状 (tile shape)、操作符类别 (operator class) 已被弃用且未使用。

  • CUTLASS 3.0 还支持一个定制的 GemmUniversal 接口。

Page 63
Page 63

CUTLASS 3.0: 主循环 API

cutlass::gemm::collective::Gemm<...>
cutlass::epilogue::collective::Epilgoue<...>

  • Collective (集体):

    • 网格 (grid) 中可以利用硬件特性进行加速通信和同步的最大线程数量。
    • CTA/CGA 范围内的操作数和工作所有权。
    • 主循环构造，如流式K (stream K)、主循环融合等。
    • 通过分发策略 (dispatch policies) 进行选择。
    • 易于编写自定义主循环并与内核层组合。

  • 下方是两个主循环的代码示例：

    1. MainloopSm70TwoStage: 2 阶段流水线，其中 1 阶段在共享内存 (smem) 中，1 阶段在寄存器文件内存 (rmem) 中。
    2. MainloopSm80Rldgs: n 缓冲流水线在共享内存中，通过 LDGSTS 异步访问。
Page 64
Page 64

集体主循环设置 (Collective Mainloop Setup)

Ampere Tensor Cores + LDGSTS + LDSW mainloop

  • 使用 GmemTiledCopy 划分 gmem 张量

    • LDGSTS 全局内存原子，带有矢量化信息。

  • 划分 tiled MMA 并分配线程片段

  • 使用 SmemTiledCopy 划分 smem 张量

    • LDSM 从源和目标进行原子布局静态检查。
    • 重新平铺 (retile) 目标 MMA rmem 张量。

  • 发布 prologue gmem 读取

  • 预取第 0 个 K_BLOCK 寄存器
Page 65
Page 65

集体主循环流水线 (Collective Mainloop Pipeline)

Ampere Tensor Cores + LDGSTS + LDSW mainloop

  • 流水线化的主循环体

    • N 个阶段在 smem 中异步获取，使用 LDGSTS。
    • 使用 LDSM 进行矢量化的、无 bank 冲突的 smem 加载。

  • 没有迭代器，只有 cute::Tensor

    • 除了需要为 tile 启用 GETT/streamK 风格的迭代策略。

  • 与 gemm/threadblock/mma_multistage.h 对比

Page 66
Page 66

结论

CUTLASS

  • CUTLASS 路线图

    • CUTLASS 2.10 已于 2022 年 8 月发布
    • CUTLASS 2.11 已于 2022 年 10 月推出
    • CUTLASS 3.0 预览版将于 2022 年 11 月推出

  • CUTLASS Python

    • 定义了用于构建 CUTLASS 内核的 Python API 和编程模型
    • 可与 Numpy 互操作

  • CUTLASS 2.x 增强功能

    • 融合的 Softmax、Layernorm 和多头注意力内核提供了显著的加速
    • 分组和深度可分离卷积降低了某些卷积工作负载的时间复杂度
    • Stream-K 动态调度算法实现了负载均衡并减少了内核选择的负担
    • NVIDIA Hopper 架构的启用加速了 FP64 并支持 FP8

  • CUTLASS 3.0 预览

    • 针对张量核心 (Tensor Cores) 的新编程模型
    • 基于 CuTe 布局代数和元编程抽象
    • 在 H100 中使用第 4 代张量核心实现最优计算
    • 将于 2022 年 11 月发布，与 CUDA 12.0 同步

https://github.com/NVIDIA/cutlass

Page 68
Page 68
Page 69
Page 69