CUTLASS: A Performant, Flexible, and Portable Way to Target Hopper Tensor Cores

Vijay Thakkar, Jack Kosaian
NVIDIA GTC 2024 | 2024/03/19

目录

1. CUTLASS 简介
2. 自 GTC'23 以来的新特性
3. CUTLASS 3 核心概念
4. CUTLASS 3 中的卷积
5. 面向大语言模型 (LLMs) 的特性
6. Epilogue Visitor Tree (EVT
7. 教程：使用 CUTLASS 编写自定义 CUDA 核函数
8. 结论与路线图

---

CUTLASS 简介

CUTLASS 是一个用于深度学习和高性能计算的 CUDA C++ 模板库。它旨在将张量计算在其所有范围和尺度上分解为其“移动部件”。

CUTLASS 的核心概念是分层抽象，每一层都封装了特定的功能：
- Device (设备层)：通用矩阵乘（GEMM）、卷积、归约（Reduction）、BLAS3 等，支持所有数据类型、SIMT 和张量核心（Tensor Cores），并兼容所有架构。
- Kernel (核函数层)：GEMM、批量 GEMM、卷积、归约、融合输出操作、融合输入操作。
- Collective (集合操作层)：CUTLASS 的时序微内核（temporal micro-kernels），用于编排空间微内核（spatial micro-kernels）的异步生产者/消费者流水线。
- Atom (原子操作层)：CuTe 空间微内核（Tiled MMA / Copy）。
- Thread (线程层)：数值转换、<functional> 风格的数组操作、complex<T>、快速数学算法。
- Architecture intrinsic (架构内置函数层)：封装架构特定的PTX指令的模板（例如 mma.cp.async, ldmatrix, cvt）。

开源社区与版本信息：
- 开源地址: https://github.com/NVIDIA/cutlass
- 社区规模: 4.4k 星标，每月 2.5M 次克隆，100+ 贡献者，以及大量活跃用户。
- 最新版本: CUTLASS 3.5
- 文档: https://github.com/NVIDIA/cutlass#documentation
- 过往 GTC 演讲: GTC'18, GTC'19, GTC'20, GTC'21, GTC'22, GTC'23

---

自 GTC'23 以来的新特性

自 GTC'23 以来，CUTLASS 引入了多项新功能，主要集中在对卷积、大语言模型（LLM）的支持以及通过访问者树（Visitor Tree）实现的 Epilogue 融合。

• CUTLASS 3 中的卷积：将卷积（CONV）实现为通用张量-张量收缩（GETT）的一种特例，而通用矩阵乘（GEMM）是 GETT 的一个子集。

• 针对 LLM 的特性：

  • 混合输入 GEMM：支持不同精度的输入，例如 FP16 和 INT4 的矩阵乘法。
  • 分组 GEMM (Grouped GEMM)：在单个核函数中执行多个独立的 GEMM 操作（A0xB0, A1xB1, A2xB2, ...）。

• 通过访问者树实现的 Epilogue 融合：提供了一种灵活的方式来融合自定义操作。例如，可以将累加、偏置（bias）和缩放等操作融合到计算的末尾，如 R = ((alpha * accum) + bias) + (beta * c)。

• 针对 Hopper 和 Ada 架构的 FP8 GEMM 和卷积：支持 E5M2 和 E4M3 两种 FP8 格式。
• 用于可组合负载均衡的 Tile 调度器：引入了 cutlass::gemm::StreamKScheduler 等新的调度策略，以实现更灵活的负载均衡。
• 改进的 Python 接口支持和 PyPI wheel 的添加：用户现在可以通过 pip install nvidia-cutlass 方便地安装和使用。
• GitHub 上的更多更新：
  • 扩展了编译器支持。
  • 支持窄对齐（Narrow-alignment）的 GEMM。
  • 改进了文档。

---

CUTLASS 3 核心概念

CUTLASS 3 概念层次结构

CUTLASS 3 对 GEMM 和卷积采用了统一的层次化抽象。

• 原子层 (Atom layer)：架构指令和相关的元信息。这是参与架构加速的指定数学/拷贝操作所必须的最小线程和值集合。
• 切片 MMA/拷贝层 (Tiled MMA/Copy)：空间微内核层。描述了一个数学/拷贝操作的完整空间切分（spatial tiling）。
• 集合层 (Collective layer)：时序微内核层。描述了计算一个输出切片的完整时序切分（temporal tiling）的数学/拷贝操作。主循环（Mainloops）负责编排拷贝/数学微内核，并进行架构特定的同步。
• 核函数层 (Kernel layer)：围绕集合操作的最外层循环。概念上是线程块/集群（threadblock/clusters）在网格（grid）中的集合。负责跨线程的负载均衡、线程调度（thread marshalling）、网格规划和参数构造。
• 设备层 (Device layer)：主机端的设置和接口。

CUTLASS 3 API 入口点

CUTLASS 3 在每个抽象层都提供了单一的 API 入口点。

• 空间微内核 (Spatial microkernels): cute::Tiled{Mma|Copy}<>

  • 在广泛的 GPU 架构上提供强大的表示能力。

• 时序微内核 (Temporal Microkernels): collective::Collective<{Mma|Conv|Epilogue}<>

  • 通过策略进行分发，这些策略定义了可以组合的核函数调度集合。

• 核函数层 (Kernel layer): kernel::{Gemm|Conv}Universal<>

  • 将 GEMM 视为一个通用的主循环和一个集合性 epilogue 的组合。
  • 每个新的核函数调度都是一个带有调度标签的特化。

• 设备层 (Device layer): device::{Gemm|Conv}UniversalAdapter<>

  • 可以与 2.x 或 3.x API 的核函数一起使用。
  • 提供一个到核函数类型的无状态句柄。

• 静态断言 (Static asserts)：在各处使用，以防止无效的组合或不正确的布局。

---

CUTLASS 3 中的卷积

卷积集合操作 (Convolution Collectives)

这是 3.x API 中为卷积新增的主要部分。

• API 与 GEMM 集合操作类似。

  • 通过主循环策略（mainloop policies）进行分发。
  • 计算单个输出切片。
  • 通过核函数策略（kernel policies）与核函数调度进行组合。

• 几乎所有卷积相关的特定更改都限制在主循环（mainloop）中。

  • 实现前向传播（fprop）、数据梯度（dgrad）和权重梯度（wgrad）。
  • 映射到 cute strides。
  • 构建张量内存访问（TMA）操作。

• Epilogue 集合操作可与 GEMM 和卷积组合使用。

  • 可与现有的自定义 epilogue 直接组合。

卷积提供熟悉的构建器 API

CUTLASS 3.x 提供了一个名为 cutlass::conv::collective::CollectiveBuilder<> 的构建器 API，以简化卷积核的创建。

用户只需提供高层级的描述，例如：

• "我只需要一个用于 NWC 布局张量的 Hopper Fprop 集合 (1-D Fprop)"：用户指定架构 (Sm90)、操作类型 (KFprop)、数据类型 (half_t)、布局 (TensorNWC) 等。
• "我想要一个用于 NDHWC 布局张量的 Hopper Dgrad 集合 (3-D Dgrad)"：同样，用户指定架构 (Sm90)、操作类型 (KDgrad)、布局 (TensorNDHWC) 等高级参数。

构建器 API 负责处理底层的复杂配置。

构建器完成繁重工作

CollectiveBuilder 自动将用户的高级描述映射到最优的主循环配置。

这个过程包括自动确定：
- 主循环类型
- 阶段计数 (stage count)
- GMMA 指令
- TMA 指令
- 共享内存布局
- 核函数调度

如上图所示，一个简洁的构建器调用（上部分代码）会被展开成一个非常复杂和高度特化的模板实例化（下部分代码），从而为特定硬件和问题配置生成最优代码。

CUTLASS 3.5 卷积 API

CUTLASS 3.5 提供了熟悉的核函数层和设备层 API 来构建完整的卷积操作。

构建过程分为几个步骤：
1. 构建 Epilogue 类型：使用 epilogue::collective::CollectiveBuilder 定义计算结束后的操作。
2. 构建 Mainloop 类型：使用 conv::collective::CollectiveBuilder 定义卷积的核心计算循环。
3. 在核函数层组合：使用 conv::kernel::ConvUniversal 将 mainloop 和 epilogue 组合成一个完整的核函数。问题形状类型会由 mainloop 推断出来。
4. 获取设备层句柄：使用 conv::device::ConvUniversalAdapter<ConvKernel> 创建一个可以从主机代码调用的句柄。

秩无关的卷积问题形状 (Rank Agnostic Conv Problem Shape)

CUTLASS 3.x 的卷积使用了一个面向用户的、与秩无关（rank agnostic）的参数。

• 背景：GEMM 内核允许用户提供任意的 cute::Shape 作为问题形状类型，非常灵活。但卷积通常不需要如此高的通用性，其参数化主要限于空间维度和算法。
• 目标：实现一个与秩无关的卷积，不需要“命名”模式。Im2Col 变换应被视为实现细节。
• 解决方案：引入一个新的问题形状类型 ConvProblemShape 用于 N 维卷积问题。这个结构是秩无关和实现选择无关的，它由卷积主循环根据空间维度和操作类型自动推断。用户只需提供不对称的填充（padding）、膨胀（dilation）和遍历步长（traversal strides）。

映射到面向核函数的步长 (Kernel Facing Strides)

如何将卷积的布局标签（layout tags）映射到 CuTe 的步长（strides）？

• CUTLASS 2.x GEMM：使用 cutlass::layout::RowMajor(ldA) 这样的布局标签，其中 ldA 是 leading dimension。
• CUTLASS 2.x CONV：使用 cutlass::layout::TensorNHWC(ldN, ldH, ldW) 这样的布局标签。

这些是布局标签，而不是直接的步长定义。CUTLASS 3.x 需要一种更通用的方式来处理步长。

分层布局的力量 (Power Of Hierarchical Layouts)

在 CUTLASS 3.x 中，GEMM 被视为张量收缩（GETTs）的一种伪装。布局和步长是分层组合的。

• 步长组件：步长被分解为多个模式（mode），例如 RowModeStrides（行模式）、ColModeStrides（列模式）、RedModeStrides（归约模式）和 BatModeStrides（批处理模式）。
• 组合：完整的张量步长（如 StrideA, StrideB）是通过将这些组件模式按照约定顺序组合而成的。
• 张量收缩：如上图右侧公式 C(mp)(l)(n) = A(mp)(kr)(n) * B(kr)(l)(n) 所示，不同的模式（批处理、行、列、归约）在输入和输出张量之间进行映射。这种分层方法使得表示复杂的张量操作变得非常灵活。

你所需要的一切就是 GETT (GETTs Are All You Need)

卷积只是通用张量收缩（GETT）的一种特殊情况。

• Im2Col 变换：Im2Col 算法将激活张量（activation tensor）的形状和步长进行变换。

• 通用性：这种方法可以推广到任何滤波器（filter）、膨胀（dilation）、遍历步长（traversal stride）和填充（padding）。

  • 滤波器形状成为激活张量的 TRS 形状（扩展了域）。
  • 遍历步长分解到 ZPQ 步长中。
  • 膨胀分解到 TRS 步长中。

• 变换之后：

  • M 多模式在逻辑上与输出 M 模式一致。
  • K 多模式在逻辑上与滤波器 K 模式一致。

• 结论：通过 Im2Col 变换，我们可以将任何卷积问题转化为一个 GETT 问题。如上图所示，激活张量（NHW x C）通过 Im2Col 变换成一个矩阵，然后与滤波器矩阵进行标准的矩阵乘法，得到输出矩阵。

Hopper TMA Im2Col

用于卷积的快速简单数据移动

• 带有 Im2Col 变换的 SIMT 内核难以实现和优化。

  • 在新扩展的收缩模式中复制输入激活张量。
  • 为填充（padding）考虑光环加载（halo loads）和越界（OOB）值。
  • 沿收缩模式对 OOB 读取进行复杂的预测。

• Hopper TMA Im2Col 使这一切变得简单。

• Im2Col TMA 执行以下操作：
  • 在张量步长（tensor strides）上执行 Im2Col 变换。
  • 预测 OOB 读取以处理填充。
  • 减少指令分发开销和寄存器压力。

CUTLASS 3.5: 卷积支持

适用于使用 TMA Im2Col 和 WGMMA 的 Hopper 架构

• CUTLASS 3.5 版本在 3.x API 中原生包含对卷积的支持。

  • Beta 版本，欢迎提供反馈！

• 以与秩无关（rank-agnostic）的方式支持 1、2 和 3 个空间维度。

• 支持 Fprop、Dgrad 和 Wgrad 算法。
• 支持非对称填充（asymmetric padding）、膨胀（dilations）和遍历步长（traversal strides）。
• 初步的性能分析器（profiler）集成。
• 未来路线图：
  • 性能提升
  • 带步长的 dGrad (Strided dGrad)
  • 将性能分析器覆盖范围扩展到一维卷积 (1D conv)

---

面向大语言模型 (LLMs) 的特性

训练和部署 LLMs 需要在 GEMM 之上进行优化

大语言模型的训练和部署需要基于通用矩阵乘法（GEMM）进行额外优化。主要体现在两个方面：专家混合（Mixture of experts）和权重化（Weight quantization）。

• 专家混合 (Mixture of experts): 该模型通过门控网络（Gating network）将输入路由到不同的专家网络，然后合并结果。这需要 Grouped GEMM，已在 CUTLASS 3.4 中为 Hopper 架构提供支持。
• 权重化 (Weight quantization): 该技术使用低精度权重（如 INT4）与高精度激活（如 FP16）进行计算，以提升效率。这需要 Mixed-input GEMM，已在 CUTLASS 3.3 中为 Hopper 和 Ampere 架构提供支持。

专家混合 (MoE) 实现计算高效的 LLMs

MoE 模型的计算需求如下：

• 期望操作:

  • 并发运行多个 GEMM。
  • 每个 GEMM 的大小可能不同。
  • 并非所有 GEMM 都会为每个输入进行计算。

• 批处理 GEMM (Batched GEMM) 不适用: 因为它要求所有 GEMM 具有相同的大小。

• 分组 GEMM (Grouped GEMM) 适用:
  • 在单个内核中计算多个 GEMM。
  • 每个 GEMM 可以有不同的大小和步长。
  • GEMM 的数量和大小无需在编译时确定。

分组 GEMM (Grouped GEMM)

传统的串行执行方式在处理多个不同大小的 GEMM 时，会导致部分计算单元（CTA）空闲，GPU 利用率不高。

分组 GEMM 通过一次内核启动（single kernel launch）来执行一组 GEMM。它使用持久化的 CTA（persistent CTAs），这些 CTA 可以在多个小的计算问题之间动态分配工作负载，从而有效利用计算资源。

• 优点:
  • 通过多个小问题使 GPU 达到饱和。
  • 只需一次内核启动。

在 Hopper 架构上使用 Grouped GEMM

以下代码片段展示了如何在 CUTLASS 中为 Hopper 架构配置 Grouped GEMM。通过使用指向布局的指针（如 LayoutC*, LayoutA*, LayoutB*），可以在运行时为组内的每个 GEMM 指定不同的大小和数据布局。这是实现 Grouped GEMM 灵活性的关键。

Hopper架构下的分组GEMM增强：可修改的TMA描述符

Hopper的张量内存加速器（Tensor Memory Accelerator, TMA）需要使用描述符来执行拷贝操作。对于不同的GEMM（通用矩阵乘法）类型，其工作方式有所不同：

• 单个GEMM：在整个内核的生命周期中，每个操作数（例如A、B）只需要一个描述符。
• 分组GEMM：需要为组中的每个GEMM提供不同的描述符字段。

工作流程对比如下：
- 单个GEMM：描述符在主机端构建，并通过启动参数复制到设备端。
- 分组GEMM：
- CUDA 12.3中引入了tensormap.replace PTX指令，用于在设备端更新TMA描述符。
- A和B的占位符TMA描述符在主机端构建，并通过内核启动参数传递。
- 每个协作线程块（CTA）在全局内存中创建这些占位符TMA描述符的副本。
- 每当遇到一个新问题（即组中的新GEMM），描述符的地址、形状和步长就会被更新。

CUTLASS中当前对分组GEMM的支持

• Ampere架构支持：自CUTLASS 2.8起提供支持，优化始于2.10版本。
• Hopper架构支持：在CUTLASS 3.4中增加了测试版，预计未来版本会提供更多优化。

• 当前应用案例：

  • TRT-LLM for Mixtral 8x7B
  • ByteDance ByteTransformer
  • PyTorch Geometric

• 示例代码：

  • examples/24_gemm_grouped
  • examples/57_hopper_grouped_gemm
  • examples/python/02_pytorch_extension_grouped_gemm

用于LLM量化的混合输入GEMM

为了显著减少内存占用和带宽，大型语言模型（LLM）的量化采用了混合输入GEMM。

• 未量化GEMM：权重（FP16）与激活值（FP16）相乘。
• 权重 量化（混合输入）：权重（INT4）与激活值（FP16）相乘。

其核心机制是在执行Tensor Core操作之前，将“窄”数据类型（如INT4）转换为“宽”数据类型（如FP16）。

混合输入GEMM节省全局和共享内存

混合输入GEMM通过在数据传输的后期阶段进行类型转换来优化内存使用。权重等“窄”类型数据以其原始紧凑格式从全局内存加载到共享内存。从共享内存移动到寄存器文件以供Tensor Core MMA（矩阵乘法累加）单元使用之前，会进行优化的类型转换。这个过程节省了宝贵的全局内存和共享内存空间。

CUTLASS对混合输入GEMM的支持

CUTLASS 3.x通过其模板化和可组合的接口，极大地简化了混合输入GEMM的实现。对于一个标准的FP16 GEMM，其CollectiveMainloop的类型定义如下所示。

要将其转换为混合输入GEMM（例如，INT4权重和FP16激活值），只需将第一个操作数的类型从cutlass::half_t更改为cutlass::int4b_t即可，如下面的代码高亮部分所示。

支持多种混合输入组合

CUTLASS支持广泛的混合输入数据类型组合，只需在CollectiveBuilder中指定ElementA和ElementB的类型即可。支持的组合包括：
- FP16 x INT8
- FP16 x INT4
- BF16 x INT8
- FP8 x INT4
- INT8 x FP16
- INT8 x BF16
- INT4 x FP8
- INT2 x FP16
- INT1 x FP8
- 等等

要求:
- 必须存在针对更宽数据类型的Tensor Core指令。
- 操作数必须满足TMA的要求。

CUTLASS中当前对混合输入GEMM的支持情况

• Hopper支持：

  • 在CUTLASS 3.3中添加。
  • 当前已在TRT-LLM中使用，相关信息可参考NVIDIA关于H200上LLM加速的博客。

• Ampere支持：

  • 在CUTLASS 3.3中贡献。
  • 感谢来自Google的Manish Gupta。相关信息可参考关于混合输入矩阵乘法性能优化的博客。

• 示例代码：

  • examples/55_hopper_mixed_dtype_gemm

---

Epilogue Visitor Tree (EVT)

常见的 workloads 调用各种 epilogues

在深度学习中，许多工作负载在GEMM计算之后会调用一系列操作，这些操作统称为“epilogue”（结尾部分）。为了获得高性能，将epilogue与GEMM融合成一个单一的内核至关重要，这可以避免将中间结果写回全局内存。常见的融合模式包括：

• GEMM + ReLU
• GEMM + Bias Add
• GEMM + Bias Add + GELU
• GEMM + GELU + Residual
• GEMM + Bias Add + ReLU + Row reduction
• 等等

为每一种融合编写新内核非常繁琐

传统的融合内核开发工作流程效率低下：
- 先前的工作流程：
- 确定放置定制逻辑的位置。
- 复制现有的GEMM和epilogue文件进行修改。
- 对每个需要添加的新融合操作重复此过程。

• 这种方法导致开发工作量大，代码重复率高。

Epilogue Visitor Tree (EVT): 用于组合融合epilogue的构建模块

Epilogue Visitor Tree (EVT) 是解决上述问题的方案。它提供了一组可以组合在一起构建复杂epilogue的原始节点。
- 节点类型：
- Load（加载）: 累加器、辅助张量、行广播、列广播、标量广播。
- Compute（计算）: 元素级、二元、三元计算。
- Store（存储）: 辅助张量、行规约、列规约、标量规约。

• 这些节点可以组合成一个树状或有向无环图（DAG）结构。下图展示了如何用EVT表示标准的GEMM输出计算 (alpha * accumulators) + (beta * C)。

在CUTLASS 3.x中向GEMM添加epilogue visitor tree (C++)

下面通过一个例子展示如何在CUTLASS 3.x中使用EVT构建一个更复杂的epilogue：ReLU((alpha * accumulators) + bias + (beta * C))。

首先，这是该操作的计算图。

接下来，我们逐步构建这个图。第一步是实现 (alpha * accumulators) + bias。这通过定义Alpha、Accum和Bias的加载节点，然后将它们输入一个MultiplyAdd计算节点来完成。

然后，我们构建图的下一部分，即添加 (beta * C)。这通过定义Beta和C的加载节点，并将它们与上一步的结果EVTCompute0一起输入到另一个MultiplyAdd计算节点EVTCompute1。

为了在 CUTLASS 3.x 中使用 C++ 为通用矩阵乘法（GEMM）添加一个 Epilogue Visitor Tree (EVT)，我们可以通过组合不同的计算节点来构建一个自定义的 epilogue 操作。以下示例演示了如何实现 ReLU(alpha * accumulators) + bias + (beta * C) 这个融合操作。

首先，定义计算图中的各个节点，如 Alpha、Accum（累加器）、Bias 等，并定义它们之间的计算关系（例如 MultiplyAdd）。

然后，使用 CollectiveBuilder 将这些节点组合成一个完整的 CollectiveEpilogue。CollectiveBuilder 负责根据指定的硬件架构（如 sm90）、操作类型、瓦片形状（TileShape）和 EVT 输出节点来生成最终的 epilogue 操作。

用于常见模式的预置别名 (Pre-baked aliases)

为了简化开发，CUTLASS 为常见的 epilogue 模式提供了预置的别名。例如，前面几步中手动构建的复杂 EVT 树，可以被一个名为 Sm90LinCombPerRowBiasEltAct 的单一别名所替代。这大大减少了样板代码，提高了可读性。

Python 中的 Epilogue Visitor Tree

CUTLASS 3.x 同样在 Python 接口中支持 Epilogue Visitor Tree，使得定义复杂的融合操作变得更加简单直观。整个过程可以分为以下五个步骤：

1. 声明一个基本的 GEMM 操作: 使用 cutlass.op.Gemm 定义一个标准的 GEMM 计算。
2. 将 epilogue 定义为一个 Python 函数: 编写一个标准的 Python 函数 my_epilogue，其内部实现了所需的融合计算逻辑，例如 D = relu(alpha * accum + beta * C + bias)。
3. 定义每个 EVT 操作数/输出的类型和形状: 为 epilogue 函数中用到的所有张量（如累加器、C、D、偏置）提供示例输入，包括它们的形状和数据类型。
4. 构造 EVT 并将其分配给 GEMM: 调用 cutlass.epilogue.trace 函数，传入 epilogue 函数和示例输入。CUTLASS 会追踪这个 Python 函数的执行，并将其转换为一个 EVT 计算图，然后将其赋值给 GEMM 计划的 epilogue_visitor 属性。
5. 编译并运行核函数: 定义核函数的运行时参数，然后调用 plan.run() 来执行编译好的、带有自定义 epilogue 的 GEMM 核函数。

EVT 如何自动编写优化的 Epilogue 循环

Epilogue Visitor Tree (EVT) 的一个核心优势是它能自动处理优化后的 epilogue 循环代码的生成，将用户定义的计算逻辑无缝注入到底层的 CUDA 核函数中。

下图展示了一个典型的消费者存储（Consumer store）warpgroup 伪代码，它负责处理 epilogue 阶段的计算和存储。这个循环遍历所有的 epilogue 子瓦片（subtiles），从共享内存（smem）加载累加器片段，执行 epilogue 计算，然后将最终结果写回到全局内存。

当使用 EVT 时，用户定义的计算图会被转换成一系列的回调函数（callbacks），这些回调函数被插入到 epilogue 循环的不同阶段。如下图高亮部分所示，callbacks.begin(), callbacks.previsit(...), callbacks.visit(...), callbacks.reduce(...) 等回调函数在循环的不同位置被调用，从而执行用户自定义的融合操作，而无需用户手动编写底层的循环和访存逻辑。

CUTLASS 中 Epilogue Visitor Tree 的支持情况

• 接口: 同时支持 C++ 和 Python 接口。
• 硬件架构: 支持 Hopper, Ada, 和 Ampere 架构。

• 稀疏 GEMM:

  • 已在 Ampere 架构上支持稀疏 GEMM，感谢开源贡献者 Aleksandar Samardžić。
  • 对 Hopper 架构的稀疏 GEMM 支持即将到来。

• 示例代码:

  • C++ 示例: examples/49_collective_builder
  • Python 示例: examples/python/04_epilogue_visitor.ipynb

---

教程：使用 CUTLASS 编写自定义 CUDA 核函数

Ampere Gather/Scatter 卷积

这是一个使用 CUTLASS 编写自定义 CUDA 核函数的教程。

• 问题陈述: 需要实现空间稀疏卷积（spatially sparse convolutions）。

• 操作:

  • 沿 NDHW 维度进行 Gather 操作。
  • 沿 NZPQ 维度进行 Scatter 操作。
  • 通道维度 KC 和滤波器张量 KTRSC 是完全稠密的。

• 应用: 这种卷积对于处理点云数据非常有用。

下图展示了空间稀疏卷积的原理，输入 x 和权重 w 进行卷积，只在特定的位置进行计算，生成稀疏的输出。

步骤 0: 核函数 API

本教程将指导如何使用 CUTLASS 编写自定义 CUDA 核函数。第一步是设置设备端代码，使其能够原生接受 CuTe 张量。下面是一个卷积函子 Conv functor 的模板定义，其 operator() 接受代表滤波器（mFlt）、激活（mAct）和输出（mOut）的 CuTe 张量作为参数。

步骤 1: Ampere 稠密 3D 卷积

利用 CUTLASS 3.x API 的组合能力，可以实现稠密 3D 卷积。

• 使用 CuTe 布局代数 (layout algebra) 将输入张量转换为 im2col 域:
  通过 make_layout 和 make_shape、make_stride 等函数，可以创建一个新的布局 xformed_act_layout，它将原始的激活张量映射到一个 (nzpq, ctrs) 的逻辑视图，这正是 im2col 转换后的形式。
• 在逻辑 MNK 问题瓦片上调用现有的 GEMM collective:
  一旦数据被布局为 im2col 形式，就可以直接调用标准的 CollectiveMma（GEMM collective）来执行矩阵乘法，从而实现卷积计算。

步骤 2: 表示 Gather/Scatter 张量

利用 CuTe 布局的强大表示能力来处理稀疏性。

• 布局即函数: CuTe 中的布局本质上是从域坐标到协域坐标的映射函数。
• 构造映射逻辑输入坐标到 gather 索引对的布局:
  创建一个 xformed_act_logical_inner 布局，它将逻辑坐标 (ctrs) 映射到一个索引对 (idx_buffer_idx, dense_offset)。
• 构造映射 gather 索引到线性偏移的布局:
  创建一个 xformed_act_gather_outer 布局，它使用一个 IndexedGather 自定义步幅，将 idx_buffer_idx 映射到内存中的实际线性偏移。
• 将它们组合起来，实现 gather 间接寻址:
  使用 composition 函数将上述两个布局组合成 xformed_act_composed_layout。这个组合后的布局直接将逻辑坐标 (nzpq, ctrs) 映射到最终的内存地址，从而实现了稀疏数据的 gather 操作。

步骤 3: 与现有的 Collective 组合

利用 CUTLASS collectives 的组合能力。

• 用组合后的 gather 布局替换仿射布局:
  在创建激活张量 gAct 时，不再使用之前为稠密卷积创建的仿射布局，而是直接使用上一步中构造的 xformed_act_composed_layout。
• 在逻辑 MNK 问题瓦片上调用现有的 GEMM collective:
  由于 CuTe 的抽象，即使底层的内存布局是稀疏的 gather 布局，上层的 CollectiveMma 仍然可以像处理稠密矩阵一样被调用。这体现了 CUTLASS 3.x 强大的组合性和可扩展性，用户只需定义好数据的布局，即可复用高效的计算原语。

步骤 4: 优化

利用领域特定信息进行性能优化。

• 假设除了图像数量外，问题形状在编译时已知: 这允许进行大量的静态优化。
• 滤波器张量的全静态全局内存布局: 使用 make_ordered_layout 创建一个完全在编译期确定的布局。

• 只有图像数量 (N) 是动态的:

  • 全静态的激活步幅。
  • 几乎全静态的激活形状。
  • 全静态的填充 (padding) 和膨胀 (dilation)。

• CuTe 布局表示消除了所有运行时索引计算:

  • 消除了对复杂越界 (OOB) 预测的需求。
  • 消除了对 Delta 表的需求。

• 约 100 行代码即可实现 SOTA 性能。

亲手实践

以下是一些使用 CUTLASS 3.x 自定义融合的案例研究，可供参考：

• Tri Dao's Flash Attention V2:

  • https://github.com/Dao-AILab/flash-attention/

• Colfax research's FP8 implementation of FA-V2:

  • https://github.com/ColfaxResearch/cutlass-kernels/tree/master/src/fmha

• 本教程中的 Gather/Scatter 示例:

  • examples/59_ampere_gather_scatter_conv

---

结论与路线图

结论：CUTLASS

• CUTLASS 3.5 与 CuTe

  • 持续为用户提供灵活的抽象，以组合定制化的内核（kernels）和集合操作（collectives）。
  • 引入一种通用的卷积方法。
  • 实现针对大型语言模型（LLMs）的新内核。
  • 通过尾声访问者树（epilogue visitor tree）抽象来促进灵活的内核融合。

• CUTLASS 路线图

  • CUTLASS 3.5 已于2024年3月发布。
  • CUTLASS 3.6 计划于2024年夏季发布。
    • 将支持Hopper架构的稀疏性。

• 需要帮助或有疑问？

  • 提交问题：https://github.com/NVIDIA/cutlass/issues
  • 联系 Matthew Nicely ([email protected])

CUTLASS GitHub: https://github.com/NVIDIA/cutlass/

路线图

注：可能会有变动

• 2024年第一季度

  • Ada FP8 支持
  • 分组 GEMM (Hopper)
  • 混合输入 GEMM
  • PyPi Wheels
  • Ptr-Array 支持

• 2024年第二季度

  • 卷积 (Fprop, D/Wgrad)
  • 分组 GEMM 优化

• 2024年第三季度

  • 卷积优化
  • 带有Python接口的稀疏性支持
  • 更多的 EVT 支持
  • 可分离/深度卷积

• 2024年第四季度

  • 文档更新
  • Conda 打包