@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 811d12ee26

[Mypy Fix] Mypy fix for "vllm/model_executor/models/[eE][fF]" (#53381)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Aug 22, 2026 at 15:41 UTC 811d12ee26fa119cf8faae44d80330becf1e7fe3

16 files changed +71 -84

tools/pre_commit/mypy.py

-2

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

112 ]

113

114 EXCLUDE = [

115 - r"vllm/model_executor/models/[eE]",

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

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

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

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

vllm/compilation/decorators.py

+4 -2

@@ -106,8 +106,10 @@ def support_torch_compile(

106 @overload

107 def support_torch_compile(

108 *,

109 - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None,

110 - mark_unbacked_dims: dict[str, int | list[int]] | None,

109 + dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None = None,

110 + mark_unbacked_dims: dict[str, int | list[int]] | None = None,

111 + enable_if: Callable[[VllmConfig], bool] | None = None,

112 + is_encoder: bool = False,

113 ) -> Callable[[type[_T]], type[_T]]: ...

114

115

vllm/model_executor/models/eagle2_5_vl.py

+2 -1

@@ -150,6 +150,7 @@ class Eagle2_5_VLForConditionalGeneration(

150 config = vllm_config.model_config.hf_config

151 quant_config = vllm_config.quant_config

152 multimodal_config = vllm_config.model_config.multimodal_config

153 + assert multimodal_config is not None

154

155 self.config = config

156 self.multimodal_config = multimodal_config

@@ -187,7 +188,7 @@ class Eagle2_5_VLForConditionalGeneration(

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

189 )

190

190 - self.img_context_token_id = None

191 + self.img_context_token_id: int | None = None

192

193 self.make_empty_intermediate_tensors = (

194 self.language_model.make_empty_intermediate_tensors

vllm/model_executor/models/ernie45_moe.py

+2 -1

@@ -127,7 +127,7 @@ class Ernie4_5_MoeMoE(nn.Module):

127 self.layer_idx = layer_idx

128 self.tp_size = get_tensor_model_parallel_world_size()

129

130 - self.moe_num_shared_experts = getattr(config, "moe_num_shared_experts", None)

130 + self.moe_num_shared_experts = getattr(config, "moe_num_shared_experts", 0)

131 self.ep_group = get_ep_group().device_group

132 self.ep_size = self.ep_group.size()

133 self.n_routed_experts: int = config.moe_num_experts

@@ -163,6 +163,7 @@ class Ernie4_5_MoeMoE(nn.Module):

163 torch.empty(config.moe_num_experts, dtype=torch.float32)

164 )

165

166 + self.shared_experts: Ernie4_5_MoeMLP | None

167 if self.has_shared_experts:

168 intermediate_size = (

169 config.moe_intermediate_size * config.moe_num_shared_experts

vllm/model_executor/models/ernie45_vl.py

+27 -20

@@ -1076,7 +1076,7 @@ class Ernie4_5VLMultiModalProcessor(BaseMultiModalProcessor[Ernie4_5_VLProcessin

1076 def _pixel_values_norm(

1077 self,

1078 pixel_values: torch.Tensor,

1079 - mm_kwargs: object,

1079 + mm_kwargs: Mapping[str, object],

1080 ) -> torch.Tensor:

1081 hf_config = self.info.get_hf_config()

1082 vision_config = hf_config.vision_config

@@ -1125,6 +1125,7 @@ class Ernie4_5VLMultiModalProcessor(BaseMultiModalProcessor[Ernie4_5_VLProcessin

1125

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

1127

1128 + mm_data = dict(mm_data)

1129 if "images" not in mm_data:

1130 mm_data["images"] = []

1131 if "videos" not in mm_data:

@@ -1136,15 +1137,15 @@ class Ernie4_5VLMultiModalProcessor(BaseMultiModalProcessor[Ernie4_5_VLProcessin

1137 hf_processor, "supports_video_metadata", False

1138 )

1139

1139 - if mm_data["videos"] and not supports_video_metadata:

1140 + videos = mm_data["videos"]

1141 + assert isinstance(videos, Sequence)

1142 + if videos and not supports_video_metadata:

1143 # Old HF processor, unwrap tuple to pure frames

1144 logger.warning_once(

1145 "HF processor doesn't support video metadata. "

1146 "Timestamps will NOT be rendered. Please upgrade the model."

1147 )

1145 - mm_data["videos"] = [

1146 - v[0] if isinstance(v, tuple) else v for v in mm_data["videos"]

1147 - ]

1148 + mm_data["videos"] = [v[0] if isinstance(v, tuple) else v for v in videos]

1149

1150 processor_output = self.info.ctx.call_hf_processor(

1151 hf_processor,

@@ -1286,6 +1287,7 @@ class Ernie4_5_VLDummyInputsBuilder(BaseDummyInputsBuilder[Ernie4_5_VLProcessing

1287

1288 image_overrides = mm_options.get("image")

1289 video_overrides = mm_options.get("video")

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

1291

1292 return {

1293 "image": self._get_dummy_images(

@@ -1541,10 +1543,14 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1543

1544 offset = mm_feature.mm_position.offset

1545 if mm_feature.modality == "image":

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

1546 + grid_thw = mm_feature.data["image_grid_thw"].data

1547 + assert isinstance(grid_thw, torch.Tensor)

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

1549 yield offset, t, h // spatial_conv_size, w // spatial_conv_size

1550 elif mm_feature.modality == "video":

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

1551 + grid_thw = mm_feature.data["video_grid_thw"].data

1552 + assert isinstance(grid_thw, torch.Tensor)

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

1554 yield (

1555 offset,

1556 t // temporal_conv_size,

@@ -1563,12 +1569,11 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1569 if pixel_values is None:

1570 return None

1571

1566 - if pixel_values is not None:

1567 - return Ernie4_5_VLImagePixelInputs(

1568 - type="pixel_values",

1569 - pixel_values=pixel_values,

1570 - image_grid_thw=image_grid_thw,

1571 - )

1572 + return Ernie4_5_VLImagePixelInputs(

1573 + type="pixel_values",

1574 + pixel_values=pixel_values,

1575 + image_grid_thw=image_grid_thw,

1576 + )

1577

1578 def _parse_and_validate_video_input(

1579 self, **kwargs: object

@@ -1579,12 +1584,11 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1584 if pixel_values_videos is None:

1585 return None

1586

1582 - if pixel_values_videos is not None:

1583 - return Ernie4_5_VLVideoPixelInputs(

1584 - type="pixel_values_videos",

1585 - pixel_values_videos=pixel_values_videos,

1586 - video_grid_thw=video_grid_thw,

1587 - )

1587 + return Ernie4_5_VLVideoPixelInputs(

1588 + type="pixel_values_videos",

1589 + pixel_values_videos=pixel_values_videos,

1590 + video_grid_thw=video_grid_thw,

1591 + )

1592

1593 # -- SupportsEncoderCudaGraph protocol methods --

1594 #

@@ -1748,6 +1752,7 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1752 # the actual batch grid_thw, then scatter the post-merge embeddings.

1753 # Ernie only uses the single "default" encoder path.

1754 output = outputs["default"]

1755 + assert batch_mm_kwargs is not None

1756 grid_thw_cpu = batch_mm_kwargs["image_grid_thw"]

1757 grid_thw = grid_thw_cpu.to(output.device, non_blocking=True)

1758 # The valid token count slices the graph output for the eager

@@ -1799,7 +1804,9 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1804 return video_embeds.split(sizes.tolist())

1805

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

1802 - modalities = {}

1807 + modalities: dict[

1808 + str, Ernie4_5_VLImageInputs | Ernie4_5_VLVideoInputs | None

1809 + ] = {}

1810

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

1812 # from the order of kwargs.

vllm/model_executor/models/ernie45_vl_moe.py

+4 -2

@@ -232,6 +232,7 @@ class Ernie4_5_VLMoeMoE(nn.Module):

232

233 assert text_moe_layer_start_index <= text_moe_layer_end_index

234

235 + self.shared_experts: Ernie4_5_VLMoeMLP | None

236 if self.has_shared_experts:

237 intermediate_size = (

238 config.moe_intermediate_size[0] * config.moe_num_shared_experts

@@ -771,9 +772,10 @@ class Ernie4_5_VLMoeForCausalLM(nn.Module, SupportsPP):

772 if is_pp_missing_parameter(name, self):

773 continue

774 # Remapping the name of FP8 kv-scale.

774 - name = maybe_remap_kv_scale_name(name, params_dict)

775 - if name is None:

775 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

776 + if remapped_name is None:

777 continue

778 + name = remapped_name

779

780 param = params_dict[name]

781

vllm/model_executor/models/exaone4_5.py

@@ -332,6 +332,7 @@ class Exaone4_5_ForConditionalGeneration(Qwen2_5_VLForConditionalGeneration):

332 config: Exaone4_5_Config = vllm_config.model_config.hf_config

333 self.vllm_config = vllm_config

334 multimodal_config = vllm_config.model_config.multimodal_config

335 + assert multimodal_config is not None

336

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

338 self.config = config

vllm/model_executor/models/exaone4_5_mtp.py

+1 -8

@@ -48,17 +48,11 @@ class Exaone4_5MultiTokenPredictor(ExaoneMoeMultiTokenPredictor):

49 model_config = vllm_config.model_config

50 quant_config = vllm_config.quant_config

51 - lora_config = vllm_config.lora_config

51 config = model_config.hf_config

52 text_config = config.text_config

54 self.config = config

56 - lora_vocab = (

57 - (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))

58 - if lora_config

59 - else 0

60 - )

61 - self.vocab_size = config.vocab_size + lora_vocab

55 + self.vocab_size = config.vocab_size

56 self.org_vocab_size = config.vocab_size

58 self.mtp_start_layer_idx = text_config.num_hidden_layers

@@ -67,7 +61,6 @@ class Exaone4_5MultiTokenPredictor(ExaoneMoeMultiTokenPredictor):

61 self.embed_tokens = VocabParallelEmbedding(

62 self.vocab_size,

63 text_config.hidden_size,

70 - org_num_embeddings=config.vocab_size,

64 )

66 self.fc = ColumnParallelLinear(

vllm/model_executor/models/exaone_moe.py

+3 -20

@@ -41,7 +41,6 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor

41 from vllm.model_executor.layers.quantization import QuantizationConfig

42 from vllm.model_executor.layers.rotary_embedding import get_rope

43 from vllm.model_executor.layers.vocab_parallel_embedding import (

44 - DEFAULT_VOCAB_PADDING_SIZE,

44 ParallelLMHead,

45 VocabParallelEmbedding,

46 )

@@ -112,6 +111,7 @@ class ExaoneMoe(nn.Module):

111 self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts

112 self.n_local_physical_experts = self.n_physical_experts // self.ep_size

113

114 + self.shared_experts: ExaoneMoeGatedMLP | None

115 if getattr(config, "num_shared_experts", 0) > 0:

116 intermediate_size = config.moe_intermediate_size * config.num_shared_experts

117 self.shared_experts = ExaoneMoeGatedMLP(

@@ -292,7 +292,7 @@ class ExaoneMoeDecoderLayer(nn.Module):

292 config: PretrainedConfig,

293 cache_config: CacheConfig | None = None,

294 quant_config: QuantizationConfig | None = None,

295 - is_mtp: bool = None,

295 + is_mtp: bool = False,

296 prefix: str = "",

297 ) -> None:

298 super().__init__()

@@ -381,26 +381,19 @@ class ExaoneMoeModel(nn.Module):

381 config = vllm_config.model_config.hf_config

382 cache_config = vllm_config.cache_config

383 quant_config = vllm_config.quant_config

384 - lora_config = vllm_config.lora_config

384 self.num_redundant_experts = (

385 vllm_config.parallel_config.eplb_config.num_redundant_experts

386 )

387

388 self.config = config

389 self.quant_config = quant_config

391 - lora_vocab = (

392 - (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))

393 - if lora_config

394 - else 0

395 - )

396 - self.vocab_size = config.vocab_size + lora_vocab

390 + self.vocab_size = config.vocab_size

391 if get_pp_group().is_first_rank or (

392 config.tie_word_embeddings and get_pp_group().is_last_rank

393 ):

394 self.embed_tokens = VocabParallelEmbedding(

395 self.vocab_size,

396 config.hidden_size,

403 - org_num_embeddings=config.vocab_size,

397 quant_config=quant_config,

398 )

399 else:

@@ -497,10 +490,8 @@ class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

490 super().__init__()

491 config = vllm_config.model_config.hf_config.get_text_config()

492 quant_config = vllm_config.quant_config

500 - lora_config = vllm_config.lora_config

493

494 self.config = config

503 - self.lora_config = lora_config

495 self.quant_config = quant_config

496

497 self.model = ExaoneMoeModel(

@@ -509,17 +500,9 @@ class ExaoneMoeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

500 )

501 if get_pp_group().is_last_rank:

502 self.unpadded_vocab_size = config.vocab_size

512 - if lora_config:

513 - self.unpadded_vocab_size += lora_config.lora_extra_vocab_size

503 self.lm_head = ParallelLMHead(

504 self.unpadded_vocab_size,

505 config.hidden_size,

517 - org_num_embeddings=config.vocab_size,

518 - padding_size=DEFAULT_VOCAB_PADDING_SIZE

519 - # We need bigger padding if using lora for kernel

520 - # compatibility

521 - if not lora_config

522 - else lora_config.lora_vocab_padding_size,

506 quant_config=quant_config,

507 )

508 if config.tie_word_embeddings:

vllm/model_executor/models/exaone_moe_mtp.py

+1 -8

@@ -46,16 +46,10 @@ class ExaoneMoeMultiTokenPredictor(nn.Module):

47 model_config = vllm_config.model_config

48 quant_config = vllm_config.quant_config

49 - lora_config = vllm_config.lora_config

49 config = model_config.hf_config

51 self.config = config

53 - lora_vocab = (

54 - (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1))

55 - if lora_config

56 - else 0

57 - )

58 - self.vocab_size = config.vocab_size + lora_vocab

52 + self.vocab_size = config.vocab_size

53 self.org_vocab_size = config.vocab_size

55 self.mtp_start_layer_idx = config.num_hidden_layers

@@ -64,7 +58,6 @@ class ExaoneMoeMultiTokenPredictor(nn.Module):

58 self.embed_tokens = VocabParallelEmbedding(

59 self.vocab_size,

60 config.hidden_size,

67 - org_num_embeddings=config.vocab_size,

61 )

63 self.fc = ColumnParallelLinear(

vllm/model_executor/models/extract_hidden_states.py

+4 -2

@@ -231,7 +231,7 @@ class CacheOnlyAttentionLayer(nn.Module, AttentionLayerBase):

231 head_size: int,

232 cache_config: CacheConfig | None = None,

233 prefix: str = "",

234 - attn_type: str = AttentionType.DECODER,

234 + attn_type: AttentionType = AttentionType.DECODER,

235 ):

236 super().__init__()

237

@@ -329,7 +329,9 @@ class ExtractHiddenStatesModel(nn.Module):

329 super().__init__()

330

331 self.vllm_config = vllm_config

332 - self.hf_config = vllm_config.speculative_config.draft_model_config.hf_config

332 + speculative_config = vllm_config.speculative_config

333 + assert speculative_config is not None

334 + self.hf_config = speculative_config.draft_model_config.hf_config

335 self.hidden_size = vllm_config.model_config.get_hidden_size()

336 self.target_num_hidden_layers = (

337 vllm_config.model_config.get_total_num_hidden_layers()

vllm/model_executor/models/falcon.py

@@ -413,6 +413,7 @@ class FalconModel(nn.Module):

413 else:

414 hidden_states = self.embed_input_ids(input_ids)

415 else:

416 + assert intermediate_tensors is not None

417 hidden_states = intermediate_tensors["hidden_states"]

418 for layer in islice(self.h, self.start_layer, self.end_layer):

419 hidden_states = layer(positions, hidden_states)

vllm/model_executor/models/fireredasr2.py

+2 -1

@@ -12,7 +12,7 @@ from transformers import (

12 )

14 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig

15 -from vllm.config.multimodal import BaseDummyOptions

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

16 from vllm.config.speech_to_text import SpeechToTextParams

17 from vllm.inputs import MultiModalDataDict, PromptType

18 from vllm.logger import init_logger

@@ -216,6 +216,7 @@ class FireRedASR2DummyInputsBuilder(BaseDummyInputsBuilder[FireRedASR2Processing

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

217

218 audio_overrides = mm_options.get("audio")

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

220

221 ret = {

222 "audio": self._get_dummy_audios(

vllm/model_executor/models/fireredlid.py

+8 -11

@@ -18,14 +18,13 @@ from __future__ import annotations

18 from collections.abc import Iterable, Mapping, Sequence

19 from typing import Annotated, Literal

21 -import numpy as np

21 import torch

22 from torch import nn

23 from transformers import BatchFeature

25 from vllm.config import ModelConfig, VllmConfig

27 -from vllm.config.multimodal import BaseDummyOptions

28 -from vllm.config.speech_to_text import SpeechToTextConfig

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

27 +from vllm.config.speech_to_text import SpeechToTextConfig, SpeechToTextParams

28 from vllm.distributed import get_tensor_model_parallel_world_size

29 from vllm.inputs import MultiModalDataDict, PromptType

30 from vllm.logger import init_logger

@@ -431,6 +430,7 @@ class FireRedLIDDummyInputsBuilder(BaseDummyInputsBuilder[FireRedLIDProcessingIn

430 audio_len = feature_extractor.chunk_length * sampling_rate

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

432 audio_overrides = mm_options.get("audio")

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

434 return {

435 "audio": self._get_dummy_audios(

436 length=audio_len,

@@ -747,13 +747,7 @@ class FireRedLIDForConditionalGeneration(

747 @classmethod

748 def get_generation_prompt(

749 cls,

750 - audio: np.ndarray,

751 - stt_config: SpeechToTextConfig,

752 - model_config: ModelConfig,

753 - language: str | None,

754 - task_type: Literal["transcribe", "translate"],

755 - request_prompt: str,

756 - to_language: str | None,

750 + stt_params: SpeechToTextParams,

751 ) -> PromptType:

752 """Build the prompt for the FireRedLID encoder-decoder model.

753

@@ -764,7 +758,10 @@ class FireRedLIDForConditionalGeneration(

758 "encoder_prompt": {

759 "prompt": "",

760 "multi_modal_data": {

767 - "audio": (audio, int(stt_config.sample_rate)),

761 + "audio": (

762 + stt_params.audio,

763 + int(stt_params.stt_config.sample_rate),

764 + ),

765 },

766 },

767 "decoder_prompt": {

vllm/model_executor/models/funasr.py

+8 -5

@@ -14,7 +14,7 @@ from transformers import (

14 )

16 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig

17 -from vllm.config.multimodal import BaseDummyOptions

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

18 from vllm.config.speech_to_text import SpeechToTextParams

19 from vllm.distributed import get_tensor_model_parallel_world_size

20 from vllm.inputs import MultiModalDataDict, PromptType

@@ -261,8 +261,8 @@ class SinusoidalPositionEncoder(torch.nn.Module):

261

262 def encode(

263 self,

264 - positions: torch.Tensor = None,

265 - depth: int = None,

264 + positions: torch.Tensor,

265 + depth: int,

266 dtype: torch.dtype = torch.float32,

267 ):

268 batch_size = positions.size(0)

@@ -573,8 +573,8 @@ class Transformer(nn.Module):

573 ]

574 )

575

576 - def forward(self, hidden_states: torch.Tensor, ilens: int = 0):

577 - max_len = max(ilens)

576 + def forward(self, hidden_states: torch.Tensor, ilens: torch.Tensor):

577 + max_len = ilens.max()

578 hidden_states = hidden_states[:, :max_len, :]

579 batch_size, seq_len, dim = hidden_states.size()

580 chunk_num = (seq_len - 1) // self.k + 1

@@ -682,6 +682,8 @@ class FunASRModel(nn.Module):

682 self.feat_permute = False

683

684 if self.feat_permute:

685 + if not isinstance(speech, torch.Tensor):

686 + raise NotImplementedError("List audio inputs cannot be permuted")

687 encoder_out, encoder_out_lens = self.encoder.audio_encoder(

688 speech.permute(0, 2, 1), speech_lengths

689 )

@@ -739,6 +741,7 @@ class FunASRDummyInputsBuilder(BaseDummyInputsBuilder[FunASRProcessingInfo]):

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

742

743 audio_overrides = mm_options.get("audio")

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

745

746 return {

747 "audio": self._get_dummy_audios(

vllm/model_executor/models/funaudiochat.py

+3 -1

@@ -26,7 +26,7 @@ from transformers.feature_extraction_utils import BatchFeature

26 from transformers.modeling_outputs import BaseModelOutput

28 from vllm.config import VllmConfig

29 -from vllm.config.multimodal import BaseDummyOptions

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

30 from vllm.inputs import MultiModalDataDict

31 from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention

32 from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear

@@ -600,6 +600,7 @@ class FunAudioChatDummyInputsBuilder(

600 num_audios = int(mm_counts.get("audio", 0))

601

602 audio_overrides = mm_options.get("audio")

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

604 return {

605 "audio": self._get_dummy_audios(

606 length=audio_len,

@@ -635,6 +636,7 @@ class FunAudioChatMultiModalProcessor(

636 processed_data = BatchFeature({"input_ids": input_ids})

637 processed_data.update(passthrough_data)

638 return processed_data

639 + assert isinstance(audios, Sequence)

640

641 feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs)

642 sr = int(feature_extractor.sampling_rate)