Merge main into dev/vllm-align for vLLM 0.31 release

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/entrypoints/openai_api/test_audio_format.py

@@ -346,3 +346,19 @@ def test_chat_audio_metadata_survives_response_serialization(stream, fmt):
             assert audio_file.samplerate == 22050
     else:
         assert len(raw) == frame_count * 2
+
+
+def test_float32_pcm_fast_path_matches_soundfile():
+    from vllm_omni.entrypoints.openai.audio_utils_mixin import _float32_to_pcm16_bytes
+
+    rng = np.random.default_rng(0)
+    audio = np.concatenate(
+        [
+            rng.uniform(-1.2, 1.2, 4096),
+            np.array([-1.0, 1.0, 0.0, -0.0, 32767 / 32768, -1 / 32768, 0.5 / 32768, -0.5 / 32768]),
+        ]
+    ).astype(np.float32)
+    with BytesIO() as buffer:
+        soundfile.write(buffer, audio, 24000, format="RAW", subtype="PCM_16")
+        expected = buffer.getvalue()
+    assert _float32_to_pcm16_bytes(audio) == expected

tests/entrypoints/openai_api/test_moss_local_prompt.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The Local-v1.5 continuation prompt fast path equals the processor layout."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import _local_continuation_prompt
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Processor:
+    """Segment API of the model's processor with a deterministic tokenizer."""
+
+    def __init__(self):
+        self.model_config = SimpleNamespace(
+            n_vq=3,
+            im_start_token_id=1,
+            im_end_token_id=2,
+            audio_start_token_id=3,
+            audio_pad_token_id=1024,
+            audio_assistant_slot_token_id=5,
+        )
+        self.encoded = []
+
+    def _encode_text(self, text):
+        self.encoded.append(text)
+        return [100 + ord(c) % 50 for c in text]
+
+    def _user_prompt_prefix_ids(self):
+        return [1] + self._encode_text("user\n") + self._encode_text("<ref>")
+
+    def _user_prompt_after_reference_ids(self, language, fields):
+        return self._encode_text(f"lang={fields.get('language', language)}")
+
+    def _assistant_prompt_prefix_ids(self):
+        return self._encode_text("</user>") + [2] + self._encode_text("\n") + [1] + self._encode_text("assistant")
+
+    def reference(self, text, language, codes):
+        # Mirrors the processor's _build_continuation_codes + _pad (one item).
+        fields = {} if language is None else {"language": language}
+        ids = (
+            self._user_prompt_prefix_ids()
+            + self._encode_text("None")
+            + self._user_prompt_after_reference_ids(language, fields)
+            + self._encode_text(text)
+            + self._assistant_prompt_prefix_ids()
+            + [3]
+        )
+        text_rows = torch.full((len(ids), 4), 1024, dtype=torch.long)
+        text_rows[:, 0] = torch.tensor(ids)
+        audio_rows = torch.full((codes.shape[0], 4), 1024, dtype=torch.long)
+        audio_rows[:, 0] = 5
+        audio_rows[:, 1:] = codes
+        unified = torch.cat([text_rows, audio_rows])
+        return unified[:, 0].tolist(), unified[:, 1:].contiguous()
+
+
+@pytest.mark.parametrize("language", [None, "English"])
+def test_matches_processor_and_caches_constant_segments(language):
+    proc = _Processor()
+    codes = torch.randint(0, 1024, (7, 3), dtype=torch.int32)
+    kwargs = {"text": "ref words target words"}
+    if language is not None:
+        kwargs["language"] = language
+    ids, audio = _local_continuation_prompt(proc, kwargs, [codes])
+    expected_ids, expected_audio = proc.reference(kwargs["text"], language, codes.long())
+    assert ids == expected_ids
+    assert audio.dtype == torch.int64 and torch.equal(audio, expected_audio)
+
+    proc.encoded.clear()
+    _local_continuation_prompt(proc, {**kwargs, "text": "other"}, [codes])
+    assert proc.encoded == ["other"]
+
+
+def test_unsupported_layouts_fall_back():
+    proc = _Processor()
+    codes = torch.zeros((2, 3), dtype=torch.long)
+    assert _local_continuation_prompt(proc, {"text": "t", "instruction": "x"}, [codes]) is None
+    assert _local_continuation_prompt(proc, {"text": "t"}, [codes, codes]) is None
+    assert _local_continuation_prompt(proc, {"text": "t"}, ["path.wav"]) is None
+    assert _local_continuation_prompt(proc, {"text": "t"}, [torch.zeros((2, 4), dtype=torch.long)]) is None

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1105,6 +1105,38 @@ class TestTTSMethods:
         assert isinstance(audio_obj, CreateAudio)
         assert audio_obj.speed == 1.0
 
