推理延迟优化：prefill与decode重叠调度实战解析

跑推理服务时间长了，你会发现一件反直觉的事：GPU的标称算力明明很高，但实时利用率曲线却像心电图一样忽高忽低。之前我在本地复刻了一套极简版推理引擎（就叫它 Mini-SGLang），核心目的就是把 SGLang 的调度逻辑剥出来，研究数据到底怎么流动、同步到底该放在哪里。研究下来，决定性能上限的往往不是某个 kernel 本身有多快，而是调度器怎么安排 prefill 和 decode 之间的重叠。prefill 计算密集，decode 访存密集，二者天然特性完全不同，串行安排必然出现大量气泡。这篇文章就从数据流和同步边界两个角度，把我在 Mini-SGLang 上调通 Overlap Scheduling 的过程完整拆一遍。

1. 先厘清：Mini-SGLang 到底在调度什么

1.1 prefill 与 decode 的算力错配是整个问题的源头

LLM 推理的每个请求都会经历两个完全不同的阶段。prefill 阶段要处理几百到几千个输入 token，主要工作是矩阵乘法和 attention 计算，计算密度非常高，GPU 上的 SM 几乎全程满载。decode 阶段则完全不同，每个 step 只需要生成一个 token，矩阵规模大幅缩小，反而变成了以显存带宽为主的访存操作，算力利用率往往只有个位数到十几个百分点。

这两者一个是 CPU 密集型的比喻，一个是 I/O 密集型的比喻，如果调度器把它们当成同一类任务排队处理，GPU 就会反复在“跑满”和“空转”之间切换。更麻烦的是，decode 阶段一旦开始，就被客户端的流式响应拴住了，不能随便中断太久；而 prefill 请求又是一下子涌进来的。传统做法就是把两者分时交错，前一秒集中做 prefill，后一秒集中做 decode，肉眼可见地浪费硬件。

1.2 从 SGLang 到 Mini-SGLang：抽掉复杂度才能看清调度

SGLang 本身是一个相当庞大的系统，有 RadixAttention 前缀树、分布式多 worker、连续批处理、自动化 KV cache 管理，还有专用编译器优化。这些问题单独拎出来每个都是深坑，混在一起根本无法判断“当前性能瓶颈到底在哪一层”。

Mini-SGLang 的思路是只保留三件套：调度队列、KV cache 管理器、执行引擎。没有前缀树复用，不做多机多卡，把 beam search 也砍掉，只支持最简单的贪心解码和随机采样。这样一个精简系统反而能清楚看到调度器的行为：请求来了进队列，调度器决定哪些 token 块进 prefill、哪些 token 块进 decode，KV cache 管理器负责分配和回收内存，执行引擎把计算图跑起来。所有瓶颈都赤裸裸地摆在你面前，没有任何高级特性帮忙掩盖。

1.3 overlap 调度要达成的三个目标

设计这套调度时，我心里有三个明确的验收标准。第一是吞吐量要明显提升，同样的 GPU 下每秒处理的 request 数要比纯串行调度高出一截。第二是 TPOT（每个输出 token 的生成时间）要尽量稳定，不能让 decode 用户感觉到“前一个请求做 prefill 时我突然卡了”。第三也是最容易被忽略的，正确性不能因为并发引入的调度变化而改变，同一条请求无论什么时候跑，输出分布应当一致。

这三个目标里，前两个靠数据流设计解决，第三个靠同步边界设计解决。很多人做 overlap 调度，精力全花在怎么把任务拆碎、怎么并发执行上，最后栽在了“数据到底什么时候对另一个执行单元可见”这个基础问题上。所以我更愿意把同步边界当作第一约束条件来设计，而不是最后再来补的补丁。

2. 顺着数据流看调度：请求在引擎里的完整流动路径

2.1 一条请求从进入到吐字的完整路径

在 Mini-SGLang 里，一条 HTTP 请求的完整路径大致是这样的：网关把文本交给调度器，调度器先做分词，把文本变成 token id 列表；接着它要检查这段 token 序列的公共前缀是否在 KV cache 里命中过，命中就直接跳过前面的计算，只计算新增部分，这是 SGLang 的看家本领，Mini 版本里我保留了一个最朴素的版本。然后执行引擎开始 prefill，把新增 token 对应的 key 和 value 写入 KV cache；之后进入 decode 循环，每步都从 KV cache 里读出历史的 K、V，和当前 query 做 attention，得到 logits 后交给采样器选出下一个 token；新的 token 又要追加写入 KV cache，同时作为下一步的输入。直到采样器吐出一个结束符，整条请求才算结束。

