prefix match、lock ref 与 cache_protected_len#

这章解决什么问题#

前面的调度与内存章节已经讲了：

  • cache 为什么能复用
  • allocator 怎样分配和回收 KV
  • ReqToTokenPool 怎样保存逻辑 token 到物理 index 的映射

但还有一层特别关键、又特别容易被一带而过的实现没有单独讲清：当一条请求命中前缀缓存时，系统究竟怎样把这段前缀保护起来、继续追加、在 chunked prefill 中更新、最后再把保护关系释放掉？

这层问题主要落在几个关键机制上：

  • match_prefix(...)
  • inc_lock_ref(...) / dec_lock_ref(...)
  • cache_unfinished_req(...)
  • cache_finished_req(...)
  • cache_protected_len

如果不把这一层讲透，读者会知道“有 prefix reuse”，却很难准确回答：

  • 命中前缀以后，哪些 KV 真正被树缓存接管
  • 哪些 KV 只是暂时挂在请求上、还没进入树
  • 为什么 page-aligned 场景下 prefix_indices 和真正受保护的长度会不完全一样

为什么这层值得单独成章#

因为 prefix reuse 的真正复杂度不在“树能不能匹配”，而在“匹配以后如何安全地持有、扩展和释放这段缓存状态”。

这正是运行时里最典型的一类复杂度：

  • 匹配本身很好理解
  • 但一旦加入 page 对齐、chunked prefill、partial page、request 生命周期和 eviction，状态所有权就会变复杂

而 cache_protected_len、lock ref 和 prefix update，正是上游用来把这类复杂度压住的几个关键抽象。

一张图：prefix reuse 的真正难点不是命中，而是保护与转交#

这张图解决的理解障碍是：很多读者会把 prefix cache 想成“命中就直接复用”，但源码里还多了一层保护与释放关系。

flowchart LR
    Match["match_prefix(...)"] --> Prefix["prefix_indices + last_node"]
    Prefix --> Lock["inc_lock_ref(last_node)"]
    Lock --> Run["request continues prefill/decode"]
    Run --> Unfinished["cache_unfinished_req(...)"]
    Unfinished --> Protect["cache_protected_len update"]
    Run --> Finished["cache_finished_req(...)"]
    Finished --> Unlock["dec_lock_ref(last_node)"]

图比纯文字多解释的一点是：prefix cache 不是“树里有一段索引就行”，而是请求在运行过程中持续和树之间做状态转交。

match_prefix(...) 真正返回的不是“命中没命中”#

base_prefix_cache.py 里 MatchResult 最值得注意的字段不是只有 device_indices，还有：

  • last_device_node
  • last_host_node
  • cache_protected_len

而在 schedule_batch.py 里，请求会把这几个字段接回自己的状态：

  • prefix_indices
  • last_node
  • cache_protected_len

这说明 prefix match 的结果不是一个布尔值，而是一整组“后续运行时还要继续依赖的保护状态”。

换句话说，match_prefix(...) 的意义不仅是给你一段可复用 KV index，还要告诉你：

  • 当前匹配停在树的哪个节点
  • 这段匹配中哪些 KV 已经受树缓存正式保护

RadixCache.match_prefix(...) 为什么会主动 split node#

radix_cache.py 里的注释写得非常清楚：如果查找在某个已存 segment 中间结束，radix tree 会 split node，把边界精确暴露出来。

这说明 prefix match 不是只读查询，它还会在必要时细化树结构，让后续匹配与保护边界更精确。

这是一个很好的系统设计点，因为它表明：

  • cache tree 不是静态索引
  • 它会在运行中根据真实匹配边界自我修整

从教学角度看，这比“它是一棵 radix tree”更有价值，因为它解释了为什么后续 lock ref 和 prefix update 能变得精确。

inc_lock_ref(...) / dec_lock_ref(...) 到底在保护什么#

这两个函数的关键不只是 lock_ref += 1 / -= 1，而是它们同时在维护：

  • evictable_size_
  • protected_size_

也就是说，lock ref 不是普通引用计数，而是“这个节点以及其祖先路径是否仍然允许被驱逐”的保护语义。

更稳的理解方式是：

  • inc_lock_ref(...)：把一段正在被请求使用的前缀从“可驱逐”切到“被保护”
  • dec_lock_ref(...)：请求结束或边界转移后，再把它从“被保护”放回“可驱逐”

这说明 lock ref 机制直接决定 eviction policy 的语义边界，而不只是树节点的小计数器。

cache_unfinished_req(...) 为什么是 prefix cache 最容易被忽略的关键函数#

很多人只会记住 cache_finished_req(...)，但真正更复杂的往往是 cache_unfinished_req(...)。因为 chunked prefill 或中间阶段请求还没结束时，系统已经希望把一部分 prefix 提前插入树并开始复用。

