生产级 LLM 推理的 KV 缓存优化

如果你正在生产环境中部署 Large Language Models（LLM）或多租户 AI Agent，很可能遇到过意料之外的 Out-Of-Memory（OOM）错误。

大多数工程团队会根据模型权重来规划 GPU 容量（例如，一个采用 FP16 的 70B 参数模型约需 140 GB），并将剩余 VRAM 视为额外的安全余量。但到了生产环境，这种假设很快就会失效。

真正悄无声息地拖垮 LLM 部署的，并不是模型权重，而是 KV（Key-Value）Cache。

1. 什么是 KV Cache？它为什么会耗尽 VRAM？

在 LLM 生成过程中，模型会将此前已处理 token 的 Key 和 Value 张量保存在 VRAM 中，从而避免每生成一个新 token，都要重新计算整个序列的 attention。

虽然这能显著加速推理，但 Cache 占用的空间会随着 batch size 和 context length 线性增长。在长期运行的 Agent 工作负载中（16K～32K+ token），KV Cache 消耗的 VRAM 往往远超模型权重本身。

问题在于：一个拥有 32K context window 的 Agent 会话，甚至还没生成第一个输出 token，就可能吞掉数十 GB 的 VRAM。再乘以并发会话数量，OOM 错误几乎不可避免。

2. 数学计算：估算 KV Cache 的内存占用

KV Cache 的内存占用严格按照以下公式增长：

Memory = 2 × batch_size × seq_len × num_layers × num_kv_heads × head_dim × precision_bytes

（最前面的 2，是因为需要分别存储 Key 和 Value 两组张量。）

真实案例：32K Context 下的 Llama 2 70B

我们代入 Llama 2 70B 的参数：

KV Heads：8（使用 Grouped-Query Attention）

精度：FP16（每个值占 2 bytes）

对于一条 context 长度为 32,768 token 的序列：

Memory = 2 × 1 × 32,768 × 80 × 8 × 128 × 2 bytes = 10,737,418,240 bytes ≈ 每条序列 10.74 GB

也就是说，单个用户就会占用 10.74 GB 的 VRAM。如果 Llama 2 70B 使用的是标准 Multi-Head Attention（64 heads），而不是 GQA（8 heads），那么每条序列的占用将飙升至 85 GB 以上。

3. KV Cache 的核心优化技术

为了避免内存枯竭，现代推理服务框架通常会结合四种主要的架构级解决方案。

4. 实战教程：使用 vLLM 实现 PagedAttention

vLLM 默认使用 PagedAttention 作为内存管理器。下面介绍如何配置并部署经过优化的推理服务器。

pip install vllm

第 2 步：使用 PagedAttention 和优化参数启动服务器

启动服务器，启用较高的内存分配比例、prefix caching 和 FP8 KV Cache 量化：

vllm serve meta-llama/Llama-2-70b-chat-hf \
  --gpu-memory-utilization 0.90 \
  --max-model-len 32768 \
  --dtype float16 \
  --enable-prefix-caching \
  --kv-cache-dtype fp8

--gpu-memory-utilization 0.90：将 GPU VRAM 的 90% 分配给模型权重和分页式 KV Cache 内存池。

--enable-prefix-caching：避免为不同请求重复计算相同的 prompt 前缀。

--kv-cache-dtype fp8：将每个 token 对应的 KV Cache 内存占用减半。

第 3 步：运行并发负载基准测试

vLLM 内置了基准测试脚本，可用于测量高负载下的吞吐量：

python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-2-70b-chat-hf &

python benchmarks/benchmark_serving.py \
  --backend vllm \
  --model meta-llama/Llama-2-70b-chat-hf \
  --num-prompts 100 \
  --request-rate 10

执行期间，通过 nvidia-smi 监控内存使用情况，确认内存能够稳定增长至设定上限，而不会导致服务崩溃。

5. 研究基准与性能影响

根据最初的 PagedAttention 研究论文（Kwon 等人，SOSP 2023）：

内存浪费：从 60%～80%（朴素的连续内存分配）降至 4% 以下。

吞吐量：在延迟相同的情况下，相比此前最先进的推理服务系统（如 Orca 和 FasterTransformer），吞吐量提高了 2～4 倍。

生产规模：真实生产部署（如 LMSYS Chatbot Arena）在承载 2～3 倍流量的同时，将所需 GPU 总量减少了 50%。

💡 总结与后续步骤

对于长 context LLM 应用来说，优化 KV Cache 并不是可选项，而是控制基础设施成本、避免因 OOM 导致服务停机的必要措施。

👉 本文最初发布于 GPUYard。你可以阅读原文，进一步深入了解 LLM 基础设施和 MLOps 优化！

如需采取进一步措施，你可以考虑屏蔽此人和／或举报滥用行为。