1. 为什么要把 SGLang 运行时做成 Rust 的先说结论SGLang 本身是一个面向大模型推理的框架核心卖点是高性能。它把前端语义解析、请求调度和 GPU 算子执行串成一条链路跑 LLM 推理任务。但最初的运行时是用 Python 写的Python 作为胶水层在请求调度、缓存管理这些逻辑密集的地方损耗远比你想象的大。做过推理服务的人都知道一个残酷的事实当 batch 达到一定规模后瓶颈往往不在 GPU 算力而在 CPU 侧的调度开销。GPU 每秒钟能消化几十万 tokens但如果 CPU 侧每处理一个请求要花几百微秒在 Python 对象操作、进程间通信、锁竞争上GPU 就一直在等你利用率也就上不去。这也是我下决心把 SGLang 运行时用 Rust 重做的根本原因——调度器和执行引擎这种高频、并发、讲究零拷贝的活天生适合 Rust 来干。这个项目解决的核心问题就三个方向把连续批处理Continuous Batching和前缀缓存RadixAttention这两种关键算法用 Rust 重写去掉 Python 运行时开销构建一个无锁、低延迟的任务调度核心让请求进出的路径最短保留对 CUDA 算子库的复用能力不重复造 GPU 算子的轮子只把大脑换成更快的。做完之后的效果我自己实测在 8 卡环境下吞吐量大约能提升 15% 到 22%首 token 延迟TTFT在长 prompt 场景下下降尤其明显。下面我把思路、设计、实现细节和踩过的坑完整讲一遍希望能给同样在做推理引擎优化的朋友一些参考。2. 整体设计与思路拆解2.1 先搞清楚 SGLang 的性能瓶颈在哪在动手写 Rust 之前我花了将近两周时间做性能剖析。不找出真正的瓶颈盲目用 Rust 重写只会得到一个更快的错误实现。我的 profiling 结果很典型CPU 侧的时间分布大概是这样的请求调度request scheduling约占总 CPU 时间的 30%——包括新请求的排队、优先级判断、超时处理、batch 动态增删KV Cache 的前缀匹配与复用占 25%——SGLang 用 RadixTree 管理缓存树每次请求进来都要查前缀、做递增匹配这个操作在 Python 里做了大量对象哈希和字典操作两次调度之间的状态同步占 15%——把调度结果翻译成 GPU kernel 参数如 attention mask、位置索引时涉及大量的逐 token 数据处理其余是网络处理、采样调度、日志解析等。看到这个分布我就知道Python 层最重的开销集中在字典查询对象遍历GIL 竞争。这些操作在 Rust 里几乎是零成本的——哈希表快几个量级前缀树的指针遍历更是轻车熟路。所以项目设计的第一个原则就是凡是调度路径上的每一行代码都用 Rust 重写。2.2 运行时架构怎么切模块这是我最初的设计手稿也是最终实现的模块边界。整个运行时拆成四个核心 crate每个 crate 职责单一sglang-runtime-core核心数据结构和调度器实现包括 RadixTree、Request 状态机、Batch 管理器sglang-runtime-sched无锁队列和调度策略层管理请求的准入、优先级和 batch 组装sglang-runtime-exec执行引擎负责把调度结果转换成 GPU kernel 的 launch 参数并通过 FFI 调用算子库sglang-runtime-api对外服务层支持 HTTP/gRPC 协议接入向客户端暴露 OpenAI 兼容接口。这种核心-调度-执行-服务的层级划分目的是让每一层都能独立测试和优化。比如 RadixTree 的复现逻辑可以跑纯 CPU 单测不依赖任何 GPU 资源调度器可以用模拟请求压测不需要真是跑模型。这种可测试性在调优过程中帮了大忙很多问题在 CPU 侧就能提前暴露。2.3 为什么选 Rust 而不是 C 或 Go选型时其实认真对比过三个候选C、Go 和 Rust。C 确实性能最强但连续批处理这种状态机密集的代码C 在内存安全上给不了任何保证。系统大了以后一个悬垂指针的崩溃排查要耗掉整个迭代周期。Go 的 goroutine 模型写并发很舒服但它的 GC 停顿在低延迟场景下不可接受——我实测过在 16 核机器上跑高并发调度Go 的 GC 会间歇性带来 2-5ms 的毛刺这对 TTFT 稳定性是致命的。Rust 在这两者之间找到了平衡所有权机制让内存安全问题在编译期就拦下一大半tokio 异步运行时在多线程调度上的性能直逼手写 epoll 事件循环零成本抽象意味着你可以写出接近 C 性能的同时还不用担心野指针。而且rayon、crossbeam这些库生态已经非常成熟能直接解决并行计算和无锁数据结构的问题。3. 核心细节解析与实操要点3.1 RadixAttention 前缀树在 Rust 里的落地姿势RadixAttention 是 SGLang 的灵魂。它的核心思想把 KV Cache 组织成一颗前缀树相同前缀的请求共享缓存节点新请求进来时从树根开始做最长前缀匹配命中部分直接复用只计算未命中的后缀部分。用 Python 实现这颗树时每个节点是dict边是字符串前缀。坦率说功能没问题但每次匹配要逐字符哈希、逐层做字符串比较长 prompt 场景下开销很大。换成 Rust 后我用Vecu8存储 token 序列每个树节点保存 token id 切片或区间索引匹配时直接按u32数组做逐段位比较。关键优化是边压缩。树节点不存储原始 token 序列而是存储base 数组 起始偏移 长度这样一个节点的匹配只需比较一段连续内存而不是拆成多个子节点逐层比对。实测下来长 prompt比如 4K token的 RadixTree 前缀匹配耗时从 Python 版本的约 1.2ms 降到约 60µs整整 20 倍提升。关于节点淘汰策略我沿用了引用计数机制。每个缓存节点记录被引用次数当某个请求结束且节点没有其他请求引用时这个节点标记为可回收但是不立刻释放——KV Cache 的成本在于显存回收次数过多会引起显存碎片。我的做法是延迟合并被标记可回收的节点先进一个垃圾桶列表等垃圾桶的显存占用超过阈值再统一做碎片整理和释放。注意RadixAttention 的 Rust 实现里最容易踩的坑是共享节点可变性问题。多个请求同时引用同一个前缀节点时谁都不能修改这个节点的任何内容否则其他请求的 cache 就污染了。必须用不可变语义约束共享部分只有独享后缀节点才能写入。我最初就是因为共享节点上直接追加了 KV导致跑批时偶尔出现记忆错乱式的输出错误。3.2 连续批处理调度器的状态机设计连续批处理Continuous Batching是推理引擎提升 GPU 利用率的核心策略。传统静态批处理要求整个 batch 全部结束才开始下一批而连续批处理允许请求动态加入和退出——某个请求生成完了立即腾出位置让新请求进来。这个机制的实现本质上是一个请求状态机。我在 Rust 里定义了清晰的枚举enum RequestState { New, // 新请求等待被调度 Prefill, // 正在执行 prefill 阶段 Decode, // 正在执行 decode 阶段 Pending, // 因显存不足被抢占等待恢复 Finished, // 生成完成等待返回并回收资源 Aborted, // 被客户端取消或超时 }状态转换的核心逻辑集中在BatchManager里。每一次 GPU 迭代之前BatchManager 做两件事一是扫描所有 New 状态请求判断显存是否允许 prefill二是扫描 Finished 和 Aborted 请求立即释放它们占用的 batch 槽位和 KV Cache 空间。调度策略里最微妙的一个参数叫做最大等待时间max_wait_ms。GPU 跑完一次 forward 后可能要等待新请求进来才划算等太久则延迟变差。我在实现中设计成空闲 GPU 每 2ms 检查一次新请求队列如果队列空则继续等待最多累积到 10ms 强制进入下一次迭代。这个值可以配置吞吐优先的场景调大延迟优先的场景调小。用 Rust 的 tokio 处理这个机制特别顺手。每个请求的生成过程封装成一个异步任务任务之间通过mpsc通道向调度器发信号。调度器本身是一个专门的tokio::task在收到触发信号后做状态推进。3.3 显存池和页式管理KV Cache 的 Rust 表达KVCache 管理是推理引擎最容易爆炸的地方。我的实现直接借鉴了 PagedAttention 的分页思路但用了 Rust 的Vec和Arc来做生命周期管理。我把显存池设计成三层结构BlockTable每个请求维护一个逻辑块表记录它拥有的显存块 ID 列表PhysicalBlockPool全局物理块池所有块的显存地址由 GPU 侧统一分配BlockPtr每个块在 Rust 里的句柄包含物理地址、大小和引用计数。一个关键 Rust 设计选择物理块在宿主侧用ArcBlockMeta管理GPU 侧使用固定的显存区域。当 BlockTable 里多个请求共享同一个前缀块时就是Arc::clone释放时Arc引用计数归零才真正归还物理池。这样避免了手动计数的悬挂指针问题。这个设计在高并发下比较稳但要注意Arc的原子操作开销。原子上面的 CPU 缓存行竞争在 32 核以上机器会非常明显。我的优化手段是最终把BlockMeta的引用计数从原子整数改成了crossbeam::atomic::AtomicCellu32因为块引用频率很高但读多写少用 relaxed 内存序就能满足需求比默认的SeqCst快了不少。4. 实操过程与核心环节实现4.1 无锁队列和无锁哈希表怎么搭调度器的核心数据结构是两个队列一个是待处理请求队列ingress queue一个是已完成请求队列completion queue。在 Go 或者 Java 里这两个队列通常会配锁。但 Rust 生态里crossbeam提供了近乎无锁的 MPMC channel底层是分段数组只在队尾竞争时才会触发原子 CAS。我在实现里用的是crossbeam_channel::unbounded作为多生产者到调度器的通道。实测 8 个 worker 同时提交请求入队吞吐大约 180 万请求/秒只是排队不含 GPU 执行完全不是瓶颈。在实现 RadixTree 时跨线程共享前缀树我采用了RwLockRadixTree。由于读多写少——绝大多数操作是查询和复用只有少数请求会创建新节点触发写锁——所以读写锁的并行度足够。这里要强调的是parking_lot的读写锁通常比标准库的RwLock快 2-3 倍因为它的实现不做公平性保证在低冲突场景下表现极好。4.2 异步任务与同步算子的协作不能全用 asyncGPU kernel 的启动CUDA launch本质是一个同步调用但真正的执行是异步的。Rust 的 async 上下文无法直接阻塞等待 GPU 完成。所以我在执行引擎里设计了一条同步边界调度器把所有待执行请求组装成 batch通过一个spawn_blocking的任务传给执行器执行器内部以同步方式启动 GPU kernel然后用stream synchronize等待完成。这个设计有一个容易被忽略的点不能在 async 任务里持有锁去等待 GPU。因为 GPU kernel 耗时可能上百毫秒如果 async 任务持锁等同步会让 tokio 调度器的所有线程卡住。正确的做法是在 async 任务里把要执行的 batch 数据准备好然后释放所有锁再交给spawn_blocking池去阻塞执行。4.3 性能剖析工具链perf、tracy、ncu 是怎么组合用的做完第一版实现后我遇到的最大问题是Rust 代码本身的效率很高但整体吞吐还是上不去。后来用perf top一看热点根本不在我写的 Rust 代码里而在两个地方一是 tokenizer 库的 C 绑定token 解析耗时占比 11%二是serde_json的序列化和反序列化每次生成的请求在返回 JSON 结果时占了约 7% 的 CPU。这个教训很有价值重写不是万能的你动手之前必须先搞清热点在哪。如果盲目优化 Rust 侧最多再提 1% 的性能而 tokenizer 和 JSON 序列化才是更大的坑。最后的解决方案tokenizer 调用从每条消息解析改成批量解析——一次接收多个输入用 rayon 并行执行 tokenizer 调用将 tokenizer 总耗时压缩了一半。JSON 序列化这快则放弃了部分字段的通用序列化手写了基于字节拼接的定长输出逻辑针对采样 logits 数组直接写 raw bytes省掉了层层嵌套的 serde 上下文。这两步组合让整体吞吐又上升了 6% 左右。下面我列一下典型 benchmark 的测试配置与结果采用的测试模型为 7B 参数模型输入输出长度比例 3:1工具脚本统一控制温度与采样策略指标原版 Python 运行时Rust 重写运行时提升幅度吞吐tokens/s8,1209,86021.5%TTFTprompt 4Kms420340-19%TPOTms/token4.54.2-6.7%CPU 调度箱时占比32%8%-75%需要说明的是这个测试是在我自己搭的容器环境里跑的GPU 是 4 卡模型虽然是 7B但是开了张量并行。不同环境下结果会有浮动但整体趋势一致。Python 侧到 Rust 侧的迁移收益主要来自调度路径的 CPU 时间大幅下降GPU 利用率因此得以提升。5. 常见问题与排查技巧实录5.1 这个问题的排查过程分享我在开发过程中遇到过最诡异的一个问题批量推理时偶尔会出现完全重复的生成内容而且概率不低大概在 0.3% 左右。排查过程花了两天最终定位到原因——在使用 RadixAttention 前缀树时我允许了共享节点附加新的 KV相当于一个全局的前缀节点被多个请求引用时其中一个请求在它上面续写了 KV Cache导致其他共享该前缀的请求读到了脏数据。复现方法其实很简单构造两个共享相同前缀的请求第二个请求的采样种子与第一个不同但生成结果居然完全一致说明第二个请求读到了第一个请求写入的 cache。定位后修复方式很直接共享节点的 KV 只能被预填充阶段的原始请求写入其他请求在读取时若发现需要续写必须新建节点做 fork。这个坑在 Rust 里靠所有权转移的思路很容易绕回来——只有持有该节点唯一引用的请求才有权限修改它。5.2 排查工具和手段清单以下是做这类引擎优化时最常用且有效的检查手段建议收藏备查性能热点优先看perf如果你的 Rust 引擎没用--features开启 debug symbolsperf 报告会很难看所以先确保debug true无锁队列的活锁检测crossbeam 无锁队列有极小概率在高争抢下触发重试。如果你观察到 CPU 定点升高而任务执行量不变用perf stat -d看context-switches和page-faults的异常变化内存碎片用jemalloc替代系统分配器在高并发调度下大量小块显存句柄分配和释放系统 malloc 的碎片率会非常惊人。启动时设置JEMALLOC_CONFbackground_thread:true,metadata_thp:auto能稳定内存占用GPU 空闲时间直接用ncu看 kernel 间间隔如果两次 kernel launch 之间超过 1ms 依赖CPU 侧调度必然有阻塞先检查是不是在持有锁的情况下做了同步等待。5.3 使用 Rust 重写时的高频注意点基于这些实践我总结了一些特别想对后来者说的高频注意点小心异步边界上的锁一定要记得在异步上下文中阻塞等待 GPU是最隐蔽的死锁来源。宁可多用spawn_blocking多占一个 tokio 线程也不要让异步线程池的线程被阻塞操作卡住Fast path 上的内存分配要最小化调度器里最忌讳在每次迭代都构造HashMap。正确思路是提前分配一个可复用的Vec(block_id, seq_id)在每次调度扫描时反复填充和清空原子操作的缓存行污染当一个结构体包含多个原子字段并且不同字段被不同线程修改时可能因为 false sharing 导致性能暴跌。解决手段是把不同类型的热字段用#[repr(align(64))]强制按缓存行对齐隔开FFI 层要用catch_unwind包一层Rust 的panic跨越 FFI 边界时C 侧根本不知情会直接 UB。我在所有 FFI 函数入口加了std::panic::catch_unwind把 panic 转换成错误码返回避免整个引擎崩溃。6. 我还在继续做的扩展方向和一点私人体会这个运行时做到现在已经能在生产环境里稳定跑了大半个月。除了一直在做的稳定性打磨我后续最想补的方向有两个。一个是异构调度——现在 GPU 上同一时刻只能执行一个大 batch我想进一步让 prefill 和 decode 任务在不同的 GPU 流上交错执行理论上还能再压榨一截显存带宽另一个是把 RadixTree 的脏节点回收做成异步后台任务让调度主循环更短更稳定。最后再多说一句实际上把一个成熟的 Python 推理框架用 Rust 重写表象是语言迁移本质是对调度路径上每一个数据结构的重新审视。这个过程你发现的问题远比安慰你的提升数字更有价值。比如我之前根本没意识到 tokenizer 的瓶颈占比那么高也没意识到 JSON 序列化会成为新的热点。Rust 把这些底层的成本暴露在你面前逼你面对它、解决它而不是继续用 Python 的胶水糊过去。这种掌控感是我做这个项目最大的收获。