@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 784cac7c42

[Mypy] Fix mypy typing for P models (#54169)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Taneem Ibrahim committed Sep 4, 2026 at 15:47 UTC 784cac7c424db2f2fdc879b672abad7e231f5698

9 files changed +126 -45

tools/pre_commit/mypy.py

-1

@@ -115,7 +115,6 @@ EXCLUDE = [

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

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

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

118 - r"vllm/model_executor/models/[pP]",

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

119 r"vllm/model_executor/models/[rR]",

120 r"vllm/model_executor/models/[sS]",

vllm/model_executor/models/paddleocr_vl.py

+20 -14

@@ -412,8 +412,8 @@ class SiglipVisionEmbeddings(nn.Module):

412

413 self.num_patches = (self.image_size // self.patch_size) ** 2

414 self.num_positions = self.num_patches

415 - self.cache_position_embedding = dict()

416 - self.cache_position_count = dict()

415 + self.cache_position_embedding: dict[tuple[int, int], torch.Tensor] = {}

416 + self.cache_position_count: dict[tuple[int, int], int] = {}

417 self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)

418

419 self.register_buffer(

@@ -469,7 +469,7 @@ class SiglipVisionEmbeddings(nn.Module):

469 if len(self.cache_position_embedding) >= max_cache:

470 min_hit_grid = min(

471 self.cache_position_count,

472 - key=self.cache_position_count.get,

472 + key=self.cache_position_count.__getitem__,

473 )

474 self.cache_position_count.pop(min_hit_grid)

475 self.cache_position_embedding.pop(min_hit_grid)

@@ -483,8 +483,7 @@ class SiglipVisionEmbeddings(nn.Module):

483 self,

484 pixel_values: torch.FloatTensor,

485 position_ids: torch.Tensor | None = None,

486 - image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]]

487 - | None = None,

486 + image_grid_thw: Sequence[tuple[int, int, int]] | None = None,

487 interpolate_pos_encoding=False,

488 ) -> torch.Tensor:

489 if pixel_values.dim() == 4:

@@ -508,6 +507,7 @@ class SiglipVisionEmbeddings(nn.Module):

507

508 start = 0

509 tmp_embeddings = list()

510 + assert image_grid_thw is not None

511 for image_grid in image_grid_thw:

512 t, h, w = image_grid

513 end = start + t * h * w

@@ -864,7 +864,7 @@ class SiglipVisionTransformer(nn.Module):

864 height_position_ids: torch.Tensor | None = None,

865 width_position_ids: torch.Tensor | None = None,

866 cu_seqlens: torch.Tensor | None = None,

867 - image_grid_thw: torch.Tensor | None = None,

867 + image_grid_thw: Sequence[tuple[int, int, int]] | None = None,

868 ) -> torch.Tensor:

869 hidden_states = self.embeddings(

870 pixel_values,

@@ -934,8 +934,7 @@ class SiglipVisionModel(nn.Module):

934 pixel_values,

935 interpolate_pos_encoding: bool = False,

936 position_ids: torch.Tensor | None = None,

937 - image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]]

938 - | None = None,

937 + image_grid_thw: Sequence[tuple[int, int, int]] | None = None,

938 cu_seqlens: torch.Tensor | None = None,

939 ) -> BaseModelOutputWithPooling:

