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, 17:15 UTC
[Misc] Remove redundant `num_embeds` (#15443)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Mar 25, 2025 at 18:27 UTC
5994430b8433b89e07bf55b266b449c1c3d6d3cd
5 files changed
+25
-64
vllm/model_executor/models/gemma3_mm.py
-16
@@ -63,9 +63,6 @@ class Gemma3ImagePixelInputs(TypedDict):
63
Shape: `(batch_size, num_images, num_embeds)`
64
"""
65
66
- num_embeds: Union[torch.Tensor, list[torch.Tensor]]
67
- """Shape: `(batch_size, num_images)`"""
68
-
66
67
Gemma3ImageInputs = Gemma3ImagePixelInputs
68
@@ -317,11 +314,6 @@ class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):
314
tokenizer.encode(image_repl, add_special_tokens=False)
315
for image_repl in image_repl_features
316
]
320
- num_embeds = [
321
- len(image_repl_feature_tokens)
322
- for image_repl_feature_tokens in image_repls_feature_tokens
323
- ]
324
- processed_outputs["num_embeds"] = torch.tensor(num_embeds)
317
318
vocab = tokenizer.get_vocab()
319
image_token_id = vocab[tokenizer.image_token]
@@ -354,7 +346,6 @@ class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):
346
"image", num_crops + 1),
347
num_crops=MultiModalFieldConfig.batched("image"),
348
embed_is_patch=MultiModalFieldConfig.batched("image"),
357
- num_embeds=MultiModalFieldConfig.batched("image"),
349
)
350
351
def _get_prompt_updates(
@@ -583,7 +574,6 @@ class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,
574
pixel_values = kwargs.pop("pixel_values", None)
575
num_crops = kwargs.pop("num_crops", None)
576
embed_is_patch = kwargs.pop("embed_is_patch", None)
586
- num_embeds = kwargs.pop("num_embeds", None)
577
image_embeds = kwargs.pop("image_embeds", None)
578
assert image_embeds is None, "Gemma3 does not support image_embeds."
579
if pixel_values is None:
@@ -601,10 +591,6 @@ class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,
591
raise ValueError("Incorrect type of embed_is_patch. "
592
f"Got type: {type(embed_is_patch)}")
593
604
- if not isinstance(num_embeds, (torch.Tensor, list)):
605
- raise ValueError("Incorrect type of num_embeds. "
606
- f"Got type: {type(num_embeds)}")
607
-
594
pixel_values = flatten_bn(pixel_values, concat=True)
595
num_crops = flatten_bn(num_crops, concat=True)
596
@@ -613,7 +599,6 @@ class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,
599
pixel_values=self._validate_pixel_values(pixel_values),
600
num_patches=num_crops + 1,
601
embed_is_patch=embed_is_patch,
616
- num_embeds=num_embeds,
602
)
603
604
def _image_pixels_to_features(
@@ -656,7 +641,6 @@ class Gemma3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP,
641
return flatten_2d_lists(
642
scatter_patch_features(*args) for args in zip(
643
image_features,
659
- image_input["num_embeds"],
644
image_input["embed_is_patch"],
645
))
646
vllm/model_executor/models/internvl.py
-14
@@ -69,9 +69,6 @@ class InternVLImagePixelInputs(TypedDict):
69
Shape: `(batch_size, num_images, num_embeds)`
70
"""
71
72
- num_embeds: Union[torch.Tensor, list[torch.Tensor]]
73
- """Shape: `(batch_size, num_images)`"""
74
-
72
73
class InternVLImageEmbeddingInputs(TypedDict):
74
type: Literal["image_embeds"]
@@ -426,7 +423,6 @@ class BaseInternVLProcessor(ABC):
423
tokenizer = self.tokenizer
424
image_token_id = self.image_token_id
425
429
- num_embeds = list[int]()
426
embed_is_patch = list[torch.Tensor]()
427
428
for pixel_values in pixel_values_lst:
@@ -438,11 +434,9 @@ class BaseInternVLProcessor(ABC):
434
add_special_tokens=False)
435
436
text = [t.replace('<image>', image_repl.full, 1) for t in text]
441
- num_embeds.append(len(feature_tokens))
437
embed_is_patch.append(
438
torch.tensor(feature_tokens) == image_token_id)
439
445
- image_inputs["num_embeds"] = torch.tensor(num_embeds)
440
image_inputs["embed_is_patch"] = embed_is_patch
441
442
text_inputs = self.tokenizer(text)
@@ -607,7 +601,6 @@ class InternVLMultiModalProcessor(BaseMultiModalProcessor[_I]):
601
"image", image_num_patches),
602
image_num_patches=MultiModalFieldConfig.batched("image"),
603
embed_is_patch=MultiModalFieldConfig.batched("image"),
610
- num_embeds=MultiModalFieldConfig.batched("image"),
604
image_embeds=MultiModalFieldConfig.batched("image"),
605
image_token_id=MultiModalFieldConfig.shared("image", num_images),
606
)
@@ -840,7 +833,6 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP):
833
pixel_values_flat = kwargs.pop("pixel_values_flat", None)
834
image_num_patches = kwargs.pop("image_num_patches", None)
835
embed_is_patch = kwargs.pop("embed_is_patch", None)
843
- num_embeds = kwargs.pop("num_embeds", None)
836
image_embeds = kwargs.pop("image_embeds", None)
837
838
if pixel_values_flat is None and image_embeds is None:
@@ -873,10 +865,6 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP):
865
raise ValueError("Incorrect type of embed_is_patch. "
866
f"Got type: {type(embed_is_patch)}")
867
876
- if not isinstance(num_embeds, (torch.Tensor, list)):
877
- raise ValueError("Incorrect type of num_embeds. "
878
- f"Got type: {type(num_embeds)}")
879
-
868
pixel_values_flat = flatten_bn(pixel_values_flat, concat=True)
869
image_num_patches = flatten_bn(image_num_patches, concat=True)
870
@@ -886,7 +874,6 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP):
874
pixel_values_flat),
875
num_patches=image_num_patches,
876
embed_is_patch=embed_is_patch,
889
- num_embeds=num_embeds,
877
)
878
879
raise AssertionError("This line should be unreachable.")
@@ -941,7 +928,6 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP):
928
return flatten_2d_lists(
929
scatter_patch_features(*args) for args in zip(
930
image_features,
944
- image_input["num_embeds"],
931
image_input["embed_is_patch"],
932
))
933
vllm/model_executor/models/llava.py
-16
@@ -76,9 +76,6 @@ class PixtralHFImagePixelInputs(TypedDict):
76
Shape: `(batch_size, num_images, num_embeds)`
77
"""
78
79
- num_embeds: Union[torch.Tensor, list[torch.Tensor]]
80
- """Shape: `(batch_size, num_images)`"""
81
-
79
80
class LlavaImageEmbeddingInputs(TypedDict):
81
type: Literal["image_embeds"]
@@ -358,15 +355,10 @@ class PixtralHFMultiModalProcessor(
355
image_height=pixel_value.shape[-2],
356
) for pixel_value in processed_outputs["pixel_values"]
357
]
361
- num_embeds = torch.tensor([(ncols + 1) * nrows
362
- for ncols, nrows in tile_sizes])
363
- # Each image may result to masks of different sizes, so we need to
364
- # later use `num_embeds` to get per-image masks.
358
embed_is_patch = [
359
torch.tensor(([True] * ncols + [False]) * nrows)
360
for ncols, nrows in tile_sizes
361
]
369
- processed_outputs["num_embeds"] = num_embeds
362
processed_outputs["embed_is_patch"] = embed_is_patch
363
364
return processed_outputs
@@ -378,7 +370,6 @@ class PixtralHFMultiModalProcessor(
370
) -> Mapping[str, MultiModalFieldConfig]:
371
return dict(
372
pixel_values=MultiModalFieldConfig.batched("image"),
381
- num_embeds=MultiModalFieldConfig.batched("image"),
373
embed_is_patch=MultiModalFieldConfig.batched("image"),
374
image_embeds=MultiModalFieldConfig.batched("image"),
375
)
@@ -627,16 +618,10 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
618
raise ValueError("Incorrect type of embed_is_patch. "
619
f"Got type: {type(embed_is_patch)}")
620
630
- num_embeds = kwargs.pop("num_embeds")
631
- if not isinstance(num_embeds, (torch.Tensor, list)):
632
- raise ValueError("Incorrect type of num_embeds. "
633
- f"Got type: {type(num_embeds)}")
634
-
621
return PixtralHFImagePixelInputs(
622
type="pixel_values_pixtral",
623
pixel_values=flatten_bn(pixel_values),
624
embed_is_patch=embed_is_patch,
639
- num_embeds=num_embeds,
625
)
626
627
return LlavaImagePixelInputs(
@@ -738,7 +723,6 @@ class LlavaForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
723
return flatten_2d_lists(
724
scatter_patch_features(*args) for args in zip(
725
vision_embeddings,
741
- image_input["num_embeds"],
726
image_input["embed_is_patch"],
727
))
728
vllm/model_executor/models/pixtral.py
-14
@@ -77,9 +77,6 @@ class PixtralImagePixelInputs(TypedDict):
77
Shape: `(batch_size, num_images, num_embeds)`
78
"""
79
80
- num_embeds: Union[torch.Tensor, list[torch.Tensor]]
81
- """Shape: `(batch_size, num_images)`"""
82
-
80
81
class PixtralProcessorAdapter:
82
"""
@@ -153,7 +150,6 @@ class PixtralProcessorAdapter:
150
images_processed = list[torch.Tensor]()
151
images_tokens = list[torch.Tensor]()
152
images_embed_is_patch = list[torch.Tensor]()
156
- images_num_embeds = list[int]()
153
154
for image in images:
155
image_inputs = self.image_processor(ImageChunk(image=image))
@@ -163,13 +159,11 @@ class PixtralProcessorAdapter:
159
images_processed.append(image_processed)
160
images_tokens.append(image_tokens)
161
images_embed_is_patch.append(image_tokens == image_token_id)
166
- images_num_embeds.append(len(image_tokens))
162
163
return {
164
"input_ids": torch.cat(images_tokens)[None].expand(len(text), -1),
165
"images": images_processed,
166
"embed_is_patch": images_embed_is_patch,
172
- "num_embeds": torch.tensor(images_num_embeds),
167
}
168
169
@@ -273,7 +267,6 @@ class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo]
267
return dict(
268
images=MultiModalFieldConfig.batched("image"),
269
embed_is_patch=MultiModalFieldConfig.batched("image"),
276
- num_embeds=MultiModalFieldConfig.batched("image"),
270
)
271
272
def _get_prompt_updates(
@@ -394,16 +387,10 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,
387
raise ValueError("Incorrect type of embed_is_patch. "
388
f"Got type: {type(embed_is_patch)}")
389
397
- num_embeds = kwargs.pop("num_embeds")
398
- if not isinstance(num_embeds, (torch.Tensor, list)):
399
- raise ValueError("Incorrect type of num_embeds. "
400
- f"Got type: {type(num_embeds)}")
401
-
390
return PixtralImagePixelInputs(
391
type="pixel_values",
392
images=flatten_bn(images),
393
embed_is_patch=embed_is_patch,
406
- num_embeds=num_embeds,
394
)
395
396
def _process_image_input(
@@ -447,7 +434,6 @@ class PixtralForConditionalGeneration(nn.Module, SupportsMultiModal,
434
return flatten_2d_lists(
435
scatter_patch_features(*args) for args in zip(
436
image_features,
450
- image_input["num_embeds"],
437
image_input["embed_is_patch"],
438
))
439
vllm/model_executor/models/vision.py
+25
-4
@@ -155,7 +155,6 @@ def resolve_visual_encoder_outputs(
155
156
def scatter_patch_features(
157
features: torch.Tensor,
158
- num_embeds: torch.Tensor,
158
embed_is_patch: torch.Tensor,
159
) -> tuple[torch.Tensor, ...]:
160
"""
@@ -168,13 +167,35 @@ def scatter_patch_features(
167
Args:
168
features: The patch features, concatenated across each image.
169
Shape: `(num_patch, feature_depth)`
171
- num_embeds: The number of image embeddings for each image.
172
- Shape: `(num_images,)`
170
embed_is_patch: A boolean mask indicating which image embeddings
171
correspond to patch tokens for each image.
172
Shape: `(num_images, num_embeds)`
173
+
174
+ Note:
175
+ The original code only considers patch tokens as feature
176
+ tokens, but our processor considers all image-related tokens
177
+ as feature tokens because the feature tokens need to be
178
+ consecutive in `input_ids`.
179
+
180
+ Example:
181
+ A simplified example for one image:
182
+
183
+ .. code-block::
184
+
185
+ Embedding tokens (from HF processor):
186
+ [<start> <patch> <patch> <col> <patch> <patch> <col> <end> ]
187
+
188
+ embed_is_patch (from HF processor):
189
+ [ False True True False True True False False ]
190
+
191
+ Encoder outputs (from model):
192
+ [ p1 p2 p3 p4 ]
193
+
194
+ The resulting embedding tensor is:
195
+ [ nan p1 p2 nan p3 p4 nan nan ]
196
"""
177
- num_embeds_per_image: list[int] = num_embeds.tolist()
197
+ num_images, num_embeds = embed_is_patch.shape
198
+ num_embeds_per_image = [num_embeds] * num_images
199
200
embeds_flat = features.new_full(
201
(sum(num_embeds_per_image), features.shape[-1]),