跳到主要内容

数据截至 (上游 commit f775db03aaa8)

02 · RadixAttention:前缀缓存的基数树

本章讲什么: SGLang 最出名的机制。KV cache 怎么在一棵树上跨请求共享,以及 match / insert / evict 三个原语的真实代码长什么样。

2.1 它要解决的小问题​

LLM 推理分两段:prefill(一口气算完输入所有 token 的 KV)和 decode(逐 token 生成)。prefill 的计算量和输入长度成正比,是首 token 延迟的大头。

真实负载里,请求之间的输入大量共享开头:

  • 同一个 system prompt 下的所有请求。
  • agent 循环:第 k 轮的输入 = 第 k-1 轮的输入 + 输出 + 新工具结果,前缀几乎全同。
  • few-shot 模板、共用文档的 RAG 问答。

如果每个请求都从零算 prefill,这些共享前缀就被白白重算。问题的本质:KV cache 是可复用的计算结果,但它挂在「一次请求」的生命周期上,请求结束就扔了。

2.2 思路:把 KV cache 挂到「token 序列前缀」上​

直觉一句话:凡是 token 序列相同的开头,它的 KV 就必然相同(因果注意力下,第 i 个 token 的 KV 只依赖前 i 个 token)。所以缓存的键不该是「请求」,该是「token 序列」。

存「序列 → KV」最直接的结构就是前缀树(trie):每条 root-to-node 路径是一个前缀。SGLang 用的是它的压缩版——基数树(radix tree):把没有分叉的连续 token 压成一条边,边上同时挂着这段前缀对应的 KV 显存索引。

                root
          ┌──────┴───────────┐
   边:「系统提示 [1..128]」   边:「系统提示B [1..96]」
   value: KV块 #10..#137      value: ...
          │
     ┌────┴─────┐
  「用户问题甲」 「用户问题乙」     ← 两条请求共享 128 token 的前缀,只存一份 KV

怎么读这张图: 从 root 沿边走到的每个节点,都代表「这段前缀的 KV 已经在显存里」。新请求来了,沿树往下比对,走到走不动为止——走过的部分免算,剩下的才 prefill。

2.3 原理演示​

把三个核心操作写成简化代码,先建立心智模型:

# 示意,非源码:基数树的三个原语
def match_prefix(root, token_ids):
    """沿树往下走,返回(已缓存的KV索引, 走到的节点)。边不匹配整条时把边劈开。"""
    node, kv, rest = root, [], token_ids
    while rest and first_unit(rest) in node.children:
        child = node.children[first_unit(rest)]
        n = shared_prefix_len(child.key, rest)
        if n < len(child.key):            # 边只命中前半截
            child = split_edge(child, n)  # 劈成两段,精确暴露边界
        kv += child.value                 # value = 这段前缀的 KV 显存索引
        node, rest = child, rest[n:]
    return kv, node

def insert(node, rest_tokens, rest_kv):
    """prefill 算完,把新一段挂到树上(同样可能先劈边)。"""
    ...

def evict_lru(root, num_tokens):
    """显存不够:从最久没用的叶子开始砍,父节点变叶子后可继续被砍。"""
    ...

重点看「劈边」(split): 它是「部分命中也能精确复用」的关键——请求命中某条边的一半时,把这条边分裂成两段,复用前半段,只 prefill 后半段。没有这个操作,共享粒度只能停在历史请求的分叉点上。

2.4 真实实现​

全部在 python/sglang/srt/mem_cache/radix_cache.py(863 行,本仓库里难得的紧凑文件)。

数据结构​

TreeNode(radix_cache.py:238)一个节点四个关键字段:

字段含义
key: RadixKey这条边上的 token 序列
value: torch.Tensor对应的 KV 显存索引(token → KV pool 槽位)
lock_ref: int引用计数,>0 表示有在跑的请求正在用,禁止驱逐
last_access_timeLRU 驱逐依据(__lt__ 直接比较它)

键类型 RadixKey(radix_cache.py:59)除了 token_ids 还带两个命名空间字段:extra_key(隔离 LoRA adapter 等)和 cache_salt(调用方主动隔离缓存)。两个请求 token 完全相同但 extra_key 不同,不会共享前缀——child_key(radix_cache.py:217)把命名空间编进子节点索引。RadixKey.match(radix_cache.py:181)比较两条序列时用倍增窗口 + 二分找第一个分叉点,避免在长前缀上逐 token 跑 Python 循环。

原语一:match_prefix​

RadixCache.match_prefix(radix_cache.py:377)是对外入口;核心循环在 _match_prefix_helper(radix_cache.py:679):

while len(key) > 0 and child_key in node.children.keys():
    child = node.children[child_key]
    prefix_len = child.key.match(key, page_size=self.page_size)
    if prefix_len < len(child.key):
        new_node = self._split_node(child.key, child, prefix_len)  # 劈边
        value.append(new_node.value)
        break
    else:
        value.append(child.value)   # 整条边命中,继续往下走
        ...

_split_node(radix_cache.py:704)的写法很省:新节点拿 child.value[:split_len].clone(),原节点 key/value 截成后半截——一次分裂,没有数据搬移(只动了索引张量)。

原语二:insert(含增量缓存)​

_insert_helper(radix_cache.py:738)沿树走、命中则劈边、走到底把剩余 (key, value) 挂成新叶子,顺手 evictable_size_ += len(key) 记账。

