flashAttention-with-cuda

从softmax变形说起

safe softmax(3-pass)

$x \in R^N$即$x$是一个长度为$N$的向量,对它进行$softmax$即有如下公式:

$$softmax(\{x_1,...,x_N\}) = \{\frac{e^{x_i}}{\sum_{j=1}^{N}{e^{x_j}}}\}^{N}_{i=1}$$

我们知道指数函数它的值随着$x$的增长是爆炸性增长的,当$x_i$大于某个值就会爆float,那么为了将它约束在fp32,即得到的值是安全的,没有上溢,可以减去一个$m(x) = max(\{x_1,...,x_N\})$,那么显然我们的$x_i-m(x)\le0$,则取到的值都是安全的.

那么整个softmax过程,可以由三个循环来表示:

• 遍历$[1,N]$,求$m(x)$,其中$m(x) = \max{\{x_i\}}$;
• 遍历$[1,N]$,求$l(x)$,其中$l(x) = \sum_{i=1}^N{e^{x_i-m(x)}}$
• 遍历$[1,N]$,求$softmax(x)$,其中$softmax(x) = [\frac{e^{x_i-m(x)}}{l(x)},...,\frac{e^{x_N-m(x)}}{l(x)}]$

online softmax(2-pass)

上面的safe softmax使用了3次循环,即需要访问gmem3次,如果可以融合中间的计算,就可以减少对gmem的访问次数,可以采用以下思路融合前两个计算:

• 遍历$[1,j]$,其中$m_{j}(x) = \max(m_{j-1}(x),x_j)$ ,其中$j\le N$;
• 遍历$[1,j]$,其中$l_j(x) = l_{j-1}(x)*e^{m_{j-1}(x) - m_j(x)} + e^{x_j - m_j(x)}$

上述式子可以进行融合,因此可以在1个循环内完成$m(x)$和$l(x)$的计算

注: 上述式子中的$l_j(x)$的计算其实把它展开来看就很好理解,因为加上一个上一轮的最大值(就是把它消掉),减去这一轮更新的最大值

上述的式子实际上就实现了分段的softmax,即softmax所要处理的一整个向量,不需要计算完全,可以逐段计算,然后利用上述思想进行合并

FlashAttentionV1

1-pass Attention

Attention计算的公式忽略对$QK^T$做scaled,以及忽略mask,可以得出简化后的Attention公式:

$$Attention = softmax(QK^{T})V$$

而MHA(Multi Head Attention)只是在这个基础上将输入的$hidden\_dimension$切成了$num\_head$份的$head\_dimension$,即内部的注意力机制的计算都是一样的,只是切成了多份进行并行计算,这里我们只考虑一个头的Attention的优化,因为别的操作一样~

其中Attention中式子产生的中间阵$S = QK^T$和$P = softmax(S)$这俩阵分别代表的是pre-softmax logits和注意力得分阵,最终期望得到的是$O = PV$.其中这里给定$Q,K,V,O \in R^{N \times d}$,$N$表示的是$sequence\_length$,$d$表示的是$head\_dimension$

根据online softmax的思路,可以描述出一个下图的2-pass的公式来计算Attention:

其中上图中的$d_i'$其实就是上述的$l_i(x)$,其中的$x_i$就是$Q$阵的第$k$行与$K^T$的第$i$行计算得到的元素,这个元素进行softmax后,作为标量$a_i$与$V$阵第$i$行进行相乘,累加到一个向量元素$o$上,最终这个值累加完后存放到结果$O$阵的第$k$行,这里需要解释的点其实就是$a_iV[i,:]$表示的含义(其实就是$P_k \in R^{1\times N}$,与$V \in R^{N \times d}$进行相乘得到的$O_k = P_kV$拆分后得到的结果),如下图所示,当i这个变量遍历完整个$sequence\_length$则底下$O^T$中的浅蓝块会变成深蓝块,即计算完毕,体现在公式则是$o_i$的更新到$o_N$:

这个2pass的Attention其实第二个循环中的值依赖于第一重循环中得到的$m_N$和$d'_N$,由此得出$a_i$,即对$x_i$做softmax后,与$V[i,:]$中对应的第i行相乘,从而更新$o_i$

如果是以$a_i$作为最后所求的值,则无法将2pass融合至1pass,但是我们所求的是$O[k,:] = o_N$,因此可以试着把中间不断迭代的$o_i$展开来看,有:

$$o_i = \sum_{j=1}^{i}{(\frac{e^{x_j-m_N}}{d'_N}V[j,:])}$$

将它分段来看,从而舍去整一行向量得到的$m_N$,$d'_N$,我们取到第i个元素的值为$m_i$,累和为$d'_i$,则原式子改写成:

$$o'_i = \sum_{j=1}^{i}{(\frac{e^{x_j-m_i}}{d'_i}V[j,:])}$$

经过裂项(裂出第i项,以寻找递推式)并整理可得:

$$o'_i = o'_{i-1}\frac{d'_{i-1}e^{m_{i-1}-m_i}}{d'_i} + \frac{e^{x_i-m_i}}{d'_i}V[i,:]$$

上面的递推式可以这么理解:将第i-1项得到的结果,将其softmax的原分母$d'_{i-1}$约去,更新内部每个元素减去的最大值,然后除于最新的累和值$d'_i$,完成前[1,i-1]的更新;并且加上此次计算得到的第i项的结果,得出[1,i]项的结果$o'_i$

根据上述递推式,可以使得Attention只需要1pass.过程如下:

Tiled Attention

这里的分块,是将QKVO进行分块,它们初始的shape都是$R^{N \times d}$,想要在1个block内的smem放下,这里假设了smem的大小是$M$,通过$B_c$,$B_r$的约束使得共享内存被最大程度的利用($/4d$是因为四个阵,它们都沿着N的方向分块,可以使得):

$K_j + V_j \lt \frac{M}{2}$

$Q_i + O_i = d * min(\frac{M}{4d},d) < \frac{M}{2}$

这样划分以尽可能多地利用共享内存

在代码实现中,由于传入的Q阵的tensor是有四个维度的数据:Tensor(b,nh,N,d)分别是batch_size,num_heads,sequence_length,head_dimension,

FlashAttentionV2

相较于V1的更新:

• 公式修改,以利用Tensor Core;
• 循环顺序调转,变成先切Q后切K,V

FlashAttentionV3

只有阅读,没有hopper

参考文件:

• [Attention优化][2w字]🔥原理&图解: 从Online-Softmax到FlashAttention V1/V2/V3
• 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
• flash-attn-minimal
• 猛猿flashattnv1-v2