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