@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9a276d6375

[Mypy Fix] Mypy fix for "vllm/model_executor/models/[cC][dD]" (#52003)

Wentao Ye committed Aug 12, 2026 at 23:33 UTC 9a276d63753a5b05de394adbdf54680d70a7d906

24 files changed +153 -66

tools/pre_commit/mypy.py

-3

@@ -111,10 +111,7 @@ SEPARATE_GROUPS = [

111 "tests/v1/worker",

112 ]

113

114 -# TODO(woosuk): Include the code from Megatron and HuggingFace.

114 EXCLUDE = [

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

117 - r"vllm/model_executor/models/[dD]",

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

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

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

vllm/model_executor/models/clip.py

@@ -142,6 +142,7 @@ class CLIPProcessingInfo(BaseProcessingInfo):

142

143 pooler_config = self.ctx.model_config.pooler_config

144 assert pooler_config is not None

145 + assert pooler_config.seq_pooling_type is not None

146

147 return get_num_selected_vision_tokens(

148 vision_encoder_info.get_num_image_tokens(

@@ -854,6 +855,7 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):

855 feature_select_strategy: VisionFeatureSelectStrategy | None = None,

856 ) -> torch.Tensor:

857 if feature_select_strategy is None:

858 + assert self.pooler_config.seq_pooling_type is not None

859 feature_select_strategy = _get_vision_feature_select_strategy(

860 self.pooler_config.seq_pooling_type

861 )

@@ -961,6 +963,7 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):

963 raise RuntimeError("PP is not supported for this model")

964

965 # Multimodal inputs

966 + assert inputs_embeds is not None

967 if not self._is_text_input:

968 return inputs_embeds

969

vllm/model_executor/models/cohere2_moe.py

@@ -178,6 +178,7 @@ class Cohere2MoeAttention(nn.Module):

178 self.max_position_embeddings = getattr(

179 config, "model_max_length", None

180 ) or getattr(config, "max_position_embeddings", 8192)

181 + assert isinstance(self.max_position_embeddings, int)

182 self.qkv_proj = QKVParallelLinear(

183 self.hidden_size,

184 self.head_dim,

@@ -283,6 +284,7 @@ class Cohere2Moe(nn.Module):

284 prefix=f"{prefix}.gate",

285 )

286

287 + self.shared_experts: Cohere2MoeMLP | None

288 if hasattr(config, "num_shared_experts") and config.num_shared_experts > 0:

289 self.shared_experts = Cohere2MoeMLP(

290 config=config,

vllm/model_executor/models/cohere2_vision.py

+1 -1

@@ -416,7 +416,7 @@ class Cohere2VisionForConditionalGeneration(

416 )

417

418 def _patch_quant_config(

419 - self, config: PretrainedConfig, quant_config: QuantizationConfig

419 + self, config: PretrainedConfig, quant_config: QuantizationConfig | None

420 ):

421 # the awq models from OpenGVLab missing `modules_to_not_convert`

422 # patch the quant_config to add `modules_to_not_convert` back

vllm/model_executor/models/cohere_asr.py

+28 -13

@@ -55,6 +55,7 @@ from vllm.transformers_utils.processors.cohere_asr import (

55 CohereASRFeatureExtractor,

56 CohereASRProcessor,

57 )

58 +from vllm.utils.collection_utils import is_list_of

59 from vllm.v1.attention.backend import (

60 AttentionType,

61 )

@@ -183,6 +184,7 @@ class CohereASRAttention(nn.Module):

184 def forward(

185 self,

186 hidden_states: torch.Tensor,

187 + encoder_hidden_states: torch.Tensor | None = None,

188 ) -> torch.Tensor:

189 qkv, _ = self.qkv_proj(hidden_states)

190 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

@@ -241,7 +243,7 @@ class CohereASRCrossAttention(CohereASRAttention):

243 def forward(

244 self,

245 hidden_states: torch.Tensor,

244 - encoder_hidden_states: torch.Tensor | None,

246 + encoder_hidden_states: torch.Tensor | None = None,

247 ) -> torch.Tensor:

248 q, _ = self.q_proj(hidden_states)

249

@@ -837,7 +839,7 @@ class CausalConv1D(nn.Conv1d):

839 out_channels: int,

840 kernel_size: int,

841 stride: int = 1,

840 - padding: str | int = 0,

842 + padding: str | int | list[int] | None = 0,

843 dilation: int = 1,

844 groups: int = 1,

845 bias: bool = True,

@@ -845,6 +847,8 @@ class CausalConv1D(nn.Conv1d):

847 device=None,

848 dtype=None,

849 ) -> None:

850 + self._left_padding: int

851 + self._right_padding: int

852 if padding is None:

853 self._left_padding = kernel_size - 1

854 self._right_padding = stride - 1

@@ -903,7 +907,7 @@ class ConformerConvolution(nn.Module):

907 d_model: int,

908 kernel_size: int,

909 norm_type: str = "batch_norm",

906 - conv_context_size: int | None = None,

910 + conv_context_size: int | list[int] | None = None,

911 pointwise_activation: str = "glu_",

912 use_bias: bool = True,

913 ) -> None:

@@ -1158,6 +1162,7 @@ class RelPositionMultiHeadAttention(CohereASRMultiHeadAttention):

1162 q, k, v = self.forward_qkv(query, key, value)

1163 q = q.transpose(1, 2) # (batch, time1, head, d_k)

1164

1165 + assert pos_emb is not None

1166 n_batch_pos = pos_emb.size(0)

1167 p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)

1168 p = p.transpose(1, 2) # (batch, head, time1, d_k)

