ThunderKittens: Simple, Fast, and Adorable Kernels

作者/机构: Benjamin F. Spector1, Simran Arora1, Aaryan Singhal1, Arjun Parthasarathy2, Daniel Y. Fu3,4, Christopher Re´1 (1 斯坦福大学, 2 哥伦比亚大学, 3 加州大学圣地亚哥分校, 4 Together AI)

A1 主要贡献

核心问题: 将人工智能（AI）架构高效地映射到GPU硬件上已成为AI发展的一个关键瓶颈。尽管付出了巨大努力，手写的定制内核（kernel）仍未能达到其理论性能阈值，即使是在线性注意力（linear attention）等成熟操作上也是如此。

研究目标: 本文探讨是否可以通过一小组关键的抽象（abstractions）来大幅简化编写高性能AI内核的过程，而不是依赖多种多样的技术来适应GPU的复杂功能。

创新点/核心贡献:
本文提出了THUNDERKITTENS (TK)，一个旨在简化高性能AI内核编写过程的框架。其核心贡献是通过三个层级的抽象来应对GPU的层级化并行结构：

1. Warp级抽象：提供带托管布局的瓦片（Tile）数据结构和类PyTorch操作。

   • 基础数据结构：使用16x16矩阵瓦片作为基本数据单元，以最大化与Tensor Core的兼容性并鼓励其使用。
   • 自动布局管理：TK自动为瓦片选择最优的内存布局，以最小化内存区域冲突（bank conflicts），同时保持与专用硬件指令的兼容性，减轻了用户的负担。
   • 类PyTorch/NumPy接口：提供了一套在瓦片上进行的并行计算原语，如逐点乘法、矩阵乘加（mma）、指数（exp）和累加和（cumsum），使得接口对机器学习研究者十分友好。

   

2. 线程块级抽象：提供异步工作流的程序模板。

   • 生产者-消费者模型：TK提供了一个通用的内核模板，用于协调线程块内不同warp之间的异步执行。开发者只需在模板中填充几个样板函数。
   • 延迟隐藏：该模板内部通过内存流水线和同步原语来隐藏延迟，例如，允许加载/存储数据的worker与在快速内存中执行计算的worker异步重叠执行。

3. 网格级抽象：优化线程块调度以隐藏开销。

   • 流水线气泡（Pipeline Bubbles）减少：TK帮助开发者减少因线程块启动和销毁开销导致的性能损失。
   • L2缓存命中率提升：通过支持持久化网格（persistent grid），在线程块边界之间重叠内存加载，从而提高数据重用和L2缓存效率。

验证:
通过实现一系列简单且多样的内核，本文验证了TK抽象的价值。这些内核在性能上达到或超过了现有技术。

• 在GEMM和注意力推断任务上，性能与CuBLAS和FlashAttention-3持平。
• 在注意力反向传播上，比最强基线快10-40%。
• 在状态空间模型上，快8倍。
• 在线性注意力上，快14倍。

A3 背景知识

2.1 GPU层级结构

GPU的软件层级结构紧密遵循其物理硬件层级（如图3所示）。以下是其几个最重要的组成部分和方面。

1. Warp：由32个连续线程组成的小组，它们在小而快的寄存器内存中操作数据。这些指令在物理执行单元上运行。单个线程可以通过指令级并行同时占用多个专用执行流水线，而不同的warp可以进一步占用可用的执行硬件：

   • (a) 加载和存储单元：用于将数据移入和移出寄存器。先进的GPU（如H100）还引入了专用的硬件加速器（如Tensor Memory Accelerator, TMA）用于HBM和共享内存之间的异步批量数据移动。
   • (b) 通用计算流水线：如用于max、min的ALU，用于乘加的FMA，以及用于exp等复杂操作的XU。不同流水线的吞吐量不同。
   • (c) 加速矩阵乘法硬件（Tensor Cores）：占据了GPU大部分的计算能力。
     线程可能因多种原因暂时停顿，包括固定的指令延迟、内存延迟、屏障、流水线节流或指令缓存未命中。

