RFC-3: Unified Hybrid Radix Cache

字段	值
作者	@yuhao
日期	2026-07-06
状态	Draft — current UHRC design overview
目标	将 sgl-jax 的 prefix cache 收敛为 `UnifiedRadixCache` + component 化数据槽（FULL / SWA / RECURRENT），并支持 linear-recurrent 模型的跨请求前缀复用
前置依赖	[[rfc_0_shared_cache_pd_infra]] 的 kv_cache_builder + registry
关联 RFC	[[rfc_1_hicache]]（HiCache 数据搬运）、[[rfc_2_pd_disaggregation]]
上游参考	sglang `UnifiedRadixCache` / Mamba component design

1. 概述

UnifiedRadixCache 将 prefix cache 收敛为一棵 radix tree + component 数据槽。核心树只管理 key、节点、锁、驱逐和 component slot；FULL / SWA / RECURRENT 分别负责自己的 device value、命中范围、锁语义和释放规则。

1.1 范围内

Component-agnostic tree core：UnifiedRadixCache 通过 component_data[ComponentType] 间接持有 FULL / SWA / RECURRENT 数据，不关心具体 pool 布局。
Unified base interface：MatchResult、MatchPrefixParams、InsertParams、EvictParams、lock/unlock result 使用同一套 typed schema，旧 RadixCache / SWARadixCache / ChunkCache 也走兼容后的接口。
Linear-recurrent prefix reuse：KDA / GDN 等模型的 RecurrentStatePool snapshot 可进入 radix tree，并通过 CoW 克隆给后续命中请求。
Paged recurrent production path：--enable-recurrent-extra-buffer 下用 page-aligned track slots 发布 recurrent snapshot，支持 page_size > 1 的生产路径。
DP cache locality：--dp-schedule-policy cache_aware 可按 cache affinity + soft load balancing 分配 dp_rank，减少 shared-prefix 在 DP rank 间重复缓存。

1.2 范围外

项目	归属
HiCache D2H/H2D + HostKVPool + L3 持久化	[[rfc_1_hicache]]
PD 分离	[[rfc_2_pd_disaggregation]]
`(FULL, SWA, REC)` 三组件同树	upstream 结构性不支持（registry 层不可构造）
Token-level recurrent checkpoint	recurrent state 是请求级，不进 RadixKey hash
MTP 接入 / DSV4 SWA L3 evict	独立 milestone

2. 当前设计状态

维度	当前设计
cache 选择	`build_kv_cache` / `registry` 当前在 unified path 构造 `(FULL,)` 或 `(FULL, REC)`；SWA 模型仍走 legacy `SWARadixCache` / `SWAChunkCache`，`(FULL, SWA)` 为 planned
tree core	`UnifiedRadixCache` 统一 match / insert / evict / lock，component 只暴露自己的命中和 value 生命周期
KV pool	MHA / MLA / SWA / RecurrentState / ReqToToken 底层布局保持不变
recurrent	`RecurrentComponent` 管理 request-owned slot、tree snapshot、CoW src/dst、extra-buffer track slots
DP routing	`cache_aware` policy 可在 DP 下优先路由到已有 prefix 的 rank，并在负载倾斜时退回 soft balancing
flags	默认旧路径；`--enable-unified-radix-tree` 和 `--enable-recurrent-extra-buffer` 显式打开对应路径

3. 架构总览

图例：实线是 component value 的持有关系；虚线是 request mapping / recurrent slot allocator 的辅助关系。UHRC 不重排底层 pool，只把 tree node 的 value 统一放进 component_data。

关键关系：

两个数值轴不可混用：① ComponentType 数组下标 FULL=0 / SWA=1 / REC=2，用于 index component_data；② eviction_priority FULL=2 > SWA=1 > REC=0（leaf 一律 0），值越高越晚被回收。已与 upstream 源码核验。
single namespace：节点只用 RadixKey 索引，recurrent state 不进 hash。
match 开销：逐 token O(K)，path 上每个节点遍历 C=3 个 component slot，开销 O(D·C)。C=3 是槽位数组的固定宽度（含 tombstone），与实际存活组件数无关。

4. 模块设计

4.1 `UnifiedRadixCache` + `UnifiedTreeNode`

职责：组件无关（component-agnostic）的 prefix tree 核心，取代现有 4 cache 类。

