Block Sparse Flash Attention

Daniel Ohayon; Itay Lamprecht; Itay Hubara; Israel Cohen; Daniel Soudry; Noam Elata

cs.LGDec 7, 2025

Block Sparse Flash Attention

Authors: Daniel Ohayon, Itay Lamprecht, Itay Hubara, Israel Cohen, Daniel Soudry, Noam Elata

Organizations: Technion – Israel Institute of Technology, Haifa, Israel · Intel - Habana Labs · Intel – Habana Labs, Tel Aviv, Israel

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1 : Accuracy-latency trade-offs on RULER for 32K (top-left), 64K (top-right), and 128K (bottom). BSFA (blue) maintains high accuracy with consistent speedups, outperforming SpargeAttention (orange). Each point is a different sparsity level (Table 6 ).

Figure 2 : BSFA mechanism. (a) Per-block maxima

s_{\max}^{(i,j)}

are computed exactly within the FlashAttention tile. The center column visualizes the calibrated threshold tensor

T_{\ell,h,i}^{(k)}

(

k\!=\!6

shown), one value per query row, with hatched

-\infty

entries where fewer than

k

off-diagonal candidates exist. (b) Off-diagonal blocks are kept (green) if

s_{\max}^{(i,j)}\geq T_{\ell,h,i}^{(k)}

. Diagonal blocks (orange) always run; skipped blocks (grey) avoid the

\mathbf{V}_{j}

load and the

\mathbf{PV}

matmul. The highlighted row

i\!=\!14

traces a single gating decision across all three panels.

Method	Param.	Accuracy $\uparrow$	Speedup $\uparrow$
Dense FlashAttention-2	All	40.24%	1.00 $\times$
BSFA	$k$ =32	39.39% (-2.1%)	1.16 $\times$
	$k$ =64	39.78% (-1.1%)	1.13 $\times$
	$k$ =96	39.88% (-0.9%)	1.10 $\times$
	$k$ =128	40.03% (-0.5%)	1.08 $\times$
SpargeAttention	$\tau$ =0.4	30.06% (-25.3%)	1.09 $\times$

Table 1 : Main results on LongBench (Llama-3.1-8B). Accuracy is on the full 27-task benchmark (

\sim

4,750 samples) via lm-evaluation-harness, and speedup is end-to-end TTFT on the 10 longest narrativeqa samples (

\sim

65K tokens each). All methods on A100 unless marked

\star

(H100, since BLASST has no public A100 implementation). Sliding Window does not run as a sparse-attention kernel against the same baseline, so we report accuracy only. Bracketed numbers show relative degradation from the dense baseline.

Method	Param.	Accuracy $\uparrow$	$<$ 4K	4–8K	8–16K	16–32K	32–65K
Dense FlashAttention-2	All	40.24%	1.00 $\times$	1.00 $\times$	1.00 $\times$	1.00 $\times$	1.00 $\times$
BSFA	$k$ =64	39.78% (-1.1%)	0.96 $\times$	0.97 $\times$	1.01 $\times$	1.05 $\times$	1.19 $\times$
MInference	default	40.06% (-0.45%)	0.17 $\times$	0.31 $\times$	0.48 $\times$	0.58 $\times$	0.99 $\times$
FlexPrefill	$\gamma$ =0.9	35.11% (-12.8%)	0.71 $\times$	0.90 $\times$	1.05 $\times$	1.18 $\times$	1.59 $\times$
FlexPrefill	$\gamma$ =0.99	38.90% (-3.3%)	0.69 $\times$	0.80 $\times$	0.93 $\times$	0.99 $\times$	1.12 $\times$
XAttention	default	39.95% (-0.7%)	0.85 $\times$	0.90 $\times$	0.97 $\times$	1.01 $\times$	1.09 $\times$

Table 2 : End-to-end TTFT speedup on a LongBench mixed-length subset (160 samples spanning 295–65,461 tokens), broken down by sequence-length bucket. Speedup is reported against dense FlashAttention-2 within each bucket. Accuracy column is the full 27-task LongBench score from Table 1 , repeated here for ease of speedup-vs-accuracy correlation. Llama-3.1-8B on A100 unless marked

\star

