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
Migrate skyworkr1v inputs to TensorSchema (#23499)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Aug 24, 2025 at 21:43 UTC
a5203d04dffcbdb095651ca4bf06589409370301
1 file changed
+36
-38
vllm/model_executor/models/skyworkr1v.py
+36
-38
@@ -8,7 +8,7 @@
8
# Licensed under The MIT License [see LICENSE for details]
9
# --------------------------------------------------------
10
from collections.abc import Iterable, Mapping, Sequence
11
-from typing import Literal, Optional, TypedDict, Union
11
+from typing import Annotated, Literal, Optional, Union
12
13
import torch
14
import torch.nn as nn
@@ -35,6 +35,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
35
from vllm.multimodal.profiling import BaseDummyInputsBuilder
36
from vllm.sequence import IntermediateTensors
37
from vllm.transformers_utils.tokenizer import AnyTokenizer
38
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
39
40
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
41
from .utils import (AutoWeightsLoader, flatten_bn, init_vllm_registered_model,
@@ -48,27 +49,42 @@ IMAGENET_MEAN = (0.485, 0.456, 0.406)
49
IMAGENET_STD = (0.229, 0.224, 0.225)
50
51
51
-class SkyworkR1VImagePixelInputs(TypedDict):
52
- type: Literal["pixel_values"]
53
- pixel_values_flat: torch.Tensor
52
+class SkyworkR1VImagePixelInputs(TensorSchema):
53
"""
55
- Shape:
56
- `(batch_size * num_images * (1 + num_patches), num_channels, height, width)`
54
+ Dimensions:
55
+ - bnp: Batch size * number of images * (1 + num_patches)
56
+ - c: Number of channels (3)
57
+ - h: Height
58
+ - w: Width
59
+ - bn: Batch size * number of images
60
"""
61
+ type: Literal["pixel_values"] = "pixel_values"
62
59
- num_patches: torch.Tensor
60
- """Shape: `(batch_size * num_images)`"""
63
+ pixel_values_flat: Annotated[
64
+ torch.Tensor,
65
+ TensorShape("bnp", 3, "h", "w"),
66
+ ]
67
68
+ num_patches: Annotated[
69
+ torch.Tensor,
70
+ TensorShape("bn"),
71
+ ]
72
63
-class SkyworkR1VImageEmbeddingInputs(TypedDict):
64
- type: Literal["image_embeds"]
65
- data: Union[torch.Tensor, list[torch.Tensor]]
66
- """
67
- A tensor of shape `(num_images, total_image_feature_size, hidden_size)`
68
- or a list of tensors of shape `(total_image_feature_size, hidden_size)`
73
70
- `hidden_size` must match the hidden size of language model backbone.
74
+class SkyworkR1VImageEmbeddingInputs(TensorSchema):
75
"""
76
+ Dimensions:
77
+ - ni: Number of images
78
+ - ifs: Image feature size
79
+ - hs: Hidden size (must match the hidden size of language model
80
+ backbone)
81
+ """
82
+ type: Literal["image_embeds"] = "image_embeds"
83
+
84
+ data: Annotated[
85
+ Union[torch.Tensor, list[torch.Tensor]],
86
+ TensorShape("ni", "ifs", "hs"),
87
+ ]
88
89
90
SkyworkR1VImageInputs = Union[SkyworkR1VImagePixelInputs,
@@ -731,26 +747,6 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):
747
vit_embeds = self.mlp1(vit_embeds)
748
return vit_embeds
749
734
- def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:
735
-
736
- h = w = self.config.vision_config.image_size
737
- expected_dims = (3, h, w)
738
-
739
- def _validate_shape(d: torch.Tensor):
740
- actual_dims = tuple(d.shape)
741
-
742
- if actual_dims != expected_dims:
743
- expected_expr = str(expected_dims)
744
- raise ValueError(
745
- "The expected shape of pixel values per image per batch "
746
- f" per patch is {expected_expr}. "
747
- f"You supplied {tuple(d.shape)}.")
748
-
749
- for d in data:
750
- _validate_shape(d)
751
-
752
- return data
753
-
750
def _parse_and_validate_image_input(
751
self, **kwargs: object) -> Optional[SkyworkR1VImageInputs]:
752
pixel_values_flat = kwargs.pop("pixel_values_flat", None)
@@ -788,10 +784,12 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):
784
785
return SkyworkR1VImagePixelInputs(
786
type="pixel_values",
791
- pixel_values_flat=self._validate_pixel_values(
792
- pixel_values_flat),
787
+ pixel_values_flat=pixel_values_flat,
788
num_patches=image_num_patches,
794
- )
789
+ resolve_bindings={
790
+ "h": self.config.vision_config.image_size,
791
+ "w": self.config.vision_config.image_size,
792
+ })
793
794
raise AssertionError("This line should be unreachable.")
795