@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / afc47e4de7

Fork 0

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

[Model] Use merge_by_field_config for MM models (M-N) (#26710)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Oct 14, 2025 at 01:27 UTC afc47e4de7a1bcce8ca8b7f767ee2994eba6a923

11 files changed +126 -330

vllm/model_executor/models/interns1.py

+10 -4

@@ -631,8 +631,11 @@ class InternS1ForConditionalGeneration(

631 )

632

633 image_token_id = kwargs["image_token_id"]

634 - assert isinstance(image_token_id, torch.Tensor)

635 - self.img_context_token_id = image_token_id.flatten().unique().item()

634 + if isinstance(image_token_id, torch.Tensor):

635 + image_token_id = image_token_id.flatten().unique().item()

636 +

637 + assert isinstance(image_token_id, int)

638 + self.img_context_token_id = image_token_id

639

640 if pixel_values is not None:

641 h, w = self.config.vision_config.image_size

@@ -665,8 +668,11 @@ class InternS1ForConditionalGeneration(

668 )

669

670 video_token_id = kwargs["video_token_id"]

668 - assert isinstance(video_token_id, torch.Tensor)

669 - self.video_context_token_id = video_token_id.flatten().unique().item()

671 + if isinstance(video_token_id, torch.Tensor):

672 + video_token_id = video_token_id.flatten().unique().item()

673 +

674 + assert isinstance(video_token_id, int)

675 + self.video_context_token_id = video_token_id

676

677 if pixel_values_flat_video is not None:

678 h, w = self.config.vision_config.image_size

vllm/model_executor/models/internvl.py

+10 -4

@@ -1232,8 +1232,11 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)

1232 )

1233

1234 image_token_id = kwargs["image_token_id"]

1235 - assert isinstance(image_token_id, torch.Tensor)

1236 - self.img_context_token_id = image_token_id.flatten().unique().item()

1235 + if isinstance(image_token_id, torch.Tensor):

1236 + image_token_id = image_token_id.flatten().unique().item()

1237 +

1238 + assert isinstance(image_token_id, int)

1239 + self.img_context_token_id = image_token_id

1240

1241 if pixel_values_flat is not None:

1242 expected_h = expected_w = self.config.vision_config.image_size

@@ -1265,8 +1268,11 @@ class InternVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA)

1268 )

1269

1270 video_token_id = kwargs["video_token_id"]

1268 - assert isinstance(video_token_id, torch.Tensor)

1269 - self.video_context_token_id = video_token_id.flatten().unique().item()

1271 + if isinstance(video_token_id, torch.Tensor):

1272 + video_token_id = video_token_id.flatten().unique().item()

1273 +

1274 + assert isinstance(video_token_id, int)

1275 + self.video_context_token_id = video_token_id

1276

1277 if pixel_values_flat_video is not None:

1278 expected_h = expected_w = self.config.vision_config.image_size

vllm/model_executor/models/midashenglm.py

+27 -53

@@ -26,7 +26,7 @@

26 import collections

27 import collections.abc

28 from collections.abc import Callable, Iterable, Mapping, Sequence

29 -from typing import Any, TypeAlias, TypedDict, cast

29 +from typing import Annotated, Any, TypeAlias, cast

31 import numpy as np

32 import torch

@@ -62,6 +62,7 @@ from vllm.multimodal.processing import (

62 from vllm.multimodal.profiling import BaseDummyInputsBuilder

63 from vllm.sequence import IntermediateTensors

64 from vllm.transformers_utils.configs.midashenglm import DashengConfig

65 +from vllm.utils.tensor_schema import TensorSchema, TensorShape

67 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

68 from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix

@@ -508,11 +509,16 @@ class AudioProjectorSubsample(nn.Module):

509

510

511 # === Audio Inputs === #

511 -class MiDashengLMAudioInputs(TypedDict):

512 - input_values: torch.Tensor

513 - """Shape: `(num_audios, num_sampling_points)`"""

514 - audio_length: torch.Tensor

515 - """Shape: `(num_audios, 1)`"""

512 +class MiDashengLMAudioInputs(TensorSchema):

513 + """

514 +

515 + Dimensions:

516 + - bn: Batch size * number of audios

517 + - p: Number of sampling points

518 + """

519 +

520 + input_values: Annotated[torch.Tensor, TensorShape("n", "p")]

521 + audio_length: Annotated[torch.Tensor, TensorShape("n")]

522

523

524 class MiDashengLMProcessingInfo(BaseProcessingInfo):

@@ -676,6 +682,8 @@ class MiDashengLMMultiModalProcessor(

682 dummy_inputs=MiDashengLMDummyInputsBuilder,

683 )

684 class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP):

685 + merge_by_field_config = True

686 +

687 packed_modules_mapping = {

688 "qkv_proj": [

689 "q_proj",

@@ -728,26 +736,6 @@ class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP):

736 self.decoder.make_empty_intermediate_tensors

737 )

738

731 - def _validate_and_reshape_mm_tensor(

732 - self, mm_input: object, name: str

733 - ) -> torch.Tensor:

734 - if not isinstance(mm_input, (torch.Tensor, list)):

735 - raise ValueError(f"Incorrect type of {name}. Got type: {type(mm_input)}")

736 - if isinstance(mm_input, torch.Tensor):

737 - return mm_input.reshape(-1, *mm_input.shape[2:])

738 -

739 - if name == "input_values":

740 - max_length = max(tensor.shape[1] for tensor in mm_input)

741 - padded_mm_input = [

742 - torch.nn.functional.pad(tensor, (0, max_length - tensor.shape[1]))

743 - if tensor.shape[1] < max_length

744 - else tensor

745 - for tensor in mm_input

746 - ]

747 - return torch.concat(padded_mm_input)

748 -

749 - return torch.concat(mm_input)

750 -

739 def _parse_and_validate_audio_input(

740 self, **kwargs: object

741 ) -> MiDashengLMAudioInputs | None:

@@ -756,16 +744,11 @@ class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP):

744

745 if input_values is None:

746 return None

759 - input_values = self._validate_and_reshape_mm_tensor(

760 - input_values, "input_values"

761 - )

762 - audio_length = self._validate_and_reshape_mm_tensor(

763 - audio_length, "audio_length"

764 - )

765 - if not isinstance(input_values, (torch.Tensor, list)):

766 - raise ValueError(

767 - "Incorrect type of audio input features. "

768 - f"Got type: {type(input_values)}"

747 +

748 + if isinstance(input_values, list):

749 + input_values = torch.nn.utils.rnn.pad_sequence(

750 + input_values,

751 + batch_first=True,

752 )

753

754 return MiDashengLMAudioInputs(

@@ -773,7 +756,10 @@ class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP):

756 audio_length=audio_length,

757 )

758

776 - def _process_audio_input(self, audio_input: MiDashengLMAudioInputs) -> torch.Tensor:

759 + def _process_audio_input(

760 + self,

761 + audio_input: MiDashengLMAudioInputs,

762 + ) -> tuple[torch.Tensor, ...]:

763 # Process audio through encoder and projector

764 input_values = audio_input["input_values"]

765 audio_length = audio_input["audio_length"]

@@ -783,17 +769,13 @@ class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP):

769 audio_embeddings = audio_embeddings.to(audio_input["input_values"].dtype)

770 batch_size, max_audio_tokens, embed_dim = audio_embeddings.shape

771

786 - audio_length_np = (

787 - audio_length.cpu().numpy()

788 - if isinstance(audio_length, torch.Tensor)

789 - else audio_length

790 - )

772 audio_output_lengths = [

773 max(1, calculate_mel_frames_dasheng(int(length))) # at least one frame

793 - for length in audio_length_np

774 + for length in audio_length.tolist()

775 ]

795 - audio_output_lengths = torch.tensor(audio_output_lengths).to(

796 - audio_embeddings.device

776 + audio_output_lengths = torch.tensor(

777 + audio_output_lengths,

778 + device=audio_embeddings.device,

779 )

780

781 audio_feature_mask = torch.arange(

@@ -826,14 +808,6 @@ class MiDashengLMModel(nn.Module, SupportsMultiModal, SupportsPP):

808 ) -> torch.Tensor | IntermediateTensors:

809 if intermediate_tensors is not None:

810 inputs_embeds = None

829 - elif inputs_embeds is None:

830 - multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)

831 - inputs_embeds = self.get_input_embeddings(

832 - input_ids,

833 - multimodal_embeddings,

834 - is_multimodal=input_ids == self.config.audio_token_id,

835 - )

836 - input_ids = None

811

