Flash Attention（Flash attention with bias）

清疚 2023-10-14 14:44 4阅读 0赞

FA主要思路还是通过tile技术减少在HBM和on-chip SRAM内存读写时间。FA在bert-large上端到端训练有15%的加速（seq length 512)， 在GPT-2有3倍加速（seq length 1024）。

![format_png][]

GPU的内存层次结构

**主要贡献点**

1.  计算softmax时候不需要全量input数据，可以分段计算。
2.  反向传播的时候，不存储attention matrix (N^2的矩阵)，而是只存储softmax归一化的系数。

--------------------

[Flash Attention with attention bias论文阅读记录 - 知乎][Flash Attention with attention bias_ -]

[速度飙升200%！Flash Attention 2一统江湖，注意力计算不再是问题！\_算法码上来的博客-CSDN博客][200_Flash Attention 2_-CSDN]

[FlashAttention：Fast and Memory-Efficient Exact Attention with IO-Awareness\_江小皮不皮的博客-CSDN博客][FlashAttention_Fast and Memory-Efficient Exact Attention with IO-Awareness_-CSDN]

[https://github.com/robotcator/flash-attention/tree/odd-len][https_github.com_robotcator_flash-attention_tree_odd-len]

[Flash Attention 推公式 - 知乎][Flash Attention _ -]

[format_png]: https://img-blog.csdnimg.cn/img_convert/1701137227caa76b04057d76f18bc24c.webp?x-oss-process=image/format,png
[Flash Attention with attention bias_ -]: https://zhuanlan.zhihu.com/p/567167376
[200_Flash Attention 2_-CSDN]: https://blog.csdn.net/God_WeiYang/article/details/131820781
[FlashAttention_Fast and Memory-Efficient Exact Attention with IO-Awareness_-CSDN]: https://blog.csdn.net/qq128252/article/details/131455000
[https_github.com_robotcator_flash-attention_tree_odd-len]: https://github.com/robotcator/flash-attention/tree/odd-len
[Flash Attention _ -]: https://zhuanlan.zhihu.com/p/646697716