flash-attention是一种算子合并(kernel fusion)的优化. 把self-attention分块, 直接在SRAM里计算, 省去了HBM来回搬运中间结果S和P的时间(如下图). self-attention由两层矩阵乘法, softmax, 和其他eltwise计算(mask, dropout)构成.
attention分块计算
forward计算: Q,K,V -> O
- 矩阵分块如上图: 都是在token的维度分块
- i iteration循环(黄色部分)
- 每次i iteration需要load不同位置的 Q, O
- j iteration循环(浅色部分)
- 每次j iteration需要load不同位置的 K, V
- i iteration循环(黄色部分)
- softmax计算:
- $softmax(x)=\frac{e^{x}}{\sum_j e^{x_j}}$ , 其中$\sum_j$ 是rowsum
- safe softmax: $softmax(x)=\frac{e^{x-m}}{\sum_j e^{x_j-m}}$
- 为了避免$\sum_j e^{x_j}$ 产生特别大的值溢出
- softmax分块:
- 对于relation S的每个分块的patch $P$, 记: 局部最大值 $m_p=\max_{x\in P}(x)$, 局部rowsum $l_p=\sum_{x\in P} e^{x-m_P}$
- 得到: $softmax(x)=\frac{e^x\cdot e^{(m_p-m)}}{\sum_P l_p\cdot e^{(m_p-m)}}$
- 写成关于$j$ 的迭代形式:(softmax和iterate i无关)
- $m_{j+1}=\max(m_j, m_{local})$
- local rowsum $l_{j+1}=l_j\cdot e^{(m_j-m_{j+1})}+l_{local}\cdot e^{(m_{local}-m_{j+1})}$
- softmax的中间结果 $\tilde P_{local}=e^{S_{local}-m_{local}}$
- $O_{j+1}=softmax(S_{j+1})V_{j+1}=\text{diag}(l_{j+1})^{-1}(\text{diag}(l_j)\cdot O_j\cdot e^{m_j-m_{j+1}}+ \tilde P_{local}V_{j+1}\cdot e^{m_{local}-m_{j+1}})$
- 对于relation S的每个分块的patch $P$, 记: 局部最大值 $m_p=\max_{x\in P}(x)$, 局部rowsum $l_p=\sum_{x\in P} e^{x-m_P}$
backward计算: O, dO, Q, K, V -> dV, dK, dQ
flash-attention里面因为不保存S, P的结果, 在backward的时候要对它们重新计算
用reverse mode 的chain rule做backward计算:
- 矩阵计算: $A=BC$
- $dC = \frac{\partial y}{\partial C}=\sum_{A_{ij}} \frac{\partial y}{\partial A_{ij}}\cdot \frac{\partial A_{ij}}{\partial C}=B^{T}dA$ , 类似的 $dB=dA\ C^T$
- softmax计算: $P=softmax(S)$, $P_{j}=\frac{e^{S_{j}}}{\sum_{x\in S} e^x}$ , 下面矩阵都是eltwise乘法$\downarrow$ \(\begin{align}dS&=\sum_{P_{j}} dP_{j}\frac{\partial P_{j}}{\partial S}&=\sum_{P_{j}} (P_{k}\mathbb 1_{j=k}-P_{j}P_{k})dP_{j}&\ (\text{根据乘除法求导法则}\frac{\partial P_j}{\partial S_k}=P_j\mathbb 1_{k=j}-\frac{e^{s_{j}}e^{s_k}}{(\sum_{x\in S} e^x)^2}=diag(P)-PP^T\text{是softmax的Jacobian})\\ &&=P_{k}dP_{k}+(\sum_{P_{j}}P_{j}dP_{j})P_{k}&=P(dP+\sum_{P_{j}}P_{j}dP_{j})\\ &&&=P(dP+ \sum_{j} P_j(\sum_i V_{ij} dO_{j}))\ (\text{因为}PV=O)\\ &&&=P(dP+ \sum_{j}O_{j}dO_{j})\ (\text{其中}\sum_j \text{就是rowsum})\end{align}\)
- backward采用和forward相同的分块
每块占用的SRAM
- forward和backward采用相同的分块, backward需要的SRAM比较大, 因此直接以backward的需求量为准
- backward:
- Q, dQ, O, dO : Br x d
- K, V, dK, dV: Bc x d
- S, P, dS, dP, : Br x Bc
- l, m: Br
- 在flash-attention文章里, 取
Bc=floor(M/4d),Br=min(floor(M/4d),d), 这样能保证SRAM够用吗❓
flash-attention v1
- 让iterate i为内循环, iterate j为外循环
flash-attention v2
- 更改iteration的顺序(一般的规则是把相关性更高的放在更内层循环)
- forward时让iterate j为内循环, 因为forward最终结果为O, 这样O的分块不用反复进出
- barkward时让iterate i为内循环, backward的最终结果为(dK, dV, dQ), 这样dK, dV不用反复进出
- softmax中间结果的存储:
- forward正常计算, 不过不存储$l,m$, 改为存$L=m+\log (l)$
- v2 forward采用j的内循环, 和$l, m$ 计算方向一致, forward的时候本来就可以inplace更新不需要存储
- backward时重新计算P, $P=diag(l)^{-1}e^{S-m}=e^{S-L}$
- forward正常计算, 不过不存储$l,m$, 改为存$L=m+\log (l)$
- 优化thread blocks的并行化计算: 把i iteration和j iteration的部分也切分进行并行计算(尤其推理时batch size=1; 或训练token长度特别长, 需要减小batch_size和attention_heads时, 并行度不够)
- forward时i iteration的相关性比较小, 所以切分i iteration更方便
- 但是每个i iteration的切分都要自己load一遍K, V了
- forward时i iteration的相关性比较小, 所以切分i iteration更方便
- 优化warp的并行计算: warp之间也是并行的
- v2 也是在i iteration切分warp的并行
reference
[1] Dao, Tri, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” arXiv, June 23, 2022. https://doi.org/10.48550/arXiv.2205.14135.
[2] Dao, Tri. “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.” arXiv, July 17, 2023. https://doi.org/10.48550/arXiv.2307.08691.
代码: https://github.com/Dao-AILab/flash-attention
其他讨论的链接
[1] https://zhuanlan.zhihu.com/p/669926191 flash-attention v1
[2] https://mp.weixin.qq.com/s/5K6yNj23NmNLcAQofHcT4Q flash-attention v2
[3] https://zhuanlan.zhihu.com/p/685020608 概括
[4] https://mp.weixin.qq.com/s/NKShFDrfDGsb0G6PAkUCGw triton的实现
[5] https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html 基于triton的flash-attention代码