[CI] Organize speculative decoding E2E tests by coverage (#50330)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Michael Goin committed
Jul 31, 2026 at 21:05 UTC
fcdc7c2e9c9b75c923751ea7d2dd79ef66572c81
39 files changed
+1941
-1863
.buildkite/intel_jobs/engine_intel.yaml
+4
-2
@@ -95,9 +95,11 @@ steps:
95
- vllm/triton_utils/
96
- vllm/utils/
97
- vllm/v1/
98
- - tests/v1/e2e/spec_decode
98
+ - tests/v1/e2e/spec_decode/
99
commands:
100
- >-
101
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
102
'cd tests &&
103
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"'
103
+ pytest -v -s
104
+ v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism
105
+ v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism'
.buildkite/intel_jobs/model_runner_v2_intel.yaml
-1
@@ -107,7 +107,6 @@ steps:
107
- vllm/v1/worker/gpu_worker.py
108
- tests/v1/spec_decode/test_max_len.py
109
- tests/v1/spec_decode/test_rejection_sampler_utils.py
110
- - tests/v1/e2e/spec_decode/test_spec_decode.py
110
commands:
111
- >-
112
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
.buildkite/test-amd.yaml
+16
-23
@@ -330,7 +330,7 @@ steps:
330
- tests/v1/e2e/spec_decode/
331
- vllm/platforms/rocm.py
332
commands:
333
- - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
333
+ - pytest -v -s v1/e2e/spec_decode/draft_model/
334
335
- label: Spec Decode Speculators + MTP # TBD
336
timeout_in_minutes: 180
@@ -348,7 +348,8 @@ steps:
348
- tests/v1/e2e/spec_decode/
349
- vllm/platforms/rocm.py
350
commands:
351
- - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
351
+ - pytest -v -s v1/e2e/spec_decode/speculators/
352
+ - pytest -v -s v1/e2e/spec_decode/mtp/
353
354
- label: V1 others (CPU) # TBD
355
timeout_in_minutes: 180
@@ -1836,7 +1837,7 @@ steps:
1837
- tests/v1/spec_decode/test_max_len.py
1838
- tests/v1/spec_decode/test_rejection_sampler_utils.py
1839
- tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
1839
- - tests/v1/e2e/spec_decode/test_spec_decode.py
1840
+ - tests/v1/e2e/spec_decode/
1841
- vllm/platforms/rocm.py
1842
commands:
1843
- set -x
@@ -1844,7 +1845,9 @@ steps:
1845
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
1846
- pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
1847
- pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
1847
- - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
1848
+ - pytest -v -s v1/e2e/spec_decode/eagle/
1849
+ - pytest -v -s v1/e2e/spec_decode/speculators/
1850
+ - pytest -v -s v1/e2e/spec_decode/mtp/
1851
1852
#------------------------------------------------------ mi300 · models / basic -------------------------------------------------------#
1853
@@ -2551,7 +2554,7 @@ steps:
2554
- tests/v1/e2e/spec_decode/
2555
- vllm/platforms/rocm.py
2556
commands:
2554
- - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
2557
+ - pytest -v -s v1/e2e/spec_decode/draft_model/
2558
2559
- label: Spec Decode Eagle # TBD
2560
timeout_in_minutes: 180
@@ -2569,7 +2572,7 @@ steps:
2572
- tests/v1/e2e/spec_decode/
2573
- vllm/platforms/rocm.py
2574
commands:
2572
- - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
2575
+ - pytest -v -s v1/e2e/spec_decode/eagle/
2576
2577
- label: Spec Decode Ngram + Suffix # TBD
2578
timeout_in_minutes: 180
@@ -2587,7 +2590,7 @@ steps:
2590
- tests/v1/e2e/spec_decode/
2591
- vllm/platforms/rocm.py
2592
commands:
2590
- - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
2593
+ - pytest -v -s v1/e2e/spec_decode/ngram_suffix/
2594
2595
- label: Spec Decode Speculators + MTP # TBD
2596
timeout_in_minutes: 180
@@ -2606,7 +2609,8 @@ steps:
2609
- tests/v1/e2e/spec_decode/
2610
- vllm/platforms/rocm.py
2611
commands:
2609
- - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
2612
+ - pytest -v -s v1/e2e/spec_decode/speculators/
2613
+ - pytest -v -s v1/e2e/spec_decode/mtp/
2614
2615
- label: Speculators Correctness # TBD
2616
timeout_in_minutes: 180
@@ -2807,7 +2811,10 @@ steps:
2811
- vllm/
2812
- tests/v1/e2e
2813
commands:
2810
- - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
2814
+ - >-
2815
+ pytest -v -s
2816
+ v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism
2817
+ v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism
2818
2819
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD
2820
timeout_in_minutes: 180
@@ -2969,20 +2976,6 @@ steps:
2976
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
2977
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
2978
2972
-- label: V1 e2e (4 GPUs) # TBD
2973
- timeout_in_minutes: 180
2974
- mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
2975
- dind: false
2976
- agent_pool: mi300_4
2977
- num_gpus: 4
2978
- optional: true
2979
- working_dir: "/vllm-workspace/tests"
2980
- source_file_dependencies:
2981
- - vllm/
2982
- - tests/v1/e2e
2983
- commands:
2984
- - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy"
2985
-
2979
- label: V1 e2e (4xH100-4xMI300) # TBD
2980
timeout_in_minutes: 180
2981
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
.buildkite/test_areas/engine.yaml
+5
-30
@@ -112,10 +112,13 @@ steps:
112
- vllm/triton_utils/
113
- vllm/utils/
114
- vllm/v1/
115
- - tests/v1/e2e/spec_decode
115
+ - tests/v1/e2e/spec_decode/
116
commands:
117
# Only run tests that need exactly 2 GPUs
118
- - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
118
+ - >-
119
+ pytest -v -s
120
+ v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_tensor_parallelism
121
+ v1/e2e/spec_decode/draft_model/test_draft_model.py::test_draft_model_engine_args_tensor_parallelism
122
mirror:
123
amd:
124
dind: false
@@ -124,34 +127,6 @@ steps:
127
depends_on:
128
- image-build-amd
129
127
-- label: V1 e2e (4 GPUs)
128
- key: v1-e2e-4-gpus
129
- timeout_in_minutes: 20 # TODO: Fix timeout after we have more confidence in the test stability
130
- optional: true
131
- num_devices: 4
132
- source_file_dependencies:
133
- - vllm/compilation/
134
- - vllm/config/
135
- - vllm/distributed/
136
- - vllm/engine/
137
- - vllm/envs.py
138
- - vllm/forward_context.py
139
- - vllm/inputs/
140
- - vllm/logger.py
141
- - vllm/logging_utils/
142
- - vllm/model_executor/
143
- - vllm/multimodal/
144
- - vllm/platforms/
145
- - vllm/sampling_params.py
146
- - vllm/transformers_utils/
147
- - vllm/triton_utils/
148
- - vllm/utils/
149
- - vllm/v1/
150
- - tests/v1/e2e/spec_decode
151
- commands:
152
- # Only run tests that need 4 GPUs
153
- - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy"
154
-
130
- label: V1 e2e (4xH100)
131
key: v1-e2e-4xh100
132
timeout_in_minutes: 35
.buildkite/test_areas/misc.yaml
-27
@@ -2,33 +2,6 @@ group: Miscellaneous
2
depends_on:
3
- image-build
4
steps:
5
-- label: V1 Spec Decode
6
- device: h200_35gb
7
- key: v1-spec-decode
8
- timeout_in_minutes: 40
9
- source_file_dependencies:
10
- - vllm/config/
11
- - vllm/distributed/
12
- - vllm/inputs/
13
- - vllm/model_executor/
14
- - vllm/platforms/
15
- - vllm/sampling_params.py
16
- - vllm/transformers_utils/
17
- - vllm/utils/
18
- - vllm/v1/
19
- - tests/v1/spec_decode
20
- commands:
21
- - export VLLM_WORKER_MULTIPROC_METHOD=spawn
22
- # TODO: create another `optional` test group for slow tests
23
- - pytest -v -s -m 'not slow_test' v1/spec_decode
24
- mirror:
25
- amd:
26
- dind: false
27
- device: mi300_1
28
- timeout_in_minutes: 50
29
- depends_on:
30
- - image-build-amd
31
-
5
- label: V1 Sample + Logits
6
key: v1-sample-logits
7
timeout_in_minutes: 83
.buildkite/test_areas/model_runner_v2.yaml
+4
-3
@@ -103,18 +103,19 @@ steps:
103
- label: Model Runner V2 Spec Decode
104
device: h200_35gb
105
key: model-runner-v2-spec-decode
106
- timeout_in_minutes: 30
106
+ timeout_in_minutes: 50
107
working_dir: "/vllm-workspace/tests"
108
source_file_dependencies:
109
- vllm/v1/worker/gpu/
110
- vllm/v1/worker/gpu_worker.py
111
- tests/v1/spec_decode/test_max_len.py
112
- tests/v1/spec_decode/test_rejection_sampler_utils.py
113
- - tests/v1/e2e/spec_decode/test_spec_decode.py
113
+ - tests/v1/e2e/spec_decode/
114
commands:
115
- set -x
116
- export VLLM_USE_V2_MODEL_RUNNER=1
117
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
118
- pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
119
- pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
120
- - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
120
+ - pytest -v -s v1/e2e/spec_decode/eagle/
121
+ - pytest -v -s v1/e2e/spec_decode/speculators/
.buildkite/test_areas/spec_decode.yaml
+78
-32
@@ -2,16 +2,43 @@ group: Spec Decode
2
depends_on:
3
- image-build
4
steps:
5
+- label: V1 Spec Decode
6
+ device: h200_35gb
7
+ key: v1-spec-decode
8
+ timeout_in_minutes: 40
9
+ source_file_dependencies:
10
+ - vllm/config/
11
+ - vllm/distributed/
12
+ - vllm/inputs/
13
+ - vllm/model_executor/
14
+ - vllm/platforms/
15
+ - vllm/sampling_params.py
16
+ - vllm/transformers_utils/
17
+ - vllm/utils/
18
+ - vllm/v1/
19
+ - tests/v1/spec_decode
20
+ commands:
21
+ - export VLLM_WORKER_MULTIPROC_METHOD=spawn
22
+ # TODO: create another `optional` test group for slow tests
23
+ - pytest -v -s -m 'not slow_test' v1/spec_decode
24
+ mirror:
25
+ amd:
26
+ dind: false
27
+ device: mi300_1
28
+ timeout_in_minutes: 50
29
+ depends_on:
30
+ - image-build-amd
31
+
32
- label: Spec Decode Eagle
33
key: spec-decode-eagle
34
timeout_in_minutes: 25
8
- device: h200_18gb
35
+ device: h200_35gb
36
source_file_dependencies:
37
- vllm/v1/spec_decode/
38
- vllm/v1/worker/gpu/spec_decode/
39
- tests/v1/e2e/spec_decode/
40
commands:
14
- - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
41
+ - pytest -v -s v1/e2e/spec_decode/eagle/
42
mirror:
43
amd:
44
dind: false
@@ -28,7 +55,7 @@ steps:
55
- tests/v1/e2e/spec_decode/
56
- vllm/platforms/rocm.py
57
31
-- label: Spec Decode Eagle Nightly B200
58
+- label: Spec Decode Eagle Nightly (B200)
59
key: spec-decode-eagle-nightly-b200
60
timeout_in_minutes: 25
61
device: b200-k8s
@@ -38,12 +65,12 @@ steps:
65
- vllm/v1/worker/gpu/spec_decode/
66
- tests/v1/e2e/spec_decode/
67
commands:
41
- - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
68
+ - pytest -v -s v1/e2e/spec_decode/eagle/
69
70
- label: Spec Decode Speculators + MTP
71
key: spec-decode-speculators-mtp
45
- timeout_in_minutes: 20
46
- device: h200_18gb
72
+ timeout_in_minutes: 50
73
+ device: h200_35gb
74
source_file_dependencies:
75
- vllm/v1/spec_decode/
76
- vllm/v1/worker/gpu/spec_decode/
@@ -51,7 +78,8 @@ steps:
78
- vllm/transformers_utils/configs/speculators/
79
- tests/v1/e2e/spec_decode/
80
commands:
54
- - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
81
+ - pytest -v -s v1/e2e/spec_decode/speculators/
82
+ - pytest -v -s v1/e2e/spec_decode/mtp/
83
mirror:
84
amd:
85
dind: false
@@ -69,7 +97,7 @@ steps:
97
- tests/v1/e2e/spec_decode/
98
- vllm/platforms/rocm.py
99
72
-- label: Spec Decode Speculators + MTP Nightly B200
100
+- label: Spec Decode Speculators + MTP Nightly (B200)
101
key: spec-decode-speculators-mtp-nightly-b200
102
timeout_in_minutes: 30
103
device: b200-k8s
@@ -80,19 +108,20 @@ steps:
108
- vllm/transformers_utils/configs/speculators/
109
- tests/v1/e2e/spec_decode/
110
commands:
83
- - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
84
-
111
+ - pytest -v -s v1/e2e/spec_decode/speculators/
112
+ - pytest -v -s v1/e2e/spec_decode/mtp/
113
+
114
- label: Spec Decode Ngram + Suffix
115
key: spec-decode-ngram-suffix
116
timeout_in_minutes: 20
88
- device: h200_18gb
117
+ device: h200_35gb
118
source_file_dependencies:
119
- vllm/v1/spec_decode/
120
- vllm/v1/worker/gpu/spec_decode/
121
- tests/v1/e2e/spec_decode/
122
- tests/spec_decode/
123
commands:
95
- - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
124
+ - pytest -v -s v1/e2e/spec_decode/ngram_suffix/
125
- python3 spec_decode/test_custom_proposer.py
126
mirror:
127
amd:
@@ -119,7 +148,7 @@ steps:
148
- vllm/v1/worker/gpu/spec_decode/
149
- tests/v1/e2e/spec_decode/
150
commands:
122
- - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
151
+ - pytest -v -s v1/e2e/spec_decode/draft_model/
152
mirror:
153
amd:
154
dind: false
@@ -136,7 +165,7 @@ steps:
165
- tests/v1/e2e/spec_decode/
166
- vllm/platforms/rocm.py
167
139
-- label: Spec Decode Draft Model Nightly B200
168
+- label: Spec Decode Draft Model Nightly (B200)
169
key: spec-decode-draft-model-nightly-b200
170
timeout_in_minutes: 40
171
device: b200-k8s
@@ -146,9 +175,9 @@ steps:
175
- vllm/v1/worker/gpu/spec_decode/
176
- tests/v1/e2e/spec_decode/
177
commands:
149
- - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
178
+ - pytest -v -s v1/e2e/spec_decode/draft_model/
179
151
-- label: Speculators Correctness
180
+- label: Speculators Correctness Nightly (H100)
181
key: speculators-correctness
182
timeout_in_minutes: 30
183
device: h100
@@ -156,25 +185,13 @@ steps:
185
num_devices: 1
186
source_file_dependencies:
187
- vllm/v1/spec_decode/
188
+ - vllm/v1/worker/gpu/spec_decode/dflash/
189
- vllm/model_executor/models/qwen3_dflash.py
190
- tests/v1/spec_decode/test_speculators_correctness.py
191
commands:
192
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
193
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
194
165
-- label: Spec Decode MTP hybrid (B200)
166
- timeout_in_minutes: 20
167
- device: b200-k8s
168
- optional: true
169
- source_file_dependencies:
170
- - vllm/v1/spec_decode/
171
- - vllm/v1/worker/gpu/spec_decode/
172
- - vllm/model_executor/models/qwen3_5.py
173
- - vllm/model_executor/models/qwen3_5_mtp.py
174
- - tests/v1/e2e/spec_decode/
175
- commands:
176
- - pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
177
-
195
- label: Spec Decode DeepSeek MTP Parallel Load (B200)
196
key: spec-decode-deepseek-mtp-parallel-load-b200
197
timeout_in_minutes: 30
@@ -191,14 +208,43 @@ steps:
208
commands:
209
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
210
194
-- label: Spec Decode Acceptance Rates Nightly
195
- key: spec-decode-acceptance-rates-nightly
211
+- label: Spec Decode AL DFlash Nightly
212
+ key: spec-decode-dflash-nightly
213
+ timeout_in_minutes: 90
214
+ device: h200_35gb
215
+ optional: true
216
+ source_file_dependencies:
217
+ - vllm/v1/spec_decode/
218
+ - vllm/v1/worker/gpu/spec_decode/
219
+ - vllm/model_executor/models/qwen3_dflash.py
220
+ - vllm/model_executor/models/laguna_dflash.py
221
+ - tests/v1/e2e/spec_decode/
222
+ commands:
223
+ - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/
224
+
225
+- label: Spec Decode AL DSpark Nightly
226
+ key: spec-decode-dspark-nightly
227
+ timeout_in_minutes: 60
228
+ device: h200_35gb
229
+ optional: true
230
+ source_file_dependencies:
231
+ - vllm/v1/spec_decode/
232
+ - vllm/v1/worker/gpu/spec_decode/
233
+ - vllm/model_executor/models/qwen3_dspark.py
234
+ - vllm/model_executor/models/gemma4_dspark.py
235
+ - tests/v1/e2e/spec_decode/
236
+ commands:
237
+ - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/
238
+
239
+- label: Spec Decode AL MTP + Other Acceptance Nightly
240
+ key: spec-decode-mtp-other-acceptance-nightly
241
timeout_in_minutes: 60
242
device: h200_35gb
243
optional: true
244
source_file_dependencies:
245
- vllm/v1/spec_decode/
246
- vllm/v1/worker/gpu/spec_decode/
247
+ - tests/v1/e2e/spec_decode/conftest.py
248
- tests/v1/e2e/spec_decode/
249
commands:
204
- - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "acceptance_rates"
250
+ - pytest -v -s v1/e2e/spec_decode/acceptance_rates/mtp_other/
tests/models/registry.py
+2
-2
@@ -1449,8 +1449,8 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1449
max_num_seqs=32,
1450
),
1451
"DFlashLagunaForCausalLM": _HfExamplesInfo(
1452
- "poolside/Laguna-XS-2.1",
1453
- speculative_model="poolside/Laguna-XS-2.1-DFlash",
1452
+ "poolside/Laguna-XS-2.1-NVFP4",
1453
+ speculative_model="poolside/Laguna-XS-2.1-DFlash-NVFP4",
1454
use_original_num_layers=True,
1455
max_model_len=8192, # Reduce max len to ensure test runs in low-VRAM CI env
1456
max_num_seqs=32,
tests/utils.py
+4
-1
@@ -1983,7 +1983,10 @@ def large_gpu_mark(min_gb: int) -> pytest.MarkDecorator:
1983
1984
return pytest.mark.skipif(
1985
memory_gb < min_gb,
1986
- reason=f"Need at least {min_gb}GB GPU memory to run the test.",
1986
+ reason=(
1987
+ f"Need at least {min_gb}GB GPU memory to run the test "
1988
+ f"(found {memory_gb}GB)."
1989
+ ),
1990
)
1991
1992
tests/v1/e2e/general/test_kv_sharing_fast_prefill.py
+1
-1
@@ -18,7 +18,7 @@ SEED = 42
18
@pytest.fixture
19
def test_prompts():
20
"""
21
- Adapted from tests/v1/e2e/spec_decode/test_spec_decode.py
21
+ Adapted from tests/v1/e2e/spec_decode/utils.py
22
"""
23
prompt_types = ["repeat", "sentence"]
24
# Setting higher num prompts increases the chance of numerics mismatch
tests/v1/e2e/spec_decode/acceptance_rates/__init__.py
tests/v1/e2e/spec_decode/acceptance_rates/dflash/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py
new
+152
@@ -0,0 +1,152 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import dataclass
5
+
6
+import pytest
7
+import torch
8
+
9
+from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline
10
+from tests.utils import single_gpu_only
11
+from vllm import LLM
12
+from vllm.distributed import cleanup_dist_env_and_memory
13
+
14
+from ...utils import compute_acceptance_len
15
+from ..utils import run_acceptance_length_eval
16
+
17
+
18
+@dataclass(frozen=True)
19
+class DFlashCorrectnessConfig:
20
+ model: str
21
+ draft_model: str
22
+ expected_accuracy: float
23
+ expected_acceptance_len: float
24
+ num_speculative_tokens: int = 16
25
+ max_model_len: int = 4096
26
+ max_num_seqs: int = 128
27
+ num_questions: int = 1319
28
+ use_chat_completions: bool = False
29
+ enforce_eager: bool = False
30
+ disable_flashinfer_sampler: bool = False
31
+
32
+
33
+QWEN3_DFLASH = DFlashCorrectnessConfig(
34
+ model="Qwen/Qwen3-8B",
35
+ draft_model="z-lab/Qwen3-8B-DFlash-b16",
36
+ expected_accuracy=0.8,
37
+ expected_acceptance_len=3.5,
38
+)
39
+
40
+LAGUNA_DFLASH_NVFP4 = DFlashCorrectnessConfig(
41
+ model="poolside/Laguna-XS-2.1-NVFP4",
42
+ draft_model="poolside/Laguna-XS-2.1-DFlash-NVFP4",
43
+ expected_accuracy=0.7, # Standard GSM8K sanity floor.
44
+ expected_acceptance_len=3.55 * 0.9,
45
+ num_speculative_tokens=15,
46
+ max_model_len=8192,
47
+ max_num_seqs=32,
48
+ num_questions=200,
49
+ use_chat_completions=True,
50
+ enforce_eager=True,
51
+ disable_flashinfer_sampler=True,
52
+)
53
+
54
+
55
+@pytest.mark.parametrize("use_mrv2", [False, True])
56
+def test_dflash_reference_acceptance_lengths(
57
+ monkeypatch: pytest.MonkeyPatch,
58
+ use_mrv2: bool,
59
+):
60
+ run_acceptance_length_eval(
61
+ monkeypatch,
62
+ spec_config={
63
+ "model": QWEN3_DFLASH.model,
64
+ "trust_remote_code": True,
65
+ "speculative_config": {
66
+ "method": "dflash",
67
+ "model": QWEN3_DFLASH.draft_model,
68
+ "num_speculative_tokens": QWEN3_DFLASH.num_speculative_tokens,
69
+ "max_model_len": 32768,
70
+ },
71
+ "max_model_len": 32768,
72
+ "max_num_seqs": 128,
73
+ "gpu_memory_utilization": 0.85,
74
+ "enforce_eager": False,
75
+ "disable_log_stats": False,
76
+ },
77
+ # Table 1 in https://arxiv.org/pdf/2602.06036.
78
+ expected_acceptance_lengths={
79
+ "mt-bench": 4.24,
80
+ "humaneval": 6.50,
81
+ "gsm8k": 6.54 * 0.975,
82
+ },
83
+ chat_template_kwargs={"enable_thinking": False},
84
+ use_mrv2=use_mrv2,
85
+ )
86
+
87
+
88
+@single_gpu_only
89
+@pytest.mark.parametrize(
90
+ ("config", "use_mrv2"),
91
+ [
92
+ pytest.param(
93
+ QWEN3_DFLASH,
94
+ False,
95
+ id="qwen3-mrv1",
96
+ ),
97
+ pytest.param(
98
+ QWEN3_DFLASH,
99
+ True,
100
+ id="qwen3-mrv2",
101
+ ),
102
+ pytest.param(
103
+ LAGUNA_DFLASH_NVFP4,
104
+ True,
105
+ id="laguna-nvfp4-mrv2",
106
+ ),
107
+ ],
108
+)
109
+def test_dflash_correctness(
110
+ monkeypatch: pytest.MonkeyPatch,
111
+ config: DFlashCorrectnessConfig,
112
+ use_mrv2: bool,
113
+):
114
+ """Guard GSM8K accuracy and batched acceptance length for DFlash models."""
115
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0")
116
+ if config.disable_flashinfer_sampler:
117
+ monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
118
+
119
+ spec_llm = LLM(
120
+ model=config.model,
121
+ trust_remote_code=True,
122
+ speculative_config={
123
+ "method": "dflash",
124
+ "model": config.draft_model,
125
+ "num_speculative_tokens": config.num_speculative_tokens,
126
+ "max_model_len": config.max_model_len,
127
+ },
128
+ max_model_len=config.max_model_len,
129
+ max_num_seqs=config.max_num_seqs,
130
+ gpu_memory_utilization=0.85,
131
+ enforce_eager=config.enforce_eager,
132
+ 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()
tests/v1/e2e/spec_decode/acceptance_rates/dspark/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py
new
+129
@@ -0,0 +1,129 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline
8
+from vllm import LLM
9
+from vllm.distributed import cleanup_dist_env_and_memory
10
+
11
+from ...utils import compute_acceptance_len, compute_acceptance_rate
12
+
13
+
14
+def test_gemma4_dspark_correctness_and_acceptance_rate(
15
+ monkeypatch: pytest.MonkeyPatch,
16
+):
17
+ """
18
+ E2E test for Gemma4 DSpark speculative decoding: acceptance rate/length
19
+ regression coverage plus GSM8K correctness, at temperature=1.0 to exercise
20
+ the probabilistic draft-sampling/rejection-sampling path (not just greedy).
21
+
22
+ gemma-4-12B is instruct-tuned, so GSM8K is run through the chat template
23
+ (use_chat_completions=True; raw few-shot completion collapses to a few
24
+ percent). Reference: measured over 5 runs of 200 GSM8K questions at
25
+ temperature=1.0 (prefix caching disabled):
26
+ accuracy: min=0.900 max=0.955 mean=0.937
27
+ acceptance_rate: min=0.578 max=0.595 mean=0.588
28
+ acceptance_len: min=5.044 max=5.167 mean=5.116
29
+ Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling
30
+ """
31
+ monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
32
+
33
+ spec_llm = LLM(
34
+ model="RedHatAI/gemma-4-12B-it-NVFP4",
35
+ trust_remote_code=True,
36
+ speculative_config={
37
+ "method": "dspark",
38
+ "model": "deepseek-ai/dspark_gemma4_12b_block7",
39
+ "num_speculative_tokens": 7,
40
+ "draft_sample_method": "probabilistic",
41
+ },
42
+ max_model_len=4096,
43
+ max_num_seqs=32,
44
+ gpu_memory_utilization=0.85,
45
+ enforce_eager=True,
46
+ enable_prefix_caching=False,
47
+ disable_log_stats=False,
48
+ )
49
+ try:
50
+ results = evaluate_gsm8k_offline(
51
+ spec_llm, num_questions=200, temperature=1.0, use_chat_completions=True
52
+ )
53
+ gsm8k_accuracy = results["accuracy"]
54
+
55
+ metrics = spec_llm.get_metrics()
56
+ acceptance_rate = compute_acceptance_rate(metrics)
57
+ acceptance_len = compute_acceptance_len(metrics)
58
+ print(
59
+ f"Gemma4 DSpark acceptance_rate={acceptance_rate:.2f}, "
60
+ f"acceptance_len={acceptance_len:.2f}, "
61
+ f"gsm8k_accuracy={gsm8k_accuracy:.3f}"
62
+ )
63
+
64
+ assert acceptance_rate >= 0.58 * 0.9
65
+ assert acceptance_len >= 5.0 * 0.9
66
+ assert gsm8k_accuracy >= 0.92 * 0.9
67
+ finally:
68
+ del spec_llm
69
+ torch.accelerator.empty_cache()
70
+ cleanup_dist_env_and_memory()
71
+
72
+
73
+@pytest.fixture
74
+def dspark_config():
75
+ target_model = "Qwen/Qwen3-4B-FP8"
76
+ draft_model = "deepseek-ai/dspark_qwen3_4b_block7"
77
+
78
+ return dict(
79
+ model=target_model,
80
+ trust_remote_code=True,
81
+ speculative_config={
82
+ "method": "dspark",
83
+ "model": draft_model,
84
+ "num_speculative_tokens": 7,
85
+ "attention_backend": "FLASH_ATTN",
86
+ "draft_sample_method": "probabilistic",
87
+ },
88
+ max_model_len=4096,
89
+ disable_log_stats=False,
90
+ )
91
+
92
+
93
+def test_dspark_correctness_and_acceptance_rate(dspark_config):
94
+ """
95
+ E2E test for DSpark speculative decoding: acceptance rate/length
96
+ regression coverage plus GSM8K correctness, at temperature=1.0 to
97
+ exercise the probabilistic draft-sampling/rejection-sampling path
98
+ (not just greedy).
99
+
100
+ Uses Qwen/Qwen3-4B-FP8 as target with the dspark_qwen3_4b_block7 draft
101
+ model. Reference: measured over 12 runs of the full GSM8K set at
102
+ temperature=1.0 (prefix caching disabled to avoid cross-run reuse):
103
+ accuracy: min=0.782 max=0.814 mean=0.801
104
+ acceptance_rate: min=0.418 max=0.434 mean=0.428
105
+ acceptance_len: min=3.928 max=4.037 mean=3.994
106
+ Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling
107
+ """
108
+ spec_llm = LLM(**dspark_config)
109
+
110
+ results = evaluate_gsm8k_offline(spec_llm, temperature=1.0)
111
+ gsm8k_accuracy = results["accuracy"]
112
+
113
+ metrics = spec_llm.get_metrics()
114
+ acceptance_rate = compute_acceptance_rate(metrics)
115
+ acceptance_len = compute_acceptance_len(metrics)
116
+
117
+ print(
118
+ f"DSpark acceptance_rate={acceptance_rate:.2f}, "
119
+ f"acceptance_len={acceptance_len:.2f}, "
120
+ f"gsm8k_accuracy={gsm8k_accuracy:.3f}"
121
+ )
122
+
123
+ assert acceptance_rate >= 0.428 * 0.9
124
+ assert acceptance_len >= 3.994 * 0.9
125
+ assert gsm8k_accuracy >= 0.801 * 0.9
126
+
127
+ del spec_llm
128
+ torch.accelerator.empty_cache()
129
+ cleanup_dist_env_and_memory()
tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py
new
+50
@@ -0,0 +1,50 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+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
9
+
10
+from ...utils import compute_acceptance_rate
11
+
12
+
13
+def test_medusa_acceptance_rate(
14
+ sampling_config: SamplingParams,
15
+):
16
+ """Verify a trained Medusa checkpoint achieves nonzero acceptance rate.
17
+
18
+ Uses the canonical FasterDecoding vicuna-7b checkpoint to confirm the
19
+ speculation path actually accepts tokens — unlike test_medusa_correctness,
20
+ which uses a random head and only validates output correctness.
21
+ """
22
+ target_model = "lmsys/vicuna-7b-v1.3"
23
+ medusa_model = "FasterDecoding/medusa-vicuna-7b-v1.3"
24
+ prompts = _build_gsm8k_prompts(num_questions=10, num_shots=1)[0]
25
+
26
+ spec_llm = LLM(
27
+ model=target_model,
28
+ speculative_config={
29
+ "method": "medusa",
30
+ "model": medusa_model,
31
+ "num_speculative_tokens": 3,
32
+ },
33
+ max_model_len=1024,
34
+ enforce_eager=True,
35
+ 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()
43
+
44
+ min_acceptance_rate = 0.198
45
+ print(f"Medusa acceptance rate: {acceptance_rate:.4f} (min {min_acceptance_rate})")
46
+
47
+ # Regression guard at 90% of the measured baseline.
48
+ assert acceptance_rate >= min_acceptance_rate, (
49
+ f"Medusa acceptance rate {acceptance_rate:.4f} below min {min_acceptance_rate}"
50
+ )
tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py
new
+36
@@ -0,0 +1,36 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from ..utils import run_acceptance_length_eval
7
+
8
+
9
+@pytest.mark.parametrize("use_mrv2", [False, True])
10
+def test_gemma4_mtp_acceptance_lengths(
11
+ monkeypatch: pytest.MonkeyPatch,
12
+ use_mrv2: bool,
13
+):
14
+ run_acceptance_length_eval(
15
+ monkeypatch,
16
+ spec_config={
17
+ "model": "google/gemma-4-E4B-it",
18
+ "trust_remote_code": True,
19
+ "speculative_config": {
20
+ "method": "mtp",
21
+ "model": "google/gemma-4-E4B-it-assistant",
22
+ "num_speculative_tokens": 2,
23
+ "max_model_len": 32768,
24
+ },
25
+ "max_model_len": 32768,
26
+ "limit_mm_per_prompt": {"image": 0, "audio": 0},
27
+ "disable_log_stats": False,
28
+ },
29
+ expected_acceptance_lengths={
30
+ "mt-bench": 2.28,
31
+ "humaneval": 2.68,
32
+ "gsm8k": 2.67 * 0.975,
33
+ },
34
+ chat_template_kwargs={},
35
+ use_mrv2=use_mrv2,
36
+ )
tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py
new
+58
@@ -0,0 +1,58 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from tests.utils import single_gpu_only
7
+from vllm import LLM, SamplingParams
8
+from vllm.distributed import cleanup_dist_env_and_memory
9
+
10
+from ...utils import compute_acceptance_len, get_test_prompts
11
+
12
+
13
+@single_gpu_only
14
+def test_synthetic_acceptance_rate():
15
+ """Verify that synthetic rejection sampling produces an acceptance
16
+ length close to the requested mean acceptance length."""
17
+ num_spec_tokens = 3
18
+ expected_acceptance_len = 1.875
19
+ tolerance = 0.15
20
+
21
+ spec_llm = LLM(
22
+ model="meta-llama/Llama-3.2-1B-Instruct",
23
+ trust_remote_code=True,
24
+ speculative_config={
25
+ "method": "eagle3",
26
+ "model": "nm-testing/Llama3_2_1B_speculator.eagle3",
27
+ "num_speculative_tokens": num_spec_tokens,
28
+ "max_model_len": 2048,
29
+ "rejection_sample_method": "synthetic",
30
+ "synthetic_acceptance_length": expected_acceptance_len,
31
+ },
32
+ max_model_len=2048,
33
+ enforce_eager=True,
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()
tests/v1/e2e/spec_decode/acceptance_rates/utils.py
new
+105
@@ -0,0 +1,105 @@
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
7
+import torch
8
+from tqdm import tqdm
9
+
10
+from vllm import LLM, SamplingParams
11
+from vllm.distributed import cleanup_dist_env_and_memory
12
+
13
+from ..utils import compute_acceptance_len
14
+
15
+
16
+def load_and_process_dataset(data_name: str):
17
+ from datasets import load_dataset
18
+
19
+ if data_name == "gsm8k":
20
+ dataset = load_dataset("openai/gsm8k", "main", split="test")
21
+ prompt_fmt = (
22
+ "{question}\nPlease reason step by step,"
23
+ " and put your final answer within \\boxed{{}}."
24
+ )
25
+ dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
26
+ elif data_name == "mt-bench":
27
+ dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train")
28
+ dataset = dataset.map(lambda x: {"turns": x["prompt"]})
29
+ elif data_name == "humaneval":
30
+ dataset = load_dataset("openai/openai_humaneval", split="test")
31
+ prompt_fmt = (
32
+ "Write a solution to the following problem and make sure"
33
+ " that it passes the tests:\n```python\n{prompt}\n```"
34
+ )
35
+ dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
36
+
37
+ return dataset
38
+
39
+
40
+def run_acceptance_length_eval(
41
+ monkeypatch: pytest.MonkeyPatch,
42
+ spec_config: dict[str, Any],
43
+ expected_acceptance_lengths: dict[str, float],
44
+ chat_template_kwargs: dict[str, Any],
45
+ use_mrv2: bool,
46
+):
47
+ """
48
+ E2E acceptance-rate validation for speculative decoding.
49
+
50
+ Drives one or more datasets (keyed in ``expected_acceptance_lengths``)
51
+ through the spec decode engine and asserts the mean acceptance length
52
+ stays within tolerance of the reference figure for each dataset.
53
+ """
54
+ 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,
74
+ )
75
+
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,
82
+ )
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/conftest.py
new
+25
@@ -0,0 +1,25 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm import SamplingParams
8
+
9
+from .utils import greedy_sampling
10
+
11
+
12
+@pytest.fixture
13
+def sampling_config() -> SamplingParams:
14
+ return greedy_sampling()
15
+
16
+
17
+@pytest.fixture
18
+def model_name() -> str:
19
+ return "meta-llama/Llama-3.1-8B-Instruct"
20
+
21
+
22
+@pytest.fixture(autouse=True)
23
+def reset_torch_dynamo():
24
+ yield
25
+ torch._dynamo.reset()
tests/v1/e2e/spec_decode/draft_model/__init__.py
tests/v1/e2e/spec_decode/draft_model/test_async.py
renamed
tests/v1/e2e/spec_decode/draft_model/test_draft_model.py
new
+388
@@ -0,0 +1,388 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import dataclass
5
+
6
+import pytest
7
+import torch
8
+
9
+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
13
+from vllm.engine.arg_utils import EngineArgs
14
+
15
+from ..utils import (
16
+ Messages,
17
+ _skip_if_insufficient_gpus_for_tp,
18
+ compute_acceptance_len,
19
+ compute_acceptance_rate,
20
+ evaluate_llm_for_gsm8k,
21
+ get_instruct_coder_messages,
22
+ get_test_prompts,
23
+ greedy_sampling,
24
+ stochastic_sampling,
25
+)
26
+
27
+
28
+class AsyncSchedulingNotEnabledError(AssertionError):
29
+ """Raised when draft-model spec decode does not enable async scheduling."""
30
+
31
+
32
+@dataclass
33
+class ArgsTest:
34
+ target_model: str
35
+ draft_model: str
36
+ sampling_config: SamplingParams
37
+ num_speculative_tokens: int
38
+ expected_acceptance_rate: float
39
+ expected_acceptance_len: float
40
+ expected_gsm8k_accuracy: float = 0.0 # skip by default
41
+ # Defaults
42
+ enforce_eager: bool = True
43
+ parallel_drafting: bool = False
44
+ target_tensor_parallel_size: int = 1
45
+ draft_tensor_parallel_size: int = 1
46
+ max_model_len: int = 2048
47
+ gpu_memory_utilization: float = 0.5
48
+ dataset: str = "test_prompts"
49
+ num_prompts: int = 100
50
+
51
+
52
+def get_messages(dataset: str, n: int) -> list[Messages]:
53
+ if dataset == "test_prompts":
54
+ return get_test_prompts(mm_enabled=False, num_prompts=n)
55
+ if dataset == "likaixin/InstructCoder":
56
+ return get_instruct_coder_messages(n=n)
57
+ raise NotImplementedError(f"Dataset '{dataset}' not implemented")
58
+
59
+
60
+def some_high_acceptance_metrics() -> dict:
61
+ return {
62
+ "sampling_config": greedy_sampling(),
63
+ "num_speculative_tokens": 3,
64
+ "expected_acceptance_len": 3.4, # ref: 3.75
65
+ "expected_acceptance_rate": 0.8, # ref: 0.9
66
+ }
67
+
68
+
69
+cases = [
70
+ # Same model for draft and target, greedy sampling.
71
+ ArgsTest(
72
+ target_model="Qwen/Qwen3-0.6B",
73
+ draft_model="Qwen/Qwen3-0.6B",
74
+ sampling_config=greedy_sampling(),
75
+ num_speculative_tokens=3, # K
76
+ expected_acceptance_len=0.98 * (3 + 1), # epsilon discount of K + 1
77
+ expected_acceptance_rate=0.98, # slight epsilon
78
+ expected_gsm8k_accuracy=0.25, # ref: 35-40%
79
+ ),
80
+ # Smaller draft model, stochastic sampling.
81
+ ArgsTest(
82
+ target_model="Qwen/Qwen3-1.7B",
83
+ draft_model="Qwen/Qwen3-0.6B",
84
+ sampling_config=stochastic_sampling(),
85
+ num_speculative_tokens=3,
86
+ expected_acceptance_len=3.4, # ref: 3.7
87
+ expected_acceptance_rate=0.80, # ref: 0.90
88
+ expected_gsm8k_accuracy=0.5, # ref: 60%. Note gsm8k always runs greedy sampling
89
+ ),
90
+]
91
+
92
+
93
+@pytest.mark.parametrize("args", cases)
94
+@pytest.mark.parametrize("enforce_eager", [True, False])
95
+@single_gpu_only
96
+# TODO: Fix async_scheduling & engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
97
+@pytest.mark.xfail(
98
+ raises=AsyncSchedulingNotEnabledError,
99
+ reason="draft_model does not yet enable async_scheduling: issue #38929",
100
+)
101
+def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
102
+ args.enforce_eager = enforce_eager
103
+ assert_draft_model_correctness(args)
104
+
105
+
106
+@single_gpu_only
107
+# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
108
+@pytest.mark.xfail(
109
+ raises=AsyncSchedulingNotEnabledError,
110
+ reason="draft_model does not yet enable async_scheduling: issue #38929",
111
+)
112
+def test_draft_model_realistic_example():
113
+ args = ArgsTest(
114
+ target_model="Qwen/Qwen3-1.7B",
115
+ draft_model="Qwen/Qwen3-0.6B",
116
+ dataset="likaixin/InstructCoder",
117
+ num_speculative_tokens=3,
118
+ sampling_config=greedy_sampling(),
119
+ enforce_eager=False,
120
+ expected_acceptance_len=2.6, # ref: 2.86
121
+ expected_acceptance_rate=0.5, # ref: 0.62
122
+ )
123
+ assert_draft_model_correctness(args)
124
+
125
+
126
+@single_gpu_only
127
+# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
128
+@pytest.mark.xfail(
129
+ raises=AsyncSchedulingNotEnabledError,
130
+ reason="draft_model does not yet enable async_scheduling: issue #38929",
131
+)
132
+def test_draft_model_parallel_drafting():
133
+ args = ArgsTest(
134
+ target_model="Qwen/Qwen3-1.7B",
135
+ draft_model="amd/PARD-Qwen3-0.6B",
136
+ dataset="likaixin/InstructCoder",
137
+ num_speculative_tokens=3,
138
+ sampling_config=greedy_sampling(),
139
+ parallel_drafting=True,
140
+ enforce_eager=False,
141
+ expected_acceptance_len=2.3, # ref: 2.52
142
+ expected_acceptance_rate=0.4, # ref: 0.51
143
+ )
144
+ assert_draft_model_correctness(args)
145
+
146
+
147
+@pytest.mark.parametrize(
148
+ "models",
149
+ [
150
+ # target_model, draft_model
151
+ ("Qwen/Qwen3-1.7B-FP8", "Qwen/Qwen3-0.6B"), # target quantized
152
+ ("Qwen/Qwen3-1.7B", "Qwen/Qwen3-0.6B-FP8"), # draft quantized
153
+ ],
154
+ ids=["target_quantized", "draft_quantized"],
155
+)
156
+@pytest.mark.parametrize("enforce_eager", [True, False])
157
+@single_gpu_only
158
+# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
159
+@pytest.mark.xfail(
160
+ raises=AsyncSchedulingNotEnabledError,
161
+ reason="draft_model does not yet enable async_scheduling: issue #38929",
162
+)
163
+def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
164
+ tgt_model, draft_model = models
165
+ sd_case = ArgsTest(
166
+ target_model=tgt_model,
167
+ draft_model=draft_model,
168
+ **some_high_acceptance_metrics(),
169
+ enforce_eager=enforce_eager,
170
+ )
171
+ assert_draft_model_correctness(sd_case)
172
+
173
+
174
+@multi_gpu_only(num_gpus=2)
175
+# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
176
+@pytest.mark.xfail(
177
+ raises=AsyncSchedulingNotEnabledError,
178
+ reason="draft_model does not yet enable async_scheduling: issue #38929",
179
+)
180
+def test_draft_model_tensor_parallelism():
181
+ """Ensure spec decode works when running with TP > 1."""
182
+ _skip_if_insufficient_gpus_for_tp(2)
183
+ sd_case = ArgsTest(
184
+ target_model="Qwen/Qwen3-1.7B",
185
+ target_tensor_parallel_size=2,
186
+ draft_model="Qwen/Qwen3-0.6B",
187
+ draft_tensor_parallel_size=2,
188
+ **some_high_acceptance_metrics(),
189
+ enforce_eager=False,
190
+ expected_gsm8k_accuracy=0.5,
191
+ )
192
+ assert_draft_model_correctness(sd_case)
193
+
194
+
195
+@multi_gpu_only(num_gpus=2)
196
+def test_draft_model_engine_args_tensor_parallelism():
197
+ """Ensure the vllm_config for the draft model is created correctly,
198
+ and independently of the target model (quantization, TP, etc.)"""
199
+ _skip_if_insufficient_gpus_for_tp(2)
200
+
201
+ engine_args = EngineArgs(
202
+ model="Qwen/Qwen3-1.7B-FP8", # <<< tgt quantized
203
+ tensor_parallel_size=2,
204
+ speculative_config={
205
+ "model": "Qwen/Qwen3-0.6B", # <<< draft not quantized
206
+ "method": "draft_model",
207
+ "num_speculative_tokens": 3,
208
+ "draft_tensor_parallel_size": 1, # <<< valid arg name
209
+ },
210
+ )
211
+ target_config: VllmConfig = engine_args.create_engine_config()
212
+ assert target_config.parallel_config.tensor_parallel_size == 2
213
+ assert target_config.quant_config.get_name() == "fp8"
214
+
215
+ speculative_config = target_config.speculative_config
216
+ draft_config: VllmConfig = replace(
217
+ target_config,
218
+ quant_config=None,
219
+ parallel_config=replace(
220
+ speculative_config.draft_parallel_config,
221
+ rank=target_config.parallel_config.rank,
222
+ ),
223
+ model_config=speculative_config.draft_model_config,
224
+ )
225
+ assert draft_config.parallel_config.tensor_parallel_size == 1
226
+ assert draft_config.quant_config is None
227
+
228
+
229
+def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig:
230
+ """Replicate SpecDecodeBaseProposer._create_draft_vllm_config logic
231
+ so we can test it without instantiating a full proposer."""
232
+ spec_cfg = vllm_config.speculative_config
233
+ if spec_cfg.moe_backend is not None:
234
+ return replace(
235
+ vllm_config,
236
+ kernel_config=replace(
237
+ vllm_config.kernel_config,
238
+ moe_backend=spec_cfg.moe_backend,
239
+ ),
240
+ )
241
+ return vllm_config
242
+
243
+
244
+def test_draft_model_moe_backend_override():
245
+ """When moe_backend is set in speculative_config, the draft VllmConfig
246
+ should use it while the target keeps its own setting."""
247
+ engine_args = EngineArgs(
248
+ model="Qwen/Qwen3-1.7B",
249
+ tensor_parallel_size=1,
250
+ moe_backend="flashinfer_trtllm",
251
+ speculative_config={
252
+ "model": "Qwen/Qwen3-0.6B",
253
+ "method": "draft_model",
254
+ "num_speculative_tokens": 3,
255
+ "moe_backend": "triton",
256
+ },
257
+ )
258
+ tgt_config: VllmConfig = engine_args.create_engine_config()
259
+ assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm"
260
+ assert tgt_config.speculative_config.moe_backend == "triton"
261
+
262
+ draft_config = _apply_draft_moe_backend(tgt_config)
263
+ assert draft_config.kernel_config.moe_backend == "triton"
264
+ # Target config must be unaffected.
265
+ assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm"
266
+
267
+
268
+def test_draft_model_moe_backend_inherits_target():
269
+ """When moe_backend is not set in speculative_config, the draft should
270
+ inherit the target's moe_backend."""
271
+ engine_args = EngineArgs(
272
+ model="Qwen/Qwen3-1.7B",
273
+ tensor_parallel_size=1,
274
+ moe_backend="flashinfer_cutlass",
275
+ speculative_config={
276
+ "model": "Qwen/Qwen3-0.6B",
277
+ "method": "draft_model",
278
+ "num_speculative_tokens": 3,
279
+ },
280
+ )
281
+ tgt_config: VllmConfig = engine_args.create_engine_config()
282
+ assert tgt_config.kernel_config.moe_backend == "flashinfer_cutlass"
283
+ assert tgt_config.speculative_config.moe_backend is None
284
+
285
+ draft_config = _apply_draft_moe_backend(tgt_config)
286
+ assert draft_config.kernel_config.moe_backend == "flashinfer_cutlass"
287
+ assert draft_config is tgt_config
288
+
289
+
290
+def test_draft_model_moe_backend_default_auto():
291
+ """When neither target nor draft set moe_backend explicitly, both should
292
+ default to 'auto'."""
293
+ engine_args = EngineArgs(
294
+ model="Qwen/Qwen3-1.7B",
295
+ tensor_parallel_size=1,
296
+ speculative_config={
297
+ "model": "Qwen/Qwen3-0.6B",
298
+ "method": "draft_model",
299
+ "num_speculative_tokens": 3,
300
+ },
301
+ )
302
+ tgt_config: VllmConfig = engine_args.create_engine_config()
303
+ assert tgt_config.kernel_config.moe_backend == "auto"
304
+ assert tgt_config.speculative_config.moe_backend is None
305
+
306
+ draft_config = _apply_draft_moe_backend(tgt_config)
307
+ assert draft_config.kernel_config.moe_backend == "auto"
308
+ assert draft_config is tgt_config
309
+
310
+
311
+def test_draft_model_engine_args_rejects_invalid_tp_argname():
312
+ """The user should pass "draft_tensor_parallel_size" rather than
313
+ "tensor_parallel_size". We enforce this with validation."""
314
+
315
+ engine_args = EngineArgs(
316
+ model="Qwen/Qwen3-1.7B",
317
+ tensor_parallel_size=1,
318
+ speculative_config={
319
+ "model": "Qwen/Qwen3-0.6B",
320
+ "method": "draft_model",
321
+ "num_speculative_tokens": 3,
322
+ "tensor_parallel_size": 1, # <<< invalid arg name
323
+ },
324
+ )
325
+ with pytest.raises(ValueError):
326
+ engine_args.create_engine_config()
327
+
328
+
329
+def assert_draft_model_correctness(args: ArgsTest):
330
+ """Compare the outputs using and not using speculative decoding.
331
+ In the greedy decoding case, the outputs must match EXACTLY."""
332
+ test_prompts: list[Messages] = get_messages(
333
+ dataset=args.dataset, n=args.num_prompts
334
+ )
335
+
336
+ spec_llm = LLM(
337
+ model=args.target_model,
338
+ speculative_config={
339
+ "model": args.draft_model,
340
+ "method": "draft_model",
341
+ "num_speculative_tokens": args.num_speculative_tokens,
342
+ "max_model_len": args.max_model_len,
343
+ "enforce_eager": args.enforce_eager,
344
+ "draft_tensor_parallel_size": args.draft_tensor_parallel_size,
345
+ "parallel_drafting": args.parallel_drafting,
346
+ },
347
+ max_num_seqs=100, # limit cudagraph capture runtime
348
+ max_model_len=args.max_model_len,
349
+ gpu_memory_utilization=args.gpu_memory_utilization,
350
+ tensor_parallel_size=args.target_tensor_parallel_size,
351
+ enforce_eager=args.enforce_eager,
352
+ disable_log_stats=False, # enables get_metrics()
353
+ )
354
+
355
+ # we don't check the outputs, only check the metrics
356
+ spec_llm.chat(test_prompts, args.sampling_config)
357
+ metrics = spec_llm.get_metrics()
358
+ acceptance_rate: float = compute_acceptance_rate(metrics)
359
+ acceptance_len: float = compute_acceptance_len(metrics)
360
+
361
+ # Need to evaluate after getting metrics to avoid polluting the AR
362
+ evaluate_llm_for_gsm8k(
363
+ spec_llm, expected_accuracy_threshold=args.expected_gsm8k_accuracy
364
+ )
365
+
366
+ print(
367
+ f"spec-decode: target={args.target_model}, draft={args.draft_model}, "
368
+ f"temperature={args.sampling_config.temperature:.2f}, "
369
+ f"acceptance_rate={acceptance_rate:.2f}, "
370
+ f"acceptance_len={acceptance_len:.2f}, "
371
+ )
372
+
373
+ assert acceptance_rate >= args.expected_acceptance_rate
374
+ assert acceptance_len >= args.expected_acceptance_len
375
+ # draft_model supports async scheduling; assert it is active by default.
376
+ # Raise AsyncSchedulingNotEnabledError (a subclass of AssertionError) so that
377
+ # @pytest.mark.xfail(raises=AsyncSchedulingNotEnabledError) catches only this
378
+ # specific failure — leaving all other assertion failures (e.g. correctness or
379
+ # acceptance-rate checks above) visible as real test failures.
380
+ has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
381
+ del spec_llm # CLEANUP
382
+ torch.accelerator.empty_cache()
383
+ cleanup_dist_env_and_memory()
384
+ if not has_async:
385
+ raise AsyncSchedulingNotEnabledError(
386
+ "Expected async_scheduling=True for draft_model spec decode, got False."
387
+ " See https://github.com/vllm-project/vllm/issues/38929"
388
+ )
tests/v1/e2e/spec_decode/draft_model/test_lora.py
renamed
tests/v1/e2e/spec_decode/eagle/__init__.py
tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py
new
+167
@@ -0,0 +1,167 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from tests.utils import (
7
+ get_attn_backend_list_based_on_platform,
8
+ single_gpu_only,
9
+)
10
+from vllm import SamplingParams
11
+from vllm.platforms import current_platform
12
+
13
+from .utils import _run_eagle_correctness
14
+
15
+
16
+@single_gpu_only
17
+@pytest.mark.skipif(
18
+ current_platform.is_device_capability_family(100),
19
+ reason="DeepSeek head_dim=192 not supported on SM100/SM110 (Blackwell)",
20
+)
21
+@pytest.mark.parametrize(
22
+ [
23
+ "model_setup",
24
+ "mm_enabled",
25
+ "enable_chunked_prefill",
26
+ "model_impl",
27
+ "expected_accuracy_threshold",
28
+ ],
29
+ [
30
+ (
31
+ (
32
+ "eagle",
33
+ "eagle618/deepseek-v3-random",
34
+ "eagle618/eagle-deepseek-v3-random",
35
+ 1,
36
+ ),
37
+ False,
38
+ False,
39
+ "auto",
40
+ 0.0,
41
+ ),
42
+ ],
43
+ ids=["deepseek_eagle"],
44
+)
45
+@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
46
+def test_eagle_correctness_light(
47
+ monkeypatch: pytest.MonkeyPatch,
48
+ sampling_config: SamplingParams,
49
+ model_setup: tuple[str, str, str, int],
50
+ mm_enabled: bool,
51
+ expected_accuracy_threshold: float,
52
+ enable_chunked_prefill: bool,
53
+ model_impl: str,
54
+ attn_backend: str,
55
+):
56
+ _run_eagle_correctness(
57
+ monkeypatch,
58
+ sampling_config,
59
+ model_setup,
60
+ mm_enabled,
61
+ expected_accuracy_threshold,
62
+ enable_chunked_prefill,
63
+ model_impl,
64
+ attn_backend,
65
+ )
66
+
67
+
68
+@single_gpu_only
69
+@pytest.mark.parametrize(
70
+ [
71
+ "model_setup",
72
+ "mm_enabled",
73
+ "enable_chunked_prefill",
74
+ "model_impl",
75
+ "expected_accuracy_threshold",
76
+ ],
77
+ [
78
+ (
79
+ ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
80
+ False,
81
+ False,
82
+ "auto",
83
+ 0.8,
84
+ ),
85
+ pytest.param(
86
+ ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
87
+ False,
88
+ False,
89
+ "transformers",
90
+ 0.8,
91
+ # 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
+ ),
95
+ ),
96
+ pytest.param(
97
+ (
98
+ "eagle3",
99
+ "Qwen/Qwen3-VL-8B-Instruct",
100
+ "taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
101
+ 1,
102
+ ),
103
+ False,
104
+ False,
105
+ "auto",
106
+ 0.8,
107
+ marks=pytest.mark.skip(
108
+ reason="architecture of its eagle3 is LlamaForCausalLMEagle3"
109
+ ),
110
+ ),
111
+ pytest.param(
112
+ (
113
+ "eagle3",
114
+ "Qwen/Qwen2.5-VL-7B-Instruct",
115
+ "Rayzl/qwen2.5-vl-7b-eagle3-sgl",
116
+ 1,
117
+ ),
118
+ False,
119
+ False,
120
+ "auto",
121
+ 0.7,
122
+ marks=pytest.mark.skip(
123
+ reason="Skipping due to its head_dim not being a multiple of 32"
124
+ ),
125
+ ),
126
+ (
127
+ (
128
+ "eagle3",
129
+ "meta-llama/Llama-3.1-8B-Instruct",
130
+ "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
131
+ 1,
132
+ ),
133
+ False,
134
+ False,
135
+ "auto",
136
+ 0.7,
137
+ ),
138
+ ],
139
+ ids=[
140
+ "qwen3_eagle3",
141
+ "qwen3_eagle3-transformers",
142
+ "qwen3_vl_eagle3",
143
+ "qwen2_5_vl_eagle3",
144
+ "llama3_eagle3",
145
+ ],
146
+)
147
+@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
148
+def test_eagle_correctness_medium(
149
+ monkeypatch: pytest.MonkeyPatch,
150
+ sampling_config: SamplingParams,
151
+ model_setup: tuple[str, str, str, int],
152
+ mm_enabled: bool,
153
+ expected_accuracy_threshold: float,
154
+ enable_chunked_prefill: bool,
155
+ model_impl: str,
156
+ attn_backend: str,
157
+):
158
+ _run_eagle_correctness(
159
+ monkeypatch,
160
+ sampling_config,
161
+ model_setup,
162
+ mm_enabled,
163
+ expected_accuracy_threshold,
164
+ enable_chunked_prefill,
165
+ model_impl,
166
+ attn_backend,
167
+ )
tests/v1/e2e/spec_decode/eagle/utils.py
new
+137
@@ -0,0 +1,137 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
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
10
+from vllm.platforms import current_platform
11
+
12
+from ..utils import (
13
+ _skip_if_insufficient_gpus_for_tp,
14
+ evaluate_llm_for_gsm8k,
15
+ get_test_prompts,
16
+)
17
+
18
+
19
+def _run_eagle_correctness(
20
+ monkeypatch: pytest.MonkeyPatch,
21
+ sampling_config: SamplingParams,
22
+ model_setup: tuple[str, str, str, int],
23
+ mm_enabled: bool,
24
+ expected_accuracy_threshold: float,
25
+ enable_chunked_prefill: bool,
26
+ model_impl: str,
27
+ attn_backend: str,
28
+):
29
+ """
30
+ Compare the outputs of an original LLM and a speculative LLM
31
+ which should be the same when using eagle speculative decoding.
32
+ """
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
+ )
44
+
45
+ test_prompts = get_test_prompts(mm_enabled)
46
+
47
+ if "Llama-4-Scout" in model_setup[1] and attn_backend == "FLASH_ATTN":
48
+ if current_platform.is_rocm():
49
+ print(
50
+ "FLASH_ATTN for spec_decode not supported on "
51
+ "ROCm currently. Changing to FLEX_ATTENTION backend."
52
+ )
53
+ attention_config = {"backend": "FLEX_ATTENTION"}
54
+ else:
55
+ attention_config = None
56
+ else:
57
+ attention_config = {"backend": attn_backend}
58
+
59
+ if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm():
60
+ pytest.skip(
61
+ "TRITON_ATTN does not support "
62
+ "multi-token eagle spec decode on current platform"
63
+ )
64
+
65
+ with monkeypatch.context() as m:
66
+ m.setenv("VLLM_MLA_DISABLE", "1")
67
+
68
+ if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
69
+ if "deepseek" in model_setup[1].lower():
70
+ m.setenv("VLLM_ROCM_USE_AITER", "1")
71
+ m.delenv("VLLM_MLA_DISABLE", raising=False)
72
+ attention_config = {"backend": "ROCM_AITER_MLA"}
73
+ else:
74
+ m.setenv("VLLM_ROCM_USE_AITER", "1")
75
+
76
+ method, model_name, spec_model_name, tp_size = model_setup
77
+ _skip_if_insufficient_gpus_for_tp(tp_size)
78
+
79
+ max_model_len = 2048
80
+ max_num_batched_tokens = 128 if enable_chunked_prefill else max_model_len
81
+
82
+ ref_llm = LLM(
83
+ model=model_name,
84
+ max_model_len=max_model_len,
85
+ tensor_parallel_size=tp_size,
86
+ 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,
101
+ trust_remote_code=True,
102
+ tensor_parallel_size=tp_size,
103
+ speculative_config={
104
+ "method": method,
105
+ "model": spec_model_name,
106
+ "num_speculative_tokens": 3,
107
+ "max_model_len": max_model_len,
108
+ },
109
+ max_model_len=max_model_len,
110
+ max_num_batched_tokens=max_num_batched_tokens,
111
+ enable_chunked_prefill=enable_chunked_prefill,
112
+ model_impl=model_impl,
113
+ attention_config=attention_config,
114
+ )
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/__init__.py
tests/v1/e2e/spec_decode/mtp/test_mtp.py
new
+155
@@ -0,0 +1,155 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import os
5
+from typing import Any
6
+
7
+import pytest
8
+import torch
9
+
10
+from tests.utils import single_gpu_only
11
+from vllm import LLM, SamplingParams
12
+from vllm.distributed import cleanup_dist_env_and_memory
13
+from vllm.platforms import current_platform
14
+
15
+from ..utils import (
16
+ _skip_if_insufficient_gpus_for_tp,
17
+ evaluate_llm_for_gsm8k,
18
+ get_test_prompts,
19
+)
20
+
21
+
22
+@pytest.mark.parametrize(
23
+ ["model_setup", "mm_enabled", "expected_accuracy_threshold"],
24
+ [
25
+ (("mtp", "XiaomiMiMo/MiMo-7B-Base", 1), False, 0.5), # ref: 65%-70%
26
+ pytest.param(
27
+ ("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1),
28
+ False,
29
+ 0.0,
30
+ marks=pytest.mark.skipif(
31
+ current_platform.is_device_capability_family(100),
32
+ reason="DeepSeek MTP: TRTLLM MoE top_k check fails on Blackwell",
33
+ ),
34
+ ), # dummy model
35
+ (
36
+ ("mtp", "Qwen/Qwen3.5-0.8B-Base", 1),
37
+ False,
38
+ 0.20,
39
+ ), # hybrid + MTP, ref: ~34%-35%
40
+ (
41
+ ("mtp", "google/gemma-4-E4B-it", 1, "google/gemma-4-E4B-it-assistant"),
42
+ False,
43
+ 0.50,
44
+ ), # gemma4 MTP with assistant model, ref: ~62%
45
+ ],
46
+ ids=["mimo", "deepseek", "qwen3_5-hybrid", "gemma4-e4b"],
47
+)
48
+@single_gpu_only
49
+def test_mtp_correctness(
50
+ monkeypatch: pytest.MonkeyPatch,
51
+ sampling_config: SamplingParams,
52
+ model_setup: tuple[str, str, int] | tuple[str, str, int, str],
53
+ mm_enabled: bool,
54
+ expected_accuracy_threshold: float,
55
+):
56
+ """
57
+ Compare the outputs of a original LLM and a speculative LLM
58
+ which should be the same when using MTP speculative decoding. Due to some variance
59
+ in the engine, it is possible for some outputs to differ, so we expect that at least
60
+ 6/10 output tokens match exactly, and that the GSM8k accuracy is above a precomputed
61
+ reference threshold for each model.
62
+ """
63
+ # Generate test prompts inside the function instead of using fixture
64
+ test_prompts = get_test_prompts(mm_enabled)
65
+ with monkeypatch.context() as m:
66
+ m.setenv("VLLM_MLA_DISABLE", "1")
67
+
68
+ if len(model_setup) == 4:
69
+ method, model_name, tp_size, draft_model = model_setup
70
+ else:
71
+ method, model_name, tp_size = model_setup
72
+ draft_model = None
73
+ _skip_if_insufficient_gpus_for_tp(tp_size)
74
+
75
+ if "Qwen3.5" in model_name and os.environ.get("VLLM_USE_V2_MODEL_RUNNER"):
76
+ pytest.skip(
77
+ "Model Runner V2 does not yet support hybrid models "
78
+ "(Qwen3.5 mixes Mamba-style GDN with attention layers)."
79
+ )
80
+
81
+ attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
82
+
83
+ # Skip multimodal profiling for models that don't need it in this test.
84
+ extra_kwargs: dict[str, Any] = {}
85
+ if "Qwen3.5" in model_name:
86
+ extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0}
87
+ elif "gemma-4" in model_name:
88
+ extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "audio": 0}
89
+
90
+ if draft_model is not None and "gemma-4" in draft_model:
91
+ import transformers
92
+ from packaging.version import Version
93
+
94
+ if Version(transformers.__version__) < Version("5.8.0"):
95
+ pytest.skip(
96
+ "Gemma4 MTP assistant requires transformers>=5.8.0, "
97
+ f"got {transformers.__version__}"
98
+ )
99
+
100
+ ref_llm = LLM(
101
+ model=model_name,
102
+ max_model_len=2048,
103
+ tensor_parallel_size=tp_size,
104
+ trust_remote_code=True,
105
+ attention_backend=attn_backend,
106
+ **extra_kwargs,
107
+ )
108
+ ref_outputs = ref_llm.chat(test_prompts, sampling_config)
109
+ evaluate_llm_for_gsm8k(
110
+ ref_llm, expected_accuracy_threshold=expected_accuracy_threshold
111
+ )
112
+ del ref_llm
113
+ torch.accelerator.empty_cache()
114
+ cleanup_dist_env_and_memory()
115
+
116
+ speculative_config: dict[str, Any] = {
117
+ "method": method,
118
+ "num_speculative_tokens": 1,
119
+ "max_model_len": 2048,
120
+ }
121
+ if draft_model is not None:
122
+ speculative_config["model"] = draft_model
123
+ speculative_config["num_speculative_tokens"] = 2
124
+
125
+ spec_llm = LLM(
126
+ model=model_name,
127
+ trust_remote_code=True,
128
+ tensor_parallel_size=tp_size,
129
+ speculative_config=speculative_config,
130
+ max_model_len=2048,
131
+ attention_backend=attn_backend,
132
+ **extra_kwargs,
133
+ )
134
+ # MTP supports async scheduling; assert it is active by default.
135
+ assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
136
+ evaluate_llm_for_gsm8k(
137
+ spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
138
+ )
139
+ spec_outputs = spec_llm.chat(test_prompts, sampling_config)
140
+ matches = 0
141
+ misses = 0
142
+ for ref_output, spec_output in zip(ref_outputs, spec_outputs):
143
+ if ref_output.outputs[0].text == spec_output.outputs[0].text:
144
+ matches += 1
145
+ else:
146
+ misses += 1
147
+ print(f"ref_output: {ref_output.outputs[0].text}")
148
+ print(f"spec_output: {spec_output.outputs[0].text}")
149
+
150
+ # Heuristic: expect at least 80% of the prompts to match exactly
151
+ # Upon failure, inspect the outputs to check for inaccuracy.
152
+ assert matches > int(0.8 * len(ref_outputs))
153
+ del spec_llm
154
+ torch.accelerator.empty_cache()
155
+ cleanup_dist_env_and_memory()
tests/v1/e2e/spec_decode/ngram_suffix/__init__.py
tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py
new
+147
@@ -0,0 +1,147 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from tests.utils import single_gpu_only
8
+from vllm import LLM, SamplingParams
9
+from vllm.config import CompilationConfig
10
+from vllm.distributed import cleanup_dist_env_and_memory
11
+from vllm.platforms import current_platform
12
+
13
+from ..utils import evaluate_llm_for_gsm8k, get_test_prompts
14
+
15
+
16
+@pytest.fixture
17
+def disable_vllm_compile_cache_on_rocm(request: pytest.FixtureRequest) -> None:
18
+ if current_platform.is_rocm():
19
+ request.getfixturevalue("disable_vllm_compile_cache")
20
+
21
+
22
+@pytest.mark.parametrize(
23
+ "speculative_config",
24
+ [
25
+ {
26
+ "method": "ngram",
27
+ "prompt_lookup_max": 5,
28
+ "prompt_lookup_min": 3,
29
+ "num_speculative_tokens": 3,
30
+ },
31
+ {
32
+ "method": "suffix",
33
+ "suffix_decoding_max_spec_factor": 2.0,
34
+ },
35
+ ],
36
+)
37
+@pytest.mark.usefixtures("disable_vllm_compile_cache_on_rocm")
38
+@single_gpu_only
39
+def test_ngram_and_suffix_correctness(
40
+ speculative_config: dict,
41
+ model_name: str,
42
+ vllm_runner,
43
+):
44
+ with vllm_runner(
45
+ model_name,
46
+ # Keep LLM defaults; VllmRunner only provides lifecycle cleanup here.
47
+ trust_remote_code=False,
48
+ enable_chunked_prefill=None,
49
+ speculative_config=speculative_config,
50
+ max_model_len=4096,
51
+ # Preserve LLM's default compilation/cudagraph configuration. Without
52
+ # this, VllmRunner injects its reduced test-only capture sizes.
53
+ compilation_config=CompilationConfig(),
54
+ ) as runner:
55
+ evaluate_llm_for_gsm8k(runner.llm)
56
+
57
+
58
+@pytest.mark.parametrize("async_scheduling", [True], ids=["async"])
59
+@single_gpu_only
60
+def test_ngram_gpu_default_with_async_scheduling(
61
+ async_scheduling: bool,
62
+):
63
+ """
64
+ Test ngram_gpu speculative decoding (k=3) correctness with and without
65
+ async scheduling, validated via GSM8K accuracy.
66
+ Uses Qwen/Qwen3-8B (ref GSM8K accuracy: 87%-92%).
67
+ """
68
+ qwen3_model = "Qwen/Qwen3-8B"
69
+ spec_llm = LLM(
70
+ model=qwen3_model,
71
+ speculative_config={
72
+ "method": "ngram_gpu",
73
+ "prompt_lookup_max": 3,
74
+ "prompt_lookup_min": 2,
75
+ "num_speculative_tokens": 2,
76
+ },
77
+ max_model_len=4096,
78
+ 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()
88
+
89
+
90
+@single_gpu_only
91
+def test_suffix_decoding_acceptance(
92
+ monkeypatch: pytest.MonkeyPatch,
93
+ sampling_config: SamplingParams,
94
+ model_name: str,
95
+):
96
+ """
97
+ Check that suffix decoding caching takes effect and improves acceptance
98
+ lengths and acceptance rates over multiple runs of the same prompts.
99
+ """
100
+ test_prompts = get_test_prompts(mm_enabled=False)
101
+
102
+ spec_llm = LLM(
103
+ model=model_name,
104
+ speculative_config={
105
+ "method": "suffix",
106
+ "suffix_decoding_max_spec_factor": 2.0,
107
+ "suffix_decoding_max_cached_requests": 1000,
108
+ },
109
+ max_model_len=1024,
110
+ 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)
125
+
126
+ # Calculate the acceptance rates for the first and last runs.
127
+ first_accept_tokens = num_accept[0]
128
+ first_draft_tokens = num_draft[0]
129
+ first_accept_rate = first_accept_tokens / first_draft_tokens
130
+
131
+ # Take the diff since the stats are cumulative.
132
+ last_accept_tokens = num_accept[-1] - num_accept[-2]
133
+ last_draft_tokens = num_draft[-1] - num_draft[-2]
134
+ last_accept_rate = last_accept_tokens / last_draft_tokens
135
+
136
+ # Expect the acceptance length to improve.
137
+ assert first_accept_tokens < last_accept_tokens
138
+
139
+ # Expect the acceptance rate to improve.
140
+ assert first_accept_rate < last_accept_rate
141
+
142
+ # 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()
tests/v1/e2e/spec_decode/speculators/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/v1/e2e/spec_decode/speculators/test_speculators.py
new
+104
@@ -0,0 +1,104 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+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
11
+
12
+
13
+@pytest.mark.parametrize(
14
+ ["model_path", "expected_accuracy_threshold"],
15
+ [
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%
18
+ ],
19
+ ids=["llama3_eagle3_speculator", "qwen3_eagle3_speculator"],
20
+)
21
+@single_gpu_only
22
+def test_speculators_model_integration(
23
+ monkeypatch: pytest.MonkeyPatch,
24
+ sampling_config: SamplingParams,
25
+ model_path: str,
26
+ expected_accuracy_threshold: float,
27
+ vllm_runner,
28
+):
29
+ """
30
+ Test that speculators models work with the simplified integration.
31
+
32
+ This verifies the `vllm serve <speculator-model>` use case where
33
+ speculative config is automatically detected from the model config
34
+ without requiring explicit --speculative-config argument.
35
+
36
+ Tests:
37
+ 1. Speculator model is correctly detected
38
+ 2. Verifier model is extracted from speculator config
39
+ 3. Speculative decoding is automatically enabled
40
+ 4. Text generation works correctly
41
+ 5. GSM8k accuracy of the model passes a sanity check when speculative decoding on
42
+ 6. Output matches reference (non-speculative) generation
43
+ """
44
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
45
+
46
+ # Generate test prompts
47
+ test_prompts = get_test_prompts(mm_enabled=False)
48
+
49
+ # First run: Direct speculator model (simplified integration)
50
+ with vllm_runner(
51
+ model_path,
52
+ trust_remote_code=False,
53
+ enable_chunked_prefill=None,
54
+ compilation_config=CompilationConfig(),
55
+ max_model_len=4096,
56
+ gpu_memory_utilization=0.92,
57
+ ) as spec_runner:
58
+ evaluate_llm_for_gsm8k(
59
+ spec_runner.llm, expected_accuracy_threshold=expected_accuracy_threshold
60
+ )
61
+ spec_outputs = spec_runner.llm.chat(test_prompts, sampling_config)
62
+
63
+ # Verify speculative config was auto-detected
64
+ assert spec_runner.llm.llm_engine.vllm_config.speculative_config is not None, (
65
+ f"Speculative config should be auto-detected for {model_path}"
66
+ )
67
+
68
+ spec_config = spec_runner.llm.llm_engine.vllm_config.speculative_config
69
+ assert spec_config.num_speculative_tokens > 0, (
70
+ f"Expected positive speculative tokens, "
71
+ f"got {spec_config.num_speculative_tokens}"
72
+ )
73
+
74
+ # Verify draft model is set to the speculator model
75
+ assert spec_config.model == model_path, (
76
+ f"Draft model should be {model_path}, got {spec_config.model}"
77
+ )
78
+
79
+ # Extract verifier model for reference run
80
+ verifier_model = spec_runner.llm.llm_engine.vllm_config.model_config.model
81
+
82
+ # Second run: Reference without speculative decoding
83
+ with vllm_runner(
84
+ verifier_model,
85
+ trust_remote_code=False,
86
+ enable_chunked_prefill=None,
87
+ compilation_config=CompilationConfig(),
88
+ max_model_len=4096,
89
+ gpu_memory_utilization=0.92,
90
+ ) as ref_runner:
91
+ ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config)
92
+
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
+
100
+ # 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."
104
+ )
tests/v1/e2e/spec_decode/test_laguna_dflash.py
deleted
-60
@@ -1,60 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import pytest
5
-import torch
6
-
7
-from tests.utils import large_gpu_mark
8
-from vllm import LLM, SamplingParams
9
-from vllm.distributed import cleanup_dist_env_and_memory
10
-
11
-
12
-def _get_counter(metrics, name: str) -> float:
13
- metric = next((m for m in metrics if m.name == name), None)
14
- assert metric is not None, f"Missing metric: {name}"
15
- return metric.value
16
-
17
-
18
-@pytest.mark.slow_test
19
-@large_gpu_mark(min_gb=80)
20
-def test_laguna_dflash_hf_pair_smoke(monkeypatch):
21
- """Smoke-test the public Laguna XS-2.1 base/DFlash checkpoint pair."""
22
- monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
23
-
24
- llm = LLM(
25
- model="poolside/Laguna-XS-2.1",
26
- trust_remote_code=True,
27
- speculative_config={
28
- "method": "dflash",
29
- "model": "poolside/Laguna-XS-2.1-DFlash",
30
- "num_speculative_tokens": 15,
31
- },
32
- max_model_len=8192,
33
- max_num_batched_tokens=65536,
34
- max_num_seqs=32,
35
- enforce_eager=True,
36
- disable_log_stats=False,
37
- )
38
-
39
- try:
40
- outputs = llm.generate(
41
- [
42
- "What is the capital of the United Kingdom?",
43
- "Write a Python function that returns the square of a number.",
44
- ],
45
- SamplingParams(temperature=0.0, max_tokens=32, ignore_eos=True),
46
- )
47
- assert len(outputs) == 2
48
- assert all(output.outputs[0].text for output in outputs)
49
-
50
- metrics = llm.get_metrics()
51
- num_drafts = _get_counter(metrics, "vllm:spec_decode_num_drafts")
52
- num_accepted = _get_counter(metrics, "vllm:spec_decode_num_accepted_tokens")
53
-
54
- assert num_drafts > 0
55
- acceptance_len = 1 + (num_accepted / num_drafts)
56
- assert acceptance_len > 1.0
57
- finally:
58
- del llm
59
- torch.accelerator.empty_cache()
60
- cleanup_dist_env_and_memory()
tests/v1/e2e/spec_decode/test_spec_decode.py
deleted
-1679
@@ -1,1679 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import os
4
-import random
5
-from collections.abc import Iterable
6
-from dataclasses import dataclass
7
-from typing import Any
8
-
9
-import pytest
10
-import torch
11
-from tqdm import tqdm
12
-
13
-from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_offline
14
-from tests.utils import (
15
- get_attn_backend_list_based_on_platform,
16
- large_gpu_mark,
17
- multi_gpu_marks,
18
- multi_gpu_only,
19
- single_gpu_only,
20
- wait_for_rocm_memory_to_settle,
21
-)
22
-from vllm import LLM, SamplingParams
23
-from vllm.assets.base import VLLM_S3_BUCKET_URL
24
-from vllm.assets.image import VLM_IMAGES_DIR
25
-from vllm.benchmarks.datasets import InstructCoderDataset
26
-from vllm.config import CompilationConfig, VllmConfig, replace
27
-from vllm.distributed import cleanup_dist_env_and_memory
28
-from vllm.engine.arg_utils import EngineArgs
29
-from vllm.platforms import current_platform
30
-from vllm.v1.metrics.reader import Metric
31
-
32
-MTP_SIMILARITY_RATE = 0.8
33
-
34
-
35
-class AsyncSchedulingNotEnabledError(AssertionError):
36
- """Raised when async_scheduling is expected to be True for draft_model
37
- spec decode but is False. Tracked in:
38
- https://github.com/vllm-project/vllm/issues/38929
39
- """
40
-
41
-
42
-def _skip_if_insufficient_gpus_for_tp(tp_size: int):
43
- """Skip test if available GPUs < tp_size on ROCm."""
44
- available_gpus = torch.accelerator.device_count()
45
- if available_gpus < tp_size:
46
- pytest.skip(
47
- f"Test requires {tp_size} GPUs, but only {available_gpus} available"
48
- )
49
-
50
-
51
-Messages = list[dict[str, Any]]
52
-
53
-
54
-def get_test_prompts(mm_enabled: bool, num_prompts: int = 100) -> list[Messages]:
55
- prompt_types = ["repeat", "gsm8k"]
56
- if mm_enabled:
57
- prompt_types.append("mm")
58
- prompts: list[Messages] = []
59
-
60
- num_repeat_prompts = num_prompts // len(prompt_types)
61
- if mm_enabled:
62
- num_gsm8k_prompts = num_prompts // len(prompt_types)
63
- num_mm_prompts = num_prompts - num_repeat_prompts - num_gsm8k_prompts
64
- else:
65
- num_mm_prompts = 0
66
- num_gsm8k_prompts = num_prompts - num_repeat_prompts
67
-
68
- # Generate a mixed batch of prompts, some of which can be easily
69
- # predicted by n-gram matching and some which likely cannot.
70
- random.seed(0)
71
- for _ in range(num_repeat_prompts):
72
- word_choices = ["test", "temp", "hello", "where"]
73
- word = random.choice(word_choices)
74
- prompts.append(
75
- [
76
- {
77
- "role": "user",
78
- "content": f"""
79
- please repeat the word '{word}' 10 times.
80
- give no other output than the word at least ten times in a row,
81
- in lowercase with spaces between each word and without quotes.
82
- """,
83
- }
84
- ]
85
- )
86
- prompts.extend(
87
- [{"role": "user", "content": prompt}]
88
- for prompt in _build_gsm8k_prompts(
89
- num_questions=num_gsm8k_prompts, num_shots=5
90
- )[0]
91
- )
92
- for _ in range(num_mm_prompts):
93
- placeholders = [
94
- {
95
- "type": "image_url",
96
- "image_url": {
97
- "url": f"{VLLM_S3_BUCKET_URL}/{VLM_IMAGES_DIR}/stop_sign.jpg"
98
- },
99
- }
100
- ]
101
- prompt = [
102
- *placeholders,
103
- {"type": "text", "text": "The meaning of the image is"},
104
- ]
105
- prompts.append([{"role": "user", "content": prompt}])
106
-
107
- return prompts
108
-
109
-
110
-def get_instruct_coder_messages(n: int) -> list[Messages]:
111
- dataset = InstructCoderDataset(
112
- dataset_path="likaixin/InstructCoder", dataset_split="train"
113
- )
114
- prompts: Iterable[str] = dataset.sample_prompts(n=n)
115
- return [[{"role": "user", "content": prompt}] for prompt in prompts]
116
-
117
-
118
-@pytest.fixture
119
-def sampling_config():
120
- return greedy_sampling()
121
-
122
-
123
-def greedy_sampling() -> SamplingParams:
124
- return SamplingParams(temperature=0, max_tokens=10, ignore_eos=False)
125
-
126
-
127
-def stochastic_sampling() -> SamplingParams:
128
- return SamplingParams(temperature=1.0, max_tokens=10, ignore_eos=False)
129
-
130
-
131
-@pytest.fixture
132
-def model_name():
133
- return "meta-llama/Llama-3.1-8B-Instruct"
134
-
135
-
136
-def evaluate_llm_for_gsm8k(llm: LLM, expected_accuracy_threshold: float = 0.70) -> None:
137
- """Evaluate the LLM on GSM8K and check that accuracy is above a sanity threshold.
138
-
139
- The default threshold assumes the LLM uses the same target model as the "model_name"
140
- fixture, with max model len == 4096. Precomputed reference value is 75% to 80%
141
- on GSM8K with greedy decoding, so we check that it's above a sanity threshold of 70%
142
- to verify that the model is correct.
143
- """
144
- if expected_accuracy_threshold <= 0.0:
145
- print("Skipping GSM8K evaluation")
146
- return
147
- results = evaluate_gsm8k_offline(llm)
148
- accuracy = results["accuracy"]
149
- print(f"GSM8K accuracy: {accuracy:.3f}")
150
- assert accuracy >= expected_accuracy_threshold, (
151
- f"Expected GSM8K accuracy >= {expected_accuracy_threshold}, got {accuracy:.3f}"
152
- )
153
-
154
-
155
-@pytest.fixture(autouse=True)
156
-def reset_torch_dynamo():
157
- """Reset torch dynamo cache before each test"""
158
- yield
159
- # Cleanup after test
160
- torch._dynamo.reset()
161
-
162
-
163
-@pytest.fixture
164
-def disable_vllm_compile_cache_on_rocm(request: pytest.FixtureRequest) -> None:
165
- if current_platform.is_rocm():
166
- request.getfixturevalue("disable_vllm_compile_cache")
167
-
168
-
169
-@pytest.mark.parametrize(
170
- "speculative_config",
171
- [
172
- {
173
- "method": "ngram",
174
- "prompt_lookup_max": 5,
175
- "prompt_lookup_min": 3,
176
- "num_speculative_tokens": 3,
177
- },
178
- {
179
- "method": "suffix",
180
- "suffix_decoding_max_spec_factor": 2.0,
181
- },
182
- ],
183
-)
184
-@pytest.mark.usefixtures("disable_vllm_compile_cache_on_rocm")
185
-@single_gpu_only
186
-@large_gpu_mark(min_gb=20)
187
-def test_ngram_and_suffix_correctness(
188
- speculative_config: dict,
189
- model_name: str,
190
- vllm_runner,
191
-):
192
- with vllm_runner(
193
- model_name,
194
- # Keep LLM defaults; VllmRunner only provides lifecycle cleanup here.
195
- trust_remote_code=False,
196
- enable_chunked_prefill=None,
197
- speculative_config=speculative_config,
198
- max_model_len=4096,
199
- # Preserve LLM's default compilation/cudagraph configuration. Without
200
- # this, VllmRunner injects its reduced test-only capture sizes.
201
- compilation_config=CompilationConfig(),
202
- ) as runner:
203
- evaluate_llm_for_gsm8k(runner.llm)
204
-
205
-
206
-@pytest.mark.parametrize("async_scheduling", [True], ids=["async"])
207
-@single_gpu_only
208
-@large_gpu_mark(min_gb=20)
209
-def test_ngram_gpu_default_with_async_scheduling(
210
- async_scheduling: bool,
211
-):
212
- """
213
- Test ngram_gpu speculative decoding (k=3) correctness with and without
214
- async scheduling, validated via GSM8K accuracy.
215
- Uses Qwen/Qwen3-8B (ref GSM8K accuracy: 87%-92%).
216
- """
217
- qwen3_model = "Qwen/Qwen3-8B"
218
- spec_llm = LLM(
219
- model=qwen3_model,
220
- speculative_config={
221
- "method": "ngram_gpu",
222
- "prompt_lookup_max": 3,
223
- "prompt_lookup_min": 2,
224
- "num_speculative_tokens": 2,
225
- },
226
- max_model_len=4096,
227
- async_scheduling=async_scheduling,
228
- )
229
- # Assert the resolved async_scheduling config matches what was requested.
230
- assert (
231
- spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
232
- == async_scheduling
233
- )
234
- evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8)
235
- del spec_llm
236
- cleanup_dist_env_and_memory()
237
-
238
-
239
-@single_gpu_only
240
-@large_gpu_mark(min_gb=20)
241
-def test_suffix_decoding_acceptance(
242
- monkeypatch: pytest.MonkeyPatch,
243
- sampling_config: SamplingParams,
244
- model_name: str,
245
-):
246
- """
247
- Check that suffix decoding caching takes effect and improves acceptance
248
- lengths and acceptance rates over multiple runs of the same prompts.
249
- """
250
- test_prompts = get_test_prompts(mm_enabled=False)
251
-
252
- spec_llm = LLM(
253
- model=model_name,
254
- speculative_config={
255
- "method": "suffix",
256
- "suffix_decoding_max_spec_factor": 2.0,
257
- "suffix_decoding_max_cached_requests": 1000,
258
- },
259
- max_model_len=1024,
260
- disable_log_stats=False,
261
- )
262
-
263
- # Run several times and check that the accepted tokens increase.
264
- num_draft = []
265
- num_accept = []
266
- for i in range(10): # Run multiple times to warm up the cache.
267
- spec_llm.chat(test_prompts, sampling_config)
268
- # Collect draft and acceptance stats.
269
- metrics = spec_llm.get_metrics()
270
- for metric in metrics:
271
- if metric.name == "vllm:spec_decode_num_draft_tokens":
272
- num_draft.append(metric.value)
273
- if metric.name == "vllm:spec_decode_num_accepted_tokens":
274
- num_accept.append(metric.value)
275
-
276
- # Calculate the acceptance rates for the first and last runs.
277
- first_accept_tokens = num_accept[0]
278
- first_draft_tokens = num_draft[0]
279
- first_accept_rate = first_accept_tokens / first_draft_tokens
280
-
281
- # Take the diff since the stats are cumulative.
282
- last_accept_tokens = num_accept[-1] - num_accept[-2]
283
- last_draft_tokens = num_draft[-1] - num_draft[-2]
284
- last_accept_rate = last_accept_tokens / last_draft_tokens
285
-
286
- # Expect the acceptance length to improve.
287
- assert first_accept_tokens < last_accept_tokens
288
-
289
- # Expect the acceptance rate to improve.
290
- assert first_accept_rate < last_accept_rate
291
-
292
- # Heuristic: expect at least 80.0% acceptance rate at the end.
293
- assert last_accept_rate > 0.80
294
-
295
- del spec_llm
296
- torch.accelerator.empty_cache()
297
- cleanup_dist_env_and_memory()
298
-
299
-
300
-@pytest.mark.slow_test
301
-@large_gpu_mark(min_gb=80)
302
-def test_gemma4_dspark_correctness_and_acceptance_rate(
303
- monkeypatch: pytest.MonkeyPatch,
304
-):
305
- """
306
- E2E test for Gemma4 DSpark speculative decoding: acceptance rate/length
307
- regression coverage plus GSM8K correctness, at temperature=1.0 to exercise
308
- the probabilistic draft-sampling/rejection-sampling path (not just greedy).
309
-
310
- Uses google/gemma-4-12B-it as target with the dspark_gemma4_12b_block7 draft
311
- model. Exercises the full Gemma4 DSpark path (draft build over the reused
312
- base classes DFlashQwen3Model / Qwen3DSparkForCausalLM / Gemma4MTP* /
313
- DSparkMarkovHead, the fused context-KV precompute, the Markov head, and
314
- rejection sampling), so it fails if any base class drifts out of sync.
315
-
316
- gemma-4-12B is instruct-tuned, so GSM8K is run through the chat template
317
- (use_chat_completions=True; raw few-shot completion collapses to a few
318
- percent). Reference: measured over 5 runs of 200 GSM8K questions at
319
- temperature=1.0 (prefix caching disabled):
320
- accuracy: min=0.900 max=0.955 mean=0.937
321
- acceptance_rate: min=0.578 max=0.595 mean=0.588
322
- acceptance_len: min=5.044 max=5.167 mean=5.116
323
- Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling
324
- """
325
- monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0")
326
-
327
- spec_llm = LLM(
328
- model="google/gemma-4-12B-it",
329
- trust_remote_code=True,
330
- speculative_config={
331
- "method": "dspark",
332
- "model": "deepseek-ai/dspark_gemma4_12b_block7",
333
- "num_speculative_tokens": 7,
334
- "draft_sample_method": "probabilistic",
335
- },
336
- max_model_len=8192,
337
- max_num_seqs=32,
338
- gpu_memory_utilization=0.8,
339
- enforce_eager=True,
340
- enable_prefix_caching=False,
341
- disable_log_stats=False,
342
- )
343
- try:
344
- results = evaluate_gsm8k_offline(
345
- spec_llm, num_questions=200, temperature=1.0, use_chat_completions=True
346
- )
347
- gsm8k_accuracy = results["accuracy"]
348
-
349
- metrics = spec_llm.get_metrics()
350
- acceptance_rate = compute_acceptance_rate(metrics)
351
- acceptance_len = compute_acceptance_len(metrics)
352
- print(
353
- f"Gemma4 DSpark acceptance_rate={acceptance_rate:.2f}, "
354
- f"acceptance_len={acceptance_len:.2f}, "
355
- f"gsm8k_accuracy={gsm8k_accuracy:.3f}"
356
- )
357
-
358
- assert acceptance_rate >= 0.588 * 0.9
359
- assert acceptance_len >= 5.116 * 0.9
360
- assert gsm8k_accuracy >= 0.937 * 0.9
361
- finally:
362
- del spec_llm
363
- torch.accelerator.empty_cache()
364
- cleanup_dist_env_and_memory()
365
-
366
-
367
-@pytest.mark.parametrize(
368
- ["model_path", "expected_accuracy_threshold"],
369
- [
370
- ("RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3", 0.7), # ref: 75%-80%
371
- ("RedHatAI/Qwen3-8B-speculator.eagle3", 0.8), # ref: 87%-92%
372
- ],
373
- ids=["llama3_eagle3_speculator", "qwen3_eagle3_speculator"],
374
-)
375
-@single_gpu_only
376
-@large_gpu_mark(min_gb=24)
377
-def test_speculators_model_integration(
378
- monkeypatch: pytest.MonkeyPatch,
379
- sampling_config: SamplingParams,
380
- model_path: str,
381
- expected_accuracy_threshold: float,
382
- vllm_runner,
383
-):
384
- """
385
- Test that speculators models work with the simplified integration.
386
-
387
- This verifies the `vllm serve <speculator-model>` use case where
388
- speculative config is automatically detected from the model config
389
- without requiring explicit --speculative-config argument.
390
-
391
- Tests:
392
- 1. Speculator model is correctly detected
393
- 2. Verifier model is extracted from speculator config
394
- 3. Speculative decoding is automatically enabled
395
- 4. Text generation works correctly
396
- 5. GSM8k accuracy of the model passes a sanity check when speculative decoding on
397
- 6. Output matches reference (non-speculative) generation
398
- """
399
- monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
400
-
401
- # Generate test prompts
402
- test_prompts = get_test_prompts(mm_enabled=False)
403
-
404
- # First run: Direct speculator model (simplified integration)
405
- with vllm_runner(
406
- model_path,
407
- trust_remote_code=False,
408
- enable_chunked_prefill=None,
409
- compilation_config=CompilationConfig(),
410
- max_model_len=4096,
411
- gpu_memory_utilization=0.92,
412
- ) as spec_runner:
413
- evaluate_llm_for_gsm8k(
414
- spec_runner.llm, expected_accuracy_threshold=expected_accuracy_threshold
415
- )
416
- spec_outputs = spec_runner.llm.chat(test_prompts, sampling_config)
417
-
418
- # Verify speculative config was auto-detected
419
- assert spec_runner.llm.llm_engine.vllm_config.speculative_config is not None, (
420
- f"Speculative config should be auto-detected for {model_path}"
421
- )
422
-
423
- spec_config = spec_runner.llm.llm_engine.vllm_config.speculative_config
424
- assert spec_config.num_speculative_tokens > 0, (
425
- f"Expected positive speculative tokens, "
426
- f"got {spec_config.num_speculative_tokens}"
427
- )
428
-
429
- # Verify draft model is set to the speculator model
430
- assert spec_config.model == model_path, (
431
- f"Draft model should be {model_path}, got {spec_config.model}"
432
- )
433
-
434
- # Extract verifier model for reference run
435
- verifier_model = spec_runner.llm.llm_engine.vllm_config.model_config.model
436
-
437
- # Second run: Reference without speculative decoding
438
- with vllm_runner(
439
- verifier_model,
440
- trust_remote_code=False,
441
- enable_chunked_prefill=None,
442
- compilation_config=CompilationConfig(),
443
- max_model_len=4096,
444
- gpu_memory_utilization=0.92,
445
- ) as ref_runner:
446
- ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config)
447
-
448
- # Compare outputs
449
- matches = sum(
450
- 1
451
- for ref, spec in zip(ref_outputs, spec_outputs)
452
- if ref.outputs[0].text == spec.outputs[0].text
453
- )
454
-
455
- # Heuristic: expect at least 66% of prompts to match exactly
456
- assert matches >= int(0.66 * len(ref_outputs)), (
457
- f"Only {matches}/{len(ref_outputs)} outputs matched. "
458
- f"Expected at least {int(0.66 * len(ref_outputs))} matches."
459
- )
460
-
461
-
462
-def _run_eagle_correctness(
463
- monkeypatch: pytest.MonkeyPatch,
464
- sampling_config: SamplingParams,
465
- model_setup: tuple[str, str, str, int],
466
- mm_enabled: bool,
467
- expected_accuracy_threshold: float,
468
- enable_chunked_prefill: bool,
469
- model_impl: str,
470
- attn_backend: str,
471
-):
472
- """
473
- Compare the outputs of an original LLM and a speculative LLM
474
- which should be the same when using eagle speculative decoding.
475
- """
476
- if model_impl == "transformers":
477
- import transformers
478
- from packaging.version import Version
479
-
480
- installed = Version(transformers.__version__)
481
- required = Version("5.0.0")
482
- if installed < required:
483
- pytest.skip(
484
- "Eagle3 with the Transformers modeling backend requires "
485
- f"transformers>={required}, but got {installed}"
486
- )
487
-
488
- test_prompts = get_test_prompts(mm_enabled)
489
-
490
- if "Llama-4-Scout" in model_setup[1] and attn_backend == "FLASH_ATTN":
491
- if current_platform.is_rocm():
492
- print(
493
- "FLASH_ATTN for spec_decode not supported on "
494
- "ROCm currently. Changing to FLEX_ATTENTION backend."
495
- )
496
- attention_config = {"backend": "FLEX_ATTENTION"}
497
- else:
498
- attention_config = None
499
- else:
500
- attention_config = {"backend": attn_backend}
501
-
502
- if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm():
503
- pytest.skip(
504
- "TRITON_ATTN does not support "
505
- "multi-token eagle spec decode on current platform"
506
- )
507
-
508
- with monkeypatch.context() as m:
509
- m.setenv("VLLM_MLA_DISABLE", "1")
510
-
511
- if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm():
512
- if "deepseek" in model_setup[1].lower():
513
- m.setenv("VLLM_ROCM_USE_AITER", "1")
514
- m.delenv("VLLM_MLA_DISABLE", raising=False)
515
- attention_config = {"backend": "ROCM_AITER_MLA"}
516
- else:
517
- m.setenv("VLLM_ROCM_USE_AITER", "1")
518
-
519
- method, model_name, spec_model_name, tp_size = model_setup
520
- _skip_if_insufficient_gpus_for_tp(tp_size)
521
-
522
- max_model_len = 2048
523
- max_num_batched_tokens = 128 if enable_chunked_prefill else max_model_len
524
-
525
- ref_llm = LLM(
526
- model=model_name,
527
- max_model_len=max_model_len,
528
- tensor_parallel_size=tp_size,
529
- attention_config=attention_config,
530
- )
531
- evaluate_llm_for_gsm8k(
532
- ref_llm, expected_accuracy_threshold=expected_accuracy_threshold
533
- )
534
- ref_outputs = ref_llm.chat(test_prompts, sampling_config)
535
- del ref_llm
536
- torch.accelerator.empty_cache()
537
- cleanup_dist_env_and_memory()
538
- # ROCm frees VRAM lazily; wait so the spec engine started right after
539
- # does not OOM on its startup memory guard.
540
- wait_for_rocm_memory_to_settle()
541
-
542
- spec_llm = LLM(
543
- model=model_name,
544
- trust_remote_code=True,
545
- tensor_parallel_size=tp_size,
546
- speculative_config={
547
- "method": method,
548
- "model": spec_model_name,
549
- "num_speculative_tokens": 3,
550
- "max_model_len": max_model_len,
551
- },
552
- max_model_len=max_model_len,
553
- max_num_batched_tokens=max_num_batched_tokens,
554
- enable_chunked_prefill=enable_chunked_prefill,
555
- model_impl=model_impl,
556
- attention_config=attention_config,
557
- )
558
- # EAGLE/EAGLE3 supports async scheduling; assert it is active by default.
559
- assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
560
- evaluate_llm_for_gsm8k(
561
- spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
562
- )
563
- spec_outputs = spec_llm.chat(test_prompts, sampling_config)
564
- matches = 0
565
- misses = 0
566
- for ref_output, spec_output in zip(ref_outputs, spec_outputs):
567
- if ref_output.outputs[0].text == spec_output.outputs[0].text:
568
- matches += 1
569
- else:
570
- misses += 1
571
- print(f"ref_output: {ref_output.outputs[0].text}")
572
- print(f"spec_output: {spec_output.outputs[0].text}")
573
-
574
- assert matches > int(0.6 * len(ref_outputs))
575
- del spec_llm
576
- torch.accelerator.empty_cache()
577
- cleanup_dist_env_and_memory()
578
- # ROCm frees VRAM lazily; wait so the next parametrization's engine does
579
- # not OOM on its startup memory guard.
580
- wait_for_rocm_memory_to_settle()
581
-
582
-
583
-@single_gpu_only
584
-@pytest.mark.skipif(
585
- current_platform.is_device_capability_family(100),
586
- reason="DeepSeek head_dim=192 not supported on SM100/SM110 (Blackwell)",
587
-)
588
-@pytest.mark.parametrize(
589
- [
590
- "model_setup",
591
- "mm_enabled",
592
- "enable_chunked_prefill",
593
- "model_impl",
594
- "expected_accuracy_threshold",
595
- ],
596
- [
597
- (
598
- (
599
- "eagle",
600
- "eagle618/deepseek-v3-random",
601
- "eagle618/eagle-deepseek-v3-random",
602
- 1,
603
- ),
604
- False,
605
- False,
606
- "auto",
607
- 0.0,
608
- ),
609
- ],
610
- ids=["deepseek_eagle"],
611
-)
612
-@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
613
-def test_eagle_correctness_light(
614
- monkeypatch: pytest.MonkeyPatch,
615
- sampling_config: SamplingParams,
616
- model_setup: tuple[str, str, str, int],
617
- mm_enabled: bool,
618
- expected_accuracy_threshold: float,
619
- enable_chunked_prefill: bool,
620
- model_impl: str,
621
- attn_backend: str,
622
-):
623
- _run_eagle_correctness(
624
- monkeypatch,
625
- sampling_config,
626
- model_setup,
627
- mm_enabled,
628
- expected_accuracy_threshold,
629
- enable_chunked_prefill,
630
- model_impl,
631
- attn_backend,
632
- )
633
-
634
-
635
-@single_gpu_only
636
-@large_gpu_mark(min_gb=24)
637
-@pytest.mark.parametrize(
638
- [
639
- "model_setup",
640
- "mm_enabled",
641
- "enable_chunked_prefill",
642
- "model_impl",
643
- "expected_accuracy_threshold",
644
- ],
645
- [
646
- (
647
- ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
648
- False,
649
- False,
650
- "auto",
651
- 0.8,
652
- ),
653
- pytest.param(
654
- ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1),
655
- False,
656
- False,
657
- "transformers",
658
- 0.8,
659
- # TODO(hmellor): figure out why memory usage is so high
660
- marks=pytest.mark.skip(
661
- reason="Feature is experimental and uses too much memory in CI",
662
- ),
663
- ),
664
- pytest.param(
665
- (
666
- "eagle3",
667
- "Qwen/Qwen3-VL-8B-Instruct",
668
- "taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
669
- 1,
670
- ),
671
- False,
672
- False,
673
- "auto",
674
- 0.8,
675
- marks=pytest.mark.skip(
676
- reason="architecture of its eagle3 is LlamaForCausalLMEagle3"
677
- ),
678
- ),
679
- pytest.param(
680
- (
681
- "eagle3",
682
- "Qwen/Qwen2.5-VL-7B-Instruct",
683
- "Rayzl/qwen2.5-vl-7b-eagle3-sgl",
684
- 1,
685
- ),
686
- False,
687
- False,
688
- "auto",
689
- 0.7,
690
- marks=pytest.mark.skip(
691
- reason="Skipping due to its head_dim not being a multiple of 32"
692
- ),
693
- ),
694
- (
695
- (
696
- "eagle3",
697
- "meta-llama/Llama-3.1-8B-Instruct",
698
- "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
699
- 1,
700
- ),
701
- False,
702
- False,
703
- "auto",
704
- 0.7,
705
- ),
706
- ],
707
- ids=[
708
- "qwen3_eagle3",
709
- "qwen3_eagle3-transformers",
710
- "qwen3_vl_eagle3",
711
- "qwen2_5_vl_eagle3",
712
- "llama3_eagle3",
713
- ],
714
-)
715
-@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
716
-def test_eagle_correctness_medium(
717
- monkeypatch: pytest.MonkeyPatch,
718
- sampling_config: SamplingParams,
719
- model_setup: tuple[str, str, str, int],
720
- mm_enabled: bool,
721
- expected_accuracy_threshold: float,
722
- enable_chunked_prefill: bool,
723
- model_impl: str,
724
- attn_backend: str,
725
-):
726
- _run_eagle_correctness(
727
- monkeypatch,
728
- sampling_config,
729
- model_setup,
730
- mm_enabled,
731
- expected_accuracy_threshold,
732
- enable_chunked_prefill,
733
- model_impl,
734
- attn_backend,
735
- )
736
-
737
-
738
-@pytest.mark.parametrize(
739
- [
740
- "model_setup",
741
- "mm_enabled",
742
- "enable_chunked_prefill",
743
- "model_impl",
744
- "expected_accuracy_threshold",
745
- ],
746
- [
747
- pytest.param(
748
- (
749
- "eagle",
750
- "meta-llama/Llama-3.1-8B-Instruct",
751
- "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
752
- 1,
753
- ),
754
- False,
755
- True,
756
- "auto",
757
- 0.7,
758
- marks=large_gpu_mark(min_gb=40),
759
- id="llama3_eagle",
760
- ),
761
- pytest.param(
762
- (
763
- "eagle",
764
- "meta-llama/Llama-4-Scout-17B-16E-Instruct",
765
- "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
766
- 4,
767
- ),
768
- False,
769
- False,
770
- "auto",
771
- 0.8,
772
- marks=[*multi_gpu_marks(num_gpus=4), large_gpu_mark(min_gb=40)],
773
- id="llama4_eagle",
774
- ),
775
- pytest.param(
776
- (
777
- "eagle",
778
- "meta-llama/Llama-4-Scout-17B-16E-Instruct",
779
- "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
780
- 4,
781
- ),
782
- True,
783
- True,
784
- "auto",
785
- 0.8,
786
- marks=[*multi_gpu_marks(num_gpus=4), large_gpu_mark(min_gb=80)],
787
- id="llama4_eagle_mm",
788
- ),
789
- ],
790
-)
791
-@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())
792
-def test_eagle_correctness_heavy(
793
- monkeypatch: pytest.MonkeyPatch,
794
- sampling_config: SamplingParams,
795
- model_setup: tuple[str, str, str, int],
796
- mm_enabled: bool,
797
- expected_accuracy_threshold: float,
798
- enable_chunked_prefill: bool,
799
- model_impl: str,
800
- attn_backend: str,
801
-):
802
- _run_eagle_correctness(
803
- monkeypatch,
804
- sampling_config,
805
- model_setup,
806
- mm_enabled,
807
- expected_accuracy_threshold,
808
- enable_chunked_prefill,
809
- model_impl,
810
- attn_backend,
811
- )
812
-
813
-
814
-@large_gpu_mark(min_gb=24)
815
-def test_medusa_acceptance_rate(
816
- sampling_config: SamplingParams,
817
-):
818
- """Verify a trained Medusa checkpoint achieves nonzero acceptance rate.
819
-
820
- Uses the canonical FasterDecoding vicuna-7b checkpoint to confirm the
821
- speculation path actually accepts tokens — unlike test_medusa_correctness,
822
- which uses a random head and only validates output correctness.
823
- """
824
- target_model = "lmsys/vicuna-7b-v1.3"
825
- medusa_model = "FasterDecoding/medusa-vicuna-7b-v1.3"
826
- prompts = _build_gsm8k_prompts(num_questions=10, num_shots=1)[0]
827
-
828
- spec_llm = LLM(
829
- model=target_model,
830
- speculative_config={
831
- "method": "medusa",
832
- "model": medusa_model,
833
- "num_speculative_tokens": 3,
834
- },
835
- max_model_len=1024,
836
- enforce_eager=True,
837
- disable_log_stats=False,
838
- )
839
- spec_llm.generate(prompts, sampling_config)
840
- metrics = spec_llm.get_metrics()
841
- acceptance_rate = compute_acceptance_rate(metrics)
842
- del spec_llm
843
- torch.accelerator.empty_cache()
844
- cleanup_dist_env_and_memory()
845
-
846
- min_acceptance_rate = 0.198
847
- print(f"Medusa acceptance rate: {acceptance_rate:.4f} (min {min_acceptance_rate})")
848
-
849
- # Regression guard at 90% of the measured baseline.
850
- assert acceptance_rate >= min_acceptance_rate, (
851
- f"Medusa acceptance rate {acceptance_rate:.4f} below min {min_acceptance_rate}"
852
- )
853
-
854
-
855
-@pytest.mark.parametrize(
856
- ["model_setup", "mm_enabled", "expected_accuracy_threshold"],
857
- [
858
- (("mtp", "XiaomiMiMo/MiMo-7B-Base", 1), False, 0.5), # ref: 65%-70%
859
- pytest.param(
860
- ("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1),
861
- False,
862
- 0.0,
863
- marks=pytest.mark.skipif(
864
- current_platform.is_device_capability_family(100),
865
- reason="DeepSeek MTP: TRTLLM MoE top_k check fails on Blackwell",
866
- ),
867
- ), # dummy model
868
- (
869
- ("mtp", "Qwen/Qwen3.5-0.8B-Base", 1),
870
- False,
871
- 0.20,
872
- ), # hybrid + MTP, ref: ~34%-35%
873
- (
874
- ("mtp", "google/gemma-4-E4B-it", 1, "google/gemma-4-E4B-it-assistant"),
875
- False,
876
- 0.50,
877
- ), # gemma4 MTP with assistant model, ref: ~62%
878
- ],
879
- ids=["mimo", "deepseek", "qwen3_5-hybrid", "gemma4-e4b"],
880
-)
881
-@single_gpu_only
882
-@large_gpu_mark(min_gb=20)
883
-def test_mtp_correctness(
884
- monkeypatch: pytest.MonkeyPatch,
885
- sampling_config: SamplingParams,
886
- model_setup: tuple[str, str, int] | tuple[str, str, int, str],
887
- mm_enabled: bool,
888
- expected_accuracy_threshold: float,
889
-):
890
- """
891
- Compare the outputs of a original LLM and a speculative LLM
892
- which should be the same when using MTP speculative decoding. Due to some variance
893
- in the engine, it is possible for some outputs to differ, so we expect that at least
894
- 6/10 output tokens match exactly, and that the GSM8k accuracy is above a precomputed
895
- reference threshold for each model.
896
- """
897
- # Generate test prompts inside the function instead of using fixture
898
- test_prompts = get_test_prompts(mm_enabled)
899
- with monkeypatch.context() as m:
900
- m.setenv("VLLM_MLA_DISABLE", "1")
901
-
902
- if len(model_setup) == 4:
903
- method, model_name, tp_size, draft_model = model_setup
904
- else:
905
- method, model_name, tp_size = model_setup
906
- draft_model = None
907
- _skip_if_insufficient_gpus_for_tp(tp_size)
908
-
909
- if "Qwen3.5" in model_name and os.environ.get("VLLM_USE_V2_MODEL_RUNNER"):
910
- pytest.skip(
911
- "Model Runner V2 does not yet support hybrid models "
912
- "(Qwen3.5 mixes Mamba-style GDN with attention layers)."
913
- )
914
-
915
- attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
916
-
917
- # Skip multimodal profiling for models that don't need it in this test.
918
- extra_kwargs: dict[str, Any] = {}
919
- if "Qwen3.5" in model_name:
920
- extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0}
921
- elif "gemma-4" in model_name:
922
- extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "audio": 0}
923
-
924
- if draft_model is not None and "gemma-4" in draft_model:
925
- import transformers
926
- from packaging.version import Version
927
-
928
- if Version(transformers.__version__) < Version("5.8.0"):
929
- pytest.skip(
930
- "Gemma4 MTP assistant requires transformers>=5.8.0, "
931
- f"got {transformers.__version__}"
932
- )
933
-
934
- ref_llm = LLM(
935
- model=model_name,
936
- max_model_len=2048,
937
- tensor_parallel_size=tp_size,
938
- trust_remote_code=True,
939
- attention_backend=attn_backend,
940
- **extra_kwargs,
941
- )
942
- ref_outputs = ref_llm.chat(test_prompts, sampling_config)
943
- evaluate_llm_for_gsm8k(
944
- ref_llm, expected_accuracy_threshold=expected_accuracy_threshold
945
- )
946
- del ref_llm
947
- torch.accelerator.empty_cache()
948
- cleanup_dist_env_and_memory()
949
-
950
- speculative_config: dict[str, Any] = {
951
- "method": method,
952
- "num_speculative_tokens": 1,
953
- "max_model_len": 2048,
954
- }
955
- if draft_model is not None:
956
- speculative_config["model"] = draft_model
957
- speculative_config["num_speculative_tokens"] = 2
958
-
959
- spec_llm = LLM(
960
- model=model_name,
961
- trust_remote_code=True,
962
- tensor_parallel_size=tp_size,
963
- speculative_config=speculative_config,
964
- max_model_len=2048,
965
- attention_backend=attn_backend,
966
- **extra_kwargs,
967
- )
968
- # MTP supports async scheduling; assert it is active by default.
969
- assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
970
- evaluate_llm_for_gsm8k(
971
- spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
972
- )
973
- spec_outputs = spec_llm.chat(test_prompts, sampling_config)
974
- matches = 0
975
- misses = 0
976
- for ref_output, spec_output in zip(ref_outputs, spec_outputs):
977
- if ref_output.outputs[0].text == spec_output.outputs[0].text:
978
- matches += 1
979
- else:
980
- misses += 1
981
- print(f"ref_output: {ref_output.outputs[0].text}")
982
- print(f"spec_output: {spec_output.outputs[0].text}")
983
-
984
- # Heuristic: expect at least 80% of the prompts to match exactly
985
- # Upon failure, inspect the outputs to check for inaccuracy.
986
- assert matches > int(MTP_SIMILARITY_RATE * len(ref_outputs))
987
- del spec_llm
988
- torch.accelerator.empty_cache()
989
- cleanup_dist_env_and_memory()
990
-
991
-
992
-@dataclass
993
-class ArgsTest:
994
- target_model: str
995
- draft_model: str
996
- sampling_config: SamplingParams
997
- num_speculative_tokens: int
998
- expected_acceptance_rate: float
999
- expected_acceptance_len: float
1000
- expected_gsm8k_accuracy: float = 0.0 # skip by default
1001
- # Defaults
1002
- enforce_eager: bool = True
1003
- parallel_drafting: bool = False
1004
- target_tensor_parallel_size: int = 1
1005
- draft_tensor_parallel_size: int = 1
1006
- max_model_len: int = 2048
1007
- gpu_memory_utilization: float = 0.5
1008
- dataset: str = "test_prompts"
1009
- num_prompts: int = 100
1010
-
1011
-
1012
-cases = [
1013
- # Same model for draft and target, greedy sampling.
1014
- ArgsTest(
1015
- target_model="Qwen/Qwen3-0.6B",
1016
- draft_model="Qwen/Qwen3-0.6B",
1017
- sampling_config=greedy_sampling(),
1018
- num_speculative_tokens=3, # K
1019
- expected_acceptance_len=0.98 * (3 + 1), # epsilon discount of K + 1
1020
- expected_acceptance_rate=0.98, # slight epsilon
1021
- expected_gsm8k_accuracy=0.25, # ref: 35-40%
1022
- ),
1023
- # Smaller draft model, stochastic sampling.
1024
- ArgsTest(
1025
- target_model="Qwen/Qwen3-1.7B",
1026
- draft_model="Qwen/Qwen3-0.6B",
1027
- sampling_config=stochastic_sampling(),
1028
- num_speculative_tokens=3,
1029
- expected_acceptance_len=3.4, # ref: 3.7
1030
- expected_acceptance_rate=0.80, # ref: 0.90
1031
- expected_gsm8k_accuracy=0.5, # ref: 60%. Note gsm8k always runs greedy sampling
1032
- ),
1033
-]
1034
-
1035
-
1036
-@pytest.mark.parametrize("args", cases)
1037
-@pytest.mark.parametrize("enforce_eager", [True, False])
1038
-@single_gpu_only
1039
-# TODO: Fix async_scheduling & engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
1040
-@pytest.mark.xfail(
1041
- raises=AsyncSchedulingNotEnabledError,
1042
- reason="draft_model does not yet enable async_scheduling: issue #38929",
1043
-)
1044
-def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
1045
- args.enforce_eager = enforce_eager
1046
- assert_draft_model_correctness(args)
1047
-
1048
-
1049
-@single_gpu_only
1050
-# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
1051
-@pytest.mark.xfail(
1052
- raises=AsyncSchedulingNotEnabledError,
1053
- reason="draft_model does not yet enable async_scheduling: issue #38929",
1054
-)
1055
-def test_draft_model_realistic_example():
1056
- args = ArgsTest(
1057
- target_model="Qwen/Qwen3-1.7B",
1058
- draft_model="Qwen/Qwen3-0.6B",
1059
- dataset="likaixin/InstructCoder",
1060
- num_speculative_tokens=3,
1061
- sampling_config=greedy_sampling(),
1062
- enforce_eager=False,
1063
- expected_acceptance_len=2.6, # ref: 2.86
1064
- expected_acceptance_rate=0.5, # ref: 0.62
1065
- )
1066
- assert_draft_model_correctness(args)
1067
-
1068
-
1069
-@single_gpu_only
1070
-# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
1071
-@pytest.mark.xfail(
1072
- raises=AsyncSchedulingNotEnabledError,
1073
- reason="draft_model does not yet enable async_scheduling: issue #38929",
1074
-)
1075
-def test_draft_model_parallel_drafting():
1076
- args = ArgsTest(
1077
- target_model="Qwen/Qwen3-1.7B",
1078
- draft_model="amd/PARD-Qwen3-0.6B",
1079
- dataset="likaixin/InstructCoder",
1080
- num_speculative_tokens=3,
1081
- sampling_config=greedy_sampling(),
1082
- parallel_drafting=True,
1083
- enforce_eager=False,
1084
- expected_acceptance_len=2.3, # ref: 2.52
1085
- expected_acceptance_rate=0.4, # ref: 0.51
1086
- )
1087
- assert_draft_model_correctness(args)
1088
-
1089
-
1090
-@pytest.mark.parametrize(
1091
- "models",
1092
- [
1093
- # target_model, draft_model
1094
- ("Qwen/Qwen3-1.7B-FP8", "Qwen/Qwen3-0.6B"), # target quantized
1095
- ("Qwen/Qwen3-1.7B", "Qwen/Qwen3-0.6B-FP8"), # draft quantized
1096
- ],
1097
- ids=["target_quantized", "draft_quantized"],
1098
-)
1099
-@pytest.mark.parametrize("enforce_eager", [True, False])
1100
-@single_gpu_only
1101
-# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
1102
-@pytest.mark.xfail(
1103
- raises=AsyncSchedulingNotEnabledError,
1104
- reason="draft_model does not yet enable async_scheduling: issue #38929",
1105
-)
1106
-def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool):
1107
- tgt_model, draft_model = models
1108
- sd_case = ArgsTest(
1109
- target_model=tgt_model,
1110
- draft_model=draft_model,
1111
- **some_high_acceptance_metrics(),
1112
- enforce_eager=enforce_eager,
1113
- )
1114
- assert_draft_model_correctness(sd_case)
1115
-
1116
-
1117
-@multi_gpu_only(num_gpus=2)
1118
-# TODO: Fix async_scheduling and engine initialization issues - see https://github.com/vllm-project/vllm/issues/38929
1119
-@pytest.mark.xfail(
1120
- raises=AsyncSchedulingNotEnabledError,
1121
- reason="draft_model does not yet enable async_scheduling: issue #38929",
1122
-)
1123
-def test_draft_model_tensor_parallelism():
1124
- """Ensure spec decode works when running with TP > 1."""
1125
- _skip_if_insufficient_gpus_for_tp(2)
1126
- sd_case = ArgsTest(
1127
- target_model="Qwen/Qwen3-1.7B",
1128
- target_tensor_parallel_size=2,
1129
- draft_model="Qwen/Qwen3-0.6B",
1130
- draft_tensor_parallel_size=2,
1131
- **some_high_acceptance_metrics(),
1132
- enforce_eager=False,
1133
- expected_gsm8k_accuracy=0.5,
1134
- )
1135
- assert_draft_model_correctness(sd_case)
1136
-
1137
-
1138
-@multi_gpu_only(num_gpus=2)
1139
-def test_draft_model_engine_args_tensor_parallelism():
1140
- """Ensure the vllm_config for the draft model is created correctly,
1141
- and independently of the target model (quantization, TP, etc.)"""
1142
- _skip_if_insufficient_gpus_for_tp(2)
1143
-
1144
- engine_args = EngineArgs(
1145
- model="Qwen/Qwen3-1.7B-FP8", # <<< tgt quantized
1146
- tensor_parallel_size=2,
1147
- speculative_config={
1148
- "model": "Qwen/Qwen3-0.6B", # <<< draft not quantized
1149
- "method": "draft_model",
1150
- "num_speculative_tokens": 3,
1151
- "draft_tensor_parallel_size": 1, # <<< valid arg name
1152
- },
1153
- )
1154
- target_config: VllmConfig = engine_args.create_engine_config()
1155
- assert target_config.parallel_config.tensor_parallel_size == 2
1156
- assert target_config.quant_config.get_name() == "fp8"
1157
-
1158
- speculative_config = target_config.speculative_config
1159
- draft_config: VllmConfig = replace(
1160
- target_config,
1161
- quant_config=None,
1162
- parallel_config=replace(
1163
- speculative_config.draft_parallel_config,
1164
- rank=target_config.parallel_config.rank,
1165
- ),
1166
- model_config=speculative_config.draft_model_config,
1167
- )
1168
- assert draft_config.parallel_config.tensor_parallel_size == 1
1169
- assert draft_config.quant_config is None
1170
-
1171
-
1172
-def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig:
1173
- """Replicate SpecDecodeBaseProposer._create_draft_vllm_config logic
1174
- so we can test it without instantiating a full proposer."""
1175
- spec_cfg = vllm_config.speculative_config
1176
- if spec_cfg.moe_backend is not None:
1177
- return replace(
1178
- vllm_config,
1179
- kernel_config=replace(
1180
- vllm_config.kernel_config,
1181
- moe_backend=spec_cfg.moe_backend,
1182
- ),
1183
- )
1184
- return vllm_config
1185
-
1186
-
1187
-def test_draft_model_moe_backend_override():
1188
- """When moe_backend is set in speculative_config, the draft VllmConfig
1189
- should use it while the target keeps its own setting."""
1190
- engine_args = EngineArgs(
1191
- model="Qwen/Qwen3-1.7B",
1192
- tensor_parallel_size=1,
1193
- moe_backend="flashinfer_trtllm",
1194
- speculative_config={
1195
- "model": "Qwen/Qwen3-0.6B",
1196
- "method": "draft_model",
1197
- "num_speculative_tokens": 3,
1198
- "moe_backend": "triton",
1199
- },
1200
- )
1201
- tgt_config: VllmConfig = engine_args.create_engine_config()
1202
- assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm"
1203
- assert tgt_config.speculative_config.moe_backend == "triton"
1204
-
1205
- draft_config = _apply_draft_moe_backend(tgt_config)
1206
- assert draft_config.kernel_config.moe_backend == "triton"
1207
- # Target config must be unaffected.
1208
- assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm"
1209
-
1210
-
1211
-def test_draft_model_moe_backend_inherits_target():
1212
- """When moe_backend is not set in speculative_config, the draft should
1213
- inherit the target's moe_backend."""
1214
- engine_args = EngineArgs(
1215
- model="Qwen/Qwen3-1.7B",
1216
- tensor_parallel_size=1,
1217
- moe_backend="flashinfer_cutlass",
1218
- speculative_config={
1219
- "model": "Qwen/Qwen3-0.6B",
1220
- "method": "draft_model",
1221
- "num_speculative_tokens": 3,
1222
- },
1223
- )
1224
- tgt_config: VllmConfig = engine_args.create_engine_config()
1225
- assert tgt_config.kernel_config.moe_backend == "flashinfer_cutlass"
1226
- assert tgt_config.speculative_config.moe_backend is None
1227
-
1228
- draft_config = _apply_draft_moe_backend(tgt_config)
1229
- assert draft_config.kernel_config.moe_backend == "flashinfer_cutlass"
1230
- assert draft_config is tgt_config
1231
-
1232
-
1233
-def test_draft_model_moe_backend_default_auto():
1234
- """When neither target nor draft set moe_backend explicitly, both should
1235
- default to 'auto'."""
1236
- engine_args = EngineArgs(
1237
- model="Qwen/Qwen3-1.7B",
1238
- tensor_parallel_size=1,
1239
- speculative_config={
1240
- "model": "Qwen/Qwen3-0.6B",
1241
- "method": "draft_model",
1242
- "num_speculative_tokens": 3,
1243
- },
1244
- )
1245
- tgt_config: VllmConfig = engine_args.create_engine_config()
1246
- assert tgt_config.kernel_config.moe_backend == "auto"
1247
- assert tgt_config.speculative_config.moe_backend is None
1248
-
1249
- draft_config = _apply_draft_moe_backend(tgt_config)
1250
- assert draft_config.kernel_config.moe_backend == "auto"
1251
- assert draft_config is tgt_config
1252
-
1253
-
1254
-def test_draft_model_engine_args_rejects_invalid_tp_argname():
1255
- """The user should pass "draft_tensor_parallel_size" rather than
1256
- "tensor_parallel_size". We enforce this with validation."""
1257
-
1258
- engine_args = EngineArgs(
1259
- model="Qwen/Qwen3-1.7B",
1260
- tensor_parallel_size=1,
1261
- speculative_config={
1262
- "model": "Qwen/Qwen3-0.6B",
1263
- "method": "draft_model",
1264
- "num_speculative_tokens": 3,
1265
- "tensor_parallel_size": 1, # <<< invalid arg name
1266
- },
1267
- )
1268
- with pytest.raises(ValueError):
1269
- engine_args.create_engine_config()
1270
-
1271
-
1272
-def assert_draft_model_correctness(args: ArgsTest):
1273
- """Compare the outputs using and not using speculative decoding.
1274
- In the greedy decoding case, the outputs must match EXACTLY."""
1275
- test_prompts: list[Messages] = get_messages(
1276
- dataset=args.dataset, n=args.num_prompts
1277
- )
1278
-
1279
- spec_llm = LLM(
1280
- model=args.target_model,
1281
- speculative_config={
1282
- "model": args.draft_model,
1283
- "method": "draft_model",
1284
- "num_speculative_tokens": args.num_speculative_tokens,
1285
- "max_model_len": args.max_model_len,
1286
- "enforce_eager": args.enforce_eager,
1287
- "draft_tensor_parallel_size": args.draft_tensor_parallel_size,
1288
- "parallel_drafting": args.parallel_drafting,
1289
- },
1290
- max_num_seqs=100, # limit cudagraph capture runtime
1291
- max_model_len=args.max_model_len,
1292
- gpu_memory_utilization=args.gpu_memory_utilization,
1293
- tensor_parallel_size=args.target_tensor_parallel_size,
1294
- enforce_eager=args.enforce_eager,
1295
- disable_log_stats=False, # enables get_metrics()
1296
- )
1297
-
1298
- # we don't check the outputs, only check the metrics
1299
- spec_llm.chat(test_prompts, args.sampling_config)
1300
- metrics = spec_llm.get_metrics()
1301
- acceptance_rate: float = compute_acceptance_rate(metrics)
1302
- acceptance_len: float = compute_acceptance_len(metrics)
1303
-
1304
- # Need to evaluate after getting metrics to avoid polluting the AR
1305
- evaluate_llm_for_gsm8k(
1306
- spec_llm, expected_accuracy_threshold=args.expected_gsm8k_accuracy
1307
- )
1308
-
1309
- print(
1310
- f"spec-decode: target={args.target_model}, draft={args.draft_model}, "
1311
- f"temperature={args.sampling_config.temperature:.2f}, "
1312
- f"acceptance_rate={acceptance_rate:.2f}, "
1313
- f"acceptance_len={acceptance_len:.2f}, "
1314
- )
1315
-
1316
- assert acceptance_rate >= args.expected_acceptance_rate
1317
- assert acceptance_len >= args.expected_acceptance_len
1318
- # draft_model supports async scheduling; assert it is active by default.
1319
- # Raise AsyncSchedulingNotEnabledError (a subclass of AssertionError) so that
1320
- # @pytest.mark.xfail(raises=AsyncSchedulingNotEnabledError) catches only this
1321
- # specific failure — leaving all other assertion failures (e.g. correctness or
1322
- # acceptance-rate checks above) visible as real test failures.
1323
- has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
1324
- del spec_llm # CLEANUP
1325
- torch.accelerator.empty_cache()
1326
- cleanup_dist_env_and_memory()
1327
- if not has_async:
1328
- raise AsyncSchedulingNotEnabledError(
1329
- "Expected async_scheduling=True for draft_model spec decode, got False."
1330
- " See https://github.com/vllm-project/vllm/issues/38929"
1331
- )
1332
-
1333
-
1334
-def get_messages(dataset: str, n: int) -> list[Messages]:
1335
- if dataset == "test_prompts":
1336
- return get_test_prompts(mm_enabled=False, num_prompts=n)
1337
- elif dataset == "likaixin/InstructCoder":
1338
- return get_instruct_coder_messages(n=n)
1339
- else:
1340
- raise NotImplementedError(f"Dataset '{dataset}' not implemented")
1341
-
1342
-
1343
-def some_high_acceptance_metrics() -> dict:
1344
- return {
1345
- "sampling_config": greedy_sampling(),
1346
- "num_speculative_tokens": 3,
1347
- "expected_acceptance_len": 3.4, # ref: 3.75
1348
- "expected_acceptance_rate": 0.8, # ref: 0.9
1349
- }
1350
-
1351
-
1352
-def compute_acceptance_rate(
1353
- metrics: list[Metric], prev_metrics: list[Metric] | None = None
1354
-) -> float:
1355
- name2metric = {metric.name: metric for metric in metrics}
1356
- n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value
1357
- if n_draft_toks == 0:
1358
- return float("nan")
1359
- n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value
1360
- if prev_metrics is not None:
1361
- prev_name2metric = {metric.name: metric for metric in prev_metrics}
1362
- n_draft_toks -= prev_name2metric["vllm:spec_decode_num_draft_tokens"].value
1363
- n_accepted_toks -= prev_name2metric[
1364
- "vllm:spec_decode_num_accepted_tokens"
1365
- ].value
1366
- if n_draft_toks <= 0:
1367
- return float("nan")
1368
- return n_accepted_toks / n_draft_toks
1369
-
1370
-
1371
-def compute_acceptance_len(
1372
- metrics: list[Metric], prev_metrics: list[Metric] | None = None
1373
-) -> float:
1374
- name2metric = {metric.name: metric for metric in metrics}
1375
- n_drafts = name2metric["vllm:spec_decode_num_drafts"].value
1376
- n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value
1377
- if n_drafts == 0:
1378
- return 1
1379
- if prev_metrics is not None:
1380
- prev_name2metric = {metric.name: metric for metric in prev_metrics}
1381
- n_drafts -= prev_name2metric["vllm:spec_decode_num_drafts"].value
1382
- n_accepted_toks -= prev_name2metric[
1383
- "vllm:spec_decode_num_accepted_tokens"
1384
- ].value
1385
- if n_drafts <= 0:
1386
- return 1
1387
- return 1 + (n_accepted_toks / n_drafts)
1388
-
1389
-
1390
-# Datasets in the format used in DFlash validations
1391
-def load_and_process_dataset(data_name: str):
1392
- from datasets import load_dataset
1393
-
1394
- if data_name == "gsm8k":
1395
- dataset = load_dataset("openai/gsm8k", "main", split="test")
1396
- prompt_fmt = (
1397
- "{question}\nPlease reason step by step,"
1398
- " and put your final answer within \\boxed{{}}."
1399
- )
1400
- dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
1401
- elif data_name == "mt-bench":
1402
- dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train")
1403
- dataset = dataset.map(lambda x: {"turns": x["prompt"]})
1404
- elif data_name == "humaneval":
1405
- dataset = load_dataset("openai/openai_humaneval", split="test")
1406
- prompt_fmt = (
1407
- "Write a solution to the following problem and make sure"
1408
- " that it passes the tests:\n```python\n{prompt}\n```"
1409
- )
1410
- dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]})
1411
-
1412
- return dataset
1413
-
1414
-
1415
-@pytest.mark.parametrize(
1416
- ["spec_config", "expected_acceptance_lengths", "chat_template_kwargs"],
1417
- [
1418
- pytest.param(
1419
- dict(
1420
- model="Qwen/Qwen3-8B",
1421
- trust_remote_code=True,
1422
- speculative_config={
1423
- "method": "dflash",
1424
- "model": "z-lab/Qwen3-8B-DFlash-b16",
1425
- "num_speculative_tokens": 16,
1426
- "max_model_len": 32768,
1427
- },
1428
- max_model_len=32768,
1429
- max_num_seqs=128,
1430
- gpu_memory_utilization=0.85,
1431
- enforce_eager=False,
1432
- disable_log_stats=False,
1433
- ),
1434
- # All scores from Table 1 in https://arxiv.org/pdf/2602.06036
1435
- {
1436
- "mt-bench": 4.24,
1437
- "humaneval": 6.50,
1438
- # runs with a subset of prompts so extra wide tol here
1439
- "gsm8k": 6.54 * 0.975,
1440
- },
1441
- {"enable_thinking": False},
1442
- id="dflash",
1443
- ),
1444
- pytest.param(
1445
- dict(
1446
- model="google/gemma-4-E4B-it",
1447
- trust_remote_code=True,
1448
- speculative_config={
1449
- "method": "mtp",
1450
- "model": "google/gemma-4-E4B-it-assistant",
1451
- "num_speculative_tokens": 2,
1452
- "max_model_len": 32768,
1453
- },
1454
- max_model_len=32768,
1455
- # Skip multimodal profiling; this is a text-only eval.
1456
- limit_mm_per_prompt={"image": 0, "audio": 0},
1457
- disable_log_stats=False,
1458
- ),
1459
- {
1460
- "mt-bench": 2.28,
1461
- "humaneval": 2.68,
1462
- # runs with a subset of prompts so extra wide tol here
1463
- "gsm8k": 2.67 * 0.975,
1464
- },
1465
- {},
1466
- id="gemma4",
1467
- ),
1468
- ],
1469
-)
1470
-@pytest.mark.parametrize("use_mrv2", [False, True])
1471
-def test_acceptance_rates(
1472
- monkeypatch: pytest.MonkeyPatch,
1473
- spec_config: dict[str, Any],
1474
- expected_acceptance_lengths: dict[str, float],
1475
- chat_template_kwargs: dict[str, Any],
1476
- use_mrv2: bool,
1477
-):
1478
- """
1479
- E2E acceptance-rate validation for speculative decoding.
1480
-
1481
- Drives one or more datasets (keyed in ``expected_acceptance_lengths``)
1482
- through the spec decode engine and asserts the mean acceptance length
1483
- stays within tolerance of the reference figure for each dataset.
1484
- """
1485
- monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0")
1486
- spec_llm = LLM(**spec_config)
1487
-
1488
- max_prompts_per_dataset = 200 # mt-bench has 80, humaneval has 164, truncates gsm8k
1489
-
1490
- tokenizer = spec_llm.get_tokenizer()
1491
- for dataset_name, expected_len in expected_acceptance_lengths.items():
1492
- dataset = load_and_process_dataset(dataset_name)
1493
- prev_metrics = None
1494
- acceptance_lengths = []
1495
- for i in tqdm(
1496
- range(min(max_prompts_per_dataset, len(dataset))),
1497
- desc=f"Processing {dataset_name}",
1498
- ):
1499
- user_content = dataset[i]["turns"][0]
1500
- prompt_text = tokenizer.apply_chat_template(
1501
- [{"role": "user", "content": user_content}],
1502
- tokenize=False,
1503
- add_generation_prompt=True,
1504
- **chat_template_kwargs,
1505
- )
1506
-
1507
- # Greedy (temp=0) so acceptance length is deterministic and comparable
1508
- # across runs.
1509
- spec_llm.generate(
1510
- [prompt_text],
1511
- SamplingParams(temperature=0, max_tokens=2048),
1512
- use_tqdm=False,
1513
- )
1514
- current_metrics = spec_llm.get_metrics()
1515
- acceptance_len = compute_acceptance_len(current_metrics, prev_metrics)
1516
- prev_metrics = current_metrics
1517
- acceptance_lengths.append(acceptance_len)
1518
-
1519
- mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths)
1520
- # Fairly tight tolerance of 95% against the reference figures,
1521
- # watching for regressions. Can be relaxed if test is flaky but be sure to
1522
- # check for genuine issues such as #40727.
1523
- expected_len = expected_len * 0.95
1524
- print(
1525
- f"acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}"
1526
- f" (expected at least {expected_len:.2f})"
1527
- )
1528
-
1529
- assert mean_acceptance_length >= expected_len, (
1530
- f"acceptance_len for {dataset_name} is below expected threshold: "
1531
- f"{mean_acceptance_length:.2f} < {expected_len:.2f}"
1532
- )
1533
-
1534
- del spec_llm
1535
- torch.accelerator.empty_cache()
1536
- cleanup_dist_env_and_memory()
1537
-
1538
-
1539
-@pytest.fixture
1540
-def dspark_config():
1541
- target_model = "Qwen/Qwen3-4B-FP8"
1542
- draft_model = "deepseek-ai/dspark_qwen3_4b_block7"
1543
-
1544
- return dict(
1545
- model=target_model,
1546
- trust_remote_code=True,
1547
- speculative_config={
1548
- "method": "dspark",
1549
- "model": draft_model,
1550
- "num_speculative_tokens": 7,
1551
- "attention_backend": "FLASH_ATTN",
1552
- "draft_sample_method": "probabilistic",
1553
- },
1554
- max_model_len=4096,
1555
- disable_log_stats=False,
1556
- )
1557
-
1558
-
1559
-@single_gpu_only
1560
-@large_gpu_mark(min_gb=24)
1561
-def test_dspark_correctness_and_acceptance_rate(dspark_config):
1562
- """
1563
- E2E test for DSpark speculative decoding: acceptance rate/length
1564
- regression coverage plus GSM8K correctness, at temperature=1.0 to
1565
- exercise the probabilistic draft-sampling/rejection-sampling path
1566
- (not just greedy).
1567
-
1568
- Uses Qwen/Qwen3-4B-FP8 as target with the dspark_qwen3_4b_block7 draft
1569
- model. Reference: measured over 12 runs of the full GSM8K set at
1570
- temperature=1.0 (prefix caching disabled to avoid cross-run reuse):
1571
- accuracy: min=0.782 max=0.814 mean=0.801
1572
- acceptance_rate: min=0.418 max=0.434 mean=0.428
1573
- acceptance_len: min=3.928 max=4.037 mean=3.994
1574
- Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling
1575
- """
1576
- spec_llm = LLM(**dspark_config)
1577
-
1578
- results = evaluate_gsm8k_offline(spec_llm, temperature=1.0)
1579
- gsm8k_accuracy = results["accuracy"]
1580
-
1581
- metrics = spec_llm.get_metrics()
1582
- acceptance_rate = compute_acceptance_rate(metrics)
1583
- acceptance_len = compute_acceptance_len(metrics)
1584
-
1585
- print(
1586
- f"DSpark acceptance_rate={acceptance_rate:.2f}, "
1587
- f"acceptance_len={acceptance_len:.2f}, "
1588
- f"gsm8k_accuracy={gsm8k_accuracy:.3f}"
1589
- )
1590
-
1591
- assert acceptance_rate >= 0.428 * 0.9
1592
- assert acceptance_len >= 3.994 * 0.9
1593
- assert gsm8k_accuracy >= 0.801 * 0.9
1594
-
1595
- del spec_llm
1596
- torch.accelerator.empty_cache()
1597
- cleanup_dist_env_and_memory()
1598
-
1599
-
1600
-@single_gpu_only
1601
-def test_synthetic_acceptance_rate():
1602
- """Verify that synthetic rejection sampling produces an acceptance
1603
- length close to the requested mean acceptance length."""
1604
- num_spec_tokens = 3
1605
- expected_acceptance_len = 1.875
1606
- tolerance = 0.15
1607
-
1608
- spec_llm = LLM(
1609
- model="meta-llama/Llama-3.2-1B-Instruct",
1610
- trust_remote_code=True,
1611
- speculative_config={
1612
- "method": "eagle3",
1613
- "model": "nm-testing/Llama3_2_1B_speculator.eagle3",
1614
- "num_speculative_tokens": num_spec_tokens,
1615
- "max_model_len": 2048,
1616
- "rejection_sample_method": "synthetic",
1617
- "synthetic_acceptance_length": expected_acceptance_len,
1618
- },
1619
- max_model_len=2048,
1620
- enforce_eager=True,
1621
- disable_log_stats=False,
1622
- )
1623
-
1624
- test_prompts = get_test_prompts(mm_enabled=False, num_prompts=50)
1625
- spec_llm.chat(
1626
- test_prompts,
1627
- SamplingParams(temperature=0, max_tokens=64, ignore_eos=True),
1628
- )
1629
-
1630
- metrics = spec_llm.get_metrics()
1631
- acceptance_len = compute_acceptance_len(metrics)
1632
-
1633
- print(
1634
- f"Synthetic acceptance length: {acceptance_len:.3f}"
1635
- f" (expected={expected_acceptance_len:.3f},"
1636
- f" tolerance=±{tolerance})"
1637
- )
1638
- assert abs(acceptance_len - expected_acceptance_len) <= tolerance, (
1639
- f"Synthetic acceptance length {acceptance_len:.3f} is not within"
1640
- f" ±{tolerance} of expected {expected_acceptance_len:.3f}"
1641
- )
1642
-
1643
- del spec_llm
1644
- torch.accelerator.empty_cache()
1645
- cleanup_dist_env_and_memory()
1646
-
1647
-
1648
-@pytest.mark.parametrize("use_mrv2", [False, True])
1649
-def test_dflash_correctness(
1650
- monkeypatch: pytest.MonkeyPatch, use_mrv2: bool, dflash_config
1651
-):
1652
- """
1653
- E2E test for DFlash (block diffusion) speculative decoding.
1654
- Ensures output correctness on GSM8k, with cudagraphs and batching on.
1655
- """
1656
- monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0")
1657
-
1658
- spec_llm = LLM(**dflash_config)
1659
-
1660
- # Evaluate GSM8k accuracy (Qwen3-8B ref: ~87-92% on GSM8k)
1661
- evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8)
1662
-
1663
- current_metrics = spec_llm.get_metrics()
1664
- acceptance_len = compute_acceptance_len(current_metrics)
1665
-
1666
- # AR is thoroughly validated in test_dflash_acceptance_rates, in a manner consistent
1667
- # with the DFlash paper. However, that test measures AL per-request and thus runs
1668
- # with a batch size of 1. To ensure that AL does not collapse with large batch sizes
1669
- # we enforce a baseline on the AL over the full lm-eval-style GSM8k test.
1670
- expected_len = 3.5 # Measured is 3.9 to 4.0
1671
- print(f"DFlash GSM8k correctness test got AL {acceptance_len}")
1672
- assert acceptance_len >= expected_len, (
1673
- "DFlash correctness check failed with"
1674
- f" {acceptance_len=}, expected at least {expected_len}"
1675
- )
1676
-
1677
- del spec_llm
1678
- torch.accelerator.empty_cache()
1679
- cleanup_dist_env_and_memory()
tests/v1/e2e/spec_decode/utils.py
new
+158
@@ -0,0 +1,158 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import random
5
+from collections.abc import Iterable
6
+from typing import Any
7
+
8
+import pytest
9
+import torch
10
+
11
+from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_offline
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.v1.metrics.reader import Metric
16
+
17
+
18
+def _skip_if_insufficient_gpus_for_tp(tp_size: int):
19
+ """Skip test if available GPUs < tp_size on ROCm."""
20
+ available_gpus = torch.accelerator.device_count()
21
+ if available_gpus < tp_size:
22
+ pytest.skip(
23
+ f"Test requires {tp_size} GPUs, but only {available_gpus} available"
24
+ )
25
+
26
+
27
+Messages = list[dict[str, Any]]
28
+
29
+
30
+def get_test_prompts(mm_enabled: bool, num_prompts: int = 100) -> list[Messages]:
31
+ prompt_types = ["repeat", "gsm8k"]
32
+ if mm_enabled:
33
+ prompt_types.append("mm")
34
+ prompts: list[Messages] = []
35
+
36
+ num_repeat_prompts = num_prompts // len(prompt_types)
37
+ if mm_enabled:
38
+ num_gsm8k_prompts = num_prompts // len(prompt_types)
39
+ num_mm_prompts = num_prompts - num_repeat_prompts - num_gsm8k_prompts
40
+ else:
41
+ num_mm_prompts = 0
42
+ num_gsm8k_prompts = num_prompts - num_repeat_prompts
43
+
44
+ # Generate a mixed batch of prompts, some of which can be easily
45
+ # predicted by n-gram matching and some which likely cannot.
46
+ random.seed(0)
47
+ for _ in range(num_repeat_prompts):
48
+ word_choices = ["test", "temp", "hello", "where"]
49
+ word = random.choice(word_choices)
50
+ prompts.append(
51
+ [
52
+ {
53
+ "role": "user",
54
+ "content": f"""
55
+ please repeat the word '{word}' 10 times.
56
+ give no other output than the word at least ten times in a row,
57
+ in lowercase with spaces between each word and without quotes.
58
+ """,
59
+ }
60
+ ]
61
+ )
62
+ prompts.extend(
63
+ [{"role": "user", "content": prompt}]
64
+ for prompt in _build_gsm8k_prompts(
65
+ num_questions=num_gsm8k_prompts, num_shots=5
66
+ )[0]
67
+ )
68
+ for _ in range(num_mm_prompts):
69
+ placeholders = [
70
+ {
71
+ "type": "image_url",
72
+ "image_url": {
73
+ "url": f"{VLLM_S3_BUCKET_URL}/{VLM_IMAGES_DIR}/stop_sign.jpg"
74
+ },
75
+ }
76
+ ]
77
+ prompt = [
78
+ *placeholders,
79
+ {"type": "text", "text": "The meaning of the image is"},
80
+ ]
81
+ prompts.append([{"role": "user", "content": prompt}])
82
+
83
+ return prompts
84
+
85
+
86
+def get_instruct_coder_messages(n: int) -> list[Messages]:
87
+ from vllm.benchmarks.datasets import InstructCoderDataset
88
+
89
+ dataset = InstructCoderDataset(
90
+ dataset_path="likaixin/InstructCoder", dataset_split="train"
91
+ )
92
+ prompts: Iterable[str] = dataset.sample_prompts(n=n)
93
+ return [[{"role": "user", "content": prompt}] for prompt in prompts]
94
+
95
+
96
+def greedy_sampling() -> SamplingParams:
97
+ return SamplingParams(temperature=0, max_tokens=10, ignore_eos=False)
98
+
99
+
100
+def stochastic_sampling() -> SamplingParams:
101
+ return SamplingParams(temperature=1.0, max_tokens=10, ignore_eos=False)
102
+
103
+
104
+def evaluate_llm_for_gsm8k(llm: LLM, expected_accuracy_threshold: float = 0.70) -> None:
105
+ """Evaluate the LLM on GSM8K and check that accuracy is above a sanity threshold.
106
+
107
+ The default threshold assumes the LLM uses the same target model as the "model_name"
108
+ fixture, with max model len == 4096. Precomputed reference value is 75% to 80%
109
+ on GSM8K with greedy decoding, so we check that it's above a sanity threshold of 70%
110
+ to verify that the model is correct.
111
+ """
112
+ if expected_accuracy_threshold <= 0.0:
113
+ print("Skipping GSM8K evaluation")
114
+ return
115
+ results = evaluate_gsm8k_offline(llm)
116
+ accuracy = results["accuracy"]
117
+ print(f"GSM8K accuracy: {accuracy:.3f}")
118
+ assert accuracy >= expected_accuracy_threshold, (
119
+ f"Expected GSM8K accuracy >= {expected_accuracy_threshold}, got {accuracy:.3f}"
120
+ )
121
+
122
+
123
+def compute_acceptance_rate(
124
+ metrics: list[Metric], prev_metrics: list[Metric] | None = None
125
+) -> float:
126
+ name2metric = {metric.name: metric for metric in metrics}
127
+ n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value
128
+ if n_draft_toks == 0:
129
+ return float("nan")
130
+ n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value
131
+ 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
137
+ if n_draft_toks <= 0:
138
+ return float("nan")
139
+ return n_accepted_toks / n_draft_toks
140
+
141
+
142
+def compute_acceptance_len(
143
+ metrics: list[Metric], prev_metrics: list[Metric] | None = None
144
+) -> 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
148
+ if n_drafts == 0:
149
+ return 1
150
+ 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
156
+ if n_drafts <= 0:
157
+ return 1
158
+ return 1 + (n_accepted_toks / n_drafts)
vllm/model_executor/models/gemma4_dspark.py
+4
-2
@@ -22,7 +22,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
22
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
23
24
from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer
25
-from .qwen3_dflash import DFlashQwen3Model
25
+from .qwen3_dflash import DFlashQwen3Model, _dflash_layer_causal
26
from .qwen3_dspark import DSparkMarkovHead, Qwen3DSparkForCausalLM
27
from .utils import extract_layer_index, maybe_prefix
28
@@ -37,7 +37,8 @@ class Gemma4DSparkAttention(Gemma4MTPAttention):
37
quant_config: QuantizationConfig | None,
38
prefix: str,
39
) -> None:
40
- is_full = config.layer_types[extract_layer_index(prefix)] == "full_attention"
40
+ layer_idx = extract_layer_index(prefix)
41
+ is_full = config.layer_types[layer_idx] == "full_attention"
42
head_dim = (
43
getattr(config, "global_head_dim", config.head_dim)
44
if is_full
@@ -62,6 +63,7 @@ class Gemma4DSparkAttention(Gemma4MTPAttention):
63
prefix=prefix,
64
)
65
self.is_kv_shared_layer = False
66
+ self.causal = _dflash_layer_causal(config, layer_idx)
67
self.use_k_eq_v = use_k_eq_v
68
self.kv_size = self.num_kv_heads * self.head_dim
69
attn_bias = getattr(config, "attention_bias", False)
vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+4
@@ -19,6 +19,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
19
20
from vllm.compilation.backends import set_model_tag
21
from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal
22
+ from vllm.model_executor.models.utils import get_draft_quant_config
23
24
draft_vllm_config = replace(
25
vllm_config,
@@ -36,6 +37,9 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
37
else vllm_config.cache_config
38
),
39
)
40
+ # VllmConfig post-init restores the target's quant config because the target
41
+ # config is retained for DSpark's target-layer metadata, so we must override it.
42
+ draft_vllm_config.quant_config = get_draft_quant_config(vllm_config)
43
44
with set_model_tag("dspark_head"):
45
draft_model = get_model(