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 14, 05:30 UTC
[Refactor] Pass `tokenizer` explicitly instead of binding to prompt update (#23542)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Aug 25, 2025 at 21:31 UTC
6879cd80ae4dba88121db226d5bfbc6a75b072ba
4 files changed
+95
-144
tests/multimodal/test_processing.py
+8
-13
@@ -243,7 +243,7 @@ def test_find_token_matches(
243
mock_tokenizer = cast(AnyTokenizer, object())
244
245
prompt_updates = {
246
- key: update_type(key, target, []).resolve(mock_tokenizer, 0)
246
+ key: update_type(key, target, []).resolve(0)
247
for key, target in target_by_key.items()
248
}
249
result = {
@@ -392,7 +392,7 @@ def test_find_text_matches(
392
mock_tokenizer = cast(AnyTokenizer, object())
393
394
prompt_updates = {
395
- key: update_type(key, target, []).resolve(mock_tokenizer, 0)
395
+ key: update_type(key, target, []).resolve(0)
396
for key, target in target_by_key.items()
397
}
398
result = {
@@ -559,10 +559,8 @@ def test_find_update_text(
559
) in expected_by_update_type_mm_count.items():
560
for mm_count, expected in expected_by_mm_count.items():
561
mm_prompt_updates = {
562
- key: [[
563
- update_type(key, target,
564
- repl_by_key[key]).resolve(mock_tokenizer, i)
565
- ] for i in range(mm_count)]
562
+ key: [[update_type(key, target, repl_by_key[key]).resolve(i)]
563
+ for i in range(mm_count)]
564
for key, target in target_by_key.items()
565
}
566
@@ -731,10 +729,8 @@ def test_find_update_tokens(
729
) in expected_by_update_type_mm_count.items():
730
for mm_count, expected in expected_by_mm_count.items():
731
mm_prompt_updates = {
734
- key: [[
735
- update_type(key, target,
736
- repl_by_key[key]).resolve(mock_tokenizer, i)
737
- ] for i in range(mm_count)]
732
+ key: [[update_type(key, target, repl_by_key[key]).resolve(i)]
733
+ for i in range(mm_count)]
734
for key, target in target_by_key.items()
735
}
736
@@ -879,12 +875,11 @@ def test_find_mm_placeholders(
875
mock_tokenizer = cast(AnyTokenizer, object())
876
877
mm_prompt_updates = {
882
- key: [[update_type(key, [], repl).resolve(mock_tokenizer, i)]
883
- for i in range(3)]
878
+ key: [[update_type(key, [], repl).resolve(i)] for i in range(3)]
879
for key, repl in repl_by_key.items()
880
}
881
887
- result = find_mm_placeholders(prompt, mm_prompt_updates)
882
+ result = find_mm_placeholders(prompt, mm_prompt_updates, mock_tokenizer)
883
884
# Only displayed on error
885
print("result:", result)
vllm/model_executor/models/gemma3_mm.py
+2
-2
@@ -28,7 +28,6 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
28
PlaceholderFeaturesInfo,
29
PromptReplacement, PromptUpdate,
30
PromptUpdateDetails,
31
- find_mm_placeholders,
31
replace_token_matches)
32
# yapf: enable
33
from vllm.multimodal.profiling import BaseDummyInputsBuilder
@@ -401,7 +400,8 @@ class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):
400
repl_token_ids.extend(repl_toks)
401
repl_orig_idxs.extend(orig_idx for _ in range(len(repl_toks)))
402
404
- repls = find_mm_placeholders(repl_token_ids, mm_prompt_updates)
403
+ repls = super()._find_mm_placeholders(repl_token_ids,
404
+ mm_prompt_updates)
405
406
return {
407
modality: [
vllm/model_executor/models/gemma3n_mm.py
+2
-2
@@ -35,7 +35,6 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
35
PlaceholderFeaturesInfo,
36
PromptReplacement, PromptUpdate,
37
PromptUpdateDetails,
38
- find_mm_placeholders,
38
replace_token_matches)
39
# yapf: enable
40
from vllm.multimodal.profiling import BaseDummyInputsBuilder
@@ -318,7 +317,8 @@ class Gemma3nMultiModalProcessor(BaseMultiModalProcessor[Gemma3nProcessingInfo]
317
repl_token_ids.extend(repl_toks)
318
repl_orig_idxs.extend(orig_idx for _ in range(len(repl_toks)))
319
321
- repls = find_mm_placeholders(repl_token_ids, mm_prompt_updates)
320
+ repls = super()._find_mm_placeholders(repl_token_ids,
321
+ mm_prompt_updates)
322
323
return {
324
modality: [
vllm/multimodal/processing.py
+83
-127
@@ -44,6 +44,44 @@ PromptSeq = Union[str, list[int]]
44
"""A token sequence (list of token IDs) or text."""
45
46
47
+@lru_cache(maxsize=2048)
48
+def _cached_encode(
49
+ tokenizer: AnyTokenizer,
50
+ text: str,
51
+ *,
52
+ add_special_tokens: Optional[bool] = None,
53
+) -> list[int]:
54
+ return encode_tokens(tokenizer,
55
+ text,
56
+ add_special_tokens=add_special_tokens)
57
+
58
+
59
+@lru_cache(maxsize=2048)
60
+def _cached_decode(
61
+ tokenizer: AnyTokenizer,
62
+ token_ids: tuple[int, ...],
63
+ *,
64
+ skip_special_tokens: Optional[bool] = None,
65
+) -> str:
66
+ return decode_tokens(tokenizer,
67
+ list(token_ids),
68
+ skip_special_tokens=skip_special_tokens)
69
+
70
+
71
+def _seq2text(tokenizer: AnyTokenizer, seq: PromptSeq) -> str:
72
+ if isinstance(seq, str):
73
+ return seq
74
+
75
+ return _cached_decode(tokenizer, tuple(seq))
76
+
77
+
78
+def _seq2tokens(tokenizer: AnyTokenizer, seq: PromptSeq) -> list[int]:
79
+ if isinstance(seq, str):
80
+ return _cached_encode(tokenizer, seq, add_special_tokens=False)
81
+
82
+ return seq
83
+
84
+
85
class _GetMatchIndex(Protocol):
86
87
def __call__(
@@ -137,7 +175,8 @@ class PromptUpdateDetails(Generic[_S]):
175
full: _S
176
"""The full content."""
177
140
- is_embed: Optional[Callable[["_BoundPromptSequence"], torch.Tensor]] = None
178
+ is_embed: Optional[Callable[[AnyTokenizer, PromptSeq],
179
+ torch.Tensor]] = None
180
"""
181
Given [`full`][vllm.multimodal.processing.PromptUpdateDetails.full],
182
return a boolean mask of shape `(len(full),)` indicating which positions
@@ -159,11 +198,12 @@ class PromptUpdateDetails(Generic[_S]):
198
embed_text: str,
199
) -> "PromptUpdateDetails[_S]":
200
162
- def is_embed(full: "_BoundPromptSequence") -> torch.Tensor:
163
- embed_token_ids = encode_tokens(full.tokenizer, embed_text)
201
+ def is_embed(tokenizer: AnyTokenizer, full: PromptSeq) -> torch.Tensor:
202
+ embed_token_ids = encode_tokens(tokenizer, embed_text)
203
+ token_ids = _seq2tokens(tokenizer, full)
204
205
return torch.isin(
166
- torch.tensor(full.token_ids),
206
+ torch.tensor(token_ids),
207
torch.tensor(embed_token_ids),
208
)
209
@@ -174,10 +214,13 @@ class PromptUpdateDetails(Generic[_S]):
214
seq: _S,
215
embed_token_id: int,
216
) -> "PromptUpdateDetails[_S]":
177
- return PromptUpdateDetails(
178
- full=seq,
179
- is_embed=lambda f: torch.tensor(f.token_ids) == embed_token_id,
180
- )
217
+
218
+ def is_embed(tokenizer: AnyTokenizer, full: PromptSeq) -> torch.Tensor:
219
+ token_ids = _seq2tokens(tokenizer, full)
220
+
221
+ return torch.tensor(token_ids) == embed_token_id
222
+
223
+ return PromptUpdateDetails(full=seq, is_embed=is_embed)
224
225
226
PromptUpdateInfo = Union[PromptSeq, PromptUpdateDetails]
@@ -230,25 +273,14 @@ class PromptUpdate(ABC):
273
"""Defines how to update the prompt."""
274
raise NotImplementedError
275
233
- def _resolve_target(
234
- self,
235
- tokenizer: AnyTokenizer,
236
- item_idx: int,
237
- ) -> Union["_BoundPromptSequence", PromptIndex]:
276
+ def _resolve_target(self, item_idx: int) -> UpdateTarget:
277
target = self.target
278
if callable(target):
279
target = target(item_idx)
280
242
- if isinstance(target, PromptIndex):
243
- return target
244
-
245
- return _BoundPromptSequence.from_seq(tokenizer, target)
281
+ return target
282
247
- def _resolve_content(
248
- self,
249
- tokenizer: AnyTokenizer,
250
- item_idx: int,
251
- ) -> "_BoundPromptContent":
283
+ def _resolve_content(self, item_idx: int) -> PromptUpdateDetails:
284
content = self.content
285
if callable(content):
286
content = content(item_idx)
@@ -256,17 +288,9 @@ class PromptUpdate(ABC):
288
if not isinstance(content, PromptUpdateDetails):
289
content = PromptUpdateDetails.from_seq(content)
290
259
- bound_full = _BoundPromptSequence.from_seq(tokenizer, content.full)
260
- bound_content = _BoundPromptContent(full=bound_full,
261
- is_embed=content.is_embed)
262
-
263
- return bound_content
291
+ return content
292
265
- def resolve(
266
- self,
267
- tokenizer: AnyTokenizer,
268
- item_idx: int,
269
- ) -> "ResolvedPromptUpdate":
293
+ def resolve(self, item_idx: int) -> "ResolvedPromptUpdate":
294
"""
295
Given the index of the processed item within
296
[`modality`][vllm.multimodal.processing.PromptUpdate.modality],
@@ -276,8 +300,8 @@ class PromptUpdate(ABC):
300
modality=self.modality,
301
item_idx=item_idx,
302
mode=self.mode,
279
- target=self._resolve_target(tokenizer, item_idx),
280
- content=self._resolve_content(tokenizer, item_idx),
303
+ target=self._resolve_target(item_idx),
304
+ content=self._resolve_content(item_idx),
305
)
306
307
@@ -424,30 +448,6 @@ class PromptReplacement(PromptUpdate):
448
return UpdateMode.REPLACE
449
450
427
-@lru_cache(maxsize=2048)
428
-def _cached_encode(
429
- tokenizer: AnyTokenizer,
430
- text: str,
431
- *,
432
- add_special_tokens: Optional[bool] = None,
433
-) -> list[int]:
434
- return encode_tokens(tokenizer,
435
- text,
436
- add_special_tokens=add_special_tokens)
437
-
438
-
439
-@lru_cache(maxsize=2048)
440
-def _cached_decode(
441
- tokenizer: AnyTokenizer,
442
- token_ids: tuple[int, ...],
443
- *,
444
- skip_special_tokens: Optional[bool] = None,
445
-) -> str:
446
- return decode_tokens(tokenizer,
447
- list(token_ids),
448
- skip_special_tokens=skip_special_tokens)
449
-
450
-
451
class _HasModalityAttr(Protocol):
452
modality: str
453
@@ -468,59 +468,6 @@ def full_groupby_modality(values: Iterable[_M]) -> ItemsView[str, list[_M]]:
468
return full_groupby(values, key=lambda x: x.modality)
469
470
471
-@dataclass
472
-class _BoundPromptSequence:
473
- """
474
- A [`_PromptSeq`][vllm.multimodal.processing.PromptSeq] bound
475
- to a tokenizer to automatically
476
- convert between token sequence and text representations.
477
- """
478
- tokenizer: AnyTokenizer = field(repr=False)
479
-
480
- _text: Optional[str]
481
- _token_ids: Optional[list[int]]
482
-
483
- @staticmethod
484
- def from_seq(
485
- tokenizer: AnyTokenizer,
486
- seq: PromptSeq,
487
- ) -> "_BoundPromptSequence":
488
- return _BoundPromptSequence(
489
- tokenizer=tokenizer,
490
- _text=seq if isinstance(seq, str) else None,
491
- _token_ids=seq if isinstance(seq, list) else None,
492
- )
493
-
494
- def __post_init__(self) -> None:
495
- if self._text is None and self._token_ids is None:
496
- raise ValueError("At least one of 'text' and 'token_ids' must be "
497
- "specified")
498
-
499
- @property
500
- def text(self) -> str:
501
- if self._text is None:
502
- assert self._token_ids is not None
503
- self._text = _cached_decode(self.tokenizer, tuple(self._token_ids))
504
-
505
- return self._text
506
-
507
- @property
508
- def token_ids(self) -> list[int]:
509
- if self._token_ids is None:
510
- assert self._text is not None
511
- self._token_ids = _cached_encode(self.tokenizer,
512
- self._text,
513
- add_special_tokens=False)
514
-
515
- return self._token_ids
516
-
517
-
518
-@dataclass
519
-class _BoundPromptContent:
520
- full: _BoundPromptSequence
521
- is_embed: Optional[Callable[["_BoundPromptSequence"], torch.Tensor]]
522
-
523
-
471
class PromptTargetMatch(NamedTuple):
472
start_idx: int
473
end_idx: int
@@ -542,10 +489,10 @@ class ResolvedPromptUpdate:
489
mode: UpdateMode
490
"""Defines how to update the prompt."""
491
545
- target: Union[_BoundPromptSequence, PromptIndex]
492
+ target: UpdateTarget
493
"""The token sequence (or text) to update."""
494
548
- content: _BoundPromptContent = field(repr=False)
495
+ content: PromptUpdateDetails = field(repr=False)
496
"""The placeholder tokens that are part of the update."""
497
498
def iter_token_matches(
@@ -565,8 +512,10 @@ class ResolvedPromptUpdate:
512
513
return
514
515
+ target_token_ids = _seq2tokens(tokenizer, target)
516
+
517
for match in iter_token_matches(prompt,
569
- target.token_ids,
518
+ target_token_ids,
519
start_idx=start_idx):
520
yield PromptTargetMatch(match.start_idx, match.end_idx)
521
@@ -587,7 +536,9 @@ class ResolvedPromptUpdate:
536
537
return
538
590
- for match in re.finditer(re.escape(target.text), prompt,
539
+ target_text = _seq2text(tokenizer, target)
540
+
541
+ for match in re.finditer(re.escape(target_text), prompt,
542
pos=start_idx):
543
yield PromptTargetMatch(match.start(), match.end())
544
@@ -779,7 +730,7 @@ def _apply_matches(
730
731
matched_update = mm_prompt_updates[modality][item_idx][
732
update_idx]
782
- matched_content = matched_update.content
733
+ matched_content = matched_update.content.full
734
735
if mode == UpdateMode.INSERT:
736
end_idx_to_insert = match.end_idx
@@ -789,8 +740,10 @@ def _apply_matches(
740
assert_never(mode)
741
742
out_seqs.append(prompt[prev_end_idx:end_idx_to_insert])
792
- out_seqs.append(matched_content.full.text if isinstance(
793
- prompt, str) else matched_content.full.token_ids)
743
+ out_seqs.append(
744
+ _seq2text(tokenizer, matched_content
745
+ ) if isinstance(prompt, str) else _seq2tokens(
746
+ tokenizer, matched_content))
747
out_result[modality][item_idx] = update_idx
748
749
# Exclude overlapping matches
@@ -842,6 +795,7 @@ def apply_text_matches(
795
def _iter_placeholders(
796
prompt: list[int],
797
mm_prompt_updates: "MultiModalPromptUpdates",
798
+ tokenizer: AnyTokenizer,
799
) -> Iterable[PlaceholderFeaturesInfo]:
800
"""
801
Yield each set of placeholder tokens found in `prompt`.
@@ -868,7 +822,7 @@ def _iter_placeholders(
822
823
for update in modality_updates[item_idx]:
824
content = update.content
871
- content_tokens_full = content.full.token_ids
825
+ content_tokens_full = _seq2tokens(tokenizer, content.full)
826
content_len_full = len(content_tokens_full)
827
end_idx_full = start_idx + content_len_full
828
@@ -878,7 +832,8 @@ def _iter_placeholders(
832
if prompt[start_idx:end_idx_full] == content_tokens_full:
833
content_is_embed = content.is_embed
834
if content_is_embed is not None:
881
- content_is_embed = content_is_embed(content.full)
835
+ content_is_embed = content_is_embed(
836
+ tokenizer, content.full)
837
838
yield PlaceholderFeaturesInfo(
839
modality=modality,
@@ -904,8 +859,9 @@ def _iter_placeholders(
859
def find_mm_placeholders(
860
prompt: list[int],
861
mm_prompt_updates: "MultiModalPromptUpdates",
862
+ tokenizer: AnyTokenizer,
863
) -> Mapping[str, list[PlaceholderFeaturesInfo]]:
908
- it = _iter_placeholders(prompt, mm_prompt_updates)
864
+ it = _iter_placeholders(prompt, mm_prompt_updates, tokenizer)
865
return dict(full_groupby_modality(it))
866
867
@@ -1160,12 +1116,9 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
1116
prompt_updates: Sequence[PromptUpdate],
1117
mm_item_counts: Mapping[str, int],
1118
) -> MultiModalPromptUpdates:
1163
- tokenizer = self.info.get_tokenizer()
1164
-
1119
return {
1166
- modality:
1167
- [[update.resolve(tokenizer, item_idx) for update in updates]
1168
- for item_idx in range(mm_item_counts.get(modality, 0))]
1120
+ modality: [[update.resolve(item_idx) for update in updates]
1121
+ for item_idx in range(mm_item_counts.get(modality, 0))]
1122
for modality, updates in full_groupby_modality(prompt_updates)
1123
}
1124
@@ -1208,7 +1161,10 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
1161
new_token_ids: list[int],
1162
mm_prompt_updates: MultiModalPromptUpdates,
1163
) -> Mapping[str, list[PlaceholderFeaturesInfo]]:
1211
- return find_mm_placeholders(new_token_ids, mm_prompt_updates)
1164
+ tokenizer = self.info.get_tokenizer()
1165
+
1166
+ return find_mm_placeholders(new_token_ids, mm_prompt_updates,
1167
+ tokenizer)
1168
1169
def _get_hf_mm_data(
1170
self,