[BugFix][MiniCPM-o] Restore audio output without async chunk (#5455)

Signed-off-by: R2-Y <ruiruyang2@gmail.com>
Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/core/sched/test_omni_scheduling_coordinator.py

@@ -119,6 +119,7 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
             ("CovoAudioForConditionalGeneration", "code2wav"),
             ("MiMoAudioModel", "code2wav"),
             ("Qwen3TTSCode2Wav", "code2wav"),
+            ("MiniCPMO45Code2Wav", "code2wav"),
             ("CosyVoice3Model", "cosyvoice3_code2wav"),
             ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
             ("DyninOmniForConditionalGeneration", "token2image"),

tests/distributed/omni_connectors/test_adapter_and_flow.py

@@ -218,7 +218,8 @@ def test_get_connectors_for_stage():
     assert stage_2_config["from_stage_1"]["spec"]["name"] == "C2"
 
 
-def test_outgoing_only_async_stage_uses_sender_connector_spec():
+@pytest.mark.parametrize("async_chunk", [False, True])
+def test_outgoing_stage_uses_sender_connector_spec(async_chunk: bool):
     config = OmniTransferConfig(
         connectors={
             ("1", "2"): ConnectorSpec(
@@ -228,12 +229,28 @@ def test_outgoing_only_async_stage_uses_sender_connector_spec():
         }
     )
 
-    spec = get_stage_connector_spec(config, stage_id=1, async_chunk=True)
+    spec = get_stage_connector_spec(config, stage_id=1, async_chunk=async_chunk)
 
     assert spec["name"] == "SharedMemoryConnector"
     assert spec["extra"] == {"codec_chunk_frames": 25, "role": "sender"}
 
 
+def test_full_payload_consumer_uses_receiver_connector_spec():
+    config = OmniTransferConfig(
+        connectors={
+            ("1", "2"): ConnectorSpec(
+                name="SharedMemoryConnector",
+                extra={"codec_chunk_frames": 25},
+            )
+        }
+    )
+
+    spec = get_stage_connector_spec(config, stage_id=2, async_chunk=False)
+
+    assert spec["name"] == "SharedMemoryConnector"
+    assert spec["extra"] == {"codec_chunk_frames": 25, "role": "receiver"}
+
+
 def test_recv_with_missing_metadata(mocker: MockerFixture):
     """Test recv when queue payload is malformed (missing metadata)."""
     # Connector expects metadata but task doesn't have it

tests/engine/test_async_omni_engine_stage_init.py

@@ -609,6 +609,29 @@ def test_build_engine_args_stage_model_overrides_parent_model():
     assert engine_args["model"] == "/stage/model"
 
 
+def test_build_engine_args_keeps_full_payload_connector_spec():
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=1,
+        stage_type="llm",
+        engine_args={"async_chunk": False},
+        default_sampling_params={},
+    )
+    connector_spec = {
+        "name": "SharedMemoryConnector",
+        "extra": {"role": "sender"},
+    }
+
+    engine_args = build_engine_args_dict(
+        stage_cfg,
+        "/parent/model",
+        stage_connector_spec=connector_spec,
+    )
+
+    assert engine_args["stage_connector_spec"] == connector_spec
+
+
 def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path):
     from vllm_omni.engine.stage_init_utils import build_engine_args_dict
 

tests/model_executor/models/minicpmo_4_5/test_pipeline.py

@@ -6,7 +6,7 @@ Covers:
   - pipeline declared in the central registry
   - lazy loader returns the expected ``PipelineConfig``
   - split 3-stage topology with no fused compatibility registration
-  - stage 1 routes through ``llm2tts`` and the async Code2Wav producer
+  - stage 1 routes through ``llm2tts`` and mode-specific Code2Wav producers
   - ``hf_architectures`` covers both the shared ``MiniCPMO`` alias and the
     explicit 4.5 arch
   - ``hf_config_predicate`` selects MiniCPM-o 4.5 only and rejects 2.6
@@ -105,7 +105,9 @@ class TestPipelineTopology:
         assert talker.engine_output_type == "latent"
         # scope KV cache / mrope sizing to talker sub-config
         assert talker.hf_config_name == "tts_config"
