[Model] Preserve MiniCPM-o Talker full attention until capacity (#6799)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

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()

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

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

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

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]
 
 

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)

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

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:

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

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):

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

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