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, 23:45 UTC
Migrate GraniteSpeechAudioInputs to TensorSchema (#21682)
Signed-off-by: Benji Beck <benjibeck@meta.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Benji Beck committed
Jul 27, 2025 at 22:37 UTC
75856bc2cb8e559dc5cd62cc0e495662be1aaae3
1 file changed
+18
-8
vllm/model_executor/models/granite_speech.py
+18
-8
@@ -25,7 +25,7 @@
25
"""Inference-only IBM Granite speech model."""
26
import math
27
from collections.abc import Iterable, Mapping
28
-from typing import Optional, TypedDict, Union
28
+from typing import Annotated, Optional, Union
29
30
import torch
31
import torch.nn.functional as F
@@ -48,6 +48,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
48
PromptUpdate)
49
from vllm.multimodal.profiling import BaseDummyInputsBuilder
50
from vllm.sequence import IntermediateTensors
51
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
52
53
from .blip2 import Blip2QFormerModel
54
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
@@ -57,16 +58,24 @@ from .utils import (AutoWeightsLoader, embed_multimodal,
58
59
60
### Audio Input
60
-class GraniteSpeechAudioInputs(TypedDict):
61
+class GraniteSpeechAudioInputs(TensorSchema):
62
+ """
63
+ Audio input features for Granite Speech model.
64
+
65
+ Dimensions:
66
+ - b: Batch size
67
+ - nf: Number of audio features (variable length)
68
+ - 160: Fixed feature dimension for Mel spectrogram features
69
+ """
70
62
- input_features: torch.Tensor
63
- """Shape: `(bsz, num_features, 160)`"""
71
+ input_features: Annotated[torch.Tensor, TensorShape("b", "nf", 160)]
72
+ """Audio input features."""
73
65
- input_features_mask: torch.Tensor
66
- """Shape: `(bsz, num_features)`"""
74
+ input_features_mask: Annotated[torch.Tensor, TensorShape("b", "nf")]
75
+ """Mask for variable length audio features."""
76
68
- audio_embed_sizes: list[int]
69
- """List of length `bsz`"""
77
+ audio_embed_sizes: Annotated[list[int], TensorShape("b")]
78
+ """List of audio embedding sizes for each item in batch."""
79
80
81
class GraniteSpeechMultiModalProcessingInfo(BaseProcessingInfo):
@@ -581,6 +590,7 @@ class GraniteSpeechForConditionalGeneration(
590
input_features = kwargs.pop("input_features", None)
591
input_features_mask = kwargs.pop("input_features_mask", None)
592
audio_embed_sizes = kwargs.pop("audio_embed_sizes", None)
593
+
594
if input_features is None:
595
return None
596