你问我为什么想不开非要用 Rust 重写一个 SGLang 里的运行时。其实不是想不开是跑线上服务跑出来的怨念。SGLang 的核心价值我很认可radix cacheRadixAttention解决长 prompt 的 KV 缓存复用、结构化输出约束解决 JSON 生成翻车的折磨、调度器把长短不齐的推理请求压进同一个 batch。但作为部署过几十个并发实例的人我越来越清楚一件事想法再好落到 Python 过程里实现层到处都是损耗。SGLang Rust 运行时不是要替代上游整个项目而是把最疼的几条热路径——前缀缓存、约束机、批调度、token 级处理——下沉到 Rust 里把 Python 当成总线而非引擎。这篇文章写给两类人一类是在思考“我是不是也该把推理中间层重写一遍”的研发另一类是纯好奇底层原理、想趁机学线程级系统设计的同学。我不会卡在脚手架阶段会直接拆到细节radix 树为什么用VecNode而非BoxNode、DFA 约束 每次解码怎么不掉链子、批量调度怎么把 Linux 内存门限定住。看完你能拿去 mock 真实方案也能避开我已经踩平的坑。1. SGLang 运行时里到底藏了什么为什么值得用 Rust 重写一遍1.1 SGLang 要做的事先说大白话。SGLang 这类推理框架解决两个核心问题第一LLM 的 KV cache 能不能尽量复用让“聊过的前缀”不再重新算第二输出能不能被约束成你要的结构而不是让模型自由发挥。围绕这两件事还需要调度器、tokenizer、显存管理、预热、日志、并发接入这些周边设施。其中真正的复杂度集中在三个组件radix attention 缓存树每条请求的前缀按 token 拆开共享前缀尽量复用遇到分叉点再生成新节点空间不够时就淘汰。结构化输出引擎把正则表达式或者 JSON Schema 编译成自动机在解码的每一个 token 上前向检查只允许那些“不会让最终结果偏离约束”的 token 进入候选集。批调度器把不同到达时刻、不同输出长度、不同优先级的请求揉进同一个 batch同时还要管理 KV cache 的占用和逐出。这三个组件在 Python 里能跑但有非常现实的性能天花板。1.2 Python 在哪些地方卡脖子Python 不是不能做高性能它更适合做“组织者”不适合做“执行者”。我列出真正让人头大的四个点对象模型开销。一整棵 radix 树如果每个节点都是 dict每个 cache 命中都要做多次哈希查表、引用计数、深拷贝这种开销在 token 级逻辑里会被放大几百倍。GIL 的存在。同一时间只有一个线程跑 Python 字节码token 校验、树更新、队列调度如果都压在 GIL 下合起来就会吃光 CPU 转圈时间。二进制数据的搬运。token id 数组、attention mask、cache 统计指标全部以 Python list 和 numpy 数组来回倒numpy 虽然快但数据进进出出要来回拷贝。调度和约束的一层层函数调用。一次解码要过is_allowed_token()、state_transition()、check_peak()这类函数函数本身不慢慢的是每一层都走一遍 Python 虚拟机和属性查找。所以 Rust 的收益就很清楚了用接近 C 的执行性能去跑这些高频小块。好比你以前是让数学博士每个数字都按计算器现在直接把公式拍到他脑子里。1.3 Rust 的收益边界别指望它让模型跑得更快这里必须先泼一盆冷水。Rust 重写运行时不会让 GPU 上的 Transformer 算子变快。矩阵乘法、flash attention、多卡通信这些真正吃时间的大头不在我们的掌控范围内。Rust 的价值在“刀口以外的地方”框架开销更低CPU 侧吞吐提升 15% 到 30%。长 prompt 多轮对话时缓存命中率和缓存回收粒度更细KV 显存使用更可以有效。结构化解码峰值更稳P99 延迟比 Python 实现有明显改善不会因为 GC 或 GIL 而抽风。如果把整条服务链路比作流水线GPU 是压路机CPU 侧是流水线入口处的分拣员。分拣员再快压路机也不能打得更快但如果分拣员双手在抖整条线就一定快不了。2. 整体架构把 Python 留在门口热路径交还给 Rust2.1 模块划分引擎、缓存、调度、约束、桥接动手之前我先把软件切块列清楚避免写到一半在模块边界上迷路。模块职责语言重要程度scheduler请求排队、batch 组、优先级处理Rust核心radix_cache前缀匹配、节点复制、引用计数、逐出Rust核心constraint_engine正则/JSON Schema 编译、token 掩码Rust核心tokenizer_bridgetokenizer 调用、token id 映射Rust/Python必要py_bridge与 Python 推理框架交互、触发执行PyO3必要monitor指标统计、日志上报Rust辅助我特别想强调一个取舍不要把所有东西都下沉到 Rust。例如显存分配策略需要和具体的部署环境、GPU 驱动调度器联动这块放 Python 里反而灵活。tokenizer 的一些特殊清理逻辑也会不断调整如果全部用 Rust 实现每次改动都变成编译和发版成本。为了一个高频、一个低频尽量把 Rust 限制在热路径内部。2.2 边界设计只让低频操作留在 Python 侧日常开发里最常见的错误是把 Rust 与 Python 的分界线划在以“函数”为单位的地方而不是以“请求”为单位。这会导致每一个操作都要走一遍图省的跨语言调用成本反而更高。我的设计是粗粒度跨语言调用submit_session(context_ids, req_id)把一次整请求的所有输入一次性送给 Rust。step()批次内推理一步触发一次完整的缓存查询、约束检查、调度决策。commit_tokens(tokens)推理框架输出 tokens 后再一次性回写缓存状态。这样一次请求只有几个跨语言边界而不是每个 token 都跨过来跨过去。2.3 数据流的一条清晰链路你可以把数据流想象成一条运输带Python 侧接收到 HTTP/WS 请求进行鉴权、session 管理。把 prompt token ids、请求元数据包成一个批次通过 PyO3 的submit_session给 Rust。Rust 执行查 radix 树 - 得到复用序列和补算序列 - 构造 batch - 约束引擎生成候选掩码 - 更新缓存统计。结果回到 PythonPython 侧只需要调用模型后端把 batch 跑掉。模型后端返回 logitsPython 再回传给 Rust 做 token 选择 和 状态推进。这套结构的好处是Rust 核心可以独立测试、独立压测。我写单元测试时甚至不需要 Python直接把 token 序列灌进 Rust 的 CLI 试验程序几分钟就能验证正确性。3. RadixAttentionprompt 缓存树的 Rust 实现3.1 为什么前缀是唯一值得信的缓存甜点LLM 推理里最贵的部分是早先的 prompt 处理和 KV cache 计算。多轮对话中后面的所有轮次都能复用前面的 KV cache。SGLang 的方案是说把每次请求的完整 token 序列按当前请求切分组合成一颗 radix 树树上每个节点代表一段共享前缀。请求 A 是 “我喜欢猫”请求 B 是 “我喜欢狗”两个请求共享 “我 喜欢” 这个前缀那么这段内容对应的 KV cache 只需计算一次。在 Rust 里的第一件事就是把树节点设计成内存紧凑的结构。3.2 用索引指针而非对象句柄的内存布局如果你直接写BoxNode建立树你会很快在可变借用上跟借用检查器打起来。我推荐用一个VecNode当对象池然后用u32索引代替指针。Rust 的Vec在堆上是连续内存遍历一级孩子节点时可以直接利用 CPU 缓存行。节点长这样#[derive(Debug, Clone)] struct RadixNode { parent: Optionu32, children: Vec(u32, u8), // (child_node_index, 首 token 的字节) kv_ref_count: u32, // 引用计数决定能否被逐出 last_accessed: u64, // 访问时间戳LRU 淘汰用 token_span: (u32, u32), // 本节点在总 token 序列中的起始和结束位置 }为什么要用(child_node_index, 首 token)这种组合因为查找某个前缀的孩子时我们期望快速按首 token 二分或哈希。如果直接拿VecNode存每次找孩子都要扫一遍太浪费。更优化的做法是孩子存在一个HashMapu32, Vecu32里首 token 作为 key。如果 Prompt 很长且孩子非常多则可以考虑直接上BTreeMap但实测里HashMap在几百个孩子时已经很够用。3.3 匹配、分叉与逐出的核心逻辑假设来了一个请求要查询前缀[我, 喜欢, 猫, 的, 毛]。整个过程从树的根节点出发查看 children 中是否存在[我]节点。存在则命中。再找这个孩子节点下是否有 token[喜欢]存在继续往下。到第三个 token[猫]时发现原本节点只到[喜欢]这里自己有两个孩子[猫]和[狗]。检查[猫]节点存在继续到 token[的]发现[猫]节点下没有对应孩子于是创建新节点承接余下部分的 token。这里的关键优化当一条长请求只命中了前 80% 时我们尽量不要打碎原有节点。处理方式是节点分裂把原有[猫]节点分裂成[猫]和[的]两个节点前者是后者的父亲。这样新请求和已缓存请求都能保留最大公共前缀。逐出策略上我用的是引用计数 时间戳双保险每个节点初始kv_ref_count 1当一个请求的缓存被引用时计数加 1请求释放时减 1。只有计数为 0 的节点才参与 LRU 淘汰。LRU 我以为用时间戳排序最直观扫描全部节点太慢于是我维护一个按last_accessed排序的 BTreeMapBTreeMap(u64, u32), u32用于快速定位最早未被引用的节点。实测一个 100 万节点的树淘汰时只要遍历映射表前几十项就够腾出内存了。淘汰的最大坑是“级联逐出”。如果一个父节点被淘汰所有子节点都失去了缓存意义必须一并释放。此时不是一个个改计数而是通过 DFS 找到整个子树后统一回收 KV 显存。我曾经在这里跳过子节点递归结果显存泄漏像漏水的水管排查了一整天才在审计日志里看到大量孤儿节点。3.4 借用检查器在树上如何唠叨写 radix 树的 Rust 代码时我大概率被 Rust 编译器教做人了VecNode存索引最大的麻烦是有时你想同时修改 node 和它的 parent。例如分裂节点时既要改原节点内容又要改父节点的 children。如果用可变引用同时访问两个元素借用检查器会直接拍桌子。我的解法是“先取索引再做批量操作”。分裂操作不直接通过引用改而是fn split_node(mut self, node_idx: u32, split_offset: u32) - u32 { let old_tail_span { let node self.nodes[node_idx]; (node.token_span.0, node.token_span.1) }; // 先构建新节点不改旧节点 let new_idx self.new_node(split_offset, old_tail_span.1, node_idx); // 再把旧节点的尾部 span 缩短 let old_node mut self.nodes[node_idx]; old_node.token_span.1 split_offset; new_idx }借用检查器其实是在保护你它逼你把“读旧数据”和“写新数据”分成两个阶段才避免了那种你在 C 语言里经常踩的迭代失效 bug。4. 结构化输出引擎从正则到基于 token 的 FSM4.1 解码约束不是要什么给什么很多人的第一反应是“要 JSON 就加个‘请输出 JSON’的 prompt 呗。” 但实际效果和在悬崖边跳舞差不多模型一激动就丢了右括号。SGLang 的约束引擎做法是硬约束每当你得到一个 logits 分布只有在约束自动机上处于“接受态”的 token 才允许被采样。这样模型只能往前走合格路径绝不允许跑到结构外。Rust 实现的核心任务是“约束自动机的状态推进”和“token 掩码计算”。4.2 自动机设计与 token 级状态推进最基础的做法是把正则表达式编译为 NFA/DFA一个字符一个字符地走。但 LLM 的 tokenizer 并非按单字符切分一个 token 可能是JSON也可能是“还可能带空格。因此我们不能只检查一个字符而要对每个候选 token 的完整字符序列进行自动机转移。面向 tokenizer 的设计是这样的为每一个 DFA 状态预先构建一个“允许字符前缀”二分表。解码时拿到一个候选 token把它拆成字符序列然后从当前状态出发依次转移。只要全部字符都在允许路径上且最终 DFA 状态是非失败态是接受态或可能扩展的中间态就保留这个 token 的采样概率。更进一步的加速是把 token 压缩成“前缀类型”。很多 token 的第一字符相同、中间字符相同其实没必要逐个字符跑。我们可以给自动机做一次离线缓存struct StateTokenTransition { state_id: u32, token_id_list: Vecu32, // 此状态下允许的所有 token next_state: Vecu32, // 每个 token 转移到的目标状态 }推理后端在采样前只需要查这张表。表太大时可以压缩成“分块索引”先按 token 的首字符建立 256 个桶再用二分查桶内的子串。实测用这种“首字符桶 子串二分”的方式128k 词表的扫描时间从几百微秒降到几十微秒。4.3 长思维链CoT场景下的约束拆解线上做 agent 的时候经常遇到模型要先输出一大段“思考过程”再输出最终 JSON。如果整个输出过程都拿正则约束住思考过程就废了模型会像被风湿一样一句话都不敢多说。我的处理方式是把整个输出分段约束前段只约束“不能有小括号未闭合”“不许提前输出 JSON 关键字”后段进入严格 JSON 模式。分段之间不共享同一个自动机但共享 token 掩码缓存。这个“分段约束”在 Rust 的实现上就是给引擎传入一个VecConstraintSegment每个 segment 有自己的编译表和起始状态。调度时引擎根据已生成的 token 数切换 segment。代码上很清爽pub enum ConstraintSegment { Free, // 无约束仅过滤不可见字符 JsonSchema(JsonSchemaId), // 严格模式 Delay(RegexId), // 宽松正则 }4.4 JSON Schema 与正则的落地差异正则做自动机容易JSON Schema 做起来就复杂。因为 JSON Schema 有嵌套对象、数组、枚举、条件 if/then/else。直接套正则等于重写一个 JSON 解析器。SGLang 原本的处理是把 Schema 转化成“结构自动机”每次解码时按当前节点类型对象、数组、字符串、数字、布尔值、null决定允许哪些 token。Rust 这边我采用“栈式 FSM”而不是单状态 FSM。每进入一个对象或数组压栈一个新的状态机每完成一个 token从顶部分弹出。比如enum JsonFSMState { StartObject, ObjectKey(StringState), ObjectValue(ValueType), ArrayElements(ValueType), EndValue, }每次采样结束后不只在单一状态上推进还要更新一个状态栈。这种实现没有正则那么“公式化”但是更可控。遇到复杂的嵌套结构时用内存空间换逻辑正确性非常划算。5. 调度、批处理与 KV 管理的协同设计5.1 为什么需要分布式 token 预算推理服务最理想的状态是批处理里的每个请求都差不多同一时刻结束。可惜现实世界永远有人闲有人忙一批请求里可能有一个在写长篇小说另外三个只需要一个词。如果不去管它长请求会无限占用 batch 位置短请求吃不到显存整个吞吐被拖垮。调度器的核心指标是“系统总 KV 内存预算”。我的设计不是每个请求独占固定显存而是有一个全局预算。Rust 调度器在每一轮step()开始前对当前批次里的所有请求做 token 配额分配。分配逻辑分两层weighted round-robin每个请求根据其还剩多少预估输出长度按权重去占 batch 的容量。KV 上限保护一旦总 KV 内存快满了就优先淘汰引用计数为零且已完成主体输出的节点。这个设计的价值在于模型后端显存不够时不是直接 OOM而是提前把一些“没说几句话但对缓存无贡献”的请求挤出 batch显存可以立刻被新请求复用。5.2 引发“过期缓存”与“重新验证”博弈如果调度器太激进把某些请求的 KV 当缓存逐出了那么这一轮的输出完了之后后续继续生成时又得重新计算前缀。这叫做“缓存过期”。没人在意这个词但代价是最直观的重新计算可能比省下的时间还要贵。所以我在调度时做了一次“重新验证”。当一个请求的 KV 节点引用计数为 0 且将要被逐出但该请求尚未结束生成调度器会先把该请求标记为 high_priority。这一轮绝不驱逐它。所谓 high_priority 判定是参考当前 batch 里所有请求的总剩余长度以及驱逐后重新计算的预估开销。这里的关键经验是不要只看完成比例要看剩余 token 数。一个已经完成 80% 但还剩 2000 token 的请求它的缓存价值远比一个完成 20% 但只剩 100 token 的请求要低。我编码成了cache_value_score。5.3 batch 突变一个请求中途插队怎么办假设当前 batch 已经填满来了一个新请求它的前缀恰好全部命中缓存。此时把它加进 batch 几乎零成本甚至能提高整批收益。调度器需要支持“前缀命中请求免排队插队”。具体实现检查新请求的前缀 radix 命中率如果命中率超过 80%进入 preempt queue。如果 batch 中有一个请求已处于收尾阶段只剩很少 token直接从 batch 中把它替换为 preempt queue 中的请求。用一次cache_value_score对比来保证插队不会亏本。这个场景很适合用 Rust 的crossbeam_channel来做队列管理把调度决策和推理执行破成两个独立线程。调度线程在一个无锁队列上取请求、做缓存查询、组 batch推理线程在 GPU/推理后端上执行。并行度高了延迟抖动自然也低。6. 通过 PyO3 把 Rust 核心接回 Python6.1 薄封装层的应用方式写 PyO3 绑定的时候我强烈建议把它做成“门后面的一层薄箱子”。不要在绑定层写复杂逻辑只做三件事数据转换、函数转发、类型转换。我用 maturin 管理项目pyproject.toml只需要一个 module 定义比如[build-system] requires [maturin1.0,2.0] build-backend maturin然后在 Rust 里定义#[pymodule] fn sglang_rs(m: Bound_, PyModule) - PyResult() { m.add_class::RadixCache()?; m.add_class::Scheduler()?; m.add_class::ConstraintEngine()?; Ok(()) }Python 侧拿到的是一个sglang_rs.RadixCache对象可以直接在业务代码里调用。但我要强调数据交换的类型尽量用Vecu32、Veci64这类原生数字数组通过PyBytes或numpy的 buffer 协议传递。它是最高效的桥接不会因为 Python int 和 Rust i64 的转换而拖后腿。我的做法是在 Rust 里导出add_sequence(seq_bytes: [u8]) - u64其中seq_bytes是打包成字节数组的 token ids。然后在 Python 侧用numpy.tobytes()生成这个输入一次性投喂。结果一次性拉回。6.2 粗粒度锁和多线程注意的点PyO3 绑定最大的问题不是性能而是 GIL 和线程安全。如果你在每个函数调用里都打开 Python 的 GIL那多线程的优势就没了。我的设计是Rust 核心持有一条内部消息队列Python 侧在给队列发请求后立刻释放 GILRust 调度线程在后台处理。Python 侧只在以下几个时机重新获取 GIL步骤结束需要拿 logits 结果时。缓存统计需要上报给监控时。有请求需要从 Python 侧获取元数据时。这样就很少会在 Rust 和 Python 之间来回横跳也不会把 GIL 拽进高频循环。不过要提醒的是PyO3 底层对象的生命周期要格外小心。如果你把某个 Python 对象传给 Rust 并留在 Rust 侧长期保存必须确保引用的持有不会造成 Python 侧对象无法回收。我大多数时候只传bytes和数字避免把复杂 Python 对象传给 Rust省得后面内存泄漏和调试到怀疑人生。7. 我踩过的一些坑与问题排查实录7.1 调优时最容易犯的低级错误现象根因解决方案缓存命中率上去了但整体吞吐不变逐出过于保守KV 显存一直顶在峰值调小缓存窗口允许动态逐出单条请求延迟很低整体批吞吐崩了每个 batch 都在等待最慢请求增加长输出请求的超时驱逐逻辑约束引擎把正确 token 拒掉了自动机状态转移实现里用了全局 BFS改成栈式 FSM压力测试后正确性提升Rust 核心不崩溃但 Python 进程暴涨缓存节点引用计数泄漏用kv_ref_count审计日志查每个节点增减是否成对7.2 做性能压测时注意别被“假快”骗了所有人都喜欢拿一条 100 token 的请求说“咱们延迟好低”。真实生产环境根本不是这样。我做压测时一定会混合三类请求超短输出、普通输出、超长 CoT 输出。并且把请求的间隔调成不均匀泊松分布模拟多个 App 用户同时发消息的节奏。更有意思的是默认情况下 Python 版本里结构化解码会非常吃 CPURust 版本的结构化约束时间可以从每 token 1ms 降到 0.1ms 以内但整体 GPU 前后效果不大。这个结论一点都不稀奇GPU 是整个系统的瓶颈CPU 侧省下的时间必须用来做“更多请求的排队”才能变成吞吐收益。别在压测报告里炫耀“约束引擎快了 10 倍”要炫耀“在相同 GPU 显存下, P99 延迟降了 40%”。7.3 树结构到底要不要内存池我一开始觉得VecNode就足够后来发现高并发下节点增删过于频繁Rust 默认分配器在高频 alloc/dealloc 时还是有点抖。优化方案是把节点回收放到一个 free list 里删除节点时不释放内存只是把节点编号压回空闲队列新节点分配时优先从空闲队列里拿。fn free_node(mut self, index: u32) { self.nodes[index].children.clear(); self.nodes[index].parent None; self.free_list.push(index); }这个改动让高性能场景下的分配抖动了大约 20%。同样的思路也用在 token 缓冲区上把频繁扩容的Vecu32预先分配好后续直接覆盖写。8. 我个人在实际操作中的体会整套折腾下来最大的感受是Rust 不会自动让系统变快但会强制你把性能问题想清楚。树、自动机、调度这些概念在 Python 里写出来很容易可一上负载就露馅因为没有人在意每一层调用到底多贵。用 Rust 重写一遍相当于把整条链路拆开、擦拭、再装回去很多以前忽略的设计缺陷都会清晰暴露出来。如果你也想做类似的事情我给的建议是不要上来就写引擎先写一个 300 行左右的 radix 缓存原型。让它在纯 CPU 环境下能吃下几万条短请求观察内存布局和 cache 命中率。等这一关过了再往里加结构化解码。一次只加一件事每件事都配好 metrics。相信我你会很快在调试日志里看到惊喜的。