SGLang 源码完全拆解（五）：KV 内存池与 RadixCache 的所有权闭环

第一层：请求行不是 KV 数据#

ReqToTokenPool 初始化一个形状为

的 int32 张量 req_to_token。其中 $N_{\text{req}}$ 是最多同时保留的请求行数，$L_{\max}$ 是单请求最大上下文长度；多出的第 0 行专门给 CUDA Graph 的 padding batch 使用，使虚拟请求的默认索引落到安全位置。这里每个元素只有一个物理槽编号，不保存 K 或 V。

设请求行号为 $r$，逻辑 token 位置为 $p$，物理槽为 $s$，映射可以写成：

注意力 Kernel 再用 $s$ 访问每一层的物理缓存：

$R$ 是间接寻址表，$K_l$ 与 $V_l$ 才是第 $l$ 层的真实缓存。分页 KV 的“非连续”发生在 $s$ 上：同一请求相邻 token 的 slot ID 可以来自不同页，但 req_to_token[r, :] 仍按逻辑序列顺序列出它们。Kernel 不必要求整条请求在显存中连续，只需读取这张表。

ReqToTokenPool.alloc(reqs) 还处理一个容易忽略的分支：chunked prefill 的后续 chunk 已经拥有 req_pool_idx，必须复用旧行，只有第一次进入的请求才从 free_slots 取新行。源码从列表尾部批量弹出行号，使复杂度与本次申请行数相关；释放时只把行号放回空闲列表，并将 req.kv.req_pool_idx 置空。

这也解释了两种“容量不足”为什么不是一回事：请求行不足表示无法再追踪更多并发序列；KV 槽不足表示还有请求行，却没有空间存下一批 token 的 K/V。前者受 batch/并发规模影响，后者受上下文长度、页大小、KV dtype、层数与 KV head 数影响。

第二层：allocator 管的是槽，KVCache 管的是字节#

token allocator 的逻辑职责是维护可分配 slot/page 集合，并提供 alloc、free、available_size 与“先驱逐树缓存再重试”等操作。实际类会随普通 MHA、MLA、滑动窗口、Mamba、分层缓存和设备后端变化，但 Scheduler 依赖的是统一契约：给定要写入的 token 数，返回不冲突的位置；释放时按 allocator 的 page 粒度归还。

MHATokenToKVPool 则创建每一层的 K/V buffer，并根据 backend 采用 NHD、HND 或平台专用布局。以普通多头注意力为例，忽略 padding 与量化元数据，单个 token 的 KV 字节数近似为：

式中 $L$ 是层数，$H_{kv}$ 是 KV head 数，$D_h$ 是 head dimension，$b$ 是每个元素的字节数；系数 2 对应 K 与 V。物理池最多容纳的 token 数近似为可用于 KV 的显存除以 $B_{\text{token}}$。GQA 减少 $H_{kv}$、KV 量化减少 $b$，都会直接扩大 slot pool；它们不会改变 ReqToTokenPool 的二维映射语义。

RadixCache：CPU 上的第二套索引#

RadixCache 的 TreeNode.value 保存一段 torch.int64 物理 KV 索引，TreeNode.key 保存与之等长的 token 序列。树结构和 token key 主要驻留 CPU；真正的 KV 仍在设备 pool 中。若路径上三条边的 key 长度分别为 $a,b,c$，那么终点代表的逻辑前缀长度就是

而对应物理位置是三段 value 的串接。树不会复制 K/V 张量，只会多保存一条可复用的寻址路径。因此，同一物理 slot 可以同时出现在某个请求行和树路径中，但必须通过锁与 cache_protected_len 明确谁负责释放。

把这四层分开后，几个名字就不会混淆：req_pool_idx 是行号，kv_indices 是物理 token 槽，prefix_indices 是旧实现对命中路径的物化，当前基线更多通过 last_node 与 path_device_indices() 延迟取得路径，K/V buffer 才是大头显存。

一次 Extend 怎样把两层索引接起来#

对一条已经命中 $p$ 个 token、本轮还要计算 $e$ 个 token 的请求，调度侧先确保它有请求行，然后把缓存路径对应的 slot 写入行前缀：

接着 allocator 为新 token 申请 $e$ 个位置 $S_{new}$，写入：

本轮输入 token、位置、序列长度与新写位置被整理进 ForwardBatch。模型每层计算出新 K/V 后，以 $S_{new}$ 为散写地址写进物理 pool；attention backend 同时根据请求行读取 $S_{prefix}$ 与本轮新位置。于是一次 extend 的注意力历史在逻辑上连续，在物理上可以完全离散。

