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 10, 10:15 UTC
[Bugfix] Fix Var Length Batched Padding in Granite Speech (#31906)
Signed-off-by: Alex-Brooks <Alex.Brooks@ibm.com>
Alex Brooks committed
Jan 9, 2026 at 03:28 UTC
dc77cb7129b9efa1c74a2a05f6359dde977f7c8b
1 file changed
+8
-3
vllm/model_executor/models/granite_speech.py
+8
-3
@@ -672,7 +672,13 @@ class GraniteSpeechForConditionalGeneration(
672
673
else:
674
# Otherwise we have a list of tensors, which are almost certainly
675
- # differing in their respective numbers of audio features;
675
+ # differing in their respective numbers of audio features; when
676
+ # passed as a batch, we expect a list of 2D var len input features
677
+ # so unsqueeze them.
678
+ input_features = [
679
+ feat.unsqueeze(dim=0) for feat in input_features if feat.ndim == 2
680
+ ]
681
+
682
# stack them into a 3D tensor of size [bsz, most_num_features, 160].
683
input_features = self._pad_and_stack_input_features(
684
input_features,
@@ -724,13 +730,12 @@ class GraniteSpeechForConditionalGeneration(
730
731
Args:
732
input_features: list[torch.Tensor]
727
- Input features to be coerced into a tensor.
733
+ 3D Input features to be coerced into a tensor.
734
Returns:
735
torch.Tensor: Tensor of shape [bsz, num_features, 160], where
736
num_features is the max number of features of any entry in the
737
batch.
738
"""
733
- # Input features are of shape [bsz, num_features, 160]
739
feat_lens = [feats.shape[1] for feats in input_features]
740
padding = [max(feat_lens) - length for length in feat_lens]
741
# TODO (Alex) - Validate that it's okay to zero pad like this;