@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2b8a5495a8

[Mypy] Fix mypy typing for N/O models (#54142)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

Taneem Ibrahim committed Sep 18, 2026 at 06:50 UTC 2b8a5495a88a9a44a9d92e8c4f4aae3b55b6b74b

19 files changed +321 -143

tools/pre_commit/mypy.py

-2

@@ -113,8 +113,6 @@ SEPARATE_GROUPS = [

113

114 EXCLUDE = [

115 r"vllm/model_executor/models/[kK]",

116 - r"vllm/model_executor/models/[nN]",

117 - r"vllm/model_executor/models/[oO]",

116 r"vllm/model_executor/models/[qQ]",

117 r"vllm/model_executor/models/[uU]",

118 r"vllm/model_executor/models/[vV]",

vllm/model_executor/layers/mamba/mamba_utils.py

+8 -2

@@ -222,12 +222,18 @@ class MambaStateShapeCalculator:

222 @classmethod

223 def append_replayssm_ring(

224 cls,

225 - base_shapes: tuple[tuple[int, ...], ...],

225 + base_shapes: tuple[tuple[int, int], tuple[int, int, int]],

226 n_groups: int,

227 tp_world_size: int,

228 logical_window: int,

229 backend: MambaBackendEnum,

230 - ) -> tuple[tuple[int, ...], ...]:

230 + ) -> tuple[

231 + tuple[int, int],

232 + tuple[int, int, int],

233 + tuple[int, int, int],

234 + tuple[int, int],

235 + tuple[int, int, int],

236 + ]:

237 """Append the physical ReplaySSM ring shapes.

238

239 ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``;

vllm/model_executor/models/aimv2.py

+14 -5

@@ -26,7 +26,10 @@ from vllm.transformers_utils.configs.ovis import AIMv2Config

27 class AIMv2SwiGLUFFN(nn.Module):

28 def __init__(

29 - self, config: AIMv2Config, quant_config: QuantizationConfig, prefix: str

29 + self,

30 + config: AIMv2Config,

31 + quant_config: QuantizationConfig | None,

32 + prefix: str,

33 ):

34 super().__init__()

35 hidden_features = config.intermediate_size

@@ -91,7 +94,10 @@ class AIMv2ViTPreprocessor(nn.Module):

95 class AIMv2Attention(nn.Module):

96 def __init__(

94 - self, config: AIMv2Config, quant_config: QuantizationConfig, prefix: str

97 + self,

98 + config: AIMv2Config,

99 + quant_config: QuantizationConfig | None,

100 + prefix: str,

101 ):

102 super().__init__()

103 self.config = config

@@ -144,7 +150,10 @@ class AIMv2Attention(nn.Module):

150

151 class AIMv2Block(nn.Module):

152 def __init__(

147 - self, config: AIMv2Config, quant_config: QuantizationConfig, prefix: str

153 + self,

154 + config: AIMv2Config,

155 + quant_config: QuantizationConfig | None,

156 + prefix: str,

157 ):

158 super().__init__()

159 self.attn = AIMv2Attention(

@@ -166,7 +175,7 @@ class AIMv2Transformer(nn.Module):

175 def __init__(

176 self,

177 config: AIMv2Config,

169 - quant_config: QuantizationConfig,

178 + quant_config: QuantizationConfig | None,

179 *,

180 require_post_norm: bool | None = None,

181 prefix: str = "",

@@ -206,7 +215,7 @@ class AIMv2Model(torch.nn.Module):

215 def __init__(

216 self,

217 config: AIMv2Config,

209 - quant_config: QuantizationConfig,

218 + quant_config: QuantizationConfig | None,

219 *,

220 require_post_norm: bool | None = None,

221 prefix: str = "",

vllm/model_executor/models/interfaces.py

@@ -91,6 +91,13 @@ MambaStateShapes: TypeAlias = (

91 tuple[int, int, int],

92 tuple[int, int, int],

93 ]

94 + | tuple[

95 + tuple[int, int],

96 + tuple[int, int, int],

97 + tuple[int, int, int],

98 + tuple[int, int],

99 + tuple[int, int, int],

100 + ]

101 )

102

103

vllm/model_executor/models/nano_nemotron_vl.py

+85 -27

@@ -12,7 +12,7 @@ import warnings

12 from collections.abc import Iterable, Mapping, Sequence

13 from functools import cached_property

14 from io import BytesIO

15 -from typing import Annotated, Literal, TypeAlias

15 +from typing import Annotated, Literal, TypeAlias, TypedDict, cast

17 import torch

18 import torch.nn as nn

@@ -20,7 +20,11 @@ from transformers import BatchFeature, PretrainedConfig

21 from vllm import envs

22 from vllm.config import VllmConfig

23 -from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions

23 +from vllm.config.multimodal import (

24 + AudioDummyOptions,

25 + BaseDummyOptions,

26 + VideoDummyOptions,

27 +)

28 from vllm.inputs import MultiModalDataDict, MultiModalInput

29 from vllm.logger import init_logger

30 from vllm.lora.layers.base import BaseLayerWithLoRA

@@ -114,7 +118,7 @@ class NanoNemotronVLAudioFeatureInputs(TensorSchema):

118 type: Literal["audio_features"] = "audio_features"

119 input_audio_features: Annotated[torch.Tensor, TensorShape("c", "t", "f")]

120 feature_attention_mask: Annotated[torch.Tensor, TensorShape("c", "t")]

117 - audio_num_clips: list[int]

121 + audio_num_clips: list[int] | torch.Tensor

122

123

124 class NanoNemotronVLImagePixelInputs(TensorSchema):

@@ -195,6 +199,12 @@ NanoNemotronVLVideoInputs: TypeAlias = (

199 )

200

201

202 +class NanoNemotronVLMultiModalInputs(TypedDict, total=False):

203 + images: NanoNemotronVLImageInputs | None

204 + videos: NanoNemotronVLVideoInputs | None

205 + audios: NanoNemotronVLAudioFeatureInputs | None

206 +

207 +

208 class NanoNemotronVLProcessingInfo(BaseProcessingInfo):

209 def get_hf_processor(self, **kwargs: object) -> NanoNemotronVLProcessor:

210 # Not accepted by NanoNemotronVLProcessor.__init__

@@ -465,8 +475,11 @@ class NanoNemotronVLMultiModalProcessor(

475 if isinstance(images, ImageEmbeddingItems):

476 feature_size = images.get_feature_size(item_idx)

477 elif self.info.is_dynamic_tiler:

468 - feature_size = out_mm_data["num_tokens_per_image"][item_idx]

478 + num_tokens_per_image = out_mm_data["num_tokens_per_image"]

479 + assert isinstance(num_tokens_per_image, list)

480 + feature_size = num_tokens_per_image[item_idx]

481 else:

482 + assert isinstance(images, ImageProcessorItems)

483 image_size = images.get_image_size(item_idx)

484 max_num_tiles = hf_processor.max_num_tiles

485 feature_size = hf_processor.get_num_image_tokens(

@@ -530,10 +543,12 @@ class NanoNemotronVLMultiModalProcessor(

543 assert isinstance(num_patches, int)

544

545 T = hf_processor.video_temporal_patch_size

546 + num_tubelets: int | None

547 if T > 1 and num_patches is not None:

548 num_tubelets = math.ceil(num_patches / T)

549 else:

550 num_tubelets = num_patches

551 + assert num_tubelets is not None

552

553 video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate

554 if video_pruning_rate is not None and video_pruning_rate > 0.0:

@@ -635,6 +650,7 @@ class NanoNemotronVLMultiModalProcessor(

650 audio_items: list[AudioItem] = []

651 has_audio: list[bool] = []

652 for idx, metadata in enumerate(metadata_list):

653 + assert metadata is not None

654 video_bytes = metadata.get("original_video_bytes")

655 if video_bytes is None or len(video_bytes) == 0:

656 raise ValueError(

@@ -661,21 +677,25 @@ class NanoNemotronVLMultiModalProcessor(

677

678 # Create a new VideoProcessorItems with metadata that does not contain

679 # the large video bytes, to avoid modifying the input `mm_items`.

664 - new_metadata_list = [

665 - {k: v for k, v in meta.items() if k != "original_video_bytes"}

666 - for meta in metadata_list

667 - ]

680 + new_metadata_list = metadata_list.copy()

681 + for idx, metadata in enumerate(metadata_list):

682 + assert metadata is not None

683 + new_metadata_list[idx] = {

684 + key: value

685 + for key, value in metadata.items()

686 + if key != "original_video_bytes"

687 + }

688 new_videos = VideoProcessorItems(data=videos.data, metadata=new_metadata_list)

689

670 - audio_parsed = {}

690 + audio_parsed = MultiModalDataItems()

691 if audio_items:

692 audio_parsed = self.data_parser.parse_mm_data({"audio": audio_items})

693

694 # Create a new MultiModalDataItems with the new video and audio items.

695 new_mm_items_dict = {**mm_items, **audio_parsed, "video": new_videos}

676 - mm_items = MultiModalDataItems(new_mm_items_dict)

696 + new_mm_items = MultiModalDataItems(new_mm_items_dict)

697

678 - return mm_items, audio_items, has_audio

698 + return new_mm_items, audio_items, has_audio

699

700 def apply(

701 self,

@@ -800,7 +820,7 @@ class NanoNemotronVLDummyInputsBuilder(

820 )

821 videos = [v.copy() for v in videos]

822

803 - video_items = []

823 + video_items: list[VideoItem] = []

824 for video in videos:

825 video_num_frames = video.shape[0]

826 video_metadata = {

@@ -869,6 +889,9 @@ class NanoNemotronVLDummyInputsBuilder(

889 target_num_frames = num_frames

890 num_videos = mm_counts.get("video", 0)

891 video_overrides = mm_options.get("video")

892 + assert video_overrides is None or isinstance(

893 + video_overrides, VideoDummyOptions

894 + )

895 dummy_video = {

896 "video": self._get_dummy_videos(

897 width=video_width,

@@ -884,6 +907,9 @@ class NanoNemotronVLDummyInputsBuilder(

907 if sound_config := self.info.sound_config:

908 num_audios = mm_counts.get("audio", 0)

909 audio_overrides = mm_options.get("audio") if mm_options else None

910 + assert audio_overrides is None or isinstance(

911 + audio_overrides, AudioDummyOptions

912 + )

913 tokens_per_audio = max(1, seq_len // max(num_audios, 1))

914 max_audio_num_samples = MAX_AUDIO_LEN_S * sound_config.sampling_rate

915 calculated_max_audio_num_samples = ParakeetExtractor.audio_length(

@@ -946,6 +972,7 @@ class NemotronH_Nano_VL_V2(

972 model_config = vllm_config.model_config

973 config = model_config.hf_config

974 multimodal_config = model_config.multimodal_config

975 + assert multimodal_config is not None

976 image_size = config.force_image_size

977 patch_size = config.patch_size

978 self.patch_size = patch_size

@@ -1138,17 +1165,26 @@ class NemotronH_Nano_VL_V2(

1165 return None

1166

1167 if self.dynamic_resolution:

1168 + assert isinstance(pixel_values_flat, (list, torch.Tensor))

1169 pixel_values_flat = DynamicResolutionImageTiler.stack(

1170 pixel_values_flat, self.patch_size

1171 )

1172 + imgs_sizes = kwargs.pop("imgs_sizes")

1173 + num_tokens_per_image = kwargs.pop("num_tokens_per_image")

1174 + assert isinstance(imgs_sizes, list)

1175 + assert isinstance(num_tokens_per_image, list)

1176 return NanoNemotronVLImagePixelInputsDynamic(

1145 - pixel_values_flat=pixel_values_flat, **kwargs

1177 + pixel_values_flat=pixel_values_flat,

1178 + imgs_sizes=imgs_sizes,

1179 + num_tokens_per_image=num_tokens_per_image,

1180 )

1181 else:

1182 + assert isinstance(pixel_values_flat, torch.Tensor)

1183 + image_num_patches = kwargs.pop("image_num_patches")

1184 + assert isinstance(image_num_patches, torch.Tensor)

1185 return NanoNemotronVLImagePixelInputs(

1186 pixel_values_flat=pixel_values_flat,

1150 - num_patches=kwargs.pop("image_num_patches"),

1151 - **kwargs,

1187 + num_patches=image_num_patches,

1188 )

1189

1190 def _process_image_input_dynamic(

@@ -1166,7 +1202,8 @@ class NemotronH_Nano_VL_V2(

1202 return image_embeds.split(num_tokens_per_image)

1203

1204 def _process_image_input(

1169 - self, image_input: NanoNemotronVLImagePixelInputs

1205 + self,

1206 + image_input: NanoNemotronVLImagePixelInputs | NanoNemotronVLVideoPixelInputs,

1207 ) -> tuple[torch.Tensor, ...]:

1208 image_embeds = self.extract_feature(image_input["pixel_values_flat"])

1209 num_patches = image_input["num_patches"]

@@ -1362,7 +1399,8 @@ class NemotronH_Nano_VL_V2(

1399 # of them than max_num_batched_tokens.

1400 # Embed them with the base weights -

1401 # the LoRA delta is undefined for tokens with no mapping entry.

1365 - embed_tokens = self.get_language_model().model.embed_tokens

1402 + language_model = cast(NemotronHForCausalLM, self.get_language_model())

1403 + embed_tokens = language_model.model.embed_tokens

1404 if isinstance(embed_tokens, BaseLayerWithLoRA):

1405 embed_tokens = embed_tokens.base_layer

1406 text_embeddings = embed_tokens(repl_token_ids)

@@ -1388,26 +1426,29 @@ class NemotronH_Nano_VL_V2(

1426 return None

1427

1428 if video_embeds is not None:

1429 + assert isinstance(video_embeds, (torch.Tensor, list))

1430 return NanoNemotronVLVideoEmbeddingInputs(

1431 type="video_embeds",

1432 data=video_embeds,

1433 )

1434

1435 if pixel_values_flat_video is not None:

1397 - if torch.is_tensor(frames_indices):

1436 + if isinstance(frames_indices, torch.Tensor):

1437 frames_indices = frames_indices.flatten()

1438 else:

1439 + assert isinstance(frames_indices, (list, tuple))

1440 frames_indices = torch.cat([f.flatten() for f in frames_indices], dim=0)

1441

1402 - if torch.is_tensor(frame_duration_ms):

1442 + if isinstance(frame_duration_ms, torch.Tensor):

1443 frame_duration_ms = frame_duration_ms.flatten()

1444 else:

1445 + assert isinstance(frame_duration_ms, (list, tuple))

1446 frame_duration_ms = torch.cat(

1447 [f.flatten() for f in frame_duration_ms], dim=0

1448 )

1449

1450 if (

1410 - torch.is_tensor(pixel_values_flat_video)

1451 + isinstance(pixel_values_flat_video, torch.Tensor)

1452 and pixel_values_flat_video.ndim == 5

1453 ):

1454 # batched._reduce_data stacked same-shape videos into

@@ -1415,9 +1456,11 @@ class NemotronH_Nano_VL_V2(

1456 # same-H,W cat path below handles it uniformly.

1457 pixel_values_flat_video = list(pixel_values_flat_video)

1458

1418 - if not torch.is_tensor(pixel_values_flat_video):

1459 + if not isinstance(pixel_values_flat_video, torch.Tensor):

1460 + assert isinstance(pixel_values_flat_video, (list, tuple))

1461 pixel_values_flat_video = torch.cat(pixel_values_flat_video, dim=0)

1462

1463 + assert isinstance(video_num_patches, torch.Tensor)

1464 expected_h = pixel_values_flat_video.shape[-2]

1465 expected_w = pixel_values_flat_video.shape[-1]

1466 num_frames = video_num_patches[0].item()

@@ -1434,8 +1477,10 @@ class NemotronH_Nano_VL_V2(

1477

1478 raise AssertionError("This line should be unreachable.")

1479

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

1438 - modalities = {}

1480 + def _parse_and_validate_multimodal_inputs(

1481 + self, **kwargs: object

1482 + ) -> NanoNemotronVLMultiModalInputs:

1483 + modalities: NanoNemotronVLMultiModalInputs = {}

1484 # Preserve the order of modalities if there are multiple of them

1485 # from the order of kwargs.

1486 for input_key in kwargs:

@@ -1458,8 +1503,17 @@ class NemotronH_Nano_VL_V2(

1503 )

1504 and "audios" not in modalities

1505 ):

1506 + input_audio_features = kwargs.get("input_audio_features")

1507 + feature_attention_mask = kwargs.get("feature_attention_mask")

1508 + audio_num_clips = kwargs.get("audio_num_clips")

1509 + assert isinstance(input_audio_features, torch.Tensor)

1510 + assert isinstance(feature_attention_mask, torch.Tensor)

1511 + assert isinstance(audio_num_clips, (list, torch.Tensor))

1512 modalities["audios"] = NanoNemotronVLAudioFeatureInputs(

1462 - **kwargs, validate=False

1513 + input_audio_features=input_audio_features,

1514 + feature_attention_mask=feature_attention_mask,

1515 + audio_num_clips=audio_num_clips,

1516 + validate=False,

1517 )

1518

1519 return modalities

@@ -1479,23 +1533,26 @@ class NemotronH_Nano_VL_V2(

1533 for modality in modalities:

1534 if modality == "images":

1535 image_input = modalities["images"]

1482 - if image_input["type"] == "image_embeds":

1536 + assert image_input is not None

1537 + if isinstance(image_input, NanoNemotronVLImageEmbeddingInputs):

1538 image_embeddings = image_input["data"]

1484 - elif self.dynamic_resolution:

1485 - assert image_input["type"] == "pixel_values_dynamic"

1539 + elif isinstance(image_input, NanoNemotronVLImagePixelInputsDynamic):

1540 image_embeddings = self._process_image_input_dynamic(image_input)

1541 else:

1542 image_embeddings = self._process_image_input(image_input)

1543 multimodal_embeddings += tuple(image_embeddings)

1544 if modality == "videos":

1545 video_input = modalities["videos"]

1546 + assert video_input is not None

1547 if video_input["type"] == "video_embeds":

1548 video_embeddings = video_input["data"]

1549 else:

1550 + assert isinstance(video_input, NanoNemotronVLVideoPixelInputs)

1551 video_embeddings = self._process_video_input(video_input)

1552 multimodal_embeddings += tuple(video_embeddings)

1553 if modality == "audios":

1554 audio_input = modalities["audios"]

1555 + assert audio_input is not None

1556 audio_embeddings = self._process_audio_input(audio_input)

1557 multimodal_embeddings += tuple(audio_embeddings)

1558

@@ -1538,6 +1595,7 @@ class NemotronH_Nano_VL_V2(

1595

1596 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

1597 mm_config = self.model_config.multimodal_config

1598 + assert mm_config is not None

1599 load_multimodal_weights = not all(

1600 mm_config.get_limit_per_prompt(modality) == 0

1601 for modality in ("image", "video", "audio")

vllm/model_executor/models/nemotron_h.py

+7 -2

@@ -68,6 +68,7 @@ from vllm.model_executor.models.interfaces import (

68 EagleModelMixin,

69 HasInnerState,

70 IsHybrid,

71 + MambaStateShapes,

72 MixtureOfExperts,

73 SupportsEagle,

74 SupportsEagle3,

@@ -138,6 +139,7 @@ class NemotronHMoE(nn.Module):

139 prefix: str = "",

140 ):

141 super().__init__()

142 + assert parallel_config is not None

143 self.tp_size = get_tensor_model_parallel_world_size()

144 self.routed_scaling_factor = config.routed_scaling_factor

145

@@ -188,6 +190,8 @@ class NemotronHMoE(nn.Module):

190 prefix=f"{prefix}.shared_experts",

191 )

192

193 + self.fc1_latent_proj: ReplicatedLinear | None

194 + self.fc2_latent_proj: ReplicatedLinear | None

195 if self.use_latent_moe:

196 self.fc1_latent_proj = ReplicatedLinear(

197 input_size=config.hidden_size,

@@ -722,7 +726,7 @@ class NemotronHForCausalLM(

726 SupportsReplaySSM,

727 ):

728 # Relevant only if self.has_moe is True

725 - is_non_gated_moe: bool = True

729 + is_non_gated_moe = True

730

731 hf_to_vllm_mapper = WeightsMapper(

732 orig_to_new_prefix={"backbone": "model", "mtp": None},

@@ -773,7 +777,7 @@ class NemotronHForCausalLM(

777 def get_mamba_state_shape_from_config(

778 cls,

779 vllm_config: "VllmConfig",

776 - ) -> tuple[tuple[int, ...], ...]:

780 + ) -> MambaStateShapes:

781 """Calculate shapes for Mamba's convolutional and state caches.

782

783 Args:

@@ -858,6 +862,7 @@ class NemotronHForCausalLM(

862 self.moe_layers.append(layer.mixer.experts)

863

864 self.num_moe_layers = len(self.moe_layers)

865 + assert example_moe is not None

866 self.num_logical_experts = example_moe.n_logical_experts

867 self.num_physical_experts = example_moe.n_physical_experts

868 self.num_local_physical_experts = example_moe.n_local_physical_experts # noqa: E501

vllm/model_executor/models/nemotron_h_mtp.py

+16 -8

@@ -92,10 +92,12 @@ class NemotronHMTPAttentionDecoderLayer(NemotronHAttentionDecoderLayer):

93 def forward(

94 self,

95 - inputs_embeds: torch.Tensor,

95 positions: torch.Tensor,

96 hidden_states: torch.Tensor,

97 residual: torch.Tensor | None = None,

98 + *,

99 + inputs_embeds: torch.Tensor | None = None,

100 + **kwargs: object,

101 ) -> tuple[torch.Tensor, torch.Tensor | None]:

102 # Start projections (Fusion)

103 if self.has_start_projections:

@@ -179,10 +181,11 @@ class NemotronHMTPMoEDecoderLayer(NemotronHMoEDecoderLayer):

181

182 def forward(

183 self,

182 - inputs_embeds: torch.Tensor,

183 - positions: torch.Tensor,

184 hidden_states: torch.Tensor,

185 residual: torch.Tensor | None = None,

186 + *,

187 + inputs_embeds: torch.Tensor | None = None,

188 + **kwargs: object,

189 ) -> tuple[torch.Tensor, torch.Tensor | None]:

190 # Start projections (Fusion)

191 if self.has_start_projections:

@@ -305,13 +308,14 @@ class NemotronHMultiTokenPredictor(nn.Module):

308

309 def forward(

310 self,

308 - input_ids: torch.Tensor,

311 + input_ids: torch.Tensor | None,

312 positions: torch.Tensor,

313 hidden_states: torch.Tensor,

314 intermediate_tensors: IntermediateTensors | None = None,

315 inputs_embeds: torch.Tensor | None = None,

316 ) -> torch.Tensor | IntermediateTensors:

317 if inputs_embeds is None:

318 + assert input_ids is not None

319 inputs_embeds = self.get_input_embeddings(input_ids)

320

321 residual = None

@@ -386,6 +390,7 @@ class NemotronHMTP(nn.Module, SupportsPP, SupportsQuant):

390 @staticmethod

391 def _find_quant_config(*args, **kwargs) -> QuantizationConfig | None:

392 vllm_config = kwargs.get("vllm_config")

393 + assert isinstance(vllm_config, VllmConfig)

394 return get_draft_quant_config(vllm_config)

395

396 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

@@ -393,14 +398,15 @@ class NemotronHMTP(nn.Module, SupportsPP, SupportsQuant):

398

399 def forward(

400 self,

396 - input_ids: torch.Tensor,

401 + input_ids: torch.Tensor | None,

402 positions: torch.Tensor,

398 - hidden_states: torch.Tensor,

403 + hidden_states: torch.Tensor | None = None,

404 intermediate_tensors: IntermediateTensors | None = None,

405 inputs_embeds: torch.Tensor | None = None,

406 **kwargs: object,

407 ) -> torch.Tensor:

408 """Forward - applies attention-based MTP."""

409 + assert hidden_states is not None

410 hidden_states = self.model(

411 input_ids,

412 positions,

@@ -473,9 +479,10 @@ class NemotronHMTP(nn.Module, SupportsPP, SupportsQuant):

479 name = name.replace("backbone.", "model.")

480

481 if "scale" in name or "zero_point" in name:

476 - name = maybe_remap_kv_scale_name(name, params_dict)

477 - if name is None:

482 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

483 + if remapped_name is None:

484 continue

485 + name = remapped_name

486

487 # Handle stacked parameters (qkv_proj) for attention layers

488 is_stacked = False

@@ -547,6 +554,7 @@ class NemotronHMTP(nn.Module, SupportsPP, SupportsQuant):

554

555 param = params_dict[name]

556 weight_loader = getattr(param, "weight_loader", default_weight_loader)

557 + assert weight_loader is not None

558 weight_loader(param, loaded_weight)

559 loaded_params.add(name)

560

vllm/model_executor/models/nemotron_vl.py

+14 -5

@@ -138,7 +138,7 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor

138 prefix=maybe_prefix(prefix, "language_model"),

139 )

140

141 - self.img_context_token_id = None

141 + self.img_context_token_id: int | None = None

142

143 self.visual_token_mask = None

144 self.make_empty_intermediate_tensors = (

@@ -146,7 +146,7 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor

146 )

147

148 def _patch_quant_config(

149 - self, config: PretrainedConfig, quant_config: QuantizationConfig

149 + self, config: PretrainedConfig, quant_config: QuantizationConfig | None

150 ):

151 # the awq models from OpenGVLab missing `modules_to_not_convert`

152 # patch the quant_config to add `modules_to_not_convert` back

@@ -518,12 +518,21 @@ class LlamaNemotronVLForEmbedding(LlamaNemotronVLChatModel, VllmModelForPooling)

518 use_head=False,

519 )

520

521 - def _init_mlp1(self, config: PretrainedConfig) -> nn.Module:

521 + def _init_mlp1(

522 + self,

523 + config: PretrainedConfig,

524 + vit_hidden_size: int | None = None,

525 + vision_projection_hidden_size: int | None = None,

526 + ) -> nn.Module:

527 """Override to use different MLP structure for embedding model."""

528 + if vit_hidden_size is None:

529 + vit_hidden_size = config.vision_config.hidden_size

530 + if vision_projection_hidden_size is None:

531 + vision_projection_hidden_size = config.get_text_config().hidden_size

532 return super()._init_mlp1(

533 config,

525 - vit_hidden_size=config.vision_config.hidden_size,

526 - vision_projection_hidden_size=config.get_text_config().hidden_size,

534 + vit_hidden_size=vit_hidden_size,

535 + vision_projection_hidden_size=vision_projection_hidden_size,

536 )

537

538 def _call_vision_model(self, pixel_values: torch.Tensor) -> torch.Tensor:

vllm/model_executor/models/nvlm_d.py

@@ -128,6 +128,7 @@ class NVLMMultiModalProcessor(BaseInternVLMultiModalProcessor[NVLMProcessingInfo

128 if isinstance(images, ImageEmbeddingItems):

129 feature_size = images.get_feature_size(item_idx)

130 else:

131 + assert isinstance(images, ImageProcessorItems)

132 image_size = images.get_image_size(item_idx)

133 feature_size = self.info.get_num_image_tokens(

134 image_width=image_size.width,

vllm/model_executor/models/olmo_hybrid.py

@@ -143,6 +143,7 @@ class OlmoHybridAttention(nn.Module):

143 rope_parameters["rope_theta"] is not None

144 )

145

146 + self.rotary_emb: nn.Module | None

147 if self._use_rope:

148 self.rotary_emb = get_rope(

149 self.head_dim,

@@ -183,6 +184,7 @@ class OlmoHybridAttention(nn.Module):

184 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

185 q, k = self._apply_qk_norm(q, k)

186 if self._use_rope:

187 + assert self.rotary_emb is not None

188 q, k = self.rotary_emb(positions, q, k)

189 attn_output = self.attn(q, k, v)

190 output, _ = self.o_proj(attn_output)

vllm/model_executor/models/opencua.py

+10 -2

@@ -7,6 +7,7 @@

8 """Inference-only OpenCUA-7B model compatible with HuggingFace weights."""

10 +import typing

11 from collections.abc import Mapping, Sequence

12 from typing import Any

@@ -27,6 +28,7 @@ from vllm.multimodal.inputs import (

28 )

29 from vllm.multimodal.parse import MultiModalDataItems

30 from vllm.multimodal.processing import (

31 + BaseDummyInputsBuilder,

32 BaseMultiModalProcessor,

33 PromptReplacement,

34 PromptUpdate,

@@ -178,7 +180,13 @@ class OpenCUAMultiModalProcessor(BaseMultiModalProcessor[OpenCUAProcessingInfo])

180 ]

181

182

181 -class OpenCUADummyInputsBuilder(Qwen2VLDummyInputsBuilder):

183 +if typing.TYPE_CHECKING:

184 + _OpenCUADummyInputsBuilderBase = BaseDummyInputsBuilder[OpenCUAProcessingInfo]

185 +else:

186 + _OpenCUADummyInputsBuilderBase = Qwen2VLDummyInputsBuilder

187 +

188 +

189 +class OpenCUADummyInputsBuilder(_OpenCUADummyInputsBuilderBase):

190 def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

191 num_images = mm_counts.get("image", 0)

192

@@ -220,7 +228,7 @@ class OpenCUAForConditionalGeneration(Qwen2_5_VLForConditionalGeneration):

228 nn.Module.__init__(self)

229 config = vllm_config.model_config.hf_config

230 quant_config = vllm_config.quant_config

223 - multimodal_config = vllm_config.model_config.multimodal_config

231 + multimodal_config = vllm_config.model_config.get_multimodal_config()

232

233 self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

234 self.config = config

vllm/model_executor/models/openpangu.py

+5 -1

@@ -175,6 +175,7 @@ class OpenPanguMoE(nn.Module):

175 self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts

176 self.n_local_physical_experts = self.n_physical_experts // self.ep_size

177

178 + self.shared_experts: OpenPanguMLP | None

179 if config.n_shared_experts is not None:

180 intermediate_size = config.moe_intermediate_size * config.n_shared_experts

181 self.shared_experts = OpenPanguMLP(

@@ -283,7 +284,7 @@ class OpenPanguMLAAttention(nn.Module):

284 )

285 self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)

286 self.q_b_proj = ColumnParallelLinear(

286 - q_lora_rank,

287 + self.q_lora_rank,

288 self.num_heads * self.qk_head_dim,

289 bias=False,

290 quant_config=quant_config,

@@ -568,6 +569,9 @@ class OpenPanguSinkAttention(nn.Module):

569 self.qk_nope_dim = getattr(config, "qk_nope_dim", None)

570 self.qk_rope_dim = getattr(config, "qk_rope_dim", None)

571 self.v_channels = getattr(config, "v_channels", None)

572 + assert self.qk_nope_dim is not None

573 + assert self.qk_rope_dim is not None

574 + assert self.v_channels is not None

575 self.head_dim = self.qk_rope_dim + self.qk_nope_dim

576 self.q_size = self.num_heads * self.head_dim

577 self.k_size = self.num_kv_heads * self.head_dim

vllm/model_executor/models/openpangu_mtp.py

+12 -5

@@ -52,7 +52,9 @@ class OpenPanguMultiTokenPredictorLayer(DeepSeekMultiTokenPredictorLayer):

52 def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:

53 nn.Module.__init__(self)

55 - config = vllm_config.speculative_config.draft_model_config.hf_config

55 + speculative_config = vllm_config.speculative_config

56 + assert speculative_config is not None

57 + config = speculative_config.draft_model_config.hf_config

58 self.config = config

59 quant_config = vllm_config.quant_config

@@ -200,10 +202,15 @@ class OpenPanguMTP(nn.Module):

202 break

203 else:

204 for mapping in expert_params_mapping:

203 - param_name, weight_name, expert_id, shard_id = mapping

204 - if weight_name not in name:

205 + (

206 + expert_param_name,

207 + expert_weight_name,

208 + expert_id,

209 + expert_shard_id,

210 + ) = mapping

211 + if expert_weight_name not in name:

212 continue

206 - name = name.replace(weight_name, param_name)

213 + name = name.replace(expert_weight_name, expert_param_name)

214

215 param = params_dict[name]

216 weight_loader = param.weight_loader

@@ -211,7 +218,7 @@ class OpenPanguMTP(nn.Module):

218 param,

219 loaded_weight,

220 name,

214 - shard_id=shard_id,

221 + shard_id=expert_shard_id,

222 expert_id=expert_id,

223 )

224 break

vllm/model_executor/models/openpangu_vl.py

+101 -74

@@ -22,7 +22,7 @@

23 from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence

24 from functools import lru_cache, partial

25 -from typing import Annotated, Literal, Optional

25 +from typing import Annotated, Literal, Optional, TypeAlias

27 import torch

28 import torch.nn as nn

@@ -627,7 +627,7 @@ class OpenPanguVisionTransformer(nn.Module):

627 pad[-(dim + 1) * 2 + 1] = pad_len

628 return F.pad(w, pad, mode="constant", value=0)

629

630 - stacked_params_mapping = [

630 + stacked_params_mapping: list[tuple[str, str, str | int]] = [

631 ("attn.qkv.", "attn.q.", "q"),

632 ("attn.qkv.", "attn.k.", "k"),

633 ("attn.qkv.", "attn.v.", "v"),

@@ -677,15 +677,11 @@ class OpenPanguVLProcessingInfo(Qwen2_5_VLProcessingInfo):

677 def get_hf_config(self):

678 return self.ctx.model_config.hf_config

679

680 - def get_hf_processor(

681 - self,

682 - *,

683 - min_pixels: int | None = None,

684 - max_pixels: int | None = None,

685 - size: dict[str, int] | None = None,

686 - fps: float | list[float] | None = None,

687 - **kwargs: object,

688 - ):

680 + def get_hf_processor(self, **kwargs: object):

681 + kwargs.pop("min_pixels", None)

682 + kwargs.pop("max_pixels", None)

683 + kwargs.pop("size", None)

684 + fps = kwargs.pop("fps", None)

685 if fps is not None:

686 kwargs["fps"] = fps

687

@@ -747,11 +743,19 @@ class OpenPanguVLVideoEmbeddingInputs(TensorSchema):

743 ]

744

745

746 +OpenPanguVLImageInputs: TypeAlias = (

747 + OpenPanguVLImagePixelInputs | OpenPanguVLImageEmbeddingInputs

748 +)

749 +OpenPanguVLVideoInputs: TypeAlias = (

750 + OpenPanguVLVideoPixelInputs | OpenPanguVLVideoEmbeddingInputs

751 +)

752 +

753 +

754 class OpenPanguVLMultiModalProcessor(Qwen2_5_VLMultiModalProcessor):

755 def _get_prompt_updates(

756 self,

757 mm_items: MultiModalDataItems,

754 - hf_processor_mm_kwargs: Mapping[str, any],

758 + hf_processor_mm_kwargs: Mapping[str, object],

759 out_mm_kwargs: MultiModalKwargsItems,

760 ) -> Sequence[PromptUpdate]:

761 hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

@@ -878,10 +882,12 @@ class OpenPanguVLForConditionalGeneration(

882 self.do_rescale = image_processor.do_rescale

883 self.rescale_factor = image_processor.rescale_factor

884 self.do_normalize = image_processor.do_normalize

881 - self.image_mean = tuple(image_processor.image_mean)

882 - self.image_std = tuple(image_processor.image_std)

885 + self.image_mean: tuple[float, ...] = tuple(image_processor.image_mean)

886 + self.image_std: tuple[float, ...] = tuple(image_processor.image_std)

887

884 - def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig):

888 + def _maybe_ignore_quant_config(

889 + self, quant_config: QuantizationConfig | None

890 + ) -> QuantizationConfig | None:

891 if isinstance(quant_config, AutoGPTQConfig):

892 return None

893 return quant_config

@@ -904,7 +910,9 @@ class OpenPanguVLForConditionalGeneration(

910 else:

911 return torch.concat(mm_input)

912

907 - def _parse_and_validate_image_input(self, **kwargs: object):

913 + def _parse_and_validate_image_input(

914 + self, **kwargs: object

915 + ) -> OpenPanguVLImageInputs | None:

916 pixel_values = kwargs.pop("pixel_values", None)

917 image_embeds = kwargs.pop("image_embeds", None)

918 image_grid_thw = kwargs.pop("image_grid_thw", None)

@@ -932,26 +940,27 @@ class OpenPanguVLForConditionalGeneration(

940 image_grid_thw=image_grid_thw,

941 )

942

935 - if image_embeds is not None:

936 - image_embeds = self._validate_and_reshape_mm_tensor(

937 - image_embeds, "image embeds"

938 - )

939 - image_grid_thw = self._validate_and_reshape_mm_tensor(

940 - image_grid_thw, "image grid_thw"

941 - )

943 + assert image_embeds is not None

944 + image_embeds = self._validate_and_reshape_mm_tensor(

945 + image_embeds, "image embeds"

946 + )

947 + image_grid_thw = self._validate_and_reshape_mm_tensor(

948 + image_grid_thw, "image grid_thw"

949 + )

950

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

944 - raise ValueError(

945 - "Incorrect type of image embeddings. "

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

947 - )

948 - return OpenPanguVLImageEmbeddingInputs(

949 - type="image_embeds",

950 - image_embeds=image_embeds,

951 - image_grid_thw=image_grid_thw,

951 + if not isinstance(image_embeds, torch.Tensor):

952 + raise ValueError(

953 + f"Incorrect type of image embeddings. Got type: {type(image_embeds)}"

954 )

955 + return OpenPanguVLImageEmbeddingInputs(

956 + type="image_embeds",

957 + image_embeds=image_embeds,

958 + image_grid_thw=image_grid_thw,

959 + )

960

954 - def _parse_and_validate_video_input(self, **kwargs: object):

961 + def _parse_and_validate_video_input(

962 + self, **kwargs: object

963 + ) -> OpenPanguVLVideoInputs | None:

964 pixel_values_videos = kwargs.pop("pixel_values_videos", None)

965 video_embeds = kwargs.pop("video_embeds", None)

966 video_grid_thw = kwargs.pop("video_grid_thw", None)

@@ -973,27 +982,30 @@ class OpenPanguVLForConditionalGeneration(

982 video_grid_thw=video_grid_thw,

983 )

984

976 - if video_embeds is not None:

977 - video_embeds = self._validate_and_reshape_mm_tensor(

978 - video_embeds, "video embeds"

979 - )

980 - video_grid_thw = self._validate_and_reshape_mm_tensor(

981 - video_grid_thw, "video grid_thw"

982 - )

985 + assert video_embeds is not None

986 + video_embeds = self._validate_and_reshape_mm_tensor(

987 + video_embeds, "video embeds"

988 + )

989 + video_grid_thw = self._validate_and_reshape_mm_tensor(

990 + video_grid_thw, "video grid_thw"

991 + )

992

984 - if not isinstance(video_embeds, torch.Tensor):

985 - raise ValueError(

986 - "Incorrect type of video embeddings. "

987 - f"Got type: {type(video_embeds)}"

988 - )

989 - return OpenPanguVLVideoEmbeddingInputs(

990 - type="video_embeds",

991 - video_embeds=video_embeds,

992 - video_grid_thw=video_grid_thw,

993 + if not isinstance(video_embeds, torch.Tensor):

994 + raise ValueError(

995 + f"Incorrect type of video embeddings. Got type: {type(video_embeds)}"

996 )

997 + return OpenPanguVLVideoEmbeddingInputs(

998 + type="video_embeds",

999 + video_embeds=video_embeds,

1000 + video_grid_thw=video_grid_thw,

1001 + )

1002

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

996 - mm_input_by_modality = {}

1003 + def _parse_and_validate_multimodal_inputs(

1004 + self, **kwargs: object

1005 + ) -> dict[str, OpenPanguVLImageInputs | OpenPanguVLVideoInputs | None]:

1006 + mm_input_by_modality: dict[

1007 + str, OpenPanguVLImageInputs | OpenPanguVLVideoInputs | None

1008 + ] = {}

1009 for input_key in kwargs:

1010 if (

1011 input_key in ("pixel_values", "image_embeds")

@@ -1019,16 +1031,19 @@ class OpenPanguVLForConditionalGeneration(

1031 multimodal_embeddings: tuple[torch.Tensor, ...] = ()

1032

1033 for modality in mm_input_by_modality:

1022 - multimodal_input = mm_input_by_modality[modality]

1034 if modality == "image":

1024 - vision_embeddings = self._process_image_input(multimodal_input)

1035 + image_input = mm_input_by_modality["image"]

1036 + assert isinstance(image_input, OpenPanguVLImageInputs)

1037 + vision_embeddings = self._process_image_input(image_input)

1038 multimodal_embeddings = (

1039 multimodal_embeddings

1040 if not vision_embeddings

1041 else (multimodal_embeddings + vision_embeddings)

1042 )

1043 if modality == "video":

1031 - video_embeddings = self._process_video_input(multimodal_input)

1044 + video_input = mm_input_by_modality["video"]

1045 + assert isinstance(video_input, OpenPanguVLVideoInputs)

1046 + video_embeddings = self._process_video_input(video_input)

1047 multimodal_embeddings = (

1048 multimodal_embeddings

1049 if not video_embeddings

@@ -1041,20 +1056,19 @@ class OpenPanguVLForConditionalGeneration(

1056 input_ids: torch.Tensor,

1057 multimodal_embeddings: MultiModalEmbeddings | None = None,

1058 ) -> torch.Tensor:

1044 - if multimodal_embeddings is None:

1045 - return self.language_model.embed_input_ids(input_ids)

1046 -

1047 - mm_token_ids = input_ids.new_tensor(

1048 - [self.config.image_token_id, self.config.video_token_id]

1049 - )

1050 - is_multimodal = torch.isin(input_ids, mm_token_ids)

1051 - return self.embed_input_ids(

1052 - input_ids,

1053 - multimodal_embeddings,

1054 - is_multimodal=is_multimodal,

1055 - )

1059 + inputs_embeds = self.language_model.embed_input_ids(input_ids)

1060 + if multimodal_embeddings is not None:

1061 + inputs_embeds = self.embed_input_ids( # type: ignore[call-overload]

1062 + input_ids,

1063 + inputs_embeds,

1064 + multimodal_embeddings,

1065 + [self.config.image_token_id, self.config.video_token_id],

1066 + )

1067 + return inputs_embeds

1068

1057 - def _process_image_input(self, image_input) -> tuple[torch.Tensor, ...]:

1069 + def _process_image_input(

1070 + self, image_input: OpenPanguVLImageInputs

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

1072 grid_thw = image_input["image_grid_thw"]

1073 if grid_thw.ndim != 2:

1074 raise ValueError(f"grid_thw.ndim must be 2, but it is {grid_thw.ndim}")

@@ -1085,7 +1099,9 @@ class OpenPanguVLForConditionalGeneration(

1099 sizes = grid_thw.prod(-1) // merge_size // merge_size

1100 return image_embeds.split(sizes.tolist())

1101

1088 - def _process_video_input(self, video_input) -> torch.Tensor:

1102 + def _process_video_input(

1103 + self, video_input: OpenPanguVLVideoInputs

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

1105 grid_thw = video_input["video_grid_thw"]

1106 if grid_thw.ndim != 2:

1107 raise ValueError(f"grid_thw.ndim must be 2, but it is {grid_thw.ndim}")

@@ -1158,8 +1174,13 @@ class OpenPanguVLForConditionalGeneration(

1174 for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):

1175 offset = mm_feature.mm_position.offset

1176 modality = mm_feature.modality

1177 + feature_data = mm_feature.data

1178 + assert feature_data is not None

1179 if modality == "image":

1162 - t, h, w = mm_feature.data["image_grid_thw"].data.tolist()

1180 + grid_thw_item = feature_data["image_grid_thw"]

1181 + grid_thw = grid_thw_item.data

1182 + assert isinstance(grid_thw, torch.Tensor)

1183 + t, h, w = grid_thw.tolist()

1184 assert t == 1, f"Image must have 1 frame, got {t}"

1185 yield (

1186 modality,

@@ -1169,7 +1190,10 @@ class OpenPanguVLForConditionalGeneration(

1190 w // spatial_merge_size,

1191 )

1192 elif modality == "video":

1172 - t, h, w = mm_feature.data["video_grid_thw"].data.tolist()

1193 + grid_thw_item = feature_data["video_grid_thw"]

1194 + grid_thw = grid_thw_item.data

1195 + assert isinstance(grid_thw, torch.Tensor)

1196 + t, h, w = grid_thw.tolist()

1197 yield (

1198 modality,

1199 offset,

@@ -1281,14 +1305,17 @@ def normalize(image, mean, std):

1305 @lru_cache(maxsize=10)

1306 def _fuse_mean_std_and_rescale_factor(

1307 do_normalize: bool | None = None,

1284 - image_mean: float | list[float] | None = None,

1285 - image_std: float | list[float] | None = None,

1308 + image_mean: float | tuple[float, ...] | None = None,

1309 + image_std: float | tuple[float, ...] | None = None,

1310 do_rescale: bool | None = None,

1311 rescale_factor: float | None = None,

1312 device: Optional["torch.device"] = None,

1313 ) -> tuple:

1314 if do_rescale and do_normalize:

1315 # Fused rescale and normalize

1316 + assert rescale_factor is not None

1317 + assert image_mean is not None

1318 + assert image_std is not None

1319 image_mean = torch.tensor(image_mean, device=device) * (1.0 / rescale_factor)

1320 image_std = torch.tensor(image_std, device=device) * (1.0 / rescale_factor)

1321 do_rescale = False

@@ -1300,8 +1327,8 @@ def rescale_and_normalize(

1327 do_rescale: bool,

1328 rescale_factor: float,

1329 do_normalize: bool,

1303 - image_mean: float | list[float],

1304 - image_std: float | list[float],

1330 + image_mean: float | tuple[float, ...],

1331 + image_std: float | tuple[float, ...],

1332 dtype: torch.dtype = torch.bfloat16,

1333 ) -> "torch.Tensor":

1334 """Rescale and normalize images."""

vllm/model_executor/models/openvla.py

@@ -349,6 +349,7 @@ class OpenVLAMultiModalProcessor(BaseMultiModalProcessor[OpenVLAProcessingInfo])

349 if isinstance(images, ImageEmbeddingItems):

350 num_image_tokens = images.get_feature_size(item_idx)

351 else:

352 + assert isinstance(images, ImageProcessorItems)

353 image_size = images.get_image_size(item_idx)

354 num_image_tokens = self.info.get_num_image_tokens(

355 image_width=image_size.width,

vllm/model_executor/models/opt.py

@@ -218,6 +218,7 @@ class OPTDecoder(nn.Module):

218 )

219

220 # Project out & in will be replicated if they exist.

221 + self.project_out: ReplicatedLinear | None

222 if config.word_embed_proj_dim != config.hidden_size:

223 self.project_out = ReplicatedLinear(

224 config.hidden_size,

@@ -229,6 +230,7 @@ class OPTDecoder(nn.Module):

230 else:

231 self.project_out = None

232

233 + self.project_in: ReplicatedLinear | None

234 if config.word_embed_proj_dim != config.hidden_size:

235 self.project_in = ReplicatedLinear(

236 config.word_embed_proj_dim,

vllm/model_executor/models/ovis.py

+10 -3

@@ -59,7 +59,12 @@ from vllm.sequence import IntermediateTensors

59 from vllm.transformers_utils.processors.ovis import OvisProcessor

60 from vllm.utils.tensor_schema import TensorSchema, TensorShape

62 -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

62 +from .interfaces import (

63 + MultiModalEmbeddings,

64 + SupportsMultiModal,

65 + SupportsPP,

66 + supports_pp,

67 +)

69 # Cannot find the following number from hf config.

70 IMAGE_TOKEN = "<image>"

@@ -281,7 +286,7 @@ class OvisProcessingInfo(BaseProcessingInfo):

286 # minus 1 for presented image token

287 return (patch_grid_length // hidden_stride) ** 2 - 1

288

284 - def get_image_pad_token(self) -> str:

289 + def get_image_pad_token(self) -> str | None:

290 hf_text_config = self.get_hf_config().get_text_config()

291 text_model_type = hf_text_config.model_type

292 return IMAGE_PAD_TOKEN_MAP.get(text_model_type)

@@ -439,8 +444,10 @@ class Ovis(nn.Module, SupportsMultiModal, SupportsPP):

444 text_model_type = self.config.get_text_config().model_type

445 self.image_pad_token_id = IMAGE_PAD_TOKEN_ID_MAP[text_model_type]

446

447 + language_model = self.get_language_model()

448 + assert supports_pp(language_model)

449 self.make_empty_intermediate_tensors = (

443 - self.get_language_model().make_empty_intermediate_tensors

450 + language_model.make_empty_intermediate_tensors

451 )

452

453 def _parse_and_validate_image_input(

vllm/model_executor/models/ovis2_5.py

+23 -6

@@ -4,14 +4,14 @@

5 from collections.abc import Iterable, Mapping

6 from functools import partial

7 -from typing import Annotated, Literal

7 +from typing import Annotated, Literal, TypedDict

9 import torch

10 import torch.nn as nn

11 from transformers import BaseImageProcessor, BatchFeature, PretrainedConfig

13 from vllm.config import VllmConfig

14 -from vllm.config.multimodal import BaseDummyOptions

14 +from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions

15 from vllm.inputs import MultiModalDataDict

16 from vllm.model_executor.layers.linear import ReplicatedLinear

17 from vllm.model_executor.layers.quantization import QuantizationConfig

@@ -40,7 +40,12 @@ from vllm.sequence import IntermediateTensors

40 from vllm.transformers_utils.processors.ovis2_5 import Ovis2_5Processor

41 from vllm.utils.tensor_schema import TensorSchema, TensorShape

43 -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

43 +from .interfaces import (

44 + MultiModalEmbeddings,

45 + SupportsMultiModal,

46 + SupportsPP,

47 + supports_pp,

48 +)

50 IMAGE_TOKEN = "<image>"

51 VIDEO_TOKEN = "<video>"

@@ -80,6 +85,11 @@ class Ovis2_5VideoPatchInputs(TensorSchema):

85 # This is used to restore the first two dimensions of `flat_data`.

88 +class Ovis2_5MultiModalInputs(TypedDict, total=False):

89 + images: Ovis2_5ImagePatchInputs | None

90 + videos: Ovis2_5VideoPatchInputs | None

91 +

92 +

93 class VisualTokenizer(torch.nn.Module):

94 """VIT."""

@@ -292,6 +302,7 @@ class Ovis2_5DummyInputsBuilder(BaseDummyInputsBuilder[Ovis2_5ProcessingInfo]):

302

303 image_overrides = mm_options.get("image")

304 video_overrides = mm_options.get("video")

305 + assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)

306

307 mm_data = {

308 "image": self._get_dummy_images(

@@ -453,8 +464,10 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

464

465 self.image_pad_token_id: int = IMAGE_PAD_TOKEN_ID

466

467 + language_model = self.get_language_model()

468 + assert supports_pp(language_model)

469 self.make_empty_intermediate_tensors = (

457 - self.get_language_model().make_empty_intermediate_tensors

470 + language_model.make_empty_intermediate_tensors

471 )

472

473 def _parse_and_validate_image_input(

@@ -562,8 +575,10 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

575 vision_embeddings.append(torch.cat(vision_embeddings_per_image, dim=0))

576 return tuple(vision_embeddings)

577

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

566 - modalities = {}

578 + def _parse_and_validate_multimodal_inputs(

579 + self, **kwargs: object

580 + ) -> Ovis2_5MultiModalInputs:

581 + modalities: Ovis2_5MultiModalInputs = {}

582

583 # Preserve the order of modalities if there are multiple of them

584 # from the order of kwargs.

@@ -593,10 +608,12 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

608 for modality in modalities:

609 if modality == "images":

610 image_input = modalities["images"]

611 + assert image_input is not None

612 image_embeddings = self._process_visual_input(image_input)

613 multimodal_embeddings += tuple(image_embeddings)

614 if modality == "videos":

615 video_input = modalities["videos"]

616 + assert video_input is not None

617 video_embeddings = self._process_visual_input(video_input)

618 multimodal_embeddings += tuple(video_embeddings)

619

vllm/transformers_utils/processors/nano_nemotron_vl.py

+3 -1

@@ -551,7 +551,9 @@ class DynamicResolutionImageTiler:

551 )

552

553 @staticmethod

554 - def stack(images: list[torch.Tensor], patch_size: int) -> torch.Tensor:

554 + def stack(

555 + images: list[torch.Tensor] | torch.Tensor, patch_size: int

556 + ) -> torch.Tensor:

557 assert len(images) > 0, "No images to stack"

558

559 def rearrange_img(x):