这个过程里存在多个数据流，最核心的是 token 数据流和 KV 数据流，两者几乎是绑定在一起的。token 流是显式的，从输入到输出，每一层网络都能看到。KV 流则是隐式的，它不断增长、不断被后续 step 读取。理解 overlap 调度的前提，就是要意识到 KV 流才是真正的“心脏血流”——token 只是表层信号，模型状态全部存放在 KV 流里。

2.2 chunk 化：让 prefill 和 decode 可重叠的关键动作

chunked prefill 这个词很多人听过，但未必理解它为什么是开启 overlap 的钥匙。一条长 prompt 如果一次性做完整 prefill，几百毫秒内 GPU 只能服务这一个请求，其他所有 decode 请求都得等着。可如果把它按 token 顺序切成几个 chunk，每个 chunk 就是一个稍小的 prefill 任务，就可以插在 decode step 的间隙里执行。

这里面的原理来自 attention 的因果掩码。每个 chunk 只需要读自己之前所有 token 的 KV，不需要等后半个 chunk 的结果。所以在计算第 N 个 chunk 时，调度器可以同时让另一条已经 prefill 到位的请求继续做 decode。两者使用不同的计算单元或不同的 CUDA stream，GPU 才能吃饱。

chunk 的大小是 overlap 调度里第一个要调的旋钮。太小了，kernel 启动开销会吃掉收益，每个 chunk 只有一两百 token，跑出来的时间还没 launch kernel 花的时间长；太大了，又回到一次性 prefill 的串行问题。我实测下来 7B 模型在 A10 上，512 到 1024 token 一个 chunk 比较合适。这个值不是一个玄学数字，它可以用 kernel 启动延迟除以每个 token 的平均计算时间反推，简单估算：如果启动一次 kernel 要 5 微秒，单 token 计算 0.1 微秒，那至少得 50 个 token 才能抵消启动开销，实际要留出几倍余量。

2.3 KV cache 是数据流的心脏也是调度器的账本

KV cache 在物理上是一块预先分配好的显存池，通常按 block 为单位划分，每个 block 能装固定数量的 token 的 K、V 张量。调度器同时要维护一张逻辑表，记录每个 block 当前属于哪条请求、已经写了多少个 token、哪些 block 被彻底释放了。

这个设计非常像操作系统里的内存分页。请求来了，调度器像内核给进程分配页一样分配 block；请求结束了，block 被回收放进空闲链表。overlap 调度带来的复杂之处在于，同一个 block 可能正在被 decode stream 读取，而调度器已经把它标记为空闲，准备分配给新请求的 prefill 写入。这就是典型的数据竞争：写入方覆盖了读取方还没读完的数据。

所以 KV cache 的分配和释放必须作为同步边界来管理。我在 Mini-SGLang 里给每个 block 加了一个状态机：空闲、写入中、只读、待回收。decode stream 只能访问“只读”状态的 block；prefill 只能写“空闲”状态的 block；“待回收”状态表示已经没有 reader 在读了，可以安全回到空闲。状态机本身不值钱，值钱的是设计者想清楚“从一个状态跳到另一个状态的时机到底由什么事件来触发”。

3. 同步边界怎么划：正确性红线和可放宽区域

3.1 三条不可跨越的同步红线

第一红线是同一个序列的 KV 读后写依赖。decode 的 step N+1 必须读到 step N 写入的 KV，这个依赖不能被任何调度技巧剪掉。Mini-SGLang 的做法是在每个序列内部维护一个 step 计数器，只有当所有执行单元都已经完成该序列的当前 step 时，才允许进入下一个 step。这不一定要用锁去保护，因为同一时刻只有一条 CUDA stream 真正执行该序列的某个 step，调度器只需要保证“不会出现两个 stream 同时执行同一序列的相邻 step”。

第二红线是 batch slot 的生命周期。调度器会把多个请求组成一个 batch 来执行，每个请求占据一个 slot，slot 里存着这个请求的当前状态、KV block 列表、输出 token 序列。如果请求在 decode 中途被终止（客户端断开或生成了结束符），slot 被回收，但另一个线程可能还在用它做采样计算。这个 bug 我在早期版本里真实遇到过，表现就是极低概率的错乱输出，排查起来极其恶心。后来加了一条铁律：slot 的状态切换必须经过调度器主线程，执行引擎只上报事件，不直接动 slot 状态。

第三红线是采样器状态。随机采样需要一个随机数生成器，如果多个请求复用同一个生成器的状态而不做隔离，输出分布会受到污染。这和数据处理里的“种子共享”问题一样，看似概率极低，一旦并发规模上来，迟早复现。

3.2 可以放心放宽同步的区域

