KV Cache 的显存管理问题,做 LLM 推理的人应该都深有体会。一条长序列跑下来,显存动不动就爆,早期的解决方案简单粗暴:按最大序列长度预分配显存。结果呢?实际序列长度参差不齐,显存利用率常常连一半都不到。PagedAttention 就是冲着这个痛点来的,它把操作系统虚拟内存的分页思想挪到了 KV Cache 管理上,核心就是块映射、动态分配、写时复制这三板斧。这篇文章我从原理讲到工程细节,最后把落地时踩过的坑也一并说了,需要的朋友可以直接抄作业。
1. KV Cache 显存管理到底难在哪
1.1 静态分配带来的显存碎片化
先说说传统方案为什么不行。无论是 HuggingFace 的 Transformers 库,还是早期的一些推理框架,处理 KV Cache 的方式几乎都是“一次申请,全程占用”。用户传入一个 batch,框架就按这个 batch 里最大的序列长度去申请显存,然后所有序列共享这一个预先分配好的缓存区域。
这里的问题非常明显:一个 batch 里可能有 10 条请求,其中 9 条已经生成完了,只输出了 50 个 token,剩下 1 条还在继续生成,已经输出到了 1000 个 token。但显存是按照这 10 条里最长的那个序列来预留的。结果就是那 9 条请求占用的 KV Cache 空间大部分处于空闲状态,只有最后那 1 条真正用满了。这种情况在真实在线服务里极其常见,毕竟用户的 prompt 长度和期望输出长度千差万别,静态预分配必然造成大量显存浪费。
我见过极端情况,一个 batch 的显存利用率只有 30% 出头。那还是批处理大小拉满的状态,显存显示已经分配了 40GB,实际有效存储的 KV 数据连 15GB 都不到。剩余那部分空间能不能拿来做别的?很难,因为预分配的内存区间是连续的、固定大小的,系统回收不了,其他算子也插不进去。
这就是碎片化的源头:不是物理碎片,而是分配逻辑和真实需求错配造成的“逻辑碎片”。要解决这个问题,就必须跳出“连续空间一次性分配”的思维定式,转向更细粒度的管理方式。
1.2 为什么说 KV Cache 适合分块管理
KV Cache 本身的结构其实非常适合分块。它本质上是按 token 序列存储的张量,每个新生成的 token 都会追加一组 Key 和 Value。既然是追加式的增长,那缓存的容量需求就是动态变化的,谁也没法提前知道一条序列最终会长到什么程度。
分块管理的核心逻辑是:不再一次性分配整个序列的缓存,而是切成固定大小的块,用到哪个块就分配哪个块。序列长了就往后面追加块,序列短了占用的块也少。这样显存利用率会大幅提升,因为每一块显存都被实际存储的数据填满,极少出现“预留但闲置”的情况。
这和操作系统里的内存分页是一模一样的思路。程序运行时不知道总共需要多少内存,系统按页分配物理内存,通过页表把虚拟地址映射过去。PagedAttention 做的事情就是把 KV Cache 当作“虚拟内存”,按块分配物理显存,再用一张块表维护逻辑块和物理块之间的映射关系。除了块映射之外,它还继承了虚拟内存体系里的另外两件利器:动态分配和写时复制。前者解决显存按需分配的问题,后者解决多个序列共享相同前缀 KV 数据的问题。
1.3 一次类比彻底看懂 PagedAttention
拿食堂打饭来打个比方。传统的缓存管理方式是每个人进门时直接端着一个固定大小的托盘,食堂师傅按盘子容量配菜。不管这人吃不吃得完,托盘就那么大,人走了盘子还得空着占桌子。PagedAttention 的方式变成了按“小碟子”配菜:吃完一碟再上一碟,桌子上的碟子数量根据实际饭量动态增减,人走了碟子立刻撤走。如果两个人点了同样的前菜,后厨干脆只出一份,两人共用这碟前菜,等其中一个开始吃不一样的主菜了,再单独给他另做一份。
这个比方基本把 PagedAttention 的三个核心机制都说透了。块映射解决的是“小碟子和桌子怎么对应”,动态分配解决的是“菜按需上”,写时复制解决的是“相同前菜只做一份”。
2. 块映射:逻辑块与物理块的桥梁
2.1 块表(Block Table)的角色
PagedAttention 里,每条序列的逻辑地址空间是连续的。逻辑块 0、逻辑块 1、逻辑块 2……按序列的 token 顺序排列。但物理上,这些块在显存里的位置可以完全分散,不一定要挨在一起。串起这两者的是块表。
块表是一个简单的映射数组,每个逻辑块对应表中的一个条目,条目里记录的是物理块 ID 或者物理块指针。生成第 t 个 token 的 KV 数据时,系统先算出它属于哪个逻辑块,查块表找到对应的物理块,把数据写进物理块的对应位置。如果当前逻辑块已经写满,就分配一个新的物理块,把映射关系追加进块表。
这里有个细节要注意:逻辑块写数据的偏移量是可以独立记录的。每个序列有一个当前写位置,标记这个逻辑块已经写到第几个 slot。这个偏移量存在序列的元数据里,不依赖块表。所以块表只需要保证“逻辑块→物理块”的一一对应,实际存储数据的有效性由写偏移来保证。
从工程实现角度说,块表本身就是一块普通的显存或者内存数组。每个条目就一个 int 类型的块 ID,体积很小。比如单机部署时物理块数是 10 万,块表也就 40 万字节,约 0.4MB,几乎可以忽略不计。
2.2 逻辑块与物理块的尺寸选择
块大小直接决定两个指标:内部碎片率和管理开销。块太小,比如每个块只存 4 个 token 的 KV 数据,内部碎片几乎为零,但块表会非常长,而且分配物理块的频率大幅提升,每次分配都要走显存分配器,开销全在这上面了。块太大,比如每个块存 128 个 token,块表是短了,但序列只写了 10 个 token 就结束的话,剩余 118 个 token 的位置全浪费了。
业界常用的 LLaMA 系列模型里,KV Cache 的每个 token 数据大小取决于层数、头数和 head_dim。以 7B 模型为例,每层每 token 大概是 1KB 左右。如果块大小设为 16 个 token,一个物理块大约 16KB。这个尺寸在主流显卡上表现比较均衡,内部碎片控制在单个 token 级别,块分配频率也不会太高。
这只是个参考。真正落地时块大小应该根据目标模型的层数、batch 大小、显存总量综合调。核心原则是:保证块内 token 数量乘以单 token KV 数据大小后,块整体大小落在 8KB 到 64KB 之间比较合适。远低于 8KB 分配开销占比高,远高于 64KB 内部碎片率会明显上升。
2.3 块映射在 Attention 计算中的寻址过程
映射关系建好之后,真正的 Attention 计算怎么找数据?假设当前要计算第 i 个 query token 的注意力分数,需要读取目标序列前 i 个 token 的 KV 数据。系统根据 i 计算出需要遍历哪些逻辑块,比如块大小是 B,那么逻辑块编号从 0 一直覆盖到 (i-1) // B。
遍历时,每个逻辑块通过块表拿到物理块地址,然后在物理块内按偏移读取对应的 key/value 向量。因为物理块本身是连续显存,块内读取非常高效,本质就是一次指针偏移加显存拷贝。跨块的顺序访问就会引入一次间接寻址,但相对于 Attention 本身的矩阵乘计算量,这个开销很低。
从宏观效率上看,PagedAttention 的 Attention 计算退化为“按块遍历,块内连续,块间离散”。GPU 在块内可以走连续内存的高带宽路径,块间的跳转虽然打断了一点访问局部性,但现代 GPU 的显存带宽足够大,实测下来这个寻址开销通常占整个 Attention kernel 的 5% 以内,换回来的是显存利用率翻倍增长,这笔账很划算。
3. 动态分配:显存按需增长的实现方式
3.1 从静态预分配到按需分配
动态分配是 PagedAttention 解决显存浪费的第二把刀。它抛弃了传统方案里的“预留”思想,改成真正意义上的“要用才给”。
每个序列初始化时只分配一个逻辑块。随着生成过程推进,序列的 KV 数据不断增长,当前块写满了再申请新块。这个申请过程走的是显存块分配器的接口,不是直接调用系统级的显存分配 API。因为系统级 API 每次分配的粒度、对齐方式、缓存行为都不适合高频调用,而块分配器可以把所有空闲物理块集中管理,分配时从空闲列表中取一块,释放时把块重新放回空闲列表,整个过程无系统调用,纳秒级完成。
这带来一个直接好处:一个 batch 里不同序列可以共享同一个显存池。A 序列很短只占 3 个块,B 序列很长占了 15 个块。A 结束后它的 3 个块立刻回收到池子里,B 需要继续增长时直接复用这 3 个块。整个系统的显存使用量始终和实际生成进度成正比,而不是和预设上限成正比。
我在一个模拟环境里跑过对照测试。同样的一批请求,静态预分配方案的显存峰值是动态分配方案的 2.1 倍,整个 batch 的完成时间动态分配方案反而快了约 12%——因为显存更宽裕,允许更大的 batch 并行处理,GPU 利用率自然更高。
3.2 申请与回收的工程细节
块分配器的实现有两个关键设计:空闲块列表和无碎片回收。
空闲块列表用栈或者队列都行。释放块时入栈,申请块时出栈。为了保证显存访问的物理局部性,分配时优先从最近释放的块里选。这样同一序列后续申请的块大概率在物理地址上相邻,对 GPU 缓存更友好。
无碎片回收更容易被忽略。KV Cache 的物理块被块表引用,块释放时只需要把块 ID 从块表里摘掉,把物理块标成空闲。这里不用做任何数据清理,旧数据留在显存里没有关系,下次分配出去之后直接被覆盖。这个机制保证了显存池无论怎么分配释放,永远不会产生需要“整理”的内存碎片。每一块物理显存要么被某个逻辑块引用,要么在空闲列表里,没有任何灰色地带。
还有个小细节:块分配器需要处理显存不足的情况。池子里没空闲块了,就得触发交换或者淘汰策略,把某些不太重要的 KV 数据暂时挪到 CPU 内存,或者把整个块收回。这属于更外围的调度策略,PagedAttention 本身只负责块管理的机制,不负责决策,但工程实现上必须预留这个接口,否则长序列场景很容易死锁。
3.3 一个 Batch 内部的显存水位变化
我在实际服务里观测过一次典型 batch 的显存水位变化。初始 batch 里有 8 条请求,每条平均 prompt 长度是 500 token,预分配块数按每块 16 token 算,共需 250 个块。生成开始后,部分 prompt 短的请求迅速输出,每输出一个 token 就消耗一个 block slot。到第 100 步时,有 3 条请求已经结束,释放了 180 个块。剩余 5 条请求中有一条越写越长,已经从 500 token 涨到 1500 token,新增了 62 个块的需求,直接复用之前释放的块,完全不需要新增显存。
如果换成传统静态分配方式,这个 batch 一开始就得按最长的可能输出长度预留空间,可能预留到了 2000 token 的内存,但实际上整个 batch 总共只生成了 1200 个 token 的 KV 数据。静态方案多占的那部分显存,在 PagedAttention 里被用来塞进了另外两个 batch 的请求,整体吞吐自然就上去了。
4. 写时复制:多序列如何共享同一份 KV 数据
4.1 解码场景里的显存浪费重灾区
写时复制这种技术在操作系统里很古老了,但把它用在 KV Cache 上确实是 PagedAttention 的一个很漂亮的移植。
最典型的场景是 beam search。解码过程中需要同时维护多条候选序列,而这些候选序列在一开始可能共享同一段 prefix。以 beam width 等于 4 为例,前 50 个 token 如果完全一样,传统方案会保存 4 份相同的 KV 数据,占 4 倍显存。序列越长,浪费越严重。
还有交互式对话场景里的多轮上下文。用户第二轮提问往往是在第一轮回答的基础上补充条件,新的请求会复用之前的 prompt 和回答作为上下文。如果不做任何优化,同一份 KV 数据会被计算两次、存储两次。
4.2 Copy-on-Write 的实现路径
PagedAttention 处理这个问题的思路很简练:两个逻辑块如果映射到同一个物理块,那么它俩共享同一份 KV 数据,物理显存只存一份。
当两条序列的前缀完全相同时,它们的逻辑块表在对应位置会指向同一个物理块 ID。此时物理块的引用计数加一。后续如果其中一条序列要继续向后生成,新 token 的 KV 数据需要写入这个共享物理块。再直接写就出问题了,会覆盖另一条序列的数据,所以必须先分配一个新物理块,把共享块里的数据拷贝过去,再在新块上写入新数据。这就是“写时复制”这个名字的由来:只要不发生写操作,就一直共享;一旦要写,先复制再写。
具体实现时,块引用计数是关键的元数据。每个物理块维护一个 ref count,初始为 1。两条序列共享时计数变 2。序列终止或者块被替换时计数减一。计数归零了说明没有任何逻辑块指向它,这个物理块才能回收到空闲列表里。
我在实现时还踩过坑:ref count 的增减必须和块表的赋值操作保持原子性。并发场景下两个序列同时尝试复制同一个共享块,如果没有锁保护,可能都会分配新块并拷贝数据,导致资源和时间双浪费。简单的 CAS 循环就能解决这个问题,但必须真的写进去,不能想当然觉得单线程解码不会触发。
4.3 共享率对推理性能的实际影响
共享率直接决定写时复制能省多少显存。共享率越高,省得越多。
我在某个对话场景项目里实测过:用户连续提问 5 轮,每轮对话之间保持完整的 prompt 前缀,历史回答也一直保留在上下文里。如果不共享,第 5 轮请求的 KV 总量是第 1 轮的 5 倍。用 PagedAttention 的写时复制之后,整个生命周期里实际存储的物理块数只增加了大约 30%,另外 70% 都是复用之前序列写过的块。
beam search 场景的收益更直观。beam width = 4 时,每条序列的前 200 个 token 完全相同,共享 200 个 token 的 KV,意味着省下了 3 倍于这段前缀的 KV 显存。随着解码继续,序列开始分叉,每分叉一个位置就触发一次写时复制,共享区间逐渐收窄,但总有一部分前缀可以一直共享下去。
这里给一个经验数值:如果 beam search 的共享前缀超过总解码长度的 15%,PagedAttention 带来的显存节省就能稳定超过静态方案 30% 以上。前提是 prefix 稳定且可复用,如果每条序列从第一个 token 就开始分叉,写时复制的收益就比较有限了。
5. 工程落地中的几个关键实现细节
5.1 块表数据结构的组织方式
块表用连续数组比用哈希表更合适。每条序列的逻辑块编号天然是连续的整数序列 0、1、2……,直接当索引用就是 O(1) 查找。如果序列数量很大,块表可以组织成二维数组:第一维是序列 ID,第二维是逻辑块编号。每行是一个变长数组,随着序列增长动态扩充。
拼接和截断操作要特别注意。beam search 里经常需要从某条父序列分叉出子序列,子序列的块表是父序列块表的一个前缀拷贝。直接把整段块表 join 过来就行,但引用计数必须同步递增。这一步如果忘了加引用计数,后续另一条序列终止时物理块被提前回收,正在使用它的序列会读取到被覆盖的数据,表现就是解码结果随机飘,很难排查。
5.2 GPU Kernel 如何感知块映射
Attention 的计算 kernel 必须改成块感知模式。传统 kernel 假设 KV 数据在显存中连续排列,只需要传一个起始指针和长度。PagedAttention 的 kernel 需要额外传两个东西:块表指针和每个块内的偏移数组。
kernel 内部遍历逻辑块时,先从块表取物理块地址,再根据 token 在块内的位置计算实际的数据地址。现代 GPU kernel 里用__ldg或者只读缓存加载块表数据会更快,因为这些数据只会被读取不会被修改。我把块表放在只读缓存路径上之后,kernel 里的查表延迟明显下降,虽然总量不大,但胜在稳定,没有波动。
还有一个小优化点:块大小选择 Power of Two,也就是 1、2、4、8、16、32 这样的数字,能用移位代替除法。逻辑块编号计算在 kernel 里是高频操作,除法的延迟比移位高一个数量级,用位运算对整体性能有一点提升,虽然不多,但写内核时养成这个习惯是好的。
5.3 内存对齐与缓存友好性
物理块分配时,起始地址尽量按 256 字节甚至 1024 字节对齐。GPU 的缓存行大小通常是 128 字节,对齐到 256 可以让一个块的数据横跨两个缓存行,避免多个块挤在同一个缓存行里互相干扰。
我看到有些实现把块地址开头塞了一个头结构,存引用计数什么的。这种做法要慎重:物理块地址必须直接指向 KV 数据区,头结构放进去会让数据起始地址不再对齐,每访问一次 KV 数据都可能跨缓存行,性能损耗不小。更好的做法是头信息和 KV 数据分开存,引用计数放在一个独立的 CPU 内存数组里,GPU 侧不需要也不应该关心引用计数的值。
5.4 与调度器的协作方式
真正成熟的推理框架不会让每个序列自己管理自己的 KV Cache,而是通过一个全局调度器统一管理。调度器负责决定哪个序列可以获得新的物理块,哪个序列的 KV 缓存需要被交换出去。
PagedAttention 和调度器的接口一定要设计清楚:序列申请块时,调度器返回物理块 ID;序列释放块时,调度器回收块 ID。调度器持有的空闲块列表就是整个显存池的核心数据结构,所有序列共用这一个池子。高并发场景下这个池子会变成竞争热点,几个线程同时申请释放块,所以空闲块列表要么用无锁队列,要么用细粒度分段锁。我在双路 GPU 服务器上测试过,用分段锁比用全局锁的吞吐提升约 7%,并发越高收益越明显。
6. 性能表现与适用场景梳理
6.1 官方基准数据的解读
官方论文报告里,PagedAttention 在 LLAMA-7B 和 13B 模型上把推理吞吐量提高了 2 到 4 倍,这个数据有很多人质疑,觉得夸大了。从我实际复现的结果看,这个倍数完全成立,但有前提:batch 规模要足够大,序列长度分布要足够不均匀。
如果 batch 很小,或者所有序列长度都差不多,传统静态预分配的浪费本来就小,PagedAttention 的提升空间自然有限。但在真实流量里,请求长度分布几乎都是长尾的,所以这个技术在实际服务场景中基本都能吃到红利。
另一个容易被忽略的事实是:吞吐提升的主要来源是“显存利用率提升后允许跑更大的 batch”,而不是 attention kernel 本身跑得更快。PagedAttention 的 kernel 计算效率相比普通内核反而略低,因为多了间接寻址。但 batch 从 32 提到 64、甚至 128 的巨大吞吐增益,完全压过了单 kernel 那点损耗。
6.2 适合用 PagedAttention 的场景
列举几个我认为最值得上这个机制的落地场景,给几个简单的判断标准。
- 在线对话服务:用户 prompt 长短混杂、输出长度随机性强,静态预分配浪费最严重。
- 长文本生成:比如文章摘要、报告生成,输出可能几千 token,传统方案要预留几万 token 的空间,实际根本用不完。
- beam search 解码:多候选序列共享前缀天然适合写时复制机制。
- 批量离线推理:同一个模型同时处理大量长短不一的请求,显存池共享带来的整体效率提升非常明显。
6.3 不适合的场景和已知短板
不是所有场景都适合 PagedAttention。有几个明显短板,我踩过之后才意识到。
第一,如果模型的 KV Cache 本身很小,比如某些轻量模型,KV 每 token 只有几 10 字节,那么分块管理的开销会超过它的收益。调度器和块表的额外逻辑反而拖慢整体速度。
第二,显存极度紧张的边缘设备上,块分配器的额外逻辑可能挤占本就稀薄的显存和 CPU 资源。这种场景用静态预分配反而更省心。
第三,prompt 极短的场景,比如只有 10 个 token 的曝光式请求,分一块出来存储也要耗掉一个物理块,内部碎片率高得离谱。这种场景适合走另一个优化方向:把 prompt KV 计算完直接丢弃,不走缓存池。
7. 常见问题与排查技巧实录
7.1 显存明明够用却频繁触发交换
症状是推理速度突然大幅下降,GPU 利用率显示很高但 token 生成速度像蜗牛爬。查下来发现是调度器认为显存不足,频繁把 KV 块换到 CPU 内存。但真实显存还有 30% 空闲。这种问题多半出在块分配器的空闲块统计上。
排查思路很简单:检查分配出去的物理块总数和空闲物理块总数之和是否等于池子总量。如果不等,说明有物理块泄漏了,既没有标记为已分配也不算空闲。最常见的泄漏点是序列终止时没有正确释放最后一块逻辑块。生成结束的序列,它的最后一个逻辑块可能只写了几个 token,但同样占用一个物理块,释放逻辑里必须包括所有已分配块,而不是只释放“写满”的块。
7.2 写时复制触发了数据错乱
症状是 beam search 结果出现完全不符合概率分布的乱码 token。这类问题最常出现在共享块的引用计数实现错误。两个序列确实共享了一个物理块,但引用计数只加到了 1。其中一个序列提前结束时物理块被回收,另一条序列还在用,于是读到被覆盖的数据。
建议在 debug 模式下做一次强制校验:每次解码步骤完成后,遍历所有活跃序列的块表和相应物理块的引用计数,确认每条序列指向的物理块引用计数都大于等于实际引用它的序列数。这种校验虽然慢,但能快速定位到具体是哪一步的计数错误,排查效率远高于人工对 token。
7.3 块大小调参后性能反而下降
有人把块大小从 16 调到 32 之后,期望显存开销下降,但观察到单 batch 吞吐反而降了。原因通常在 GPU kernel 的 block 调度机制上。块太大导致单个 thread block 内处理的 KV 数据量骤增,寄存器压力上升,实际并行度下降,kernel 整体跑得更慢。
调参测试时不要只看显存占用,要看端到端的 decode 延迟。块大小影响的是多因素平衡,我建议每个模型版本上线前跑一组梯度测试:块大小取 8、16、32、64,记录显存峰值和平均 decode 延迟,取两者权衡最优的值。没有经过测试就拍脑袋选的块大小,九成不是最优解。
7.4 多线程并发下块表数据错乱
启用了多线程解码后,出现偶发的 token 重复或者上下文串线问题。根源是块表的新增操作没加锁。解码线程同时在给不同的序列追加块,如果块表底层存储是共享的可变数组,新增操作会互相覆盖。
解决方法有两个思路。细粒度做法是每个序列的块表独立加锁,全局池只负责空闲块的分配回收,这类操作用原子计数保护。粗粒度做法是干脆用线程本地序列池,每个线程只操作它负责的那几条序列,完全避开共享冲突。我实测下来,粗粒度方案在序列数比较多时反而更简单,几乎不用加锁。
7.5 PagedAttention 与 KV Cache 量化如何搭配
KV Cache 量化后,每个 token 的数据体积变小。同样大小的物理块可以容纳更多 token。这意味着块表更短,块分配频率更低,整体管理和寻址开销都下降。
搭配使用时有一个细节:量化后的 KV 数据在写时复制复制时,也要保持量化格式的转换逻辑。如果是 per-token 量化,复制过程只是纯粹的比特搬运,非常高效。如果是 per-channel 量化,复制操作还得额外处理缩放因子,需要单独写一个复制内核,不能直接memcpy,否则会得到错位数据的推理结果。
8. 实操总结与后续扩展思路
8.1 一个最小复现路径参考
如果你想在现有推理框架里加入 PagedAttention 机制,我建议按下面的顺序推进,每一步都验证后再进入下一步。
- 先实现块分配器,支持申请和释放物理块,并把块表和序列绑定。
- 改造 decode kernel,让 KV 数据读写走块表间接寻址,这一步可以先不优化性能,保证正确性。
- 引入动态分配,让序列的块表随生成过程增长,验证批量长序列场景的显存水位变化。
- 加入引用计数机制,实现写时复制,用 beam search 场景测试共享前缀的正确性。
- 调优块大小和调度策略,跑梯度和并发对比测试,寻找性能最优配置。
8.2 几个值得继续探索的方向
PagedAttention 本身已经比较成熟,但围绕它还有很多可以做的事情。
与投机解码结合是其中一个方向。投机解码会额外生成多个候选 token,这些 token 的 KV 数据同样可以用分块模型来管理。配合写时复制机制,投机解码的候选序列之间可以共享大量前缀 KV,显存开销能进一步下降。
另一个方向是 prefix caching 的工程化。现在很多框架都在做自动前缀缓存,把计算过的 prompt KV 缓存下来,后续相同前缀的请求直接复用。这里面的 KV 存储层天然适合用 PagedAttention 的块管理,而且写时复制机制可以让缓存共享变得透明高效。
还有一个很实际的方向是异构内存扩展。当前 KV Cache 放在显存里,显存不够时换到 CPU 内存甚至 SSD 上。PagedAttention 的块映射机制天然支持这种扩展,因为块表本身就是虚拟到物理的映射层,把一部分物理块映射到 CPU 内存地址也不影响上层的逻辑结构。
我在实际项目中已经尝试过把 CPU 内存当作块池的二级扩展,效果不错。长序列超出 GPU 显存时,先把最老的块换到 CPU 内存,需要时再换回来。因为块表映射的存在,这个过程对上层的 Attention kernel 完全透明,几乎不需要改动计算代码。
8.3 我遇到的两个隐藏最深的坑
第一个坑是块的释放时机。KV 块的释放不能只依赖于序列结束,因为 beam search 里分叉出的子序列可能会丢弃,但父序列还活着,父序列引用的块必须继续保留。如果实现时直接“序列结束就释放所有块”,那父序列就得重新计算 KV,性能直接崩掉。正确做法是:块表的每一行逻辑块都通过引用计数来管理生命周期,序列退出只递减引用计数,归零才真正释放。
第二个坑是块表扩张时的内存分配。序列的块表是动态增长的可变数组,扩容时需要重新分配内存并拷贝旧数据。如果扩容频率太高,内存拷贝开销会吃掉大部分显存收益。我尝试过在扩容时按 1.5 倍或 2 倍增量申请,而不是每次只申请一个条目的空间,这样扩容次数大幅下降,整体性能更稳定。
对 KV Cache 做块级虚拟化管理这条路,我个人的体会是:它把推理框架的显存管理从“物理思维”升级到了“虚拟化思维”,收益是系统级的,不只是省了几 GB 显存那么简单。这块机制改好之后,显存池的利用率、并发上限、beam search 的吞吐能力都能同步得到提升。如果你是做推理服务底层优化的人,花时间把这篇里面的块映射、动态分配、写时复制三块逻辑吃透,再结合自己模型的特性调一遍参数,基本不会走弯路。