源码精读:RadixCache 实现细节
谁该读这一篇? 已经读过
02-core-concepts/01-radix-attention.md,想啃源码每行细节的工程师。 前置阅读: 上面那篇必读。再看一遍 §2 的TreeNode字段会有帮助。 耗时: 50 分钟 学完能: 1. 在radix_cache.py和unified_radix_cache.py里定位每个函数的角色; 2. 读懂_match_prefix_helper递归逻辑; 3. 解释 lock_ref / host_ref_counter 双计数器在什么场景下分别用; 4. 在 evict 时能预测哪个节点会被先淘汰; 5. 改写淘汰策略时不会踩 race 坑。
1. 文件分布
mem_cache/
├── base_prefix_cache.py ← 抽象基类
├── radix_cache.py ← 基础 RadixCache 实现
├── unified_radix_cache.py ← 生产路径的统一设备/主机缓存
├── unified_cache/ ← 树核心、组件与 HiCache L3 适配
├── chunk_cache.py ← chunk-only 模式(关闭 radix)
├── memory_pool.py ← KV 物理存储
├── allocator/ ← slot/page 分配
├── evict_policy.py ← 淘汰策略接口
├── cache_init_params.py ← 初始化参数
└── cpp_radix_tree/ ← C++ 加速版(可选)
本章只精读 radix_cache.py,其它顺带提及。
2. 主要类(按行号)
| 类 | 行 | 作用 |
|---|---|---|
RadixKey |
59 | 复合 key:token ids + extra_key |
TreeNode |
238 | 基础树节点 |
RadixCache |
303 | 基础实现 |
UnifiedRadixCache |
148 | 统一缓存 facade,组合 tree/components |
KVCacheEventMixin |
(从其它文件) | 事件通知(HiCache 用) |
BasePrefixCache |
(base_prefix_cache.py) | 抽象基类 |
3. RadixKey(66+)
@dataclass
class RadixKey:
token_ids: List[int]
extra_key: Optional[str] = None # LoRA id / 命名空间 / cache version
def page_aligned(self, page_size):
...
def maybe_to_bigram_view(self, is_eagle):
# EAGLE 投机解码用:把 token 序列看作 bigram
...
设计意图:
- 不直接用
List[int]而是包装 → 可以加 metadata。 extra_key让"相同 token 但语义不同"的请求隔离,举例:不同 LoRA 的 KV 数值不同,必须分树。
__hash__ / __eq__ 都依赖 (tuple(token_ids), extra_key),所以 extra_key 不同的 key 严格不等。
4. TreeNode(208+)
class TreeNode:
counter = 0
def __init__(self, id=None, priority=0):
self.children = defaultdict(TreeNode) # 按首 token id 索引
self.parent: TreeNode = None
self.key: RadixKey = None # 这条边上的 token 序列
self.value: Optional[torch.Tensor] = None # device KV slot 索引
self.lock_ref = 0 # device 引用计数
self.last_access_time = time.monotonic()
self.creation_time = time.monotonic()
self.hit_count = 0 # 用于 LFU
self.host_ref_counter = 0 # CPU 镜像引用计数
self.host_value: Optional[torch.Tensor] = None # CPU 镜像
self.hash_value: Optional[List[str]] = None # HiCache 用,每 page 一个 hash
self.priority = priority # priority-aware eviction
self.id = TreeNode.counter if id is None else id
TreeNode.counter += 1
双计数器(lock_ref / host_ref_counter)的设计:
lock_ref:当前有请求在用 device 上的 KV(不能从 GPU evict)。host_ref_counter:当前有 HiCache 操作引用 CPU 镜像(不能从 CPU 释放)。
两个独立,避免"device 已经 evict 但 host 还要用"或反过来。
hit_count:每次 match_prefix 命中时 ++,可用作 LFU 排序。SGLang 默认 LRU,但可切到 LFU(evict_policy.py)。
5. RadixCache.__init__(303+)
class RadixCache(KVCacheEventMixin, BasePrefixCache):
def __init__(self, params: CacheInitParams):
self.token_to_kv_pool_allocator = params.allocator
self.page_size = params.page_size
self.disable = params.disable
self.is_eagle = params.is_eagle
self.root_node = TreeNode() # 根
self._empty_match_result = MatchResult(...)
# ... 锁 / events / mixin 状态
root_node是空节点,没有 key、没有 value,只是所有子树的父。_empty_match_result预建空结果,避免每次 alloc。disable=True时所有操作 no-op(用于关闭 prefix caching 跑对照基准)。
6. match_prefix 主路径(361+)
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
key = params.key
key, _ = key.maybe_to_bigram_view(self.is_eagle)
if self.disable or len(key) == 0:
return self._empty_match_result
key = key.page_aligned(self.page_size)
if len(key) == 0:
return self._empty_match_result
value, last_node = self._match_prefix_helper(self.root_node, key)
if value:
value = torch.cat(value)
else:
value = self._empty_match_result.device_indices
return MatchResult(
device_indices=value,
last_device_node=last_node,
last_host_node=last_node,
best_match_node=last_node,
)
注意点:
- EAGLE 视角:投机解码模式下 key 要做 bigram 变换。
- page 对齐:截掉零头,避免命中跨 page 半截。
- 真正递归在
_match_prefix_helper。
6.1 _match_prefix_helper(递归核心)
伪代码(实现细节复杂,简化):
def _match_prefix_helper(self, node, key):
"""从 node 出发,匹配 key 的最长前缀。"""
if len(key) == 0:
return [], node
first_token = key.first_token()
if first_token not in node.children:
return [], node # 当前节点子树里没匹配
child = node.children[first_token]
common_prefix_len = longest_common_prefix(child.key, key)
if common_prefix_len == len(child.key):
# 完全消耗 child 的边,递归下钻
remaining_key = key.slice_from(common_prefix_len)
child.last_access_time = time.monotonic()
child.hit_count += 1
sub_values, last_node = self._match_prefix_helper(child, remaining_key)
return [child.value] + sub_values, last_node
else:
# 部分匹配 → 分裂 child
new_node = self._split_node(child, common_prefix_len)
return [new_node.value], new_node
6.2 _split_node:就地分裂
def _split_node(self, child, split_pos):
"""把 child 在 split_pos 处一分为二。
Returns 上半部分(也是新的中间节点)。"""
head_key = child.key.slice_to(split_pos)
tail_key = child.key.slice_from(split_pos)
head_node = TreeNode()
head_node.key = head_key
head_node.value = child.value[:split_pos] # KV 切片(只是改 tensor view,不复制)
head_node.parent = child.parent
# 改 child
child.key = tail_key
child.value = child.value[split_pos:]
child.parent = head_node
head_node.children[tail_key.first_token()] = child
# 父节点的指针改向新 head
child.parent_old = child.parent # 暂存
head_node.parent.children[head_key.first_token()] = head_node
return head_node
物理 KV slot 不动,只是索引张量被切片。
7. insert 主路径(421+)
def insert(self, params: InsertParams) -> InsertResult:
"""把 (key, value=device_indices) 挂到树上。"""
# 1. 先尝试 match,命中部分跳过
matched_node, depth = self._find_attach_point(params.key)
# 2. 剩余部分建新节点
remaining_key = params.key.slice_from(depth)
remaining_value = params.value[depth:]
new_node = TreeNode()
new_node.key = remaining_key
new_node.value = remaining_value
new_node.parent = matched_node
matched_node.children[remaining_key.first_token()] = new_node
# 3. 引用计数(让插入的节点至少有一个 ref,防止立刻被 evict)
new_node.lock_ref += 1
return InsertResult(node=new_node, ...)
_find_attach_point 类似 _match_prefix_helper,但带"如果没分裂就在当前节点尾部挂新子节点"的逻辑。
8. evict 主路径(561+)
def evict(self, params: EvictParams) -> EvictResult:
"""释放 params.num_tokens 个 token 的 KV slot."""
leaves = self._collect_evictable_leaves() # lock_ref == 0 的叶子
leaves.sort(key=lambda n: (n.priority, n.last_access_time)) # LRU + priority
freed = 0
freed_slots = []
while freed < params.num_tokens and leaves:
node = leaves.pop(0)
freed_slots.append(node.value)
freed += node.value.shape[0]
self._detach(node) # 从父节点移除
# 父节点是否变成新叶子?是的话加进 queue
parent = node.parent
if parent is not self.root_node and len(parent.children) == 0 and parent.lock_ref == 0:
insort(leaves, parent, key=lambda n: (n.priority, n.last_access_time))
self.token_to_kv_pool_allocator.free(torch.cat(freed_slots))
return EvictResult(num_freed_tokens=freed)
关键不变量:
- 只 evict 叶子。
- evict 后父节点可能变叶子,要重新插入 candidate 队列。
- priority 字段让"业务标 priority=高"的节点最后才被淘汰。
9. 并发安全
实际上 RadixCache 在 Scheduler 主循环里只在一个线程访问(Scheduler 是单线程 while True)。 所以 RadixCache 内部没有锁。
但有几个 race 风险:
- HiCache 异步线程操作 host_ref_counter / host_value,要用专门的锁。
- 多 TP rank 的 RadixCache 是各自独立的;它们通过 Scheduler 的 broadcast 保持视图一致(rank 0 操作,rank > 0 跟着做相同操作)。
- 不要在 RadixCache 外部直接改 TreeNode 字段(debug 时也别)。
10. 内部节点能不能 evict?
直觉上"如果内部节点的所有子节点都被 evict 了,它就没用了"——可以 evict。 SGLang 实现里:
- evict 时父节点的子树清空后,父节点也被视为可 evict 叶子(见 §8 的循环)。
- 但 evict 父节点的"KV value"时要注意 —— 这块 KV 是真实存在的(之前某段 prefix 用过)。
- 实现上把"父节点也是叶子"的判断写在
_collect_evictable_leaves里。
11. UnifiedRadixCache 与 HiCache L3
当前生产路径不是旧版 HiRadixCache(RadixCache) 继承,而是
UnifiedRadixCache(unified_radix_cache.py:148)组合统一树核心和多个组件:
- 设备侧 full/SWA/Mamba 状态由组件分别管理;树核心负责 prefix 匹配、引用和淘汰顺序。
-
match_prefix、insert、evict的统一 facade 分别位于unified_radix_cache.py:496、:515、:534。 -
HiCache L3 通过
HiCacheStorage/storage attachment 将完整 KV page 持久化到 host 或外部存储;命中但设备页不存在时由 prefetch/restore 流程恢复,而不是 直接把一个host_value字段当作全部状态。
启用层级缓存后,HTTP 控制面可通过 PUT /hicache/storage-backend 动态挂载 storage,
用 DELETE 同一路径卸载、GET 查询状态;内部 Scheduler 消息仍叫
AttachHiCacheStorageReqInput。旧的 /attach_hicache_storage 不是当前 HTTP 路由。
监控中会出现 sglang:hicache_host_used_tokens 与
sglang:hicache_host_total_tokens。详见 unified_cache/
和 http_server.py。
12. 关键源码索引
| 内容 | 文件:行 |
|---|---|
RadixKey |
radix_cache.py:59 |
TreeNode |
radix_cache.py:238 |
RadixCache.__init__ |
radix_cache.py:303 |
match_prefix |
radix_cache.py:377 |
insert |
radix_cache.py:437 |
evict |
radix_cache.py:593 |
evictable_size |
radix_cache.py:659 |
UnifiedRadixCache |
unified_radix_cache.py:148、unified_cache/ |
| C++ 加速 | cpp_radix_tree/ |
13. 小结
- 基础
RadixCache主要由match_prefix/insert/evict组成;生产路径通过UnifiedRadixCache复用同一组操作语义。 _match_prefix_helper用递归 + 就地分裂,KV 物理 slot 不移动只切索引。- 双引用计数(lock_ref / host_ref_counter)分别保护 device / host KV。
- 内部节点无子时也可 evict(被循环捡进队列)。
- HiCache L3 是 storage attachment + restore/prefetch 流程,设备树与 host/storage 生命周期分开管理。
14. 自检
_split_node时 KV 物理 slot 有没有被复制?为什么?
答案
**没有**。`_split_node` 只改父子指针 + 把原节点的 `value`(device_indices: int64 tensor)切两段,物理 KV 张量在 memory_pool 里原地不动。 原因:KV 张量很大(几 GB),复制代价不可接受;而 device_indices 只是 int64 数组,切片是张量 view 操作 O(1)。 后果:split 极廉价,可以频繁做。- 双引用计数为什么不能合一?
答案
`lock_ref`(device)和 host/storage 引用独立,是因为 UnifiedRadixCache 的 设备页、host 页和 L3 storage 生命周期可以不同。设备页被淘汰时,异步 backup/restore 仍可能持有 host/storage 引用;反过来设备页在使用时,host 副本可以回收。具体引用由 unified tree/component 与 storage attachment 管理, 不能假定所有版本都有一个统一的 `host_value` 字段。- 一个孤立的内部节点(无子、无 ref)会发生什么?
答案
evict 时它**变成新叶子**被捡进 candidate 队列继续被淘汰。`evict` 实现里循环检查"父节点是否变叶子"([`radix_cache.py:593`](../sglang/python/sglang/srt/mem_cache/radix_cache.py)),逐层往上回收。 极端情况:一连串内部节点都变叶子,从下往上一路 evict 到 root 的直接子节点(root 自己不会被 evict)。 树最终如果完全空,只剩 root,等价于 RadixCache 初始状态。- HiCache 的 prefetch 触发条件是什么?
答案
`match_prefix` 命中设备缺页但 host/storage 有副本时,UnifiedRadixCache 会通过 HiCache storage attachment 的 restore/prefetch 流程申请设备页并异步回填。请求 是否等待恢复由调度器和 cache action 决定;不要直接操作 `TreeNode.host_value`。 详见 [`unified_radix_cache.py`](../sglang/python/sglang/srt/mem_cache/unified_radix_cache.py) 和 [`unified_cache/`](../sglang/python/sglang/srt/mem_cache/unified_cache/)。- 多 TP 时 RadixCache 怎么保持视图一致?
答案
**TP rank 0 是主,rank 1+ 跟随**。 主流程:rank 0 的 Scheduler 持有"权威 RadixCache",所有 match_prefix / insert / evict 决策都在 rank 0 做;其它 rank 接收 rank 0 通过 broadcast 发来的 ScheduleBatch 元数据(含 page_indices 等),各自的本地 RadixCache 视图按指令同步更新。 关键不变量:**所有 rank 持有相同的 RadixCache 拓扑**——KV 物理上每 rank 各持一份切片(按 head 维),但树结构 / 节点元数据完全一致。这样 attention kernel 拿到的 page_indices 在所有 rank 上指向自己 GPU 上对应的 KV 切片。 实现:决策只在 rank 0 跑,其它 rank 跑 `tp_worker.py` 接收命令;用 NCCL broadcast 同步 batch 信息。15. 下一步
06-attention-backends.md— RadixCache 怎么被 attention kernel 用。02-core-concepts/03-cache-aware-scheduling.md— Scheduler 怎么用 RadixCache。02-core-concepts/05-token-attention.md— page=1 的设计选择。- 源码:
radix_cache.py、unified_cache/。