Chengshu@skadai · 2026.09.30
1,563 字 · 2,217 词 · 约 10 分钟

走进 vLLM(三):speculative decoding 与 disaggregated P/D

一次大模型 forward 最多换回 k+1 个 token,而且分布和逐个采样完全等价——speculative decoding 是怎么做到的?vLLM V1 为什么不用小模型做 draft,而用 n-gram / EAGLE / Medusa?以及为什么 prefill 和 decode 值得被拆到两批实例上跑。

本篇属于系列 走进 vLLM:高吞吐 LLM 推理系统解剖 · 第 3 篇

原文 · English中文译文

原文:Inside vLLM: Anatomy of a High-Throughput LLM Inference System,作者 Aleksa Gordic,2025-09-05 发布于 vLLM 官方博客。

本文是系列《走进 vLLM:高吞吐 LLM 推理系统解剖》的第 3 篇,共 5 篇。承接第 2 篇(chunked prefill、prefix caching 与 guided decoding)。

本站中英对照排版:左栏英文原文,右栏中文译文;不好翻译的术语保留英文写法。

Speculative Decoding

Speculative Decoding

In autoregressive generation, each new token requires a forward pass of the large LM. This is expensive — every step reloads and applies all model weights just to compute a single token! (assuming batch size == 1, in general it’s B)

在自回归生成里,每产出一个新 token 都需要大模型做一次 forward pass。这很贵——每一步都要把全部模型权重重新读一遍、算一遍,只为了算出一个 token!(这里假设 batch size == 1,一般来说是 B。)

Speculative decoding [8] speeds this up by introducing a smaller draft LM. The draft proposes k tokens cheaply. But we don’t ultimately want to sample from the smaller model — it’s only there to guess candidate continuations. The large model still decides what’s valid.

Speculative decoding [8] 用一个更小的 draft LM 来加速:draft 模型廉价地提出 k 个 token。但我们最终并不想从这个小模型采样——它只是用来猜候选续写的。哪些是有效的,仍然由大模型说了算。

Here are the steps:

步骤如下:

  1. Draft: run the small model on the current context and propose k tokens
  2. Verify: run the large model once on context + k draft tokens. This produces probabilities for those k positions plus one extra (so we get k+1 candidates)
  3. Accept/reject: going from left to right over the k draft tokens:
  1. Draft:在当前 context 上跑小模型,提出 k 个 token
  2. Verify:把 context + k 个 draft token 一起喂给大模型跑一次。这会给出这 k 个位置的(以及额外一个位置的)概率(于是我们拿到了 k+1 个候选)
  3. Accept/reject:对这 k 个 draft token 从左到右逐个判断:
  • If the large model's probability for the draft token ≥ the draft's probability, accept it
  • Otherwise, accept it with probability p_large(token)/p_draft(token)
  • Stop at the first rejection, or accept all k draft tokens
    • If all k draft tokens are accepted, also sample the extra (k+1)-th token "for free" from the large model (we already computed that distribution)
    • If there was a rejection create a new rebalanced distribution at that position (p_large - p_draft, clamp min at 0, normalize to sum to 1) and sample the last token from it
  • 如果大模型给出该 draft token 的概率 ≥ draft 模型的概率,接受它
  • 否则以 p_large(token)/p_draft(token) 的概率接受它
  • 遇到第一次拒绝就停下,或者接受全部 k 个 draft token
    • 如果 k 个 draft token 全被接受,就顺便从大模型里“免费”采样出多出来的第 (k+1) 个 token(那个分布我们已经算好了)
    • 如果中途被拒,就在那个位置上构造一个新的、重新平衡过的分布(p_large - p_draft,下界截到 0,再归一化到和为 1),从它里面采样出最后一个 token

Why this works: Although we use the small model to propose candidates, the accept/reject rule guarantees that in expectation the sequence is distributed exactly as if we had sampled token by token from the large model. This means speculative decoding is statistically equivalent to standard autoregressive decoding — but potentially much faster, since a single large-model pass can yield up to k+1 tokens.

为什么这是对的:虽然我们是用小模型来提候选,但这条接受/拒绝规则保证了:期望上,序列的分布和逐个从大模型采样出来的分布完全一致。也就是说,speculative decoding 在统计上等价于标准的自回归解码——但可能快得多,因为一次大模型 forward 最多能换来 k+1 个 token。

说明

I recommend looking at gpt-fast for a simple implementation, and the original paper for the math details and the proof of equivalence to sampling from the full model.

说明

想了解简单实现,推荐看 gpt-fast;想看数学细节以及“与从完整模型采样等价”的证明,读原始论文。

