实现一个 Scheduler:让多个请求在同一个推理引擎里流动起来

1. 这一篇要解决什么问题

上一篇文章里,我们讨论了静态 batching 为什么不适合直接用于在线 LLM serving。

单请求推理吃不满 GPU,所以我们需要 batching。

但是固定 batch 又会带来新的问题:请求长度不同会造成 padding 浪费,生成长度不同会让 batch 逐渐变空,新请求无法及时加入,长请求还可能阻塞短请求。

所以我们需要一种更灵活的方式。

它不是提前固定一批请求,然后让它们一起跑到结束。

它是在每一轮推理之前,重新决定这一轮应该执行哪些请求。

这就是 Scheduler 要做的事情。

这一篇我们开始实现 mini vLLM 的第一个调度器。

它暂时不会很复杂,但会具备一个推理引擎调度器最核心的结构:

waiting 队列:还没有做 prefill 的请求
running 队列:已经完成 prefill,正在 decode 的请求
finished 集合:已经结束的请求

每一轮 engine step,Scheduler 都要做一个决策:

哪些 waiting 请求进入 prefill
哪些 running 请求继续 decode
这一轮总 token 数是否超过预算
这一轮总 sequence 数是否超过限制
哪些请求应该保留到下一轮
哪些请求应该结束并释放资源

到这一篇结束时,我们的 mini vLLM 会从“能处理一个 Sequence”,变成“能管理多个 Sequence”。

这一步非常关键。

因为只要有了 Scheduler,后面的 continuous batching、PagedAttention、KV cache block manager、prefix caching,都会有一个可以接入的中心位置。

2. Scheduler 不只是一个队列

很多系统里,调度器看起来像一个队列。

请求先进来,排队,然后一个个被处理。

如果只是普通后端服务,这种理解可能够用。

但 LLM 推理里的 Scheduler 要复杂得多。

因为一个请求不是执行一次就结束。

一个 LLM 请求通常会经历这样的过程:

进入 waiting
执行 prefill
进入 running
执行 decode
继续 decode
继续 decode
直到 finished

也就是说,同一个请求会被调度很多次。

它不是一次性消费掉的任务,而是一个会持续存在、不断推进的状态对象。

这也是为什么上一篇我们要先引入 Sequence。

Scheduler 调度的不是一个普通函数调用,而是一个个可暂停、可恢复、可继续推进的 Sequence。

每一轮调度时,Scheduler 都要回答两个问题。

第一个问题是:新请求要不要进来?

这对应 waiting 队列里的 prefill。

第二个问题是:老请求要不要继续生成?

这对应 running 队列里的 decode。

这两个问题都很重要。

如果一直处理新请求,已经开始流式输出的请求就会卡住。

如果一直处理老请求,新请求就迟迟拿不到首 token。

所以 Scheduler 本质上是在平衡两件事:

让新请求尽快开始
让老请求持续输出

这就是 LLM serving 调度的核心矛盾之一。

3. Engine、Scheduler、ModelRunner 的边界

在写 Scheduler 之前,我们先明确几个模块的职责。

Engine 是总控。

它接收请求,创建 Sequence,调用 Scheduler,调用 ModelRunner,然后把模型输出写回 Sequence。

Scheduler 负责决策。

它不直接调用模型,也不直接做采样。它只决定这一轮哪些 Sequence 要运行,以及它们应该以什么阶段运行。

ModelRunner 负责执行。

它接收 Scheduler 给出的 batch,构造模型输入,执行 forward,返回 logits 和 KV cache。

这三个模块的关系大概是:

Engine
  调用 Scheduler.schedule()
  得到本轮要执行的 prefill sequences 和 decode sequences
  调用 ModelRunner.execute()
  得到模型输出
  更新 Sequence 状态

为什么要这样拆?

因为 Scheduler 不应该知道模型内部细节。

它不应该关心 attention 怎么算,也不应该关心 logits 怎么采样。

它只关心资源和状态。

