據最新報道，Flash Attention 的共同作者 Tri Dao 與普林斯頓大學的兩位博士生聯合推出了一個名爲 QuACK 的新內核庫，令人矚目的是，他們僅使用 Python 和 CuTe-DSL 開發，完全沒有涉及 CUDA C++ 代碼。這一創新不僅打破了傳統的編程框架，還在強大的 H100顯卡上實現了比 PyTorch 中的 torch.compile 和 Liger 等庫快33%-50% 的速度提升。

Tri Dao 表示，內存密集型的內核達到高效運作並不是一個難以實現的 “祕密”，而是依賴於對一些關鍵細節的精準處理。他強調，現代加速器的線程與內存層級結構的理解是至關重要的。隨着對 GPU 性能優化的不斷深入，利用 CuTe-DSL 這一基於 Python 的領域特定語言，開發者們能夠在更友好的環境中實現性能的大幅提升。

image.png

這項成果迅速引起了衆多業內專家的關注。英偉達 CUTLASS 團隊的資深架構師 Vijay 對此表示讚賞，並強調 CuTe-DSL 的設計使得像 Tri Dao 這樣的專家能夠輕鬆地實現 GPU 的高效運行。他還透露，今年將會有更多關於這方面的精彩內容發佈。同時，PyTorch 團隊成員 Horace He 也對這一創新表示了極大的興趣，尤其認爲對於長序列處理有着顯著的優勢。

爲了讓更多開發者受益，QuACK 的作者們還撰寫了一篇詳細教程，介紹了實現的具體步驟及代碼，便於大家直接使用。文章強調，要在 GPU 的模型訓練和推理過程中達到高效運行，既要優化計算密集型內核，也要兼顧內存密集型內核。在過去的工作中，矩陣乘法和注意力機制的優化已經非常成熟，因此本次研究將焦點放在了內存密集型內核上。

作者解釋，內存密集型內核的算術強度較低，因此吞吐量更依賴於每秒傳輸的數據量。通過巧妙利用 GPU 的內存層級結構以及硬件特性，作者們成功將內存密集型內核的性能提升至近乎 “光速” 的水平。