3、KV Cache优化：从原理到实战
聊到推理优化，KV Cache 是个绕不开的话题。我刚开始接触大模型推理时，第一反应是：这玩意儿不就是把计算过的 Key 和 Value 存起来吗？有什么好讲的？
后来在实际项目中，显存爆了、推理慢了，我才意识到——KV Cache 的优化，直接决定了你的模型能跑多快、能支持多长的上下文。
3.1 KV Cache 原理：为什么需要它？
先说说最基础的东西。Transformer 做自回归生成时，每次只生成一个 token。但生成第 t 个 token 时，需要用到前面所有 token 的注意力信息。
如果不做缓存，每次生成都要重新计算前面所有 token 的 Key 和 Value。你想想看，生成第 1000 个 token 时，前面 999 个 token 的 K、V 要重新算一遍——这计算量，谁受得了？
核心思路：把已经算好的 Key 和 Value 存起来，下次直接用。
具体来说，KV Cache 就是维护两个矩阵：
- K_cache：形状为 [batch_size, num_heads, seq_len, head_dim]
- V_cache：形状同上
每次生成新 token 时，只计算当前 token 的 K 和 V，然后拼接到缓存后面。这样注意力计算就变成了：
# 伪代码示意
K_full = torch.cat([K_cache, K_new], dim=-2) # 拼接历史 + 当前
V_full = torch.cat([V_cache, V_new], dim=-2)
attn_output = scaled_dot_product_attention(Q, K_full, V_full)嗯，这里要注意：Q 只来自当前 token，而 K、V 来自全部历史。这就是 KV Cache 能省计算的根本原因。
3.2 共享 KV Cache：多轮对话的利器
我在做对话系统时遇到一个场景：用户连续问多个问题，每个问题都基于之前的对话历史。如果每次新问题都重新计算整个历史的 KV Cache，那延迟会直线上升。
共享 KV Cache 的思路很简单——把历史对话的 KV Cache 保留下来，新问题来了直接复用。
我的经验：在实现共享 KV Cache 时，要注意处理「系统提示词」和「用户消息」的边界。我曾经踩过一个坑：把系统提示词的 KV Cache 和用户消息的混在一起，结果模型回答时上下文错乱了。
共享 KV Cache 的实现要点：
- 维护一个全局的 KV Cache 列表，按对话轮次分段
- 每轮对话结束后，只保留本轮新增的 KV Cache
- 新轮次开始时，从缓存中加载所有历史的 KV Cache
3.3 Multi-Query Attention (MQA)：显存杀手锏
标准的多头注意力（MHA）中，每个头都有自己的 K 和 V。假设有 32 个头，每个头的 head_dim 是 128，那一个 token 的 KV Cache 大小就是 32 × 128 × 2 = 8192 个元素。
对于 4096 的序列长度，batch_size=1 时，KV Cache 就是 4096 × 8192 ≈ 33M 个元素。如果 batch_size=64，那就是 2.1B 个元素——显存直接爆炸。
MQA 的做法很暴力：所有 Query 头共享同一组 Key 和 Value。
# MQA 结构示意
# Query: [batch, num_heads, seq_len, head_dim] # 多个头
# Key: [batch, 1, seq_len, head_dim] # 只有一个头
# Value: [batch, 1, seq_len, head_dim] # 只有一个头这样一来，KV Cache 的大小直接缩小到原来的 1/num_heads。比如 32 个头，KV Cache 就变成原来的 1/32。
注意：MQA 虽然省显存，但会损失模型表达能力。我在项目中测试过，某些任务上 MQA 比 MHA 掉了 1-2 个点的准确率。需要根据实际场景权衡。
3.4 Grouped-Query Attention (GQA)：折中之道
MQA 太激进，MHA 太费显存。GQA 就是那个「刚刚好」的方案。
GQA 的思路：把 Query 头分成若干组，每组共享一组 Key 和 Value。
举个例子：
- 假设有 32 个 Query 头
- 分成 8 组，每组 4 个 Query 头
- 每组共享 1 个 Key 头和 1 个 Value 头
- 总共 8 个 Key 头、8 个 Value 头
KV Cache 大小 = 8/32 = 1/4 的 MHA。既省了显存，又保留了足够的表达能力。
我个人习惯：在部署时，如果显存紧张但模型质量要求高，我会优先选择 GQA。比如 32 个 Query 头配 8 个 KV 头，这个比例在大多数场景下效果都不错。
3.5 三种方案的对比
| 方案 | KV Cache 大小 | 模型质量 | 适用场景 |
|---|---|---|---|
| MHA | 大（基准） | 最高 | 训练、高质量推理 |
| MQA | 小（1/num_heads） | 较低 | 显存极度受限 |
| GQA | 中（可调） | 接近 MHA | 大多数推理场景 |
3.6 知识体系结构图
下面这张图展示了 KV Cache 优化的核心逻辑：
3.7 实战建议
说了这么多，到底怎么选？我给出几条实际经验：
- 训练阶段：用 MHA，别省这点显存。模型质量第一。
- 推理阶段，显存充足：用 GQA，8 组 KV 头是个不错的起点。
- 推理阶段，显存紧张：上 MQA，配合共享 KV Cache，能省下大量显存。
- 多轮对话场景：一定要实现共享 KV Cache，不然每次新问题都要重新计算历史，延迟受不了。
避坑指南：我曾经在实现 GQA 时，把分组逻辑写错了——把 Query 头分成了 4 组，但每组内的 Query 头数量不一致，导致注意力计算时维度不匹配。调试了半天才发现。建议在实现时先打印出每个头的索引和分组情况，确认无误后再跑推理。
好了，KV Cache 优化这块就聊到这儿。说白了就是：用空间换时间，再用技巧省空间。MHA、MQA、GQA 这三兄弟，你根据实际情况选一个就行。
想要更多资料，联系可可