ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

KV Cache原理与显存优化:大模型推理加速的核心技术

KV Cache原理与显存优化:大模型推理加速的核心技术

1. 为什么需要 KV Cache

1.1 自回归生成里的隐形成本

先说一个很多刚接触大模型的朋友都会困惑的点:大模型推理的时候,CPU和GPU到底在忙什么?如果只是调 API,你看到的是一个 prompt 丢进去,几秒后一串 token 吐出来,好像没什么特别的。但实际你每看到"一个字"生成完,模型背后已经跑了完整的一次前向传播。

大模型的生成方式是自回归,也就是把上一次生成的 token 拼到输入序列后面,再预测下一个 token。这个流程听起来很简单,但问题来了:序列是逐次变长的。从第 1 个 token 到第 N 个 token,序列长度从 1 涨到 N。如果每次生成时都把整个序列从头算一遍注意力,那成本的增长速度是平方级的。

举个例子。你输入 1000 个字的 prompt,让模型生成 1000 个字的回答。朴素的做法是:生成第 1 个字时,对 1000 个 token 做注意力计算;生成第 2 个字时,把 1000 个 token 和刚生成的 1 个 token 凑成 1001 个,再整体算一遍注意力。到生成第 500 个字时,你每往前推一步,都要对 1500 个 token 重新做全量计算。这一路算下来,大量计算花在了重复处理旧 token 上,真正对新 token 的注意力计算反而被淹没。

这个浪费有多夸张?如果序列总长度是 L,朴素实现每一步要算的注意力复杂度是 O(L),而整个过程要跑 L 步,加起来是 O(L^2)。但每一步真正新增的信息只有一个 token,理论上每一步只需要 O(1) 的计算量——前提是你把历史信息想办法留住了。这就是 KV Cache 出现的根本原因。

1.2 关键洞察:哪些计算可以复用

要理解 KV Cache,得先看注意力机制里每一部分具体在算什么。

自注意力层的核心公式是:

Attention(Q, K, V) = softmax(QK^T / √d) V

这里 Q、K、V 分别来自输入序列和三个权重矩阵的乘法。对于一个长度为 n 的序列,Q、K、V 都是 n×d 的矩阵(先忽略多头拆分)。

我们在生成第 t 个 token 的时候,Q、K、V 是什么状态?

Q 只跟当前 token 有关。每一步生成新的 token,Q 就是新 token 通过 W_Q 映射得到的那一行。它只需算一次,用完即弃,不存在复不复用的问题。

但 K 和 V 不一样。计算第 t 个 token 的输出时,需要拿 Q_t 和之前所有位置的 K_1, K_2, ..., K_t 做点积,再和所有位置的 V_1, ..., V_t 做加权求和。也就是说,K、V 里那前 t-1 行,在计算后续每一步时都会被反复用到。

最关键的发现就在这里:历史 token 对应的 K 和 V,是只依赖历史 token 本身的。它们不会因为当前生成的是哪个 token 而变化。既然它们不会变,那为什么不把第一轮算好的结果存下来,后面直接取用,非要每次从头再算一遍?

这就是 KV Cache 的本质思路:用显存换算力。把已经算过的 K、V 矩阵缓存住,生成每一步时只需要计算当前 token 的 K、V,然后追加到缓存尾部,再基于完整的缓存做注意力计算。整个过程的计算复杂度从 O(L^2) 降到了 O(L),代价是额外占用显存来存放这些缓存数据。

打个生活化的比方。你做一道很长的数学证明题,如果每推导一步都要把题目条件和前面所有推导过程重新抄一遍,那做完一整道题的时间是没法看的。KV Cache 就是你的草稿纸:前面已经推出来的中间结论,写在草稿纸上,后面直接用,不用重推。

2. KV Cache 到底存了什么

2.1 K、V 矩阵的维度与生命周期

先回答一个新手最容易搞混的问题:KV Cache 里存的到底是啥?是文本吗?是 embedding 吗?

都不是。它存的是 K 和 V 矩阵,而且是每一层 Transformer 层各自一份。

