调度器里的前缀树

主线位置：一个请求的一生里「调度 → KV 块分配 / 前缀命中」这一段。SGLang 调度器每收一个请求要查一次前缀树，请求结束要插回去，显存不够要驱逐；从 v0.5.21 起，这棵树默认是 Rust 的。

版本：vLLM v0.32.0、SGLang v0.5.22

这是两个引擎里唯一默认跑在调度器热路径上的第一方 Rust。分工很清楚：

• Rust（sglang-radix-tree）管树本身：匹配和插入时的遍历、节点的分配与回收、引用计数锁、LRU 次序、驱逐时挑哪个节点、KV 事件队列。
• Python（UnifiedRadixCache）管编排：拿 Rust 返回的「要释放的 KV 下标」去调分配器，做 HiCache 的设备 ↔ 主机拷贝、等异步确认，再告诉 Rust 继续。

README 第一段就是这句话：树的一侧在 Rust，缓存编排留在 Python（README.md#L3）。

基数树本身怎么用来复用 KV、和 vLLM 按块哈希的前缀缓存有什么不同，见 RadixAttention。这一页只讲 Rust 在里面的位置。

flowchart LR
  subgraph PY["Python：UnifiedRadixCache"]
    m["match_prefix"]
    ins["insert：泵一个可恢复的遍历"]
    ev["evict：逐节点推进驱逐"]
    app["应用动作<br/>释放 KV、拷贝、等确认"]
  end
  subgraph RS["Rust：TreeCore（PyO3）"]
    walk["匹配 插入遍历"]
    arena["节点 arena"]
    lock["lock_ref 引用计数"]
    lru["LRU 链表与驱逐堆"]
    evq["KV 事件队列"]
  end
  m --> walk
  ins --> walk
  ev --> lru
  walk --> app
  lru --> app
  app --> alloc["KV 分配器 / HiCache 控制器"]

选择发生在哪里

UnifiedRadixCache 构造时问一次注册表用哪个树核心，再按名字构造（unified_radix_cache.py#L220）：

# python/sglang/srt/mem_cache/unified_radix_cache.py
self._tree_core_backend = select_tree_core_backend(params)
# ...
self.tree_core = create_tree_core(
    name=self._tree_core_backend,
    params=tree_params,
    components=self.components,
)

名字来自缓存实例上的覆盖值，没有就读 SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND，默认 "rust"（environ.py#L688，tree_core_registry.py#L124）。注册表里只有两个后端："python" 构造纯 Python 的 UnifiedTreeCore，"rust" 构造 RustUnifiedTreeCore 适配器（tree_core_registry.py#L171）。两者实现同一个抽象类 UnifiedTreeCoreInterface（unified_tree_core_interface.py#L162）。

前提是前缀缓存选的是 UnifiedRadixCache。它是默认的前缀缓存（registry.py#L146），LMCache 的那个也是它的子类；但纯滑窗模型走 PureSWARadixCache（registry.py#L127），打开 FlexKV 走 FlexKV 自己的实现，这两种都不经过树核心。

回退矩阵

选到 "rust" 之后，_rust_fallback_reason 按顺序检查，第一个命中的条件让它换成 "python" 并记一条 info 日志（tree_core_registry.py#L44）。显式设成 rust 也走同一套检查（tree_core_registry.py#L109）。

第 4 到第 8 条之间还有两处提前返回「用 Rust」的捷径：

• 扩展模块已经 import 过，或者 SGLANG_RUST_BUILD_MODE 不是 auto / never / force 之一，直接返回 Rust，让加载器去报它自己的错（#L76）。
• 不是 force，且（never 模式或不是源码检出），且有打包好的扩展：直接用，不查工具链（#L86）。正式 wheel 里的用户走的就是这条。

SGLANG_RUST_BUILD_MODE 默认 auto：缺扩展就从源码构建（environ.py#L1905）。也就是说，在源码检出上第一次起服务，可能要先编译这个 crate，产物放进按指纹区分的本地缓存，不写进 Python 包目录。

flowchart TB
  start["后端名 = rust"] --> s1{"会话缓存?"}
  s1 -->|是| pyb["Python TreeCore"]
  s1 -->|否| s2{"组件只有 FULL SWA MAMBA<br/>且无自定义覆盖?"}
  s2 -->|否| pyb
  s2 -->|是| s3{"Linux?"}
  s3 -->|否| pyb
  s3 -->|是| s4{"torch 2.11 到 2.14<br/>设备 cpu 或 cuda?"}
  s4 -->|否| pyb
  s4 -->|是| s5{"已加载 或 有打包扩展且不必构建?"}
  s5 -->|是| rsb["Rust TreeCore"]
  s5 -->|否| s6{"有源码 且 cargo rustc 可用?"}
  s6 -->|是| rsb
  s6 -->|否| pyb

