[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]