-        assert talker.custom_process_next_stage_input_func is None
+        assert talker.custom_process_next_stage_input_func == (
+            "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_full_payload"
+        )
         assert talker.async_chunk_process_next_stage_input_func == (
             "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_async_chunk"
         )
@@ -130,7 +132,9 @@ class TestPipelineTopology:
         assert code2wav.final_output_type == "audio"
         assert code2wav.engine_output_type == "audio"
         assert code2wav.model_arch == "MiniCPMO45Code2Wav"
-        assert code2wav.sync_process_input_func is None
+        assert code2wav.sync_process_input_func == (
+            "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_token_only"
+        )
 
 
 class TestDeployTopology:
@@ -156,6 +160,8 @@ class TestDeployTopology:
         assert connector["extra"]["codec_left_context_frames"] == 3
         assert connector["extra"]["connector_get_max_wait_first_chunk"] == 3000
         assert connector["extra"]["connector_get_max_wait"] == 300
+        expected_processor = "tts2code2wav_async_chunk" if deploy.async_chunk else "tts2code2wav_full_payload"
+        assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(expected_processor)
         if filename == "minicpmo_4_5.yaml":
             assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 4, 4]
             assert [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] == [
@@ -171,14 +177,21 @@ class TestDeployTopology:
                 0.35,
             ]
 
-    def test_pipeline_exposes_no_full_payload_or_token_placeholder_hooks(self) -> None:
+    def test_pipeline_exposes_full_and_async_payload_hooks(self) -> None:
         pipeline = OMNI_PIPELINES[_PIPELINE_KEY]
         talker = pipeline.get_stage(1)
         code2wav = pipeline.get_stage(2)
         assert talker is not None
         assert code2wav is not None
-        assert talker.custom_process_next_stage_input_func is None
-        assert code2wav.sync_process_input_func is None
+        assert talker.custom_process_next_stage_input_func == (
+            "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_full_payload"
+        )
+        assert talker.async_chunk_process_next_stage_input_func == (
+            "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_async_chunk"
+        )
+        assert code2wav.sync_process_input_func == (
+            "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_token_only"
+        )
 
 
 def test_code2wav_model_is_lazily_registered() -> None:

tests/model_executor/stage_input_processors/test_minicpmo_4_5_async_chunk.py

@@ -10,6 +10,8 @@ from vllm.v1.request import RequestStatus
 
 from vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni import (
     tts2code2wav_async_chunk,
+    tts2code2wav_full_payload,
+    tts2code2wav_token_only,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -132,6 +134,67 @@ def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> None:
     assert payload.meta.turn_end is True
 
 
+def test_full_payload_forwards_all_codes_and_request_metadata() -> None:
+    manager = _manager()
+    request = _request("req")
+    request.additional_information = {
+        "codes": {"ref": [0.1, -0.1]},
+        "meta": {
+            "ref_audio_sr": 16000,
+            "native_duplex_segment_text": "hello",
+            "segment_end": True,
+            "turn_end": True,
+        },
+        "duplex": {"epoch": 3, "model_turn_id": 7},
+    }
+
+    payload = tts2code2wav_full_payload(
+        transfer_manager=manager,
+        pooling_output={
+            "codes.audio": torch.arange(7, dtype=torch.long).reshape(-1, 1),
+            "meta.finished": torch.tensor(True),
+        },
+        request=request,
+    )
+
+    assert _codes(payload) == [4218, 4218, 4218, *range(7)]
+    assert payload.codes.ref.tolist() == pytest.approx([0.1, -0.1])
+    assert payload.meta.request_id == "req"
+    assert payload.meta.chunk_seq == 0
+    assert payload.meta.code_flat_numel == 10
+    assert payload.meta.codec_chunk_frames == 7
+    assert payload.meta.codec_left_context_frames == 3
+    assert payload.meta.left_context_size == 3
+    assert payload.meta.last_chunk is True
+    assert payload.meta.finished.item() is True
+    assert payload.meta.ref_audio_sr == 16000
+    assert payload.meta.native_duplex_segment_text == "hello"
+    assert payload.meta.duplex_epoch == 3
+    assert payload.meta.duplex_turn_id == 7
+    assert payload.meta.segment_end is True
+    assert payload.meta.turn_end is True
+
+
+def test_sync_token_only_reserves_codec_and_silence_slots() -> None:
+    output = SimpleNamespace(
+        finished=True,
+        outputs=[
+            SimpleNamespace(
+                multimodal_output={
+                    "codes.audio": torch.arange(7, dtype=torch.long).reshape(-1, 1),
+                    "meta.finished": torch.tensor(True),
+                }
+            )
+        ],
+    )
+
+    prompts = tts2code2wav_token_only([output])
+
+    assert len(prompts) == 1
+    assert prompts[0]["prompt_token_ids"] == [0] * 10
+    assert prompts[0]["additional_information"] is None
+
+
 def test_empty_final_releases_wait_gate_once() -> None:
     manager = _manager()
     request = _request("req")

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -47,6 +47,8 @@ _FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
         # qwen3_tts: Qwen3TTSTalkerForConditionalGeneration (Stage 0)
         # -> Qwen3TTSCode2Wav (Stage 1).  Stage 1 is the consumer.
         ("Qwen3TTSCode2Wav", "code2wav"),
+        # minicpmo_4_5: Talker (Stage 1) -> Code2Wav (Stage 2).
+        ("MiniCPMO45Code2Wav", "code2wav"),
         # cosyvoice3: cosyvoice3_talker (Stage 0) -> cosyvoice3_code2wav (Stage 1).
         ("CosyVoice3Model", "cosyvoice3_code2wav"),
         # indextts2: indextts2_talker (Stage 0) -> indextts2_s2mel_decoder

vllm_omni/engine/stage_init_utils.py

@@ -713,7 +713,7 @@ def build_engine_args_dict(
     # Stage id must come from stage config instead of inherited CLI kwargs
     # (e.g. `--stage-id` defaulting to None).
     engine_args_dict["stage_id"] = stage_id
-    if engine_args_dict.get("async_chunk", False):
+    if stage_connector_spec:
         engine_args_dict["stage_connector_spec"] = dict(stage_connector_spec or {})
 
     if stage_type == "diffusion":
@@ -1054,19 +1054,17 @@ def get_stage_connector_spec(
     stage_id: int,
     async_chunk: bool,
 ) -> dict[str, Any]:
-    """Return the first connector spec for the stage when async chunking is enabled."""
+    """Return the first connector spec for a stage data-plane edge."""
     from vllm_omni.distributed.omni_connectors import get_stage_connector_config
 
-    if not async_chunk:
-        return {}
-
     stage_connectors_cfg = get_stage_connector_config(omni_transfer_config, stage_id)
     for cfg in stage_connectors_cfg.values():
         return dict(cfg.get("spec", {}))
 
-    # A stage can be an async producer without consuming chunks itself. Keep
-    # its connector for save_async(), but mark it sender-only so the scheduler
-    # does not park normal orchestrator-provided inputs waiting for a chunk.
+    # A producer does not consume connector data itself. Keep its connector
+    # for both async-chunk and terminal full-payload sends, but mark it
+    # sender-only so the scheduler does not park orchestrator-provided inputs
+    # waiting for an upstream payload.
     target_stage = str(stage_id)
     for (from_stage, _to_stage), spec in getattr(omni_transfer_config, "connectors", {}).items():
         if from_stage == target_stage:

vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py

@@ -60,6 +60,7 @@ MINICPMO_4_5_PIPELINE = PipelineConfig(
             hf_config_name="tts_config",
             engine_output_type="latent",
             custom_process_input_func=f"{_PROC}.llm2tts",
+            custom_process_next_stage_input_func=f"{_PROC}.tts2code2wav_full_payload",
             async_chunk_process_next_stage_input_func=f"{_PROC}.tts2code2wav_async_chunk",
             sampling_constraints={"detokenize": False},
         ),
@@ -72,6 +73,7 @@ MINICPMO_4_5_PIPELINE = PipelineConfig(
             final_output_type="audio",
             engine_output_type="audio",
             model_arch="MiniCPMO45Code2Wav",
+            sync_process_input_func=f"{_PROC}.tts2code2wav_token_only",
             sampling_constraints={"detokenize": True},
         ),
     ),

vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py

@@ -167,7 +167,7 @@ def _extract_codec_delta(pooling_output: Any, request_id: str) -> list[int]:
         if routed_id is not None and str(routed_id) != request_id:
             return []
         codes = pooling_output.get("codes")
-        audio = codes.get("audio") if isinstance(codes, Mapping) else None
+        audio = codes.get("audio") if isinstance(codes, Mapping) else pooling_output.get("codes.audio")
         return _codec_scalars(audio)
     if isinstance(pooling_output, Sequence) and not isinstance(
         pooling_output,
@@ -341,6 +341,103 @@ def tts2code2wav_async_chunk(
     )
 
 
+def tts2code2wav_full_payload(
+    transfer_manager: Any,
+    pooling_output: Any,
+    request: Any,
+) -> OmniPayloadStruct:
+    """Build one terminal Code2Wav payload when async chunks are disabled."""
+    external_id = getattr(request, "external_req_id", None)
+    internal_id = getattr(request, "request_id", None)
+    request_id = str(external_id if external_id is not None else internal_id)
+    codes = _extract_codec_delta(pooling_output, request_id)
+    _, left_context_frames = _codec_config(transfer_manager)
+    context = [_MINICPMO45_SILENCE_CODE] * left_context_frames if codes else []
+    output_codes = [*context, *codes]
+
+    request_info = getattr(request, "additional_information", None)
+    if not isinstance(request_info, Mapping):
+        request_info = {}
+    codes_info = request_info.get("codes")
+    if not isinstance(codes_info, Mapping):
+        codes_info = {}
+    meta_info = request_info.get("meta")
+    if not isinstance(meta_info, Mapping):
+        meta_info = {}
+    duplex_info = request_info.get("duplex")
+    if not isinstance(duplex_info, Mapping):
+        duplex_info = {}
+
+    ref_audio = codes_info.get("ref")
+    finished = torch.tensor(True, dtype=torch.bool)
+    return OmniPayloadStruct(
+        codes=CodesStruct(
+            audio=torch.tensor(output_codes, dtype=torch.long),
+            ref=torch.as_tensor(ref_audio, dtype=torch.float32).reshape(-1) if ref_audio is not None else None,
+        ),
+        meta=_MiniCPMO45MetaStruct(
+            request_id=request_id,
+            chunk_seq=0,
+            cache_epoch=0,
+            code_flat_numel=len(output_codes),
+            codec_chunk_frames=len(codes),
+            codec_left_context_frames=len(context),
+            left_context_size=len(context),
+            last_chunk=True,
+            stream_finished=finished,
+            finished=finished,
+            req_id=[request_id],
+            ref_audio_sr=_coerce_int(meta_info.get("ref_audio_sr")),
+            native_duplex_segment_text=(
+                str(meta_info["native_duplex_segment_text"])
+                if isinstance(meta_info.get("native_duplex_segment_text"), str)
+                else None
+            ),
+            duplex_turn_id=_coerce_int(duplex_info.get("model_turn_id", duplex_info.get("turn_id"))),
+            duplex_epoch=_coerce_int(duplex_info.get("epoch")),
+            segment_end=bool(meta_info.get("segment_end", False)),
+            turn_end=bool(meta_info.get("turn_end", False)),
+            tts_is_last_chunk=True,
+        ),
+        request_id=request_id,
+    )
+
+
+def tts2code2wav_token_only(
+    source_outputs: list[Any],
+    _prompt: Any = None,
+    _requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    """Build the sync Code2Wav placeholder; codec data arrives by connector."""
+    code2wav_inputs: list[OmniTokensPrompt] = []
+    for talker_output in source_outputs:
+        if not talker_output.finished:
+            continue
+        output = talker_output.outputs[0]
+        multimodal_output = getattr(output, "multimodal_output", None)
+        codes = multimodal_output.get("codes") if isinstance(multimodal_output, Mapping) else None
+        audio = (
+            codes.get("audio")
+            if isinstance(codes, Mapping)
+            else multimodal_output.get("codes.audio")
+            if isinstance(multimodal_output, Mapping)
+            else None
+        )
+        codec_codes = _codec_scalars(audio)
+        if not codec_codes:
+            continue
+        prompt_len = len(codec_codes) + 3
+        code2wav_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * prompt_len,
+                additional_information=None,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return code2wav_inputs
+
+
 def _special_token_ids_from_mm_output(mm_output):
     if not isinstance(mm_output, Mapping):
         return {}