同样,ModelRunner 也不应该知道请求排队策略。

它只负责把给定 batch 跑出来。

这种边界划分会让后续扩展更容易。

例如,将来要改调度策略,只需要改 Scheduler。

将来要换 attention backend,只需要改 ModelRunner 或更底层的执行模块。

将来要接入 block manager,也主要是在 Scheduler 和 cache manager 之间增加资源检查。

4. Scheduler 的输入和输出

Scheduler 的输入是什么?

最直接的输入就是它内部维护的队列状态:

waiting sequences
running sequences
finished sequences

除此之外,它还需要一些配置。

例如:

max_num_batched_tokens
max_num_seqs
max_num_prefill_seqs

这些配置用来限制每一轮最多能调度多少工作。

Scheduler 的输出是什么?

它应该输出一个调度结果,而不是直接修改所有东西。

可以抽象成:

class SchedulerOutputs:
    prefill_sequences: list[Sequence]
    decode_sequences: list[Sequence]

更完整一点,还可以包含:

本轮使用了多少 token
本轮调度了多少 sequence
有哪些请求被跳过
有哪些请求因为资源不足不能进入

第一版可以先简单一些,只返回 prefill 和 decode 两类 Sequence。

这已经足够让 Engine 知道下一步该做什么。

5. 为什么要分 prefill batch 和 decode batch

在 mini vLLM 第一版里,我们可以把 prefill 和 decode 分开处理。

也就是说,Scheduler 输出两个列表:

prefill batch
decode batch

prefill batch 里的 Sequence 还没有 KV cache,需要输入完整 prompt。

decode batch 里的 Sequence 已经有 KV cache,只需要输入 last token。

这两类请求的模型输入形状不同。

prefill 的输入长度可能是:

[batch_size, prompt_length]

decode 的输入长度通常是:

[batch_size, 1]

它们的成本模型也不同。

一个 prompt 长度为 4096 的 prefill 请求,可能比几百个 decode token 都重。

所以 Scheduler 必须知道每个 Sequence 当前处于什么阶段。

不能把所有请求都当成同一种任务。

更深一点看,prefill 和 decode 对用户体验的影响也不同。

prefill 决定首 token 延迟。

decode 决定流式输出是否连续。

如果 prefill 太少,新请求迟迟没有响应。

如果 decode 太少,正在输出的请求会一顿一顿。

一个好的调度器,必须能同时照顾这两类体验。

6. 第一版调度策略

第一版 Scheduler 可以先采用一个简单策略:

优先调度 running 请求做 decode
如果还有 token budget,再调度 waiting 请求做 prefill

这个策略的直觉是:

已经开始生成的请求,应该尽量保持输出连续。

用户看到模型开始输出之后,会期待它持续输出。如果中途卡顿,体验会很明显变差。

所以我们先让 running 请求继续 decode。

然后,如果这一轮还有剩余 token budget,再从 waiting 队列里挑新请求做 prefill。

伪代码大概是:

def schedule():
    budget = SchedulingBudget(max_num_batched_tokens, max_num_seqs)

    decode_sequences = schedule_running_sequences(budget)
    prefill_sequences = schedule_waiting_sequences(budget)

    return SchedulerOutputs(
        prefill_sequences=prefill_sequences,
        decode_sequences=decode_sequences,
    )

这个策略不一定最优,但它是一个非常好的起点。

因为它简单、稳定,而且能体现调度器的基本结构。

后面我们可以在这个基础上继续讨论饥饿、抢占、chunked prefill 和更复杂的优先级策略。

7. SchedulingBudget:调度器的资源账本

Scheduler 需要一个资源账本。

否则它不知道这一轮还能不能继续往 batch 里塞请求。

这个资源账本可以叫 SchedulingBudget

它至少记录两个值:

当前已经使用了多少 token
当前已经调度了多少 sequence

同时它有两个上限:

max_num_batched_tokens
max_num_seqs

一个简化结构可以是:

