Rust in Everything

RadixAttention：SGLang 的基数树

主线位置：一个请求的一生里「调度 → 前缀命中 / KV 分配」这一段，以及预填充之后、请求结束时把 KV 交回缓存的两步。全部发生在 SGLang 的调度器进程里：编排是 Python；树核心从 v0.5.21 起默认是 Rust（sglang-radix-tree，PyO3 扩展，进程内调用），满足不了条件时回退到 Python。vLLM 没有这棵树，它的做法见 前缀缓存。

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

结论先说：

  • 默认的前缀缓存是 UnifiedRadixCache。它管编排：调用树核心、执行树核心返回的「释放这些 KV」动作、和分配器打交道。树本身的结构、锁、LRU、驱逐遍历在「TreeCore」里。
  • 树的键是 token 序列，按 token 匹配（page_size 默认 1）。LoRA、cache_salt 不同的请求，第一层子节点的键就不同，互不命中。
  • 匹配不是只读的：命中停在某个节点中间时，当场把节点拆成两段。
  • 插入发生两次：预填充之后，和请求结束时。树里已有的那段，请求自己算的那份 KV 立即释放。
  • 正在用的路径加锁，锁一直加到根。驱逐只动没上锁的叶子，按最近访问时间从旧到新，叶子删完父节点变成叶子再接着排。

谁是默认的缓存

python/sglang/srt/mem_cache/registry.py

def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache:
    # ...
    if ctx.disable_radix_cache:
        # ...
        return create_unified_radix_cache(ctx)
    if get_memory().enable_lmcache:
        # ...
    if is_pure_swa:
        # ...
        return PureSWARadixCache(params=params)
    # ...
    return create_unified_radix_cache(ctx)

除了纯滑动窗口模型、LMCache、FlexKV 几个分支，都落到 UnifiedRadixCache。--disable-radix-cache 默认 False（memory.py）；打开它也仍然建一个 UnifiedRadixCache，只是处于禁用状态，匹配永远返回空。早先用来打开它的 SGLANG_ENABLE_UNIFIED_RADIX_TREE 已经标为废弃（environ.py）。

「Unified」指一棵树同时挂几种组件：全注意力的 KV（FULL）、滑动窗口的 KV（SWA）、Mamba 状态（MAMBA）。本页只讲普通模型用到的 FULL。

编排与树核心

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）。选了 Rust 也要先过 _rust_fallback_reason，以下任何一条成立就改用 Python 的 UnifiedTreeCore，日志里打一行 Using the Python TreeCore: ...：

条件出处
打开了会话感知缓存（enable_session_radix_cache）L45
组件不止 FULL / SWA / MAMBA，或用了自定义组件L52、L54
不是 Linux（macOS 总是回退）L55
PyTorch 不在 2.11–2.14L68、torch_build.py
设备不是 CPU 或 CUDAL73
既没有编好的扩展，也没有源码L91
要从源码编，但 cargo / rustc 不可用L105

两个实现满足同一个接口 UnifiedTreeCoreInterface。下文的代码引用 Python 版 UnifiedTreeCore，因为它更好读、行为与 Rust 版对齐；Rust 版的对应函数如 match_prefix。Rust 一侧怎么组织节点、怎么跨 PyO3 边界、为什么挪到 Rust，见 调度器里的前缀树。

分工的要点：树核心不直接释放显存。它返回一串「动作」（比如 FreeDeviceKV(槽位)），由 Python 侧的 UnifiedRadixCache 交给分配器执行（_apply_cache_actions）。这样 Rust 核心只碰元数据，不碰 CUDA 张量的生命周期。

键：token 序列，加命名空间

python/sglang/srt/mem_cache/radix_cache.py

    def child_key_at(self, offset: int, page_size: int = 1):
        t = self.token_ids
        # ...
            plain = (
                t[offset] if page_size == 1 else tuple(t[offset : offset + page_size])
            )
        if self.cache_salt is not None:
            return ((self.extra_key, self.cache_salt), plain)
        return plain if self.extra_key is None else (self.extra_key, plain)

