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, 21:00 UTC
[V1] Do not detokenize if sampling param detokenize is False (#14224)
Signed-off-by: Himanshu Jaju <hj@mistral.ai> Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com>
Himanshu Jaju committed
Mar 6, 2025 at 19:40 UTC
cd579352bfef4828bb9a1e40221ff757495ccbdd
4 files changed
+69
-27
tests/v1/sample/test_sampling_params_e2e.py
+28
-1
@@ -14,7 +14,10 @@ PROMPT = "Hello my name is Robert and I"
14
15
@pytest.fixture(scope="module")
16
def model() -> LLM:
17
- return LLM(MODEL, enforce_eager=True)
17
+ # Disable prefix caching so that we can test prompt logprobs.
18
+ # TODO remove this after https://github.com/vllm-project/vllm/pull/13949
19
+ # is merged
20
+ return LLM(MODEL, enforce_eager=True, enable_prefix_caching=False)
21
22
23
def test_n_gt_1(model):
@@ -87,9 +90,33 @@ def test_stop_token_ids(model):
90
91
stop_token_ids = [stop_token_id_0, stop_token_id_1]
92
params = SamplingParams(temperature=0, stop_token_ids=stop_token_ids)
93
+ output = model.generate(PROMPT, params)
94
assert output[0].outputs[0].token_ids[-1] == stop_token_id_0
95
96
97
+def test_detokenize_false(model):
98
+ """Check that detokenize=False option works."""
99
+
100
+ output = model.generate(PROMPT, SamplingParams(detokenize=False))
101
+ assert len(output[0].outputs[0].token_ids) > 0
102
+ assert len(output[0].outputs[0].text) == 0
103
+
104
+ output = model.generate(
105
+ PROMPT, SamplingParams(detokenize=False, logprobs=3,
106
+ prompt_logprobs=3))
107
+ assert len(output[0].outputs[0].token_ids) > 0
108
+ assert len(output[0].outputs[0].text) == 0
109
+
110
+ prompt_logprobs = output[0].prompt_logprobs
111
+ sampled_logprobs = output[0].outputs[0].logprobs
112
+ assert len(prompt_logprobs) > 1
113
+ assert len(sampled_logprobs) > 1
114
+ for all_logprobs in (prompt_logprobs[1:], sampled_logprobs):
115
+ for logprobs in all_logprobs:
116
+ assert 3 <= len(logprobs) <= 4
117
+ assert all(lp.decoded_token is None for lp in logprobs.values())
118
+
119
+
120
def test_bad_words(model):
121
"""Check that we respect bad words."""
122
vllm/v1/engine/detokenizer.py
+24
-16
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-from dataclasses import dataclass
3
+from dataclasses import dataclass, field
4
from typing import Optional
5
6
from vllm.engine.output_processor.stop_checker import StopChecker
@@ -16,41 +16,46 @@ logger = init_logger(__name__)
16
class IncrementalDetokenizer:
17
18
# Generation data
19
- output_text: str
20
- tokens: list[str]
19
token_ids: list[int]
22
- prompt_len: int
20
+ output_text: str = ""
21
+ tokens: list[str] = field(default_factory=list)
22
+ prompt_len: int = 0
23
24
# Stop strings
25
- stop: list[str]
26
- include_stop_str_in_output: bool
25
+ stop: list[str] = field(default_factory=list)
26
+ include_stop_str_in_output: bool = False
27
28
# Metadata for incremental detokenization
29
- prefix_offset: int
30
- read_offset: int
29
+ prefix_offset: int = 0
30
+ read_offset: int = 0
31
32
# Parameters for detokenization
33
- skip_special_tokens: bool
34
- spaces_between_special_tokens: bool
33
+ skip_special_tokens: bool = True
34
+ spaces_between_special_tokens: bool = True
35
36
- # Tokenizer for this request
37
- tokenizer: AnyTokenizer
36
+ # Tokenizer for this request,
37
+ # None if detokenization is disabled.
38
+ tokenizer: Optional[AnyTokenizer] = None
39
40
# Accounting for stop string buffering
40
- stop_buffer_length: int
41
+ stop_buffer_length: int = 0
42
_last_output_text_offset: int = 0
43
44
@property
45
def output_token_ids(self) -> list[int]:
45
- return self.token_ids[self.prompt_len:]
46
+ return self.token_ids if not self.prompt_len else (
47
+ self.token_ids[self.prompt_len:])
48
49
@classmethod
50
def from_new_request(
51
cls,
50
- tokenizer: AnyTokenizer,
52
+ tokenizer: Optional[AnyTokenizer],
53
request: EngineCoreRequest,
54
) -> "IncrementalDetokenizer":
55
56
+ if tokenizer is None:
57
+ return cls(token_ids=[])
58
+
59
tokens, prefix_offset, read_offset = convert_prompt_ids_to_tokens(
60
tokenizer=tokenizer,
61
prompt_ids=request.prompt_token_ids,
@@ -66,7 +71,6 @@ class IncrementalDetokenizer:
71
stop_buffer_length = 0
72
73
return cls(
69
- output_text="",
74
tokens=tokens,
75
# Detokenizer mutates this list, so need a unique copy.
76
# NOTE(Nick): could we take ownership of it though?
@@ -93,6 +97,10 @@ class IncrementalDetokenizer:
97
Return matched stop string or None.
98
"""
99
100
+ if self.tokenizer is None:
101
+ self.token_ids.extend(new_token_ids)
102
+ return None
103
+
104
# 1) Detokenize the new token ids incrementally.
105
# TODO(woosuk): This method becomes very inefficient when the number of
106
# new_token_ids is more than 1. We need to optimize this.
vllm/v1/engine/logprobs.py
+15
-10
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
import itertools
4
+from collections.abc import Iterable
5
from dataclasses import dataclass
6
from typing import Optional
7
@@ -13,12 +14,15 @@ from vllm.v1.outputs import LogprobsLists, LogprobsTensors
14
15
logger = init_logger(__name__)
16
17
+NONES = itertools.repeat(None)
18
+
19
20
@dataclass
21
class LogprobsProcessor:
22
20
- # Tokenizer for this request
21
- tokenizer: AnyTokenizer
23
+ # Tokenizer for this request,
24
+ # None if detokenization is disabled.
25
+ tokenizer: Optional[AnyTokenizer]
26
27
# Logprobs for this request
28
logprobs: Optional[SampleLogprobs]
@@ -30,7 +34,7 @@ class LogprobsProcessor:
34
@classmethod
35
def from_new_request(
36
cls,
33
- tokenizer: AnyTokenizer,
37
+ tokenizer: Optional[AnyTokenizer],
38
request: EngineCoreRequest,
39
) -> "LogprobsProcessor":
40
num_logprobs = request.sampling_params.logprobs
@@ -66,8 +70,8 @@ class LogprobsProcessor:
70
token_ids_lst):
71
72
# Detokenize (non-incrementally).
69
- decoded_tokens = convert_ids_list_to_tokens(
70
- self.tokenizer, token_ids)
73
+ decoded_tokens = NONES if self.tokenizer is None else (
74
+ convert_ids_list_to_tokens(self.tokenizer, token_ids))
75
76
# Sampler puts the sampled logprob in first.
77
sampled_token_logprob = logprobs[0]
@@ -103,9 +107,9 @@ class LogprobsProcessor:
107
108
# Detokenize non-incrementally.
109
# Output is flat: [num_tok, num_lps] -> [num_tok * num_lps]
106
- decoded_tokens = convert_ids_list_to_tokens(
107
- self.tokenizer,
108
- token_ids.flatten().tolist())
110
+ decoded_tokens = None if self.tokenizer is None else (
111
+ convert_ids_list_to_tokens(self.tokenizer,
112
+ token_ids.flatten().tolist()))
113
114
# Recover shapes.
115
num_prompt_tokens, num_logprobs = logprobs.shape
@@ -121,7 +125,8 @@ class LogprobsProcessor:
125
# Handle flattening.
126
offset = pos * num_logprobs
127
offset_end = offset + num_logprobs
124
- decoded_tokens_for_pos = decoded_tokens[offset:offset_end]
128
+ decoded_tokens_for_pos = NONES \
129
+ if decoded_tokens is None else decoded_tokens[offset:offset_end]
130
131
# Update with the Logprob dictionary for this pos.
132
self.prompt_logprobs.append(
@@ -153,7 +158,7 @@ class LogprobsProcessor:
158
def _make_logprob_dict(
159
logprobs: list[float],
160
logprob_token_ids: list[int],
156
- decoded_tokens: list[str],
161
+ decoded_tokens: Iterable[Optional[str]],
162
rank: int,
163
num_logprobs: int,
164
) -> dict[int, Logprob]:
vllm/v1/engine/output_processor.py
+2
@@ -68,6 +68,8 @@ class RequestState:
68
queue: Optional[asyncio.Queue[RequestOutput]],
69
log_stats: bool,
70
) -> "RequestState":
71
+ if not request.sampling_params.detokenize:
72
+ tokenizer = None
73
return cls(
74
request_id=request.request_id,
75
parent_req=parent_req,