什么时候被调?两个时机,这是 RadixAttention 区别于「请求结束才写缓存」方案的要害:

  • 请求运行中:cache_unfinished_req(radix_cache.py:516)——chunked prefill 每跑完一块就把这块的 KV 插进树(由 maybe_cache_unfinished_req 在 python/sglang/srt/mem_cache/common.py:109 包一层调用)。后续请求立刻能命中这半截。
  • 请求结束时:cache_finished_req(radix_cache.py:459)——把整条 origin_input_ids + output_ids 插进去。

增量插入带来一个重复区间问题:这段 KV 里有一部分上次已经插过了。处理方式是插入返回已存在的最长前缀长度,把重复段的显存槽释放掉,再把请求的 token→KV 映射改指向树里的槽位(cache_unfinished_req 里 free_segment + req_to_token_pool.write 那几行)。req.cache_protected_len 记录「页对齐尾部的保护段」,防止 page_size>1 时不完整的页被提前释放造成泄漏(注释见 :565-568)。

原语三:evict​

evict(radix_cache.py:593)按驱逐策略建一个叶子堆,逐个弹出直到凑够 num_tokens:

leaves = list(self.evictable_leaves)
eviction_heap = [(self.eviction_strategy.get_priority(node), node) for node in leaves]
heapq.heapify(eviction_heap)
while num_evicted < num_tokens and len(eviction_heap):
    _priority, x = heapq.heappop(eviction_heap)
    self.token_to_kv_pool_allocator.free_segment(x.value, start_pos=0)
    self._delete_leaf(x)
    if len(x.parent.children) == 0 and x.parent.lock_ref == 0:
        heapq.heappush(eviction_heap, (...))  # 父节点变新叶子,可继续砍

策略可插拔:evict_policy.py 里 LRU(默认,比 last_access_time)/ LFU / FIFO / MRU / FILO 五种(python/sglang/srt/mem_cache/evict_policy.py:10-37)。

配套机制:lock_ref​

正在被跑的请求,其前缀路径必须钉在显存里。inc_lock_ref / dec_lock_ref(radix_cache.py:623 / :638)沿节点一路爬到 root,给路径上每个节点增减 lock_ref,并同步调整 evictable_size_ / protected_size_ 两本账。调度器组批时(match_prefix_for_req,python/sglang/srt/managers/schedule_policy.py:138)对命中的 last_node 加锁,请求结束或换节点时解锁。

2.5 关键细节与坑​

  • 页对齐截断。 page_size > 1 时(paged KV pool),match_prefix 先把 key 截到页整数倍(page_aligned,radix_cache.py:150),尾巴不足一页不参与共享。换来的好处是 KV 分配按页走、碎片少。
  • 投机采样(EAGLE)用 bigram 键。 RadixKey(is_bigram=True) 把相邻 token 对当逻辑单元(radix_cache.py:60 的 docstring),因为 draft 树的验证语义下单 token 前缀不够精确。
  • ignore_eos 请求不共享。 这类请求(常见于基准测试)的输出会被当作下文继续喂,复用 KV 不安全;PrefillAdder.add_one_req 遇到 ignore_eos 且树被禁用时走单独分支(schedule_policy.py:1220 附近)。
  • 驱逐只砍叶子。 内部节点被砍会让整棵子树失联,所以驱逐堆只放叶子;砍完若父节点变叶子再补进堆。这保证任意时刻树连通。
  • 「命中数」也有自己的账。 TreeNode.hit_count 给 LFU/统计用;chunked 请求不计数,防止一个请求被自己上一轮插入的节点抬高热度(_inc_hit_count 的注释,:731-737)。
  • 这套树有好几个变体。 选哪棵由 default_radix_cache_factory(python/sglang/srt/mem_cache/registry.py:80)决定:分层缓存(GPU→CPU 的 HiCache)走 UnifiedRadixCache、混合滑动窗口模型走 SWARadixCache、实验性 C++ 实现是 RadixCacheCpp、纯关闭前缀缓存时是 ChunkCache。本文讲的 RadixCache 是所有变体的思想原型。
  • 看不出的一点: 树的操作全部在 CPU 侧用 Python 完成(操作的是索引不是 KV 本体),单线程跑在 scheduler 主循环里;库的注释声称调度器足够轻,但如果前缀极长、每步都劈边,树操作本身的开销代码里没有直接的兜底限流,只有 LPM 排序在 >128 队列时降级(见第 3 章)这一类间接保护。

2.6 本章代码地图​

主题文件路径符号名
树节点python/sglang/srt/mem_cache/radix_cache.pyTreeNode
缓存键与命名空间python/sglang/srt/mem_cache/radix_cache.pyRadixKey、RadixKey.match、child_key
最长前缀匹配python/sglang/srt/mem_cache/radix_cache.pyRadixCache.match_prefix、_match_prefix_helper
劈边python/sglang/srt/mem_cache/radix_cache.pyRadixCache._split_node
插入python/sglang/srt/mem_cache/radix_cache.pyRadixCache.insert、_insert_helper
运行中/结束时缓存python/sglang/srt/mem_cache/radix_cache.pycache_unfinished_req、cache_finished_req
驱逐python/sglang/srt/mem_cache/radix_cache.pyRadixCache.evict
驱逐策略python/sglang/srt/mem_cache/evict_policy.pyLRUStrategy、LFUStrategy 等
防驱逐锁python/sglang/srt/mem_cache/radix_cache.pyinc_lock_ref、dec_lock_ref
调度器侧调用点python/sglang/srt/managers/schedule_policy.pymatch_prefix_for_req
变体选择python/sglang/srt/mem_cache/registry.pydefault_radix_cache_factory、create_tree_cache