在大型语言模型（LLMs）中，Key-Value (KV) Cache 是一种至关重要的优化技术，主要用于加速模型的推理过程，特别是在自回归生成（token by token generation）任务中。它通过存储和重用注意力机制中的中间计算结果来显著提高效率。
要理解 KV Cache，首先需要回顾 Transformer 架构中的自注意力（Self-Attention）机制。在 Transformer 模型中，每个输入的 token 都会被转换成三个向量：
- Query (Q)：表示当前 token 的“查询”内容，用于寻找相关信息。
- Key (K)：表示序列中所有 token 的“键”，用于与查询进行匹配。
- Value (V)：表示序列中所有 token 的“值”，包含实际的信息内容。
自注意力机制的计算过程大致如下：
- 计算注意力分数：Query 向量与所有 Key 向量进行点积（或更复杂的兼容性函数），得到注意力分数。这些分数表示当前 Query 对每个 Key 的关注程度。
- 归一化注意力分数：将注意力分数通过 Softmax 函数进行归一化，得到注意力权重。
- 加权求和 Value：将注意力权重与对应的 Value 向量进行加权求和，得到当前 Query 的输出向量。这个输出向量包含了序列中所有相关信息，并根据注意力权重进行了聚合。
KV Cache 的作用
在自回归生成任务中（例如，生成下一个单词），模型会一个 token 一个 token 地生成输出。对于每个新生成的 token，模型都需要重新计算整个序列的注意力。例如，当模型生成第二个 token 时，它需要考虑第一个 token 的信息；当生成第三个 token 时，它需要考虑第一个和第二个 token 的信息，以此类推。
如果没有 KV Cache：
每次生成新 token 时，模型都需要重新计算整个历史序列（包括当前正在生成的 token 之前的序列）的 Q、K、V 向量，并执行完整的注意力计算。这会导致大量的重复计算，尤其是在序列长度较长时，计算成本会呈平方级增长（因为注意力计算的复杂度与序列长度的平方成正比）。
有了 KV Cache：
KV Cache 的核心思想是存储已经计算过的 Key 和 Value 向量，并在生成后续 token 时复用它们，避免重复计算。具体来说：
- 首次生成：当模型处理第一个输入（例如，提示词）时，它会计算所有 token 对应的 Key 和 Value 向量，并将它们存储在 KV Cache 中。
- 后续生成：当模型生成下一个 token 时，它只需要计算新生成 token 的 Query 向量。然后，它会将这个新的 Query 向量与 KV Cache 中存储的所有历史 Key 向量以及新生成 token 的 Key 向量进行注意力计算。同样地，Value 向量也只考虑新的 Value 向量和 KV Cache 中已有的 Value 向量。
通过这种方式，KV Cache 显著减少了每次生成新 token 时需要进行的计算量，从而大幅提升了推理速度。
KV Cache 的优势和挑战
优势：
- 加速推理：这是最主要的优势。对于长序列生成，KV Cache 可以将计算复杂度从 $O(N^2)$ 降低到近似 $O(N)$（其中 $N$ 是序列长度），从而显著提高生成速度。
- 提高吞吐量：通过减少单次推理时间，可以更快地处理更多请求，提高系统整体的吞吐量。
挑战：
- 内存消耗：KV Cache 需要存储所有层、所有注意力头、所有已处理 token 的 Key 和 Value 向量。随着序列长度的增加，KV Cache 的内存占用会线性增长。对于大型模型和长上下文场景，KV Cache 可能会占用大量的 GPU 显存，成为推理的瓶颈。
- 例如，一个 130 亿参数的模型，单个 token 的 KV Cache 可能需要 800KB 左右的内存，如果处理 2048 个 token 的序列，单个请求的 KV Cache 就能达到 1.6GB。当前 GPU 显存有限，这会严重限制并发请求的数量。
- 内存管理：在实际部署中，需要高效地管理 KV Cache 的内存分配、回收和重用，以最大化 GPU 利用率。
优化 KV Cache 的技术
为了应对 KV Cache 带来的内存挑战，研究人员和开发者提出了多种优化技术：
量化 (Quantization)：
- 将 KV Cache 中的 Key 和 Value 向量从浮点数（如 FP16）量化为更低精度的数据类型（如 INT8 或 INT4）。这可以直接减少 KV Cache 的内存占用，例如，从 FP16 到 INT4 可以将内存占用减少 4 倍。
- 常见的量化方法包括 per-head per-token 的非对称量化。
- 挑战在于需要在内存节省和模型精度损失之间找到平衡。
多查询注意力 (Multi-Query Attention, MQA) 和分组查询注意力 (Grouped-Query Attention, GQA)：
- MQA：传统的多头注意力（Multi-Head Attention, MHA）中，每个注意力头都有自己独立的 Key 和 Value 投影。MQA 提出让所有查询头共享一套 Key 和 Value 投影。这样，KV Cache 的大小就不再随着查询头的数量线性增长，从而显著减少内存占用。
- GQA：MQA 的一个推广，将查询头分成若干组，每组查询头共享一套 Key 和 Value 投影。这在 MQA 的极致内存节省和 MHA 的更高表达能力之间取得了平衡。
- 这些技术在模型架构层面减少了 Key 和 Value 向量的数量，从而直接降低了 KV Cache 的大小。
跨层注意力 (Cross-Layer Attention, CLA)：
- CLA 提出在 Transformer 的不同层之间共享 Key 和 Value 投影。例如，每 2 层或 3 层共享一套 KV 投影。
- 这进一步减少了 KV Cache 的内存占用，因为不需要为每一层都存储独立的 K 和 V 向量。
KV Cache 压缩：
- 例如，Finch 等方法通过识别并保留对提示词最相关的 K/V 对，丢弃不重要的 K/V 对，从而压缩 KV Cache。这在长上下文场景中特别有用。
内存分块和分页 (Paging)：
- 类似于操作系统中的内存管理，将 KV Cache 划分为更小的块（或页），根据实际需求动态分配和释放这些块，而不是预先分配整个最大长度的 KV Cache。这可以提高内存利用率，尤其是在处理不同长度的序列时。
离线 KV Cache：
- 在某些场景下，可以将 KV Cache 存储到磁盘或其他存储介质中，以处理超长序列或进行持久化，但在推理时会引入额外的加载延迟。
总结
KV Cache 是现代大型语言模型高效推理的关键技术之一。它通过缓存注意力机制中的 Key 和 Value 向量，避免了重复计算，从而显著提升了生成速度。然而，KV Cache 也带来了巨大的内存开销。因此，KV Cache 的优化是 LLM 部署和服务中的一个重要研究方向，各种技术如量化、MQA/GQA、CLA 和 KV Cache 压缩等都在不断发展，以在速度、内存和精度之间取得最佳平衡。