[Feat]audio streaming input for async chunk (#3614)

Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>

examples/online_serving/qwen3_omni/README.md

@@ -18,10 +18,6 @@ Asynchronous chunk streaming operates as **enabled by default** within this bund
 Additionally, NPU, ROCm, and XPU per-platform configuration deltas are deterministically merged from the
 `platforms`: section of the corresponding YAML.
 
-**Note:** The OpenAI-style **`/v1/realtime`** WebSocket interface (facilitating streaming PCM audio input alongside audio and transcription output)
-is currently **unsupported** while the `async_chunk` configuration attribute is enabled.
-It is requisite to instantiate the default omni architecture or utilize a deployment configuration specifying `async_chunk: false` to facilitate real-time streaming sessions.
-
 To explicitly utilize a custom deployment YAML, mandate the configuration path accordingly:
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
@@ -289,7 +285,7 @@ python openai_realtime_client.py \
 | `--num-requests` | `1` | Number of sequential sessions (see `--concurrency`) |
 | `--concurrency` | `1` | Max concurrent WebSocket sessions when `--num-requests` > 1 |
 
-Ensure the server is running **without** `async_chunk` if you use `/v1/realtime`, for example:
+Ensure the server is running, for example:
 
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -0,0 +1,92 @@
+"""Unit tests for Omni AR streaming-session async placeholder handling."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request before
+# Request / StreamingUpdate are bound in this module. Ruff isort would reorder them.
+# isort: off
+import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
+from vllm.sampling_params import SamplingParams
+from vllm.v1.request import Request, RequestStatus, StreamingUpdate
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
+    sched = OmniARScheduler.__new__(OmniARScheduler)
+    sched._new_prompt_len_snapshot = {}
+    sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id))
+    sched.num_waiting_for_streaming_input = 0
+    sched.log_stats = False
+    sched.chunk_transfer_adapter = None
+    return sched
+
+
+def _make_request() -> Request:
+    return Request(
+        request_id="req-ar-streaming-test",
+        prompt_token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=8),
+        pooling_params=None,
+        arrival_time=100.0,
+        block_hasher=None,
+    )
+
+
+def _make_update(prompt_token_ids: list[int] | None = None) -> StreamingUpdate:
+    return StreamingUpdate(
+        mm_features=None,
+        prompt_token_ids=[10, 20] if prompt_token_ids is None else prompt_token_ids,
+        max_tokens=32,
+        arrival_time=200.0,
+        sampling_params=SamplingParams(max_tokens=16),
+    )
+
+
+def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
+    sched = _make_scheduler(stage_id=0)
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    session.append_output_token_ids([7, 8, 9])
+    session.num_computed_tokens = 6
+    session.num_output_placeholders = 1
+    session.spec_token_ids = [-1]
+
+    sched._update_request_as_session(session, _make_update([10, 20]))
+
+    assert session.discard_latest_async_tokens is True
+    assert session.num_output_placeholders == 0
+    assert session.spec_token_ids == []
+    # The async placeholder makes token 9 unconfirmed, so only 7 and 8 are
+    # carried into the next streaming prompt before the new chunk tokens.
+    assert session.prompt_token_ids == [1, 2, 3, 7, 8, 10, 20]
+    assert list(session._all_token_ids) == [1, 2, 3, 7, 8, 10, 20]
+    assert session._output_token_ids == []
+    assert session.num_prompt_tokens == 7
+    assert sched._new_prompt_len_snapshot[session.request_id] == 2
+
+
+def test_stage0_streaming_update_keeps_all_computed_tokens_without_placeholder() -> None:
+    sched = _make_scheduler(stage_id=0)
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    session.append_output_token_ids([7, 8, 9])
+    session.num_computed_tokens = 6
+    session.num_output_placeholders = 0
+
+    sched._update_request_as_session(session, _make_update([10, 20]))
+
+    assert session.discard_latest_async_tokens is False
+    assert session.num_output_placeholders == 0
+    assert session.prompt_token_ids == [1, 2, 3, 7, 8, 9, 10, 20]
+    assert list(session._all_token_ids) == [1, 2, 3, 7, 8, 9, 10, 20]
+    assert session._output_token_ids == []
+    assert session.num_prompt_tokens == 8
+    assert sched._new_prompt_len_snapshot[session.request_id] == 2

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -1,6 +1,6 @@
-"""Unit tests for OmniSchedulerMixin streaming session replacement.
+"""Unit tests for generation streaming session replacement.
 
-These tests pin the behavior of `_replace_session_with_streaming_update` against
+These tests pin the behavior of `_update_request_as_session` against
 current vLLM `Request` / `StreamingUpdate` (and Omni patches). When upgrading
 vLLM, failures here should highlight incompatible changes to request state or
 update payloads early.
@@ -19,19 +19,34 @@ import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine import EngineCoreEventType
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
-from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 
 # isort: on
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-class _SchedulerStub(OmniSchedulerMixin):
-    """Minimal scheduler surface required by OmniSchedulerMixin."""
+class _SkippedWaitingStub:
+    def __contains__(self, request: Request) -> bool:
+        return False
+
+
+class _ChunkTransferAdapterStub:
+    def __init__(self) -> None:
+        self.segment_finished_requests: set[str] = set()
+
+
+class _SchedulerStub(OmniGenerationScheduler):
+    """Minimal scheduler surface required by OmniGenerationScheduler."""
 
     def __init__(self, *, log_stats: bool = False) -> None:
         self.num_waiting_for_streaming_input = 0
         self.log_stats = log_stats
+        self.chunk_transfer_adapter = _ChunkTransferAdapterStub()
+        self.skipped_waiting = _SkippedWaitingStub()
+
+    def _enqueue_waiting_request(self, session: Request) -> None:
+        raise AssertionError("unexpected enqueue for skipped_waiting miss")
 
 
 def _make_request(**kwargs) -> Request:
@@ -71,7 +86,7 @@ class TestReplaceSessionWithStreamingUpdate:
 
         update = _make_update(prompt_token_ids=[40, 41, 42])
         sched.num_waiting_for_streaming_input = 3
-        sched._replace_session_with_streaming_update(session, update)
+        sched._update_request_as_session(session, update)
 
         assert session._output_token_ids == []
         assert list(session._all_token_ids) == [40, 41, 42]
@@ -88,7 +103,7 @@ class TestReplaceSessionWithStreamingUpdate:
         session = _make_request()
         session.status = RequestStatus.RUNNING
         update = _make_update(prompt_token_ids=None)
-        sched._replace_session_with_streaming_update(session, update)
+        sched._update_request_as_session(session, update)
 
         assert session.prompt_token_ids == ()
         assert list(session._all_token_ids) == []
@@ -108,7 +123,7 @@ class TestReplaceSessionWithStreamingUpdate:
             pytest.skip("StreamingUpdate has no additional_information (Omni patch inactive?)")
         update = replace(base, additional_information=None)
 
-        sched._replace_session_with_streaming_update(session, update)
+        sched._update_request_as_session(session, update)
         assert session.additional_information is None
 
     def test_does_not_decrement_waiting_when_not_streaming_status(self) -> None:
@@ -116,14 +131,14 @@ class TestReplaceSessionWithStreamingUpdate:
         session = _make_request()
         session.status = RequestStatus.RUNNING
         sched.num_waiting_for_streaming_input = 5
-        sched._replace_session_with_streaming_update(session, _make_update())
+        sched._update_request_as_session(session, _make_update())
         assert sched.num_waiting_for_streaming_input == 5
 
     def test_records_queued_event_when_log_stats_enabled(self) -> None:
         sched = _SchedulerStub(log_stats=True)
         session = _make_request()
         session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
-        sched._replace_session_with_streaming_update(session, _make_update())
+        sched._update_request_as_session(session, _make_update())
 
         assert session.events
         assert session.events[-1].type == EngineCoreEventType.QUEUED

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -230,9 +230,11 @@ def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, m
     cleanup_calls = []
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
 
-    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": False})
+    adapter._send_single_request(
+        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
+    )
 
-    assert cleanup_calls == []  # no terminal cleanup; meta.finished is unobservable
+    assert cleanup_calls == []  # no terminal cleanup; meta.finished is false
 
 
 def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatch):
@@ -248,12 +250,50 @@ def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatc
     adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
 
-    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": False})
+    adapter._send_single_request(
+        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
+    )
 
     assert connector.put.called
     sent_payload = connector.put.call_args.kwargs["data"]
     assert isinstance(sent_payload, OmniPayloadStruct)
-    assert sent_payload.meta is None  # confirms it's the empty struct on the wire
+    assert sent_payload.meta.finished.item() is False
+    assert sent_payload.meta.is_segment_finished.item() is False
+
+
+def test_send_single_request_struct_preserves_segment_finished(build_adapter, monkeypatch):
+    adapter, connector = build_adapter(stage_id=1)
+    request = _req("req-segment", RequestStatus.WAITING, external_req_id="ext-segment")
+
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct()
+    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
+
+    adapter._send_single_request(
+        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
+    )
+
+    sent_payload = connector.put.call_args.kwargs["data"]
+    assert sent_payload.meta.finished.item() is False
+    assert sent_payload.meta.is_segment_finished.item() is True
+
+
+def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
+    request.resumable = True
+    request.num_computed_tokens = 0
+    adapter.requests_num_chunks_sent["ext-stream"] = 111
+
+    adapter.save_async(pooling_output=None, request=request, is_segment_finished=False)
+
+    assert len(adapter._pending_save_reqs) == 0
+    assert adapter.requests_num_chunks_sent["ext-stream"] == 111
+
+    adapter.requests_num_chunks_sent.pop("ext-stream")
+    adapter.save_async(pooling_output=None, request=request, is_segment_finished=False)
+
+    assert len(adapter._pending_save_reqs) == 1
+    assert adapter.requests_num_chunks_sent["ext-stream"] == 0
 
 
 def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
@@ -266,7 +306,9 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon
     cleanup_calls = []
     monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
 
-    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": True})
+    adapter._send_single_request(
+        {"pooling_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
+    )
 
     assert len(cleanup_calls) == 1
     args, _ = cleanup_calls[0]

tests/engine/test_orchestrator.py

@@ -964,6 +964,7 @@ async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> Non
 
     pool = RecordingPool()
     orchestrator = object.__new__(Orchestrator)
+    orchestrator.async_chunk = False
     orchestrator.request_states = {
         "req-stream": OrchestratorRequestState(
             request_id="req-stream",

tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py

@@ -44,7 +44,6 @@ realtime_server_params = [
             model=MODEL,
             stage_config_path=default_stage_config,
             use_stage_cli=True,
-            server_args=["--no-async-chunk"],
         ),
         id="default",
     ),

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -20,7 +20,7 @@ def _streaming_context() -> SimpleNamespace:
 
 
 def test_get_streaming_talker_tokens_first_segment(_streaming_context: SimpleNamespace) -> None:
-    inc_p, inc_o, merged, thinker_in = q3._get_streaming_talker_tokens(
+    inc_p, inc_o = q3._get_streaming_talker_tokens(
         "r1",
         [1, 2],
         [10, 11],
@@ -28,13 +28,11 @@ def test_get_streaming_talker_tokens_first_segment(_streaming_context: SimpleNam
     )
     assert inc_p == [1, 2]
     assert inc_o == [10, 11]
-    assert merged == [1, 2, 10, 11]
-    assert thinker_in == [1, 2]
 
 
 def test_get_streaming_talker_tokens_second_segment_accumulates(_streaming_context: SimpleNamespace) -> None:
     q3._get_streaming_talker_tokens("r2", [1, 2], [10, 11], streaming_context=_streaming_context)
-    inc_p, inc_o, merged, thinker_in = q3._get_streaming_talker_tokens(
+    inc_p, inc_o = q3._get_streaming_talker_tokens(
         "r2",
         [1, 2, 3, 4],
         [10, 11, 12, 13],
@@ -42,14 +40,12 @@ def test_get_streaming_talker_tokens_second_segment_accumulates(_streaming_conte
     )
     assert inc_p == [3, 4]
     assert inc_o == [12, 13]
-    assert merged == [1, 2, 10, 3, 4, 12, 13]
-    assert thinker_in == [1, 2, 10, 3, 4]
 
 
 def test_get_streaming_talker_tokens_new_prompt_len_snapshot_truncates(
     _streaming_context: SimpleNamespace,
 ) -> None:
-    inc_p, inc_o, merged, thinker_in = q3._get_streaming_talker_tokens(
+    inc_p, inc_o = q3._get_streaming_talker_tokens(
         "r3",
         [1, 2, 3, 4, 5, 6],
         [10],
@@ -58,8 +54,6 @@ def test_get_streaming_talker_tokens_new_prompt_len_snapshot_truncates(
     )
     assert inc_p == [1, 2, 3, 4]
     assert inc_o == [10]
-    assert merged == [1, 2, 3, 4, 10]
-    assert thinker_in == [1, 2, 3, 4]
 
 
 def test_get_streaming_talker_tokens_clear_state(_streaming_context: SimpleNamespace) -> None:
@@ -82,6 +76,72 @@ def test_get_streaming_codec_delta_len_increments_and_finishes(_streaming_contex
     assert state.talker2code2wav_last_seq_len == 0
 
 
+def test_streaming_input_prefill_chunk_is_cached() -> None:
+    transfer_manager = SimpleNamespace(_pending_streaming_prefills={})
+    request = SimpleNamespace(
+        external_req_id="rt-1",
+        output_token_ids=[100],
+        all_token_ids=[151644, 872, 100],
+        prompt_token_ids=[151644, 872],
+        resumable=True,
+        additional_information=None,
+    )
+    thinker_emb = torch.ones(2, 3)
+    thinker_hid = torch.full((2, 3), 2.0)
+
+    payload = q3._construct_thinker2talker_streaming_input_async_chunk(
+        False,
+        request,
+        thinker_emb,
+        thinker_hid,
+        transfer_manager,
+    )
+
+    assert payload is None
+    cached = transfer_manager._pending_streaming_prefills["rt-1"]
+    assert torch.equal(cached["embed"]["prefill"], thinker_emb)
+    assert torch.equal(cached["hidden_states"]["output"], thinker_hid)
+    assert cached["ids"]["all"] == request.all_token_ids
+    assert cached["ids"]["prompt"] == request.prompt_token_ids
+
+
+def test_streaming_input_prefill_flushes_with_next_decode_chunk() -> None:
+    transfer_manager = SimpleNamespace(
+        _pending_streaming_prefills={
+            "rt-2": {
+                "embed": {"prefill": torch.ones(2, 3)},
+                "hidden_states": {"output": torch.full((2, 3), 2.0)},
+                "ids": {"all": [151644, 872, 100], "prompt": [151644, 872]},
+            }
+        }
+    )
+    request = SimpleNamespace(
+        external_req_id="rt-2",
+        output_token_ids=[101],
+        all_token_ids=[151644, 872, 100, 101],
+        prompt_token_ids=[151644, 872],
+        resumable=True,
+        additional_information=None,
+    )
+    thinker_emb = torch.full((1, 3), 3.0)
+    thinker_hid = torch.full((1, 3), 4.0)
+
+    payload = q3._construct_thinker2talker_streaming_input_async_chunk(
+        False,
+        request,
+        thinker_emb,
+        thinker_hid,
+        transfer_manager,
+    )
+
+    assert payload is not None
+    assert payload.embed.prefill.shape == (3, 3)
+    assert payload.hidden_states.output.shape == (3, 3)
+    assert payload.ids.all == [151644, 872, 100]
+    assert payload.ids.prompt == [151644, 872]
+    assert "rt-2" not in transfer_manager._pending_streaming_prefills
+
+
 def test_talker2code2wav_full_payload_filters_by_output_token_ids() -> None:
     request = SimpleNamespace(
         request_id="codec",

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -13,7 +13,7 @@ from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMSchedu
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
 from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
@@ -426,8 +426,19 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # Capture finish_reason BEFORE _handle_stopped_request, which may
                 # reset the status to WAITING for streaming requests that continue.
                 finish_reason = request.get_finished_reason()
-                is_segment_finished = request.is_finished() and request.resumable
+                is_segment_finished = True
                 finished = self._handle_stopped_request(request)
+                if not finished:
+                    # for streaming input request only
+                    if self.chunk_transfer_adapter:
+                        if self.vllm_config.model_config.stage_id != 0:
+                            # Downstream async-chunk stages receive real payloads from the
+                            # connector. This update only resumes polling for the next segment.
+                            self.chunk_transfer_adapter.segment_finished_requests.discard(request.request_id)
+                    request.discard_latest_async_tokens = True
+                    request.num_output_placeholders = 0
+                    request.spec_token_ids = []
+                    request._output_token_ids.clear()
                 if finished:
                     kv_transfer_params = self._free_request(request)
                 if status_before_stop == RequestStatus.RUNNING:
@@ -471,7 +482,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     )
                 )
                 if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.save_async(pooler_output, request)
+                    self.chunk_transfer_adapter.save_async(pooler_output, request, is_segment_finished)
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
@@ -622,11 +633,36 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         """
         req_id = session.request_id
         self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids)
