Research map / Architectures and efficiency / Attention and transformers
FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes FlexPrefill, a training-free sparse prefilling attention that per-head, per-input chooses between query-aware and vertical-slash patterns via Jensen-Shannon divergence and sets per-head budgets via cumulative attention thresholds.
Problem
Quadratic cost of attention during long-sequence LLM prefilling; existing sparse methods use fixed patterns or offline-searched patterns and ratios that do not adapt to each input and attention head.
Approach
Per attention head, estimate block-pooled attention from the last query block and compare it to the true distribution with square-root JS divergence; if the estimate is accurate use query-aware block selection, otherwise fall back to a vertical-slash pattern. Then select the smallest set of blocks/lines whose cumulative normalized attention exceeds a threshold gamma, giving each head and input its
Key results
- On RULER with Llama-3.1-8B, FlexPrefill (gamma=0.95) averages 89.18 vs 88.70 for full attention and 85.42 for MInference, with 2.43x attention speedup at 128k.
- On InfiniteBench, FlexPrefill averages 47.14 on Llama and 44.83 on GLM, close to full attention (48.06 and 41.98).
- Lowering gamma to 0.9 on Llama-3.1-8B gives a 3.49x speedup at 128k with average 88.38 on RULER.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use dynamic rather than fixed computational budgets across attention heads (ablation) | FlexPrefill · RULER | score and inference speed | Helped |
| Enable Query-Aware attention with an appropriate threshold (tau=0.1 rather than tau=0) (ablation) | FlexPrefill on Qwen · RULER Avg | score: 69.47 -> 70.75 | Helped |
| Raise the Query-Aware threshold to tau=0.3 from tau=0.1 (ablation) | FlexPrefill on LLaMA · RULER Avg | score: 89.18 -> 88.38 | Hurt |
| Change Triton block size from 128 to 64 (ablation) | FlexPrefill on GLM · RULER Avg | score: 89.34 -> 89.34 | No clear effect |
| Use middle rather than last representative queries for sparse pattern determination (ablation) | FlexPrefill on LLaMA-3.1 · RULER Avg | score: 89.18 -> 89.02 | No clear effect |
| Use middle rather than last representative queries for vertical-slash index selection (ablation) | FlexPrefill on LLaMA-3.1 · RULER Avg | score: 89.18 -> 38.18 | Hurt |
| Use query-wise index selection instead of flattening the estimated attention map (ablation) | FlexPrefill on LLaMA · RULER Avg | score: 89.18 -> 89.3 | No clear effect |
| Add a minimum computational budget of 1024 tokens rather than no minimum at gamma=0.9 (ablation) | FlexPrefill on LLaMA · RULER Avg | score: 84.78 -> 88.38 | Helped |
| Increase the minimum budget from 1024 to 2048 tokens at gamma=0.9 (ablation) | FlexPrefill on LLaMA · RULER Avg | score: 88.38 -> 88.82 | No clear effect |
| Cap maximum computational budget at 8k rather than leave it uncapped (ablation) | FlexPrefill on LLaMA · RULER Avg | score: 89.18 -> 86.98 | Hurt |
| Cap maximum computational budget at 32k rather than leave it uncapped (ablation) | FlexPrefill on GLM · RULER Avg | score: 89.34 -> 89.54 | Helped |
| Lower cumulative-attention threshold gamma from 0.95 to 0.85 (ablation) | FlexPrefill on LLaMA-3.1 · RULER Avg | score: 89.18 -> 87.10 | Hurt |
Surprising findings
- Sparse attention sometimes exceeds full attention on RULER (e.g., Qwen 70.75 vs 68.83 and GLM 89.34 vs 89.06), which the authors attribute to enhanced capability.
- Restricting the maximum per-head budget hurts Llama at 128k but maintains or improves GLM performance.
Limitations
- Speedup is measured on attention only, not end-to-end inference, and results depend on gamma and tau hyperparameters tuned per model.
- Evaluated on only four models and two benchmarks, with RULER and InfiniteBench being mainly synthetic-heavy.
- Fallback to vertical-slash relies on a heuristic JS threshold tau; a too-large tau misclassifies heads with inaccurate estimates.
Experimental setup
Datasets: RULER InfiniteBench
Benchmarks: RULER InfiniteBench
Metrics: accuracy attention latency (ms) speedup
Baselines: FlashAttention StreamingLLM MInference MoA InfLLM LM-Infinite HiP
Builds on: MInference StreamingLLM FlashAttention
Closest papers by meaning
- Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps (NeurIPS 2026 · similarity 0.92)
- QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference (ICML 2024 · similarity 0.91)
- MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention (NeurIPS 2024 · similarity 0.91)
- Star Attention: Efficient LLM Inference over Long Sequences (ICML 2025 · similarity 0.91)
- SparQ Attention: Bandwidth-Efficient LLM Inference (ICML 2024 · similarity 0.91)
- Efficient Low Rank Attention for Long-Context Inference in Large Language Models (NeurIPS 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.