vLLM 高性能 LLM 推理
bash pip install vllm
如果想从源码编译以获得最新功能或定制,请参照官方仓库指引。
对于初学者,直接使用 `pip install vllm` 即可得到一个功能完整的推理引擎。
> **注意**:vLLM 依赖 FlashAttention 等底层库以获得最佳性能,pip 包中已内置兼容版本。如果在某些环境中遇到问题,可以尝试安装 `flash-attn` 库并确保 CUDA 路径正确。
## 快速开始:离线推理
只需几行代码,就能加载模型并生成文本。以下示例加载 `facebook/opt-125m` 小型模型进行演示,实际应用中请替换为你想使用的模型路径或 HuggingFace ID。
```python
from vllm import LLM, SamplingParams
# 初始化引擎,指定模型名称和使用的 GPU 数量
llm = LLM(model="facebook/opt-125m", trust_remote_code=True)
# 定义生成参数
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=100)
# 准备提示词列表
prompts = [
"The future of AI is",
"Once upon a time",
]
# 执行批量推理
outputs = llm.generate(prompts, sampling_params)
# 打印结果
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}")
print(f"Generated: {generated_text!r}\n")
SamplingParams 支持大部分 HuggingFace generate 方法的参数,如 temperature、top_k、top_p、max_tokens 等,同时还支持 stop 字符串列表和 frequency_penalty 等。
启动与 OpenAI 兼容的 API 服务
vLLM 可以快速部署为一个提供 HTTP API 的服务器,接口与 OpenAI Completions API 一致,方便现有应用无缝切换。
启动命令
python -m vllm.entrypoints.openai.api_server \
--model meta-llama/Llama-2-7b-chat-hf \
--tensor-parallel-size 1 \
--dtype auto
--model:模型名称或本地路径。--tensor-parallel-size:张量并行所用的 GPU 数量(例如 2 表示在两块 GPU 上切分模型)。--dtype auto:自动选择合适的数据类型(FP16 或 BF16)。- 还可以添加
--quantization awq或--quantization gptq来使用量化模型。
服务启动后,默认监听 http://localhost:8000。可以通过以下方式测试:
curl http://localhost:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{
"model": "meta-llama/Llama-2-7b-chat-hf",
"prompt": "San Francisco is a",
"max_tokens": 50,
"temperature": 0.7
}'
使用 OpenAI Python 库连接
将 openai 库的 base_url 指向 vLLM 服务,即可像调用标准 OpenAI 接口一样使用:
from openai import OpenAI
client = OpenAI(
base_url="http://localhost:8000/v1",
api_key="not-needed" # vLLM 默认不验证 api_key
)
response = client.completions.create(
model="meta-llama/Llama-2-7b-chat-hf",
prompt="The capital of France is",
max_tokens=20,
temperature=0.5,
)
print(response.choices[0].text)
PagedAttention:高吞吐的秘密
传统 KV 缓存内存管理的问题
Transformer 在解码时每生成一个 token,都需要对历史 token 的键和值进行注意力计算。这些键和值会被缓存,称为 KV 缓存。传统的实现会为每个请求预先分配一块足够容纳最大长度序列的连续内存:
- 内部碎片:如果请求生成长度比预分配短,内存被浪费。
- 外部碎片:不同请求间内存释放后无法被合并利用,导致内存利用率低,batch size 受限。
- 无法共享:即使采用 Beam Search 等算法,多个候选序列也必须重复存储相同的 KV 缓存部分。
PagedAttention 如何解决
vLLM 将 KV 缓存划分为固定大小的块(页面),每个块可存储固定数量 token 的键和值。注意力计算时,通过页面表将逻辑位置映射到物理块,缓存不再需要连续存储:
- 请求只需要为当前生成的 token 分配需要的页面,内存浪费极小。
- 页面可以在不同请求之间共享(例如 Prefix Caching),进一步减少冗余。
- 这种设计使得 GPU 内存利用率超过 95%,从而支持更大的 batch size,吞吐量大幅提升。
高级配置与性能调优
张量并行与多卡推理
对于大型模型,单卡可能容不下全部参数。借助张量并行,可以将模型的每层参数切分到多张 GPU 上并行计算。
在 LLM 初始化时设置 tensor_parallel_size:
llm = LLM(model="meta-llama/Llama-2-70b-hf", tensor_parallel_size=4)
API 服务器中通过 --tensor-parallel-size 4 指定。vLLM 会自动将模型权重分割并分配到不同的设备,同时插入必要的通信操作。
量化模型支持
vLLM 支持 AWQ、GPTQ、SqueezeLLM 等量化方法,能以更小的内存占用和更低的显存带宽需求运行模型,同时保持精度损失极小。使用量化模型只需在初始化或启动时指定 --quantization 参数。
例如,加载 TheBloke 的 AWQ 量化 LLaMA-2 模型:
python -m vllm.entrypoints.openai.api_server \
--model TheBloke/Llama-2-7B-Chat-AWQ \
--quantization awq
或者在代码中:
llm = LLM(model="TheBloke/Llama-2-7B-Chat-AWQ", quantization="awq")