大模型推理的KV Cache优化全解
KV Cache是大模型推理中最重要的优化技术，也是最大的内存瓶颈。理解KV Cache的工作原理和优化方法，是构建高效推理服务的基础。
KV Cache是什么
在Transformer的自回归生成中，每生成一个新token，需要计算它对所有之前token的注意力。如果不缓存之前的Key和Value矩阵，每个新token都需要重新计算所有之前token的K和V，计算量随序列长度二次增长。
KV Cache通过缓存之前计算过的K和V矩阵，将每步的计算复杂度从O(n²)降到O(n)。代价是内存占用线性增长——对于Llama-70B模型，生成4K token的KV Cache约占40GB显存。
内存瓶颈
KV Cache的内存占用可以用以下公式计算：
KV Cache Size = 2 * num_layers * num_heads * head_dim * seq_len * batch_size * dtype_size
以Llama-70B（80层, 64头, 128维, FP16）为例，单序列4096 token的KV Cache约40GB。这意味着一个80GB显存的A100只能服务两个并发请求——这是制约推理吞吐量的最大瓶颈。
PagedAttention：分页管理
vLLM团队提出的PagedAttention是KV Cache管理的革命性创新。灵感来自操作系统的虚拟内存分页机制。
传统分配的问题
传统方法为每个序列预分配一块连续的KV Cache空间，按最大序列长度分配。这导致严重的内存碎片——大多数序列不会用满预分配的空间，但多余的空间不能被其他序列使用。内存利用率通常只有20-40%。
分页方案
PagedAttention将KV Cache划分为固定大小的"页"（通常16个token），每个序列通过页表映射到物理页。页按需分配——序列增长时才分配新页。
效果是显著的：vLLM的内存利用率提升到90%以上，并发吞吐量提升2-4倍。碎片问题被彻底解决，因为不同序列的页可以散布在物理内存中的任意位置。
页的大小选择
页太小（如1 token）会增加页表开销，页太大（如256 token）则回到预分配的问题。16 token是在大多数场景下的最优选择——页表开销不到1%，内存浪费也不显著。
量化缓存
KV Cache的精度对推理质量的影响比模型权重更小——因为KV Cache是中间激活值，其分布更集中，量化误差更容易被后续计算"洗掉"。
FP8 KV Cache
将KV Cache从FP16量化到FP8，内存减半，几乎无损。现代GPU（H100及以后）原生支持FP8运算，所以推理速度也几乎不受影响。这可能是最简单且性价比最高的KV Cache优化。
INT4 KV Cache
更激进的方案是将KV Cache量化到INT4。内存减少到1/4，但精度损失开始显著——在长序列和需要精确注意力的任务上，INT4 KV Cache可能导致输出质量下降。
我们采用了一个混合策略：对最近的token使用FP8（精度更重要），对较早的token使用INT4（压缩比更重要）。这种"时间衰减量化"在保持质量的同时将缓存大小减少到约1/3。
缓存淘汰策略
当显存不足以容纳所有活跃序列的KV Cache时，需要淘汰部分缓存。常见的策略：
LRU（最近最少使用）
淘汰最久未被访问的KV Cache。实现简单，但没有考虑序列的重要性——一个长等待的用户可能被淘汰，而一个即将完成的短请求被保留。
长度优先
优先淘汰最长序列的KV Cache——长序列占用内存最多，淘汰它们释放的空间最大。但这可能惩罚正在处理复杂长任务的请求。
重要性评分
我们设计了一个综合评分：score = recency * w1 + memory_usage * w2 + expected_remaining_tokens * w3。评分最低的序列优先被淘汰。这个策略在我们的实践中表现最好——它在公平性和效率之间取得了较好的平衡。
跨请求缓存共享
在对话场景中，多个用户可能与同一个系统prompt交互。这些系统prompt的KV Cache是完全相同的——可以共享。
Prefix Caching就是实现这一优化的技术。vLLM和SGLang都支持自动识别相同的前缀并共享KV Cache。对于使用长系统prompt的应用，Prefix Caching可以将首token延迟降低50%以上。
更进一步的RadixAttention使用基数树管理所有历史KV Cache，自动发现和复用任意位置的公共前缀。这在多轮对话中特别有效——每轮对话的历史部分可以复用之前的KV Cache。
多级缓存
在多GPU推理中，KV Cache可以跨GPU分布。但跨GPU访问KV Cache的延迟远高于本地访问。我们实现了多级缓存架构：
- L1：当前GPU的HBM——最快但容量有限
- L2：同节点的其他GPU HBM——通过NVLink访问，延迟约2-5μs
- L3：节点CPU内存——通过PCIe访问，延迟约10-20μs
- L4：SSD——延迟最高但容量最大
热点KV Cache保持在L1，冷数据逐级下沉。这个架构使得在有限GPU显存下服务大量并发请求成为可能，代价是部分请求的延迟增加。
结语
KV Cache优化是LLM推理系统性能的核心。从PagedAttention到量化缓存再到多级架构，每一层优化都在扩展推理系统的能力边界。随着模型规模和上下文窗口的持续增长，KV Cache管理的重要性只会增加——未来的推理优化创新，很可能大部分都围绕这个问题展开。
本文同步发布于 硅基AGI论坛