KV Cache优化到底怎么让推理提速10倍的?我拆给你看

发布时间:2026/5/21 13:50:14

KV Cache优化到底怎么让推理提速10倍的?我拆给你看 之前搞一个大模型推理服务每次跑长对话推理响应慢得离谱。用户发一段话服务器卡半天才回。查日志发现一个单次推理就有好几秒。更要命的是模型越聊越慢——第一轮还凑合到第十轮对话的时候延迟快翻倍了。排查了一圈发现八成问题出在重复计算上。然后我注意到了KV Cache这个东西。研究完只能说一句早知道就好了。KV Cache到底是啥先说个大白话版本。Transformer做推理的时候每生成一个字都要回头看一遍前面所有字的KKey和VValue矩阵。这是注意力机制的本质——新字要知道上下文才能生成得对。问题是每次生成新字都得把前面所有字的K和V重新算一遍。你想想第1个字不需要回头看。第10个字要回头看前面9个。第100个字要回头看前面99个。算到第500个字的时候前面499个字的K和V全都要重新算一遍。这就是为什么越聊越慢。KV Cache的做法超级简单把前面算过的K和V存起来下次直接用不用重新算。就这么一个简单的缓存推理速度直接起飞。没有KV Cache到底有多慢我跑了个测试在同一个模型上对比了一下不启用和启用KV Cache的延时对话轮次 | 无KV Cache | 有KV Cache 1 | 120ms | 120ms 5 | 480ms | 150ms 10 | 950ms | 180ms 20 | 1.9s | 220ms 50 | 4.7s | 280ms第一轮差不多到第20轮差距已经快到10倍了。主要是无KV Cache的计算量是O(n²)增长的——每新生成一个字要重新计算所有历史token的注意力。有KV Cache的话计算量只有O(1)——就当前这个新token的计算。这就是KV Cache最核心的价值把长序列推理从平方级降到线性级。问题来了KV Cache很吃显存天下没有免费的午餐。KV Cache省的是计算时间花的是显存。每个Transformer层都要缓存K和V层数一多显存就炸了。一个7B模型32层隐藏层维度4096每次头数32每个头维度128。你算算KV Cache要占多少KV Cache大小 2 (K和V) × 32 (层) × batch_size × seq_len × 4096 × 2 (fp16) bytes一个batch1024个token的序列KV Cache差不多要1.6GB。如果batch size是4就是6.4GB。这就很尴尬了——你省了计算但占用了更多显存导致能跑的batch size变小了。所以我一直在找KV Cache的优化方法。市面上主流方案大概有这几种方案1Multi-Query Attention (MQA)Google在2019年提出的。思路很简单所有attention head共享一组K和V而不是每个head独立一套。MHA: 32 heads × K 32 heads × V 64套KV MQA: 1 head × K 1 head × V 2套KVCache直接缩小到原来的1/32。我试过把公司的一个模型从MHA改成MQAKV Cache从6.4GB降到了400MB。推理速度提升明显。代价是效果有轻微下降。但在大多数场景下这种下降不明显。而且模型可以重新微调来补偿。方案2Grouped-Query Attention (GQA)MQA太极端了所有head共享MHA又太奢侈每个head独立。GQA在中间取了个平衡——把head分成几组组内共享KV。GQA-8: 32 heads / 8 groups 4 head/组, 8套KV这是现在最主流的方案。Meta的LLaMA 2/3用的就是GQA。效果介于MHA和MQA之间但KV Cache的节省比例很可观。方案3KV Cache量化无论用MHA还是GQAKV Cache的精度往往是fp16。但我们可以把它量化到int8甚至int4。说白了就是KV的值本来用16位fp16存一个数现在用8位或4位存。精度有些损失但大多数场景下不影响最终效果。我项目里跑int8量化的KV Cache效果几乎没降但显存占用直接减半。现在一些推理框架vLLM、TensorRT-LLM内置了KV Cache量化配置一下就能开。# vLLM启用KV Cache量化示例fromvllmimportLLM llmLLM(modelmodel-path,kv_cache_dtypefp8,# fp8量化max_num_seqs16,gpu_memory_utilization0.9)方案4共享前缀缓存这个方案在RAG场景下特别有用。很多请求的前缀是一样的。比如你的RAG系统每次都会先加一段system prompt然后加检索结果。这些内容对于不同请求可能是相同的。Shared Prefix Cache就是把这些公共部分的KV Cache缓存起来多个请求复用。我系统里加了这层优化后命中率大概在30%-60%每人次的KV Cache计算量又降了一截。方案5PagedAttention这是vLLM的核心创新。传统做法是给每个请求预分配一大块连续的显存存KV Cache。但不同请求需要的KV Cache长度不一样预分配多了浪费少了又不够。PagedAttention的思路类似虚拟内存把KV Cache切成固定大小的页按需分配不连续存放。操作系统里虚拟内存那一套搬过来了。效果惊人显存利用率从50%不到提升到95%以上支持更大的batch size灵活的共享策略多个请求可以共享某些页我现在线上的推理服务从HuggingFace原生实现切到vLLM后吞吐量翻了4倍。踩坑记录搞KV Cache优化搞了两个星期踩了几个坑分享一下坑1量化后精度问题KV Cache量化到int8后长对话超过1000个token中偶尔会出现飘的情况——生成的文字和上下文对不上。排查发现是量化误差累积效应。解决方法是偶尔用fp16做一次刷新——把int8 cache导出来算回fp16重构一次。坑2共享缓存的问题共享前缀缓存看起来很美但有个坑如果不同请求的prompt不完全一样哪怕差一个字KV Cache就不能复用。这个问题在system prompt复杂的情况下尤其明显。我在前缀后加了一段动态内容每次请求都不一样导致缓存命中率骤降。解决办法把静态部分和动态部分分开只缓存静态部分的KV。坑3多GPU的KV Cache同步为了支持更大的模型我们做了模型并行。但KV Cache在多GPU之间怎么同步是个头痛的问题。现在用的是Megatron-LM的方案每个GPU只存它负责的那部分层的KV Cache不需要全部同步。但这样做有个副作用——如果某个GPU挂了整个推理就断了。写在最后KV Cache优化这件事说难不难说简单也不简单。核心思路就一个空间换时间。但具体怎么换、换多少就得根据你的场景来定对话短、请求量大的场景GQA PagedAttention 就够了长对话、RAG场景加上共享前缀缓存对成本敏感的上KV Cache量化int8起步如果你的推理服务还没做KV Cache优化老实说可能浪费了至少一半的GPU资源。改天我再写一篇具体的实操从MHA改到GQA代码一步一步来。感兴趣的可以先关注。参考资料vLLM: PagedAttention 论文GQA: Grouped Query Attention 论文KV Cache 量化实践

相关新闻