2. 线程块 (Thread blocks)：由多个warp组成，共同在称为流式多处理器（Streaming Multiprocessor, SM）的物理核心上执行一个内核。尽管每个SM只有四个物理执行单元，但最多可以同时运行64个软件warp（称为“占用率”或occupancy）。这些共存的warp通常会竞争硬件资源，如寄存器、共享内存、发射槽和计算流水线，但它们共同有助于在每个执行单元内同时运行多个工作流。Warp在屏障处同步，期间不能发布新任务。

   • 共享内存 (SMEM)：同一块内的warp可以通过专用的共享内存（SMEM，在H100上为227 KB，33 TB/s）快速通信。为提高带宽，SMEM被分为32个物理“banks”，可同时提供内存服务。然而，如果不同线程同时访问同一个bank（称为“bank冲突”），它们的访问必须串行化，这会增加访问延迟并降低可用带宽。
   • 寄存器与L1缓存：Hopper架构每个线程最多可使用255个寄存器，超出部分会“溢出”（spill）到L1缓存。SMEM可以重新配置为L1缓存，用于快速访问常用内存，如溢出的寄存器。

3. 网格 (Grids)：由多个线程块组成，用于运行内核。H100 SXM GPU有132个物理SM，可以同时运行线程块。虽然SM能够共存多个线程块，但大多数AI内核通过在单个线程块内增加warp数量（提高占用率）即可实现高性能。同一GPU上的线程块共享公共内存资源：

   • 高带宽内存 (HBM)：容量大但速度慢（80 GB, 3 TB/s），是GPU内存中延迟最高、带宽最低的。
   • L2缓存：容量较小但速度更快（50 MB, 12 TB/s），由硬件管理。如果线程块重用相同数据，L2缓存有助于降低内存延迟并增加带宽。
     调度线程块存在开销。首先，块启动会产生设置成本。其次，如果网格大小设置不当，会产生“尾部效应”（tail effect）。例如，在一个有132个SM的H100上执行一个包含133个块的内核，需要两波执行，第一波效率全满，第二波效率低于1%。使用独立的CUDA流等更高级的调度可以缓解这些尾部效应，如最近关于异步张量并行调度的研究【10，Flux: Fast software-based communication overlap on gpus through kernel fusion，2024，https: //http://arxiv.org/abs/2406.06858】。

2.2 成本模型

GPU并行性的简化成本模型。本文提出了一个受roofline模型【49，Roofline: An insightful visual performance model for floating-point programs and multicore architectures，2008】启发的简化成本模型。内核总执行时间 $C_{Overall}$ 是以下成本的总和，其中内存成本是延迟和带宽的组合，计算成本是延迟和吞吐量的组合。

$$ \text{C}_{\text{Overall}} = \max \left( \underbrace{\text{C}_{\text{HBM}}, \text{C}_{\text{L2}}, \text{C}_{\text{L1}}, \text{C}_{\text{Shared}}}_{\textbf{Memory}}, \underbrace{\text{C}_{\text{Tensor}}, \text{C}_{\text{ALU}}, \text{C}_{\text{FMA}}, \text{C}_{\text{XU}}}_{\textbf{Compute}} \right) + \underbrace{\text{C}_{\text{Setup}} + \text{C}_{\text{Sync}}}_{\textbf{Overhead}} $$

该模型代表了内存、计算和Tensor Core成本之间完美重叠的理想情况。内核的实际成本将介于各项成本的最大值与总和之间，具体取决于工作负载的属性（例如，某些操作本质上是顺序的）和实现效率。作者的目标是（1）降低这些个体成本，以及（2）改善它们的集体重叠。

2.3 GPU编程框架

现有AI内核开发框架的启发。本文的灵感来源于多个简化AI内核开发的相关工作，如NVIDIA CUTLASS/CuTe【31，Cuda templates for linear algebra subroutines，2017】和Triton【42，Triton: an intermediate language and compiler for tiled neural network computations，2019】。

• CUTLASS：其大量的嵌套CUDA模板为高度优化的AI内核提供了动力【39，Flashattention-3: Fast and accurate attention with asynchrony and low-precision，2024；8，A case study in cuda kernel fusion: Implementing flashattention-2 on nvidia hopper architecture using the cutlass library，2023a；9，Developing cuda kernels for accelerated matrix multiplication on nvidia hopper architecture using the cutlass library，2023b】。由于TK和CUTLASS都是嵌入式库，它们在表达能力上是等价的，都允许用户使用C++的全部功能。TK采取了互补的方法，选择了更为“固执己见”（opinionated）的抽象，旨在探索用一小组模板能走多远，以及简洁性是否会牺牲性能。这种方法有望提高AI研究人员的可访问性，因为充分利用CUTLASS的功能可能具有挑战性。作者发现，即使是像FlashAttention-3这样用CUTLASS编写的工业级内核，也存在可避免的bank冲突问题。TK旨在通过抽象来为用户管理这些问题。

