[Bugfix][MiniCPM-o] Bound full-duplex Talker context with sliding recompute (#6626)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>

tests/model_executor/models/minicpmo_4_5/test_pipeline.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for the MiniCPM-o 4.5 pipeline registration.
 
 Covers:
@@ -19,6 +19,8 @@ from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
+from transformers import PretrainedConfig
+from vllm.config import ModelConfig
 
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.stage_config import (
@@ -167,8 +169,11 @@ class TestDeployTopology:
         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 "hf_overrides" not in stages[1].yaml_engine_args
         if filename == "minicpmo_4_5.yaml":
             assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 4, 4]
+            assert stages[1].yaml_engine_args["hf_overrides"] == {"tts_config": {"attention_type": "sliding_recompute"}}
             memory_utilizations = [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages]
             assert memory_utilizations == [
                 0.55,
@@ -189,6 +194,18 @@ class TestDeployTopology:
                 0.35,
             ]
 
+    def test_duplex_talker_hf_override_resolves_nested_attention_policy(self) -> None:
+        deploy = load_deploy_config(_DEPLOY_DIR / "minicpmo_4_5.yaml")
+        stages = merge_pipeline_deploy(OMNI_PIPELINES[deploy.pipeline], deploy)
+        overrides = stages[1].yaml_engine_args["hf_overrides"]
+        hf_config = PretrainedConfig()
+        hf_config.tts_config = PretrainedConfig(attention_type="full_attention")
+
+        model_config = object.__new__(ModelConfig)
+        model_config._apply_dict_overrides(hf_config, overrides)
+
+        assert hf_config.tts_config.attention_type == "sliding_recompute"
+
     @pytest.mark.parametrize(
         "filename",
         [

tests/model_executor/models/minicpmo_4_5/test_talker_batching.py

@@ -102,6 +102,7 @@ def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration:
     talker._pending_force_eos_rows = None
     talker._penalty_histories = None
     talker._request_audio_states = {}
+    talker._request_condition_states = {}
     talker._deferred_cleanup_ids = set()
     talker._tts_config = ConditionalChatTTSConfig()
     talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)])
@@ -833,11 +834,164 @@ def test_native_duplex_condition_matches_official_text_plus_audio_bos() -> None:
     assert condition.shape[0] == token_ids.shape[0] + 1
 
 
