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, 19:45 UTC
[V1][Bugfix][Spec Decode] Fix incorrect outputs in V1 speculative decoding due to batch indexing (#14645)
Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai>
Benjamin Chislett committed
Mar 12, 2025 at 01:12 UTC
5c538c37b26f8ebb9250dfc6a1866ec4aaac4299
2 files changed
+50
-15
tests/v1/e2e/test_ngram_spec_decode.py
+49
-12
@@ -1,4 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+import random
3
+
4
import pytest
5
6
from vllm import LLM, SamplingParams
@@ -6,16 +8,41 @@ from vllm import LLM, SamplingParams
8
9
@pytest.fixture
10
def test_prompts():
9
- return [
10
- "Can you repeat the sentence ten times, this is a sentence.",
11
- "Can you repeat the sentence ten times, this is a test.",
12
- ]
11
+ prompt_types = ["repeat", "sentence"]
12
+ num_prompts = 100
13
+ prompts = []
14
+
15
+ random.seed(0)
16
+ random_prompt_type_choices = random.choices(prompt_types, k=num_prompts)
17
+
18
+ # Generate a mixed batch of prompts, some of which can be easily
19
+ # predicted by n-gram matching and some which likely cannot.
20
+ for kind in random_prompt_type_choices:
21
+ word_choices = ["test", "temp", "hello", "where"]
22
+ word = random.choice(word_choices)
23
+ if kind == "repeat":
24
+ prompt = f"""
25
+ please repeat the word '{word}' 10 times.
26
+ give no other output than the word at least ten times in a row,
27
+ in lowercase with spaces between each word and without quotes.
28
+ """
29
+ elif kind == "sentence":
30
+ prompt = f"""
31
+ please give a ten-word sentence that
32
+ uses the word {word} at least once.
33
+ give no other output than that simple sentence without quotes.
34
+ """
35
+ else:
36
+ raise ValueError(f"Unknown prompt type: {kind}")
37
+ prompts.append([{"role": "user", "content": prompt}])
38
+
39
+ return prompts
40
41
42
@pytest.fixture
43
def sampling_config():
44
# Only support greedy for now
18
- return SamplingParams(temperature=0, max_tokens=30, ignore_eos=False)
45
+ return SamplingParams(temperature=0, max_tokens=10, ignore_eos=False)
46
47
48
@pytest.fixture
@@ -32,18 +59,28 @@ def test_ngram_correctness(monkeypatch, test_prompts, sampling_config,
59
with monkeypatch.context() as m:
60
m.setenv("VLLM_USE_V1", "1")
61
35
- ref_llm = LLM(model=model_name)
36
- ref_outputs = ref_llm.generate(test_prompts, sampling_config)
62
+ ref_llm = LLM(model=model_name, max_model_len=1024)
63
+ ref_outputs = ref_llm.chat(test_prompts, sampling_config)
64
del ref_llm
65
66
spec_llm = LLM(model=model_name,
67
speculative_model='[ngram]',
68
ngram_prompt_lookup_max=5,
69
ngram_prompt_lookup_min=3,
43
- num_speculative_tokens=3)
44
- spec_outputs = spec_llm.generate(test_prompts, sampling_config)
70
+ num_speculative_tokens=3,
71
+ max_model_len=1024)
72
+ spec_outputs = spec_llm.chat(test_prompts, sampling_config)
73
+ matches = 0
74
+ misses = 0
75
for ref_output, spec_output in zip(ref_outputs, spec_outputs):
46
- assert ref_output.outputs[0].text == spec_output.outputs[0].text, \
47
- (f"ref_output: {ref_output.outputs[0].text},"
48
- f"spec_output: {spec_output.outputs[0].text}")
76
+ if ref_output.outputs[0].text == spec_output.outputs[0].text:
77
+ matches += 1
78
+ else:
79
+ misses += 1
80
+ print(f"ref_output: {ref_output.outputs[0].text}")
81
+ print(f"spec_output: {spec_output.outputs[0].text}")
82
+
83
+ # Heuristic: expect at least 70% of the prompts to match exactly
84
+ # Upon failure, inspect the outputs to check for inaccuracy.
85
+ assert matches > int(0.7 * len(ref_outputs))
86
del spec_llm
vllm/v1/worker/gpu_model_runner.py
+1
-3
@@ -1015,11 +1015,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1015
else:
1016
target_probs = self.model.sampler.compute_probs(
1017
logits, sampling_metadata)
1018
- scheduled_request_ids = scheduler_output.num_scheduled_tokens.keys(
1019
- )
1018
draft_token_ids = [
1019
scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])
1022
- for req_id in scheduled_request_ids
1020
+ for req_id in self.input_batch.req_ids
1021
]
1022
sampler_output = self.rejection_sampler(draft_token_ids,
1023
target_probs,