-        if self.vllm_config.model_config.stage_id != 0:
-            self._replace_session_with_streaming_update(session, update)
+        outstanding_async_tokens = getattr(session, "num_output_placeholders", 0)
+        if outstanding_async_tokens > 0:
+            # Async scheduling may already have sampled the previous
+            # segment's next token. Drop that late token instead of
+            # appending it to the new streaming segment.
+            session.discard_latest_async_tokens = True
+            session.num_computed_tokens -= session.num_output_placeholders
+            session.num_output_placeholders = 0
+            session.spec_token_ids = []
+        if self.chunk_transfer_adapter:
+            self.chunk_transfer_adapter.requests_num_chunks_sent.pop(session.external_req_id, None)
+            if self.vllm_config.model_config.stage_id != 0:
+                # Downstream async-chunk stages receive real payloads from the
+                # connector. This update only resumes polling for the next segment.
+                self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id)
+                # Do not replace prompt/additional_information here; the next
+                # upstream chunk will populate them in chunk transfer adapter.
+                session.arrival_time = update.arrival_time
+                session.sampling_params = update.sampling_params
+                if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
+                    self.num_waiting_for_streaming_input -= 1
+                session.status = RequestStatus.WAITING
+                if session in self.skipped_waiting:
+                    self.skipped_waiting.remove_requests((session,))
+                    self._enqueue_waiting_request(session)
 
