资讯动态

KV Cache 原理、显存计算与推理优化实战

发布时间:2026/10/20 16:24:50 来源:尧图企业网站定制
1. 为什么大模型推理绕不开 KV Cache第一次接触 Transformer 推理优化的人十有八九会在某个时刻盯着显存占用曲线发呆明明模型权重才几个 G怎么一跑长文本生成显存就蹭蹭往上涨batch 稍微开大一点直接 OOM。这个问题背后最大的“元凶”之一就是注意力机制里那个看似不起眼、实则决定生死的东西——KV Cache。先把话说直白点KV Cache 不是什么高深的新算法它本质上是一种用显存换时间的工程手段。Transformer 在自回归生成时每生成一个新 token都要拿当前 token 的 Query 去和前面所有 token 的 Key、Value 做注意力计算。如果不做任何缓存那每生成一步前面所有 token 的 K 和 V 都要重新算一遍。生成第 N 个 token 时计算量是 O(N)生成整个长度为 N 的序列累计就是 O(N²)。而有了 KV Cache前面算过的 K、V 直接存下来复用每步只需要算当前这一个 token 的 K、V单步计算量降到 O(N)整体复杂度从平方级压到线性级。这个差别有多大我拿一个实际场景举例。假设你在部署一个 7B 级别的对话模型输出长度 512 个 token。不开 KV Cache注意力部分的计算量大概是开了缓存的几十倍量级首 token 之后的每个 token 延迟会肉眼可见地飙升用户体感就是“打字机卡成 PPT”。所以但凡你看到某个推理框架号称“高性能”它一定在 KV Cache 上做了大量文章。这篇文章我打算把 KV Cache 从原理到落地讲透。包括它到底缓存了什么、显存怎么算、为什么会有 MQA/GQA 这些变体、PagedAttention 又是怎么解决碎片问题的、长上下文场景下有哪些坑。适合正在做推理部署、想优化显存和吞吐的工程师也适合想搞明白“为什么我的显存不够用”的算法同学。看完你至少能做到两件事第一能自己估算一个模型在给定并发下的 KV Cache 显存开销第二知道在什么场景下该选哪种注意力变体和缓存策略。2. KV Cache 的核心原理与显存账本2.1 自回归生成到底重复算了什么要理解 KV Cache得先回到 Transformer 的注意力公式。标准的缩放点积注意力是这样的Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V在训练阶段整个序列是并行喂进去的Q、K、V 都是一次性算出来的不存在重复计算的问题。但推理阶段是自回归的你给一个 prompt模型生成第一个 token然后把这个 token 拼回输入再生成第二个如此循环。问题就出在这个循环里。生成第 t 个 token 时模型需要用到前 t-1 个 token 的 K 和 V 来计算注意力。如果每步都重新算那第 t 步就要重新计算前 t-1 个位置的 K、V 投影。这些投影涉及的是同样的输入、同样的权重矩阵结果完全一样纯属重复劳动。KV Cache 的思路非常朴素既然算过一次那就存起来。每生成一个新 token只计算这个新 token 对应的 K、V然后追加到缓存里注意力计算时直接读取整个缓存。这样每步的 K、V 投影计算量从 O(t) 降到 O(1)注意力打分部分仍然是 O(t)因为要和所有历史位置算相似度但整体已经从平方级降下来了。注意KV Cache 只缓存 Key 和 Value不缓存 Query。因为 Query 只跟当前 token 有关历史位置的 Query 在后续步骤里根本用不到。这一点很多人第一次会搞混。2.2 显存开销的精确计算公式KV Cache 省了计算代价是显存。这笔账必须算清楚否则你永远不知道自己为什么 OOM。单个 token 的 KV Cache 大小可以用下面这个公式估算单 token KV 大小 2 × num_layers × num_kv_heads × head_dim × dtype_bytes这里的 2 是因为要存 K 和 V 两份。num_kv_heads是关键它不一定等于注意力头数后面讲 MQA/GQA 时会展开。dtype_bytes取决于你用 FP16 还是 INT8FP16 就是 2 字节。我拿一个具体的配置算一遍方便你对量级有感觉。假设一个模型有 32 层隐藏维度 409632 个注意力头head_dim 是 128用 FP16 推理单 token KV 大小 2 × 32 × 32 × 128 × 2 字节 524288 字节 ≈ 0.5 MB看起来不大但你要乘以序列长度和并发数。序列长度 4096、并发 16 的情况下总 KV Cache 0.5 MB × 4096 × 16 ≈ 32 GB32 GB 什么概念一张 24G 的卡直接放不下还没算模型权重本身。这就是为什么长上下文 高并发是显存杀手。如果你把num_kv_heads通过 GQA 降到 8单 token 大小直接变成原来的四分之一总占用降到 8 GB一张卡就能扛住了。这个对比就是 GQA 存在的全部意义。2.3 为什么说它是“推理加速的核心机制”有人可能会问KV Cache 省的是计算跟“加速”到底什么关系。这里要区分两个指标计算量和内存带宽。在大模型推理里尤其是 decode 阶段瓶颈往往不是算力而是显存带宽。每生成一个 token都要把整个 KV Cache 从显存读进计算单元这个读取量正比于缓存大小。所以 KV Cache 的设计直接决定了两个东西一是你能开多大的 batch 和多长的上下文显存容量二是每个 token 的生成速度显存带宽。这也是为什么后来的优化几乎都围绕“如何让 KV Cache 更小、读取更快”展开——MQA、GQA 是压缩头数PagedAttention 是减少碎片量化是压缩精度。理解了这一点你就理解了整个推理优化领域的主线。3. 从 MHA 到 GQA注意力变体怎么选3.1 MHA、MQA、GQA 的本质区别原始的 Multi-Head AttentionMHA里每个注意力头都有自己独立的 K、V 投影头数等于 KV 头数。这样表达能力最强但 KV Cache 也最大。为了压缩缓存研究者提出了两条路线。第一条是 Multi-Query AttentionMQA所有注意力头共享同一组 K、V。也就是说num_kv_heads 1。KV Cache 直接缩小到原来的 1/num_heads压缩比非常夸张。但代价是表达能力下降训练时容易不稳定生成质量在某些任务上会掉点。第二条是 Grouped-Query AttentionGQA算是 MHA 和 MQA 的折中把注意力头分成若干组每组共享一组 K、V。比如 32 个注意力头分成 8 组那num_kv_heads 8KV Cache 缩小到 1/4质量损失比 MQA 小很多。现在主流的中大规模模型基本都默认用 GQA因为它把压缩比和质量平衡得比较好。变体KV 头数KV Cache 相对大小表达能力典型适用场景MHA等于注意力头数1x最强小模型、质量优先GQA注意力头数的 1/4 到 1/81/4 到 1/8接近 MHA主流大模型默认MQA11/num_heads最弱极致吞吐、边缘部署3.2 选型时到底看什么指标选哪种变体不是拍脑袋得看你的部署目标。我一般按下面这个顺序判断。先看质量容忍度。如果是面向用户的对话产品生成质量掉一点用户马上能感知那就优先 GQA别为了省显存上 MQA。如果是内部批处理任务比如批量摘要、分类质量要求没那么苛刻MQA 换来的吞吐提升很划算。再看并发规模。并发越高KV Cache 越吃显存压缩收益越明显。单用户低并发场景MHA 和 GQA 的差距可能感知不到但你要做几十路并发GQA 省下的显存就是能不能多塞几个请求的区别。最后看硬件条件。显存小的卡上GQA 几乎是必选项。我见过有人在 16G 卡上硬跑 MHA 的长上下文模型结果上下文只能开到 2K再长就 OOM体验极差。换成 GQA 版本后上下文直接翻倍。实操心得模型选型阶段就要把 KV Cache 开销算进去别等部署了才发现显存不够。一个简单的判断方法是用 2.2 节的公式算出单 token 大小乘以你预期的最大上下文和并发数看是否超过显存的一半。超过一半就要警惕因为还要留空间给权重和激活值。3.3 一个容易忽略的细节head_dim 的影响很多人只关注头数忽略了head_dim。实际上 KV Cache 大小和head_dim是线性关系。有些模型为了提升表达能力会把head_dim设得比较大比如 128 甚至 256这时候即使头数不多单 token 缓存也不小。举个例子同样是 32 层、FP16配置 A 是 32 头 × head_dim 64配置 B 是 16 头 × head_dim 128。两者隐藏维度都是 2048但单 token KV 大小配置 A2 × 32 × 32 × 64 × 2 262144 字节配置 B2 × 32 × 16 × 128 × 2 262144 字节诶居然一样。因为隐藏维度相同的情况下头数 × head_dim 是定值KV Cache 总量只跟隐藏维度有关。所以真正决定缓存大小的是隐藏维度 × 层数头数只是通过 GQA 的分组比例间接影响。这个认知能帮你在看模型配置时快速估算不用被头数迷惑。4. PagedAttention 与缓存管理的工程实践4.1 传统 KV Cache 的显存碎片问题早期推理框架给每个请求预分配一块连续的 KV Cache 空间大小按最大可能长度算。这带来两个浪费一是内部碎片请求实际只生成了 100 个 token但预分配了 2048 的空间剩下 1948 全浪费二是外部碎片不同请求的块大小不一显存里出现很多无法利用的小空洞。在低并发时这个问题不明显但一旦并发上来碎片能吃掉 30% 到 60% 的显存。我实测过一个场景理论能开 32 路并发实际因为碎片只能开 18 路吞吐直接腰斩。这就是 PagedAttention 要解决的核心痛点。4.2 分页机制是怎么工作的PagedAttention 的思路借鉴了操作系统的虚拟内存分页。它把 KV Cache 切成固定大小的块block比如每块存 16 个 token 的 KV。每个请求的缓存不再要求连续而是由一组物理块通过块表block table映射起来逻辑上连续、物理上可以离散。这样做的好处很直接块是固定大小的分配和回收都以块为单位不会产生外部碎片请求按需申请块生成到哪申请到哪不会预分配浪费内部碎片最多浪费最后一个块的部分空间。实测下来显存利用率能从 60% 左右提到 90% 以上同样的卡能多扛将近一半的并发。块表的存在还带来一个额外好处前缀共享。如果多个请求有相同的 system prompt它们的开头 KV 块可以指向同一份物理块不用各存一份。这在多轮对话和批量推理场景里省下的显存相当可观。4.3 落地时的参数调优用支持 PagedAttention 的框架时有几个参数值得调。block size是最关键的。块太小块表本身占内存管理开销大块太大最后一个块的浪费变多。常见取值是 16我试过 8 和 3216 在大多数场景下比较均衡。如果你的请求长度普遍很短可以调到 8如果都是长文本32 也行。gpu_memory_utilization决定了框架能用多少显存做 KV Cache。默认一般是 0.9意思是留 10% 给其他开销。这个值别设太满设到 0.95 以上容易在峰值时 OOM。我一般设 0.85 到 0.9留点余量更稳。max_num_seqs控制最大并发请求数。这个不是越大越好得配合你的显存和上下文长度。设太大框架会尝试调度更多请求但 KV Cache 不够时反而触发频繁的换入换出吞吐不升反降。建议先用小值压测逐步往上加找到吞吐拐点。注意开启前缀共享后如果请求的 prompt 前缀差异很大共享收益会很低甚至因为块表查找带来额外开销。这种场景下可以考虑关掉或者把 system prompt 统一化来提高共享率。5. 长上下文与量化场景下的实战避坑5.1 长上下文为什么让 KV Cache 失控上下文从 4K 拉到 32KKV Cache 是线性增长的8 倍长度就是 8 倍显存。但问题不止于此长上下文下注意力计算本身也变重decode 阶段每步要读取的缓存更大显存带宽压力陡增单 token 延迟会明显上升。我踩过的一个坑是以为只要显存够就能开长上下文结果发现延迟高到没法用。后来才意识到长上下文场景下 KV Cache 的读取带宽才是瓶颈。解决办法有几个一是用 GQA 压缩缓存二是对 KV Cache 做量化三是用滑动窗口注意力只保留最近的一部分缓存。滑动窗口是个很实用的技巧。很多任务其实不需要关注全部历史只保留最近 N 个 token 的 KV 就够了。这样缓存大小被固定住不随序列增长。代价是丢失了远距离信息适合那些局部依赖为主的任务比如流式对话、代码补全。5.2 KV Cache 量化省显存但别踩精度坑把 KV Cache 从 FP16 量化到 INT8显存直接减半这是长上下文场景的救命稻草。但量化不是无脑开有几个细节要注意。首先是量化粒度。per-tensor 量化最简单但精度损失大per-channel 或 per-token 量化精度好很多实现也复杂一些。我一般优先选 per-token因为 K 和 V 在不同 token 上的分布差异比较大逐 token 量化能更好地保留信息。其次是哪些层可以量化。实测下来浅层的 KV 对量化更敏感深层相对鲁棒。如果精度掉得厉害可以只量化后半部分的层前面几层保持 FP16。这个策略能在省显存和保质量之间找到不错的平衡。最后是校准数据。量化需要校准来确定缩放因子校准数据的分布要贴近实际推理场景。用通用语料校准结果在特定领域任务上掉点这种情况我遇到过不止一次。建议用你实际业务的一小批样本做校准。量化方案显存节省精度损失实现复杂度推荐场景FP16 原始0%无低短上下文、质量优先INT8 per-tensor50%中等低对精度不敏感的任务INT8 per-token50%较小中长上下文通用场景分层混合量化30% 到 50%小高精度要求高的长上下文5.3 常见问题速查实际部署中KV Cache 相关的问题翻来覆去就那么几类。我整理成一张表方便你对照排查。现象可能原因排查方向解决手段显存随生成持续上涨缓存未释放或泄漏检查请求结束后块是否回收升级框架版本检查块表释放逻辑并发上不去碎片或预分配浪费看显存利用率指标启用 PagedAttention调小 block size长上下文延迟高缓存读取带宽瓶颈测单 token 延迟随长度变化用 GQA、量化、滑动窗口量化后质量掉点校准数据不匹配对比量化前后输出换 per-token用业务数据校准前缀共享没效果prompt 前缀差异大统计共享块命中率统一 system prompt 或关闭共享实操心得排查 KV Cache 问题时先看显存曲线。如果曲线是阶梯状上升然后突然下降说明块在按请求生命周期分配回收正常如果只升不降基本可以确定是泄漏。另外压测时一定要用真实长度的请求用固定短请求压出来的并发数没有参考价值。6. 我个人的一些部署体会KV Cache 这个东西原理讲起来半小时能讲完但真正在工程里用好靠的是对显存账本的敏感度和对场景的准确判断。我自己最大的体会是不要等到 OOM 了才去关注它。在模型选型和容量规划阶段就应该把 KV Cache 开销算进去把它当成和模型权重同等重要的显存预算项。另一个体会是没有万能的配置。GQA 不是永远比 MHA 好量化不是永远划算PagedAttention 也不是所有场景都能提吞吐。关键是想清楚你的瓶颈在哪是显存容量不够还是带宽不够还是并发调度效率低。瓶颈不同优化手段完全不同。我见过有人盲目上量化结果质量掉了显存其实本来够用纯属白折腾。最后分享一个小技巧如果你在做容量规划可以先用一个简化的公式快速估算——显存需求 ≈ 权重 2 × 层数 × 隐藏维度 × 序列长度 × 并发 × 精度字节。这个公式忽略了头数细节但量级上足够准能帮你在几分钟内判断一个配置是否可行比跑一遍压测快得多。等粗略估算通过后再上真实压测做精细调优。

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

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

免费获取报价 →
↑