@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1a1727330a

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

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>

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