vLLM V1 does not support the LLM draft model method, instead it implements faster—but less accurate—proposal schemes: n-gram, EAGLE [9], and Medusa [10].

vLLM V1 不支持用 LLM 做 draft model 的方案,而是实现了更快、但精度略低的候选方案:n-gram、EAGLE [9] 和 Medusa [10]。

One-liners on each:

各用一句话说清:

  • n-gram: take the last prompt_lookup_max tokens; find a prior match in the sequence; if found, propose the k tokens that followed that match; otherwise decrement the window and retry down to prompt_lookup_min
  • n-gram:取最后 prompt_lookup_max 个 token,在序列里找一次历史匹配;找到了就提议匹配之后跟随的 k 个 token;没找到就缩小窗口重试,一直降到 prompt_lookup_min

说明

The current implementation returns k tokens after the first match. It feels more natural to introduce a recency bias and reverse the search direction? (i.e. last match)

说明

当前实现返回第一次匹配之后的 k 个 token。引入一点 recency bias、把搜索方向反过来(也就是取最后一次匹配)是不是更自然?

  • Eagle: perform “model surgery” on the large LM—keep embeddings and LM head, replace the transformer stack with a lightweight MLP; fine-tune that as a cheap draft
  • Eagle:对大 LM 做“模型手术”——保留 embedding 和 LM head,把 transformer 堆栈换成一个轻量 MLP,再微调它作为廉价的 draft
  • Medusa: train auxiliary linear heads on top (embeddings before LM head) of the large model to predict the next k tokens in parallel; use these heads to propose tokens more efficiently than running a separate small LM
  • Medusa:在大模型之上(LM head 之前的 embedding 处)训练若干辅助线性头,并行预测接下来 k 个 token;用这些头来提 token,比单独跑一个小 LM 更高效

Here’s how to invoke speculative decoding in vLLM using ngram as the draft method:

这是用 ngram 作为 draft method 在 vLLM 里调用 speculative decoding 的方式:

from vllm import LLM, SamplingParams

prompts = [
    "Hello, my name is",
    "The president of the United States is",
]

sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

speculative_config={
    "method": "ngram",
    "prompt_lookup_max": 5,
    "prompt_lookup_min": 3,
    "num_speculative_tokens": 3,
}

def main():
    llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", speculative_config=speculative_config)

    outputs = llm.generate(prompts, sampling_params)

if __name__ == "__main__":
    main()

How does this work in vLLM?

在 vLLM 里这是怎么实现的?

Setup (during engine construction):

Setup(engine 构造期间):

  1. Init device: create a drafter (draft model, e.g., NgramProposer) and a rejection_sampler (parts of it are written in Triton).
  2. Load model: load draft model weights (no-op for n-gram).
  1. Init device:创建一个 drafter(draft 模型,比如 NgramProposer)和一个 rejection_sampler(它的一部分是用 Triton 写的)。
  2. Load model:加载 draft 模型权重(对 n-gram 来说是空操作)。

After that in the generate function (assume we get a brand new request):

之后在 generate 函数里(假设来的是一个全新请求):

  1. Run the regular prefill step with the large model.
  2. After the forward pass and standard sampling, call propose_draft_token_ids(k) to sample k draft tokens from the draft model.
  3. Store these in request.spec_token_ids (update the request metadata).
  4. On the next engine step, when the request is in the running queue, add len(request.spec_token_ids) to the “new tokens” count so allocate_slots reserves sufficient KV blocks for the fwd pass.
  5. Copy spec_token_ids into input_batch.token_ids_cpu to form (context + draft) tokens.
  6. Compute metadata via _calc_spec_decode_metadata (this copies over tokens from input_batch.token_ids_cpu, prepares logits, etc.), then run a large-model forward pass over the draft tokens.
  7. Instead of regular sampling from logits, use the rejection_sampler to accept/reject left-to-right and produce output_token_ids.
  8. Repeat steps 2-7 until a stop condition is met.
  1. 用大模型跑一次常规的 prefill step。
  2. forward pass 和标准采样之后,调用 propose_draft_token_ids(k),从 draft 模型采样出 k 个 draft token。
  3. 把它们存进 request.spec_token_ids(更新请求 metadata)。
  4. 下一个 engine step,当这个请求在 running 队列里时,把 len(request.spec_token_ids) 加到“新 token”计数上,这样 allocate_slots 就会为 forward pass 预留足够的 KV block。
  5. 把 spec_token_ids 拷进 input_batch.token_ids_cpu,拼成(context + draft)token。
  6. 通过 _calc_spec_decode_metadata 计算 metadata(它会从 input_batch.token_ids_cpu 拷贝 token、准备 logits 等),然后对这批 draft token 跑一次大模型的 forward pass。
  7. 不再从 logits 做常规采样,而是用 rejection_sampler 从左到右接受/拒绝,产出 output_token_ids。
  8. 重复步骤 2–7,直到满足某个停止条件。

