从FlashAttention到FlashDecoding 原理详解

04290629

已于 2024-06-18 23:48:49 修改

阅读量1.6k

点赞数 25

分类专栏： LLM笔记文章标签： transformer 人工智能深度学习

于 2024-06-18 23:47:17 首次发布

本文链接：https://blog.csdn.net/weixin_45817413/article/details/139786333

版权

LLM笔记专栏收录该内容

3 篇文章

订阅专栏

FlashAttention 原理

Attention 计算 $softmax(Q_i, K_j^T)V_j$ ，

softmax 对某个 $q_i$ 与所有 $k_j$ 的 attention score( $q_i\cdot k_1,q_i\cdot k_2,...,q_i\cdot k_n$ )做归一化。

原生 softmax 计算

$softmax(x_1,x_2,...,x_n)=(\dfrac{e^{x_1}}{e^{x_1}+e^{x_2}+...+e^{x_n}},\dfrac{e^{x_2}}{e^{x_1}+e^{x_2}+...+e^{x_n}},...,\dfrac{e^{x_n}}{e^{x_1}+e^{x_2}+...+e^{x_n}})$

由于指数计算容易溢出，safe-softmax 相对于原生的 softmax 计算在每个指数项都减去一个 max 值。

$max=max(x_1,x_2,...,x_n)$

$safe-softmax(x_1,x_2,...,x_n)=(\dfrac{e^{x_1-max}}{e^{x_1-max}+e^{x_2-max}+...+e^{x_n-max}},\dfrac{e^{x_2-max}}{e^{x_1-max}+e^{x_2-max}+...+e^{x_n-max}},...,\dfrac{e^{x_n-max}}{e^{x_1-max}+e^{x_2-max}+...+e^{x_n-max}})$

由于 $ma x$ 需要全局信息，需要遍历所有 $X$ 结果以后得到，因此原生的 safe softmax 计算需要遍历三遍 $X$ 。如果 SRAM 没有足够的空间存储 $X$ ，那么需要每次都取出对应的 $q_i,k_j$ 来 recompute $x_{j}$ ，因此 I/O 开销大。

原生3-pass safe softmax

为此，Online softmax 改进了 safe softmax，将 3 个 pass 减为 2 个 pass，减少了一遍 I/O。Online softmax 的原理是在遍历 $x_i$ 的过程中动态更新当前遇到过的局部最大值 $m_{i}=max(x_1,x_2,...,x_i)$ ，在每轮迭代时通过修正将前面用到的局部最大值 $m_{i-1}$ 代替为 $m_{i}$ ，修正方式是把上一轮的结果乘以因数 $e^{m_{i-1}-m_{i}}$ 。

2-pass Online Softmax

基于上面online的思想，FlashAttention 把 Attention 计算优化为 1 pass。

在 Attention 计算中，某个 $q$ 向量与所有 $k_j,v_j$ 计算得到结果向量 $o$

$o=\sum_{j=1}^{n}\dfrac{e^{x_j-m_n}}{d^{'}_{n}}v_j$

在遍历时，定义第 $i$ 轮的结果为：

$o_i=\sum_{j=1}^{i}\dfrac{e^{x_j-m_i}}{d^{'}_{i}}v_j$

当第 n 轮计算完成时， $o_n=o$ 。第 i 轮的结果与第 $i - 1$ 轮的结果之间的关系如下，因此可以达到 1 pass 就计算出 $o$ 的效果。

FlashAttention 中间结果的迭代关系

在此基础上，FlashAttention-2 采用分块计算，把 $K, V$ 在 sequence length 维度上分成 $T_c$ 个大小相同的 $K_i,V_i$ 块，把 $Q$ 在 sequence length 维度上分成 $T_r$ 个大小相同的 $K_i,V_i$ 块。

FlashAttention2 计算过程示意图

计算时，先 load $Q_i$ ，然后内层循环每轮计算 load 一对 $K_i,V_i$ 块。这样每个 $Q_i$ 只需要 load 一次，并且不同的 $Q_i$ 之间可以没有干扰地做并行计算。

在这里插入图片描述

FlashDecoding 原理

FlashAttention 在 $Q$ 的 sequence length 维度上并行，因此在 prefill 阶段可以良好地并行，但是在 decode 阶段只有上一轮新生成的 token 的 $q$ ，因此无法在 $Q$ 的 sequence length 维度上并行。为此，Flash-Decoding 让 LLM 在 decode 阶段能够在 $K, V$ 的 sequence length 维度上并行。FlashDecoding 将 $K, V$ 分块，每个块分别与 $q$ 做 FlashAttention，分块之间可以并行计算，分块得到的结果通过 reduce 得到最终结果。Reduce 时，只需要用到每个分块计算出的 $o_j(1\times D),l_j(1\times 1),m_j(1\times 1)$ ，Reduce 操作的时间复杂度非常低： $O(D\cdot T_c)$ 。

FlashDecoding：分块计算+Reduce

FlashAttention 和 FlashDecoding 在 vLLM 中的应用

vLLM 在 prefill 阶段调用 xformers 的 flash-attn 后端，间接使用 FlashAttention 进行推理；在 decode 阶段，当单个 GPU 上 Attention Head 较少时，此时 Head 维度上的并行性低，因此 vLLM 的 PagedAttention V2 使用了 FlashDecoding 的思路进行优化，在 $K, V$ 的 sequence length 维度上增加并行性。

“PagedAttention V2 (#1348) implements a similar idea to boost the performance when the batch size or the number of attention heads per GPU is small.”
Does vLLM support flash attention? · vllm-project/vllm · Discussion #425
does vllm use Flash-Decoding? · Issue #1362 · vllm-project/vllm