★ FEATURED ARTICLE 做推理优化的人迟早都会撞上KV Cache这个词。大模型LLM生成文本时是逐字挤出来的每生成一个新token都要重新看一眼前面所有已经生成的token重新算一遍注意力。KV Cache就是在这个过程里把已经算好的Key和Value缓存下来避免每生成一个字都把整段历史的注意力从头算一遍。没有它一个7B量级的模型推理速度会慢到让人崩溃有了它每次生成只需要做一次单token级别的前向计算。这篇文章我想从头到尾拆一遍KV Cache的核心原理、显存账本、长上下文里常见的lost in middle现象以及Prompt Cache到底缓存了什么最后还会聊聊我在工程里实测过的优化手段和踩过的坑。适合正在做推理部署、服务优化或者想搞懂注意力机制背后成本的同学。1. 从自回归生成说起KV Cache到底在解决什么问题1.1 Decoder-only模型的排队吐字机制现在主流的大模型基本都是Decoder-only架构生成方式叫自回归Autoregressive。意思是模型每次都只预测下一个token然后把这个token拼到输入末尾再预测下一个。你可以把它想象成一个人在排队报数报完一个数字记住了前面所有数字才能报下一个。这个机制决定了推理天然是串行的第一个token必须等prompt全部处理完后面的每个token都依赖前一个的输出。很多刚接触大模型的人会奇怪为什么GPU那么强生成速度还是每秒只有几十个token原因就在这里——不是GPU算不动而是每一步都要被前一步的结果卡住。如果完全没有KV Cache第N步生成时模型需要把前N-1个token的Key和Value全部重新计算一遍。这就相当于一个学生每天复习时都把整本课本重新抄一遍而不是只看笔记效率低到离谱。1.2 朴素实现的计算量与重复劳动我们看一个具体的例子。假设用户输入了1000个token的prompt模型要生成500个token。用最朴素的实现完全不缓存任何东西第1次生成时需要处理完整的1000个token计算量来自这1000个token的注意力。第2次生成时输入变成1001个token又要重新计算这1001个token的注意力。第3次生成时输入变成1002个token继续从头算。以此类推第500次生成时模型要处理1499个token。把整个过程加起来实际算过的token量大约是[ 1000 1001 1002 ... 1499 ]等差数列求和总数大约375000个token的处理量。但仔细想一下真正需要从头算的只有第一个新token吗并不是。其实前1000个token的Key和Value在第2次、第3次、第4次计算时计算结果完全是一样的因为这些token根本没有变化。唯一变化的是新追加的那个token。也就是说朴素实现里绝大部分计算都是同一个东西的重复计算白烧了大量GPU算力。KV Cache的思路简单粗暴把这些计算结果存下来下次直接用。1.3 KV Cache的本质用空间换时间的记忆KV Cache本质上是一种典型的空间换时间策略。它把已经计算过的历史token的Key矩阵和Value矩阵存在显存里在后续的自回归过程中直接拼接使用省掉了海量的重复矩阵乘法。这里有一个很微妙的地方每次生成新token时我们其实只需要计算当前这个新token的Query、Key、Value。当前token的Query需要去跟所有历史token的Key做点积得到注意力分数然后再拿这些分数去加权所有历史token的Value。如果我们把历史Key和Value都缓存在显存里这一步就只需要加载缓存、拼接新token的KV然后做一次完整的Attention计算即可。这个记忆的代价是显存而且会随着序列长度增加和并发请求数增加而线性膨胀。后面我专门开一节讲显存账本因为这是KV Cache在实际部署中最大的坎。2. 为什么缓存的是KV而不是QKV2.1 Attention的四步流程很多人第一次接触KV Cache时会有一个很自然的疑问为什么不把Q也一起缓存了毕竟Q也是算出来的缓存起来不是更省事吗要回答这个问题得先看清Attention的计算流程。以标准的缩放点积注意力为例一个token的参与过程是通过三个权重矩阵 ( W_Q, W_K, W_V ) 分别算出当前token的Query、Key、Value。用当前token的Query去跟序列里所有token的Key做点积再缩放、softmax得到注意力权重。用注意力权重去加权所有token的Value得到当前的输出向量。这个流程里每个token的Key和Value代表的是一种被查询时的属性Key是内容索引Value是内容本身。一旦确定下来在整段序列生成完毕之前都不会变。而Query代表的是查询者的意图它的使命就是去寻找当前正在处理的位置应该关注哪些历史内容。每一步生成的token不同它的Query就要重新计算而且只服务于当前这一步。旧token的Query在生成它自己的时候已经把该做的事情做完了之后再也没有任何用处。2.2 Q的作用域只活在当下我们用一个生活化的类比来理解。想象你在一间资料室查资料桌上摊着一堆档案每份档案旁边贴着两个标签一个是检索标签Key说明档案涉及什么主题一个是正文内容Value档案的具体文字。你每次来查资料都会带着一个查询请求这个请求就是Query。你要做的是拿着查询请求扫一遍所有档案的检索标签挑出最匹配的几份然后打开正文阅读摘录。在这个场景里档案馆的检索标签和正文内容是不会变的变的是你每次带来的查询请求。如果你明天又来查一次没有任何理由把今天的查询请求也存着因为明天的请求大概率不同而且今天这份请求在今天的查询里已经用完使命了。注意力机制一模一样。Key就像档案的检索标签Value就像档案正文Query就是每一次新的查询请求。所以在推理时缓存KV完全说得通而缓存Q反而没有意义因为旧的Q在生成旧token时已经用过了后续步骤根本不会再用到它。这就是为什么需要KV Cache而不是QKV Cache的根本答案。2.3 QKV三者视角的直觉类比再换个角度补充一下。在标准的自注意力里一个token要扮演两个角色作为被关注者它要被后续token的Query所查询这时需要提供自己的Key和Value。作为关注者它要去查询前面的所有token这时需要使用自己的Query。假设序列里已经生成了1000个历史token现在要生成第1001个token。第1001个token作为关注者必须用自己的Query去和前1000个token的Key做匹配。而前1000个token作为被关注者它们的Key和Value早就确定且不会再变了。缓存Q意味着把前1000个token的查询请求都留着。可它们的查询请求已经在它们各自被生成的时候执行完了后续根本不存在它们再次作为关注者去查询别人的场景。所以工程上从没人做QKV Cache最多只会做KV Cache。3. 显存占用公式KV Cache如何吃掉你的GPU3.1 公式推导KV Cache最大的坑不在原理而在显存。它的占用是可以精确算出来的我先把公式列出来。对于每个token每一层注意力都要缓存一组Key和Value每组的维度是[ num_kv_heads \times head_dim ]其中 ( num_kv_heads ) 是KV头的数量( head_dim ) 是每个注意力头的维度。注意这里不是用Query头的数量而是用KV头的数量。传统多头注意力MHA里KV头数等于Query头数但像LLaMA 2-70B这类模型用了GQAKV头数会少很多这个我们后面细说。每个token、每一层需要存储的字节数大约是[ 2 \times num_kv_heads \times head_dim \times 2 \text{ bytes} ]乘以2是因为Key和Value各一份再乘以2是因为一般用FP16存储每个数占2字节。整个模型有 ( L ) 层序列长度是 ( S )并发请求数是 ( B )那么KV Cache总显存大约是[ 4 \times L \times num_kv_heads \times head_dim \times S \times B ]注意这里已经把Key和Value两份、FP16两个字节都折算进去了。如果要换成GB再除以 ( 1024^3 )。3.2 不同规模模型的量化对比我用几个常见的模型配置来算一下大家感受会更直观。先看LLaMA-7B的经典配置32层32个Query头每个头维度128也就是32个KV头Head Dim 128FP16存储。当batch size为1、序列长度为2048时[ 4 \times 32 \times 32 \times 128 \times 2048 \times 1 4 \times 32 \times 32 \times 128 \times 2048 1,073,741,824 \text{ bytes} ]也就是约1GB。你没看错一个7B参数的模型光KV Cache就占了1GB显存这还没算模型权重本身和中间激活值。再看LLaMA-13B的典型配置40层40个Query头KV头数也是40Head Dim 128。同样batch1、序列2048[ 4 \times 40 \times 40 \times 128 \times 2048 \approx 1.67 \text{ GB} ]注意13B参数只比7B多了一倍不到KV Cache也差不多是1.67倍基本跟层数和头数线性挂钩。最夸张的是长上下文和并发场景。如果序列长度从2048拉到3276832KKV Cache直接涨16倍。如果同时有8个并发请求再乘以8。一个7B模型光KV Cache就能冲到128GB以上这已经超过很多单卡的全部显存了。下面这个表可以帮你快速估算模型规模层数KV头数Head Dim2048长度单请求KV Cache8192长度单请求KV Cache7B (MHA)3232128约1GB约4GB13B (MHA)4040128约1.67GB约6.7GB70B (GQA)808128约0.67GB约2.7GB注意70B因为用了GQAKV头数大幅缩减KV Cache反而比13B的全量MHA更低。这个对比已经能说明很多事了。3.3 与上下文长度的关系长度是放大器从公式看KV Cache跟序列长度是明确的线性正比关系。所以现在各家都在卷上下文长度8K、32K、128K、1M每提升一档KV Cache的显存压力也跟着提升一档。很多人只盯着模型支不支持长上下文没有意识到长上下文的成本不是模型参数而是KV Cache。举个实际案例用7B模型、batch1、FP16部署显存预算假设12GB。模型权重大约占14GB7B × 2字节其实单卡都装不下通常要量化到INT4后权重约3.5GB才勉强腾出空间。序列长度如果跑到32KKV Cache就要16GB直接把显存爆掉。这也是为什么纯靠加显存做长上下文不现实必须靠工程手段去压KV Cache的体积。后面第5章我会展开讲几个真正有效的优化方向。4. 长上下文的两大现象Lost in Middle 与 Prompt Cache4.1 Lost in the Middle到底是什么长上下文场景里有一个很有名的现象叫Lost in the Middle中文经常翻译成中间迷失。意思是模型在处理长文本时对开头的部分和结尾的部分利用得比较好但对夹在中间的内容经常视而不见导致检索或者问答任务里相关信息如果落在长文本中段模型反而不容易答对。这个现象最早在2023年的论文《Lost in the Middle: How Language Models Use Long Contexts》里被系统性地描述出来。论文在做开放域问答和多项选择任务时发现无论上下文里放了多少篇文档模型答对问题的概率都跟正确答案放在哪个位置强相关放在开头或者结尾准确率高放在中间准确率显著下降。你可能会想这不就是人的注意力也有首因效应和近因效应吗模型确实也有一点类似倾向。但更深层的原因包括位置编码的远程衰减效应。目前很多模型用旋转位置编码RoPE一类的位置编码某些设计天然让远处的token在注意力分数上容易被拉低中间位置有时反而处于尴尬区。长上下文里注意力被稀释。序列越长softmax越容易变得平滑大量token分散注意力分数中间token更难以脱颖而出。训练数据里信息密度不均衡。如果训练时很少有正确答案藏在长文中间的样本模型天然不会特别擅长这种模式。4.2 Prompt Cache / Prefix Caching 又是什么这个话题下经常有另一个高频问题Prompt Cache是什么跟KV Cache有什么关系我在网上看到很多人把这俩混为一谈其实一个是底层机制一个是工程复用策略。KV Cache是个通用机制只要你用自回归模型推理每一步生成的KV都会缓存。Prefix Caching也叫Prompt Cache则是更上层的优化策略它把多个请求之间重复的那部分前缀在KV Cache层面做复用。举个例子。你架设的服务里system prompt是固定的一大段系统设定然后每个用户会在这段设定后面追加自己的问题。没有Prefix Caching时每次请求进来即使system prompt完全一样模型也会重新为这段system prompt算一遍KV。有了Prefix Caching后第一请求处理完系统会把system prompt对应的KV Cache按前缀做缓存第二个请求进来如果前缀完全一致直接复用缓存只需要计算用户新追加的那一小段内容的KV。注意这里有一个容易误导的点表面的Prompt Cache名称会让人觉得是在缓存提示词字符串但其实底层缓存的载体仍然是KV Cache只是按照前缀去索引和复用。现在vLLM、SGLang这类主流推理框架都有类似机制有些框架直接叫Automatic Prefix Caching原理是一样的。4.3 两者之间的关联Lost in the Middle和Prefix Caching看起来是两个不相关的话题但是放在长上下文工程化里它们其实是两条交织的主线。一方面上下文越长KV Cache越大服务商越需要Prefix Caching来降低重复计算的成本否则多轮对话和共享system prompt的高并发场景根本跑不动。另一方面上下文越长Lost in the Middle问题越明显用户如果通过RAG往文档中间塞了一堆相关资料模型可能压根没看见这时你再怎么优化KV Cache的显存也不代表答案质量会变好。显存是成本问题Lost in the Middle是效果问题两个都要解决。工程上常见的做法是文档重排把最相关的片段放在上下文开头和结尾或者至少避开正中间的位置。有人觉得这有点玄学但我实测下来确实有效尤其是10段以上文档同时塞入的场景重排带来的准确率提升经常有5到10个点。你可以在自己的RAG管线里直接试一下这个思路。5. 工程优化实战把KV Cache按进显存的几招5.1 架构层面GQA和MQA先说说从模型架构层面解决KV Cache膨胀的思路。前面提到过KV头数直接影响KV Cache大小。传统多头注意力MHA里每个Query头配一组独立的KV头KV头数很多缓存自然大。MQAMulti-Query Attention的思路是所有Query头共享同一组Key和ValueKV头数直接降到1。这样KV Cache能缩小几倍到几十倍但质量会有一定损失因为所有Query头都用一套KV表达能力受限。GQAGrouped Query Attention分组查询注意力则是MQA和MHA的折中把Query头分成若干组每组共享一套KV头。比如LLaMA 2-70B用8个KV头但Query头有64个相当于每8个Query头共享1个KV头。KV Cache相比全量MHA直接缩小8倍。从效果上看GQA的损失比MQA小得多这也是为什么LLaMA 2中大规模模型都选择了GQA。你在微调或者选型时如果想压推理显存优先选用了GQA的底座模型比什么优化技巧都省心。5.2 存储层面KV Cache量化第二个方向是降低单个KV的存储字节数。常规FP16是2字节如果量化到INT8单个数值降到1字节KV Cache直接减半再激进一点压到INT4甚至FP8还能再缩。KV Cache量化的难点在于精度控制。Key和Value的数值分布跟权重不太一样直接简单粗暴地量化会导致注意力分数漂移生成质量下降。现在主流做法通常是对每个head或者每个channel单独统计缩放因子有些方案还会按token动态校准。实测下来INT8量化的KV Cache在大多数任务上质量损失很小但显存收益非常明显——8K上下文、16并发这种场景能救回好几个GB。我的建议是优先启用INT8 KV Cache如果任务对精度不太敏感比如代码生成、格式化输出可以再试FP8或INT4。但如果任务本身非常吃细节比如长文档摘要、法律文本问答量化级别先保守一点。5.3 管理层面PagedAttention与高效内存调度第三个方向是显存管理。KV Cache是动态增长的每个请求的长度你看不到提前分配一整块固定显存要么不够用要么浪费严重。这块最出名的方案是PagedAttention也就是vLLM的核心技术。PagedAttention的思路非常像操作系统里的虚拟内存分页。传统做法给每个请求预留一大块连续显存预留少了放不下长序列预留多了其他请求又用不到白白浪费。PagedAttention把KV Cache切成固定大小的块按需分配逻辑上连续物理上可以散落在显存各处。这样就能实现接近用多少占多少的效果显存利用率大幅提升。我自己的经验是换了vLLM这类带分页管理的框架后同样一张卡能撑住的并发数明显增加。所以做在线推理服务尽量别自己手写KV Cache管理直接用成熟框架踩坑成本低很多。5.4 稀疏化与滑动窗口第四个方向是少存一点。既然长上下文里很多历史信息其实用不到那干脆只保留最近的一部分KV。滑动窗口注意力是其中一种思路只维护最近W个token的KV超过窗口的历史信息直接丢掉。这样做Prompt很长时KV Cache的上限被锁死在窗口大小不会随着序列无限增长。代价是模型对窗口之外的早期信息失忆如果任务需要引用很早期的内容效果会打折。另一种思路是流式处理StreamingLLM它发现即使是滑动窗口丢掉最开始的几个token的KV也会导致性能崩盘原因在于注意力softmax需要一些稳定的锚点。方案是把序列最初几个token的KV单独保留窗口只滑动中间和末尾的部分这样既控制显存又能保住关键位置。如果你的业务场景是持续的对话流而不是一次性超长文档滑动窗口是性价比很高的选择。但如果用户喜欢把超长上下文整段丢进来做全局分析那滑动窗口就不是最优解得回到GQA和量化这套组合。5.5 服务层批次调度与Prefix Caching最后再补一个服务层的优化就是上一节说的Prefix Caching。在推理服务里多轮对话和同类system prompt的重复率非常高。打开Prefix Caching之后相同前缀的KV被复用prefill阶段的计算量能缩掉一大截首token延迟也会明显下降。这里有一个细节Prefix Caching通常要求前缀完全一致哪怕差一个token前面所有缓存都不能完整复用有些实现支持部分匹配。所以设计system prompt时最好保持严格稳定不要在固定前缀里插入随机时间戳之类的动态内容否则会把缓存命中率拉得很低。我见过有人把user id直接拼在system prompt开头缓存基本全部失效这纯属给自己增加成本。6. 常见问题与排查技巧实录6.1 问题速查表我在实际部署和调优中遇到过不少问题整理成下面这张速查表方便你直接对照现象常见原因排查方向推理时报显存OOMKV Cache增长超出显存预算降低max_seq_len、减少并发、启用KV Cache量化并发稍高就崩每个请求独立KV Cache累计占用过高换PagedAttention、开Prefix Caching、限制最大并发首token延迟高Prefill阶段要计算完整prompt的KV开Prefix Caching、精简system prompt长度长上下文生成速度反而更慢KV Cache太大访存带宽成瓶颈减少KV头数GQA、量化KV、考虑稀疏注意力回答总忽略中间文档内容Lost in the MiddleRAG结果重排、把关键信息放到开头或结尾Prefix Caching命中率低前缀里有动态内容固定前缀、去掉时间戳/user id等动态字段KV Cache量化后质量下降量化的缩放因子校准不充分按head/channel单独校准、改用更高精度量化6.2 实操心得与避坑指南第一不要一上来就调KV Cache。先看权重占了多少、激活值占了多少把tensorboard或者nvidia-smi显存监控打开观察推理过程中显存的增长曲线。如果显存增长明显跟序列长度线性相关基本就是KV Cache在吃显存这时候再针对性优化。第二动手优化前先确认模型是否用了GQA。如果你的底座模型本身就GQAKV Cache压力已经小了很多如果是老式MHA模型优先考虑换底座或微调时过渡到GQA结构而不是在推理层疯狂塞量化。第三测性能和显存要分开测。测显存时用最大序列长度的单请求来压上限测吞吐时一般用小并发中等长度模拟线上真实流量。很多人只测了并发结果上线后遇到超长输入直接OOM其实早点拿个超长序列压一下就知道了。第四长上下文场景下把KV Cache优化和RAG重排序结合起来做。显存是否够看的是KV Cache优化到不到位答案质量看的是Lost in the Middle有没有被规避。只有两头都顾到长上下文才会真正落地可用。第五监控工具一定要有KV Cache指标。vLLM里有专门的metrics能看到每个请求的KV Cache使用量别只看GPU利用率。我遇到过GPU利用率看起来跑满了实际上大部分时间都在重复计算前缀KV的情况开了前缀缓存后利用率反而下降因为不用重复prefill了但整体吞吐和延迟都变好了。最后再分享一个小技巧调试KV Cache相关问题时写一个不带任何缓存的基线推理脚本用来和带缓存的版本做对比。很多时候你以为的显存问题、速度问题其实是别的原因比如位置编码实现有bug、attention mask传错有基线脚本在定位起来快很多。我自己每次搭新框架第一件事就是跑一遍同序列长度、同batch的朴素版和缓存版对照这套路子帮我省掉了大量排查时间。
阅读完成 · 觉得有帮助？