812 return self.decoder.model(

813 input_ids,

vllm/model_executor/models/minicpmo.py

+5 -42

@@ -71,7 +71,7 @@ from .minicpmv import (

71 MiniCPMVProcessingInfo,

72 _minicpmv_field_config,

73 )

74 -from .utils import AutoWeightsLoader, cast_overflow_tensors, flatten_bn, maybe_prefix

74 +from .utils import AutoWeightsLoader, cast_overflow_tensors, maybe_prefix

76 CPU_DEVICE = torch.device("cpu")

@@ -132,15 +132,11 @@ MiniCPMOAudioInputs: TypeAlias = (

132

133

134 def _minicpmo_field_config(hf_inputs: Mapping[str, torch.Tensor]):

135 - audio_features = hf_inputs.get("audio_features", torch.empty(0))

136 - num_audios = len(audio_features)

137 -

135 return dict(

136 **_minicpmv_field_config(hf_inputs),

137 audio_features=MultiModalFieldConfig.batched("audio"),

138 audio_feature_lens=MultiModalFieldConfig.batched("audio"),

139 audio_embeds=MultiModalFieldConfig.batched("audio"),

143 - audio_token_id=MultiModalFieldConfig.shared("audio", num_audios),

140 )

141

142

@@ -332,10 +328,6 @@ class MiniCPMOMultiModalProcessor(MiniCPMVMultiModalProcessor[MiniCPMOProcessing

328 ]

329 audio_inputs["audio_features"] = unpadded_audio_features

330

335 - tokenizer = self.info.get_tokenizer()

336 - unk_token_id = tokenizer.get_vocab()["<unk>"]

337 - audio_inputs["audio_token_id"] = torch.tensor(unk_token_id)

338 -

331 return audio_inputs

332

333 def process_mm_inputs(

@@ -436,12 +428,10 @@ class MiniCPMWhisperEncoderLayer(nn.Module):

428 attention_mask: torch.Tensor,

429 ) -> torch.Tensor:

430 residual = hidden_states

439 - past_key_values = None

431 hidden_states = self.self_attn_layer_norm(hidden_states)

441 - hidden_states, attn_weights, past_key_values = self.self_attn(

432 + hidden_states, _ = self.self_attn(

433 hidden_states=hidden_states,

434 attention_mask=attention_mask,

444 - past_key_value=past_key_values,

435 )

436 hidden_states = nn.functional.dropout(

437 hidden_states, p=self.dropout, training=self.training

@@ -567,8 +557,6 @@ class MiniCPMO(MiniCPMV2_6):

557 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "apm")

558 )

559

570 - self.audio_token_id = None

571 -

560 def init_audio_module(self, *, vllm_config: VllmConfig, prefix: str = ""):

561 # Do not use parameters temporarily

562 audio_config = self.config.audio_config

@@ -731,43 +719,18 @@ class MiniCPMO(MiniCPMV2_6):

719 if audio_features is None and audio_embeds is None:

720 return None

721

734 - audio_token_id = kwargs.pop("audio_token_id")

735 - if audio_token_id is not None:

736 - assert isinstance(audio_token_id, torch.Tensor)

737 - self.mm_token_ids.add(audio_token_id.flatten().unique().item())

738 -

722 if audio_embeds is not None:

740 - if not isinstance(audio_embeds, (torch.Tensor, list)):

741 - raise ValueError(

742 - f"Incorrect type of audio_embeds. Got type: {type(audio_embeds)}"

743 - )

744 -

745 - audio_embeds_flat = flatten_bn(audio_embeds)

746 -

723 return MiniCPMOAudioEmbeddingInputs(

724 type="audio_embeds",

749 - audio_embeds=audio_embeds_flat,

750 - )

751 -

752 - if not isinstance(audio_features, (torch.Tensor, list)):

753 - raise ValueError(

754 - f"Incorrect type of audio_features. Got type: {type(audio_features)}"

725 + audio_embeds=audio_embeds,

726 )

727

728 audio_feature_lens = kwargs.pop("audio_feature_lens")

758 - if not isinstance(audio_feature_lens, (torch.Tensor, list)):

759 - raise ValueError(

760 - "Incorrect type of audio_feature_lens. "

761 - f"Got type: {type(audio_feature_lens)}"

762 - )

763 -

764 - audio_features_flat = flatten_bn(audio_features)

765 - audio_feature_lens_flat = flatten_bn(audio_feature_lens)

729

730 return MiniCPMOAudioFeatureInputs(

731 type="audio_features",

769 - audio_features=audio_features_flat,

770 - audio_feature_lens=audio_feature_lens_flat,

732 + audio_features=audio_features,

733 + audio_feature_lens=audio_feature_lens,

734 )

735

736 def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

vllm/model_executor/models/minicpmv.py

+8 -56

@@ -114,7 +114,7 @@ class MiniCPMVImagePixelInputs(TensorSchema):

114

115 type: Literal["pixel_values"] = "pixel_values"

116

117 - # Note that the image size may vary, so we pass it as a list instead of a

117 + # Note that the patch size may vary, so we pass it as a list instead of a

118 # batched tensor.

119 pixel_values: Annotated[

120 list[torch.Tensor],

@@ -453,12 +453,6 @@ def get_version_by_config(config: PretrainedConfig) -> tuple[int, ...]:

453

454

455 def _minicpmv_field_config(hf_inputs: Mapping[str, torch.Tensor]):

456 - pixel_values = hf_inputs.get("pixel_values", torch.empty(0))

457 - num_images = len(pixel_values)

458 -

459 - video_pixel_values = hf_inputs.get("video_pixel_values", torch.empty(0))

460 - num_videos = len(video_pixel_values)

461 -

456 return dict(

457 pixel_values=MultiModalFieldConfig.batched("image"),

458 image_sizes=MultiModalFieldConfig.batched("image"),

@@ -468,8 +462,6 @@ def _minicpmv_field_config(hf_inputs: Mapping[str, torch.Tensor]):

462 video_image_sizes=MultiModalFieldConfig.batched("video"),

463 video_tgt_sizes=MultiModalFieldConfig.batched("video"),

464 video_embeds=MultiModalFieldConfig.batched("video"),

471 - image_token_id=MultiModalFieldConfig.shared("image", num_images),

472 - video_token_id=MultiModalFieldConfig.shared("video", num_videos),

465 )

466

467

@@ -792,10 +784,6 @@ class MiniCPMVMultiModalProcessor(BaseMultiModalProcessor[_I]):

784 out_keys={"pixel_values", "image_sizes", "tgt_sizes"},

785 )

786

795 - tokenizer = self.info.get_tokenizer()

796 - unk_token_id = tokenizer.get_vocab()["<unk>"]

797 - image_inputs["image_token_id"] = torch.tensor(unk_token_id)

798 -

787 return image_inputs

788

789 def process_videos(

@@ -831,10 +819,6 @@ class MiniCPMVMultiModalProcessor(BaseMultiModalProcessor[_I]):

819

820 video_inputs = {f"video_{k}": v for k, v in video_inputs.items()}

821

834 - tokenizer = self.info.get_tokenizer()

835 - unk_token_id = tokenizer.get_vocab()["<unk>"]

836 - video_inputs["video_token_id"] = torch.tensor(unk_token_id)

837 -

822 return video_inputs

823

824 def process_mm_inputs(

@@ -1021,6 +1005,8 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1005 instantiated.

1006 """

1007

1008 + merge_by_field_config = True

1009 +

1010 supports_encoder_tp_data = True

1011

1012 @classmethod

@@ -1066,7 +1052,6 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1052 prefix=maybe_prefix(prefix, "resampler"),

1053 )

1054

1069 - self.mm_token_ids = set[int]()

1055 self.make_empty_intermediate_tensors = self.llm.make_empty_intermediate_tensors

1056

1057 def _parse_and_validate_vision_input(

@@ -1080,43 +1065,17 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1065 if pixel_values is None and image_embeds is None:

1066 return None

1067

1083 - image_token_id = kwargs.pop("image_token_id")

1084 - if image_token_id is not None:

1085 - assert isinstance(image_token_id, torch.Tensor)

1086 - self.mm_token_ids.add(image_token_id.flatten().unique().item())

1087 -

1068 if image_embeds is not None:

1089 - if not isinstance(image_embeds, (torch.Tensor, list)):

1090 - raise ValueError(

1091 - f"Incorrect type of image_embeds for {modality=}. "

1092 - f"Got type: {type(image_embeds)}"

1093 - )

1094 -

1095 - image_embeds_flat = flatten_bn(image_embeds)

1096 -

1069 return MiniCPMVImageEmbeddingInputs(

1070 type="image_embeds",

1099 - image_embeds=image_embeds_flat,

1100 - )

1101 -

1102 - if not isinstance(pixel_values, (torch.Tensor, list)):

1103 - raise ValueError(

1104 - f"Incorrect type of pixel_values for {modality=}. "

1105 - f"Got type: {type(pixel_values)}"

1071 + image_embeds=image_embeds,

1072 )

1073

1074 tgt_sizes = kwargs.pop("tgt_sizes")

1109 - if not isinstance(tgt_sizes, (torch.Tensor, list)):

1110 - raise ValueError(

1111 - f"Incorrect type of tgt_sizes for {modality=}. "

1112 - f"Got type: {type(tgt_sizes)}"

1113 - )

1114 -

1115 - num_slices = [[len(p) for p in ps] for ps in pixel_values]

1116 - num_slices_flat = flatten_bn(torch.tensor(num_slices))

1075

1118 - pixel_values_flat = flatten_bn(flatten_2d_lists(pixel_values))

1119 - tgt_sizes_flat = flatten_bn(flatten_2d_lists(tgt_sizes), concat=True)

1076 + num_slices_flat = torch.tensor([len(ps) for ps in pixel_values])

1077 + pixel_values_flat = flatten_bn(pixel_values)

1078 + tgt_sizes_flat = flatten_bn(tgt_sizes, concat=True)

1079

1080 return MiniCPMVImagePixelInputs(

1081 type="pixel_values",

@@ -1142,15 +1101,8 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1101 input_key in ("video_pixel_values", "video_embeds")

1102 and "videos" not in modalities

1103 ):

1145 -

1146 - def _image_key(video_key: str):

1147 - if video_key == "video_token_id":

1148 - return "image_token_id"

1149 -

1150 - return video_key.removeprefix("video_")

1151 -

1104 modalities["videos"] = self._parse_and_validate_vision_input(

1153 - "videos", **{_image_key(k): v for k, v in kwargs.items()}

1105 + "videos", **{k.removeprefix("video_"): v for k, v in kwargs.items()}

1106 )

1107

1108 return modalities

vllm/model_executor/models/mllama4.py

+10 -13

@@ -71,7 +71,7 @@ from .interfaces import (

71 SupportsPP,

72 )

73 from .llama4 import Llama4ForCausalLM

74 -from .utils import AutoWeightsLoader, flatten_bn, maybe_prefix

74 +from .utils import AutoWeightsLoader, maybe_prefix

75 from .vision import run_dp_sharded_vision_model

@@ -86,7 +86,7 @@ class Llama4ImagePatchInputs(TensorSchema):

87 type: Literal["pixel_values"] = "pixel_values"

89 - flat_data: Annotated[

89 + pixel_values: Annotated[

90 torch.Tensor,

91 TensorShape("total_num_chunks", "num_channels", "image_size", "image_size"),

92 ]

@@ -96,7 +96,7 @@ class Llama4ImagePatchInputs(TensorSchema):

96 The number of total patches for each image in the batch.

98 This is used to split the embeddings which has the first two dimensions

99 - flattened just like `flat_data`.

99 + flattened just like `pixel_values`.

100 """

101

102 aspect_ratios: Annotated[torch.Tensor, TensorShape("batch_size", 2)]

@@ -725,6 +725,8 @@ class Mllama4DummyInputsBuilder(BaseDummyInputsBuilder[Mllama4ProcessingInfo]):

725 class Llama4ForConditionalGeneration(

726 nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3

727 ):

728 + merge_by_field_config = True

729 +

730 packed_modules_mapping = {

731 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

732 "gate_up_proj": ["gate_proj", "up_proj"],

@@ -798,17 +800,12 @@ class Llama4ForConditionalGeneration(

800 if pixel_values is None:

801 return None

802

801 - # num_images x num_chunks, channel, image_size, image_size

802 - # TODO: confirm handling for variable lengths

803 - flat_pixel_values = flatten_bn(pixel_values, concat=True)

804 - patches_per_image = flatten_bn(kwargs.pop("patches_per_image"))

803 + patches_per_image = kwargs.pop("patches_per_image")

804 aspect_ratios = kwargs.pop("aspect_ratios")

806 - if aspect_ratios.ndim == 3:

807 - aspect_ratios = aspect_ratios.squeeze(1)

805

806 return Llama4ImagePatchInputs(

807 type="pixel_values",

811 - flat_data=flat_pixel_values,

808 + pixel_values=pixel_values,

809 patches_per_image=patches_per_image,

810 aspect_ratios=aspect_ratios,

811 )

@@ -817,16 +814,16 @@ class Llama4ForConditionalGeneration(

814 self, image_input: Llama4ImagePatchInputs

815 ) -> MultiModalEmbeddings:

816 assert self.vision_model and self.multi_modal_projector

820 - flat_data = image_input["flat_data"]

817 + pixel_values = image_input["pixel_values"]

818 patches_per_image = image_input["patches_per_image"].tolist()

819

820 # shard image input

821 if self.use_data_parallel:

822 vision_embeddings_flat = run_dp_sharded_vision_model(

826 - flat_data, self.vision_model

823 + pixel_values, self.vision_model

824 )

825 else:

829 - vision_embeddings_flat = self.vision_model(flat_data)

826 + vision_embeddings_flat = self.vision_model(pixel_values)

827

828 vision_embeddings_flat = self.multi_modal_projector(vision_embeddings_flat)

829

vllm/model_executor/models/molmo.py

+18 -40

@@ -75,7 +75,6 @@ from .interfaces import (

75 from .utils import (

76 AutoWeightsLoader,

77 WeightsMapper,

78 - flatten_bn,

78 is_pp_missing_parameter,

79 make_empty_intermediate_tensors_factory,

80 make_layers,

@@ -97,28 +96,19 @@ class MolmoImageInputs(TensorSchema):

96 """

97 Dimensions:

98 - bn: Batch size * number of images

100 - - nc: Number of crops (dynamic)

99 + - bnc: Batch size * number of images * number of crops (dynamic)

100 - np: Number of patches

101 - tp: Token sequence positions

102 - pd: Patch dimension

103 """

104

106 - images: Annotated[

107 - torch.Tensor | list[torch.Tensor],

108 - TensorShape("bn", "nc", "np", "pd", dynamic_dims={"nc"}),

109 - ]

110 - # Number of crops may vary per batch and image, so pass it as a list.

105 + images: Annotated[torch.Tensor, TensorShape("bnc", "np", "pd")]

106

112 - image_masks: Annotated[

113 - torch.Tensor | list[torch.Tensor] | None,

114 - TensorShape("bn", "nc", "np", dynamic_dims={"nc"}),

115 - ]

107 + image_masks: Annotated[torch.Tensor | None, TensorShape("bnc", "np")]

108 +

109 + image_input_idx: Annotated[torch.Tensor, TensorShape("bnc", "tp")]

110 + """An index tensor that maps image features to their corresponding patch tokens."""

111

117 - image_input_idx: Annotated[

118 - torch.Tensor | list[torch.Tensor],

119 - TensorShape("bn", "nc", "tp", dynamic_dims={"nc"}),

120 - ]

121 - # An index tensor that maps image features to their corresponding patch tokens.

112 num_crops: Annotated[torch.Tensor, TensorShape("bn")]

113

114

@@ -1363,6 +1353,8 @@ class MolmoMultiModalProcessor(BaseMultiModalProcessor[MolmoProcessingInfo]):

1353 class MolmoForCausalLM(

1354 nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA, SupportsQuant

1355 ):

1356 + merge_by_field_config = True

1357 +

1358 hf_to_vllm_mapper = WeightsMapper(

1359 orig_to_new_substr={

1360 # vision backbone mapping

@@ -1451,18 +1443,12 @@ class MolmoForCausalLM(

1443 if images is None:

1444 return None

1445

1454 - if not isinstance(num_crops, (torch.Tensor, list)):

1455 - raise ValueError(

1456 - f"Incorrect type of num_crops. Got type: {type(num_crops)}"

1457 - )

1458 - num_crops = flatten_bn(num_crops, concat=True)

1459 -

1446 img_patch_id = kwargs.pop("img_patch_id", None)

1461 - if not isinstance(img_patch_id, torch.Tensor):

1462 - raise ValueError(

1463 - f"Incorrect type of img_patch_id. Got type: {type(img_patch_id)}"

1464 - )

1465 - self.img_patch_id = img_patch_id.flatten().unique().item()

1447 + if isinstance(img_patch_id, torch.Tensor):

1448 + img_patch_id = img_patch_id.item()

1449 +

1450 + assert isinstance(img_patch_id, int)

1451 + self.img_patch_id = img_patch_id

1452

1453 return MolmoImageInputs(

1454 images=images,

@@ -1481,17 +1467,9 @@ class MolmoForCausalLM(

1467 num_crops = image_input["num_crops"]

1468

1469 # Call the vision backbone on the whole batch at once

1484 - images_flat = flatten_bn(images, concat=True)

1485 - image_masks_flat = (

1486 - None if image_masks is None else flatten_bn(image_masks, concat=True)

1487 - )

1488 - image_input_idx_flat = flatten_bn(image_input_idx, concat=True)

1489 -

1490 - image_features_flat = self.vision_backbone(

1491 - images=images_flat.unsqueeze(0),

1492 - image_masks=(

1493 - None if image_masks_flat is None else image_masks_flat.unsqueeze(0)

1494 - ),

1470 + image_features = self.vision_backbone(

1471 + images=images.unsqueeze(0),

1472 + image_masks=None if image_masks is None else image_masks.unsqueeze(0),

1473 ).squeeze(0)

1474

1475 # Only the features corresponding to patch tokens are relevant

@@ -1499,8 +1477,8 @@ class MolmoForCausalLM(

1477 results = []

1478 num_crops_list = num_crops.tolist()

1479 for feats, img_idx in zip(

1502 - image_features_flat.split(num_crops_list),

1503 - image_input_idx_flat.split(num_crops_list),

1480 + image_features.split(num_crops_list),

1481 + image_input_idx.split(num_crops_list),

1482 ):

1483 is_valid = img_idx >= 0

1484 valid_img_idx = img_idx[is_valid]

vllm/model_executor/models/nano_nemotron_vl.py

+27 -113

@@ -11,7 +11,7 @@ import copy

11 import warnings

12 from abc import ABC, abstractmethod

13 from collections.abc import Iterable, Mapping, Sequence

14 -from typing import Annotated, Any, Literal, TypeAlias, TypedDict, TypeVar

14 +from typing import Annotated, Any, Literal, TypeAlias, TypeVar

16 import numpy.typing as npt

17 import torch

@@ -40,7 +40,6 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys

40 from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM

41 from vllm.model_executor.models.radio import RadioModel

42 from vllm.model_executor.models.utils import (

43 - flatten_bn,

43 init_vllm_registered_model,

44 maybe_prefix,

45 )

@@ -96,31 +95,35 @@ MAX_FRAMES = 16

95 DEFAULT_NUM_TILES = 12

99 -class NanoNemotronVLImagePixelInputs(TypedDict):

100 - type: Literal["pixel_values"]

101 - pixel_values_flat: torch.Tensor

98 +class NanoNemotronVLImagePixelInputs(TensorSchema):

99 """

103 - Shape:

104 - `(batch_size * num_images * (1 + num_patches), num_channels, height, width)`

100 + Dimensions:

101 + - bn: Batch size * number of images

102 + - bnp: Batch size * number of images * (1 + num_patches)

103 + - c: Number of channels (3)

104 + - h: Height of each image patch

105 + - w: Width of each image patch

106 """

107

107 - num_patches: torch.Tensor

108 - """Shape: `(batch_size * num_images)`"""

109 -

108 + type: Literal["pixel_values"]

109 + pixel_values_flat: Annotated[torch.Tensor, TensorShape("bnp", 3, "h", "w")]

110 + num_patches: Annotated[torch.Tensor, TensorShape("bn")]

111

111 -class NanoNemotronVLImageEmbeddinInputs(TypedDict):

112 - type: Literal["image_embeds"]

113 - data: torch.Tensor | list[torch.Tensor]

114 - """

115 - A tensor of shape `(num_images, total_image_feature_size, hidden_size)`

116 - or a list of tensors of shape `(total_image_feature_size, hidden_size)`

112

118 - `hidden_size` must match the hidden size of language model backbone.

113 +class NanoNemotronVLImageEmbeddingInputs(TensorSchema):

114 + """

115 + Dimensions:

116 + - n: Number of images

117 + - f: Total image feature size

118 + - h: Hidden size (must match the hidden size of language model backbone)

119 """

120

121 + type: Literal["image_embeds"]

122 + data: Annotated[torch.Tensor | list[torch.Tensor], TensorShape("n", "f", "h")]

123 +

124

125 NanoNemotronVLImageInputs: TypeAlias = (

123 - NanoNemotronVLImagePixelInputs | NanoNemotronVLImageEmbeddinInputs

126 + NanoNemotronVLImagePixelInputs | NanoNemotronVLImageEmbeddingInputs

127 )

128

129

@@ -710,37 +713,12 @@ class NanoNemotronVLProcessingInfo(BaseNanoNemotronVLProcessingInfo):

713 class NanoNemotronBaseVLMultiModalProcessor(BaseMultiModalProcessor[_I]):

714 """Basic image-only MultiModalProcessor for InternVL-style models."""

715

713 - def _call_hf_processor(

714 - self,

715 - prompt: str,

716 - mm_data: Mapping[str, object],

717 - mm_kwargs: Mapping[str, object],

718 - tok_kwargs: Mapping[str, object],

719 - ) -> BatchFeature:

720 - processed_outputs = super()._call_hf_processor(

721 - prompt=prompt,

722 - mm_data=mm_data,

723 - mm_kwargs=mm_kwargs,

724 - tok_kwargs=tok_kwargs,

725 - )

726 -

727 - hf_processor = self.info.get_hf_processor(**mm_kwargs)

728 - image_token_id = hf_processor.image_token_id

729 -

730 - # Since there may be extra tokens in the feature placeholders,

731 - # we need to pass the image token ID to the model to select the

732 - # tokens to merge from the vision encoder outputs

733 - processed_outputs["image_token_id"] = torch.tensor(image_token_id)

734 -

735 - return processed_outputs

736 -

716 def _get_mm_fields_config(

717 self,

718 hf_inputs: BatchFeature,

719 hf_processor_mm_kwargs: Mapping[str, object],

720 ) -> Mapping[str, MultiModalFieldConfig]:

721 image_num_patches = hf_inputs.get("image_num_patches", torch.empty(0))

743 - num_images = len(image_num_patches)

722

723 return dict(

724 pixel_values_flat=MultiModalFieldConfig.flat_from_sizes(

@@ -748,7 +726,6 @@ class NanoNemotronBaseVLMultiModalProcessor(BaseMultiModalProcessor[_I]):

726 ),

727 image_num_patches=MultiModalFieldConfig.batched("image"),

728 image_embeds=MultiModalFieldConfig.batched("image"),

751 - image_token_id=MultiModalFieldConfig.shared("image", num_images),

729 )

730

731 def _get_prompt_updates(

@@ -814,25 +791,6 @@ class NanoNemotronVLMultiModalProcessor(

791 ):

792 """MultiModalProcessor extended for video support"""

793

817 - def _call_hf_processor(

818 - self,

819 - prompt: str,

820 - mm_data: Mapping[str, object],

821 - mm_kwargs: Mapping[str, object],

822 - tok_kwargs: Mapping[str, object],

823 - ) -> BatchFeature:

824 - processed_outputs = super()._call_hf_processor(

825 - prompt, mm_data, mm_kwargs, tok_kwargs

826 - )

827 -

828 - hf_processor = self.info.get_hf_processor(**mm_kwargs)

829 - if (

830 - self.info.supports_video

831 - and (video_token_id := hf_processor.video_token_id) is not None

832 - ):

833 - processed_outputs["video_token_id"] = torch.tensor(video_token_id)

834 - return processed_outputs

835 -

794 def _get_mm_fields_config(

795 self,

796 hf_inputs: BatchFeature,

@@ -841,13 +799,12 @@ class NanoNemotronVLMultiModalProcessor(

799 image_fields = super()._get_mm_fields_config(hf_inputs, hf_processor_mm_kwargs)

800 if self.info.supports_video:

801 video_num_patches = hf_inputs.get("video_num_patches", torch.empty(0))

844 - num_videos = len(video_num_patches)

802 +

803 video_fields = dict(

804 pixel_values_flat_video=MultiModalFieldConfig.flat_from_sizes(

805 "video", video_num_patches

806 ),

807 video_num_patches=MultiModalFieldConfig.batched("video"),

850 - video_token_id=MultiModalFieldConfig.shared("video", num_videos),

808 )

809 else:

810 video_fields = {}

@@ -999,6 +956,8 @@ class NanoNemotronVLDummyInputsBuilder(

956 class NemotronH_Nano_VL_V2(

957 nn.Module, HasInnerState, IsHybrid, SupportsMultiModal, SupportsMultiModalPruning

958 ):

959 + merge_by_field_config = True

960 +

961 @classmethod

962 def get_placeholder_str(cls, modality: str, i: int) -> str | None:

963 if modality.startswith("image"):

@@ -1051,8 +1010,6 @@ class NemotronH_Nano_VL_V2(

1010 )

1011 self.mlp1 = self.mlp1.to(self.language_model.config.torch_dtype)

1012

1054 - self.img_context_token_id = None

1055 - self.video_context_token_id = None

1013 self.config = config

1014 self.model_config = vllm_config.model_config

1015

@@ -1106,37 +1063,12 @@ class NemotronH_Nano_VL_V2(

1063 return None

1064

1065 if image_embeds is not None:

1109 - if not isinstance(image_embeds, (torch.Tensor, list)):

1110 - raise ValueError(

1111 - "Incorrect type of image embeddings. "

1112 - f"Got type: {type(image_embeds)}"

1113 - )

1114 -

1115 - return NanoNemotronVLImageEmbeddinInputs(

1066 + return NanoNemotronVLImageEmbeddingInputs(

1067 type="image_embeds",

1117 - data=flatten_bn(image_embeds),

1068 + data=image_embeds,

1069 )

1070

1120 - image_token_id = kwargs["image_token_id"]

1121 - assert isinstance(image_token_id, torch.Tensor)

1122 - self.img_context_token_id = image_token_id.flatten().unique().item()

1123 -

1071 if pixel_values_flat is not None:

1125 - if not isinstance(pixel_values_flat, (torch.Tensor, list)):

1126 - raise ValueError(

1127 - "Incorrect type of pixel values. "

1128 - f"Got type: {type(pixel_values_flat)}"

1129 - )

1130 -

1131 - if not isinstance(image_num_patches, (torch.Tensor, list)):

1132 - raise ValueError(

1133 - "Incorrect type of image_num_patches. "

1134 - f"Got type: {type(image_num_patches)}"

1135 - )

1136 -

1137 - pixel_values_flat = flatten_bn(pixel_values_flat, concat=True)

1138 - image_num_patches = flatten_bn(image_num_patches, concat=True)

1139 -

1072 return NanoNemotronVLImagePixelInputs(

1073 type="pixel_values",

1074 pixel_values_flat=pixel_values_flat,

@@ -1285,28 +1217,10 @@ class NemotronH_Nano_VL_V2(

1217 if video_embeds is not None:

1218 return NanoNemotronVLVideoEmbeddingInputs(

1219 type="video_embeds",

1288 - data=flatten_bn(video_embeds),

1220 + data=video_embeds,

1221 )

1222

1291 - video_token_id = kwargs["video_token_id"]

1292 - assert isinstance(video_token_id, torch.Tensor)

1293 - self.video_context_token_id = video_token_id.flatten().unique().item()

1294 -

1223 if pixel_values_flat_video is not None:

1296 - if not isinstance(pixel_values_flat_video, (torch.Tensor, list)):

1297 - raise ValueError(

1298 - "Incorrect type of pixel values. "

1299 - f"Got type: {type(pixel_values_flat_video)}"

1300 - )

1301 -

1302 - if not isinstance(video_num_patches, (torch.Tensor, list)):

1303 - raise ValueError(

1304 - "Incorrect type of image_num_patches. "

1305 - f"Got type: {type(video_num_patches)}"

1306 - )

1307 -

1308 - pixel_values_flat_video = flatten_bn(pixel_values_flat_video, concat=True)

1309 - video_num_patches = flatten_bn(video_num_patches, concat=True)

1224 expected_h = expected_w = self.config.force_image_size

1225 resolve_bindings = {"h": expected_h, "w": expected_w}

1226

vllm/model_executor/models/nemotron_vl.py

+5 -2

@@ -496,8 +496,11 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor

496 )

497

498 image_token_id = kwargs["image_token_id"]

499 - assert isinstance(image_token_id, torch.Tensor)

500 - self.img_context_token_id = image_token_id.flatten().unique().item()

499 + if isinstance(image_token_id, torch.Tensor):

500 + image_token_id = image_token_id.flatten().unique().item()

501 +

502 + assert isinstance(image_token_id, int)

503 + self.img_context_token_id = image_token_id

504

505 if pixel_values_flat is not None:

506 return InternVLImagePixelInputs(

vllm/model_executor/models/skyworkr1v.py

+5 -2

@@ -814,8 +814,11 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):

814 )

815

816 image_token_id = kwargs["image_token_id"]

817 - assert isinstance(image_token_id, torch.Tensor)

818 - self.img_context_token_id = image_token_id.flatten().unique().item()

817 + if isinstance(image_token_id, torch.Tensor):

818 + image_token_id = image_token_id.flatten().unique().item()

819 +

820 + assert isinstance(image_token_id, int)

821 + self.img_context_token_id = image_token_id

822

823 if pixel_values_flat is not None:

824 return SkyworkR1VImagePixelInputs(

vllm/multimodal/utils.py

+1 -1

@@ -432,7 +432,7 @@ def group_mm_kwargs_by_modality(

432

433 if device is not None:

434 mm_kwargs_group = json_map_leaves(

435 - lambda x: x.to(device=device),

435 + lambda x: x.to(device=device) if isinstance(x, torch.Tensor) else x,

436 mm_kwargs_group,

437 )

438 else: