arXiv:2409.15097v1 公告类型: 交叉 摘要: 变压器在各种应用中被广泛使用，其中许多应用产生了稀疏或部分填充的注意力矩阵。例如，设计用于减少注意力二次复杂性的注意力掩码、序列打包技术以及最近在MEDUSA中用于快速验证的树掩码等创新。尽管这些矩阵具有固有的稀疏性，但最先进的算法Flash Attention仍然以二次复杂性处理它们，仿佛它们是密集的。在本文中，我们引入了\textbf{二进制块掩码}，这是一种高效的修改，通过使其具有掩码感知能力来增强Flash Attention。我们进一步提出了两种优化：一种针对具有连续非零模式的掩码，另一种针对极度稀疏的掩码。我们在从现实场景中得出的注意力掩码上的实验表明，运行时间提高了多达9倍。该实现将公开发布，以促进进一步的研究和应用。