vLLM 显存优化：如何让推理速度提升 20 倍

在大模型部署中，显存利用率（Memory Utilization）和吞吐量（Throughput）是两个核心指标。vLLM 作为一个开源的高性能推理库，通过 PagedAttention 技术，极大地解决了显存碎片化的问题。

PagedAttention 原理

传统的 KV Cache 显存分配是预先分配连续的内存块，这导致了大量的内部碎片和外部碎片。PagedAttention 借鉴了操作系统中虚拟内存分页（Paging）的思想，将 KV Cache 切分成固定大小的 Block，允许非连续的物理显存分配。

安装与使用

vLLM 的安装非常简单：

pip install vllm

使用 Python API 进行离线推理：

from vllm import LLM, SamplingParams

prompts = [
    "Hello, my name is",
    "The president of the United States is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

llm = LLM(model="facebook/opt-125m")
outputs = llm.generate(prompts, sampling_params)

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

性能对比

在 NVIDIA A100 上，相比于 HuggingFace Transformers，vLLM 的吞吐量通常能提升 2-4 倍。在并发请求较高的情况下，这一优势甚至能扩大到 20 倍以上。这主要得益于它高效的显存管理和 Continuous Batching 机制。

最佳实践

1. 设置适当的 block_size：通常推荐 16 或 32。
2. 启用量化：配合 AWQ 或 GPTQ 量化，可以进一步降低显存占用。
3. 使用 Tensor Parallelism：对于超大模型，利用多卡并行推理。