资讯详情

KV Cache显存优化:共享内存池与分页机制实践

📅 2026/10/17 11:33:38 | 华诺云谱 👁 阅读
KV Cache显存优化:共享内存池与分页机制实践
做推理服务优化有一段时间的人应该都有过这种经历跑压测的时候显卡算力明明没打满显存先爆了线上请求量稍微一波动某个实例就开始OOM重启打开监控一看KV Cache占了成片显存实际利用率却低得让人心慌。这篇文章就围绕我做的一个优化项目展开核心就三件事把KV Cache的显存占用量算明白、用共享内存池接管缓存分配、再用分页机制把缓存利用率拉上去。整个过程落地后服务的并发吞吐和显存碎片问题都有了实打实的改善如果你也在搞大模型推理部署或者做服务性能优化这篇文章里的思路和坑应该能直接用上。1. KV Cache凭什么成了推理服务的大头开销1.1 自回归解码的“记账本”机制大模型生成token的方式是逐字往外蹦的每生成一个新token模型都要重新计算一次所有历史token的注意力分数。问题在于如果不缓存历史token的K和V矩阵每一步解码都得把历史所有token的Key和Value重算一遍计算量随序列长度线性增长根本没法做长文本交互。所以KV Cache本质上就是给自回归过程设了一个“记账本”第n步生成完后把这一轮产生的Key矩阵和Value矩阵存下来下一步只计算新token对应的Q和K/V再和历史缓存拼起来做注意力计算。这个做法把单次解码的计算量从序列长度级降到了常数级代价就是显存里得一直留着一块逐token增长的区域。这就带来一个直接后果KV Cache的显存占用量和序列长度强相关而且在大规模并发场景下增长得极快。我见过不少初上线的推理服务模型权重也就占了十几GB显存结果KV Cache预留区直接划走了几十GB调度器还会按请求的最大长度去预留哪怕当前请求只有几百token预留空间一分不少。1.2 算一笔账一个token的KV Cache到底多大给KV Cache计算显存占用的公式不复杂每个token对应的KV数据量是单个token的KV字节数 2K和V两个矩阵× 层数 × 隐藏层维度 × 数据类型字节数以常见的7B规模模型为例32层Transformer隐藏层维度4096FP16精度下每个token的显存占用是2 × 32 × 4096 × 2字节 512KB。听着好像不大但一个请求上下文长4096 token单请求的KV Cache就需要512KB × 4096 2GB。如果模型服务同时扛着32个请求、每个平均2000 tokenKV Cache就吃掉32 × 2GB 64GB比模型权重本身还大得多。13B模型的参数又不一样40层、隐藏层5120同样FP16每个token占用约800KB。把这些数字写进容量规划文档后再去看显存扩容申请就心里有底了KV Cache才是并发瓶颈的核心变量不是模型权重。1.3 静态预留方案的三宗罪很多推理框架的默认做法是启动时按“最大并发数 × 最大序列长度”一次性把KV Cache显存预留出来。这个方案实现简单但问题非常明显。第一是显存碎片化严重。不同请求的生成长度不一样KV Cache随着解码不断增长请求结束后释放的内存块大小又各不相同。预留区里很快就堆满了大大小小的空洞新请求想分配一大块连续显存时分配不出来但显存总量明明是够的。我在一次压测中就观察到碎片率超过40%GPU看起来显存还有空余请求却报无法分配缓存。第二是长尾请求拖垮整体。遇到长上下文请求时它的KV Cache会一路增长到预留上限而旁边很多短请求只用了很小一块空间。调度器为了保证“预留充足”只能按最坏情况配置结果是硬件的钱花了不少真正跑到的并发量只有设计值的六成左右。第三是跨请求的缓存复用完全没法做。同一个对话前缀如果被多个请求共用比如公共Prompt、历史会话的前半段静态预留方案下每个请求都存了一份完整的K和V白白浪费好几倍的显存。这部分冗余在共享内存池方案里可以被规规矩矩地吃掉。2. 共享内存池方案把KV Cache握在自己手里2.1 不绕开cudaMalloc优化就无从谈起最先需要改变的思路是别再把KV Cache的内存请求全部交给运行时去动态分配。虽然有些场景下动态分配也能跑通但在多请求高并发的推理服务里大量kv cache分配释放操作会让显存管理成为新的瓶颈而且每次分配的大小不一致显存碎片率很快就失控了。换成共享内存池之后逻辑就顺了服务启动时一次性从系统申请一大块物理内存后续KV Cache的分配和释放全部由我们自己管理的分配器来接管不再频繁调用运行时申请函数。这个思路合租房的逻辑一样与其每次来人就新开一套公寓不如包下一整栋楼自己分房间控制权回到手里才能做精细化管理。2.2 分页缓存逻辑连续物理分散KV Cache在模型计算时看起来是按序列顺序连续访问的但物理上我们并不需要它真的连续存放。类比操作系统的虚拟内存分页机制逻辑地址空间连续物理页面可以散落在不同位置通过页表建立映射。KV Cache也按固定大小切成块每个块容纳固定数量的token调度器在逻辑上维护一个连续的缓存视图底层用块表把逻辑块映射到物理块。这就是分页式KV Cache的核心思路。它带来两个直接好处物理内存分配粒度变小后碎片率大幅下降释放一个块就能立刻给别的请求用请求增长时按需取块不再是开工时就锁死一整段最大长度的空间。2.3 多进程共享带来的额外收益共享内存这个名字里最容易被忽略的词是“共享”。推理服务经常以多进程或者多副本的形式部署每个Worker进程各维护一份KV Cache以前是各管各的。一旦KV Cache所在的内存区域允许进程间共享几个事就顺理成章地发生了。公共Prompt的KV缓存可以只计算一次多个请求通过引用计数共享同一个物理块只有写入不同分支时才做复制某个进程空闲时释放的缓存块另一个进程可以直接接管物理内存的利用不再以进程为边界对请求做前缀匹配时干净利落的缓存命中能直接把响应延迟打下来一大截。2.4 整体架构的四个部件落地这套方案时我的代码里核心就是四个部件。调度器负责给请求分配逻辑上的缓存区间并记录每个逻辑块对应的物理块块表维护逻辑块到物理块的映射、引用计数、访问状态共享内存分配器管着物理块的分配、回收、合并最后还有一个后台整理线程定期整理空闲块适当地做内存规整减少物理块的离散程度。这里的部件分工看起来平平无奇但每一层都有各自的取舍。块表条目太大浪费显存做索引太小单次块查找开销就上去了。分配器若不支持回收合并共享内存池用久了还是会慢慢碎裂。下一章展开讲这些细节。3. 核心细节与实操要点3.1 块大小的选取没有唯一答案块大小是整个分页方案里影响最深远的参数。每个块容纳的token数决定了块表的大小、物理块的分配粒度和块内浪费的比例。常见的选择是每个块装16个或32个token太小会让块表特别长查找开销变大太大又会在请求末尾浪费不少空间。我给一个直观对比假设单个token KV数据是512KB如果每块装16个token单块大小8MB块表每项8字节的话支持1M块只需要8MB索引空间完全可以忽略。如果每块只装4个token块表规模直接翻四倍分配器管理更细碎收益却只有末尾那一两个块的空间不划算。块内浪费公式也很简单引用块数向上取整平均浪费约等于半个块所以块越大浪费越严重。我的做法是先按目标并发数和最大序列长度估算峰值块总数再反推块内token数并让单个物理块的大小保持在对齐友好的区间比如2MB到16MB之间稠密的请求场景用16token每块碎片敏感的长尾场景再降到8。3.2 共享内存池的初始化参数怎么定共享内存池的容量设计有个不太严谨但很实用的估算方式峰值并发数 × 平均请求长度 × 每token KV字节数再乘上1.3到1.5的冗余系数。这里的冗余不是浪费而是要留出块表、管理结构、临时分配需要的空间。初始化流程一般分三步先按估算值创建共享内存区域映射到进程地址空间再把整块区域头部的管理元数据写进去包括总块数、空闲块链表、块状态数组最后把所有物理块串进空闲链表。之后所有KV Cache的分配都从这个池子里走不再碰系统级申请。注意第一块区域头部需要预留几十KB做元数据别把第一个块直接从0号地址开始用否则后面写元数据就是自己覆盖自己。3.3 块表与复制写机制的设计要点块表是整个分页方案的命根子。每个逻辑块在块表里至少要有三项内容物理块索引、引用计数、标志位。引用计数指向当前有多少请求在用它只有计数归零才能回收。标志位用来标记块是否被写入过、是否只读共享。公共前缀的请求命中同一个物理块时多个请求的块表条目里同时写着同一个物理块索引引用计数相应增加。新请求继续生成时会往块里写新的KV数据这时候就不能直接改了否则会污染其他请求的缓存必须分配一个新的物理块把原块内容复制过去再写入——这就是复制写机制。它很巧妙代价是复制一次整块数据所以复制写只在分支发生时才触发一次总体上比重复计算省得多。3.4 并发控制不能一把大锁锁死多进程共享内存池的并发控制一开始很容易图省事用一把全局锁实测开发够用一上压测就完蛋。KV Cache的分配和释放频率远高于显存管理全局锁会让所有请求串行排队吞吐直接崩。后来我把控制粒度按逻辑块做了分片。块表条目用每块独立的原子引用计数来保证并发安全物理块的分配回收用了一个无锁的空闲队列待分配时从队头取块释放时还回队尾块表整体只读的部分走无锁查询只有写入新映射时才短暂加锁。再配合每进程一个本地缓存池小块的分配优先走本地不够才向全局池要能大幅度减少跨进程竞争。这套改完并发表现从“勉强能用”变成了“压测下保持稳定”。3.5 池子快满时的驱逐策略共享内存池满的时候不能直接让新请求失败得有驱逐策略。我实现了两级驱逐先按引用计数找零的块直接回收这是最便宜的如果没有空闲块再找那些只被一个请求占用且这个请求暂停等待的块写回或直接作废把空间腾出来。需要特别小心的是不能为了腾空间把一个正在推理中途的请求的核心块替换掉那会导致这个请求后续解码速度断崖式下跌。所以驱逐优先级我设定为空闲块 长空闲时间的只读共享块 等待状态请求的块 活跃解码请求的块最后一级实际上不会触发真到了那一步说明容量设计已经脱离实际情况了。4. 实操过程与核心环节实现4.1 共享内存池初始化完整流程直接上操作序列这是我实际项目中跑通的流程。第一步按前面算好的总容量调用共享内存接口创建并映射一块共享内存区域映射标志要带读写权限同时建议带上MAP_POPULATE之类的标志让物理页提前分配好避免访问时缺页抖动。第二步把共享内存区域首地址按64字节对齐在头部初始化元数据区。元数据区包含魔数、版本号、池容量、块大小、总块数、空闲队列指针。第三步物理块区从元数据区末尾开始把每个块都串进空闲链表。块控制结构可以内嵌在每个物理块的头部几个字节里标记块状态和归属进程ID。第四步调用一个注册函数把池句柄挂在全局状态里。后续所有分配请求都从这个池句柄出发。初始化完成后再让另一个Worker进程join进来它只需要通过同样的共享内存句柄把池子映射到自己的地址空间就能看到完全一致的元数据和空闲队列不需要重复创建。这一步是整个共享方案的精髓。4.2 请求维度的KV Cache分配流程当一个请求进入推理服务调度器给它分配KV Cache的流程大概是这样的。首先根据请求的输入长度计算需要多少逻辑块然后逐块查块表看是否命中已有的公共前缀缓存命中的情况直接把引用计数加一未命中的就从空闲队列取一个物理块注册到块表的对应逻辑条目里。生成过程中每解码一步请求的内容长度增加如果当前最后一个逻辑块满了就再取一个新物理块追加到块表尾部。请求结束后遍历该请求涉及的所有逻辑块把引用计数逐一减一计数归零的物理块还回空闲队列。这里有一个容易被忽略的问题请求提前终止或者因超时被杀死时释放逻辑必须放在析构函数里保证执行否则就会漏块。我踩过的坑是异常分支里忘了释放物理块跑了一天后池子满了所有新请求开始排队等缓存。后来分成两步处理正常完成走统一释放入口异常终止走强制回收入口加日志告警。4.3 块表查询与缓存命中的实现块表的具体结构我用了一层哈希加一层数组。逻辑块编号可以直接作为数组索引每个表项存物理块索引和引用计数另外单独维护一个从token序列哈希到逻辑块起始的哈希表用来做前缀匹配。命中流程是这样的新请求的Prompt先算一个序列哈希查哈希表如果命中从命中的逻辑块开始遍历后续块看物理块索引是否连续可映射全部对上就直接继承引用计数。对不上的部分再从空闲队列补块。这样公共前缀的解析开销只有一次哈希计算加几次数组访问比完全从头解码快几个数量级。4.4 关键伪代码与解释下面这段是分配器核心逻辑的简化版伪代码展示从空闲队列取物理块的核心路径// 从共享内存池取一个空闲物理块 Block* alloc_block(SharedPool* pool) { // 先查本进程本地缓存 Block* b pop_local_free(pool); if (b) return init_block(pool, b); // 本地没有走全局无锁队列 b pop_global_free(pool); if (b) return init_block(pool, b); // 全局也没有尝试驱逐 if (evict_idle_block(pool)) { b pop_global_free(pool); if (b) return init_block(pool, b); } return nullptr; // 池满外层做排队等待 } // 请求结束释放一个物理块 void free_block(SharedPool* pool, Block* b) { if (atomic_dec(b-refcount) 0) return; // 还有其他请求在引用 push_local_free(pool, b); // 放回本地缓存减少跨进程竞争 }init_block里要做的事情包括清零块状态、设置归属句柄、初始化写入游标。驱逐函数里用的是3.5描述的两级策略具体代码里要加锁保护块表写入区但物理块的分配释放走无锁队列所以大部分路径不会碰到锁。4.5 压测结果与效果对比换上这套方案后我在同样的压测脚本下跑了对比。压测环境是两张同规格的GPU卡同一个7B模型一个用原来的静态预留方式另一个用共享内存加分页管理请求内容统一用长公共Prompt开头。指标静态预留方案共享内存池 分页方案KV Cache显存利用率52%94%显存碎片率38%6%并发请求吞吐TPS3149单请求P95首token时延2.8s1.7s单请求P95平均token生成速度41 token/s55 token/s利用率提升来自两个地方短请求不再被预留大段空间公共前缀缓存被多请求共享重复计算基本消失。时延下降则是缓存命中的直接红利很多请求的前半段解码根本不需要真正逐token走一遍直接拿到公共缓存结果接着往下算就行。5. 常见问题与排查技巧5.1 共享内存越界访问的定位方法这类问题很典型某进程写越界污染了池元数据另一个进程随即崩掉错误看起来完全不相干。我排查时用过的一套组合拳是第一元数据区放一个结构魔数每轮调度循环里校验魔数变了立刻打日志锁定崩溃波次第二每个物理块头写入块魔数和归属请求ID崩溃时通过共享内存转储能直接看到污染源第三临时开启守护模式每次块释放时全量校验块内容长度是否匹配抓非法写入。这套代码平时是关闭的只在排查问题时打开校验逻辑本身还会消耗一些CPU不常驻。5.2 多进程并发死锁的实战案例有一次改造后压测一开始就出现所有进程全部卡死的情况。排查过程分三步走第一步看进程状态发现全部阻塞在共享内存互斥锁上第二步在锁的入口处打点记录申请者ID发现释放块时持有锁还去驱赶活跃请求而活跃请求正持锁等着分配新块互相等成了闭环第三步的修复方式是把驱逐流程改成先标记后执行标记不持有锁真正执行驱逐时单独加锁绝不在持有块表锁的同时去等其他请求释放锁。那次之后我形成了一条规矩任何锁的持有函数里只做自己的事别调用可能申请新块或释放旧块的逻辑一旦出现跨锁调用就必须重构。5.3 池容量规划偏差的诊断经验池容量估少了压测时表现是吞吐曲线先平后陡然下降伴随大量分配失败日志估多了显存闲置浪费利用率数据非常难看。诊断时看两个指标池分配请求峰值和历史分配请求平均值。前者决定池子的硬上限后者决定冗余系数该留多少。如果峰值是均值的两倍以上说明请求长度波动太大驱逐策略要更积极如果两者接近冗余系数可以从1.5降到1.2省下来的显存可以扩batch。5.4 精度量化对KV Cache字节数的影响模型量化不只影响权重也会影响KV Cache的存储。KV Cache用FP16存还是INT8存每token的字节数直接差一倍。切换成量化精度后共享内存池的容量和块大小都需要重新设计。我保留的通用做法是容量按量化后字节数计算块大小尽量保持物理块数量不变而不是保持字节数不变这样块表规模不变命中逻辑不用动。量化后KV Cache的精度损失会导致输出质量轻微波动如果拿不准就选带量化和不带量化两套配置做AB对比别直接拍脑袋切。5.5 几个值得长期保留的调优心得总结几条调优经验。第一块大小不是拍出来的是根据每token字节数和目标块数量反推的每token字节数变了就要同步检查第二本地缓存池的容量很关键太大会让全局空闲块分布不均太小又发挥不了减少锁竞争的作用我这边每个进程缓存4个块是均衡点第三规整线程的运行频率不要太高一般几百毫秒一次就够太频繁了CPU空转反而拉低整体吞吐第四监控指标里加一项“引用计数峰值”这个数能直接反映公共缓存的复用程度也是判断驱逐策略是否合理的参考。结尾一个值得长期投入的方向这套优化方案落地后我最大的感受是KV Cache问题不能只在推理框架层面打补丁得从底层内存管理的思想去重新设计。共享内存加的分页机制本质上就是把操作系统管物理内存的那套成熟方法论搬到了模型推理场景里效果立竿见影。最后再分享一个还在试验中的方向把块表也放进共享内存多个实体推理引擎共享同一份块表元数据这样前端路由在做分发时就能直接判断后端缓存命中情况从请求进来到缓存命中的整个链路都能进一步缩短。
📝

华诺云谱内容团队

资深建站顾问 · 行业研究员

10年+企业数字化服务经验,专注智能建站、SEO优化与品牌营销,持续输出建站技巧、行业洞察与营销干货,已帮助5000+企业实现数字化增长。

你可能需要的服务

订阅华诺云谱资讯周报

每周一封,精选建站技巧、SEO与营销干货,直达邮箱。已有 8,000+ 企业主订阅,助你少走弯路。

↑