Internal Wiki

主题

调研：sglang 与 tpu-inference 的多级 KV Cache ​

目的：为 sgl-jax 实现 HiCache（多级 KV 缓存）与 UnifiedRadixCache（统一前缀缓存）提供上游对照。控制面以 sglang 为主参考（成熟、接口稳定），数据面以 tpu-inference 的 JAX/Pallas API 为主参考（TPU 上的原语选型）。

代码基准：

  • sglang：origin/main HEAD f04c52253（2026-05-18）
  • tpu-inference：main HEAD（2026-05-18）

配套：本文专注「多级 KV Cache + UnifiedRadixCache」；PD 分离见 sglang-tpu-inference-pd-disaggregation.md；tpu-inference 单边的 PD 实现细节见 gap-pd-disaggregation.md。

摘要 / TL;DR ​

议题sglang 现状tpu-inference 现状
前缀缓存抽象BasePrefixCache ABC + 5 个实现（RadixCache / SWARadixCache / MambaRadixCache / ChunkCache / UnifiedRadixCache）无前缀缓存（vLLM v1 风格 paged KV 调度）
UnifiedRadixCacheorigin/main 1960 行，含 D↔H HiCache 完整集成；通过 SGLANG_ENABLE_UNIFIED_RADIX_TREE env var 启用—
TreeComponent ABC14 个 hook + 3 个 HiCache hook；Full/SWA/Mamba 三 component 全部实装—
HiCache L1↔L2 控制面HiCacheController（write_stream / load_stream / backup_thread / prefetch_thread）+ HostKVCache 4 变种无 HiCache 概念（offload_connector 只到 host）
HiCache L1↔L2 数据面CUDA stream + event；layer-wise overlapjax.device_put + memory_kind="pinned_host" + block_until_ready；无 layer-wise overlap（XLA 静态编译）
HiCache L2↔L3HiCacheStorage ABC + 7 个 backend（File / NIXL / Mooncake / HF3FS / Aibrix / EIC / SiMM）无 L3
Tree cache 创建kv_cache_builder.build_kv_cache()（318 行工厂）—
PD ↔ HiCache 主路径耦合70% 独立（UnifiedRadixCache 1960 行 0 处提及 PD）—
PD ↔ HiCache 共享底座30%（Mooncake TransferEngine 同进程可被 PD + L3 复用）—
GPU MPMD vs TPU SPMD每 TP rank 独立维护 host pool + radix tree，all_reduce(MIN) 同步队列SPMD（单 scheduler + thread-based worker）

GPU/TPU 关键差异、HiCache 控制面与数据面的 sgl-jax 适配方向见 [[rfc_1_hicache]]（路线由 RFC 决定）。

零、术语与全景 ​

0.1 术语 ​

术语含义
L1 / DeviceTPU HBM 上的 paged KV cache（计算时直接读写）
L2 / HostHost DRAM 上的 KV 镜像（pinned_host memory_kind）
L3 / Storage持久化存储（本地文件 / Mooncake / NIXL / HF3FS 等）
HiCache"Hierarchical Cache" — L1/L2/L3 的统一调度抽象
UnifiedRadixCachesglang 新一代前缀缓存（统一 RadixCache + SWA + Mamba + HiCache）
TreeComponentUnifiedRadixCache 的可插拔组件（Full / SWA / Mamba）
Prefix cache / Tree cache基于 token prefix 的 KV 复用结构（trie / radix tree）
write_through / write_backD→H 写策略：同步立即写 vs 延迟批量写
load_backH→D 从 L2 加载回 L1
prefetch_from_storageL3→L2 异步预取

0.2 模块全景 ​

text
                            ┌─────────────────────────────────────┐
                            │  Scheduler (event loop)             │
                            │  仅入口:                             │
                            │   init → kv_cache_builder.build()   │
                            │   step → tree_cache.check_*_events  │
                            │   step → tree_cache.flush_acks      │
                            └─────────────────┬───────────────────┘
                                              │
                              ┌───────────────▼───────────────┐
                              │ kv_cache_builder              │
                              │ build_kv_cache(args)          │
                              │ ┌───────────────────────────┐ │
                              │ │ 8-way dispatch:            │ │
                              │ │ ChunkCache / RadixCache /  │ │
                              │ │ HiRadixCache / SWA / ...   │ │
                              │ │ UnifiedRadixCache ★        │ │
                              │ └───────────────────────────┘ │
                              └─────────────────┬─────────────┘
                                                │
        ┌───────────────────────────────────────┴─────────────────────┐
        │                                                             │
   ┌────▼──────────┐                              ┌──────────────────▼──────┐
   │ ChunkCache    │                              │ UnifiedRadixCache       │
   │ (无 prefix    │                              │ (port from sglang)      │
   │  复用)         │                              │  • match_prefix          │
   └───────────────┘                              │  • insert / evict        │
                                                  │  • cache_*_req           │
                                                  │  • init_load_back        │
                                                  │  • check_hicache_events  │
                                                  │  • flush_write_through_  │
                                                  │      acks                │
                                                  │  • ready_to_load_host_   │
                                                  │      cache               │
                                                  │                          │
                                                  │  tree_components: list ──┤
                                                  └────────┬─────────────────┘
                                                           │
                                  ┌────────────────────────▼───────────────────┐
                                  │ TreeComponent ABC                          │
                                  │  • 14 个核心 hook                           │
                                  │  • 3 个 HiCache hook                        │
                                  │    - build_hicache_transfers                │
                                  │    - commit_hicache_transfer                │
                                  │    - drive_host_eviction                    │
                                  └──┬──────────────┬────────────┬─────────────┘
                                     │              │            │
                              ┌──────▼─────┐ ┌──────▼─────┐ ┌────▼────────────┐
                              │ Full       │ │ SWA        │ │ Mamba           │
                              │ Component  │ │ Component  │ │ Component       │
                              └────────────┘ └────────────┘ └─────────────────┘

  ─────── 共享底座（被 cache / HiCache / PD 不同程度复用） ────────────

  ┌────────────────────┐   ┌──────────────────────┐   ┌──────────────────┐
  │ MemoryPools (pytree)│  │ HostKVPool            │   │ HostMemoryAllocator│
  │ • req_to_token_pool │  │ (memory_pool_host.py) │   │                   │
  │ • kv_pool (HBM)     │  │ • pinned_host buffer  │   └──────────────────┘
  │ • [NEW] host_pool   │  │ • LRU / Queue 两风格  │
  └────────────────────┘   └──────────────────────┘
                              ┌──────────────────────────┐
                              │ HiCacheController         │
                              │ • write_stream (D→H)      │
                              │ • load_stream (H→D)       │
                              │ • backup_thread (→ L3)    │
                              │ • prefetch_thread (← L3)  │
                              └──────────────────────────┘
                              ┌──────────────────────────┐
                              │ HiCacheStorage ABC (L3)   │
                              │ • File / NIXL / Mooncake  │
                              │ • HF3FS / Aibrix / EIC    │
                              │ • SiMM / LMCache          │
                              └──────────────────────────┘

