Compare commits

...
Author SHA1 Message Date
Sy03 5740ca8e01 [Perf][TTS] Add deadline-aware audio scheduling
Signed-off-by: Sy03 <1370724210@qq.com>
2026-08-25 12:35:35 +08:00
8 changed files with 559 additions and 19 deletions

No files matched your search

@@ -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-<field> <value>`.
- 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.
---
@@ -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()
@@ -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)
@@ -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.
@@ -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
+8
View File
@@ -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
@@ -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
@@ -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