@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6c0baee610

Fork 0

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 8, 08:45 UTC

[Voxtral Realtime] Refactor & Improve buffering logic (#34428)

Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Patrick von Platen committed Feb 12, 2026 at 18:46 UTC 6c0baee61025f258c6d56830d0150feab34c45ab

7 files changed +135 -238

requirements/common.txt

+2 -2

@@ -31,7 +31,7 @@ partial-json-parser # used for parsing partial JSON outputs

31 pyzmq >= 25.0.0

32 msgspec

33 gguf >= 0.17.0

34 -mistral_common[image] >= 1.9.0

34 +mistral_common[image] >= 1.9.1

35 opencv-python-headless >= 4.13.0 # required for video IO

36 pyyaml

37 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12

@@ -52,4 +52,4 @@ anthropic >= 0.71.0

52 model-hosting-container-standards >= 0.1.13, < 1.0.0

53 mcp

54 grpcio

55 -grpcio-reflection

\ No newline at end of file

55 +grpcio-reflection

requirements/nightly_torch_test.txt

+1 -1

@@ -23,7 +23,7 @@ jiwer # required for audio tests

23 timm # required for internvl test

24 transformers_stream_generator # required for qwen-vl test

25 matplotlib # required for qwen-vl test

26 -mistral_common[image,audio] >= 1.9.0 # required for voxtral test

26 +mistral_common[image,audio] >= 1.9.1 # required for voxtral test

27 num2words # required for smolvlm test

28 opencv-python-headless >= 4.13.0 # required for video test

29 datamodel_code_generator # required for minicpm3 test

requirements/test.in

+1 -1

@@ -30,7 +30,7 @@ torchaudio==2.10.0

30 torchvision==0.25.0

31 transformers_stream_generator # required for qwen-vl test

32 matplotlib # required for qwen-vl test

33 -mistral_common[image,audio] >= 1.9.0 # required for voxtral test

33 +mistral_common[image,audio] >= 1.9.1 # required for voxtral test

34 num2words # required for smolvlm test

35 open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py

36 opencv-python-headless >= 4.13.0 # required for video test

requirements/test.txt

+1 -1

@@ -499,7 +499,7 @@ mbstrdecoder==1.1.3

499 # typepy

500 mdurl==0.1.2

501 # via markdown-it-py

502 -mistral-common==1.9.0

502 +mistral-common==1.9.1

503 # via -r requirements/test.in

504 mlflow==2.22.0

505 # via terratorch

tests/models/multimodal/generation/test_voxtral_realtime.py

+15 -113

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -import asyncio

3 from dataclasses import asdict

5 import pytest

@@ -10,14 +9,13 @@ from mistral_common.protocol.transcription.request import (

9 StreamingMode,

10 TranscriptionRequest,

11 )

13 -from mistral_common.tokens.tokenizers.audio import AudioConfig

12 from mistral_common.tokens.tokenizers.mistral import MistralTokenizer

14 from vllm import LLM, EngineArgs, SamplingParams

15 from vllm.assets.audio import AudioAsset

16 from vllm.engine.arg_utils import AsyncEngineArgs

19 -from vllm.inputs.data import TokensPrompt

20 -from vllm.v1.engine.async_llm import AsyncLLM, StreamingInput

17 +from vllm.model_executor.models.voxtral_realtime import VoxtralRealtimeBuffer

18 +from vllm.v1.engine.async_llm import AsyncLLM

20 MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"

21 ENGINE_CONFIG = dict(

@@ -114,136 +112,40 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):

112 assert texts == EXPECTED_TEXT

113

114

117 -class RealTimeAudioInput:

118 - """

119 - This class is used to stream an audio file just as

120 - if it would be streamed in real-time.

121 - """

122 -

123 - def __init__(self, tokenizer: MistralTokenizer) -> None:

124 - self._tokenizer = tokenizer

125 - self._config: AudioConfig = (

126 - self._tokenizer.instruct_tokenizer.audio_encoder.audio_config

127 - )

128 -

129 - self._look_ahead_in_ms = self._config.streaming_look_ahead_ms

130 - self._look_back_in_ms = self._config.streaming_look_back_ms

131 -

132 - self._sampling_rate = self._config.sampling_rate

133 -

134 - self._audio: Audio | None = None

135 -

136 - # mutable objects

137 - self._start = 0

138 -

139 - n_left_pad_samples = (

140 - self._config.raw_audio_length_per_tok * self._config.n_left_pad_tokens

141 - )

142 - self._end = self.streaming_delay + n_left_pad_samples + self.streaming_size

143 - self._queue: asyncio.Queue[StreamingInput | None] = asyncio.Queue()

144 -

145 - @classmethod

146 - async def create(cls, audio: Audio, tokenizer: MistralTokenizer):

147 - self = cls(tokenizer)

148 -

149 - # we're doing "OFFLINE" encoding here to right & left pad the audio since

150 - # we have access to the whole audio

151 - # if we'd do an actual online realtime streaming application we

152 - # should instead pass `StreamingMode.ONLINE`

153 - req = TranscriptionRequest(

154 - streaming=StreamingMode.OFFLINE,

155 - audio=RawAudio.from_audio(audio),

156 - language=None,

157 - )

158 - audio_enc = self._tokenizer.encode_transcription(req)

159 - self._audio = audio_enc.audios[0]

160 -

161 - # add first request

162 - await self.add_tokens(audio_enc.tokens)

163 -

164 - return self

165 -

166 - @property

167 - def look_ahead(self) -> int:

168 - return self._get_len_in_samples(self._look_ahead_in_ms)

169 -

170 - @property

171 - def look_back(self) -> int:

172 - return self._get_len_in_samples(self._look_back_in_ms)

173 -

174 - @property

175 - def streaming_delay(self) -> int:

176 - return self._get_len_in_samples(self._config.transcription_delay_ms)

177 -

178 - @property

179 - def streaming_size(self) -> int:

180 - stream_size_in_ms = 1000 / self._config.frame_rate

181 - return self._get_len_in_samples(stream_size_in_ms)

182 -

183 - def _get_len_in_samples(self, len_in_ms: float) -> int:

184 - _len_in_s = self._sampling_rate * len_in_ms / 1000

185 - assert _len_in_s.is_integer(), _len_in_s

186 - len_in_s = int(_len_in_s)

187 -

188 - return len_in_s

189 -

190 - async def add_tokens(self, tokens: list[int]) -> None:

191 - assert self._audio is not None

192 - if self._start >= len(self._audio.audio_array):

193 - self.stop()

194 - return

195 -

196 - _end = self._end + self.look_ahead

197 - _start = max(0, self._start - self.look_back)

198 -

199 - multi_modal_data = {"audio": (self._audio.audio_array[_start:_end], None)}

200 -

201 - prompt = TokensPrompt(

202 - prompt_token_ids=tokens, multi_modal_data=multi_modal_data

203 - )

204 -

205 - await self._queue.put(StreamingInput(prompt))

206 -

207 - # increase

208 - self._start = self._end

209 - self._end = self._end + self.streaming_size

210 -

211 - def stop(self):

212 - self._queue.put_nowait(None)

213 -

214 - async def generator(self):

215 - while (item := await self._queue.get()) is not None:

216 - yield item

217 -

218 -

115 @pytest.mark.asyncio

116 async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine):

