[WIP][Bugfix][MiniCPM-o] Fix Full duplex bugs for minicpm-challenge (#5496)

Signed-off-by: Sy03 <1370724210@qq.com>

tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py

@@ -19,10 +19,12 @@ class _FakeEncoder(nn.Module):
     def __init__(self):
         super().__init__()
         self.calls: list[int] = []
+        self.last_chunk_calls: list[bool] = []
 
     def forward_chunk(self, xs, last_chunk=False, cnn_cache=None, att_cache=None):
         batch, length, _ = xs.shape
         self.calls.append(batch)
+        self.last_chunk_calls.append(last_chunk)
         old_length = 0 if att_cache is None else att_cache.shape[3]
         output = xs[:, : max(1, length - 1)]
         cnn = xs[:, :1, :].transpose(1, 2).contiguous()
@@ -156,6 +158,35 @@ def _model():
     return model, token2wav
 
 
+def test_code2wav_resolves_hf_model_id_for_assets(mocker, tmp_path):
+    resolved_root = tmp_path / "snapshot"
+    resolved_root.mkdir()
+    config = _config()
+    config.model_config.model = "openbmb/MiniCPM-o-4_5"
+    config.model_config.revision = "test-revision"
+    config.model_config.stage_connector_config["extra"].pop("prompt_wav")
+    config.load_config = SimpleNamespace(download_dir="/model-cache")
+    model = MiniCPMO45Code2Wav(vllm_config=config)
+    mock_download = mocker.patch(
+        "vllm_omni.model_executor.model_loader.weight_utils.download_weights_from_hf_specific",
+        return_value=str(resolved_root),
+    )
+
+    assert model._resolve_model_root() == resolved_root
+    assert model.model_path == str(resolved_root)
+    assert model._default_prompt_wav == str(resolved_root / "assets" / "HT_ref_audio.wav")
+    mock_download.assert_called_once_with(
+        "openbmb/MiniCPM-o-4_5",
+        "/model-cache",
+        allow_patterns=[
+            "assets/HT_ref_audio.wav",
+            "assets/token2wav/*",
+        ],
+        revision="test-revision",
+        require_all=True,
+    )
+
+
 def _info(
     request_id: str,
     chunk_seq: int,
@@ -212,6 +243,18 @@ def test_adapter_runs_true_batch_cfg_and_splits_request_caches():
     assert cache1[0, 0, 0, 0, 0].item() == 20
 
 
+def test_fade_in_out_limits_overlap_to_available_previous_audio():
+    speech = torch.arange(6, dtype=torch.float32).reshape(1, -1)
+    previous = torch.full((1, 3), 2.0)
+    window = torch.hamming_window(8, periodic=False)
+
+    actual = BatchedToken2Wav._fade_in_out(speech, previous, window)
+
+    expected = speech.clone()
+    expected[..., :3] = speech[..., :3] * window[:3] + previous * window[-3:]
+    torch.testing.assert_close(actual, expected)
+
+
 def test_estimator_cache_stack_split_round_trip_preserves_cfg_rows():
     token2wav = _FakeToken2Wav()
     adapter = BatchedToken2Wav(token2wav)
@@ -259,6 +302,166 @@ def test_model_preserves_output_slots_and_prefers_runtime_codes():
     assert token2wav.flow.encoder.calls[-1] == 2
 
 
+def test_code2wav_projects_duplex_metadata_to_final_audio_output():
+    model, token2wav = _model()
+    segment = _info("duplex", 0, [10, 11])
+    segment_text_utf8 = torch.tensor(list(b"hello"), dtype=torch.uint8)
+    segment["meta"].update(
+        {
+            "duplex_epoch": 3,
+            "duplex_turn_id": 7,
+            "llm_output_text_utf8": segment_text_utf8,
+            "tts_is_last_chunk": True,
+            "turn_end": False,
+        }
+    )
+
+    segment_output = _forward(model, [segment])
+
+    assert segment_output.multimodal_outputs["meta.turn_end"][0].item() is False
+    # A Talker unit boundary only drains pending codec tokens. The official
+    # streaming path keeps Token2wav open until the assistant turn ends.
+    assert token2wav.flow.encoder.last_chunk_calls[-1] is False
+    assert "duplex" in model._states
+
+    final = _info("duplex", 1, [12, 13], last_chunk=True)
+    final["meta"].update(segment["meta"])
+    final["meta"]["chunk_seq"] = 1
+    final["meta"]["last_chunk"] = True
+    final["meta"]["turn_end"] = True
+    output = _forward(model, [final])
+
+    payload = output.multimodal_outputs
+    assert "meta" not in payload
+    assert payload["meta.duplex_epoch"][0].item() == 3
+    assert payload["meta.duplex_turn_id"][0].item() == 7
+    torch.testing.assert_close(
+        payload["meta.llm_output_text_utf8"][0],
+        segment_text_utf8,
+    )
+    assert payload["meta.tts_is_last_chunk"][0].item() is True
+    assert payload["meta.turn_end"][0].item() is True
+    assert token2wav.flow.encoder.last_chunk_calls[-1] is True
+    assert "duplex" not in model._states
+
+
+def test_initial_empty_segment_marker_initializes_stream_without_audio():
+    model, token2wav = _model()
+    boundary = _info("duplex", 0, [])
+    boundary["meta"].update(
+        {
+            "code_flat_numel": 0,
+            "tts_is_last_chunk": True,
+            "turn_end": False,
+        }
+    )
+
+    output = _forward(model, [boundary])
+
+    assert output.multimodal_outputs["model_outputs"][0].numel() == 0
+    assert "duplex" in model._states
+    assert token2wav.hift.calls == []
+
+    resumed = _info(
+        "duplex",
+        1,
+        [4218, 4218, 4218, 10, 11, 12, 13, 14],
+    )
+    output = _forward(model, [resumed])
+
+    assert output.multimodal_outputs["model_outputs"][0].numel() > 0
+    assert "duplex" in model._states
+
+
+def test_shared_runtime_prompt_recreates_missing_file_before_second_owner(tmp_path, monkeypatch):
+    monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
+    model, _ = _model()
+    reference = torch.tensor([0.0, 0.25, -0.25, 0.0])
+
+    first = _info("voice-a", 0, [10, 11])
+    first["codes"]["ref"] = reference
+    first["meta"]["ref_audio_sr"] = 16000
+    first["meta"].pop("prompt_cache_id")
+    _forward(model, [first], request_ids=["internal-a"])
+
+    prompt_key = model._request_prompt_keys["internal-a"]
+    prompt_path = Path(model._runtime_prompts[prompt_key].path)
+    prompt_path.unlink()
+
+    second = _info("voice-b", 0, [12, 13])
+    second["codes"]["ref"] = reference
+    second["meta"]["ref_audio_sr"] = 16000
+    second["meta"].pop("prompt_cache_id")
+    _forward(model, [second], request_ids=["internal-b"])
+
+    assert prompt_path.is_file()
+    assert model._runtime_prompts[prompt_key].owners == {"internal-a", "internal-b"}
+
+    model.on_requests_finished(["internal-a"])
+    assert prompt_path.is_file()
+    assert model._runtime_prompts[prompt_key].owners == {"internal-b"}
+
+    model.on_requests_finished(["internal-b"])
+    assert not prompt_path.exists()
+    assert prompt_key not in model._runtime_prompts
+
+
+def test_runtime_prompt_write_failure_does_not_publish_partial_file(tmp_path, monkeypatch):
+    monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
+    model, _ = _model()
+    reference = torch.tensor([0.0, 0.25, -0.25, 0.0])
+
+    def fail_after_partial_write(path, *_args, **_kwargs):
+        Path(path).write_bytes(b"partial")
+        raise OSError("simulated write failure")
+
+    monkeypatch.setattr(
+        "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav.sf.write",
+        fail_after_partial_write,
+    )
+
+    with pytest.raises(OSError, match="simulated write failure"):
+        model._materialize_runtime_prompt(reference, 16000)
+
+    assert len(model._runtime_prompts) == 1
+    entry = next(iter(model._runtime_prompts.values()))
+    assert not Path(entry.path).exists()
+    assert list(Path(entry.path).parent.iterdir()) == []
+
+
+def test_runtime_prompt_files_are_isolated_between_model_instances(tmp_path, monkeypatch):
+    monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
+    first_model, _ = _model()
+    second_model, _ = _model()
+    reference = torch.tensor([0.0, 0.25, -0.25, 0.0])
+
+    def runtime_ref_info(request_id: str):
+        info = _info(request_id, 0, [10, 11])
+        info["codes"]["ref"] = reference
+        info["meta"]["ref_audio_sr"] = 16000
+        info["meta"].pop("prompt_cache_id")
+        return info
+
+    _forward(first_model, [runtime_ref_info("voice-a")], request_ids=["internal-a"])
+    _forward(second_model, [runtime_ref_info("voice-b")], request_ids=["internal-b"])
+
+    first_key = first_model._request_prompt_keys["internal-a"]
+    second_key = second_model._request_prompt_keys["internal-b"]
+    first_path = Path(first_model._runtime_prompts[first_key].path)
+    second_path = Path(second_model._runtime_prompts[second_key].path)
+    assert first_key == second_key
+    assert first_path != second_path
+    assert first_path.is_file()
+    assert second_path.is_file()
+
+    first_model.on_requests_finished(["internal-a"])
+    assert not first_path.exists()
+    assert second_path.is_file()
+
+    second_model.on_requests_finished(["internal-b"])
+    assert not second_path.exists()
+
+
 def test_mixed_final_exact_buckets_keep_order_and_release_only_final_states():
     model, _ = _model()
     _forward(
@@ -485,27 +688,37 @@ def test_reference_voice_and_duplex_metadata_follow_request_lifecycle():
     model, _ = _model()
     first = _info("voice-a", 0, [1, 2])
     first["codes"]["ref"] = torch.linspace(-0.1, 0.1, 160)
+    segment_text_utf8 = torch.tensor(list(b"hello"), dtype=torch.uint8)
     first["meta"].update(
         ref_audio_sr=16000,
-        native_duplex_segment_text="hello",
+        llm_output_text_utf8=segment_text_utf8,
         duplex_turn_id=7,
         duplex_epoch=3,
     )
     first["meta"].pop("prompt_cache_id")
 
     output = _forward(model, [first])
-    prompt_cache_id, prompt_wav = model._owned_prompt_wavs["voice-a"]
-    assert prompt_cache_id.startswith("voice-a:")
+    prompt_key = model._request_prompt_keys["voice-a"]
+    prompt = model._runtime_prompts[prompt_key]
+    prompt_cache_id, prompt_wav = prompt.cache_id, prompt.path
+    assert prompt_cache_id.startswith("runtime-ref-")
     assert Path(prompt_wav).is_file()
-    assert bytes(output.multimodal_outputs["meta.llm_output_text_utf8"][0].tolist()).decode() == "hello"
+    torch.testing.assert_close(
+        output.multimodal_outputs["meta.llm_output_text_utf8"][0],
+        segment_text_utf8,
+    )
     assert output.multimodal_outputs["meta.duplex_turn_id"][0].item() == 7
     assert output.multimodal_outputs["meta.duplex_epoch"][0].item() == 3
 
     final = _info("voice-a", 1, [3, 4], last_chunk=True)
     final["meta"].pop("prompt_cache_id")
+    final["meta"]["tts_is_last_chunk"] = True
     output = _forward(model, [final])
 
     assert output.multimodal_outputs["meta.tts_is_last_chunk"][0].item() is True
-    assert "voice-a" not in model._owned_prompt_wavs
+    assert model._request_prompt_keys["voice-a"] == prompt_key
+    model.on_requests_finished(["voice-a"])
+    assert "voice-a" not in model._request_prompt_keys
+    assert prompt_key not in model._runtime_prompts
     assert not Path(prompt_wav).exists()
     assert (prompt_cache_id, prompt_wav) not in model.backend._prompt_features

tests/model_executor/models/minicpmo_4_5/test_talker_batching.py

@@ -14,6 +14,7 @@ from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
 from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
     MiniCPMO45OmniTTSForConditionalGeneration,
     _max_audio_tokens,
+    _restore_weight_norm_weight,
 )
 from vllm_omni.utils.mm_outputs import to_payload_element
 
@@ -56,6 +57,8 @@ def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration:
     talker._request_generators = {}
     talker._request_audio_states = {}
     talker._deferred_cleanup_ids = set()
+    talker._codec_min_tokens = 50
+    talker._codec_seed = 42
     return talker
 
 
@@ -81,6 +84,23 @@ def test_audio_token_limit_scales_with_condition_length(
     assert _max_audio_tokens(condition_tokens) == expected
 
 
+def test_weight_norm_restore_matches_checkpoint_parametrization_in_bfloat16() -> None:
+    generator = torch.Generator().manual_seed(42)
+    weight_v = torch.randn(8, 16, generator=generator, dtype=torch.bfloat16)
+    weight_g = torch.rand(8, 1, generator=generator, dtype=torch.bfloat16)
+    linear = nn.utils.parametrizations.weight_norm(
+        nn.Linear(16, 8, bias=False, dtype=torch.bfloat16),
+        dim=0,
+    )
+    with torch.no_grad():
+        linear.parametrizations.weight.original0.copy_(weight_g)
+        linear.parametrizations.weight.original1.copy_(weight_v)
+
+    restored = _restore_weight_norm_weight(weight_g, weight_v)
+
+    assert torch.equal(restored, linear.weight)
+
+
 def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> None:
     talker = _make_talker()
     seen: list[tuple[str, list[float], list[int]]] = []
@@ -114,9 +134,73 @@ def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> N
     assert _routed(output, 0)["codes"]["audio"].tolist() == [[2]]
     assert _routed(output, 1)["codes"]["audio"].tolist() == [[3]]
     assert _routed(output, 0)["meta"]["finished"].item() is False
+    assert set(output.multimodal_outputs["meta"]) == {"finished"}
     assert talker.compute_logits(output.text_hidden_states).argmax(dim=-1).tolist() == [0, 0]
 
 
+def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None:
+    talker = _make_talker()
+    mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2))
+    infos = [
+        {
+            "request_id": "req-a",
+            "native_duplex": True,
+            "duplex": {"epoch": 3, "turn_id": 7},
+            "ids": {"tts": [41]},
+            "meta": {
+                "native_duplex_segment_text": "first",
+                "turn_eos_token_id": 99,
+            },
+            "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)},
+        },
+        {
+            "request_id": "req-b",
+            "native_duplex": True,
+            "duplex": {"epoch": 4, "turn_id": 8},
+            "ids": {"tts": [42, 99]},
+            "meta": {
+                "native_duplex_segment_text": "second",
+                "turn_eos_token_id": 99,
+            },
+            "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)},
+        },
+    ]
+
+    output = talker.make_omni_output(
+        torch.ones(2, 2),
+        model_intermediate_buffer=infos,
+        request_token_spans=[(0, 1), (1, 2)],
+    )
+
+    meta = output.multimodal_outputs["meta"]
+    assert [value.item() for value in meta["native_duplex"]] == [True, True]
+    assert [value.item() for value in meta["duplex_epoch"]] == [3, 4]
+    assert [value.item() for value in meta["duplex_turn_id"]] == [7, 8]
+    assert "native_duplex_segment_text" not in meta
+    assert [bytes(value.tolist()).decode("utf-8") for value in meta["llm_output_text_utf8"]] == [
+        "first",
+        "second",
+    ]
+    assert [value.item() for value in meta["turn_end"]] == [False, True]
+
+
+def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None:
+    talker = _make_talker()
+    mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2))
+    info = {
+        "request_id": "req-missing-fence",
+        "native_duplex": True,
+        "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)},
+    }
+
+    with pytest.raises(RuntimeError, match="requires non-negative integer epoch and turn_id"):
+        talker.make_omni_output(
+            torch.ones(1, 2),
+            model_intermediate_buffer=[info],
+            request_token_spans=[(0, 1)],
+        )
+
+
 def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) -> None:
     talker = _make_talker()
     sample = mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2))