class SchedulingBudget:
    max_num_batched_tokens: int
    max_num_seqs: int
    num_batched_tokens: int
    num_seqs: int

它需要提供几个判断方法:

can_schedule(num_new_tokens, num_new_seqs)
add(num_new_tokens, num_new_seqs)

为什么需要 token budget?

因为 batch size 不足以描述 LLM 推理成本。

一个 sequence 可能是一个 decode token。

也可能是一个 4096 tokens 的 prefill prompt。

它们都是一个 sequence,但成本完全不同。

所以 Scheduler 不能只限制 sequence 数量,还必须限制 token 数量。

这就是 token budget 的意义。

8. decode 请求如何消耗 budget

对于 running sequence 来说,每一轮 decode 通常只消耗一个输入 token。

因为它只需要把 last token 输入模型。

所以每个 running sequence 的本轮 token 成本可以近似认为是 1。

假设现在有 64 个 running sequence。

那么这一轮 decode 的 token 成本大概是:

64 tokens

这就是 decode 的好处。

单个请求 decode 很小,GPU 吃不满。

但大量请求一起 decode,就可以形成一个比较大的 batch。

Scheduler 要做的,就是尽量把这些 running sequence 聚合起来。

伪代码可以是:

for sequence in running_queue:
    if budget.can_schedule(num_new_tokens=1, num_new_seqs=1):
        decode_sequences.append(sequence)
        budget.add(num_new_tokens=1, num_new_seqs=1)

这里先用最简单的 FIFO 方式遍历 running 队列。

后面如果要引入优先级,可以在这里改。

9. prefill 请求如何消耗 budget

prefill 请求不一样。

对于 waiting sequence 来说,本轮要处理的是完整 prompt。

如果 prompt 长度是 1024,那么它会消耗 1024 个 token budget。

如果 prompt 长度是 8192,那么它可能一口气占满整个 budget。

所以 prefill 对调度器的影响要大得多。

伪代码可以是:

for sequence in waiting_queue:
    prompt_len = sequence.get_prompt_len()

    if budget.can_schedule(num_new_tokens=prompt_len, num_new_seqs=1):
        prefill_sequences.append(sequence)
        budget.add(num_new_tokens=prompt_len, num_new_seqs=1)

这里会出现一个问题。

如果队首 waiting sequence 特别长,超过了剩余 budget,后面短请求要不要跳过它先执行?

这就是调度策略问题。

严格 FIFO 的做法是:队首请求放不下,就停止调度 waiting 队列。

更激进的做法是:跳过长请求,继续寻找后面能放下的短请求。

前者更公平。

后者平均延迟可能更好。

第一版建议先采用严格 FIFO。

原因不是它最好,而是它最容易解释,也最容易验证正确性。

等我们有了 benchmark,再讨论是否要优化这个策略。

10. 为什么 running 通常优先于 waiting

我们刚才选择了“running 优先”的策略。

这背后有几个考虑。

第一,running 请求已经占用了 KV cache。

如果它们长时间不 decode,资源就会被占着但没有产生输出。

第二,running 请求通常正在流式返回。

用户已经看到模型开始说话,如果后续 token 迟迟不来,会明显感觉卡顿。

第三,decode 每个请求只消耗一个 token budget。

相比长 prefill,它对单轮 budget 的冲击更小。

第四,running 请求继续 decode 之后,可能很快结束并释放资源。

例如某个请求只剩几个 token 就结束了,让它继续跑完,反而能尽快释放 KV cache。

所以第一版优先 running 是合理的。

但它也有风险。

如果 running 队列长期很大,waiting 请求可能一直进不来,首 token 延迟会变差。

这就是饥饿问题。

所以这个策略后面一定要改进。

但在 mini vLLM 的早期版本里,我们先接受这个限制,重点把调度链路跑通。

11. waiting 请求饥饿问题

什么是饥饿?

饥饿就是某些请求虽然一直在队列里,但长期得不到调度。

在我们的第一版策略里,饥饿最可能发生在 waiting 队列。

