[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>
2026-09-06 22:25:35 +08:00
1 parent 715b8b8746
commit ff6e906a25
@@ -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()