Flash Attention 的共同作者 Tri Dao 与普林斯顿大学的两位博士生最近联合推出了一个名为 QuACK 的新内核库。这一创新的内核库引起了广泛关注,尤其是在高性能计算领域。
QuACK 的开发背景
QuACK 的开发完全基于 Python 和 CuTe-DSL,令人瞩目的是,它不涉及任何 CUDA C++ 代码。这一设计理念打破了传统的编程框架,使得开发者能够在更友好的环境中进行高效的 GPU 编程。
性能优势
在强大的 H100 显卡上,QuACK 实现了比 PyTorch 中的 torch.compile 和 Liger 等库快 33%-50% 的速度提升。Tri Dao 强调,内存密集型内核的高效运作并不是一个难以实现的“秘密”,而是依赖于对一些关键细节的精准处理。他指出,理解现代加速器的线程与内存层级结构是至关重要的。
🧩 为什么这么快?
Tri Dao 的总结:
“让内存密集内核达到光速并不神秘,只要把 bank-conflict-free shared-memory layout、pipeline 延迟隐藏、L2 预取 这几个细节做对即可。”
CuTe-DSL 把上述硬件级技巧抽象成 Python 表达式,开发者无需写 PTX / CUDA 也能榨干 GPU。
🔧 上手示例
from quack import layer_norm
y = layer_norm(x, weight=gamma, bias=beta, eps=1e-5)
一行 Python,背后自动生成 TMA + Hopper WGMMA 指令流。
📦 资源与路线图
- 开源地址:https://github.com/Dao-AILab/quack