vLLM 与 PagedAttention 详解：显存分页的高吞吐推理框架

Posted on | In , | 1.3k | 5 分钟

介绍 vLLM 这一开源大模型推理加速框架：以 PagedAttention 高效管理注意力缓存，把吞吐量做到 HuggingFace Transformers 的 14 到 24 倍，并给出基于 Qwen1.5-7B-Chat 的推理代码示例。

GitHub | Documentation | Paper | Blog

简介

vLLM是一个开源的大模型推理加速框架，通过PagedAttention高效地管理attention中缓存的张量，实现了比HuggingFace Transformers高14-24倍的吞吐量。

vLLM 的吞吐量比 HF 高 14 倍 - 24 倍，比 TGI 高 2.2 倍 - 2.5 倍

PagedAttention 是 vLLM 的核心技术，它解决了LLM服务中内存的瓶颈问题。传统的注意力算法在自回归解码过程中，需要将所有输入Token的注意力键和值张量存储在GPU内存中，以生成下一个Token。这些缓存的键和值张量通常被称为KV缓存。

主要特性

  • 通过PagedAttention对 KV Cache 的有效管理
  • 传入请求的continus batching，而不是static batching
  • 支持张量并行推理
  • 支持流式输出
  • 兼容 OpenAI 的接口服务
  • 与 HuggingFace 模型无缝集成

    PagedAttention

PagedAttention，这是一种受操作系统中虚拟内存和分页的经典思想启发的注意力算法。与传统的注意力算法不同，PagedAttention 允许在不连续的内存空间中存储连续的键和值。具体来说，PagedAttention 将每个序列的 KV 缓存划分为块，每个块包含固定数量 token 的键和值。在注意力计算过程中，PagedAttention 内核有效地识别并获取这些块。

⚠️ 图片缺失（Notion 导出时未包含）：PagedAttention：KV Cache 被划分为块。块在内存空间中不需要是连续的（notion_ac1263e5f7.gif）

因为 blocks 在内存中不需要是连续的，所以我们可以像在操作系统的虚拟内存中一样以更灵活的方式管理键和值：可以将 blocks 视为页面，将 token 视为字节，将 sequences 视为进程。序列的连续logical blocks *通过块表映射到非连续 *physical blocks。当 tokens 生成时，物理块会按需分配。

KV cache

⚠️ 图片缺失（Notion 导出时未包含）：使用 PagedAttention 的请求生成过程示例（notion_8ce2253278.gif）

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

PagedAttention 还有另一个关键优势：高效的内存共享。例如，在 parallel sampling 中，从同一提示生成多个输出序列。在这种情况下，提示的计算和内存可以在输出序列之间共享。

⚠️ 图片缺失（Notion 导出时未包含）：并行采样示例（notion_dbf5eba4aa.gif）

PagedAttention 自然可以通过其块表实现内存共享。与进程共享物理页的方式类似，PagedAttention 中的不同序列可以通过将其逻辑块映射到同一物理块来共享块。为了确保安全共享，PagedAttention 跟踪物理块的引用计数并实现 Copy-on-Write 机制。

⚠️ 图片缺失（Notion 导出时未包含）：对多个输出进行采样的请求的示例生成过程（notion_cad70eb63b.gif）

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

vLLM inference

离线批量推理

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams

# Initialize the tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-7B-Chat")

# Pass the default decoding hyperparameters of Qwen1.5-7B-Chat
# max_tokens is for the maximum length for generation.
sampling_params = SamplingParams(temperature=0.7, top_p=0.8, repetition_penalty=1.05, max_tokens=512)

# Input the model name or path. Can be GPTQ or AWQ models.
llm = LLM(model="Qwen/Qwen1.5-7B-Chat")

# Prepare your prompts
prompt = "Tell me something about large language models."
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

# generate outputs
outputs = llm.generate([text], sampling_params)

# Print the outputs.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

OpenAI-API兼容API服务

1
2
python -m vllm.entrypoints.openai.api_server \
    --model Qwen/Qwen1.5-7B-Chat

通过输入prompt调用服务

1
2
3
4
5
6
7
curl http://localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d '{
    "model": "Qwen/Qwen1.5-7B-Chat",
    "messages": [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Tell me something about large language models."}
    ]
    }'

通过代码调用

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from openai import OpenAI
# Set OpenAI's API key and API base to use vLLM's API server.
openai_api_key = "EMPTY"
openai_api_base = "http://localhost:8000/v1"

client = OpenAI(
    api_key=openai_api_key,
    base_url=openai_api_base,
)

chat_response = client.chat.completions.create(
    model="Qwen/Qwen1.5-7B-Chat",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Tell me something about large language models."},
    ]
)
print("Chat response:", chat_response)

多GPU分布式服务

1
2
from vllmimport LLM, SamplingParams
llm= LLM(model="Qwen/Qwen1.5-72B-Chat", tensor_parallel_size=4)

或者

1
2
3
python -m vllm.entrypoints.api_server \
    --model Qwen/Qwen1.5-72B-Chat \
    --tensor-parallel-size4

量化模型服务

1
2
3
4
from vllmimport LLM, SamplingParams
llm= LLM(model="Qwen/Qwen1.5-7B-Chat-AWQ", quantization="awq")

# llm = LLM(model="Qwen/Qwen1.5-7B-Chat-GPTQ-Int4", quantization="gptq")

或者

1
2
3
4
5
6
7
python -m vllm.entrypoints.openai.api_server \
    --model Qwen/Qwen1.5-7B-Chat-AWQ \
    --quantization awq
    
python -m vllm.entrypoints.openai.api_server \
    --model Qwen/Qwen1.5-7B-Chat-GPTQ-Int8 \
    --quantization gptq

另外，vLLM还支持AWQ或GPTQ模型与KV缓存量化的组合，即FP8 E5M2 KV Cache。

1
llm = LLM(model="Qwen/Qwen1.5-7B-Chat-GPTQ-Int8", quantization="gptq", kv_cache_dtype="fp8_e5m2")
1
2
3
4
python -m vllm.entrypoints.openai.api_server \
    --model Qwen/Qwen1.5-7B-Chat-GPTQ-Int8 \
    --quantization gptq \
    --kv-cache-dtype fp8_e5m2
本文结束 感谢您的阅读