这段逻辑大致做了这些事：

  1. 把当前 fill_ids 对应的 KV indices 拿出来
  2. 尝试 insert 到 radix tree
  3. 对重复部分 free 掉已被树接管的 KV
  4. 再 match_prefix(...) 一次，拿到新的 prefix_indices 和 new_last_node
  5. 更新 req.cache_protected_len
  6. dec_lock_ref(req.last_node)，再 inc_lock_ref(new_last_node)

这说明 unfinished cache insertion 不是一次“写树”动作，而是一次完整的保护边界重计算。

为什么需要 cache_protected_len#

这是这章最该单独讲透的字段。

radix_cache.py 里注释已经直接解释了它存在的原因：

  • 当 page_size > 1 时，prefix_indices 里可能会包含尾部的 partial page
  • 但这部分 partial page 并没有真正进入 radix tree

如果此时你把 len(prefix_indices) 直接当成“都已经被树保护住了”，后续 free 时就会出错。于是系统引入：

  • cache_protected_len

来表达：

  • 哪一段 prefix indices 已经被树正式接管
  • 哪一段虽然已经挂在请求上，但仍需在下一次 unfinished/finished cache 时继续正确释放

换句话说，prefix_indices 是“当前请求认为可复用的前缀视图”，而 cache_protected_len 才是“已经被树保护住的真实边界”。

这是一个非常值得技术书强调的区别。

为什么 cache_protected_len 和 prefix_indices 长度可能不一样#

在 page-aligned 或 bigram key 场景下，源码会出现这样的情况：

  • prefix_indices 可能比真正插入树的那部分更长
  • 因为尾部 partial page 或某些额外索引还留在请求视图里

这也是为什么：

  • cache_unfinished_req(...) 会重新写 req.req_to_token_pool
  • 但仍然要单独维护 cache_protected_len

对读者来说，这里最重要的 takeaway 是：

不要把“请求当前看见的前缀长度”和“树缓存已经真正拥有的前缀长度”混为一谈。

cache_finished_req(...) 怎样完成最后的所有权转交#

finished 路径里，系统会：

  1. 根据 kv_committed_len 取出最终 token ids 与 kv indices
  2. 再次尝试 insert
  3. 对已经被树接管的重复部分做 free
  4. free 未对齐尾部
  5. 最后 dec_lock_ref(req.last_node)

这条顺序非常值得记住，因为它说明 finished cache insertion 的重点不是“终于插树成功”，而是“请求和树之间的所有权关系终于结清了”。

也就是说：

  • unfinished 路径更像中途保护边界重算
  • finished 路径更像最终交账

这层对 ScheduleBatch 有什么直接影响#

schedule_batch.py 里请求会在 prefix match 后持有：

  • prefix_indices
  • last_node
  • cache_protected_len

后续：

  • set_extend_input_len(...)
  • chunked prefill
  • PrefillAdder.add_chunked_req(...)

都会依赖这些字段。

这说明 prefix cache 不是一个独立旁路优化，而是直接参与 batch 成形和 extend 长度判断的主链路机制。

这层最容易出现的误判#

1. 以为 prefix match 返回的是静态命中结果#

实际上它会返回后续请求仍需依赖的节点和保护边界。

2. 以为 lock_ref 只是树节点引用计数#

它实际上在维护“可驱逐 / 被保护”这条缓存状态边界。

3. 以为 len(prefix_indices) 就等于已经安全进入树缓存的长度#

在 page-aligned / partial page 场景下，这常常不成立，必须看 cache_protected_len。

如果 prefix reuse 行为异常，先怎么查#

建议按这个顺序：

  1. 看 match_prefix(...) 到底返回了多长的 device_indices
  2. 看请求上的 last_node 和 cache_protected_len 是否合理
  3. 看 cache_unfinished_req(...) 是否在 chunked prefill 场景下正确更新了保护边界
  4. 看 inc_lock_ref(...) / dec_lock_ref(...) 是否让 protected_size_ 与 evictable_size_ 保持一致
  5. 最后再看 eviction 或 allocator 是否只是放大了前面的边界错误

这样查能避免一种很常见的误区：prefix reuse 看起来失效，就立刻去怀疑 allocator；但问题往往先出在“树和请求之间的保护边界”。

小结#

这一章真正要补齐的，是前缀复用这条主线里最容易被略写的一层：

  • prefix 命中之后，系统还要继续维护树节点保护关系
  • cache_protected_len 用来区分“请求前缀视图”和“树已正式接管的部分”
  • unfinished / finished 两条路径本质上都在做缓存所有权转交

到这里，调度与内存部分对 prefix reuse 的解释就不再只停留在“有 radix tree 能匹配”，而是真正落到了运行时怎样安全持有和释放这段前缀状态。