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 MiniCPMOAudioInputs to TensorSchema (#21847)
Signed-off-by: Benji Beck <benjibeck@meta.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Benji Beck committed
Aug 22, 2025 at 01:43 UTC
998720859caadd8a8d2a3e2af8b3e6e34a42e8da
1 file changed
+35
-17
vllm/model_executor/models/minicpmo.py
+35
-17
@@ -24,7 +24,7 @@
24
# limitations under the License.
25
"""Inference-only MiniCPM-O model compatible with HuggingFace weights."""
26
from collections.abc import Iterable, Mapping, Sequence
27
-from typing import Any, Callable, Literal, Optional, TypedDict, Union
27
+from typing import Annotated, Any, Callable, Literal, Optional, Union
28
29
import torch
30
from torch import nn
@@ -49,6 +49,7 @@ from vllm.multimodal.parse import (AudioItem, AudioProcessorItems,
49
MultiModalDataParser)
50
from vllm.multimodal.processing import (PromptReplacement, PromptUpdate,
51
PromptUpdateDetails)
52
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
53
54
from .minicpmv import (_MAX_FRAMES_PER_VIDEO, MiniCPMV2_6,
55
MiniCPMVDummyInputsBuilder,
@@ -61,35 +62,52 @@ from .utils import (AutoWeightsLoader, cast_overflow_tensors, flatten_bn,
62
CPU_DEVICE = torch.device("cpu")
63
64
64
-class MiniCPMOAudioFeatureInputs(TypedDict):
65
- type: Literal["audio_features"]
66
- audio_features: Union[torch.Tensor, list[torch.Tensor]]
65
+class MiniCPMOAudioFeatureInputs(TensorSchema):
66
"""
68
- Shape: `(batch_size * num_audios * num_slices, num_channels, length)`
69
- Slice here means chunk. Audio that is too long will be split into slices,
70
- which is the same as image.
71
- Padding is used therefore `audio_features` is `torch.Tensor`.
67
+ Dimensions:
68
+ - bns: Batch size * number of audios * number of slices
69
+ - bn: Batch size * number of audios
70
+ - c: Number of channels
71
+ - l: Length
72
+ - s: Number of slices
73
"""
74
+ type: Literal["audio_features"] = "audio_features"
75
74
- audio_feature_lens: Union[torch.Tensor, list[torch.Tensor]]
76
+ audio_features: Annotated[
77
+ Union[torch.Tensor, list[torch.Tensor]],
78
+ TensorShape("bns", "c", "l", dynamic_dims={"l"}),
79
+ ]
80
+ """
81
+ Slice here means chunk. Audio that is too long will be split into slices,
82
+ which is the same as image. Padding is used therefore `audio_features` is
83
+ `torch.Tensor`.
84
"""
76
- Shape: `(batch_size * num_audios, num_slices)`
85
86
+ audio_feature_lens: Annotated[
87
+ Union[torch.Tensor, list[torch.Tensor]],
88
+ TensorShape("bn", "s"),
89
+ ]
90
+ """
91
This should be feature length of each audio slice,
92
which equals to `audio_features.shape[-1]`
93
"""
94
95
83
-class MiniCPMOAudioEmbeddingInputs(TypedDict):
84
- type: Literal["audio_embeds"]
85
- audio_embeds: Union[torch.Tensor, list[torch.Tensor]]
96
+class MiniCPMOAudioEmbeddingInputs(TensorSchema):
97
"""
87
- Shape: `(batch_size * num_audios, num_slices, hidden_size)`
88
-
89
- `hidden_size` must match the hidden size of language model backbone.
90
- instead of a batched tensor.
98
+ Dimensions:
99
+ - bn: Batch size * number of audios
100
+ - s: Number of slices
101
+ - h: Hidden size (must match language model backbone)
102
+
103
Length of each slice may vary, so pass it as a list.
104
"""
105
+ type: Literal["audio_embeds"] = "audio_embeds"
106
+
107
+ audio_embeds: Annotated[
108
+ Union[torch.Tensor, list[torch.Tensor]],
109
+ TensorShape("bn", "s", "h", dynamic_dims={"s"}),
110
+ ]
111
112
113
MiniCPMOAudioInputs = Union[MiniCPMOAudioFeatureInputs,