+    @pytest.mark.asyncio
+    async def test_audio_synthesis_error_flag_raises_tts_generation_error(
+        self,
+        speech_server,
+        mocker: MockerFixture,
+    ):
+        """A model-flagged synthesis failure surfaces as a non-retryable
+        TTSGenerationError, so create_speech answers 500 instead of shipping a
+        zero-length WAV or crashing on an unpacked Response."""
+
+        async def mock_generate():
+            yield create_mock_audio_output_for_test()
+
+        mocker.patch.object(
+            speech_server,
+            "_prepare_speech_generation",
+            new=mocker.AsyncMock(return_value=("speech-synth-err", mock_generate(), {})),
+        )
+        adapter = mocker.MagicMock()
+        adapter.collect_response_metadata = lambda _audio_output, collect: collect.__setitem__(
+            "audio_synthesis_error", True
+        )
+        mocker.patch.object(speech_server, "_get_tts_adapter", return_value=adapter)
+
+        with pytest.raises(TTSGenerationError, match="failed to synthesize audio") as exc_info:
+            await speech_server._generate_audio_bytes(
+                OpenAICreateSpeechRequest(input="Hello"),
+                collect={},
+            )
+
+        assert exc_info.value.retryable is False
+
     def test_is_tts_detection_with_tts_stage(self, mocker: MockerFixture):
         """Test TTS model detection when TTS stage exists."""
         mock_engine_client = mocker.MagicMock()

tests/entrypoints/openai_api/test_serving_video_stream.py

@@ -7,10 +7,12 @@ from __future__ import annotations
 import asyncio
 import base64
 import gc
+import hashlib
 import io
 import json
 import threading
 import weakref
+from types import SimpleNamespace
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock
 
@@ -26,7 +28,9 @@ from vllm_omni.entrypoints.openai.serving_video_stream import (
     QwenOmniStreamingVideoHandler,
     StreamingVideoSessionConfig,
 )
-from vllm_omni.entrypoints.openai.video_stream_base import OmniStreamingVideoHandler
+from vllm_omni.entrypoints.openai.video_stream_base import (
+    OmniStreamingVideoHandler,
+)
 from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -506,17 +510,9 @@ async def test_video_frames_consumed_excludes_bad_frame_in_query_snapshot(
         pytest.param(2, ["f0", "f1", "f2"], True, [0, 2], id="repeated-image"),
     ],
 )
-async def test_video_frame_selection_preserves_uncached_frames(
-    monkeypatch, num_frames, frame_ids, repeat_image, expected_indices
-):
-    """Prewarm misses still use image URLs, with metadata paired by buffer position."""
+async def test_video_frame_selection_pairs_metadata_by_position(num_frames, frame_ids, repeat_image, expected_indices):
+    """Queries wait for prewarmed PIL frames, with metadata paired by buffer position."""
     engine_prompts = []
-    pending_decode = asyncio.Event()
-
-    async def blocked_to_thread(*args, **kwargs):
-        await pending_decode.wait()
-
-    monkeypatch.setattr(asyncio, "to_thread", blocked_to_thread)
 
     class OneOutputEngine:
         async def generate(self, **kwargs):
