Prefix Caching:复用相同前缀的 KV cache

前面几篇文章里,我们已经把 mini vLLM 的核心推理链路搭起来了。

请求进入系统后会变成 Sequence。Scheduler 决定每一轮做 prefill 还是 decode。Block Manager 把 KV cache 切成 blocks,并为每个 Sequence 维护 block table。PagedAttention 根据 block table 读取历史 KV。Sampler 从 logits 中选择下一个 token。

到这里,系统已经解决了一个很关键的问题:每个请求内部可以复用自己的历史 KV cache。

也就是说,同一个请求在 decode 阶段不需要反复计算 prompt 和历史生成 token。

但在线服务里,还有另一类浪费非常常见:

不同请求之间,可能有相同的前缀。

例如很多 ChatGPT 风格的服务都会给每个请求加一个系统提示词。

You are a helpful assistant.

或者一个代码助手会给每个请求拼接同样的仓库说明、规范、工具描述。

又或者一个 RAG 应用中,多个用户问题可能共享同一段检索文档。

如果这些前缀完全相同,那么每个请求都重新 prefill 一遍,就会浪费大量计算。

Prefix caching 要解决的就是这个问题。

它的目标是:如果一个请求的前缀已经被其他请求计算过,那么新的请求可以复用这部分前缀对应的 KV cache,而不是从头计算。

这听起来像一个简单缓存,但放到 LLM serving 系统里,它会牵涉 token 匹配、KV cache block 共享、引用计数、copy on write、调度准入和缓存淘汰。

这一篇我们先实现一个教学版 prefix caching,把核心思想讲清楚。

Prefix caching 复用的到底是什么

要理解 prefix caching,先要明确它缓存的不是文本,也不是模型输出,而是某段 token 前缀对应的 KV cache。

假设有两个请求:

请求 A:
system prompt + user question A

请求 B:
system prompt + user question B

它们的 system prompt 完全一样。

如果 system prompt 有 1000 个 token,那么这 1000 个 token 在同一个模型、同一套权重、同一套位置编码和同样推理配置下,经过 prefill 得到的 KV cache 是一样的。

所以请求 B 不需要再次计算这 1000 个 token 的 K 和 V。

它可以直接复用请求 A 已经计算好的前缀 KV cache。

然后只需要继续 prefill 自己不同的后缀部分,也就是 user question B。

从概念上看,请求 B 的 prefill 可以变成:

复用 prefix KV cache
只计算 suffix tokens
把 suffix KV 追加到 cache
然后采样第一个 generated token

这和普通 prefill 不同。

普通 prefill 是:

输入完整 prompt
计算完整 prompt 的 KV cache

prefix caching 后变成:

找到已缓存前缀
加载前缀 block table
只 prefill 未命中的后缀

这就是 prefix caching 的核心。

它减少的是 prefill 计算。

它不会减少 decode 阶段对完整上下文的访问,因为生成后续 token 时,模型仍然需要看到完整历史上下文。只是历史上下文里有一部分 KV cache 是复用来的。

为什么 prefix caching 和 block manager 天然相关

如果 KV cache 是一整段连续张量,prefix caching 会很麻烦。

因为两个请求共享一段前缀,但后续 suffix 和 generated tokens 又不同。

假设请求 A 和请求 B 共享前 1000 个 token。

后面 A 继续生成 A 的内容,B 继续生成 B 的内容。

如果 cache 是连续数组,那么共享前缀、追加不同后缀、释放不同请求,都会变得很复杂。

而 block manager 正好让这件事变得自然。

因为 KV cache 已经被切成 blocks。

如果某个前缀正好覆盖若干个完整 block,那么多个 Sequence 可以共享这些 physical blocks。

例如 block size 是 16,某个公共前缀长度是 64。

它正好占 4 个 block。

prefix blocks:
[3, 8, 1, 5]

请求 A 的 block table 可以是:

[3, 8, 1, 5, 7, 9]

请求 B 的 block table 可以是:

[3, 8, 1, 5, 2, 6]

前 4 个 block 是共享的。

后面的 block 是各自私有的。

这样,PagedAttention 读取历史 KV 时完全不需要知道这些 block 是共享的还是私有的。它只看 block table。

对于请求 A,logical block 0 到 3 指向 prefix blocks。

对于请求 B,logical block 0 到 3 也指向同样的 prefix blocks。

attention 语义没有变化。

只是底层物理 blocks 被多个 Sequence 引用了。

