vllm 学习笔记

Appearance

使用小模型（draft model）快速生成候选 token，再由大模型（target model）并行验证，在不损失精度的前提下加速推理。

为什么需要 Speculative Decoding ​

自回归解码每步只能生成一个 token，GPU 利用率极低（memory-bound）。Speculative Decoding 利用 draft model 一次生成 K 个候选 token，target model 通过一次 forward 验证所有候选，接受匹配的 token 并拒绝不匹配的。平均接受率较高时，等效每步生成多个 token，显著降低延迟。

核心原理 ​

  • Draft-then-Verify：Draft model 自回归生成 K 个 token，Target model 对 K+1 个位置做一次 forward，得到每个位置的概率分布。
  • 拒绝采样：对每个候选 token，若 Target model 的概率 >= Draft model 的概率则接受；否则按比例概率拒绝并以 Target model 的分布采样替代。
  • 无损保证：数学上可证明输出分布与仅用 Target model 完全一致。
  • 多种 Draft 来源：小型 LLM、Medusa head、EAGLE、n-gram 猜测等都可作为 draft 来源。

在源码中的实现 ​

  • vllm/v1/spec_decode/ — 推测解码的草稿生成（proposer）实现目录：eagle.py、medusa.py、dflash.py、ngram_proposer.py、vocab_mapping.py（异构词表 TLI）等。
  • vllm/v1/spec_decode/draft_model.py — 驱动 draft 模型多步生成候选 token。
  • vllm/v1/worker/gpu/spec_decode/ — Target worker 验证候选 token 的主逻辑（speculator.py、eagle/、autoregressive/、dspark/（DSpark 并行起草）、multi_module_mtp/、adaptive_verification.py 等）。
  • 拒绝采样验证分两套：V1 runner 用 vllm/v1/sample/rejection_sampler.py（nn.Module + Triton 内核）；MRV2 (V2) runner 用 vllm/v1/worker/gpu/spec_decode/rejection_sampler.py + rejection_sampler_utils.py——use_block_verification 与 joint block 验证逻辑只在此 V2 版本中。由于 MRV2 已是全部模型默认（#53183），多数部署命中后者。
  • vllm/config/speculative.py — SpeculativeConfig 定义 draft model、num_speculative_tokens 等参数。

关键配置（v0.26 起） ​

  • rejection_sample_method（三选一）：'standard'（Leviathan et al. 2023，逐 token 概率比 p(x) > u·q(x)）、'synthetic'（按 synthetic_acceptance_rates 校准的衰减接受率）、'block'（block verification，Sun et al. 2024 / arXiv 2403.10444，维护联合前缀比 p_i = Π q(x_j)/p(x_j) 整块判定接受长度，不可与 synthetic 同用；仅 V2 sampler 支持）。
  • use_heterogeneous_vocab（异构词表推测，TLI）：允许 draft/target 词表不同，仅 method='draft_model' + draft_sample_method='greedy' 可用。初始化时归一化两侧 token 串求交集，采样前 constrain_draft_logits 把非交集位置置 -inf，采样后 map_draft_to_target_ids 翻译 ID（见 vocab_mapping.py），从而在词表不同的草稿/目标模型间保持贪心解码无损。
  • kv_cache_dtype：可单独指定草稿模型的 KV 缓存精度，默认（None）继承目标模型的 --kv-cache-dtype。
  • disable_eagle_block_drop（#53388，实验）：禁用 EAGLE 类方法对尾部（含未验证 draft token 的 volatile）prefix-cache 块的丢弃，用于度量尾块复用对接受率的影响。
  • --per-request-spec-decode-metrics（ObservabilityConfig，#48915）：OpenAI 响应内返回每请求接受长度/接受率直方图（summary/detailed），仅 n==1。

相关概念 ​

  • continuous-batching — Speculative decoding 与 continuous batching 结合管理 draft/verify 请求
  • kv-cache — Draft 和 Target 模型各自维护 KV Cache
  • cuda-graph — Draft 模型的小 forward 可用 CUDA Graph 加速
  • kv-cache-offloading — Draft model 的 KV Cache 可卸载以节省显存

反向链接