From 4124a1f8589b979b2833a2cff8816e81f7988436 Mon Sep 17 00:00:00 2001 From: ChenWenjing <54166744+Shirley125@users.noreply.github.com> Date: Sat, 18 Apr 2026 15:42:41 +0800 Subject: [PATCH] [Feature]Support audio streaming input and output-phase2 (#2581) Signed-off-by: CHEN <116010019@link.cuhk.edu.cn> --- README.md | 2 + examples/online_serving/qwen3_omni/README.md | 33 +- .../qwen3_omni/openai_realtime_client.py | 404 +++++++++++++----- tests/core/sched/test_omni_scheduler_mixin.py | 129 ++++++ .../test_qwen3_omni_realtime_websocket.py | 208 +++++++++ tests/engine/test_orchestrator.py | 2 +- .../test_realtime_connection_helpers.py | 86 ++++ .../test_qwen3_omni_streaming_helpers.py | 81 ++++ vllm_omni/core/sched/omni_ar_scheduler.py | 27 +- .../core/sched/omni_generation_scheduler.py | 21 +- vllm_omni/core/sched/omni_scheduler_mixin.py | 33 ++ vllm_omni/engine/__init__.py | 4 + vllm_omni/engine/orchestrator.py | 63 ++- vllm_omni/engine/stage_engine_core_client.py | 10 + vllm_omni/entrypoints/async_omni.py | 1 - vllm_omni/entrypoints/openai/api_server.py | 18 +- .../entrypoints/openai/realtime_connection.py | 193 +++++++++ .../models/qwen3_omni/qwen3_omni.py | 5 + .../stage_input_processors/qwen3_omni.py | 139 +++++- vllm_omni/patch.py | 5 +- vllm_omni/request.py | 34 ++ vllm_omni/worker/gpu_model_runner.py | 28 ++ 22 files changed, 1382 insertions(+), 144 deletions(-) create mode 100644 tests/core/sched/test_omni_scheduler_mixin.py create mode 100644 tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py create mode 100644 tests/entrypoints/test_realtime_connection_helpers.py create mode 100644 tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py create mode 100644 vllm_omni/core/sched/omni_scheduler_mixin.py create mode 100644 vllm_omni/entrypoints/openai/realtime_connection.py diff --git a/README.md b/README.md index 6726c51f8..dd53c2fb1 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,8 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn - [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/) - [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/) +**Qwen3-Omni realtime (WebSocket):** stream PCM through the OpenAI-style `/v1/realtime` endpoint using the example client and server notes in [examples/online_serving/qwen3_omni/README.md](examples/online_serving/qwen3_omni/README.md) (requires `async_chunk` disabled on the server). + ## Contributing We welcome and value any contributions and collaborations. diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index ff0264224..4c3f9595f 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/examples/online_serving/qwen3_omni/README.md @@ -18,6 +18,8 @@ If you want to open async chunking for qwen3-omni, launch the server with comman vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml ``` +**Note:** The OpenAI-style **`/v1/realtime`** WebSocket (streaming PCM audio in, audio + transcription out) is **not supported** when `async_chunk` is enabled. Use the default omni layout or a stage config with `async_chunk: false` for realtime sessions. + If you have custom stage configs file, launch the server with command below ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file @@ -38,36 +40,43 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene #### Realtime WebSocket client (`openai_realtime_client.py`) -[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://:/v1/realtime`**, uploads a local audio file as **PCM16 mono @ 16 kHz** chunks (OpenAI-style `input_audio_buffer.append` / `commit`), and prints **streaming transcription** (`transcription.delta` / `transcription.done`). +[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://:/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks (OpenAI-style `input_audio_buffer.append` / `commit`), and receives **`response.audio.delta`** (incremental PCM for the reply) plus **`transcription.*`** events. By default it concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging. + +Streaming input works well for translation-style use cases; if the Thinker runs while input is still incomplete, consider limiting **`max_tokens`** in your session / server defaults to avoid over-generation. **Dependencies:** ```bash -pip install websockets numpy +pip install websockets ``` **From this directory** (`examples/online_serving/qwen3_omni`): ```bash python openai_realtime_client.py \ - --host localhost \ - --port 8091 \ + --url ws://localhost:8091/v1/realtime \ --model Qwen/Qwen3-Omni-30B-A3B-Instruct \ - --audio_path /path/to/your.wav + --input-wav /path/to/input_16k_mono.wav \ + --output-wav realtime_output.wav \ + --delta-dump-dir ./rt_delta_wavs ``` -If `--audio_path` is omitted, the script uses a bundled default clip (`mary_had_lamb` via vLLM assets). - **Arguments:** | Flag | Default | Description | |------|---------|-------------| -| `--host` | `localhost` | API server host | -| `--port` | `8000` | API server port (match your `vllm serve` port, e.g. `8091`) | -| `--model` | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | Must match the served model (also sent in `session.update`) | -| `--audio_path` | *(optional)* | Path to input audio; resampled to 16 kHz mono inside the client | +| `--url` | `ws://localhost:8091/v1/realtime` | Full WebSocket URL including path | +| `--model` | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | Must match the served model (sent in `session.update`) | +| `--input-wav` | *(required)* | Input WAV: mono, 16-bit PCM, **16 kHz** | +| `--output-wav` | `realtime_output.wav` | Output path for concatenated reply audio | +| `--output-text` | *(optional)* | If set, write final transcription text to this path | +| `--chunk-ms` | `200` | Size of each uploaded audio chunk (milliseconds of audio) | +| `--send-delay-ms` | `0` | Delay between chunk sends (simulate realtime upload) | +| `--delta-dump-dir` | *(optional)* | Directory to write per-`response.audio.delta` WAV files | +| `--num-requests` | `1` | Number of sequential sessions (see `--concurrency`) | +| `--concurrency` | `1` | Max concurrent WebSocket sessions when `--num-requests` > 1 | -Ensure the vLLM-Omni server is running with realtime support for this endpoint, for example: +Ensure the server is running **without** `async_chunk` if you use `/v1/realtime`, for example: ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 diff --git a/examples/online_serving/qwen3_omni/openai_realtime_client.py b/examples/online_serving/qwen3_omni/openai_realtime_client.py index 660e4ac33..79e30a3f5 100644 --- a/examples/online_serving/qwen3_omni/openai_realtime_client.py +++ b/examples/online_serving/qwen3_omni/openai_realtime_client.py @@ -1,81 +1,118 @@ +"""Realtime client for vLLM-Omni /v1/realtime (audio + text events). + +This client: +1) Reads a local WAV file (must be mono, 16-bit PCM, 16kHz), +2) Streams PCM16 chunks to /v1/realtime with OpenAI-style events, +3) Receives response.audio.* and transcription.* events, +4) Saves synthesized audio to an output WAV file and optional text file. + +By default each ``response.audio.delta`` is treated as an **incremental PCM** +chunk and all chunks are concatenated into the final ``--output-wav``. + +Optional debugging: pass ``--delta-dump-dir DIR`` to write every +``response.audio.delta`` payload as ``delta_000001.wav``, ``delta_000002.wav``, … + +Usage: + python openai_realtime_client.py \ + --url ws://localhost:8091/v1/realtime \ + --model Qwen/Qwen3-Omni-30B-A3B-Instruct \ + --input-wav input_16k_mono.wav \ + --output-wav realtime_output.wav \ + --delta-dump-dir ./rt_delta_wavs + +Dependencies: + pip install websockets """ -This script demonstrates how to use the vLLM-Omni Realtime WebSocket API to perform -audio transcription by uploading an audio file. -Before running this script, you must start the vLLM-Omni server with a realtime-capable -model, for example: - - vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni - -Requirements: -- vllm with audio support -- websockets -- soundfile -- numpy - -The script: -1. Connects to the Realtime WebSocket endpoint -2. Converts an audio file to PCM16 @ 16kHz -3. Sends audio chunks to the server -4. Receives and prints transcription as it streams -""" +from __future__ import annotations import argparse import asyncio import base64 import json +import wave +from pathlib import Path -import numpy as np -import websockets -from vllm.assets.audio import AudioAsset -from vllm.multimodal.media.audio import load_audio +try: + import websockets +except ImportError: + print("Please install websockets: pip install websockets") + raise SystemExit(1) -def audio_to_pcm16_base64(audio_path: str) -> str: - """ - Load an audio file and convert it to base64-encoded PCM16 @ 16kHz. - """ - # Load audio and resample to 16kHz mono - audio, _ = load_audio(audio_path, sr=16000, mono=True) - # Convert to PCM16 - pcm16 = (audio * 32767).astype(np.int16) - # Encode as base64 - return base64.b64encode(pcm16.tobytes()).decode("utf-8") +def _read_wav_pcm16(path: Path) -> bytes: + with wave.open(str(path), "rb") as wf: + nchannels = wf.getnchannels() + sampwidth = wf.getsampwidth() + framerate = wf.getframerate() + comptype = wf.getcomptype() + nframes = wf.getnframes() + + if nchannels != 1: + raise ValueError(f"Input WAV must be mono (got {nchannels} channels).") + if sampwidth != 2: + raise ValueError(f"Input WAV must be 16-bit PCM (got sample width={sampwidth}).") + if framerate != 16000: + raise ValueError(f"Input WAV must be 16kHz (got {framerate} Hz).") + if comptype != "NONE": + raise ValueError(f"Input WAV must be uncompressed PCM (got comptype={comptype}).") + if nframes <= 0: + raise ValueError("Input WAV has no audio frames.") + + return wf.readframes(nframes) -async def realtime_transcribe(audio_path: str, host: str, port: int, model: str): - """ - Connect to the Realtime API and transcribe an audio file. - """ - uri = f"ws://{host}:{port}/v1/realtime" +def _write_wav_pcm16(path: Path, pcm16_bytes: bytes, sample_rate_hz: int) -> None: + with wave.open(str(path), "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate_hz) + wf.writeframes(pcm16_bytes) - async with websockets.connect(uri) as ws: - # Wait for session.created - response = json.loads(await ws.recv()) - if response["type"] == "session.created": - print(f"Session created: {response['id']}") - else: - print(f"Unexpected response: {response}") - return - # Validate model - await ws.send(json.dumps({"type": "session.update", "model": model})) +async def run_client( + url: str, + model: str, + input_wav: Path, + output_wav: Path, + output_text: Path | None, + chunk_ms: int, + send_delay_ms: int, + delta_dump_dir: Path | None, + request_idx: int = 1, + total_requests: int = 1, +) -> None: + log_prefix = f"[req {request_idx:02d}/{total_requests:02d}] " if total_requests > 1 else "" + pcm16 = _read_wav_pcm16(input_wav) + bytes_per_ms = 16000 * 2 // 1000 # mono PCM16 at 16kHz + chunk_bytes = max(bytes_per_ms * chunk_ms, 2) - # Signal ready to start - await ws.send(json.dumps({"type": "input_audio_buffer.commit"})) + incremental_pcm_parts: list[bytes] = [] + output_sample_rate = 24000 + delta_index = 0 + text_chunks: list[str] = [] + final_text: str = "" - # Convert audio file to base64 PCM16 - print(f"Loading audio from: {audio_path}") - audio_base64 = audio_to_pcm16_base64(audio_path) + if delta_dump_dir is not None: + delta_dump_dir.mkdir(parents=True, exist_ok=True) - # Send audio in chunks (4KB of raw audio = ~8KB base64) - chunk_size = 4096 - audio_bytes = base64.b64decode(audio_base64) - total_chunks = (len(audio_bytes) + chunk_size - 1) // chunk_size + async with websockets.connect(url, max_size=64 * 1024 * 1024) as ws: + # 1) Validate model. + await ws.send( + json.dumps( + { + "type": "session.update", + "model": model, + } + ) + ) - print(f"Sending {total_chunks} audio chunks...") - for i in range(0, len(audio_bytes), chunk_size): - chunk = audio_bytes[i : i + chunk_size] + # 2) Start generation once (non-final commit). + await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False})) + + # 3) Stream audio chunks. + for i in range(0, len(pcm16), chunk_bytes): + chunk = pcm16[i : i + chunk_bytes] await ws.send( json.dumps( { @@ -84,63 +121,212 @@ async def realtime_transcribe(audio_path: str, host: str, port: int, model: str) } ) ) + if send_delay_ms > 0: + await asyncio.sleep(send_delay_ms / 1000.0) - # Signal all audio is sent + # 4) Final commit closes input stream. await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True})) - print("Audio sent. Waiting for transcription...\n") - # Receive transcription - print("Transcription: ", end="", flush=True) + # 5) Receive server events until audio done. while True: - response = json.loads(await ws.recv()) - if response["type"] == "transcription.delta": - print(response["delta"], end="", flush=True) - elif response["type"] == "transcription.done": - print(f"\n\nFinal transcription: {response['text']}") - if response.get("usage"): - print(f"Usage: {response['usage']}") - break - elif response["type"] == "error": - print(f"\nError: {response['error']}") + message = await ws.recv() + if isinstance(message, bytes): + # We only expect JSON text frames. + continue + + event = json.loads(message) + event_type = event.get("type") + + if event_type == "session.created": + continue + + if event_type == "response.audio.delta": + sr = event.get("sample_rate_hz") + if isinstance(sr, int) and sr > 0: + output_sample_rate = sr + audio_b64 = event.get("audio", "") + if audio_b64: + pcm_delta = base64.b64decode(audio_b64) + incremental_pcm_parts.append(pcm_delta) + if delta_dump_dir is not None and pcm_delta: + delta_index += 1 + dump_path = delta_dump_dir / f"delta_{delta_index:06d}.wav" + _write_wav_pcm16(dump_path, pcm_delta, output_sample_rate) + print( + f"{log_prefix}delta dump #{delta_index}: {dump_path} " + f"(pcm bytes={len(pcm_delta)}, sr={output_sample_rate})" + ) + continue + + if event_type == "transcription.delta": + delta = event.get("delta", "") + if delta: + text_chunks.append(delta) + print(delta, end="", flush=True) + continue + + if event_type == "transcription.done": + final_text = event.get("text", "") or "".join(text_chunks) + usage = event.get("usage") + final_text_with_tag = f"Final transcription: {final_text}" + if text_chunks: + print() + print(f"{log_prefix}{final_text_with_tag}") + if usage: + print(f"{log_prefix}text usage: {usage}") + continue + + if event_type == "response.audio.done": break + if event_type == "error": + raise RuntimeError(f"Server error: {event}") -def main(args): - if args.audio_path: - audio_path = args.audio_path + all_pcm16 = b"".join(incremental_pcm_parts) + if not all_pcm16: + raise RuntimeError("No audio received from server.") + + output_wav.parent.mkdir(parents=True, exist_ok=True) + _write_wav_pcm16(output_wav, all_pcm16, output_sample_rate) + print(f"{log_prefix}Saved realtime audio to: {output_wav} (incremental chunks joined)") + + if output_text is not None: + text_to_save = final_text if final_text else "".join(text_chunks) + output_text.parent.mkdir(parents=True, exist_ok=True) + output_text.write_text(text_to_save, encoding="utf-8") + print(f"{log_prefix}Saved realtime text to: {output_text}") + + +def _indexed_output_path(path: Path | None, index: int, total: int) -> Path | None: + if path is None or total <= 1: + return path + return path.with_name(f"{path.stem}_{index:02d}{path.suffix}") + + +async def run_clients_concurrent( + *, + url: str, + model: str, + input_wav: Path, + output_wav: Path, + output_text: Path | None, + chunk_ms: int, + send_delay_ms: int, + delta_dump_dir: Path | None, + num_requests: int, + concurrency: int, +) -> None: + sem = asyncio.Semaphore(concurrency) + + async def _run_one(index: int) -> tuple[int, bool, str | None]: + per_output_wav = _indexed_output_path(output_wav, index, num_requests) + per_output_text = _indexed_output_path(output_text, index, num_requests) + per_delta_dir = None + if delta_dump_dir is not None: + per_delta_dir = delta_dump_dir / f"req_{index:02d}" + async with sem: + try: + await run_client( + url=url, + model=model, + input_wav=input_wav, + output_wav=per_output_wav, + output_text=per_output_text, + chunk_ms=chunk_ms, + send_delay_ms=send_delay_ms, + delta_dump_dir=per_delta_dir, + request_idx=index, + total_requests=num_requests, + ) + return index, True, None + except Exception as exc: + return index, False, str(exc) + + tasks = [asyncio.create_task(_run_one(i), name=f"rt-client-{i}") for i in range(1, num_requests + 1)] + results = await asyncio.gather(*tasks) + + failed = [(idx, err) for idx, ok, err in results if not ok] + succeeded = num_requests - len(failed) + print(f"[summary] succeeded={succeeded}, failed={len(failed)}, total={num_requests}") + if failed: + for idx, err in failed: + print(f"[summary] req {idx:02d} failed: {err}") + raise RuntimeError(f"{len(failed)} concurrent request(s) failed") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Realtime audio/text client for vLLM-Omni") + parser.add_argument("--url", default="ws://localhost:8091/v1/realtime", help="WebSocket URL") + parser.add_argument( + "--model", + default="Qwen/Qwen3-Omni-30B-A3B-Instruct", + help="Model name for session.update", + ) + parser.add_argument("--input-wav", required=True, type=Path, help="Input WAV (mono, PCM16, 16kHz)") + parser.add_argument("--output-wav", default=Path("realtime_output.wav"), type=Path, help="Output WAV path") + parser.add_argument( + "--output-text", + default=None, + type=Path, + help="Optional output text path for final transcription", + ) + parser.add_argument("--chunk-ms", type=int, default=200, help="Input chunk size in milliseconds") + parser.add_argument( + "--send-delay-ms", + type=int, + default=0, + help="Delay between chunk sends; set >0 to simulate realtime upload", + ) + parser.add_argument( + "--delta-dump-dir", + type=Path, + default=None, + help="If set, each response.audio.delta is saved as delta_NNNNNN.wav under this directory", + ) + parser.add_argument("--num-requests", type=int, default=1, help="Total number of requests to send") + parser.add_argument( + "--concurrency", + type=int, + default=1, + help="Maximum number of concurrent websocket requests", + ) + args = parser.parse_args() + + if args.num_requests <= 0: + raise ValueError("--num-requests must be >= 1") + if args.concurrency <= 0: + raise ValueError("--concurrency must be >= 1") + concurrency = min(args.concurrency, args.num_requests) + + if args.num_requests == 1: + asyncio.run( + run_client( + url=args.url, + model=args.model, + input_wav=args.input_wav, + output_wav=args.output_wav, + output_text=args.output_text, + chunk_ms=args.chunk_ms, + send_delay_ms=args.send_delay_ms, + delta_dump_dir=args.delta_dump_dir, + ) + ) else: - # Use default audio asset - audio_path = str(AudioAsset("mary_had_lamb").get_local_path()) - print(f"No audio path provided, using default: {audio_path}") - - asyncio.run(realtime_transcribe(audio_path, args.host, args.port, args.model)) + asyncio.run( + run_clients_concurrent( + url=args.url, + model=args.model, + input_wav=args.input_wav, + output_wav=args.output_wav, + output_text=args.output_text, + chunk_ms=args.chunk_ms, + send_delay_ms=args.send_delay_ms, + delta_dump_dir=args.delta_dump_dir, + num_requests=args.num_requests, + concurrency=concurrency, + ) + ) if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Realtime WebSocket Transcription Client") - parser.add_argument( - "--model", - type=str, - default="Qwen/Qwen3-Omni-30B-A3B-Instruct", - help="Model that is served and should be pinged.", - ) - parser.add_argument( - "--audio_path", - type=str, - default=None, - help="Path to the audio file to transcribe.", - ) - parser.add_argument( - "--host", - type=str, - default="localhost", - help="vLLM-Omni server host (default: localhost)", - ) - parser.add_argument( - "--port", - type=int, - default=8000, - help="vLLM-Omni server port (default: 8000)", - ) - args = parser.parse_args() - main(args) + main() diff --git a/tests/core/sched/test_omni_scheduler_mixin.py b/tests/core/sched/test_omni_scheduler_mixin.py new file mode 100644 index 000000000..e04a9c39f --- /dev/null +++ b/tests/core/sched/test_omni_scheduler_mixin.py @@ -0,0 +1,129 @@ +"""Unit tests for OmniSchedulerMixin streaming session replacement. + +These tests pin the behavior of `_replace_session_with_streaming_update` against +current vLLM `Request` / `StreamingUpdate` (and Omni patches). When upgrading +vLLM, failures here should highlight incompatible changes to request state or +update payloads early. +""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest + +# Imports must run in this order: vllm_omni applies patches to vllm.v1.request before +# Request / StreamingUpdate are bound in this module. Ruff isort would reorder them. +# isort: off +import vllm_omni # noqa: F401 - import for side effects (patch vLLM) +from vllm.sampling_params import SamplingParams +from vllm.v1.engine import EngineCoreEventType +from vllm.v1.request import Request, RequestStatus, StreamingUpdate +from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin + +# isort: on + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _SchedulerStub(OmniSchedulerMixin): + """Minimal scheduler surface required by OmniSchedulerMixin.""" + + def __init__(self, *, log_stats: bool = False) -> None: + self.num_waiting_for_streaming_input = 0 + self.log_stats = log_stats + + +def _make_request(**kwargs) -> Request: + sp = SamplingParams(max_tokens=8) + defaults = dict( + request_id="req-mixin-test", + prompt_token_ids=[1, 2, 3], + sampling_params=sp, + pooling_params=None, + arrival_time=100.0, + block_hasher=None, + ) + defaults.update(kwargs) + return Request(**defaults) + + +def _make_update(**kwargs) -> StreamingUpdate: + sp_new = SamplingParams(max_tokens=16) + defaults = dict( + mm_features=None, + prompt_token_ids=[10, 20], + max_tokens=32, + arrival_time=200.0, + sampling_params=sp_new, + ) + defaults.update(kwargs) + return StreamingUpdate(**defaults) + + +class TestReplaceSessionWithStreamingUpdate: + def test_resets_tokens_and_prompt_from_update(self) -> None: + sched = _SchedulerStub() + session = _make_request() + session.append_output_token_ids([7, 8]) + session.num_computed_tokens = 99 + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + + update = _make_update(prompt_token_ids=[40, 41, 42]) + sched.num_waiting_for_streaming_input = 3 + sched._replace_session_with_streaming_update(session, update) + + assert session._output_token_ids == [] + assert list(session._all_token_ids) == [40, 41, 42] + assert session.prompt_token_ids == [40, 41, 42] + assert session.num_computed_tokens == 0 + assert session.num_prompt_tokens == 3 + assert session.arrival_time == 200.0 + assert session.sampling_params is update.sampling_params + assert session.status == RequestStatus.WAITING + assert sched.num_waiting_for_streaming_input == 2 + + def test_none_prompt_token_ids_becomes_empty(self) -> None: + sched = _SchedulerStub() + session = _make_request() + session.status = RequestStatus.RUNNING + update = _make_update(prompt_token_ids=None) + sched._replace_session_with_streaming_update(session, update) + + assert session.prompt_token_ids == () + assert list(session._all_token_ids) == [] + assert session.num_prompt_tokens == 0 + assert sched.num_waiting_for_streaming_input == 0 + + def test_additional_information_cleared_when_update_omits_it(self) -> None: + sched = _SchedulerStub() + session = _make_request() + if not hasattr(session, "additional_information"): + pytest.skip("Request has no additional_information (Omni patch inactive?)") + session.additional_information = {"keep": True} + session.status = RequestStatus.RUNNING + + base = _make_update() + if not hasattr(base, "additional_information"): + pytest.skip("StreamingUpdate has no additional_information (Omni patch inactive?)") + update = replace(base, additional_information=None) + + sched._replace_session_with_streaming_update(session, update) + assert session.additional_information is None + + def test_does_not_decrement_waiting_when_not_streaming_status(self) -> None: + sched = _SchedulerStub() + session = _make_request() + session.status = RequestStatus.RUNNING + sched.num_waiting_for_streaming_input = 5 + sched._replace_session_with_streaming_update(session, _make_update()) + assert sched.num_waiting_for_streaming_input == 5 + + def test_records_queued_event_when_log_stats_enabled(self) -> None: + sched = _SchedulerStub(log_stats=True) + session = _make_request() + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + sched._replace_session_with_streaming_update(session, _make_update()) + + assert session.events + assert session.events[-1].type == EngineCoreEventType.QUEUED diff --git a/tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py b/tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py new file mode 100644 index 000000000..c0f700fc4 --- /dev/null +++ b/tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E online tests for Qwen3-Omni /v1/realtime WebSocket (streaming PCM in, audio out). +""" + +from __future__ import annotations + +import asyncio +import base64 +import io +import json +import os +import wave +from pathlib import Path + +import pytest +import websockets + +from tests.conftest import ( + OmniServerParams, + convert_audio_bytes_to_text, + cosine_similarity_text, + generate_synthetic_audio, + modify_stage_config, +) +from tests.utils import hardware_test +from vllm_omni.platforms import current_omni_platform + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct" + +default_stage_config = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml") +if current_omni_platform.is_xpu(): + default_stage_config = str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml") + + +def _realtime_stage_config_path() -> str: + """CI omni layout without async_chunk; stage 0 thinker max_tokens=10.""" + return modify_stage_config( + default_stage_config, + updates={"stage_args": {0: {"default_sampling_params.max_tokens": 10}}}, + ) + + +realtime_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=_realtime_stage_config_path(), + use_stage_cli=True, + ), + id="thinker_max_tokens_10", + ), +] + + +def _pcm16_mono_16k_from_wav_bytes(wav_bytes: bytes) -> bytes: + with wave.open(io.BytesIO(wav_bytes), "rb") as wf: + if wf.getnchannels() != 1: + raise ValueError(f"Expected mono WAV, got {wf.getnchannels()} channels") + if wf.getsampwidth() != 2: + raise ValueError(f"Expected 16-bit PCM, sampwidth={wf.getsampwidth()}") + if wf.getframerate() != 16000: + raise ValueError(f"Expected 16 kHz input for /v1/realtime, got {wf.getframerate()} Hz") + if wf.getcomptype() != "NONE": + raise ValueError(f"Expected uncompressed PCM, comptype={wf.getcomptype()!r}") + return wf.readframes(wf.getnframes()) + + +def _wav_bytes_from_pcm16(pcm: bytes, sample_rate_hz: int) -> bytes: + buf = io.BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate_hz) + wf.writeframes(pcm) + return buf.getvalue() + + +async def _run_realtime_audio_roundtrip( + host: str, + port: int, + model: str, + pcm16: bytes, + *, + chunk_ms: int = 100, +) -> dict: + uri = f"ws://{host}:{port}/v1/realtime" + incremental: list[bytes] = [] + output_sr = 24000 + text_chunks: list[str] = [] + final_text = "" + delta_events = 0 + + bytes_per_ms = 16000 * 2 // 1000 + chunk_bytes = max(bytes_per_ms * chunk_ms, 2) + + async with websockets.connect(uri, max_size=64 * 1024 * 1024) as ws: + await ws.send(json.dumps({"type": "session.update", "model": model})) + await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False})) + + for i in range(0, len(pcm16), chunk_bytes): + chunk = pcm16[i : i + chunk_bytes] + await ws.send( + json.dumps( + { + "type": "input_audio_buffer.append", + "audio": base64.b64encode(chunk).decode("utf-8"), + } + ) + ) + + await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True})) + + while True: + message = await asyncio.wait_for(ws.recv(), timeout=600) + if isinstance(message, bytes): + continue + + event = json.loads(message) + event_type = event.get("type") + + if event_type == "session.created": + continue + + if event_type == "response.audio.delta": + delta_events += 1 + sr = event.get("sample_rate_hz") + if isinstance(sr, int) and sr > 0: + output_sr = sr + audio_b64 = event.get("audio", "") + if audio_b64: + incremental.append(base64.b64decode(audio_b64)) + continue + + if event_type == "transcription.delta": + d = event.get("delta", "") + if d: + text_chunks.append(d) + continue + + if event_type == "transcription.done": + final_text = event.get("text", "") or "".join(text_chunks) + continue + + if event_type == "response.audio.done": + break + + if event_type == "error": + raise AssertionError(f"WebSocket error: {event}") + + raise AssertionError(f"Unexpected WebSocket event: {event}") + + out_pcm = b"".join(incremental) + return { + "output_pcm": out_pcm, + "output_sample_rate": output_sr, + "transcription_text": final_text if final_text else "".join(text_chunks), + "delta_events": delta_events, + } + + +class TestQwen3OmniRealtimeWebSocket: + @pytest.mark.advanced_model + @pytest.mark.omni + @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) + @pytest.mark.parametrize("omni_server", realtime_server_params, indirect=True) + def test_streaming_audio_input_pcm_output(self, omni_server) -> None: + """ + Short streamed 16 kHz mono PCM16 input; expect streamed PCM16 audio deltas and + transcription. Verify Whisper(output audio) aligns with model text (same idea + as multimodal omni e2e). + """ + syn = generate_synthetic_audio(10, 1, sample_rate=16000) + wav_bytes = base64.b64decode(syn["base64"]) + pcm16 = _pcm16_mono_16k_from_wav_bytes(wav_bytes) + + result = asyncio.run( + _run_realtime_audio_roundtrip( + omni_server.host, + omni_server.port, + omni_server.model, + pcm16, + chunk_ms=100, + ) + ) + + out_pcm = result["output_pcm"] + assert result["delta_events"] >= 1 + assert out_pcm, "No output PCM from response.audio.delta" + assert len(out_pcm) % 2 == 0 + assert len(out_pcm) >= 4096, "Output audio unexpectedly small" + assert result["output_sample_rate"] > 0 + + final_text = (result["transcription_text"] or "").strip() + assert final_text, "Expected non-empty transcription (model text stream)" + + wav_out = _wav_bytes_from_pcm16(out_pcm, result["output_sample_rate"]) + whisper_text = convert_audio_bytes_to_text(wav_out).strip() + assert whisper_text, "Whisper returned empty string for synthesized output audio" + + sim = cosine_similarity_text(whisper_text.lower(), final_text.lower()) + assert sim > 0.9, ( + f"Output audio transcript should match model text (sim={sim:.3f}): " + f"whisper={whisper_text!r}, model_text={final_text!r}" + ) diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 7bf2eccf7..0b549f58e 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -70,7 +70,7 @@ class FakeStageClient: def set_engine_outputs(self, outputs) -> None: return None - def process_engine_inputs(self, stage_list, prompt=None): + def process_engine_inputs(self, stage_list, prompt=None, streaming_context=None): return list(self.next_inputs) async def abort_requests_async(self, request_ids: list[str]) -> None: diff --git a/tests/entrypoints/test_realtime_connection_helpers.py b/tests/entrypoints/test_realtime_connection_helpers.py new file mode 100644 index 000000000..e795aa92d --- /dev/null +++ b/tests/entrypoints/test_realtime_connection_helpers.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for realtime streaming helpers (PR #2581 /v1/realtime path).""" + +from __future__ import annotations + +import base64 + +import numpy as np +import pytest +import torch +from vllm.sampling_params import RequestOutputKind, SamplingParams + +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.openai.realtime_connection import RealtimeConnection + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def realtime_conn() -> RealtimeConnection: + return RealtimeConnection.__new__(RealtimeConnection) + + +class TestRealtimeConnectionTensorAndPcm: + def test_tensor_to_numpy_none(self) -> None: + assert RealtimeConnection._tensor_to_numpy(None) is None + + def test_tensor_to_numpy_1d_numpy(self) -> None: + arr = np.array([1.0, 2.0], dtype=np.float64) + out = RealtimeConnection._tensor_to_numpy(arr) + assert out is not None + assert out.dtype == np.float32 + assert out.shape == (2,) + + def test_tensor_to_numpy_2d_numpy_flattened(self) -> None: + arr = np.array([[0.5], [-0.5]], dtype=np.float32) + out = RealtimeConnection._tensor_to_numpy(arr) + assert out is not None + assert out.shape == (2,) + + def test_tensor_to_numpy_torch(self) -> None: + t = torch.tensor([[0.25, -0.25]], dtype=torch.float32) + out = RealtimeConnection._tensor_to_numpy(t) + assert out is not None + assert out.shape == (2,) + np.testing.assert_allclose(out, [0.25, -0.25], rtol=1e-5) + + def test_pcm16_b64_roundtrip(self) -> None: + audio = np.array([0.0, 1.0, -1.0], dtype=np.float32) + b64 = RealtimeConnection._pcm16_b64(audio) + raw = base64.b64decode(b64) + assert len(raw) == 6 + pcm = np.frombuffer(raw, dtype=np.int16) + assert pcm[0] == 0 + assert pcm[1] == 32767 + assert pcm[2] == -32767 + + +class TestAsyncOmniStreamingParamsValidation: + def test_accepts_streaming_friendly_params(self) -> None: + p = SamplingParams( + n=1, + stop=[], + output_kind=RequestOutputKind.DELTA, + ) + AsyncOmni._validate_streaming_input_sampling_params(p) + + def test_rejects_non_sampling_params(self) -> None: + with pytest.raises(ValueError, match="Input streaming"): + AsyncOmni._validate_streaming_input_sampling_params(object()) # type: ignore[arg-type] + + def test_rejects_n_greater_than_one(self) -> None: + p = SamplingParams(n=2, stop=[], output_kind=RequestOutputKind.DELTA) + with pytest.raises(ValueError, match="Input streaming"): + AsyncOmni._validate_streaming_input_sampling_params(p) + + def test_rejects_final_only(self) -> None: + p = SamplingParams(n=1, stop=[], output_kind=RequestOutputKind.FINAL_ONLY) + with pytest.raises(ValueError, match="Input streaming"): + AsyncOmni._validate_streaming_input_sampling_params(p) + + def test_rejects_stop_strings(self) -> None: + p = SamplingParams(n=1, stop=["\n"], output_kind=RequestOutputKind.DELTA) + with pytest.raises(ValueError, match="Input streaming"): + AsyncOmni._validate_streaming_input_sampling_params(p) diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py new file mode 100644 index 000000000..18972c91d --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for Qwen3-Omni streaming thinker→talker / talker→codec helpers (PR #2581).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3 + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def _streaming_context() -> SimpleNamespace: + return SimpleNamespace(bridge_states={}) + + +def test_get_streaming_talker_tokens_first_segment(_streaming_context: SimpleNamespace) -> None: + inc_p, inc_o, merged, thinker_in = q3._get_streaming_talker_tokens( + "r1", + [1, 2], + [10, 11], + streaming_context=_streaming_context, + ) + assert inc_p == [1, 2] + assert inc_o == [10, 11] + assert merged == [1, 2, 10, 11] + assert thinker_in == [1, 2] + + +def test_get_streaming_talker_tokens_second_segment_accumulates(_streaming_context: SimpleNamespace) -> None: + q3._get_streaming_talker_tokens("r2", [1, 2], [10, 11], streaming_context=_streaming_context) + inc_p, inc_o, merged, thinker_in = q3._get_streaming_talker_tokens( + "r2", + [1, 2, 3, 4], + [10, 11, 12, 13], + streaming_context=_streaming_context, + ) + assert inc_p == [3, 4] + assert inc_o == [12, 13] + assert merged == [1, 2, 10, 3, 4, 12, 13] + assert thinker_in == [1, 2, 10, 3, 4] + + +def test_get_streaming_talker_tokens_new_prompt_len_snapshot_truncates( + _streaming_context: SimpleNamespace, +) -> None: + inc_p, inc_o, merged, thinker_in = q3._get_streaming_talker_tokens( + "r3", + [1, 2, 3, 4, 5, 6], + [10], + new_prompt_len_snapshot=2, + streaming_context=_streaming_context, + ) + assert inc_p == [1, 2, 3, 4] + assert inc_o == [10] + assert merged == [1, 2, 3, 4, 10] + assert thinker_in == [1, 2, 3, 4] + + +def test_get_streaming_talker_tokens_clear_state(_streaming_context: SimpleNamespace) -> None: + q3._get_streaming_talker_tokens("r4", [1], [2], streaming_context=_streaming_context, clear_state=True) + state = q3._get_qwen3_streaming_state("r4", _streaming_context).thinker2talker + assert state.last_prompt_len == 0 + assert state.last_output_len == 0 + assert state.merged_sequences == [] + + +def test_get_streaming_codec_delta_len_increments_and_finishes(_streaming_context: SimpleNamespace) -> None: + d1 = q3._get_streaming_codec_delta_len(5, "c1", SimpleNamespace(finished=False), _streaming_context) + assert d1 == 5 + d2 = q3._get_streaming_codec_delta_len(8, "c1", SimpleNamespace(finished=False), _streaming_context) + assert d2 == 2 + # After d2, stored cursor is cur_seq_len + 1 == 9; next delta uses new cur_seq_len - 9. + d3 = q3._get_streaming_codec_delta_len(10, "c1", SimpleNamespace(finished=True), _streaming_context) + assert d3 == 1 + state = q3._get_qwen3_streaming_state("c1", _streaming_context) + assert state.talker2code2wav_last_seq_len == 0 diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index c5031ad11..a5579dd46 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -15,9 +15,10 @@ from vllm.v1.core.sched.utils import remove_all from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs from vllm.v1.metrics.perf import PerfStats from vllm.v1.outputs import ModelRunnerOutput -from vllm.v1.request import Request, RequestStatus +from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.metrics import SpecDecodingStats +from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin from vllm_omni.core.sched.output import OmniSchedulerOutput from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( OmniChunkTransferAdapter, @@ -38,7 +39,7 @@ class KVCacheTransferData: return asdict(self) -class OmniARScheduler(VLLMScheduler): +class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): """ OmniARScheduler: Scheduler for vLLM-Omni multimodal processing. @@ -76,6 +77,8 @@ class OmniARScheduler(VLLMScheduler): self.chunk_transfer_adapter = None if getattr(model_config, "async_chunk", False): self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config) + # Snapshot prompt length for each streaming input update + self._new_prompt_len_snapshot: dict[str, int] = {} def _get_kv_transfer_criteria(self) -> dict | None: # Note: vllm_config is available in Scheduler after super().__init__ @@ -338,6 +341,7 @@ class OmniARScheduler(VLLMScheduler): ) stopped = False + is_segment_finished = False new_logprobs = None new_token_ids = generated_token_ids pooler_output = pooler_outputs[req_index] if pooler_outputs else None @@ -366,6 +370,7 @@ class OmniARScheduler(VLLMScheduler): # Capture finish_reason BEFORE _handle_stopped_request, which may # reset the status to WAITING for streaming requests that continue. finish_reason = request.get_finished_reason() + is_segment_finished = request.is_finished() and request.resumable finished = self._handle_stopped_request(request) if finished: kv_transfer_params = self._free_request(request) @@ -418,6 +423,8 @@ class OmniARScheduler(VLLMScheduler): num_external_computed_tokens=request.num_external_computed_tokens, routed_experts=routed_experts, num_nans_in_logits=request.num_nans_in_logits, + is_segment_finished=is_segment_finished, + new_prompt_len_snapshot=self._new_prompt_len_snapshot.get(req_id, None), ) ) if self.chunk_transfer_adapter is not None: @@ -557,6 +564,21 @@ class OmniARScheduler(VLLMScheduler): return super().finish_requests(request_ids, finished_status) + def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None: + """ + Override: Only extend prompt at stage 0, and replace + the existing session with the next streaming update at other stages. + + Discards the last sampled output token from the prior input chunk at stage 0. + """ + req_id = session.request_id + self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids) + if self.vllm_config.model_config.stage_id != 0: + self._replace_session_with_streaming_update(session, update) + + else: + super()._update_request_as_session(session, update) + def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None: # TODO(wzliu)! for offline mode, we should not end process until all data is transferred """Mark a request as finished and free its resources.""" @@ -570,6 +592,7 @@ class OmniARScheduler(VLLMScheduler): self.encoder_cache_manager.free(request) request_id = request.request_id self.finished_req_ids.add(request_id) + self._new_prompt_len_snapshot.pop(request_id, None) if self.finished_req_ids_dict is not None: self.finished_req_ids_dict[request.client_index].add(request_id) diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 87f2d282a..81f0b7fc2 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import time from collections import defaultdict @@ -11,11 +13,16 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.request_queue import create_request_queue from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.core.sched.utils import remove_all -from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs +from vllm.v1.engine import ( + EngineCoreEventType, + EngineCoreOutput, + EngineCoreOutputs, +) from vllm.v1.metrics.perf import PerfStats -from vllm.v1.request import Request, RequestStatus +from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.metrics import SpecDecodingStats +from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( OmniChunkTransferAdapter, @@ -25,7 +32,7 @@ from vllm_omni.outputs import OmniModelRunnerOutput logger = init_logger(__name__) -class OmniGenerationScheduler(VLLMScheduler): +class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) model_config = self.vllm_config.model_config @@ -599,3 +606,11 @@ class OmniGenerationScheduler(VLLMScheduler): eco.scheduler_stats = stats return engine_core_outputs + + def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None: + """ + Override: Just replace the existing session with the next streaming update. + + Do not expend prompt id using update. + """ + self._replace_session_with_streaming_update(session, update) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py new file mode 100644 index 000000000..36080e63a --- /dev/null +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from vllm.v1.engine import EngineCoreEventType +from vllm.v1.request import Request, RequestStatus, StreamingUpdate + + +class OmniSchedulerMixin: + """Shared scheduler helpers for omni-specific request handling.""" + + def _replace_session_with_streaming_update( + self, + session: Request, + update: StreamingUpdate, + ) -> None: + """For streaming input: Replace an existing streaming session payload with the latest update.""" + session._output_token_ids.clear() + session._all_token_ids.clear() + new_prompt = update.prompt_token_ids or () + session._all_token_ids.extend(new_prompt) + session.num_computed_tokens = 0 + session.prompt_token_ids = update.prompt_token_ids or () + session.additional_information = update.additional_information or None + # Update block hashes for the new tokens. + session.update_block_hashes() + session.num_prompt_tokens = len(session.prompt_token_ids) + session.arrival_time = update.arrival_time + session.sampling_params = update.sampling_params + if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + self.num_waiting_for_streaming_input -= 1 + session.status = RequestStatus.WAITING + + if self.log_stats: + session.record_event(EngineCoreEventType.QUEUED) diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index c8a96e6d2..6c92d7952 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -79,6 +79,10 @@ class OmniEngineCoreRequest(EngineCoreRequest): class OmniEngineCoreOutput(EngineCoreOutput): pooling_output: dict[str, torch.Tensor] | None = None + # Finished flag for streaming input segment + is_segment_finished: bool | None = False + # Streaming update prompt length + new_prompt_len_snapshot: int | None = None class OmniEngineCoreOutputs(EngineCoreOutputs): diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 8204d70e6..a5d7ad032 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -42,6 +42,7 @@ def build_engine_core_request_from_tokens( params: SamplingParams | PoolingParams, arrival_time: float | None = None, model_config: ModelConfig | None = None, + resumable: bool = False, mm_features: list | None = None, ) -> OmniEngineCoreRequest: """Build an OmniEngineCoreRequest directly from an OmniTokensPrompt. @@ -85,6 +86,7 @@ def build_engine_core_request_from_tokens( cache_salt=None, data_parallel_rank=None, prompt_embeds=prompt_embeds, + resumable=resumable, additional_information=additional_info_payload, ) @@ -108,6 +110,20 @@ class OrchestratorRequestState: mm_processor_kwargs: dict | None = None mm_features: list | None = None + streaming: StreamingInputState = field(default_factory=lambda: StreamingInputState()) + + +@dataclass +class StreamingInputState: + # Flag of streaming input request + enabled: bool = False + # Flag of segment of streaming input finished + segment_finished: bool = False + # Streaming update prompt length + new_prompt_len_snapshot: int | None = None + # Model/bridge-specific runtime states (e.g., thinker->talker) + bridge_states: dict[str, Any] = field(default_factory=dict) + class Orchestrator: """Runs inside a background thread's asyncio event loop. @@ -391,15 +407,36 @@ class Orchestrator: ) if ( - finished + (finished or (req_state.streaming.enabled and req_state.streaming.segment_finished)) and stage_id < req_state.final_stage_id and not self.async_chunk - and not self._next_stage_already_submitted(stage_id, req_state) + and (not self._next_stage_already_submitted(stage_id, req_state) or req_state.streaming.enabled) ): - if self._cfg_tracker.has_companions(req_id) and not self._cfg_tracker.all_companions_done(req_id): + if ( + finished + and self._cfg_tracker.has_companions(req_id) + and not self._cfg_tracker.all_companions_done(req_id) + ): self._cfg_tracker.defer_parent(req_id, output, stage_id) else: - await self._forward_to_next_stage(req_id, stage_id, output, req_state) + await self._forward_to_next_stage( + req_id, + stage_id, + output, + req_state, + is_streaming_session=req_state.streaming.enabled, + is_final_update=False, + ) + if req_state.streaming.enabled and finished: + # For streaming sessions, send the terminal (resumable=False) update only on a finish + await self._forward_to_next_stage( + req_id, + stage_id, + output, + req_state, + is_streaming_session=True, + is_final_update=True, + ) if finished and stage_id == req_state.final_stage_id: # PD: clean up any lingering KV params for this request @@ -573,6 +610,9 @@ class Orchestrator: stage_id: int, output: Any, req_state: OrchestratorRequestState, + *, + is_streaming_session: bool = False, + is_final_update: bool = False, ) -> None: """Forward output from current stage to the next stage. @@ -582,6 +622,7 @@ class Orchestrator: next_stage_id = stage_id + 1 next_client = self.stage_clients[next_stage_id] params = req_state.sampling_params_list[next_stage_id] + next_stage_resumable = is_streaming_session and not is_final_update if next_client.stage_type == "diffusion": self.stage_clients[stage_id].set_engine_outputs([output]) @@ -671,6 +712,7 @@ class Orchestrator: next_inputs = next_client.process_engine_inputs( stage_list=self.stage_clients, prompt=req_state.prompt, + streaming_context=(req_state.streaming if req_state.streaming.enabled else None), ) except Exception: logger.exception( @@ -692,6 +734,7 @@ class Orchestrator: params=params, model_config=self.stage_vllm_configs[next_stage_id].model_config, mm_features=_mm_features, + resumable=next_stage_resumable, ) # TODO: Here we directly use the req id to assign. @@ -733,6 +776,13 @@ class Orchestrator: raw_outputs.timestamp, None, ) + for eco in raw_outputs.outputs: + if not hasattr(eco, "request_id"): + continue + req_state = self.request_states.get(eco.request_id) + if req_state: + req_state.streaming.segment_finished = eco.is_segment_finished + req_state.streaming.new_prompt_len_snapshot = eco.new_prompt_len_snapshot if processed.reqs_to_abort: await self.stage_clients[stage_id].abort_requests_async(processed.reqs_to_abort) @@ -768,6 +818,8 @@ class Orchestrator: # Track request state - use original_prompt so downstream stages # (e.g. thinker2talker) can access the raw dict with multi_modal_data. + request = prompt + is_streaming = bool(getattr(request, "resumable", False)) req_state = OrchestratorRequestState( request_id=request_id, prompt=original_prompt, @@ -775,13 +827,13 @@ class Orchestrator: final_stage_id=final_stage_id, mm_features=getattr(prompt, "mm_features", None), # Save mm_features for PD ) + req_state.streaming.enabled = is_streaming req_state.stage_submit_ts[stage_id] = _time.time() self.request_states[request_id] = req_state # Stage-0 prompt is already a fully-formed OmniEngineCoreRequest # (pre-processed by AsyncOmniEngine.add_request, output processor # already registered there) - submit directly. - request = prompt stage_client = self.stage_clients[stage_id] if stage_client.stage_type == "diffusion": if isinstance(prompt, list): @@ -817,6 +869,7 @@ class Orchestrator: if "sampling_params_list" in msg and msg["sampling_params_list"]: req_state.sampling_params_list = msg["sampling_params_list"] + req_state.streaming.enabled = True req_state.stage_submit_ts[stage_id] = _time.time() stage_client = self.stage_clients[stage_id] diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index ab2de757b..66bae22a5 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -302,11 +302,21 @@ class StageEngineCoreClientBase: self, stage_list: list[Any], prompt: OmniTokensPrompt | list[OmniTokensPrompt] | None = None, + streaming_context: Any | None = None, ) -> list[OmniTokensPrompt]: """Process inputs from upstream stages.""" from vllm_omni.inputs.data import OmniTokensPrompt if self.custom_process_input_func is not None: + # Keep legacy arg call for non-streaming processors. + if bool(getattr(streaming_context, "enabled", False)): + return self.custom_process_input_func( + stage_list, + self.engine_input_source, + prompt, + self.requires_multimodal_data, + streaming_context, + ) return self.custom_process_input_func( stage_list, self.engine_input_source, diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 5823aa4ab..9606cc80d 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -309,7 +309,6 @@ class AsyncOmni(EngineClient, OmniBase): if not stage0_params.skip_clone: stage0_params = stage0_params.clone() stage0_params.skip_clone = True - stage0_params.output_kind = RequestOutputKind.DELTA has_submitted_first_chunk = False diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index b45c096d0..1c3761aba 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -53,7 +53,6 @@ from vllm.entrypoints.openai.engine.protocol import ( from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.orca_metrics import metrics_header -from vllm.entrypoints.openai.realtime.connection import RealtimeConnection from vllm.entrypoints.openai.realtime.serving import OpenAIServingRealtime from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses from vllm.entrypoints.openai.server_utils import get_uvicorn_log_config @@ -108,6 +107,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( VideoListResponse, VideoResponse, ) +from vllm_omni.entrypoints.openai.realtime_connection import RealtimeConnection from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler @@ -1204,6 +1204,22 @@ async def streaming_speech(websocket: WebSocket): @router.websocket("/v1/realtime") async def realtime_websocket(websocket: WebSocket): """WebSocket endpoint for OpenAI-style realtime interactions.""" + engine_client = getattr(websocket.app.state, "engine_client", None) + if engine_client is not None and getattr(engine_client, "async_chunk", False): + await websocket.accept() + await websocket.send_json( + { + "type": "error", + "error": ( + "The /v1/realtime API is not supported when async_chunk is enabled on the server. " + "Use a stage configuration with async_chunk disabled and restart the server before using " + "this endpoint." + ), + "code": "unsupported", + } + ) + await websocket.close() + return serving = getattr(websocket.app.state, "openai_serving_realtime", None) if serving is None: await websocket.accept() diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py new file mode 100644 index 000000000..1d5470f56 --- /dev/null +++ b/vllm_omni/entrypoints/openai/realtime_connection.py @@ -0,0 +1,193 @@ +from __future__ import annotations + +import asyncio +import base64 +import json +from collections.abc import AsyncGenerator +from uuid import uuid4 + +import numpy as np +from vllm.entrypoints.openai.engine.protocol import UsageInfo +from vllm.entrypoints.openai.realtime.connection import RealtimeConnection as VllmRealtimeConnection +from vllm.entrypoints.openai.realtime.protocol import TranscriptionDelta, TranscriptionDone +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +class RealtimeConnection(VllmRealtimeConnection): + """Omni realtime connection with audio-only server events. + + Reuses upstream vLLM websocket/session lifecycle and only customizes + generation output handling to emit audio deltas. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Last audio buffer seen for this realtime generation (cumulative or concatenation + # of increments); used to turn server cumulative PCM into true deltas. + self._realtime_audio_ref: np.ndarray | None = None + + async def start_generation(self): + await super().start_generation() + + @staticmethod + def _tensor_to_numpy(value) -> np.ndarray | None: + if value is None: + return None + if isinstance(value, np.ndarray): + arr = value + elif hasattr(value, "detach"): + arr = value.detach().float().cpu().numpy() + else: + try: + arr = np.asarray(value) + except Exception: + return None + if arr.ndim > 1: + arr = arr.reshape(-1) + return arr.astype(np.float32, copy=False) + + @staticmethod + def _numpy_audio_prefix_match(prev: np.ndarray, curr: np.ndarray) -> bool: + n = prev.shape[0] + if n == 0: + return True + if curr.shape[0] < n: + return False + return bool(np.allclose(curr[:n], prev, rtol=1e-3, atol=2e-4)) + + def _raw_waveform_to_deltas(self, arr: np.ndarray) -> list[np.ndarray]: + """Convert one streaming PCM f32 chunk into incremental piece(s) for the client. + + Some engine paths emit a growing cumulative waveform each step; others emit + true per-step deltas. We support both without duplicating audio on the client. + """ + if arr.size == 0: + return [] + ref = self._realtime_audio_ref + if ref is None: + self._realtime_audio_ref = arr.copy() + return [arr] + if self._numpy_audio_prefix_match(ref, arr): + delta = arr[ref.shape[0] :] + self._realtime_audio_ref = arr.copy() + return [delta] if delta.size > 0 else [] + # True per-step delta (not a prefix extension of what we have seen). + self._realtime_audio_ref = np.concatenate([ref, arr]) + return [arr] + + def _extract_audio_chunks(self, output) -> tuple[list[np.ndarray], int]: + mm = getattr(output, "multimodal_output", None) + if not isinstance(mm, dict): + return [], 24000 + + sr = mm.get("sr") or mm.get("sample_rate") or mm.get("audio_sample_rate") or 24000 + key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None) + if key is None: + return [], int(sr) + + raw_audio = mm.get(key) + chunks: list[np.ndarray] = [] + if isinstance(raw_audio, (list, tuple)): + if len(raw_audio) > 0: + arr = self._tensor_to_numpy(raw_audio[-1]) + if arr is not None and arr.size > 0: + chunks.extend(self._raw_waveform_to_deltas(arr)) + else: + arr = self._tensor_to_numpy(raw_audio) + if arr is not None and arr.size > 0: + chunks.extend(self._raw_waveform_to_deltas(arr)) + return chunks, int(sr) + + @staticmethod + def _pcm16_b64(audio_f32: np.ndarray) -> str: + clipped = np.clip(audio_f32, -1.0, 1.0) + pcm16 = (clipped * 32767.0).astype(np.int16) + return base64.b64encode(pcm16.tobytes()).decode("utf-8") + + async def _run_generation( + self, + streaming_input_gen: AsyncGenerator, + input_stream: asyncio.Queue[list[int]], + ): + request_id = f"rt-{self.connection_id}-{uuid4()}" + sent_audio = False + audio_done_sent = False + full_text = "" + sent_text_len = 0 + prompt_token_ids_len = 0 + completion_tokens_len = 0 + self._realtime_audio_ref = None + + try: + result_gen = self.serving.engine_client.generate( + prompt=streaming_input_gen, + request_id=request_id, + ) + + async for output in result_gen: + if output.outputs and len(output.outputs) > 0: + output0 = output.outputs[0] + token_ids = list(output0.token_ids) + if token_ids: + input_stream.put_nowait(token_ids) + # token_ids are cumulative per request + completion_tokens_len = len(token_ids) + if not prompt_token_ids_len and output.prompt_token_ids: + prompt_token_ids_len = len(output.prompt_token_ids) + cumulative_text = output0.text or "" + if cumulative_text: + if len(cumulative_text) >= sent_text_len: + delta_text = cumulative_text[sent_text_len:] + else: + delta_text = cumulative_text + sent_text_len = len(cumulative_text) + full_text = cumulative_text + else: + delta_text = "" + + if delta_text: + await self.send(TranscriptionDelta(delta=delta_text)) + + audio_chunks, sample_rate = self._extract_audio_chunks(output) + + for chunk in audio_chunks: + sent_audio = True + await self.send_json( + { + "type": "response.audio.delta", + "audio": self._pcm16_b64(chunk), + "format": "pcm16", + "sample_rate_hz": sample_rate, + } + ) + + if not self._is_connected: + break + + usage = UsageInfo( + prompt_tokens=prompt_token_ids_len, + completion_tokens=completion_tokens_len, + total_tokens=prompt_token_ids_len + completion_tokens_len, + ) + await self.send(TranscriptionDone(text=full_text, usage=usage)) + + if sent_audio: + await self.send_json({"type": "response.audio.done", "has_audio": True}) + audio_done_sent = True + except Exception as e: + logger.exception("Error in generation: %s", e) + await self.send_error(str(e), "processing_error") + finally: + # Always send terminal event so clients don't hang forever. + if self._is_connected and not audio_done_sent: + try: + await self.send_json({"type": "response.audio.done", "has_audio": sent_audio}) + except Exception: + logger.exception("Failed to send response.audio.done") + while not self.audio_queue.empty(): + self.audio_queue.get_nowait() + + async def send_json(self, payload: dict): + await self.websocket.send_text(json.dumps(payload)) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index cd8d147ca..f06ecf41d 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -180,6 +180,11 @@ class Qwen3OmniMoeForConditionalGeneration( "trailing_text_hidden", "tts_pad_embed_projected", } + # Keys that need to be accumulated across streaming inputs + self.streaming_accumulated_keys: set[str] = { + "thinker_prefill_embeddings", + "thinker_hidden_states", + } elif self.model_stage == "code2wav": self.enable_update_additional_information = True diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index c502041fe..699e4b194 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -4,6 +4,7 @@ """Stage input processor for Qwen3 Omni MoE: Thinker → Talker transition.""" import logging +from dataclasses import dataclass, field from typing import Any import torch @@ -180,6 +181,102 @@ def _resolve_tts_token_embedding( return val.detach().to(device=device, dtype=torch.float) if val is not None else None +# ========================= +# Streaming input helpers +# ========================= + + +@dataclass +class _Thinker2TalkerStreamingState: + last_prompt_len: int = 0 + last_output_len: int = 0 + merged_sequences: list[int] = field(default_factory=list) + + +@dataclass +class _Qwen3OmniStreamingState: + thinker2talker: _Thinker2TalkerStreamingState = field(default_factory=_Thinker2TalkerStreamingState) + talker2code2wav_last_seq_len: int = 0 + + +def _get_qwen3_streaming_state( + request_id: str, + streaming_context: Any | None, +) -> _Qwen3OmniStreamingState: + bridge_states = getattr(streaming_context, "bridge_states", None) + per_model_state = bridge_states.setdefault("qwen3_omni", {}) + state = per_model_state.get(request_id) + if state is None: + state = _Qwen3OmniStreamingState() + per_model_state[request_id] = state + return state + + +def _get_streaming_talker_tokens( + request_id: str, + prompt_token_ids: list[int], + output_token_ids: list[int], + new_prompt_len_snapshot: int | None = None, + streaming_context: Any | None = None, + *, + clear_state: bool = False, +) -> tuple[list[int], list[int], list[int], list[int]]: + """Return streaming token slices and merged token views for thinker->talker. + e.g. For the second streaming input request: + merged_sequences: [input_prompt 1, output_tokens 1[:-1], input_prompt 2, output_tokens 2] + thinker_input_ids: [input_prompt 1, output_tokens 1[:-1], input_prompt 2] + Returns: + inc_prompt: prompt token delta for this segment. + inc_output: output token delta for this segment. + merged_sequences: full thinker_sequences to send downstream. + thinker_input_ids: full thinker_input_ids paired with merged_sequences. + """ + state = _get_qwen3_streaming_state(request_id, streaming_context).thinker2talker + if new_prompt_len_snapshot: + prompt_token_ids = prompt_token_ids[:-new_prompt_len_snapshot] + cur_prompt_len = len(prompt_token_ids) + cur_output_len = len(output_token_ids) + + inc_prompt = prompt_token_ids[state.last_prompt_len :] + inc_output = output_token_ids[state.last_output_len :] + delta_sequences = inc_prompt + inc_output + cached_sequences = state.merged_sequences + + merged_sequences = cached_sequences + delta_sequences + thinker_input_ids = cached_sequences + inc_prompt + + # Persist history for next segment. Drop the latest sampled token to keep + # thinker_input_ids / thinker_sequences alignment with next-step append. + cached_sequences.extend(delta_sequences[:-1]) + + state.last_prompt_len = cur_prompt_len + state.last_output_len = cur_output_len + + if clear_state: + state.last_prompt_len = 0 + state.last_output_len = 0 + state.merged_sequences.clear() + + return inc_prompt, inc_output, merged_sequences, thinker_input_ids + + +def _get_streaming_codec_delta_len( + cur_seq_len: int, + request_id: str, + talker_output: Any, + streaming_context: Any | None = None, +) -> int: + """Return newly added seq_len for talker->code2wav in streaming mode.""" + state = _get_qwen3_streaming_state(request_id, streaming_context) + prev_seq_len = state.talker2code2wav_last_seq_len + seq_len = cur_seq_len - prev_seq_len + state.talker2code2wav_last_seq_len = cur_seq_len + 1 + if bool(getattr(talker_output, "finished", False)): + # Final segment: clear history to avoid cross-session carry-over. + state.talker2code2wav_last_seq_len = 0 + return seq_len + + # ========================= # Thinker -> Talker # ========================= @@ -272,6 +369,7 @@ def thinker2talker( engine_input_source: list[int], prompt: OmniTokensPrompt | TextPrompt | None = None, requires_multimodal_data: bool = False, + streaming_context: Any | None = None, ) -> list[OmniTokensPrompt]: """ Process thinker outputs to create talker inputs. @@ -305,18 +403,37 @@ def thinker2talker( # Process each thinker output for i, thinker_output in enumerate(thinker_outputs): output = thinker_output.outputs[0] + req_id = str(getattr(thinker_output, "request_id", f"idx-{i}")) + prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids) + output_ids = _ensure_list(output.token_ids) + is_streaming_session = bool(getattr(streaming_context, "enabled", False)) + if is_streaming_session: + prompt_token_ids, output_ids, thinker_sequences, thinker_input_ids = _get_streaming_talker_tokens( + req_id, + prompt_token_ids, + output_ids, + getattr(streaming_context, "new_prompt_len_snapshot", None), + streaming_context, + clear_state=bool(getattr(thinker_output, "finished", False)), + ) + else: + thinker_sequences = prompt_token_ids + output_ids + thinker_input_ids = prompt_token_ids + # For streaming input, just send incremental prefill and hidden states tensor to talker + # Equally applicable to non-streaming cases. + new_seq_length = len(prompt_token_ids + output_ids) - 1 thinker_mm = output.multimodal_output # Full thinker embedding sequence for the talker: single thinker engine in the # non-PD path; after optional merge with prefill-side tensors in PD mode. - thinker_emb = thinker_mm[_EMBED_LAYER_KEY].detach().to(device=device, dtype=torch.float) - thinker_hid = thinker_mm[_HIDDEN_LAYER_KEY].detach().to(device=device, dtype=torch.float) + thinker_emb = thinker_mm[_EMBED_LAYER_KEY].detach().to(device=device, dtype=torch.float)[-new_seq_length:] + thinker_hid = thinker_mm[_HIDDEN_LAYER_KEY].detach().to(device=device, dtype=torch.float)[-new_seq_length:] prefill_mm: dict[str, Any] | None = None if prefill_stage is not None: prefill_mm = _get_prefill_multimodal_output(prefill_stage, i) if prefill_mm is not None: - expected_total = len(thinker_output.prompt_token_ids) + len(output.token_ids) + expected_total = len(prompt_token_ids) + len(output_ids) try: thinker_emb, thinker_hid = _merge_pd_embeddings( thinker_emb, thinker_hid, prefill_mm, device, expected_total=expected_total @@ -327,10 +444,8 @@ def thinker2talker( info = { "thinker_prefill_embeddings": thinker_emb, "thinker_hidden_states": thinker_hid, - "thinker_sequences": ( - thinker_output.prompt_token_ids + output.token_ids - ), # the thinker_sequences is the whole ids - "thinker_input_ids": thinker_output.prompt_token_ids, + "thinker_sequences": thinker_sequences, # the thinker_sequences is the whole ids + "thinker_input_ids": thinker_input_ids, # Provide thinker-side TTS token embeddings for talker projection "tts_bos_embed": _resolve_tts_token_embedding( "tts_bos_embed", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device @@ -441,6 +556,7 @@ def talker2code2wav( engine_input_source: list[int], prompt: OmniTokensPrompt | TextPrompt | None = None, requires_multimodal_data: bool = False, + streaming_context: Any | None = None, ) -> list[OmniTokensPrompt]: """ Process talker outputs to create code2wav inputs. @@ -462,9 +578,14 @@ def talker2code2wav( talker_outputs = _validate_stage_inputs(stage_list, engine_input_source) code2wav_inputs: list[OmniTokensPrompt] = [] # Process each talker output - for talker_output in talker_outputs: + for i, talker_output in enumerate(talker_outputs): output = talker_output.outputs[0] - seq_len = len(output.token_ids) - 1 + req_id = str(getattr(talker_output, "request_id", f"idx-{i}")) + cur_seq_len = len(output.token_ids) - 1 + seq_len = cur_seq_len + is_streaming_session = bool(getattr(streaming_context, "enabled", False)) + if is_streaming_session: + seq_len = _get_streaming_codec_delta_len(cur_seq_len, req_id, talker_output, streaming_context) # Extract codec codes from talker output # Expected shape: [8, seq_len] (8-layer RVQ codes) codec_codes = ( diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index d4ab78f13..f6c483a92 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -12,12 +12,13 @@ from vllm.v1.engine import EngineCoreOutputs as _OriginalEngineCoreOutputs from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest from vllm.v1.request import Request as _OriginalRequest from vllm.v1.request import RequestStatus +from vllm.v1.request import StreamingUpdate as _OriginalStreamingUpdate import vllm_omni.logger # noqa: F401 from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs, OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.layers.rotary_embedding import OmniMRotaryEmbedding -from vllm_omni.request import OmniRequest +from vllm_omni.request import OmniRequest, OmniStreamingUpdate # ============================================================================= # Patch ModelConfig.is_mm_prefix_lm to support omni-specific models @@ -115,5 +116,7 @@ for module_name, module in sys.modules.items(): module.MRotaryEmbedding = OmniMRotaryEmbedding if hasattr(module, "Request") and module.Request == _OriginalRequest: module.Request = OmniRequest + if hasattr(module, "StreamingUpdate") and module.StreamingUpdate == _OriginalStreamingUpdate: + module.StreamingUpdate = OmniStreamingUpdate if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest: module.EngineCoreRequest = OmniEngineCoreRequest diff --git a/vllm_omni/request.py b/vllm_omni/request.py index 3ec325316..48cbf9b31 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -1,8 +1,11 @@ from collections.abc import Callable +from dataclasses import dataclass from typing import TYPE_CHECKING import numpy as np import torch +from vllm.multimodal.inputs import MultiModalFeatureSpec +from vllm.sampling_params import SamplingParams from vllm.v1.request import Request if TYPE_CHECKING: @@ -92,3 +95,34 @@ class OmniRequest(Request): resumable=request.resumable, reasoning_ended=request.reasoning_ended, ) + + +@dataclass +class OmniStreamingUpdate: + """ + Override: add additional information + Lightweight data for streaming session continuation. + + Contains only the fields needed to update an existing streaming session + with new input data. + """ + + mm_features: list[MultiModalFeatureSpec] | None + prompt_token_ids: list[int] | None + max_tokens: int + arrival_time: float + sampling_params: SamplingParams | None + additional_information: AdditionalInformationPayload | None = None + + @classmethod + def from_request(cls, request: "Request") -> "OmniStreamingUpdate | None": + if not request.resumable: + return None + return cls( + mm_features=request.mm_features, + prompt_token_ids=request.prompt_token_ids, + max_tokens=request.max_tokens, + arrival_time=request.arrival_time, + sampling_params=request.sampling_params, + additional_information=request.additional_information, + ) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index de78011c7..d1c15eac6 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -308,6 +308,7 @@ class OmniGPUModelRunner(GPUModelRunner): for new_req_data in scheduler_output.scheduled_new_reqs: req_id = new_req_data.req_id if req_id in self.requests: + self._update_streaming_input_additional_info(new_req_data, req_id) req_state = self._update_streaming_request(req_id, new_req_data) reqs_to_add.append(req_state) continue @@ -1414,3 +1415,30 @@ class OmniGPUModelRunner(GPUModelRunner): def _merge_additional_information_update(self, req_id, upd): logger.warning_once("_merge_additional_information_update is deprecated, use _update_intermediate_buffer") return self._update_intermediate_buffer(req_id, upd) + + def _update_streaming_input_additional_info(self, new_req_data, req_id): + # For streaming input prefill case only. Update buffer from last segment input + cached_additional_info = self.model_intermediate_buffer.get(req_id, {}) + if cached_additional_info: + payload_info = getattr(new_req_data, "additional_information", None) + inc_info = deserialize_additional_information(payload_info) + if isinstance(inc_info, dict) and inc_info: + merged_info = dict(cached_additional_info) + for key, value in inc_info.items(): + accumulated_keys: set[str] = set() + if hasattr(self, "model") and hasattr(self.model, "streaming_accumulated_keys"): + accumulated_keys = self.model.streaming_accumulated_keys + if key in accumulated_keys and isinstance(value, torch.Tensor): + inc_tensor = value.detach().to("cpu").contiguous() + old_tensor = merged_info.get(key) + if old_tensor is None: + merged_info[key] = inc_tensor + else: + merged_info[key] = torch.cat((old_tensor, inc_tensor), dim=0) + continue + + # Default for other keys: latest value. + merged_info[key] = value + merged_info["num_processed_tokens"] = 0 + self.model_intermediate_buffer[req_id] = merged_info + setattr(self.requests[req_id], "additional_information_cpu", merged_info)