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 6, 12:00 UTC
chunk parakeet into 30s clips to prevent OOMs on long audios (#36671)
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>
Netanel Haber committed
Mar 18, 2026 at 23:22 UTC
6ae4c8d6fc0483ab736243045e529aa397693d4b
3 files changed
+73
-65
vllm/model_executor/models/nano_nemotron_vl.py
+29
-35
@@ -99,15 +99,16 @@ MAX_AUDIO_LEN_S = 10 * 60 # 10 minutes
99
class NanoNemotronVLAudioFeatureInputs(TensorSchema):
100
"""
101
Dimensions:
102
- - b: Number of audio clips
102
+ - c: Number of audio clips (possibly flattened across audio items)
103
+ - b: Number of original audio items
104
- t: Audio feature length
105
- f: Feature size (mel bins)
106
"""
107
108
type: Literal["audio_features"] = "audio_features"
108
- input_audio_features: Annotated[torch.Tensor, TensorShape("b", "t", "f")]
109
- feature_attention_mask: Annotated[torch.Tensor, TensorShape("b", "t")]
110
- audio_feature_lengths: Annotated[torch.Tensor, TensorShape("b")]
109
+ input_audio_features: Annotated[torch.Tensor, TensorShape("c", "t", "f")]
110
+ feature_attention_mask: Annotated[torch.Tensor, TensorShape("c", "t")]
111
+ audio_num_clips: list[int]
112
113
114
class NanoNemotronVLImagePixelInputs(TensorSchema):
@@ -548,10 +549,17 @@ class NanoNemotronVLMultiModalProcessor(
549
video_fields = {}
550
551
if self.info.audio_extractor is not None:
552
+ audio_num_clips = torch.as_tensor(hf_inputs["audio_num_clips"])
553
audio_fields = dict(
552
- input_audio_features=MultiModalFieldConfig.batched("audio"),
553
- feature_attention_mask=MultiModalFieldConfig.batched("audio"),
554
- audio_feature_lengths=MultiModalFieldConfig.batched("audio"),
554
+ input_audio_features=MultiModalFieldConfig.flat_from_sizes(
555
+ "audio", audio_num_clips
556
+ ),
557
+ feature_attention_mask=MultiModalFieldConfig.flat_from_sizes(
558
+ "audio", audio_num_clips
559
+ ),
560
+ audio_num_clips=MultiModalFieldConfig.batched(
561
+ "audio", keep_on_cpu=True
562
+ ),
563
)
564
else:
565
audio_fields = {}
@@ -1095,28 +1103,9 @@ class NemotronH_Nano_VL_V2(
1103
assert self.sound_encoder is not None
1104
input_audio_features = audio_input.input_audio_features
1105
feature_attention_mask = audio_input.feature_attention_mask
1106
+ audio_num_clips = audio_input.audio_num_clips
1107
target_device = next(self.sound_encoder.parameters()).device
1108
1100
- # When cross-request batching combines audio clips with different
1101
- # time dimensions, _reduce_data returns a list instead of a stacked
1102
- # tensor. Pad to the max time dim and stack; the attention mask
1103
- # already marks valid positions so zero-padding is safe.
1104
- if isinstance(input_audio_features, list):
1105
- feature_sizes = [f.shape[-2] for f in input_audio_features]
1106
- max_t = max(feature_sizes)
1107
- padded_feats = [
1108
- torch.nn.functional.pad(feat, (0, 0, 0, max_t - feat_size))
1109
- for feat, feat_size in zip(
1110
- input_audio_features, feature_sizes, strict=True
1111
- )
1112
- ]
1113
- padded_masks = [
1114
- torch.nn.functional.pad(mask, (0, max_t - mask.shape[-1]))
1115
- for mask in feature_attention_mask
1116
- ]
1117
- input_audio_features = torch.stack(padded_feats)
1118
- feature_attention_mask = torch.stack(padded_masks)
1119
-
1109
input_audio_features = input_audio_features.to(
1110
dtype=self.llm_dtype, device=target_device
1111
)
@@ -1126,13 +1115,18 @@ class NemotronH_Nano_VL_V2(
1115
valid_input_lens = feature_attention_mask.sum(dim=1)
1116
valid_output_lens = self.sound_encoder.encoder._get_subsampling_output_length(
1117
valid_input_lens
1129
- )
1130
- truncated_embeds = []
1131
- for i in range(sound_embeds.shape[0]):
1132
- valid_len = valid_output_lens[i].item()
1133
- truncated_embeds.append(sound_embeds[i, :valid_len])
1134
-
1135
- return tuple(truncated_embeds)
1118
+ ).tolist()
1119
+ grouped_embeds = []
1120
+ clip_offset = 0
1121
+ for num_clips in audio_num_clips:
1122
+ embeds = []
1123
+ for clip_idx in range(clip_offset, clip_offset + num_clips):
1124
+ valid_len = valid_output_lens[clip_idx]
1125
+ embeds.append(sound_embeds[clip_idx, :valid_len])
1126
+ grouped_embeds.append(torch.cat(embeds, dim=0))
1127
+ clip_offset += num_clips
1128
+
1129
+ return tuple(grouped_embeds)
1130
1131
def _create_final_video_embeddings(
1132
self,
@@ -1246,7 +1240,7 @@ class NemotronH_Nano_VL_V2(
1240
in (
1241
"input_audio_features",
1242
"feature_attention_mask",
1249
- "audio_feature_lengths",
1243
+ "audio_num_clips",
1244
)
1245
and "audios" not in modalities
1246
):
vllm/model_executor/models/parakeet.py
+40
-23
@@ -114,33 +114,50 @@ class ParakeetExtractor(ParakeetFeatureExtractor):
114
round(self.config.clip_min_duration_s * self.sampling_rate)
115
)
116
117
- def _normalize_audio_length(self, audio_len: int) -> int:
118
- # Match mcore's compute_params() logic for clip/minduration handling.
119
- target_len = max(audio_len, self._tail_min_samples)
120
- tail_remainder = target_len % self._clip_target_samples
121
- if 0 < tail_remainder < self._tail_min_samples:
122
- padding = self._tail_min_samples - tail_remainder
123
- target_len += padding
124
- assert isinstance(target_len, int)
125
- return target_len
117
+ def _clip_sizes(self, audio_len: int) -> list[int]:
118
+ audio_len = max(audio_len, self._tail_min_samples)
119
+ num_full_clips, remainder = divmod(audio_len, self._clip_target_samples)
120
+ clip_sizes = [self._clip_target_samples] * num_full_clips
121
+ if remainder > 0:
122
+ clip_sizes.append(max(remainder, self._tail_min_samples))
123
+ return clip_sizes
124
125
def audio_token_count(self, audio_len: int) -> int:
128
- audio_len = self._normalize_audio_length(audio_len)
129
- num_frames = audio_len // self.hop_length
130
- n_tokens = HFParakeetEncoder._get_subsampling_output_length(
131
- self, torch.tensor([num_frames], dtype=torch.float)
132
- )
133
- return max(1, n_tokens.item())
126
+ total_tokens = 0
127
+ for clip_size in self._clip_sizes(audio_len):
128
+ num_frames = clip_size // self.hop_length
129
+ n_tokens = HFParakeetEncoder._get_subsampling_output_length(
130
+ self, torch.tensor([num_frames], dtype=torch.float)
131
+ )
132
+ total_tokens += int(n_tokens.item())
133
+ return max(1, total_tokens)
134
+
135
+ def split_audio_into_clips(self, audio: np.ndarray) -> list[np.ndarray]:
136
+ assert audio.ndim == 1
137
+ audio_len = int(audio.shape[0])
138
+ clip_sizes = self._clip_sizes(audio_len)
139
+ target_len = sum(clip_sizes)
140
+ if audio_len < target_len:
141
+ audio = np.pad(audio, (0, target_len - audio_len))
142
+
143
+ clips = list[np.ndarray]()
144
+ offset = 0
145
+ for clip_size in clip_sizes:
146
+ clips.append(audio[offset : offset + clip_size])
147
+ offset += clip_size
148
+ return clips
149
150
def __call__(self, raw_speech: list[np.ndarray], *args, **kwargs):
136
- padded = []
137
- for p in raw_speech:
138
- assert p.ndim == 1
139
- audio_len = int(p.shape[0])
140
- target_len = self._normalize_audio_length(audio_len)
141
- p = np.pad(p, (0, target_len - audio_len))
142
- padded.append(p)
143
- return super().__call__(padded, *args, **kwargs)
151
+ audio_clips = list[np.ndarray]()
152
+ audio_num_clips = list[int]()
153
+ for audio in raw_speech:
154
+ clips = self.split_audio_into_clips(audio)
155
+ audio_clips.extend(clips)
156
+ audio_num_clips.append(len(clips))
157
+
158
+ outputs = super().__call__(audio_clips, *args, **kwargs)
159
+ outputs["audio_num_clips"] = audio_num_clips
160
+ return outputs
161
162
def audio_length(self, audio_tokens: int) -> int:
163
return int(audio_tokens * self.config.subsampling_factor * self.hop_length)
vllm/transformers_utils/processors/nano_nemotron_vl.py
+4
-7
@@ -845,7 +845,7 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
845
audios: list[npt.NDArray],
846
) -> tuple[list[str], dict[str, Any]]:
847
if len(audios) == 0:
848
- return text, {}
848
+ return text, {"audio_num_clips": []}
849
850
assert self.audio_extractor is not None
851
extractor = self.audio_extractor
@@ -869,13 +869,10 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
869
sampling_rate=extractor.sampling_rate,
870
return_tensors="pt",
871
)
872
- input_audio_features = audio_inputs.input_features
873
- feature_attention_mask = audio_inputs.attention_mask
874
- audio_feature_lengths = feature_attention_mask.sum(dim=1)
872
audio_inputs = {
876
- "input_audio_features": input_audio_features,
877
- "feature_attention_mask": feature_attention_mask,
878
- "audio_feature_lengths": audio_feature_lengths,
873
+ "input_audio_features": audio_inputs.input_features,
874
+ "feature_attention_mask": audio_inputs.attention_mask,
875
+ "audio_num_clips": audio_inputs.audio_num_clips,
876
}
877
878
return text, audio_inputs