RFC-1: HiCache 多级 KV Cache

字段	值
作者	@john
日期	2026-05-20
状态	Draft
目标	在 sgl-jax 实装三级 KV cache（L1=HBM / L2=host pinned / L3=持久存储），以 UnifiedRadixCache 为入口，复用 RFC-0 builder
前置依赖	[[rfc_0_shared_cache_pd_infra]] M0 完成
关联 RFC	[[rfc_3_uhrc]]（UnifiedRadixCache + TreeComponent 重构在此细化：#1280 Stage 1 + Stage 5；本 RFC 的 Tree 入口依赖其交付）、[[rfc_2_pd_disaggregation]]（独立但同启时需协调 host memory）
关联 reference	[[sglang-tpu-inference-multi-level-cache]]

1. 概述

• L1（HBM）：现有 device-side KV pool，行为不变
• L2（host pinned）：新增 HostKVPool，基于 jax.device_put(..., memory_kind='pinned_host') 提供 async future API
• L3（持久存储）：新增 HiCacheStorage ABC，第一版实现 LocalFileStorage（page-granular + parent-hash chain）
• Tree 入口：移植 sglang UnifiedRadixCache + TreeComponent ABC，作为带 HiCache hooks 的 prefix tree（该重构的组件化设计在 [[rfc_3_uhrc]] 细化，本 RFC 复用其交付并叠加 HiCache hooks）
• Builder 集成：复用 RFC-0 kv_cache_builder 注册 unified_radix_cache

不引入：layer-wise overlap（XLA 不支持）、GPU 专用 storage backend（mooncake / NIXL 等）、替换 L1 device pool。

2. 背景与现状 gap

维度	sgl-jax 现状	sglang 已稳定	第一版缺什么
L1 device	KVCache + RadixCache	同上 + UnifiedRadixCache + HiCache hooks	UnifiedRadixCache port
L2 host	无	mooncake / 自有 host pool	HostKVPool（pinned + async）
L3 持久	无	mooncake / NIXL / aibrix 等 7+ backend	HiCacheStorage ABC + LocalFileStorage
Tree 入口	RadixCache（单注意力类型）	TreeComponent ABC（Full / SWA / Mamba 三组件全实装）	port Full + 后续 SWA/Mamba

UnifiedRadixCache 在 sglang origin/main HEAD f04c52253 已稳定（含 init_hicache() 内部装配 + 3 个 hook 全实装于三个组件）。详见 [[sglang-tpu-inference-multi-level-cache]] §1。

TPU 与 GPU 的本质差异：

维度	GPU	TPU
D2H 通路	cudaMemcpyAsync + stream	jax.device_put(..., memory_kind='pinned_host')（C++ 调用释放 GIL）
Layer-wise overlap	CUDA stream 与 kernel 重叠	XLA 静态编译，stream 不可控，无法 layer-wise overlap（改做 step-level）
GPU 专用库	mooncake / NIXL / aibrix / hf3fs / eic 等	全部不可用（依赖 NVLink / RDMA / CUDA）
Host pool 池化	cudaHostAlloc 自由分配	JAX 无 first-class pinned buffer pool API，需自建

控制面 ABC + hook 语义可与 sglang 对齐；L2/L3 实现路径需要为 TPU 重新设计。

L3 必要性：L2 host pool 受单机 host 内存上限制（典型 v6e 单机 ~1.5 TB），长上下文 / 多租户会驱逐。L3 是跨实例 / 跨重启命中率的长期保障。第一版用 LocalFileStorage 验证 ABC 与端到端通路，未来按需补 backend。

3. 架构总览

关键关系：

• UnifiedRadixCache 唯一 cache 入口（注册到 RFC-0 builder）
• 三个 TreeComponent（Full / SWA / Mamba）分别处理三种注意力类型，全部实装 3 个 hook（write_through / write_back / prefetch）
• HiCacheController 是 L1↔L2 控制器，唯一持有 D2H/H2D 调度权
• HiCacheStorage ABC 屏蔽 L3 实现细节

4. 模块设计

4.1 UnifiedRadixCache + TreeComponent ABC

职责：取代现有 RadixCache，作为带 HiCache hooks 的 prefix tree 入口。

接口	说明
match_prefix(tokens)	查 prefix，返回命中节点
insert(tokens, kv_indices)	插入新 prefix，触发 write_through hook
evict(node)	驱逐节点，触发 write_back hook
lock_ref(node) / unlock_ref(node)	引用计数

