第 8 章 — 新一代 KV Cache 管理

《LLM Primer VI: Scaling AI Systems》章节走读的第八篇。这一章要说的是:把操作系统那份分页的洞见搬到推理 engine 里来 —— 把 KV cache 从一大板预留字节,变成一份可共享、可淘汰、可前缀缓存的资源。

这一章为什么存在

第 7 章留下 continuous batching 手里一笔债。序列每一次迭代都在进出;朴素布局给每个槽一大板按最大长度分的内存,大部分是浪费;batching 的吞吐胜利被还回去一半。这笔债就是内部碎片化,操作系统在 1960 年代靠分页解决过的同一个失败模式。第 8 章讲的就是把那份解决方案应用到 LLM serving:把 KV cache 切成小的物理 block,用一张 page table 把它们和逻辑 token 位置解耦,再让淘汰和缓存策略跨序列回收或共享 block。PagedAttention 是奠基那一步;H2O 和 InfiniGen 是淘汰策略;前缀缓存是那个让一小把 GPU 能撑起几百万 agent 请求的技术。

一句话:把 KV cache 当成分页虚拟内存来对待 —— 小 block、page table、引用计数共享 —— continuous batching 的内存问题就不再是问题。

8.1 PagedAttention 是给 KV cache 的虚拟内存

vLLM 那篇 PagedAttention 论文(2023)把 OS 的洞见原样搬进 engine。KV cache 被切成固定大小的 block —— 常见是 16 token 一块 —— 放在一份扁平的物理池子里。一条序列由一份 block table 表示:一份指针数组,把逻辑位置映射到物理 block ID。attention kernel 把 block table 当额外输入,按索引间接取 K、V,而不是按连续切片。Hopper 上 L2 缓存能把这份随机访问模式吸收得足够好,kernel 跑在板式版本的几个百分点以内。收益很大。内部碎片化从 60–80% 掉到大约 4%(每条序列一个不满的尾 block),可用并发涨 2–4×。基于引用计数的 block 共享让复杂采样几乎白送 —— 2,000 token prompt 上的 best-of-8 从 16,800 token-block 掉到 2,800 —— 它同时是前缀缓存建在上面的那块地基。

8.2 H2O 和 InfiniGen 淘汰那些无关紧要的 token

PagedAttention 解决了碎片化,但没解决"上下文长到任何布局都塞不下 VRAM"。200,000 token 的 Llama-3-70B 上下文,除权重之外还要 60 GB KV。H2O("Heavy Hitter Oracle")观察到:decoding 期间的 attention 权重集中在很小一批源位置 —— 近期 token、开头那些 attention-sink token,以及少量内容相关的命中 —— 而大多数历史位置基本上是零权重。engine 按位置累计 attention 分数;序列的 KV 预算接近上限时,淘汰得分最低的那些位置,只保留一个保证保留的近期和 sink 窗口。收益很大;代价是永久性 —— 如果后面某个查询要用被淘汰的位置,engine 找不回来。InfiniGen 用动态、可恢复的选择细化这份取舍:不是直接丢 token,而是把它们的 KV 卸载到 CPU 内存,如果 attention 开始重新集中到它们上面,再分页回来。用哪种淘汰策略,取决于工作负载在多大程度上重复查询自己的长历史;agent 型工作负载惩罚永久淘汰,奖励 InfiniGen 那种可恢复。

8.3 前缀缓存是 PagedAttention 解锁的最高回报的杠杆

真实流量里,多数 prompt 的前几千 token 在跨请求间是相同的。一个 chat 服务,每一段对话都复用同一份系统 prompt。一个 RAG 服务,每一个搜到同样内容的用户都把同样一段检索文本拼进 prompt。一个 agent 每一步都注入同样的工具描述和推理脚手架。PagedAttention 让共享变得机械化:按 block 大小对 prompt 哈希;如果哈希在全局 block cache 里,就把新请求的 block table 指向缓存 block,跳过这段前缀的 prefill;如果不在,就跑 prefill,然后把产出的 block 注册。生产命中率数字很惊人 —— chat 服务的系统 prompt 上超过 99%、RAG 检索前缀 30–60%、agent 脚手架接近 1.0。SGLang 的 RadixAttention 把这个思路又推一步,用一棵基数树按任意长度而不是 block 对齐的长度做前缀键。前缀缓存是那个最常见的把超预算 serving 集群拉回预算内的开关。

值得记住:如果你的工作负载里有共享的系统 prompt、共享的工具脚手架、或稳定的检索语料,前缀缓存很可能是你能拨的最高 ROI 那个开关。在假设 engine 已经开着它之前,先量一下命中率;那个数字告诉你有多少 prefill 你在付两次钱。

第 8 章为下面铺了什么

分页、淘汰、前缀缓存缩小了每 token 的 KV 占用,让 engine 在高并发下的内存行为变得可控。但这些都没解决 decoding 本身那道顺序依赖:每一步、每一条序列产出一个 token,槽再多也没用。第 9 章拿投机解码去打这个约束 —— 那族技术用一个便宜的 draft 提前预测几个 token,在一次昂贵的前向传播里一起验证,给那些对用户感受延迟最重要的序列打破单 token 每步的地板。

下一章 — 第 9 章:投机解码。Draft、verify,以及投机什么时候划算的那道算术。

想看完整的全貌?书里这一章包括能跑的 paged KV 数据结构代码、vLLM 和 SGLang 打开前缀缓存的配置开关、H2O 淘汰的调参旋钮、InfiniGen 卸载再恢复的机制,以及把上面四样合成一份生产 KV 栈的参考架构。在 Amazon 上查看 LLM Primer VI →