极智AI | 大模型优化之KV Cache

欢迎关注我的公众号 [极智视界],获取我的更多经验分享

大家好,我是极智视界,本文来介绍一下 大模型优化之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

相关推荐
梦想画家6 分钟前
PyTorch梯度:深度学习的引擎与实战解析
人工智能·pytorch·python
政安晨2 小时前
【嵌入式人工智能产品开发实战】(二十)—— 政安晨:小智AI嵌入式终端代码解读:【B】小智AI嵌入式终端OTA升级功能深度解析
c++·人工智能·嵌入式·ota·小智ai·代码解读·ai聊天助手
pzx_0012 小时前
【深度学习】自定义实现DataSet和DataLoader
开发语言·人工智能·python·深度学习·算法
努力向上的年轻人2 小时前
DeepSeek 与开源:肥沃土壤孕育 AI 硕果
人工智能·gitee·开源·deepseek·飞桨paddlepaddle·昇思mindspore
人工智能教学实践4 小时前
【深度学习】神经网络中,BatchNormalization(批量归一化)层
人工智能·深度学习·神经网络
起个破名想半天了8 小时前
Sklearn入门之datasets的基本用法
人工智能·python·机器学习·sklearn
微小冷8 小时前
微软出品的AI Toolkit,在VS Code中使用DeepSeek
人工智能·microsoft·插件·vs code·deepseek
IT古董9 小时前
【漫话机器学习系列】197.外核(Out of Core)
人工智能·机器学习
AIGC方案9 小时前
免费下载 | 2025天津大学:智能制造与数字孪生技术:面向可持续制造方向发展
人工智能·制造
yumuing9 小时前
告别 “生成废图”!UNO 让你一键掌控多物体图像生成: Less-to-More Generalization
aigc·计算机图形学