深度解析：KV Cache Offloading

在大型语言模型（LLM）推理的优化领域，如何高效管理显存资源一直是核心议题。随着上下文窗口（Context Window）的不断扩大，KV Cache（键值缓存） 的管理策略变得尤为关键。

本文将深入探讨 KV Cache Offloading 技术，解析它如何通过将数据从 GPU 迁移至低成本存储层，来解决显存瓶颈，从而显著提升 LLM 的推理效率与扩展性。

什么是 KV Cache Offloading？

KV Cache Offloading 指的是将 Attention 层的 Key/Value 数据从珍贵的 GPU 显存迁移至成本更低的存储介质（如 CPU 内存或磁盘）的过程。

这一技术的核心价值在于：它在释放 GPU 资源的同时，保留了在需要时恢复推理的能力，无需重新计算（Recomputation）。这对于平衡 LLM 工作负载的性能与内存成本至关重要。

为什么 KV Cache很重要？

在大语言模型推理（LLM Inference）过程中，KV Cache 是一项关键的内存优化技术。它通过缓存输入序列中每个 Token 的模型 Attention 的键（Key, K）和值（Value, V）矩阵，避免模型在生成每个新Token时重复计算整个文本序列，从而显著提升推理速度。如何最大化 KV Cache 的命中率，是影响 LLM 推理性能的关键因素之一。

为什么 KV Cache 会成为瓶颈？

然而，随着提示长度和对话复杂度的提升，KV Cache的大小会线性增长。为了确保访问速度，它通常需要保存在GPU内存中——而GPU内存既昂贵又有限。这使得KV Cache很容易成为推理性能的瓶颈。

1. 线性增长的显存消耗：随着上下文窗口的增加，KV Cache 的大小与序列长度呈线性增长。在长文本（Long Context）场景下，这会迅速耗尽 GPU 显存。
2. 闲置资源的浪费：在实际应用中，用户与 LLM 的交互往往是间歇性的（例如用户输入时的停顿、多轮对话的间隔）。在这些“空闲时间”里，KV Cache 依然占据着显存，无法被释放给其他活跃请求使用。

这种低效的内存使用限制了系统的并发用户数（Concurrency）和整体吞吐量（Throughput）。

KV Cache Offloading 的工作机制

为了解决上述问题，KV Cache Offloading 采取了以下策略：

• 卸载（Offload）：将非活跃或访问频率较低的 Cache 数据从 GPU 显存移动到 CPU RAM、本地 SSD 或远程对象存储中。
• 按需加载（Reload）：当用户恢复交互或请求相同的上下文时，系统将数据重新加载回 GPU。

这种机制避免了昂贵的重新计算开销，同时让 GPU 显存专注于处理当前的活跃计算任务。

适用场景：何时使用 Offloading？

KV Cache Offloading 并非银弹，但在以下场景中效果显著：

• 长上下文推理：部署具有超长 Context Window 的模型，且显存容易溢出时。
• 多轮交互与上下文共享：多个用户或 Agent 需要跨会话访问相同的基础内容（例如：开发者在 IDE 中反复查询同一段代码）。
• 资源受限环境：由于预算限制无法扩容 GPU，或需要优化基础设施成本时。
• 分布式推理：在大规模分布式 Worker 节点上，GPU 资源相对稀缺。
• 间歇性工作负载：用户会话存在明显的空闲时间，长期占用显存极度浪费。

核心优势

1. 提升资源利用率：通过腾出显存空间，单张 GPU 可以服务更多的并发用户或处理更长的序列。
2. 降低计算成本：利用廉价的 CPU 内存或磁盘代替昂贵的 HBM（高带宽内存），减少了对高端 GPU 的过度依赖。
3. 降低延迟（Latency）：相比于从头重新计算 KV Cache，从 CPU/磁盘加载数据通常更快。NVIDIA 报告显示，对于长输入序列，Offloading 技术可将首词延迟（TTFT）降低高达 14倍。

性能权衡（Trade-offs）

虽然 Offloading 优势明显，但其性能高度依赖于卸载目标存储层的速度。

• 传输开销 vs. 重算开销：如果存储介质（如机械硬盘）读写过慢，将数据拷回 GPU 的时间可能比直接重新计算还要长。
• 关键指标：确保 数据传输成本 < 重新计算成本。这种情况在长文本、多轮对话中尤为常见，因为此时重算的计算量非常大。

技术深钻：KV Cache 大小计算

KV Cache 显存计算

