A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 18:30 UTC
[V1][Usage] Refactor speculative decoding configuration and tests (#14434)
Signed-off-by: Shangming Cai <caishangming@linux.alibaba.com>
shangmingc committed
Mar 23, 2025 at 13:28 UTC
50c9636d8745987f5691ed96f7dc21733740d53d
20 files changed
+1061
-808
docs/source/features/spec_decode.md
+25
-14
@@ -30,8 +30,10 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
30
llm = LLM(
31
model="facebook/opt-6.7b",
32
tensor_parallel_size=1,
33
- speculative_model="facebook/opt-125m",
34
- num_speculative_tokens=5,
33
+ speculative_config={
34
+ "model": "facebook/opt-125m",
35
+ "num_speculative_tokens": 5,
36
+ },
37
)
38
outputs = llm.generate(prompts, sampling_params)
39
@@ -45,10 +47,14 @@ To perform the same with an online mode launch the server:
47
48
```bash
49
python -m vllm.entrypoints.openai.api_server --host 0.0.0.0 --port 8000 --model facebook/opt-6.7b \
48
- --seed 42 -tp 1 --speculative_model facebook/opt-125m \
49
- --num_speculative_tokens 5 --gpu_memory_utilization 0.8
50
+ --seed 42 -tp 1 --gpu_memory_utilization 0.8 \
51
+ --speculative_config '{"model": "facebook/opt-125m", "num_speculative_tokens": 5}'
52
```
53
54
+:::{warning}
55
+Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately will be deprecated in the next release.
56
+:::
57
+
58
Then use a client:
59
60
```python
@@ -101,9 +107,11 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
107
llm = LLM(
108
model="facebook/opt-6.7b",
109
tensor_parallel_size=1,
104
- speculative_model="[ngram]",
105
- num_speculative_tokens=5,
106
- ngram_prompt_lookup_max=4,
110
+ speculative_config={
111
+ "method": "ngram",
112
+ "num_speculative_tokens": 5,
113
+ "prompt_lookup_max": 4,
114
+ },
115
)
116
outputs = llm.generate(prompts, sampling_params)
117
@@ -131,8 +139,10 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
139
llm = LLM(
140
model="meta-llama/Meta-Llama-3.1-70B-Instruct",
141
tensor_parallel_size=4,
134
- speculative_model="ibm-ai-platform/llama3-70b-accelerator",
135
- speculative_draft_tensor_parallel_size=1,
142
+ speculative_config={
143
+ "model": "ibm-ai-platform/llama3-70b-accelerator",
144
+ "draft_tensor_parallel_size": 1,
145
+ },
146
)
147
outputs = llm.generate(prompts, sampling_params)
148
@@ -175,8 +185,10 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
185
llm = LLM(
186
model="meta-llama/Meta-Llama-3-8B-Instruct",
187
tensor_parallel_size=4,
178
- speculative_model="yuhuili/EAGLE-LLaMA3-Instruct-8B",
179
- speculative_draft_tensor_parallel_size=1,
188
+ speculative_config={
189
+ "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
190
+ "draft_tensor_parallel_size": 1,
191
+ },
192
)
193
194
outputs = llm.generate(prompts, sampling_params)
@@ -194,11 +206,10 @@ A few important things to consider when using the EAGLE based draft models:
206
be able to be loaded and used directly by vLLM after [PR 12304](https://github.com/vllm-project/vllm/pull/12304).
207
If you are using vllm version before [PR 12304](https://github.com/vllm-project/vllm/pull/12304), please use the
208
[script](https://gist.github.com/abhigoyal1997/1e7a4109ccb7704fbc67f625e86b2d6d) to convert the speculative model,
197
- and specify `speculative_model="path/to/modified/eagle/model"`. If weight-loading problems still occur when using
198
- the latest version of vLLM, please leave a comment or raise an issue.
209
+ and specify `"model": "path/to/modified/eagle/model"` in `speculative_config`. If weight-loading problems still occur when using the latest version of vLLM, please leave a comment or raise an issue.
210
211
2. The EAGLE based draft models need to be run without tensor parallelism
201
- (i.e. speculative_draft_tensor_parallel_size is set to 1), although
212
+ (i.e. draft_tensor_parallel_size is set to 1 in `speculative_config`), although
213
it is possible to run the main model using tensor parallelism (see example above).
214
215
3. When using EAGLE-based speculators with vLLM, the observed speedup is lower than what is
examples/offline_inference/mlpspeculator.py
+3
-1
@@ -50,7 +50,9 @@ if __name__ == "__main__":
50
# Create an LLM with spec decoding
51
llm = LLM(
52
model="meta-llama/Llama-2-13b-chat-hf",
53
- speculative_model="ibm-ai-platform/llama-13b-accelerator",
53
+ speculative_config={
54
+ "model": "ibm-ai-platform/llama-13b-accelerator",
55
+ },
56
)
57
58
print("With speculation")
tests/spec_decode/e2e/conftest.py
+1
-1
@@ -56,7 +56,7 @@ def test_llm_generator(common_llm_kwargs, per_test_common_llm_kwargs,
56
def maybe_assert_ngram_worker(llm):
57
# Verify the proposer worker is ngram if ngram is specified.
58
if (llm.llm_engine.speculative_config is not None
59
- and llm.llm_engine.speculative_config.ngram_prompt_lookup_max > 0):
59
+ and llm.llm_engine.speculative_config.method == "ngram"):
60
from vllm.spec_decode.ngram_worker import NGramWorker
61
assert isinstance(
62
llm.llm_engine.model_executor.driver_worker.proposer_worker,
tests/spec_decode/e2e/test_compatibility.py
+20
-9
@@ -7,28 +7,39 @@ from vllm import SamplingParams
7
from .conftest import get_output_from_llm_generator
8
9
10
-@pytest.mark.parametrize("common_llm_kwargs", [{
11
- "model": "meta-llama/Llama-3.2-1B-Instruct",
12
- "speculative_model": "JackFram/llama-68m",
13
- "num_speculative_tokens": 5,
14
-}])
10
+@pytest.mark.parametrize("common_llm_kwargs",
11
+ [{
12
+ "model": "meta-llama/Llama-3.2-1B-Instruct",
13
+ }])
14
@pytest.mark.parametrize(
15
"per_test_common_llm_kwargs",
16
[
17
{
18
# Speculative max model len > overridden max model len should raise.
19
+ "speculative_config": {
20
+ "model": "JackFram/llama-68m",
21
+ "num_speculative_tokens": 5,
22
+ "max_model_len": 129,
23
+ },
24
"max_model_len": 128,
21
- "speculative_max_model_len": 129,
25
},
26
{
27
# Speculative max model len > draft max model len should raise.
28
# https://huggingface.co/JackFram/llama-68m/blob/3b606af5198a0b26762d589a3ee3d26ee6fa6c85/config.json#L12
26
- "speculative_max_model_len": 2048 + 1,
29
+ "speculative_config": {
30
+ "model": "JackFram/llama-68m",
31
+ "num_speculative_tokens": 5,
32
+ "max_model_len": 2048 + 1,
33
+ },
34
},
35
{
36
# Speculative max model len > target max model len should raise.
30
- # https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/blob/9213176726f574b556790deb65791e0c5aa438b6/config.json#L18
31
- "speculative_max_model_len": 131072 + 1,
37
+ # https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct/blob/9213176726f574b556790deb65791e0c5aa438b6/config.json#L18
38
+ "speculative_config": {
39
+ "model": "JackFram/llama-68m",
40
+ "num_speculative_tokens": 5,
41
+ "max_model_len": 131072 + 1,
42
+ },
43
},
44
])
45
@pytest.mark.parametrize("test_llm_kwargs", [{}])
tests/spec_decode/e2e/test_eagle_correctness.py
+57
-40
@@ -57,8 +57,10 @@ PRECISION = "float32"
57
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
58
@pytest.mark.parametrize("test_llm_kwargs", [
59
{
60
- "speculative_model": SPEC_MODEL,
61
- "num_speculative_tokens": MAX_SPEC_TOKENS,
60
+ "speculative_config": {
61
+ "model": SPEC_MODEL,
62
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
63
+ },
64
},
65
])
66
@pytest.mark.parametrize("output_len", [
@@ -95,18 +97,19 @@ def test_eagle_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
97
}])
98
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
99
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
98
-@pytest.mark.parametrize("test_llm_kwargs", [
99
- {
100
- "speculative_model": SPEC_MODEL,
100
+@pytest.mark.parametrize("test_llm_kwargs", [{
101
+ "speculative_config": {
102
+ "model": SPEC_MODEL,
103
"num_speculative_tokens": MAX_SPEC_TOKENS,
102
- "disable_logprobs_during_spec_decoding": False,
104
+ "disable_logprobs": False,
105
},
104
- {
105
- "speculative_model": SPEC_MODEL,
106
+}, {
107
+ "speculative_config": {
108
+ "model": SPEC_MODEL,
109
"num_speculative_tokens": MAX_SPEC_TOKENS,
107
- "disable_logprobs_during_spec_decoding": True,
110
+ "disable_logprobs": True,
111
},
109
-])
112
+}])
113
@pytest.mark.parametrize("output_len", [
114
128,
115
])
@@ -119,18 +122,19 @@ def test_eagle_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
122
batch_size: int, output_len: int, seed: int,
123
logprobs: int):
124
122
- run_equality_correctness_test(vllm_runner,
123
- common_llm_kwargs,
124
- per_test_common_llm_kwargs,
125
- baseline_llm_kwargs,
126
- test_llm_kwargs,
127
- batch_size,
128
- output_len,
129
- seed,
130
- logprobs=logprobs,
131
- prompt_logprobs=logprobs,
132
- disable_logprobs=test_llm_kwargs[
133
- 'disable_logprobs_during_spec_decoding'])
125
+ run_equality_correctness_test(
126
+ vllm_runner,
127
+ common_llm_kwargs,
128
+ per_test_common_llm_kwargs,
129
+ baseline_llm_kwargs,
130
+ test_llm_kwargs,
131
+ batch_size,
132
+ output_len,
133
+ seed,
134
+ logprobs=logprobs,
135
+ prompt_logprobs=logprobs,
136
+ disable_logprobs=test_llm_kwargs["speculative_config"]
137
+ ["disable_logprobs"])
138
139
140
@pytest.mark.parametrize(
@@ -151,8 +155,10 @@ def test_eagle_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
155
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
156
@pytest.mark.parametrize("test_llm_kwargs", [
157
{
154
- "speculative_model": SPEC_MODEL,
155
- "num_speculative_tokens": MAX_SPEC_TOKENS,
158
+ "speculative_config": {
159
+ "model": SPEC_MODEL,
160
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
161
+ },
162
},
163
])
164
@pytest.mark.parametrize("output_len", [
@@ -193,8 +199,10 @@ def test_eagle_e2e_greedy_correctness_cuda_graph(
199
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
200
@pytest.mark.parametrize("test_llm_kwargs", [
201
{
196
- "speculative_model": SPEC_MODEL,
197
- "num_speculative_tokens": MAX_SPEC_TOKENS,
202
+ "speculative_config": {
203
+ "model": SPEC_MODEL,
204
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
205
+ },
206
},
207
])
208
@pytest.mark.parametrize(
@@ -236,8 +244,10 @@ def test_eagle_e2e_greedy_correctness_with_preemption(
244
"test_llm_kwargs",
245
[
246
{
239
- "speculative_model": SPEC_MODEL,
240
- "num_speculative_tokens": k,
247
+ "speculative_config": {
248
+ "model": SPEC_MODEL,
249
+ "num_speculative_tokens": k,
250
+ },
251
}
252
# Try a range of num. speculative tokens
253
for k in range(1, 1 + MAX_SPEC_TOKENS)
@@ -277,12 +287,13 @@ def test_eagle_different_k(vllm_runner, common_llm_kwargs,
287
}])
288
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
289
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
280
-@pytest.mark.parametrize("test_llm_kwargs",
281
- [{
282
- "speculative_model": SPEC_MODEL,
283
- "num_speculative_tokens": MAX_SPEC_TOKENS,
284
- "speculative_disable_by_batch_size": 4
285
- }])
290
+@pytest.mark.parametrize("test_llm_kwargs", [{
291
+ "speculative_config": {
292
+ "model": SPEC_MODEL,
293
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
294
+ "disable_by_batch_size": 4,
295
+ },
296
+}])
297
@pytest.mark.parametrize("batch_size", [1, 5])
298
@pytest.mark.parametrize(
299
"output_len",
@@ -324,8 +335,10 @@ def test_eagle_disable_queue(vllm_runner, common_llm_kwargs,
335
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
336
@pytest.mark.parametrize("test_llm_kwargs", [
337
{
327
- "speculative_model": "yuhuili/EAGLE-llama2-chat-7B",
328
- "num_speculative_tokens": MAX_SPEC_TOKENS,
338
+ "speculative_config": {
339
+ "model": "yuhuili/EAGLE-llama2-chat-7B",
340
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
341
+ },
342
},
343
])
344
@pytest.mark.parametrize(
@@ -372,8 +385,10 @@ def test_llama2_eagle_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
385
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
386
@pytest.mark.parametrize("test_llm_kwargs", [
387
{
375
- "speculative_model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
376
- "num_speculative_tokens": MAX_SPEC_TOKENS,
388
+ "speculative_config": {
389
+ "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
390
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
391
+ },
392
},
393
])
394
@pytest.mark.parametrize(
@@ -420,8 +435,10 @@ def test_llama3_eagle_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
435
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
436
@pytest.mark.parametrize("test_llm_kwargs", [
437
{
423
- "speculative_model": "yuhuili/EAGLE-Qwen2-7B-Instruct",
424
- "num_speculative_tokens": MAX_SPEC_TOKENS,
438
+ "speculative_config": {
439
+ "model": "yuhuili/EAGLE-Qwen2-7B-Instruct",
440
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
441
+ },
442
},
443
])
444
@pytest.mark.parametrize(
tests/spec_decode/e2e/test_integration.py
+28
-18
@@ -23,8 +23,10 @@ MAIN_MODEL = "JackFram/llama-68m"
23
[
24
{
25
# Identical models.
26
- "speculative_model": "JackFram/llama-68m",
27
- "num_speculative_tokens": 5,
26
+ "speculative_config": {
27
+ "model": "JackFram/llama-68m",
28
+ "num_speculative_tokens": 5,
29
+ },
30
},
31
])
32
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@@ -57,26 +59,33 @@ def test_spec_decode_cuda_graph(vllm_runner, common_llm_kwargs,
59
# Skip cuda graph recording for fast test.
60
"enforce_eager": True,
61
}])
60
-@pytest.mark.parametrize("per_test_common_llm_kwargs", [
61
- {
62
- "speculative_model": "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit",
63
- "num_speculative_tokens": 5,
64
- },
65
-])
62
+@pytest.mark.parametrize("per_test_common_llm_kwargs", [])
63
@pytest.mark.parametrize(
64
"test_llm_kwargs",
65
[
66
# Explicitly specify draft model quantization
67
{
71
- "speculative_model_quantization": "gptq",
68
+ "speculative_config": {
69
+ "model": "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit",
70
+ "num_speculative_tokens": 5,
71
+ "quantization": "gptq",
72
+ },
73
},
74
# Explicitly specify GPTQ-based draft model to use marlin quantization
75
{
75
- "speculative_model_quantization": "marlin",
76
+ "speculative_config": {
77
+ "model": "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit",
78
+ "num_speculative_tokens": 5,
79
+ "quantization": "marlin",
80
+ },
81
},
82
# Not explicitly specify draft model quantization
83
{
79
- "speculative_model_quantization": None,
84
+ "speculative_config": {
85
+ "model": "LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit",
86
+ "num_speculative_tokens": 5,
87
+ "quantization": None,
88
+ },
89
},
90
])
91
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
@@ -107,15 +116,16 @@ def test_speculative_model_quantization_config(vllm_runner, common_llm_kwargs,
116
117
# Skip cuda graph recording for fast test.
118
"enforce_eager": True,
110
- "speculative_model": "JackFram/llama-68m",
111
- "num_speculative_tokens": 3,
119
}])
120
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
121
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
115
-@pytest.mark.parametrize("test_llm_kwargs",
116
- [{
117
- "speculative_disable_mqa_scorer": True,
118
- }])
122
+@pytest.mark.parametrize("test_llm_kwargs", [{
123
+ "speculative_config": {
124
+ "model": "JackFram/llama-68m",
125
+ "num_speculative_tokens": 3,
126
+ "disable_mqa_scorer": True,
127
+ },
128
+}])
129
@pytest.mark.parametrize("batch_size", [1, 5])
130
@pytest.mark.parametrize(
131
"output_len",
@@ -127,7 +137,7 @@ def test_speculative_model_quantization_config(vllm_runner, common_llm_kwargs,
137
def test_mqa_scorer(vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
138
baseline_llm_kwargs, test_llm_kwargs, batch_size: int,
139
output_len: int, seed: int):
130
- """Verify that ngram speculative decoding generates the same output
140
+ """Verify that speculative decoding generates the same output
141
with batch expansion scorer and mqa scorer.
142
"""
143
run_equality_correctness_test(vllm_runner,
tests/spec_decode/e2e/test_integration_dist_tp2.py
+40
-37
@@ -27,18 +27,19 @@ from .conftest import run_equality_correctness_test_tp
27
@pytest.mark.parametrize("baseline_llm_kwargs", [[]])
28
@pytest.mark.parametrize("test_llm_kwargs", [
29
[
30
- "--speculative-model",
31
- "JackFram/llama-68m",
32
- "--num-speculative-tokens",
33
- "3",
30
+ "--speculative_config",
31
+ str({
32
+ "model": "JackFram/llama-68m",
33
+ "num_speculative_tokens": 3,
34
+ }),
35
],
36
[
36
- "--speculative-model",
37
- "[ngram]",
38
- "--num-speculative-tokens",
39
- "5",
40
- "--ngram-prompt-lookup-max",
41
- "3",
37
+ "--speculative_config",
38
+ str({
39
+ "model": "ngram",
40
+ "num_speculative_tokens": 5,
41
+ "prompt_lookup_max": 3,
42
+ }),
43
],
44
])
45
@pytest.mark.parametrize("batch_size", [2])
@@ -83,23 +84,24 @@ def test_target_model_tp_gt_1(common_llm_kwargs, per_test_common_llm_kwargs,
84
]])
85
@pytest.mark.parametrize("per_test_common_llm_kwargs", [[]])
86
@pytest.mark.parametrize("baseline_llm_kwargs", [[]])
86
-@pytest.mark.parametrize("model, test_llm_kwargs",
87
- [("JackFram/llama-68m", [
88
- "--speculative-model",
89
- "JackFram/llama-68m",
90
- "--num_speculative-tokens",
91
- "5",
92
- "--speculative-draft-tensor-parallel-size",
93
- "1",
94
- ]),
95
- ("ibm-granite/granite-3b-code-instruct", [
96
- "--speculative-model",
97
- "ibm-granite/granite-3b-code-instruct",
98
- "--num_speculative-tokens",
99
- "5",
100
- "--speculative-draft-tensor-parallel-size",
101
- "1",
102
- ])])
87
+@pytest.mark.parametrize(
88
+ "model, test_llm_kwargs",
89
+ [("JackFram/llama-68m", [
90
+ "--speculative_config",
91
+ str({
92
+ "model": "JackFram/llama-68m",
93
+ "num_speculative_tokens": 5,
94
+ "draft_tensor_parallel_size": 1,
95
+ }),
96
+ ]),
97
+ ("ibm-granite/granite-3b-code-instruct", [
98
+ "--speculative_config",
99
+ str({
100
+ "model": "ibm-granite/granite-3b-code-instruct",
101
+ "num_speculative_tokens": 5,
102
+ "draft_tensor_parallel_size": 1,
103
+ }),
104
+ ])])
105
@pytest.mark.parametrize("batch_size", [2])
106
@pytest.mark.parametrize("seed", [1])
107
def test_draft_model_tp_lt_target_model_tp2(model, common_llm_kwargs,
@@ -144,18 +146,19 @@ def test_draft_model_tp_lt_target_model_tp2(model, common_llm_kwargs,
146
@pytest.mark.parametrize("baseline_llm_kwargs", [[]])
147
@pytest.mark.parametrize("model, test_llm_kwargs",
148
[("JackFram/llama-68m", [
147
- "--speculative-model",
148
- "JackFram/llama-68m",
149
- "--num_speculative-tokens",
150
- "3",
149
+ "--speculative_config",
150
+ str({
151
+ "model": "JackFram/llama-68m",
152
+ "num_speculative_tokens": 3,
153
+ }),
154
]),
155
("JackFram/llama-68m", [
153
- "--speculative-model",
154
- "JackFram/llama-68m",
155
- "--num_speculative-tokens",
156
- "3",
157
- "--speculative-draft-tensor-parallel-size",
158
- "1",
156
+ "--speculative_config",
157
+ str({
158
+ "model": "JackFram/llama-68m",
159
+ "num_speculative_tokens": 3,
160
+ "draft_tensor_parallel_size": 1,
161
+ }),
162
])])
163
@pytest.mark.parametrize("logprobs", [None, 2])
164
@pytest.mark.parametrize("batch_size", [2])
tests/spec_decode/e2e/test_integration_dist_tp4.py
+13
-15
@@ -24,12 +24,7 @@ SPEC_MODEL = "JackFram/llama-68m"
24
"4",
25
]])
26
@pytest.mark.parametrize("per_test_common_llm_kwargs", [
27
- [
28
- "--speculative-model",
29
- f"{SPEC_MODEL}",
30
- "--num-speculative-tokens",
31
- "5",
32
- ],
27
+ [],
28
])
29
@pytest.mark.parametrize("baseline_llm_kwargs", [[]])
30
@pytest.mark.parametrize(
@@ -37,8 +32,12 @@ SPEC_MODEL = "JackFram/llama-68m"
32
[
33
#TODO(wooyeon): add spec_draft_dp=2 case
34
[
40
- "--speculative-draft-tensor-parallel-size",
41
- "1",
35
+ "--speculative_config",
36
+ str({
37
+ "model": f"{SPEC_MODEL}",
38
+ "num_speculative_tokens": 5,
39
+ "draft_tensor_parallel_size": 1,
40
+ }),
41
],
42
])
43
@pytest.mark.parametrize("batch_size", [2])
@@ -78,15 +77,14 @@ def test_draft_model_tp_lt_target_model_tp4(common_llm_kwargs,
77
"test_llm_kwargs",
78
[
79
[
81
- "--speculative-model",
82
- f"{SPEC_MODEL}",
83
- "--num-speculative-tokens",
84
- "5",
85
-
80
# Artificially limit the draft model max model len; this forces vLLM
81
# to skip speculation once the sequences grow beyond 32-k tokens.
88
- "--speculative-max-model-len",
89
- "32",
82
+ "--speculative_config",
83
+ str({
84
+ "model": f"{SPEC_MODEL}",
85
+ "num_speculative_tokens": 5,
86
+ "max_model_len": 32,
87
+ }),
88
],
89
])
90
@pytest.mark.parametrize("batch_size", [8])
tests/spec_decode/e2e/test_logprobs.py
+102
-88
@@ -20,16 +20,19 @@ from .conftest import run_equality_correctness_test
20
}])
21
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
22
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
23
-@pytest.mark.parametrize("test_llm_kwargs",
24
- [{
25
- "speculative_model": "JackFram/llama-68m",
26
- "num_speculative_tokens": 3,
27
- "disable_logprobs_during_spec_decoding": False,
28
- }, {
29
- "speculative_model": "JackFram/llama-68m",
30
- "num_speculative_tokens": 3,
31
- "disable_logprobs_during_spec_decoding": True,
32
- }])
23
+@pytest.mark.parametrize("test_llm_kwargs", [{
24
+ "speculative_config": {
25
+ "model": "JackFram/llama-68m",
26
+ "num_speculative_tokens": 3,
27
+ "disable_logprobs": False,
28
+ },
29
+}, {
30
+ "speculative_config": {
31
+ "model": "JackFram/llama-68m",
32
+ "num_speculative_tokens": 3,
33
+ "disable_logprobs": True,
34
+ },
35
+}])
36
@pytest.mark.parametrize("batch_size", [8])
37
@pytest.mark.parametrize(
38
"output_len",
@@ -48,19 +51,20 @@ def test_logprobs_equality(vllm_runner, common_llm_kwargs,
51
as well as with and without chunked prefill.
52
"""
53
maybe_enable_chunked_prefill(prefill_chunk_size, common_llm_kwargs)
51
- run_equality_correctness_test(vllm_runner,
52
- common_llm_kwargs,
53
- per_test_common_llm_kwargs,
54
- baseline_llm_kwargs,
55
- test_llm_kwargs,
56
- batch_size,
57
- output_len,
58
- seed,
59
- temperature=0.0,
60
- logprobs=logprobs,
61
- prompt_logprobs=logprobs,
62
- disable_logprobs=test_llm_kwargs[
63
- 'disable_logprobs_during_spec_decoding'])
54
+ run_equality_correctness_test(
55
+ vllm_runner,
56
+ common_llm_kwargs,
57
+ per_test_common_llm_kwargs,
58
+ baseline_llm_kwargs,
59
+ test_llm_kwargs,
60
+ batch_size,
61
+ output_len,
62
+ seed,
63
+ temperature=0.0,
64
+ logprobs=logprobs,
65
+ prompt_logprobs=logprobs,
66
+ disable_logprobs=test_llm_kwargs["speculative_config"]
67
+ ["disable_logprobs"])
68
69
70
@pytest.mark.parametrize(
@@ -73,16 +77,19 @@ def test_logprobs_equality(vllm_runner, common_llm_kwargs,
77
}])
78
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
79
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
76
-@pytest.mark.parametrize("test_llm_kwargs",
77
- [{
78
- "speculative_model": "JackFram/llama-160m",
79
- "num_speculative_tokens": 3,
80
- "disable_logprobs_during_spec_decoding": False,
81
- }, {
82
- "speculative_model": "JackFram/llama-160m",
83
- "num_speculative_tokens": 6,
84
- "disable_logprobs_during_spec_decoding": False,
85
- }])
80
+@pytest.mark.parametrize("test_llm_kwargs", [{
81
+ "speculative_config": {
82
+ "model": "JackFram/llama-160m",
83
+ "num_speculative_tokens": 3,
84
+ "disable_logprobs": False,
85
+ },
86
+}, {
87
+ "speculative_config": {
88
+ "model": "JackFram/llama-160m",
89
+ "num_speculative_tokens": 6,
90
+ "disable_logprobs": False,
91
+ },
92
+}])
93
@pytest.mark.parametrize("batch_size", [8])
94
@pytest.mark.parametrize(
95
"output_len",
@@ -98,18 +105,19 @@ def test_logprobs_different_k(vllm_runner, common_llm_kwargs,
105
output_len: int, seed: int, logprobs: int):
106
"""Veriy logprob greedy equality with different speculation lens.
107
"""
101
- run_equality_correctness_test(vllm_runner,
102
- common_llm_kwargs,
103
- per_test_common_llm_kwargs,
104
- baseline_llm_kwargs,
105
- test_llm_kwargs,
106
- batch_size,
107
- output_len,
108
- seed,
109
- temperature=0.0,
110
- logprobs=logprobs,
111
- disable_logprobs=test_llm_kwargs[
112
- 'disable_logprobs_during_spec_decoding'])
108
+ run_equality_correctness_test(
109
+ vllm_runner,
110
+ common_llm_kwargs,
111
+ per_test_common_llm_kwargs,
112
+ baseline_llm_kwargs,
113
+ test_llm_kwargs,
114
+ batch_size,
115
+ output_len,
116
+ seed,
117
+ temperature=0.0,
118
+ logprobs=logprobs,
119
+ disable_logprobs=test_llm_kwargs["speculative_config"]
120
+ ["disable_logprobs"])
121
122
123
@pytest.mark.parametrize(
@@ -125,13 +133,15 @@ def test_logprobs_different_k(vllm_runner, common_llm_kwargs,
133
@pytest.mark.parametrize(
134
"test_llm_kwargs",
135
[{
128
- "speculative_model": "JackFram/llama-160m",
129
- "num_speculative_tokens": 3,
130
- "disable_logprobs_during_spec_decoding": False,
131
-
132
- # Artificially limit the draft model max model len; this forces vLLM
133
- # to skip speculation once the sequences grow beyond 32-k tokens.
134
- "speculative_max_model_len": 32,
136
+ "speculative_config": {
137
+ "model": "JackFram/llama-160m",
138
+ "num_speculative_tokens": 3,
139
+ "disable_logprobs": False,
140
+ # Artificially limit the draft model max model len; this forces
141
+ # vLLM to skip speculation once the sequences grow beyond 32-k
142
+ # tokens.
143
+ "max_model_len": 32,
144
+ },
145
}])
146
@pytest.mark.parametrize("batch_size", [8])
147
@pytest.mark.parametrize(
@@ -149,18 +159,19 @@ def test_logprobs_when_skip_speculation(vllm_runner, common_llm_kwargs,
159
seed: int, logprobs: int):
160
"""Verify logprobs greedy equality when some sequences skip speculation.
161
"""
152
- run_equality_correctness_test(vllm_runner,
153
- common_llm_kwargs,
154
- per_test_common_llm_kwargs,
155
- baseline_llm_kwargs,
156
- test_llm_kwargs,
157
- batch_size,
158
- output_len,
159
- seed,
160
- temperature=0.0,
161
- logprobs=logprobs,
162
- disable_logprobs=test_llm_kwargs[
163
- 'disable_logprobs_during_spec_decoding'])
162
+ run_equality_correctness_test(
163
+ vllm_runner,
164
+ common_llm_kwargs,
165
+ per_test_common_llm_kwargs,
166
+ baseline_llm_kwargs,
167
+ test_llm_kwargs,
168
+ batch_size,
169
+ output_len,
170
+ seed,
171
+ temperature=0.0,
172
+ logprobs=logprobs,
173
+ disable_logprobs=test_llm_kwargs["speculative_config"]
174
+ ["disable_logprobs"])
175
176
177
@pytest.mark.parametrize(
@@ -173,12 +184,13 @@ def test_logprobs_when_skip_speculation(vllm_runner, common_llm_kwargs,
184
}])
185
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
186
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
176
-@pytest.mark.parametrize("test_llm_kwargs",
177
- [{
178
- "speculative_model": "JackFram/llama-160m",
179
- "num_speculative_tokens": 3,
180
- "disable_logprobs_during_spec_decoding": False,
181
- }])
187
+@pytest.mark.parametrize("test_llm_kwargs", [{
188
+ "speculative_config": {
189
+ "model": "JackFram/llama-160m",
190
+ "num_speculative_tokens": 3,
191
+ "disable_logprobs": False,
192
+ },
193
+}])
194
@pytest.mark.parametrize("batch_size", [1])
195
@pytest.mark.parametrize(
196
"output_len",
@@ -248,12 +260,13 @@ def test_logprobs_temp_1(vllm_runner, common_llm_kwargs,
260
}])
261
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
262
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
251
-@pytest.mark.parametrize("test_llm_kwargs",
252
- [{
253
- "speculative_model": "JackFram/llama-68m",
254
- "num_speculative_tokens": 3,
255
- "disable_logprobs_during_spec_decoding": True,
256
- }])
263
+@pytest.mark.parametrize("test_llm_kwargs", [{
264
+ "speculative_config": {
265
+ "model": "JackFram/llama-68m",
266
+ "num_speculative_tokens": 3,
267
+ "disable_logprobs": True,
268
+ },
269
+}])
270
@pytest.mark.parametrize("seed", [1])
271
@pytest.mark.parametrize("batch_size", [4])
272
@pytest.mark.parametrize(
@@ -270,15 +283,16 @@ def test_logprobs_disabled(vllm_runner, common_llm_kwargs,
283
"""Check the behavior when logprobs are disabled.
284
Token choices should match with the base model.
285
"""
273
- run_equality_correctness_test(vllm_runner,
274
- common_llm_kwargs,
275
- per_test_common_llm_kwargs,
276
- baseline_llm_kwargs,
277
- test_llm_kwargs,
278
- batch_size,
279
- output_len,
280
- seed,
281
- temperature=0.0,
282
- logprobs=logprobs,
283
- disable_logprobs=test_llm_kwargs[
284
- 'disable_logprobs_during_spec_decoding'])
286
+ run_equality_correctness_test(
287
+ vllm_runner,
288
+ common_llm_kwargs,
289
+ per_test_common_llm_kwargs,
290
+ baseline_llm_kwargs,
291
+ test_llm_kwargs,
292
+ batch_size,
293
+ output_len,
294
+ seed,
295
+ temperature=0.0,
296
+ logprobs=logprobs,
297
+ disable_logprobs=test_llm_kwargs["speculative_config"]
298
+ ["disable_logprobs"])
tests/spec_decode/e2e/test_medusa_correctness.py
+55
-40
@@ -60,8 +60,10 @@ PRECISION = "float32"
60
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
61
@pytest.mark.parametrize("test_llm_kwargs", [
62
{
63
- "speculative_model": SPEC_MODEL,
64
- "num_speculative_tokens": MAX_SPEC_TOKENS,
63
+ "speculative_config": {
64
+ "model": SPEC_MODEL,
65
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
66
+ },
67
},
68
])
69
@pytest.mark.parametrize("output_len", [
@@ -107,14 +109,18 @@ def test_medusa_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
109
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
110
@pytest.mark.parametrize("test_llm_kwargs", [
111
{
110
- "speculative_model": SPEC_MODEL,
111
- "num_speculative_tokens": MAX_SPEC_TOKENS,
112
- "disable_logprobs_during_spec_decoding": False,
112
+ "speculative_config": {
113
+ "model": SPEC_MODEL,
114
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
115
+ "disable_logprobs": False,
116
+ },
117
},
118
{
115
- "speculative_model": SPEC_MODEL,
116
- "num_speculative_tokens": MAX_SPEC_TOKENS,
117
- "disable_logprobs_during_spec_decoding": True,
119
+ "speculative_config": {
120
+ "model": SPEC_MODEL,
121
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
122
+ "disable_logprobs": True,
123
+ },
124
},
125
])
126
@pytest.mark.parametrize("output_len", [
@@ -132,19 +138,20 @@ def test_medusa_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
138
prefill_chunk_size: int):
139
"""Verify greedy equality with different batch size."""
140
maybe_enable_chunked_prefill(prefill_chunk_size, test_llm_kwargs)
135
- run_equality_correctness_test(vllm_runner,
136
- common_llm_kwargs,
137
- per_test_common_llm_kwargs,
138
- baseline_llm_kwargs,
139
- test_llm_kwargs,
140
- batch_size,
141
- max_output_len=output_len,
142
- seed=seed,
143
- temperature=0.0,
144
- logprobs=logprobs,
145
- prompt_logprobs=logprobs,
146
- disable_logprobs=test_llm_kwargs[
147
- 'disable_logprobs_during_spec_decoding'])
141
+ run_equality_correctness_test(
142
+ vllm_runner,
143
+ common_llm_kwargs,
144
+ per_test_common_llm_kwargs,
145
+ baseline_llm_kwargs,
146
+ test_llm_kwargs,
147
+ batch_size,
148
+ max_output_len=output_len,
149
+ seed=seed,
150
+ temperature=0.0,
151
+ logprobs=logprobs,
152
+ prompt_logprobs=logprobs,
153
+ disable_logprobs=test_llm_kwargs["speculative_config"]
154
+ ["disable_logprobs"])
155
156
157
@pytest.mark.parametrize(
@@ -165,8 +172,10 @@ def test_medusa_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
172
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
173
@pytest.mark.parametrize("test_llm_kwargs", [
174
{
168
- "speculative_model": SPEC_MODEL,
169
- "num_speculative_tokens": MAX_SPEC_TOKENS,
175
+ "speculative_config": {
176
+ "model": SPEC_MODEL,
177
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
178
+ },
179
},
180
])
181
@pytest.mark.parametrize("output_len", [
@@ -214,8 +223,10 @@ def test_medusa_e2e_greedy_correctness_cuda_graph(
223
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
224
@pytest.mark.parametrize("test_llm_kwargs", [
225
{
217
- "speculative_model": SPEC_MODEL,
218
- "num_speculative_tokens": MAX_SPEC_TOKENS,
226
+ "speculative_config": {
227
+ "model": SPEC_MODEL,
228
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
229
+ },
230
},
231
])
232
@pytest.mark.parametrize(
@@ -264,8 +275,10 @@ def test_medusa_e2e_greedy_correctness_with_preemption(
275
"test_llm_kwargs",
276
[
277
{
267
- "speculative_model": SPEC_MODEL,
268
- "num_speculative_tokens": k,
278
+ "speculative_config": {
279
+ "model": SPEC_MODEL,
280
+ "num_speculative_tokens": k,
281
+ },
282
}
283
# Try a range of num. speculative tokens
284
for k in range(1, 1 + MAX_SPEC_TOKENS)
@@ -312,12 +325,13 @@ def test_medusa_different_k(vllm_runner, common_llm_kwargs,
325
}])
326
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
327
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
315
-@pytest.mark.parametrize("test_llm_kwargs",
316
- [{
317
- "speculative_model": SPEC_MODEL,
318
- "num_speculative_tokens": MAX_SPEC_TOKENS,
319
- "speculative_disable_by_batch_size": 4
320
- }])
328
+@pytest.mark.parametrize("test_llm_kwargs", [{
329
+ "speculative_config": {
330
+ "model": SPEC_MODEL,
331
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
332
+ "disable_by_batch_size": 4,
333
+ },
334
+}])
335
@pytest.mark.parametrize("batch_size", [1, 5])
336
@pytest.mark.parametrize(
337
"output_len",
@@ -359,16 +373,17 @@ def test_medusa_disable_queue(vllm_runner, common_llm_kwargs,
373
374
# Main model
375
"model_name": MAIN_MODEL,
362
- "speculative_model": SPEC_MODEL,
363
- "num_speculative_tokens": MAX_SPEC_TOKENS,
364
- "speculative_disable_by_batch_size": 4
376
}])
377
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
378
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
368
-@pytest.mark.parametrize("test_llm_kwargs",
369
- [{
370
- "speculative_disable_mqa_scorer": True,
371
- }])
379
+@pytest.mark.parametrize("test_llm_kwargs", [{
380
+ "speculative_config": {
381
+ "model": SPEC_MODEL,
382
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
383
+ "disable_by_batch_size": 4,
384
+ "disable_mqa_scorer": True,
385
+ },
386
+}])
387
@pytest.mark.parametrize("batch_size", [1, 5])
388
@pytest.mark.parametrize(
389
"output_len",
tests/spec_decode/e2e/test_mlp_correctness.py
+54
-35
@@ -62,7 +62,9 @@ PRECISION = "float32"
62
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
63
@pytest.mark.parametrize("test_llm_kwargs", [
64
{
65
- "speculative_model": SPEC_MODEL,
65
+ "speculative_config": {
66
+ "model": SPEC_MODEL,
67
+ },
68
},
69
])
70
@pytest.mark.parametrize("output_len", [
@@ -108,12 +110,16 @@ def test_mlp_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
110
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
111
@pytest.mark.parametrize("test_llm_kwargs", [
112
{
111
- "speculative_model": SPEC_MODEL,
112
- "disable_logprobs_during_spec_decoding": False,
113
+ "speculative_config": {
114
+ "model": SPEC_MODEL,
115
+ "disable_logprobs": False,
116
+ },
117
},
118
{
115
- "speculative_model": SPEC_MODEL,
116
- "disable_logprobs_during_spec_decoding": True,
119
+ "speculative_config": {
120
+ "model": SPEC_MODEL,
121
+ "disable_logprobs": True,
122
+ },
123
},
124
])
125
@pytest.mark.parametrize("output_len", [8])
@@ -133,19 +139,20 @@ def test_mlp_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
139
# up sampling different tokens at the tail (ie top tokens don't change).
140
# TL;DR: sd+cp == org+cp but sd+cp != org..is this expected?
141
maybe_enable_chunked_prefill(prefill_chunk_size, baseline_llm_kwargs)
136
- run_equality_correctness_test(vllm_runner,
137
- common_llm_kwargs,
138
- per_test_common_llm_kwargs,
139
- baseline_llm_kwargs,
140
- test_llm_kwargs,
141
- batch_size,
142
- max_output_len=output_len,
143
- seed=seed,
144
- temperature=0.0,
145
- logprobs=logprobs,
146
- prompt_logprobs=logprobs,
147
- disable_logprobs=test_llm_kwargs[
148
- 'disable_logprobs_during_spec_decoding'])
142
+ run_equality_correctness_test(
143
+ vllm_runner,
144
+ common_llm_kwargs,
145
+ per_test_common_llm_kwargs,
146
+ baseline_llm_kwargs,
147
+ test_llm_kwargs,
148
+ batch_size,
149
+ max_output_len=output_len,
150
+ seed=seed,
151
+ temperature=0.0,
152
+ logprobs=logprobs,
153
+ prompt_logprobs=logprobs,
154
+ disable_logprobs=test_llm_kwargs["speculative_config"]
155
+ ["disable_logprobs"])
156
157
158
@pytest.mark.parametrize(
@@ -167,7 +174,9 @@ def test_mlp_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
174
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
175
@pytest.mark.parametrize("test_llm_kwargs", [
176
{
170
- "speculative_model": SPEC_MODEL,
177
+ "speculative_config": {
178
+ "model": SPEC_MODEL,
179
+ },
180
},
181
])
182
@pytest.mark.parametrize("output_len", [2048])
@@ -209,8 +218,10 @@ def test_mlp_e2e_acceptance_rate(vllm_runner, common_llm_kwargs,
218
# Main model
219
"model_name": MAIN_MODEL,
220
212
- # Speculative model
213
- "speculative_model": SPEC_MODEL,
221
+ # Speculative config
222
+ "speculative_config": {
223
+ "model": SPEC_MODEL,
224
+ },
225
}])
226
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
227
@pytest.mark.parametrize("baseline_llm_kwargs", [{"seed": 1}])
@@ -274,7 +285,9 @@ def test_mlp_e2e_seeded_correctness(vllm_runner, common_llm_kwargs,
285
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
286
@pytest.mark.parametrize("test_llm_kwargs", [
287
{
277
- "speculative_model": SPEC_MODEL,
288
+ "speculative_config": {
289
+ "model": SPEC_MODEL,
290
+ },
291
},
292
])
293
@pytest.mark.parametrize(
@@ -326,7 +339,9 @@ def test_mlp_e2e_greedy_correctness_with_preemption(
339
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
340
@pytest.mark.parametrize("test_llm_kwargs", [
341
{
329
- "speculative_model": SPEC_MODEL,
342
+ "speculative_config": {
343
+ "model": SPEC_MODEL,
344
+ },
345
},
346
])
347
@pytest.mark.parametrize(
@@ -382,8 +397,10 @@ def test_mlp_e2e_greedy_correctness_with_padding(
397
"test_llm_kwargs",
398
[
399
{
385
- "speculative_model": SPEC_MODEL,
386
- "num_speculative_tokens": k,
400
+ "speculative_config": {
401
+ "model": SPEC_MODEL,
402
+ "num_speculative_tokens": k,
403
+ },
404
}
405
# Try a range of num. speculative tokens
406
for k in range(1, 1 + MAX_SPEC_TOKENS)
@@ -430,11 +447,12 @@ def test_mlp_different_k(vllm_runner, common_llm_kwargs,
447
}])
448
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
449
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
433
-@pytest.mark.parametrize("test_llm_kwargs",
434
- [{
435
- "speculative_model": SPEC_MODEL,
436
- "speculative_disable_by_batch_size": 4
437
- }])
450
+@pytest.mark.parametrize("test_llm_kwargs", [{
451
+ "speculative_config": {
452
+ "model": SPEC_MODEL,
453
+ "disable_by_batch_size": 4,
454
+ },
455
+}])
456
@pytest.mark.parametrize("batch_size", [1, 5])
457
@pytest.mark.parametrize(
458
"output_len",
@@ -475,14 +493,15 @@ def test_mlp_disable_queue(vllm_runner, common_llm_kwargs,
493
494
# Skip cuda graph recording for fast test.
495
"enforce_eager": True,
478
- "speculative_model": SPEC_MODEL,
496
}])
497
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
498
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
482
-@pytest.mark.parametrize("test_llm_kwargs",
483
- [{
484
- "speculative_disable_mqa_scorer": True,
485
- }])
499
+@pytest.mark.parametrize("test_llm_kwargs", [{
500
+ "speculative_config": {
501
+ "model": SPEC_MODEL,
502
+ "disable_mqa_scorer": True,
503
+ },
504
+}])
505
@pytest.mark.parametrize("batch_size", [1, 5])
506
@pytest.mark.parametrize(
507
"output_len",
tests/spec_decode/e2e/test_mtp_correctness.py
+39
-25
@@ -57,7 +57,9 @@ PRECISION = "bfloat16"
57
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
58
@pytest.mark.parametrize("test_llm_kwargs", [
59
{
60
- "num_speculative_tokens": MAX_SPEC_TOKENS,
60
+ "speculative_config": {
61
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
62
+ },
63
},
64
])
65
@pytest.mark.parametrize("output_len", [
@@ -99,12 +101,16 @@ def test_mtp_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
101
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
102
@pytest.mark.parametrize("test_llm_kwargs", [
103
{
102
- "num_speculative_tokens": MAX_SPEC_TOKENS,
103
- "disable_logprobs_during_spec_decoding": False,
104
+ "speculative_config": {
105
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
106
+ "disable_logprobs": False,
107
+ },
108
},
109
{
106
- "num_speculative_tokens": MAX_SPEC_TOKENS,
107
- "disable_logprobs_during_spec_decoding": True,
110
+ "speculative_config": {
111
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
112
+ "disable_logprobs": True,
113
+ },
114
},
115
])
116
@pytest.mark.parametrize("output_len", [
@@ -119,18 +125,19 @@ def test_mtp_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
125
batch_size: int, output_len: int, seed: int,
126
logprobs: int):
127
122
- run_equality_correctness_test(vllm_runner,
123
- common_llm_kwargs,
124
- per_test_common_llm_kwargs,
125
- baseline_llm_kwargs,
126
- test_llm_kwargs,
127
- batch_size,
128
- output_len,
129
- seed,
130
- logprobs=logprobs,
131
- prompt_logprobs=logprobs,
132
- disable_logprobs=test_llm_kwargs[
133
- 'disable_logprobs_during_spec_decoding'])
128
+ run_equality_correctness_test(
129
+ vllm_runner,
130
+ common_llm_kwargs,
131
+ per_test_common_llm_kwargs,
132
+ baseline_llm_kwargs,
133
+ test_llm_kwargs,
134
+ batch_size,
135
+ output_len,
136
+ seed,
137
+ logprobs=logprobs,
138
+ prompt_logprobs=logprobs,
139
+ disable_logprobs=test_llm_kwargs["speculative_config"]
140
+ ["disable_logprobs"])
141
142
143
@pytest.mark.parametrize(
@@ -152,7 +159,9 @@ def test_mtp_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
159
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
160
@pytest.mark.parametrize("test_llm_kwargs", [
161
{
155
- "num_speculative_tokens": MAX_SPEC_TOKENS,
162
+ "speculative_config": {
163
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
164
+ },
165
},
166
])
167
@pytest.mark.parametrize("output_len", [
@@ -198,7 +207,9 @@ def test_mtp_e2e_greedy_correctness_cuda_graph(vllm_runner, common_llm_kwargs,
207
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
208
@pytest.mark.parametrize("test_llm_kwargs", [
209
{
201
- "num_speculative_tokens": MAX_SPEC_TOKENS,
210
+ "speculative_config": {
211
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
212
+ },
213
},
214
])
215
@pytest.mark.parametrize(
@@ -243,7 +254,9 @@ def test_mtp_e2e_greedy_correctness_with_preemption(
254
"test_llm_kwargs",
255
[
256
{
246
- "num_speculative_tokens": k,
257
+ "speculative_config": {
258
+ "num_speculative_tokens": k,
259
+ },
260
}
261
# Try a range of num. speculative tokens
262
for k in range(1, 1 + MAX_SPEC_TOKENS)
@@ -286,11 +299,12 @@ def test_mtp_different_k(vllm_runner, common_llm_kwargs,
299
}])
300
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
301
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
289
-@pytest.mark.parametrize("test_llm_kwargs",
290
- [{
291
- "num_speculative_tokens": MAX_SPEC_TOKENS,
292
- "speculative_disable_by_batch_size": 4
293
- }])
302
+@pytest.mark.parametrize("test_llm_kwargs", [{
303
+ "speculative_config": {
304
+ "num_speculative_tokens": MAX_SPEC_TOKENS,
305
+ "disable_by_batch_size": 4
306
+ },
307
+}])
308
@pytest.mark.parametrize("batch_size", [1, 5])
309
@pytest.mark.parametrize(
310
"output_len",
tests/spec_decode/e2e/test_multistep_correctness.py
+112
-65
@@ -61,15 +61,19 @@ from .conftest import (get_output_from_llm_generator,
61
"per_test_common_llm_kwargs",
62
[
63
{
64
- "speculative_model": "JackFram/llama-68m",
65
- "num_speculative_tokens": 5,
64
+ "speculative_config": {
65
+ "model": "JackFram/llama-68m",
66
+ "num_speculative_tokens": 5,
67
+ },
68
"enable_chunked_prefill": False,
69
},
70
{
71
# Chunked prefill enabled with small value
72
# to make sure we get mixed batches.
71
- "speculative_model": "JackFram/llama-68m",
72
- "num_speculative_tokens": 5,
73
+ "speculative_config": {
74
+ "model": "JackFram/llama-68m",
75
+ "num_speculative_tokens": 5,
76
+ },
77
"enable_chunked_prefill": True,
78
"max_num_batched_tokens": 4,
79
"max_num_seqs": 4
@@ -148,20 +152,23 @@ def test_spec_decode_e2e_with_detokenization(test_llm_generator,
152
},
153
])
154
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
151
-@pytest.mark.parametrize("test_llm_kwargs",
152
- [{
153
- "speculative_model": "JackFram/llama-68m",
154
- "num_speculative_tokens": 5,
155
- "enable_chunked_prefill": False,
156
- "disable_logprobs_during_spec_decoding": False
157
- }, {
158
- "speculative_model": "JackFram/llama-68m",
159
- "num_speculative_tokens": 3,
160
- "enable_chunked_prefill": True,
161
- "max_num_batched_tokens": 4,
162
- "max_num_seqs": 4,
163
- "disable_logprobs_during_spec_decoding": False
164
- }])
155
+@pytest.mark.parametrize("test_llm_kwargs", [{
156
+ "speculative_config": {
157
+ "model": "JackFram/llama-68m",
158
+ "num_speculative_tokens": 5,
159
+ "disable_logprobs": False,
160
+ },
161
+ "enable_chunked_prefill": False,
162
+}, {
163
+ "speculative_config": {
164
+ "model": "JackFram/llama-68m",
165
+ "num_speculative_tokens": 3,
166
+ "disable_logprobs": False,
167
+ },
168
+ "enable_chunked_prefill": True,
169
+ "max_num_batched_tokens": 4,
170
+ "max_num_seqs": 4,
171
+}])
172
@pytest.mark.parametrize(
173
"output_len",
174
[
@@ -184,7 +191,7 @@ def test_spec_decode_e2e_greedy_correctness_tiny_model_bs1(
191
whether all speculative tokens are accepted.
192
"""
193
ensure_all_accepted = per_test_common_llm_kwargs.get(
187
- "model_name") == test_llm_kwargs.get("speculative_model")
194
+ "model_name") == test_llm_kwargs.get("speculative_config")["model"]
195
run_equality_correctness_test(vllm_runner,
196
common_llm_kwargs,
197
per_test_common_llm_kwargs,
@@ -224,13 +231,17 @@ def test_spec_decode_e2e_greedy_correctness_tiny_model_bs1(
231
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
232
@pytest.mark.parametrize("test_llm_kwargs", [
233
{
227
- "speculative_model": "JackFram/llama-68m",
228
- "num_speculative_tokens": 5,
234
+ "speculative_config": {
235
+ "model": "JackFram/llama-68m",
236
+ "num_speculative_tokens": 5,
237
+ },
238
"enable_chunked_prefill": False,
239
},
240
{
232
- "speculative_model": "JackFram/llama-68m",
233
- "num_speculative_tokens": 5,
241
+ "speculative_config": {
242
+ "model": "JackFram/llama-68m",
243
+ "num_speculative_tokens": 5,
244
+ },
245
"enable_chunked_prefill": True,
246
"max_num_batched_tokens": 4,
247
"max_num_seqs": 4
@@ -283,13 +294,17 @@ def test_spec_decode_e2e_greedy_correctness_tiny_model_large_bs(
294
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
295
@pytest.mark.parametrize("test_llm_kwargs", [
296
{
286
- "speculative_model": "JackFram/llama-68m",
287
- "num_speculative_tokens": 5,
297
+ "speculative_config": {
298
+ "model": "JackFram/llama-68m",
299
+ "num_speculative_tokens": 5,
300
+ },
301
"enable_chunked_prefill": False,
302
},
303
{
291
- "speculative_model": "JackFram/llama-68m",
292
- "num_speculative_tokens": 5,
304
+ "speculative_config": {
305
+ "model": "JackFram/llama-68m",
306
+ "num_speculative_tokens": 5,
307
+ },
308
"enable_chunked_prefill": True,
309
"max_num_batched_tokens": 4,
310
"max_num_seqs": 4
@@ -336,13 +351,17 @@ def test_spec_decode_e2e_greedy_correctness_tiny_model_large_bs_diff_output_len(
351
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
352
@pytest.mark.parametrize("test_llm_kwargs", [
353
{
339
- "speculative_model": "JackFram/llama-68m",
340
- "num_speculative_tokens": 5,
354
+ "speculative_config": {
355
+ "model": "JackFram/llama-68m",
356
+ "num_speculative_tokens": 5,
357
+ },
358
"enable_chunked_prefill": False,
359
},
360
{
344
- "speculative_model": "JackFram/llama-68m",
345
- "num_speculative_tokens": 5,
361
+ "speculative_config": {
362
+ "model": "JackFram/llama-68m",
363
+ "num_speculative_tokens": 5,
364
+ },
365
"enable_chunked_prefill": True,
366
"max_num_batched_tokens": 4,
367
"max_num_seqs": 4
@@ -391,13 +410,17 @@ def test_spec_decode_e2e_greedy_correctness_real_model_bs1(
410
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
411
@pytest.mark.parametrize("test_llm_kwargs", [
412
{
394
- "speculative_model": "JackFram/llama-68m",
395
- "num_speculative_tokens": 5,
413
+ "speculative_config": {
414
+ "model": "JackFram/llama-68m",
415
+ "num_speculative_tokens": 5,
416
+ },
417
"enable_chunked_prefill": False,
418
},
419
{
399
- "speculative_model": "JackFram/llama-68m",
400
- "num_speculative_tokens": 5,
420
+ "speculative_config": {
421
+ "model": "JackFram/llama-68m",
422
+ "num_speculative_tokens": 5,
423
+ },
424
"enable_chunked_prefill": True,
425
"max_num_batched_tokens": 4,
426
"max_num_seqs": 4
@@ -449,13 +472,17 @@ def test_spec_decode_e2e_greedy_correctness_real_model_large_bs(
472
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
473
@pytest.mark.parametrize("test_llm_kwargs", [
474
{
452
- "speculative_model": "JackFram/llama-68m",
453
- "num_speculative_tokens": 5,
475
+ "speculative_config": {
476
+ "model": "JackFram/llama-68m",
477
+ "num_speculative_tokens": 5,
478
+ },
479
"enable_chunked_prefill": False,
480
},
481
{
457
- "speculative_model": "JackFram/llama-68m",
458
- "num_speculative_tokens": 5,
482
+ "speculative_config": {
483
+ "model": "JackFram/llama-68m",
484
+ "num_speculative_tokens": 5,
485
+ },
486
"enable_chunked_prefill": True,
487
"max_num_batched_tokens": 4,
488
"max_num_seqs": 4
@@ -514,13 +541,17 @@ def test_spec_decode_e2e_greedy_correctness_with_preemption(
541
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
542
@pytest.mark.parametrize("test_llm_kwargs", [
543
{
517
- "speculative_model": "JackFram/llama-68m",
518
- "num_speculative_tokens": 5,
544
+ "speculative_config": {
545
+ "model": "JackFram/llama-68m",
546
+ "num_speculative_tokens": 5,
547
+ },
548
"enable_chunked_prefill": False,
549
},
550
{
522
- "speculative_model": "JackFram/llama-68m",
523
- "num_speculative_tokens": 5,
551
+ "speculative_config": {
552
+ "model": "JackFram/llama-68m",
553
+ "num_speculative_tokens": 5,
554
+ },
555
"enable_chunked_prefill": True,
556
"max_num_batched_tokens": 4,
557
"max_num_seqs": 4
@@ -567,21 +598,25 @@ def test_spec_decode_different_block_size(vllm_runner, common_llm_kwargs,
598
"test_llm_kwargs",
599
[
600
{
570
- "speculative_model": "JackFram/llama-68m",
571
- "num_speculative_tokens": 5,
601
602
# Artificially limit the draft model max model len; this forces vLLM
603
# to skip speculation once the sequences grow beyond 32-k tokens.
575
- "speculative_max_model_len": 32,
604
+ "speculative_config": {
605
+ "model": "JackFram/llama-68m",
606
+ "num_speculative_tokens": 5,
607
+ "max_model_len": 32,
608
+ },
609
"enable_chunked_prefill": False,
610
},
611
{
579
- "speculative_model": "JackFram/llama-68m",
580
- "num_speculative_tokens": 5,
612
+ "speculative_config": {
613
+ "model": "JackFram/llama-68m",
614
+ "num_speculative_tokens": 5,
615
+ "max_model_len": 32,
616
+ },
617
"enable_chunked_prefill": True,
618
"max_num_batched_tokens": 4,
619
"max_num_seqs": 4,
584
- "speculative_max_model_len": 32,
620
},
621
])
622
@pytest.mark.parametrize("batch_size", [8])
@@ -627,15 +662,19 @@ def test_skip_speculation(vllm_runner, common_llm_kwargs,
662
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
663
@pytest.mark.parametrize("test_llm_kwargs", [
664
{
630
- "speculative_model": "JackFram/llama-68m",
631
- "num_speculative_tokens": 5,
632
- "speculative_disable_by_batch_size": 2,
665
+ "speculative_config": {
666
+ "model": "JackFram/llama-68m",
667
+ "num_speculative_tokens": 5,
668
+ "disable_by_batch_size": 2,
669
+ },
670
"enable_chunked_prefill": False,
671
},
672
{
636
- "speculative_model": "JackFram/llama-68m",
637
- "num_speculative_tokens": 5,
638
- "speculative_disable_by_batch_size": 2,
673
+ "speculative_config": {
674
+ "model": "JackFram/llama-68m",
675
+ "num_speculative_tokens": 5,
676
+ "disable_by_batch_size": 2,
677
+ },
678
"enable_chunked_prefill": True,
679
"max_num_batched_tokens": 4,
680
"max_num_seqs": 4,
@@ -676,15 +715,19 @@ def test_disable_speculation(vllm_runner, common_llm_kwargs,
715
"test_llm_kwargs",
716
[
717
{
679
- "speculative_model": "JackFram/llama-68m",
680
- "num_speculative_tokens": k,
718
+ "speculative_config": {
719
+ "model": "JackFram/llama-68m",
720
+ "num_speculative_tokens": k,
721
+ },
722
"enable_chunked_prefill": False,
723
}
724
# Try a range of common k, as well as large speculation.
725
for k in [1, 2, 3, 4, 5, 6, 7, 8, 9, 63]
726
] + [{
686
- "speculative_model": "JackFram/llama-68m",
687
- "num_speculative_tokens": k,
727
+ "speculative_config": {
728
+ "model": "JackFram/llama-68m",
729
+ "num_speculative_tokens": k,
730
+ },
731
"enable_chunked_prefill": True,
732
"max_num_batched_tokens": 4,
733
"max_num_seqs": 4,
@@ -729,17 +772,21 @@ def test_many_k(vllm_runner, common_llm_kwargs, per_test_common_llm_kwargs,
772
"test_llm_kwargs",
773
[
774
{
732
- "speculative_model": "JackFram/llama-68m",
733
- "num_speculative_tokens": k,
734
- "spec_decoding_acceptance_method": "typical_acceptance_sampler",
775
+ "speculative_config": {
776
+ "model": "JackFram/llama-68m",
777
+ "num_speculative_tokens": k,
778
+ "acceptance_method": "typical_acceptance_sampler",
779
+ },
780
"enable_chunked_prefill": False
781
}
782
# Try a range of common k.
783
for k in [1, 2, 3]
784
] + [{
740
- "speculative_model": "JackFram/llama-68m",
741
- "num_speculative_tokens": k,
742
- "spec_decoding_acceptance_method": "typical_acceptance_sampler",
785
+ "speculative_config": {
786
+ "model": "JackFram/llama-68m",
787
+ "num_speculative_tokens": k,
788
+ "acceptance_method": "typical_acceptance_sampler",
789
+ },
790
"enable_chunked_prefill": True,
791
"max_num_batched_tokens": 4,
792
"max_num_seqs": 4
tests/spec_decode/e2e/test_ngram_correctness.py
+88
-69
@@ -48,16 +48,20 @@ from .conftest import run_equality_correctness_test
48
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
49
@pytest.mark.parametrize("test_llm_kwargs", [
50
{
51
- "speculative_model": "[ngram]",
52
- "num_speculative_tokens": 5,
53
- "ngram_prompt_lookup_max": 3,
54
- "speculative_disable_mqa_scorer": False,
51
+ "speculative_config": {
52
+ "method": "ngram",
53
+ "num_speculative_tokens": 5,
54
+ "prompt_lookup_max": 3,
55
+ "disable_mqa_scorer": False,
56
+ },
57
},
58
{
57
- "speculative_model": "[ngram]",
58
- "num_speculative_tokens": 5,
59
- "ngram_prompt_lookup_max": 3,
60
- "speculative_disable_mqa_scorer": True,
59
+ "speculative_config": {
60
+ "method": "ngram",
61
+ "num_speculative_tokens": 5,
62
+ "prompt_lookup_max": 3,
63
+ "disable_mqa_scorer": True,
64
+ },
65
},
66
])
67
@pytest.mark.parametrize("output_len", [
@@ -101,16 +105,20 @@ def test_ngram_e2e_greedy_correctness(vllm_runner, common_llm_kwargs,
105
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
106
@pytest.mark.parametrize("test_llm_kwargs", [
107
{
104
- "speculative_model": "[ngram]",
105
- "num_speculative_tokens": 5,
106
- "ngram_prompt_lookup_max": 3,
107
- "disable_logprobs_during_spec_decoding": False,
108
+ "speculative_config": {
109
+ "method": "ngram",
110
+ "num_speculative_tokens": 5,
111
+ "prompt_lookup_max": 3,
112
+ "disable_logprobs": False,
113
+ },
114
},
115
{
110
- "speculative_model": "[ngram]",
111
- "num_speculative_tokens": 5,
112
- "ngram_prompt_lookup_max": 3,
113
- "disable_logprobs_during_spec_decoding": True,
116
+ "speculative_config": {
117
+ "method": "ngram",
118
+ "num_speculative_tokens": 5,
119
+ "prompt_lookup_max": 3,
120
+ "disable_logprobs": True,
121
+ },
122
},
123
])
124
@pytest.mark.parametrize("output_len", [
@@ -125,19 +133,20 @@ def test_ngram_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
133
batch_size: int, output_len: int, seed: int,
134
logprobs: int):
135
"""Verify greedy equality on a tiny model with different batch size."""
128
- run_equality_correctness_test(vllm_runner,
129
- common_llm_kwargs,
130
- per_test_common_llm_kwargs,
131
- baseline_llm_kwargs,
132
- test_llm_kwargs,
133
- batch_size,
134
- max_output_len=output_len,
135
- seed=seed,
136
- temperature=0.0,
137
- logprobs=logprobs,
138
- prompt_logprobs=logprobs,
139
- disable_logprobs=test_llm_kwargs[
140
- 'disable_logprobs_during_spec_decoding'])
136
+ run_equality_correctness_test(
137
+ vllm_runner,
138
+ common_llm_kwargs,
139
+ per_test_common_llm_kwargs,
140
+ baseline_llm_kwargs,
141
+ test_llm_kwargs,
142
+ batch_size,
143
+ max_output_len=output_len,
144
+ seed=seed,
145
+ temperature=0.0,
146
+ logprobs=logprobs,
147
+ prompt_logprobs=logprobs,
148
+ disable_logprobs=test_llm_kwargs["speculative_config"]
149
+ ["disable_logprobs"])
150
151
152
@pytest.mark.parametrize(
@@ -159,17 +168,21 @@ def test_ngram_e2e_greedy_logprobs(vllm_runner, common_llm_kwargs,
168
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
169
@pytest.mark.parametrize("test_llm_kwargs", [
170
{
162
- "speculative_model": "[ngram]",
163
- "num_speculative_tokens": 5,
164
- "ngram_prompt_lookup_max": 3,
171
+ "speculative_config": {
172
+ "method": "ngram",
173
+ "num_speculative_tokens": 5,
174
+ "prompt_lookup_max": 3,
175
+ },
176
"enable_chunked_prefill": False,
177
},
178
{
168
- "speculative_model": "[ngram]",
169
- "num_speculative_tokens": 5,
170
- "ngram_prompt_lookup_max": 3,
179
+ "speculative_config": {
180
+ "method": "ngram",
181
+ "num_speculative_tokens": 5,
182
+ "prompt_lookup_max": 3,
183
+ "disable_mqa_scorer": True,
184
+ },
185
"enable_chunked_prefill": True,
172
- "speculative_disable_mqa_scorer": True,
186
"max_num_batched_tokens": 4,
187
"max_num_seqs": 4
188
},
@@ -214,17 +227,21 @@ def test_ngram_e2e_greedy_correctness_with_preemption(
227
"test_llm_kwargs",
228
[
229
{
217
- "speculative_model": "[ngram]",
218
- "num_speculative_tokens": k,
219
- "ngram_prompt_lookup_max": 3,
230
+ "speculative_config": {
231
+ "method": "ngram",
232
+ "num_speculative_tokens": k,
233
+ "prompt_lookup_max": 3,
234
+ },
235
}
236
# Try a range of common k, as well as large speculation.
237
for k in [1, 3, 5]
238
] + [
239
{
225
- "speculative_model": "[ngram]",
226
- "num_speculative_tokens": k,
227
- "ngram_prompt_lookup_max": 1,
240
+ "speculative_config": {
241
+ "method": "ngram",
242
+ "num_speculative_tokens": k,
243
+ "prompt_lookup_max": 1,
244
+ },
245
}
246
# Try a range of common k, as well as large speculation.
247
for k in [1, 3, 5]
@@ -243,7 +260,7 @@ def test_ngram_different_k(vllm_runner, common_llm_kwargs,
260
seed: int):
261
"""Verify that ngram speculative decoding produces exact equality
262
to without spec decode with many different values of k and
246
- different ngram_prompt_lookup_max.
263
+ different ngram prompt_lookup_max.
264
"""
265
run_equality_correctness_test(vllm_runner,
266
common_llm_kwargs,
@@ -266,22 +283,25 @@ def test_ngram_different_k(vllm_runner, common_llm_kwargs,
283
}])
284
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
285
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
269
-@pytest.mark.parametrize("test_llm_kwargs",
270
- [{
271
- "speculative_model": "[ngram]",
272
- "num_speculative_tokens": 5,
273
- "ngram_prompt_lookup_max": 3,
274
- "speculative_disable_by_batch_size": 4
275
- }, {
276
- "speculative_model": "[ngram]",
277
- "num_speculative_tokens": 5,
278
- "ngram_prompt_lookup_max": 3,
279
- "speculative_disable_by_batch_size": 4,
280
- "enable_chunked_prefill": True,
281
- "speculative_disable_mqa_scorer": True,
282
- "max_num_batched_tokens": 4,
283
- "max_num_seqs": 4
284
- }])
286
+@pytest.mark.parametrize("test_llm_kwargs", [{
287
+ "speculative_config": {
288
+ "method": "ngram",
289
+ "num_speculative_tokens": 5,
290
+ "prompt_lookup_max": 3,
291
+ "disable_by_batch_size": 4
292
+ },
293
+}, {
294
+ "speculative_config": {
295
+ "method": "ngram",
296
+ "num_speculative_tokens": 5,
297
+ "prompt_lookup_max": 3,
298
+ "disable_by_batch_size": 4,
299
+ "disable_mqa_scorer": True,
300
+ },
301
+ "enable_chunked_prefill": True,
302
+ "max_num_batched_tokens": 4,
303
+ "max_num_seqs": 4
304
+}])
305
@pytest.mark.parametrize("batch_size", [1, 5])
306
@pytest.mark.parametrize(
307
"output_len",
@@ -296,7 +316,7 @@ def test_ngram_disable_queue(vllm_runner, common_llm_kwargs,
316
seed: int):
317
"""Verify that ngram speculative decoding produces exact equality
318
to without spec decode with many different values of k and
299
- different ngram_prompt_lookup_max.
319
+ different ngram prompt_lookup_max.
320
"""
321
run_equality_correctness_test(vllm_runner,
322
common_llm_kwargs,
@@ -316,18 +336,17 @@ def test_ngram_disable_queue(vllm_runner, common_llm_kwargs,
336
337
# Skip cuda graph recording for fast test.
338
"enforce_eager": True,
319
-
320
- # Required for spec decode.
321
- "speculative_model": "[ngram]",
322
- "num_speculative_tokens": 5,
323
- "ngram_prompt_lookup_max": 3,
339
}])
340
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
341
@pytest.mark.parametrize("baseline_llm_kwargs", [{}])
327
-@pytest.mark.parametrize("test_llm_kwargs",
328
- [{
329
- "speculative_disable_mqa_scorer": True,
330
- }])
342
+@pytest.mark.parametrize("test_llm_kwargs", [{
343
+ "speculative_config": {
344
+ "method": "ngram",
345
+ "num_speculative_tokens": 5,
346
+ "prompt_lookup_max": 3,
347
+ "disable_mqa_scorer": True,
348
+ },
349
+}])
350
@pytest.mark.parametrize("batch_size", [1, 5])
351
@pytest.mark.parametrize(
352
"output_len",
tests/spec_decode/e2e/test_seed.py
+5
-5
@@ -19,11 +19,11 @@ SPEC_MODEL = "JackFram/llama-160m"
19
# Skip cuda graph recording for fast test.
20
"enforce_eager": True,
21
22
- # speculative model
23
- "speculative_model": "JackFram/llama-160m",
24
-
25
- # num speculative tokens
26
- "num_speculative_tokens": 3,
22
+ # speculative config
23
+ "speculative_config": {
24
+ "model": "JackFram/llama-160m",
25
+ "num_speculative_tokens": 3,
26
+ },
27
}])
28
@pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
29
@pytest.mark.parametrize("baseline_llm_kwargs", [{"seed": 1}])
tests/v1/e2e/test_ngram_spec_decode.py
+10
-6
@@ -70,12 +70,16 @@ def test_ngram_correctness(
70
ref_outputs = ref_llm.chat(test_prompts, sampling_config)
71
del ref_llm
72
73
- spec_llm = LLM(model=model_name,
74
- speculative_model='[ngram]',
75
- ngram_prompt_lookup_max=5,
76
- ngram_prompt_lookup_min=3,
77
- num_speculative_tokens=3,
78
- max_model_len=1024)
73
+ spec_llm = LLM(
74
+ model=model_name,
75
+ speculative_config={
76
+ "method": "ngram",
77
+ "prompt_lookup_max": 5,
78
+ "prompt_lookup_min": 3,
79
+ "num_speculative_tokens": 3,
80
+ },
81
+ max_model_len=1024,
82
+ )
83
spec_outputs = spec_llm.chat(test_prompts, sampling_config)
84
matches = 0
85
misses = 0
vllm/config.py
+307
-301
@@ -1810,12 +1810,139 @@ class DeviceConfig:
1810
self.device = torch.device(self.device_type)
1811
1812
1813
+@dataclass
1814
class SpeculativeConfig:
1814
- """Configuration for speculative decoding.
1815
+ """
1816
+ Configuration for speculative decoding.
1817
+ Configurable parameters include:
1818
+ - General Speculative Decoding Control:
1819
+ - num_speculative_tokens (int): The number of speculative
1820
+ tokens, if provided. It will default to the number in the draft
1821
+ model config if present, otherwise, it is required.
1822
+ - model (Optional[str]): The name of the draft model, eagle head,
1823
+ or additional weights, if provided.
1824
+ - method (Optional[str]): The name of the speculative method to use.
1825
+ If users provide and set the `model` param, the speculative method
1826
+ type will be detected automatically if possible, if `model` param
1827
+ is not provided, the method name must be provided.
1828
+ - Possible values:
1829
+ - ngram
1830
+ Related additional configuration:
1831
+ - prompt_lookup_max (Optional[int]):
1832
+ Maximum size of ngram token window when using Ngram
1833
+ proposer, required when method is set to ngram.
1834
+ - prompt_lookup_min (Optional[int]):
1835
+ Minimum size of ngram token window when using Ngram
1836
+ proposer, if provided. Defaults to 1.
1837
+ - eagle
1838
+ - medusa
1839
+ - mlp_speculator
1840
+ - draft_model
1841
+ - acceptance_method (str): The method to use for accepting draft
1842
+ tokens. This can take two possible values: 'rejection_sampler' and
1843
+ 'typical_acceptance_sampler' for RejectionSampler and
1844
+ TypicalAcceptanceSampler respectively. If not specified, it
1845
+ defaults to 'rejection_sampler'.
1846
+ - Possible values:
1847
+ - rejection_sampler
1848
+ - typical_acceptance_sampler
1849
+ Related additional configuration:
1850
+ - posterior_threshold (Optional[float]):
1851
+ A threshold value that sets a lower bound on the
1852
+ posterior probability of a token in the target model
1853
+ for it to be accepted. This threshold is used only
1854
+ when we use the TypicalAcceptanceSampler for token
1855
+ acceptance.
1856
+ - posterior_alpha (Optional[float]):
1857
+ Scaling factor for entropy-based threshold, applied
1858
+ when using TypicalAcceptanceSampler.
1859
+ - draft_tensor_parallel_size (Optional[int]): The degree of the tensor
1860
+ parallelism for the draft model. Can only be 1 or the same as the
1861
+ target model's tensor parallel size.
1862
+ - disable_logprobs (bool): If set to True, token log probabilities are
1863
+ not returned during speculative decoding. If set to False, token
1864
+ log probabilities are returned according to the log probability
1865
+ settings in SamplingParams. If not specified, it defaults to True.
1866
+
1867
+ - Draft Model Configuration:
1868
+ - quantization (Optional[str]): Quantization method that was used to
1869
+ quantize the draft model weights. If None, we assume the
1870
+ model weights are not quantized. Note that it only takes effect
1871
+ when using the draft model-based speculative method.
1872
+ - max_model_len (Optional[int]): The maximum model length of the
1873
+ draft model. Used when testing the ability to skip
1874
+ speculation for some sequences.
1875
+ - revision: The specific model version to use for the draft model. It
1876
+ can be a branch name, a tag name, or a commit id. If unspecified,
1877
+ will use the default version.
1878
+ - code_revision: The specific revision to use for the draft model code
1879
+ on Hugging Face Hub. It can be a branch name, a tag name, or a
1880
+ commit id. If unspecified, will use the default version.
1881
1816
- The configuration is currently specialized to draft-model speculative
1817
- decoding with top-1 proposals.
1882
+ - Advanced Control:
1883
+ - disable_mqa_scorer (bool): Disable the MQA scorer and fall back to
1884
+ batch expansion for scoring proposals. If not specified, it
1885
+ defaults to False.
1886
+ - disable_by_batch_size (Optional[int]): Disable speculative decoding
1887
+ for new incoming requests when the number of enqueued requests is
1888
+ larger than this value, if provided.
1889
+
1890
+ Although the parameters above are structured hierarchically, there is no
1891
+ need to nest them during configuration.
1892
+
1893
+ Non-configurable internal parameters include:
1894
+ - Model Configuration:
1895
+ - target_model_config (ModelConfig): The configuration of the target
1896
+ model.
1897
+ - draft_model_config (ModelConfig): The configuration of the draft
1898
+ model initialized internal.
1899
+ - Parallelism Configuration:
1900
+ - target_parallel_config (ParallelConfig): The parallel configuration
1901
+ for the target model.
1902
+ - draft_parallel_config (ParallelConfig): The parallel configuration
1903
+ for the draft model initialized internal.
1904
+ - Execution Control:
1905
+ - enable_chunked_prefill (bool): Whether vLLM is configured to use
1906
+ chunked prefill or not. Used for raising an error since it's not
1907
+ yet compatible with speculative decode.
1908
+ - disable_log_stats (bool): Whether to disable the periodic printing of
1909
+ stage times in speculative decoding.
1910
"""
1911
+ # speculative configs from cli args
1912
+ num_speculative_tokens: int = field(default=None,
1913
+ init=True) # type: ignore
1914
+ method: Optional[str] = None
1915
+ acceptance_method: str = "rejection_sampler"
1916
+ draft_tensor_parallel_size: Optional[int] = None
1917
+ disable_logprobs: bool = True
1918
+
1919
+ model: Optional[str] = None
1920
+ quantization: Optional[str] = None
1921
+ max_model_len: Optional[int] = None
1922
+ revision: Optional[str] = None
1923
+ code_revision: Optional[str] = None
1924
+
1925
+ disable_mqa_scorer: bool = False
1926
+ disable_by_batch_size: Optional[int] = None
1927
+ prompt_lookup_max: Optional[int] = None
1928
+ prompt_lookup_min: Optional[int] = None
1929
+ posterior_threshold: Optional[float] = None
1930
+ posterior_alpha: Optional[float] = None
1931
+
1932
+ # required configuration params passed from engine
1933
+ target_model_config: ModelConfig = field(default=None,
1934
+ init=True) # type: ignore
1935
+ target_parallel_config: ParallelConfig = field(default=None,
1936
+ init=True) # type: ignore
1937
+ enable_chunked_prefill: bool = field(default=None,
1938
+ init=True) # type: ignore
1939
+ disable_log_stats: bool = field(default=None, init=True) # type: ignore
1940
+
1941
+ # params generated in the post-init stage
1942
+ draft_model_config: ModelConfig = field(default=None,
1943
+ init=True) # type: ignore
1944
+ draft_parallel_config: ParallelConfig = field(default=None,
1945
+ init=True) # type: ignore
1946
1947
def compute_hash(self) -> str:
1948
"""
@@ -1835,6 +1962,11 @@ class SpeculativeConfig:
1962
hash_str = hashlib.md5(str(factors).encode()).hexdigest()
1963
return hash_str
1964
1965
+ @classmethod
1966
+ def from_dict(cls, dict_value: dict) -> "SpeculativeConfig":
1967
+ """Parse the CLI value for the speculative config."""
1968
+ return cls(**dict_value)
1969
+
1970
@staticmethod
1971
def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
1972
if hf_config.model_type == "deepseek_v3":
@@ -1847,230 +1979,160 @@ class SpeculativeConfig:
1979
})
1980
return hf_config
1981
1850
- @staticmethod
1851
- def maybe_create_spec_config(
1852
- target_model_config: ModelConfig,
1853
- target_parallel_config: ParallelConfig,
1854
- target_dtype: str,
1855
- speculative_model: Optional[str],
1856
- speculative_model_quantization: Optional[str],
1857
- speculative_draft_tensor_parallel_size: Optional[int],
1858
- num_speculative_tokens: Optional[int],
1859
- speculative_disable_mqa_scorer: Optional[bool],
1860
- speculative_max_model_len: Optional[int],
1861
- enable_chunked_prefill: bool,
1862
- disable_log_stats: bool,
1863
- speculative_disable_by_batch_size: Optional[int],
1864
- ngram_prompt_lookup_max: Optional[int],
1865
- ngram_prompt_lookup_min: Optional[int],
1866
- draft_token_acceptance_method: str,
1867
- typical_acceptance_sampler_posterior_threshold: Optional[float],
1868
- typical_acceptance_sampler_posterior_alpha: Optional[float],
1869
- disable_logprobs: Optional[bool],
1870
- ) -> Optional["SpeculativeConfig"]:
1871
- """Create a SpeculativeConfig if possible, else return None.
1872
-
1873
- This function attempts to create a SpeculativeConfig object based on the
1874
- provided parameters. If the necessary conditions are met, it returns an
1875
- instance of SpeculativeConfig. Otherwise, it returns None.
1876
-
1877
- Args:
1878
- target_model_config (ModelConfig): The configuration of the target
1879
- model.
1880
- target_parallel_config (ParallelConfig): The parallel configuration
1881
- for the target model.
1882
- target_dtype (str): The data type used for the target model.
1883
- speculative_model (Optional[str]): The name of the speculative
1884
- model, if provided.
1885
- speculative_model_quantization (Optional[str]): Quantization method
1886
- that was used to quantize the speculative model weights. If
1887
- None, we assume the model weights are not quantized.
1888
- speculative_draft_tensor_parallel_size (Optional[int]): The degree
1889
- of the tensor parallelism for the draft model.
1890
- num_speculative_tokens (Optional[int]): The number of speculative
1891
- tokens, if provided. Will default to the number in the draft
1892
- model config if present, otherwise is required.
1893
- speculative_disable_mqa_scorer (Optional[bool]): Disable the MQA
1894
- scorer for the speculative model and fall back to batch
1895
- expansion for scoring.
1896
- speculative_max_model_len (Optional[int]): The maximum model len of
1897
- the speculative model. Used when testing the ability to skip
1898
- speculation for some sequences.
1899
- enable_chunked_prefill (bool): Whether vLLM is configured to use
1900
- chunked prefill or not. Used for raising an error since its not
1901
- yet compatible with spec decode.
1902
- speculative_disable_by_batch_size (Optional[int]): Disable
1903
- speculative decoding for new incoming requests when the number
1904
- of enqueue requests is larger than this value, if provided.
1905
- ngram_prompt_lookup_max (Optional[int]): Max size of ngram token
1906
- window, if provided.
1907
- ngram_prompt_lookup_min (Optional[int]): Min size of ngram token
1908
- window, if provided.
1909
- draft_token_acceptance_method (str): The method to use for
1910
- accepting draft tokens. This can take two possible
1911
- values 'rejection_sampler' and 'typical_acceptance_sampler'
1912
- for RejectionSampler and TypicalAcceptanceSampler
1913
- respectively.
1914
- typical_acceptance_sampler_posterior_threshold (Optional[float]):
1915
- A threshold value that sets a lower bound on the posterior
1916
- probability of a token in the target model for it to be
1917
- accepted. This threshold is used only when we use the
1918
- TypicalAcceptanceSampler for token acceptance.
1919
- typical_acceptance_sampler_posterior_alpha (Optional[float]):
1920
- A scaling factor for the entropy-based threshold in the
1921
- TypicalAcceptanceSampler.
1922
- disable_logprobs (Optional[bool]): If set to True, token log
1923
- probabilities are not returned during speculative decoding.
1924
- If set to False, token log probabilities are returned
1925
- according to the log probability settings in SamplingParams.
1926
- If not specified, it defaults to True.
1982
+ def __post_init__(self):
1983
1928
- Returns:
1929
- Optional["SpeculativeConfig"]: An instance of SpeculativeConfig if
1930
- the necessary conditions are met, else None.
1931
- """
1932
- if speculative_model is None:
1933
- if num_speculative_tokens is not None:
1934
- if target_model_config.hf_text_config.model_type \
1984
+ # Note: After next release, the method parameter will be used to
1985
+ # specify the speculative method, which helps to extend the
1986
+ # configuration of non-model-based proposers, and the model parameter
1987
+ # will be used when the draft model or head is needed.
1988
+ # If users do not specify the method, the speculative method will
1989
+ # be detected automatically if possible. If the speculative method can
1990
+ # not be detected, it will be considered as the draft-model-based
1991
+ # method by default.
1992
+
1993
+ if self.model is None and self.num_speculative_tokens is not None:
1994
+ # TODO(Shangming): Refactor mtp configuration logic when supporting
1995
+ # mtp acceleration for more models besides deepseek_v3
1996
+ if self.target_model_config.hf_text_config.model_type \
1997
== "deepseek_v3":
1936
- # use the draft model from the same model:
1937
- speculative_model = target_model_config.model
1938
- else:
1939
- raise ValueError(
1940
- "num_speculative_tokens was provided without "
1941
- "speculative_model.")
1998
+ # use the draft model from the same model:
1999
+ self.model = self.target_model_config.model
2000
+ elif self.method in ("ngram", "[ngram]"):
2001
+ self.model = "ngram"
2002
else:
1943
- return None
1944
-
1945
- if (speculative_disable_by_batch_size is not None
1946
- and speculative_disable_by_batch_size < 2):
1947
- raise ValueError("Expect the batch size threshold of disabling "
1948
- "speculative decoding is > 1, but got "
1949
- f"{speculative_disable_by_batch_size=}")
1950
- if (enable_chunked_prefill and speculative_model == "eagle"):
1951
- raise ValueError("Chunked prefill and EAGLE are not compatible.")
1952
- # TODO: The user should be able to specify revision/max model len
1953
- # for the draft model. It is not currently supported.
1954
- draft_revision = None
1955
- draft_code_revision = None
1956
- draft_quantization = speculative_model_quantization
1957
-
1958
- if speculative_model == "[ngram]":
1959
- if ngram_prompt_lookup_min is None:
1960
- ngram_prompt_lookup_min = 1
1961
- if ngram_prompt_lookup_max is None or ngram_prompt_lookup_max < 1:
1962
- raise ValueError(f"{ngram_prompt_lookup_max=} must be > 0")
1963
- if ngram_prompt_lookup_min < 1:
1964
- raise ValueError(f"{ngram_prompt_lookup_min=} must be > 0")
1965
- if ngram_prompt_lookup_min > ngram_prompt_lookup_max:
1966
- raise ValueError(f"{ngram_prompt_lookup_min=} cannot be "
1967
- f"larger than {ngram_prompt_lookup_max=}")
2003
+ raise ValueError("num_speculative_tokens was provided without "
2004
+ "speculative model.")
2005
+
2006
+ # Automatically configure the ngram method during configuration
2007
+ # refactoring to ensure a smooth transition.
2008
+ if self.method is None and (self.model is not None
2009
+ and self.model in ("ngram", "[ngram]")):
2010
+ self.method = "ngram"
2011
+
2012
+ if self.method in ("ngram", "[ngram]"):
2013
+ # Unified to "ngram" internally
2014
+ self.method = "ngram"
2015
+ if self.prompt_lookup_min is None:
2016
+ self.prompt_lookup_min = 1
2017
+ if self.prompt_lookup_max is None or self.prompt_lookup_max < 1:
2018
+ raise ValueError("prompt_lookup_max="
2019
+ f"{self.prompt_lookup_max} must be > 0")
2020
+ if self.prompt_lookup_min < 1:
2021
+ raise ValueError("prompt_lookup_min="
2022
+ f"{self.prompt_lookup_min} must be > 0")
2023
+ if self.prompt_lookup_min > self.prompt_lookup_max:
2024
+ raise ValueError(f"prompt_lookup_min={self.prompt_lookup_min} "
2025
+ "cannot be larger than prompt_lookup_max="
2026
+ f"{self.prompt_lookup_max}")
2027
2028
# TODO: current we still need extract vocab_size from target model
2029
# config, in future, we may try refactor it out, and set
2030
# draft related config as None here.
1972
- draft_model_config = target_model_config
1973
- draft_parallel_config = target_parallel_config
2031
+ self.draft_model_config = self.target_model_config
2032
+ self.draft_parallel_config = self.target_parallel_config
2033
else:
1975
- ngram_prompt_lookup_max = 0
1976
- ngram_prompt_lookup_min = 0
1977
- draft_model_config = ModelConfig(
1978
- model=speculative_model,
1979
- task="draft",
1980
- tokenizer=target_model_config.tokenizer,
1981
- tokenizer_mode=target_model_config.tokenizer_mode,
1982
- trust_remote_code=target_model_config.trust_remote_code,
1983
- allowed_local_media_path=target_model_config.
1984
- allowed_local_media_path,
1985
- dtype=target_model_config.dtype,
1986
- seed=target_model_config.seed,
1987
- revision=draft_revision,
1988
- code_revision=draft_code_revision,
1989
- tokenizer_revision=target_model_config.tokenizer_revision,
1990
- max_model_len=None,
1991
- spec_target_max_model_len=target_model_config.max_model_len,
1992
- quantization=draft_quantization,
1993
- enforce_eager=target_model_config.enforce_eager,
1994
- max_seq_len_to_capture=target_model_config.
1995
- max_seq_len_to_capture,
1996
- max_logprobs=target_model_config.max_logprobs,
1997
- hf_overrides=SpeculativeConfig.hf_config_override,
1998
- )
1999
-
2000
- draft_hf_config = draft_model_config.hf_config
2034
+ self.prompt_lookup_max = 0
2035
+ self.prompt_lookup_min = 0
2036
+
2037
+ if self.model is not None:
2038
+ self.draft_model_config = ModelConfig(
2039
+ model=self.model,
2040
+ task="draft",
2041
+ tokenizer=self.target_model_config.tokenizer,
2042
+ tokenizer_mode=self.target_model_config.tokenizer_mode,
2043
+ trust_remote_code=self.target_model_config.
2044
+ trust_remote_code,
2045
+ allowed_local_media_path=self.target_model_config.
2046
+ allowed_local_media_path,
2047
+ dtype=self.target_model_config.dtype,
2048
+ seed=self.target_model_config.seed,
2049
+ revision=self.revision,
2050
+ code_revision=self.code_revision,
2051
+ tokenizer_revision=self.target_model_config.
2052
+ tokenizer_revision,
2053
+ max_model_len=None,
2054
+ spec_target_max_model_len=self.target_model_config.
2055
+ max_model_len,
2056
+ quantization=self.quantization,
2057
+ enforce_eager=self.target_model_config.enforce_eager,
2058
+ max_seq_len_to_capture=self.target_model_config.
2059
+ max_seq_len_to_capture,
2060
+ max_logprobs=self.target_model_config.max_logprobs,
2061
+ hf_overrides=SpeculativeConfig.hf_config_override,
2062
+ )
2063
2002
- # Detect EAGLE prefix to replace hf_config for EAGLE draft_model
2003
- if "eagle-" in draft_model_config.model.lower():
2004
- from vllm.transformers_utils.configs.eagle import EAGLEConfig
2005
- if isinstance(draft_model_config.hf_config, EAGLEConfig):
2006
- pass
2064
+ # Automatically detect the method
2065
+ if "eagle-" in self.draft_model_config.model.lower():
2066
+ self.method = "eagle"
2067
+ elif self.draft_model_config.hf_config.model_type == "medusa":
2068
+ self.method = "medusa"
2069
+ elif (self.draft_model_config.hf_config.model_type ==
2070
+ "mlp_speculator"):
2071
+ self.method = "mlp_speculator"
2072
else:
2008
- eagle_config = EAGLEConfig(draft_model_config.hf_config)
2009
- draft_model_config.hf_config = eagle_config
2010
-
2011
- if (num_speculative_tokens is not None
2012
- and hasattr(draft_hf_config, "num_lookahead_tokens")):
2013
- draft_hf_config.num_lookahead_tokens = num_speculative_tokens
2014
- n_predict = getattr(draft_hf_config, "n_predict", None)
2015
- if n_predict is not None:
2016
- if num_speculative_tokens is None:
2017
- # Default to max value defined in draft model config.
2018
- num_speculative_tokens = n_predict
2019
- elif num_speculative_tokens > n_predict and \
2020
- num_speculative_tokens % n_predict != 0:
2021
- # Ensure divisibility for MTP module reuse.
2022
- raise ValueError(
2023
- f"{num_speculative_tokens=} must be divisible by "
2024
- f"{n_predict=}")
2025
-
2026
- speculative_draft_tensor_parallel_size = \
2027
- SpeculativeConfig._verify_and_get_draft_model_tensor_parallel_size(
2028
- target_parallel_config,
2029
- speculative_draft_tensor_parallel_size,
2030
- draft_hf_config
2031
- )
2073
+ self.method = "draft_model"
2074
+
2075
+ # Replace hf_config for EAGLE draft_model
2076
+ if self.method == "eagle":
2077
+ if self.enable_chunked_prefill:
2078
+ raise ValueError(
2079
+ "Chunked prefill and EAGLE are not compatible.")
2080
+
2081
+ from vllm.transformers_utils.configs.eagle import (
2082
+ EAGLEConfig)
2083
+ if isinstance(self.draft_model_config.hf_config,
2084
+ EAGLEConfig):
2085
+ pass
2086
+ else:
2087
+ eagle_config = EAGLEConfig(
2088
+ self.draft_model_config.hf_config)
2089
+ self.draft_model_config.hf_config = eagle_config
2090
+
2091
+ if (self.num_speculative_tokens is not None
2092
+ and hasattr(self.draft_model_config.hf_config,
2093
+ "num_lookahead_tokens")):
2094
+ self.draft_model_config.hf_config.num_lookahead_tokens = \
2095
+ self.num_speculative_tokens
2096
+
2097
+ n_predict = getattr(self.draft_model_config.hf_config,
2098
+ "n_predict", None)
2099
+ if n_predict is not None:
2100
+ if self.num_speculative_tokens is None:
2101
+ # Default to max value defined in draft model config.
2102
+ self.num_speculative_tokens = n_predict
2103
+ elif self.num_speculative_tokens > n_predict and \
2104
+ self.num_speculative_tokens % n_predict != 0:
2105
+ # Ensure divisibility for MTP module reuse.
2106
+ raise ValueError(
2107
+ f"num_speculative_tokens:{self.num_speculative_tokens}"
2108
+ f" must be divisible by {n_predict=}")
2109
+
2110
+ self.draft_tensor_parallel_size = \
2111
+ SpeculativeConfig._verify_and_get_draft_tp(
2112
+ self.target_parallel_config,
2113
+ self.draft_tensor_parallel_size,
2114
+ self.draft_model_config.hf_config
2115
+ )
2116
2033
- draft_model_config.max_model_len = (
2034
- SpeculativeConfig._maybe_override_draft_max_model_len(
2035
- speculative_max_model_len,
2036
- draft_model_config.max_model_len,
2037
- target_model_config.max_model_len,
2038
- ))
2117
+ self.draft_model_config.max_model_len = (
2118
+ SpeculativeConfig._maybe_override_draft_max_model_len(
2119
+ self.max_model_len,
2120
+ self.draft_model_config.max_model_len,
2121
+ self.target_model_config.max_model_len,
2122
+ ))
2123
2040
- draft_parallel_config = (
2041
- SpeculativeConfig.create_draft_parallel_config(
2042
- target_parallel_config,
2043
- speculative_draft_tensor_parallel_size, draft_hf_config))
2124
+ self.draft_parallel_config = (
2125
+ SpeculativeConfig.create_draft_parallel_config(
2126
+ self.target_parallel_config,
2127
+ self.draft_tensor_parallel_size))
2128
2045
- if num_speculative_tokens is None:
2046
- raise ValueError(
2047
- "num_speculative_tokens must be provided with "
2048
- "speculative_model unless the draft model config contains an "
2049
- "n_predict parameter.")
2129
+ if self.acceptance_method == "typical_acceptance_sampler":
2130
+ if self.posterior_threshold is None:
2131
+ self.posterior_threshold = 0.09
2132
+ if self.posterior_alpha is None:
2133
+ self.posterior_alpha = 0.3
2134
2051
- if typical_acceptance_sampler_posterior_threshold is None:
2052
- typical_acceptance_sampler_posterior_threshold = 0.09
2053
- if typical_acceptance_sampler_posterior_alpha is None:
2054
- typical_acceptance_sampler_posterior_alpha = 0.3
2055
- if disable_logprobs is None:
2056
- disable_logprobs = True
2057
-
2058
- return SpeculativeConfig(
2059
- draft_model_config,
2060
- draft_parallel_config,
2061
- num_speculative_tokens,
2062
- speculative_disable_mqa_scorer,
2063
- speculative_disable_by_batch_size,
2064
- ngram_prompt_lookup_max,
2065
- ngram_prompt_lookup_min,
2066
- draft_token_acceptance_method=draft_token_acceptance_method,
2067
- typical_acceptance_sampler_posterior_threshold=\
2068
- typical_acceptance_sampler_posterior_threshold,
2069
- typical_acceptance_sampler_posterior_alpha=\
2070
- typical_acceptance_sampler_posterior_alpha,
2071
- disable_logprobs=disable_logprobs,
2072
- disable_log_stats=disable_log_stats,
2073
- )
2135
+ self._verify_args()
2136
2137
@staticmethod
2138
def _maybe_override_draft_max_model_len(
@@ -2108,7 +2170,7 @@ class SpeculativeConfig:
2170
)
2171
2172
@staticmethod
2111
- def _verify_and_get_draft_model_tensor_parallel_size(
2173
+ def _verify_and_get_draft_tp(
2174
target_parallel_config: ParallelConfig,
2175
speculative_draft_tensor_parallel_size: Optional[int],
2176
draft_hf_config: PretrainedConfig) -> int:
@@ -2140,7 +2202,6 @@ class SpeculativeConfig:
2202
def create_draft_parallel_config(
2203
target_parallel_config: ParallelConfig,
2204
speculative_draft_tensor_parallel_size: int,
2143
- draft_hf_config: PretrainedConfig,
2205
) -> ParallelConfig:
2206
"""Create a parallel config for use by the draft worker.
2207
@@ -2164,74 +2225,13 @@ class SpeculativeConfig:
2225
2226
return draft_parallel_config
2227
2167
- def __init__(
2168
- self,
2169
- draft_model_config: ModelConfig,
2170
- draft_parallel_config: ParallelConfig,
2171
- num_speculative_tokens: int,
2172
- speculative_disable_mqa_scorer: Optional[bool],
2173
- speculative_disable_by_batch_size: Optional[int],
2174
- ngram_prompt_lookup_max: Optional[int],
2175
- ngram_prompt_lookup_min: Optional[int],
2176
- draft_token_acceptance_method: str,
2177
- typical_acceptance_sampler_posterior_threshold: float,
2178
- typical_acceptance_sampler_posterior_alpha: float,
2179
- disable_logprobs: bool,
2180
- disable_log_stats: bool,
2181
- ):
2182
- """Create a SpeculativeConfig object.
2183
-
2184
- Args:
2185
- draft_model_config: ModelConfig for the draft model.
2186
- draft_parallel_config: ParallelConfig for the draft model.
2187
- num_speculative_tokens: The number of tokens to sample from the
2188
- draft model before scoring with the target model.
2189
- speculative_disable_by_batch_size: Disable speculative
2190
- decoding for new incoming requests when the number of
2191
- enqueue requests is larger than this value.
2192
- ngram_prompt_lookup_max: Max size of ngram token window.
2193
- ngram_prompt_lookup_min: Min size of ngram token window.
2194
- draft_token_acceptance_method (str): The method to use for
2195
- accepting draft tokens. This can take two possible
2196
- values 'rejection_sampler' and 'typical_acceptance_sampler'
2197
- for RejectionSampler and TypicalAcceptanceSampler
2198
- respectively.
2199
- typical_acceptance_sampler_posterior_threshold (Optional[float]):
2200
- A threshold value that sets a lower bound on the posterior
2201
- probability of a token in the target model for it to be
2202
- accepted. This threshold is used only when we use the
2203
- TypicalAcceptanceSampler for token acceptance.
2204
- typical_acceptance_sampler_posterior_alpha (Optional[float]):
2205
- A scaling factor for the entropy-based threshold in the
2206
- TypicalAcceptanceSampler.
2207
- disable_logprobs: If set to True, token log probabilities will not
2208
- be returned even if requested by sampling parameters. This
2209
- reduces latency by skipping logprob calculation in proposal
2210
- sampling, target sampling, and after accepted tokens are
2211
- determined. If set to False, log probabilities will be
2212
- returned.
2213
- disable_log_stats: Whether to disable periodic printing of stage
2214
- times in speculative decoding.
2215
- """
2216
- self.draft_model_config = draft_model_config
2217
- self.draft_parallel_config = draft_parallel_config
2218
- self.num_speculative_tokens = num_speculative_tokens
2219
- self.speculative_disable_mqa_scorer = speculative_disable_mqa_scorer
2220
- self.speculative_disable_by_batch_size = \
2221
- speculative_disable_by_batch_size
2222
- self.ngram_prompt_lookup_max = ngram_prompt_lookup_max or 0
2223
- self.ngram_prompt_lookup_min = ngram_prompt_lookup_min or 0
2224
- self.draft_token_acceptance_method = draft_token_acceptance_method
2225
- self.typical_acceptance_sampler_posterior_threshold = \
2226
- typical_acceptance_sampler_posterior_threshold
2227
- self.typical_acceptance_sampler_posterior_alpha = \
2228
- typical_acceptance_sampler_posterior_alpha
2229
- self.disable_logprobs = disable_logprobs
2230
- self.disable_log_stats = disable_log_stats
2231
-
2232
- self._verify_args()
2233
-
2228
def _verify_args(self) -> None:
2229
+ if self.num_speculative_tokens is None:
2230
+ raise ValueError(
2231
+ "num_speculative_tokens must be provided with "
2232
+ "speculative model unless the draft model config contains an "
2233
+ "n_predict parameter.")
2234
+
2235
if self.num_speculative_tokens <= 0:
2236
raise ValueError("Expected num_speculative_tokens to be greater "
2237
f"than zero ({self.num_speculative_tokens}).")
@@ -2241,29 +2241,34 @@ class SpeculativeConfig:
2241
self.draft_parallel_config)
2242
# Validate and set draft token acceptance related settings.
2243
2244
- if (self.draft_token_acceptance_method is None):
2245
- raise ValueError("draft_token_acceptance_method is not set. "
2244
+ if self.acceptance_method is None:
2245
+ raise ValueError("acceptance_method is not set. "
2246
"Expected values are rejection_sampler or "
2247
"typical_acceptance_sampler.")
2248
2249
- if (self.draft_token_acceptance_method != 'rejection_sampler'
2250
- and self.draft_token_acceptance_method
2251
- != 'typical_acceptance_sampler'):
2249
+ if (self.acceptance_method != 'rejection_sampler'
2250
+ and self.acceptance_method != 'typical_acceptance_sampler'):
2251
raise ValueError(
2253
- "Expected draft_token_acceptance_method to be either "
2252
+ "Expected acceptance_method to be either "
2253
"rejection_sampler or typical_acceptance_sampler. Instead it "
2255
- f"is {self.draft_token_acceptance_method}")
2254
+ f"is {self.acceptance_method}")
2255
2257
- if (self.typical_acceptance_sampler_posterior_threshold < 0
2258
- or self.typical_acceptance_sampler_posterior_alpha < 0):
2256
+ if self.acceptance_method == "typical_acceptance_sampler" and (
2257
+ (self.posterior_threshold is not None
2258
+ and self.posterior_threshold < 0) or
2259
+ (self.posterior_alpha is not None and self.posterior_alpha < 0)):
2260
raise ValueError(
2260
- "Expected typical_acceptance_sampler_posterior_threshold "
2261
- "and typical_acceptance_sampler_posterior_alpha to be > 0. "
2262
- "Instead found "
2263
- f"typical_acceptance_sampler_posterior_threshold = "
2264
- f"{self.typical_acceptance_sampler_posterior_threshold} and "
2265
- f"typical_acceptance_sampler_posterior_alpha = "
2266
- f"{self.typical_acceptance_sampler_posterior_alpha}")
2261
+ "Expected the posterior_threshold and posterior_alpha of "
2262
+ "typical_acceptance_sampler to be > 0. "
2263
+ "Instead found posterior_threshold = "
2264
+ f"{self.posterior_threshold} and posterior_alpha = "
2265
+ f"{self.posterior_alpha}")
2266
+
2267
+ if (self.disable_by_batch_size is not None
2268
+ and self.disable_by_batch_size < 2):
2269
+ raise ValueError("Expect the batch size threshold of disabling "
2270
+ "speculative decoding is > 1, but got "
2271
+ f"{self.disable_by_batch_size=}")
2272
2273
@property
2274
def num_lookahead_slots(self) -> int:
@@ -2276,8 +2281,8 @@ class SpeculativeConfig:
2281
return self.num_speculative_tokens
2282
2283
def __repr__(self) -> str:
2279
- if self.ngram_prompt_lookup_max > 0:
2280
- draft_model = "[ngram]"
2284
+ if self.prompt_lookup_max is not None and self.prompt_lookup_max > 0:
2285
+ draft_model = "ngram"
2286
else:
2287
draft_model = self.draft_model_config.model
2288
num_spec_tokens = self.num_speculative_tokens
@@ -3285,7 +3290,8 @@ class VllmConfig:
3290
init=True) # type: ignore
3291
load_config: LoadConfig = field(default=None, init=True) # type: ignore
3292
lora_config: Optional[LoRAConfig] = None
3288
- speculative_config: Optional[SpeculativeConfig] = None
3293
+ speculative_config: SpeculativeConfig = field(default=None,
3294
+ init=True) # type: ignore
3295
decoding_config: Optional[DecodingConfig] = None
3296
observability_config: Optional[ObservabilityConfig] = None
3297
prompt_adapter_config: Optional[PromptAdapterConfig] = None
vllm/engine/arg_utils.py
+92
-26
@@ -177,7 +177,10 @@ class EngineArgs:
177
178
guided_decoding_backend: str = 'xgrammar'
179
logits_processor_pattern: Optional[str] = None
180
- # Speculative decoding configuration.
180
+
181
+ speculative_config: Optional[Union[str, Dict[str, Any]]] = None
182
+
183
+ # TODO(Shangming): Deprecate these out-of-date params after next release
184
speculative_model: Optional[str] = None
185
speculative_model_quantization: Optional[str] = None
186
speculative_draft_tensor_parallel_size: Optional[int] = None
@@ -190,9 +193,9 @@ class EngineArgs:
193
spec_decoding_acceptance_method: str = 'rejection_sampler'
194
typical_acceptance_sampler_posterior_threshold: Optional[float] = None
195
typical_acceptance_sampler_posterior_alpha: Optional[float] = None
193
- qlora_adapter_name_or_path: Optional[str] = None
196
disable_logprobs_during_spec_decoding: Optional[bool] = None
197
198
+ qlora_adapter_name_or_path: Optional[str] = None
199
show_hidden_metrics_for_version: Optional[str] = None
200
otlp_traces_endpoint: Optional[str] = None
201
collect_detailed_traces: Optional[str] = None
@@ -780,7 +783,11 @@ class EngineArgs:
783
const="True",
784
help='If set, the prefill requests can be chunked based on the '
785
'max_num_batched_tokens.')
783
-
786
+ parser.add_argument('--speculative-config',
787
+ type=nullable_str,
788
+ default=None,
789
+ help='The configurations for speculative decoding.'
790
+ ' Should be a JSON string.')
791
parser.add_argument(
792
'--speculative-model',
793
type=nullable_str,
@@ -1192,6 +1199,82 @@ class EngineArgs:
1199
use_tqdm_on_load=self.use_tqdm_on_load,
1200
)
1201
1202
+ def create_speculative_config(
1203
+ self,
1204
+ target_model_config: ModelConfig,
1205
+ target_parallel_config: ParallelConfig,
1206
+ enable_chunked_prefill: bool,
1207
+ disable_log_stats: bool,
1208
+ ) -> Optional["SpeculativeConfig"]:
1209
+ """Initializes and returns a SpeculativeConfig object based on
1210
+ `speculative_config`.
1211
+
1212
+ This function utilizes `speculative_config` to create a
1213
+ SpeculativeConfig object. The `speculative_config` can either be
1214
+ provided as a JSON string input via CLI arguments or directly as a
1215
+ dictionary from the engine. If `speculative_config` is not set, this
1216
+ function will attempt to construct a configuration dictionary using
1217
+ certain parameters, which are scheduled for deprecation in the next
1218
+ release. Note that in next releases, `speculative_config` must be
1219
+ provided, and the deprecated standalone speculative-related parameters
1220
+ will be removed.
1221
+ """
1222
+ if self.speculative_config is None:
1223
+ if (self.speculative_model is None
1224
+ and self.num_speculative_tokens is None):
1225
+ return None
1226
+
1227
+ # TODO(Shangming): Deprecate this way of setting SpeculativeConfig,
1228
+ # only allow '--speculative-config' after next release
1229
+ logger.warning_once(
1230
+ "Please use '--speculative-config' to set all configurations "
1231
+ "related to speculative decoding. The current method of "
1232
+ "specifying the model through '--speculative-model' and "
1233
+ "adding related parameters (e.g., '--num-speculative-tokens') "
1234
+ "separately will be deprecated in the next release.")
1235
+
1236
+ spec_config_dict = {
1237
+ "model": self.speculative_model,
1238
+ "quantization": self.speculative_model_quantization,
1239
+ "max_model_len": self.speculative_max_model_len,
1240
+ "draft_tensor_parallel_size":
1241
+ self.speculative_draft_tensor_parallel_size,
1242
+ "num_speculative_tokens": self.num_speculative_tokens,
1243
+ "disable_mqa_scorer": self.speculative_disable_mqa_scorer,
1244
+ "disable_by_batch_size":
1245
+ self.speculative_disable_by_batch_size,
1246
+ "prompt_lookup_max": self.ngram_prompt_lookup_max,
1247
+ "prompt_lookup_min": self.ngram_prompt_lookup_min,
1248
+ "acceptance_method": self.spec_decoding_acceptance_method,
1249
+ "posterior_threshold":
1250
+ self.typical_acceptance_sampler_posterior_threshold,
1251
+ "posterior_alpha":
1252
+ self.typical_acceptance_sampler_posterior_alpha,
1253
+ "disable_logprobs": self.disable_logprobs_during_spec_decoding,
1254
+ }
1255
+
1256
+ self.speculative_config = spec_config_dict
1257
+ else:
1258
+ if isinstance(self.speculative_config, str):
1259
+ import ast
1260
+ self.speculative_config = ast.literal_eval(
1261
+ self.speculative_config)
1262
+ # Note(Shangming): These parameters are not obtained from the cli arg
1263
+ # '--speculative-config' and must be passed in when creating the engine
1264
+ # config.
1265
+
1266
+ assert isinstance(self.speculative_config, dict)
1267
+ self.speculative_config.update({
1268
+ "target_model_config": target_model_config,
1269
+ "target_parallel_config": target_parallel_config,
1270
+ "enable_chunked_prefill": enable_chunked_prefill,
1271
+ "disable_log_stats": disable_log_stats,
1272
+ })
1273
+ speculative_config = SpeculativeConfig.from_dict(
1274
+ self.speculative_config)
1275
+
1276
+ return speculative_config
1277
+
1278
def create_engine_config(
1279
self,
1280
usage_context: Optional[UsageContext] = None,
@@ -1238,6 +1321,8 @@ class EngineArgs:
1321
else:
1322
self._set_default_args_v0(model_config)
1323
1324
+ assert self.enable_chunked_prefill is not None
1325
+
1326
cache_config = CacheConfig(
1327
block_size=self.block_size,
1328
gpu_memory_utilization=self.gpu_memory_utilization,
@@ -1280,31 +1365,11 @@ class EngineArgs:
1365
worker_extension_cls=self.worker_extension_cls,
1366
)
1367
1283
- speculative_config = SpeculativeConfig.maybe_create_spec_config(
1368
+ speculative_config = self.create_speculative_config(
1369
target_model_config=model_config,
1370
target_parallel_config=parallel_config,
1286
- target_dtype=self.dtype,
1287
- speculative_model=self.speculative_model,
1288
- speculative_model_quantization = \
1289
- self.speculative_model_quantization,
1290
- speculative_draft_tensor_parallel_size = \
1291
- self.speculative_draft_tensor_parallel_size,
1292
- num_speculative_tokens=self.num_speculative_tokens,
1293
- speculative_disable_mqa_scorer=self.speculative_disable_mqa_scorer,
1294
- speculative_disable_by_batch_size=self.
1295
- speculative_disable_by_batch_size,
1296
- speculative_max_model_len=self.speculative_max_model_len,
1371
enable_chunked_prefill=self.enable_chunked_prefill,
1372
disable_log_stats=self.disable_log_stats,
1299
- ngram_prompt_lookup_max=self.ngram_prompt_lookup_max,
1300
- ngram_prompt_lookup_min=self.ngram_prompt_lookup_min,
1301
- draft_token_acceptance_method=\
1302
- self.spec_decoding_acceptance_method,
1303
- typical_acceptance_sampler_posterior_threshold=self.
1304
- typical_acceptance_sampler_posterior_threshold,
1305
- typical_acceptance_sampler_posterior_alpha=self.
1306
- typical_acceptance_sampler_posterior_alpha,
1307
- disable_logprobs=self.disable_logprobs_during_spec_decoding,
1373
)
1374
1375
# Reminder: Please update docs/source/features/compatibility_matrix.md
@@ -1569,7 +1634,7 @@ class EngineArgs:
1634
if (self.speculative_model is not None
1635
or self.num_speculative_tokens is not None):
1636
# This is supported but experimental (handled below).
1572
- if self.speculative_model == "[ngram]":
1637
+ if self.speculative_model in ("ngram", "[ngram]"):
1638
pass
1639
else:
1640
_raise_or_fallback(feature_name="Speculative Decoding",
@@ -1617,7 +1682,8 @@ class EngineArgs:
1682
return False
1683
1684
# ngram is supported on V1, but off by default for now.
1620
- if self.speculative_model == "[ngram]" and _warn_or_fallback("ngram"):
1685
+ if self.speculative_model in (
1686
+ "ngram", "[ngram]") and _warn_or_fallback("ngram"):
1687
return False
1688
1689
# Non-CUDA is supported on V1, but off by default for now.
vllm/spec_decode/spec_decode_worker.py
+7
-9
@@ -92,22 +92,20 @@ def create_spec_worker(*args, **kwargs) -> "SpecDecodeWorker":
92
# Override draft-model specific worker args.
93
draft_worker_kwargs.update(
94
vllm_config=draft_worker_config,
95
- ngram_prompt_lookup_max=speculative_config.ngram_prompt_lookup_max,
96
- ngram_prompt_lookup_min=speculative_config.ngram_prompt_lookup_min,
95
+ ngram_prompt_lookup_max=speculative_config.prompt_lookup_max,
96
+ ngram_prompt_lookup_min=speculative_config.prompt_lookup_min,
97
)
98
99
spec_decode_worker = SpecDecodeWorker.create_worker(
100
scorer_worker=target_worker,
101
draft_worker_kwargs=draft_worker_kwargs,
102
- disable_mqa_scorer=speculative_config.speculative_disable_mqa_scorer,
103
- disable_by_batch_size=speculative_config.
104
- speculative_disable_by_batch_size,
105
- draft_token_acceptance_method=speculative_config.
106
- draft_token_acceptance_method,
102
+ disable_mqa_scorer=speculative_config.disable_mqa_scorer,
103
+ disable_by_batch_size=speculative_config.disable_by_batch_size,
104
+ draft_token_acceptance_method=speculative_config.acceptance_method,
105
typical_acceptance_sampler_posterior_threshold=speculative_config.
108
- typical_acceptance_sampler_posterior_threshold,
106
+ posterior_threshold,
107
typical_acceptance_sampler_posterior_alpha=speculative_config.
110
- typical_acceptance_sampler_posterior_alpha,
108
+ posterior_alpha,
109
disable_logprobs=speculative_config.disable_logprobs,
110
disable_log_stats=speculative_config.disable_log_stats,
111
num_speculative_tokens=speculative_config.num_speculative_tokens,
vllm/v1/worker/gpu_model_runner.py
+3
-4
@@ -151,8 +151,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
151
self.use_spec_decode = False
152
if self.speculative_config:
153
self.use_spec_decode = True
154
- # TODO: find a better way to check if we are using ngram.
155
- assert self.speculative_config.ngram_prompt_lookup_min, \
154
+ assert self.speculative_config.method == "ngram", \
155
"Currently, only ngram spec decode is supported in V1."
156
if get_pp_group().is_last_rank:
157
self.drafter = NgramProposer()
@@ -160,7 +159,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
159
# This usually takes less than 1 second.
160
self.drafter.propose(
161
np.zeros(1024, dtype=np.int32),
163
- self.speculative_config.ngram_prompt_lookup_min,
162
+ self.speculative_config.prompt_lookup_min,
163
self.speculative_config.num_speculative_tokens,
164
)
165
self.rejection_sampler = RejectionSampler()
@@ -1155,7 +1154,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1154
self.input_batch.token_ids_cpu[i, start_idx:end_idx] = sampled_ids
1155
drafter_output = self.drafter.propose(
1156
self.input_batch.token_ids_cpu[i, :end_idx],
1158
- self.speculative_config.ngram_prompt_lookup_min,
1157
+ self.speculative_config.prompt_lookup_min,
1158
self.speculative_config.num_speculative_tokens,
1159
)
1160
if drafter_output is None or len(drafter_output) == 0: