ARTICLE DETAIL

建站实战干货

来自一线的建站与推广经验沉淀,每一条都经过真实交付验证。

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

2026/9/16 17:57:01 拓冰建站 浏览量
推理延迟优化: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 块进 decodeKV 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 streamGPU 才能吃饱。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 只能访问“只读”状态的 blockprefill 只能写“空闲”状态的 block“待回收”状态表示已经没有 reader 在读了可以安全回到空闲。状态机本身不值钱值钱的是设计者想清楚“从一个状态跳到另一个状态的时机到底由什么事件来触发”。3. 同步边界怎么划正确性红线和可放宽区域3.1 三条不可跨越的同步红线第一红线是同一个序列的 KV 读后写依赖。decode 的 step N1 必须读到 step N 写入的 KV这个依赖不能被任何调度技巧剪掉。Mini-SGLang 的做法是在每个序列内部维护一个 step 计数器只有当所有执行单元都已经完成该序列的当前 step 时才允许进入下一个 step。这不一定要用锁去保护因为同一时刻只有一条 CUDA stream 真正执行该序列的某个 step调度器只需要保证“不会出现两个 stream 同时执行同一序列的相邻 step”。第二红线是 batch slot 的生命周期。调度器会把多个请求组成一个 batch 来执行每个请求占据一个 slotslot 里存着这个请求的当前状态、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 同步操作的成本量级选边界就是选开销为了在设计时对“该不该加一个同步”有感觉我整理了一张大致的成本表全是当前主流硬件和驱动版本下的粗测值同步方式开销量级适用场景线程锁互斥量几十纳秒到几微秒保护短小临界区如状态标志切换原子操作几纳秒到几十纳秒计数器增减、引用计数cudaEventRecord cudaStreamWaitEvent微秒级同一设备内不同 stream 的依赖关系cudaDeviceSynchronize几百微秒到几毫秒性能灾难只用于调试CPU 与 GPU 间的 memcpy十几微秒到数百微秒仅传输小批量结果数据顺着这张表往下推就能得出几条设计原则。能用原子操作解决的同步绝不用锁能用 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 kernelkernel 结束后记录一个 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 的条件下我测过几组典型参数。这里不放全量数据挑最有代表性的三行配置chunk_sizemax_inflight吞吐req/sTPOTms串行调度不切分13.285overlap 但 chunk 太小12844.873overlap 且 chunk 适中76846.166overlap 且同步过度76814.579“overlap 且同步过度”那一行很能说明问题明明 chunk 选对了但因为我在每个步骤之间都加了严格的 CPU 等待把并发 pipeline 又打回了串行。Max_inflight 从 4 降到 1等于把 pipeline 拆掉了吞吐直接掉了四分之一。同步边界不是越多越好这句话我在这个数据上体会得很深。调参的顺序建议是先固定 max_inflight4把 chunk_size 从 128 逐步加大到 2048观察吞吐峰值然后再动态调 max_inflight观察 TPOT 稳定性。最后用 nsys 抓一帧时间线确认 prefill 和 decode 的 kernel 在时间轴上确实有重叠区间而不是自认为有重叠。我在最后想分享一个体会Mini-SGLang 这个项目做到后面真正的复杂度不在“如何并行”而在“如何让所有并行的执行单元对数据的时间达成共识”。数据流确定了任务之间天然的逻辑依赖同步边界则把这些依赖翻译成系统里的等待事件。对每一项同步都问一句“这笔等待值不值”是让推理引擎性能再上一个台阶最值得投入的精力所在。