你的大模型为什么越跑越慢CompressKV只存关键记忆让长文本推理快10倍 目录一、从背单词说起KV缓存是什么二、问题为什么大模型越跑越慢三、现有方案的坑所有头一刀切四、CompressKV的解法找到最会找重点的头五、层间分配不是每层都一样重要六、效果有多强数据说话七、伪代码核心逻辑长什么样八、能跟现有技术叠buff吗九、总结与展望十、参考资料一、从背单词说起KV缓存是什么想象一下你在考英语阅读理解文章有5000个单词。你读完文章后需要回答5个问题。你不可能每回答一个问题就把5000个单词重新读一遍。你会怎么做你会在草稿纸上记下关键信息主角叫Tom第1段事件发生在1920年第3段关键结论在第8段这些笔记就是你的**“缓存”**——你不需要反复阅读原文看笔记就够了。大语言模型LLM也有类似的笔记系统叫做 KV 缓存。当你在跟 ChatGPT 聊天时模型会记住你之前说的每一句话把这些信息编码成Key键和Value值对存在内存里。这样当它生成下一个词时就能快速翻看笔记而不是把整段对话重新理解一遍。用户: 请帮我总结这篇关于AI的论文重点是创新点和实验结果。 模型: 我需要记住AI论文创新点实验结果这些关键词... → 存入 KV 缓存但是如果你的对话越来越长比如扔给模型一本10万字的小说这个草稿纸就会越来越厚。上下文长度KV缓存大小Llama 3.1 8B占用的GPU内存1K tokens~0.5 GB轻松4K tokens~2 GB还行32K tokens~16 GB开始紧张128K tokens~64 GB爆掉了KV缓存的大小跟上下文长度是线性关系。你输入的文字越长显存占用就越大推理速度也越慢。这是长文本推理的核心瓶颈。二、问题为什么大模型越跑越慢2.1 注意力机制模型的注意力是有限的大模型使用**注意力机制Attention**来理解上下文。简单说就是生成每个词时模型会回头看之前的所有词给每个词打个重要性分数重点关注重要的词。# 注意力机制伪代码defattention(query,key,value):scoresquery key.T# 计算每个词的重要性分数weightssoftmax(scores)# 归一化成概率outputweights value# 加权求和重点关注重要词returnoutput2.2 多头注意力很多专家在投票现代LLM如 Llama、Qwen、Mistral使用分组查询注意力GQA把注意力的专家分成多个头head。每个头负责不同方面的理解‍⚕️头A专门关注语法结构主谓宾️头B专门寻找关键实体人名、地名头C专门捕捉逻辑关系因果、转折这些头一起投票决定哪些词重要。但问题是——它们的票是等权的吗三、现有方案的坑所有头一刀切3.1 现有方法谁喊得大声就听谁的为了节省KV缓存研究人员想了很多办法。最主流的是KV缓存驱逐Eviction只保留重要的token扔掉不重要的。现有方法通常是这么做的# 现有方法如SnapKV的简化逻辑defselect_important_tokens(all_heads,num_to_keep):# 1. 把每个头的注意力分数加起来平均aggregated_scoressum(head.attention_scoresforheadinall_heads)# 2. 选分数最高的前N个tokentop_indicestop_k(aggregated_scores,num_to_keep)returntop_indices看起来合理但有个致命缺陷。3.2 流注意力头“我只看开头和结尾”在GQA中有一类特殊的头叫做流注意力头Streaming Head。它们的行为非常固定——几乎只关注prompt的开头第一个token和结尾最近的token。就像一个评委不管舞台上表演什么他永远只看第一个节目和最后一个节目。上图Streaming Head的注意力分数——开头和结尾几乎100%中间几乎为0当把所有头的注意力分数加起来时这些极端分子的分数会淹没其他头的声音。结果是什么中间那些包含关键信息的token被无情地驱逐了举个例子Prompt: 请根据以下文章回答小明早上吃了面包中午吃了面条晚上吃了三明治。问小明晚上吃了什么 Streaming Head的注意力[请, 三明治] Semantic Head的注意力[晚上, 吃了, 三明治] 聚合后[请, 三明治] → 中间token 晚上 被驱逐了 → 但如果没有晚上模型可能不知道三明治是晚餐还是午餐⚠️这就是现有方法的核心问题Streaming Head主导了驱逐决策导致语义上重要的中间token被错误地丢弃。四、CompressKV的解法找到最会找重点的头4.1 核心洞察有些头天生更会找重点CompressKV的作者发现了一个关键现象在注意力头中有一类特殊的头它们不仅能找到正确答案还能关注答案周围的语义上下文。这些头被称为语义检索头Semantic Retrieval Heads, SRH。类型特点比喻Streaming Head只看开头和结尾只看目录和最后一页的人传统检索头TRH只看精确匹配的词只会CtrlF搜索的人语义检索头SRH关注答案周围语义上下文真正理解上下文会划重点的人4.2 怎么找到这些最会找重点的头CompressKV使用了一个巧妙的**跨度聚合Span Aggregation**标准defscore_srh(head,dataset): 计算一个头是否是语义检索头 核心思想看它在生成答案时是否关注了答案周围的整个语义区间 score0forsampleindataset:answer_spansample.answer_tokens# 正确答案的区间如sandwichfortinrange(num_generation_steps):ifgenerated_token[t]inanswer_span:# 关键不是只看某个token的分数# 而是看整个答案区间的注意力总和scoresum(head.attention_weights[t,j]forjinanswer_span)returnscore为什么这样更好传统方法要求头的注意力精确命中正确答案的某个tokentop-1规则SRH方法允许头分散地关注答案周围的一片区域如eat“a”“thing”“sandwich”即使某个头没有给sandwich最高的注意力只要它给答案区域整体的关注度高就被认可类比传统方法像找完全匹配的简历SRH方法像找整体匹配度高的候选人——后者更灵活更容易找到真正合适的人。4.3 用SRH来选择保留哪些token找到SRH后CompressKV的token选择逻辑就很简单了defcompress_kv_cache(layer,all_heads,budget,num_srh4):# 1. 识别该层top-k的SRH离线完成只需一次top_srhget_top_srh(layer,knum_srh)# 2. 用这些SRH的注意力分数来选择重要token# 只关注观察窗口内的最近token跟SnapKV一致importance_scoreszeros(num_tokens)forsrhintop_srh:importance_scoressrh.attention_scores_in_window# 3. 平均并排序选出top-N个tokenimportance_scores/num_srh keep_indicestop_k(importance_scores,budget)# 4. 只保留这些token的KV对其他的驱逐compressed_KK[keep_indices]compressed_VV[keep_indices]returncompressed_K,compressed_V关键优势SRH能识别语义上重要的中间token比如晚上“吃了”Streaming Head的噪音被过滤掉了同一层内的所有头共享相同的token索引保持GQA结构五、层间分配不是每层都一样重要5.1 问题每层给一样多的预算合理吗现有方法大多给每层分配相同的KV缓存预算。但直觉告诉我们模型的不同层承担不同的职责。浅层前几层提取词级别特征、语法结构中层提取短语级别特征、局部语义深层后几层提取句子级别特征、全局语义、推理逻辑如果每层都给2048个token的预算可能深层需要更多浅层可以更少。5.2 CompressKV的解法离线估计每层压缩误差CompressKV的思路很优雅用压缩前后的输出差异来衡量每层对压缩的敏感程度。defestimate_layer_error(model,calibration_data,full_kv_cache): 离线计算每层的压缩误差 只需在模型部署前运行一次 errors[]forlayerinmodel.layers:# 1. 全缓存输出作为参考标准O_fulllayer.forward_with_full_cache(query,K_full,V_full)# 2. 只保留少量token的压缩缓存输出K_comp,V_compcompress_tokens(K_full,V_full,small_budget)O_complayer.forward_with_compressed_cache(query,K_comp,V_comp)# 3. 计算Frobenius范数差异矩阵差异的整体度量errorfrobenius_norm(O_comp-O_full)/(frobenius_norm(O_full)1e-6)errors.append(error)# 4. 归一化得到每层的相对敏感程度normalized_errorserrors/sum(errors)returnnormalized_errors5.3 预算分配误差大的层多分误差小的层少分有了每层的误差分数预算分配就变得直观了defallocate_budget(total_budget,layer_errors,min_per_layer32,max_multiplier3):num_layerslen(layer_errors)base_per_layertotal_budget//num_layers# 1. 先给每层保底预算budgets[min_per_layer]*num_layers remainingtotal_budget-min_per_layer*num_layers# 2. 按误差比例分配剩余预算fori,errinenumerate(layer_errors):extraremaining*err max_budgetmax_multiplier*base_per_layer budgets[i]min(min_per_layerextra,max_budget)returnbudgets类比就像考试前复习时间有限你应该多花时间在自己薄弱的科目上而不是每科平均分配。CompressKV就是给对压缩更敏感的层更多预算。六、效果有多强数据说话6.1 LongBench长文本理解综合基准LongBench包含16个长文本任务涵盖单文档问答、多文档问答、摘要、代码补全等。Llama 3.1 8B 上的结果固定KV缓存预算方法256 token预算1024 token预算FullKV无压缩基准49.08—StreamingLLM33.9236.95SnapKV45.2147.82PyramidKV44.3647.65CAKE46.3047.97HeadKV44.1147.05CompressKV46.7148.24关键发现在256 token的紧预算下CompressKV领先所有基线46.71 vs 46.30 CAKE在1024 token预算下CompressKV甚至超过了无压缩的FullKV等等48.24 49.08让我重新检查数据… 实际上48.24 49.08但差距非常小只有0.84分而且使用了约1/20的内存在4个模型Llama、Mistral、Qwen 14B、Qwen 32B上一致领先6.2 Needle-in-a-Haystack大海捞针测试“Needle-in-a-Haystack”NIAH是个经典测试在超长文本中藏一个关键信息比如小明最喜欢的数字是4826然后问模型这个数字是什么。Llama 3.1 8B 的 NIAH 结果关键数据2048 token KV缓存全缓存的约5%近无损性能256 token KV缓存全缓存的约0.7%仍保持90%的原始准确率相比之下AdaKV和HeadKV在低预算下表现较差这意味着什么你可以把KV缓存压缩到原来的1/100还能保持90%的准确率6.3 只选4个SRH就够了CompressKV的消融实验证明了一个惊人的事实每层只需选择top-4个SRH就能达到最佳性能。每层SRH数量平均准确率相比Top-4Top-244.33-0.63Top-444.96基准Top-644.79-0.17Top-1244.960.00Top-2444.30-0.66为什么Top-24反而变差因为选太多SRH会引入噪音头稀释了真正重要的头的信号。这符合少即是多的直觉。6.4 推理效率内存和延迟内存在固定KV预算下CompressKV的峰值内存与全缓存相比大幅降低且与上下文长度无关延迟驱逐方法的解码延迟保持恒定不随上下文增长而全缓存的延迟线性增长首token时间TTFT所有方法都随上下文增加这是prefilling阶段的固有成本 ** CompressKV 在 128K 上下文下只使用 1024-token KV缓存内存占用仅为全缓存的约1.5%而推理速度大幅提升。**七、伪代码核心逻辑长什么样以下是CompressKV的完整推理流程伪代码classCompressKV:def__init__(self,model,calibration_data,num_srh4):self.modelmodel self.num_srhnum_srh# 离线阶段只需运行一次# 1. 识别每层的语义检索头self.srh_per_layerself.identify_srh(calibration_data)# 2. 估计每层的压缩误差self.layer_errorsself.estimate_layer_errors(calibration_data)defidentify_srh(self,calibration_data):识别语义检索头公式1srh_scores{}forlayerinself.model.layers:forheadinlayer.attention_heads:score0forsampleincalibration_data:answer_spansample.answer_token_indicesforstepinrange(sample.num_generation_steps):ifsample.generated_token[step]inanswer_span:# 跨度聚合关注整个答案区间scoresum(head.attention[step,idx]foridxinanswer_span)srh_scores[head]score# 选top-ktop_headssorted(srh_scores,keysrh_scores.get,reverseTrue)[:self.num_srh]srh_per_layer[layer]top_headsreturnsrh_per_layerdefestimate_layer_errors(self,calibration_data):离线计算每层压缩误差errors[]forlayerinself.model.layers:# 全缓存 vs 压缩缓存的输出差异O_fulllayer.forward(query,K_full,V_full)K_comp,V_compself.compress(layer,K_full,V_full,small_budget128)O_complayer.forward(query,K_comp,V_comp)errorfrobenius_norm(O_comp-O_full)/frobenius_norm(O_full)errors.append(error)# 归一化return[e/sum(errors)foreinerrors]defcompress(self,layer,K,V,budget):压缩KV缓存只保留SRH认为重要的tokensrhsself.srh_per_layer[layer]# 聚合SRH的注意力分数在观察窗口内scoreszeros(K.shape[0])forsrhinsrhs:scoressrh.get_attention_scores(window_size64)scores/len(srhs)# 选top-Nkeep_indicestop_k(scores,budget)returnK[keep_indices],V[keep_indices]defallocate_budget(self,total_budget):按误差比例分配层间预算num_layerslen(self.model.layers)basetotal_budget//num_layers budgets[]forerrinself.layer_errors:# 保底32个token最多3倍基础预算bmax(32,min(3*base,total_budget*err))budgets.append(int(b))returnbudgetsdefgenerate(self,prompt,total_budget1024):推理入口# 1. 预填充计算初始KV缓存K_cache,V_cacheself.model.prefill(prompt)# 2. 按层分配预算layer_budgetsself.allocate_budget(total_budget)# 3. 逐层压缩fori,layerinenumerate(self.model.layers):K_cache[i],V_cache[i]self.compress(layer,K_cache[i],V_cache[i],layer_budgets[i])# 4. 解码生成使用压缩后的KV缓存output[]forstepinrange(max_new_tokens):next_tokenself.model.decode_step(K_cache,V_cache)output.append(next_token)return.join(output)八、能跟现有技术叠buff吗8.1 与Prefilling加速器MInference/XAttention这些技术优化的是预填充阶段把长文本变成KV缓存的过程而CompressKV优化的是解码阶段使用KV缓存生成回复的过程。它们是完全正交的可以一起使用既加速预填充又减少解码内存。8.2 与KV缓存量化KIVIKIVI把KV缓存的值从16-bit压缩到2-bit甚至1-bit但保留所有token。CompressKV保留全精度但只保留重要token。两者可以组合使用方案KV内存占用准确率FullKV16-bit100%基准KIVI 2-bit~12.5%接近基准KIVI 1-bit~6.25%大幅下降CompressKV~3%接近基准CompressKV KIVI 2-bit~1.6%仍保持强劲性能1.6%的内存占用意味着什么原本需要64GB显存才能跑128K上下文现在只需要约1GB8.3 与头级分配方法HeadKV/AdaKVHeadCompressKV CompressKV的token选择 HeadKV的头级预算分配AdaCompressKV CompressKV的token选择 误差感知层间分配 AdaKV实验表明这些组合变体在LongBench上提升近2分在NIAH上提升达11分紧内存下九、总结与展望核心贡献语义检索头SRH识别那些真正会找重点的注意力头避免Streaming Head的噪音主导驱逐决策离线误差感知层间分配通过测量压缩前后的输出差异智能地在不同层之间分配KV缓存预算极简高效只需每层4个SRH无需在线计算部署零额外开销关键数据回顾指标数值LongBench问答仅3% KV缓存保留97%全缓存性能NIAH仅0.7% KV缓存90%准确率每层所需SRH4个与KIVI 2-bit组合后内存占用约1.6%对开发者的启示如果你在做RAG检索增强生成或长文档处理CompressKV的思路值得借鉴不是所有信息都重要找到会找重点的机制是关键如果你在部署LLM到资源受限环境边缘设备、移动AppKV缓存压缩是必修课离线分析 在线推理的架构模式很优雅把计算移到预处理阶段运行时几乎零开销未来方向将SRH的概念扩展到其他注意力变体如稀疏注意力、线性注意力动态调整SRH目前是一次性离线识别与更激进的量化方案如1-bit结合进一步压缩内存十、参考资料论文原文: CompressKV: Semantic-Retrieval-Guided KV-Cache Compression for Resource-Efficient Long-Context LLM Inference代码仓库: https://github.com/TUDa-HWAI/CompressKV相关论文:StreamingLLM: arXiv:2309.17453SnapKV: arXiv:2404.14469HeadKV: arXiv:2406.07018KIVI: arXiv:2402.02750作者: Marst.zhang | 整理日期: 2026-06-25如果这篇博客对你有帮助欢迎点赞、收藏、转发有任何问题欢迎在评论区交流