Transformer KV Cache：推理加速的收益和显存代价

Transformer KV Cache：推理加速的收益和显存代价

自回归 Transformer 推理时，KV Cache 是核心优化。没有缓存，每生成一个 token 都要重新计算前面所有 token 的 key 和 value；有了缓存，模型只处理新增 token，大幅减少重复计算。但 KV Cache 不是免费午餐，它会占用显存，并且随 batch size、层数、头数、上下文长度增长。

理解 KV Cache，有助于解释为什么长上下文推理显存压力很大，也能帮助评估服务端并发。

一、KV Cache 缓存了什么

flowchart TD A[Prompt Tokens] --> B[Key Value Projection] B --> C[KV Cache] D[Next Token] --> E[Attention With Cache] C --> E E --> F[Generate Token] F --> G[Append New KV] G --> C

每一层注意力都会保存历史 token 的 key 和 value。后续生成时，只需要为新 token 计算新的 query、key、value，并与缓存交互。

二、显存估算要写清维度

KV Cache 大小通常与 batch、层数、序列长度、hidden size 和 dtype 相关。

kv_cache_bytes ≈ batch_size × seq_len × num_layers × 2 × hidden_size × bytes_per_element

这里的2表示 key 和 value。实际实现还会受 attention head、分组查询注意力、内存布局和框架优化影响，但这个估算足够帮助建立直觉。

三、长上下文会压缩并发

同一张 GPU 上，短请求可以同时服务很多个；长上下文请求会占用大量 KV Cache，导致可并发数量下降。

serving_tradeoff: short_chat: context: 1024 concurrency: high long_document: context: 32768 concurrency: low

因此服务端不能只按请求数限流，还要按 token 预算限流。长 prompt 和长输出都应计入资源估算。

容量规划时可以把请求换算为 token budget。两个请求数量相同的服务，如果一个平均上下文为 1k，另一个平均上下文为 16k，对 GPU 显存和调度的压力完全不同。

四、优化方向要看瓶颈

如果瓶颈是计算，KV Cache 帮助很大；如果瓶颈是显存，可能需要分页缓存、量化 KV、限制上下文或做请求调度。

optimization_options ├── paged KV cache ├── grouped query attention ├── kv cache quantization ├── max context policy └── request batching and preemption

不同优化会影响延迟、吞吐和质量。比如压缩或量化 KV Cache，需要评估对输出质量的影响。

五、总结

KV Cache 通过缓存历史 token 的 key 和 value，减少自回归推理中的重复计算，是大模型推理加速的重要机制。但它会显著消耗显存，并限制长上下文场景下的并发。

评估推理服务时，要把 KV Cache 显存纳入容量规划。推理性能不是只看模型参数量，还要看上下文长度和并发形态。

在长上下文业务里，限制最大输入长度通常不是产品保守，而是服务稳定性的必要条件。