Generalized Neighborhood Attention: Multi-dimensional Sparse Attention at the Speed of Light

要約

近隣の注意などの多くのまばらな注意メカニズムは、通常、自己注意ベースライン上で一貫してスピードアップを提供することができませんでした。
これは、主に注意インフラストラクチャの複雑さのレベルと、AIハードウェアアーキテクチャの急速な進化によるものです。
同時に、特にコンピュータービジョンにおける最先端の基礎モデルの多くは、注意に大きく縛られており、O（n^2）の複雑さを逃れるために信頼できるスパース性が必要です。
この論文では、地域に焦点を当てた有望なまばらな注意メカニズムのクラスを研究し、パフォーマンスの改善のより良い分析モデルを開発することを目指しています。
まず、一般化された近隣の注意（GNA）を紹介します。これは、スライドウィンドウを記述し、伸びたスライドウィンドウを記述し、注意をブロックすることができます。
次に、これらのアプローチを実装する際に可能な設計の選択を検討し、特定の設定ではるかにリアルなスピードアップ上限を提供できるシミュレーターを作成します。
最後に、CutlassのNvidia Blackwell Architecture向けに設計された最先端の融合マルチヘッド注意（FMHA）カーネルの上にGNAを実装します。
私たちの実装は、多くの完全にブロックスパルスのケースで理論的に可能な最大スピードアップを完全に実現し、FP16で1.3 PetaFlops/secondの効果的な利用を達成することができます。
さらに、さまざまなGNA構成をCosmos-7B、Hunyuanvideo、Fluxなどの既製の生成モデルに接続し、微調整なしでB200で28％〜46％のエンドツーエンドのスピードアップを提供できることを示しています。
Nattenプロジェクトを通じて、シミュレーターとブラックウェルカーネルを直接オープンします。

要約(オリジナル)

Many sparse attention mechanisms such as Neighborhood Attention have typically failed to consistently deliver speedup over the self attention baseline. This is largely due to the level of complexity in attention infrastructure, and the rapid evolution of AI hardware architecture. At the same time, many state-of-the-art foundational models, particularly in computer vision, are heavily bound by attention, and need reliable sparsity to escape the O(n^2) complexity. In this paper, we study a class of promising sparse attention mechanisms that focus on locality, and aim to develop a better analytical model of their performance improvements. We first introduce Generalized Neighborhood Attention (GNA), which can describe sliding window, strided sliding window, and blocked attention. We then consider possible design choices in implementing these approaches, and create a simulator that can provide much more realistic speedup upper bounds for any given setting. Finally, we implement GNA on top of a state-of-the-art fused multi-headed attention (FMHA) kernel designed for the NVIDIA Blackwell architecture in CUTLASS. Our implementation can fully realize the maximum speedup theoretically possible in many perfectly block-sparse cases, and achieves an effective utilization of 1.3 petaFLOPs/second in FP16. In addition, we plug various GNA configurations into off-the-shelf generative models, such as Cosmos-7B, HunyuanVideo, and FLUX, and show that it can deliver 28% to 46% end-to-end speedup on B200 without any fine-tuning. We will open source our simulator and Blackwell kernels directly through the NATTEN project.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google