Token 的一生：Radix Cache —— Token 前缀共享的秘密
原创
zouyee
zouyee
DCOS
2026年8月20日 08:00
上海
在小说阅读器读本章
去阅读
系列定位：进阶。理解 SGLang 的核心优化：基于前缀树的 KV Cache 共享，以及 RadixKey 的设计细节。
Summary
本篇揭示 SGLang 高效复用 KV 计算的核心机制。内容涵盖：为什么大量请求共享前缀时 Radix Cache 能带来 2-5x 吞吐提升；TreeNode 和 RadixKey 的数据结构设计；match_prefix 的指数搜索+二分查找加速；节点分裂（split）保证精确匹配边界；lock_ref 引用计数保护正在使用的 KV；LRU/LFU/FIFO/Priority 等驱逐策略；extra_key 实现 LoRA 命名空间隔离；page 对齐保证 KV 分配原子性；以及 Session Radix Cache 对多轮对话的支持。
1. 为什么需要 Radix Cache？
大模型推理的一个常见规律：大量请求共享相同的前缀。
请求 A: [system_prompt] + \"Tell me a joke\"
请求 B: [system_prompt] + \"What is the capital of France?\"
请求 C: [system_prompt] + \"Explain transformers\"
system_prompt 可能有几百甚至几千个 token。如果每个请求都重新计算一遍这段 KV，是巨大的浪费。
Radix Cache 的核心思想：用一棵\\前缀树（Trie/Radix Tree）\\管理所有已计算的 KV Cache，相同的 token 前缀只计算一次，后续请求直接复用。
2. 数据结构
TreeNode
前缀树的每个节点对应一段 token 序列：
class TreeNode:
key: RadixKey # 这段 token 序列（从父节点边到当前节点）
value: Optional[torch.Tensor] # 对应的 KV Cache 槽位索引（device，GPU 上）
host_value: Optional[torch.Tensor] # 对应的 CPU 内存槽位索引（层级缓存）
children: defaultdict(TreeNode) # 子节点（defaultdict）
lock_ref: int # 引用计数（被几个请求正在使用）
last_access_time: float # 最近访问时间（用于 LRU 驱逐）
priority: int # 用于 priority 驱逐
hash_value: Optional[List[str]] # 内容哈希（用于验证）
RadixKey
RadixKey 封装了一段 token ID 序列，是前缀树的查找键：
class RadixKey:
token_ids: array[int] # token ID 序列
extra_key: Optional[str] # 命名空间标签（如 lora_id、cache_salt）
is_bigram: bool # EAGLE 投机解码用的双 token 视图
limit: Optional[int] # 截断上限（O(1) 切片，避免拷贝）
extra_key 的设计很巧妙：相同 token 序列但不同 LoRA 适配器的请求，前缀完全不同，不应共享 KV。通过 extra_key 把它们隔离在不同的命名空间里。
3. 前缀匹配：match\_prefix
当一个新请求到来，调度器调用 match_prefix 找最长已缓存前缀：
# python/sglang/srt/mem_cache/radix_cache.py
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
\"\"\"
返回：
device_indices: 已缓存的 KV 槽位索引（torch.int64 tensor）
last_device_node: 命中的最深树节点
\"\"\"
匹配过程（从根节点开始）：
图 1：架构与流程示意（点击图片查看原图）
\\节点分裂（split）\\是 Radix Tree 的核心操作：当匹配在节点中间结束时，把节点一分为二，暴露出精确的匹配边界，提升后续匹配效率。
快速前缀比较
对于很长的共享前缀，逐 token 比较很慢。SGLang 用指数搜索（exponential search）+ 二分查找加速：
def match(self, other: RadixKey, page_size: int = 1) -> int:
\"\"\"找到两个 key 的最长公共前缀长度。\"\"\"
# 指数跳跃：1, 2, 4, 8, 16, ... 直到找到第一个不同位置的区间
# 再在区间内二分搜索
# 每次比较是 C 级别的 array slice 比较，不走 Python 循环
对于 10,000 token 的共享前缀，这把 O(n) 的逐 token 比较变成 O(log n) 的跳跃比较。
4. 前缀插入：insert
请求完成后，把已计算的 KV 插入缓存：
def insert(self, params: InsertParams) -> InsertResult:
# key: 完整的 token 序列
# value: 对应的 KV Cache 槽位索引
# 沿着树路径插入，遇到已有节点合并，新节点添加
注意：不是每个请求的每个 token 都会被插入。SGLang 有以下限制：
- 只插入到
cache_protected_len（kv_committed_len的一部分），thinking 阶段的 KV 可能被排除（strip_thinking_cache） - 正在被其他请求使用（
lock_ref > 0）的节点不会被驱逐
5. 引用计数与驱逐
lock\_ref（引用计数）
每当一个请求命中某个节点并开始使用其 KV，lock_ref 递增；请求结束时递减。
lock_ref > 0 的节点是\"受保护\"的，不能被驱逐。
# 请求开始使用前缀：
def inc_lock_ref(self, node):
node.lock_ref += 1
# 从可驱逐集合中移除
# 请求结束，释放前缀：
def dec_lock_ref(self, node):
node.lock_ref -= 1
if node.lock_ref == 0:
self.evictable_leaves.add(node) # 加入可驱逐集合
驱逐策略
当 KV Cache 满了需要腾出空间时，从可驱逐的叶节点开始：
# 支持多种驱逐策略，通过 --radix-eviction-policy 配置（默认 lru）
eviction_strategy = get_eviction_strategy(eviction_policy)
# \"lru\" — 最近最少使用（默认）
# \"lfu\" — 最不常使用
# \"fifo\" — 先进先出
# \"mru\" — 最近最多使用
# \"filo\" — 后进先出
# \"slru\" — 分段 LRU
# \"priority\" — 按 request priority 字段驱逐（--enable-session-radix-cache 必须用这个）
驱逐从叶节点开始，沿树向上，直到释放出足够空间。
6. RadixKey 的 page 对齐
GPU KV Cache 通常以 page（页） 为单位分配（类似虚拟内存的分页），默认 page_size=1 token，但也可以配置更大的 page（如 16 tokens），减少内存碎片。
def page_aligned(self, page_size: int) -> RadixKey:
\"\"\"截断到 page_size 的整数倍。\"\"\"
aligned_len = len(self) // page_size * page_size
return self[:aligned_len]
前缀匹配的结果会对齐到 page 边界，保证 KV Cache 分配的原子性。
7. Session Radix Cache
对于多轮对话，SGLang 扩展了 Radix Cache 支持 Session 概念：
图 2：架构与流程示意（点击图片查看原图）
同一会话的不同轮次都缓存在同一棵子树里，完整的对话历史只需计算一次。
8. 实际效果
在典型的 RAG（检索增强生成）场景下，假设：
- system prompt：500 tokens
- 每个请求的 prompt 中有 500 tokens 的相同检索文档
没有 Radix Cache：每个请求需要计算 1000+ tokens 的 prefill 有 Radix Cache（第 2 个及以后的相同请求）：跳过全部 1000 tokens 的 prefill，直接进入 decode
吞吐量提升可达 2-5x，TTFT 大幅降低。
小结
Radix Cache 工作流程：
新请求到来
↓
match_prefix(token_ids)
↓
返回 prefix_indices（已缓存的 KV 槽位）
↓
调度器设置 extend_range = [prefix_end, input_end]
↓
TP Worker 只计算 extend_range 范围内的注意力
↓
请求完成后，insert(token_ids, kv_indices) 存入缓存
| 机制 | 作用 |
|---|---|
| 前缀树（Radix Tree） | O(log n) 前缀匹配 |
| lock\_ref 引用计数 | 保护正在使用的 KV，安全驱逐 |
| extra\_key 命名空间 | LoRA/Cache-Salt 隔离 |
| page 对齐 | 保证 KV 分配原子性 |
| 指数搜索 + 二分查找 | 长前缀快速比较 |
下一篇：Token 生成的最后一步 —— 采样。温度、Top-P、重复惩罚……这些参数究竟在 token 层面做了什么？
版权声明
本站仅做备份收录，仅供研究与教学参考之用。
读者将信息用于其他用途的，全部法律及连带责任由读者自行承担，本站不承担任何责任。
评论