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

摘要

长上下文建模是大语言模型的一项关键能力，然而注意力机制的二次复杂度仍然是关键瓶颈，尤其是在计算密集的预填充阶段。我们先前的工作 FlashPrefill 通过即时模式发现和基于最大值的动态阈值化来缓解这一开销；然而，它仍是一个算法原型，距离生产部署尚有较大差距。本文提出 FlashPrefill V2，沿三个维度将 FlashPrefill 从原型推向实用的长上下文服务。首先，我们引入均值修正项，有效抑制近似误差，即使在极端稀疏度下也能将性能退化控制在可接受范围内。其次，我们利用 PackGQA 内存访问、warp 特化和乒乓流水线重新设计了稀疏注意力算子，使其与最新的 FlashAttention-3/4 实现完全对齐，并支持 FP8 推理以满足实际量化需求。第三，FlashPrefill V2 原生支持分页 KV 缓存和连续批处理，可作为注意力后端集成到 SGLang 等现代推理框架中。在 NVIDIA H20 GPU（最广泛部署的推理加速器之一）上进行的大量实验表明，在 128K 上下文长度下，FlashPrefill V2 在 FP8 和 BF16 精度下分别较 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.