@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 21fa2c5a2a

[Mypy] Fix mypy typing for model interfaces and H/I models (#54079)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Taneem Ibrahim committed Aug 28, 2026 at 15:49 UTC 21fa2c5a2a471b9c978a5a3b7208a4c7fe88cefa

23 files changed +222 -137

tools/pre_commit/mypy.py

-2

@@ -112,8 +112,6 @@ SEPARATE_GROUPS = [

112 ]

113

114 EXCLUDE = [

115 - r"vllm/model_executor/models/[hH]",

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

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

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

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

vllm/model_executor/models/chatglm.py

+1 -1

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

430

431

432 class ChatGLMForCausalLM(ChatGLMBaseModel, SupportsLoRA, SupportsPP, SupportsQuant):

433 - packed_modules_mapping = {

433 + packed_modules_mapping: dict[str, list[str]] = {

434 "query_key_value": ["query_key_value"],

435 "dense_h_to_4h": ["dense_h_to_4h"],

436 }

vllm/model_executor/models/commandr.py

+1 -1

@@ -417,7 +417,7 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):

417 orig_to_new_substr={"_quantizer.": None},

418 orig_to_new_prefix={"lm_head": None},

419 )

420 - packed_modules_mapping = {

420 + packed_modules_mapping: dict[str, list[str]] = {

421 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

422 "gate_up_proj": ["gate_proj", "up_proj"],

423 }

vllm/model_executor/models/gemma.py

+1 -1

@@ -339,7 +339,7 @@ class GemmaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):

339 ".up_proj": (".gate_up_proj", 1),

340 }

341 )

