Source-linked AI summary
FlashPrefill V2: Block-Sparse Prefill Attention for Long-Context LLM Serving
Qihang Fan, Huaibo Huang, Zhiying Wu, Bingning Wang, Ran He
TL;DR
Long-context prefilling is bottlenecked by quadratic self-attention, while the earlier FlashPrefill prototype remained limited in accuracy robustness, kernel efficiency, and production compatibility. FlashPrefill V2 adds mean correction, a FlashAttention-3/4-aligned sparse operator with FP8 support, and native paged-cache and continuous-batching integration. At 128K context on NVIDIA H20, it reaches 47.26× FP8 and 27.19× BF16 speedups over FlashAttention-2, while achieving 30.49× over an FA3/4-aligned dense baseline in FP8.
Problem
Quadratic self-attention imposes prohibitive overhead on long-context sequences, especially during compute-intensive prefilling, and FlashPrefill remained distant from production deployment.
Method
FlashPrefill V2 combines mean correction for pruned blocks, a FlashAttention-3/4-aligned sparse operator supporting FP8, and native paged KV-cache and continuous-batching support.
Results
47.26× FP8 and 27.19× BF16 speedups over FlashAttention-2 are achieved at 128K context on NVIDIA H20, with 30.49× FP8 speedup over an FA3/4-aligned dense baseline.
Takeaways & Limitations
FlashPrefill V2 evolves sparse prefilling toward practical long-context serving with manageable accuracy degradation and compatibility with modern inference frameworks.
Takeaways & Limitations
The prior FlashPrefill prototype had accuracy degradation under aggressive sparsity, a FlashAttention-2-based kernel, and contiguous-KV incompatibility with paged-cache serving frameworks.
Abstract
from arXiv · showhide
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.
1. Introduction
FlashPrefill V2 addresses FlashPrefill’s remaining accuracy, kernel-efficiency, and serving-integration gaps for long-context prefilling. It combines mean correction, a hardware-aligned sparse operator, and production-oriented cache and batching support, with large H20 speedups at 128K context.
- Quadratic self-attention creates prohibitive overhead for long-context sequences, especially during compute-intensive prefilling.
- FlashPrefill reduces selection overhead through instantaneous pattern discovery and max-based dynamic thresholding but remains distant from production deployment.
- FlashPrefill V2 adds a mean correction term that keeps accuracy degradation manageable under extreme sparsity.The correction compensates pruned blocks with pooled K/V statistics inside attention computation.
- FlashPrefill V2’s redesigned sparse operator uses PackGQA, warp specialization, pingpong pipelining, FlashAttention-3/4 alignment, and FP8 support.
- 47.26× FP8 and 27.19× BF16 speedups over FlashAttention-2 are achieved at 128K context on NVIDIA H20 GPUs.Against an FA3/4-aligned dense baseline, the corresponding FP8 and BF16 speedups are 30.49× and 17.54×.
- Native paged KV cache and continuous batching support enables integration as an attention backend in serving frameworks such as SGLang.The introduction reports up to 4.8× lower end-to-end time-to-first-token at 128K with marginal accuracy loss.
2. Related Works
Related work spans sparse-attention pattern selection, hardware-aware dense attention kernels, and production serving systems. FlashPrefill V2 connects these strands by combining sparse computation with modern kernel execution and serving-system compatibility.
- Sparse Attention: Sparse-attention methods use fixed layouts, learnable patterns, or inference-time estimation to reduce attention’s quadratic cost.
- Efficient Attention Kernels: FlashAttention-2, FlashAttention-3, and FlashAttention-4 progressively emphasize parallelism, Hopper asynchrony, low-precision arithmetic, and kernel pipelining.
- LLM Serving Systems: Production LLM serving increasingly relies on paged KV-cache management, scheduling, and continuous batching to support throughput and latency objectives.
- FlashPrefill V2: FlashPrefill V2’s pipeline combines fused block selection and mean-statistic pooling with exact sparse attention and correction paths merged in online softmax.
- LLM Serving Systems: Unlike research-oriented sparse kernels assuming contiguous KV layouts, FlashPrefill V2 is designed for paged KV cache and continuous batching from the outset.
3. Method
FlashPrefill V2 combines block selection with mean correction and a Hopper-optimized sparse attention operator. Its two-stage design reduces memory and computation while supporting FP8 execution, grouped-query access, continuous batching, and SGLang integration.
- FlashPrefill framework: FlashPrefill partitions keys and values into blocks, estimates block salience with pooled keys, and executes attention only over selected blocks.The framework uses fused block-level score computation, max-based thresholding, and causal sink, window, and recent-block selections.
- FlashPrefill framework: O((L/B)^2) memory replaces O(L^2/B) memory for block-level score estimation without materializing the full score matrix.The fused kernel computes scores per query tile and key block, followed by global rescaling and row normalization.
- Mean correction: Mean correction restores discarded probability mass by representing each pruned block with pooled key and value statistics in the softmax numerator and denominator.Each surrogate contribution represents |B_J| tokens without per-token computation, addressing the otherwise neglected discarded mass under extreme sparsity.
- Mean correction: The correction’s mass surrogate is second-order accurate, while its numerator surrogate retains a first-order covariance term governed by within-block value dispersion.The per-block deviation is a covariance that vanishes under within-block independence and is bounded by dispersion via Cauchy–Schwarz.
- Mean correction: 6.2 points is the maximum FP8 accuracy loss without correction at 128K, while the corrected pipeline remains within 1.5 points of full attention in BF16.The correction is applied only to fully visible pruned blocks; salient blocks and diagonal-band selections follow the exact path.
- Sparse operator implementation: PackGQA shares staged KV blocks across grouped query heads, while warp specialization and pingpong pipelining overlap data movement, matrix multiplication, and softmax work.The operator uses a producer warpgroup for asynchronous loads and consumer warpgroups for wgmma, aligning execution with Hopper-oriented dense implementations.
- Sparse operator implementation: FP8 operands are dequantized on the fly inside online softmax, with layout-specific handling for the second GEMM and register-only probability-fragment rearrangement.Per-tensor scales support FP8-e4m3 execution while the softmax offset cancels in the probability ratio.
4. Experiments
Experiments on NVIDIA H20 GPUs evaluate FlashPrefill V2 across accuracy, operator efficiency, serving latency, open-loop load, chunked prefill, and implementation overhead. Across these settings, V2 maintains accuracy close to full attention while delivering substantial speedups, including under FP8 and production-oriented serving conditions.
- Accuracy Results: FlashPrefill V2 stays within 1.1 points of full attention on average across three models, with FP8 adding only 0.4–1.2 points of loss.At 128K, where fewer than 5% of blocks survive, the RULER gap remains within 1.8 points.
- Accuracy Results: FlashPrefill V2 achieves the highest average LongBench score among sparse operators on every model, staying within 0.9 points of full attention.Its advantage is strongest on retrieval-sensitive synthetic tasks, including English passage retrieval and counting.
- Efficiency Results: 47.3× FP8 and 27.2× BF16 speedups over FlashAttention-2 are reached at 128K, while speedups over the FA3/4-aligned dense kernel reach 30.5× and 17.5×.The density falls from about 70% at 4K to about 5% at 128K, contributing to the growing speedup.
- Efficiency Results: At 128K, FlashPrefill V2 reduces end-to-end TTFT by 2.1×–3.4× in BF16 and 3.7×–4.8× in FP8 across models and batch sizes.For Qwen3-30B-A3B at batch size 16, TTFT falls from 123.2 s to 36.2 s in BF16 and 25.5 s in FP8.
- Open-Loop Serving: In open-loop serving, FP8 reaches 0.88–1.34 req/s and lowers P50 TTFT to 2.8–33 s, while also improving decoding time and output throughput.The gains persist under load, though they are largest at the lowest arrival rate because faster prefills reduce queueing.
- Compatibility with Chunked Prefill: Chunked prefill reduces speedups because selection is rerun per chunk and mandatory tail blocks increase effective density, but V2 remains faster than the dense backend.At 16K chunks, most of the margin is recovered; the experiments recommend chunk sizes of at least 16K, though the supplied passage truncates the final recommendation.
- Comparison with the Production Kernel: FlashPrefill V2 is 6%–7% faster than the production HPC-Ops BSA kernel across sparsity levels in FP8 at 64K.The comparison uses comparable dense references, supporting the attribution of the gain to sparse execution.
- Ablation Study: Mean correction adds at most 12.2 ms in the dense limit and 3.6–4.5 ms at 90% sparsity, with relative overhead reaching 18% in BF16 and 27% in FP8 only at 90% sparsity.The correction recovers accuracy degradation that otherwise grows at longer contexts and is especially important under FP8 quantization.
5. Conclusion
FlashPrefill V2 advances sparse attention toward practical long-context prefilling through error correction, hardware-aligned execution, and serving-system integration. It preserves accuracy at extreme sparsity while delivering substantial GPU and end-to-end serving speedups.
- A zero-order mean correction term keeps accuracy loss within about one point on RULER and LongBench averages.At 128K on RULER, loss remains within 1.8 points while fewer than 5% of blocks are computed.
- The redesigned sparse operator aligns with FlashAttention-3/4 execution and supports FP8 execution for quantized deployment.It uses PackGQA memory access, warp-specialized producer-consumer pipelines, and intra-warpgroup pingpong overlap.
- 27.19× BF16 and 47.26× FP8 speedups over FlashAttention-2 are achieved at 128K context length.
- 30.49× speedup over an FA3/4-aligned dense baseline is achieved in FP8 at 128K context length.
- Native paged KV cache and continuous batching support enables SGLang integration without model-side changes.Under continuous batching at 128K, end-to-end time-to-first-token is reduced by up to 3.4× in BF16 and 4.8× in FP8.