每个节点的 children 是一个字典，键是子节点第一个 token（page_size 为 1）或第一页 token 组成的元组。带 extra_key 或 cache_salt 的请求，这个字典键前面多一层命名空间，所以从根下第一层起就走到不同的子树。extra_key 由调用方给出，LoRA 请求会把 lora_id 拼上去（schedule_batch.py）。多模态的处理方式不同，见 前缀缓存。

page_size 大于 1 时，匹配和插入前都先把键截到整页（page_aligned），比较结果也向下取整到整页（match_at）。这时命中粒度和 vLLM 的块一样。

匹配：沿树走，停在节点中间就拆

python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py

        while key_offset < len(key) and child_key in node.children:
            child = node.children[child_key]
            # ...
            prefix_len = child.key.match_at(key, key_offset, page_size=self.page_size)
            full_kv_hit_length += prefix_len
            if prefix_len < len(child.key):
                node, action = self._split_node(child.key, child, prefix_len)
                # ...
                break
            # ...
            node = child
            key_offset += prefix_len
            # ...

走到某个子节点只匹配了前一部分时，_split_node 把它拆成「已匹配的前段」和「剩下的后段」两个节点，前段成为后段的父节点，KV 槽位也跟着切开。这样匹配结果总能落在一个完整节点的末尾，后面加锁和插入都以节点为单位。

匹配结束后，从命中的最深节点一路往上到根，把每个节点的 last_access_time 改成当前时刻（L1028），这是 LRU 的「访问」。

调度器拿到命中长度后再截一刀：最多复用 输入长度 - 1 个 token，最后一个 token 必须重算才能得到 logits（_compute_max_prefix_len，调用处在 match_kv_cache）。要返回 prompt 的 logprob 时还会截得更短。

插入：预填充后一次，结束时一次

请求在两处把 KV 交给树：

  1. 每轮预填充算完，checkpoint_kv_cache 把已经算出的部分插进去。这样同时在跑的其他请求马上就能命中这段前缀，不用等它解码完。
  2. 请求结束，release_kv_cache 把 prompt 加输出整段插进去，然后释放超出部分、解锁。

插入沿树走，和匹配一样会拆节点；走过的节点刷新访问时间。对「树里已经有」的那一段，请求自己那份 KV 是重复的，当场释放：

python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py

            dup_start = max(0, state.params.prev_prefix_len - state.total_prefix_length)
            if dup_start < consumed_from:
                dup = value_slice[dup_start:consumed_from]
                # ...
                if swa_already_freed < dup.numel():
                    step_actions.append(FreeDeviceKV([dup[swa_already_freed:]]))

这种重复在并发时常见：两个请求同时预填充同一段系统提示，后插的那个发现树里已经有了，就交出自己的槽位。剩下树里没有的部分，挂成一个新叶子（_add_new_node），槽位归树所有。vLLM 在这一点上相反：同一哈希的多份块都留着，不去重。

锁：从节点一路加到根

请求命中某个节点后要持有一把锁，防止它正在用的 KV 被驱逐：

python/sglang/srt/mem_cache/unified_cache/components/full.py

        # Lock the device-on segment up to root
        delta = 0
        while cur is not root:
            cd = cur.component_data[ct]
            # ...
            if cd.lock_ref == 0:
                key_len = self.reclaimable_tokens(cur)
                self.tree_core.component_evictable_size_[ct] -= key_len
                self.tree_core.component_protected_size_[ct] += key_len
                delta += key_len
            cd.lock_ref += 1
            self.tree_core.evictable_device_leaves.discard(cur)
            cur = cur.parent

锁是计数，从被锁的节点一直加到根。某个节点第一次被锁，它的 token 数从「可驱逐」挪到「受保护」，调度器靠这两个数估算还能腾出多少显存。解锁（release_component_lock）反过来，计数归零的节点如果是叶子，重新进入可驱逐集合。节点被拆开时，新的父节点继承原节点的锁计数（full.py L148），路径上的计数保持一致。

插入之后，请求的锁从原来命中的节点挪到插入停下的节点（checkpoint_into_tree）。

驱逐：叶子最小堆

分配槽位前，分配器先算缺口，缺多少就请树驱逐多少：

python/sglang/srt/mem_cache/allocator/base.py

        shortfall = num_tokens - self.available_size()
        if shortfall > 0:
            tree_cache.evict_for_alloc(EvictParams(num_tokens=shortfall))

