斯坦福“雷猫”:让 GPU 高速运转,比 FlashAttention2 还快!

AI 算力资源越来越紧张,如何提升 GPU 的运行效率成为了热门话题。斯坦福大学的研究人员近日推出了一个名为 ThunderKittens(暂且译为“雷猫”)的嵌入式 CUDA DSL 工具,它能显著提升 GPU 的性能,甚至比目前最快的 FlashAttention-2 还要快 30%。

“雷猫”的秘密武器:小型张量块

“雷猫”的强大之处在于它巧妙地利用了 GPU 的硬件特性。它将 AI 内核的编写简化,并充分利用底层硬件能力。具体来说,它主要抽象了寄存器和共享内存中的小型张量块(tile),这与目前 GPU 中对小矩阵乘法的优化相匹配。通过操作这些 tile,开发者可以相对简单地编写代码,充分利用张量核心、异步数据传输和共享内存等硬件特性。

H100:如何榨干它的潜力?

为了更好地理解“雷猫”的原理,我们以英伟达最新的 H100 GPU 为例。H100 拥有强大的计算能力,但要充分发挥它的潜力,需要克服一些挑战。

  • WGMMA 指令:H100 引入了一套新的指令集,名为“warp group matrix multiply accumulate”。这些指令对于充分发挥 H100 的计算能力至关重要,但它们的使用也颇为复杂,需要精心控制内存请求的合并和避免 bank conflicts。
  • 共享内存:共享内存的速度并不如预期快,使用时需要格外注意。处理不当可能会引发 bank conflicts,显著拖慢内核速度。
  • 地址生成:生成地址的成本较高,特别是加入复杂的交错或重排模式时,这种情况更为明显。
  • 占用率:保持高占用率对于提升性能是有益的,寄存器至关重要。

“雷猫”:简化代码,提升性能

为了解决这些挑战,斯坦福的研究人员开发了“雷猫”。它包含四种模板类型:寄存器 tiles、寄存器向量、共享 tiles 和共享向量。开发者可以使用这些模板类型和一系列操作来处理张量,并充分利用 GPU 的硬件特性。

“雷猫”的优势:

  • 简化代码:使用“雷猫”编写的内核代码量更少,复杂性更低。
  • 提升性能:“雷猫”可以显著提高 GPU 的硬件利用率,性能超过直接使用底层库(如 Cutlass)。
  • 适应性强:“雷猫”作为一个嵌入到 CUDA 中的库,其提供的抽象层在遇到不支持的功能时能够很好地处理。

未来展望:

斯坦福的研究人员认为,小型张量块的设计理念符合 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

0 0 投票数
Article Rating
订阅评论
提醒
0 评论
最多投票
最新 最旧
内联反馈
查看所有评论
人生梦想 - 关注前沿的计算机技术 acejoy.com
0
希望看到您的想法,请您发表评论x