这说明 prefix caching 和 paged KV cache 是天然契合的。

如果没有 block table,prefix caching 很难做得优雅。

有了 block table,共享前缀就是共享 block table 的一部分。

缓存粒度:token、block 还是整段 prompt

接下来要决定一个问题:prefix cache 的粒度是什么?

最直观的做法是整段 prompt 级缓存。

如果某个完整 prompt 已经出现过,就直接复用完整 prompt 的 KV cache。

这很简单,但命中率很低。

真实场景中,两个请求完整 prompt 完全相同的概率不一定高。更常见的是它们共享前面一部分,但后面用户问题不同。

所以我们需要前缀级缓存。

那么前缀缓存到多细?

可以按 token 粒度,也可以按 block 粒度。

token 粒度最灵活,但管理复杂。

block 粒度更适合我们当前的系统,因为 KV cache 已经按 block 管理。

所以教学版可以先采用 block 级 prefix cache。

也就是说,只有完整 block 才进入 prefix cache。

如果 block size 是 16,一个前缀有 50 个 token,那么前 48 个 token 对应的 3 个完整 block 可以缓存,剩下 2 个 token 先不缓存,作为 suffix 重新计算。

这样做会损失一点复用空间,但实现简单很多。

更重要的是,它避免共享一个只写了一部分的 block。

共享未满 block 会带来很多麻烦。

因为后续不同请求可能会往这个 block 的剩余位置写入不同 tokens。如果共享这个未满 block,就会发生写冲突。

所以第一版 prefix caching 应该只共享已经写满且不会再被修改的 blocks。

这条原则非常重要:

只共享不可变的完整 prefix blocks

这样可以避免复杂的 copy on write。

后面如果要支持更精细的缓存,再考虑共享部分 block 或做 copy on write。

第一版先把语义做稳。

Prefix cache key 应该怎么设计

缓存必须有 key。

prefix caching 的 key 不能直接使用原始文本。

原因是模型真正看到的是 token ids,而不是字符串。

两个文本看起来相同,但如果 tokenizer 或 chat template 不同,得到的 token ids 可能不同。

相反,两个文本形式略有不同,也可能在某些 tokenizer 下产生相同 token 序列。

所以 prefix cache key 应该基于 token ids。

一种简单做法是对 token ids 的前缀做 hash。

例如:

hash(token_ids[0:16])
hash(token_ids[0:32])
hash(token_ids[0:48])
...

每个完整 block 边界都可以有一个 prefix hash。

这样,当新请求进来时,我们可以计算它在每个 block 边界上的 prefix hash,然后查找最长命中的前缀。

举个例子,block size 是 16。

新请求 prompt 长度是 80。

可以计算这些前缀:

前 16 tokens 的 hash
前 32 tokens 的 hash
前 48 tokens 的 hash
前 64 tokens 的 hash
前 80 tokens 的 hash

如果 cache 中命中了前 64 tokens,那么请求只需要从 token 64 之后继续 prefill。

当然,真实系统还需要把模型 id、tokenizer 版本、LoRA adapter、系统配置等信息放进 key。因为不同模型或不同 adapter 下,同样 token ids 对应的 KV cache 不能混用。

教学版可以先把 key 简化成:

prefix_hash = hash(tuple(prefix_token_ids))

但文章里要明确:生产系统不能只看 token ids,还要把会影响模型 hidden states 的因素纳入 key。

如何查找最长公共前缀

有了 prefix hash,就可以做最长前缀匹配。

当新请求进来时,它的 prompt token ids 已经确定。

系统从最长完整 block 前缀开始查找。

例如 prompt 长度是 100,block size 是 16。

完整 block 前缀长度有:

16, 32, 48, 64, 80, 96

查找时可以从 96 开始往下找。

是否命中前 96 tokens?
是否命中前 80 tokens?
是否命中前 64 tokens?
...

找到第一个命中,就是最长可复用前缀。

如果命中了前 64 tokens,那么这个请求可以复用前 64 tokens 对应的 blocks。

剩下的 token 64 到 token 99 需要继续 prefill。

为什么要找最长前缀?

因为命中越长,省掉的 prefill 计算越多。

但是这里要注意,最长前缀必须对应完整 blocks。

如果 prompt 长度是 100,最多只能匹配到 96,而不是 100。因为前面说过,第一版只共享完整 block。

prefix cache entry 可以保存:

prefix length
block ids
hash
reference count
last access time

