发布于 2025 年 10 月 7 日
背景与合作概述
在大语言模型推理（LLM Inference）过程中，KV Cache 是一项关键的内存优化技术。它通过缓存模型的键（Key, K）和值（Value, V）矩阵，避免模型在生成每个新令牌（Token）时重复计算整个文本序列，从而显著提升推理速度。如何最大化 KV Cache 的命中率，是影响 LLM 推理性能的关键因素之一。然而，KV Cache 会占用大量 GPU 高速显存（HBM），而且这种占用会随着系统提示词（System Prompt）长度的增加而迅速增长。这一内存瓶颈直接限制了模型的上下文长度、并发能力以及整体吞吐量。我们的目标是通过一种分层的节点本地存储解决方案，有效扩展 GPU HBM，从而最大化 KV Cache 的命中率。为此，我们与 LMCache 团队（Kuntai Du, Jiayi Yao, Yihua Cheng，他们在合作开发 LMCache 后成立了名为 Tensormesh 的公司）合作，利用他们的智能软件在 GKE 上创建了这一解决方案。
分层存储方法
LMCache 允许将 KV Cache 从 GPU 高速显存（HBM，Tier 1）扩展到更大、更具成本效益的存储层，如 CPU RAM 和本地 SSD。这种分层存储架构显著提升了总缓存容量，使更多数据能够被保存在加速器节点本地，从而提升缓存命中率和整体推理性能。对于 GKE 用户而言，这意味着即使在支持超大上下文窗口的模型场景下，也能维持高效的性能表现。
性能基准测试与结果
我们设计了多组实验，用于评估不同分层 KV Cache 配置下的性能表现。通过调整总缓存大小，使工作负载填满每一层存储（HBM、CPU RAM、本地 SSD），并使用多种上下文长度（1k、5k、10k、50k、100k 令牌）进行基准测试。
主要性能指标包括：
- TTFT（首令牌生成延迟）
- 令牌输入吞吐量
- 端到端延迟
不同系统提示词长度的用例示例如下：
| 提示词长度 | 典型应用场景 |
|---|---|
| 1k–5k | 高保真角色设定、复杂指令 |
| 10k | 普通用户查询、小型 RAG 或网页文章内容 |
| 50k | 提示词填充（Prompt Stuffing） |
| 100k | 长篇文档或整本书的内容 |
接下来的结果展示了每种 KV Cache 大小配置下性能最佳的存储设置及其带来的改进。
实验设置
我们在一台 A3 mega 机器上部署了 vLLM 服务器，并使用本地 SSD 作为临时存储（通过 emptyDir 挂载）。测试环境配置如下：
- 硬件: 8 × nvidia-h100-mega-80gb GPUs
- 模型: Llama-3.3-70B-Instruct
- LMCache 版本: v0.3.3
- 缓存配置
- 仅 HBM
- HBM + CPU RAM
- HBM + CPU RAM + 本地 SSD
- 存储资源: HBM: 640 G, CPU RAM: 1T, 本地 SSD 5T
- 基准测试工具: SGLang
bench_serving - 请求: 测试使用了多种系统提示词长度（1k、5k、10k、50k 和 100k 令牌）。每个系统提示词（System Prompt）为一个批次的 20 个推理请求提供共享上下文。虽然系统提示词是共享的，但每个请求都包含不同的 256 令牌输入，并生成 512 令牌的输出。
示例命令
python3 sglang/bench_serving.py \
--host=${IP} \
--port=${PORT} \
--dataset-name='generated-shared-prefix' \
--model=$MODEL \
--tokenizer=$MODEL \
--backend=vllm \
--gsp-num-groups=80 \
--gsp-prompts-per-group=20 \
--gsp-system-prompt-len=1000 \
--gsp-question-len=256 \
--gsp-output-len=512 \
--request-rate=800 \
--max-concurrency=200
基准测试结果
我们针对不同的总 KV Cache 大小创建了测试。以下结果展示了每种 KV Cache 大小下性能最佳的存储设置及其带来的改进：
测试 1: 总缓存 1.1M–1.3M 令牌（完全存储于 HBM 内）
结果: 在此缓存规模下，HBM 足以容纳全部 KV Cache，增加更慢的存储层无性能收益。因此，仅使用 HBM 的配置为最优方案。
测试 2: 总缓存 4.0M–4.3M 令牌（超出 HBM 容量，但小于 HBM + CPU RAM 总容量）
结果:
| 系统提示词长度 | 性能最佳的存储设置 | 平均 TTFT (ms) 变化 (%) (与仅 HBM 设置相比) | 输入吞吐量变化 (%) (与仅 HBM 设置相比) | 平均端到端延迟变化 (%) (与仅 HBM 设置相比) |
|---|---|---|---|---|
| 1000 | HBM | 0% | 0% | 0% |
| 5000 | HBM + CPU RAM | -18% | +16% | -14% |
| 10000 | HBM + CPU RAM | -44% | +50% | -33% |
| 50000 | HBM + CPU RAM + 本地 SSD | -68% | +179% | -64% |
| 100000 | HBM + CPU RAM + 本地 SSD | -79% | +264% | -73% |
测试 3: 总缓存 12.6M–13.7M 令牌（HBM 与 CPU RAM 饱和，部分数据存入 SSD，但是总缓存大小小于 HBM + CPU RAM + 本地 SSD 的总容量）
结果:
| 系统提示词长度 | 性能最佳的存储设置 | 平均 TTFT (ms) 变化 (%) (与仅 HBM 设置相比) | 输入吞吐量变化 (%) (与仅 HBM 设置相比) | 平均端到端延迟变化 (%) (与仅 HBM 设置相比) |
|---|---|---|---|---|
| 1000 | HBM + CPU RAM | +5% | +1% | -1% |
| 5000 | HBM + CPU RAM | -6% | +27% | -21% |
| 10000 | HBM + CPU RAM + 本地 SSD | +121% | +23% | -19% |
| 50000 | HBM + CPU RAM + 本地 SSD | +48% | +69% | -41% |
| 100000 | HBM + CPU RAM + 本地 SSD | -3% | +130% | -57% |
总结与后续步骤
实验结果表明，LMCache 的分层存储方案能够充分利用节点本地资源，有效提升 LLM 推理性能，尤其在长提示词导致 KV Cache 规模较大的场景中表现突出。
优化 LLM 推理性能是一项复杂的系统工程，需要存储、计算和网络等多个基础设施组件协同配合。我们的工作是优化整个端到端推理栈的一部分，包括从推理网关的智能负载均衡，到模型服务器内部的缓存策略优化。
接下来，我们将继续探索将更多远程存储解决方案与 LMCache 集成，以进一步提升系统性能与可扩展性。
进一步了解：
- 您可以在 GKE 教程 上尝试相同的部署设置。
- 欢迎关注 LLM-D Inference Stack 的最新进展。
发表评论