@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 234 lines 8.41 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # Copyright 2026 The Qwen team.
4 # Copyright 2023 The vLLM team.
5 #
6 # Licensed under the Apache License, Version 2.0 (the "License");
7 # you may not use this file except in compliance with the License.
8 # You may obtain a copy of the License at
9 #
10 # http://www.apache.org/licenses/LICENSE-2.0
11 #
12 # Unless required by applicable law or agreed to in writing, software
13 # distributed under the License is distributed on an "AS IS" BASIS,
14 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15 # See the License for the specific language governing permissions and
16 # limitations under the License.
17 """Inference-only Qwen3-ASR realtime model."""
18
19 import asyncio
20 from collections.abc import AsyncGenerator, Mapping
21
22 import numpy as np
23 import torch
24
25 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
26 from vllm.inputs import PromptType, TokensPrompt
27 from vllm.logger import init_logger
28 from vllm.multimodal import MULTIMODAL_REGISTRY
29 from vllm.multimodal.parse import MultiModalDataItems
30 from vllm.multimodal.processing import ProcessorInputs, TimingContext
31 from vllm.multimodal.processing.processor import (
32 MultiModalProcessingResult,
33 PlaceholderFeaturesInfo,
34 cached_encode,
35 )
36 from vllm.tokenizers import cached_tokenizer_from_config
37 from vllm.transformers_utils.processor import cached_processor_from_config
38
39 from .interfaces import SupportsRealtime
40 from .qwen3_asr import (
41 Qwen3ASRDummyInputsBuilder,
42 Qwen3ASRForConditionalGeneration,
43 Qwen3ASRMultiModalProcessor,
44 Qwen3ASRProcessingInfo,
45 _get_feat_extract_output_lengths,
46 )
47
48 logger = init_logger(__name__)
49
50 _PRE_ALLOCATE_BUFFER_SIZE_IN_S = 60
51
52
53 class Qwen3ASRRealtimeBuffer:
54 """Audio buffer for Qwen3-ASR realtime streaming.
55
56 Accumulates audio samples and yields segments when enough
57 audio has been buffered for processing.
58 """
59
60 def __init__(self, sampling_rate: int, segment_duration_s: float = 5.0):
61 self._sampling_rate = sampling_rate
62 self._segment_size = int(segment_duration_s * sampling_rate)
63
64 self._buffer_size = _PRE_ALLOCATE_BUFFER_SIZE_IN_S * sampling_rate
65 self._buffer: np.ndarray = np.empty(self._buffer_size, dtype=np.float32)
66 self._filled_len = 0
67
68 def write_audio(self, audio: np.ndarray) -> None:
69 put_end = self._filled_len + len(audio)
70 if put_end > self._buffer_size:
71 new_size = max(self._buffer_size * 2, put_end)
72 new_buffer = np.empty(new_size, dtype=np.float32)
73 new_buffer[: self._filled_len] = self._buffer[: self._filled_len]
74 self._buffer = new_buffer
75 self._buffer_size = new_size
76
77 self._buffer[self._filled_len : put_end] = audio
78 self._filled_len = put_end
79
80 def read_audio(self) -> np.ndarray | None:
81 if self._filled_len < self._segment_size:
82 return None
83
84 segment = self._buffer[: self._segment_size].copy()
85 remaining = self._filled_len - self._segment_size
86 if remaining > 0:
87 self._buffer[:remaining] = self._buffer[
88 self._segment_size : self._filled_len
89 ]
90 self._filled_len = remaining
91 return segment
92
93 def flush(self) -> np.ndarray | None:
94 if self._filled_len == 0:
95 return None
96 audio = self._buffer[: self._filled_len].copy()
97 self._filled_len = 0
98 return audio
99
100
101 class Qwen3ASRRealtimeMultiModalProcessor(Qwen3ASRMultiModalProcessor):
102 def _cached_apply_hf_processor(
103 self,
104 inputs: ProcessorInputs,
105 timing_ctx: TimingContext,
106 ) -> MultiModalProcessingResult:
107 # realtime can't make use of a cache yet
108 return self._apply_hf_processor(inputs, timing_ctx)
109
110 def _maybe_apply_prompt_updates(
111 self,
112 mm_items: MultiModalDataItems,
113 mm_res: MultiModalProcessingResult,
114 ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]:
115 mm_kwargs = mm_res.kwargs
116
117 audios = mm_kwargs.get("audio", [])
118 assert len(audios) == 1, (
119 f"Expected only one audio input for realtime, got {len(audios)}"
120 )
121
122 audio_data = audios[0]
123 assert audio_data is not None
124 audio_feature_lengths = audio_data.get("audio_feature_lengths")
125 if audio_feature_lengths is not None:
126 if isinstance(audio_feature_lengths.data, torch.Tensor):
127 audio_len = _get_feat_extract_output_lengths(
128 audio_feature_lengths.data
129 ).item()
130 else:
131 audio_len = int(
132 _get_feat_extract_output_lengths(
133 torch.tensor(audio_feature_lengths.data)
134 ).item()
135 )
136 else:
137 audio_len = 0
138
139 # Get audio_pad token ID and expand placeholder in prompt_ids
140 # so that MRoPE position computation matches seq_len.
141 tokenizer = self.info.get_tokenizer()
142 audio_pad_id = tokenizer.convert_tokens_to_ids("<|audio_pad|>")
143
144 # Find the audio_pad token position and expand it to audio_len tokens
145 expanded_ids = list[int]()
146 pad_start_idx = -1
147 for i, tid in enumerate(mm_res.prompt_ids):
148 if tid == audio_pad_id and pad_start_idx == -1:
149 pad_start_idx = i
150 expanded_ids.extend([audio_pad_id] * audio_len)
151 else:
152 expanded_ids.append(tid)
153
154 if pad_start_idx == -1:
155 pad_start_idx = 0
156
157 features_info = PlaceholderFeaturesInfo(
158 modality="audio",
159 item_idx=0,
160 start_idx=pad_start_idx,
161 tokens=audio_len * [audio_pad_id],
162 is_embed=None,
163 )
164 return expanded_ids, {"audio": [features_info]}
165
166
167 # NOTE: A separate model class is required here because the multimodal
168 # processor registry binds one processor per model class. The realtime
169 # endpoint needs a different processor (Qwen3ASRRealtimeMultiModalProcessor)
170 # than the base transcription endpoint, so we register it on this subclass.
171 @MULTIMODAL_REGISTRY.register_processor(
172 Qwen3ASRRealtimeMultiModalProcessor,
173 info=Qwen3ASRProcessingInfo,
174 dummy_inputs=Qwen3ASRDummyInputsBuilder,
175 )
176 class Qwen3ASRRealtimeGeneration(Qwen3ASRForConditionalGeneration, SupportsRealtime):
177 realtime_max_tokens = 64
178
179 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
180 super().__init__(vllm_config=vllm_config, prefix=prefix)
181
182 @classmethod
183 async def buffer_realtime_audio(
184 cls,
185 audio_stream: AsyncGenerator[np.ndarray, None],
186 input_stream: asyncio.Queue[list[int]],
187 model_config: ModelConfig,
188 ) -> AsyncGenerator[PromptType, None]:
189 processor = cached_processor_from_config(model_config)
190 feature_extractor = processor.feature_extractor
191 sampling_rate = feature_extractor.sampling_rate
192 tokenizer = cached_tokenizer_from_config(model_config)
193
194 # Use a small segment size for low-latency streaming.
195 segment_duration_s = 5.0
196 buffer = Qwen3ASRRealtimeBuffer(
197 sampling_rate=sampling_rate,
198 segment_duration_s=segment_duration_s,
199 )
200
201 audio_placeholder = cls.get_placeholder_str("audio", 0)
202 prompt_template = (
203 f"<|im_start|>user\n{audio_placeholder}<|im_end|>\n<|im_start|>assistant\n"
204 )
205
206 prompt_token_ids = cached_encode(tokenizer, prompt_template)
207
208 async for audio_chunk in audio_stream:
209 buffer.write_audio(audio_chunk)
210
211 while (segment := buffer.read_audio()) is not None:
212 yield TokensPrompt(
213 prompt_token_ids=prompt_token_ids,
214 multi_modal_data={"audio": segment},
215 )
216
217 remaining = buffer.flush()
218 if remaining is not None and len(remaining) > 0:
219 yield TokensPrompt(
220 prompt_token_ids=prompt_token_ids,
221 multi_modal_data={"audio": remaining},
222 )
223
224 @classmethod
225 def get_speech_to_text_config(
226 cls, model_config: ModelConfig, task_type: str
227 ) -> SpeechToTextConfig:
228 processor = cached_processor_from_config(model_config)
229 feature_extractor = processor.feature_extractor
230 return SpeechToTextConfig(
231 max_audio_clip_s=None,
232 sample_rate=feature_extractor.sampling_rate,
233 min_energy_split_window_size=None,
234 )