假设 running 队列一直很满,每一轮都耗尽 max_num_seqs。

那么 waiting 请求就没有机会做 prefill。

用户会一直等不到第一个 token。

这在服务里是不可接受的。

所以真实系统通常需要一些保护机制。

例如:

限制每轮 decode 最多占用多少 budget
为 prefill 保留一部分 budget
设置 waiting 请求最大等待时间
定期强制调度一部分 waiting 请求
把长 prefill 切成 chunk

第一版 Scheduler 不一定马上实现这些机制,但文章里必须把问题说清楚。

因为这正是调度器设计中最有思考价值的地方。

一个策略解决了某个问题,往往会引入另一个问题。

优先 running 可以让输出更平滑,但可能牺牲首 token 延迟。

优先 waiting 可以让新请求更快开始,但可能让老请求流式输出卡顿。

调度器真正难的地方,不是把请求放进列表,而是在这些目标之间做权衡。

12. 为什么暂时不做抢占

真实推理系统可能会做抢占。

例如,当 KV cache 不够时,可以把某些低优先级请求暂停,把它们的资源释放给高优先级请求。

但第一版 mini vLLM 不建议做抢占。

原因有三个。

第一,抢占会增加状态复杂度。

一个请求不再只是 waiting、running、finished,还可能出现 swapped、paused、preempted 等状态。

第二,抢占和 KV cache 管理强相关。

如果没有 block manager,抢占之后如何释放和恢复 cache 会变得很别扭。

第三,抢占不是理解调度器的第一步。

我们当前最重要的是跑通:

waiting 进入 prefill
prefill 后进入 running
running 持续 decode
结束后进入 finished

把这条主线跑通,比过早引入复杂状态更重要。

后面实现 block manager 之后,再讨论抢占会更自然。

13. Scheduler 如何管理 running 队列

running 队列里的 Sequence 每一轮都可能发生变化。

某个 Sequence decode 后,如果还没结束,就继续留在 running。

如果已经结束,就移出 running,进入 finished。

所以 Engine 在执行完 ModelRunner 后,需要把结果写回 Sequence,然后更新 Scheduler 的队列。

一种简单做法是:

Scheduler 只负责选出本轮要执行的 running sequences
Engine 执行并更新它们
Engine 再通知 Scheduler 哪些完成了
Scheduler 维护 running 队列

也可以让 Scheduler 提供一个方法:

def free_finished_sequences():
    ...

但要注意职责边界。

判断 Sequence 是否 finished,通常应该由 Sequence 自己完成。

Scheduler 只根据状态移动队列。

例如:

for sequence in scheduled_sequences:
    if sequence.is_finished():
        move_to_finished(sequence)
    else:
        keep_in_running(sequence)

这会让逻辑更清晰。

14. prefill 之后为什么要进入 running

waiting sequence 被调度做 prefill 之后,会发生几件事。

模型处理完整 prompt。

生成初始 KV cache。

根据最后一个位置的 logits 采样出第一个 token。

把这个 token 加入 generated tokens。

如果请求还没结束,它就进入 running。

如果它刚好结束,比如 max_new_tokens 等于 1,或者第一个 token 就是 EOS,它可以直接进入 finished。

所以 prefill 之后不是无条件进入 running。

更准确的逻辑是:

prefill 完成
append 第一个 token
检查停止条件

如果 finished:
    放入 finished
否则:
    放入 running

这个细节很重要。

因为用户可能设置很小的 max_new_tokens。

也可能模型第一步就输出 EOS。

调度器和 engine 都不能假设 prefill 后一定还有 decode。

15. decode 之后为什么可能释放资源

decode 每轮会追加一个 token。

追加之后,Sequence 可能结束。

结束原因可能是:

达到 max_new_tokens
遇到 EOS
命中 stop words
用户取消

一旦结束,它就不应该继续留在 running。

更重要的是,它占用的 KV cache 也应该被释放。

