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 14, 05:30 UTC
Migrate ultravox inputs to TensorSchema (#23503)
Signed-off-by: Benji Beck <benjibeck@meta.com>
Benji Beck committed
Sep 3, 2025 at 23:09 UTC
cb55ad86fe47759c22de06b3b3b93b3aa28690a2
1 file changed
+28
-30
vllm/model_executor/models/ultravox.py
+28
-30
@@ -4,7 +4,7 @@
4
# Adapted from https://github.com/fixie-ai/ultravox/blob/ecd58c4041030bae2ad15aa6bcf04ab43199ea02/ultravox/model/ultravox_model.py
5
"""PyTorch Ultravox model."""
6
from collections.abc import Iterable, Mapping, Sequence
7
-from typing import Any, Literal, Optional, TypedDict, Union
7
+from typing import Annotated, Any, Literal, Optional, Union
8
9
import torch
10
from torch import nn
@@ -31,6 +31,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
31
from vllm.multimodal.profiling import BaseDummyInputsBuilder
32
from vllm.sequence import IntermediateTensors
33
from vllm.transformers_utils.configs.ultravox import UltravoxConfig
34
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
35
36
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
37
SupportsMultiModal, SupportsPP)
@@ -43,26 +44,37 @@ _AUDIO_PLACEHOLDER_OVERRIDE = "<|audio|>"
44
_MAX_ENCODER_BATCH_SIZE = 16
45
46
46
-class UltravoxAudioFeatureInputs(TypedDict):
47
- type: Literal["audio_features"]
48
- data: Union[torch.Tensor, list[torch.Tensor], list[list[torch.Tensor]]]
49
- """Shape: `(batch_size, num_chunks, 80, M)`"""
50
- lens: Union[torch.Tensor, list[torch.Tensor]]
47
+class UltravoxAudioFeatureInputs(TensorSchema):
48
"""
52
- Length of the audio frames. Used for attention mask in WhisperEncoder.
53
- Shape: `(batch_size, num_chunks)`
49
+ Dimensions:
50
+ - b: batch size
51
+ - n: number of chunks
52
+ - t: Time frames (M)
53
+ - nmb: Number of mel bins
54
"""
55
- token_len: Union[torch.Tensor, list[torch.Tensor]]
55
+ type: Literal["audio_features"]
56
+ data: Annotated[Union[torch.Tensor, list[torch.Tensor],
57
+ list[list[torch.Tensor]]],
58
+ TensorShape("b", "n", "nmb", "t", dynamic_dims={"n"})]
59
+ lens: Annotated[Union[torch.Tensor, list[torch.Tensor]],
60
+ TensorShape("b", "n", dynamic_dims={"n"})]
61
+ """Length of the audio frames. Used for attention mask in WhisperEncoder."""
62
+ token_len: Annotated[Union[torch.Tensor, list[torch.Tensor]],
63
+ TensorShape("b", "n", dynamic_dims={"n"})]
64
+ """Length of the audio tokens. Used for flattening the audio features."""
65
+
66
+
67
+class UltravoxAudioEmbeddingInputs(TensorSchema):
68
"""
57
- Length of the audio tokens. Used for flattening the audio features.
58
- Shape: `(batch_size, num_chunks)`
69
+ Dimensions:
70
+ - b: batch size
71
+ - na: number of audios
72
+ - afs: audio feature size
73
+ - hs: hidden size
74
"""
60
-
61
-
62
-class UltravoxAudioEmbeddingInputs(TypedDict):
75
type: Literal["audio_embeds"]
64
- data: NestedTensors
65
- """Shape: `(batch_size, num_audios, audio_feature_size, hidden_size)`"""
76
+ data: Annotated[Union[torch.Tensor, list[torch.Tensor]],
77
+ TensorShape("b", "na", "afs", "hs")]
78
79
80
UltravoxAudioInputs = Union[UltravoxAudioFeatureInputs,
@@ -484,26 +496,12 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
496
return None
497
498
if audio_features is not None:
487
- if not isinstance(audio_features, (torch.Tensor, list)):
488
- raise ValueError("Incorrect type of audio features. "
489
- f"Got type: {type(audio_features)}")
490
- if not isinstance(audio_lens, (torch.Tensor, list)):
491
- raise ValueError("Incorrect type of audio_lens. "
492
- f"Got type: {type(audio_features)}")
493
- if not isinstance(audio_token_len, (torch.Tensor, list)):
494
- raise ValueError("Incorrect type of audio_token_len. "
495
- f"Got type: {type(audio_features)}")
496
-
499
return UltravoxAudioFeatureInputs(type="audio_features",
500
data=audio_features,
501
lens=audio_lens,
502
token_len=audio_token_len)
503
504
if audio_embeds is not None:
503
- if not isinstance(audio_embeds, (torch.Tensor, list)):
504
- raise ValueError("Incorrect type of audio embeds. "
505
- f"Got type: {type(audio_embeds)}")
506
-
505
return UltravoxAudioEmbeddingInputs(type="audio_embeds",
506
data=audio_embeds)
507