回退之后的失败不再兜底。README 说得明确：受支持配置里的构建、导入、运行时错误仍然是错误（README.md#L20）。

PyO3 接口

扩展模块名 sglang.srt.mem_cache.rust_tree_core.mem_cache，在 crate 的元数据里声明，python/setup.py 自动发现（Cargo.toml#L9）。它不在 rust/ 工作区里（rust/Cargo.toml#L13），因为依赖 tch，构建时要链 libtorch：

# rust/sglang-radix-tree/Cargo.toml
[lib]
name = "mem_cache"
crate-type = ["cdylib"]

# Keep this crate in its own workspace: tch needs libtorch at build time,
# which the rest of the rust/ workspace does not.
[workspace]

[dependencies]
hashbrown = "0.16"
pyo3 = { version = "0.29", optional = true }
sha2 = "0.10"
tch = "=0.24.0"
# ...

要 libtorch 的原因：节点上存的值是 KV 池下标的 torch.Tensor，Rust 一侧以 tch::Tensor 不透明地持有，不复制成 Vec。

对 Python 暴露两个类，用同一个宏生成，区别只在子节点的键（python_bindings.rs#L3979）：

// rust/sglang-radix-tree/src/python_bindings.rs
tree_core_binding!(
    /// The UnifiedTreeCore Python adapter over single-token (unigram) child keys.
    RustUnifiedTreeCoreBinding,
    Vec<i64>
);

tree_core_binding!(
    /// The UnifiedTreeCore Python adapter over bigram (EAGLE) child keys; keys
    /// cross the boundary as raw token ids and pair up rust-side.
    RustBigramUnifiedTreeCoreBinding,
    Vec<(i64, i64)>
);

EAGLE 投机解码时用二元组做键，Python 适配器按 is_eagle 选类（adapter.py#L456）。

每个方法都长一个样：外层 catch_native_panic，里面取锁、在锁里干活，干活时释放 GIL（python_bindings.rs#L2789，python_bindings.rs#L1130）：

// rust/sglang-radix-tree/src/python_bindings.rs
// 宏生成的 pymethods
fn match_prefix(
    &self,
    py: Python<'_>,
    params: &MatchParamsBinding,
) -> PyResult<MatchResultBinding> {
    catch_native_panic(|| self.inner.match_prefix(py, params))
}

// TreeCoreBinding<K>::match_prefix
let result = py.detach(|| self.core().match_prefix(&params));
MatchResultBinding::from_match_result(py, result)

宏里一共生成一百多个方法，按用途分几组：

驱逐和插入都拆成了「开始 / 下一步 / 结束」，这是分工的关键。Python 的驱逐循环（unified_radix_cache.py#L940）每次让 Rust 挑一个节点，拿到要释放的下标后自己去释放；遇到 Mamba 或滑窗状态需要先备份到主机内存时，Rust 返回一个「要备份的节点」，Python 做拷贝、等确认，再调 finish_*_state_eviction 让 Rust 继续（unified_radix_cache.py#L795）。README 的说法：这样 I/O 留在 Rust 树的互斥锁之外（README.md#L57）。

sequenceDiagram
  participant P as Python UnifiedRadixCache
  participant R as Rust TreeCore
  participant A as KV 分配器 / HiCache
  P->>R: evict_device_start(组件, 配额)
  loop 直到够数或无可驱逐
    P->>R: evict_device_next_node
    alt 普通叶子
      R-->>P: 节点与要释放的下标
      P->>A: free(下标)
    else 需先备份状态
      R-->>P: 要备份的节点
      P->>A: 设备到主机拷贝，等确认
      P->>R: finish_*_state_eviction
      R-->>P: 要释放的下标
      P->>A: free(下标)
    end
  end
  P->>R: evict_device_end

数据结构

节点 arena

所有节点住在一个 Vec<Option<Node>> 里，父子之间用下标而不是指针相连（node.rs#L1090）：

// rust/sglang-radix-tree/src/node.rs
pub struct NodeArena<K: ChildKeyType> {
    /// Node store indexed by `NodeIdx_`; `None` marks a freed slot.
    nodes: Vec<Option<Node<K>>>,
    /// Freed slot ids available for reuse.
    free: Vec<NodeIdx_>,
    /// External handle -> live slot; a freed `NodeId` leaves the map, so a
    /// stale handle can never alias a recycled slot.
    id_map: HashMap<NodeId, NodeIdx_>,
    /// Next external handle; monotonic, never recycled (survives `reset`).
    next_id: NodeId,
    // ...
}

两套编号：

• NodeIdx_：arena 里的槽位，回收后会复用，只在 Rust 内部用。
• NodeId：交给 Python 的句柄，单调递增，永不复用，连 reset 也不清零。Python 拿着一个已经被驱逐的节点的 NodeId 来调，id_map 里查不到，得到的是 NodeAccessError，而不是悄悄指到一个新节点上。