接口	说明
`match_prefix(key)`	返回 `MatchResult`；命中 = 各组件各自判定命中后取交集
`insert(key, …)`	插入 prefix，按命中组件写 `component_data` 槽位
`evict(EvictParams)`	统一 evict，跨组件按 eviction_priority 级联
`inc_lock_ref` / `dec_lock_ref`	引用计数（FULL/SWA path-to-root；REC single-node）

UnifiedTreeNode 关键字段：

key：RadixKey（带 dp_rank；recurrent state 不进 hash）
component_data：长度 3 的槽位数组，按 ComponentType 下标索引（FULL=0 / SWA=1 / REC=2）
last_access_time：LRU 用的 np.float64 单调递增计数器，不是 wall-clock 时间戳

命名决定：组件正式符号为 RECURRENT / RecurrentComponent，与本仓命名（LinearRecurrentAttnBackend + RecurrentStatePool）对齐。从 upstream 移植 MambaComponent 时，把符号 MAMBA 一律改名为 RECURRENT：代价是每次移植多一步重命名，收益是代码术语与本仓一致。

4.2 `ComponentData` 槽位

tree node 上单个组件的数据槽位，按 ComponentType 下标排布，不带 component_type 字段（与 upstream 一致）。

字段	说明
`value`	该组件在此节点持有的 device 数据：KV indices（FULL/SWA）或 recurrent slot index（REC）
`lock_ref`	该组件的引用计数
`host_value`	本RFC 预留（HiCache 才用），永远为 None
`host_lock_ref`	本RFC 预留

tombstone 规则：每个节点创建时就带 3 个槽位（[ComponentData()×3]）。槽位 value=None 表示"这里没数据"，即 tombstone——core 遇到就跳过。

哪些节点会出现 tombstone？取决于命中时各组件刷新 last_access_time 的范围——未刷新的祖先最终被 LRU 驱逐，value 清为 None：

组件	命中时刷新范围	internal tombstone
FULL	path-to-root（全部祖先）	无——祖先 `last_access_time` 永远 ≥ 子孙
SWA	window 内祖先	有——window 外祖先不刷新；window 边界拆节点时父节点 SWA 数据也被清空
REC	仅命中节点自身	有——所有祖先都不刷新，最容易被驱逐

4.3 `FullComponent`

Full attention（MHA / GQA / MLA），直接复用现有 RadixCache 的 trie / 引用计数 / eviction 算法。

eviction_priority：内部节点 2，leaf 0
path-to-root 锁语义（与 SWA 同，区别于 REC single-node）
host_value/host_lock_ref 字段预留，本RFC 不启用 HiCache 搬运；复用现有字段定义，不新增代码

4.4 `SWAComponent`（planned）

Sliding-window attention。inc_lock_ref 只锁 window 范围内的祖先；window uuid / metadata 需要在 SWA 接入 unified path 时补进 component schema 或等价结构。技术上无阻塞（component-agnostic core 与 upstream 均已支持 (FULL, SWA)），但 scheduler 侧 SWA 能力判断需要 trait 化。

eviction_priority：内部节点 1，leaf 0
planned window metadata 存 window uuid（沿用现有整型 UUID，单线程无竞态）
命中 = SWA 与其他组件各自判定命中后取交集

4.5 `RecurrentComponent`

4.5.1 设计：slot 角色、锁、后端选择

职责：将 linear-recurrent state 纳入 tree node 生命周期，是本RFC核心业务价值。覆盖三个后端（均派生自 LinearRecurrentAttnBackend，共用 RecurrentStatePool）：

KDA（Kimi-Linear，配 MLA）—— 优先级最高，已接入、路径最完整
GDN（Qwen3.5 Gated DeltaNet）
GLA（BailingMoE-v2.5，LightningAttnBackend）

三者互斥，由 HF config 自动选择（非 CLI flag）。

为什么需要 CoW：recurrent state 是会被原地改写的 hidden state（h_t = f(h_{t-1}, x_t)），不像 KV 只追加可只读共享。两个请求在某前缀分叉时，必须 clone state 给新分支，否则写入会污染原请求。CoW 带来两处难点：① 克隆期间须 inc_lock_ref 源节点，防止被 evict；② device 侧克隆——slot→slot 拷贝 state。

