vLLM——使用PagedAttention加速推理

news2024/9/24 21:17:40

参考自https://blog.vllm.ai/2023/06/20/vllm.html

介绍

vLLM是一个用于快速LLM推理和服务的开源库。vLLM 利用PagedAttention，可以有效地管理注意力键和值。PagedAttention 的 vLLM 重新定义了 LLM 服务的最新水平：它提供了比 HuggingFace Transformers 高 24 倍的吞吐量，而无需任何模型架构更改。

PagedAttention

在 vLLM 中，我们发现 LLM 服务的性能瓶颈在于内存。在自回归解码过程中，LLM 的所有输入标记都会产生其注意键和值张量，这些张量保存在 GPU 内存中以生成下一个标记。这些缓存的键和值张量通常称为 KV 缓存。KV 缓存是

大：LLaMA-13B中的单个序列占用高达1.7B
动态：其大小取决于序列长度，序列长度变化很大且不可预测。因此，有效管理 KV 缓存是一项重大挑战。我们发现现有系统由于碎片化和过度预留而浪费了**60% - 80%**的内存。

为了解决这个问题，我们引入了PagedAttention，这是一种注意力算法，其灵感来自操作系统中虚拟内存和分页的经典思想。与传统的注意力算法不同，PagedAttention 允许在非连续的内存空间中存储连续的键和值。具体来说，PagedAttention 将每个序列的 KV 缓存划分为块，每个块包含固定数量 token 的键和值。在注意力计算过程中，PagedAttention 内核会高效地识别和获取这些块。
在这里插入图片描述
由于块在内存中不需要连续，因此我们可以像在操作系统的虚拟内存中一样以更灵活的方式管理键和值：可以将块视为页面，将令牌视为字节，将序列视为进程。序列的连续逻辑块通过块表映射到非连续物理块。物理块在生成新令牌时按需分配。
在这里插入图片描述
在 PagedAttention 中，内存浪费仅发生在序列的最后一个块中。实际上，这会导致接近最佳的内存使用率，浪费率仅为 4% 以下。内存效率的提高被证明是非常有益的：它允许系统将更多序列批量处理在一起，提高 GPU 利用率，从而显著提高吞吐量，如上图性能结果所示。

PagedAttention 还有一个关键优势：高效的内存共享。例如，在并行采样中，同一个提示会生成多个输出序列。在这种情况下，输出序列之间可以共享提示的计算和内存。
在这里插入图片描述
PagedAttention 通过其块表自然地实现了内存共享。与进程共享物理页面的方式类似，PagedAttention 中的不同序列可以通过将其逻辑块映射到同一物理块来共享块。为了确保安全共享，PagedAttention 会跟踪物理块的引用计数并实现写时复制机制。
在这里插入图片描述
PageAttention 的内存共享功能大大降低了复杂采样算法（例如并行采样和波束搜索）的内存开销，最多可减少 55% 的内存使用量。这可以转化为高达 2.2 倍的吞吐量提升。这使得此类采样方法在 LLM 服务中变得实用。

$ pip install vllm

使用vLLM离线推理

from vllm import LLM

prompts = ["Hello, my name is", "The capital of France is"]  # Sample prompts.
llm = LLM(model="lmsys/vicuna-7b-v1.3")  # Create an LLM.
outputs = llm.generate(prompts)  # Generate texts from the prompts.

使用vLLM在线服务

 python -m vllm.entrypoints.openai.api_server --model lmsys/vicuna-7b-v1.3

可以使用与 OpenAI API 相同的格式查询服务器：

curl http://localhost:8000/v1/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "lmsys/vicuna-7b-v1.3",
        "prompt": "San Francisco is a",
        "max_tokens": 7,
        "temperature": 0
    }'