LLMs

主题切换

LLM（大语言模型）部署加速方法——PagedAttention篇 ​

来源：AiGC面试宝典 作者：宁静致远 日期：2023年09月29日

一、vLLM 用于大模型并行推理加速，存在什么问题？ ​

vLLM 用于大模型并行推理加速，其中核心改进是 PagedAttention 算法。在 vLLM 中，我们发现 LLM 服务的性能受到内存瓶颈的限制。在自回归解码过程中，LLM 的所有输入 token 都会生成其 key 和 value 张量，并且这些张量保存在 GPU 内存中以生成下一个 token。这些缓存的 key 和 value 张量通常称为 KV 缓存。KV 缓存的特点是：

  • 占用空间大：LLaMA-13B 中的单个序列最多占用 1.7GB
  • 动态变化：其大小取决于序列长度，序列长度变化很大且不可预测

因此，有效管理 KV 缓存提出了重大挑战。我们发现现有系统由于内存碎片和过度预留而浪费了 60% - 80% 的内存。

📝 通俗解释：想象你在一家餐厅工作，客人（token）陆陆续续来点菜。每个客人坐下后，你需要把所有点过的菜（key-value缓存）都记在脑子里。但是传统方法要求你把所有菜名写在一张很大的纸上，而且这张纸必须是一整块连续的空间——如果中途有人离开腾出中间的位置，这张纸也不能重复利用，导致大量浪费。PagedAttention就像允许你用很多小纸条来记菜名，每张小纸条可以放在不同位置，需要时再拼起来看，大大减少了浪费。

二、vLLM 如何优化大模型并行推理加速？ ​

vLLM 引入了 PagedAttention，这是一种受操作系统中虚拟内存和分页经典思想启发的注意力算法。

📝 通俗解释：这个创意其实来自计算机操作系统。操作系统管理内存时，不需要把一个程序的所有数据都放在连续的一大块内存里，而是可以把程序分成很多"页面"放在不同位置，需要时通过页表把逻辑地址映射到物理地址。PagedAttention把这个思想用到了AI模型处理文字上。

三、什么是 PagedAttention？ ​

与传统的注意力算法不同，PagedAttention 允许在不连续的内存空间中存储连续的 key 和 value。

📝 通俗解释：传统的注意力机制要求把一段文字的所有"记忆"整齐地排在一起，就像一本连续的书。但PagedAttention允许这些"记忆"分散存放在不同地方，就像图书馆里的一本书，内容被分散保存在不同书架上，但通过目录（块表）可以快速找到任何一页。

四、PagedAttention 如何存储连续的 key 和 value？ ​

具体来说，PagedAttention 将每个序列的 KV 缓存划分为块（Block），每个块包含固定数量 token 的 key 和 value。在注意力计算过程中，PagedAttention 内核有效地识别并获取这些块。

图一：PagedAttention 内容描述

图片展示了一个注意力机制的示意图：

  • 左侧有一个 Query vector 指向单词 "for"
  • 右侧是一个表格结构，标题为 Key and value vectors
  • 表格被划分为不同的块（Block），这些块在逻辑上是不连续的：
    • Block 1：包含单词 "computer", "scientist", "and", "mathematician"
    • Block 2：包含单词 "renowned", "for"
    • Block 0：包含单词 "Alan", "Turing", "is", "a"
  • 箭头从左侧的 "for" 指向右侧表格中的不同块，表示注意力机制在非连续内存块中查找对应的 key/value

因为块在内存中不需要是连续的，所以我们可以像在操作系统的虚拟内存中一样以更灵活的方式管理 key 和 value：可以将块视为页面（Page），将 token 视为字节（Byte），将序列视为进程（Process）。序列的连续逻辑块通过块表（Block Table）映射到非连续物理块。当新 token 生成时，物理块会按需分配。

📝 通俗解释：想象一个图书馆管理系统：

  • 每个"块（Block）"就像一个固定大小的书架
  • 每个"token"就像一本书
  • "序列"就像一个读者借阅的书籍清单
  • "块表"就像一个目录，记录了某本书（逻辑位置）实际放在哪个书架（物理位置）上

这样，读者不需要书都连续地放在一个书架上，可以分散存放，而且可以随时添加新书（生成新token）而不需要重新整理整个图书馆。

五、PagedAttention 技术细节 ​

  1. 内存浪费极低：在 PagedAttention 中，内存浪费仅发生在序列的最后一个块中。实际上，这会导致内存使用接近最佳，浪费率低于 4%。事实证明，内存效率的提高非常有益：它允许系统将更多序列一起批处理，提高 GPU 利用率，从而显著提高吞吐量。

  2. 高效的内存共享：PagedAttention 还有另一个关键优势——高效的内存共享。例如，在并行采样中，从同一提示（Prompt）生成多个输出序列。在这种情况下，提示的计算和内存可以在输出序列之间共享。

