杨艺韬讲堂

Appearance

第11章 分块预填充与混合批处理 ​

"The art of scheduling is the art of saying 'not all at once'."

本章要点

  • 理解预填充阻塞问题：为什么长 Prompt 会影响解码请求的延迟
  • 掌握分块预填充的工作原理：将预填充拆分为多个可控大小的块
  • 深入 V1 统一调度如何天然支持分块预填充
  • 理解混合批处理在 FlashAttention3 中的实现
  • 认识分块大小的选择策略及其对延迟/吞吐的影响

11.1 预填充阻塞问题 ​

回忆第 3 章的内容：预填充和解码对 GPU 的使用模式截然不同。预填充是计算密集型——一次处理大量 Token；解码是内存带宽密集型——每次只处理 1 个 Token。

问题出在它们共享 GPU 时间。假设一个批次中有：

  • 1 个新请求，需要预填充 4096 个 Token
  • 50 个老请求，每个需要解码 1 个 Token

如果不做分块，这一步要处理 4096 + 50 = 4146 个 Token。预填充的 4096 Token 主导了计算时间——可能需要 200ms。在这 200ms 内，50 个正在解码的请求全部被"冻结"，用户感到输出中断了 200ms。

分块预填充的解法：将 4096 Token 切成 4 块（每块 1024），每步只处理一块。

现在每步只需 ~80ms，解码请求不再被长时间阻塞。代价是新请求的首 Token 延迟从 200ms 增加到了 4 × 80ms = 320ms——但这通常是可接受的折中。

11.2 V1 的自然实现 ​

V1 统一 Token 调度的美妙之处在于：分块预填充不需要任何特殊处理。

源码版本：本节基于 vLLM v0.8.5，核心文件 vllm/v1/core/sched/scheduler.py。

让我们看调度器中实现分块预填充的真实代码（scheduler.py:176-185）：

python
# vllm/v1/core/sched/scheduler.py:176-185
while req_index < len(self.running) and token_budget > 0:
    request = self.running[req_index]
    num_new_tokens = (request.num_tokens_with_spec -
                      request.num_computed_tokens)

    # 长预填充阈值：超过这个长度就分块
    if (0 < self.scheduler_config.long_prefill_token_threshold <
            num_new_tokens):
        num_new_tokens = (
            self.scheduler_config.long_prefill_token_threshold)

    # 核心：用 min 限制不超过剩余 budget
    num_new_tokens = min(num_new_tokens, token_budget)

    # ... 调度该请求 ...
    token_budget -= num_new_tokens

关键在 min(num_new_tokens, token_budget) 这一行——如果一个请求需要 4096 token 但 budget 只剩 1024，就只调度 1024 个 token。下一步从 num_computed_tokens + 1024 处继续。整个分块逻辑就是这一个 min 调用——没有特殊的分块代码路径。

回忆调度器的输出：{request_id: num_tokens}。对于上面的场景：

python
# Step 1
{"new_req": 1024, "req_1": 1, "req_2": 1, ..., "req_50": 1}
# Step 2
{"new_req": 1024, "req_1": 1, "req_2": 1, ..., "req_50": 1}
# Step 3
{"new_req": 1024, "req_1": 1, "req_2": 1, ..., "req_50": 1}
# Step 4
{"new_req": 1024, "req_1": 1, ..., "req_50": 1}  # 预填充完毕，下一步 new_req 也是 1

调度器只是限制了每步给新请求的 Token 数量。Worker 端不需要知道"这 1024 Token 是预填充的前半段还是后半段"——它只管计算这 1024 个 Token 的注意力。请求的 num_computed_tokens 记录了已处理到哪里，下次从断点继续。

这就是统一抽象的力量——分块预填充、全量预填充、纯解码，对 Worker 来说都是"处理 N 个 Token"，没有区别。

11.3 断点续传机制 ​

分块预填充的一个关键问题是：如何记住一个请求"已经计算到哪里了"？

V1 的 Request 对象维护了一个 num_computed_tokens 字段：

python
# vllm/v1/request.py（简化）
class Request:
    prompt_token_ids: list[int]     # 完整的 Prompt Token 序列
    num_computed_tokens: int = 0     # 已完成预填充的 Token 数

    @property
    def num_remaining_tokens(self):
        return len(self.prompt_token_ids) - self.num_computed_tokens

每次调度器给这个请求分配 N 个 Token 的预算：

  1. 取 prompt_token_ids[num_computed_tokens : num_computed_tokens + N] 送入 GPU
  2. GPU 计算完成后，num_computed_tokens += N
  3. 当 num_computed_tokens == len(prompt_token_ids) 时，预填充完成，请求进入解码阶段

这个机制非常简洁——不需要保存任何中间计算结果（KV Cache 已经存在于 GPU 显存中），只需要一个整数记录进度。

被抢占后的恢复 ​

如果一个正在分块预填充的请求被抢占（显存不足），它的部分 KV Cache 块可能被回收。恢复时有两种情况：

