深入解析 vLLM: PagedAttention 与高性能推理

摘要: 在 LLM 推理场景中，显存带宽（Memory Bandwidth）往往是性能的决定性瓶颈。vLLM 通过引入操作系统的虚拟内存分页思想（PagedAttention），解决了 KV Cache 的显存碎片化问题，将推理吞吐量提升了 2-4 倍。

1. 推理的痛点：显存碎片与 KV Cache

在大模型推理的自回归（Auto-regressive）过程中，模型会逐个生成 token。为了避免重复计算，我们会将之前的计算结果（Key 和 Value 矩阵）缓存下来，这被称为 KV Cache。

显存浪费的根源

传统的推理系统（如 HuggingFace Transformers）通常要求 KV Cache 在显存中是连续的。这导致了严重的内存碎片：

1. 预分配浪费: 为了防止溢出，必须按最大序列长度（如 2048）预分配显存，但实际可能只用了 100。
2. 碎片化: 即使显存有空闲，但如果不是连续的大块，也无法被新的请求利用。

据统计，在传统系统中，60% - 80% 的显存实际上是被“浪费”的（Internal/External Fragmentation）。

---

2. 核心解法：PagedAttention

vLLM 的核心创新在于引入了操作系统中 虚拟内存（Virtual Memory） 和 分页（Paging） 的概念。

2.1 算法原理

PagedAttention 允许 KV Cache 存储在非连续的物理显存空间中。

• 逻辑块 (Logical Block): 类似于 OS 的虚拟页（Page）。
• 物理块 (Physical Block): 类似于 OS 的物理页帧（Page Frame）。
• 块表 (Block Table): 记录逻辑块到物理块的映射关系。

2.2 优势

1. 零浪费: 只有当确实需要存储数据时，才会分配物理块。
2. 灵活共享: 类似于 OS 的 Copy-on-Write，不同的序列（如 Parallel Sampling 或 Beam Search）可以共享相同的物理块（Prompt 部分），极大节省显存。

---

3. 另一大杀器：Continuous Batching

除了显存管理，vLLM 还实现了 Continuous Batching（也称 Iteration-level Scheduling）。

• 传统 Static Batching: 一个 Batch 中必须等待所有请求都生成完毕，才能处理下一批。短请求被迫等待长请求（Head-of-line blocking）。
• Continuous Batching: 一旦某个请求生成结束（遇到了 EOS），立即将该槽位释放给新的请求，无需等待其他请求结束。

这使得 GPU 的利用率始终保持在高位。

---

4. 实战：vLLM 安装与使用

4.1 环境准备

• OS: Linux
• Python: 3.8+
• CUDA: 11.8 或 12.1 (推荐)

# 推荐使用 pip 安装预编译版本
pip install vllm

# 验证安装
python -c "import vllm; print(vllm.__version__)"

4.2 离线推理 (Offline Inference)

适用于批处理任务，如对数据集进行打分或生成。

from vllm import LLM, SamplingParams

# 1. 初始化引擎 (自动下载模型)
#由于 vLLM 对显存管理极其激进，建议手动指定 gpu_memory_utilization
llm = LLM(model="facebook/opt-125m", gpu_memory_utilization=0.9)

# 2. 定义采样参数
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=100)

# 3. 准备输入
prompts = [
    "Hello, my name is",
    "The capital of France is",
]

# 4. 执行推理
outputs = llm.generate(prompts, sampling_params)

# 5. 打印结果
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

4.3 构建 OpenAI 兼容服务 (Online Serving)

vLLM 自带了一个高性能的 API Server，完全兼容 OpenAI 的接口协议。

启动服务：

python -m vllm.entrypoints.openai.api_server \
    --model Qwen/Qwen2.5-7B-Instruct \
    --trust-remote-code \
    --port 8000

客户端调用 (使用 openai 库)：

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="EMPTY", # vLLM 默认不需要 Key
)

completion = client.chat.completions.create(
    model="Qwen/Qwen2.5-7B-Instruct",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "介绍一下 PagedAttention 的原理"}
    ]
)

print(completion.choices[0].message.content)

---

5. 性能基准测试

在 NVIDIA A100 (80GB) 上，与 HuggingFace Transformers 相比，vLLM 的吞吐量（Throughput）通常能提升 2x - 4x。

Framework	Throughput (req/s)	Latency (ms)
HF Transformers	2.5	450
HF Text Gen Interface (TGI)	5.8	120
vLLM	12.4	110

(数据仅供参考，实际性能取决于模型大小与 Batch Size)

6. 总结

vLLM 不仅仅是一个推理库，它展示了如何将 HPC 系统设计（虚拟内存、调度算法）应用到 AI 算法 中解决实际瓶颈。对于构建大规模 AI 生产系统的工程师来说，理解 vLLM 的原理是必修课。