@@ -1213,7 +1218,7 @@ class ConformerLayer(torch.nn.Module):

1218 n_heads: int = 4,

1219 conv_kernel_size: int = 31,

1220 conv_norm_type: str = "batch_norm",

1216 - conv_context_size: int | None = None,

1221 + conv_context_size: int | list[int] | None = None,

1222 pos_bias_u: nn.Parameter | torch.Tensor | None = None,

1223 pos_bias_v: nn.Parameter | torch.Tensor | None = None,

1224 att_context_size: list[int] | None = None,

@@ -1382,6 +1387,7 @@ class ConformerEncoder(nn.Module):

1387 conv_kernel_size=conv_kernel_size,

1388 )

1389

1390 + self.xscale: float | None

1391 if xscaling:

1392 self.xscale = math.sqrt(d_model)

1393 else:

@@ -1633,9 +1639,12 @@ class ConformerEncoder(nn.Module):

1639 ) -> tuple[list[list[int]], list[int], list[float], list[int]]:

1640 # convert att_context_size to a standard list of lists

1641 if att_context_size:

1636 - att_context_size_all = list(att_context_size)

1637 - if isinstance(att_context_size_all[0], int):

1638 - att_context_size_all = [att_context_size_all]

1642 + if is_list_of(att_context_size, int, check="all"):

1643 + att_context_size_all = [att_context_size]

1644 + elif is_list_of(att_context_size, list, check="all"):

1645 + att_context_size_all = att_context_size

1646 + else:

1647 + raise ValueError("att_context_size cannot mix nested and flat values")

1648 for i, att_cs in enumerate(att_context_size_all):

1649 if att_context_style == "chunked_limited":

1650 if att_cs[0] > 0 and att_cs[0] % (att_cs[1] + 1) > 0:

@@ -1684,6 +1693,7 @@ class ConformerEncoder(nn.Module):

1693 if conv_context_size == "causal":

1694 conv_context_size = [conv_kernel_size - 1, 0]

1695 else:

1696 + assert isinstance(conv_context_size, list)

1697 total = conv_context_size[0] + conv_context_size[1] + 1

1698 if total != conv_kernel_size:

1699 raise ValueError(

@@ -1940,7 +1950,7 @@ class CohereASRMultiModalProcessor(EncDecMultiModalProcessor[CohereASRProcessing

1950 ):

1951 if mm_data:

1952 feature_extractor = self.info.get_feature_extractor(**mm_kwargs)

1943 - mm_data = dict(audio=mm_data.pop("audios"))

1953 + mm_data = dict(audio=dict(mm_data).pop("audios"))

1954 mm_kwargs = dict(

1955 **mm_kwargs,

1956 sampling_rate=feature_extractor.sampling_rate,

@@ -2045,7 +2055,6 @@ class CohereAsrForConditionalGeneration(

2055 # (which is different from "_") and encode("▁ABC") ignores the first token

2056 # so the prompt_text is unreliable. However, prompt_token_ids can be used

2057 # to get prompt_text but it wont have the first token "▁".

2048 - prompt_text = None

2058 prompt_token_ids = cls._get_default_prompt_token_ids(

2059 tokenizer,

2060 model_config,

@@ -2053,7 +2062,6 @@ class CohereAsrForConditionalGeneration(

2062 )

2063

2064 return TokensPrompt(

2056 - prompt=prompt_text,

2065 prompt_token_ids=prompt_token_ids,

2066 multi_modal_data={"audio": (audio, stt_config.sample_rate)},

2067 )

@@ -2199,7 +2207,9 @@ class CohereAsrForConditionalGeneration(

2207 audio_input, seq_lens = self._parse_and_validate_audio_input(**kwargs)

2208

2209 if hasattr(audio_input, "input_features"):

2202 - out = self.model.get_encoder_outputs(audio_input["input_features"])

2210 + out = self.model.get_encoder_outputs(

2211 + audio_input["input_features"], seq_lens

2212 + )

2213 else:

2214 out = self.model.get_encoder_outputs(audio_input, seq_lens)

2215

@@ -2219,12 +2229,17 @@ class CohereAsrForConditionalGeneration(

2229 f"Incorrect type of audio features. Got type: {type(input_features)}"

2230 )

2231

2232 + if not isinstance(length, torch.Tensor):

2233 + raise ValueError("Audio feature lengths must be a tensor.")

2234 +

2235 if isinstance(input_features, torch.Tensor):

2236 seq_lens = length.reshape(-1)

2237 else:

2238 + if not is_list_of(input_features, torch.Tensor, check="all"):

2239 + raise ValueError("All audio features must be tensors.")

2240 input_features = [

2226 - feat.to(self.dtype).squeeze(0).transpose(1, 0)

2227 - for feat in input_features

2241 + feature.to(self.dtype).squeeze(0).transpose(1, 0)

2242 + for feature in input_features

2243 ]

2244 seq_lens = length.reshape(-1)

2245 input_features = torch.nn.utils.rnn.pad_sequence(

vllm/model_executor/models/cohere_eagle.py

+3 -1

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

58 start_layer_id: int = 0,

59 ) -> None:

60 super().__init__()

61 + assert vllm_config.speculative_config is not None

62 self.config = vllm_config.speculative_config.draft_model_config.hf_config

63 self.quant_config = get_draft_quant_config(vllm_config)

@@ -137,6 +138,7 @@ class CohereEagleModel(nn.Module):

138 class EagleCohereForCausalLM(CohereForCausalLM):

139 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

140 nn.Module.__init__(self)

141 + assert vllm_config.speculative_config is not None

142 self.config = vllm_config.speculative_config.draft_model_config.hf_config

143 # Flags checked by the speculative proposer to decide whether to share

144 # embed_tokens / lm_head with the target model. Cohere EAGLE checkpoints

@@ -160,7 +162,7 @@ class EagleCohereForCausalLM(CohereForCausalLM):

162 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

163 return self.model.embed_input_ids(input_ids)

164

163 - def forward(

165 + def forward( # type: ignore[override]

166 self,

167 input_ids: torch.Tensor,

168 positions: torch.Tensor,

vllm/model_executor/models/colmodernvbert.py

@@ -33,6 +33,7 @@ from vllm.multimodal.processing import (

33 PromptUpdate,

34 )

35 from vllm.sequence import IntermediateTensors

36 +from vllm.tokenizers.hf import HfTokenizer

37 from vllm.transformers_utils.configs.colmodernvbert import ColModernVBertConfig

39 from .interfaces import (

@@ -161,6 +162,7 @@ class ColModernVBertMultiModalProcessor(

162 tok_kwargs: Mapping[str, object],

163 ) -> BatchFeature:

164 tokenizer = self.info.get_tokenizer()

165 + assert isinstance(tokenizer, HfTokenizer)

166 text_encoding = tokenizer(

167 prompt,

168 return_tensors="pt",

vllm/model_executor/models/colpali.py

+4 -1

@@ -25,6 +25,7 @@ from transformers import BatchFeature, PaliGemmaProcessor

26 from vllm.config import VllmConfig

27 from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed

28 +from vllm.model_executor.layers.pooler.tokwise.heads import TokenPoolerHead

29 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

30 from vllm.multimodal import MULTIMODAL_REGISTRY

@@ -232,7 +233,9 @@ class ColPaliModel(

233 loaded.add(f"custom_text_proj.{param_name}")

234

235 # Update pooler projector for the lazy-creation path

235 - self.pooler.head.projector = self.custom_text_proj

236 + head = self.pooler.head

237 + assert isinstance(head, TokenPoolerHead)

238 + head.projector = self.custom_text_proj

239

240 # Mark pooler projector params as loaded

241 if hasattr(self, "pooler") and hasattr(self.pooler, "head"):

vllm/model_executor/models/commandr.py

@@ -154,6 +154,7 @@ class CohereAttention(nn.Module):

154 self.max_position_embeddings = getattr(

155 config, "model_max_length", None

156 ) or getattr(config, "max_position_embeddings", 8192)

157 + assert isinstance(self.max_position_embeddings, int)

158 self.use_qk_norm = getattr(config, "use_qk_norm", False)

159 self.qkv_proj = QKVParallelLinear(

160 self.hidden_size,

vllm/model_executor/models/config.py

+11 -2

@@ -9,6 +9,7 @@ if TYPE_CHECKING:

9 from transformers import PretrainedConfig

11 from vllm.config import CacheConfig, ModelConfig, VllmConfig

12 + from vllm.config.cache import MambaDType

15 logger = init_logger(__name__)

@@ -452,6 +453,7 @@ class JambaForSequenceClassificationConfig(VerifyAndUpdateConfig):

453 @staticmethod

454 def verify_and_update_model_config(model_config: "ModelConfig") -> None:

455 pooler_config = model_config.pooler_config

456 + assert pooler_config is not None

457 if pooler_config.use_activation is None:

458 pooler_config.use_activation = False

459

@@ -518,6 +520,7 @@ class JinaVLForSequenceClassificationConfig(VerifyAndUpdateConfig):

520 config = model_config.hf_config

521 config.num_labels = 1

522 pooler_config = model_config.pooler_config

523 + assert pooler_config is not None

524 if pooler_config.logit_mean is None:

525 pooler_config.logit_mean = 2.65

526

@@ -540,6 +543,7 @@ class LlamaBidirectionalConfig(VerifyAndUpdateConfig):

543 if pooling_type is None:

544 raise ValueError(f"pool_type {hf_config.pooling!r} not supported")

545

546 + assert model_config.pooler_config is not None

547 model_config.pooler_config.seq_pooling_type = pooling_type

548

549

@@ -572,10 +576,13 @@ class LlamaNemotronVLConfig(VerifyAndUpdateConfig):

576 if pooling is None and hasattr(hf_config, "llm_config"):

577 pooling = getattr(hf_config.llm_config, "pooling", "avg")

578

579 + assert isinstance(pooling, str)

580 +

581 pooling_type = pooling_type_map.get(pooling)

582 if pooling_type is None:

583 raise ValueError(f"pool_type {pooling!r} not supported")

584

585 + assert model_config.pooler_config is not None

586 model_config.pooler_config.seq_pooling_type = pooling_type

587

588

@@ -631,7 +638,7 @@ class MambaModelConfig(VerifyAndUpdateConfig):

638

639

640 class NemotronHForCausalLMConfig(VerifyAndUpdateConfig):

634 - DEFAULT_MAMBA_SSM_CACHE_DTYPE = "float32"

641 + DEFAULT_MAMBA_SSM_CACHE_DTYPE: "MambaDType" = "float32"

642 """Only `float32` is known to have no accuracy issues by default."""

643

644 @classmethod

@@ -643,7 +650,7 @@ class NemotronHForCausalLMConfig(VerifyAndUpdateConfig):

650 `float32` if not specified.

651 """

652 if cache_config.mamba_ssm_cache_dtype == "auto":

646 - mamba_ssm_cache_dtype = getattr(

653 + mamba_ssm_cache_dtype: MambaDType = getattr(

654 hf_config, "mamba_ssm_cache_dtype", cls.DEFAULT_MAMBA_SSM_CACHE_DTYPE

655 )

656 logger.info(

@@ -729,6 +736,7 @@ class Qwen2ForProcessRewardModelConfig(VerifyAndUpdateConfig):

736 @staticmethod

737 def verify_and_update_model_config(model_config: "ModelConfig") -> None:

738 pooler_config = model_config.pooler_config

739 + assert pooler_config is not None

740

741 if pooler_config.step_tag_id is None:

742 pooler_config.step_tag_id = 151651

@@ -738,6 +746,7 @@ class Qwen2ForRewardModelConfig(VerifyAndUpdateConfig):

746 @staticmethod

747 def verify_and_update_model_config(model_config: "ModelConfig") -> None:

748 pooler_config = model_config.pooler_config

749 + assert pooler_config is not None

750

751 if pooler_config.use_activation is None:

752 pooler_config.use_activation = False

vllm/model_executor/models/cosmos3_edge.py

+10 -6

@@ -66,7 +66,7 @@ class Cosmos3EdgeVisionEncoder(Siglip2VisionTransformer):

66 def dtype(self) -> torch.dtype:

67 return self.embeddings.patch_embedding.weight.dtype

69 - def forward(

69 + def encode(

70 self,

71 pixel_values: torch.Tensor,

72 grid_thw: torch.Tensor,

@@ -253,7 +253,7 @@ class Cosmos3EdgeVisionModel(nn.Module):

253 grid_thw: torch.Tensor | list[list[int]],

254 ) -> torch.Tensor:

255 grid_thw = torch.as_tensor(grid_thw, dtype=torch.int64, device="cpu")

256 - image_embeds = self.encoder(pixel_values.type(self.dtype), grid_thw=grid_thw)

256 + image_embeds = self.encoder.encode(pixel_values.type(self.dtype), grid_thw)

257 image_embeds = patch_merging_by_param(

258 image_embeds,

259 grid_thw,

@@ -315,11 +315,10 @@ class Cosmos3EdgeAttention(NemotronHAttention):

315 rope_parameters=config.rope_parameters,

316 )

317

318 - def forward(

318 + def forward_with_positions(

319 self,

320 positions: torch.Tensor,

321 hidden_states: torch.Tensor,

322 - **kwargs,

322 ) -> torch.Tensor:

323 qkv, _ = self.qkv_proj(hidden_states)

324 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

@@ -367,7 +366,7 @@ class Cosmos3EdgeAttentionDecoderLayer(nn.Module):

366 hidden_states = self.norm(hidden_states)

367 else:

368 hidden_states, residual = self.norm(hidden_states, residual)

370 - hidden_states = self.mixer(positions=positions, hidden_states=hidden_states)

369 + hidden_states = self.mixer.forward_with_positions(positions, hidden_states)

370 return hidden_states, residual

371

372

@@ -562,6 +561,7 @@ class Cosmos3EdgeForConditionalGeneration(

561 config = vllm_config.model_config.hf_config

562 quant_config = vllm_config.quant_config

563 multimodal_config = vllm_config.model_config.multimodal_config

564 + assert multimodal_config is not None

565

566 self.config = config

567 self.multimodal_config = multimodal_config

@@ -619,6 +619,8 @@ class Cosmos3EdgeForConditionalGeneration(

619 image_grid_thw=image_grid_thw,

620 )

621

622 + raise AssertionError

623 +

624 def _parse_and_validate_video_input(

625 self, **kwargs: object

626 ) -> Qwen2_5_VLVideoInputs | None:

@@ -645,6 +647,8 @@ class Cosmos3EdgeForConditionalGeneration(

647 video_grid_thw=video_grid_thw,

648 )

649

650 + raise AssertionError

651 +

652 def _process_image_input(

653 self, image_input: Qwen2_5_VLImageInputs

654 ) -> tuple[torch.Tensor, ...]:

@@ -684,7 +688,7 @@ class Cosmos3EdgeForConditionalGeneration(

688 return self._get_image_features(pixel_values_videos, grid_thw)

689

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

687 - modalities = {}

691 + modalities: dict = {}

692 image_input = self._parse_and_validate_image_input(**kwargs)

693 if image_input is not None:

694 modalities["image"] = image_input

vllm/model_executor/models/dbrx.py

+10 -4

@@ -3,6 +3,7 @@

4 from collections.abc import Iterable

5 from itertools import islice

6 +from typing import Any

8 import torch

9 import torch.nn as nn

@@ -93,8 +94,11 @@ class DbrxExperts(RoutedExperts):

94 param: nn.Parameter,

95 loaded_weight: torch.Tensor,

96 weight_name: str,

96 - param_name: str,

97 - ):

97 + shard_id: str,

98 + expert_id: int | None = None,

99 + return_success: Any = False,

100 + ) -> Any:

101 + param_name = shard_id

102 tp_rank = get_tensor_model_parallel_rank()

103 param_data = param.data

104 shard_size = self.intermediate_size

@@ -134,6 +138,7 @@ class DbrxExperts(RoutedExperts):

138 param_data[:] = loaded_weight[:, :, shard]

139 else:

140 param_data[:] = loaded_weight

141 + return True if return_success else None

142

143

144 class DbrxMoE(nn.Module):

@@ -409,9 +414,10 @@ class DbrxModel(nn.Module):

414 if is_pp_missing_parameter(name, self):

415 continue

416 # Remapping the name of FP8 kv-scale.

412 - name = maybe_remap_kv_scale_name(name, params_dict)

413 - if name is None:

417 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

418 + if remapped_name is None:

419 continue

420 + name = remapped_name

421 param = params_dict[name]

422 weight_loader = getattr(param, "weight_loader", default_weight_loader)

423 weight_loader(param, loaded_weight)

vllm/model_executor/models/deepencoder.py

+4 -3

@@ -10,7 +10,7 @@

10 # This source code is licensed under the license found in the

11 # LICENSE file in the root directory of this source tree.

12 import math

13 -from collections.abc import Iterable

13 +from collections.abc import Callable, Iterable

14 from functools import partial

16 import torch

@@ -73,7 +73,7 @@ class ImageEncoderViT(nn.Module):

73 mlp_ratio: float = 4.0,

74 out_chans: int = 256,

75 qkv_bias: bool = True,

76 - norm_layer: type[nn.Module] = nn.LayerNorm,

76 + norm_layer: Callable[[int], nn.Module] = nn.LayerNorm,

77 act_layer: type[nn.Module] = nn.GELU,

78 use_abs_pos: bool = True,

79 use_rel_pos: bool = False,

@@ -205,7 +205,7 @@ class Block(nn.Module):

205 num_heads: int,

206 mlp_ratio: float = 4.0,

207 qkv_bias: bool = True,

208 - norm_layer: type[nn.Module] = nn.LayerNorm,

208 + norm_layer: Callable[[int], nn.Module] = nn.LayerNorm,

209 act_layer: type[nn.Module] = nn.GELU,

210 use_rel_pos: bool = False,

211 rel_pos_zero_init: bool = True,

@@ -326,6 +326,7 @@ class RelPosAttention(PluggableLayer):

326 v = v.view(B, self.num_heads, H * W, -1)

327

328 if self.use_rel_pos:

329 + assert rel_h is not None and rel_w is not None

330 rel_h = rel_h.view(

331 B, self.num_heads, rel_h.size(1), rel_h.size(2), rel_h.size(3)

332 )

vllm/model_executor/models/deepencoder2.py

+3 -3

@@ -15,6 +15,7 @@ from functools import lru_cache

15 import torch

16 import torch.nn as nn

17 import transformers

18 +from transformers.models.qwen2.modeling_qwen2 import Qwen2Model

20 from vllm.model_executor.custom_op import PluggableLayer

@@ -49,7 +50,6 @@ class CustomQwen2Decoder(PluggableLayer):

50 super().__init__()

52 # load

52 - Qwen2Model = transformers.models.qwen2.modeling_qwen2.Qwen2Model

53 Qwen2Config = transformers.Qwen2Config

55 # config

@@ -70,11 +70,11 @@ class CustomQwen2Decoder(PluggableLayer):

70 )

72 #

73 - self.model = self._create_custom_model(Qwen2Model, config)

73 + self.model = self._create_custom_model(config)

75 del self.model.embed_tokens

77 - def _create_custom_model(self, Qwen2Model, config):

77 + def _create_custom_model(self, config):

78 """Qwen2Model"""

80 class CustomQwen2ModelInner(Qwen2Model):

vllm/model_executor/models/deepseek_eagle.py

+8 -5

@@ -39,6 +39,7 @@ class DeepseekV2Model(nn.Module):

39 start_layer_id: int = 0,

40 ) -> None:

41 super().__init__()

42 + assert vllm_config.speculative_config is not None

43 self.config = vllm_config.speculative_config.draft_model_config.hf_config

44 quant_config = vllm_config.quant_config

45 self.vocab_size = self.config.vocab_size

@@ -155,7 +156,7 @@ class DeepseekV2Model(nn.Module):

156 break

157 else:

158 for mapping in expert_params_mapping:

158 - param_name, weight_name, expert_id, shard_id = mapping

159 + param_name, weight_name, expert_id, expert_shard_id = mapping

160 if weight_name not in name:

161 continue

162 name = name.replace(weight_name, param_name)

@@ -166,7 +167,7 @@ class DeepseekV2Model(nn.Module):

167 param,

168 loaded_weight,

169 name,

169 - shard_id=shard_id,

170 + shard_id=expert_shard_id,

171 expert_id=expert_id,

172 )

173 break

@@ -176,9 +177,10 @@ class DeepseekV2Model(nn.Module):

177 continue

178

179 # Remapping the name of FP8 kv-scale.

179 - name = maybe_remap_kv_scale_name(name, params_dict)

180 - if name is None:

180 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

181 + if remapped_name is None:

182 continue

183 + name = remapped_name

184

185 param = params_dict[name]

186 weight_loader = getattr(

@@ -192,6 +194,7 @@ class DeepseekV2Model(nn.Module):

194 class EagleDeepseekV3ForCausalLM(DeepseekV3ForCausalLM):

195 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

196 nn.Module.__init__(self)

197 + assert vllm_config.speculative_config is not None

198 self.config = vllm_config.speculative_config.draft_model_config.hf_config

199 quant_config = vllm_config.quant_config

200 target_layer_num = vllm_config.model_config.get_num_layers(

@@ -222,7 +225,7 @@ class EagleDeepseekV3ForCausalLM(DeepseekV3ForCausalLM):

225 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

226 return self.model.embed_input_ids(input_ids)

227

225 - def forward(

228 + def forward( # type: ignore[override]

229 self,

230 input_ids: torch.Tensor,

231 positions: torch.Tensor,

vllm/model_executor/models/deepseek_eagle3.py

+3 -1

@@ -159,6 +159,7 @@ class DeepseekV2Eagle3Model(nn.Module):

159 prefix: str = "",

160 ) -> None:

161 super().__init__()

162 + assert vllm_config.speculative_config is not None

163 self.config = vllm_config.speculative_config.draft_model_config.hf_config

164 self.vocab_size = self.config.vocab_size

165

@@ -276,6 +277,7 @@ class Eagle3DeepseekV2ForCausalLM(LocalArgmaxMixin, DeepseekV2ForCausalLM):

277

278 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

279 nn.Module.__init__(self)

280 + assert vllm_config.speculative_config is not None

281 self.config = vllm_config.speculative_config.draft_model_config.hf_config

282

283 # Ensure draft_vocab_size is set

@@ -318,7 +320,7 @@ class Eagle3DeepseekV2ForCausalLM(LocalArgmaxMixin, DeepseekV2ForCausalLM):

320 ) -> torch.Tensor:

321 return self.model.embed_input_ids(input_ids)

322

321 - def forward(

323 + def forward( # type: ignore[override]

324 self,

325 input_ids: torch.Tensor,

326 positions: torch.Tensor,

vllm/model_executor/models/deepseek_mtp.py

+6 -4

@@ -68,6 +68,7 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):

68 def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:

69 super().__init__()

71 + assert vllm_config.speculative_config is not None

72 config = vllm_config.speculative_config.draft_model_config.hf_config

73 self.config = config

74 quant_config = vllm_config.quant_config

@@ -423,7 +424,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):

424 # with expert_id.

425 is_expert_weight = False

426 for mapping in expert_params_mapping:

426 - param_name, weight_name, expert_id, shard_id = mapping

427 + param_name, weight_name, expert_id, expert_shard_id = mapping

428 if weight_name not in chunk_name:

429 continue

430

@@ -446,7 +447,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):

447 param,

448 weight_to_load,

449 name_mapped,

449 - shard_id=shard_id,

450 + shard_id=expert_shard_id,

451 expert_id=expert_id,

452 return_success=True,

453 )

@@ -467,9 +468,10 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):

468 if name.endswith(".bias") and name not in params_dict:

469 continue

470

470 - name = maybe_remap_kv_scale_name(name, params_dict)

471 - if name is None:

471 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

472 + if remapped_name is None:

473 continue

474 + name = remapped_name

475

476 # According to DeepSeek-V3 Technical Report, MTP modules

477 # shares embedding layer. We only load the first weights.

vllm/model_executor/models/deepseek_ocr.py

+11 -2

@@ -49,6 +49,7 @@ from vllm.multimodal.processing import (

49 from vllm.sampling_params import SamplingParams

50 from vllm.sequence import IntermediateTensors

51 from vllm.tokenizers import cached_tokenizer_from_config

52 +from vllm.tokenizers.hf import HfTokenizer

53 from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config

54 from vllm.transformers_utils.processors.deepseek_ocr import (

55 BASE_SIZE,

@@ -184,6 +185,7 @@ class NGramPerReqLogitsProcessor(AdapterLogitsProcessor):

185 )

186 if ngram_size is None:

187 return None

188 + assert isinstance(ngram_size, int) and isinstance(window_size, int)

189

190 whitelist_token_ids = set(whitelist_token_ids) if whitelist_token_ids else None

191 return NoRepeatNGramLogitsProcessor(

@@ -301,6 +303,7 @@ class DeepseekOCRMultiModalProcessor(

303

304 else:

305 tokenizer = self.info.get_tokenizer()

306 + assert isinstance(tokenizer, HfTokenizer)

307 processed_outputs = tokenizer(

308 prompt, add_special_tokens=True, return_tensors="pt"

309 )

@@ -342,6 +345,7 @@ class DeepseekOCRMultiModalProcessor(

345 if isinstance(images, ImageEmbeddingItems):

346 num_image_tokens = images.get_feature_size(item_idx)

347 else:

348 + assert isinstance(images, ImageProcessorItems)

349 size = images.get_image_size(item_idx)

350

351 num_image_tokens = self.info.get_num_image_tokens(

@@ -459,7 +463,11 @@ class DeepseekOCRForCausalLM(

463 images_spatial_crop = kwargs.pop("images_spatial_crop", None)

464 images_crop = kwargs.pop("images_crop", None)

465

462 - if pixel_values is None or torch.sum(pixel_values).item() == 0:

466 + if pixel_values is None:

467 + return None

468 + assert isinstance(pixel_values, torch.Tensor)

469 + assert images_crop is None or isinstance(images_crop, torch.Tensor)

470 + if torch.sum(pixel_values).item() == 0:

471 return None

472

473 # Use actual tensor spatial dim instead of hardcoded

@@ -675,7 +683,8 @@ class DeepseekOCRForCausalLM(

683 num_input_tokens = global_input_side**2

684

685 if is_tiled:

678 - num_patches = image_spatial_crop.prod(dim=-1)

686 + assert image_spatial_crop is not None

687 + num_patches = int(image_spatial_crop.prod().item())

688 num_input_tokens += num_patches * (local_input_side**2)

689

690 global_output_token = self.global_image_output_token

vllm/model_executor/models/deepseek_ocr2.py

+4 -1

@@ -47,6 +47,7 @@ from vllm.multimodal.processing import (

47 )

48 from vllm.sequence import IntermediateTensors

49 from vllm.tokenizers import cached_tokenizer_from_config

50 +from vllm.tokenizers.hf import HfTokenizer

51 from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config

52 from vllm.transformers_utils.processors.deepseek_ocr import (

53 BASE_SIZE,

@@ -173,6 +174,7 @@ class DeepseekOCR2MultiModalProcessor(

174

175 else:

176 tokenizer = self.info.get_tokenizer()

177 + assert isinstance(tokenizer, HfTokenizer)

178 processed_outputs = tokenizer(

179 prompt, add_special_tokens=True, return_tensors="pt"

180 )

@@ -214,6 +216,7 @@ class DeepseekOCR2MultiModalProcessor(

216 if isinstance(images, ImageEmbeddingItems):

217 num_image_tokens = images.get_feature_size(item_idx)

218 else:

219 + assert isinstance(images, ImageProcessorItems)

220 size = images.get_image_size(item_idx)

221

222 num_image_tokens = self.info.get_num_image_tokens(

@@ -284,7 +287,7 @@ class DeepseekOCR2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, Support

287 patch_size=16,

288 qkv_bias=True,

289 use_rel_pos=True,

287 - global_attn_indexes=[2, 5, 8, 11],

290 + global_attn_indexes=(2, 5, 8, 11),

291 window_size=14,

292 out_chans=256,

293 last_conv_output=896,

vllm/model_executor/models/deepseek_v2.py

+13 -6

@@ -626,7 +626,7 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):

626

627 def forward(self): ...

628

629 - def get_attn_backend(self) -> AttentionBackend:

629 + def get_attn_backend(self) -> type[AttentionBackend]:

630 return DeepseekV32IndexerBackend

631

632

@@ -681,6 +681,7 @@ class Indexer(nn.Module):

681 # NOTE: (zyongye) we use fp8 naive cache,

682 # where we store value in fp8 and scale in fp32

683 # per self.quant_block_size element

684 + assert cache_config is not None

685 self.k_cache = DeepseekV32IndexerCache(

686 head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,

687 dtype=torch.uint8,

@@ -1107,6 +1108,10 @@ class DeepseekV2MLAAttention(nn.Module):

1108 )

1109

1110 if self.is_v32 and (not _skip_topk or is_mtp_layer):

1111 + assert q_lora_rank is not None

1112 + assert cache_config is not None

1113 + self.indexer_rope_emb: nn.Module | None

1114 + self.indexer: Indexer | None

1115 self.indexer_rope_emb = get_rope(

1116 qk_rope_head_dim,

1117 max_position=max_position_embeddings,

@@ -1218,6 +1223,7 @@ class DeepseekV2DecoderLayer(nn.Module):

1223

1224 self.use_mha = use_mha

1225

1226 + attn_cls: typing.Any

1227 if use_mha:

1228 attn_cls = DeepseekAttention

1229 elif model_config.use_mla:

@@ -1508,7 +1514,7 @@ class DeepseekV2Model(nn.Module):

1514 rocm_aiter_moe_shared_expert_enabled = (

1515 rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()

1516 )

1511 - stacked_params_mapping = [

1517 + stacked_params_mapping: list[tuple[str, str, int | str]] = [

1518 # (param_name, shard_name, shard_id)

1519 ("gate_up_proj", "gate_proj", 0),

1520 ("gate_up_proj", "up_proj", 1),

@@ -1677,7 +1683,7 @@ class DeepseekV2Model(nn.Module):

1683 # param and delegate to its expert-aware weight_loader

1684 # with expert_id.

1685 for mapping in expert_params_mapping:

1680 - param_name, weight_name, expert_id, shard_id = mapping

1686 + param_name, weight_name, expert_id, expert_shard_id = mapping

1687 if weight_name not in chunk_name:

1688 continue

1689

@@ -1703,7 +1709,7 @@ class DeepseekV2Model(nn.Module):

1709 param,

1710 weight_to_load,

1711 name_mapped,

1706 - shard_id=shard_id,

1712 + shard_id=expert_shard_id,

1713 expert_id=expert_id,

1714 return_success=True,

1715 )

@@ -1725,9 +1731,10 @@ class DeepseekV2Model(nn.Module):

1731 continue

1732

1733 # Remapping the name of FP8 kv-scale.

1728 - name = maybe_remap_kv_scale_name(name, params_dict)

1729 - if name is None:

1734 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

1735 + if remapped_name is None:

1736 continue

1737 + name = remapped_name

1738

1739 if is_pp_missing_parameter(name, self):

1740 continue

vllm/model_executor/models/deepseek_vl2.py

+6 -1

@@ -43,6 +43,7 @@ from vllm.multimodal.processing.processor import (

43 )

44 from vllm.sequence import IntermediateTensors

45 from vllm.tokenizers import cached_tokenizer_from_config

46 +from vllm.tokenizers.hf import HfTokenizer

47 from vllm.transformers_utils.configs.deepseek_vl2 import (

48 DeepseekVLV2Config,

49 MlpProjectorConfig,

@@ -244,6 +245,7 @@ class DeepseekVL2MultiModalProcessor(

245 ) -> BatchFeature:

246 if not mm_data:

247 tokenizer = self.info.get_tokenizer()

248 + assert isinstance(tokenizer, HfTokenizer)

249 return tokenizer(prompt, add_special_tokens=True, return_tensors="pt")

250

251 processed_outputs = super()._call_hf_processor(

@@ -291,6 +293,7 @@ class DeepseekVL2MultiModalProcessor(

293 if isinstance(images, ImageEmbeddingItems):

294 num_image_tokens = images.get_feature_size(item_idx)

295 else:

296 + assert isinstance(images, ImageProcessorItems)

297 image_size = images.get_image_size(item_idx)

298

299 num_image_tokens = self.info.get_num_image_tokens(

@@ -406,7 +409,9 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

409 return parent, names[-1]

410

411 # patch for timm ViT instance to support tensor parallel

409 - def patch_vit_for_tp(self, vit: torch.nn.Module, quant_config: QuantizationConfig):

412 + def patch_vit_for_tp(

413 + self, vit: torch.nn.Module, quant_config: QuantizationConfig | None

414 + ):

415 try:

416 import timm

417 except ImportError as e:

vllm/model_executor/models/diffusion_gemma.py

+4 -1

@@ -32,6 +32,7 @@ from vllm.distributed.parallel_state import get_tp_group

32 from vllm.logger import init_logger

33 from vllm.model_executor.layers.layernorm import RMSNorm

34 from vllm.model_executor.layers.logits_processor import LogitsProcessor

35 +from vllm.model_executor.layers.quantization import QuantizationConfig

36 from vllm.model_executor.layers.vocab_parallel_embedding import (

37 ParallelLMHead,

38 )

@@ -195,8 +196,10 @@ class DiffusionGemmaForConditionalGeneration(

196

197 # ---- Vision tower ----

198 vision_config = getattr(config, "vision_config", None)

199 + self.embed_vision: Gemma4MultimodalEmbedder | None

200 if vision_config is not None:

201 quant_config = vllm_config.quant_config

202 + tower_quant: QuantizationConfig | None

203 if quant_config and quant_config.get_name() in [

204 "bitsandbytes",

205 "torchao",

@@ -1048,7 +1051,7 @@ class DiffusionSampler:

1051 sampler: Any,

1052 diffusion_config: Any,

1053 vocab_size: int,

1051 - diffusion_states: DiffusionGemmaRequestStates | None = None,

1054 + diffusion_states: DiffusionGemmaRequestStates,

1055 *,

1056 confidence_threshold: float,

1057 t_min: float,

vllm/model_executor/models/dots_ocr.py

+14 -7

@@ -149,7 +149,7 @@ class DotsOCRProcessingInfo(Qwen2VLProcessingInfo):

149 self,

150 **kwargs: object,

151 ) -> Qwen2VLProcessor:

152 - self.get_tokenizer().image_token = IMAGE_TOKEN # Ensure image token is set

152 + setattr(self.get_tokenizer(), "image_token", IMAGE_TOKEN) # noqa: B010

153 processor = self.ctx.get_hf_processor(

154 Qwen2VLProcessor,

155 **kwargs,

@@ -430,7 +430,7 @@ class DotsVisionBlock(nn.Module):

430 *,

431 cu_seqlens: torch.Tensor,

432 rotary_pos_emb: torch.Tensor,

433 - max_seqlen: int | None = None,

433 + max_seqlen: torch.Tensor | None = None,

434 ) -> torch.Tensor:

435 hidden_states = hidden_states + self.attn(

436 self.norm1(hidden_states),

@@ -483,6 +483,7 @@ class DotsVisionTransformer(nn.Module):

483 )

484 if require_post_norm is None:

485 require_post_norm = len(self.blocks) == config.num_hidden_layers

486 + self.post_trunk_norm: RMSNorm | None

487 if require_post_norm and self.config.post_norm:

488 self.post_trunk_norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)

489 else:

@@ -536,7 +537,7 @@ class DotsVisionTransformer(nn.Module):

537 rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)

538 return rotary_pos_emb

539

539 - def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> int | None:

540 + def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> torch.Tensor | None:

541 max_seqlen = None

542 if self.attn_backend in {

543 AttentionBackendEnum.FLASH_ATTN,

@@ -552,15 +553,17 @@ class DotsVisionTransformer(nn.Module):

553 rotary_pos_emb = self.rot_pos_emb(grid_thw)

554

555 # Convert grid_thw to tensor (always expecting list format now)

555 - grid_thw = torch.tensor(grid_thw, device=hidden_states.device, dtype=torch.long)

556 + grid_thw_tensor = torch.tensor(

557 + grid_thw, device=hidden_states.device, dtype=torch.long

558 + )

559 hidden_states = hidden_states.to(self.dtype)

557 - hidden_states = self.patch_embed(hidden_states, grid_thw)

560 + hidden_states = self.patch_embed(hidden_states, grid_thw_tensor)

561

562 cu_seqlens = torch.repeat_interleave(

560 - grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]

563 + grid_thw_tensor[:, 1] * grid_thw_tensor[:, 2], grid_thw_tensor[:, 0]

564 ).cumsum(

565 dim=0,

563 - dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,

566 + dtype=grid_thw_tensor.dtype if torch.jit.is_tracing() else torch.int32,

567 )

568 cu_seqlens = torch.cat([cu_seqlens.new_zeros(1), cu_seqlens])

569

@@ -620,6 +623,7 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA

623 def get_placeholder_str(cls, modality: str, i: int) -> str | None:

624 if modality.startswith("image"):

625 return "<|img|><|imgpad|><|endofimg|>"

626 + return None

627

628 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

629 super().__init__()

@@ -627,6 +631,7 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA

631 self.config: DotsOCRConfig = vllm_config.model_config.hf_config

632 self.quant_config = vllm_config.quant_config

633 multimodal_config = vllm_config.model_config.multimodal_config

634 + assert multimodal_config is not None

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

636 if isinstance(self.config.vision_config, dict):

637 vision_config = DotsVisionConfig(**self.config.vision_config)

@@ -677,6 +682,8 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA

682 image_grid_thw=image_grid_thw,

683 )

684

685 + raise AssertionError

686 +

687 def _process_image_input(

688 self, image_input: DotsOCRImageInputs

689 ) -> tuple[torch.Tensor, ...]:

vllm/multimodal/processing/dummy_inputs.py

+2 -1

@@ -118,8 +118,9 @@ class BaseDummyInputsBuilder(ABC, Generic[_I]):

118 width: int,

119 height: int,

120 num_images: int,

121 - overrides: ImageDummyOptions | None = None,

121 + overrides: BaseDummyOptions | None = None,

122 ) -> list[Image.Image]:

123 + assert overrides is None or isinstance(overrides, ImageDummyOptions)

124 if num_images == 0:

125 return []

126 if overrides: