From 4f7c08c0eb0dd2dd47a1baa70384a46c7536afd5 Mon Sep 17 00:00:00 2001 From: NATURE Date: Fri, 21 Aug 2026 16:37:21 +0800 Subject: [PATCH] [Bugfix][MiniCPM-o] Fix async-chunk snapshot replacement and prompt cleanup (#6406) Signed-off-by: natureofnature Co-authored-by: amy-why-3459 --- .../sched/test_omni_ar_scheduler_streaming.py | 49 ++++++- .../test_chunk_transfer_adapter.py | 135 ++++++++++++++++++ .../minicpmo_4_5/test_code2wav_batching.py | 9 ++ .../test_minicpmo_4_5_async_chunk.py | 10 ++ .../npu/test_minicpmo_code2wav_npugraph.py | 23 +++ tests/worker/test_omni_gpu_model_runner.py | 74 ++++++++++ vllm_omni/core/sched/omni_ar_scheduler.py | 1 + vllm_omni/core/sched/omni_scheduler_mixin.py | 47 +++++- vllm_omni/data_entry_keys.py | 2 + .../distributed/omni_connectors/adapter.py | 12 +- .../chunk_transfer_adapter.py | 46 +++++- .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 17 ++- .../minicpmo_4_5_omni.py | 12 +- vllm_omni/worker/gpu_model_runner.py | 40 +++++- 14 files changed, 458 insertions(+), 19 deletions(-) diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index ca80323a7..50fcbcf4e 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/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 == {} diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index fcfd284bf..d0fd00376 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/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") diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py index a1ef3b2ec..9579a9aac 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py +++ b/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. {}, ], diff --git a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py index f0da872b9..7959efa1f 100644 --- a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py +++ b/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 diff --git a/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py b/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py index 9a069fb8f..4ae544380 100644 --- a/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py +++ b/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() diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index ca282994a..cd5c559d7 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/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) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 66da231a6..f24edae38 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 05ab26e06..9920e3a27 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/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() diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index 98575b689..62594bd19 100644 --- a/vllm_omni/data_entry_keys.py +++ b/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 diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index f1e224b15..3870155fe 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/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 diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 4a5abcbe5..9bcc74257 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -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: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index 686b39c5b..4ebe0b28c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index cd5860786..5bbbfbc4a 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -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) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 7b3bb2cf9..a4716d682 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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 {}