可以把一次写入的正确性条件写成：

并且对本轮新分配的集合要求：

$N_{slot}$ 是 pool 容量，$S_{live}$ 是其他未释放数据占用的槽集合。命中前缀允许与别的请求行重叠，所以不能粗暴要求整行 slot 全局唯一；真正必须互斥的是新分配写位置。很多 KV 错误排查失败，正是因为把“共享读地址合法”和“重复写地址非法”混成同一条去重规则。

这种间接寻址也解释了为什么 request row 必须在模型前向期间稳定。Kernel 可能异步读取 req_pool_indices、sequence length 和页表；CPU 若在 GPU 尚未完成时释放并复用该行，新请求会覆盖旧映射。普通 Scheduler 通过同步提交边界避免此事，Overlap Scheduler 则需要更精细地区分已发出、已完成与可回收批次，下一篇会专门展开。

为什么不让树直接取代请求行#

树适合表达跨请求共享前缀，却不适合作为每轮 Kernel 的批量页表。首先，活动请求的最新后缀可能尚未达到 checkpoint 边界，根本不在树里；其次，从每个 last_node 沿 parent 链重建路径会产生不规则 CPU 工作；最后，CUDA Graph 希望输入 tensor 的形状和地址尽量稳定。ReqToTokenPool 用固定二维表提供设备友好的随机访问，RadixCache 用 CPU 树处理动态共享，两者是针对不同访问模式的互补结构。

由此也能理解内存开销的层次。请求表的字节数约为

因为元素为 int32；RadixCache 主要增加 token key、索引副本和 Python/C++ 节点元数据；物理 KV 则随层数、KV head 和 dtype 成倍增长。工程上常常值得用较小的索引开销换取大量 KV 计算复用，但 request table 的 $N_{req}\times L_{max}$ 仍需纳入超长上下文部署的显存预算。

压缩树的边存一段 token，不是一个 token#

普通 trie 的每条边只存一个符号，长 prompt 会产生大量只有一个孩子的节点。radix tree 把连续单分支压成一条边，因此 TreeNode.key 与 TreeNode.value 都是一段序列。节点还保存 children、parent、lock_ref、last_access_time、hit_count 与 eviction priority。children 的字典键不是完整前缀，而是子边第一个逻辑单元；确定候选子边后，再逐项比较整段 child.key。

把查询 key 记为 $Q$，当前边 key 记为 $E$，两者最长公共前缀长度为

匹配循环有三种结果：

• 没有以 Q.child_key(page_size) 开头的子边，匹配在当前节点结束；
• $m=|E|$，整条边命中，把 child.value 加入结果，继续用 $Q[m:]$ 向下；
• $0<m<|E|$，查询在压缩边内部结束，需要先把边切出一个真实节点。

最后一种情况说明 match_prefix() 并非纯只读操作。假设树中已有：

1
root -- [A B C D E] / [10 11 12 13 14] --> leaf

新查询只有 [A B C]。若直接返回旧 leaf，last_node 会代表五个 token，锁定和释放边界都会多出 [D E]。_split_node() 因而改写为：

1
root -- [A B C] / [10 11 12] --> middle
2
middle -- [D E] / [13 14] --> leaf

新节点继承原节点的 priority、hit_count 与 lock_ref，旧节点保留后半段 key/value。物理 KV 没有复制；.clone() 复制的是很小的索引 tensor。分裂的目标是让命中终点、锁终点与可淘汰边界一致。

match_prefix() 返回的 MatchResult 中，普通单层 RadixCache 的 last_device_node、last_host_node 与 best_match_node 是同一个节点。接口之所以保留三个字段，是为了兼容 HiCache 等分层实现：最长逻辑命中可能在 host，而已经位于 GPU 的命中终点更短。Scheduler 应依赖接口语义，不能把当前三者相等写成永久假设。

为什么只返回终点，不长期复制整条路径#

当前基线用 last_node 保存命中终点，需要物理索引时调用 path_device_indices()：从终点沿 parent 回到 root，收集各段 value，反转后拼接。这样请求长期保存的是一个树节点句柄，而不是每轮都复制随上下文增长的完整索引数组。

代价是节点句柄依赖树结构稳定。只要请求可能使用这条路径，就必须加锁阻止相关节点被驱逐；匹配结果在下一次 allocation 或 eviction 后是否仍有效，也由 BasePrefixCache 的接口注释明确约束。换句话说，last_node 不是普通缓存查询结果，而是一张需要生命周期管理的借用凭证。

