☰
PagedAttention与前缀缓存:大模型KV Cache显存优化实战解析
2026/11/3 16:12:28 网站建设 项目流程

大模型推理真正跑起来以后,很多人的第一个感觉是:模型权重明明装得下,怎么一跑长上下文显存就报警?原因很简单——权重只是显存开销的一部分,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 存下来,后续生成时只读取、不再重复计算。于是每生成一个 token,KV Cache 就变长一段,而且是所有层同时变长,不是某一层的单独行为。

理解这一步很关键,因为它决定了显存增长的基本形态:增量式、不可控、且和序列长度强相关。连续分配模式下,你没法预料 100 个并发请求各自最终会生成多长,为了不让程序中途爆掉,只能在请求开始时按最大长度预留空间。这就是显存浪费的第一个来源。

1.2 算一笔账:7B 模型的 KV Cache 有多大

很多初学者觉得 KV Cache 不过是一堆小矩阵,实际数字会吓人一跳。假设一个 7B 规模模型,32 层,GQA 配置下 8 个 KV 头,每个头的维度是 128,权重用 FP16(每个元素 2 字节)。

每生成 1 个 token,KV Cache 增加的数据量是:

  • 每层每 token 存 Key 和 Value 两份:2 × 8 × 128 = 2048 个元素
  • 每层 4KB(FP16:2048 × 2),32 层就是 128KB
  • 上下文 4096 token 时,单条请求的 KV Cache = 128KB × 4096 = 512MB

如果 batch 里有 32 条并发请求,光 KV Cache 就占用 16GB。24GB 显卡上模型权重已经占了 14GB,剩下空间根本不够用。这还没算 MHA 配置,如果是 32 个注意力头且不做 GQA,KV 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 序列做哈希,作为 key,value 指向一个可复用的 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 为什么选迷你框架而不是直接读 vLLM

vLLM 的工程代码极其完善,但也极其庞大。变化无穷的调度逻辑、多 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 table:hash 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 源码,会觉得每条注释背后都是你踩过的同一个坑,整个学习过程就顺了。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询