xLLM Notes / notes / linear-state-prefix-cache-internals
Knowledge · Linear-State Cache

Linear-State Prefix Cache 内部机制

GDN 类线性注意力模型的 recurrent state 无法像 KV block 那样零拷贝共享。这篇讲 prefix cache 如何用「零拷贝换手(promotion)save + worker 侧拷贝 restore」这对不对称操作绕过限制,以及 scheduler↔worker 之间那个四字段契约结构体每个字段的必要性。

Note GDN / Linear Attention Prefix Cache Recurrent State

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。这带来一个硬约束:

recurrent state 永不共享、单 owner。如果两个序列指向同一个 slot,其中一个 forward 一步就会把另一个的状态冲掉。所以 KV 那套「加个引用指过去」的零拷贝共享在这里根本走不通。

于是 linear-state prefix cache 用了一对不对称的操作:

KV prefix cache(熟悉的)linear-state prefix cache(本文)
存进 cacheinsert()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 真拷贝——两半不对称,是本文所有细节的根源。

本文对应已发布的设计报告 Qwen3.5 Linear-State Prefix Cache(PR #1839); block / slot 管理背景见 KV Cache Block 管理架构; 整体执行流见 xLLM 推理执行主干

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
  1. move 出、不 free:把 slot 7 从序列身上 move 出来(reset_linear_state_slot),但不还回 free list
  2. 晋升为 checkpointinsert_checkpoint(hash, slot7),slot 7 进入 cache 表,可被后续请求命中。
  3. 换新 slot + 标冷:序列改持预分配好的 slot 9,并把冷暖位重置为冷(见 §5)。
第①步「move 出但不 free」是关键:如果用普通的「erase 后 free」语义,会先把 slot 还回 free list,随即可能被别的序列分走并覆写——checkpoint 内容就被腐蚀了。零拷贝 save 依赖「所有权转移而非释放」。

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,标记为暖
为什么 SKIPPED 必须"不动"而不是"写 0":续跑(decode 中)请求的 live slot 本来就是暖的,它没请求 restore 恰恰是因为它不需要拷贝。若 SKIPPED 也写 0,会把暖续跑请求错误地冷启,输出立刻错乱。三路不能合并,是正确性问题而非代码风格。

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;
为什么不能删(硬反证):删掉这个 gate,restore 条件退化成「只要块对齐就 restore」。decode 阶段 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
};
两组对称字段的由来:restore 和 save 是镜像操作,各配一个 hash(逻辑 / 查表)+ 一个 slot_id(物理 / 执行)。但两组的「拷贝方向」和「谁来拷」完全不同——restore 在 worker 真拷贝,save 在 scheduler 零拷贝换手。这正是它们省不掉、也不完全对称的根因。

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_idmatch() 查出,是历史 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」是两回事。