图源：SGLang 论文 Figure 3。图中每条边承载一段 token，绿色表示新插入，蓝色表示本时刻访问的缓存，红色表示被淘汰的节点。第 7 步为了暴露共享 few-shot 前缀而分裂节点；第 8、9 步则从最久未访问的叶子回收空间。源码中的 _split_node() 与 _delete_leaf() 正是这两类结构变化的落点。

用伪代码保留主干#

真实实现还要处理 page alignment、EAGLE bigram、事件 hash 与访问优先级。去掉这些分支后，最长匹配主干只有以下几步：

1
def match(node, key):
2
    values = []
3
    while key and first_unit(key) in node.children:
4
        child = node.children[first_unit(key)]
5
        n = common_prefix_len(child.key, key)
6

7
        if n < len(child.key):
8
            node = split(child, n)
9
            values.append(node.value)
10
            break
11

12
        values.append(child.value)
13
        node = child
14
        key = key[n:]
15

16
    return concat(values), node

若树深为 $D$、实际比较了 $M$ 个 token，则搜索代价可写成 $O(D+M)$；压缩边减少了 Python 节点跳转，却不能省掉必要的 token 比较。长上下文下，CPU 端重复走完整前缀仍可能成为瓶颈，这也是后续 C++/统一缓存实现持续优化树遍历的原因。

第一步：从请求行截取 token 与 slot#

checkpoint 先取 full_untruncated_fill_ids[:up_to]，再从

1
req_to_token_pool.req_to_token[req.kv.req_pool_idx, :len(token_ids)]

取出同长度的 kv_indices。二者必须位置对齐：第 $i$ 个 token 的 key 与第 $i$ 个 slot value 共同组成一条树边。随后按 page size 截齐，并复制索引 slice。复制索引的必要性在于请求行稍后会回收和复用，树不能保留指向可变行的 view。

第二步：插入时识别重复物理 KV#

insert() 沿已有路径比较 key。若某段前缀早已在树中，返回 prefix_len；剩余后缀才创建新 TreeNode。这里存在两份“内容相同但槽不同”的 KV：树中旧前缀指向槽 $S_{tree}$，本请求刚计算的相同前缀可能指向槽 $S_{req}$。保留两份没有意义，_insert_cache() 会释放请求私有区间内被现有树前缀覆盖的槽：

1
[0, cache_protected_len)       之前已由树保护，不归请求释放
2
[cache_protected_len, matched) 本轮发现为重复，释放请求副本
3
[matched, key_len)             新插入树的后缀，树接管
4
[key_len, allocated_len)       部分页或预分配尾部，仍由请求管理

cache_inserted_len 用来区分当前请求此前已经发布的部分与真正的新命中，避免统计和事件重复。对 chunked prefill 来说，每个 chunk 都会推进该边界；这不是“每轮重新缓存完整 prompt”，而是把已经算完的连续前缀逐步发布。

第三步：重新匹配并改写请求行#

插入之后，checkpoint 再调用一次 match_prefix()，得到规范化后的 new_last_node，通过 path_device_indices() 取得树选择的整条物理路径。随后只改写请求行中从旧 cache_protected_len 到新路径末端的部分。

为什么不直接继续使用插入前的 kv_indices？因为插入可能发现同 key 已由另一条请求先发布，规范路径会选择树中已有 slot，并释放当前请求的重复 slot。请求行必须改指向树的 canonical copy，否则下一轮 attention 仍可能读取已释放位置。

最后，cache_protected_len 更新为树路径长度，旧锁释放，新锁锚定 new_last_node。这三个动作必须共同提交：只改行不换锁会让新后缀可被驱逐；只换锁不改行会让行仍引用重复或已释放槽；只推进长度不处理部分页，会在请求结束时越过所有权边界。

当 page size 大于 1 时，尾部不满一页的 token 仍留在请求行，却不进入树。cache_protected_len 因而总是页对齐，并准确标出“左侧由树持有、右侧仍由请求持有”的分界。

可用、可驱逐与受保护三本账#

读日志时应区分：

• allocator.available_size()：此刻已经空闲、可立即分配的容量；
• tree_cache.evictable_size()：仍装着缓存，但可通过 eviction 变成空闲的容量；
• tree_cache.protected_size()：树中被活动请求锁住、不能用于本轮回收的容量。

在普通单组件池中，可以把潜在可分配上界近似写成：