FULL 组件的驱逐遍历先把所有可驱逐叶子（没上锁、没有子节点）按策略的优先级建成最小堆，然后一次弹一个：

python/sglang/srt/mem_cache/unified_cache/components/full.py

    def _evict_device_start(self, request_cnt: int) -> None:
        # ...
        self._evict_device_heap = [
            (self.session_ref_eviction_strategy(n), n)
            for n in self.tree_core.evictable_device_leaves
        ]
        heapq.heapify(self._evict_device_heap)

    def _evict_device_next_node(self, tracker, device_frees, host_frees):
        # ...
        if (
            lv is not None
            and lv.parent is not None
            and lv.parent in self.tree_core.evictable_device_leaves
        ):
            heapq.heappush(
                self._evict_device_heap,
                (self.session_ref_eviction_strategy(lv.parent), lv.parent),
            )
        # ...
        while tracker[ct] < self._evict_device_request_cnt and self._evict_device_heap:
            _, x = heapq.heappop(self._evict_device_heap)
            # ...
            return x.id

上一个被删的叶子如果让父节点变成了叶子，父节点进堆接着排。删一个节点就释放它的全部槽位，所以驱逐的粒度是整个叶子：缺 10 个 token、最旧的叶子有 300 个 token，也是 300 个一起走。

优先级由 --radix-eviction-policy 决定，默认 lru（memory.py），可选 lfu、slru、priority、tlru（choices.py）。LRU 的优先级就是节点的 last_access_time（LRUStrategy）。

一个请求和树打交道的全过程

sequenceDiagram
  participant Sch as 调度器（Python）
  participant URC as UnifiedRadixCache（Python）
  participant TC as TreeCore（默认 Rust）
  participant Alloc as 槽位分配器（Python）
  Sch->>URC: match_prefix(token 序列)
  URC->>TC: 沿树匹配，必要时拆节点
  TC-->>URC: 命中长度、最深节点
  Sch->>URC: lock(节点)：锁计数加到根
  Sch->>Alloc: 要 输入长度 - 命中 个槽位
  Alloc->>URC: 不够：evict_for_alloc(缺口)
  URC->>TC: 从叶子堆里按 LRU 弹节点
  TC-->>URC: FreeDeviceKV 动作
  URC->>Alloc: 释放这些槽位
  Note over Sch: 预填充
  Sch->>URC: checkpoint：插入已算部分，重复的 KV 释放
  Note over Sch: 解码直到结束
  Sch->>URC: release：插入整段，解锁

拿一个具体的例子：两个请求共用一段 40 个 token 的系统提示，各自再跟 10 个 token，各生成 4 个 token。按上面的规则走一遍（本卷的 CPU 模拟，见 前缀缓存），第二个请求命中 40 个 token，树变成下面这样：

flowchart TB
  root(("根")) --> sys["系统提示 40"]
  sys --> a1["请求 0 的后缀 10"]
  a1 --> a2["请求 0 的输出 4"]
  sys --> b1["请求 1 的后缀 10"]
  b1 --> b2["请求 1 的输出 4"]

请求 0 预填充后插入了 50 个 token 的一个节点，结束时又挂上 4 个 token 的输出；请求 1 匹配时在第 40 个 token 处把那个 50 的节点拆成 40 和 10。同样的两个请求在 vLLM 里只命中 32 个 token，原因见下一页。

小结

  • SGLang 默认的前缀缓存是 UnifiedRadixCache：Python 编排，树核心默认 Rust，七类条件回退 Python。树核心只改元数据，释放显存的动作交回 Python 执行。
  • 键是 token 序列，extra_key（含 LoRA）和 cache_salt 在第一层就把树分开。page_size 默认 1，命中按 token。
  • 匹配会拆节点并刷新路径的访问时间；调度器最多复用 输入长度 - 1 个 token。
  • 插入在预填充后和请求结束时各一次，树里已有的部分当场释放请求自己那份。
  • 锁沿路径加到根；驱逐从未上锁叶子的最小堆里按 LRU 弹，粒度是整个叶子。
  • 树操作本身的开销（Python 对 Rust）见 调度器里的前缀树；命中率与 TTFT 待实测（C3-3）。

💬 关联讨论

发起讨论
加载中...