当前阶段我们还没有实现自己的 cache manager,所以释放动作可以先只是概念上的。

但从 Scheduler 的角度,要形成这个习惯:

finished sequence 不再参与调度
finished sequence 应该释放推理资源

后面有了 block manager,这一步会变成真正的:

释放它持有的 KV cache blocks
归还给 free block pool

这就是为什么 Scheduler 和 KV cache 管理会天然耦合。

调度决定谁能继续运行。

运行产生 cache 增长。

结束触发 cache 释放。

16. Engine.step 应该如何组织

现在我们可以把 Engine 的 step 流程串起来。

每一轮 step 大概包含这些阶段:

调用 Scheduler 生成调度结果
根据 prefill_sequences 构造 prefill batch
根据 decode_sequences 构造 decode batch
调用 ModelRunner 执行模型
对每个 Sequence 采样新 token
更新 Sequence 的 KV cache 和状态
把新增文本包装成输出
更新 Scheduler 队列
返回本轮 outputs

伪代码可以是:

def step():
    scheduler_outputs = scheduler.schedule()

    model_outputs = model_runner.execute(scheduler_outputs)

    request_outputs = process_model_outputs(model_outputs)

    scheduler.update_after_step()

    return request_outputs

这个 step 是整个推理引擎的心跳。

以后所有能力都会挂在这个循环上。

continuous batching 是让每一轮 schedule 都能动态改变 batch。

PagedAttention 是让每个 Sequence 的 KV cache 可以用 block table 管起来。

streaming output 是让每一轮 step 的增量 token 立即返回给用户。

benchmark 是测这个循环的吞吐和延迟。

所以 Engine.step 的结构一定要设计清楚。

17. 调度结果为什么不要直接是一个 tensor

有人可能会觉得,Scheduler 最终不就是要给模型一个 tensor 吗?

为什么不让 Scheduler 直接返回 input_ids?

原因是 Scheduler 不应该直接构造模型输入。

模型输入的构造会涉及很多模型相关细节。

例如:

input_ids
position_ids
attention_mask
past_key_values
seq_lens
block_tables
slot_mapping

这些东西会随着执行 backend 变化。

今天我们用 Hugging Face 的 past_key_values

后面我们会用自己的 block manager。

再后面可能会接入 Triton kernel 或 FlashAttention backend。

如果 Scheduler 直接构造 tensor,它就会和模型执行细节强耦合。

更好的方式是,Scheduler 返回 Sequence 列表和阶段信息。

ModelRunner 再根据这些信息构造真正的模型输入。

也就是说:

Scheduler 决定谁运行
ModelRunner 决定怎么运行

这个边界非常重要。

18. 第一版 Scheduler 的核心接口

第一版 Scheduler 可以有几个核心接口。

添加请求:

scheduler.add_sequence(sequence)

执行调度:

outputs = scheduler.schedule()

更新状态:

scheduler.update(sequence)

判断是否还有未完成请求:

scheduler.has_unfinished_sequences()

内部维护三类容器:

waiting_queue
running_queue
finished_sequences

这里不需要追求复杂。

重点是让请求状态流动起来。

一开始,一个新 Sequence 被加入 waiting。

调度后,它从 waiting 移出。

prefill 成功后,如果没结束,进入 running。

decode 若干轮后,进入 finished。

这条流动链路比接口数量更重要。

19. 一个简单调度过程示例

假设系统里来了三个请求。

A: prompt 100 tokens,最多生成 4 tokens
B: prompt 200 tokens,最多生成 2 tokens
C: prompt 50 tokens,最多生成 3 tokens

假设每轮 token budget 是 256,max_num_seqs 是 4。

初始状态:

waiting: A, B, C
running: empty
finished: empty

第 1 轮调度:

running 为空,所以先看 waiting。

A 的 prompt 长度 100,可以放入。

B 的 prompt 长度 200,如果再放入,总 token 会变成 300,超过 budget,所以不能放。

