@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2619b6b72b

[Mypy] Fix mypy typing for Voxtral and vision models (#58251)

Signed-off-by: Ashraf Bhuiyan <mbhuiyan@redhat.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Ashraf Bhuiyan committed Sep 28, 2026 at 15:52 UTC 2619b6b72b24a8c1963ef6b29427c53b826023d1

5 files changed +34 -20

tests/models/multimodal/generation/test_voxtral_realtime.py

+11 -1

@@ -16,7 +16,9 @@ from vllm import LLM, SamplingParams

16 from vllm.assets.audio import AudioAsset

17 from vllm.config import CUDAGraphMode

18 from vllm.engine.arg_utils import AsyncEngineArgs

19 +from vllm.engine.protocol import StreamingInput

20 from vllm.platforms import current_platform

21 +from vllm.renderers.inputs.preprocess import parse_model_prompt

22 from vllm.utils.math_utils import cdiv

23 from vllm.v1.engine.async_llm import AsyncLLM

24 from vllm.v1.kv_cache_interface import SlidingWindowSpec

@@ -222,6 +224,14 @@ async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine)

224 sampling_params = SamplingParams(temperature=0.0, max_tokens=1)

225 audio_config = tokenizer.instruct_tokenizer.audio_encoder.audio_config

226

227 + async def input_stream(audio_buffer):

228 + async for prompt in audio_buffer.get_input_stream():

229 + parsed_prompt = parse_model_prompt(async_engine.model_config, prompt)

230 + (engine_input,) = await async_engine.renderer.render_cmpl_async(

231 + [parsed_prompt]

232 + )

233 + yield StreamingInput(prompt=engine_input)

234 +

235 output_tokens_list = []

236 for i, audio_asset in enumerate(audio_assets):

237 output_tokens = []

@@ -241,7 +251,7 @@ async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine)

251 request_id = f"session-{i}"

252

253 async for resp in async_engine.generate(

244 - prompt=buffer.get_input_stream(),

254 + prompt=input_stream(buffer),

255 sampling_params=sampling_params,

256 request_id=request_id,

257 ):

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/[vV]",

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

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

118 ]

vllm/model_executor/models/vision.py

+7 -4

@@ -5,7 +5,7 @@ import itertools

5 import math

6 from abc import ABC, abstractmethod

7 from collections.abc import Callable

8 -from typing import Final, Generic, Literal, Protocol, TypeAlias, TypeVar

8 +from typing import Generic, Literal, Protocol, TypeAlias, TypeVar

10 import torch

11 from transformers import PreTrainedConfig

@@ -24,10 +24,12 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum

24 logger = init_logger(__name__)

26 _C = TypeVar("_C", bound=PreTrainedConfig)

27 +_C_co = TypeVar("_C_co", bound=PreTrainedConfig, covariant=True)

29 -class _RootConfig(Protocol[_C]):

30 - vision_config: _C

30 +class _RootConfig(Protocol[_C_co]):

31 + @property

32 + def vision_config(self) -> _C_co: ...

35 class VisionEncoderInfo(ABC, Generic[_C]):

@@ -60,7 +62,8 @@ class VisionEncoderInfo(ABC, Generic[_C]):

64 class VisionLanguageConfig(Protocol):

63 - vision_config: Final[PreTrainedConfig]

65 + @property

66 + def vision_config(self) -> PreTrainedConfig: ...

69 def get_vision_encoder_info(hf_config: VisionLanguageConfig) -> VisionEncoderInfo:

vllm/model_executor/models/voxtral.py

+4 -3

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

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import math

5 -from collections.abc import Iterable, Mapping, Sequence

5 +from collections.abc import Callable, Iterable, Mapping, Sequence

6 from functools import partial

7 from typing import cast

