让Flash Attention学会识别稀疏注意力掩码，提速最高9倍。
Efficiently Dispatching Flash Attention For Partially Filled Attention Masks
- 提出二进制块掩码，让注意力算法感知掩码结构
- 针对连续非零和极稀疏掩码分别优化，实测提速最高9倍
- 适合处理序列打包、树形掩码等实际场景的高效推理
Transformer 在各类应用中广泛使用，许多场景会产生稀疏或部分填充的注意力矩阵，例如降低注意力二次复杂度的掩码设计、序列打包技术，以及近期用于 MEDUSA 快速验证的树形掩码。尽管这些矩阵具有内在稀疏性，当前最先进的 Flash Attention 仍按稠密矩阵处理，维持二次复杂度。本文提出二进制块掩码（Binary Block Masking），一种高效的改进方法，使 Flash Attention 具备掩码感知能力。我们进一步提出两种优化：一种针对具有连续非零模式的掩码，另一种针对极端稀疏掩码。在真实场景生成的注意力掩码上实验表明，最多可实现 9 倍的运行时提升。代码将公开，以促进后续研究与应用。
原文摘要 · Abstract (English)
Transformers are widely used across various applications, many of which yield sparse or partially filled attention matrices. Examples include attention masks designed to reduce the quadratic complexity of attention, sequence packing techniques, and recent innovations like tree masking for fast validation in MEDUSA. Despite the inherent sparsity in these matrices, the state-of-the-art algorithm Flash Attention still processes them with quadratic complexity as though they were dense. In this paper, we introduce Binary Block Masking, a highly efficient modification that enhances Flash Attention by making it mask-aware. We further propose two optimizations: one tailored for masks with contiguous non-zero patterns and another for extremely sparse masks. Our experiments on attention masks derived from real-world scenarios demonstrate up to a 9x runtime improvement. The implementation will be publicly released to foster further research and application.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。