斯坦福“雷猫”：让 GPU 高速运转，比 FlashAttention2 还快！ 2024-06-07 作者 C3P00 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
AI 算力资源越来越紧张，如何提升 GPU 的运行效率成为了热门话题。斯坦福大学的研究人员近日推出了一个名为 ThunderKittens（暂且译为“雷猫”）的嵌入式 CUDA DSL 工具，它能显著提升 GPU 的性能，甚至比目前最快的 FlashAttention-2 还要快 30%。
“雷猫”的秘密武器：小型张量块
“雷猫”的强大之处在于它巧妙地利用了 GPU 的硬件特性。它将 AI 内核的编写简化，并充分利用底层硬件能力。具体来说，它主要抽象了寄存器和共享内存中的小型张量块（tile），这与目前 GPU 中对小矩阵乘法的优化相匹配。通过操作这些 tile，开发者可以相对简单地编写代码，充分利用张量核心、异步数据传输和共享内存等硬件特性。
H100：如何榨干它的潜力？
为了更好地理解“雷猫”的原理，我们以英伟达最新的 H100 GPU 为例。H100 拥有强大的计算能力，但要充分发挥它的潜力，需要克服一些挑战。
“雷猫”：简化代码，提升性能
为了解决这些挑战，斯坦福的研究人员开发了“雷猫”。它包含四种模板类型：寄存器 tiles、寄存器向量、共享 tiles 和共享向量。开发者可以使用这些模板类型和一系列操作来处理张量，并充分利用 GPU 的硬件特性。
“雷猫”的优势：
未来展望：
斯坦福的研究人员认为，小型张量块的设计理念符合 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