Token的一生:RadixCache——Token前缀共享的秘密

admin 2026-08-21 05:13:54 网络安全文章 来源:ZONE.CI 全球网 0 阅读模式

文章总结: 本文深入解析SGLang中RadixCache的核心机制,通过前缀树实现KVCache共享,显著提升大模型推理吞吐量2-5倍。文章详细介绍了TreeNode和RadixKey数据结构、matchprefix的指数搜索与二分查找加速、节点分裂、lockref引用计数保护、多种驱逐策略以及page对齐等关键技术,并展示了在RAG场景下的实际效果。 综合评分: 88 文章分类: AI安全,技术标准,解决方案


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_lenkv_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 层面做了什么?


免责声明:

本文所载程序、技术方法仅面向合法合规的安全研究与教学场景,旨在提升网络安全防护能力,具有明确的技术研究属性。

任何单位或个人未经授权,将本文内容用于攻击、破坏等非法用途的,由此引发的全部法律责任、民事赔偿及连带责任,均由行为人独立承担,本站不承担任何连带责任。

本站内容均为技术交流与知识分享目的发布,若存在版权侵权或其他异议,请通过邮件联系处理,具体联系方式可点击页面上方的联系我

本文转载自:DCOS zouyee zouyee《Token 的一生:Radix Cache —— Token 前缀共享的秘密》

京东-安全研发面经 网络安全文章

京东-安全研发面经

文章总结: 本文分享了京东安全研发岗位的面试经验,涵盖面试流程、技术考察重点及准备建议,为求职者提供实战参考。 综合评分: 75 文章分类: 实战经验京东-安全
评论:0   参与:  0