修炼AI技能
AI 教程 / 大模型应用 / 大模型推理部署：vLLM、KV Cache 与量化选型

大模型推理部署：vLLM、KV Cache 与量化选型

把训练好的大模型跑起来服务线上流量，比训练本身更考验工程。推理服务的核心矛盾是：自回归生成需要把历史 token 的注意力键值（KV）缓存下来，而这部分显存随序列长度与并发数增长，极易成为吞吐瓶颈。本文讲清 KV Cache 的本质、PagedAttention 的突破，以及主流推理引擎与量化方案的选型逻辑。

KV Cache 与显存瓶颈

Transformer 自回归解码时，每个新 token 都要与之前所有 token 计算注意力。为避免重复计算，系统把每一层、每个位置的 Key/Value 缓存在显存中，即 KV Cache。其大小约为：

KV_size ≈ 2 × batch × seq_len × n_layers × n_heads × d_head × bytes_per_elem

可见它随并发（batch）与序列长度（seq_len）线性膨胀。当同时处理大量请求、且各请求长度不一、动态增长时，传统实现会为每个请求预分配最长可能序列的连续显存，造成严重的内部碎片（未用满）与外部碎片（无法拼接），进而被迫减小批大小、拉低吞吐。

vLLM 与 PagedAttention

vLLM（Kwon et al., 2023）提出了 PagedAttention：借鉴操作系统的虚拟内存分页思想，把 KV Cache 切成固定大小的”页（block）“，不要求物理上连续。请求动态增长时按需分配页，不同请求之间可共享相同前缀的页（如相同系统提示、Few-shot 示例），从而把 KV 显存浪费降到接近零，并支持更灵活的批处理（continuous batching）。实测在同等延迟下，vLLM 吞吐较 FasterTransformer、Orca 等系统提升 2–4 倍，序列越长、模型越大收益越明显。

# 用 vLLM 拉起一个 OpenAI 兼容服务
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-3-8B-Instruct \
  --tensor-parallel-size 1 \
  --gpu-memory-utilization 0.9

吞吐与延迟的权衡

部署时要明确两项指标：吞吐（每秒生成 token 数，决定成本）与延迟（首 token 延迟 TTFT、单 token 间隔 ITL，决定体验）。增大批大小能抬高吞吐，却会拉长单请求的延迟；continuous batching 通过随时并入新请求，在二者之间取得更好平衡。此外，投机解码（Speculative Decoding）、前缀缓存（Prefix Caching）也是提升体验与复用计算的关键手段。

量化：用精度换显存与速度

量化通过降低权重/激活的数值精度来压缩模型、加速推理。两类最常用：

  • GPTQ（Frantar et al., 2022）：训练后量化（PTQ）方法，基于近似二阶信息，一次性把权重压到 4-bit/3-bit，几乎不损精度。可在单张消费级显卡上跑起 175B 级模型；
  • AWQ（Lin et al., 2023）：激活感知权重量化，发现仅保护约 1% 的”显著权重通道”即可大幅降低量化误差，通过等效缩放保护显著通道，泛化性好、支持多模态。

二者都属于 weight-only 量化，配合高效 kernel 在 4-bit 下仍能获得明显加速。另有更激进的 2-bit/三元量化，以及激活也量化的 W8A8/W4A16 方案，需在精度与速度间折中。

# 以 AutoGPTQ 加载 4-bit GPTQ 模型
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Llama-3-8B-GPTQ",
    device_map="auto",
    trust_remote_code=False,
)

主流推理引擎选型

  • vLLM：吞吐优先、OpenAI 兼容、PagedAttention 标杆，社区活跃，适合高并发在线服务；
  • TensorRT-LLM：NVIDIA 官方、深度绑定 GPU 与 TensorRT，极致延迟/吞吐优化，但构建与部署较重、硬件绑定强；
  • llama.cpp：纯 C++、支持 GGUF 格式与 CPU/GPU 混合推理，跨平台（含 macOS、树莓派），是本地/边缘部署首选；
  • TGI（Text Generation Inference）：Hugging Face 出品，工程成熟、内置量化与流式输出，适合 HF 生态下的生产部署。

选择时可按场景对号入座：云端高并发选 vLLM/TGI；追求极致性能且是 NVIDIA 栈选 TensorRT-LLM；消费级本地、笔记本、边缘设备选 llama.cpp；需要 GGUF 量化权重时 llama.cpp 生态最完整。

小结

大模型服务化的关键在显存管理：KV Cache 的碎片问题被 vLLM 的 PagedAttention 系统性解决，带来数倍吞吐提升。量化（GPTQ、AWQ）用 4-bit 权重把大模型塞进更小显存并获得加速。vLLM、TensorRT-LLM、llama.cpp、TGI 各有定位，选型应围绕吞吐/延迟目标、硬件与部署形态来决定。

参考与延伸阅读

  1. Kwon, W. et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180 (2023).
  2. Frantar, E. et al. GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers. arXiv:2210.17323 (2022).
  3. Lin, J. et al. AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration. arXiv:2306.00978 (2023).
  4. vLLM 文档: https://docs.vllm.ai/
  5. llama.cpp 仓库: https://github.com/ggerganov/llama.cpp
  6. Hugging Face TGI 仓库: https://github.com/huggingface/text-generation-inference
本文累计阅读 — 次