anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[perf] Minicpm o4.5 optimization (#5228)

Browse Source 
Signed-off-by: R2-Y <ruiruyang2@gmail.com>
Co-authored-by: Sy03 <1370724210@qq.com>
This commit is contained in:
rein yangandSy03 authored and GitHub committed 2026-07-24 17:20:39 +08:00
1 parent acf941da4d
commit ad3a162eb8
54 files changed
+3861 -3997
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/worker/test_native_duplex_hooks.py
-363
View File
@@ -1,8 +1,6 @@
from __future__ import annotations
import inspect
import os
import threading
from types import SimpleNamespace
import numpy as np
@@ -544,220 +542,6 @@ def test_minicpmo_stage0_rejects_invalid_resolved_ref_audio():
)
def test_minicpmo_tts_native_duplex_exports_segment_text_not_accumulated_condition_text():
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
MiniCPMO45OmniForConditionalGeneration,
)
class _Talker:
_ar_last_chunk_flags = [False]
_ar_last_emitted_text = "你好，你有什莫想聊的吗？"
def __call__(self, **kwargs):
del kwargs
return None, torch.zeros(8, dtype=torch.float32)
model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
model.model_stage = "tts"
model.config = SimpleNamespace(hidden_size=4)
model.talker = _Talker()
output = model.forward(
input_ids=torch.zeros(1, dtype=torch.long),
positions=torch.zeros(1, dtype=torch.long),
runtime_additional_information=[
{
"native_duplex": True,
"llm_output_text": ["你好，你有什莫想聊的吗？你好，你有什莫想聊的吗？"],
"meta": {
"native_duplex_segment_text": "你好，你有什莫想聊的吗？",
},
}
],
)
text_bytes = output.multimodal_outputs["meta.llm_output_text_utf8"].detach().cpu().tolist()
assert bytes(text_bytes).decode("utf-8") == "你好，你有什莫想聊的吗？"
assert int(output.multimodal_outputs["meta.audio_text_total_chars"].item()) == len("你好，你有什莫想聊的吗？")
def test_minicpmo_tts_native_duplex_exports_model_turn_end_metadata():
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
MiniCPMO45OmniForConditionalGeneration,
)
class _Talker:
_ar_last_chunk_flags = [True]
_ar_turn_end_flags = [True]
def __call__(self, **kwargs):
del kwargs
return None, torch.zeros(0, dtype=torch.float32)
model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
model.model_stage = "tts"
model.config = SimpleNamespace(hidden_size=4)
model.talker = _Talker()
output = model.forward(
input_ids=torch.zeros(1, dtype=torch.long),
positions=torch.zeros(1, dtype=torch.long),
runtime_additional_information=[
{
"native_duplex": True,
"meta": {"native_duplex_segment_text": ""},
}
],
)
assert int(output.multimodal_outputs["meta.turn_end"].item()) == 1
def test_minicpmo_tts_routes_batched_requests_and_terminal_flags_independently():
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
MiniCPMO45OmniForConditionalGeneration,
)
class _Talker:
def __init__(self):
self.calls = []
self._ar_last_chunk_flags = [True]
self._ar_turn_end_flags = [False]
self._ar_last_emitted_text = ""
def __call__(self, **kwargs):
info = kwargs["additional_information"]
request_id = info["global_request_id"][0]
self.calls.append(request_id)
is_last = request_id == "req-b"
self._ar_last_chunk_flags = [is_last]
self._ar_turn_end_flags = [is_last]
self._ar_last_emitted_text = request_id
value = 1.0 if request_id == "req-a" else 2.0
return None, torch.full((2,), value, dtype=torch.float32)
model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
model.model_stage = "tts"
model.config = SimpleNamespace(hidden_size=4)
model.talker = _Talker()
output = model.forward(
input_ids=torch.zeros(3, dtype=torch.long),
positions=torch.zeros(3, dtype=torch.long),
runtime_additional_information=[
{
"global_request_id": ["req-a"],
"native_duplex": True,
"duplex": {"epoch": 3, "turn_id": 7},
"meta": {"native_duplex_segment_text": "first"},
},
{
"global_request_id": ["req-idle"],
"native_duplex": True,
"duplex": {"epoch": 3, "turn_id": 9},
"meta": {"native_duplex_segment_text": "idle"},
},
{
"global_request_id": ["req-b"],
"native_duplex": True,
"duplex": {"epoch": 4, "turn_id": 8},
"meta": {"native_duplex_segment_text": "second"},
},
],
request_token_spans=[(0, 2), (2, 2), (2, 3)],
)
assert model.talker.calls == ["req-a", "req-b"]
assert model.talker._ar_last_chunk_flags == [False, False, True]
assert model.talker._ar_turn_end_flags == [False, False, True]
assert output.multimodal_outputs["meta.req_id"] == ["req-a", "req-b"]
assert output.multimodal_outputs["meta.sparse_audio"] == ["1"]
assert [waveform.tolist() for waveform in output.multimodal_outputs["model_outputs"]] == [
[1.0, 1.0],
[2.0, 2.0],
]
assert [int(value.item()) for value in output.multimodal_outputs["meta.tts_is_last_chunk"]] == [0, 1]
assert [int(value.item()) for value in output.multimodal_outputs["meta.turn_end"]] == [0, 1]
assert [int(value.item()) for value in output.multimodal_outputs["meta.duplex_epoch"]] == [3, 4]
assert [int(value.item()) for value in output.multimodal_outputs["meta.duplex_turn_id"]] == [7, 8]
@pytest.mark.parametrize(
("runtime_info", "request_token_spans", "expected_rows", "expected_ids"),
[
(
[
{
"global_request_id": ["req-a"],
"native_duplex": True,
"duplex": {"epoch": 3, "turn_id": 7},
"meta": {"native_duplex_segment_text": "first"},
}
],
[(0, 3)],
[3],
["req-a"],
),
(
[
{
"global_request_id": ["req-a"],
"native_duplex": True,
"duplex": {"epoch": 3, "turn_id": 7},
"meta": {"native_duplex_segment_text": "first"},
},
{
"global_request_id": ["req-b"],
"native_duplex": True,
"duplex": {"epoch": 4, "turn_id": 8},
"meta": {"native_duplex_segment_text": "second"},
},
],
[(0, 2), (2, 3)],
[2, 1],
["req-a", "req-b"],
),
],
ids=["single-request", "batched-requests"],
)
def test_minicpmo_tts_ignores_cuda_graph_padding_after_request_spans(
runtime_info,
request_token_spans,
expected_rows,
expected_ids,
):
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
MiniCPMO45OmniForConditionalGeneration,
)
class _Talker:
def __init__(self):
self.input_rows = []
self._ar_last_chunk_flags = [True]
self._ar_turn_end_flags = [False]
self._ar_last_emitted_text = ""
def __call__(self, **kwargs):
self.input_rows.append(int(kwargs["input_ids"].shape[0]))
return None, torch.ones(2, dtype=torch.float32)
model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration)
model.model_stage = "tts"
model.config = SimpleNamespace(hidden_size=4)
model.talker = _Talker()
output = model.forward(
input_ids=torch.zeros(4, dtype=torch.long),
positions=torch.zeros(4, dtype=torch.long),
runtime_additional_information=runtime_info,
request_token_spans=request_token_spans,
)
assert model.talker.input_rows == expected_rows
assert output.text_hidden_states.shape == (4, 4)
assert output.multimodal_outputs["meta.req_id"] == expected_ids
def test_minicpmo_stage0_special_token_ids_are_tokenizer_derived():
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
MiniCPMO45Stage0DuplexRuntime,
@@ -2427,150 +2211,3 @@ def test_minicpmo_stage0_session_context_includes_resolved_ref_audio():
assert state.context_token_ids == [1, 2, 3, 151683, 151683, 4, 5]
assert len(state.context_embeds) == 6
def test_minicpmo_tts_ref_audio_lru_evicts_matching_vocoder_base_cache(tmp_path):
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
MiniCPMO45OmniTTSForConditionalGeneration,
MiniCPMO45TTSRuntimeConfig,
)
model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration)
model._runtime_config = MiniCPMO45TTSRuntimeConfig(ref_audio_file_cache_size=1)
model._t2w_base_caches = {}
model._token2wav_state_lock = threading.RLock()
first_path = model._write_ref_audio_prompt_wav(np.array([0.1, -0.1], dtype=np.float32), 16_000)
assert first_path is not None
model._t2w_base_caches[first_path] = (torch.tensor([1]), {"cache": torch.tensor([2])})
second_path = model._write_ref_audio_prompt_wav(np.array([0.2, -0.2], dtype=np.float32), 16_000)
assert second_path is not None
assert second_path != first_path
assert first_path not in model._t2w_base_caches
assert not os.path.exists(first_path)
def test_minicpmo_tts_prompt_wav_path_does_not_default_to_model_asset(tmp_path):
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
MiniCPMO45OmniTTSForConditionalGeneration,
)
model_dir = tmp_path / "MiniCPM-o-4_5"
assets_dir = model_dir / "assets"
assets_dir.mkdir(parents=True)
(assets_dir / "HT_ref_audio.wav").write_bytes(b"placeholder")
model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration)
model.vllm_config = SimpleNamespace(model_config=SimpleNamespace(model=str(model_dir)))
model._write_ref_audio_prompt_wav = lambda ref_audio, ref_audio_sr: None
assert model._resolve_prompt_wav_path(None, None) == (None, None)
def test_minicpmo_tts_stream_turn_id_prefers_duplex_identity():
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
MiniCPMO45OmniTTSForConditionalGeneration,
)
model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration)
assert (
model._stream_turn_id(
{
"duplex": {"turn_id": 0, "model_turn_id": 4},
"meta": {"turn_id": 3},
}
)
== 4
)
assert (
model._stream_turn_id(
{
"duplex": {"turn_id": 0},
"meta": {"turn_id": 3},
}
)
== 0
)
assert model._stream_turn_id({"meta": {"turn_id": 3}, "turn_id": 2}) == 3
assert model._stream_turn_id({"turn_id": 2}) == 2
def test_minicpmo_tts_no_ref_audio_initializes_empty_vocoder_cache():
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
MiniCPMO45OmniTTSForConditionalGeneration,
_TalkerTurnState,
)
class AudioTokenizer:
def set_stream_cache(self, prompt_wav_path):
raise AssertionError(f"unexpected reference-audio cache load: {prompt_wav_path!r}")
model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration)
model._t2w_base_caches = {}
model._token2wav_state_lock = threading.RLock()
model.audio_tokenizer = AudioTokenizer()
state = _TalkerTurnState(None, None)
model._begin_turn_vocoder_cache(None, state=state)
assert state.stream_cache is None
assert state.hift_cache_dict == {}
assert state.vocoder_initialized is True
assert model._t2w_base_caches[""] == (None, {})
def test_minicpmo_tts_native_duplex_merges_and_left_pads_nonfinal_unit_audio():
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
_native_duplex_unit_waveform,
)
waveform = _native_duplex_unit_waveform(
[torch.tensor([1.0, 2.0]), torch.tensor([3.0])],
turn_end=False,
target_samples=6,
)
assert waveform is not None
assert waveform.tolist() == [0.0, 0.0, 0.0, 1.0, 2.0, 3.0]
def test_minicpmo_tts_native_duplex_does_not_pad_final_unit_audio():
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
_native_duplex_unit_waveform,
)
waveform = _native_duplex_unit_waveform(
[torch.tensor([1.0, 2.0]), torch.tensor([3.0])],
turn_end=True,
target_samples=6,
)
assert waveform is not None
assert waveform.tolist() == [1.0, 2.0, 3.0]
def test_minicpmo_tts_turn_cleanup_removes_deleted_ref_audio_vocoder_cache(tmp_path):
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
MiniCPMO45OmniTTSForConditionalGeneration,
_TalkerTurnState,
)
prompt_path = tmp_path / "session-ref.wav"
prompt_path.write_bytes(b"wav")
model = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration)
model._talker_turn_states = {
"request": _TalkerTurnState(str(prompt_path), str(prompt_path)),
}
model._talker_consumed_tokens = {"request": 1}
model._talker_request_keys = {}
model._t2w_base_caches = {
str(prompt_path): (torch.tensor([1]), {"cache": torch.tensor([2])}),
}
model._token2wav_state_lock = threading.RLock()
model.audio_tokenizer = None
assert model._close_turn_state("request") is True
assert str(prompt_path) not in model._t2w_base_caches
assert not prompt_path.exists()