现在的 LLM 基本都是多层 Transformer 结构,比如一个 7B 参数的模型往往有 32 层。每一层都有自己的多头注意力,也就有自己的 W_K、W_V 权重矩阵。输入 token 经过第 1 层算出的 K、V,和第 20 层算出的 K、V 是完全不同的东西。所以 KV Cache 不是一个整体,而是"每一层都有一份",各层互不通用。

具体到单个层来说,K、V 矩阵的维度通常是:

  • batch size:当前处理的样本数
  • num heads:注意力头数
  • seq_len:缓存的 token 数量
  • head_dim:每个注意力头的维度

也就是你有一个 (batch, num_heads, seq_len, head_dim) 的四维张量。K 和 V 各一份,所以一共两个这样的张量。

随着生成过程的推进,每生成一个 token,这个 token 在当前层的 K 和 V 就被算出来,然后追加到缓存矩阵的 seq_len 那一维末尾。也就是说,KV Cache 是动态增长的。如果你生成 1024 个 token,那缓存里最终会有 1024 个位置;如果你在服务端开了一个长对话,缓存里存的就是整段对话历史上所有 token 的 K、V。

这里有个细节:Q 为什么不缓存?上面解释过,因为每一步的 Q 只和当前生成的 token 有关,用完就没了。每一步只需要算当前 Q 对全部缓存的 K、V 的注意力,完全不需要把历史 Q 留住。

2.2 KV Cache 和模型架构参数的关系

不同模型之间的 KV Cache 大小差异很大,原因就在于架构参数不同。

第一是层数。32 层的模型,KV Cache 是 32 份;80 层的模型,80 份。层数越多,缓存总量越大。

第二是注意力头数。这里要区分一个概念:Q 的头数和 K、V 的头数未必相等。传统多头注意力 MHA 里,Q、K、V 头数完全一致,每个 Q 头配一个 K 头和一个 V 头。但后来的模型里大量使用 GQA(分组查询注意力)和 MQA(多查询注意力),让多个 Q 头共享同一组 K、V 头,这会在很大程度上减少 K、V 的存储量。

第三是 head_dim,也就是每个注意力头的向量维度。常见的是 128 或者 64。这个数值比较固定,因为注意力分数计算要对 head_dim 开根号做缩放,改它的空间不大。

这三个参数乘起来,再乘以 2(K 和 V 两份),再乘以序列长度,就是 KV Cache 的体量。

一个直观感受:同样是 4096 上下文长度,传统 MHA 架构模型的 KV Cache 动不动就上 GB,而 GQA 架构的模型可能只要几百 MB。这个差异直接决定了你能不能把模型塞进一张 24GB 的卡里跑长文本。

3. KV Cache 的显存占用到底有多离谱

3.1 一个公式算出你的缓存有多少

显存估算其实很简单,就是算一下缓存矩阵里有多少个元素,再乘以每个元素的字节数。

通用公式:

KV Cache 大小 = 2 × num_layers × batch_size × seq_len × num_kv_heads × head_dim × dtype_size

这个 2 就是 K、V 各一份。dtype_size 如果是 FP16/BF16,就是 2 字节;FP32 是 4 字节;INT8 是 1 字节;INT4 量化后平均 0.5 字节。

拿一个经典 7B 模型来举例。假设它 32 层、32 个 Q 头(这里先按 MHA,也就是 KV 头也是 32)、head_dim 128,FP16 精度,batch size 1,上下文长度 4096:

单个 token 需要缓存一层里所有 KV 头的 KV 对:32 头 × 128 维 × 2(K 和 V)× 2 字节 = 16384 字节。32 层全部加起来:16384 × 32 = 524288 字节,约 0.5 MB。

注意这是每个 token 的开销。4096 个 token 就是 0.5 MB × 4096 = 2048 MB,正好约 2GB。

这个数字和模型权重放一起看就更清楚了。7B 模型 BF16 权重约 14GB,在 24GB 显卡上已经占了大部分显存,再加上 2GB 的 KV Cache 以及激活值、临时缓冲区,跑 4096 上下文基本就是贴着上限了。如果上下文再拉到 8192,KV Cache 直接翻倍到 4GB,一张 4090 可能连模型加载都有风险。

