iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[CI] Solidify speculative decoding E2E coverage (#50713)

Browse Source 
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Andreas Karatzas
2026-08-10 19:56:39 -06:00
committed by GitHub
co-authored by OpenAI Codex
parent 48bada6ea4
commit 1a1727330a
18 changed files with 737 additions and 459 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/spec_decode.yaml
+98 -6
View File
@@ -29,8 +29,8 @@ steps:
depends_on:
- image-build-amd
- label: Spec Decode Eagle
key: spec-decode-eagle
- label: "Spec Decode Eagle 1: DeepSeek + Qwen"
key: spec-decode-eagle-1-deepseek-qwen
timeout_in_minutes: 25
device: h200_35gb
source_file_dependencies:
@@ -38,12 +38,38 @@ steps:
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode/eagle/
- pytest -v -s v1/e2e/spec_decode/eagle/ -k "deepseek_eagle or qwen3_eagle3"
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 55
timeout_in_minutes: 25
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/model_loader/
- vllm/v1/sample/
- vllm/model_executor/layers/
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
- label: "Spec Decode Eagle 2: Llama 3 + Qwen VL + Other"
key: spec-decode-eagle-2-llama3-qwen-vl-other
timeout_in_minutes: 25
device: h200_35gb
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode/eagle/ -k "not deepseek_eagle and not qwen3_eagle3"
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 25
depends_on:
- image-build-amd
source_file_dependencies:
@@ -192,8 +218,8 @@ steps:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
- label: Spec Decode DeepSeek MTP Parallel Load (B200)
key: spec-decode-deepseek-mtp-parallel-load-b200
- label: Spec Decode DeepSeek MTP Parallel Load (2xB200-2xMI300)
key: spec-decode-deepseek-mtp-parallel-load-2xb200-2xmi300
timeout_in_minutes: 30
device: b200-k8s
optional: true
@@ -207,6 +233,26 @@ steps:
- tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
commands:
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
mirror:
amd:
dind: false
device: mi300_2
timeout_in_minutes: 45
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/distributed/
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/
- vllm/model_executor/models/deepseek_mtp.py
- vllm/model_executor/models/deepseek_v2.py
- vllm/v1/attention/
- vllm/v1/engine/
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/v1/worker/gpu_model_runner.py
- tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
- vllm/platforms/rocm.py
- label: Spec Decode AL DFlash Nightly
key: spec-decode-dflash-nightly
@@ -221,6 +267,27 @@ steps:
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 120
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/v1/attention/
- vllm/v1/sample/
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/
- vllm/model_executor/models/qwen2.py
- vllm/model_executor/models/qwen3.py
- vllm/model_executor/models/qwen3_dflash.py
- vllm/model_executor/models/laguna.py
- vllm/model_executor/models/laguna_dflash.py
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
- label: Spec Decode AL DSpark Nightly
key: spec-decode-dspark-nightly
@@ -235,6 +302,31 @@ steps:
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/
mirror:
amd:
dind: false
device: mi300_1
timeout_in_minutes: 90
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/v1/attention/
- vllm/v1/sample/
- vllm/model_executor/model_loader/
- vllm/model_executor/layers/
- vllm/model_executor/models/qwen2.py
- vllm/model_executor/models/qwen3.py
- vllm/model_executor/models/qwen3_dflash.py
- vllm/model_executor/models/qwen3_dspark.py
- vllm/model_executor/models/gemma4.py
- vllm/model_executor/models/gemma4_mm.py
- vllm/model_executor/models/gemma4_unified.py
- vllm/model_executor/models/gemma4_mtp.py
- vllm/model_executor/models/gemma4_dspark.py
- tests/v1/e2e/spec_decode/
- vllm/platforms/rocm.py
- label: Spec Decode AL MTP + Other Acceptance Nightly
key: spec-decode-mtp-other-acceptance-nightly
tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py
+33 -24
View File
@@ -4,12 +4,10 @@
from dataclasses import dataclass
import pytest
import torch
from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline
from tests.utils import single_gpu_only
from vllm import LLM
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.config import CompilationConfig
from ...utils import compute_acceptance_len
from ..utils import run_acceptance_length_eval
@@ -56,9 +54,11 @@ LAGUNA_DFLASH_NVFP4 = DFlashCorrectnessConfig(
def test_dflash_reference_acceptance_lengths(
monkeypatch: pytest.MonkeyPatch,
use_mrv2: bool,
vllm_runner,
):
run_acceptance_length_eval(
monkeypatch,
vllm_runner,
spec_config={
"model": QWEN3_DFLASH.model,
"trust_remote_code": True,
@@ -110,14 +110,16 @@ def test_dflash_correctness(
monkeypatch: pytest.MonkeyPatch,
config: DFlashCorrectnessConfig,
use_mrv2: bool,
vllm_runner,
):
"""Guard GSM8K accuracy and batched acceptance length for DFlash models."""
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0")
if config.disable_flashinfer_sampler:
monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
spec_llm = LLM(
model=config.model,
with vllm_runner(
config.model,
block_size=None,
trust_remote_code=True,
speculative_config={
"method": "dflash",
@@ -130,23 +132,30 @@ def test_dflash_correctness(
gpu_memory_utilization=0.85,
enforce_eager=config.enforce_eager,
disable_log_stats=False,
)
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
) as spec_runner:
spec_llm = spec_runner.llm
results = evaluate_gsm8k_offline(
spec_llm,
num_questions=config.num_questions,
use_chat_completions=config.use_chat_completions,
)
accuracy = results["accuracy"]
acceptance_len = compute_acceptance_len(spec_llm.get_metrics())
context = (
f"DFlash target={config.model}, draft={config.draft_model}, MRV2={use_mrv2}"
)
print(
f"{context}: GSM8K accuracy={accuracy:.3f}, "
f"acceptance_len={acceptance_len:.2f}"
)
results = evaluate_gsm8k_offline(
spec_llm,
num_questions=config.num_questions,
use_chat_completions=config.use_chat_completions,
)
accuracy = results["accuracy"]
acceptance_len = compute_acceptance_len(spec_llm.get_metrics())
print(
f"{config.model}: GSM8K accuracy={accuracy:.3f}, "
f"acceptance_len={acceptance_len:.2f}"
)
assert accuracy >= config.expected_accuracy
assert acceptance_len >= config.expected_acceptance_len
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
assert accuracy >= config.expected_accuracy, (
f"{context}: GSM8K accuracy {accuracy:.3f} is below "
f"{config.expected_accuracy:.3f}; acceptance_len={acceptance_len:.3f}"
)
assert acceptance_len >= config.expected_acceptance_len, (
f"{context}: acceptance_len {acceptance_len:.3f} is below "
f"{config.expected_acceptance_len:.3f}; accuracy={accuracy:.3f}"
)
tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py
+32 -21
View File
@@ -4,11 +4,10 @@
from dataclasses import dataclass
import pytest
import torch
from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline
from vllm import LLM
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.config import CompilationConfig
from vllm.platforms import current_platform
from ...utils import compute_acceptance_len, compute_acceptance_rate
@@ -44,6 +43,7 @@ QWEN3_DSPARK_DEEPSPEC = DSparkCorrectnessConfig(
reference_accuracy=0.801,
reference_acceptance_rate=0.428,
reference_acceptance_len=3.994,
attention_backend="FLASH_ATTN",
)
# References from five 200-question runs at temperature 1.0: accuracy
@@ -58,6 +58,9 @@ GEMMA4_DSPARK_DEEPSPEC = DSparkCorrectnessConfig(
num_questions=200,
use_chat_completions=True,
gpu_memory_utilization=0.85,
enforce_eager=True,
disable_flashinfer_sampler=True,
language_model_only=True,
)
# Reference from 200 GSM8K questions at temperature 1.0.
@@ -88,34 +91,40 @@ QWEN3_6_DSPARK_SPECULATORS = DSparkCorrectnessConfig(
def test_dspark_correctness_and_acceptance_rate(
monkeypatch: pytest.MonkeyPatch,
config: DSparkCorrectnessConfig,
vllm_runner,
):
"""Guard GSM8K accuracy and acceptance metrics for DSpark models."""
if config.disable_flashinfer_sampler:
monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
speculative_config = {
"method": "dspark",
"model": config.draft_model,
"num_speculative_tokens": config.num_speculative_tokens,
"draft_sample_method": "probabilistic",
}
if config.attention_backend is not None:
if config.attention_backend is not None and current_platform.is_cuda():
speculative_config["attention_backend"] = config.attention_backend
llm_kwargs = {
"model": config.model,
runner_config = {
"block_size": None,
"trust_remote_code": True,
"speculative_config": speculative_config,
"max_model_len": config.max_model_len,
"gpu_memory_utilization": config.gpu_memory_utilization,
"enforce_eager": config.enforce_eager,
"enable_chunked_prefill": None,
"enable_prefix_caching": False,
"disable_log_stats": False,
"compilation_config": CompilationConfig(),
}
if config.max_num_seqs is not None:
llm_kwargs["max_num_seqs"] = config.max_num_seqs
runner_config["max_num_seqs"] = config.max_num_seqs
if config.language_model_only:
llm_kwargs["language_model_only"] = True
runner_config["language_model_only"] = True
spec_llm = LLM(**llm_kwargs)
try:
with vllm_runner(config.model, **runner_config) as spec_runner:
spec_llm = spec_runner.llm
results = evaluate_gsm8k_offline(
spec_llm,
num_questions=config.num_questions,
@@ -128,18 +137,20 @@ def test_dspark_correctness_and_acceptance_rate(
metrics = spec_llm.get_metrics()
acceptance_rate = compute_acceptance_rate(metrics)
acceptance_len = compute_acceptance_len(metrics)
print(
f"{config.model}: GSM8K accuracy={accuracy:.3f}, "
f"acceptance_rate={acceptance_rate:.2f}, "
f"acceptance_len={acceptance_len:.2f}"
context = f"DSpark target={config.model}, draft={config.draft_model}"
metrics_summary = (
f"gsm8k_accuracy={accuracy:.3f}, "
f"acceptance_rate={acceptance_rate:.3f}, "
f"acceptance_len={acceptance_len:.3f}"
)
print(f"{context}: {metrics_summary}")
assert accuracy >= config.reference_accuracy * REGRESSION_TOLERANCE
assert accuracy >= config.reference_accuracy * REGRESSION_TOLERANCE, (
f"{context}: {metrics_summary}"
)
assert (
acceptance_rate >= config.reference_acceptance_rate * REGRESSION_TOLERANCE
)
assert acceptance_len >= config.reference_acceptance_len * REGRESSION_TOLERANCE
finally:
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
), f"{context}: {metrics_summary}"
assert (
acceptance_len >= config.reference_acceptance_len * REGRESSION_TOLERANCE
), f"{context}: {metrics_summary}"
tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py
+13 -13
View File
@@ -1,17 +1,16 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm import SamplingParams
from vllm.config import CompilationConfig
from ...utils import compute_acceptance_rate
def test_medusa_acceptance_rate(
sampling_config: SamplingParams,
vllm_runner,
):
"""Verify a trained Medusa checkpoint achieves nonzero acceptance rate.
@@ -23,8 +22,10 @@ def test_medusa_acceptance_rate(
medusa_model = "FasterDecoding/medusa-vicuna-7b-v1.3"
prompts = _build_gsm8k_prompts(num_questions=10, num_shots=1)[0]
spec_llm = LLM(
model=target_model,
with vllm_runner(
target_model,
block_size=None,
trust_remote_code=False,
speculative_config={
"method": "medusa",
"model": medusa_model,
@@ -32,14 +33,13 @@ def test_medusa_acceptance_rate(
},
max_model_len=1024,
enforce_eager=True,
enable_chunked_prefill=None,
disable_log_stats=False,
)
spec_llm.generate(prompts, sampling_config)
metrics = spec_llm.get_metrics()
acceptance_rate = compute_acceptance_rate(metrics)
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
compilation_config=CompilationConfig(),
) as spec_runner:
spec_runner.llm.generate(prompts, sampling_config)
metrics = spec_runner.llm.get_metrics()
acceptance_rate = compute_acceptance_rate(metrics)
min_acceptance_rate = 0.198
print(f"Medusa acceptance rate: {acceptance_rate:.4f} (min {min_acceptance_rate})")
tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py
+2
View File
@@ -10,9 +10,11 @@ from ..utils import run_acceptance_length_eval
def test_gemma4_mtp_acceptance_lengths(
monkeypatch: pytest.MonkeyPatch,
use_mrv2: bool,
vllm_runner,
):
run_acceptance_length_eval(
monkeypatch,
vllm_runner,
spec_config={
"model": "google/gemma-4-E4B-it",
"trust_remote_code": True,
tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py
+26 -28
View File
@@ -1,25 +1,24 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from tests.utils import single_gpu_only
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm import SamplingParams
from vllm.config import CompilationConfig
from ...utils import compute_acceptance_len, get_test_prompts
@single_gpu_only
def test_synthetic_acceptance_rate():
def test_synthetic_acceptance_rate(vllm_runner):
"""Verify that synthetic rejection sampling produces an acceptance
length close to the requested mean acceptance length."""
num_spec_tokens = 3
expected_acceptance_len = 1.875
tolerance = 0.15
spec_llm = LLM(
model="meta-llama/Llama-3.2-1B-Instruct",
with vllm_runner(
"meta-llama/Llama-3.2-1B-Instruct",
block_size=None,
trust_remote_code=True,
speculative_config={
"method": "eagle3",
@@ -31,28 +30,27 @@ def test_synthetic_acceptance_rate():
},
max_model_len=2048,
enforce_eager=True,
enable_chunked_prefill=None,
disable_log_stats=False,
)
compilation_config=CompilationConfig(),
) as spec_runner:
spec_llm = spec_runner.llm
test_prompts = get_test_prompts(mm_enabled=False, num_prompts=50)
spec_llm.chat(
test_prompts,
SamplingParams(temperature=0, max_tokens=64, ignore_eos=True),
)
test_prompts = get_test_prompts(mm_enabled=False, num_prompts=50)
spec_llm.chat(
test_prompts,
SamplingParams(temperature=0, max_tokens=64, ignore_eos=True),
)
metrics = spec_llm.get_metrics()
acceptance_len = compute_acceptance_len(metrics)
metrics = spec_llm.get_metrics()
acceptance_len = compute_acceptance_len(metrics)
print(
f"Synthetic acceptance length: {acceptance_len:.3f}"
f" (expected={expected_acceptance_len:.3f},"
f" tolerance=±{tolerance})"
)
assert abs(acceptance_len - expected_acceptance_len) <= tolerance, (
f"Synthetic acceptance length {acceptance_len:.3f} is not within"
f" ±{tolerance} of expected {expected_acceptance_len:.3f}"
)
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
print(
f"Synthetic acceptance length: {acceptance_len:.3f}"
f" (expected={expected_acceptance_len:.3f},"
f" tolerance=±{tolerance})"
)
assert abs(acceptance_len - expected_acceptance_len) <= tolerance, (
f"Synthetic acceptance length {acceptance_len:.3f} is not within"
f" ±{tolerance} of expected {expected_acceptance_len:.3f}"
)
tests/v1/e2e/spec_decode/acceptance_rates/utils.py
+48 -50
View File
@@ -4,11 +4,9 @@
from typing import Any
import pytest
import torch
from tqdm import tqdm
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm import SamplingParams
from vllm.config import CompilationConfig
from ..utils import compute_acceptance_len
@@ -39,6 +37,7 @@ def load_and_process_dataset(data_name: str):
def run_acceptance_length_eval(
monkeypatch: pytest.MonkeyPatch,
vllm_runner,
spec_config: dict[str, Any],
expected_acceptance_lengths: dict[str, float],
chat_template_kwargs: dict[str, Any],
@@ -52,54 +51,53 @@ def run_acceptance_length_eval(
stays within tolerance of the reference figure for each dataset.
"""
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0")
spec_llm = LLM(**spec_config)
runner_config = spec_config.copy()
model = runner_config.pop("model")
runner_config.setdefault("block_size", None)
runner_config.setdefault("enable_chunked_prefill", None)
runner_config.setdefault("compilation_config", CompilationConfig())
with vllm_runner(model, **runner_config) as spec_runner:
spec_llm = spec_runner.llm
# MT-Bench has 80 prompts and HumanEval has 164; truncate GSM8K.
max_prompts_per_dataset = 200
max_prompts_per_dataset = 200 # mt-bench has 80, humaneval has 164, truncates gsm8k
tokenizer = spec_llm.get_tokenizer()
for dataset_name, expected_len in expected_acceptance_lengths.items():
dataset = load_and_process_dataset(dataset_name)
prev_metrics = None
acceptance_lengths = []
for i in range(min(max_prompts_per_dataset, len(dataset))):
user_content = dataset[i]["turns"][0]
prompt_text = tokenizer.apply_chat_template(
[{"role": "user", "content": user_content}],
tokenize=False,
add_generation_prompt=True,
**chat_template_kwargs,
)
tokenizer = spec_llm.get_tokenizer()
for dataset_name, expected_len in expected_acceptance_lengths.items():
dataset = load_and_process_dataset(dataset_name)
prev_metrics = None
acceptance_lengths = []
for i in tqdm(
range(min(max_prompts_per_dataset, len(dataset))),
desc=f"Processing {dataset_name}",
):
user_content = dataset[i]["turns"][0]
prompt_text = tokenizer.apply_chat_template(
[{"role": "user", "content": user_content}],
tokenize=False,
add_generation_prompt=True,
**chat_template_kwargs,
# Greedy (temp=0) so acceptance length is deterministic and comparable
# across runs.
spec_llm.generate(
[prompt_text],
SamplingParams(temperature=0, max_tokens=2048),
use_tqdm=False,
)
current_metrics = spec_llm.get_metrics()
acceptance_len = compute_acceptance_len(current_metrics, prev_metrics)
prev_metrics = current_metrics
acceptance_lengths.append(acceptance_len)
mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths)
# Fairly tight tolerance of 95% against the reference figures,
# watching for regressions. Can be relaxed if test is flaky but be sure to
# check for genuine issues such as #40727.
expected_len = expected_len * 0.95
print(
f"acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}"
f" (expected at least {expected_len:.2f})"
)
# Greedy (temp=0) so acceptance length is deterministic and comparable
# across runs.
spec_llm.generate(
[prompt_text],
SamplingParams(temperature=0, max_tokens=2048),
use_tqdm=False,
assert mean_acceptance_length >= expected_len, (
f"acceptance_len for {dataset_name} is below expected threshold: "
f"{mean_acceptance_length:.2f} < {expected_len:.2f}"
)
current_metrics = spec_llm.get_metrics()
acceptance_len = compute_acceptance_len(current_metrics, prev_metrics)
prev_metrics = current_metrics
acceptance_lengths.append(acceptance_len)
mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths)
# Fairly tight tolerance of 95% against the reference figures,
# watching for regressions. Can be relaxed if test is flaky but be sure to
# check for genuine issues such as #40727.
expected_len = expected_len * 0.95
print(
f"acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}"
f" (expected at least {expected_len:.2f})"
)
assert mean_acceptance_length >= expected_len, (
f"acceptance_len for {dataset_name} is below expected threshold: "
f"{mean_acceptance_length:.2f} < {expected_len:.2f}"
)
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
tests/v1/e2e/spec_decode/draft_model/test_async.py
+29 -19
View File
@@ -102,14 +102,15 @@ def test_no_sync_with_spec_decode(
spec_model: str,
method: str,
num_spec_tokens: int,
vllm_runner,
):
"""
Test that no implicit GPU-CPU sync occurs during speculative decoding
generation.
"""
# Import vLLM AFTER sync_tracker fixture has applied the patch
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm import SamplingParams
from vllm.config import CompilationConfig
# Qwen3.5 is a VLM; without this, profile_run runs the ViT warmup
# and peaks well above the 18GB MIG slice used by one of the CI lanes.
@@ -119,8 +120,10 @@ def test_no_sync_with_spec_decode(
if "Qwen3.5" in model:
extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0}
llm = LLM(
model=model,
with vllm_runner(
model,
block_size=None,
trust_remote_code=False,
max_model_len=256,
speculative_config={
"method": method,
@@ -129,24 +132,31 @@ def test_no_sync_with_spec_decode(
},
enforce_eager=True,
async_scheduling=True,
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
**extra_kwargs,
)
) as runner:
llm = runner.llm
# Assert async scheduling is actually active before running inference.
assert llm.llm_engine.vllm_config.scheduler_config.async_scheduling, (
f"Expected async_scheduling=True for spec decode, got False. method={method}"
)
# Assert async scheduling is actually active before running inference.
assert llm.llm_engine.vllm_config.scheduler_config.async_scheduling, (
f"Expected async_scheduling=True for spec decode, got False. "
f"method={method}, target={model}, draft={spec_model}"
)
outputs = llm.generate(
["Hello, my name is"],
SamplingParams(temperature=0, max_tokens=10),
)
outputs = llm.generate(
["Hello, my name is"],
SamplingParams(temperature=0, max_tokens=10),
)
assert len(outputs) == 1
assert len(outputs[0].outputs[0].text) > 0
del llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
assert len(outputs) == 1, (
f"{method} target={model}: expected one request output, got {len(outputs)}"
)
assert outputs[0].outputs, (
f"{method} target={model}: request output has no completion candidates"
)
assert outputs[0].outputs[0].text, (
f"{method} target={model}: generated completion text is empty"
)
sync_tracker.assert_no_sync()
tests/v1/e2e/spec_decode/draft_model/test_draft_model.py
+79 -53
View File
@@ -4,17 +4,16 @@
from dataclasses import dataclass
import pytest
import torch
from tests.utils import multi_gpu_only, single_gpu_only
from vllm import LLM, SamplingParams
from vllm.config import VllmConfig, replace
from vllm.distributed import cleanup_dist_env_and_memory
from vllm import SamplingParams
from vllm.config import CompilationConfig, VllmConfig, replace
from vllm.config.kernel import MoEBackend
from vllm.engine.arg_utils import EngineArgs
from vllm.platforms import current_platform
from ..utils import (
Messages,
_skip_if_insufficient_gpus_for_tp,
compute_acceptance_len,
compute_acceptance_rate,
evaluate_llm_for_gsm8k,
@@ -42,7 +41,8 @@ class ArgsTest:
max_model_len: int = 2048
gpu_memory_utilization: float = 0.5
dataset: str = "test_prompts"
num_prompts: int = 100
# Doubling 100 to 200 reduces sampling standard error by about 29%.
num_prompts: int = 200
def get_messages(dataset: str, n: int) -> list[Messages]:
@@ -89,13 +89,13 @@ cases = [
@pytest.mark.parametrize("args", cases)
@pytest.mark.parametrize("enforce_eager", [True, False])
@single_gpu_only
def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool, vllm_runner):
args.enforce_eager = enforce_eager
assert_draft_model_correctness(args)
assert_draft_model_correctness(args, vllm_runner)
@single_gpu_only
def test_draft_model_realistic_example():
def test_draft_model_realistic_example(vllm_runner):
args = ArgsTest(
target_model="Qwen/Qwen3-1.7B",
draft_model="Qwen/Qwen3-0.6B",
@@ -106,11 +106,11 @@ def test_draft_model_realistic_example():
expected_acceptance_len=2.6, # ref: 2.86
expected_acceptance_rate=0.5, # ref: 0.62
)
assert_draft_model_correctness(args)
assert_draft_model_correctness(args, vllm_runner)
@single_gpu_only
def test_draft_model_parallel_drafting():
def test_draft_model_parallel_drafting(vllm_runner):
args = ArgsTest(
target_model="Qwen/Qwen3-1.7B",
draft_model="amd/PARD-Qwen3-0.6B",
@@ -122,7 +122,7 @@ def test_draft_model_parallel_drafting():
expected_acceptance_len=2.3, # ref: 2.52
expected_acceptance_rate=0.4, # ref: 0.51
)
assert_draft_model_correctness(args)
assert_draft_model_correctness(args, vllm_runner)
@pytest.mark.parametrize(
@@ -136,7 +136,9 @@ def test_draft_model_parallel_drafting():
)
@pytest.mark.parametrize("enforce_eager", [True, False])
@single_gpu_only
def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
def test_draft_model_quantization(
models: tuple[str, str], enforce_eager: bool, vllm_runner
):
tgt_model, draft_model = models
sd_case = ArgsTest(
target_model=tgt_model,
@@ -144,13 +146,12 @@ def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
**some_high_acceptance_metrics(),
enforce_eager=enforce_eager,
)
assert_draft_model_correctness(sd_case)
assert_draft_model_correctness(sd_case, vllm_runner)
@multi_gpu_only(num_gpus=2)
def test_draft_model_tensor_parallelism():
def test_draft_model_tensor_parallelism(vllm_runner):
"""Ensure spec decode works when running with TP > 1."""
_skip_if_insufficient_gpus_for_tp(2)
sd_case = ArgsTest(
target_model="Qwen/Qwen3-1.7B",
target_tensor_parallel_size=2,
@@ -160,15 +161,13 @@ def test_draft_model_tensor_parallelism():
enforce_eager=False,
expected_gsm8k_accuracy=0.5,
)
assert_draft_model_correctness(sd_case)
assert_draft_model_correctness(sd_case, vllm_runner)
@multi_gpu_only(num_gpus=2)
def test_draft_model_engine_args_tensor_parallelism():
"""Ensure the vllm_config for the draft model is created correctly,
and independently of the target model (quantization, TP, etc.)"""
_skip_if_insufficient_gpus_for_tp(2)
engine_args = EngineArgs(
model="Qwen/Qwen3-1.7B-FP8", # <<< tgt quantized
tensor_parallel_size=2,
@@ -212,13 +211,24 @@ def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig:
return vllm_config
def test_draft_model_moe_backend_override():
def _platform_moe_backend(monkeypatch: pytest.MonkeyPatch) -> MoEBackend:
if current_platform.is_rocm():
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
return "aiter"
if current_platform.is_cuda():
return "flashinfer_trtllm"
return "auto"
def test_draft_model_moe_backend_override(monkeypatch: pytest.MonkeyPatch):
"""When moe_backend is set in speculative_config, the draft VllmConfig
should use it while the target keeps its own setting."""
target_moe_backend = _platform_moe_backend(monkeypatch)
engine_args = EngineArgs(
model="Qwen/Qwen3-1.7B",
tensor_parallel_size=1,
moe_backend="flashinfer_trtllm",
moe_backend=target_moe_backend,
speculative_config={
"model": "Qwen/Qwen3-0.6B",
"method": "draft_model",
@@ -227,22 +237,23 @@ def test_draft_model_moe_backend_override():
},
)
tgt_config: VllmConfig = engine_args.create_engine_config()
assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm"
assert tgt_config.kernel_config.moe_backend == target_moe_backend
assert tgt_config.speculative_config.moe_backend == "triton"
draft_config = _apply_draft_moe_backend(tgt_config)
assert draft_config.kernel_config.moe_backend == "triton"
# Target config must be unaffected.
assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm"
assert tgt_config.kernel_config.moe_backend == target_moe_backend
def test_draft_model_moe_backend_inherits_target():
def test_draft_model_moe_backend_inherits_target(monkeypatch: pytest.MonkeyPatch):
"""When moe_backend is not set in speculative_config, the draft should
inherit the target's moe_backend."""
target_moe_backend = _platform_moe_backend(monkeypatch)
engine_args = EngineArgs(
model="Qwen/Qwen3-1.7B",
tensor_parallel_size=1,
moe_backend="flashinfer_cutlass",
moe_backend=target_moe_backend,
speculative_config={
"model": "Qwen/Qwen3-0.6B",
"method": "draft_model",
@@ -250,11 +261,11 @@ def test_draft_model_moe_backend_inherits_target():
},
)
tgt_config: VllmConfig = engine_args.create_engine_config()
assert tgt_config.kernel_config.moe_backend == "flashinfer_cutlass"
assert tgt_config.kernel_config.moe_backend == target_moe_backend
assert tgt_config.speculative_config.moe_backend is None
draft_config = _apply_draft_moe_backend(tgt_config)
assert draft_config.kernel_config.moe_backend == "flashinfer_cutlass"
assert draft_config.kernel_config.moe_backend == target_moe_backend
assert draft_config is tgt_config
@@ -293,19 +304,23 @@ def test_draft_model_engine_args_rejects_invalid_tp_argname():
"tensor_parallel_size": 1, # <<< invalid arg name
},
)
with pytest.raises(ValueError):
with pytest.raises(ValueError, match="draft_tensor_parallel_size"):
engine_args.create_engine_config()
def assert_draft_model_correctness(args: ArgsTest):
def assert_draft_model_correctness(args: ArgsTest, vllm_runner):
"""Compare the outputs using and not using speculative decoding.
In the greedy decoding case, the outputs must match EXACTLY."""
test_prompts: list[Messages] = get_messages(
dataset=args.dataset, n=args.num_prompts
)
spec_llm = LLM(
model=args.target_model,
with vllm_runner(
args.target_model,
block_size=None,
trust_remote_code=False,
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
speculative_config={
"model": args.draft_model,
"method": "draft_model",
@@ -321,30 +336,41 @@ def assert_draft_model_correctness(args: ArgsTest):
tensor_parallel_size=args.target_tensor_parallel_size,
enforce_eager=args.enforce_eager,
disable_log_stats=False, # enables get_metrics()
)
) as spec_runner:
spec_llm = spec_runner.llm
# we don't check the outputs, only check the metrics
spec_llm.chat(test_prompts, args.sampling_config)
metrics = spec_llm.get_metrics()
acceptance_rate: float = compute_acceptance_rate(metrics)
acceptance_len: float = compute_acceptance_len(metrics)
# we don't check the outputs, only check the metrics
spec_llm.chat(test_prompts, args.sampling_config)
metrics = spec_llm.get_metrics()
acceptance_rate: float = compute_acceptance_rate(metrics)
acceptance_len: float = compute_acceptance_len(metrics)
# Need to evaluate after getting metrics to avoid polluting the AR
evaluate_llm_for_gsm8k(
spec_llm, expected_accuracy_threshold=args.expected_gsm8k_accuracy
)
# Need to evaluate after getting metrics to avoid polluting the AR
evaluate_llm_for_gsm8k(
spec_llm, expected_accuracy_threshold=args.expected_gsm8k_accuracy
)
print(
f"spec-decode: target={args.target_model}, draft={args.draft_model}, "
f"temperature={args.sampling_config.temperature:.2f}, "
f"acceptance_rate={acceptance_rate:.2f}, "
f"acceptance_len={acceptance_len:.2f}, "
)
print(
f"spec-decode: target={args.target_model}, draft={args.draft_model}, "
f"temperature={args.sampling_config.temperature:.2f}, "
f"acceptance_rate={acceptance_rate:.2f}, "
f"acceptance_len={acceptance_len:.2f}, "
)
context = (
f"draft_model target={args.target_model}, draft={args.draft_model}, "
f"eager={args.enforce_eager}, "
f"acceptance_rate={acceptance_rate:.3f}, "
f"acceptance_len={acceptance_len:.3f}"
)
assert acceptance_rate >= args.expected_acceptance_rate, (
f"{context}; expected acceptance_rate >= "
f"{args.expected_acceptance_rate:.3f}"
)
assert acceptance_len >= args.expected_acceptance_len, (
f"{context}; expected acceptance_len >= {args.expected_acceptance_len:.3f}"
)
# draft_model supports async scheduling; assert it is active by default.
has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
assert acceptance_rate >= args.expected_acceptance_rate
assert acceptance_len >= args.expected_acceptance_len
has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
del spec_llm # CLEANUP
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
assert has_async, "Expected async_scheduling=True for draft_model spec decode"
tests/v1/e2e/spec_decode/draft_model/test_lora.py
+38 -41
View File
@@ -8,12 +8,14 @@ This script contains:
import pytest
import torch
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm import SamplingParams
from vllm.config import CompilationConfig
from vllm.lora.request import LoRARequest
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
from ..utils import assert_request_outputs_match
LORA_TEST_PROMPT_MAP: dict[str, str] = {}
LORA_TEST_PROMPT_MAP["premjatin/qwen-linear-algebra-coder"] = """
@@ -36,7 +38,6 @@ Numbers should be represented as integers only.
SEED = 42
@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
@pytest.mark.parametrize(
"model_setup",
[
@@ -49,9 +50,13 @@ SEED = 42
)
],
)
@pytest.mark.skipif(
not current_platform.is_cuda_alike(), reason="Requires CUDA or ROCm"
)
def test_batch_inference_correctness(
monkeypatch: pytest.MonkeyPatch,
model_setup: tuple[str, str, str, str, int],
vllm_runner,
):
"""
Compare the outputs of a LLM with only Lora and a LLM with both SD and Lora.
@@ -60,10 +65,11 @@ def test_batch_inference_correctness(
"""
with monkeypatch.context() as m:
# Disable randomness
m.setenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
if current_platform.is_cuda():
m.setenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
set_random_seed(SEED)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
m.setattr(torch.backends.cudnn, "benchmark", False)
m.setattr(torch.backends.cudnn, "deterministic", True)
method, model_name, spec_model_name, lora_path, tp_size = model_setup
@@ -74,8 +80,9 @@ def test_batch_inference_correctness(
)
# without speculative decoding
ref_llm = LLM(
model=model_name,
with vllm_runner(
model_name,
block_size=None,
trust_remote_code=True,
tensor_parallel_size=tp_size,
max_model_len=2048,
@@ -84,18 +91,16 @@ def test_batch_inference_correctness(
max_loras=1,
max_cpu_loras=1,
max_lora_rank=16,
)
try:
ref_outputs = ref_llm.generate(
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
) as ref_runner:
ref_outputs = ref_runner.llm.generate(
prompts, sampling_params, lora_request=lora_request
)
finally:
del ref_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
lora_spec_llm = LLM(
model=model_name,
with vllm_runner(
model_name,
block_size=None,
trust_remote_code=True,
tensor_parallel_size=tp_size,
speculative_config={
@@ -110,30 +115,22 @@ def test_batch_inference_correctness(
max_loras=1,
max_cpu_loras=1,
max_lora_rank=16,
)
try:
lora_spec_outputs = lora_spec_llm.generate(
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
) as spec_runner:
lora_spec_outputs = spec_runner.llm.generate(
prompts, sampling_params, lora_request=lora_request
)
matches = 0
for ref_output, spec_output in zip(ref_outputs, lora_spec_outputs):
if ref_output.outputs[0].text == spec_output.outputs[0].text:
matches += 1
else:
print(f"ref_output: {ref_output.outputs[0].text}")
print(f"spec_output: {spec_output.outputs[0].text}")
# Heuristic threshold: under greedy verification, the spec-decode
# output should equal the non-spec output (modulo FP noise from the
# target's verify-path matmul running at seqlen
# num_speculative_tokens+1 vs 1). 90% leaves slack for that noise.
threshold = int(0.90 * len(ref_outputs))
print(f"match ratio: {matches}/{len(ref_outputs)}")
assert matches > threshold, (
f"match ratio {matches}/{len(ref_outputs)} <= {threshold}"
)
finally:
del lora_spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
# Under greedy verification, the spec-decode output should equal the
# non-spec output (modulo FP noise from the target's verify-path matmul
# running at seqlen num_speculative_tokens+1 vs 1). 90% leaves slack.
assert_request_outputs_match(
ref_outputs,
lora_spec_outputs,
required_matches=int(0.90 * len(ref_outputs)) + 1,
context=(
f"LoRA + {method}: target={model_name}, draft={spec_model_name}, "
f"adapter={lora_path}"
),
)
tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py
+12 -7
View File
@@ -52,6 +52,7 @@ def test_eagle_correctness_light(
enable_chunked_prefill: bool,
model_impl: str,
attn_backend: str,
vllm_runner,
):
_run_eagle_correctness(
monkeypatch,
@@ -62,6 +63,7 @@ def test_eagle_correctness_light(
enable_chunked_prefill,
model_impl,
attn_backend,
vllm_runner,
)
@@ -89,11 +91,12 @@ def test_eagle_correctness_light(
"transformers",
0.8,
# TODO(hmellor): figure out why memory usage is so high
marks=pytest.mark.skip(
reason="Feature is experimental and uses too much memory in CI",
marks=pytest.mark.skipif(
not current_platform.is_rocm(),
reason="Feature is experimental and uses too much memory in CUDA CI",
),
),
pytest.param(
(
(
"eagle3",
"Qwen/Qwen3-VL-8B-Instruct",
@@ -104,9 +107,6 @@ def test_eagle_correctness_light(
False,
"auto",
0.8,
marks=pytest.mark.skip(
reason="architecture of its eagle3 is LlamaForCausalLMEagle3"
),
),
pytest.param(
(
@@ -119,8 +119,11 @@ def test_eagle_correctness_light(
False,
"auto",
0.7,
# TODO: Re-measure the reference threshold when re-enabling this case.
# Text-only mode starts but loses target correctness with ROCm/TRITON;
# full multimodal profiling currently fails during engine startup.
marks=pytest.mark.skip(
reason="Skipping due to its head_dim not being a multiple of 32"
reason="Qwen2.5-VL Eagle3 is not yet reliable in this test setup"
),
),
(
@@ -154,6 +157,7 @@ def test_eagle_correctness_medium(
enable_chunked_prefill: bool,
model_impl: str,
attn_backend: str,
vllm_runner,
):
_run_eagle_correctness(
monkeypatch,
@@ -164,4 +168,5 @@ def test_eagle_correctness_medium(
enable_chunked_prefill,
model_impl,
attn_backend,
vllm_runner,
)
tests/v1/e2e/spec_decode/eagle/utils.py
+58 -59
View File
@@ -1,16 +1,17 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from typing import Any
from tests.utils import wait_for_rocm_memory_to_settle
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
import pytest
from vllm import SamplingParams
from vllm.config import CompilationConfig
from vllm.platforms import current_platform
from ..utils import (
_skip_if_insufficient_gpus_for_tp,
assert_request_outputs_match,
evaluate_llm_for_gsm8k,
get_test_prompts,
)
@@ -25,26 +26,24 @@ def _run_eagle_correctness(
enable_chunked_prefill: bool,
model_impl: str,
attn_backend: str,
vllm_runner,
):
"""
Compare the outputs of an original LLM and a speculative LLM
which should be the same when using eagle speculative decoding.
"""
if model_impl == "transformers":
import transformers
from packaging.version import Version
installed = Version(transformers.__version__)
required = Version("5.0.0")
if installed < required:
pytest.skip(
"Eagle3 with the Transformers modeling backend requires "
f"transformers>={required}, but got {installed}"
)
method, model_name, spec_model_name, tp_size = model_setup
_skip_if_insufficient_gpus_for_tp(tp_size)
test_prompts = get_test_prompts(mm_enabled)
if "Llama-4-Scout" in model_setup[1] and attn_backend == "FLASH_ATTN":
extra_kwargs: dict[str, Any] = {}
if not mm_enabled and "Qwen3-VL" in model_name:
# These cases only exercise text generation. Avoid profiling an unused
# vision tower, which adds substantial memory to both reference runs.
extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0}
if "Llama-4-Scout" in model_name and attn_backend == "FLASH_ATTN":
if current_platform.is_rocm():
print(
"FLASH_ATTN for spec_decode not supported on "
@@ -66,38 +65,36 @@ def _run_eagle_correctness(
m.setenv("VLLM_MLA_DISABLE", "1")
if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
if "deepseek" in model_setup[1].lower():
if "deepseek" in model_name.lower():
m.setenv("VLLM_ROCM_USE_AITER", "1")
m.delenv("VLLM_MLA_DISABLE", raising=False)
attention_config = {"backend": "ROCM_AITER_MLA"}
else:
m.setenv("VLLM_ROCM_USE_AITER", "1")
method, model_name, spec_model_name, tp_size = model_setup
_skip_if_insufficient_gpus_for_tp(tp_size)
max_model_len = 2048
max_num_batched_tokens = 128 if enable_chunked_prefill else max_model_len
ref_llm = LLM(
model=model_name,
with vllm_runner(
model_name,
block_size=None,
trust_remote_code=False,
max_model_len=max_model_len,
tensor_parallel_size=tp_size,
attention_config=attention_config,
)
evaluate_llm_for_gsm8k(
ref_llm, expected_accuracy_threshold=expected_accuracy_threshold
)
ref_outputs = ref_llm.chat(test_prompts, sampling_config)
del ref_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
# ROCm frees VRAM lazily; wait so the spec engine started right after
# does not OOM on its startup memory guard.
wait_for_rocm_memory_to_settle()
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
**extra_kwargs,
) as ref_runner:
evaluate_llm_for_gsm8k(
ref_runner.llm,
expected_accuracy_threshold=expected_accuracy_threshold,
)
ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config)
spec_llm = LLM(
model=model_name,
with vllm_runner(
model_name,
block_size=None,
trust_remote_code=True,
tensor_parallel_size=tp_size,
speculative_config={
@@ -111,27 +108,29 @@ def _run_eagle_correctness(
enable_chunked_prefill=enable_chunked_prefill,
model_impl=model_impl,
attention_config=attention_config,
)
# EAGLE/EAGLE3 supports async scheduling; assert it is active by default.
assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
evaluate_llm_for_gsm8k(
spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
)
spec_outputs = spec_llm.chat(test_prompts, sampling_config)
matches = 0
misses = 0
for ref_output, spec_output in zip(ref_outputs, spec_outputs):
if ref_output.outputs[0].text == spec_output.outputs[0].text:
matches += 1
else:
misses += 1
print(f"ref_output: {ref_output.outputs[0].text}")
print(f"spec_output: {spec_output.outputs[0].text}")
compilation_config=CompilationConfig(),
**extra_kwargs,
) as spec_runner:
# EAGLE/EAGLE3 supports async scheduling by default.
has_async = (
spec_runner.llm.llm_engine.vllm_config.scheduler_config.async_scheduling
)
assert has_async, (
f"Expected async scheduling for {method}: target={model_name}, "
f"draft={spec_model_name}, backend={attn_backend}; got {has_async}"
)
evaluate_llm_for_gsm8k(
spec_runner.llm,
expected_accuracy_threshold=expected_accuracy_threshold,
)
spec_outputs = spec_runner.llm.chat(test_prompts, sampling_config)
assert matches > int(0.6 * len(ref_outputs))
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
# ROCm frees VRAM lazily; wait so the next parametrization's engine does
# not OOM on its startup memory guard.
wait_for_rocm_memory_to_settle()
assert_request_outputs_match(
ref_outputs,
spec_outputs,
required_matches=int(0.6 * len(ref_outputs)) + 1,
context=(
f"{method} target={model_name}, draft={spec_model_name}, "
f"backend={attn_backend}"
),
)
tests/v1/e2e/spec_decode/mtp/test_mtp.py
+48 -46
View File
@@ -4,15 +4,15 @@
from typing import Any
import pytest
import torch
from tests.utils import single_gpu_only
from vllm import LLM, SamplingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm import SamplingParams
from vllm.config import CompilationConfig
from vllm.platforms import current_platform
from ..utils import (
_skip_if_insufficient_gpus_for_tp,
assert_request_outputs_match,
evaluate_llm_for_gsm8k,
get_test_prompts,
)
@@ -21,9 +21,10 @@ from ..utils import (
@pytest.mark.parametrize(
["model_setup", "mm_enabled", "expected_accuracy_threshold"],
[
(("mtp", "XiaomiMiMo/MiMo-7B-Base", 1), False, 0.5), # ref: 65%-70%
# Reference accuracy: 65%-70%.
(("mtp", "XiaomiMiMo/MiMo-7B-Base", 1, None), False, 0.5),
pytest.param(
("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1),
("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1, None),
False,
0.0,
marks=pytest.mark.skipif(
@@ -32,7 +33,7 @@ from ..utils import (
),
), # dummy model
(
("mtp", "Qwen/Qwen3.5-0.8B-Base", 1),
("mtp", "Qwen/Qwen3.5-0.8B-Base", 1, None),
False,
0.20,
), # hybrid + MTP, ref: ~34%-35%
@@ -48,9 +49,10 @@ from ..utils import (
def test_mtp_correctness(
monkeypatch: pytest.MonkeyPatch,
sampling_config: SamplingParams,
model_setup: tuple[str, str, int] | tuple[str, str, int, str],
model_setup: tuple[str, str, int, str | None],
mm_enabled: bool,
expected_accuracy_threshold: float,
vllm_runner,
):
"""
Compare the outputs of a original LLM and a speculative LLM
@@ -59,18 +61,14 @@ def test_mtp_correctness(
6/10 output tokens match exactly, and that the GSM8k accuracy is above a precomputed
reference threshold for each model.
"""
method, model_name, tp_size, draft_model = model_setup
_skip_if_insufficient_gpus_for_tp(tp_size)
# Generate test prompts inside the function instead of using fixture
test_prompts = get_test_prompts(mm_enabled)
with monkeypatch.context() as m:
m.setenv("VLLM_MLA_DISABLE", "1")
if len(model_setup) == 4:
method, model_name, tp_size, draft_model = model_setup
else:
method, model_name, tp_size = model_setup
draft_model = None
_skip_if_insufficient_gpus_for_tp(tp_size)
attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
# Skip multimodal profiling for models that don't need it in this test.
@@ -90,21 +88,22 @@ def test_mtp_correctness(
f"got {transformers.__version__}"
)
ref_llm = LLM(
model=model_name,
with vllm_runner(
model_name,
block_size=None,
max_model_len=2048,
tensor_parallel_size=tp_size,
trust_remote_code=True,
attention_backend=attn_backend,
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
**extra_kwargs,
)
ref_outputs = ref_llm.chat(test_prompts, sampling_config)
evaluate_llm_for_gsm8k(
ref_llm, expected_accuracy_threshold=expected_accuracy_threshold
)
del ref_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
) as ref_runner:
ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config)
evaluate_llm_for_gsm8k(
ref_runner.llm,
expected_accuracy_threshold=expected_accuracy_threshold,
)
speculative_config: dict[str, Any] = {
"method": method,
@@ -115,34 +114,37 @@ def test_mtp_correctness(
speculative_config["model"] = draft_model
speculative_config["num_speculative_tokens"] = 2
spec_llm = LLM(
model=model_name,
with vllm_runner(
model_name,
block_size=None,
trust_remote_code=True,
tensor_parallel_size=tp_size,
speculative_config=speculative_config,
max_model_len=2048,
attention_backend=attn_backend,
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
**extra_kwargs,
)
# MTP supports async scheduling; assert it is active by default.
assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
evaluate_llm_for_gsm8k(
spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
)
spec_outputs = spec_llm.chat(test_prompts, sampling_config)
matches = 0
misses = 0
for ref_output, spec_output in zip(ref_outputs, spec_outputs):
if ref_output.outputs[0].text == spec_output.outputs[0].text:
matches += 1
else:
misses += 1
print(f"ref_output: {ref_output.outputs[0].text}")
print(f"spec_output: {spec_output.outputs[0].text}")
) as spec_runner:
# MTP supports async scheduling by default.
has_async = (
spec_runner.llm.llm_engine.vllm_config.scheduler_config.async_scheduling
)
assert has_async, (
f"Expected async scheduling for {method}: target={model_name}, "
f"draft={draft_model}; got {has_async}"
)
evaluate_llm_for_gsm8k(
spec_runner.llm,
expected_accuracy_threshold=expected_accuracy_threshold,
)
spec_outputs = spec_runner.llm.chat(test_prompts, sampling_config)
# Heuristic: expect at least 80% of the prompts to match exactly
# Upon failure, inspect the outputs to check for inaccuracy.
assert matches > int(0.8 * len(ref_outputs))
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
assert_request_outputs_match(
ref_outputs,
spec_outputs,
required_matches=int(0.8 * len(ref_outputs)) + 1,
context=f"{method} target={model_name}, draft={draft_model}",
)
tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py
+67 -39
View File
@@ -2,15 +2,17 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from tests.utils import single_gpu_only
from vllm import LLM, SamplingParams
from vllm import SamplingParams
from vllm.config import CompilationConfig
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.platforms import current_platform
from ..utils import evaluate_llm_for_gsm8k, get_test_prompts
from ..utils import (
evaluate_llm_for_gsm8k,
get_spec_decode_metric_value,
get_test_prompts,
)
@pytest.fixture
@@ -43,6 +45,7 @@ def test_ngram_and_suffix_correctness(
):
with vllm_runner(
model_name,
block_size=None,
# Keep LLM defaults; VllmRunner only provides lifecycle cleanup here.
trust_remote_code=False,
enable_chunked_prefill=None,
@@ -59,6 +62,7 @@ def test_ngram_and_suffix_correctness(
@single_gpu_only
def test_ngram_gpu_default_with_async_scheduling(
async_scheduling: bool,
vllm_runner,
):
"""
Test ngram_gpu speculative decoding (k=3) correctness with and without
@@ -66,8 +70,10 @@ def test_ngram_gpu_default_with_async_scheduling(
Uses Qwen/Qwen3-8B (ref GSM8K accuracy: 87%-92%).
"""
qwen3_model = "Qwen/Qwen3-8B"
spec_llm = LLM(
model=qwen3_model,
with vllm_runner(
qwen3_model,
block_size=None,
trust_remote_code=False,
speculative_config={
"method": "ngram_gpu",
"prompt_lookup_max": 3,
@@ -76,22 +82,22 @@ def test_ngram_gpu_default_with_async_scheduling(
},
max_model_len=4096,
async_scheduling=async_scheduling,
)
# Assert the resolved async_scheduling config matches what was requested.
assert (
spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
== async_scheduling
)
evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8)
del spec_llm
cleanup_dist_env_and_memory()
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
) as spec_runner:
# Assert the resolved async_scheduling config matches what was requested.
assert (
spec_runner.llm.llm_engine.vllm_config.scheduler_config.async_scheduling
== async_scheduling
)
evaluate_llm_for_gsm8k(spec_runner.llm, expected_accuracy_threshold=0.8)
@single_gpu_only
def test_suffix_decoding_acceptance(
monkeypatch: pytest.MonkeyPatch,
sampling_config: SamplingParams,
model_name: str,
vllm_runner,
):
"""
Check that suffix decoding caching takes effect and improves acceptance
@@ -99,8 +105,10 @@ def test_suffix_decoding_acceptance(
"""
test_prompts = get_test_prompts(mm_enabled=False)
spec_llm = LLM(
model=model_name,
with vllm_runner(
model_name,
block_size=None,
trust_remote_code=False,
speculative_config={
"method": "suffix",
"suffix_decoding_max_spec_factor": 2.0,
@@ -108,40 +116,60 @@ def test_suffix_decoding_acceptance(
},
max_model_len=1024,
disable_log_stats=False,
)
# Run several times and check that the accepted tokens increase.
num_draft = []
num_accept = []
for i in range(10): # Run multiple times to warm up the cache.
spec_llm.chat(test_prompts, sampling_config)
# Collect draft and acceptance stats.
metrics = spec_llm.get_metrics()
for metric in metrics:
if metric.name == "vllm:spec_decode_num_draft_tokens":
num_draft.append(metric.value)
if metric.name == "vllm:spec_decode_num_accepted_tokens":
num_accept.append(metric.value)
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
) as spec_runner:
# Run several times and check that the accepted tokens increase.
num_draft = []
num_accept = []
for _ in range(10): # Run multiple times to warm up the cache.
spec_runner.llm.chat(test_prompts, sampling_config)
# Collect draft and acceptance stats.
metrics = spec_runner.llm.get_metrics()
num_draft.append(
get_spec_decode_metric_value(
metrics, "vllm:spec_decode_num_draft_tokens"
)
)
num_accept.append(
get_spec_decode_metric_value(
metrics, "vllm:spec_decode_num_accepted_tokens"
)
)
# Calculate the acceptance rates for the first and last runs.
first_accept_tokens = num_accept[0]
first_draft_tokens = num_draft[0]
assert first_draft_tokens > 0, (
"Suffix decoder produced no draft tokens on the first run: "
f"accepted={first_accept_tokens}, drafted={first_draft_tokens}"
)
first_accept_rate = first_accept_tokens / first_draft_tokens
# Take the diff since the stats are cumulative.
last_accept_tokens = num_accept[-1] - num_accept[-2]
last_draft_tokens = num_draft[-1] - num_draft[-2]
assert last_draft_tokens > 0, (
"Suffix decoder produced no draft tokens on the last run: "
f"accepted_delta={last_accept_tokens}, drafted_delta={last_draft_tokens}; "
f"cumulative_drafted={num_draft[-2:]}"
)
last_accept_rate = last_accept_tokens / last_draft_tokens
summary = (
f"first accepted/drafted={first_accept_tokens}/{first_draft_tokens} "
f"(rate={first_accept_rate:.3f}); last delta accepted/drafted="
f"{last_accept_tokens}/{last_draft_tokens} (rate={last_accept_rate:.3f})"
)
# Expect the acceptance length to improve.
assert first_accept_tokens < last_accept_tokens
assert first_accept_tokens < last_accept_tokens, (
f"Expected accepted tokens to increase after cache warmup; {summary}"
)
# Expect the acceptance rate to improve.
assert first_accept_rate < last_accept_rate
assert first_accept_rate < last_accept_rate, (
f"Expected acceptance rate to increase after cache warmup; {summary}"
)
# Heuristic: expect at least 80.0% acceptance rate at the end.
assert last_accept_rate > 0.80
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
assert last_accept_rate > 0.80, f"Expected final acceptance rate > 0.80; {summary}"
tests/v1/e2e/spec_decode/speculators/test_speculators.py
+16 -13
View File
@@ -7,14 +7,20 @@ from tests.utils import single_gpu_only
from vllm import SamplingParams
from vllm.config import CompilationConfig
from ..utils import evaluate_llm_for_gsm8k, get_test_prompts
from ..utils import (
assert_request_outputs_match,
evaluate_llm_for_gsm8k,
get_test_prompts,
)
@pytest.mark.parametrize(
["model_path", "expected_accuracy_threshold"],
[
("RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3", 0.7), # ref: 75%-80%
("RedHatAI/Qwen3-8B-speculator.eagle3", 0.8), # ref: 87%-92%
# Measured reference: 75%-80%.
("RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3", 0.72),
# Measured reference: 87%-92%.
("RedHatAI/Qwen3-8B-speculator.eagle3", 0.84),
],
ids=["llama3_eagle3_speculator", "qwen3_eagle3_speculator"],
)
@@ -49,6 +55,7 @@ def test_speculators_model_integration(
# First run: Direct speculator model (simplified integration)
with vllm_runner(
model_path,
block_size=None,
trust_remote_code=False,
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
@@ -82,6 +89,7 @@ def test_speculators_model_integration(
# Second run: Reference without speculative decoding
with vllm_runner(
verifier_model,
block_size=None,
trust_remote_code=False,
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
@@ -90,15 +98,10 @@ def test_speculators_model_integration(
) as ref_runner:
ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config)
# Compare outputs
matches = sum(
1
for ref, spec in zip(ref_outputs, spec_outputs)
if ref.outputs[0].text == spec.outputs[0].text
)
# Heuristic: expect at least 66% of prompts to match exactly
assert matches >= int(0.66 * len(ref_outputs)), (
f"Only {matches}/{len(ref_outputs)} outputs matched. "
f"Expected at least {int(0.66 * len(ref_outputs))} matches."
assert_request_outputs_match(
ref_outputs,
spec_outputs,
required_matches=int(0.66 * len(ref_outputs)),
context=f"speculator={model_path}, verifier={verifier_model}",
)
tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
+46 -20
View File
@@ -11,8 +11,10 @@ import torch
from tests.utils import multi_gpu_marks
from vllm import LLM, SamplingParams
from vllm.config import CompilationConfig
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.platforms import current_platform
from vllm.sampling_params import RequestOutputKind
from vllm.v1.engine.async_llm import AsyncLLM
from vllm.v1.metrics.reader import Metric
@@ -59,7 +61,7 @@ INLINE_CONFIGS = [
pp_size=2,
skip_reason=(
"DeepSeek MTP pipeline-parallel support is in flight upstream; "
"see https://github.com/vllm-project/vllm/pull/38104"
"see https://github.com/vllm-project/vllm/pull/44698"
),
),
]
@@ -109,7 +111,11 @@ def _generate_token_ids_inline(llm: LLM) -> tuple[int, ...]:
SamplingParams(temperature=0.0, max_tokens=MAX_TOKENS, ignore_eos=True),
)
assert outputs and outputs[0].outputs, "expected one completion"
return tuple(outputs[0].outputs[0].token_ids)
token_ids = tuple(outputs[0].outputs[0].token_ids)
assert len(token_ids) == MAX_TOKENS, (
f"expected {MAX_TOKENS} generated tokens, got {len(token_ids)}: {token_ids}"
)
return token_ids
def _spec_decode_num_drafts(metrics: list[Metric]) -> int:
@@ -121,6 +127,16 @@ def _spec_decode_num_drafts(metrics: list[Metric]) -> int:
return int(counter.value)
def _enable_batch_invariance(monkeypatch: pytest.MonkeyPatch) -> None:
# This test exercises DeepSeek MLA prefill. ROCm's upstream FlashAttention
# varlen API lacks the num_splits argument used by batch invariance, so
# explicitly override any ambient setting on that path.
if current_platform.is_cuda():
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
else:
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0")
@pytest.mark.parametrize(
"config",
[
@@ -131,6 +147,7 @@ def _spec_decode_num_drafts(metrics: list[Metric]) -> int:
def test_deepseek_mtp_load_inline(
monkeypatch: pytest.MonkeyPatch,
config: InlineConfig,
vllm_runner,
):
"""MTP loads and drafts under TP/EP/EPLB; spec output matches no-spec greedy."""
if config.skip_reason is not None:
@@ -138,24 +155,30 @@ def test_deepseek_mtp_load_inline(
# Reduces run-to-run nondeterminism so spec / no-spec stay close;
# see tests/v1/distributed/test_eagle_dp.py for the same pattern.
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
_enable_batch_invariance(monkeypatch)
spec_llm = LLM(**_inline_kwargs(config, with_spec=True))
try:
spec_tokens = _generate_token_ids_inline(spec_llm)
n_drafts = _spec_decode_num_drafts(spec_llm.get_metrics())
finally:
del spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
spec_kwargs = _inline_kwargs(config, with_spec=True)
spec_model = spec_kwargs.pop("model")
with vllm_runner(
spec_model,
block_size=None,
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
**spec_kwargs,
) as spec_runner:
spec_tokens = _generate_token_ids_inline(spec_runner.llm)
n_drafts = _spec_decode_num_drafts(spec_runner.llm.get_metrics())
no_spec_llm = LLM(**_inline_kwargs(config, with_spec=False))
try:
no_spec_tokens = _generate_token_ids_inline(no_spec_llm)
finally:
del no_spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
no_spec_kwargs = _inline_kwargs(config, with_spec=False)
no_spec_model = no_spec_kwargs.pop("model")
with vllm_runner(
no_spec_model,
block_size=None,
enable_chunked_prefill=None,
compilation_config=CompilationConfig(),
**no_spec_kwargs,
) as no_spec_runner:
no_spec_tokens = _generate_token_ids_inline(no_spec_runner.llm)
# Non-vacuity: a silently-broken MTP drafter that falls back to
# verifier-only decoding would still produce matching output below.
@@ -185,7 +208,7 @@ async def test_deepseek_mtp_load_dp(monkeypatch: pytest.MonkeyPatch, dp_size: in
if torch.accelerator.device_count() < dp_size:
pytest.skip(f"dp{dp_size} requires at least {dp_size} GPUs")
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
_enable_batch_invariance(monkeypatch)
base_args = AsyncEngineArgs(
model=DEEPSEEK_MTP_MAIN_RANDOM,
@@ -225,7 +248,10 @@ async def test_deepseek_mtp_load_dp(monkeypatch: pytest.MonkeyPatch, dp_size: in
sampling_params=sampling_params,
):
token_ids = tuple(out.outputs[0].token_ids)
assert len(token_ids) == MAX_TOKENS
assert len(token_ids) == MAX_TOKENS, (
f"{request_id}: expected {MAX_TOKENS} generated tokens, "
f"got {len(token_ids)}: {token_ids}"
)
return token_ids
raise AssertionError("AsyncLLM produced no output")
finally:
tests/v1/e2e/spec_decode/utils.py
+91 -19
View File
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import random
from collections.abc import Iterable
from collections.abc import Iterable, Sequence
from typing import Any
import pytest
@@ -12,11 +12,13 @@ from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_of
from vllm import LLM, SamplingParams
from vllm.assets.base import VLLM_S3_BUCKET_URL
from vllm.assets.image import VLM_IMAGES_DIR
from vllm.outputs import RequestOutput
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.metrics.reader import Metric
def _skip_if_insufficient_gpus_for_tp(tp_size: int):
"""Skip test if available GPUs < tp_size on ROCm."""
"""Skip if fewer than ``tp_size`` accelerators are available."""
available_gpus = torch.accelerator.device_count()
if available_gpus < tp_size:
pytest.skip(
@@ -43,7 +45,7 @@ def get_test_prompts(mm_enabled: bool, num_prompts: int = 100) -> list[Messages]
# Generate a mixed batch of prompts, some of which can be easily
# predicted by n-gram matching and some which likely cannot.
random.seed(0)
set_random_seed(0)
for _ in range(num_repeat_prompts):
word_choices = ["test", "temp", "hello", "where"]
word = random.choice(word_choices)
@@ -120,20 +122,88 @@ def evaluate_llm_for_gsm8k(llm: LLM, expected_accuracy_threshold: float = 0.70)
)
def assert_request_outputs_match(
ref_outputs: Sequence[RequestOutput],
spec_outputs: Sequence[RequestOutput],
*,
required_matches: int,
context: str,
max_mismatches: int = 3,
) -> None:
"""Assert a thresholded exact-text match with bounded failure diagnostics."""
assert ref_outputs, f"{context}: no reference outputs"
assert len(ref_outputs) == len(spec_outputs), (
f"{context}: output count differs: "
f"reference={len(ref_outputs)}, speculative={len(spec_outputs)}"
)
assert 0 <= required_matches <= len(ref_outputs), (
f"{context}: invalid required_matches={required_matches} for "
f"{len(ref_outputs)} outputs"
)
mismatches: list[str] = []
matches = 0
for index, (ref_output, spec_output) in enumerate(zip(ref_outputs, spec_outputs)):
assert ref_output.outputs, (
f"{context}: reference output {index} has no candidate"
)
assert spec_output.outputs, (
f"{context}: speculative output {index} has no candidate"
)
ref_candidate = ref_output.outputs[0]
spec_candidate = spec_output.outputs[0]
if ref_candidate.text == spec_candidate.text:
matches += 1
elif len(mismatches) < max_mismatches:
mismatches.append(
f"[{index}] ref_text={ref_candidate.text[:240]!r}, "
f"spec_text={spec_candidate.text[:240]!r}\n"
f" ref_token_ids={list(ref_candidate.token_ids)[:64]}\n"
f" spec_token_ids={list(spec_candidate.token_ids)[:64]}"
)
print(
f"{context}: exact text matches={matches}/{len(ref_outputs)} "
f"(required={required_matches})"
)
mismatch_summary = "\n".join(mismatches) or "no mismatches captured"
assert matches >= required_matches, (
f"{context}: only {matches}/{len(ref_outputs)} outputs matched; "
f"required at least {required_matches}. First mismatches:\n"
f"{mismatch_summary}"
)
def get_spec_decode_metric_value(metrics: Sequence[Metric], metric_name: str) -> float:
"""Get a spec-decode metric with an actionable error when stats are absent."""
name2metric = {metric.name: metric for metric in metrics}
metric = name2metric.get(metric_name)
assert metric is not None, (
f"Missing metric {metric_name!r}. Ensure disable_log_stats=False. "
"Available spec-decode metrics: "
f"{sorted(name for name in name2metric if 'spec_decode' in name) or ['<none>']}"
)
return float(metric.value)
def compute_acceptance_rate(
metrics: list[Metric], prev_metrics: list[Metric] | None = None
) -> float:
name2metric = {metric.name: metric for metric in metrics}
n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value
n_draft_toks = get_spec_decode_metric_value(
metrics, "vllm:spec_decode_num_draft_tokens"
)
if n_draft_toks == 0:
return float("nan")
n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value
n_accepted_toks = get_spec_decode_metric_value(
metrics, "vllm:spec_decode_num_accepted_tokens"
)
if prev_metrics is not None:
prev_name2metric = {metric.name: metric for metric in prev_metrics}
n_draft_toks -= prev_name2metric["vllm:spec_decode_num_draft_tokens"].value
n_accepted_toks -= prev_name2metric[
"vllm:spec_decode_num_accepted_tokens"
].value
n_draft_toks -= get_spec_decode_metric_value(
prev_metrics, "vllm:spec_decode_num_draft_tokens"
)
n_accepted_toks -= get_spec_decode_metric_value(
prev_metrics, "vllm:spec_decode_num_accepted_tokens"
)
if n_draft_toks <= 0:
return float("nan")
return n_accepted_toks / n_draft_toks
@@ -142,17 +212,19 @@ def compute_acceptance_rate(
def compute_acceptance_len(
metrics: list[Metric], prev_metrics: list[Metric] | None = None
) -> float:
name2metric = {metric.name: metric for metric in metrics}
n_drafts = name2metric["vllm:spec_decode_num_drafts"].value
n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value
n_drafts = get_spec_decode_metric_value(metrics, "vllm:spec_decode_num_drafts")
n_accepted_toks = get_spec_decode_metric_value(
metrics, "vllm:spec_decode_num_accepted_tokens"
)
if n_drafts == 0:
return 1
if prev_metrics is not None:
prev_name2metric = {metric.name: metric for metric in prev_metrics}
n_drafts -= prev_name2metric["vllm:spec_decode_num_drafts"].value
n_accepted_toks -= prev_name2metric[
"vllm:spec_decode_num_accepted_tokens"
].value
n_drafts -= get_spec_decode_metric_value(
prev_metrics, "vllm:spec_decode_num_drafts"
)
n_accepted_toks -= get_spec_decode_metric_value(
prev_metrics, "vllm:spec_decode_num_accepted_tokens"
)
if n_drafts <= 0:
return 1
return 1 + (n_accepted_toks / n_drafts)
vllm/config/speculative.py
+1 -1
View File
@@ -121,7 +121,7 @@ class SpeculativeConfig:
attention_backend: AttentionBackendEnum | None = None
"""Attention backend to use for the draft model. When `None`, the backend is
automatically selected. Useful when the drafter requires a different attention
backend (e.g. DFlash needs a non-causal-capable backend like FLASH_ATTN)."""
backend (e.g. DFlash needs a backend that supports non-causal attention)."""
kv_cache_dtype: CacheDType | None = None
"""KV cache dtype for the draft model. When `None`, the draft inherits the
target model's `--kv-cache-dtype`."""