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

主要观点总结

文章介绍了斯坦福大学研究人员如何通过设计嵌入式CUDA DSL工具ThunderKittens（雷猫）来提升AI内核的性能。雷猫简化了AI内核的编写，并充分利用底层硬件能力。它通过操作小型张量块（tile），开发者可以相对简单地编写代码，并利用张量核心、异步数据传输和共享内存等硬件特性。使用雷猫实现的注意力机制内核，代码量少且能实现很高的硬件利用率，性能超过直接使用底层库。文章还探讨了H100 GPU的硬件细节，并分析了影响GPU性能的因素。研究人员发现，要保持张量核心持续运行，必须注意WGMMA指令的使用、共享内存的速度、地址生成的成本以及保持高占用率。

关键观点总结

关键观点1: 设计嵌入式CUDA DSL工具ThunderKittens（雷猫）

雷猫简化了AI内核的编写，并充分利用底层硬件能力。它通过操作小型张量块（tile），开发者可以相对简单地编写代码，并利用张量核心、异步数据传输和共享内存等硬件特性。

关键观点2: 使用雷猫实现高性能内核

使用雷猫实现的注意力机制内核，代码量少且能实现很高的硬件利用率，性能超过直接使用底层库。

关键观点3: 分析H100 GPU的硬件细节

文章探讨了H100 GPU的硬件细节，包括内存、缓存和流式多处理器等，并分析了影响GPU性能的因素。

关键观点4: 保持张量核心持续运行的注意事项

研究人员发现，要保持张量核心持续运行，必须注意WGMMA指令的使用、共享内存的速度、地址生成的成本以及保持高占用率。

关键观点5: 未来研究方向

研究人员认为，应该根据硬件特性来重新定义AI的设计理念，并利用对硬件的了解来帮助设计与之匹配的AI。