Flash Attention

op融合

原始方法：

痛点：多次读取、写入显存。

解决：中间结果不保存，1个kernel顺序算完多个操作。

痛点：参与计算的矩阵，特别是和句子长度相关的S矩阵，太大，无法放入Shared Memory;

解决：分块矩阵计算；

问：反向传播时用到这些中间结果要求导，怎么办？

答：类似activation checkpointing，重新计算一遍。

通过矩阵的分块计算：1. 达到了融合多个操作，减少了IO加快了速度；2. 消除了句子长度平方项矩阵，减少了显存占用。

计算步骤

K、V是外循环，Q、O是内循环。

1.1

1.2

1.3

2.1

2.2

2.3

S的softmax怎么办

为了不让e^x溢出，safe_softmax的做法是，分子分母同时除以e^m，m为最大的x：

每次重新做归一化即可：

l和m占用显存非常小。

原文算法：

M/4d里的4怎么来的？
SharedMemory里，只需要保存4个子矩阵：Q、K、V、O。（S*V可以用乘加来跳过中间结果矩阵的存储。但是S矩阵不需要占空间吗？）

IO分析：

原始方法，读写几次Q、K、V、O矩阵，复杂度是O(Nd)；读写各1次S矩阵，复杂度是O(N^2)。

FlastAttention，K、V矩阵大小是Nd，SharedMemory大小是M，则外层循环次数是Nd/M；内存循环读写Q、O矩阵，矩阵大小是O(Nd)；总的读写复杂度是:前者*后者=O(Nd/M*Nd)=O(N^2*d^2/M)

因为，Multi-head attention里，一般d是64或128，M一般是100KB，因此d^2/M远小于1，因此，FlashAttention的复杂度远小于原始方法。

FlashAttention-2

第3条，说明训练时，也要加Mask的，句子里，靠前面的token的Q，不能和靠后面token的K和V进行计算。