• Triton及其他编译器：Triton、PyTorch【33，Pytorch: An imperative style, high-performance deep learning library，2019】、TVM【11，{TVM}: An automated {End-to-End} optimizing compiler for deep learning，2018】、TensorFlow XLA【1，Tensorflow: Large-scale machine learning on heterogeneous distributed systems，2016】等从编译器角度解决问题。这些框架不是C++嵌入式的，因此使用不受支持的专用硬件指令可能很困难。在高级框架中管理异步执行和寄存器使用也可能很困难。TK探索的是保留类似PyTorch的简洁感，同时实现最高性能的途径。

A2 方法细节

3.1 使用熟悉的数据结构和操作实现Warp级并行

核心抽象。THUNDERKITTENS (TK) 的核心建立在两个基本抽象之上：内存层级中每一级的瓦片（tile）数据结构，以及类似于PyTorch和NumPy中熟悉的操作套件的批量操作数。

编程抽象。TK的设计深受PyTorch和NumPy的启发，旨在为ML用户提供熟悉感。它提供了一套简洁的并行计算操作（如图2所示），由一个“worker”（一个warp或一个warpgroup，即4个warp）执行，这些worker协同拥有并操作一块数据。TK使用一个16×16的矩阵瓦片作为其基本数据结构，旨在最大化与Tensor Core的兼容性。TK为内存层级的每一级提供了瓦片：
1. 寄存器瓦片和向量：通过类型、形状和布局进行模板化。例如，在图2中，初始化了一个bfloat16类型、列主序布局、高度16、宽度64的瓦片。对寄存器内存的显式控制有助于用户减少第2节中提到的 $C_{Memory}$。
2. 共享瓦片和向量：通过类型和形状进行模板化。
3. 全局布局描述符：将HBM的加载和存储设置为对4D张量的索引，其维度可以在运行时或编译时确定（从而节省宝贵的寄存器）。
这些基于瓦片的抽象的一个优点是，它们使TK能够静态检查布局和操作，这对于难以调试的GPU内核非常重要。例如，一个寄存器内的Tensor Core乘法 mma AB 要求A为行主序布局，B为列主序布局，如果这些条件不满足，TK可以在编译时报错。

选择内存布局。布局指定了逻辑数据元素如何映射到物理线程的所有权。不同的瓦片大小、类型和硬件加速指令受益于不同的布局，而某些布局会导致bank冲突。TK的目标是：
* 让寄存器瓦片（最快的内存）保持Tensor Core（最快的计算单元）所使用的布局。如图1（左）所示，每种颜色代表不同线程对数据的所有权。这些格式使用起来很困难，如图4进一步所示。
* 支持使用硬件加速指令（例如，异步矩阵乘法和批量复制指令），这些指令也需要特定的共享内存布局。
在TK中，作者将布局简化为3种——在32、64和128字节边界上进行“swizzle”（交错），并根据瓦片的大小和类型自动为其分配能最小化bank冲突的布局。这种方法有助于最小化冲突，从而降低第2节中提到的 $C_{Shared}$。相比之下，即使是使用CUTLASS模板编写的FlashAttention-3内核也面临bank冲突，影响了性能。

3.2 使用通用异步模板实现块级并行

通过协调异步执行减少开销。THUNDERKITTENS通过协调线程块中的worker如何异步重叠执行，帮助开发者减少开销。尽管GPU的层级结构可能暗示需要多种技术，但本文提出了一个单一的简洁模板，该模板在广泛的AI工作负载上都能实现高性能。这个模板被称为LCSF（Load-Compute-Store-Finish），它建立在经典的生产者-消费者范式【17，Co-operating sequential processes，1968；7，Cudadma: Optimizing gpu memory bandwidth via warp specialization，2011】之上。

