Speculative decoding

Speculative decoding#

Author: Chang Liu
Knowledge level: Intermediate

This tutorial demonstrates how to enable speculative decoding in LLM serving and measure its effect. It uses vLLM, one of the most commonly used open-source LLM frameworks, as the serving engine. The tutorial conducts all the benchmark and analysis on AMD Instinct™ MI300X GPUs and the AMD ROCm software stack. For a basic understanding of speculative decoding, including usage guidelines, see the vLLM Speculative Decoding blog.

In this tutorial, you’ll use Llama-3.1 70B as the base model and Llama-3.2 1B as the draft model, and serve the same workload with and without speculative decoding enabled. The tutorial provides detailed steps to measure the difference on AMD MI300X GPUs.

Prerequisites#

This tutorial was developed and tested using the following setup.

Operating system#

  • Ubuntu 22.04: Ensure your system is running Ubuntu version 22.04.

Hardware#

  • AMD Instinct™ GPU(s): This tutorial has been tested on AMD Instinct MI300X GPUs. Ensure you are using an AMD Instinct GPU or compatible hardware with ROCm support and that your system meets the official requirements.

Software#

  • ROCm 6.2+: Install and verify ROCm by following the ROCm install guide. After installation, confirm your setup using:

    amd-smi

    This command lists your AMD GPUs with relevant details.

    Note: For ROCm 6.4 and earlier, use the rocm-smi command instead.

  • Docker: Ensure Docker is installed and configured correctly. Follow the Docker installation guide for your operating system.

    Note: Ensure the Docker permissions are correctly configured. To configure permissions to allow non-root access, run the following commands:

    sudo usermod -aG docker $USER
newgrp docker

    Verify Docker is working correctly:

    docker run hello-world

Hugging Face API access#

  • Obtain an API token from Hugging Face for downloading models.

  • Ensure the Hugging Face API token has the necessary permissions and approval to access the Meta Llama checkpoints.

Data preparation#

This tutorial uses a sample dataset from Hugging Face, which is prepared during the setup steps.

Preparing the serving environment#

1. Pull the Docker image#

Ensure your system meets the system requirements.

Pull the Docker image required for this tutorial:

docker pull vllm/vllm-openai-rocm:v0.31.0

2. Launch the Docker container#

The vLLM servers and the benchmark client all run inside this container. Replace /path/to/model/weights with a directory with room for the model weights (about 75 GB):

docker run -it --rm \
  --network=host \
  --ipc=host \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  --cap-add=SYS_PTRACE \
  --security-opt seccomp=unconfined \
  --shm-size 16G \
  -v /path/to/model/weights:/models \
  -v $(pwd):/workspace/notebooks \
  -w /workspace/notebooks \
  --entrypoint /bin/bash \
  vllm/vllm-openai-rocm:v0.31.0

3. Install and launch Jupyter#

Inside the Docker container, install Jupyter using the following command:

pip install jupyter

Start the Jupyter server:

jupyter-lab --ip=0.0.0.0 --port=8888 --no-browser --allow-root

Note: Ensure port 8888 is not already in use on your system before running the above command. If it is, you can specify a different port by replacing --port=8888 with another port number, for example, --port=8890.

4. Provide your Hugging Face token#

You’ll require a Hugging Face API token to access Llama-3.1. Generate your token at Hugging Face Tokens and request access for Llama-3.1 8B. Tokens typically start with “hf_”.

There are two ways to set up the Hugging Face token: by running a Python script or by running a bash command.

By running a Python script, you can run the following interactive block in your Jupyter notebook to set up the token:

Note: Uncheck the “Add token as Git credential” option.

from huggingface_hub import notebook_login, HfApi

# Prompt the user to log in
notebook_login()

Verify that your token was accepted correctly. You should see output that begins with “Token validated successfully”.

from huggingface_hub import HfApi

try:
    api = HfApi()
    user_info = api.whoami()
    print(f"Token validated successfully! Logged in as: {user_info['name']}")
except Exception as e:
    print(f"Token validation failed. Error: {e}")

Downloading the model weights#

Download the model weights to /models using the commands below. This tutorial uses the AMD-optimized Llama-3.1 70B as the base model and Llama-3.2 1B as the speculative draft model. The servers below load the models from these local paths, so run this step first.

%%bash
set -e
# Download Llama-3.1 70B Instruct FP8 KV model weights in local
hf download \
  amd/Llama-3.1-70B-Instruct-FP8-KV \
  --local-dir /models/amd--Llama-3.1-70B-Instruct-FP8-KV

# Download Llama-3.1 1B Instruct FP8 KV model weights in local
hf download \
  amd/Llama-3.2-1B-Instruct-FP8-KV \
  --local-dir /models/amd--Llama-3.2-1B-Instruct-FP8-KV

Starting the serving#

You can use the online serving mode of vLLM to compare the performance difference between disabling and enabling speculative decoding. Each benchmark below sends NUM_PROMPTS requests at one request per second, so it takes a little over that many seconds:

