Qwen3.5 Linear-State Prefix Cache

1. 执行摘要

xLLM 的前缀缓存原本只服务标准 KV attention:相同前缀的 token 可零拷贝共享物理 KV block。但 Qwen3.5 引入了 GDN(Gated Delta Net)linear-attention 层——以 9B 为例,32 层里有 24 层是 linear、8 层是 full attention。linear 层不保存逐 token 的 KV,而是维护一份被累积压缩的 recurrent state。这份 state 每步被 kernel 原地覆盖,无法像 KV 那样按 block 引用计数共享。

本 PR 为 linear-state 建立一套独立的 checkpoint 缓存:在稀疏的 prefill chunk 边界把当时的 recurrent state 提交为一个 checkpoint slot,新请求命中同前缀时,把最深的那个 checkpoint 拷回 live slot 作为热启动起点,从而跳过对已知前缀的重复 prefill。KV 与 linear 两条 leaf 各自维护缓存,由 CompositeBlockManager 统一编排复用长度。

2. 问题与约束

2.1 为什么不能照搬 KV 引用计数

KV 前缀缓存能零拷贝共享,前提是 KV block 一旦写入就不可变:序列 A 和序列 B 共享同一段前缀 block,谁都不会改写它。而 GDN 的 recurrent state 相反——kernel 每一步都 index_put_ 原地覆盖 slot 内容。若两个序列共享同一 slot,后写者会污染先写者。因此 linear-state 的复用只能是拷贝语义:命中时把 checkpoint 的内容拷贝进请求私有的 live slot,而非共享指针。

2.2 累积压缩带来的"单点恢复"

KV attention 逐 token 独立,可以在任意 block 边界冷启动——前面复用的 block 是精确值。但 recurrent state 是把历史累积压缩进固定大小的 state;它只在被显式 checkpoint 的那些边界上是"精确快照"。因此:

• 复用只能落在 checkpoint 存在的稀疏 chunk 边界上,不能落在任意 KV block 边界。
• 命中多个 checkpoint 时,只挂载最深的一个——它的 state 已包含之前所有 checkpoint 的信息。
• 若强行按 KV 匹配长度冷启动到一个没有 checkpoint 的边界,recurrent state 会损坏(KV 前缀是对的,但 state 是错的),输出语义漂移。

2.3 slot id 空间约束

id 池是统一的:slot id [1, num_slots) 既服务 LIVE slot(被某序列持有、正在原地更新),也服务已提交的 CHECKPOINT slot(被 prefix_cache_ 拥有),两者在引用计数下可互换。slot 0 是继承来的 padding block。GDN live slot id 是通用的、且被 embedding_cache 共享,所以重构 linear cache 不能把 checkpoint id 扩进这个受限空间之外。

3. 系统架构

3.1 核心组件

3.2 架构图

flowchart TB
  subgraph sched["调度层 (Scheduler / Composite)"]
    seq["Sequence<br/>linear_state_hashes_"]
    composite["CompositeBlockManager<br/>跨 leaf min + chunk 对齐 + KV clamp"]
  end

  subgraph leaves["Block Manager Leaves"]
    kv["KV leaf<br/>BlockManagerImpl<br/>零拷贝共享"]
    linear["LINEAR leaf<br/>LinearStateBlockManager<br/>checkpoint 拷贝"]
    lidx{"LinearStatePrefixCache<br/>checkpoint 索引 (LRU)"}
  end

  subgraph worker["Worker 执行层"]
    op["LinearStateCacheOp<br/>live id / restore / src id"]
    restore["restore_linear_state_slots<br/>src -> live 逐层拷贝"]
    forward["Model Forward<br/>has_initial_state = warm/cold"]
  end

  seq --> composite
  composite -- allocate_shared --> kv
  composite -- "allocate_shared token/hash" --> linear
  linear --> lidx
  lidx -- "最深 checkpoint" --> linear
  composite -- "min(kv_len, linear_len)" --> op
  op --> restore --> forward

  classDef schedNode fill:#eef0fd,stroke:#3b4bd8,color:#1b2130;
  classDef kvNode fill:#e7ecff,stroke:#232a52,color:#1b2130;
  classDef linNode fill:#f6e5df,stroke:#c1583b,color:#1b2130;
  classDef wkNode fill:#f3f5f9,stroke:#5a6478,color:#1b2130;
  classDef idxNode fill:#fbe9e7,stroke:#b3261e,color:#1b2130;

  class seq,composite schedNode;
  class kv kvNode;
  class linear,restore linNode;
  class lidx idxNode;
  class op,forward wkNode;

3.3 设计原则

4. 独立 Hash 域

这是本 PR 最关键、也最早踩坑的设计点。早期实现让 linear checkpoint 复用 KV 的 by-block hash,结果是"存了却查不到":checkpoint 按 KV block 边界的 hash 存入,但 linear 的复用探测发生在 chunk 边界,两套边界不重合,匹配恒失败。

正确方向是让 linear 拥有自己的 hash domain:

• 步长不同:KV 按 KV block(每 block_size 个 token)链式 hash;linear 按 prefill chunk(chunk_stride 个 token)链式 hash。
• 单源缓存:链式 hash 在 Sequence::linear_state_hashes_ 上预计算一次,随 prefill 推进由 update_linear_state_hashes(chunk_stride) 增量补齐,match 时经 block_hashes 参数直接传入,不重算——与 KV 传 seq->block_hashes() 对称。
• 提交用记录 hash:checkpoint 插入索引时用当时边界的 recorded hash(insert_with_recorded_hash),而非用序列当前 token 重算——否则会 key 到更深的当前边界,checkpoint 永远匹配不上。
• KV 复用取 min:最终复用长度是 min(kv_match, linear_reach),由 composite 完成。linear 只回答"我最远能恢复到哪(L)",KV 复用被 clamp 到 L。

5. 核心工作流

5.1 命中与恢复(restore)

1. 调度器为新请求准备 token span,update_linear_state_hashes(chunk_stride) 补齐链式 hash。
2. composite 对每个 leaf 调 allocate_shared:KV leaf 返回可共享 block 并计入序列;LINEAR leaf 走 match() 只探测最深 checkpoint,不计入 admission。
3. LINEAR 的 match 用 contains() 探测(不动 LRU),命中后用 find() 取回最深 checkpoint block(refcount+1、LRU 提升),作为 restore 源暂存,不进序列的 LINEAR block 向量。
4. composite 取 min(kv_len, linear_reach),做 chunk 对齐与 KV-block clamp,生成每个序列的 LinearStateCacheOp(live id、restore_requested、restore_src_slot_id)。
5. worker 侧在 forward 前调 restore_linear_state_slots:按 op 把 restore_src_slot_id 跨所有 linear 层拷进 linear_state_id,并回写 has_initial_state(RESTORED→warm=1,COLD_START→cold=0)。
6. forward 消费 warm state 从 checkpoint 边界继续,跳过已恢复前缀的 prefill。

5.2 保存(save / checkpoint 提交)

LINEAR leaf 的每步入口 allocate_for_sequence 拥有完整 slot 生命周期,save 与 KV 只 append 不同,是 replace-vs-append 的非对称:

1. 先执行上一步延迟的 save:序列的 live slot 此刻正好持有刚算完的边界 state,把它提升进 checkpoint 索引(insert_with_recorded_hash),并把序列原地轮转到一个新的 live slot(replace_block)。
2. 若序列此刻还缺 live slot,再 acquire 一个。

save 通过 replace_block 原地轮转、不返回新 block;composite 只 append step 2 新 acquire 的。save 是纯 host 侧簿记,不需要 worker 拷贝。

5.3 时序图

sequenceDiagram
  participant Sched as Scheduler / Composite
  participant KV as KV leaf
  participant LIN as LINEAR leaf
  participant IDX as Checkpoint Index
  participant WK as Worker (restore + forward)

  Sched->>Sched: update_linear_state_hashes(chunk_stride)
  Sched->>KV: allocate_shared(token_ids, block_hashes)
  KV-->>Sched: shared blocks, kv_len
  Sched->>LIN: allocate_shared(token_ids, linear_hashes)
  activate LIN
  LIN->>IDX: match() probe (contains, no LRU)
  IDX-->>LIN: 最深 checkpoint hit / miss
  LIN->>IDX: find() 取回 (refcount+1)
  LIN-->>Sched: restore src (不入序列), linear_reach
  deactivate LIN
  Sched->>Sched: min(kv_len, linear_reach) + chunk 对齐 + clamp
  Sched->>WK: LinearStateCacheOp[]  (live / restore / src)
  activate WK
  WK->>WK: restore_linear_state_slots: src -> live 逐层拷贝
  WK->>WK: has_initial_state = warm / cold
  WK->>WK: model forward 从 checkpoint 继续
  WK-->>Sched: step 完成
  deactivate WK

5.4 prefill 分配的 gate

get_prefill_allocate_tokens 里改用 num_tokens() 的分支只在 has_linear_attention_layers_ && enable_prefix_cache_ 时走;纯 KV 路径逐字节回退到主分支的 kv + step 逻辑,避免对非 linear 模型产生行为漂移。

6. 关键数据结构与接口

6.1 LinearStateCacheOp(调度器 → worker)

struct LinearStateCacheOp {
  // 序列推进 recurrent state 的 live slot。
  int32_t linear_state_id = -1;

  // restore 请求位 + 调度器解析出的 checkpoint 源 slot。
  // worker 把 restore_src_slot_id 拷进 linear_state_id。
  // 保留为独立 bool(不从 src_id>=0 推导),让
  // "请求了 restore 但没解析到源 -> 强制冷启动" 能跨 IPC 存活。
  bool    restore_requested = false;
  int32_t restore_src_slot_id = -1;
};

6.2 LinearStatePrefixCache::match

// 把 token_ids 按 chunk_stride 切成链式 per-chunk key,
// 对 STRICTLY BELOW final token 的每个整 chunk 探测索引,
// 找最远可恢复前缀。每个命中把 reach 扩到
// (chunk_index + 1) * chunk_stride。允许 gap:后面的命中可
// 越过中间的 miss。matched_tokens 写回该 reach(chunk_stride 的整数倍;
// 全 miss 时为 0)。因累积压缩,只挂载最深的单个命中。
std::vector<Block> match(
    const Slice<int32_t>& token_ids,
    const Slice<Block>&   existed_shared_blocks = {},
    const MMData&          mm_data = MMData(),
    const Slice<XXH3Key>&  block_hashes = {},   // 序列预算的 linear 链式 hash
    size_t*                matched_tokens = nullptr) override;

• probe 用 contains(),不动 LRU 时序。
• 命中后用 find() 取回最深 checkpoint(refcount+1、LRU 提升),pin 住直到 restore。
• leaf 保持 KV-blind:只在自己 token/chunk 域计数,不看 KV 匹配长度。

6.3 LinearStateBlockManager

class LinearStateBlockManager final : public BlockManagerImpl {
 public:
  // chunk_stride: checkpoint 步长(一个 prefill chunk 的 token 数),
  // 转发给索引作为 hash 域步长。默认 -1(禁用探测),
  // 供只驱动 slot 池、不走 token match 的 block 级单测。
  LinearStateBlockManager(uint32_t num_slots,
                          int32_t  kv_block_size,
                          int32_t  chunk_stride = -1);

  // 每步入口:先执行上一步延迟的 save(提升 + 原地轮转),
  // 再按需 acquire live slot。replace-vs-append 非对称。
  std::optional<std::vector<Block>> allocate_for_sequence(
      Sequence* seq, size_t num_tokens) override;

  // 探测(不 admit)checkpoint 索引取最远前缀。
  // 与 KV 走同一 virtual,但返回的是 restore SOURCE,
  // 不入序列 block 向量、不计入 admission。
  std::vector<Block> allocate_shared(
      const Slice<int32_t>& token_ids,
      const Slice<Block>&   existed_shared_blocks = {},
      const MMData&          mm_data = MMData(),
      const Slice<XXH3Key>&  block_hashes = {},
      size_t*                matched_tokens = nullptr) override;
};

6.4 restore(worker 侧)

// 逐 op in-place:把 restore_src_slot_id 跨所有 linear 层拷进
// linear_state_id,并回写 has_initial_state:
//   RESTORED   -> 1 (warm)
//   COLD_START -> 0 (miss / 越界 / 拷贝失败)
//   无 restore 请求 -> 不动(continued request 保持 kv-cache 默认)
void restore_linear_state_slots(
    std::vector<KVCache>&                     kv_caches,
    const std::vector<LinearStateCacheOp>&    cache_ops,
    std::vector<int64_t>&                     has_initial_state);

slot 归属与 LRU 驱逐在调度器侧的 LinearStateBlockManager;此 helper 纯粹是 cache 指示的 worker 侧 restore 拷贝。save 不需要 worker 拷贝(由调度器侧提升完成),故此 API 只覆盖 restore。

7. 并发与正确性

7.1 restore 的 stream 归属

restore 拷贝入队在调用点的当前 stream。修复后的落地方案:restore 移到 worker 线程 step_for_schedule_overlap 入口,走 compute_stream;单线程 ThreadPool + 同一 stream-ordered 语义即自动排序——同 stream 上的 copy_ 与 forward 依赖天然有序,不需要 synchronize_default_stream 这种 stop-the-world。

7.2 save barrier 可删

worker_service.cpp 里的 wait_for_linear_state_save 等的是 NO_SYNC future、不等 device,而 save 本身是纯 host 簿记,因此该 barrier 可直接删。在 overlap + graph 组合下实测验证通过。

7.3 MTP / graph 兼容性

• MTP + graph 已知不兼容:测 MTP 必须关 graph。回归时把「无 MTP + graph」与「MTP + 无 graph」两种组合分开跑。
• MTP model path:主模型与 draft 必须同规格(混用不同尺寸会因 rms_norm shape 不匹配而 crash)。
• overlap restore 的 override 保留:MTP 内层 worker 关 overlap 走 prepare restore、外层无 kv_caches_,若强行搬到基类既无效又会破坏 LLM 的 NO_SYNC 语义——保留该 override。

7.4 admission flag

per-sequence 的 LINEAR leaf 一律 participates_in_admission=false(判据是语义而非 block_size 数值)。checkpoint 是 restore 源,不占 admission 预算。此外 LINEAR leaf 的 block_size 是真实 KV 值,不再压成魔法值 1。

8. 测试与验证

8.1 单元测试

8.2 端到端回归

• 回归覆盖串行 + 并发两种 workload,避免只验串行而漏过并发下的语义漂移。
• 性能结论以预热后的客户端侧 TTFT 为准,服务端 request 日志仅作辅助排查。
• 验证用 Qwen3.5-9B 即可(24 linear + 8 full),不必跑更大规格。
• 全量单元测试通过。

9. 风险与未决问题