编程抽象。典型的AI内核流程是从HBM加载瓦片到SRAM，在快速内存中执行计算，将结果瓦片存回HBM，然后为下一个瓦片重复此过程。要使用LCSF模板，开发者需要编写四个函数：
1. Load函数：指定加载worker应从HBM加载到共享内存的数据，以及何时通知计算worker该内存已准备就绪。
2. Compute函数：指定计算worker应执行的内核指令，使用3.1节的瓦片数据结构和操作原语。
3. Store函数：指定worker需要存储到HBM的数据。
4. Finish函数：在内核结束时，worker存储任何最终状态并退出。
TK提供了以下抽象来帮助开发者管理worker的重叠和同步：
1. 多级缓冲区：模板在共享内存中维护N级流水线缓冲区，用于从HBM加载和存储。加载/存储worker根据计算worker的状态从缓冲区中添加/移除数据瓦片。使用2级缓冲区可以隐藏HBM加载（或存储）的延迟，因为当计算worker处理当前瓦片时，下一个瓦片可以异步加载。更深的缓冲区可以减少计算worker之间的同步需求，允许它们同时操作多个瓦片。用户只需设置一个数字来指定级数，TK会为用户管理这些缓冲区的设置和使用。如表1所示，改变流水线级数N会显著影响GEMM内核的性能。

表1：流水线缓冲级数。我们测量GEMM内核在TK模板中改变流水线缓冲级数时的效率（TFLOPS）。

1. 同步屏障：加载/存储worker需要通知计算worker新内存已写入输入缓冲区。计算worker也需要通知加载/存储worker瓦片已写入输出缓冲区，或输入瓦片可以被替换。在TK模板中，提供了一个 arrive 函数，供worker发信号表示他们已完成其阶段。

2. 异步I/O：TK将同步和异步的加载和存储指令（包括 cp.async 和TMA）封装在同一个接口中。它还为全局布局描述符（gl）自动创建用于TMA硬件加速地址生成的张量映射描述符。

占用率与效率的权衡。TK参数化了加载/存储和计算worker的数量（即占用率），为开发者提供了一种简单的内核调优方式。如第2节所述，更高的占用率可以增加重叠，但也会导致对有限硬件资源（如寄存器）的竞争。寄存器减少，worker需要操作更小的数据瓦片，导致更多指令发射、SRAM到寄存器的I/O，以及可能因数据分区增多而带来的更高同步成本。
图6展示了注意力内核的占用率权衡。作者比较了（1）一个只使用warp级并行的简单内核和（2）一个用LCSF模板编写的内核。虽然两种内核的性能都随着占用率的增加而提升，直到资源竞争成为主导，但LCSF模板将朴素内核的帕累托前沿（Pareto frontier）向外扩展，实现了更高的性能。

3.3 通过块启动调度实现网格级并行

简化网格布局和调度。TK使用户能够更容易地尝试不同的网格布局和协调线程块的启动。这有助于减少每个线程块的设置和拆卸成本（第2节中的 $C_{Setup}$），并促进线程块之间的内存重用，以避免缓慢的HBM访问（第2节中的 $C_{HBM}$）。

块启动成本。TK提供了围绕“持久化网格”（persistent grid）的优化来最小化启动成本。在这种模式下，线程块在所有SM上预先启动一次，然后只需在现有块内加载内核的下一个任务。通过让加载/存储worker预测下一个任务并预加载内存，同时计算worker为前一个任务运行 finish 阶段，进一步消除了流水线气泡。表2展示了这些优化对矩阵乘法的影响。

表2：TK GEMM内核在不同矩阵维度K下，使用（yes）和不使用（no）持久化启动时的TFLOPS。

L2重用和块启动顺序。线程块之间需要通过HBM通信。当线程块重用内存时，数据通常可以在比HBM快得多的L2缓存中找到。然而，由于缓存替换策略，这种重用性取决于块的启动顺序。表3总结了在注意力（Attention）和通用矩阵乘法（GEMM）内核中改变块启动顺序对效率的影响。块启动顺序显著影响L2重用（通过HBM带宽测量），进而控制内核性能。

表3：L2重用：我们改变块启动顺序，并测量从HBM消耗的带宽（GB/s）和效率（TFLOPS）。对于注意力，我们比较了一个优化的内核（内部瓦片化为8行块）和一个朴素的内核（按行主序调度块）。对于注意力，我们比较了块顺序 (1) 序列长度N、头H、最外层批次B vs. (2) 最内层B、H、然后最外层N。块启动顺序具有显著的性能影响。

