PaperlayerView exact revision on arXiv

Source-linked AI summary

FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference

Xunhao Lai, Jianqiao Lu, Yao Luo, Yiyuan Ma, Xun Zhou

arXiv:2502.20766v1cs.LGcs.CL

TL;DR

Long-sequence LLM inference is bottlenecked by quadratic attention pre-filling and limited adaptability in prior sparse methods. FlexPrefill dynamically selects attention patterns and query-key indexes per input and head, and experiments report improved speed and accuracy while preserving or enhancing model performance across long-context tasks.

  • Problem

    Quadratic full-attention cost makes long-sequence pre-filling inefficient, while prior sparse methods have limited flexibility across inputs and attention heads.

  • Method

    FlexPrefill uses divergence-based Query-Aware or Vertical-Slash pattern selection and cumulative-attention index selection to adapt sparse computation per input and head.

  • Results

    FlexPrefill consistently preserves or enhances model performance while significantly improving computational efficiency across long-context models, benchmarks, tasks, and context lengths.

  • Takeaways & Limitations

    Dynamic pattern and budget allocation provides a better balance between inference speed and model effectiveness than static allocation across the reported scenarios.

  • Takeaways & Limitations

    The evaluation compares methods using sparse pre-filling but dense computation during decoding.

Abstract

from arXiv · show

Large language models (LLMs) encounter computational challenges during long-sequence inference, especially in the attention pre-filling phase, where the complexity grows quadratically with the prompt length. Previous efforts to mitigate these challenges have relied on fixed sparse attention patterns or identifying sparse attention patterns based on limited cases. However, these methods lacked the flexibility to efficiently adapt to varying input demands. In this paper, we introduce FlexPrefill, a Flexible sparse Pre-filling mechanism that dynamically adjusts sparse attention patterns and computational budget in real-time to meet the specific requirements of each input and attention head. The flexibility of our method is demonstrated through two key innovations: 1) Query-Aware Sparse Pattern Determination: By measuring Jensen-Shannon divergence, this component adaptively switches between query-specific diverse attention patterns and predefined attention patterns. 2) Cumulative-Attention Based Index Selection: This component dynamically selects query-key indexes to be computed based on different attention patterns, ensuring the sum of attention scores meets a predefined threshold. FlexPrefill adaptively optimizes the sparse pattern and sparse ratio of each attention head based on the prompt, enhancing efficiency in long-sequence inference tasks. Experimental results show significant improvements in both speed and accuracy over prior methods, providing a more flexible and efficient solution for LLM inference.

1 INTRODUCTION

Long-context inference makes attention pre-filling a bottleneck because full attention scales quadratically with sequence length. FlexPrefill addresses limited adaptability in prior sparse methods by dynamically tailoring patterns and computation to inputs and attention heads.

  • Full attention becomes inefficient during long-sequence pre-filling because its computational complexity grows quadratically with sequence length.
  • Predefined sparse patterns restrict flexibility and often require additional training or fine-tuning.
  • FlexPrefill dynamically adapts sparse attention patterns in real time to each input and attention head.
  • Its query-aware component switches between query-specific Diverse patterns and consistent Structured patterns using Jensen-Shannon divergence.
  • Experiments across state-of-the-art LLMs and long-context benchmarks show improvements in speed and accuracy while preserving or enhancing performance.

2 SPARSE ATTENTION

Attention sparsity varies across heads and inputs, making fixed patterns and sparsity ratios inadequate. FlexPrefill therefore targets adaptive pattern and ratio selection while balancing attention effectiveness against computation.

  • Sparse attention computes only selected query-key pairs to reduce computation relative to full attention.
  • The dynamic sparse-attention goal minimizes both the difference from full attention and the size of the selected subset.
  • Attention heads exhibit either Diverse query-dependent patterns or Structured patterns that remain consistent across queries.
  • Structured heads may allow a subset of the attention map to estimate the entire sparse index set.
  • Different prompt complexities require different sparsity ratios, with long-range dependencies favoring lower ratios and local dependencies permitting higher ratios.
  • FlexPrefill dynamically adjusts each sample’s sparse pattern and sparsity ratio.

3 METHOD