+def test_native_duplex_rollover_matches_official_sliding_recompute(mocker) -> None:
+    talker = _make_talker()
+    talker.emb_text = nn.Embedding(1, 2)
+    talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)])
+    with torch.no_grad():
+        talker.emb_code[0].weight.copy_(torch.arange(16, dtype=torch.float32).reshape(8, 2))
+
+    previous_condition = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    current_condition = torch.tensor([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]])
+    build_condition = mocker.patch.object(
+        talker,
+        "_build_condition_embeddings",
+        side_effect=[previous_condition, current_condition, current_condition],
+    )
+
+    talker.preprocess(
+        torch.zeros(2, dtype=torch.long),
+        None,
+        _omni_is_prefill=True,
+        _omni_prompt_len=2,
+        request_id="req-window",
+        native_duplex=True,
+        meta={"turn_start": True, "streaming_condition_seq": 7},
+        tts_token_ids=torch.tensor([1]),
+        tts_hidden_states=torch.ones(1, 2),
+    )
+
+    older_codes = [3, 4, 5]
+    talker._request_condition_states["req-window"]["base_recent_codes"] = tuple(older_codes)
+    previous_codes = [1, 2]
+    recompute_meta = {
+        "streaming_condition_seq": 8,
+        "streaming_prompt_recompute": True,
+    }
+    expected = torch.cat(
+        [
+            previous_condition,
+            talker.emb_code[0](torch.tensor(previous_codes)),
+            current_condition,
+        ],
+        dim=0,
+    )
+    _, full_embeds, _ = talker.preprocess(
+        torch.zeros(expected.shape[0], dtype=torch.long),
+        None,
+        _omni_is_prefill=True,
+        _omni_prompt_len=expected.shape[0],
+        request_id="req-window",
+        native_duplex=True,
+        ids={"streaming_prompt_previous_codes": previous_codes},
+        meta=recompute_meta,
+        tts_token_ids=torch.tensor([2]),
+        tts_hidden_states=torch.ones(1, 2),
+    )
+
+    assert torch.equal(full_embeds, expected)
+    assert talker._request_audio_states["req-window"]["recent_codes"] == older_codes + previous_codes
+
+    # Chunked-prefill/preemption retries reuse the immutable recipe for the
+    # same condition instead of advancing the one-chunk window twice.
+    _, retry_slice, _ = talker.preprocess(
+        torch.zeros(3, dtype=torch.long),
+        None,
+        _omni_is_prefill=True,
+        _omni_num_computed_tokens=2,
+        _omni_prompt_len=expected.shape[0],
+        request_id="req-window",
+        native_duplex=True,
+        ids={"streaming_prompt_previous_codes": previous_codes},
+        meta=recompute_meta,
+        tts_token_ids=torch.tensor([2]),
+        tts_hidden_states=torch.ones(1, 2),
+    )
+    assert torch.equal(retry_slice, expected[2:5])
+    assert talker._request_audio_states["req-window"]["recent_codes"] == older_codes + previous_codes
+    assert build_condition.call_count == 3
+
+
+def test_native_duplex_condition_advance_requires_recompute_marker(mocker) -> None:
+    talker = _make_talker()
+    condition = torch.ones(2, 2)
+    mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition)
+    common = {
+        "_omni_is_prefill": True,
+        "_omni_prompt_len": 2,
+        "request_id": "req-history",
+        "native_duplex": True,
+        "tts_token_ids": torch.tensor([1]),
+        "tts_hidden_states": torch.ones(1, 2),
+    }
+
+    talker.preprocess(torch.zeros(2, dtype=torch.long), None, meta={"streaming_condition_seq": 0}, **common)
+    with pytest.raises(ValueError, match="advanced without its streaming recompute marker"):
+        talker.preprocess(torch.zeros(2, dtype=torch.long), None, meta={"streaming_condition_seq": 1}, **common)
+
+
+@pytest.mark.parametrize(
+    ("state", "previous_codes", "error"),
+    [
+        (None, [1], "missing the previous Talker condition"),
+        (
+            {"condition_seq": 0, "condition": torch.ones(2, 2), "base_recent_codes": ()},
+            [7],
+            "invalid or terminal token",
+        ),
+    ],
+)
+def test_native_duplex_rollover_rejects_unverifiable_history(state, previous_codes, error) -> None:
+    talker = _make_talker()
+    talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)])
+    if state is not None:
+        talker._request_condition_states["req-invalid-window"] = state
+
+    with pytest.raises(ValueError, match=error):
+        talker._build_streaming_recompute_embeddings(
+            torch.ones(3, 2),
+            request_id="req-invalid-window",
+            info_dict={"ids": {"streaming_prompt_previous_codes": previous_codes}},
+            meta={
+                "streaming_condition_seq": 1,
+                "streaming_prompt_recompute": True,
+            },
+        )
+
+
+def test_native_duplex_turn_start_resets_rollover_history() -> None:
+    talker = _make_talker()
+    talker._request_condition_states["req-turn"] = {
+        "condition_seq": 4,
+        "condition": torch.full((2, 2), 9.0),
+        "active_embeddings": torch.full((5, 2), 9.0),
+    }
+    current = torch.ones(3, 2)
+
+    result = talker._build_streaming_recompute_embeddings(
+        current,
+        request_id="req-turn",
+        info_dict={},
+        meta={"turn_start": True, "streaming_condition_seq": 5},
+    )
+
+    assert torch.equal(result, current)
+    state = talker._request_condition_states["req-turn"]
+    assert state["condition_seq"] == 5
+    assert torch.equal(state["condition"], current)
+    assert "active_embeddings" not in state
+
+
 def test_request_cleanup_evicts_decode_state() -> None:
     talker = _make_talker()
     talker._request_audio_states["req-done"] = {"finished": False}
+    talker._request_condition_states["req-done"] = {
+        "condition_seq": 1,
+        "condition": torch.ones(2, 2),
+    }
 
     talker.on_requests_finished(["req-done"])
     talker._flush_deferred_cleanup()
 
     assert "req-done" not in talker._request_audio_states
+    assert "req-done" not in talker._request_condition_states