diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py index 3685fa5a2..37794edae 100644 --- a/tests/core/sched/test_omni_scheduling_coordinator.py +++ b/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"), diff --git a/tests/distributed/omni_connectors/test_adapter_and_flow.py b/tests/distributed/omni_connectors/test_adapter_and_flow.py index 471d3108a..1cabaf3b3 100644 --- a/tests/distributed/omni_connectors/test_adapter_and_flow.py +++ b/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 diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index e98c58d74..db7d53b79 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/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 diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 968df04a9..2442cad44 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/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: 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 4f17b2c64..25124ed6f 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 @@ -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") diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index fc21be991..59df59a26 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/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 diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 9861d7b38..3b1bb93af 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/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: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index 31d02c7cc..b8fed05f1 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/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}, ), ), 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 4f052b68d..950d5434d 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 @@ -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 {}