FlexPrefill determines an attention head’s sparse pattern and then selects the smallest query-key subsets whose cumulative normalized attention reaches a threshold. It uses divergence-based switching and pattern-specific index selection to allocate computation adaptively.

  • 3 METHOD: FlexPrefill combines Query-Aware Sparse Pattern Determination with Cumulative-Attention Based Index Selection.
  • Query-Aware Sparse Pattern Determination: The method classifies heads as Query-Aware for query-dependent patterns or Vertical-Slash for common structured patterns.
  • Query-Aware Sparse Pattern Determination: It compares estimated and true attention distributions using the square root of Jensen-Shannon divergence, switching patterns according to threshold τ.
  • Query-Aware Sparse Pattern Determination: When estimation is unreliable, FlexPrefill falls back to a Vertical-Slash pattern based on a subset of queries.
  • Cumulative-Attention Based Index Selection: For each query position, it selects the smallest index subset whose normalized attention scores reach cumulative threshold γ.
  • Cumulative-Attention Based Index Selection: Query-Aware selection ranks blocks, while Vertical-Slash selection ranks vertical and slash lines before extending them across the attention matrix.
  • Algorithm: The overall algorithm determines patterns, obtains sparse index sets, and computes sparse attention using Q, K, V, τ, and γ.

4 EXPERIMENT

FlexPrefill is evaluated on long-context models and benchmarks using sparse prefill implementations and comparisons with established attention methods. Results show preserved or improved task performance, favorable latency–quality trade-offs, and benefits from dynamic budgets and calibrated thresholds.

  • Method implementation: FlexPrefill dynamically selects sparse patterns and index sets before computing the final sparse attention output.The pipeline determines a pattern from threshold τ, selects indexes using the corresponding search procedure, and applies sparse attention.
  • Settings: Experiments use four long-context LLMs, including LLaMA, GLM, Yi, and Qwen, evaluated on RULER and Infinite Bench.The setup includes both instruct and pretrained models across synthetic and real-world long-context tasks.
  • Main results: FlexPrefill consistently preserves model performance across RULER context lengths while accelerating computation, whereas StreamingLLM and MInference show deterioration or suboptimal results in some settings.The reported RULER comparison covers multiple models and context lengths; detailed latency comparisons are provided separately.
  • Main results: On Infinite Bench, FlexPrefill preserves most retrieval and question-answering performance while remaining effective on complex mathematical and coding tasks.The comparison spans multiple models and tasks.
  • Ablations: Dynamic budgets improve performance and balance inference speed against model effectiveness compared with static allocation.The method also achieves better performance with lower latency than MInference in the reported performance–latency comparison.
  • Ablations: An appropriate τ strengthens performance without increasing computation, while overly large τ can misclassify inaccurately estimated heads and degrade results.A minimum budget threshold also enhances performance and prevents collapse at extremely high sparsity ratios.

5 RELATED WORKS

Related work addresses long-context modeling and inference acceleration through model adaptation, hardware-aware attention optimization, context reduction, and KV-cache techniques.

  • Long Context LLMs: Some long-context methods extend model context by collecting long-text datasets and persistently pre-training or fine-tuning models.These approaches target the model’s ability to handle longer contexts through additional training.
  • LLM Inference Acceleration: Inference acceleration methods optimize attention computation with hardware-aware algorithms such as FlashAttention and RingAttention.These methods improve the original attention computation rather than necessarily reducing the represented context.
  • LLM Inference Acceleration: Other approaches accelerate long-context inference by reducing context length or managing decoding KV caches through optimized computation, quantization, or token merging.The related work distinguishes prefill-oriented attention acceleration from decoding-phase KV-cache optimization.

6 CONCLUSION

FlexPrefill combines real-time sparse-pattern adaptation with cumulative-attention-based subset selection to reduce prefill computation while preserving or enhancing performance. Its dual formulation formalizes the trade-off between selected-subset size and retained attention mass.

  • 6 CONCLUSION: FlexPrefill dynamically optimizes sparse patterns and ratios for each attention head based on the input during long-sequence pre-filling.The method combines Query-Aware Sparse Pattern Determination with Cumulative-Attention Based Index Selection.
  • 6 CONCLUSION: The sparse-attention objective minimizes attention difference under a constraint on the size of the selected query-key subset.A tolerance rate γS controls the maximum selected-subset size.
  • 6 CONCLUSION: Maximizing normalized attention mass in the selected subset minimizes an upper bound on the error between full and sparse attention.The analysis bounds the contribution of omitted positions using the attention mass retained by the selected subset.
  • 6 CONCLUSION: Duality transforms the primal objective into minimizing subset size while requiring the selected subset’s normalized attention scores to meet threshold γa.The equivalence follows from strong duality and the zero duality gap of the convex relaxation.
  • 6 CONCLUSION: The resulting formulation establishes a theoretical trade-off between computational subset size and retained attention scores.This connection provides theoretical support for the sparse-attention optimization goal.

C DETAILED LATENCY OF DIFFERENT METHODS

On RULER, FlexPrefill is evaluated for single-attention-call latency at 64k and 128k tokens and for average latency across sequence lengths, with better performance and lower latency than competing methods.

  • C DETAILED LATENCY OF DIFFERENT METHODS: FlexPrefill achieves better performance while maintaining lower latency than competing methods in RULER latency comparisons.Table 3 covers various models and sequence lengths, including long sequences of 64k and 128k tokens.

