Merge branch 'main' into wentao-fix-mypy-models-a-b

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

tools/pre_commit/mypy.py

@@ -104,8 +104,6 @@ SEPARATE_GROUPS = [
 
 # TODO(woosuk): Include the code from Megatron and HuggingFace.
 EXCLUDE = [
-    r"vllm/model_executor/models/[aA]",
-    r"vllm/model_executor/models/[bB]",
     r"vllm/model_executor/models/[cC]",
     r"vllm/model_executor/models/[dD]",
     r"vllm/model_executor/models/[eE]",

vllm/model_executor/layers/mamba/mamba_utils.py

@@ -144,7 +144,7 @@ class MambaStateShapeCalculator:
         num_heads: int,
         tp_size: int,
         head_dim: int,
-    ) -> tuple[tuple[int, int, int], ...]:
+    ) -> tuple[tuple[int, int, int]]:
         state_shape = (num_heads // tp_size, head_dim, head_dim)
         return (state_shape,)
 

vllm/model_executor/models/AXK1.py

@@ -110,8 +110,12 @@ class AXK1MoE(nn.Module):
         self.ep_group = get_ep_group().device_group
         self.ep_rank = get_ep_group().rank_in_group
         self.ep_size = self.ep_group.size()
+        assert config.n_routed_experts is not None
+        assert config.num_experts_per_tok is not None
+        assert config.scoring_func is not None
+        assert config.hidden_act is not None
         self.n_routed_experts: int = config.n_routed_experts
-        self.n_shared_experts: int = config.n_shared_experts
+        self.n_shared_experts: int | None = config.n_shared_experts
 
         self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
 
@@ -245,7 +249,7 @@ class AXK1Attention(nn.Module):
         qk_nope_head_dim: int,
         qk_rope_head_dim: int,
         v_head_dim: int,
-        q_lora_rank: int,
+        q_lora_rank: int | None,
         kv_lora_rank: int,
         max_position_embeddings: int = 8192,
         cache_config: CacheConfig | None = None,
@@ -282,7 +286,7 @@ class AXK1Attention(nn.Module):
             )
             self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
             self.q_b_proj = ColumnParallelLinear(
-                q_lora_rank,
+                self.q_lora_rank,
                 self.num_heads * self.qk_head_dim,
                 bias=False,
                 quant_config=quant_config,
@@ -320,6 +324,7 @@ class AXK1Attention(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.o_proj",
         )
+        assert config.rope_parameters is not None
         if config.rope_parameters["rope_type"] != "default":
             config.rope_parameters["rope_type"] = (
                 "deepseek_yarn"
@@ -492,6 +497,7 @@ class AXK1MLAAttention(nn.Module):
             prefix=f"{prefix}.o_proj",
         )
 
+        assert config.rope_parameters is not None
         if config.rope_parameters["rope_type"] != "default":
             config.rope_parameters["rope_type"] = (
                 "deepseek_yarn"
@@ -573,6 +579,13 @@ class AXK1DecoderLayer(nn.Module):
         parallel_config = vllm_config.parallel_config
         self.config = config
 
+        assert config.max_position_embeddings is not None
+        assert config.hidden_act is not None
+        assert config.routed_scaling_factor is not None
+        assert config.qk_nope_head_dim is not None
+        assert config.qk_rope_head_dim is not None
+        assert config.v_head_dim is not None
+        assert config.kv_lora_rank is not None
         self.hidden_size = config.hidden_size
         max_position_embeddings = config.max_position_embeddings
         # DecoderLayers are created with `make_layers` which passes the prefix
@@ -588,6 +601,7 @@ class AXK1DecoderLayer(nn.Module):
         use_mha = all(dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim))
         self.use_mha = use_mha
 
+        attn_cls: type[nn.Module]
         if use_mha:
             attn_cls = DeepseekAttention
         elif model_config.use_mla:
@@ -635,6 +649,7 @@ class AXK1DecoderLayer(nn.Module):
         self.routed_scaling_factor = config.routed_scaling_factor
 
     def _is_layer_sparse(self) -> bool:
+        assert self.config.moe_layer_freq is not None
         return (
             self.config.n_routed_experts is not None
             and self.layer_idx >= self.config.first_k_dense_replace
@@ -789,7 +804,7 @@ class AXK1Model(nn.Module):
         rocm_aiter_moe_shared_expert_enabled = (
             rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
         )
-        stacked_params_mapping = [
+        stacked_params_mapping: list[tuple[str, str, int | str]] = [
             # (param_name, shard_name, shard_id)
             ("gate_up_proj", "gate_proj", 0),
             ("gate_up_proj", "up_proj", 1),
@@ -810,6 +825,7 @@ class AXK1Model(nn.Module):
 
         # Params for weights, fp8 weight scales, fp8 activation scales
         # (param_name, weight_name, expert_id, shard_id)
+        assert self.config.n_routed_experts is not None
         expert_params_mapping = fused_moe_make_expert_params_mapping(
             self,
             ckpt_gate_proj_name="gate_proj",
@@ -817,7 +833,7 @@ class AXK1Model(nn.Module):
             ckpt_up_proj_name="up_proj",
             num_experts=self.config.n_routed_experts
             + (
-                self.config.n_shared_experts
+                (self.config.n_shared_experts or 0)
                 if rocm_aiter_moe_shared_expert_enabled
                 else 0
             ),
@@ -926,7 +942,7 @@ class AXK1Model(nn.Module):
                     # param and delegate to its expert-aware weight_loader
                     # with expert_id.
                     for mapping in expert_params_mapping:
-                        param_name, weight_name, expert_id, shard_id = mapping
+                        param_name, weight_name, expert_id, expert_shard_id = mapping
                         if weight_name not in chunk_name:
                             continue
 
@@ -952,7 +968,7 @@ class AXK1Model(nn.Module):
                             param,
                             weight_to_load,
                             name_mapped,
-                            shard_id=shard_id,
+                            shard_id=expert_shard_id,
                             expert_id=expert_id,
                             return_success=True,
                         )
@@ -974,9 +990,10 @@ class AXK1Model(nn.Module):
                             continue
 
                         # Remapping the name of FP8 kv-scale.
-                        name = maybe_remap_kv_scale_name(name, params_dict)
-                        if name is None:
+                        remapped_name = maybe_remap_kv_scale_name(name, params_dict)
+                        if remapped_name is None:
                             continue
+                        name = remapped_name
 
                         if is_pp_missing_parameter(name, self):
                             continue
@@ -1014,7 +1031,7 @@ class AXK1MixtureOfExperts(MixtureOfExperts):
             self.num_physical_experts = example_moe.n_physical_experts
             self.num_local_physical_experts = example_moe.n_local_physical_experts
             self.num_routed_experts = example_moe.n_routed_experts
-            self.num_shared_experts = example_moe.n_shared_experts
+            self.num_shared_experts = example_moe.n_shared_experts or 0
             self.num_redundant_experts = example_moe.n_redundant_experts
 
     def update_physical_experts_metadata(
@@ -1132,6 +1149,7 @@ class AXK1ForCausalLM(
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         # Params for weights, fp8 weight scales, fp8 activation scales
         # (param_name, weight_name, expert_id, shard_id)
+        assert self.config.n_routed_experts is not None
         return fused_moe_make_expert_params_mapping(
             self,
             ckpt_gate_proj_name="gate_proj",

vllm/model_executor/models/adapters.py

@@ -4,7 +4,7 @@
 import itertools
 from collections.abc import Iterable
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, TypeVar, cast
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import torch
 import torch.nn as nn
@@ -25,7 +25,7 @@ if TYPE_CHECKING:
     from vllm.config import ModelConfig, VllmConfig
     from vllm.model_executor.layers.pooler import Pooler
 
-_T = TypeVar("_T", bound=type[nn.Module])
+_T = TypeVar("_T", bound=nn.Module)
 
 logger = init_logger(__name__)
 
@@ -45,7 +45,7 @@ def _load_st_projector(model_config: "ModelConfig") -> nn.Module | None:
     )
 
     if dense_modules is None:
-        return
+        return None
 
     try:
         layers = []
@@ -126,14 +126,14 @@ def _get_pooling_model_name(orig_model_name: str, pooling_suffix: str) -> str:
     return model_name + pooling_suffix
 
 
-def _create_pooling_model_cls(orig_cls: _T) -> _T:
+def _create_pooling_model_cls(orig_cls: type[_T]) -> type[_T]:
     # Lazy import
     from vllm.model_executor.layers.logits_processor import LogitsProcessor
     from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 
     from .utils import AutoWeightsLoader, StageMissingLayer, no_init_weights
 
-    class ModelForPooling(orig_cls, VllmModelForPooling):
+    class ModelForPooling(orig_cls, VllmModelForPooling):  # type: ignore[valid-type,misc]
         is_pooling_model = True
 
         def __init__(
@@ -148,16 +148,19 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T:
                 lambda mod: StageMissingLayer("output", mod),
                 targets=(LogitsProcessor, ParallelLMHead),
             ):
-                super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
+                super().__init__(  # type: ignore[safe-super]
+                    vllm_config=vllm_config, prefix=prefix, **kwargs
+                )
 
             # Used by SEQ_CLS_LOAD_METHODS
             self.vllm_config = vllm_config
 
             # If the model already defines a pooler instance, don't overwrite it
             pooler = getattr(self, "pooler", None)
-            if not pooler and supports_multimodal(self):
+            multimodal_model: object = self
+            if not pooler and supports_multimodal(multimodal_model):
                 # Try to get the pooler from the LM backbone
-                language_model = self.get_language_model()
+                language_model = multimodal_model.get_language_model()
                 if hasattr(language_model, "pooler"):
                     pooler = language_model.pooler
 
@@ -173,7 +176,9 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T:
         ) -> "Pooler":
             raise NotImplementedError
 
-        def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        def _load_pooling_model_weights(
+            self, weights: Iterable[tuple[str, torch.Tensor]]
+        ):
             params_dict = dict(self.named_parameters())
 
             # We support loading from both `*ForCausalLM` and `*Model`
@@ -224,10 +229,13 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T:
             load_weights = getattr(super(), "load_weights", default_load_weights)
             return load_weights(mapped_weights)
 
-    return ModelForPooling  # type: ignore
+        def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+            return self._load_pooling_model_weights(weights)
+
+    return ModelForPooling
 
 
-def as_embedding_model(cls: _T) -> _T:
+def as_embedding_model(cls: type[_T]) -> type[_T]:
     """
     Subclass an existing vLLM model to support embeddings.
 
@@ -245,7 +253,7 @@ def as_embedding_model(cls: _T) -> _T:
     # Lazy import
     from vllm.model_executor.layers.pooler import DispatchPooler
 
-    class ModelForEmbedding(_create_pooling_model_cls(cls)):
+    class ModelForEmbedding(_create_pooling_model_cls(cls)):  # type: ignore[misc]
         def _init_pooler(
             self,
             vllm_config: "VllmConfig",
@@ -258,10 +266,10 @@ def as_embedding_model(cls: _T) -> _T:
 
     ModelForEmbedding.__name__ = _get_pooling_model_name(cls.__name__, "ForEmbedding")
 
-    return ModelForEmbedding  # type: ignore
+    return ModelForEmbedding
 
 
-def as_seq_cls_model(cls: _T) -> _T:
+def as_seq_cls_model(cls: type[_T]) -> type[_T]:
     """
     Subclass an existing vLLM model to support classify and score tasks.
 
@@ -285,7 +293,8 @@ def as_seq_cls_model(cls: _T) -> _T:
     from .utils import maybe_prefix
 
     class ModelForSequenceClassification(
-        _create_pooling_model_cls(cls), SupportsCrossEncoding
+        _create_pooling_model_cls(cls),  # type: ignore[misc]
+        SupportsCrossEncoding,
     ):
         def _init_pooler(
             self,
@@ -366,7 +375,7 @@ def as_seq_cls_model(cls: _T) -> _T:
         cls.__name__, "ForSequenceClassification"
     )
 
-    return ModelForSequenceClassification  # type: ignore
+    return ModelForSequenceClassification
 
 
 class SequenceClassificationConfig(VerifyAndUpdateConfig):
@@ -400,14 +409,15 @@ class SequenceClassificationConfig(VerifyAndUpdateConfig):
         text_config.use_sep_token = use_sep_token
 
 
-def _get_language_model_for_seq_cls(model) -> nn.Module:
+def _get_language_model_for_seq_cls(model: nn.Module) -> nn.Module:
     """
     Get the language model component for sequence classification conversion.
     For VLMs, returns the inner language model. For standard LLMs, returns model itself.
     """
-    if supports_multimodal(model):
+    multimodal_model: object = model
+    if supports_multimodal(multimodal_model):
         try:
-            lm = model.get_language_model()
+            lm = multimodal_model.get_language_model()
             if lm is not model:
                 return lm
         except Exception:
@@ -481,12 +491,11 @@ def load_weights_using_from_2_way_softmax(
     hf_config = model.config
     text_config = hf_config.get_text_config()
 
-    tokens = getattr(
+    tokens: list[str] = getattr(
         hf_config,
         "classifier_from_token",
         getattr(text_config, "classifier_from_token", []),
     )
-    tokens = cast(list[int], tokens)
     assert len(tokens) == 2
 
     language_model = _get_language_model_for_seq_cls(model)
@@ -510,12 +519,7 @@ def load_weights_using_from_2_way_softmax(
         language_model.lm_head = language_model.lm_head.tie_weights(embed_tokens)
 
     with _disable_seq_cls_loading_on_inner_model(language_model, is_vlm):
-        # ModelForPooling is dynamically defined inside the _create_pooling_model_cls
-        # function, so we need use this hacky method to obtain it.
-        pooling_model_cls = next(
-            x for x in type(model).__mro__ if x.__name__ == "ModelForPooling"
-        )
-        loaded_weights = pooling_model_cls.load_weights(model, weights)
+        loaded_weights = model._load_pooling_model_weights(weights)
 
     from vllm.tokenizers import get_tokenizer
 
@@ -559,8 +563,7 @@ def load_weights_no_post_processing(model, weights: Iterable[tuple[str, torch.Te
     model_config = model.vllm_config.model_config
     text_config = model.config.get_text_config()
 
-    tokens = getattr(text_config, "classifier_from_token", [])
-    tokens = cast(list[int], tokens)
+    tokens: list[str] = getattr(text_config, "classifier_from_token", [])
     assert len(tokens) > 0
 
     language_model = _get_language_model_for_seq_cls(model)
@@ -584,11 +587,8 @@ def load_weights_no_post_processing(model, weights: Iterable[tuple[str, torch.Te
         language_model.lm_head = language_model.lm_head.tie_weights(embed_tokens)
 
     with _disable_seq_cls_loading_on_inner_model(language_model, is_vlm):
-        pooling_model_cls = next(
-            x for x in type(model).__mro__ if x.__name__ == "ModelForPooling"
-        )
-        # Skip ModelForSequenceClassification in MRO to avoid infinite recursion
-        loaded_weights = pooling_model_cls.load_weights(model, weights)
+        # Bypass ModelForSequenceClassification to avoid infinite recursion
+        loaded_weights = model._load_pooling_model_weights(weights)
 
     from vllm.tokenizers import get_tokenizer
 

vllm/model_executor/models/afmoe.py

@@ -47,7 +47,6 @@ from vllm.model_executor.models.utils import (
     PPMissingLayer,
     WeightsMapper,
     extract_layer_index,
-    make_empty_intermediate_tensors_factory,
     make_layers,
     maybe_prefix,
 )
@@ -232,7 +231,7 @@ class AfmoeAttention(nn.Module):
 
         # Only create rotary embeddings for local attention
         if self.is_local_attention:
-            self.rotary_emb = get_rope(
+            self.rotary_emb: nn.Module | None = get_rope(
                 self.head_dim,
                 max_position=max_position_embeddings,
                 rope_parameters=config.rope_parameters,
@@ -409,10 +408,6 @@ class AfmoeModel(nn.Module, EagleModelMixin):
         else:
             self.norm = PPMissingLayer()
 
-        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
-            ["hidden_states", "residual"], config.hidden_size
-        )
-
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)
 

vllm/model_executor/models/aimv2.py

@@ -180,7 +180,9 @@ class AIMv2Transformer(nn.Module):
             ]
         )
         if require_post_norm:
-            self.post_trunk_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+            self.post_trunk_norm: RMSNorm | None = RMSNorm(
+                config.hidden_size, eps=config.rms_norm_eps
+            )
         else:
             self.post_trunk_norm = None
 

vllm/model_executor/models/arctic.py

@@ -501,7 +501,7 @@ class ArcticModel(nn.Module):
                 weight_loader(param, loaded_weight, shard_id)
                 break
             else:
-                for param_name, weight_name, shard_id in mlp_params_mapping:
+                for param_name, weight_name, mlp_shard_id in mlp_params_mapping:
                     if weight_name not in name:
                         continue
                     name = name.replace(weight_name, param_name)
@@ -509,10 +509,10 @@ class ArcticModel(nn.Module):
                         continue
                     param = params_dict[name]
                     weight_loader = param.weight_loader
-                    weight_loader(param, loaded_weight, shard_id)
+                    weight_loader(param, loaded_weight, mlp_shard_id)
                     break
                 else:
-                    for param_name, weight_name, shard_id in expert_params_mapping:
+                    for param_name, weight_name, expert_id in expert_params_mapping:
                         if weight_name not in name:
                             continue
                         name = name.replace(weight_name, param_name)
@@ -521,7 +521,7 @@ class ArcticModel(nn.Module):
                         param = params_dict[name]
                         weight_loader = param.weight_loader
                         weight_loader(
-                            param, loaded_weight, weight_name, expert_id=shard_id
+                            param, loaded_weight, weight_name, expert_id=expert_id
                         )
                         break
                     else:

vllm/model_executor/models/aria.py

@@ -10,7 +10,7 @@ from transformers.models.aria.modeling_aria import AriaCrossAttention
 from transformers.models.aria.processing_aria import AriaProcessor
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.fused_moe import FusedMoE
@@ -317,7 +317,8 @@ class AriaDummyInputsBuilder(BaseDummyInputsBuilder[AriaProcessingInfo]):
         num_images = mm_counts.get("image", 0)
 
         processor = self.info.get_hf_processor()
-        image_token: str = processor.tokenizer.image_token  # type: ignore
+        image_token = getattr(processor.tokenizer, "image_token", None)
+        assert isinstance(image_token, str)
 
         return image_token * num_images
 
@@ -333,6 +334,7 @@ class AriaDummyInputsBuilder(BaseDummyInputsBuilder[AriaProcessingInfo]):
         num_images = mm_counts.get("image", 0)
 
         image_overrides = mm_options.get("image")
+        assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)
 
         return {
             "image": self._get_dummy_images(

vllm/model_executor/models/audioflamingo3.py

@@ -17,7 +17,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence, Sized
 from typing import Annotated, Any, Literal, TypeAlias
 
 import torch
@@ -30,7 +30,7 @@ from transformers.models.audioflamingo3 import (
 from transformers.models.qwen2_audio import Qwen2AudioEncoder
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import AudioDummyOptions, BaseDummyOptions
 from vllm.inputs import ModalityData, MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.models.module_mapping import MultiModelKeys
@@ -226,6 +226,7 @@ class AudioFlamingo3DummyInputsBuilder(
         audio_len = int(hf_processor.max_audio_len * sampling_rate)
         num_audios = mm_counts.get("audio", 0)
         audio_overrides = mm_options.get("audio")
+        assert audio_overrides is None or isinstance(audio_overrides, AudioDummyOptions)
 
         return {
             "audio": self._get_dummy_audios(
@@ -317,7 +318,7 @@ def _count_audio_tokens_from_mask(
         if isinstance(chunk_counts, torch.Tensor):
             counts = chunk_counts.tolist()
         elif chunk_counts and isinstance(chunk_counts[0], torch.Tensor):
-            counts = [count.item() for count in chunk_counts]
+            counts = [int(count) for count in chunk_counts]
         else:
             counts = chunk_counts
 
@@ -375,7 +376,7 @@ class AudioFlamingo3MultiModalProcessor(
     def _call_hf_processor(
         self,
         prompt: str,
-        mm_data: dict[str, object],
+        mm_data: Mapping[str, object],
         mm_kwargs: Mapping[str, Any],
         tok_kwargs: Mapping[str, object],
     ) -> BatchFeature:
@@ -412,7 +413,8 @@ class AudioFlamingo3MultiModalProcessor(
         chunk_counts = []
         for audio in audio_list:
             # audio is numpy array or list
-            n_samples = len(audio) if isinstance(audio, list) else audio.shape[0]
+            assert isinstance(audio, Sized)
+            n_samples = len(audio)
 
             n_win = max(1, (n_samples + window_size - 1) // window_size)
             if n_win > max_windows:
@@ -457,6 +459,7 @@ class AudioFlamingo3MultiModalProcessor(
                 )
             else:
                 audio_embeds = out_mm_data["audio_embeds"][item_idx]
+                assert isinstance(audio_embeds, torch.Tensor)
                 num_features = audio_embeds.shape[0]
 
             if num_features == 0:
@@ -556,8 +559,8 @@ class AudioFlamingo3ForConditionalGeneration(
     def _process_audio_input(
         self, audio_input: AudioFlamingo3Inputs
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
-        if audio_input["type"] == "audio_embeds":
-            audio_embeds = audio_input["audio_embeds"]
+        if audio_input.type == "audio_embeds":
+            audio_embeds = audio_input.audio_embeds
             return tuple(audio_embeds)
 
         (

vllm/model_executor/models/bagel.py

@@ -14,7 +14,7 @@ import torch
 import torch.nn as nn
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import get_act_fn
@@ -258,6 +258,7 @@ class BagelDummyInputsBuilder(BaseDummyInputsBuilder[BagelProcessingInfo]):
         # Use the configured image size
         image_size = vit_config.image_size
         image_overrides = mm_options.get("image")
+        assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)
 
         return {
             "image": self._get_dummy_images(

vllm/model_executor/models/bailing_moe.py

@@ -182,7 +182,7 @@ class BailingMLP(nn.Module):
         intermediate_size: int,
         config: PretrainedConfig,
         quant_config: QuantizationConfig | None = None,
-        reduce_results: bool | None = True,
+        reduce_results: bool = True,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -279,7 +279,7 @@ class BailingMoE(nn.Module):
             else:
                 intermediate_size = config.moe_intermediate_size
             intermediate_size *= config.num_shared_experts
-            self.shared_experts = BailingMLP(
+            self.shared_experts: BailingMLP | None = BailingMLP(
                 intermediate_size=intermediate_size,
                 config=config,
                 quant_config=quant_config,

vllm/model_executor/models/bailing_moe_linear.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import copy
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 
 import torch
 import torch.nn as nn
@@ -133,19 +133,21 @@ class BailingMoeV25MLAAttention(nn.Module):
 
         if self.q_lora_rank is not None:
             # Use fused_qkv_a_proj when q_lora_rank is set
-            self.fused_qkv_a_proj = MergedColumnParallelLinear(
-                self.hidden_size,
-                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
-                bias=False,
-                quant_config=quant_config,
-                prefix=f"{prefix}.fused_qkv_a_proj",
-                disable_tp=True,
+            self.fused_qkv_a_proj: MergedColumnParallelLinear | None = (
+                MergedColumnParallelLinear(
+                    self.hidden_size,
+                    [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
+                    bias=False,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.fused_qkv_a_proj",
+                    disable_tp=True,
+                )
             )
-            self.q_a_layernorm = RMSNorm(
+            self.q_a_layernorm: RMSNorm | None = RMSNorm(
                 self.q_lora_rank,
                 eps=config.rms_norm_eps,
             )
-            self.q_b_proj = ColumnParallelLinear(
+            self.q_b_proj: ColumnParallelLinear | None = ColumnParallelLinear(
                 self.q_lora_rank,
                 self.num_heads * self.qk_head_dim,
                 bias=False,
@@ -311,6 +313,8 @@ class BailingMoeV25(nn.Module):
                 "score_function and correction_bias should be "
                 "(softmax, None) or (sigmoid, not None)"
             )
+        else:
+            self.score_function = "softmax"
 
         # Shared experts (using BailingMLP)
         if self.num_shared_experts > 0:
@@ -319,7 +323,7 @@ class BailingMoeV25(nn.Module):
             else:
                 intermediate_size = config.moe_intermediate_size
             intermediate_size *= config.num_shared_experts
-            self.shared_experts = BailingMLP(
+            self.shared_experts: BailingMLP | None = BailingMLP(
                 intermediate_size=intermediate_size,
                 config=config,
                 quant_config=quant_config,
@@ -613,7 +617,9 @@ class BailingMoeV25Model(nn.Module):
                 return False
 
             param = params_dict[name]
-            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader: Callable[..., None] = getattr(
+                param, "weight_loader", default_weight_loader
+            )
 
             if shard_id is None:
                 weight_loader(param, tensor)
@@ -691,11 +697,16 @@ class BailingMoeV25Model(nn.Module):
                         continue
 
                 # Routed experts
-                for param_name, weight_name, expert_id, shard_id in expert_mappings:
+                for (
+                    param_name,
+                    weight_name,
+                    expert_id,
+                    expert_shard_id,
+                ) in expert_mappings:
                     if weight_name not in norm_name:
                         continue
                     mapped = norm_name.replace(weight_name, param_name)
-                    if load_param(mapped, weight, (expert_id, shard_id)):
+                    if load_param(mapped, weight, (expert_id, expert_shard_id)):
                         break
                 continue
 
@@ -783,7 +794,7 @@ class BailingMoeV25ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):
     def get_mamba_state_shape_from_config(
         cls,
         vllm_config: VllmConfig,
-    ) -> tuple[tuple[int, ...], ...]:
+    ) -> tuple[tuple[int, int, int]]:
         """Calculate shape for linear attention cache."""
         config = vllm_config.model_config.hf_config
         tp_size = vllm_config.parallel_config.tensor_parallel_size

vllm/model_executor/models/bailing_moe_mtp.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Inference-only Bailing MoE v2.5 MTP model."""
 
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 
 import torch
 import torch.nn as nn
@@ -268,14 +268,17 @@ class BailingMoeV25MTPModel(nn.Module):
             loaded_weight: torch.Tensor,
             shard_id=None,
         ) -> bool:
-            name = maybe_remap_kv_scale_name(name, params_dict)
-            if name is None:
+            remapped_name = maybe_remap_kv_scale_name(name, params_dict)
+            if remapped_name is None:
                 return False
+            name = remapped_name
             if name not in params_dict or is_pp_missing_parameter(name, self):
                 return False
 
             param = params_dict[name]
-            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader: Callable[..., None] = getattr(
+                param, "weight_loader", default_weight_loader
+            )
             if shard_id is None:
                 weight_loader(param, loaded_weight)
             elif isinstance(shard_id, int):
@@ -351,14 +354,14 @@ class BailingMoeV25MTPModel(nn.Module):
 
             if "mlp.experts" in name:
                 for mapping in expert_params_mapping:
-                    param_name, weight_name, expert_id, shard_id = mapping
+                    param_name, weight_name, expert_id, expert_shard_id = mapping
                     if weight_name not in name:
                         continue
                     mapped_name = name.replace(weight_name, param_name)
                     if load_param(
                         mapped_name,
                         loaded_weight,
-                        (expert_id, shard_id),
+                        (expert_id, expert_shard_id),
                     ):
                         loaded = True
                         break

vllm/model_executor/models/bee.py

@@ -8,7 +8,7 @@ import torch.nn as nn
 from transformers.activations import GELUActivation
 
 from vllm.config import VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.multimodal import MULTIMODAL_REGISTRY
 
@@ -97,6 +97,7 @@ class BeeDummyInputsBuilder(LlavaDummyInputsBuilder[BeeProcessingInfo]):
         target_width, target_height = self.info.get_image_size_with_most_features()
 
         image_overrides = mm_options.get("image")
+        assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)
 
         return {
             "image": self._get_dummy_images(

vllm/model_executor/models/bert.py

@@ -97,6 +97,7 @@ class BertPooler(SequencePooler):
     def __init__(self, model_config: ModelConfig):
         pooler_config = model_config.pooler_config
         assert pooler_config is not None
+        assert pooler_config.seq_pooling_type is not None
 
         config: BertConfig = model_config.hf_config
 

vllm/model_executor/models/bert_with_rope.py

@@ -57,8 +57,8 @@ class BertWithRopeEmbedding(nn.Module):
             config.vocab_size, config.hidden_size
         )
         if config.type_vocab_size > 0:
-            self.token_type_embeddings = VocabParallelEmbedding(
-                config.type_vocab_size, config.hidden_size
+            self.token_type_embeddings: VocabParallelEmbedding | None = (
+                VocabParallelEmbedding(config.type_vocab_size, config.hidden_size)
             )
         else:
             self.token_type_embeddings = None
@@ -127,6 +127,7 @@ class BertWithRopeAttention(nn.Module):
             prefix=f"{prefix}.qkv_proj",
         )
 
+        assert rotary_kwargs is not None
         self.rotary_emb = get_rope(**rotary_kwargs)
 
         self.attn = EncoderOnlyAttention(
@@ -466,7 +467,7 @@ class BertWithRope(nn.Module, SupportsQuant):
         )
 
         if add_pooling_layer:
-            self.pooler = BertPooler(vllm_config.model_config)
+            self.pooler: BertPooler | None = BertPooler(vllm_config.model_config)
         else:
             self.pooler = None
 
@@ -632,7 +633,7 @@ class JinaRobertaModel(BertWithRope):
         scaling = self.config.lora_alpha / self.config.lora_rank
         device = self.vllm_config.device_config.device
 
-        weights = {name: weight for name, weight in weights}
+        weights_dict = {name: weight for name, weight in weights}
 
         o = ".original"
         a = ".0.lora_A"
@@ -641,34 +642,34 @@ class JinaRobertaModel(BertWithRope):
         # text-matching
         i = -1
 
-        for name in list(weights.keys()):
+        for name in list(weights_dict.keys()):
             if o in name:
-                dtype = weights[name].dtype
-                shape = weights[name].shape
+                dtype = weights_dict[name].dtype
+                shape = weights_dict[name].shape
                 weight_name = name[: -len(o)]
 
                 if "embeddings" in weight_name:
-                    B = weights[weight_name + a][i].to(device).float()
-                    A = weights[weight_name + b][i].to(device).float()
+                    B = weights_dict[weight_name + a][i].to(device).float()
+                    A = weights_dict[weight_name + b][i].to(device).float()
                 else:
-                    B = weights[weight_name + b][i].to(device).float()
-                    A = weights[weight_name + a][i].to(device).float()
+                    B = weights_dict[weight_name + b][i].to(device).float()
+                    A = weights_dict[weight_name + a][i].to(device).float()
 
                 weight = (
-                    weights[weight_name + o].to(device)
+                    weights_dict[weight_name + o].to(device)
                     + torch.matmul(B, A).view(shape) * scaling
                 )
                 weight = weight.cpu().to(dtype)
 
-                weights[weight_name.replace(".parametrizations", "")] = weight
+                weights_dict[weight_name.replace(".parametrizations", "")] = weight
 
                 del (
-                    weights[weight_name + o],
-                    weights[weight_name + a],
-                    weights[weight_name + b],
+                    weights_dict[weight_name + o],
+                    weights_dict[weight_name + a],
+                    weights_dict[weight_name + b],
                 )
 
-        return [(name, weight) for name, weight in weights.items()]
+        return [(name, weight) for name, weight in weights_dict.items()]
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         weights = self.jina_merge_lora_weights(weights)

vllm/model_executor/models/blip2.py

@@ -14,7 +14,7 @@ from transformers import (
 )
 
 from vllm.config import CacheConfig, VllmConfig
-from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
 from vllm.inputs import MultiModalDataDict
 from vllm.model_executor.layers.activation import get_act_fn
 from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -454,6 +454,7 @@ class Blip2DummyInputsBuilder(BaseDummyInputsBuilder[Blip2ProcessingInfo]):
         num_images = mm_counts.get("image", 0)
 
         image_overrides = mm_options.get("image")
+        assert image_overrides is None or isinstance(image_overrides, ImageDummyOptions)
 
         return {
             "image": self._get_dummy_images(
@@ -620,8 +621,8 @@ class Blip2ForConditionalGeneration(
         return self._image_pixels_to_features(self.vision_model, pixel_values)
 
     def _process_image_input(self, image_input: Blip2ImageInputs) -> torch.Tensor:
-        if image_input["type"] == "image_embeds":
-            return image_input["data"]
+        if image_input.type == "image_embeds":
+            return image_input.data
 
         image_features = self._process_image_pixels(image_input)
 

vllm/model_executor/models/idefics2_vision_model.py

@@ -19,6 +19,7 @@
 """PyTorch Idefics2 model."""
 
 from collections.abc import Iterable
+from typing import ClassVar
 
 import torch
 from torch import nn
@@ -352,7 +353,7 @@ class Idefics2Encoder(nn.Module):
 
 
 class Idefics2VisionTransformer(nn.Module):
-    hf_to_vllm_mapper = WeightsMapper(
+    hf_to_vllm_mapper: ClassVar[WeightsMapper] = WeightsMapper(
         orig_to_new_stacked={
             ".q_proj": (".qkv_proj", "q"),
             ".k_proj": (".qkv_proj", "k"),

vllm/model_executor/models/interfaces.py

@@ -70,6 +70,13 @@ The output embeddings must be one of the following formats:
 - A single 3D tensor, with the batch dimension grouping the 2D tensors.
 """
 
+MambaStateShapes: TypeAlias = (
+    tuple[tuple[int, int]]
+    | tuple[tuple[int, int, int]]
+    | tuple[tuple[int, int], tuple[int, int]]
+    | tuple[tuple[int, int], tuple[int, int, int]]
+)
+
 
 class StreamingTranscriptionPostProcessor:
     """Stateful streaming post-processor for transcription deltas."""
@@ -566,7 +573,7 @@ class SupportsLoRA(Protocol):
     # The `embedding_module` and `embedding_padding_modules`
     # are empty by default.
     embedding_modules: ClassVar[dict[str, str]] = {}
-    packed_modules_mapping: dict[str, list[str]] = {}
+    packed_modules_mapping: ClassVar[dict[str, list[str]]] = {}
     # Module prefixes to skip during LoRA loading (e.g., ["mtp."] for MTP layers)
     lora_skip_prefixes: ClassVar[list[str]] = []
     lora_manager: "LoRAModelManager | None"
@@ -628,6 +635,15 @@ def _supports_lora(model: type[object] | object) -> bool:
     return isinstance(model, SupportsLoRA)
 
 
+class _MakeEmptyIntermediateTensors(Protocol):
+    def __call__(
+        self,
+        batch_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> "IntermediateTensors": ...
+
+
 @runtime_checkable
 class SupportsPP(Protocol):
     """The interface required for all models that support pipeline parallel."""
@@ -641,14 +657,8 @@ class SupportsPP(Protocol):
         MRO of your model class.
     """
 
-    def make_empty_intermediate_tensors(
-        self,
-        batch_size: int,
-        dtype: torch.dtype,
-        device: torch.device,
-    ) -> "IntermediateTensors":
-        """Called when PP rank > 0 for profiling purposes."""
-        ...
+    make_empty_intermediate_tensors: _MakeEmptyIntermediateTensors
+    """Called when PP rank > 0 for profiling purposes."""
 
     def forward(
         self,
@@ -656,7 +666,7 @@ class SupportsPP(Protocol):
         positions: Tensor,
         *,
         intermediate_tensors: "IntermediateTensors | None",
-    ) -> "IntermediateTensors | None":
+    ) -> "Tensor | IntermediateTensors | tuple[Tensor, list[Tensor]]":
         """
         Accept [`IntermediateTensors`][vllm.sequence.IntermediateTensors] when
         PP rank > 0.
@@ -673,12 +683,7 @@ class SupportsPP(Protocol):
 class _SupportsPPType(Protocol):
     supports_pp: Literal[True]
 
-    def make_empty_intermediate_tensors(
-        self,
-        batch_size: int,
-        dtype: torch.dtype,
-        device: torch.device,
-    ) -> "IntermediateTensors": ...
+    make_empty_intermediate_tensors: _MakeEmptyIntermediateTensors
 
     def forward(
         self,
@@ -686,7 +691,7 @@ class _SupportsPPType(Protocol):
         positions: Tensor,
         *,
         intermediate_tensors: "IntermediateTensors | None",
-    ) -> "Tensor | IntermediateTensors": ...
+    ) -> "Tensor | IntermediateTensors | tuple[Tensor, list[Tensor]]": ...
 
 
 @overload
@@ -817,16 +822,14 @@ class IsHybrid(Protocol):
     def get_mamba_state_shape_from_config(
         cls,
         vllm_config: "VllmConfig",
-    ) -> tuple[tuple[int, int], tuple[int, int, int]]:
+    ) -> MambaStateShapes:
         """Calculate shapes for Mamba's convolutional and state caches.
 
         Args:
             vllm_config: vLLM config
 
         Returns:
-            Tuple containing:
-            - conv_state_shape: Shape for convolutional state cache
-            - temporal_state_shape: Shape for state space model cache
+            Shapes for each state cache used by the model.
         """
         ...
 
@@ -1040,7 +1043,7 @@ class SupportsQuant:
     """The interface required for all models that support quantization."""
 
     hf_to_vllm_mapper: ClassVar["WeightsMapper | None"] = None
-    packed_modules_mapping: ClassVar[dict[str, list[str]] | None] = None
+    packed_modules_mapping: ClassVar[dict[str, list[str]]]
     quant_config: QuantizationConfig | None = None
 
     def __new__(cls, *args, **kwargs) -> Self:
@@ -1075,8 +1078,8 @@ class SupportsQuant:
         if (hf_to_vllm_mapper := self.hf_to_vllm_mapper) is not None:
             unstacked_mapper = hf_to_vllm_mapper.get_unstacked_mapper()
             self.quant_config.apply_vllm_mapper(unstacked_mapper)
-        if self.packed_modules_mapping is not None:
-            self.quant_config.packed_modules_mapping.update(self.packed_modules_mapping)
+        if packed_modules_mapping := getattr(self, "packed_modules_mapping", None):
+            self.quant_config.packed_modules_mapping.update(packed_modules_mapping)
 
 
 @runtime_checkable

vllm/model_executor/models/llama.py

@@ -26,6 +26,7 @@
 
 from collections.abc import Iterable
 from itertools import islice
+from typing import ClassVar
 
 import torch
 from torch import nn
@@ -342,7 +343,7 @@ class LlamaDecoderLayer(nn.Module):
     },
 )
 class LlamaModel(nn.Module, EagleModelMixin):
-    hf_to_vllm_mapper = WeightsMapper(
+    hf_to_vllm_mapper: ClassVar[WeightsMapper] = WeightsMapper(
         orig_to_new_stacked={
             # weight_name: (param_name, shard_id)
             ".q_proj": (".qkv_proj", "q"),

vllm/model_executor/models/llava_next.py

@@ -204,9 +204,7 @@ class BaseLlavaNextMultiModalProcessor(BaseLlavaMultiModalProcessor[_I]):
         raise NotImplementedError
 
 
-class LlavaNextMultiModalProcessor(
-    BaseLlavaNextMultiModalProcessor[LlavaNextProcessingInfo]
-):
+class LlavaNextMultiModalProcessor(BaseLlavaNextMultiModalProcessor[_I]):
     def _get_mm_fields_config(
         self,
         hf_inputs: BatchFeature,

vllm/models/inkling/nvidia/ops/fa4_rel_attention.py

@@ -130,13 +130,20 @@ def inkling_fa4_rel_attention(
     cute_window = (None, None) if window_size == (-1, -1) else window_size
 
     rel_logits = rel_logits.contiguous()
+    flash_attn_varlen_func: Callable[..., Any]
     if _use_sheared_bias():
-        from vllm.third_party.tml_fa4 import flash_attn_varlen_func
+        from vllm.third_party.tml_fa4 import (
+            flash_attn_varlen_func as tml_flash_attn_varlen_func,
+        )
 
+        flash_attn_varlen_func = tml_flash_attn_varlen_func
         bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits}
     else:
-        from vllm.vllm_flash_attn.cute import flash_attn_varlen_func
+        from vllm.vllm_flash_attn.cute import (
+            flash_attn_varlen_func as cute_flash_attn_varlen_func,
+        )
 
+        flash_attn_varlen_func = cute_flash_attn_varlen_func
         bias_kwargs = {
             "score_mod": _get_score_mod(rel_extent),
             "aux_tensors": [rel_logits],