@@ -550,13 +546,13 @@ async def test_video_frame_selection_preserves_uncached_frames(
     assert not [message for message in ws.sent if message["type"] == "error"]
     assert len(engine_prompts) == 1
     content = engine_prompts[0]["messages"][-1]["content"]
-    assert content == [
-        *[
-            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{frames[index]}"}}
-            for index in expected_indices
-        ],
-        {"type": "text", "text": "describe"},
-    ]
+    assert content[-1] == {"type": "text", "text": "describe"}
+    image_parts = content[:-1]
+    assert [part["type"] for part in image_parts] == ["image_pil"] * len(expected_indices)
+    for part, index in zip(image_parts, expected_indices):
+        raw = base64.b64decode(frames[index])
+        assert part["image_pil"].tobytes() == Image.open(io.BytesIO(raw)).convert("RGB").tobytes()
+        assert part["uuid"] == hashlib.md5(raw, usedforsecurity=False).hexdigest()
     consumed_events = [message for message in ws.sent if message["type"] == "video.frames.consumed"]
     if not any(frame_ids):
         assert consumed_events == []
@@ -909,6 +905,9 @@ async def test_interrupted_queries_wait_for_abort_without_fixed_delay(monkeypatc
 
     class BlockingEngine:
         async def generate(self, *, request_id, **kwargs):
+            submitted = kwargs.get("submitted")
+            if submitted is not None and not submitted.done():
+                submitted.set_result(None)
             if request_ids:
                 assert aborted[-1] == request_ids[-1]
             request_ids.append(request_id)
@@ -1004,7 +1003,7 @@ async def test_video_done_waits_for_in_flight_query():
 
 
 @pytest.mark.asyncio
-async def test_frame_prewarm_does_not_block_following_query(monkeypatch):
+async def test_query_waits_for_frame_prewarm(monkeypatch):
     decode_started = threading.Event()
     release_decode = threading.Event()
     query_started = asyncio.Event()
@@ -1035,14 +1034,80 @@ async def test_frame_prewarm_does_not_block_following_query(monkeypatch):
     assert decode_started.is_set()
 
     ws.put({"type": "video.query", "text": "describe"})
-    await asyncio.wait_for(query_started.wait(), timeout=2.0)
+    await asyncio.sleep(0.05)
+    assert not query_started.is_set()
 
     release_decode.set()
+    await asyncio.wait_for(query_started.wait(), timeout=2.0)
     ws.put({"type": "video.done"})
     await asyncio.wait_for(task, timeout=2.0)
     assert "session.done" in ws.sent_types()
 
 
+@pytest.mark.asyncio
+async def test_legacy_query_does_not_wait_for_unselected_frame(monkeypatch):
+    """APC off and num_frames=1: a blocked older frame must not delay the query."""
+    frame_a = _make_jpeg(1, 0, 0)
+    frame_b = _make_jpeg(0, 1, 0)
+    decode_a_started = threading.Event()
+    release_a = threading.Event()
+    decode_b_done = threading.Event()
+    query_started = asyncio.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        if raw_bytes == frame_a:
+            decode_a_started.set()
+            release_a.wait(timeout=5.0)
+        image = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+        if raw_bytes == frame_b:
+            decode_b_done.set()
+        return image
+
+    class RecordingHandler(QwenOmniStreamingVideoHandler):
+        async def _process_query(self, *args, **kwargs):
+            query_started.set()
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    ws = TimedWebSocket()
+    handler = RecordingHandler(chat_service=object(), engine_client=object(), idle_timeout=5.0)
+    task = asyncio.create_task(handler.handle_session(ws))
+    try:
+        ws.put(
+            {
+                "type": "session.config",
+                "model": "test",
+                "modalities": ["text"],
+                "num_frames": 1,
+                "enable_frame_filter": False,
+            }
+        )
+        await asyncio.sleep(0)
+        ws.put({"type": "video.frame", "data": _b64(frame_a), "frame_id": "A"})
+        for _ in range(100):
+            if decode_a_started.is_set():
+                break
+            await asyncio.sleep(0.01)
+        assert decode_a_started.is_set()
+        ws.put({"type": "video.frame", "data": _b64(frame_b), "frame_id": "B"})
+        for _ in range(200):
+            if decode_b_done.is_set():
+                break
+            await asyncio.sleep(0.01)
+        assert decode_b_done.is_set()
+
+        ws.put({"type": "video.query", "text": "describe"})
+        await asyncio.wait_for(query_started.wait(), timeout=2.0)
+        assert not release_a.is_set()
+        release_a.set()
+        ws.put({"type": "video.done"})
+        await asyncio.wait_for(task, timeout=2.0)
+    finally:
+        release_a.set()
+        if not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("max_frames", [1, 2], ids=["evicted", "retained"])
 async def test_frame_prewarm_only_keeps_retained_images(monkeypatch, max_frames):
@@ -1272,3 +1337,452 @@ def test_build_messages_keeps_recent_history_text_only():
     assert messages[1] == {"role": "assistant", "content": "recent answer"}
     assert messages[2] == user_message
     assert user_message["content"][-1] == {"type": "text", "text": "current question"}
+
+
+# ---------------------------------------------------------------------------
+# Incremental prefill (warmup) lifecycle
+# ---------------------------------------------------------------------------
+
+
+class RecordingReuseEngine:
+    """Engine stub with stage-0 prefix caching on; records generate() and abort()."""
+
+    def __init__(self, warmup_delay: float = 0.0, query_delay: float = 0.0):
+        self.requests: list[dict[str, Any]] = []
+        self.aborted: list[str] = []
+        self._warmup_delay = warmup_delay
+        self._query_delay = query_delay
+        self.engine = SimpleNamespace(
+            stage_vllm_configs=[SimpleNamespace(cache_config=SimpleNamespace(enable_prefix_caching=True))]
+        )
+
+    def warmup_ids(self) -> list[str]:
+        return [r["request_id"] for r in self.requests if r["request_id"].startswith("video-warmup-")]
+
+    def query_ids(self) -> list[str]:
+        return [r["request_id"] for r in self.requests if not r["request_id"].startswith("video-warmup-")]
+
+    def generate(self, **kwargs):
+        self.requests.append(kwargs)
+        submitted = kwargs.get("submitted")
+        if submitted is not None and not submitted.done():
+            submitted.set_result(None)
+        is_warmup = kwargs["request_id"].startswith("video-warmup-")
+        delay = self._warmup_delay if is_warmup else self._query_delay
+
+        async def _gen():
+            if delay:
+                await asyncio.sleep(delay)
+            yield _text_result("ok")
+
+        return _gen()
+
+    async def abort(self, request_id: str) -> None:
+        self.aborted.append(request_id)
+
+
+class PromptRecordingHandler(QwenOmniStreamingVideoHandler):
+    """Records each query turn's user-message content and processor kwargs."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.query_contents: list[list[dict[str, Any]]] = []
+        self.warmup_contents: list[list[dict[str, Any]]] = []
+        self.mm_processor_kwargs: list[Any] = []
+
+    async def _preprocess_to_engine_prompt(self, request):
+        self.mm_processor_kwargs.append(getattr(request, "mm_processor_kwargs", None))
+        return {"prompt": "p"}
+
+    def build_engine_prompt(
+        self, config, frame_buffer, audio_buffer, message_history, query_text, prewarmed_frames, *, frame_indices=None
+    ):
+        messages, user_message = super().build_engine_prompt(
+            config,
+            frame_buffer,
+            audio_buffer,
+            message_history,
+            query_text,
+            prewarmed_frames,
+            frame_indices=frame_indices,
+        )
+        self.query_contents.append(list(user_message["content"]))
+        return messages, user_message
+
+    def build_engine_prompt_prefix(self, config, frame_buffer, message_history, prewarmed_frames):
+        messages = super().build_engine_prompt_prefix(config, frame_buffer, message_history, prewarmed_frames)
+        if messages:
+            self.warmup_contents.append(list(messages[-1]["content"]))
+        return messages
+
+
+async def _poll(predicate, timeout: float = 3.0) -> bool:
+    for _ in range(int(timeout / 0.01)):
+        if predicate():
+            return True
+        await asyncio.sleep(0.01)
+    return bool(predicate())
+
+
+@pytest.mark.asyncio
+async def test_warmup_fires_on_frame_and_query_cancels_it():
+    engine = RecordingReuseEngine(warmup_delay=10.0, query_delay=0.3)
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+
+    assert await _poll(lambda: engine.warmup_ids())
+    warmup_id = engine.warmup_ids()[0]
+    assert engine.requests[0]["sampling_params"].max_tokens == 1
+
+    ws.put({"type": "video.query", "text": "describe"})
+    assert await _poll(lambda: engine.query_ids())
+    # in-flight warmup was cancelled + aborted, and not restarted mid-query
+    assert warmup_id in engine.aborted
+    assert len(engine.warmup_ids()) == 1
+    assert handler.mm_processor_kwargs[:2] == [
+        {"use_audio_in_video": True},
+        {"use_audio_in_video": True},
+    ]
+
+    assert await _poll(lambda: "response.text.done" in ws.sent_types())
+    # the committed turn changed the context -> a fresh warmup fires
+    assert await _poll(lambda: len(engine.warmup_ids()) == 2)
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_cancelled_warmup_reaches_stage0_before_abort():
+    """Cancel during submit must not abort a warmup the engine never received."""
+
+    class _HoldSubmitEngine(RecordingReuseEngine):
+        def __init__(self):
+            super().__init__()
+            self.started = asyncio.Event()
+            self.release_submit = asyncio.Event()
+
+        def generate(self, **kwargs):
+            submitted = kwargs.get("submitted")
+
+            async def _gen():
+                self.started.set()
+                await self.release_submit.wait()
+                self.requests.append(kwargs)
+                if submitted is not None and not submitted.done():
+                    submitted.set_result(None)
+                yield _text_result("ok")
+
+            return _gen()
+
+    engine = _HoldSubmitEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    assert await _poll(engine.started.is_set)
+    ws.put({"type": "video.query", "text": "describe"})
+    await asyncio.sleep(0.05)
+    assert engine.requests == []
+    assert engine.aborted == []
+    engine.release_submit.set()
+    assert await _poll(lambda: engine.aborted)
+    assert engine.warmup_ids()
+    assert engine.aborted[0] == engine.warmup_ids()[0]
+    assert await _poll(lambda: engine.query_ids())
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_cancelled_query_reaches_stage0_before_abort():
+    """Cancel during query submit must not abort a request the engine never received."""
+
+    class _HoldQueryEngine(RecordingReuseEngine):
+        def __init__(self):
+            super().__init__()
+            self.query_started = asyncio.Event()
+            self.release_submit = asyncio.Event()
+
+        def generate(self, **kwargs):
+            if str(kwargs.get("request_id", "")).startswith("video-warmup-"):
+                return super().generate(**kwargs)
+            submitted = kwargs.get("submitted")
+
+            async def _gen():
+                self.query_started.set()
+                await self.release_submit.wait()
+                self.requests.append(kwargs)
+                if submitted is not None and not submitted.done():
+                    submitted.set_result(None)
+                yield _text_result("ok")
+
+            return _gen()
+
+    engine = _HoldQueryEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    assert await _poll(lambda: engine.warmup_ids())
+    ws.put({"type": "video.query", "text": "first"})
+    assert await _poll(engine.query_started.is_set)
+    ws.put({"type": "video.query", "text": "second"})
+    await asyncio.sleep(0.05)
+    assert engine.query_ids() == []
+    assert engine.aborted == []
+    engine.release_submit.set()
+    assert await _poll(lambda: engine.query_ids())
+    first_query = engine.query_ids()[0]
+    assert await _poll(lambda: first_query in engine.aborted)
+    assert await _poll(lambda: len(engine.query_ids()) == 2)
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_audio_session_never_starts_warmup():
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text", "audio"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg(9, 9, 9))})
+    await asyncio.sleep(0.3)
+
+    assert engine.warmup_ids() == []
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_warmup_starts_only_after_frames_ready(monkeypatch):
+    decode_started = threading.Event()
+    release_decode = threading.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        decode_started.set()
+        release_decode.wait(timeout=5.0)
+        return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
+    assert await _poll(decode_started.is_set)
+    await asyncio.sleep(0.1)
+    assert engine.warmup_ids() == []
+
+    release_decode.set()
+    assert await _poll(lambda: engine.warmup_ids())
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_query_cancelling_warmup_keeps_prewarm_alive(monkeypatch):
+    """Cancel warmup during prefill must not drop a later frame still decoding."""
+    frame1 = _b64(_make_jpeg(1, 1, 1))
+    frame2 = _b64(_make_jpeg(2, 2, 2))
+    raw2 = base64.b64decode(frame2)
+    decode_started = threading.Event()
+    release_decode = threading.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        if raw_bytes == raw2:
+            decode_started.set()
+            release_decode.wait(timeout=5.0)
+        return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    engine = RecordingReuseEngine(warmup_delay=10.0)
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False})
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": frame1})
+    assert await _poll(lambda: engine.warmup_ids())
+
+    ws.put({"type": "video.frame", "data": frame2})
+    assert await _poll(decode_started.is_set)
+
+    ws.put({"type": "video.query", "text": "describe"})
+    await asyncio.sleep(0.1)
+    release_decode.set()
+
+    assert await _poll(lambda: "response.text.done" in ws.sent_types())
+    assert len(handler.query_contents) == 1
+    image_parts = [p for p in handler.query_contents[0] if p.get("type") == "image_pil"]
+    assert len(image_parts) == 2
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_query_keeps_frames_evicted_while_pinned(monkeypatch):
+    """Regression: frames FIFO-evicted while an in-flight query awaits prewarm
+    must keep their PIL entries and stay in the query prompt."""
+    frame_a = _b64(_make_jpeg(1, 1, 1))
+    frame_a2 = _b64(_make_jpeg(2, 2, 2))
+    frame_b = _b64(_make_jpeg(3, 3, 3))
+    raw_a = base64.b64decode(frame_a)
+    decoded: list[bytes] = []
+    release_others = threading.Event()
+
+    def gated_decode(raw_bytes: bytes):
+        if raw_bytes != raw_a:
+            release_others.wait(timeout=5.0)
+        decoded.append(raw_bytes)
+        return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+
+    monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode)
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": ["text"],
+            "max_frames": 2,
+            "enable_frame_filter": False,
+        }
+    )
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": frame_a})
+    assert await _poll(lambda: raw_a in decoded)
+    ws.put({"type": "video.frame", "data": frame_a2})  # decode gated -> query will wait
+
+    ws.put({"type": "video.query", "text": "describe"})  # pins [a, a2]
+    await asyncio.sleep(0.1)
+    ws.put({"type": "video.frame", "data": frame_b})  # buffer full -> evicts pinned a
+    await asyncio.sleep(0.1)
+    release_others.set()
+
+    assert await _poll(lambda: "response.text.done" in ws.sent_types())
+    assert len(handler.query_contents) == 1
+    uuids = [p["uuid"] for p in handler.query_contents[0] if p.get("type") == "image_pil"]
+    expected = [hashlib.md5(base64.b64decode(f), usedforsecurity=False).hexdigest() for f in (frame_a, frame_a2)]
+    assert uuids == expected
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+
+
+async def _query_image_uuids(modalities: list[str], frames: list[str], *, max_frames: int) -> list[str]:
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": modalities,
+            "max_frames": max_frames,
+            "num_frames": max_frames,
+            "enable_frame_filter": False,
+        }
+    )
+    await asyncio.sleep(0)
+    for frame in frames:
+        ws.put({"type": "video.frame", "data": frame})
+    ws.put({"type": "video.query", "text": "describe"})
+    assert await _poll(lambda: handler.query_contents)
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+    return [part["uuid"] for part in handler.query_contents[0] if part.get("type") == "image_pil"]
+
+
+def _frame_uuid(frame_b64: str) -> str:
+    return hashlib.md5(base64.b64decode(frame_b64), usedforsecurity=False).hexdigest()
+
+
+@pytest.mark.asyncio
+async def test_below_max_frames_incremental_matches_legacy_window():
+    frames = [_b64(_make_jpeg(i, 0, 0)) for i in range(3)]
+    expected = [_frame_uuid(frame) for frame in frames]
+    assert await _query_image_uuids(["text"], frames, max_frames=4) == expected
+    assert await _query_image_uuids(["text", "audio"], frames, max_frames=4) == expected
+
+
+@pytest.mark.asyncio
+async def test_incremental_full_window_drops_oldest_half():
+    frames = [_b64(_make_jpeg(i, 1, 0)) for i in range(5)]
+    kept = await _query_image_uuids(["text"], frames, max_frames=4)
+    legacy = await _query_image_uuids(["text", "audio"], frames, max_frames=4)
+    assert kept == [_frame_uuid(frame) for frame in frames[2:]]
+    assert legacy == [_frame_uuid(frame) for frame in frames[1:]]
+
+
+async def _run_frame_audio_query(modalities: list[str], frame: str, pcm: bytes) -> PromptRecordingHandler:
+    engine = RecordingReuseEngine()
+    handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0)
+    ws = TimedWebSocket()
+    task = asyncio.create_task(handler.handle_session(ws))
+    ws.put(
+        {
+            "type": "session.config",
+            "model": "test",
+            "modalities": modalities,
+            "enable_frame_filter": False,
+        }
+    )
+    await asyncio.sleep(0)
+    ws.put({"type": "video.frame", "data": frame})
+    if "audio" not in modalities:
+        assert await _poll(lambda: engine.warmup_ids())
+    else:
+        await asyncio.sleep(0.05)
+        assert engine.warmup_ids() == []
+    ws.put({"type": "audio.chunk", "data": _b64(pcm)})
+    ws.put({"type": "video.query", "text": "describe"})
+    assert await _poll(lambda: engine.query_ids())
+    ws.put({"type": "video.done"})
+    await asyncio.wait_for(task, timeout=2.0)
+    return handler
+
+
+@pytest.mark.asyncio
+async def test_reuse_query_with_audio_matches_legacy_prompt():
+    """Incremental and legacy query prompts match; both set the processor kwarg."""
+    frame = _b64(_make_jpeg())
+    pcm = b"\x01\x00\x02\x00"
+    reuse = await _run_frame_audio_query(["text"], frame, pcm)
+    legacy = await _run_frame_audio_query(["text", "audio"], frame, pcm)
+
+    assert reuse.warmup_contents
+    assert all(part.get("type") != "input_audio" for part in reuse.warmup_contents[0])
+    assert [part.get("type") for part in reuse.warmup_contents[0]] == ["image_pil"]
+
+    assert reuse.query_contents and legacy.query_contents
+    reuse_q = reuse.query_contents[0]
+    legacy_q = legacy.query_contents[0]
+    assert [part.get("type") for part in reuse_q] == ["image_pil", "input_audio", "text"]
+    assert [part.get("type") for part in legacy_q] == ["image_pil", "input_audio", "text"]
+    assert reuse_q[0]["uuid"] == legacy_q[0]["uuid"]
+    assert reuse_q[1] == legacy_q[1]
+    assert reuse_q[2] == legacy_q[2] == {"type": "text", "text": "describe"}
+
+    assert reuse.mm_processor_kwargs[0] == {"use_audio_in_video": True}
+    assert reuse.mm_processor_kwargs[1] == {"use_audio_in_video": True}
+    assert legacy.mm_processor_kwargs[0] == {"use_audio_in_video": True}

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -691,6 +691,45 @@ def test_moss_get_processor_places_audio_tokenizer_on_codec_stage_device(mocker)
     audio_tokenizer.to.assert_called_once()
 
 
+@pytest.mark.parametrize(
+    "variant,graphs_env,prepares",
+    [("local", None, False), ("local", "1", True), ("local", "0", False), ("realtime", "1", False)],
+)
+def test_moss_warmup_prepares_the_reference_encoder(variant, graphs_env, prepares, mocker, monkeypatch):
+    if graphs_env is None:
+        monkeypatch.delenv("VLLM_OMNI_MOSS_REF_GRAPHS", raising=False)
+    else:
+        monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", graphs_env)
+    adapter = _moss_adapter_with_stage_configs([], mocker)
+    adapter._moss_variant = variant
+    encoder = mocker.Mock()
+    mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
+    asyncio.run(adapter.warmup())
+    assert encoder.prepare.called is prepares
+
+
+def test_moss_warmup_failure_does_not_stop_startup(mocker, monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1")
+    adapter = _moss_adapter_with_stage_configs([], mocker)
+    encoder = mocker.Mock()
+    encoder.prepare.side_effect = RuntimeError("capture failed")
+    mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
+    asyncio.run(adapter.warmup())
+    encoder.prepare.assert_called_once_with()
+
+
+def test_moss_shared_encoder_startup_failure_stops_api_startup(mocker, monkeypatch):
+    monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1")
+    from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import SharedReferenceEncoderStartupError
+
+    adapter = _moss_adapter_with_stage_configs([], mocker)
+    encoder = mocker.Mock()
+    encoder.prepare.side_effect = SharedReferenceEncoderStartupError("still compiling")
+    mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
+    with pytest.raises(SharedReferenceEncoderStartupError):
+        asyncio.run(adapter.warmup())
+
+
 def test_qwen3_tts_metadata():
     assert Qwen3TTSAdapter.backend == "ar"
     assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)

tests/entrypoints/openai_api/test_tts_detection.py

@@ -183,6 +183,7 @@ _PIPELINE_STAGES = [
     "token2wav",
     "tts",
     "encoder",
+    "yue2",
 ]
 
 _STAGES = [*_PIPELINE_STAGES, None, "vae", "not_a_real_stage"]

tests/entrypoints/openai_api/test_video_audio_delta.py

@@ -174,7 +174,9 @@ async def test_real_producer_deltas_reach_video_client(
         produced.append(OmniRequestOutput.from_stage_output(result, final_output_type="audio"))
 
     class Engine:
-        async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None):
+        async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None, submitted=None):
+            if submitted is not None and not submitted.done():
+                submitted.set_result(None)
             if explicit_sampling:
                 assert sampling_params_list is not None
                 assert sampling_params_list[0].output_kind == RequestOutputKind.DELTA

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -11,6 +11,7 @@ from pydantic import ValidationError
 from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig, extract_diffusion_stage_config_kwargs
 from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config
+from vllm_omni.config.stage_config import build_stage_runtime_overrides
 from vllm_omni.diffusion.data import AttentionConfig, OmniDiffusionConfig
 from vllm_omni.engine import stage_init_utils
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
@@ -945,3 +946,33 @@ def test_generic_diffusion_structured_stage_reaches_standard_startup(mocker):
     assert launched["model"] == "generic-diffusion"
     assert launched["stage_id"] == 0
     assert runtime.stage_pools[0].clients == [client]
+
+
+def test_default_stage_config_preserves_pre_sharded_hsdp_weight_load() -> None:
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"hsdp_weight_load_strategy": "pre_sharded"})[0]
+
+    assert stage_cfg["engine_args"]["hsdp_weight_load_strategy"] == "pre_sharded"
+
+
+@pytest.mark.parametrize("prefix", ["", "stage_0_"])
+def test_pre_sharded_hsdp_weight_load_survives_stage_overrides(prefix: str) -> None:
+    overrides = build_stage_runtime_overrides(0, {f"{prefix}hsdp_weight_load_strategy": "pre_sharded"})
+
+    assert overrides["hsdp_weight_load_strategy"] == "pre_sharded"
+
+
+@pytest.mark.parametrize("strategy", ["full", "pre_sharded"])
+def test_typed_default_stage_preserves_hsdp_weight_load_strategy(strategy: str) -> None:
+    config = StageConfigFactory.create_typed_default_diffusion(
+        "generic-diffusion", {"hsdp_weight_load_strategy": strategy}
+    )
+
+    assert config.stage_configs[0].diffusion_config.hsdp_weight_load_strategy == strategy
+
+
+@pytest.mark.parametrize("strategy", ["memory_limited", "invalid"])
+def test_removed_hsdp_loading_strategies_are_rejected(strategy):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    with pytest.raises(ValueError, match="hsdp_weight_load_strategy"):
+        OmniDiffusionConfig(hsdp_weight_load_strategy=strategy)

tests/entrypoints/test_pd_disaggregation.py

@@ -288,6 +288,9 @@ def _setup_log_mocks(monkeypatch):
         def build_and_log_summary(self):
             return "Fake summary"
 
+        def log_timing_summary(self):
+            pass
+
     monkeypatch.setattr(
         "vllm_omni.entrypoints.omni.OrchestratorAggregator",
         _FakeOrchestratorAggregator,