从源码读 vLLM Scheduler：一轮调度如何分配请求与 KV Cache

vLLM 的 Scheduler 并不执行模型，也不实现 attention kernel。它更像推理引擎的调度台：每一轮决定哪些请求前进、各前进多少 token，以及有限的 KV Cache 是否还能容纳它们。

vLLM Scheduler 调度闭环

这张图刻意只保留主干。waiting 与 running 是请求状态；KVCacheManager 是内存准入者；worker 只执行 Scheduler 写好的“派工单”。下文对应 vLLM commit 7fbd44c。

先建立一个正确的心智模型

一个请求的状态可以先简化为：

1	waiting → running → finished

waiting：请求已被接收，但尚未获得计算和 KV Cache 资源；
running：请求已拥有 KV block，可以继续 prefill 或 decode；
finished：命中 EOS、stop 条件、max_tokens，或被取消；相关 KV block 会被回收。

这里最反直觉的一点在 Scheduler.schedule()：它没有严格的“prefill 阶段”和“decode 阶段”。每个请求只维护已经计算到的位置 num_computed_tokens，以及需要追上的位置 num_tokens_with_spec。调度器的目标只是让前者向后者推进。

# scheduler.py
num_new_tokens = (
    request.num_tokens_with_spec
    + request.num_output_placeholders
    - request.num_computed_tokens
)
num_new_tokens = min(num_new_tokens, token_budget, input_budget - draft_slots)

因此，长 prompt 的 chunked prefill、常规 decode、推测解码，本质上都是“这个请求本轮还需要推进多少 token”的不同来源。

`schedule()`：先保住正在运行的请求

调度函数开头建立两类预算：

1 2	token_budget = self.max_num_scheduled_tokens input_budget = self.scheduler_config.max_num_batched_tokens

接着先遍历 self.running：

# First, schedule the RUNNING requests.
while req_index < len(self.running) and token_budget > 0:
    request = self.running[req_index]
    ...

这解释了 vLLM 的一个核心取舍：已经开始生成的请求通常优先继续前进，以降低 token 间延迟；剩余预算才用于接纳新请求。

假设每轮预算是 8 token：

A：10-token prompt，已计算 8 token，位于 running
B：2-token prompt，仍在 waiting

本轮：先给 A 2 token；剩余 6 token 才轮到 B。

如果 A、B 都已经完成 prefill，它们通常会在后续轮次各推进一个或少量 decode token。Scheduler 关心的是 token 总量，而不是“本轮只能跑 prefill”或“本轮只能跑 decode”。

新请求如何从 waiting 进入 running

对于 waiting 队首，scheduler 不会立即出队，而是先 peek：确认并发上限、LoRA 限制、token 预算和 KV Cache 都满足，再真正 pop 并加入 running。

request = request_queue.peek_request()
...
new_blocks = self.kv_cache_manager.allocate_slots(...)
if new_blocks is None:
    break

request = request_queue.pop_request()
self.running.append(request)
request.status = RequestStatus.RUNNING

“取出请求”不是删除请求，而是状态迁移：

1 2	waiting = [A, B] → waiting = [B] running = [] → running = [A]

Prefix Cache：复用的是 KV，不是答案

新请求入场前先走 _get_local_prefix_cache_hit()，其默认路径会调用 KVCacheManager.get_computed_blocks()。

假设 block size 为 4，两个请求拥有同一段 8-token 系统提示词：

1 2	A：[系统提示词 8 token][解释 TCP] B：[系统提示词 8 token][解释 HTTP]

A 已执行后，系统提示词对应的两个 KV block 已缓存为 P1、P2。B 到达时，流程是：

B 的 token blocks
  → 计算前缀 hash
  → 在 BlockPool 的 hash 索引中查找
  → 命中 P1、P2
  → B 复用 P1、P2，只计算“解释 HTTP”

这里 hash 的对象是输入 token block 与其前缀上下文，而不是庞大的 K/V tensor。hash 仅用于索引；命中后复用的才是真实 K/V memory。这样既避免逐元素比较 tensor，也避免把相同系统提示词重复 forward。

