斯坦福新工具“雷猫”：让 GPU 性能飙升，比 FlashAttention2 更快！

AI 算力资源越来越紧张，如何让 GPU 充分发挥潜力成为一大难题。最近，斯坦福大学的研究人员发布了一个名为“雷猫”（ThunderKittens）的工具，它可以显著提升 GPU 的运行效率，甚至比目前最快的 FlashAttention2 还要快 30%！

雷猫的秘密武器在于它对 GPU 硬件的深度理解。研究人员从“硬件实际需要什么？如何满足这些需求？”这两个问题出发，设计了一个嵌入式 CUDA DSL 工具。雷猫通过操作小型张量块（tile）来简化 AI 内核的编写，并充分利用张量核心、异步数据传输和共享内存等硬件特性。

H100：性能之王，如何榨干它的潜力？

研究人员以英伟达最新的 H100 GPU 为例，深入探讨了如何优化 GPU。H100 拥有 989 TFLOPs 的半精度矩阵乘法计算能力，但要充分发挥它的能力，关键是保持张量核心持续运算。

然而，要做到这一点并不容易。研究人员发现，H100 硬件具有一些特性，对于保持矩阵乘法的运行至关重要：

• WGMMA 指令： H100 引入了新的指令集 WGMMA，它允许 128 个线程跨 SM 所有子单元协作同步，并从共享内存及寄存器异步启动矩阵乘法。这些指令对于充分利用 H100 的计算能力是必不可少的，没有它们，GPU 的峰值利用率会损失 37%。
• 共享内存： 共享内存的延迟虽然看似不多，但对于高速运行的张量核心来说已经足够显著。此外，共享内存的存储单元独立，处理不当会导致 bank conflicts，显著拖慢内核速度。
• 地址生成： H100 的张量核心和内存速度极快，仅生成用于获取数据的内存地址就占用了芯片的大量资源。英伟达提供的张量内存加速器（TMA）可以异步提取张量的一部分，大大节省了地址生成的开销。
• 占用率： 占用率指的是 GPU 上同时调度的线程数。高占用率可以隐藏缺陷或同步问题，但一个设计良好的 pipeline 即使在占用率不高的情况下也能运行得相当快。

雷猫：小而精，性能大提升

为了更轻松地编写所需的内核类型，并充分发挥硬件的全部潜力，雷猫应运而生。它是一个嵌入在 CUDA 中的 DSL，包含四种模板类型：寄存器 tiles、寄存器向量、共享 tiles 和共享向量。雷猫提供了一系列操作来处理这些张量，既可在 warp 级别使用，也可用于多个 warp 协作。

雷猫的优势在于它的小而简单，功能有限，但它符合 AI 和硬件的发展趋势。研究人员认为，AI 的设计理念应该根据硬件特性进行重新定义，例如循环状态应该足够大以适应一个 SM，计算的密度应该不低于硬件的需求。

未来展望：AI 与硬件的深度融合

雷猫的出现，标志着 AI 和硬件深度融合的新方向。研究人员认为，未来需要利用对硬件的了解，设计与之匹配的 AI，从而进一步提升 GPU 的性能，释放 AI 的潜力。

参考文献：

[1] https://hazyresearch.stanford.edu/blog/2024-05-12-tk

[2] https://github.com/HazyResearch/ThunderKittens

[3] https://news.ycombinator.com/item?id=40337936

---

斯坦福大学的研究人员最近发布了一个名为“雷猫”（ThunderKittens）的工具，它可以显著提升 GPU 的运行效率，比目前最快的 FlashAttention2 还要快30% [1]。雷猫是一个嵌入在 CUDA 中的 DSL 工具，通过操作小型张量块来简化 AI 内核的编写，并充分利用 GPU 硬件的特性，如张量核心、异步数据传输和共享内存等 [1]。

下面是关于雷猫和 GPU 性能提升的一些关键信息：

雷猫：GPU 性能提升的秘密武器

• 雷猫是一个嵌入在 CUDA 中的 DSL 工具，通过操作小型张量块来简化 AI 内核的编写，并充分利用 GPU 硬件的特性 [1].
• 雷猫提供了四种模板类型：寄存器 tiles、寄存器向量、共享 tiles 和共享向量，用于处理张量操作 [1].
• 雷猫的优势在于它的小而简单，功能有限，但符合 AI 和硬件的发展趋势 [1].

GPU 性能提升的关键因素

• GPU 硬件具有一些特性，对于保持矩阵乘法的运行至关重要 [1]:
• WGMMA 指令：H100 GPU 引入了新的指令集 WGMMA，允许线程跨 SM 所有子单元协作同步，并从共享内存及寄存器异步启动矩阵乘法 [1].
• 共享内存：共享内存的延迟对于高速运行的张量核心来说已经足够显著，处理不当会导致 bank conflicts，拖慢内核速度 [1].
• 地址生成：GPU 的张量核心和内存速度极快，但生成用于获取数据的内存地址占用了大量资源，张量内存加速器（TMA）可以异步提取张量的一部分，节省了地址生成的开销 [1].
• 占用率：高占用率可以隐藏缺陷或同步问题，但设计良好的 pipeline 即使在占用率不高的情况下也能运行得相当快 [1].

未来展望：AI 与硬件的深度融合

• 雷猫的出现标志着 AI 和硬件深度融合的新方向，未来需要利用对硬件的了解，设计与之匹配的 AI，进一步提升 GPU 的性能，释放 AI 的潜力 [1].

参考文献：

1. 比标准Attention提速5-9倍，大模型都在用的FlashAttention v2来了 | 机器之心 [1]
2. 斯坦福大学发布的ThunderKittens工具 [1]
3. https://news.ycombinator.com/item?id=40337936 [1]

---

Learn more:

1. 比标准Attention提速5-9倍，大模型都在用的FlashAttention v2来了 | 机器之心
2. 通透理解FlashAttention与FlashAttention2：全面降低显存读写、加快计算速度-CSDN博客
3. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning | Princeton NLP Group