当命中某个 prefix 时,新 Sequence 会把这些 block ids 复制到自己的 block table 前面,并增加这些 blocks 的引用计数。

这一步不是复制 KV cache 内容,只是复制 block 引用。

这正是 block table 的优势。

引用计数:共享 block 如何释放

一旦多个 Sequence 可以共享同一个 physical block,就不能再简单地认为一个 block 只属于一个 Sequence。

在第十篇的 Block Manager 里,我们曾经说过一个不变量:一个 physical block 同一时间只能属于一个 Sequence。

引入 prefix caching 后,这条不变量要修改。

更准确地说:

普通可写 block 同一时间只能属于一个 Sequence
不可变 prefix block 可以被多个 Sequence 共享

这意味着 Block Manager 需要为 block 增加引用计数。

ref_count[block_id]

当 prefix cache 保存某些 blocks 时,ref count 增加。

当新 Sequence 命中 prefix 并引用这些 blocks 时,ref count 继续增加。

当 Sequence 结束时,它释放自己的 block table。对于每个 block,ref count 减 1。

只有当 ref count 变成 0 时,block 才能真正回到 free list。

否则,这个 block 仍然被其他 Sequence 或 prefix cache entry 使用,不能释放。

这会改变 free_sequence 的逻辑。

以前释放是:

把 sequence.block_table 里的 blocks 全部放回 free list

现在释放是:

对 sequence.block_table 里的每个 block:
  ref_count -= 1
  如果 ref_count == 0:
    放回 free list

如果没有引用计数,共享 block 很容易被提前释放。其他 Sequence 还在使用这个 block,PagedAttention 却读到了被复用或覆盖的内容,结果会非常危险。

所以 prefix caching 的第一条正确性原则是:

共享 block 必须有引用计数保护

为什么共享 block 必须是只读的

引用计数解决了释放问题,但还没有解决写入问题。

如果两个 Sequence 共享同一个 block,其中一个 Sequence 往这个 block 里写新 token,会影响另一个 Sequence。

这显然不允许。

所以共享 prefix blocks 必须是只读的。

它们代表已经完成的、不可变的前缀 KV cache。

Sequence 在 decode 或继续 prefill suffix 时,不能向共享 prefix block 写入新的 token。

这也是为什么第一版只共享完整 block。

一个完整 block 对应的 token 全部已经写入,不需要再修改。

后续新 token 一定会写到新的私有 block 中。

如果要共享未满 block,就需要 copy on write。

也就是当某个 Sequence 想往共享 block 的剩余位置写入时,必须先复制这个 block,得到自己的私有副本,然后再写。

copy on write 能提高复用粒度,但会让实现复杂很多。

教学版 prefix caching 可以先避免它。

我们的规则是:

只有完整 block 可以进入 prefix cache
共享 block 不允许写入
新的 suffix 和 generated tokens 使用私有 block

这个规则简单,但足够实现一个正确可用的 prefix cache。

命中 prefix 后,prefill 如何变化

假设 block size 是 16。

新请求 prompt 有 80 个 token。

Prefix cache 命中了前 48 个 token。

那么 Sequence 初始化时,不再从空 block table 开始。

它会先引用命中的 prefix blocks。

block_table = cached_prefix_blocks
cached_token_count = 48

剩下的 token 48 到 79 是 suffix,需要继续 prefill。

这时 prefill 的输入不再是完整 prompt,而是 suffix tokens。

但是这里要非常小心:suffix tokens 的 position ids 不能从 0 开始。

它们在完整上下文中的位置是 48 到 79。

所以 prefill suffix 时,模型需要知道它们接在已有 prefix KV cache 后面。

这在实现上意味着 ModelRunner 要支持“带 prefix cache 的 prefill”。

它的输入不只是 suffix token ids,还包括:

已有 prefix 的 block table
prefix length
suffix token ids
suffix 写入 slots
position ids 从 prefix length 开始

attention 计算时,suffix tokens 需要能看到 prefix tokens。

也就是说,suffix prefill 不只是单独处理后缀,它需要在 attention 中读取 prefix KV cache。

这和普通 prefill 有差异。

普通 prefill 一次性处理完整 prompt,prompt 内部 attention 由 causal mask 处理。

prefix cache 命中后,suffix prefill 要做的是:

suffix tokens 之间做 causal attention
suffix tokens 还要 attend 到 prefix KV cache

在教学版里,为了降低复杂度,可以先选择一种简化策略:

命中 prefix 后,仍然用框架或简单路径处理 suffix,并把 prefix 作为 past KV 传入。

如果我们已经有 PagedAttention,则可以在 suffix prefill 中通过 block table 读取 prefix KV。

这部分实现会比 decode 更复杂,因为 suffix 可能不止一个 token。它介于完整 prefill 和单 token decode 之间。

所以第一版 prefix caching 可以先做更保守的实现:只在 prompt 完全命中某个 cached prefix 并且 suffix 很短时走复用,或者先把 suffix prefill 写成容易理解的版本,不追求性能。

这里的重点是理解:prefix caching 不只是跳过前缀 tokens,它还要保证后续 suffix 的 position 和 attention 语义正确。

Prefix cache 对 Scheduler 的影响

引入 prefix cache 后,Scheduler 的估算也会变化。

以前一个 waiting 请求的 prefill 成本等于完整 prompt 长度。

现在,如果它命中 prefix cache,真正需要计算的只有 suffix 长度。

例如 prompt 长度是 4096,但命中了前 3072 tokens。

那么本轮 prefill 计算成本大概只剩 1024 tokens。

这会影响 token budget。

也会影响 block 分配。

前 3072 tokens 对应的 blocks 是共享引用,不需要重新计算,也不需要重新写入。

但引用这些 blocks 会增加 ref count。

后面的 suffix tokens 需要分配新的私有 blocks。

所以 Scheduler 和 Block Manager 要配合 Prefix Cache 做准入判断:

先查 prefix cache
得到 matched_prefix_len
计算 suffix_len
根据 suffix_len 估算 token budget
根据 suffix_len 分配私有 blocks
对 matched prefix blocks 增加引用计数

这会让调度更复杂,但也更高效。

因为一个长 prompt 请求如果命中了大部分 prefix,它就不应该再被当成完整长 prompt 对待。

这对系统吞吐和首 token 延迟都有明显帮助。

尤其是系统 prompt 很长、工具描述很长、RAG 文档复用率高的场景,prefix caching 价值很大。

缓存淘汰:cache 不能无限增长

Prefix cache 不能无限保存。

否则所有历史前缀都会占用 KV cache blocks,最终 free blocks 变少,影响新请求运行。

所以 prefix cache 必须有淘汰策略。

最简单的是 LRU。

每个 cache entry 记录最后访问时间。

当 free blocks 不够,或者 prefix cache 总占用超过阈值时,就淘汰最久未使用的 entry。

淘汰一个 entry 时,不是立刻把 blocks 放回 free list,而是减少这些 blocks 的 ref count。

如果某些 blocks 仍然被正在运行的 Sequence 使用,它们不会释放。

只有 ref count 变成 0,才真正进入 free list。

这说明 prefix cache entry 和 Sequence 都可能引用同一个 block。

block 的生命周期由 ref count 决定,而不是由某一个对象决定。

淘汰逻辑大概是:

选择一个 cache entry
从 prefix cache 中删除它
对 entry.blocks 逐个 decrement ref count
ref count 为 0 的 block 回到 free list

这也带来一个设计问题:prefix cache 占用的是 KV cache 空间,它和正在 running 的请求共享同一个 block pool。

如果缓存太激进,会挤占活跃请求空间。

如果缓存太保守,命中率低。

所以 prefix cache 的容量策略也需要结合 workload。

第一版可以设置一个简单上限,比如最多允许 prefix cache 占用一定比例的 blocks。

后续 benchmark 时,再观察命中率和显存压力之间的关系。

Prefix cache 的正确性风险

Prefix caching 的 bug 往往比普通 KV cache bug 更隐蔽。

因为它涉及跨请求共享。

一个请求生成不对,可能不是这个请求自己的 block table 错了,而是它共享的 prefix block 被其他请求错误写入、提前释放或错误复用。

所以实现时要特别注意几个正确性条件。

第一,cache key 必须足够严格。

不能把不同模型、不同 tokenizer、不同 adapter、不同 position encoding 配置下的 KV cache 混用。

第二,共享 blocks 必须只读。

任何写入共享 prefix block 的行为都应该被禁止或触发 copy on write。

第三,引用计数必须正确。

block 只要还被 Sequence 或 cache entry 引用,就不能释放。

第四,prefix length 必须和 block table 对齐。

第一版只共享完整 blocks,避免半块共享。

第五,suffix prefill 的 position ids 必须正确。

命中 prefix 后,suffix token 的位置不能从 0 开始,而要从 prefix length 开始。