但它不是永远精确的恒等式。page size 会带来向上取整，滑动窗口和混合模型可能有多类 allocator，分层缓存还涉及 host/device load-back；当前源码甚至会在 page-interleave sharding 下循环驱逐，直到 available_size() 真正达到要求。安全做法是把 available_size() 当分配事实，把 evictable_size() 当潜在供给，最终仍由 allocator 的实际申请结果裁决。

A 第一次进入#

树为空，A 的最长命中为 0。ReqToTokenPool 给 A 分配行 $r_A$，allocator 为前四个 token 分配一页物理槽 [100,101,102,103]。第一次 chunk 执行后：

1
R[r_A, 0:4] = [100, 101, 102, 103]

checkpoint 把 key [10,11,12,13] 与 value [100,101,102,103] 插入树，设置 cache_protected_len=4，并锁住该节点。A 的剩余 [20,21] 在后续 forward 中获得其他槽；不足一整页的尾部暂时仍属请求。

B 到达并命中#

B 的匹配上界是 $|S_B|-1=4$，恰好命中树中的完整页，得到 prefix_len=4 与相同物理路径 [100,101,102,103]。B 获得自己的请求行 $r_B$，但前四个位置写入同一组 slot：

1
R[r_A, 0:4] = [100, 101, 102, 103]
2
R[r_B, 0:4] = [100, 101, 102, 103]

这不是两份 KV，而是两张逻辑行共享四个物理位置。B 锁住命中节点后，路径的 lock_ref 再加一；即使 A 结束并解锁，B 仍在运行，页也不能被驱逐。

A 完成#

A 的 release 先把页对齐的有效后缀 checkpoint 到树，再释放未进入树的尾部或过量预分配，解锁路径并归还请求行 $r_A$。此时 slot [100..103] 仍由树与 B 引用，不会进入 allocator 空闲表；“请求行释放”和“KV 页释放”完全是两件事。

B 完成后发生内存压力#

B checkpoint 自己的 [30] 分支，但 page size 为 4 时不足一页，是否进入树取决于后续补齐与有效长度。B 解锁共享路径并释放行后，已缓存且无锁的叶子进入 evictable_leaves。新 batch 申请空间不足，evict 才真正把某个冷叶子的 value 归还 allocator。

这条时间线揭示三个不变量：

三个区间可以为空，但不能重叠释放。checkpoint 推进第一条边界，正常/异常 release 结算第二条边界，投机或批量预分配清理第三条边界。

图源：SGLang 论文 Figure 8。左侧比较不同缓存策略的命中率，右侧比较关闭 RadixAttention 等组件后的吞吐。论文结果说明前缀复用的收益来自“少算 token”与“调度维持命中”共同作用；这些数据基于论文版本、模型与工作负载，不能直接外推到当前提交或任意线上流量。

先沿所有权走，不要先钻进张量布局#

建议用一个两请求、单 GPU、无投机、无 HiCache 的最小例子，按以下顺序打断点：

1
Req.init_next_round_input
2
  -> match_kv_cache
3
  -> RadixCache.match_prefix
4

5
PrefillAdder / ScheduleBatch.prepare_for_extend
6
  -> lock matched node
7
  -> ReqToTokenPool.alloc
8
  -> token allocator alloc
9
  -> write req_to_token row
10

11
Batch result commit
12
  -> checkpoint_kv_cache
13
  -> RadixCache.checkpoint
14
  -> _insert_cache / insert / match_prefix
15

16
Finish / abort / retract
17
  -> release_kv_cache
18
  -> free_kv_row
19
  -> unlock
20
  -> ReqToTokenPool.free

第一遍只记录长度和索引，不查看所有层的 K/V 数值。推荐的观察表如下：

源码导航可以直接使用：

1
rg -n "class (ReqToTokenPool|MHATokenToKVPool)" \
2
  python/sglang/srt/mem_cache/memory_pool.py
3

4
rg -n "class (RadixKey|TreeNode|RadixCache)" \
5
  python/sglang/srt/mem_cache/radix_cache.py
6

7
rg -n "def (match_kv_cache|checkpoint_kv_cache|release_kv_cache)" \
8
  python/sglang/srt/mem_cache/common.py

用断言定位所有权错误#

若遇到 illegal memory access、生成内容随机损坏或 KV 泄漏，不要只盯 attention Kernel。可以按以下顺序缩小范围：

1. 在写 KV 前检查每个 slot 是否位于 pool 合法范围，且本 batch 内需要唯一的位置没有重复；
2. checkpoint 前检查 token key 与 kv_indices 长度一致；
3. checkpoint 后检查请求行的保护区与 path_device_indices(last_node) 完全相同；
4. release 前把三个所有权区间打印出来，确认无负长度、无交叠；
5. eviction 时确认候选是无锁叶子，释放长度等于 len(node.value)；
6. 请求行复用时比较 req_generation，排除旧异步读者把新请求的同一行当成旧状态。