克隆范围（正确性关键）：一个 slot 的 state = temporal/recurrent buffer + conv buffer，且是全部 recurrent 层各一份。只克隆 temporal 会导致 cache hit 后 logits 不等价。

三种 slot 角色：

角色	所有权	分配时机	释放时机
active slot	请求	请求开始；命中前缀时由 CoW 克隆	请求结束转给树，或 abort 时 free
tree snapshot	树	请求 finish 时 commit 入叶子节点（对齐 upstream `commit_insert_component_data`）	仅 evict 时释放
CoW dst slot	瞬态	命中 snapshot 后克隆到新 slot	立即成为 active

slot 守恒不变量：任意时刻 active + tree + free == 池容量。

锁 = single-node-lock（区别于 FULL/SWA 的 path-to-root）：克隆时只锁 source 节点（即 §5.1 的 best_match_node），与 upstream MambaComponent.finalize_match_result 一致。

RecurrentStatePool 寻址：head-sharded，slot 经 component_data[2].value 间接持有。RecurrentStatePool 本身只持物理 state buffer；slot 的 free-list 与 request mapping 由 HybridReqToTokenPool 管理。

请求结束时，cacheable recurrent state 不立即释放，而是写入命中的树节点（commit_insert_component_data）成为 tree snapshot。后续请求命中后从 snapshot CoW 克隆一个 active slot；真正不再需要时才在 evict 阶段释放。

4.5.2 CoW device 侧克隆与 page_size

路径	机制
page=1 base	`RecurrentStatePool.copy_slots` 做 temporal+conv 全 state 的 slot→slot CoW；finish 时把 request-owned running slot donate/commit 给 tree snapshot
page>=128 production path	`--enable-recurrent-extra-buffer` 启用 page-aligned ping-pong track slots；scheduler 在 track boundary 切分 prefill；linear backend masked scatter 当前 forward 的最终 state 到 track slot；`RecurrentComponent` 再把 track slot commit 到树
pool sizing / admission	`request_owned_slots` 区分 running slot 与 ping-pong track slots；auto sizing 预留 snapshot headroom；scheduler 在 free+evictable recurrent slots 不足时 backpressure，而不是让 `alloc_req_slots` 崩溃

生产路径不在 page 内提取 mid-forward state，而是保证每次可发布 snapshot 都来自 forward 的最终 recurrent state。--recurrent-track-interval 默认为 --chunked-prefill-size（没有 chunked prefill 时退回 --page-size），且必须是 page_size 的正倍数。

page_size：page=1 保留为 recurrent radix base/debug path；生产路径要求 --page-size > 1 且显式打开 --enable-recurrent-extra-buffer。当前 extra-buffer 不支持 speculative decoding、mixed chunked prefill，也不支持 GLA/Lightning（缺 masked recurrent track scatter）。FULL 的 paged path 仍复用已有 PagedTokenToKVPoolAllocator / paged matching；recurrent snapshot 按 track boundary 对齐发布。

4.6 typed-params + base-interface 重构

unified API 把分散的 evict/match/lock 签名收敛为一组 typed params：EvictParams(num_tokens, swa_num_tokens, dp_rank, recurrent_num)（upstream 字段名 mamba_num）、MatchPrefixParams / InsertParams / IncLockRefResult / DecLockRefParams / EvictResult；并把 MatchResult 从 4 字段扩到含 best_match_node（recurrent CoW 命中条件 §5.1、HiCache anchor 都依赖它，对齐 upstream base_prefix_cache.py:154 的 9 字段 NamedTuple）。

这是 base-interface 层的统一，不受 --enable-unified-radix-tree 门控；flag 只决定是否构造 UnifiedRadixCache。旧 RadixCache、SWARadixCache、ChunkCache 及调用方都按字段名访问 MatchResult，以保持 flag-off 路径稳定。

4.7 组件选择与 pool 装配

组件选择：在 registry 按模型类型拼出 tree_components——[FULL] 或 [FULL]+[SWA|REC]。新增组件只需在此多一个分支。

pool 装配：本 RFC 的 pool 装配沿用现状（_maybe_wrap_hybrid_kv_pool 等），不引入 HiCache host/device pool strategy。

允许组合矩阵（组件选择的权威约束）：

