第 7 章 · 02 kv_prefix.h KV 前缀复用

第 7 章 · 02 kv_prefix.h KV 前缀复用

本节摘要:本节深潜 c/kv_prefix.h——Colibrì 在同一进程内让"第 N+1 轮"复用"第 N 轮"已经算好的 KV 状态的机制。聊天客户端每轮把全文重发是普遍习惯,如果新 prompt 的开头恰好是上一轮引擎已经处理过的 token 序列,那状态就是这个位置的状态——跳过重置,只 prefill 新增尾部。无快照、无回卷,任何位置分叉就从头来。这一节用源码注释 + API 拆解它为什么用"记录"而不是"计数器",以及什么是 tainted。

内容来源:原项目源码 c/kv_prefix.h(注释 33 行 + 实现 97 行,本节重点贴注释精读)

⚠️ 注意:kv_prefix.h 的注释写得比代码还密——这是 Colibrì 的风格:机制设计原理直接写在头文件里。本节会大段引用注释原文,因为它们本身就是设计文档。

学习目标

1. 读懂头文件 33 行注释中 THE IDEA、WHY RECORD NOT COUNTER、INVARIANT、TAINT 四段的设计动机。
2. 理解 kv_prefix_reuse() 的全部失败短路条件,以及为什么"需要一个新 token"。
3. 解释 tainted 标记为何能让 inkling 的音频帧永不复用。
4. 看清 kv_prefix_record() 与 KV 缓冲区生命周期的绑死关系。
5. 用 DeepSeek V4 的 320s→61s 实测理解"流式引擎主导成本是磁盘 IO"。

一、THE IDEA:跳过重置只 prefill 尾部

我们直接读头文件的开场注释,这是全章最值得逐句读的段落之一:

1   /* kv_prefix.h — reuse the attention state a previous turn already built.
2    *
3    * A chat client resends the whole transcript every turn. colibri.c has pinned
4    * each conversation to a KV slot since #639, so its turn N prefills only the
5    * new text; inkling.c, kimi_k3.c and deepseek were written without it and
6    * re-processed turns 1..N-1 from scratch. The cost of a message therefore grew
7    * with the length of the conversation, and every replayed position pulled its
8    * experts off disk again — on a streaming engine that is the dominant cost, not
9    * a rounding error. Measured on DeepSeek V4, a second turn that reused 82% of
10   * its prompt took 61s instead of 320s.
11   *
12   * THE IDEA. After a turn, the engine's state covers some number of positions.
13   * If the next prompt BEGINS with the token sequence that produced them, that
14   * state already IS the state at that position: skip the reset and prefill only
15   * the tail. No snapshot, no rewind — a prompt that diverges anywhere starts
16   * over.
​

要点拆解:

• 聊天客户端(包括 OpenAI 协议客户端)的普遍习惯是每 turn 把整个对话历史重新发一遍——不是增量。
• colibri.c 从 PR #639 起,把每个对话钉在一个 KV slot 上,所以第 N 轮只 prefill 新增文本。但 inkling.c、kimi_k3.c、老的 deepseek 实现没有这套,每轮从 1..N-1 重新 prefill。
• 在流式 MoE 引擎里,每重放一个位置就要把它的路由专家从磁盘读一遍。这不是"四舍五入"的损耗,而是主导成本。实测 DeepSeek V4 第二轮复用 82% 的 prompt,从 320 秒降到 61 秒。

THE IDEA 那段(12-16 行)是机制核心:turn 之后引擎状态覆盖了若干 position,如果下一 prompt 的开头恰好是产生这些状态的 token 序列,那这个状态就是这个位置的状态——直接跳过 reset,只 prefill 尾部新内容。无快照、无回卷——一旦 prompt 在任何位置分叉,就从头来。

💡 深潜要点:这套机制故意做得乐观——它假设大多数对话下一轮的开头就是上一轮的全历史 + 新增,所以"复用"是常态。但它绝不承诺一定能复用:失败立刻退化成"从头 prefill",正确性永远在。

二、WHY RECORD NOT COUNTER:错不会崩但会静默说错话

最反直觉的设计决策是"为什么不用计数器"。注释给得很直白:

18   * WHY A RECORD AND NOT A COUNTER. It is tempting to derive the reusable length
19   * from the caller's own bookkeeping ("prompt_count + generated - 1"). That
20   * invariant differs per engine — whether the last sampled token was fed back,
21   * whether a chunked prefill ran to completion, whether generation stopped early
22   * — and getting it wrong does not crash: it silently answers from a state that
23   * belongs to a different conversation. So the ids are recorded WHERE THEY ARE
24   * FED, and the record is the only description of the state anyone consults.
​

诱惑是这样的:既然知道"用户发了多少 token + 引擎生成了多少",似乎能算出复用长度。但每个引擎的细节不同:

• 最后一个采样的 token 是否被 fed back 进 KV?
• 分块(chunked)prefill 是否真的跑完了?
• 生成是否 early stop?

这些 per-engine 不变量稍有差池,算出的"复用长度"就会和真实状态错位。而错位不会崩溃——引擎会从"另一个对话的状态"开始回答,产出看似合理但语义已偏的输出。这是最危险的一类 bug:静默错误。

所以 Colibrì 选择在 token 被 fed 进 KV 的地方记录它们(kv_prefix_record),让这份记录成为唯一描述状态的内容。任何"上层 bookkeeping"都不得参与复用判断。

三、INVARIANT 与 TAINT:状态描述的两条铁律

26   * INVARIANT: fed[0..len-1] are exactly the token ids the current state was
27   * built from, in position order. Everything else follows from it.
28   *
29   * TAINT. Some inputs are not described by their token ids. Inkling's audio
30   * frames all carry the same id (c->audio_tok) while the mel payload differs, so
31   * an id-only comparison would cheerfully "match" two different clips. A state
32   * that consumed such an input is marked tainted and is never reused.
​

INVARIANT:fed[0..len-1] 恰好是当前状态由之构建的 token id,按位置顺序排列。所有其他逻辑都从这条不变量推出。

TAINT:有的输入无法用 token id 描述。inkling.c 是 Colibrì 的多模态引擎,音频帧在 token 流里都共用同一个 id c->audio_tok,但每帧的 mel(梅尔频谱)载荷不同。如果只比对 id,系统会兴高采烈地"匹配"两段不同的音频——这会从错误的状态回答。所以一旦状态消费了这类输入,就打上 tainted=1 标记,永不复用。

这两条是"复用安全"的全部根基。

四、API 与 KV 缓冲区的生命期绑死

来看实现。结构体只有四个字段:

40  typedef struct {
41      int *fed;      /* token ids at positions 0..len-1 */
42      int  len;      /* positions the state currently covers */
43      int  cap;      /* allocated positions */
44      int  tainted;  /* state consumed something token ids cannot describe */
45  } kv_prefix;
​

最关键的 API 是 kv_prefix_alloc 和 kv_prefix_grow——它们与 KV 缓冲区的(重)分配严格配对:

47  /* Size the record to the KV it describes. Call it wherever the KV buffers are
48   * (re)allocated: growing them discards the positions fed[] refers to, so the
49   * record has to be dropped at the same moment. ...
52  static inline int kv_prefix_alloc(kv_prefix *p, int cap) {
53      if (!p) return 0;
54      free(p->fed);
55      p->fed = (cap > 0) ? (int *)calloc((size_t)cap, sizeof(int)) : NULL;
56      p->cap = p->fed ? cap : 0;
57      p->len = 0;
58      p->tainted = 0;
59      return p->fed != NULL;
60  }
​

注释点破:调用它的位置必须就是 KV 缓冲区(重)分配的位置。因为 KV 缓冲区一旦长大,fed[] 引用的旧 position 就被丢弃了,记录必须同时清掉。配对错位 = 静默错误。

kv_prefix_grow()(86-98 行)处理"KV 缓冲区是 copy-grow"的引擎——这种情况 position 内容被保留,fed[] 也得一并 copy 保留:

86  static inline int kv_prefix_grow(kv_prefix *p, int cap, int keep) {
87      if (!p || cap <= 0) return 0;
88      int *grown = (int *)calloc((size_t)cap, sizeof(int));
89      if (!grown) { kv_prefix_free(p); return 0; }
90      if (keep > p->len) keep = p->len;
91      if (keep > cap)    keep = cap;
92      if (keep > 0 && p->fed) memcpy(grown, p->fed, (size_t)keep * sizeof(int));
93      free(p->fed);
94      p->fed = grown;
95      p->cap = cap;
96      p->len = keep > 0 ? keep : 0;
97      return 1;
98  }
​

keep 是"扩容时保留前几个 position"。注释明确:这种情况正是复用机制存在的理由——一个对话的 prompt 每轮都在变长。如果 grow 失败,记录直接 free 成空,而不是留一份 stale 的——kv_prefix_free(p); return 0; 这一行就是诚实假设的代码化:复用机制失败绝不拖累 turn。

kv_prefix_record() 在 token fed 进 KV 时被调用:

103  static inline void kv_prefix_record(kv_prefix *p, const int *ids,
104                                      int pos0, int n) {
105      if (!p || !p->fed || !ids || n <= 0 || pos0 < 0) return;
106      if (pos0 + n > p->cap) { p->len = 0; return; }
107      memcpy(p->fed + pos0, ids, (size_t)n * sizeof(int));
108      if (pos0 + n > p->len) p->len = pos0 + n;
109  }
​

注意第 106 行:写超范围就直接清空 len,绝不截断成部分记录——因为部分记录会"声称覆盖状态实际没有的 position"。

五、reuse 的全部失败条件

kv_prefix_reuse() 的代码极简,但每一条 return 0 都是一道防线:

123  static inline int kv_prefix_reuse(const kv_prefix *p, const int *ids, int n) {
124      if (!p || !p->fed || !ids) return 0;
125      if (p->tainted || p->len <= 0 || p->len >= n) return 0;
126      if (memcmp(p->fed, ids, (size_t)p->len * sizeof(int)) != 0) return 0;
127      return p->len;
128  }
​

四个失败条件:

1. 记录不存在(!p || !p->fed || !ids);
2. tainted 或 len 为零(无状态可复用);
3. p->len >= n——记录长度大于等于新 prompt 长度。这一条对应注释 119-122 行:如果新 prompt 是记录的"前缀"或"相等",意味着需要回卷(rewind)状态——这套机制做不到。prefill 零个 token 也没意义(没有最终隐状态可采样)。所以至少要有一个新 token(len < n);
4. memcmp 不为零——开头任一 token 分叉,从头来。

只有四关都过,才返回可复用长度 p->len。调用方据此跳过 reset,从 position len 开始 prefill 尾部 n - len 个 token。

💡 深潜要点:kv_prefix_reuse 的失败处理体现了 Colibrì 的整套工程心法——优化是乐观的、安全是悲观的。复用能省则省,但任何不确定立刻退化成"从头 prefill",代价是性能,绝不搭上语义。

本节要点回顾

• 聊天客户端每 turn 重发全文是普遍习惯;colibri.c 自 PR #639 起每对话钉一个 KV slot,只 prefill 新增文本。复用 82% prompt 时 DeepSeek V4 从 320s 降到 61s。
• THE IDEA:新 prompt 开头若等于上一轮已处理的 token 序列,状态就是该位置状态,跳过 reset 只 prefill 尾部。无快照、无回卷,任意分叉从头来。
• WHY RECORD NOT COUNTER:per-engine 不变量不同,计数器算错会静默从错误对话状态回答。所以 token 在 fed 进 KV 的位置记录,reuse 只信这份记录。
• INVARIANT:fed[0..len-1] 恰是当前状态的 token id;TAINT:inkling 音频帧同 id 不同 mel,tainted 状态永不复用。
• kv_prefix_reuse 四道防线:记录存在、未 tainted、至少一个新 token、memcmp 完全相等——任一不过就从头 prefill。

下一节我们把第 7 章收口:把 token 精确前向验证、MLA 57× 压缩、KV 持久化、kv_prefix 复用四件事串起来,讲清 Colibrì 的 faithful DSA 与"压缩状态不是另一个模型"的边界。