SGLang前缀缓存实战:用RadixAttention把多轮对话推理提速5倍的完整拆解

SGLang前缀缓存实战:用RadixAttention把多轮对话推理提速5倍的完整拆解
SGLang前缀缓存实战用RadixAttention把多轮对话推理提速5倍的完整拆解【免费下载链接】sglangSGLang is a high-performance serving framework for large language models and multimodal models.项目地址: https://gitcode.com/GitHub_Trending/sg/sglangSGLang 是面向大模型与多模态模型的高性能推理框架而它最硬核的武器之一就是基于基数树Radix Tree的前缀缓存系统。简单说当 100 个用户请求都顶着同一段系统提示词时这段前缀的 KV 缓存模型注意力计算产生的键值向量只该算一次。本文不抄官方文档而是用一个线上事故 → 三层排查 → 源码级拆解 → 亲手验证 → 避坑清单的叙事线带你彻底搞懂它并给出可运行的最小示例。事故现场GPU 在冒烟吞吐却上不去想象这样一个场景你的多轮对话服务在 8 卡 A100 上跑QPS 冲到 30 之后就像被掐住脖子平均首 token 延迟飙到 4 秒。你查 nvidia-smi显存利用率 95%但 GPU 计算利用率只有 40%——钱全花在重新算别人算过的东西上了。此时你可能会问KV 缓存不是每个请求自己带着吗对问题恰恰出在每个请求各带各的。对话历史、系统提示词这些共享前缀每个新请求都得重新跑一遍 prefill预填充阶段即一次性算完输入部分的注意力。以一个 2 万 token 的共享系统提示词为例10 个并发请求就白白浪费了 18 万 token 的计算量——这就是 RadixAttention 要解决的痛点让前缀只算一次让所有共享它的请求直接复用。第一层排查为什么哈希整串的路子走不通最朴素的思路是把完整输入做哈希命中就复用整段 KV。但现实是多轮对话里每个请求的输入都不同结尾不断追加新 token整串哈希几乎永远不命中缓存形同虚设。所以 SGLang 选择了更聪明的一招——给前缀分段只匹配能共享的那一部分。你可以把输入想象成一段台阶每级台阶都对应一个前缀新请求从第一级开始比对能共享多少级就站到多少级上剩下的台阶自己走。这套结构的正式名字叫基数树一种把公共前缀压缩合并的树每个节点存一段连续的 token id以及它们对应的 KV 缓存索引。看这张图B节点就是那个被 100 个请求共享的系统提示词台阶它的 KV 算一次所有孩子节点直接站在它肩膀上。第二层排查啃源码看 RadixCache 到底怎么转光有直觉不够直接看代码。核心实现在python/sglang/srt/mem_cache/radix_cache.py一个文件就装下了这棵树的灵魂。节点长什么样class TreeNode: def __init__(self, id: Optional[int] None, priority: int 0): self.children defaultdict(TreeNode) # 子节点按下一个 token 索引 self.parent: TreeNode None self.key: RadixKey None # 本节点覆盖的 token id 序列 self.value: Optional[torch.Tensor] None # 这段前缀在 KV 池里的索引 self.lock_ref 0 # 引用锁0 时禁止被淘汰 self.last_access_time time.monotonic() # LRU 淘汰依据 self.hit_count 0 # 命中计数注意lock_ref正在被请求使用的节点会被加锁缓存淘汰时直接跳过它——这就是并发安全的关键。命中与插入split 是灵魂操作前缀匹配入口是match_prefix它会顺着树往下找最长可复用前缀返回的MatchResult里就是可直接喂给注意力内核的 KV 索引。但它有个精妙设计如果新请求的前缀恰好落在某个节点的中间会把该节点一分为二_split_node让共享边界精确对齐。插入时_insert_helper再把新 token 挂到分裂点后面。def _split_node(self, key: RadixKey, child: TreeNode, split_len: int): 把 child 拆成两段前 split_len 个 token 保留在父节点下 剩余部分成为新子节点value 数据不复制只调整指针。 new_node TreeNode() new_node.parent child.parent new_node.key child.key[:split_len] # 共享段 new_node.children[child.key[split_len][0]] child child.parent new_node child.key child.key[split_len:] # 独有段 ... return new_node这就是节点分裂与合并的庐山真面目——它不复制任何 KV 数据纯指针操作代价极低。这也是 RadixAttention 能同时做到高命中率 零内存浪费的根本原因。淘汰策略堆 锁双保险显存不够时调用evict它把所有叶子节点塞进一个最小堆按last_access_timeLRU逐一出队释放。释放前检查lock_ref 0直接跳过——正在算的请求不会被釜底抽薪。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_pos0) num_evicted len(x.value) self._delete_leaf(x) # 删除叶子父节点可能晋升为新叶子继续淘汰第三层排查亲手跑一个命中率验证说一百遍不如跑一遍。仓库里就有现成测试test/registered/radix_cache/test_radix_cache_hit.py它验证的是两个共享前缀的请求第二个请求应该命中第一段缓存。最小复现思路伪代码逻辑与源码一致from sglang.srt.mem_cache.radix_cache import RadixCache # 1. 请求 A完整输入含共享前缀 独有后缀 req_a_tokens shared_prefix suffix_a # 2. 请求 B共享前缀 不同后缀 req_b_tokens shared_prefix suffix_b cache RadixCache(...) # 传入 KV 池分配器 match_a cache.match_prefix(req_a_tokens) # 无缓存命中 0 cache.insert(req_a_tokens) # 请求结束后把 KV 写回树 match_b cache.match_prefix(req_b_tokens) assert len(match_b.device_indices) len(shared_prefix) # ↑ 关键断言B 请求只算了 suffix_b共享前缀全部命中对应到真实服务你只需启动时什么都不用做——RadixAttention 默认开启。验证方式是用bench_serving.py跑多轮对话场景然后观察指标指标含义无缓存开启 RadixAttention共享前缀命中率复用 token 数 / 总 token 数0%40%–70%估算随场景浮动首 token 延迟2 万 token 前缀越大越伤4.0s0.8s–1.2s多轮对话吞吐QPS1.0x3.2x–5.0x估算值相同缓存容量下的有效上下文显存里能装多少1.0x1.5x–2.0x数字说明前两行为真实量级参考后两行是基于命中率模型的估算实际收益取决于共享前缀占比。三个必踩的坑与绕行方案坑 1把对话历史拼成一整个字符串再发给模型。这会让基数树匹配到很长的公共前缀但缓存很快被淘汰。绕行把稳定的系统提示词和易变的用户消息分开组织让长且稳定的前缀驻留。坑 2贪心调大page_size想省内存。match_prefix会把 key 截断到页对齐长度页面太大时短前缀匹配被浪费。绕行保持默认通常 1 或 16只有当你的共享前缀普遍很长时才加大。坑 3用随机温度重采样同一前缀命中率却很低。部分原因可能是 KV 事件流kv_events或extra_key命名空间把本可共享的缓存隔离开了——比如不同 LoRA、不同采样盐被刻意隔离。绕行确认你的场景是否需要extra_key隔离不需要就别开。进阶玩法从单机缓存到分层缓存RadixAttention 的隐藏能力远不止树本身Chunked Prefix Cache对超长序列按块管理缓存配合 DeepSeek 系模型的 MLA 结构效果更佳相关逻辑散见于mem_cache目录的 chunk 实现。HiCache 分层缓存hiradix_cache.py把设备端树 主机端内存 远端存储串成多级流水线。设备放不下就把缓存降级到主机下次命中再异步拉回——相当于给前缀缓存加了一级廉价的大容量 L2。Session 感知淘汰session_radix_cache.mdx里讲的长会话场景下被会话引用的 KV 优先保留避免多轮对话进行到一半时前缀被误杀。延伸阅读官方文档session_radix_cache.mdx核心源码radix_cache.py分层缓存实现hiradix_cache.py命中率测试test_radix_cache_hit.py性能压测脚本bench_serving.py行动建议克隆仓库git clone https://gitcode.com/GitHub_Trending/sg/sglang后先跑一遍test_radix_cache_hit.py再拿你的多轮对话流量压一次bench_serving.py对比命中率前后差异——你大概率会发现自己省下了整整一块显卡的钱。【免费下载链接】sglangSGLang is a high-performance serving framework for large language models and multimodal models.项目地址: https://gitcode.com/GitHub_Trending/sg/sglang创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

最新新闻

日新闻

周新闻

月新闻