Flash Attention Softmax计算分块梳理

July-Second

已于 2024-09-03 22:44:01 修改

阅读量935

点赞数 16

分类专栏：人工智能文章标签： transformer

于 2024-09-03 22:36:49 首次发布

本文链接：https://blog.csdn.net/M_H_T__/article/details/141868507

版权

人工智能专栏收录该内容

2 篇文章

订阅专栏

Flash Attention Softmax分块计算梳理

在Transformer中,对于一个 $M\in R^{(2n,2n)}$ 的矩阵，Softmax是按行计算的。
假设 $M$ 中某一行 $X=[x_1,x_2,x_3,...,x_{2n}]$

经典Softmax

$softmax([x_1,x_2,...,x_{2n}])= \{{\frac {e^{x_i}}{\sum_j^{2n}e^{x_j}}}\}_{i=1}^{2n}$

半精度Softmax(safe_softmax)

为防止计算溢出，需要对指数进行缩放
$m=max([x_1,x_2,...x_{2n}])$
$safe\_softmax([x_1,x_2,...,x_{2n}])= \{{\frac {e^{x_i}/e^m}{\sum_j^{2n}(e^{x_j}/e^m)}}\}_{i=1}^{2n} = \{{\frac {e^{x_i-m}}{\sum_j^{2n}e^{x_j-m}}}\}_{i=1}^{2n}$

Flash Attention Softmax

首先将safe_softmax计算过程拆解

对于 $X=[x_1,x_2,x_3,...,x_{2n}]$

$ma x (X) = m$
$fun(X) = [e^{x_1-m},e^{x_2-m},...,e^{x_{2n}-m}]$
$s (X) = s u m (f u n (X))$
$saft\_softmax(X) = \frac{fun(X)}{s(X)}$

分块计算

将 $X=[x_1,x_2,x_3,...,x_{2n}]$ 划分为 $X_1=[x_1,x_2,x_3,...,x_{n}]$ 和 $X_2=[x_{n+1},x_{n+2},x_{n+3},...,x_{2n}]$

对 $X_1,X_2$ 分别开展计算：

分别计算 $X_1,X_2$ 的最大值记为 $m_1 = max(X_1)$ , $m_2 = max(X_2)$
$f_1 = fun(X_1)$ , $f_2 = fun(X_2)$
比较两组数据中的最大值： $m = max(m_1,m_2)$
根据最大值 $m$ 更新两组数据的计算结果 $fun(X) = [e^{m_1-m}f_1,e^{m_2-m}f_2]$
求和 $s (X) = s u m (f u n (X))$
$safe\_softmax(X) = \frac{func(X)}{s(X)}$

简单举例，对于 $X = [1, 2, 3, 4, 5, 6]$ :

$X_1=[1,2,3],X_2=[4,5,6]$
$m_1=max([1,2,3])=3,m_2=max(4,5,6)=6$
$f_1=[e^{(1-m_1)},e^{(2-m_1)},e^{(3-m_1)}] = [e^{-2},e^{-1},e^{0}] \\ f_2=[e^{(4-m_2)},e^{(5-m_2)},e^{(6-m_2)}] = [e^{-2},e^{-1},e^{0}]$
$m = max(m_1,m_2) = 6$
$fun(X)=[e^{m_1-m}f_1,e^{m_2-m}f_2]=[e^{-3}f_1,e^0f_2]\\ =[e^{-5},e^{-4},e^{-3},e^{-2},e^{-1},e^{0}]$
$s(X) = sum(fun(X)) = e^{-5}+e^{-4}+e^{-3}+e^{-2}+e^{-1}+e^{0}$
$safe\_softmax(X) = \frac{func(X)}{s(X)}$ 此步计算结果与safe_softmax一致