Source-linked AI summary

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

Qihang Fan, Huaibo Huang, Zhiying Wu, Juqiu Wang, Bingning Wang, Ran He

arXiv:2603.06199v1cs.CLcs.AI

TL;DR

Long-context prefilling is hindered by quadratic attention and sparse methods' search latency or insufficient sparsity. FlashPrefill combines instantaneous block-based pattern discovery with Max-based Dynamic Thresholding, achieving substantial speedups while maintaining model performance. Evaluations across models and benchmarks demonstrate improved prefill efficiency, including 27.78× operator speedup at 256K and 1.71× at 4K.

  • Problem

    Long-context prefilling faces quadratic attention overhead, while existing sparse methods incur search latency and insufficient sparsity from sorting, accumulation, and long-tail distributions.

  • Method

    FlashPrefill combines instantaneous pattern discovery with block-approximation kernels and Max-based Dynamic Thresholding to identify salient patterns without sorting or accumulation.

  • Results

    FlashPrefill improves prefill efficiency across LLMs, VLMs, and benchmarks, reaching 27.78× speedup at 256K and 1.71× at 4K while maintaining model performance.

  • Takeaways & Limitations

    FlashPrefill provides efficient long-context prefilling across varying sequence scales while preserving nearly identical model performance with negligible accuracy loss.

  • Takeaways & Limitations

    Top-k and Top-p comparisons are constrained by long-tail attention distributions that can force inclusion of insignificant blocks and limit pruning efficiency.

Abstract

from arXiv · show

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. While various sparse attention mechanisms have been explored, they typically suffer from either significant search latency or insufficient sparsity. In this paper, we propose FlashPrefill, a framework enabling ultra-fast prefilling via instantaneous pattern discovery and thresholding. FlashPrefill leverages a fast block-searching technique to simultaneously locate dynamic vertical, slash, and block-sparse attention patterns. Crucially, it introduces a dynamic thresholding mechanism that bypasses the prohibitive overhead of sorting or accumulating attention scores while effectively eliminating the long-tail distribution to enhance sparsity. Extensive evaluations demonstrate that FlashPrefill achieves a substantial leap in efficiency, delivering an unprecedented 27.78x speedup on 256K sequences. Notably, unlike existing methods that incur efficiency degradation on shorter contexts, FlashPrefill maintains a 1.71x speedup even at a 4K context length, demonstrating its robustness and practical utility across varying sequence scales.

1. Introduction

FlashPrefill targets long-context prefill bottlenecks caused by quadratic attention and costly sparse-pattern selection. It combines instantaneous pattern discovery with Max-based Dynamic Thresholding and reports strong speedups while preserving model accuracy.

  • Quadratic self-attention creates prohibitive overhead for long-context sequences, especially during compute-intensive prefilling.
  • Existing sparse attention methods incur latency from coarse score estimation, sorting, cumulative summation, and incomplete pruning of long-tail tokens.
  • FlashPrefill identifies vertical, slash, and block-wise attention patterns through instantaneous pattern discovery optimized with block-level approximation.
  • Max-based Dynamic Thresholding avoids exhaustive sorting while filtering redundant blocks and mitigating long-tail attention distributions.
  • 27.78× speedup is achieved at 256K sequences, while 1.71× speedup remains at 4K and maximum end-to-end TTFT speedup reaches 7.22×.
  • FlashPrefill maintains nearly identical model performance with negligible accuracy loss on the Qwen3 Needle In A Haystack evaluation.

2. Related Works

Related work addresses attention's long-context cost through sparse attention mechanisms that reduce computation. These approaches include methods requiring training and methods that select salient structures before fine-grained attention.

  • Attention captures dependencies in Transformer models, but its quadratic complexity creates substantial computational overhead in long-context scenarios.
  • Sparse attention mechanisms have been proposed to reduce the computational cost of the attention layer.
  • A portion of sparse-attention approaches requires explicit model training.

3. Method

FlashPrefill discovers vertical, slash, and block-sparse patterns using block-level approximations, then accelerates selection with dynamic thresholding. Its fused kernels reduce memory overhead while preserving effective pattern ranking and improving sparsity.

  • Instantaneous Pattern Discovery: FlashPrefill uses uniformly distributed query probes to identify vertical, slash, and block-sparse attention patterns concurrently.Vertical patterns reflect column-wise invariance, while slash patterns exhibit translational symmetry across the sequence.
  • Instantaneous Pattern Discovery: Average-pooled key blocks serve as proxies for token-level attention because localized tokens exhibit semantic coherence and redundant attention patterns.The probing score uses block-level aggregation to estimate each block’s significance.
  • Instantaneous Pattern Discovery: Low intra-block variance preserves block rank ordering when the geometric mean acts as a monotonic proxy for the arithmetic mean.A secondary averaging operation across queries within each block produces aggregate block-level significance.
  • Block Approximation for Kernel Optimization: The fused 2D-reduction kernel replaces explicit compute-then-pool processing with a single-pass block-wise approximation, reducing memory from O(L·L/B) to O((L/B)^2).The approximation bypasses the massive L × (L/B) intermediate matrix and reduces global memory traffic.
  • Comparison with Previous Methods: On RULER with top-8 block selection, the block-approximation method achieves the best balance between efficiency and effectiveness.Mean pooling causes performance degradation, whereas the non-approximated method incurs substantial memory-access overhead.
  • Max-based Dynamic Thresholding: Max-based dynamic thresholding avoids sorting and cumulative computations by deriving a pruning threshold from each query block’s peak score.Blocks below the threshold are discarded, mitigating long-tail effects and enabling higher sparsity than fixed Top-k or Top-p selection.

