[Mypy] Fix mypy typing for Ultravox and Unlimited-OCR models (#58239)
Signed-off-by: Ashraf Bhuiyan <mbhuiyan@redhat.com>
Ashraf Bhuiyan committed
Sep 27, 2026 at 21:19 UTC
b9e2903b89361a8650540100e120a4dcad6395b1
6 files changed
+36
-27
tools/pre_commit/mypy.py
-1
@@ -113,7 +113,6 @@ SEPARATE_GROUPS = [
113
114
EXCLUDE = [
115
r"vllm/model_executor/models/[kK]",
116
- r"vllm/model_executor/models/[uU]",
116
r"vllm/model_executor/models/[vV]",
117
r"vllm/model_executor/models/[wW]",
118
r"vllm/model_executor/models/[zZ]",
vllm/model_executor/model_loader/reload/torchao_decorator.py
+13
-10
@@ -1,10 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from collections.abc import Iterable
4
+from collections.abc import Callable
5
from functools import wraps
6
-from types import FunctionType
7
-from typing import TYPE_CHECKING
6
+from typing import TYPE_CHECKING, Concatenate, ParamSpec, TypeVar
7
8
import torch
9
@@ -20,13 +19,18 @@ if TYPE_CHECKING:
19
20
__all__ = ["set_torchao_reload_attrs", "support_quantized_model_reload_from_hp_weights"]
21
22
+_P = ParamSpec("_P")
23
+_R = TypeVar("_R")
24
+
25
26
def set_torchao_reload_attrs(model: torch.nn.Module, model_config: ModelConfig):
27
model._do_torchao_reload = True
28
model._model_config = model_config
29
30
29
-def support_quantized_model_reload_from_hp_weights(original_load_weights: FunctionType):
31
+def support_quantized_model_reload_from_hp_weights(
32
+ original_load_weights: Callable[Concatenate["AutoWeightsLoader", _P], _R],
33
+) -> Callable[Concatenate["AutoWeightsLoader", _P], _R]:
34
"""Decorator for `load_weights` method for AutoWeightsLoader.load_weights to support
35
reloading high precision (bfloat16/float16/float32) weight for an already quantized
36
model, this involves restoring the weights to a high precision weights and
@@ -39,17 +43,16 @@ def support_quantized_model_reload_from_hp_weights(original_load_weights: Functi
43
@wraps(original_load_weights)
44
def patched_model_load_weights(
45
self: "AutoWeightsLoader",
42
- weights: Iterable[tuple[str, torch.Tensor]],
43
- *args,
44
- **kwargs,
45
- ):
46
+ *args: _P.args,
47
+ **kwargs: _P.kwargs,
48
+ ) -> _R:
49
model = self.module
50
51
if not getattr(model, "_do_torchao_reload", False):
49
- return original_load_weights(self, weights, *args, **kwargs)
52
+ return original_load_weights(self, *args, **kwargs)
53
54
initialize_layerwise_reload(model)
52
- loaded_weights = original_load_weights(self, weights, *args, **kwargs)
55
+ loaded_weights = original_load_weights(self, *args, **kwargs)
56
finalize_layerwise_reload(model, model._model_config)
57
58
return loaded_weights
vllm/model_executor/models/phi4mm.py
+1
-1
@@ -1297,7 +1297,7 @@ class Phi4MMForCausalLM(nn.Module, SupportsLoRA, SupportsMultiModal):
1297
logits = self.logits_processor(self.lm_head, hidden_states)
1298
return logits
1299
1300
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None:
1300
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1301
loader = AutoWeightsLoader(self)
1302
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1303
vllm/model_executor/models/ultravox.py
+13
-8
@@ -282,10 +282,14 @@ class UltravoxMultiModalProcessor(BaseMultiModalProcessor[UltravoxProcessingInfo
282
283
def get_replacement_ultravox(item_idx: int):
284
if "audio_num_tokens" in out_mm_data:
285
- return [replacement_id] * int(out_mm_data["audio_num_tokens"][item_idx])
285
+ audio_num_tokens = out_mm_data["audio_num_tokens"]
286
+ assert isinstance(audio_num_tokens, torch.Tensor)
287
+ return [replacement_id] * int(audio_num_tokens[item_idx])
288
start = chunks_start_idx[item_idx]
289
end = chunks_start_idx[item_idx + 1]
288
- audio_token_len = out_mm_data["audio_token_len"][start:end].sum()
290
+ audio_token_lens = out_mm_data["audio_token_len"]
291
+ assert isinstance(audio_token_lens, torch.Tensor)
292
+ audio_token_len = audio_token_lens[start:end].sum()
293
return [replacement_id] * int(audio_token_len) # type: ignore
294
295
return [
@@ -566,7 +570,8 @@ class UltravoxWhisperEncoder(WhisperEncoder):
570
) -> torch.Tensor:
571
return (input_lengths - 1) // 2 + 1
572
569
- def forward(
573
+ # Ultravox requires chunk lengths in addition to Whisper's input features.
574
+ def forward( # type: ignore[override]
575
self,
576
input_features: torch.Tensor,
577
audio_lens: torch.Tensor,
@@ -960,12 +965,12 @@ def pad_and_concat_to_dim3(
965
966
return features
967
963
- features = [pad_and_concat_to_dim3(f) for f in features]
968
+ tensors = [pad_and_concat_to_dim3(f) for f in features]
969
965
- max_len = max(f.shape[-1] for f in features)
970
+ max_len = max(f.shape[-1] for f in tensors)
971
# Ensure all features have dim=3
967
- features = [f.view(-1, *f.shape[-2:]) for f in features]
972
+ tensors = [f.view(-1, *f.shape[-2:]) for f in tensors]
973
# Pad and concatenate:
974
# [[B1, 80, M1], [B2, 80, M2]] -> [B1+B2, 80, max(M1, M2)]
970
- features = [F.pad(f, (0, max_len - f.shape[-1])) for f in features]
971
- return torch.cat(features)
975
+ tensors = [F.pad(f, (0, max_len - f.shape[-1])) for f in tensors]
976
+ return torch.cat(tensors)
vllm/model_executor/models/unlimited_ocr.py
+1
@@ -199,6 +199,7 @@ class UnlimitedOCRMultiModalProcessor(DeepseekOCRMultiModalProcessor):
199
if isinstance(images, ImageEmbeddingItems):
200
num_image_tokens = images.get_feature_size(item_idx)
201
else:
202
+ assert isinstance(images, ImageProcessorItems)
203
size = images.get_image_size(item_idx)
204
205
# Disable crop mode for multi-image input.
vllm/model_executor/models/utils.py
+8
-7
@@ -35,7 +35,6 @@ if TYPE_CHECKING:
35
from transformers import PreTrainedConfig
36
from transformers.conversion_mapping import WeightRenaming
37
38
- from vllm.config.model import ModelConfig
38
from vllm.model_executor.layers.quantization import QuantizationConfig
39
40
logger = init_logger(__name__)
@@ -337,7 +336,7 @@ class AutoWeightsLoader:
336
"""
337
# Add persistent registered buffers.
338
# Non-persistent buffers are excluded, matching PyTorch state_dict().
340
- non_persistent = getattr(module, "_non_persistent_buffers_set", set())
339
+ non_persistent: set[str] = getattr(module, "_non_persistent_buffers_set", set())
340
for buf_name, buf in module.named_buffers(recurse=False):
341
if buf_name not in child_params and buf_name not in non_persistent:
342
child_params[buf_name] = buf
@@ -558,7 +557,7 @@ def maybe_fuse_shared_experts(
557
558
559
def get_spec_layer_idx_from_weight_name(
561
- config: "ModelConfig", weight_name: str
560
+ config: "PreTrainedConfig", weight_name: str
561
) -> int | None:
562
"""Return the MTP layer index a weight belongs to, or None.
563
@@ -587,7 +586,7 @@ def get_spec_layer_idx_from_weight_name(
586
587
588
def skip_spec_layers(
590
- weights: Iterable[tuple[str, torch.Tensor]], config: "ModelConfig"
589
+ weights: Iterable[tuple[str, torch.Tensor]], config: "PreTrainedConfig"
590
) -> Iterable[tuple[str, torch.Tensor]]:
591
"""Drop MTP spec-layer weights (loaded by the MTP head, not the base model).
592
@@ -629,11 +628,11 @@ def init_vllm_registered_model(
628
629
630
@overload
632
-def flatten_bn(x: torch.Tensor) -> torch.Tensor: ...
631
+def flatten_bn(x: list[torch.Tensor]) -> list[torch.Tensor]: ...
632
633
634
@overload
636
-def flatten_bn(x: list[torch.Tensor]) -> list[torch.Tensor]: ...
635
+def flatten_bn(x: torch.Tensor) -> torch.Tensor: ...
636
637
638
@overload
@@ -973,7 +972,9 @@ def get_draft_quant_config(vllm_config: VllmConfig) -> "QuantizationConfig | Non
972
The draft model's config if available, None otherwise.
973
974
"""
976
- draft_model_config = vllm_config.speculative_config.draft_model_config
975
+ speculative_config = vllm_config.speculative_config
976
+ assert speculative_config is not None
977
+ draft_model_config = speculative_config.draft_model_config
978
draft_load_config = vllm_config.load_config
979
980
return (