最新の報道によると、Flash Attention の共同著者である Tri Dao は、プリンストン大学の2人の博士課程生と共同で、QuACK という新しいカーネルライブラリをリリースしました。注目すべきは、彼らが Python と CuTe-DSL を使用して開発したことで、CUDA C++ のコードは一切使っていない点です。この革新は従来のプログラミングフレームワークを打ち破り、強力な H100 グラフィックカード上で PyTorch の torch.compile や Liger などのライブラリよりも33〜50％も高速化する結果を出しました。

Tri Dao は、メモリ集約型のカーネルが効率的に動作することは「秘密」ではなく、いくつかの重要な詳細を正確に処理することに依存していると述べました。彼は、現代のアクセラレータのスレッドとメモリ階層構造の理解が非常に重要であると強調しました。GPU性能最適化の進展とともに、Pythonベースのドメイン固有言語である CuTe-DSL を利用することで、開発者はより使いやすい環境でパフォーマンスの大幅向上を実現できるようになりました。

image.png

この成果は、多くの業界専門家から注目を集めています。NVIDIA CUTLASS チームのシニアアーキテクトである Vijay はこれを称賛し、CuTe-DSL の設計が Tri Dao のような専門家にとって GPU の効率的な動作を簡単に実現可能にしていると強調しました。彼はまた、今年中にこれに関連する興味深い内容がさらに公開される予定であることを明かしました。一方、PyTorch チームのメンバーである Horace He もこのイノベーションに強い関心を示し、特に長期間のシーケンス処理において大きな利点があると考えています。

多くの開発者に恩恵をもたらすため、QuACK の作成者たちは詳細なチュートリアルを書きました。このチュートリアルでは、具体的なステップとコードが紹介されており、誰でも直接利用できるようにしています。記事では、GPUでのモデルトレーニングや推論プロセスで効率的な動作を達成するには、計算集約型カーネルだけでなく、メモリ集約型カーネルの最適化も考慮しなければならないと強調しています。これまでの研究では、行列乗算やアテンションメカニズムの最適化は非常に成熟していますので、今回の研究はメモリ集約型カーネルに焦点を当てています。

作成者らは、メモリ集約型カーネルの算術強度が低いため、トランザクション量（データ転送量）に強く依存すると説明しています。GPUのメモリ階層構造およびハードウェアの特性を巧みに活用することで、メモリ集約型カーネルのパフォーマンスをほぼ「光速」レベルまで引き上げることに成功しました。