[Bugfix][MiniCPM-o] Fix async-chunk snapshot replacement and prompt cleanup (#6406)

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

@@ -33,6 +33,9 @@ def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
     sched.log_stats = False
     sched.chunk_transfer_adapter = None
     sched.skipped_waiting = set()
+    sched._free_request_blocks = MagicMock()
+    sched.encoder_cache_manager = MagicMock()
+    sched._inflight_prefills = set()
     return sched
 
 
@@ -276,29 +279,33 @@ def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None:
         "meta": {"turn_eos_token_id": 99},
     }
     assert session.status == RequestStatus.WAITING
+    sched._free_request_blocks.assert_called_once_with(session)
+    sched.encoder_cache_manager.free.assert_called_once_with(session)
 
 
-def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None:
+def test_explicit_model_intermediate_prompt_replacement_releases_cache_and_watermark() -> None:
     sched = _make_scheduler(stage_id=1)
+    session = _make_request()
     sched.chunk_transfer_adapter = SimpleNamespace(
         receives_chunks=False,
         segment_finished_requests=set(),
+        requests_num_chunks_sent={session.external_req_id: 59},
     )
-    session = _make_request()
     session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
     session.prompt_token_ids = [0] * 59
     session._all_token_ids.clear()
     session._all_token_ids.extend(session.prompt_token_ids)
     session.num_prompt_tokens = 59
     session.num_computed_tokens = 59
+    session.num_in_flight_tokens = 2
     update = _make_update([0] * 10)
     update.additional_information = None
     update.model_intermediate_buffer = {
         "ids": {"tts": list(range(8))},
         "hidden_states": {"tts": [[0.0]] * 8},
         "meta": {
-            "replace_streaming_prompt": True,
             "next_stage_prompt_len": 10,
+            "replace_streaming_prompt": True,
         },
     }
 
@@ -308,6 +315,42 @@ def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None
     assert list(session._all_token_ids) == [0] * 10
     assert session.num_prompt_tokens == 10
     assert session.num_computed_tokens == 0
+    assert session.num_stale_output_tokens == 2
     assert session.additional_information is None
     assert session.model_intermediate_buffer == update.model_intermediate_buffer
     assert session.status == RequestStatus.WAITING
+    assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
+    sched._free_request_blocks.assert_called_once_with(session)
+    sched.encoder_cache_manager.free.assert_called_once_with(session)
+
+
+def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
+    sched = _make_scheduler(stage_id=1)
+    session = _make_request()
+    session.external_req_id = "external-ar-streaming-test"
+    session.num_in_flight_tokens = 2
+    # _update_request_as_session() may have already fenced this frame before
+    # the connector marks the explicit replacement ready.
+    session.num_stale_output_tokens = 2
+    session.num_output_placeholders = 2
+    session.spec_token_ids = [-1, -1]
+    sched.requests = {session.request_id: session}
+    sched._inflight_prefills.add(session)
+    sched.chunk_transfer_adapter = SimpleNamespace(
+        replaced_streaming_prompt_ids={session.request_id},
+        requests_with_ready_chunks={session.request_id},
+        requests_num_chunks_sent={session.external_req_id: 4090},
+    )
+
+    sched._reset_ready_async_chunk_replacements()
+    sched._reset_ready_async_chunk_replacements()
+
+    sched._free_request_blocks.assert_called_once_with(session)
+    sched.encoder_cache_manager.free.assert_called_once_with(session)
+    assert session not in sched._inflight_prefills
+    assert session.num_stale_output_tokens == 2
+    assert session.num_output_placeholders == 0
+    assert session.spec_token_ids == []
+    assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set()
+    assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id}
+    assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -208,6 +208,55 @@ def test_load_poll(build_adapter):
     assert "req-1" not in adapter._pending_load_reqs
 
 
