【深度学习】FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness【一】

论文：https://arxiv.org/abs/2205.14135

论文主体。

摘要

Transformers 在长序列上运行缓慢且内存需求高，因为自注意力的时间和内存复杂度是序列长度的平方。近似注意力方法试图通过牺牲模型质量来减少计算复杂度，但通常无法实现实际的速度提升。我们认为缺失的原则是使注意力算法具备 IO 感知能力——考虑在 GPU 内存级别之间的读写操作。我们提出了 FlashAttention，一种 IO 感知的精确注意力算法，它使用分块技术来减少 GPU 高带宽内存（HBM）和 GPU 片上 SRAM 之间的内存读写次数。我们分析了 FlashAttention 的 IO 复杂度，显示其比标准注意力需要更少的 HBM 访问次数，并且在一定范围的 SRAM 大小下是最优的。我们还将 FlashAttention 扩展到块稀疏注意力，得到了一种比现有任何近似注意力方法都更快的近似注意力算法。FlashAttention 使 Transformers 的训练速度超过现有基准：在 BERT-large（序列长度 512）上相比 MLPerf 1.1 训练速度记录快 15%，在 GPT-2（序列长度 1K）上快 3 倍，在 long-range arena（序列长度 1