课程 11 · 进阶机制
投机解码、PD 分离、结构化输出、异步调度、睡眠唤醒——这些看起来互不相干的功能,在 vLLM 里全部落在同一个 SchedulerOutput 上。这一课讲它们怎么被「塞进」那套 token 预算模型。
投机解码的做法是:先用一个小模型(或 n-gram、EAGLE head)猜出接下来 5 个 token,再用大模型一次前向验证这 5 个,能接受几个就接受几个。关键在于:这 5 个 draft token 也要占 KV 块。
scheduler.py · L727–729spec = self.vllm_config.speculative_config
draft_slots = spec.max_num_new_slots_for_drafting if spec is not None else 0
input_budget = self.scheduler_config.max_num_batched_tokens
...
num_new_tokens = min(num_new_tokens, token_budget, input_budget - draft_slots)
注意 input_budget - draft_slots:正常请求的输入预算要先扣掉留给 draft 的槽位。否则一个长 prompt 会把预算吃光,投机就没地方放了。
v1/core/sched/output.py · L261–264scheduled_spec_decode_tokens: dict[str, list[int]]
# req_id -> spec_token_ids(没有 draft 的请求不会出现在字典里)
num_scheduled_tokens是这条请求本轮主模型输入总数,已经包含本轮draft验证位置;scheduled_spec_decode_tokens标识其中哪些为draft。例如上轮最后token加3个draft,总数就是4。返回后按实际接受情况修正computed账本。拒绝采样的理想分布等价也不同于跨batch逐位确定性;详见27课算例。
因为 vLLM 的调度单位是「token」而不是「阶段」。投机解码只是让某条请求的 num_tokens_with_spec 比它的实际 token 多几个——第 08 课那句注释说的「general enough」就是这个意思。
Prefill 是算力密集,Decode 是带宽密集。把两者拆到不同机器上,各自用最合适的硬件配比,这就是 Prefill-Decode Disaggregation。
代价是:decode 节点需要的 KV,必须由 prefill 节点传过来。vLLM 用 kv_transfer 模块 + KV Connector 抽象来做这件事[v1/distributed/kv_transfer]。
kv_cache_manager.py · L379–395num_external_computed_tokens=0, # 从 connector 拿到的 KV 对应的 token 数
delay_cache_blocks=False, # 传输还没完成,先占块但不登记前缀缓存
delay_cache_blocks 是最能说明问题的一个参数:块已经分配了,但里面的数据还在网络上飞。这时候如果把它登记进前缀缓存,别的请求命中就会读到垃圾——所以必须延迟登记,等传输完成再标记。
调度器还要处理「传输失败」「对端超时」,这就是为什么 kv_connector_meta 会出现在 SchedulerOutput 里、并且每个 step 都可能带 has_sync_kv_loads 这样的标志。
v1/engine/core.py · L720–729scheduler_output = self.scheduler.schedule(...)
future = self.model_executor.execute_model(scheduler_output, non_block=True)
grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output) # 与 GPU 并行
...
model_output = self._wait_for_output(future)
if model_output is None:
model_output = self.model_executor.sample_tokens(grammar_output)
位置很讲究:bitmask 计算被放在「等 GPU」的窗口里做。语法解析(正则 / JSON Schema / EBNF)是纯 CPU 的字符串状态机,正好可以蹭 GPU 的延迟。
这也是第 08 课说「把 sample_tokens 拆出来」的原因之一:只有在采样这一步才需要 bitmask,采样又必须在拿到 GPU 输出之后。拆开之后,两者可以并行。
普通流程是一问一答:调度 → 跑模型 → 等结果 → 再调度。异步调度把「等结果」这段时间利用起来:在 GPU 还在跑这一步的时候,先把下一步的调度做完。
难点在于状态一致性:下一步调度时,上一步的采样结果还没回来,调度器并不知道某条请求会不会结束。vLLM 的解法是引入 num_output_placeholders——先假设这些 token 会产生,等真实结果回来再校正。抢占时则要标记 in-flight 输出为 stale。
收益很直接:decode 一步 20 ms 时,CPU 侧的调度、tokenize、输出处理可以完全藏进这 20 ms 里。
v1/engine/llm_engine.py · L426–449def sleep(self, level: int = 1, mode: PauseMode = "abort"): ...
def release_kv_cache_memory(self): ...
def wake_up(self, tags: list[str] | None = None) -> bool: ...
这是为强化学习训练场景设计的:训练需要显卡时,推理引擎把 KV cache 和权重释放掉,把显存让出来;训练完再唤醒并重新加载。没有这一层,训练和推理只能分时复用两套机器。
| 特性 | 落在调度的哪里 | 落在执行的哪里 |
|---|---|---|
| 投机解码 | draft_slots、scheduled_spec_decode_tokens | 验证前向 + 拒绝采样 |
| 前缀缓存 | get_computed_blocks | 命中块直接读 KV |
| PD 分离 | num_external_computed_tokens、delay_cache_blocks | KV connector 传输 |
| 结构化输出 | get_grammar_bitmask | sample_tokens(grammar_output) |
| 异步调度 | num_output_placeholders、stale 标记 | future + batch queue |
| 睡眠唤醒 | 暂停状态 PauseState | 释放 / 重建 KV 与权重 |
假设你要给 nano-vllm 加一个「最简单的投机解码」:用一个 4-gram 查表猜 3 个 token,然后一次前向验证。列出你需要改的三个地方,以及每处的关键风险。
① Sequence:增加 spec_token_ids,让 num_tokens 与实际待算 token 数分离(对应 vLLM 的 num_tokens_with_spec)。风险:num_blocks、last_block_num_tokens 全部要按「含 draft」计算,否则块表错位。
② prepare_decode:把 input_ids 从 1 个变成 1+k 个,slot_mapping 相应扩展。风险:position 与 context_lens 也要跟着变,否则 RoPE 错位。
③ postprocess:接受/拒绝逻辑——逐位置比较 draft 与实际采样结果,截断到第一个拒绝点,把 KV 块按接受长度回收。风险:被拒绝的 token 在 KV cache 里留下的垃圾必须能覆盖写,否则下次会读到。
「投机解码为什么不会改变输出分布?」
「因为验证阶段用的是拒绝采样:对每个 draft token,以 min(1, p_target / p_draft) 的概率接受;一旦拒绝,就从修正后的分布 norm(max(0, p_target - p_draft)) 里重新采一个。数学上可以证明最终样本的分布严格等于从目标模型采样。所以它只影响速度,不影响质量。」
可以问:「PD 分离什么时候值得做?」——需要估算带宽、集群规模、TTFT 目标,我们可以拿你未来的实验数据一起算。