[Perf][TTS] Bounded-K active-stream window for Stage 1 (RFC #3535) (#3592)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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-<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.
 
 ---
 

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

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",

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

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",
 )
 

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

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,