KV Cache 原理
KV Cache 是自回归解码(decode)阶段最重要的加速手段之一:缓存历史 token 的 Key/Value,避免每一步重复计算与重复读取。
为什么会慢:从“每步重算”说起
在第 (t) 步解码时,模型输入长度是 (t)。如果你每一步都对整个长度重新计算 attention,那么总成本大致是:
[ \sum_{t=1}^{T} O(t^2) = O(T^3) ]
这在长输出时不可接受(实际实现会有细节差异,但“重复计算随步数累积爆炸”的趋势是一样的)。
KV Cache 的做法
对每一层注意力,把历史 token 的 (K,V) 缓存起来:
prefill(处理 prompt)阶段:一次性算出所有位置的 (K,V) 并缓存
decode 阶段:每来一个新 token,只算这个新 token 的 (K_t,V_t),并 append 到缓存
在第 (t) 步:
(Q_t):只对当前 token 计算
(K_{1:t},V_{1:t}):直接从 cache 读取
于是每步 attention 成本从“重算全部”变成:
[ O(t) \quad (\text{与序列长度线性相关}) ]
总成本约为:
[ \sum_{t=1}^T O(t) = O(T^2) ]
这就是 KV Cache 在推理中必不可少的原因。
KV Cache 的代价:显存与带宽
KV cache 的显存大致与以下因素成正比:
batch size
序列长度(prompt + 已生成)
层数 (L)
头数与每头维度(以及是否使用 GQA/MQA)
dtype(fp16/bf16/int8 等)
很多推理系统的瓶颈不是算力而是 KV cache 带宽(读写 cache 的开销)。
常见工程优化
GQA/MQA:减少 K/V 头数,显著降低 cache 体积与带宽压力
Paged KV Cache:把 cache 做成分页/块状管理,减少碎片并支持高并发变长请求(vLLM 等常用)
量化 KV:把 KV cache 低比特存储以换显存(需要权衡精度与吞吐)
FlashAttention:在训练/prefill 阶段减少中间张量与提升带宽利用率
与“并发/吞吐/延迟”的关系
在服务端推理中,KV cache 直接决定:
能同时服务多少请求(并发)
每个 token 的生成速度(吞吐)
是否能进行高效 batching(动态 batch / continuous batching)
因此你会看到:推理系统的很多工程设计(batching、调度、memory manager)都围绕 KV cache 展开。