RadixAttention：用 Radix Tree 管理前缀缓存

谁该读这一篇？ 想真的弄懂 SGLang 招牌算法、能讲明白和 PagedAttention 的设计差异的工程师。 前置阅读： 01-overview/00-prerequisites.md（KV cache、prefix cache 的基本概念）；最好对 PagedAttention 有印象（vllm-learning/02-core-concepts/01-paged-attention.md）。 耗时： 45 分钟 学完能： 1. 画出 RadixTree 的核心数据结构，说出每个字段为什么存在； 2. 跑通"插入一个请求 / 匹配前缀 / 引用计数变化 / 淘汰"这一整套生命周期； 3. 解释为什么 SGLang 用 token 级 + 树状管理，vLLM 用 block 级 + hash，各自的代价； 4. 给定一个 multi-tenant + LoRA 的复杂场景，能讲出 extra_key 的作用； 5. 在源码里定位到核心函数（match_prefix / insert / evict）的入口和关键行号。

1. 问题陈述

回到 00-prerequisites.md §5 那张图：

请求 A: "system: 你是助手 ... user: 1+1=?"     (prompt 1500 token)
请求 B: "system: 你是助手 ... user: 2+2=?"     (prompt 1500 token)
        └────共享 1490 token────┘└── 各自 10 ──┘

我们要在引擎里维护一个"已经算好的 KV"池，给新请求做最长前缀匹配。问题是：怎么组织这个池？

业界两条路线：

各有取舍。SGLang 的选择是为了在"复杂程序复用"场景做到极致：当一个请求只命中前缀 1500 token 中的 1490 个，hash 表做不到（最后那部分 block 凑不齐 16 token 会被算两次），radix tree 可以精确止住。

2. RadixTree 数据结构

源码：sglang/python/sglang/srt/mem_cache/radix_cache.py:208（class TreeNode）。

class TreeNode:
    def __init__(self, id=None, priority=0):
        self.children = defaultdict(TreeNode)      # 子节点：用第一个 token 索引
        self.parent: TreeNode = None
        self.key: RadixKey = None                  # 这条边上保存的 token 序列
        self.value: Optional[torch.Tensor] = None  # 设备上的 KV 索引
        self.lock_ref = 0                          # 引用计数（>0 时不能淘汰）
        self.last_access_time = time.monotonic()   # LRU 用
        self.hit_count = 0
        self.host_ref_counter = 0                  # CPU 镜像的引用计数
        self.host_value: Optional[torch.Tensor] = None  # CPU 镜像
        self.hash_value: Optional[List[str]] = None     # 分页 hash（HiCache 等用）
        self.priority = priority                   # 优先感知淘汰
        ...

关键设计：

• 边上存 token 序列（key: RadixKey），不是节点上存。这样"分裂"操作（一条边被中途打断）就是把一个节点裂成两个。
• value 存的是 KV cache 的物理索引（device_indices: int64 tensor），不是直接存 KV 张量。真正的 KV 在另一块 memory pool 里，这里只是指针。
• lock_ref 和 host_ref_counter 分别管"设备 KV 在用"和"CPU 镜像在用"——SGLang 支持 KV 异步 offload 到 CPU 内存（HiCache）。
• defaultdict(TreeNode) 让插入时不用判 if key in children，直接 children[token_id] 就有了，简化代码。

树的形状（举例）

                          root
                            │
              ┌─────────────┼──────────────┐
              │             │              │
        "Sys.X User:"    "Sys.Y User:"   "Code: ```python"
        ┌──────┴──────┐     │                │
     "1+1=?"      "2+2=?"  "hi"           "def fib"
       (A)         (B)     (C)              (D)

• 根节点是 root（空 key）。
• 边上保存 token 序列。
• 每个叶子节点对应一个"被某请求引用过的最长前缀"。

3. 三个核心操作

3.1 match_prefix：找最长命中

入口：radix_cache.py:361。

