@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e7e3e6d263

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 02:15 UTC

Voxtral (#20970)

Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

Patrick von Platen committed Jul 15, 2025 at 16:35 UTC e7e3e6d2636f6cd012c7ffeff773b20b3c90b958

14 files changed +913 -47

examples/offline_inference/audio_language.py

+75 -10

@@ -10,7 +10,7 @@ on HuggingFace model repository.

11 import os

12 from dataclasses import asdict

13 -from typing import NamedTuple, Optional

13 +from typing import Any, NamedTuple, Optional

15 from huggingface_hub import snapshot_download

16 from transformers import AutoTokenizer

@@ -30,7 +30,9 @@ question_per_audio_count = {

31 class ModelRequestData(NamedTuple):

32 engine_args: EngineArgs

33 - prompt: str

33 + prompt: Optional[str] = None

34 + prompt_token_ids: Optional[dict[str, list[int]]] = None

35 + multi_modal_data: Optional[dict[str, Any]] = None

36 stop_token_ids: Optional[list[int]] = None

37 lora_requests: Optional[list[LoRARequest]] = None

@@ -40,6 +42,60 @@ class ModelRequestData(NamedTuple):

42 # Unless specified, these settings have been tested to work on a single L4.

45 +# Voxtral

46 +def run_voxtral(question: str, audio_count: int) -> ModelRequestData:

47 + from mistral_common.audio import Audio

48 + from mistral_common.protocol.instruct.messages import (

49 + AudioChunk,

50 + RawAudio,

51 + TextChunk,

52 + UserMessage,

53 + )

54 + from mistral_common.protocol.instruct.request import ChatCompletionRequest

55 + from mistral_common.tokens.tokenizers.mistral import MistralTokenizer

56 +

57 + model_name = "mistralai/Voxtral-Mini-3B-2507"

58 + tokenizer = MistralTokenizer.from_hf_hub(model_name)

59 +

60 + engine_args = EngineArgs(

61 + model=model_name,

62 + max_model_len=8192,

63 + max_num_seqs=2,

64 + limit_mm_per_prompt={"audio": audio_count},

65 + config_format="mistral",

66 + load_format="mistral",

67 + tokenizer_mode="mistral",

68 + enforce_eager=True,

69 + enable_chunked_prefill=False,

70 + )

71 +

72 + text_chunk = TextChunk(text=question)

73 + audios = [

74 + Audio.from_file(str(audio_assets[i].get_local_path()), strict=False)

75 + for i in range(audio_count)

76 + ]

77 + audio_chunks = [

78 + AudioChunk(input_audio=RawAudio.from_audio(audio)) for audio in audios

79 + ]

80 +

81 + messages = [UserMessage(content=[*audio_chunks, text_chunk])]

82 +

83 + req = ChatCompletionRequest(messages=messages, model=model_name)

84 +

85 + tokens = tokenizer.encode_chat_completion(req)

86 + prompt_ids, audios = tokens.tokens, tokens.audios

87 +

88 + audios_and_sr = [(au.audio_array, au.sampling_rate) for au in audios]

89 +

90 + multi_modal_data = {"audio": audios_and_sr}

91 +

92 + return ModelRequestData(

93 + engine_args=engine_args,

94 + prompt_token_ids=prompt_ids,

95 + multi_modal_data=multi_modal_data,

96 + )

97 +

98 +

99 # Granite Speech

100 def run_granite_speech(question: str, audio_count: int) -> ModelRequestData:

101 # NOTE - the setting in this example are somehat different than what is

@@ -243,6 +299,7 @@ def run_whisper(question: str, audio_count: int) -> ModelRequestData:

299

300

301 model_example_map = {

302 + "voxtral": run_voxtral,

303 "granite_speech": run_granite_speech,

304 "minicpmo": run_minicpmo,

305 "phi4_mm": run_phi4mm,

@@ -311,16 +368,24 @@ def main(args):

368 temperature=0.2, max_tokens=64, stop_token_ids=req_data.stop_token_ids

369 )

370

314 - mm_data = {}

315 - if audio_count > 0:

316 - mm_data = {

317 - "audio": [

318 - asset.audio_and_sample_rate for asset in audio_assets[:audio_count]

319 - ]

320 - }

371 + mm_data = req_data.multi_modal_data

372 + if not mm_data:

373 + mm_data = {}

374 + if audio_count > 0:

375 + mm_data = {

376 + "audio": [

377 + asset.audio_and_sample_rate for asset in audio_assets[:audio_count]

378 + ]

379 + }

380

381 assert args.num_prompts > 0

323 - inputs = {"prompt": req_data.prompt, "multi_modal_data": mm_data}

382 + inputs = {"multi_modal_data": mm_data}

383 +

384 + if req_data.prompt:

385 + inputs["prompt"] = req_data.prompt

386 + else:

387 + inputs["prompt_token_ids"] = req_data.prompt_token_ids

388 +

389 if args.num_prompts > 1:

390 # Batch inference

391 inputs = [inputs] * args.num_prompts

requirements/common.txt

+1 -1

@@ -33,7 +33,7 @@ pyzmq >= 25.0.0

33 msgspec

34 gguf >= 0.13.0

35 importlib_metadata; python_version < '3.10'

36 -mistral_common[opencv] >= 1.6.2

36 +mistral_common[opencv] >= 1.8.0

37 opencv-python-headless >= 4.11.0 # required for video IO

38 pyyaml

39 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12

requirements/nightly_torch_test.txt

+1 -1

@@ -23,7 +23,7 @@ jiwer # required for audio tests

23 timm # required for internvl test

24 transformers_stream_generator # required for qwen-vl test

25 matplotlib # required for qwen-vl test

26 -mistral_common[opencv] >= 1.6.2 # required for pixtral test

26 +mistral_common[opencv] >= 1.8.0 # required for voxtral test

27 num2words # required for smolvlm test

28 opencv-python-headless >= 4.11.0 # required for video test

29 datamodel_code_generator # required for minicpm3 test

requirements/test.in

+1 -1

@@ -28,7 +28,7 @@ torchvision==0.22.0

28 transformers_stream_generator # required for qwen-vl test

29 mamba_ssm # required for plamo2 test

30 matplotlib # required for qwen-vl test

31 -mistral_common[opencv] >= 1.7.0 # required for pixtral test

31 +mistral_common[opencv] >= 1.8.0 # required for voxtral test

32 num2words # required for smolvlm test

33 opencv-python-headless >= 4.11.0 # required for video test

34 datamodel_code_generator # required for minicpm3 test

requirements/test.txt

+7 -1

@@ -305,7 +305,7 @@ mbstrdecoder==1.1.3

305 # typepy

306 mdurl==0.1.2

307 # via markdown-it-py

308 -mistral-common==1.7.0

308 +mistral-common==1.8.0

309 # via -r requirements/test.in

310 more-itertools==10.5.0

311 # via lm-eval

@@ -518,6 +518,8 @@ pyasn1-modules==0.4.2

518 # via google-auth

519 pybind11==2.13.6

520 # via lm-eval

521 +pycountry==24.6.1

522 + # via pydantic-extra-types

523 pycparser==2.22

524 # via cffi

525 pycryptodomex==3.22.0

@@ -528,9 +530,12 @@ pydantic==2.11.5

530 # datamodel-code-generator

531 # mistral-common

532 # mteb

533 + # pydantic-extra-types

534 # ray

535 pydantic-core==2.33.2

536 # via pydantic

537 +pydantic-extra-types==2.10.5

538 + # via mistral-common

539 pygments==2.18.0

540 # via rich

541 pyparsing==3.2.0

@@ -835,6 +840,7 @@ typing-extensions==4.12.2

840 # pqdm

841 # pydantic

842 # pydantic-core

843 + # pydantic-extra-types

844 # torch

845 # typer

846 # typing-inspection

setup.py

+2 -1

@@ -692,7 +692,8 @@ setup(

692 "tensorizer": ["tensorizer==2.10.1"],

693 "fastsafetensors": ["fastsafetensors >= 0.1.10"],

694 "runai": ["runai-model-streamer", "runai-model-streamer-s3", "boto3"],

695 - "audio": ["librosa", "soundfile"], # Required for audio processing

695 + "audio": ["librosa", "soundfile",

696 + "mistral_common[audio]"], # Required for audio processing

697 "video": [] # Kept for backwards compatibility

698 },

699 cmdclass=cmdclass,

tests/entrypoints/openai/test_transcription_validation.py

+23 -5

@@ -17,6 +17,11 @@ from vllm.assets.audio import AudioAsset

18 from ...utils import RemoteOpenAIServer

20 +MISTRAL_FORMAT_ARGS = [

21 + "--tokenizer_mode", "mistral", "--config_format", "mistral",

22 + "--load_format", "mistral"

23 +]

24 +

26 @pytest.fixture

27 def mary_had_lamb():

@@ -33,9 +38,18 @@ def winning_call():

40 @pytest.mark.asyncio

36 -async def test_basic_audio(mary_had_lamb):

37 - model_name = "openai/whisper-large-v3-turbo"

41 +@pytest.mark.parametrize(

42 + "model_name",

43 + ["openai/whisper-large-v3-turbo", "mistralai/Voxtral-Mini-3B-2507"])

44 +async def test_basic_audio(mary_had_lamb, model_name):

45 server_args = ["--enforce-eager"]

46 +

47 + if model_name.startswith("mistralai"):

48 + server_args += MISTRAL_FORMAT_ARGS

49 +

50 + # TODO(PATRICK) - REMOVE AFTER RELEASE

51 + return # skip for now

52 +

53 # Based on https://github.com/openai/openai-cookbook/blob/main/examples/Whisper_prompting_guide.ipynb.

54 with RemoteOpenAIServer(model_name, server_args) as remote_server:

55 client = remote_server.get_async_client()

@@ -65,10 +79,13 @@ async def test_bad_requests(mary_had_lamb):

81 @pytest.mark.asyncio

68 -async def test_long_audio_request(mary_had_lamb):

69 - model_name = "openai/whisper-large-v3-turbo"

82 +@pytest.mark.parametrize("model_name", ["openai/whisper-large-v3-turbo"])

83 +async def test_long_audio_request(mary_had_lamb, model_name):

84 server_args = ["--enforce-eager"]

86 + if model_name.startswith("openai"):

87 + return

88 +

89 mary_had_lamb.seek(0)

90 audio, sr = librosa.load(mary_had_lamb)

91 # Add small silence after each audio for repeatability in the split process

@@ -87,7 +104,8 @@ async def test_long_audio_request(mary_had_lamb):

104 response_format="text",

105 temperature=0.0)

106 out = json.loads(transcription)['text']

90 - assert out.count("Mary had a little lamb") == 10

107 + counts = out.count("Mary had a little lamb")

108 + assert counts == 10, counts

109

110

111 @pytest.mark.asyncio

tests/models/registry.py

+2 -1

@@ -440,6 +440,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {

440 tokenizer="Isotr0py/Florence-2-tokenizer", # noqa: E501

441 trust_remote_code=True), # noqa: E501

442 "MllamaForConditionalGeneration": _HfExamplesInfo("meta-llama/Llama-3.2-11B-Vision-Instruct"), # noqa: E501

443 + "VoxtralForConditionalGeneration": _HfExamplesInfo("mistralai/Voxtral-Mini-3B-2507", is_available_online=False, tokenizer_mode="mistral"), # noqa: E501

444 "WhisperForConditionalGeneration": _HfExamplesInfo("openai/whisper-large-v3"), # noqa: E501

445

446 # [Cross-encoder]

@@ -513,4 +514,4 @@ class HfExampleModels:

514 raise ValueError(f"No example model defined for {model_id}")

515

516

516 -HF_EXAMPLE_MODELS = HfExampleModels(_EXAMPLE_MODELS)

\ No newline at end of file

517 +HF_EXAMPLE_MODELS = HfExampleModels(_EXAMPLE_MODELS)

vllm/entrypoints/openai/speech_to_text.py

@@ -112,6 +112,7 @@ class OpenAISpeechToText(OpenAIServing):

112 prompt = self.model_cls.get_generation_prompt(

113 audio=chunk,

114 stt_config=self.asr_config,

115 + model_config=self.model_config,

116 language=lang,

117 task_type=self.task_type,

118 request_prompt=request.prompt)

vllm/model_executor/models/interfaces.py

+2 -1

@@ -722,7 +722,8 @@ class SupportsTranscription(Protocol):

722

723 @classmethod

724 def get_generation_prompt(cls, audio: np.ndarray,

725 - stt_config: SpeechToTextConfig, language: str,

725 + stt_config: SpeechToTextConfig,

726 + model_config: ModelConfig, language: str,

727 task_type: str,

728 request_prompt: str) -> PromptType:

729 """Get the prompt for the ASR model.

vllm/model_executor/models/registry.py

@@ -231,6 +231,7 @@ _MULTIMODAL_MODELS = {

231 "Phi4MMForCausalLM": ("phi4mm", "Phi4MMForCausalLM"),

232 "TarsierForConditionalGeneration": ("tarsier", "TarsierForConditionalGeneration"), # noqa: E501

233 "Tarsier2ForConditionalGeneration": ("qwen2_vl", "Tarsier2ForConditionalGeneration"), # noqa: E501

234 + "VoxtralForConditionalGeneration": ("voxtral", "VoxtralForConditionalGeneration"), # noqa: E501

235 # [Encoder-decoder]

236 "Florence2ForConditionalGeneration": ("florence2", "Florence2ForConditionalGeneration"), # noqa: E501

237 "MllamaForConditionalGeneration": ("mllama", "MllamaForConditionalGeneration"), # noqa: E501

vllm/model_executor/models/voxtral.py new

+691

@@ -0,0 +1,691 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +

4 +import math

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

6 +from functools import cached_property

7 +from math import ceil

8 +from typing import Optional, Union, cast

9 +

10 +import numpy as np

11 +import regex as re

12 +import torch

13 +import torch.nn as nn

14 +from mistral_common.audio import mel_filter_bank

15 +from mistral_common.protocol.instruct.messages import (AudioChunk, RawAudio,

16 + TextChunk, UserMessage)

17 +from mistral_common.protocol.instruct.request import ChatCompletionRequest

18 +from mistral_common.protocol.transcription.request import TranscriptionRequest

19 +from mistral_common.tokens.tokenizers.audio import Audio, AudioEncoder

20 +from transformers import TensorType, WhisperConfig

21 +from transformers.tokenization_utils_base import TextInput

22 +

23 +from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig

24 +from vllm.inputs.data import PromptType

25 +from vllm.logger import init_logger

26 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

27 +from vllm.model_executor.models import SupportsPP

28 +# yapf: disable

29 +from vllm.model_executor.models.whisper import (

30 + WhisperEncoder, WhisperForConditionalGeneration)

31 +# yapf: enable

32 +from vllm.model_executor.sampling_metadata import SamplingMetadata

33 +from vllm.multimodal import MULTIMODAL_REGISTRY

34 +from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,

35 + MultiModalKwargs, NestedTensors)

36 +from vllm.multimodal.parse import (AudioProcessorItems, MultiModalDataItems,

37 + MultiModalDataParser)

38 +from vllm.multimodal.processing import (BaseMultiModalProcessor,

39 + BaseProcessingInfo, MultiModalHashes,

40 + PromptReplacement, PromptUpdate)

41 +from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs

42 +from vllm.sequence import IntermediateTensors

43 +from vllm.transformers_utils.tokenizer import (MistralTokenizer,

44 + cached_tokenizer_from_config)

45 +

46 +from .interfaces import (MultiModalEmbeddings, SupportsMultiModal,

47 + SupportsTranscription)

48 +from .utils import (flatten_bn, init_vllm_registered_model, maybe_prefix,

49 + merge_multimodal_embeddings)

50 +

51 +logger = init_logger(__name__)

52 +

53 +

54 +class VoxtralProcessorAdapter:

55 + """

56 + Provide a HF-compatible interface for

57 + :class:`mistral_common.tokens.tokenizers.multimodal.AudioEncoder`.

58 + """

59 +

60 + def __init__(self, tokenizer: MistralTokenizer) -> None:

61 + super().__init__()

62 + self.tokenizer = tokenizer

63 +

64 + @cached_property

65 + def _audio_processor(self) -> AudioEncoder:

66 + audio_encoder = self.tokenizer.instruct.audio_encoder

67 + assert isinstance(audio_encoder, AudioEncoder)

68 + return audio_encoder

69 +

70 + @cached_property

71 + def audio_token_id(self) -> int:

72 + return self._audio_processor.special_ids.audio

73 +

74 + @cached_property

75 + def begin_audio_token_id(self) -> int:

76 + return self._audio_processor.special_ids.begin_audio

77 +

78 + # @cached_property

79 + # def begin_transcript_token_id(self) -> int:

80 + # return self._audio_processor.special_ids.begin_transcript

81 +

82 + # @cached_property

83 + # def end_transcript_token_id(self) -> int:

84 + # return self._audio_processor.special_ids.end_transcript

85 +

86 + @cached_property

87 + def sampling_rate(self) -> int:

88 + return self._audio_processor.audio_config.sampling_rate

89 +

90 + @cached_property

91 + def frame_rate(self) -> float:

92 + return self._audio_processor.audio_config.frame_rate

93 +

94 + def get_num_audio_tokens(

95 + self,

96 + audio_length: int,

97 + ) -> int:

98 + pad_audio_length = self._audio_processor.next_multiple_of_chunk_frames(

99 + audio_length, self.sampling_rate)

100 + return ceil(pad_audio_length / (self.sampling_rate // self.frame_rate))

101 +

102 + def __call__(

103 + self,

104 + text: Optional[Union[TextInput, list[TextInput]]] = None,

105 + audios: Optional[Union[np.ndarray, list[np.ndarray]]] = None,

106 + return_tensors: Optional[Union[str, TensorType]] = None,

107 + **kwargs,

108 + ) -> Mapping[str, NestedTensors]:

109 + if text is None:

110 + text = []

111 + if not isinstance(text, list):

112 + text = [text]

113 + if audios is None:

114 + audios = []

115 + if not isinstance(audios, list):

116 + audios = [audios]

117 +

118 + if not audios:

119 + input_ids = self.tokenizer(text).input_ids

120 + return {"input_ids": torch.tensor(input_ids)}

121 +

122 + # Allow dummy text, which is used for profiling as well as token inputs

123 + if any(len(t) > 0 for t in text):

124 + raise ValueError(

125 + "You've passed text inputs instead of token inputs. "

126 + "Make sure to process your input via `mistral_common`'s "

127 + "tokenizer or pass a chat completion request. "

128 + "For more info, see: "

129 + "https://github.com/vllm-project/vllm/issues/8411.")

130 +

131 + audios_tokens = list[torch.Tensor]()

132 + audios_processed = list[torch.Tensor]()

133 + for audio in audios:

134 + assert isinstance(audio, np.ndarray)

135 + assert audio.ndim == 1

136 +

137 + # pad if necessary

138 + audio = self._audio_processor.pad(audio, self.sampling_rate)

139 +

140 + audio_tokens = [

141 + self.begin_audio_token_id

142 + ] + [self.audio_token_id] * self.get_num_audio_tokens(len(audio))

143 +

144 + audios_tokens.append(torch.tensor(audio_tokens))

145 + audios_processed.append(torch.tensor(audio))

146 +

147 + return {

148 + "input_ids": torch.cat(audios_tokens)[None].expand(len(text), -1),

149 + "audio_arrays": audios_processed,

150 + }

151 +

152 +

153 +class VoxtralProcessingInfo(BaseProcessingInfo):

154 +

155 + def get_tokenizer(self) -> MistralTokenizer:

156 + tokenizer = cached_tokenizer_from_config(self.ctx.model_config)

157 + if not isinstance(tokenizer, MistralTokenizer):

158 + raise ValueError("This model requires `--tokenizer-mode mistral`")

159 +

160 + return tokenizer

161 +

162 + def get_hf_processor(self) -> VoxtralProcessorAdapter:

163 + return VoxtralProcessorAdapter(self.get_tokenizer())

164 +

165 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

166 + return {"audio": 5} # Performance tends to degrade after 5

167 +

168 + def get_mm_max_tokens_per_item(

169 + self,

170 + seq_len: int,

171 + mm_counts: Mapping[str, int],

172 + ) -> Mapping[str, int]:

173 + return {"audio": self.get_max_audio_tokens()}

174 +

175 + def get_max_audio_tokens(self) -> int:

176 + return self.ctx.model_config.max_model_len

177 +

178 + def get_max_audio_array_len(self) -> int:

179 + processor = self.get_hf_processor()

180 + return self.get_max_audio_tokens() * int(

181 + processor.sampling_rate // processor.frame_rate)

182 +

183 +

184 +class VoxtralDummyInputsBuilder(BaseDummyInputsBuilder[VoxtralProcessingInfo]):

185 +

186 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

187 + return ""

188 +

189 + def get_dummy_mm_data(

190 + self,

191 + seq_len: int,

192 + mm_counts: Mapping[str, int],

193 + ) -> MultiModalDataDict:

194 + num_audios = mm_counts.get("audio", 0)

195 +

196 + target_length = self.info.get_max_audio_array_len()

197 +

198 + return {

199 + "audio":

200 + self._get_dummy_audios(length=target_length, num_audios=num_audios)

201 + }

202 +

203 + def get_dummy_processor_inputs(

204 + self,

205 + seq_len: int,

206 + mm_counts: Mapping[str, int],

207 + ) -> ProcessorInputs:

208 + tokenizer = self.info.get_tokenizer()

209 +

210 + dummy_text = self.get_dummy_text(mm_counts)

211 + dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts)

212 + dummy_audios = dummy_mm_data.get("audio", [])

213 +

214 + audio_chunks: list[AudioChunk] = []

215 + format = "wav"

216 + for audio in dummy_audios:

217 + audio_item = Audio(

218 + audio_array=audio,

219 + sampling_rate=self.info.get_hf_processor().sampling_rate,

220 + format=format,

221 + )

222 + chunk = AudioChunk(input_audio=RawAudio.from_audio(audio_item))

223 + audio_chunks.append(chunk)

224 +

225 + request = ChatCompletionRequest(messages=[

226 + UserMessage(content=[TextChunk(text=dummy_text), *audio_chunks]),

227 + ])

228 + res = tokenizer.mistral.encode_chat_completion(request)

229 + dummy_tokens = res.tokens

230 + # whixtral tokenizer adds padding to the audio

231 + # so we need to update the audio arrays

232 + dummy_mm_data["audio"] = [a.audio_array for a in res.audios]

233 +

234 + return ProcessorInputs(prompt=dummy_tokens, mm_data=dummy_mm_data)

235 +

236 +

237 +class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo]

238 + ):

239 +

240 + def _get_mm_fields_config(

241 + self,

242 + hf_inputs: Mapping[str, NestedTensors],

243 + hf_processor_mm_kwargs: Mapping[str, object],

244 + ) -> Mapping[str, MultiModalFieldConfig]:

245 + return dict(audio_arrays=MultiModalFieldConfig.batched("audio"))

246 +

247 + def _get_prompt_updates(

248 + self,

249 + mm_items: MultiModalDataItems,

250 + hf_processor_mm_kwargs: Mapping[str, object],

251 + out_mm_kwargs: MultiModalKwargs,

252 + ) -> Sequence[PromptUpdate]:

253 + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

254 +

255 + audio_id = processor.audio_token_id

256 +

257 + def get_replacement(item_idx: int):

258 + audios = mm_items.get_items("audio", AudioProcessorItems)

259 + audio_len = audios.get_audio_length(item_idx)

260 +

261 + nb_audio_tokens = processor.get_num_audio_tokens(audio_len)

262 +

263 + return [audio_id] * nb_audio_tokens

264 +

265 + return [

266 + PromptReplacement(

267 + modality="audio",

268 + target="", # Never match the prompt (see below note)

269 + replacement=get_replacement,

270 + ),

271 + ]

272 +

273 + def _cached_apply_hf_processor(

274 + self,

275 + prompt: Union[str, list[int]],

276 + mm_data_items: MultiModalDataItems,

277 + hf_processor_mm_kwargs: Mapping[str, object],

278 + tokenization_kwargs: Mapping[str, object],

279 + *,

280 + return_mm_hashes: bool,

281 + ) -> tuple[list[int], MultiModalKwargs, Optional[MultiModalHashes], bool]:

282 + prompt_ids, mm_kwargs, mm_hashes, _ = super(

283 + )._cached_apply_hf_processor(

284 + prompt=prompt,

285 + mm_data_items=mm_data_items,

286 + hf_processor_mm_kwargs=hf_processor_mm_kwargs,

287 + tokenization_kwargs=tokenization_kwargs,

288 + return_mm_hashes=return_mm_hashes,

289 + )

290 +

291 + # NOTE: The tokens are already inserted by the chat template

292 + return prompt_ids, mm_kwargs, mm_hashes, True

293 +

294 + def _get_data_parser(self) -> MultiModalDataParser:

295 + sampling_rate = self.info.get_hf_processor().sampling_rate

296 + return MultiModalDataParser(target_sr=sampling_rate)

297 +

298 +

299 +@MULTIMODAL_REGISTRY.register_processor(VoxtralMultiModalProcessor,

300 + info=VoxtralProcessingInfo,

301 + dummy_inputs=VoxtralDummyInputsBuilder)

302 +class VoxtralForConditionalGeneration(nn.Module, SupportsMultiModal,

303 + SupportsPP, SupportsTranscription):

304 +

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

306 + super().__init__()

307 + self.tokenizer = cached_tokenizer_from_config(vllm_config.model_config)

308 +

309 + config = vllm_config.model_config.hf_config

310 + self.config = config

311 + self.downsample_factor = self.config.audio_config.downsample_factor

312 +

313 + self.language_model = init_vllm_registered_model(

314 + vllm_config=vllm_config,

315 + hf_config=config.text_config,

316 + prefix=maybe_prefix(prefix, "language_model"),

317 + )

318 + self.whisper_encoder = VoxtralEncoderModel(

319 + vllm_config.with_hf_config(config.audio_config),

320 + prefix=maybe_prefix(prefix, "whisper_encoder"),

321 + )

322 + self.audio_language_adapter = AudioLanguageAdapter(

323 + hidden_size=config.audio_config.d_model * self.downsample_factor,

324 + dim=config.text_config.hidden_size,

325 + )

326 +

327 + def get_language_model(self) -> torch.nn.Module:

328 + return self.language_model

329 +

330 + def forward(

331 + self,

332 + input_ids: torch.Tensor,

333 + positions: torch.Tensor,

334 + intermediate_tensors: Optional[IntermediateTensors] = None,

335 + inputs_embeds: Optional[torch.Tensor] = None,

336 + **kwargs: object,

337 + ) -> Union[torch.Tensor, IntermediateTensors]:

338 + if intermediate_tensors is not None:

339 + inputs_embeds = None

340 +

341 + # NOTE: In v1, inputs_embeds is always generated at model runner, this

342 + # condition is for v0 compatibility.

343 + elif inputs_embeds is None:

344 + audio_embeddings = self.get_multimodal_embeddings(**kwargs)

345 + inputs_embeds = self.get_input_embeddings(input_ids,

346 + audio_embeddings)

347 + input_ids = None

348 +

349 + hidden_states = self.language_model.model(input_ids,

350 + positions,

351 + intermediate_tensors,

352 + inputs_embeds=inputs_embeds)

353 +

354 + return hidden_states

355 +

356 + def get_multimodal_embeddings(

357 + self, **kwargs

358 + ) -> Union[list[torch.Tensor], torch.Tensor, tuple[torch.Tensor, ...],

359 + None]:

360 + audio_inputs = self._parse_and_validate_audio_arrays(**kwargs)

361 + if audio_inputs is None:

362 + return None

363 +

364 + audio_embeddings = self.whisper_encoder(audio_inputs)

365 +

366 + for i, audio_embedding in enumerate(audio_embeddings):

367 + seq_len, dim = audio_embedding.shape

368 + # Pad such that seq_len is divisible by downsample_factor

369 + target_seq_len = self.downsample_factor * math.ceil(

370 + seq_len / self.downsample_factor)

371 + audio_embedding = torch.nn.functional.pad(

372 + audio_embedding,

373 + (0, 0, 0, target_seq_len - seq_len),

374 + )

375 + audio_embeddings[i] = audio_embedding.reshape(

376 + target_seq_len // self.downsample_factor,

377 + dim * self.downsample_factor)

378 +

379 + # Concat, project and resplit

380 + audio_embeddings_packed = torch.cat(audio_embeddings, dim=0)

381 + audio_embeddings_packed = self.audio_language_adapter(

382 + audio_embeddings_packed)

383 + audio_embeddings = torch.split(audio_embeddings_packed,

384 + [a.shape[0] for a in audio_embeddings],

385 + dim=0)

386 +

387 + return audio_embeddings

388 +

389 + def get_input_embeddings(

390 + self,

391 + input_ids: torch.Tensor,

392 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

393 + ) -> torch.Tensor:

394 + audio_encoder = self.tokenizer.instruct.audio_encoder

395 + audio_tok_id = audio_encoder.audio_token

396 +

397 + inputs_embeds = self.language_model.get_input_embeddings(input_ids)

398 + if multimodal_embeddings is not None:

399 + inputs_embeds = merge_multimodal_embeddings(

400 + input_ids, inputs_embeds, multimodal_embeddings, audio_tok_id)

401 + return inputs_embeds

402 +

403 + def _parse_and_validate_audio_arrays(

404 + self, **kwargs: object) -> Union[list[torch.Tensor], None]:

405 + audio_arrays = kwargs.pop("audio_arrays", None)

406 + if audio_arrays is None:

407 + return None

408 +

409 + if not isinstance(audio_arrays, (torch.Tensor, list)):

410 + raise ValueError("Incorrect type of audio_arrays. "

411 + f"Got type: {type(audio_arrays)}")

412 +

413 + audio_arrays = flatten_bn(audio_arrays)

414 + if isinstance(audio_arrays, torch.Tensor):

415 + audio_arrays = list(audio_arrays.unbind(0))

416 + return audio_arrays

417 +

418 + def compute_logits(

419 + self,

420 + hidden_states: torch.Tensor,

421 + sampling_metadata: SamplingMetadata,

422 + ) -> Optional[torch.Tensor]:

423 + return self.language_model.compute_logits(hidden_states,

424 + sampling_metadata)

425 +

426 + @classmethod

427 + def get_speech_to_text_config(cls, model_config: ModelConfig,

428 + task_type: str) -> SpeechToTextConfig:

429 + tokenizer = cached_tokenizer_from_config(model_config)

430 + audio_config = tokenizer.instruct.audio_encoder.audio_config

431 + max_audio_clip_s = audio_config.chunk_length_s

432 + sample_rate = audio_config.sampling_rate

433 + return SpeechToTextConfig(

434 + max_audio_clip_s=max_audio_clip_s,

435 + sample_rate=sample_rate,

436 + # mistral_common and whisper encoder take care of chunking

437 + min_energy_split_window_size=None,

438 + )

439 +

440 + @classmethod

441 + # for speech-to-text transcription

442 + def get_generation_prompt(cls, audio: np.ndarray,

443 + model_config: ModelConfig,

444 + stt_config: SpeechToTextConfig, language: str,

445 + task_type: str,

446 + request_prompt: str) -> PromptType:

447 + tokenizer = cached_tokenizer_from_config(model_config)

448 + audio = Audio(audio, int(stt_config.sample_rate),

449 + format="wav") # lossless

450 + req = TranscriptionRequest(model=model_config.model,

451 + audio=RawAudio.from_audio(audio),

452 + language=language)

453 +

454 + tokenized = tokenizer.instruct.encode_transcription(req)

455 + audio = (tokenized.audios[0].audio_array, stt_config.sample_rate)

456 + prompts_dict = {"multi_modal_data": {"audio": audio}}

457 + prompts_dict["prompt_token_ids"] = tokenized.tokens

458 + return cast(PromptType, prompts_dict)

459 +

460 + @classmethod

461 + def validate_language(cls, language: str) -> bool:

462 + # same as whisper

463 + return WhisperForConditionalGeneration.validate_language(language)

464 +

465 + @classmethod

466 + def get_num_audio_tokens(cls, audio_duration_s: float,

467 + stt_config: SpeechToTextConfig,

468 + model_config: ModelConfig) -> Optional[int]:

469 + """

470 + Map from audio duration to number of audio tokens produced by the ASR

471 + model, without running a forward pass.

472 + This is used for estimating the amount of processing for this audio.

473 + """

474 + tokenizer = cached_tokenizer_from_config(model_config)

475 + adapter = VoxtralProcessorAdapter(tokenizer)

476 + return adapter.get_num_audio_tokens(

477 + int(audio_duration_s * stt_config.sample_rate))

478 +

479 + def load_weights(self, weights: Iterable[tuple[str,

480 + torch.Tensor]]) -> set[str]:

481 + # fmt: off

482 + remapping_rules = [

483 + (r"mm_whisper_embeddings\.(.*)", r"\1"),

484 + (r"audio_language_projection\.(.*)", r"audio_language_adapter.\1"),

485 + (r"audio_language_adapter\.0\.weight", r"audio_language_adapter.w_in.weight"), # noqa: E501

486 + (r"audio_language_adapter\.2\.weight", r"audio_language_adapter.w_out.weight"), # noqa: E501

487 + ]

488 + # fmt: on

489 +

490 + audio_params = dict(

491 + nn.ModuleDict({

492 + "audio_language_adapter":

493 + self.audio_language_adapter,

494 + }).named_parameters())

495 +

496 + loaded_weights = set()

497 +

498 + def llm_weights_generator():

499 + nonlocal loaded_weights

500 + for name, w in weights:

501 + is_encoder = (

502 + name.startswith("mm_whisper_embeddings") and

503 + not name.startswith("mm_whisper_embeddings.tok_embeddings")

504 + and not name.startswith(

505 + "mm_whisper_embeddings.audio_language_projection"))

506 +

507 + for pattern, repl in remapping_rules:

508 + if re.fullmatch(pattern, name):

509 + name = re.sub(pattern, repl, name)

510 +

511 + if is_encoder:

512 + name = self.whisper_encoder.load_weight((name, w))

513 + loaded_weights.add(f"whisper_encoder.{name}")

514 + continue

515 +

516 + if name in audio_params:

517 + param = audio_params[name]

518 + with torch.no_grad():

519 + default_weight_loader(param, w)

520 + loaded_weights.add(name)

521 + else:

522 + yield (name, w)

523 +

524 + for name in self.language_model.load_weights(llm_weights_generator()):

525 + loaded_weights.add(f"language_model.{name}")

526 +

527 + # potentially manually add position embeddings

528 + sin_key = "whisper_encoder.whisper_encoder.embed_positions.weight"

529 + if sin_key not in loaded_weights:

530 + # make sure we don't hit an error here

531 + loaded_weights.add(sin_key)

532 +

533 + return loaded_weights

534 +

535 +

536 +class AudioLanguageAdapter(nn.Module):

537 +

538 + def __init__(self, hidden_size: int, dim: int) -> None:

539 + super().__init__()

540 + self.w_in = nn.Linear(hidden_size, dim, bias=False)

541 + self.gelu = nn.GELU()

542 + self.w_out = nn.Linear(dim, dim, bias=False)

543 +

544 + def forward(self, x: torch.Tensor) -> torch.Tensor:

545 + return self.w_out(self.gelu(self.w_in(x)))

546 +

547 +

548 +class VoxtralEncoderModel(nn.Module):

549 + packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}

550 +

551 + # fmt: off

552 + mistral_remapping = [

553 + (r"whisper_encoder\.conv_layers\.0\.(weight|bias)", r"whisper_encoder.conv1.\1"), # noqa: E501

554 + (r"whisper_encoder\.conv_layers\.1\.(weight|bias)", r"whisper_encoder.conv2.\1"), # noqa: E501

555 + (r"whisper_encoder\.transformer\.layers\.(\d+)\.attention\.w([qkv])\.(weight|bias)", r"whisper_encoder.layers.\1.self_attn.\2_proj.\3"), # noqa: E501

556 + (r"whisper_encoder\.transformer\.layers\.(\d+)\.attention\.wo\.(weight|bias)", r"whisper_encoder.layers.\1.self_attn.out_proj.\2"), # noqa: E501

557 + (r"whisper_encoder\.transformer\.layers\.(\d+)\.attention_norm\.(weight|bias)", r"whisper_encoder.layers.\1.self_attn_layer_norm.\2"), # noqa: E501

558 + (r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w1\.(weight|bias)", r"whisper_encoder.layers.\1.mlp.fc1.\2"), # noqa: E501

559 + (r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w2\.(weight|bias)", r"whisper_encoder.layers.\1.mlp.fc2.\2"), # noqa: E501

560 + (r"whisper_encoder\.transformer\.layers\.(\d+)\.ffn_norm\.(weight|bias)", r"whisper_encoder.layers.\1.final_layer_norm.\2"), # noqa: E501

561 + (r"whisper_encoder\.transformer\.norm\.(weight|bias)", r"whisper_encoder.layer_norm.\1"), # noqa: E501

562 + ]

563 + # fmt: on

564 +

565 + def __init__(

566 + self,

567 + vllm_config: VllmConfig,

568 + *,

569 + prefix: str = "",

570 + ) -> None:

571 + super().__init__()

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

573 + self.dtype: torch.dtype = vllm_config.model_config.dtype

574 + self.whisper_encoder = WhisperEncoder(vllm_config=vllm_config,

575 + prefix=maybe_prefix(

576 + prefix, "whisper_encoder"),

577 + is_standalone_encoder=True,

578 + init_in_fp32=True)

579 + mel_filters = mel_filter_bank(

580 + num_frequency_bins=1 + self.config.window_size // 2,

581 + num_mel_bins=self.config.num_mel_bins,

582 + min_frequency=0.0,

583 + max_frequency=8000.0,

584 + sampling_rate=self.config.sampling_rate,

585 + )

586 + self.mel_filters = torch.tensor(mel_filters, dtype=torch.float32)

587 +

588 + def compute_whisper_melspec(

589 + self,

590 + audio_waveforms: torch.Tensor,

591 + ) -> torch.Tensor:

592 + input_dtype = audio_waveforms.dtype

593 + window = torch.hann_window(self.config.window_size).to(

594 + audio_waveforms.device)

595 + stft = torch.stft(

596 + audio_waveforms,

597 + self.config.window_size,

598 + self.config.hop_length,

599 + window=window,

600 + return_complex=True,

601 + )

602 + magnitudes = stft[..., :-1].abs()**2

603 + mel_spec = self.mel_filters.T @ magnitudes

604 + log_spec = torch.clamp(mel_spec, min=1e-10).log10()

605 + log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)

606 + log_spec = (log_spec + 4.0) / 4.0

607 + return log_spec.to(input_dtype)

608 +

609 + @property

610 + def downsample_factor(self) -> int:

611 + return self.whisper_encoder.conv1.stride[

612 + 0] * self.whisper_encoder.conv2.stride[0]

613 +

614 + @property

615 + def chunk_size(self) -> int:

616 + return self.config.max_source_positions * self.downsample_factor

617 +

618 + def prepare_inputs_for_conv(

619 + self,

620 + audio_waveforms: list[torch.Tensor],

621 + ) -> tuple[torch.Tensor, list[int]]:

622 + assert isinstance(audio_waveforms, list)

623 + # list[num_mel_bins, seq_len]

624 + input_features = [

625 + self.compute_whisper_melspec(audio).to(self.dtype)

626 + for audio in audio_waveforms

627 + ]

628 +

629 + chunked_features: list[torch.Tensor] = []

630 + chunks_per_example: list[int] = []

631 + for feature in input_features:

632 + chunks = feature.split(self.chunk_size, dim=-1)

633 + chunked_features += chunks

634 + chunks_per_example.append(len(chunks))

635 +

636 + # [total_num_chunks, num_mel_bins, chunk_size]

637 + return torch.stack(chunked_features), chunks_per_example

638 +

639 + def forward(

640 + self, input_features: Union[torch.Tensor, list[torch.Tensor]]

641 + ) -> list[torch.Tensor]:

642 + if not isinstance(input_features, list):

643 + input_features = [input_features]

644 +

645 + # Split long inputs into chunks

646 + input_embeds, chunks_per_example = (

647 + self.prepare_inputs_for_conv(input_features))

648 +

649 + # [total_num_chunks, ceil(chunk_size / downsample_factor), hidden_size]

650 + out = self.whisper_encoder([input_embeds])

651 +

652 + # Re-concatenate the chunks

653 + chunk_idx = 0

654 + results = []

655 + for n_chunks in chunks_per_example:

656 + result = out[chunk_idx:chunk_idx + n_chunks].flatten(0, 1)

657 + results.append(result)

658 + chunk_idx += n_chunks

659 +

660 + return results

661 +

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

663 + stacked_params_mapping = [

664 + # (param_name, shard_name, shard_id)

665 + ("qkv_proj", "q_proj", "q"),

666 + ("qkv_proj", "k_proj", "k"),

667 + ("qkv_proj", "v_proj", "v"),

668 + ]

669 + params_dict = dict(self.named_parameters())

670 +

671 + name, loaded_weight = weight

672 + for pattern, repl in self.mistral_remapping:

673 + if re.fullmatch(pattern, name):

674 + name = re.sub(pattern, repl, name)

675 +

676 + for (param_name, weight_name, shard_id) in stacked_params_mapping:

677 + if weight_name not in name:

678 + continue

679 + name = name.replace(weight_name, param_name)

680 +

681 + param = params_dict[name]

682 + weight_loader = param.weight_loader

683 + weight_loader(param, loaded_weight, shard_id)

684 + break

685 + else:

686 + param = params_dict[name]

687 + weight_loader = getattr(param, "weight_loader",

688 + default_weight_loader)

689 + weight_loader(param, loaded_weight)

690 +

691 + return name

vllm/model_executor/models/whisper.py

+59 -22

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

4 import math

5 from collections.abc import Iterable, Mapping, Sequence

6 +from contextlib import nullcontext

7 from typing import Optional, TypedDict, Union, cast

9 import numpy as np

@@ -13,6 +14,7 @@ from transformers import (BatchFeature, WhisperConfig, WhisperFeatureExtractor,

14 from transformers.models.whisper.modeling_whisper import sinusoids

16 from vllm.attention import Attention, AttentionType

17 +from vllm.attention.layer import MultiHeadAttention

18 from vllm.config import (CacheConfig, ModelConfig, SpeechToTextConfig,

19 VllmConfig)

20 from vllm.distributed import get_tensor_model_parallel_world_size

@@ -26,6 +28,7 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor

28 from vllm.model_executor.layers.quantization.base_config import (

29 QuantizationConfig)

30 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

31 +from vllm.model_executor.model_loader.utils import set_default_torch_dtype

32 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

33 from vllm.model_executor.sampling_metadata import SamplingMetadata

34 from vllm.multimodal import MULTIMODAL_REGISTRY, NestedTensors

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

181 cache_config: Optional[CacheConfig] = None,

182 quant_config: Optional[QuantizationConfig] = None,

183 prefix: str = "",

184 + standalone_encoder: bool = False,

185 ):

186 super().__init__()

187 self.embed_dim = embed_dim

@@ -213,16 +217,24 @@ class WhisperAttention(nn.Module):

217 quant_config=quant_config,

218 prefix=f"{prefix}.out_proj",

219 )

216 - self.attn = Attention(

217 - self.num_heads,

218 - self.head_dim,

219 - self.scaling,

220 - num_kv_heads=self.num_kv_heads,

221 - cache_config=cache_config,

222 - quant_config=quant_config,

223 - prefix=f"{prefix}.attn",

224 - attn_type=self.attn_type,

225 - )

220 + if standalone_encoder:

221 + self.attn = MultiHeadAttention(

222 + self.num_heads,

223 + self.head_dim,

224 + self.scaling,

225 + num_kv_heads=self.num_kv_heads,

226 + )

227 + else:

228 + self.attn = Attention(

229 + self.num_heads,

230 + self.head_dim,

231 + self.scaling,

232 + num_kv_heads=self.num_kv_heads,

233 + cache_config=cache_config,

234 + quant_config=quant_config,

235 + prefix=f"{prefix}.attn",

236 + attn_type=self.attn_type,

237 + )

238

239 def _init_qkv(

240 self,

@@ -357,7 +369,11 @@ class WhisperMLP(nn.Module):

369

370 class WhisperEncoderLayer(nn.Module):

371

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

372 + def __init__(self,

373 + *,

374 + vllm_config: VllmConfig,

375 + prefix: str = "",

376 + is_standalone_encoder: bool = False):

377 super().__init__()

378 config = vllm_config.model_config.hf_config

379 cache_config = vllm_config.cache_config

@@ -371,6 +387,7 @@ class WhisperEncoderLayer(nn.Module):

387 cache_config=cache_config,

388 quant_config=quant_config,

389 prefix=f"{prefix}.self_attn",

390 + standalone_encoder=is_standalone_encoder,

391 )

392 self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)

393 self.mlp = WhisperMLP(

@@ -462,10 +479,16 @@ class WhisperDecoderLayer(nn.Module):

479

480 class WhisperEncoder(nn.Module):

481

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

482 + def __init__(self,

483 + *,

484 + vllm_config: VllmConfig,

485 + prefix: str = "",

486 + is_standalone_encoder: bool = False,

487 + init_in_fp32: bool = False):

488 super().__init__()

489 config = vllm_config.model_config.hf_config

490 embed_dim = config.d_model

491 + self.is_standalone_encoder = is_standalone_encoder

492 self.num_mel_bins = config.num_mel_bins

493 self.max_source_positions = config.max_source_positions

494 self.embed_scale = (math.sqrt(embed_dim)

@@ -480,17 +503,25 @@ class WhisperEncoder(nn.Module):

503 kernel_size=3,

504 stride=2,

505 padding=1)

483 - self.embed_positions = nn.Embedding(self.max_source_positions,

484 - embed_dim)

506 self.start_layer, self.end_layer, self.layers = make_layers(

507 config.encoder_layers,

508 lambda prefix: WhisperEncoderLayer(vllm_config=vllm_config,

488 - prefix=f"{prefix}.layers"),

509 + prefix=f"{prefix}.layers",

510 + is_standalone_encoder=

511 + is_standalone_encoder),

512 prefix=f"{prefix}.layers",

513 )

514 self.layer_norm = nn.LayerNorm(config.d_model)

515

493 - with torch.no_grad():

516 + maybe_fp32_init_ctx = set_default_torch_dtype(

517 + torch.float32) if init_in_fp32 else nullcontext()

518 +

519 + with (

520 + torch.no_grad(),

521 + maybe_fp32_init_ctx,

522 + ):

523 + self.embed_positions = nn.Embedding(self.max_source_positions,

524 + embed_dim)

525 self.embed_positions.weight.copy_(

526 sinusoids(*self.embed_positions.weight.shape))

527

@@ -499,8 +530,10 @@ class WhisperEncoder(nn.Module):

530 for features in input_features:

531 embeds = nn.functional.gelu(self.conv1(features))

532 embeds = nn.functional.gelu(self.conv2(embeds))

502 - embeds = embeds.permute(1, 0)

503 - embeds = embeds + self.embed_positions.weight[:embeds.size(0), :]

533 + embeds = embeds.transpose(-1, -2)

534 + embeds = (embeds +

535 + self.embed_positions.weight[:embeds.size(-2), :]).to(

536 + embeds.dtype)

537 hidden_states.append(embeds)

538 hidden_states = torch.cat(hidden_states)

539

@@ -792,10 +825,14 @@ class WhisperForConditionalGeneration(nn.Module, SupportsTranscription,

825 f"or {list(ISO639_1_OTHER_LANGS.values())}")

826

827 @classmethod

795 - def get_generation_prompt(cls, audio: np.ndarray,

796 - stt_config: SpeechToTextConfig, language: str,

797 - task_type: str,

798 - request_prompt: str) -> PromptType:

828 + def get_generation_prompt(

829 + cls,

830 + audio: np.ndarray,

831 + model_config: ModelConfig, # not needed here

832 + stt_config: SpeechToTextConfig,

833 + language: str,

834 + task_type: str,

835 + request_prompt: str) -> PromptType:

836 prompt = {

837 "encoder_prompt": {

838 # Whisper does not support encoder prompt.

vllm/transformers_utils/configs/mistral.py

+47 -3

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

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

3 from typing import Any

5 -from transformers import PretrainedConfig

5 +from transformers import PretrainedConfig, WhisperConfig

7 from vllm.logger import init_logger

@@ -24,9 +24,21 @@ def adapt_config_dict(config_dict: dict[str, Any],

25 if bool(config_dict.get("yarn")):

26 config_dict = _remap_mistral_yarn_args(config_dict)

27 - if bool((config_dict.get("multimodal") or {}).get("vision_encoder_args")

28 - or config_dict.get("vision_encoder")):

27 +

28 + is_vision = ((config_dict.get("multimodal")

29 + or {}).get("vision_encoder_args")

30 + or config_dict.get("vision_encoder"))

31 + is_audio = bool(

32 + ((config_dict.get("multimodal") or {}).get("whisper_model_args")

33 + or {}).get("encoder_args"))

34 +

35 + assert not (is_vision and is_audio), \

36 + "Vision and audio are mutually exclusive"

37 +

38 + if is_vision:

39 config_dict = _remap_mistral_vision_args(config_dict)

40 + if is_audio:

41 + config_dict = _remap_mistral_audio_args(config_dict)

43 config = PretrainedConfig.from_dict(config_dict)

@@ -118,3 +130,35 @@ def _remap_mistral_quantization_args(config: dict) -> dict:

130 config["quantization_config"] = quantization_config

131

132 return config

133 +

134 +

135 +def _remap_mistral_audio_args(config: dict) -> dict:

136 + whisper_args = config["multimodal"].pop("whisper_model_args")

137 + encoder_args = whisper_args["encoder_args"]

138 + downsample_args = whisper_args["downsample_args"]

139 +

140 + quant_config = config.get("quantization_config")

141 + config = {

142 + "model_type":

143 + "whixtral",

144 + "architectures": ["VoxtralForConditionalGeneration"],

145 + "text_config":

146 + PretrainedConfig.from_dict(config),

147 + "audio_config":

148 + WhisperConfig(

149 + num_mel_bins=encoder_args["audio_encoding_args"]["num_mel_bins"],

150 + window_size=encoder_args["audio_encoding_args"]["window_size"],

151 + sampling_rate=encoder_args["audio_encoding_args"]["sampling_rate"],

152 + hop_length=encoder_args["audio_encoding_args"]["hop_length"],

153 + downsample_factor=downsample_args["downsample_factor"],

154 + d_model=encoder_args["dim"],

155 + encoder_layers=encoder_args["n_layers"],

156 + encoder_ffn_dim=encoder_args["hidden_dim"],

157 + encoder_attention_heads=encoder_args["n_heads"],

158 + vocab_size=encoder_args["vocab_size"],

159 + max_source_positions=encoder_args["max_source_positions"],

160 + )

161 + }

162 + if quant_config:

163 + config["quantization_config"] = quant_config

164 + return config