无需编写CUDA代码！H100带来33%-50%的加速性能，Flash Attention的创作者再度推出备受关注的创新作品
发布时间：2025-07-11 17:06:41 | 责任编辑：张毅 | 浏览量：21 次
据最新报道，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 的内存层级结构以及硬件特性，作者们成功将内存密集型内核的性能提升至近乎 “光速” 的水平。
Flash Attention 作者之一 Tri Dao 与两位普林斯顿 CS 博士生提出了一个名叫 QuACK 的新 SOL 内存绑定内核库，无需 CUDA 代码，完全用 Python 编写，借助 CuTe-DSL 实现。在带宽为 3TB/s 的 H100 上，QuACK 的速度比 PyTorch 的 torch.compile、Liger 等深度优化的库还要快 33%-50%。
技术细节
-
内存层级优化：H100 的内存层级结构复杂，访问延迟和带宽各不相同。QuACK 通过精心设计内存访问策略，将大部分本地归约操作分配在较高的内存层级上，只将少量经过本地归约后的值传递到下一个内存层级。
-
硬件感知的加载与存储策略：QuACK 在启动内核之前，通过特定的线程-值布局（TV-layout）对输入数据进行分区，确保每次加载操作在硬件上连续地传输最大数量的 bits。
-
归约策略：QuACK 采用了集群归约，多个 SM 协同工作，共享各自的资源，组成一个“超级”SM。
测试结果
QuACK 在 H100 上的测试结果显示，在归约维度大于 4k 时，内存吞吐量一般能稳定在 3TB/s 左右，接近峰值的 90%。例如，归约维度为 262k 时，FP32 的 softmax 吞吐量能达到 3.01TB/s，而 torch.compile 只有 1.89TB/s，快了近 50%。
行业影响
-
英伟达 CUTLASS 团队：资深架构师 Vijay 转发并表示 CuTe-DSL 把各种细节都打磨得很好。
-
PyTorch 团队：成员 Horace He 夸赞 QuACK“太酷了”，并提出了一些优化 torch.compile 性能的建议。
未来展望
Tri Dao 表示，高效的 GPU 内核开发流程是可以自动化的，未来或许只需调用“LLM.compile”就能生成高度优化的 GPU 内核。
本网站(https://aigc.izzi.cn)刊载的所有内容，包括文字、图片、音频、视频等均在网上搜集。
访问者可将本网站提供的内容或服务用于个人学习、研究或欣赏，以及其他非商业性或非盈利性用途，但同时应遵守著作权法及其他相关法律的规定，不得侵犯本网站及相关权利人的合法权利。除此以外，将本网站任何内容或服务用于其他用途时，须征得本网站及相关权利人的书面许可，并支付报酬。
本网站内容原作者如不愿意在本网站刊登内容，请及时通知本站，予以删除。
