iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Taneem Ibrahim
2026-09-04 20:47:11 +00:00
committed by GitHub
co-authored by Harry Mellor
parent 5093e4844a
commit 784cac7c42
9 changed files with 126 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/pre_commit/mypy.py
-1
View File
@@ -115,7 +115,6 @@ EXCLUDE = [
r"vllm/model_executor/models/[kK]",
r"vllm/model_executor/models/[nN]",
r"vllm/model_executor/models/[oO]",
r"vllm/model_executor/models/[pP]",
r"vllm/model_executor/models/[qQ]",
r"vllm/model_executor/models/[rR]",
r"vllm/model_executor/models/[sS]",
vllm/model_executor/models/paddleocr_vl.py
+20 -14
View File
@@ -412,8 +412,8 @@ class SiglipVisionEmbeddings(nn.Module):
self.num_patches = (self.image_size // self.patch_size) ** 2
self.num_positions = self.num_patches
self.cache_position_embedding = dict()
self.cache_position_count = dict()
self.cache_position_embedding: dict[tuple[int, int], torch.Tensor] = {}
self.cache_position_count: dict[tuple[int, int], int] = {}
self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
self.register_buffer(
@@ -469,7 +469,7 @@ class SiglipVisionEmbeddings(nn.Module):
if len(self.cache_position_embedding) >= max_cache:
min_hit_grid = min(
self.cache_position_count,
key=self.cache_position_count.get,
key=self.cache_position_count.__getitem__,
)
self.cache_position_count.pop(min_hit_grid)
self.cache_position_embedding.pop(min_hit_grid)
@@ -483,8 +483,7 @@ class SiglipVisionEmbeddings(nn.Module):
self,
pixel_values: torch.FloatTensor,
position_ids: torch.Tensor | None = None,
image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]]
| None = None,
image_grid_thw: Sequence[tuple[int, int, int]] | None = None,
interpolate_pos_encoding=False,
) -> torch.Tensor:
if pixel_values.dim() == 4:
@@ -508,6 +507,7 @@ class SiglipVisionEmbeddings(nn.Module):
start = 0
tmp_embeddings = list()
assert image_grid_thw is not None
for image_grid in image_grid_thw:
t, h, w = image_grid
end = start + t * h * w
@@ -864,7 +864,7 @@ class SiglipVisionTransformer(nn.Module):
height_position_ids: torch.Tensor | None = None,
width_position_ids: torch.Tensor | None = None,
cu_seqlens: torch.Tensor | None = None,
image_grid_thw: torch.Tensor | None = None,
image_grid_thw: Sequence[tuple[int, int, int]] | None = None,
) -> torch.Tensor:
hidden_states = self.embeddings(
pixel_values,
@@ -934,8 +934,7 @@ class SiglipVisionModel(nn.Module):
pixel_values,
interpolate_pos_encoding: bool = False,
position_ids: torch.Tensor | None = None,
image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]]
| None = None,
image_grid_thw: Sequence[tuple[int, int, int]] | None = None,
cu_seqlens: torch.Tensor | None = None,
) -> BaseModelOutputWithPooling:
return self.vision_model(
@@ -1029,17 +1028,24 @@ class PaddleOCRVLForConditionalGeneration(nn.Module, SupportsMultiModal, Support
tokens_per_second = getattr(self.config.vision_config, "tokens_per_second", 1.0)
for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):
offset = mm_feature.mm_position.offset
feature_data = mm_feature.data
assert feature_data is not None
if mm_feature.modality == "image":
t, h, w = mm_feature.data["image_grid_thw"].data.tolist()
grid_data = feature_data["image_grid_thw"].data
assert isinstance(grid_data, torch.Tensor)
t, h, w = grid_data.tolist()
assert t == 1, f"Image must have 1 frame, got {t}"
yield offset, 1, h // spatial_merge_size, w // spatial_merge_size, 1.0
elif mm_feature.modality == "video":
t, h, w = mm_feature.data["video_grid_thw"].data.tolist()
grid_data = feature_data["video_grid_thw"].data
assert isinstance(grid_data, torch.Tensor)
t, h, w = grid_data.tolist()
second_per_grid_ts = 1.0
if mm_feature.data.get("second_per_grid_ts", None):
second_per_grid_ts = mm_feature.data[
"second_per_grid_ts"
].data.item()
second_per_grid_item = feature_data.get("second_per_grid_ts")
if second_per_grid_item is not None:
second_per_grid_data = second_per_grid_item.data
assert isinstance(second_per_grid_data, torch.Tensor)
second_per_grid_ts = second_per_grid_data.item()
t_factor = second_per_grid_ts * tokens_per_second
yield (
offset,
vllm/model_executor/models/paligemma.py
+4 -1
View File
@@ -181,6 +181,7 @@ class PaliGemmaMultiModalProcessor(BaseMultiModalProcessor[PaliGemmaProcessingIn
if isinstance(images, ImageEmbeddingItems):
num_image_tokens = images.get_feature_size(item_idx)
else:
assert isinstance(images, ImageProcessorItems)
image_size = images.get_image_size(item_idx)
num_image_tokens = self.info.get_num_image_tokens(
image_width=image_size.width,
@@ -197,11 +198,13 @@ class PaliGemmaMultiModalProcessor(BaseMultiModalProcessor[PaliGemmaProcessingIn
# Paligemma 1 and 2 have different tokenizer.add_bos_token
# Insert <image>*n + <bos> after <bos> for Paligemma 1
# Insert <image>*n + <bos> for Paligemma 2
add_bos_token = getattr(tokenizer, "add_bos_token", None)
assert isinstance(add_bos_token, bool)
return [
PromptInsertion(
modality="image",
target=PromptIndexTargets.prefix(
[bos_token_id] if tokenizer.add_bos_token else []
[bos_token_id] if add_bos_token else []
),
insertion=get_insertion,
)
vllm/model_executor/models/param2moe.py
+4 -1
View File
@@ -34,6 +34,7 @@ from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.fused_moe import (
FusedMoEFactory,
MoERunner,
)
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
@@ -344,7 +345,7 @@ class Param2MoEMoEBlock(nn.Module):
routed_scaling_factor=self.routed_scaling_factor,
)
def maybe_get_fused_moe(self) -> FusedMoEFactory:
def maybe_get_fused_moe(self) -> MoERunner:
return self.experts
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -560,6 +561,8 @@ class Param2MoEModel(nn.Module):
class Param2MoEMixtureOfExperts(MixtureOfExperts):
"""Implements the vLLM MixtureOfExperts protocol for Param2MoE."""
moe_mlp_layers: list[Param2MoEMoEBlock]
def extract_moe_parameters(self, example_moe: Param2MoEMoEBlock | None) -> None:
if example_moe is None:
raise RuntimeError(
vllm/model_executor/models/phi3v.py
+1
View File
@@ -453,6 +453,7 @@ class Phi3VMultiModalProcessor(BaseMultiModalProcessor[Phi3VProcessingInfo]):
if isinstance(images, ImageEmbeddingItems):
num_image_tokens = images.get_feature_size(item_idx)
else:
assert isinstance(images, ImageProcessorItems)
image_size = images.get_image_size(item_idx)
num_image_tokens = self.info.get_num_image_tokens(
image_width=image_size.width,
vllm/model_executor/models/phi4mm.py
+34 -11
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
import typing
from collections.abc import Iterable, Mapping, Sequence
from typing import Annotated, Any, Literal, TypeAlias
@@ -16,7 +17,7 @@ from transformers import (
)
from vllm.config import VllmConfig
from vllm.config.multimodal import BaseDummyOptions
from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
from vllm.distributed import get_pp_group
from vllm.inputs import MultiModalDataDict
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -533,6 +534,11 @@ class Phi4MMAudioEmbeddingInputs(TensorSchema):
Phi4MMAudioInputs: TypeAlias = Phi4MMAudioFeatureInputs | Phi4MMAudioEmbeddingInputs
class Phi4MMParsedModalities(typing.TypedDict, total=False):
images: Phi4MMImagePixelInputs
audios: Phi4MMAudioInputs
def stack_with_pad(
tensors: torch.Tensor | list[torch.Tensor],
padding_value: int | float = 0,
@@ -611,13 +617,13 @@ class Phi4MMProcessingInfo(BaseProcessingInfo):
aspect_ratio = orig_width / orig_height
# calculate the existing image aspect ratio
target_ratios = set(
target_ratio_candidates = set(
(i, j)
for i in range(1, max_num + 1)
for j in range(1, max_num + 1)
if i * j <= max_num and i * j >= min_num
)
target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
target_ratios = sorted(target_ratio_candidates, key=lambda x: x[0] * x[1])
# find the closest aspect ratio to the target
image_processor = self.get_hf_processor().image_processor
@@ -785,7 +791,8 @@ class Phi4MMProcessingInfo(BaseProcessingInfo):
# Resample to 16000 or 8000 if needed
if sr > 16000:
audio_len //= sr // 16000
resample_ratio = int(sr // 16000)
audio_len //= resample_ratio
elif 8000 <= sr < 16000:
# We'll resample to 16K from 8K
audio_len *= 2
@@ -850,6 +857,7 @@ class Phi4MMDummyInputsBuilder(BaseDummyInputsBuilder[Phi4MMProcessingInfo]):
image_overrides = mm_options.get("image")
audio_overrides = mm_options.get("audio")
assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
mm_data = {
"image": self._get_dummy_images(
@@ -885,12 +893,20 @@ class Phi4MMMultiModalProcessor(BaseMultiModalProcessor[Phi4MMProcessingInfo]):
prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts())
sr = self.info.get_feature_extractor(**hf_processor_mm_kwargs).sampling_rate
if audio_data := mm_data.get("audios", []):
mm_data["audios"] = [(data, sr) for data in audio_data]
raw_audio_data = mm_data.get("audios", [])
assert isinstance(raw_audio_data, Sequence)
audio_data = []
for data in raw_audio_data:
assert isinstance(data, (list, np.ndarray, torch.Tensor))
audio_data.append(data)
processor_data = dict(text=prompt_text, **mm_data)
if audio_data:
processor_data["audios"] = [(data, sr) for data in audio_data]
processed_data = self.info.ctx.call_hf_processor(
self.info.get_hf_processor(**hf_processor_mm_kwargs),
dict(text=prompt_text, **mm_data),
processor_data,
hf_processor_mm_kwargs,
)
@@ -960,6 +976,7 @@ class Phi4MMMultiModalProcessor(BaseMultiModalProcessor[Phi4MMProcessingInfo]):
if isinstance(images, ImageEmbeddingItems):
num_image_tokens = images.get_feature_size(item_idx)
else:
assert isinstance(images, ImageProcessorItems)
image_size = images.get_image_size(item_idx)
num_image_tokens = self.info.get_num_image_tokens(
image_width=image_size.width,
@@ -1201,8 +1218,10 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
num_img_tokens=num_img_tokens,
)
def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
modalities = {}
def _parse_and_validate_multimodal_inputs(
self, **kwargs: object
) -> Phi4MMParsedModalities:
modalities = Phi4MMParsedModalities()
# Preserve the order of modalities if there are multiple of them
# from the order of kwargs.
@@ -1211,12 +1230,16 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
input_key in ("input_image_embeds", "image_embeds")
and "images" not in modalities
):
modalities["images"] = self._parse_and_validate_image_input(**kwargs)
image_input = self._parse_and_validate_image_input(**kwargs)
assert image_input is not None
modalities["images"] = image_input
if (
input_key in ("input_audio_embeds", "audio_embeds")
and "audios" not in modalities
):
modalities["audios"] = self._parse_and_validate_audio_input(**kwargs)
audio_input = self._parse_and_validate_audio_input(**kwargs)
assert audio_input is not None
modalities["audios"] = audio_input
return modalities
vllm/model_executor/models/phi4mm_audio.py
+48 -10
View File
@@ -7,6 +7,7 @@
#!/usr/bin/env python3
import abc
import math
import typing
from typing import Any, Literal
import numpy as np
@@ -152,7 +153,7 @@ class ConformerEncoderLayer(nn.Module):
batch_norm: bool = False,
activation: str = "relu",
chunk_se: int = 0,
chunk_size: int = 18,
chunk_size: int | list[int] = 18,
conv_activation: str = "relu",
conv_glu_type: str = "sigmoid",
bias_in_glu: bool = True,
@@ -392,6 +393,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
else None
)
if self.relative_attention_bias_type == "t5":
assert relative_attention_bias_args is not None
assert self.num_heads % self.attention_group_size == 0, (
"attention_group_size must divide n_head"
)
@@ -405,6 +407,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
else:
raise NotImplementedError
assert self.encoder_embedding_config is not None
self.encoder_embedding = MeanVarianceNormLayer(
self.encoder_embedding_config["input_size"]
)
@@ -421,6 +424,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
the cases and return one of them. Torchscript does support that.
"""
if self.input_layer == "nemo_conv":
assert self.nemo_conv_settings is not None
# Handle the special causal case
subsampling_causal_cond = self.nemo_conv_settings.get(
"subsampling", "dw_striding"
@@ -431,19 +435,19 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
]
is_causal = self.nemo_conv_settings.get("is_causal", False)
if is_causal and subsampling_causal_cond:
lens_change = (
torch.ceil(feature_lens / self.time_reduction).long()
if isinstance(feature_lens, Tensor)
else math.ceil(feature_lens / self.time_reduction)
)
feature_lens_remainder = feature_lens % self.time_reduction
if isinstance(feature_lens, Tensor):
lens_change = torch.ceil(feature_lens / self.time_reduction).long()
lens_change[feature_lens_remainder != 1] += 1
elif feature_lens_remainder != 1:
return lens_change
lens_change = math.ceil(feature_lens / self.time_reduction)
if feature_lens_remainder != 1:
lens_change += 1
return lens_change
ceil_func = math.ceil if isinstance(feature_lens, int) else torch.ceil
return ceil_func(feature_lens / self.time_reduction)
raise ValueError("unknown input_layer: " + self.input_layer)
@abc.abstractmethod
def forward(self) -> Any:
@@ -476,6 +480,8 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
)
left_chunk_train_eff = left_chunk[chunk_size_index]
else:
assert chunk_size is not None
assert isinstance(left_chunk, int)
chunk_size_train_eff = chunk_size
left_chunk_train_eff = left_chunk
@@ -516,7 +522,7 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
seq_len: int,
batch_size: int,
chunk_size: int | list[int],
left_chunk: int | list[int],
left_chunk: int | list[int] | None,
) -> torch.Tensor:
chunk_size_train_eff, left_chunk_train_eff = self._chunk_size_selection(
chunk_size, left_chunk
@@ -535,6 +541,37 @@ class TransformerEncoderBase(abc.ABC, nn.Module):
)
return enc_streaming_mask
@typing.overload
def forward_embeddings(
self,
xs_pad: torch.Tensor,
masks: torch.Tensor,
chunk_size_nc: None = None,
left_chunk_nc: int | list[int] | None = None,
) -> tuple[
torch.Tensor,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor,
torch.Tensor,
]: ...
@typing.overload
def forward_embeddings(
self,
xs_pad: torch.Tensor,
masks: torch.Tensor,
chunk_size_nc: int | list[int],
left_chunk_nc: int | list[int] | None = None,
) -> tuple[
torch.Tensor,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor,
torch.Tensor,
torch.Tensor,
]: ...
def forward_embeddings(
self,
xs_pad: torch.Tensor,
@@ -911,6 +948,7 @@ class ConformerEncoder(TransformerEncoderBase):
) -> torch.Tensor | None:
if self.relative_attention_bias_layer:
return self.relative_attention_bias_layer(input_tensor)
return None
def calculate_hs_mask(
self, xs_pad: torch.Tensor, device: torch.device, mask: torch.Tensor | None
@@ -1148,7 +1186,7 @@ class AudioEmbedding(nn.Module):
if kwargs.get("use_qformer", False):
qformer_config = kwargs.get("qformer_config", {})
qformer_config["attention_dim"] = audio_dim_out
self.qformer = WindowQformer(**qformer_config)
self.qformer: WindowQformer | None = WindowQformer(**qformer_config)
else:
self.qformer = None
@@ -1175,7 +1213,7 @@ class AudioEmbedding(nn.Module):
"{i} should be specified outside of the NeMo dictionary"
)
self.conv_ds = NemoConvSubsampling(
self.conv_ds: NemoConvSubsampling | None = NemoConvSubsampling(
**default_nemo_conv_settings,
)
else:
vllm/model_executor/models/phi4mm_utils.py
+9 -5
View File
@@ -355,7 +355,7 @@ class ConvModule(nn.Module):
causal: bool = False,
batch_norm: bool = False,
chunk_se: int = 0,
chunk_size: int = 18,
chunk_size: int | list[int] = 18,
activation: str = "relu",
glu_type: str = "sigmoid",
bias_in_glu: bool = True,
@@ -790,7 +790,7 @@ class AbsolutePositionalEncoding(nn.Module):
self.d_model = d_model
self.xscale = math.sqrt(self.d_model)
self.dropout = torch.nn.Dropout(p=dropout_rate)
self.pe = None
self.pe: torch.Tensor | None = None
self.extend_pe(torch.tensor(0.0).expand(1, max_len))
self._register_load_state_dict_pre_hook(_pre_hook)
@@ -826,6 +826,7 @@ class AbsolutePositionalEncoding(nn.Module):
"""
self.extend_pe(x)
assert self.pe is not None
x = x * self.xscale + self.pe[:, : x.size(1)]
return self.dropout(x)
@@ -880,7 +881,7 @@ class CausalConv1D(nn.Conv1d):
out_channels: int,
kernel_size: int,
stride: int = 1,
padding: str | int = 0,
padding: str | int | list[int] | None = 0,
dilation: int = 1,
groups: int = 1,
bias: bool = True,
@@ -888,7 +889,9 @@ class CausalConv1D(nn.Conv1d):
device=None,
dtype=None,
) -> None:
self.cache_drop_size = None
self.cache_drop_size: int | None = None
self._left_padding: int
self._right_padding: int
if padding is None:
self._left_padding = kernel_size - 1
self._right_padding = stride - 1
@@ -933,6 +936,7 @@ class CausalConv1D(nn.Conv1d):
else:
new_x = F.pad(x, pad=(0, self._right_padding))
new_x = torch.cat([cache, new_x], dim=-1)
assert self.cache_drop_size is not None
if self.cache_drop_size > 0:
next_cache = new_x[:, :, : -self.cache_drop_size]
else:
@@ -965,7 +969,7 @@ class CausalConv2D(nn.Conv2d):
out_channels: int,
kernel_size: int,
stride: int = 1,
padding: str | int = 0,
padding: str | int | None = 0,
dilation: int = 1,
groups: int = 1,
bias: bool = True,
vllm/model_executor/models/plamo3.py
+6 -2
View File
@@ -218,10 +218,14 @@ class Plamo3AttentionMixer(nn.Module):
q_shape = q.shape
q = q.reshape(q_shape[:-1] + (q_shape[-1] // self.head_dim, self.head_dim))
q = self.q_norm.forward_native(q).reshape(q_shape)
normalized_q = self.q_norm.forward_native(q)
assert isinstance(normalized_q, torch.Tensor)
q = normalized_q.reshape(q_shape)
k_shape = k.shape
k = k.reshape(k_shape[:-1] + (k_shape[-1] // self.head_dim, self.head_dim))
k = self.k_norm.forward_native(k).reshape(k_shape)
normalized_k = self.k_norm.forward_native(k)
assert isinstance(normalized_k, torch.Tensor)
k = normalized_k.reshape(k_shape)
q, k = self.rotary_emb(positions, q, k)
attn_output = self.attn(q, k, v)