-        else:
-            super()._update_request_as_session(session, update)
+                if self.log_stats:
+                    session.record_event(EngineCoreEventType.QUEUED)
+                return
+        super()._update_request_as_session(session, update)
 
     def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None:
         # TODO(wzliu)! for offline mode, we should not end process until all data is transferred

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -108,9 +108,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 break
             # async_chunk: don't schedule placeholder tokens when no new chunk is available.
             if required_tokens <= 0:
-                if (
-                    self.chunk_transfer_adapter is not None
-                    and request.request_id in self.chunk_transfer_adapter.finished_requests
+                if self.chunk_transfer_adapter is not None and self.chunk_transfer_adapter.is_done_receiving_chunks(
+                    request.request_id
                 ):
                     self._pending_finish_reqs.append(request)
                 req_index += 1
@@ -160,7 +159,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             # async_chunk: wait for the first upstream chunk (don't start with placeholders).
             if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0:
-                if request.request_id in self.chunk_transfer_adapter.finished_requests:
+                if self.chunk_transfer_adapter.is_done_receiving_chunks(request.request_id):
                     self.waiting.pop_request()
                     self._pending_finish_reqs.append(request)
                     continue
@@ -489,7 +488,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 or (self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens)
                 or (
                     self.chunk_transfer_adapter is not None
-                    and request.request_id in self.chunk_transfer_adapter.finished_requests
+                    and self.chunk_transfer_adapter.is_done_receiving_chunks(request.request_id)
                     and request.num_computed_tokens >= len(request.prompt_token_ids)
                 )
             ):