一、UnifiedRadixCache（前缀缓存基础） ​

1.1 sglang 控制面 ​

1.1.1 BasePrefixCache ABC ​

python/sglang/srt/mem_cache/base_prefix_cache.py:196（含 HiCache hooks 默认签名）：

python
class BasePrefixCache(ABC, PrefixCacheTrait):
    # ─── 核心 (所有 cache 必须 override) ───
    def match_prefix(self, params: MatchPrefixParams) -> MatchResult: ...
    def cache_finished_req(self, req, **kwargs): ...
    def cache_unfinished_req(self, req, chunked, **kwargs): ...
    def insert(self, params: InsertParams) -> InsertResult: ...
    def evict(self, params: EvictParams) -> EvictResult: ...
    def inc_lock_ref(self, node) -> IncLockRefResult: ...
    def dec_lock_ref(self, node, params): ...

    # ─── HiCache hooks (L273-302, 默认 raise NotImplementedError) ───
    def init_load_back(self, params: InitLoadBackParams): ...       # L273
    def ready_to_load_host_cache(self): ...                         # L282
    def flush_write_through_acks(self): ...                         # L288
    def check_hicache_events(self): ...                             # L296
    def take_events(self): ...                                      # L302

    # ─── 能力探测 (L305-336) ───
    def supports_swa(self) -> bool: ...                             # L305
    def supports_mamba(self) -> bool: ...                           # L308
    def supports_streaming_session(self) -> bool: ...               # L311
    def is_chunk_cache(self) -> bool: ...                           # L332
    def is_tree_cache(self) -> bool: ...                            # L335

关联 dataclass（base_prefix_cache.py L38-192）：

类型行号说明
MatchPrefixParamsL39匹配请求参数
InsertParams / InsertResultL50 / L69插入参数与结果
EvictParamsL77驱逐参数
InitLoadBackParamsL136HiCache load_back 参数
MatchResult (NamedTuple)L145匹配结果（含 last_host_node L173 / best_match_node L174 / host_hit_length L175，为 HiCache 预留）

1.1.2 UnifiedRadixCache 关键方法 ​

python/sglang/srt/mem_cache/unified_radix_cache.py:201 class UnifiedRadixCache(BasePrefixCache) — 平行于 RadixCache 和 HiRadixCache，通过组合 tree_components: list[TreeComponent] 支持多形态。

方法行号签名 / 用途
__init__202(self, params: CacheInitParams)
init_hicache298(self, server_args, params) -> None（启用 D↔H HiCache）
register_sidecar_pool338(self, spec: SidecarPoolSpec) -> None
match_prefix341(self, params: MatchPrefixParams) -> MatchResult
insert368(self, params: InsertParams) -> InsertResult
evict384(self, params: EvictParams) -> EvictResult
evict_host1035host pool 驱逐
write_backup1160(self, node, write_back: bool = False) -> int（D→H）
load_back1223H→D
writing_check1368(self, write_back: bool = False) -> None
loading_check1415(self) -> None
init_load_back1432(self, params: InitLoadBackParams) -> tuple[torch.Tensor, UnifiedTreeNode]
check_hicache_events1478内部调 writing_check + loading_check
flush_write_through_acks1483scheduler 每 step 调一次
ready_to_load_host_cache1487(self) -> int（剩余可加载 token 数）

MatchResult 字段（base_prefix_cache.py L145-175）：

text
device_indices          # L1 命中的 token 索引
last_device_node        # L1 命中链终端
last_host_node          # L173: L2 命中链终端 (HiCache)
best_match_node         # L174: L1+L2 综合最佳匹配点 (load_back 锚点)
host_hit_length         # L175: L2 命中长度
mamba_branching_seqlen  # Mamba 分支起点
cache_protected_len     # 不可驱逐区间

1.1.3 TreeComponent ABC + HiCache hooks ​

python/sglang/srt/mem_cache/unified_cache_components/tree_component.py（364 行）：

python
class TreeComponent(ABC):
    # ─── 14 个核心 hook ───
    #   create_match_validator / build_match_segment /
    #   resolve_match / commit_match / redistribute_on_node_split /
    #   acquire_component_lock / release_component_lock /
    #   evict_component / drive_eviction / ...

    # ─── 3 个 HiCache hook (L342-358) ───
    def build_hicache_transfers(node, phase, **kw) -> Optional[list[PoolTransfer]]:
        """根据 CacheTransferPhase 构造从 host 或 device 拉取的 PoolTransfer 描述符"""

    def commit_hicache_transfer(node, phase, transfers=()) -> None:
        """传输完成后回填 device_indices / host_indices 到 ComponentData"""

    def drive_host_eviction(num_tokens, tracker) -> None:
        """host pool 驱逐时由该 component 决定丢弃哪些 host 节点"""

CacheTransferPhase 四态：

Phase触发当前是否实装
BACKUP_HOSTD→H 写回✅ 全部 component 实装
LOAD_BACKH→D 加载✅ 全部 component 实装
BACKUP_STORAGEH→L3❌ UnifiedRadixCache 未实装（L3 集成缺失）
PREFETCHL3→H 预取❌ 未实装

1.1.4 三个 component 的 HiCache hook 实装摘要 ​

full_component.py（282 行）— 最简：FULL KV 走主流程

Hook行号行为
drive_host_eviction142heap 出 evictable host leaf 调 _evict_host_leaf
build_hicache_transfers213BACKUP 阶段返回 None（主流程处理）；LOAD_BACK 沿 evicted 链上溯收集 host_value，生成单条 PoolTransfer(name=KV)
commit_hicache_transfer253device_indices 切片写回各 node cd.value + 更新 evictable leaf 集合

swa_component.py（537 行）— 双池 + 滑动窗口

Hook行号行为
build_hicache_transfers427BACKUP 返回 PoolName.SWA；LOAD_BACK 在 sliding_window_size 窗口内只收 host-only 节点
commit_hicache_transfer482调 _restore_device_value + allocator.set_full_to_swa_mapping，重建 full↔swa 映射
drive_host_eviction517host LRU 驱动，区分 leaf（_evict_host_leaf）vs internal（tombstone + cascade）

mamba_component.py（448 行）— 单节点 + CoW

Hook行号行为
build_hicache_transfers341BACKUP 返回 PoolName.MAMBA；LOAD_BACK 支持单节点 restore + per-request CoW（按需 mamba_pool.alloc(1)，OOM 触发 evict(mamba_num=1)）
commit_hicache_transfer397写回 cd.value，host LRU → device LRU 迁移
drive_host_eviction426同 SWA 模式，host LRU 驱动 tombstone + cascade

1.1.5 互斥关系矩阵 ​

