从KV Cache到AI内存系统：大模型推理架构演进与优化实践

大模型推理这件事表面上看是算力问题实际上越往深里做越会发现真正卡脖子的往往是内存。我接触过不少做推理服务的团队大家一开始都盯着 GPU 的算力利用率觉得把 FLOPS 跑满就万事大吉了结果上线之后才发现显存先爆了请求排队排到天上去。这个项目标题“从 KV Cache 到 AI 内存系统大模型推理架构的演进”说的其实就是这么一条主线——我们是怎么从最初只想着缓存几个 Key-Value 向量一步步被逼着重新设计整个内存系统的。这篇文章我想聊的不是某个具体框架的 API 怎么调而是把这条演进脉络讲清楚KV Cache 到底缓存了什么、为什么它成了推理的显存黑洞、业界围绕它做了哪些优化、以及当这些优化都做到头之后为什么大家开始把目光投向更上层的“AI 内存系统”。如果你正在做推理部署、模型服务或者只是好奇为什么一张 80G 的卡跑个长上下文就喘不过气这篇应该能给你一些可以直接参考的思路和踩坑经验。1. 从 Transformer 推理的计算特征说起1.1 自回归解码为什么天然是内存瓶颈要理解 KV Cache得先回到 Transformer 解码的本质。Transformer 的自回归生成是逐 token 进行的每生成一个新 token都要拿它去和前面所有已经生成的 token 做注意力计算。这里有个关键点——注意力的计算需要当前 token 的 Query 去和历史上所有 token 的 Key、Value 做点积。如果什么都不缓存那每生成一个 token都要把前面所有 token 的 Key 和 Value 重新算一遍。假设序列长度是 n生成第 n 个 token 时前面 n-1 个 token 的 K、V 都要重算整体复杂度就是 O(n²) 次前向计算。这在工程上完全不可接受因为生成 1000 个 token 就要做接近百万次量级的重复矩阵运算。KV Cache 的思路非常直接既然前面 token 的 Key 和 Value 在生成过程中不会变那就把它们算一次存下来后续每步直接读取。这样每步只需要计算当前 token 的 Q、K、V然后拿新的 Q 去和缓存里的全部 K 做注意力。计算量从 O(n²) 降到了 O(n)代价是显存里多存了一份随序列长度线性增长的数据。提示这里说的“不变”是有前提的。标准自注意力里历史 token 的 K、V 确实不依赖后续 token所以可以缓存。但如果你用的是某些带双向信息的结构或者做了特殊的注意力掩码缓存策略就要重新考虑。1.2 KV Cache 的显存占用到底怎么算很多人对 KV Cache 的显存占用只有一个模糊印象觉得“挺大的”但具体多大、怎么算说不清楚。我把它拆成一个可以直接套的公式。单个 token 的 KV Cache 占用 2 × 层数 × 注意力头数 × 头维度 × 数据类型字节数。这里的 2 是因为要存 Key 和 Value 两份。注意头数乘以头维度通常等于隐藏维度所以也可以简化成 2 × 层数 × 隐藏维度 × 字节数。拿一个常见的 7B 模型举例假设 32 层、隐藏维度 4096、FP16 存储2 字节。那么单个 token 的 KV Cache 就是 2 × 32 × 4096 × 2 524288 字节约 0.5 MB。看起来不大但序列长度一上去就吓人了4096 长度的上下文单个请求就是 2 GB如果并发 16 个请求直接 32 GB 没了。而 7B 模型本身的权重在 FP16 下也才 14 GB 左右。也就是说长上下文场景下KV Cache 的显存占用可以轻松超过模型权重本身。这就是为什么很多团队发现模型明明能加载进去一跑长文本就 OOM。问题不在权重在缓存。1.3 内存带宽才是真正的隐形天花板显存容量只是第一道坎第二道坎是内存带宽。自回归解码每一步的计算量其实很小——就是几个矩阵向量乘但每一步都要把整个 KV Cache 从显存里读一遍。这意味着解码阶段的性能瓶颈往往不是算力而是显存带宽。我实测过一个现象在同样的硬件上把 batch size 从 1 加到 8吞吐量提升远没有线性增长甚至在某些配置下几乎不涨。原因就是 KV Cache 的读取把带宽吃满了GPU 的计算单元在等数据。这也是为什么后来大家开始研究 PagedAttention、量化 KV、稀疏注意力这些方向——本质上都是在想办法减少每步要搬运的数据量。理解了这一层你就能明白为什么“AI 内存系统”这个概念会被提出来。当计算优化做到极致之后内存的容量、带宽、管理方式就成了决定推理成本和吞吐的核心变量。2. KV Cache 优化的几条主流技术路线2.1 PagedAttention把操作系统分页思想搬进显存PagedAttention 是我认为过去两年对推理架构影响最大的一个设计。它的核心洞察是KV Cache 的显存浪费很大程度上来自内存碎片和过度预留。传统做法是给每个请求按最大可能长度预分配一块连续显存。但实际生成时很多请求根本到不了最大长度预分配的空间就浪费了而且不同请求长度不一连续分配会产生大量碎片导致明明总显存够却分配不出足够大的连续块。PagedAttention 借用了操作系统虚拟内存分页的思路把 KV Cache 切成固定大小的块block每个块存固定数量 token 的 K、V。请求的 KV Cache 不再要求物理连续而是通过一张块表来映射逻辑顺序到物理块。这样一来显存按需分配碎片问题基本消失实测能把显存利用率从百分之五六十提到百分之九十以上。注意块大小是个需要调的参数。块太小块表本身的开销和寻址次数上去了块太大内部碎片又回来了。常见取值在 16 到 32 个 token 之间具体要看你的序列长度分布。2.2 KV Cache 量化用精度换容量量化是另一条很实用的路线。既然 KV Cache 占显存那就把它从 FP16 压到 INT8 甚至 INT4。INT8 量化理论上能把缓存占用直接砍半INT4 再砍一半。但 KV 量化和权重量化不太一样坑更多。Key 的分布通常比 Value 更集中量化相对安全Value 的分布更分散激进量化容易掉点。我自己的经验是INT8 的 KV 量化在大多数模型上几乎无损可以直接上INT4 就要谨慎最好针对 Key 和 Value 用不同的量化粒度并且一定要在自己的评测集上验证。还有一个容易被忽略的点量化之后反量化本身也要消耗计算和带宽。如果反量化的开销吃掉了省下来的带宽收益那就得不偿失。所以量化 KV 一定要配合 kernel 层面的优化不能只在框架层做。2.3 稀疏化与淘汰不是所有 token 都值得留长上下文里很多历史 token 对当前生成的贡献其实很小。基于这个观察一批工作开始做 KV Cache 的稀疏化和淘汰。思路大致分两类。一类是静态的比如只保留最近的滑动窗口加上少数几个“重要”token另一类是动态的根据注意力分数实时决定哪些 KV 保留、哪些丢弃。H2O 这类方法就是典型代表它维护一个累计注意力分数把分数低的 token 淘汰掉。这类方法在特定任务上效果不错但通用性是个问题。有些任务恰恰依赖远处的关键信息你一淘汰就出错。我的建议是如果你的业务场景比较固定比如都是短问答或者都是摘要可以大胆用如果是开放域长文本理解淘汰策略要非常保守或者干脆不用。2.4 前缀共享与多请求复用在实际服务里经常有大量请求共享相同的前缀比如同一个系统提示词、同一段背景文档。这时候如果每个请求都独立存一份 KV Cache就是巨大的浪费。前缀共享的思路是把公共前缀的 KV Cache 只存一份多个请求通过引用共享。PagedAttention 的块表机制天然支持这种共享——不同请求的块表可以指向相同的物理块只要这些块是只读的。这在 RAG 场景和 Agent 场景里收益特别明显因为系统提示词往往很长且固定。实现上要注意写时复制一旦某个请求要修改共享块比如继续生成导致块内容变化就得先复制一份再改不能直接改共享块。这个逻辑和操作系统的 copy-on-write 是一模一样的。3. 从单点优化到系统级内存管理3.1 为什么零散优化会撞到天花板上面这些优化单独拿出来都有效但组合起来用的时候问题就来了。PagedAttention 管的是块分配量化管的是数据精度淘汰管的是保留策略前缀共享管的是复用——它们各自维护一套状态彼此之间缺乏协调。举个我遇到过的真实情况一个请求的 KV Cache 被量化成 INT8 存着同时又被淘汰策略标记为“低优先级”还被分页机制拆成了不连续的块。当这个请求突然需要用到那些被标记的 token 时系统要同时做反量化、块寻址、以及可能的重新计算延迟一下就上去了。这种跨层协调的缺失就是零散优化的天花板。更深层的问题是这些优化都是被动响应式的——显存不够了才淘汰请求来了才分配。它们没有一个全局视角去决定“这块内存到底该给谁、以什么形式给、给多久”。3.2 AI 内存系统要解决的三个核心问题当我把上面这些问题摊开看发现所谓“AI 内存系统”本质上要回答三个问题。第一是容量问题怎么在有限的显存里塞下更多的有效 KV 数据。这涉及量化、压缩、淘汰、以及把冷数据换出到主机内存甚至更远的存储层。第二是带宽问题怎么减少每步解码要搬运的数据量。这涉及稀疏注意力、块级别的按需加载、以及计算和访存的 overlap。第三是调度问题怎么在多个请求、多个层级之间做全局最优的分配。这涉及优先级、生命周期管理、以及跨请求的复用决策。这三个问题不是独立的。你为了省容量做了激进量化可能增加带宽压力你为了省带宽做了稀疏化可能影响调度时的复用机会。真正的内存系统必须把这三者放在一个框架里权衡。3.3 分层内存把主机内存和存储拉进来一个很自然的演进方向是分层。显存最快但最贵最小主机内存慢一些但大得多再往下还有本地存储。KV Cache 完全可以按“热度”分层存放。热数据也就是当前正在生成、马上要用的 KV放显存温数据比如还在活跃但暂时不用的请求的 KV放主机内存冷数据比如长时间挂起的会话可以落到存储上。需要的时候再按层换入。这个思路听起来简单工程上难点在于换入换出的时机和粒度。换得太频繁PCIe 带宽成为新瓶颈换得太粗又起不到省显存的作用。我见过比较靠谱的做法是按块做换出并且用预测性的预取——根据生成进度提前把下一步可能用到的块拉回显存。提示分层内存的收益高度依赖你的请求模式。如果是短请求为主分层几乎没用如果是长会话、多轮对话为主分层能显著提升并发数。3.4 内存系统与调度的协同设计再往深一层内存系统不能和请求调度分开设计。传统调度器只看请求的优先级和到达时间不看内存状态。但内存紧张的时候调度器应该主动限制并发、或者优先处理那些 KV 复用率高的请求。我试过一个简单的策略在调度时给每个请求估算一个“内存成本”包括它需要的 KV 容量和预计的带宽消耗然后按成本收益比排序。实测下来在混合负载场景里这种内存感知的调度比纯 FIFO 能提升百分之二三十的有效吞吐。当然这需要调度器和内存管理器之间有清晰的接口不能各自为政。4. 实操搭一个带 KV Cache 管理的推理服务原型4.1 环境与依赖准备光讲原理不够我把自己搭原型的过程整理一下你可以照着复现。这里不绑定具体框架讲的是通用结构你用 vLLM、TensorRT-LLM 还是自己写都能对应上。基础环境是一台带 GPU 的机器CUDA 和 PyTorch 装好。核心依赖是你要有一个能暴露 KV Cache 的推理后端。如果用的是现成框架确认它支持分页或至少支持自定义缓存管理如果是自己写那就要实现注意力层的 KV 存取接口。我建议先用一个小模型起步比如 1B 到 3B 级别的这样单卡就能跑调试周期短。等逻辑跑通了再换大模型验证。4.2 实现一个最简块管理器块管理器是整个系统的核心。我用 Python 写过一个简化版逻辑大概是这样的。class BlockManager: def __init__(self, num_blocks, block_size): self.block_size block_size self.free_blocks list(range(num_blocks)) self.block_tables {} # request_id - list of block ids 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 MemoryError(not enough blocks) blocks [self.free_blocks.pop() for _ in range(num_needed)] self.block_tables[request_id] blocks return blocks def free(self, request_id): blocks self.block_tables.pop(request_id, []) self.free_blocks.extend(blocks)这段代码很粗糙但把核心逻辑说清楚了按块分配、用块表记录映射、释放时归还。真实系统里还要处理块内碎片、共享块的引用计数、以及并发安全。注意块表本身也要占内存。如果块很小、请求很多块表的开销不能忽略。我一般会把块大小设成 16然后在压测里观察块表占用必要时调大块尺寸。4.3 接入注意力层并验证正确性块管理器写好后要把它接到注意力计算里。关键改动是计算注意力时不再假设 KV 是连续存储的而是通过块表逐块读取。这一步最容易出错的地方是位置编码。因为块是乱序分配的读取时必须按逻辑顺序还原位置编码也要对应正确。我踩过的坑是块表顺序和位置编码对不上导致生成结果乱码排查了半天才发现是块表拼接时顺序错了。验证正确性的方法很直接用同样的输入对比开启和关闭分页管理的输出。如果两者生成的 token 完全一致说明逻辑没问题。如果不一致先检查块表顺序再检查位置编码。4.4 压测与关键指标观测跑通之后就要压测。我关注的指标有这么几个。指标含义观测方式显存利用率实际用于 KV 的显存占比显存监控工具首 token 延迟请求到第一个 token 的时间请求日志打点每 token 延迟解码阶段单步耗时请求日志打点有效吞吐单位时间完成的 token 数压测工具统计块碎片率空闲但无法利用的块比例块管理器统计压测时我建议分两组一组固定并发看延迟一组固定延迟看最大并发。这样能看出系统的容量边界在哪里。我实测下来分页管理在长序列、高并发场景下优势最明显短序列低并发时反而因为块表开销略慢一点这是正常的。5. 常见问题与排查实录5.1 显存够但分配失败这是最典型的问题。监控显示显存还有余量但请求就是分配不到块。原因通常是碎片——空闲块不连续或者被块表、临时张量占着。排查思路是先看块管理器的空闲块数再看显存的实际碎片情况。如果空闲块够但分配失败多半是块表或引用计数出了问题。我遇到过一次是共享块的引用计数没减导致块永远不释放查了好久。5.2 长上下文生成质量下降如果用了量化或淘汰长上下文质量下降是常见副作用。排查时先关掉所有优化确认基线质量然后逐个开启定位是哪个环节导致的。量化导致的下降通常表现为局部用词异常淘汰导致的下降通常表现为丢失远处信息。两者特征不同可以据此判断。5.3 吞吐上不去但 GPU 利用率低这说明瓶颈不在算力在内存带宽或调度。先看 KV Cache 的读取量如果每步读取量很大考虑量化或稀疏化如果读取量不大但利用率还是低看是不是调度器把请求串行化了或者块换入换出太频繁。5.4 多请求前缀共享失效前缀共享失效通常是因为前缀匹配没做对。检查两点一是前缀的哈希或标识是否一致二是共享块的写时复制是否正确触发。我见过因为系统提示词里有个时间戳导致每次前缀都不同的情况这种要提前规范化。问题现象可能原因排查方向显存够但分配失败碎片或引用计数泄漏查块管理器和引用计数长上下文质量下降量化或淘汰过度逐项关闭优化定位吞吐低 GPU 利用率低带宽瓶颈或调度串行查读取量和调度日志前缀共享失效前缀不一致或写时复制错误查前缀标识和复制逻辑6. 我对这条演进路线的一些判断做了这么久推理优化我越来越觉得 KV Cache 只是冰山一角。它之所以重要是因为它把“内存”这个变量第一次推到了推理架构的中心位置。以前我们优化推理想的是怎么把计算图压得更紧、kernel 写得更快现在想的是内存怎么分配、怎么复用、怎么分层、怎么和调度协同。这个转变带来的直接影响是推理系统的设计越来越像操作系统。块管理像内存分页前缀共享像写时复制分层存储像缓存层次调度器像进程调度。如果你有系统背景会发现很多老思路可以直接迁移过来这是我觉得特别有意思的地方。另一个判断是未来的 AI 内存系统不会只服务推理。训练时的激活值、优化器状态本质上也是内存问题只是访问模式不同。当推理和训练的内存管理开始共享同一套抽象整个 AI 基础设施的形态可能会再变一次。最后分享一个我自己的小习惯每次做内存优化我都会先画一张数据流图标清楚每一份数据从哪来、到哪去、活多久、被谁读。这张图往往比任何 profiling 工具都更能帮我找到优化点。KV Cache 的很多问题画完图就一目了然了。