SGLang RadixAttention：基数树驱动的 KV Cache 自动复用与 LRU 驱逐策略
LLM 推理过程中，KV Cache 的管理效率直接决定吞吐量。vLLM 的 PagedAttention 用固定大小的 block 解决了显存碎片问题，但留下了一个遗憾：前缀缓存必须按 block 边界对齐，语义上完整的前缀被 block 尺寸强行截断。SGLang 的 RadixAttention 用一棵变长的基数树绕开了这个限制，把缓存匹配精度从 block 级提升到 token 级。
这篇从 PagedAttention 的对齐瓶颈讲起，拆解 RadixAttention 的基数树结构、最长前缀匹配机制、驱逐策略的设计取舍。
1 问题的起点：KV Cache 的前缀冗余
LLM 的自回归生成依赖 KV Cache 来避免重复计算历史 token 的 Key 和 Value 向量。不同请求之间经常共享相同的前缀：system prompt、few-shot 示例、多轮对话历史。这些前缀的 KV Cache 算一次就够，后续请求可以直接复用。
vLLM 的 PagedAttention 把 GPU 显存中的 KV Cache 切成固定大小的 block，每个 block 存储固定数量 token（比如 16 个）的 KV 向量。这个设计类似操作系统的虚拟内存分页，好处是管理简单、碎片少。
问题出在前缀缓存的匹配粒度上。假设 system prompt 有 100 个 token，block 大小是 16。前 96 个 token 占满 6 个完整 block，可以被缓存和复用。剩余 4 个 token 不构成完整 block，无法被纳入前缀缓存。更麻烦的是，PagedAttention 的前缀缓存默认关闭，需要手动传 --enable-prefix-caching 参数才能启用，匹配依据是 block 的哈希值，灵活性受限。
block 对齐带来的本质矛盾是：用户 prompt 的长度几乎不会是 block size 的整数倍，而前缀的语义边界（一句话、一段指令、一个工具描述）跟 block 的物理边界没有对应关系。
2 基数树：RadixAttention 的数据结构基础
RadixAttention 用基数树（Radix Tree）替代固定 block 来组织 KV Cache。先简要介绍基数树本身。
基数树是前缀树（Trie）的压缩版本。标准 Trie 的每个节点只存一个字符，从根到叶子的路径拼出完整的字符串。基数树允许一个节点存储多个字符组成的序列，减少了树的深度和节点数量，同时保留了"从根到任意节点的路径代表一个唯一前缀"这一核心性质。在路由表查找（IP 最长前缀匹配）和字符串索引等场景中，基数树是经典数据结构。
SGLang 把基数树应用到了 KV Cache 管理上。树的每个 TreeNode 存储三个信息：
- key 是一段变长的 token 序列，长度不受固定 block 约束
- value 是这段 token 对应的 KV Cache tensor 在 GPU 显存中的存储位置
- ref_count（引用计数）记录当前有多少个活跃请求正在使用这段 KV Cache
变长 key 是 RadixAttention 与 PagedAttention 的核心区别。PagedAttention 的每个 block 只能存固定数量的 token，而基数树节点的 key 可以容纳任意长度的 token 序列，按语义自然分段。
下面是一棵典型的 Radix Tree 结构。根节点为空，每个子节点代表一段可复用的前缀，路径越深前缀越长：
root
+-- [system: "你是一个AI助手" (500 token)] -> GPU 上 500 token 的 KV tensor
| +-- [user: "帮我写一首诗" (8 token)] -> GPU 上 8 token 的 KV tensor
| | +-- [assistant: "春风拂面来..." (50 token)] -> GPU 上 50 token 的 KV tensor
| +-- [user: "解释量子力学" (7 token)] -> GPU 上 7 token 的 KV tensor
+-- [system: "You are a translator" (300 token)] -> GPU 上 300 token 的 KV tensor两个不同的 system prompt 分别从 root 延伸出独立的子树。共享同一 system prompt 的请求会复用第一层的 500 token 节点，在此基础上各自追加不同的 user message 分支。
3 最长前缀匹配：找到能复用的最长缓存
新请求到达时，SGLang 在基数树中执行最长前缀匹配（Longest Prefix Match）。从 root 出发，沿着树向下查找，找到与新请求的 token 序列匹配的最长路径。匹配到的部分直接复用其 KV Cache，只有未匹配到的尾部 token 需要重新计算。
用一个具体例子说明匹配过程。假设新请求的 token 序列是 500 token 的 system prompt 加上 6 token 的 user message：
新请求: [system prompt (500 token)] + [user: "写一段代码" (6 token)]
匹配流程:
1. 从 root 出发
2. 命中子节点 "你是一个AI助手" (500 token)，完全匹配
3. 在该节点的子树中查找 "写一段代码" (6 token)，无匹配
4. 结果: 前 500 token 复用已有 KV Cache，仅后 6 token 需要新计算506 个 token 的 prefill 被压缩为 6 个 token 的计算量，节省约 99% 的 prefill 算力。这在 few-shot、multi-turn chat、agent 工具描述等前缀高度重叠的场景中收益巨大。
4 RadixAttention 与 PagedAttention 的精确对比
两者管理 KV Cache 的方式可以用一张对齐图来直观比较：
PagedAttention (block 对齐, block_size=16):
[===block 1===][===block 2===][===block 3===][partial 4]
命中 命中 命中 丢失
16 tok 16 tok 16 tok 4 tok
RadixAttention (token 级精确):
[======== system prompt (500 token) ========][== user msg (6 tok) ==]
全量命中 未命中PagedAttention 的缓存边界由 block 的物理尺寸决定，尾部不足一个 block 的 token 直接丢弃。RadixAttention 的缓存边界由 token 序列的实际长度决定，不存在对齐损失。
差异可以归纳为三个维度。匹配粒度上，PagedAttention 以 block 为单位，RadixAttention 以 token 为单位。开关机制上，vLLM 需要手动传 --enable-prefix-caching 启用前缀缓存，SGLang 默认自动启用。匹配方式上，PagedAttention 基于 block 的哈希值做精确匹配，RadixAttention 沿基数树做最长前缀匹配，能命中任意长度的公共前缀。
在 few-shot learning、多轮对话、agent 工具调用、RAG 文档检索等前缀大量重叠的场景下，SGLang 相比 vLLM 实现了最高 5 倍的吞吐提升。
5 驱逐策略：显存不够时淘汰谁
基数树中缓存的节点会越来越多，GPU 显存终归有限。当新请求需要分配显存而空间不足时，必须从树中淘汰一部分不再活跃的缓存节点。淘汰策略需要在"释放空间"和"保留热点数据"之间取得平衡。
5.1 引用计数：区分可驱逐与不可驱逐
每个 TreeNode 的 ref_count 字段记录当前有多少个活跃请求正在使用该节点的 KV Cache。请求开始使用某节点时 ref_count 加 1，请求完成或不再需要时减 1。
这个字段划出了一条硬边界：ref_count 大于 0 的节点绝对不能驱逐，否则正在执行的请求会丢失中间计算状态，导致输出错误。只有 ref_count 等于 0 的节点才进入驱逐候选池。
5.2 LRU：优先淘汰最久未被访问的节点
SGLang 默认的驱逐策略是 LRU（Least Recently Used）。在 ref_count 等于 0 的所有节点中，按最后访问时间排序，优先驱逐最久没被访问的节点，释放其 KV tensor 占用的显存，直到腾出足够空间。
选择 LRU 而非 FIFO（先进先出）有具体的理由。基数树中存在明显的热点数据：system prompt 对应的节点可能很早就进入缓存，但几乎每个请求都会命中它。LRU 保护了这类热点，只要节点持续被访问，其最后访问时间就会不断更新，不会被排到淘汰队列的前面。FIFO 只看进入时间，不管后续访问频率，可能把仍在被频繁使用的热点节点误驱逐。
5.3 LFU：按访问频率淘汰
SGLang 也支持 LFU（Least Frequently Used）作为可选策略。LFU 根据节点的累计访问次数排序，优先驱逐被访问次数最少的节点。
两种策略各有适用场景。请求模式相对稳定、热点前缀长期活跃时，LRU 表现更好，实现也更简单。访问模式频繁变化时，LRU 可能因为一次偶然访问而保留本应被淘汰的节点，LFU 的频率统计能更准确地区分冷热数据。代价是 LFU 需要额外的频率计数开销，且历史累计频率在新模式出现后可能产生误导。SGLang 默认使用 LRU，在大多数生产场景中效果足够好。
6 一个可迁移的工程判断
缓存匹配粒度的选择是一个在系统设计中反复出现的问题。PagedAttention 和 RadixAttention 的分歧本质上是 block 粒度与 token 粒度的选择。
判断方法可以看数据的自然边界与管理单元是否对齐。如果数据天然以固定大小为单位（如磁盘扇区、数据库页、网络数据包），粗粒度方案简单高效。如果数据长度变化大、语义边界不规则（如自然语言文本、变长协议消息），细粒度匹配能避免对齐损失，值得引入更复杂的数据结构来管理。
这个判断在数据库 buffer pool 的页大小选择、文件系统的 block size 调优、内存管理中的大页与小页选择中反复出现。关键不是粗好还是细好，而是先确认数据的自然边界在哪里，再决定管理单元的粒度。