FlashPrefill: 超高速長文脈プレフィリングのための瞬時パターン発見と閾値処理

要旨

長文脈モデリングは大規模言語モデルにおける重要な能力であるが、アテンションの二次計算量が依然として重大なボトルネックとなっており、特に計算集約的なプリフィル段階で顕著である。様々な疎アテンション手法が提案されているが、それらは検索遅延が大きいか、十分な疎性が得られないという課題を抱えている。本論文では、瞬時のパターン発見と閾値処理による超高速プリフィルを実現するFlashPrefillを提案する。FlashPrefillは高速ブロック検索技術を活用し、動的な垂直方向・斜め方向・ブロック疎アテンションパターンを同時に特定する。特に重要なのは、ソートやアテンションスコアの累積という過大なオーバーヘッドを回避しつつ、長裾分布を効果的に除去して疎性を高める動的閾値処理機構を導入した点である。大規模な評価により、FlashPrefillが256Kトークン系列で前例のない27.78倍の高速化を達成し、効率性において飛躍的進歩を実現することが実証された。既存手法が短文脈では効率低下するのとは異なり、FlashPrefillは4Kトークンという短い文脈長でも1.71倍の高速化を維持し、様々な系列長にわたる頑健性と実用性を証明している。

English

Long-context modeling is a pivotal capability for Large Language Models, yet the quadratic complexity of attention remains a critical bottleneck, particularly during the compute-intensive prefilling phase. While various sparse attention mechanisms have been explored, they typically suffer from either significant search latency or insufficient sparsity. In this paper, we propose FlashPrefill, a framework enabling ultra-fast prefilling via instantaneous pattern discovery and thresholding. FlashPrefill leverages a fast block-searching technique to simultaneously locate dynamic vertical, slash, and block-sparse attention patterns. Crucially, it introduces a dynamic thresholding mechanism that bypasses the prohibitive overhead of sorting or accumulating attention scores while effectively eliminating the long-tail distribution to enhance sparsity. Extensive evaluations demonstrate that FlashPrefill achieves a substantial leap in efficiency, delivering an unprecedented 27.78x speedup on 256K sequences. Notably, unlike existing methods that incur efficiency degradation on shorter contexts, FlashPrefill maintains a 1.71x speedup even at a 4K context length, demonstrating its robustness and practical utility across varying sequence scales.

FlashPrefill: 超高速長文脈プレフィリングのための瞬時パターン発見と閾値処理

FlashPrefill: Instantaneous Pattern Discovery and Thresholding for Ultra-Fast Long-Context Prefilling

要旨

Support