组合是否共存说明
UnifiedRadixCache ↔ RadixCache互斥（builder 二选一）SGLANG_ENABLE_UNIFIED_RADIX_TREE env var
UnifiedRadixCache ↔ HiRadixCache互斥（builder 二选一）UnifiedRadixCache 优先级更高，内部已含 HiCache
UnifiedRadixCache ↔ HiCache不互斥（内部已实装 init_hicache）通过 init_hicache(server_args, params) 启用
HiRadixCache ↔ LMCache互斥dispatch 二选一
任何 tree_cache ↔ StreamingSession共存（装饰器）—
任何 tree_cache ↔ PD共存（独立维度）详见 §五

1.2 tpu-inference 视角 ​

1.2.1 tpu-inference 没有 UnifiedRadixCache ​

tpu-inference 走 vLLM v1 调度路径，KV cache 是无前缀树的 paged pool：

python
# tpu-inference 的 KV cache 标准 sharding
device_sharding = NamedSharding(
    mesh,
    PartitionSpec(None, None, "model"),     # 仅 KV head 维度沿 TP 轴切分
    memory_kind="device",
)
# Tensor shape: [num_blocks, block_size, num_head, 2, head_dim]
# 实际分片：dim[2] 沿 mesh axis "model"

KV cache 容器与调度逻辑分离：tpu-inference 的 kv_cache_manager.py 只管 paged 分配/释放（参考 vLLM v1 风格），不维护 prefix tree。

1.2.2 sgl-jax 现状对照 ​

组件sgl-jax 路径状态
BasePrefixCachepython/sgl_jax/srt/mem_cache/base_prefix_cache.py:31✅ 已存在，含 HiCache 钩子签名
MatchResult (含 last_host_node)同上 L12✅ 已有 host_hit_length / last_host_node 字段
RadixCacheradix_cache.py:145✅ 已有
SWARadixCacheswa_radix_cache.py:281✅ 已有（双 LRUList）
ChunkCachechunk_cache.py:15✅ 已有
UnifiedRadixCache❌待 port
TreeComponent ABC + 3 component❌待 port
MemoryPools (pytree)memory_pool.py:1360✅ 已有，可加 host pool 同级条目
KV pool shardingmemory_pool.py:410 NamedSharding(mesh, P("data", None, "tensor", None, None))✅ 已有；只需加 memory_kind="pinned_host" 即可派生 host 版本

1.3 GPU MPMD vs TPU SPMD 对前缀缓存的影响 ​

维度sglang / GPU（MPMD）sgl-jax / TPU（SPMD）
进程模型每个 TP rank 独立进程单 scheduler + thread-based worker
Radix tree 维护每 rank 独立维护一棵 tree单 tree（scheduler 内）
HiCache 队列同步all_reduce(op=MIN) 跨 rank 同步队列消费数量不需要跨进程同步
决策广播每 rank 独立决策（隐式一致）单点决策，worker 收 block_id 列表执行
优势容错好、CPU 并行决策控制面简单、状态唯一
劣势实现复杂、SPMD 不变量难维持单点瓶颈风险（但 scheduler 决策本就轻量）

关键事实：sgl-jax 当前的 scheduler 单进程架构已经天然兼容 SPMD，HiCache port 时不需要引入跨 rank 同步代码。

二、HiCache L1 ↔ L2（HBM ↔ Host DRAM） ​

2.1 sglang 控制面 ​

2.1.1 HiCacheController ​

python/sglang/srt/managers/cache_controller.py：HiCache 的 I/O 调度器，管理 CUDA stream 和后台线程。

python
class HiCacheController:
    # ─── CUDA streams (GPU 上独立 stream，避免与 forward 抢资源) ───
    write_stream         # 专用 D2H 写 stream
    load_stream          # 专用 H2D 加载 stream

    # ─── 后台线程 ───
    backup_thread        # 异步写入 L3 storage
    prefetch_thread      # 异步从 L3 预取到 L2

    # ─── 关键方法 ───
    write(device_indices, node_id) -> host_indices       # 分配 host 内存 + 入队
    start_writing()                                       # 发起 D2H DMA
    load(host_indices, node_id) -> device_indices         # 分配 device 内存 + 入队
    start_loading() -> producer_id                        # 发起 H2D DMA (per-layer)

2.1.2 HostKVCache 与 4 个变种 ​

python/sglang/srt/mem_cache/memory_pool_host.py：

text
HostKVCache (ABC)
    │
    ├── MHATokenToKVPoolHost
    ├── MLATokenToKVPoolHost
    ├── MambaPoolHost
    └── NSAIndexerPoolHost

HostPoolGroup / PoolEntry  ← 多 pool 统一管理（HiCache v2 用）

2.1.3 KVCache D2H/H2D 接口 ​

python/sglang/srt/mem_cache/memory_pool.py:668：

python
class KVCache(abc.ABC):
    def get_key_buffer(self, layer_id) -> Tensor: ...       # L730
    def get_value_buffer(self, layer_id) -> Tensor: ...     # L734
    def get_kv_buffer(self, layer_id) -> Tuple: ...         # L738
    def set_kv_buffer(self, ...): ...                       # L742

    # ─── HiCache 复用 ───
    def get_cpu_copy(self, indices) -> Tensor: ...          # D→H 实际拷贝
    def load_cpu_copy(self, cpu_tensor, indices): ...       # H→D 实际拷贝

    # ─── HiCache layer-wise overlap 用 ───
    def register_layer_transfer_counter(self, counter): ... # L751

2.1.4 写入策略 ​

策略行为默认
write_throughfinished_req 命中点立即同步 D→H✅ 默认
write_backfinished_req 累积后批量 D→H（延迟降低，命中略增）opt-in
write_through_selective仅写 hit prefix 部分opt-in，复杂

2.1.5 Scheduler 集成点 ​

时机方法scheduler.py 行号
初始化创建 HiRadixCache / HiMambaRadixCache（通过 kv_cache_builder.build_kv_cache）见 §四
请求入队tree_cache.prefetch_from_storage()L2035
主循环每 steptree_cache.check_hicache_events()L2442
Batch 准备tree_cache.ready_to_load_host_cache()L2627
处理 acktree_cache.flush_write_through_acks()L2708
L3 attach（动态）tree_cache.attach_storage_backend(...)L3233（前置探测 L3226 hasattr）

HTTP 控制面：/clear_hicache_storage_backend、HiRadixCache.attach_storage_backend() / detach_storage_backend() 支持运行时挂载后端。

2.1.6 HiCache 事件钩子（hooks）的语义 ​

"事件钩子"是指 BasePrefixCache 上一组默认 raise NotImplementedError、由 HiCache 实现类（如 HiRadixCache / UnifiedRadixCache.init_hicache 后）填充的方法。它们的本质是 scheduler 主循环与 HiCache 的异步事件回调约定——scheduler 不直接调用 HiCacheController 的内部线程或队列，而是周期性调用这几个 hook 让 cache 把"有什么完成了 / 有什么准备好了"反馈回主循环：

