资讯动态

KV Cache详解:大模型推理加速与显存优化的核心机制

发布时间:2026/10/26 8:38:40 来源:尧图企业网站定制
刚接触大模型推理的时候多半会被一个现象搞懵模型明明只生成一个新字显卡却要把整段历史都“重新读一遍”结果就是越往后生成越慢、显存越吃越凶。等你自己动手写生成循环或者去翻推理框架的源码会遇到一个高频词——KV Cache。它几乎是所有主流大模型推理加速方案的核心底座也是理解大模型生成、上下文窗口、显存占用这些概念绕不开的一环。这篇文章我想把 KV Cache 讲透它为什么必须存在、缓存里到底放的什么东西、显存是怎么被吃掉的、工程上怎么落地、以及在真实跑模型时会遇到哪些坑。适合刚看完 Transformer 结构、想搞懂推理优化的同学也适合已经在用 vLLM 这类框架、想弄明白背后机制的人。1. 为什么没有 KV Cache 的大模型生成慢到劝退从自回归解码说起1.1 大模型生成 token 的本质是逐字造句要说清楚 KV Cache得先回到大模型生成任务最基本的设定自回归解码。现在的 LLM 本质上是一个“给定前面所有 token预测下一个 token”的函数。你给它“我喜欢吃”它算出来下一个字最可能是“苹”然后你把“苹”接在后面喂回去再预测下一个字是“果”如此往复一个词一个词地把内容“挤”出来。这个过程很像一个特别健忘的抄写员每次只允许他写一个词但他每次动笔之前都要把已经写好的整篇文章从头到尾读一遍才肯落笔。原文5个字的时候还好写到2000字的时候每写一个字都要重读2000字效率自然崩了。Transformer 的自回归解码就是这样工作的——每生成一个新 token注意力层需要让最新的 token 和之前所有 token 进行交互。这种设计带来了一个最重要的自然后果下一个 token 的预测完全依赖前面的全部历史信息。所以模型每走一步都必须拿到“截至当前时刻”的完整上下文。问题只在于这个“拿到”的方式到底是重新算一遍还是从某个地方直接读取。1.2 没有缓存时每一步都在重复算一笔冤枉账没有 KV Cache 的时代或者说你在写代码时没注意、直接做朴素实现的时候生成第 N 个 token 的流程是这样把从 1 到 N 的全部 token 拼成完整序列一次性喂给模型跑完整的前向传播拿到第 N 步的预测结果。这样做在逻辑上没错但有个致命浪费模型在计算第 N-1 个 token 时其实已经把序列 1 到 N-1 的注意力结果都算过一遍了。算第 N 个 token 的时候序列 1 到 N-1 的部分不会因为你多了一个新 token 就发生任何改变。它们对应的 Key 向量和 Value 向量在数学上是完全确定的——可是朴素实现里这些向量被扔进了“计算垃圾桶”下一次从头再算一遍。我们拿一个 7B 规模的模型直观感受一下。假设隐藏层维度为 4096共 32 层。没有缓存时上下文长度为 N每一层注意力都要计算一个 N×N 的注意力分数矩阵这还不包括后面更大的加权求和矩阵乘。每次前向传播的复杂度是 O(N²) 量级的而自回归生成要调用 N 次前向传播所以朴素实现的整体复杂度是恐怖的 O(N³)。N32 的时候你完全没感觉N1024 的时候已经有点卡了N4096 时每生成一个 token 都在做几百万次不必要的矩阵运算。我最早自己写过一个“纯新手版”生成函数跑一个 1.5B 模型生成 500 个 token 等了快十分钟就是因为在每一步把整个前文重新过了一遍模型。当时我还以为是模型太大、机器太慢后来才意识到模型确实不大是算法结构里藏着三倍的浪费。2. KV Cache 到底缓存在哪、缓存了什么原理与显存账本2.1 原理简单到令人发指把历史 K、V 留在内存里新 token 只算增量KV Cache 的思路朴素得跟记账差不多把每层注意力算出来的 Key 矩阵和 Value 矩阵缓存在显存里等生成下一个 token 时直接拿这个 token 对应的 Query 去跟缓存的 Key 做点积再对缓存的 Value 做加权求和。举个例子假设当前序列是 [A, B, C]模型正在生成 D。没有缓存时你要把 [A, B, C, D] 整个序列再算一遍注意力其中 A、B、C 的 Key 和 Value 又算了一次。有缓存时你在生成 C 的那一步已经把 Key_A、Value_A、Key_B、Value_B、Key_C、Value_C 都存好了。现在生成 D你只需要用 D 的输入向量算 Query_D、Key_D、Value_D用 Query_D 和缓存的 Key_{A,B,C} 算注意力分数用注意力分数对缓存的 Value_{A,B,C} 加权求和然后顺手把 Key_D、Value_D 追加到缓存里供下一个 token 使用。这套流程里历史部分的矩阵运算被完全省略了。每一步的计算量从“跟整个序列长度成正比”变成“跟一个 token 的处理成本成正比”这就是 KV Cache 能大幅提速的根本原因。这里有个值得注意的细节缓存的是 Key 和 Value而不是 Query。原因也很直接——每一个 token 在自回归生成中只会作为“当前需要预测的 token”出现一次。当 D 被生成后它马上变成了历史序列的一部分不再需要以“当前 token”的身份再次参与计算。所以 Query 用完即弃Key 和 Value 却要一直留在缓存里提供给后续所有新的 token 使用。2.2 显存占用的数学账一个公式算出缓存大小KV Cache 不是免费的午餐它用显存换速度。缓存大小有一个非常清晰的公式KV Cache 字节数 2Key 和 Value 两份 × 层数 L × 序列长度 S × 隐藏维度 d_model × 每元素字节数为什么是 2因为每个 token 在每层都要存一份 Key、一份 Value。FP16 精度下一个元素占 2 字节。我拿几个常见配置算过一笔账直接看数字最直观模型配置d_model × 层数序列长度KV Cache 显存占用7B 级别4096 × 322048约 1 GB7B 级别4096 × 328192约 4 GB7B 级别4096 × 3232768约 16 GB13B 级别5120 × 404096约 3.3 GB70B 级别8192 × 804096约 10.7 GB70B 级别8192 × 8065536约 171 GB看见没有序列长度和维度都是线性增长但实际翻起来非常吓人。一个 7B 模型权重在 FP16 下大概是 14 GB模型的 KV Cache 在 32K 上下文下反而比权重还要大。这也是为什么长上下文推理如此吃显存、为什么很多推理框架要在 KV Cache 上做那么多文章的核心原因。我见过很多同学的困惑不是模型权重占显存最多吗实际上在批量推理和长上下文的场景下KV Cache 经常占据总显存的一半甚至三分之二权重反而成了“小头”。规划显存预算时如果只盯着参数量大概率会在实际运行时 OOM。2.3 为什么这些缓存值不会算错因果掩码与单向注意力的红利有人可能会疑惑Transformer 每个 token 不是会看到序列里的所有 token 吗缓存历史 K、V再把新的 Query 跟它们做注意力会不会跟完整计算不一致不会。关键在于生成场景用的是因果注意力Causal Attention也就是带掩码的注意力第 i 个 token 只能看到它自己和它之前的 token看不到后面的。这个掩码在模型训练时就决定了文本生成时天然满足这个规则。当前要生成的 token 处于序列末尾它本来就应该看到所有历史 token——这正好就是缓存里存着的全部 Key 和 Value。所以无论你在生成第几个 token只要有缓存结果都和完整重计算严格一致是数学上等价的优化不是近似。之所以强调这一点是因为很多做工程的人会有一种“缓存会不会带来误差”的顾虑。我可以确定地告诉你不会。KV Cache 不是量化不是剪枝它没有任何精度损失纯粹是把你算过的结果留下来复用。唯一的代价是显存换速度以及代码里多出来的一堆 cache 管理逻辑。另外KV Cache 的好处不止体现在推理速度上也体现在理解上——你一旦把“每层存 K、V按序列位置追加”这个模型在脑子里建立起来后面看 GQA分组查询注意力、MLA多头潜在注意力、PagedAttention 这些优化时都会轻松很多它们都是在“存什么、怎么存、怎么省着存”这几个维度上做文章的。3. 动手实现 KV Cache从接口设计到显存优化三板斧3.1 最朴素的实现思路一张图看懂生成循环先看代码层面的最小实现思路。你可以把带 KV Cache 的生成循环理解成这样一个分批过程# 伪代码用于说明 KV Cache 的核心逻辑 past_key_values None # 初始时缓存为空 for step in range(max_new_tokens): # 输入当前要预测的 token第一步是整个 prompt current_input prompt if step 0 else last_token_embedding # 前向传播模型除了返回 logits还返回更新后的缓存 logits, past_key_values model( current_input, past_key_valuespast_key_values # 每层只会对新增 token 做完整计算 ) # 取最后一个位置的 logits挑一个 token next_token sample(logits[:, -1, :]) last_token_embedding embedding(next_token)这里是理解整个机制的关键第一次调用时模型要把整个 prompt 完整过一遍这个阶段叫做prefill预填充。prefill 阶段一次性把 prompt 里所有 token 的 Key 和 Value 都算好并存入缓存。之后每走一步模型只输入一个 token叫做decode解码阶段模型只对这个新 token 做增量计算再更新缓存。有一个非常典型的误区是以为 KV Cache 是模型帮你自动加的。还真不是。模型结构里有这个能力但你是否把 past_key_values 传进去、传得对不对完全取决于你的代码。HuggingFace Transformers 的 generate 方法默认帮你管理了这些所以你平时不太感知得到但如果你自己写循环、或者用纯模型 forward 方法不传缓存模型就会退化成分批重算的朴素模式。很多人在微调时发现推理慢得离谱往往就是在这里没接上。3.2 prefill 和 decode 的分工为什么 decode 阶段吃不满显卡算力把生成拆成 prefill 和 decode 两个阶段不只是为了好叙述它们对计算资源的消耗特征完全相反。prefill 阶段你要一次性处理整个输入 prompt假设长度是 1024。这一阶段是典型的计算密集矩阵乘很大GPU 的算力能充分跑满。所以 prompt 越长prefill 耗时越长但这个阶段总归只发生一次。decode 阶段就完全不一样了每一步只处理一个 token。虽然 KV Cache 省掉了历史计算但这一阶段每个 token 对应的矩阵乘都很小GPU 的空闲反而更多瓶颈变成了内存带宽——每一层都要把缓存的 KV可能几 GB从显存里读出来参与注意力计算还要把新的 KV 写回去。这个阶段每个 token 的耗时主要由“缓存大小”决定而不是由“算力高低”决定。这也是为什么你换了更贵的显卡单 token 生成速度未必能快很多因为大家都在抢内存带宽这个东西。如果用一句话概括设计思路prefill 阶段给你“装满弹匣”decode 阶段用 KV Cache 做“快速连发”。前者要的是吞吐后者要的是延迟。很多推理框架对这两个阶段采取完全不同的调度策略和并行方案也正是基于这个区别。3.3 显存优化三板斧PagedAttention、前缀共享、缓存量化KV Cache 的显存账让人头大所以业界围绕它发展出了一整套优化方案。我挑三个在实操中最常遇到的展开。第一板斧是PagedAttention。这个技术的核心贡献是引入了“按页管理”的思路。在它之前KV Cache 要为每个序列预分配一块连续的显存但一个序列到底会生成多长在生成前谁也说不准。分配少了不够用分配多了浪费而且显存碎片化之后明明总量够却因为不连续而无法分配。PagedAttention 把 KV Cache 切成固定大小的块像操作系统的内存分页一样按需分配可以散落在显存各处。vLLM 的显存复用率大幅提升靠的就是这个。第二板斧是前缀共享Prefix Caching。很多应用场景里不同请求的前缀是完全相同的比如系统提示词、few-shot 示例、长文档的固定开头。如果每次新请求都从头算一遍这些前缀的 KV那就是在做重复劳动。实现上你可以在服务层按内容 hash 前缀的 KV Cache命中时直接继承只对新增部分做计算。这个优化在 RAG、多轮对话场景里尤其明显实测往往能省下 30% 以上的 prefill 时间。第三板斧是KV Cache 量化。既然缓存值都是 FP16那能不能用 INT8 或者 FP8 来存可以。KV Cache 量化就是对缓存做数值压缩常见做法是把 Key 和 Value 分别做分组缩放。INT8 量化可以让缓存显存减半FP8 目前在 Hopper 及更新架构上有硬件级支持速度损失很小。这块的选择要看你对精度损失的容忍度。一般来说K 的量化比 V 更敏感因为 K 直接参与注意力分数的计算误差会被 softmax 放大V 的敏感度稍低。我对精度要求高的生成任务一般先试 INT8 加分组量化效果不满意就退回 FP16。这三板斧不互斥可以叠加使用。实际生产中在保证精度的前提下KV Cache 的显存可以被压到原来的四分之一甚至更低配上长上下文任务才真正变得可用。3.4 多序列生成时的缓存管理beam search 的隐藏复杂度很多人自己实现生成逻辑时都是从贪心解码开始的。一旦改用 beam search 或者 batch 多序列生成KV Cache 的管理复杂度会瞬间上一个台阶。beam search 的每一步会保留 beam width 个候选序列。每一条序列都有自己的历史 token相当于每条 beam 都有自己独立的 KV Cache。当某条 beam 因为得分不好被剪枝或者两条 beam 发生合并时缓存需要跟着做丢弃或复制。更麻烦的是beam 之间的长度可能不一样导致 batch 内 padding 的 token 也会进入缓存白白占显存。这块我没有手写过完整实现但我在调试推理框架时认真看过里面的逻辑。如果你不是非要跟 beam search 死磕我的建议是直接用现成框架的生成接口而不是自己造轮子。真正需要你动手实现 KV Cache 的场景通常是做模型结构定制、或者特殊的数据流控制那种情况下也尽量保持“batch 内序列长度对齐 只缓存有效 token”这个基本原则能少踩很多坑。4. 实测与踩坑KV Cache 的显存账单和常见故障排查4.1 一份真实的显存账本模型权重与 KV Cache 的比例关系纸上谈兵不如直接对账。我以一个 FP16 的 7B 模型为例感受一下各部分显存的真实占比。7B 参数在 FP16 下仅权重就需要约 14 GB如果加上梯度训练场景就是 28 GB再加优化器状态就更夸张了。我们只看推理单条序列、上下文长度 4096 的场景下KV Cache 约 2 GB同样配置下上下文长度 32768 时KV Cache 约 16 GB。这时总显存里权重 14 GB 缓存 16 GB缓存占比超过一半。如果再开 batch size 为 8每条序列都是 32K 上下文KV Cache 直接冲到 128 GB——单卡根本放不下。所以你会看到大模型服务商在宣传“超长上下文”的时候背后的显存成本是成倍增加的。你在本地跑 7B 模型32K 上下文确实能跑可一旦设置过大的 max_lengthOOM 几乎是必然的。我自己跑 qwen 系列时最常遇到的就是 max_model_length 设大了、显存直接炸掉。我建议你养成一个好习惯部署模型之前先用公式算好 KV Cache 的显存上限再结合显卡可用显存确定“上下文长度 × batch size”的组合。不要等爆了再调那个过程很折磨。4.2 高频踩坑实录从缓存长度不匹配到位置编码偏移KV Cache 的坑绝不比它的收益少。我按出现频率排一排基本都是我在调模型时真真切切遇到过的。第一个坑是缓存长度与位置编码不匹配。RoPE旋转位置编码是现在的主流位置编码方案它的特点是“位置信息直接编码在 K、Q 向量里”。使用 KV Cache 时如果新 token 计算时输入的 position_id 是从 0 重新开始算的而不是接着缓存的长度继续递增那么算出来的注意力分数会错得非常离谱。表现就是前几个 token 没问题后面输出越来越乱、逐渐开始复读或无意义内容。排查时看到这种“越到后面越崩”的现象优先检查位置 id 的递增逻辑。第二个坑是多轮对话没清缓存。做多轮对话你需要把历史上所有轮次的 token 都拼接进上下文。如果不小心在新一轮只传入“最新一轮的用户输入”但缓存里还挂着上一轮的 KV模型会丢失前面的上下文回答驴唇不对马嘴。我见过最隐蔽的情况是手动拼接 prompt 时新旧轮次之间的分隔符处理错了导致缓存记录的位置和你预期的不一致。第三个坑是显存碎片化导致的 OOM 假象。有时候你看显存明明还有不少空间但一申请大块连续缓存就 OOM。这是因为 KV Cache 预分配是一整块连续显存。用 vLLM 之后这个问题会好很多因为分页机制把连续需求打散了。如果你自己管理缓存预留显存时一定要按最大可能长度去估而不是当前实际长度去估。第四个坑藏在padding 策略里。batch 推理时不同请求长度不一样短的需要 padding 到跟最长的一样。如果不做 attention mask 掩码这些 padding token 也会被算进缓存和注意力里既浪费显存又污染计算结果。还有一个我差点忘了说half precision 下的精度陷阱。KV Cache 在 FP16 或 BF16 下序列非常长时Key 向量累加带来的误差可能让注意力分布轻微漂移。这不是 bug是数值精度问题。遇到长文本输出质量微降时可以试试把缓存保持为 FP32或者换 BF16有时效果好得出奇。4.3 缓存有没有生效快速判断的三个办法怎么知道你用的推理代码到底有没有真正吃到 KV Cache 的红利不用看源码三个土办法很管用。第一个办法观察每 token 生成耗时随序列长度的变化。如果 KV Cache 生效生成新 token 的耗时应该基本稳定不随历史长度剧烈增长如果没生效序列越长每步耗时增长得越明显到了几千 token 之后卡到你怀疑人生。第二个办法看显存曲线。KV Cache 生效时显存占用会随着生成步数阶梯式往上走因为每一层都在追加新的 K、V。如果显存从一开始就整个 prompt 长度对应的常量说明大概率在重复计算整个序列。这个特征在 nvidia-smi 里看得很清楚。第三个办法对比 logits。同一 prompt 分别跑“缓存开启”和“缓存关闭”两个版本自己写两版代码即可对比每个生成步的 logits 是否完全一致。缓存正确实现时两者应该严格相等任何不一致都说明某个环节写错了比如位置编码、mask、缓存内容更新顺序等。这个办法最可靠因为它直接验证了 KV Cache 的数学等价性。5. 顺着 KV Cache 往下想推理优化的几个延伸方向5.1 长上下文下的有界缓存Sliding Window 与 StreamingLLM显存是有限的上下文是无限的KV Cache 无界增长这件事迟早会碰到天花板。为了在超长文本场景下控制缓存规模学术界和工业界想了很多办法。滑动窗口Sliding Window的思路最直观只保留最近 W 个 token 的 KV超过窗口的历史全部丢弃。代价是模型看不到太久远的信息但很多场景下“最近的上下文”确实比“久远的历史”更重要。StreamingLLM 在此基础上做了改良它发现保留最开始的几个 token 的 KV能很好地稳定注意力分布防止“失去锚点”导致的生成崩溃。把“开头锚点 滑动窗口”组合起来就实现了理论上无限长度、但缓存大小恒定的推理。我自己在长文档问答场景里试过这类方案效果很实用。如果你要处理几十万字的小说或者会议纪要又只有单卡与其硬撑完整 KV Cache不如考虑这类“有损但可控”的缓存策略。5.2 投机采样里的 KV Cache 复用少算一步是一步投机采样Speculative Decoding是最近非常火的推理加速方法它让一个小模型先草拟好几步 token再让大模型一次性验证。这个过程中小模型的 KV Cache 和大模型的 KV Cache 是各管各的但两者之间的协调也很有意思。更值得关注的是动态缓存复用如果验证时草拟的 token 被接受了那么大模型可以直接接受这些 token 的 KV 结果不需要重新计算如果某些位置被拒绝缓存要能快速回滚到试点之前的版本。这要求缓存支持高效的“丢弃尾部”操作。我在看一些推理框架的实现文档时发现这块逻辑写得好的能大幅减少大模型实际前向传播的步数。5.3 设计角度把 KV Cache 当作推理系统的“内存管理问题”说了这么多最后想换个角度看待 KV Cache它本质上不是模型结构问题而是推理系统的内存管理问题。模型结构决定了需要缓存什么而如何高效存储、调度、压缩和复用这些缓存决定了推理系统的性能上限。所以业界才有那么多工程创新围绕 KV Cache 展开vLLM 把它做成显存分页TensorRT-LLM 把它做成高效的内存池SGLang 把它做成了自动前缀缓存路由。这些都是“以 KV Cache 为中心”的系统设计的典型代表。如果你以后要设计自己的推理服务我建议从第一天起就把 KV Cache 当作一等公民来规划估算显存时单独给它建模设计接口时把 prefill/decode 分离做性能优化时优先处理缓存瓶颈。这个视角一旦建立起来你会发现对大模型推理性能的很多感性认识都能对上了。有一点是我在实际操作中反复体会到的KV Cache 本身不难懂难的是在“不加缓存会慢死、加了缓存会撑爆显存”的两难局面里做出合理的取舍。我的习惯是任何推理方案落地前先问自己三个问题——最长上下文会到多少batch 多大显存够不够缓存和权重同时放下这三个问题有了明确答案再动手写代码基本不会跑偏。最后再分享一个实用小技巧如果你不确定某次生成到底有没有命中缓存就在生成前后各打一次显存占用快照如果增长符合“每 token 增量 ≈ 层数 × 隐藏维度 × 2 字节 × 2”这个公式那缓存就是真正在为你工作了。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