📝 通俗解释：

  • 关于内存浪费：就像用固定大小的盒子装东西，最后一个盒子可能没装满，但其他盒子都是满的，所以空间利用率非常高（超过96%）
  • 关于内存共享：就像老师给全班同学布置作文题，大家都要先抄题目。传统方法每个学生都要自己抄一遍题目，浪费时间和纸张。PagedAttention允许大家共享同一个题目内容，只需要抄一份就够了，省时又省力。

六、PagedAttention 如何实现安全共享？ ​

动机 ​

PagedAttention 中的不同序列可以通过将其逻辑块映射到同一物理块来共享块。这时候就涉及到如何安全共享的问题。

思路 ​

PagedAttention 跟踪物理块的引用计数（Reference Count），并实现 Copy-on-Write（写时复制）机制。

  • 当多个序列共享同一物理块时，系统会记录有多少序列正在使用这个块（引用计数）
  • 当某个序列需要修改共享块的内容时，系统会先检查引用计数：
    • 如果只有一个序列使用，则直接修改
    • 如果有多个序列使用，则先复制一份新的物理块给该序列，再进行修改

这样既保证了内存共享的效率，又避免了不同序列之间的数据冲突。

📝 通俗解释：这就像几个人共用一本画册：

  • 引用计数：就像记录有多少人正在借阅这本画册
  • 写时复制：当有人想在画册上画画时，如果只有他一个人借，就直接画；如果还有其他人借着看，就先复印一本新的，再在复印件上画

这样既节省了资源（不用每人一本），又保证了每个人都能修改自己的副本，不会互相干扰。

图三说明：多个输出采样的共享示例 ​

该图展示了两个序列（Seq A 和 Seq B）共享同一个物理 KV 缓存块的过程：

  • 中间表格（Physical KV cache blocks）：包含具体的 token 数据
  • 左侧表格（Seq A - Logical KV cache blocks）：逻辑块映射到中间物理块
  • 右侧表格（Seq B - Logical KV cache blocks）：逻辑块同样映射到中间相同的物理块

PagedAttention 的内存共享极大地降低了复杂采样算法的内存开销，例如并行采样和波束搜索，将其内存占用降低高达 55%。这可以将吞吐量提高高达 2.2 倍。

📝 通俗解释：比如让AI同时生成多个回答（并行采样），传统方法每个回答都要独立存储一份"记忆"，很占内存。现在有了共享机制，相同的开头部分大家共用一份内存，只有在开始"分道扬镳"产生不同内容时才复制新的内存，这样内存占用可以减少一半以上，生成速度也能快两倍多。

七、PagedAttention 源码解析 ​

PagedAttention 是 vLLM 背后的核心技术。vLLM 是 LLM 推理和服务引擎，支持各种模型的高性能推理。

以下是从 vLLM 源码中提取的 LLaMA 模型实现的核心部分，展示了 PagedAttention 是如何与 HuggingFace 模型进行集成的：

python
"""Inference-only LLaMA model compatible with HuggingFace weights.

The input of the model is flattened to a 1D tensor of tokens. The model uses
InputMetadata to extract the original 2D shape of the input.
"""
from typing import Dict, List, Optional, Tuple

import torch
from torch import nn
from transformers import LlamaConfig

from vllm.model_executor.input_metadata import InputMetadata
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.attention import PagedAttentionWithRoPE  # 核心：PagedAttention实现
from vllm.model_executor.layers.sampler import Sampler
from vllm.model_executor.weight_utils import (hf_model_weights_iterator,
                                              load_tensor_parallel_weights)
from vllm.model_executor.parallel_utils.parallel_state import (
    get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)
from vllm.model_executor.parallel_utils.tensor_parallel import (
    VocabParallelEmbedding, ColumnParallelLinear, RowParallelLinear)
from vllm.sequence import SequenceOutputs

KVCache = Tuple[torch.Tensor, torch.Tensor]  # Key-Value缓存类型定义

📝 通俗解释：这段代码展示了vLLM如何"借用"HuggingFace已经训练好的模型权重，但在推理（生成文字）时使用自己优化的PagedAttention技术。关键在于PagedAttentionWithRoPE这个类，它实现了高效的分页注意力机制，替代了传统的标准注意力实现。

总结 ​

PagedAttention 通过借鉴操作系统的虚拟内存分页思想，实现了：

  1. 高效的内存管理：将 KV 缓存分块存储，减少内存碎片和浪费（浪费率 < 4%）
  2. 灵活的内存共享：通过块表和引用计数实现安全的多序列共享
  3. 显著的性能提升：吞吐量提升高达 2.2 倍

这项技术是 vLLM 实现高性能 LLM 推理的核心基础。

📝 最终通俗解释：PagedAttention的本质是让AI模型在生成文字时，能像操作系统管理内存一样灵活地管理它的"记忆"。传统方法要求模型一次性预留大块连续内存，即使只用一点点也会浪费很多。PagedAttention把"记忆"切成小块，按需分配，用多少占多少，而且还能让多个任务共享相同的"记忆"部分。这样既省内存又能跑得更快，是大语言模型部署加速的关键技术。

基于 MIT 许可发布

版权所有 © 2024 LLMs