Hook触发时机（scheduler 主循环侧）语义
prefetch_from_storage(req)请求入队时提示 cache："给这个请求的 prefix 启动 L3→L2 预取"（异步派发，不阻塞）
check_hicache_events()每个 step 开头让 cache 推进内部状态机：消化已完成的 D→H / H→D 传输、释放队列槽位、推进 LRU
ready_to_load_host_cache()batch 准备阶段询问 cache："还能再触发多少 token 的 H→D 加载"（返回剩余容量，scheduler 据此决定本 step 加载量）
flush_write_through_acks()每个 step 末尾让 cache 处理 write_through 的 ACK——把已成功落盘到 host 的节点标为"可被驱逐 device"、可继续 evict
init_load_back(params)命中需要 H→D 时由 cache 分配 device 侧 KV slot + 启动 H→D，返回 (device_indices, anchor_node)供 scheduler 后续 attention 用
take_events()监控/调试需要时取出累积的 HiCache 事件（命中率、传输延迟等）供 metrics 上报

为什么是"钩子"而不是直接调内部线程：

  • HiCache 内部使用后台线程 + CUDA stream 异步执行 D2H/H2D，scheduler 不能直接 await 这些操作
  • 钩子提供一个"轮询 + 推进"的接口：scheduler 在主循环固定时机调用，cache 在钩子内同步处理"上次派发但已完成"的事 → 不需要 scheduler 自己管理线程
  • 钩子签名稳定（在 BasePrefixCache 上）→ 即使具体 cache 实现换（HiRadixCache → UnifiedRadixCache），scheduler 一行不用改

2.2 tpu-inference 数据面（JAX API 选型） ​

2.2.1 D2H 标准路径：jax.device_put + pinned_host ​

python
# offload/tpu_offload_connector.py:1763-1770
# host_sharding 在 register_runner 时构造（line 1274-1317）：
host_sharding = NamedSharding(
    mesh,
    P(None, None, "model"),        # 只在 model (TP) 轴 partition
    memory_kind="pinned_host",     # ← 关键：告诉 JAX 放到 pinned host memory
)

# D2H 调用
chunks_on_cpu = []
for i in range(total_num_blocks):
    chunks_on_cpu.append(
        jax.device_put(flat_kv_caches_tpu[i], host_sharding)
    )
jax.block_until_ready(chunks_on_cpu)

机制：JAX runtime 知道每个 device 的 host affinity，自动把 shard 路由到对应 host 的 pinned CPU 内存。不需要显式多节点协调代码。

2.2.2 pinned_host 内存的生命周期与回收 ​

memory_kind="pinned_host" 在 JAX/XLA 中映射到 PJRT 的 pinned host memory space —— OS 层面是 mlock 锁页内存（参考 cudaMallocHost 之于 GPU）。理解其生命周期对 HiCache L2 容器设计至关重要：

  1. 分配：jax.device_put(arr, host_sharding) 由 PJRT 调用底层（TPU runtime / CUDA runtime / CPU mock）的 pinned allocator 申请一块连续 host 内存，并返回一个新的 jax.Array。该 jax.Array 是 Python 层对底层 buffer 的 handle。
  2. 保活：底层 pinned buffer 的释放完全由 jax.Array 的 Python 引用计数驱动 —— 只要 jax.Array 还在内存中（被任何 list / dict / 局部变量持有），底层 pinned buffer 就不会归还给 OS。
  3. 释放路径：del arr 或所有引用消失 → Python GC 触发 jax.Array.__del__ → PJRT 把 buffer 标记为 free → buffer 进入 XLA runtime 内部的 buffer pool（不是立即 munlock） → 后续相同 shape/sharding 的分配优先复用 → pool 收缩时才真正归还给 OS。
  4. 与 block_until_ready 的关系：jax.device_put 异步派发，返回的 jax.Array 在传输完成前是"尚未就绪"的 handle。HiCache 必须在 block_until_ready 之后才能把 handle 放入 L2 容器，否则后续读到的可能是空 buffer。
  5. 多 host 视角：每个 host 上分配的 pinned 内存只对应本 host 的 device shard；跨 host 检索 L2 时不能直接拷贝 jax.Array，要经 sharding 路由。

HiCache 容器实现要点：

  • L2 容器（如 sglang 的 OrderedDict[hash, jax.Array] / MHATokenToKVPoolHost）应是 owns-jax.Array 的，而不是把 array 字段拆出来存
  • LRU 淘汰时必须 del 掉 jax.Array 引用（或从 OrderedDict 中弹出），否则 pinned memory 永不释放
  • 用 psutil/cgroups 观察实际 RSS 与 PJRT pool size，可能与 Python 持有的 jax.Array 总和有时间差（pool 不立即收缩）
  • pinned 容量在 OS 层有上限（默认 ulimit -l，容器场景需要显式提权），超出会导致 device_put 失败 —— 必须在 host pool 容量配置层做硬上限
  • 不要依赖 GC 时机驱逐 —— 显式管理生命周期（pop() / del），不要写 weakref 类的隐式回收

两种 host pool 组织模式：

JAX 上构造 host pinned 内存有两种风格，差异主要在"谁负责池化"：

模式分配粒度池化责任方代表实现
A. per-block device_put每个 KV block 一次 device_put → 一个独立 jax.ArrayPJRT 内部 buffer pool（黑盒）tpu-inference LocalCPUBackend（OrderedDict[chunk_id, jax.Array]）
B. 预分配大 Array + slot 索引启动时一次 device_put(zeros(num_slots, *block_shape)) → 一个大 jax.Array，Python 侧维护 free_list: list[int]应用层（自管 slot allocator）sglang MHATokenToKVPoolHost（GPU 路径用 torch.empty(pin_memory=True) + index）的 JAX 对应

模式 B 在 JAX 中的写入需要 jit kernel（slot 索引参与计算）：

python
# 启动时一次性预分配 L2 整池
self.host_pool = jax.device_put(
    jnp.zeros((num_slots, *block_shape), dtype),
    host_sharding,                              # memory_kind="pinned_host"
)
self.free_list = list(range(num_slots))         # Python 侧 slot allocator

# D2H 写入指定 slot（不再 device_put，而是 jit 的 dynamic_update_slice）
@jax.jit
def write_slot(host_pool, kv_block, slot_id):
    return jax.lax.dynamic_update_slice(
        host_pool, kv_block, (slot_id, 0, 0, 0)
    )

self.host_pool = write_slot(self.host_pool, kv_block, slot_id)

两种模式的权衡：

