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

文章预览

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