From c0674850306b05f969ad1ff3cc19785da4dd2894 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Mon, 1 Jun 2026 06:59:10 -0700 Subject: [PATCH] [Perf][TTS] Bounded-K active-stream window for Stage 1 (RFC #3535) (#3592) Signed-off-by: Yueqian Lin Co-authored-by: Hongsheng Liu --- .../examples/online_serving/text_to_speech.md | 1 + .../test_chunk_transfer_adapter.py | 78 ++++++++- tests/test_config_factory.py | 1 + vllm_omni/config/model.py | 2 + vllm_omni/config/stage_config.py | 4 + .../chunk_transfer_adapter.py | 154 +++++++++++++++++- vllm_omni/engine/arg_utils.py | 5 + 7 files changed, 240 insertions(+), 5 deletions(-) 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 bc1499b2d..9ee6b33fb 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -447,6 +447,7 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000 - Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above. - 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. --- diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index aa72133c9..f8eb6af8b 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -50,6 +50,7 @@ def build_adapter(monkeypatch, mocker: MockerFixture): stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2, + active_stream_window: int = 0, connector_extra: dict | None = None, ): connector = mocker.MagicMock() @@ -76,7 +77,11 @@ def build_adapter(monkeypatch, mocker: MockerFixture): classmethod(lambda cls, _model_config: connector), ) - model_config = SimpleNamespace(worker_type=model_mode) + model_config = SimpleNamespace( + worker_type=model_mode, + max_num_seqs=max_num_seqs, + active_stream_window=active_stream_window, + ) scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs) adapter = OmniChunkTransferAdapter( SimpleNamespace(model_config=model_config, scheduler_config=scheduler_config) @@ -396,6 +401,75 @@ def test_process_and_restore_queues(build_adapter): assert adapter.waiting_for_chunk_running_requests == deque() +def test_fifo_promotion(build_adapter): + adapter, _ = build_adapter(stage_id=1, max_num_seqs=2, active_stream_window=2) + reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)] + waiting_queue = DummyWaitingQueue(reqs) + running_queue = [] + + adapter.process_pending_chunks(waiting_queue, running_queue) + + assert list(adapter._active_streams) == ["req-1", "req-2"] + assert waiting_queue == reqs[2:] + assert reqs[0].status == RequestStatus.WAITING_FOR_CHUNK + assert reqs[1].status == RequestStatus.WAITING_FOR_CHUNK + assert reqs[2].status == RequestStatus.WAITING + + adapter.finished_requests.add("req-1") + # Eviction is deferred to postprocess_scheduler_output in the runtime path + # (commit c4d95fd9 — otherwise the terminal chunk deadlocks at c=8 K=2). + # Simulate it here so promotion can pick up the freed slot. + adapter._evict_finished_active_streams({"req-1"}) + adapter.process_pending_chunks(waiting_queue, running_queue) + + assert list(adapter._active_streams) == ["req-2", "req-3"] + # Promotion + chunk processing happen in the same call, so req-3 is + # already WAITING_FOR_CHUNK by the time we check. + assert reqs[2].status == RequestStatus.WAITING_FOR_CHUNK + + +def test_legacy_k0(build_adapter): + adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=0) + waiting_req = _req("waiting", RequestStatus.WAITING) + running_req_1 = _req("running-1", RequestStatus.RUNNING) + running_req_2 = _req("running-2", RequestStatus.RUNNING) + waiting_queue = DummyWaitingQueue([waiting_req]) + running_queue = [running_req_1, running_req_2] + + adapter.requests_with_ready_chunks.update({"running-1", "running-2"}) + adapter.process_pending_chunks(waiting_queue, running_queue) + + assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK + assert adapter.waiting_for_chunk_waiting_requests == deque([waiting_req]) + assert running_queue == [running_req_1] + assert waiting_queue == [running_req_2] + assert running_req_2.status == RequestStatus.PREEMPTED + assert adapter._active_streams == {} + + +def test_finished_releases_slot(build_adapter): + adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=1) + req_1 = _req("req-1", RequestStatus.WAITING) + req_2 = _req("req-2", RequestStatus.WAITING) + waiting_queue = DummyWaitingQueue([req_1, req_2]) + running_queue = [] + + adapter.process_pending_chunks(waiting_queue, running_queue) + assert list(adapter._active_streams) == ["req-1"] + assert waiting_queue == [req_2] + + adapter.finished_requests.add("req-1") + # Eviction is deferred to postprocess_scheduler_output in the runtime path + # (commit c4d95fd9). Simulate it so promotion can pick up the freed slot. + adapter._evict_finished_active_streams({"req-1"}) + adapter.process_pending_chunks(waiting_queue, running_queue) + + assert list(adapter._active_streams) == ["req-2"] + # Promotion + chunk processing happen in the same call, so req-2 is + # already WAITING_FOR_CHUNK by the time we check. + assert req_2.status == RequestStatus.WAITING_FOR_CHUNK + + def test_postprocess_scheduler_output(build_adapter): adapter, _ = build_adapter() adapter.requests_with_ready_chunks = {"new-ready", "cached-ready", "leftover"} @@ -422,6 +496,7 @@ def test_postprocess_scheduler_output(build_adapter): def _populate_adapter_state(adapter, req_id="req-1", ext_id="ext-1"): """Fill every per-request structure so cleanup can be verified.""" adapter.finished_requests.add(req_id) + adapter._active_streams[req_id] = SimpleNamespace(request_id=req_id) adapter.get_req_chunk[req_id] = 3 adapter.requests_with_ready_chunks.add(req_id) adapter.request_ids_mapping[req_id] = ext_id @@ -442,6 +517,7 @@ def test_cleanup_clears_all_state(build_adapter): adapter.cleanup(req_id, ext_id) assert req_id not in adapter.finished_requests + assert req_id not in adapter._active_streams assert req_id not in adapter.get_req_chunk assert req_id not in adapter.requests_with_ready_chunks assert req_id not in adapter.request_ids_mapping diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 01facfba1..ed0a6fa1d 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -1020,6 +1020,7 @@ class TestPipelineRegistry: class TestDeployConfigLoading: def test_deploy_override_fields_include_deploy_schema_fields(self): expected_fields = { + "active_stream_window", "async_chunk", # StageDeployConfig: stage placement and runtime fields. "devices", diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 9097f59bf..a5f1376e8 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -119,6 +119,8 @@ class OmniModelConfig(ModelConfig): stage_id: int = 0 async_chunk: bool = False + # Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin. + active_stream_window: int = 0 model_stage: str = "thinker" model_arch: str | None = None worker_type: str | None = None diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index cd0fe5dd5..5142ba234 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -500,6 +500,8 @@ class DeployConfig: """ async_chunk: bool = True + # Stage-1 active stream slots; 0 preserves legacy all-stream cycling. + active_stream_window: int = 0 connectors: dict[str, Any] | None = None edges: list[dict[str, Any]] | None = None stages: list[StageDeployConfig] = field(default_factory=list) @@ -677,6 +679,7 @@ def load_deploy_config(path: str | Path) -> DeployConfig: kwargs: dict[str, Any] = { "async_chunk": raw_dict.get("async_chunk", True), + "active_stream_window": int(raw_dict.get("active_stream_window", 0) or 0), "connectors": raw_dict.get("connectors", None), "edges": raw_dict.get("edges", None), "stages": stages, @@ -816,6 +819,7 @@ _PIPELINE_WIDE_ENGINE_FIELDS: tuple[str, ...] = ( "enable_chunked_prefill", "data_parallel_size", "pipeline_parallel_size", + "active_stream_window", "custom_voice_dir", ) 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 0b6e0b572..04fb81bbb 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 @@ -41,6 +41,20 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): def __init__(self, vllm_config: Any): model_config = vllm_config.model_config 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 + if self._active_window > 0: + logger.info( + "Bounded active-stream window enabled: K=%d. " + "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.connector = self.create_connector(model_config) super().__init__(model_config) self.model_mode = getattr(model_config, "worker_type", None) or "ar" @@ -64,6 +78,14 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.waiting_for_chunk_running_requests: deque[Any] = deque() self.requests_with_ready_chunks = set() self.requests_origin_status = {} + self._active_streams: dict[str, Any] = {} + # Private hold-queue for non-active running requests. Restored to + # running_queue inside restore_queues(). Avoids calling + # waiting_queue.prepend_requests mid-step, which trips vllm's + # per-step LogitsProcessor invariant + # ("Cannot register new removed request after self.removed has + # been read"). + self._held_non_active: deque[Any] = deque() self.requests_num_chunks_sent: dict[str, int] = defaultdict(int) self._pending_streaming_prefills: dict[str, dict] = {} @@ -362,6 +384,10 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): Idempotent: calling with an already-cleaned or unknown id is safe. """ + if request_id in self.finished_requests: + self._evict_finished_active_streams({request_id}) + else: + self._active_streams.pop(request_id, None) self.finished_requests.discard(request_id) self.segment_finished_requests.discard(request_id) self.get_req_chunk.pop(request_id, None) @@ -424,16 +450,112 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): """ if self.connector.stage_id == 0: return + if self._active_window <= 0: + 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, + ) + while len(running_queue) > self.scheduler_max_num_seqs: + request = running_queue.pop() + request.status = RequestStatus.PREEMPTED + waiting_queue.prepend_requests([request]) + return + + self._promote_active_streams(running_queue) + self._promote_active_streams(waiting_queue) self._process_chunk_queue( waiting_queue, self.waiting_for_chunk_waiting_requests, RequestStatus.WAITING, self._finished_load_reqs ) self._process_chunk_queue( running_queue, self.waiting_for_chunk_running_requests, RequestStatus.RUNNING, self._finished_load_reqs ) - while len(running_queue) > self.scheduler_max_num_seqs: - request = running_queue.pop() - request.status = RequestStatus.PREEMPTED - waiting_queue.prepend_requests([request]) + self._promote_active_streams(waiting_queue) + self._preempt_non_active_running(waiting_queue, running_queue) + + def _evict_finished_active_streams(self, request_ids: set[str] | None = None) -> None: + for request_id in list(self._active_streams): + if request_ids is not None and request_id not in request_ids: + continue + if request_id in self.finished_requests: + self._active_streams.pop(request_id, None) + + def _promote_active_streams(self, queue: Any) -> None: + if len(self._active_streams) >= self._active_window: + return + for request in list(queue): + if len(self._active_streams) >= self._active_window: + return + request_id = request.request_id + if request_id in self._active_streams or request_id in self.finished_requests: + continue + # Iterating the existing queue preserves FIFO admission. + self._active_streams[request_id] = request + + def _ensure_active_stream(self, request: Request) -> bool: + if self._active_window <= 0: + return True + request_id = request.request_id + if request_id in self._active_streams: + self._active_streams[request_id] = request + return True + if request_id in self.finished_requests or len(self._active_streams) >= self._active_window: + return False + self._active_streams[request_id] = request + return True + + def _preempt_non_active_running(self, waiting_queue: Any, running_queue: list[Request]) -> None: + # Hold non-active running requests in a private deque rather than + # routing them back through waiting_queue. Routing through the + # vllm RequestQueue mid-step triggers + # "Cannot register new removed request after self.removed has + # been read" + # in vllm.v1.sample.logits_processor.state when the persistent + # batch was already snapshotted. They are returned to + # running_queue in restore_queues() so the next scheduler tick + # re-evaluates them through _promote_active_streams. + index = len(running_queue) - 1 + while index >= 0: + request = running_queue[index] + if request.request_id in self._active_streams: + index -= 1 + continue + request = running_queue.pop(index) + self._held_non_active.append(request) + index -= 1 + + def _process_chunk_queue_legacy( + self, + queue: Any, + waiting_for_chunk_list: deque[Any], + target_status: RequestStatus, + finished_load_reqs: set[str], + ) -> None: + queue_snapshot = list(queue) + for request in queue_snapshot: + if request.status != RequestStatus.WAITING_FOR_CHUNK: + if request.request_id in self.requests_with_ready_chunks: + # Requests that have loaded chunk from last round + # of schedule, but have not scheduled + continue + if request.request_id in self.finished_requests: + continue + # Requests that waiting for chunk + self.load_async(request) + request.status = RequestStatus.WAITING_FOR_CHUNK + else: + if request.request_id in finished_load_reqs: + request.status = target_status + finished_load_reqs.remove(request.request_id) + self.requests_with_ready_chunks.add(request.request_id) + continue + queue.remove(request) + self.requests_origin_status[request.request_id] = target_status + waiting_for_chunk_list.append(request) def restore_queues( self, @@ -459,6 +581,10 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): running_queue.extend(live_running_requests) self.waiting_for_chunk_running_requests = deque() + if self._held_non_active: + running_queue.extend(self._held_non_active) + self._held_non_active = deque() + def postprocess_scheduler_output( self, scheduler_output: Any, @@ -475,7 +601,25 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if requests is not None: self.attach_cached_additional_information(scheduler_output, requests) + scheduled_req_ids = self._scheduled_request_ids(scheduler_output) self._clear_chunk_ready(scheduler_output) + if scheduled_req_ids: + # Terminal chunks must stay active until they are scheduled once. + self._evict_finished_active_streams(scheduled_req_ids) + + @staticmethod + def _scheduled_request_ids(scheduler_output: Any) -> set[str]: + req_ids: set[str] = set() + if scheduler_output.scheduled_new_reqs: + for req_data in scheduler_output.scheduled_new_reqs: + req_id = getattr(req_data, "req_id", None) + if req_id: + req_ids.add(req_id) + if scheduler_output.scheduled_cached_reqs: + for req_id in scheduler_output.scheduled_cached_reqs.req_ids: + if req_id: + req_ids.add(req_id) + return req_ids @staticmethod def attach_cached_additional_information(scheduler_output: Any, requests: dict[str, Request]) -> None: @@ -500,6 +644,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): ) -> None: queue_snapshot = list(queue) for request in queue_snapshot: + if not self._ensure_active_stream(request): + continue if request.status != RequestStatus.WAITING_FOR_CHUNK: if request.request_id in self.requests_with_ready_chunks: # Requests that have loaded chunk from last round diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 8d4388b5e..c4298298e 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -139,6 +139,10 @@ class OmniEngineArgs(EngineArgs): stage_connector_spec: dict[str, Any] = field(default_factory=dict) subtalker_sampling_params: dict[str, Any] | None = None async_chunk: bool = False + # WS-A: Stage-1 active stream slots. 0 = legacy preempt-everything. + # Must be declared here so engine_args dict propagation does not silently + # drop the value when constructing OmniEngineArgs from kwargs. + active_stream_window: int = 0 omni_kv_config: dict | None = None quantization_config: Any | None = None force_cutlass_fp8: bool | None = None @@ -339,6 +343,7 @@ class OmniEngineArgs(EngineArgs): # All kwargs below are Omni specific stage_id=self.stage_id, async_chunk=self.async_chunk, + active_stream_window=self.active_stream_window, model_stage=self.model_stage, model_arch=self.model_arch, worker_type=self.worker_type,