ReqToTokenPool.clear() 特意不把 req_generation 清零，因为某些读者用 generation 判断行是否已经复用。若清零，经历足够少的重启/清理后，新请求可能恰好得到与旧请求相同的 generation，迟到消息就无法识别 ABA 式复用。这类字段看起来与 KV 算法无关，却是异步系统正确性的组成部分。

看懂日志里的“内存够不够”#

排查 OOM 时至少同时打印四组量：request row 的 free count、allocator 的 available token/page、tree 的 evictable/protected size、每条请求的 kv_allocated_len 与 cache_protected_len。只看 torch.cuda.mem_get_info() 解释不了预分配 pool 内部是否有槽；只看 pool available 又解释不了多少冷缓存可以驱逐。

若 available=0、evictable>0，说明可通过淘汰扩容；若二者都为 0 但 protected 很大，说明活跃 batch 或临时接纳锁住了容量；若 request rows 先耗尽，增加 KV 显存也无法提升并发。把这三类压力分开，才能决定该调 batch、上下文上限、页大小还是缓存策略。

误读一：RadixCache 自己保存 K/V tensor#

树节点只保存 token key 与物理 slot 索引。物理 K/V 由 KVCache 持有，allocator 管理 slot 的生死。树的 CPU 内存开销和 KV 的 GPU 显存开销不在一个数量级，也不能用删除 Python 节点替代 allocator free。

误读二：匹配是只读查询#

查询若停在压缩边内部，会调用 _split_node() 暴露精确终点；每次查询还刷新访问时间。并发化或迁移到 C++ 时必须保留这些写语义，不能把 match 当无锁只读函数随意并行。

误读三：命中前缀不占容量#

命中避免新分配与重复计算，却会把命中路径从 evictable 变成 protected。对单条请求，它几乎不增加物理占用；对接纳器，它会减少可承诺的驱逐容量。两句话同时成立。

误读四：lock_ref 只加终点节点#

终点依赖所有祖先边。锁必须沿 parent 链传播到 root，否则淘汰祖先会让终点成为不可达的悬空路径。解锁也必须沿同一路径对称执行。

误读五：checkpoint 会复制整段 KV#

普通设备 RadixCache 主要复制的是 slot 索引，并在必要时让请求行改指向既有 canonical slot；真实 K/V 通常仍在原物理池位置。HiCache load-back、设备迁移或特殊组件可能发生数据搬运，但那不是 radix 插入的普遍语义。

误读六：请求结束就应立即 free 全部 slot#

可复用前缀要由树继续持有。结束路径先 checkpoint，再只释放树未接管的范围，最后释放请求行。立即 free 全部 slot 会让树留下指向空闲甚至已被重用位置的悬空索引，产生静默错误。

误读七：禁用 RadixCache 后可以删除缓存接口#

Scheduler 仍需要统一的匹配、checkpoint、release 与容量调用。禁用前缀共享时可用空实现或 chunk cache 保持接口契约，这比在所有调用点散布 if disable_radix 更容易维护；chunked prefill 仍需跨轮记录已分配 KV。

误读八：page size 只影响性能，不影响正确性#

page size 同时决定匹配截断、插入长度、free 粒度与 cache_protected_len 边界。任何一处用 token 粒度、另一处用 page 粒度，都可能产生重复释放或半页悬挂。性能上的内部碎片只是表象，底层首先是所有权原子性。

还需单独审查的分支#

• LoRA 或其他改变权重的上下文必须通过 extra_key 隔离；
• 自定义 positional embedding override 不能只凭 token ID 共享 KV，当前请求会禁用前缀匹配；
• EAGLE bigram key 的逻辑长度与原始 token 数不同，key/value 截断必须一致；
• speculative verify 可能一次预分配多个位置，并在接受长度确定后回收尾部；
• sliding-window attention 已经丢弃窗口外状态，不能把普通全注意力的前缀规则原样套用；
• Mamba/混合模型除 token KV 外还有 recurrent state，checkpoint 可能需要转移组件状态；
• HiCache 的 host hit 不等于 device hit，load-back 完成前不能把它计入可执行 prefix_len；
• P/D 分离会让 KV 所有权跨节点传输，远端完成与本地释放之间需要额外握手；
• 多模态输入即使 token 占位符相同，底层图像 embedding 不同也不能错误共享。