提升长文本生成速度,实现高效低延迟推理。
FlashPrefill V2: Block-Sparse Prefill Attention for Long-Context LLM Serving
- 采用稀疏注意力与动态阈值优化计算效率
- 在128K上下文下比基线快30倍以上
- 适合需要高吞吐长文本服务的场景
长上下文建模是大语言模型的关键能力,但注意力机制的二次复杂度仍是预填充阶段的瓶颈。此前工作FlashPrefill通过即时模式发现和基于最大值的动态阈值缓解开销,但仍为算法原型,难以投入生产。本文提出FlashPrefill V2,从原型走向实用化,在三个维度实现演进:首先引入均值修正项,有效抑制近似误差,即使在极端稀疏下性能下降仍可控;其次重新设计稀疏注意力算子,支持PackGQA内存访问、波段特化和乒乓流水,与FlashAttention-3/4对齐,并支持FP8推理以满足实际量化需求;第三,原生支持分页KV缓存和连续批处理,可作为现代推理框架(如SGLang)的注意力后端。在广泛部署的NVIDIA H20 GPU上评估显示,FlashPrefill V2在128K上下文长度下,相比FlashAttention-2分别实现高达47.26x(FP8)和27.19x(BF16)的加速,且在FP8下仍较对齐的稠密基线快30.49x。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。