【大模型理论篇】自注意力机制计算加速工程优化技巧

源泉的小广场

已于 2024-09-22 18:13:45 修改

阅读量992

点赞数 26

分类专栏：大模型文章标签：大模型性能优化大模型内存优化大模型计算加速 llm 分页注意力 flashattention

于 2024-08-28 10:46:15 首次发布

本文链接：https://blog.csdn.net/weixin_65514978/article/details/141609201

版权

大模型专栏收录该内容

35 篇文章 8 订阅

订阅专栏

自注意力机制的工程优化一直是提高Transformer模型效率的关键部分，特别是在处理长序列时。Transformer架构的扩展在很大程度上受到自注意力机制的瓶颈限制，自注意力机制具有二次增长的时间和内存复杂度，这点我们在《Transformer KV Cache原理深入浅出》中进行了分析，可以参考。因此，业内提出了不同的优化方案，包括采用《LLaMA3结构关键模块分析》文章中提到的分组查询注意力（GQA）以及KV Cache技术。除此之外，本文将继续分享针对自注意力机制的工程优化方法，包括Flash Attention和PagedAttention。

1. Flash Attention

FlashAttention【1，2，3】是一种提高Transformer模型效率的方法。针对提高注意力操作的效率，模型能够以更快的效率训练和推理。FlashAttention通过重新排列注意力计算顺序，并利用分块和重计算来显著加速计算的算法，它将内存使用量从序列长度的二次方降低到线性。

Flash Attention 不是在每次计算迭代中多次加载查询、键和值或中间计算结果，而是仅加载一次所有数据（查询、键和值）。并且使用分块技术，将输入数据从 HBM（GPU 内存）加载到 SRAM（高速缓存），然后针对该块执行注意力计算，并在 HBM 中更新输出。通过避免将大型中间注意力矩阵写入 HBM，减少了内存读写的次数，从而实现了 2-4 倍的时钟时间加速。通过策略性地最小化不同类型内存之间的数据来回传输，Flash Attention 能够显著优化资源利用率。关键策略包括 “内核融合”，它将多个计算步骤合并为一个操作，减少了重复数据传输的需求，从而降低了开销。另一个关键策略是 “分块”，它涉及将输入数据分成较小的块以促进并行处理，优化内存使用，为具有更大输入大小的模型提供了可扩展的解决方案。

高带宽内存（HBM）提供大的内存容量，但处理速度较慢。静态随机存取存储器（SRAM）是一种提供快速数据访问的存储器类型，处理速度更快，但与 HBM 相比通常容量有限。

在标准注意力机制中，HBM 用于存储、读取和写入注意力计算中使用的键、查询和值。然而，注意力计算中涉及的操作经常导致 HBM 和片上 SRAM 之间频繁的数据传输。例如，在计算过程中，键、查询和值从 HBM 加载到片上 SRAM 进行处理，中间结果和最终输出在注意力机制的每个步骤后都写回到 HBM。HBM 和 SRAM 之间频繁的数据移动由于在数据传输和处理上花费的时间而导致高开销【4】。带着下述公式【5】来理解：

这里展示 FlashAttention 前向传播的示意图：通过分块和 softmax 重缩放，以块为单位操作，避免了从 HBM 进行读写，同时无需近似便可获得正确的输出。

Flash Attention 通过减少冗余的读和写来优化 HBM 和片上 SRAM 之间的数据移动。Flash Attention 不是为每个单独的注意力步骤执行这些操作，而是仅加载一次键、查询和值，合并或 “融合” 注意力机制的操作，然后将结果写回内存，减少总体计算开销并提高处理效率。

FlashAttention算法步骤【1】：

输入要求：

矩阵 Q, K, V 位于 HBM（高带宽内存）中，大小为 $N \times d$ 。
芯片上 SRAM（高速缓存）大小为 M。

步骤1：设置块大小

