RFC-0: sgl-jax 共享 Cache 与 PD 基础设施

字段	值
作者	@john
日期	2026-05-20
状态	Draft
目标	定义 HiCache（RFC-1）与 PD 分离（RFC-2）共同依赖的最小底座：cache 构建工厂、ChunkCache 一等公民地位
前置依赖	无
后续 RFC	[[rfc_1_hicache]]、[[rfc_2_pd_disaggregation]]、[[rfc_3_uhrc]]
关联 reference	[[sglang-tpu-inference-multi-level-cache]]、[[sglang-tpu-inference-pd-disaggregation]]

1. 概述

本 RFC 定义 RFC-1（HiCache）与 RFC-2（PD 分离）共同依赖的两项底层约定：

cache 构建工厂 kv_cache_builder：替换 scheduler 当前的 if-elif 选择逻辑，让新增 cache 类型不再侵入 scheduler 主流程。
ChunkCache 一等公民地位：作为 PD D 节点第一版默认 cache，显式列入配置候选。

这两项是 cache 子系统与 PD 子系统的最小公共底座。任何一项推迟，会让 RFC-1/2 在各自范围内重复实现。本 RFC 不引入新的 ABC，不改变 KV 内存模型，不影响推理正确性。

2. 背景与现状 gap

2.1 scheduler 当前 cache 选择

sgl-jax 当前 scheduler 在 _init_tree_cache 入口用 if-elif 选 cache 类型（RadixCache / ChunkCache / 其他）。这种结构对新增 cache 类型不友好——RFC-1 引入 UnifiedRadixCache、RFC-2 强化 ChunkCache 的使用场景，都会在这个入口堆叠分支。

sglang 已在 origin/main HEAD f04c52253 把同样的入口重构为 builder 模式（kv_cache_builder.build_kv_cache），详见 [[sglang-tpu-inference-multi-level-cache]] §4。

2.2 ChunkCache 在 sglang 与 sgl-jax 的当前地位

sglang 默认让 D 节点走 ChunkCache（不复用前缀）。原因：PD 场景下 P 已完成前缀计算并直接传输 KV 给 D，D 没有跨请求复用前缀的机会；prefix-tree 在 D 节点反而是一份纯成本（lookup 永远 miss、占内存）。

sgl-jax ChunkCache 实现已存在，但未作为"一等公民"对待：缺少配置入口、缺少与 PD 的对接说明、disaggregation_mode=decode 时没有默认值绑定。

2.3 为什么放在 RFC-0

两个 gap 单独看都不大，但 RFC-1 与 RFC-2 都会触碰：

RFC-1 注册 UnifiedRadixCache 需要工厂入口
RFC-2 D 节点启动需要明确 ChunkCache 是默认 cache，且不依赖 RFC-1 已完成

把两件事提取到 RFC-0，让 RFC-1 与 RFC-2 在 M0 完成后并行启动，互不阻塞。

3. 架构总览

要点：

kv_cache_builder 是 scheduler 唯一的 cache 创建入口。
注册表机制让 RFC-1 的 UnifiedRadixCache 注册不侵入 builder 本体。
ChunkCache 在 D 节点路径上是默认值。
HiCache（如启用）由 UnifiedRadixCache 内部装配，不出现在 builder 注册表里。

4. 模块设计

4.1 `kv_cache_builder`

职责：单一入口，根据 ServerArgs 与依赖项构造 cache 实例。

模块边界：

输入	输出
`ServerArgs`（含 `cache_type`、`disaggregation_mode`）	cache 实例（`BasePrefixCache` 子类）
依赖容器：token_allocator、device_pool、page_size、tp_rank、tp_size	—

选择逻辑：

优先读 ServerArgs.cache_type（显式配置）
未显式配置时，根据 disaggregation_mode + 模型类取默认值：
- disaggregation_mode=decode → chunk_cache
- 其他：按模型类分流
  - Full attention 模型（MHA / GQA / MLA）：[[rfc_1_hicache]] H1 落地后 unified_radix_cache；H1 前 radix_cache
  - SWA 模型：[[rfc_1_hicache]] H3 SWAComponent 落地后 unified_radix_cache；H3 前走现有 SWA cache 路径（不动）
  - Hybrid recurrent 模型（Mamba / KDA）：[[rfc_1_hicache]] H3 RecurrentStateComponent 落地后 unified_radix_cache；H3 前继续 disable_radix_cache 路径不变