如果你用的是更大的模型,比如 70B 级别,层数更多,KV Cache 的增速更夸张。这也是为什么很多服务端跑长文本时要专门做显存规划,甚至多卡分布缓存。

3.2 不同架构下的显存对照

上面说的 7B MHA 模型只是一个参照,实际工程中很少拿一个纯 MHA 的大模型直接上长上下文。现在主流模型纷纷转向 GQA,KV Cache 能够显著瘦身。

GQA 的意思是,比如 32 个 Q 头分成 4 组,每组共享 1 个 K 头和 1 个 V 头,那么 KV 头数就从 32 降到了 4。KV Cache 直接缩到原来的 1/8。像一些 70B 级别模型使用 GQA 后,KV Cache 的显存压力大幅下降,部署门槛就低了很多。

下面这张表可以看得更清楚(假设都是 batch 1、FP16、seq_len 4096):

模型规模层数KV 头数head_dimKV Cache 大小
7B MHA32321282048 MB
7B GQA(4组)328128512 MB
70B MHA806412810240 MB
70B GQA(8组)8081281280 MB

这个对比非常直观。同样 4096 长度,GQA 架构的 70B 模型 KV Cache 只有 MHA 架构 70B 模型的 1/8。所以在现在的开源模型榜单上,GQA 已经成了标配。

另外要注意:KV Cache 是动态增长的。显存规划一定要按最大上下文长度估算,而不是按当前 prompt 长度。很多朋友部署后跑短文本没问题,一拉长就 OOM,原因就在这。

4. 工程上怎么给它瘦身和调优

4.1 模型侧的减负:GQA/MQA 的取舍

先说说为什么 GQA 能成立。注意力机制里,多个 Q 头本身就是在不同子空间里做查询,但 K、V 承载的信息其实有大量冗余。让一组 Q 头共享同一份 K、V,在理论上会损失一定表达能力,但实际效果远小于参数量的节省。尤其对于推理友好的部署场景,省下的显存可以换来更大的 batch 或者更长的上下文,综合收益很高。

MQA 更进一步,所有 Q 头共享一个 K 头和一个 V 头。这种极限压缩在模型效果上往往有可感知的下降,所以现在主流模型基本都用 GQA,MQA 只在早期一些模型里出现过。

选型建议:如果你要自己训练或者微调一个模型,优先选择带 GQA 的底座模型。如果只是调用现成模型服务,那不用管架构,框架层面已经处理好了。但你要知道服务端为什么有时候限制上下文长度——很大程度上不是模型不聪明,而是 KV Cache 装不下。

4.2 系统侧的优化:量化、PagedAttention 和前缀复用

模型侧的 GQA 是从根源减少 KV Cache 总量,工程侧的优化可以让缓存更高效地跑起来。

第一个思路是 KV Cache 量化。既然缓存里存的是矩阵,那就可以像量化权重一样,把 KV 从 FP16 压到 INT8 甚至 INT4。INT8 量化 KV Cache 通常能让缓存体积减半,而精度损失在绝大多数场景下可接受。INT4 更激进,但需要谨慎验证效果衰减,尤其是在长上下文 RAG 场景里,信息密度低,量化误差可能会被放大。

第二个思路是 PagedAttention。这是 vLLM 等推理框架的核心优化。传统做法给每个请求预分配一块完整显存来存 KV Cache,但生成的实际 token 数不确定,导致内部碎片严重。PagedAttention 把缓存按固定大小的块来管理,不连续的内存块可以用索引串联,显存利用率大幅提升,最重要的是它允许在同一批请求之间共享缓存。比如在 RAG 场景里,多个请求共用同一段检索文档前缀,这时候共用 KV Cache 能让吞吐翻倍。我当时把这个优化接入一个知识库问答服务之后,在同样的显存下并发能力直接提了接近一倍。

第三个思路是前缀复用。如果系统里经常有相同的 prompt 前缀(比如 system prompt 很长、多轮对话里历史消息固定),可以在服务端把这部分对应的 KV Cache 提前算好缓存起来。新请求进来,只要前缀命中,这部分 K、V 就不需要重算。这个优化在小 batch 高并发的场景里效果明显,因为很多对话应用的前缀部分是固定的。

