FlashAttention 公式推导

本文目前只介绍关于FlashAttention的公式推导，相关背景可参考：

• paper
• blog

一、分块下的softmax如何计算

对于向量$x\in R^B$
$$\begin{gathered} x=[x_1,x_2,...,x_B] \\ m(x):=\underset{i}{\max }{x}_i=max(x_1,x_2,...,x_B) \\ f(x):=[e^{x_1-m(x)},e^{x_2-m(x)},...,e^{x_B-m(x)}] \\ l(x):=\sum _i^{B}f(x{)}_i=e^{x_1-m(x)}+e^{x_2-m(x)}+...+e^{x_B-m(x)} \\ softmax(x):=\frac{f(x)}{l(x)}=[\frac{e^{x_1-m(x)}}{l(x)},\frac{e^{x_2-m(x)}}{l(x)},...,\frac{e^{x_B-m(x)}}{l(x)}] \end{gathered}$$
则对于向量$x^{(1)},x^{(2)}\in R^B$, $x=[x^{(1)},x^{(2)}]\in R^{2B}$
$$\begin{gathered} x^{(1)}=[x_1^{(1)},x_2^{(1)},...,x_B^{(1)}],x^{(2)}=[x_1^{(2)},x_2^{(2)},...,x_B^{(2)}] \\ m(x)=m([x^{(1)},x^{(2)}])=max(m(x^{(1)}),m(x^{(2)})) \\ f(x)=[e^{m(x^{(1)})-m(x)}f(x^{(1)}),e^{m(x^{(2)})-m(x)}f(x^{(2)})] \\ l(x)=l([x^{(1)},x^{(2)}])=e^{m(x^{(1)})-m(x)}l(x^{(1)})+e^{m(x^{(2)})-m(x)}l(x^{(2)}) \\ softmax(x)=\frac{f(x)}{l(x)} \end{gathered}$$
下面来推导$f(x)$
$$\begin{gathered} f(x)=[e^{x_1^{(1)}-m(x)},e^{x_2^{(1)}-m(x)},...,e^{x_B^{(1)}-m(x)},e^{x_1^{(2)}-m(x)},e^{x_2^{(2)}-m(x)},...,e^{x_B^{(2)}-m(x)}] \\ \to f(x)=[e^{x_1^{(1)}-m(x^{(1)})+m(x^{(1)})-m(x)},e^{x_2^{(1)}-m(x^{(1)})+m(x^{(1)})-m(x)},...] \\ \to f(x)=[e^{x_1^{(1)}-m(x^{(1)})}\ast e^{m(x^{(1)})-m(x)},e^{x_2^{(1)}-m(x^{(1)})}\ast e^{m(x^{(1)})-m(x)},...] \\ \to f(x)=[e^{m(x^{(1)})-m(x)}f(x^{(1)}),e^{m(x^{(2)})-m(x)}f(x^{(2)})] \end{gathered}$$
下面来推导$l(x)$
$$\begin{gathered} l(x)=\sum _i^B{e}^{m(x^{(1)})-m(x)}f(x^{(1)})+\sum _i^B{e}^{m(x^{(2)})-m(x)}f(x^{(2)}) \\ \to l(x)=e^{m(x^{(1)})-m(x)}l(x^{(1)})+e^{m(x^{(2)})-m(x)}l(x^{(2)}) \end{gathered}$$
则$$softmax(x)=\frac{f(x)}{l(x)}=[\frac{e^{m(x^{(1)})-m(x)}f(x^{(1)})}{e^{m(x^{(1)})-m(x)}l(x^{(1)})+e^{m(x^{(2)})-m(x)}l(x^{(2)})},\frac{e^{m(x^{(2)})-m(x)}f(x^{(2)})}{e^{m(x^{(1)})-m(x)}l(x^{(1)})+e^{m(x^{(2)})-m(x)}l(x^{(2)})}]$$
有了当前的公式基础，我们可以开始FlashAttention的公式推导了

二、FlashAttention

下面是FlashAttention的算法描述：

下面我们逐行解释算法：

• 0、假设矩阵$Q、K、V\in R^{N\times d}$位于HBM(GPU global memory)，on-chip SRAM(GPU share memory)的内存大小为 M。
• 1、设置块大小为 $B_c=\lceil \frac{M}{4d}\rceil ，B_r=min(\lceil \frac{M}{4d}\rceil ,d)$
• 2、初始化 $[N\times d]$ 输出矩阵 O 全为0
        初始化 $N$ 维向量 $l$ 全为0。存储 softmax 的累积分母——指数分数的总和
        初始化 $N$ 维向量 $m$ 全为 $-\infty$。存储按行最大分数
• 3、使用步骤1中的块大小将 Q、K、V 分块。
        Q 按 $B_r$ 分块 $Q_1,...,Q_{T_r}$，每个块的维度是 $[B_r\times d]$, Q的块数为$T_r=\lceil \frac{N}{B_r}\rceil$。
        K、V 按 $B_c$ 分块为 $K_1,...,K_{T_c}$ 和 $V_1,...,V_{T_c}$，每个块的维度是 $[B_c\times d]$，K、V 的块数为$T_c=\lceil \frac{N}{B_c}\rceil$。
