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
[Bugfix] Added `embed_is_patch` mask for fuyu model (#15731)
Signed-off-by: Kyle Huang <kylhuang@nvidia.com>
kYLe committed
Mar 30, 2025 at 05:45 UTC
bb103b29bf3d26fa831458a8ca5af1bb820c4712
1 file changed
+36
-5
vllm/model_executor/models/fuyu.py
+36
-5
@@ -18,7 +18,7 @@
18
""" PyTorch Fuyu model."""
19
import math
20
from collections.abc import Iterable, Mapping, Sequence
21
-from typing import Literal, Optional, Set, Tuple, TypedDict
21
+from typing import Literal, Optional, Set, Tuple, TypedDict, Union
22
23
import torch
24
import torch.nn as nn
@@ -39,10 +39,12 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
39
PromptUpdate, PromptUpdateDetails)
40
from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
41
from vllm.sequence import IntermediateTensors
42
+from vllm.utils import flatten_2d_lists
43
44
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
45
from .utils import (AutoWeightsLoader, flatten_bn, maybe_prefix,
46
merge_multimodal_embeddings)
47
+from .vision import scatter_patch_features, select_patch_features
48
49
# Cannot find the following 2 numbers from hf config.
50
_IMAGE_TOKEN_ID = 71011
@@ -64,6 +66,11 @@ class FuyuImagePatchInputs(TypedDict):
66
This is used to split the embeddings which has the first two dimensions
67
flattened just like `flat_data`.
68
"""
69
+ embed_is_patch: Union[torch.Tensor, list[torch.Tensor]]
70
+ """
71
+ A boolean mask indicating which image embeddings correspond
72
+ to patch tokens.
73
+ """
74
75
76
class FuyuProcessingInfo(BaseProcessingInfo):
@@ -183,6 +190,19 @@ class FuyuMultiModalProcessor(BaseMultiModalProcessor[FuyuProcessingInfo]):
190
191
processed_outputs["image_patches"] = image_patches[0]
192
193
+ # get patch grid size for each image
194
+ embed_is_patch = []
195
+ for image in images:
196
+ ncols, nrows = self.info.get_image_feature_grid_size(
197
+ image_width=image.width,
198
+ image_height=image.height,
199
+ )
200
+
201
+ mask = torch.tensor(([True] * ncols + [False]) * nrows)
202
+ embed_is_patch.append(mask)
203
+
204
+ processed_outputs["embed_is_patch"] = embed_is_patch
205
+
206
return processed_outputs
207
208
def _apply_hf_processor_tokens_only(
@@ -202,7 +222,8 @@ class FuyuMultiModalProcessor(BaseMultiModalProcessor[FuyuProcessingInfo]):
222
hf_inputs: BatchFeature,
223
hf_processor_mm_kwargs: Mapping[str, object],
224
) -> Mapping[str, MultiModalFieldConfig]:
205
- return dict(image_patches=MultiModalFieldConfig.batched("image"))
225
+ return dict(image_patches=MultiModalFieldConfig.batched("image"),
226
+ embed_is_patch=MultiModalFieldConfig.batched("image"))
227
228
def _get_prompt_updates(
229
self,
@@ -301,11 +322,15 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
322
def _parse_and_validate_image_input(
323
self, **kwargs: object) -> Optional[FuyuImagePatchInputs]:
324
image_patches = kwargs.pop("image_patches", None)
325
+ embed_is_patch = kwargs.pop("embed_is_patch", None)
326
if image_patches is not None:
327
if not isinstance(image_patches, (torch.Tensor, list)):
328
raise ValueError("Incorrect type of image patches. "
329
f"Got type: {type(image_patches)}")
330
331
+ if not isinstance(embed_is_patch, (torch.Tensor, list)):
332
+ raise ValueError("Incorrect type of embed_is_patch. "
333
+ f"Got type: {type(embed_is_patch)}")
334
image_patches_flat = flatten_bn(image_patches)
335
336
return FuyuImagePatchInputs(
@@ -313,6 +338,7 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
338
flat_data=self._validate_pixel_values(
339
flatten_bn(image_patches_flat, concat=True)),
340
patches_per_image=[x.size(0) for x in image_patches_flat],
341
+ embed_is_patch=embed_is_patch,
342
)
343
344
return None
@@ -333,7 +359,12 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
359
if image_input is None:
360
return None
361
vision_embeddings = self._process_image_input(image_input)
336
- return vision_embeddings
362
+ #return vision_embeddings
363
+ return flatten_2d_lists(
364
+ scatter_patch_features(*args) for args in zip(
365
+ vision_embeddings,
366
+ image_input["embed_is_patch"],
367
+ ))
368
369
def get_input_embeddings(
370
self,
@@ -343,8 +374,8 @@ class FuyuForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
374
inputs_embeds = self.language_model.get_input_embeddings(input_ids)
375
if multimodal_embeddings is not None:
376
inputs_embeds = merge_multimodal_embeddings(
346
- input_ids, inputs_embeds, multimodal_embeddings,
347
- _IMAGE_TOKEN_ID)
377
+ input_ids, inputs_embeds,
378
+ select_patch_features(multimodal_embeddings), _IMAGE_TOKEN_ID)
379
return inputs_embeds
380
381
def forward(