FlashPrefill V2がブロック疎注意機構をSGLangに統合、128KプリフィルのTTFTを最大4.8倍短縮

FlashPrefill V2は、もともとアルゴリズム検証寄りだったFlashPrefillを、長文コンテキストのサービングスタックに組み込めるプリフィル注意機構バックエンドへと刷新した。システムはまずブロック単位の近似スコアで重要なK/V領域を選び、それらのブロックだけに厳密な注意機構を適用する。選択されなかったブロックは単純にゼロ化せず、そのK/V平均値をonline softmaxにゼロ次補償として加えることで、高い疎性による出力のずれを抑える。インデックス作成段階ではCSR疎インデックスと補償統計を一度に生成し、Top-kソートを回避する。

カーネルはCUDA/CuTeを用いてHopper SM90向けに書き直され、PackGQAメモリレイアウト、warp specialization、producer-consumer ping-pong pipelineを採用し、BF16とFP8をサポートする。連続したK/Vしか受け付けない研究用カーネルが多いなか、これはpaged KV cacheを直接読み取り、可変長リクエストとcontinuous batchingを処理できる。チームはさらにSGLang 0.5.10ブランチへ独立したprefill backendを追加しており、デコードには引き続きFA3を使用できる。

NVIDIA H20を4基使用し、batch size 4、128Kコンテキストで行った演算子テストでは、FP8とBF16はFlashAttention-2に対してそれぞれ最大47.26倍、27.19倍の高速化を達成した。より高性能なFA3/4を密な注意機構の比較対象にしても、なお30.49倍、17.54倍だった。エンドツーエンドの結果はより現実的で、論文によると128Kでのtime-to-first-tokenは最大4.8倍改善した。open-loopテストでも、長いプリフィルを短縮することで、ほかのリクエストのデコードを妨げる時間が減ることが示された。

ただし、エンジニアリングチームは演算子単体の高速化倍率を、そのままサービング全体の高速化倍率と見なすべきではない。4K、BF16では疎性が約30％にとどまり、性能はおおむね密な注意機構のカーネルと同等になる。chunked prefillではインデックス作成が繰り返され、実効密度も高くなるため、効果が縮小する。公開実装が対応するのはH20／H100系のSM90 GPUのみで、SGLang対応も上流で正式採用されたバックエンドではなく、完全なソースツリーとして同梱されたものだ。今後は上流への統合、異なるHopperモデルでの再現結果、そしてより多くのモデルと実トラフィックにおいて長文コンテキストの品質が安定するかを見極める必要がある。