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, 23:45 UTC
Migrate Idefics3ImagePixelInputs and Idefics3ImageEmbeddingInputs to … (#21683)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Jul 27, 2025 at 22:37 UTC
3ea57a56d9113ffb81673918931a69058fce4ae1
1 file changed
+28
-41
vllm/model_executor/models/idefics3.py
+28
-41
@@ -18,7 +18,7 @@
18
19
import math
20
from collections.abc import Iterable, Mapping, Sequence
21
-from typing import Literal, Optional, TypedDict, Union
21
+from typing import Annotated, Literal, Optional, Union
22
23
import torch
24
from torch import nn
@@ -45,6 +45,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
45
# yapf: enable
46
from vllm.multimodal.profiling import BaseDummyInputsBuilder
47
from vllm.sequence import IntermediateTensors
48
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
49
50
# yapf: disable
51
from .idefics2_vision_model import (
@@ -56,26 +57,30 @@ from .utils import (AutoWeightsLoader, flatten_bn, maybe_prefix,
57
merge_multimodal_embeddings)
58
59
59
-class Idefics3ImagePixelInputs(TypedDict):
60
- type: Literal["pixel_values"]
61
- pixel_values: torch.Tensor
60
+class Idefics3ImagePixelInputs(TensorSchema):
61
"""
63
- Shape: `(batch_size * num_images * num_patches,
64
- num_channels, height, width)`
62
+ Dimensions:
63
+ - bn: Batch size * number of images
64
+ - bnp: Batch size * number of images * number of patches
65
+ - c: Number of channels (3)
66
+ - h: Height
67
+ - w: Width
68
"""
69
+ type: Literal["pixel_values"]
70
+ pixel_values: Annotated[torch.Tensor, TensorShape("bnp", 3, "h", "w")]
71
pixel_attention_mask: torch.Tensor
67
-
68
- num_patches: torch.Tensor
69
- """Shape: `(batch_size * num_images)`"""
72
+ num_patches: Annotated[torch.Tensor, TensorShape("bn")]
73
74
72
-class Idefics3ImageEmbeddingInputs(TypedDict):
73
- type: Literal["image_embeds"]
74
- data: torch.Tensor
75
+class Idefics3ImageEmbeddingInputs(TensorSchema):
76
"""
76
- Shape: `(batch_size * num_images, image_feature_size, hidden_size)`
77
- `hidden_size` must match the hidden size of language model backbone.
77
+ Dimensions:
78
+ - bn: Batch size * number of images
79
+ - f: Image feature size
80
+ - h: Hidden size (must match the hidden size of language model backbone)
81
"""
82
+ type: Literal["image_embeds"]
83
+ data: Annotated[torch.Tensor, TensorShape("bn", "f", "h")]
84
85
86
ImageInputs = Union[Idefics3ImagePixelInputs, Idefics3ImageEmbeddingInputs]
@@ -614,25 +619,6 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal,
619
self.lm_head.weight = self.model.text_model.wte.weight
620
self.logits_processor = LogitsProcessor(config.text_config.vocab_size)
621
617
- def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
618
- h = w = self.config.vision_config.image_size
619
- expected_dims = (3, h, w)
620
-
621
- def _validate_shape(d: torch.Tensor):
622
- actual_dims = tuple(d.shape)
623
-
624
- if actual_dims != expected_dims:
625
- expected_expr = str(expected_dims)
626
- raise ValueError(
627
- "The expected shape of pixel values per image per batch "
628
- f" per patch is {expected_expr}. "
629
- f"You supplied {tuple(d.shape)}.")
630
-
631
- for d in data:
632
- _validate_shape(d)
633
-
634
- return data
635
-
622
def _parse_and_validate_image_input(
623
self, **kwargs: object) -> Optional[ImageInputs]:
624
pixel_values = kwargs.pop("pixel_values", None)
@@ -666,16 +652,17 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal,
652
raise ValueError("Incorrect type of num_patches. "
653
f"Got type: {type(num_patches)}")
654
669
- pixel_values = flatten_bn(pixel_values, concat=True)
670
- pixel_attention_mask = flatten_bn(pixel_attention_mask,
671
- concat=True)
672
- num_patches = flatten_bn(num_patches, concat=True)
673
-
655
+ expected_h = expected_w = self.config.vision_config.image_size
656
return Idefics3ImagePixelInputs(
657
type="pixel_values",
676
- pixel_values=self._validate_pixel_values(pixel_values),
677
- pixel_attention_mask=pixel_attention_mask,
678
- num_patches=num_patches,
658
+ pixel_values=flatten_bn(pixel_values, concat=True),
659
+ pixel_attention_mask=flatten_bn(pixel_attention_mask,
660
+ concat=True),
661
+ num_patches=flatten_bn(num_patches, concat=True),
662
+ resolve_bindings={
663
+ "h": expected_h,
664
+ "w": expected_w
665
+ },
666
)
667
668
raise AssertionError("This line should be unreachable.")