1. 显存账本连续KV Cache为什么会拖垮高并发推理1.1 从一次推理算一笔显存账PagedAttention 这个名字第一次撞进我视野是在给一个并发推理服务做显存优化的时候。当时服务一上量GPU 显存就像漏了底的水桶请求队列越拉越长明明显卡不算差却只能同时跑几个请求很多显存被白白占住却没有产出。后来把 KV Cache 从连续分配改成块式分配同样的卡能多扛三倍以上的并发请求。这套机制看着不起眼实际上是把大语言模型推理服务里的“内存碎片”问题彻底换了个解法。先别急着看实现把账算清楚才知道它到底解决了什么。推理阶段每生成一个 token都要读取前面所有 token 对应的 Key 和 Value这个数组叫 KV Cache。它不参与训练却几乎占据推理时的大部分显存。假设一个模型有 32 层 Transformer每层 32 个注意力头每个头维度 128使用 float16 存储那么一个 token 对应的 KV Cache 大小是32 层 × 2Key 和 Value 两份× 32 个头 × 128 维 × 2 字节 512KB这不是夸张一个 2048 token 的序列光 KV Cache 就是 1GB4096 token 就是 2GB。一个并发 32 的推理服务如果每个请求都是 4K 长度光 KV Cache 就要吃掉 64GB 显存。真实场景里还没算模型权重、激活值和 CUDA 上下文KV Cache 确实是最大的显存消耗项。很多早期的推理实现走了最简单的路线每个请求按最大生成长度预分配一整块连续显存生成到哪算哪。这个方案的毛病不是一步出来的是慢性和急性混合发作。1.2 连续分配、预留浪费和外部碎片连续分配最大的问题有两个一个是预留浪费一个是碎片化。先看预留浪费。调度器不知道一个请求到底会生成多少 token只能按上限猜比如 2048 或者 4096。但真实请求往往是长短混合的一大半请求实际几百个 token 就结束了预留的长长度空间中后半段全空着。显存被占了却没有任何计算价值这就是内部碎片。打个生活化的比方你每天往一个饭盒里装饭菜饭盒按最大食量买但你实际每顿只吃一半剩下的空间每天都在浪费。再看外部碎片。并发场景下每个请求的长度不同先到的请求占了一块连续区间中间某个请求提前结束释放出来的空洞是零散的。后到的长请求需要一大块连续地址很可能就塞不进这些空洞分配器要么把请求排在后面要么触发显存整理。整理这件事在 CPU 内存里已经很麻烦了在 GPU 显存上更难受因为要暂停所有计算然后拷数据对在线延迟是致命的。连续分配还带来一个隐性约束新请求的 KV Cache 必须地址连续这导致并发的上限被“空洞”卡住而不是被实际的可用显存大小卡住。明明显存里还有很多字节空闲但因为它们不连续反而没法用这大概是推理服务最容易遇到也最让人恼火的问题。后来我把目光放到块式分配思路上的原因也很简单既然操作系统能靠分页把物理内存碎片藏起来用户进程看到的是一个连续的虚拟地址空间那 KV Cache 为什么不能做成同样的机制于是就有了接下来的块映射、动态分配和写时复制这三个核心部件。2. 块映射一张块表把逻辑块和物理地址串起来2.1 逻辑块与物理块两个不同的视角PagedAttention 的全部秘密可以概括成一句话不再为每个请求分配一段连续的显存而是把 KV Cache 切成大小固定的小块每个请求通过一张块表把“逻辑上的连续序列”映射到“物理上可能不连续的块”。这里涉及两个视角。从模型推理的角度看序列是从第 0 个 token 一路排到第 L 个 token 的连续时间步注意力计算需要逐个读取前面所有 token 的 Key 和 Value。从显存分配的角度看这些 token 的 KV 数据存在一组物理块里物理块在显存中的地址不一定相邻二次分配时也完全可以分布在显存各处。逻辑连续不等于物理连续这就是分页系统的核心思想。一个块通常包含固定数量 token 的 KV 数据比如 16 个 token 或者 32 个 token。每个物理块对应一段实际显存块与块之间不需要连续。请求持有的是一个块表Block Table它是一个数组数组里第 i 项存的是逻辑块 i 对应的物理块地址或物理块 ID。查找某个 token 对应的 KV 数据时先算出它在哪个逻辑块的第几个槽位再通过块表把逻辑块号换成物理块号最后定位到具体显存地址。举一个很简单的例子块大小 B4一个序列总共有 9 个 token逻辑块 0 包含 token 03逻辑块 1 包含 token 47逻辑块 2 包含 token 8块表为 [物理块 7, 物理块 1, 物理块 9]也就是说模型眼里的 token 08 是连续的但实际数据可能散落在物理块 7、物理块 1、物理块 9 里面。物理块 7 后面紧跟着物理块 8但序列的下一个逻辑块却映射到了物理块 1完全不挨着。这有个很直接的收益当某个请求结束后归还给分配器的是单个物理块而不是一大段连续内存其他请求可以马上复用这些零散的块。我举个例子说明这个收益。假设物理块 2 和物理块 5 都是空闲的虽然它们在地址上隔得很远但新请求的逻辑块 3 可以映射到物理块 2逻辑块 4 可以映射到物理块 5。只要块表里记清楚对应关系哪怕物理地址再乱也能组合出逻辑上连续的序列。外部碎片问题被彻底绕开了剩下的只是单个物理块内部最后一小块浪费。2.2 内核侧的块遍历和索引开销块映射不只是换了个存储方式注意力计算本身也得跟着改。以前 KV Cache 是一整块连续显存内核可以直接按线性地址从 0 扫到 L访问跟数组遍历一样简单。改成块式存储后内核拿到的不再是“起始地址 长度”而是一个块表数组加上序列长度。计算某个 token 时内核需要先算出逻辑块号block_idx token_idx // B以及块内偏移offset token_idx % B然后用 block_idx 去索引块表得到物理块地址最后按 offset 定位到块内 KV。伪代码大致长这样block_table [p0, p1, p2] # 逻辑块索引 - 物理块地址 BLOCK_SIZE 32 def load_kv(token_idx: int): block_id token_idx // BLOCK_SIZE offset token_idx % BLOCK_SIZE physical_addr block_table[block_id] return read_kv(physical_addr, offset)在 CUDA kernel 里这个逻辑通常会展开成两层循环。外层循环遍历块表里的每个物理块内层循环处理块内的 token。为了减少分支内核一般会一次性加载整个块的 KV然后通过掩码判断哪些 token 是有效位置哪些位置属于 Padding。第一次改的时候我踩了一个认知上的坑总觉得多了一次查表就会让注意力计算变慢。实际上 GPU 上多做的不是分配动作而是内存寻址动作数据还是从全局显存读计算也不变复杂。真正需要关注的开销在两层一是块表本身如果存放在 CPU 内存每次 kernel 启动前要上传到 GPU二是块表条目数量太多时L2 缓存可能放不下反复命中全局内存会拖慢一点读取速度。一张块表大概多大一般一个物理块的编号用 4 字节或 8 字节存储假设一个服务管理 100 万个物理块块表数据大约 8MB。相比几十 GB 的 KV Cache这个开销可以忽略。真正不能忽略的是块粒度带来的内层计算效率后面我会专门讲块大小怎么选。2.3 块表在主存还是在 GPU 上更合适工程实现时一定会遇到一个问题块表放哪。我的建议是分配和管理逻辑放 CPU把“当前请求”的块表作为轻量张量传入 GPU。原因有三个。第一分配器运行在 CPU 侧它需要在每次生成完一个 token 后决定是否分配新块、是否触发写时复制这些决策依赖 CPU 端的数据结构如果块表放在 GPU 上CPU 每一步都要发起跨端查询延迟没法接受。第二同一物理块可能同时被多个请求引用引用计数和映射关系的变化频率很高这种几十万次级别的指针操作明显更适合 CPU 内存完成。第三kernel 侧并不需要修改块表它只要读取所以 CPU 在生成完新块之后把块表同步到 GPU 就行了不需要复杂的双端互斥。如果块表放 GPU也不是完全不行但它会让分配器变成瓶颈。所有空闲块的管理、引用计数的增减、写时复制的判断都要在 GPU 上原子操作调试起来非常痛苦。做工程不是追求极端而是要找到简单稳定的方案块表放 CPU、按请求上传到 GPU 是我实测下来最省事的做法。3. 动态分配按需给块而不是按最大长度给满3.1 按需分配的生命周期和空闲块池块映射解决的是“地址不连续能不能用”的问题动态分配解决的是“块到底什么时候给”的问题。传统预分配是请求一开始就按最大长度占满显存动态分配则遵循一条朴素原则用多少分多少生成一个 token 的 KV才占一个 token 的空间。一个请求从进入服务到离开块分配大概会经历这样几个阶段请求到达初始化 KV Cache此时序列还没生成任何新 token但为了写第一个 token 的 KV需要先分配一个物理块。每生成一个 token把它对应的 Key 和 Value 写入当前块的槽位。当前块写满了分配器再从空闲池里拿出一个新块把它加到这个请求的块表末尾。请求结束或者被抢占它名下的物理块逐个归还到空闲池。这个流程里最关键的一点是分配动作的触发点是“写入需求”不是“长度预期”。一个请求生成到第 2 个 token 就结束那它从头到尾只需要一个物理块一个请求生成 5000 个 token块表才会一次次增长。显存的使用量和实际计算量严格同步不再有任何预留浪费。分配器内部一般维护一个空闲块池。物理块被释放后不会真正还给操作系统而是放进池子里等下一个请求拿走。这样做的好处是避免频繁申请显存毕竟一次显存分配可能涉及同步整个设备代价比 CPU 内存分配大得多。我见过的实现里空闲块池是一个简单数组或链表分配时取一个块释放时把块地址塞回来操作本身不复杂难的是并发环境下的多个请求同时分配时要保证正确性。如果你自己写这块逻辑我建议做一个非常朴素的模拟器先在 CPU 内存里用 Python 把块表、空闲池、请求生命周期跑通再考虑 CUDA kernel 的细节。我后来在公司项目里就是先用模拟器复现各种边界情况才避免了后面改 GPU 代码时反复返工。3.2 最大可分配块数和高负载下的取舍动态分配不等于无限分配。物理块的数目会受到显存容量约束我们可以用一条公式控制全局上限max_num_blocks kv_cache_allocated_size / (block_size * bytes_per_token)其中 bytes_per_token 是每个 token 的 KV 大小block_size 是块内 token 数量。比如某台 GPU 上给 KV Cache 留了 40GB 显存块大小 32每个 token 512KB那么最大物理块数就是 40 × 1024³ / (32 × 512 × 1024) 2560 个左右。这个上限必须在服务启动前算清楚否则高并发下所有块都被占满新的请求就会互相挤占内存。块被占满以后怎么办策略通常有两种排队或者抢占。排队最简单新请求等着等前面的请求释放块再进但对在线服务来说排队意味着延迟上升用户体验变差。抢占则是把某个正在运行的请求暂停把它已经生成的物理块回收优先给新请求使用等显存有裕量了再重新恢复被抢占的请求。实现抢占时最麻烦的是恢复不过因为有了块表和动态分配恢复过程只需要重建逻辑块到物理块的新映射不需要搬移原始 KV 数据。这一块是分页思想最漂亮的地方物理块可以重新映射逻辑连续性始终由块表保证。还有一个细节值得注意全局上限并不是越小越好。上限设置得太低服务会频繁触发抢占GPU 的大部分时间花在切换请求上真正算注意力的时间反而变少。上限设置太高又可能挤压模型权重和激活值的显存空间。比较稳妥的做法是先跑几个典型长度的压测请求观察实际 KV Cache 峰值再留 20% 的余量。不要拍脑袋定数字用数据说话。4. 写时复制让多个分支共享前缀 KV代价降到最低4.1 并行解码带来的前缀共享需求如果说块映射和动态分配解决的是“怎么省显存”那写时复制解决的就是“怎么让同一个前缀被更多人共用”。这个需求是从解码算法里长出来的。在自回归生成里经常需要同时探索多个候选序列。最常见的两个场景一是束搜索多个候选之间共享同一个开头只是后段分叉二是并行采样同一个 prompt 同时生成多个不同的续写结果。这两个场景有一个共同特点就是大量序列在很长一段前缀里完全相同。拿束搜索举例。假设两个候选序列都来自同一个父序列父序列已经生成了 4000 个 token接下来要分支成 4 个候选。如果不做共享每个候选都要存一份前面 4000 个 token 的 KV四个候选就是四份一模一样的显存而它们的内容完全相同纯粹是重复存储。按前面算过的一 token 512KB 来算四个候选光前缀就多占 4 × 4000 × 512KB 8GB 左右这还没算后续生成的部分。如果把这些相同前缀的 KV 存在同一组物理块里四个候选共享同一块显存那么 8GB 就可以降到 2GB因为只需要一份 4000 token 的 KV 加上四个候选各自新生成部分的 KV。这就是前缀共享的核心收益。共享的粒度自然是物理块而不是单个 token原因很简单块表可以在多个请求之间保持一致的映射申请新块时按块登记引用即可管理成本低得多。4.2 写时复制触发条件和引用计数共享了一个物理块之后最怕的事情是一个请求往块里写新数据把其他请求的前缀数据冲掉。写时复制Copy-on-Write就是用来防止这个问题的。需要区分两种物理块完整的块和不完整的块。一个完整的块意味着它已经写满 B 个 token内容确定不变多个请求共享它是绝对安全的因为没有人会往里面写新东西。真正危险的是最后一个未写满的块后续分支需要往这个块的剩余槽位写入新 token如果多个分支都往同一个物理块写那就互相覆盖了。那怎么处理未满块实现上引用计数是标准做法。每个物理块维护一个 ref_cnt表示有多少个请求正在引用它。当一个请求需要继续写入最后一个块时先检查这个块的 ref_cnt如果 ref_cnt 1说明只有当前请求在用放心写入。如果 ref_cnt 1说明还有其他分支在共享这个块不能直接写否则会覆盖别人。此时分配一个新物理块把旧块已有的 KV 数据复制过去更新当前请求的块表指向新块然后把旧块的 ref_cnt 减一新块的 ref_cnt 初始为一。这就是写时复制的完整动作复制发生在我们真正要“写”的那个时间点而不是分支创建的那一刻。分支创建时所有分支的块表都指向同一组物理块一个字节都没移动过。只有当前分支要往后接新 token才把最后那个未满块单独复制出来。之前所有完整的块继续共享。从实现上理解这段话可以看这个简化逻辑def append_token_kv(seq, token_kv): last_block seq.block_table[-1] if last_block.full: new_block allocator.allocate() seq.block_table.append(new_block) write_token(new_block, 0, token_kv) elif last_block.ref_cnt 1: new_block allocator.allocate() copy_kv(new_block, last_block) seq.block_table[-1] new_block last_block.ref_cnt - 1 new_block.ref_cnt 1 write_token(new_block, last_block.used_slots, token_kv) seq.update_gpu_block_table() else: write_token(last_block, last_block.used_slots, token_kv)这里有个工程细节容易忽略块表更新后必须同步到 GPU不是只改 CPU 端结构就能让 kernel 拿到新地址。如果你的实现里每次生成一个 token 都要同步一次块表同步本身会成为性能热点。更常见的做法是每一步生成时批量处理多个 token或者只在交换到新块的那一步同步让块表更新的次数远小于 token 数量。4.3 完整共享块的不可变性与安全性写时复制看起来很美好但安全前提只有一个完整的共享块必须是不可变的。这个前提一旦被破坏共享的所有收益都会变成 bug。我见过一个新手实现会犯的错为了省一次分配在某个完整块里把一个 token 的 KV 覆盖成新值。这样做的后果是整个服务的结果随机的因为其他分支读到的 KV 已经被污染了。要避免这种问题设计上必须把“写”动作严格限制在当前请求的最后一个未满块内。完整块一旦写满就把它视为只读数据任何请求都不能再去修改它。因为完整块是只读的多个 kernel 同时读取同一个物理块完全不需要加锁。GPU 端的读操作天然不会互相干扰而写操作又被调度器限制为“每个序列每步最多写一个 token”所以不会出现两个 warp 同时写同一个共享块的竞争。这个设计之所以能保持正确性不是靠复杂的原子操作而是靠调度层面的约束写时复制发生在 CPU 侧kernel 侧只读共享块。这里再补一个容易踩的坑引用计数的回收顺序。当请求结束时要遍历它的块表对每个物理块做 ref_cnt - 1只有 ref_cnt 归零的块才真正进入空闲池。如果请求在运行中被抢占它的块也要参与引用计数回收否则会出现“块被释放了但 ref_cnt 还挂在树上”的泄漏。调试这种问题最好的办法是在分配器和释放器里打日志每次 ref_cnt 变化都输出一行对照请求生命周期看谁多减了一次、谁漏减了一次。别觉得日志慢上线前抓三天 bug 的成本比写日志高得多。5. 工程里的关键参数块大小、碎片率和选型判断5.1 块大小不是越大越好也不是越小越好块大小是整个机制里最微妙的参数它直接决定内部碎片、块表长度和 GPU 计算效率三者的平衡。先看内部碎片。动态分配只做到“按块”分配没有做到“按 token”分配。如果一个块能装 32 个 token但当前请求只写了 5 个 token 就结束那这 27 个 slot 就浪费了。这个浪费对每个请求平均来说大概是半个块块越大浪费越多。比如平均每个请求 500 token块大小 64尾部浪费约 6%块大小 32 时尾部浪费约 3%块大小 16 时只有约 1.5%。在 40GB 的 KV Cache 里这百分之零点几都是几百 MB 的显存。再看块表长度。块越小同样容量的 KV Cache 需要的块数越多块表也越长kernel 遍历块表的开销越大。更重要的是GPU 里注意力计算如果每个块只装 8 个 token那么在同一个线程块内能做的并行计算变少算术强度和访存效率都会下降。块越大单次加载连续数据的效率越高kernel 越容易吃饱。我个人的实测经验是通用场景下 32 或 64 是比较稳的选择。短请求多、平均长度在 100500 的对话场景16 或 32 更合适因为尾部浪费能降低不少长文本场景平均长度几千甚至上万64 甚至 128 都可以接受因为尾部浪费占比很小而大块的 kernel 吞吐更高。具体怎么定最好用真实请求长度分布跑一次压测不要只试一个值。5.2 碎片率怎么观测很多人一听动态分配就觉得碎片问题消失了其实只是“外部碎片”消失了内部碎片还一直存在。观测碎片率最直接的指标叫做块利用率定义是block_utilization 实际写入的有效 token 数 / (已分配块数 × 块大小)块利用率接近 1说明每个块都装得很满显存没有任何浪费如果只有 0.4说明有六成显存其实是空的。服务上线后我建议把这个指标和并发数一起监控。正常情况下每个请求只有最后一个块可能不满所以块利用率应该接近 1 - 块大小 / (2 × 平均序列长度)。如果实测值明显低于这个理论值说明有大量请求没走完就被抢占或提前结束分配器把块分得太碎可能需要调整抢占策略。除了块利用率还有一块表命中率值得看如果块表太大GPU 的 L2 缓存放不下每次取物理地址都可能访问全局内存。症状是压力测试时计算时间没有明显增加但显存带宽被拖得很高。这时要么把块大小调大要么调整缓存让块表尽量留在 L2 里。单独处理这个问题时最容易走弯路我建议先确认瓶颈是计算还是带宽再决定动手方向。5.3 什么样的服务值得接入这套机制最后说一个反直觉的判断不是所有推理服务都应该用 PagedAttention。如果服务场景是单条长文本、并发很低、长度也很稳定比如一次只跑一个离线任务传统连续预分配就够用接入块式分配反而会引入额外的块表和分配器复杂度。PagedAttention 最大的优势体现在高并发、变长请求、大 batch 的服务上。判断标准可以很朴素看并发乘平均长度的积也就是 KV Cache 总需求量。几千兆显存规模的场景值得接只有几百兆的场景优化收益有限简单的连续缓存加预分配更容易维护。还有一个场景我觉得特别适合写时复制就是 beam search 这类天然需要共享前缀的生成算法它带来的显存节省量是整个机制里性价比最高的部分。6. 几个容易被忽略的实现细节6.1 分配器要区分“保留块”和“共享块”我在第一次实现的时候把所有块都放同一个空闲池结果出现一个隐蔽问题写时复制复制的块被当成普通块使用导致提前打满显存。后来把块分为两类状态一类是普通独占块一类是被共享引用的块。释放时只有 ref_cnt 减到 0 的块才回池子这对独占块和共享块都适用但分配时优先从独占块空闲池里拿避免共享块刚被释放又被复用导致引用计数历史错乱。这是很小的设计决策却能避免很多莫名其妙的 bug。6.2 抢占恢复时的块表重建抢占时被抢请求的块可能已经映射到物理块恢复时要重新分配物理块并把块表写回。这里容易发生一个问题如果直接复用原来的物理块新请求可能也在用数据会被覆盖。正确做法是恢复时重新走一遍分配逻辑为这个请求申请新的物理块然后从原来的 KV 数据拷贝回来。是的恢复过程会引入一点额外拷贝但相比把整个 KV Cache 做大几倍这个拷贝成本完全可接受。6.3 把调试开关留好再上线上线前一定要在分配器里留一个 debug 模式输出每个物理块的分配时间、请求 ID、ref_cnt 增减时间和释放时间。这套机制的所有 bug 都和生命周期的错乱有关靠定位代码很难一眼看出问题但日志会告诉你谁拿了块没还、谁减了两次 ref_cnt。事发当时觉得打日志麻烦出 bug 以后才知道这开关能救命。我不建议任何人一步到位自己从头写全套 CUDA 实现。按我之前说的先在 CPU 上写一个模拟器验证块映射、动态分配和写时复制的逻辑甚至可以让模拟器统计内存占用、块利用率、复制次数再动手写真正的高性能 kernel。这一步看起来多花了几天时间实际上省下的是后面几周在 GPU 时间和显存脏数据上的挣扎。把 KV Cache 当成一个小型虚拟内存系统来看耐心一点这套机制其实不难理解难的是在工程细节上不给错误留机会。
1. 显存账本连续KV Cache为什么会拖垮高并发推理1.1 从一次推理算一笔显存账PagedAttention 这个名字第一次撞进我视野是在给一个并发推理服务做显存优化的时候。当时服务一上量GPU 显存就像漏了底的水桶请求队列越拉越长明明显卡不算差却只能同时跑几个请求很多显存被白白占住却没有产出。后来把 KV Cache 从连续分配改成块式分配同样的卡能多扛三倍以上的并发请求。这套机制看着不起眼实际上是把大语言模型推理服务里的“内存碎片”问题彻底换了个解法。先别急着看实现把账算清楚才知道它到底解决了什么。推理阶段每生成一个 token都要读取前面所有 token 对应的 Key 和 Value这个数组叫 KV Cache。它不参与训练却几乎占据推理时的大部分显存。假设一个模型有 32 层 Transformer每层 32 个注意力头每个头维度 128使用 float16 存储那么一个 token 对应的 KV Cache 大小是32 层 × 2Key 和 Value 两份× 32 个头 × 128 维 × 2 字节 512KB这不是夸张一个 2048 token 的序列光 KV Cache 就是 1GB4096 token 就是 2GB。一个并发 32 的推理服务如果每个请求都是 4K 长度光 KV Cache 就要吃掉 64GB 显存。真实场景里还没算模型权重、激活值和 CUDA 上下文KV Cache 确实是最大的显存消耗项。很多早期的推理实现走了最简单的路线每个请求按最大生成长度预分配一整块连续显存生成到哪算哪。这个方案的毛病不是一步出来的是慢性和急性混合发作。1.2 连续分配、预留浪费和外部碎片连续分配最大的问题有两个一个是预留浪费一个是碎片化。先看预留浪费。调度器不知道一个请求到底会生成多少 token只能按上限猜比如 2048 或者 4096。但真实请求往往是长短混合的一大半请求实际几百个 token 就结束了预留的长长度空间中后半段全空着。显存被占了却没有任何计算价值这就是内部碎片。打个生活化的比方你每天往一个饭盒里装饭菜饭盒按最大食量买但你实际每顿只吃一半剩下的空间每天都在浪费。再看外部碎片。并发场景下每个请求的长度不同先到的请求占了一块连续区间中间某个请求提前结束释放出来的空洞是零散的。后到的长请求需要一大块连续地址很可能就塞不进这些空洞分配器要么把请求排在后面要么触发显存整理。整理这件事在 CPU 内存里已经很麻烦了在 GPU 显存上更难受因为要暂停所有计算然后拷数据对在线延迟是致命的。连续分配还带来一个隐性约束新请求的 KV Cache 必须地址连续这导致并发的上限被“空洞”卡住而不是被实际的可用显存大小卡住。明明显存里还有很多字节空闲但因为它们不连续反而没法用这大概是推理服务最容易遇到也最让人恼火的问题。后来我把目光放到块式分配思路上的原因也很简单既然操作系统能靠分页把物理内存碎片藏起来用户进程看到的是一个连续的虚拟地址空间那 KV Cache 为什么不能做成同样的机制于是就有了接下来的块映射、动态分配和写时复制这三个核心部件。2. 块映射一张块表把逻辑块和物理地址串起来2.1 逻辑块与物理块两个不同的视角PagedAttention 的全部秘密可以概括成一句话不再为每个请求分配一段连续的显存而是把 KV Cache 切成大小固定的小块每个请求通过一张块表把“逻辑上的连续序列”映射到“物理上可能不连续的块”。这里涉及两个视角。从模型推理的角度看序列是从第 0 个 token 一路排到第 L 个 token 的连续时间步注意力计算需要逐个读取前面所有 token 的 Key 和 Value。从显存分配的角度看这些 token 的 KV 数据存在一组物理块里物理块在显存中的地址不一定相邻二次分配时也完全可以分布在显存各处。逻辑连续不等于物理连续这就是分页系统的核心思想。一个块通常包含固定数量 token 的 KV 数据比如 16 个 token 或者 32 个 token。每个物理块对应一段实际显存块与块之间不需要连续。请求持有的是一个块表Block Table它是一个数组数组里第 i 项存的是逻辑块 i 对应的物理块地址或物理块 ID。查找某个 token 对应的 KV 数据时先算出它在哪个逻辑块的第几个槽位再通过块表把逻辑块号换成物理块号最后定位到具体显存地址。举一个很简单的例子块大小 B4一个序列总共有 9 个 token逻辑块 0 包含 token 03逻辑块 1 包含 token 47逻辑块 2 包含 token 8块表为 [物理块 7, 物理块 1, 物理块 9]也就是说模型眼里的 token 08 是连续的但实际数据可能散落在物理块 7、物理块 1、物理块 9 里面。物理块 7 后面紧跟着物理块 8但序列的下一个逻辑块却映射到了物理块 1完全不挨着。这有个很直接的收益当某个请求结束后归还给分配器的是单个物理块而不是一大段连续内存其他请求可以马上复用这些零散的块。我举个例子说明这个收益。假设物理块 2 和物理块 5 都是空闲的虽然它们在地址上隔得很远但新请求的逻辑块 3 可以映射到物理块 2逻辑块 4 可以映射到物理块 5。只要块表里记清楚对应关系哪怕物理地址再乱也能组合出逻辑上连续的序列。外部碎片问题被彻底绕开了剩下的只是单个物理块内部最后一小块浪费。2.2 内核侧的块遍历和索引开销块映射不只是换了个存储方式注意力计算本身也得跟着改。以前 KV Cache 是一整块连续显存内核可以直接按线性地址从 0 扫到 L访问跟数组遍历一样简单。改成块式存储后内核拿到的不再是“起始地址 长度”而是一个块表数组加上序列长度。计算某个 token 时内核需要先算出逻辑块号block_idx token_idx // B以及块内偏移offset token_idx % B然后用 block_idx 去索引块表得到物理块地址最后按 offset 定位到块内 KV。伪代码大致长这样block_table [p0, p1, p2] # 逻辑块索引 - 物理块地址 BLOCK_SIZE 32 def load_kv(token_idx: int): block_id token_idx // BLOCK_SIZE offset token_idx % BLOCK_SIZE physical_addr block_table[block_id] return read_kv(physical_addr, offset)在 CUDA kernel 里这个逻辑通常会展开成两层循环。外层循环遍历块表里的每个物理块内层循环处理块内的 token。为了减少分支内核一般会一次性加载整个块的 KV然后通过掩码判断哪些 token 是有效位置哪些位置属于 Padding。第一次改的时候我踩了一个认知上的坑总觉得多了一次查表就会让注意力计算变慢。实际上 GPU 上多做的不是分配动作而是内存寻址动作数据还是从全局显存读计算也不变复杂。真正需要关注的开销在两层一是块表本身如果存放在 CPU 内存每次 kernel 启动前要上传到 GPU二是块表条目数量太多时L2 缓存可能放不下反复命中全局内存会拖慢一点读取速度。一张块表大概多大一般一个物理块的编号用 4 字节或 8 字节存储假设一个服务管理 100 万个物理块块表数据大约 8MB。相比几十 GB 的 KV Cache这个开销可以忽略。真正不能忽略的是块粒度带来的内层计算效率后面我会专门讲块大小怎么选。2.3 块表在主存还是在 GPU 上更合适工程实现时一定会遇到一个问题块表放哪。我的建议是分配和管理逻辑放 CPU把“当前请求”的块表作为轻量张量传入 GPU。原因有三个。第一分配器运行在 CPU 侧它需要在每次生成完一个 token 后决定是否分配新块、是否触发写时复制这些决策依赖 CPU 端的数据结构如果块表放在 GPU 上CPU 每一步都要发起跨端查询延迟没法接受。第二同一物理块可能同时被多个请求引用引用计数和映射关系的变化频率很高这种几十万次级别的指针操作明显更适合 CPU 内存完成。第三kernel 侧并不需要修改块表它只要读取所以 CPU 在生成完新块之后把块表同步到 GPU 就行了不需要复杂的双端互斥。如果块表放 GPU也不是完全不行但它会让分配器变成瓶颈。所有空闲块的管理、引用计数的增减、写时复制的判断都要在 GPU 上原子操作调试起来非常痛苦。做工程不是追求极端而是要找到简单稳定的方案块表放 CPU、按请求上传到 GPU 是我实测下来最省事的做法。3. 动态分配按需给块而不是按最大长度给满3.1 按需分配的生命周期和空闲块池块映射解决的是“地址不连续能不能用”的问题动态分配解决的是“块到底什么时候给”的问题。传统预分配是请求一开始就按最大长度占满显存动态分配则遵循一条朴素原则用多少分多少生成一个 token 的 KV才占一个 token 的空间。一个请求从进入服务到离开块分配大概会经历这样几个阶段请求到达初始化 KV Cache此时序列还没生成任何新 token但为了写第一个 token 的 KV需要先分配一个物理块。每生成一个 token把它对应的 Key 和 Value 写入当前块的槽位。当前块写满了分配器再从空闲池里拿出一个新块把它加到这个请求的块表末尾。请求结束或者被抢占它名下的物理块逐个归还到空闲池。这个流程里最关键的一点是分配动作的触发点是“写入需求”不是“长度预期”。一个请求生成到第 2 个 token 就结束那它从头到尾只需要一个物理块一个请求生成 5000 个 token块表才会一次次增长。显存的使用量和实际计算量严格同步不再有任何预留浪费。分配器内部一般维护一个空闲块池。物理块被释放后不会真正还给操作系统而是放进池子里等下一个请求拿走。这样做的好处是避免频繁申请显存毕竟一次显存分配可能涉及同步整个设备代价比 CPU 内存分配大得多。我见过的实现里空闲块池是一个简单数组或链表分配时取一个块释放时把块地址塞回来操作本身不复杂难的是并发环境下的多个请求同时分配时要保证正确性。如果你自己写这块逻辑我建议做一个非常朴素的模拟器先在 CPU 内存里用 Python 把块表、空闲池、请求生命周期跑通再考虑 CUDA kernel 的细节。我后来在公司项目里就是先用模拟器复现各种边界情况才避免了后面改 GPU 代码时反复返工。3.2 最大可分配块数和高负载下的取舍动态分配不等于无限分配。物理块的数目会受到显存容量约束我们可以用一条公式控制全局上限max_num_blocks kv_cache_allocated_size / (block_size * bytes_per_token)其中 bytes_per_token 是每个 token 的 KV 大小block_size 是块内 token 数量。比如某台 GPU 上给 KV Cache 留了 40GB 显存块大小 32每个 token 512KB那么最大物理块数就是 40 × 1024³ / (32 × 512 × 1024) 2560 个左右。这个上限必须在服务启动前算清楚否则高并发下所有块都被占满新的请求就会互相挤占内存。块被占满以后怎么办策略通常有两种排队或者抢占。排队最简单新请求等着等前面的请求释放块再进但对在线服务来说排队意味着延迟上升用户体验变差。抢占则是把某个正在运行的请求暂停把它已经生成的物理块回收优先给新请求使用等显存有裕量了再重新恢复被抢占的请求。实现抢占时最麻烦的是恢复不过因为有了块表和动态分配恢复过程只需要重建逻辑块到物理块的新映射不需要搬移原始 KV 数据。这一块是分页思想最漂亮的地方物理块可以重新映射逻辑连续性始终由块表保证。还有一个细节值得注意全局上限并不是越小越好。上限设置得太低服务会频繁触发抢占GPU 的大部分时间花在切换请求上真正算注意力的时间反而变少。上限设置太高又可能挤压模型权重和激活值的显存空间。比较稳妥的做法是先跑几个典型长度的压测请求观察实际 KV Cache 峰值再留 20% 的余量。不要拍脑袋定数字用数据说话。4. 写时复制让多个分支共享前缀 KV代价降到最低4.1 并行解码带来的前缀共享需求如果说块映射和动态分配解决的是“怎么省显存”那写时复制解决的就是“怎么让同一个前缀被更多人共用”。这个需求是从解码算法里长出来的。在自回归生成里经常需要同时探索多个候选序列。最常见的两个场景一是束搜索多个候选之间共享同一个开头只是后段分叉二是并行采样同一个 prompt 同时生成多个不同的续写结果。这两个场景有一个共同特点就是大量序列在很长一段前缀里完全相同。拿束搜索举例。假设两个候选序列都来自同一个父序列父序列已经生成了 4000 个 token接下来要分支成 4 个候选。如果不做共享每个候选都要存一份前面 4000 个 token 的 KV四个候选就是四份一模一样的显存而它们的内容完全相同纯粹是重复存储。按前面算过的一 token 512KB 来算四个候选光前缀就多占 4 × 4000 × 512KB 8GB 左右这还没算后续生成的部分。如果把这些相同前缀的 KV 存在同一组物理块里四个候选共享同一块显存那么 8GB 就可以降到 2GB因为只需要一份 4000 token 的 KV 加上四个候选各自新生成部分的 KV。这就是前缀共享的核心收益。共享的粒度自然是物理块而不是单个 token原因很简单块表可以在多个请求之间保持一致的映射申请新块时按块登记引用即可管理成本低得多。4.2 写时复制触发条件和引用计数共享了一个物理块之后最怕的事情是一个请求往块里写新数据把其他请求的前缀数据冲掉。写时复制Copy-on-Write就是用来防止这个问题的。需要区分两种物理块完整的块和不完整的块。一个完整的块意味着它已经写满 B 个 token内容确定不变多个请求共享它是绝对安全的因为没有人会往里面写新东西。真正危险的是最后一个未写满的块后续分支需要往这个块的剩余槽位写入新 token如果多个分支都往同一个物理块写那就互相覆盖了。那怎么处理未满块实现上引用计数是标准做法。每个物理块维护一个 ref_cnt表示有多少个请求正在引用它。当一个请求需要继续写入最后一个块时先检查这个块的 ref_cnt如果 ref_cnt 1说明只有当前请求在用放心写入。如果 ref_cnt 1说明还有其他分支在共享这个块不能直接写否则会覆盖别人。此时分配一个新物理块把旧块已有的 KV 数据复制过去更新当前请求的块表指向新块然后把旧块的 ref_cnt 减一新块的 ref_cnt 初始为一。这就是写时复制的完整动作复制发生在我们真正要“写”的那个时间点而不是分支创建的那一刻。分支创建时所有分支的块表都指向同一组物理块一个字节都没移动过。只有当前分支要往后接新 token才把最后那个未满块单独复制出来。之前所有完整的块继续共享。从实现上理解这段话可以看这个简化逻辑def append_token_kv(seq, token_kv): last_block seq.block_table[-1] if last_block.full: new_block allocator.allocate() seq.block_table.append(new_block) write_token(new_block, 0, token_kv) elif last_block.ref_cnt 1: new_block allocator.allocate() copy_kv(new_block, last_block) seq.block_table[-1] new_block last_block.ref_cnt - 1 new_block.ref_cnt 1 write_token(new_block, last_block.used_slots, token_kv) seq.update_gpu_block_table() else: write_token(last_block, last_block.used_slots, token_kv)这里有个工程细节容易忽略块表更新后必须同步到 GPU不是只改 CPU 端结构就能让 kernel 拿到新地址。如果你的实现里每次生成一个 token 都要同步一次块表同步本身会成为性能热点。更常见的做法是每一步生成时批量处理多个 token或者只在交换到新块的那一步同步让块表更新的次数远小于 token 数量。4.3 完整共享块的不可变性与安全性写时复制看起来很美好但安全前提只有一个完整的共享块必须是不可变的。这个前提一旦被破坏共享的所有收益都会变成 bug。我见过一个新手实现会犯的错为了省一次分配在某个完整块里把一个 token 的 KV 覆盖成新值。这样做的后果是整个服务的结果随机的因为其他分支读到的 KV 已经被污染了。要避免这种问题设计上必须把“写”动作严格限制在当前请求的最后一个未满块内。完整块一旦写满就把它视为只读数据任何请求都不能再去修改它。因为完整块是只读的多个 kernel 同时读取同一个物理块完全不需要加锁。GPU 端的读操作天然不会互相干扰而写操作又被调度器限制为“每个序列每步最多写一个 token”所以不会出现两个 warp 同时写同一个共享块的竞争。这个设计之所以能保持正确性不是靠复杂的原子操作而是靠调度层面的约束写时复制发生在 CPU 侧kernel 侧只读共享块。这里再补一个容易踩的坑引用计数的回收顺序。当请求结束时要遍历它的块表对每个物理块做 ref_cnt - 1只有 ref_cnt 归零的块才真正进入空闲池。如果请求在运行中被抢占它的块也要参与引用计数回收否则会出现“块被释放了但 ref_cnt 还挂在树上”的泄漏。调试这种问题最好的办法是在分配器和释放器里打日志每次 ref_cnt 变化都输出一行对照请求生命周期看谁多减了一次、谁漏减了一次。别觉得日志慢上线前抓三天 bug 的成本比写日志高得多。5. 工程里的关键参数块大小、碎片率和选型判断5.1 块大小不是越大越好也不是越小越好块大小是整个机制里最微妙的参数它直接决定内部碎片、块表长度和 GPU 计算效率三者的平衡。先看内部碎片。动态分配只做到“按块”分配没有做到“按 token”分配。如果一个块能装 32 个 token但当前请求只写了 5 个 token 就结束那这 27 个 slot 就浪费了。这个浪费对每个请求平均来说大概是半个块块越大浪费越多。比如平均每个请求 500 token块大小 64尾部浪费约 6%块大小 32 时尾部浪费约 3%块大小 16 时只有约 1.5%。在 40GB 的 KV Cache 里这百分之零点几都是几百 MB 的显存。再看块表长度。块越小同样容量的 KV Cache 需要的块数越多块表也越长kernel 遍历块表的开销越大。更重要的是GPU 里注意力计算如果每个块只装 8 个 token那么在同一个线程块内能做的并行计算变少算术强度和访存效率都会下降。块越大单次加载连续数据的效率越高kernel 越容易吃饱。我个人的实测经验是通用场景下 32 或 64 是比较稳的选择。短请求多、平均长度在 100500 的对话场景16 或 32 更合适因为尾部浪费能降低不少长文本场景平均长度几千甚至上万64 甚至 128 都可以接受因为尾部浪费占比很小而大块的 kernel 吞吐更高。具体怎么定最好用真实请求长度分布跑一次压测不要只试一个值。5.2 碎片率怎么观测很多人一听动态分配就觉得碎片问题消失了其实只是“外部碎片”消失了内部碎片还一直存在。观测碎片率最直接的指标叫做块利用率定义是block_utilization 实际写入的有效 token 数 / (已分配块数 × 块大小)块利用率接近 1说明每个块都装得很满显存没有任何浪费如果只有 0.4说明有六成显存其实是空的。服务上线后我建议把这个指标和并发数一起监控。正常情况下每个请求只有最后一个块可能不满所以块利用率应该接近 1 - 块大小 / (2 × 平均序列长度)。如果实测值明显低于这个理论值说明有大量请求没走完就被抢占或提前结束分配器把块分得太碎可能需要调整抢占策略。除了块利用率还有一块表命中率值得看如果块表太大GPU 的 L2 缓存放不下每次取物理地址都可能访问全局内存。症状是压力测试时计算时间没有明显增加但显存带宽被拖得很高。这时要么把块大小调大要么调整缓存让块表尽量留在 L2 里。单独处理这个问题时最容易走弯路我建议先确认瓶颈是计算还是带宽再决定动手方向。5.3 什么样的服务值得接入这套机制最后说一个反直觉的判断不是所有推理服务都应该用 PagedAttention。如果服务场景是单条长文本、并发很低、长度也很稳定比如一次只跑一个离线任务传统连续预分配就够用接入块式分配反而会引入额外的块表和分配器复杂度。PagedAttention 最大的优势体现在高并发、变长请求、大 batch 的服务上。判断标准可以很朴素看并发乘平均长度的积也就是 KV Cache 总需求量。几千兆显存规模的场景值得接只有几百兆的场景优化收益有限简单的连续缓存加预分配更容易维护。还有一个场景我觉得特别适合写时复制就是 beam search 这类天然需要共享前缀的生成算法它带来的显存节省量是整个机制里性价比最高的部分。6. 几个容易被忽略的实现细节6.1 分配器要区分“保留块”和“共享块”我在第一次实现的时候把所有块都放同一个空闲池结果出现一个隐蔽问题写时复制复制的块被当成普通块使用导致提前打满显存。后来把块分为两类状态一类是普通独占块一类是被共享引用的块。释放时只有 ref_cnt 减到 0 的块才回池子这对独占块和共享块都适用但分配时优先从独占块空闲池里拿避免共享块刚被释放又被复用导致引用计数历史错乱。这是很小的设计决策却能避免很多莫名其妙的 bug。6.2 抢占恢复时的块表重建抢占时被抢请求的块可能已经映射到物理块恢复时要重新分配物理块并把块表写回。这里容易发生一个问题如果直接复用原来的物理块新请求可能也在用数据会被覆盖。正确做法是恢复时重新走一遍分配逻辑为这个请求申请新的物理块然后从原来的 KV 数据拷贝回来。是的恢复过程会引入一点额外拷贝但相比把整个 KV Cache 做大几倍这个拷贝成本完全可接受。6.3 把调试开关留好再上线上线前一定要在分配器里留一个 debug 模式输出每个物理块的分配时间、请求 ID、ref_cnt 增减时间和释放时间。这套机制的所有 bug 都和生命周期的错乱有关靠定位代码很难一眼看出问题但日志会告诉你谁拿了块没还、谁减了两次 ref_cnt。事发当时觉得打日志麻烦出 bug 以后才知道这开关能救命。我不建议任何人一步到位自己从头写全套 CUDA 实现。按我之前说的先在 CPU 上写一个模拟器验证块映射、动态分配和写时复制的逻辑甚至可以让模拟器统计内存占用、块利用率、复制次数再动手写真正的高性能 kernel。这一步看起来多花了几天时间实际上省下的是后面几周在 GPU 时间和显存脏数据上的挣扎。把 KV Cache 当成一个小型虚拟内存系统来看耐心一点这套机制其实不难理解难的是在工程细节上不给错误留机会。
平方复杂度无法被常规优化消除｜TaoToken 统一 Key 通道实测)
