ThunderKittens: Simple, Fast, and Adorable AI Kernels

📄 arXiv: 2410.20399v1 📥 PDF

作者: Benjamin F. Spector, Simran Arora, Aaryan Singhal, Daniel Y. Fu, Christopher Ré

分类: cs.LG, cs.AI

发布日期: 2024-10-27

💡 一句话要点

ThunderKittens：一种简单、快速且易于维护的AI Kernel框架，提升GPU利用率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: GPU加速 AI Kernel CUDA 深度学习 高性能计算 线性注意力 并行计算

📋 核心要点

  1. 现有手写CUDA kernel难以充分利用GPU硬件特性，导致AI模型性能受限，尤其是在线性注意力等操作上。
  2. ThunderKittens (TK) 框架通过warp、thread-block和grid三个层级的抽象，简化了高性能AI kernel的编写。
  3. 实验表明，TK在GEMM和attention推理上与CuBLAS和FlashAttention-3性能相当，并在其他任务上显著超越现有方法。

📝 摘要（中文）

将AI架构映射到GPU硬件的挑战正在成为AI进步的关键瓶颈。尽管付出了巨大的努力，但即使在线性注意力等成熟操作上，手写的自定义kernel也未能达到其理论性能阈值。GPU多样化的硬件能力可能表明我们需要各种各样的技术来实现高性能。然而，我们的工作探索了少量关键抽象是否可以大大简化这一过程。我们提出了ThunderKittens (TK)，这是一个用于编写高性能AI kernel的框架，同时保持易于使用和维护。我们的抽象映射到GPU层次结构的三个级别：(1) 在warp级别，我们提供16x16矩阵tiles作为基本数据结构，以及类似PyTorch的并行计算操作；(2) 在thread-block级别，我们提供一个模板，用于跨并行warps重叠异步操作；(3) 在grid级别，我们提供支持以帮助隐藏block启动和关闭以及内存成本。我们通过提供kernel来展示TK的价值，这些kernel在各种AI操作上匹配或优于先前的kernel。我们在GEMM和attention推理性能上与CuBLAS和FlashAttention-3相匹配，并在attention反向传播上优于最强的基线10-40％，在状态空间模型上优于8倍，在线性注意力上优于14倍。

🔬 方法详解

问题定义：论文旨在解决AI模型在GPU上部署时，手动编写CUDA kernel效率低下的问题。现有方法难以充分利用GPU的硬件特性，导致性能瓶颈，且开发和维护成本高昂。即使是像线性注意力这样相对成熟的操作，也难以达到理论性能上限。

核心思路：论文的核心思路是通过引入一系列抽象层，将GPU的硬件层次结构（warp、thread-block、grid）暴露给开发者，同时提供易于使用的编程接口，类似于PyTorch。这样，开发者可以更方便地编写高性能的AI kernel，而无需深入了解底层的CUDA细节。

技术框架：ThunderKittens (TK) 框架包含以下三个主要层级： 1. Warp级别：提供16x16矩阵tiles作为基本数据结构，并提供类似PyTorch的并行计算操作，方便warp内的并行计算。 2. Thread-block级别：提供一个模板，用于跨并行warps重叠异步操作，从而提高GPU的利用率。 3. Grid级别：提供支持以帮助隐藏block启动和关闭以及内存成本，简化grid级别的编程。

关键创新：TK的关键创新在于其分层抽象的设计，它将GPU的硬件特性以一种易于理解和使用的方式暴露给开发者。这种抽象不仅简化了kernel的编写过程，还允许开发者更好地利用GPU的并行计算能力。与传统的手写CUDA kernel相比，TK降低了开发难度，提高了开发效率，同时保证了性能。

关键设计：TK的关键设计包括： 1. 16x16矩阵tiles：选择16x16作为warp级别的数据结构，是为了充分利用GPU的shared memory和register file，提高数据访问效率。 2. 异步操作模板：通过提供异步操作模板，TK允许开发者在不同的warps之间重叠计算和数据传输，从而隐藏延迟，提高GPU的利用率。 3. Grid级别优化：TK通过优化block启动和关闭以及内存成本，减少了grid级别的开销，进一步提高了性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，ThunderKittens在多个AI任务上表现出色。在GEMM和attention推理性能上，TK与CuBLAS和FlashAttention-3相当。更重要的是，TK在attention反向传播上优于最强的基线10-40％，在状态空间模型上优于8倍，在线性注意力上优于14倍，展示了其在特定任务上的显著优势。

🎯 应用场景

ThunderKittens框架可广泛应用于深度学习模型的GPU加速，尤其是在需要高性能自定义kernel的场景下，例如Transformer模型的优化、新型神经网络架构的探索以及特定硬件平台的性能优化。该框架降低了AI kernel开发的门槛，加速了AI算法的落地和部署。

📄 摘要（原文）

The challenge of mapping AI architectures to GPU hardware is creating a critical bottleneck in AI progress. Despite substantial efforts, hand-written custom kernels fail to meet their theoretical performance thresholds, even on well-established operations like linear attention. The diverse hardware capabilities of GPUs might suggest that we need a wide variety of techniques to achieve high performance. However, our work explores whether a small number of key abstractions can drastically simplify the process. We present ThunderKittens (TK), a framework for writing performant AI kernels while remaining easy to use and maintain. Our abstractions map to the three levels of the GPU hierarchy: (1) at the warp-level, we provide 16x16 matrix tiles as basic data structures and PyTorch-like parallel compute operations over tiles, (2) at the thread-block level, we provide a template for overlapping asynchronous operations across parallel warps, and (3) at the grid-level, we provide support to help hide the block launch and tear-down, and memory costs. We show the value of TK by providing kernels that match or outperform prior kernels for a range of AI operations. We match CuBLAS and FlashAttention-3 on GEMM and attention inference performance and outperform the strongest baselines by $10-40\%$ on attention backwards, $8\times$ on state space models, and $14\times$ on linear attention.