Block Sparse Flash Attention

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× end-to-end speedup on LongBench with only a 1.1% accuracy drop, and up to 1.24× on Needle-in-a-Haystack retrieval at a 1% accuracy drop. The attention kernel itself accelerates by up to 1.38×. 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://eur01.safelinks.protection.outlook.com/?url=https%3A%2F%2Fprotect.checkpoint.com%2Fv2%2Fr02%2F___https%3A%2F%2Fgithub.com%2FDanielohayon%2FBlock-Sparse-Flash-Attention___.YzJlOnRlY2huaW9uOmM6bzoxMDY4ZWQ0MWZiMDMyMTBiYWQ3MGQ0M2Y0MDI2YmI3OTo3OjY0MjU6ZjE3YzUyM2ExM2JkYzRmYjI5Njg5N2Y0OTVmYmNiYjAzNTU3YmM5NjgyNmI1NjgzMjdhOGQ5YTMxOTJjYzYxMTpwOlQ6Rg&data=05%7C02%7Cece.ad.coord%40technion.ac.il%7C647b309ea6994274b60808df24757378%7Cf1502c4cee2e411c9715c855f6753b84%7C1%7C0%7C639269759009409259%7CUnknown%7CTWFpbGZsb3d8eyJFbXB0eU1hcGkiOnRydWUsIlYiOiIwLjAuMDAwMCIsIlAiOiJXaW4zMiIsIkFOIjoiTWFpbCIsIldUIjoyfQ%3D%3D%7C0%7C%7C%7C&sdata=EtBx86vRvGBPpVFsuRqDmzc4PB9t%2Fdjuq7gN5SCnctQ%3D&reserved=0

M.Sc. student under the supervision of Prof. Israel Cohen.

Meeting ID: 941 9864 3052