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 后，我用Vec<u8>存储 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 设计选择：物理块在宿主侧用Arc<BlockMeta>管理，GPU 侧使用固定的显存区域。当 BlockTable 里多个请求共享同一个前缀块时，就是Arc::clone，释放时Arc引用计数归零才真正归还物理池。这样避免了手动计数的悬挂指针问题。

这个设计在高并发下比较稳，但要注意Arc的原子操作开销。原子上面的 CPU 缓存行竞争在 32 核以上机器会非常明显。我的优化手段是最终把BlockMeta的引用计数从原子整数改成了crossbeam::atomic::AtomicCell<u32>，因为块引用频率很高，但读多写少，用 relaxed 内存序就能满足需求，比默认的SeqCst快了不少。

4. 实操过程与核心环节实现

4.1 无锁队列和无锁哈希表怎么搭

调度器的核心数据结构是两个队列：一个是待处理请求队列（ingress queue），一个是已完成请求队列（completion queue）。在 Go 或者 Java 里，这两个队列通常会配锁。但 Rust 生态里，crossbeam提供了近乎无锁的 MPMC channel，底层是分段数组，只在队尾竞争时才会触发原子 CAS。

我在实现里用的是crossbeam_channel::unbounded作为多生产者到调度器的通道。实测 8 个 worker 同时提交请求，入队吞吐大约 180 万请求/秒（只是排队，不含 GPU 执行），完全不是瓶颈。

在实现 RadixTree 时，跨线程共享前缀树我采用了RwLock<RadixTree>。由于读多写少——绝大多数操作是查询和复用，只有少数请求会创建新节点触发写锁——所以读写锁的并行度足够。这里要强调的是：parking_lot的读写锁通常比标准库的RwLock快 2-3 倍，因为它的实现不做公平性保证，在低冲突场景下表现极好。

4.2 异步任务与同步算子的协作：不能全用 async

GPU 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，工具脚本统一控制温度与采样策略：

需要说明的是，这个测试是在我自己搭的容器环境里跑的，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 symbols，perf 报告会很难看，所以先确保debug = true；
• 无锁队列的活锁检测：crossbeam 无锁队列有极小概率在高争抢下触发重试。如果你观察到 CPU 定点升高而任务执行量不变，用perf stat -d看context-switches和page-faults的异常变化；
• 内存碎片用jemalloc替代系统分配器：在高并发调度下，大量小块显存句柄分配和释放，系统 malloc 的碎片率会非常惊人。启动时设置JEMALLOC_CONF=background_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 的"胶水"糊过去。这种掌控感，是我做这个项目最大的收获。