第六,attention 必须能看到 prefix。

suffix prefill 和后续 decode 都要能通过 block table 访问共享 prefix blocks。

这些条件任何一个错了,模型输出都可能出问题。

因此,prefix caching 的测试也要非常认真。

如何验证 prefix caching 正确

最重要的验证方法是和无 prefix cache 的结果对齐。

给定同一个 prompt,用两种方式运行:

方式 A:
  不使用 prefix cache,完整 prefill

方式 B:
  先构造 prefix cache,再复用 prefix,只 prefill suffix

在 greedy decoding 下,两者输出应该一致。

为什么用 greedy?

因为 greedy 没有随机性,更容易对齐。

测试时可以构造几个场景。

第一个场景是完整 block 命中。

例如 block size 是 4,prefix 长度是 8,刚好两个 block。

第二个场景是 prompt 长度不是 block size 的整数倍。

例如 prompt 长度是 10,只缓存前 8 个 token,剩下 2 个 token 作为 suffix。

第三个场景是多个请求共享同一个 prefix,同时生成不同后续。

这可以验证共享 block 不被错误写入。

第四个场景是请求结束释放。

一个请求结束后,共享 prefix block 不应该被释放,因为另一个请求还在使用。

第五个场景是 cache entry 淘汰。

淘汰 prefix cache entry 后,如果还有 Sequence 引用这些 blocks,blocks 不能回到 free list。

这些测试比普通 sampling 测试更重要。

因为 prefix caching 涉及资源共享,任何生命周期错误都可能导致难以复现的问题。

第一版 prefix caching 应该怎么取舍

对于 mini vLLM,第一版 prefix caching 不应该做太复杂。

建议采用这些限制。

只缓存完整 blocks。

不共享未满 block。

不实现 copy on write。

只支持相同模型、相同 tokenizer、相同配置下的 cache 复用。

cache key 基于 token ids 的 block 边界 prefix hash。

使用引用计数保护共享 blocks。

使用简单 LRU 淘汰 cache entry。

先支持 greedy 对齐测试,再考虑 sampling 场景。

这些限制看起来保守,但能保证主线清楚。

等这版跑通后,再逐步扩展。

例如支持更细粒度的 prefix match。

支持 copy on write。

支持多 LoRA adapter 的 key 隔离。

支持更复杂的缓存淘汰策略。

支持跨请求的 prompt tree 或 trie 结构。

但这些都不是第一版的重点。

第一版的重点是让读者理解:

prefix caching = 复用已经计算好的前缀 KV cache
paged KV cache = prefix caching 能优雅实现的基础
引用计数 = 共享 block 正确释放的关键
只读完整 block = 避免写冲突的简单策略

小结

这一篇我们实现了 prefix caching 的核心设计。

它解决的是跨请求重复 prefill 的问题。很多在线服务请求会共享相同的系统提示词、模板、工具描述或 RAG 文档前缀。如果每个请求都重新计算这些前缀,就会浪费大量计算。

Prefix caching 缓存的是前缀 token 对应的 KV cache,而不是最终文本输出。

有了 block manager 和 block table 后,prefix caching 变得非常自然。多个 Sequence 可以在自己的 block table 中引用同一组不可变 prefix blocks,从而复用已经计算好的 KV cache。

第一版设计中,我们只共享完整 blocks,不共享未满 block,也不做 copy on write。共享 blocks 通过引用计数保护。只有当 prefix cache entry 和所有 Sequence 都不再引用某个 block 时,它才能回到 free list。

Prefix cache 会影响 Scheduler 的成本估算。命中前缀后,一个长 prompt 的实际 prefill 成本变成 suffix 长度,而不是完整 prompt 长度。这能降低首 token 延迟,也能提高系统吞吐。

但 prefix caching 也带来正确性风险。cache key 必须严格,引用计数必须正确,共享 block 必须只读,suffix 的 position ids 必须从 prefix length 开始。验证时应该和无缓存版本做 greedy 对齐测试。

到这里,mini vLLM 已经开始具备跨请求复用能力。

下一篇文章,我们会继续处理另一个和 prefill 相关的问题:长 prompt 会阻塞 decode。

Prefix caching 可以减少重复前缀计算,但如果一个请求真的有很长的新 prompt,prefill 仍然可能占用大量 GPU 时间,让正在流式输出的请求卡顿。

为了解决这个问题,我们需要引入 chunked prefill。