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

Affiliations: WeChat · Institute of Automation, Chinese Academy of Sciences (CASIA) References: arXiv:2603.06199 · GitHub: qhfan/FlashPrefill

> Core idea: Make block-importance discovery and selection nearly zero-cost, prune the attention long tail with a dynamic threshold, and use index compaction to physically skip irrelevant blocks — achieving a 27.78x operator speedup at 256K context.

Background and Motivation

Prefill-stage bottleneck:

• Prefill must traverse the whole prompt to build the KV cache at \(O(L^2)\) complexity.
• Long contexts cause time-to-first-token (TTFT) to explode, degrading user experience.

Pain points in existing methods:

• High pattern-discovery latency: estimating block importance is itself computationally heavy.
• Expensive selection strategies: top-k sorting and top-p accumulation are hard to parallelize.
• Incomplete sparsity: under long-tail distributions, top-k/top-p force in many low-value blocks just to fill the quota.

Method

1. Instantaneous Pattern Discovery

• Block-level approximation with mean-proxy ranking.
• Low-variance assumption: tokens within a block are semantically similar.
• 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

\[\text{Threshold} = \alpha \times \max(\text{Score})\]
• Extremely cheap: only a single max-reduction is needed.
• Prunes the long tail entirely instead of padding to a fixed k or probability mass.
• Safety-net mechanisms: always keep Attention Sinks (first 256 tokens) and the Local Window (last 512 tokens).

3. Index Compaction for Physical Jumps

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

Experimental Results

Speed:

• 27.78x prefill operator speedup on Qwen3-30B at 256K context.
• Sparsity comparison at 256K (retained blocks): FlashPrefill 3.5%, FlexPrefill 8.4%, XAttention 18.5%.

Accuracy preservation:

• RULER (Qwen3-30B): 92.68 vs. 93.28 full attention (minor drop).
• InfiniteBench (Qwen2.5-7B): 24.93 vs. 23.87 full attention (slightly better).
• Needle-in-a-Haystack: retrieval nearly lossless from 2K to 256K.
• End-to-end TTFT is significantly reduced when integrated with vLLM.

Takeaway

Through approximate computation, dynamic thresholding, and physical index-based jumps, FlashPrefill achieves near-zero-cost sparsification of long-context prefill, substantially improving inference efficiency.