Speculative decoding in vLLM

Warning

Please note that speculative decoding in vLLM is not yet optimized and does not usually yield inter-token latency reductions for all prompt datasets or sampling parameters. The work to optimize it is ongoing and can be followed in this issue.

This document shows how to use Speculative Decoding with vLLM. Speculative decoding is a technique which improves inter-token latency in memory-bound LLM inference.

Speculating with a draft model

The following code configures vLLM in an offline mode to use speculative decoding with a draft model, speculating 5 tokens at a time.

  1. from vllm import LLM, SamplingParams
  3. prompts = [
  4.     "The future of AI is",
  5. ]
  6. sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
  8. llm = LLM(
  9.     model="facebook/opt-6.7b",
  10.     tensor_parallel_size=1,
  11.     speculative_model="facebook/opt-125m",
  12.     num_speculative_tokens=5,
  13.     use_v2_block_manager=True,
  14. )
  15. outputs = llm.generate(prompts, sampling_params)
  17. for output in outputs:
  18.     prompt = output.prompt
  19.     generated_text = output.outputs[0].text
  20.     print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

To perform the same with an online mode launch the server:

  1.    python -m vllm.entrypoints.openai.api_server --host 0.0.0.0 --port 8000 --model facebook/opt-6.7b \
  2.    --seed 42 -tp 1 --speculative_model facebook/opt-125m --use-v2-block-manager \
  3.    --num_speculative_tokens 5 --gpu_memory_utilization 0.8
  5. Then use a client:
  1. from openai import OpenAI
  3. # Modify OpenAI's API key and API base to use vLLM's API server.
  4. openai_api_key = "EMPTY"
  5. openai_api_base = "http://localhost:8000/v1"
  7. client = OpenAI(
  8.     # defaults to os.environ.get("OPENAI_API_KEY")
  9.     api_key=openai_api_key,
  10.     base_url=openai_api_base,
  11. )
  13. models = client.models.list()
  14. model = models.data[0].id
  16. # Completion API
  17. stream = False
  18. completion = client.completions.create(
  19.     model=model,
  20.     prompt="The future of AI is",
  21.     echo=False,
  22.     n=1,
  23.     stream=stream,
  24. )
  26. print("Completion results:")
  27. if stream:
  28.     for c in completion:
  29.         print(c)
  30. else:
  31.     print(completion)

Speculating by matching n-grams in the prompt

The following code configures vLLM to use speculative decoding where proposals are generated by matching n-grams in the prompt. For more information read this thread.

  1. from vllm import LLM, SamplingParams
  3. prompts = [
  4.     "The future of AI is",
  5. ]
  6. sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
  8. llm = LLM(
  9.     model="facebook/opt-6.7b",
  10.     tensor_parallel_size=1,
  11.     speculative_model="[ngram]",
  12.     num_speculative_tokens=5,
  13.     ngram_prompt_lookup_max=4,
  14.     use_v2_block_manager=True,
  15. )
  16. outputs = llm.generate(prompts, sampling_params)
  18. for output in outputs:
  19.     prompt = output.prompt
  20.     generated_text = output.outputs[0].text
  21.     print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

Speculating using MLP speculators

The following code configures vLLM to use speculative decoding where proposals are generated by draft models that conditioning draft predictions on both context vectors and sampled tokens. For more information see this blog or this technical report.

  1. from vllm import LLM, SamplingParams
  3. prompts = [
  4.     "The future of AI is",
  5. ]
  6. sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
  8. llm = LLM(
  9.     model="meta-llama/Meta-Llama-3.1-70B-Instruct",
  10.     tensor_parallel_size=4,
  11.     speculative_model="ibm-fms/llama3-70b-accelerator",
  12.     speculative_draft_tensor_parallel_size=1,
  13.     use_v2_block_manager=True,
  14. )
  15. outputs = llm.generate(prompts, sampling_params)
  17. for output in outputs:
  18.     prompt = output.prompt
  19.     generated_text = output.outputs[0].text
  20.     print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

Note that these speculative models currently need to be run without tensor parallelism, although it is possible to run the main model using tensor parallelism (see example above). Since the speculative models are relatively small, we still see significant speedups. However, this limitation will be fixed in a future release.

A variety of speculative models of this type are available on HF hub:

Resources for vLLM contributors