| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import asyncio |
| 5 | import math |
| 6 | from collections.abc import AsyncGenerator, Iterable, Iterator, Mapping |
| 7 | |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | from mistral_common.protocol.transcription.request import ( |
| 11 | StreamingMode, |
| 12 | TranscriptionRequest, |
| 13 | ) |
| 14 | from mistral_common.tokens.tokenizers.audio import Audio, AudioConfig |
| 15 | |
| 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.inputs import PromptType, TokensPrompt |
| 20 | from vllm.logger import init_logger |
| 21 | from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsRealtime |
| 22 | from vllm.model_executor.models.voxtral import ( |
| 23 | VoxtralDummyInputsBuilder, |
| 24 | VoxtralForConditionalGeneration, |
| 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 |
| 32 | from vllm.multimodal.processing.processor import ( |
| 33 | MultiModalProcessingResult, |
| 34 | PlaceholderFeaturesInfo, |
| 35 | ) |
| 36 | from vllm.sequence import IntermediateTensors |
| 37 | from vllm.tokenizers import cached_tokenizer_from_config |
| 38 | from vllm.utils.torch_utils import is_torch_equal_or_newer |
| 39 | |
| 40 | from .utils import _flatten_embeddings |
| 41 | |
| 42 | logger = init_logger(__name__) |
| 43 | |
| 44 | |
| 45 | class VoxtralRealtimeMultiModalProcessor(VoxtralMultiModalProcessor): |
| 46 | def _cached_apply_hf_processor( |
| 47 | self, |
| 48 | inputs: ProcessorInputs, |
| 49 | timing_ctx: TimingContext, |
| 50 | ) -> MultiModalProcessingResult: |
| 51 | # realtime can't make use of a cache yet |
| 52 | return self._apply_hf_processor(inputs, timing_ctx) |
| 53 | |
| 54 | def _maybe_apply_prompt_updates( |
| 55 | self, |
| 56 | mm_items: MultiModalDataItems, |
| 57 | mm_res: MultiModalProcessingResult, |
| 58 | ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: |
| 59 | mm_kwargs = mm_res.kwargs.require_data() |
| 60 | |
| 61 | # there are no placeholder audio tokens for streaming |
| 62 | # so we need to build the place placeholder positions manually |
| 63 | |
| 64 | # in realtime there is always only one audio input |
| 65 | audios = mm_kwargs.get("audio", []) |
| 66 | assert len(audios) == 1, ( |
| 67 | f"Expected only one audio input for realtime, got {mm_kwargs=}" |
| 68 | ) |
| 69 | tokenizer = self.info.get_tokenizer() |
| 70 | audio_config = tokenizer.instruct.audio_encoder.audio_config |
| 71 | |
| 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) |
| 76 | |
| 77 | features_info = PlaceholderFeaturesInfo( |
| 78 | modality="audio", |
| 79 | item_idx=0, |
| 80 | start_idx=0, |
| 81 | tokens=length |
| 82 | * [0], # only used for length computation, so we can take dummy inputs |
| 83 | is_embed=None, |
| 84 | ) |
| 85 | return mm_res.prompt_ids, {"audio": [features_info]} |
| 86 | |
| 87 | |
| 88 | class TimeEmbedding(torch.nn.Module): |
| 89 | """Sinusoidal Embedding for encoding time.""" |
| 90 | |
| 91 | def __init__(self, dim: int, theta: float = 10000.0) -> None: |
| 92 | super().__init__() |
| 93 | self.dim = dim |
| 94 | self.theta = theta |
| 95 | inv_freq = torch.exp( |
| 96 | -math.log(self.theta) |
| 97 | * torch.arange(self.dim // 2).float() |
| 98 | / (self.dim // 2) |
| 99 | ) |
| 100 | self.register_buffer("inv_freq", inv_freq, persistent=False) |
| 101 | |
| 102 | def forward(self, t: torch.Tensor) -> torch.Tensor: |
| 103 | t = t[..., None] # (B,) -> (B, 1) or (B, T) -> (B, T, 1) |
| 104 | inv_freq = self.inv_freq.to(device=t.device, dtype=t.dtype) |
| 105 | emb = ( |
| 106 | t * inv_freq |
| 107 | ) # (B, 1) x (D/2,) -> (B, D/2) or (B, T, 1) x (D/2,) -> (B, T, D/2) |
| 108 | return torch.cat((emb.cos(), emb.sin()), dim=-1) # (B, D) or (B, T, D) |
| 109 | |
| 110 | |
| 111 | def _expand_tensor(input_tensor: torch.Tensor, scaling: int) -> torch.Tensor: |
| 112 | # 1. Multiply by the scaling factor (e.g. 4) |
| 113 | base = input_tensor * scaling |
| 114 | |
| 115 | # 2. Create the offsets, e.g. [0, 1, 2, 3] |
| 116 | offsets = torch.arange(scaling, device=input_tensor.device) |
| 117 | |
| 118 | # 3. Use broadcasting, e.g. (N, 1) + (4,) results in (N, 4) |
| 119 | # Then flatten back to 1D |
| 120 | return (base.unsqueeze(1) + offsets).view(-1) |
| 121 | |
| 122 | |
| 123 | class VoxtralRealtimeBuffer: |
| 124 | def __init__(self, config: AudioConfig, prompt_tokens: list[int]) -> None: |
| 125 | self._config = config |
| 126 | |
| 127 | _look_ahead_in_ms = self._config.streaming_look_ahead_ms |
| 128 | _look_back_in_ms = self._config.streaming_look_back_ms |
| 129 | self._look_ahead_in_samples = self._ms_to_samples(_look_ahead_in_ms) |
| 130 | self._look_back_in_samples = self._ms_to_samples(_look_back_in_ms) |
| 131 | |
| 132 | # None signals the end |
| 133 | self._audio_queue: asyncio.Queue[np.ndarray | None] = asyncio.Queue() |
| 134 | self._leftover: np.ndarray | None = None |
| 135 | self._token_queue: asyncio.Queue[int] = asyncio.Queue() |
| 136 | |
| 137 | self._initial_end = len(prompt_tokens) * self._config.raw_audio_length_per_tok |
| 138 | for token in prompt_tokens: |
| 139 | self._token_queue.put_nowait(token) |
| 140 | |
| 141 | def _generate_frame_size_and_num_tokens(self) -> Iterator[tuple[int, int]]: |
| 142 | streaming_step_size = self._ms_to_samples(1000 / self._config.frame_rate) |
| 143 | start = 0 |
| 144 | end = self._initial_end |
| 145 | while True: |
| 146 | frame_start = max(start - self._look_back_in_samples, 0) |
| 147 | frame_end = end + self._look_ahead_in_samples |
| 148 | frame_size = frame_end - frame_start |
| 149 | num_tokens = (end - start) / self._config.raw_audio_length_per_tok |
| 150 | assert num_tokens.is_integer() |
| 151 | yield frame_size, int(num_tokens) |
| 152 | start = end |
| 153 | end += streaming_step_size |
| 154 | |
| 155 | def _ms_to_samples(self, ms: float) -> int: |
| 156 | len_ = self._config.sampling_rate * ms / 1000 |
| 157 | assert len_.is_integer(), len_ |
| 158 | return int(len_) |
| 159 | |
| 160 | async def append_audio(self, audio_array: np.ndarray | None) -> None: |
| 161 | await self._audio_queue.put(audio_array) |
| 162 | |
| 163 | async def append_tokens(self, tokens: Iterable[int]) -> None: |
| 164 | for token in tokens: |
| 165 | await self._token_queue.put(token) |
| 166 | |
| 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 | |
| 171 | audio_arrays: list[np.ndarray] = ( |
| 172 | [self._leftover] if self._leftover is not None else [] |
| 173 | ) |
| 174 | while sum(len(arr) for arr in audio_arrays) < frame_size: |
| 175 | arr = await self._audio_queue.get() |
| 176 | if arr is None: |
| 177 | return |
| 178 | audio_arrays.append(arr) |
| 179 | |
| 180 | audio_array = np.concatenate(audio_arrays) |
| 181 | frame = audio_array[:frame_size] |
| 182 | |
| 183 | # The current stride took look_ahead_in_samples audio of the next sample |
| 184 | # In addition the next sample will take look_back_in_samples audio of |
| 185 | # the current sample => So let's put both of this into the leftover |
| 186 | stride = ( |
| 187 | frame_size - self._look_ahead_in_samples - self._look_back_in_samples |
| 188 | ) |
| 189 | assert stride > 0, f"{stride=} must be positive" |
| 190 | |
| 191 | self._leftover = audio_array[stride:] |
| 192 | |
| 193 | yield TokensPrompt( |
| 194 | prompt_token_ids=next_tokens, |
| 195 | multi_modal_data={"audio": (frame, None)}, |
| 196 | ) |
| 197 | |
| 198 | |
| 199 | @MULTIMODAL_REGISTRY.register_processor( |
| 200 | VoxtralRealtimeMultiModalProcessor, |
| 201 | info=VoxtralProcessingInfo, |
| 202 | dummy_inputs=VoxtralDummyInputsBuilder, |
| 203 | ) |
| 204 | @support_torch_compile |
| 205 | class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtime): |
| 206 | requires_raw_input_tokens = True |
| 207 | # transformers' currently has limited support for MistralCommon backend |
| 208 | # and cached_get_processor. Let's skip until fixed |
| 209 | skip_warmup_audio_preprocessing = True |
| 210 | |
| 211 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 212 | super().__init__(vllm_config=vllm_config, prefix=prefix) |
| 213 | |
| 214 | # Full cudagraphs are supported for decode-only batches (the encoder's |
| 215 | # block-pooling attention builder is capture-safe for uniform |
| 216 | # single-token decode). Mixed/prefill batches are not, so block only |
| 217 | # pure FULL; FULL_DECODE_ONLY and FULL_AND_PIECEWISE are allowed. |
| 218 | cudagraph_mode = vllm_config.compilation_config.cudagraph_mode |
| 219 | assert not cudagraph_mode.mixed_mode().has_full_cudagraphs(), ( |
| 220 | "Voxtral realtime supports full cudagraphs for decode-only batches. " |
| 221 | "Use cudagraph_mode=FULL_DECODE_ONLY or FULL_AND_PIECEWISE (not FULL)." |
| 222 | ) |
| 223 | |
| 224 | self.time_embedding: TimeEmbedding = TimeEmbedding( |
| 225 | dim=self.config.text_config.hidden_size |
| 226 | ) |
| 227 | |
| 228 | audio_config = self.tokenizer.instruct.audio_encoder.audio_config |
| 229 | self.n_delay_tokens = audio_config.get_num_delay_tokens() |
| 230 | |
| 231 | # for realtime transcription |
| 232 | @classmethod |
| 233 | async def buffer_realtime_audio( |
| 234 | cls, |
| 235 | audio_stream: AsyncGenerator[np.ndarray, None], |
| 236 | input_stream: asyncio.Queue[list[int]], |
| 237 | model_config: ModelConfig, |
| 238 | ) -> AsyncGenerator[PromptType, None]: |
| 239 | tokenizer = cached_tokenizer_from_config(model_config) |
| 240 | audio_encoder = tokenizer.instruct.audio_encoder |
| 241 | config = audio_encoder.audio_config |
| 242 | |
| 243 | # Get prompt tokens (streaming prefix tokens) without encoding audio |
| 244 | prompt_tokens = ( |
| 245 | tokenizer.instruct.start() + audio_encoder.encode_streaming_tokens() |
| 246 | ) |
| 247 | |
| 248 | # Get left/right padding audio |
| 249 | left_pad, right_pad = audio_encoder.get_padding_audio() |
| 250 | |
| 251 | buffer = VoxtralRealtimeBuffer(config, prompt_tokens) |
| 252 | |
| 253 | # Feed audio with padding into buffer in background |
| 254 | async def feed_audio(): |
| 255 | yielded_first_chunk = False |
| 256 | async for audio_chunk in audio_stream: |
| 257 | if not yielded_first_chunk: |
| 258 | yielded_first_chunk = True |
| 259 | # Prepend left padding before first real audio |
| 260 | await buffer.append_audio(left_pad.audio_array) |
| 261 | await buffer.append_audio(audio_chunk) |
| 262 | # Append right padding at the end |
| 263 | await buffer.append_audio(right_pad.audio_array) |
| 264 | await buffer.append_audio(None) # signal end |
| 265 | |
| 266 | # Feed output tokens back into the buffer. Idle waits are normal here; |
| 267 | # request cleanup still cancels this task through the finally block. |
| 268 | async def feed_tokens(): |
| 269 | while True: |
| 270 | all_outputs = await input_stream.get() |
| 271 | await buffer.append_tokens(all_outputs[-1:]) |
| 272 | |
| 273 | audio_task = asyncio.create_task(feed_audio()) |
| 274 | token_task = asyncio.create_task(feed_tokens()) |
| 275 | |
| 276 | try: |
| 277 | async for prompt in buffer.get_input_stream(): |
| 278 | yield prompt |
| 279 | finally: |
| 280 | audio_task.cancel() |
| 281 | token_task.cancel() |
| 282 | |
| 283 | @property |
| 284 | def audio_config(self): |
| 285 | return self.tokenizer.instruct.audio_encoder.audio_config |
| 286 | |
| 287 | def embed_input_ids( |
| 288 | self, |
| 289 | input_ids: torch.Tensor, |
| 290 | multimodal_embeddings: MultiModalEmbeddings | None = None, |
| 291 | *, |
| 292 | is_multimodal: torch.Tensor | None = None, |
| 293 | # Multi-modal token ID may exceed vocab size |
| 294 | ) -> torch.Tensor: |
| 295 | """Pass post-conv embeddings directly as input. |
| 296 | |
| 297 | For realtime models, multimodal embeddings are required at every |
| 298 | decode step. If they are missing (e.g. due to an empty audio |
| 299 | commit, encoder-cache eviction under GPU memory pressure, or a |
| 300 | client disconnect), return zero embeddings instead of crashing |
| 301 | the engine so that all other in-flight requests stay alive. |
| 302 | """ |
| 303 | if multimodal_embeddings is None or len(multimodal_embeddings) == 0: |
| 304 | logger.warning( |
| 305 | "Realtime model received empty multimodal embeddings " |
| 306 | "for %d input tokens. Returning zero embeddings to " |
| 307 | "avoid engine crash.", |
| 308 | input_ids.shape[0], |
| 309 | ) |
| 310 | pool_size = self.config.audio_config.block_pool_size |
| 311 | embed_dim = self.config.audio_config.d_model * pool_size |
| 312 | return torch.zeros( |
| 313 | input_ids.shape[0], |
| 314 | embed_dim, |
| 315 | dtype=self.whisper_encoder.dtype, |
| 316 | device=input_ids.device, |
| 317 | ) |
| 318 | mm_embeds_flat = _flatten_embeddings(multimodal_embeddings) |
| 319 | return mm_embeds_flat |
| 320 | |
| 321 | def forward( |
| 322 | self, |
| 323 | input_ids: torch.Tensor | None, |
| 324 | positions: torch.Tensor, |
| 325 | intermediate_tensors: IntermediateTensors | None = None, |
| 326 | inputs_embeds: torch.Tensor | None = None, |
| 327 | **kwargs: object, |
| 328 | ) -> torch.Tensor | IntermediateTensors: |
| 329 | assert inputs_embeds is not None |
| 330 | assert input_ids is not None |
| 331 | |
| 332 | pool_size = self.config.audio_config.block_pool_size |
| 333 | if is_torch_equal_or_newer("2.11"): |
| 334 | inputs_embeds = inputs_embeds.view( |
| 335 | inputs_embeds.shape[0] * pool_size, inputs_embeds.shape[1] // pool_size |
| 336 | ) |
| 337 | else: |
| 338 | # TODO Use reshape + clone to break the view chain and avoid output |
| 339 | # aliasing input bug in torch.compile's AOT autograd cache. |
| 340 | # Without clone(), if any downstream operation returns a view that's |
| 341 | # connected to this view of inputs_embeds, the AOT autograd cache |
| 342 | # fails to pickle the ViewMetaSequence containing SymInt shapes. |
| 343 | # This will be fixed in pytorch 2.11 and beyond. |
| 344 | # issue: https://github.com/pytorch/pytorch/issues/174299 |
| 345 | inputs_embeds = inputs_embeds.reshape( |
| 346 | inputs_embeds.shape[0] * pool_size, inputs_embeds.shape[1] // pool_size |
| 347 | ).clone() |
| 348 | |
| 349 | whisper_positions = _expand_tensor(positions, pool_size) |
| 350 | audio_hidden_states = self.whisper_encoder.whisper_encoder( |
| 351 | inputs_embeds, whisper_positions |
| 352 | ) |
| 353 | |
| 354 | num_tokens, audio_hidden_size = audio_hidden_states.shape |
| 355 | assert num_tokens % self.downsample_factor == 0 |
| 356 | audio_hidden_states = audio_hidden_states.reshape( |
| 357 | num_tokens // self.downsample_factor, |
| 358 | audio_hidden_size * self.downsample_factor, |
| 359 | ) |
| 360 | audio_text_embeds = self.audio_language_adapter(audio_hidden_states) |
| 361 | |
| 362 | text_embeds = self.language_model.embed_input_ids(input_ids) |
| 363 | |
| 364 | # sum pool text and audio embeddings |
| 365 | inputs_embeds = audio_text_embeds + text_embeds |
| 366 | |
| 367 | time_tensor = torch.full( |
| 368 | (1,), |
| 369 | fill_value=self.n_delay_tokens, |
| 370 | device=inputs_embeds.device, |
| 371 | dtype=inputs_embeds.dtype, |
| 372 | ) |
| 373 | t_cond = self.time_embedding(time_tensor) |
| 374 | |
| 375 | hidden_states = self.language_model.model( |
| 376 | input_ids, |
| 377 | positions, |
| 378 | intermediate_tensors, |
| 379 | inputs_embeds=inputs_embeds, |
| 380 | t_cond=t_cond, |
| 381 | ) |
| 382 | |
| 383 | return hidden_states |
| 384 | |
| 385 | def embed_multimodal( |
| 386 | self, **kwargs |
| 387 | ) -> list[torch.Tensor] | torch.Tensor | tuple[torch.Tensor, ...] | None: |
| 388 | """Transform audio waveforms -> initial whisper post-conv embeddings.""" |
| 389 | audio_inputs = self._parse_and_validate_audio_arrays(**kwargs) |
| 390 | |
| 391 | if audio_inputs is None: |
| 392 | logger.warning( |
| 393 | "Realtime model received no audio inputs in " |
| 394 | "embed_multimodal. Returning empty embeddings." |
| 395 | ) |
| 396 | return [] |
| 397 | |
| 398 | def _truncate_left( |
| 399 | sample: torch.Tensor, mult_of: int, pos: int |
| 400 | ) -> torch.Tensor: |
| 401 | assert pos in [0, 1], pos |
| 402 | if (ctx := sample.shape[pos] % mult_of) != 0: |
| 403 | sample = sample[ctx:] if pos == 0 else sample[:, ctx:] |
| 404 | assert sample.shape[pos] > 0, ( |
| 405 | f"Sample is empty after truncation with ctx {ctx}" |
| 406 | ) |
| 407 | |
| 408 | return sample |
| 409 | |
| 410 | mel_features = [ |
| 411 | self.whisper_encoder.compute_whisper_melspec(audio).to( |
| 412 | self.whisper_encoder.dtype |
| 413 | ) |
| 414 | for audio in audio_inputs |
| 415 | ] |
| 416 | |
| 417 | # we truncate the left most mel feature |
| 418 | # if the sequence length in impair |
| 419 | mel_features = [_truncate_left(mel, 2, 1) for mel in mel_features] |
| 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 | ) |
| 427 | conv_stride = self.whisper_encoder.whisper_encoder.total_stride |
| 428 | audio_embeddings_per_sample = audio_embeddings.split( |
| 429 | [s // conv_stride for s in seq_lens], dim=0 |
| 430 | ) |
| 431 | |
| 432 | # audio_embeddings per sample need to be divisible by 4 |
| 433 | pool_size = self.config.audio_config.block_pool_size |
| 434 | |
| 435 | audio_embeddings_per_sample = [ |
| 436 | _truncate_left(sample, pool_size, 0) |
| 437 | for sample in audio_embeddings_per_sample |
| 438 | ] |
| 439 | |
| 440 | audio_embeddings_per_sample = [ |
| 441 | e.view(e.shape[0] // pool_size, e.shape[1] * pool_size) |
| 442 | for e in audio_embeddings_per_sample |
| 443 | ] |
| 444 | return audio_embeddings_per_sample |
| 445 | |
| 446 | @classmethod |
| 447 | def get_speech_to_text_config( |
| 448 | cls, model_config: ModelConfig, task_type: str |
| 449 | ) -> SpeechToTextConfig: |
| 450 | tokenizer = cached_tokenizer_from_config(model_config) |
| 451 | audio_config = tokenizer.instruct.audio_encoder.audio_config |
| 452 | sample_rate = audio_config.sampling_rate |
| 453 | return SpeechToTextConfig( |
| 454 | max_audio_clip_s=None, # only limited by memory |
| 455 | sample_rate=sample_rate, |
| 456 | min_energy_split_window_size=None, |
| 457 | ) |
| 458 | |
| 459 | @classmethod |
| 460 | # for speech-to-text transcription |
| 461 | def get_generation_prompt( |
| 462 | cls, |
| 463 | stt_params: SpeechToTextParams, |
| 464 | ) -> PromptType: |
| 465 | audio = stt_params.audio |
| 466 | model_config = stt_params.model_config |
| 467 | stt_config = stt_params.stt_config |
| 468 | language = stt_params.language |
| 469 | |
| 470 | tokenizer = cached_tokenizer_from_config(model_config) |
| 471 | audio = Audio(audio, int(stt_config.sample_rate), format="wav") # lossless |
| 472 | |
| 473 | req = TranscriptionRequest( |
| 474 | model=model_config.model, |
| 475 | audio=audio.to_base64(audio.format), |
| 476 | language=language, |
| 477 | streaming=StreamingMode.OFFLINE, |
| 478 | ) |
| 479 | |
| 480 | tokenized = tokenizer.instruct.encode_transcription(req) |
| 481 | |
| 482 | return TokensPrompt( |
| 483 | prompt_token_ids=tokenized.tokens, |
| 484 | multi_modal_data={ |
| 485 | "audio": (tokenized.audios[0].audio_array, stt_config.sample_rate) |
| 486 | }, |
| 487 | ) |