The best way to internalize this is to fire up your debugger and step through the codebase, but this section hopefully gives you a taste for it. This as well:

把这一切内化的最好方式是打开 debugger、在代码库里一步步走。但希望这一节已经让你尝到味道了。下面这张图也能帮忙:

Speculative decoding:draft 与 verify 的示意

图 11:Speculative decoding
图 11:Speculative decoding

Figure 11: Speculative decoding

Disaggregated P/D

Disaggregated P/D

I’ve already previously hinted at the motivation behind disaggregated P/D (prefill/decode).

前面我已经暗示过 disaggregated P/D(prefill/decode 分离)的动机。

Prefill and decode have very different performance profiles (compute-bound vs. memory-bandwidth-bound), so separating their execution is a sensible design. It gives tighter control over latency — both TFTT (time-to-first-token) and ITL (inter-token latency) — more on this in the benchmarking section.

Prefill 和 decode 的性能画像差别极大(compute-bound 对 memory-bandwidth-bound),所以把它们的执行分开是个很自然的设计。这样能更精细地控制延迟——包括 TTFT(time-to-first-token)和 ITL(inter-token latency)——benchmark 那一节会展开讲。

In practice, we run N vLLM prefill instances and M vLLM decode instances, autoscaling them based on the live request mix. Prefill workers write KV to a dedicated KV-cache service; decode workers read from it. This isolates long, bursty prefill from steady, latency-sensitive decode.

实践中,我们会跑 N 个 vLLM prefill 实例和 M 个 vLLM decode 实例,并根据实时请求构成自动伸缩。prefill worker 把 KV 写到一个专用的 KV-cache 服务,decode worker 从那里读。这样就把又长又突发的 prefill,与稳定、对延迟敏感的 decode 隔离开来。

How does this work in vLLM?

在 vLLM 里这是怎么实现的?

For clarity, the example below relies on SharedStorageConnector, a debugging connector implementation used to illustrate the mechanics.

为了讲清楚,下面的例子用了 SharedStorageConnector——一个用来演示机制、偏调试用途的 connector 实现。

说明

Connector is vLLM’s abstraction for handling the exchange of KVs between instances. Connector interface is not yet stable, there are some near-term improvements planned which will involve changes, some potentially breaking.

说明

Connector 是 vLLM 用来处理实例之间 KV 交换的抽象。Connector 接口目前还不稳定,近期有一些改进计划,会带来变更,其中一部分可能是不兼容变更。

We launch 2 vLLM instances (GPU 0 for prefill and GPU 1 for decode), and then transfer the KV cache between them:

我们启动两个 vLLM 实例(GPU 0 做 prefill,GPU 1 做 decode),然后在它们之间搬运 KV cache:

import os
import time
from multiprocessing import Event, Process
import multiprocessing as mp

from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig

prompts = [
    "Hello, my name is",
    "The president of the United States is",
]

def run_prefill(prefill_done):
  os.environ["CUDA_VISIBLE_DEVICES"] = "0"

  sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)

  ktc=KVTransferConfig(
      kv_connector="SharedStorageConnector",
      kv_role="kv_both",
      kv_connector_extra_config={"shared_storage_path": "local_storage"},
  )

  llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", kv_transfer_config=ktc)
  llm.generate(prompts, sampling_params)

  prefill_done.set()  # notify decode instance that KV cache is ready

  # To keep the prefill node running in case the decode node is not done;
  # otherwise, the script might exit prematurely, causing incomplete decoding.
  try:
      while True:
          time.sleep(1)
  except KeyboardInterrupt:
      print("Script stopped by user.")

def run_decode(prefill_done):
  os.environ["CUDA_VISIBLE_DEVICES"] = "1"

  sampling_params = SamplingParams(temperature=0, top_p=0.95)

  ktc=KVTransferConfig(
      kv_connector="SharedStorageConnector",
      kv_role="kv_both",
      kv_connector_extra_config={"shared_storage_path": "local_storage"},
  )

  llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", kv_transfer_config=ktc)

  prefill_done.wait()  # block waiting for KV cache from prefill instance

  # Internally it'll first fetch KV cache before starting the decoding loop
  outputs = llm.generate(prompts, sampling_params)