@@ -507,6 +506,10 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     routed_experts = model_runner_output.routed_experts_dict[req_id]
                 finish_reason = request.get_finished_reason()
                 finished = self._handle_stopped_request(request)
+                if not finished:
+                    # for streaming input request only
+                    if self.chunk_transfer_adapter:
+                        self.chunk_transfer_adapter.segment_finished_requests.discard(req_id)
                 if finished:
                     kv_transfer_params = self._free_request(request)
                     if self.chunk_transfer_adapter is not None:
@@ -669,4 +672,26 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         Do not expend prompt id using update.
         """
-        self._replace_session_with_streaming_update(session, update)
+        if self.chunk_transfer_adapter:
+            self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id)
+        session._output_token_ids.clear()
+        session._all_token_ids.clear()
+        new_prompt = update.prompt_token_ids or ()
+        session._all_token_ids.extend(new_prompt)
+        session.num_computed_tokens = 0
+        session.prompt_token_ids = update.prompt_token_ids or ()
+        session.additional_information = update.additional_information or None
+        # Update block hashes for the new tokens.
+        session.update_block_hashes()
+        session.num_prompt_tokens = len(session.prompt_token_ids)
+        session.arrival_time = update.arrival_time
+        session.sampling_params = update.sampling_params
+        if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
+            self.num_waiting_for_streaming_input -= 1
+        session.status = RequestStatus.WAITING
+        if session in self.skipped_waiting:
+            self.skipped_waiting.remove_requests((session,))
+            self._enqueue_waiting_request(session)
+
+        if self.log_stats:
+            session.record_event(EngineCoreEventType.QUEUED)

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -5,8 +5,7 @@ from typing import Any
 
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.engine import EngineCoreEventType
-from vllm.v1.request import Request, RequestStatus, StreamingUpdate
+from vllm.v1.request import RequestStatus
 
 from vllm_omni.core.sched.output import OmniChunkRecvHandle, OmniSchedulerOutput
 
@@ -146,28 +145,3 @@ class OmniSchedulerMixin:
             finished_requests_needing_kv_transfer=finished_requests_needing_kv_transfer or {},
             pending_input_registrations=pending_input_registrations,
         )
-
-    def _replace_session_with_streaming_update(
-        self,
-        session: Request,
-        update: StreamingUpdate,
-    ) -> None:
-        """For streaming input: Replace an existing streaming session payload with the latest update."""
-        session._output_token_ids.clear()
-        session._all_token_ids.clear()
-        new_prompt = update.prompt_token_ids or ()
-        session._all_token_ids.extend(new_prompt)
-        session.num_computed_tokens = 0
-        session.prompt_token_ids = update.prompt_token_ids or ()
-        session.additional_information = update.additional_information or None
-        # Update block hashes for the new tokens.
-        session.update_block_hashes()
-        session.num_prompt_tokens = len(session.prompt_token_ids)
-        session.arrival_time = update.arrival_time
-        session.sampling_params = update.sampling_params
-        if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
-            self.num_waiting_for_streaming_input -= 1
-        session.status = RequestStatus.WAITING
-
-        if self.log_stats:
-            session.record_event(EngineCoreEventType.QUEUED)

vllm_omni/data_entry_keys.py

@@ -59,6 +59,7 @@ class Ids(TypedDict, total=False):
 
 class OmniPayloadMeta(TypedDict, total=False):
     finished: torch.Tensor
+    is_segment_finished: torch.Tensor
     stream_finished: torch.Tensor
     req_id: list[str]
     left_context_size: int
@@ -142,6 +143,7 @@ class IdsStruct(_StructBase):
 
 class MetaStruct(_StructBase):
     finished: torch.Tensor | None = None
+    is_segment_finished: torch.Tensor | None = None
     stream_finished: torch.Tensor | None = None
     req_id: list[str] | None = None
     left_context_size: int | None = None

vllm_omni/distributed/omni_connectors/adapter.py

@@ -214,3 +214,34 @@ def compute_talker_prompt_ids_length(prompt_ids: list[int]) -> int:
             pass
 
     return sum_user_len + assistant_len
+
+
+def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], request: Any) -> None:
+    """Update a downstream streaming request prompt from connector payload ids.
+
+    Async-chunk downstream stages are prewarmed before the real Talker prompt is
+    known. When a Thinker payload carries ``ids.prompt``, this helper:
+
+    * Preserves ``num_computed_tokens`` (the scheduler token watermark).
+    * Moves already-computed output tokens into ``prompt_token_ids``.
+    * Appends a new placeholder prompt slice sized from the upstream ids.
+    * Refreshes block hashes so the scheduler allocates KV slots for the
+      extended prompt without discarding prior computed state.
+    """
+    ids = payload_data.get("ids", {})
+    prompt_token_ids = ids.get("prompt", None)
+    if not prompt_token_ids:
+        return
+    num_computed_tokens = request.num_computed_tokens
+    kept_output_tokens = request._all_token_ids[request.num_prompt_tokens : num_computed_tokens]
+    del request._all_token_ids[num_computed_tokens:]
+    request._output_token_ids.clear()
+    assert request.prompt_token_ids is not None
+    # Extend prompt with kept output tokens.
+    request.prompt_token_ids.extend(kept_output_tokens)
+    next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
+    new_prompt = [0] * next_prompt_len
+    request._all_token_ids.extend(new_prompt or ())
+    request.prompt_token_ids.extend(new_prompt or ())
+    request.update_block_hashes()
+    request.num_prompt_tokens = len(request.prompt_token_ids)

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -9,8 +9,9 @@ from typing import Any
 import torch
 from vllm.v1.request import Request, RequestStatus
 
-from vllm_omni.data_entry_keys import OmniPayloadStruct, unflatten_payload
+from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload
 
+from ..adapter import construct_next_stage_streaming_input_prompt
 from ..factory import OmniConnectorFactory
 from ..utils.config import ConnectorSpec
 from ..utils.logging import get_connector_logger
@@ -54,6 +55,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.put_req_chunk: dict[str, int] = defaultdict(int)
         self.get_req_chunk: dict[str, int] = defaultdict(int)
         self.finished_requests: set[str] = set()
+        self.segment_finished_requests: set[str] = set()
         self.request_payload = {}
         self.code_prompt_token_ids: dict[str, list[torch.Tensor]] = defaultdict(list)
         self.request_ids_mapping: dict[str, str] = {}
@@ -63,6 +65,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.requests_with_ready_chunks = set()
         self.requests_origin_status = {}
         self.requests_num_chunks_sent: dict[str, int] = defaultdict(int)
+        self._pending_streaming_prefills: dict[str, dict] = {}
 
     @staticmethod
     def _is_truthy_scalar(value: Any) -> bool:
@@ -123,16 +126,26 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self,
         pooling_output: torch.Tensor | None = None,
         request: Request | None = None,
+        is_segment_finished: bool = False,
     ):
         """Build and enqueue one chunk for asynchronous sending.
 
         Payload extraction happens in ``_send_single_request`` on the
         background save_loop thread.
 
