1. KV缓存基础
各位同学好,我是这门课的主讲。今天咱们聊聊KV缓存——这个让大模型推理快上几倍甚至几十倍的核心技术。
说实话,我第一次接触KV缓存时,也觉得它不过是个简单的“存中间结果”技巧。直到我在一个线上推理服务里,亲眼看着QPS从50飙到300,才真正意识到它的威力。嗯,这节课我们就从最基础的东西开始,把KV缓存的来龙去脉彻底讲清楚。
1.1 自注意力机制回顾
要理解KV缓存,得先回到Transformer的心脏——自注意力机制。说白了,自注意力就是让每个词都能“看到”序列里的所有其他词。
公式很简单,但很重要:
Attention(Q, K, V) = softmax(Q × K^T / √d_k) × V
其中:
- Q(Query):当前词的查询向量,相当于“我在找什么”
- K(Key):所有词的键向量,相当于“我有什么信息”
- V(Value):所有词的值向量,相当于“我的实际内容”
计算过程分三步:
- Q和K做点积,得到注意力分数矩阵
- 用softmax归一化,得到注意力权重
- 权重乘以V,加权求和得到输出
我在优化一个对话模型时发现,很多人把注意力机制想得太复杂。其实你想想看,它就是个“信息检索”的过程——Q是搜索词,K是索引,V是文档内容。
1.2 为什么需要KV缓存
好,问题来了。在自回归生成时,模型是一个词一个词往外蹦的。每次生成新词,都要重新计算整个序列的注意力。
举个例子:生成“今天天气真好”这句话:
- 第1步:输入“今天”,计算“今天”的注意力
- 第2步:输入“今天天气”,重新计算“今天”和“天气”的注意力
- 第3步:输入“今天天气真”,重新计算前面3个词的注意力
- ……
你发现问题了吗?每次生成新词,前面所有词的K和V都要重新算一遍。这太浪费了!
核心洞察:在自回归生成中,已生成词的K和V是固定不变的。只有新词的Q需要计算注意力。所以,把历史K和V缓存起来,就能避免重复计算。
我曾经在一个项目里,没做KV缓存,生成128个token的序列,单次推理耗时从5ms涨到了80ms。做了缓存后,稳定在6ms左右。差距就是这么明显。
1.3 KV缓存的核心思想
KV缓存的思想极其简单,但效果惊人:
把每一步计算出的K和V存起来,供后续步骤复用。
具体来说:
- 生成第1个token时,正常计算K和V,并存入缓存
- 生成第2个token时,只计算新token的K和V,从缓存读取历史K和V
- 以此类推,每次只计算新token的K和V
我的建议:刚开始实现KV缓存时,可以用Python字典来存,key是层索引,value是(K, V)元组。虽然效率不高,但逻辑清晰,方便调试。
1.4 KV缓存的计算流程
下面这张图展示了带KV缓存的完整推理流程:
流程其实不复杂:
- 新token进来,做embedding
- 进入多头注意力层:新token生成Q,同时生成自己的K和V
- 从缓存读取历史K和V,和新K、V拼接
- 用Q和完整的K、V计算注意力
- 把新K、V追加到缓存
- 继续走FFN、LayerNorm,输出结果
我曾经踩过的坑:在实现KV缓存时,忘记把新token的K和V追加到缓存里。结果第二次推理时,缓存里只有第一个token的K和V,后面的token全乱了。排查了半天才发现是缓存更新逻辑写错了。
1.5 缓存带来的性能提升
咱们用数据说话。假设模型有12层,每层hidden_size=768,生成100个token:
| 对比项 | 无KV缓存 | 有KV缓存 | 提升倍数 |
|---|---|---|---|
| 每步计算量 | O(n²) | O(n) | n倍 |
| 生成100 token总计算量 | O(100³) | O(100²) | ~100倍 |
| 内存占用 | 低(不存历史) | 高(存K和V) | 空间换时间 |
你想想看,生成100个token时,无缓存的计算量是100³,有缓存是100²。差了整整100倍!这就是为什么所有工业级推理框架都标配KV缓存。
1.6 缓存的内存开销
当然,天下没有免费的午餐。KV缓存占用的内存也不小:
缓存大小 = 2 × batch_size × seq_len × hidden_size × num_layers × dtype_bytes
举个例子:
- batch_size=1, seq_len=2048, hidden_size=4096, num_layers=32, float16
- 缓存大小 = 2 × 1 × 2048 × 4096 × 32 × 2 = 1,073,741,824 bytes ≈ 1GB
嗯,单条请求就要1GB显存。如果batch_size=32,那就是32GB。所以后面我们会讲各种缓存优化技巧,比如内存复用、量化压缩、窗口缓存等。
我的经验:在实际部署时,我习惯先估算最大序列长度下的缓存大小,然后预留20%的余量。不然遇到长文本生成,显存突然爆了,服务就挂了。
1.7 小结
这节课我们讲了:
- 自注意力机制的本质:Q找K,V给内容
- KV缓存的核心思想:存历史K和V,避免重复计算
- 计算流程:新token算Q和新K、V,和历史拼接后做注意力
- 性能收益:计算量从O(n²)降到O(n)
- 代价:显存占用增加,需要做优化
说白了,KV缓存就是用显存换速度。在推理场景下,这个交换非常划算。后面的课程,我们会深入各种缓存优化技术,把显存占用降下来,把推理速度提上去。