同步边界不是越多越好，锁和事件都是成本。有三类数据我建议直接放宽同步，不要为了“绝对一致”把性能拖死。

第一类是指标统计。比如每轮迭代的 GPU 利用率、cache 命中率、平均 decode 延迟，这些数据晚几十毫秒更新完全无所谓。我在主循环里让各个 worker 线程把指标写进本地 buffer，主线程每秒汇总一次，绝不实时去读 worker 的统计变量。做 profiling 的时候单独开一个开关，打开后才走精确路径，默认跑异步路径。

第二类是前缀缓存的后台更新。RadixAttention 的树结构在每次请求进来自动更新当然好，但锁开销很高。Mini-SGLang 里配额是“请求完成后再异步更新前缀树”，新请求已经到达时宁可少命中一些前缀，也不阻塞在缓存更新的锁上。实测下来缓存命中率下降了不到 3%，但调度抖动明显减少。

第三类是日志和事件追踪。日志 IO 在推理路径上绝对是毒瘤，一旦磁盘抖一下，整个调度周期都被拖住。我的方案是把日志内容先堆到内存环形缓冲区，满了就丢最旧的，由独立线程缓慢落盘。启动时加环境变量开关，需要 debug 再打开完整日志。

3.3 同步操作的成本量级：选边界就是选开销

为了在设计时对“该不该加一个同步”有感觉，我整理了一张大致的成本表，全是当前主流硬件和驱动版本下的粗测值：

顺着这张表往下推，就能得出几条设计原则。能用原子操作解决的同步，绝不用锁；能用 CUDA event 解决的跨 stream 依赖，绝不用 device synchronize；能在 CPU 侧延迟处理的统计，绝不进 GPU 同步区间。同步边界的本质就是“数据依赖必须成立的地方才有同步”，而不是“我担心并发问题就到处加锁”。

4. 工程落地：用 CUDA Stream 和事件循环把 overlap 搭出来

4.1 双 Stream 框架：计算和拷贝各走各的轨道

Mini-SGLang 的执行引擎在 GPU 上维护两类 CUDA stream。计算 stream 负责跑 attention、MLP 这类核心算子；拷贝 stream 负责把采样结果从 GPU 拷回 CPU，以及把新来的输入 token 从 CPU 拷进 GPU。这两类 stream 之间通过 cudaEvent 建立依赖，而不是靠强制串行。

具体来说，每次 decode 迭代开始时，调度器先在计算 stream 上 launch 一个 attention kernel；kernel 结束后记录一个 event，拷贝 stream 在等待这个 event 之后才执行结果回传。回传的数据反而是下一轮迭代要用的，所以下一轮的计算 stream 又要等拷贝 stream 上的另一个 event。这样就形成了一条 pipeline：计算、拷贝、再计算、再拷贝，每一轮的等待时间被另一段工作填满，GPU 和 GPU 之间的空闲链路被压缩到很短。

这里有一个很多教程不会提的细节：CUDA stream 内部的 kernel 默认是串行的，所以要 overlap 的是在不同 stream 上的 kernel。如果你只想 shuffle 一下代码、把所有算子放在同一个 stream 里，无论怎么写都不可能重叠。

4.2 事件循环调度器：生产消费与背压控制

调度器本身是一个事件循环，它维护三条队列：等待队列（新请求还没分配 KV block）、就绪队列（已经 prefilled 完，可以进入 decode 或继续 chunk prefill）、完成队列（本轮已经跑完等待回收 slot 的请求）。执行引擎是消费者，每完成一个 batch 的计算，就把这些请求的状态递交给调度器。

这里最关键的工程决策是背压控制。如果允许调度器无限地向执行引擎提交任务，GPU 队列会被塞满，显存也可能溢出。我在 Mini-SGLang 里用一个最大在途任务数来控制，默认设成 4。也就是说，调度器最多同时积累 4 个“已提交但还没执行完毕”的迭代任务。每结束一个，再补充一个，保持 pipeline 是满的，但不会狂泄。

实现上，我用了类似信号量的计数器和条件变量。执行引擎每完成一个 batch，调用一个回调，主线程在回调里把计数器减一，然后唤醒调度线程继续派发新任务。这个模式的优点是控制粒度非常细，任务级别而不是请求级别；缺点是回调函数本身不能做重活儿，型号错了、日志、数据搬运都不能放在回调里，只允许更新计数器和推指针。

4.3 让 overlap 真正发生的三个细节

第一个细节是不要过度依赖“自动并发”。PyTorch 的 CUDA 图或者推理框架的图优化有时候会自动分配 stream，但 Mini-SGLang 的目标是手动控制，所以我显式地指定了每个算子在哪个 stream 上跑，并且用 export_graph 之后用 nsys 检查流之间的依赖是否和设计一致。

