S2-Attention: Hardware-Aware Context Sharding Among Attention Heads

要約

コンテキストでトークンのサブセットに選択的に注意を払うまばらな注意は、効率的であると想定されていました。
ただし、フロップの理論的な減少は、フラッシュアットのようなハードウェアに認識された最適化がないため、密集した注意事項にわたって壁1杯のスピードアップに変換されることはめったにありません。
一方、まばらな注意がモデルの品質を今日の大規模な言語モデル（LLM）の規模で維持できるかどうかは不明のままです。
このホワイトペーパーでは、まばらなシェード（S2）の注意を紹介します。これは、ヘッドごととコンテキストごとのレベルの両方でカスタマイズ可能なスパース注意のカーネル最適化を提供するトリトンライブラリです。
S2アテンションにより、さまざまなモデルスケールでの広範なスパースアテンションデザインにわたる広範なアブレーションを通じて、斬新で高性能のまばらな注意技術の探索が可能になります。
これらの洞察から、私たちはいくつかの基本的なガイドラインを提示して、実用的な効率の改善だけでなく、下流のパフォーマンスも強力なものを達成できるまばらな注意を設計します。
高い並列化と最適化されたメモリIOを実現するには、まばらな注意は注意ヘッド全体でコンテキストを不均一に破壊する必要があります。そこでは、各ヘッドが完全なコンテキストをまとめてカバーしながら、トークンの異なるサブセットに注意してください。
一方、私たちは、実際には、まばらで密集した注意を組み合わせたハイブリッドアーキテクチャを見つけます。
S2アテンションは、8.79x、15.87x、25.3倍の壁1クロックスピードアップを達成し、強力なフラッシュアット2ベースラインと比較して、128kのコンテキスト長で完全な注意と完全な検索パフォーマンスを備えた強力なダウンストリームパフォーマンスを備えています。
推論では、7Bモデルの場合、S2アテンションカーネルの助けを借りて、私たちのモデルは、密集したカウンターパートと比較して4.5倍のスピードアップを達成します。
S2アテンションは、メガトロンとVLLMで直接使用するために、簡単に顧客化できるAPIでリリースされます。

要約(オリジナル)

Sparse attention, which selectively attends to a subset of tokens in the context was supposed to be efficient. However, its theoretical reduction in FLOPs has rarely translated into wall-clock speed-up over its dense attention counterparts due to the lack of hardware-aware optimizations like FlashAttention. Meanwhile, it remains unclear whether sparse attention can maintain the model’s quality at a scale of today’s large language models (LLMs) and how. This paper presents Sparsely-Sharded(S2) Attention, a Triton library that provides kernel optimization for sparse attention customizable at both per-head and per-context-range levels. S2-Attention enables the exploration of novel and high-performance sparse attention techniques, which we demonstrate through extensive ablations across a wide range of sparse attention designs at various model scales. From these insights, we present several basic guidelines to design sparse attention that can achieve not only practical efficiency improvements, but also strong downstream performance. To achieve high parallelization and optimized memory IO, sparse attention should shard the context heterogeneously across attention heads, where each head attends to a different subset of tokens while collectively covering the full context. Meanwhile, we find hybrid architectures combining sparse and dense attention particularly beneficial in practice. S2-Attention achieves wall-clock speedup of 8.79X, 15.87X, 25.3X compared to the strong FlashAttention-2 baseline with strong downstream performance on-par with full attention and perfect retrieval performance at a 128k context length. At inference, for 7B models, our model, with the help of our S2-Attention kernel, achieves 4.5x speed-up compared to dense counterparts. S2-Attention is released with easy-to-customize APIs for direct usage in Megatron and vLLM.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google