FlashAttention 笔记

Standard Self-Attention

$$\begin{gathered} S=Q{K}^T,(N\times N) \\ P=softmax(S),(N\times N) \\ O=PV,(N\times d) \end{gathered}$$

Safe Softmax

$$safe-softmax=\frac{e^{x_i-m}}{\sum _{j=1}^N{e}^{x_j-m}}$$

计算 safe-softmax 的时候需要对 [1, N] 重复三次，需要访问 Q 和 K 三次， 并实时重新计算 x，很低效，现在希望将计算进行合并。

现在定义

$$d_i^{{}^{\prime }}=\sum _{j=1}^i{e}^{x_j-m_i}$$

当 i = N 的时候，有：

$$d_N=d_N^{{}^{\prime }}=\sum _{j=1}^i{e}^{x_j-m_N}$$

$d_i^{{}^{\prime }}$ 和 $d_{i-1}^{{}^{\prime }}$ 有如下递归关系：

$$\begin{gathered} d_i^{{}^{\prime }}=\sum _{j=1}^i{e}^{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}^{{}^{\prime }}{e}^{m_{i-1}-m_i}+e^{x_i-m_i} \end{gathered}$$

递归关系只依赖 $m_{i-1}$ 和 $m_i$，于是可以把 $d_i^{{}^{\prime }}$ 和 $m_i$ 放在同一个循环中。

FlashAttention V1

online softmax 最多只有一个 2-pass 的算法，不存在 1-pass 算法，但是 Attention 可以有 1-pass 算法。基于上述的 online softmax 可以得到一个 1-pass 的 Attention 算法。

重点在第二个循环：

$$\begin{gathered} a_i=\frac{e^{x_i-m_N}}{d_N^{{}^{\prime }}} \\ {o}_i=o_{i-1}+a_{i}V[i,:] \end{gathered}$$

推导 1-pass 版本的 FlashAttention:

$$\begin{gathered} o_i^{{}^{\prime }}=(\sum _{j=1}^i(\frac{e^{x_j-m_i}}{d_i^{{}^{\prime }}})V[j,:]) \\ \mathrm{当}i=N\mathrm{时}\mathrm{有}： \\ {o}_N^{{}^{\prime }}=o_N=(\sum _{j=1}^i(\frac{e^{x_j-m_N}}{d_N^{{}^{\prime }}})V[j,:]) \end{gathered}$$

推导 $o_i{}^{\prime }$ 和 $o_{i-1}{}^{\prime }$ 之间的关系：

$$\begin{gathered} o_i^{{}^{\prime }}=(\sum _{j=1}^i(\frac{e^{x_j-m_i}}{d_i^{{}^{\prime }}})V[j,:]) \\ =(\sum _{j=1}^{i-1}(\frac{e^{x_j-m_i}}{d_i^{{}^{\prime }}})V[j,:])+(\frac{e^{x_i-m_i}}{d_i^{{}^{\prime }}})V[i,:] \\ =(\sum _{j=1}^{i-1}(\frac{e^{x_j-m_{i-1}}}{d_i^{{}^{\prime }}})\frac{e^{x_j-m_i}}{e^{x_j-m_{i-1}}}\frac{d_{i-1}^{{}^{\prime }}}{d_i^{{}^{\prime }}}V[j,:])+(\frac{e^{x_i-m_i}}{d_i^{{}^{\prime }}})V[i,:] \\ =o_{i-1}^{{}^{\prime }}\frac{d_{i-1}^{{}^{\prime }}{e}^{m_{i-1}-m_i}}{d_i^{{}^{\prime }}}+(\frac{e^{x_i-m_i}}{d_i^{{}^{\prime }}})V[i,:] \end{gathered}$$

可以看到 $o_i^{{}^{\prime }}$ 和 $o_{i-1}^{{}^{\prime }}$ 递归关系不依赖 $m_n$，因此可以将第二个循环完全合并到第一个循环中去。

FlashAttention V2

主要做了工程上的优化：

• 减少大量非矩阵乘的冗余计算，增加 Tensor Core 的计算比例
• forward pass/backward pass 均增加 seqlen 维度的并行，forward pass 交替 Q，K，V 循环顺序
• 更好的 Warp Partitioning 策略，避免 Split-K

在 Tri Dao 的 FlashAttention 的 Triton 实现中，使用了 LSE(LogSumExp) 来做 smooth maximum。

$$\begin{gathered} softmax=\frac{exp(x_i-c)}{\sum _{i}exp(x_i-c)} \\ c=max(x_1,...,x_n) \end{gathered}$$

LSE 可以被用于作为 approximation of max。

$$\begin{gathered} LSE(x_1,...,x_n)=log(exp(x_1)+...+exp(x_n)) \\ =c+log\sum _{n}exp(x_i-c) \\ c=max(x_1,...,x_n) \end{gathered}$$

Triton 中的算法可以被描述为:

等式 $(2)$ 中 lse_i 被用于作为最大值的一个近似。

在等式 $(9)$ 中：

$$\begin{gathered} exp(ls{e}_i-m_{ij})=exp(c_{old}+log\sum _{n}exp(x_i-c_{old})-c_{new})+l_{ij} \\ =exp((c_{old}-c_{new}+log\sum _{n}exp(x_i-c_{old}))) \\ =exp(c_{old}-c_{new})\sum _{n}exp(x_i-c_{old}) \end{gathered}$$

在等式 $(11)$ 中的 o_scale 作为 softmax 函数的分母：

$$\begin{gathered} m_i-ls{e}_i=exp(c-c-log\sum _{n}exp(x_i-c)) \\ =exp(-log\sum _{n}exp(x_i-c)) \\ =exp(log\frac{1}{\sum _{n}exp(x_i-c)}) \\ =\frac{1}{\sum _{n}exp(x_i-c)} \end{gathered}$$