情况 1：前缀缓存命中——被回收的块仍然在前缀缓存中（引用计数降为 0 但未被驱逐）。KV Cache Manager 通过哈希链找到这些块，直接将引用计数加回来。num_computed_tokens 不需要回退——因为 KV Cache 数据仍然有效。

情况 2：块已被驱逐——KV Cache 数据已被覆盖。num_computed_tokens 回退到仍然有效的最后一个块的末尾位置，然后从那里重新开始预填充。

这就是分块预填充与前缀缓存协同工作的优雅之处——抢占的代价不再是"从头来"，而是"从断点来"。

11.4 FlashAttention3 的混合批处理 ​

分块预填充的一个技术挑战是：预填充 Token 和解码 Token 在同一步中如何高效地执行注意力计算？

预填充 Token 需要对一大段 Token 做自注意力（算力密集），解码 Token 需要与长序列做 KV Cache 注意力（带宽密集）。两者的计算特征差异很大。

FlashAttention3（vllm/v1/attention/backends/flash_attn.py）支持变长序列的混合批处理——在一次内核调用中处理不同长度的多个序列。它通过 cu_seqlens（累积序列长度）数组标记每个序列的边界，内核内部根据边界选择不同的计算路径。

具体来说，FlashAttention3 的输入是：

query:      [total_tokens, num_heads, head_dim]    # 所有请求的 Q 拼接
key:        分页 KV Cache（通过块表间接访问）
value:      分页 KV Cache（通过块表间接访问）
cu_seqlens_q: [0, 1024, 1025, 1026, 1090, ...]    # 每个请求的 Q 长度累积和
cu_seqlens_k: [0, 1024, 2048, 4096, 1090, ...]    # 每个请求的 KV 长度累积和

cu_seqlens_q 和 cu_seqlens_k 的差异体现了预填充和解码的混合：

  • 预填充请求：q_len = 1024（本次处理 1024 个 Token），kv_len = 1024（已有的 KV Cache 长度）
  • 解码请求：q_len = 1（本次只处理 1 个新 Token），kv_len = 2048（该请求之前的所有 Token）

FlashAttention3 内核会根据每个序列的 q_len 选择不同的计算策略——长 Q 序列使用更多的线程并行，短 Q 序列（解码）使用更少的线程但更高的内存带宽利用。

11.5 定量分析：分块预填充的延迟影响 ​

让我们用具体数字理解分块预填充的效果。

场景：A100 GPU，Llama-2-70B（TP=4），50 个并发解码请求 + 1 个新请求（4096 Token Prompt）。

不分块：

步骤处理 TokenGPU 时间解码请求延迟
Step 14096 + 50 = 4146~200ms200ms（被阻塞）
Step 251~25ms25ms（恢复正常）

解码请求在 Step 1 遭遇 200ms 的延迟尖峰——用户感觉输出"卡"了一下。

分块（chunk_size=1024）：

步骤处理 TokenGPU 时间解码请求延迟
Step 11024 + 50 = 1074~60ms60ms
Step 21024 + 50 = 1074~60ms60ms
Step 31024 + 50 = 1074~60ms60ms
Step 41024 + 50 = 1074~60ms60ms
Step 551~25ms25ms

分块后，单步最大延迟从 200ms 降到 60ms（3.3x 改善），但新请求的 TTFT 从 200ms 增到 240ms（4 步 × 60ms）。这是经典的延迟尖峰 vs TTFT 权衡。

long_prefill_token_threshold

vLLM v0.8.5 新增了 long_prefill_token_threshold 配置（scheduler.py:181-184）。只有当请求的待处理 token 数超过这个阈值时才强制分块——短请求直接全量预填充，避免不必要的分块开销：

python
# scheduler.py:181-184
if (0 < self.scheduler_config.long_prefill_token_threshold <
        num_new_tokens):
    num_new_tokens = self.scheduler_config.long_prefill_token_threshold

11.6 分块大小的选择 ​

max_num_batched_tokens 参数控制了每步的最大 Token 数，间接决定了分块大小：

  • 较大值（如 8192）→ 新请求的预填充块更大，首 Token 延迟（TTFT）更低，但解码请求被阻塞的时间更长
  • 较小值（如 512）→ 解码请求的延迟更稳定，但新请求的 TTFT 更高（需要更多步才能完成预填充）

最佳值取决于工作负载特征。如果大部分请求是短对话（Prompt < 256 Token），分块没有必要——全量预填充的延迟本就很低。如果有大量长文档（Prompt > 4096 Token），分块的收益就很明显。

11.6 本章小结 ​

  • 预填充阻塞——长 Prompt 独占 GPU 会冻结解码请求
  • 分块预填充——将长 Prompt 切块，与解码请求混合处理
  • V1 自然支持——统一 Token 调度 + num_computed_tokens 断点续传
  • FlashAttention3——单次内核调用处理混合批次
  • 分块大小——在 TTFT 和解码延迟之间权衡

源码导航

  • 调度器（分块逻辑）：vllm/v1/core/sched/scheduler.py
  • FlashAttention3 后端：vllm/v1/attention/backends/flash_attn.py

基于 VitePress 构建

Copyright © 2024-present 杨艺韬