如何通过KV缓存显著加速LLMs的推理速度🚀 1️⃣ KV缓存-20250215205924

如何通过KV缓存显著加速LLMs的推理速度🚀

1️⃣ KV缓存概述：KV缓存是一种用于加快LLM推理的技术。通过一个简单的比较（图1视频），我们可以看到KV缓存的强大之处：启用KV缓存的推理速度为9秒，而未启用时需要42秒，几乎慢了5倍！

2️⃣ LLMs的工作原理：要理解KV缓存，我们首先需要了解LLMs是如何输出令牌的。Transformer模型为所有Tokens产生隐藏状态，这些隐藏状态被映射到词汇空间，利用最后一个令牌的逻辑回归值生成下一个令牌，如此重复。

3️⃣ KV缓存的关键：生成新tokens时，我们实际上只需要最近一个token的隐藏状态。在Transformer层内部，通过注意力机制计算最后一个隐藏状态，每次注意力操作仅需最后一个token的查询向量以及所有的键值向量。

4️⃣ KV缓存的实现：当我们生成新token时，所有之前token用到的KV向量不会改变。因此，我们只需为前一步生成的令牌生成一个KV向量，其他的KV向量可以从缓存中检索，以节省计算资源和时间。

这项技术可以显著提高了推理效率。但同时也要注意其内存开销。比如对于Llama3-70B ：
- 总层数 = 80
- 隐藏大小 = 8k
- 最大输出大小 = 4k

这里：
- 每个token在 KV 缓存中占用约 2.5 MB。
- 4k 个token将占用 10.5 GB。

附图说明
图1：使用KV缓存和不使用KV缓存的推理速度对比，9秒 vs. 43秒
图2，LLM如何输出下一个Tokens，Transformer为所有Tokens生成隐藏状态->投射到词汇空间->Logits生成下一个Token->循环往复
图3，为了生成新的Token，只需计算最新Token
图4，注意力机制的QKV计算，可以看出注意力机制最后一行只依赖于最后一个向量和其他关键向量
图5，图4的流程可变部分（去掉可缓存部分）
图6，可从之前计算步骤复用部分（可缓存部分）
图7，汇总示意图

附图来自：Avi Chawla

#ai创造营# #deepseek# #程序员#