设置块大小 $B_c = \left\lceil \frac{M}{4d} \right\rceil$ 和 $B_r = \min \left(\left\lceil \frac{M}{4d} \right\rceil, d\right)$ 。
这里 $B_c$ 和 $B_r$ 代表分块后的大小。
这一步确保分块的大小适合芯片内存，并最大化使用高速缓存。

步骤2：初始化输出和中间变量

初始化 $O = 0_{N \times d}$ 在 HBM 中。
初始化 $\ell = 0_{N}$ 和 $m = (-\infty)_{N}$ 在 HBM 中。
O 是最终输出矩阵。
$\ell$ 和 m 是用于Softmax缩放和数值稳定的辅助变量。

步骤3：对 Q, K, V 进行分块

将 Q 分成 $T_r = \left\lceil \frac{N}{B_r} \right\rceil$ 个块 $Q_1, \dots, Q_{T_r}$ ，每块大小为 $B_r \times d$ 。
将 K 和 V 分成 $T_c = \left\lceil \frac{N}{B_c} \right\rceil$ 个块 $K_1, \dots, K_{T_c}$ 和 $V_1, \dots, V_{T_c}$ ，每块大小为 $B_c \times d$ 。

步骤4：对 O 进行分块

将 O 分成 $T_r$ 个块 $O_1, \dots, O_{T_r}$ ，每块大小为 $B_r \times d$ 。
将 $\ell$ 分成 $T_r$ 个块 $\ell_1, \dots, \ell_{T_r}$ ，每块大小为 $B_r$ 。
将 m 分成 $T_r$ 个块 $m_1, \dots, m_{T_r}$ ，每块大小为 $B_r$ 。

步骤5-6：加载 $K_j$ 和 $V_j$ 到芯片上的 SRAM

对于 $1 \leq j \leq T_c$ ，从 HBM 中将 $K_j, V_j$ 加载到芯片上的 SRAM。

步骤7-14：对于每个 $Q_i$ ，计算注意力得分和输出

对于 $1 \leq i \leq T_r$ ，执行以下操作：
将 $Q_i, O_i, \ell_i, m_i$ 从 HBM 加载到芯片上的 SRAM。
在芯片上计算注意力得分矩阵： $S_{ij} = Q_i K_j^T$ ，大小为 $B_r \times B_c$ 。
计算 $\tilde{m}_{ij}$ （当前块的最大值）和 $\tilde{P}_{ij}$ （softmax计算的中间结果）。
更新数值稳定性变量 $m_{\text{new}} = \max(m_i, \tilde{m}_{ij})$ 和 $\ell_{\text{new}}$ 。
将结果 O_i 和 $\ell_i, m_i$ 写回 HBM。

步骤15：返回最终输出

返回最终的结果 O。

关于这里为什么采用 $\left\lceil \frac{M}{4d} \right\rceil$ ,是因为查询（query）、键（key）和值（value）向量都是 d 维的，并且还需要将它们组合成输出的 d 维向量。所以，这个大小基本上允许在使用查询、键、值和输出向量时最大化使用SRAM的容量【5】。

一个简单的例子：假设M = 10000，d = 50。在这个例子中，块的大小是 $\frac{10000}{4 \times 50} = 50$ 。所以在这个例子中，每次会加载50个查询、键、值和输出向量的块，以确保减少HBM和SRAM之间的读写次数。

2. PagedAttention

大模型在生成过程中面临内存限制。在生成的解码部分，为先前token生成的所有注意力键和值都存储在 GPU 内存中以供重复使用。这被称为 KV Cache，在《Transformer KV Cache原理深入浅出》有做过详细说明。所谓KV Cache，这里再简单解释下，在基于Transformer的大模型解码过程中，随着每个输入token被处理，模型生成相应的注意力键和值张量。这些键和值张量编码了关于当前输入及其与上下文关系的重要信息。模型不是在解码过程的每个步骤中从头重新计算这些与注意力相关的张量，而是将它们存储在 GPU 内存中。