第二个细节是 kernel 选择上要避开“巨无霸”。一个占用 10 毫秒的大 kernel 会把整个 stream 占死，其他 stream 只能干等。overlap 调度的偏好是“很多个中等大小的 kernel 交替发射”，这样才有可能把一个 stream 的空隙塞进另一个 stream 的工作。所以实际操作里，我会把大的 prefill 算子进一步拆分到多个 chunk 里执行，每个 chunk 对应一到两个 kernel，再插入 decode 的计算。

第三个细节是采样器要放在 CPU 端执行。GPU 上生成 logits 后直接拷回 CPU 做采样，采样一个 token 的耗时在微秒级，CPU 完全能扛。如果强行在 GPU 上用自定义 kernel 做采样，反倒把 GPU 占住了。这种“让 CPU 分担一部分非核心工作”的思路，也是 overlap 的一种——它重叠的是 CPU 和 GPU 的工作，而不是 GPU 和 GPU 的工作。

5. 实测与踩坑：同步边界和数据流打架的真实案例

5.1 坑一：KV block 复用导致的“幽灵回复”

早期版本里，我为了追求极致的 overlap，把 block 回收做成异步：请求结束时，block 直接标为空闲，立刻可以分配给新请求。结果在一个并发稍高的压测脚本里，大概跑了两万多次请求后出现了一次诡异输出：客户端明明问的是“今天天气怎么样”，回复的却是另一条完全无关的话。

排查链路是这样的：先怀疑数据集污染，检查输入输出对不上，发现那个回复确实来自于另一条请求。然后怀疑采样器种子交叉，查了一遍也没有问题。最后把 KV 分配日志打开，才发现那条被误回复的请求在 decode 阶段刚开始时，它读取的某个 KV block 已经被新请求的 prefill 写入覆盖了。根因就是 block 回收太快，没有确认“之后不会再被读取”。

修复方式就是我前面提到的状态机：block 从“只读”到“空闲”必须经过“待回收”状态，等待所有正在读取它的 stream 都越过对应的 CUDA event 之后才真正释放。修复后这个 bug 再也没出现过。这里我给读者的建议是：宁可多等一个迭代的显存浪费，也不要让回收动作跨越同步边界。

5.2 坑二：Python 侧的 Stream Callback 几乎抵消了全部收益

我在一开始实现“batch 完成回调”时图方便，直接在 Python 里写了回调函数，用 torch.cuda.CUDAPluggableContext 之类的接口注册。结果压测出来的吞吐量不仅没涨，反而比串行调度还低了 20%。查证了 Utils 发现，Python 层回调在执行时会上 GIL，而且跨线程调用 Python 对象会有额外的同步开销，等于我在“优化后的并发结构”外面套了一层性能枷锁。

这个坑的修复办法是把回调挪到 C++ 扩展层实现，在 C++ 里只做计数和指针交换，不碰 Python 对象。攒够一批结果后，统一交给 Python 层处理。这一步改动直接让测试吞吐量回升了近一半——Overlap 的收益又回来了。

给同类项目的建议是：能用 native 层完成的回调逻辑，千万不要为了省事拉到 Python 层。Python 的灵活性和推理引擎的性能要求之间有一条需要靠工程经验划出的线。

5.3 参数调节和可参考的实测数据

在 A10 GPU、7B 模型、并发 16 的条件下，我测过几组典型参数。这里不放全量数据，挑最有代表性的三行：

“overlap 且同步过度”那一行很能说明问题：明明 chunk 选对了，但因为我在每个步骤之间都加了严格的 CPU 等待，把并发 pipeline 又打回了串行。Max_inflight 从 4 降到 1，等于把 pipeline 拆掉了，吞吐直接掉了四分之一。同步边界不是越多越好，这句话我在这个数据上体会得很深。

调参的顺序建议是先固定 max_inflight=4，把 chunk_size 从 128 逐步加大到 2048，观察吞吐峰值；然后再动态调 max_inflight，观察 TPOT 稳定性。最后用 nsys 抓一帧时间线，确认 prefill 和 decode 的 kernel 在时间轴上确实有重叠区间，而不是自认为有重叠。

我在最后想分享一个体会：Mini-SGLang 这个项目做到后面，真正的复杂度不在“如何并行”，而在“如何让所有并行的执行单元对数据的时间达成共识”。数据流确定了任务之间天然的逻辑依赖，同步边界则把这些依赖翻译成系统里的等待事件。对每一项同步都问一句“这笔等待值不值”，是让推理引擎性能再上一个台阶最值得投入的精力所在。