维度A. per-block device_putB. 预分配大 Array + slot
实现复杂度低中（需要 jit slot writer + 分桶）
分配延迟可预测性受 PJRT pool 状态影响（黑盒）启动时一次开销，运行时仅 free_list.pop()
碎片表现PJRT pool 内部碎片不可控定长 slot，零碎片
L2 容量上限受 ulimit -l + PJRT pool 收缩策略影响严格 = num_slots × block_size
jit 编译成本0写入路径必须 jit；num_slots / block_shape 变化要分桶
donate_argnames 需求不需要必须 donate=True 才能 in-place 更新整池，配合 optimization_barrier（详见 §2.2.5）
sharding 灵活性每个 Array 独立 sharding整池 sharding spec 必须与 device KV pool 兼容
与 sglang MHATokenToKVPoolHost 语义对齐✗✓

当前 JAX API 的限制（事实）：

JAX/PJRT 当前没有第三种模式（"从已分配的 host 指针零拷贝构造 TPU pinned_host jax.Array"）：

  • jax.dlpack.from_dlpack 只映射 kDLCPU / kDLCUDA / kDLROCM 三种 device，TPU 不在内；即使走 numpy + mlock + kDLCPU 旁路，拿到的也是 JAX CPU backend 的 Array，不是 TPU pinned_host memory space，到 TPU 仍要重拷贝
  • jax.make_array_from_single_device_arrays 要求输入已是 jax.Array，不能从原始 host 指针构造
  • jax.experimental.transfer 名字易误导，实际是 DCN 跨切片传输（87 行），与 host pool 无关
  • PJRT Python 层未暴露 HostBufferAllocator / register_pinned_buffer / BufferFromHostBuffer(ZeroCopy) 接口；pinned_host 在 IR 中以 custom_call @annotate_device_placement {_xla_buffer_placement = "pinned_host"} 形式存在，allocator 由 XLA TPU runtime 自管

所以"预分配 + 内部 slab 管理"在 JAX 上只能通过模式 B（预分配大 Array + slot index）实现，不能像 sglang GPU 路径那样直接传 raw pinned 指针给底层拷贝原语。

2.2.3 H2D 路径：jax.device_put + 后续 D2D scatter ​

python
# offload/tpu_offload_connector.py:2077-2082
raw_chunked_kv_on_tpu = []
for i in range(num_blocks_to_load):
    raw_chunked_kv_on_tpu.append(
        jax.device_put(assembled_kv_on_cpu[i], device_sharding)
    )
jax.block_until_ready(raw_chunked_kv_on_tpu)

完整 H2D 流程：

text
LocalCPUBackend (OrderedDict[chunk_id, jax.Array])
    ↓ Python dict lookup (无数据搬运)
jax.Array (pinned_host)
    ↓ jax.device_put(..., device_sharding)     ← H2D 实际传输
jax.Array (device HBM, "staging")
    ↓ multi_layer_copy (Pallas DMA)            ← Scatter 到 paged KV cache
jax.Array (device HBM, paged pool 内具体 blocks)

注意 H2D 后还需要一次 D2D scatter（multi_layer_copy）把数据写入 paged KV cache 的指定 block 位置。

2.2.4 异步执行：JAX 默认派发 + 可选线程模型 ​

JAX 的 device_put / JIT / Pallas kernel 都是异步派发：调用立即返回 handle，实际计算在 PJRT 后端进行，直到 block_until_ready() 才真正同步。这是 JAX 与 CUDA 的根本差异——不需要显式 stream/event 管理。

基于此，HiCache 的 D2H/H2D 可以有两种实现风格，选择哪一种取决于 scheduler 的线程模型：

风格 A：scheduler 自己是多线程 / ThreadPoolExecutor 后台 worker（tpu-inference 当前的做法）

python
self.save_executor = ThreadPoolExecutor(max_workers=4)

# 提交后台任务（不阻塞主线程）
future = self.save_executor.submit(self._async_transfer_task, args...)

def _async_transfer_task(self, ...):
    chunks_on_cpu = jax.device_put(..., host_sharding)
    jax.block_until_ready(chunks_on_cpu)   # 仅阻塞后台线程
    self.cpu_backend.add(chunk_id, chunks_on_cpu)

依赖 GIL 行为：

操作持有 GIL？
jax.device_put() 派发短暂持有，C++ 层释放
jax.block_until_ready() 等待释放 GIL
JIT 编译后的 model forward释放 GIL
Python dict 操作持有 GIL

主线程 forward 与后台线程 D2H 可以真正并行（两边都在 GIL 之外）。sglang 走的是类似路径（CUDA stream + Python 后台线程）。

风格 B：scheduler 单线程 + 显式管理 jax.Array 的 ready 状态（"future-style"）

由于 jax.device_put 返回的 jax.Array 自身就是一个"未来值"——可用 arr.is_ready() 非阻塞查询完成、用 block_until_ready 显式同步——单线程 scheduler 也能在 event loop 中编排异步 D2H：

python
# 风格 B 伪代码：单线程 scheduler 编排
pending_handles: list[tuple[ChunkId, jax.Array]] = []

def step():
    # ... model forward ...
    # 派发新的 D2H（立即返回，未完成）
    h = jax.device_put(kv_to_offload, host_sharding)
    pending_handles.append((chunk_id, h))

    # 检查上一轮派发的 handle 是否就绪，就绪的入 L2
    ready, still_pending = [], []
    for cid, h in pending_handles:
        if h.is_ready():
            ready.append((cid, h))
        else:
            still_pending.append((cid, h))
    for cid, h in ready:
        self.cpu_backend.add(cid, h)
    pending_handles = still_pending

风格 B 的优势：避免线程同步复杂性、与 SPMD 模型天然契合、与 sgl-jax 中正在进行的 scheduler 单线程化 refactor 方向一致。

风格选择影响：

维度风格 A（ThreadPoolExecutor）风格 B（future API）
与 scheduler 线程模型需要 scheduler 容忍多线程与单线程 scheduler 兼容
编程复杂度需要处理线程同步、生命周期event loop 内直接编排
并行度主线程 forward + 后台 D2H 真正并行主线程派发 + JAX runtime 后台执行（同样真正并行）
与 sglang 对齐✓（HiCacheController 走多线程）✗（与 sglang HiCacheController 风格不同）
调试多线程，需要看 thread state单线程，event loop 易追踪

sgl-jax 适配点：sgl-jax 当前 scheduler 处于单线程化 refactor 阶段；最终选风格 A 还是 B 应结合 refactor 终态决定，本文不做硬性推荐。两种风格都能用 JAX 原语实现，关键是底层异步派发 + block_until_ready/is_ready 语义不变。

2.2.5 donate_argnames + optimization_barrier ​

HiCache 的 KV gather 场景必须配套使用，否则 XLA 可能把 gather 重排到 buffer 被覆盖之后：

