1. 先搞懂KV Cache到底在解什么问题做过大模型推理的朋友都知道自回归生成是一个非常“笨”的过程模型一次只能吐出一个token然后把新生成的token拼到已有序列后面再继续预测下一个token。这个过程本质上是一个无限长的循环而每一步都要对整个输入序列做一次完整的前向计算。如果不做任何缓存优化第1个token要算1次第2个token要算2次第3个token要算3次……等到生成第1000个token时前面999个token的Key和Value都要重新算一遍。时间复杂度近似O(n^2)到O(n^3)序列稍微一长GPU就被拖垮了。KV Cache解决的就是这个重复计算问题。Attention机制里每个token会生成Query、Key、Value三个向量。在计算当前token对历史token的注意力时真正需要的是历史token的Key和Value而不是历史token的原始输入。既然Key和Value在一次前向计算后已经拿到了把它们暂存在显存里后续生成阶段直接复用就能把每次decode的时间复杂度从O(seq_len)压缩到O(1)单步或O(seq_len)增量计算。这就是KV Cache最初的动机用显存换算力。但KV Cache不是免费的午餐。它最大的代价就是显存占用会随着序列长度、并发请求数、模型层数线性增长。我见过很多项目模型权重只有十几GB但跑到几百并发、上下文稍微长一点显存直接爆掉结果不是算力不够而是KV Cache把显存吃光了。这个问题在长上下文和持续对话场景里尤其突出甚至比模型权重本身更让人头疼。所以理解KV Cache的工作原理、显存计算方式以及围绕它的各种优化手段是做推理加速绕不开的一课。1.1 自回归生成中的重复计算是如何发生的先看一个具体的计算流程。假设现在有一个Transformer模型输入序列长度为L模型有N层。在做一次完整的前向传播时每一层都会对输入序列中的所有token计算Q、K、V然后做Attention。如果输入序列是整段文本比如Prompt这部分的Attention计算是不可避免的因为每个token都需要和所有其他token交互我们把这一阶段叫作Prefill预填充阶段。Prefill阶段并行度高因为所有token的输入是已知的可以一次算出所有位置的K和V。问题是生成阶段。模型每生成一个新token新的Q和新的K/V只来自新token这个位置。但Attention计算需要当前Q与所有历史token的K做点积。如果没有缓存就必须把整个历史序列重新过一次前向网络重新生成每个token的K和V这样算出来的结果虽然和缓存计算在数学上完全一致但重复做了大量无用功。假设已经生成了M个token那么到第M1步时前面M个token的K和V都是可以通过历史输入重新计算出来的这种“回溯式”计算让每一步的代价都随M增长。生成到几百token以后单步延时会越来越明显用户感知就是模型越说越慢。有了KV Cache之后流程变成这样Prefill阶段正常计算完整输入序列的K和V全部存到显存里Decode阶段每来一个新token只计算这个token自己的K和V然后追加到缓存末尾同时使用缓存里的历史K/V与当前Q算Attention。这样每一步的计算量只跟前缀长度相关不会因为历史越来越长而重复计算整个前缀。从工程角度说KV Cache是一个典型的增量数据结构它把完整的Attention计算过程拆成了可复用的中间结果换来的是显存和调度复杂度的上升。1.2 KV Cache的显存占用公式与实例推算要怎么估算KV Cache占多少显存公式并不复杂但很多人凭感觉拍脑袋最后在线上的时候OOM。先给一个通用公式KV Cache显存大小字节 2 × batch_size × sequence_length × num_layers × num_kv_heads × head_dim × precision_bytes公式里的2代表Key和Value两份。注意这里用的是num_kv_heads而不是num_attention_heads因为如果模型用了GQA或MQAKV head数量会远小于Q head数量。如果模型是原始的MHA结构那num_kv_heads就等于num_attention_heads。举一个具体例子。假设一个模型有32层KV head数是8head_dim是128每个token的KV占8 × 128 × 2 × 2字节假设是FP16 8KB。那4096个token的KV Cache就是32MB。看起来不多但如果同时服务32个并发请求每个请求上下文都是4096KV Cache直接到1GB。如果把context拉长到32K单请求的KV Cache是256MB10个并发就能到2.5GB。再算上中间激活、模型权重、CUDA context显存很快就吃紧了。更直观的对比是模型权重。一个7B参数的FP16模型权重占14GB。但如果支持128并发、4K上下文KV Cache可能轻松超过15GB。也就是说在在线推理服务里KV Cache往往比模型权重占用更高。这也是为什么现在很多推理框架都在KV Cache的分配和复用上做了大量文章因为这里是性价比最高的优化空间。2. 围绕KV Cache的优化手段该怎么选当你理解了KV Cache的本质就会发现优化思路其实分三个层面第一是让缓存本身更小第二是让缓存分配更高效第三是让缓存精度更低但效果可接受。这三个层面不是互斥的实际项目中通常叠加使用但顺序和力度要根据场景来定。先说一个基本判断。如果只是本机跑个demo、生成不追求高并发KV Cache优化意义不大PyTorch默认的缓存逻辑就够了。但只要上生产、做对外服务或者部署在T4、A10这类显存较小的卡上KV Cache优化就是刚需。我个人做优化时的优先级是先算清显存账再开分页分配接着考虑量化最后看模型架构是否要换成GQA版本。架构的改动成本最高一般放在模型选型阶段就定下来运行时优化才是日常手段。2.1 架构瘦身从MHA到GQA/MQA的取舍原始Transformer用的都是多头注意力MHA每个注意力头都有独立的一组K/V。比如32个Q头那就有32组K/V。这样表达能力强但KV Cache开销也大。后来出现MQA多查询注意力让所有Q头共享一组K/VKV Cache直接降为原来的1/32。再后来GQA分组查询注意力作为折中方案把Q头分成若干组每组共享一组K/V既能保留一定的表达能力又能显著压缩KV Cache。GQA是目前最主流的中间态。它有一个关键参数叫query-to-kv-head-ratio比如Q头32个、KV头8个ratio就是4:1。GQA压缩后KV Cache占用降为原来的1/ratio。在解码阶段每个KV头会被多个Q头复用计算时需要用向量广播但实际计算量并没有额外增加太多。这里有一个值得注意的细节GQA的模型在量化KV Cache时由于KV头数量本来就少每个缓存块被频繁访问量化误差可能会被放大所以要做量化评估时尤其要关注这类模型。MQA现在用得少了因为极端压缩下输出质量会有损失。GQA是个兼顾质量和性能的选择。很多开源模型在7B级别就开始用GQA说明已经在为推理效率妥协。如果是自己训练模型务必要在训练阶段就决定attention结构不要在训练后用转换脚本强行合并K/V容易引入质量损失。2.2 运行时显存管理从连续缓冲到PagedAttention的思路传统框架里KV Cache通常按最大序列长度预分配一整块连续显存。一个请求如果最多支持2048个token那每个token的位置都预留给它。缺点显而易见请求实际只生成200个token时剩下1800个token的缓存就白白被占用。另外显存碎片化严重因为不同请求的长度是动态变化的预分配方式很难做到灵活动态扩展。PagedAttention的思路和操作系统虚拟内存管理很像。把KV Cache拆成固定大小的物理块block每个block能容纳固定数量的token比如16个token。逻辑上每个请求有一个连续的KV Cache“虚拟空间”但物理上这些block可以分散在显存的各个位置通过一个block table来映射。当请求长度超过当前已分配的block时只需要额外申请一个block而不是移动整块数据。这样既避免了内部碎片又实现了真正按需分配。这个设计的另一个好处是支持前缀共享。很多多轮对话场景里多个请求共用同一段系统提示词或历史上下文如果这些KV Cache block能被多个请求共享显存占用能省掉一大截。我在实际项目中跑过一组对比开PagedAttention block共享后相同显存下并发能力能提升3倍以上。代价是block table查询会带来少量CPU开销但整体收益远大于损耗。2.3 数值压缩KV Cache量化的几种主流方案KV Cache量化是近两年很热的加速方向思路是把缓存中的Key和Value从FP16压缩成更低精度比如INT8、FP8甚至INT4。因为KV Cache远大于权重压缩它带来的显存收益非常直接。先讲最简单的INT8量化。每个token的K/V用absmax对称量化存一个scale值运行时先反量化再参与计算。这种方式实现简单但把整段缓存用同一个scale如果某个token数值特别大其他token的量化误差会变大。于是有了per-channel量化针对每个通道分别算scale精度好一些但存储开销也变大。还有per-group量化比如128个元素一组各算各的scale精度和开销的平衡更好。现在很多推理框架都支持KV Cache量化参数通常类似kv_cache_quant_method可以选择fp8、int8、int4等还会有kv_cache_quant_scale选项。我自己实测下来INT8 KV Cache量化对大部分任务质量影响很小显存占用减半。FP8受限某些GPU硬件支持效果也不错。INT4虽然更省但长序列下容易出现质量退化特别是代码生成、数学推理这类对数值敏感的任务。还要注意KV Cache量化并不是每一层都适合。部分推理框架支持按层跳过量化比如前几层或某些attention层保留高精度后面用低精度。这个方法能在质量几乎不掉的情况下进一步压显存。3. 实操给一个对话推理服务做KV Cache优化这一部分记录一次完整的优化过程基于一个模拟项目某公司的客服对话机器人使用7B模型推理卡为单张24GB显存目标是支持100并发、平均上下文长度约2K单token延迟控制在80ms以内。初始实现用的是最朴素的transformer推理逻辑一上压测就OOM或者并发一上去就超时。下面是我按步骤做的优化每一步都可以直接对应到你的项目。3.1 第一步建立显存占用基线任何优化都不能靠猜。第一步就是把KV Cache的显存占用模型拉出来和实际观测值做对比。我写了下面这个简单脚本用于估算峰值KV Cache占用。def estimate_kv_cache_bytes( batch_size, seq_len, layers, kv_heads, head_dim, dtype_bits16, groups1, ): # 每个KV元素占dtype_bits/8字节Key和Value两份 per_token_bytes 2 * layers * kv_heads * head_dim * dtype_bits // 8 total_bytes batch_size * seq_len * per_token_bytes total_gb total_bytes / (1024**3) return round(total_gb, 2) # 示例7B模型32层8个KV头head_dim128FP16100并发2K上下文 print(estimate_kv_cache_bytes(100, 2048, 32, 8, 128, 16)) # 输出16.0 GB这个模型权重FP16约14GB光KV Cache就16GB24GB的卡显然放不下。从这一步就能得出结论必须做优化而且不能只做一种。我在真正压测前还会用推理框架自带的监控接口看实际的KV Cache峰值和OOM时的报错日志把理论估算和实际观察对齐。这里还有一个容易被忽略的点GPU显存不是只有模型权重和KV Cache还有CUDA context、中间激活值、临时缓冲区。通常要预留10%到20%的显存给这些开销。按24GB卡算可用KV Cache空间可能只有15~18GB。如果照满算必然OOM。我一般用保守系数0.8乘以总显存再减去权重余量。多轮迭代后把每个环节的实际占用都记录下来后面调参数才会有依据。3.2 第二步启用分页式KV Cache并调整调度参数基线确认后开分页式KV Cache。以某开源推理服务框架为例支持类似PagedAttention的机制需要改几个关键配置。首先是把推理引擎切换到支持分页分配的版本同时将GPU显存利用率从默认值调低一点比如gpu_memory_utilization0.85给中间激活和临时数据留出余量。然后是max_num_seqs最大并发序列数和max_num_batched_tokens单次批处理的最大token数。刚开始我习惯把max_num_seqs设成目标并发数但发现这样容易把显存压满。后来调整为先算出一批请求的KV Cache预估占用再折算并发上限。还有一个参数是block_size也就是每个物理块装多少个token。默认值是16通常够用。block_size太小调度开销高太大又增加内部碎片。如果场景是短上下文、高并发block_size设成8可能更省显存如果长上下文为主16或32会更好。这个参数没有标准答案最好跑一组A/B对比。我在模拟项目中对2K上下文场景试了block_size8和16两种都开了量化后16的吞吐略高一点显存峰值差不大最终保留16。启用分页后原有的连续预分配方式要关掉否则两套逻辑会冲突。我在切换时遇到过明明开了分页显存占用还是按最大长度预分配的情况后来发现是框架里上游参数仍指定了max_model_len且未配套新的显存估算逻辑。把max_model_len适当调小到业务实际需要的最大值比如2048而不是4096或8192显存压力立刻小了一大截。不要觉得设大了没坏处KV Cache是线性的长度上限直接决定最坏情况下的显存占用。3.3 第三步对KV Cache做FP8或INT8量化在分页分配已经缓解显存碎片的前提下再上量化。这个模拟项目里我先后测了FP16、FP8、INT8和INT4四种精度。先解释一下框架里的量化参数。对于KV Cache通常有类似kv_cache_dtype的选项可设为auto、fp8或int8。如果选择int8还可以指定量化粒度和scale计算方式常见选项为per_channel、per_group、per_token。我一般默认推荐per_groupgroup size取128或64精度和存储的平衡都比较理想。实操时我写了个脚本先做质量回归避免只看显存不看效果。用一组固定问题集分别在这四种精度下跑相同Prompt生成200~500个token比较生成结果的ROUGE-L分数和人工评分。结果如下KV Cache精度KV Cache显存占比相对FP16生成质量单Token平均延迟FP1616GB1.055msFP88GB0.9952msINT8 per-group8GB0.9850msINT4 per-group4GB0.9147msFP8和INT8的表现非常接近质量损失都在可接受范围内显存直接减半。INT4虽然显存更省但连续生成超过300个token后部分回答开始出现重复和逻辑断裂质量退化明显。最终在这个项目中选了INT8 per-group量化配合分页分配后100并发的KV Cache占用降到了7.1GB左右单卡承担100路请求已经没问题。这里有几个实操注意事项。第一量化后的KV Cache在解码时可能存在反复反量化开销实际延迟下降没有显存收益那么明显需要靠批处理来摊薄。第二如果模型层数很多不要一次性把所有层都量化。我后续实验发现把最后4层的KV Cache保持FP16其他层用INT8显存只增加约0.5GB但长文本回答的稳定性好了很多。原因是靠近输出层的注意力模式对细节更敏感量化误差会被放大。3.4 第四步结合模型架构替换做进一步压显存分页和量化做完之后KV Cache占用从16GB降到约7.1GB已经很可观了。但如果还想继续压就要从模型架构入手。这个7B模拟模型其实用的是MHAQ头和KV头都是32个。我在对比测试中换了一个结构基本一致但采用GQA的7B模型KV头数为8个KV Cache占用理论上是原来的1/4。实际跑下来在相同batch和上下文长度下使用GQA模型后KV Cache从7.1GB进一步降到约1.8GB。此时KV Cache不再是瓶颈单卡可以支撑400路并发以上单Token延迟受批处理影响略有上升但整体吞吐提升显著。质量方面用同一组测试集对比MHA版和GQA版在常见对话任务上几乎没有可感知差异只在某些需要长链条推理的任务上GQA略有下降。考虑到客服场景对实时性和并发的需求这个取舍完全可接受。架构替换的工程成本比前面两步高因为要重新训练或微调模型不是改配置就能完成。所以在项目规划时就要想清楚是将来的扩展阶段再换GQA还是一开始就用。对一个新项目我的建议是直接用GQA模型前提是下游任务质量能满足要求。如果是老项目可以先通过分页和量化顶一阵把架构升级放到迭代计划里。4. 常见问题与排查技巧实录这几个月做推理加速遇到的问题比想象中多尤其是KV Cache优化和显存优化纠缠在一起的时候很多参数牵一发动全身。我把实践中踩过的坑和排查思路整理一下都是框架文档里不太会写的内容。4.1 为什么开了优化还是会OOM我遇到过最典型的情况是明明开好了PagedAttention和量化压测时打到80并发就OOM。排查后发现问题不在当前请求而在于框架有两个隐藏的内存水位。第一个是max_model_len前面说过它决定每个block最多能分配多少缓存如果某个请求上下文特别长最坏情况下的显存占用会把其他并发请求的buffer挤爆。第二个是prefill阶段临时激活值在单次请求输入非常长比如几千token时激活值可能瞬间占掉好几GB。我的解决办法是做一次最坏情况估算而不是按平均长度估算。把最大上下文长度、最大并发数、最大初始Prompt长度分别乘起来看总占用是否会超过剔除权重后的剩余显存。如果会超就不要把所有buffer都设成最大值可以改成按长度分池长短请求走不同的显存池。但这个方案不是每个框架都支持所以更常见的做法是限制单请求最大输入长度比如超过2048的部分做截断或滑动窗口。还有一个很少被提到的问题CUDA kernel启动后显存Caching Allocator并不会立刻把释放的空间交还给OS而是留在缓存池里复用。如果压测前的模型权重加载方式不对比如重复热加载缓存池会被打得很碎。遇到OOM但显存监控却显示还有空闲可以试试显存碎片整理或者重启一个干净的worker进程。我在项目中就把一次性的推理服务拆成了常驻worker并定期根据监控做worker回收OOM概率明显降低。4.2 KV Cache量化后效果反而变差量化后生成质量下降最常见的原因是量化粒度太粗。我之前贪图方便用per-token的INT8量化短文本没问题但长上下文一开始模型频繁出现注意力分散。排查时发现per-token量化对每个token单独scale没有考虑到通道维度的异常值一旦某个通道的数值范围特别大量化后信息损失就会集中在某些头部导致注意力权重偏移。换成per-group后质量问题明显缓解。如果你框架不支持per-group也可以改用per-channel。我在一些老版本库上为了兼顾性能和精度还试过混合量化策略第一层和最后一层的KV Cache保留高精度中间层用低精度。这种做法在开源框架里可以用skip-layer参数实现效果不错实现成本也不高。除了量化粒度还要注意输出温度的设置。同一组测试在高温度比如1.0下量化模型表现更不稳定低温度比如0.2下差距变小。所以量化上线时建议同时在推理参数上做一个保守调整尤其对于对话系统把temperature从0.8降到0.5附近可以掩盖一部分量化带来的随机性损失。4.3 并发上去但吞吐不升反降KV Cache优化通常能提升并发但如果光开并发而忽视批处理策略可能出现“同时处理的请求很多但GPU利用率不高、延迟反而变长”的情况。原因有两类第一每个请求长度差异太大短请求和长请求混排时为了等长请求处理完短请求一直在block队列里等待造成调度阻塞第二量化后的KV Cache在反量化时需要额外计算如果kernel没有将量化操作融合进Attention计算会带来大量小kernel执行开销。针对第一点我把请求队列改成按预估长度分段短请求走低延迟通道长请求走高吞吐通道两个通道共享同一个显存池。针对第二点确实需要升级到更新版本的推理框架因为老框架对量化KV Cache的kernel融合做得差。我在模拟项目里对比过新旧版本融合了量化操作的版本在同样并发下整体吞吐高出约25%这个优化效果不亚于把KV Cache从FP16换成INT8。所以买新卡不如检查框架版本内核优化带来的收益常常被低估。4.4 前缀共享无效的处理PagedAttention支持前缀共享核心逻辑是多个请求如果拥有相同的前缀比如系统提示词、固定历史可以直接复用KV Cache block而不是各自为这段前缀分配缓存。但在我的项目里一开始并没有生效排查后发现是因为请求里带了细微差异比如一个请求在系统提示词末尾多了个空格另一个没有框架的block哈希就无法匹配共享失效。解决办法是标准化入站文本在tokenize前做一次统一的格式处理确保共享前缀完全一致。同时把prompt cache的判定阈值调低一点允许部分block可以跨请求共享。这里还要注意共享KV Cache是有代价的如果一个请求需要更新共享block会触发写时复制导致另一个请求的缓存被复制一份显存可能反而增加。所以共享不要把池子占满通常控制共享块比例在50%以下比较安全。5. 工具选型与工程落地的一些个人心得KV Cache优化虽然在理论上很清晰但落地时要依赖具体推理框架的特性。我在不同阶段用过几类工具这里不指名具体产品只讲选型逻辑。如果你是做快速验证直接用大模型库的默认缓存逻辑就好。它支持简单的缓存拼接但不会自动做分页、量化和调度。这种模式下能跑通Pipeline就比较满足不追求高并发。如果你想上生产寻找一个支持分页式KV Cache、量化KV Cache、continuous batching等特性的开源推理服务框架。判断标准很简单看它是否允许你在运行时观察每个请求的KV Cache block分配情况以及是否支持按层配置量化精度。这两个能力的支持程度基本决定了你能在这个框架上压多狠。还有一种情况是垂直场景比如只需要固定长度的多轮对话不需要超长上下文。这种场景可以绕过通用框架的复杂调度使用手写推理脚本把KV Cache预分配成一个静态数组按轮次覆盖旧数据显存效率非常高。我在做某个特殊项目时就是用这个方式把原本需要多卡的推理服务压到单卡上效果出乎意料地好。但要小心手写方案需要自己管理Attention mask和增量位置编码对调试能力要求高建议只在应急或高度定制场景使用。另外有一个工程细节我反复强调KV Cache优化绝不等于只调一个参数。所有推理框架都会提供一组和缓存、调度相关的配置它们互相影响。比如你开了KV Cache量化但block size没有跟着调小反量化计算的batchsize可能会让kernel效率下降你开了前缀共享但不做输入规范化共享命中率偏低反而增加额外开销。我的习惯是每次只改一个变量观察显存峰值、P99延迟、吞吐三个指标然后用小步迭代的方式逼近最优配置。不要试图一次把十个参数都搬到生产环境出了问题你都不知道是哪个参数导致的。5.1 显存性能测试的前置检查清单这里给出一份我自己每次做推理性能测试前都会过的检查清单照着做基本能避免低级的配置错误。确认模型权重和KV Cache量化精度是否一致避免权重FP16但缓存INT8时出现数值对齐问题。确认GPU是否支持目标精度如FP8不支持的硬件上不要强行开启会导致kernel退化和未知错误。确认max_model_len是否超过业务最长输入最长生成的汇总长度而不是只看训练时的序列长度。确认显存监控指标里有没有包含CUDA context和临时缓存没有的话额外预留5%~10%。确认压测时的请求长度分布覆盖了P50、P95、P99而不是只用固定长度测平均指标。确认服务启动时是否复用了已加载权重的缓存避免重复加载而导致显存浪费用。这些都是很基础的检查项但往往最基础的步骤出错率最高。我见过不止一次团队花了几周调KV Cache参数最后发现是上游一个size参数类型用错了导致所有请求都按最大长度走。先跑一遍基线、检查日志、确认配置确实生效再开始优化能节省大量时间。5.2 关于长上下文场景的额外建议如果你的业务是文档问答、代码仓库问答这类动辄几十K上下文的长文本场景KV Cache的挑战会比普通对话大一个数量级。此时要特别注意KV Cache显存随长度线性增长而Attention计算量随长度平方增长两者叠加起来长文本下最终瓶颈往往是计算量而不是显存。不过这不代表KV Cache优化不重要它让你的显存能承载更长上下文是撑起特殊场景的前提。长上下文场景下我还有一个建议不要把首个Prompt一次性灌进去。好的做法是先做一个信息抽取或摘要把无关的文本剪掉保留真正有用的上下文。很多团队为了“无损”把所有聊天历史全部拼进去结果KV Cache急剧膨胀延迟变长效果却没有提升。通过文本压缩和关键片段检索可以把有效上下文缩短一半以上这部分减少的显存可以直接转化成并发能力。淘汰无意义的历史比无脑扩大显存更经济。收尾的一点体会经过这几个月的反复调试我最大的感受是KV Cache优化不是单一技巧而是一整套围绕“中间结果复用”和“显存分配效率”的工程体系。一开始我以为只是把缓存打开后来发现它背后牵扯到模型架构、输入长度控制、调度策略、量化精度等多个因素的联动。每次调优都是先解决最吃显存的部分再找下一个瓶颈。把KV Cache从16GB压到2GB以内之后同一套GPU能承担的任务量完全不一样了。如果后面再做类似项目我会在模型选型阶段就明确KV Cache的预算然后反向约束并发数量和上下文长度在推理阶段则优先确保缓存分配的灵活性再考虑数值压缩。这样能在保证输出质量的前提下把硬件性价比用到极致。最后想提醒你任何优化方案上线前都要做充分的质量回归尤其是KV Cache量化最容易出现“指标好看、效果翻车”的问题。在显存和效果之间找到平衡点才是长期稳定的解法。