@@ -317,7 +317,7 @@ class VoxtralForConditionalGeneration(

317

318 # update quant config to so that ignored module and target module names

319 # match the vLLM model names

320 - if hasattr(vllm_config, "quant_config"):

320 + if vllm_config.quant_config is not None:

321 vllm_config.quant_config = self.maybe_update_quant_config(

322 vllm_config.quant_config

323 )

@@ -724,6 +724,7 @@ class VoxtralEncoderModel(nn.Module):

724 self.config = cast(WhisperConfig, vllm_config.model_config.hf_config)

725 self.dtype: torch.dtype = vllm_config.model_config.dtype

726 self.is_causal = getattr(self.config, "is_causal", False)

727 + WhisperEncoderCls: Callable[..., WhisperEncoder | WhisperCausalEncoder]

728 if self.is_causal:

729 WhisperEncoderCls = WhisperCausalEncoder

730 else:

@@ -830,7 +831,7 @@ class VoxtralEncoderModel(nn.Module):

831 return results

832

833 def load_weight(self, weight: tuple[str, torch.Tensor]) -> str:

833 - stacked_params_mapping = [

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

835 # (param_name, shard_name, shard_id)

836 ("qkv_proj", "q_proj", "q"),

837 ("qkv_proj", "k_proj", "k"),

vllm/model_executor/models/voxtral_realtime.py

+12 -11

@@ -16,7 +16,6 @@ from mistral_common.tokens.tokenizers.audio import Audio, AudioConfig

16 from vllm.compilation.decorators import support_torch_compile

17 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig

18 from vllm.config.speech_to_text import SpeechToTextParams

19 -from vllm.engine.protocol import StreamingInput

19 from vllm.inputs import PromptType, TokensPrompt

20 from vllm.logger import init_logger

21 from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsRealtime

@@ -26,6 +25,7 @@ from vllm.model_executor.models.voxtral import (

25 VoxtralMultiModalProcessor,

26 VoxtralProcessingInfo,

27 )

28 +from vllm.model_executor.models.whisper_causal import WhisperCausalEncoder

29 from vllm.multimodal import MULTIMODAL_REGISTRY

30 from vllm.multimodal.parse import MultiModalDataItems

31 from vllm.multimodal.processing import ProcessorInputs, TimingContext

@@ -56,7 +56,7 @@ class VoxtralRealtimeMultiModalProcessor(VoxtralMultiModalProcessor):

56 mm_items: MultiModalDataItems,

57 mm_res: MultiModalProcessingResult,

58 ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:

59 - mm_kwargs = mm_res.kwargs

59 + mm_kwargs = mm_res.kwargs.require_data()

61 # there are no placeholder audio tokens for streaming

62 # so we need to build the place placeholder positions manually

@@ -69,7 +69,9 @@ class VoxtralRealtimeMultiModalProcessor(VoxtralMultiModalProcessor):

69 tokenizer = self.info.get_tokenizer()

70 audio_config = tokenizer.instruct.audio_encoder.audio_config

72 - num_audio_samples = audios[0]["audio_arrays"].data.shape[0]

72 + audio_array = audios[0]["audio_arrays"].data

73 + assert isinstance(audio_array, (torch.Tensor, np.ndarray))

74 + num_audio_samples = audio_array.shape[0]

75 length = audio_config.num_audio_tokens(num_audio_samples)

77 features_info = PlaceholderFeaturesInfo(

@@ -162,7 +164,7 @@ class VoxtralRealtimeBuffer:

164 for token in tokens:

165 await self._token_queue.put(token)

166

165 - async def get_input_stream(self) -> AsyncGenerator[StreamingInput]:

167 + async def get_input_stream(self) -> AsyncGenerator[TokensPrompt]:

168 for frame_size, num_tokens in self._generate_frame_size_and_num_tokens():

169 next_tokens = [await self._token_queue.get() for _ in range(num_tokens)]

170

@@ -188,11 +190,9 @@ class VoxtralRealtimeBuffer:

190

191 self._leftover = audio_array[stride:]

192

191 - yield StreamingInput(

192 - TokensPrompt(

193 - prompt_token_ids=next_tokens,

194 - multi_modal_data={"audio": (frame, None)},

195 - )

193 + yield TokensPrompt(

194 + prompt_token_ids=next_tokens,

195 + multi_modal_data={"audio": (frame, None)},

196 )

197

198

@@ -274,8 +274,8 @@ class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtim

274 token_task = asyncio.create_task(feed_tokens())

275

276 try:

277 - async for streaming_input in buffer.get_input_stream():

278 - yield streaming_input.prompt

277 + async for prompt in buffer.get_input_stream():

278 + yield prompt

279 finally:

280 audio_task.cancel()

281 token_task.cancel()

@@ -420,6 +420,7 @@ class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtim

420

421 seq_lens = [mel.shape[1] for mel in mel_features]

422 # [total_num_20ms_frames, hidden_size]

423 + assert isinstance(self.whisper_encoder.whisper_encoder, WhisperCausalEncoder)

424 audio_embeddings = self.whisper_encoder.whisper_encoder.forward_conv(

425 mel_features

426 )