FlashPrefill V2: Block-Sparse Prefill Attention for Long-Context LLM Serving
Qihang Fan, Huaibo Huang, Zhiying Wu, Bingning Wang, Ran He
cs.CL
2026-08-20
FlashPrefill V2 adds mean correction, a Hopper sparse kernel, and paged KV, then plugs into SGLang. At 128K on H20 it is 47.26× vs FA2 in FP8, about 1 RULER point down.
Prefill attention is still quadratic. Training-free sparse attention estimates scores, keeps the salient blocks, and skips the rest. Most of those methods stay prototypes: aggressive sparsity wrecks accuracy; the kernel still sits on FlashAttention-2 and underuses Hopper TMA and async pipelines; a contiguous KV layout will not plug into paged cache and continuous batching in vLLM or SGLang.
The same group's FlashPrefill already did instantaneous pattern discovery and max-based dynamic thresholding, so it does not pay for a global Top-k / Top-p sort. It still fails those three production checks. FlashPrefill V2's job is narrower: turn block-sparse prefill into an attention backend that a serving stack can actually load.
Prefill stays two-stage. Probe queries and mean-pooled keys score each query tile against key blocks. The keep threshold is \(\alpha\) times that tile's peak score (default \(\alpha=0.1\)), plus forced sink, local window, and recent blocks. The \(L \times (L/B)\) score matrix is never materialized; memory drops from \(O(L^2/B)\) to \(O((L/B)^2)\).
Under extreme sparsity the discarded softmax mass is no longer tiny. Mean correction injects a zero-order stand-in for every pruned block: the block means \(\bar kJ, \bar vJ\) contribute \(|BJ| e^{\bar sJ} \bar vJ\) to the numerator and \(|BJ| e^{\bar sJ}\) to the denominator. Mass is second-order accurate; the numerator still carries a within-block covariance. Max-based thresholding caps the mass share of each pruned block, which is why the error stays bounded. Inside the kernel a corrected block is one extra iteration with a \(\log|BJ|\) shift on the logit, reusing the same MMA and online softmax.
The Hopper-aligned sparse operator is what makes this deployable. PackGQA folds the \(g\) query heads of one KV group into a single tile so each staged KV block is consumed once, and the sparse index is stored per KV head. A persistent kernel splits a producer warpgroup (TMA / cp.async for K/V) from two consumer warpgroups (wgmma), with QK GEMM, PV GEMM, and online softmax in pingpong overlap. Selected blocks are walked backwards from a CSR list; producer and consumer compute the same sequence, so skipping a block needs no extra sync. Addresses go through the page table. A varlen scheduler enumerates (batch, head, tile) work units from per-request lengths, which is continuous batching. The FP8 path dequantizes on the fly and maps probabilities into e4m3's [0, 256] range, cancelled in the softmax ratio. Decode has one query token per step, so it falls back to dense attention.
SGLang integration swaps only the extend (prefill) backend. Model definition, KV layout, and scheduling stay put. Index workspace is cached per stream so the steady state does not hit the host.
All numbers are on NVIDIA H20 with Llama-3.1-8B-Instruct, Qwen3-4B-Instruct-2507, and Qwen3-30B-A3B-Instruct-2507. One config for every accuracy run: block size 128, 256 sink tokens, a 512-token window, \(\alpha=0.1\). Operator latency is a single H20. End-to-end TTFT is one SGLang server, TP=4 on four H20s, with decode fixed to an FA3/4-aligned dense kernel.
| Setting | Metric | FlashPrefill V2 | Baseline |
| Qwen3-30B-A3B, 128K, BF16 | Op vs FA2 | 27.19× | V1 is 18.67× |
| Same, FP8 | Op vs FA2 | 47.26× | 30.49× vs FA3/4-aligned dense |
| Three-model RULER avg | vs Full | -0.29 to -1.03 pts | 5% density at 128K |
| Qwen3-30B, BS=16, 128K | TTFT | 36.21 s / 25.51 s (FP8) | FA3/4 is 123.23 s |
RULER averages for V2 are 87.79 / 86.23 / 91.76 against Full at 88.82 / 87.06 / 92.05. At 128K, density falls to 4.6%–4.9% and V2 stays within 1.8 points of Full. Llama FP8 at 128K is the ugly cell: 73.82 down to 67.78, measured with online quantization and no corrected weights. On LongBench's 21 tasks, V2 is the best sparse method (49.31 / 46.96 / 50.73), 0.45–0.90 points behind Full.
Drop mean correction on Qwen3-4B RULER and BF16 loses 0.46 points on average; FP8 loses 2.33, and 6.2 at 128K. The extra latency at 64K and 90% sparsity is about 3.6–4.5 ms. Pushing \(\alpha\) from 0.1 to 0.2 at 64K FP8 cuts density from 9.2% to 5.2%; with correction the score is still 80.12 against Full at 82.81.
In open-loop serving (Poisson arrivals, mixed 4K–128K prompts) FA3/4 request throughput caps at 0.31–0.37 req/s. V2 roughly doubles that; FP8 reaches 0.88–1.34 req/s. Chunked prefill eats the speedup because every chunk re-runs selection and short chunks inflate density with mandatory tail blocks. The paper's practical advice is a chunk of at least 8K.
This is sparse attention moving from a paper number into a serving framework, and the move is mostly systems work, not a new block-picking heuristic. H20 is a common inference GPU in production. 47× is wall-clock against FlashAttention-2. The fairer dense baseline is their own FA3/4-aligned kernel, where 128K FP8 is still 30.49×, and end-to-end TTFT improves by up to about 4.8×. Near 4K the density is still around 70%, and BF16 is roughly even with FA3/4. The win shows up when context is long and attention dominates prefill.
Native paged KV and continuous batching, plus a drop-in SGLang backend, is more useful than another CUDA kernel that assumes a contiguous layout. Decode stays dense, so this attacks the prefill wall, not the full generate path.
There is no Limitations section in the paper. Every number is H20; A100, H100, and Blackwell are absent, and the kernel edge may shrink on another architecture. The speed tables use needle-in-a-haystack inputs, whose attention maps are more structured than messy real documents, so 5% density may not transfer to a repo or multi-hop retrieval. Llama FP8 accuracy carries an asterisk: online quantization without corrected weights, and the 128K drop is steeper than BF16. Mean correction is zero-order; if logits and values covary inside a block, a first-order residual remains. The authors lean on low intra-block variance and do not stress high-variance maps separately. Small chunks, which production engines like for TPOT, claw back the speedup. The repo is public. Whether a given model and quant scheme drop in cleanly is still an integration question.