vllm 中 async scheduling - 学习让我快乐

当前 vLLM 的异步调度（async scheduling）实现仍较为复杂. 在深入剖析其具体代码之前, 本文将以一条具体的查询请求为例, 手动模拟调度器 scheduler 与 worker 的执行流程, 并标注其中的关键状态变化. 通过这种方式纲举目张, 为后续深入理解代码细节提供一个清晰的参照框架.

S0.schedule

S0 对应着发生在 R 上的第 0 次 step; S0.schedule 对应着该 step 的 schedule 部分. R.S0.num_computed_tokens 为在 S0 schedule 时看到的 R.num_computed_tokens 值. R.S0.num_new_tokens 表明在 S0 请求 R 被调度待计算的 token 数.

R.S0.num_computed_tokens = R.num_computed_tokens; R.S0.num_new_tokens;
allocate_slots(R, R.S0.num_new_tokens); cache_blocks(R, R.S0.num_computed_tokens + R.S0.num_new_tokens)
NewRequestData.from_request(R), computed=R.S0.num_computed_tokens
R.num_computed_tokens = R.S0.num_computed_tokens + R.S0.num_new_tokens < R.num_prompt_tokens,

S1.schedule

在 async scheduling 开启时, vllm 会紧接着开始调度发生在 R 上第 1 次 step.

R.S1.num_computed_tokens = R.num_computed_tokens; R.S1.num_new_tokens = R.num_prompt_tokens - R.S1.num_computed_tokens;
allocate_slots(R, R.S1.num_new_tokens); cache_blocks(R, R.S1.num_computed_tokens + R.S1.num_new_tokens = R.num_prompt_tokens)
R CachedRequestData.computed = R.S1.num_computed_tokens; R CachedRequestData num_output_tokens=0
R.num_computed_tokens = R.S1.num_computed_tokens + R.S1.num_new_tokens = R.num_prompt_tokens.
R.num_output_placeholders = 1; R.spec_token_ids = [-1] * self.num_spec_tokens. 这里模拟 step S1 执行结束之后 R 上发生的状态变化. num_output_placeholders = 1 意味着假设了 step S1 结束之后会为 R 生成一个 output token, 当然这个不是假设, 这个是明确的事实, S1 肯定会生成 1 个 output token 得. R.spec_token_ids = [-1] * self.num_spec_tokens, 假设了 step S1 结束之后会为 R 生成 num_spec_tokens 个 draft tokens; 这个确实是假设, 实际上 step S1 可能并不为 R 生成任何 draft tokens.

S0.update_from_output

之后 vllm 不会在尝试调度 step S2, 而是会等待 step S0 结束. 即在 step L 调度时, step L - 2 刚刚执行结束, step L - 1 刚刚下发给 worker 尚未执行结束; 也即此时 scheduler 侧看到的状态是 step [0, L - 2] 步的真实, 已更新状态; 以及第 L−1 步的 假设状态, 既然是假设, 那就有可能是错的, vllm scheduler/worker 侧都有根据实际结果矫正的逻辑.

R 无状态变更, 尚未完成 prefill, R 没有对应的 new_token_ids.

S2.schedule

之后假设 S1 生成 num_spec_tokens 个 draft tokens, 开始 step S2 的调度.

R.S2.num_computed_tokens = R.num_computed_tokens = R.num_prompt_tokens; R.S2.num_new_tokens = 1 + s2; 这里 假设了 S1 生成 num_spec_tokens 个 draft tokens, 由于 token budget 或者 structured output mask 等原因这里只调度了 s2 个 draft tokens, 0 <= s2 <= num_spec_tokens.
allocate_slots(R, R.S2.num_new_tokens); cache_blocks(R, R.num_prompt_tokens); 这里 cache_blocks 被截断, 对应注释 ‘ensuring only “finalized” tokens are cached’, 也比较合理, 因为这时 S1 尚未执行结束. R.S2 new tokens 都不知道是啥呢, block hash 都无法计算, 自然无法 cache.
R CachedRequestData computed = R.S2.num_computed_tokens; R CachedRequestData num_output_tokens=1
R.num_computed_tokens = R.S2.num_computed_tokens + R.S2.num_new_tokens = R.num_prompt_tokens + 1 + s2
R.num_output_placeholders = 1 + 1 + s2; R.spec_token_ids = [-1] * self.num_spec_tokens. 这里模拟 step S2 执行结束之后 R 上发生的状态变化. 假设了 S2 accept 所有 draft tokens, 即生成了 1 + s2 个 output tokens; 同时 假设了 了 step S2 生成了 num_spec_tokens 个 draft tokens.

S1.update_from_output

step S1 执行结束了, 根据 step S1 实际情况矫正 R 上若干状态.

R._output_token_ids = [t1]
R.num_output_placeholders = 1 + s2;
cache_blocks(R, R.num_prompt_tokens)

S3.schedule