组合	允许	说明
`(FULL,)`	Yes	纯 full attention（MHA/GQA/MLA）
`(FULL, SWA)`	Planned	SWA 混合模型；当前仍走 legacy SWA cache
`(FULL, REC)`	Yes	linear-recurrent 模型，对应 upstream `(FULL, MAMBA)`
`(FULL, SWA, REC)`	No	upstream 结构性不可构造（`registry` 选择链 SWA 与 REC 互斥，只取其一）

5. 关键数据流

5.1 请求生命周期

各阶段 per-component 行为：

阶段	FULL	SWA	REC
① 命中条件	节点 value 非空	累计 len ≥ window	`best_match_node.component_data[REC].value` 非空（即该节点有可供 CoW 的 snapshot）
② Alloc	path-to-root 锁 + alloc_extend	path-to-root 至 window 边界	single-node CoW，alloc 1 slot（上游对应 `mamba_pool.alloc(1)`；本仓经 `HybridReqToTokenPool` free-list）
③ Forward	set_kv_buffer in-place（Pallas）	同 FULL	state.replace_buffer + donate（sglang-jax#1110）
④ Finish	value=kv_indices	在 window 边界拆节点，父节点 SWA `value` 清为 None	commit 入树为 snapshot；仅 abort/未缓存才 free

HiCache hook 触发点（本RFC 全为占位）：Arrive 命中 host-only 时 prefetch · Finish 时 write_through · Evict 时 write_back。真实数据搬运归 [[rfc_1_hicache]]。

5.2 cascade eviction

text

evict(EvictParams)
   ├─ 从低优先级到高优先级依次回收：REC(0) → SWA(1) → FULL(2)
   │     leaf 一律优先级 0，首先被回收
   ├─ 各 component 选 victim（LRU），跳过 tombstone 槽位
   └─ 跨组件按 component_data 下标遍历

同树的组件集合在装配期定死（合法组合见 §4.7），故 cascade 只需按固定优先级遍历存在的组件、跳过 value=None 槽位。

6. 验证面

验证面	覆盖内容
tree core	match / insert / split / evict / tombstone / `best_match_node` / dp-rank namespace
flag-off equivalence	`--enable-unified-radix-tree` 关闭时旧 `RadixCache` / `SWARadixCache` / `ChunkCache` 行为保持稳定
recurrent CoW	temporal+conv 全 state copy、CoW src lock、abort/retract、tree snapshot lifecycle
extra-buffer	track metadata round-trip、boundary split、masked track scatter、split equivalence、hit/flush determinism
sizing / admission	`request_owned_slots`、snapshot headroom、undersized fail-fast、over-capacity backpressure
DP routing	`cache_aware` affinity-vs-load policy、pending load accounting、DP namespace isolation
serving / accuracy / perf	cache-hit vs no-cache determinism、representative accuracy sanity、TTFT/throughput A/B

Nightly 或长跑验证继续覆盖 RSS / long-run leak、larger multi-host topology、以及新模型或新 recurrent_track_interval bucket。

7. 配置项

text

--enable-unified-radix-tree    # env SGLANG_JAX_ENABLE_UNIFIED_RADIX_TREE，默认 off
                               # 对齐 upstream environ.py 默认 False
--enable-recurrent-extra-buffer # recurrent production path: page_size > 1 的 track-slot snapshot
--recurrent-track-interval     # 默认 chunked_prefill_size；必须是 page_size 的正倍数
--dp-schedule-policy cache_aware # cache affinity + soft load balancing；默认仍是 min_running_queue
page_size                      # page=1 base/debug；page>=128 production path 需 extra-buffer
ComponentType                  # RECURRENT 为正式符号；移植 upstream 时 MAMBA→RECURRENT
eviction_priority              # FULL=2 > SWA=1 > REC=0（注意：不是 component_data 下标）

不引入：D + RadixCache opt-in flag / (FULL, SWA, REC) 三合一 / HiCache 搬运开关。

8. 剩余工作

项目	说明
SWA component	需要把 scheduler 中对 `SWARadixCache` / `ChunkCache` 的具体类判断改为能力判断，再接入 `(FULL, SWA)` 的 unified path
HiCache hooks	`host_value` / `host_lock_ref` 字段已预留；真实 D2H/H2D、host-only 命中、write-through / write-back 归 [[rfc_1_hicache]]
read-only prefix probe	`cache_aware` 当前用 `match_prefix` probe，会刷新 LRU 并可能 split partial node；后续应提供 `peek_prefix_len`
long-run validation	RSS / long-run leak、新模型 bucket、larger multi-host topology 继续放在 nightly 或专项验证

