跳转至

实验二 Bonus：在 ai_daemon 中实现最小前缀缓存

Bonus 说明

本页对应实验二的选做 Bonus。Bonus 满分 5 分，需要助教进行人工检查，不计入平台自动测试。

建议你先完成基础实验，再考虑这个 Bonus。如果基础实验还没有通过，就不建议过早开始做Bonus。（Bonus实验需要自行理解概念，建议使用搜索引擎或者AI了解Prefill以及KV cache的相关概念后开始实验）

这个 Bonus 在做什么

基础实验的重点是把 HW2 的推理后端封装成 AI Service。这个 Bonus 则是在此基础上，继续为 ai_daemon 增加一个最小可用的 prefix cache，让它在遇到共享 prompt prefix 的请求时，尽量复用已有中间状态，减少重复 prefill。

开始前你需要知道什么

如果你已经完成了 HW2，那么你应该已经知道一次推理大致可以分成 Prefill 和 Decode 两个阶段。这个 Bonus 主要关心的是前者，也就是“把整段 prompt 读进去，并建立后续生成所需上下文状态”的那一段过程。

之所以把注意力放在 Prefill，是因为它会把同一段前缀反复算很多次。假设两次请求分别是：

0 1 2 3
0 1 2 4

它们前面的 0 1 2 是共享的。如果每次都从头开始做完整 Prefill，那么这段共享前缀对应的计算和缓存写入就会被重复执行。prefix cache 想解决的，就是这类“共享前缀被重复计算”的问题。

这里要特别注意：prefix cache 复用的不是“最终生成结果”，而是前缀对应的中间状态。在当前代码里，这些可复用状态最核心的就是已经建立好的 KV cache 以及与前缀长度相关的运行时进度。也正因如此，恢复缓存后，程序可以从某个中间位置继续 forward。

你在代码里会看到一个很关键的边界：很多地方只打算把前缀保存到 prompt_n - 1，而不是整个 prompt。这样做的原因是，最后一个 prompt token 往往和“产出第一个生成 token”直接相连。如果要把整个 prompt 都完整缓存下来，通常还需要额外保存更细的运行时状态，例如最后一个 prompt token 对应的隐藏状态或 logits。对于课程实验来说，这会明显抬高复杂度，因此当前骨架刻意把目标收缩到“先把前 prompt_n - 1 个 token 的复用做起来”。

cold run 和 warm run 的区别也可以这样理解：cold run 表示系统里还没有可用的 prefix snapshot，只能从头做完整路径；warm run 表示系统已经保存过一份可复用前缀，第二次请求可以尝试恢复并从中间继续。你在 aitest --cacheprobe 里看到的对比，本质上就是在观察这两种路径的差异。

最后，当前代码把这个 Bonus 先限制在 Smol 和简单同步路径上，不包含Qwen模型相关。

为什么设计 Bonus

如果你对 prefix cache 还没有概念，可以先把它理解成一件很朴素的事：

很多请求的开头其实都一样，只有最后一小段不同。既然前面那一大段已经算过了，就没必要每次都从头再算一遍。

从这个角度看，prefix cache 有点像“把公共前缀的计算进度先存下来”。下次再遇到同样的开头时，就直接从这个保存点继续，而不是回到起点重跑。

下面几个常见场景里，这个想法都会很自然地出现。

1. Few-shot learning：示例相同，最后的问题不同

在少样本学习里，我们经常会先给模型一组示例，再在最后附上新的问题。例如：

示例 1
示例 2
示例 3
现在请回答这个新问题

如果你连续问很多个类似问题，前面的“示例 1/2/3”往往都不变，变化的只有最后那一句新问题。对模型来说，这意味着前面那段示例部分会被反复读、反复算。prefix cache 的作用，就是把这段公共示例对应的状态保存下来，后面换问题时直接复用。

2. Self-consistency：同一个问题，重复采样多次

有些方法会对同一个问题让模型多回答几次，得到多条不同的推理路径，然后再从这些结果里选一个最一致的答案。

对你来说，这像是在“同一个题目上多次试答”；对模型来说，这些请求的题目部分其实完全一样，不同的只是后面采样出来的推理过程和答案。于是，问题本身这段前缀就很适合做缓存，每次都直接复用前面的问题部分，只重新计算后面分叉出来的答案部分。

3. Multi-turn chat：聊天历史会一轮一轮累积

多轮对话更容易理解。假设你和聊天机器人已经聊了很多轮，那么下一轮请求通常长这样：

历史对话 1
历史对话 2
历史对话 3
……
用户最新一句话

随着对话变长，前面的历史聊天内容会一遍又一遍地出现在新请求里。如果每来一轮新消息，都让模型把整段聊天历史重新从头处理一次，代价会越来越高。prefix cache 的意义就是：把前面已经存在的聊天历史缓存起来，下一轮只在这个基础上继续处理新增的问答内容。

4. Tree-of-thought：不同分支共享同一段“思考起点”

在更复杂的推理任务里，我们有时会把一个问题拆成很多思考分支。你可以把它想成一棵树：

  • 前面有一段公共的分析过程；
  • 后面从某个节点开始，分成不同分支继续探索。

这些分支虽然最后走向不同，但它们往往共享同样的“树干”。对模型来说，这个树干部分就是公共前缀。prefix cache 可以让所有分支复用这段共同历史，只对各自分叉出去的内容做增量计算。

这四个场景的共同点

不管是 few-shot、self-consistency、多轮对话，还是 tree-of-thought，它们的共同点都是：很多请求拥有相同的前缀，只有后半段不同。
只要这个共同前缀足够长、被复用得足够多，prefix cache 就会变得很有价值。

目标

这个 Bonus 的目标不是设计一个完整、泛化的缓存系统，而是完成一个最小可运行的前缀缓存原型。完成后，你应当能够：

  • 理解哪些中间状态值得被缓存；
  • 识别“共享 prompt prefix”的请求；
  • 在合适的场景下恢复前缀对应状态，而不是每次都从冷启动路径重做；
  • 观察 cold run 和 warm run 的差异。

代码入口

建议重点查看 user/ai_daemon.c 中和前缀缓存直接相关的这几个位置（当然，如果你觉得实验设计的不是很好，也可以自由发挥，逻辑自洽即可）：

  • prefix_cache_try_restore()
  • prefix_cache_save_after_prefill()
  • run_decode()

从代码组织上看，这个 Bonus 的主线很清楚：

  1. 在收到新请求时，先判断是否适合走前缀复用；
  2. 如果可以复用，就尝试恢复之前保存的 prefix 状态；
  3. 如果不能复用，就走原来的冷启动路径；
  4. 在合适的时机保存新的 prefix 状态，供下一次请求使用。

实现限制

当前 Bonus 的限制

  • 当前骨架只把这个 Bonus 开放给 Smol 模型。
  • 建议只围绕简单同步路径完成它，不要一开始就和 --async、--batch 等模式混在一起。

你在实现时可以重点把握下面几个问题：

  • 什么时候应该直接绕过缓存；
  • 什么时候两次请求可以被认为共享 prefix；
  • 恢复后应从哪个位置继续 forward；
  • 保存缓存时，究竟应该保存到哪个阶段的状态。

建议测试方法

和基础实验一样，这里只给最小测试入口

下面这些命令足够帮助你先看出 Bonus 是否基本生效，不代表助教侧完整检查项。（Bonus实验最主要还是考察是否对KV cache建立了理解，实现的逻辑是否合理）

建议先做一条最小测试：

aitest --cacheprobe --predict 1 0 1 2 3

如果你已经把Bonus实验完成，这条命令通常会展示一组 cold run 与 warm run 的对比结果。

如果你想再补一条简单检查，可以试一个不适合复用的单 token prompt：

aitest --cacheprobe --predict 1 7

你也可以自己构造两个具有共享前缀、但末尾不同的 prompt，观察前缀缓存是否带来行为变化。不过对于课程实验来说，只要你能清楚说明自己的缓存条件与结果现象，就已经足够。

检查建议

助教人工检查时，通常更关心你是否真正理解了这个优化的边界和价值，而不只是“跑出了一组数字”。你需要可以说明下面几件事：

  1. 你具体缓存了什么状态；
  2. 你在什么条件下允许恢复缓存；
  3. 为什么 prefix cache 更像“减少重复 prefill”的优化；