main
py 359 lines 13.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import sys
4 from abc import ABC, abstractmethod
5
6 import tokenizers.decoders
7 from tokenizers import Tokenizer
8 from transformers import TokenizersBackend
9
10 from vllm.logger import init_logger
11 from vllm.tokenizers import TokenizerLike
12 from vllm.tokenizers.detokenizer_utils import (
13 convert_prompt_ids_to_tokens,
14 detokenize_incrementally,
15 )
16 from vllm.utils import length_from_prompt_token_ids_or_embeds
17 from vllm.v1.engine import EngineCoreRequest
18
19 logger = init_logger(__name__)
20
21 # Error string from https://github.com/huggingface/tokenizers/blob/909fdde2a4ffedd9295206f705eb612be2a91b12/tokenizers/src/tokenizer/mod.rs#L1042
22 INVALID_PREFIX_ERR_MSG = "Invalid prefix encountered"
23
24
25 def uses_fast_detokenizer(tokenizer: TokenizerLike) -> bool:
26 """Whether the engine detokenizes with `FastIncrementalDetokenizer`."""
27 return isinstance(tokenizer, TokenizersBackend)
28
29
30 class IncrementalDetokenizer:
31 def __init__(self):
32 self.token_ids: list[int] = []
33
34 @property
35 def output_token_ids(self) -> list[int]:
36 return self.token_ids
37
38 def num_output_tokens(self) -> int:
39 return len(self.token_ids)
40
41 def update(self, new_token_ids: list[int], stop_terminated: bool) -> str | None:
42 self.token_ids.extend(new_token_ids)
43 return None
44
45 def get_next_output_text(self, finished: bool, delta: bool) -> str:
46 return ""
47
48 @classmethod
49 def from_new_request(
50 cls,
51 tokenizer: TokenizerLike | None,
52 request: EngineCoreRequest,
53 ) -> "IncrementalDetokenizer":
54 assert request.sampling_params is not None
55
56 if tokenizer is None:
57 # No tokenizer => skipping detokenization.
58 return IncrementalDetokenizer()
59
60 if uses_fast_detokenizer(tokenizer):
61 # Fast tokenizer => use tokenizers library DecodeStream.
62 return FastIncrementalDetokenizer(tokenizer, request)
63
64 # Fall back to slow python-based incremental detokenization.
65 return SlowIncrementalDetokenizer(tokenizer, request)
66
67
68 class BaseIncrementalDetokenizer(IncrementalDetokenizer, ABC):
69 def __init__(self, request: EngineCoreRequest):
70 super().__init__()
71
72 # Stop strings
73 params = request.sampling_params
74 assert params is not None
75 if params.stop is None:
76 self.stop = []
77 elif isinstance(params.stop, str):
78 self.stop = [params.stop]
79 else:
80 self.stop = params.stop
81 self.min_tokens = params.min_tokens
82 self.include_stop_str_in_output = params.include_stop_str_in_output
83
84 # Number of chars to hold back when stop strings are to be excluded
85 # from streamed output.
86 if self.stop and not self.include_stop_str_in_output:
87 self.stop_buffer_length = max(len(s) for s in self.stop) - 1
88 else:
89 self.stop_buffer_length = 0
90 self._last_output_text_offset: int = 0
91
92 # Generation data
93 self.output_text = ""
94
95 def update(self, new_token_ids: list[int], stop_terminated: bool) -> str | None:
96 """Update RequestState for the request_id by:
97 1) Detokenize the new token ids incrementally.
98 2) Evaluate stop criteria.
99
100 Return matched stop string or None.
101 """
102 if not new_token_ids:
103 # Skip detokenization if no new token ids.
104 return None
105
106 if stop_terminated and not self.include_stop_str_in_output:
107 # If stop-terminated, exclude last token from detokenization
108 # based on include_stop_str_in_output parameter.
109 skipped_stop_token_id = new_token_ids[-1]
110 new_token_ids = new_token_ids[:-1]
111 else:
112 skipped_stop_token_id = None
113
114 # 1) Detokenize the new token ids incrementally.
115 stop_check_offset = len(self.output_text)
116 for new_token_id in new_token_ids:
117 self.token_ids.append(new_token_id)
118 self.output_text += self.decode_next(new_token_id)
119 # Support min_tokens, see https://github.com/vllm-project/vllm/pull/22014
120 if self.min_tokens and self.num_output_tokens() <= self.min_tokens:
121 stop_check_offset = len(self.output_text)
122
123 if skipped_stop_token_id is not None:
124 # Cleanup after skipping detokenization.
125 self.token_ids.append(skipped_stop_token_id)
126
127 # 2) Evaluate stop strings.
128 stop_string = None
129 if self.stop and self.num_output_tokens() > self.min_tokens:
130 stop = check_stop_strings(
131 output_text=self.output_text,
132 new_char_count=len(self.output_text) - stop_check_offset,
133 stop=self.stop,
134 include_in_output=self.include_stop_str_in_output,
135 )
136 if stop is not None:
137 stop_string, truncate_to = stop
138 if truncate_to != -1:
139 self.output_text = self.output_text[:truncate_to]
140
141 return stop_string
142
143 @abstractmethod
144 def decode_next(self, next_token_id: int) -> str:
145 raise NotImplementedError
146
147 def get_next_output_text(self, finished: bool, delta: bool) -> str:
148 """If delta is True, only new text since the last call to
149 this method is returned"""
150 # We return the full output text if the sequence is finished.
151 buffer_length = 0 if finished else self.stop_buffer_length
152 if not delta:
153 if not buffer_length:
154 return self.output_text
155 return self.output_text[:-buffer_length]
156
157 length = len(self.output_text) - buffer_length
158 last_offset = self._last_output_text_offset
159 if last_offset < length:
160 self._last_output_text_offset = length
161 return self.output_text[last_offset:length]
162 return ""
163
164
165 class FastIncrementalDetokenizer(BaseIncrementalDetokenizer):
166 def __init__(self, tokenizer: TokenizersBackend, request: EngineCoreRequest):
167 super().__init__(request)
168
169 sampling_params = request.sampling_params
170 assert sampling_params is not None
171
172 self.request_id = request.request_id
173 self.skip_special_tokens = sampling_params.skip_special_tokens
174
175 self.tokenizer: Tokenizer = tokenizer._tokenizer
176
177 # Use native prefill to prime the decode stream with prompt tokens.
178 # Look up DecodeStream on the module so backend patches (e.g. the
179 # fastokens shim that replaces ``tokenizers.decoders.DecodeStream``)
180 # are honored regardless of import order.
181 self.stream = tokenizers.decoders.DecodeStream(
182 ids=request.prompt_token_ids,
183 skip_special_tokens=self.skip_special_tokens,
184 )
185
186 self.spaces_between_special_tokens = (
187 sampling_params.skip_special_tokens
188 or sampling_params.spaces_between_special_tokens
189 )
190
191 if not self.spaces_between_special_tokens:
192 # Store dict of added token ids so that we can suppress
193 # the spaces between them.
194 added_token_ids = getattr(self.tokenizer, "added_token_ids", None)
195 if added_token_ids is None:
196 self.tokenizer.added_token_ids = added_token_ids = {
197 tid: tok.content
198 for tid, tok in self.tokenizer.get_added_tokens_decoder().items()
199 }
200
201 if added_token_ids:
202 self.last_special = False
203 self.added_token_ids = added_token_ids
204 else:
205 # No added tokens.
206 self.spaces_between_special_tokens = True
207
208 def decode_next(self, next_token_id: int) -> str:
209 token = self._protected_step(next_token_id)
210
211 if not self.spaces_between_special_tokens:
212 special_token = self.added_token_ids.get(next_token_id)
213 is_special = special_token is not None
214 if is_special and self.last_special:
215 # Return raw token string without any prefixed spaces.
216 token = special_token
217 self.last_special = is_special
218
219 return token or ""
220
221 def _protected_step(self, next_token_id: int) -> str | None:
222 try:
223 token = self.stream.step(self.tokenizer, next_token_id)
224 except (OverflowError, TypeError):
225 # Handle rare observed overflow, still to be diagnosed.
226 # See https://github.com/vllm-project/vllm/issues/21951.
227 logger.exception("Encountered invalid token id: %r", next_token_id)
228 token = None
229 except Exception as e:
230 if not str(e).startswith(INVALID_PREFIX_ERR_MSG):
231 raise e
232 # Recover from edge case where tokenizer can produce non-monotonic,
233 # invalid UTF-8 output, which breaks the internal state of
234 # tokenizers' DecodeStream.
235 # See https://github.com/vllm-project/vllm/issues/17448.
236 logger.warning(
237 "Encountered invalid prefix detokenization error"
238 " for request %s, resetting decode stream.",
239 self.request_id,
240 )
241 self.stream = tokenizers.decoders.DecodeStream(
242 skip_special_tokens=self.skip_special_tokens
243 )
244 token = self.stream.step(self.tokenizer, next_token_id)
245 return token
246
247
248 class SlowIncrementalDetokenizer(BaseIncrementalDetokenizer):
249 def __init__(self, tokenizer: TokenizerLike, request: EngineCoreRequest):
250 super().__init__(request)
251
252 self.tokenizer = tokenizer
253 params = request.sampling_params
254 assert params is not None
255
256 self.prompt_len = length_from_prompt_token_ids_or_embeds(
257 request.prompt_token_ids, request.prompt_embeds
258 )
259
260 # Metadata for incremental detokenization.
261 if request.prompt_token_ids is not None:
262 self.tokens, self.prefix_offset, self.read_offset = (
263 convert_prompt_ids_to_tokens(
264 tokenizer=tokenizer,
265 prompt_ids=request.prompt_token_ids,
266 skip_special_tokens=params.skip_special_tokens,
267 )
268 )
269 else:
270 # Prompt embedding requests cannot be detokenized, in general.
271 self.tokens = [""] * self.prompt_len
272 self.prefix_offset = 0
273 self.read_offset = 0
274
275 self.token_ids.extend(request.prompt_token_ids or [0] * self.prompt_len)
276
277 self.skip_special_tokens = params.skip_special_tokens
278 self.spaces_between_special_tokens = params.spaces_between_special_tokens
279
280 @property
281 def output_token_ids(self) -> list[int]:
282 if self.prompt_len:
283 return self.token_ids[self.prompt_len :]
284 return self.token_ids
285
286 def num_output_tokens(self) -> int:
287 return len(self.token_ids) - self.prompt_len
288
289 def decode_next(self, next_token_id: int) -> str:
290 new_tokens, decoded_text, prefix_offset, read_offset = detokenize_incrementally(
291 tokenizer=self.tokenizer,
292 all_input_ids=self.token_ids,
293 prev_tokens=self.tokens,
294 prefix_offset=self.prefix_offset,
295 read_offset=self.read_offset,
296 skip_special_tokens=self.skip_special_tokens,
297 spaces_between_special_tokens=self.spaces_between_special_tokens,
298 )
299
300 self.tokens.extend(new_tokens)
301 self.prefix_offset = prefix_offset
302 self.read_offset = read_offset
303
304 return decoded_text
305
306
307 def check_stop_strings(
308 output_text: str,
309 new_char_count: int,
310 stop: list[str],
311 include_in_output: bool,
312 ) -> tuple[str, int] | None:
313 """Check if any stop strings are matched and truncate sequence
314 output text accordingly.
315
316 Returns tuple (stop_string, offset) if matched or else None.
317
318 Where stop_string is the matched stop string and offset is the
319 length to which output_text should be truncated, or -1 for no
320 truncation.
321
322 When several stop strings match within the newly generated text (for
323 example when speculative decoding appends multiple tokens in a single
324 step), the stop string that completes earliest in the text is selected,
325 so the result matches appending one token at a time. Ties are broken by
326 stop-list order.
327 """
328 if not new_char_count or not stop:
329 return None
330
331 best_stop_str: str | None = None
332 best_stop_index = 0
333 best_end = sys.maxsize
334 for stop_str in stop:
335 stop_string_len = len(stop_str)
336 # Avoid searching already-searched text.
337 stop_index = output_text.find(stop_str, 1 - new_char_count - stop_string_len)
338 if stop_index == -1:
339 continue
340
341 # Prefer the stop string that completes earliest in the text.
342 end = stop_index + stop_string_len
343 if end < best_end:
344 best_stop_str = stop_str
345 best_stop_index = stop_index
346 best_end = end
347
348 if best_stop_str is None:
349 return None
350
351 if include_in_output:
352 # Truncate to end of stop string.
353 if best_end >= len(output_text):
354 # No truncation required.
355 return best_stop_str, -1
356 return best_stop_str, best_end
357
358 # Truncate the output text to the beginning of the stop string.
359 return best_stop_str, best_stop_index