• 4、将O、l、m 按 $B_r$ 分块。
        O(矩阵) 分成 $O_1,...,O_{T_r}$，每个块大小为 $[B_r\times d]$；
        l(向量) 分成 $l_1,...,l_{T_r}$，每个块大小为 $B_r$
        m(向量) 分成 $m_1,...,m_{T_r}$，每个块大小为 $B_r$
• 5、outloop 遍历 $for 1 <= j <= T_c $，即遍历 Key/Value 向量
• 6、从 HBM(global memory) 加载 $K_i,V_i$ 到 on-chip SRAM(share memory).由于我们构建块大小的方式，此时 SRAM 仍有至少 50%未被占用（用于 Q 和 O）。
• 7、innerloop 遍历 $for 1<= i <= T_r $，即对 Query 向量进行循环
• 8、从 HBM 加载 $Q_i,O_i,l_i,m_i$ 到 on-chip SRAM。
• 9、计算 $S_{ij}=Q_i{K}_j^T\in R^{B_r\times B_c}$
• 10、使用上一步的 $S_{ij}$ 计算 $m_{ij},l_{ij},P_{ij}$
  $$\begin{gathered} m_{ij}=rowmax(S_{ij})\in R^{B_r} \\ {P}_{ij}=exp(S_{ij}-m_{ij})\in R^{B_r\times B_c} \\ {l}_{ij}=rowsum(p_{ij})\in R^{B_r} \end{gathered}$$
• 11、计算 $$\begin{gathered} m_i^{new}=max(m_i,m_{ij}) \\ {l}_i^{new}=e^{m_i-m_i^{new}}{l}_i+e^{m_{ij-m_i^{new}}}{l}_{ij} \end{gathered}$$
• 12、
  $$Write{O}_i\leftarrow diag(l_i^{new}{)}^{-1}(diag(l_i)e^{m_i-m_i^{new}}{O}_i+e^{m_{ij}-m_i^{new}}{P}_{ij}{V}_j)$$
  上面的过程都很好理解，这里是最难理解的一步，我们来推导一下(我们不考虑$V_i$)：
  只第一个块时，第一个块的softmax输出:
  $$\begin{gathered} O_i\leftarrow softmax(x^{(1)})=\frac{f(x^{(1)})}{l(x^{(1)})} \\ 从而O_i=[\frac{f(x^{(1)})}{l(x^{(1)})},0,0,0,...] \\ (请记住这里O_i是向量，且其余块的值都为0) \end{gathered}$$
  第一、二个块时，第一个块的softmax输出:
  $$\begin{gathered} softmax(x^{(1)})=\frac{e^{m(x^{(1)})-m(x)}f(x^{(1)})}{e^{m(x^{(1)})-m(x)}l(x^{(1)})+e^{m(x^{(2)})-m(x)}l(x^{(2)})} \\ \to softmax(x^{(1)})=\frac{e^{m(x^{(1)})-m(x)}{O}_{i}l(x^{(1)})}{e^{m(x^{(1)})-m(x)}l(x^{(1)})+e^{m(x^{(2)})-m(x)}l(x^{(2)})} \\ \to softmax(x^{(1)})=\frac{e^{m(x^{(1)})-m(x)}{O}_{i}l{(}_i)}{l_i^{new}} \end{gathered}$$
  第二个块的softmax输出：
  $$\begin{gathered} softmax(x^{(2)})=\frac{e^{m(x^{(2)})-m(x)}f(x^{(2)})}{e^{m(x^{(1)})-m(x)}l(x^{(1)})+e^{m(x^{(2)})-m(x)}l(x^{(2)})} \\ \to softmax(x^{(2)})=\frac{e^{m(x^{(2)})-m(x)}f(x^{(2)})}{l_i^{new}} \end{gathered}$$
  则理论上$O_i$就应该是：
  $$O_{i_{new}}=[\frac{e^{m(x^{(1)})-m(x)}{O}_{i}l{(}_i)}{l_i^{new}},\frac{e^{m(x^{(2)})-m(x)}f(x^{(2)})}{l_i^{new}},0,0,0,...]$$

这就是每个块softmax结果的递推过程了。有了这些，我们就可以拆解一下步骤12中的公式：
1、$diag(l_i^{new}{)}^{-1}$ 即将 $l_i^{new}$ 作为分母
2、$diag(l_i)e^{m_i-m_i^{new}}{O}_i$ 不难发现是我们上面推导的 $softmax(x^{(1)})=\frac{e^{m(x^{(1)})-m(x)}{O}_{i}l{(}_i)}{l_i^{new}}$ 的分子部分
3、而$e^{m_{ij}-m_i^{new}}{P}_{ij}$ 就是上面 $softmax(x^{(2)})=\frac{e^{m(x^{(2)})-m(x)}f(x^{(2)})}{l_i^{new}}$ 的分子部分

• 13、
  $$Write{l}_i\leftarrow l_i^{new},m_i\leftarrow m_i^{new}toHBM$$
• 14、 end for innerloop
• 15、end for outloop
• 16、Return O