+        For streaming input request ``is_segment_finished`` marks the end
+        of the current realtime input segment. It is intentionally separate
+        from ``request.is_finished()``: a resumable `/v1/realtime` session
+        can finish one audio segment and later continue with another segment
+        under the same external request id. For other requests, it is the same
+        as ``request.is_finished()``.
+
         Args:
             pooling_output: Partial pooling output dictionary
             request: Request object
+            is_segment_finished: whether the segment of request is finished
         """
+        is_finished = request.is_finished() and not request.resumable
 
         confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request)
 
@@ -150,7 +163,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         task = {
             "pooling_output": pooling_output,
             "request": request,
-            "is_finished": request.is_finished(),
+            "is_finished": is_finished,
+            "is_segment_finished": is_segment_finished,
         }
         self._pending_save_reqs.append(task)
         with self._save_cond:
@@ -185,13 +199,24 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
             meta = payload_data.get("meta", {})
             payload_finished = self._is_truthy_scalar(meta.get("finished"))
+            payload_segment_finished = self._is_truthy_scalar(meta.get("is_segment_finished"))
             if self.model_mode == "ar":
                 request.additional_information = payload_data
+                if chunk_id > 0 and request.resumable:
+                    # For new streaming input segment, we should update prompt from payload
+                    construct_next_stage_streaming_input_prompt(payload_data, request)
+
                 if payload_finished:
                     self.finished_requests.add(req_id)
+                    request.resumable = False
+                if payload_segment_finished:
+                    self.segment_finished_requests.add(req_id)
             else:
                 if payload_finished:
                     self.finished_requests.add(req_id)
+                    request.resumable = False
+                if payload_segment_finished:
+                    self.segment_finished_requests.add(req_id)
 
                 new_ids = payload_data.get("codes", {}).get("audio")
                 has_tensor_codes = isinstance(new_ids, torch.Tensor)
@@ -249,6 +274,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         pooling_output = unflatten_payload(raw_po) if isinstance(raw_po, dict) else raw_po
         request = task["request"]
         is_finished = task["is_finished"]
+        is_segment_finished = task["is_segment_finished"]
         stage_id = self.connector.stage_id
         next_stage_id = stage_id + 1
         external_req_id = request.external_req_id
@@ -262,14 +288,23 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                     transfer_manager=self,
                     pooling_output=pooling_output,
                     request=request,
-                    is_finished=is_finished,
+                    # Existing processors use is_finished as a flush signal.
+                    is_finished=is_segment_finished,
                 )
 
             except Exception as e:
                 logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}")
 
         if payload_data is None:
-            return
+            if not (is_segment_finished or is_finished):
+                return
+            # Segment/request finish markers must still reach downstream even when
+            # the processor has no tensor payload.
+            payload_data = OmniPayloadStruct()
+        if payload_data.meta is None:
+            payload_data.meta = MetaStruct()
+        payload_data.meta.finished = torch.tensor(is_finished, dtype=torch.bool)
+        payload_data.meta.is_segment_finished = torch.tensor(is_segment_finished, dtype=torch.bool)
 
         success, size, metadata = self.connector.put(
             from_stage=str(stage_id),
@@ -298,13 +333,22 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             if is_payload_finished:
                 self.cleanup(request.request_id, external_req_id)
 
-        if is_finished:
+        if is_segment_finished:
             self.code_prompt_token_ids.pop(external_req_id, None)
             self.requests_num_chunks_sent.pop(external_req_id, None)
             cached_ic = getattr(self, "_cached_ic", None)
             if cached_ic is not None:
                 cached_ic.pop(external_req_id, None)
 
+    def is_done_receiving_chunks(self, request_id: str) -> bool:
+        """Return True if the request should stop polling upstream chunks.
+
+        Covers both the whole-request finish marker (``finished_requests``) and
+        the per-segment finish marker (``segment_finished_requests``) used while
+        waiting for the next streaming input slice.
+        """
+        return request_id in self.finished_requests or request_id in self.segment_finished_requests
+
     ########################################################################
     # Cleanup
     ########################################################################
@@ -319,6 +363,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         Idempotent: calling with an already-cleaned or unknown id is safe.
         """
         self.finished_requests.discard(request_id)
+        self.segment_finished_requests.discard(request_id)
         self.get_req_chunk.pop(request_id, None)
         self.requests_with_ready_chunks.discard(request_id)
         self.request_ids_mapping.pop(request_id, None)
@@ -339,6 +384,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.request_payload.pop(external_req_id, None)
         self.code_prompt_token_ids.pop(external_req_id, None)
         self.requests_num_chunks_sent.pop(external_req_id, None)
+        self._pending_streaming_prefills.pop(external_req_id, None)
 
         cached_ic = getattr(self, "_cached_ic", None)
         if cached_ic is not None:
@@ -459,7 +505,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                     # Requests that have loaded chunk from last round
                     # of schedule, but have not scheduled
                     continue
-                if request.request_id in self.finished_requests:
+                if self.is_done_receiving_chunks(request.request_id):
+                    request.additional_information = None
                     continue
                 # Requests that waiting for chunk
                 self.load_async(request)

