当前位置：首页 > news >正文

flash attention2 计算过程的探索和学习

news 2025/9/24 11:18:46

之前初步探索了flash attention相比标准attention性能提升的原因。

https://blog.csdn.net/liliang199/article/details/151789333

这里尝试进一步探索，解析flash attention2相比flash attention的性能提升的原因。

首先了解稳定版softmax、以及基于稳定版softmax的标准attention，然后参考论文，使用公式表示flash attention计算的计算过程，进而引出flash attention2的计算过程。

1 稳定版softmax

考虑到数值过大，指数运算后可能会溢出，针对如下所示的softmax计算过程，一般会减去x中最大元素的值。

处理前后的softmax公式如下

$softmax(x) = \frac{e^{x_i}}{\sum e^{x_j}} = \frac{e^{x_i-x_{max}}}{\sum e^{x_j-x_{max}}}$

由于是减去x_max后做指数运算，相当于分子和分母同时除以exp(x_max)，不改变softmax特性。

2 标准attention

softmax是attention最核心的运算，这里设定seq_len=2，查询Q、键K=[K1, K2]和值V。

针对K1和K2，S1和S2表示如下。

$S^{(1)} = Q (K^{(1)})^{T}, S^{(2)} = Q (K^{(2)})^{T}$

基于稳定版softmax，attention计算过程如下。

$m = max(rowmax(S^{(1)}), rowmax(S^{(2)})) \in \mathbb{R}^{B_r} \\ l = rowsum(e^{S^{(1)}-m}) + rowsum(e^{S^{(2)}-m}) \in \mathbb{R}^{B_r} \\ P = \left[P^{(1)}, P^{(2)} \right] = diag(l)^{-1} \left[e^{S^{(1)}-m}, e^{S^{(2)}-m}\right] \in \mathbb{R}^{B_r \times 2B_c} \\ O = \left[P^{(1)}, P^{(2)} \right] \begin{bmatrix} V^{(1)}\\V^{(2)} \end{bmatrix} \\= diag(l)^{-1} (e^{S^{(1)}-m} V^{(1)} + e^{S^{(2)}-m} V^{(2)}) \in \mathbb{R}^{B_r \times 2B_c}$

3 flash attention

flash attention的计算过程示例如下

$\displaystyle m^{(1)} = rowmax(S^{(1)}) \in \mathbb{R}^{B_r} \\ l^{(1)} = rowsum(e^{S^{(1)}-m^{(1)}}) \in \mathbb{R}^{B_r} \\ \hat{P} ^{(1)} = diag(l^{(1)})^{-1} e^{S^{(1)}-m^{(1)}} \in \mathbb{R}^{B_r \times B_c} \\O^{(1)} = \hat{P} ^{(1)} V^{(1)} = diag(l^{(1)})^{-1}e^{S^{(1)}-m^{(1)}}V^{(1)} \in \mathbb{R}^{B_r \times d} \\ m^{(2)} = max(m^{(1)}, rowmax(S^{(2)})) = m \\ l^{(2)} = e^{m^{(1)} - m^{(2)}} l^{(1)} + rowsum(e^{S^{(2)}-m^{(2)}}) \\ \hat{P} ^{(2)} = diag(l^{(2)})^{-1} e^{S^{(2)}-m^{(2)}} \\ O^{(2)} = diag(l^{(1)}/l^{(2)}) diag(e^{m^{(1)}-m^{(2)}})O^{(1)} + \hat{P} ^{(2)} V^{(2)} \\= diag(l^{(2)})^{-1} e^{S^{(1)}-m}V^{(1)} + diag(l^{(2)})^{-1} e^{S^{(2)}-m} V^{(2)} = O$

如果不考虑减去最大值，flash attention的计算过程如下所示。

4 flash attention2

flash attention在计算的过程中，每计算一个块，需要使用最新m更新l。

参考softmax公式，l主要是对一化。

flash attention2采用了归一化后移的处理方法，即在计算过程中暂时不做归一化，将其放到最后。

计算过程示例如下。

$\displaystyle m^{(1)} = rowmax(S^{(1)}) \in \mathbb{R}^{B_r} \\ l^{(1)} = rowsum(e^{S^{(1)}-m^{(1)}}) \in \mathbb{R}^{B_r} \\ \hat{P} ^{(1)} = diag(l^{(1)})^{-1} e^{S^{(1)}-m^{(1)}} \in \mathbb{R}^{B_r \times B_c} \\ \hat{O}^{(1)} = e^{S^{(1)}-m^{(1)}}V^{(1)} \in \mathbb{R}^{B_r \times d} \\ m^{(2)} = max(m^{(1)}, rowmax(S^{(2)})) = m \\ l^{(2)} = e^{m^{(1)} - m^{(2)}} l^{(1)} + rowsum(e^{S^{(2)}-m^{(2)}}) =l \\ \hat{P} ^{(2)} = diag(l^{(2)})^{-1} e^{S^{(2)}-m^{(2)}} \\ \hat{O}^{(2)} = diag(e^{m^{(1)}-m^{(2)}})\hat{O}^{(1)} + \hat{P} ^{(2)} V^{(2)} \\= e^{S^{(1)}-m}V^{(1)} + e^{S^{(2)}-m} V^{(2)} \\O^{(2)} = diag(l^{(2)})^{-1} \hat{O}^{(2)} = O$

这种处理方式，有效减少了如此可降低计算复杂程度，减少了中间数据P1的存储需求。

以下是flash attention2的算法示例

附录

1 diag的含义

diag(a，b，c…)表示一个对角矩阵，即除主对角线外的元素都为0，diag有如下特性

$diag(l)^{-1} = diag(l^{-1})$

示例程序如下

import numpy as np
a = np.arange(1, 4)
diag_a = np.diag(a)
diag_a_reverse = 1 / np.diag(a)
print(f"diag_a: {diag_a}\ndiag_a_reverse: {diag_a_reverse}")

输入如下