vLLM 与 KV Cache

vLLM 是由UC Berkeley开发的高性能LLM推理与服务引擎。它的核心创新在于 PagedAttention 算法，通过借鉴操作系统虚拟内存管理的思想来高效管理 KV Cache，从而大幅提升推理吞吐量。

---

KV Cache：为什么需要缓存？

自回归生成的重复计算问题

LLM生成文本时采用自回归（autoregressive）方式，即逐个token生成。每生成一个新token，都需要对之前所有token做 Attention 计算。如果不做任何优化，生成第 \(t\) 个token时需要重新计算前 \(t-1\) 个token的 Key 和 Value，导致大量重复计算。

KV Cache 的核心思想很简单：把已经计算过的 Key 和 Value 缓存起来，每次只计算新token的 K、V，然后拼接到缓存中。

不使用KV Cache：生成第t个token → 重新计算所有t个token的K, V → O(t^2)总计算量
使用KV Cache：  生成第t个token → 只计算第t个token的K, V → O(t)增量计算

KV Cache 的显存问题

KV Cache虽然减少了计算量，但会占用大量显存。对于一个 Transformer 模型，每一层、每一个 Attention Head 都需要存储 K 和 V 张量。

单个请求的 KV Cache 显存估算：

\[ \text{KV Cache Size} = 2 \times L \times H \times d \times s \times \text{sizeof(dtype)} \]

其中 \(L\) = 层数，\(H\) = 注意力头数，\(d\) = 每头维度，\(s\) = 序列长度。

以 LLaMA-13B 为例（40层，40头，128维，fp16），一个2048 token的请求就需要约 800MB 的 KV Cache。当并发请求数增多或序列变长时，KV Cache 成为显存的最大瓶颈。

传统方案的浪费

在 vLLM 出现之前，大多数推理引擎为每个请求预分配一块连续内存来存储 KV Cache。由于无法预知生成长度，通常按最大长度预分配，导致：

• 内部碎片（Internal Fragmentation）：实际生成的序列往往远短于最大长度，预分配的多余空间被浪费
• 外部碎片（External Fragmentation）：不同请求释放内存后，剩余空间不连续，无法被新请求利用
• 预留浪费（Reservation Waste）：为了安全预留，通常过度分配

研究表明，传统方案中 KV Cache 的显存利用率仅约 20-40%。

---

PagedAttention：vLLM 的核心创新

虚拟内存类比

PagedAttention 的灵感直接来自操作系统的虚拟内存与分页机制：

操作系统概念	PagedAttention 对应
虚拟页（Virtual Page）	逻辑上的 KV Block
物理页帧（Physical Frame）	GPU 显存中的固定大小块
页表（Page Table）	Block Table（记录逻辑→物理映射）
按需分页（Demand Paging）	Token生成时才分配新块

工作原理

1. 分块存储：将每个序列的 KV Cache 切分成固定大小的 Block（默认16个token一块），而非预分配一整块连续内存
2. 动态映射：通过 Block Table 记录每个序列的逻辑块到物理块的映射，物理块可以分散在显存的任意位置
3. 按需分配：只在生成新token需要新块时才分配物理块，用完可以立即释放
4. 共享机制：当多个请求共享相同的前缀（如相同的 system prompt）时，可以通过 Copy-on-Write 共享物理块

这种设计将 KV Cache 的显存浪费降到接近零（仅有最后一个块可能存在少量内部碎片），显存利用率接近 100%。

---

Continuous Batching（连续批处理）

传统的 Static Batching 在一个 batch 中等所有请求都完成后才处理下一个 batch。问题是不同请求的生成长度差异很大——有的生成10个token就结束了，有的要生成500个token。短请求完成后只能空等，GPU利用率低。

vLLM 采用 Continuous Batching（也称 Iteration-level Batching）：

• 每次迭代（生成一个token）时动态调度
• 某个请求完成后，立即用新请求填补空位
• 不需要等待整个 batch 全部完成

效果：在高并发场景下，吞吐量相比 Static Batching 可提升 2-4倍。

---

使用方式

Python API（离线推理）

from vllm import LLM, SamplingParams

# 初始化模型
llm = LLM(model="meta-llama/Llama-2-7b-chat-hf",
          tensor_parallel_size=1,   # GPU并行数
          gpu_memory_utilization=0.9)

# 设置采样参数
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=256
)

# 批量推理
prompts = ["What is deep learning?", "Explain attention mechanism."]
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    print(output.outputs[0].text)

Server 模式（在线服务）

vLLM 提供 OpenAI 兼容的 API Server，可以直接替换 OpenAI API：

# 启动服务
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-7b-chat-hf \
    --port 8000 \
    --tensor-parallel-size 2

客户端调用（与 OpenAI API 完全兼容）：

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")

response = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-chat-hf",
    messages=[{"role": "user", "content": "Hello!"}],
    max_tokens=128
)
print(response.choices[0].message.content)

---

与其他推理引擎的对比

特性	vLLM	TGI (HuggingFace)	TensorRT-LLM (NVIDIA)	llama.cpp
核心优化	PagedAttention	Continuous Batching	图优化 + Kernel融合	量化 + CPU优化
KV Cache管理	分页式，接近零浪费	预分配	预分配 + 优化	简单连续分配
吞吐量	很高	高	最高（特定硬件）	较低
易用性	很好（Python原生）	好（Docker部署）	较复杂（需编译）	很好（单文件）
适用场景	云端高并发服务	云端服务	生产环境极致性能	本地/边缘推理
量化支持	GPTQ, AWQ, FP8	GPTQ, BnB	FP8, INT8, INT4	GGUF (2-8bit)
硬件要求	NVIDIA GPU	NVIDIA GPU	NVIDIA GPU	CPU / GPU均可

选型建议：

• 高并发在线服务 → vLLM（开箱即用，吞吐量优秀）
• 极致延迟优化 → TensorRT-LLM（需要投入编译优化的工程成本）
• 快速原型验证 → TGI（HuggingFace生态集成好）
• 本地个人使用 → llama.cpp / Ollama（见 本地推理部署）

---

参考

• Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention", SOSP 2023
• vLLM GitHub
• vLLM 官方文档

---

评论