老实说第一次跑完那份火焰图之后我盯着屏幕安静了很长时间。早高峰模拟压测时服务端的首字延迟就像心电图一样乱跳平时一百多毫秒压力一上来就飙到四五百毫秒个别请求甚至接近一秒。点开采样数据一看排在最前面的那些符号和 CUDA kernel 没关系也不是显存拷贝而是一堆 Python 层的函数调度步、前缀树匹配、块表查找、队列弹出。那会儿我心里基本有数了——这个基于 SGLang 推理框架搭起来的服务问题不在模型推理本身而在请求调度的运行时。所以我最后做了一个决定单独写一个高性能的 SGLang Rust 运行时把调度这一整层从 Python 里挪出来用 Rust 重写一遍。这个项目解决的是非常具体的问题在不换模型、不换核心 inference kernel 的前提下怎么把高并发场景下的调度开销按数量级往下压。它不替代 SGLang 的前端接口也不重新发明 KV cache 管理而是把请求从进入队列到完成解码这条路径上所有“运行时逻辑”用 Rust 实现。适合那些已经在用 SGLang、但对延迟抖动和 CPU 占用不满意、又不想把整个推理栈推翻重来的人参考也适合想了解 Rust 在 LLM serving 调度层到底能做什么的读者。1. 为什么把 SGLang 的运行时换掉一次延迟漂移排查1.1 火焰图最上面那几层项目起因说起来很简单。我当时在维护一个基于 SGLang 的文本生成服务模型是一个 7B 量级的对话模型流量模型很典型白天线上使用每天有几次明显的并发波峰。最初各项指标都还行但随着调用方从几个内部系统扩到几十个问题开始暴露——并发一高服务整体延迟并没有按预期平滑上升而是直接出现“断崖式恶化”P95 比 P50 高出好几倍。这种症状懂点推理服务的人第一反应可能是显存瓶颈或者 GPU 利用率打满。但我查了一圈GPU 利用率其实不算离谱显存也没有爆倒是 CPU 利用率有点反常。然后我开始采集线程栈一轮压测跑五分钟采样出来看调度线程的栈顶几乎全部集中在几个固定函数上——请求入队的锁竞争、RadixAttention 树节点匹配、batch 状态合并、tokenizer 调用。其中某个锁等待函数的热度高得让人有点尴尬。更直白地说请求一旦多起来调度循环本身成了瓶颈。模型 forward 一次可能只要几十毫秒但调度器在决定“下一步跑哪些请求、用哪些缓存块、怎么重组 batch”这件事上花掉的时间超过了模型推理本身。这就是为什么首字延迟会乱跳一个请求在队列里等了很久不是因为 GPU 忙而是因为 CPU 在调度其他请求时就快撑不住了。1.2 重写边界只碰调度层不碰模型层排查到这一步我最初的想法其实是“调参数”。SGLang 本身能调的东西不少batch 大小上限、并发窗口、预填充策略、排队算法。但我在本地反复试了几组配置后发现问题依旧存在只是把延迟峰往后挪了挪。这个判断很重要你没法靠参数调优去解决“调度代码本身太贵”的问题参数调整只是用空间换时间底层的 Python 开销还在那里。所以我开始考虑重写。但重写不等于把整个 SGLang 都推倒那样工作量太大了。我的边界非常清晰模型权重、CUDA kernel、Attention 计算逻辑、前缀树中 KV cache 的实际存储这些全部保留我要替换的只是调度循环、请求状态管理、树节点检索、块表操作以及队列这批偏向“逻辑控制”的部分。换句话说GPU 负责算Rust 负责管。这里有个容易误判的点很多人以为 Rust 运行时就得从 GPU 底层一路写上去其实不必。Rust 做得最好、风险最低的部分恰恰是那些和硬件计算无关、但需要极低延迟和高并发的控制逻辑。SGLang 的调度器本质上是一个复杂的事件驱动状态机这种形态放到 Rust 里编译后就是一组紧凑的循环和数据结构性能会比 Python 解释器下的同类实现好一个量级。1.3 选 Rust 而不是 C / Go 的理由既然要重写选型问题一直绕不开。我确实考虑过 C。C 做调度引擎完全可行SGLang 底层很多 kernel 本来就是 C/CUDA 写的直接用 C 还能省掉一层 FFI。但项目开发速度、内存安全边界、以及后续维护成本方面Rust 的优势非常明显。真正让我决定用 Rust 的其实是三个点。第一是所有权模型对这颗“共享缓存树”非常友好RadixAttention 是一棵多请求共享的树节点被多个请求同时引用释放顺序一旦错了就会造成 UAFuse-after-free。在 C 里这要靠精心设计的引用计数在 Rust 里我用带引用计数的共享指针加显式所有权编译器直接帮我挡住了一批低级错误。第二是 Rust 的标准工具链对“高吞吐小数据结构”支持很好Arc、RwLock、mpsc、crossbeam 这些库写出来就是无 GC、低锁冲突的风格。第三是测试方式Rust 的特性测试和基准测试非常方便我可以把调度器从 GPU 环境里完全剥离出来只喂 token 序列做精准的单元回归——这种东西在 Python 里做起来就麻烦得多。实际做完之后回头看Go 也不是不能做但它在内存布局和零开销抽象上的表现不如 Rust 贴近需求加上我们的核心代码要跟 C 扩展互操作Rust 的 FFI 设计更直截了当。当然选择 Rust 也带来了一些代价后面我会专门讲。2. 第一版架构重写了哪几条关键路径2.1 调度循环和请求状态机第一版跑通的时候我做了个很简单的请求生命周期定义新请求进来是 Waiting匹配到前缀后进入 Prefill然后每解码出一个 token 进入一次 Running 状态请求满足停止条件后进入 Finished。你会发现和 SGLang 原始的状态机基本一致这是有意的——我需要遵循同样的语义才能做行为兼容测试。调度循环是核心。每个“调度步”里运行时需要做四件事从外部请求队列取出新请求把已经在跑但还没结束的请求合并成下一轮 batch检查这批请求里哪些前缀可以复用缓存决定 prefill 组和 decode 组各自的组成。这些步骤在 Python 版里分散在好几个类和方法中每次都要经过属性读取和对象方法调用累计开销非常高。我在 Rust 里把四个步骤写成一个紧凑的循环内部用 enum 表示请求状态用索引数组代替散落的对象引用这样一来单步调度耗时从毫秒级降到了几十微秒级。这个改造的收益不在于单次调度快了多少而在于它让“调度频率”敢设计得更高。SGLang 默认的调度步通常受限于事件循环轮询粒度Rust 版本可以毫无压力地每 5 毫秒做一次调度决策。我实测下来调度步频的提升对首字延迟的改善非常明显——请求不用再干等一个较长的事件循环周期。2.2 RadixAttention 树缓存的 arena 实现RadixAttention 是 SGLang 调度器的精髓它把历史请求的前缀 token 存成 radix tree新请求只要匹配到公共前缀就可以直接复用之前的 KV cache 块省掉重复的 prefill 计算。这个设计本身很优雅但在 Python 里实现时节点是对象边是字典映射查找一个前缀路径要先构造 tuple 再做哈希内存开销和 GC 压力都不小。我在 Rust 里用 arena 方式重写了这棵树。简单说所有树节点分配在一块连续内存里节点之间不用指针链接而是用索引子节点查找用基于 token id 的哈希索引。这样做有两个好处局部性好遍历树的时候 CPU 缓存命中率高很多节点生命周期明确arena 本身可以整体回收不像散落对象那样需要依赖 GC 回收。树节点释放顺序也是个坑。每个请求运行结束后它占用的路径上的引用计数要减一当计数归零这个节点连同它对应的 KV cache 块要释放。但释放必须自底向上先子后父否则父节点会访问到已经释放的子节点。我在最初版本用一个递归释放函数结果在极深路径下栈溢出后来改成显式栈迭代安全得多。2.3 KV cache 块表的页表语义SGLang 的 KV cache 也是分块管理的逻辑块到物理块的映射很像操作系统的页表。在并发场景下这个映射的访问频率极高因为每调度一个请求都要查询或更新它。Python 实现里它是一个多级字典嵌套读一次要经历多次哈希写一次还要维护空余块位图。我的做法是把块表整体改成连续的数组逻辑块号直接作为索引物理块号作为数组元素空余块用一个简单的空闲链表维护。这样一次块分配就是一个数组读加一个链表头移动复杂度是 O(1) 并且没有哈希抖动。同时我把“块引用计数”直接放进块表数据结构里避免额外的哈希映射。这里有一个取舍固定用逻辑块号做索引会浪费一些内存因为逻辑块号可能稀疏。但在实际流量下块号连续度相当高这点浪费换来的稳定低延迟非常划算。如果你要做更碎片化的场景可以用两级页表但第一版完全不建议那么做。2.4 GPU worker 的薄封装设计调度循环在 Rust 里但模型 forward 还在原来的 GPU 侧。所以整体架构上是“Rust 大脑 GPU 手”的形态Rust 侧负责任务排队、前缀匹配、块表分配、batch 规划决定好本轮要执行的请求后把 batch 的 token id 序列和 query 信息通过共享内存发出去由一组极薄的 GPU worker 进程执行 forward。这组 worker 是故意的“薄”只保留三件事加载模型权重、接收 Rust side 的 batch 指令、执行 forward 并回传 token id 和新写入的 KV cache 块元数据。原本模型加载、采样 kernel 这些代码直接沿用不改动。这种架构带来的额外好处是即使调度器崩溃GPU 进程持有的模型上下文不会被拖垮重启调度器就能快速恢复。代价是跨进程通信有开销所以我只传“计划”而不是传张量token id 都是整数一次调度指令总共也就几十到几百字节。有人可能会问为什么不用 Rust 直接加载模型执行 forward。答案是没必要也暂时不现实现有的 CUDA kernel 和量化逻辑都是 Python 扩展包装好的我用 Rust 重写调度层已经是收益最大的部分如果把 GPU 执行栈也换掉工程量和风险都会失控第一版完全不该那么做。3. 把性能真正压出来的关键优化3.1 无锁队列与批处理节奏调度器对外接收请求的入口第一版我用了互斥锁队列。调了几轮后发现请求到达的 QPS 一高锁竞争会成为新的瓶颈。后来换成了 crossbeam 的无锁 MPSC 队列情况好转很多再配合一个专门的“领取线程”批量地把请求从队列里取出来攒够一定数量或到达时间阈值后统一交给调度循环效果更好。这个小设计的原理并不复杂与其让每个请求到达都立刻惊动调度器不如让调度器按自己的节拍批量消费。原来单个请求入队就要触发一次上下文切换现在每秒几十甚至上百个请求只在调度步边界整批被拉取。批处理节奏本身也是在给调度循环减负——每次调度面对的是一个批量集合而不是一个不断变化的动态集合决策逻辑简单得多。无锁队列带来的坑是 ABA 问题和内存回收。crossbeam 的 epoch 机制解决了大部分回收问题但在部分场景下内存占用峰值会略高。我用的另一个方案是固定容量的环形队列变体只在队列满时才退化成锁实测性能稳定空锁大部分时间不存在。如果你也想做类似的东西不要一上来就自己手写无锁队列先用成熟库把调度循环跑通再根据火焰图决定要不要优化到这一层。3.2 显式分离 prefill 和 decodeSGLang 的调度在处理 prefill 和 decode 时会做混合编排但混合编排的代价是 prefill 阶段的长序列会卡住 decode 阶段的短输出反之亦然。我在 Rust 运行时里做的最关键的设计之一就是把这两种状态彻底拆开维护两个独立的调度组prefill group 和 decode group。一个请求只有 prefill 完成后才进入 decode group不再混在一个 batch 里跑。这个改动带来两个直接效果。第一decode 组的 batch 大小可以保持稳定不再因为 prefill 大序列的插入而剧烈波动输出 token 的生成速率更平滑第二prefill 组可以更激进地向前填充把预填充吞吐提高因为不再受 decode 组的节奏拖累。代价是实现上要处理两组之间的缓存块转移不能让一个组在切换时丢掉已有 KV cache 引用。如果你把 SGLang 原版的调度参数调到特定组合也能近似达到这个效果但 Rust 版本把这种分离做成了架构上的默认而不是参数上的偶然。压力测试时这个改动对 P95 尾延迟的改善非常可观我的经验是比任何锁优化都要值得优先投入。3.3 内存池化与节点释放顺序调度器的 CPU 开销很大一部分来自内存分配请求对象要分配树节点要分配batch 要分配块表要分配。Python 里这些对象创建和销毁的主导因素是分配器开销与 GCRust 里虽然分配快但高频分配仍然会产生碎片和系统调用。我的做法是老一套但很有效内存池。请求对象从一个预分配的对象池里取用完归还树节点从 arena 分配arena 满了再整体扩容块表本身就是数组天然复用。实测下来系统调用次数从每秒几万降到了几百分配器带来的抖动几乎消失。还有一个值得提的细节树节点的释放不能简单依赖引用计数归零。由于多个请求可能共享同一棵树的祖先节点真正安全的释放必须等整条路径都不再被任何请求引用。我在计数器里存了“节点被多少活跃请求引用”清零时再做迭代式释放。曾经有一版我用了递归释放在深度超过一千层的树路径上直接爆栈。改成迭代后问题消失。这类细节看着小实际运行起来影响极大一旦处理错就是内存损坏或泄漏。4. 测试结果同样流量下的前后对比4.1 测试环境和流量构造为了公平对比我做了一套固定的测试环境同一个 7B 对话模型、同样的权重文件、同样的一张显卡分别跑原生 SGLang 服务和 Rust 运行时。流量方面用了个简单的模拟压测脚本混合了长短不同的请求一部分是十几 token 的短问题一部分是几百 token 的长上下文输出长度也从几十到几百不等。压测持续五分钟前六十秒做预热之后统计。客户端并发模拟为 64 个会话每个会话连续发请求间隔服从近似指数分布。这里要强调一点本地压测的绝对数字跟部署环境关系很大你照跑未必是同一个结果。它反映的是一个趋势同一个 GPU、同一份权重调度层变更后服务端的表现差异。因此下面表格里的数字我更希望读者关注变化比例而不是绝对值。4.2 指标变化TTFT、TPOT 与吞吐测下来的核心指标如下指标原生 SGLang 调度层Rust 运行时变化TTFT P50182 ms141 ms-22.5%TTFT P95415 ms238 ms-42.7%TPOT P5038 ms31 ms-18.4%TPOT P9576 ms49 ms-35.5%总吞吐328 req/s419 req/s27.7%调度循环单核占用~31%~4%-87.1%按我的理解TTFT 改善主要来自两个地方调度步频提升后请求等待下一个调度周期的时间变短prefill/decode 分离后prefill 组不再被 decode 慢请求拖住。TPOT 改善则主要来自 decode 组 batch 大小稳定后的连续性提升明显感觉到输出 token 的间隔更均匀了。尾延迟的改善比中位数更突出因为原生调度锁竞争最激烈的时刻恰恰是流量高峰的瞬间把竞争从热路径上移除后高峰期的排队等待被大幅削减。P95 从 415 毫秒降到 238 毫秒这种级别的变化对用户体感来说是“从明显卡顿到几乎无感”的跨度。4.3 尾延迟和批次分布的变化除了汇总指标我还统计了调度批次大小的分布。原生运行时的 batch 大小波动很大某些瞬间 batch 会冲到很大的值另一些瞬间又几乎空转Rust 运行时由于显式分离batch 大小分布更集中在几个稳定的档位。这个分布直接影响延迟抖动batch 过大会导致单次迭代时间突然拉长batch 过小又浪费 GPU 计算资源。另一个观察是空闲周期。原生调度在低负载时会让调度线程空转等待事件循环的轮询频率不高Rust 运行时因为调度循环本身极轻我可以让它在负载低时快速进入等待有请求到达时立刻再次调度。这个差异对首字延迟的低负载表现也有改善。当然29% 的吞吐提升并不是每类负载都一样。长上下文、高复用前缀场景下提升更明显因为前缀匹配的 CPU 开销在原生实现里更大短请求、前缀完全互不共享的场景下提升就相对有限大概在百分之十几。做这个对比也让我想清楚了一件事这个运行时真正的价值在大量请求共享前缀、调度压力大的生产流量里。5. 接入方式、兼容层和踩坑记录5.1 OpenAI 兼容 API 桥接接入现有服务时我最担心的是客户端改版。所以首先保证对外 API 保持 OpenAI 兼容仍然是 POST /v1/chat/completions仍然返回流式 SSE业务方代码一行不用改。内部把请求解析成统一的请求结构体直接送入无锁队列返回路径同样通过 SSE 把 token 流式吐回。Rust 侧直接用异步 HTTP 框架挂起监听每个请求连接对应一个任务。注意这里有个细节异步任务里的“挂起等待”必须和调度循环解耦不能让 tokio 的调度和 GPU batch 调度互锁。我的做法是请求进来先写进队列返回一个响应通道调度循环每轮把完成的输出 token 写入对应通道异步任务从通道取到 token 后立刻发送给客户端。这个桥接层虽然代码量不大但踩的坑不少。最典型的流式响应时部分客户端会等待首包才建立连接如果调度器有一段时间没有输出客户端会超时。解决方案是在等待 batch 调度之前先发送一个空事件或注释行保持连接活跃再把真实 token 推出去。5.2 tokenizer 和采样逻辑的处理tokenizer 的处理也很关键。原生 SGLang 里 tokenizer 用 Python 库跑每次前处理都要经过 Python 层我在 Rust 侧直接复用 Hugging Face tokenizers 的 Rust 绑定加载同一个 tokenizer.json在请求入队时完成分词不再跨进程调用 Python。这样请求进入调度器之前就已经是 token id 数组省掉了大量跨语言开销。采样器是我一开始没想清楚的部分。第一版我直接沿用原本 Python 侧的采样逻辑Rust 调度器决定 batchGPU worker 执行 forward 后把 logits 交回 Python 侧采样。这意味着跨进程回传的是一整块 logits 张量通信量瞬间变大性能反而掉了。后来我把采样器也搬到 Rust 侧GPU worker 只回传 token id 和必要的概率向量摘要信息Rust 本地完成 top-k/top-p 采样。greedy 和 top-k/top-p 这两种模式在 Rust 里各写一遍逻辑代码量不大但效果一下子好起来了。这个改动也带来一个取舍目前我的 Rust 采样器只支持常见的确定性采样和随机采样beam search 这类复杂采样还没实现。如果业务里重度使用 beam search这块需要额外开发。我的建议是先明确你的业务主要用哪种采样策略再决定要不要把采样器也下沉到 Rust。5.3 过程里最值得记下的几个坑第一个坑是内存布局假设。Rust 侧我把块表设计成连续数组但 GPU worker 那边的 KV cache 块索引是从原版逻辑映射过来的两边索引语义不完全一致。调试时经常出现“Rust 以为块被占用、GPU worker 却复用同一个物理块”的诡异问题。最后我统一了物理块号的来源所有块分配都以 Rust 侧为准GPU worker 只提供块数量的限制问题才彻底消失。第二个坑是 RadixAttention 的前缀匹配语义。原版在匹配公共前缀时有一些边界处理比如某些请求的 prompt 完全被之前的 prompt 包含、某些后缀 token 并非有效前缀等。照搬状态机定义不够必须跑一批回归用例跟原版逐 token 对比。我做了个自动对比工具输入同一请求序列比对两边的输出 token 序列和 KV cache 导出信息确保行为等价。哪怕有一丁点不一致代码就得修。第三个坑是堆栈深度。上面提过的递归释放导致栈溢出是我第一版里最头疼的问题。换成迭代式释放后实测最深的 tree 路径也不会崩。这个问题提醒我核心数据结构即使在内存分配上很安全递归深度仍旧是不可忽视的隐患。第四个坑也是最容易被人忽略的基准测试本身的可重复性。做性能对比时如果两边随机种子不一致输出内容不同缓存的复用率就不同指标对比毫无意义。我最后用了固定 token 序列、固定采样种子并且关闭了任何影响前缀匹配的随机性才能得到相对可信的数字。6. 如果让我重做一次我会怎么调整如果从头再写一遍我不会先碰 GPU也不会第一时间写 HTTP API。第一步应该写一个纯 CPU 的调度模拟器读入一份真实请求 trace模拟 RadixAttention 和 batch 调度输出标准的调度事件日志。把这些日志和原版 SGLang 在同样 trace 下产生的日志做 diff先验证调度逻辑完全等价再接入真实 GPU 处理。当初我跳过了这一步导致后面大量时间花在排查“是调度逻辑 bug 还是 GPU 通信 bug”上很痛苦。性能调试的顺序也应该反过来。我一开始就盯着无锁队列和块表优化后来发现收益的大头并不在那里而是 prefill/decode 分离和采样器下沉。如果下次做同类项目我会先把 trace 测出来统计每次调度循环真正的时间分布再挑占比最高的部分优化。还有一个细节所有共享内存的数据结构都应该加上结构化日志和内存校验。跨进程场景下Rust 侧状态和 GPU worker 侧的 KV cache 状态一旦不一致错误会延迟到几十步之后才暴露出来。有了日志定位这类问题会快很多。这个项目做到这里我自己的体会是SGLang 这类框架的性能瓶颈并不总在“显存”和“kernel”而经常在调度那层看起来不起眼的软件逻辑里。Rust 运行时只是把这层逻辑重新实现了一遍没有改变模型本身的分毫。从实际效果看这个方向是值得的。如果你也正在被 LLM serving 的高延迟和高 CPU 占用困扰不妨先跑一份火焰图仔细看看再做决定。