调注册表中对应类型的构造函数

注册表机制：

builder 本体不知道具体 cache 类型，只知道注册表的 key
每个 cache 实现在自己的 module init 时调 register_kv_cache("radix_cache", RadixCacheCtor) 注册
新增 cache 类型不需要修改 builder 本体（RFC-1 加 UnifiedRadixCache 时就是新增一行注册）

与 RFC-1/2 的契约：

RFC	在 RFC-0 之上需要做什么
RFC-1	注册 `unified_radix_cache`，并把 HiCache 装配进 `UnifiedRadixCache` 的初始化路径
RFC-2	不引入新 cache 类型；启动 D 节点时只需依赖 RFC-0 已把 `chunk_cache` 设为 `decode` 模式的默认值

4.2 ChunkCache 升格为一等公民

改动内容：

在 cache_type 候选值列表中显式列入 chunk_cache
在 PD D 节点启动模式（disaggregation_mode=decode）下作为默认值
配置文档与代码注释明确："ChunkCache 不做前缀复用，专为 D 节点设计"
启动时拒绝 disaggregation_mode=decode + cache_type != chunk_cache 组合（builder 阶段直接报错，与 §6 "第一版不开放 override" 一致）

不改变的部分：

ChunkCache 的实现本身（已存在，行为不变）
非 PD 路径的 ChunkCache 使用（保留作为 debug / 特殊场景选项）

与 sglang 的对齐：

sglang D 节点同样默认 ChunkCache。
sglang 提供 opt-in flag 让 D 节点切到 RadixCache（用于 prefix locality routing 等高级场景）。sgl-jax 第一版不引入该 flag，未来如有需求按 sglang 模式补加。

5. 关键数据流

5.1 scheduler 启动期 cache 实例化时序

text

ServerArgs (cache_type, disaggregation_mode, ...)
    │
    ▼
kv_cache_builder.build(args, deps)
    │
    ├─ 解析 cache_type（显式或按 disaggregation_mode 取默认）
    │
    ├─ 查注册表 cache_type → ctor
    │
    ▼
RadixCache | ChunkCache | UnifiedRadixCache 实例
    │
    ▼
Scheduler 持有 tree_cache 引用

无并发、无重试、无 fallback——本 RFC 范围内的数据流就这一条。

5.2 PD + cache 组合矩阵（第一版）

disaggregation_mode	默认 cache_type	可选 override（第一版）
`null`（无 PD）	按模型类分流（见 §4.1）：Full attention H1 后 `unified_radix_cache`；SWA H3 后 `unified_radix_cache`；Hybrid recurrent H3 后 `unified_radix_cache`；各阶段前走原路径	`chunk_cache`
`prefill`	同上分流	`radix_cache` / `chunk_cache`
`decode`	`chunk_cache`	第一版不开放 override

按模型类分流的理由：[[rfc_1_hicache]] H1 只实装 FullComponent，SWA / RecurrentStateComponent 到 H3 才落地；过渡期 SWA / Hybrid recurrent 模型如果切到 unified_radix_cache 会缺组件失败，必须维持原有 cache 路径（SWA 现有 cache / disable_radix_cache）。

6. 配置项

新增配置：

text

--cache-type {radix_cache, chunk_cache, unified_radix_cache}
    Cache 类型显式选择。
    未指定时由 disaggregation_mode + 模型类取默认值：
      decode                                 → chunk_cache
      其他 + Full attention (MHA/GQA/MLA)    → H1 后 unified_radix_cache，H1 前 radix_cache
      其他 + SWA                             → H3 后 unified_radix_cache，H3 前现有 SWA cache 路径
      其他 + Hybrid recurrent (Mamba/KDA)    → H3 后 unified_radix_cache，H3 前 disable_radix_cache 路径不变

