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 5, 06:00 UTC
ParakeetExtractor performance and UX enhancements (#39423)
Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>
Netanel Haber committed
Apr 14, 2026 at 00:37 UTC
6adacfcb65c1d40aa5e6d99249abba9cf3e350da
3 files changed
+177
-31
vllm/model_executor/models/parakeet.py
+162
-15
@@ -5,19 +5,24 @@ Modules below used for the audio encoder component in: models/nano_nemotron_vl.p
5
"""
6
7
from collections.abc import Iterable
8
-from dataclasses import asdict
8
+from functools import cache
9
+from typing import Any
10
11
import numpy as np
12
import torch
13
import torch.nn as nn
14
from transformers import ParakeetEncoder as HFParakeetEncoder
14
-from transformers import ParakeetFeatureExtractor, PretrainedConfig
15
+from transformers import PretrainedConfig
16
+from transformers.audio_utils import mel_filter_bank
17
18
+from vllm.logger import init_logger
19
from vllm.model_executor.layers.activation import ReLUSquaredActivation
20
from vllm.model_executor.layers.layernorm import RMSNorm
21
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
22
from vllm.transformers_utils.configs.parakeet import ExtractorConfig, ParakeetConfig
23
24
+logger = init_logger(__name__)
25
+
26
27
class ParakeetProjection(nn.Module):
28
def __init__(self, config: ParakeetConfig) -> None:
@@ -103,16 +108,124 @@ class ProjectedParakeet(nn.Module):
108
return loaded_params
109
110
106
-class ParakeetExtractor(ParakeetFeatureExtractor):
111
+EPSILON = 1e-5
112
+LOG_ZERO_GUARD_VALUE = 2**-24
113
+
114
+
115
+class ParakeetExtractor:
116
def __init__(self, config: PretrainedConfig) -> None:
117
self.config = ExtractorConfig.from_hf_config(config)
109
- super().__init__(**asdict(self.config))
118
+ """`config` is named *exactly* for `._get_subsampling_output_length` below"""
119
self._clip_target_samples = int(
111
- round(self.config.clip_duration_s * self.sampling_rate)
120
+ round(self.config.clip_duration_s * self.config.sampling_rate)
121
)
122
self._tail_min_samples = int(
114
- round(self.config.clip_min_duration_s * self.sampling_rate)
123
+ round(self.config.clip_min_duration_s * self.config.sampling_rate)
124
+ )
125
+
126
+ @staticmethod
127
+ @cache
128
+ def _get_window(win_length: int, device: str) -> torch.Tensor:
129
+ return torch.hann_window(win_length, periodic=False, device=device)
130
+
131
+ @staticmethod
132
+ @cache
133
+ def _get_mel_filters(
134
+ feature_size: int, sampling_rate: int, n_fft: int, device: str
135
+ ) -> torch.Tensor:
136
+ filter_bank = mel_filter_bank(
137
+ num_frequency_bins=n_fft // 2 + 1,
138
+ num_mel_filters=feature_size,
139
+ min_frequency=0.0,
140
+ max_frequency=sampling_rate / 2,
141
+ sampling_rate=sampling_rate,
142
+ norm="slaney",
143
+ mel_scale="slaney",
144
+ )
145
+ return torch.from_numpy(filter_bank.T).to(device=device, dtype=torch.float32)
146
+
147
+ def _torch_extract_fbank_features(self, waveform: torch.Tensor, device: str):
148
+ # spectrogram
149
+ device = str(torch.device(device))
150
+ cfg = self.config
151
+ window = self._get_window(cfg.win_length, device)
152
+ stft = torch.stft(
153
+ waveform,
154
+ self.config.n_fft,
155
+ hop_length=cfg.hop_length,
156
+ win_length=cfg.win_length,
157
+ window=window,
158
+ return_complex=True,
159
+ pad_mode="constant",
160
+ )
161
+ mel_filters = self._get_mel_filters(
162
+ cfg.feature_size, cfg.sampling_rate, cfg.n_fft, device
163
+ )
164
+ return self._apply_mel_filters(stft, mel_filters)
165
+
166
+ @torch.compile(dynamic=True)
167
+ def _apply_mel_filters(
168
+ self, stft_output: torch.Tensor, mel_filters: torch.Tensor
169
+ ) -> torch.Tensor:
170
+ magnitudes = stft_output.real.square() + stft_output.imag.square()
171
+ mel_spec = mel_filters @ magnitudes
172
+ mel_spec = torch.log(mel_spec + LOG_ZERO_GUARD_VALUE)
173
+ return mel_spec.permute(0, 2, 1)
174
+
175
+ @torch.compile(dynamic=True)
176
+ def _apply_preemphasis(
177
+ self, input_features: torch.Tensor, audio_lengths: torch.Tensor
178
+ ) -> torch.Tensor:
179
+ timemask = torch.arange(
180
+ input_features.shape[1], device=input_features.device
181
+ ).unsqueeze(0) < audio_lengths.unsqueeze(1)
182
+ input_features = torch.cat(
183
+ [
184
+ input_features[:, :1],
185
+ input_features[:, 1:]
186
+ - self.config.preemphasis * input_features[:, :-1],
187
+ ],
188
+ dim=1,
189
+ )
190
+ input_features = input_features.masked_fill(~timemask, 0.0)
191
+ return input_features
192
+
193
+ @torch.compile(dynamic=True)
194
+ def _normalize_mel_features(
195
+ self, mel_features: torch.Tensor, audio_lengths: torch.Tensor
196
+ ) -> tuple[torch.Tensor, torch.Tensor]:
197
+ features_lengths = torch.floor_divide(
198
+ audio_lengths + self.config.n_fft // 2 * 2 - self.config.n_fft,
199
+ self.config.hop_length,
200
+ )
201
+ attention_mask = (
202
+ torch.arange(mel_features.shape[1], device=mel_features.device)[None, :]
203
+ < features_lengths[:, None]
204
)
205
+ mask = attention_mask.unsqueeze(-1)
206
+ lengths = attention_mask.sum(dim=1)
207
+ mel_features_masked = mel_features * mask
208
+ mean = (mel_features_masked.sum(dim=1) / lengths.unsqueeze(-1)).unsqueeze(1)
209
+ variance = ((mel_features_masked - mean) ** 2 * mask).sum(dim=1) / (
210
+ lengths - 1
211
+ ).unsqueeze(-1)
212
+ std = torch.sqrt(variance).unsqueeze(1)
213
+ return (mel_features - mean) / (std + EPSILON) * mask, attention_mask
214
+
215
+ def _pad_raw_speech(
216
+ self, raw_speech: list[torch.Tensor], max_len: int, device: str
217
+ ) -> torch.Tensor:
218
+ output = torch.full(
219
+ (len(raw_speech), max_len),
220
+ self.config.padding_value,
221
+ device=device,
222
+ dtype=torch.float32,
223
+ )
224
+ dsts = [output[i, : raw_speech[i].shape[0]] for i in range(len(raw_speech))]
225
+ srcs = [s.squeeze(-1) for s in raw_speech]
226
+ # single kernel horizontal fusion
227
+ torch._foreach_copy_(dsts, srcs)
228
+ return output
229
230
def _clip_sizes(self, audio_len: int) -> list[int]:
231
audio_len = max(audio_len, self._tail_min_samples)
@@ -125,39 +238,73 @@ class ParakeetExtractor(ParakeetFeatureExtractor):
238
def audio_token_count(self, audio_len: int) -> int:
239
total_tokens = 0
240
for clip_size in self._clip_sizes(audio_len):
128
- num_frames = clip_size // self.hop_length
241
+ num_frames = clip_size // self.config.hop_length
242
n_tokens = HFParakeetEncoder._get_subsampling_output_length(
243
self, torch.tensor([num_frames], dtype=torch.float)
244
)
245
total_tokens += int(n_tokens.item())
246
return max(1, total_tokens)
247
135
- def split_audio_into_clips(self, audio: np.ndarray) -> list[np.ndarray]:
248
+ def split_audio_into_clips(self, audio: torch.Tensor) -> list[torch.Tensor]:
249
assert audio.ndim == 1
250
audio_len = int(audio.shape[0])
251
clip_sizes = self._clip_sizes(audio_len)
252
target_len = sum(clip_sizes)
253
if audio_len < target_len:
141
- audio = np.pad(audio, (0, target_len - audio_len))
254
+ audio = torch.nn.functional.pad(audio, (0, target_len - audio_len))
255
143
- clips = list[np.ndarray]()
256
+ clips = list[torch.Tensor]()
257
offset = 0
258
for clip_size in clip_sizes:
259
clips.append(audio[offset : offset + clip_size])
260
offset += clip_size
261
return clips
262
150
- def __call__(self, raw_speech: list[np.ndarray], *args, **kwargs):
151
- audio_clips = list[np.ndarray]()
263
+ def __call__(
264
+ self,
265
+ raw_speech: list[np.ndarray],
266
+ *,
267
+ device: str = "cpu",
268
+ ) -> dict[str, Any]:
269
+ raw_speech = [
270
+ torch.as_tensor(speech, device=device, dtype=torch.float32)
271
+ for speech in raw_speech
272
+ ]
273
+
274
+ for i, speech in enumerate(raw_speech):
275
+ if len(speech.shape) > 1:
276
+ logger.warning(
277
+ "Only mono-channel audio is supported for input to %s. "
278
+ "We will take the mean of the channels to convert to mono.",
279
+ self.__class__.__name__,
280
+ )
281
+ raw_speech[i] = speech.mean(-1)
282
+
283
+ audio_clips = list[torch.Tensor]()
284
audio_num_clips = list[int]()
285
for audio in raw_speech:
286
clips = self.split_audio_into_clips(audio)
287
audio_clips.extend(clips)
288
audio_num_clips.append(len(clips))
289
+ raw_speech = audio_clips
290
158
- outputs = super().__call__(audio_clips, *args, **kwargs)
159
- outputs["audio_num_clips"] = audio_num_clips
160
- return outputs
291
+ audio_lengths = torch.tensor(
292
+ [len(speech) for speech in raw_speech], dtype=torch.long, device=device
293
+ )
294
+
295
+ max_length = max(len(speech) for speech in raw_speech)
296
+ input_features = self._pad_raw_speech(raw_speech, max_length, device)
297
+ input_features = self._apply_preemphasis(input_features, audio_lengths)
298
+ input_features = self._torch_extract_fbank_features(input_features, device)
299
+ input_features, attention_mask = self._normalize_mel_features(
300
+ input_features, audio_lengths
301
+ )
302
+
303
+ return {
304
+ "input_audio_features": input_features,
305
+ "feature_attention_mask": attention_mask,
306
+ "audio_num_clips": audio_num_clips,
307
+ }
308
309
@staticmethod
310
def audio_length(raw_config: PretrainedConfig, audio_tokens: int) -> int:
vllm/transformers_utils/configs/parakeet.py
+14
-5
@@ -49,15 +49,24 @@ class ExtractorConfig:
49
clip_duration_s: int = 30
50
clip_min_duration_s: float = 0.1
51
52
- @staticmethod
53
- def from_hf_config(config: PretrainedConfig) -> "ExtractorConfig":
52
+ win_length: int = 400
53
+ preemphasis: float = 0.97
54
+ n_fft: int = 512
55
+ padding_value: float = 0.0
56
+
57
+ @classmethod
58
+ def from_hf_config(cls, config: PretrainedConfig) -> "ExtractorConfig":
59
assert isinstance(config, PretrainedConfig)
55
- hop_length = int(getattr(config, "hop_length", ExtractorConfig.hop_length))
56
- return ExtractorConfig(
60
+ defaults = ("hop_length", "win_length", "preemphasis", "n_fft", "padding_value")
61
+ optional_kwargs = {
62
+ name: getattr(config, name) for name in defaults if hasattr(config, name)
63
+ }
64
+
65
+ return cls(
66
feature_size=config.num_mel_bins,
67
sampling_rate=config.sampling_rate,
59
- hop_length=hop_length,
68
subsampling_factor=config.subsampling_factor,
69
subsampling_conv_kernel_size=config.subsampling_conv_kernel_size,
70
subsampling_conv_stride=config.subsampling_conv_stride,
71
+ **optional_kwargs,
72
)
vllm/transformers_utils/processors/nano_nemotron_vl.py
+1
-11
@@ -992,17 +992,7 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
992
parts[idx] = audio_repl.full
993
audio_index += 1
994
text = ["".join(parts)]
995
- audio_inputs = extractor(
996
- audios,
997
- sampling_rate=extractor.sampling_rate,
998
- return_tensors="pt",
999
- )
1000
- audio_inputs = {
1001
- "input_audio_features": audio_inputs.input_features,
1002
- "feature_attention_mask": audio_inputs.attention_mask,
1003
- "audio_num_clips": audio_inputs.audio_num_clips,
1004
- }
1005
-
995
+ audio_inputs = extractor(audios)
996
return text, audio_inputs
997
998
def __call__(