0. 为什么 linear-state 不能照搬 KV prefix cache
KV prefix cache 之所以能省算力,是因为 KV block 可以多 owner 共享:同一段前缀的 KV block 本就在显存里,缓存命中时只要把引用计数加一、指过去即可,零拷贝。
GDN(gated delta net)这类线性注意力层不同。它每一层维护一份 recurrent state(conv + ssm 两块张量),按 slot 存储,而且 kernel 每一步都会原地覆写这个 slot。这带来一个硬约束:
于是 linear-state prefix cache 用了一对不对称的操作:
| KV prefix cache(熟悉的) | linear-state prefix cache(本文) | |
|---|---|---|
| 存进 cache | insert() | promotion(晋升 / 换手) |
| 缓存对象 | KV block,可多 owner 共享 | recurrent state slot,永不共享、单 owner |
| 存的方式 | block 本就在,加 ref 指过去,零拷贝 | 把序列当前 live slot 整个晋升为 checkpoint,序列换新空 slot,也零拷贝 |
| 取出 | match() 返回共享 block(加 ref) | match() 返回 pin,worker 拷贝进新 live slot(不可共享,必须拷) |
一句话:save 侧靠 scheduler 零拷贝换手,restore 侧靠 worker 真拷贝——两半不对称,是本文所有细节的根源。
1. save / restore 闭环全景
LinearStatePrefixCache(scheduler 侧,每 dp_rank 一个)
内部一个 SingleBlockManager 管全部物理 slot [0, num_slots)
┌─────────────────────────────────────────────────────────┐
│ free / live / checkpoint 三态共享同一 id 空间 │
└─────────────────────────────────────────────────────────┘
▲ allocate_live_slot │ match(hash) 命中
│ 分配空 slot 给运行序列 ▼ pin 住 checkpoint
┌──────────┐ set_slot ┌──────────┐ promotion ┌──────────────┐
│ 序列 A │◄─────────│ live slot │────────────►│ checkpoint │
│ 运行中 │ 持有 │ 私有, │ 晋升(换手) │ 公共,冻结, │
└──────────┘ │ 原地覆写 │ │ LRU 可回收 │
│ └──────────┘ └──────────────┘
│ 第一次 forward 且复用块对齐前缀 │
│ → restore: worker 把 checkpoint 拷进我的 live slot │
└────────────────────────────────────────────────────┘
restore 是真拷贝,因为 GDN state 不可共享
图里 promotion(live slot ──► checkpoint 那条向右箭头)是 save,走 scheduler 侧零拷贝换手;restore(底部那条回到 live slot 的回环)走 worker 侧跨层拷贝。中间那个 SingleBlockManager 用同一片 id 空间管理三种状态的 slot,靠引用计数和所属容器区分。
2. 两个 slot 概念,务必分清
| 概念 | 谁持有 | 特征 |
|---|---|---|
| live slot | 运行序列(Sequence::linear_state_slot_) |
worker 每步原地覆写;私有 |
| checkpoint slot | prefix cache(cached_slots_) |
冻结某前缀边界的 state;可被 LRU 回收、可被后来请求 restore |
两者共享同一物理 id 空间(同一个 SingleBlockManager),靠 ref_count + 所属容器区分状态。理解「promotion 是把一个 slot 从 live 身份换成 checkpoint 身份」是读懂后面所有内容的前提。
3. save 路径:promotion 零拷贝换手
"promotion(晋升)"的字面含义:一个 slot 本是某序列私有的 live slot(正被覆写),prefix cache 把它"晋升"为公共的 committed checkpoint(冻结、可被 restore、可被 LRU 回收)。身份从「私有运行态」升格为「公共缓存态」。
3.1 为什么必须延迟到 forward 之后
promotion 不能在调度阶段就换手,因为那时 slot 里还是上一步的内容。这一步的 forward 才会写入「块对齐边界的最终 recurrent state」——那才是要冻结的 checkpoint。所以拆成两阶段:resolve 阶段(forward 前)只预定哪些 slot 要晋升;commit 阶段(forward 后)才真正换手。
一个 step 的生命周期
════════════════════════════════════════════════════════════════
[调度阶段] [forward 阶段] [forward 之后]
resolve_cache_ops worker 跑模型 commit_reservations
覆写 live slot
──────────────────────►─────────────────────►──────────────────────► 时间
决定: slot 7 本步末尾 模型把"块对齐边界的 此刻 slot 7 里 =
要晋升成 checkpoint 最终 recurrent state" "本步结束的正确内容"
① 预分配新空 slot 9 写进 slot 7 ↓ 换手(见 3.2):
② 记 Promotion{hash, - slot7 晋升 checkpoint
seq, live_slot=7, - 序列改持 slot9
replacement=9} - slot9 标冷
暂不动序列!
3.2 换手三步
换手前: 换手后:
序列 ──持有──→ [slot 7] 序列 ──持有──→ [slot 9](新,空,冷)
↑刚算好的
块对齐 state [slot 7] ←─pin── prefix cache
committed checkpoint, hash=H
- move 出、不 free:把 slot 7 从序列身上 move 出来(
reset_linear_state_slot),但不还回 free list。 - 晋升为 checkpoint:
insert_checkpoint(hash, slot7),slot 7 进入 cache 表,可被后续请求命中。 - 换新 slot + 标冷:序列改持预分配好的 slot 9,并把冷暖位重置为冷(见 §5)。
3.3 resolve 阶段会放弃哪些 save
「想 save」不等于「真 save」。resolve 阶段会跳过这几类:
- 去重:同一批里多个序列想存同一 hash,只晋升一次。
- 已在 cache:该前缀已有 committed checkpoint,跳过。
- 晋升条件不满足:
can_promote判定失败。 - 无空闲替补 slot:预分配 replacement 失败则跳过这次 save——保持 cache 稀疏,而不是强行 LRU 驱逐现有 checkpoint。
promotion 不是可选优化层,就是 save 机制本身。不要 promotion 就意味着要么不存 checkpoint(prefix cache 形同虚设),要么用一次额外 device copy 来存(正是 promotion 要省掉的那次拷贝)。它和 KV 零拷贝 insert 同哲学,只因 GDN 不可共享而多了「换手」这一步。
4. restore 路径:worker 侧跨层拷贝
save 是 scheduler 零拷贝换手,restore 却必须在 worker 侧做真拷贝——因为 recurrent state 不可共享,命中的 checkpoint 内容得实打实复制进当前请求的 live slot。
┌── scheduler 侧(先跑)──────────────┐ ┌── worker 侧 ─────────────┐
│ resolve_cache_ops │ │ restore_linear_state_slots│
│ - 算 prefix hash → 查表 │ → │ - discover_num_slots │
│ - 命中: 把 src_slot_id 写进 op │ │ - 跨层 narrow/copy_ │
│ - 不命中: src_slot_id 保持 -1 │ │ - 就地写 has_initial_state│
└─────────────────────────────────────┘ └───────────────────────────┘
4.1 has_initial_state:给 kernel 的冷/暖开关
GDN kernel 靠一个 has_initial_state 向量决定是否把 recurrent initial_state 清零:1 = 暖(用 slot 里已有 state),0 = 冷(从零起算)。它的默认值预填为「这一行是否复用了 KV cache token」——复用了 KV 的请求默认被当作暖。restore 函数只覆盖它必须改的项。
4.2 三路语义(正确性核心)
逐 op 判定,走三条路之一:
| 路 | 触发条件 | 对 has_initial_state 的动作 | 语义 |
|---|---|---|---|
| SKIPPED | 没请求 restore 且无源 slot | 不动(保留默认) | 续跑请求 / 零前缀冷启:无 checkpoint 可恢复,保留 KV cache 默认值 |
| COLD_START | 请求了 restore 但源 slot 无效 / 越界 / 拷贝失败 | = 0 | 强制冷启,避免把复用的 KV block 误当暖 recurrent state |
| RESTORED | 拷贝成功 | = 1 | checkpoint 已拷进 live slot,标记为暖 |
4.3 跨层拷贝与跨流衔接
拷贝遍历所有层:对每个 linear 层,把 conv cache 的 src slot 拷进 dst slot;ssm 按 checkpoint stride 做 narrow().copy_()。纯 attention 层(conv/ssm 未定义)跳过。若一层都没拷到,说明「有 linear 层」这个不变量破了,当作软失败回退冷启,不崩服务。
这条拷贝 enqueue 在 worker 的 prepare stream 上,不在此处做 host sync;caller 负责在 forward 消费这些 slot 前插一个 stream-event barrier。这与 KV cache 侧的 H2D 写入同构(都在 prepare stream、都靠 event 与 forward 衔接),不是 linear-state 独有的额外等待。
5. linear_state_initialized:live slot 的冷/暖位
这是一个绑在序列上的 bool,表示当前 live slot 里是否已装有效 recurrent state。注意它不是「slot 是否存在」(那是看 Block 是否 valid),而是「slot 内容是否有效」——slot 一分配就有合法 id,但内容是空的,这两件事必须分开。
| 时机 | 操作 | 含义 |
|---|---|---|
| 创建 / reset | 置 false | 新 slot,冷 |
| 每次 forward 末尾 | 置 true | 之后 slot 有 state,别再 restore |
| promotion 换新 slot 后 | 置 false | 换了新空 slot,又冷 |
它唯一的用途是 gate「cold-start restore 只发一次」:
needs_restore = !linear_state_initialized
&& n_kv_cache_tokens > 0
&& n_kv_cache_tokens % block_size == 0;
n_kv_cache_tokens 持续增长,每隔 block_size 就块对齐一次 → 每次都发 restore → worker 把 live slot 覆盖成 checkpoint 旧 state → decode 已累积的 recurrent state 被冲掉 → 输出错乱。
为什么不能用「第几次 forward」之类序列级信号替代:promotion 那条路堵死了。promotion 给序列换新空 slot 时,序列早已 forward 多次,但新 slot 物理上是冷的。任何序列级信号都无法表达「换了 slot 所以又冷」——只有绑定 slot 内容的标志能做到。这个 bool 编码的是 slot 的物理属性,跟着 slot 走。
6. scheduler↔worker 契约:四字段结构体
LinearStateCacheOp 是每个 batch row 一个的小结构体,是 scheduler 把「这一行该不该 restore / 该不该 save、源/目标物理 slot 是谁」传给 worker 的唯一载体。它在 scheduler 侧构造,随 model input params 发到 worker 执行。四个核心字段是两组对称的 hash + slot_id:
struct LinearStateCacheOp {
// 这一行的 live slot(本步覆写目标)
int32_t linear_state_id = -1;
// restore: 从哪个前缀 hash 恢复,scheduler 解析出的源 slot
// worker 把 restore_src_slot_id 拷进 linear_state_id
LinearStatePrefixHash restore_prefix_hash{}; // 非零 → 请求 restore;查表 key
int32_t restore_src_slot_id = -1; // resolve 查出的源 slot
// save: 要 checkpoint 的前缀 hash,scheduler 决定 promote 后的旧 live slot
LinearStatePrefixHash save_prefix_hash{}; // 非零 → 请求 save;cache key
int32_t save_dst_slot_id = -1; // resolve 决定 promote 后的旧 live slot
};
6.1 为什么 hash 省不掉
高频疑问:「能不能只留 slot_id,把两个 hash 删掉?」不能,三条独立理由:
-
restore:hash 是因,slot_id 是果。build 阶段只有 token_ids,能算出
restore_prefix_hash(前缀内容指纹),但「哪个物理 slot 存了这段前缀?」只有 prefix cache 的hash → slot映射知道。必须拿 hash 去match()才能解析出restore_src_slot_id。删掉 hash,没有任何东西能查出源 slot。 -
save:hash 就是 cache 的 key。commit 阶段
insert_checkpoint(hash, slot)用它作索引存进 cache。prefix cache 的本质是内容寻址:未来请求用前缀内容的 hash 找物理资源。只存 slot_id 不存 hash,cache 退化成「一堆没索引的 slot」,等于没有 cache。 -
hash 还兼任「是否请求该操作」的标志。worker 判 restore 是否发生看
restore_prefix_hash是否非零,判 save 看save_prefix_hash非零(作副条件)。这个「零值约定」让单个 op 能携带 restore-only / save-only / 两者 / 皆无四种组合,无需额外 enum。
反过来也不能「只留 hash 删 slot_id」:worker 是 thin executor,跨层拷贝要的是具体整数 slot 号,而且 worker 不持有 hash → slot 映射(那在 scheduler 侧)。必须把 scheduler 解析好的 slot_id 传过去。
hash = 逻辑 / 内容层标识,跨请求跨时间稳定,是 cache 的键,scheduler 用它内容寻址;slot_id = 物理层标识,某次具体分配的瞬时结果,worker 用它做物理操作。scheduler(逻辑寻址)与 worker(物理执行)是分离主体,两层标识都得有。
6.2 唯一可商榷的字段:save_dst_slot_id
四字段里只有这个有「冗余味」,但它不是死字段。它的值在主路径恒等于 linear_state_id(命名 dst 有误导:worker 没把它当拷贝目标,save 也没有 worker 拷贝)。它真正的用途是被 worker 读作「resolve 确认了 save 真会发生」的标志。
这个语义不能用 save_prefix_hash != 0 替代——hash 非零只代表 build 阶段「想 save」,而 resolve 可能放弃这次 save(见 §3.3 的四类跳过)。这些情况下 hash 仍非零,但 promotion 没建、save_dst_slot_id 保持 -1。worker 靠 save_dst_slot_id >= 0 判「本步真有 promotion」。所以它承载的是 resolve 的决策结果,理论上可以用一个 bool will_promote 表达,当前设计复用了 slot_id 字段的对称模式来携带这个信号。
7. 四阶段流水线:字段在哪一棒出生、哪一棒消费
一个 step 内,这个结构体被四个阶段接力读写。hash 在最早一棒(build)就得填,那时 slot_id 还全是 -1——这是「hash 不能省」最直接的物理事实。
阶段① BUILD 阶段② RESOLVE 阶段③ FORWARD 阶段④ COMMIT
scheduler/主线程 scheduler/主线程 worker/子进程 scheduler/主线程
────────────────── ────────────────── ──────────────── ──────────────────
只有 token_ids, 拿 hash 去查/占槽, 只认整数槽号, forward 写完旧槽后,
能算 hash 不知 slot 把结果回填 slot_id 做物理拷贝 真正换手 promote
linear_state_id ──┬────(读,校验)──────►(读,拷贝 dst)───►(读,校验)
│
restore_hash ≠0 ───┼─► match(hash) ──────►(读,restore? 标志)
│ │
│ └─► restore_src_slot_id = 命中槽 ──►(读,拷贝 src)
│ (阶段②才出生)
│
save_hash ≠0 ──────┴─► 查重 ─────────────────────────────────► insert_checkpoint(
│ hash, 旧槽)
└─► save_dst_slot_id = 旧 live 槽 ──►(读,"本步有 save"标志)
(阶段②才出生)
hash 是阶段①(build)就必须带上的输入,因为那是唯一知道前缀内容、却还不知道物理 slot 的时刻;两个 slot_id 是阶段②(resolve)用 hash 查表后才回填的输出。先有 hash(因),后有 slot_id(果)。
7.1 关键非对称,为什么不能机械对称地省字段
-
restore_src_slot_id由match()查出,是历史 checkpoint slot;save_dst_slot_id= 序列已持有的旧 live slot,不是新分配的。save 真正新分配的是替补 slot,它不进 op、只进 promotion 预定,commit 时才换给序列。 -
restore 在 worker 真做 device 拷贝;save 在 scheduler 零拷贝换手,worker 不拷。这就是为什么
save_dst_slot_id在 worker 侧只当标志用。
7.2 结论
hash 是 scheduler 做内容寻址的查询键 / cache 键(build 阶段唯一能提供的信息,且兼任 restore/save 请求标志),slot_id 是 scheduler 用 hash 查表后回填、传给 worker 做物理拷贝的执行参数——因(hash)与果(slot_id)、逻辑层与物理层、调度主体与执行主体,缺一不可。真要精简,只有 save_dst_slot_id 像个可优化的 flag,但那和「用 slot id 取代 hash」是两回事。