NUM_PROMPTS = 500  # requests per benchmark

These helpers start a server in the background, wait until it answers (or stop with its log if it fails), and stop it again:

Hide code cell source

 
import os
import signal
import subprocess
import time
import urllib.request


def server_is_up(url, headers=None):
    try:
        with urllib.request.urlopen(urllib.request.Request(url, headers=headers or {}),
                                    timeout=5):
            return True
    except Exception:
        return False


def ensure_server(command, url, log_path, timeout=1200, headers=None):
    """Start `command` in the background (its output goes to log_path) unless
    url already answers, then wait until it does. Returns the server process,
    or None if a server was already running. Raises, with the end of the log,
    if the server exits or is not ready within timeout seconds."""
    if server_is_up(url, headers):
        print(f"Already running: {url}")
        return None
    log = open(log_path, "w")
    server = subprocess.Popen(command, shell=True, stdout=log, stderr=subprocess.STDOUT,
                              start_new_session=True)
    deadline = time.time() + timeout
    while time.time() < deadline and server.poll() is None:
        if server_is_up(url, headers):
            print(f"Ready after {timeout - int(deadline - time.time())}s: {url}")
            return server
        time.sleep(5)
    tail = "".join(open(log_path).readlines()[-40:])
    state = f"exited with code {server.returncode}" if server.poll() is not None else "timed out"
    raise RuntimeError(f"Server {state} before {url} answered. Last log lines:\n{tail}")


def stop_server(server):
    """Stop a server started by ensure_server, and everything it started."""
    if server is not None and server.poll() is None:
        os.killpg(server.pid, signal.SIGTERM)
        try:
            server.wait(timeout=60)
        except subprocess.TimeoutExpired:
            os.killpg(server.pid, signal.SIGKILL)

Start the server without speculative decoding. Modify the port if port 8001 is in use. The settings are for vLLM 0.31. Two of them select AMD’s kernels:

  • VLLM_ROCM_USE_AITER=1 turns on AMD’s AITER kernels, which vLLM 0.31 leaves off by default.

  • --attention-backend ROCM_AITER_UNIFIED_ATTN selects AITER’s unified attention kernel, which handles long prompts and the several tokens per request that speculative decoding verifies at each step.

SERVE holds the command both servers use, so the only difference between the two runs is speculative decoding:

BASE_MODEL = "/models/amd--Llama-3.1-70B-Instruct-FP8-KV"
DRAFT_MODEL = "/models/amd--Llama-3.2-1B-Instruct-FP8-KV"

ENV = ("PYTORCH_TUNABLEOP_ENABLED=0 PYTORCH_TUNABLEOP_TUNING=0"
       " PYTORCH_TUNABLEOP_MAX_TUNING_DURATION_MS=100 PYTORCH_TUNABLEOP_MAX_WARMUP_DURATION_MS=10"
       " PYTORCH_TUNABLEOP_ROTATING_BUFFER_SIZE=1024 PYTORCH_TUNABLEOP_FILENAME=afo_tune_device_%d_full.csv"
       " HIP_FORCE_DEV_KERNARG=1 TOKENIZERS_PARALLELISM=false RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1"
       " NCCL_MIN_NCHANNELS=112 VLLM_ROCM_USE_AITER=1")

SERVE = (ENV + f" vllm serve {BASE_MODEL}"
         " --tensor-parallel-size 1 --distributed-executor-backend mp"
         " --dtype float16 --quantization fp8 --kv-cache-dtype fp8 --no-enable-chunked-prefill"
         " --attention-backend ROCM_AITER_UNIFIED_ATTN"
         " --max-num-seqs 300 --port 8001")

vllm_server = ensure_server(SERVE, "http://localhost:8001/v1/models", "vllm_base.log", timeout=3600)

The server loads the model weights from /models, the directory mounted into the container. Loading the 70B model takes several minutes; the cell above waits until the server answers. If it does not start, the error and the end of vllm_base.log are shown.

Check that the server lists the model:

!curl http://localhost:8001/v1/models

Setting the client#

After starting the server, send it requests and measure the serving performance. The client is vllm bench serve, the benchmark client that comes with vLLM, so it runs in the same container as the server. It is set up to send this workload:

  • random prompts with input lengths spread evenly between 1 and 8192 tokens, asking for between 1 and 256 output tokens (--random-input-len 4096 --random-output-len 128 --random-range-ratio 0.9997, since vllm bench serve spreads lengths around the value you give);

  • greedy decoding that always produces the requested number of tokens (--temperature 0 --ignore-eos);

  • one request per second, reporting the end-to-end latency of each request (--percentile-metrics ttft,tpot,itl,e2el) along with the time to the first token and per output token.

Before the measurement, a short burst of requests warms the server up. A server compiles some of its GPU kernels the first time it sees a batch of a given size, and a request that waits for that compilation would be counted as slow; the burst reaches the batch sizes the measurement does, so those compilations happen first. Both servers get the same warm-up.