117 sampling_params = SamplingParams(temperature=0.0, max_tokens=1)

118 + audio_config = tokenizer.instruct_tokenizer.audio_encoder.audio_config

119

120 output_tokens_list = []

121 for i, audio_asset in enumerate(audio_assets):

122 output_tokens = []

123 audio = Audio.from_file(audio_asset.get_local_path(), strict=False)

227 - streaming_input = await RealTimeAudioInput.create(

228 - audio=audio, tokenizer=tokenizer

124 +

125 + req = TranscriptionRequest(

126 + streaming=StreamingMode.OFFLINE,

127 + audio=RawAudio.from_audio(audio),

128 + language=None,

129 )

130 + audio_enc = tokenizer.encode_transcription(req)

131 +

132 + buffer = VoxtralRealtimeBuffer(audio_config, audio_enc.tokens)

133 + await buffer.append_audio(audio_enc.audios[0].audio_array)

134 + await buffer.append_audio(None)

135

136 request_id = f"session-{i}"

137

138 async for resp in async_engine.generate(

234 - prompt=streaming_input.generator(),

139 + prompt=buffer.get_input_stream(),

140 sampling_params=sampling_params,

141 request_id=request_id,

142 ):

143 tokens = resp.outputs[0].token_ids[-1:]

239 -

144 output_tokens.extend(tokens)

241 - await streaming_input.add_tokens(tokens)

145 + await buffer.append_tokens(tokens)

146

147 output_tokens_list.append(output_tokens)

148

149 texts = [tokenizer.decode(output_tokens) for output_tokens in output_tokens_list]

246 -

150 texts[1] = texts[1].replace("a base hit", "OBS").replace("oh my", "oh, my")

248 -

151 assert texts == EXPECTED_TEXT

vllm/model_executor/models/voxtral.py

+1 -3

@@ -155,9 +155,7 @@ class VoxtralProcessorAdapter:

155 assert audio.ndim == 1

156

157 if not self._audio_processor.audio_config.is_streaming:

158 - audio = self._audio_processor.pad(

159 - audio, self.sampling_rate, is_online_streaming=False

160 - )

158 + audio = self._audio_processor.pad(audio, self.sampling_rate)

159

160 audio_tokens = [self.begin_audio_token_id] + [

161 self.audio_token_id

vllm/model_executor/models/voxtral_realtime.py

+114 -117

@@ -3,7 +3,7 @@

4 import asyncio

5 import math

6 -from collections.abc import AsyncGenerator, Mapping

6 +from collections.abc import AsyncGenerator, Iterable, Iterator, Mapping

7 from typing import Literal

9 import numpy as np

@@ -18,7 +18,7 @@ from mistral_common.tokens.tokenizers.audio import Audio, AudioConfig

18 from vllm.compilation.decorators import support_torch_compile

19 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig

20 from vllm.envs import VLLM_ENGINE_ITERATION_TIMEOUT_S

21 -from vllm.inputs.data import PromptType, TokensPrompt

21 +from vllm.inputs.data import PromptType, StreamingInput, TokensPrompt

22 from vllm.logger import init_logger

23 from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsRealtime

24 from vllm.model_executor.models.voxtral import (

@@ -47,8 +47,6 @@ from .utils import (

48 logger = init_logger(__name__)

50 -_PRE_ALLOCATE_BUFFER_SIZE_IN_S = 30

51 -

51 class VoxtralRealtimeMultiModalProcessor(VoxtralMultiModalProcessor):

52 def __init__(

@@ -130,84 +128,81 @@ def _expand_tensor(input_tensor: torch.Tensor, scaling: int) -> torch.Tensor:

128

129

130 class VoxtralRealtimeBuffer:

133 - def __init__(self, config: AudioConfig) -> None:

131 + def __init__(self, config: AudioConfig, prompt_tokens: list[int]) -> None:

132 self._config = config

133

136 - self._look_ahead_in_ms = config.streaming_look_ahead_ms

137 - self._look_back_in_ms = config.streaming_look_back_ms

138 -

139 - self._sampling_rate = self._config.sampling_rate

140 -

141 - self._look_ahead = self._get_len_in_samples(self._look_ahead_in_ms)

142 - self._look_back = self._get_len_in_samples(self._look_back_in_ms)

143 - self._streaming_size = self._get_len_in_samples(1000 / self._config.frame_rate)

144 -

145 - # mutable objects

146 - streaming_delay = self._get_len_in_samples(self._config.transcription_delay_ms)

147 - self._start = 0

148 - self._end = streaming_delay + self._streaming_size

149 -

150 - # always pre-allocate 30 second buffers

151 - self._buffer_size = _PRE_ALLOCATE_BUFFER_SIZE_IN_S * self._sampling_rate

152 - self._buffer: np.ndarray = np.empty(self._buffer_size, dtype=np.float32)

153 - self._filled_buffer_len = 0

154 -

155 - @property

156 - def start_idx(self):

157 - return max(self._start - self._look_back, 0)

158 -

159 - @property

160 - def end_idx(self):

161 - return self._end + self._look_ahead

162 -

163 - @property

164 - def is_audio_complete(self) -> bool:

165 - return self._filled_buffer_len >= self.end_idx

166 -

167 - def _get_len_in_samples(self, len_in_ms: float) -> int:

168 - _len_in_s = self._sampling_rate * len_in_ms / 1000

169 - assert _len_in_s.is_integer(), _len_in_s

170 - len_in_s = int(_len_in_s)

171 -

172 - return len_in_s

173 -

174 - def _allocate_new_buffer(self) -> None:

175 - # allocate new buffer

176 - new_buffer = np.empty(self._buffer_size, dtype=np.float32)

177 - left_to_copy = max(self._filled_buffer_len - self.start_idx, 0)

178 -

179 - if left_to_copy > 0:

180 - new_buffer[:left_to_copy] = self._buffer[

181 - self.start_idx : self._filled_buffer_len

182 - ]

183 -

184 - del self._buffer

185 - self._buffer = new_buffer

186 -

187 - self._filled_buffer_len = left_to_copy

188 - self._start = self._look_back

189 - self._end = self._start + self._streaming_size

190 -

191 - def write_audio(self, audio: np.ndarray) -> None:

192 - put_end_idx = self._filled_buffer_len + len(audio)

193 -

194 - if put_end_idx > self._buffer_size:

195 - self._allocate_new_buffer()

196 -

197 - self._buffer[self._filled_buffer_len : self._filled_buffer_len + len(audio)] = (

198 - audio

199 - )

200 - self._filled_buffer_len += len(audio)

201 -

202 - def read_audio(self) -> np.ndarray | None:

203 - if not self.is_audio_complete:

204 - return None

134 + _look_ahead_in_ms = self._config.streaming_look_ahead_ms

135 + _look_back_in_ms = self._config.streaming_look_back_ms

136 + self._look_ahead_in_samples = self._ms_to_samples(_look_ahead_in_ms)

137 + self._look_back_in_samples = self._ms_to_samples(_look_back_in_ms)

138 +

139 + # None signals the end

140 + self._audio_queue: asyncio.Queue[np.ndarray | None] = asyncio.Queue()

141 + self._leftover: np.ndarray | None = None

142 + self._token_queue: asyncio.Queue[int] = asyncio.Queue()

143 +

144 + self._initial_end = len(prompt_tokens) * self._config.raw_audio_length_per_tok

145 + for token in prompt_tokens:

146 + self._token_queue.put_nowait(token)

147 +

148 + def _generate_frame_size_and_num_tokens(self) -> Iterator[tuple[int, int]]:

149 + streaming_step_size = self._ms_to_samples(1000 / self._config.frame_rate)

150 + start = 0

151 + end = self._initial_end

152 + while True:

153 + frame_start = max(start - self._look_back_in_samples, 0)

154 + frame_end = end + self._look_ahead_in_samples

155 + frame_size = frame_end - frame_start

156 + num_tokens = (end - start) / self._config.raw_audio_length_per_tok

157 + assert num_tokens.is_integer()

158 + yield frame_size, int(num_tokens)

159 + start = end

160 + end += streaming_step_size

161 +

162 + def _ms_to_samples(self, ms: float) -> int:

163 + len_ = self._config.sampling_rate * ms / 1000

164 + assert len_.is_integer(), len_

165 + return int(len_)

166 +

167 + async def append_audio(self, audio_array: np.ndarray | None) -> None:

168 + await self._audio_queue.put(audio_array)

169 +

170 + async def append_tokens(self, tokens: Iterable[int]) -> None:

171 + for token in tokens:

172 + await self._token_queue.put(token)

173 +

174 + async def get_input_stream(self) -> AsyncGenerator[StreamingInput]:

175 + for frame_size, num_tokens in self._generate_frame_size_and_num_tokens():

176 + next_tokens = [await self._token_queue.get() for _ in range(num_tokens)]

177 +

178 + audio_arrays: list[np.ndarray] = (

179 + [self._leftover] if self._leftover is not None else []

180 + )

181 + while sum(len(arr) for arr in audio_arrays) < frame_size:

182 + arr = await self._audio_queue.get()

183 + if arr is None:

184 + return

185 + audio_arrays.append(arr)

186 +

187 + audio_array = np.concatenate(audio_arrays)

188 + frame = audio_array[:frame_size]

189 +

190 + # The current stride took look_ahead_in_samples audio of the next sample

191 + # In addition the next sample will take look_back_in_samples audio of

192 + # the current sample => So let's put both of this into the leftover

193 + stride = (

194 + frame_size - self._look_ahead_in_samples - self._look_back_in_samples

195 + )

196 + assert stride > 0, f"{stride=} must be positive"

197

206 - audio = self._buffer[self.start_idx : self.end_idx]

207 - self._start = self._end

208 - self._end += self._streaming_size

198 + self._leftover = audio_array[stride:]

199

210 - return audio

200 + yield StreamingInput(

201 + TokensPrompt(

202 + prompt_token_ids=next_tokens,

203 + multi_modal_data={"audio": (frame, None)},

204 + )

205 + )

206

207

208 @MULTIMODAL_REGISTRY.register_processor(

@@ -234,7 +229,7 @@ class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtim

229 )

230

231 audio_config = self.tokenizer.instruct.audio_encoder.audio_config

237 - self.n_delay_tokens = audio_config.num_delay_tokens

232 + self.n_delay_tokens = audio_config.get_num_delay_tokens()

233

234 # for realtime transcription

235 @classmethod

@@ -248,45 +243,47 @@ class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtim

243 audio_encoder = tokenizer.instruct.audio_encoder

244 config = audio_encoder.audio_config

245

251 - buffer = VoxtralRealtimeBuffer(config)

252 - is_first_yield = True

253 -

254 - async for audio in audio_stream:

255 - buffer.write_audio(audio)

256 -

257 - while (new_audio := buffer.read_audio()) is not None:

258 - if is_first_yield:

259 - # make sure that input_stream is empty

260 - assert input_stream.empty()

261 -

262 - audio = Audio(new_audio, config.sampling_rate, format="wav")

263 -

264 - request = TranscriptionRequest(

265 - streaming=StreamingMode.ONLINE,

266 - audio=RawAudio.from_audio(audio),

267 - language=None,

268 - )

269 - # mistral tokenizer takes care

270 - # of preparing the first prompt inputs

271 - # and does some left-silence padding

272 - # for improved performance

273 - audio_enc = tokenizer.mistral.encode_transcription(request)

274 -

275 - token_ids = audio_enc.tokens

276 - new_audio = audio_enc.audios[0].audio_array

277 -

278 - is_first_yield = False

279 - else:

280 - # pop last element from input_stream

281 - all_outputs = await asyncio.wait_for(

282 - input_stream.get(), timeout=VLLM_ENGINE_ITERATION_TIMEOUT_S

283 - )

284 - token_ids = all_outputs[-1:]

285 -

286 - multi_modal_data = {"audio": (new_audio, None)}

287 - yield TokensPrompt(

288 - prompt_token_ids=token_ids, multi_modal_data=multi_modal_data

246 + # Get prompt tokens (streaming prefix tokens) without encoding audio

247 + prompt_tokens = (

248 + tokenizer.instruct.start() + audio_encoder.encode_streaming_tokens()

249 + )

250 +

251 + # Get left/right padding audio

252 + left_pad, right_pad = audio_encoder.get_padding_audio()

253 +

254 + buffer = VoxtralRealtimeBuffer(config, prompt_tokens)

255 +

256 + # Feed audio with padding into buffer in background

257 + async def feed_audio():

258 + yielded_first_chunk = False

259 + async for audio_chunk in audio_stream:

260 + if not yielded_first_chunk:

261 + yielded_first_chunk = True

262 + # Prepend left padding before first real audio

263 + await buffer.append_audio(left_pad.audio_array)

264 + await buffer.append_audio(audio_chunk)

265 + # Append right padding at the end

266 + await buffer.append_audio(right_pad.audio_array)

267 + await buffer.append_audio(None) # signal end

268 +

269 + # Feed output tokens back into buffer in background

270 + async def feed_tokens():

271 + while True:

272 + all_outputs = await asyncio.wait_for(

273 + input_stream.get(),

274 + timeout=VLLM_ENGINE_ITERATION_TIMEOUT_S,

275 )

276 + await buffer.append_tokens(all_outputs[-1:])

277 +

278 + audio_task = asyncio.create_task(feed_audio())

279 + token_task = asyncio.create_task(feed_tokens())

280 +

281 + try:

282 + async for streaming_input in buffer.get_input_stream():

283 + yield streaming_input.prompt

284 + finally:

285 + audio_task.cancel()

286 + token_task.cancel()

287

288 @property

289 def audio_config(self):