C 排在 B 后面,严格 FIFO 下暂时不跳过 B。

所以第 1 轮:

prefill: A
decode: empty

A prefill 后生成第一个 token,进入 running。

第 2 轮调度:

running 中有 A。

先调度 A decode,消耗 1 token。

剩余 budget 还有 255。

waiting 队首是 B,长度 200,可以放入。

所以第 2 轮:

decode: A
prefill: B

A decode 后生成第二个 token。

B prefill 后生成第一个 token,进入 running。

第 3 轮调度:

running 中有 A、B。

它们各 decode 一个 token,消耗 2 token。

waiting 队首是 C,长度 50,可以放入。

所以第 3 轮:

decode: A, B
prefill: C

这时三个请求都进入系统了。

后面的轮次里,A、B、C 会根据各自 max tokens 陆续结束。

这个例子展示了 continuous batching 的雏形。

每一轮 batch 都不一样。

有时只有 prefill。

有时只有 decode。

有时 prefill 和 decode 同时存在。

有的请求进入,有的请求继续,有的请求退出。

这就是在线 LLM 推理和静态 batch 的根本区别。

20. 第一版策略的不足

第一版 Scheduler 虽然能工作,但它有很多不足。

第一个不足是长 prompt 队头阻塞。

如果 waiting 队首是超长 prompt,它可能因为 budget 不够而阻塞后面的短请求。

第二个不足是 waiting 饥饿。

如果 running 队列持续占满 max_num_seqs,waiting 请求可能很难进入。

第三个不足是没有显式 KV cache 资源判断。

当前版本只看 token budget 和 sequence 数量,但真实系统还要看 KV cache block 是否足够。

第四个不足是没有 chunked prefill。

长 prompt 必须一次性 prefill 完,这可能会造成 decode 卡顿。

第五个不足是没有优先级。

所有请求都按 FIFO 处理,无法区分普通请求和高优先级请求。

这些不足不是失败。

它们恰恰是后续文章要继续演进的方向。

好的系统不是一开始就完美,而是在每一版中清楚地暴露问题,然后围绕问题做下一轮设计。

21. 调度器的设计思路总结

写 Scheduler 时,最容易犯的错误是只把它当成“从队列里拿请求”。

但在 LLM 推理系统里,Scheduler 至少要同时考虑四件事。

第一,请求状态。

waiting 请求需要 prefill,running 请求需要 decode,finished 请求需要移除。

第二,计算预算。

每一轮不能无限塞 token,否则延迟和显存都会失控。

第三,用户体验。

新请求需要尽快拿到首 token,老请求需要稳定输出。

第四,资源生命周期。

请求开始后会占用 KV cache,请求结束后要释放 KV cache。

所以 Scheduler 不是一个简单容器。

它是整个推理系统的决策中心。

第一版 Scheduler 不需要复杂,但必须从一开始就保留这些思考位置。

22. 小结

这一篇我们实现了 mini vLLM 的第一个 Scheduler 设计。

它维护三类请求状态:

waiting
running
finished

它在每一轮 step 中构造两类 batch:

prefill batch
decode batch

它使用 token budget 和 sequence 数量限制本轮工作量。

它采用一个简单策略:

优先调度 running 做 decode
剩余预算调度 waiting 做 prefill

这个策略让已经开始输出的请求尽量保持流畅,同时也允许新请求在预算允许时进入系统。

不过,这个 Scheduler 仍然只是第一版。

它还没有解决长 prompt 队头阻塞、waiting 饥饿、KV cache 显存判断、chunked prefill 和优先级调度等问题。

但它已经完成了最重要的一步:

让多个 Sequence 在同一个推理引擎里流动起来。

下一篇文章,我们会进一步拆分 prefill 和 decode 的执行路径。

我们会看到,prefill 和 decode 不只是调度状态不同,它们在模型输入构造、KV cache 更新、性能瓶颈和用户体验上的关注点都不同。

理解这一点之后,mini vLLM 的 engine 结构会变得更加清晰。