python
@functools.partial(
    jax.jit,
    static_argnames=['num_blocks'],
    donate_argnames=('kv_caches',),       # 关键：允许 XLA 复用输入 buffer
)
def stack_kv_cache_cross_layers(kv_caches, block_ids, num_blocks):
    def _gather_blocks(layer_kv_cache):
        return layer_kv_cache.at[block_ids].get()

    gathered_kv_layers = jax.tree.map(_gather_blocks, kv_caches)
    stacked_blocks = jnp.stack(gathered_kv_layers, axis=1)
    split_blocks = jnp.split(stacked_blocks, num_blocks, axis=0)
    kv_caches = jax.lax.optimization_barrier(kv_caches)   # 关键：防止 XLA 重排
    return kv_caches, split_blocks

2.2.6 Pallas copy_to_host 作为高性能替代（PD 用） ​

python
# distributed/kv_transfer.py:416 (函数签名简化)
def copy_to_host(src: jax.Array, dest: jax.Array, ...) -> jax.Array:
    # Pallas kernel 内部：
    # async_copy = pltpu.make_async_copy(src_ref, dest_ref_in_pltpu.HOST, semaphore_ref)
    # async_copy.start(); async_copy.wait()

特点：

  • Pallas kernel 直接操作 DMA 引擎，逐 layer 拷贝
  • 目标是预分配的 pinned host buffer（HostKVPool），避免每次传输的 MapDmaBuffer 系统调用开销
  • 配合 is_ready() 轮询用作传输完成检查（非 block_until_ready）

为什么 PD 用 Pallas、HiCache offload 用 jax.device_put：

维度HiCache (offload)PD
是否在请求关键路径❌ 后台异步✅ 请求关键路径
延迟敏感中高
实现复杂度低（device_put 简单稳定）高（Pallas + 预分配 + is_ready 轮询）
建议第一版jax.device_putcopy_to_host（或 D2H staging 模式参考）

2.2.7 Pallas DMA 用于 D2D scatter（multi_layer_copy） ​

python
def multi_layer_copy(
    *,
    src_array: list[jax.Array],
    dest_array: list[jax.Array],
    src_offsets: jax.Array,
    dest_offsets: jax.Array,
    chunk_sizes: jax.Array,
    num_chunks: jax.Array | None = None,
    mesh: Mesh | None = None,
    src_sharding_spec: P | None = None,
    dest_sharding_spec: P | None = None,
    replicated_sharding_spec: P | None = None,
)

实现层级：

text
multi_layer_copy                  (kv_transfer.py:296)
  └─ _async_copy_jit              (kv_transfer.py:187, JIT 入口)
       └─ jax.shard_map 内部:
            ├─ _start_chunked_copy_kernel      ← Pallas kernel
            │    └─ pltpu.make_async_copy(...).start()
            │
            └─ _wait_for_chunked_copy_kernel   ← Pallas kernel
                 └─ pltpu.make_async_copy(...).wait()

pltpu.SemaphoreType.DMA 用于同步；多个 make_async_copy 可以重叠执行——这是 KV scatter 性能的关键。

2.2.8 pinned_host vs unpinned_host ​

memory_kind物理位置DMA 友好何时用
"device"TPU HBM✓默认计算时
"pinned_host"OS 锁页内存（不可 swap）✓HiCache L2 / PD host buffer
"unpinned_host"可 swap 的普通 host memory✗offload 容量优先场景（TPU_OFFLOAD_USE_UNPINNED_HOST=true）

坑点：unpinned_host 在 swap 时性能急剧下降；KV 容量远小于 host RAM 总量时，永远用 pinned_host。

2.3 关键差异 ​

2.3.1 CUDA stream/event vs JAX 默认异步 ​

SGLang (CUDA)JAX 等价
cudaStream_t (write_stream / load_stream)不需要——JAX 默认异步派发
cudaEvent_t.record()不需要——jax.Array 隐含完成事件
cudaEventQuery()arr.is_ready()（Pallas async copy 产物）或 future.done()
cudaEventSynchronize()jax.block_until_ready(arr)
cudaMemcpyAsync(D2H)jax.device_put(arr, host_sharding)
cudaMemcpyAsync(H2D)jax.device_put(arr, device_sharding)
<<<grid, block>>> kernel launchjit_fn(args)
自定义 CUDA kernelPallas kernel (@pl.kernel + pltpu.make_async_copy 等)
Pinned memory (cudaMallocHost)memory_kind="pinned_host"
cudaStreamWaitEvent (跨 stream 同步)不需要——JAX 自动管理依赖

2.3.2 Layer-wise overlap：TPU 上不可行 ​

sglang HiCache 的核心性能优化之一是 layer-wise H2D overlap：传输完 layer N 立即开始 layer N 的 forward 计算，与 layer N+1 的传输重叠。

TPU/XLA 下不可行，原因：

  1. XLA 静态编译：jax.jit 把整个 model forward 编译成一个完整 HLO 图，中间无法插 host 同步点
  2. 无 CUDA stream 等价物：TPU 没有多 stream 概念，DMA 引擎和 MXU 调度由 XLA 统一管理
  3. 改造代价极高：把 jit forward 拆成 per-layer 小函数会破坏算子融合、内存规划、编译时间爆炸

替代方案：

  • Step-level overlap：request A 的 forward 与 request B 的 H2D 并行
  • Request-level overlap：批量预先加载满足"可加载"条件的请求

三、HiCache L2 ↔ L3（Host DRAM ↔ Persistent Storage） ​

3.1 sglang 控制面 ​

3.1.1 HiCacheStorage ABC ​

python/sglang/srt/mem_cache/hicache_storage.py:132：

python
class HiCacheStorage(ABC):
    def register_mem_pool_host(self, host): ...                     # L139
    def register_mem_host_pool_v2(self, host_pool, pool_name): ...  # L142

    # ─── v2 API（多 pool, 用 PoolTransfer 描述符）★ 新代码用这套 ───
    def batch_exists_v2(self, ...): ...                             # L147
    def batch_get_v2(self, transfers: list[PoolTransfer]): ...      # L180
    def batch_set_v2(self, transfers: list[PoolTransfer]): ...      # L191

    # ─── v1 API（单 pool）★ 兼容性保留 ───
    def batch_get_v1(self, keys, host_indices): ...                 # L202
    def batch_set_v1(self, keys, host_indices): ...                 # L214

    # ─── 单条 / 通用 ───
    def get(key) / batch_get(keys) / set / batch_set                # L227-282
    def exists / batch_exists / clear / get_stats                   # L283-308

配套类（hicache_storage.py L21-128）：

类行号用途
HiCacheStorageConfigL21backend 配置
HiCacheStorageExtraInfoL38额外元数据
PrefetchTimeoutConfigL44预取超时配置
PoolName(Enum)L52KV / SWA / MAMBA / ...
PoolHitPolicy(Enum)L72命中策略
PoolTransferL84传输描述符（key + host_indices + pool_name）
SidecarPoolSpecL102sidecar pool 规格
PoolTransferResultL111传输结果

具体 backend 实装：HiCacheFile 在 L311（同文件）。

3.1.2 7 个 backend + LMCache 集成 ​