4.3 上下文长度限制的真实原因

很多人在用大模型时会问:为什么模型有上下文窗口上限?比如 4096、8192、128K。是不是模型权重里固定了序列长度?

不完全是。严格来说,注意力层本身对长度没有天然限制,你给它多长的序列它都能算。真正的限制来自两个地方:一是训练时的位置编码分布。模型在预训练时见过的序列长度是有限的,超过这个范围,位置编码的外推能力再强也会有衰减。二是 KV Cache 显存。就算位置编码能外推到更长,服务端也得有显存把那么长的 KV Cache 装下来。

这两个限制叠加在一起,就构成了实际的上下文上限。所以当你看到某个模型宣称支持 128K 上下文时,除了看位置编码方案,还要留意它对 KV Cache 做了什么优化。如果它没有量化和显存调度优化,可能只是模型结构上能支持,真部署起来单用户就要吃掉几十 GB 显存,性价比很低。

4.4 解码策略与缓存的生命周期管理

这一节容易被忽略。beam search、温度采样、多轮对话这些看起来和 KV Cache 关系不大,实际上影响很大。

beam search 中,每个 beam 有自己的生成轨迹,也就是每个 beam 需要一套独立的 KV Cache。如果 beam size 为 4,那么 KV Cache 就翻 4 倍。许多框架优化的目标之一就是最大化各个 beam 之间的 KV 共享,尽量把相同前缀的部分合并计算。

温度采样和 top-p 则不影响 KV Cache 的大小,但会影响缓存的生命周期。因为采样引入了随机性,你不能预判最终会生成哪些 token,只能在每步生成后动态追加,然后把不被采样的分支的缓存丢弃。如果要做多个候选结果重新排序,那每个候选的 KV Cache 都要保留到最终选完才释放。

多轮对话场景也要特别注意:上一轮的 KV Cache 要不要带进下一轮?理论上应该带,因为对话历史是连续的,复用上一轮的缓存可以省掉重新计算整个聊天记录的代价。但显存有限,对话轮数多了之后缓存会持续累积。这时候要么做截断,只保留最近 N 轮历史;要么做摘要,把旧对话压缩成摘要文本参与后续计算。

我自己在做一个长对话应用时踩过这个坑:开了缓存持续累积,结果跑到第 20 轮对话时显存直接被打满。后来改成"权重高的历史消息保留 KV Cache,更早的对话压成摘要",问题就解决了。这个思路在很多生产环境里都是标配。

5. 常见问题与排查实录

5.1 显存爆掉,先查谁?

推理服务 OOM 的时候,很多人的第一反应是"模型太大了",然后去换小模型。但我建议先查 KV Cache 的占用。

用一个简单的公式估算:模型权重占用的显存是固定的,你真正要关注的是权重之外的部分。常见排查步骤如下:

  1. 先看模型权重多大:参数量 × 精度字节数。
  2. 再算激活值:前向传播过程中每层产生的中间矩阵,这部分往往被很多人忽略。
  3. 最后算 KV Cache:按上面给的公式,代入实际上下文长度。
  4. 对比显卡总显存,看哪个占比最高。

通常在短文本 batch 大的场景里,占显存大头的是激活值;在长文本 batch 小的场景里,KV Cache 是大头。OOM 的时候如果发现 KV Cache 占比很高,优先考虑调整上下文长度、做 KV Cache 量化或者开 PagedAttention,而不是急着换模型。

5.2 为什么生成阶段比预填充阶段快?

这是刚从业者最容易困惑的问题。同样是模型前向传播,为什么输入 prompt 之后输出 token 的速度远高于处理 prompt 的速度?

答案是 KV Cache 起了关键作用。

处理 prompt 的阶段通常叫预填充阶段。这时没有任何缓存,模型需要一次性把整个 prompt 的 K、V 算出来,并完成初始的注意力计算。因为此时没有历史缓存可用,计算量就是跟 prompt 长度成比例的,prompt 越长,这个阶段越慢。

而生成阶段每一轮只需要算新 token 的 Q、K、V,再拿这个 Q 去查 KV Cache 里完整的 K、V 做注意力。单步的计算量不随历史长度线性增长,所以看起来每吐一个 token 都很快。

