大模型--FlashAttention V2 原理--27
1. 参考
https://zhuanlan.zhihu.com/p/642962397
2. 回顾 FlashAttention V1
softmax稳定版
FlashAttention 的整体思路
-
先算block1的softmax
-
保存全局变量：m --全局的最大元素 l --全局的EXP指数求和项
-
算block2的softmax
-
更新block2的全局softmax
-
同理将上面的前3项换成x1的，实现对block1的softmax进行更新
-
算block3的softmax 更新block2 block1
...
直到把所有的block算完
如图所示
这里会发现一个问题每次算一个block需要更新前面的所有block 这复杂度挺高的 有没有办法优化
V2
局部向量Block2 的更新逻辑是乘以
也就是乘上原来的分母 
，然后再除以新的全局EXP求和项
。
其实可以在每一次分块计算完毕后不去除以此时的EXP求和项，只需要等到最后去直接除以最终的
每一次迭代过程中，不再除以EXP求和项
直到处理完最后一个分块后，直接用此时的全局EXP求和项来做分母即可。
简单的例子：
假设我们要计算从1到10的平均值，但每次只能算两个数。
FlashAttention v1的计算方法类似于：
第一次算出前两个数的平均值，并记录下当前已经计算过的数字数量N，
第二次算出前三个数的平均值，需要使用之前的N（N=2）来更新，
...
FlashAttention v2的方法类似于：
第一次算出前两个数的和，并记录下当前已经计算过的数字数量N：
, N = 2。
第二次算出前三个数的和，并记录下当前已经计算过的数字数量N：
, N = 3。
计算完毕后，只需要将当前的总和除以N即可。
FlashAttention v2的优势在于少了原来每一步的乘法和除法。
Efficient Memory Attention
这一节介绍另一种常用的self-attention加速算法：EMA（Efficient Memory Attention）。正如其名，EMA原本主要为解决self-attention的空间复杂度问题而设计。Attention加速库xformers对EMA进一步进行了速度上的优化，在后来被大量LLM所使用。本小节主要关注EMA自身的改进方法。
我们首先将self-attention的计算公式重新写为向量形式。
分子、分母都是一个求和项，并且它们都是从索引1开始，求和至索引 i，
那么一种优化后的计算方法是：对分子和分母分别使用一个临时变量来保存之前n项的求和结果；当需要计算n+1项的结果时，只需要单独计算n+1项的值，并加上之前的求和结果即可。
例如：
可以使用以下两个变量来分别保存分子和分母截至到索引 m 的求和结果：
当需要计算 m+1 的分子和分母时，只需要做如下更新：
xformers
xformers是一个attention加速库，它的文档中提到了它的关键特性之一正是Efficient Memory Attention。但是需要注意的是，xformers所谓的Efficient Memory Attention并不特指4.1小节中介绍的EMA方法，它指一系列对存储友好的attention变体方法，包括FlashAttention。
在xformers中使用Efficient Memory Attention时，它会自动根据用户的输入来选择使用何种attention算子。这些候选算子中，FlashAttention的优先级最高。
