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 6, 09:30 UTC
[Model] Refactor Step3-VL processor to HF style (#37579)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Mar 20, 2026 at 14:05 UTC
30108fc8b063e741bb3d590a1d2b608dc650e67c
4 files changed
+228
-160
vllm/model_executor/models/step3_vl.py
+27
-29
@@ -39,7 +39,11 @@ from vllm.multimodal.processing import (
39
)
40
from vllm.sequence import IntermediateTensors
41
from vllm.transformers_utils.configs.step3_vl import Step3VisionEncoderConfig
42
-from vllm.transformers_utils.processors.step3_vl import Step3VLProcessor
42
+from vllm.transformers_utils.processors.step3_vl import (
43
+ MAX_IMAGE_SIZE,
44
+ Step3VLImageProcessor,
45
+ Step3VLProcessor,
46
+)
47
from vllm.utils.tensor_schema import TensorSchema, TensorShape
48
49
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
@@ -86,21 +90,30 @@ Step3VLImageInputs: TypeAlias = Step3VLImagePixelInputs | Step3VLImageEmbeddingI
90
91
92
class Step3VLProcessingInfo(BaseProcessingInfo):
93
+ def get_image_processor(self, **kwargs):
94
+ config = self.get_hf_config()
95
+
96
+ kwargs.setdefault(
97
+ "enable_patch",
98
+ getattr(config.vision_config, "enable_patch", True),
99
+ )
100
+
101
+ return Step3VLImageProcessor(**kwargs)
102
+
103
def get_hf_processor(self) -> Step3VLProcessor:
104
return Step3VLProcessor(
91
- self.get_hf_config(),
92
- self.get_tokenizer(),
105
+ tokenizer=self.get_tokenizer(),
106
+ image_processor=self.get_image_processor(),
107
)
108
109
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
110
return {"image": None}
111
112
def get_max_image_tokens(self) -> int:
99
- hf_processor = self.get_hf_processor()
100
- return hf_processor.get_num_image_tokens(
101
- self.get_image_size_with_most_features().width,
102
- self.get_image_size_with_most_features().height,
103
- )
113
+ image_processor = self.get_image_processor()
114
+ target_width, target_height = self.get_image_size_with_most_features()
115
+
116
+ return image_processor.get_num_image_tokens(target_width, target_height)
117
118
def get_mm_max_tokens_per_item(
119
self,
@@ -110,20 +123,7 @@ class Step3VLProcessingInfo(BaseProcessingInfo):
123
return {"image": self.get_max_image_tokens()}
124
125
def get_image_size_with_most_features(self) -> ImageSize:
113
- return ImageSize(3024, 3024)
114
-
115
- def get_num_mm_tokens(self, mm_data: MultiModalDataDict) -> int:
116
- if len(mm_data) != 1 or "image" not in mm_data:
117
- raise ValueError("mm_data could only contain one key 'image' for steo1o")
118
-
119
- image_data = mm_data["image"]
120
- if not isinstance(image_data, (list, tuple)):
121
- image_data = [image_data]
122
-
123
- return sum(
124
- self.get_hf_processor().get_num_image_tokens(img.width, img.height)
125
- for img in image_data
126
- )
126
+ return ImageSize(MAX_IMAGE_SIZE, MAX_IMAGE_SIZE)
127
128
129
class Step3VLDummyInputsBuilder(BaseDummyInputsBuilder[Step3VLProcessingInfo]):
@@ -165,13 +165,11 @@ class Step3VLMultiModalProcessor(BaseMultiModalProcessor[Step3VLProcessingInfo])
165
def get_replacement_step1o(item_idx: int):
166
out_item = out_mm_kwargs["image"][item_idx]
167
num_patches = int(out_item["num_patches"].data)
168
- if num_patches > 0:
169
- patch_newline_mask = out_item["patch_newline_mask"].data
170
- image_repl_ids = hf_processor._get_image_repl_features(
171
- 1, num_patches, patch_newline_mask.tolist()
172
- )[1]
173
- else:
174
- image_repl_ids = hf_processor._get_image_repl_features(1, 0, None)[1]
168
+ patch_newline_mask = out_item["patch_newline_mask"].data
169
+ image_repl_ids = hf_processor.get_image_repl_feature_ids(
170
+ 1, num_patches, patch_newline_mask.tolist()
171
+ )
172
+
173
return PromptUpdateDetails.select_token_id(
174
seq=image_repl_ids,
175
embed_token_id=image_placeholder_token_id,
vllm/transformers_utils/processors/internvl.py
+4
-3
@@ -558,6 +558,7 @@ class InternVLProcessor(ProcessorMixin):
558
else:
559
text_inputs = {}
560
561
- combined_outputs = {**text_inputs, **image_inputs, **video_inputs}
562
-
563
- return BatchFeature(combined_outputs, tensor_type=return_tensors)
561
+ return BatchFeature(
562
+ data={**text_inputs, **image_inputs, **video_inputs},
563
+ tensor_type=return_tensors,
564
+ )
vllm/transformers_utils/processors/kimi_k25.py
-1
@@ -19,7 +19,6 @@ class KimiK25Processor(ProcessorMixin):
19
self.media_token_id = media_token_id
20
assert self.media_token_id is not None
21
22
- # We do not support str input for text here
22
def __call__(
23
self,
24
vision_chunks: list[VisionChunk] | None = None,
vllm/transformers_utils/processors/step3_vl.py
+197
-127
@@ -8,13 +8,13 @@ import torch
8
from PIL import Image
9
from torchvision import transforms
10
from torchvision.transforms.functional import InterpolationMode
11
-from transformers import BatchFeature, PretrainedConfig, TensorType
11
+from transformers import BatchFeature, ProcessorMixin, TensorType
12
13
from vllm.tokenizers import TokenizerLike
14
15
MAX_IMAGE_SIZE: int = 3024
16
17
-ImageWithPatches = tuple[Image.Image, list[Image.Image], list[bool] | None]
17
+ImageWithPatches = tuple[Image.Image, list[Image.Image], list[bool]]
18
19
20
class Step3VisionProcessor:
@@ -185,7 +185,7 @@ class ImagePatcher:
185
186
def __call__(
187
self, img: Image.Image
188
- ) -> tuple[Image.Image, list[Image.Image], list[bool] | None]:
188
+ ) -> tuple[Image.Image, list[Image.Image], list[bool]]:
189
img_width, img_height = img.size
190
new_img_width, new_img_height = self.get_image_size_for_padding(
191
img_width, img_height
@@ -203,7 +203,7 @@ class ImagePatcher:
203
)
204
205
if window_size == 0 or not self.enable_patch:
206
- return img, [], None
206
+ return img, [], []
207
else:
208
new_img_width, new_img_height = self.get_image_size_for_crop(
209
new_img_width, new_img_height, window_size
@@ -236,43 +236,28 @@ class ImagePatcher:
236
return (
237
img,
238
patches,
239
- [i in newlines for i in range(len(patches))]
240
- if len(patches) > 0
241
- else None,
239
+ [i in newlines for i in range(len(patches))],
240
)
241
242
245
-class Step3VLProcessor:
243
+class Step3VLImageProcessor:
244
def __init__(
245
self,
248
- config: PretrainedConfig,
249
- tokenizer: TokenizerLike,
246
+ image_size: int = 728,
247
+ patch_size: int = 504,
248
+ num_image_feature_size: int = 169,
249
+ num_patch_feature_size: int = 81,
250
+ enable_patch: bool = True,
251
) -> None:
251
- super().__init__()
252
-
253
- self.config = config
254
- self.tokenizer = tokenizer
255
- self.image_size = 728
256
- self.patch_size = 504
252
+ self.image_size = image_size
253
+ self.patch_size = patch_size
254
+ self.num_image_feature_size = num_image_feature_size
255
+ self.num_patch_feature_size = num_patch_feature_size
256
self.image_preprocessor = Step3VisionProcessor(
258
- self.image_size, "bilinear", self.patch_size
257
+ image_size, "bilinear", patch_size
258
)
260
-
261
- self.num_image_feature_size = 169
262
- self.num_patch_feature_size = 81
263
- self.image_token = "<im_patch>"
264
- self.image_feature_placeholder = self.image_token * self.num_image_feature_size
265
- self.patch_feature_placeholder = self.image_token * self.num_patch_feature_size
266
-
267
- # Respect vision config switch to enable/disable patch extraction.
268
- # For video understanding, it's preferable to disable patch.
269
- enable_patch = getattr(self.config.vision_config, "enable_patch", True)
259
self.patcher = ImagePatcher(enable_patch=enable_patch)
260
272
- @property
273
- def image_token_id(self) -> int:
274
- return self.tokenizer.get_vocab()[self.image_token]
275
-
261
def get_num_image_tokens(self, img_width: int, img_height: int) -> int:
262
num_patches, num_newlines = self.patcher.get_num_patches(img_width, img_height)
263
@@ -299,58 +284,168 @@ class Step3VLProcessor:
284
for img in images
285
]
286
302
- def _get_patch_repl(
287
+ def __call__(
288
+ self,
289
+ images: Image.Image | list[Image.Image] | None = None,
290
+ return_tensors: str | TensorType | None = None,
291
+ ) -> BatchFeature:
292
+ if images is None:
293
+ images = []
294
+ if not isinstance(images, list):
295
+ images = [images]
296
+
297
+ split_images_data = self._split_images(images)
298
+ pixel_values_lst = []
299
+ patch_pixel_values_lst = []
300
+ patch_newline_mask_lst = []
301
+ num_patches = []
302
+ for raw_img, img_patches, patch_newline_mask in split_images_data:
303
+ pixel_values_lst.extend(self._convert_images_to_pixel_values([raw_img]))
304
+ num_patches.append(len(img_patches))
305
+ patch_pixel_values_lst.extend(
306
+ self._convert_images_to_pixel_values(img_patches, is_patch=True)
307
+ )
308
+ patch_newline_mask_lst.extend(patch_newline_mask)
309
+
310
+ pixel_values = torch.cat(pixel_values_lst)
311
+ patch_size = self.patch_size
312
+ image_inputs = {
313
+ "pixel_values": pixel_values,
314
+ "num_patches": num_patches,
315
+ "patch_pixel_values": (
316
+ torch.cat(patch_pixel_values_lst)
317
+ if patch_pixel_values_lst
318
+ else pixel_values.new_empty((0, 3, patch_size, patch_size))
319
+ ),
320
+ "patch_newline_mask": torch.tensor(
321
+ patch_newline_mask_lst, dtype=torch.bool
322
+ ),
323
+ }
324
+ return BatchFeature(image_inputs, tensor_type=return_tensors)
325
+
326
+
327
+class Step3VLProcessor(ProcessorMixin):
328
+ attributes = ["image_processor", "tokenizer"]
329
+
330
+ def __init__(
331
+ self,
332
+ image_processor: Step3VLImageProcessor,
333
+ tokenizer: TokenizerLike,
334
+ ) -> None:
335
+ self.image_processor = image_processor
336
+ self.tokenizer = tokenizer
337
+
338
+ self.image_start_token = image_start_token = "<im_start>"
339
+ self.image_end_token = image_end_token = "<im_end>"
340
+ self.patch_start_token = patch_start_token = "<patch_start>"
341
+ self.patch_end_token = patch_end_token = "<patch_end>"
342
+ self.patch_newline_token = patch_newline_token = "<patch_newline>"
343
+ self.image_start_token_id = tokenizer.convert_tokens_to_ids(image_start_token)
344
+ self.image_end_token_id = tokenizer.convert_tokens_to_ids(image_end_token)
345
+ self.patch_start_token_id = tokenizer.convert_tokens_to_ids(patch_start_token)
346
+ self.patch_end_token_id = tokenizer.convert_tokens_to_ids(patch_end_token)
347
+ self.patch_newline_token_id = tokenizer.convert_tokens_to_ids(
348
+ patch_newline_token
349
+ )
350
+
351
+ self.image_token = image_token = "<im_patch>"
352
+ self.image_feature_tokens = image_token * image_processor.num_image_feature_size
353
+ self.patch_feature_tokens = image_token * image_processor.num_patch_feature_size
354
+
355
+ self.image_token_id = image_token_id = tokenizer.convert_tokens_to_ids(
356
+ image_token
357
+ )
358
+ self.image_feature_token_ids = [
359
+ image_token_id
360
+ ] * image_processor.num_image_feature_size
361
+ self.patch_feature_token_ids = [
362
+ image_token_id
363
+ ] * image_processor.num_patch_feature_size
364
+
365
+ def _get_patch_repl_text(
366
self,
367
num_patches: int,
305
- patch_newline_mask: list[bool] | None,
306
- ) -> tuple[str, list[int]]:
307
- text = ""
308
- token_ids = []
368
+ patch_newline_mask: list[bool],
369
+ ) -> str:
370
+ assert len(patch_newline_mask) == num_patches
371
+
372
+ parts = []
373
for i in range(num_patches):
310
- assert (
311
- patch_newline_mask is not None
312
- and len(patch_newline_mask) == num_patches
374
+ parts.extend(
375
+ [
376
+ self.patch_start_token,
377
+ self.patch_feature_tokens,
378
+ self.patch_end_token,
379
+ ]
380
)
314
- text += f"<patch_start>{self.patch_feature_placeholder}<patch_end>"
315
- token_ids.extend(
316
- [self.tokenizer.convert_tokens_to_ids("<patch_start>")]
317
- + [self.image_token_id] * self.num_patch_feature_size
318
- + [self.tokenizer.convert_tokens_to_ids("<patch_end>")]
381
+ if patch_newline_mask[i]:
382
+ parts.append(self.patch_newline_token)
383
+
384
+ return "".join(parts)
385
+
386
+ def _get_patch_repl_ids(
387
+ self,
388
+ num_patches: int,
389
+ patch_newline_mask: list[bool],
390
+ ) -> list[int]:
391
+ assert len(patch_newline_mask) == num_patches
392
+
393
+ parts = []
394
+ for i in range(num_patches):
395
+ parts.extend(
396
+ [
397
+ self.patch_start_token_id,
398
+ *self.patch_feature_token_ids,
399
+ self.patch_end_token_id,
400
+ ]
401
)
320
- if patch_newline_mask and patch_newline_mask[i]:
321
- text += "<patch_newline>"
322
- token_ids.append(
323
- self.tokenizer.convert_tokens_to_ids("<patch_newline>")
324
- )
325
- return text, token_ids
402
+ if patch_newline_mask[i]:
403
+ parts.append(self.patch_newline_token_id)
404
327
- def _get_image_repl(
405
+ return parts
406
+
407
+ def _get_image_repl_text(
408
self,
409
num_images: int,
330
- ) -> tuple[str, list[int]]:
331
- text = f"<im_start>{self.image_feature_placeholder}<im_end>"
332
- token_ids = (
333
- [self.tokenizer.convert_tokens_to_ids("<im_start>")]
334
- + [self.image_token_id] * self.num_image_feature_size
335
- + [self.tokenizer.convert_tokens_to_ids("<im_end>")]
336
- )
337
- return text * num_images, token_ids * num_images
410
+ ) -> str:
411
+ parts = [
412
+ self.image_start_token,
413
+ self.image_feature_tokens,
414
+ self.image_end_token,
415
+ ] * num_images
416
+
417
+ return "".join(parts)
418
+
419
+ def _get_image_repl_ids(
420
+ self,
421
+ num_images: int,
422
+ ) -> list[int]:
423
+ part = [
424
+ self.image_start_token_id,
425
+ *self.image_feature_token_ids,
426
+ self.image_end_token_id,
427
+ ]
428
+ return part * num_images
429
339
- def _get_image_repl_features(
430
+ def get_image_repl_feature_text(
431
self,
432
num_images: int,
433
num_patches: int,
343
- patch_new_line_idx: list[bool] | None,
344
- ) -> tuple[str, list[int]]:
345
- if num_patches > 0:
346
- patch_repl, patch_repl_ids = self._get_patch_repl(
347
- num_patches, patch_new_line_idx
348
- )
349
- else:
350
- patch_repl = ""
351
- patch_repl_ids = []
352
- image_repl, image_repl_ids = self._get_image_repl(num_images)
353
- return patch_repl + image_repl, patch_repl_ids + image_repl_ids
434
+ patch_new_line_idx: list[bool],
435
+ ) -> str:
436
+ patch_repl = self._get_patch_repl_text(num_patches, patch_new_line_idx)
437
+ image_repl = self._get_image_repl_text(num_images)
438
+ return patch_repl + image_repl
439
+
440
+ def get_image_repl_feature_ids(
441
+ self,
442
+ num_images: int,
443
+ num_patches: int,
444
+ patch_new_line_idx: list[bool],
445
+ ) -> list[int]:
446
+ patch_repl = self._get_patch_repl_ids(num_patches, patch_new_line_idx)
447
+ image_repl = self._get_image_repl_ids(num_images)
448
+ return patch_repl + image_repl
449
450
def replace_placeholder(self, text: str, placeholder: str, repls: list[str]) -> str:
451
parts = text.split(placeholder)
@@ -373,69 +468,44 @@ class Step3VLProcessor:
468
images: Image.Image | list[Image.Image] | None = None,
469
return_tensors: str | TensorType | None = None,
470
) -> BatchFeature:
376
- if text is None:
377
- text = []
378
- if not isinstance(text, list):
379
- text = [text]
380
- if images is None:
381
- images = []
382
- if not isinstance(images, list):
383
- images = [images]
384
-
385
- if len(images) == 0:
386
- image_inputs = {}
387
- text_inputs = self.tokenizer(text)
471
+ if images is not None:
472
+ image_inputs = self.image_processor(
473
+ images=images,
474
+ return_tensors=return_tensors,
475
+ )
476
+ num_patches = image_inputs["num_patches"]
477
+ patch_newline_mask = image_inputs["patch_newline_mask"]
478
else:
389
- split_images_data = self._split_images(images)
390
- pixel_values_lst = []
391
- patch_pixel_values_lst = []
392
- patch_newline_mask_lst = []
393
- image_repl_str_lst = []
394
- image_repl_ids_lst = []
479
+ image_inputs = {}
480
num_patches = []
396
- for raw_img, img_patches, patch_newline_mask in split_images_data:
397
- pixel_values_lst.extend(self._convert_images_to_pixel_values([raw_img]))
398
-
399
- if len(img_patches) > 0:
400
- patch_pixel_values_lst.extend(
401
- self._convert_images_to_pixel_values(img_patches, is_patch=True)
481
+ patch_newline_mask = []
482
+
483
+ if text is not None:
484
+ if not isinstance(text, list):
485
+ text = [text]
486
+
487
+ if image_inputs:
488
+ image_token = self.image_token
489
+ image_repl_str_lst = []
490
+ start = 0
491
+ for n_patches in num_patches:
492
+ image_repl_str = self.get_image_repl_feature_text(
493
+ 1, n_patches, patch_newline_mask[start : start + n_patches]
494
)
403
- num_patches.append(len(img_patches))
495
+ image_repl_str_lst.append(image_repl_str)
496
405
- image_repl_str, image_repl_ids = self._get_image_repl_features(
406
- 1, len(img_patches), patch_newline_mask
407
- )
408
- image_repl_str_lst.append(image_repl_str)
409
- image_repl_ids_lst.extend(image_repl_ids)
410
-
411
- if patch_newline_mask is not None:
412
- patch_newline_mask_lst.extend(patch_newline_mask)
413
-
414
- pixel_values = torch.cat(pixel_values_lst)
415
- patch_size = self.patch_size
416
- image_inputs = {
417
- "pixel_values": pixel_values,
418
- "num_patches": num_patches,
419
- "patch_pixel_values": (
420
- torch.cat(patch_pixel_values_lst)
421
- if patch_pixel_values_lst
422
- else pixel_values.new_empty((0, 3, patch_size, patch_size))
423
- ),
424
- "patch_newline_mask": torch.tensor(
425
- patch_newline_mask_lst, dtype=torch.bool
426
- ),
427
- }
497
+ start += n_patches
498
+
499
+ text = [
500
+ self.replace_placeholder(t, image_token, image_repl_str_lst)
501
+ for t in text
502
+ ]
503
429
- text = [
430
- self.replace_placeholder(t, self.image_token, image_repl_str_lst)
431
- for t in text
432
- ]
504
text_inputs = self.tokenizer(text)
505
+ else:
506
+ text_inputs = {}
507
508
return BatchFeature(
436
- {
437
- **text_inputs,
438
- **image_inputs,
439
- },
509
+ data={**text_inputs, **image_inputs},
510
tensor_type=return_tensors,
511
)