diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index a1aa7351f..bca7d095c 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -19,6 +19,9 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler +from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( + OmniChunkTransferAdapter, +) # isort: on @@ -62,6 +65,43 @@ def _make_update(prompt_token_ids: list[int] | None = None) -> StreamingUpdate: ) +def _make_talker_adapter( + *, + max_model_len: int = 100, + recompute_on_capacity: bool = True, +) -> OmniChunkTransferAdapter: + adapter = OmniChunkTransferAdapter.__new__(OmniChunkTransferAdapter) + adapter.receives_chunks = False + adapter._max_model_len = max_model_len + adapter._streaming_prompt_previous_chunks = 1 + adapter._streaming_prompt_recompute_on_capacity = recompute_on_capacity + adapter._streaming_condition_lengths = {} + adapter._streaming_condition_seqs = {} + adapter.segment_finished_requests = set() + adapter.requests_num_chunks_sent = {} + return adapter + + +def _make_talker_update( + condition_len: int, + *, + reserve: int = 10, + condition_seq: int | None, +) -> StreamingUpdate: + update = _make_update([0] * condition_len) + update.model_intermediate_buffer = { + "native_duplex": True, + "ids": {"prompt": [1]}, + "meta": { + "next_stage_prompt_len": condition_len, + "next_stage_generation_tokens": reserve, + }, + } + if condition_seq is not None: + update.model_intermediate_buffer["meta"]["streaming_condition_seq"] = condition_seq + return update + + def _run_resumable_segment_stop( session: Request, *, @@ -538,6 +578,288 @@ def test_explicit_model_intermediate_prompt_replacement_releases_cache_and_water sched.encoder_cache_manager.free.assert_called_once_with(session) +def test_talker_capacity_exact_fit_extends_from_declared_length_without_ids_prompt() -> None: + sched = _make_scheduler(stage_id=1) + adapter = _make_talker_adapter(max_model_len=100) + sched.chunk_transfer_adapter = adapter + session = _make_request() + session.external_req_id = "external-capacity-exact-fit" + session.prompt_token_ids = [0] * 70 + session._all_token_ids.clear() + session._all_token_ids.extend(session.prompt_token_ids) + session._output_token_ids.clear() + session.append_output_token_ids([7, 8, 9, 10]) + session.num_prompt_tokens = 70 + session.num_computed_tokens = 74 + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + sched.num_waiting_for_streaming_input = 1 + adapter._streaming_condition_lengths[session.request_id] = 20 + adapter._streaming_condition_seqs[session.request_id] = 0 + update = _make_talker_update(16, condition_seq=1) + update.model_intermediate_buffer["ids"] = {"tts": [1, 2, 3]} + + sched._update_request_as_session(session, update) + + assert session.num_computed_tokens == 74 + assert session.num_prompt_tokens == 90 + assert list(session._all_token_ids[-20:]) == [7, 8, 9, 10] + [0] * 16 + assert update.model_intermediate_buffer["meta"] == { + "next_stage_prompt_len": 16, + "next_stage_generation_tokens": 10, + "streaming_condition_seq": 1, + "streaming_prompt_recompute": False, + } + assert adapter._streaming_condition_lengths[session.request_id] == 16 + assert adapter._streaming_condition_seqs[session.request_id] == 1 + assert session._omni_segment_generation == 1 + assert session.status == RequestStatus.WAITING + assert sched.num_waiting_for_streaming_input == 0 + sched._free_request_blocks.assert_not_called() + sched.encoder_cache_manager.free.assert_not_called() + + +def test_talker_first_update_seeds_sender_condition_tracking() -> None: + sched = _make_scheduler(stage_id=1) + adapter = _make_talker_adapter(max_model_len=100) + sched.chunk_transfer_adapter = adapter + session = _make_request() + session.external_req_id = "external-capacity-first-update" + session.prompt_token_ids = [0] * 20 + session._all_token_ids.clear() + session._all_token_ids.extend(session.prompt_token_ids) + session._output_token_ids.clear() + session.append_output_token_ids([7, 8, 9]) + session.num_prompt_tokens = 20 + session.num_computed_tokens = 23 + session.model_intermediate_buffer = { + "native_duplex": True, + "meta": { + "next_stage_prompt_len": 20, + "next_stage_generation_tokens": 10, + "streaming_condition_seq": 0, + }, + } + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + sched.num_waiting_for_streaming_input = 1 + update = _make_talker_update(16, condition_seq=1) + + sched._update_request_as_session(session, update) + + assert update.model_intermediate_buffer["meta"]["streaming_condition_seq"] == 1 + assert update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is False + assert adapter._streaming_condition_lengths[session.request_id] == 16 + assert adapter._streaming_condition_seqs[session.request_id] == 1 + assert session.num_computed_tokens == 23 + assert session.num_prompt_tokens == 39 + sched._free_request_blocks.assert_not_called() + sched.encoder_cache_manager.free.assert_not_called() + + +def test_talker_first_sliding_update_recomputes_from_seeded_condition() -> None: + sched = _make_scheduler(stage_id=1) + adapter = _make_talker_adapter(recompute_on_capacity=False) + sched.chunk_transfer_adapter = adapter + session = _make_request() + session.external_req_id = "external-sliding-first-update" + session.prompt_token_ids = [0] * 20 + session._all_token_ids.clear() + session._all_token_ids.extend(session.prompt_token_ids) + session._output_token_ids.clear() + previous_codes = [7, 8, 9] + session.append_output_token_ids(previous_codes) + session.num_prompt_tokens = 20 + session.num_computed_tokens = 23 + session.model_intermediate_buffer = { + "native_duplex": True, + "meta": { + "next_stage_prompt_len": 20, + "next_stage_generation_tokens": 10, + "streaming_condition_seq": 0, + }, + } + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + sched.num_waiting_for_streaming_input = 1 + update = _make_talker_update(16, condition_seq=1) + + sched._update_request_as_session(session, update) + + assert session.num_computed_tokens == 0 + assert session.num_prompt_tokens == 39 + assert update.model_intermediate_buffer["ids"]["streaming_prompt_previous_codes"] == previous_codes + assert update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is True + assert adapter._streaming_condition_lengths[session.request_id] == 16 + assert adapter._streaming_condition_seqs[session.request_id] == 1 + sched._free_request_blocks.assert_called_once_with(session) + sched.encoder_cache_manager.free.assert_called_once_with(session) + + +@pytest.mark.parametrize( + ("condition_seq", "error"), + [ + (None, "missing streaming_condition_seq"), + (0, "expected=1, received=0"), + (2, "expected=1, received=2"), + ], +) +def test_talker_invalid_condition_sequence_does_not_advance_tracking(mocker, condition_seq, error) -> None: + sched = _make_scheduler(stage_id=1) + adapter = _make_talker_adapter() + adapter.record_receive_failure = mocker.MagicMock() + adapter._streaming_condition_lengths["req-ar-streaming-test"] = 20 + adapter._streaming_condition_seqs["req-ar-streaming-test"] = 0 + sched.chunk_transfer_adapter = adapter + session = _make_request() + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + sched.num_waiting_for_streaming_input = 1 + update = _make_talker_update(16, condition_seq=condition_seq) + + sched._update_request_as_session(session, update) + + failure = adapter.record_receive_failure.call_args.args + assert failure[0] == session.request_id + assert error in failure[1] + assert adapter._streaming_condition_lengths[session.request_id] == 20 + assert adapter._streaming_condition_seqs[session.request_id] == 0 + assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ + + +def test_talker_first_update_that_cannot_fit_window_is_request_local(mocker) -> None: + sched = _make_scheduler(stage_id=1) + adapter = _make_talker_adapter() + adapter.record_receive_failure = mocker.MagicMock() + sched.chunk_transfer_adapter = adapter + session = _make_request() + session.prompt_token_ids = [0] * 20 + session._all_token_ids.clear() + session._all_token_ids.extend(session.prompt_token_ids) + session._output_token_ids.clear() + session.append_output_token_ids(list(range(9))) + session.num_prompt_tokens = 20 + session.num_computed_tokens = 29 + session.model_intermediate_buffer = { + "native_duplex": True, + "meta": { + "next_stage_prompt_len": 20, + "next_stage_generation_tokens": 10, + "streaming_condition_seq": 0, + }, + } + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + sched.num_waiting_for_streaming_input = 1 + update = _make_talker_update(62, condition_seq=1) + + sched._update_request_as_session(session, update) + + failure = adapter.record_receive_failure.call_args.args + assert failure[0] == session.request_id + assert "sliding streaming prompt plus generation reserve exceeds max_model_len" in failure[1] + assert adapter._streaming_condition_lengths == {} + assert adapter._streaming_condition_seqs == {} + assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ + + +def test_talker_capacity_overflow_recomputes_then_resumes_accumulation() -> None: + sched = _make_scheduler(stage_id=1) + adapter = _make_talker_adapter(max_model_len=100) + sched.chunk_transfer_adapter = adapter + session = _make_request() + session.external_req_id = "external-capacity-rollover" + session.prompt_token_ids = [0] * 70 + session._all_token_ids.clear() + session._all_token_ids.extend(session.prompt_token_ids) + session._output_token_ids.clear() + previous_codes = [7, 8, 9, 10] + session.append_output_token_ids(previous_codes) + session.num_prompt_tokens = 70 + session.num_computed_tokens = 74 + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + sched.num_waiting_for_streaming_input = 1 + adapter._streaming_condition_lengths[session.request_id] = 20 + adapter._streaming_condition_seqs[session.request_id] = 0 + adapter.segment_finished_requests.add(session.request_id) + adapter.requests_num_chunks_sent[session.external_req_id] = 74 + overflow_update = _make_talker_update(17, condition_seq=1) + + sched._update_request_as_session(session, overflow_update) + + assert session.num_computed_tokens == 0 + assert session.num_prompt_tokens == 41 + assert session.prompt_token_ids == [0] * 41 + assert overflow_update.model_intermediate_buffer["ids"]["streaming_prompt_previous_codes"] == previous_codes + assert overflow_update.model_intermediate_buffer["meta"]["streaming_condition_seq"] == 1 + assert overflow_update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is True + assert adapter.segment_finished_requests == set() + assert adapter.requests_num_chunks_sent == {} + assert session._omni_segment_generation == 1 + sched._free_request_blocks.assert_called_once_with(session) + sched.encoder_cache_manager.free.assert_called_once_with(session) + + # Once the replacement has been recomputed, later conditions append to + # the new full-attention prefix until capacity is approached again. + next_codes = [21, 22, 23, 24, 25, 26] + session.append_output_token_ids(next_codes) + session.num_computed_tokens = 47 + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + sched.num_waiting_for_streaming_input = 1 + append_update = _make_talker_update(20, condition_seq=2) + + sched._update_request_as_session(session, append_update) + + assert session.num_computed_tokens == 47 + assert session.num_prompt_tokens == 67 + assert list(session._all_token_ids[-26:]) == next_codes + [0] * 20 + assert append_update.model_intermediate_buffer["meta"]["streaming_condition_seq"] == 2 + assert append_update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is False + assert session._omni_segment_generation == 2 + assert sched._free_request_blocks.call_count == 1 + assert sched.encoder_cache_manager.free.call_count == 1 + + # A later overflow replaces the accumulated prefix again rather than + # reverting to recompute-on-every-condition behavior. + second_rollover_codes = [31, 32, 33, 34, 35, 36] + session.num_computed_tokens = session.num_prompt_tokens + session.append_output_token_ids(second_rollover_codes) + session.num_computed_tokens += len(second_rollover_codes) + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + sched.num_waiting_for_streaming_input = 1 + second_overflow_update = _make_talker_update(18, condition_seq=3) + + sched._update_request_as_session(session, second_overflow_update) + + assert session.num_computed_tokens == 0 + assert session.num_prompt_tokens == 44 + assert second_overflow_update.model_intermediate_buffer["ids"]["streaming_prompt_previous_codes"] == ( + second_rollover_codes + ) + assert second_overflow_update.model_intermediate_buffer["meta"]["streaming_condition_seq"] == 3 + assert second_overflow_update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is True + assert session._omni_segment_generation == 3 + assert sched._free_request_blocks.call_count == 2 + assert sched.encoder_cache_manager.free.call_count == 2 + + +def test_talker_invalid_capacity_update_is_reported_per_request(mocker) -> None: + sched = _make_scheduler(stage_id=1) + adapter = _make_talker_adapter(max_model_len=100) + adapter.record_receive_failure = mocker.MagicMock() + sched.chunk_transfer_adapter = adapter + session = _make_request() + session.status = RequestStatus.WAITING_FOR_STREAMING_REQ + sched.num_waiting_for_streaming_input = 1 + update = _make_talker_update(91, condition_seq=0) + + sched._update_request_as_session(session, update) + + adapter.record_receive_failure.assert_called_once_with( + session.request_id, + "fresh streaming prompt plus generation reserve exceeds max_model_len: prompt=91, reserve=10, limit=100", + ) + assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ + assert sched.num_waiting_for_streaming_input == 1 + sched._free_request_blocks.assert_not_called() + sched.encoder_cache_manager.free.assert_not_called() + + def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None: sched = _make_scheduler(stage_id=1) session = _make_request() diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 0e3194190..40d29e969 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -142,7 +142,7 @@ def test_streaming_prompt_rejects_invalid_generation_reserve( 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: +def test_capacity_managed_streaming_window_at_limit_still_appends(mocker: MockerFixture) -> None: request = _streaming_request(mocker, num_computed_tokens=4060) payload = { "ids": {"prompt": [1]}, @@ -152,13 +152,53 @@ def test_capacity_managed_streaming_prompt_at_limit_still_appends(mocker: Mocker }, } - replaced = construct_next_stage_streaming_input_prompt(payload, request, max_model_len=4096) + replaced = construct_next_stage_streaming_input_prompt( + payload, + request, + max_model_len=4096, + previous_condition_len=10, + previous_condition_seq=0, + condition_seq=1, + recompute_previous_chunks=1, + recompute_on_capacity=True, + ) assert replaced is False assert request.num_computed_tokens == 4060 assert request.num_prompt_tokens == 4070 +def test_capacity_managed_streaming_prompt_appends_from_declared_length_without_ids_prompt( + mocker: MockerFixture, +) -> None: + request = _streaming_request(mocker, num_computed_tokens=20) + payload = { + "ids": {"tts": [1, 2, 3]}, + "hidden_states": {"tts": [[0.1], [0.2], [0.3]]}, + "meta": { + "next_stage_prompt_len": 4, + "next_stage_generation_tokens": 26, + }, + } + + replaced = construct_next_stage_streaming_input_prompt( + payload, + request, + max_model_len=4096, + previous_condition_len=4, + previous_condition_seq=0, + condition_seq=1, + recompute_previous_chunks=1, + recompute_on_capacity=True, + ) + + assert replaced is False + assert request.num_computed_tokens == 20 + assert request.num_prompt_tokens == 24 + assert request.prompt_token_ids == [0] * 24 + request.update_block_hashes.assert_called_once_with() + + def test_streaming_window_builds_one_chunk_recompute_recipe(mocker: MockerFixture) -> None: previous_condition_len = 10 previous_codes = list(range(25)) @@ -203,13 +243,41 @@ def test_streaming_window_builds_one_chunk_recompute_recipe(mocker: MockerFixtur request.update_block_hashes.assert_called_once_with() -def test_streaming_window_recomputes_each_condition_without_accumulating(mocker: MockerFixture) -> None: +def test_capacity_policy_turn_start_replacement_does_not_require_generation_reserve( + 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, + }, + } + + 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, + recompute_on_capacity=True, + ) + + assert replaced is True + assert request.num_computed_tokens == 0 + assert request.num_prompt_tokens == 10 + + +def test_streaming_window_appends_until_capacity_then_recomputes(mocker: MockerFixture) -> None: request = SimpleNamespace( - _all_token_ids=[0] * 10 + [101, 102], + _all_token_ids=[0] * 4000 + [101, 102], _output_token_ids=[101, 102], - prompt_token_ids=[0] * 10, - num_computed_tokens=12, - num_prompt_tokens=10, + prompt_token_ids=[0] * 4000, + num_computed_tokens=4002, + num_prompt_tokens=4000, num_output_placeholders=0, update_block_hashes=mocker.Mock(), ) @@ -221,7 +289,7 @@ def test_streaming_window_recomputes_each_condition_without_accumulating(mocker: }, } - assert construct_next_stage_streaming_input_prompt( + assert not construct_next_stage_streaming_input_prompt( second, request, max_model_len=4096, @@ -229,17 +297,20 @@ def test_streaming_window_recomputes_each_condition_without_accumulating(mocker: previous_condition_seq=0, condition_seq=1, recompute_previous_chunks=1, + recompute_on_capacity=True, ) - assert request.num_prompt_tokens == 24 - assert second["ids"]["streaming_prompt_previous_codes"] == [101, 102] + assert request.num_computed_tokens == 4002 + assert request.num_prompt_tokens == 4014 + assert "streaming_prompt_previous_codes" not in second["ids"] + assert second["meta"]["streaming_prompt_recompute"] is False request._all_token_ids.extend([201, 202, 203]) request._output_token_ids.extend([201, 202, 203]) - request.num_computed_tokens = 27 + request.num_computed_tokens = 4017 third = { "ids": {"prompt": [2]}, "meta": { - "next_stage_prompt_len": 8, + "next_stage_prompt_len": 60, "next_stage_generation_tokens": 26, }, } @@ -252,11 +323,39 @@ def test_streaming_window_recomputes_each_condition_without_accumulating(mocker: previous_condition_seq=1, condition_seq=2, recompute_previous_chunks=1, + recompute_on_capacity=True, ) - assert request.num_prompt_tokens == 23 + assert request.num_prompt_tokens == 75 + assert request.num_computed_tokens == 0 assert third["ids"]["streaming_prompt_previous_codes"] == [201, 202, 203] assert third["meta"]["streaming_condition_seq"] == 2 + request._all_token_ids.extend([301]) + request._output_token_ids.extend([301]) + request.num_computed_tokens = 76 + fourth = { + "ids": {"prompt": [3]}, + "meta": { + "next_stage_prompt_len": 8, + "next_stage_generation_tokens": 26, + "streaming_condition_seq": 3, + }, + } + + assert not construct_next_stage_streaming_input_prompt( + fourth, + request, + max_model_len=4096, + previous_condition_len=60, + previous_condition_seq=2, + condition_seq=3, + recompute_previous_chunks=1, + recompute_on_capacity=True, + ) + assert request.num_computed_tokens == 76 + assert request.num_prompt_tokens == 84 + assert fourth["meta"]["streaming_prompt_recompute"] is False + def test_capacity_rollover_requires_explicit_window_contract(mocker: MockerFixture) -> None: request = _streaming_request(mocker, num_computed_tokens=4064) @@ -353,13 +452,14 @@ def build_adapter(monkeypatch, mocker: MockerFixture): ) hf_config = SimpleNamespace( - model_type="minicpmtts", + model_type="conditional_chattts", max_position_embeddings=tts_max_model_len, attention_type=tts_attention_type, ) if not flat_tts_config: hf_config = SimpleNamespace( tts_config=SimpleNamespace( + model_type="minicpmtts", max_position_embeddings=tts_max_model_len, attention_type=tts_attention_type, ) @@ -369,6 +469,7 @@ def build_adapter(monkeypatch, mocker: MockerFixture): max_num_seqs=max_num_seqs, max_model_len=max_model_len, hf_config=hf_config, + hf_config_name="tts_config" if stage_id == 1 else None, active_stream_window=active_stream_window, stage_connector_config={ "name": "SharedMemoryConnector", @@ -396,13 +497,24 @@ def _dequeue_load_entry(adapter, request): @pytest.mark.parametrize( - ("attention_type", "expected_previous_chunks"), - [("full_attention", 0), ("sliding_recompute", 1)], + ("attention_type", "expected_previous_chunks", "expected_capacity_trigger"), + [("full_attention", 1, True), ("sliding_recompute", 1, False), ("other", 0, False)], ) -def test_talker_attention_policy_controls_streaming_recompute(build_adapter, attention_type, expected_previous_chunks): +def test_talker_attention_policy_controls_streaming_recompute( + build_adapter, attention_type, expected_previous_chunks, expected_capacity_trigger +): adapter, _ = build_adapter(tts_attention_type=attention_type) assert adapter._streaming_prompt_previous_chunks == expected_previous_chunks + assert adapter._streaming_prompt_recompute_on_capacity is expected_capacity_trigger + + +@pytest.mark.parametrize("stage_id", [0, 2]) +def test_talker_streaming_policy_is_scoped_to_tts_stage(build_adapter, stage_id: int) -> None: + adapter, _ = build_adapter(stage_id=stage_id, tts_attention_type="full_attention") + + assert adapter._streaming_prompt_previous_chunks == 0 + assert adapter._streaming_prompt_recompute_on_capacity is False @pytest.mark.parametrize( @@ -570,23 +682,23 @@ def test_load_poll_ar_requeues_explicitly_replaced_running_prompt(build_adapter) @pytest.mark.parametrize("flat_tts_config", [False, True]) -def test_load_poll_ar_recomputes_native_duplex_prompt_each_condition(build_adapter, flat_tts_config): +def test_load_poll_ar_recomputes_native_duplex_prompt_over_capacity(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", + tts_attention_type="full_attention", 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.prompt_token_ids = [0] * 4039 + request._all_token_ids = [0] * 4039 + previous_codes request._output_token_ids = previous_codes.copy() - request.num_prompt_tokens = 10 - request.num_computed_tokens = 35 + request.num_prompt_tokens = 4039 + request.num_computed_tokens = 4064 request.num_output_placeholders = 0 request.update_block_hashes = Mock() adapter.get_req_chunk[request.request_id] = 1 @@ -626,6 +738,35 @@ def test_load_poll_ar_recomputes_native_duplex_prompt_each_condition(build_adapt assert waiting_queue == [request] +def test_load_poll_ar_keeps_generic_prompt_extension_for_non_tts_stage(build_adapter) -> None: + adapter, connector = build_adapter(stage_id=2, model_mode="ar") + request = _req("req-generic-extension", RequestStatus.RUNNING) + request.resumable = True + request.prompt_token_ids = [1, 2, 3] + request._all_token_ids = [1, 2, 3, 7] + request._output_token_ids = [7] + request.num_prompt_tokens = 3 + request.num_computed_tokens = 4 + request.update_block_hashes = Mock() + adapter.get_req_chunk[request.request_id] = 1 + connector.get.return_value = ( + { + "ids": {"prompt": [1]}, + "meta": {"next_stage_prompt_len": 2}, + }, + 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}) + + assert request.num_computed_tokens == 4 + assert request.prompt_token_ids == [1, 2, 3, 7, 0, 0] + assert request._all_token_ids == [1, 2, 3, 7, 0, 0] + assert request.num_prompt_tokens == 6 + + def test_window_condition_sequence_ignores_control_only_chunks(build_adapter) -> None: adapter, connector = build_adapter( stage_id=1, @@ -644,6 +785,7 @@ def test_window_condition_sequence_ignores_control_only_chunks(build_adapter) -> first = receive( { + "native_duplex": True, "ids": {"prompt": [1]}, "meta": {"next_stage_prompt_len": 10}, } @@ -657,6 +799,7 @@ def test_window_condition_sequence_ignores_control_only_chunks(build_adapter) -> second = receive( { + "native_duplex": True, "ids": {"prompt": [2]}, "meta": {"next_stage_prompt_len": 12}, } @@ -676,6 +819,7 @@ def test_load_poll_ar_reports_invalid_capacity_metadata(build_adapter): "meta": { "next_stage_prompt_len": 4096, "next_stage_generation_tokens": 26, + "replace_streaming_prompt": True, "finished": False, }, }, 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 450246696..f8fdfdcd2 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -19,8 +19,6 @@ 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 ( @@ -168,11 +166,9 @@ 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 + 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, @@ -193,18 +189,6 @@ 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 95d845686..0c5b97ae2 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 @@ -912,10 +912,18 @@ def test_native_duplex_rollover_matches_official_sliding_recompute(mocker) -> No assert build_condition.call_count == 3 -def test_native_duplex_condition_advance_requires_recompute_marker(mocker) -> None: +def test_native_duplex_condition_advance_without_rollover_updates_window_state(mocker) -> None: talker = _make_talker() - condition = torch.ones(2, 2) - mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition) + talker.emb_text = nn.Embedding(1, 2) + talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)]) + first_condition = torch.ones(2, 2) + second_condition = torch.full((2, 2), 2.0) + third_condition = torch.full((2, 2), 3.0) + mocker.patch.object( + talker, + "_build_condition_embeddings", + side_effect=[first_condition, second_condition, second_condition, third_condition], + ) common = { "_omni_is_prefill": True, "_omni_prompt_len": 2, @@ -925,11 +933,105 @@ def test_native_duplex_condition_advance_requires_recompute_marker(mocker) -> No "tts_hidden_states": torch.ones(1, 2), } + talker.preprocess(torch.zeros(2, dtype=torch.long), None, meta={"streaming_condition_seq": 0}, **common) + initial_state = talker._request_condition_states["req-history"] + assert initial_state["condition_seq"] == 0 + assert torch.equal(initial_state["condition"], first_condition) + talker._request_audio_states["req-history"]["recent_codes"] = [1, 2, 3] + + _, embeds, _ = talker.preprocess( + torch.zeros(4, dtype=torch.long), + None, + _omni_prompt_len=4, + meta={"streaming_condition_seq": 1}, + **{key: value for key, value in common.items() if key != "_omni_prompt_len"}, + ) + + assert torch.equal(embeds[-2:], second_condition) + state = talker._request_condition_states["req-history"] + assert state["condition_seq"] == 1 + assert torch.equal(state["condition"], second_condition) + assert state["base_recent_codes"] == (1, 2, 3) + assert talker._request_audio_states["req-history"]["recent_codes"] == [1, 2, 3] + + # A chunked-prefill retry for the same condition is idempotent. + _, retry_embeds, _ = talker.preprocess( + torch.zeros(2, dtype=torch.long), + None, + _omni_num_computed_tokens=2, + _omni_prompt_len=4, + meta={"streaming_condition_seq": 1}, + **{key: value for key, value in common.items() if key != "_omni_prompt_len"}, + ) + assert torch.equal(retry_embeds, second_condition) + assert talker._request_condition_states["req-history"]["base_recent_codes"] == (1, 2, 3) + + previous_codes = [4, 5] + expected_rollover = torch.cat( + [ + second_condition, + talker.emb_code[0](torch.tensor(previous_codes)), + third_condition, + ], + dim=0, + ) + _, rollover_embeds, _ = talker.preprocess( + torch.zeros(expected_rollover.shape[0], dtype=torch.long), + None, + _omni_prompt_len=expected_rollover.shape[0], + ids={"streaming_prompt_previous_codes": previous_codes}, + meta={ + "streaming_condition_seq": 2, + "streaming_prompt_recompute": True, + }, + **{key: value for key, value in common.items() if key != "_omni_prompt_len"}, + ) + assert torch.equal(rollover_embeds, expected_rollover) + rollover_state = talker._request_condition_states["req-history"] + assert rollover_state["base_recent_codes"] == (1, 2, 3, 4, 5) + + +def test_native_duplex_sliding_condition_advance_requires_recompute_marker(mocker) -> None: + talker = _make_talker() + talker._tts_config.attention_type = "sliding_recompute" + condition = torch.ones(2, 2) + mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition) + common = { + "_omni_is_prefill": True, + "_omni_prompt_len": 2, + "_omni_num_computed_tokens": 0, + "request_id": "req-sliding", + "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) +def test_native_duplex_disabled_recompute_policy_accepts_condition_advance(mocker) -> None: + talker = _make_talker() + talker._tts_config.attention_type = "other" + condition = torch.ones(2, 2) + mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition) + common = { + "_omni_is_prefill": True, + "_omni_prompt_len": 2, + "_omni_num_computed_tokens": 0, + "request_id": "req-disabled-policy", + "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) + talker.preprocess(torch.zeros(2, dtype=torch.long), None, meta={"streaming_condition_seq": 1}, **common) + + assert talker._request_condition_states["req-disabled-policy"]["condition_seq"] == 1 + + @pytest.mark.parametrize( ("state", "previous_codes", "error"), [ 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 795047f1c..27037f029 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 @@ -21,6 +21,7 @@ def _output( latent: torch.Tensor, multimodal_output: dict | None = None, token_list: list[int] | None = None, + request_id: str = "req-1", ): mm_output = dict(multimodal_output or {}) mm_output["latent"] = latent @@ -30,7 +31,7 @@ def _output( multimodal_output=mm_output, ) return SimpleNamespace( - request_id="req-1", + request_id=request_id, prompt_token_ids=prompt_ids, outputs=[completion], ) @@ -193,21 +194,47 @@ def test_native_duplex_continuation_appends_only_new_talker_condition() -> None: ) first_input = llm2tts([first], prompt=[{}], _streaming_context=context)[0] + replayed_inputs = llm2tts([first], prompt=[{}], _streaming_context=context) second_input = llm2tts([second], prompt=[{}], _streaming_context=context)[0] context.bridge_states["duplex"]["model_turn_id"] = 8 third_input = llm2tts([third], prompt=[{}], _streaming_context=context)[0] + restarted_input = llm2tts( + [ + _output( + prompt_ids=prompt_ids, + output_ids=first_ids, + latent=torch.arange(24, dtype=torch.float32).reshape(6, 4), + multimodal_output={ + "duplex_prompt_token_ids": prompt_ids, + "meta": token_ids, + }, + request_id="req-2", + ) + ], + prompt=[{}], + _streaming_context=context, + )[0] assert first_input["model_intermediate_buffer"]["ids"]["tts"] == [21, 22] + assert replayed_inputs == [] assert second_input["model_intermediate_buffer"]["ids"]["tts"] == [23, 24] assert third_input["model_intermediate_buffer"]["ids"]["tts"] == [25, 26] assert first_input["model_intermediate_buffer"]["meta"]["turn_start"] is True assert second_input["model_intermediate_buffer"]["meta"]["turn_start"] is False assert third_input["model_intermediate_buffer"]["meta"]["turn_start"] is True + assert first_input["model_intermediate_buffer"]["meta"]["streaming_condition_seq"] == 0 + assert second_input["model_intermediate_buffer"]["meta"]["streaming_condition_seq"] == 1 + # A new turn keeps the same stage request, so its condition sequence stays + # monotonic; a new request/incarnation starts again from zero. + assert third_input["model_intermediate_buffer"]["meta"]["streaming_condition_seq"] == 2 + assert restarted_input["model_intermediate_buffer"]["meta"]["streaming_condition_seq"] == 0 assert first_input["model_intermediate_buffer"]["meta"]["replace_streaming_prompt"] is True 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 first_input["model_intermediate_buffer"]["meta"]["next_stage_generation_tokens"] == 26 assert second_input["model_intermediate_buffer"]["meta"]["next_stage_generation_tokens"] == 26 + assert third_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_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index ad7baf764..329e8ad81 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -6,6 +6,7 @@ from __future__ import annotations import time from collections import defaultdict from collections.abc import Iterable, Iterator +from dataclasses import replace from typing import Any import numpy as np @@ -763,11 +764,15 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): return engine_core_outputs def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None: - """ - Override: Only extend prompt at stage 0, and replace - the existing session with the next streaming update at other stages. + """Apply the next streaming update to a persistent session. - Discards the last sampled output token from the prior input chunk at stage 0. + Stage 0 uses upstream prompt extension. A MiniCPM Talker preserves its + accumulated prompt while it fits, then rebuilds a bounded window and + re-enters admission. Other downstream stages retain their existing + replacement or connector-polling behavior. + + Discards the last sampled output token from the prior input chunk at + stage 0. """ req_id = session.request_id self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids) @@ -818,6 +823,56 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.log_stats: session.record_event(EngineCoreEventType.QUEUED) return + streaming_prompt_payload = next( + ( + info + for info in update_infos + if isinstance(info, dict) + and isinstance(info.get("meta"), dict) + and "next_stage_prompt_len" in info["meta"] + ), + None, + ) + update_streaming_prompt = getattr( + self.chunk_transfer_adapter, + "update_streaming_prompt_for_condition", + None, + ) + if stage_id != 0 and streaming_prompt_payload is not None and callable(update_streaming_prompt): + mm_feature_base = session.num_computed_tokens + try: + replaced = update_streaming_prompt( + streaming_prompt_payload, + session, + update_prompt=True, + ) + except ValueError as exc: + # This streaming update has already been dequeued. Report the + # permanent contract failure so the next scheduling pass + # finishes only this request instead of crashing EngineCore. + self.chunk_transfer_adapter.record_receive_failure(req_id, str(exc)) + return + if replaced is not None: + if replaced: + # The window recipe needs the old prompt and confirmed + # codec ids, so it runs before their KV/encoder state is + # released. The rebuilt prompt is then admitted from zero. + self._release_replaced_streaming_prompt_cache(session) + self._reset_streaming_session_replacement_state(session) + else: + session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1 + if update.mm_features: + # Match upstream streaming-session extension semantics. + # The helper has already appended this condition, so use + # the pre-append confirmed length as the MM offset base. + for mm_feature in update.mm_features: + mm_feature.mm_position = replace( + mm_feature.mm_position, + offset=mm_feature.mm_position.offset + mm_feature_base, + ) + session.mm_features.extend(update.mm_features) + self._finish_streaming_session_update(session, update) + return replace_streaming_prompt = any( isinstance(info, dict) and isinstance(info.get("meta"), dict) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 58c3546cd..aa71245a0 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -147,8 +147,8 @@ class OmniSchedulerMixin: if input_coordinator is not None: input_coordinator.free_finished_request(request_id) - def _replace_streaming_session(self, session: Request, update: StreamingUpdate) -> None: - """Replace a downstream stage's placeholder with its next payload.""" + def _reset_streaming_session_replacement_state(self, session: Request) -> None: + """Fence transport and async state that belongs to an old prompt.""" session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1 adapter = getattr(self, "chunk_transfer_adapter", None) if adapter is not None: @@ -156,8 +156,6 @@ class OmniSchedulerMixin: watermark = getattr(adapter, "requests_num_chunks_sent", None) if watermark is not None: watermark.pop(session.external_req_id, None) - session._output_token_ids.clear() - session._all_token_ids.clear() # In-flight outputs from the previous segment were optimistically # scheduled (async lookahead). Mark them stale so update_from_output # drops them instead of underflowing num_output_placeholders @@ -172,18 +170,15 @@ class OmniSchedulerMixin: session.num_stale_output_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0) session.num_output_placeholders = 0 session.spec_token_ids = [] - new_prompt = update.prompt_token_ids or () - session._all_token_ids.extend(new_prompt) - session.num_computed_tokens = 0 - session.prompt_token_ids = new_prompt + + def _finish_streaming_session_update(self, session: Request, update: StreamingUpdate) -> None: + """Install payload metadata and send an updated session to admission.""" session.additional_information = update.additional_information or None session.model_intermediate_buffer = getattr( update, "model_intermediate_buffer", None, ) - session.update_block_hashes() - session.num_prompt_tokens = len(new_prompt) session.arrival_time = update.arrival_time session.sampling_params = update.sampling_params if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ: @@ -195,6 +190,19 @@ class OmniSchedulerMixin: if self.log_stats: session.record_event(EngineCoreEventType.QUEUED) + def _replace_streaming_session(self, session: Request, update: StreamingUpdate) -> None: + """Replace a downstream stage's placeholder with its next payload.""" + self._reset_streaming_session_replacement_state(session) + 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 = new_prompt + session.update_block_hashes() + session.num_prompt_tokens = len(new_prompt) + self._finish_streaming_session_update(session, update) + def _release_replaced_streaming_prompt_cache(self, session: Request) -> None: """Discard cache state that belongs to a replaced prompt.""" # A prompt replacement is not a normal streaming extension: none of @@ -272,7 +280,7 @@ class OmniSchedulerMixin: self._log_failed_chunk_sends() def _process_chunk_receive_failures(self) -> None: - """Fail requests whose consumed async chunk violated its contract.""" + """Fail requests whose consumed streaming input violated its contract.""" adapter = getattr(self, "chunk_transfer_adapter", None) collector = getattr(adapter, "collect_failed_receive_request_ids", None) if collector is None: @@ -282,7 +290,7 @@ class OmniSchedulerMixin: if not present_ids: return logger.error( - "Marking %d request(s) as FINISHED_ERROR after invalid connector input: %s", + "Marking %d request(s) as FINISHED_ERROR after invalid streaming input: %s", len(present_ids), {request_id: failures[request_id] for request_id in sorted(present_ids)}, ) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 023101223..4e9fd5a77 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -64,13 +64,8 @@ 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 + # Preserve the checkpoint's full-attention Talker prompt while it fits; + # native duplex rolls over to a one-chunk window at context capacity. # Leave device-level headroom for the colocated HiFi-GAN cuDNN workspace. gpu_memory_utilization: 0.15 enforce_eager: false @@ -118,8 +113,10 @@ platforms: cuda: stages: - stage_id: 1 - # Four Talker streams need much less KV than automatic sizing. - # Leave single-GPU headroom for concurrent Stage 2 diffusion batches. + # Full-attention occupancy is bounded by max_num_seqs * max_model_len + # (4 * 4096 = 16384 tokens). Keep the reported CUDA KV token capacity + # above that bound because physical-preemption replay is unsupported. + # The explicit cap leaves device headroom for concurrent Stage 2. kv_cache_memory_bytes: 2147483648 # 2 GiB npu: diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index d585b3ba1..79a1d7887 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -227,16 +227,21 @@ def construct_next_stage_streaming_input_prompt( previous_condition_seq: int | None = None, condition_seq: int | None = None, recompute_previous_chunks: int = 0, + recompute_on_capacity: bool = False, ) -> bool: - """Update a downstream streaming request prompt from connector payload ids. + """Update a downstream streaming request prompt from connector payload. Async-chunk downstream stages are prewarmed before the real Talker prompt is - known. When a Thinker payload carries ``ids.prompt``, this helper: + known. When a Thinker payload carries ``ids.prompt`` or declares the + placeholder length in ``meta.next_stage_prompt_len``, this helper: * Preserves ``num_computed_tokens`` while extending a non-window prompt. - Explicit replacements and one-condition window recomputes reset it. + Explicit replacements and one-condition window recomputes reset it. A + full-attention Talker recomputes only when the next condition would + exceed its context limit. * Moves already-computed output tokens into ``prompt_token_ids``. - * Appends a new placeholder prompt slice sized from the upstream ids. + * Appends a placeholder slice using the declared downstream prompt length, + or derives its length from ``ids.prompt`` for legacy producers. * Refreshes block hashes so the scheduler allocates KV slots for the extended prompt without discarding prior computed state. """ @@ -285,13 +290,29 @@ def construct_next_stage_streaming_input_prompt( 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") + if not isinstance(recompute_on_capacity, bool): + raise ValueError("recompute_on_capacity must be a bool") has_window_contract = recompute_previous_chunks == 1 - window_recompute = not explicit_replacement and has_window_contract and previous_condition_seq is not None + if recompute_on_capacity and not has_window_contract: + raise ValueError("capacity-triggered streaming prompt recompute requires window_size=1") + if ( + recompute_on_capacity + and not explicit_replacement + and previous_condition_seq is not None + and generation_reserve == 0 + ): + raise ValueError("capacity-triggered streaming prompt recompute requires a positive generation reserve") 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 ) + window_recompute = ( + not explicit_replacement + and has_window_contract + and previous_condition_seq is not None + and (exceeds_accumulated_capacity if recompute_on_capacity else True) + ) 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 @@ -369,7 +390,12 @@ def construct_next_stage_streaming_input_prompt( request.update_block_hashes() return True prompt_token_ids = ids.get("prompt", None) - if not prompt_token_ids: + has_declared_prompt_len = ( + isinstance(next_stage_prompt_len, int) + and not isinstance(next_stage_prompt_len, bool) + and next_stage_prompt_len > 0 + ) + if not has_declared_prompt_len and not prompt_token_ids: return False num_computed_tokens = request.num_computed_tokens kept_output_tokens = request._all_token_ids[request.num_prompt_tokens : num_computed_tokens] @@ -378,7 +404,7 @@ def construct_next_stage_streaming_input_prompt( assert request.prompt_token_ids is not None # Extend prompt with kept output tokens. request.prompt_token_ids.extend(kept_output_tokens) - if isinstance(next_stage_prompt_len, int) and next_stage_prompt_len > 0: + if has_declared_prompt_len: next_prompt_len = next_stage_prompt_len else: next_prompt_len = max(1, compute_talker_prompt_ids_length(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 b9b03c9c5..d71d55bf1 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 @@ -51,13 +51,21 @@ class _LoadEntry: 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.""" +def _resolve_talker_streaming_prompt_config(model_config: Any) -> tuple[int, int, bool]: + """Return the Talker context limit and streaming recompute policy.""" max_model_len = int(getattr(model_config, "max_model_len", 0) or 0) + if getattr(model_config, "hf_config_name", None) != "tts_config": + return max_model_len, 0, False 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": + if tts_config is None: tts_config = hf_config + if isinstance(tts_config, Mapping): + tts_model_type = tts_config.get("model_type") + else: + tts_model_type = getattr(tts_config, "model_type", None) + if tts_model_type not in {"conditional_chattts", "minicpmtts"}: + return max_model_len, 0, False tts_max_model_len = ( tts_config.get("max_position_embeddings", 0) if isinstance(tts_config, Mapping) @@ -71,11 +79,17 @@ def _resolve_talker_streaming_prompt_config(model_config: Any) -> tuple[int, int 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 + if tts_config is not None and attention_type is None: + attention_type = "full_attention" + # ``full_attention`` keeps the complete prompt until its next condition + # would exceed the Talker context. It then falls back to the same supported + # one-previous-chunk recompute recipe. An explicitly configured + # ``sliding_recompute`` applies that recipe at every condition boundary. + if attention_type == "sliding_recompute": + return max_model_len, 1, False + if attention_type == "full_attention": + return max_model_len, 1, True + return max_model_len, 0, False class OmniChunkTransferAdapter(OmniTransferAdapterBase): @@ -113,9 +127,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): 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 - ) + ( + self._max_model_len, + self._streaming_prompt_previous_chunks, + self._streaming_prompt_recompute_on_capacity, + ) = _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: @@ -962,6 +978,99 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self._promote_active_streams(waiting_queue) self._preempt_non_active_running(waiting_queue, running_queue) + def update_streaming_prompt_for_condition( + self, + payload_data: dict[str, Any], + request: Request, + *, + update_prompt: bool, + ) -> bool | None: + """Track one Talker condition and optionally update its prompt. + + Returns ``None`` when the payload is not governed by the Talker + streaming-window contract. Otherwise returns whether the prompt was + replaced, as opposed to extended in place. + """ + if self._streaming_prompt_previous_chunks != 1: + return None + meta = payload_data.get("meta") + if not isinstance(meta, dict): + return None + condition_len = meta.get("next_stage_prompt_len") + if not isinstance(condition_len, int) or isinstance(condition_len, bool) or condition_len <= 0: + return None + + request_id = request.request_id + previous_condition_seq = self._streaming_condition_seqs.get(request_id) + previous_condition_len = self._streaming_condition_lengths.get(request_id) + if (previous_condition_seq is None) != (previous_condition_len is None): + raise ValueError("streaming Talker condition tracking is incomplete") + if previous_condition_seq is None and not self.receives_chunks: + # A sender-only Talker is admitted with its first real condition, + # before this adapter observes any StreamingUpdate. Seed that + # condition as sequence zero from the active request so its first + # continuation follows the same lifecycle as receiver stages. + initial_info = getattr(request, "model_intermediate_buffer", None) + initial_meta = initial_info.get("meta") if isinstance(initial_info, dict) else None + initial_condition_len = ( + initial_meta.get("next_stage_prompt_len") if isinstance(initial_meta, dict) else None + ) + initial_condition_seq = ( + initial_meta.get("streaming_condition_seq") if isinstance(initial_meta, dict) else None + ) + if ( + isinstance(initial_info, dict) + and initial_info.get("native_duplex") is True + and isinstance(initial_condition_len, int) + and not isinstance(initial_condition_len, bool) + and initial_condition_len > 0 + and isinstance(initial_condition_seq, int) + and not isinstance(initial_condition_seq, bool) + and initial_condition_seq == 0 + ): + if initial_condition_len != request.num_prompt_tokens: + raise ValueError( + "initial streaming Talker condition length does not match its prompt: " + f"condition={initial_condition_len}, prompt={request.num_prompt_tokens}" + ) + previous_condition_seq = initial_condition_seq + previous_condition_len = initial_condition_len + supplied_condition_seq = meta.get("streaming_condition_seq") + if supplied_condition_seq is None: + if payload_data.get("native_duplex") is True and not self.receives_chunks: + raise ValueError("native duplex Talker condition is missing streaming_condition_seq") + condition_seq = previous_condition_seq + 1 if previous_condition_seq is not None else 0 + meta["streaming_condition_seq"] = condition_seq + elif ( + not isinstance(supplied_condition_seq, int) + or isinstance(supplied_condition_seq, bool) + or supplied_condition_seq < 0 + ): + raise ValueError("streaming_condition_seq must be a non-negative integer") + else: + condition_seq = supplied_condition_seq + expected_condition_seq = previous_condition_seq + 1 if previous_condition_seq is not None else 0 + if condition_seq != expected_condition_seq: + raise ValueError( + "native duplex Talker condition sequence is not contiguous: " + f"expected={expected_condition_seq}, received={condition_seq}" + ) + replaced = False + if update_prompt: + replaced = construct_next_stage_streaming_input_prompt( + payload_data, + request, + max_model_len=self._max_model_len, + previous_condition_len=previous_condition_len, + previous_condition_seq=previous_condition_seq, + condition_seq=condition_seq, + recompute_previous_chunks=self._streaming_prompt_previous_chunks, + recompute_on_capacity=self._streaming_prompt_recompute_on_capacity, + ) + self._streaming_condition_lengths[request_id] = condition_len + self._streaming_condition_seqs[request_id] = condition_seq + return replaced + 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): @@ -981,31 +1090,31 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): 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: + previous_condition_seq = self._streaming_condition_seqs.get(request_id) + meta = payload_data.get("meta") + assert isinstance(meta, dict) 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 = self.update_streaming_prompt_for_condition( + payload_data, + request, + update_prompt=update_prompt, + ) + if replaced is None and update_prompt: + # Non-window AR receivers still use the generic + # prompt-extension contract from the base adapter. 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) + 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 @@ -1013,10 +1122,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): 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/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 7c87f1a8b..bdc418f05 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 @@ -319,10 +319,26 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): 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}" - ) + attention_type = getattr(self._tts_config, "attention_type", "full_attention") + if attention_type == "sliding_recompute": + raise ValueError( + "a native duplex Talker condition advanced without its streaming recompute marker: " + f"current={condition_seq}, previous={previous_seq}" + ) + audio_state = self._request_audio_states.get(request_id) + recent_codes = audio_state.get("recent_codes") if isinstance(audio_state, dict) else None + if isinstance(recent_codes, list): + base_recent_codes = tuple(int(code_id) for code_id in recent_codes[-_CODEC_PENALTY_WINDOW:]) + else: + base_recent_codes = state.get("base_recent_codes") + if not isinstance(base_recent_codes, tuple): + raise ValueError("streaming Talker condition lost its frozen codec history") + states[request_id] = { + "condition_seq": condition_seq, + "condition": current_condition.detach().clone(), + "base_recent_codes": base_recent_codes, + } + return current_condition if "active_embeddings" in state: raise ValueError("an active streaming recompute was replayed without its recompute marker") return current_condition 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 c1cbd14bc..80995f92e 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 @@ -637,6 +637,7 @@ def _native_duplex_segment_output_ids( state["request_id"] = request_id state["sent_output_len"] = 0 state["sent_output_ids"] = [] + state["condition_seq"] = -1 previous_turn_id = state.get("turn_id") sent_len = state.get("sent_output_len", 0) prev_output_ids = state.get("sent_output_ids", []) @@ -1000,6 +1001,8 @@ def llm2tts( if native_turn_end_handoff: model_intermediate_buffer.setdefault("meta", {})["turn_end"] = True + condition_sequence_state = None + condition_sequence_value = None if handoff_ids is not None and handoff_hidden is not None: condition_suffix_length = 1 if is_native_duplex_handoff else 2 condition_length = max(len(handoff_ids), len(handoff_hidden)) + condition_suffix_length @@ -1011,6 +1014,20 @@ def llm2tts( 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 + bridge_states = getattr(_streaming_context, "bridge_states", None) + handoff_state = bridge_states.get("minicpmo45_tts_handoff") if isinstance(bridge_states, dict) else None + if not isinstance(handoff_state, dict) or handoff_state.get("request_id") != str(llm_output.request_id): + raise ValueError("native duplex Talker handoff is missing request-local streaming state") + previous_condition_seq = handoff_state.get("condition_seq", -1) + if ( + not isinstance(previous_condition_seq, int) + or isinstance(previous_condition_seq, bool) + or previous_condition_seq < -1 + ): + raise ValueError("native duplex Talker condition sequence state is invalid") + condition_sequence_state = handoff_state + condition_sequence_value = previous_condition_seq + 1 + handoff_meta["streaming_condition_seq"] = condition_sequence_value # 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: @@ -1033,6 +1050,8 @@ def llm2tts( mm_processor_kwargs=None, ) ) + if condition_sequence_state is not None: + condition_sequence_state["condition_seq"] = condition_sequence_value if native_turn_end_handoff: bridge_states = getattr(_streaming_context, "bridge_states", None) duplex_state = bridge_states.get("duplex") if isinstance(bridge_states, dict) else None