From 87736a3cc50a96b5c6e83e631c7b4d3a256fcae0 Mon Sep 17 00:00:00 2001 From: NATURE Date: Sun, 30 Aug 2026 00:49:42 +0800 Subject: [PATCH] [Bugfix][MiniCPM-o] Bound full-duplex Talker context with sliding recompute (#6626) Signed-off-by: natureofnature --- .../test_omni_scheduler_mixin_timeouts.py | 34 +- .../test_chunk_transfer_adapter.py | 661 +++++++++++++++++- .../models/minicpmo_4_5/test_pipeline.py | 19 +- .../minicpmo_4_5/test_talker_batching.py | 154 ++++ .../test_minicpmo_4_5_omni.py | 5 + vllm_omni/core/sched/omni_scheduler_mixin.py | 23 + vllm_omni/data_entry_keys.py | 8 + vllm_omni/deploy/minicpmo_4_5.yaml | 7 + .../distributed/omni_connectors/adapter.py | 144 +++- .../omni_connectors/transfer_adapter/base.py | 15 +- .../chunk_transfer_adapter.py | 235 ++++++- .../models/minicpmo_4_5/__init__.py | 5 +- .../minicpmo_4_5/minicpmo_4_5_omni_tts.py | 146 +++- .../minicpmo_4_5_omni.py | 5 +- 14 files changed, 1389 insertions(+), 72 deletions(-) diff --git a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py index 1824d2809..e7d126cd9 100644 --- a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py +++ b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit coverage for the two stage-input timeout hooks. ``_process_pending_input_timeouts`` covers the full-payload transport; @@ -253,6 +253,38 @@ def test_log_failed_chunk_sends_noop_without_adapter(): scheduler._log_failed_chunk_sends() +class _FakeReceiveFailureAdapter: + def __init__(self, failures): + self._failures = dict(failures) + + def collect_failed_receive_request_ids(self): + failures, self._failures = self._failures, {} + return failures + + +class _FakeReceiveFailureScheduler(OmniSchedulerMixin): + def __init__(self, requests, adapter): + self.requests = requests + self.chunk_transfer_adapter = adapter + self.finish_calls = [] + + def finish_requests(self, req_ids, status): + self.finish_calls.append((set(req_ids), status)) + + +def test_process_chunk_receive_failures_delegates_to_finish_requests(): + request_id = "invalid-chunk" + adapter = _FakeReceiveFailureAdapter({request_id: "prompt exceeds limit"}) + scheduler = _FakeReceiveFailureScheduler({request_id: SimpleNamespace(request_id=request_id)}, adapter) + + scheduler._process_chunk_receive_failures() + + assert len(scheduler.finish_calls) == 1 + finished_ids, status = scheduler.finish_calls[0] + assert finished_ids == {request_id} + assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR") + + # --- R1.4: VLLM_OMNI_INPUT_WAIT_TIMEOUT_S is parsed at import time, so each of # these reloads the module under a different environment. One knob now gates the # deadline on both transports, which is why a silent 0 matters more than it did diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index ca518bd79..59780d5d1 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -57,6 +57,19 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None) return request +def _streaming_request(mocker: MockerFixture, num_computed_tokens: int) -> SimpleNamespace: + prompt_token_ids = [0] * num_computed_tokens + return SimpleNamespace( + _all_token_ids=prompt_token_ids.copy(), + _output_token_ids=[], + prompt_token_ids=prompt_token_ids, + num_computed_tokens=num_computed_tokens, + num_prompt_tokens=num_computed_tokens, + num_output_placeholders=0, + update_block_hashes=mocker.Mock(), + ) + + def test_streaming_payload_can_replace_placeholder_prompt(mocker: MockerFixture) -> None: request = SimpleNamespace( _all_token_ids=[0, 0, 7, 8], @@ -84,30 +97,219 @@ def test_streaming_payload_can_replace_placeholder_prompt(mocker: MockerFixture) request.update_block_hashes.assert_called_once_with() -def test_streaming_payload_can_append_exact_prompt_length(mocker: MockerFixture) -> None: +def test_turn_start_replacement_ignores_accumulated_prompt_capacity(mocker: MockerFixture) -> None: + request = _streaming_request(mocker, num_computed_tokens=4064) + payload = { + "ids": {"prompt": [1]}, + "meta": { + "replace_streaming_prompt": True, + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + }, + } + + replaced = construct_next_stage_streaming_input_prompt( + payload, + request, + max_model_len=4096, + previous_condition_len=10, + previous_condition_seq=7, + condition_seq=8, + recompute_previous_chunks=1, + ) + + assert replaced is True + assert request.prompt_token_ids == [0] * 10 + assert payload["meta"]["streaming_prompt_recompute"] is False + assert "streaming_prompt_previous_codes" not in payload["ids"] + + +@pytest.mark.parametrize("reserve", [True, 0, -1, 1.5, "26", None]) +def test_streaming_prompt_rejects_invalid_generation_reserve( + mocker: MockerFixture, + reserve: object, +) -> None: + request = _streaming_request(mocker, num_computed_tokens=10) + payload = { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": reserve, + }, + } + + with pytest.raises(ValueError, match="next_stage_generation_tokens must be a positive integer"): + construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096) + + +def test_capacity_managed_streaming_prompt_at_limit_still_appends(mocker: MockerFixture) -> None: + request = _streaming_request(mocker, num_computed_tokens=4060) + payload = { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + }, + } + + replaced = construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096) + + assert replaced is False + assert request.num_computed_tokens == 4060 + assert request.num_prompt_tokens == 4070 + + +def test_streaming_window_builds_one_chunk_recompute_recipe(mocker: MockerFixture) -> None: + previous_condition_len = 10 + previous_codes = list(range(25)) request = SimpleNamespace( - _all_token_ids=[0, 0, 7, 8], - _output_token_ids=[7, 8], - prompt_token_ids=[0, 0], - num_computed_tokens=4, - num_prompt_tokens=2, + _all_token_ids=[0] * 4039 + previous_codes + [999], + _output_token_ids=previous_codes + [999], + prompt_token_ids=[0] * 4039, + num_computed_tokens=4065, + num_prompt_tokens=4039, + num_output_placeholders=1, update_block_hashes=mocker.Mock(), ) payload = { - "ids": {"prompt": [1, 2, 3]}, - "meta": {"next_stage_prompt_len": 3}, + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + }, } - construct_next_stage_streaming_input_prompt(payload, request) + replaced = construct_next_stage_streaming_input_prompt( + payload, + request, + max_model_len=4096, + previous_condition_len=previous_condition_len, + previous_condition_seq=7, + condition_seq=8, + recompute_previous_chunks=1, + ) - assert request.prompt_token_ids == [0, 0, 7, 8, 0, 0, 0] - assert request._all_token_ids == [0, 0, 7, 8, 0, 0, 0] - assert request._output_token_ids == [] - assert request.num_computed_tokens == 4 - assert request.num_prompt_tokens == 7 + assert replaced is True + assert request.prompt_token_ids == [0] * 45 + assert request.num_prompt_tokens == 45 + assert request.num_computed_tokens == 0 + assert payload["ids"]["streaming_prompt_previous_codes"] == previous_codes + assert payload["meta"] == { + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + "streaming_prompt_recompute": True, + "streaming_condition_seq": 8, + } request.update_block_hashes.assert_called_once_with() +def test_streaming_window_recomputes_each_condition_without_accumulating(mocker: MockerFixture) -> None: + request = SimpleNamespace( + _all_token_ids=[0] * 10 + [101, 102], + _output_token_ids=[101, 102], + prompt_token_ids=[0] * 10, + num_computed_tokens=12, + num_prompt_tokens=10, + num_output_placeholders=0, + update_block_hashes=mocker.Mock(), + ) + second = { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 12, + "next_stage_generation_tokens": 26, + }, + } + + assert construct_next_stage_streaming_input_prompt( + second, + request, + max_model_len=4096, + previous_condition_len=10, + previous_condition_seq=0, + condition_seq=1, + recompute_previous_chunks=1, + ) + assert request.num_prompt_tokens == 24 + assert second["ids"]["streaming_prompt_previous_codes"] == [101, 102] + + request._all_token_ids.extend([201, 202, 203]) + request._output_token_ids.extend([201, 202, 203]) + request.num_computed_tokens = 27 + third = { + "ids": {"prompt": [2]}, + "meta": { + "next_stage_prompt_len": 8, + "next_stage_generation_tokens": 26, + }, + } + + assert construct_next_stage_streaming_input_prompt( + third, + request, + max_model_len=4096, + previous_condition_len=12, + previous_condition_seq=1, + condition_seq=2, + recompute_previous_chunks=1, + ) + assert request.num_prompt_tokens == 23 + assert third["ids"]["streaming_prompt_previous_codes"] == [201, 202, 203] + assert third["meta"]["streaming_condition_seq"] == 2 + + +def test_capacity_rollover_requires_explicit_window_contract(mocker: MockerFixture) -> None: + request = _streaming_request(mocker, num_computed_tokens=4064) + payload = { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + }, + } + + with pytest.raises(ValueError, match="requires window_size=1"): + construct_next_stage_streaming_input_prompt( + payload, + request, + max_model_len=4096, + previous_condition_len=10, + previous_condition_seq=0, + condition_seq=1, + ) + + +@pytest.mark.parametrize("next_stage_prompt_len", [0, -1]) +def test_capacity_managed_streaming_prompt_rejects_nonpositive_length( + mocker: MockerFixture, next_stage_prompt_len: int +) -> None: + request = _streaming_request(mocker, num_computed_tokens=0) + payload = { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": next_stage_prompt_len, + "next_stage_generation_tokens": 26, + }, + } + + with pytest.raises(ValueError, match="positive next_stage_prompt_len"): + construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096) + + +def test_capacity_managed_streaming_prompt_rejects_fresh_overflow(mocker: MockerFixture) -> None: + request = _streaming_request(mocker, num_computed_tokens=0) + payload = { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 4071, + "next_stage_generation_tokens": 26, + }, + } + + with pytest.raises(ValueError, match="exceeds max_model_len"): + construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096) + + @pytest.fixture def build_adapter(monkeypatch, mocker: MockerFixture): def _build( @@ -115,6 +317,10 @@ def build_adapter(monkeypatch, mocker: MockerFixture): stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2, + max_model_len: int = 0, + tts_max_model_len: int = 0, + tts_attention_type: str = "full_attention", + flat_tts_config: bool = False, active_stream_window: int = 0, connector_extra: dict | None = None, ): @@ -131,6 +337,10 @@ def build_adapter(monkeypatch, mocker: MockerFixture): self._cancelled_load_reqs = set() self._pending_save_reqs = deque() self._finished_save_reqs = set() + self._send_failures = {} + self._send_failure_lock = threading.Lock() + self._receive_failures = {} + self._receive_failure_lock = threading.Lock() self.stop_event = threading.Event() self._recv_cond = threading.Condition() self._save_cond = threading.Condition() @@ -142,9 +352,23 @@ def build_adapter(monkeypatch, mocker: MockerFixture): classmethod(lambda cls, _model_config: connector), ) + hf_config = SimpleNamespace( + model_type="minicpmtts", + max_position_embeddings=tts_max_model_len, + attention_type=tts_attention_type, + ) + if not flat_tts_config: + hf_config = SimpleNamespace( + tts_config=SimpleNamespace( + max_position_embeddings=tts_max_model_len, + attention_type=tts_attention_type, + ) + ) model_config = SimpleNamespace( worker_type=model_mode, max_num_seqs=max_num_seqs, + max_model_len=max_model_len, + hf_config=hf_config, active_stream_window=active_stream_window, stage_connector_config={ "name": "SharedMemoryConnector", @@ -160,6 +384,27 @@ def build_adapter(monkeypatch, mocker: MockerFixture): return _build +def _dequeue_load_entry(adapter, request): + """Register and take the entry as the background recv loop would.""" + adapter.load_async(request) + entry = adapter._registered_load_entries[request.request_id] + try: + adapter._pending_load_reqs.remove(entry) + except ValueError: + pass + return entry + + +@pytest.mark.parametrize( + ("attention_type", "expected_previous_chunks"), + [("full_attention", 0), ("sliding_recompute", 1)], +) +def test_talker_attention_policy_controls_streaming_recompute(build_adapter, attention_type, expected_previous_chunks): + adapter, _ = build_adapter(tts_attention_type=attention_type) + + assert adapter._streaming_prompt_previous_chunks == expected_previous_chunks + + @pytest.mark.parametrize( ("raw_cfg", "expected_name", "expected_extra"), [ @@ -200,15 +445,78 @@ def test_load_poll(build_adapter): "meta": {"finished": torch.tensor(True, dtype=torch.bool)}, } connector.get.return_value = (payload, 16) - adapter._poll_single_request(request) + adapter._poll_single_request(_dequeue_load_entry(adapter, request)) + assert request.additional_information is None + assert "req-1" in adapter._finished_load_reqs + request.status = RequestStatus.WAITING_FOR_CHUNK + adapter.process_pending_chunks( + DummyWaitingQueue(), + [request], + scheduler_requests={request.request_id: request}, + ) assert request.additional_information == payload assert adapter.get_req_chunk["req-1"] == 1 - assert "req-1" in adapter._finished_load_reqs + assert "req-1" not in adapter._finished_load_reqs assert "req-1" in adapter.upstream_exhausted_requests assert "req-1" not in adapter._pending_load_reqs +def test_load_async_does_not_requeue_registered_inflight_request(build_adapter): + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-inflight", RequestStatus.RUNNING, external_req_id="ext-inflight") + + entry = _dequeue_load_entry(adapter, request) + + # The recv loop owns the popped request until polling completes. A + # scheduler-side duplicate registration must not append a second copy. + adapter.load_async(request) + assert adapter._pending_load_reqs == deque() + + connector.get.return_value = ( + { + "ids": {"prompt": [1]}, + "meta": {"finished": True}, + }, + 1, + ) + assert adapter._poll_single_request(entry) is True + + # Successful polling releases the registration for the next chunk. + adapter.load_async(request) + assert [queued.request for queued in adapter._pending_load_reqs] == [request] + + +def test_cleanup_then_reregister_drops_stale_queued_entry(build_adapter): + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-resume", RequestStatus.WAITING, external_req_id="ext-resume") + + adapter.load_async(request) + stale_entry = adapter._pending_load_reqs[0] + adapter.cleanup_receiver(request.request_id) + adapter.load_async(request) + fresh_entry = adapter._pending_load_reqs[-1] + assert len(adapter._pending_load_reqs) == 2 + assert adapter._pending_load_reqs.popleft() is stale_entry + assert adapter._pending_load_reqs.popleft() is fresh_entry + + connector.get.return_value = ( + { + "ids": {"prompt": [1]}, + "meta": {"finished": True}, + }, + 1, + ) + assert adapter._poll_single_request(stale_entry) is True + connector.get.assert_not_called() + assert adapter._registered_load_entries[request.request_id] is fresh_entry + + assert adapter._poll_single_request(fresh_entry) is True + connector.get.assert_called_once() + assert request.request_id not in adapter._registered_load_entries + assert adapter.get_req_chunk[request.request_id] == 1 + + def test_load_poll_ar_requeues_explicitly_replaced_running_prompt(build_adapter): adapter, connector = build_adapter(stage_id=1, model_mode="ar") request = _req("req-replace", RequestStatus.RUNNING, external_req_id="external-replace") @@ -234,15 +542,18 @@ def test_load_poll_ar_requeues_explicitly_replaced_running_prompt(build_adapter) 1, ) - assert adapter._poll_single_request(request) is True - assert request.num_computed_tokens == 0 - assert request.prompt_token_ids == [0] * 10 - assert request.request_id in adapter.replaced_streaming_prompt_ids - + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True request.status = RequestStatus.WAITING_FOR_CHUNK running_queue = [request] waiting_queue = DummyWaitingQueue() - adapter.process_pending_chunks(waiting_queue, running_queue) + adapter.process_pending_chunks( + waiting_queue, + running_queue, + scheduler_requests={request.request_id: request}, + ) + assert request.num_computed_tokens == 0 + assert request.prompt_token_ids == [0] * 10 + assert request.request_id in adapter.replaced_streaming_prompt_ids assert running_queue == [] assert waiting_queue == [request] assert request.status == RequestStatus.WAITING @@ -258,6 +569,134 @@ def test_load_poll_ar_requeues_explicitly_replaced_running_prompt(build_adapter) assert request.external_req_id not in adapter.requests_num_chunks_sent +@pytest.mark.parametrize("flat_tts_config", [False, True]) +def test_load_poll_ar_recomputes_native_duplex_prompt_each_condition(build_adapter, flat_tts_config): + adapter, connector = build_adapter( + stage_id=1, + model_mode="ar", + max_model_len=8192, + tts_max_model_len=4096, + tts_attention_type="sliding_recompute", + flat_tts_config=flat_tts_config, + ) + request = _req("req-rollover", RequestStatus.RUNNING, external_req_id="external-rollover") + request.resumable = True + previous_codes = list(range(25)) + request.prompt_token_ids = [0] * 10 + request._all_token_ids = [0] * 10 + previous_codes + request._output_token_ids = previous_codes.copy() + request.num_prompt_tokens = 10 + request.num_computed_tokens = 35 + request.num_output_placeholders = 0 + request.update_block_hashes = Mock() + adapter.get_req_chunk[request.request_id] = 1 + adapter.request_ids_mapping[request.request_id] = request.external_req_id + adapter._streaming_condition_lengths[request.request_id] = 10 + adapter._streaming_condition_seqs[request.request_id] = 0 + connector.get.return_value = ( + { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + "finished": True, + }, + }, + 1, + ) + + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True + request.status = RequestStatus.WAITING_FOR_CHUNK + running_queue = [request] + waiting_queue = DummyWaitingQueue() + adapter.process_pending_chunks( + waiting_queue, + running_queue, + scheduler_requests={request.request_id: request}, + ) + + assert request.num_computed_tokens == 0 + assert request.prompt_token_ids == [0] * 45 + assert request.additional_information["ids"]["streaming_prompt_previous_codes"] == previous_codes + assert request.additional_information["meta"]["streaming_prompt_recompute"] is True + assert request.resumable is False + assert request.request_id in adapter.replaced_streaming_prompt_ids + assert adapter._max_model_len == 4096 + assert running_queue == [] + assert waiting_queue == [request] + + +def test_window_condition_sequence_ignores_control_only_chunks(build_adapter) -> None: + adapter, connector = build_adapter( + stage_id=1, + model_mode="ar", + tts_attention_type="sliding_recompute", + ) + request = _req("req-condition-seq", RequestStatus.RUNNING) + + def receive(payload): + connector.get.return_value = (payload, 1) + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True + adapter.requests_with_ready_chunks.add(request.request_id) + adapter._apply_pending_ar_prompt_updates({request.request_id: request}) + adapter.requests_with_ready_chunks.discard(request.request_id) + return request.additional_information + + first = receive( + { + "ids": {"prompt": [1]}, + "meta": {"next_stage_prompt_len": 10}, + } + ) + assert first["meta"]["streaming_condition_seq"] == 0 + assert adapter._streaming_condition_seqs[request.request_id] == 0 + + control = receive({"meta": {"is_segment_finished": True}}) + assert "streaming_condition_seq" not in control["meta"] + assert adapter._streaming_condition_seqs[request.request_id] == 0 + + second = receive( + { + "ids": {"prompt": [2]}, + "meta": {"next_stage_prompt_len": 12}, + } + ) + assert second["meta"]["streaming_condition_seq"] == 1 + assert adapter._streaming_condition_lengths[request.request_id] == 12 + + +def test_load_poll_ar_reports_invalid_capacity_metadata(build_adapter): + adapter, connector = build_adapter(stage_id=1, model_mode="ar", max_model_len=4096) + request = _req("req-invalid-capacity", RequestStatus.RUNNING) + request.resumable = True + adapter.get_req_chunk[request.request_id] = 1 + connector.get.return_value = ( + { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 4096, + "next_stage_generation_tokens": 26, + "finished": False, + }, + }, + 1, + ) + + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True + request.status = RequestStatus.WAITING_FOR_CHUNK + running_queue = [request] + adapter.process_pending_chunks( + DummyWaitingQueue(), + running_queue, + scheduler_requests={request.request_id: request}, + ) + assert adapter.collect_failed_receive_request_ids() == { + request.request_id: "fresh streaming prompt plus generation reserve exceeds max_model_len: " + "prompt=4096, reserve=26, limit=4096" + } + assert request.request_id not in adapter._finished_load_reqs + + def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter): adapter, connector = build_adapter(stage_id=1, model_mode="generation") request = _req("req-tensor", RequestStatus.WAITING, external_req_id="external-tensor") @@ -272,7 +711,7 @@ def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter) } connector.get.return_value = (payload, 16) - adapter._poll_single_request(request) + adapter._poll_single_request(_dequeue_load_entry(adapter, request)) assert request.prompt_token_ids == [0] assert request.num_computed_tokens == 0 @@ -302,12 +741,13 @@ def test_load_poll_generation_empty_nonterminal_chunk_keeps_polling(build_adapte } connector.get.side_effect = [(empty_payload, 16), (ready_payload, 16)] - assert adapter._poll_single_request(request) is False + entry = _dequeue_load_entry(adapter, request) + assert adapter._poll_single_request(entry) is False assert request.request_id not in adapter._finished_load_reqs assert request.request_id not in adapter.requests_with_ready_chunks assert adapter.get_req_chunk[request.request_id] == 1 - assert adapter._poll_single_request(request) is True + assert adapter._poll_single_request(entry) is True assert request.request_id in adapter._finished_load_reqs assert torch.equal(request.additional_information["codes"]["audio"], ready_payload["codes"]["audio"]) assert adapter.get_req_chunk[request.request_id] == 2 @@ -701,6 +1141,77 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon assert args[1] == "ext-finished" +def test_old_turn_terminal_send_preserves_new_turn_state_and_sender(build_adapter): + adapter, connector = build_adapter( + stage_id=1, + model_mode="ar", + max_model_len=4096, + tts_attention_type="sliding_recompute", + ) + request = _req("req-live", RequestStatus.RUNNING, external_req_id="ext-live") + request.resumable = True + request._all_token_ids = [] + request._output_token_ids = [] + request.update_block_hashes = Mock() + adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct( + meta=MetaStruct( + finished=torch.tensor(True, dtype=torch.bool), + codec_streaming=True, + ) + ) + + put_started = threading.Event() + release_put = threading.Event() + put_keys = [] + + def blocking_first_put(*_args, **kwargs): + put_keys.append(kwargs["put_key"]) + if len(put_keys) == 1: + put_started.set() + assert release_put.wait(timeout=5) + return True, 1, {} + + connector.put.side_effect = blocking_first_put + adapter.save_async(multimodal_output=None, request=request) + first_task = adapter._pending_save_reqs.popleft() + first_sender_token = first_task["sender_token"] + send_thread = threading.Thread(target=lambda: adapter._send_single_request(first_task)) + send_thread.start() + assert put_started.wait(timeout=5) + + entry = _dequeue_load_entry(adapter, request) + connector.get.return_value = ( + { + "ids": {"prompt": [1]}, + "meta": { + "replace_streaming_prompt": True, + "next_stage_prompt_len": 10, + "next_stage_generation_tokens": 26, + }, + }, + 1, + ) + assert adapter._poll_single_request(entry) is True + adapter.requests_with_ready_chunks.add(request.request_id) + adapter._apply_pending_ar_prompt_updates({request.request_id: request}) + assert adapter._streaming_condition_seqs[request.request_id] == 0 + + request.num_computed_tokens = 1 + adapter.save_async(multimodal_output=None, request=request) + second_task = adapter._pending_save_reqs.popleft() + assert second_task["sender_token"] is first_sender_token + + release_put.set() + send_thread.join(timeout=5) + assert not send_thread.is_alive() + assert adapter._streaming_condition_seqs[request.request_id] == 0 + assert adapter._sender_tokens[request.external_req_id] is first_sender_token + assert first_sender_token.cancelled is False + + adapter._send_single_request(second_task) + assert put_keys == ["ext-live_1_0", "ext-live_1_1"] + + def test_load_poll_non_ar_merges_into_existing_additional_information(build_adapter): adapter, connector = build_adapter(stage_id=2, model_mode="diffusion") request = _req("req-non-ar", RequestStatus.WAITING, external_req_id="ext-non-ar") @@ -720,7 +1231,7 @@ def test_load_poll_non_ar_merges_into_existing_additional_information(build_adap } connector.get.return_value = (payload, 8) - assert adapter._poll_single_request(request) is True + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True assert request.prompt_token_ids == [7, 8] assert request.num_computed_tokens == 0 @@ -761,7 +1272,7 @@ def test_load_poll_generation_segment_marker_replaces_previous_chunk(build_adapt 1, ) - assert adapter._poll_single_request(request) is True + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True assert request.prompt_token_ids == [7, 8] assert "audio" not in request.additional_information["codes"] @@ -790,7 +1301,7 @@ def test_load_poll_generation_empty_replacement_snapshot_is_ready(build_adapter) 1, ) - assert adapter._poll_single_request(request) is True + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True assert request.prompt_token_ids == [0] assert "codes" not in request.additional_information @@ -816,7 +1327,7 @@ def test_load_poll_generation_without_snapshot_marker_keeps_incremental_state(bu 1, ) - assert adapter._poll_single_request(request) is False + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is False assert torch.equal(request.additional_information["codes"]["audio"], torch.tensor([1, 2])) assert request.additional_information["meta"]["cache_epoch"] == 3 @@ -840,9 +1351,17 @@ def test_load_poll_ar_request_additional_information_concats_tensors(build_adapt } connector.get.return_value = (payload, 8) - adapter._poll_single_request(request) + adapter._poll_single_request(_dequeue_load_entry(adapter, request)) - # AR mode now forwards the latest payload directly. + # The recv thread records the latest payload; scheduler-side queue + # restoration publishes it onto the Request. + assert torch.equal(request.additional_information["hidden_states"]["output"], torch.tensor([[1.0]])) + request.status = RequestStatus.WAITING_FOR_CHUNK + adapter.process_pending_chunks( + DummyWaitingQueue(), + [request], + scheduler_requests={request.request_id: request}, + ) assert request.additional_information == payload assert request.additional_information["meta"]["finished"].item() is True @@ -860,7 +1379,7 @@ def test_non_ar_poll_reinitializes_prefill_stats_for_later_chunks(build_adapter) }, 8, ) - assert adapter._poll_single_request(request) is True + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True OmniGenerationScheduler._record_prefill_stats(request) first_chunk_stats = request.prefill_stats request.prefill_stats = None @@ -872,7 +1391,7 @@ def test_non_ar_poll_reinitializes_prefill_stats_for_later_chunks(build_adapter) }, 8, ) - assert adapter._poll_single_request(request) is True + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True assert isinstance(request.prefill_stats, PrefillStats) OmniGenerationScheduler._record_prefill_stats(request) second_chunk_stats = request.prefill_stats @@ -971,7 +1490,7 @@ def test_non_active_waiting_request_is_held_off_scheduler(build_adapter): assert waiting_queue == [] assert active.status == RequestStatus.WAITING_FOR_CHUNK assert non_active.status == RequestStatus.WAITING - assert list(adapter._pending_load_reqs) == [active] + assert [entry.request for entry in adapter._pending_load_reqs] == [active] assert list(adapter.waiting_for_chunk_waiting_requests) == [active, non_active] @@ -1122,8 +1641,19 @@ def _populate_adapter_state(adapter, req_id="req-1", ext_id="ext-1"): adapter.get_req_chunk[req_id] = 3 adapter.requests_with_ready_chunks.add(req_id) adapter.request_ids_mapping[req_id] = ext_id - adapter._pending_load_reqs.append(SimpleNamespace(request_id=req_id)) + entry = SimpleNamespace(request_id=req_id) + adapter._pending_load_reqs.append(entry) + adapter._registered_load_entries[req_id] = entry adapter._finished_load_reqs.add(req_id) + adapter._pending_ar_prompt_updates[req_id] = ( + SimpleNamespace(request_id=req_id), + {}, + False, + False, + False, + ) + adapter._streaming_condition_lengths[req_id] = 10 + adapter._streaming_condition_seqs[req_id] = 2 adapter.put_req_chunk[ext_id] = 5 adapter.request_payload[ext_id] = {"hidden": [1, 2]} @@ -1145,12 +1675,71 @@ def test_cleanup_clears_all_state(build_adapter): assert req_id not in adapter.request_ids_mapping assert req_id in adapter._cancelled_load_reqs assert req_id not in adapter._finished_load_reqs + assert req_id not in adapter._pending_ar_prompt_updates + assert req_id not in adapter._streaming_condition_lengths + assert req_id not in adapter._streaming_condition_seqs + assert req_id not in adapter._registered_load_entries assert ext_id not in adapter.put_req_chunk assert ext_id not in adapter.request_payload assert ext_id not in adapter.code_prompt_token_ids +def test_cleanup_and_reregister_discards_inflight_old_segment_chunk(build_adapter): + """A late old-segment chunk cannot overwrite a resumed registration.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-late", RequestStatus.RUNNING, external_req_id="ext-late") + request.resumable = True + entry = _dequeue_load_entry(adapter, request) + + get_started = threading.Event() + release_get = threading.Event() + + def blocking_get(*_args): + get_started.set() + assert release_get.wait(timeout=5) + return ( + { + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": 10, + }, + }, + 1, + ) + + connector.get.side_effect = blocking_get + poll_results = [] + poll_thread = threading.Thread(target=lambda: poll_results.append(adapter._poll_single_request(entry))) + poll_thread.start() + assert get_started.wait(timeout=5) + + adapter.cleanup_receiver(request.request_id) + fresh_entry = _dequeue_load_entry(adapter, request) + release_get.set() + poll_thread.join(timeout=5) + + assert not poll_thread.is_alive() + assert poll_results == [True] + assert request.additional_information is None + assert request.request_id not in adapter._pending_ar_prompt_updates + assert request.request_id not in adapter._streaming_condition_seqs + assert adapter.get_req_chunk[request.request_id] == 0 + assert adapter._registered_load_entries[request.request_id] is fresh_entry + + connector.get.side_effect = None + connector.get.return_value = ( + { + "ids": {"prompt": [2]}, + "meta": {"finished": True}, + }, + 1, + ) + assert adapter._poll_single_request(fresh_entry) is True + assert adapter.get_req_chunk[request.request_id] == 1 + assert request.request_id not in adapter._registered_load_entries + + def test_cleanup_infers_external_id(build_adapter): """When external_req_id is None, cleanup should look it up from the mapping.""" adapter, _ = build_adapter(stage_id=1) @@ -1299,7 +1888,7 @@ def test_cleanup_after_poll_flow(build_adapter): "meta": {"finished": torch.tensor(True, dtype=torch.bool)}, } connector.get.return_value = (payload, 8) - adapter._poll_single_request(request) + adapter._poll_single_request(_dequeue_load_entry(adapter, request)) assert "req-flow" in adapter.upstream_exhausted_requests assert adapter.get_req_chunk["req-flow"] == 1 diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 1c33aaa56..396d50fd1 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for the MiniCPM-o 4.5 pipeline registration. Covers: @@ -19,6 +19,8 @@ from pathlib import Path from types import SimpleNamespace import pytest +from transformers import PretrainedConfig +from vllm.config import ModelConfig from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import ( @@ -167,8 +169,11 @@ class TestDeployTopology: assert connector["extra"]["connector_get_max_wait"] == 300 expected_processor = "tts2code2wav_async_chunk" if deploy.async_chunk else "tts2code2wav_full_payload" assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(expected_processor) + if filename != "minicpmo_4_5.yaml": + assert "hf_overrides" not in stages[1].yaml_engine_args if filename == "minicpmo_4_5.yaml": assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 4, 4] + assert stages[1].yaml_engine_args["hf_overrides"] == {"tts_config": {"attention_type": "sliding_recompute"}} memory_utilizations = [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] assert memory_utilizations == [ 0.55, @@ -189,6 +194,18 @@ class TestDeployTopology: 0.35, ] + def test_duplex_talker_hf_override_resolves_nested_attention_policy(self) -> None: + deploy = load_deploy_config(_DEPLOY_DIR / "minicpmo_4_5.yaml") + stages = merge_pipeline_deploy(OMNI_PIPELINES[deploy.pipeline], deploy) + overrides = stages[1].yaml_engine_args["hf_overrides"] + hf_config = PretrainedConfig() + hf_config.tts_config = PretrainedConfig(attention_type="full_attention") + + model_config = object.__new__(ModelConfig) + model_config._apply_dict_overrides(hf_config, overrides) + + assert hf_config.tts_config.attention_type == "sliding_recompute" + @pytest.mark.parametrize( "filename", [ diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py index 5a3e7d67d..95d845686 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py @@ -102,6 +102,7 @@ def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration: talker._pending_force_eos_rows = None talker._penalty_histories = None talker._request_audio_states = {} + talker._request_condition_states = {} talker._deferred_cleanup_ids = set() talker._tts_config = ConditionalChatTTSConfig() talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)]) @@ -833,11 +834,164 @@ def test_native_duplex_condition_matches_official_text_plus_audio_bos() -> None: assert condition.shape[0] == token_ids.shape[0] + 1 +def test_native_duplex_rollover_matches_official_sliding_recompute(mocker) -> None: + talker = _make_talker() + talker.emb_text = nn.Embedding(1, 2) + talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)]) + with torch.no_grad(): + talker.emb_code[0].weight.copy_(torch.arange(16, dtype=torch.float32).reshape(8, 2)) + + previous_condition = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + current_condition = torch.tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]) + build_condition = mocker.patch.object( + talker, + "_build_condition_embeddings", + side_effect=[previous_condition, current_condition, current_condition], + ) + + talker.preprocess( + torch.zeros(2, dtype=torch.long), + None, + _omni_is_prefill=True, + _omni_prompt_len=2, + request_id="req-window", + native_duplex=True, + meta={"turn_start": True, "streaming_condition_seq": 7}, + tts_token_ids=torch.tensor([1]), + tts_hidden_states=torch.ones(1, 2), + ) + + older_codes = [3, 4, 5] + talker._request_condition_states["req-window"]["base_recent_codes"] = tuple(older_codes) + previous_codes = [1, 2] + recompute_meta = { + "streaming_condition_seq": 8, + "streaming_prompt_recompute": True, + } + expected = torch.cat( + [ + previous_condition, + talker.emb_code[0](torch.tensor(previous_codes)), + current_condition, + ], + dim=0, + ) + _, full_embeds, _ = talker.preprocess( + torch.zeros(expected.shape[0], dtype=torch.long), + None, + _omni_is_prefill=True, + _omni_prompt_len=expected.shape[0], + request_id="req-window", + native_duplex=True, + ids={"streaming_prompt_previous_codes": previous_codes}, + meta=recompute_meta, + tts_token_ids=torch.tensor([2]), + tts_hidden_states=torch.ones(1, 2), + ) + + assert torch.equal(full_embeds, expected) + assert talker._request_audio_states["req-window"]["recent_codes"] == older_codes + previous_codes + + # Chunked-prefill/preemption retries reuse the immutable recipe for the + # same condition instead of advancing the one-chunk window twice. + _, retry_slice, _ = talker.preprocess( + torch.zeros(3, dtype=torch.long), + None, + _omni_is_prefill=True, + _omni_num_computed_tokens=2, + _omni_prompt_len=expected.shape[0], + request_id="req-window", + native_duplex=True, + ids={"streaming_prompt_previous_codes": previous_codes}, + meta=recompute_meta, + tts_token_ids=torch.tensor([2]), + tts_hidden_states=torch.ones(1, 2), + ) + assert torch.equal(retry_slice, expected[2:5]) + assert talker._request_audio_states["req-window"]["recent_codes"] == older_codes + previous_codes + assert build_condition.call_count == 3 + + +def test_native_duplex_condition_advance_requires_recompute_marker(mocker) -> None: + talker = _make_talker() + condition = torch.ones(2, 2) + mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition) + common = { + "_omni_is_prefill": True, + "_omni_prompt_len": 2, + "request_id": "req-history", + "native_duplex": True, + "tts_token_ids": torch.tensor([1]), + "tts_hidden_states": torch.ones(1, 2), + } + + talker.preprocess(torch.zeros(2, dtype=torch.long), None, meta={"streaming_condition_seq": 0}, **common) + with pytest.raises(ValueError, match="advanced without its streaming recompute marker"): + talker.preprocess(torch.zeros(2, dtype=torch.long), None, meta={"streaming_condition_seq": 1}, **common) + + +@pytest.mark.parametrize( + ("state", "previous_codes", "error"), + [ + (None, [1], "missing the previous Talker condition"), + ( + {"condition_seq": 0, "condition": torch.ones(2, 2), "base_recent_codes": ()}, + [7], + "invalid or terminal token", + ), + ], +) +def test_native_duplex_rollover_rejects_unverifiable_history(state, previous_codes, error) -> None: + talker = _make_talker() + talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)]) + if state is not None: + talker._request_condition_states["req-invalid-window"] = state + + with pytest.raises(ValueError, match=error): + talker._build_streaming_recompute_embeddings( + torch.ones(3, 2), + request_id="req-invalid-window", + info_dict={"ids": {"streaming_prompt_previous_codes": previous_codes}}, + meta={ + "streaming_condition_seq": 1, + "streaming_prompt_recompute": True, + }, + ) + + +def test_native_duplex_turn_start_resets_rollover_history() -> None: + talker = _make_talker() + talker._request_condition_states["req-turn"] = { + "condition_seq": 4, + "condition": torch.full((2, 2), 9.0), + "active_embeddings": torch.full((5, 2), 9.0), + } + current = torch.ones(3, 2) + + result = talker._build_streaming_recompute_embeddings( + current, + request_id="req-turn", + info_dict={}, + meta={"turn_start": True, "streaming_condition_seq": 5}, + ) + + assert torch.equal(result, current) + state = talker._request_condition_states["req-turn"] + assert state["condition_seq"] == 5 + assert torch.equal(state["condition"], current) + assert "active_embeddings" not in state + + def test_request_cleanup_evicts_decode_state() -> None: talker = _make_talker() talker._request_audio_states["req-done"] = {"finished": False} + talker._request_condition_states["req-done"] = { + "condition_seq": 1, + "condition": torch.ones(2, 2), + } talker.on_requests_finished(["req-done"]) talker._flush_deferred_cleanup() assert "req-done" not in talker._request_audio_states + assert "req-done" not in talker._request_condition_states diff --git a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py index 652215a08..795047f1c 100644 --- a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py +++ b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from types import SimpleNamespace import pytest @@ -131,6 +134,7 @@ def test_native_duplex_speak_segment_reaches_split_talker() -> None: assert converted["prompt_token_ids"] == [0, 0, 0] assert info["meta"]["replace_streaming_prompt"] is True assert info["meta"]["next_stage_prompt_len"] == 3 + assert info["meta"]["next_stage_generation_tokens"] == 26 assert info["meta"]["turn_start"] is True assert info["meta"]["segment_end"] is True assert info["duplex"]["epoch"] == 3 @@ -203,6 +207,7 @@ def test_native_duplex_continuation_appends_only_new_talker_condition() -> None: assert "replace_streaming_prompt" not in second_input["model_intermediate_buffer"]["meta"] assert third_input["model_intermediate_buffer"]["meta"]["replace_streaming_prompt"] is True assert second_input["model_intermediate_buffer"]["meta"]["next_stage_prompt_len"] == 3 + assert second_input["model_intermediate_buffer"]["meta"]["next_stage_generation_tokens"] == 26 assert second_input["prompt_token_ids"] == [0, 0, 0] diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index e7c533dcd..2832b54f6 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import math @@ -260,10 +263,30 @@ class OmniSchedulerMixin: self.running, scheduler_requests=self.requests, ) + # Prompt-window contracts are finalized on this scheduler thread + # and may surface a permanent receive failure. + self._process_chunk_receive_failures() self._reset_ready_async_chunk_replacements() self._process_pending_chunk_timeouts() self._log_failed_chunk_sends() + def _process_chunk_receive_failures(self) -> None: + """Fail requests whose consumed async chunk violated its contract.""" + adapter = getattr(self, "chunk_transfer_adapter", None) + collector = getattr(adapter, "collect_failed_receive_request_ids", None) + if collector is None: + return + failures = collector() + present_ids = {request_id for request_id in failures if request_id in self.requests} + if not present_ids: + return + logger.error( + "Marking %d request(s) as FINISHED_ERROR after invalid connector input: %s", + len(present_ids), + {request_id: failures[request_id] for request_id in sorted(present_ids)}, + ) + self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR) + def _restore_omni_wait_queues(self) -> None: """Restore requests temporarily parked by Omni input gates.""" if self.chunk_transfer_adapter: diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index 685558a00..2ce1eed70 100644 --- a/vllm_omni/data_entry_keys.py +++ b/vllm_omni/data_entry_keys.py @@ -59,6 +59,7 @@ class Ids(TypedDict, total=False): output: list[int] speech_token: list[int] prior_image: list[int] + streaming_prompt_previous_codes: list[int] class OmniPayloadMeta(TypedDict, total=False): @@ -78,7 +79,10 @@ class OmniPayloadMeta(TypedDict, total=False): override_keys: list[tuple[str, str]] num_processed_tokens: int next_stage_prompt_len: int + next_stage_generation_tokens: int replace_streaming_prompt: bool + streaming_prompt_recompute: bool + streaming_condition_seq: int replace_runtime_additional_information: bool ar_width: int eol_token_id: int @@ -167,6 +171,7 @@ class IdsStruct(_StructBase): output: list[int] | None = None speech_token: list[int] | None = None prior_image: list[int] | None = None + streaming_prompt_previous_codes: list[int] | None = None class MetaStruct(_StructBase): @@ -183,7 +188,10 @@ class MetaStruct(_StructBase): override_keys: list[tuple[str, str]] | None = None num_processed_tokens: int | None = None next_stage_prompt_len: int | None = None + next_stage_generation_tokens: int | None = None replace_streaming_prompt: bool | None = None + streaming_prompt_recompute: bool | None = None + streaming_condition_seq: int | None = None replace_runtime_additional_information: bool | None = None ar_width: int | None = None eol_token_id: int | None = None diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 6903fcae0..023101223 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -64,6 +64,13 @@ stages: - stage_id: 1 max_num_seqs: 4 + # Native duplex uses MiniCPMTTS.generate_with_buffer's supported + # sliding_recompute policy: retain one completed condition/codec chunk + # when the next condition arrives. Other deploy profiles continue to use + # the checkpoint's full_attention policy. + hf_overrides: + tts_config: + attention_type: sliding_recompute # Leave device-level headroom for the colocated HiFi-GAN cuDNN workspace. gpu_memory_utilization: 0.15 enforce_eager: false diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index 3870155fe..d585b3ba1 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import time from collections.abc import Callable from typing import Any +from vllm.v1.request import Request + from vllm_omni.metrics import OrchestratorAggregator from .utils.logging import get_connector_logger @@ -216,34 +218,154 @@ def compute_talker_prompt_ids_length(prompt_ids: list[int]) -> int: return sum_user_len + assistant_len -def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], request: Any) -> bool: +def construct_next_stage_streaming_input_prompt( + payload_data: dict[str, Any], + request: Request, + *, + max_model_len: int | None = None, + previous_condition_len: int | None = None, + previous_condition_seq: int | None = None, + condition_seq: int | None = None, + recompute_previous_chunks: int = 0, +) -> bool: """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`` while extending the current prompt. - Explicit replacements reset it with the prompt. + * Preserves ``num_computed_tokens`` while extending a non-window prompt. + Explicit replacements and one-condition window recomputes reset it. * 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", {}) - meta = payload_data.get("meta", {}) - next_stage_prompt_len = meta.get("next_stage_prompt_len") if isinstance(meta, dict) else None - replace_prompt = isinstance(meta, dict) and meta.get("replace_streaming_prompt") is True - if replace_prompt and isinstance(next_stage_prompt_len, int) and next_stage_prompt_len > 0: + ids = payload_data.get("ids") + if not isinstance(ids, dict): + ids = {} + payload_data["ids"] = ids + meta = payload_data.get("meta") + if not isinstance(meta, dict): + meta = {} + payload_data["meta"] = meta + # This flag is transported through a merge-based runtime buffer. Set it on + # every condition so a prior rollover cannot be replayed by a later append. + meta["streaming_prompt_recompute"] = False + next_stage_prompt_len = meta.get("next_stage_prompt_len") + if "next_stage_generation_tokens" in meta: + next_generation_tokens = meta["next_stage_generation_tokens"] + if ( + not isinstance(next_generation_tokens, int) + or isinstance(next_generation_tokens, bool) + or next_generation_tokens <= 0 + ): + raise ValueError("next_stage_generation_tokens must be a positive integer") + generation_reserve = next_generation_tokens + else: + generation_reserve = 0 + capacity_limit = ( + max_model_len + if isinstance(max_model_len, int) and not isinstance(max_model_len, bool) and max_model_len > 0 + else None + ) + managed_prompt_len: int | None = None + if generation_reserve > 0 and capacity_limit is not None: + if ( + not isinstance(next_stage_prompt_len, int) + or isinstance(next_stage_prompt_len, bool) + or next_stage_prompt_len <= 0 + ): + raise ValueError("capacity-managed streaming prompt requires a positive next_stage_prompt_len") + managed_prompt_len = next_stage_prompt_len + if managed_prompt_len + generation_reserve > capacity_limit: + raise ValueError( + "fresh streaming prompt plus generation reserve exceeds max_model_len: " + f"prompt={managed_prompt_len}, reserve={generation_reserve}, limit={capacity_limit}" + ) + explicit_replacement = meta.get("replace_streaming_prompt") is True + if isinstance(recompute_previous_chunks, bool) or recompute_previous_chunks not in (0, 1): + raise ValueError("streaming prompt recompute supports exactly one previous chunk") + has_window_contract = recompute_previous_chunks == 1 + window_recompute = not explicit_replacement and has_window_contract and previous_condition_seq is not None + exceeds_accumulated_capacity = ( + managed_prompt_len is not None + and capacity_limit is not None + and request.num_computed_tokens + managed_prompt_len + generation_reserve > capacity_limit + ) + if not explicit_replacement and exceeds_accumulated_capacity and not window_recompute: + raise ValueError("capacity-managed streaming prompt rollover requires window_size=1") + replacement_prompt_len = next_stage_prompt_len + if window_recompute: + if managed_prompt_len is None or capacity_limit is None: + raise ValueError("streaming prompt window requires a positive generation reserve and max_model_len") + if ( + not isinstance(previous_condition_len, int) + or isinstance(previous_condition_len, bool) + or previous_condition_len <= 0 + or not isinstance(previous_condition_seq, int) + or isinstance(previous_condition_seq, bool) + or previous_condition_seq < 0 + or not isinstance(condition_seq, int) + or isinstance(condition_seq, bool) + or condition_seq < 0 + ): + raise ValueError("streaming prompt recompute requires condition lengths and monotonic sequence metadata") + if condition_seq != previous_condition_seq + 1: + raise ValueError( + "streaming prompt recompute skipped a Talker condition: " + f"previous={previous_condition_seq}, current={condition_seq}" + ) + + num_placeholders = int(getattr(request, "num_output_placeholders", 0) or 0) + confirmed_num_computed_tokens = max(0, int(request.num_computed_tokens) - num_placeholders) + if confirmed_num_computed_tokens < request.num_prompt_tokens: + raise ValueError( + "confirmed streaming output ends before the current condition: " + f"confirmed={confirmed_num_computed_tokens}, prompt={request.num_prompt_tokens}" + ) + previous_codec_ids = list(request._all_token_ids[request.num_prompt_tokens : confirmed_num_computed_tokens]) + max_confirmed_codec_tokens = generation_reserve - 1 + if len(previous_codec_ids) > max_confirmed_codec_tokens: + raise ValueError( + "streaming prompt recompute retained too many codec tokens: " + f"retained={len(previous_codec_ids)}, limit={max_confirmed_codec_tokens}" + ) + + replacement_prompt_len = previous_condition_len + len(previous_codec_ids) + managed_prompt_len + if replacement_prompt_len + generation_reserve > capacity_limit: + raise ValueError( + "sliding streaming prompt plus generation reserve exceeds max_model_len: " + f"previous={previous_condition_len}, codec={len(previous_codec_ids)}, " + f"current={managed_prompt_len}, reserve={generation_reserve}, limit={capacity_limit}" + ) + ids["streaming_prompt_previous_codes"] = previous_codec_ids + meta.update( + { + "streaming_prompt_recompute": True, + "streaming_condition_seq": condition_seq, + } + ) + logger.debug( + "Recomputing a one-chunk streaming prompt window: previous=%d, codec=%d, current=%d, reserve=%d, limit=%d", + previous_condition_len, + len(previous_codec_ids), + managed_prompt_len, + generation_reserve, + capacity_limit, + ) + + replace_prompt = explicit_replacement or window_recompute + if replace_prompt and isinstance(replacement_prompt_len, int) and replacement_prompt_len > 0: # Some downstream stages consume complete, independently conditioned # segments instead of extending an existing KV prefix. The producer # declares that transport behavior explicitly in payload metadata. - new_prompt = [0] * next_stage_prompt_len + new_prompt = [0] * replacement_prompt_len request._output_token_ids.clear() request._all_token_ids.clear() request._all_token_ids.extend(new_prompt) request.prompt_token_ids = new_prompt request.num_computed_tokens = 0 - request.num_prompt_tokens = next_stage_prompt_len + request.num_prompt_tokens = replacement_prompt_len request.update_block_hashes() return True prompt_token_ids = ids.get("prompt", None) diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py index 578fbb233..3420dcb94 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/base.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import threading from collections import deque @@ -37,6 +37,8 @@ class OmniTransferAdapterBase: self._save_cond = threading.Condition() self._send_failures: dict[str, str] = {} self._send_failure_lock = threading.Lock() + self._receive_failures: dict[str, str] = {} + self._receive_failure_lock = threading.Lock() self.recv_thread = threading.Thread(target=self.recv_loop, daemon=True) self.recv_thread.start() @@ -102,6 +104,17 @@ class OmniTransferAdapterBase: failures, self._send_failures = self._send_failures, {} return failures + def record_receive_failure(self, request_id: str, reason: str) -> None: + """Record an invalid received chunk that cannot be retried.""" + with self._receive_failure_lock: + self._receive_failures.setdefault(request_id, reason) + + def collect_failed_receive_request_ids(self) -> dict[str, str]: + """Drain and return requests whose received chunk was invalid.""" + with self._receive_failure_lock: + failures, self._receive_failures = self._receive_failures, {} + return failures + def save_loop(self): """Loop to send outgoing data.""" while not self.stop_event.is_set(): 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 c57b0c7b2..0979956e1 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 @@ -34,6 +34,50 @@ class _SenderGeneration: self.cancelled = False +class _LoadEntry: + """Identify one receiver registration across queue and I/O boundaries.""" + + __slots__ = ("request",) + + def __init__(self, request: Request) -> None: + self.request = request + + @property + def request_id(self) -> str: + return self.request.request_id + + @property + def external_req_id(self) -> str: + return self.request.external_req_id + + +def _resolve_talker_streaming_prompt_config(model_config: Any) -> tuple[int, int]: + """Return the effective Talker context limit and recompute window.""" + max_model_len = int(getattr(model_config, "max_model_len", 0) or 0) + hf_config = getattr(model_config, "hf_config", None) + tts_config = getattr(hf_config, "tts_config", None) + if tts_config is None and getattr(hf_config, "model_type", None) == "minicpmtts": + tts_config = hf_config + tts_max_model_len = ( + tts_config.get("max_position_embeddings", 0) + if isinstance(tts_config, Mapping) + else getattr(tts_config, "max_position_embeddings", 0) + ) + tts_max_model_len = int(tts_max_model_len or 0) + if tts_max_model_len > 0: + max_model_len = min(max_model_len, tts_max_model_len) if max_model_len else tts_max_model_len + attention_type = ( + tts_config.get("attention_type") + if isinstance(tts_config, Mapping) + else getattr(tts_config, "attention_type", None) + ) + # Official MiniCPMTTS supports both full_attention and + # sliding_recompute. Only an explicit Talker-stage policy enables the + # latter; native duplex alone must not override checkpoint semantics. + previous_chunks = 1 if attention_type == "sliding_recompute" else 0 + return max_model_len, previous_chunks + + class OmniChunkTransferAdapter(OmniTransferAdapterBase): """Chunk-level transfer adapter for Omni connector pipelines. @@ -59,8 +103,19 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # lock only protects this map and short state transitions. Slow payload # construction and connector writes never hold it. self._sender_state_lock = threading.Lock() + # Serialize receiver registration/cleanup with the commit of a chunk + # returned by connector.get(). The get itself stays outside the lock so + # scheduler cleanup never waits for connector I/O. + self._receiver_state_lock = threading.Lock() + # Includes requests queued for polling and the request currently being + # polled. Per-registration identity lets the receiver reject an old + # segment's queued or in-flight work after the same request id resumes. + self._registered_load_entries: dict[str, _LoadEntry] = {} self._sender_tokens: dict[str, _SenderGeneration] = {} self.scheduler_max_num_seqs = vllm_config.scheduler_config.max_num_seqs + self._max_model_len, self._streaming_prompt_previous_chunks = _resolve_talker_streaming_prompt_config( + model_config + ) active_stream_window = int(getattr(model_config, "active_stream_window", 0) or 0) model_max_num_seqs = int(getattr(model_config, "max_num_seqs", self.scheduler_max_num_seqs) or 0) if model_max_num_seqs <= 0: @@ -116,6 +171,15 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self._held_non_active: deque[Any] = deque() self.requests_num_chunks_sent: dict[str, int] = defaultdict(int) self._pending_streaming_prefills: dict[str, dict] = {} + # The recv thread only records AR payloads. Prompt mutation is finalized + # by the scheduler after it observes the ready marker, where the request + # token counters form a consistent snapshot. + self._pending_ar_prompt_updates: dict[ + str, + tuple[Request, dict[str, Any], bool, int | None, bool], + ] = {} + self._streaming_condition_lengths: dict[str, int] = {} + self._streaming_condition_seqs: dict[str, int] = {} # Monotonic timestamp of when each request last began waiting for a # chunk, refreshed every time one arrives. Read by # collect_timed_out_request_ids() so a stream that stops advancing @@ -211,8 +275,14 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): return if not hasattr(request, "additional_information"): request.additional_information = None - self._cancelled_load_reqs.discard(request.request_id) - self._pending_load_reqs.append(request) + with self._receiver_state_lock: + request_id = request.request_id + if request_id in self._registered_load_entries: + return + entry = _LoadEntry(request) + self._cancelled_load_reqs.discard(request_id) + self._registered_load_entries[request_id] = entry + self._pending_load_reqs.append(entry) with self._recv_cond: self._recv_cond.notify() @@ -309,12 +379,16 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): with self._save_cond: self._save_cond.notify() - def _poll_single_request(self, request: Request): + def _poll_single_request(self, entry: _LoadEntry): + request = entry.request stage_id = self.connector.stage_id target_stage_id = stage_id - 1 req_id = request.request_id - chunk_id = self.get_req_chunk[req_id] - external_req_id = self.request_ids_mapping.get(req_id, req_id) + with self._receiver_state_lock: + if self._registered_load_entries.get(req_id) is not entry: + return True + chunk_id = self.get_req_chunk[req_id] + external_req_id = self.request_ids_mapping.get(req_id, req_id) connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}" # Use timeout=0 for non-blocking poll @@ -326,10 +400,46 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): ) except Exception as e: logger.error(f"SharedMemoryConnector get failed for req {connector_get_key}: {e}") + with self._receiver_state_lock: + if self._registered_load_entries.get(req_id) is not entry: + return True return False if result is None: + with self._receiver_state_lock: + if self._registered_load_entries.get(req_id) is not entry: + return True return False + + with self._receiver_state_lock: + # cleanup_receiver() can run while connector.get() is in flight. + # Treat cleanup as a commit barrier: a late chunk must not recreate + # prompt/window state for the removed request. + if self._registered_load_entries.get(req_id) is not entry: + return True + is_success = self._commit_received_chunk( + request, + result, + stage_id=stage_id, + req_id=req_id, + chunk_id=chunk_id, + connector_get_key=connector_get_key, + ) + if is_success: + self._registered_load_entries.pop(req_id, None) + return is_success + + def _commit_received_chunk( + self, + request: Request, + result: tuple[dict[str, Any], int], + *, + stage_id: int, + req_id: str, + chunk_id: int, + connector_get_key: str, + ) -> bool: + """Commit a received connector chunk while receiver state is locked.""" payload_data, size = result if payload_data: @@ -337,16 +447,38 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.get_req_chunk[req_id] += 1 meta = payload_data.get("meta", {}) + if not isinstance(meta, dict): + meta = {} + payload_data["meta"] = 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 + was_resumable = bool(getattr(request, "resumable", False)) + meta["streaming_prompt_recompute"] = False + prompt_len = meta.get("next_stage_prompt_len") replace_prompt = meta.get("replace_streaming_prompt") is True - if getattr(request, "resumable", False) and (chunk_id > 0 or replace_prompt): - # For new streaming input segment, we should update prompt from payload - replaced = construct_next_stage_streaming_input_prompt(payload_data, request) - if replaced: - self.replaced_streaming_prompt_ids.add(req_id) + payload_ids = payload_data.get("ids") + has_prompt_payload = bool( + replace_prompt + or (isinstance(prompt_len, int) and not isinstance(prompt_len, bool) and prompt_len > 0) + or (isinstance(payload_ids, dict) and payload_ids.get("prompt")) + ) + window_condition_len = ( + prompt_len + if self._streaming_prompt_previous_chunks == 1 + and isinstance(prompt_len, int) + and not isinstance(prompt_len, bool) + and prompt_len > 0 + else None + ) + update_prompt = bool(was_resumable and has_prompt_payload and (chunk_id > 0 or replace_prompt)) + self._pending_ar_prompt_updates[req_id] = ( + request, + payload_data, + update_prompt, + window_condition_len, + was_resumable, + ) if payload_finished: self.upstream_exhausted_requests.add(req_id) @@ -559,9 +691,12 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): elif finished_flag is not None: is_payload_finished = bool(finished_flag) - # Reclaim per-request async state only after the terminal payload - # has been sent successfully. This avoids cleanup->save races. - if is_payload_finished: + # Processor-controlled codec ``finished`` can mark the end of one + # native-duplex turn while the persistent request remains live. + # Only the scheduler's whole-request terminal may reclaim shared + # sender/receiver state; otherwise a late old-turn put can erase a + # newly admitted turn for the same request. + if is_payload_finished and is_finished: self.cleanup(request.request_id, external_req_id) else: # R1.2 of #4855. connector.put returning False is a silent drop: no @@ -607,9 +742,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): def cleanup_receiver(self, request_id: str) -> None: """Reclaim receiver-side per-request state (keyed by internal id). - Safe to call from the scheduler even when ``save_async()`` has - enqueued work that the background thread has not yet processed, - because it only touches receiver-side dictionaries. + Safe to call from the scheduler while ``connector.get()`` is in + flight: cleanup and the post-get state commit are serialized, and a + cancelled request cannot publish the late chunk. Must also purge the request from the chunk-parking deques (``waiting_for_chunk_waiting_requests`` / ``_running_requests`` / @@ -623,6 +758,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): Idempotent: calling with an already-cleaned or unknown id is safe. """ + with self._receiver_state_lock: + self._clear_receiver_state_locked(request_id) + + def _clear_receiver_state_locked(self, request_id: str) -> None: + """Clear receiver state while ``_receiver_state_lock`` is held.""" self._active_streams.pop(request_id, None) self.upstream_exhausted_requests.discard(request_id) self.segment_finished_requests.discard(request_id) @@ -631,6 +771,10 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.replaced_streaming_prompt_ids.discard(request_id) self.request_ids_mapping.pop(request_id, None) self.requests_origin_status.pop(request_id, None) + self._pending_ar_prompt_updates.pop(request_id, None) + self._streaming_condition_lengths.pop(request_id, None) + self._streaming_condition_seqs.pop(request_id, None) + self._registered_load_entries.pop(request_id, None) self._discard_from_chunk_deque(self.waiting_for_chunk_waiting_requests, request_id) self._discard_from_chunk_deque(self.waiting_for_chunk_running_requests, request_id) self._discard_from_chunk_deque(self._held_non_active, request_id) @@ -757,6 +901,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): RequestStatus.RUNNING, self._finished_load_reqs, ) + self._apply_pending_ar_prompt_updates(scheduler_requests) self._requeue_replaced_prompts(waiting_queue, running_queue) while len(running_queue) > self.scheduler_max_num_seqs: request = running_queue.pop() @@ -772,10 +917,66 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self._process_chunk_queue( running_queue, self.waiting_for_chunk_running_requests, RequestStatus.RUNNING, self._finished_load_reqs ) + self._apply_pending_ar_prompt_updates(scheduler_requests) self._requeue_replaced_prompts(waiting_queue, running_queue) self._promote_active_streams(waiting_queue) self._preempt_non_active_running(waiting_queue, running_queue) + def _apply_pending_ar_prompt_updates(self, scheduler_requests: dict[str, Request] | None) -> None: + """Finalize ready AR prompt updates on the scheduler thread.""" + for request_id in tuple(self.requests_with_ready_chunks): + with self._receiver_state_lock: + pending = self._pending_ar_prompt_updates.pop(request_id, None) + if pending is None: + continue + ( + fallback_request, + payload_data, + update_prompt, + window_condition_len, + was_resumable, + ) = pending + request = scheduler_requests.get(request_id) if scheduler_requests is not None else fallback_request + if request is None: + continue + + request.additional_information = payload_data + meta = payload_data.get("meta") + meta = meta if isinstance(meta, dict) else {} + previous_condition_seq = self._streaming_condition_seqs.get(request_id) + condition_seq = previous_condition_seq + 1 if previous_condition_seq is not None else 0 + is_window_condition = window_condition_len is not None + if is_window_condition: + meta["streaming_condition_seq"] = condition_seq + if is_window_condition: + update_prompt = bool( + was_resumable + and (previous_condition_seq is not None or meta.get("replace_streaming_prompt") is True) + ) + try: + if update_prompt: + replaced = construct_next_stage_streaming_input_prompt( + payload_data, + request, + max_model_len=self._max_model_len, + previous_condition_len=self._streaming_condition_lengths.get(request_id), + previous_condition_seq=previous_condition_seq, + condition_seq=condition_seq, + recompute_previous_chunks=self._streaming_prompt_previous_chunks, + ) + if replaced: + self.replaced_streaming_prompt_ids.add(request_id) + except ValueError as exc: + # The connector chunk was consumed, so retrying would skip + # this transition and desynchronize Talker. Let the + # scheduler fail it. + self.record_receive_failure(request_id, str(exc)) + continue + + if window_condition_len is not None: + self._streaming_condition_lengths[request_id] = window_condition_len + self._streaming_condition_seqs[request_id] = condition_seq + def _requeue_replaced_prompts(self, waiting_queue: Any, running_queue: list[Request]) -> None: """Move a replaced running prompt back through scheduler admission.""" for request in list(running_queue): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py b/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py index fb6738a29..6e8677cf2 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py @@ -1,8 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # NOTE: Do not import model classes in this file. Importing any # submodule in this package triggers __init__.py execution, and # both the model registry and pipeline registry import submodules # directly — heavy imports here would be loaded as a side effect # even though nothing depends on these re-exports. + +# One MiniCPMTTS.generate_chunk emits 25 codec frames plus its terminator. +MINICPMO45_DUPLEX_CODEC_TOKENS_PER_CHUNK = 26 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 68c794478..369219317 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -28,6 +28,7 @@ from vllm.model_executor.models.utils import maybe_prefix from vllm.v1.sample.sampler import Sampler from vllm_omni.experimental.fullduplex.engine.intermediate import get_tts_handoff +from vllm_omni.model_executor.models.minicpmo_4_5 import MINICPMO45_DUPLEX_CODEC_TOKENS_PER_CHUNK from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform @@ -46,7 +47,7 @@ _OFFLINE_CODEC_MAX_NEW_TOKENS = 2048 # 25 codec frames (``codec_chunk_frames``) plus the terminating sample. # Without this, the single-vocab Sampler keeps the stage-1 request alive # until codec EOS / 4096 and Thinker never starts the next model turn. -_DUPLEX_CODEC_TOKENS_PER_CHUNK = 26 +_DUPLEX_CODEC_TOKENS_PER_CHUNK = MINICPMO45_DUPLEX_CODEC_TOKENS_PER_CHUNK def _native_duplex_chunk_budget(meta: Mapping[str, Any] | None) -> tuple[int, int]: @@ -158,6 +159,9 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): self._pending_force_eos_rows: list[bool] | None = None self._penalty_histories: list[torch.Tensor] | None = None self._request_audio_states: dict[str, dict[str, Any]] = {} + # Mirrors upstream TTSStreamingGenerator._chunk_info: one committed + # condition plus, during a rollover, one immutable recompute recipe. + self._request_condition_states: dict[str, dict[str, Any]] = {} self._deferred_cleanup_ids: set[str] = set() tts_config = getattr(config, "tts_config", None) @@ -262,6 +266,116 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): return torch.cat([condition, audio_bos], dim=0) return torch.cat([condition, self._boundary_embeddings()], dim=0) + def _build_streaming_recompute_embeddings( + self, + current_condition: torch.Tensor, + *, + request_id: str, + info_dict: Mapping[str, Any], + meta: Mapping[str, Any], + ) -> torch.Tensor: + """Return the official one-previous-chunk sliding-recompute window.""" + condition_seq = meta.get("streaming_condition_seq") + if not isinstance(condition_seq, int) or isinstance(condition_seq, bool): + if meta.get("streaming_prompt_recompute") is True: + raise ValueError("streaming prompt recompute is missing streaming_condition_seq") + # Direct model tests and non-connector callers do not participate in + # the persistent async-chunk lifecycle, so they need no window state. + return current_condition + + turn_start = bool(meta.get("turn_start")) + recompute = meta.get("streaming_prompt_recompute") is True + states = self._request_condition_states + state = states.get(request_id) + if turn_start: + if recompute: + raise ValueError("streaming prompt recompute cannot cross a native duplex turn boundary") + states[request_id] = { + "condition_seq": condition_seq, + "condition": current_condition.detach().clone(), + "base_recent_codes": (), + } + return current_condition + + if state is None: + if recompute: + raise ValueError("streaming prompt recompute is missing the previous Talker condition") + states[request_id] = { + "condition_seq": condition_seq, + "condition": current_condition.detach().clone(), + "base_recent_codes": (), + } + return current_condition + + previous_seq = state.get("condition_seq") + if not isinstance(previous_seq, int) or condition_seq < previous_seq: + raise ValueError( + f"stale native duplex Talker condition sequence: current={condition_seq}, previous={previous_seq}" + ) + if condition_seq > previous_seq + 1: + raise ValueError( + f"native duplex Talker skipped a condition sequence: current={condition_seq}, previous={previous_seq}" + ) + + if not recompute: + if condition_seq > previous_seq: + raise ValueError( + "a native duplex Talker condition advanced without its streaming recompute marker: " + f"current={condition_seq}, previous={previous_seq}" + ) + if "active_embeddings" in state: + raise ValueError("an active streaming recompute was replayed without its recompute marker") + return current_condition + + if condition_seq == previous_seq: + active_embeddings = state.get("active_embeddings") + if not isinstance(active_embeddings, torch.Tensor): + raise ValueError("streaming prompt window lost its cached recompute embeddings") + return active_embeddings + + if condition_seq != previous_seq + 1: + raise ValueError( + "streaming prompt recompute skipped a Talker condition: " + f"previous={previous_seq}, current={condition_seq}" + ) + + previous_condition = state.get("condition") + if not isinstance(previous_condition, torch.Tensor): + raise ValueError("streaming prompt recompute lost the previous Talker condition") + + ids = info_dict.get("ids") + previous_codes = ids.get("streaming_prompt_previous_codes") if isinstance(ids, Mapping) else None + if isinstance(previous_codes, torch.Tensor): + code_ids = previous_codes.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(-1) + elif isinstance(previous_codes, (list, tuple)): + code_ids = torch.as_tensor(previous_codes, device=self.emb_code[0].weight.device, dtype=torch.long) + else: + raise ValueError("streaming prompt recompute is missing confirmed codec ids") + if code_ids.numel() > _DUPLEX_CODEC_TOKENS_PER_CHUNK - 1: + raise ValueError(f"streaming prompt recompute has too many codec ids: {code_ids.numel()}") + if code_ids.numel() and bool(((code_ids < 0) | (code_ids >= self._codec_eos_id)).any()): + raise ValueError("streaming prompt recompute codec ids include an invalid or terminal token") + + parts = [previous_condition] + if code_ids.numel(): + parts.append(self.emb_code[0](code_ids)) + parts.append(current_condition) + full_embeddings = torch.cat(parts, dim=0) + previous_code_ids = tuple(int(code_id) for code_id in code_ids.tolist()) + previous_base_codes = state.get("base_recent_codes") + if not isinstance(previous_base_codes, tuple): + raise ValueError("streaming Talker condition lost its frozen codec history") + states[request_id] = { + "condition_seq": condition_seq, + "condition": current_condition.detach().clone(), + "active_embeddings": full_embeddings.detach().clone(), + # Official generate_with_buffer keeps all_generated_tokens across + # sliding recomputes, so the first sample in this chunk still sees + # the previous chunk's repetition-penalty window. + "base_recent_codes": (*previous_base_codes, *previous_code_ids)[-_CODEC_PENALTY_WINDOW:], + } + return full_embeddings + def preprocess( self, input_ids: torch.Tensor, @@ -274,6 +388,7 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): is_prefill = bool(info_dict.get("_omni_is_prefill", False)) state = info_dict.get("audio_state") first_call = not isinstance(state, dict) + request_id = str(info_dict.get("request_id", "0")) if is_prefill or first_call: token_ids, hidden_states = get_tts_handoff(info_dict) @@ -309,13 +424,35 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): hidden_states, native_duplex=native_duplex, ) + if native_duplex: + full_embeds = self._build_streaming_recompute_embeddings( + full_embeds, + request_id=request_id, + info_dict=info_dict, + meta=meta if isinstance(meta, Mapping) else {}, + ) + retained_codes: list[int] = [] + condition_seq = meta.get("streaming_condition_seq") if isinstance(meta, Mapping) else None + if native_duplex and isinstance(condition_seq, int) and not isinstance(condition_seq, bool): + condition_state = self._request_condition_states.get(request_id) + base_recent_codes = ( + condition_state.get("base_recent_codes") if isinstance(condition_state, dict) else None + ) + if not isinstance(base_recent_codes, tuple): + raise ValueError("streaming Talker condition lost its frozen codec history") + retained_codes = list(base_recent_codes) offset = int(info_dict.get("_omni_num_computed_tokens", 0)) - request_id = str(info_dict.get("request_id", "0")) # The handoff rebuilds only the tail-aligned Talker condition. # Materialize zero-token embeddings for any scheduler prompt # prefix so chunked prefill can slice from a non-zero offset. prompt_len = info_dict.get("_omni_prompt_len") target_len = int(prompt_len) if prompt_len is not None else offset + span_len + if native_duplex and isinstance(meta, Mapping) and meta.get("streaming_prompt_recompute") is True: + if target_len != full_embeds.shape[0]: + raise ValueError( + "streaming prompt recompute length mismatch: " + f"scheduler={target_len}, model={full_embeds.shape[0]}" + ) prefix_len = target_len - full_embeds.shape[0] if prefix_len > 0: placeholder_ids = torch.zeros( @@ -348,6 +485,8 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): "max_tokens": max_tokens, "min_tokens": min_tokens, } + if retained_codes: + state["recent_codes"] = retained_codes request_states = getattr(self, "_request_audio_states", None) if request_states is None: request_states = {} @@ -365,7 +504,6 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): }, ) - request_id = str(info_dict.get("request_id", "0")) stored = self._request_audio_states.get(request_id) if isinstance(stored, dict): state = stored @@ -533,8 +671,10 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): def _flush_deferred_cleanup(self) -> None: request_audio_states = getattr(self, "_request_audio_states", {}) + request_condition_states = getattr(self, "_request_condition_states", {}) for request_id in self._deferred_cleanup_ids: request_audio_states.pop(request_id, None) + request_condition_states.pop(request_id, None) self._deferred_cleanup_ids.clear() def _dummy_hidden_states( diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index 9f094dada..e303774ff 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniCPM-o 4.5 Thinker-to-Talker and Talker-to-Code2Wav bridges.""" import logging @@ -16,6 +16,7 @@ from vllm_omni.experimental.fullduplex.engine.intermediate import ( set_tts_handoff, ) from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.models.minicpmo_4_5 import MINICPMO45_DUPLEX_CODEC_TOKENS_PER_CHUNK logger = logging.getLogger(__name__) _MINICPMO45_ASYNC_STATE = "_minicpmo45_async_codec_state" @@ -989,6 +990,8 @@ def llm2tts( scheduler_prompt_token_ids = [0] * condition_length handoff_meta = model_intermediate_buffer.setdefault("meta", {}) handoff_meta["next_stage_prompt_len"] = condition_length + if is_native_duplex_handoff: + handoff_meta["next_stage_generation_tokens"] = MINICPMO45_DUPLEX_CODEC_TOKENS_PER_CHUNK # Native duplex resumes one Talker request within a turn, but a new # assistant turn must discard the previous turn's prompt and KV. if not is_native_duplex_handoff or native_turn_start: