欢迎重视我的大众号 [极智视界]，获取我的更多经验分享
大家好，我是极智视界，本文来介绍一下 大模型优化之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 | 大模型优化之KV Cache》
畅享人工智能的科技魅力，让好玩的AI项目不难玩。邀请您参加我的常识星球， 星球内我精心整备了大量好玩的AI项目，皆以工程源码方式敞开运用，涵盖人脸、检测、分割、多模态、AIGC、自动驾驶、工业等。不敢说会对你学习有所协助，但一定十分好玩，并继续更新愈加有趣的项目。 t.zsxq.com/0aiNxERDq
