大模型推理真正跑起来以后很多人的第一个感觉是模型权重明明装得下怎么一跑长上下文显存就报警原因很简单——权重只是显存开销的一部分KV Cache 才是那个随着生成不断膨胀的变量。同一个 batch 里每条样本长度还不一样连续显存的分配方式在混合长短序列的场景下会产生大量空洞。业界解决这个问题的两条主线一条是 PagedAttention把连续的 KV 缓存切成物理块按需分配另一条是前缀缓存把重复计算过的前缀直接跳过不复算。这篇文章我想把这两套机制拆开讲透再结合 nano-vllm 这类迷你推理框架聊一聊怎么用最小代码量把关键功能亲手实现一遍。最近我照着《大模型训练与推理加速实战基于 CUDA 计算平台的 Python 版》的思路在本地复现这些逻辑发现比直接硬啃 vLLM 上万行源码舒服得多。1. 先搞明白你的显存到底被谁吃掉了1.1 KV Cache 是怎么出现的Transformer 里的注意力计算需要把当前 token 的 Query 和历史所有 token 的 Key、Value 做运算。如果不做缓存每生成一个新 token就要把整段历史重新过一遍网络计算量随序列长度呈平方级增长这在实际推理中完全不可接受。KV Cache 做的事情很朴素用显存换算力。每算过一个 token就把这个 token 在各层的 Key 和 Value 存下来后续生成时只读取、不再重复计算。于是每生成一个 tokenKV Cache 就变长一段而且是所有层同时变长不是某一层的单独行为。理解这一步很关键因为它决定了显存增长的基本形态增量式、不可控、且和序列长度强相关。连续分配模式下你没法预料 100 个并发请求各自最终会生成多长为了不让程序中途爆掉只能在请求开始时按最大长度预留空间。这就是显存浪费的第一个来源。1.2 算一笔账7B 模型的 KV Cache 有多大很多初学者觉得 KV Cache 不过是一堆小矩阵实际数字会吓人一跳。假设一个 7B 规模模型32 层GQA 配置下 8 个 KV 头每个头的维度是 128权重用 FP16每个元素 2 字节。每生成 1 个 tokenKV Cache 增加的数据量是每层每 token 存 Key 和 Value 两份2 × 8 × 128 2048 个元素每层 4KBFP162048 × 232 层就是 128KB上下文 4096 token 时单条请求的 KV Cache 128KB × 4096 512MB如果 batch 里有 32 条并发请求光 KV Cache 就占用 16GB。24GB 显卡上模型权重已经占了 14GB剩下空间根本不够用。这还没算 MHA 配置如果是 32 个注意力头且不做 GQAKV Cache 直接翻四倍2048 上下文的压力就等同 4096。所以真实推理系统里 KV Cache 的优化空间非常大这也是 PagedAttention 能成为关键技术的根本原因。1.3 连续分配模式下的显存碎片问题早期基于 Transformers 的推理实现很直接每个序列整体持有一个连续的 KV 张量形状是 (2, num_layers, num_heads, max_len, head_dim)。问题集中在两点。内部碎片最明显。请求开始前按 max_len 预留但实际生成往往只用了预留的几分之一预留部分就一直空着。外部碎片同样严重——显存里空闲块不是一块整的这个请求要 7 个块、那个要 3 个来回分配释放之后大的空闲块被拆成小碎片新请求明明总量够用却申请不到足够大的连续区域。你可以把这种模式想象成搬家时非要找一块刚好放得下所有家具的空房间结果发现仓库里散落的小空间加起来很多但拼不成一整块。PagedAttention 的思路就是放弃“整块连续”转而用分页表管理“小而散”的物理块利用效率自然就上来了。2. PagedAttention把连续的 KV 块拆成物理页面2.1 核心思路来自操作系统的虚拟内存PagedAttention 这个名字直白地告诉你它参考了什么——操作系统的分页机制。进程看到的地址空间是连续的虚拟地址实际物理内存却是按页分配的通过页表做映射。KV Cache 同理逻辑上每个序列的缓存是连续增长的物理上则拆成固定大小的 block由 block table 记录映射关系。这套解耦带来三个直接好处block 可以按需分配序列增长多少就分配多少空闲 block 不再要求物理连续有零散就能用block 可以跨序列共享这是后面要说的 beam search、前缀缓存的基础。vLLM 也不是第一个意识到 KV Cache 要分块的但它是第一个把“操作系统式分页”完整搬到 GPU 显存管理里的方案且配套定制了注意力 kernel所以效果非常突出。2.2 Block 尺寸和内部数据排布实际工程里 block 默认包含 16 个 token 的 KV 数据。为什么不把 block 设成 1 个 token粒度越小越灵活但元数据和调度开销会爆炸GPU kernel 启动也没法充分利用计算单元。设太大又回到预分配浪费的旧问题——每个序列最后一块不可能恰好填满。16 是性能和碎片率之间的一个平衡点vLLM 实测默认值在多数场景表现都不错。Block 内部的排布在 kernel 优化里很讲究。KV 数据通常组织成 (num_blocks, num_layers, 2, num_kv_heads, block_size, head_dim) 这种结构确保 GPU 读取时尽量走连续内存。看起来只是形状问题实际影响带宽利用。推理是访存密集型任务Kernel 能不能快速读入数据往往比浮点计算本身更能拉开性能差距。2.3 引用计数与共享Beam Search 的显存救星PagedAttention 另一个被低估的设计是 block 的引用计数机制。一个 block 可以被多个序列共享每多一个序列引用ref count 加一只有引用归零时才真正释放。这个机制在 beam search 里效果特别明显。模型要维护 N 条候选序列它们前面大部分 token 是一样的只是最后分支不同。如果每条序列都独立持有完整的 KV Cache前缀部分就被重复存了 N 份。有了 block 共享公共前缀只用一份物理块各序列的 block table 都指向它。只有分支写入新 token 时才需要复制受影响的 block也就是 copy-on-write。具体到一个场景序列 A 已经生成了 50 个 token序列 B 从同一个 prompt 出发在第 30 个 token 处分叉。第 30 到 50 之间的 block 会被 A 和 B 共享只有当 B 继续写出新 token、需要修改某个共享块里的位置时系统才复制该块并让 B 指向副本。前缀部分的显存开销直接减半。2.4 KV Kernel 的适配方式显存管理变了注意力计算也得跟着改。标准 attention 拿着连续的 (seq_len, num_heads, head_dim) 张量直接做矩阵乘法但分页后 KV 散落在不连续的块里Kernel 必须通过 block table 找到每个 token 对应的物理位置。实现上大致是kernel 先读取当前 sequence 的 block table逐个 block 把 K、V 从显存搬运到共享内存block 内部做分块 attention再用衰减系数把跨 block 的结果合并起来。这个思路和 FlashAttention 的分块软化是同一套数学原理区别在于 KV 的来源变成了按块寻址的物理显存。我自己实现的时候建议先别急着优化 kernel用最简单的 for 循环把正确性跑通再慢慢对照高性能实现学 kernel 里的 tiling 和访存优化。一步到位很容易被指针和 offset 计算搞晕。2.5 实际收益的判断标准PagedAttention 不是让显存变大而是把原本浪费掉的空洞用起来。对比连续分配比较合理的预期是 KV 内存减少 60%-80%系统吞吐提升数倍。不过这个数字跟模型结构、batch 大小、序列长度分布相关不要盲目拿别人 benchmark 当自己的 KPI。要验证效果最直接的方式是跑同样一组请求分别统计torch.cuda.max_memory_allocated()和实际生成速度。观察到的提升来自两个部分空闲 block 利用率提升带来的显存容量释放以及避免请求间为了凑连续空间而产生的数据拷贝。前者是存储收益后者是时间收益。3. 前缀缓存让重复的前缀不再从头计算3.1 你每次都在重复算同样的前缀PagedAttention 解决的是 KV Cache“怎么放”的问题前缀缓存解决的是“哪些能不放”的问题。日常服务里大量请求的前缀完全一样Agent 场景里每个请求都带着同一个系统提示词和工具定义可能占两三千 token多轮对话里后半段轮次一直在重复前半段的上下文编程助手里同一段代码上下文往往要带出多轮续写。这些前缀每次都重新走一遍 prefill算完再丢掉纯浪费。前缀缓存的思路简单粗暴既然 token 序列完全一致那么它们的 KV Cache 理论上也完全一致。把算好的前缀 KV 存下来下次遇到相同前缀直接复用只对新产生的 token 做计算。3.2 哈希查表是最简单的实现方式实现前缀缓存最直接的做法是哈希表。对请求的 token ID 序列做哈希作为 keyvalue 指向一个可复用的 block 片段。新请求进来时从长到短尝试匹配前缀命中后用缓存 block 填充 block table。有个容易犯的错误是只存完整请求的哈希。一次性哈希匹配完全前缀虽然实现最简单但多轮对话中每轮上下文长度都不同前缀逐渐变长每次都匹配不上缓存命中率会非常低。正确做法是把序列按固定 chunk 长度分开哈希支持局部匹配。具体设计可以是每 64 或 128 个 token 划分一个 hash 区间缓存条目记录区间内最后一个 token 的全局位置这样任何长度大于等于该区间的请求都能复用前半部分。代价是哈希和索引的代码更复杂了收益是前缀复用从“非赢即输”变成“赢多少算多少”。3.3 从哈希表到 Radix Tree 的进化哈希表处理固定 chunk 还好遇到树状的前缀关系就吃力了。比如多轮对话历史是 A→A-B→A-C哈希表里要分别存 A、A-B、A-C 三个 key公共部分 A 被重复存了两次更新某一支新对话时又要重算整个链路的哈希。SGLang 的 RadixAttention 用前缀树radix tree解决了这个结构问题。每条路径代表一条历史对话序列相同前缀自然合并成树上的公共节点新分支只需要在叶子处生长。树上的节点同样有引用计数和 LRU 淘汰淘汰叶子节点不影响父节点复用。相比哈希表它更适合处理有多轮分支的复用场景。对学习来说建议先用哈希表把原理吃透再去看 RadixAttention 的实现。树的代码量比哈希表大不少一上来直接啃容易迷失在指针操作里误把“数据结构”当成“核心思想”去背。3.4 哪些场景最值得上前缀缓存前缀缓存的收益取决于前缀占比。如果总 token 里 80% 是公共前缀收益就非常可观如果每条请求的前缀千差万别收益就很有限。最典型的是 Agent 工具调用场景。系统提示 工具 schema 可能有几千 token用户每次调用同一工具时这部分的 prefill 是固定的。批量做 100 次相同请求缓存后 prefill 的计算量从 100 份降成 1 份。另一个是 RAG 检索问答。很多场景下召回文档被拼进上下文问题不同但文档相同前缀缓存能对公共文档部分做复用。从工程落地顺序看先接入哈希前缀缓存跑通日志里看到命中率数字再决定要不要上 Radix Tree。一上来就搞最复杂方案既容易返工也看不出效果在哪儿。4. 动手路线用 nano-vllm 把关键功能复现一遍4.1 为什么选迷你框架而不是直接读 vLLMvLLM 的工程代码极其完善但也极其庞大。变化无穷的调度逻辑、多 block 管理、服务化封装、多 GPU 通信一个新读者进去很难分清哪些是核心机制、哪些是生产环境补丁。nano-vllm 这类项目的价值在于把架构删到只剩骨架。没有服务层、没有多卡管理通常只有几条关键路径一个简化的 scheduler、一个 block manager、一个 attention 实现外加显存统计日志。你可以在上面亲手验证 PagedAttention 和前缀缓存的效果改一行看一行这对建立直觉非常重要。《大模型训练与推理加速实战基于 CUDA 计算平台的 Python 版》里的实验也是这个路线先用 Python 实现基础逻辑再慢慢替换成 CUDA 内核。名字里虽然带 CUDA但真正核心的显存管理思想完全可以在纯 Python 层面先复现。4.2 写一个最小可用的 BlockManager我建议的第一个动手项目是 BlockManager。它的职责是维护哪些 block 空闲、哪些 block 被哪个序列占用以及每个序列的 block table。一个最简版本大概长这样class BlockManager: def __init__(self, num_blocks, block_size): self.block_size block_size self.free_blocks list(range(num_blocks)) self.ref_counts [0] * num_blocks self.block_tables {} # seq_id - [block_id, ...] def allocate(self, seq_id, num_blocks): assert len(self.free_blocks) num_blocks blocks self.free_blocks[:num_blocks] del self.free_blocks[:num_blocks] for blk in blocks: self.ref_counts[blk] 1 self.block_tables[seq_id] blocks def get_block_ids(self, seq_id): return self.block_tables[seq_id] def free(self, seq_id): blocks self.block_tables.pop(seq_id) for blk in blocks: self.ref_counts[blk] - 1 if self.ref_counts[blk] 0: self.free_blocks.append(blk) def append_token(self, seq_id): table self.block_tables[seq_id] last table[-1] if self._used_slots(last) self.block_size: self._used_slots[last] 1 else: new_block self.free_blocks.pop() self.ref_counts[new_block] 1 self._used_slots[new_block] 1 table.append(new_block)先不用管显存封装把它跑成一个纯 Python 的模块打印每个 seq 的 block table观察 allocate、free、append竟然不会出错。这一步做完你对“分页”的理解会比读十篇论文都牢。4.3 在 BlockManager 之上加前缀缓存第二步给 BlockManager 接上前缀缓存。核心工作是 cache tablehash chunk 对应的 KV block 引用。最简单的方式是维护一个(prefix_hash) - (chunk_start_block_id, chunk_end_block_id)的字典新请求到达时先查缓存命中就把它对应的 block 加入新 seq 的 block table同时把 ref_counts 加一。注意这里有个关键点缓存的 block 是共享的新序列一旦要修改块内未缓存位置必须走 copy-on-write。如果你的实现里忘了共享 ref前缀复用必然导致显存泄漏或数据串写。可以先做一个“只读”版本只允许复用已满的 block不做块内写入这样代码逻辑最简单。跑通后再加写时复制改造成本低很多。4.4 用对比实验验证效果逻辑实现完下一步是量化收益。设计很简单的对照组同一批请求跑连续缓存版本和分页版本分别记录torch.cuda.max_memory_allocated()和峰值显存。再测一轮带前缀缓存和不带前缀缓存的 prefill 耗时。我实测下来最直观的信号其实是nvidia-smi里 memory-used 曲线的平滑度。连续分配时显存曲线会出现锯齿状跳变因为要不断分配大块分页版本曲线则平滑上移。把截图放进笔记里比任何文字描述都更有说服力。4.5 日志设计你能看见每一步分配调试这类系统最重要的一件事是加日志。分配发生时打印 seq_id、block_id、ref_count、当前空闲 block 数释放时同样打印。只要日志打印到位显存异常的原因一眼就能定位某条 seq 没释放会显示 ref_count 永远大于零前缀缓存共享未递增 ref_count 会显示同一个 block 被释放了两次。建议日志级固定保留 info线上才用 error。学习阶段日志多到刷屏完全没关系那正是你理解系统行为的最佳窗口。5. 常见坑与排查笔记5.1 显存碎片不减反增现象用了分页后峰值内存没降下来甚至比连续分配还高。排查方向free block 列表是否产生空洞ref_count 是否有泄漏。常见原因是对 block 共享的引用计数管理不正确某条序列释放时没有把共享 block 的 ref_count 减到零导致这部分显存永远无法回收。5.2 前缀命中但结果错误现象命中前缀缓存后模型输出质量明显下降或出现奇怪重复。这是哈希和 KV 数据没对齐的典型症状。排查要确认三件事哈希确实是按 token ID 计算而不是按字符串hash 对应块的 token 区间长度正确模型配置层数、头数、精度和缓存数据当时的配置一致。配置变更后必须清理缓存别偷懒。5.3 吞吐反而下降现象显存占用下来了但生成速度变慢。大概率是注意力 kernel 没有适配分页数据代码里对 KV 做了 gather 拷贝再计算。拷贝过程额外耗费了显存带宽。PagedAttention 的收益来自按块计算如果实现是先把不连续的块 gather 成完整 KV 再算基本等于放弃了分页优势。学习阶段可以先接受慢但要知道慢的原因。5.4 缓存命中率很低现象日志显示大量 miss。多数情况是 chunk 切分太大或哈希字段没做好。比如系统提示词变化了几个字节哈希值就全变了。工程上常见的处理是只对首个 chunk 使用宽松匹配或者干脆把系统提示做单独缓存段和用户输入的尾巴切开。5.5 Beam Search 中数据串写现象多路 Beam 的输出出现互相污染。这是 copy-on-write 没做或没做对。当多个序列共享同一个 block 时任何一个序列要新增 token 写入这个块都必须先复制该块并让当前序列指向副本。如果发现可能有串写先把所有共享 block 标记成只读强制任何写入都走复制路径再逐步优化复制条件。还有一个容易忽略的细节是 block_size 调完之后的联动。block_size 改了注意力核心里面的偏移计算、block table 的槽位数量、前缀缓存的 chunk 映射都要同步改。它不是一个独立参数牵一发动全身。5.6 排查经验小结遇到显存类问题我先看三件事日志里 ref_count 有没有稳步收敛到初始值block table 里是否有同一个 block 被多个 seq 使用却没显示共享记录torch.cuda.memory_summary()的输出里 cached 和 allocated 的差值是否持续扩大。三者都能快速定位到大多数显存管理 bug。另一个值得养成的习惯是把真实跑日志固化下来。每次修改后跑同样一组 prompt把 block 分配序列和显存曲线存下来对比通过 diff 看清楚每次改动到底改变了什么。6. 从这套机制里我学到的排错直觉真把 nano-vllm 这套逻辑跑通之后我最大的感受是显存优化不是黑魔法它就是把资源形态和组织方式一层层拆开重新设计。PagedAttention 解决空间碎片和共享问题前缀缓存解决重复计算问题两者配合时显存空间的利用率和算力的有效使用率是同时提升的不是此消彼长。如果你想快速建立体感我真的建议先别碰大工程。找一个迷你实现把 BlockManager 的 print 日志打开用 nvidia-smi 盯着显存变化亲手把连续分配改成按块分配再亲手加上一个最简单的哈希前缀缓存。每一层改动带来的数值变化都会变成你排障时最可靠的直觉来源。学习顺序上先跑通最简单的 decode 流程再引入 block 管理然后加前缀缓存最后再看高性能 kernel。这三个阶段都能独立验证收益任何一步出问题都能定位。等你把这条链路走完再回头读 vLLM 源码会觉得每条注释背后都是你踩过的同一个坑整个学习过程就顺了。
