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

要約

コンテキスト内のトークンのサブセットに選択的に注意を向ける、まばらな注意が効率的であると考えられていました。
ただし、FlashAttendant のようなハードウェアを意識した最適化が欠如しているため、FLOP の理論的な削減が、対応する高密度アテンションよりも実時間の高速化につながることはほとんどありません。
一方、注意力がまばらな場合でも、今日の大規模言語モデル (LLM) の規模でモデルの品質を維持できるかどうか、またその方法は不明のままです。
このペーパーでは、ヘッドごとおよびコンテキスト範囲ごとの両方のレベルでカスタマイズ可能なスパース アテンションのカーネル最適化を提供する Triton ライブラリである Sparsely-Sharded(S2) Attendance について説明します。
S2-tention は、斬新で高性能なスパース アテンション技術の探求を可能にし、さまざまなモデル スケールでの幅広いスパース アテンション設計にわたる広範なアブレーションを通じて実証します。
これらの洞察に基づいて、実際的な効率の向上だけでなく、強力なダウンストリームのパフォーマンスも達成できるまばらな注意を設計するためのいくつかの基本的なガイドラインを示します。
高度な並列化と最適化されたメモリ IO を実現するには、スパース アテンションはアテンション ヘッド間でコンテキストを異種にシャーディングする必要があります。各ヘッドは集合的に完全なコンテキストをカバーしながら、トークンの異なるサブセットに対応します。
一方、実際には、疎な注意と密な注意を組み合わせたハイブリッド アーキテクチャが特に有益であることがわかりました。
S2-tention は、フル アテンションと同等の強力なダウンストリーム パフォーマンスと 128k コンテキスト長での完璧な取得パフォーマンスにより、強力な FlashAttendant ベースラインと比較して 8.79 倍、15.87 倍、25.3 倍の実時間の高速化を実現します。
推論すると、7B モデルの場合、S2-tention カーネルの助けを借りて、私たちのモデルは、高密度の対応物と比較して 4.5 倍の高速化を達成します。
S2-Attention は、Megatron および 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