1. 显存账本：连续KV Cache为什么会拖垮高并发推理
1.1 从一次推理算一笔显存账
PagedAttention 这个名字第一次撞进我视野，是在给一个并发推理服务做显存优化的时候。当时服务一上量，GPU 显存就像漏了底的水桶，请求队列越拉越长，明明显卡不算差，却只能同时跑几个请求，很多显存被白白占住却没有产出。后来把 KV Cache 从连续分配改成块式分配，同样的卡能多扛三倍以上的并发请求。这套机制看着不起眼，实际上是把大语言模型推理服务里的“内存碎片”问题彻底换了个解法。
先别急着看实现，把账算清楚才知道它到底解决了什么。推理阶段每生成一个 token，都要读取前面所有 token 对应的 Key 和 Value，这个数组叫 KV Cache。它不参与训练，却几乎占据推理时的大部分显存。假设一个模型有 32 层 Transformer，每层 32 个注意力头，每个头维度 128，使用 float16 存储，那么一个 token 对应的 KV Cache 大小是：
32 层 × 2（Key 和 Value 两份）× 32 个头 × 128 维 × 2 字节 = 512KB
这不是夸张，一个 2048 token 的序列光 KV Cache 就是 1GB，4096 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 数据时，先算出它在哪个逻辑块的第几个槽位，再通过块表把逻辑块号换成物理块号，最后定位到具体显存地址。
举一个很简单的例子：
- 块大小 B=4，一个序列总共有 9 个 token
- 逻辑块 0 包含 token 0～3
- 逻辑块 1 包含 token 4～7
- 逻辑块 2 包含 token 8
- 块表为 [物理块 7, 物理块 1, 物理块 9]
也就是说，模型眼里的 token 0～8 是连续的，但实际数据可能散落在物理块 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 是比较稳的选择。短请求多、平均长度在 100～500 的对话场景，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 当成一个小型虚拟内存系统来看，耐心一点，这套机制其实不难理解，难的是在工程细节上不给错误留机会。
