@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 76f3249fbd

[Mypy] Fix typing for M models (#54262)

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

Taneem Ibrahim committed Sep 1, 2026 at 10:26 UTC 76f3249fbdff761aa897218913050de0526ccca2

26 files changed +513 -201

tools/pre_commit/mypy.py

-1

@@ -114,7 +114,6 @@ SEPARATE_GROUPS = [

114 EXCLUDE = [

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

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

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

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

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

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

vllm/model_executor/models/mamba.py

+2 -2

@@ -52,14 +52,14 @@ class MambaDecoderLayer(nn.Module):

52 model_config: ModelConfig | None = None,

53 cache_config: CacheConfig | None = None,

54 quant_config: QuantizationConfig | None = None,

55 - is_lora_enabled: bool | None = False,

55 + is_lora_enabled: bool = False,

56 prefix: str = "",

57 ) -> None:

58 super().__init__()

59 self.config = config

60 self.is_falcon_mamba = config.model_type == "falcon_mamba"

61 self.is_lora_enabled = is_lora_enabled

62 - mixer_rms_eps = config.mixer_rms_eps if self.is_falcon_mamba else None

62 + mixer_rms_eps = config.mixer_rms_eps if self.is_falcon_mamba else 1e-5

63 self.mixer = MambaMixer(

64 hidden_size=config.hidden_size,

65 ssm_state_size=config.state_size,

vllm/model_executor/models/medusa.py

+3 -1

@@ -54,7 +54,9 @@ class Medusa(nn.Module):

54 needs to have truncated_vocab_size (=k) as an attribute."""

56 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

57 - config = vllm_config.speculative_config.draft_model_config.hf_config

57 + speculative_config = vllm_config.speculative_config

58 + assert speculative_config is not None

59 + config = speculative_config.draft_model_config.hf_config

60 super().__init__()

61 self.config = config

62 self.blocks = nn.ModuleList(

vllm/model_executor/models/mellum.py

+24 -18

@@ -92,24 +92,30 @@ class MellumAttention(Qwen3MoeAttention):

92 rope_parameters=rope_parameters,

93 dual_chunk_attention_config=dual_chunk_attention_config,

94 )

95 - self.attn = Attention(

96 - self.num_heads,

97 - self.head_dim,

98 - self.scaling,

99 - num_kv_heads=self.num_kv_heads,

100 - cache_config=cache_config,

101 - quant_config=quant_config,

102 - per_layer_sliding_window=per_layer_sliding_window,

103 - prefix=f"{prefix}.attn",

104 - **(

105 - {

106 - "layer_idx": extract_layer_index(prefix),

107 - "dual_chunk_attention_config": dual_chunk_attention_config,

108 - }

109 - if dual_chunk_attention_config

110 - else {}

111 - ),

112 - )

95 + if dual_chunk_attention_config:

96 + self.attn = Attention(

97 + self.num_heads,

98 + self.head_dim,

99 + self.scaling,

100 + num_kv_heads=self.num_kv_heads,

101 + cache_config=cache_config,

102 + quant_config=quant_config,

103 + per_layer_sliding_window=per_layer_sliding_window,

104 + prefix=f"{prefix}.attn",

105 + layer_idx=extract_layer_index(prefix),

106 + dual_chunk_attention_config=dual_chunk_attention_config,

107 + )

108 + else:

109 + self.attn = Attention(

110 + self.num_heads,

111 + self.head_dim,

112 + self.scaling,

113 + num_kv_heads=self.num_kv_heads,

114 + cache_config=cache_config,

115 + quant_config=quant_config,

116 + per_layer_sliding_window=per_layer_sliding_window,

117 + prefix=f"{prefix}.attn",

118 + )

119

120 self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

121 self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

vllm/model_executor/models/midashenglm.py

+11 -5

@@ -36,7 +36,7 @@ from torch.nn.functional import scaled_dot_product_attention

36 from transformers import BatchFeature

38 from vllm.config import VllmConfig

39 -from vllm.config.multimodal import BaseDummyOptions

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

40 from vllm.distributed import get_tensor_model_parallel_world_size

41 from vllm.inputs import MultiModalDataDict

42 from vllm.model_executor.layers.activation import get_act_fn

@@ -570,6 +570,7 @@ class MiDashengLMDummyInputsBuilder(BaseDummyInputsBuilder[MiDashengLMProcessing

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

571

572 audio_overrides = mm_options.get("audio")

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

574

575 return {

576 "audio": self._get_dummy_audios(

@@ -593,6 +594,7 @@ class MiDashengLMMultiModalProcessor(

594 ) -> tuple[Mapping[str, object], Mapping[str, object]]:

595 mm_data = dict(mm_data)

596 audios = mm_data.pop("audios", [])

597 + assert isinstance(audios, list)

598

599 min_audio_len = self.info.get_min_audio_len()

600 processed_audios = [

@@ -637,6 +639,7 @@ class MiDashengLMMultiModalProcessor(

639

640 out_mm_data = out_mm_kwargs.get_data()

641 audio_length = out_mm_data.get("audio_length")

642 + audio_output_lengths: list[int]

643 if audio_length is None:

644 audio_output_lengths = []

645 else:

@@ -645,10 +648,13 @@ class MiDashengLMMultiModalProcessor(

648 if isinstance(audio_length, torch.Tensor)

649 else audio_length

650 )

648 - audio_output_lengths = [

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

650 - for length in audio_length_np

651 - ]

651 + assert isinstance(audio_length_np, (list, tuple, np.ndarray))

652 + audio_output_lengths = []

653 + for length in audio_length_np:

654 + assert isinstance(length, (int, np.integer))

655 + audio_output_lengths.append(

656 + max(1, calculate_mel_frames_dasheng(int(length)))

657 + )

658

659 def get_replacement_midashenglm(item_idx: int):

660 num_features = audio_output_lengths[item_idx]

vllm/model_executor/models/mimo_audio.py

+18 -11

@@ -531,18 +531,18 @@ class MiMoAudioTokenizerConfig(PretrainedConfig):

531 kernel_size: int = 3,

532 activation_function: str = "gelu",

533 encoder_layers: int = 8,

534 - encoder_skip_layer_id: int = None,

534 + encoder_skip_layer_id: int | None = None,

535 encoder_attention_heads: int = 12,

536 encoder_ffn_dim: int = 3072,

537 encoder_causal: bool = False,

538 - encoder_attn_window_size: list = None,

538 + encoder_attn_window_size: list[int] | None = None,

539 decoder_layers: int = 8,

540 decoder_attention_heads: int = 12,

541 decoder_ffn_dim: int = 3072,

542 decoder_kernel_size: int = 3,

543 decoder_stride_size: int = 2,

544 decoder_causal: bool = True,

545 - decoder_attn_window_size: list = None,

545 + decoder_attn_window_size: list[int] | None = None,

546 nfft: int = 1024,

547 vocoder_dim: int = 512,

548 vocoder_intermediate_dim: int = 4096,

@@ -553,16 +553,16 @@ class MiMoAudioTokenizerConfig(PretrainedConfig):

553 window_size: int = 1024,

554 vocoder_padding: str = "same",

555 fmin: int = 0,

556 - fmax: int = None,

556 + fmax: int | None = None,

557 num_quantizers: int = 12,

558 - codebook_size: list = None,

558 + codebook_size: list[int] | None = None,

559 threshold_ema_dead_code: int = 10,

560 position_embedding_type: str = "rope",

561 rope_theta: int = 10000,

562 rope_type: str = "default",

563 ln_type: str = "LayerNorm",

564 vocoder_attention_heads: int = 4,

565 - vocoder_attn_window_size: list = None,

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

566 use_istft_only: bool = False,

567 hybrid_attention: bool = False,

568 hybrid_block_size: int = 8,

@@ -808,17 +808,22 @@ class AudioEncoder(nn.Module):

808 config.rope_type,

809 )

810

811 - attn_window_sizes = []

811 + encoder_attn_window_size = config.encoder_attn_window_size

812 + assert encoder_attn_window_size is not None

813 + assert len(encoder_attn_window_size) == 2

814 + attn_window_size = (

815 + encoder_attn_window_size[0],

816 + encoder_attn_window_size[1],

817 + )

818 + attn_window_sizes: list[tuple[int, int]] = []

819 if config.hybrid_attention:

820 for i in range(config.encoder_layers):

821 if i % config.swa_per_block < config.swa_per_block - 1:

815 - attn_window_sizes.append(tuple(config.encoder_attn_window_size))

822 + attn_window_sizes.append(attn_window_size)

823 else:

824 attn_window_sizes.append((-1, -1))

825 else:

819 - attn_window_sizes = [

820 - tuple(config.encoder_attn_window_size)

821 - ] * config.encoder_layers

826 + attn_window_sizes = [attn_window_size] * config.encoder_layers

827

828 self.layers = nn.ModuleList(

829 [

@@ -848,6 +853,7 @@ class AudioEncoder(nn.Module):

853 else:

854 self.down_sample_layer = None

855

856 + self.quantizer: ResidualVectorQuantizer | None

857 if config.num_quantizers != 0:

858 self.quantizer = ResidualVectorQuantizer(

859 dimension=config.d_model,

@@ -976,6 +982,7 @@ class AudioEncoder(nn.Module):

982

983 @torch.no_grad()

984 def decode_vq(self, codes):

985 + assert self.quantizer is not None

986 self.quantizer.float()

987 return self.quantizer.decode(codes)

988

vllm/model_executor/models/mimo_v2.py

+5 -3

@@ -303,6 +303,7 @@ class MiMoV2Attention(nn.Module):

303 backend_enum = requested

304 else:

305 fa_backend = AttentionBackendEnum.FLASH_ATTN_DIFFKV.get_class()

306 + assert hasattr(fa_backend, "is_supported_on_current_device")

307 if fa_backend.is_supported_on_current_device(

308 head_size=self.head_dim,

309 head_size_v=self.v_head_dim,

@@ -312,6 +313,7 @@ class MiMoV2Attention(nn.Module):

313 else:

314 backend_enum = AttentionBackendEnum.TRITON_ATTN_DIFFKV

315 attn_backend = backend_enum.get_class()

316 + assert hasattr(attn_backend, "set_head_size_v")

317 attn_backend.set_head_size_v(self.v_head_dim)

318 logger.info_once("Using %s for attention.", attn_backend.get_name())

319 else:

@@ -636,7 +638,7 @@ class MiMoV2Model(nn.Module, EagleModelMixin):

638 )

639

640 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

639 - stacked_params_mapping = [

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

642 # (param_name, shard_name, shard_id)

643 ("qkv_proj", "q_proj", "q"),

644 ("qkv_proj", "k_proj", "k"),

@@ -709,7 +711,7 @@ class MiMoV2Model(nn.Module, EagleModelMixin):

711 ):

712 continue

713 stacked_matched = False

712 - for param_name, weight_name, shard_id in stacked_params_mapping:

714 + for param_name, weight_name, stacked_shard_id in stacked_params_mapping:

715 if weight_name not in name:

716 continue

717 name_rewritten = name.replace(weight_name, param_name)

@@ -728,7 +730,7 @@ class MiMoV2Model(nn.Module, EagleModelMixin):

730

731 param = params_dict[name_rewritten]

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

731 - weight_loader(param, loaded_weight, shard_id)

733 + weight_loader(param, loaded_weight, stacked_shard_id)

734 loaded_params.add(name_rewritten)

735

736 stacked_matched = True

vllm/model_executor/models/mimo_v2_mtp.py

+10 -5

@@ -255,7 +255,7 @@ class MiMoV2MTP(nn.Module):

255 tp_rank = get_tensor_model_parallel_rank()

256 tp_size = get_tensor_model_parallel_world_size()

257

258 - stacked_params_mapping = [

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

259 ("gate_up_proj", "gate_proj", 0),

260 ("gate_up_proj", "up_proj", 1),

261 # Flash format: separate projections → fused qkv_proj

@@ -308,8 +308,11 @@ class MiMoV2MTP(nn.Module):

308 if name_rewritten not in params_dict:

309 continue

310 param = params_dict[name_rewritten]

311 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

312 - weight_loader(param, loaded_weight, shard_id)

311 + weight_loader = getattr(param, "weight_loader", None)

312 + if weight_loader is None:

313 + default_weight_loader(param, loaded_weight)

314 + else:

315 + weight_loader(param, loaded_weight, shard_id)

316 loaded_params.add(name_rewritten)

317 stacked_matched = True

318 break

@@ -331,8 +334,10 @@ class MiMoV2MTP(nn.Module):

334 0, tp_rank * heads_per_rank, heads_per_rank

335 )

336

334 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

335 - weight_loader(param, loaded_weight)

337 + direct_weight_loader = getattr(

338 + param, "weight_loader", default_weight_loader

339 + )

340 + direct_weight_loader(param, loaded_weight)

341 loaded_params.add(name)

342

343 return loaded_params

vllm/model_executor/models/mimo_v2_omni.py

+119 -43

@@ -12,9 +12,14 @@ import torch.nn as nn

12 import torch.nn.functional as F

13 from transformers import BatchFeature, PretrainedConfig

14 from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize

15 +from typing_extensions import TypedDict

17 from vllm.config import VllmConfig

17 -from vllm.config.multimodal import BaseDummyOptions

18 +from vllm.config.multimodal import (

19 + BaseDummyOptions,

20 + ImageDummyOptions,

21 + VideoDummyOptions,

22 +)

23 from vllm.distributed import parallel_state

24 from vllm.distributed import utils as dist_utils

25 from vllm.inputs import MultiModalDataDict

@@ -70,6 +75,18 @@ from .qwen2_vl import _create_qwen2vl_field_factory

75 from .utils import AutoWeightsLoader, IntermediateTensors, WeightsMapper, maybe_prefix

78 +class MiMoV2OmniAudioInputs(TypedDict):

79 + type: str

80 + audio_features: torch.Tensor | list[torch.Tensor]

81 + audio_token_lens: torch.Tensor | list[int] | None

82 +

83 +

84 +class MiMoV2OmniMultiModalInputs(TypedDict, total=False):

85 + image: Qwen2_5_VLImageInputs | None

86 + video: Qwen2_5_VLVideoInputs | None

87 + audio: MiMoV2OmniAudioInputs | None

88 +

89 +

90 class MiMoVisionMLP(Qwen2_5_VisionMLP):

91 pass

@@ -937,11 +954,14 @@ class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessing

954 # Handle video_audio items: convert video part to (TCHW, timestamps) tuple

955 if "video_audio" in mm_data:

956 va_converted: list[VideoAudioInput] = []

940 - for va_item in mm_data["video_audio"]:

957 + video_audio_items = mm_data["video_audio"]

958 + assert isinstance(video_audio_items, Sequence)

959 + for va_item in video_audio_items:

960 if isinstance(va_item, VideoAudioInput):

961 vid = va_item.video

962 else:

963 # Expect (video_frames, audio_source) tuple

964 + assert isinstance(va_item, (list, tuple)) and len(va_item) == 2

965 vid, audio_src = va_item

966 va_item = VideoAudioInput(video=vid, audio=audio_src)

967 vid = vid

@@ -976,7 +996,9 @@ class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessing

996

997 if "videos" in mm_data:

998 converted: list[tuple[torch.Tensor, torch.Tensor]] = []

979 - for video in mm_data["videos"]:

999 + videos = mm_data["videos"]

1000 + assert isinstance(videos, Sequence)

1001 + for video in videos:

1002 if (

1003 isinstance(video, tuple)

1004 and len(video) == 2

@@ -1037,30 +1059,47 @@ class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessing

1059 video_end_id = p.video_end_token_id

1060 audio_start_id = p.audio_start_token_id

1061 audio_end_id = p.audio_end_token_id

1062 + assert isinstance(vision_start_id, int)

1063 + assert isinstance(vision_end_id, int)

1064 + assert isinstance(video_start_id, int)

1065 + assert isinstance(video_end_id, int)

1066

1067 def get_image_replacement(item_idx: int) -> PromptUpdateDetails:

1068 out_item = out_mm_kwargs["image"][item_idx]

1069 grid_thw = out_item["image_grid_thw"].data

1070 + assert isinstance(grid_thw, torch.Tensor)

1071 n_tokens = int(grid_thw.prod()) // merge_size**2

1045 - return [image_pad_id] * n_tokens

1072 + return PromptUpdateDetails.from_seq([image_pad_id] * n_tokens)

1073

1074 def get_video_replacement(item_idx: int) -> PromptUpdateDetails:

1075 out_item = out_mm_kwargs["video"][item_idx]

1076 grid_thw = out_item["video_grid_thw"].data

1050 - spt = float(out_item["second_per_grid_ts"].data)

1051 - start = float(out_item["video_start_times"].data)

1077 + second_per_grid_ts = out_item["second_per_grid_ts"].data

1078 + video_start_time = out_item["video_start_times"].data

1079 + assert isinstance(grid_thw, torch.Tensor)

1080 + assert isinstance(second_per_grid_ts, (int, float, torch.Tensor))

1081 + assert isinstance(video_start_time, (int, float, torch.Tensor))

1082 + spt = float(second_per_grid_ts)

1083 + start = float(video_start_time)

1084

1085 T, H, W = map(int, grid_thw)

1086 n_per_grid = H * W // (merge_size * merge_size)

1087

1088 # Check if this is a video_audio item

1089 n_segs_field = out_item.get("video_audio_n_segs")

1058 - n_segs_val = int(n_segs_field.data) if n_segs_field is not None else 0

1090 + if n_segs_field is None:

1091 + n_segs_val = 0

1092 + else:

1093 + n_segs_data = n_segs_field.data

1094 + assert isinstance(n_segs_data, (int, torch.Tensor))

1095 + n_segs_val = int(n_segs_data)

1096 va_seg_lens: list[int] | None = None

1097 if n_segs_val > 0:

1098 seg_lens_field = out_item.get("video_audio_seg_lens")

1099 if seg_lens_field is not None:

1063 - va_seg_lens = seg_lens_field.data[:n_segs_val].tolist()

1100 + seg_lens_data = seg_lens_field.data

1101 + assert isinstance(seg_lens_data, torch.Tensor)

1102 + va_seg_lens = seg_lens_data[:n_segs_val].tolist()

1103

1104 full: list[int] = [video_start_id]

1105 is_embed_mask: list[bool] = [False]

@@ -1099,6 +1138,9 @@ class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessing

1138 is_embed_mask.append(False)

1139 # Audio tokens for this group

1140 seg_len = va_seg_lens[g]

1141 + assert audio_start_id is not None

1142 + assert audio_pad_id is not None

1143 + assert audio_end_id is not None

1144 full.append(audio_start_id)

1145 is_embed_mask.append(False)

1146 full.extend([audio_pad_id] * seg_len)

@@ -1117,8 +1159,11 @@ class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessing

1159

1160 def get_audio_replacement(item_idx: int) -> PromptUpdateDetails:

1161 out_item = out_mm_kwargs["audio"][item_idx]

1120 - tok_len = int(out_item["audio_token_lens"].data)

1121 - return [audio_pad_id] * tok_len

1162 + audio_token_lens = out_item["audio_token_lens"].data

1163 + assert isinstance(audio_token_lens, (int, torch.Tensor))

1164 + assert audio_pad_id is not None

1165 + tok_len = int(audio_token_lens)

1166 + return PromptUpdateDetails.from_seq([audio_pad_id] * tok_len)

1167

1168 updates: list[PromptUpdate] = [

1169 PromptReplacement(

@@ -1166,20 +1211,24 @@ class MiMoV2OmniDummyInputsBuilder(BaseDummyInputsBuilder[MiMoV2OmniProcessingIn

1211 target_num_frames = self.info.get_num_frames_with_most_features(

1212 seq_len, mm_counts

1213 )

1214 + image_overrides = mm_options.get("image")

1215 + video_overrides = mm_options.get("video")

1216 + assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)

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

1218

1219 return {

1220 "image": self._get_dummy_images(

1221 width=target_width,

1222 height=target_height,

1223 num_images=num_images,

1175 - overrides=mm_options.get("image"),

1224 + overrides=image_overrides,

1225 ),

1226 "video": self._get_dummy_videos(

1227 width=target_width,

1228 height=target_height,

1229 num_frames=target_num_frames,

1230 num_videos=num_videos,

1182 - overrides=mm_options.get("video"),

1231 + overrides=video_overrides,

1232 ),

1233 }

1234

@@ -1235,6 +1284,7 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ

1284 )

1285 audio_config = getattr(config, "audio_config", None)

1286 model_path = vllm_config.model_config.model

1287 + self.audio_encoder: MimoAudioEncoder | None

1288 if audio_config is not None:

1289 with self._mark_tower_model(vllm_config, "audio"):

1290 self.audio_encoder = MimoAudioEncoder(

@@ -1269,12 +1319,12 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ

1319 image_grid_thw=image_grid_thw,

1320 )

1321

1272 - if image_embeds is not None:

1273 - return Qwen2_5_VLImageEmbeddingInputs(

1274 - type="image_embeds",

1275 - image_embeds=image_embeds,

1276 - image_grid_thw=image_grid_thw,

1277 - )

1322 + assert image_embeds is not None

1323 + return Qwen2_5_VLImageEmbeddingInputs(

1324 + type="image_embeds",

1325 + image_embeds=image_embeds,

1326 + image_grid_thw=image_grid_thw,

1327 + )

1328

1329 def _parse_and_validate_video_input(

1330 self, **kwargs: object

@@ -1295,13 +1345,13 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ

1345 second_per_grid_ts=second_per_grid_ts,

1346 )

1347

1298 - if video_embeds is not None:

1299 - return Qwen2_5_VLVideoEmbeddingInputs(

1300 - type="video_embeds",

1301 - video_embeds=video_embeds,

1302 - video_grid_thw=video_grid_thw,

1303 - second_per_grid_ts=second_per_grid_ts,

1304 - )

1348 + assert video_embeds is not None

1349 + return Qwen2_5_VLVideoEmbeddingInputs(

1350 + type="video_embeds",

1351 + video_embeds=video_embeds,

1352 + video_grid_thw=video_grid_thw,

1353 + second_per_grid_ts=second_per_grid_ts,

1354 + )

1355

1356 def _process_image_input(

1357 self, image_input: Qwen2_5_VLImageInputs

@@ -1339,19 +1389,29 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ

1389 sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

1390 return video_embeds.split(sizes)

1391

1342 - def _parse_and_validate_audio_input(self, **kwargs: object) -> dict | None:

1392 + def _parse_and_validate_audio_input(

1393 + self, **kwargs: object

1394 + ) -> MiMoV2OmniAudioInputs | None:

1395 audio_features = kwargs.pop("audio_features", None)

1396 audio_token_lens = kwargs.pop("audio_token_lens", None)

1397 if audio_features is None:

1398 return None

1399 + assert isinstance(audio_features, (torch.Tensor, list))

1400 + if isinstance(audio_features, list):

1401 + assert all(isinstance(item, torch.Tensor) for item in audio_features)

1402 + assert audio_token_lens is None or isinstance(

1403 + audio_token_lens, (torch.Tensor, list)

1404 + )

1405 return {

1406 "type": "audio",

1407 "audio_features": audio_features,

1408 "audio_token_lens": audio_token_lens,

1409 }

1410

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

1354 - mm_input_by_modality = {}

1411 + def _parse_and_validate_multimodal_inputs(

1412 + self, **kwargs: object

1413 + ) -> MiMoV2OmniMultiModalInputs:

1414 + mm_input_by_modality: MiMoV2OmniMultiModalInputs = {}

1415

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

1417 # from the order of kwargs.

@@ -1376,7 +1436,9 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ

1436 )

1437 return mm_input_by_modality

1438

1379 - def _process_audio_input(self, audio_input: dict) -> tuple[torch.Tensor, ...]:

1439 + def _process_audio_input(

1440 + self, audio_input: MiMoV2OmniAudioInputs

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

1442 mel_specs = audio_input["audio_features"]

1443 if self.audio_encoder is None:

1444 return ()

@@ -1407,38 +1469,52 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ

1469

1470 # Pre-process va audio: one mel spec per va video → per-video audio embeddings

1471 # keyed by va video index (0-based among va videos only)

1410 - va_audio_embs_list: list[tuple[torch.Tensor, ...]] = []

1472 + va_audio_embs_list: list[torch.Tensor] = []

1473 if va_audio_features is not None and self.audio_encoder is not None:

1412 - mel_list = (

1413 - list(va_audio_features)

1414 - if isinstance(va_audio_features, torch.Tensor)

1415 - else list(va_audio_features)

1416 - )

1474 + assert isinstance(va_audio_features, (torch.Tensor, list, tuple))

1475 + mel_list: list[torch.Tensor] = []

1476 + for mel_spec in va_audio_features:

1477 + assert isinstance(mel_spec, torch.Tensor)

1478 + mel_list.append(mel_spec)

1479 for mel_spec in mel_list:

1480 embs, tok_lens = self.audio_encoder.get_audio_feature([mel_spec])

1481 # tok_lens is a list/tensor with one entry (total tokens for this mel)

1482 va_audio_embs_list.append(embs) # shape (total_tok, hidden)

1483

1484 va_cursor = 0 # index into va_audio_embs_list

1485 + assert video_audio_n_segs is None or isinstance(

1486 + video_audio_n_segs, torch.Tensor

1487 + )

1488 + assert video_audio_seg_lens is None or isinstance(

1489 + video_audio_seg_lens, torch.Tensor

1490 + )

1491

1492 # NOTE: Iterate in dict insertion order to preserve token sequence order.

1493 for modality in mm_input_by_modality:

1426 - multimodal_input = mm_input_by_modality[modality]

1494 if modality == "image":

1428 - multimodal_embeddings.extend(

1429 - self._process_image_input(multimodal_input)

1495 + image_input = mm_input_by_modality["image"]

1496 + assert isinstance(

1497 + image_input,

1498 + (Qwen2_5_VLImagePixelInputs, Qwen2_5_VLImageEmbeddingInputs),

1499 )

1500 + multimodal_embeddings.extend(self._process_image_input(image_input))

1501 elif modality == "video":

1432 - video_embs_tuple = self._process_video_input(multimodal_input)

1502 + video_input = mm_input_by_modality["video"]

1503 + assert isinstance(

1504 + video_input,

1505 + (Qwen2_5_VLVideoPixelInputs, Qwen2_5_VLVideoEmbeddingInputs),

1506 + )

1507 + video_embs_tuple = self._process_video_input(video_input)

1508 if video_audio_n_segs is None:

1509 multimodal_embeddings.extend(video_embs_tuple)

1510 else:

1436 - grid_thw = multimodal_input["video_grid_thw"]

1511 + grid_thw = video_input["video_grid_thw"]

1512 for i, vid_embs in enumerate(video_embs_tuple):

1513 n_segs = int(video_audio_n_segs[i])

1514 if n_segs == 0 or not va_audio_embs_list:

1515 multimodal_embeddings.append(vid_embs)

1516 else:

1517 + assert video_audio_seg_lens is not None

1518 T = int(grid_thw[i][0])

1519 n_per_grid = vid_embs.shape[0] // T

1520 frames = list(vid_embs.split(n_per_grid, dim=0))

@@ -1459,9 +1535,9 @@ class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQ

1535 )

1536 multimodal_embeddings.append(group_audio_embs[g])

1537 elif modality == "audio":

1462 - multimodal_embeddings.extend(

1463 - self._process_audio_input(multimodal_input)

1464 - )

1538 + audio_input = mm_input_by_modality["audio"]

1539 + assert audio_input is not None

1540 + multimodal_embeddings.extend(self._process_audio_input(audio_input))

1541 return tuple(multimodal_embeddings)

1542

1543 def forward(

vllm/model_executor/models/minicpm.py

@@ -456,6 +456,7 @@ class MiniCPMModel(nn.Module, EagleModelMixin):

456 hidden_states = self.embed_input_ids(input_ids)

457 residual = None

458 else:

459 + assert intermediate_tensors is not None

460 hidden_states = intermediate_tensors["hidden_states"]

461 residual = intermediate_tensors["residual"]

462

vllm/model_executor/models/minicpm_eagle.py

+18 -3

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

27 import math

28 from collections.abc import Iterable

29 +from typing import TYPE_CHECKING

31 import torch

32 from torch import nn

@@ -55,6 +56,14 @@ from .utils import (

56 process_eagle_weight,

57 )

59 +if TYPE_CHECKING:

60 +

61 + class _EagleMiniCPMSupportsPP:

62 + pass

63 +

64 +else:

65 + _EagleMiniCPMSupportsPP = SupportsPP

66 +

68 class EagleMiniCPMDecoderLayer(nn.Module):

69 def __init__(

@@ -149,7 +158,9 @@ class EagleMiniCPMModel(nn.Module):

158 ):

159 super().__init__()

160

152 - config = vllm_config.speculative_config.draft_model_config.hf_config

161 + speculative_config = vllm_config.speculative_config

162 + assert speculative_config is not None

163 + config = speculative_config.draft_model_config.hf_config

164 cache_config = vllm_config.cache_config

165 quant_config = vllm_config.quant_config

166

@@ -291,7 +302,9 @@ class EagleMiniCPMModel(nn.Module):

302 return loaded_params

303

304

294 -class EagleMiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle):

305 +class EagleMiniCPMForCausalLM(

306 + nn.Module, SupportsLoRA, _EagleMiniCPMSupportsPP, SupportsEagle

307 +):

308 packed_modules_mapping = {

309 "qkv_proj": [

310 "q_proj",

@@ -312,7 +325,9 @@ class EagleMiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle

325

326 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

327 super().__init__()

315 - config = vllm_config.speculative_config.draft_model_config.hf_config

328 + speculative_config = vllm_config.speculative_config

329 + assert speculative_config is not None

330 + config = speculative_config.draft_model_config.hf_config

331 cache_config = vllm_config.cache_config

332 quant_config = vllm_config.quant_config

333

vllm/model_executor/models/minicpmo.py

+62 -18

@@ -40,7 +40,7 @@ from transformers.models.whisper.modeling_whisper import (

40 )

42 from vllm.config import VllmConfig

43 -from vllm.config.multimodal import BaseDummyOptions

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

44 from vllm.inputs import ModalityData, MultiModalDataDict

45 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems

46 from vllm.multimodal.inputs import (

@@ -66,8 +66,12 @@ from .minicpmv import (

66 _MAX_FRAMES_PER_VIDEO,

67 MiniCPMV2_6,

68 MiniCPMV4_5,

69 + MiniCPMVBaseModel,

70 MiniCPMVDummyInputsBuilder,

71 + MiniCPMVImageEmbeddingInputs,

72 + MiniCPMVImagePixelInputs,

73 MiniCPMVMultiModalDataParser,

74 + MiniCPMVMultiModalInputs,

75 MiniCPMVMultiModalProcessor,

76 MiniCPMVProcessingInfo,

77 _minicpmv_field_config,

@@ -181,6 +185,10 @@ MiniCPMOAudioInputs: TypeAlias = (

185 )

186

187

188 +class MiniCPMOMultiModalInputs(MiniCPMVMultiModalInputs, total=False):

189 + audios: MiniCPMOAudioInputs | None

190 +

191 +

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

193 audio_features = hf_inputs.get("audio_features")

194 audio_feature_lens = hf_inputs.get("audio_feature_lens")

@@ -405,6 +413,7 @@ class MiniCPMODummyInputsBuilder(MiniCPMVDummyInputsBuilder[MiniCPMOProcessingIn

413 )

414

415 audio_overrides = mm_options.get("audio")

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

417

418 audio_mm_data = {

419 "audio": self._get_dummy_audios(

@@ -469,7 +478,8 @@ class MiniCPMOMultiModalProcessor(MiniCPMVMultiModalProcessor[MiniCPMOProcessing

478 if isinstance(lens, torch.Tensor):

479 flat_feature_lens.extend(lens.flatten().tolist())

480 else:

472 - flat_feature_lens.append(int(lens))

481 + assert isinstance(lens, int)

482 + flat_feature_lens.append(lens)

483 unpadded_audio_features = [

484 feat[:, :length]

485 for feat, length in zip(

@@ -515,6 +525,9 @@ class MiniCPMOMultiModalProcessor(MiniCPMVMultiModalProcessor[MiniCPMOProcessing

525 audios = mm_items.get_items(

526 "audio", (MiniCPMOAudioEmbeddingItems, AudioProcessorItems)

527 )

528 + assert isinstance(

529 + audios, (MiniCPMOAudioEmbeddingItems, AudioProcessorItems)

530 + )

531

532 if isinstance(audios, MiniCPMOAudioEmbeddingItems):

533 single_audio_embeds = audios.get(item_idx)["audio_embeds"]

@@ -652,7 +665,7 @@ class MiniCPMWhisperEncoder(WhisperEncoder):

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

666 )

667

655 - encoder_states = ()

668 + encoder_states: tuple[torch.Tensor, ...] = ()

669

670 for idx, encoder_layer in enumerate(self.layers):

671 encoder_states = encoder_states + (hidden_states,)

@@ -682,7 +695,13 @@ class MiniCPMWhisperEncoder(WhisperEncoder):

695 )

696

697

685 -class MiniCPMOBaseModel:

698 +if TYPE_CHECKING:

699 + _MiniCPMOBaseModelBase = MiniCPMVBaseModel

700 +else:

701 + _MiniCPMOBaseModelBase = object

702 +

703 +

704 +class MiniCPMOBaseModel(_MiniCPMOBaseModelBase):

705 """Base mixin class for MiniCPM-O models with audio support."""

706

707 # Unlike the vision-only MiniCPM-V models, audio weights are loaded here.

@@ -897,8 +916,12 @@ class MiniCPMOBaseModel:

916 audio_feature_lens=audio_feature_lens,

917 )

918

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

901 - modalities = super()._parse_and_validate_multimodal_inputs(**kwargs)

919 + def _parse_and_validate_multimodal_inputs(

920 + self, **kwargs: object

921 + ) -> MiniCPMOMultiModalInputs:

922 + modalities = MiniCPMOMultiModalInputs(

923 + **super()._parse_and_validate_multimodal_inputs(**kwargs)

924 + )

925

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

927 # from the order of kwargs.

@@ -918,16 +941,33 @@ class MiniCPMOBaseModel:

941 if audio_input["type"] == "audio_embeds":

942 return audio_input["audio_embeds"]

943

944 + assert isinstance(audio_input, MiniCPMOAudioFeatureInputs)

945 return self.get_audio_hidden_states(audio_input)

946

923 - def _process_multimodal_inputs(self, modalities: dict):

924 - multimodal_embeddings = super()._process_multimodal_inputs(modalities)

925 -

926 - for modality in modalities:

927 - if modality == "audios":

928 - audio_input = modalities["audios"]

929 - audio_embeddings = self._process_audio_input(audio_input)

930 - multimodal_embeddings += tuple(audio_embeddings)

947 + def _process_multimodal_inputs(self, modalities: Mapping[str, object]):

948 + base_modalities: MiniCPMVMultiModalInputs = {}

949 + image_input = modalities.get("images")

950 + video_input = modalities.get("videos")

951 + assert image_input is None or isinstance(

952 + image_input, (MiniCPMVImagePixelInputs, MiniCPMVImageEmbeddingInputs)

953 + )

954 + assert video_input is None or isinstance(

955 + video_input, (MiniCPMVImagePixelInputs, MiniCPMVImageEmbeddingInputs)

956 + )

957 + if "images" in modalities:

958 + base_modalities["images"] = image_input

959 + if "videos" in modalities:

960 + base_modalities["videos"] = video_input

961 + multimodal_embeddings = super()._process_multimodal_inputs(base_modalities)

962 +

963 + audio_input = modalities.get("audios")

964 + if audio_input is not None:

965 + assert isinstance(

966 + audio_input,

967 + (MiniCPMOAudioFeatureInputs, MiniCPMOAudioEmbeddingInputs),

968 + )

969 + audio_embeddings = self._process_audio_input(audio_input)

970 + multimodal_embeddings += tuple(audio_embeddings)

971

972 return multimodal_embeddings

973

@@ -983,7 +1023,7 @@ class MiniCPMO(MiniCPMOBaseModel, MiniCPMV2_6):

1023 try:

1024 version_str = str(config.version)

1025 version_parts = version_str.split(".")

986 - version = tuple(int(x) for x in version_parts[:2])

1026 + version_values = tuple(int(x) for x in version_parts[:2])

1027 except (ValueError, TypeError) as e:

1028 raise ValueError(

1029 f"Invalid model version format in config: {config.version}. "

@@ -991,17 +1031,21 @@ class MiniCPMO(MiniCPMOBaseModel, MiniCPMV2_6):

1031 ) from e

1032 else:

1033 # Default to 2.6 for backward compatibility

994 - version = (2, 6)

1034 + version_values = (2, 6)

1035

1036 # Dispatch class based on version

997 - instance_cls = _MINICPMO_SUPPORT_VERSION.get(version)

1037 + if len(version_values) == 2:

1038 + version_key = (version_values[0], version_values[1])

1039 + instance_cls = _MINICPMO_SUPPORT_VERSION.get(version_key)

1040 + else:

1041 + instance_cls = None

1042 if instance_cls is None:

1043 supported_versions = ", ".join(

1044 [f"{v[0]}.{v[1]}" for v in sorted(_MINICPMO_SUPPORT_VERSION.keys())]

1045 )

1046 raise ValueError(

1047 f"Currently, MiniCPMO only supports versions "

1004 - f"{supported_versions}. Got version: {version}"

1048 + f"{supported_versions}. Got version: {version_values}"

1049 )

1050

1051 return instance_cls(vllm_config=vllm_config, prefix=prefix)

vllm/model_executor/models/minicpmv.py

+43 -13

@@ -41,7 +41,7 @@ from transformers.dynamic_module_utils import (

41 get_class_from_dynamic_module,

42 resolve_trust_remote_code,

43 )

44 -from typing_extensions import TypeVar

44 +from typing_extensions import TypedDict, TypeVar

46 from vllm.config import VllmConfig

47 from vllm.config.multimodal import (

@@ -97,7 +97,7 @@ from vllm.transformers_utils.processor import (

97 cached_get_image_processor,

98 )

99 from vllm.transformers_utils.utils import convert_model_repo_to_path

100 -from vllm.utils.collection_utils import flatten_2d_lists

100 +from vllm.utils.collection_utils import flatten_2d_lists, is_list_of

101 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

102 from vllm.utils.tensor_schema import TensorSchema, TensorShape

103 from vllm.utils.torch_utils import set_default_torch_dtype

@@ -386,6 +386,7 @@ class Resampler4_5(Resampler2_5):

386 for i in range(bs):

387 tgt_h, tgt_w = tgt_sizes[i]

388 if temporal_pos_emb:

389 + assert temporal_ids_flatten is not None

390 if temporal_ids_flatten[i] == -1:

391 pos_embed_temporal.append(

392 torch.zeros(self.embed_dim, dtype=dtype, device=device)

@@ -549,6 +550,11 @@ class MiniCPMVVideoEmbeddingItems(DictEmbeddingItems):

550 return len(self.get(index)["video_image_sizes"])

551

552

553 +class MiniCPMVMultiModalInputs(TypedDict, total=False):

554 + images: MiniCPMVImageInputs | None

555 + videos: MiniCPMVImageInputs | None

556 +

557 +

558 class MiniCPMVMultiModalDataParser(MultiModalDataParser):

559 # The per-item sizes are what size the placeholder range.

560 embedding_fields = {

@@ -833,6 +839,7 @@ class MiniCPMVDummyInputsBuilder(BaseDummyInputsBuilder[_I]):

839

840 image_overrides = mm_options.get("image")

841 video_overrides = mm_options.get("video")

842 + assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)

843

844 # Convert video overrides to image overrides for per-frame image generation,

845 # and apply num_frames override to num_video_frames.

@@ -842,6 +849,7 @@ class MiniCPMVDummyInputsBuilder(BaseDummyInputsBuilder[_I]):

849 num_video_frames = min(num_video_frames, video_overrides.num_frames)

850 if video_overrides.width or video_overrides.height:

851 video_frame_overrides = ImageDummyOptions(

852 + count=video_overrides.count,

853 width=video_overrides.width,

854 height=video_overrides.height,

855 )

@@ -1081,6 +1089,9 @@ class MiniCPMVMultiModalProcessor(BaseMultiModalProcessor[_I]):

1089 images = mm_items.get_items(

1090 "image", (MiniCPMVImageEmbeddingItems, ImageProcessorItems)

1091 )

1092 + assert isinstance(

1093 + images, (MiniCPMVImageEmbeddingItems, ImageProcessorItems)

1094 + )

1095

1096 image_size = images.get_image_size(item_idx)

1097

@@ -1097,6 +1108,9 @@ class MiniCPMVMultiModalProcessor(BaseMultiModalProcessor[_I]):

1108 videos = mm_items.get_items(

1109 "video", (MiniCPMVVideoEmbeddingItems, VideoProcessorItems)

1110 )

1111 + assert isinstance(

1112 + videos, (MiniCPMVVideoEmbeddingItems, VideoProcessorItems)

1113 + )

1114

1115 frame_size = videos.get_frame_size(item_idx)

1116 num_frames = videos.get_num_frames(item_idx)

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

1213 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

1214 config = vllm_config.model_config.hf_config

1215 multimodal_config = vllm_config.model_config.multimodal_config

1216 + assert multimodal_config is not None

1217 quant_config = vllm_config.quant_config

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

1219 super().__init__()

@@ -1261,10 +1276,15 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1276 image_embeds=image_embeds,

1277 )

1278

1279 + assert is_list_of(pixel_values, list, check="all")

1280 + pixel_values_list: list[list[torch.Tensor]] = []

1281 + for pixel_value in pixel_values:

1282 + assert is_list_of(pixel_value, torch.Tensor, check="all")

1283 + pixel_values_list.append(pixel_value)

1284 tgt_sizes = kwargs.pop("tgt_sizes")

1285

1266 - num_slices_flat = torch.tensor([len(ps) for ps in pixel_values])

1267 - pixel_values_flat = flatten_bn(pixel_values)

1286 + num_slices_flat = torch.tensor([len(ps) for ps in pixel_values_list])

1287 + pixel_values_flat = flatten_2d_lists(pixel_values_list)

1288 tgt_sizes_flat = flatten_bn(tgt_sizes, concat=True)

1289

1290 return MiniCPMVImagePixelInputs(

@@ -1274,8 +1294,11 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1294 num_slices=num_slices_flat,

1295 )

1296

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

1278 - modalities = {}

1297 + def _parse_and_validate_multimodal_inputs(

1298 + self, **kwargs: object

1299 + ) -> MiniCPMVMultiModalInputs:

1300 + kwargs.pop("modality", None)

1301 + modalities: MiniCPMVMultiModalInputs = {}

1302

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

1304 # from the order of kwargs.

@@ -1304,12 +1327,13 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1327 if image_input["type"] == "image_embeds":

1328 return image_input["image_embeds"]

1329

1330 + assert isinstance(image_input, MiniCPMVImagePixelInputs)

1331 image_features_flat = self.get_vision_hidden_states(image_input)

1332

1333 num_slices = image_input["num_slices"]

1334 return [e.flatten(0, 1) for e in image_features_flat.split(num_slices.tolist())]

1335

1312 - def _process_multimodal_inputs(self, modalities: dict):

1336 + def _process_multimodal_inputs(self, modalities: MiniCPMVMultiModalInputs):

1337 # The result multimodal_embeddings is tuple of tensors, with each

1338 # tensor corresponding to a multimodal data item (image or video).

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

@@ -1319,10 +1343,12 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):

1343 for modality in modalities:

1344 if modality == "images":

1345 image_input = modalities["images"]

1346 + assert image_input is not None

1347 image_embeddings = self._process_vision_input(image_input)

1348 multimodal_embeddings += tuple(image_embeddings)

1349 if modality == "videos":

1350 video_input = modalities["videos"]

1351 + assert video_input is not None

1352 video_embeddings = self._process_vision_input(video_input)

1353 multimodal_embeddings += tuple(video_embeddings)

1354

@@ -1915,23 +1941,27 @@ class MiniCPMV(MiniCPMVBaseModel, SupportsMultiModal, SupportsLoRA):

1941

1942 def __new__(cls, *, vllm_config: VllmConfig, prefix: str = ""):

1943 config = vllm_config.model_config.hf_config

1944 + version_values: tuple[int, ...]

1945 if not hasattr(config, "version"):

1946 if config.hidden_size == 2304 and config.query_num == 64:

1920 - version = (2, 0)

1947 + version_values = (2, 0)

1948 else:

1922 - version = (2, 5)

1949 + version_values = (2, 5)

1950 else:

1924 - version = str(config.version).split(".")

1925 - version = tuple([int(x) for x in version])

1951 + version_values = tuple(int(x) for x in str(config.version).split("."))

1952 # Dispatch class based on version

1927 - instance_cls = _SUPPORT_VERSION.get(version)

1953 + if len(version_values) == 2:

1954 + version_key = (version_values[0], version_values[1])

1955 + instance_cls = _SUPPORT_VERSION.get(version_key)

1956 + else:

1957 + instance_cls = None

1958 if instance_cls is None:

1959 supported_versions = ", ".join(

1960 [f"{v[0]}.{v[1]}" for v in sorted(_SUPPORT_VERSION.keys())]

1961 )

1962 raise ValueError(

1963 f"Currently, MiniCPMV only supports versions "

1934 - f"{supported_versions}. Got version: {version}"

1964 + f"{supported_versions}. Got version: {version_values}"

1965 )

1966

1967 # quant_config references base class members,

vllm/model_executor/models/minicpmv4_6.py

+9 -2

@@ -98,7 +98,9 @@ class MiniCPMV4_6MultiModalProcessor(MiniCPMVMultiModalProcessor):

98 ds = mm_kwargs.get("downsample_mode")

99 if ds is not None:

100 return str(ds)

101 - return self.info._get_downsample_mode()

101 + info = self.info

102 + assert isinstance(info, MiniCPMV4_6ProcessingInfo)

103 + return info._get_downsample_mode()

104

105 def get_image_prompt_texts(

106 self,

@@ -106,7 +108,9 @@ class MiniCPMV4_6MultiModalProcessor(MiniCPMVMultiModalProcessor):

108 image_idx: int = 0,

109 downsample_mode: str | None = None,

110 ) -> str:

109 - return self.info.get_slice_image_placeholder(

111 + info = self.info

112 + assert isinstance(info, MiniCPMV4_6ProcessingInfo)

113 + return info.get_slice_image_placeholder(

114 image_size,

115 image_idx=image_idx,

116 downsample_mode=downsample_mode,

@@ -128,6 +132,7 @@ class MiniCPMV4_6MultiModalProcessor(MiniCPMVMultiModalProcessor):

132 # which one is used. Using image_token for video silently produces

133 # garbage descriptions.

134 info = self.info

135 + assert isinstance(info, MiniCPMV4_6ProcessingInfo)

136 grids, source_tokens, patch_tokens = info._compute_visual_tokens(

137 image_size,

138 max_slice_nums=info.get_video_max_slice_num(),

@@ -1016,6 +1021,7 @@ class MiniCPMV4_6ForConditionalGeneration(

1021 config: MiniCPMV4_6Config = vllm_config.model_config.hf_config

1022 quant_config = vllm_config.quant_config

1023 multimodal_config = vllm_config.model_config.multimodal_config

1024 + assert multimodal_config is not None

1025

1026 self.config = config

1027 self.multimodal_config = multimodal_config

@@ -1078,6 +1084,7 @@ class MiniCPMV4_6ForConditionalGeneration(

1084 image_embeds=image_embeds,

1085 )

1086

1087 + assert isinstance(pixel_values, torch.Tensor | list)

1088 tgt_sizes = kwargs.pop("tgt_sizes")

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

1090 pixel_values_flat = flatten_bn(pixel_values)

vllm/model_executor/models/mistral.py

+8 -2

@@ -221,15 +221,21 @@ class MistralModel(LlamaModel):

221 intermediate_tensors: IntermediateTensors | None,

222 inputs_embeds: torch.Tensor | None = None,

223 t_cond: torch.Tensor | None = None,

224 + **extra_layer_kwargs: object,

225 ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:

226 return super().forward(

226 - input_ids, positions, intermediate_tensors, inputs_embeds, t_cond=t_cond

227 + input_ids,

228 + positions,

229 + intermediate_tensors,

230 + inputs_embeds,

231 + t_cond=t_cond,

232 + **extra_layer_kwargs,

233 )

234

235

236 class MistralForCausalLM(LlamaForCausalLM):

237 # Mistral: We don't support LoRA on the embedding layers.

232 - embedding_modules: dict[str, str] = {}

238 + embedding_modules = {}

239

240 # Mistral/Llama models can also be loaded with --load-format mistral

241 # from consolidated.safetensors checkpoints

vllm/model_executor/models/mistral3.py

@@ -196,6 +196,7 @@ class Mistral3HFEncoderInfo(PixtralHFEncoderInfo):

196 image_height = math.floor(image_height / ratio)

197

198 patch_size = self.vision_config.patch_size

199 + assert isinstance(self.hf_config, Mistral3Config)

200 spatial_merge_size = self.hf_config.spatial_merge_size

201

202 # The HF processor rounds each dimension up to the vision patch size

vllm/model_executor/models/mistral_eagle.py

+25 -4

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

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

4 from collections.abc import Iterable

5 +from typing import TYPE_CHECKING

7 import torch

8 import torch.nn as nn

@@ -33,6 +34,20 @@ from vllm.model_executor.models.utils import (

34 logger = init_logger(__name__)

37 +if TYPE_CHECKING:

38 +

39 + class _EagleMistralModelBase(nn.Module):

40 + pass

41 +

42 + class _EagleMistralForCausalLMBase(nn.Module):

43 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

44 + raise NotImplementedError

45 +

46 +else:

47 + _EagleMistralModelBase = MistralModel

48 + _EagleMistralForCausalLMBase = MistralForCausalLM

49 +

50 +

51 class EagleMistralDecoderLayer(MistralDecoderLayer):

52 def __init__(

53 self,

@@ -47,7 +62,7 @@ class EagleMistralDecoderLayer(MistralDecoderLayer):

64 @support_torch_compile

50 -class EagleMistralModel(MistralModel):

65 +class EagleMistralModel(_EagleMistralModelBase):

66 def __init__(

67 self,

68 *,

@@ -58,7 +73,9 @@ class EagleMistralModel(MistralModel):

73 # Bypass MistralModel.__init__ to avoid creating duplicate attention

74 # layer entries in the global context.

75 nn.Module.__init__(self)

61 - self.config = vllm_config.speculative_config.draft_model_config.hf_config

76 + speculative_config = vllm_config.speculative_config

77 + assert speculative_config is not None

78 + self.config = speculative_config.draft_model_config.hf_config

79 self.vocab_size = self.config.vocab_size

80 # Get drafter's quantization config

81 self.quant_config = get_draft_quant_config(vllm_config)

@@ -112,7 +129,9 @@ class EagleMistralModel(MistralModel):

129 return hidden_states, hidden_states

130

131

115 -class EagleMistralForCausalLM(MistralForCausalLM, SupportsMultiModalEmbeddings):

132 +class EagleMistralForCausalLM(

133 + _EagleMistralForCausalLMBase, SupportsMultiModalEmbeddings

134 +):

135 mistral_mapping = MistralForCausalLM.mistral_mapping | {

136 "eagle_linear": "model.fc",

137 }

@@ -121,7 +140,9 @@ class EagleMistralForCausalLM(MistralForCausalLM, SupportsMultiModalEmbeddings):

140 # Bypass MistralForCausalLM.__init__ to use the draft model config

141 # and to avoid creating an lm_head.

142 nn.Module.__init__(self)

124 - self.config = vllm_config.speculative_config.draft_model_config.hf_config

143 + speculative_config = vllm_config.speculative_config

144 + assert speculative_config is not None

145 + self.config = speculative_config.draft_model_config.hf_config

146 target_layer_num = vllm_config.model_config.get_num_layers(

147 vllm_config.parallel_config

148 )

vllm/model_executor/models/mistral_large_3_eagle.py

+15 -3

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

4 import copy

5 from collections.abc import Iterable

6 from functools import partial

7 +from typing import TYPE_CHECKING

9 import regex

10 import torch

@@ -28,6 +29,15 @@ from .utils import WeightsMapper, make_empty_intermediate_tensors_factory, maybe

29 logger = init_logger(__name__)

32 +if TYPE_CHECKING:

33 +

34 + class _EagleMistralLarge3ForCausalLMBase(nn.Module):

35 + pass

36 +

37 +else:

38 + _EagleMistralLarge3ForCausalLMBase = MistralLarge3ForCausalLM

39 +

40 +

41 @support_torch_compile

42 class EagleMistralLarge3Model(DeepseekV2Model):

43 def __init__(

@@ -108,7 +118,7 @@ class EagleMistralLarge3Model(DeepseekV2Model):

118

119

120 class EagleMistralLarge3ForCausalLM(

111 - MistralLarge3ForCausalLM, SupportsMultiModalEmbeddings

121 + _EagleMistralLarge3ForCausalLMBase, SupportsMultiModalEmbeddings

122 ):

123 hf_to_vllm_mapper = MistralLarge3ForCausalLM.hf_to_vllm_mapper | WeightsMapper(

124 orig_to_new_regex={

@@ -119,13 +129,15 @@ class EagleMistralLarge3ForCausalLM(

129 )

130

131 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

132 + speculative_config = vllm_config.speculative_config

133 + assert speculative_config is not None

134 target_layer_num = vllm_config.model_config.get_num_layers(

135 vllm_config.parallel_config

136 )

125 - vllm_config.model_config = vllm_config.speculative_config.draft_model_config

137 + vllm_config.model_config = speculative_config.draft_model_config

138 # draft model quantization config may differ from target model

139 self.quant_config = VllmConfig.get_quantization_config(

128 - vllm_config.speculative_config.draft_model_config, vllm_config.load_config

140 + speculative_config.draft_model_config, vllm_config.load_config

141 )

142 vllm_config.quant_config = self.quant_config

143 self.model_cls = partial(

vllm/model_executor/models/mllama4.py

+10 -7

@@ -620,14 +620,16 @@ class Mllama4MultiModalProcessor(BaseMultiModalProcessor[Mllama4ProcessingInfo])

620 max_num_chunks=self.info.get_max_num_tiles(),

621 patch_size=SizeDict(height=tile_size, width=tile_size),

622 )

623 - best_fit_sizes = [

624 - get_best_fit(

625 - (image.size[1], image.size[0]),

626 - torch.tensor(possible_resolutions),

627 - resize_to_max_canvas=image_processor.resize_to_max_canvas,

623 + best_fit_sizes = []

624 + for image in parsed_images:

625 + assert image is not None

626 + best_fit_sizes.append(

627 + get_best_fit(

628 + (image.size[1], image.size[0]),

629 + torch.tensor(possible_resolutions),

630 + resize_to_max_canvas=image_processor.resize_to_max_canvas,

631 + )

632 )

629 - for image in parsed_images

630 - ]

633 # TODO tile height/width do not necessarily need to match

634 aspect_ratios = [

635 (image_size[0] // tile_size, image_size[1] // tile_size)

@@ -760,6 +762,7 @@ class Llama4ForConditionalGeneration(

762 config = vllm_config.model_config.hf_config

763 quant_config = vllm_config.quant_config

764 multimodal_config = vllm_config.model_config.multimodal_config

765 + assert multimodal_config is not None

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

767

768 self.vllm_config = vllm_config

vllm/model_executor/models/module_mapping.py

+4 -4

@@ -17,10 +17,10 @@ class MultiModelKeys:

18 @staticmethod

19 def from_string_field(

20 - language_model: str | list[str] = None,

21 - connector: str | list[str] = None,

22 - tower_model: str | list[str] = None,

23 - generator: str | list[str] = None,

20 + language_model: str | list[str] | None = None,

21 + connector: str | list[str] | None = None,

22 + tower_model: str | list[str] | None = None,

23 + generator: str | list[str] | None = None,

24 **kwargs,

25 ) -> "MultiModelKeys":

26 def to_list(value):

vllm/model_executor/models/molmo.py

+7 -4

@@ -332,7 +332,7 @@ class VisionTransformer(nn.Module):

332 ):

333 super().__init__()

334 scale = config.image_emb_dim**-0.5

335 - self.patch_num = config.image_num_patch

335 + self.patch_num: tuple[int, int] = config.image_num_patch

336 self.class_embedding = nn.Parameter(torch.randn(config.image_emb_dim) * scale)

337 self.num_prefix_tokens: int = NUM_PREFIX_TOKENS

338 self.positional_embedding = nn.Parameter(

@@ -349,7 +349,7 @@ class VisionTransformer(nn.Module):

349 config, quant_config, prefix=f"{prefix}.transformer"

350 )

351

352 - def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor:

352 + def add_pos_emb(self, x: torch.Tensor, patch_num: tuple[int, int]) -> torch.Tensor:

353 cls_emb = self.positional_embedding[0:1]

354 pos_emb = self.positional_embedding[1:]

355

@@ -380,7 +380,7 @@ class VisionTransformer(nn.Module):

380 return x

381

382 def forward(

383 - self, x: torch.Tensor, patch_num: int | None = None

383 + self, x: torch.Tensor, patch_num: tuple[int, int] | None = None

384 ) -> list[torch.Tensor]:

385 """

386 : param x: (batch_size, num_patch, n_pixels)

@@ -484,12 +484,15 @@ class MolmoAttention(nn.Module):

484 def _apply_qk_norm(

485 self, q: torch.Tensor, k: torch.Tensor

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

487 + assert self.q_norm is not None

488 + assert self.k_norm is not None

489 if self.tp_size > 1:

490 q = tensor_model_parallel_all_gather(q.contiguous())

491 k = tensor_model_parallel_all_gather(k.contiguous())

492 q = self.q_norm(q)

493 k = self.k_norm(k)

494 if self.tp_size > 1:

495 + assert self.tp_rank is not None

496 splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size)

497 q = splitter(q)[self.tp_rank]

498 k = splitter(k)[self.tp_rank]

@@ -1345,7 +1348,7 @@ class MolmoForCausalLM(

1348 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

1349 )

1350

1348 - self.img_patch_id = None

1351 + self.img_patch_id: int | None = None

1352

1353 self.lm_head = ParallelLMHead(

1354 config.embedding_size or config.vocab_size,

vllm/model_executor/models/molmo2.py

+41 -9

@@ -22,6 +22,7 @@ from transformers import (

22 )

23 from transformers.image_utils import ImageInput

24 from transformers.video_utils import VideoMetadata

25 +from typing_extensions import TypedDict

27 from vllm.compilation.decorators import support_torch_compile

28 from vllm.config import CacheConfig, VllmConfig

@@ -157,6 +158,11 @@ class Molmo2VideoInputs(TensorSchema):

158 num_video_tokens: Annotated[torch.Tensor, TensorShape("nv")]

159

160

161 +class Molmo2MultiModalInputs(TypedDict, total=False):

162 + images: Molmo2ImageInputs | None

163 + videos: Molmo2VideoInputs | None

164 +

165 +

166 @dataclass

167 class VitConfig:

168 """Config for a vision transformer"""

@@ -474,7 +480,7 @@ class Molmo2VisionTransformer(nn.Module):

480 super().__init__()

481 scale = config.hidden_size**-0.5

482 self.num_prefix_tokens: int = 0 # no class embeddings

477 - self.patch_num = config.image_num_patch

483 + self.patch_num: tuple[int, int] = config.image_num_patch

484 self.positional_embedding = nn.Parameter(

485 torch.randn(config.image_num_pos, config.hidden_size) * scale,

486 )

@@ -490,7 +496,7 @@ class Molmo2VisionTransformer(nn.Module):

496 prefix=maybe_prefix(prefix, "transformer"),

497 )

498

493 - def add_pos_emb(self, x: torch.Tensor, patch_num: int) -> torch.Tensor:

499 + def add_pos_emb(self, x: torch.Tensor, patch_num: tuple[int, int]) -> torch.Tensor:

500 pos_emb = self.positional_embedding

501

502 pos_emb = pos_emb.reshape(

@@ -522,7 +528,7 @@ class Molmo2VisionTransformer(nn.Module):

528 def forward(

529 self,

530 x: torch.Tensor,

525 - patch_num: int | None = None,

531 + patch_num: tuple[int, int] | None = None,

532 ) -> list[torch.Tensor]:

533 """

534 : param x: (batch_size, num_patch, n_pixels)

@@ -652,6 +658,7 @@ class ImagePoolingAttention(nn.Module):

658 if self.use_pytorch_sdpa:

659 output = self.forward_sdpa(xq, xk, xv, attn_mask)

660 else:

661 + assert self.attn is not None

662 output = self.attn(xq, xk, xv)

663

664 output, _ = self.o_proj(output)

@@ -954,12 +961,15 @@ class Molmo2Attention(nn.Module):

961 q: torch.Tensor,

962 k: torch.Tensor,

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

964 + assert self.q_norm is not None

965 + assert self.k_norm is not None

966 if self.tp_size > 1:

967 q = tensor_model_parallel_all_gather(q.contiguous())

968 k = tensor_model_parallel_all_gather(k.contiguous())

969 q = self.q_norm(q)

970 k = self.k_norm(k)

971 if self.tp_size > 1:

972 + assert self.tp_rank is not None

973 splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size)

974 q = splitter(q)[self.tp_rank]

975 k = splitter(k)[self.tp_rank]

@@ -1911,10 +1921,10 @@ class Molmo2DummyInputsBuilder(BaseDummyInputsBuilder[Molmo2ProcessingInfo]):

1921 )

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

1923

1914 - video_items = []

1924 + video_items: list[VideoItem] = []

1925 for video in videos:

1926 video_num_frames = video.shape[0]

1917 - video_metadata = {

1927 + video_metadata: dict[str, int | float | str | bool | list[int]] = {

1928 "fps": 2.0,

1929 "duration": video_num_frames / 2.0,

1930 "total_num_frames": video_num_frames,

@@ -1972,6 +1982,7 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

1982 image_processor = hf_processor.image_processor

1983

1984 if videos := mm_data.pop("videos", []):

1985 + assert isinstance(videos, Sequence)

1986 bos_token_id = tokenizer.bos_token_id or tokenizer.eos_token_id

1987

1988 pixel_values_videos_lst = []

@@ -1983,7 +1994,9 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

1994 num_video_tokens_lst = []

1995

1996 for item in videos:

1997 + assert isinstance(item, tuple) and len(item) == 2

1998 video_array, metadata = item

1999 + assert isinstance(metadata, Mapping)

2000

2001 # NOTE: metadata.frames_indices indicates

2002 # the sampled frames indices of pre-sampled videos, which is

@@ -2240,6 +2253,7 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

2253 def get_video_replacement_molmo2(item_idx: int):

2254 video, metadata = mm_items["video"][item_idx]

2255 do_sample_frames = hf_processor_mm_kwargs.get("do_sample_frames")

2256 + assert do_sample_frames is None or isinstance(do_sample_frames, bool)

2257

2258 timestamps = self.info._get_video_second_idx(metadata, do_sample_frames)

2259 ncols, nrows = self.info.get_base_grid_size(video_processor)

@@ -2408,10 +2422,17 @@ class Molmo2ForConditionalGeneration(

2422 image_tokens = kwargs.pop("image_tokens", None)

2423 num_image_tokens = kwargs.pop("num_image_tokens", None)

2424

2425 + assert isinstance(token_pooling, torch.Tensor)

2426 + assert isinstance(num_pooled_patches, torch.Tensor)

2427 + assert isinstance(num_patches, torch.Tensor)

2428 + assert isinstance(image_tokens, torch.Tensor)

2429 + assert isinstance(num_image_tokens, torch.Tensor)

2430 +

2431 accum_patches = [0] + num_patches.cumsum(dim=0)[:-1].tolist()

2432 patch_offset = 0

2433 new_token_pooling = token_pooling.clone()

2414 - for i, n in enumerate(num_pooled_patches):

2434 + for i, n_tensor in enumerate(num_pooled_patches):

2435 + n = int(n_tensor)

2436 cur_slice = token_pooling[patch_offset : patch_offset + n]

2437 index_offset = int(accum_patches[i])

2438 new_token_pooling[patch_offset : patch_offset + n] = torch.where(

@@ -2443,10 +2464,17 @@ class Molmo2ForConditionalGeneration(

2464 video_tokens = kwargs.pop("video_tokens", None)

2465 num_video_tokens = kwargs.pop("num_video_tokens", None)

2466

2467 + assert isinstance(token_pooling, torch.Tensor)

2468 + assert isinstance(num_pooled_patches, torch.Tensor)

2469 + assert isinstance(num_patches, torch.Tensor)

2470 + assert isinstance(video_tokens, torch.Tensor)

2471 + assert isinstance(num_video_tokens, torch.Tensor)

2472 +

2473 accum_patches = [0] + num_patches.cumsum(dim=0)[:-1].tolist()

2474 patch_offset = 0

2475 new_token_pooling = token_pooling.clone()

2449 - for i, n in enumerate(num_pooled_patches):

2476 + for i, n_tensor in enumerate(num_pooled_patches):

2477 + n = int(n_tensor)

2478 cur_slice = token_pooling[patch_offset : patch_offset + n]

2479 index_offset = int(accum_patches[i])

2480 new_token_pooling[patch_offset : patch_offset + n] = torch.where(

@@ -2464,8 +2492,10 @@ class Molmo2ForConditionalGeneration(

2492 num_video_tokens=num_video_tokens,

2493 )

2494

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

2468 - modalities = {}

2495 + def _parse_and_validate_multimodal_inputs(

2496 + self, **kwargs: object

2497 + ) -> Molmo2MultiModalInputs:

2498 + modalities: Molmo2MultiModalInputs = {}

2499

2500 for input_key in kwargs:

2501 if input_key in ("pixel_values",) and "images" not in modalities:

@@ -2544,10 +2574,12 @@ class Molmo2ForConditionalGeneration(

2574 for modality in modalities:

2575 if modality == "images":

2576 image_input = modalities["images"]

2577 + assert image_input is not None

2578 image_embeddings = self._process_image_input(image_input)

2579 multimodal_embeddings += image_embeddings

2580 if modality == "videos":

2581 video_input = modalities["videos"]

2582 + assert video_input is not None

2583 video_embeddings = self._process_video_input(video_input)

2584 multimodal_embeddings += video_embeddings

2585

vllm/model_executor/models/moonvit.py

+3 -1

@@ -551,9 +551,11 @@ class MoonVitEncoder(nn.Module):

551 max_seqlen: torch.Tensor | None = None,

552 ) -> torch.Tensor:

553 if rope_freqs_cis is None:

554 + assert grid_hw is not None

555 rope_freqs_cis = self.rope_2d.get_freqs_cis_by_seqlens(grid_hws=grid_hw)

556

557 if cu_seqlens is None:

558 + assert grid_hw is not None

559 lengths = torch.cat(

560 (

561 torch.zeros(1, device=hidden_states.device, dtype=grid_hw.dtype),

@@ -578,7 +580,7 @@ class MoonVitEncoder(nn.Module):

580 def patch_merger(

581 x: torch.Tensor,

582 grid_hw: torch.Tensor,

581 - merge_kernel_size: list[int, int] = (2, 2),

583 + merge_kernel_size: tuple[int, int] = (2, 2),

584 ) -> list[torch.Tensor]:

585 d_model = x.size(-1)

586

vllm/model_executor/models/moss_audio.py

+47 -32

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

5 import math

6 from collections.abc import Iterable, Mapping, Sequence

7 from dataclasses import dataclass, field

8 +from functools import partial

9 from typing import Annotated, Any

11 import numpy as np

@@ -14,10 +15,11 @@ import torch.nn.functional as F

15 from torch import nn

16 from transformers import BatchFeature, PretrainedConfig, Qwen3Config

17 from transformers.models.whisper import WhisperFeatureExtractor

18 +from typing_extensions import TypedDict

20 from vllm.compilation.decorators import support_torch_compile

21 from vllm.config import VllmConfig

20 -from vllm.config.multimodal import BaseDummyOptions

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

23 from vllm.distributed import (

24 get_pp_group,

25 get_tensor_model_parallel_world_size,

@@ -89,6 +91,16 @@ DEFAULT_MOSS_AUDIO_MEL_CONFIG = {

91 "mel_hop_length": 160,

92 "mel_n_fft": 400,

93 }

94 +

95 +

96 +class MossAudioProcessorKwargs(TypedDict):

97 + audio_token_id: int

98 + audio_start_id: int

99 + audio_end_id: int

100 + enable_time_marker: bool

101 + mel_config: dict[str, int]

102 +

103 +

104 MOSS_AUDIO_PLACEHOLDER = (

105 f"{MOSS_AUDIO_BOS_TOKEN}{MOSS_AUDIO_TOKEN}{MOSS_AUDIO_EOS_TOKEN}"

106 )

@@ -170,7 +182,9 @@ def _extract_moss_audio_mel_config(

182 for target_key, source_keys in aliases.items():

183 for source_key in source_keys:

184 if source_key in mel_config:

173 - config[target_key] = int(mel_config[source_key])

185 + value = mel_config[source_key]

186 + assert isinstance(value, int)

187 + config[target_key] = value

188 break

189

190 return config

@@ -963,6 +977,9 @@ class MossAudioProcessor:

977 enable_time_marker: bool = False,

978 mel_config: Mapping[str, object] | None = None,

979 ) -> None:

980 + assert hasattr(tokenizer, "encode")

981 + assert hasattr(tokenizer, "decode")

982 + assert hasattr(tokenizer, "batch_decode")

983 self.tokenizer = tokenizer

984 self.audio_token_id = int(audio_token_id)

985 self.audio_start_id = int(audio_start_id)

@@ -1229,15 +1246,19 @@ class MossAudioProcessingInfo(BaseProcessingInfo):

1246

1247 @staticmethod

1248 def _get_processor_cache_key(kwargs: Mapping[str, object]) -> tuple[object, ...]:

1232 - mel_config = _normalize_moss_audio_mel_config(

1233 - kwargs.get("mel_config")

1234 - if isinstance(kwargs.get("mel_config"), Mapping)

1235 - else None

1236 - )

1249 + mel_config_arg = kwargs.get("mel_config")

1250 + assert mel_config_arg is None or isinstance(mel_config_arg, Mapping)

1251 + mel_config = _normalize_moss_audio_mel_config(mel_config_arg)

1252 + audio_token_id = kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID)

1253 + audio_start_id = kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID)

1254 + audio_end_id = kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID)

1255 + assert isinstance(audio_token_id, int)

1256 + assert isinstance(audio_start_id, int)

1257 + assert isinstance(audio_end_id, int)

1258 return (

1238 - int(kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID)),

1239 - int(kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID)),

1240 - int(kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID)),

1259 + audio_token_id,

1260 + audio_start_id,

1261 + audio_end_id,

1262 bool(kwargs.get("enable_time_marker", False)),

1263 tuple(sorted(mel_config.items())),

1264 )

@@ -1248,21 +1269,19 @@ class MossAudioProcessingInfo(BaseProcessingInfo):

1269 self.ctx.get_merged_mm_kwargs({}),

1270 kwargs,

1271 )

1251 - mel_config = _normalize_moss_audio_mel_config(

1252 - merged_kwargs.get("mel_config")

1253 - if isinstance(merged_kwargs.get("mel_config"), Mapping)

1254 - else None

1255 - )

1256 - processor_kwargs = {

1257 - "audio_token_id": int(

1258 - merged_kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID)

1259 - ),

1260 - "audio_start_id": int(

1261 - merged_kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID)

1262 - ),

1263 - "audio_end_id": int(

1264 - merged_kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID)

1265 - ),

1272 + mel_config_arg = merged_kwargs.get("mel_config")

1273 + assert mel_config_arg is None or isinstance(mel_config_arg, Mapping)

1274 + mel_config = _normalize_moss_audio_mel_config(mel_config_arg)

1275 + audio_token_id = merged_kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID)

1276 + audio_start_id = merged_kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID)

1277 + audio_end_id = merged_kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID)

1278 + assert isinstance(audio_token_id, int)

1279 + assert isinstance(audio_start_id, int)

1280 + assert isinstance(audio_end_id, int)

1281 + processor_kwargs: MossAudioProcessorKwargs = {

1282 + "audio_token_id": audio_token_id,

1283 + "audio_start_id": audio_start_id,

1284 + "audio_end_id": audio_end_id,

1285 "enable_time_marker": bool(merged_kwargs.get("enable_time_marker", False)),

1286 "mel_config": mel_config,

1287 }

@@ -1326,6 +1345,7 @@ class MossAudioDummyInputsBuilder(BaseDummyInputsBuilder[MossAudioProcessingInfo

1345 ) -> MultiModalDataDict:

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

1347 audio_overrides = mm_options.get("audio")

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

1349 return {

1350 "audio": self._get_dummy_audios(

1351 length=16000,

@@ -1423,7 +1443,7 @@ class MossAudioMultiModalProcessor(BaseMultiModalProcessor[MossAudioProcessingIn

1443 ),

1444 )

1445

1426 - prompt_update_specs = [

1446 + prompt_update_specs: list[tuple[list[int], list[int]]] = [

1447 (

1448 [

1449 processor.audio_start_id,

@@ -1452,12 +1472,7 @@ class MossAudioMultiModalProcessor(BaseMultiModalProcessor[MossAudioProcessingIn

1472 PromptReplacement(

1473 modality="audio",

1474 target=target,

1455 - replacement=(

1456 - lambda item_idx, suffix_token_ids=suffix_token_ids: get_replacement(

1457 - item_idx,

1458 - suffix_token_ids,

1459 - )

1460 - ),

1475 + replacement=partial(get_replacement, suffix_token_ids=suffix_token_ids),

1476 )

1477 for target, suffix_token_ids in prompt_update_specs

1478 ]

vllm/model_executor/models/moss_transcribe_diarize.py

+6 -3

@@ -20,7 +20,7 @@ from torch import nn

20 from transformers import BatchFeature

22 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig

23 -from vllm.config.multimodal import BaseDummyOptions

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

24 from vllm.config.speech_to_text import SpeechToTextParams

25 from vllm.inputs import ModalityData, MultiModalDataDict, PromptType, TextPrompt

26 from vllm.model_executor.models.interfaces import (

@@ -299,8 +299,9 @@ class MossTranscribeDiarizeWhisperEncoder(WhisperEncoder):

299 def forward(

300 self,

301 input_features: torch.Tensor,

302 - audio_feature_lengths: torch.Tensor,

302 + audio_feature_lengths: torch.Tensor | None = None,

303 ) -> torch.Tensor:

304 + assert audio_feature_lengths is not None

305 if input_features.numel() == 0:

306 return input_features.new_empty((1, 0, self.conv1.out_channels))

307 device = self.conv1.weight.device

@@ -431,11 +432,13 @@ class MossTranscribeDiarizeDummyInputsBuilder(

432 return {}

433

434 feature_extractor = self.info.get_feature_extractor()

435 + audio_overrides = mm_options.get("audio")

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

437 return {

438 "audio": self._get_dummy_audios(

439 length=_get_max_audio_samples(feature_extractor),

440 num_audios=num_audios,

438 - overrides=mm_options.get("audio"),

441 + overrides=audio_overrides,

442 )

443 }

444

vllm/model_executor/models/muse_glimmer.py

+21 -7

@@ -45,7 +45,11 @@ from transformers import BatchFeature

46 from vllm.compilation.decorators import support_torch_compile

47 from vllm.config import CacheConfig, VllmConfig

48 -from vllm.config.multimodal import BaseDummyOptions

48 +from vllm.config.multimodal import (

49 + BaseDummyOptions,

50 + ImageDummyOptions,

51 + VideoDummyOptions,

52 +)

53 from vllm.distributed import (

54 divide,

55 get_pp_group,

@@ -241,12 +245,16 @@ class MuseGlimmerDummyInputsBuilder(BaseDummyInputsBuilder[MuseGlimmerProcessing

245 * int(video_processor.downsample_factor)

246 * video_grid

247 )

248 + image_overrides = mm_options.get("image")

249 + video_overrides = mm_options.get("video")

250 + assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)

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

252 return {

253 "image": self._get_dummy_images(

254 width=image_width,

255 height=image_height,

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

249 - overrides=mm_options.get("image"),

257 + overrides=image_overrides,

258 ),

259 "video": self._get_dummy_videos(

260 width=video_size,

@@ -255,7 +263,7 @@ class MuseGlimmerDummyInputsBuilder(BaseDummyInputsBuilder[MuseGlimmerProcessing

263 seq_len, mm_counts

264 ),

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

258 - overrides=mm_options.get("video"),

266 + overrides=video_overrides,

267 ),

268 }

269

@@ -370,7 +378,7 @@ class MuseGlimmerMultiModalProcessor(

378 hf_inputs: BatchFeature,

379 hf_processor_mm_kwargs: Mapping[str, object],

380 ) -> Mapping[str, MultiModalFieldConfig]:

373 - fields = {}

381 + fields: dict[str, MultiModalFieldConfig] = {}

382 if "image_pixel_values" in hf_inputs:

383 fields.update(

384 image_pixel_values=MultiModalFieldConfig.batched("image"),

@@ -410,7 +418,9 @@ class MuseGlimmerMultiModalProcessor(

418

419 def image_replacement(item_idx: int) -> PromptUpdateDetails:

420 out_item = out_mm_kwargs["image"][item_idx]

413 - num_tokens = int(out_item["image_feature_sizes"].data)

421 + feature_size = out_item["image_feature_sizes"].data

422 + assert isinstance(feature_size, torch.Tensor)

423 + num_tokens = int(feature_size.item())

424 replacement = (

425 [image_start_id] + [config.image_token_id] * num_tokens + [image_end_id]

426 )

@@ -427,7 +437,9 @@ class MuseGlimmerMultiModalProcessor(

437 "MuseGlimmer video must contain at least one frame group"

438 )

439

430 - num_tokens = int(out_item["video_feature_sizes"].data)

440 + feature_size = out_item["video_feature_sizes"].data

441 + assert isinstance(feature_size, torch.Tensor)

442 + num_tokens = int(feature_size.item())

443 tokens_per_group, remainder = divmod(num_tokens, num_groups)

444 if remainder:

445 raise ValueError(

@@ -1477,6 +1489,8 @@ class MuseGlimmerForCausalLM(

1489 prefix=maybe_prefix(prefix, "model"),

1490 )

1491

1492 + self.vision_encoder: MuseGlimmerVisionEncoder | None

1493 + self.vision_adapter: MuseGlimmerVisionAdapter | None

1494 if self.has_vision:

1495 with self._mark_tower_model(vllm_config, {"image", "video"}):

1496 self.vision_encoder = MuseGlimmerVisionEncoder(

@@ -1616,7 +1630,7 @@ class MuseGlimmerForCausalLM(

1630 image_input = self._parse_and_validate_image_input(**kwargs)

1631 video_input = self._parse_and_validate_video_input(**kwargs)

1632

1619 - embeddings: MultiModalEmbeddings = []

1633 + embeddings: list[torch.Tensor] = []

1634 for key in kwargs:

1635 if key == "image_pixel_values" and image_input is not None:

1636 embeddings.extend(self._process_image_input(image_input))