『大模型笔记』FlashAttention: 具有IO意识的快速且内存高效的精确注意力机制！

Transformer在处理长序列时速度慢且内存消耗大，因为自注意力机制的时间和内存复杂度与序列长度呈二次方关系。近似注意力方法试图通过降低计算复杂度来解决这个问题，但通常无法在实际中实现速度提升。我们认为一个缺失的原则是使注意力算法具有IO感知能力——考虑在GPU内存各层之间的读写操作。我们提出了FlashAttention，这是一种具有IO感知能力的精确注意力算法，通过使用分块技术减少GPU高带宽内存（HBM）和GPU片上SRAM之间的内存读写次数。