在 Python 实现里，这类「悬空引用」靠对象引用天然避免；换成 arena 之后要显式地防，id_map 就是做这个的。

节点（node.rs#L159）的主要字段：

• children：hashbrown 表，键是（命名空间，子节点第一页的键）。命名空间由 extra_key 和 cache_salt 组成，不同 LoRA 或加了盐的请求互不共享前缀。
• values：每个（组件 × 层级）一格，组件是 FULL / SWA / MAMBA，层级是设备 / 主机。每格是 ValueState { value: Option<Tensor>, lock_ref: u32 }（node.rs#L715）。value 为 None 表示墓碑（已驱逐或滑出窗口），lock_ref 是正在用这段 KV 的请求数。
• last_access_counter、hit_count、priority：给各种驱逐策略算优先级用。
• hash_value：每页的链式哈希，给 L3 存储和 KV 事件用。

锁

这里有两层「锁」，不要混：

1. 引用计数锁 lock_ref：请求匹配到一段前缀后 inc_lock_ref，这段路径上的节点就不能被驱逐；请求结束 dec_lock_ref。这是缓存语义上的锁，Python 版也有。
2. 互斥锁 Mutex<UnifiedTreeCore>：整个树核心包在一把 std::sync::Mutex 里（python_bindings.rs#L1018）。注释说它的作用是让只有 Send 的核心满足 pyclass 要求的 Sync，并且因为有了它，调用时才能用 py.detach 释放 GIL。调度器是单线程调它，所以这把锁基本不会有争用。

互斥锁带来一个后果：Rust 里 panic 时，守卫在展开时毒化这把锁；之后每次调用 core() 都拒绝复用状态（python_bindings.rs#L1079）：

// rust/sglang-radix-tree/src/python_bindings.rs
/// Lock the core for one adapter call. A panic can leave a mutation half-applied,
/// so a poisoned core is never reused.
fn core(&self) -> std::sync::MutexGuard<'_, UnifiedTreeCore<K>> {
    self.core.lock().unwrap_or_else(|_| {
        panic!("Rust TreeCore mutex poisoned; refusing to reuse state after an earlier panic")
    })
}

外层的 catch_native_panic 把 panic 变成 Python 的 RuntimeError（python_bindings.rs#L30），让 Python 的崩溃处理器去协调退出。也就是说：一次 panic 等于这个调度器的前缀缓存报废，进程该退出。crate 顶部整体放开了 clippy::unwrap_used、clippy::panic 等检查，并留了一条 TODO，要把可恢复的 panic 换成显式错误（lib.rs#L2）。

LRU 与驱逐

树核心的状态（unified_tree_core.rs#L613）里和驱逐相关的有三样：

• evictable_device_leaves / evictable_host_leaves：当前可驱逐的叶子集合。
• lru_lists：每个（组件 × 层级）一条 LRU 链表。链表不是指针链，而是一张以 NodeIdx_ + 2 为下标的格子表，0 和 1 号格是头尾哨兵（unified_lru_list.rs#L35）。
• full_evict_device_heap：FULL 组件在设备上没有 LRU 链表，而是在每次驱逐开始时把所有可驱逐叶子按策略算出的优先级放进一个最小堆（full.rs#L139）。每弹出一个叶子，下一步时检查它的父节点是否变成了叶子，是就重新入堆。

优先级由可插拔的策略算，构造时校验，支持 lru、lfu、fifo、mru、filo、priority、slru、tlru（python_bindings.rs#L1058）。组件组合只接受四种：[Full]、[Full, Swa]、[Full, Mamba]、[Full, Swa, Mamba]（python_bindings.rs#L1035），和回退矩阵第 2 条对应。

flowchart TB
  subgraph core["UnifiedTreeCore（Mutex 内）"]
    ar["NodeArena<br/>Vec of Option Node<br/>free 列表 + id_map"]
    lv["可驱逐叶子集合<br/>设备 / 主机"]
    ll["LRU 链表 × 组件 × 层级<br/>格子表 + 哨兵"]
    hp["FULL 驱逐最小堆"]
    eq["KV 事件队列"]
  end
  ar --> lv
  lv --> hp
  ar --> ll
  ar --> eq
  eq -->|take_events| pub["Python KV 事件发布"]

KV 事件：树结构变动时往 kv_event_queue 里追加 BlockStored / BlockRemoved，Python 定期 take_events 取走再发布（unified_tree_core.rs#L654）。还有一个开关让事件用 Dynamo 的格式。KV 感知路由订阅的就是这些事件，见路由层。

规模

sglang-radix-tree 35,473 行 Rust code（总 39,134，含 src/tests/；tokei 口径同本部分导读）。最大的两个文件是 unified_tree_core.rs（5,942 行）和 python_bindings.rs（4,068 行），这两个数是 wc -l 的物理行。对照的纯 Python 实现 unified_tree_core.py 是 3,375 行物理行，加上 1,180 行的 Rust 适配器 adapter.py；两边共用一套 parity 测试。

成为默认的经过

默认值的变化可以直接在三个 tag 的 environ.py 上看到：v0.5.19 和 v0.5.20 里是 EnvStr("python")，v0.5.21 起是 EnvStr("rust")。从加入到成为默认不到一个月，中间隔了两个发布。

Python 实现没有删，它是回退目标，也是 parity 测试的参照：同一套 UnifiedRadixCache 测试分别以两个后端各跑一遍（例如 test_rust_unified_radix_cache_bench.py 只是把 test_unified_radix_cache_bench.py 的后端切到 rust 再跑）。

实测：两个树核心的 CPU 基准

这是 C3-8 里不需要 GPU 的那一半：用上游自带的基准脚本，在同一组合成请求序列上比两个树核心的单项操作耗时。

环境：笔记本上的 Linux arm64 容器（M2 Max 笔记本，容器 8 CPU、16 GB）。放在 Linux 容器里，是因为 Rust 树核心只支持 Linux（上文回退条件第 4 条，macOS 上永远回退 Python）；那台 4 核 x86 Linux 机器的 3 GB 容器编不过它，编译 torch-sys 生成的 C++ 时内存不足被杀。软件：Python 3.12，torch 2.14.1+cpu，SGLang v0.5.22 源码（经 PYTHONPATH 引入）；Rust 树核心由 SGLang 自己的加载器从源码构建（SGLANG_RUST_BUILD_MODE=auto，工具链是 rust/rust-toolchain.toml 指定的 1.92）。SGLANG_USE_CPU_ENGINE=1 走 CPU 平台，SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND 切 python / rust。

sgl_kernel 只有 CUDA 版的 wheel，装不上。这里用一个只在测试里用的占位模块顶替：里面的函数一调用就报错。树基准没有调用它，两个后端用的是同一个占位模块。

命令：test/registered/unit/mem_cache/test_unified_radix_cache_bench.py --bench --components full --num-seqs 5000（chunk_len 256，kv_size 500000，page_size 1）。每个后端跑 3 轮，两个后端交替。另用 --verify（num_seqs 1000）跑了一次 Rust 后端，校验通过。

下表是 3 轮的 ops/s 范围，括号里是 p50 单次耗时的范围：

拿两边 ops/s 区间的端点相除，Rust 对 Python：insert 5,958 / 3,709 ≈ 1.6 到 6,427 / 3,101 ≈ 2.1 倍；match_prefix 74,648 / 33,613 ≈ 2.2 到 80,933 / 30,131 ≈ 2.7 倍；evict 5,639 / 3,950 ≈ 1.4 到 6,156 / 2,528 ≈ 2.4 倍；lock_unlock 272,441 / 50,571 ≈ 5.4 到 337,167 / 43,062 ≈ 7.8 倍；release 4,960 / 2,628 ≈ 1.9 到 5,433 / 2,166 ≈ 2.5 倍。

这些是单项树操作在合成序列上的耗时，不是调度器每一步的开销：一步里调几次树、和分配器与 GPU 前向比占多少，要在真实服务里测。机器是 arm64 的容器，不是部署常见的 x86 + CUDA。

待实测

• 真实服务里每个调度步的树操作开销，Rust 对 Python：随 GPU 机器那一批做，待实测（C3-8 的剩余部分）。
• 真实负载下前缀命中率和 TTFT 与树核心无关（两个实现语义相同），归在 RadixAttention 页，待实测（C3-3）。

小结

• SGLang 的前缀树核心从 v0.5.21 起默认是 Rust，这是两个引擎里唯一默认跑在调度器热路径上的第一方 Rust。
• 分工：Rust 管树（arena、引用计数锁、LRU 与驱逐堆、KV 事件），Python 管编排（释放 KV、HiCache 拷贝、等确认）。驱逐和插入拆成多步，I/O 不进 Rust 的互斥锁。
• 接口：两个 pyclass（单 token 键、EAGLE 二元组键），每次调用释放 GIL；panic 变 RuntimeError 并毒化锁，缓存不再可用。
• arena 用两套编号：可复用的槽位和永不复用的外部句柄，防止 Python 拿着旧句柄指到新节点。
• 八条回退条件里，macOS、torch 版本、非 CPU/CUDA 设备、会话缓存最常见；回退之后的错误不再兜底。
• 笔记本上的 Linux arm64 容器里，上游基准测得 Rust 树核心的单项操作吞吐是 Python 的约 1.4–7.8 倍（match_prefix 2.2–2.7 倍，lock_unlock 5.4–7.8 倍）；每个调度步的真实开销待实测（C3-8 的剩余部分）。