首页 / 资讯中心 / 文章详情

KV Cache 原理与优化:从显存计算到 PagedAttention 与量化实战

KV Cache 原理与优化:从显存计算到 PagedAttention 与量化实战 ★ FEATURED ARTICLE
1. 从一次对话卡顿说起KV Cache 到底在解决什么如果你用过大语言模型做多轮对话大概率遇到过这种情况第一句话回复得飞快聊到第十轮、第二十轮输出速度肉眼可见地变慢首字延迟也越来越高。很多人第一反应是模型是不是累了服务器是不是在限流其实都不是。真正的原因藏在 Transformer 的注意力计算机制里而 KV Cache 就是针对这个问题最核心的优化手段。先把结论摆出来KV Cache 是一种用显存换时间的缓存机制它把自回归生成过程中已经算过的 Key 和 Value 张量存下来避免每一步都重复计算历史 token 的注意力。没有它生成第 N 个 token 的计算量会随序列长度线性增长聊得越久越慢有了它单步计算量被压到接近常数级别但代价是显存占用随上下文长度线性膨胀。这就是越聊越慢和长上下文很贵这两件事的共同根源。这篇文章适合三类人看一是正在做 LLM 推理部署、被延迟和显存折磨的工程师二是想搞懂长上下文模型成本结构的产品和技术决策者三是准备面试、需要把 KV Cache 讲清楚的同学。我会从原理推导讲到工程实操包括显存怎么算、PagedAttention 怎么省、量化怎么选、长上下文为什么贵得离谱尽量把每个为什么都讲透。需要先明确一个前提本文讨论的是**自回归解码autoregressive decoding**场景也就是 GPT 类模型一个 token 一个 token 往外吐的过程。编码器类模型或者一次性前向的场景KV Cache 的意义完全不同不在讨论范围内。2. 自回归生成的计算冗余为什么不做缓存会越聊越慢2.1 注意力机制里到底重复算了什么要理解 KV Cache得先回到 Transformer 的注意力公式。对于每个 token模型会生成三个向量QueryQ、KeyK、ValueV。注意力计算本质上是Attention(Q, K, V) softmax(QK^T / sqrt(d)) V在生成第 t 个 token 时这个 token 的 Query 需要和前面所有 t 个 token 的 Key 做点积得到注意力权重再对所有 Value 加权求和。问题来了第 t 步要用到的前 t-1 个 token 的 K 和 V在第 t-1 步、第 t-2 步……其实都已经算过了。如果不做任何缓存每生成一个新 token都要把整个历史序列重新跑一遍前向重新计算所有历史 token 的 K 和 V。这意味着生成第 t 个 token 时K/V 的计算量是 O(t)整个序列生成完的总计算量就是 O(t²)。序列越长平方级增长越恐怖。我做过一个粗略的实测对比用同一个 7B 模型、同样的硬件生成 512 个 token方案首 token 延迟平均每 token 延迟总耗时无缓存每步重算全序列约 180ms随长度递增末段约 400ms约 150s开启 KV Cache约 180ms稳定在 35ms 左右约 18s差距接近一个数量级而且序列越长差距越大。这就是为什么所有正经的推理框架默认都开 KV Cache——不开根本没法用。2.2 缓存带来的复杂度变化开启 KV Cache 后第 t 步只需要计算当前这个新 token 的 Q、K、V然后把新的 K、V 追加到缓存里用当前 Q 去和缓存中全部 K 做注意力。单步的注意力计算量从重算全序列变成了只算新 token 对历史的注意力虽然注意力本身还是 O(t)因为要和 t 个历史 K 做点积但省掉了最贵的线性投影部分。这里有个容易混淆的点KV Cache 并没有把注意力复杂度从 O(t²) 降到 O(t)它降的是重复计算那部分。单步注意力仍然是 O(t)所以总复杂度还是 O(t²)只是常数因子小了很多。真正把注意力降到 O(t) 的是另一类技术比如线性注意力、稀疏注意力那是另一个话题。理解这一点很重要因为它解释了为什么长上下文依然慢——即使有 KV Cache注意力那部分开销还是随长度增长只是没有缓存的话会更慢。2.3 一个生活化的类比把 KV Cache 想象成做菜。你要做一道需要不断加料的汤每加一次新料新 token都要尝一遍整锅汤的味道注意力。如果不做记录每次加料你都得把之前所有料重新放一遍再尝——这就是无缓存。KV Cache 相当于你把每次加的料单独留了一份样本K 和 V 缓存加新料时直接拿新料和所有样本比对不用重新熬整锅。但样本越攒越多你的台面显存就越占越满。这就是缓存的两面性省了时间占了空间。3. KV Cache 的显存账本长上下文真正的成本在哪3.1 显存占用的精确计算公式很多人知道长上下文吃显存但具体吃多少、怎么算说不清楚。这里给一个可以直接套用的公式。单个序列的 KV Cache 显存占用为KV Cache 字节数 2 × L × H_kv × d_head × S × B × P其中各参数含义2K 和 V 两份LTransformer 层数H_kvKV 头数注意不是注意力头数GQA/MQA 下会更少d_head每个头的维度S序列长度上下文长度Bbatch sizeP每个数值的字节数FP16 是 2INT8 是 1FP8 是 1举个具体例子。一个 7B 级别的模型假设 L32H_kv32标准 MHAd_head128用 FP16 推理batch1上下文 40962 × 32 × 32 × 128 × 4096 × 1 × 2 2,147,483,648 字节 ≈ 2GB光是 KV Cache 就 2GB。如果把上下文拉到 128K同样的模型2 × 32 × 32 × 128 × 131072 × 1 × 2 ≈ 64GB64GB 的 KV Cache比模型权重本身7B FP16 约 14GB还大好几倍。这就是长上下文贵的本质——不是模型变大了是缓存爆炸了。3.2 为什么长上下文成本是超线性的单看公式KV Cache 随序列长度 S 是线性的。但实际部署中长上下文的成本远不止线性原因有三层第一层注意力计算本身是 O(S²)。即使有 KV Cache每个新 token 都要和全部历史做注意力序列翻倍单步注意力计算量也翻倍总计算量翻四倍。第二层显存压力导致 batch 被迫缩小。KV Cache 占的显存多了能同时处理的并发请求就少了。原本能跑 batch32 的卡长上下文下可能只能跑 batch4吞吐量直接掉到八分之一。第三层显存碎片和调度开销。传统做法要为每个请求预留最大长度的连续显存实际用不满就浪费用超了就得换出。长上下文下这种浪费被放大。把这三层叠起来长上下文的实际成本增长远快于线性。这也是为什么很多厂商对长上下文单独定价而且价格曲线很陡。3.3 不同上下文长度下的显存实测对照下面这张表是我在一张 80GB 显存的卡上用同一个模型GQAH_kv8L32d_head128FP16跑出来的 KV Cache 占用估算上下文长度单序列 KV Cache理论最大并发仅算缓存4K约 0.5GB约 15032K约 4GB约 18128K约 16GB约 41M约 128GB单序列都放不下注意这里用的是 GQA分组查询注意力KV 头数从 32 降到 8缓存直接省了 4 倍。如果换成 MQAH_kv1还能再省 8 倍。这就是为什么现在主流模型几乎都用 GQA——它是在几乎不损失效果的前提下把 KV Cache 砍掉数倍的最直接手段。提示算显存时千万别漏掉 batch 维度。很多人在单序列上算得好好的一上并发就 OOM就是因为忘了乘 B。4. 从 MHA 到 GQA/MQA用架构设计给缓存瘦身4.1 多头注意力的缓存冗余问题标准的多头注意力MHA里每个注意力头都有自己独立的 K 和 V。假设有 32 个注意力头那就有 32 份 K 和 32 份 V 要缓存。但研究发现这些头之间的 K/V 其实高度冗余很多头的 Key 分布非常接近各存一份纯属浪费。这个观察催生了 MQAMulti-Query Attention所有注意力头共享同一份 K 和 V只保留多份 Query。缓存直接降到原来的 1/HH 是头数。代价是效果会掉一点训练稳定性也差一些。4.2 GQA 的折中分组共享MQA 省得太狠效果损失明显。GQAGrouped-Query Attention取了个中间值把注意力头分成若干组每组共享一份 K/V。比如 32 个 Query 头分成 8 组每组 4 个头共享一份 K/V那 KV 头数就是 8缓存降到 1/4。这个折中非常划算。实测下来GQA 在绝大多数任务上效果和 MHA 几乎无差别但缓存和计算都省了数倍。现在开源模型里Llama 系列、Qwen 系列、Mistral 系列基本都默认 GQA。注意力类型KV 头数以 32 头为例相对缓存大小效果损失MHA321x基准GQA8 组81/4几乎无损GQA4 组41/8轻微MQA11/32可感知4.3 选型时的实际考量如果你在选模型或者做微调关于注意力类型有几点经验推理部署优先选 GQA 模型这是当前性价比最高的方案没有之一。MQA 适合极端显存受限场景比如边缘设备但要接受效果打折。微调时不要轻易改注意力结构除非你有充足的算力和数据重新训练否则直接用原模型的结构最稳。我见过有人为了省显存把一个 MHA 模型硬改成 MQA 再微调结果效果崩了最后发现是训练不充分导致的。改结构这件事成本比想象中高。5. PagedAttention 与缓存量化工程侧的两把省钱利器5.1 传统缓存分配的浪费在哪前面算的显存是理论占用实际部署中浪费更严重。传统做法是每个请求进来就按最大可能长度比如 4096预留一块连续显存。但实际生成时很多请求可能只用了 500 个 token 就结束了剩下 3500 个 token 的显存白白占着不能用。这种浪费在长上下文下极其致命。假设预留 128K 的连续空间实际平均只用 10K浪费率高达 90% 以上。而且连续显存分配本身就会产生碎片跑久了碎片化严重明明总显存够就是分配不出来。5.2 PagedAttention 的分页思路PagedAttention 借用了操作系统虚拟内存分页的思想把 KV Cache 切成固定大小的块block不再要求连续用一张块表来映射逻辑位置和物理块。每个请求按需分配块用多少给多少块之间不需要连续。这个改动带来几个直接好处内部碎片大幅减少最后一个块可能用不满但浪费最多一个块的大小而不是整个预留空间。显存共享成为可能多个请求如果前缀相同比如相同的 system prompt可以共享同一批物理块这在批量推理和 few-shot 场景下省得非常多。碎片问题缓解块是固定大小的分配和回收都规整不会产生难以利用的碎片。实测下来PagedAttention 能把显存利用率从传统的 20%-40% 提升到 90% 以上吞吐量提升 2-4 倍是常态。这也是 vLLM 这类框架能火起来的核心技术之一。5.3 KV Cache 量化精度换空间除了分页另一个省钱手段是量化。KV Cache 默认用 FP16 存每个数值 2 字节。如果量化到 INT8直接省一半量化到 FP8也是省一半但精度更好。量化方案每数值字节相对 FP16 省精度影响FP162基准无FP8 (E4M3)150%极小INT8150%小INT40.575%可感知需谨慎实操经验FP8 是目前最推荐的 KV Cache 量化方案省一半显存精度损失基本测不出来而且新一代硬件对 FP8 有原生支持速度还更快。INT8 也可以但要注意 per-channel 还是 per-token 量化选错了精度掉得厉害。INT4 除非显存实在紧张否则不建议长上下文下误差会累积。注意KV Cache 量化和权重量化是两回事可以独立选择。权重用 INT4、KV Cache 用 FP8 是很常见的组合。5.4 一个组合优化的实际配置把 GQA PagedAttention FP8 KV Cache 组合起来效果是叠加的。以那个 7B GQA 模型为例128K 上下文原始 MHA FP16 连续分配约 64GB实际可用并发几乎为 0GQA省 4 倍 FP8省 2 倍 分页利用率 90%约 8GB能跑好几个并发从单序列都费劲到能跑并发这就是工程优化的价值。这些手段不需要重新训练模型纯推理侧改动性价比极高。6. 长上下文推理的实战调优与常见坑6.1 显存不够时的排查顺序线上遇到 OOM别急着加卡按这个顺序排查先看是不是 KV Cache 爆了。用 profiling 工具看显存分布如果 KV Cache 占了大头那就是上下文或并发的问题。检查是否开了 GQA 和分页。如果还在用 MHA 连续分配先换框架往往立竿见影。考虑 KV Cache 量化。FP16 换 FP8直接省一半。限制最大上下文和并发。给请求设一个合理的 max_len 上限别让单个请求吃掉全部显存。最后才考虑加卡或换更大显存的卡。我见过太多团队一 OOM 就加卡结果发现是配置问题白白多花钱。6.2 长上下文下的延迟优化显存解决了延迟还是可能高。长上下文的首 token 延迟prefill 阶段和后续 token 延迟decode 阶段要分开看Prefill 阶段处理整个输入 prompt计算量大是 O(S²)。优化手段是 chunked prefill分块预填充把长 prompt 切成小块逐步处理避免单次占用过多资源。Decode 阶段逐 token 生成每步都要读全部 KV Cache。优化手段是 KV Cache 量化减少带宽压力以及用更快的注意力 kernel如 FlashAttention 系列。实测中decode 阶段往往是长上下文延迟的主要来源因为每生成一个 token 都要把几十 GB 的 KV Cache 读一遍显存带宽成了瓶颈。这也是为什么 KV Cache 量化对长上下文延迟改善特别明显——它直接减少了要读的数据量。6.3 几个容易踩的坑坑一以为 KV Cache 能无限复用。不同请求的 KV Cache 不能随便混用除非前缀完全一致。前缀缓存prefix caching是有条件的别想当然。坑二忽略 prefill 和 decode 的资源竞争。长 prompt 的 prefill 会瞬间吃满算力把正在 decode 的请求卡住。生产环境要做调度隔离或者用 chunked prefill 平滑掉峰值。坑三量化后没做精度回归。KV Cache 量化后一定要跑一遍评测集确认精度没掉。有些任务对 KV 精度特别敏感比如需要精确回忆长文档细节的任务。坑四max_len 设得过大。为了支持长上下文把 max_len 设成模型上限结果每个请求都预留巨量显存。正确做法是按实际业务分布设一个合理的值长请求单独走特殊通道。6.4 一个可复现的调优 checklist把上面的经验整理成一个可以直接用的清单[ ] 确认模型用的是 GQA 或 MQA不是 MHA[ ] 推理框架支持 PagedAttention 或类似的分页机制[ ] KV Cache 开启 FP8 量化并跑过精度回归[ ] 配置了合理的 max_len 和最大并发[ ] 开启了 prefix caching如果业务有大量相同前缀[ ] prefill 和 decode 有调度隔离或 chunked prefill[ ] 有显存和延迟的监控能定位到 KV Cache 占用这套组合拳打下来同样的硬件能支撑的长上下文并发通常能提升 3-5 倍延迟也能明显下降。7. 我对 KV Cache 这件事的几点个人体会做了几年推理优化KV Cache 是我觉得最反直觉的一个点。表面上它是个缓存实际上它决定了整个推理系统的成本结构。很多人优化模型权重、优化算子却忽略了 KV Cache 才是长上下文场景下的真正瓶颈。我的第一个体会是架构层面的优化GQA比工程层面的优化量化、分页收益更大但工程优化更容易落地。如果你能选模型优先选 GQA 的如果模型已经定了那就把工程手段用满。第二个体会是长上下文的成本不是线性的别用线性思维做预算。上下文从 4K 到 128K 是 32 倍但实际成本可能是几百倍。做产品定价和容量规划时一定要按实测数据来别拍脑袋。第三个体会是KV Cache 量化是当前性价比最高的优化没有之一。FP8 省一半显存、几乎无损、硬件还加速这种好事不多。如果你的框架还不支持值得花时间升级。最后一个实用建议上线前一定要做长上下文的压力测试而且要测真实分布不是测极限值。很多问题碎片、调度竞争、量化误差累积只有在真实负载下才会暴露。我踩过的坑里有一半是实验室跑得好好的一上量就崩。KV Cache 这个话题往下还能挖很多比如投机解码怎么和 KV Cache 配合、多轮对话的前缀复用怎么做、分布式推理下 KV Cache 怎么切分。每一个都够写一篇。但把上面这些搞明白你已经能解决 80% 的实际问题了。
阅读完成 · 觉得有帮助?
咨询建站