了解 KV Cache 的显存占用对于规划 Offloading 策略至关重要。下面公式用于估算大语言模型（LLM）推理过程中 KV Cache（键值缓存）所占用的显存大小（以 GB 为单位）。 $$ \text{KV Cache Size (GB)} = \frac{2 \times B \times S \times L \times H \times D \times P_{\text{byte} } }{1024^3} $$ 其中：

• 2: 分别存储 Key 和 Value 向量

  • Transformer 的自注意力机制中，每个 token 都会生成对应的 Key 和 Value 向量，并在后续生成过程中被缓存起来以避免重复计算。
  • 因此，总缓存包含 Key 和 Value 两部分，所以乘以 2。

• B: Batch Size（批大小）

  • 表示一次推理中处理的样本数量（即批大小）。

• S: Sequence Length（序列长度）

  • 表示已生成的 token 序列长度（即上下文长度）。
  • 在自回归生成过程中，每生成一个新 token，序列长度加 1。KV Cache 需要存储从第 1 个到当前第 S 个 token 的所有 K 和 V。

• L: Number of Layers（层数）

  • 表示模型的层数（即 Transformer block 的数量）。
  • 每一层都有独立的自注意力模块，因此每一层都需要自己的 KV Cache。

• H: Attention Heads（注意力头数）

  • 表示多头注意力机制中的头数。
  • 每个注意力头都会生成自己的 K 和 V 向量，因此总维度是 头数 × 每头维度。

• D: Head Dimension（注意力头维度）

  • 表示每个注意力头的隐藏维度（即每个 head 的 key/value 向量长度）, 有时也记作 ( d_k ) 或 ( d_v )。
  • 注意：整个模型的隐藏维度（通常记作 ( d_{model} )）等于 ( H \times D )。

• P_byte: 每个参数的字节数（FP16 为 2，FP32 为 4）

• 1024³

  • 将总字节数转换为 GB（Gigabytes）。

  • ( 1024^3 = 2^{30} ≈ 1,073,741,824 ) 字节 = 1 GiB（二进制 GB）。

注：如果已知模型的 Hidden Size，通常 $H \times D = \text{Hidden Size}$。

补充说明

• 为什么没有 Query（Q）？ Query 是当前 token 临时计算的，不需要缓存；而 Key 和 Value 需要在后续所有步骤中重复使用，所以只缓存 K 和 V。

• KV Cache 是推理阶段显存瓶颈之一，尤其在长上下文（如 32k、128k tokens）场景下，其占用可能远超模型参数本身。

• 优化手段包括：

  • KV Cache 量化（如 INT8、FP8）
  • 分页 KV Cache（如 vLLM 等系统）
  • 稀疏注意力或滑动窗口（如 Mistral 的 sliding window）

示例计算

示例1：

假设 Batch Size=1，序列长度=1024，32层，32个头，头维度128，使用 FP16 (2字节) 精度： $$ \text{KV Cache Size} \approx \frac{2 \times 1 \times 1024 \times 32 \times 32 \times 128 \times 2}{1024^3} = 0.5 \text{ GB} $$ 可以看到，显存消耗随序列长度线性增长。对于 128k 甚至更长的上下文，单请求的 Cache 就可能达到几十 GB。

示例2：

下面以 Qwen25-32B 模型的配置为例，Batch Size 设为1，推理长度设为 128K，来计算推理时 KV Cache 的大小

• ( B = 1 )
• ( S = 131072 )
• ( L = 64 )
• ( H = 40 )
• ( D = 128 )
• 使用 FP16（P_byte=2）

代入公式：

$$ \text{KV Cache Size} = \frac{2 \times 1 \times 131072 \times 64 \times 40 \times 128 \times 2}{1024^3} $$

结果： $$ \text{KV Cache Size} \approx 160 \text{ GB} $$

可以看到，对于 128k 的上下文，单请求的 Cache 就可能达到 160 GB。

实践方案：使用 LMCache

LMCache 是一个专为 LLM 服务设计的扩展引擎，旨在优化长上下文场景下的推理性能。它支持将 KV Cache 分层存储（GPU -> CPU DRAM -> 本地磁盘），并支持跨引擎实例复用 Cache。

LMCache 的集成现状：

• vLLM: 集成 LMCache 以支持 CPU Offloading 和请求间的 Cache 共享，实现存算分离的预填充（Prefill-decode disaggregation）。
• KServe: 利用 LMCache 降低推理成本，确保大规模服务下的 SLO（服务等级目标）。
• llm-d: 通过 LMCache 将数据卸载至网络磁盘等更廉价的存储中。

在 Benchmark 测试中，结合 LMCache 与 vLLM 可在多种用例中实现 3倍至10倍 的延迟降低。

参考链接：

• LMCache Documentation
• NVIDIA Technical Blog on KV Cache Reuse