LLM 优化技术(2)——paged_attention 原理

1 从 flash attention 到 paged attention

1.1 Flash Attention 基本原理

Flash Attention 计算 self-attention 的关键是有效地硬件利用。GPU 的典型操作方式是使用大量的线程来执行一个操作，这个操作被称为内核。输入从HBM加载到寄存器和SRAM，并在计算后写回HBM。

Flash attention 的优化思路：

考虑到 self-attention 计算过程中，对 HBM 的重复读写是一个主要瓶颈。针对这个问题，Flash Attention 通过减少对 HBM 的访问次数来提高性能：

• 在不访问整个输入的情况下计算 softmax
• 不为反向传播存储巨大的中间 attention 矩阵

为此 FlashAttention 提出了两种方法来分布解决上述问题：tiing 和 recomputation：

• tiling，将输入分割成块，通过在输入块上进行多次传递来递增地计算 softmax
• recomputation，存储来自前向的 softmax 归一化因子，以便在反向中快速重新计算芯片上的 attention，这比从HBM读取中间矩阵的标准注意力方法更快

通过重新计算，会导致总的 Flops 增加，但是这样可以减少大量的 HBM 访问，从而有效提高 Flash Attention 的计算速度。

Flash Attention 的核心思想是分割输入，将它们从慢速HBM加载到快速SRAM，然后计算这些块的 attention 输出。在将每个块的输出相加之前，将其按正确的归一化因子进行缩放，从而得到正确的结果。

1.2 Tiling 与前向计算

分块计算注意力的关键部分是 softmax 的分块计算。向量的 softmax 可以计算为：

$$\begin{array}{rlr}& m(x)=max(x)& \\ & f(x)=[e^{x_1-m(x)},\cdots ,e^{x_B-m(x)}]& \\ & l(x)=\sum _{i}f(x{)}_i& \\ & softmax(x)=\frac{f(x)}{l(x)}& \end{array}$$

其中 $x$ 可以分解为 $x=[x^{(1)}{x}^{(2)}]\in {\mathbb{R}}^{2B},x^{(1)},x^{(2)}\in {\mathbb{R}}^B$，那么有：

$$\begin{array}{rlr}m(x)=& m([x^{(1)}{x}^{(2)}])& \\ =& max(m(x^{(1)}),m(x^{(2)}))& \end{array}$$

那么可以通过如下构造的方式，使得 $f(x)$ 的结果与分块前保持统一：

$$\begin{array}{rlr}f(x)=& [e^{m(x^{(1)})-m(x)}f(x^{(1)})e^{m(x^{(1)})-m(x)}f(x^{(1)})]& \\ =& [e^{m(x^{(1)})-m(x)}[e^{x_1^{(1)}-m(x^{(1)})}\cdots e^{x_B^{(1)}-m(x^{(1)})}]e^{m(x^{(2)})-m(x)}[e^{x_1^{(2)}-m(x^{(2)})}\cdots e^{x_B^{(2)}-m(x^{(2)})}]]& \\ =& [[e^{x_1^{(1)}-m(x)}\cdots e^{x_B^{(1)}-m(x)}][e^{x_1^{(2)}-m(x)}\cdots e^{x_B^{(2)}-m(x)}]]& \\ =& [e^{x_1-m(x)}\cdots e^{x_B-m(x)}]& \end{array}$$

$l(x)$ 的构造方式同理：

$$\begin{array}{rlr}l(x)=& l([x^{(1)}{x}^{(2)}])& \\ =& e^{m(x^{(1)})-m(x)}l(x^{(1)})+e^{m(x^{(2)})-m(x)}l(x^{(2)})& \end{array}$$

从而可以得到 $softmax$ 结果的形式，也可保持与分块前一致：

$$softmax(x)=\frac{f(x)}{l(x)}$$

softmax 的分块计算解决后，其他部分的分块计算就简单了。

完整的 decoder 的 self-attention 的计算过程：

$$\begin{array}{rlr}S=& \frac{Q{K}^T}{\sqrt{d}}\in {\mathbb{R}}^{N\times N}& \\ S^{masked}=& MASK(S)\in {\mathbb{R}}^{N\times N}& \\ P=& softmax(S^{masked})\in {\mathbb{R}}^{N\times N}& \\ p^{dropped}=& dropout(P,p_{drop})\in {\mathbb{R}}^{N\times N}& \\ O=& p^{dropped}V\in {\mathbb{R}}^{N\times d}& \end{array}$$

