From 76f3249fbdff761aa897218913050de0526ccca2 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Tue, 1 Sep 2026 10:26:00 -0500 Subject: [PATCH] [Mypy] Fix typing for M models (#54262) Signed-off-by: Taneem Ibrahim --- tools/pre_commit/mypy.py | 1 - vllm/model_executor/models/mamba.py | 4 +- vllm/model_executor/models/medusa.py | 4 +- vllm/model_executor/models/mellum.py | 42 +++-- vllm/model_executor/models/midashenglm.py | 16 +- vllm/model_executor/models/mimo_audio.py | 29 ++-- vllm/model_executor/models/mimo_v2.py | 8 +- vllm/model_executor/models/mimo_v2_mtp.py | 15 +- vllm/model_executor/models/mimo_v2_omni.py | 162 +++++++++++++----- vllm/model_executor/models/minicpm.py | 1 + vllm/model_executor/models/minicpm_eagle.py | 21 ++- vllm/model_executor/models/minicpmo.py | 78 +++++++-- vllm/model_executor/models/minicpmv.py | 56 ++++-- vllm/model_executor/models/minicpmv4_6.py | 11 +- vllm/model_executor/models/mistral.py | 10 +- vllm/model_executor/models/mistral3.py | 1 + vllm/model_executor/models/mistral_eagle.py | 29 +++- .../models/mistral_large_3_eagle.py | 18 +- vllm/model_executor/models/mllama4.py | 17 +- vllm/model_executor/models/module_mapping.py | 8 +- vllm/model_executor/models/molmo.py | 11 +- vllm/model_executor/models/molmo2.py | 50 +++++- vllm/model_executor/models/moonvit.py | 4 +- vllm/model_executor/models/moss_audio.py | 79 +++++---- .../models/moss_transcribe_diarize.py | 9 +- vllm/model_executor/models/muse_glimmer.py | 28 ++- 26 files changed, 512 insertions(+), 200 deletions(-) diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index 5e0dd0fe84..55ebd6ef18 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -114,7 +114,6 @@ SEPARATE_GROUPS = [ EXCLUDE = [ r"vllm/model_executor/models/[kK]", r"vllm/model_executor/models/[lL]", - r"vllm/model_executor/models/[mM]", r"vllm/model_executor/models/[nN]", r"vllm/model_executor/models/[oO]", r"vllm/model_executor/models/[pP]", diff --git a/vllm/model_executor/models/mamba.py b/vllm/model_executor/models/mamba.py index 72893c574f..286096abf1 100644 --- a/vllm/model_executor/models/mamba.py +++ b/vllm/model_executor/models/mamba.py @@ -52,14 +52,14 @@ class MambaDecoderLayer(nn.Module): model_config: ModelConfig | None = None, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, - is_lora_enabled: bool | None = False, + is_lora_enabled: bool = False, prefix: str = "", ) -> None: super().__init__() self.config = config self.is_falcon_mamba = config.model_type == "falcon_mamba" self.is_lora_enabled = is_lora_enabled - mixer_rms_eps = config.mixer_rms_eps if self.is_falcon_mamba else None + mixer_rms_eps = config.mixer_rms_eps if self.is_falcon_mamba else 1e-5 self.mixer = MambaMixer( hidden_size=config.hidden_size, ssm_state_size=config.state_size, diff --git a/vllm/model_executor/models/medusa.py b/vllm/model_executor/models/medusa.py index a857533ab4..6adf3b7669 100644 --- a/vllm/model_executor/models/medusa.py +++ b/vllm/model_executor/models/medusa.py @@ -54,7 +54,9 @@ class Medusa(nn.Module): needs to have truncated_vocab_size (=k) as an attribute.""" def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - config = vllm_config.speculative_config.draft_model_config.hf_config + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + config = speculative_config.draft_model_config.hf_config super().__init__() self.config = config self.blocks = nn.ModuleList( diff --git a/vllm/model_executor/models/mellum.py b/vllm/model_executor/models/mellum.py index 2eb35687d7..ebfbaf93fd 100644 --- a/vllm/model_executor/models/mellum.py +++ b/vllm/model_executor/models/mellum.py @@ -92,24 +92,30 @@ class MellumAttention(Qwen3MoeAttention): rope_parameters=rope_parameters, dual_chunk_attention_config=dual_chunk_attention_config, ) - self.attn = Attention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - cache_config=cache_config, - quant_config=quant_config, - per_layer_sliding_window=per_layer_sliding_window, - prefix=f"{prefix}.attn", - **( - { - "layer_idx": extract_layer_index(prefix), - "dual_chunk_attention_config": dual_chunk_attention_config, - } - if dual_chunk_attention_config - else {} - ), - ) + if dual_chunk_attention_config: + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=per_layer_sliding_window, + prefix=f"{prefix}.attn", + layer_idx=extract_layer_index(prefix), + dual_chunk_attention_config=dual_chunk_attention_config, + ) + else: + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=per_layer_sliding_window, + prefix=f"{prefix}.attn", + ) self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) diff --git a/vllm/model_executor/models/midashenglm.py b/vllm/model_executor/models/midashenglm.py index f84294f04f..1129a0e4b2 100644 --- a/vllm/model_executor/models/midashenglm.py +++ b/vllm/model_executor/models/midashenglm.py @@ -36,7 +36,7 @@ from torch.nn.functional import scaled_dot_product_attention from transformers import BatchFeature from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions from vllm.distributed import get_tensor_model_parallel_world_size from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.activation import get_act_fn @@ -570,6 +570,7 @@ class MiDashengLMDummyInputsBuilder(BaseDummyInputsBuilder[MiDashengLMProcessing num_audios = mm_counts.get("audio", 0) audio_overrides = mm_options.get("audio") + assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) return { "audio": self._get_dummy_audios( @@ -593,6 +594,7 @@ class MiDashengLMMultiModalProcessor( ) -> tuple[Mapping[str, object], Mapping[str, object]]: mm_data = dict(mm_data) audios = mm_data.pop("audios", []) + assert isinstance(audios, list) min_audio_len = self.info.get_min_audio_len() processed_audios = [ @@ -637,6 +639,7 @@ class MiDashengLMMultiModalProcessor( out_mm_data = out_mm_kwargs.get_data() audio_length = out_mm_data.get("audio_length") + audio_output_lengths: list[int] if audio_length is None: audio_output_lengths = [] else: @@ -645,10 +648,13 @@ class MiDashengLMMultiModalProcessor( if isinstance(audio_length, torch.Tensor) else audio_length ) - audio_output_lengths = [ - max(1, calculate_mel_frames_dasheng(int(length))) # at least one frame - for length in audio_length_np - ] + assert isinstance(audio_length_np, (list, tuple, np.ndarray)) + audio_output_lengths = [] + for length in audio_length_np: + assert isinstance(length, (int, np.integer)) + audio_output_lengths.append( + max(1, calculate_mel_frames_dasheng(int(length))) + ) def get_replacement_midashenglm(item_idx: int): num_features = audio_output_lengths[item_idx] diff --git a/vllm/model_executor/models/mimo_audio.py b/vllm/model_executor/models/mimo_audio.py index 91d46b1cea..f1f901418a 100644 --- a/vllm/model_executor/models/mimo_audio.py +++ b/vllm/model_executor/models/mimo_audio.py @@ -531,18 +531,18 @@ class MiMoAudioTokenizerConfig(PretrainedConfig): kernel_size: int = 3, activation_function: str = "gelu", encoder_layers: int = 8, - encoder_skip_layer_id: int = None, + encoder_skip_layer_id: int | None = None, encoder_attention_heads: int = 12, encoder_ffn_dim: int = 3072, encoder_causal: bool = False, - encoder_attn_window_size: list = None, + encoder_attn_window_size: list[int] | None = None, decoder_layers: int = 8, decoder_attention_heads: int = 12, decoder_ffn_dim: int = 3072, decoder_kernel_size: int = 3, decoder_stride_size: int = 2, decoder_causal: bool = True, - decoder_attn_window_size: list = None, + decoder_attn_window_size: list[int] | None = None, nfft: int = 1024, vocoder_dim: int = 512, vocoder_intermediate_dim: int = 4096, @@ -553,16 +553,16 @@ class MiMoAudioTokenizerConfig(PretrainedConfig): window_size: int = 1024, vocoder_padding: str = "same", fmin: int = 0, - fmax: int = None, + fmax: int | None = None, num_quantizers: int = 12, - codebook_size: list = None, + codebook_size: list[int] | None = None, threshold_ema_dead_code: int = 10, position_embedding_type: str = "rope", rope_theta: int = 10000, rope_type: str = "default", ln_type: str = "LayerNorm", vocoder_attention_heads: int = 4, - vocoder_attn_window_size: list = None, + vocoder_attn_window_size: list[int] | None = None, use_istft_only: bool = False, hybrid_attention: bool = False, hybrid_block_size: int = 8, @@ -808,17 +808,22 @@ class AudioEncoder(nn.Module): config.rope_type, ) - attn_window_sizes = [] + encoder_attn_window_size = config.encoder_attn_window_size + assert encoder_attn_window_size is not None + assert len(encoder_attn_window_size) == 2 + attn_window_size = ( + encoder_attn_window_size[0], + encoder_attn_window_size[1], + ) + attn_window_sizes: list[tuple[int, int]] = [] if config.hybrid_attention: for i in range(config.encoder_layers): if i % config.swa_per_block < config.swa_per_block - 1: - attn_window_sizes.append(tuple(config.encoder_attn_window_size)) + attn_window_sizes.append(attn_window_size) else: attn_window_sizes.append((-1, -1)) else: - attn_window_sizes = [ - tuple(config.encoder_attn_window_size) - ] * config.encoder_layers + attn_window_sizes = [attn_window_size] * config.encoder_layers self.layers = nn.ModuleList( [ @@ -848,6 +853,7 @@ class AudioEncoder(nn.Module): else: self.down_sample_layer = None + self.quantizer: ResidualVectorQuantizer | None if config.num_quantizers != 0: self.quantizer = ResidualVectorQuantizer( dimension=config.d_model, @@ -976,6 +982,7 @@ class AudioEncoder(nn.Module): @torch.no_grad() def decode_vq(self, codes): + assert self.quantizer is not None self.quantizer.float() return self.quantizer.decode(codes) diff --git a/vllm/model_executor/models/mimo_v2.py b/vllm/model_executor/models/mimo_v2.py index 8df6a4c0e8..c7ba39bb0d 100644 --- a/vllm/model_executor/models/mimo_v2.py +++ b/vllm/model_executor/models/mimo_v2.py @@ -303,6 +303,7 @@ class MiMoV2Attention(nn.Module): backend_enum = requested else: fa_backend = AttentionBackendEnum.FLASH_ATTN_DIFFKV.get_class() + assert hasattr(fa_backend, "is_supported_on_current_device") if fa_backend.is_supported_on_current_device( head_size=self.head_dim, head_size_v=self.v_head_dim, @@ -312,6 +313,7 @@ class MiMoV2Attention(nn.Module): else: backend_enum = AttentionBackendEnum.TRITON_ATTN_DIFFKV attn_backend = backend_enum.get_class() + assert hasattr(attn_backend, "set_head_size_v") attn_backend.set_head_size_v(self.v_head_dim) logger.info_once("Using %s for attention.", attn_backend.get_name()) else: @@ -636,7 +638,7 @@ class MiMoV2Model(nn.Module, EagleModelMixin): ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ + stacked_params_mapping: list[tuple[str, str, str | int]] = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), ("qkv_proj", "k_proj", "k"), @@ -709,7 +711,7 @@ class MiMoV2Model(nn.Module, EagleModelMixin): ): continue stacked_matched = False - for param_name, weight_name, shard_id in stacked_params_mapping: + for param_name, weight_name, stacked_shard_id in stacked_params_mapping: if weight_name not in name: continue name_rewritten = name.replace(weight_name, param_name) @@ -728,7 +730,7 @@ class MiMoV2Model(nn.Module, EagleModelMixin): param = params_dict[name_rewritten] weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight, shard_id) + weight_loader(param, loaded_weight, stacked_shard_id) loaded_params.add(name_rewritten) stacked_matched = True diff --git a/vllm/model_executor/models/mimo_v2_mtp.py b/vllm/model_executor/models/mimo_v2_mtp.py index c863cedaeb..fa2cabe288 100644 --- a/vllm/model_executor/models/mimo_v2_mtp.py +++ b/vllm/model_executor/models/mimo_v2_mtp.py @@ -255,7 +255,7 @@ class MiMoV2MTP(nn.Module): tp_rank = get_tensor_model_parallel_rank() tp_size = get_tensor_model_parallel_world_size() - stacked_params_mapping = [ + stacked_params_mapping: list[tuple[str, str, str | int]] = [ ("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1), # Flash format: separate projections → fused qkv_proj @@ -308,8 +308,11 @@ class MiMoV2MTP(nn.Module): if name_rewritten not in params_dict: continue param = params_dict[name_rewritten] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight, shard_id) + weight_loader = getattr(param, "weight_loader", None) + if weight_loader is None: + default_weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) loaded_params.add(name_rewritten) stacked_matched = True break @@ -331,8 +334,10 @@ class MiMoV2MTP(nn.Module): 0, tp_rank * heads_per_rank, heads_per_rank ) - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) + direct_weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + direct_weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index d7b7f1e85b..93bd6a44b5 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -12,9 +12,14 @@ import torch.nn as nn import torch.nn.functional as F from transformers import BatchFeature, PretrainedConfig from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize +from typing_extensions import TypedDict from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import ( + BaseDummyOptions, + ImageDummyOptions, + VideoDummyOptions, +) from vllm.distributed import parallel_state from vllm.distributed import utils as dist_utils from vllm.inputs import MultiModalDataDict @@ -70,6 +75,18 @@ from .qwen2_vl import _create_qwen2vl_field_factory from .utils import AutoWeightsLoader, IntermediateTensors, WeightsMapper, maybe_prefix +class MiMoV2OmniAudioInputs(TypedDict): + type: str + audio_features: torch.Tensor | list[torch.Tensor] + audio_token_lens: torch.Tensor | list[int] | None + + +class MiMoV2OmniMultiModalInputs(TypedDict, total=False): + image: Qwen2_5_VLImageInputs | None + video: Qwen2_5_VLVideoInputs | None + audio: MiMoV2OmniAudioInputs | None + + class MiMoVisionMLP(Qwen2_5_VisionMLP): pass @@ -937,11 +954,14 @@ class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessing # Handle video_audio items: convert video part to (TCHW, timestamps) tuple if "video_audio" in mm_data: va_converted: list[VideoAudioInput] = [] - for va_item in mm_data["video_audio"]: + video_audio_items = mm_data["video_audio"] + assert isinstance(video_audio_items, Sequence) + for va_item in video_audio_items: if isinstance(va_item, VideoAudioInput): vid = va_item.video else: # Expect (video_frames, audio_source) tuple + assert isinstance(va_item, (list, tuple)) and len(va_item) == 2 vid, audio_src = va_item va_item = VideoAudioInput(video=vid, audio=audio_src) vid = vid @@ -976,7 +996,9 @@ class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessing if "videos" in mm_data: converted: list[tuple[torch.Tensor, torch.Tensor]] = [] - for video in mm_data["videos"]: + videos = mm_data["videos"] + assert isinstance(videos, Sequence) + for video in videos: if ( isinstance(video, tuple) and len(video) == 2 @@ -1037,30 +1059,47 @@ class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessing video_end_id = p.video_end_token_id audio_start_id = p.audio_start_token_id audio_end_id = p.audio_end_token_id + assert isinstance(vision_start_id, int) + assert isinstance(vision_end_id, int) + assert isinstance(video_start_id, int) + assert isinstance(video_end_id, int) def get_image_replacement(item_idx: int) -> PromptUpdateDetails: out_item = out_mm_kwargs["image"][item_idx] grid_thw = out_item["image_grid_thw"].data + assert isinstance(grid_thw, torch.Tensor) n_tokens = int(grid_thw.prod()) // merge_size**2 - return [image_pad_id] * n_tokens + return PromptUpdateDetails.from_seq([image_pad_id] * n_tokens) def get_video_replacement(item_idx: int) -> PromptUpdateDetails: out_item = out_mm_kwargs["video"][item_idx] grid_thw = out_item["video_grid_thw"].data - spt = float(out_item["second_per_grid_ts"].data) - start = float(out_item["video_start_times"].data) + second_per_grid_ts = out_item["second_per_grid_ts"].data + video_start_time = out_item["video_start_times"].data + assert isinstance(grid_thw, torch.Tensor) + assert isinstance(second_per_grid_ts, (int, float, torch.Tensor)) + assert isinstance(video_start_time, (int, float, torch.Tensor)) + spt = float(second_per_grid_ts) + start = float(video_start_time) T, H, W = map(int, grid_thw) n_per_grid = H * W // (merge_size * merge_size) # Check if this is a video_audio item n_segs_field = out_item.get("video_audio_n_segs") - n_segs_val = int(n_segs_field.data) if n_segs_field is not None else 0 + if n_segs_field is None: + n_segs_val = 0 + else: + n_segs_data = n_segs_field.data + assert isinstance(n_segs_data, (int, torch.Tensor)) + n_segs_val = int(n_segs_data) va_seg_lens: list[int] | None = None if n_segs_val > 0: seg_lens_field = out_item.get("video_audio_seg_lens") if seg_lens_field is not None: - va_seg_lens = seg_lens_field.data[:n_segs_val].tolist() + seg_lens_data = seg_lens_field.data + assert isinstance(seg_lens_data, torch.Tensor) + va_seg_lens = seg_lens_data[:n_segs_val].tolist() full: list[int] = [video_start_id] is_embed_mask: list[bool] = [False] @@ -1099,6 +1138,9 @@ class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessing is_embed_mask.append(False) # Audio tokens for this group seg_len = va_seg_lens[g] + assert audio_start_id is not None + assert audio_pad_id is not None + assert audio_end_id is not None full.append(audio_start_id) is_embed_mask.append(False) full.extend([audio_pad_id] * seg_len) @@ -1117,8 +1159,11 @@ class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessing def get_audio_replacement(item_idx: int) -> PromptUpdateDetails: out_item = out_mm_kwargs["audio"][item_idx] - tok_len = int(out_item["audio_token_lens"].data) - return [audio_pad_id] * tok_len + audio_token_lens = out_item["audio_token_lens"].data + assert isinstance(audio_token_lens, (int, torch.Tensor)) + assert audio_pad_id is not None + tok_len = int(audio_token_lens) + return PromptUpdateDetails.from_seq([audio_pad_id] * tok_len) updates: list[PromptUpdate] = [ PromptReplacement( @@ -1166,20 +1211,24 @@ class MiMoV2OmniDummyInputsBuilder(BaseDummyInputsBuilder[MiMoV2OmniProcessingIn target_num_frames = self.info.get_num_frames_with_most_features( seq_len, mm_counts ) + image_overrides = mm_options.get("image") + video_overrides = mm_options.get("video") + assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions) + assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions) return { "image": self._get_dummy_images( width=target_width, height=target_height, num_images=num_images, - overrides=mm_options.get("image"), + overrides=image_overrides, ), "video": self._get_dummy_videos( width=target_width, height=target_height, num_frames=target_num_frames, num_videos=num_videos, - overrides=mm_options.get("video"), + overrides=video_overrides, ), } @@ -1235,6 +1284,7 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ ) audio_config = getattr(config, "audio_config", None) model_path = vllm_config.model_config.model + self.audio_encoder: MimoAudioEncoder | None if audio_config is not None: with self._mark_tower_model(vllm_config, "audio"): self.audio_encoder = MimoAudioEncoder( @@ -1269,12 +1319,12 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ image_grid_thw=image_grid_thw, ) - if image_embeds is not None: - return Qwen2_5_VLImageEmbeddingInputs( - type="image_embeds", - image_embeds=image_embeds, - image_grid_thw=image_grid_thw, - ) + assert image_embeds is not None + return Qwen2_5_VLImageEmbeddingInputs( + type="image_embeds", + image_embeds=image_embeds, + image_grid_thw=image_grid_thw, + ) def _parse_and_validate_video_input( self, **kwargs: object @@ -1295,13 +1345,13 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ second_per_grid_ts=second_per_grid_ts, ) - if video_embeds is not None: - return Qwen2_5_VLVideoEmbeddingInputs( - type="video_embeds", - video_embeds=video_embeds, - video_grid_thw=video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - ) + assert video_embeds is not None + return Qwen2_5_VLVideoEmbeddingInputs( + type="video_embeds", + video_embeds=video_embeds, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + ) def _process_image_input( self, image_input: Qwen2_5_VLImageInputs @@ -1339,19 +1389,29 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist() return video_embeds.split(sizes) - def _parse_and_validate_audio_input(self, **kwargs: object) -> dict | None: + def _parse_and_validate_audio_input( + self, **kwargs: object + ) -> MiMoV2OmniAudioInputs | None: audio_features = kwargs.pop("audio_features", None) audio_token_lens = kwargs.pop("audio_token_lens", None) if audio_features is None: return None + assert isinstance(audio_features, (torch.Tensor, list)) + if isinstance(audio_features, list): + assert all(isinstance(item, torch.Tensor) for item in audio_features) + assert audio_token_lens is None or isinstance( + audio_token_lens, (torch.Tensor, list) + ) return { "type": "audio", "audio_features": audio_features, "audio_token_lens": audio_token_lens, } - def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: - mm_input_by_modality = {} + def _parse_and_validate_multimodal_inputs( + self, **kwargs: object + ) -> MiMoV2OmniMultiModalInputs: + mm_input_by_modality: MiMoV2OmniMultiModalInputs = {} # Preserve the order of modalities if there are multiple of them # from the order of kwargs. @@ -1376,7 +1436,9 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ ) return mm_input_by_modality - def _process_audio_input(self, audio_input: dict) -> tuple[torch.Tensor, ...]: + def _process_audio_input( + self, audio_input: MiMoV2OmniAudioInputs + ) -> tuple[torch.Tensor, ...]: mel_specs = audio_input["audio_features"] if self.audio_encoder is None: return () @@ -1407,38 +1469,52 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ # Pre-process va audio: one mel spec per va video → per-video audio embeddings # keyed by va video index (0-based among va videos only) - va_audio_embs_list: list[tuple[torch.Tensor, ...]] = [] + va_audio_embs_list: list[torch.Tensor] = [] if va_audio_features is not None and self.audio_encoder is not None: - mel_list = ( - list(va_audio_features) - if isinstance(va_audio_features, torch.Tensor) - else list(va_audio_features) - ) + assert isinstance(va_audio_features, (torch.Tensor, list, tuple)) + mel_list: list[torch.Tensor] = [] + for mel_spec in va_audio_features: + assert isinstance(mel_spec, torch.Tensor) + mel_list.append(mel_spec) for mel_spec in mel_list: embs, tok_lens = self.audio_encoder.get_audio_feature([mel_spec]) # tok_lens is a list/tensor with one entry (total tokens for this mel) va_audio_embs_list.append(embs) # shape (total_tok, hidden) va_cursor = 0 # index into va_audio_embs_list + assert video_audio_n_segs is None or isinstance( + video_audio_n_segs, torch.Tensor + ) + assert video_audio_seg_lens is None or isinstance( + video_audio_seg_lens, torch.Tensor + ) # NOTE: Iterate in dict insertion order to preserve token sequence order. for modality in mm_input_by_modality: - multimodal_input = mm_input_by_modality[modality] if modality == "image": - multimodal_embeddings.extend( - self._process_image_input(multimodal_input) + image_input = mm_input_by_modality["image"] + assert isinstance( + image_input, + (Qwen2_5_VLImagePixelInputs, Qwen2_5_VLImageEmbeddingInputs), ) + multimodal_embeddings.extend(self._process_image_input(image_input)) elif modality == "video": - video_embs_tuple = self._process_video_input(multimodal_input) + video_input = mm_input_by_modality["video"] + assert isinstance( + video_input, + (Qwen2_5_VLVideoPixelInputs, Qwen2_5_VLVideoEmbeddingInputs), + ) + video_embs_tuple = self._process_video_input(video_input) if video_audio_n_segs is None: multimodal_embeddings.extend(video_embs_tuple) else: - grid_thw = multimodal_input["video_grid_thw"] + grid_thw = video_input["video_grid_thw"] for i, vid_embs in enumerate(video_embs_tuple): n_segs = int(video_audio_n_segs[i]) if n_segs == 0 or not va_audio_embs_list: multimodal_embeddings.append(vid_embs) else: + assert video_audio_seg_lens is not None T = int(grid_thw[i][0]) n_per_grid = vid_embs.shape[0] // T frames = list(vid_embs.split(n_per_grid, dim=0)) @@ -1459,9 +1535,9 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ ) multimodal_embeddings.append(group_audio_embs[g]) elif modality == "audio": - multimodal_embeddings.extend( - self._process_audio_input(multimodal_input) - ) + audio_input = mm_input_by_modality["audio"] + assert audio_input is not None + multimodal_embeddings.extend(self._process_audio_input(audio_input)) return tuple(multimodal_embeddings) def forward( diff --git a/vllm/model_executor/models/minicpm.py b/vllm/model_executor/models/minicpm.py index 699e003592..0d7da44c3b 100644 --- a/vllm/model_executor/models/minicpm.py +++ b/vllm/model_executor/models/minicpm.py @@ -456,6 +456,7 @@ class MiniCPMModel(nn.Module, EagleModelMixin): hidden_states = self.embed_input_ids(input_ids) residual = None else: + assert intermediate_tensors is not None hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] diff --git a/vllm/model_executor/models/minicpm_eagle.py b/vllm/model_executor/models/minicpm_eagle.py index 89ce01d611..a495fed22c 100644 --- a/vllm/model_executor/models/minicpm_eagle.py +++ b/vllm/model_executor/models/minicpm_eagle.py @@ -26,6 +26,7 @@ import math from collections.abc import Iterable +from typing import TYPE_CHECKING import torch from torch import nn @@ -55,6 +56,14 @@ from .utils import ( process_eagle_weight, ) +if TYPE_CHECKING: + + class _EagleMiniCPMSupportsPP: + pass + +else: + _EagleMiniCPMSupportsPP = SupportsPP + class EagleMiniCPMDecoderLayer(nn.Module): def __init__( @@ -149,7 +158,9 @@ class EagleMiniCPMModel(nn.Module): ): super().__init__() - config = vllm_config.speculative_config.draft_model_config.hf_config + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + config = speculative_config.draft_model_config.hf_config cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config @@ -291,7 +302,9 @@ class EagleMiniCPMModel(nn.Module): return loaded_params -class EagleMiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle): +class EagleMiniCPMForCausalLM( + nn.Module, SupportsLoRA, _EagleMiniCPMSupportsPP, SupportsEagle +): packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -312,7 +325,9 @@ class EagleMiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() - config = vllm_config.speculative_config.draft_model_config.hf_config + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + config = speculative_config.draft_model_config.hf_config cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config diff --git a/vllm/model_executor/models/minicpmo.py b/vllm/model_executor/models/minicpmo.py index d4aaeac6ad..c8b46a45f1 100644 --- a/vllm/model_executor/models/minicpmo.py +++ b/vllm/model_executor/models/minicpmo.py @@ -40,7 +40,7 @@ from transformers.models.whisper.modeling_whisper import ( ) from vllm.config import VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions from vllm.inputs import ModalityData, MultiModalDataDict from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems from vllm.multimodal.inputs import ( @@ -66,8 +66,12 @@ from .minicpmv import ( _MAX_FRAMES_PER_VIDEO, MiniCPMV2_6, MiniCPMV4_5, + MiniCPMVBaseModel, MiniCPMVDummyInputsBuilder, + MiniCPMVImageEmbeddingInputs, + MiniCPMVImagePixelInputs, MiniCPMVMultiModalDataParser, + MiniCPMVMultiModalInputs, MiniCPMVMultiModalProcessor, MiniCPMVProcessingInfo, _minicpmv_field_config, @@ -181,6 +185,10 @@ MiniCPMOAudioInputs: TypeAlias = ( ) +class MiniCPMOMultiModalInputs(MiniCPMVMultiModalInputs, total=False): + audios: MiniCPMOAudioInputs | None + + def _minicpmo_field_config(hf_inputs: Mapping[str, torch.Tensor]): audio_features = hf_inputs.get("audio_features") audio_feature_lens = hf_inputs.get("audio_feature_lens") @@ -405,6 +413,7 @@ class MiniCPMODummyInputsBuilder(MiniCPMVDummyInputsBuilder[MiniCPMOProcessingIn ) audio_overrides = mm_options.get("audio") + assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) audio_mm_data = { "audio": self._get_dummy_audios( @@ -469,7 +478,8 @@ class MiniCPMOMultiModalProcessor(MiniCPMVMultiModalProcessor[MiniCPMOProcessing if isinstance(lens, torch.Tensor): flat_feature_lens.extend(lens.flatten().tolist()) else: - flat_feature_lens.append(int(lens)) + assert isinstance(lens, int) + flat_feature_lens.append(lens) unpadded_audio_features = [ feat[:, :length] for feat, length in zip( @@ -515,6 +525,9 @@ class MiniCPMOMultiModalProcessor(MiniCPMVMultiModalProcessor[MiniCPMOProcessing audios = mm_items.get_items( "audio", (MiniCPMOAudioEmbeddingItems, AudioProcessorItems) ) + assert isinstance( + audios, (MiniCPMOAudioEmbeddingItems, AudioProcessorItems) + ) if isinstance(audios, MiniCPMOAudioEmbeddingItems): single_audio_embeds = audios.get(item_idx)["audio_embeds"] @@ -652,7 +665,7 @@ class MiniCPMWhisperEncoder(WhisperEncoder): hidden_states, p=self.dropout, training=self.training ) - encoder_states = () + encoder_states: tuple[torch.Tensor, ...] = () for idx, encoder_layer in enumerate(self.layers): encoder_states = encoder_states + (hidden_states,) @@ -682,7 +695,13 @@ class MiniCPMWhisperEncoder(WhisperEncoder): ) -class MiniCPMOBaseModel: +if TYPE_CHECKING: + _MiniCPMOBaseModelBase = MiniCPMVBaseModel +else: + _MiniCPMOBaseModelBase = object + + +class MiniCPMOBaseModel(_MiniCPMOBaseModelBase): """Base mixin class for MiniCPM-O models with audio support.""" # Unlike the vision-only MiniCPM-V models, audio weights are loaded here. @@ -897,8 +916,12 @@ class MiniCPMOBaseModel: audio_feature_lens=audio_feature_lens, ) - def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: - modalities = super()._parse_and_validate_multimodal_inputs(**kwargs) + def _parse_and_validate_multimodal_inputs( + self, **kwargs: object + ) -> MiniCPMOMultiModalInputs: + modalities = MiniCPMOMultiModalInputs( + **super()._parse_and_validate_multimodal_inputs(**kwargs) + ) # Preserve the order of modalities if there are multiple of them # from the order of kwargs. @@ -918,16 +941,33 @@ class MiniCPMOBaseModel: if audio_input["type"] == "audio_embeds": return audio_input["audio_embeds"] + assert isinstance(audio_input, MiniCPMOAudioFeatureInputs) return self.get_audio_hidden_states(audio_input) - def _process_multimodal_inputs(self, modalities: dict): - multimodal_embeddings = super()._process_multimodal_inputs(modalities) + def _process_multimodal_inputs(self, modalities: Mapping[str, object]): + base_modalities: MiniCPMVMultiModalInputs = {} + image_input = modalities.get("images") + video_input = modalities.get("videos") + assert image_input is None or isinstance( + image_input, (MiniCPMVImagePixelInputs, MiniCPMVImageEmbeddingInputs) + ) + assert video_input is None or isinstance( + video_input, (MiniCPMVImagePixelInputs, MiniCPMVImageEmbeddingInputs) + ) + if "images" in modalities: + base_modalities["images"] = image_input + if "videos" in modalities: + base_modalities["videos"] = video_input + multimodal_embeddings = super()._process_multimodal_inputs(base_modalities) - for modality in modalities: - if modality == "audios": - audio_input = modalities["audios"] - audio_embeddings = self._process_audio_input(audio_input) - multimodal_embeddings += tuple(audio_embeddings) + audio_input = modalities.get("audios") + if audio_input is not None: + assert isinstance( + audio_input, + (MiniCPMOAudioFeatureInputs, MiniCPMOAudioEmbeddingInputs), + ) + audio_embeddings = self._process_audio_input(audio_input) + multimodal_embeddings += tuple(audio_embeddings) return multimodal_embeddings @@ -983,7 +1023,7 @@ class MiniCPMO(MiniCPMOBaseModel, MiniCPMV2_6): try: version_str = str(config.version) version_parts = version_str.split(".") - version = tuple(int(x) for x in version_parts[:2]) + version_values = tuple(int(x) for x in version_parts[:2]) except (ValueError, TypeError) as e: raise ValueError( f"Invalid model version format in config: {config.version}. " @@ -991,17 +1031,21 @@ class MiniCPMO(MiniCPMOBaseModel, MiniCPMV2_6): ) from e else: # Default to 2.6 for backward compatibility - version = (2, 6) + version_values = (2, 6) # Dispatch class based on version - instance_cls = _MINICPMO_SUPPORT_VERSION.get(version) + if len(version_values) == 2: + version_key = (version_values[0], version_values[1]) + instance_cls = _MINICPMO_SUPPORT_VERSION.get(version_key) + else: + instance_cls = None if instance_cls is None: supported_versions = ", ".join( [f"{v[0]}.{v[1]}" for v in sorted(_MINICPMO_SUPPORT_VERSION.keys())] ) raise ValueError( f"Currently, MiniCPMO only supports versions " - f"{supported_versions}. Got version: {version}" + f"{supported_versions}. Got version: {version_values}" ) return instance_cls(vllm_config=vllm_config, prefix=prefix) diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index 91a74e581a..f8602983e5 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -41,7 +41,7 @@ from transformers.dynamic_module_utils import ( get_class_from_dynamic_module, resolve_trust_remote_code, ) -from typing_extensions import TypeVar +from typing_extensions import TypedDict, TypeVar from vllm.config import VllmConfig from vllm.config.multimodal import ( @@ -97,7 +97,7 @@ from vllm.transformers_utils.processor import ( cached_get_image_processor, ) from vllm.transformers_utils.utils import convert_model_repo_to_path -from vllm.utils.collection_utils import flatten_2d_lists +from vllm.utils.collection_utils import flatten_2d_lists, is_list_of from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.utils.torch_utils import set_default_torch_dtype @@ -386,6 +386,7 @@ class Resampler4_5(Resampler2_5): for i in range(bs): tgt_h, tgt_w = tgt_sizes[i] if temporal_pos_emb: + assert temporal_ids_flatten is not None if temporal_ids_flatten[i] == -1: pos_embed_temporal.append( torch.zeros(self.embed_dim, dtype=dtype, device=device) @@ -549,6 +550,11 @@ class MiniCPMVVideoEmbeddingItems(DictEmbeddingItems): return len(self.get(index)["video_image_sizes"]) +class MiniCPMVMultiModalInputs(TypedDict, total=False): + images: MiniCPMVImageInputs | None + videos: MiniCPMVImageInputs | None + + class MiniCPMVMultiModalDataParser(MultiModalDataParser): # The per-item sizes are what size the placeholder range. embedding_fields = { @@ -833,6 +839,7 @@ class MiniCPMVDummyInputsBuilder(BaseDummyInputsBuilder[_I]): image_overrides = mm_options.get("image") video_overrides = mm_options.get("video") + assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions) # Convert video overrides to image overrides for per-frame image generation, # and apply num_frames override to num_video_frames. @@ -842,6 +849,7 @@ class MiniCPMVDummyInputsBuilder(BaseDummyInputsBuilder[_I]): num_video_frames = min(num_video_frames, video_overrides.num_frames) if video_overrides.width or video_overrides.height: video_frame_overrides = ImageDummyOptions( + count=video_overrides.count, width=video_overrides.width, height=video_overrides.height, ) @@ -1081,6 +1089,9 @@ class MiniCPMVMultiModalProcessor(BaseMultiModalProcessor[_I]): images = mm_items.get_items( "image", (MiniCPMVImageEmbeddingItems, ImageProcessorItems) ) + assert isinstance( + images, (MiniCPMVImageEmbeddingItems, ImageProcessorItems) + ) image_size = images.get_image_size(item_idx) @@ -1097,6 +1108,9 @@ class MiniCPMVMultiModalProcessor(BaseMultiModalProcessor[_I]): videos = mm_items.get_items( "video", (MiniCPMVVideoEmbeddingItems, VideoProcessorItems) ) + assert isinstance( + videos, (MiniCPMVVideoEmbeddingItems, VideoProcessorItems) + ) frame_size = videos.get_frame_size(item_idx) num_frames = videos.get_num_frames(item_idx) @@ -1199,6 +1213,7 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config multimodal_config = vllm_config.model_config.multimodal_config + assert multimodal_config is not None quant_config = vllm_config.quant_config self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" super().__init__() @@ -1261,10 +1276,15 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP): image_embeds=image_embeds, ) + assert is_list_of(pixel_values, list, check="all") + pixel_values_list: list[list[torch.Tensor]] = [] + for pixel_value in pixel_values: + assert is_list_of(pixel_value, torch.Tensor, check="all") + pixel_values_list.append(pixel_value) tgt_sizes = kwargs.pop("tgt_sizes") - num_slices_flat = torch.tensor([len(ps) for ps in pixel_values]) - pixel_values_flat = flatten_bn(pixel_values) + num_slices_flat = torch.tensor([len(ps) for ps in pixel_values_list]) + pixel_values_flat = flatten_2d_lists(pixel_values_list) tgt_sizes_flat = flatten_bn(tgt_sizes, concat=True) return MiniCPMVImagePixelInputs( @@ -1274,8 +1294,11 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP): num_slices=num_slices_flat, ) - def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: - modalities = {} + def _parse_and_validate_multimodal_inputs( + self, **kwargs: object + ) -> MiniCPMVMultiModalInputs: + kwargs.pop("modality", None) + modalities: MiniCPMVMultiModalInputs = {} # Preserve the order of modalities if there are multiple of them # from the order of kwargs. @@ -1304,12 +1327,13 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP): if image_input["type"] == "image_embeds": return image_input["image_embeds"] + assert isinstance(image_input, MiniCPMVImagePixelInputs) image_features_flat = self.get_vision_hidden_states(image_input) num_slices = image_input["num_slices"] return [e.flatten(0, 1) for e in image_features_flat.split(num_slices.tolist())] - def _process_multimodal_inputs(self, modalities: dict): + def _process_multimodal_inputs(self, modalities: MiniCPMVMultiModalInputs): # The result multimodal_embeddings is tuple of tensors, with each # tensor corresponding to a multimodal data item (image or video). multimodal_embeddings: tuple[torch.Tensor, ...] = () @@ -1319,10 +1343,12 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP): for modality in modalities: if modality == "images": image_input = modalities["images"] + assert image_input is not None image_embeddings = self._process_vision_input(image_input) multimodal_embeddings += tuple(image_embeddings) if modality == "videos": video_input = modalities["videos"] + assert video_input is not None video_embeddings = self._process_vision_input(video_input) multimodal_embeddings += tuple(video_embeddings) @@ -1915,23 +1941,27 @@ class MiniCPMV(MiniCPMVBaseModel, SupportsMultiModal, SupportsLoRA): def __new__(cls, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_config + version_values: tuple[int, ...] if not hasattr(config, "version"): if config.hidden_size == 2304 and config.query_num == 64: - version = (2, 0) + version_values = (2, 0) else: - version = (2, 5) + version_values = (2, 5) else: - version = str(config.version).split(".") - version = tuple([int(x) for x in version]) + version_values = tuple(int(x) for x in str(config.version).split(".")) # Dispatch class based on version - instance_cls = _SUPPORT_VERSION.get(version) + if len(version_values) == 2: + version_key = (version_values[0], version_values[1]) + instance_cls = _SUPPORT_VERSION.get(version_key) + else: + instance_cls = None if instance_cls is None: supported_versions = ", ".join( [f"{v[0]}.{v[1]}" for v in sorted(_SUPPORT_VERSION.keys())] ) raise ValueError( f"Currently, MiniCPMV only supports versions " - f"{supported_versions}. Got version: {version}" + f"{supported_versions}. Got version: {version_values}" ) # quant_config references base class members, diff --git a/vllm/model_executor/models/minicpmv4_6.py b/vllm/model_executor/models/minicpmv4_6.py index 716e5dd1f5..e43f13f79f 100644 --- a/vllm/model_executor/models/minicpmv4_6.py +++ b/vllm/model_executor/models/minicpmv4_6.py @@ -98,7 +98,9 @@ class MiniCPMV4_6MultiModalProcessor(MiniCPMVMultiModalProcessor): ds = mm_kwargs.get("downsample_mode") if ds is not None: return str(ds) - return self.info._get_downsample_mode() + info = self.info + assert isinstance(info, MiniCPMV4_6ProcessingInfo) + return info._get_downsample_mode() def get_image_prompt_texts( self, @@ -106,7 +108,9 @@ class MiniCPMV4_6MultiModalProcessor(MiniCPMVMultiModalProcessor): image_idx: int = 0, downsample_mode: str | None = None, ) -> str: - return self.info.get_slice_image_placeholder( + info = self.info + assert isinstance(info, MiniCPMV4_6ProcessingInfo) + return info.get_slice_image_placeholder( image_size, image_idx=image_idx, downsample_mode=downsample_mode, @@ -128,6 +132,7 @@ class MiniCPMV4_6MultiModalProcessor(MiniCPMVMultiModalProcessor): # which one is used. Using image_token for video silently produces # garbage descriptions. info = self.info + assert isinstance(info, MiniCPMV4_6ProcessingInfo) grids, source_tokens, patch_tokens = info._compute_visual_tokens( image_size, max_slice_nums=info.get_video_max_slice_num(), @@ -1016,6 +1021,7 @@ class MiniCPMV4_6ForConditionalGeneration( config: MiniCPMV4_6Config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config multimodal_config = vllm_config.model_config.multimodal_config + assert multimodal_config is not None self.config = config self.multimodal_config = multimodal_config @@ -1078,6 +1084,7 @@ class MiniCPMV4_6ForConditionalGeneration( image_embeds=image_embeds, ) + assert isinstance(pixel_values, torch.Tensor | list) tgt_sizes = kwargs.pop("tgt_sizes") num_slices_flat = torch.tensor([len(ps) for ps in pixel_values]) pixel_values_flat = flatten_bn(pixel_values) diff --git a/vllm/model_executor/models/mistral.py b/vllm/model_executor/models/mistral.py index 86c4b171e5..d2d6e0f39a 100644 --- a/vllm/model_executor/models/mistral.py +++ b/vllm/model_executor/models/mistral.py @@ -221,15 +221,21 @@ class MistralModel(LlamaModel): intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, t_cond: torch.Tensor | None = None, + **extra_layer_kwargs: object, ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]: return super().forward( - input_ids, positions, intermediate_tensors, inputs_embeds, t_cond=t_cond + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + t_cond=t_cond, + **extra_layer_kwargs, ) class MistralForCausalLM(LlamaForCausalLM): # Mistral: We don't support LoRA on the embedding layers. - embedding_modules: dict[str, str] = {} + embedding_modules = {} # Mistral/Llama models can also be loaded with --load-format mistral # from consolidated.safetensors checkpoints diff --git a/vllm/model_executor/models/mistral3.py b/vllm/model_executor/models/mistral3.py index 23c6323619..3463e59e76 100644 --- a/vllm/model_executor/models/mistral3.py +++ b/vllm/model_executor/models/mistral3.py @@ -196,6 +196,7 @@ class Mistral3HFEncoderInfo(PixtralHFEncoderInfo): image_height = math.floor(image_height / ratio) patch_size = self.vision_config.patch_size + assert isinstance(self.hf_config, Mistral3Config) spatial_merge_size = self.hf_config.spatial_merge_size # The HF processor rounds each dimension up to the vision patch size diff --git a/vllm/model_executor/models/mistral_eagle.py b/vllm/model_executor/models/mistral_eagle.py index 6a85f970ae..04b37bb4d2 100644 --- a/vllm/model_executor/models/mistral_eagle.py +++ b/vllm/model_executor/models/mistral_eagle.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable +from typing import TYPE_CHECKING import torch import torch.nn as nn @@ -33,6 +34,20 @@ from vllm.model_executor.models.utils import ( logger = init_logger(__name__) +if TYPE_CHECKING: + + class _EagleMistralModelBase(nn.Module): + pass + + class _EagleMistralForCausalLMBase(nn.Module): + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + +else: + _EagleMistralModelBase = MistralModel + _EagleMistralForCausalLMBase = MistralForCausalLM + + class EagleMistralDecoderLayer(MistralDecoderLayer): def __init__( self, @@ -47,7 +62,7 @@ class EagleMistralDecoderLayer(MistralDecoderLayer): @support_torch_compile -class EagleMistralModel(MistralModel): +class EagleMistralModel(_EagleMistralModelBase): def __init__( self, *, @@ -58,7 +73,9 @@ class EagleMistralModel(MistralModel): # Bypass MistralModel.__init__ to avoid creating duplicate attention # layer entries in the global context. nn.Module.__init__(self) - self.config = vllm_config.speculative_config.draft_model_config.hf_config + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + self.config = speculative_config.draft_model_config.hf_config self.vocab_size = self.config.vocab_size # Get drafter's quantization config self.quant_config = get_draft_quant_config(vllm_config) @@ -112,7 +129,9 @@ class EagleMistralModel(MistralModel): return hidden_states, hidden_states -class EagleMistralForCausalLM(MistralForCausalLM, SupportsMultiModalEmbeddings): +class EagleMistralForCausalLM( + _EagleMistralForCausalLMBase, SupportsMultiModalEmbeddings +): mistral_mapping = MistralForCausalLM.mistral_mapping | { "eagle_linear": "model.fc", } @@ -121,7 +140,9 @@ class EagleMistralForCausalLM(MistralForCausalLM, SupportsMultiModalEmbeddings): # Bypass MistralForCausalLM.__init__ to use the draft model config # and to avoid creating an lm_head. nn.Module.__init__(self) - self.config = vllm_config.speculative_config.draft_model_config.hf_config + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + self.config = speculative_config.draft_model_config.hf_config target_layer_num = vllm_config.model_config.get_num_layers( vllm_config.parallel_config ) diff --git a/vllm/model_executor/models/mistral_large_3_eagle.py b/vllm/model_executor/models/mistral_large_3_eagle.py index d194536beb..d2e263f3e7 100644 --- a/vllm/model_executor/models/mistral_large_3_eagle.py +++ b/vllm/model_executor/models/mistral_large_3_eagle.py @@ -4,6 +4,7 @@ import copy from collections.abc import Iterable from functools import partial +from typing import TYPE_CHECKING import regex import torch @@ -28,6 +29,15 @@ from .utils import WeightsMapper, make_empty_intermediate_tensors_factory, maybe logger = init_logger(__name__) +if TYPE_CHECKING: + + class _EagleMistralLarge3ForCausalLMBase(nn.Module): + pass + +else: + _EagleMistralLarge3ForCausalLMBase = MistralLarge3ForCausalLM + + @support_torch_compile class EagleMistralLarge3Model(DeepseekV2Model): def __init__( @@ -108,7 +118,7 @@ class EagleMistralLarge3Model(DeepseekV2Model): class EagleMistralLarge3ForCausalLM( - MistralLarge3ForCausalLM, SupportsMultiModalEmbeddings + _EagleMistralLarge3ForCausalLMBase, SupportsMultiModalEmbeddings ): hf_to_vllm_mapper = MistralLarge3ForCausalLM.hf_to_vllm_mapper | WeightsMapper( orig_to_new_regex={ @@ -119,13 +129,15 @@ class EagleMistralLarge3ForCausalLM( ) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + speculative_config = vllm_config.speculative_config + assert speculative_config is not None target_layer_num = vllm_config.model_config.get_num_layers( vllm_config.parallel_config ) - vllm_config.model_config = vllm_config.speculative_config.draft_model_config + vllm_config.model_config = speculative_config.draft_model_config # draft model quantization config may differ from target model self.quant_config = VllmConfig.get_quantization_config( - vllm_config.speculative_config.draft_model_config, vllm_config.load_config + speculative_config.draft_model_config, vllm_config.load_config ) vllm_config.quant_config = self.quant_config self.model_cls = partial( diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index b9edbad3d0..6b40cd0c86 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -620,14 +620,16 @@ class Mllama4MultiModalProcessor(BaseMultiModalProcessor[Mllama4ProcessingInfo]) max_num_chunks=self.info.get_max_num_tiles(), patch_size=SizeDict(height=tile_size, width=tile_size), ) - best_fit_sizes = [ - get_best_fit( - (image.size[1], image.size[0]), - torch.tensor(possible_resolutions), - resize_to_max_canvas=image_processor.resize_to_max_canvas, + best_fit_sizes = [] + for image in parsed_images: + assert image is not None + best_fit_sizes.append( + get_best_fit( + (image.size[1], image.size[0]), + torch.tensor(possible_resolutions), + resize_to_max_canvas=image_processor.resize_to_max_canvas, + ) ) - for image in parsed_images - ] # TODO tile height/width do not necessarily need to match aspect_ratios = [ (image_size[0] // tile_size, image_size[1] // tile_size) @@ -760,6 +762,7 @@ class Llama4ForConditionalGeneration( config = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config multimodal_config = vllm_config.model_config.multimodal_config + assert multimodal_config is not None self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" self.vllm_config = vllm_config diff --git a/vllm/model_executor/models/module_mapping.py b/vllm/model_executor/models/module_mapping.py index 5cb2461751..42efb0de8e 100644 --- a/vllm/model_executor/models/module_mapping.py +++ b/vllm/model_executor/models/module_mapping.py @@ -17,10 +17,10 @@ class MultiModelKeys: @staticmethod def from_string_field( - language_model: str | list[str] = None, - connector: str | list[str] = None, - tower_model: str | list[str] = None, - generator: str | list[str] = None, + language_model: str | list[str] | None = None, + connector: str | list[str] | None = None, + tower_model: str | list[str] | None = None, + generator: str | list[str] | None = None, **kwargs, ) -> "MultiModelKeys": def to_list(value): diff --git a/vllm/model_executor/models/molmo.py b/vllm/model_executor/models/molmo.py index c00ca71982..502b8a5ce6 100644 --- a/vllm/model_executor/models/molmo.py +++ b/vllm/model_executor/models/molmo.py @@ -332,7 +332,7 @@ class VisionTransformer(nn.Module): ): super().__init__() scale = config.image_emb_dim**-0.5 - self.patch_num = config.image_num_patch + self.patch_num: tuple[int, int] = config.image_num_patch self.class_embedding = nn.Parameter(torch.randn(config.image_emb_dim) * scale) self.num_prefix_tokens: int = NUM_PREFIX_TOKENS self.positional_embedding = nn.Parameter( @@ -349,7 +349,7 @@ class VisionTransformer(nn.Module): config, quant_config, prefix=f"{prefix}.transformer" ) - def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor: + def add_pos_emb(self, x: torch.Tensor, patch_num: tuple[int, int]) -> torch.Tensor: cls_emb = self.positional_embedding[0:1] pos_emb = self.positional_embedding[1:] @@ -380,7 +380,7 @@ class VisionTransformer(nn.Module): return x def forward( - self, x: torch.Tensor, patch_num: int | None = None + self, x: torch.Tensor, patch_num: tuple[int, int] | None = None ) -> list[torch.Tensor]: """ : param x: (batch_size, num_patch, n_pixels) @@ -484,12 +484,15 @@ class MolmoAttention(nn.Module): def _apply_qk_norm( self, q: torch.Tensor, k: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: + assert self.q_norm is not None + assert self.k_norm is not None if self.tp_size > 1: q = tensor_model_parallel_all_gather(q.contiguous()) k = tensor_model_parallel_all_gather(k.contiguous()) q = self.q_norm(q) k = self.k_norm(k) if self.tp_size > 1: + assert self.tp_rank is not None splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size) q = splitter(q)[self.tp_rank] k = splitter(k)[self.tp_rank] @@ -1345,7 +1348,7 @@ class MolmoForCausalLM( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) - self.img_patch_id = None + self.img_patch_id: int | None = None self.lm_head = ParallelLMHead( config.embedding_size or config.vocab_size, diff --git a/vllm/model_executor/models/molmo2.py b/vllm/model_executor/models/molmo2.py index 0306529fe3..66de399294 100644 --- a/vllm/model_executor/models/molmo2.py +++ b/vllm/model_executor/models/molmo2.py @@ -22,6 +22,7 @@ from transformers import ( ) from transformers.image_utils import ImageInput from transformers.video_utils import VideoMetadata +from typing_extensions import TypedDict from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig @@ -157,6 +158,11 @@ class Molmo2VideoInputs(TensorSchema): num_video_tokens: Annotated[torch.Tensor, TensorShape("nv")] +class Molmo2MultiModalInputs(TypedDict, total=False): + images: Molmo2ImageInputs | None + videos: Molmo2VideoInputs | None + + @dataclass class VitConfig: """Config for a vision transformer""" @@ -474,7 +480,7 @@ class Molmo2VisionTransformer(nn.Module): super().__init__() scale = config.hidden_size**-0.5 self.num_prefix_tokens: int = 0 # no class embeddings - self.patch_num = config.image_num_patch + self.patch_num: tuple[int, int] = config.image_num_patch self.positional_embedding = nn.Parameter( torch.randn(config.image_num_pos, config.hidden_size) * scale, ) @@ -490,7 +496,7 @@ class Molmo2VisionTransformer(nn.Module): prefix=maybe_prefix(prefix, "transformer"), ) - def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor: + def add_pos_emb(self, x: torch.Tensor, patch_num: tuple[int, int]) -> torch.Tensor: pos_emb = self.positional_embedding pos_emb = pos_emb.reshape( @@ -522,7 +528,7 @@ class Molmo2VisionTransformer(nn.Module): def forward( self, x: torch.Tensor, - patch_num: int | None = None, + patch_num: tuple[int, int] | None = None, ) -> list[torch.Tensor]: """ : param x: (batch_size, num_patch, n_pixels) @@ -652,6 +658,7 @@ class ImagePoolingAttention(nn.Module): if self.use_pytorch_sdpa: output = self.forward_sdpa(xq, xk, xv, attn_mask) else: + assert self.attn is not None output = self.attn(xq, xk, xv) output, _ = self.o_proj(output) @@ -954,12 +961,15 @@ class Molmo2Attention(nn.Module): q: torch.Tensor, k: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: + assert self.q_norm is not None + assert self.k_norm is not None if self.tp_size > 1: q = tensor_model_parallel_all_gather(q.contiguous()) k = tensor_model_parallel_all_gather(k.contiguous()) q = self.q_norm(q) k = self.k_norm(k) if self.tp_size > 1: + assert self.tp_rank is not None splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size) q = splitter(q)[self.tp_rank] k = splitter(k)[self.tp_rank] @@ -1911,10 +1921,10 @@ class Molmo2DummyInputsBuilder(BaseDummyInputsBuilder[Molmo2ProcessingInfo]): ) videos = [v.copy() for v in videos] - video_items = [] + video_items: list[VideoItem] = [] for video in videos: video_num_frames = video.shape[0] - video_metadata = { + video_metadata: dict[str, int | float | str | bool | list[int]] = { "fps": 2.0, "duration": video_num_frames / 2.0, "total_num_frames": video_num_frames, @@ -1972,6 +1982,7 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]): image_processor = hf_processor.image_processor if videos := mm_data.pop("videos", []): + assert isinstance(videos, Sequence) bos_token_id = tokenizer.bos_token_id or tokenizer.eos_token_id pixel_values_videos_lst = [] @@ -1983,7 +1994,9 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]): num_video_tokens_lst = [] for item in videos: + assert isinstance(item, tuple) and len(item) == 2 video_array, metadata = item + assert isinstance(metadata, Mapping) # NOTE: metadata.frames_indices indicates # the sampled frames indices of pre-sampled videos, which is @@ -2240,6 +2253,7 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]): def get_video_replacement_molmo2(item_idx: int): video, metadata = mm_items["video"][item_idx] do_sample_frames = hf_processor_mm_kwargs.get("do_sample_frames") + assert do_sample_frames is None or isinstance(do_sample_frames, bool) timestamps = self.info._get_video_second_idx(metadata, do_sample_frames) ncols, nrows = self.info.get_base_grid_size(video_processor) @@ -2408,10 +2422,17 @@ class Molmo2ForConditionalGeneration( image_tokens = kwargs.pop("image_tokens", None) num_image_tokens = kwargs.pop("num_image_tokens", None) + assert isinstance(token_pooling, torch.Tensor) + assert isinstance(num_pooled_patches, torch.Tensor) + assert isinstance(num_patches, torch.Tensor) + assert isinstance(image_tokens, torch.Tensor) + assert isinstance(num_image_tokens, torch.Tensor) + accum_patches = [0] + num_patches.cumsum(dim=0)[:-1].tolist() patch_offset = 0 new_token_pooling = token_pooling.clone() - for i, n in enumerate(num_pooled_patches): + for i, n_tensor in enumerate(num_pooled_patches): + n = int(n_tensor) cur_slice = token_pooling[patch_offset : patch_offset + n] index_offset = int(accum_patches[i]) new_token_pooling[patch_offset : patch_offset + n] = torch.where( @@ -2443,10 +2464,17 @@ class Molmo2ForConditionalGeneration( video_tokens = kwargs.pop("video_tokens", None) num_video_tokens = kwargs.pop("num_video_tokens", None) + assert isinstance(token_pooling, torch.Tensor) + assert isinstance(num_pooled_patches, torch.Tensor) + assert isinstance(num_patches, torch.Tensor) + assert isinstance(video_tokens, torch.Tensor) + assert isinstance(num_video_tokens, torch.Tensor) + accum_patches = [0] + num_patches.cumsum(dim=0)[:-1].tolist() patch_offset = 0 new_token_pooling = token_pooling.clone() - for i, n in enumerate(num_pooled_patches): + for i, n_tensor in enumerate(num_pooled_patches): + n = int(n_tensor) cur_slice = token_pooling[patch_offset : patch_offset + n] index_offset = int(accum_patches[i]) new_token_pooling[patch_offset : patch_offset + n] = torch.where( @@ -2464,8 +2492,10 @@ class Molmo2ForConditionalGeneration( num_video_tokens=num_video_tokens, ) - def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: - modalities = {} + def _parse_and_validate_multimodal_inputs( + self, **kwargs: object + ) -> Molmo2MultiModalInputs: + modalities: Molmo2MultiModalInputs = {} for input_key in kwargs: if input_key in ("pixel_values",) and "images" not in modalities: @@ -2544,10 +2574,12 @@ class Molmo2ForConditionalGeneration( for modality in modalities: if modality == "images": image_input = modalities["images"] + assert image_input is not None image_embeddings = self._process_image_input(image_input) multimodal_embeddings += image_embeddings if modality == "videos": video_input = modalities["videos"] + assert video_input is not None video_embeddings = self._process_video_input(video_input) multimodal_embeddings += video_embeddings diff --git a/vllm/model_executor/models/moonvit.py b/vllm/model_executor/models/moonvit.py index 29400f597b..9269cd2976 100644 --- a/vllm/model_executor/models/moonvit.py +++ b/vllm/model_executor/models/moonvit.py @@ -551,9 +551,11 @@ class MoonVitEncoder(nn.Module): max_seqlen: torch.Tensor | None = None, ) -> torch.Tensor: if rope_freqs_cis is None: + assert grid_hw is not None rope_freqs_cis = self.rope_2d.get_freqs_cis_by_seqlens(grid_hws=grid_hw) if cu_seqlens is None: + assert grid_hw is not None lengths = torch.cat( ( torch.zeros(1, device=hidden_states.device, dtype=grid_hw.dtype), @@ -578,7 +580,7 @@ class MoonVitEncoder(nn.Module): def patch_merger( x: torch.Tensor, grid_hw: torch.Tensor, - merge_kernel_size: list[int, int] = (2, 2), + merge_kernel_size: tuple[int, int] = (2, 2), ) -> list[torch.Tensor]: d_model = x.size(-1) diff --git a/vllm/model_executor/models/moss_audio.py b/vllm/model_executor/models/moss_audio.py index 70ce6ac4e2..d0bd8510a8 100644 --- a/vllm/model_executor/models/moss_audio.py +++ b/vllm/model_executor/models/moss_audio.py @@ -5,6 +5,7 @@ import math from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, field +from functools import partial from typing import Annotated, Any import numpy as np @@ -14,10 +15,11 @@ import torch.nn.functional as F from torch import nn from transformers import BatchFeature, PretrainedConfig, Qwen3Config from transformers.models.whisper import WhisperFeatureExtractor +from typing_extensions import TypedDict from vllm.compilation.decorators import support_torch_compile 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, get_tensor_model_parallel_world_size, @@ -89,6 +91,16 @@ DEFAULT_MOSS_AUDIO_MEL_CONFIG = { "mel_hop_length": 160, "mel_n_fft": 400, } + + +class MossAudioProcessorKwargs(TypedDict): + audio_token_id: int + audio_start_id: int + audio_end_id: int + enable_time_marker: bool + mel_config: dict[str, int] + + MOSS_AUDIO_PLACEHOLDER = ( f"{MOSS_AUDIO_BOS_TOKEN}{MOSS_AUDIO_TOKEN}{MOSS_AUDIO_EOS_TOKEN}" ) @@ -170,7 +182,9 @@ def _extract_moss_audio_mel_config( for target_key, source_keys in aliases.items(): for source_key in source_keys: if source_key in mel_config: - config[target_key] = int(mel_config[source_key]) + value = mel_config[source_key] + assert isinstance(value, int) + config[target_key] = value break return config @@ -963,6 +977,9 @@ class MossAudioProcessor: enable_time_marker: bool = False, mel_config: Mapping[str, object] | None = None, ) -> None: + assert hasattr(tokenizer, "encode") + assert hasattr(tokenizer, "decode") + assert hasattr(tokenizer, "batch_decode") self.tokenizer = tokenizer self.audio_token_id = int(audio_token_id) self.audio_start_id = int(audio_start_id) @@ -1229,15 +1246,19 @@ class MossAudioProcessingInfo(BaseProcessingInfo): @staticmethod def _get_processor_cache_key(kwargs: Mapping[str, object]) -> tuple[object, ...]: - mel_config = _normalize_moss_audio_mel_config( - kwargs.get("mel_config") - if isinstance(kwargs.get("mel_config"), Mapping) - else None - ) + mel_config_arg = kwargs.get("mel_config") + assert mel_config_arg is None or isinstance(mel_config_arg, Mapping) + mel_config = _normalize_moss_audio_mel_config(mel_config_arg) + audio_token_id = kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID) + audio_start_id = kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID) + audio_end_id = kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID) + assert isinstance(audio_token_id, int) + assert isinstance(audio_start_id, int) + assert isinstance(audio_end_id, int) return ( - int(kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID)), - int(kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID)), - int(kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID)), + audio_token_id, + audio_start_id, + audio_end_id, bool(kwargs.get("enable_time_marker", False)), tuple(sorted(mel_config.items())), ) @@ -1248,21 +1269,19 @@ class MossAudioProcessingInfo(BaseProcessingInfo): self.ctx.get_merged_mm_kwargs({}), kwargs, ) - mel_config = _normalize_moss_audio_mel_config( - merged_kwargs.get("mel_config") - if isinstance(merged_kwargs.get("mel_config"), Mapping) - else None - ) - processor_kwargs = { - "audio_token_id": int( - merged_kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID) - ), - "audio_start_id": int( - merged_kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID) - ), - "audio_end_id": int( - merged_kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID) - ), + mel_config_arg = merged_kwargs.get("mel_config") + assert mel_config_arg is None or isinstance(mel_config_arg, Mapping) + mel_config = _normalize_moss_audio_mel_config(mel_config_arg) + audio_token_id = merged_kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID) + audio_start_id = merged_kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID) + audio_end_id = merged_kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID) + assert isinstance(audio_token_id, int) + assert isinstance(audio_start_id, int) + assert isinstance(audio_end_id, int) + processor_kwargs: MossAudioProcessorKwargs = { + "audio_token_id": audio_token_id, + "audio_start_id": audio_start_id, + "audio_end_id": audio_end_id, "enable_time_marker": bool(merged_kwargs.get("enable_time_marker", False)), "mel_config": mel_config, } @@ -1326,6 +1345,7 @@ class MossAudioDummyInputsBuilder(BaseDummyInputsBuilder[MossAudioProcessingInfo ) -> MultiModalDataDict: num_audios = mm_counts.get("audio", 0) audio_overrides = mm_options.get("audio") + assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) return { "audio": self._get_dummy_audios( length=16000, @@ -1423,7 +1443,7 @@ class MossAudioMultiModalProcessor(BaseMultiModalProcessor[MossAudioProcessingIn ), ) - prompt_update_specs = [ + prompt_update_specs: list[tuple[list[int], list[int]]] = [ ( [ processor.audio_start_id, @@ -1452,12 +1472,7 @@ class MossAudioMultiModalProcessor(BaseMultiModalProcessor[MossAudioProcessingIn PromptReplacement( modality="audio", target=target, - replacement=( - lambda item_idx, suffix_token_ids=suffix_token_ids: get_replacement( - item_idx, - suffix_token_ids, - ) - ), + replacement=partial(get_replacement, suffix_token_ids=suffix_token_ids), ) for target, suffix_token_ids in prompt_update_specs ] diff --git a/vllm/model_executor/models/moss_transcribe_diarize.py b/vllm/model_executor/models/moss_transcribe_diarize.py index c26ebf5206..4dc3bdb3d7 100644 --- a/vllm/model_executor/models/moss_transcribe_diarize.py +++ b/vllm/model_executor/models/moss_transcribe_diarize.py @@ -20,7 +20,7 @@ from torch import nn from transformers import BatchFeature from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions from vllm.config.speech_to_text import SpeechToTextParams from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TextPrompt from vllm.model_executor.models.interfaces import ( @@ -299,8 +299,9 @@ class MossTranscribeDiarizeWhisperEncoder(WhisperEncoder): def forward( self, input_features: torch.Tensor, - audio_feature_lengths: torch.Tensor, + audio_feature_lengths: torch.Tensor | None = None, ) -> torch.Tensor: + assert audio_feature_lengths is not None if input_features.numel() == 0: return input_features.new_empty((1, 0, self.conv1.out_channels)) device = self.conv1.weight.device @@ -431,11 +432,13 @@ class MossTranscribeDiarizeDummyInputsBuilder( return {} feature_extractor = self.info.get_feature_extractor() + audio_overrides = mm_options.get("audio") + assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions) return { "audio": self._get_dummy_audios( length=_get_max_audio_samples(feature_extractor), num_audios=num_audios, - overrides=mm_options.get("audio"), + overrides=audio_overrides, ) } diff --git a/vllm/model_executor/models/muse_glimmer.py b/vllm/model_executor/models/muse_glimmer.py index 6b6530e9a6..ade38cc703 100644 --- a/vllm/model_executor/models/muse_glimmer.py +++ b/vllm/model_executor/models/muse_glimmer.py @@ -45,7 +45,11 @@ from transformers import BatchFeature from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig -from vllm.config.multimodal import BaseDummyOptions +from vllm.config.multimodal import ( + BaseDummyOptions, + ImageDummyOptions, + VideoDummyOptions, +) from vllm.distributed import ( divide, get_pp_group, @@ -241,12 +245,16 @@ class MuseGlimmerDummyInputsBuilder(BaseDummyInputsBuilder[MuseGlimmerProcessing * int(video_processor.downsample_factor) * video_grid ) + image_overrides = mm_options.get("image") + video_overrides = mm_options.get("video") + assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions) + assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions) return { "image": self._get_dummy_images( width=image_width, height=image_height, num_images=mm_counts.get("image", 0), - overrides=mm_options.get("image"), + overrides=image_overrides, ), "video": self._get_dummy_videos( width=video_size, @@ -255,7 +263,7 @@ class MuseGlimmerDummyInputsBuilder(BaseDummyInputsBuilder[MuseGlimmerProcessing seq_len, mm_counts ), num_videos=mm_counts.get("video", 0), - overrides=mm_options.get("video"), + overrides=video_overrides, ), } @@ -370,7 +378,7 @@ class MuseGlimmerMultiModalProcessor( hf_inputs: BatchFeature, hf_processor_mm_kwargs: Mapping[str, object], ) -> Mapping[str, MultiModalFieldConfig]: - fields = {} + fields: dict[str, MultiModalFieldConfig] = {} if "image_pixel_values" in hf_inputs: fields.update( image_pixel_values=MultiModalFieldConfig.batched("image"), @@ -410,7 +418,9 @@ class MuseGlimmerMultiModalProcessor( def image_replacement(item_idx: int) -> PromptUpdateDetails: out_item = out_mm_kwargs["image"][item_idx] - num_tokens = int(out_item["image_feature_sizes"].data) + feature_size = out_item["image_feature_sizes"].data + assert isinstance(feature_size, torch.Tensor) + num_tokens = int(feature_size.item()) replacement = ( [image_start_id] + [config.image_token_id] * num_tokens + [image_end_id] ) @@ -427,7 +437,9 @@ class MuseGlimmerMultiModalProcessor( "MuseGlimmer video must contain at least one frame group" ) - num_tokens = int(out_item["video_feature_sizes"].data) + feature_size = out_item["video_feature_sizes"].data + assert isinstance(feature_size, torch.Tensor) + num_tokens = int(feature_size.item()) tokens_per_group, remainder = divmod(num_tokens, num_groups) if remainder: raise ValueError( @@ -1477,6 +1489,8 @@ class MuseGlimmerForCausalLM( prefix=maybe_prefix(prefix, "model"), ) + self.vision_encoder: MuseGlimmerVisionEncoder | None + self.vision_adapter: MuseGlimmerVisionAdapter | None if self.has_vision: with self._mark_tower_model(vllm_config, {"image", "video"}): self.vision_encoder = MuseGlimmerVisionEncoder( @@ -1616,7 +1630,7 @@ class MuseGlimmerForCausalLM( image_input = self._parse_and_validate_image_input(**kwargs) video_input = self._parse_and_validate_video_input(**kwargs) - embeddings: MultiModalEmbeddings = [] + embeddings: list[torch.Tensor] = [] for key in kwargs: if key == "image_pixel_values" and image_input is not None: embeddings.extend(self._process_image_input(image_input))