vLLM Chunked Prefill 策略源码精读

为什么需要 Chunked Prefill

在 Default 调度策略中，prefill 和 decode 是互斥的：一旦有 prefill 请求被调度，decode 请求就不会被安排进同一个 batch。这导致了两个问题：

GPU 利用率低：长 prompt 的 prefill 独占整个 batch，此时已在运行的 decode 请求必须等待
Inter-token latency 高：decode 请求被 prefill 阻塞，用户感知到的生成速度下降

Chunked Prefill 的核心思想是：将 prefill 和 decode 请求放入同一个 batch。长 prompt 不再一次性处理完，而是被拆成多个 chunk，每个 chunk 与当前正在运行的 decode 请求一起组成 batch。

核心设计原则

Chunked Prefill 的优先级顺序为：Decode → Swap-in → New Prefill。先保证已在运行的请求不被饿死，再处理 swap 恢复，最后安排新的 prefill 请求。这与 Default 策略的 Prefill → Decode → Swap-in 顺序完全相反。

整体调度流程

_schedule_chunked_prefill 方法（L885-992）按以下四步执行：

graph TD A["创建 SchedulingBudget"] --> B["Step 1: _schedule_running
调度 decode + 已有 chunked prefill"] B --> C{"发生了抢占?"} C -->|"是"| E["跳过 swap-in"] C -->|"否"| D["Step 2: _schedule_swapped
调度 swap-in 请求"] D --> F["Step 3: _schedule_prefills
调度新 prefill 请求"] E --> F F --> G["更新三队列状态"] G --> H["返回 SchedulerOutputs"]

scheduler.py — _schedule_chunked_prefill 入口 L885-916

def _schedule_chunked_prefill(self):
    """Schedule queued requests.

    Chunked prefill allows to chunk prefill requests, batch them together
    with decode requests. This policy 1. schedule as many decoding requests
    as possible. 2. schedule chunked prefill requests that are not
    finished. 3. schedule swapped request. 4. schedule new prefill
    requests.
    """
    budget = SchedulingBudget(
        token_budget=self.scheduler_config.max_num_batched_tokens,
        max_num_seqs=self.scheduler_config.max_num_seqs,
    )
    curr_loras: Set[int] = set()

    remaining_waiting, prefills = (self.waiting,
                                   SchedulerPrefillOutputs.create_empty())
    remaining_running, running_scheduled = (
        self.running, SchedulerRunningOutputs.create_empty())
    remaining_swapped, swapped_in = (
        self.swapped, SchedulerSwappedInOutputs.create_empty())

    # Decoding should be always scheduled first by fcfs.
    fcfs_policy = PolicyFactory.get_policy(policy_name="fcfs")

Budget 初始化的差异

注意 Chunked Prefill 创建 budget 时不预先注册 running 请求的 seq 数，而 Default 策略会先 budget.add_num_seqs() 遍历所有 running 请求。这是因为 Chunked Prefill 在 _schedule_running 调用时会自行在 budget 中注册（enable_chunking=True 分支）。

Step 1: 调度 Running 请求

第一步调度所有正在运行的请求（包括 decode 和被分块但尚未完成的 prefill）：

scheduler.py — 调度 running 请求 L920-929

# Decoding should be always scheduled first by fcfs.
remaining_running, running_scheduled = self._schedule_running(
    self.running,
    budget,
    curr_loras,
    fcfs_policy,
    enable_chunking=True)  # 关键: enable_chunking=True

关键参数 enable_chunking=True 的影响贯穿整个调度链路：

在 _schedule_running（L382-519）中，enable_chunking=True 会走 budget.add_num_seqs() 分支（L503-507），因为此时 budget 尚未注册这些 running 请求
在 _get_num_new_tokens（L1278-1306）中，enable_chunking=True 允许将 token 数裁剪为 budget.remaining_token_budget()
Running 队列中可能包含上一轮未完成的 chunked prefill（seq_group.is_prefill() 为 True），它们被归入 prefill_seq_groups 而非 decode_seq_groups

_schedule_running 按 FCFS 优先级排序后逐个处理。对每个 seq_group，先检查是否能分配新 slot（_can_append_slots）：

如果不能 → 触发抢占，从队尾弹出低优先级请求
如果可以 → 调用 _append_slots 分配 block，将请求加入调度结果

Step 2: 调度 Swapped 请求

只有当 Step 1 没有发生抢占时，才会尝试 swap-in：

scheduler.py — 条件性调度 swap-in L933-942

# Schedule swapped out requests.
# If preemption happens, it means we don't have space for swap-in.
if len(running_scheduled.preempted) + len(
        running_scheduled.swapped_out) == 0:
    remaining_swapped, swapped_in = self._schedule_swapped(
        self.swapped, budget, curr_loras, fcfs_policy)

这个条件检查的逻辑很直观：如果连当前正在运行的请求都需要被抢占（说明 GPU 显存紧张），那么从 CPU swap 回 GPU 显然是不可行的。

_schedule_swapped（L521-635）的处理流程：

按 FCFS 优先级排序 swapped 队列
对每个请求检查 block_manager.can_swap_in()
如果返回 AllocStatus.LATER → 停止（没有足够的 GPU block）
如果返回 AllocStatus.NEVER → 标记为 FINISHED_IGNORED
检查 budget 是否足够 → 调用 _swap_in 执行实际的块交换

Step 3: 调度新 Prefill 请求

最后调度 waiting 队列中的新请求：

scheduler.py — 调度新 prefill L947-952

# Schedule new prefills.
remaining_waiting, prefills = self._schedule_prefills(
    self.waiting, budget, curr_loras, enable_chunking=True)

注意这里同样传入 enable_chunking=True，这意味着：

新 prefill 请求不需要一次性获得足够的 token budget 来处理整个 prompt
如果 budget 剩余空间不足以放下整个 prompt，会只取 remaining_token_budget() 个 token
未处理完的 token 在下一轮 step 中继续处理（此时 seq_group 已在 running 队列中）

与 Default 策略的关键区别

Default 策略中 enable_chunking=False，要求 num_new_tokens == num_prompt_tokens（L700-701），即整个 prompt 必须一次性调度。而 Chunked Prefill 允许部分调度，这是两者最本质的差异。

Token 分块机制

分块的核心逻辑在 _get_num_new_tokens 方法（L1278-1306）：

scheduler.py — _get_num_new_tokens L1278-1306

def _get_num_new_tokens(self, seq_group: SequenceGroup,
                        status: SequenceStatus, enable_chunking: bool,
                        budget: SchedulingBudget) -> int:
    num_new_tokens = 0
    seqs = seq_group.get_seqs(status=status)
    for seq in seqs:
        num_new_tokens += seq.get_num_new_tokens()
    assert num_new_tokens > 0
    # Chunk if a running request cannot fit in.
    # If number of seq > 1, it means it is doing beam search in a
    # decode phase. Do not chunk in that case.
    if enable_chunking and len(seqs) == 1:
        num_new_tokens = min(num_new_tokens,
                             budget.remaining_token_budget())
        if num_new_tokens == budget.remaining_token_budget() \
                and seq_group.is_prefill():
            # floor and align to 256
            num_new_tokens = (num_new_tokens // 256) * 256
    return num_new_tokens

分块规则：

只对单序列请求分块：如果 len(seqs) > 1（beam search），不分块
按 budget 剩余量裁剪：min(num_new_tokens, remaining_token_budget())
256 对齐：如果恰好用完剩余 budget 且是 prefill，向下对齐到 256 的倍数。这是为了让 GPU kernel 更高效地执行

为什么要 256 对齐

GPU 的矩阵运算（如 FlashAttention）在 token 数是 2 的幂次或 256 的倍数时效率更高。对齐可以避免 padding 开销，提升计算效率。

举个例子：假设 max_num_batched_tokens=2048，当前已调度 1500 个 decode token，剩余 budget 为 548。一个新 prefill 请求有 4096 个 prompt token：

num_new_tokens = min(4096, 548) = 548
因为 548 == remaining_budget 且 is_prefill → 对齐到 256：(548 // 256) * 256 = 512
本轮只处理 512 个 token，剩余 3584 个在后续轮次处理

队列状态更新

调度完成后，更新三个队列的状态（L958-974）：

scheduler.py — 队列更新 L958-974

# Update waiting requests.
self.waiting = remaining_waiting
self.waiting.extendleft(running_scheduled.preempted)
# Update new running requests.
self.running = remaining_running
self.running.extend([s.seq_group for s in prefills.seq_groups])
self.running.extend(
    [s.seq_group for s in running_scheduled.decode_seq_groups])
self.running.extend(
    [s.seq_group for s in running_scheduled.prefill_seq_groups])
self.running.extend(
    [s.seq_group for s in swapped_in.decode_seq_groups])
self.running.extend(
    [s.seq_group for s in swapped_in.prefill_seq_groups])
# Update swapped requests.
self.swapped = remaining_swapped
self.swapped.extend(running_scheduled.swapped_out)

关键区别：Chunked Prefill 的 running 队列中会同时包含 decode_seq_groups 和 prefill_seq_groups。而 Default 策略的 running 队列只有 decode 请求（因为 assert len(running_scheduled.prefill_seq_groups) == 0）。

被抢占的请求放回 waiting 队列头部（extendleft），保证下次优先被重新调度。

与 Default 策略对比

维度	Default 策略	Chunked Prefill 策略
调度优先级	Prefill → Decode → Swap-in	Decode → Swap-in → Prefill
Prefill/Decode 混合	互斥，不混合	可在同一 batch 中混合
enable_chunking	False	True
长 prompt 处理	必须一次性处理完	可拆分为多个 chunk
Budget 初始化	预先注册 running 的 seq 数	不预注册，在 _schedule_running 中动态注册
Running 队列内容	只有 decode 请求	decode + 未完成的 chunked prefill
SchedulerOutputs 中的 prefill_seq_groups	只来自 waiting 队列	来自 waiting + running + swapped
适用场景	短 prompt、批量离线推理	长 prompt、在线服务、低延迟要求

设计权衡

Chunked Prefill 优先调度 decode 请求，确保正在生成的请求不被阻塞。代价是新请求的首 token 延迟（TTFT）可能略高，因为 prefill 被推迟了。但在实际在线服务中，稳定的 inter-token latency 比快速的 TTFT 更重要。

do_sample 控制

Chunked Prefill 引入了一个重要的控制逻辑：只有当 prefill 完全完成时才需要采样。

scheduler.py — do_sample 判断 L1063-1075

do_sample = True
if seq_group.is_prefill():
    seqs = seq_group.get_seqs()
    # Prefill has only 1 sequence.
    assert len(seqs) == 1
    # In the next iteration, all prompt tokens are not computed.
    # It means the prefill is chunked, and we don't need sampling.
    # NOTE: We use get_len instead of get_prompt_len because when
    # a sequence is preempted, prefill includes previous generated
    # output tokens.
    if (token_chunk_size + seqs[0].data.get_num_computed_tokens() <
            seqs[0].data.get_len()):
        do_sample = False

逻辑解释：

token_chunk_size：本轮要处理的 token 数
get_num_computed_tokens()：已经处理过的 token 数
get_len()：总 token 数（包括 prompt + 已生成的 output）
如果 已处理 + 本轮处理 < 总量，说明 prefill 还没完成，不需要采样

这个设计避免了在中间 chunk 上做无意义的采样，节省了计算资源。do_sample=False 会传递给 ModelRunner，使其跳过 Sampler 层。

配置与调优

启用 Chunked Prefill 的关键配置：

scheduler.py — 策略选择 L994-999

def _schedule(self) -> SchedulerOutputs:
    """Schedule queued requests."""
    if self.scheduler_config.chunked_prefill_enabled:
        return self._schedule_chunked_prefill()
    else:
        return self._schedule_default()

参数	作用	建议值
`enable_chunked_prefill`	是否启用 chunked prefill	`True`（在线服务推荐）
`max_num_batched_tokens`	每个 batch 的最大 token 数（即 token budget）	2048-8192，取决于显存
`max_num_seqs`	每个 batch 的最大序列数	256（默认值）