From d123d2b00ad522b2a6fe175ea82b3d0bb46a4cc6 Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Fri, 17 Jul 2026 19:01:27 +0000 Subject: [PATCH] wentao-fix-mypy-models-a-b Signed-off-by: yewentao256 --- tools/pre_commit/mypy.py | 2 - vllm/model_executor/models/AXK1.py | 40 ++++++++++++++----- vllm/model_executor/models/adapters.py | 33 ++++++++------- vllm/model_executor/models/afmoe.py | 12 +++--- vllm/model_executor/models/aimv2.py | 4 +- vllm/model_executor/models/apertus.py | 2 +- vllm/model_executor/models/arcee.py | 2 +- vllm/model_executor/models/arctic.py | 10 ++--- vllm/model_executor/models/aria.py | 15 ++++--- vllm/model_executor/models/audioflamingo3.py | 15 ++++--- vllm/model_executor/models/bagel.py | 4 +- vllm/model_executor/models/bailing_moe.py | 10 ++--- .../models/bailing_moe_linear.py | 38 +++++++++++------- vllm/model_executor/models/bailing_moe_mtp.py | 12 +++--- vllm/model_executor/models/bee.py | 4 +- vllm/model_executor/models/bert.py | 1 + vllm/model_executor/models/bert_with_rope.py | 35 ++++++++-------- vllm/model_executor/models/blip2.py | 8 ++-- vllm/model_executor/models/bloom.py | 2 +- .../inkling/nvidia/ops/fa4_rel_attention.py | 11 ++++- 20 files changed, 156 insertions(+), 104 deletions(-) diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index c52e0e7ce63..48a054fbc4d 100755 --- a/tools/pre_commit/mypy.py +++ b/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]", diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index a465c6b5632..73ff5d6300f 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -109,8 +109,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 @@ -244,7 +248,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, @@ -281,7 +285,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, @@ -319,6 +323,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" @@ -491,6 +496,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" @@ -572,6 +578,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 @@ -587,6 +600,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: @@ -634,6 +648,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 @@ -788,7 +803,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), @@ -809,6 +824,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", @@ -816,7 +832,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 ), @@ -925,7 +941,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 @@ -951,7 +967,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, ) @@ -973,9 +989,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 @@ -1013,7 +1030,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( @@ -1078,7 +1095,7 @@ class AXK1ForCausalLM( else: self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.model.make_empty_intermediate_tensors ) # Set MoE hyperparameters @@ -1131,6 +1148,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", diff --git a/vllm/model_executor/models/adapters.py b/vllm/model_executor/models/adapters.py index cccd849681f..676f6096ad8 100644 --- a/vllm/model_executor/models/adapters.py +++ b/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 @@ -45,7 +45,7 @@ def _load_st_projector(model_config: "ModelConfig") -> nn.Module | None: ) if dense_modules is None: - return + return None try: layers = [] @@ -133,7 +133,7 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T: 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,7 +148,9 @@ 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 @@ -157,7 +159,7 @@ def _create_pooling_model_cls(orig_cls: _T) -> _T: pooler = getattr(self, "pooler", None) if not pooler and supports_multimodal(self): # Try to get the pooler from the LM backbone - language_model = self.get_language_model() + language_model = self.get_language_model() # type: ignore[call-arg] if hasattr(language_model, "pooler"): pooler = language_model.pooler @@ -245,7 +247,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", @@ -285,7 +287,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, @@ -407,7 +410,7 @@ def _get_language_model_for_seq_cls(model) -> nn.Module: """ if supports_multimodal(model): try: - lm = model.get_language_model() + lm = model.get_language_model() # type: ignore[call-arg] if lm is not model: return lm except Exception: @@ -481,12 +484,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) @@ -515,7 +517,9 @@ def load_weights_using_from_2_way_softmax( 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 = pooling_model_cls.load_weights( # type: ignore[attr-defined] + model, 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) @@ -588,7 +591,9 @@ def load_weights_no_post_processing(model, weights: Iterable[tuple[str, torch.Te 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) + loaded_weights = pooling_model_cls.load_weights( # type: ignore[attr-defined] + model, weights + ) from vllm.tokenizers import get_tokenizer diff --git a/vllm/model_executor/models/afmoe.py b/vllm/model_executor/models/afmoe.py index 0122d019588..f456229b2f7 100644 --- a/vllm/model_executor/models/afmoe.py +++ b/vllm/model_executor/models/afmoe.py @@ -232,7 +232,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,8 +409,10 @@ 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 + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: @@ -525,7 +527,7 @@ class AfmoeForCausalLM( else: self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.model.make_empty_intermediate_tensors ) # Set MoE hyperparameters @@ -576,7 +578,7 @@ class AfmoeForCausalLM( def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) - def forward( + def forward( # type: ignore[override] self, input_ids: torch.Tensor | None, positions: torch.Tensor, diff --git a/vllm/model_executor/models/aimv2.py b/vllm/model_executor/models/aimv2.py index bc4b3e8cd08..eaff05c4f0d 100644 --- a/vllm/model_executor/models/aimv2.py +++ b/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 diff --git a/vllm/model_executor/models/apertus.py b/vllm/model_executor/models/apertus.py index b997e153a99..ffda816bb47 100644 --- a/vllm/model_executor/models/apertus.py +++ b/vllm/model_executor/models/apertus.py @@ -454,7 +454,7 @@ class ApertusForCausalLM( else: self.lm_head = PPMissingLayer() - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.model.make_empty_intermediate_tensors ) diff --git a/vllm/model_executor/models/arcee.py b/vllm/model_executor/models/arcee.py index c32a903bba8..185509eb7bc 100644 --- a/vllm/model_executor/models/arcee.py +++ b/vllm/model_executor/models/arcee.py @@ -326,7 +326,7 @@ class ArceeForCausalLM( # Placeholder for lm_head on non-last ranks self.lm_head = PPMissingLayer() - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.model.make_empty_intermediate_tensors ) diff --git a/vllm/model_executor/models/arctic.py b/vllm/model_executor/models/arctic.py index fee9eaaac8d..3922c1025a4 100644 --- a/vllm/model_executor/models/arctic.py +++ b/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: @@ -562,7 +562,7 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant): self.num_experts_per_tok = config.num_experts_per_tok self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.model.make_empty_intermediate_tensors ) diff --git a/vllm/model_executor/models/aria.py b/vllm/model_executor/models/aria.py index 6b723883423..da9a4470ca9 100644 --- a/vllm/model_executor/models/aria.py +++ b/vllm/model_executor/models/aria.py @@ -79,7 +79,9 @@ class AriaImagePixelInputs(TensorSchema): ] -class AriaVisionTransformer(Idefics3VisionTransformer, SupportsQuant): +class AriaVisionTransformer( # type: ignore[misc] + Idefics3VisionTransformer, SupportsQuant +): packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} def __init__( @@ -218,7 +220,7 @@ class AriaProjector(nn.Module): class AriaRoutedExperts(RoutedExperts): - def weight_loader( + def weight_loader( # type: ignore[override] self, param: nn.Parameter, loaded_weight: torch.Tensor, shard_id: str ) -> None: # Override the weight_loader to handle the expert weights in the Aria @@ -326,7 +328,7 @@ class AriaTextDecoderLayer(LlamaDecoderLayer): ) -class AriaTextModel(LlamaModel, SupportsQuant): +class AriaTextModel(LlamaModel, SupportsQuant): # type: ignore[misc] """ Custom LlamaModel for the AriaMoE model which modifies the standard LlamaModel by replacing the `LlamaDecoderLayer` with `MoEDecoderLayer`. @@ -386,9 +388,10 @@ class AriaTextModel(LlamaModel, SupportsQuant): if name.endswith(".bias") and name not in params_dict: 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 @@ -445,7 +448,7 @@ class AriaDummyInputsBuilder(BaseDummyInputsBuilder[AriaProcessingInfo]): width=max_image_size, height=max_image_size, num_images=num_images, - overrides=image_overrides, + overrides=image_overrides, # type: ignore[arg-type] ) } diff --git a/vllm/model_executor/models/audioflamingo3.py b/vllm/model_executor/models/audioflamingo3.py index e03baeda0cd..5a471bfaf4c 100644 --- a/vllm/model_executor/models/audioflamingo3.py +++ b/vllm/model_executor/models/audioflamingo3.py @@ -231,7 +231,7 @@ class AudioFlamingo3DummyInputsBuilder( "audio": self._get_dummy_audios( length=audio_len, num_audios=num_audios, - overrides=audio_overrides, + overrides=audio_overrides, # type: ignore[arg-type] ) } @@ -317,7 +317,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 +375,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 +412,7 @@ 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] + n_samples = len(audio) # type: ignore[arg-type] n_win = max(1, (n_samples + window_size - 1) // window_size) if n_win > max_windows: @@ -457,6 +457,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: @@ -523,7 +524,7 @@ class AudioFlamingo3ForConditionalGeneration( architectures=["Qwen2ForCausalLM"], ) - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.language_model.make_empty_intermediate_tensors ) @@ -564,7 +565,9 @@ class AudioFlamingo3ForConditionalGeneration( input_features, feature_attention_mask, chunk_counts, - ) = self._normalize_audio_feature_inputs(audio_input) + ) = self._normalize_audio_feature_inputs( + audio_input # type: ignore[arg-type] + ) audio_hidden_states = self._encode_audio_features( input_features, feature_attention_mask, diff --git a/vllm/model_executor/models/bagel.py b/vllm/model_executor/models/bagel.py index 97cfe75e67d..fc18751ec86 100644 --- a/vllm/model_executor/models/bagel.py +++ b/vllm/model_executor/models/bagel.py @@ -264,7 +264,7 @@ class BagelDummyInputsBuilder(BaseDummyInputsBuilder[BagelProcessingInfo]): width=image_size, height=image_size, num_images=num_images, - overrides=image_overrides, + overrides=image_overrides, # type: ignore[arg-type] ), } @@ -430,7 +430,7 @@ class BagelForConditionalGeneration( self.connector = StageMissingLayer("image_tower") self.vit_pos_embed = StageMissingLayer("image_tower") - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.language_model.make_empty_intermediate_tensors ) diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py index a45d0ca81de..e73f715d1a8 100644 --- a/vllm/model_executor/models/bailing_moe.py +++ b/vllm/model_executor/models/bailing_moe.py @@ -186,7 +186,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__() @@ -283,7 +283,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, @@ -516,7 +516,7 @@ class BailingMoeModel(nn.Module): break else: 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 name = name.replace(weight_name, param_name) @@ -531,7 +531,7 @@ class BailingMoeModel(nn.Module): param, loaded_weight, name, - shard_id=shard_id, + shard_id=expert_shard_id, expert_id=expert_id, ) break @@ -596,7 +596,7 @@ class BailingMoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA): else: self.lm_head = PPMissingLayer() - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.model.make_empty_intermediate_tensors ) diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index 5cc057c3b45..c7658294e81 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy from collections.abc import Iterable +from typing import Any import torch import torch.nn as nn @@ -133,19 +134,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 +314,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 +324,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 +618,7 @@ class BailingMoeV25Model(nn.Module): return False param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader: Any = getattr(param, "weight_loader", default_weight_loader) if shard_id is None: weight_loader(param, tensor) @@ -691,11 +696,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 @@ -780,7 +790,7 @@ class BailingMoeV25ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP): ) @classmethod - def get_mamba_state_shape_from_config( + def get_mamba_state_shape_from_config( # type: ignore[override] cls, vllm_config: VllmConfig, ) -> tuple[tuple[int, ...], ...]: diff --git a/vllm/model_executor/models/bailing_moe_mtp.py b/vllm/model_executor/models/bailing_moe_mtp.py index da6b1ddb8b6..a5aed9d245a 100644 --- a/vllm/model_executor/models/bailing_moe_mtp.py +++ b/vllm/model_executor/models/bailing_moe_mtp.py @@ -3,6 +3,7 @@ """Inference-only Bailing MoE v2.5 MTP model.""" from collections.abc import Iterable +from typing import Any import torch import torch.nn as nn @@ -265,14 +266,15 @@ 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: Any = getattr(param, "weight_loader", default_weight_loader) if shard_id is None: weight_loader(param, loaded_weight) elif isinstance(shard_id, int): @@ -348,14 +350,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 diff --git a/vllm/model_executor/models/bee.py b/vllm/model_executor/models/bee.py index af44c34f48e..06e90d486d1 100644 --- a/vllm/model_executor/models/bee.py +++ b/vllm/model_executor/models/bee.py @@ -103,7 +103,7 @@ class BeeDummyInputsBuilder(LlavaDummyInputsBuilder[BeeProcessingInfo]): width=target_width, height=target_height, num_images=num_images, - overrides=image_overrides, + overrides=image_overrides, # type: ignore[arg-type] ), } @@ -133,7 +133,7 @@ class BeeMultiModalProjector(nn.Module): return hidden_states -@MULTIMODAL_REGISTRY.register_processor( +@MULTIMODAL_REGISTRY.register_processor( # type: ignore[misc] LlavaNextMultiModalProcessor, info=BeeProcessingInfo, dummy_inputs=BeeDummyInputsBuilder, diff --git a/vllm/model_executor/models/bert.py b/vllm/model_executor/models/bert.py index 225f6c2bda7..67301130545 100644 --- a/vllm/model_executor/models/bert.py +++ b/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 diff --git a/vllm/model_executor/models/bert_with_rope.py b/vllm/model_executor/models/bert_with_rope.py index 020ef993474..97f7fc1e786 100644 --- a/vllm/model_executor/models/bert_with_rope.py +++ b/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) diff --git a/vllm/model_executor/models/blip2.py b/vllm/model_executor/models/blip2.py index 49d2a5a89f5..b14c73f4d16 100644 --- a/vllm/model_executor/models/blip2.py +++ b/vllm/model_executor/models/blip2.py @@ -460,7 +460,7 @@ class Blip2DummyInputsBuilder(BaseDummyInputsBuilder[Blip2ProcessingInfo]): width=max_image_size, height=max_image_size, num_images=num_images, - overrides=image_overrides, + overrides=image_overrides, # type: ignore[arg-type] ) } @@ -523,7 +523,7 @@ class Blip2MultiModalProcessor(BaseMultiModalProcessor[Blip2ProcessingInfo]): info=Blip2ProcessingInfo, dummy_inputs=Blip2DummyInputsBuilder, ) -class Blip2ForConditionalGeneration( +class Blip2ForConditionalGeneration( # type: ignore[misc] nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP, SupportsQuant ): @classmethod @@ -576,7 +576,7 @@ class Blip2ForConditionalGeneration( prefix=maybe_prefix(prefix, "language_model"), ) - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.language_model.make_empty_intermediate_tensors ) @@ -623,7 +623,7 @@ class Blip2ForConditionalGeneration( if image_input["type"] == "image_embeds": return image_input["data"] - image_features = self._process_image_pixels(image_input) + image_features = self._process_image_pixels(image_input) # type: ignore[arg-type] query_tokens = self.query_tokens.expand(image_features.shape[0], -1, -1) query_output = self.qformer( diff --git a/vllm/model_executor/models/bloom.py b/vllm/model_executor/models/bloom.py index 233028a905f..728a1ea55e0 100644 --- a/vllm/model_executor/models/bloom.py +++ b/vllm/model_executor/models/bloom.py @@ -349,7 +349,7 @@ class BloomForCausalLM(nn.Module, SupportsPP, SupportsQuant): ) self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.transformer.make_empty_intermediate_tensors ) diff --git a/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py b/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py index f547f650629..883545e3f2d 100644 --- a/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py +++ b/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],