@@ -179,7 +263,7 @@ def test_eos_is_terminal_once_and_never_enters_codec_history(mocker) -> None:
     assert talker.compute_logits(second.text_hidden_states).argmax(dim=-1).tolist() == [1]
 
 
-def test_max_token_terminal_includes_only_new_codec_delta(mocker) -> None:
+def test_max_token_terminal_drops_unconsumed_codec_delta(mocker) -> None:
     talker = _make_talker()
     mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(3))
     info = {
@@ -194,8 +278,10 @@ def test_max_token_terminal_includes_only_new_codec_delta(mocker) -> None:
         request_token_spans=[(0, 1)],
     )
 
-    assert info["audio_codes"]["accumulated"].tolist() == [4, 5, 3]
-    assert output.multimodal_outputs["codes"]["audio"][0].tolist() == [[3]]
+    # MiniCPMTTS.generate_chunk samples once at the max-token boundary to
+    # advance RNG state, but the sampled code is not fed into KV or returned.
+    assert info["audio_codes"]["accumulated"].tolist() == [4, 5]
+    assert output.multimodal_outputs["codes"]["audio"][0].shape == (0, 1)
     assert output.multimodal_outputs["meta"]["finished"][0].item() is True
     assert talker.compute_logits(output.text_hidden_states).argmax(dim=-1).tolist() == [1]
 