940 return self.vision_model(

@@ -1029,17 +1028,24 @@ class PaddleOCRVLForConditionalGeneration(nn.Module, SupportsMultiModal, Support

1028 tokens_per_second = getattr(self.config.vision_config, "tokens_per_second", 1.0)

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

1030 offset = mm_feature.mm_position.offset

1031 + feature_data = mm_feature.data

1032 + assert feature_data is not None

1033 if mm_feature.modality == "image":

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

1034 + grid_data = feature_data["image_grid_thw"].data

1035 + assert isinstance(grid_data, torch.Tensor)

1036 + t, h, w = grid_data.tolist()

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

1038 yield offset, 1, h // spatial_merge_size, w // spatial_merge_size, 1.0

1039 elif mm_feature.modality == "video":

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

1040 + grid_data = feature_data["video_grid_thw"].data

1041 + assert isinstance(grid_data, torch.Tensor)

1042 + t, h, w = grid_data.tolist()

1043 second_per_grid_ts = 1.0

1039 - if mm_feature.data.get("second_per_grid_ts", None):

1040 - second_per_grid_ts = mm_feature.data[

1041 - "second_per_grid_ts"

1042 - ].data.item()

1044 + second_per_grid_item = feature_data.get("second_per_grid_ts")

1045 + if second_per_grid_item is not None:

1046 + second_per_grid_data = second_per_grid_item.data

1047 + assert isinstance(second_per_grid_data, torch.Tensor)

1048 + second_per_grid_ts = second_per_grid_data.item()

1049 t_factor = second_per_grid_ts * tokens_per_second

1050 yield (

1051 offset,

vllm/model_executor/models/paligemma.py

+4 -1

@@ -181,6 +181,7 @@ class PaliGemmaMultiModalProcessor(BaseMultiModalProcessor[PaliGemmaProcessingIn

181 if isinstance(images, ImageEmbeddingItems):

182 num_image_tokens = images.get_feature_size(item_idx)

183 else:

184 + assert isinstance(images, ImageProcessorItems)

185 image_size = images.get_image_size(item_idx)

186 num_image_tokens = self.info.get_num_image_tokens(

187 image_width=image_size.width,

@@ -197,11 +198,13 @@ class PaliGemmaMultiModalProcessor(BaseMultiModalProcessor[PaliGemmaProcessingIn

198 # Paligemma 1 and 2 have different tokenizer.add_bos_token

199 # Insert <image>*n + <bos> after <bos> for Paligemma 1

200 # Insert <image>*n + <bos> for Paligemma 2

201 + add_bos_token = getattr(tokenizer, "add_bos_token", None)

202 + assert isinstance(add_bos_token, bool)

203 return [

204 PromptInsertion(

205 modality="image",

206 target=PromptIndexTargets.prefix(

204 - [bos_token_id] if tokenizer.add_bos_token else []

207 + [bos_token_id] if add_bos_token else []

208 ),

209 insertion=get_insertion,

210 )

vllm/model_executor/models/param2moe.py

+4 -1

@@ -34,6 +34,7 @@ from vllm.model_executor.layers.activation import SiluAndMul

34 from vllm.model_executor.layers.attention import Attention

35 from vllm.model_executor.layers.fused_moe import (

36 FusedMoEFactory,

37 + MoERunner,

38 )

39 from vllm.model_executor.layers.layernorm import RMSNorm

40 from vllm.model_executor.layers.linear import (

@@ -344,7 +345,7 @@ class Param2MoEMoEBlock(nn.Module):

345 routed_scaling_factor=self.routed_scaling_factor,

346 )

347

347 - def maybe_get_fused_moe(self) -> FusedMoEFactory:

348 + def maybe_get_fused_moe(self) -> MoERunner:

349 return self.experts

350

351 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

@@ -560,6 +561,8 @@ class Param2MoEModel(nn.Module):

561 class Param2MoEMixtureOfExperts(MixtureOfExperts):

562 """Implements the vLLM MixtureOfExperts protocol for Param2MoE."""

563

564 + moe_mlp_layers: list[Param2MoEMoEBlock]

565 +

566 def extract_moe_parameters(self, example_moe: Param2MoEMoEBlock | None) -> None:

567 if example_moe is None:

568 raise RuntimeError(

vllm/model_executor/models/phi3v.py

@@ -453,6 +453,7 @@ class Phi3VMultiModalProcessor(BaseMultiModalProcessor[Phi3VProcessingInfo]):

453 if isinstance(images, ImageEmbeddingItems):

454 num_image_tokens = images.get_feature_size(item_idx)

455 else:

456 + assert isinstance(images, ImageProcessorItems)

457 image_size = images.get_image_size(item_idx)

458 num_image_tokens = self.info.get_num_image_tokens(

459 image_width=image_size.width,

vllm/model_executor/models/phi4mm.py

+34 -11

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

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import math

4 +import typing

5 from collections.abc import Iterable, Mapping, Sequence

6 from typing import Annotated, Any, Literal, TypeAlias

@@ -16,7 +17,7 @@ from transformers import (

17 )

19 from vllm.config import VllmConfig

19 -from vllm.config.multimodal import BaseDummyOptions

20 +from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions

21 from vllm.distributed import get_pp_group

22 from vllm.inputs import MultiModalDataDict

23 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -533,6 +534,11 @@ class Phi4MMAudioEmbeddingInputs(TensorSchema):

534 Phi4MMAudioInputs: TypeAlias = Phi4MMAudioFeatureInputs | Phi4MMAudioEmbeddingInputs

535

536

537 +class Phi4MMParsedModalities(typing.TypedDict, total=False):

538 + images: Phi4MMImagePixelInputs

539 + audios: Phi4MMAudioInputs

540 +

541 +

542 def stack_with_pad(

543 tensors: torch.Tensor | list[torch.Tensor],

544 padding_value: int | float = 0,

@@ -611,13 +617,13 @@ class Phi4MMProcessingInfo(BaseProcessingInfo):

617 aspect_ratio = orig_width / orig_height

618

619 # calculate the existing image aspect ratio

614 - target_ratios = set(

620 + target_ratio_candidates = set(

621 (i, j)

622 for i in range(1, max_num + 1)

623 for j in range(1, max_num + 1)

624 if i * j <= max_num and i * j >= min_num

625 )

620 - target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

626 + target_ratios = sorted(target_ratio_candidates, key=lambda x: x[0] * x[1])

627

628 # find the closest aspect ratio to the target

629 image_processor = self.get_hf_processor().image_processor

@@ -785,7 +791,8 @@ class Phi4MMProcessingInfo(BaseProcessingInfo):

791

792 # Resample to 16000 or 8000 if needed

793 if sr > 16000:

788 - audio_len //= sr // 16000

794 + resample_ratio = int(sr // 16000)

795 + audio_len //= resample_ratio

796 elif 8000 <= sr < 16000:

797 # We'll resample to 16K from 8K

798 audio_len *= 2

@@ -850,6 +857,7 @@ class Phi4MMDummyInputsBuilder(BaseDummyInputsBuilder[Phi4MMProcessingInfo]):

857

858 image_overrides = mm_options.get("image")

859 audio_overrides = mm_options.get("audio")

860 + assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)

861

862 mm_data = {

863 "image": self._get_dummy_images(

@@ -885,12 +893,20 @@ class Phi4MMMultiModalProcessor(BaseMultiModalProcessor[Phi4MMProcessingInfo]):

893 prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())

894

895 sr = self.info.get_feature_extractor(**hf_processor_mm_kwargs).sampling_rate

888 - if audio_data := mm_data.get("audios", []):

889 - mm_data["audios"] = [(data, sr) for data in audio_data]

896 + raw_audio_data = mm_data.get("audios", [])

897 + assert isinstance(raw_audio_data, Sequence)

898 + audio_data = []

899 + for data in raw_audio_data:

900 + assert isinstance(data, (list, np.ndarray, torch.Tensor))

901 + audio_data.append(data)

902 +

903 + processor_data = dict(text=prompt_text, **mm_data)

904 + if audio_data:

905 + processor_data["audios"] = [(data, sr) for data in audio_data]

906

907 processed_data = self.info.ctx.call_hf_processor(

908 self.info.get_hf_processor(**hf_processor_mm_kwargs),

893 - dict(text=prompt_text, **mm_data),

909 + processor_data,

910 hf_processor_mm_kwargs,

911 )

912

@@ -960,6 +976,7 @@ class Phi4MMMultiModalProcessor(BaseMultiModalProcessor[Phi4MMProcessingInfo]):

976 if isinstance(images, ImageEmbeddingItems):

977 num_image_tokens = images.get_feature_size(item_idx)

978 else:

979 + assert isinstance(images, ImageProcessorItems)

980 image_size = images.get_image_size(item_idx)

981 num_image_tokens = self.info.get_num_image_tokens(

982 image_width=image_size.width,

@@ -1201,8 +1218,10 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1218 num_img_tokens=num_img_tokens,

1219 )

1220

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

1205 - modalities = {}

1221 + def _parse_and_validate_multimodal_inputs(

1222 + self, **kwargs: object

1223 + ) -> Phi4MMParsedModalities:

1224 + modalities = Phi4MMParsedModalities()

1225

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

1227 # from the order of kwargs.

@@ -1211,12 +1230,16 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):

1230 input_key in ("input_image_embeds", "image_embeds")

1231 and "images" not in modalities

1232 ):

1214 - modalities["images"] = self._parse_and_validate_image_input(**kwargs)

1233 + image_input = self._parse_and_validate_image_input(**kwargs)

1234 + assert image_input is not None

1235 + modalities["images"] = image_input

1236 if (

1237 input_key in ("input_audio_embeds", "audio_embeds")

1238 and "audios" not in modalities

1239 ):

1219 - modalities["audios"] = self._parse_and_validate_audio_input(**kwargs)

1240 + audio_input = self._parse_and_validate_audio_input(**kwargs)

1241 + assert audio_input is not None

1242 + modalities["audios"] = audio_input

1243

1244 return modalities

1245

vllm/model_executor/models/phi4mm_audio.py

+48 -10

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

7 #!/usr/bin/env python3

8 import abc

9 import math

10 +import typing

11 from typing import Any, Literal

13 import numpy as np

@@ -152,7 +153,7 @@ class ConformerEncoderLayer(nn.Module):

153 batch_norm: bool = False,

154 activation: str = "relu",

155 chunk_se: int = 0,

155 - chunk_size: int = 18,

156 + chunk_size: int | list[int] = 18,

157 conv_activation: str = "relu",

158 conv_glu_type: str = "sigmoid",

159 bias_in_glu: bool = True,

@@ -392,6 +393,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

393 else None

394 )

395 if self.relative_attention_bias_type == "t5":

396 + assert relative_attention_bias_args is not None

397 assert self.num_heads % self.attention_group_size == 0, (

398 "attention_group_size must divide n_head"

399 )

@@ -405,6 +407,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

407 else:

408 raise NotImplementedError

409

410 + assert self.encoder_embedding_config is not None

411 self.encoder_embedding = MeanVarianceNormLayer(

412 self.encoder_embedding_config["input_size"]

413 )

@@ -421,6 +424,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

424 the cases and return one of them. Torchscript does support that.

425 """

426 if self.input_layer == "nemo_conv":

427 + assert self.nemo_conv_settings is not None

428 # Handle the special causal case

429 subsampling_causal_cond = self.nemo_conv_settings.get(

430 "subsampling", "dw_striding"

@@ -431,19 +435,19 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

435 ]

436 is_causal = self.nemo_conv_settings.get("is_causal", False)

437 if is_causal and subsampling_causal_cond:

434 - lens_change = (

435 - torch.ceil(feature_lens / self.time_reduction).long()

436 - if isinstance(feature_lens, Tensor)

437 - else math.ceil(feature_lens / self.time_reduction)

438 - )

438 feature_lens_remainder = feature_lens % self.time_reduction

439 if isinstance(feature_lens, Tensor):

440 + lens_change = torch.ceil(feature_lens / self.time_reduction).long()

441 lens_change[feature_lens_remainder != 1] += 1

442 - elif feature_lens_remainder != 1:

442 + return lens_change

443 +

444 + lens_change = math.ceil(feature_lens / self.time_reduction)

445 + if feature_lens_remainder != 1:

446 lens_change += 1

447 return lens_change

448 ceil_func = math.ceil if isinstance(feature_lens, int) else torch.ceil

449 return ceil_func(feature_lens / self.time_reduction)

450 + raise ValueError("unknown input_layer: " + self.input_layer)

451

452 @abc.abstractmethod

453 def forward(self) -> Any:

@@ -476,6 +480,8 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

480 )

481 left_chunk_train_eff = left_chunk[chunk_size_index]

482 else:

483 + assert chunk_size is not None

484 + assert isinstance(left_chunk, int)

485 chunk_size_train_eff = chunk_size

486 left_chunk_train_eff = left_chunk

487

@@ -516,7 +522,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

522 seq_len: int,

523 batch_size: int,

524 chunk_size: int | list[int],

519 - left_chunk: int | list[int],

525 + left_chunk: int | list[int] | None,

526 ) -> torch.Tensor:

527 chunk_size_train_eff, left_chunk_train_eff = self._chunk_size_selection(

528 chunk_size, left_chunk

@@ -535,6 +541,37 @@ class TransformerEncoderBase(abc.ABC, nn.Module):

541 )

542 return enc_streaming_mask

543

544 + @typing.overload

545 + def forward_embeddings(

546 + self,

547 + xs_pad: torch.Tensor,

548 + masks: torch.Tensor,

549 + chunk_size_nc: None = None,

550 + left_chunk_nc: int | list[int] | None = None,

551 + ) -> tuple[

552 + torch.Tensor,

553 + torch.Tensor | None,

554 + torch.Tensor | None,

555 + torch.Tensor,

556 + torch.Tensor,

557 + ]: ...

558 +

559 + @typing.overload

560 + def forward_embeddings(

561 + self,

562 + xs_pad: torch.Tensor,

563 + masks: torch.Tensor,

564 + chunk_size_nc: int | list[int],

565 + left_chunk_nc: int | list[int] | None = None,

566 + ) -> tuple[

567 + torch.Tensor,

568 + torch.Tensor | None,

569 + torch.Tensor | None,

570 + torch.Tensor,

571 + torch.Tensor,

572 + torch.Tensor,

573 + ]: ...

574 +

575 def forward_embeddings(

576 self,

577 xs_pad: torch.Tensor,

@@ -911,6 +948,7 @@ class ConformerEncoder(TransformerEncoderBase):

948 ) -> torch.Tensor | None:

949 if self.relative_attention_bias_layer:

950 return self.relative_attention_bias_layer(input_tensor)

951 + return None

952

953 def calculate_hs_mask(

954 self, xs_pad: torch.Tensor, device: torch.device, mask: torch.Tensor | None

@@ -1148,7 +1186,7 @@ class AudioEmbedding(nn.Module):

1186 if kwargs.get("use_qformer", False):

1187 qformer_config = kwargs.get("qformer_config", {})

1188 qformer_config["attention_dim"] = audio_dim_out

1151 - self.qformer = WindowQformer(**qformer_config)

1189 + self.qformer: WindowQformer | None = WindowQformer(**qformer_config)

1190 else:

1191 self.qformer = None

1192

@@ -1175,7 +1213,7 @@ class AudioEmbedding(nn.Module):

1213 "{i} should be specified outside of the NeMo dictionary"

1214 )

1215

1178 - self.conv_ds = NemoConvSubsampling(

1216 + self.conv_ds: NemoConvSubsampling | None = NemoConvSubsampling(

1217 **default_nemo_conv_settings,

1218 )

1219 else:

vllm/model_executor/models/phi4mm_utils.py

+9 -5

@@ -355,7 +355,7 @@ class ConvModule(nn.Module):

355 causal: bool = False,

356 batch_norm: bool = False,

357 chunk_se: int = 0,

358 - chunk_size: int = 18,

358 + chunk_size: int | list[int] = 18,

359 activation: str = "relu",

360 glu_type: str = "sigmoid",

361 bias_in_glu: bool = True,

@@ -790,7 +790,7 @@ class AbsolutePositionalEncoding(nn.Module):

790 self.d_model = d_model

791 self.xscale = math.sqrt(self.d_model)

792 self.dropout = torch.nn.Dropout(p=dropout_rate)

793 - self.pe = None

793 + self.pe: torch.Tensor | None = None

794 self.extend_pe(torch.tensor(0.0).expand(1, max_len))

795 self._register_load_state_dict_pre_hook(_pre_hook)

796

@@ -826,6 +826,7 @@ class AbsolutePositionalEncoding(nn.Module):

826

827 """

828 self.extend_pe(x)

829 + assert self.pe is not None

830 x = x * self.xscale + self.pe[:, : x.size(1)]

831 return self.dropout(x)

832

@@ -880,7 +881,7 @@ class CausalConv1D(nn.Conv1d):

881 out_channels: int,

882 kernel_size: int,

883 stride: int = 1,

883 - padding: str | int = 0,

884 + padding: str | int | list[int] | None = 0,

885 dilation: int = 1,

886 groups: int = 1,

887 bias: bool = True,

@@ -888,7 +889,9 @@ class CausalConv1D(nn.Conv1d):

889 device=None,

890 dtype=None,

891 ) -> None:

891 - self.cache_drop_size = None

892 + self.cache_drop_size: int | None = None

893 + self._left_padding: int

894 + self._right_padding: int

895 if padding is None:

896 self._left_padding = kernel_size - 1

897 self._right_padding = stride - 1

@@ -933,6 +936,7 @@ class CausalConv1D(nn.Conv1d):

936 else:

937 new_x = F.pad(x, pad=(0, self._right_padding))

938 new_x = torch.cat([cache, new_x], dim=-1)

939 + assert self.cache_drop_size is not None

940 if self.cache_drop_size > 0:

941 next_cache = new_x[:, :, : -self.cache_drop_size]

942 else:

@@ -965,7 +969,7 @@ class CausalConv2D(nn.Conv2d):

969 out_channels: int,

970 kernel_size: int,

971 stride: int = 1,

968 - padding: str | int = 0,

972 + padding: str | int | None = 0,

973 dilation: int = 1,

974 groups: int = 1,

975 bias: bool = True,

vllm/model_executor/models/plamo3.py

+6 -2

@@ -218,10 +218,14 @@ class Plamo3AttentionMixer(nn.Module):

218

219 q_shape = q.shape

220 q = q.reshape(q_shape[:-1] + (q_shape[-1] // self.head_dim, self.head_dim))

221 - q = self.q_norm.forward_native(q).reshape(q_shape)

221 + normalized_q = self.q_norm.forward_native(q)

222 + assert isinstance(normalized_q, torch.Tensor)

223 + q = normalized_q.reshape(q_shape)

224 k_shape = k.shape

225 k = k.reshape(k_shape[:-1] + (k_shape[-1] // self.head_dim, self.head_dim))

224 - k = self.k_norm.forward_native(k).reshape(k_shape)

226 + normalized_k = self.k_norm.forward_native(k)

227 + assert isinstance(normalized_k, torch.Tensor)

228 + k = normalized_k.reshape(k_shape)

229

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

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