[Mypy] Fix mypy typing for Whisper models (#58254)
Signed-off-by: Ashraf Bhuiyan <mbhuiyan@redhat.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Ashraf Bhuiyan committed
Sep 29, 2026 at 12:10 UTC
2454b5a4f0c4c65109fff40e9c60e94bc0245bd2
4 files changed
+30
-16
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/[wW]",
116
r"vllm/model_executor/models/[zZ]",
117
]
118
vllm/model_executor/models/interfaces.py
+3
-2
@@ -61,6 +61,7 @@ if TYPE_CHECKING:
61
from vllm.multimodal.registry import _ProcessorFactories
62
from vllm.sequence import IntermediateTensors
63
from vllm.tasks import ScoreType
64
+ from vllm.tokenizers import TokenizerLike
65
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
66
from vllm.v1.worker.encoder_cudagraph_defs import (
67
EncoderCudaGraphCaptureInputs,
@@ -1449,7 +1450,7 @@ class SupportsTranscription(Protocol):
1450
def parse_language_detection_output(
1451
cls,
1452
token_ids: list[int],
1452
- tokenizer: object,
1453
+ tokenizer: "TokenizerLike",
1454
) -> str:
1455
"""Parse the detected language from model output token IDs.
1456
@@ -1461,7 +1462,7 @@ class SupportsTranscription(Protocol):
1462
@classmethod
1463
def get_language_token_ids(
1464
cls,
1464
- tokenizer: object,
1465
+ tokenizer: "TokenizerLike",
1466
) -> list[int] | None:
1467
"""Return token IDs that represent valid language tokens.
1468
vllm/model_executor/models/whisper.py
+11
-5
@@ -63,6 +63,7 @@ from vllm.multimodal.processing import (
63
)
64
from vllm.multimodal.processing.processor import HFMultiModalInputs
65
from vllm.renderers import TokenizeParams
66
+from vllm.tokenizers import TokenizerLike
67
from vllm.transformers_utils.processor import cached_processor_from_config
68
from vllm.utils.jsontree import json_map_leaves
69
from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -325,7 +326,7 @@ class WhisperCrossAttention(WhisperAttention):
326
prefix=f"{prefix}.kv_proj",
327
)
328
328
- def forward(
329
+ def forward( # type: ignore[override]
330
self,
331
hidden_states: torch.Tensor,
332
encoder_hidden_states: torch.Tensor | None,
@@ -872,7 +873,7 @@ class WhisperForConditionalGeneration(
873
@classmethod
874
def get_language_token_ids(
875
cls,
875
- tokenizer: object,
876
+ tokenizer: TokenizerLike,
877
) -> list[int]:
878
"""Return token IDs for all supported language tokens.
879
@@ -908,8 +909,8 @@ class WhisperForConditionalGeneration(
909
def parse_language_detection_output(
910
cls,
911
token_ids: list[int],
911
- tokenizer: object,
912
- ) -> str | None:
912
+ tokenizer: TokenizerLike,
913
+ ) -> str:
914
"""Parse the language token predicted by Whisper.
915
916
Decodes the first token ID and extracts the language code from the
@@ -1004,6 +1005,7 @@ class WhisperForConditionalGeneration(
1005
audio_input = self._parse_and_validate_audio_input(**kwargs)
1006
# Split concatenated encoder outputs into one tensor per audio input
1007
enc_output = self.model.get_encoder_outputs(audio_input["input_features"])
1008
+ assert enc_output is not None
1009
# The assumption is we can only process whole mm items (audios)
1010
return enc_output.unbind(dim=0)
1011
@@ -1021,8 +1023,12 @@ class WhisperForConditionalGeneration(
1023
def _parse_and_validate_audio_input(self, **kwargs: object) -> WhisperAudioInputs:
1024
input_features = kwargs.pop("input_features", None)
1025
1026
+ def to_dtype(value: object) -> torch.Tensor:
1027
+ assert isinstance(value, torch.Tensor)
1028
+ return value.to(self.dtype)
1029
+
1030
if input_features is not None:
1025
- input_features = json_map_leaves(lambda x: x.to(self.dtype), input_features)
1031
+ input_features = json_map_leaves(to_dtype, input_features)
1032
1033
return WhisperAudioInputs(input_features=input_features)
1034
vllm/model_executor/models/whisper_causal.py
+16
-8
@@ -6,7 +6,7 @@ import math
6
from dataclasses import replace
7
from fractions import Fraction
8
from functools import partial
9
-from typing import ClassVar
9
+from typing import ClassVar, Protocol
10
11
import torch
12
import torch.nn.functional as F
@@ -35,11 +35,6 @@ from vllm.v1.attention.backend import (
35
)
36
from vllm.v1.attention.backends.cpu_attn import CPUAttentionBackend
37
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
38
-
39
-try:
40
- from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
41
-except ImportError:
42
- AiterFlashAttentionBackend = None
38
from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
39
from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
40
from vllm.v1.attention.selector import get_attn_backend
@@ -47,11 +42,24 @@ from vllm.v1.kv_cache_interface import AttentionSpec, SlidingWindowSpec
42
43
from .utils import make_layers
44
45
+AiterFlashAttentionBackend: type[AttentionBackend] | None = None
46
+try:
47
+ from vllm.v1.attention.backends import rocm_aiter_fa
48
+
49
+ AiterFlashAttentionBackend = rocm_aiter_fa.AiterFlashAttentionBackend
50
+except ImportError:
51
+ pass
52
+
53
logger = logging.getLogger(__name__)
54
55
CausalRMSNorm = partial(RMSNorm, eps=1e-5)
56
57
58
+class _BlockPoolingAttentionMetadata(Protocol):
59
+ @property
60
+ def slot_mapping(self) -> torch.Tensor: ...
61
+
62
+
63
def _pad1d(
64
x: torch.Tensor,
65
paddings: tuple[int, int],
@@ -115,7 +123,7 @@ class WhisperCausalConv1d(nn.Conv1d):
123
124
@functools.lru_cache
125
def create_whisper_attention_backend_with_block_pooling(
118
- underlying_attn_backend: AttentionBackend,
126
+ underlying_attn_backend: type[AttentionBackend],
127
block_pool_size: int,
128
sliding_window: int | None = None,
129
) -> type[AttentionBackend]:
@@ -270,7 +278,7 @@ def create_whisper_attention_backend_with_block_pooling(
278
key: torch.Tensor,
279
value: torch.Tensor,
280
kv_cache: torch.Tensor,
273
- attn_metadata: AttentionMetadata,
281
+ attn_metadata: _BlockPoolingAttentionMetadata | None,
282
output: torch.Tensor,
283
output_scale: torch.Tensor | None = None,
284
output_block_scale: torch.Tensor | None = None,