📚 课程目录第3章 KV Cache基础

3. KV Cache基础：什么是KV Cache，如何存储和利用

3.1 从自回归解码说起

大语言模型（LLM）在推理时采用自回归（Autoregressive）方式逐token生成输出。以生成“人工智能”为例，模型先根据输入“人工”预测下一个token“智”，再将“智”拼接到输入末尾，预测“能”，如此循环。每一次预测都依赖完整的上文序列，这意味着如果不做任何优化，第 t 步需要重新计算前 t-1 个token的注意力分数，计算量随序列长度呈二次方增长。

KV Cache的核心思想是：将已经计算过的Key和Value向量缓存下来，避免重复计算。在每一步解码时，只需计算当前新token的Query、Key、Value，然后用缓存的KV与当前Query做注意力计算即可。这一机制将单步推理的复杂度从 O(t²) 降为 O(t)，是LLM高效推理的基石。

核心公式对比：

无KV Cache：Attention(Q_all, K_all, V_all) = softmax(Q_all·K_all^T / √d)·V_all
  每次生成新token都要重新计算全部位置的K、V

有KV Cache：Attention(Q_new, K_cache, V_cache) = softmax(Q_new·K_cache^T / √d)·V_cache
  只需计算新token的Q，K/V从缓存中读取

3.2 KV Cache的存储结构

KV Cache在推理过程中以张量（Tensor）形式存储在GPU显存中。其核心维度包括：

层数（L）：Transformer的层数，每层都有独立的KV Cache。
批次大小（B）：当前推理的序列数量。
序列长度（S）：已生成的token数量，随解码过程动态增长。
注意力头数（H）：多头注意力中的头数。
头维度（D）：每个注意力头的Key/Value向量维度。

# KV Cache张量形状（以PyTorch为例）
# 假设：L=32层, B=1, S=1024, H=32, D=128
key_cache_shape   = (L, B, H, S, D)  # 例如 (32, 1, 32, 1024, 128)
value_cache_shape = (L, B, H, S, D)  # 与Key形状相同

# 显存占用估算（FP16，每个元素2字节）
memory_bytes = 2 * L * B * H * S * D * 2  # 乘以2是因为K和V各一份
memory_gb = memory_bytes / (1024**3)
print(f"显存占用约: {memory_gb:.2f} GB")

参数	典型值	说明
层数 L	32 / 40 / 80	LLaMA-7B为32层，LLaMA-70B为80层
注意力头数 H	32 / 40 / 64	与模型隐藏维度相关
头维度 D	64 / 128	通常为隐藏维度除以头数
KV Cache总量	2 × L × H × S × D	K和V各一份，FP16存储

3.3 实际推理中的KV Cache利用流程

在推理引擎（如vLLM、TensorRT-LLM）中，KV Cache的利用遵循以下步骤：

预填充（Prefill）阶段：处理完整的输入提示词，计算所有输入token的K/V并存入缓存。
解码（Decode）阶段：每步只处理一个新token，计算其K/V，追加到缓存末尾。
注意力计算：当前token的Query与缓存中所有K做点积，再与V加权求和。
缓存更新：将新token的K/V追加到缓存，序列长度S增加1。

工程实践要点：

KV Cache通常以预分配内存块的方式管理，避免动态扩容带来的显存碎片。
在批处理（Batching）场景下，不同序列长度不同，需要Padding或使用PagedAttention（后续章节详述）来高效管理。
FP16存储可节省一半显存，但需注意精度损失；INT8量化可进一步压缩。

3.4 案例：LLaMA-7B的KV Cache显存估算

以LLaMA-7B为例，其配置为：L=32层，H=32头，D=128，隐藏维度4096。假设输入序列长度为2048，批次大小为1，使用FP16存储：

# 单序列KV Cache显存计算
L, H, D, S = 32, 32, 128, 2048
bytes_per_element = 2  # FP16

# 每层KV大小 = 2 (K和V) × H × S × D
per_layer = 2 * H * S * D * bytes_per_element
total_bytes = L * per_layer

print(f"每层KV Cache: {per_layer / 1024**2:.1f} MB")
print(f"总KV Cache: {total_bytes / 1024**3:.2f} GB")
# 输出：每层KV Cache: 16.0 MB
#       总KV Cache: 0.50 GB

当序列长度增长到8192时，KV Cache将占用约2GB显存。对于70B模型（L=80, H=64, D=128），同样序列长度下KV Cache将超过16GB，这解释了为什么长序列推理对显存压力巨大，也凸显了高效KV Cache管理（如PagedAttention）的必要性。

3.5 常见误区与复盘

误区1：KV Cache是模型参数的一部分
KV Cache是推理过程中动态生成的临时缓存，不属于模型权重。模型权重在推理时保持不变，而KV Cache随输入和生成过程变化。

误区2：KV Cache只存Key不存Value
注意力计算需要同时使用K和V，两者缺一不可。只缓存K会导致Value重复计算，无法达到加速目的。

误区3：KV Cache越大越好
KV Cache占用显存，过大的缓存会挤占模型权重和计算所需空间，甚至导致OOM。需要根据显存容量和序列长度合理设置最大缓存大小。

复盘总结：KV Cache是LLM推理加速的核心机制，通过空间换时间，将自回归解码的重复计算转化为增量计算。理解其存储结构（L×B×H×S×D）和显存占用规律，是后续学习PagedAttention、连续批处理等高级优化技术的基础。在实际部署中，需要根据模型规模、显存容量和业务场景（如长文档问答、多轮对话）合理配置KV Cache策略。