ThunderKittens：简单、快速且可爱的 AI 内核
机器学习
2024-10-29 v1 人工智能
摘要
将 AI 架构映射到 GPU 硬件的挑战正在成为 AI 进步的关键瓶颈。尽管付出了大量努力，手写的自定义内核仍未达到其理论性能阈值，即使在线性注意力等成熟操作上也是如此。GPU 的多样化硬件能力可能表明我们需要各种各样的技术来实现高性能。然而，我们的工作探讨了少量关键抽象是否能极大地简化这一过程。我们提出了 ThunderKittens (TK)，一个用于编写高性能 AI 内核同时保持易于使用和维护的框架。我们的抽象映射到 GPU 层次结构的三个层级：(1) 在 warp 级别，我们提供 16x16 矩阵块作为基本数据结构以及类似于 PyTorch 的基于块的并行计算操作，(2) 在线程块级别，我们提供了一个模板用于在并行 warp 之间重叠异步操作，(3) 在网格级别，我们提供支持以帮助隐藏块启动和拆卸以及内存开销。我们通过提供一系列 AI 操作的内核来展示 TK 的价值，这些内核匹配或优于先前操作的内核。我们在 GEMM 和注意力推理性能上匹配了 CuBLAS 和 FlashAttention-3，并在注意力反向传播上优于最强基线 ，在状态空间模型上优于 ，在线性注意力上优于 。
引用
@article{arxiv.2410.20399,
title = {ThunderKittens: Simple, Fast, and Adorable AI Kernels},
author = {Benjamin F. Spector and Simran Arora and Aaryan Singhal and Daniel Y. Fu and Christopher Ré},
journal= {arXiv preprint arXiv:2410.20399},
year = {2024}
}