@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3faa8bee57
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 11, 00:00 UTC

adapt voxtral (#31095)

Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com>

Patrick von Platen committed Dec 23, 2025 at 14:31 UTC 3faa8bee57982d08cae5adf733f92e27708f15ce
12 files changed +740 -99
tests/models/multimodal/generation/test_voxtral.py
+1
@@ -111,4 +111,5 @@ async def test_online_serving(client, audio_assets: AudioTestAssets):
111
112 assert len(chat_completion.choices) == 1
113 choice = chat_completion.choices[0]
114 + assert choice.message.content == "In the first audio clip, you hear a brief"
115 assert choice.finish_reason == "length"
tests/models/registry.py
+5
@@ -860,6 +860,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
860 # disable this temporarily until we support HF format
861 is_available_online=False,
862 ),
863 + "VoxtralStreamingGeneration": _HfExamplesInfo(
864 + "<place-holder>",
865 + # disable this temporarily until we support HF format
866 + is_available_online=False,
867 + ),
868 # [Encoder-decoder]
869 "WhisperForConditionalGeneration": _HfExamplesInfo(
870 "openai/whisper-large-v3-turbo",
vllm/config/model.py
+4
@@ -1542,6 +1542,10 @@ class ModelConfig:
1542 def is_multimodal_raw_input_only_model(self) -> bool:
1543 return self._model_info.supports_multimodal_raw_input_only
1544
1545 + @property
1546 + def requires_raw_input_tokens(self) -> bool:
1547 + return self._model_info.requires_raw_input_tokens
1548 +
1549 @property
1550 def is_cross_encoder(self) -> bool:
1551 return (
vllm/model_executor/models/interfaces.py
+10
@@ -94,6 +94,12 @@ class SupportsMultiModal(Protocol):
94 `multimodal_config.mm_encoder_tp_mode="data"`.
95 """
96
97 + requires_raw_input_tokens: ClassVar[bool] = False
98 + """
99 + A flag that indicates this model processes input id tokens
100 + in their raw form and not input embeddings.
101 + """
102 +
103 merge_by_field_config: ClassVar[bool | None] = None
104 """
105 [DEPRECATED] A flag that indicates which implementation of
@@ -306,6 +312,10 @@ def supports_multimodal_raw_input_only(model: type[object] | object) -> bool:
312 return getattr(model, "supports_multimodal_raw_input_only", False)
313
314
315 +def requires_raw_input_tokens(model: type[object] | object) -> bool:
316 + return getattr(model, "requires_raw_input_tokens", False)
317 +
318 +
319 def supports_multimodal_encoder_tp_data(model: type[object] | object) -> bool:
320 return getattr(model, "supports_encoder_tp_data", False)
321
vllm/model_executor/models/registry.py
+4
@@ -46,6 +46,7 @@ from .interfaces import (
46 has_noops,
47 is_attention_free,
48 is_hybrid,
49 + requires_raw_input_tokens,
50 supports_cross_encoding,
51 supports_mamba_prefix_caching,
52 supports_multimodal,
@@ -422,6 +423,7 @@ _MULTIMODAL_MODELS = {
423 ),
424 "UltravoxModel": ("ultravox", "UltravoxModel"),
425 "VoxtralForConditionalGeneration": ("voxtral", "VoxtralForConditionalGeneration"), # noqa: E501
426 + "VoxtralStreamingGeneration": ("voxtral_streaming", "VoxtralStreamingGeneration"), # noqa: E501
427 # [Encoder-decoder]
428 "WhisperForConditionalGeneration": ("whisper", "WhisperForConditionalGeneration"), # noqa: E501
429 }
@@ -539,6 +541,7 @@ class _ModelInfo:
541 supports_cross_encoding: bool
542 supports_multimodal: bool
543 supports_multimodal_raw_input_only: bool
544 + requires_raw_input_tokens: bool
545 supports_multimodal_encoder_tp_data: bool
546 supports_pp: bool
547 has_inner_state: bool
@@ -562,6 +565,7 @@ class _ModelInfo:
565 supports_multimodal_raw_input_only=supports_multimodal_raw_input_only(
566 model
567 ),
568 + requires_raw_input_tokens=requires_raw_input_tokens(model),
569 supports_multimodal_encoder_tp_data=supports_multimodal_encoder_tp_data(
570 model
571 ),
vllm/model_executor/models/voxtral.py
+29 -11
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import inspect
5 import math
6 from collections.abc import Iterable, Mapping, Sequence
7 from functools import cached_property
@@ -116,10 +117,7 @@ class VoxtralProcessorAdapter:
117 self,
118 audio_length: int,
119 ) -> int:
119 - pad_audio_length = self._audio_processor.next_multiple_of_chunk_frames(
120 - audio_length, self.sampling_rate
121 - )
122 - return ceil(pad_audio_length / (self.sampling_rate // self.frame_rate))
120 + return ceil(audio_length / (self.sampling_rate // self.frame_rate))
121
122 def __call__(
123 self,
@@ -158,7 +156,14 @@ class VoxtralProcessorAdapter:
156 assert audio.ndim == 1
157
158 # pad if necessary
161 - audio = self._audio_processor.pad(audio, self.sampling_rate)
159 + # TODO(Patrick) - remove once mistral-common is bumped
160 + sig = inspect.signature(self._audio_processor.pad)
161 + if "is_online_streaming" in sig.parameters:
162 + audio = self._audio_processor.pad(
163 + audio, self.sampling_rate, is_online_streaming=False
164 + )
165 + else:
166 + audio = self._audio_processor.pad(audio, self.sampling_rate)
167
168 audio_tokens = [self.begin_audio_token_id] + [
169 self.audio_token_id
@@ -510,6 +515,7 @@ class VoxtralForConditionalGeneration(
515
516 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
517 remapping_rules = [
518 + (r"mm_streams_embeddings.embedding_module\.(.*)", r"\1"),
519 (r"mm_whisper_embeddings\.(.*)", r"\1"),
520 (r"audio_language_projection\.(.*)", r"audio_language_adapter.\1"),
521 (
@@ -535,13 +541,16 @@ class VoxtralForConditionalGeneration(
541 def llm_weights_generator():
542 nonlocal loaded_weights
543 for name, w in weights:
538 - is_encoder = (
539 - name.startswith("mm_whisper_embeddings")
540 - and not name.startswith("mm_whisper_embeddings.tok_embeddings")
541 - and not name.startswith(
542 - "mm_whisper_embeddings.audio_language_projection"
544 + is_encoder = False
545 + for k in [
546 + "mm_whisper_embeddings",
547 + "mm_streams_embeddings.embedding_module",
548 + ]:
549 + is_encoder |= (
550 + name.startswith(k)
551 + and not name.startswith(f"{k}.tok_embeddings")
552 + and not name.startswith(f"{k}.audio_language_projection")
553 )
544 - )
554
555 for pattern, repl in remapping_rules:
556 if re.fullmatch(pattern, name):
@@ -676,6 +685,7 @@ class VoxtralEncoderModel(nn.Module):
685 packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
686
687 mistral_remapping = [
688 + (r"mm_streams_embeddings.embedding_module\.(.*)", r"\1"),
689 (
690 r"whisper_encoder\.conv_layers\.0\.(weight|bias)",
691 r"whisper_encoder.conv1.\1",
@@ -684,6 +694,14 @@ class VoxtralEncoderModel(nn.Module):
694 r"whisper_encoder\.conv_layers\.1\.(weight|bias)",
695 r"whisper_encoder.conv2.\1",
696 ),
697 + (
698 + r"whisper_encoder\.conv_layers\.0\.conv\.(weight|bias)",
699 + r"whisper_encoder.conv1.\1",
700 + ), # noqa: E501
701 + (
702 + r"whisper_encoder\.conv_layers\.1\.conv\.(weight|bias)",
703 + r"whisper_encoder.conv2.\1",
704 + ), # noqa: E501
705 (
706 r"whisper_encoder\.transformer\.layers\.(\d+)\.attention\.w([qkv])\.(weight|bias)", # noqa: E501
707 r"whisper_encoder.layers.\1.self_attn.\2_proj.\3",
vllm/model_executor/models/voxtral_streaming.py new
+243
@@ -0,0 +1,243 @@
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 Mapping
6 +
7 +import torch
8 +
9 +from vllm.config.vllm import VllmConfig
10 +from vllm.logger import init_logger
11 +from vllm.model_executor.models.interfaces import MultiModalEmbeddings
12 +from vllm.model_executor.models.voxtral import (
13 + VoxtralDummyInputsBuilder,
14 + VoxtralForConditionalGeneration,
15 + VoxtralMultiModalProcessor,
16 + VoxtralProcessingInfo,
17 +)
18 +from vllm.multimodal import MULTIMODAL_REGISTRY
19 +from vllm.multimodal.cache import _I, BaseMultiModalProcessorCache
20 +from vllm.multimodal.inputs import (
21 + MultiModalKwargsOptionalItems,
22 +)
23 +from vllm.multimodal.parse import MultiModalDataItems
24 +from vllm.multimodal.processing import (
25 + MultiModalPromptUpdates,
26 + PlaceholderFeaturesInfo,
27 +)
28 +from vllm.multimodal.profiling import BaseDummyInputsBuilder
29 +from vllm.sequence import IntermediateTensors
30 +
31 +from .utils import (
32 + _flatten_embeddings,
33 +)
34 +
35 +logger = init_logger(__name__)
36 +
37 +
38 +class VoxtralStreamingMultiModalProcessor(VoxtralMultiModalProcessor):
39 + def __init__(
40 + self,
41 + info: _I,
42 + dummy_inputs: BaseDummyInputsBuilder[_I],
43 + *,
44 + cache: BaseMultiModalProcessorCache | None = None,
45 + ) -> None:
46 + # streaming can't make use of a cache yet
47 + super().__init__(info, dummy_inputs, cache=None)
48 +
49 + def _maybe_apply_prompt_updates(
50 + self,
51 + mm_items: MultiModalDataItems,
52 + prompt_ids: list[int],
53 + mm_kwargs: MultiModalKwargsOptionalItems,
54 + mm_prompt_updates: MultiModalPromptUpdates,
55 + is_update_applied: bool,
56 + ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
57 + # there are no placeholder audio tokens for streaming
58 + # so we need to build the place placeholder positions manually
59 +
60 + # in streaming there is always only one audio input
61 + audios = mm_kwargs.get("audio", [])
62 + assert len(audios) == 1, (
63 + f"Expected only one audio input for streaming, got {mm_kwargs=}"
64 + )
65 + tokenizer = self.info.get_tokenizer()
66 + audio_config = tokenizer.instruct.audio_encoder.audio_config
67 +
68 + num_audio_samples = audios[0]["audio_arrays"].data.shape[0]
69 + length = audio_config.num_audio_tokens(num_audio_samples)
70 +
71 + features_info = PlaceholderFeaturesInfo(
72 + modality="audio",
73 + item_idx=0,
74 + start_idx=0,
75 + tokens=length
76 + * [0], # only used for length computation, so we can take dummy inputs
77 + is_embed=None,
78 + )
79 + return prompt_ids, {"audio": [features_info]}
80 +
81 +
82 +class TimeEmbedding(torch.nn.Module):
83 + """Sinusoidal Embedding for encoding time"""
84 +
85 + def __init__(self, dim: int, theta: float = 10000.0) -> None:
86 + super().__init__()
87 + self.dim = dim
88 + self.theta = theta
89 + inv_freq = torch.exp(
90 + -math.log(self.theta)
91 + * torch.arange(self.dim // 2).float()
92 + / (self.dim // 2)
93 + )
94 + self.register_buffer("inv_freq", inv_freq, persistent=False)
95 +
96 + def forward(self, t: torch.Tensor) -> torch.Tensor:
97 + t = t[..., None] # (B,) -> (B, 1) or (B, T) -> (B, T, 1)
98 + inv_freq = self.inv_freq.to(device=t.device, dtype=t.dtype)
99 + emb = (
100 + t * inv_freq
101 + ) # (B, 1) x (D/2,) -> (B, D/2) or (B, T, 1) x (D/2,) -> (B, T, D/2)
102 + return torch.cat((emb.cos(), emb.sin()), dim=-1) # (B, D) or (B, T, D)
103 +
104 +
105 +@MULTIMODAL_REGISTRY.register_processor(
106 + VoxtralStreamingMultiModalProcessor,
107 + info=VoxtralProcessingInfo,
108 + dummy_inputs=VoxtralDummyInputsBuilder,
109 +)
110 +class VoxtralStreamingGeneration(VoxtralForConditionalGeneration):
111 + requires_raw_input_tokens = True
112 +
113 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
114 + super().__init__(vllm_config=vllm_config, prefix=prefix)
115 + self.time_embedding: TimeEmbedding = TimeEmbedding(
116 + dim=self.config.text_config.hidden_size
117 + )
118 +
119 + audio_config = self.tokenizer.instruct.audio_encoder.audio_config
120 + _n_delay_tokens = (
121 + audio_config.frame_rate * audio_config.transcription_delay_ms / 1000
122 + )
123 + assert _n_delay_tokens.is_integer(), (
124 + f"n_delay_tokens must be integer, got {_n_delay_tokens}"
125 + )
126 +
127 + self.n_delay_tokens = int(_n_delay_tokens)
128 +
129 + @property
130 + def audio_config(self):
131 + return self.tokenizer.instruct.audio_encoder.audio_config
132 +
133 + def embed_input_ids(
134 + self,
135 + input_ids: torch.Tensor,
136 + multimodal_embeddings: MultiModalEmbeddings | None = None,
137 + *,
138 + is_multimodal: torch.Tensor | None = None,
139 + # Multi-modal token ID may exceed vocab size
140 + handle_oov_mm_token: bool = True,
141 + ) -> torch.Tensor:
142 + """Pass post-conv embeddings directly as input"""
143 + # for streaming we simply flatten the multimodal embeddings
144 + # to be in tensor format, we treat the input ids later
145 + assert multimodal_embeddings is not None
146 + assert len(multimodal_embeddings) > 0, (
147 + "For streaming you must provide a multimodal_embedding at every step."
148 + )
149 + mm_embeds_flat = _flatten_embeddings(multimodal_embeddings)
150 + return mm_embeds_flat
151 +
152 + def forward(
153 + self,
154 + input_ids: torch.Tensor,
155 + positions: torch.Tensor,
156 + intermediate_tensors: IntermediateTensors | None = None,
157 + inputs_embeds: torch.Tensor | None = None,
158 + **kwargs: object,
159 + ) -> torch.Tensor | IntermediateTensors:
160 + assert inputs_embeds is not None
161 + assert input_ids is not None
162 +
163 + pool_size = self.config.audio_config.block_pool_size
164 + inputs_embeds = inputs_embeds.view(
165 + inputs_embeds.shape[0] * pool_size, inputs_embeds.shape[1] // pool_size
166 + )
167 +
168 + audio_hidden_states = self.whisper_encoder.whisper_encoder.forward_layers(
169 + inputs_embeds
170 + )
171 +
172 + num_tokens, audio_hidden_size = audio_hidden_states.shape
173 + assert num_tokens % self.downsample_factor == 0
174 + audio_hidden_states = audio_hidden_states.reshape(
175 + num_tokens // self.downsample_factor,
176 + audio_hidden_size * self.downsample_factor,
177 + )
178 + audio_text_embeds = self.audio_language_adapter(audio_hidden_states)
179 +
180 + text_embeds = self.language_model.embed_input_ids(input_ids)
181 +
182 + # sum pool text and audio embeddings
183 + inputs_embeds = audio_text_embeds + text_embeds
184 +
185 + time_tensor = torch.tensor(
186 + [self.n_delay_tokens],
187 + device=inputs_embeds.device,
188 + dtype=inputs_embeds.dtype,
189 + )
190 + inputs_embeds = inputs_embeds + self.time_embedding(time_tensor)
191 +
192 + hidden_states = self.language_model.model(
193 + input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
194 + )
195 +
196 + return hidden_states
197 +
198 + def embed_multimodal(
199 + self, **kwargs
200 + ) -> list[torch.Tensor] | torch.Tensor | tuple[torch.Tensor, ...] | None:
201 + """Transform audio waveforms -> initial whisper post-conv embeddings"""
202 + audio_inputs = self._parse_and_validate_audio_arrays(**kwargs)
203 +
204 + assert audio_inputs is not None, (
205 + "For streaming you must provide an audio input at every step."
206 + )
207 +
208 + multiple_of = self.audio_config.raw_audio_length_per_tok
209 + assert all(
210 + (this_audio := audio.shape[0]) % multiple_of == 0 for audio in audio_inputs
211 + ), (
212 + f"Every input audio waveform has to be a multiple of {multiple_of}, but"
213 + f" one is {this_audio} with {(this_audio / multiple_of)=}."
214 + )
215 +
216 + mel_features = [
217 + self.whisper_encoder.compute_whisper_melspec(audio).to(
218 + self.whisper_encoder.dtype
219 + )
220 + for audio in audio_inputs
221 + ]
222 + seq_lens = [mel.shape[1] for mel in mel_features]
223 + # [total_num_20ms_frames, hidden_size]
224 + audio_embeddings = self.whisper_encoder.whisper_encoder.forward_conv(
225 + mel_features
226 + )[0]
227 + conv_stride = self.whisper_encoder.whisper_encoder.total_stride
228 + audio_embeddings_per_sample = audio_embeddings.split(
229 + [s // conv_stride for s in seq_lens], dim=0
230 + )
231 +
232 + # audio_embeddings per sample need to be divisible by 4
233 + pool_size = self.config.audio_config.block_pool_size
234 + assert all(
235 + (this_shape := sample.shape[0]) % pool_size == 0
236 + for sample in audio_embeddings_per_sample
237 + ), f"Every audio embedding has to be a multiple of 4, but one is {this_shape}."
238 +
239 + audio_embeddings_per_sample = [
240 + e.view(e.shape[0] // pool_size, e.shape[1] * pool_size)
241 + for e in audio_embeddings_per_sample
242 + ]
243 + return audio_embeddings_per_sample
vllm/model_executor/models/whisper.py
+91 -81
@@ -1,9 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import enum
5 import math
6 from collections.abc import Iterable, Mapping, Sequence
7 from contextlib import nullcontext
8 +from functools import partial
9 from typing import Annotated, Literal, cast
10
11 import numpy as np
@@ -16,7 +18,10 @@ from transformers import (
18 )
19 from transformers.models.whisper.modeling_whisper import sinusoids
20
19 -from vllm.attention.layer import Attention, AttentionType
21 +from vllm.attention.backends.abstract import (
22 + AttentionType,
23 +)
24 +from vllm.attention.layer import Attention
25 from vllm.attention.layers.cross_attention import CrossAttention
26 from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
27 from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig
@@ -34,6 +39,11 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
39 from vllm.model_executor.layers.quantization import QuantizationConfig
40 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
41 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
42 +from vllm.model_executor.models.whisper_utils import (
43 + ISO639_1_SUPPORTED_LANGS,
44 + WhisperAttentionWithBlockPooling,
45 + WhisperCausalConv1d,
46 +)
47 from vllm.multimodal import MULTIMODAL_REGISTRY
48 from vllm.multimodal.inputs import (
49 MultiModalDataDict,
@@ -64,67 +74,11 @@ from .utils import (
74
75 logger = init_logger(__name__)
76
67 -# From https://platform.openai.com/docs/guides/speech-to-text/supported-languages
68 -
69 -ISO639_1_SUPPORTED_LANGS = {
70 - "af": "Afrikaans",
71 - "ar": "Arabic",
72 - "hy": "Armenian",
73 - "az": "Azerbaijani",
74 - "be": "Belarusian",
75 - "bs": "Bosnian",
76 - "bg": "Bulgarian",
77 - "ca": "Catalan",
78 - "zh": "Chinese",
79 - "hr": "Croatian",
80 - "cs": "Czech",
81 - "da": "Danish",
82 - "nl": "Dutch",
83 - "en": "English",
84 - "et": "Estonian",
85 - "fi": "Finnish",
86 - "fr": "French",
87 - "gl": "Galician",
88 - "de": "German",
89 - "el": "Greek",
90 - "he": "Hebrew",
91 - "hi": "Hindi",
92 - "hu": "Hungarian",
93 - "is": "Icelandic",
94 - "id": "Indonesian",
95 - "it": "Italian",
96 - "ja": "Japanese",
97 - "kn": "Kannada",
98 - "kk": "Kazakh",
99 - "ko": "Korean",
100 - "lv": "Latvian",
101 - "lt": "Lithuanian",
102 - "mk": "Macedonian",
103 - "ms": "Malay",
104 - "mr": "Marathi",
105 - "mi": "Maori",
106 - "ne": "Nepali",
107 - "no": "Norwegian",
108 - "fa": "Persian",
109 - "pl": "Polish",
110 - "pt": "Portuguese",
111 - "ro": "Romanian",
112 - "ru": "Russian",
113 - "sr": "Serbian",
114 - "sk": "Slovak",
115 - "sl": "Slovenian",
116 - "es": "Spanish",
117 - "sw": "Swahili",
118 - "sv": "Swedish",
119 - "tl": "Tagalog",
120 - "ta": "Tamil",
121 - "th": "Thai",
122 - "tr": "Turkish",
123 - "uk": "Ukrainian",
124 - "ur": "Urdu",
125 - "vi": "Vietnamese",
126 - "cy": "Welsh",
127 -}
77 +
78 +class WhisperPosEmbedType(enum.Enum):
79 + SINUSOIDAL = "sinusoidal"
80 + NOPE = "nope"
81 + LEARNED = "learned"
82
83
84 class WhisperAudioInputs(TensorSchema):
@@ -184,6 +138,8 @@ class WhisperAttention(nn.Module):
138 num_heads: int,
139 bias: bool = True,
140 attn_type: AttentionType = AttentionType.DECODER,
141 + per_layer_sliding_window: int | None = None,
142 + block_pool_size: int = 1,
143 cache_config: CacheConfig | None = None,
144 quant_config: QuantizationConfig | None = None,
145 prefix: str = "",
@@ -242,7 +198,14 @@ class WhisperAttention(nn.Module):
198 attn_type=self.attn_type,
199 )
200 else: # AttentionType.DECODER (regular decoder self-attention)
245 - self.attn = Attention(
201 + if block_pool_size > 1:
202 + attn_cls = partial(
203 + WhisperAttentionWithBlockPooling, block_pool_size=block_pool_size
204 + )
205 + else:
206 + attn_cls = Attention
207 +
208 + self.attn = attn_cls(
209 self.num_heads,
210 self.head_dim,
211 self.scaling,
@@ -251,6 +214,7 @@ class WhisperAttention(nn.Module):
214 quant_config=quant_config,
215 prefix=f"{prefix}.attn",
216 attn_type=self.attn_type,
217 + per_layer_sliding_window=per_layer_sliding_window,
218 )
219
220 def _init_qkv(
@@ -386,6 +350,9 @@ class WhisperEncoderLayer(nn.Module):
350 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
351 super().__init__()
352 config = vllm_config.model_config.hf_config
353 + is_causal = getattr(config, "is_causal", False)
354 + sliding_window = getattr(config, "sliding_window", None)
355 + block_pool_size = getattr(config, "block_pool_size", 1)
356 cache_config = vllm_config.cache_config
357 quant_config = vllm_config.quant_config
358
@@ -393,7 +360,9 @@ class WhisperEncoderLayer(nn.Module):
360 self.self_attn = WhisperAttention(
361 embed_dim=self.embed_dim,
362 num_heads=config.encoder_attention_heads,
396 - attn_type=AttentionType.ENCODER,
363 + attn_type=AttentionType.DECODER if is_causal else AttentionType.ENCODER,
364 + block_pool_size=block_pool_size,
365 + per_layer_sliding_window=sliding_window,
366 cache_config=cache_config,
367 quant_config=quant_config,
368 prefix=f"{prefix}.self_attn",
@@ -492,12 +461,21 @@ class WhisperEncoder(nn.Module):
461 super().__init__()
462 config = vllm_config.model_config.hf_config
463 embed_dim = config.d_model
464 +
465 + self.pos_embed_type = WhisperPosEmbedType(
466 + getattr(config, "pos_embed", "sinusoidal")
467 + )
468 self.num_mel_bins = config.num_mel_bins
469 self.max_source_positions = config.max_source_positions
470 self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
471
499 - self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1)
500 - self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1)
472 + is_causal = getattr(config, "is_causal", False)
473 + Conv1d = WhisperCausalConv1d if is_causal else partial(nn.Conv1d, padding=1)
474 +
475 + self.conv1 = Conv1d(self.num_mel_bins, embed_dim, kernel_size=3)
476 + self.conv2 = Conv1d(embed_dim, embed_dim, stride=2, kernel_size=3)
477 +
478 + self.total_stride = self.conv1.stride[0] * self.conv2.stride[0]
479 self.start_layer, self.end_layer, self.layers = make_layers(
480 config.encoder_layers,
481 lambda prefix: WhisperEncoderLayer(
@@ -507,29 +485,54 @@ class WhisperEncoder(nn.Module):
485 )
486 self.layer_norm = nn.LayerNorm(config.d_model)
487
510 - maybe_fp32_init_ctx = (
511 - set_default_torch_dtype(torch.float32) if init_in_fp32 else nullcontext()
512 - )
513 -
514 - with (
515 - torch.no_grad(),
516 - maybe_fp32_init_ctx,
488 + if is_causal and self.pos_embed_type != WhisperPosEmbedType.NOPE:
489 + raise ValueError(
490 + "Only NOPE position embeddings are supported "
491 + f"for causal models, but got {self.pos_embed_type}"
492 + )
493 + elif self.pos_embed_type in (
494 + WhisperPosEmbedType.SINUSOIDAL,
495 + WhisperPosEmbedType.LEARNED,
496 ):
518 - self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim)
519 - self.embed_positions.weight.copy_(
520 - sinusoids(*self.embed_positions.weight.shape)
497 + maybe_fp32_init_ctx = (
498 + set_default_torch_dtype(torch.float32)
499 + if init_in_fp32
500 + else nullcontext()
501 )
502
523 - def forward(self, input_features: torch.Tensor | list[torch.Tensor]):
503 + with (
504 + torch.no_grad(),
505 + maybe_fp32_init_ctx,
506 + ):
507 + self.embed_positions = nn.Embedding(
508 + self.max_source_positions, embed_dim
509 + )
510 + self.embed_positions.weight.copy_(
511 + sinusoids(*self.embed_positions.weight.shape)
512 + )
513 +
514 + def forward_conv(
515 + self, input_features: torch.Tensor | list[torch.Tensor]
516 + ) -> torch.Tensor:
517 hidden_states = []
518 input_is_batched = False
519 for features in input_features:
520 embeds = nn.functional.gelu(self.conv1(features))
521 embeds = nn.functional.gelu(self.conv2(embeds))
529 - embeds = embeds.transpose(-1, -2)
530 - embeds = (embeds + self.embed_positions.weight[: embeds.size(-2), :]).to(
531 - embeds.dtype
532 - )
522 +
523 + if self.pos_embed_type in (
524 + WhisperPosEmbedType.SINUSOIDAL,
525 + WhisperPosEmbedType.LEARNED,
526 + ):
527 + embeds = embeds.transpose(-1, -2)
528 + embeds = (
529 + embeds + self.embed_positions.weight[: embeds.size(-2), :]
530 + ).to(embeds.dtype)
531 + elif self.pos_embed_type == WhisperPosEmbedType.NOPE:
532 + embeds = embeds.transpose(-1, -2).to(embeds.dtype)
533 + else:
534 + raise ValueError(f"Unknown pos_embed_type: {self.pos_embed_type}")
535 +
536 hidden_states.append(embeds)
537 input_is_batched = embeds.ndim > 2
538 # Input to MHA must be B x T x D
@@ -539,12 +542,19 @@ class WhisperEncoder(nn.Module):
542 else:
543 hidden_states = torch.stack(hidden_states, dim=0)
544
545 + return hidden_states
546 +
547 + def forward_layers(self, hidden_states: torch.Tensor) -> torch.Tensor:
548 for encoder_layer in self.layers:
549 hidden_states = encoder_layer(hidden_states)
550
551 hidden_states = self.layer_norm(hidden_states)
552 return hidden_states
553
554 + def forward(self, input_features: torch.Tensor | list[torch.Tensor]):
555 + hidden_states = self.forward_conv(input_features)
556 + return self.forward_layers(hidden_states)
557 +
558
559 class WhisperDecoder(nn.Module):
560 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
vllm/model_executor/models/whisper_utils.py new
+299
@@ -0,0 +1,299 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import copy
4 +import functools
5 +import math
6 +from dataclasses import replace
7 +
8 +import torch
9 +import torch.nn.functional as F
10 +from torch import nn
11 +
12 +from vllm.attention.backends.abstract import (
13 + AttentionBackend,
14 + AttentionMetadata,
15 + AttentionType,
16 +)
17 +from vllm.attention.layer import Attention
18 +from vllm.attention.selector import get_attn_backend
19 +from vllm.config import CacheConfig, VllmConfig
20 +from vllm.model_executor.layers.quantization import QuantizationConfig
21 +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
22 +from vllm.v1.attention.backends.utils import (
23 + CommonAttentionMetadata,
24 + subclass_attention_backend_with_overrides,
25 +)
26 +from vllm.v1.kv_cache_interface import AttentionSpec
27 +
28 +# From https://platform.openai.com/docs/guides/speech-to-text/supported-languages
29 +ISO639_1_SUPPORTED_LANGS = {
30 + "af": "Afrikaans",
31 + "ar": "Arabic",
32 + "hy": "Armenian",
33 + "az": "Azerbaijani",
34 + "be": "Belarusian",
35 + "bs": "Bosnian",
36 + "bg": "Bulgarian",
37 + "ca": "Catalan",
38 + "zh": "Chinese",
39 + "hr": "Croatian",
40 + "cs": "Czech",
41 + "da": "Danish",
42 + "nl": "Dutch",
43 + "en": "English",
44 + "et": "Estonian",
45 + "fi": "Finnish",
46 + "fr": "French",
47 + "gl": "Galician",
48 + "de": "German",
49 + "el": "Greek",
50 + "he": "Hebrew",
51 + "hi": "Hindi",
52 + "hu": "Hungarian",
53 + "is": "Icelandic",
54 + "id": "Indonesian",
55 + "it": "Italian",
56 + "ja": "Japanese",
57 + "kn": "Kannada",
58 + "kk": "Kazakh",
59 + "ko": "Korean",
60 + "lv": "Latvian",
61 + "lt": "Lithuanian",
62 + "mk": "Macedonian",
63 + "ms": "Malay",
64 + "mr": "Marathi",
65 + "mi": "Maori",
66 + "ne": "Nepali",
67 + "no": "Norwegian",
68 + "fa": "Persian",
69 + "pl": "Polish",
70 + "pt": "Portuguese",
71 + "ro": "Romanian",
72 + "ru": "Russian",
73 + "sr": "Serbian",
74 + "sk": "Slovak",
75 + "sl": "Slovenian",
76 + "es": "Spanish",
77 + "sw": "Swahili",
78 + "sv": "Swedish",
79 + "tl": "Tagalog",
80 + "ta": "Tamil",
81 + "th": "Thai",
82 + "tr": "Turkish",
83 + "uk": "Ukrainian",
84 + "ur": "Urdu",
85 + "vi": "Vietnamese",
86 + "cy": "Welsh",
87 +}
88 +
89 +
90 +def _pad1d(
91 + x: torch.Tensor,
92 + paddings: tuple[int, int],
93 + mode: str = "constant",
94 + value: float = 0.0,
95 +) -> torch.Tensor:
96 + """Tiny wrapper around F.pad, just to allow for
97 + reflect padding on small input.
98 + If this is the case, we insert extra 0 padding
99 + to the right before the reflection happen.
100 + """
101 + length = x.shape[-1]
102 + padding_left, padding_right = paddings
103 + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
104 + if mode == "reflect":
105 + max_pad = max(padding_left, padding_right)
106 + extra_pad = 0
107 + if length <= max_pad:
108 + extra_pad = max_pad - length + 1
109 + x = F.pad(x, (0, extra_pad))
110 + padded = F.pad(x, paddings, mode, value)
111 + end = padded.shape[-1] - extra_pad
112 + return padded[..., :end]
113 + else:
114 + return F.pad(x, paddings, mode, value)
115 +
116 +
117 +class WhisperCausalConv1d(nn.Conv1d):
118 + def __init__(
119 + self,
120 + in_channels: int,
121 + out_channels: int,
122 + kernel_size: int,
123 + stride: int = 1,
124 + padding: int = 0,
125 + bias: bool = True,
126 + ) -> None:
127 + super().__init__(
128 + in_channels,
129 + out_channels,
130 + kernel_size,
131 + stride=stride,
132 + padding=padding,
133 + bias=bias,
134 + )
135 + self._stride = self.stride[0]
136 + self._effective_kernel_size = (kernel_size - 1) * self.dilation[0] + 1
137 + self._padding_total = self._effective_kernel_size - self._stride
138 +
139 + def forward(self, x: torch.Tensor) -> torch.Tensor:
140 + n_frames = (
141 + x.shape[-1] - self._effective_kernel_size + self._padding_total
142 + ) / self._stride + 1
143 + target_length = (math.ceil(n_frames) - 1) * self._stride + (
144 + self._effective_kernel_size - self._padding_total
145 + )
146 + extra_padding = target_length - x.shape[-1]
147 + x = _pad1d(x, (self._padding_total, extra_padding), mode="constant")
148 + return super().forward(x)
149 +
150 +
151 +@functools.lru_cache
152 +def create_whisper_attention_backend_with_block_pooling(
153 + underlying_attn_backend: AttentionBackend, block_pool_size: int
154 +) -> type[AttentionBackend]:
155 + prefix = "WhisperAttentionWithBlockPooling_"
156 + underlying_builder = underlying_attn_backend.get_builder_cls()
157 +
158 + class WhisperAttentionWithBlockPoolingBuilder(underlying_builder): # type: ignore
159 + def __init__(
160 + self,
161 + kv_cache_spec: AttentionSpec,
162 + layer_names: list[str],
163 + vllm_config: VllmConfig,
164 + device: torch.device,
165 + ):
166 + assert kv_cache_spec.num_kv_heads % block_pool_size == 0
167 + kv_cache_spec = replace(
168 + kv_cache_spec,
169 + block_size=kv_cache_spec.block_size * block_pool_size,
170 + num_kv_heads=kv_cache_spec.num_kv_heads // block_pool_size,
171 + )
172 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
173 +
174 + def build(
175 + self,
176 + common_prefix_len: int,
177 + common_attn_metadata: CommonAttentionMetadata,
178 + fast_build: bool = False,
179 + ) -> AttentionMetadata:
180 + new_common_attn_metadata = copy.deepcopy(common_attn_metadata)
181 + new_common_attn_metadata.query_start_loc *= block_pool_size
182 + new_common_attn_metadata.query_start_loc_cpu *= block_pool_size
183 + new_common_attn_metadata.seq_lens *= block_pool_size
184 + new_common_attn_metadata._seq_lens_cpu *= block_pool_size
185 + new_common_attn_metadata._num_computed_tokens_cpu *= block_pool_size
186 + new_common_attn_metadata.num_actual_tokens *= block_pool_size
187 + new_common_attn_metadata.max_query_len *= block_pool_size
188 + new_common_attn_metadata.max_seq_len *= block_pool_size
189 + original_slot_mapping = common_attn_metadata.slot_mapping
190 + common_prefix_len *= block_pool_size
191 + new_common_attn_metadata.slot_mapping = (
192 + (
193 + original_slot_mapping.unsqueeze(1) * block_pool_size
194 + + torch.arange(block_pool_size, device=original_slot_mapping.device)
195 + )
196 + .flatten()
197 + .clamp(min=-1)
198 + )
199 + return super().build(
200 + common_prefix_len, new_common_attn_metadata, fast_build
201 + )
202 +
203 + if not issubclass(underlying_attn_backend, FlashAttentionBackend):
204 + raise NotImplementedError(
205 + f"{underlying_attn_backend} is not yet supported."
206 + "Contributions to support more backends are much "
207 + "appreciated."
208 + )
209 +
210 + attn_backend = subclass_attention_backend_with_overrides(
211 + name_prefix=prefix,
212 + attention_backend_cls=underlying_attn_backend,
213 + overrides={
214 + "get_builder_cls": lambda: WhisperAttentionWithBlockPoolingBuilder,
215 + "get_kv_cache_shape": lambda num_blocks,
216 + block_size,
217 + num_kv_heads,
218 + head_size,
219 + cache_dtype_str: (
220 + 2,
221 + num_blocks,
222 + # we stretch each block by `block_pool_size`
223 + block_size * block_pool_size,
224 + num_kv_heads // block_pool_size,
225 + head_size,
226 + ), # TODO: generalize to other backends
227 + },
228 + )
229 +
230 + return attn_backend
231 +
232 +
233 +class WhisperAttentionWithBlockPooling(Attention):
234 + """Attention layer with block pooling."""
235 +
236 + def __init__(
237 + self,
238 + num_heads: int,
239 + head_size: int,
240 + scale: float,
241 + num_kv_heads: int | None = None,
242 + alibi_slopes: list[float] | None = None,
243 + cache_config: CacheConfig | None = None,
244 + quant_config: QuantizationConfig | None = None,
245 + logits_soft_cap: float | None = None,
246 + per_layer_sliding_window: int | None = None,
247 + prefix: str = "",
248 + attn_type: str = AttentionType.DECODER,
249 + kv_sharing_target_layer_name: str | None = None,
250 + block_pool_size: int = 1,
251 + attn_backend: type[AttentionBackend] | None = None,
252 + **extra_impl_args,
253 + ) -> None:
254 + self.block_pool_size = block_pool_size
255 + dtype = torch.get_default_dtype()
256 +
257 + if cache_config is not None:
258 + kv_cache_dtype = cache_config.cache_dtype
259 + block_size = cache_config.block_size
260 + else:
261 + kv_cache_dtype = "auto"
262 + block_size = 16
263 +
264 + underlying_attn_backend = get_attn_backend(
265 + head_size,
266 + dtype,
267 + kv_cache_dtype,
268 + block_size,
269 + attn_type=attn_type,
270 + )
271 + attn_backend = create_whisper_attention_backend_with_block_pooling(
272 + underlying_attn_backend, block_pool_size
273 + )
274 +
275 + super().__init__(
276 + num_heads=num_heads,
277 + head_size=head_size,
278 + scale=scale,
279 + num_kv_heads=num_kv_heads,
280 + alibi_slopes=alibi_slopes,
281 + cache_config=cache_config,
282 + quant_config=quant_config,
283 + logits_soft_cap=logits_soft_cap,
284 + per_layer_sliding_window=per_layer_sliding_window,
285 + prefix=prefix,
286 + attn_type=attn_type,
287 + kv_sharing_target_layer_name=kv_sharing_target_layer_name,
288 + attn_backend=attn_backend,
289 + **extra_impl_args,
290 + )
291 +
292 + def get_kv_cache_spec(self, vllm_config: VllmConfig):
293 + kv_cache_spec = super().get_kv_cache_spec(vllm_config)
294 + assert isinstance(kv_cache_spec, AttentionSpec)
295 + kv_cache_spec = replace(
296 + kv_cache_spec,
297 + num_kv_heads=self.block_pool_size * kv_cache_spec.num_kv_heads,
298 + )
299 + return kv_cache_spec
vllm/transformers_utils/configs/mistral.py
+31 -3
@@ -184,18 +184,42 @@ def _remap_mistral_audio_args(config: dict) -> dict:
184 whisper_args = config["multimodal"].pop("whisper_model_args")
185 encoder_args = whisper_args["encoder_args"]
186 downsample_args = whisper_args["downsample_args"]
187 + downsample_factor = downsample_args["downsample_factor"]
188 +
189 + # make sure that k/v blocks can be allocated with
190 + # unified k/v cache class and pool whisper k/v cache blocks
191 + # with downsample_factor:1 ratio
192 + if encoder_args.get("causal"):
193 + block_pool_size = downsample_factor
194 + config["projection_size"] = downsample_factor * encoder_args["dim"]
195 + else:
196 + block_pool_size = 1
197 +
198 + _maybe_sliding_window = encoder_args.get("ragged_attention", None)
199 + if _maybe_sliding_window is None:
200 + sliding_window = None
201 + elif _maybe_sliding_window.isdigit():
202 + sliding_window = int(_maybe_sliding_window)
203 + else:
204 + raise NotImplementedError(f"Unsupported: {_maybe_sliding_window=}")
205 +
206 + architecture = (
207 + "VoxtralStreamingGeneration"
208 + if encoder_args.get("causal")
209 + else "VoxtralForConditionalGeneration"
210 + )
211
212 quant_config = config.get("quantization_config")
213 config = {
190 - "model_type": "whixtral",
191 - "architectures": ["VoxtralForConditionalGeneration"],
214 + "model_type": "voxtral",
215 + "architectures": [architecture],
216 "text_config": PretrainedConfig.from_dict(config),
217 "audio_config": WhisperConfig(
218 num_mel_bins=encoder_args["audio_encoding_args"]["num_mel_bins"],
219 window_size=encoder_args["audio_encoding_args"]["window_size"],
220 sampling_rate=encoder_args["audio_encoding_args"]["sampling_rate"],
221 hop_length=encoder_args["audio_encoding_args"]["hop_length"],
198 - downsample_factor=downsample_args["downsample_factor"],
222 + downsample_factor=downsample_factor,
223 d_model=encoder_args["dim"],
224 encoder_layers=encoder_args["n_layers"],
225 encoder_ffn_dim=encoder_args["hidden_dim"],
@@ -203,6 +227,10 @@ def _remap_mistral_audio_args(config: dict) -> dict:
227 vocab_size=encoder_args["vocab_size"],
228 max_source_positions=encoder_args["max_source_positions"],
229 is_encoder_decoder=False, # Override WhisperConfig default
230 + is_causal=encoder_args.get("causal", False),
231 + sliding_window=sliding_window,
232 + block_pool_size=block_pool_size,
233 + pos_embed=encoder_args.get("pos_embed", "sinusoidal"),
234 ),
235 }
236 if quant_config:
vllm/v1/attention/backends/utils.py
+9
@@ -835,6 +835,15 @@ def subclass_attention_backend(
835 )
836
837
838 +def subclass_attention_backend_with_overrides(
839 + name_prefix: str,
840 + attention_backend_cls: type[AttentionBackend],
841 + overrides: dict[str, Any],
842 +) -> type[AttentionBackend]:
843 + name: str = name_prefix + attention_backend_cls.__name__ # type: ignore
844 + return type(name, (attention_backend_cls,), overrides)
845 +
846 +
847 def split_decodes_prefills_and_extends(
848 common_attn_metadata: CommonAttentionMetadata,
849 decode_threshold: int = 1,
vllm/v1/worker/gpu_model_runner.py
+14 -4
@@ -2457,6 +2457,17 @@ class GPUModelRunner(
2457 return round_up(num_scheduled_tokens, tp_size)
2458 return num_scheduled_tokens
2459
2460 + def _prepare_mm_inputs(
2461 + self, num_tokens: int
2462 + ) -> tuple[torch.Tensor | None, torch.Tensor]:
2463 + if self.model.requires_raw_input_tokens:
2464 + input_ids = self.input_ids.gpu[:num_tokens]
2465 + else:
2466 + input_ids = None
2467 +
2468 + inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
2469 + return input_ids, inputs_embeds
2470 +
2471 def _preprocess(
2472 self,
2473 scheduler_output: "SchedulerOutput",
@@ -2499,8 +2510,7 @@ class GPUModelRunner(
2510 # TODO(woosuk): Avoid the copy. Optimize.
2511 self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
2512
2502 - input_ids = None
2503 - inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens]
2513 + input_ids, inputs_embeds = self._prepare_mm_inputs(num_input_tokens)
2514 model_kwargs = {
2515 **self._init_model_kwargs(num_scheduled_tokens),
2516 **self._extract_mm_kwargs(scheduler_output),
@@ -4220,8 +4230,8 @@ class GPUModelRunner(
4230 assert num_tokens_padded <= self.max_num_tokens
4231 model_kwargs = self._init_model_kwargs(num_tokens_padded)
4232 if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
4223 - input_ids = None
4224 - inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]
4233 + input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded)
4234 +
4235 model_kwargs = {
4236 **model_kwargs,
4237 **self._dummy_mm_kwargs(num_reqs),