flash attention 的算法没有增加额外的计算量，只是将大块的操作拆分成多个小块逐步计算，因此算法的复杂度不变，仍为 $O(N^{2}d)$。另外，增加了中间变量 $l,m$，因此空间复杂度增加 $O(N)$。

1.3 Recomputation 与反向计算

Flash Attention 的后向传播需要 $S$ 和 $P$ 矩阵来计算 $Q,K,V$ 的梯度，空间复杂度是 $O(N^2)$，$S$ 和 $P$ 没有显式存储。

解决方法是使用输出 $O$ 和 $softmax$ 归一化统计 $(m,l)$，可以利用 SRAM 中的 $Q,K,V$ 重新计算 $S$ 和 $P$ 矩阵。这个过程会增加额外的计算量，但是减少了对 HBM 的访存，整体上加快了后向传播的速度。

标准 self-attention 的反向的求导过程：

注意力机制 $O=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$，反向传播时，需要根据损失函数 $\Phi$ 对模块输出的导数 $dO$ (即 $\frac{\partial \Phi }{\partial O}$)，求出其对输入的导数 $dQ,dK,dV$(即 $\frac{\partial \Phi }{\partial Q},\frac{\partial \Phi }{\partial K},\frac{\partial \Phi }{\partial V}$)。

令 $P=softmax(S),S=\frac{Q{K}^T}{\sqrt{d_k}},O=PV$，则其求导过程如下：

由 $\frac{\partial O}{\partial V}=P^T$ 得到：

$$\frac{\partial \Phi }{\partial V}=P^T\frac{\partial \Phi }{\partial O}$$

即：

$$dV=P^{T}dO$$

其余导数计算类似。

Flash Attention 将 $dV$ 的计算拆分成若干个子块 $d{V}_j$ 的计算，从而可以使用 Tiling 技术：

$$\begin{array}{rlr}d{V}_j=& (P^T{)}_{J}dO& \\ =& \sum _i(P^T{)}_{ji}d{O}_i& \\ =& \sum _i{p}_{ij}^{T}d{O}_i& \end{array}$$

$P$ 矩阵的大小是 $O(N^2)$，为了减小内存的消耗和对 HBM 的访存，FlashAttention 在反向传播时重新计算
$P$ 而非在前向传播时保存：

$$\begin{array}{rlr}{P}_{ij}=& (softmax(S){)}_{ij}& \\ =& \frac{e^{S_{ij}}-{\delta }^T(S_i^{max},C_N)}{{\delta }^T(S_i^{sum},C_N)}& \end{array}$$

其中 $S_i^{max},S_i^{sum}$ 是前向的时候保存的，且 $S_{ij}=\frac{Q_i(K^T{)}_j}{\sqrt{d_k}}=\frac{Q_i{K}_j^T}{\sqrt{d_k}}$

根据前面的推导，同理可得：

$$dP=dO{V}^T$$

考虑到 $s_i=\frac{q_i{K}^T}{\sqrt{d_k}},p_i=softmax(s_i)$，可以得到：

$$\begin{array}{rlr}d{q}_i=& d{s}_i\frac{\partial s_i}{\partial q_i}& \\ =& d{s}_i\frac{K}{\sqrt{d_k}}& \end{array}$$

$$\begin{array}{rlr}d{s}_i=& d{p}_i\frac{\partial p_i}{\partial s_i}& \\ =& d{p}_i(diag(p_i)-p_i^T{p}_i)& \\ =& d{p}_i\odot p_i-d{o}_i{o}_i^T{p}_i& \end{array}$$

拓展 $d{q}_i$ 到 $d{Q}_i$:

$$d{Q}_i=d{S}_i\frac{K}{\sqrt{d_k}}=\sum _{j}d{S}_{ij}\frac{K_j}{d_k}$$

计算 $d{S}_{ij}$：