if __name__ == "__main__":
  prefill_done = Event()
  prefill_process = Process(target=run_prefill, args=(prefill_done,))
  decode_process = Process(target=run_decode, args=(prefill_done,))

  prefill_process.start()
  decode_process.start()

  decode_process.join()
  prefill_process.terminate()

说明

I’ve also experimented with LMCache [11], the fastest production-ready connector (uses NVIDIA’s NIXL as the backend), but it’s still at the bleeding edge and I ran into some bugs. Since much of its complexity lives in an external repo, SharedStorageConnector is a better choice for explanation.

说明

我也试过 LMCache [11]——目前最快的、生产可用的 connector(后端用 NVIDIA 的 NIXL),但它还处在非常前沿的位置,我踩到了一些 bug。考虑到它的大部分复杂度都在外部仓库里,用 SharedStorageConnector 来讲原理更合适。

These are the steps in vLLM:

vLLM 里的步骤如下:

  1. Instantiation — During engine construction, connectors are created in two places:
  1. Instantiation —— engine 构造期间,connector 会在两个地方被创建:
  • Inside the worker’s init device procedure (under init worker distributed environment function), with role “worker”.
  • Inside the scheduler constructor, with role “scheduler”.
  • worker 的 init device 流程内部(在 init worker distributed environment 函数里),角色是 “worker”。
  • scheduler 构造函数内部,角色是 “scheduler”。
  1. Cache lookup — When the scheduler processes prefill requests from the waiting queue (after local prefix-cache checks), it calls connector’s get_num_new_matched_tokens. This checks for externally cached tokens in the KV-cache server. Prefill always sees 0 here; decode may have a cache hit. The result is added to the local count before calling allocate_slots.
  2. State update — The scheduler then calls connector.update_state_after_alloc, which records requests that had a cache (no-op for prefill).
  3. Build metadata object — At the end of scheduling, the scheduler calls meta = connector.build_connector_meta:
  1. Cache lookup —— scheduler 处理 waiting 队列里的 prefill 请求时(在本地 prefix-cache 检查之后),会调用 connector 的 get_num_new_matched_tokens,去 KV-cache 服务里查外部缓存了多少 token。prefill 这里永远是 0;decode 可能命中。结果会在调用 allocate_slots 之前加到本地计数上。
  2. State update —— 接着 scheduler 调用 connector.update_state_after_alloc,记录下命中了缓存的那些请求(对 prefill 是空操作)。
  3. Build metadata object —— 调度结束时,scheduler 调用 meta = connector.build_connector_meta:
  • Prefill adds all requests with is_store=True (to upload KV).
  • Decode adds requests with is_store=False (to fetch KV).
  • prefill 把 is_store=True 的请求加进去(要上传 KV)。
  • decode 把 is_store=False 的请求加进去(要拉取 KV)。
  1. Context manager — Before the forward pass, the engine enters a KV-connector context manager:
  1. Context manager —— forward pass 之前,engine 进入一个 KV-connector context manager:
  • On enter: kv_connector.start_load_kv is called. For decode, this loads KV from the external server and injects it into paged memory. For prefill, it’s a no-op.
  • On exit: kv_connector.wait_for_save is called. For prefill, this blocks until KV is uploaded to the external server. For decode, it’s a no-op.
  • 进入时:调用 kv_connector.start_load_kv。对 decode,这会把 KV 从外部服务加载进来并注入 paged memory;对 prefill 是空操作。
  • 退出时:调用 kv_connector.wait_for_save。对 prefill,这会阻塞到 KV 上传到外部服务为止;对 decode 是空操作。

Here is a visual example:

一个可视化的例子:

图 12:disaggregated P/D
图 12:disaggregated P/D

Figure 12: disaggregated P/D

Additional notes:

  • For SharedStorageConnector “external server” is just a local file system.
  • Depending on configuration, KV transfers can also be done layer-by-layer (before/after each attention layer).
  • Decode loads external KV only once, on the first step of its requests; afterwards it computes/stores locally.

补充说明

  • 对 SharedStorageConnector 来说,“外部服务”就是本地文件系统。
  • 视配置而定,KV 传输也可以逐层进行(在每个 attention 层之前/之后)。
  • decode 只在它那些请求的第一步加载一次外部 KV;之后就在本地计算/存储了。

注释

  1. “Accelerating Large Language Model Decoding with Speculative Sampling” —— https://arxiv.org/abs/2302.01318
  2. “EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty” —— https://arxiv.org/abs/2401.15077
  3. “Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads” —— https://arxiv.org/abs/2401.10774
  4. LMCache —— https://github.com/LMCache/LMCache

讨论

这里是静态站点,没有内嵌评论区。如果这篇文章对你有用,欢迎通过 RSS 订阅后续更新。