A4 实验环境

• 硬件配置:

  • 主要平台: NVIDIA H100 80GB SXM GPU
  • 其他平台: NVIDIA RTX 4090, Apple M2 Pro

• 软件配置:

  • CUDA: 12.6
  • Triton: 3.00
  • PyTorch: 2.4

• 测试的内核/模型:

  • GEMM (通用矩阵乘法)
  • Attention (因果、非因果、分组查询注意力)
  • Linear Attention (基于多项式和基于学习的特征图)
  • State Space Models (长卷积, Mamba-2)
  • Fused dropout-residual-layernorm
  • Rotary (旋转位置编码)

• 数据集: 论文未指定具体数据集，实验重点是评估各种AI计算原语（primitives）的性能，而非端到端的模型训练。

• 基准测试协议:
  • 为确保公平比较，所有内核在10次预热迭代后，使用cudaEvents测量10次基准测试迭代的总执行时间，并报告平均性能。
  • PyTorch基线使用torch.compile进行优化。
  • 基线内核（如CuBLAS、Triton）会通过网格搜索或自动调优工具进行调优，以报告其最佳性能。

A4 实验结果

4.1 TK赋能简单且高性能的AI内核

AI主力内核 (Workhorse Kernels):

• GEMM (通用矩阵乘法): 如图7所示，一个仅有40行设备代码的TK GEMM内核，其性能与业界最强的基线CuBLAS【32，cuBLAS，2023】和CUTLASS【31，Cuda templates for linear algebra subroutines，2017】相当。

  

• Attention: 如图8所示，TK支持多种注意力变体。

  • 前向传播: 在非因果注意力上，TK与同期的最强基线FlashAttention-3 (FA3)【39，Flashattention-3: Fast and accurate attention with asynchrony and low-precision，2024】性能相当。
  • 反向传播: 在因果和非因果注意力的反向传播中，TK比FA3快10%到40%以上，尤其在短序列上优势更明显。

  
  
  
  

新兴AI架构内核:

如图9所示，TK在多种新兴AI工作负载上也表现出色：

• Linear Attention:

  • 对于基于多项式的线性注意力，TK比流行的Flash Linear Attention (FLA)【51，Fla: A triton-based library for hardware-efficient implementations of linear attention mechanism，2024】（基于Triton）快14倍。
  • 对于基于学习特征图的线性注意力，TK比FLA快6.5倍。

• State Space Models (SSM):

  • 长卷积: TK比FlashFFTConv【22，Flashfftconv: Efficient convolutions for long sequences with tensor cores，2023c】快4.7倍（序列长度4096）至7.9倍（序列长度1024）。
  • Mamba-2: TK比Mamba-2的官方Triton内核【14，Transformers are ssms: Generalized models and efficient algorithms through structured state space duality，2024】快超过3倍，这主要得益于TK中融合复杂操作的便捷性。

• 其他内核: 在融合的dropout-residual-layernorm和旋转位置编码（rotary）等内存密集型操作上，TK同样比流行的Triton内核更高效。

4.2 内核实现对比分析

通过NVIDIA NSight Compute (NCU)工具进行的性能剖析揭示了TK性能优势的来源（见表4）：
* 长卷积: 与FlashFFTConv相比，TK通过其模板和warpgroup操作，实现了更好的worker重叠（更高的指令发射槽利用率和更少的内存停顿）和显著更高的Tensor Core利用率（提升4.1倍）。
* 注意力反向传播: 与FA3相比，TK的占用率调优更好（指令发射槽利用率更高），HBM吞吐量更高，等待HBM的停顿周期减少了10%。最显著的是，TK的共享内存停顿周期减少了85%，因为它避免了FA3中存在的高达9.6路的bank冲突。

表4：1) FlashAttention-3 (Shah et al., 2024) 与 TK 的注意力反向传播内核的 NCU 剖析；2) FlashFFTConv (Fu et al., 2023c) 与 TK 的长卷积内核的 NCU 剖析。

A5 结论