@@ -268,6 +354,94 @@ def test_empty_speech_segment_finishes_without_sampling_codes() -> None:
     assert decode_embeds.shape == (1, 4)
 
 
+def test_chunked_prefill_tail_aligns_condition_with_prompt_length(mocker) -> None:
+    talker = _make_talker()
+    talker.emb_text = nn.Embedding(1, 2)
+    condition = torch.arange(18, dtype=torch.float32).reshape(9, 2)
+    mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition)
+
+    _, embeds, _ = talker.preprocess(
+        torch.zeros(9, dtype=torch.long),
+        None,
+        _omni_is_prefill=True,
+        _omni_num_computed_tokens=59,
+        _omni_prompt_len=68,
+        request_id="req-chunked-prefill",
+        tts_token_ids=torch.tensor([1]),
+        tts_hidden_states=torch.ones(1, 2),
+    )
+
+    assert torch.equal(embeds, condition)
+    state = talker._request_audio_states["req-chunked-prefill"]
+    assert state["min_tokens"] == 50
+    assert state["max_tokens"] == 64
+
+
+@pytest.mark.parametrize(
+    ("meta", "expected_min_tokens"),
+    [
+        ({"turn_start": True}, 0),
+        ({}, 26),
+        ({"turn_end": True}, 0),
+    ],
+)
+def test_native_duplex_prefill_uses_official_chunk_limits(
+    mocker,
+    meta,
+    expected_min_tokens,
+) -> None:
+    talker = _make_talker()
+    talker.emb_text = nn.Embedding(1, 2)
+    mocker.patch.object(
+        talker,
+        "_build_condition_embeddings",
+        return_value=torch.ones(3, 2),
+    )
+
+    talker.preprocess(
+        torch.zeros(3, dtype=torch.long),
+        None,
+        _omni_is_prefill=True,
+        request_id="req-duplex-chunk",
+        native_duplex=True,
+        meta=meta,
+        tts_token_ids=torch.tensor([1]),
+        tts_hidden_states=torch.ones(1, 2),
+    )
+
+    state = talker._request_audio_states["req-duplex-chunk"]
+    assert state["min_tokens"] == expected_min_tokens
+    assert state["max_tokens"] == 26
+
+
+def test_native_duplex_condition_matches_official_text_plus_audio_bos() -> None:
+    talker = _make_talker()
+    talker.emb_text = nn.Embedding(16, 2)
+    talker.projector_semantic = nn.Identity()
+    talker._normalize = False
+    talker._text_eos_id = 14
+    talker._tts_bos_id = 15
+    with torch.no_grad():
+        talker.emb_text.weight.copy_(torch.arange(32, dtype=torch.float32).reshape(16, 2))
+
+    token_ids = torch.tensor([2, 3])
+    hidden_states = torch.tensor([[0.5, 1.0], [1.5, 2.0]])
+
+    condition = talker._build_condition_embeddings(
+        token_ids,
+        hidden_states,
+        native_duplex=True,
+    )
+
+    expected_text = talker.emb_text(token_ids) + hidden_states
+    expected = torch.cat(
+        [expected_text, talker.emb_text(torch.tensor([talker._tts_bos_id]))],
+        dim=0,
+    )
+    assert torch.equal(condition, expected)
+    assert condition.shape[0] == token_ids.shape[0] + 1
+
+
 def test_request_cleanup_evicts_ar_rng_and_decode_state() -> None:
     talker = _make_talker()
     talker._request_generators["req-done"] = torch.Generator()