通过块稀疏注意力加速长文本推理，不降精度还更快。
Block Sparse Flash Attention
- 按块计算注意力得分，动态筛选关键值块
- 跳过约50%计算与内存传输，实测提速1.24倍
- 无需训练，仅需小数据集校准阈值，适合部署优化
现代大语言模型在推理和多文档任务中需要更长上下文，但注意力机制的二次复杂度带来严重计算瓶颈。我们提出块稀疏闪光注意力（BSFA），一种无需修改模型即可替代FlashAttention的加速方案。不同于预判重要性的方法，BSFA先计算精确的查询-键相似度，为每个查询选择最重要的k个值块。通过将每块最大得分与校准阈值比较，可跳过约50%的计算与内存传输。该训练无关方法仅需在小数据集上进行一次阈值校准，以学习各层各头的注意力得分分布。我们提供了可直接替换FlashAttention的CUDA内核实现。在Llama-3.1-8B上，BSFA在真实推理任务中最高提速1.10倍，在针堆找针任务中最高提速1.24倍，同时保持99%以上基线准确率，部分配置甚至因聚焦相关内容而提升精度，显著优于现有稀疏注意力方法。
原文摘要 · Abstract (English)
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 FlashAttention (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.10x speedup on real-world reasoning benchmarks and up to 1.24x for needle-in-a-haystack retrieval tasks while maintaining above 99% baseline accuracy, with certain configurations even improving accuracy by focusing on the most relevant content, substantially outperforming existing sparse attention methods. The implementation is available at https://github.com/Danielohayon/Block-Sparse-Flash-Attention
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。