【多模态大模型】FlashAttention in NeurIPS 2022

一、引言

论文： FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
作者： Stanford University
代码： FlashAttention
特点： 该方法提出将Q、K、V拆分为若干小块，使执行注意力时不需要频繁进行读写操作，而是每个小块只进行一次读写，从而提升注意力的执行速度。

⚠️ 在学习该方法前，建议补充Attention的相关知识。

二、详情

GPU中SRAM和HBM的计算和存储能力如下图：

可见，SRAM计算能力强（17TB/s），HBM的存储容量大（40GB）。因此，GPU的运算通常在SRAM上进行，如果运算结果的内存占用太大，系统会把运算结果先写入HBM，然后从HBM读出来再在SRAM上进行下一步的运算。

于是，我们就得到原始Attention的执行过程：

其中，Q、K、V分别是Query、Key、Value矩阵，S是相似度矩阵，P是权重矩阵，O是输出矩阵。

这里没写除以$\sqrt{d_k}$的操作，不过无伤大雅，因为它对运算的影响并不大。

可见，计算S、P、O时都要进行读取，计算完成后也都要进行写入。然而，运算速度领先于读写速度导致SRAM运算完了要等数据过来才能进行下一步运算，这就拖慢了整体的速度。

2.1 拆分

FlashAttention提出将Q、K、V拆分成若干小块，这样每个小块的S、P矩阵不至于太大到需要写入HBM中，这样就能只在最开始读取Q、K、V、O（之前的运算结果），在SRAM中完成所有运算后，再将新的O写入HBM。

如果没有SoftMax操作，该过程很容易实现，如下图：

分别循环Q和K、V的小块，循环结果求和就是我们所有期望的O。但是，SoftMax阻碍了它的实现，回顾原始SoftMax公式：
$$softmax(\mathbf{s}{)}_j=\frac{e^{s_j}}{{\sum }_{k=1}^N{e}^{s_k}}$$

可见，它要把相似度矩阵S的每一行转为一个概率分布。但是分块策略无法一次性获得完整的S中的行，于是FlashAttention在SoftMax中引入了$m(\mathbf{s})$，新的SoftMax公式如下：
$$softmax(\mathbf{s}{)}_i=\frac{e^{s_i-m(\mathbf{s})}}{{\sum }_{j=1}^N{e}^{s_j-m(\mathbf{s})}}=\frac{f_i}{l(\mathbf{s})}$$

其中，最大值$m(\mathbf{s})={\max }_i{s}_i$，指数和$l(\mathbf{s})={\sum }_i{f}_i$。事实上，该操作不会影响SoftMax的结果，如下：
$$\begin{gathered} softmax([1,2,3,10])=[\frac{e^1}{e^1+e^2+e^3+e^{10}},\frac{e^2}{e^1+e^2+e^3+e^{10}},\frac{e^3}{e^1+e^2+e^3+e^{10}},\frac{e^{10}}{e^1+e^2+e^3+e^{10}}] \\ =[\frac{e^{1-10}}{e^{1-10}+e^{2-10}+e^{3-10}+e^{10-10}},\frac{e^{2-10}}{e^{1-10}+e^{2-10}+e^{3-10}+e^{10-10}},\frac{e^{3-10}}{e^{1-10}+e^{2-10}+e^{3-10}+e^{10-10}},\frac{e^{10-10}}{e^{1-10}+e^{2-10}+e^{3-10}+e^{10-10}}] \end{gathered}$$

可见，上下同乘$e^{10}$即可还原为原公式。

此时，我们分$T_r=2$块分别计算上述SoftMax，有：
$$\begin{gathered} softmax([1,2])=[\frac{e^{1-m_1}}{e^{1-m_1}+e^{2-m_1}},\frac{e^{2-m_1}}{e^{1-m_1}+e^{2-m_1}}]=[\frac{f_1}{l_1},\frac{f_2}{l_1}],m_1=2 \\ softmax([3,10])=[\frac{e^{3-m_2}}{e^{3-m_2}+e^{10-m_2}},\frac{e^{10-m_2}}{e^{3-m_2}+e^{10-m_2}}]=[\frac{f_3}{l_2},\frac{f_4}{l_2}],m_2=10 \end{gathered}$$

其中，每个小块里减去的是当前块的最大值，记为$m_i$；当前块的分子，记为${\mathbf{p}}_i$（是多个$f_i$组成的向量）；当前块的分母指数和，记为$l_i$。对应地，当前块的输出${\mathbf{p}}_i/l_i$，记为$\mathbf{o}$。

在不同块的遍历计算过程中，我们可以不断更新最大值$m(\mathbf{s})$（初始为负无穷）、指数和$l(\mathbf{s})$（初始为0）。

对于$m(\mathbf{s})$，更新公式为$m(\mathbf{s}{)}^{new}=\max (m(\mathbf{s}),m_i)$。
对于$l(\mathbf{s})$，更新公式为$l(\mathbf{s}{)}^{new}=e^{m(\mathbf{s})-m(\mathbf{s}{)}^{new}}\times l(\mathbf{s})+e^{m_i-m(\mathbf{s}{)}^{new}}\times l_i$。

在第一块中，

• $m(\mathbf{s}{)}^{new}=\max (-\inf ,m_1)=2$
• $l(\mathbf{s}{)}^{new}=e^{m(\mathbf{s})-m(\mathbf{s}{)}^{new}}\times l(\mathbf{s})+e^{m_1-m(\mathbf{s}{)}^{new}}\times l_1=e^{-\inf -2}\times 0+e^{2-2}\times (e^{1-2}+e^{2-2})$
• 令$m(\mathbf{s})\leftarrow m(\mathbf{s}{)}^{new}$，$l(\mathbf{s})\leftarrow l(\mathbf{s}{)}^{new}$

在第二块中，

• $m(\mathbf{s}{)}^{new}=\max (2,m_2)=10$
• $l(\mathbf{s}{)}^{new}=e^{m(\mathbf{s})-m(\mathbf{s}{)}^{new}}\times l(\mathbf{s})+e^{m_2-m(\mathbf{s}{)}^{new}}\times l_2$
  $=e^{2-10}\times (e^{1-2}+e^{2-2})+e^{10-10}\times (e^{3-10}+e^{10-10})=e^{1-10}+e^{2-10}+e^{3-10}+e^{10-10}$

可见，最后的输出结果$m(\mathbf{s})$和$l(\mathbf{s})$已经与实际$softmax([1,2,3,10])$中的一致。

$m(\mathbf{s})$的更新公式能使$m(\mathbf{s}{)}^{new}$始终为当前行的最大值，$l(\mathbf{s})$的更新公式能使$l(\mathbf{s}{)}^{new}$的指数项始终减的是$m(\mathbf{s}{)}^{new}$。

同样地，在遍历过程中，我们也可以根据新的$m(\mathbf{s})$和$l(\mathbf{s})$计算和更新当前的$\mathbf{o}$（初始为0向量）。

对于$\mathbf{o}$，更新公式为
$${\mathbf{o}}^{new}=\frac{l(\mathbf{s})\times e^{m(\mathbf{s})-m(\mathbf{s}{)}^{new}}\times \mathbf{o}+e^{m_i-m(\mathbf{s}{)}^{new}}\times {\mathbf{p}}_i\times {\mathbf{V}}_i}{l(\mathbf{s}{)}^{new}}$$

其中，${\mathbf{p}}_i=[f_{i\ast Br},\cdots ,f_{(i+1)\ast B_r}]$，${\mathbf{V}}_i$为V矩阵的第$i$块。

我们假设$\mathbf{V}=[[1,2],[3,4],[5,6],[7,8]]$，则有

在第一块中，

• $m(\mathbf{s}{)}^{new}=2$
• $l(\mathbf{s}{)}^{new}=e^{-\inf -2}\times 0+e^{2-2}\times (e^{1-2}+e^{2-2})=e^{1-2}+e^{2-2}$
• ${\mathbf{o}}^{new}=\frac{0\times e^{-\inf -2}\times 0+e^{2-2}\times [e^{1-2},e^{2-2}]\times \left[\begin{array}{cc}1& 2\\ 3& 4\end{array}\right]}{e^{-\inf -2}\times 0+e^{2-2}\times (e^{1-2}+e^{2-2})}=\frac{[e^{1-2},e^{2-2}]\times \left[\begin{array}{cc}1& 2\\ 3& 4\end{array}\right]}{(e^{1-2}+e^{2-2})}$
• 令$m(\mathbf{s})\leftarrow m(\mathbf{s}{)}^{new}$，$l(\mathbf{s})\leftarrow l(\mathbf{s}{)}^{new}$，$\mathbf{o}\leftarrow {\mathbf{o}}^{new}$

在第二块中，

• $m(\mathbf{s}{)}^{new}=10$
• $l(\mathbf{s}{)}^{new}=e^{1-10}+e^{2-10}+e^{3-10}+e^{10-10}$
• $$\begin{gathered} {\mathbf{o}}^{new}=\frac{(e^{1-2}+e^{2-2})\times e^{2-10}\times \frac{[e^{1-2},e^{2-2}]\times \left[\begin{array}{cc}1& 2\\ 3& 4\end{array}\right]}{(e^{1-2}+e^{2-2})}+e^{10-10}\times [e^{3-10},e^{10-10}]\times \left[\begin{array}{cc}5& 6\\ 7& 8\end{array}\right]}{e^{1-10}+e^{2-10}+e^{3-10}+e^{10-10}} \\ =\frac{[e^{1-10},e^{2-10}]\times \left[\begin{array}{cc}1& 2\\ 3& 4\end{array}\right]+[e^{3-10},e^{10-10}]\times \left[\begin{array}{cc}5& 6\\ 7& 8\end{array}\right]}{e^{1-10}+e^{2-10}+e^{3-10}+e^{10-10}} \end{gathered}$$

可见，最后的结果已经与实际$softmax([1,2,3,10])\times \mathbf{V}$一致。

$\mathbf{o}$的更新公式能使各块分子指数项上减去最新的$m(\mathbf{s}{)}^{new}$，并使各块的最新的指数和合并。

致谢：

本博客仅做记录使用，无任何商业用途，参考内容如下：
Flash Attention 为什么那么快？原理讲解
Flash Attention论文解读