AI 推論系統
FlashPrefill V2 Brings Block-Sparse Attention to SGLang, Cutting 128K-Prefill TTFT by Up to 4.8×
Teams from Tencent WeChat and the Chinese Academy of Sciences rewrote the Hopper sparse-attention kernel, adding FP8, paged KV cache, continuous batching, and mean compensation for pruned blocks. In H20 tests, kernel-level speedups reached 47.26×, but end-to-end gains were smaller, and the current integration depends on the SGLang fork bundled with the project.
FlashPrefill V2 transforms the original FlashPrefill, which was geared more toward algorithm validation, into a prefill-attention backend suitable for long-context serving stacks. The system first uses block-level approximate scores to select important K/V regions, then performs exact attention only on those blocks. Unselected blocks are not simply zeroed out; instead, their mean K/V values are incorporated into the online softmax as a zeroth-order compensation term, reducing output errors at high sparsity levels. The indexing stage generates both CSR sparse indices and compensation statistics in a single pass, avoiding Top-k sorting.
The kernel was rewritten in CUDA/CuTe for Hopper SM90, using a PackGQA memory layout, warp specialization, and a producer-consumer ping-pong pipeline, with support for both BF16 and FP8. Unlike many research kernels that accept only contiguous K/V tensors, it can read from a paged KV cache and handle variable-length requests and continuous batching. The team also added a standalone prefill backend to an SGLang 0.5.10 fork, while decoding can continue to use FA3.
In kernel benchmarks on four NVIDIA H20 GPUs with a batch size of 4 and a 128K context, FP8 and BF16 achieved maximum speedups of 47.26× and 27.19×, respectively, over FlashAttention-2. Even when compared against the stronger FA3/4-aligned dense kernel, the corresponding gains were still 30.49× and 17.54×. End-to-end results were more modest: the paper reports up to a 4.8× improvement in time to first token at 128K. Open-loop tests also showed that shortening long prefills reduces the time they block the decoding of other requests.
Engineering teams should still avoid treating kernel-level speedups as equivalent to serving-level speedups. At 4K with BF16, sparsity is only about 30%, leaving performance roughly on par with dense kernels. Chunked prefill reruns indexing and increases effective density, narrowing the gains. The public implementation supports only SM90 GPUs such as the H20 and H100, and SGLang support is provided as a bundled full source tree rather than an officially upstreamed backend. The next points to watch are upstream integration, reproducibility across different Hopper GPU models, and whether long-context quality remains stable across more models and real-world traffic.