SGLang HiCache 写路径：GPU KV 如何写入主机内存和外部存储

HiCache 写路径回答的是：一段已经在 GPU 上算出的前缀 KV 缓存，何时、如何离开 GPU，并在离开后仍然可以复用。

这个问题的难点在于，写路径里有两次“成功”：一次是 GPU KV 备份到主机内存，另一次是主机缓存页写入存储后端。只看其中一次，就会误判后续请求能否实现远端命中。

源码版本与范围

在 SGLang v0.5.14 中，主线入口在：

python/sglang/srt/mem_cache/unified_radix_cache.py
  write_backup
  write_backup_storage

python/sglang/srt/mem_cache/hicache_storage.py
  PoolTransfer
  HiCacheStorage.batch_set_v1/v2

python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py
  MooncakeStore.batch_set_v1/v2

关键调用链：

write_backup(node)
  -> 收集 GPU device_indices
  -> 构造 PoolTransfer
  -> cache_controller.write(...)
  -> commit_hicache_transfer(...)
  -> 节点获得主机内存备份
  -> write_backup_storage(...)
  -> HiCacheStorage.batch_set_v1/v2
  -> MooncakeStore.batch_set_v1/v2
  -> 按缓存页汇总 Store 对象的写入结果

本文只分析 HiCache 写入路径和 Mooncake 后端的页与对象边界，不覆盖完整淘汰策略、GPU 计算内核和分布式部署基准测试，也不声称某种写入策略在所有负载下都更优。

写路径分两段

不要把 HiCache 写入理解为一次 set(key, value)。它至少分成两段：

D->H 备份:
  GPU KV 池 -> 主机 KV 池

H->Storage 备份:
  主机 KV 池 -> 存储后端

第一段让 GPU KV 可以被驱逐而不丢失，第二段则允许释放主机 KV，并在后续请求中从远端存储重新预取。

GPU 到主机内存：write_backup

write_backup(node) 的输入是一个基数树节点。这个节点对应一段前缀，其 GPU KV 索引保存在基础组件的 value 中。

写入主机内存前，代码会先维护一个重要不变量：

非回写场景下，父节点必须先完成备份。

原因是前缀缓存是一棵树。如果子节点已经写入主机内存，而父节点没有，后续只保留子节点会破坏前缀的连续可恢复性。因此，write_backup 会在必要时递归备份父节点。

核心流程可以压成：

write_backup(node)
  -> 获取节点的 GPU device_indices
  -> 构造 PoolTransfer(name=KV, device_indices=...)
  -> 为 Mamba/SWA/DSA 等组件构造额外传输项
  -> 主机内存不足时调用 evict_host
  -> cache_controller.write(...)
  -> 得到 host_indices
  -> commit_hicache_transfer(...)
  -> 标记节点已经有主机内存备份

这里的 cache_controller.write 才是真正执行 D->H 复制的地方。write_backup 更多是在组织传输描述、分配主机内存和提交节点状态。

PoolTransfer 是跨层描述符

PoolTransfer 同时服务 D->H、H->D、H->Storage 和 Storage->H。它不是 Mooncake 专属结构，而是 HiCache 用来描述“某个内存池中的一批页如何搬运”的统一对象。

写入主机内存时，KV 传输主要关心：

name=PoolName.KV
device_indices=<GPU token 索引>

写入存储层时，同一个逻辑会变成：

name=PoolName.KV
host_indices=<主机 token/页面索引>
keys=<page key>

这就是 PoolTransfer 同时包含 device_indices、host_indices 和 keys 的原因：它需要跨越多个缓存层级。

写穿和写回的区别

HiCache 中的写穿与回写很容易混淆。

写穿更主动。一个节点被命中到一定次数，或者策略明确要求写穿时，HiCache 会提前把它备份到主机内存，甚至继续写入存储层。这样，后续从 GPU 驱逐时的成本更低。

回写更像驱逐前的保护。GPU 空间不足、需要驱逐某个节点，而该节点还没有主机内存备份时，先执行 write_backup(node, write_back=True)，再回收 GPU 槽位。

区别可以这样看：

写穿：
  缓存仍在 GPU，但提前复制一份到主机内存或存储层

