FlashAttention 原理之伪代码解释

FlashAttention 原理之伪代码解释

$$\mathrm{Attn}(Q,K,V)=\underset{\text{大小为 }N\times N}{\underbrace{\mathrm{softmax}\left(Q{K}^T\right)}}\times V,$$

其中 $Q,K,V\in {\mathbb{R}}^{N\times d}$，$N$ 常常很大，直接构造和存储 $\mathrm{softmax}(Q{K}^T)$ 需要 $O(N^2)$ 的空间，并且读写代价很高。

FlashAttention 通过将 Q, K, V 分别切成小的块（Block）并在片上(SRAM)完成大部分计算，从而避免在外部大内存(HBM)中存储或频繁读写大规模中间矩阵。算法中用到了按行块切的 query（大小为 $B_r\times d$）和按列块切的 key/value（大小为 $B_c\times d$），然后在两重循环中分块累加结果。

输入：

• 矩阵 $Q,K,V\in {\mathbb{R}}^{N\times d}$ 存放于外部大内存(HBM)。
• 片上(on-chip)有大小为 $M$ 的高速缓存(SRAM)。
• 还会在HBM/片上分别为中间变量分配空间，如输出 $O\in {\mathbb{R}}^{N\times d}$、部分归一化因子 $\ell \in {\mathbb{R}}^N$、以及部分最大值 $m\in {\mathbb{R}}^N$（用于数值稳定）。

切分：

1. 将 Q 按“行方向”分块，得到 $T_r=\lceil \frac{N}{B_r}\rceil$ 个子块 $Q_1,\dots ,Q_{T_r}$，每块大小约为 $B_r\times d$。
2. 将 K, V 按“列方向”分块，得到 $T_c=\lceil \frac{N}{B_c}\rceil$ 个子块 $(K_1,V_1),\dots ,(K_{T_c},V_{T_c})$，每块大小约为 $B_c\times d$。
3. 相应地，也把输出 O、部分归一化因子 $\ell$、和部分最大值 $m$ 按同样的行块大小 $B_r$ 切分成 $O_1,\dots ,O_{T_r},{\ell }_1,\dots ,{\ell }_{T_r},m_1,\dots ,m_{T_r}$。

双层循环：

• 外层循环 (第 5 行)：遍历 Key/Value 的列块索引 $j=1\dots T_c$。
• 将第 $j$ 个块 $(K_j,V_j)$ 从外部内存(HBM)加载到片上SRAM。
• 内层循环 (第 7 行)：遍历 Query 的行块索引 $i=1\dots T_r$。
• 将 $Q_i,O_i,{\ell }_i,m_i$ 从 HBM 加载到片上SRAM（它们都是大小为 $B_r$ 或 $B_r\times d$ 级别）。
• 计算局部注意力得分 $S_{ij}=Q_i{K}_j^T$ （大小为 $(B_r\times d)\times (B_c\times d)\to B_r\times B_c$）。
• 做数值稳定的行级最大值 $\tilde{m}ij=\mathrm{rowmax}(Sij)$，再计算 $\tilde{P}ij=\exp (Sij-{\tilde{m}}_{ij})$。
• 行和 ${\ell }_{ij}=\mathrm{rowsum}\left({\tilde{P}}_{ij}\right)$。
• 用 $\tilde{m}ij$ 和 $\ell ij$ 来更新行块级的最大值 $m_i^{new}$ 和部分归一化因子 ${\ell }_i^{new}$，并将新的输出块 $O_i$ 与新块 ${\tilde{P}}_{ij}{V}_j$ 合并。

计算$O_i$的结果

$$m_i^{new}=\max (m_i,{\tilde{m}}_{ij}),{\ell }_i^{new}=e^{m_i-m_i^{new}}{\ell }_i+e^{{\tilde{m}}_{ij}-m_i^{new}}{\ell }_{ij}.$$
若 ${\tilde{m}}_{ij}$ 比原先 $m_i$ 大，就把整条“数值刻度”往上挪到 $\tilde{m}ij$；否则保留原先的刻度 $m_i$。
${\ell }_i$ 和 ${\ell }_{ij}$ 都要根据新的刻度 $m_i^{new}$ 进行相应的指数缩放，然后相加得到新的部分归一化因子。

现在我们得到了新的刻度 $m_i^{new}$ 和新的部分归一化因子 ${\ell }_i^{new}$。剩下的，就是把老的输出 $O_i$ 和当前块新增的输出 ${\tilde{P}}_{ij}{V}_j$ 合在一起，放到一个统一的“刻度”下。第 12 行的写法大致如下（作一点数学化展开）：

$$\begin{array}{rl}{O}_i^{(\mathrm{new})}& =\underbrace{[\mathrm{diag}({\ell }_i^{new}){]}^{-1}}\text{相当于每行除以 }{\ell }_i^{new}(\underbrace{\mathrm{diag}({\ell }_i)e^{m_i-m_i^{new}}{O}_i}\text{"重刻度"后的旧输出}+\underset{\text{“重刻度”后的新块输出}}{\underbrace{e^{\tilde{m}ij-m_i^{new}}\tilde{P}ij{V}_j}}).\end{array}$$

旧输出 $O_i$ 的重刻度：

1. 先将 $O_i$ 乘上原先“局部归一化”用到的因子 ${\ell }_i$，并且再乘上 $e^{m_i-m_i^{new}}$，把它转移到新的刻度 ${\ell }_i^{new}$ 所在的坐标系。
2. 新块的输出 ${\tilde{P}}_{ij}{V}_j$：
   这是在本次块的局部最大值 $\tilde{m}ij$ 上做的 exponent，所以要乘上 $e^{\tilde{m}ij-m_i^{new}}$ 去对齐到新的刻度。
3. 再除以新的 ${\ell }_i^{new}$：
   由于我们想让最终的 $O_i$ 处在 $“输出=\frac{1}{\text{(部分Softmax分母)}}\times (加权和)”$ 的形式，所以最后整体还要除以 ${\ell }_i^{new}$。
   代码里用 $[\mathrm{diag}({\ell }_i^{new}){]}^{-1}$ 的写法，是因为在实现中 $O_i$ 维度是 $B_r\times d$，而 ${\ell }_i^{new}$ 是长度为 $B_r$ 的向量，需要对每个行分别做除法。

最终把这个新的 $O_i^{(\mathrm{new})}$ 写回HBM，并且更新 ${\ell }_i\leftarrow {\ell }_i^{new}$、$m_i\leftarrow m_i^{new}$（见第 13 行），就完成了对第 $j$ 块 Key/Value 的处理。

小结: 第 12 行就是做了一次“加权合并”，把“老的注意力累加输出”和“新的块注意力输出”整合到同一个数值基准，并更新归一化，使得当所有块都处理完后，O_i 就等价于 $\mathrm{softmax}(Q_i{K}^T)V$ 中“i 行块”的那部分结果。

FlashAttention公式推理：https://gordicaleksa.medium.com/eli5-flash-attention-5c44017022ad