斯坦福让“GPU 高速运转”的新工具火了，比 FlashAttention2 更快

图片地址：https://img.ithome.com/newsuploadfiles/2024/6/b89e5d5f-0592-4007-9b88-966f2deff4c3.png?x-bce-process=image/format,f_auto

AI 算力资源越发紧张的当下，斯坦福新研究将 GPU 运行效率再提升一波 —— 内核只有 100 行代码，让 H100 比使用 FlashAttention-2，性能还要提升 30%。怎么做到的？研究人员从“硬件实际需要什么？如何满足这些需求？”这两个问题出发，设计了 一个嵌入式 CUDA DSL 工具，名为 ThunderKittens（暂且译为雷猫）。雷猫可简化 AI 内核的编写，同时充分利用底层硬件能力。具体来说，雷猫的主要抽象是寄存器和共享内存中的小型张量块（tile），和目前 GPU 中对小矩阵乘法的优化相匹配。通过操作这些 tile，开发者可相对简单地编写代码，充分利用张量核心、异步数据传输和共享内存等硬件特性。使用雷猫实现的注意力机制内核，代码量少且能实现很高的硬件利用率，性能超过直接使用底层库（如 Cutlass）。详细讨论过程以及雷猫是怎么设计出的，研究人员以“GPUs Go Brrr”为题，发在了斯坦福 Ha