9. 未决问题

问题	默认处理
`cache_aware` probe 是否需要只读接口	是。新增 `peek_prefix_len`，避免 probe 刷新 LRU 或 split partial node
SWA 是否纳入当前 unified path	不阻塞 FULL / RECURRENT；待 scheduler SWA trait 化后接入
JAX / libtpu 升级后是否需要重新验 donated-buffer 路径	是。重跑 recurrent track scatter、split equivalence、serving determinism

10. 风险登记册

已被当前实现和验证覆盖的风险不在本表保留。

ID	风险	等级	缓解
R1	`cache_aware` probe 仍使用 `match_prefix`，会刷新 LRU 并可能 split partial node	LOW	增加 read-only `peek_prefix_len`；在此之前该行为只影响缓存统计/局部 LRU，不影响正确性
R2	JAX / libtpu 升级后，donated buffer + `optimization_barrier` 语义可能漂移	MED	升级后重跑 recurrent track scatter、split equivalence、serving determinism
R3	larger multi-host / `ep>1` 拓扑下 recurrent resume 可能只保持 bf16 数值等价，非 byte-identical	MED	保持 no-cache matched control；把 multi-host recurrent determinism / accuracy 放入 nightly
R4	长跑 RSS / host-side 引用环仍需覆盖更长 workload	LOW	保留 1h / 24h RSS 或 nightly leak gate
R5	SWA unified path 未接入前，SWA 仍依赖旧 cache path	LOW	scheduler trait 化后接入 `(FULL, SWA)` 并补 SWA component equivalence tests

文档设计

tpu-inference Q2 对照

sglang 代码分析

RFC-3: Unified Hybrid Radix Cache

1. 概述

1.1 范围内

1.2 范围外

2. 当前设计状态

3. 架构总览

4. 模块设计

4.1 `UnifiedRadixCache` + `UnifiedTreeNode`

4.2 `ComponentData` 槽位

4.3 `FullComponent`

4.4 `SWAComponent`（planned）

4.5 `RecurrentComponent`

4.5.1 设计：slot 角色、锁、后端选择

4.5.2 CoW device 侧克隆与 page_size

4.6 typed-params + base-interface 重构

4.7 组件选择与 pool 装配

5. 关键数据流

5.1 请求生命周期

5.2 cascade eviction

6. 验证面

7. 配置项

8. 剩余工作

9. 未决问题

10. 风险登记册

tpu-inference Q2 对照

RFC-3: Unified Hybrid Radix Cache ​

1. 概述 ​

1.1 范围内 ​

1.2 范围外 ​

2. 当前设计状态 ​

3. 架构总览 ​

4. 模块设计 ​

4.1 UnifiedRadixCache + UnifiedTreeNode ​

4.2 ComponentData 槽位 ​

4.3 FullComponent ​

4.4 SWAComponent（planned） ​

4.5 RecurrentComponent ​

4.5.1 设计：slot 角色、锁、后端选择 ​

4.5.2 CoW device 侧克隆与 page_size ​

4.6 typed-params + base-interface 重构 ​

4.7 组件选择与 pool 装配 ​

5. 关键数据流 ​

5.1 请求生命周期 ​

5.2 cascade eviction ​

6. 验证面 ​

7. 配置项 ​

8. 剩余工作 ​

9. 未决问题 ​

10. 风险登记册 ​

RFC-3: Unified Hybrid Radix Cache

1. 概述

1.1 范围内

1.2 范围外

2. 当前设计状态

3. 架构总览

4. 模块设计

4.1 `UnifiedRadixCache` + `UnifiedTreeNode`

4.2 `ComponentData` 槽位

4.3 `FullComponent`

4.4 `SWAComponent`（planned）

4.5 `RecurrentComponent`

4.5.1 设计：slot 角色、锁、后端选择

4.5.2 CoW device 侧克隆与 page_size

4.6 typed-params + base-interface 重构

4.7 组件选择与 pool 装配

5. 关键数据流

5.1 请求生命周期

5.2 cascade eviction

6. 验证面

7. 配置项

8. 剩余工作

9. 未决问题

10. 风险登记册