diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index 09432b4494..e05930b0f3 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -106,7 +106,8 @@ def tokenizer() -> MistralTokenizer: @pytest_asyncio.fixture -async def async_engine(): +async def async_engine(request): + engine_config = {**ENGINE_CONFIG, **getattr(request, "param", {})} gpu_memory_utilization = ENGINE_CONFIG.get("gpu_memory_utilization", 0.9) from vllm.platforms import current_platform @@ -115,7 +116,7 @@ async def async_engine(): wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) - engine_args = AsyncEngineArgs(**ENGINE_CONFIG) + engine_args = AsyncEngineArgs(**engine_config) llm = AsyncLLM.from_engine_args(engine_args) try: yield llm @@ -275,3 +276,82 @@ async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine) f" got: {got!r}\n" f" expected: {expected!r}" ) + + +@pytest.mark.asyncio +async def test_voxtral_realtime_buffer_stops_at_max_model_len(tokenizer): + from vllm.model_executor.models.voxtral_realtime import VoxtralRealtimeBuffer + + audio_config = tokenizer.instruct_tokenizer.audio_encoder.audio_config + audio = Audio.from_file(AudioAsset("winning_call").get_local_path(), strict=False) + req = TranscriptionRequest( + streaming=StreamingMode.OFFLINE, + audio=audio.to_base64(audio.format), + language=None, + ) + audio_enc = tokenizer.encode_transcription(req) + + max_model_len = 64 + buffer = VoxtralRealtimeBuffer( + audio_config, audio_enc.tokens, max_model_len=max_model_len + ) + await buffer.append_audio(audio_enc.audios[0].audio_array) + await buffer.append_audio(None) + # Unlimited feedback tokens so only the length cap can stop the stream. + await buffer.append_tokens([0] * max_model_len) + + n_tokens = 0 + async for prompt in buffer.get_input_stream(): + n_tokens += len(prompt["prompt_token_ids"]) + + assert len(audio_enc.tokens) < n_tokens < max_model_len + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_engine", [{"max_model_len": 128}], indirect=True) +async def test_voxtral_realtime_generator_max_model_len( + audio_assets, tokenizer, async_engine +): + # Lazy import to avoid CUDA-reinitialization error + from vllm.model_executor.models.voxtral_realtime import VoxtralRealtimeBuffer + + max_model_len = async_engine.model_config.max_model_len + sampling_params = SamplingParams(temperature=0.0, max_tokens=1) + audio_config = tokenizer.instruct_tokenizer.audio_encoder.audio_config + + async def input_stream(audio_buffer): + async for prompt in audio_buffer.get_input_stream(): + parsed_prompt = parse_model_prompt(async_engine.model_config, prompt) + (engine_input,) = await async_engine.renderer.render_cmpl_async( + [parsed_prompt] + ) + yield StreamingInput(prompt=engine_input) + + # winning_call is ~300 tokens long, well past max_model_len. + audio = Audio.from_file(audio_assets[1].get_local_path(), strict=False) + req = TranscriptionRequest( + streaming=StreamingMode.OFFLINE, + audio=audio.to_base64(audio.format), + language=None, + ) + audio_enc = tokenizer.encode_transcription(req) + buffer = VoxtralRealtimeBuffer( + audio_config, audio_enc.tokens, max_model_len=max_model_len + ) + await buffer.append_audio(audio_enc.audios[0].audio_array) + await buffer.append_audio(None) + + output_tokens = [] + async for resp in async_engine.generate( + prompt=input_stream(buffer), + sampling_params=sampling_params, + request_id="session-max-len", + ): + tokens = resp.outputs[0].token_ids[-1:] + output_tokens.extend(tokens) + await buffer.append_tokens(tokens) + + assert resp.finished + assert 0 < len(output_tokens) < max_model_len + # The engine is still healthy after the session hit the cap. + assert not async_engine.errored diff --git a/vllm/entrypoints/speech_to_text/realtime/connection.py b/vllm/entrypoints/speech_to_text/realtime/connection.py index ed8ec7f862..4a5e6cef9a 100644 --- a/vllm/entrypoints/speech_to_text/realtime/connection.py +++ b/vllm/entrypoints/speech_to_text/realtime/connection.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio +import contextlib import json from collections.abc import AsyncGenerator from http import HTTPStatus @@ -48,6 +49,7 @@ class RealtimeConnection: self.serving = serving self.audio_queue: asyncio.Queue[np.ndarray | None] = asyncio.Queue() self.generation_task: asyncio.Task | None = None + self.request_id: str | None = None self._is_connected = False self._is_model_validated = False @@ -204,6 +206,7 @@ class RealtimeConnection: 6. Cleans up the audio queue """ request_id = f"rt-{self.connection_id}-{uuid4()}" + self.request_id = request_id full_text = "" prompt_token_ids_len: int = 0 @@ -263,7 +266,9 @@ class RealtimeConnection: except Exception as e: logger.exception("Error in generation: %s", e) - await self.send_error(sanitize_message(str(e)), "processing_error") + # The socket may already be closed; nothing more we can do. + with contextlib.suppress(Exception): + await self.send_error(sanitize_message(str(e)), "processing_error") async def send( self, event: SessionCreated | TranscriptionDelta | TranscriptionDone @@ -282,8 +287,28 @@ class RealtimeConnection: # Signal audio stream to stop self.audio_queue.put_nowait(None) - # Cancel generation task if running + # Abort the engine request. Cancelling the asyncio task alone does NOT + # abort the underlying engine request: a request parked in the + # VAD-waiting state keeps occupying a slot and generating into a now + # dead websocket (an orphaned request that never frees and starves + # other streams). Aborting also unblocks the generation task, which is + # otherwise stuck awaiting the next engine output. + if self.request_id is not None: + try: + await self.serving.engine_client.abort(self.request_id) + except Exception: + logger.exception("Failed to abort request %s", self.request_id) + self.request_id = None + + # Cancel the generation task if still running, and await it so the + # cancellation (and any generator cleanup) actually completes. if self.generation_task and not self.generation_task.done(): self.generation_task.cancel() + try: + await self.generation_task + except asyncio.CancelledError: + pass + except Exception: + logger.exception("Error awaiting cancelled generation task") logger.debug("Connection cleanup complete: %s", self.connection_id) diff --git a/vllm/model_executor/models/voxtral_realtime.py b/vllm/model_executor/models/voxtral_realtime.py index 8531d7d18f..72318976f1 100644 --- a/vllm/model_executor/models/voxtral_realtime.py +++ b/vllm/model_executor/models/voxtral_realtime.py @@ -120,9 +120,35 @@ def _expand_tensor(input_tensor: torch.Tensor, scaling: int) -> torch.Tensor: return (base.unsqueeze(1) + offsets).view(-1) +class VoxtralRealtimeProcessingInfo(VoxtralProcessingInfo): + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + # Each streaming update carries a single audio frame. + return {"audio": 1} + + def get_max_audio_array_len(self) -> int: + # Profile with the largest streaming chunk, the first one: it covers + # the prompt prefix (BOS + left padding + delay), later chunks are + # ~1 token. get_max_audio_tokens keeps max_model_len so the encoder + # cache stays as large as before: each paused session keeps its last + # chunk cached until it resumes, and a cache sized to one chunk + # deadlocks a few concurrent sessions. + tokenizer = self.get_tokenizer() + audio_encoder = tokenizer.instruct.audio_encoder + prompt_tokens = ( + tokenizer.instruct.start() + audio_encoder.encode_streaming_tokens() + ) + return len(prompt_tokens) * audio_encoder.audio_config.raw_audio_length_per_tok + + class VoxtralRealtimeBuffer: - def __init__(self, config: AudioConfig, prompt_tokens: list[int]) -> None: + def __init__( + self, + config: AudioConfig, + prompt_tokens: list[int], + max_model_len: int | None = None, + ) -> None: self._config = config + self._max_model_len = max_model_len _look_ahead_in_ms = self._config.streaming_look_ahead_ms _look_back_in_ms = self._config.streaming_look_back_ms @@ -165,7 +191,16 @@ class VoxtralRealtimeBuffer: await self._token_queue.put(token) async def get_input_stream(self) -> AsyncGenerator[TokensPrompt]: + n_tokens = 0 for frame_size, num_tokens in self._generate_frame_size_and_num_tokens(): + n_tokens += num_tokens + if self._max_model_len is not None and n_tokens >= self._max_model_len: + logger.warning( + "Stopping realtime audio stream: reached max_model_len (%d).", + self._max_model_len, + ) + return + next_tokens = [await self._token_queue.get() for _ in range(num_tokens)] audio_arrays: list[np.ndarray] = ( @@ -198,7 +233,7 @@ class VoxtralRealtimeBuffer: @MULTIMODAL_REGISTRY.register_processor( VoxtralRealtimeMultiModalProcessor, - info=VoxtralProcessingInfo, + info=VoxtralRealtimeProcessingInfo, dummy_inputs=VoxtralDummyInputsBuilder, ) @support_torch_compile @@ -248,7 +283,9 @@ class VoxtralRealtimeGeneration(VoxtralForConditionalGeneration, SupportsRealtim # Get left/right padding audio left_pad, right_pad = audio_encoder.get_padding_audio() - buffer = VoxtralRealtimeBuffer(config, prompt_tokens) + buffer = VoxtralRealtimeBuffer( + config, prompt_tokens, max_model_len=model_config.max_model_len + ) # Feed audio with padding into buffer in background async def feed_audio():