像 GPT - 4 这样的大模型可以有数万亿个参数，使得它们能力非常强大，但在服务期间进行推理时也极其消耗内存。内存的主要瓶颈是由于 KV 缓存。对于大模型和长序列，会占用大量内存。对于一个具有 130 亿参数的大模型，KV 缓存的大小可以达到占比30%以上。另外，KV 缓存的大小也是动态且不可预测的，取决于输入序列的长度。

此外，现有系统中的 KV 缓存内存管理，存在三种类型的内存浪费 —— 预留内存、内部碎片和外部碎片。每个内存槽中的token代表其 KV 缓存。相同的token在不同位置时可以具有不同的 KV 缓存。

因此，Kwon 等人【6】提出PagedAttention技术，大幅减少大模型 KV 缓存的内存占用，从而使得大模型更高效利用内存且易于访问。PagedAttention背后的核心思想是将每个序列的 KV 缓存分割成更小、更易于管理的 “页” 或块。每个块包含固定数量token的键值向量。这样，在注意力计算期间可以更高效地加载和访问 KV 缓存。

分页注意力算法的图示，其中注意力键和值向量作为非连续的块存储在内存中

PagedAttention管理用于存储键和值向量的内存的方法类似于操作系统处理虚拟内存的方式。分页注意力将数据的逻辑组织与其物理存储解耦【8，9】。属于一个序列的逻辑块使用块表映射到可能不连续的物理块。有利于内存利用，在生成新token时可以分配新的物理块。 PagedAttention通过将 KV 缓存分割成通过查找表访问的块来优化内存使用。因此，KV 缓存不需要存储在连续的内存中，并且可以根据需要分配块。内存效率可以提高内存受限工作负载下的 GPU 利用率，因此可以支持更多的推理批次。使用查找表访问内存块也有助于进行 KV 共享，对于并行采样有帮助，在并行采样中，对于相同的提示同时生成多个输出。在这种情况下，缓存的 KV 块可以在多次生成之间共享。

PagedAttention处理流程：

分割 KV 缓存：将 KV 缓存分成固定大小的块或 “页”。每个块包含原始缓存中的一部分键值对。
构建查找表：构建并维护一个查找表，该表将查询键映射到存储相应值的特定页。这个表将每个可能的查询键映射到存储相应值的特定页。实现快速分配和检索。
选择性加载：在推理期间，模型仅加载处理当前输入序列所需的页。与加载整个 KV 缓存相比，可以大幅减少总体内存占用。
注意力计算：加载相关页后，模型可以像之前一样进行注意力计算，根据需要从加载的页中访问键值对。

PagedAttention已被 vLLM 采用，vLLM 是由加州伯克利开发的用于快速 LLM 推理和服务的开源库。vLLM 支持在线服务和离线推理，可以通过Python 库直接在 Python 脚本中使用。

pip install vllm

这里参考【10】给出llama的使用示例：

from vllm import LLM
from transformers import AutoTokenizer

# 加载 Llama 3 8B model 到 vLLM
model = LLM("meta-llama/Meta-Llama-3-8B-Instruct")
# 模型推理
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
messages = [{"role": "user", "content": "What is the capital of France?"}]
formatted_prompt =  tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
output = model.generate(formatted_prompt)
print(output)

PagedAttention的优势与挑战：

好处：
接近最优的内存利用率，在分页注意力中，内存浪费主要局限于输入序列的最后一个块。接近最优的内存使用，平均只有大约 4% 的浪费。这种内存效率的大幅提升使系统能够将更多序列批处理在一起，从而提高 GPU 利用率并显著增加总体吞吐量。

挑战：

查找表的开销，用于将查询键映射到 KV 缓存页的查找表在推理期间会增加一些计算开销。模型需要在加载相关页之前执行此查找，如果实现效率不高，可能会影响推理速度。根据具体的实现和页的大小，可能会引入一些延迟。例如，如果页太小，频繁的页加载和卸载可能会引入延迟。