vllm_omni/engine/orchestrator.py

@@ -416,7 +416,7 @@ class Orchestrator:
         stage_id = 0
         request_id = msg.request_id
         request = msg.prompt
-
+        final_stage_id = msg.final_stage_id
         req_state = self.request_states.get(request_id)
         if req_state is None:
             logger.warning(
@@ -449,6 +449,9 @@ class Orchestrator:
             prompt_text=msg.output_prompt_text,
         )
 
+        if self.async_chunk and stage_id == 0 and final_stage_id > 0:
+            await self._prewarm_async_chunk_stages(request_id, request, req_state)
+
     async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None:
         """Handle an add_companion_request message: submit companion to stage 0."""
         companion_id = msg.companion_id
@@ -1160,12 +1163,13 @@ class Orchestrator:
                 base_input["prompt_token_ids"] = [0] * next_prompt_len
                 base_input["multi_modal_data"] = None
                 base_input["mm_processor_kwargs"] = None
-
+                downstream_resumable = bool(getattr(stage0_request, "resumable", req_state.streaming.enabled))
                 request = build_engine_core_request_from_tokens(
                     request_id=request_id,
                     prompt=base_input,
                     params=params,
                     model_config=next_pool.stage_vllm_config.model_config,
+                    resumable=downstream_resumable,
                 )
                 request.external_req_id = request.request_id
                 await next_pool.submit_initial(request_id, req_state, request, prompt_text=None)

vllm_omni/entrypoints/openai/api_server.py

@@ -1381,22 +1381,6 @@ async def streaming_video_chat(websocket: WebSocket):
 @router.websocket("/v1/realtime")
 async def realtime_websocket(websocket: WebSocket):
     """WebSocket endpoint for OpenAI-style realtime interactions."""
-    engine_client = getattr(websocket.app.state, "engine_client", None)
-    if engine_client is not None and getattr(engine_client, "async_chunk", False):
-        await websocket.accept()
-        await websocket.send_json(
-            {
-                "type": "error",
-                "error": (
-                    "The /v1/realtime API is not supported when async_chunk is enabled on the server. "
-                    "Use a stage configuration with async_chunk disabled and restart the server before using "
-                    "this endpoint."
-                ),
-                "code": "unsupported",
-            }
-        )
-        await websocket.close()
-        return
     serving = getattr(websocket.app.state, "openai_serving_realtime", None)
     if serving is None:
         await websocket.accept()

vllm_omni/entrypoints/openai/realtime_connection.py

@@ -120,7 +120,6 @@ class RealtimeConnection(VllmRealtimeConnection):
         full_text = ""
         prompt_token_ids_len = 0
         completion_tokens_len = 0
-        last_prompt_token_ids_len = 0  # detect Stage-0 segment rollover
         self._realtime_audio_ref = None
 
         # Coerce cumulative outputs to delta outputs; this ensures
@@ -139,40 +138,26 @@ class RealtimeConnection(VllmRealtimeConnection):
             )
 
             async for output in result_gen:
-                # Handle delta texts; this is very similar to the client from vLLM
-                if output.outputs and len(output.outputs) > 0:
+                stage_id = getattr(output, "stage_id", None)
+                if stage_id == 0 and output.outputs:
                     first_output = output.outputs[0]
                     new_token_ids = list(first_output.token_ids)
-                    new_tokens_len = len(new_token_ids)
-
-                    cur_prompt_token_ids_len = len(output.prompt_token_ids or [])
-                    # Stage-0 segment rollover: buffer_realtime_audio may yield
-                    # multiple TokensPrompt segments and the second segment's
-                    # prompt_token_ids include the first segment's decoded output.
-                    # Clear accumulated full_text at the boundary so the
-                    # carried-over prefix is not duplicated in the final
-                    # TranscriptionDone.text.
-                    stage_id = getattr(output, "stage_id", None)
-                    if stage_id == 0 and cur_prompt_token_ids_len > last_prompt_token_ids_len > 0:
-                        full_text = ""
-                    last_prompt_token_ids_len = cur_prompt_token_ids_len
-
-                    if not prompt_token_ids_len and output.prompt_token_ids:
-                        prompt_token_ids_len = len(output.prompt_token_ids)
-
-                    if new_tokens_len:
+                    if new_token_ids:
                         input_stream.put_nowait(new_token_ids)
 
+                    if output.prompt_token_ids:
+                        prompt_token_ids_len = max(
+                            prompt_token_ids_len,
+                            len(output.prompt_token_ids),
+                        )
+
                     delta_text = first_output.text or ""
                     full_text += delta_text
+                    completion_tokens_len += len(new_token_ids)
 
-                    # append output to input if there was any delta text
                     if delta_text:
                         await self.send(TranscriptionDelta(delta=delta_text))
 
-                    completion_tokens_len += new_tokens_len
-
-                # Handle audio chunking; this is Omni specific
                 audio_chunks, sample_rate = self._extract_audio_chunks(output)
 
                 for chunk in audio_chunks:

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -183,11 +183,6 @@ class Qwen3OmniMoeForConditionalGeneration(
                 # sync stall; build_mm_cpu handles the eventual D2H at payload time.
                 ("codes", "audio"),
             }
-            # Keys that need to be accumulated across streaming inputs
-            self.streaming_accumulated_keys: set[tuple[str, str]] = {
-                ("embed", "prefill"),
-                ("hidden_states", "output"),
-            }
 
         elif self.model_stage == "code2wav":
             multimodal_config.skip_mm_profiling = True
