karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

wentao-fix-mypy-models-a-b

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
yewentao256
2026-07-17 19:01:27 +00:00
parent c4dd6d78fd
commit d123d2b00a
20 changed files with 156 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/pre_commit/mypy.py
-2
View File
@@ -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/models/AXK1.py
+29 -11
View File
@@ -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",
vllm/model_executor/models/adapters.py
+19 -14
View File
@@ -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
vllm/model_executor/models/afmoe.py
+7 -5
View File
@@ -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,
vllm/model_executor/models/aimv2.py
+3 -1
View File
@@ -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/apertus.py
+1 -1
View File
@@ -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
)
vllm/model_executor/models/arcee.py
+1 -1
View File
@@ -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
)
vllm/model_executor/models/arctic.py
+5 -5
View File
@@ -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
)
vllm/model_executor/models/aria.py
+9 -6
View File
@@ -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]
)
}
vllm/model_executor/models/audioflamingo3.py
+9 -6
View File
@@ -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,
vllm/model_executor/models/bagel.py
+2 -2
View File
@@ -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
)
vllm/model_executor/models/bailing_moe.py
+5 -5
View File
@@ -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
)
vllm/model_executor/models/bailing_moe_linear.py
+24 -14
View File
@@ -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, ...], ...]:
vllm/model_executor/models/bailing_moe_mtp.py
+7 -5
View File
@@ -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
vllm/model_executor/models/bee.py
+2 -2
View File
@@ -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,
vllm/model_executor/models/bert.py
+1
View File
@@ -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
+18 -17
View File
@@ -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
+4 -4
View File
@@ -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(
vllm/model_executor/models/bloom.py
+1 -1
View File
@@ -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
)
vllm/models/inkling/nvidia/ops/fa4_rel_attention.py
+9 -2
View File
@@ -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],