def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
    """Find the longest cached prefix of `key` in the radix tree."""
    key = params.key.page_aligned(self.page_size)
    value, last_node = self._match_prefix_helper(self.root_node, key)
    return MatchResult(
        device_indices=torch.cat(value) if value else empty,
        last_device_node=last_node,
        ...
    )

算法：从 root 出发，按新请求的 token 序列逐节点下钻。

1. 在当前节点的 children 里找以新请求下一个 token 开头的子节点。
2. 子节点边上的 token 序列与请求逐字节比较，找最长公共前缀。
3. 如果完全匹配子节点的边 → 进入子节点，继续递归；如果只部分匹配 → 就地分裂子节点，命中止于分裂处。
4. 如果没有匹配的子节点 → 命中结束。

关键细节 — 命中时的"就地分裂"：

分裂前：    ──"hello world there"──[node A]
                                       │
                                  ┌────┴────┐
                                  …         …

新请求带的 token: "hello world how are you"
                            ▲
                            └─ 在 "world" 后分叉

分裂后：    ──"hello world "──[node A']──"there"──[node A]
                                  │              │
                                  ▼              └─ A 的原子树
                              新分支
                              "how are you"
                                  │
                                  └─ [新建 node B]

分裂只重排父子关系，不复制 KV —— 物理 KV 在 memory pool 里，索引张量被切成两段就好。

extra_key 的妙用：RadixKey 除了 token id 还带 extra_key（默认 None）。不同 extra_key 的请求会被强制不共享前缀。典型用法：

• 不同 LoRA adapter（KV 数值不同，强制隔离）。
• 不同 sampling salt / cache version。
• RAG 上下文不同但 prompt 相同的请求（如果业务要求严格隔离）。

源码 docstring 直接说了这点（radix_cache.py:368）：

Entries that share identical leading token ids but have different extra_key values are intentionally kept disjoint and never share prefix nodes.

3.2 insert：把新算好的 KV 挂进树

入口：radix_cache.py:421。

请求 prefill 完成后，新算的 KV 物理索引需要"挂"到 radix tree 上，供后续请求复用。 algorithm 类似 match_prefix：

1. 从 root 下钻找最长匹配的 path。
2. 遇到部分匹配 → 分裂。
3. 走到底以后，把剩余的 token 序列建一个新节点挂上去，value = 新算的 device_indices。
4. 新节点 lock_ref += 1（当前请求还在用）。

写时复用：因为 prefill 之前已经做过 match_prefix，命中的部分 KV 是直接复用的，只有"miss 的尾部"是真正算出来的；insert 只把这段尾部挂上去。

3.3 evict：显存压力下淘汰

入口：radix_cache.py:561。

显存满了怎么腾地方？SGLang 默认是 LRU + 引用计数保护：

1. 在树里找所有 lock_ref == 0 的叶子节点。
2. 按 last_access_time 升序排序。
3. 从最旧的开始释放，直到腾出足够槽位。
4. 释放一个叶子之后，如果父节点的 lock_ref == 0 且没有其他子节点，父节点也变成可淘汰叶子。

注意：只能淘汰叶子——内部节点被淘汰会使下游 KV 失去 path，无法寻址。 这是 radix tree 比 hash 表更难 evict 的核心原因。

SGLang 还支持 优先级感知淘汰（TreeNode.priority）和 HiCache 异步 CPU 镜像（前面提到的 host_value / host_ref_counter）—— 不在本章展开，详见 03-code-walkthrough/05-radix-tree.md。

4. 一次完整生命周期演练

设定：page_size = 1（每个 token 一个 KV 槽，简化讨论）。 两个请求：

A: "X Y Z 1 2 3"
B: "X Y Z 4 5 6"

T0：树是空的，只有 root。

T1：A 进来。match_prefix("X Y Z 1 2 3") → 命中 0，全 miss。 A 走完整 prefill，6 个 token 的 KV 全部新算。 insert(...) 把 X Y Z 1 2 3 作为一个分支挂到 root 下，lock_ref = 1（A 在用）。

root ──"X Y Z 1 2 3"──[node A] (lock=1)

T2：B 进来。match_prefix("X Y Z 4 5 6") → 命中 X Y Z（3 token），止于 node A 边上的第 3 个位置。 就地分裂：

root ──"X Y Z"──[node P] (lock=0)──"1 2 3"──[node A] (lock=1)

B 复用前 3 token 的 KV，只 prefill 后 3 个 token，然后 insert：

root ──"X Y Z"──[node P] (lock=0)
                     │
                     ├──"1 2 3"──[node A] (lock=1)
                     └──"4 5 6"──[node B] (lock=1)

node P 是分裂出来的内部节点，自己没有 lock_ref（没被任何活动请求专属持有），但它支撑了两个子节点。

T3：A 完成生成、释放。node A.lock_ref -= 1（变 0）。

root ──"X Y Z"──[node P] (lock=0)
                     │
                     ├──"1 2 3"──[node A] (lock=0)   ← 可被 evict
                     └──"4 5 6"──[node B] (lock=1)

注意 node P 的 lock_ref 始终是 0，但因为它有子节点，evict 时不会考虑它（只考虑叶子）。

T4：B 完成生成、释放，同时显存吃紧。node B.lock_ref -= 1 → 0。 现在 node A 和 node B 都是叶子且 lock_ref=0。按 LRU，先 evict node A（更久未访问）。 evict 之后：

root ──"X Y Z"──[node P] (lock=0)
                     │
                     └──"4 5 6"──[node B] (lock=0)

node P 现在只有一个子节点 node B，可以做反分裂优化（把 P 和 B 合并），但实现里通常不做，简化逻辑。 继续 evict node B，然后 node P 变成无子叶子，也可以 evict。最终回到空树。

5. 与 PagedAttention 的关键差异

直觉：

• vLLM 把"缓存是不是命中"这件事粗化到 block 粒度，逻辑简单、但漏掉一些复用机会。
• SGLang 把粒度做到 token 级 + 显式树结构，复用机会全收，代价是 evict / split 等操作明显复杂。

业务上的影响：

• chat / RAG（系统 prompt 长共享）：两者差距不大，vLLM 也能把整段 prompt 命中。
• agent / 多分支推理（fork 出多个分支）：SGLang 显著优——s.fork(N) 在 DSL 层直接告诉运行时分支关系，radix tree 一行 insert 就完成共享。
• JSON Schema 严格输出：与本章无关，是压缩 FSM 的功劳。

6. Page Size > 1 的特殊情况

实际部署中 SGLang 也允许 page_size > 1（多 token 一槽，提升 attention kernel 效率）。 此时 match_prefix 会先把 key 长度对齐到 page_size 的倍数：

key = key.page_aligned(self.page_size)   # 截掉零头

代价：尾部不到一个 page 的部分不参与匹配。 收益：每条边的 KV 索引更整齐，attention kernel 可以用更大的 cascade。

PagedAttention 的 block_size 在概念上对应 SGLang 的 page_size，但 vLLM 默认 16、SGLang 默认 1。后者依赖 FlashInfer 的 cascade 在 page_size=1 时也能高效。

7. 性能直觉（理论上限 vs 现实）

记 $L_{\text{prompt}}$ 为请求 prompt 长度，$L_{\text{hit}}$ 为命中的前缀长度，$L_{\text{out}}$ 为生成长度。

朴素无缓存：每请求 prefill 算力 $\propto L_{\text{prompt}}^2$。 有 prefix cache：prefill 算力 $\propto (L_{\text{prompt}} - L_{\text{hit}})^2$。

TTFT 改善：

当 $L_{\text{hit}} / L_{\text{prompt}} = 0.95$（系统 prompt 占大头）：TTFT 降到 0.25%。 这就是 SGLang 论文里那些 "5×、10× 提升" 的数学来源。

但有几个收益不能拿到的现实情况：

1. 没人复用：每个用户独立 system prompt → 几乎 0 命中。
2. 碎片化命中：用户 prompt 中间一段一样、前后不同 → radix tree 也救不了（树只擅长前缀，不擅长任意子序列）。
3. 显存吃紧 evict 频繁：刚算好的 KV 还没被复用就被 evict 了。

工程上要关注前缀复用率监控指标——SGLang 暴露 cache_hit_rate，盯这个指标盯到吐。

8. 与调度的耦合：为什么是"RadixAttention"而不是"Radix Cache"

读到这里你可能觉得这只是个聪明的缓存数据结构，为什么叫 "Attention"？

因为：这棵树不仅决定了 KV 怎么存，还决定了 attention kernel 怎么算。

FlashInfer 的 Cascade Inference 支持"一棵 prefix tree 内部多个分支共享 K/V 张量做 attention" —— 不需要把共享部分复制 N 份，attention kernel 直接按 tree 拓扑算。 SGLang 的 radix tree 正好是这棵 prefix tree 的天然存在。

flowchart LR
    A[Shared prefix<br/>K, V tensors] --> B1[Branch 1<br/>Q1]
    A --> B2[Branch 2<br/>Q2]
    A --> B3[Branch 3<br/>Q3]
    B1 --> O1[Output 1]
    B2 --> O2[Output 2]
    B3 --> O3[Output 3]
    classDef shared fill:#fff3e0,stroke:#b35900;
    classDef branch fill:#eff5ff,stroke:#2563eb;
    class A shared;
    class B1,B2,B3 branch;

所以叫 RadixAttention —— radix tree 是数据结构，cascade attention 是 kernel，两个一起才完整。 FlashInfer 后端的实现细节见 03-code-walkthrough/06-attention-backends.md。

9. 关键源码索引（开始啃源码用）

建议在 IDE 里把 radix_cache.py 打开放在一个 split，本章放在另一个 split，对照读最直观。

10. 小结

• RadixAttention 把所有请求的 prefix 用一棵 radix tree 组织起来，节点之间用 token 序列作为边。
• 三个核心操作：match_prefix（最长前缀匹配，带就地分裂）、insert（挂新算的 KV）、evict（LRU + 引用计数保护下淘汰叶子）。
• 相比 vLLM 的 block 级 hash，粒度更细、对分支天然友好、与 FlashInfer cascade kernel 深度耦合；代价是实现复杂。
• 业务前缀复用率越高，收益越大；纯独立请求场景与 vLLM 差距不大。

11. 自检

不看上文写答案：

1. 为什么 radix tree 用 token 作为边，而不是节点上存 token？分裂操作的代价分别是什么？

2. 一棵 radix tree 里能不能 evict 内部节点？为什么？

3. extra_key 在什么场景下必须用？不用会发生什么？

4. 命中率 95% 的请求和命中率 5% 的请求混在一起跑，RadixCache 会偏向谁？这和 cache-aware scheduling 什么关系？

5. 给定 prompt = 4096 token、命中前缀 = 3900 token，估算 prefill 算力节省的比例。

12. 下一步

• 看代码：03-code-walkthrough/05-radix-tree.md 把每个函数的实现细节过一遍。
• 理解调度怎么用这棵树：02-core-concepts/03-cache-aware-scheduling.md。
• 理解 attention kernel 怎么用这棵树：03-code-walkthrough/06-attention-backends.md（FlashInfer cascade）。
• 横向对比 vLLM：vllm-learning/02-core-concepts/01-paged-attention.md、04-prefix-caching.md。
• 动手实验：07-hands-on/03-debug-scheduler.md，把 print 加在 match_prefix 里看树怎么长大。

上游源码起点：sglang/python/sglang/srt/mem_cache/radix_cache.py。