[Bugfix][Model] Voxtral realtime: fix boot OOM and engine crash at max_model_len (#45022)

Signed-off-by: Damien Laine <damien.laine@gmail.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

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

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)

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():