TreeComponent ABC（3 个子类）：

组件	注意力类型	sgl-jax 实装
FullComponent	Full attention（MHA / GQA / MLA）	H1
SWAComponent	Sliding-window attention	H3
RecurrentStateComponent	Mamba / KDA 等 recurrent state	H3

每个 component 实装 3 个 hook：

• write_through：节点 insert 时触发，把 kv_indices 加入 controller 写队列（不阻塞 forward）
• write_back：节点 evict 时触发，把节点对应 host 数据上交 L3（如启用）
• prefetch：lookup 命中但数据在 L2/L3 时触发，触发 H2D 拷贝并阻塞到完成

与现有模块复用：trie 结构 / 引用计数 / eviction 算法与 sglang UnifiedRadixCache 一致（与设备无关，无需重写）；device-side KV 操作调 sgl-jax KVCache 已有接口；HiCacheController 引用替换为本 RFC §4.2 实现。详见 [[sglang-tpu-inference-multi-level-cache]] §1/§2.1.6。

4.2 HiCacheController（L1↔L2 控制面）

职责：调度 device pool 与 host pool 之间的 D2H/H2D 拷贝，管理 host pool 内的 allocator 与 LRU 驱逐。

接口	输入 / 输出
write(host_indices, device_indices)	future（非阻塞）；paged-to-paged 一对一拷贝：device page i → host page j，按 index 配对，不需要 packed 中间形态（与 sglang HostKVCache 一致）
load(device_indices, host_indices)	future（非阻塞）；反向 paged-to-paged 一对一拷贝
evict_callback(host_indices)	None
get_cpu_copy / load_cpu_copy	仅诊断 / dump 用，非主路径（见关键约定）

依赖：device_pool + kv_sharding + device_allocator（sgl-jax 现有 KVCache / TokenToKVPoolAllocator）+ host_pool（§4.3）+ storage（可选，§4.4）+ tree_cache_evict_callback（由 UnifiedRadixCache 注入）。

关键约定：

• D2H/H2D 用新增 pinned-host async API，不复用 sgl-jax KVCache.get_cpu_copy / load_cpu_copy——这两者是同步 jax.device_get 到 numpy ndarray 再 jax.device_put（memory_pool.py:540, 1216），不是 pinned host、不是 async future，套到 HiCache 主路径会让 L2 退化为普通 CPU 内存。HiCache 主路径走 §4.3 HostKVPool 提供的 jax.device_put(x, NamedSharding(mesh, P(...), memory_kind='pinned_host'))，返回的 jax.Array 是 async future（可 .is_ready() 轮询）
• 不做 layer-wise overlap：XLA 不支持层级 stream 控制；改做 step-level overlap（forward N 期间后台跑 N-1 的 D2H）
• 驱逐：HostKVPool 用 LRU，命中 promote 到 LRU 头

详见 [[sglang-tpu-inference-multi-level-cache]] §2.1。

4.3 HostKVPool（CPU pinned 内存模型）

布局：paged（与 sglang HostKVCache 一致，memory_pool_host.py:155+，host buffer 也按 page_size 切片访问）；write/load 是 device page ↔ host page 一对一拷贝（按 index 配对），不需要 packed 中间形态。这跟 PD 传输的 paged → packed → paged 路径（RFC-2 §5.2）不同——HiCache L1↔L2 没有跨进程序列化需求，paged-to-paged 直接拷贝最简单且与 sglang HiCache 实现对齐。

前提：所有方案的底座都是

pinned_host_sharding = NamedSharding(mesh, P(...), memory_kind='pinned_host')
host_array = jax.device_put(device_data, pinned_host_sharding)
# host_array 是 async future jax.Array (非 numpy)；is_ready() 可轮询

JAX 没有 first-class API 暴露 pinned buffer pool——jax.device_put 每次返回新的 jax.Array，引用归零即被回收。GPU 那种"预分配大块、按 slot 写"的模式 host 端做不到。两种实现路径：

4.3.1 方案 A：per-block device_put（每次新建）

write(device_idx, host_idx):
    host_array = jax.device_put(device_pool.get(device_idx), pinned_host_sharding)
    host_slot_table[host_idx] = host_array  # 持引用避免回收

简单、与 JAX 主流用法一致；缺点：每 block 一份 jax.Array，slot 间无连续内存，驱逐只能 del 让 runtime 回收。

4.3.2 方案 B：预分配大 jax.Array + slot index 写入（接近 GPU 模式）

init: host_pool = jax.device_put(jnp.zeros([N_slots, ...]), pinned_host_sharding)
write: host_pool = host_pool.at[host_idx].set(device_data)  # functional update

布局可控、slot 表是 int 索引；缺点：.at[].set() 是 functional update，JAX 可能复制整个 buffer，需要 donate_argnames + optimization_barrier 强制 in-place，但对 host 内存语义未文档化保证。

4.3.3 第一版选择

优先方案 A——HiCache 不在延迟关键路径，方案 A 对小 host pool（<100 GB）影响可控；方案 B 的 in-place 保证依赖 JAX 内部细节，第一版避免。H1 实测后视情况升级 B。详见 [[sglang-tpu-inference-multi-level-cache]] §2.2.2。

模块边界（不受方案 A/B 影响）：

接口	说明
alloc(n_slots) / free(slot_indices)	分配 / 释放 slot
get(slot_indices) / put(slot_indices, data)	读写 slot
evict(n_slots)	LRU 驱逐 n 个 slot，返回被驱逐索引

4.4 HiCacheStorage ABC（L2↔L3 协议）

职责：屏蔽 L3 持久存储的实现细节。

接口	说明
get(keys) -> List[Optional[bytes]]	批量取 page
set(keys, values)	批量写
exists(keys) -> List[bool]	批量探测
delete(keys)	批量删除

Key 命名约定：KV page 依赖前缀上下文 + position——相同 token 序列在不同 prefix 下生成的 K/V 不同。sha256(tokens) 直接当 key 会让两条不同 prefix 撞 key 互相污染（correctness bug）。参考 sglang radix_cache.py:195 hash_page + :274 compute_node_hash_values，page hash 必须链上 parent hash + 多维度隔离：

page_hash = sha256(
    parent_hash || model_id || model_version || page_size
    || kv_dtype || layout_id || extra_key || tokens[start:end]
)

parent_hash 在 UnifiedRadixCache 节点 insert 时由 tree 提供（parent.hash_value[-1]），HiCacheStorage 不感知 tree 结构。目录结构：<root>/<model_id>/<page_size>/<sha256_prefix_2>/<sha256>.npy。

与 GPU 库的隔离：ABC 不引入 GPU 专用类型（cudaIpcMem、rdma_handle 等），所有数据交换走 bytes + 调用方负责 H2D；未来接入 GPU 库新增独立 ABC。

4.5 LocalFileStorage（L3 第一版实现）

维度	选择
粒度	page-granular（每 page 一文件；hook 触发粒度本就是 page，per-token 无信息增益）
原子性	写临时文件 + POSIX rename
序列化	numpy .npy（小、原生支持）
目录	<root>/<model_id>/<page_size>/<sha256_prefix_2>/<sha256>.npy
并发	不同 key 无冲突；同 key rename 兜底（后写覆盖）
容量管理	第一版无内置 LRU（语义在跨进程 / 跨实例下模糊），外部脚本清理；分布式 backend 统一在新 backend 层做

4.6 异步执行模型

L1↔L2 D2H/H2D 需要异步（forward 不阻塞）。两方案都基于 §4.3 pinned-host async API，区别在"谁发起 jax.device_put(..., pinned_host)"：

方案	实现	优点	缺点
A: ThreadPoolExecutor	controller.write_queue.submit(lambda: do_d2h(...))	简单；JAX C++ 释放 GIL，与后台 D2H 真正并行	引入多线程模型；与 sgl-jax single-thread + future 主体风格不一致
B: single-thread + future	主线程 jax.device_put(...) → 立即返回 future → 每 step is_ready() 轮询 pending list	与 sgl-jax 主体风格一致；与 scheduler refactor（如转 single-thread + future）天然对齐	is_ready() 是 JAX 私有 API（实验性）；每 step 末尾要轮询

第一版：两方案在 RFC 内并行保留，H1 PR 评审时根据 scheduler refactor 状态选定（refactor 已转 future → 倾向 B；否则 A；is_ready() 稳定性恶化也倾向 A）。详见 [[sglang-tpu-inference-multi-level-cache]] §2.2.4。

5. 关键数据流

5.1 write_through + write_back（forward 后 + evict 时）

forward(batch) → device KV (L1) → UnifiedRadixCache.insert
    │
    ├─ tree 节点 insert
    └─ HiCache hook: write_through → HiCacheController.write(host_idx, device_idx)
                                         │
                                         ├─ alloc host slot
                                         └─ enqueue D2H job (异步，不阻塞 forward)

[后台 / 下个 step 头] D2H 完成 → host_pool[host_idx] 就绪 → slot 可读

[节点 evict (LRU 触发)] → HiCache hook: write_back
    │
    ├─ storage.set(keys, host_pool.get(host_indices))   # L3 持久化（如启用）
    └─ host_pool.free(host_indices)

5.2 prefetch + L3 miss

match_prefix(tokens)
    │
    ├─ L1 命中（device）          → 直接返回 device_indices
    ├─ L2 命中（host）             → prefetch hook → HiCacheController.load → H2D（阻塞或同步等待）
    ├─ L3 命中（storage.exists）   → storage.get → host_pool.put → load
    └─ 全 miss                     → 走完整 forward → 后续 write_through 进 L1

6. 配置项

--hicache-storage {none, file}         # L3 backend，第一版仅 file 与 none
--hicache-ratio FLOAT                  # Host pool 大小 = ratio × device pool；默认 2.0
--hicache-storage-path PATH            # LocalFileStorage 根目录，默认 ~/.cache/sgl-jax/hicache
--hicache-write-policy {write_through, write_through_and_back}    # 不引入 selective (YAGNI)
--hicache-async-impl {threadpool, future}   # L1↔L2 异步方案（§4.6），默认值 H1 PR 评审定

不引入：layer-wise overlap 开关（XLA 不支持）、多 backend 同时启用、selective write。

7. 实施路线

阶段	交付	验收	依赖
H0	HostKVPool + allocator + 单测	alloc/free/get/put/evict 单测全通过；不接入 scheduler	RFC-0 M0
H1	HiCacheController + D2H/H2D 端到端 + UnifiedRadixCache port（FullComponent）	启用 --hicache-storage none 跑 Llama/Qwen，与现有 RadixCache KL 等价；forward latency 不增加 > 2%	H0。前置 UnifiedRadixCache + FullComponent + toggle flag 已交付（sglang-jax#1347，Stage 1, merged 2026-06-12）；SWA/RecurrentComponent 在 [[rfc_3_uhrc]] M2/M3
H2	LocalFileStorage + 重启复用验证	写 → 重启 → match_prefix 命中率达预期；page 文件可见可校验	H1
H3	SWAComponent + RecurrentStateComponent + tuning	MiMo-V2 / KDA 模型 KL 等价；不同 hicache-ratio 下命中率与 latency 测试报告	H1，与 [[rfc_2_pd_disaggregation]] D3 同期

H1 决策点：方案 A vs B（host pool 实现 + 异步执行）以 H1 PR 评审形式定。H3 与 RFC-2 D3 同期：D + HiCache 同启时的 HBM / host memory 切分需要两 RFC 实施者协同 benchmark。

8. 范围外 / 未来工作

项目	原因
Layer-wise H2D overlap	XLA 不支持
Mooncake / NIXL / aibrix / hf3fs / eic 等 storage backend	GPU 专用
write_through_selective 策略	YAGNI；H3 评估命中率后决定
LMCache 集成	与 HiCache 互斥的替代方案
跨进程共享 host pool	sglang 也没做
L3 LRU 驱逐 / 容量管理	第一版用外部脚本
DecodeKVCacheOffloadManager（D + L3）	与 RFC-2 同启时考虑
多 storage backend 同时启用	第一版单 backend
Pallas 自定义 D2H/H2D 算子	profile 后视情况

9. 未决问题

问题	实施时定
HostKVPool 是否在 H1 升级方案 B	默认 A，H1 实测后评估
L1↔L2 异步选 ThreadPool 还是 future	H1 PR 评审（依赖 scheduler refactor 状态）
--hicache-ratio 默认值 2.0	H3 benchmark 后定
LocalFileStorage 是否需要内置容量上限	H2（若外部脚本足够则不引入）
extra_key 在多 LoRA / 多模态场景的绑定方式	H2（key scheme §4.4 已定，仅调用方约定未定）
H3 是否纳入 DecodeKVCacheOffloadManager	视 [[rfc_2_pd_disaggregation]] D3 联合决策