トークン順列による疎なブロックスパースアテンション

要旨

大規模言語モデル（LLM）のコンテキスト長の拡大は大きな利点をもたらすが、計算コストが高い課題がある。このコストの主な原因は、シーケンス長に対してO(N^2)の計算量を要するセルフアテンション機構であり、メモリとレイテンシの両面で重大なボトルネックとなっている。幸い、アテンション行列は特に長いシーケンスにおいて疎となる傾向があり、最適化の余地が示唆されている。ブロック疎アテンションは、シーケンスをブロック単位に分割し、一部ブロックの計算を省略する有望な解決策として登場した。しかし、この手法の効果は基盤となるアテンションパターンに強く依存し、ブロックレベルの疎性が最適とは言えない場合がある。例えば、単一ブロック内のクエリにとって重要なキートークンが多数の他ブロックに分散している場合、計算の重複が生じる。本研究では、アテンションの置換特性を活用してブロックレベルの疎性を高め、LLMのプリフィリングの計算効率を向上させるプラグアンドプレイ方式「Permuted Block-Sparse Attention（PBS-Attn）」を提案する。実世界の難易度の高い長文コンテキストデータセットで包括的な実験を行った結果、PBS-Attnが既存のブロック疎アテンション手法をモデル精度で一貫して上回り、完全アテンションベースラインに極めて近い性能を達成することを実証した。カスタム開発したpermuted-FlashAttentionカーネルにより、PBS-Attnは長文プリフィリングにおいて最大2.75倍のエンドツーエンド高速化を実現し、実用性を確認した。コードはhttps://github.com/xinghaow99/pbs-attnで公開している。

English

Scaling the context length of large language models (LLMs) offers significant benefits but is computationally expensive. This expense stems primarily from the self-attention mechanism, whose O(N^2) complexity with respect to sequence length presents a major bottleneck for both memory and latency. Fortunately, the attention matrix is often sparse, particularly for long sequences, suggesting an opportunity for optimization. Block-sparse attention has emerged as a promising solution that partitions sequences into blocks and skips computation for a subset of these blocks. However, the effectiveness of this method is highly dependent on the underlying attention patterns, which can lead to sub-optimal block-level sparsity. For instance, important key tokens for queries within a single block may be scattered across numerous other blocks, leading to computational redundancy. In this work, we propose Permuted Block-Sparse Attention (PBS-Attn), a plug-and-play method that leverages the permutation properties of attention to increase block-level sparsity and enhance the computational efficiency of LLM prefilling. We conduct comprehensive experiments on challenging real-world long-context datasets, demonstrating that PBS-Attn consistently outperforms existing block-sparse attention methods in model accuracy and closely matches the full attention baseline. Powered by our custom permuted-FlashAttention kernels, PBS-Attn achieves an end-to-end speedup of up to 2.75times in long-context prefilling, confirming its practical viability. Code available at https://github.com/xinghaow99/pbs-attn

トークン順列による疎なブロックスパースアテンション

Sparser Block-Sparse Attention via Token Permutation

要旨

Support