diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index a012cf21f..a13ab52df 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/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 diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py new file mode 100644 index 000000000..f6539a40a --- /dev/null +++ b/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 diff --git a/tests/core/sched/test_omni_scheduler_mixin.py b/tests/core/sched/test_omni_generation_scheduler_update_session.py similarity index 79% rename from tests/core/sched/test_omni_scheduler_mixin.py rename to tests/core/sched/test_omni_generation_scheduler_update_session.py index e04a9c39f..cf492e7cc 100644 --- a/tests/core/sched/test_omni_scheduler_mixin.py +++ b/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 diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 24ba819b1..8bfe837e2 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/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] diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 9e050a3f8..55e3edc41 100644 --- a/tests/engine/test_orchestrator.py +++ b/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", diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index 00a6565d9..39b27103e 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/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", ), diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py index a04a75de8..4b2bdf3fe 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py +++ b/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", diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 3d51765a9..5c0c489d1 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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 diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 8d6655742..dae3be8ae 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/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) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 570fa5545..c4b1e2429 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/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) diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index b239ff5ed..09316b55e 100644 --- a/vllm_omni/data_entry_keys.py +++ b/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 diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index 79a824aa3..db253ebad 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/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) 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 5686c7696..0b6e0b572 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 @@ -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) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 48026abab..791b6d5d9 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index d8d59285d..c13001063 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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() diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py index f2ce75df2..e04e8cdf9 100644 --- a/vllm_omni/entrypoints/openai/realtime_connection.py +++ b/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: diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 1f637fcd4..6bd8bf305 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/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( diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py index ca756b794..c9a0f266d 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index b671951b2..aa98c2141 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/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(): diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index a97cc2b65..0663673d2 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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)