R.S3.num_computed_tokens = R.num_prompt_tokens + 1 + s2; R.S3.num_new_tokens = 1 + s3; 这里 假设了 S2 生成 num_spec_tokens 个 draft tokens, 只实际调度了 s3 个 draft tokens
allocate_slots(R, R.S3.num_new_tokens); cache_blocks(R, R.num_prompt_tokens + 1);
R CachedRequestData computed = R.S3.num_computed_tokens; R CachedRequestData num_output_tokens=1 + 1 + s2;
- 这里 R.S3.num_computed_tokens 是 假设了 S2 accept 所有 draft token. worker 侧会根据 S2 实际 accept 情况进行修正.
R.num_computed_tokens = R.S3.num_computed_tokens + R.S3.num_new_tokens = R.num_prompt_tokens + 1 + s2 + 1 + s3
R.num_output_placeholders = 1 + s2 + 1 + s3; R.spec_token_ids = [-1] * self.num_spec_tokens.

S2.update_from_output

step S2 执行结束了, 设这里 step S2 接受了 R.S2.num_accepted 个 draft tokens, 拒绝了剩下的 R.S2.num_rejected 个 draft tokens; R.S2.num_accepted + R.S2.num_rejected = s2.

R.num_computed_tokens -= R.S2.num_rejected = R.num_prompt_tokens + 1 + R.S2.num_accepted + 1 + s3; R.num_output_placeholders -= R.S2.num_rejected = 1 + R.S2.num_accepted + 1 + s3
R._output_token_ids = [t1 d2 d3 t4], d2 d3 是 draft model 生成且被 accept 的 token, 这里设 R.S2.num_accepted = 2. t1, t4 是 target model 生成.
R.num_output_placeholders -= R.S2.num_accepted + 1 = 1 + s3;
cache_blocks(R, R.num_prompt_tokens + 1 + R.S2.num_accepted)

在有了这个执行流程之后, 相信后续再具体深入代码验证时就不至于太茫然迷失了.

3 个 bug

哎, 我个人觉得 vllm async scheduling 实现还是有点复杂了, 尤其是其为了进一步最大化 GPU 利用率完全移除掉了 step 之间的 gpu sync; 导致现在 scheduler 侧发起一次 step 之前, 其看到的状态都是 “假” 的, worker 侧在 execute_model 之前其看到的状态也都是 “假” 的, 需要仔细做好状态的修正工作. 我个人在学习这块代码时就发现了 3 个 bug: bug1, bug2, bug3, 主要发生在上面执行链路部分假设失效而且缺少必需的状态矫正导致.

想了下, 之所以实现这么复杂原因应该还是在 vllm v1 最初立项时的想法, vllm v1 统一了 prefill, chunked prefill, decode 的处理; 我还记得在 vllm v1 刚推出时, 当前 scheduler 代码非常简洁清晰, 就是确定当前 computed token num, 待计算的 new token ids, 之后下发给 model 执行就好了; scheduler 侧完全没有了 prefill/chunk prefill/decode 这些概念.

之后在 async scheduling 引入的第一个版本, 这里 async scheduling 尚不支持 spec decoding, 那么一切还算是清晰, 毕竟只要没有 spec decoding 那么 worker 侧行为就是确定的, 每次只生成 1 个 output token, 每次也只需要接受这个刚刚生成的 output token; 虽然在 async scheduling 开启时, 我们无法拿到上一个 step 生成的 output token, 那也无所谓, 填个 -1 带给 worker, 之后 worker 再根据实际 output token 矫正一下就行.

但是在尝试 async scheduling + spec decoding 之后, worker 侧行为完全不确定了; 无法确定 worker 侧接受了多少个 draft tokens 生成了多少 output tokens, 也无法确定 worker 侧生成了多少 draft tokens; 情况就有点恶化了; 此时 scheduler 被迫在一系列假设基础上构造 CachedRequestData, 这种情况下 CachedRequestData 中除了 req_id 以及 new_block_ids 之外其他信息都是不可信的, worker 必须在执行前进行状态矫正.

class CachedRequestData:
  req_ids: list[str]
  resumed_req_ids: set[str]
  new_token_ids: list[list[int]]
  all_token_ids: dict[str, list[int]]
  new_block_ids: list[tuple[list[int], ...] | None]
  num_computed_tokens: list[int]
  num_output_tokens: list[int]

那么应该怎么改进这种情况呢? 从第一性原理出发, 既然 scheduler 侧无法在 async scheduling 时维护准确的状态, 那么索性就不维护好了. 对于已经处于 decoding 阶段的请求, scheduler 侧只维护 num_output_tokens 这一个状态, 每次 step 时只下发 reqid, new_block_ids 给 worker; 而 worker 侧其状态始终是准确的, 拿来直接用 prepare input, start forward 就行了. 这其实就是早期我们内部的推理框架 BladeLLM 做法=

3 个 bug

CATALOG