写回：
  为了释放 GPU，先保存到主机内存，再移除 GPU 引用

两者都会让 KV 离开纯 GPU 状态，但触发时机不同。

主机内存到外部存储：write_backup_storage

当节点已经有主机内存备份，并且启用了存储后端时，write_backup_storage(node) 会把主机缓存页写入 L3。

核心流程：

write_backup_storage(node)
  -> 检查 enable_storage、cache_controller、node.backuped
  -> 可选生成 prefix_keys
  -> 为额外组件构造 BACKUP_STORAGE 传输项
  -> 构造 KV PoolTransfer(host_indices=node.host_value, keys=node.hash_value)
  -> cache_controller.write_storage(...)
  -> 记录 ongoing_backup，并持有主机内存锁

注意，这里写入存储层的键是 node.hash_value，而不是原始 token 列表。HiCache 会按页计算哈希键，让 L3 存储可以按页查询和复用。

为什么要持有主机内存锁

存储层写入是异步的。write_backup_storage 发起后，主机 KV 缓存页不能立刻回收，否则后端可能仍在读取这块主机缓冲区。

因此，代码会在 ongoing_backup 中记录操作编号，并通过主机内存锁保护节点：

ongoing_backup[operation_id] = (node, host_lock_params)

锁要等存储写入完成后再释放。这个机制可以防止一种常见错误：上层认为写入已经发起便复用主机缓冲区，最终导致远端保存了错误内容。

v1 和 v2 写入

普通 KV 路径走 v1：

batch_set_v1(keys, host_indices)

混合缓存走 v2：

batch_set_v2(transfers)

v2 的意义在于，多个内存池必须共同构成可复用前缀。例如，某些模型除了 KV 还需要辅助状态；如果 KV 写入成功而辅助内存池没有写入，下一次前缀命中也无法完整复用。

因此，v2 会按内存池返回结果：

{
  "kv": [true, true, false],
  "mamba": [true, true],
  "swa": [true]
}

最终可用的前缀长度，要取各内存池共同支持的命中边界。

接到 Mooncake 时发生什么

当后端为 Mooncake 时，batch_set_v1 会根据 page key 生成 Mooncake 对象键。

MHA 通常是：

page_key
  -> page_key_<tp_or_head_suffix>_k
  -> page_key_<tp_or_head_suffix>_v

MLA 通常是：

page_key
  -> page_key_<tp_or_rank_suffix>_k

然后 Mooncake 后端从主机内存池取得指针和长度：

mem_pool_host.get_page_buffer_meta(host_indices)

最后调用零拷贝写入接口：

batch_put_from(...)
batch_put_from_multi_buffers(...)

所以写入 Mooncake 的不是 Python 字节串，也不是 GPU 张量，而是一组主机 KV 页指针。

写路径的失败边界

写路径可能在几个位置失败：

• 主机内存不足，执行 evict_host 后仍无法分配。
• D->H 传输失败，cache_controller.write 返回 None。
• 存储后端写入失败，对应缓存页的结果为 false。
• 混合缓存中的某个辅助内存池写入失败，导致整体可用前缀缩短。
• 主机内存锁或节点拆分处理不当，可能造成写入中的节点状态不一致。

HiCache 通过锁、进行中的操作记录和每个缓存页的写入结果，尽量把失败限制在可回退状态，而不是直接把基数树标记为“全部可用”。

这条路径能证明什么

结论：写入成功必须按页提交

HiCache 写入路径最容易误读的地方，是把“某个缓冲区已经传输完成”当成“这个前缀以后可以复用”。更稳妥的读法是按缓存页检查提交条件：

父前缀可以恢复
  -> GPU 页已有主机内存备份
  -> 写入存储层期间，主机页保持锁定
  -> 每个必要内存池或对象都返回成功
  -> 逻辑页结果为真

所以 write_backup 要关注父节点，write_backup_storage 要关注主机缓存页生命周期，v2 写入路径要关注辅助内存池的共同结果。对于 Mooncake 后端，写入的最终目标不是对一批对象调用完 put，而是让一个逻辑页在后续读取路径中能够被完整地判定为远端命中。