D COMPARISON WITH ADDITIONAL BASELINES

FlexPrefill is compared with additional baselines on RULER using specified Llama model variants and consistently outperforms them in both performance and inference speed.

  • D COMPARISON WITH ADDITIONAL BASELINES: FlexPrefill consistently outperforms LM-Infinite, InfLLM, MoA, and HIP in both performance and inference speed on RULER.The comparisons use Llama-3-8B-Instruct-262k for three baselines and Meta-Llama-3.1-8B-Instruct-128k for HIP.

E PERFORMANCE-LATENCY TRADE-OFF WITH DIFFERENT γ

FlexPrefill exposes a speed-quality trade-off through γ and performs better with dynamic than static budget allocation. Additional experiments examine threshold, block-size, and representative-query choices.

  • E PERFORMANCE-LATENCY TRADE-OFF WITH DIFFERENT γ: Decreasing γ accelerates processing, whereas increasing γ preserves model quality, and FlexPrefill achieves better performance with lower latency than MInference.Figure 4 compares model performance with average prefill time for a single attention head.
  • E PERFORMANCE-LATENCY TRADE-OFF WITH DIFFERENT γ: Model performance is evaluated under different Query-Aware threshold τ values and context lengths.These settings are reported in Table 6 for multiple models on RULER.
  • E PERFORMANCE-LATENCY TRADE-OFF WITH DIFFERENT γ: Different Triton block sizes can be selected flexibly because block sizes of 64 and 128 do not significantly affect model performance.The comparison is reported on the RULER dataset.
  • E PERFORMANCE-LATENCY TRADE-OFF WITH DIFFERENT γ: Changing the representative subset for sparse-pattern determination has little effect, while changing the subset for vertical-slash index selection significantly reduces performance.The ablation compares representative query vectors taken from the sequence end versus the middle.

F.4 ALTERNATIVE IMPLEMENTATION FOR QUERY-AWARE INDEX SEARCH

FlexPrefill’s Query-Aware index selection offers a global implementation that preserves performance while improving implementation efficiency. Its adaptive budgets and sparse computation introduce overhead that is offset by savings relative to dense attention.

  • Alternative Query-Aware Index Search: The global index-selection approach chooses the minimum computational budget needed for cumulative attention scores to exceed threshold γ.A query-wise alternative instead requires selected key-block scores to exceed γ for each query block.
  • Alternative Query-Aware Index Search: The global implementation was selected because it simplifies the attention mechanism while maintaining performance parity with query-wise selection.Table 9 compares the alternative Query-Aware index-search implementations.
  • Budget Ablations: A minimum budget is unnecessary at high γ but significantly improves performance at smaller γ, when some heads select too few tokens.The ablation uses Llama-3.1-8B-Instruct and GLM-4-9B-Chat on RULER.
  • Budget Ablations: Maximum-budget caps can harm LLaMA performance, whereas GLM performance is maintained or improved under the same type of constraint.The effect of maximum computational budgets differs across models on RULER.
  • Complexity and Latency: FlexPrefill’s sparse attention computation costs approximately O(αn^2d), while representative score computation costs O(bnd), pattern search O(bn), and index construction O(n log n).Here α denotes the fraction of dense-attention computations performed, while b is block size, d hidden dimension, and n sequence length.
  • Complexity and Latency: Compared with dense attention’s O(n^2d), FlexPrefill adds approximately O(αn^2d) + O(n log n) + O(bnd), offset by sparsity savings.At shorter inputs, non-attention overheads are larger; as inputs grow, sparse-attention computation occupies a larger latency share.

J SPARSITY RATIO

FlexPrefill finds highly dynamic sparsity patterns and ratios across attention heads, samples, task types, and context lengths. Longer inputs generally exhibit higher sparsity ratios, while Query-Aware and Vertical-Slash patterns vary across heads and layers.

  • Sparsity Ratios: Different samples require varying sparsity rates and show inconsistent sparsity distributions across attention heads.Input length also affects sparsity, with longer inputs showing higher sparsity ratios.
  • Attention Patterns: Most attention heads use Vertical-Slash patterns, while fewer Query-Aware patterns occur mainly in the model’s first layer.The distribution varies with task type and context length.
  • Attention Patterns: Query-Aware heads contain diverse blocks that may depart from a specific pattern, although some still exhibit the Vertical-Slash pattern.The visualization covers Llama-3.1-8B-Instruct attention heads.
  • Sparsity Ratios: Longer contexts of 256k show higher overall sparsity ratios than shorter 64k contexts across the illustrated sample types.Figure 10 uses heatmaps in which darker colors indicate lower sparsity.
Loading 2502.20766v1…