text
HiCacheStorage (ABC)  ← hicache_storage.py:98
    │
    ├── HiCacheFile          (local disk)
    ├── HiCacheNixl          (NVIDIA NIXL plugin)
    ├── MooncakeStore        (Mooncake)
    ├── HiCacheHF3FS         (3FS)
    ├── AibrixKVCacheStorage
    ├── EICStorage
    └── HiCacheSiMM

StorageBackendFactory  ← storage/backend_factory.py
    └── 懒加载注册表 + dynamic 动态加载

LMCache 是替代方案：通过 LMCRadixCache 走另一条 dispatch 路径，不实现 HiCacheStorage ABC。

3.1.3 UnifiedRadixCache 当前 L3 集成缺失 ​

当前事实：UnifiedRadixCache 只覆盖 L1↔L2 (D↔H)。prefetch_from_storage / attach_storage_backend / detach_storage_backend 仍只在 HiRadixCache / HiMambaRadixCache 中存在。

Scheduler 通过 hasattr(self.tree_cache, "attach_storage_backend") 做能力探测（python/sglang/srt/managers/scheduler.py:3226，命中后 :3233 调用）。Issue #20415 中「L3 support」复选框未勾。

对 sgl-jax 的意义：L3 prefetch 是 upstream 的真正窗口期，但短期不在 sgl-jax 关键路径——第一版 HiCache 可只做 L1↔L2。

3.2 tpu-inference 视角 ​

tpu-inference 当前没有 L3 实现。offload_connector.py 只到 host：

  • LocalCPUBackend (OrderedDict[chunk_id, jax.Array]) 作为 L2 容器
  • jax.Array 本身持有 pinned host memory 引用，无需额外内存管理

3.3 TPU 的 L3 backend 现实约束 ​

BackendTPU 可用性原因
File（local disk）✅纯 Python file I/O，与硬件无关
Mooncake⚠️ 受限主要依赖 RDMA，TPU 集群一般无 RDMA
NIXL⚠️ 受限NVIDIA-specific
HF3FS⚠️ 受限3FS 部署成本高
Aibrix❓ 待验证设计可能跨硬件
EIC❓ 待验证—
SiMM❌GPU shared memory，TPU 无对应物
LMCache❓ 待验证—

事实：上述 7 个 backend 中，只有 File 不依赖 GPU/RDMA 软件栈；其余 6 个 backend 在 TPU 集群下要么不可用、要么需要可观的额外部署成本。这是把 HiCacheStorage ABC 与具体 backend 解耦的客观依据。

四、Tree Cache 创建：Builder 模式 ​

4.0 什么是 Builder 模式 / 解决了什么问题 ​

sglang 支持的 BasePrefixCache 子类越来越多（RadixCache / SWARadixCache / MambaRadixCache / ChunkCache / SWAChunkCache / RadixCacheCpp / UnifiedRadixCache / HiRadixCache / HiMambaRadixCache / LMCRadixCache，共 10 种）。选哪一种取决于多个正交开关：

  • disable_radix_cache（关闭 prefix 复用，走 ChunkCache）
  • enable_hierarchical_cache（启用 HiCache D↔H）
  • is_hybrid_swa / is_hybrid_ssm（混合架构）
  • enable_lmcache（LMCache 替代方案）
  • 实验性 env var：SGLANG_EXPERIMENTAL_CPP_RADIX_TREE / SGLANG_ENABLE_UNIFIED_RADIX_TREE

历史上，这些 dispatch 逻辑直接散落在 Scheduler.__init__ 内部（一长串 if-elif），随着 cache 实现增多，scheduler 主类不断膨胀且测试困难。

sglang 当前的解法是把这套 dispatch 整体迁移到一个独立的 builder 函数：

  • python/sglang/srt/mem_cache/kv_cache_builder.py 中导出 build_kv_cache(params, server_args, ...) -> KVCacheBuildResult
  • Scheduler 只需要一行 result = build_kv_cache(...) + self.tree_cache = result.tree_cache
  • Builder 内部封装：根据 ServerArgs 选 cache 类、构造对应 component / host pool / hicache controller、返回打包结果

简单说：builder 模式 = "在 scheduler 之外、由一个专门函数负责'根据配置选一种 tree_cache 子类并完成所有依赖装配'"。它解决的不是性能问题，而是控制面代码组织问题——避免 scheduler 主类承担越来越多的 cache 类型 dispatch 责任。

对 sgl-jax 的意义：sgl-jax 当前的 scheduler 中也有类似的 cache 选择逻辑（RadixCache / SWARadixCache / ChunkCache）；如果未来要加 UnifiedRadixCache / HiCache，dispatch 矩阵会快速膨胀。提前引入 builder 模式可以让 scheduler 主类保持精简。

4.1 sglang kv_cache_builder.build_kv_cache() ​

python/sglang/srt/mem_cache/kv_cache_builder.py（318 行）。入口：build_kv_cache() 行 132；返回 KVCacheBuildResult（L12 定义）。

Dispatch 优先级（基于 builder 中的实际控制流）：

text
1. (disable_radix_cache & chunked_prefill) → ChunkCache / SWAChunkCache       (L230/L234)
2. SGLANG_EXPERIMENTAL_CPP_RADIX_TREE      → RadixCacheCpp                    (L241)
3. SGLANG_ENABLE_UNIFIED_RADIX_TREE        → UnifiedRadixCache                (L256) ★
   └ 按 hybrid 类型注入 tree_components(FULL / +SWA / +MAMBA)
   └ 若 enable_hierarchical_cache → tree_cache.init_hicache(server_args, params)
4. enable_hierarchical_cache               → HiMambaRadixCache / HiRadixCache (L268/L272)
5. is_hybrid_swa                           → SWARadixCache                    (L279)
6. is_hybrid_ssm                           → MambaRadixCache                  (L283)
7. enable_lmcache                          → LMCRadixCache                    (L289)
8. default                                 → RadixCache

关键事实：

  • UnifiedRadixCache 优先级高于 HiRadixCache——开了 SGLANG_ENABLE_UNIFIED_RADIX_TREE，HiCache 走 Unified 路径
  • UnifiedRadixCache 仍默认关闭（env var gated），与旧实现平行存在
  • Issue #20415 的「移除其他实现」目标尚未达成

4.2 Scheduler 主类不再做内部 dispatch ​

python
# scheduler.py 中 init_memory_pool_and_cache 区
result = kv_cache_builder.build_kv_cache(...)
self.tree_cache = result.tree_cache

scheduler.py 中事件循环 dispatch（独立于 tree_cache 选择）：

python
if disaggregation_mode == NULL:
    if enable_pdmux:        event_loop_pdmux()
    elif pp_size > 1:       event_loop_pp()
    elif enable_overlap:    event_loop_overlap()
    else:                   event_loop_normal()
elif disaggregation_mode == PREFILL:
    if pp_size > 1:         event_loop_pp_disagg_prefill()
    elif enable_overlap:    event_loop_overlap_disagg_prefill()
    else:                   event_loop_normal_disagg_prefill()
