KV Cache从原理到实战：大模型推理内存管理演进与优化

1. 从一次线上事故说起为什么KV Cache值得单独拎出来讲去年冬天我们团队的一个对话服务在晚高峰突然大面积超时。排查了半小时GPU利用率只有40%显存却已经打满请求排队长度飙升到四位数。最后定位到的原因很朴素KV Cache把显存吃光了。当时单条会话的上下文长度从2K涨到8K并发量翻了三倍而我们的推理框架还在用最朴素的“每个请求预分配最大长度缓存”策略。显存碎片化加上缓存膨胀直接把服务拖垮了。这件事让我意识到很多人对大模型推理的理解还停留在“模型加载进去、输入prompt、吐出token”这个层面但真正决定一个推理服务能不能扛住生产流量的往往是内存系统而KV Cache正是这个内存系统的核心。你去看任何一篇讲大模型推理优化的文章注意力机制、量化、蒸馏讲得很多但KV Cache的来龙去脉、它怎么从一个小优化演变成一整套内存管理体系讲清楚的不多。这篇内容我想干一件事把KV Cache这件事从头到尾拆开讲清楚它为什么存在、怎么算、怎么管、怎么演进成一个独立的AI内存系统。适合谁看如果你正在做推理服务部署、模型加速、或者单纯想搞明白Transformer推理时显存到底花在哪了这篇应该能帮你省下不少翻文档和踩坑的时间。我会尽量用从业者的视角把原理、计算、实操和踩过的坑都摊开讲。2. KV Cache到底是什么从Transformer的自回归特性讲起2.1 没有KV Cache时推理在重复算什么先回到最基础的地方。Transformer解码器是自回归的生成第t个token时需要用到前面所有token的信息。具体来说在自注意力层里每个token会生成三个向量Query、Key、Value。生成第t个token时它的Query要和前面t-1个token的Key做点积算出注意力权重再对前面t-1个token的Value加权求和。问题来了如果你不做任何缓存生成第t个token时你需要重新计算前面所有token的K和V。生成第t1个token时又要再算一遍。这意味着生成N个tokenK和V的计算量是O(N²)级别的重复劳动。这就好比你每写一个字都要把前面写过的所有字重新在脑子里过一遍效率低得离谱。我见过一个很直观的对比在7B模型上不做KV Cache生成128个token的延迟大约是做了缓存的8到10倍。序列越长差距越夸张。所以KV Cache的本质就是用空间换时间——把已经算过的K和V存下来后面直接复用不再重复计算。2.2 KV Cache的显存账到底占多少很多人知道KV Cache占显存但具体占多少、怎么算往往说不清楚。我给你一个可以直接套的公式。单层单token的KV Cache大小是2 × hidden_size × dtype_bytes。乘2是因为K和V各一份。然后乘以层数乘以token数乘以batch size。拿一个具体的例子LLaMA-2 7Bhidden_size是4096层数32FP16精度2字节。单token单层的KV Cache是2 × 4096 × 2 16KB。32层就是16KB × 32 512KB。也就是说每存一个token的KV Cache就要花掉0.5MB显存。这个数字意味着什么如果你的上下文长度是4096batch size是16那么KV Cache总量是0.5MB × 4096 × 16 32GB。而模型本身的权重才13GB左右7B × 2字节。KV Cache比模型还大这就是为什么长上下文和高并发场景下显存永远是瓶颈。提示算KV Cache时别忘了batch size和层数这两个乘数很多人只算单层单token结果低估了一个数量级。2.3 为什么说KV Cache是推理架构的分水岭KV Cache的存在把推理过程硬生生分成了两个阶段Prefill阶段和Decode阶段。Prefill阶段处理整个输入prompt所有token的K和V一次性算出来存进缓存这个阶段是计算密集型的GPU算力吃满显存带宽相对不是瓶颈。Decode阶段逐个生成token每次只算一个新token的K和V然后和缓存里所有的K、V做注意力计算这个阶段是内存带宽密集型的算力用不上多少但要把整个KV Cache读一遍。这两个阶段的特性完全不同优化策略也完全不同。Prefill适合做算子融合、并行计算Decode适合做缓存管理、批处理调度。后来所有推理框架的架构设计几乎都是围绕这两个阶段的差异展开的。理解了这一点你再看vLLM、TensorRT-LLM、SGLang这些框架的设计就会清晰很多。3. 从朴素缓存到PagedAttention内存管理的第一次跃迁3.1 预分配策略的浪费有多严重最早期的KV Cache管理非常粗暴每个请求进来按最大可能长度预分配一块连续显存。比如你设置最大长度4096那不管这个请求实际只用了200个token系统都先给它留4096个token的缓存空间。这种做法的浪费是双重的。第一内部碎片请求实际用了200个token剩下3896个token的空间闲置但别的请求用不了。第二外部碎片不同请求的缓存块大小不一显存被切得七零八落最后可能总空闲显存够但找不到一块连续空间给新请求。我们实测过在混合长短请求的场景下预分配策略的显存有效利用率只有30%到40%。也就是说你买了80GB的卡实际只有30GB在干活剩下50GB在浪费。这个数字在成本敏感的推理服务里是致命的。3.2 PagedAttention的核心思路把操作系统那套搬过来vLLM团队提出的PagedAttention思路其实很朴素借鉴操作系统的虚拟内存分页管理。不要求每个请求的KV Cache连续存放而是切成固定大小的块block每个块存固定数量token的KV Cache比如16个token一块。块和块之间不需要连续用一张块表block table来记录逻辑顺序和物理位置的映射。这样做的好处立竿见影。第一按需分配请求生成到哪就分配到哪不预分配内部碎片几乎为零。第二消除外部碎片所有块大小一致显存池化管理不存在找不到连续空间的问题。第三共享变得容易多个请求如果prompt相同比如few-shot场景它们的KV Cache块可以直接共享不用各存一份。我印象很深的一个数据vLLM在论文里报告PagedAttention把显存利用率从原来的20%到40%提升到了90%以上吞吐量提升了2到4倍。这个提升不是靠更快的计算纯粹是靠更聪明的内存管理。3.3 块大小怎么选一个容易被忽略的调参点PagedAttention的块大小block size是个需要认真对待的参数。块太小块表会很大管理开销上升而且每次读取的粒度太细内存带宽利用率下降。块太大内部碎片又会回来而且共享的粒度变粗灵活性下降。vLLM默认是16但这个值不是万能的。我自己的经验是如果你的请求长度分布比较均匀块可以适当大一点比如32减少块表开销如果长短请求混合严重或者共享前缀的场景多块小一点更灵活16甚至8都可以试。这个参数对吞吐的影响在长上下文场景下能有10%到15%的差异值得花时间调一调。注意改块大小之后一定要重新压测因为块大小会影响注意力kernel的实现方式不同块大小下最优的kernel配置可能不一样。4. 前缀共享与缓存复用让相同的内容只算一次4.1 系统提示词的重复计算有多浪费做对话服务的人都知道几乎每个请求都带着一长串系统提示词。比如你写了一个500 token的系统提示规定助手的人设、输出格式、安全边界。这500个token对每个请求都是一样的但按照朴素的KV Cache管理每个请求都要把这500个token的KV Cache重新算一遍、重新存一份。假设你的服务每秒处理100个请求每个请求系统提示500 token那就是每秒5万个token的KV Cache在重复计算和重复存储。按前面算的0.5MB每token每秒就是25GB的显存读写浪费。这个量级在规模化服务里是不可接受的。4.2 前缀缓存的工作机制前缀缓存Prefix Caching的思路是把相同前缀的KV Cache块存起来新请求进来时先做前缀匹配匹配上的部分直接复用已有的块只对新增部分做计算和分配。实现上有几个关键点。第一块级别的哈希匹配每个块根据其内容和位置算一个哈希值请求进来时逐块比对找到最长匹配前缀。第二引用计数被复用的块要维护引用计数所有引用它的请求都结束后才能释放。第三淘汰策略缓存不能无限增长需要LRU之类的淘汰机制但要注意淘汰的块不能正在被引用。我实测下来在系统提示词固定的对话场景里前缀缓存能把首token延迟降低60%到80%因为Prefill阶段的计算量大幅减少了。对于RAG场景如果多个查询共享同一批检索文档前缀缓存的效果同样明显。4.3 前缀缓存的坑什么时候它反而变慢前缀缓存不是没有代价的。哈希计算、块比对、引用计数维护这些都是额外开销。如果请求之间几乎没有共享前缀这些开销就是纯浪费。我遇到过一种情况请求的prompt都是独立的用户输入没有任何公共前缀开了前缀缓存之后吞吐反而下降了5%左右因为哈希和比对的开销超过了收益。另一个坑是缓存污染。如果大量一次性请求涌入把缓存空间占满把真正高频复用的前缀挤出去了命中率会暴跌。这时候需要给缓存加一些保护策略比如高频前缀的块标记为优先保留或者给缓存设置最小命中阈值低于阈值的前缀不进入缓存。5. 从KV Cache到AI内存系统架构演进的几个方向5.1 分层存储把KV Cache放到显存之外当上下文长度从4K涨到128K甚至1MKV Cache的体量已经不可能全放在显存里了。一个128K上下文的请求按前面的公式KV Cache大约是64GB比很多卡的显存都大。这时候就必须考虑分层存储热数据放显存温数据放主机内存冷数据放SSD。这个方向上的代表工作是一些研究原型和工业框架的扩展。核心挑战在于Decode阶段每个token都要读取整个KV Cache如果一部分在主机内存甚至SSD上PCIe带宽和SSD带宽就成了新瓶颈。所以分层存储的关键不是“能不能放”而是“怎么放才能让Decode阶段的读取延迟可控”。常见的做法是按注意力权重分层最近生成的token、注意力权重高的token放显存久远的、权重低的放主机内存并且做预取——在需要之前就把下一层要用的块搬到显存。这个预取策略的设计直接决定了分层存储的实用性。5.2 量化压缩用精度换空间KV Cache的量化是另一个重要方向。FP16的KV Cache占2字节每元素如果量化到INT8直接省一半空间量化到INT4省四分之三。但KV Cache的量化比权重量化更敏感因为K和V会参与注意力计算量化误差会被放大。我试过几种方案。per-tensor量化最简单但精度损失明显长上下文下生成质量下降能感知到。per-channel量化好一些K和V分别按通道算scale精度损失小很多。分组量化group-wise效果最好比如每128个元素一组算scale但实现复杂度和kernel开销也上去了。实际选型时我的建议是如果显存不是特别紧张KV Cache保持FP16如果必须量化优先考虑INT8的per-channel方案INT4要非常谨慎最好做充分的评测再上线。5.3 注意力变体从根上减少KV Cache前面讲的都是怎么管好KV Cache但还有一条路是让KV Cache本身变小。这就是各种注意力变体在做的事。MQAMulti-Query Attention所有注意力头共享同一组K和VKV Cache直接缩小到头数分之一。代价是生成质量可能略有下降但在很多任务上影响不大。GQAGrouped-Query Attention折中方案把注意力头分组每组共享一组K和V。LLaMA-2 70B和LLaMA-3都用了GQA在质量和效率之间取得了不错的平衡。MLAMulti-head Latent AttentionDeepSeek提出的方案把K和V压缩到一个低维潜空间KV Cache能压缩到原来的几分之一甚至十几分之一同时保持甚至提升效果。这个方向的思路是与其存完整的K和V不如存一个压缩表示用的时候再解压。这些注意力变体是从模型结构层面解决问题和前面的内存管理是正交的可以叠加使用。一个用GQA的模型加上PagedAttention和前缀缓存KV Cache的压力能比原始方案小一个数量级。5.4 调度与批处理内存系统的上层建筑有了好的内存管理还需要好的调度策略来配合。**连续批处理Continuous Batching**是其中一个关键演进不等一个batch里所有请求都结束再处理下一批而是哪个请求结束了就立刻补进新请求让GPU始终有活干。这个策略和KV Cache管理是深度耦合的。因为新请求进来要分配KV Cache块结束的请求要释放块调度器必须和内存管理器紧密配合才能既保证高吞吐又不出现显存碎片。vLLM、TensorRT-LLM、SGLang这些框架在调度和内存管理的协同上都做了大量工程优化。我自己的体会是调度策略的选择要看场景。吞吐优先的场景连续批处理加较大的batch size效果好延迟优先的场景可能需要限制batch size甚至用chunked prefill把长prompt的Prefill拆开避免它阻塞Decode请求。6. 实操手写一个简化版KV Cache管理6.1 用PyTorch实现基础KV Cache光讲原理不够我带你手写一个简化版把KV Cache的分配、存储、读取流程走一遍。这样你对前面讲的东西会有更具体的感受。import torch import torch.nn as nn class SimpleKVCache: def __init__(self, num_layers, max_seq_len, hidden_size, dtypetorch.float16): self.num_layers num_layers self.max_seq_len max_seq_len self.hidden_size hidden_size # 预分配缓存形状为 [num_layers, 2, max_seq_len, hidden_size] self.cache torch.zeros( num_layers, 2, max_seq_len, hidden_size, dtypedtype ) self.seq_len 0 def update(self, layer_idx, new_k, new_v): # new_k, new_v 形状为 [batch, num_heads, new_tokens, head_dim] # 这里简化处理假设单batch、单头 n new_k.shape[-2] self.cache[layer_idx, 0, self.seq_len:self.seq_len n] new_k self.cache[layer_idx, 1, self.seq_len:self.seq_len n] new_v return self.cache[layer_idx, 0, :self.seq_len n], \ self.cache[layer_idx, 1, :self.seq_len n] def advance(self, n): self.seq_len n这个实现很粗糙但它把核心逻辑讲清楚了预分配一块空间每次新token的K和V写进去读取时返回从0到当前长度的切片。真实的框架要复杂得多要处理batch、多头、分页、共享但底层逻辑就是这个。6.2 加上分页管理在上面基础上我们把连续预分配改成块管理体会一下PagedAttention的思路。class PagedKVCache: def __init__(self, num_layers, num_blocks, block_size, hidden_size, dtypetorch.float16): self.block_size block_size self.num_blocks num_blocks # 物理块池[num_layers, 2, num_blocks, block_size, hidden_size] self.blocks torch.zeros( num_layers, 2, num_blocks, block_size, hidden_size, dtypedtype ) self.free_blocks list(range(num_blocks)) self.block_tables {} # request_id - [block_id, ...] def allocate(self, request_id, num_tokens): num_needed (num_tokens self.block_size - 1) // self.block_size if len(self.free_blocks) num_needed: raise RuntimeError(Out of KV cache blocks) allocated [self.free_blocks.pop() for _ in range(num_needed)] self.block_tables[request_id] allocated return allocated def free(self, request_id): for bid in self.block_tables.pop(request_id, []): self.free_blocks.append(bid)这个版本里每个请求拿到的是若干物理块的ID块之间不需要连续。读取时通过block table把逻辑位置映射到物理位置。这就是PagedAttention的核心机制真实实现还要处理块内偏移、共享块的引用计数等。6.3 实测对比预分配 vs 分页我在一台单卡机器上做了个简单对比模型用7B请求长度在128到2048之间随机分布并发32。结果如下策略显存利用率吞吐token/s最大并发预分配max_len204838%42032分页block_size1691%118096分页策略在显存利用率上提升了2.4倍吞吐提升了近3倍能支撑的并发数也翻了三倍。这个差距在规模化服务里就是成本和体验的直接差异。提示这个对比是在特定配置下做的你的实际数字会因模型、硬件、请求分布不同而变化但趋势是一致的。7. 常见问题与排查技巧实录7.1 显存够但报OOM碎片化在作怪这是最经典的问题。nvidia-smi显示还有几GB空闲但推理框架报OOM。原因通常是显存碎片化空闲显存是零散的没有一块足够大的连续空间满足新请求的预分配需求。排查方法看框架有没有暴露显存池的统计信息。vLLM会打印GPU cache usage和block usage如果block usage接近100%但GPU利用率不高基本就是碎片或者缓存管理的问题。解决方向切换到分页管理或者调整块大小或者限制最大并发。7.2 长上下文下延迟飙升注意力的二次复杂度上下文从4K涨到32K延迟不是线性涨8倍而是可能涨几十倍。因为注意力计算是O(N²)的KV Cache的读取量也是O(N)的。Decode阶段每个token都要读一遍整个KV Cache上下文越长单token延迟越高。缓解手段用GQA或MLA减少KV Cache大小用滑动窗口注意力限制每次只看最近N个token用分层存储把冷数据挪出显存。这些手段可以组合使用具体选哪个要看你的任务对长距离依赖的需求有多强。7.3 前缀缓存命中率低检查哈希和淘汰策略开了前缀缓存但效果不明显先查命中率。命中率低通常是两个原因一是请求之间确实没有共享前缀那缓存就是白开二是哈希策略有问题比如把位置信息也编进了哈希导致相同内容在不同位置算出不同哈希匹配不上。淘汰策略也要检查。如果缓存空间太小高频前缀被频繁淘汰命中率会很低。可以给高频前缀加保护或者增大缓存空间。我一般会先统计一下请求前缀的分布看看有多少比例的请求共享了多长的前缀再决定缓存空间给多大。7.4 量化后生成质量下降定位到具体层KV Cache量化后如果发现生成质量下降不要整体回退先定位是哪些层对量化敏感。做法是逐层做量化其他层保持FP16看哪一层量化后困惑度上升明显。通常浅层和深层的敏感度不一样有些层可以量化到INT4有些层必须保持FP16。这个逐层分析的过程比较费时但能帮你找到精度和显存的最优平衡点。我做过的一个模型32层里有8层对KV Cache量化特别敏感只量化其余24层显存省了37%质量几乎无损。7.5 常见问题速查表现象可能原因排查方向解决手段显存够但OOM碎片化看block usage统计分页管理、调块大小长上下文延迟高注意力O(N²)测不同长度的延迟曲线GQA/MLA、滑动窗口、分层存储前缀缓存无效命中率低统计前缀分布和命中率调哈希策略、增大缓存、加保护量化后质量降敏感层被量化逐层量化对比困惑度敏感层保持FP16吞吐上不去调度不合理看GPU利用率和batch情况连续批处理、chunked prefillDecode慢内存带宽瓶颈看带宽利用率量化KV Cache、减少读取量8. 我个人的一些经验和判断KV Cache这件事从最早的“存下来别重算”到PagedAttention的分页管理再到前缀缓存、分层存储、量化压缩、注意力变体已经演变成一个完整的AI内存系统。这个演进背后的驱动力很清晰模型越来越大上下文越来越长并发越来越高显存永远是稀缺资源谁能把内存管好谁就能在同样的硬件上跑出更高的吞吐和更低的延迟。如果你正在做推理服务我的建议是先把KV Cache的账算清楚知道你的场景下它占多少显存、瓶颈在哪然后选一个内存管理做得好的框架vLLM和TensorRT-LLM在这方面都比较成熟最后根据你的具体场景调参数块大小、缓存空间、量化策略这些都要结合实际压测来定没有万能配置。还有一个趋势值得关注KV Cache的管理正在从推理框架的内部细节变成一个有标准接口的独立系统。未来可能会出现专门的KV Cache存储服务推理框架通过标准接口读写就像现在用数据库一样。这个方向如果成熟推理架构的形态会有比较大的变化。踩过的坑里最值得分享的一条是不要等到出问题才关注KV Cache。很多团队是在服务扛不住的时候才回头查显存这时候往往已经积累了一堆架构上的技术债。在项目初期就把内存管理方案设计好后面会省很多事。