From 90c890ad6aa771f595a8769495c2123d5667c6aa Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Tue, 25 Aug 2026 12:32:25 +0800 Subject: [PATCH] [Perf][TTS] Add deadline-aware audio scheduling --- .../examples/online_serving/text_to_speech.md | 1 + ..._generation_scheduler_playback_deadline.py | 39 ++++ .../test_chunk_transfer_adapter.py | 214 ++++++++++++++++++ .../test_qwen3_tts_async_chunk.py | 116 ++++++++++ .../core/sched/omni_generation_scheduler.py | 12 + vllm_omni/deploy/qwen3_tts.yaml | 8 + .../chunk_transfer_adapter.py | 140 +++++++++++- .../stage_input_processors/qwen3_tts.py | 48 ++-- 8 files changed, 559 insertions(+), 19 deletions(-) create mode 100644 tests/core/sched/test_omni_generation_scheduler_playback_deadline.py diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index 5d4ccbf3b..d14acf03d 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -457,6 +457,7 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000 - With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request. - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N- `. - Under vocoder-bound overload (single-stream `rtf_p99 ≥ 1` at the target concurrency), set `active_stream_window: 2` at the top of the deploy yaml to cap simultaneously active Stage 1 streams. Off by default; trades TTFP for streaming continuity. See [#3592](https://github.com/vllm-project/vllm-omni/pull/3592) for the mechanism and tradeoff numbers. +- For a bounded active-stream deployment, `active_stream_policy: edf` can prioritize streams whose server-estimated playback buffer is closest to exhaustion. Configure it under the connector's `extra` block and use a `codec_chunk_ramp`; this is opt-in, estimates playback from server-side PCM emission (no client playback acknowledgements), and targets first-audio latency and continuity rather than single-request completion latency. --- diff --git a/tests/core/sched/test_omni_generation_scheduler_playback_deadline.py b/tests/core/sched/test_omni_generation_scheduler_playback_deadline.py new file mode 100644 index 000000000..e63620f22 --- /dev/null +++ b/tests/core/sched/test_omni_generation_scheduler_playback_deadline.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_generation_scheduler_records_audio_output_for_edf(mocker): + scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler) + scheduler.chunk_transfer_adapter = SimpleNamespace( + playback_deadline_enabled=True, + record_audio_output=mocker.Mock(), + ) + multimodal_output = { + "model_outputs": torch.zeros(24000), + "sr": torch.tensor(24000), + } + + scheduler._record_playback_deadline_output("req", multimodal_output) + + scheduler.chunk_transfer_adapter.record_audio_output.assert_called_once_with("req", 24000, 24000) + + +def test_generation_scheduler_skips_audio_accounting_when_edf_is_disabled(mocker): + scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler) + scheduler.chunk_transfer_adapter = SimpleNamespace( + playback_deadline_enabled=False, + record_audio_output=mocker.Mock(), + ) + + scheduler._record_playback_deadline_output("req", {"model_outputs": torch.zeros(24000)}) + + scheduler.chunk_transfer_adapter.record_audio_output.assert_not_called() diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index d0fd00376..3de96d610 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -113,6 +113,7 @@ def build_adapter(monkeypatch, mocker: MockerFixture): *, stage_id: int = 1, model_mode: str = "ar", + engine_output_type: str | None = None, max_num_seqs: int = 2, active_stream_window: int = 0, connector_extra: dict | None = None, @@ -143,6 +144,7 @@ def build_adapter(monkeypatch, mocker: MockerFixture): model_config = SimpleNamespace( worker_type=model_mode, + engine_output_type=engine_output_type or ("audio" if model_mode == "generation" else None), max_num_seqs=max_num_seqs, active_stream_window=active_stream_window, stage_connector_config={ @@ -843,6 +845,218 @@ def test_fifo_promotion(build_adapter): assert reqs[2].status == RequestStatus.WAITING_FOR_CHUNK +def test_edf_requires_active_window(build_adapter): + with pytest.raises(ValueError, match="requires active_stream_window"): + build_adapter( + stage_id=1, + model_mode="generation", + active_stream_window=0, + connector_extra={"active_stream_policy": "edf"}, + ) + + +def test_edf_rejects_unknown_policy(build_adapter): + with pytest.raises(ValueError, match="active_stream_policy"): + build_adapter( + stage_id=1, + model_mode="generation", + active_stream_window=1, + connector_extra={"active_stream_policy": "unknown"}, + ) + + +@pytest.mark.parametrize( + ("key", "value", "message"), + [ + ("playback_deadline_first_chunk_ms", float("nan"), "first_chunk"), + ("playback_deadline_safety_margin_ms", float("inf"), "safety_margin"), + ], +) +def test_edf_rejects_non_finite_timing_config(build_adapter, key, value, message): + with pytest.raises(ValueError, match=message): + build_adapter( + stage_id=1, + model_mode="generation", + active_stream_window=1, + connector_extra={"active_stream_policy": "edf", key: value}, + ) + + +def test_edf_selects_earliest_playback_deadline(build_adapter, monkeypatch): + monkeypatch.setattr( + "vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter.time.monotonic", + lambda: 100.0, + ) + adapter, _ = build_adapter( + stage_id=1, + model_mode="generation", + max_num_seqs=2, + active_stream_window=1, + connector_extra={"active_stream_policy": "edf"}, + ) + later = _req("later", RequestStatus.WAITING) + urgent = _req("urgent", RequestStatus.WAITING) + adapter.record_audio_output(later.request_id, 5, 1, emitted_at=100.0) + adapter.record_audio_output(urgent.request_id, 2, 1, emitted_at=100.0) + adapter.requests_with_ready_chunks.update({later.request_id, urgent.request_id}) + waiting_queue = DummyWaitingQueue([later, urgent]) + + adapter.process_pending_chunks(waiting_queue, []) + + assert list(adapter._active_streams) == [urgent.request_id] + assert waiting_queue == [urgent] + assert list(adapter.waiting_for_chunk_waiting_requests) == [later] + + +def test_edf_balances_first_audio_against_playback_deadline(build_adapter, monkeypatch): + monkeypatch.setattr( + "vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter.time.monotonic", + lambda: 100.0, + ) + adapter, _ = build_adapter( + stage_id=1, + model_mode="generation", + max_num_seqs=2, + active_stream_window=1, + connector_extra={ + "active_stream_policy": "edf", + "playback_deadline_first_chunk_ms": 100, + "playback_deadline_safety_margin_ms": 0, + }, + ) + new_request = _req("new", RequestStatus.WAITING) + buffered = _req("buffered", RequestStatus.WAITING) + adapter.record_audio_output(buffered.request_id, 5, 1, emitted_at=100.0) + adapter.requests_with_ready_chunks.update({new_request.request_id, buffered.request_id}) + waiting_queue = DummyWaitingQueue([buffered, new_request]) + + adapter.process_pending_chunks(waiting_queue, []) + + assert list(adapter._active_streams) == [new_request.request_id] + + adapter.cleanup_receiver(new_request.request_id) + adapter.restore_queues(waiting_queue, []) + starving = _req("starving", RequestStatus.WAITING) + adapter.record_audio_output(starving.request_id, 1, 1, emitted_at=98.0) + adapter.requests_with_ready_chunks.update({new_request.request_id, starving.request_id}) + waiting_queue = DummyWaitingQueue([new_request, starving]) + + adapter.process_pending_chunks(waiting_queue, []) + + assert list(adapter._active_streams) == [starving.request_id] + + +def test_edf_protects_near_underrun_stream_from_overdue_first_audio(build_adapter, monkeypatch): + monkeypatch.setattr( + "vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter.time.monotonic", + lambda: 100.0, + ) + adapter, _ = build_adapter( + stage_id=1, + model_mode="generation", + max_num_seqs=2, + active_stream_window=1, + connector_extra={ + "active_stream_policy": "edf", + "playback_deadline_first_chunk_ms": 100, + "playback_deadline_safety_margin_ms": 250, + }, + ) + overdue_startup = _req("startup", RequestStatus.WAITING) + near_underrun = _req("playing", RequestStatus.WAITING) + adapter._playback_deadline_key(overdue_startup, order=0, now=90.0) + adapter.record_audio_output(near_underrun.request_id, 100, 1000, emitted_at=100.0) + adapter.requests_with_ready_chunks.update({overdue_startup.request_id, near_underrun.request_id}) + waiting_queue = DummyWaitingQueue([overdue_startup, near_underrun]) + + adapter.process_pending_chunks(waiting_queue, []) + + assert list(adapter._active_streams) == [near_underrun.request_id] + + +def test_edf_reselects_at_chunk_boundaries(build_adapter, monkeypatch): + monkeypatch.setattr( + "vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter.time.monotonic", + lambda: 100.0, + ) + adapter, _ = build_adapter( + stage_id=1, + model_mode="generation", + max_num_seqs=2, + active_stream_window=1, + connector_extra={"active_stream_policy": "edf", "playback_deadline_safety_margin_ms": 0}, + ) + first = _req("first", RequestStatus.WAITING) + second = _req("second", RequestStatus.WAITING) + adapter.record_audio_output(first.request_id, 1, 1, emitted_at=100.0) + adapter.record_audio_output(second.request_id, 2, 1, emitted_at=100.0) + adapter.requests_with_ready_chunks.update({first.request_id, second.request_id}) + waiting_queue = DummyWaitingQueue([first, second]) + + adapter.process_pending_chunks(waiting_queue, []) + assert list(adapter._active_streams) == [first.request_id] + + scheduler_output = SimpleNamespace( + scheduled_new_reqs=[SimpleNamespace(req_id=first.request_id)], + scheduled_cached_reqs=None, + ) + adapter.postprocess_scheduler_output(scheduler_output) + adapter.restore_queues(waiting_queue, []) + adapter.record_audio_output(first.request_id, 5, 1, emitted_at=100.0) + adapter.requests_with_ready_chunks.add(first.request_id) + + adapter.process_pending_chunks(waiting_queue, []) + + assert list(adapter._active_streams) == [second.request_id] + + +def test_edf_k2_orders_ready_running_and_waiting_requests(build_adapter, monkeypatch): + monkeypatch.setattr( + "vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter.time.monotonic", + lambda: 100.0, + ) + adapter, _ = build_adapter( + stage_id=1, + model_mode="generation", + max_num_seqs=4, + active_stream_window=2, + connector_extra={"active_stream_policy": "edf", "playback_deadline_safety_margin_ms": 0}, + ) + later_running = _req("later-running", RequestStatus.RUNNING) + urgent_waiting = _req("urgent-waiting", RequestStatus.WAITING) + middle_running = _req("middle-running", RequestStatus.RUNNING) + adapter.record_audio_output(later_running.request_id, 5, 1, emitted_at=100.0) + adapter.record_audio_output(urgent_waiting.request_id, 1, 1, emitted_at=100.0) + adapter.record_audio_output(middle_running.request_id, 2, 1, emitted_at=100.0) + adapter.requests_with_ready_chunks.update( + {later_running.request_id, urgent_waiting.request_id, middle_running.request_id} + ) + waiting_queue = DummyWaitingQueue([urgent_waiting]) + running_queue = [later_running, middle_running] + + adapter.process_pending_chunks(waiting_queue, running_queue) + + assert list(adapter._active_streams) == [urgent_waiting.request_id, middle_running.request_id] + assert waiting_queue == [urgent_waiting] + assert running_queue == [middle_running] + assert list(adapter._held_non_active) == [later_running] + + +def test_edf_cleanup_drops_playback_state(build_adapter): + adapter, _ = build_adapter( + stage_id=1, + model_mode="generation", + active_stream_window=1, + connector_extra={"active_stream_policy": "edf"}, + ) + adapter.record_audio_output("req", 24000, 24000, emitted_at=100.0) + assert "req" in adapter._playback_deadlines + + adapter.cleanup_receiver("req") + + assert "req" not in adapter._playback_deadlines + + def test_non_active_waiting_request_is_held_off_scheduler(build_adapter): adapter, _ = build_adapter(stage_id=2, max_num_seqs=2, active_stream_window=1) active = _req("req-active", RequestStatus.WAITING) diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index 013c0b99e..3f986beb0 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -105,6 +105,122 @@ def test_flush_on_finish(): assert len(p.codes.audio) == _Q * 24 +def _record_successful_chunk(tm, rid): + tm.put_req_chunk[rid] += 1 + tm.ramp_chunk_count[rid] += 1 + + +def test_fixed_chunk_terminal_after_exact_initial_boundary_emits_no_audio(): + tm = _tm(initial_chunk_frames=10) + rid = "fixed-exact-initial" + + first = _call(tm, rid, n_frames=10, req_ic=10) + assert first is not None + _record_successful_chunk(tm, rid) + + terminal = _call(tm, rid, n_frames=10, finished=True, req_ic=10) + + assert terminal is not None + assert terminal.meta.finished.item() is True + assert terminal.codes.audio.numel() == 0 + + +def test_fixed_chunk_terminal_after_exact_steady_boundary_emits_no_audio(): + tm = _tm(initial_chunk_frames=10) + rid = "fixed-exact-steady" + + first = _call(tm, rid, n_frames=10, req_ic=10) + assert first is not None + _record_successful_chunk(tm, rid) + steady = _call(tm, rid, n_frames=35, req_ic=10) + assert steady is not None + _record_successful_chunk(tm, rid) + + terminal = _call(tm, rid, n_frames=35, finished=True, req_ic=10) + + assert terminal is not None + assert terminal.meta.finished.item() is True + assert terminal.codes.audio.numel() == 0 + + +def test_fixed_chunk_terminal_flushes_exact_boundary_not_previously_sent(): + tm = _tm(initial_chunk_frames=10) + rid = "fixed-unsent-boundary" + + first = _call(tm, rid, n_frames=10, req_ic=10) + assert first is not None + _record_successful_chunk(tm, rid) + + terminal = _call(tm, rid, n_frames=35, finished=True, req_ic=10) + + assert terminal is not None + assert terminal.meta.finished.item() is True + assert terminal.codes.audio.numel() == _Q * 25 + + +def test_fixed_chunk_without_initial_terminal_after_exact_boundary_emits_no_audio(): + tm = _tm(chunk_frames=4, initial_chunk_frames=4) + rid = "fixed-no-initial-exact" + + first = _call(tm, rid, n_frames=4, req_ic=4) + assert first is not None + _record_successful_chunk(tm, rid) + second = _call(tm, rid, n_frames=8, req_ic=4) + assert second is not None + _record_successful_chunk(tm, rid) + + terminal = _call(tm, rid, n_frames=8, finished=True, req_ic=4) + + assert terminal is not None + assert terminal.codes.audio.numel() == 0 + + +def test_fixed_chunk_terminal_after_exact_boundary_flushes_only_new_tail(): + tm = _tm(initial_chunk_frames=10) + rid = "fixed-tail-after-boundary" + + first = _call(tm, rid, n_frames=10, req_ic=10) + assert first is not None + _record_successful_chunk(tm, rid) + steady = _call(tm, rid, n_frames=35, req_ic=10) + assert steady is not None + _record_successful_chunk(tm, rid) + + terminal = _call(tm, rid, n_frames=40, finished=True, req_ic=10) + + assert terminal is not None + assert terminal.codes.audio.numel() == _Q * 5 + + +def test_fixed_chunk_stream_sends_every_codec_frame_exactly_once(): + tm = _tm(chunk_frames=4, initial_chunk_frames=2) + rid = "fixed-unique" + frames = [[frame * 10 + quantizer for quantizer in range(_Q)] for frame in range(10)] + emitted = [] + + for length in (2, 6, 10): + tm.code_prompt_token_ids[rid] = frames[:length] + payload = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output={"codes": {"audio": torch.zeros((0,))}}, + request=_req(rid, finished=False, initial_codec_chunk_frames=2), + ) + assert payload is not None + emitted.extend(payload.codes.audio.reshape(_Q, -1).transpose(0, 1).tolist()) + _record_successful_chunk(tm, rid) + + terminal = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output=None, + request=_req(rid, finished=True, initial_codec_chunk_frames=2), + is_finished=True, + ) + + assert terminal is not None + assert terminal.codes.audio.numel() == 0 + assert emitted == frames + + _CASES = [ # ── IC boundary rule ────────────────────────────────────────────── # initial_codec_chunk_frames only controls the first emitted chunk. diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index ae19a41a3..3828d7908 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -24,6 +24,8 @@ 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.core.sched.utils import omni_routed_experts_for_request +from vllm_omni.metrics import definitions as metric_defs +from vllm_omni.metrics.utils import count_audio_frames from vllm_omni.outputs import OmniModelRunnerOutput logger = init_logger(__name__) @@ -70,6 +72,14 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): return False return super()._handle_stopped_request(request) + def _record_playback_deadline_output(self, request_id: str, multimodal_output: Any) -> None: + adapter = self.chunk_transfer_adapter + if adapter is None or not getattr(adapter, "playback_deadline_enabled", False): + return + audio_frames = count_audio_frames(multimodal_output) + sample_rate = metric_defs.resolve_audio_sample_rate(multimodal_output) + adapter.record_audio_output(request_id, audio_frames, sample_rate) + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: """One-shot generation fast path: - Feed all input tokens of the request at once @@ -440,6 +450,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): kv_transfer_params = None pooler_output = pooler_outputs[req_index] if pooler_outputs else None mm_output = mm_outputs[req_index] if mm_outputs else None + if mm_output is not None: + self._record_playback_deadline_output(req_id, mm_output) status_before_stop = request.status finish_reason = None is_segment_finished = False diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 58b191f77..ae510453e 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -26,6 +26,14 @@ connectors: codec_left_context_frames: 72 # Emit only the first audio chunk early, then return to codec_chunk_frames. initial_codec_chunk_frames: 1 + # Optional Stage-1 EDF scheduler. Requires top-level + # active_stream_window > 0. It estimates each stream's playback deadline + # from actual PCM emitted by Code2Wav and reselects the K most urgent + # ready chunks at every scheduling turn. Use it with codec_chunk_ramp or + # an initial chunk large enough to cover the first follow-up decode. + # active_stream_policy: edf + # playback_deadline_first_chunk_ms: 100 + # playback_deadline_safety_margin_ms: 250 # Chunk ramp (opt-in): gradual size increase for first N chunks to # eliminate the IC→steady cliff (1→25). Uncomment to enable. # Each entry is the size of that chunk index; indices past the table diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 9bcc74257..a585a8729 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -2,9 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import importlib +import math import time from collections import defaultdict, deque from collections.abc import Callable, Mapping +from dataclasses import dataclass from typing import Any import torch @@ -22,6 +24,13 @@ from .base import OmniTransferAdapterBase logger = get_connector_logger(__name__) +@dataclass +class _PlaybackDeadlineState: + first_seen_at: float + playback_deadline: float | None = None + emitted_audio_s: float = 0.0 + + class OmniChunkTransferAdapter(OmniTransferAdapterBase): """Chunk-level transfer adapter for Omni connector pipelines. @@ -42,25 +51,54 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): def __init__(self, vllm_config: Any): model_config = vllm_config.model_config + self.model_mode = getattr(model_config, "worker_type", None) or "ar" + self.engine_output_type = getattr(model_config, "engine_output_type", None) self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs active_stream_window = int(getattr(model_config, "active_stream_window", 0) or 0) model_max_num_seqs = int(getattr(model_config, "max_num_seqs", self.scheduler_max_num_seqs) or 0) if model_max_num_seqs <= 0: model_max_num_seqs = self.scheduler_max_num_seqs self._active_window = min(active_stream_window, model_max_num_seqs) if active_stream_window > 0 else 0 + connector_config = getattr(model_config, "stage_connector_config", None) + raw_extra = ( + connector_config.get("extra", connector_config) + if isinstance(connector_config, dict) + else getattr(connector_config, "extra", None) + ) + extra = raw_extra if isinstance(raw_extra, dict) else {} + self._active_stream_policy = str(extra.get("active_stream_policy", "fifo")).strip().lower() + if self._active_stream_policy not in {"fifo", "edf"}: + raise ValueError(f"active_stream_policy must be 'fifo' or 'edf', got {self._active_stream_policy!r}") + if self._active_stream_policy == "edf" and self._active_window <= 0: + raise ValueError("active_stream_policy='edf' requires active_stream_window > 0") + self._playback_first_chunk_slo_s = float(extra.get("playback_deadline_first_chunk_ms", 100.0)) / 1000.0 + self._playback_safety_margin_s = float(extra.get("playback_deadline_safety_margin_ms", 50.0)) / 1000.0 + if not math.isfinite(self._playback_first_chunk_slo_s) or self._playback_first_chunk_slo_s <= 0: + raise ValueError("playback_deadline_first_chunk_ms must be positive") + if not math.isfinite(self._playback_safety_margin_s) or self._playback_safety_margin_s < 0: + raise ValueError("playback_deadline_safety_margin_ms must be non-negative") + self._playback_deadlines: dict[str, _PlaybackDeadlineState] = {} if self._active_window > 0: logger.info( - "Bounded active-stream window enabled: K=%d. " + "Bounded active-stream window enabled: K=%d policy=%s. " "Multi-replica deployments require sticky per-stream routing across Stage 1 " "replicas (each replica owns an independent active-set; without sticky routing, " "a stream can be active on one replica and non-active on another and both will " "race to evict it).", self._active_window, + self._active_stream_policy, ) self.connector = self.create_connector(model_config) self.receives_chunks = stage_receives_chunks(model_config) + if self.playback_deadline_enabled and not extra.get("codec_chunk_ramp"): + initial_frames = int(extra.get("initial_codec_chunk_frames") or 0) + if 0 < initial_frames < 4: + logger.warning( + "Playback EDF is enabled with initial_codec_chunk_frames=%d and no codec_chunk_ramp; " + "the first audio buffer may be too short to avoid a chunk-0 underrun.", + initial_frames, + ) super().__init__(model_config) - self.model_mode = getattr(model_config, "worker_type", None) or "ar" # State specific to Chunk management self.custom_process_next_stage_input_func: Callable[..., OmniPayloadStruct | None] | None = None custom_process_next_stage_input_func = getattr(model_config, "custom_process_next_stage_input_func", None) @@ -103,6 +141,51 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # OmniSchedulingCoordinator._waiting_since on the full-payload path. self._waiting_since: dict[str, float] = {} + @property + def playback_deadline_enabled(self) -> bool: + return ( + self._active_stream_policy == "edf" + and self.model_mode == "generation" + and self.engine_output_type == "audio" + ) + + def record_audio_output( + self, + request_id: str, + audio_frames: int, + sample_rate: int, + *, + emitted_at: float | None = None, + ) -> None: + """Advance the server-estimated playback deadline after one PCM chunk.""" + if not self.playback_deadline_enabled or audio_frames <= 0 or sample_rate <= 0: + return + now = time.monotonic() if emitted_at is None else float(emitted_at) + state = self._playback_deadlines.setdefault(request_id, _PlaybackDeadlineState(first_seen_at=now)) + duration_s = float(audio_frames) / float(sample_rate) + state.emitted_audio_s += duration_s + # A late chunk restarts a starved player at its server emission time; + # otherwise append its duration to the buffered playback horizon. + playback_base = max(state.playback_deadline or now, now) + state.playback_deadline = playback_base + duration_s + + def _playback_deadline_key(self, request: Request, order: int, now: float) -> tuple[int, float, float, int]: + state = self._playback_deadlines.setdefault( + request.request_id, + _PlaybackDeadlineState(first_seen_at=now), + ) + if state.playback_deadline is None: + deadline = state.first_seen_at + self._playback_first_chunk_slo_s + protected_playback = False + else: + deadline = state.playback_deadline + protected_playback = deadline - now <= self._playback_safety_margin_s + # A stream whose buffered playback is close to exhaustion must not be + # preempted by a large backlog of overdue first-audio deadlines. Once + # established streams have enough headroom, ordinary EDF ordering lets + # new requests consume the available slots. + return (0 if protected_playback else 1), deadline, state.first_seen_at, order + @staticmethod def _is_truthy_scalar(value: Any) -> bool: if isinstance(value, torch.Tensor): @@ -465,6 +548,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): Idempotent: calling with an already-cleaned or unknown id is safe. """ self._active_streams.pop(request_id, None) + self._playback_deadlines.pop(request_id, None) self.upstream_exhausted_requests.discard(request_id) self.segment_finished_requests.discard(request_id) self.get_req_chunk.pop(request_id, None) @@ -586,6 +670,19 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): waiting_queue.prepend_requests([request]) return + if self.playback_deadline_enabled: + self._process_chunk_queue_legacy( + waiting_queue, self.waiting_for_chunk_waiting_requests, RequestStatus.WAITING, self._finished_load_reqs + ) + self._process_chunk_queue_legacy( + running_queue, + self.waiting_for_chunk_running_requests, + RequestStatus.RUNNING, + self._finished_load_reqs, + ) + self._select_playback_deadline_streams(waiting_queue, running_queue) + return + self._promote_active_streams(running_queue) self._promote_active_streams(waiting_queue) self._process_chunk_queue( @@ -624,6 +721,45 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # Iterating the existing queue preserves FIFO admission. self._active_streams[request_id] = request + def _select_playback_deadline_streams(self, waiting_queue: Any, running_queue: list[Request]) -> None: + """Select up to K ready chunks by server-estimated playback deadline.""" + now = time.monotonic() + candidates = [ + request + for request in (*running_queue, *list(waiting_queue)) + if request.request_id in self.requests_with_ready_chunks + ] + ordered = sorted( + enumerate(candidates), + key=lambda item: self._playback_deadline_key(item[1], item[0], now), + ) + selected = [request for _, request in ordered[: self._active_window]] + selected_ids = {request.request_id for request in selected} + self._active_streams = {request.request_id: request for request in selected} + + for request in list(running_queue): + if request.request_id in self.requests_with_ready_chunks and request.request_id not in selected_ids: + running_queue.remove(request) + self._held_non_active.append(request) + + held_waiting = [] + for request in list(waiting_queue): + if request.request_id in self.requests_with_ready_chunks and request.request_id not in selected_ids: + waiting_queue.remove(request) + self.requests_origin_status[request.request_id] = RequestStatus.WAITING + self.waiting_for_chunk_waiting_requests.append(request) + elif request.request_id in selected_ids: + held_waiting.append(request) + + # Keep selected running requests in EDF order so a tight token budget + # still admits the most urgent chunk first. + running_order = {request.request_id: index for index, request in enumerate(selected)} + running_queue.sort(key=lambda request: running_order.get(request.request_id, len(running_order))) + if len(held_waiting) > 1: + held_waiting.sort(key=lambda request: running_order[request.request_id]) + waiting_queue.remove_requests(held_waiting) + waiting_queue.prepend_requests(held_waiting) + def _ensure_active_stream(self, request: Request) -> bool: if self._active_window <= 0: return True diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index da5fa840d..501296a27 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -29,6 +29,17 @@ from vllm_omni.model_executor.stage_input_processors.tts_utils import ( logger = init_logger(__name__) +def _empty_finished_payload(request_id: str) -> OmniPayloadStruct: + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct( + request_id=request_id, + left_context_size=0, + finished=torch.tensor(True, dtype=torch.bool), + ), + ) + + def _qwen3_tts_degenerate_finished_payload(): """Single-placeholder-frame finished payload for a degenerate talker take. @@ -162,14 +173,7 @@ def talker2code2wav_async_chunk( if length <= 0: if finished: - return OmniPayloadStruct( - codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), - meta=MetaStruct( - request_id=request_id, - left_context_size=0, - finished=torch.tensor(True, dtype=torch.bool), - ), - ) + return _empty_finished_payload(request_id) return None if ramp is not None: @@ -179,18 +183,28 @@ def talker2code2wav_async_chunk( if not emit: return None if context_length == 0: - return OmniPayloadStruct( - codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), - meta=MetaStruct( - request_id=request_id, - left_context_size=0, - finished=torch.tensor(True, dtype=torch.bool), - ), - ) + return _empty_finished_payload(request_id) else: - first_chunk = int(transfer_manager.put_req_chunk.get(request_id, 0)) <= 0 + chunk_index = int(transfer_manager.ramp_chunk_count.get(request_id, 0)) + first_chunk = chunk_index <= 0 use_first_chunk = initial_chunk_size > 0 and initial_chunk_size < chunk_size + if use_first_chunk: + emitted_frames = 0 if first_chunk else initial_chunk_size + (chunk_index - 1) * chunk_size + else: + emitted_frames = chunk_index * chunk_size + if finished and length <= emitted_frames: + # EOS is emitted on a separate talker step. If the preceding codec + # frame completed an exact chunk boundary, that chunk has already + # been sent; terminal must carry only an empty finish marker. + logger.debug( + "Qwen3-TTS terminal has no unsent codec frames: request=%s length=%d emitted=%d", + request_id, + length, + emitted_frames, + ) + return _empty_finished_payload(request_id) + if use_first_chunk and length <= initial_chunk_size: if not finished and length < initial_chunk_size: return None