342 - packed_modules_mapping = {

342 + packed_modules_mapping: dict[str, list[str]] = {

343 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

344 "gate_up_proj": ["gate_proj", "up_proj"],

345 }

vllm/model_executor/models/gemma4_mm.py

+1 -1

@@ -1019,7 +1019,7 @@ class Gemma4ForConditionalGeneration(

1019 mm_prefix_clamp_sliding_window: bool = True

1020 supports_tower_connector_lora = True

1021

1022 - packed_modules_mapping = {

1022 + packed_modules_mapping: dict[str, list[str]] = {

1023 "qkv_proj": [

1024 "q_proj",

1025 "k_proj",

vllm/model_executor/models/granite.py

+1 -1

@@ -395,7 +395,7 @@ class GraniteModel(nn.Module):

395 class GraniteForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):

396 hf_to_vllm_mapper = GraniteModel.hf_to_vllm_mapper

397 # LoRA specific attributes

398 - packed_modules_mapping = {

398 + packed_modules_mapping: dict[str, list[str]] = {

399 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

400 "gate_up_proj": ["gate_proj", "up_proj"],

401 }

vllm/model_executor/models/granitemoehybrid.py

+1 -1

@@ -587,7 +587,7 @@ class GraniteMoeHybridForCausalLM(

587 SupportsQuant,

588 SupportsMambaPrefixCaching,

589 ):

590 - packed_modules_mapping = {

590 + packed_modules_mapping: dict[str, list[str]] = {

591 "qkv_proj": [

592 "q_proj",

593 "k_proj",

vllm/model_executor/models/h2ovl.py

@@ -114,6 +114,7 @@ class H2OVLMultiModalProcessor(BaseInternVLMultiModalProcessor[H2OVLProcessingIn

114 if isinstance(images, ImageEmbeddingItems):

115 feature_size = images.get_feature_size(item_idx)

116 else:

117 + assert isinstance(images, ImageProcessorItems)

118 image_size = images.get_image_size(item_idx)

119 feature_size = self.info.get_num_image_tokens(

120 image_width=image_size.width,

vllm/model_executor/models/hrm_text.py

+2 -2

@@ -248,7 +248,7 @@ class HrmTextDecoderLayer(nn.Module):

248 self,

249 config: PretrainedConfig,

250 layer_idx_in_stack: int,

251 - stack_kind: str,

251 + stack_kind: Literal["L", "H"],

252 cache_config: CacheConfig | None = None,

253 quant_config: QuantizationConfig | None = None,

254 prefix: str = "",

@@ -312,7 +312,7 @@ class HrmTextStack(nn.Module):

312 def __init__(

313 self,

314 config: PretrainedConfig,

315 - stack_kind: str,

315 + stack_kind: Literal["L", "H"],

316 cache_config: CacheConfig | None = None,

317 quant_config: QuantizationConfig | None = None,

318 prefix: str = "",

vllm/model_executor/models/hy_v3.py

+12 -11

@@ -161,6 +161,7 @@ class HYV3MoEFused(nn.Module):

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

162 )

163

164 + self.shared_mlp: HYV3FeedForward | None

165 if config.num_shared_experts > 0:

166 self.shared_mlp = HYV3FeedForward(

167 hidden_size=config.hidden_size,

@@ -467,28 +468,27 @@ class HYV3Model(nn.Module, MixtureOfExperts):

468 # Set MoE hyperparameters

469 self.num_expert_groups = 1

470 self.moe_layers = []

470 - example_layer = None

471 + example_layer: HYV3MoEFused | None = None

472 for layer in self.layers:

473 if isinstance(layer, PPMissingLayer):

474 continue

475

476 assert isinstance(layer, HYV3DecoderLayer)

477 if layer.block_type == "moe":

478 + assert isinstance(layer.mlp, HYV3MoEFused)

479 example_layer = layer.mlp

478 - self.moe_layers.append(layer.mlp.experts)

480 + self.moe_layers.append(example_layer.experts)

481

482 if example_layer is None:

483 self.num_moe_layers = 0

484 raise RuntimeError("No MoE layer found in model.layers.")

485

486 self.num_moe_layers = len(self.moe_layers)

485 - self.num_logical_experts = getattr(example_layer, "n_logical_experts", None)

486 - self.num_physical_experts = getattr(example_layer, "n_physical_experts", None)

487 - self.num_local_physical_experts = getattr(

488 - example_layer, "n_local_physical_experts", None

489 - )

490 - self.num_routed_experts = getattr(example_layer, "n_routed_experts", None)

491 - self.num_redundant_experts = getattr(example_layer, "n_redundant_experts", None)

487 + self.num_logical_experts = example_layer.n_logical_experts

488 + self.num_physical_experts = example_layer.n_physical_experts

489 + self.num_local_physical_experts = example_layer.n_local_physical_experts

490 + self.num_routed_experts = example_layer.n_routed_experts

491 + self.num_redundant_experts = example_layer.n_redundant_experts

492

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

494 return self.embed_tokens(input_ids)

@@ -572,9 +572,10 @@ class HYV3Model(nn.Module, MixtureOfExperts):

572 continue

573 if "scale" in name:

574 # Remapping the name of FP8 kv-scale.

575 - name = maybe_remap_kv_scale_name(name, params_dict)

576 - if name is None:

575 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

576 + if remapped_name is None:

577 continue

578 + name = remapped_name

579 is_found = False

580 for param_name, weight_name, shard_id in stacked_params_mapping:

581 if weight_name not in name:

vllm/model_executor/models/hy_v3_mtp.py

+14 -5

@@ -24,7 +24,7 @@

24 # limitations under the License.

25 """Inference-only HY V3 MTP model compatible with HuggingFace weights."""

27 -from collections.abc import Iterable

27 +from collections.abc import Callable, Iterable

29 import regex as re

30 import torch

@@ -288,7 +288,15 @@ class HYV3MTP(nn.Module):

288 num_kv_heads = getattr(

289 self.config, "num_key_value_heads", self.config.num_attention_heads

290 )

291 - split_params_mapping = [

291 + split_params_mapping: list[

292 + tuple[

293 + str,

294 + str,

295 + int,

296 + list[tuple[str | int, int]],

297 + Callable[[torch.Tensor], torch.Tensor] | None,

298 + ]

299 + ] = [

300 (".gate_up_proj", ".gate_and_up_proj", 2, [(1, 1), (0, 1)], None),

301 (

302 ".qkv_proj",

@@ -308,7 +316,7 @@ class HYV3MTP(nn.Module):

316 num_experts=self.config.num_experts,

317 )

318 else:

311 - expert_params_mapping = {}

319 + expert_params_mapping = []

320

321 params_dict = dict(self.named_parameters())

322

@@ -354,9 +362,10 @@ class HYV3MTP(nn.Module):

362 if name == "__skip__":

363 continue

364 if "scale" in name:

357 - name = maybe_remap_kv_scale_name(name, params_dict)

358 - if name is None:

365 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

366 + if remapped_name is None:

367 continue

368 + name = remapped_name

369 is_found = False

370

371 for param_name, weight_name, shard_id in stacked_params_mapping:

vllm/model_executor/models/hyperclovax_vision_v2.py

+12 -3

@@ -12,7 +12,7 @@ Supports:

13 from collections.abc import Iterable, Mapping, Sequence

14 from functools import partial

15 -from typing import Annotated, Literal

15 +from typing import Annotated, Literal, TypedDict

17 import torch

18 import torch.nn as nn

@@ -122,6 +122,11 @@ class HCXVisionV2VideoEmbeddingInputs(TensorSchema):

122 HCXVisionV2VideoInputs = HCXVisionV2VideoPixelInputs | HCXVisionV2VideoEmbeddingInputs

123

124

125 +class HCXVisionV2MultiModalInputs(TypedDict, total=False):

126 + image: HCXVisionV2ImageInputs | None

127 + video: HCXVisionV2VideoInputs | None

128 +

129 +

130 class HCXVisionV2ProcessingInfo(BaseProcessingInfo):

131 """Processing info for HyperCLOVAX V2 (32B Think model)."""

132

@@ -329,6 +334,7 @@ class HCXVisionV2MultiModalProcessor(

334 if grid_thw_elem is not None:

335 # Access .data to get the actual tensor from MultiModalFieldElem

336 grid_thw = grid_thw_elem.data

337 + assert isinstance(grid_thw, torch.Tensor)

338 # Qwen2.5-VL style calculation

339 h, w = grid_thw[1].item(), grid_thw[2].item()

340 num_tokens = (h * w) // (merge_size**2)

@@ -340,6 +346,7 @@ class HCXVisionV2MultiModalProcessor(

346 if grid_thw_elem is not None:

347 # Access .data to get the actual tensor from MultiModalFieldElem

348 grid_thw = grid_thw_elem.data

349 + assert isinstance(grid_thw, torch.Tensor)

350 t, h, w = grid_thw[0].item(), grid_thw[1].item(), grid_thw[2].item()

351 num_tokens = (t * h * w) // (merge_size**2)

352 else:

@@ -606,8 +613,10 @@ class HCXVisionV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

613 sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

614 return video_embeds.split(sizes)

615

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

610 - modalities = {}

616 + def _parse_and_validate_multimodal_inputs(

617 + self, **kwargs: object

618 + ) -> HCXVisionV2MultiModalInputs:

619 + modalities: HCXVisionV2MultiModalInputs = {}

620

621 for input_key in kwargs:

622 if (

vllm/model_executor/models/idefics3.py

+2 -1

@@ -625,6 +625,7 @@ class Idefics3ForConditionalGeneration(

625 if image_input["type"] == "image_embeds":

626 return image_input["data"]

627

628 + assert isinstance(image_input, Idefics3ImagePixelInputs)

629 image_features = self._process_image_pixels(image_input)

630 image_features = self.model.connector(image_features)

631

@@ -788,7 +789,7 @@ class Idefics3ForConditionalGeneration(

789 path: str = "default",

790 ) -> torch.Tensor:

791 image_input = self._parse_and_validate_image_input(**mm_kwargs)

791 - assert image_input is not None and image_input["type"] == "pixel_values"

792 + assert isinstance(image_input, Idefics3ImagePixelInputs)

793 image_features = self._process_image_pixels(image_input)

794 image_features = self.model.connector(image_features)

795 return image_features.flatten(0, 1)

vllm/model_executor/models/interfaces.py

+70 -44

@@ -32,6 +32,7 @@ from typing_extensions import Self, TypeIs

33 from vllm.logger import init_logger

34 from vllm.model_executor.layers.quantization import QuantizationConfig

35 +from vllm.model_executor.models.interfaces_base import VllmModel

36 from vllm.utils.collection_utils import common_prefix

37 from vllm.utils.func_utils import supports_kw

@@ -46,8 +47,9 @@ if TYPE_CHECKING:

47 from vllm.inputs import PromptType, TokensPrompt

48 from vllm.lora.model_manager import LoRAModelManager

49 from vllm.model_executor.layers.fused_moe import MoERunner

50 + from vllm.model_executor.layers.logits_processor import LogitsProcessor

51 from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc

50 - from vllm.model_executor.models.interfaces_base import VllmModel

52 + from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

53 from vllm.model_executor.models.utils import WeightsMapper

54 from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItem

55 from vllm.multimodal.registry import _ProcessorFactories

@@ -128,14 +130,24 @@ class SupportsMultiModalEmbeddings(Protocol):

130

131 supports_multimodal_embeddings: ClassVar[Literal[True]] = True

132

133 + @overload

134 def embed_input_ids(

135 self,

136 input_ids: Tensor,

134 - multimodal_embeddings: MultiModalEmbeddings | None = None,

137 + multimodal_embeddings: None = None,

138 *,

139 is_multimodal: Tensor | None = None,

140 ) -> Tensor: ...

141

142 + @overload

143 + def embed_input_ids(

144 + self,

145 + input_ids: Tensor,

146 + multimodal_embeddings: MultiModalEmbeddings,

147 + *,

148 + is_multimodal: Tensor,

149 + ) -> Tensor: ...

150 +

151

152 @runtime_checkable

153 class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):

@@ -243,26 +255,30 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):

255 torch.nn.Module: The core language model component.

256 """

257 # Cached

246 - if self in _language_model_by_module:

247 - return _language_model_by_module[self]

258 + assert isinstance(self, nn.Module)

259 + self_module = self

260 + if self_module in _language_model_by_module:

261 + return _language_model_by_module[self_module]

262

263 if self._language_model_names:

250 - mod = self

264 + mod = self_module

265 for attr in common_prefix(

266 [name.split(".") for name in self._language_model_names]

267 ):

268 if attr:

269 mod = getattr(mod, attr)

270

257 - if mod is not self and hasattr(mod, "embed_input_ids"):

258 - _language_model_by_module[self] = mod

259 - return mod

271 + candidate: object = mod

272 + if candidate is not self_module and isinstance(candidate, VllmModel):

273 + _language_model_by_module[self_module] = candidate

274 + return candidate

275

276 # Fallback

262 - for mod in self.children():

263 - if hasattr(mod, "embed_input_ids"):

264 - _language_model_by_module[self] = mod

265 - return mod

277 + for mod in self_module.children():

278 + candidate = mod

279 + if isinstance(candidate, VllmModel):

280 + _language_model_by_module[self_module] = candidate

281 + return candidate

282

283 raise NotImplementedError(

284 f"No language model found in {type(self).__name__}! "

@@ -289,7 +305,7 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):

305 """

306 from .utils import StageMissingLayer, collect_children, no_init_weights

307

292 - mm_config = vllm_config.model_config.multimodal_config

308 + mm_config = vllm_config.model_config.get_multimodal_config()

309

310 with collect_children(self, targets=targets) as children_names: # noqa: SIM117

311 with (

@@ -339,7 +355,7 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):

355 else:

356 stage_name = "_".join([*modalities, "tower"])

357

342 - mm_config = vllm_config.model_config.multimodal_config

358 + mm_config = vllm_config.model_config.get_multimodal_config()

359

360 with collect_children(self, targets=targets) as children_names: # noqa: SIM117

361 with (

@@ -361,7 +377,8 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):

377 offloader = get_offloader()

378 if offloader.supports_tower_offload:

379 for name in children_names:

364 - offloader.wrap_modules(iter([attrgetter(name)(self)]), prefix=name)

380 + modules = (module for module in [attrgetter(name)(self)])

381 + offloader.wrap_modules(modules, prefix=name)

382

383 @contextmanager

384 def _mark_composite_model(

@@ -434,18 +451,6 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):

451 num_connector_tokens if isinstance(num_connector_tokens, int) else None,

452 )

453

437 - @overload

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

439 -

440 - @overload

441 - def embed_input_ids(

442 - self,

443 - input_ids: Tensor,

444 - multimodal_embeddings: MultiModalEmbeddings,

445 - *,

446 - is_multimodal: torch.Tensor,

447 - ) -> Tensor: ...

448 -

454 def _embed_text_input_ids(

455 self,

456 input_ids: Tensor,

@@ -465,6 +470,24 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):

470

471 return embed_input_ids(input_ids)

472

473 + @overload

474 + def embed_input_ids(

475 + self,

476 + input_ids: Tensor,

477 + multimodal_embeddings: None = None,

478 + *,

479 + is_multimodal: Tensor | None = None,

480 + ) -> Tensor: ...

481 +

482 + @overload

483 + def embed_input_ids(

484 + self,

485 + input_ids: Tensor,

486 + multimodal_embeddings: MultiModalEmbeddings,

487 + *,

488 + is_multimodal: Tensor,

489 + ) -> Tensor: ...

490 +

491 def embed_input_ids(

492 self,

493 input_ids: Tensor,

@@ -671,7 +694,7 @@ class SupportsLoRA(Protocol):

694 # The `embedding_module` and `embedding_padding_modules`

695 # are empty by default.

696 embedding_modules: ClassVar[dict[str, str]] = {}

674 - packed_modules_mapping: ClassVar[dict[str, list[str]]] = {}

697 + packed_modules_mapping: dict[str, list[str]] = {}

698 # Module prefixes to skip during LoRA loading (e.g., ["mtp."] for MTP layers)

699 lora_skip_prefixes: ClassVar[list[str]] = []

700 lora_manager: "LoRAModelManager | None"

@@ -864,11 +887,11 @@ class HasInnerState(Protocol):

887

888

889 @overload

867 -def has_inner_state(model: object) -> TypeIs[HasInnerState]: ...

890 +def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]: ...

891

892

893 @overload

871 -def has_inner_state(model: type[object]) -> TypeIs[type[HasInnerState]]: ...

894 +def has_inner_state(model: object) -> TypeIs[HasInnerState]: ...

895

896

897 def has_inner_state(

@@ -891,11 +914,11 @@ class IsAttentionFree(Protocol):

914

915

916 @overload

894 -def is_attention_free(model: object) -> TypeIs[IsAttentionFree]: ...

917 +def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]: ...

918

919

920 @overload

898 -def is_attention_free(model: type[object]) -> TypeIs[type[IsAttentionFree]]: ...

921 +def is_attention_free(model: object) -> TypeIs[IsAttentionFree]: ...

922

923

924 def is_attention_free(

@@ -946,11 +969,11 @@ class IsHybrid(Protocol):

969

970

971 @overload

949 -def is_hybrid(model: object) -> TypeIs[IsHybrid]: ...

972 +def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]: ...

973

974

975 @overload

953 -def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]: ...

976 +def is_hybrid(model: object) -> TypeIs[IsHybrid]: ...

977

978

979 def is_hybrid(

@@ -1025,7 +1048,7 @@ class MixtureOfExperts(Protocol):

1048 self.expert_weights = []

1049 for layer_idx, layer in enumerate(self.moe_layers):

1050 # Register the expert weights.

1028 - self.expert_weights.append(layer.get_expert_weights())

1051 + self.expert_weights.append(list(layer.get_expert_weights()))

1052 layer.set_eplb_state(

1053 moe_layer_idx=layer_idx,

1054 expert_load_view=expert_load_view,

@@ -1083,11 +1106,11 @@ class HasNoOps(Protocol):

1106

1107

1108 @overload

1086 -def has_noops(model: object) -> TypeIs[HasNoOps]: ...

1109 +def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]: ...

1110

1111

1112 @overload

1090 -def has_noops(model: type[object]) -> TypeIs[type[HasNoOps]]: ...

1113 +def has_noops(model: object) -> TypeIs[HasNoOps]: ...

1114

1115

1116 def has_noops(

@@ -1108,14 +1131,14 @@ class SupportsMambaPrefixCaching(Protocol):

1131

1132 @overload

1133 def supports_mamba_prefix_caching(

1111 - model: object,

1112 -) -> TypeIs[SupportsMambaPrefixCaching]: ...

1134 + model: type[object],

1135 +) -> TypeIs[type[SupportsMambaPrefixCaching]]: ...

1136

1137

1138 @overload

1139 def supports_mamba_prefix_caching(

1117 - model: type[object],

1118 -) -> TypeIs[type[SupportsMambaPrefixCaching]]: ...

1140 + model: object,

1141 +) -> TypeIs[SupportsMambaPrefixCaching]: ...

1142

1143

1144 def supports_mamba_prefix_caching(

@@ -1136,11 +1159,11 @@ class SupportsReplaySSM(Protocol):

1159

1160

1161 @overload

1139 -def supports_replayssm(model: object) -> TypeIs[SupportsReplaySSM]: ...

1162 +def supports_replayssm(model: type[object]) -> TypeIs[type[SupportsReplaySSM]]: ...

1163

1164

1165 @overload

1143 -def supports_replayssm(model: type[object]) -> TypeIs[type[SupportsReplaySSM]]: ...

1166 +def supports_replayssm(model: object) -> TypeIs[SupportsReplaySSM]: ...

1167

1168

1169 def supports_replayssm(

@@ -1172,7 +1195,7 @@ class SupportsQuant:

1195 """The interface required for all models that support quantization."""

1196

1197 hf_to_vllm_mapper: "WeightsMapper | None" = None

1175 - packed_modules_mapping: ClassVar[dict[str, list[str]]]

1198 + packed_modules_mapping: dict[str, list[str]]

1199 quant_config: QuantizationConfig | None = None

1200

1201 def __new__(cls, *args, **kwargs) -> Self:

@@ -1504,6 +1527,9 @@ class LocalArgmaxMixin:

1527 ``self.draft_id_to_target_id`` (optional): nn.Parameter

1528 """

1529

1530 + logits_processor: "LogitsProcessor"

1531 + lm_head: "ParallelLMHead"

1532 +

1533 def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor:

1534 """Vocab-parallel argmax with optional D2T remapping."""

1535 top = self.logits_processor.get_top_tokens(

vllm/model_executor/models/interfaces_base.py

+6 -5

@@ -57,7 +57,7 @@ class VllmModel(Protocol[T_co]):

59 def _check_vllm_model_init(model: type[object] | object) -> bool:

60 - model_init = model.__init__

60 + model_init = getattr(model, "__init__") # noqa: B009

61 return supports_kw(model_init, "vllm_config")

@@ -274,10 +274,11 @@ def get_attn_type(model: type[object] | object) -> AttnTypeStr:

274

275

276 def get_score_type(model: type[object] | object) -> ScoreType:

277 - score_types = set()

278 - for m in model.__mro__:

279 - score_type = getattr(m, "score_type", "bi-encoder")

277 + score_types: set[ScoreType] = set()

278 + model_cls = model if isinstance(model, type) else type(model)

279 + for m in model_cls.__mro__:

280 + score_type: ScoreType = getattr(m, "score_type", "bi-encoder")

281 if score_type != "bi-encoder":

282 score_types.add(score_type)

283 assert len(score_types) < 2

283 - return "bi-encoder" if not score_types else list(score_types)[0]

284 + return "bi-encoder" if not score_types else next(iter(score_types))

vllm/model_executor/models/internlm2.py

+1 -1

@@ -323,7 +323,7 @@ class InternLM2Model(nn.Module):

323

324 class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA, SupportsQuant):

325 hf_to_vllm_mapper = InternLM2Model.hf_to_vllm_mapper

326 - packed_modules_mapping = {

326 + packed_modules_mapping: dict[str, list[str]] = {

327 "wqkv": ["wqkv"],

328 "gate_up_proj": ["w1", "w3"],

329 }

vllm/model_executor/models/interns1.py

+31 -19

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

7 # Licensed under The MIT License [see LICENSE for details]

8 # --------------------------------------------------------

9 from collections.abc import Iterable, Mapping, Sequence

10 -from typing import Annotated, Literal, TypeAlias

10 +from typing import Annotated, Literal, TypeAlias, TypedDict

12 import regex as re

13 import torch

@@ -22,7 +22,7 @@ from transformers.models.internvl.video_processing_internvl import (

22 )

24 from vllm.config import VllmConfig

25 -from vllm.config.multimodal import BaseDummyOptions

25 +from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions

26 from vllm.inputs import MultiModalDataDict

27 from vllm.model_executor.layers.quantization import QuantizationConfig

28 from vllm.model_executor.models.interns1_vit import InternS1VisionModel

@@ -148,6 +148,11 @@ class InternS1VideoEmbeddingInputs(TensorSchema):

148 InternS1VideoInputs: TypeAlias = InternS1VideoPixelInputs | InternS1VideoEmbeddingInputs

149

150

151 +class InternS1MultiModalInputs(TypedDict, total=False):

152 + images: InternS1ImageInputs | None

153 + videos: InternS1VideoInputs | None

154 +

155 +

156 def resolve_interns1_min_max_num(

157 min_dynamic_patch: int,

158 max_dynamic_patch: int,

@@ -312,6 +317,7 @@ class InternS1DummyInputsBuilder(BaseDummyInputsBuilder[InternS1ProcessingInfo])

317

318 image_overrides = mm_options.get("image")

319 video_overrides = mm_options.get("video")

320 + assert video_overrides is None or isinstance(video_overrides, VideoDummyOptions)

321

322 return {

323 "image": self._get_dummy_images(

@@ -474,22 +480,24 @@ class InternS1MultiModalProcessor(BaseMultiModalProcessor[InternS1ProcessingInfo

480 video_token = hf_processor.video_token

481

482 tokenizer = self.info.get_tokenizer()

477 - video_token_id = cached_encode(tokenizer, video_token, add_special_tokens=False)

478 - assert len(video_token_id) == 1

479 - video_token_id = video_token_id[0]

483 + video_token_ids = cached_encode(

484 + tokenizer, video_token, add_special_tokens=False

485 + )

486 + assert len(video_token_ids) == 1

487 + video_token_id = video_token_ids[0]

488

489 out_mm_data = out_mm_kwargs.get_data()

490 if "video_num_patches" in out_mm_data:

483 - video_num_patches = out_mm_data["video_num_patches"]

484 - assert isinstance(video_num_patches, torch.Tensor)

485 - video_num_patches = video_num_patches.tolist()

491 + video_num_patches_tensor = out_mm_data["video_num_patches"]

492 + assert isinstance(video_num_patches_tensor, torch.Tensor)

493 + video_num_patches: list[int] = video_num_patches_tensor.tolist()

494 else:

495 video_num_patches = []

496

497 if "image_num_patches" in out_mm_data:

490 - image_num_patches = out_mm_data["image_num_patches"]

491 - assert isinstance(image_num_patches, torch.Tensor)

492 - image_num_patches = image_num_patches.tolist()

498 + image_num_patches_tensor = out_mm_data["image_num_patches"]

499 + assert isinstance(image_num_patches_tensor, torch.Tensor)

500 + image_num_patches: list[int] = image_num_patches_tensor.tolist()

501 else:

502 image_num_patches = []

503

@@ -602,8 +610,8 @@ class InternS1ForConditionalGeneration(

610 prefix=maybe_prefix(prefix, "language_model"),

611 )

612

605 - self.img_context_token_id = None

606 - self.video_context_token_id = None

613 + self.img_context_token_id: int | None = None

614 + self.video_context_token_id: int | None = None

615

616 self.visual_token_mask = None

617 self.make_empty_intermediate_tensors = (

@@ -756,8 +764,10 @@ class InternS1ForConditionalGeneration(

764 ]

765 return image_embeds.split(image_feature_sizes)

766

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

760 - modalities = {}

767 + def _parse_and_validate_multimodal_inputs(

768 + self, **kwargs: object

769 + ) -> InternS1MultiModalInputs:

770 + modalities: InternS1MultiModalInputs = {}

771

772 # Preserve the order of modalities if there are multiple of them

773 # from the order of kwargs.

@@ -789,12 +799,14 @@ class InternS1ForConditionalGeneration(

799 for modality in modalities:

800 if modality == "images":

801 image_input = modalities["images"]

792 - image_embeddings = self._process_vision_input(image_input)

793 - multimodal_embeddings += tuple(image_embeddings)

802 + if image_input is not None:

803 + image_embeddings = self._process_vision_input(image_input)

804 + multimodal_embeddings += tuple(image_embeddings)

805 if modality == "videos":

806 video_input = modalities["videos"]

796 - video_embeddings = self._process_vision_input(video_input)

797 - multimodal_embeddings += tuple(video_embeddings)

807 + if video_input is not None:

808 + video_embeddings = self._process_vision_input(video_input)

809 + multimodal_embeddings += tuple(video_embeddings)

810

811 return multimodal_embeddings

812

vllm/model_executor/models/interns1_pro.py

+18 -12

@@ -26,7 +26,7 @@

27 import functools

28 from collections.abc import Iterable

29 -from typing import Any

29 +from typing import Any, ClassVar

31 import torch

32 from torch import nn

@@ -321,6 +321,14 @@ class InternS1ProMoeAttention(nn.Module):

321 dual_chunk_attention_config=dual_chunk_attention_config,

322 )

323

324 + attn_kwargs: dict[str, Any] = (

325 + {

326 + "layer_idx": extract_layer_index(prefix),

327 + "dual_chunk_attention_config": dual_chunk_attention_config,

328 + }

329 + if dual_chunk_attention_config

330 + else {}

331 + )

332 self.attn = Attention(

333 self.num_heads,

334 self.head_dim,

@@ -329,12 +337,7 @@ class InternS1ProMoeAttention(nn.Module):

337 cache_config=cache_config,

338 quant_config=quant_config,

339 prefix=f"{prefix}.attn",

332 - **{

333 - "layer_idx": extract_layer_index(prefix),

334 - "dual_chunk_attention_config": dual_chunk_attention_config,

335 - }

336 - if dual_chunk_attention_config

337 - else {},

340 + **attn_kwargs,

341 )

342

343 self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

@@ -489,6 +492,9 @@ class InternS1ProMoeLLMForCausalLM(Qwen3MoeForCausalLM):

492

493

494 class InternS1ProMoeMixtureOfExperts(MixtureOfExperts):

495 + language_model: InternS1ProMoeLLMForCausalLM

496 + num_local_physical_experts: int

497 +

498 def update_physical_experts_metadata(

499 self,

500 num_physical_experts: int,

@@ -538,8 +544,8 @@ class InternS1ProMoeMixtureOfExperts(MixtureOfExperts):

544 class InternS1ProForConditionalGeneration(

545 Qwen3VLForConditionalGeneration, InternS1ProMoeMixtureOfExperts

546 ):

541 - is_3d_moe_weight: bool = True

542 - packed_modules_mapping = {

547 + is_3d_moe_weight: ClassVar[bool] = True

548 + packed_modules_mapping: dict[str, list[str]] = {

549 "qkv_proj": [

550 "q_proj",

551 "k_proj",

@@ -559,7 +565,7 @@ class InternS1ProForConditionalGeneration(

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

566 super(Qwen3VLForConditionalGeneration, self).__init__()

567 config: PretrainedConfig = vllm_config.model_config.hf_config

562 - multimodal_config = vllm_config.model_config.multimodal_config

568 + multimodal_config = vllm_config.model_config.get_multimodal_config()

569

570 self.config = config

571 self.multimodal_config = multimodal_config

@@ -601,8 +607,8 @@ class InternS1ProForConditionalGeneration(

607 # Set MoE hyperparameters

608 self.set_moe_parameters()

609

604 - def get_frope_params_map(self) -> str:

605 - mapper = {}

610 + def get_frope_params_map(self) -> dict[str, str]:

611 + mapper: dict[str, str] = {}

612 for name, params in self.language_model.model.named_parameters():

613 if "rotary_emb.sin_coef" in name:

614 mapper["language_model.model.rotary_emb.sin_coef"] = (

vllm/model_executor/models/interns2_mobius.py

+1 -1

@@ -476,7 +476,7 @@ class InternS2MobiusForConditionalGeneration(InternS2PreviewForConditionalGenera

476 if getattr(config, "image_token_index", None) is None:

477 config.image_token_index = config.image_token_id

478 quant_config = vllm_config.quant_config

479 - multimodal_config = vllm_config.model_config.multimodal_config

479 + multimodal_config = vllm_config.model_config.get_multimodal_config()

480

481 self.config = config

482 self.model_config = vllm_config.model_config

vllm/model_executor/models/internvl.py

+29 -16

@@ -10,14 +10,14 @@

10 from abc import abstractmethod

11 from collections.abc import Iterable, Mapping, Sequence

12 from functools import cached_property

13 -from typing import Annotated, Any, Literal, TypeAlias, TypeVar

13 +from typing import Annotated, Any, Literal, TypeAlias, TypedDict, TypeVar

15 import torch

16 import torch.nn as nn

17 from transformers import BatchFeature, PretrainedConfig

19 from vllm.config import VllmConfig

20 -from vllm.config.multimodal import BaseDummyOptions

20 +from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions

21 from vllm.inputs import MultiModalDataDict

22 from vllm.model_executor.layers.quantization import QuantizationConfig

23 from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig

@@ -128,6 +128,11 @@ class InternVLVideoEmbeddingInputs(TensorSchema):

128 InternVLVideoInputs: TypeAlias = InternVLVideoPixelInputs | InternVLVideoEmbeddingInputs

129

130

131 +class InternVLMultiModalInputs(TypedDict, total=False):

132 + images: InternVLImageInputs | None

133 + videos: InternVLVideoInputs | None

134 +

135 +

136 class BaseInternVLProcessingInfo(BaseProcessingInfo):

137 """Basic image-only ProcessingInfo for InternVL-style models."""

138

@@ -286,6 +291,7 @@ class BaseInternVLMultiModalProcessor(BaseMultiModalProcessor[_I]):

291 if isinstance(images, ImageEmbeddingItems):

292 feature_size = images.get_feature_size(item_idx)

293 else:

294 + assert isinstance(images, ImageProcessorItems)

295 image_size = images.get_image_size(item_idx)

296 feature_size = self.info.get_num_image_tokens(

297 image_width=image_size.width,

@@ -433,6 +439,9 @@ class InternVLDummyInputsBuilder(

439 )

440 num_videos = mm_counts.get("video", 0)

441 video_overrides = mm_options.get("video")

442 + assert video_overrides is None or isinstance(

443 + video_overrides, VideoDummyOptions

444 + )

445 dummy_video = {

446 "video": self._get_dummy_videos(

447 width=image_size,

@@ -504,9 +513,9 @@ class InternVLMultiModalProcessor(

513 tokenizer = self.info.get_tokenizer()

514

515 if "video_num_patches" in out_mm_data:

507 - video_num_patches = out_mm_data["video_num_patches"]

508 - assert isinstance(video_num_patches, torch.Tensor)

509 - video_num_patches = video_num_patches.tolist()

516 + video_num_patches_tensor = out_mm_data["video_num_patches"]

517 + assert isinstance(video_num_patches_tensor, torch.Tensor)

518 + video_num_patches: list[int] = video_num_patches_tensor.tolist()

519 else:

520 video_num_patches = []

521

@@ -591,7 +600,7 @@ class InternVLChatModel(

600

601 config = vllm_config.model_config.hf_config

602 quant_config = vllm_config.quant_config

594 - multimodal_config = vllm_config.model_config.multimodal_config

603 + multimodal_config = vllm_config.model_config.get_multimodal_config()

604

605 self.config = config

606 self.multimodal_config = multimodal_config

@@ -621,15 +630,15 @@ class InternVLChatModel(

630 prefix=maybe_prefix(prefix, "language_model"),

631 )

632

624 - self.img_context_token_id = None

625 - self.video_context_token_id = None

633 + self.img_context_token_id: int | None = None

634 + self.video_context_token_id: int | None = None

635

636 self.make_empty_intermediate_tensors = (

637 self.language_model.make_empty_intermediate_tensors

638 )

639

640 def _patch_quant_config(

632 - self, config: PretrainedConfig, quant_config: QuantizationConfig

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

642 ):

643 # the awq models from OpenGVLab missing `modules_to_not_convert`

644 # patch the quant_config to add `modules_to_not_convert` back

@@ -743,7 +752,7 @@ class InternVLChatModel(

752

753 def _parse_and_validate_video_input(

754 self, **kwargs: object

746 - ) -> InternVLVideoPixelInputs | None:

755 + ) -> InternVLVideoInputs | None:

756 pixel_values_flat_video = kwargs.pop("pixel_values_flat_video", None)

757 video_num_patches = kwargs.pop("video_num_patches", None)

758 video_embeds = kwargs.pop("image_embeds", None)

@@ -804,8 +813,10 @@ class InternVLChatModel(

813 ]

814 return image_embeds.split(image_feature_sizes)

815

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

808 - modalities = {}

816 + def _parse_and_validate_multimodal_inputs(

817 + self, **kwargs: object

818 + ) -> InternVLMultiModalInputs:

819 + modalities: InternVLMultiModalInputs = {}

820

821 # Preserve the order of modalities if there are multiple of them

822 # from the order of kwargs.

@@ -834,12 +845,14 @@ class InternVLChatModel(

845 for modality in modalities:

846 if modality == "images":

847 image_input = modalities["images"]

837 - image_embeddings = self._process_vision_input(image_input)

838 - multimodal_embeddings += tuple(image_embeddings)

848 + if image_input is not None:

849 + image_embeddings = self._process_vision_input(image_input)

850 + multimodal_embeddings += tuple(image_embeddings)

851 if modality == "videos":

852 video_input = modalities["videos"]

841 - video_embeddings = self._process_vision_input(video_input)

842 - multimodal_embeddings += tuple(video_embeddings)

853 + if video_input is not None:

854 + video_embeddings = self._process_vision_input(video_input)

855 + multimodal_embeddings += tuple(video_embeddings)

856

857 return multimodal_embeddings

858

vllm/model_executor/models/iquest_loopcoder.py

+10 -6

@@ -16,7 +16,7 @@

17 from __future__ import annotations

19 -from collections.abc import Iterable

19 +from collections.abc import Callable, Iterable

20 from dataclasses import replace

21 from typing import Any

@@ -487,11 +487,14 @@ class IQuestLoopCoderModel(nn.Module):

487 continue

488 if name.endswith("scale"):

489 # Remapping the name of FP8 kv-scale.

490 - name = maybe_remap_kv_scale_name(name, params_dict)

491 - if name is None:

490 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

491 + if remapped_name is None:

492 continue

493 + name = remapped_name

494 param = params_dict[name]

494 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

495 + weight_loader: Callable[..., None] = getattr(

496 + param, "weight_loader", default_weight_loader

497 + )

498 if weight_loader == default_weight_loader:

499 weight_loader(param, loaded_weight)

500 else:

@@ -518,9 +521,10 @@ class IQuestLoopCoderModel(nn.Module):

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

522 continue

523 # Remapping the name of FP8 kv-scale.

521 - name = maybe_remap_kv_scale_name(name, params_dict)

522 - if name is None:

524 + remapped_name = maybe_remap_kv_scale_name(name, params_dict)

525 + if remapped_name is None:

526 continue

527 + name = remapped_name

528 param = params_dict[name]

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

530 weight_loader(param, loaded_weight)

vllm/model_executor/models/isaac.py

+6 -2

@@ -901,7 +901,11 @@ class IsaacForConditionalGeneration(

901 for mm_feature in sorted(mm_features, key=lambda f: f.mm_position.offset):

902 offset = mm_feature.mm_position.offset

903 if mm_feature.modality == "image":

904 - t, h, w = mm_feature.data["image_grid_thw"].data.tolist()

904 + mm_data = mm_feature.data

905 + assert mm_data is not None

906 + grid_thw = mm_data["image_grid_thw"].data

907 + assert isinstance(grid_thw, torch.Tensor)

908 + t, h, w = grid_thw.tolist()

909 assert t == 1, f"Image must have 1 frame, got {t}"

910 yield offset, h // spatial_merge_size, w // spatial_merge_size

911 else:

@@ -912,7 +916,7 @@ class IsaacForConditionalGeneration(

916 input_tokens: list[int],

917 mm_features: list[MultiModalFeatureSpec],

918 ) -> tuple[torch.Tensor, int]:

915 - llm_pos_ids_list = []

919 + llm_pos_ids_list: list[np.ndarray] = []

920 st = 0

921 for offset, llm_grid_h, llm_grid_w in self.iter_mm_grid_hw(

922 input_tokens, mm_features

vllm/model_executor/models/qwen3_vl.py

+1 -1

@@ -1753,7 +1753,7 @@ class Qwen3VLForConditionalGeneration(

1753 SupportsEagle3,

1754 SupportsMultiModalPruning,

1755 ):

1756 - packed_modules_mapping = {

1756 + packed_modules_mapping: dict[str, list[str]] = {

1757 "qkv_proj": [

1758 "q_proj",

1759 "k_proj",