The results are saved to vllm_base.json:

!vllm bench serve --backend vllm --model {BASE_MODEL} --port 8001 \
    --dataset-name random --random-input-len 4096 --random-output-len 64 \
    --temperature 0 --ignore-eos --num-warmups 0 --num-prompts 48 --request-rate 4 > /dev/null
!vllm bench serve --backend vllm --model {BASE_MODEL} --port 8001 \
    --dataset-name random --random-input-len 4096 --random-output-len 128 --random-range-ratio 0.9997 \
    --temperature 0 --ignore-eos --num-warmups 1 --percentile-metrics ttft,tpot,itl,e2el \
    --num-prompts {NUM_PROMPTS} --request-rate 1.0 \
    --save-result --result-filename vllm_base.json | tee vllm_base_log

The benchmark prints its results, keeps them in the log file vllm_base_log, and saves them to vllm_base.json for the comparison at the end.

Enabling speculative decoding in serving#

First stop the server without speculative decoding:

stop_server(vllm_server)

# The old server can take a moment to release the port after it stops.
deadline = time.time() + 120
while server_is_up("http://localhost:8001/v1/models") and time.time() < deadline:
    time.sleep(2)

Then start the server with speculative decoding enabled, using the same SERVE command. --speculative-config names the draft model, the number of tokens it proposes at each step, and its quantization; it replaces the --speculative-model, --num_speculative_tokens and --speculative-model-quantization options of earlier vLLM releases.

The draft model proposes three tokens at each step. Each proposed token costs a draft-model step and makes the base model’s verification step a little longer, and the later tokens of a proposal are accepted less often, so proposing more tokens is not always better. Try other values of num_speculative_tokens to see the trade-off on your own workload:

import json

SPEC_CONFIG = json.dumps({"model": DRAFT_MODEL, "num_speculative_tokens": 3, "quantization": "fp8"})
vllm_server = ensure_server(SERVE + f" --speculative-config '{SPEC_CONFIG}'",
                            "http://localhost:8001/v1/models", "vllm_spec_dec.log", timeout=3600)

As before, the cell waits until the server answers; its log is vllm_spec_dec.log.

Check that the server lists the model:

!curl http://localhost:8001/v1/models

Sending requests#

Run the same benchmark against the server with speculative decoding:

!vllm bench serve --backend vllm --model {BASE_MODEL} --port 8001 \
    --dataset-name random --random-input-len 4096 --random-output-len 64 \
    --temperature 0 --ignore-eos --num-warmups 0 --num-prompts 48 --request-rate 4 > /dev/null
!vllm bench serve --backend vllm --model {BASE_MODEL} --port 8001 \
    --dataset-name random --random-input-len 4096 --random-output-len 128 --random-range-ratio 0.9997 \
    --temperature 0 --ignore-eos --num-warmups 1 --percentile-metrics ttft,tpot,itl,e2el \
    --num-prompts {NUM_PROMPTS} --request-rate 1.0 \
    --save-result --result-filename vllm_spec_dec.json | tee vllm_spec_dec_log

Compare the two runs and stop the server:

import json

base = json.load(open("vllm_base.json"))
spec = json.load(open("vllm_spec_dec.json"))
assert base["completed"] == spec["completed"] == NUM_PROMPTS, "some benchmark requests failed"
print(f"{'':34} {'baseline':>10} {'speculative':>12}")
for name, key in [("Mean end-to-end latency (ms)", "mean_e2el_ms"),
                  ("Mean time per output token (ms)", "mean_tpot_ms"),
                  ("Mean time to first token (ms)", "mean_ttft_ms"),
                  ("Output token throughput (tok/s)", "output_throughput")]:
    print(f"{name:34} {base[key]:10.1f} {spec[key]:12.1f}")
print(f"Tokens per base-model step (accepted draft tokens + 1): {spec['spec_decode_acceptance_length']:.2f}")
assert spec["spec_decode_acceptance_length"] > 1, "the base model accepted none of the draft tokens"

stop_server(vllm_server)

The table shows both runs side by side, and the last line shows how many tokens each base-model step produced: one plus the draft tokens the base model accepted. A value above one means the draft model’s proposals are being accepted.

When you compare the runs, keep in mind what speculative decoding changes. It affects generating tokens, not reading the prompt, so look at the time per output token and the end-to-end latency rather than the time to the first token. The output token throughput reflects the request rate: the benchmark sends one request per second to either server.

The effect of speculative decoding depends on the vLLM release, the workload and the hardware. With many requests in flight, the base model’s steps are already full, and verifying proposed tokens can cost more than it saves.

To explore speculative decoding more thoroughly, you can modify the value of each parameter (for example, --request-rate in the benchmark commands, or num_speculative_tokens in SPEC_CONFIG) to observe the performance changes. For more information, see the study in the Speculative Decoding - Deep Dive blog.

Contents