FlashPrefill V2: 長文脈LLMサービングのためのブロックスパース・プリフィル・アテンション

要旨

長文コンテキストモデリングは大規模言語モデルの中核となる能力であるが、アテンションの二次計算量は、特に計算集約的なプレフィル段階において、依然として重大なボトルネックとなっている。我々の以前の研究であるFlashPrefillは、瞬時パターン発見とmaxベースの動的閾値処理によってこのコストを軽減するものの、本番導入にはほど遠いアルゴリズムの試作実装に留まっている。本稿では、FlashPrefill V2を提案する。これは、FlashPrefillを試作実装から実用的な長文コンテキストサービングへと、以下の3つの側面で発展させるものである。第一に、近似誤差を効果的に抑制する平均補正項を導入し、極端なスパース性レベルでも性能劣化を許容範囲内に保つ。第二に、スパースアテンション演算子をPackGQAメモリアクセス、ワープ特殊化、ピンポンパイプライン処理によって再設計し、最新のFlashAttention-3/4実装と完全に整合させるとともに、実用的な量子化要件を満たすFP8推論をサポートする。第三に、FlashPrefill V2はページ管理されたKVキャッシュと連続バッチ処理をネイティブにサポートし、SGLangなどの最新推論フレームワークにおいてアテンションバックエンドとして統合することを可能にする。最も広く展開されている推論アクセラレータの一つであるNVIDIA H20 GPU上での広範な評価により、FlashPrefill V2はFP8およびBF16精度においてコンテキスト長128KでFlashAttention-2と比較してそれぞれ最大47.26倍および27.19倍の高速化を達成し、さらにFP8ではFA3/4に整合した高密度ベースラインに対しても30.49倍の高速化を達成することを実証した。

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. Our previous work, FlashPrefill, mitigates this cost through instantaneous pattern discovery and max-based dynamic thresholding; however, it remains an algorithmic prototype that is still distant from production deployment. In this paper, we present FlashPrefill V2, which evolves FlashPrefill from a prototype toward practical long-context serving along three dimensions. First, we introduce a mean correction term that effectively suppresses the approximation error, keeping performance degradation manageable even at extreme sparsity levels. Second, we redesign the sparse attention operator with PackGQA memory access, warp specialization, and pingpong pipelining, fully aligning with the latest FlashAttention-3/4 implementations and supporting FP8 inference to meet practical quantization requirements. Third, FlashPrefill V2 natively supports paged KV cache and continuous batching, allowing integration as an attention backend in modern inference frameworks such as SGLang. Extensive evaluations on NVIDIA H20 GPUs---among the most widely deployed inference accelerators---demonstrate that FlashPrefill V2 delivers up to 47.26x and 27.19x speedups over FlashAttention-2 at 128K context length under FP8 and BF16 precision, respectively, and, in FP8, still achieves a 30.49x speedup against an FA3/4-aligned dense baseline.