KV Cache原理与显存计算:大模型推理优化核心指南

KV Cache原理与显存计算:大模型推理优化核心指南 在 Transformer 架构的大模型推理场景里KV Cache 几乎无处不在地影响着吞吐量、延迟和显存占用。面试聊到大模型推理优化、长文本生成、甚至在线服务的成本控制时KV Cache 都是绕不开的核心概念。很多人知道“生成阶段会缓存 K 和 V”但追问下去就发现几个容易混淆的点预填充阶段为什么不缓存、缓存为什么只存 K 和 V 而不存 Q、显存和内存中的缓存如何管理、多轮对话里的 token 复用到底复用了什么。这篇博客把这些链条拆开讲清楚包含计算公式、实际推理流程、显存占用分析、常见面试追问和生产环境里的缓存治理思路。本文适合这几类读者准备大模型推理相关岗位面试的开发者正在做推理服务性能优化或成本控制的工程师以及想搞懂“为什么同一个模型、同一批请求加上缓存后吞吐量差异巨大”的大模型应用开发者。读完能理清 KV Cache 的前因后果能自己估算一个模型需要多少显存能定位缓存命中率低、显存不足、响应变慢这类问题的排查方向。1. 先从一次完整的 token 生成看起搞清楚缓存到底缓在哪一步要理解 KV Cache 的作用先要理解自回归生成。大模型生成文本不是一个词一个词独立预测的而是把已经生成的 token 再次作为输入预测下一个 token如此循环。这个过程可以拆成两个阶段预填充阶段和逐 token 生成阶段。1.1 自回归生成的两阶段到底发生了什么以 Llama、Qwen 这类 decoder-only Transformer 为例。用户输入“介绍一下缓存机制”之后模型先把这个句子切分成 token 序列假设得到 8 个 token。第一次前向计算时这 8 个 token 会一起进入模型。模型基于这 8 个 token 的上下文预测第 9 个 token。这一整次计算就是预填充阶段。从第 9 个 token 开始模型每生成一个 token都需要把“原始输入 已经生成的历史 token”拼接起来再做一次前向计算。假设模型已经生成了 100 个 token生成第 101 个 token 时理想情况下要重新计算前 100 个 token 对应的注意力结果。如果每一次生成都重新从第一层开始算计算量会随着生成长度线性膨胀。实际推理框架不会这么做它们会使用一种缓存机制把历史 token 在注意力计算中产生的中间结果保存下来。这个中间结果就是 KV Cache。所以从功能上看KV Cache 解决的是重复计算问题。没有缓存时每生成一个 token 都要重新计算全部历史 token 的注意力有缓存后每生成一个 token 只需要计算当前这个新 token 的注意力历史部分直接复用缓存。1.2 Q、K、V 里为什么只缓存 K 和 VTransformer 的注意力计算可以写成这样Attention(Q, K, V) softmax(Q * K^T / sqrt(d_k)) * V其中 Q 是查询向量K 是键向量V 是值向量。对于一个新 token 来说它的 Q 只用来和所有历史 token 的 K 做匹配匹配结果再和所有历史 token 的 V 做加权求和。这意味着新 token 需要看到历史 token 的 K 和 V但不需要重新计算历史 token 的 Q。具体到推理生成第 101 个 token 时历史 100 个 token 不再作为 query 去查询其他 token它们只需要提供 K 和 V 供新 token 查询。因此缓存可以只保存每一层、每一个历史 token 的 K 和 V。Q 是当前 token 自己计算出来的用完即走不需要缓存。这个设计的本质在于自回归模型中token 之间的依赖方向是单向的每个 token 只能关注自己和之前的 token。已经过去的 token 不会再作为查询方重新参与注意力计算。1.3 缓存存储的基本单元按层、按请求存储KV Cache 不是一张简单的表。对于一个 L 层的 Transformer 模型每一层都有独立的 K 和 V 缓存。每一层的 KV 缓存形状可以近似表示为[batch_size, seq_len, num_key_value_heads, head_dim]其中 batch_size 是当前请求批次里的样本数seq_len 是已经处理的历史 token 数量num_key_value_heads 是 KV 头的数量head_dim 是每个注意力头的维度。把这三者相乘再乘以层数、乘以每个元素的字节数就得到一份 KV Cache 占用的显存大小。需要区分的是预填充阶段产生历史 token 的 K 和 V然后整个生成阶段逐步追加新 token 的 K 和 V。因此 KV Cache 的增长是按 token 变化的每生成一个 token每个 Transformer 层会追加一组 K 和 V。2. KV Cache 由哪些指标决定模型参数、上下文长度、批次大小、精度面试里最容易考到的并不是“什么是 KV Cache”而是给出模型配置让你计算 KV Cache 占多少显存。要会算这个先要对模型结构中的几组关键数字做拆解。2.1 决定 KV Cache 大小的核心变量KV Cache 的显存占用可以拆成四个因子层数、序列长度、KV 头数量、头的维度。还有数据类型在内通常是 2 字节的 BF16 或 FP16。如果使用 FP8 或 INT8 量化单字节数会减半。先看一组常见参数。假设一个模型的配置如下参数示例值说明hidden_size4096隐藏层维度num_attention_heads32注意力头数num_key_value_heads8KV 头数少于 Q 头数时使用 GQAnum_hidden_layers32Transformer 层数max_position_embeddings32768最大上下文长度数据类型BF16每个元素 2 字节KV Cache 每层每组 K 和 V 的大小为2 * batch_size * seq_len * num_key_value_heads * head_dimhead_dim 在很多模型中等于 hidden_size 除以 num_attention_heads。上面的配置里head_dim 是 4096 / 32 128。把这个值代入公式假设 batch_size 为 1seq_len 为 32768得到单层 KV Cache 大小2 * 1 * 32768 * 8 * 128 67108864 个元素每个元素 BF16 占 2 字节所以单层占用 128 MB。再乘以 32 层总占用约 4 GB。这里的核心观察是这只是 batch_size 1 的情况。把 batch 提高到 16KV Cache 会直接乘 16变成 64 GB。2.2 减掉 KV 头数量的 GQA 与 MHA 对比早期 Transformer 模型普遍使用多头注意力每个 Q 头都对应独立的 K 和 V 头。以 32 个 Q 头为例KV 头也是 32KV Cache 会相当大。后来出现 GQA即分组查询注意力。多个 Q 头共享一组 K 和 V。比如 32 个 Q 头对应 8 个 KV 头每个 KV 头服务 4 个 Q 头。由于 KV 头数量从 32 降到 8KV Cache 直接降到原来的四分之一。这也是为什么很多大模型推理框架对推理性能如此敏感MHA 和 GQA 的缓存占用差异非常大。注意力类型Q 头数KV 头数相对 KV Cache 大小典型场景MHA3232基准早期模型精度较高缓存膨胀快MQA3211/32缓存最小表达力受限GQA3281/4兼顾缓存和表达力常见主力方案2.3 显存预算为什么要单独看 KV Cache很多人计算显存时只算模型权重。一个 7B 参数模型以 BF16 加载需要约 14 GB 显存。把优化器状态、激活值、KV Cache 算进去之后实际占用会明显更高。按上面配置一个 7B 级别模型在 32K 上下文、batch1 时大约需要 4 GB 的 KV Cache。如果再开多并发KV Cache 会成为显存里的主要增量。这也是生产环境中经常遇到“模型加载正常跑几个并发请求后 OOM”的原因。模型权重是固定的KV Cache 才是动态变化的部分。动态意味着需要专门的管理机制预分配、按需分配、上限控制、淘汰策略。3. 预填充阶段与生成阶段的差异决定缓存策略KV Cache 不是只跟生成阶段有关。预填充阶段也在计算 KV但它计算出来的结果是整个输入序列的 KV Cache 的初始部分。为什么预填充阶段不能直接用缓存为什么它计算量巨大需要结合这两个阶段的计算特征来看。3.1 预填充阶段并行计算完整输入的 KV预填充阶段输入的是用户提示词和系统提示词这一段通常在几百到几千 token。因为 token 之间还没有生成关系模型可以并行处理这批 token一次性得到它们的 K 和 V。预填充结束时模型已经为这些 token 建好了 KV Cache 的初始段。这个阶段的计算特点和生成阶段不同。预填充阶段的矩阵乘法和注意力计算是高度并行的GPU 利用率通常较高。生成阶段一次只计算一个 token并行度低GPU 利用率反而容易掉下来。这也是为什么很多推理框架会优先优化预填充阶段的算子融合和显存分配。3.2 生成阶段逐 token 追加 KV Cache进入生成阶段后每生成一个新 token需要执行一次完整的前向计算。当前 token 在各层计算出 Q、K、V 后K 和 V 会追加到该层已有的缓存末尾然后用 Q 去查询包括历史缓存和当前 token 的完整 K。这样计算量从“重新处理所有历史 token”变成“只处理当前 token 访问历史缓存”。整个生成阶段对 KV Cache 的访问模式是典型的追加写和顺序读。追加写意味着缓存大小是动态增长的顺序读意味着显存带宽和缓存管理策略会直接影响速度。生成阶段要循环很多次以输出 512 个 token 为例模型要执行 512 次前向计算每次都追加一次 KV。3.3 为什么“不缓存 Q”不会丢失信息有人会问当前 token 的 Q 不缓存后面的 token 查询不到它怎么办实际上后面 token 查询时当前 token 已经作为历史 token 存在于 KV Cache 中。当前 token 自身的 Q 只作用于它作为查询方的这一次计算。当它作为被查询方时它提供的是 K 和 V不是 Q。这正是自回归注意力中“当前 token 只能关注它自己以及它前面的 token”这一单向结构的体现。因此缓存的设计逻辑和注意力计算逻辑是严格对应的。缓存只保存后续 token 还需要复用的部分当前 token 已经执行完的查询部分不需要保留。4. 为什么说 KV Cache 直接影响并发、吞吐和成本在线推理服务想要高吞吐就要尽可能把多个请求放到同一个 batch 里并行处理。但 KV Cache 会吃掉大量显存限制 batch 大小。KV Cache 虽然看不到业务日志却能决定服务能同时支撑多少请求。4.1 显存被 KV Cache 挤压后并发能力下降一个 7B 模型在 BF16 下模型权重占约 14 GB。假设单卡 80 GB剩余约 66 GB 可用于 KV Cache 和激活值。如果每个请求的 KV Cache 在 32K 上下文下占 4 GB那么最理想情况下也只能容纳 16 个左右的长上下文请求。这里的瓶颈不是算力而是显存。如果把上下文长度压到 4KKV Cache 占用会缩小到原来的八分之一同样显存能容纳的并发请求数大幅提高。上下文长度和并发之间是强相关的权衡。这也是 KV Cache 计算在容量规划里非常重要的原因。4.2 缓存命中率对 token 复用的意义多轮对话中用户追问“上面提到的缓存机制再展开讲讲内存版本”。如果系统把完整历史对话都重新发到模型每次请求都要重复计算历史回答对应的 KV。使用前缀复用后模型只需计算追加的新 token历史部分的 KV Cache 直接从缓存中读取。这里的 token 复用逻辑可以这样理解同样的历史上下文第一次计算后缓存在显存里后续请求如果包含相同前缀就可以跳过前缀的预填充计算。所以缓存命中率越高无效预填充越少服务的有效吞吐越高。这也是 vLLM 等推理框架越来越重视前缀缓存和自动前缀缓存的原因。4.3 请求越分散缓存收益越低前缀缓存命中的前提是请求之间存在公共前缀。同一个知识库问答系统里如果所有请求都共享固定的系统提示词前缀命中率会很高。如果每个请求都是独立业务前缀差异很大缓存命中率就会很低。所以生产环境做 KV Cache 优化前先要看请求分布。不能只看模型支持多长上下文要看实际业务里公共前缀占比。缓存命中率低时再好的缓存框架也发挥不出作用。5. 从零实现一个可运行的最小 KV Cache 计算流程为了把上面这些原理串起来可以写一个非常小的模拟案例。这个示例不追求真实 Transformer 的完整精度只用来展示“历史 token 的 K、V 会被保存当前 token 只计算自己的 Q、K、V 并复用缓存”这个过程。代码用 Python 和 NumPy 实现适合本地快速跑通。5.1 准备最小环境推荐使用 Python 3.9 或以上版本安装 NumPy 即可。如果本机没有 NumPy执行pip install numpy这个环境不需要 GPU也不依赖任何大模型权重。目的是把 KV Cache 的追加逻辑和注意力计算拆开观察。5.2 最小模拟代码写入缓存再查询import numpy as np def softmax(x): exp_x np.exp(x - np.max(x, axis-1, keepdimsTrue)) return exp_x / np.sum(exp_x, axis-1, keepdimsTrue) class KVCacheDemo: def __init__(self, num_layers2, num_heads2, head_dim4): self.tokens [] self.layer_cache { layer: {k: [], v: []} for layer in range(num_layers) } self.num_layers num_layers self.num_heads num_heads self.head_dim head_dim def project(self, token_id): # 简化把 token 映射成一个向量再拆分为 Q、K、V rng np.random.default_rng(token_id 1000) base rng.normal(size(self.num_heads, self.head_dim * 3)) qkv np.split(base, 3, axis-1) return qkv # q, k, v def add_token_and_generate(self, token_id): self.tokens.append(token_id) q, k, v self.project(token_id) pos len(self.tokens) - 1 logits_by_layer [] for layer in range(self.num_layers): # 新 token 的 K、V 追加到缓存 self.layer_cache[layer][k].append(k[layer]) self.layer_cache[layer][v].append(v[layer]) K np.array(self.layer_cache[layer][k]) V np.array(self.layer_cache[layer][v]) # 当前 token 只用自己这一层的 Q 去查缓存 attention_scores np.sum(q[layer] * K, axis-1) attention_scores attention_scores / np.sqrt(self.head_dim) weights softmax(attention_scores) attended np.sum(weights[:, None] * V, axis0) logits_by_layer.append(attended) # 合并各层结果并输出一个模拟分数 final_score float(np.mean([np.sum(logit) for logit in logits_by_layer])) return final_score kv_cache KVCacheDemo(num_layers2, num_heads2, head_dim4) for token_id in [1, 2, 3]: score kv_cache.add_token_and_generate(token_id) print(f处理 token {token_id}当前缓存 token 总数 {len(kv_cache.tokens)}模拟输出得分 {score:.4f})这个示例里每处理一个新 token代码会先把它当前这一层的 K 和 V 追加到缓存列表K 和 V 涵盖所有历史 token。历史 token 不需要重新计算 K 和 V它们已经在缓存里。当前 token 只用自己的 Q 去点积缓存里的所有 K再按权重聚合 V。如果不用缓存每次处理新 token 时都要从历史 token 开始重新执行 project 和注意力计算。示例很小看不出性能差异但把 token 数量放大到几千差异就会非常明显。5.3 验证缓存逻辑是否成立运行后可以看到每处理一个 token缓存 token 总数增加 1。检查layer_cache里的k列表它们的长度始终等于当前已处理 token 数。这就是“历史 K 一直存在当前 K 不断追加”的最小体现。在实际 Transformer 推理中中间还会经过 LayerNorm、残差、FFN、RoPE 位置编码等过程KV 也不是原始的投影结果但缓存的核心结构就是每个 Transformer 层维护一份不断追加的 K、V 列表。这个最小示例帮助你建立直觉而不是代替真实推理框架。6. 生产环境里的 KV Cache 管理内存分配、前缀缓存和淘汰策略推理框架不会简单地把 K、V 放进 Python 列表。生产环境要把显存当资源池管理需要预分配、追踪、复用和淘汰。6.1 PagedAttention 和显存碎片vLLM 提出的 PagedAttention 是 KV Cache 管理中比较关键的设计之一。传统做法为每个请求预先分配最大长度的连续显存长度不确定时很容易出现内部碎片和预留浪费。PagedAttention 把 KV Cache 划分成固定大小的块类似操作系统的分页。请求需要缓存时按块分配不要求物理连续从而减少碎片。这个设计对实际生产的意义很直接同样是 100 GB 显存按最大长度预分配可能只能同时处理几十个请求按块按需分配可以承载更多并发请求能有效提升吞吐。6.2 前缀缓存的命中判断只看 token id 序列前缀缓存的实现通常把已计算的 token 前缀哈希后作为索引。哈希键一般包括模型标识、token id 序列、以及影响计算的配置。当新请求到达时如果它的前缀和历史请求的前缀一致就能复用历史前缀对应的 KV Cache。生产环境里要提升命中率通常有三种手段设计固定的系统提示词并放在最前面避免每条请求前缀不同限制自由对话中随意插入随机前缀调整服务端的缓存池大小和淘汰策略。命中率低时可以先看请求前缀的重复度而不是盲目加大显存。6.3 淘汰策略不能只看 LRUKV Cache 缓存池总大小有限。达到上限后要淘汰旧块常见的策略有 LRU、LFU 以及更细的按块维度的策略。但要注意KV Cache 的“块”和传统缓存对象不同一个块的后续 token 依赖前面的 token。淘汰一个中间块会导致后续块的 prefix 不再完整可能引发级联失效。因此生产级框架会结合前缀树或哈希链来做淘汰不能简单按访问时间排序。缓存维度考虑因素常见策略显存分配碎片率、对齐、预分配比例PagedAttention、块池预分配前缀命中token 前缀重复度、系统提示词稳定性自动前缀缓存淘汰策略长前缀连续性、块依赖LRU、LFU、前缀树淘汰并发管理KV 显存上限、最大并发数调度器限流、排队6.4 学习环境与生产环境的差异学习环境里KV Cache 主要用来理解原理单条请求、短上下文显存压力不大。开发调试时可以观察每个请求的 KV Cache 大小日志确认长上下文消耗。测试环境可以模拟真实请求分布验证前缀缓存命中率。生产环境要考虑的问题更多配置外置化把最大上下文长度、显存比例、缓存池大小放到配置中心。监控统计每请求 KV Cache 峰值、缓存命中率、显存碎片率。安全限制单用户最大上下文长度防止恶意长对话占满显存。回滚推理框架版本升级会改变缓存格式或分配策略要有回滚方案。异常处理显存不足时不能让请求直接崩溃要降级为重新计算或排队等待。7. 常见问题排查从现象定位 KV Cache 相关问题KV Cache 相关故障不一定直接报“KV Cache 失败”更多时候表现为显存不足、响应变慢、并发上不去、首次返回慢但后续稳定。下面按排查顺序列出典型问题。7.1 现象一请求量稍大就 OOM现象单条请求正常并发增加到某个阈值后进程报 CUDA OOM。检查顺序看模型权重和 KV Cache 各自占用确认瓶颈。计算单请求 KV Cache 上限公式为2 * max_seq_len * num_layers * num_kv_heads * head_dim * 字节数。查看当前批量大小。长上下文下batch 略增KV Cache 会成倍增长。检查是否对单请求最大生成长度做了限制。解决方案降低最大上下文长度或限制单用户生成长度。使用 GQA 或 KV 量化降低单元素字节数。使用 PagedAttention 或块式分配减少碎片。增加显存或减少并发必要时加入限流排队。7.2 现象二响应延迟整体变慢但 CPU 和 GPU 占用率不高现象模型没到显存上限但响应速度明显下降GPU 使用率不稳定。检查方向确认是否在生成阶段频繁重建完整 KV而不是复用缓存。查看请求之间前缀是否一致。前缀不一致时缓存命中率低。检查缓存池是否频繁淘汰。池太小会导致大段缓存不断被清除后续请求被迫重新预填充。查看是否使用了预填充和生成并行调度。长预填充请求会抢占资源。处理建议固定系统提示词顺序让更多请求共享前缀。调大缓存池或调整淘汰阈值。把长上下文请求和短上下文请求拆分到不同服务或实例。7.3 现象三首 token 延迟很高之后生成速度正常现象第一次响应等待时间长后续 token 生成较快。原因预填充阶段需要处理完整输入序列输入越长首 token 延迟越高。这是正常现象但如果持续偏高可能是预填充算子或输入长度控制不佳。处理建议缩短用户输入或限制上下文长度。使用前缀缓存避免相同历史对话反复预填充。对预填充阶段启用连续批处理和算子优化。7.4 现象四多轮对话越长后续请求越慢现象多轮对话进行到中后段时每次追问都需要较长等待。原因如果框架没有启用前缀缓存每轮都要把完整历史重新预填充。启用前缀缓存后只有新追加的 token 才需要真正计算。排查方式查看服务日志中的 prefill token 数和 decode token 数。对比相同对话第二轮和第三轮的 prefill 耗时。确认缓存池是否足够容纳历史对话的 KV。8. 面试常问的 KV Cache 问题与可复用清单把常见问题整理成清单既方便面试复习也方便日常做代码审查和容量规划时对照。8.1 高频面试题面试题关键回答方向为什么需要 KV Cache避免每生成一个 token 就重复计算历史 token 的注意力为什么只缓存 K 和 V历史 token 只作为被查询方不再使用 Q 做查询KV Cache 占用如何计算层数、序列长度、KV 头数、头维度、字节数相乘GQA 为什么省显存多个 Q 头共享 KV 头KV 头数量减少预填充阶段和生成阶段的差异预填充并行处理生成阶段逐 token 追加显存不足怎么优化限制上下文、KV 量化、GQA、PagedAttention、限流前缀缓存是什么复用相同 token 前缀的 KV Cache避免重复预填充KV Cache 淘汰要注意什么大块缓存有顺序依赖淘汰中间块可能导致级联失效8.2 容量评估清单在下发大模型推理服务前先按下面几步做评估确认模型层数、KV 头数、hidden_size、最大上下文长度。确认权重精度和 KV Cache 精度BF16 每个元素按 2 字节计。计算单请求在最大上下文长度下的 KV Cache 占用。估算目标并发下 KV Cache 总量。为激活值、临时缓冲和框架自身开销预留 20% 到 30% 显存余量。根据业务请求前缀重复度判断是否值得配置前缀缓存。设置单用户上下文上限和最大并发数避免极端请求占满显存。上线后监控 KV Cache 峰值、命中率、显存碎片率和 OOM 次数。8.3 坑点汇总第一个坑只关心模型参数量忽略 KV Cache。7B 模型权重看着不大但长上下文、高并发下KV Cache 会远超权重成为显存主要瓶颈。第二个坑开启前缀缓存后不关注请求分布。请求前缀完全随机时缓存命中率可能极低加显存也提升不了命中率。要先规范公共前缀。第三个坑用连续显存预分配模拟 KV Cache。生产环境应使用按块分配或等价机制否则长尾请求会让显存碎片快速上升短请求也会白白占用大块预留空间。第四个坑只测首 token 延迟不测持续生成长文本。长文本生成时 KV Cache 会持续增长可能到生成中后段才出现 OOM。压测必须覆盖完整生成流程。第五个坑没有监控缓存命中率和淘汰次数。缓存是否生效、是否频繁重建必须有日志或指标反馈否则性能问题很难定位。9. 扩展方向KV Cache 量化、投机采样和上下文工程KV Cache 的优化方向已经不只是“加显存”或“加缓存”。在工程落地时有几条路径值得深入了解。9.1 降低 KV Cache 精度KV Cache 可以使用 8-bit 或 4-bit 量化存储。量化后单元素字节数从 2 字节降到 1 字节或 0.5 字节显存占用大幅下降。代价是可能带来轻微精度损失。实际是否可用要看业务对输出质量的敏感度以及具体模型对 KV 量化的兼容程度。落地前要在自己的数据集上做对比实验不能只看显存节省。9.2 投机采样减少生成步数投机采样用一个小模型先草拟多个 token再用大模型一次验证。验证过程中如果可以一次确认多个 token实际需要的解码步数会减少KV Cache 的追加数量也随之下降。这样做能降低延迟也不改变最终输出分布。实现复杂度较高但收益可观。9.3 上下文管理减少无效历史业务层面对上下文做压缩、摘要、裁剪能够减少进入模型的 token 数量从源头降低 KV Cache 需求。固定系统提示词、定期总结老对话、限制历史轮数都是可落地的工程手段。表面看是提示词工程实际直接影响 KV Cache 总量和缓存命中率。9.4 从面试到生产的一条主线如果把 KV Cache 相关的知识串成一条主线可以这样理解Transformer 生成是自回归的历史 token 需要保留中间结果。保留 K 和 V不保留 Q是按注意力计算规律推导出来的最小必要性。KV Cache 占用和模型层数、上下文长度、KV 头数、量化精度直接相关。高并发场景下显存是 KV Cache 的主要限制。前缀缓存、PagedAttention、KV 量化都能缓解显存和命中率问题。性能优化必须配合监控和容量规划否则很难定位瓶颈。实际项目里不建议一上来就追求复杂优化。先跑通基础推理做好显存统计再根据监控数据决定是加深上下文限制、做前缀缓存、上 KV 量化还是调整并发策略。把 KV Cache 的账算清楚推理服务的容量和成本问题就已经解决了一大半。