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, 22:30 UTC
[Model] Add support for transformer-based Ultravox v0.7 projector (#30089)
Signed-off-by: Peter Salas <peter@fixie.ai>
Peter Salas committed
Dec 5, 2025 at 20:55 UTC
e858bc4d1490808e23bbc32e17202ebbf8713a76
2 files changed
+91
-7
vllm/model_executor/models/ultravox.py
+89
-7
@@ -4,15 +4,21 @@
4
# Adapted from https://github.com/fixie-ai/ultravox/blob/ecd58c4041030bae2ad15aa6bcf04ab43199ea02/ultravox/model/ultravox_model.py
5
"""PyTorch Ultravox model."""
6
7
+import copy
8
from collections.abc import Iterable, Mapping, Sequence
9
+from types import SimpleNamespace
10
from typing import Annotated, Any, Literal, TypeAlias
11
12
import torch
13
from torch import nn
14
from torch.nn import functional as F
15
from transformers import BatchFeature, ProcessorMixin
16
+from transformers.modeling_utils import ModuleUtilsMixin
17
from transformers.models.whisper import WhisperFeatureExtractor
15
-from transformers.models.whisper.modeling_whisper import WhisperEncoder
18
+from transformers.models.whisper.modeling_whisper import (
19
+ WhisperEncoder,
20
+ WhisperEncoderLayer,
21
+)
22
23
from vllm.config import VllmConfig
24
from vllm.config.multimodal import BaseDummyOptions
@@ -282,7 +288,7 @@ class StackAudioFrames(nn.Module):
288
return audio_embeds
289
290
285
-class UltravoxProjector(nn.Module):
291
+class UltravoxFeedForwardProjector(nn.Module):
292
def __init__(self, config: UltravoxConfig):
293
super().__init__()
294
self.hidden_dim = config.hidden_size
@@ -310,7 +316,9 @@ class UltravoxProjector(nn.Module):
316
self.ln_mid = nn.Identity()
317
self.ln_post = RMSNorm(dim_out)
318
313
- def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
319
+ def forward(
320
+ self, audio_features: torch.Tensor, audio_token_len: torch.Tensor
321
+ ) -> torch.Tensor:
322
audio_features = self._pad_and_stack(audio_features)
323
audio_features = self.ln_pre(audio_features)
324
hidden_states = self.linear_1(audio_features)
@@ -321,6 +329,70 @@ class UltravoxProjector(nn.Module):
329
return hidden_states
330
331
332
+class UltravoxTransformerProjector(nn.Module, ModuleUtilsMixin):
333
+ def __init__(self, config: UltravoxConfig):
334
+ super().__init__()
335
+ self.config = SimpleNamespace(is_decoder=False)
336
+
337
+ self._pad_and_stack = StackAudioFrames(config.stack_factor)
338
+ dim_in = config.audio_config.hidden_size * config.stack_factor
339
+
340
+ projector_audio_config = copy.deepcopy(config.audio_config)
341
+
342
+ self.ln_pre = RMSNorm(dim_in)
343
+ self.linear_in = nn.Linear(dim_in, projector_audio_config.d_model)
344
+
345
+ self.embed_positions = nn.Embedding(
346
+ projector_audio_config.max_source_positions,
347
+ projector_audio_config.d_model,
348
+ )
349
+
350
+ self.layers = nn.ModuleList(
351
+ [
352
+ WhisperEncoderLayer(projector_audio_config)
353
+ for _ in range(config.num_projector_layers)
354
+ ]
355
+ )
356
+
357
+ self.ln_post = RMSNorm(projector_audio_config.d_model)
358
+ self.linear_out = nn.Linear(
359
+ projector_audio_config.d_model, config.text_config.hidden_size
360
+ )
361
+
362
+ def forward(
363
+ self, audio_features: torch.Tensor, audio_token_len: torch.Tensor
364
+ ) -> torch.Tensor:
365
+ audio_features = self._pad_and_stack(audio_features)
366
+
367
+ max_len_stacked = audio_features.shape[1]
368
+ attention_mask = torch.arange(max_len_stacked, device=audio_features.device)[
369
+ None, :
370
+ ].lt(audio_token_len[:, None])
371
+ extended_attention_mask = self.get_extended_attention_mask(
372
+ attention_mask, attention_mask.shape, audio_features.dtype
373
+ )
374
+
375
+ hidden_states = self.ln_pre(audio_features)
376
+ hidden_states = self.linear_in(hidden_states)
377
+
378
+ positions = self.embed_positions(
379
+ torch.arange(hidden_states.size(1), device=hidden_states.device)
380
+ )
381
+ hidden_states = hidden_states + positions
382
+
383
+ for layer in self.layers:
384
+ layer_outputs = layer(
385
+ hidden_states,
386
+ attention_mask=extended_attention_mask,
387
+ layer_head_mask=None,
388
+ )
389
+ hidden_states = layer_outputs[0]
390
+
391
+ hidden_states = self.ln_post(hidden_states)
392
+ hidden_states = self.linear_out(hidden_states)
393
+ return hidden_states
394
+
395
+
396
class ModifiedWhisperEncoder(WhisperEncoder):
397
"""
398
Encoder portion of OpenAI's Whisper model.
@@ -464,7 +536,10 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
536
prefix="audio_tower.",
537
)
538
)
467
- self.multi_modal_projector = UltravoxProjector(config)
539
+ if config.num_projector_layers > 0:
540
+ self.multi_modal_projector = UltravoxTransformerProjector(config)
541
+ else:
542
+ self.multi_modal_projector = UltravoxFeedForwardProjector(config)
543
self.language_model = init_vllm_registered_model(
544
vllm_config=vllm_config,
545
hf_config=config.wrapped_model_config,
@@ -496,7 +571,10 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
571
)
572
573
def _audio_features_to_embeddings(
499
- self, input_features: torch.Tensor, audio_lens: torch.Tensor
574
+ self,
575
+ input_features: torch.Tensor,
576
+ audio_lens: torch.Tensor,
577
+ audio_token_len: torch.Tensor,
578
) -> torch.Tensor:
579
audio_features = input_features.to(self.audio_tower.dtype)
580
batch_size = audio_features.size(0)
@@ -512,7 +590,9 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
590
batch_features = batch_features.to(self.audio_tower.dtype)
591
592
# Process through projector
515
- batch_embeddings = self.multi_modal_projector(batch_features)
593
+ batch_embeddings = self.multi_modal_projector(
594
+ batch_features, audio_token_len[start:end]
595
+ )
596
audio_embeddings.append(batch_embeddings)
597
598
# Concatenate results
@@ -559,7 +639,9 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
639
audio_lens = audio_input["lens"]
640
audio_token_len = audio_input["token_len"]
641
562
- embeddings = self._audio_features_to_embeddings(audio_features, audio_lens)
642
+ embeddings = self._audio_features_to_embeddings(
643
+ audio_features, audio_lens, audio_token_len
644
+ )
645
646
# We should flatten and concatenate embeddings based on token lengths
647
# For example, with token_len = [4, 2, 3], flattened_embeddings will be
vllm/transformers_utils/configs/ultravox.py
+2
@@ -61,6 +61,7 @@ class UltravoxConfig(transformers.PretrainedConfig):
61
norm_init: float = 0.4,
62
projector_act: str = "swiglu",
63
projector_ln_mid: bool = False,
64
+ num_projector_layers: int = 0,
65
**kwargs,
66
):
67
self.ignore_index = ignore_index
@@ -71,6 +72,7 @@ class UltravoxConfig(transformers.PretrainedConfig):
72
self.norm_init = norm_init
73
self.projector_act = projector_act
74
self.projector_ln_mid = projector_ln_mid
75
+ self.num_projector_layers = num_projector_layers
76
77
# N.B. May set the wrapped_model_config below.
78
self.text_model_id = text_model_id