elif disaggregation_mode == DECODE:
    ...

意义：tree_cache 选择与 PD mode 选择正交——sgl-jax port builder 时不需要把 PD 逻辑混进 builder。

4.3 tpu-inference 没有对应物 ​

tpu-inference 调度器走 vLLM v1 路径，KV cache 创建直接由 kv_cache_manager.py 完成，无前缀树工厂。

五、PD ↔ HiCache 的真实耦合（30% 共享 / 70% 独立） ​

直觉上「P→D 的 KV 传输」与「HiCache L3 存储」是同构操作——P 把 KV "存"到某处，D "取"回。但代码层的耦合度只成立 30%。

5.1 共享的（30%）— 底层传输引擎 ​

Mooncake（同进程可共享 TransferEngine 实例）：

路径文件用的 mooncake API
PDdisaggregation/mooncake/conn.py:184 MooncakeKVManagerget_mooncake_transfer_engine() → MooncakeTransferEngine（distributed/device_communicators/mooncake_transfer_engine.py:93），走 transfer_sync_*（点对点 RDMA）
HiCache L3mem_cache/storage/mooncake_store/mooncake_store.py:297 MooncakeStore (导入 MooncakeDistributedStore 在 L252)from mooncake.store import MooncakeDistributedStore，K/V Object Store API (put/get)
复用机制mooncake_store.py:362-380显式尝试 get_mooncake_transfer_engine() 复用底层 TransferEngine 实例

NIXL（各自独立 agent）：

  • PD：disaggregation/nixl/conn.py:216 NixlKVManager，在 L226-248 创建 nixl_agent
  • HiCache L3：mem_cache/storage/nixl/hicache_nixl.py:35 HiCacheNixl，在 L75-77 创建独立的 nixl_agent，名字以 hicache_nixl_ 开头
  • 不共享，仅同包不同实例

结论：一个进程里 PD + HiCache-mooncake 可以共享底层传输引擎，但走的是两个不同的上层抽象（P2P RDMA vs K/V Store）。

5.2 完全独立的（70%）— 上层协议路径 ​

UnifiedRadixCache 不知道 PD 存在：

  • unified_radix_cache.py 1960 行
  • grep -E "disagg|Disagg|kv_send|kv_recv|KVSender|KVReceiver|bootstrap|PD" 0 命中

PD 完全绕过 tree_cache 做 KV 传输：

python
# prefill.py:755+ (send_kv_chunk)
kv_indices = (
    self.req_to_token_pool.req_to_token[req.req_pool_idx, start_idx:end_idx]
    .cpu()
    .numpy()
)
# ... 选 KV blocks 后
req.disagg_kv_sender.send(page_indices, state_indices)

直接读 req_to_token_pool、调 KVSender.send，与 tree 完全解耦。

D 端接收 KV 后不调 tree_cache.insert：

  • DecodeTransferQueue._commit_transfer_to_req (decode.py:1383) 只写 metadata_buffers
  • tree 插入由 D 完成 decode 后通过 cache_finished_req 走标准路径

5.3 唯一现存的 PD ↔ HiCache 融合实践 ​

disaggregation/decode_kvcache_offload_manager.py 是当前唯一把 PD 与 HiCache L3 真正打通的代码：

维度内容
位置python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
启用disaggregation_decode_enable_offload_kvcache=True（要求 disaggregation_mode=="decode" + hicache_storage_backend != None）
行为D 侧实例化自己的 MHATokenToKVPoolHost + HiCacheController + storage_backend，把 D 自己生成的增量 KV 写到 L3（offload_kv_cache L109-180）
语义是 "D offload 到 L3"，不是 "D 从 L3 拉取 P 写入的 KV"
不改变P→D 的初始 prefix KV 仍然走 MooncakeKVManager 的 RDMA 直传，与 L3 无关

5.4 协议层抽象差异（为什么独立不是缺陷） ​

维度PD 协议L3 协议
寻址方式bootstrap 房间号 + req_id（点对点路由）内容哈希（content-addressed）
同步模型同步点对点（P 等 D 来 pull）异步 put/get
生命周期请求级（一次性）prefix 级（长期，为 future hit）
数据形态完整 prefill KV（一次大块发送）按 page 切分（细粒度寻址）
触发时机prefill 完成立即发起由 write-through/write-back 策略决定

核心结论：PD 和 L3 解耦不是技术缺陷，是协议层抽象不同。复用底层传输引擎可行，但上层协议必须分开。

5.5 CLI 同时启用 PD + HiCache 的兼容性 ​

  • server_args.py:3424 _handle_hicache 只规范化 layout/IO
  • _handle_cache_compatibility 仅检查 enable_hierarchical_cache ⊥ disable_radix_cache
  • 没有禁止"PD + HiCache 同启"的检查
  • 真正的耦合 flag 是 disaggregation_decode_enable_offload_kvcache：强制要求 disaggregation_mode == "decode" 且 hicache_storage_backend is not None

参考文件索引 ​

sglang（origin/main HEAD f04c52253） ​

文件行数内容
mem_cache/base_prefix_cache.py—ABC + dataclass + HiCache hook 签名
mem_cache/unified_radix_cache.py1960UnifiedRadixCache + D↔H HiCache 完整集成
mem_cache/unified_cache_components/tree_component.py364TreeComponent ABC
mem_cache/unified_cache_components/full_component.py282Full component（标准 MHA）
mem_cache/unified_cache_components/swa_component.py537SWA component（滑动窗口）
mem_cache/unified_cache_components/mamba_component.py448Mamba component（Recurrent state）
mem_cache/hicache_storage.py—HiCacheStorage ABC + v1/v2 API + dataclass
mem_cache/memory_pool.py—KVCache ABC + get_cpu_copy / load_cpu_copy
mem_cache/memory_pool_host.py—HostKVCache ABC + 4 个变种 + HostPoolGroup
mem_cache/storage/*—7 个 L3 backend + lmcache 集成
mem_cache/kv_cache_builder.py318build_kv_cache() 工厂
managers/cache_controller.py—HiCacheController（write/load stream + 后台线程）
managers/scheduler.py—HiCache 集成点（L834/L2050/L2447/L2628/L2681 等）
disaggregation/decode_kvcache_offload_manager.py—唯一 PD↔HiCache 融合点

tpu-inference（main HEAD 2026-05-18） ​

文件内容
offload/tpu_offload_connector.pyD2H/H2D 主路径（jax.device_put + pinned_host）
offload/utils.pystack_kv_cache_cross_layers（gather + donate_argnames + optimization_barrier）
distributed/kv_transfer.pymulti_layer_copy（Pallas DMA + 跨 layer 流水线）+ copy_to_host
runner/kv_cache_manager.pypaged KV cache 管理 + transfer_kv_cache（跨 mesh）

配套文档 ​