Organizations: Department of Electronic Engineering, Tsinghua University, Beijing, China · Intern at Tongyi Lab, Alibaba Group · Tongyi Lab, Alibaba Group, Shanghai, China · University of Science and Technology of China, Hefei, China · Shanghai Artificial Intelligence Laboratory, Shanghai, China
Long-context inference has become a defining capability of large language models, but exact dense attention remains costly due to its quadratic scaling with sequence length. Block-sparse attention offers a hardware-friendly alternative by routing each query block to a small set of relevant key blocks, yet accurate training-free block routing remains difficult. Existing routers often pool post-RoPE token representations, which entangles semantic aggregation with RoPE-induced geometry and attenuates local positional cues through high-frequency phase cancellation. To resolve this mismatch, we propose Semantic-Geometric Decoupled Routing, a training-free block routing framework that shifts semantic aggregation to the pre-RoPE space and reconstructs geometric bias with an offline structural prior and relative block distances. This decomposition yields an explicit closed-form block routing score without token-level search or post-hoc calibration. Experiments on long-context text and video tasks show that our method approaches full-attention accuracy across 4K--128K contexts, keeps routing overhead below 3.4 ms, and achieves a 5.03× speedup over FlashAttn at a 128K context length.
Figures & tables
Figure 1: RMS energy decay across frequency dimensions for different pooling methods. Traditional Post-RoPE pooling (blue) causes severe high-frequency energy loss, shown by the shaded area. Our Pre-RoPE pooling (red) mitigates this loss, closely matching the original token energy.
Method
Block-only
4K
8K
16K
32K
64K
128K
Avg.
Llama-3.1-8B-Instruct
Flash Attn
✘
95.42
94.38
93.38
87.98
84.72
77.77
88.94
StreamLLM
✔
95.34
67.86
54.29
39.02
21.53
21.01
49.84
MInference
✘
95.43
94.46
93.42
87.22
83.07
71.04
87.44
FlexPrefill
✘
93.80
92.44
93.28
87.92
84.74
72.41
87.43
XAttention
✘
95.17
94.3
93.28
89.06
82.31
70.52
87.44
Table 1: Performance comparison on RULER across different context lengths. “Block-only” denotes methods operating ONLY on block-level features. Bold scores indicate the best results, and underlines denote the second best among sparse attention baselines. Our method achieves competitive performance across both short and long context scenarios compared to strong sparse baselines.
Method
Block-only
Single-Doc QA
Multi-Doc QA
Summarization
Few-shot Learning
Code
Synthetic
Avg.
Llama-3.1-8B-Instruct
Flash Attn
✘
47.51
43.28
25.90
45.92
18.01
68.18
41.47
StreamLLM
✔
34.89
27.47
24.81
46.67
17.18
26.24
29.54
MInference
✘
47.42
42.54
25.85
45.58
17.84
67.60
41.14
FlexPrefill
✘
46.13
41.49
25.85
46.63
17.68
25.61
33.90
XAttention
✘
45.89
41.56
26.18
45.86
19.24
59.32
39.68
Table 2: Performance comparison on LongBench across different task categories. Although MInference performs SOTA, it needs extra heuristic search. Our method shows competitive performance with higher efficiency and matches Prism using only block features.
Method
4K
8K
16K
32K
64K
128K
Avg.
FlashAttn
27.01
25.84
25.19
24.81
24.42
24.22
25.25
StreamLLM
28.77
29.07
29.55
30.09
30.71
33.37
30.26
MInference
27.02
25.93
25.44
25.30
25.18
25.93
25.80
FlexPrefill
29.63
28.25
27.33
26.98
26.44
26.04
27.45
XAttention
27.02
25.79
25.04
24.56
24.15
23.97
25.09
Prism
27.10
25.89
25.18
24.73
24.37
24.23
25.25
Table 3: Language modeling performance across different context lengths, measured by perplexity (PPL). Lower PPL indicates better performance.
Method
VideoMME
LVB
Short
Med.
Long
Overall
Overall
FlashAttn
79.89
70.67
63.11
71.22
65.00
StreamLLM
64.33
50.11
50.67
55.03
50.93
MInference
79.44
70.00
62.44
70.63
61.48
FlexPrefill
77.67
70.67
62.67
70.34
64.10
XAttention
79.22
69.78
63.44
70.81
64.25
Table 4: Performance comparison on VideoMME and LongVideoBench (LVB) datasets, reporting multiple-choice accuracy across various video duration subsets.
Figure 2: Prefill latency (bars, left axis) and speedup (lines, right axis) from 8K to 128K sequence lengths. Our method achieves the lowest latency and highest speedup across all settings, reaching a 5.03 × acceleration over FlashAttn at 128K. Although MInference attains comparable speed, its accuracy degrades at 128K, as shown in Tab. 1 .
Figure 3: Estimation time overhead comparison. Our method keeps setup time near zero (under 3.4 ms) across all lengths, showing no major growth as sequences grow.
Figure 4: Visual comparison of block routing scores. (a) Pure semantics loses positional awareness. (b) Post-RoPE dilutes context via high-frequency interference. (c) Center RoPE overly isolates the diagonal. (d) Ours perfectly restores a robust diagonal and natural local context using structural priors.
Variant / Configuration
4K
8K
16K
32K
64K
128K
Avg.
Ours (Full Model)
94.97
92.05
88.72
87.67
79.02
74.11
86.09
w/o Offline Profiling ( Λ^f,θ^f=1,0 )
88.34
84.40
82.10
79.06
72.72
66.63
78.87
w/o Macro Geometry ( ωfBΔM=0 )
84.92
78.96
75.06
71.01
69.51
63.58
73.84
w/o Online Phase ( Φf=0 )
86.27
82.12
76.67
76.24
67.19
61.17
74.94
w/o All Positional Info (Semantic Only)
76.64
70.60
66.51
64.13
66.63
59.03
67.26
Table 5: Ablation study on the RULER benchmark using Qwen-3-8B across various context lengths (from 4K to 128K). We report the comprehensive metric accuracy for our model and four ablation variants.
Modern large language models increasingly require long contexts for reasoning and multi-document tasks, but attention's quadratic complexity creates a severe computational bottleneck. We present Block Sparse Flash Attention (BSFA), a drop-in replacement that accelerates long-context inference while preserving model quality. Unlike methods that predict importance before computing scores, BSFA computes exact query-key similarities to select the top-k most important value blocks for each query. By comparing per-block maximum scores against calibrated thresholds, we skip approximately 50% of the computation and memory transfers for pruned blocks. Our training-free approach requires only a one-time threshold calibration on a small dataset to learn the per-layer and per-head attention score distributions. We provide a CUDA kernel implementation that can be used as a drop-in replacement for FlashAttention. On Llama-3.1-8B, BSFA achieves up to 1.13x end-to-end speedup on LongBench with only a 1.1% accuracy drop, and up to 1.24x on Needle-in-a-Haystack retrieval at a 1% accuracy drop. The attention kernel itself accelerates by up to 1.38x. We compare BSFA against five recent sparse attention baselines (SpargeAttention, MInference, FlexPrefill, XAttention, and BLASST), and verify the method on Qwen2.5-7B and on A6000 and H100 GPUs. The implementation is available at https://github.com/Danielohayon/Block-Sparse-Flash-Attention.
Daniel Ohayon, Itay Lamprecht, Itay Hubara +3
Technion – Israel Institute of Technology, Haifa, Israel · Intel - Habana Labs · Intel – Habana Labs, Tel Aviv, Israel
Long contexts have become standard in pretrained LLMs, yet they remain expensive to run: prefill compute grows quadratically with sequence length, and every decode step re-reads a key-value cache that grows linearly with it. Sparse attention cuts these costs by attending only to a relevant subset of past tokens, but selecting that subset is itself expensive. We present SpotAttention, a lightweight selector that attaches to a frozen pretrained transformer and learns by KL distillation to estimate its attention distribution. The selector picks the top-K keys each query attends to, and because its estimate is a calibrated distribution, a dual top-p rule reads the per-query, per-layer budget directly from it. Across Qwen3 (dense, 4B-32B) and Qwen3.5 (hybrid linear/full attention, 4B-9B), SpotAttention matches dense accuracy at contexts up to 128K tokens, eight times the training length. Decode at L=128K runs 3.9x faster than FlashAttention and 1.8x faster than Twilight, the strongest training-free baseline. Quantizing the selector's K-cache to INT4 or FP4 microscale shrinks it 3.5x at no accuracy cost.
Long-context large language model inference is increasingly limited by prefill, where dense self-attention processes the entire prompt before generation begins. Sparse block selection can reduce this cost, but a block centroid may hide a highly relevant token among many irrelevant ones. We call this failure mode mean dilution and propose RBS-Attention, a training-free sparse-prefill method with two complementary selection branches. A centroid base branch captures average relevance, while a rescue branch uses the maximum key-block radius and its prompt-, layer-, and head-dependent distribution to identify blocks at risk of underestimation. Independently thresholding the two branches and combining their masks controls the contribution of rescue blocks while preserving regular block-sparse FlashAttention execution. On H100 GPUs, RBS-Attention achieves 20.65× standalone prefill-attention speedup, 11.92× vLLM prefill-attention speedup, and 5.97× end-to-end time-to-first-token speedup at 128K on Qwen3-30B-A3B-Instruct-2507-FP8. On the dense Qwen3-32B model, it obtains 88.65 overall RULER accuracy versus 89.52 for dense attention; LongBench-v2, InfiniteBench, and Video-MME provide additional quality evaluation. Supporting experiments measure actual retention, compare selectors at matched density, and characterize block-size, threshold, and memory behavior. Together, these results support radius-adaptive dual-branch selection as an effective approach to long-context prefill.
Chuxu Song, Jiuqi Wei, Zhencan Peng
Department of Computer Science, Rutgers University · OceanBase, Ant Group