FlashPrefill: Near-Zero-Cost Sparsity Discovery for Ultra-Fast Long-Context Prefilling

FlashPrefill: Instantaneous Pattern Discovery and Thresholding for Ultra-Fast Long-Context Prefilling

• Authors: WeChat · Institute of Automation, Chinese Academy of Sciences
• arXiv:2603.06199 | GitHub: qhfan/FlashPrefill

Core Idea

FlashPrefill makes "finding important blocks" and "selection strategy" nearly zero-cost, uses dynamic thresholds to prune the attention long tail, and compacts indices to enable physical jumps — achieving a 27.78x operator speedup at 256K context.

Motivation

• Prefill must traverse the whole prompt to build the KV cache with \(O(L^2)\) complexity.
• Long contexts cause explosive growth in time-to-first-token (TTFT) and a poor user experience.

Pain points of existing sparse-attention methods:

1. High pattern-discovery latency — estimating block importance is itself compute-heavy. 2. Expensive selection strategies — Top-k sorting and Top-p accumulation are hard to parallelize. 3. Incomplete sparsity — under long-tail distributions,凑 enough K blocks / P probability introduces many useless blocks.

Method

1. Instantaneous Pattern Discovery

• Block-level approximation with mean-proxy ranking, under a low-variance assumption (blocks are semantically similar internally).
• Key/Query pooling plus global reweighting produces a global "attention map".
• Data movement drops from \(L \times L/B\) to \((L/B)^2\).

2. Dynamic Threshold Pruning

Threshold = α × max(Score)

• Extremely light: only one max-reduction needed.
• Prunes the long tail entirely instead of padding to a fixed count.
• Safety net: always keep Attention Sinks (first 256 blocks) and the Local Window (last 512).

3. Index Compaction → Physical Jumps

• Problem: logical skipping (if mask=0) still iterates all blocks with heavy branch overhead.
• Solution: compact valid block indices into a contiguous list; the kernel's inner loop only visits valid indices. Logical skipping becomes physical jumping, with more concentrated memory access and far fewer loop iterations.

Results

• Speed: 27.78x prefill operator speedup (Qwen3-30B @ 256K).
• Sparsity at 256K (blocks retained): FlashPrefill 3.5% vs. FlexPrefill 8.4% vs. XAttention 18.5%.
• Accuracy:
• RULER (Qwen3-30B): 92.68 vs. 93.28 full attention (minor drop)
• InfiniteBench (Qwen2.5-7B): 24.93 vs. 23.87 (slightly better)
• Needle-in-a-Haystack: retrieval capability from 2K to 256K is nearly lossless.
• End-to-end TTFT is significantly reduced when integrated into vLLM.

Summary

FlashPrefill achieves near-zero-cost sparsification for long-context prefilling through approximate computation, dynamic thresholds, and physical jumps, greatly improving inference efficiency.