Standard Attention
标准Attention计算可以简化为：
\[O = softmax(QK^T)V \tag{1} \]此处忽略了Attention Mask和维度归一化因子\(1/\sqrt{d}\)。
公式(1)的标准计算方式是分解成三步：
\[S = QK^T \tag{2} \]\[P=softmax(S) \tag{3} \]\[O = PV \tag{4} \]但这样做的问题在于，假设\(Q,K,V \in R^{N\times d}\)，其中\(N\)为序列长度，\(D\)为注意力头的维度，那么输出\(O \in R^{N\times d}\)，\(S,P\in R^{N\times N}\)。由于在标准实现下，\(S,P\)都需要从HBM中读写，因此构成了\(O(N^2)\)的内存复杂度。一般情况下\(N\gg D\)，例如GPT-2中，\(N=1024\)，\(d=64\)，因此\(S\)和\(P\)的\(O(N^2)\)显存开销是远大于\(Q,K,V,O\)的\(O(Nd)\)的。
一个朴素的想法是：我们能否不进行\(O(N^2)\)的HBM读写，通过避免这一频繁的读写操作来大大提升Attention的计算效率。
Online Softmax
假设我们不将\(S,P\)写回HBM，那么就得将其放在片上SRAM中。但是这里的问题是片上SRAM受限于容量，一般无法一次性完整的计算Attention，因此我们必须采用分块（Tiling）操作，使得分块后的内存需求不超过SRAM的大小。但计算softmax的时候，其归一化因子（分母）需要所有的输入数据，因此进行分块计算的难度较大。
考虑softmax的公式，对于输入序列\(x\)：
\[x=[x_1,x_2,...,x_d] \tag{5} \]原生softmax函数为：
\[softmax(x_i)=\frac{e^{x_i}}{\sum_{j=1}^{d}e^{x_j}} \tag{6} \]为了避免数值溢出的问题，现在一般采用safe softmax的方式，即定义：
\[m(x)=max([x_1,x_2,...,x_d]) \tag{7} \]safe softmax函数在e指数上减去\(m(x)\)，使得所有的e指数项的值分布在0到1之间（因为\(x_i-m(x)\leq 0\)），从而规避数值溢出的问题，此外还能提升数值稳定性，加快计算速度。改造后的函数为：
\[softmax(X)=\frac{e^{x_i-m(x)}}{\sum_{j=1}^{d}e^{x_j-m(x)}} \tag{8} \]接下来我们需要研究如何对safe softmax应用分块策略来计算，即所谓的online softmax。
标准的softmax情况下，算法为：
for i = 1 to N do:
\[m_i\leftarrow max(m_{i-1},x_i) \tag{9} \]end
for i = 1 to N do:
\[d_i \leftarrow d_{i-1}+e^{x_i-m_N} \tag{10} \]end
for i = 1 to N do:
\[a_i \leftarrow \frac{e^{x_i-m_N}}{d_N} \tag{11} \]end
可以看到这个计算过程需要三次1到N的循环，在Attention中，这里的\(x_i\)来自于\(QK^T\)，由于我们没办法在SRAM中装下\(Q\)和\(K\)，因此我们需要从内存中访问他们三次。假如我们能够想办法将\((9)\)到\((11)\)放到一个循环中，我们就能将访存从三次减少到一次。然而，由于\((9)\)和\((10)\)之间存在依赖，因为\((10)\)中包含一个不到最后一次循环就无法获知的\(m_N\)，因此我们很难将它们合并起来。
我们可以构造一个\(d_i^{'}=\sum_{j=1}^ie^{x_j-m_i}\)替代原有的\(d_i=\sum_{j=1}^ie^{x_j-m_N}\)以取消其对\(m_N\)的全局依赖，并且只要一达到\(i=N\)，我们自然而然的就有\(d_N^{'}=d_N\)，因此我们可以用\(d_N^{'}\)来替换\((11)\)中的\(d_N\)。并且我们可以求得\(d_{i}^{'}\)和\(d_{i-1}^{'}\)之间的递推关系：
\[d_i^{'}=\sum_{j=1}^ie^{x_j-m_i} =(\sum_{j=1}^{i-1}e^{x_j-m_i})+e^{x_i-m_i}=(\sum_{j=1}^{i-1}e^{x_j-m_{i-1}})e^{m_{i-1}-m_i}+e^{x_i-m_i}=d_{i-1}^{'}e^{m_{i-1}-m_i}+e^{x_i-m_i} \tag{12} \]可以看到这里的公式依赖于\(m_{i-1}\)和\(m_i\)。因此我们可以把\((9)\)和\((10)\)放进一个循环中：
for i = 1 to N do:
\[m_i\leftarrow max(m_{i-1},x_i) \tag{13} \]\[d_i^{'} \leftarrow d_{i-1}^{'}e^{m_{i-1}-m_i}+e^{x_i-m_i} \tag{14} \]end
for i = 1 to N do:
\[a_i \leftarrow \frac{e^{x_i-m_N}}{d_N^{'}} \tag{15} \]end
这样我们实现了3次循环到2次循环的合并，从而减少了1/3的内存访问。但是我们能否进一步直接合并到一步循环内呢。对于softmax来说，很不幸的是不可能的。但对于我们要求的Attention来说，这是可以实现的。
FlashAttention V1
对于Attention来说，我们最终要获得的并非是softmax后得出的矩阵\(P\)，而是输出矩阵\(O=PV\)，因此我们的目标是尝试找到一个一步循环求得\(O\)的方法。
我们先来看应用了online softmax的Attention计算过程：
for i = 1 to N do:
\[x_i \leftarrow Q[k,:]K^T[:,i] \]\[m_i\leftarrow max(m_{i-1},x_i) \]\[d_i^{'} \leftarrow d_{i-1}^{'}e^{m_{i-1}-m_i}+e^{x_i-m_i} \]end
for i = 1 to N do:
\[a_i \leftarrow \frac{e^{x_i-m_N}}{d_N^{'}} \tag{16} \]\[o_i\leftarrow o_{i-1}+a_iV[i,:] \tag{17} \]end
\[O[k,:]\leftarrow o_N \]我们将\((17)\)中的\(a_i\)替换成定义式\((16)\)，从而有：
\[o_i=(\sum_{j=1}^i\frac{e^{x_j-m_N}}{d_N^{'}}V[j,:]) \tag{18} \]这里可以看到依赖于两个全局值\(m_N\)和\(d_N^{'}\)。我们可以应用和online softmax推导时类似的技巧，先构造一个\(o_i^{'}\)：
\[o_i^{'}=(\sum_{j=1}^i\frac{e^{x_j-m_i}}{d_i^{'}}V[j,:]) \]只要达到\(i=N\)，我们就有\(o_N^{'}=o_N\)，并且我们可以求出一个\(o_{i-1}^{'}\)到\(o_i^{'}\)之间的递推公式：
\[o_i^{'}=(\sum_{j=1}^i\frac{e^{x_j-m_i}}{d_i^{'}}V[j,:])=(\sum_{j=1}^{i-1}\frac{e^{x_j-m_i}}{d_i^{'}}V[j,:])+\frac{e^{x_i-m_i}}{d_i^{'}}V[i,:]\\ =(\sum_{j=1}^{i-1}\frac{e^{x_j-m_{i-1}}}{d_{i-1}^{'}}\frac{e^{x_j-m_i}}{e^{x_j-m_{i-1}}}\frac{d_{i-1}^{'}}{d_i^{'}}V[j,:])+\frac{e^{x_i-m_i}}{d_i^{'}}V[i,:] \\ =(\sum_{j=1}^{i-1}\frac{e^{x_j-m_{i-1}}}{d_{i-1}^{'}}V[j,:])\frac{d_{i-1}^{'}}{d_i^{'}}e^{m_{i-1}-m_i}+\frac{e^{x_i-m_i}}{d_i^{'}}V[i,:] \\ = o_{i-1}^{'}\frac{d_{i-1}^{'}}{d_i^{'}}e^{m_{i-1}-m_i}+\frac{e^{x_i-m_i}}{d_i^{'}}V[i,:] \tag{19} \]可以看到这里不再依赖任何一个全局值，因此我们可以得到Flash Attention的算法：
for i = 1 to N do:
\[x_i \leftarrow Q[k,:]K^T[:,i] \]\[m_i\leftarrow max(m_{i-1},x_i) \]\[d_i^{'} \leftarrow d_{i-1}^{'}e^{m_{i-1}-m_i}+e^{x_i-m_i} \]\[o_i^{'}=o_{i-1}^{'}\frac{d_{i-1}^{'}}{d_i^{'}}e^{m_{i-1}-m_i}+\frac{e^{x_i-m_i}}{d_i^{'}}V[i,:] \]end
\[O[k,:]\leftarrow o_N^{'} \]我们可以进一步对这个算法应用分块（tiling），假定tile的大小为\(b\)，共分块\(\#tiles\)个。那么\(x_i\)为存储\([(i-1)b:ib]\)的\(QK^T\)值的向量。\(m_i^{(local)}\)为向量\(x_i\)的局部最大值。那么对于每个tile，有：
for i = 1 to #tiles do:
\[x_i \leftarrow Q[k,:]K^T[:,(i-1)b:ib] \]\[m_i^{(local)}\leftarrow max_{j=1}^b(x_i[j]) \]\[m_i \leftarrow max(m_{i-1},m_i^{(local)}) \]\[d_i^{'} \leftarrow d_{i-1}^{'}e^{m_{i-1}-m_i}+\sum_{j=1}^b e^{x_i[j]-m_i} \]\[o_i^{'}=o_{i-1}^{'}\frac{d_{i-1}^{'}}{d_i^{'}}e^{m_{i-1}-m_i}+\sum_{j=1}^{b}\frac{e^{x_i[j]-m_i}}{d_i^{'}}V[j+(i-1)b,:] \]end
\[O[k,:]\leftarrow o_{N/b}^{'} \]形象的理解如下图所示：
最后我们来看效果，由于\(S\)和\(P\)的计算完全在SRAM上完成（之前做不到的原因在这节开头时说了，想要完整的把\(S\)，\(P\)放上去，片上SRAM的容量不够，但是采用分块迭代策略后就ok了）而不需要对HBM做写回。因此在Standard Attention一节我们分析的，\(O(N^2)\)的\(S\)，\(P\)的HBM读写开销就没有了，只有\(Q\)，\(K\)，\(V\)，\(O\)的\(O(Nd)\)的开销，但我们之前也分析过，由于\(N\gg d\)，所以\(N^2\gg Nd\)，我们可以进一步的当作现在的显存开销变成了只有与\(N\)线性相关，而非二次相关的\(O(N)\)。从\(O(N^2)\)到\(O(N)\)，这显然是一个非常显著的改进。
FlashAttention V2
在V1的基础上，我们来看V2的一个insight。从硬件的角度来说，GPU计算矩阵乘加的算力是远高于其他的运算的。具体来说，以A100为例，FP16/BF16的矩阵乘法可以达到312TFLOPS，但是对于非矩阵乘法的FP32，其算力只有19.5TFLOPS，差了一个数量级（16x）。因此一个明显的改进思路是减少FlashAttention中的非矩阵乘加运算。
观察公式\((19)\)，一个切入点是每个循环计算\(O\)时进行了两次除法，即：
\[o_i^{'}=o_{i-1}^{'}\frac{d_{i-1}^{'}}{d_i^{'}}e^{m_{i-1}-m_i}+\frac{e^{x_i-m_i}}{d_i^{'}}V[i,:] \]两项都需要除以\(d_i^{'}\)。因此相当于是进行了2N次的除法。但实际上这个除法操作可以提取到循环外，即每次更新\(o_i^{'}\)时，采用：
\[\widetilde{o}_i^{'}=\widetilde{o}_{i-1}^{'}e^{m_{i-1}-m_i}+e^{x_i-m_i}V[i,:] \]因此每次更新时可以只维护未缩放的\(\widetilde{o}_i^{'}\)。当\(i=N\)时，利用\(o_N^{'}=\widetilde{o}_N^{'}/d_N^{'}\)，可以将之前每次循环中的2N次除法提出，变成循环结束后进行一次除法，从而大大减少除法的计算量（从2N次变为1次）。
即：
for i = 1 to N do:
\[x_i \leftarrow Q[k,:]K^T[:,i] \]\[m_i\leftarrow max(m_{i-1},x_i) \]\[d_i^{'} \leftarrow d_{i-1}^{'}e^{m_{i-1}-m_i}+e^{x_i-m_i} \]\[\widetilde{o}_i^{'}=\widetilde{o}_{i-1}^{'}e^{m_{i-1}-m_i}+e^{x_i-m_i}V[i,:] \]end
\[o_N^{'}=\frac{\widetilde{o}_N^{'}}{d_N^{'}} \]\[O[k,:]\leftarrow o_N^{'} \]最本质的原因其实在于在迭代计算时，实际上每一次\(o_i^{'}\)的缩放项\(d_{i-1}^{'}/d_i^{'}\)都可以把上一次\(o_{i-1}^{'}\)的共分母\(d_{i-1}^{'}\)给吸收掉。因此也可以在迭代时直接丢弃这个冗余的运算（不妨联想一下反向传播的链式法则，有一定的相似性）。
V2为了应对训练时的需求，在前向计算的循环中也会暂存维护一个变量，不过我们这里不做详细讨论。此外V2在算法上也根据GPU特性更改了内外层循环的顺序来提高并行度，但这里就不去做详细介绍了，可以看论文以及其他的博客理解。
FlashAttention V3
现在来看V3。在V2的基础上，为了提升Flash Attention算法在H100 GPU上的利用率，V3做了几件事，首先将GEMM操作以Producer & Consumer的形式进行了异步化，随后通过Ping-Pong操作将softmax操作隐藏到GEMM操作中（GEMM-softmax流水线），最后应用了更低精度的FP8数制GEMM操作来实现性能提升。
Producer和Consumer的理解其实很简单，Producer的目的是从HBM中加载计算所需的\(Q\)，\(K\)，\(V\)，而Consumer的内容和V2的公式完全一样，主要起到消耗掉Producer提供的\(Q\)，\(K\)，\(V\)并计算\(O\)然后写回。通过Ping-Pong调度这两个部分，可以把慢速的softmax操作隐藏到分段的GEMM操作中。具体来说，以下图为例，当一个Warpgroup在进行GEMM操作时，另一个Warpgroup在进行前一批GEMM操作后的softmax操作中去。
更一步的，在一个Warpgroup中，我们可以将一些softmax的指令与GEMM的指令进行并行来进一步提高吞吐率。如下图所示，可以将一些Softmax的指令隐藏到GEMM的指令执行时间中去。
具体的算法上和V2实际上没有发生什么变化。
参考文献
From Online Softmax to FlashAttention
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
标签：frac,推导,leftarrow,sum,Attention,逐代,tag,softmax,FlashAttention From： https://www.cnblogs.com/sasasatori/p/18474946