(H100). SpargeAttention was not measured on this subset. Bold cells highlight regressions below dense (

<\!1.00\times

Param.	Accuracy $\uparrow$	E2E Speedup $\uparrow$	Attn-only Speedup $\uparrow$
SDPA	84.96%	1.00 $\times$	1.00 $\times$
$k$ =192	83.08% (-2.2%)	1.13 $\times$	1.38 $\times$
$k$ =256	83.39% (-1.8%)	1.09 $\times$	1.30 $\times$
$k$ =384	84.24% (-0.8%)	1.05 $\times$	1.19 $\times$

Table 3 : Attention-kernel-only versus end-to-end speedup at RULER 64K (BSFA, Llama-3.1-8B, A100). Accuracy is on RULER 64K with relative degradation from dense in brackets. Kernel-only timing isolates the attention call from work BSFA cannot accelerate (Q/K/V/O projections, RoPE, GQA expansion, MLP, the LM head).

Figure 3 : Accuracy vs. Time-to-First-Token (TTFT) speedup on the 64K Needle-in-a-Haystack Single-key task. BSFA retains 99% of baseline accuracy at $k{=}32$ with a $1.24\times$ speedup, exploiting that retrieval concentrates signal in few blocks. SpargeAttention is shown as a multi-operating-point reference. The remaining baselines (MInference, FlexPrefill, XAttention, BLASST) were not evaluated on this task and are compared on LongBench in Tables 1 – 2 .

Table 7

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Seq Length	Method	Sparsity Param.	Accuracy $\uparrow$	Pred./Meas. Density	Speedup $\uparrow$
32K	Dense FlashAttention-2	All	85.98%	1.00	1.00 $\times$
	BSFA	$k$ =64	84.94% (-1.2%)	0.24/0.28 $\pm$ 0.04	1.07 $\times$
		$k$ =96	85.52% (-0.5%)	0.35/0.38 $\pm$ 0.05	1.04 $\times$
		$k$ =128	86.00% (+0.0%)	0.45/0.47 $\pm$ 0.05	1.03 $\times$
		$k$ =192	86.42% (+0.5%)	0.62/0.62 $\pm$ 0.05	1.00 $\times$
	SpargeAttention	$\tau$ =0.6	75.62% (-12.1%)	-	0.91 $\times$

Appendix

Table 6 : Full RULER results across multiple sparsity levels. Accuracy column shows absolute accuracy with relative degradation from baseline in brackets. Pred./Meas. Density shows predicted (Eq. 3 ) vs. measured density during inference (predicted/measured

\pm

std). Speedup is end-to-end TTFT vs. dense FlashAttention-2 on the same samples.

Method	Param.	Speedup $\uparrow$
Dense FA-2	All	1.00 $\times$
BSFA	$k$ =192	1.13 $\times$
	$k$ =256	1.09 $\times$
	$k$ =384	1.05 $\times$
SpargeAttention	$\tau$ =0.5	1.03 $\times$
MInference	default	1.15 $\times$

Appendix

Table 7 : End-to-end TTFT speedup on RULER 64K (Llama-3.1-8B). All methods on A100 unless marked

\star

(H100). XAttention was not measured on RULER.

Explore similar work

FSA: An Alternative Efficient Implementation of Native Sparse Attention Kernel
Aug 25, 2025Ran Yan, Youhe Jiang, Zhuoming Chen +3Dynamic Sparse Attention Alternative
1. The Hong Kong University of Science and Technology · 2. Carnegie Mellon University
SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers
Jun 22, 2026Huzama Ahmad, Se-Young YunDynamic Sparse Attention Efficient Long-Context Inference
KAIST
Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference
Dec 18, 2025Dhruv Deshmukh, Saurabh Goyal, Nipun Kwatra +1Dynamic Sparse Attention LLM Inference Optimization
Microsoft Research India.

Aug 25, 2025, cs.DCJ/K move · Enter open · S save

FSA: An Alternative Efficient Implementation of Native Sparse Attention Kernel

Ran Yan, Youhe Jiang, Zhuoming Chen, Haohui Mai+2

1. The Hong Kong University of Science and Technology · 2. Carnegie Mellon University

Dynamic Sparse Attention Alternative

Open paper AskPDF arXiv

Recent advances in sparse attention mechanisms have demonstrated strong potential for reducing the computational cost of long-context training and inference in large language models (LLMs). Native Sparse Attention (NSA), one state-of-the-art approach, introduces natively trainable, hardware-aligned sparse attention that delivers substantial system-level performance boosts while maintaining accuracy comparable to full attention. However, the kernel implementation of NSA forces a loop order that is only efficient with a relatively large number of query heads in each Grouped Query Attention (GQA) group, whereas existing LLMs widely adopt a much smaller number of query heads in each GQA group -- such an inconsistency significantly limits the applicability of this sparse algorithmic advance. In this work, we propose Flash Sparse Attention (FSA), an alternative kernel implementation that enables efficient NSA computation across a wide range of popular LLMs with a varied, smaller number of heads in each GQA group on modern GPUs. Compared to vanilla NSA kernel implementation, our empirical evaluation demonstrates that FSA achieves (i) up to 3.5x and on average 1.6x kernel-level latency reduction, (ii) up to 1.25x and 1.09x on average end-to-end training speedup on state-of-the-art LLMs, and (iii) up to 1.36x and 1.11x on average for prefill-phase speedup in LLM generative inference. The source code is open-sourced and publicly available at https://github.com/Relaxed-System-Lab/Flash-Sparse-Attention.

Abstract

Explore similar work

FSA: An Alternative Efficient Implementation of Native Sparse Attention Kernel

SpotAttention: Plug-In Block-Sparse Routing for Pretrained Long-Context Transformers

Kascade: A Practical Sparse Attention Method for Long-Context LLM Inference