文章

KV Cache Engineering：LLM 推理真正的内存瓶颈，不只是模型权重

本文系统拆解 KV cache 在 LLM 推理中的作用、显存瓶颈、PagedAttention、Prefix Caching、RadixAttention 与 KV cache quantization，帮助工程团队从内存账本角度优化长上下文、batch size、TTFT 和吞吐量。

为什么 KV cache 值得单独写一篇

很多团队优化 LLM 推理时，第一反应是模型权重量化、tensor parallel、batch size 或更快的 GPU。但在真实服务中，另一个更隐蔽的瓶颈经常先出现：KV cache。

KV cache 保存的是 Transformer 自注意力层中历史 token 的 key/value 张量。没有它，模型每生成一个新 token，都要重新计算前面所有 token 的注意力表示。有了它，decode 阶段只需要读取历史 key/value，再处理当前 token。

问题在于，KV cache 不是免费的。上下文越长、并发越高、batch size 越大，它占用的显存越多。对长上下文 Agent、RAG、多轮对话和代码生成服务来说，KV cache 往往会变成真正的运行时内存账本。

一句话理解 KV cache

标准自回归生成可以简化成：

输入 prompt → 生成 token 1 → 生成 token 2 → 生成 token 3 → ...

每一步都需要关注前面所有 token。KV cache 的作用是把历史 token 的 key/value 存起来，后续 token 直接复用：

  • Prefill：计算 prompt 的 K/V，并写入 cache
  • Decode：只计算新 token 的 K/V，并读取历史 cache

因此，KV cache 同时影响两件事：

  1. 速度：减少重复计算。
  2. 显存：保存每一层、每个 token 的 key/value。

这就是它既是加速器，也是瓶颈的原因。

核心知识结构

可以把 KV cache 工程拆成 6 个模块。

1. Prefill 与 Decode

Prefill 处理输入 prompt，通常计算量大、并行度高。Decode 逐 token 生成，通常更受内存带宽和 cache 读取影响。

KV cache 在 prefill 阶段被写入，在 decode 阶段被持续读取和追加。

2. 显存增长模型

KV cache 大小通常随这些因素增长：

layers × sequence_length × batch_size × kv_heads × head_dim × precision

这意味着同一个模型，在短 prompt、低并发下看似很轻；一旦换成长上下文或高 batch，就可能因为 KV cache 爆显存。

3. PagedAttention

PagedAttention 把 KV cache 拆成固定大小 block，并允许这些 block 在显存中非连续存储。这样可以减少为不同长度请求预留连续大块显存造成的浪费。

4. Prefix Caching

如果多个请求共享相同前缀，比如相同 system prompt、相同 few-shot examples、相同文档开头，就可以复用已计算好的 KV-cache blocks，避免重复 prefill。

5. RadixAttention

RadixAttention 使用 radix tree 管理 token prefix 与 KV cache tensor 的映射，适合复杂 LLM 程序中多次调用、分支推理和共享上下文复用。

6. KV Cache Quantization

KV cache quantization 用更低精度保存 key/value，从而减少显存和带宽压力。它适合长上下文和高 batch，但必须验证准确率与任务稳定性。

KV cache 生命周期：从写入到复用

一个生产 LLM 服务中的 KV cache 生命周期大致如下：

  1. 用户请求进入调度器
  2. Prefill 计算 prompt 的 key/value
  3. KV blocks 写入 paged cache
  4. Decode 阶段逐 token 读取历史 blocks
  5. 如果新请求命中相同 prefix，直接复用已有 blocks
  6. 如果显存不足，按 eviction policy 回收低价值 blocks
  7. 如果启用 quantization，以低精度保存部分或全部 KV cache
  8. 监控 cache hit rate、TTFT、ITL、显存碎片和吞吐量

这里最容易被忽略的是第 5 到第 8 步。很多 Demo 只验证单请求能不能跑，但生产服务真正的收益来自缓存复用、内存布局和调度策略。

为什么只看模型权重大小会误导

假设一个模型权重量化后可以放进显存，服务仍然可能跑不起来。原因是模型权重是相对固定的，而 KV cache 是随请求动态增长的。

典型误判包括：

  • 只用短 prompt 测试，没测长上下文。
  • 只测 batch size = 1，没测并发。
  • 只测 prefill，没测 decode 中后段。
  • 只看 tokens/s，没看 TTFT 和 inter-token latency。
  • 只看显存峰值，没看碎片和 cache eviction。

生产中更合理的测试方式是把 KV cache 单独作为资源统计。

metrics = {
    "kv_cache_used_gb": current_kv_cache_memory(),
    "kv_cache_hit_rate": prefix_cache_hits / prefix_cache_lookups,
    "evicted_blocks": evicted_block_count,
    "fragmentation_ratio": internal_fragmentation / allocated_cache,
    "ttft_ms": time_to_first_token,
    "itl_ms": inter_token_latency,
}

Prefix Caching：最接近免费午餐的优化

很多业务请求天然共享前缀。例如：

  • 相同 system prompt
  • 相同安全规则
  • 相同工具说明
  • 相同 RAG 文档头部
  • 相同 few-shot examples
  • 相同多轮对话历史

