KV Cache：大语言模型推理加速的核心机制详解

🔍 KV Cache：大语言模型推理加速的核心机制详解

一、什么是 KV Cache？

在大语言模型（LLM）的自回归生成过程中，为了提升推理效率，KV Cache（Key/Value Cache） 是一个至关重要的优化机制。

简单定义：

KV Cache 是一种用于缓存 Transformer 模型中注意力机制所需的 Key 和 Value 向量的结构。

它允许模型在逐词生成时复用之前 token 的 K/V 值，从而避免重复计算，提高推理速度和资源利用率。

二、为什么需要 KV Cache？

在传统 Transformer 解码过程中，每一步生成新 token 都要重新计算整个序列的 attention 中的 Key 和 Value 向量，这会带来大量冗余计算。

例如，在生成句子 "The cat sat on the mat" 时：

第一次前向传播：输入 "The"，输出 "cat"
第二次前向传播：输入 "The cat"，输出 "sat"

如果每次都从头开始计算，效率非常低。而 KV Cache 的出现解决了这个问题。

三、KV Cache 的工作原理

1. 注意力机制回顾

Transformer 中的标准注意力公式如下：

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

其中：

$ Q $：当前 token 的 Query 向量
$ K $：历史 token 的 Key 向量
$ V $：历史 token 的 Value 向量
$ d_k $：Key 维度

2. KV Cache 如何运作？

KV Cache 的核心思想是：将已生成 token 的 Key 和 Value 缓存起来，后续只需计算当前 token 的 Query 向量即可完成 attention 运算。

示例流程：

Step 1: 输入 "The" → 生成 K1, V1 → 缓存
Step 2: 输入 "cat" → 生成 Q2 → 使用 K1, V1 计算 attention
Step 3: 输入 "sat" → 生成 Q3 → 使用 K1/K2, V1/V2 计算 attention
...
最终：每次只需计算当前 token 的 Q，其余 K/V 从 cache 中读取

四、KV Cache 的结构与存储方式

KV Cache 的结构通常是多层、多头、按时间步缓存的形式：

[[layer_0_k_cache, layer_0_v_cache],[layer_1_k_cache, layer_1_v_cache],...
]

每个 cache 层的数据形状为：

(batch_size, num_heads, seq_len, head_dim)

💡 图解示意如下：

[Layer 0] → [K Cache (seq_len=5), V Cache (seq_len=5)]
[Layer 1] → [K Cache (seq_len=5), V Cache (seq_len=5)]
...

随着生成过程进行，seq_len 不断增长，KV Cache 动态扩展。

五、KV Cache 的作用与优势

优势	描述
提高推理速度	避免重复计算历史 token 的 K/V 向量
节省显存	相比于重新计算，KV Cache 占用空间更小（虽然仍较大）
支持流式生成	实现逐词生成的同时保持上下文一致性
支持批量处理	多个请求可以并行处理而不冲突

六、KV Cache 的挑战与优化方法

1. 显存占用问题

KV Cache 的大小与以下几个因素有关：

参数	默认值	显存公式
Batch Size	B	$ B \times n_{layers} \times n_{heads} \times seq_len \times head_dim $
序列长度	L	同上
注意力头数	H	同上
向量维度	D	同上
数据类型	float16 / bfloat16	每个数值占 2 字节

例如，一个 32 层、每层 32 头、head_dim=128 的模型，在生成长度为 1024 的文本时，KV Cache 占用的显存约为：

$\text{Size} = 1 \times 32 \times 32 \times 1024 \times 128 \times 2 \approx 268MB$

如果是并发服务多个用户（如 batch size=8），则需要 2GB+，这就是为什么 KV Cache 优化如此重要。

七、KV Cache 的优化方法

技术名称	描述
Multi-Query Attention (MQA)	所有注意力头共享相同的 Key 和 Value 向量，极大减少 KV Cache 占用
Grouped Query Attention (GQA)	将 Query 分组，每组共享一组 Key/Value，是 MQA 的扩展形式
PagedAttention（vLLM 使用）	类似操作系统的分页机制，将 KV 缓存分成块，支持动态长度和高效内存利用
KV Cache 压缩	使用 INT8 或 FP8 等量化技术压缩缓存内容

这些技术都能有效降低 KV Cache 的内存占用，从而实现更大 batch size、更高并发数、更长上下文支持。

八、KV Cache 在实际中的应用（以 Med-R1 为例）

在你提供的论文《Med-R1》中，作者采用了 GRPO（Group Relative Policy Optimization）来训练视觉-语言模型。在这个过程中，KV Cache 的管理优化对于提升推理吞吐量和降低延迟起到了重要作用。

尽管 Med-R1 模型参数仅为 2B，但通过高效的 KV Cache 管理和 GRPO 强化学习策略，其推理性能甚至超过了 72B 的 Qwen2-VL-72B 模型。

九、KV Cache 的可视化图解

+---------------------------+
|        用户输入           |
|   "The cat sat ..."       |
+-------------+-------------+↓
+-------------+-------------+
|     KV Cache Manager      |
| 存储所有 token 的 K/V 向量 |
+-------------+-------------+↓
+-------------+-------------+
|     Attention Module      |
| 利用当前 Q 与历史 K/V 计算 attention |
+-------------+-------------+↓
+-------------+-------------+
|         输出下一个词       |
+---------------------------+

十、如何查看和控制 KV Cache 使用？

在 PyTorch 或 vLLM 中，可以通过如下方式监控和控制 KV Cache 的使用：

PyTorch 示例（伪代码）：

with torch.no_grad():for step in range(max_length):outputs = model(input_ids=prompt_ids, past_key_values=past_kv)next_token = outputs.logits.argmax(-1)prompt_ids = torch.cat([prompt_ids, next_token], dim=-1)past_kv = outputs.past_key_values

vLLM 示例：

from vllm import LLM, SamplingParamsllm = LLM(model="meta-llama/Llama-3-8B")
sampling_params = SamplingParams(max_tokens=50)outputs = llm.generate(["Explain quantum computing"], sampling_params)

vLLM 内部自动管理 KV Cache，无需手动维护。

十一、总结

关键点	内容
KV Cache 是什么？	存储每个 token 的 Key 和 Value 向量，用于 attention 计算
为什么要用？	避免重复计算，提升推理速度
有什么缺点？	显存占用高，尤其在长文本和多用户场景下
如何优化？	使用 GQA、MQA、PagedAttention 等技术