重叠调度：藏起 CPU 开销¶
一个小模型 decode 一步，GPU 只要几毫秒；而调度器每一轮在 CPU 上要做的事——收消息、组 batch、准备元数据、处理上一轮的结果、发消息——同样要几毫秒。如果两者轮流进行，GPU 有将近一半的时间在等 CPU。重叠调度（出自 NanoFlow，SGLang 0.4 引入）让 CPU 在 GPU 计算第 N+1 轮的同时处理第 N 轮的结果。mini-sglang 的实现只多了十几行，但它改变了"什么时候知道什么"，由此引出了本书在复刻时发现的四个问题。
自测：能答上来就可以跳过本章
- 发射第 N+1 轮时，CPU 还不知道第 N 轮采样出的 token。第 N+1 轮的输入从哪来？
- 一个请求在第 N 轮采样出了 EOS，它会不会出现在第 N+1 轮？
- 为什么调度器和引擎要用两条不同的 CUDA stream？
- 重叠调度下，处理第 N 轮结果时，
req.device_len反映的是第几轮之后的状态？
自测参考答案（先自己答，再展开对照）
- 从 GPU 上的 token pool：第 N 轮在 GPU 上采样出的 token 直接写进 token pool，第 N+1 轮的前向从 token pool 读取输入，CPU 不需要知道它是什么。
- 可能会：发射第 N+1 轮时 CPU 还没处理第 N 轮的结果，不知道它已经结束，所以多调度了一轮；处理结果时丢弃这个多余的 token。
- 让调度器准备下一轮（拷贝元数据、组 batch）的工作和引擎的计算在不同的 stream 上并行，并用事件建立依赖；都放在一条 stream 上就又串行了。
- 第 N+1 轮之后：CPU 上的请求状态在发射时就提前推进了，比正在处理的结果领先一轮。
本章要写的文件：scheduler/scheduler.py 中的 overlap_loop、run_forever，以及 _process_last_data、_free_req_resources、_process_one_msg 里与重叠有关的部分。
minisgl/
├── __init__.py
├── core.py
├── env.py
├── attention/ （4 个文件，不变）
├── distributed/ （3 个文件，不变）
├── engine/ （5 个文件，不变）
├── kernel/ （2 个文件，不变）
├── kvcache/ （5 个文件，不变）
├── layers/ （8 个文件，不变）
├── llm/ （2 个文件，不变）
├── message/ （4 个文件，不变）
├── models/ （7 个文件，不变）
├── moe/ （2 个文件，不变）
├── scheduler/
│ ├── table.py
│ ├── __init__.py
│ ├── config.py
│ ├── utils.py
│ ├── decode.py
│ ├── prefill.py
│ ├── scheduler.py ← 本章修改
│ ├── cache.py
│ ├── io.py
├── utils/ （8 个文件，不变）
examples/ch11_overlap.py——它只用到上面这些文件；python tools/steps.py check 会逐章搭出这棵树、跑这个 main。
两种循环¶
def overlap_loop(self, last_data: ForwardData | None) -> ForwardData | None:
"""重叠循环：先发射第 N 轮的计算，再在 CPU 上处理第 N-1 轮的结果。
GPU 执行第 N 轮的同时，CPU 在做第 N-1 轮的结果处理和第 N+1 轮的调度，CPU 开销被藏起来。
能这样做的前提是：第 N 轮的输入 token 不需要 CPU 知道——它由第 N-1 轮的采样结果在 GPU 上
直接写进 token_pool（见 _forward）。
"""
self._defer_slot_free = True
if last_data is None: # 没有在途的 batch：之前推迟释放的请求槽可以放心复用了
self._release_deferred_slots()
blocking = not (last_data is not None or self.prefill_manager.runnable
or self.decode_manager.runnable)
for msg in self.receive_msg(blocking=blocking):
self._process_one_msg(msg)
forward_input = self._schedule_next_batch()
ongoing = None
if forward_input is not None:
with self.engine_stream_ctx:
self.engine.stream.wait_stream(self.stream) # 等本轮的元数据准备好
ongoing = (forward_input, self._forward(forward_input))
self._process_last_data(last_data)
return ongoing
normal_loop 比，只有一处不同：先发射本轮，再处理上一轮。overlap_loop 接收上一轮的数据、返回本轮的数据，run_forever 把它们串起来：
@torch.inference_mode()
def run_forever(self) -> NoReturn:
if ENV.DISABLE_OVERLAP_SCHEDULING:
with self.engine_stream_ctx:
self.engine.stream.wait_stream(self.stream)
while True:
self.normal_loop()
else:
data = None
while True:
data = self.overlap_loop(data)
在 GPU 上，"发射"只是把 kernel 放进 stream 的队列，立即返回；"处理"要等结果拷回 CPU。重叠循环里，处理第 1 轮时 GPU 已经在算第 2 轮，CPU 的开销被藏在 GPU 计算的时间里：
三个前提¶
能这样做，依赖前面几章早已埋下的三个设计。
第一，下一轮的输入不需要经过 CPU。 第 N+1 轮 decode 的输入是第 N 轮采样出的 token，发射第 N+1 轮时 CPU 还不知道它是什么。但 _forward 把采样结果直接在 GPU 上写进了 token pool 的下一个位置，第 N+1 轮再从 token pool 里取输入（第 7 章的 write_tuple 和 input_tuple）。GPU 按 stream 顺序执行，读一定发生在写之后。
第二，调度所需的状态在 CPU 上提前推进。 Engine.forward_batch 在发射后立即对每个请求调用 complete_one()，所以调度第 N+1 轮时，请求的 cached_len、device_len 已经是"第 N 轮之后"的值，可以正确地算出位置、分配 KV 页。唯一不知道的是第 N 轮的采样结果——也就是请求是否在第 N 轮遇到了 EOS。所以一个在第 N 轮遇到 EOS 的请求，会被多调度一轮。
第三，两条 stream。 调度器在自己的 stream 上准备元数据（把位置、下标等从锁页内存异步拷到 GPU），引擎在另一条 stream 上计算。发射前 self.engine.stream.wait_stream(self.stream) 让计算等元数据拷完；而调度器 stream 上下一轮的准备工作不必等当前的计算。ForwardInput 把这一轮用到的所有张量（包括下标）保存到处理结果之后，防止它们在 GPU 还没用完时被释放、内存被复用。
由此带来的四个问题¶
"CPU 状态提前一轮"和"请求可能被多调度一轮"这两点，让处理结果变得微妙。官方实现处理了其中一种情况（防止重复释放），本书在复刻时又发现了四个问题，示例和测试都能复现：
"""第 11 章：重叠调度的事件顺序，以及按官方逻辑判断结束时出现的两个问题。"""
from typing import List, Set
import torch
from minisgl.core import Req, SamplingParams
from minisgl.env import ENV
from minisgl.llm import LLM
from minisgl.message import DetokenizeMsg
from minisgl.scheduler.prefill import ChunkedReq
PATH = "models/Qwen3-0.6B"
KW = dict(dtype=torch.float32, max_running_req=8, num_page_override=512, max_seq_len_override=128)
def trace_events(overlap: bool) -> None:
ENV.DISABLE_OVERLAP_SCHEDULING.value = not overlap
llm = LLM(PATH, **KW)
step = {"n": 0}
orig_forward, orig_process = llm._forward, llm._process_last_data
def forward(fi):
step["n"] += 1
fi.batch.step = step["n"]
print(f" 发射第 {step['n']} 轮（{fi.batch.phase}）")
return orig_forward(fi)
def process(data):
if data is not None:
print(f" 处理第 {data[0].batch.step} 轮的结果")
return orig_process(data)
llm._forward, llm._process_last_data = forward, process
llm.generate(["The capital of France is"], SamplingParams(max_tokens=3, ignore_eos=True))
llm.shutdown()
print("普通循环：")
trace_events(overlap=False)
print("重叠循环：")
trace_events(overlap=True)
class UpstreamLogicLLM(LLM):
"""把 _process_last_data 换成官方的判断方式（其余不变），用来复现问题。"""
def _process_last_data(self, last_data) -> None:
if last_data is None:
return
batch, (_, next_tokens_cpu, copy_done) = last_data[0].batch, last_data[1]
copy_done.synchronize()
reply: List[DetokenizeMsg] = []
new_finished: Set[Req] = set()
with self.cache_manager.lazy_free_region():
for i, req in enumerate(batch.reqs):
if isinstance(req, ChunkedReq):
continue
req.append_host(next_tokens_cpu[i].unsqueeze(0))
token = int(next_tokens_cpu[i].item())
finished = not req.can_decode # 官方：按 can_decode 判断
if not req.sampling_params.ignore_eos:
finished |= token == self.eos_token_id
reply.append(DetokenizeMsg(uid=req.uid, next_token=token, finished=finished))
if finished and req not in self.finished_reqs: # 官方：只防重复释放
self.decode_manager.remove_req(req)
self._free_req_resources(req)
new_finished.add(req)
elif batch.is_prefill:
self.cache_manager.cache_req(req, finished=False)
self.finished_reqs = new_finished
self.send_result(reply)
def replies_of(cls, prompt: str, params: SamplingParams, eos: int | None = None):
ENV.DISABLE_OVERLAP_SCHEDULING.value = False
llm = cls(PATH, **KW)
if eos is not None:
llm.eos_token_id = eos
replies = []
orig = llm.offline_send_result
llm.send_result = lambda reply: (replies.extend(reply), orig(reply))
llm.generate([prompt], params)
llm.shutdown()
return [(m.next_token, m.finished) for m in replies]
print("\n问题一：max_tokens=3，发给 detokenizer 的 (token, finished)：")
for cls in (UpstreamLogicLLM, LLM):
print(f" {cls.__name__:<17}", replies_of(cls, "The capital of France is",
SamplingParams(max_tokens=3, ignore_eos=True)))
print("问题二：把 eos 设成第 3 个输出 token（576），max_tokens=8：")
for cls in (UpstreamLogicLLM, LLM):
print(f" {cls.__name__:<17}", replies_of(cls, "The capital of France is",
SamplingParams(max_tokens=8), eos=576))
普通循环：
发射第 1 轮（prefill）
处理第 1 轮的结果
发射第 2 轮（decode）
处理第 2 轮的结果
发射第 3 轮（decode）
处理第 3 轮的结果
重叠循环：
发射第 1 轮（prefill）
发射第 2 轮（decode）
处理第 1 轮的结果
发射第 3 轮（decode）
处理第 2 轮的结果
处理第 3 轮的结果
问题一：max_tokens=3，发给 detokenizer 的 (token, finished)：
UpstreamLogicLLM [(12095, False), (13, True), (576, True)]
LLM [(12095, False), (13, False), (576, True)]
问题二：把 eos 设成第 3 个输出 token（576），max_tokens=8：
UpstreamLogicLLM [(12095, False), (13, False), (576, True), (6722, False)]
LLM [(12095, False), (13, False), (576, True)]
问题一：结束标记提前了一个 token¶
官方用 finished = not req.can_decode 判断是否达到 max_tokens。但处理第 N 轮时，第 N+1 轮已经发射，device_len 已经被推进了一步，can_decode 反映的是"第 N+1 轮之后"的状态。于是 max_tokens=3 的请求在第 2 个 token 时就被标记为结束（上面的 (13, True)），第 3 个 token 又发了一条带结束标记的消息。
离线接口恰好不受影响（它记下了每一条消息的 token），但在线服务的前端收到第一条 finished=True 就结束响应，客户端只会收到 max_tokens - 1 个 token。修正方法是改用与调度节奏无关的量——CPU 上已经收到的 token 数：
问题二：EOS 之后多发一条过期消息¶
一个在第 N 轮遇到 EOS 的请求会被多调度一轮（前提二）。官方用 finished_reqs 记录上一轮结束的请求，处理第 N+1 轮时跳过它们的资源释放，但仍然为它们生成了回复消息——上面的 (6722, False)。前端已经删除了这个请求，会忽略这条消息；detokenizer 却会为它新建一个解码状态，而且永远等不到结束，造成一点内存泄漏。修正：在循环开头直接跳过这些请求的一切处理。
def _process_last_data(self, last_data: ForwardData | None) -> None:
if last_data is None:
return
batch, (_, next_tokens_cpu, copy_done) = last_data[0].batch, last_data[1]
copy_done.synchronize() # 等采样结果拷回 CPU；这也意味着这个 batch 以及之前的 batch 都算完了
self._release_deferred_slots()
reply: List[DetokenizeMsg] = []
new_finished: Set[Req] = set()
with self.cache_manager.lazy_free_region():
for i, req in enumerate(batch.reqs):
if isinstance(req, ChunkedReq): # 分块中的请求这一轮没有输出
continue
if req in self.finished_reqs:
# 重叠调度下，请求在上一轮已经结束（遇到 EOS），但结束的消息还没处理，这一轮就已经
# 把它带上了。这一轮的结果是多余的，直接丢弃（官方实现会多发一条消息，前端恰好会忽略）。
continue
next_token = next_tokens_cpu[i]
req.append_host(next_token.unsqueeze(0))
token = int(next_token.item())
# 判断是否达到 max_tokens 不能用 req.can_decode：重叠调度下，处理第 N 轮结果时
# 第 N+1 轮已经发射，req 的 device_len 已经被推进了一步，会提前一个 token 判定结束。
# 用 CPU 侧已经收到的 token 数判断，与调度节奏无关。
finished = len(req.input_ids) >= req.max_device_len
if not req.sampling_params.ignore_eos:
finished |= token == self.eos_token_id
reply.append(DetokenizeMsg(uid=req.uid, next_token=token, finished=finished))
if finished:
self.decode_manager.remove_req(req)
self._free_req_resources(req)
new_finished.add(req)
elif batch.is_prefill: # prefill 刚结束：立刻把提示词的 KV 交给前缀缓存
self.cache_manager.cache_req(req, finished=False)
self.finished_reqs = new_finished
self.send_result(reply)
问题三：请求槽过早复用¶
请求在第 N 轮遇到 EOS，处理第 N 轮结果时它的资源被释放，其中包括它在 page table / token pool 里占的那一行。但此刻第 N+1 轮正在 GPU 上跑，其中还有这个请求：它会在引擎 stream 上把采样结果写进 token pool 的这一行。如果下一轮调度立即把这一行分给新请求，调度器 stream 上对这一行的写入（新请求的提示词）就和引擎 stream 上的旧写入没有先后保证，新请求的输入可能被覆盖。
KV 页没有这个问题：新请求写 KV 的 batch 在引擎 stream 上一定排在第 N+1 轮之后。请求槽的问题在于它被调度器 stream写入。修正：重叠模式下，释放的请求槽推迟到"下一次处理结果"时才真正归还——那时 copy_done.synchronize() 已经保证在途的 batch 算完了。没有在途 batch 时（last_data is None）立即归还，否则只有一个请求槽的交互模式会卡住。
def _free_req_resources(self, req: Req) -> None:
# 请求槽（page table / token pool 的一行）：重叠调度下，GPU 上可能还有一个包含这个请求的
# batch 在跑（它会往 token_pool 的这一行写采样结果）。如果此刻把这一行分给新请求，调度器流上
# 对这一行的写入会与引擎流上的旧 batch 竞争。所以推迟到下一次处理结果、确认在途 batch 算完之后。
# KV 页不需要推迟：新请求写 KV 的 batch 在引擎流上一定排在旧 batch 之后。
if self._defer_slot_free:
self._deferred_slots.append(req.table_idx)
else:
self.table_manager.free(req.table_idx)
self.cache_manager.cache_req(req, finished=True)
问题四：prefill 在途时收到 abort¶
客户端断开连接会触发 abort。如果 abort 到达时，这个请求的 prefill 恰好已经发射、结果还没处理，官方的流程是：abort 释放请求（cache_req(finished=True)，提示词的 KV 进入 Radix Cache）；随后处理 prefill 结果时，请求没有结束，于是再调用一次 cache_req(finished=False)——对一个已经释放的请求。第二次插入发现前缀已在缓存中，把"重复的那份"释放掉，而这些页正是缓存里那份：同一批页既在空闲列表里又在缓存树里。之后它们可能被分给新请求，覆盖掉缓存的 KV；调度器空闲时的完整性检查也会报错（Radix Cache 下复现为 free_pages(1024) + cache_pages(5) != num_pages(1024)）。
修正只需一行：abort 时把请求加入 finished_reqs，在途 batch 的结果就会被问题二的逻辑整体丢弃。
与官方的差异
以上四处都是对官方重叠调度的修正，每一处都有对应的测试：把修正改回官方写法，测试就会失败。它们在 CPU 上的复现依赖"处理第 N 轮时第 N+1 轮已经发射"这个时序，与是否真的有 GPU 并发无关——问题三的测试直接检查"分配请求槽时，它不属于任何在途 batch 中的请求"这个不变量。
官方实现
scheduler/scheduler.py · Scheduler.overlap_loop（第 83～106 行）scheduler/scheduler.py · Scheduler._process_last_data（第 138～167 行）（finished = not req.can_decode、if finished and req not in self.finished_reqs）scheduler/scheduler.py · Scheduler._process_one_msg（第 169～198 行）（abort 分支）- 关闭重叠调度：环境变量
MINISGL_DISABLE_OVERLAP_SCHEDULING=1，官方 README 里的消融实验就用它
在 CPU 上意味着什么¶
CPU 上没有异步执行：NullStream、NullEvent 都是空操作，"发射"时计算就已经完成。所以在 CPU 上重叠调度不会更快，但调度逻辑完全相同，上面四个问题照样出现、照样被测试抓住。性能上的收益要在 GPU 上测：官方 README 建议用 MINISGL_DISABLE_OVERLAP_SCHEDULING=1 做消融。模型越小、每步 GPU 时间越短，CPU 开销的占比越大，重叠的收益越明显。第 21 章给出在 GPU 上做这个消融的方法。
测试¶
@pytest.mark.parametrize("cache,chunk,page_size", [("radix", 8192, 1), ("naive", 5, 1), ("radix", 7, 4)])
def test_overlap_equals_normal_and_hf(cache, chunk, page_size):
results = {}
for overlap in (True, False):
ENV.DISABLE_OVERLAP_SCHEDULING.value = not overlap
llm = build_llm(QWEN3, cache_type=cache, max_extend_tokens=chunk, page_size=page_size,
num_page_override=1024 // page_size)
results[overlap] = [o["token_ids"] for o in llm.generate(PROMPTS, greedy(10))]
llm.cache_manager.check_integrity()
llm.shutdown()
assert results[True] == results[False] == hf_greedy(QWEN3, PROMPTS, 10)
def test_abort_while_prefill_is_in_flight():
"""prefill 已经发射、结果还没处理时收到 abort：请求被释放，在途 batch 的结果必须丢弃。"""
from minisgl.message import AbortBackendMsg
ENV.DISABLE_OVERLAP_SCHEDULING.value = False
llm = build_llm(QWEN3, cache_type="radix")
calls = {"n": 0}
orig = llm.offline_receive_msg
def receive(blocking=False):
calls["n"] += 1
if calls["n"] == 2: # 第二轮开始时（prefill 正在"GPU 上"）中止它
return [AbortBackendMsg(uid=0)]
return orig(blocking)
llm.receive_msg = receive
replies = []
llm.send_result = lambda reply: replies.extend(reply)
llm.generate(PROMPTS[:1], greedy(8))
assert replies == [] # 被中止的请求没有任何输出
llm.cache_manager.check_integrity()
assert llm.table_manager.available_size == 8
llm.shutdown()
怎么讲清楚
讲重叠调度：先发射第 N+1 轮，再在 CPU 上处理第 N 轮的结果，调度、组批、反分词这些 CPU 开销就藏在 GPU 计算后面。前提有三：第 N+1 轮的输入 token 在 GPU 上由第 N 轮的采样结果直接写入（CPU 不必知道具体值）；请求状态在发射时提前推进；调度器和引擎用两条 stream。代价是 CPU 的状态领先 GPU 一轮：在第 N 轮采样出 EOS 的请求，可能已经被放进了第 N+1 轮，要丢弃它多算的结果；结束判断要用已经收到的 token 数，请求槽要等在途的 batch 算完再复用。收益在小模型、小 batch（GPU 一步很短）时最大。
练习¶
- 在
overlap_loop里，如果把"发射本轮"和"处理上一轮"的顺序对调回来，但保留last_data的传递方式，会得到什么？ - 问题二里的"多调度一轮"浪费了一次计算。能不能在调度第 N+1 轮之前就知道第 N 轮有没有 EOS？代价是什么？
- 为什么问题三只影响请求槽，而 KV 页可以立即复用？如果新请求命中了前缀缓存，调度器 stream 会往 page table 写入命中的位置，这会不会引出类似的问题？
参考答案
- 等价于普通循环（发射后立即等待结果），失去重叠的收益，但逻辑仍然正确。
- 可以在发射第 N+1 轮之前
synchronize第 N 轮的结果，但这正是重叠调度要避免的等待。另一种思路是接受这一轮浪费：一个请求最多多算一个 token，对吞吐的影响只和"每轮结束的请求数"成正比，通常可以忽略。 - KV 池只被引擎 stream 上的 kernel 写入，新请求的 KV 写入排在旧 batch 之后；请求槽对应的 token pool 行和 page table 行则会被调度器 stream 写入（新请求的提示词、命中的位置、新分配的页）。命中前缀时写的是新请求自己的那一行，只要这一行不是刚释放、仍被在途 batch 使用的，就没有问题——这正是推迟释放请求槽所保证的。
小结¶
- 重叠调度先发射本轮、再处理上一轮，CPU 开销藏在 GPU 计算之后。
- 前提：输入在 GPU 上由上一轮的采样结果写入；请求状态在发射时提前推进；调度器与引擎用两条 stream。
- 代价：CPU 状态领先一轮，请求可能在 EOS 之后被多调度一轮。
- 本书修正了官方的四个问题：结束判断改用已收到的 token 数；丢弃过期结果；请求槽推迟到在途 batch 算完再复用；abort 的请求加入
finished_reqs。
发现错误或有更好的写法？在 GitHub 上编辑此页提交问题