如果每个请求都重新 prefill，这些共享部分会被反复计算。Prefix Caching 的价值就是把已经处理过的前缀 KV-cache blocks 保存下来，新请求命中相同 prefix 时直接复用。

vLLM 的做法是对 KV-cache block 做 hash。hash 组成通常包括父 block hash、当前 block token，以及 LoRA、多模态输入、cache salt 等额外因素。这样既能复用相同前缀，又能避免不同租户或不同模型配置误共享。

工程上要注意 3 点：

  1. Prefix caching 不改变模型输出，但会改变系统调度和缓存命中表现。
  2. 命中率取决于 prompt 结构是否稳定。
  3. 多租户场景需要 cache isolation，避免跨用户上下文泄露。

PagedAttention：把显存当分页系统管理

传统 KV cache 管理容易出现两类浪费：

  • 内部碎片：为了最大输出长度预留空间，但实际没用完。
  • 外部碎片：不同请求长度差异太大，连续内存块难以高效分配。

PagedAttention 借鉴操作系统分页思想，把每个请求的 KV cache 切成固定 token block，再用 block table 记录逻辑 token 到物理 block 的映射。

这样做的好处是：

  • 不要求每个请求占用连续显存
  • 按需追加新的 cache blocks
  • 更容易支持不同长度请求混合 batch
  • 减少预分配造成的浪费
  • 提升可承载 batch size

对服务端来说，PagedAttention 的意义不是单纯”更快”，而是让 KV cache 从静态大块分配变成动态资源管理。

RadixAttention：复杂 LLM 程序里的自动复用

普通 prefix caching 多用于共享 system prompt 或固定模板。但 Agent、Tree-of-Thought、自一致性、多轮对话、批量抽取等场景中，复用模式更复杂。

SGLang 的 RadixAttention 将 token sequence 到 KV cache tensor 的映射放入 radix tree 中。它可以在运行时查找共享 prefix、插入新路径，并根据 LRU 策略回收低价值叶子节点。

更关键的是，RadixAttention 不只是缓存结构，还会配合 cache-aware scheduling。也就是说，调度器会尽量把更容易共享缓存的请求组织在一起，从而提高实际命中率。

KV Cache Quantization：用精度换上下文和并发

当 KV cache 成为显存瓶颈时，可以考虑把 key/value 张量用更低精度保存。

它带来的收益包括：

  • 更长上下文。
  • 更大 batch size。
  • 更高 cache hit rate。
  • 更低 memory bandwidth 压力。
  • 更低 TTFT 或 decode latency。

但它不是无风险优化。KV cache 直接参与 attention 计算，如果量化误差过大，可能影响长上下文检索、代码生成、数学推理和多步逻辑任务。

上线前至少应验证：

  • 短文本对话质量
  • 长上下文 needle retrieval
  • 代码生成 benchmark
  • RAG 问答准确率
  • 工具调用 JSON 稳定性
  • 采样输出分布变化

生产落地 Checklist

上线前建议检查这些点：

  • 是否单独统计 KV cache 显存，而不是只看模型权重。
  • 是否区分 prefill latency、TTFT、ITL 和整体吞吐量。
  • 是否测试长上下文和高 batch size 的组合。
  • 是否启用或评估 prefix caching。
  • Prompt 模板是否稳定，能否提升 prefix hit rate。
  • 是否有多租户 cache isolation。
  • 是否使用 paged layout 降低碎片。
  • 是否有明确 eviction policy。
  • 是否记录 cache hit rate 和 evicted blocks。
  • 是否评估 KV cache quantization 对准确率的影响。
  • 是否为 Agent、RAG、代码生成分别建测试集。
  • 是否把 scheduler 与 cache policy 联合优化。

结论

KV cache 是 LLM 推理系统里的运行时内存账本。

模型权重决定”模型能不能装下”，而 KV cache 决定”长上下文、高并发和持续生成能不能稳定跑”。生产级推理优化不能只看权重量化和 GPU 算力，还要系统设计：

  • Prefix Caching 提高复用率
  • PagedAttention 降低碎片
  • RadixAttention 管理复杂共享前缀
  • KV Cache Quantization 扩大上下文和 batch 空间
  • Scheduler 决定缓存命中是否真实发生
  • Metrics 证明优化是否有效

对工程团队来说，最实用的原则是：先把 KV cache 变成可观测资源，再谈优化。否则，很多所谓推理加速只是在短 prompt、低并发、无缓存压力的测试里成立。

常见问题

为什么 KV cache 会成为 LLM 推理瓶颈？
因为每个请求在生成过程中都要保存历史 token 的 key/value 张量，序列越长、batch size 越大，占用显存越多，最终会限制吞吐量、上下文长度和并发能力。
PagedAttention 解决的是什么问题？
PagedAttention 借鉴操作系统分页思想，把 KV cache 拆成固定大小 block，通过 block table 管理非连续内存，减少不同请求长度造成的显存浪费。
KV cache quantization 是否一定安全？
不一定。它能降低显存和带宽压力，但必须在目标模型、任务类型和长上下文 benchmark 上验证准确率、代码生成稳定性和检索任务表现。
特效配置