源码还刻意限制命中不超过 request.num_tokens - 1：即使整段 prompt 都命中，最后一个位置仍要重新计算以获得下一个 token 的 logits。

`allocate_slots()`：KV Cache 的准入闸门

KVCacheManager.allocate_slots() 是 Scheduler 与物理 KV Cache 之间最关键的接口。它不是简单地“分配 N 个 block”，而是按以下顺序工作：

结合请求已计算 token、prefix 命中 token、新 token 和 speculative lookahead，算出需要多少 slot；
先释放 sliding-window attention 已不再需要的旧 block；
计算真正缺少的新 block 数；
在保留 watermark 后检查空闲 block 是否足够；
成功则把复用 block 挂到请求上，再分配新 block，并把已确定的完整 block 写入 prefix cache。

仍以上面的 B 为例：

复用前缀：P1、P2，共 8 token
新后缀：2 token
block size：4

B 的逻辑 block 表： [P1][P2][B3]
真正新分配的只有：B3

核心准入判断非常直接：

available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks
required_blocks = num_blocks_to_allocate + watermark_blocks
if required_blocks > available_blocks:
    return None

None 的含义不是请求出错，而是“现在不能安全地给它安排这轮计算”。新请求通常继续留在 waiting；对已经 running 的请求，scheduler 才可能尝试抢占。

KV 不够时，抢占的到底是什么

当 running 请求申请新 slot 得到 None，schedule() 会挑选一个 running 请求作为牺牲者，然后调用 _preempt_request()。

self._free_request_blocks(request)
request.status = RequestStatus.PREEMPTED
request.num_computed_tokens = 0
self.waiting.prepend_request(request)

它做的不是取消用户请求，而是：释放该请求当前占用的私有 KV block，标记为 PREEMPTED，再放回 waiting 队列前部等待恢复。恢复时会重新走 prefix-cache 查询，因此仍被缓存的完整前缀可复用，私有尾部则可能需要重算。

抢占对象由 scheduler_config.policy 决定：priority 策略选择优先级更低的请求；同优先级时，较晚到达者先被抢占。非 priority 路径则从 running 列表尾部挑选。

`SchedulerOutput`：调度器与 worker 的边界

调度器不直接调用模型。它在函数末尾组装 SchedulerOutput：

scheduled_new_reqs：worker 第一次见到的请求，携带完整信息
scheduled_cached_reqs：worker 已认识的请求，只发送增量
num_scheduled_tokens：每个请求本轮推进多少 token
finished_req_ids：上一轮后可清理的请求

以 Metal 后端为例，runner 先导入新请求，再更新已有请求的 block table，最后将 prefill 与 decode 工作合成执行 batch。调度决策与 kernel 实现由此解耦。

为什么还有 `num_in_flight_tokens`

schedule() 返回前会调用 _update_after_schedule()，先乐观推进请求状态：

1 2	request.num_computed_tokens += num_scheduled_token request.num_in_flight_tokens += num_scheduled_token

前者表示 scheduler 已经安排到的位置，后者表示仍在 GPU 上、结果尚未返回的工作量。等 worker 返回 ModelRunnerOutput 后，update_from_output() 会减少 num_in_flight_tokens、追加采样 token、检查 EOS/stop/max_tokens；完成请求则从 running 移除并释放 KV block。

这组状态让 scheduler 可以在异步执行时提前准备下一轮，而不会把同一段 token 重复派发。

结语

理解 vLLM Scheduler，先抓住这条闭环即可：

waiting / running 请求
  → token 预算与 KV 准入
  → SchedulerOutput
  → worker 执行
  → 输出回写、完成或继续

推测解码、多模态 encoder cache、分布式 KV transfer、Mamba/hybrid cache 和 LoRA 都是在这条主线上增加约束或额外状态；不需要先理解它们，已经可以读懂调度器的核心行为。

先建立一个正确的心智模型

schedule()：先保住正在运行的请求