$$\begin{array}{rlr}d{S}_i=& d{P}_i\odot P_i-\alpha ({\beta }^T(d{O}_i\odot O_i),N)\odot P_i& \\ =& (d{P}_i-\alpha ({\beta }^T(d{O}_i\odot O_i),N))\odot P_i& \\ & & \\ d{S}_{ij}=& (d{P}_{ij}-\alpha ({\beta }^T(d{O}_i\odot O_i),C_N))\odot P_{ij}& \end{array}$$

这里，$P_{ij}$ 是重计算得到的：

$$d{P}_{ij}=d{O}_i(V^T{)}_j=d{O}_i{V}_j^T$$

计算 $d{K}_j$:

$$\begin{array}{rl}dK=& (d{K}^T{)}^T=d{S}^T\frac{Q}{\sqrt{d_k}}\\ d{K}_j=& (d{S}^T{)}_j\frac{Q}{\sqrt{d_k}}=\sum _i(d{S}^T{)}_{ji}\frac{Q}{\sqrt{d_k}}\\ =& \sum _{i}d{S}_{ij}^T\frac{Q}{\sqrt{d_k}}\end{array}$$

1.4 PagedAttention

vLLM 主要用于快速 LLM 推理和服务，其核心是 PagedAttention，这是一种新颖的注意力算法，将在操作系统的虚拟内存中分页的经典思想引入到 LLM 服务中。在无需任何模型架构修改的情况下，可以做到比 HuggingFace Transformers 提供高达 24 倍的 Throughput。

1.4.1 PagedAttention 的基本原理

vLLM 的特点：

• 非常高的服务吞吐性能

• 使用 paged Attention 优化 KV cache 内存管理

• 动态 batch

• 优化的 cuda Kernel

• 各种 decoder 算法，包括并行采样、beam search 等

• 张量并行(TP)以支持分布式推理

• 流输出

• 兼容 OpenAI 的 API 服务

KV cache 占用巨大的显存，在 LLaMA-13B 中，缓存单个序列最多需要 1.7GB 显存；另外，KV cache 动态变化，：KV 缓存的大小取决于序列长度，这是高度可变和不可预测的。由于碎片化和过度保留，现有系统浪费了 60% - 80% 的显存。为解决这样的问题，paged Attention 出现了。

PagedAttention 是一种受操作系统中虚拟内存和分页经典思想启发的注意力算法。与传统的注意力算法不同，PagedAttention 允许在非连续的内存空间中存储连续的 key 和 value 。

具体来说，PagedAttention 将每个序列的 KV cache 划分为块，每个块包含固定数量 token 的键和值。在注意力计算期间，PagedAttention 内核可以有效地识别和获取这些块。

因为块在内存中不需要连续，因而可以用一种更加灵活的方式管理 key 和 value ，就像在操作系统的虚拟内存中一样：可以将块视为页面，将 token 视为字节，将序列视为进程。序列的连续逻辑块通过块表映射到非连续物理块中。物理块在生成新 token 时按需分配。

在 PagedAttention 中，内存浪费只会发生在序列的最后一个块中。这使得在实践中可以实现接近最佳的内存使用，仅浪费不到 4%。这种内存效率的提升被证明非常有用，允许系统将更多序列进行批处理，提高 GPU 使用率，显著提升吞吐量。

PagedAttention 还有另一个关键优势 —— 高效的内存共享。例如在并行采样中，多个输出序列是由同一个 prompt 生成的。在这种情况下，prompt 的计算和内存可以在输出序列中共享。

并行采样：

PagedAttention 自然地通过其块表格来启动内存共享。与进程共享物理页面的方式类似，PagedAttention 中的不同序列可以通过将它们的逻辑块映射到同一个物理块的方式来共享块。为了确保安全共享，PagedAttention 会对物理块的引用计数进行跟踪，并实现写时复制（Copy-on-Write）机制。

PageAttention 的内存共享大大减少了复杂采样算法的内存开销，例如并行采样和集束搜索的内存使用量降低了 55%。这可以转化为高达 2.2 倍的吞吐量提升。这种采样方法也在 LLM 服务中变得实用起来。

1.4.2 attention_ops.single_query_cached_kv_attention op

vLLM 的核心是 PagedAttention ，PagedAttention 核心是 attention_ops.single_query_cached_kv_attention op，完整代码vllm/tests/kernels/test_attention.py 。