这也解释了为什么有的 API 服务在输入很长文档时响应慢,而输出却很快。实际就是在预填充阶段要把整篇文档的 K、V 都算一遍、缓存下来。优化预填充阶段的关键在于并行度:prompt 里的 token 之间没有依赖关系,可以一次性并行处理,所以很多框架专门优化了这个阶段的矩阵运算调度。

5.3 一个容易被忽略的缓存正确性问题

如果你是自己手写推理代码,而不是用现成框架,最容易踩的坑是缓存中混入了来自不同前缀的 K、V。

举个例子。服务端的 KV Cache 按 token 序列位置存放。如果两个请求的前缀都命中了同一段 prompt,而另一个请求的前缀只有一部分相同,那它们的缓存尾部不能直接共享。实现时必须为每个请求维护独立的缓存索引,否则一个请求的生成结果会静默地受到另一个请求历史的影响。

这个问题在引入共享前缀优化后尤其容易发生。我之前遇到过一次:先对 A 请求生成了一半,然后 B 请求以相同的前缀进来复用了 A 的缓存,但 A 生成出来的 token 并不适合作为 B 的上下文,结果 B 的生成质量明显变差。后来在框架层面对缓存 key 做了严格管理,把"前缀扩展后重新计算"和"前缀命中后直接复用"两类情况明确区分开,才彻底解决。

5.4 长上下文下 KV Cache 的精度衰减

还有一个细节值得单独说:KV Cache 量化之后,长上下文的精度衰减会更明显。短文本下 INT8 量化的效果衰减可能感觉不出来,但序列拉到 32K 甚至更长,注意力分数经过多次 softmax 和加权后,量化误差会一层层累积,最终表现为生成质量下降、复读、甚至逻辑混乱。

如果业务场景就是长文档问答、长代码分析这类对信息保真度要求高的任务,我建议先做一轮对照组测试:同样的 prompt,分别用 FP16 和 INT8 的 KV Cache 跑,人工对比关键事实的召回率。大多数情况下 INT8 是安全的,但不要默认安全。至于 INT4,除非你的框架对压缩后的 KV 有特殊补偿机制,否则长文本场景要非常谨慎。

6. 实操中的几点建议

如果这篇文章能帮你建立起一个核心认知,我希望是这个:大模型推理速度优化的核心,很大程度上就是围绕 KV Cache 做文章的。它决定了你能跑多长的上下文、能撑起多大的并发、甚至决定了某个模型适不适合在特定硬件上部署。

几个可以直接落地的小建议。

选型阶段:优先考虑使用 GQA 架构的模型。在显存有限的情况下,GQA 带来的收益比单纯增加参数量更实在。部署阶段:不要等 OOM 了才去优化,先按公式估算一下你的 KV Cache 峰值,再决定要不要量化、要不要开 PagedAttention。上线阶段:监控里一定要有 KV Cache 的实时占用指标,这个指标比单纯的显存使用率更能反映长文本场景的风险。

我最初从 7B MHA 模型切到 GQA 架构模型时,最直观的感受是:同样的 4090,以前 4096 长度就贴着显存上限跑,现在可以轻松跑到 8192,而且多路并发的余量也出来了。后来再加上 KV Cache 量化和前缀复用,吞吐又上了一个台阶。这些都不是什么玄学技巧,就是把上面的原理一条条落到工程里。

最后提一个容易被忽略的小点:KV Cache 的复用和 release 时机也很重要。在多轮对话里,如果用户新输入的内容立刻触发了历史缓存的截断(比如超过最大长度),你会明显感觉到首 token 变慢。有些框架支持在截断后重新构建前缀缓存,但这需要额外的时间和显存交换。设计服务端策略时,最好把"多少轮后做摘要压缩"作为可配置参数而不是拍脑袋写死。这个度需要根据你的显存预算和业务场景反复调,我一般建议初始值设置为能覆盖用户 80% 对话历史的最大轮数,不要盲目追求缓存全量覆盖,因为那往往意味着你需要为极少数超长会话预留大量显存,整体不划算。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进