@@ -685,7 +680,8 @@ class Qwen3OmniMoeForConditionalGeneration(
         is_prefill = bool(payload.get("_omni_is_prefill", span_len > 1))
         if is_prefill:
             num_computed_tokens = payload.get("_omni_num_computed_tokens")
-            if num_computed_tokens is not None:
+            request_resumable = meta.get("resumable", False)
+            if num_computed_tokens is not None and not request_resumable:
                 meta["num_processed_tokens"] = int(num_computed_tokens)
             input_ids, input_embeds, update_dict = self.talker_preprocess_prefill(input_ids, input_embeds, payload)
             code_predictor_codes = torch.zeros(

vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py

@@ -893,12 +893,12 @@ class Qwen3TTSPromptEmbedsBuilder:
 
             * ``talker_prompt``: ``[prompt_len, hidden]`` prefill embedding.
             * ``trailing_text_hidden``: ``[T, hidden]`` queue of text-side
-              embeddings consumed one-per-decode-step (streaming mode) or a
-              single pad row (non-streaming).
+                embeddings consumed one-per-decode-step (streaming mode) or a
+                single pad row (non-streaming).
             * ``ref_code_len``: Number of ref codec frames consumed when the
-              Base task ran in in-context mode (``None`` otherwise).
+                Base task ran in in-context mode (``None`` otherwise).
             * ``ref_code``: The ``[T, Q]`` ref codec tensor used for ICL
-              (``None`` outside Base / ICL).
+                (``None`` outside Base / ICL).
         """
         config = self._config
         talker_config = self._talker_config

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -261,6 +261,89 @@ def _resolve_tts_token_embedding(
 # =========================
 
 
+def _construct_thinker2talker_streaming_input_async_chunk(
+    is_finished: bool,
+    request,
+    thinker_emb,
+    thinker_hid,
+    transfer_manager,
+) -> OmniPayloadStruct | None:
+    """Build Thinker -> Talker payloads for realtime streaming input chunks.
+
+    A resumable realtime request reuses the same logical request id across
+    audio segments. The first streaming prefill chunk is cached and returns ``None`` so the
+    connector does not emit an incomplete downstream chunk. The following
+    decode chunk flushes that cached prefill together with the current Thinker
+    output, keeping Talker ids and tensor rows aligned.
+    """
+    request_id = request.external_req_id
+    output_token_ids = request.output_token_ids
+    # Convert ConstantList to regular list for OmniSerializer serialization
+    output_token_ids = _ensure_list(output_token_ids)
+    speaker = extract_speaker_from_request(request)
+    language = extract_language_from_request(request)
+    finished = torch.tensor(is_finished, dtype=torch.bool)
+    emb_cpu = thinker_emb.detach().cpu()
+    hid_cpu = thinker_hid.detach().cpu()
+
+    if output_token_ids:
+        if thinker_emb.shape[0] > 1:
+            # if thinker_emb.shape[0] > 1, new streaming input segment is added
+            # and will transfer prefill embeddings and hidden states to talker.
+            new_prompt_len = thinker_emb.shape[0]
+            payload = OmniPayloadStruct(
+                meta=MetaStruct(finished=finished),
+                embed=EmbeddingsStruct(prefill=emb_cpu),
+                hidden_states=HiddenStatesStruct(output=hid_cpu),
+                ids=IdsStruct(
+                    all=_ensure_list(request.all_token_ids[-new_prompt_len - 1 :]),
+                    prompt=_ensure_list(request.prompt_token_ids[-new_prompt_len:]),
+                ),
+                speaker=speaker,
+                language=language,
+            )
+            transfer_manager._pending_streaming_prefills[request_id] = to_dict(payload)
+            return None
+        else:
+            save_payload = transfer_manager._pending_streaming_prefills.pop(request_id, None)
+            if save_payload is not None:
+                saved_prefill = save_payload.get("embed", {}).get("prefill")
+                saved_output = save_payload.get("hidden_states", {}).get("output")
+                if isinstance(saved_prefill, torch.Tensor) and isinstance(saved_output, torch.Tensor):
+                    return OmniPayloadStruct(
+                        meta=MetaStruct(finished=finished),
+                        embed=EmbeddingsStruct(prefill=torch.cat((saved_prefill, emb_cpu), dim=0)),
+                        hidden_states=HiddenStatesStruct(output=torch.cat((saved_output, hid_cpu), dim=0)),
+                        ids=IdsStruct(
+                            all=save_payload.get("ids", {}).get("all"),
+                            prompt=save_payload.get("ids", {}).get("prompt"),
+                        ),
+                        speaker=speaker,
+                        language=language,
+                    )
+            return OmniPayloadStruct(
+                meta=MetaStruct(
+                    finished=finished,
+                ),
+                embed=EmbeddingsStruct(decode=emb_cpu),
+                hidden_states=HiddenStatesStruct(output=hid_cpu),
+                ids=IdsStruct(output=output_token_ids),
+                speaker=speaker,
+                language=language,
+            )
+    else:
+        if not is_finished:
+            # do not send async chunk mode placeholder token or embedding/hidden of the stop token
+            return None
+        return OmniPayloadStruct(
+            meta=MetaStruct(finished=finished),
+            embed=EmbeddingsStruct(decode=emb_cpu),
+            hidden_states=HiddenStatesStruct(output=hid_cpu),
+            speaker=speaker,
+            language=language,
+        )
+
+
 @dataclass
 class _Thinker2TalkerStreamingState:
     last_prompt_len: int = 0
@@ -295,16 +378,16 @@ def _get_streaming_talker_tokens(
     streaming_context: Any | None = None,
     *,
     clear_state: bool = False,
-) -> tuple[list[int], list[int], list[int], list[int]]:
-    """Return streaming token slices and merged token views for thinker->talker.
-       e.g. For the second streaming input request:
-       merged_sequences: [input_prompt 1, output_tokens 1[:-1], input_prompt 2, output_tokens 2]
-      thinker_input_ids: [input_prompt 1, output_tokens 1[:-1], input_prompt 2]
+) -> tuple[list[int], list[int]]:
+    """Return prompt/output token deltas for the current streaming segment.
+
+    In non-async-chunk streaming, Thinker's prompt may already include the
+    next input segment. Remove that new prompt tail before building the Talker
+    delta for the previous segment.
+
     Returns:
         inc_prompt: prompt token delta for this segment.
         inc_output: output token delta for this segment.
-        merged_sequences: full thinker_sequences to send downstream.
-        thinker_input_ids: full thinker_input_ids paired with merged_sequences.
     """
     state = _get_qwen3_streaming_state(request_id, streaming_context).thinker2talker
     if new_prompt_len_snapshot:
@@ -314,15 +397,6 @@ def _get_streaming_talker_tokens(
 
     inc_prompt = prompt_token_ids[state.last_prompt_len :]
     inc_output = output_token_ids[state.last_output_len :]
-    delta_sequences = inc_prompt + inc_output
-    cached_sequences = state.merged_sequences
-
-    merged_sequences = cached_sequences + delta_sequences
-    thinker_input_ids = cached_sequences + inc_prompt
-
-    # Persist history for next segment. Drop the latest sampled token to keep
-    # thinker_input_ids / thinker_sequences alignment with next-step append.
-    cached_sequences.extend(delta_sequences[:-1])
 
     state.last_prompt_len = cur_prompt_len
     state.last_output_len = cur_output_len
@@ -332,7 +406,7 @@ def _get_streaming_talker_tokens(
         state.last_output_len = 0
         state.merged_sequences.clear()
 
-    return inc_prompt, inc_output, merged_sequences, thinker_input_ids
+    return inc_prompt, inc_output
 
 
 def _get_streaming_codec_delta_len(
@@ -429,6 +503,10 @@ def thinker2talker_async_chunk(
                 transfer_manager.request_payload[request_id] = to_dict(payload)
                 return None
     else:
+        if request.resumable:
+            return _construct_thinker2talker_streaming_input_async_chunk(
+                is_finished, request, thinker_emb, thinker_hid, transfer_manager
+            )
         if thinker_emb.shape[0] > 1:
             logger.warning(
                 "Unexpected multiple embeddings in thinker2talker_async_chunk for chunk_id %d: "
@@ -567,7 +645,7 @@ def thinker2talker(
         output_ids = _ensure_list(output.cumulative_token_ids)
         is_streaming_session = bool(getattr(streaming_context, "enabled", False))
         if is_streaming_session:
-            prompt_token_ids, output_ids, thinker_sequences, thinker_input_ids = _get_streaming_talker_tokens(
+            prompt_token_ids, output_ids = _get_streaming_talker_tokens(
                 req_id,
                 prompt_token_ids,
                 output_ids,
@@ -575,9 +653,8 @@ def thinker2talker(
                 streaming_context,
                 clear_state=bool(getattr(thinker_output, "finished", False)),
             )
-        else:
-            thinker_sequences = prompt_token_ids + output_ids
-            thinker_input_ids = prompt_token_ids
+        thinker_sequences = prompt_token_ids + output_ids
+        thinker_input_ids = prompt_token_ids
         new_seq_length = len(prompt_token_ids + output_ids) - 1
         thinker_mm_raw = getattr(output, "multimodal_output", None)
         if not isinstance(thinker_mm_raw, dict):
@@ -683,7 +760,7 @@ def thinker2talker_token_only(
         output_ids = _ensure_list(output.cumulative_token_ids)
         is_streaming_session = bool(getattr(streaming_context, "enabled", False))
         if is_streaming_session:
-            prompt_token_ids, output_ids, thinker_sequences, thinker_input_ids = _get_streaming_talker_tokens(
+            prompt_token_ids, output_ids = _get_streaming_talker_tokens(
                 req_id,
                 prompt_token_ids,
                 output_ids,
@@ -691,9 +768,8 @@ def thinker2talker_token_only(
                 streaming_context,
                 clear_state=bool(getattr(thinker_output, "finished", False)),
             )
-        else:
-            thinker_sequences = prompt_token_ids + output_ids
-            thinker_input_ids = prompt_token_ids
+        thinker_sequences = prompt_token_ids + output_ids
+        thinker_input_ids = prompt_token_ids
         info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}}
         prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu")
 
@@ -759,6 +835,15 @@ def talker2code2wav_async_chunk(
     if isinstance(code_predictor_codes, torch.Tensor):
         if not code_predictor_codes.any():
             return None
+        sampling_params = getattr(request, "sampling_params", None)
+        stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or [])
+        stop_token_id = getattr(sampling_params, "stop_token_id", None)
+        if stop_token_id is not None:
+            stop_token_ids.add(stop_token_id)
+        first_codebook = int(code_predictor_codes[0, 0].item())
+        if first_codebook in stop_token_ids:
+            logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook)
+            return None
     else:
         code_tensor = torch.tensor(code_predictor_codes, dtype=torch.long)
         if not code_tensor.any():

vllm_omni/worker/gpu_model_runner.py

@@ -437,7 +437,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         for new_req_data in scheduler_output.scheduled_new_reqs:
             req_id = new_req_data.req_id
             if req_id in self.requests:
-                self._update_streaming_input_additional_info(new_req_data, req_id)
+                self._update_streaming_input_additional_info(req_id)
                 req_state = self._update_streaming_request(req_id, new_req_data)
                 reqs_to_add.append(req_state)
                 continue
@@ -1806,34 +1806,12 @@ class OmniGPUModelRunner(GPUModelRunner):
         logger.warning_once("_merge_additional_information_update is deprecated, use _update_intermediate_buffer")
         return self._update_intermediate_buffer(req_id, upd)
 
-    def _update_streaming_input_additional_info(self, new_req_data, req_id):
-        # For streaming input prefill case only. Update buffer from last segment input
+    def _update_streaming_input_additional_info(self, req_id):
+        # For streaming input prefill case only. Set num processed tokens = 0 for new segment input
         cached_additional_info = self.model_intermediate_buffer.get(req_id, {})
         if cached_additional_info:
-            payload_info = getattr(new_req_data, "additional_information", None)
-            inc_info = deserialize_additional_information(payload_info)
-            if isinstance(inc_info, dict) and inc_info:
-                accumulated_keys: set[tuple[str, str]] = set()
-                if hasattr(self, "model") and hasattr(self.model, "streaming_accumulated_keys"):
-                    accumulated_keys = self.model.streaming_accumulated_keys
-                merged_info = dict(cached_additional_info)
-                for key, value in inc_info.items():
-                    if isinstance(value, dict):
-                        existing_sub = merged_info.get(key)
-                        merged_sub = dict(existing_sub) if isinstance(existing_sub, dict) else {}
-                        for sk, sv in value.items():
-                            if (key, sk) in accumulated_keys and isinstance(sv, torch.Tensor):
-                                inc_tensor = sv.detach().to("cpu").contiguous()
-                                old_tensor = merged_sub.get(sk)
-                                if old_tensor is None:
-                                    merged_sub[sk] = inc_tensor
-                                else:
-                                    merged_sub[sk] = torch.cat((old_tensor, inc_tensor), dim=0)
-                            else:
-                                merged_sub[sk] = sv
-                        merged_info[key] = merged_sub
-                    else:
-                        merged_info[key] = value
-                merged_info.setdefault("meta", {})["num_processed_tokens"] = 0
-                self.model_intermediate_buffer[req_id] = merged_info
-                setattr(self.requests[req_id], "additional_information_cpu", merged_info)
+            merged_info = dict(cached_additional_info)
+            merged_info.setdefault("meta", {})["num_processed_tokens"] = 0
+            merged_info.setdefault("meta", {})["resumable"] = True
+            self.model_intermediate_buffer[req_id] = merged_info
+            setattr(self.requests[req_id], "additional_information_cpu", merged_info)