KV Cache: 大模型自回归生成的关键优化技术KV Cache键值缓存是 Transformer 解码器在自回归推理阶段普遍采用的一种工程优化。它通过缓存每一层注意力计算中历史 token 的Key和Value矩阵避免每一步都重复计算显著降低推理延迟是 LLM 高性能服务的基础组件。1. 背景自回归生成中的重复计算问题大型语言模型LLM在生成文本时采用自回归方式逐个生成 token并将新 token 拼接到输入序列末尾作为下一步的输入。假设已经生成了t − 1 t-1t−1个 token当前输入序列长度为t − 1 t-1t−1下一步需要生成第t tt个 token。在未使用 KV Cache 的情况下每一步都要重新计算整个序列[ x 1 , x 2 , . . . , x t − 1 , x t ] [x_1, x_2, ..., x_{t-1}, x_t][x1,x2,...,xt−1,xt]的所有注意力 K 和 V。而实际上前t − 1 t-1t−1个 token 的 K 和 V 在前几步已经计算过完全可以直接复用。不加缓存时第t tt步的计算复杂度为O ( t 2 ) \mathcal{O}(t^2)O(t2)因为需要计算t × t t \times tt×t的注意力矩阵。随着序列长度t tt增大计算开销呈平方增长使推理变得极其缓慢。2. KV Cache 的核心思想只计算当前新 token 的 QueryQ、KeyK、ValueV并复用之前所有 token 已缓存的 K 和 V。缓存内容每一层 Transformer 中所有历史 token 的 Key 矩阵和 Value 矩阵。更新方式每生成一个新 token计算其 K、V 并追加到缓存中。注意力计算使用当前 token 的 Q 与缓存中所有历史 K 计算注意力分数再用该分数加权缓存中所有历史 V。这样第t tt步的计算复杂度从O ( t 2 ) \mathcal{O}(t^2)O(t2)降为O ( t ) \mathcal{O}(t)O(t)因为只需要对新 token 计算一次 Q并与长度为t tt的 K 序列做点积即Q K T Q K^TQKT的其中一行。3. 工作机制与流程以单层自注意力为例。假设已经生成了t − 1 t-1t−1个 token缓存了K cache [ k 1 , k 2 , … , k t − 1 ] K_{\text{cache}} [k_1, k_2, \dots, k_{t-1}]Kcache[k1,k2,…,kt−1]V cache [ v 1 , v 2 , … , v t − 1 ] V_{\text{cache}} [v_1, v_2, \dots, v_{t-1}]Vcache[v1,v2,…,vt−1]第t tt步生成新 tokenx t x_txt的过程将x t x_txt通过 embedding 层和位置编码得到向量表示。通过线性变换计算该 token 的查询q t q_tqt、键k t k_tkt、值v t v_tvtq t x t W Q , k t x t W K , v t x t W V q_t x_t W_Q,\quad k_t x_t W_K,\quad v_t x_t W_VqtxtWQ,ktxtWK,vtxtWV更新缓存K cache ← concat ( K cache , k t ) K_{\text{cache}} \leftarrow \text{concat}(K_{\text{cache}}, k_t)Kcache←concat(Kcache,kt)V cache ← concat ( V cache , v t ) V_{\text{cache}} \leftarrow \text{concat}(V_{\text{cache}}, v_t)Vcache←concat(Vcache,vt)计算当前 token 的注意力输出注意力分数s t q t ⋅ [ k 1 , k 2 , … , k t ] T s_t q_t \cdot [k_1, k_2, \dots, k_t]^Tstqt⋅[k1,k2,…,kt]T长度为t tt的行向量。权重a t softmax ( s t / d k ) a_t \text{softmax}(s_t / \sqrt{d_k})atsoftmax(st/dk)其中d k d_kdk是 head 维度。输出o t a t ⋅ V cache o_t a_t \cdot V_{\text{cache}}otat⋅Vcache。经过 FFN 等后续层得到输出 logits采样得到下一个 tokenx t 1 x_{t1}xt1。每一步只计算一个q t q_tqt因此注意力计算的复杂度为O ( t ) \mathcal{O}(t)O(t)。注意第一个 tokenprefill 阶段仍需计算完整序列这一步无法加速。4. 为什么可以安全地缓存 K 和 V在 Transformer Decoder 中通过因果注意力掩码causal attention mask确保每个 token 只能看到它之前包括自己的 token不能看到未来 token。因此历史 token 的 K 和 V 不会依赖任何未来的信息它们是静态的。无论后续生成多少 token已有 token 的 K、V 都保持不变这为缓存提供了理论依据。5. 内存开销分析KV Cache 消耗的显存与batch_size、seq_length、层数、注意力头数等因素成正比。具体公式如下Memory batch_size × seq_len × num_layers × num_heads × head_dim × 2 × bytes_per_elem \text{Memory} \text{batch\_size} \times \text{seq\_len} \times \text{num\_layers} \times \text{num\_heads} \times \text{head\_dim} \times 2 \times \text{bytes\_per\_elem}Memorybatch_size×seq_len×num_layers×num_heads×head_dim×2×bytes_per_elem其中2 22表示同时缓存 K 和 V。以LLaMA 7B为例num_layers 32num_heads 32head_dim 128每 token 每层的 KV 参数量 2 × 32 × 128 8192 2 \times 32 \times 128 81922×32×1288192使用 FP162 bytes/param时每 token 每层占用8192 × 2 16 KB 8192 \times 2 16\ \text{KB}8192×216KB32 层共32 × 16 KB 512 KB 32 \times 16\ \text{KB} 512\ \text{KB}32×16KB512KB每 token当batch_size 64seq_len 2048时Memory 64 × 2048 × 512 KB 64 GB \text{Memory} 64 \times 2048 \times 512\ \text{KB} 64\ \text{GB}Memory64×2048×512KB64GB这远超模型参数本身的内存7B FP16 约 14 GB。因此KV Cache 是 LLM 推理时的显存瓶颈尤其在长上下文和大 batch 场景下。6. 实现细节与代码示例在工程实现中KV Cache 通常以past_key_values的形式在各个解码层之间传递。以下是一个极简的 PyTorch 风格伪代码classDecoderLayer:def__init__(self,d_model,n_head):self.Wqnn.Linear(d_model,d_model)self.Wknn.Linear(d_model,d_model)self.Wvnn.Linear(d_model,d_model)# ... 其他组件defforward(self,x,past_kvNone):qself.Wq(x)# [seq_len, d_model]kself.Wk(x)vself.Wv(x)ifpast_kvisnotNone:past_k,past_vpast_kv ktorch.cat([past_k,k],dim0)# 沿序列长度拼接vtorch.cat([past_v,v],dim0)# 注意力计算带因果掩码attn_outscaled_dot_product_attention(q,k,v,maskcausal_mask)returnattn_out,(k,v)# 推理循环past_kvs[None]*num_layers# 每层单独缓存for_inrange(max_new_tokens):fori,layerinenumerate(layers):x,past_kvs[i]layer(x,past_kvpast_kvs[i])# 采样下一个 tokennext_tokensample(x[-1])xnext_token.unsqueeze(0)# 下一步仅输入新 token在实际框架如 Hugging Face Transformers中past_key_values 是一个长度为 num_layers 的tuple每个元素是(K,V)两个张量形状为[batch_size,num_heads,past_seq_len,head_dim]。7. 优化变体降低 KV Cache 内存为了缓解显存压力工业界和学术界提出了多种优化方案。7.1 Multi-Query Attention (MQA)所有 Q 头共享同一组 K 和 V。因此 KV Cache 大小降为原来的1 num_heads \frac{1}{\text{num\_heads}}num_heads1。例如 PaLM、Falcon 等模型采用 MQA推理速度提升明显但可能带来轻微的质量损失。7.2 Grouped-Query Attention (GQA)将 Q 头分组每组共享一个 K、V 头。这是 MQA 和标准多头注意力的折中方案。例如 LLaMA 2/3 中7B 模型使用 32 个 Q 头、8 个 KV 头缓存大小降至1 4 \frac{1}{4}41。Gemma、DeepSeek 等模型也广泛采用 GQA。7.3 PagedAttention将 KV Cache 分割成固定大小的“页”以非连续方式存储。这可以消除显存碎片并支持变长序列的高效批处理。由 vLLM 项目提出目前是生产环境推理的标准技术之一。7.4 量化缓存对 KV Cache 使用 INT8 甚至 INT4 量化。例如 KIVI、KVQuant 等工作表明适度量化每 channel 或 per-token可以在保持生成质量的同时减少 50%~75% 的缓存占用。7.5 滑动窗口与动态缓存只保留最近w ww个 token 的 K、V丢弃更早的信息。适用于局部依赖较强的任务如代码补全、流式生成。Longformer、BigBird 等模型原生支持窗口注意力但会丧失长距离建模能力。8. 局限性显存消耗长序列如 100k tokens即使单 batch 推理KV Cache 也可能填满多张 80GB A100。批处理不友好不同请求的序列长度差异导致缓存形状不规则需借助 PagedAttention 等技术。首 token 延迟第一个 tokenprefill仍需计算完整序列的 K、V无法加速。9. 总结方面说明作用消除自回归生成中历史 token 的重复 K/V 计算将单步复杂度从O ( t 2 ) \mathcal{O}(t^2)O(t2)降至O ( t ) \mathcal{O}(t)O(t)代价显存占用随序列长度线性增长成为长上下文推理的主要瓶颈实现形式Hugging Facepast_key_values或手工维护每层(K, V)张量优化方向MQA/GQA、量化、PagedAttention、滑动窗口适用范围所有基于 Transformer Decoder 的自回归模型GPT、LLaMA、ChatGLM、Qwen 等理解 KV Cache 是分析 LLM 推理速度、显存占用以及设计高效部署方案的基础。结合 FlashAttention、Continuous Batching 等技术可以进一步释放硬件潜力。参考文献与延伸阅读Efficiently Scaling Transformer Inference (Google Research, 2022)GQA: Training Generalized Multi-Query Transformer Models (arXiv, 2023)PagedAttention: vLLM (SOSP, 2023)Hugging Face Transformers: past_key_values 文档
KV Cache 详解
发布时间：2026/5/28 2:04:49
KV Cache: 大模型自回归生成的关键优化技术KV Cache键值缓存是 Transformer 解码器在自回归推理阶段普遍采用的一种工程优化。它通过缓存每一层注意力计算中历史 token 的Key和Value矩阵避免每一步都重复计算显著降低推理延迟是 LLM 高性能服务的基础组件。1. 背景自回归生成中的重复计算问题大型语言模型LLM在生成文本时采用自回归方式逐个生成 token并将新 token 拼接到输入序列末尾作为下一步的输入。假设已经生成了t − 1 t-1t−1个 token当前输入序列长度为t − 1 t-1t−1下一步需要生成第t tt个 token。在未使用 KV Cache 的情况下每一步都要重新计算整个序列[ x 1 , x 2 , . . . , x t − 1 , x t ] [x_1, x_2, ..., x_{t-1}, x_t][x1,x2,...,xt−1,xt]的所有注意力 K 和 V。而实际上前t − 1 t-1t−1个 token 的 K 和 V 在前几步已经计算过完全可以直接复用。不加缓存时第t tt步的计算复杂度为O ( t 2 ) \mathcal{O}(t^2)O(t2)因为需要计算t × t t \times tt×t的注意力矩阵。随着序列长度t tt增大计算开销呈平方增长使推理变得极其缓慢。2. KV Cache 的核心思想只计算当前新 token 的 QueryQ、KeyK、ValueV并复用之前所有 token 已缓存的 K 和 V。缓存内容每一层 Transformer 中所有历史 token 的 Key 矩阵和 Value 矩阵。更新方式每生成一个新 token计算其 K、V 并追加到缓存中。注意力计算使用当前 token 的 Q 与缓存中所有历史 K 计算注意力分数再用该分数加权缓存中所有历史 V。这样第t tt步的计算复杂度从O ( t 2 ) \mathcal{O}(t^2)O(t2)降为O ( t ) \mathcal{O}(t)O(t)因为只需要对新 token 计算一次 Q并与长度为t tt的 K 序列做点积即Q K T Q K^TQKT的其中一行。3. 工作机制与流程以单层自注意力为例。假设已经生成了t − 1 t-1t−1个 token缓存了K cache [ k 1 , k 2 , … , k t − 1 ] K_{\text{cache}} [k_1, k_2, \dots, k_{t-1}]Kcache[k1,k2,…,kt−1]V cache [ v 1 , v 2 , … , v t − 1 ] V_{\text{cache}} [v_1, v_2, \dots, v_{t-1}]Vcache[v1,v2,…,vt−1]第t tt步生成新 tokenx t x_txt的过程将x t x_txt通过 embedding 层和位置编码得到向量表示。通过线性变换计算该 token 的查询q t q_tqt、键k t k_tkt、值v t v_tvtq t x t W Q , k t x t W K , v t x t W V q_t x_t W_Q,\quad k_t x_t W_K,\quad v_t x_t W_VqtxtWQ,ktxtWK,vtxtWV更新缓存K cache ← concat ( K cache , k t ) K_{\text{cache}} \leftarrow \text{concat}(K_{\text{cache}}, k_t)Kcache←concat(Kcache,kt)V cache ← concat ( V cache , v t ) V_{\text{cache}} \leftarrow \text{concat}(V_{\text{cache}}, v_t)Vcache←concat(Vcache,vt)计算当前 token 的注意力输出注意力分数s t q t ⋅ [ k 1 , k 2 , … , k t ] T s_t q_t \cdot [k_1, k_2, \dots, k_t]^Tstqt⋅[k1,k2,…,kt]T长度为t tt的行向量。权重a t softmax ( s t / d k ) a_t \text{softmax}(s_t / \sqrt{d_k})atsoftmax(st/dk)其中d k d_kdk是 head 维度。输出o t a t ⋅ V cache o_t a_t \cdot V_{\text{cache}}otat⋅Vcache。经过 FFN 等后续层得到输出 logits采样得到下一个 tokenx t 1 x_{t1}xt1。每一步只计算一个q t q_tqt因此注意力计算的复杂度为O ( t ) \mathcal{O}(t)O(t)。注意第一个 tokenprefill 阶段仍需计算完整序列这一步无法加速。4. 为什么可以安全地缓存 K 和 V在 Transformer Decoder 中通过因果注意力掩码causal attention mask确保每个 token 只能看到它之前包括自己的 token不能看到未来 token。因此历史 token 的 K 和 V 不会依赖任何未来的信息它们是静态的。无论后续生成多少 token已有 token 的 K、V 都保持不变这为缓存提供了理论依据。5. 内存开销分析KV Cache 消耗的显存与batch_size、seq_length、层数、注意力头数等因素成正比。具体公式如下Memory batch_size × seq_len × num_layers × num_heads × head_dim × 2 × bytes_per_elem \text{Memory} \text{batch\_size} \times \text{seq\_len} \times \text{num\_layers} \times \text{num\_heads} \times \text{head\_dim} \times 2 \times \text{bytes\_per\_elem}Memorybatch_size×seq_len×num_layers×num_heads×head_dim×2×bytes_per_elem其中2 22表示同时缓存 K 和 V。以LLaMA 7B为例num_layers 32num_heads 32head_dim 128每 token 每层的 KV 参数量 2 × 32 × 128 8192 2 \times 32 \times 128 81922×32×1288192使用 FP162 bytes/param时每 token 每层占用8192 × 2 16 KB 8192 \times 2 16\ \text{KB}8192×216KB32 层共32 × 16 KB 512 KB 32 \times 16\ \text{KB} 512\ \text{KB}32×16KB512KB每 token当batch_size 64seq_len 2048时Memory 64 × 2048 × 512 KB 64 GB \text{Memory} 64 \times 2048 \times 512\ \text{KB} 64\ \text{GB}Memory64×2048×512KB64GB这远超模型参数本身的内存7B FP16 约 14 GB。因此KV Cache 是 LLM 推理时的显存瓶颈尤其在长上下文和大 batch 场景下。6. 实现细节与代码示例在工程实现中KV Cache 通常以past_key_values的形式在各个解码层之间传递。以下是一个极简的 PyTorch 风格伪代码classDecoderLayer:def__init__(self,d_model,n_head):self.Wqnn.Linear(d_model,d_model)self.Wknn.Linear(d_model,d_model)self.Wvnn.Linear(d_model,d_model)# ... 其他组件defforward(self,x,past_kvNone):qself.Wq(x)# [seq_len, d_model]kself.Wk(x)vself.Wv(x)ifpast_kvisnotNone:past_k,past_vpast_kv ktorch.cat([past_k,k],dim0)# 沿序列长度拼接vtorch.cat([past_v,v],dim0)# 注意力计算带因果掩码attn_outscaled_dot_product_attention(q,k,v,maskcausal_mask)returnattn_out,(k,v)# 推理循环past_kvs[None]*num_layers# 每层单独缓存for_inrange(max_new_tokens):fori,layerinenumerate(layers):x,past_kvs[i]layer(x,past_kvpast_kvs[i])# 采样下一个 tokennext_tokensample(x[-1])xnext_token.unsqueeze(0)# 下一步仅输入新 token在实际框架如 Hugging Face Transformers中past_key_values 是一个长度为 num_layers 的tuple每个元素是(K,V)两个张量形状为[batch_size,num_heads,past_seq_len,head_dim]。7. 优化变体降低 KV Cache 内存为了缓解显存压力工业界和学术界提出了多种优化方案。7.1 Multi-Query Attention (MQA)所有 Q 头共享同一组 K 和 V。因此 KV Cache 大小降为原来的1 num_heads \frac{1}{\text{num\_heads}}num_heads1。例如 PaLM、Falcon 等模型采用 MQA推理速度提升明显但可能带来轻微的质量损失。7.2 Grouped-Query Attention (GQA)将 Q 头分组每组共享一个 K、V 头。这是 MQA 和标准多头注意力的折中方案。例如 LLaMA 2/3 中7B 模型使用 32 个 Q 头、8 个 KV 头缓存大小降至1 4 \frac{1}{4}41。Gemma、DeepSeek 等模型也广泛采用 GQA。7.3 PagedAttention将 KV Cache 分割成固定大小的“页”以非连续方式存储。这可以消除显存碎片并支持变长序列的高效批处理。由 vLLM 项目提出目前是生产环境推理的标准技术之一。7.4 量化缓存对 KV Cache 使用 INT8 甚至 INT4 量化。例如 KIVI、KVQuant 等工作表明适度量化每 channel 或 per-token可以在保持生成质量的同时减少 50%~75% 的缓存占用。7.5 滑动窗口与动态缓存只保留最近w ww个 token 的 K、V丢弃更早的信息。适用于局部依赖较强的任务如代码补全、流式生成。Longformer、BigBird 等模型原生支持窗口注意力但会丧失长距离建模能力。8. 局限性显存消耗长序列如 100k tokens即使单 batch 推理KV Cache 也可能填满多张 80GB A100。批处理不友好不同请求的序列长度差异导致缓存形状不规则需借助 PagedAttention 等技术。首 token 延迟第一个 tokenprefill仍需计算完整序列的 K、V无法加速。9. 总结方面说明作用消除自回归生成中历史 token 的重复 K/V 计算将单步复杂度从O ( t 2 ) \mathcal{O}(t^2)O(t2)降至O ( t ) \mathcal{O}(t)O(t)代价显存占用随序列长度线性增长成为长上下文推理的主要瓶颈实现形式Hugging Facepast_key_values或手工维护每层(K, V)张量优化方向MQA/GQA、量化、PagedAttention、滑动窗口适用范围所有基于 Transformer Decoder 的自回归模型GPT、LLaMA、ChatGLM、Qwen 等理解 KV Cache 是分析 LLM 推理速度、显存占用以及设计高效部署方案的基础。结合 FlashAttention、Continuous Batching 等技术可以进一步释放硬件潜力。参考文献与延伸阅读Efficiently Scaling Transformer Inference (Google Research, 2022)GQA: Training Generalized Multi-Query Transformer Models (arXiv, 2023)PagedAttention: vLLM (SOSP, 2023)Hugging Face Transformers: past_key_values 文档
