FlashPrefill V2：面向长上下文LLM服务的块稀疏预填充注意力

摘要

長上下文建模是大型語言模型的關鍵能力，然而注意力機制的二次複雜度仍然是嚴重的效能瓶頸，特別是在計算密集的預填充階段。我們先前的工作 FlashPrefill 透過即時模式發現與基於最大值的動態閾值方法來降低此成本；然而，它仍僅是一個演算法原型，距離生產環境部署尚有相當距離。在本文中，我們提出 FlashPrefill V2，沿三個維度將 FlashPrefill 從原型推向實際的長上下文服務。首先，我們引入均值修正項，有效抑制近似誤差，即使在極端稀疏度下仍能將效能衰減控制在可接受範圍內。其次，我們以 PackGQA 記憶體存取、warp 專用化及乒乓流水線重新設計稀疏注意力運算子，使其完全對齊最新的 FlashAttention-3/4 實作，並支援 FP8 推論以滿足實際量化需求。第三，FlashPrefill V2 原生支援分頁 KV 快取與連續批次處理，使其能夠作為注意力後端整合至 SGLang 等現代推論框架中。在 NVIDIA H20 GPU（最廣泛部署的推論加速器之一）上的大量評測結果顯示，FlashPrefill V2 在 128K 上下文長度下，相較於 FlashAttention-2 於 FP8 與 BF16 精度分別可達到最高 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.