没有引入新的开关：

不引入 --enable-builder（builder 是唯一入口，无需开关）
不引入 D + RadixCache 的 opt-in flag（第一版不开放）

7. 实施路线

本 RFC 只有一个阶段 M0：

阶段	交付物	验收标准
M0	`kv_cache_builder` 工厂、注册表、ChunkCache 配置入口、scheduler 入口重构、配置文档	1) 现有 RadixCache 行为完全不变（回归 KL 等价）；2) `--cache-type chunk_cache` 可正常 serve；3) `--disaggregation-mode decode` 启动时自动选用 ChunkCache；4) 单测覆盖注册表注册与查询

前置依赖：无。

对后续 RFC 的约束：M0 必须完成，RFC-1 与 RFC-2 才能启动；M0 完成后 RFC-1 与 RFC-2 可并行。

8. 范围外 / 未来工作

项目	原因
`DecodeKVCacheOffloadManager`（D 节点把生成的 KV offload 到 L3）	RFC-1 + RFC-2 同启时的高级优化，跨 RFC，本 RFC 不涉及；纳入 RFC-1 H3 阶段考虑
跨进程共享 host pool	sglang 也未实现，复杂度高
D + UnifiedRadixCache 的 opt-in flag	sgl-jax 第一版不引入，未来按 sglang 模式补加
Cache 算法改动（trie 结构、eviction 策略）	不在本 RFC 范围
多种 storage backend 选择	属于 RFC-1 HiCacheStorage ABC 范围
多种 transfer backend 选择	属于 RFC-2 KVTransferEngine ABC 范围

9. 未决问题

本 RFC 范围紧凑，所有决策在写作时已能定下，无遗留未决问题。

如果实施时发现以下情况，需要追加讨论：

注册表机制是否需要支持运行时动态注册（vs 仅 import-time 注册）——目前假设 import-time 足够
ChunkCache 在 prefill 模式下是否需要作为 fallback——目前不引入，保持简单

文档设计

tpu-inference Q2 对照

sglang 代码分析

RFC-0: sgl-jax 共享 Cache 与 PD 基础设施

1. 概述

2. 背景与现状 gap

2.1 scheduler 当前 cache 选择

2.2 ChunkCache 在 sglang 与 sgl-jax 的当前地位

2.3 为什么放在 RFC-0

3. 架构总览

4. 模块设计

4.1 `kv_cache_builder`

4.2 ChunkCache 升格为一等公民

5. 关键数据流

5.1 scheduler 启动期 cache 实例化时序

5.2 PD + cache 组合矩阵（第一版）

6. 配置项

7. 实施路线

8. 范围外 / 未来工作

9. 未决问题

tpu-inference Q2 对照

RFC-0: sgl-jax 共享 Cache 与 PD 基础设施 ​

1. 概述 ​

2. 背景与现状 gap ​

2.1 scheduler 当前 cache 选择 ​

2.2 ChunkCache 在 sglang 与 sgl-jax 的当前地位 ​

2.3 为什么放在 RFC-0 ​

3. 架构总览 ​

4. 模块设计 ​

4.1 kv_cache_builder ​

4.2 ChunkCache 升格为一等公民 ​

5. 关键数据流 ​

5.1 scheduler 启动期 cache 实例化时序 ​

5.2 PD + cache 组合矩阵（第一版） ​

6. 配置项 ​

7. 实施路线 ​

8. 范围外 / 未来工作 ​

9. 未决问题 ​

RFC-0: sgl-jax 共享 Cache 与 PD 基础设施

1. 概述

2. 背景与现状 gap

2.1 scheduler 当前 cache 选择

2.2 ChunkCache 在 sglang 与 sgl-jax 的当前地位

2.3 为什么放在 RFC-0

3. 架构总览

4. 模块设计

4.1 `kv_cache_builder`

4.2 ChunkCache 升格为一等公民

5. 关键数据流

5.1 scheduler 启动期 cache 实例化时序

5.2 PD + cache 组合矩阵（第一版）

6. 配置项

7. 实施路线

8. 范围外 / 未来工作

9. 未决问题