面对将AI架构映射到GPU硬件的挑战，本文探讨了使用一小组易于使用的编程抽象能达到的效果。THUNDERKITTENS (TK) 在GPU层级的每个层面都提供了抽象：在worker层面是带托管布局的瓦片（tile），在线程块层面是异步的LCSF模板。同时，本文还强调了在网格层面进行持久化块启动和优化L2重用的权衡。实验结果令人振奋：使用这些有限的抽象编写的内核不仅通用，而且在性能上始终达到或超过了当前最先进的水平。这表明，开发易于使用的AI硬件编程方式具有巨大潜力，可以在不牺牲性能的前提下，提高内核开发的可及性。

A6 附录

B.2 分析TK在不同工作负载下的可扩展性

FP8精度支持。TK支持FP8精度。图10展示了TK FP8 GEMM内核与CuBLAS的对比，两者性能相当。输入和输出均为FP8精度，累加使用FP32。

填充瓦片以支持非对齐工作负载。虽然TK默认使用16x16瓦片以鼓励利用Tensor Core和合并加载，但它也支持非对齐的工作负载。TK通过填充瓦片来处理此问题。在NVIDIA 4090和H100 GPU上实现的非对齐维度注意力内核，其性能特征与对齐内核保持一致。

B.3 分析TK在不同硬件平台上的可扩展性

消费级硬件：NVIDIA 4090 GPU。如图11所示，在NVIDIA 4090上，TK实现的非因果注意力内核与流行的参考内核FlashAttention-2【13，FlashAttention-2: Faster attention with better parallelism and work partitioning，2024】相比具有竞争力。

个人硬件：Apple M2芯片。如图12所示，在Apple M2芯片上，TK实现的非因果注意力和GEMM内核与Apple MLX框架的示例内核相比也具有竞争力。

图13和图14展示了在NVIDIA 4090和Apple M2上的注意力内核代码，突显了两者实现上的高度相似性。

B.4 分析TK的简洁性

库大小: 如表5所示，TK库的大小远小于CUTLASS和Triton，这反映了其轻量级的设计理念。

表5：各种CUDA库的大小。

代码行数 (LoC): 如表6所示，TK内核的平均代码行数少于200行，通常比相应的最先进基线内核更简洁，同时还能提供性能加速。

表6：TK H100内核、最先进的非TK内核的代码行数（LoC），以及TK相对于参考基线在第4节评估的输入维度上的加速比。

C THUNDERKITTENS内核代码清单

附录C提供了使用TK LCSF模板和瓦片抽象编写的内核代码示例。

C.1 矩阵乘法:
图15展示了一个与CuBLAS性能相当的模板化矩阵乘法内核。该内核的核心逻辑是：每个计算warpgroup负责计算输出矩阵的一个64x64的块。加载worker负责加载输入矩阵的块，计算worker执行异步矩阵乘法，最后存储worker将结果写回。内核的复杂性主要在于设置网格参数以优化L2缓存局部性。

C.2 长卷积:
图16展示了用于序列长度4096的长卷积内核，它使用FFT卷积算法，性能优于FlashFFTConv。

C.3 注意力:
图17展示了用于头维度64和128的非因果注意力内核，其性能与FlashAttention-3相当。

C.4 旋转位置编码:
图18展示了用于头维度128的旋转位置编码内核，其性能优于流行的Triton基线。

D 库实现细节

D.1 瓦片数据结构。
* 精度: TK可扩展支持FP32, FP16, BF16, FP8等数据类型。它通过在库内部处理不同数据类型所需的内存布局和打包格式的差异，对用户隐藏了复杂性，使用户可以使用统一的API（如mma, exp）操作不同精度的瓦片。
* 填充: 对于形状不是16倍数的工作负载，TK提供了填充机制。加载和存储操作可以进行边界检查，越界加载填充零，越界存储则不执行。此外，还提供了 top_fill, bottom_fill, left_fill, right_fill, triu, tril 等函数来处理瓦片内部的数据，以防止不正确的计算。

D.2 共享内存布局。
为了解决bank冲突问题，TK评估了多种共享内存布局，包括朴素的行主序布局（图19）、填充布局（图20）和多种交错（swizzled）布局（图21-23）。

THUNDERKITTENS的方法：结论是，支持HGMMA和UTMA指令的后三种交错布局（32字节、64字节、128字节）对于实现高性能至关重要。因此，TK在编译时根据瓦片的宽度选择可能达到的最高级别的交错，以最大程度地减少bank冲突。