P07 · Scheduler：continuous batching + chunked prefill + 抢占
nano-vllm 的核心：token budget、并发上限、KV block 上限三重约束下，每步决定谁跑、跑几个 token——还要支持 chunked prefill 和 recompute 抢占。
mine/p07_scheduler.py
验收 ./check.sh p07
提示 A-16 · 参考解 C-16 solutions/p07_scheduler.py
题面速览
- 在系统中的位置：每个 LLM stage 进程的 busy loop 每步调用一次
schedule()，输出“本步每个序列算几个 token”交给 model runner；算完后update()把采样结果写回。 - 三重约束：
max_num_batched_tokens（本步 token 总数上限，决定激活显存和单步延迟）；max_num_seqs（同时 running 的序列数上限）；KV block：持有 n 个 token（已算 + 本步要算）的序列需要ceil(n / block_size)个 block，总数不超过num_blocks。 - 调度规则：
budget = max_num_batched_tokens
1) running（按顺序）：n = min(剩余未算 token, budget) # decode 时为 1；chunked prefill 未完时更多
block 不够 → 抢占 running 中最后一个（最晚加入）：释放 block、num_computed=0、放回 waiting 队首
→ 重试；如果被抢占的就是它自己，则跳过它
2) waiting（严格 FCFS）：len(running) < max_num_seqs 且 budget > 0 时，n = min(剩余, budget)
第一个 block 放不下的请求就停止（后面的不许插队）
- update：每个被调度的序列
num_computed += n；只有当它算完了全部 token（prefill 最后一块或 decode）才调用 sample_fn 追加一个 token；遇到 eos 或达到 max_tokens 就结束并释放 block。被抢占的序列下次重新计算 prompt + 已生成的 token（recompute），已生成的输出保留不变。
例子（budget=16，三个 prompt 长 10）
step1: s0:10 s1:6 # s1 只做了一部分 prefill
step2: s0:1 s1:4 s2:10 # s0 decode；s1 prefill 剩 4；s2 整段
step3: s0:1 s1:1 s2:1
验收要点：token budget 与 chunked prefill 的精确调度结果 · max_num_seqs 限制 · FCFS：队首放不下时后面的小请求也不能插队 · KV 极小时发生抢占，所有请求最终正确完成且输出不变 · EOS 提前结束并释放 block。
直觉与类比
waiting = 门口排队的客人（严格按到店顺序）；running = 已坐下的桌子；max_num_seqs = 一共几张桌子；KV block = 桌面面积（越吃越久餐具越多，只增不减，直到他走）；token budget = 厨房每一轮能出多少个菜——一桌 10 人的大团一轮出不完就先出一部分（chunked prefill）。抢占 = 桌面实在不够时，请最后坐下的那桌先回到队伍最前面：他们吃掉的（已生成的 token）不用吐出来，但餐具被收走了（KV 释放），回来得重新摆一遍（recompute）。
抢占不是取消请求：丢的只是缓存（KV），不是进度（已生成的 token）。用户完全看不出发生过抢占：最终输出与从不抢占时逐 token 相同，只是慢了一点——这是本题最容易写错、也最值得自己验证的性质。
框架视角 · 两个阶段 · 三本预算 · 六条不变量
P07 是单线程纯状态机，难点是预算记账和队列搬动必须同时正确。
| 预算 | 谁持有 / 怎么计量 | 扣 | 退 |
|---|---|---|---|
| token budget | 局部变量，每次 schedule() 重置 | 每成功调度一个序列 budget -= n，n = min(num_tokens − num_computed, budget) |
本步不退；下一步重新满额 |
| running 槽位 | len(self.running)，上限 max_num_seqs |
阶段 2 接纳一个就 append 到末尾（末尾 = 最低优先级 = 抢占首选） | 抢占或结束时 running.remove |
| KV block | blocks_used 与每个 seq.num_blocks 两处同时记 |
_try_grow(seq, num_computed + n)：算出增量 need，放得下才同时改两处，否则一点不改返回 False |
抢占 / 结束：blocks_used -= seq.num_blocks; seq.num_blocks = 0 |
| 阶段 | 遍历什么 | 放不下怎么办 | 终止条件 |
|---|---|---|---|
| 阶段 1 running | self.running，按下标 while（抢占会改列表） |
抢占 running[-1] 然后重试同一个序列；受害者是自己 → 跳过、下标不前进 | 下标到末尾，或 budget == 0 |
| 阶段 2 waiting | 只看 self.waiting[0]（FCFS） |
break，不许跳过它 | waiting 空 / budget == 0 / running 满 |
| 不变量 | 违反时的症状 |
|---|---|
I1 blocks_used == sum(s.num_blocks for s in running) ≤ num_blocks；全部结束后为 0 |
block 泄漏：后面谁都调度不进来，或超发导致真实系统 OOM |
I2 sum(schedule().values()) ≤ max_num_batched_tokens，且每个 n ≥ 1 |
batch 超预算 → 激活显存爆；n == 0 的序列白占一行 |
I3 0 ≤ num_computed ≤ num_tokens；只有相等时才允许采样 |
chunked prefill 中途就采样 → 多出 token、提前达到 max_tokens |
| I4 抢占只重置 num_computed、num_blocks、status；output 必须保留 | 最终输出与无抢占时不同（最难查） |
| I5 被抢占者进 waiting 队首；waiting 内部始终是到达顺序 | 被抢占的请求饿死；压力下顺序乱掉，延迟长尾爆炸 |
| I6 每个序列同时只在 waiting、running、finished 之一 | 同一请求被调度两次，或结束后仍被调度 |
易错点
| bug | 症状 |
|---|---|
| 按 num_tokens 而不是 num_computed+n 算 block | chunked prefill 时提前占满 block，无谓抢占 |
| 被抢占后没把 num_computed 归零 | 重新进入时只调度 1 个 token：KV 已随 block 释放却没有重算（模型读到别人的 KV，不报错） |
| 抢占时把 output 也清空了 | 最终 output 与无抢占时不同 |
| 放回 waiting 尾部 | 被抢占的序列饿死，或在压力下顺序乱掉 |
| prefill 未完成也调用 sample_fn | output 多出 token，max_tokens 提前达到 |
| 阶段 2 用 continue 跳过放不下的 | 违反 FCFS：例 3 中 a3 会被调度 |
| 遍历 running 用 for | 抢占时 pop 让迭代器跳元素或访问到已被抢占的序列 |
| update 不知道本步调度了谁 | update 没有参数，必须把 schedule() 结果存进实例字段（返回副本） |
| 先判 eos 再 append | eos 本身也是输出的一部分：先 append 再判 tok == eos or len(output) >= max_tokens |
逐步图：抢占别人（例 2）
{a1: 2, a2: 3}：a1 1 块、a2 2 块 → blocks_used = 3（满了）
blocks：a1a2a2；waiting []
a1 要持有 3 个 token → 需要 2 块，+1 → 4 > 3 ✗ → 抢占 running[-1] = a2：释放 2 块、num_computed=0、放回 waiting 队首，num_preemptions = 1 → 重试 a1 成功。
阶段 2：a2 需要 ceil(4/2)=2 块，2+2 > 3 → break。返回 {a1: 1}。
blocks：a1a1·；waiting [a2]
{a1: 1}：a1 第 3 个输出，结束并释放 → blocks_used = 0。a2 仍放不下（同一步 a1 还占着）。
{a2: 4}：prompt 3 + 已生成 1 = 4 个 token 一次算完，采样出第 2 个输出，结束。
最终 a1.output = [100, 101, 102]，a2.output = [200, 201]——与没有抢占时完全一样。
三级提示
L1I/O 示例与关键问题I/O 示例 + 关键问题 + 细节（轻度剧透）
例 1：token budget + chunked prefill（num_blocks=100, block_size=4, max_num_seqs=8, budget=8；a1(prompt 6, max 3)、a2(prompt 5, max 2)、a3(prompt 3, max 1)；每步 schedule 后立刻 update）：
| step | schedule() | update 后 (num_computed / num_tokens / num_blocks) | blocks_used | 本步结束 |
|---|---|---|---|---|
| 1 | {a1: 6, a2: 2} | a1 6/7/2；a2 2/5/1（prefill 没完，不采样）；a3 仍 waiting | 3 | - |
| 2 | {a1: 1, a2: 3, a3: 3} | a1 7/8/2；a2 5/6/2（这一步才采样第一个 token）；a3 采样 1 个即达 max | 4 | a3 |
| 3 | {a1: 1, a2: 1} | a1 输出 3 个、a2 输出 2 个，均达 max | 0 | a1, a2 |
step 1：a1 先拿 6，剩 budget 2 给 a2。step 2：running 先排：a1 decode 1，a2 剩余 3；budget 还剩 4，a3 整段进入。block 数按“已算 + 本步要算”算：step 1 后 a2 只持有 2 个 token → 1 块。
例 2：见逐步图。例 3：FCFS 队首阻塞（num_blocks=5, block_size=4, budget=100；a1(9)、a2(9)、a3(1)）：schedule() → {"a1": 9}，blocks_used 3，waiting ["a2", "a3"]——a2 需要 3 块但只剩 2 块 → 停；a3 只要 1 块、放得下，但也不许插队。
关键问题
- 一个正在 decode 的序列（刚采样了一个 token）下一步
num_tokens − num_computed等于多少？chunked prefill 做了一半的呢？ - 判断 block 放不放得下时按多少个 token 算：num_tokens 还是 num_computed + n？例 1 step 1 的 a2 分别几块？
- running 中的 X 放不下时被抢占的是谁？抢占后继续处理 X 还是跳过？被抢占的恰好是 X 自己呢——下标要不要前进？
- 被抢占的序列哪些字段重置、哪些保留？重新进入时要算多少个 token？放到 waiting 哪一端？
- 阶段 2 遇到第一个放不下的：continue 还是 break？max_num_seqs 限制哪个阶段？
- update 里什么情况下被调度的序列不调用 sample_fn？结束时除了改 status 还要做哪两件事？
- update 没有参数，“上一步调度了谁、各几个 token”存哪？
细节
num_tokens = prompt_len + len(output)；decode 时剩余就是 1（刚采样的那个还没算 KV）——decode 和 prefill 用同一个公式，这就是 vLLM v1 调度器统一的地方。- block 增量
need = ceil((num_computed + n) / bs) − seq.num_blocks（≤0 不用动）。 - 两个阶段共用一个 budget；budget 用完后后面的 running 序列本步不被调度（不是抢占）。
- 被抢占的是最晚进来的，但仍比 waiting 里所有序列都早，所以放队首；它在同一步的阶段 2 也会被尝试。
- 结束：status=FINISHED、移出 running、释放 block、放进 finished。
L2思路与常见坑思路拆分与常见坑（中度剧透）
- 先写三个 helper：
_blocks_for(n)（向上取整）、_try_grow(seq, n_tokens) -> bool（放得下就同时更新 blocks_used 和 seq.num_blocks，否则什么都不改）、_preempt(seq)（从 running 移除、释放 block、num_computed=0、status=WAITING、waiting.appendleft、计数+1）。 - 阶段 1 用下标 while；内层
while not _try_grow(seq, ...)：抢占 running[-1]；若它就是 seq 则跳出并标记“本序列被跳过”。只有成功才写 scheduled、扣 budget、下标+1。（Python 的while … else很适合） - 阶段 2：
while waiting and budget > 0 and len(running) < max_num_seqs：看队首，_try_grow 失败就 break；成功则 popleft、RUNNING、append、记录 n、扣 budget。 - 把 schedule 结果存进
self._scheduled（返回副本）；update 遍历它，结束后清空。update 中先建{seq_id: seq}，因为遍历中会把结束的序列从 running 删除。
L3详解：答案与追踪详解：答案、逐步追踪、骨架提示（重度剧透）
关键问题答案
- decode：1；chunked prefill：prompt 剩下没算的部分。
- num_computed + n。例 1 step 1 的 a2：按 num_tokens=5 是 2 块，按 0+2=2 是 1 块；正确是 1 块。
- running 的最后一个。抢占别人后重试 X；X 自己被抢占则跳过，下标不前进（running 已变短）。
- 重置：num_computed=0、num_blocks=0（并从 blocks_used 减掉）、status=WAITING；保留：output。重新算 prompt_len + len(output) 个 token。放队首。
- break。max_num_seqs 只在阶段 2 作为准入条件。
- num_computed 加完后仍 < num_tokens（prefill 块没做完）时不采样。结束时：从 running 移除并释放 block、放进 finished（并返回）。
- 实例字段，例如
self._scheduled。
追踪：抢占自己（num_blocks=3, block_size=2, budget=64；a1、a2 都是 prompt 2、max 3）：
| step | 阶段 1 | 阶段 2 | 返回 | blocks_used（update 后） | a2.output |
|---|---|---|---|---|---|
| 1 | running 为空 | a1 n=2 → 1 块；a2 n=2 → 1 块 | {a1:2, a2:2} | 2 | [200] |
| 2 | a1 n=1，需 2 块，+1 → 3 ✓；a2 n=1，需 2 块，+1 → 4 ✗ → 抢占 running[-1] = ?? | ?? | ?? | 2 | ?? |
| 3 | a1 n=1，不需新块 | a2 仍放不下 | {a1:1} | 0（a1 结束） | [200] |
| 4 | running 为空 | a2 n=?? | ?? | 2 | ?? |
| 5 | a2 decode | - | {a2:1} | 0 | [200, 201, 202] |
填空答案
step 2 victim 是 a2 自己 → 跳过；阶段 2 中 a2 n=3（2+1），需 2 块，2+2 > 3 → break；返回 {a1: 1}；a2.output 仍是 [200]。step 4：n=3，返回 {a2: 3}，output 变为 [200, 201]。num_preemptions = 1。
自检：把例 1、例 2、例 3 和上面的追踪各写成几行脚本跑一遍；每步断言 blocks_used <= num_blocks 且 blocks_used == sum(s.num_blocks for s in running)；所有序列结束后 blocks_used 回到 0。
面试官追问 & 对照源码
- recompute 抢占 vs swap 抢占的取舍？
- 多模态请求的 encoder 计算如何计入 budget？
- async scheduling 下 num_computed 里包含了 placeholder token，update 要怎么改？
对照源码：vllm/v1/core/sched/scheduler.py