欢迎关注我的公众号 [极智视界]，获取我的更多经验分享
大家好，我是极智视界，本文来介绍一下 大模型优化之KV Cache。
邀您加入我的知识星球「极智视界」，星球内有超多好玩的项目实战源码下载，链接：t.zsxq.com/0aiNxERDq
在大模型的优化中经常会听到的一个技术叫 KV Cache，要理解这个技术需要你有基本的大模型概念以及 Transformer 基础。KV Cache 的全称是 key-value cache，可以简单理解为对大模型推理过程中的 key-value 缓存的优化。如果不对大模型推理过程中的 key-value 缓存进行优化，它会随着对话的增加而不断增加，也就是所占用的内存会不断动态增加，而且这种动态内存增加是不太可控的，这就像软件里的内存溢出一样。不去做 KV Cache，会对大规模推理造成几个压力：(1) 频繁的开辟增加内存；(2) 内存很快就不够用了；(3) 增加了很多冗余矩阵计算量。所以进行 KV 缓存优化是大规模训练、推理里很重要的一个环节。KV Cache 采用以空间换时间的思想，复用上次推理的 KV 缓存，可以极大降低内存压力、提高推理性能，而且不会影响任何计算精度。
以 GPT 为代表的一个 token 一个 token 往外蹦的 AIGC 大模型为例，里面最主要的结构就是 transformer 中的 self-attention 结构的堆叠，实质是将之前计算过的 key-value 以及当前的 query 来生成下一个 token。
这个时候如果不做优化，就会产生重复计算，用下面的图来进行说明。
下面两张图是没有 KV Cache 生成 token 的流程，假如 prompt 是 “我热爱打”，首先会生成对应的 Q 向量、K 向量和 V 向量，然后通过 Self-Attention 模块以及 MLP&LMH 模块，然后就往外吐生成 “热爱打篮”。而这里面真正有效的是生成的第四个 token – “篮”，当生成 “篮” 之后呢就要把这个 “篮” 接到 “我热爱打” 的后面，形成下一个输入 “我热爱打篮”，也就是下面的第二张示意图，依次类推。这个过程中其实会发现，其实中间一部分的 token 是没有必要每次都去生成的，最后咱们只是想要最后一个 token，前面的 token 计算其实是多余的。这样就考虑有了下面的 KV Cache 优化。
下面两张图展示了 KV Cache 优化，下面第一张图表示生成第一个 token，也就是预填充阶段，其实和上面的第一张图一致，多出来的是将 key-value Cache 预存起来。在下面第二张图生成第二个 token 的时候，就使用了 KV Cache 优化，Self-Attention 中用到的 K、V 的前四个向量 0、1、2、3 是直接从第一张图中预存的 KV Cache 中直接拿过来拼接起来用的，这样直接就生成了第五个 token “篮”，前面的不用重复计算，大大提升推理速度和降低内存占用。然后把这个阶段的 KV 给 Cache 起来，用于生成下一个 token，以此类推。
所以 KV Cache 优化其实主要包含两个阶段：
- 预填充阶段；=> 正常推理，预存 key-value cache；=> compute-bound 计算；
- KV Cache阶段；=> 预取上一阶段的 key-value cache，只需计算最后一个 token，将 gemm 变为 gemv 操作，最后预存追加后的 key-value cache 给下一阶段用；=> memory-bound 计算；
好了，以上分享了 大模型优化之KV Cache，希望我的分享能对你的学习有一点帮助。
【公众号传送】
畅享人工智能的科技魅力，让好玩的AI项目不难玩。邀请您加入我的知识星球， 星球内我精心整备了大量好玩的AI项目，皆以工程源码形式开放使用，涵盖人脸、检测、分割、多模态、AIGC、自动驾驶、工业等。不敢说会对你学习有所帮助，但一定非常好玩，并持续更新更加有趣的项目。 t.zsxq.com/0aiNxERDq