+def test_load_poll_ar_requeues_explicitly_replaced_running_prompt(build_adapter):
+    adapter, connector = build_adapter(stage_id=1, model_mode="ar")
+    request = _req("req-replace", RequestStatus.RUNNING, external_req_id="external-replace")
+    request.resumable = True
+    request.prompt_token_ids = [0] * 4
+    request._all_token_ids = [0] * 4
+    request._output_token_ids = []
+    request.num_prompt_tokens = 4
+    request.num_computed_tokens = 4
+    request.update_block_hashes = Mock()
+    adapter.get_req_chunk[request.request_id] = 1
+    adapter.requests_num_chunks_sent[request.external_req_id] = 4
+    adapter.request_ids_mapping[request.request_id] = request.external_req_id
+    connector.get.return_value = (
+        {
+            "ids": {"prompt": [1]},
+            "meta": {
+                "next_stage_prompt_len": 10,
+                "replace_streaming_prompt": True,
+                "finished": False,
+            },
+        },
+        1,
+    )
+
+    assert adapter._poll_single_request(request) is True
+    assert request.num_computed_tokens == 0
+    assert request.prompt_token_ids == [0] * 10
+    assert request.request_id in adapter.replaced_streaming_prompt_ids
+
+    request.status = RequestStatus.WAITING_FOR_CHUNK
+    running_queue = [request]
+    waiting_queue = DummyWaitingQueue()
+    adapter.process_pending_chunks(waiting_queue, running_queue)
+    assert running_queue == []
+    assert waiting_queue == [request]
+    assert request.status == RequestStatus.WAITING
+
+    adapter.requests_num_chunks_sent[request.external_req_id] = 9
+    adapter.postprocess_scheduler_output(
+        SimpleNamespace(
+            scheduled_new_reqs=[SimpleNamespace(req_id=request.request_id)],
+            scheduled_cached_reqs=SimpleNamespace(req_ids=[]),
+        )
+    )
+    assert request.request_id not in adapter.replaced_streaming_prompt_ids
+    assert request.external_req_id not in adapter.requests_num_chunks_sent
+
+
 def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter):
     adapter, connector = build_adapter(stage_id=1, model_mode="generation")
     request = _req("req-tensor", RequestStatus.WAITING, external_req_id="external-tensor")
@@ -573,6 +622,92 @@ def test_load_poll_non_ar_merges_into_existing_additional_information(build_adap
     assert "req-non-ar" in adapter.upstream_exhausted_requests
 
 
+def test_load_poll_generation_segment_marker_replaces_previous_chunk(build_adapter):
+    adapter, connector = build_adapter(stage_id=2, model_mode="generation")
+    request = _req("req-marker", RequestStatus.WAITING, external_req_id="external-marker")
+    request.additional_information = {
+        "codes": {"audio": torch.tensor([1, 2])},
+        "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
+    }
+    connector.get.return_value = (
+        {
+            "codes": {
+                "audio": torch.tensor([7, 8], dtype=torch.long),
+                "ref": torch.tensor([0.1, -0.1]),
+            },
+            "meta": {
+                "finished": torch.tensor(False, dtype=torch.bool),
+                "is_segment_finished": torch.tensor(True, dtype=torch.bool),
+                "request_finished": torch.tensor(False, dtype=torch.bool),
+                "replace_runtime_additional_information": True,
+            },
+        },
+        1,
+    )
+
+    assert adapter._poll_single_request(request) is True
+
+    assert request.prompt_token_ids == [7, 8]
+    assert "audio" not in request.additional_information["codes"]
+    torch.testing.assert_close(
+        request.additional_information["codes"]["ref"],
+        torch.tensor([0.1, -0.1]),
+    )
+    assert not {"cache_epoch", "chunk_seq", "last_chunk"}.intersection(request.additional_information["meta"])
+    assert request.request_id in adapter.segment_finished_requests
+
+
+def test_load_poll_generation_empty_replacement_snapshot_is_ready(build_adapter):
+    adapter, connector = build_adapter(stage_id=2, model_mode="generation")
+    request = _req("req-empty-marker", RequestStatus.WAITING, external_req_id="external-empty-marker")
+    request.additional_information = {
+        "codes": {"audio": torch.tensor([1, 2])},
+        "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
+    }
+    connector.get.return_value = (
+        {
+            "meta": {
+                "is_segment_finished": torch.tensor(True, dtype=torch.bool),
+                "replace_runtime_additional_information": True,
+            }
+        },
+        1,
+    )
+
+    assert adapter._poll_single_request(request) is True
+
+    assert request.prompt_token_ids == [0]
+    assert "codes" not in request.additional_information
+    assert request.additional_information["meta"]["replace_runtime_additional_information"] is True
+    assert request.request_id in adapter.segment_finished_requests
+    assert request.request_id in adapter._finished_load_reqs
+
+
+def test_load_poll_generation_without_snapshot_marker_keeps_incremental_state(build_adapter):
+    adapter, connector = build_adapter(stage_id=2, model_mode="generation")
+    request = _req("req-incremental", RequestStatus.WAITING, external_req_id="external-incremental")
+    request.additional_information = {
+        "codes": {"audio": torch.tensor([1, 2])},
+        "meta": {"cache_epoch": 3, "chunk_seq": 2},
+    }
+    connector.get.return_value = (
+        {
+            "meta": {
+                "finished": torch.tensor(False, dtype=torch.bool),
+                "phase": "decode",
+            }
+        },
+        1,
+    )
+
+    assert adapter._poll_single_request(request) is False
+
+    assert torch.equal(request.additional_information["codes"]["audio"], torch.tensor([1, 2]))
+    assert request.additional_information["meta"]["cache_epoch"] == 3
+    assert request.additional_information["meta"]["chunk_seq"] == 2
+    assert request.additional_information["meta"]["phase"] == "decode"
+
+
 def test_load_poll_ar_request_additional_information_concats_tensors(build_adapter):
     adapter, connector = build_adapter(stage_id=2, model_mode="ar")
     request = _req("req-merged", RequestStatus.WAITING, external_req_id="ext-merged")

tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py

@@ -537,6 +537,15 @@ def test_empty_final_ignores_generation_scheduler_placeholder_token():
         # The runner injects the engine request id on every step (GPU
         # _preprocess, NPU _gather_runtime_additional_information)...
         {"request_id": "a", "meta": {"request_id": "a"}},
+        # Runtime snapshot bookkeeping alone is not a Talker payload.
+        {
+            "request_id": "a",
+            "meta": {
+                "request_id": "a",
+                "num_processed_tokens": 0,
+                "resumable": True,
+            },
+        },
         # ...but a pre-warm step can also reach the model with nothing at all.
         {},
     ],

tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py

@@ -423,6 +423,12 @@ def test_duplex_turn_end_closes_epoch_and_next_turn_restarts_sequence() -> None:
         request,
         True,
     )
+    duplicate_boundary = tts2code2wav_async_chunk(
+        manager,
+        _duplex_delta(turn_id=7, turn_end=True),
+        request,
+        True,
+    )
     next_turn = tts2code2wav_async_chunk(
         manager,
         _duplex_delta(20, turn_id=8),
@@ -435,6 +441,10 @@ def test_duplex_turn_end_closes_epoch_and_next_turn_restarts_sequence() -> None:
     assert turn_end.meta.last_chunk is True
     assert turn_end.meta.turn_end is True
     assert turn_end.meta.is_segment_finished.item() is True
+    assert turn_end.meta.replace_runtime_additional_information is True
+    assert duplicate_boundary is not None
+    assert duplicate_boundary.meta.is_segment_finished.item() is True
+    assert duplicate_boundary.meta.replace_runtime_additional_information is True
     assert next_turn.meta.cache_epoch == turn_end.meta.cache_epoch + 1
     assert next_turn.meta.chunk_seq == 0
     assert next_turn.meta.last_chunk is False

tests/platforms/npu/test_minicpmo_code2wav_npugraph.py

@@ -9,6 +9,8 @@ import pytest
 import torch
 import torch.nn as nn
 
+pytest.importorskip("vllm_ascend")
+
 from tests.helpers.mark import hardware_marks
 from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import (
     BatchedToken2Wav,
@@ -21,6 +23,8 @@ from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import (
 from vllm_omni.platforms.npu.models.step_audio2_token2wav import (
     npu_token2wav_sdpa_context,
 )
+from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 
 pytestmark = [
     pytest.mark.core_model,
@@ -47,6 +51,25 @@ class _Token2Wav:
         self.speech_window = torch.ones(4, device="npu")
 
 
+def test_npu_runner_inherits_runtime_snapshot_replacement():
+    assert issubclass(OmniNPUModelRunner, OmniGPUModelRunner)
+    assert OmniNPUModelRunner._replace_intermediate_buffer is OmniGPUModelRunner._replace_intermediate_buffer
+    runner = object.__new__(OmniNPUModelRunner)
+    request = SimpleNamespace(additional_information_cpu=None)
+    runner.requests = {"req": request}
+    runner.model_intermediate_buffer = {
+        "req": {
+            "codes": {"audio": torch.tensor([1, 2])},
+            "meta": {"last_chunk": True},
+        }
+    }
+
+    runner._replace_intermediate_buffer("req", {"meta": {"is_segment_finished": True}})
+
+    assert runner.model_intermediate_buffer["req"] == {"meta": {"is_segment_finished": True}}
+    assert request.additional_information_cpu == runner.model_intermediate_buffer["req"]
+
+
 def test_minicpmo_cfm_estimator_npugraph_matches_eager():
     decoder_dit = pytest.importorskip("cosyvoice2.flow.decoder_dit")
     prepare_code2wav_graph_runtime()

tests/worker/test_omni_gpu_model_runner.py

@@ -543,6 +543,80 @@ def test_update_additional_information_deserializes_new_request_payload():
     )
 
 
+def test_streaming_new_request_marker_replaces_terminal_chunk_snapshot():
+    from vllm_omni.engine.serialization import serialize_additional_information
+
+    runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
+    runner.model.replace_runtime_additional_information = True
+    terminal = {
+        "codes": {"audio": torch.tensor([1, 2])},
+        "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
+    }
+    peer = {
+        "codes": {"audio": torch.tensor([9])},
+        "meta": {"cache_epoch": 3, "chunk_seq": 1, "last_chunk": False},
+    }
+    runner.model_intermediate_buffer.update(r1=terminal, r2=peer)
+    marker = {
+        "meta": {
+            "finished": False,
+            "is_segment_finished": True,
+            "request_finished": False,
+            "replace_runtime_additional_information": True,
+        }
+    }
+    new_req = SimpleNamespace(
+        req_id="r1",
+        model_intermediate_buffer=marker,
+        additional_information=serialize_additional_information(terminal),
+    )
+
+    OmniGPUModelRunner._update_streaming_input_additional_info(runner, new_req, "r1")
+    OmniGPUModelRunner._update_additional_information(
+        runner,
+        SimpleNamespace(
+            scheduled_new_reqs=[new_req],
+            scheduled_cached_reqs=SimpleNamespace(),
+        ),
+    )
+
+    info = runner.model_intermediate_buffer["r1"]
+    assert "codes" not in info
+    assert info["meta"] == {
+        **marker["meta"],
+        "num_processed_tokens": 0,
+        "resumable": True,
+    }
+    assert runner.requests["r1"].additional_information_cpu == info
+    assert runner.model_intermediate_buffer["r2"] == peer
+
+
+def test_cached_empty_marker_replaces_terminal_chunk_snapshot():
+    runner = _make_runner(req_ids=("r1",), hidden_size=4)
+    runner.model.replace_runtime_additional_information = True
+    runner.model_intermediate_buffer["r1"] = {
+        "codes": {"audio": torch.tensor([1, 2])},
+        "meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
+    }
+    marker = {
+        "meta": {
+            "is_segment_finished": torch.tensor(True, dtype=torch.bool),
+            "replace_runtime_additional_information": True,
+        }
+    }
+
+    OmniGPUModelRunner._update_additional_information(
+        runner,
+        SimpleNamespace(
+            scheduled_new_reqs=[],
+            scheduled_cached_reqs=SimpleNamespace(additional_information={"r1": marker}),
+        ),
+    )
+
+    assert runner.model_intermediate_buffer["r1"] == marker
+    assert runner.requests["r1"].additional_information_cpu == marker
+
+
 def test_update_intermediate_buffer_skips_empty_update():
     """Validate that an empty update dict is a no-op."""
     runner = _make_runner(req_ids=("r1",), hidden_size=4)

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -679,6 +679,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             for info in update_infos
         )
         if replace_streaming_prompt:
+            self._release_replaced_streaming_prompt_cache(session)
             self._replace_streaming_session(session, update)
             return
         super()._update_request_as_session(session, update)

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -149,6 +149,9 @@ class OmniSchedulerMixin:
         adapter = getattr(self, "chunk_transfer_adapter", None)
         if adapter is not None:
             adapter.segment_finished_requests.discard(session.request_id)
+            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
@@ -159,7 +162,10 @@ class OmniSchedulerMixin:
         # pre-replacement frame will drain, so the counter reaches exactly
         # zero; a placeholder-based seed swallowed valid new-segment frames
         # whenever placeholder counts diverged from scheduled counts.
-        session.num_stale_output_tokens += int(getattr(session, "num_in_flight_tokens", 0) or 0)
+        # num_in_flight_tokens already includes any undrained stale share.
+        # Assign instead of accumulating so callers that fenced the same
+        # rollover before entering this helper do not count it twice.
+        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 ()
@@ -185,6 +191,44 @@ class OmniSchedulerMixin:
         if self.log_stats:
             session.record_event(EngineCoreEventType.QUEUED)
 
+    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
+        # the old KV blocks or encoder state is valid for the new prompt. Use
+        # the scheduler's block-free path so an in-flight GPU step is fenced
+        # correctly before the blocks return to the pool.
+        self._free_request_blocks(session)
+        self.encoder_cache_manager.free(session)
+        getattr(self, "_inflight_prefills", set()).discard(session)
+
+    def _reset_ready_async_chunk_replacements(self) -> None:
+        """Release stale cache state after an async-chunk prompt rollover."""
+        adapter = getattr(self, "chunk_transfer_adapter", None)
+        if adapter is None:
+            return
+        replaced_ids = getattr(adapter, "replaced_streaming_prompt_ids", None)
+        ready_ids = getattr(adapter, "requests_with_ready_chunks", None)
+        if not replaced_ids or not ready_ids:
+            return
+
+        for request_id in tuple(replaced_ids & ready_ids):
+            request = self.requests.get(request_id)
+            if request is None:
+                replaced_ids.discard(request_id)
+                continue
+            # The streaming update may already have fenced this same in-flight
+            # frame. Seed idempotently so the replacement does not count it twice.
+            request.num_stale_output_tokens = int(getattr(request, "num_in_flight_tokens", 0) or 0)
+            request.num_output_placeholders = 0
+            request.spec_token_ids = []
+            self._release_replaced_streaming_prompt_cache(request)
+            watermark = getattr(adapter, "requests_num_chunks_sent", None)
+            if watermark is not None:
+                watermark.pop(request.external_req_id, None)
+            # Consume this marker after the one-time cache reset. The separate
+            # ready-chunk marker remains until scheduler admission succeeds.
+            replaced_ids.discard(request_id)
+
     def _consume_pending_connector_output(self, model_mode: str) -> None:
         """Drain ``self._latest_omni_connector_output`` into the coordinator.
 
@@ -216,6 +260,7 @@ class OmniSchedulerMixin:
                 self.running,
                 scheduler_requests=self.requests,
             )
+            self._reset_ready_async_chunk_replacements()
             self._process_pending_chunk_timeouts()
             self._log_failed_chunk_sends()
 

vllm_omni/data_entry_keys.py

@@ -75,6 +75,7 @@ class OmniPayloadMeta(TypedDict, total=False):
     num_processed_tokens: int
     next_stage_prompt_len: int
     replace_streaming_prompt: bool
+    replace_runtime_additional_information: bool
     ar_width: int
     eol_token_id: int
     visual_token_start_id: int
@@ -175,6 +176,7 @@ class MetaStruct(_StructBase):
     num_processed_tokens: int | None = None
     next_stage_prompt_len: int | None = None
     replace_streaming_prompt: bool | None = None
+    replace_runtime_additional_information: bool | None = None
     ar_width: int | None = None
     eol_token_id: int | None = None
     visual_token_start_id: int | None = None

vllm_omni/distributed/omni_connectors/adapter.py

@@ -216,13 +216,14 @@ def compute_talker_prompt_ids_length(prompt_ids: list[int]) -> int:
     return sum_user_len + assistant_len
 
 
-def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], request: Any) -> None:
+def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], request: Any) -> bool:
     """Update a downstream streaming request prompt from connector payload ids.
 
     Async-chunk downstream stages are prewarmed before the real Talker prompt is
     known. When a Thinker payload carries ``ids.prompt``, this helper:
 
-    * Preserves ``num_computed_tokens`` (the scheduler token watermark).
+    * Preserves ``num_computed_tokens`` while extending the current prompt.
+      Explicit replacements reset it with the prompt.
     * Moves already-computed output tokens into ``prompt_token_ids``.
     * Appends a new placeholder prompt slice sized from the upstream ids.
     * Refreshes block hashes so the scheduler allocates KV slots for the
@@ -230,8 +231,8 @@ def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], re
     """
     ids = payload_data.get("ids", {})
     meta = payload_data.get("meta", {})
-    replace_prompt = isinstance(meta, dict) and meta.get("replace_streaming_prompt") is True
     next_stage_prompt_len = meta.get("next_stage_prompt_len") if isinstance(meta, dict) else None
+    replace_prompt = isinstance(meta, dict) and meta.get("replace_streaming_prompt") is True
     if replace_prompt and isinstance(next_stage_prompt_len, int) and next_stage_prompt_len > 0:
         # Some downstream stages consume complete, independently conditioned
         # segments instead of extending an existing KV prefix. The producer
@@ -244,10 +245,10 @@ def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], re
         request.num_computed_tokens = 0
         request.num_prompt_tokens = next_stage_prompt_len
         request.update_block_hashes()
-        return
+        return True
     prompt_token_ids = ids.get("prompt", None)
     if not prompt_token_ids:
-        return
+        return False
     num_computed_tokens = request.num_computed_tokens
     kept_output_tokens = request._all_token_ids[request.num_prompt_tokens : num_computed_tokens]
     del request._all_token_ids[num_computed_tokens:]
@@ -264,3 +265,4 @@ def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], re
     request.prompt_token_ids.extend(new_prompt or ())
     request.update_block_hashes()
     request.num_prompt_tokens = len(request.prompt_token_ids)
+    return False

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -84,6 +84,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.waiting_for_chunk_waiting_requests: deque[Any] = deque()
         self.waiting_for_chunk_running_requests: deque[Any] = deque()
         self.requests_with_ready_chunks = set()
+        self.replaced_streaming_prompt_ids: set[str] = set()
         self.requests_origin_status = {}
         self._active_streams: dict[str, Any] = {}
         # Private hold-queue for non-active running requests. Restored to
@@ -246,7 +247,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 replace_prompt = meta.get("replace_streaming_prompt") is True
                 if getattr(request, "resumable", False) and (chunk_id > 0 or replace_prompt):
                     # For new streaming input segment, we should update prompt from payload
-                    construct_next_stage_streaming_input_prompt(payload_data, request)
+                    replaced = construct_next_stage_streaming_input_prompt(payload_data, request)
+                    if replaced:
+                        self.replaced_streaming_prompt_ids.add(req_id)
 
                 if payload_finished:
                     self.upstream_exhausted_requests.add(req_id)
@@ -275,15 +278,26 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 else:
                     prompt_token_ids = new_ids
                 request.prompt_token_ids = prompt_token_ids
+                # Full-snapshot producers opt in explicitly; generation and
+                # diffusion models keep their existing incremental merge.
                 prev_info = getattr(request, "additional_information", None)
-                info = dict(prev_info) if isinstance(prev_info, dict) else {}
+                replace_snapshot = meta.get("replace_runtime_additional_information") is True
+                info = {} if replace_snapshot else (dict(prev_info) if isinstance(prev_info, dict) else {})
                 for key, value in payload_data.items():
                     if key == "codes":
-                        if use_tensor_codes and isinstance(value, dict):
+                        if isinstance(value, dict):
                             existing_sub = info.get(key)
                             merged_sub = dict(existing_sub) if isinstance(existing_sub, dict) else {}
-                            merged_sub.update(value)
-                            info[key] = merged_sub
+                            for subkey, subvalue in value.items():
+                                # A 1-D audio tensor is represented by the
+                                # placeholder prompt above, but sibling fields
+                                # such as the reference voice still belong in
+                                # the current runtime snapshot.
+                                if subkey == "audio" and not use_tensor_codes:
+                                    continue
+                                merged_sub[subkey] = subvalue
+                            if merged_sub:
+                                info[key] = merged_sub
                         continue
                     if isinstance(value, dict):
                         existing_sub = info.get(key)
@@ -455,6 +469,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.segment_finished_requests.discard(request_id)
         self.get_req_chunk.pop(request_id, None)
         self.requests_with_ready_chunks.discard(request_id)
+        self.replaced_streaming_prompt_ids.discard(request_id)
         self.request_ids_mapping.pop(request_id, None)
         self.requests_origin_status.pop(request_id, None)
         self._discard_from_chunk_deque(self.waiting_for_chunk_waiting_requests, request_id)
@@ -564,6 +579,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 RequestStatus.RUNNING,
                 self._finished_load_reqs,
             )
+            self._requeue_replaced_prompts(waiting_queue, running_queue)
             while len(running_queue) > self.scheduler_max_num_seqs:
                 request = running_queue.pop()
                 request.status = RequestStatus.PREEMPTED
@@ -578,9 +594,24 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self._process_chunk_queue(
             running_queue, self.waiting_for_chunk_running_requests, RequestStatus.RUNNING, self._finished_load_reqs
         )
+        self._requeue_replaced_prompts(waiting_queue, running_queue)
         self._promote_active_streams(waiting_queue)
         self._preempt_non_active_running(waiting_queue, running_queue)
 
+    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):
+            request_id = request.request_id
+            if (
+                request_id not in self.replaced_streaming_prompt_ids
+                or request_id not in self.requests_with_ready_chunks
+            ):
+                continue
+            running_queue.remove(request)
+            request.status = RequestStatus.WAITING
+            self.requests_origin_status[request_id] = RequestStatus.WAITING
+            waiting_queue.add_request(request)
+
     def _promote_active_streams(self, queue: Any) -> None:
         if len(self._active_streams) >= self._active_window:
             return
@@ -852,6 +883,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             for req_data in scheduler_output.scheduled_new_reqs:
                 if req_data.req_id in self.requests_with_ready_chunks:
                     self.requests_with_ready_chunks.remove(req_data.req_id)
+                if req_data.req_id in self.replaced_streaming_prompt_ids:
+                    external_req_id = self.request_ids_mapping.get(req_data.req_id)
+                    if external_req_id is not None:
+                        self.requests_num_chunks_sent.pop(external_req_id, None)
+                    self.replaced_streaming_prompt_ids.remove(req_data.req_id)
 
         if scheduler_output.scheduled_cached_reqs:
             for req_id in scheduler_output.scheduled_cached_reqs.req_ids:

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py

@@ -69,6 +69,17 @@ def _codec_tensor(value: Any, fallback: torch.Tensor) -> torch.Tensor:
 # OmniGPUModelRunner._preprocess and the NPU _gather_runtime_additional_information
 # override). A step carrying only these has no producer payload at all.
 _RUNNER_STAMPED_KEYS = frozenset({"request_id", "req_id", "generated_len", "meta"})
+_PRODUCER_META_KEYS = frozenset(
+    {
+        "cache_epoch",
+        "chunk_seq",
+        "code_flat_numel",
+        "last_chunk",
+        "prompt_cache_id",
+        "prompt_wav",
+        "ref_audio_sr",
+    }
+)
 
 
 def _carries_stage_payload(info: Mapping[str, Any], meta: Mapping[str, Any]) -> bool:
@@ -77,9 +88,12 @@ def _carries_stage_payload(info: Mapping[str, Any], meta: Mapping[str, Any]) ->
     Any real async-chunk payload brings producer metadata along, whether the
     transport delivers it nested under ``meta`` or as flattened ``meta.*`` keys.
     """
+    codes = info.get("codes")
+    if isinstance(codes, Mapping) and any(value is not None for value in codes.values()):
+        return True
     if any(key not in _RUNNER_STAMPED_KEYS for key in info):
         return True
-    return meta is not info and any(key not in _RUNNER_STAMPED_KEYS for key in meta)
+    return any(meta.get(key) is not None for key in _PRODUCER_META_KEYS)
 
 
 @dataclass(frozen=True)
@@ -126,6 +140,7 @@ class MiniCPMO45Code2Wav(nn.Module):
     input_modalities = "audio"
     have_multimodal_outputs = True
     enable_update_additional_information = True
+    replace_runtime_additional_information = True
     requires_raw_input_tokens = True
     requires_request_ids = True
     has_preprocess = False

vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py

@@ -260,7 +260,15 @@ def tts2code2wav_async_chunk(
         return None
 
     if native_duplex and turn_end and record.get("last_terminal_turn") == duplex_turn_key:
-        return None
+        # Emit an empty replacement snapshot so Code2Wav cannot replay the
+        # prior terminal audio when this control-only boundary arrives.
+        return OmniPayloadStruct(
+            meta=_MiniCPMO45MetaStruct(
+                is_segment_finished=torch.tensor(True, dtype=torch.bool),
+                replace_runtime_additional_information=True,
+            ),
+            request_id=request_id,
+        )
 
     state = container.get(_MINICPMO45_ASYNC_STATE)
     if not isinstance(state, dict):
@@ -373,6 +381,7 @@ def tts2code2wav_async_chunk(
             llm_output_text_utf8=segment_text_utf8,
             tts_is_last_chunk=flush_pending,
             turn_end=turn_end and last_chunk,
+            replace_runtime_additional_information=True,
             ref_audio_sr=ref_audio_sr,
         ),
         request_id=request_id,
@@ -432,6 +441,7 @@ def tts2code2wav_full_payload(
             finished=finished,
             req_id=[request_id],
             ref_audio_sr=_coerce_int(meta_info.get("ref_audio_sr")),
+            replace_runtime_additional_information=True,
             native_duplex_segment_text=(
                 str(meta_info["native_duplex_segment_text"])
                 if isinstance(meta_info.get("native_duplex_segment_text"), str)

vllm_omni/worker/gpu_model_runner.py

@@ -1494,14 +1494,18 @@ class OmniGPUModelRunner(GPUModelRunner):
                 self.inputs_embeds.gpu[start_offset : start_offset + overlay_len].copy_(src)
 
     def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> None:
+        replace = getattr(self.model, "replace_runtime_additional_information", False)
+        update_buffer = self._replace_intermediate_buffer if replace else self._update_intermediate_buffer
         for new_req in scheduler_output.scheduled_new_reqs:
             model_buffer = getattr(new_req, "model_intermediate_buffer", None)
             if isinstance(model_buffer, dict) and model_buffer:
-                self._update_intermediate_buffer(new_req.req_id, model_buffer)
+                update_buffer(new_req.req_id, model_buffer)
+                if replace:
+                    continue
             payload_info = getattr(new_req, "additional_information", None)
             decoded_info = deserialize_additional_information(payload_info)
             if decoded_info:
-                self._update_intermediate_buffer(new_req.req_id, decoded_info)
+                update_buffer(new_req.req_id, decoded_info)
 
         if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"):
             cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {})
@@ -1509,7 +1513,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 for req_id, req_infos in cached_infos.items():
                     decoded_info = deserialize_additional_information(req_infos)
                     if decoded_info:
-                        self._update_intermediate_buffer(req_id, decoded_info)
+                        update_buffer(req_id, decoded_info)
 
     def _maybe_attach_mimo_audio_req_infos(
         self,
@@ -2039,6 +2043,27 @@ class OmniGPUModelRunner(GPUModelRunner):
         # Backward compatible: mirror to old setattr location
         setattr(req_state, "additional_information_cpu", existing)
 
+    def _replace_intermediate_buffer(self, req_id: str, upd: dict) -> None:
+        """Replace one request's runtime payload with the current chunk."""
+        if not isinstance(upd, dict):
+            return
+        req_state = self.requests.get(req_id)
+        if req_state is None:
+            return
+        snapshot = dict(upd)
+        previous = self.model_intermediate_buffer.get(req_id)
+        previous_meta = previous.get("meta") if isinstance(previous, dict) else None
+        incoming_meta = snapshot.get("meta")
+        meta = dict(incoming_meta) if isinstance(incoming_meta, dict) else {}
+        if isinstance(previous_meta, dict):
+            for key in ("num_processed_tokens", "resumable"):
+                if key not in meta and key in previous_meta:
+                    meta[key] = previous_meta[key]
+        if meta or "meta" in snapshot:
+            snapshot["meta"] = meta
+        self.model_intermediate_buffer[req_id] = snapshot
+        setattr(req_state, "additional_information_cpu", snapshot)
+
     def _update_streaming_input_additional_info(self, new_req_data, req_id):
         # For streaming input prefill case only. Update buffer from last segment input.
         cached_additional_info = self.model_intermediate_buffer.get(req_id, {})
@@ -2046,6 +2071,15 @@ class OmniGPUModelRunner(GPUModelRunner):
         if not isinstance(inc_info, dict) or not inc_info:
             payload_info = getattr(new_req_data, "additional_information", None)
             inc_info = deserialize_additional_information(payload_info)
+        if getattr(self.model, "replace_runtime_additional_information", False):
+            snapshot = dict(inc_info) if isinstance(inc_info, dict) else {}
+            meta = snapshot.get("meta")
+            meta = dict(meta) if isinstance(meta, dict) else {}
+            meta["num_processed_tokens"] = 0
+            meta["resumable"] = True
+            snapshot["meta"] = meta
+            self._replace_intermediate_buffer(req_id, snapshot)
+            return
         if not isinstance(inc_info, dict) or not inc_info:
             return
         merged_info = dict(cached_additional_info) if isinstance(cached_additional_info, dict) else {}