diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index 585585e352..0feac74a52 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -113,8 +113,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/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index e3361ca7dc..50f6e059c0 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/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,) diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index 862604921e..a95b701baf 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/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", diff --git a/vllm/model_executor/models/adapters.py b/vllm/model_executor/models/adapters.py index 33b80fcc76..8f269fcabc 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 @@ -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()) # Use model-owned remapping only for the probe membership check # (e.g. Gemma4 model.language_model.* -> language_model.model.*) @@ -236,10 +241,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. @@ -257,7 +265,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", @@ -270,7 +278,7 @@ def as_embedding_model(cls: _T) -> _T: ModelForEmbedding.__name__ = _get_pooling_model_name(cls.__name__, "ForEmbedding") - return ModelForEmbedding # type: ignore + return ModelForEmbedding def _resolve_num_labels(hf_config: Any, text_config: Any) -> int: @@ -296,7 +304,7 @@ def _resolve_num_labels(hf_config: Any, text_config: Any) -> int: return text_config.num_labels -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. @@ -320,7 +328,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, @@ -401,7 +410,7 @@ def as_seq_cls_model(cls: _T) -> _T: cls.__name__, "ForSequenceClassification" ) - return ModelForSequenceClassification # type: ignore + return ModelForSequenceClassification class SequenceClassificationConfig(VerifyAndUpdateConfig): @@ -435,14 +444,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: @@ -516,12 +526,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) @@ -545,12 +554,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 @@ -594,8 +598,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) @@ -619,11 +622,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 diff --git a/vllm/model_executor/models/afmoe.py b/vllm/model_executor/models/afmoe.py index 6ed67861e9..6ffc571674 100644 --- a/vllm/model_executor/models/afmoe.py +++ b/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) diff --git a/vllm/model_executor/models/aimv2.py b/vllm/model_executor/models/aimv2.py index bc4b3e8cd0..eaff05c4f0 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/arctic.py b/vllm/model_executor/models/arctic.py index fee9eaaac8..b7e8e6796f 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: diff --git a/vllm/model_executor/models/aria.py b/vllm/model_executor/models/aria.py index d4423e83cf..f35a32fab5 100644 --- a/vllm/model_executor/models/aria.py +++ b/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 ( @@ -319,7 +319,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 @@ -335,6 +336,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( diff --git a/vllm/model_executor/models/audioflamingo3.py b/vllm/model_executor/models/audioflamingo3.py index e03baeda0c..e64d9e551b 100644 --- a/vllm/model_executor/models/audioflamingo3.py +++ b/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) ( diff --git a/vllm/model_executor/models/bagel.py b/vllm/model_executor/models/bagel.py index 97cfe75e67..e592dd3728 100644 --- a/vllm/model_executor/models/bagel.py +++ b/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( diff --git a/vllm/model_executor/models/bailing_moe.py b/vllm/model_executor/models/bailing_moe.py index 926f4b829e..0834de5ee5 100644 --- a/vllm/model_executor/models/bailing_moe.py +++ b/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, diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index 87e1d4dba0..8731e91c5d 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/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 diff --git a/vllm/model_executor/models/bailing_moe_mtp.py b/vllm/model_executor/models/bailing_moe_mtp.py index 53679edf89..a403244a17 100644 --- a/vllm/model_executor/models/bailing_moe_mtp.py +++ b/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 diff --git a/vllm/model_executor/models/bailing_moe_v3.py b/vllm/model_executor/models/bailing_moe_v3.py index 1bd3d2e14d..2085a192a9 100644 --- a/vllm/model_executor/models/bailing_moe_v3.py +++ b/vllm/model_executor/models/bailing_moe_v3.py @@ -9,7 +9,7 @@ vLLM's parallel linear layers, MLA kernel, KDA kernel and fused MoE loader. """ import copy -from collections.abc import Iterable +from collections.abc import Callable, Iterable import torch import torch.nn as nn @@ -75,7 +75,6 @@ from vllm.third_party.flash_linear_attention.ops.kda import ( fused_recurrent_kda_fwd, ) from vllm.utils.torch_utils import direct_register_custom_op -from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum @@ -136,7 +135,7 @@ def _is_kda_layer( def _get_kda_state_shape_for_config( vllm_config: VllmConfig, -) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]: +) -> tuple[tuple[int, int], tuple[int, int, int]]: config = vllm_config.model_config.hf_config num_spec = ( vllm_config.speculative_config.num_speculative_tokens @@ -205,7 +204,7 @@ class BailingMoeV3MLP(nn.Module): intermediate_size: int, config: PretrainedConfig, quant_config: QuantizationConfig | None = None, - reduce_results: bool | None = True, + reduce_results: bool = True, prefix: str = "", swiglu_limit: float | None = None, ) -> None: @@ -226,9 +225,7 @@ class BailingMoeV3MLP(nn.Module): prefix=f"{prefix}.down_proj", ) self.act_fn = ( - SwigluStepAndMul(limit=swiglu_limit) - if swiglu_limit not in (None, 0) - else SiluAndMul() + SwigluStepAndMul(limit=swiglu_limit) if swiglu_limit else SiluAndMul() ) def forward(self, x): @@ -264,6 +261,8 @@ class BailingMoeV3MLAAttention(nn.Module): self.num_local_heads = self.num_heads // tp_size self.scaling = self.qk_head_dim**-0.5 + self.q_proj: ColumnParallelLinear | None + self.kv_a_proj_with_mqa: ReplicatedLinear | None if self.q_lora_rank is None: self.q_proj = ColumnParallelLinear( self.hidden_size, @@ -400,7 +399,7 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase): def get_state_dtype( self, - ) -> tuple[torch.dtype, torch.dtype, torch.dtype, torch.dtype]: + ) -> tuple[torch.dtype, torch.dtype]: if self.model_config is None or self.cache_config is None: raise ValueError("model_config and cache_config must be set") return MambaStateDtypeCalculator.kda_state_dtype( @@ -409,7 +408,7 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase): def get_state_shape( self, - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + ) -> tuple[tuple[int, int], tuple[int, int, int]]: return MambaStateShapeCalculator.kda_state_shape( self.tp_size, self.num_heads, @@ -579,12 +578,12 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase): core_attn_out: torch.Tensor, ) -> None: forward_context = get_forward_context() - attn_metadata: AttentionMetadata = forward_context.attn_metadata - if attn_metadata is None: + attn_metadata_map = forward_context.attn_metadata + if attn_metadata_map is None: return - assert isinstance(attn_metadata, dict) - attn_metadata = attn_metadata[self.prefix] + assert isinstance(attn_metadata_map, dict) + attn_metadata = attn_metadata_map[self.prefix] assert isinstance(attn_metadata, GDNAttentionMetadata) has_initial_state = attn_metadata.has_initial_state spec_query_start_loc = attn_metadata.spec_query_start_loc @@ -825,6 +824,7 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase): assert q is not None and k is not None and v is not None assert g1_non_spec is not None and beta_non_spec is not None assert state_indices is not None + assert has_initial_state is not None zero_idx = state_indices[~has_initial_state] recurrent_state[zero_idx] = 0 initial_state = recurrent_state_active[state_indices].contiguous() @@ -846,6 +846,7 @@ class BailingMoeV3KimiDeltaAttention(PluggableLayer, MambaBase): assert q is not None and k is not None and v is not None assert g1_non_spec is not None and beta_non_spec is not None assert state_indices is not None + assert query_start_loc is not None out, _ = fused_recurrent_kda( q=q, k=k, @@ -1205,13 +1206,13 @@ class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP): @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: VllmConfig - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + ) -> tuple[tuple[int, int], tuple[int, int, int]]: return _get_kda_state_shape_for_config(vllm_config) @classmethod def get_mamba_state_dtype_from_config( cls, vllm_config: VllmConfig - ) -> tuple[torch.dtype, torch.dtype, torch.dtype, torch.dtype]: + ) -> tuple[torch.dtype, torch.dtype]: return MambaStateDtypeCalculator.kda_state_dtype( vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype, @@ -1243,7 +1244,9 @@ class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP): 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, tensor) elif isinstance(shard_id, int): @@ -1268,10 +1271,10 @@ class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP): if name is None: continue loaded = False - for param_suf, weight_suf, shard_id in stacked_mappings: + for param_suf, weight_suf, stacked_shard_id in stacked_mappings: if weight_suf in name: mapped = name.replace(weight_suf, param_suf) - if load_param(mapped, weight, shard_id): + if load_param(mapped, weight, stacked_shard_id): loaded = True break if loaded: diff --git a/vllm/model_executor/models/bailing_moe_v3_mtp.py b/vllm/model_executor/models/bailing_moe_v3_mtp.py index c89007b5b3..48c8c9d7a9 100644 --- a/vllm/model_executor/models/bailing_moe_v3_mtp.py +++ b/vllm/model_executor/models/bailing_moe_v3_mtp.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Bailing MoE v3 MTP model.""" -from collections.abc import Iterable +from collections.abc import Callable, Iterable import torch import torch.nn as nn @@ -233,11 +233,12 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP): self, input_ids: torch.Tensor | None, positions: torch.Tensor, - hidden_states: torch.Tensor, + hidden_states: torch.Tensor | None = None, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, spec_step_idx: int = 0, ) -> torch.Tensor: + assert hidden_states is not None del intermediate_tensors return self.model( input_ids, @@ -280,14 +281,17 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP): 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): @@ -347,13 +351,13 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP): name = normalize_name(name) loaded = False - for param_name, weight_name, shard_id in stacked_params_mapping: + for param_name, weight_name, stacked_shard_id in stacked_params_mapping: if weight_name not in name: continue if "mlp.experts." in name and name not in params_dict: continue mapped_name = name.replace(weight_name, param_name) - if load_param(mapped_name, loaded_weight, shard_id): + if load_param(mapped_name, loaded_weight, stacked_shard_id): loaded = True break if loaded: diff --git a/vllm/model_executor/models/bee.py b/vllm/model_executor/models/bee.py index af44c34f48..4a57413b6e 100644 --- a/vllm/model_executor/models/bee.py +++ b/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( diff --git a/vllm/model_executor/models/bert.py b/vllm/model_executor/models/bert.py index 225f6c2bda..6730113054 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 02ba333ca6..ba6b20ec5f 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 49d2a5a89f..baa1a16a59 100644 --- a/vllm/model_executor/models/blip2.py +++ b/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) diff --git a/vllm/model_executor/models/idefics2_vision_model.py b/vllm/model_executor/models/idefics2_vision_model.py index 81e7e51eb9..b6f5531074 100644 --- a/vllm/model_executor/models/idefics2_vision_model.py +++ b/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"), diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 66ac662dfa..8f365359ee 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -71,6 +71,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]] +) + @dataclass(frozen=True, slots=True) class DiarizedTranscriptionSegment: @@ -616,7 +623,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" @@ -678,6 +685,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.""" @@ -691,14 +707,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, @@ -706,7 +716,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. @@ -723,12 +733,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, @@ -736,7 +741,7 @@ class _SupportsPPType(Protocol): positions: Tensor, *, intermediate_tensors: "IntermediateTensors | None", - ) -> "Tensor | IntermediateTensors": ... + ) -> "Tensor | IntermediateTensors | tuple[Tensor, list[Tensor]]": ... @overload @@ -785,7 +790,7 @@ def supports_pp( def _supports_pp_attributes(model: type[object] | object) -> bool: if isinstance(model, type): - return isinstance(model, _SupportsPPType) + return SupportsPP in model.__mro__ or isinstance(model, _SupportsPPType) return isinstance(model, SupportsPP) @@ -867,16 +872,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. """ ... @@ -1090,7 +1093,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: @@ -1125,8 +1128,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 diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py index bb223a3114..cbf75593e0 100644 --- a/vllm/model_executor/models/llama.py +++ b/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"), diff --git a/vllm/model_executor/models/llava_next.py b/vllm/model_executor/models/llava_next.py index fdd1d41af5..44d50f434a 100644 --- a/vllm/model_executor/models/llava_next.py +++ b/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,