Block Sparse Flash Attention
Daniel Ohayon ⋅ Itay Lamprecht ⋅ Itay Hubara ⋅ Israel Cohen ⋅ Daniel Soudry ⋅ Noam Elata
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.13\times$ end-to-end speedup on LongBench with only a $1.1\%$ accuracy drop, and up to $1.24\times$ on Needle-in-a-Haystack retrieval at a $1\%$ accuracy drop. The attention kernel itself accelerates by up to $1.38\times$. 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/Anonymous44414/Block-Sparse-Flash-Attention.
Chat is not available.
Successful Page Load