KV Cache是 Transformer 架构大模型推理加速的核心机制，其内存占用与调度策略直接决定推理服务的吞吐量和延迟表现。在自回归解码阶段，模型每生成一个 token 都需重算历史 token 的 Key 和 Value 矩阵，KV Cache 通过缓存这些矩阵避免重复计算，使推理速度提升数倍，但同时也带来了显著的内存增长与碎片化问题。本文从参数识别、内存估算、优化策略与工程落地四个维度，系统拆解 KV Cache 对推理性能的影响路径。
KV Cache 在推理过程中存储于 GPU 显存，其生命周期从首个 token 进入模型开始，到序列生成结束（遇到结束符或达到最大长度）才被释放。在实际生产中，多路并发请求各自持有独立的 KV Cache 空间，显存的瞬时压力往往由并发度与序列长度的乘积决定，而非简单的模型参数量。
KV Cache 优化的核心价值在于将 Transformer 解码的复杂度从 O(n²) 降为 O(n)，其中 n 为序列长度。没有 KV Cache 时，每生成一个新 token，模型都要重新计算此前所有 token 的注意力权重，计算量随序列长度线性增长，推理时延呈二次方恶化。以 GPT 系列模型为例，当序列长度达到 2048 token 时，无 Cache 情况下的计算量是头 128 token 的 256 倍，实际推理延迟会从毫秒级恶化到秒级——这在任何在线服务场景下都是不可接受的。
引入 KV Cache 后，模型只需要处理当前 token 的 Query，通过缓存的历史 Key/Value 计算注意力分布。这一改动将单步生成的计算复杂度降为 O(1)（仅对当前 token），整体解码复杂度降为 O(n)。据 NVIDIA 《LLM Inference Performance Engineering》白皮书数据显示，KV Cache 能将 GPT-3 175B 模型在 A100 GPU 上的推理吞吐量提升约 3.2 倍，同时将首 token 延迟降低 40% 以上。
KV Cache 的显存占用公式为：每层内存 = 2（Key 与 Value 两个矩阵）× 隐藏层维度 × 序列长度 × 精度字节数（FP16 为 2 字节）。总占用 = 每层内存 × 层数 × 批大小（batch size）。以 LLaMA-7B 模型（32 层、隐藏维度 4096、FP16 精度）为例，单条 2048 token 的序列需要 2×4096×2048×2×32 ≈ 1 GiB 显存。这还只是单请求的占用，当 batch size 达到 32 时，KV Cache 占用将接近 32 GiB，远超模型本身的 14 GiB 参数显存。
KV Cache 在真实推理系统中会引发显存碎片化、无效内存浪费和调度延迟三类问题，它们相互交织并显著拉高整体推理成本。
vLLM 提出的 PagedAttention 借鉴操作系统虚拟内存的分页思想，将 KV Cache 切分为固定大小的块（默认 16 token/块），按需分配而非按最大长度预分配。这种机制将显存利用率从约 40% 提升至 90% 以上，同时允许物理上不连续的显存块通过页表逻辑关联，从根源上消除了碎片化问题。
除 PagedAttention 外，KV Cache 的压缩、量化与淘汰是降低内存压力的三条主流技术路径，各有适用的场景与代价。
| 优化技术 | 核心思路 | 内存降幅 | 精度影响 |
|---|---|---|---|
| GQA（分组查询注意力） | 多个 Query 头共享同一组 Key/Value 头 | 内存降至 MHA 的 1/8（8 个 Query 共享 1 个 KV 头） | 几乎无损 |
| INT8/INT4 量化 | 将 Key/Value 矩阵从 FP16 降为 INT8 或 INT4 | 内存降至 75%（INT8）或 87.5%（INT4） | INT8 近乎无损，INT4 需校准 |
| KV 淘汰策略 | 移除序列中对后续预测贡献极低的 token 的 KV 条目 | 可动态减少 20%-50% | 需根据注意力分数设定阈值 |
LLaMA-2 70B 采用 GQA（8 个 Query 头共享 1 组 KV 头），与标准的 MHA（多头注意力）相比，在推理时 KV Cache 占用减少约 87.5%，而模型质量损失几乎可以忽略（在 MMLU 基准上差距小于 0.3%）。MQA（多查询注意力）更进一步，所有 Query 共享单组 KV 头，内存节省更显著但会对长序列任务（如代码生成）的精度产生可达 2% 的可测量下降——因此目前主流模型普遍选择 GQA 作为折中方案。
对 KV Cache 做量化时，需要区分 Key 和 Value 的不同敏感性。Key 矩阵对量化误差更敏感，直接使用对称量化可能导致注意力分布偏移，建议使用按行 scale 的非对称量化；Value 矩阵相对鲁棒，可使用逐张量（per-tensor）量化。以下是一个简化版的 PyTorch 实现示例，展示对 Key/Value Cache 进行 INT8 量化的基本流程：
import torch
def quantize_fp16_to_int8(tensor):
# 按行计算 scale 和 zero_point，保留更多数值分布信息
scale = tensor.abs().max(dim=-1, keepdim=True).values / 127.0
scale = scale.clamp(min=1e-8)
quantized = torch.round(tensor / scale).clamp(-128, 127).to(torch.int8)
return quantized, scale
def inference_with_quantized_kv_cache(key_cache, value_cache, query):
# 反量化 Key 缓存用于注意力计算
key_fp16 = key_cache[0] * key_cache[1] # 反量化
attn_scores = query @ key_fp16.transpose(-2, -1) / (query.size(-1) ** 0.5)
attn_probs = torch.softmax(attn_scores, dim=-1)
return attn_probs @ value_cache在实际部署中，建议先用小规模评估集验证量化后模型的困惑度变化。当 INT8 量化导致困惑度上升超过 0.5 时，应退回 FP16 或改用混合精度策略（Key 保持 FP16，仅量化 Value）。
不同层的 KV Cache 对模型输出的重要性并不均等。研究表明，浅层（靠近输入的层）的 KV Cache 承载更多全局语义信息，深层（靠近输出的层）则更多编码局部语法模式。因此，可对深层 KV Cache 使用更大压缩比（如 INT4），浅层保持 INT8 或 FP16。这种 Layer-wise 混合精度策略在保持端到端精度损失小于 1% 的前提下，可将 KV Cache 总内存再降低约 40%。
以字节跳动豆包大模型团队发布的工程实践为参考，在其推理服务中应用 PagedAttention 与 INT8 量化组合后，单机并发能力从 16 路提升至 64 路，KV Cache 显存占用从 40 GiB 降至 12 GiB，整体吞吐量提升 4 倍，同时保持平均单 token 延迟低于 30ms。这一案例验证了组合策略在真实生产环境中的有效性。
当前主流推理框架（vLLM、TensorRT-LLM、SGLang）均内置了 PagedAttention 或类似的分页管理机制，但 API 风格与硬件利用效率存在差异。vLLM 适合快速部署，社区生态成熟；TensorRT-LLM 在 NVIDIA GPU 上的算子融合更极致，适合追求极限性能的线上服务；SGLang 则在多模态模型的支持上更完善。选型时应以实际负载的序列长度分布、并发度要求和硬件型号为依据，建议统一做一次线性扫描压测对比。
KV Cache 的下一步演进将聚焦于更精准的语义剪枝与跨模型共享。具体而言，基于注意力分数分布的动态淘汰算法（如 H2O，即 Heavy Hitter Oracle）通过识别序列中对后续生成贡献最大的关键 token（Heavy Hitters），仅保留这些 token 的 KV 条目，从而在 3000 token 长度下将内存降低约 50% 且保持困惑度基本不变。此外，为某些 KV 块引入 CPU 显存分层（冷热数据迁移）也是平衡成本与性能的实用方向——例如将历史较远的 KV 块迁移至 CPU 内存，在进行长对话时按需回传。
优化 KV Cache 本质上是让有限的显存承载更长的上下文和更高的并发度。从工程角度出发，先通过 GQA 控制模型结构带来的固有消耗，再以 INT8 量化压缩数值精度冗余，最后用 PagedAttention 解决动态分配的碎片问题——三条路径组合使用通常能获得 3-5 倍的整体显存收益，且对生成质量的损耗保持在可接受范围内。
评论
加载中...