4. Experiments

Experiments evaluate FlashPrefill across LLMs, VLMs, benchmarks, density, end-to-end TTFT, and pattern-selection overhead. The results report broad performance advantages, substantial speedups, reduced density, and faster token identification.

  • Evaluation setup: FlashPrefill was evaluated on RULER, InfiniteBench, and VideoMME, including both LLM and VLM settings.The study also reports ablations of individual modules and uses multiple model architectures, including dense and Mixture-of-Experts models.
  • Benchmark results: FlashPrefill consistently outperforms other baselines on InfiniteBench across dense and Mixture-of-Experts models.
  • Benchmark results: 22.67×, 16.87×, and 18.67× speedups are achieved at 128K on three representative models in RULER, substantially exceeding existing methods.
  • Benchmark results: FlashPrefill achieves superior results on VideoMME compared with existing sparse attention approaches.
  • Density and sparsity: FlashPrefill substantially reduces attention density as sequence length increases while mitigating long-tail effects.The density study reports reduced density relative to the other two methods as effective information becomes relatively sparser at longer sequences.
  • End-to-end efficiency: 5.02× TTFT speedup is achieved at 128K on Qwen3-30B-A3B-Instruct-2507 after integration into vLLM.The end-to-end measurement covers three language models, including dense and Mixture-of-Experts architectures, and compares against Full Attention.
  • Pattern discovery and thresholding: FlashPrefill identifies important tokens significantly faster than other methods during combined pattern discovery and thresholding.
  • Pattern discovery and thresholding: Max-based Dynamic Thresholding significantly reduces computational density while preserving most model performance.Compared with Top-k and Top-p approaches, it removes long-tail effects and achieves more thorough sparsity in long-context settings.

5. Conclusion

FlashPrefill accelerates long-context LLM prefilling through instantaneous pattern discovery, block approximation, and Max-based Dynamic Thresholding. Evaluations across LLMs, VLMs, and benchmarks report improved prefill efficiency while maintaining superior performance.

  • FlashPrefill accelerates long-context prefilling with instantaneous pattern discovery and a block-approximation kernel implementation.
  • Max-based Dynamic Thresholding avoids sorting and accumulation while mitigating long-tail distributions to achieve higher sparsity.
  • Evaluations across LLMs, VLMs, and diverse benchmarks show improved prefill efficiency with superior performance maintained.

A. Hyperparameter Configurations

FlashPrefill calibrates one dynamic-threshold hyperparameter using computational density at 4K, while keeping block and structural-retention settings fixed across models.

  • α is the sole FlashPrefill hyperparameter requiring adjustment for the dynamic threshold.It is calibrated using computational density observed at a 4K sequence length.
  • The threshold factor α is regulated to maintain approximately 70% computational density for 4K sequences.
  • A uniform block size of 128 is used across all models.
  • Attention sinks and a local window are explicitly retained with sizes of 256 and 512 tokens, respectively.These blocks are retained in addition to those selected through Max-based Dynamic Thresholding.

B. Detailed Implementation

FlashPrefill is organized into instantaneous pattern discovery, Max-based dynamic thresholding with block selection, and block-sparse attention. The implementation uses block approximations, dynamic and structural masks, compressed indices, and sparse online-softmax computation.

  • Framework stages: FlashPrefill comprises Instantaneous Pattern Discovery, Max-based Dynamic Thresholding and Block Selection, and Block Sparse Attention.The three stages are formalized in Algorithms 1, 2, and 3.
  • Instantaneous Pattern Discovery: Instantaneous Pattern Discovery pre-pools key blocks and computes block-level importance through fused block approximation.The procedure loads query tiles, compares them with causally valid pooled key blocks, and applies global normalization.
  • Dynamic Thresholding and Block Selection: Max-based Dynamic Thresholding selects blocks whose scores reach α times the per-query-block maximum.Sink, window, and causal masks are combined with the score mask before active-block indices are compressed.
  • Block Sparse Attention: Block Sparse Attention iterates only over identified salient blocks and applies causal masking where required.The kernel uses online softmax updates and final normalization to produce the output.
Loading 2603.06199v1…