据最新报道，Flash Attention 的共同作者 Tri Dao 与普林斯顿大学的两位博士生联合推出了一个名为 QuACK 的新内核库，令人瞩目的是，他们仅使用 Python 和 CuTe-DSL 开发，完全没有涉及 CUDA C++ 代码。这一创新不仅打破了传统的编程框架，还在强大的 H100显卡上实现了比 PyTorch 中的 torch.compile 和 Liger 等库快33%-50% 的速度提升。
Tri Dao 表示，内存密集型的内核达到高效运作并不是一个难以实现的 “秘密”，而是依赖于对一些关键细节的精准处理。他强调，现代加速器的线程与内存层级结构的理解是至关重要的。随着对 GPU 性能优化的不断深入，利用 CuTe-DSL 这一基于 Python 的领域特定语言，开发者们能够在更友好的环境中实现性能的大幅提升。
这项成果迅速引起了众多业内专家的关注。英伟达 CUTLASS 团队的资深架构师 Vijay 对此表示赞赏，并强调 CuTe-DSL 的设计使得像 Tri Dao 这样的专家能够轻松地实现 GPU 的高效运行。他还透露，今年将会有更多关于这方面的精彩内容发布。同时，PyTorch 团队成员 Horace He 也对这一创新表示了极大的兴趣，尤其认为对于长序列处理有着显著的优势。
为了让更多开发者受益，QuACK 的作者们还撰写了一篇详细教程，介绍了实现的具体步骤及代码，便于大家直接使用。文章强调，要在 GPU 的模型训练和推理过程中达到高效运行，既要优化计算密集型内核，也要兼顾内存密集型内核。在过去的工作中，矩阵乘法和注意力机制的优化已经非常成熟，因此本次研究将焦点放在了内存密集型内核上。
作者解释，内存密集型内核的算术强度较低，因此吞吐量更依赖于每秒传输的数据量。通过巧妙利用 GPU 的内存层级结构以及硬件特性，作者们成功将内存密集型内核的性能提升至近乎 “光速” 的水平。