KV Cache:加速LLM推理的关键
1. KV 缓存是什么?
KV 缓存(KV Cache)是一种优化大语言模型(LLM, Large Language Models)推理速度的关键技术。最近爆火的DeepSeek,其首创的MLA技术,使得KV Cache降低了93%,在大语言模型的训练和推理上有效降低了对高性能GPU的性能要求(比如原先需要H100才能训练和推理,现在H20就能满足)。
在生成文本时,GPT 这样的模型需要不断计算自注意力(Self-Attention)。但实际上,我们可以缓存 Key 和 Value(K/V)向量,避免重复计算,从而显著加速推理。
先来看一个直观的对比:
- 不开启 KV 缓存 → 42 秒
- 使用 KV 缓存 → 9 秒,推理速度提升约 5 倍
为什么 KV 缓存能让推理速度提升这么多?让我们深入解析。
1. 什么是 KV Cache?
大语言模型在生成文本时,模型需要在每一步计算 注意力(Attention)。标准的 Transformer 计算 自注意力 时,每次都需要重新计算 Key(K) 和 Value(V) 矩阵,并与新的 Query(Q) 进行注意力计算:
但是,Key 和 Value 其实是静态的(不会变化),它们在过去的时间步已经计算过了。因此,我们可以缓存(Cache)这些 Key/Value,避免重复计算,提高生成效率。
2. 传统 Transformer 计算的性能瓶颈
标准 Transformer 计算 Attention(没有 KV Cache)
-
每次生成新 token xt 时:
- 重新计算所有 K 和 V (包括所有历史 token)。
- 计算新的 Query Qt。
- 计算注意力 softmax(QtKT≤t)V≤t。
-
计算复杂度:
- 每个 Token 计算 Attention 需要 O(tdk)(因为每个新 Query 需要和 t 个 Key 进行点积)。
- 生成 T 个 Token 的总复杂度为:O(T2dk)
- 随着生成长度增加,计算量呈二次增长!
3. KV Cache 的优化
KV Cache 机制
-
缓存 Key/Value 矩阵:
- 每个 Token 计算 Key/Value 后 存起来,而不是在每个时间步重新计算。
- 只需要存储形状为 (t,dk) 和 (t,dv) 的矩阵。
-
下一个 Token 计算时:
- 只需计算新的 Query Qt+1。
- 直接使用缓存的 Key/Value 进行注意力计算:At+1=softmax(Qt+1KT≤t√dk)V≤t
- 然后只计算 Kt+1 和 Vt+1,并追加到缓存中。
-
计算复杂度优化:
- 由于只计算新的 K,V,每个 token 的计算变为 O(tdk)。
- 生成 T 个 Token 的总复杂度变为:O(Tdk)
- 从二次复杂度 O(T2dk) 降到线性复杂度 O(Tdk),极大提升推理速度。
4. KV Cache 如何提升推理速度
(1)减少重复计算
-
在 没有 KV Cache 的情况下:
- 每生成一个新 token,都需要重新计算所有 Key/Value。
- 计算量随着 Token 长度增加 呈二次增长,导致长文本生成非常慢。
-
使用 KV Cache 后:
- 只需要计算当前新 token 的 Key/Value,并追加到缓存中。
- 计算量随 Token 长度增加仅为 线性增长,显著加速推理。
(2)减少显存(VRAM)占用
- 标准 Transformer 需要存储完整的 Attention 计算历史,占用大量显存。
- KV Cache 仅存储 Key/Value 矩阵,相比于存储整个计算图,显存占用大幅降低。
(3)加速 GPU 计算
-
减少矩阵运算规模:
- 原始 Transformer 需要计算 O(t2dk) 规模的矩阵乘法。
- KV Cache 只需要 O(tdk),能更好地利用 GPU 并行计算能力。
-
提升批处理(Batch Processing)性能:
- 由于不需要重复计算过去的 Key/Value,GPU 可以更快地处理多个请求,适用于大规模文本生成任务(如 ChatGPT)。
5. KV Cache 在实际推理中的应用
在实际使用中,KV Cache 通过 动态扩展缓存,让模型在生成过程中不断追加 Key/Value,而不是重新计算整个序列。在优化框架中,通常:
- 缓存格式:
- 多头注意力情况下,缓存维度为:(batch_size,num_heads,seq_len,dk)
- 多头注意力情况下,缓存维度为:
- 缓存更新策略:
- 计算 Qt。
- 读取之前存储的 Key/Value。
- 计算 Attention 得到输出。
- 计算新的 Key/Value 并追加到缓存。
示意流程:
- 计算 Query Qt。
- 直接从 KV Cache 读取 K≤t 和 V≤t。
- 计算 Attention 权重 At。
- 计算 Attention 输出 Ot。
- 计算新 Key/Value 并存入 Cache,供下一个 Token 使用。
6. 结论
KV Cache 极大地优化了 GPT 模型的推理速度,其主要贡献包括:
- 减少 Key/Value 计算量(避免重复计算)。
- 降低计算复杂度(从 O(T2dk) 降到 O(Tdk))。
- 减少显存占用(仅存储 Key/Value,不存储完整计算图)。
- 加速 GPU 计算(减少不必要的矩阵计算)。
- 提升大规模推理能力(适用于 ChatGPT、文档生成等场景)。
这个优化方案使得 GPT 模型能够高效地进行长文本生成,是现代大语言模型推理的关键技术之一。
这就是 KV 缓存的核心原理,也是 GPT 等 LLM 运行高效的关键。
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】凌霞软件回馈社区,博客园 & 1Panel & Halo 联合会员上线
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】博客园社区专享云产品让利特惠,阿里云新客6.5折上折
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· DeepSeek “源神”启动!「GitHub 热点速览」
· 我与微信审核的“相爱相杀”看个人小程序副业
· 微软正式发布.NET 10 Preview 1:开启下一代开发框架新篇章
· 如何使用 Uni-app 实现视频聊天(源码,支持安卓、iOS)
· C# 集成 DeepSeek 模型实现 AI 私有化(本地部署与 API 调用教程)