在以Transformer架构的模型推理阶段，KV Cache是一种为了大幅提升大语言模型(LLM)生成速度而采用的关键技术，核心目标是：避免在生成一个新的Token时，重复计算之前已经生成过的Token的注意力信息。

1 KV Cache是什么？

在Transformer模型中，注意力机制的核心公式是：

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

其中：

  • Q (Query)：当前正在处理的向量
  • K (Key)：用来匹配相关性的键向量
  • V (Value)：用来提取信息的价值向量

当你输入一段提示词（Prompt），让模型生成后续文本时，这是一个自回归（Autoregressive）过程：当生成N+1个词时，需要依赖前N个词的信息。
而KV Cache的本质就是将之前所有Token计算出来的K(key)和V(value)矩阵缓存下来，这样在生成下一个词时，模型不需要从头计算从1到N个词的K和V，直接读取缓存中的数据即可。

2 为什么需要KV Cache?

如果Transformer在推理阶段没有KV Cache：

  • 每生成一个词，都要重新计算整个序列的Attention
  • 当生成第N个Token时，模型需要把之前的N-1个Token全部塞进模型，重新进行一遍矩阵乘法，这会导致计算量随序列长度的增加呈平方级增长(O(N^2))，推理效率会随着序列长度增加而急剧下降

推理阶段有了KV Cache：

  • 只需要计算最后一个Token的 Q、K、V
  • 将新产生的K和V追加到已有的Cache列表末尾
  • 利用缓存中之前的K和V与当前的Q进行运算
  • 计算效率大幅提升，虽然显存占用增加（用于存储K和V的缓存），但计算复杂度从 O(N^2) 降低到了O(N)，对于长文本生成，这相当于把原本几秒的推理过程缩短到了毫秒级。

3 KV Cache的缺点

虽然KV Cache极大地提升了模型推理速度，但是也存在缺点：

  • 显存占用大：如果上下文窗口很大，KV Cache会占用大量的显存
  • Batch处理限制：在处理多个用户并发请求时，不同用户序列长度不一致，导致Cache的内存管理变得复杂

4 KV Cache的后续优化

为了解决KV Cache的上述缺点，目前主要有一下的优化技术：

  • PagedAttention(vLLM)：像操作系统管理内存一样管理KV Cache，避免显存碎片化
  • Grouped Query Attention：通过多个Query共享一组Key和Value，减少缓存大小
  • FlashAttention：从底层算子层面优化计算逻辑，降低缓存读写的开销