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

Signed-off-by: Sy03 <1370724210@qq.com>
This commit is contained in:
Sy03 authored and GitHub committed 2026-07-28 15:59:26 +08:00
1 parent fb89ab43ea
commit ffa9a6af69
60 files changed
+3696 -634

No files matched your search

+25
View File
@@ -165,6 +165,31 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Omni · MiniCPM-o 4.5 Duplex Test"
source_file_dependencies:
- tests/assets/minicpmo_4_5/response_required_16k.wav
- tests/assets/minicpmo_4_5/soft_interrupt_16k.wav
- tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
- tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
- tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- tests/helpers/minicpmo_4_5_duplex.py
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
- vllm_omni/experimental/fullduplex/
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
- vllm_omni/core/sched/
- vllm_omni/outputs/multimodal_accumulation.py
- vllm_omni/engine/stage_init_utils.py
- vllm_omni/engine/async_omni_engine.py
- vllm_omni/entrypoints/cli/serve.py
timeout_in_minutes: 50
commands:
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
mirror_hardwares: h100_2
- label: "TTS · VoxCPM2 Test"
source_file_dependencies:
- tests/e2e/offline_inference/test_voxcpm2_tts.py
+7 -1
View File
@@ -15,7 +15,13 @@ steps:
- label: ":full_moon: Omni · Function Test with H100"
timeout_in_minutes: 120
commands:
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
mirror_hardwares: h100_2
- label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test"
timeout_in_minutes: 50
commands:
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni" --run-level "full_model"
mirror_hardwares: h100_2
- label: ":full_moon: Omni · Doc Test with L4"
+27
View File
@@ -210,6 +210,33 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Omni · MiniCPM-o 4.5 Duplex Test"
source_file_dependencies:
- tests/assets/minicpmo_4_5/response_required_16k.wav
- tests/assets/minicpmo_4_5/soft_interrupt_16k.wav
- tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
- tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
- tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
- tests/helpers/minicpmo_4_5_duplex.py
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
- vllm_omni/experimental/fullduplex/
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
- vllm_omni/core/sched/
- vllm_omni/outputs/multimodal_accumulation.py
- vllm_omni/engine/stage_init_utils.py
- vllm_omni/engine/async_omni_engine.py
- vllm_omni/entrypoints/cli/serve.py
commands:
- |
timeout 20m bash -c "
pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and cuda' --run-level 'core_model'
"
mirror_hardwares: h100_2
- label: "TTS · VoxCPM2 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_voxcpm2_tts.py
@@ -6,6 +6,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
this.playedFrames = 0;
this.underrunFrames = 0;
this.drain = null;
this.terminalFadeFrames = 0;
this.started = false;
this.activeResponseId = null;
this.initialBufferFrames = Math.round(sampleRate * 0.2);
@@ -31,6 +32,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
}
} else if (message.type === 'drain') {
this.drain = { responseId: message.responseId || null };
this.terminalFadeFrames = Math.min(this.fadeFrames, this.bufferedFrames());
if (!this.started && this.bufferedFrames() > 0) {
this.bufferWaitFrames = 0;
this.startPlayback();
@@ -42,6 +44,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
this.playedFrames = 0;
this.underrunFrames = 0;
this.drain = null;
this.terminalFadeFrames = 0;
this.started = false;
this.activeResponseId = null;
this.bufferWaitFrames = this.initialBufferFrames;
@@ -76,6 +79,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
this.playedFrames = 0;
this.underrunFrames = 0;
this.drain = null;
this.terminalFadeFrames = 0;
this.started = false;
this.activeResponseId = null;
this.bufferWaitFrames = this.initialBufferFrames;
@@ -116,6 +120,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
while (target < output.length && this.queue.length > 0) {
const pcm = this.queue[0];
const count = Math.min(output.length - target, pcm.length - this.offset);
const remainingBeforeChunk = this.bufferedFrames();
for (let index = 0; index < count; index += 1) {
let sample = pcm[this.offset + index] / 32768;
if (this.fadeInFrames > 0) {
@@ -123,6 +128,16 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
sample *= elapsed / this.fadeFrames;
this.fadeInFrames -= 1;
}
const remainingFrames = remainingBeforeChunk - index;
if (
this.drain
&& this.terminalFadeFrames > 0
&& remainingFrames <= this.terminalFadeFrames
) {
sample *= this.terminalFadeFrames === 1
? 0
: (remainingFrames - 1) / (this.terminalFadeFrames - 1);
}
output[target + index] = sample;
}
target += count;
Binary file not shown.
Binary file not shown.
@@ -242,9 +242,52 @@ def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None:
"tts_token_ids": [10, 20],
"meta": {"replace_streaming_prompt": True},
}
update.model_intermediate_buffer = {
"ids": {"tts": [41, 42, 99]},
"meta": {"turn_eos_token_id": 99},
}
sched._update_request_as_session(session, update)
assert session.prompt_token_ids == [10, 20]
assert session.additional_information == update.additional_information
assert session.model_intermediate_buffer == {
"ids": {"tts": [41, 42, 99]},
"meta": {"turn_eos_token_id": 99},
}
assert session.status == RequestStatus.WAITING
def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None:
sched = _make_scheduler(stage_id=1)
sched.chunk_transfer_adapter = SimpleNamespace(
receives_chunks=False,
segment_finished_requests=set(),
)
session = _make_request()
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
session.prompt_token_ids = [0] * 59
session._all_token_ids.clear()
session._all_token_ids.extend(session.prompt_token_ids)
session.num_prompt_tokens = 59
session.num_computed_tokens = 59
update = _make_update([0] * 10)
update.additional_information = None
update.model_intermediate_buffer = {
"ids": {"tts": list(range(8))},
"hidden_states": {"tts": [[0.0]] * 8},
"meta": {
"replace_streaming_prompt": True,
"next_stage_prompt_len": 10,
},
}
sched._update_request_as_session(session, update)
assert session.prompt_token_ids == [0] * 10
assert list(session._all_token_ids) == [0] * 10
assert session.num_prompt_tokens == 10
assert session.num_computed_tokens == 0
assert session.additional_information is None
assert session.model_intermediate_buffer == update.model_intermediate_buffer
assert session.status == RequestStatus.WAITING
@@ -54,6 +54,18 @@ class _SchedulerStub(OmniGenerationScheduler):
raise AssertionError("unexpected enqueue for skipped_waiting miss")
class _AsyncChunkStopSchedulerStub(OmniGenerationScheduler):
def __init__(self) -> None:
self.num_waiting_for_streaming_input = 0
self.chunk_transfer_adapter = SimpleNamespace(receives_chunks=True)
self.enqueued: list[Request] = []
self.enqueued_statuses: list[RequestStatus] = []
def _enqueue_waiting_request(self, session: Request) -> None:
self.enqueued.append(session)
self.enqueued_statuses.append(session.status)
def _make_request(**kwargs) -> Request:
sp = SamplingParams(max_tokens=8)
defaults = dict(
@@ -133,6 +145,21 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
assert output.is_segment_finished is True
def test_async_chunk_resumable_stop_rearms_connector_polling() -> None:
sched = _AsyncChunkStopSchedulerStub()
session = _make_request(request_id="req-async-chunk-next-segment")
session.status = RequestStatus.FINISHED_STOPPED
session.resumable = True
finished = sched._handle_stopped_request(session)
assert finished is False
assert sched.enqueued == [session]
assert sched.enqueued_statuses == [RequestStatus.WAITING]
assert session.status == RequestStatus.WAITING
assert sched.num_waiting_for_streaming_input == 0
class TestReplaceSessionWithStreamingUpdate:
def test_resets_tokens_and_prompt_from_update(self) -> None:
sched = _SchedulerStub()
@@ -22,10 +22,7 @@ _MINICPMO_SERVER = [
model="openbmb/MiniCPM-o-4_5",
stage_config_path=_MINICPMO_DEPLOY,
use_stage_cli=True,
server_args=[
"--trust-remote-code",
"--no-async-chunk",
],
server_args=["--trust-remote-code"],
),
id="minicpmo_4_5",
)
@@ -71,6 +71,30 @@ def test_streaming_payload_can_replace_placeholder_prompt(mocker: MockerFixture)
request.update_block_hashes.assert_called_once_with()
def test_streaming_payload_can_append_exact_prompt_length(mocker: MockerFixture) -> None:
request = SimpleNamespace(
_all_token_ids=[0, 0, 7, 8],
_output_token_ids=[7, 8],
prompt_token_ids=[0, 0],
num_computed_tokens=4,
num_prompt_tokens=2,
update_block_hashes=mocker.Mock(),
)
payload = {
"ids": {"prompt": [1, 2, 3]},
"meta": {"next_stage_prompt_len": 3},
}
construct_next_stage_streaming_input_prompt(payload, request)
assert request.prompt_token_ids == [0, 0, 7, 8, 0, 0, 0]
assert request._all_token_ids == [0, 0, 7, 8, 0, 0, 0]
assert request._output_token_ids == []
assert request.num_computed_tokens == 4
assert request.num_prompt_tokens == 7
request.update_block_hashes.assert_called_once_with()
@pytest.fixture
def build_adapter(monkeypatch, mocker: MockerFixture):
def _build(
@@ -252,6 +276,37 @@ def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_ad
assert len(adapter._pending_save_reqs) == 1
def test_send_single_request_terminal_chunk_still_flushes_processor(build_adapter, monkeypatch):
"""A terminal stop is not a segment boundary (#5383), but the producer-side
processor must still receive the flush signal on the terminal chunk.
Passing only ``is_segment_finished`` starved processors of their final
accumulated payload once terminal stops stopped setting it (#5413: the
downstream stage got ``meta.finished`` with the tail data missing).
"""
adapter, connector = build_adapter(stage_id=0)
request = _req("req-terminal", RequestStatus.FINISHED_STOPPED, external_req_id="ext-terminal")
seen_flush_flags = []
def recording_processor(**kwargs):
seen_flush_flags.append(kwargs["is_finished"])
return OmniPayloadStruct(
codes=CodesStruct(audio=torch.tensor([1, 2, 3], dtype=torch.long)),
)
adapter.custom_process_next_stage_input_func = recording_processor
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
adapter._send_single_request(
{"multimodal_output": None, "request": request, "is_finished": True, "is_segment_finished": False}
)
assert seen_flush_flags == [True]
sent_payload = connector.put.call_args.kwargs["data"]
assert bool(sent_payload.meta.finished.item()) is True
assert bool(sent_payload.meta.is_segment_finished.item()) is False
def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, monkeypatch):
"""Producer may return a struct with ``meta=None`` (e.g. payload that
carries only ``embed`` or ``codes``). The sender's ``meta is not None``
@@ -314,6 +369,23 @@ def test_send_single_request_struct_preserves_segment_finished(build_adapter, mo
assert sent_payload.meta.is_segment_finished.item() is True
def test_send_single_request_respects_processor_receiver_boundary(build_adapter, monkeypatch):
adapter, connector = build_adapter(stage_id=1)
request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
meta=MetaStruct(is_segment_finished=torch.tensor(False, dtype=torch.bool))
)
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
adapter._send_single_request(
{"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
)
sent_payload = connector.put.call_args.kwargs["data"]
assert sent_payload.meta.is_segment_finished.item() is False
def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
adapter, _ = build_adapter(stage_id=1)
request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
@@ -463,7 +535,7 @@ def test_fifo_promotion(build_adapter):
adapter.process_pending_chunks(waiting_queue, running_queue)
assert list(adapter._active_streams) == ["req-1", "req-2"]
assert waiting_queue == reqs[2:]
assert waiting_queue == []
assert reqs[0].status == RequestStatus.WAITING_FOR_CHUNK
assert reqs[1].status == RequestStatus.WAITING_FOR_CHUNK
assert reqs[2].status == RequestStatus.WAITING
@@ -473,6 +545,7 @@ def test_fifo_promotion(build_adapter):
# (commit c4d95fd9 — otherwise the terminal chunk deadlocks at c=8 K=2).
# Simulate it here so promotion can pick up the freed slot.
adapter._evict_finished_active_streams({"req-1"})
adapter.restore_queues(waiting_queue, running_queue)
adapter.process_pending_chunks(waiting_queue, running_queue)
assert list(adapter._active_streams) == ["req-2", "req-3"]
@@ -481,6 +554,22 @@ def test_fifo_promotion(build_adapter):
assert reqs[2].status == RequestStatus.WAITING_FOR_CHUNK
def test_non_active_waiting_request_is_held_off_scheduler(build_adapter):
adapter, _ = build_adapter(stage_id=2, max_num_seqs=2, active_stream_window=1)
active = _req("req-active", RequestStatus.WAITING)
non_active = _req("req-non-active", RequestStatus.WAITING)
waiting_queue = DummyWaitingQueue([active, non_active])
running_queue = []
adapter.process_pending_chunks(waiting_queue, running_queue)
assert waiting_queue == []
assert active.status == RequestStatus.WAITING_FOR_CHUNK
assert non_active.status == RequestStatus.WAITING
assert list(adapter._pending_load_reqs) == [active]
assert list(adapter.waiting_for_chunk_waiting_requests) == [active, non_active]
def test_legacy_k0(build_adapter):
adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=0)
waiting_req = _req("waiting", RequestStatus.WAITING)
@@ -509,12 +598,13 @@ def test_finished_releases_slot(build_adapter):
adapter.process_pending_chunks(waiting_queue, running_queue)
assert list(adapter._active_streams) == ["req-1"]
assert waiting_queue == [req_2]
assert waiting_queue == []
adapter.finished_requests.add("req-1")
# Eviction is deferred to postprocess_scheduler_output in the runtime path
# (commit c4d95fd9). Simulate it so promotion can pick up the freed slot.
adapter._evict_finished_active_streams({"req-1"})
adapter.restore_queues(waiting_queue, running_queue)
adapter.process_pending_chunks(waiting_queue, running_queue)
assert list(adapter._active_streams) == ["req-2"]
@@ -83,7 +83,7 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None:
@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
"""Test processing text, generating audio output through the talker token2wav path."""
"""Test processing text and generating audio through Talker and Code2Wav."""
request_config = {"prompts": get_question("text"), "modalities": ["audio"]}
omni_runner_handler.send_omni_request(request_config)
@@ -1174,10 +1174,13 @@ async def _send_listen_only_overlap_pair(
before_second_listen = state.model_listen_count
overlap_started_while_active = not state.response_done(first_response_id)
model_unit_ready_while_active = False
second_unit_event_index: int | None = None
def record_model_unit_ready() -> None:
nonlocal model_unit_ready_while_active
nonlocal model_unit_ready_while_active, second_unit_event_index
model_unit_ready_while_active = not state.response_done(first_response_id)
if model_unit_ready_while_active:
second_unit_event_index = len(state.events)
await _send_pcm16(
ws,
@@ -1239,14 +1242,13 @@ async def _send_listen_only_overlap_pair(
label=transcripts[1],
)
second_speech_index = _nth_event_index(state, "input_audio_buffer.speech_started", 2)
first_done_index = _event_index_for_response(state, "response.done", first_response_id)
overlap_ok = (
overlap_started_while_active
and model_unit_ready_while_active
and second_speech_index is not None
and second_unit_event_index is not None
and first_done_index is not None
and second_speech_index < first_done_index
and second_unit_event_index < first_done_index
)
return [first_response_id, second_response_id], ["speak", second_outcome], overlap_ok
@@ -1380,6 +1382,17 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
await ws.send(json.dumps({"type": "session.close"}))
await _wait_for(state, lambda: state.count("session.closed") > 0, timeout_s=20, label="session.closed")
finally:
if state.count("session.created") > 0 and state.count("session.closed") == 0:
try:
await ws.send(json.dumps({"type": "session.close"}))
await _wait_for(
state,
lambda: state.count("session.closed") > 0,
timeout_s=min(args.timeout_s, 5),
label="session.closed during cleanup",
)
except (ConnectionClosed, TimeoutError):
pass
stop.set()
reader.cancel()
try:
@@ -19,7 +19,11 @@ SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
sys.path.insert(0, str(SCRIPT_DIR))
from minicpmo_realtime_duplex_scenarios import _url_with_model, run_demo # noqa: E402
from minicpmo_realtime_duplex_scenarios import ( # noqa: E402
_ref_audio_data_url,
_url_with_model,
run_demo,
)
def _with_resume_mode(url: str) -> str:
@@ -114,7 +118,12 @@ async def _open_admission_session(
args: argparse.Namespace,
session_id: str,
) -> tuple[object, dict[str, object]]:
url = _url_with_model(args.url, args.model, session_id=session_id)
url = _url_with_model(
args.url,
args.model,
autostart=False if getattr(args, "ref_audio", None) else None,
session_id=session_id,
)
ws = await websockets.connect(url, max_size=64 * 1024 * 1024)
await ws.send(
json.dumps(
@@ -125,6 +134,7 @@ async def _open_admission_session(
"model": args.model,
"modalities": ["audio", "text"],
"extra_body": {"minicpmo45_native_duplex": True},
**({"ref_audio": _ref_audio_data_url(args.ref_audio)} if getattr(args, "ref_audio", None) else {}),
},
}
)
@@ -151,7 +161,12 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str,
accepted.append(await _open_admission_session(args, f"{prefix}-accepted-{index}"))
overflow_id = f"{prefix}-overflow"
overflow_url = _url_with_model(args.url, args.model, session_id=overflow_id)
overflow_url = _url_with_model(
args.url,
args.model,
autostart=False if getattr(args, "ref_audio", None) else None,
session_id=overflow_id,
)
async with websockets.connect(overflow_url, max_size=64 * 1024 * 1024) as overflow:
await overflow.send(
json.dumps(
@@ -162,6 +177,11 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str,
"model": args.model,
"modalities": ["audio", "text"],
"extra_body": {"minicpmo45_native_duplex": True},
**(
{"ref_audio": _ref_audio_data_url(args.ref_audio)}
if getattr(args, "ref_audio", None)
else {}
),
},
}
)
@@ -207,16 +227,27 @@ async def _resume_probe(
session_id: str,
expect_expired: bool = False,
) -> dict[str, object]:
url = _url_with_model(args.url, args.model, session_id=session_id)
url = _url_with_model(
args.url,
args.model,
autostart=False if getattr(args, "ref_audio", None) else None,
session_id=session_id,
)
async with websockets.connect(url, max_size=64 * 1024 * 1024) as first:
await first.send(
json.dumps(
{
"type": "session.update",
"session": {
"session_id": session_id,
"model": args.model,
"modalities": ["audio", "text"],
"extra_body": {"minicpmo45_native_duplex": True},
**(
{"ref_audio": _ref_audio_data_url(args.ref_audio)}
if getattr(args, "ref_audio", None)
else {}
),
},
}
)
@@ -306,7 +337,12 @@ async def _takeover_probe(
*,
session_id: str,
) -> dict[str, object]:
url = _url_with_model(args.url, args.model, session_id=session_id)
url = _url_with_model(
args.url,
args.model,
autostart=False if getattr(args, "ref_audio", None) else None,
session_id=session_id,
)
first = await websockets.connect(url, max_size=64 * 1024 * 1024)
second = None
try:
@@ -315,9 +351,15 @@ async def _takeover_probe(
{
"type": "session.update",
"session": {
"session_id": session_id,
"model": args.model,
"modalities": ["audio", "text"],
"extra_body": {"minicpmo45_native_duplex": True},
**(
{"ref_audio": _ref_audio_data_url(args.ref_audio)}
if getattr(args, "ref_audio", None)
else {}
),
},
}
)
@@ -451,20 +493,7 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]:
args,
session_id=f"takeover-{args.takeover_session_index}-{uuid.uuid4().hex}",
)
expiry_result = None
if args.expire_session_index is not None:
if not 0 <= args.expire_session_index < args.sessions:
raise ValueError("--expire-session-index is outside the session range")
expiry_result = await _resume_probe(
args,
session_id=f"expire-{args.expire_session_index}-{uuid.uuid4().hex}",
expect_expired=True,
)
admission_result = (
await _admission_probe(args, limit=args.verify_admission_limit)
if args.verify_admission_limit is not None
else None
)
lifecycle_result = await run_lifecycle_probes(args)
session_results = await asyncio.gather(
*(run_demo(_demo_args(args, index)) for index in range(args.sessions)),
@@ -487,16 +516,15 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]:
and semantic_isolation_ok
and (resume_result is None or resume_result.get("ok") is True)
and (takeover_result is None or takeover_result.get("ok") is True)
and (expiry_result is None or expiry_result.get("ok") is True)
and (admission_result is None or admission_result.get("ok") is True)
and lifecycle_result["ok"] is True
),
"session_count": args.sessions,
"identity_isolation_ok": identity_isolation_ok,
"semantic_isolation_ok": semantic_isolation_ok,
"resume": resume_result,
"takeover": takeover_result,
"expiry": expiry_result,
"admission": admission_result,
"expiry": lifecycle_result["expiry"],
"admission": lifecycle_result["admission"],
"failures": failures,
"sessions": completed,
}
@@ -507,6 +535,34 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]:
return result
async def run_lifecycle_probes(args: argparse.Namespace) -> dict[str, object]:
"""Run expiry and admission probes without requiring model output."""
if args.sessions < 1:
raise ValueError("--sessions must be positive")
expiry_result = None
if args.expire_session_index is not None:
if not 0 <= args.expire_session_index < args.sessions:
raise ValueError("--expire-session-index is outside the session range")
expiry_result = await _resume_probe(
args,
session_id=f"expire-{args.expire_session_index}-{uuid.uuid4().hex}",
expect_expired=True,
)
admission_result = (
await _admission_probe(args, limit=args.verify_admission_limit)
if args.verify_admission_limit is not None
else None
)
return {
"ok": (
(expiry_result is None or expiry_result.get("ok") is True)
and (admission_result is None or admission_result.get("ok") is True)
),
"expiry": expiry_result,
"admission": admission_result,
}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--url", default="ws://127.0.0.1:8113/v1/realtime?duplex=1")
@@ -2,8 +2,8 @@
This E2E driver runs the public ``realtime_duplex_demo.py`` against a live
duplex backend. Arbitrary audio defaults to the model-policy lifecycle contract.
The stronger response-required mode binds the two-response contract to a known
input checksum and expected second response.
The stronger response-required mode binds the multi-response contract to a
known input checksum and expected follow-up response.
"""
from __future__ import annotations
@@ -39,6 +39,17 @@ def _validate_input_sha256(path: Path, expected: str) -> str:
return actual
def _input_duration_s(path: Path) -> float:
with wave.open(str(path), "rb") as wav_file:
return wav_file.getnframes() / wav_file.getframerate()
def _client_process_timeout_s(input_wav: Path, protocol_timeout_s: float) -> float:
# The child first streams the WAV in real time, then can independently
# exhaust its post-commit and session-close protocol waits.
return _input_duration_s(input_wav) + 2 * protocol_timeout_s + 30.0
def _read_json(path: Path) -> dict[str, object]:
if not path.is_file():
return {}
@@ -211,7 +222,7 @@ def summarize_artifacts(
validation_mode: str,
min_responses: int,
min_audio_deltas_per_response: int,
expect_second_response_substring: str | None,
expect_followup_response_substring: str | None,
) -> dict[str, object]:
if validation_mode not in {"model-policy", "response-required"}:
raise ValueError(f"unsupported validation mode: {validation_mode}")
@@ -274,12 +285,20 @@ def summarize_artifacts(
)
final_listen_after_commit = commit_index is not None and any(index > commit_index for index in listen_indices)
transcript = "".join(str(summary.get("transcript") or "") for summary in response_summaries)
second_response_transcript = (
str(response_summaries[1].get("transcript") or "") if len(response_summaries) >= 2 else ""
followup_response_transcripts = [
str(summary.get("transcript") or "")
for summary in response_summaries[1:]
if commit_index is not None
and isinstance(summary.get("created_index"), int)
and summary["created_index"] < commit_index
]
followup_response_transcript_ok = any(
bool(_normalize_text(transcript)) for transcript in followup_response_transcripts
)
followup_response_transcript_expectation_ok = not expect_followup_response_substring or any(
_normalize_text(expect_followup_response_substring) in _normalize_text(transcript)
for transcript in followup_response_transcripts
)
second_response_transcript_expectation_ok = not expect_second_response_substring or _normalize_text(
expect_second_response_substring
) in _normalize_text(second_response_transcript)
error_events = [
event
for event in events
@@ -310,9 +329,8 @@ def summarize_artifacts(
mode_contract_ok = validation_mode == "model-policy" or (
second_response_before_final_commit
and listen_before_first_response
and listen_between_responses
and listen_after_last_done
and second_response_transcript_expectation_ok
and followup_response_transcript_ok
)
ok = common_contract_ok and mode_contract_ok
return {
@@ -337,9 +355,10 @@ def summarize_artifacts(
"listen_after_response_before_commit": listen_after_response_before_commit,
"final_listen_after_commit": final_listen_after_commit,
"transcript": transcript,
"second_response_transcript": second_response_transcript,
"expect_second_response_substring": expect_second_response_substring,
"second_response_transcript_expectation_ok": second_response_transcript_expectation_ok,
"followup_response_transcripts": followup_response_transcripts,
"followup_response_transcript_ok": followup_response_transcript_ok,
"expect_followup_response_substring": expect_followup_response_substring,
"followup_response_transcript_expectation_ok": followup_response_transcript_expectation_ok,
"error_count": len(error_events),
"cancelled_count": cancelled_count,
"compact_sequence": _compact_sequence(events),
@@ -388,7 +407,7 @@ async def run_soft_interrupt(args: argparse.Namespace) -> dict[str, object]:
try:
stdout_bytes, stderr_bytes = await asyncio.wait_for(
process.communicate(),
timeout=args.timeout_s + 30.0,
timeout=_client_process_timeout_s(input_wav, args.timeout_s),
)
except TimeoutError:
process.kill()
@@ -406,7 +425,7 @@ async def run_soft_interrupt(args: argparse.Namespace) -> dict[str, object]:
validation_mode=args.validation_mode,
min_responses=args.min_responses,
min_audio_deltas_per_response=args.min_audio_deltas_per_response,
expect_second_response_substring=args.expect_second_response_substring,
expect_followup_response_substring=args.expect_followup_response_substring,
)
summary.update(
{
@@ -444,14 +463,12 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--min-responses", type=int, default=2)
parser.add_argument("--min-audio-deltas-per-response", type=int, default=2)
parser.add_argument("--input-sha256")
parser.add_argument("--expect-second-response-substring")
parser.add_argument("--expect-followup-response-substring")
args = parser.parse_args()
if args.validation_mode == "response-required" and args.min_responses < 2:
parser.error("--min-responses must be at least 2")
if args.validation_mode == "response-required" and not args.input_sha256:
parser.error("--input-sha256 is required in response-required mode")
if args.validation_mode == "response-required" and not args.expect_second_response_substring:
parser.error("--expect-second-response-substring is required in response-required mode")
if args.min_audio_deltas_per_response < 1:
parser.error("--min-audio-deltas-per-response must be positive")
return args
+7 -11
View File
@@ -1,9 +1,7 @@
"""
E2E Online tests for MiniCPM-o 4.5 model with multimodal input and audio / text output.
"""E2E online tests for MiniCPM-o 4.5 multimodal input and audio/text output.
MiniCPM-o 4.5 has ``async_chunk: false``, ``max_num_seqs: 1`` on both stages,
and the vocoder runs in-process inside the talker stage rather than as a separate
Code2Wav stage.
The batching deployment uses async chunk transfer across separate Thinker,
Talker, and Code2Wav stages.
"""
import os
@@ -26,10 +24,7 @@ test_params = [
model=_MODEL,
stage_config_path=_CI_DEPLOY,
use_stage_cli=True,
server_args=[
"--trust-remote-code",
"--no-async-chunk",
],
server_args=["--trust-remote-code"],
),
id="default",
)
@@ -99,7 +94,7 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
def test_text_to_audio_001(omni_server, openai_client) -> None:
"""
Test text-only input generating text + audio output via OpenAI API.
This exercises the talker TTS region detection and in-process token2wav vocoder.
This exercises Talker TTS region detection and the Code2Wav stage.
Deploy Setting: default 2GPU
Input Modal: text
Output Modal: text + audio
@@ -133,13 +128,14 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
audio_data_url=audio_data_url,
content_text=get_prompt("mix"),
content_text=get_prompt(),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"key_words": {"text": ["Beijing"]},
}
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
@@ -0,0 +1,147 @@
"""CI coverage for the MiniCPM-o 4.5 native-duplex Realtime API."""
from __future__ import annotations
import asyncio
import json
import uuid
from pathlib import Path
import pytest
import websockets
from tests.e2e.online_serving.minicpmo_realtime_duplex_scenarios import (
_ref_audio_data_url,
run_demo,
)
from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import (
run_multi_session,
)
from tests.helpers.mark import hardware_test
from tests.helpers.minicpmo_4_5_duplex import (
CORE_SERVER_PARAMS,
SERVER_PARAMS,
demo_args,
multi_session_args,
realtime_url,
resolve_ref_audio,
validated_input_wav,
)
from vllm_omni.experimental.fullduplex.client import build_realtime_url
pytestmark = pytest.mark.omni
async def _receive_protocol_events(ws, required_types: set[str], *, timeout_s: float) -> list[dict[str, object]]:
async def receive() -> list[dict[str, object]]:
events: list[dict[str, object]] = []
seen: set[str] = set()
while not required_types.issubset(seen):
raw = await ws.recv()
if not isinstance(raw, str):
continue
event = json.loads(raw)
if not isinstance(event, dict):
continue
events.append(event)
event_type = event.get("type")
if event_type == "error":
raise AssertionError(f"WebSocket protocol smoke received an error: {event}")
if isinstance(event_type, str):
seen.add(event_type)
return events
return await asyncio.wait_for(receive(), timeout=timeout_s)
async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[dict[str, object]]:
session_id = f"duplex-ci-protocol-{uuid.uuid4().hex}"
websocket_url = build_realtime_url(url, model, autostart=False, session_id=session_id)
async with websockets.connect(websocket_url, max_size=64 * 1024 * 1024) as ws:
await ws.send(
json.dumps(
{
"type": "session.update",
"session": {
"session_id": session_id,
"model": model,
"modalities": ["audio", "text"],
"ref_audio": _ref_audio_data_url(str(ref_audio)),
"extra_body": {"minicpmo45_native_duplex": True},
},
}
)
)
events = await _receive_protocol_events(
ws,
{"session.created", "session.updated"},
timeout_s=60,
)
await ws.send(json.dumps({"type": "session.close"}))
events.extend(await _receive_protocol_events(ws, {"session.closed"}, timeout_s=60))
return events
@pytest.mark.core_model
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server", CORE_SERVER_PARAMS, indirect=True)
def test_duplex_websocket_protocol_smoke(omni_server, model_prefix: str) -> None:
ref_audio = resolve_ref_audio(model_prefix)
events = asyncio.run(
_run_protocol_smoke(
url=realtime_url(omni_server),
model=omni_server.model,
ref_audio=ref_audio,
)
)
event_types = [event.get("type") for event in events]
assert "session.created" in event_types
assert "session.updated" in event_types
assert event_types[-1] == "session.closed"
@pytest.mark.advanced_model
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
def test_duplex_single_session_response_required(omni_server, model_prefix: str, tmp_path: Path) -> None:
result = asyncio.run(
run_demo(
demo_args(
omni_server=omni_server,
input_wav=validated_input_wav(),
ref_audio=resolve_ref_audio(model_prefix),
output_dir=tmp_path / "single_session",
)
)
)
assert result["ok"] is True
assert result["audio_delta_count"] > 0
assert result["done_count"] == 1
assert result["error_count"] == 0
assert result["all_audio_responses_have_transcript"] is True
assert result["transcript_delta_done_ok"] is True
@pytest.mark.advanced_model
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
def test_duplex_two_sessions_resume_and_takeover(omni_server, model_prefix: str, tmp_path: Path) -> None:
result = asyncio.run(
run_multi_session(
multi_session_args(
omni_server=omni_server,
input_wav=validated_input_wav(),
ref_audio=resolve_ref_audio(model_prefix),
output_dir=tmp_path / "multi_session",
response_required=True,
)
)
)
assert result["ok"] is True
assert result["session_count"] == 2
assert result["resume"]["ok"] is True
assert result["takeover"]["ok"] is True
assert not result["failures"]
assert all(session["audio_delta_count"] > 0 for session in result["sessions"])
assert all(session["done_count"] == 1 for session in result["sessions"])
assert all(session["error_count"] == 0 for session in result["sessions"])
@@ -0,0 +1,86 @@
"""Nightly lifecycle coverage for the MiniCPM-o 4.5 native-duplex API."""
from __future__ import annotations
import asyncio
import json
from pathlib import Path
from types import SimpleNamespace
import pytest
from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import (
run_lifecycle_probes,
)
from tests.e2e.online_serving.run_minicpmo_realtime_duplex_soft_interrupt import (
run_soft_interrupt,
)
from tests.helpers.mark import hardware_test
from tests.helpers.minicpmo_4_5_duplex import (
SERVER_PARAMS,
SOFT_INTERRUPT_SHA256,
multi_session_args,
realtime_url,
resolve_ref_audio,
validated_input_wav,
validated_soft_interrupt_wav,
)
pytestmark = [pytest.mark.full_model, pytest.mark.omni]
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
def test_duplex_admission_and_expiry_reaper(omni_server, model_prefix: str, tmp_path: Path) -> None:
args = multi_session_args(
omni_server=omni_server,
input_wav=validated_input_wav(),
ref_audio=resolve_ref_audio(model_prefix),
output_dir=tmp_path / "admission_expiry",
response_required=False,
)
args.sessions = 1
args.disconnect_session_index = None
args.takeover_session_index = None
args.expire_session_index = 0
args.verify_admission_limit = 2
result = asyncio.run(run_lifecycle_probes(args))
assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
assert result["expiry"]["ok"] is True
assert result["expiry"]["error_code"] == "session_resume_expired"
assert result["admission"]["ok"] is True
assert result["admission"]["overflow_error_code"] == "resource_exhausted"
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) -> None:
input_wav = validated_soft_interrupt_wav()
result = asyncio.run(
run_soft_interrupt(
SimpleNamespace(
url=realtime_url(omni_server),
model=omni_server.model,
input_wav=str(input_wav),
ref_audio=str(resolve_ref_audio(model_prefix)),
output_dir=str(tmp_path / "soft_interrupt"),
summary_output=None,
chunk_ms=200,
timeout_s=180.0,
require_audio=True,
no_realtime_pacing=False,
validation_mode="response-required",
min_responses=2,
min_audio_deltas_per_response=2,
input_sha256=SOFT_INTERRUPT_SHA256,
expect_followup_response_substring=None,
)
)
)
assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
assert result["error_count"] == 0
assert result["response_lifecycle_ok"] is True
assert result["response_audio_contract_ok"] is True
assert result["followup_response_transcript_ok"] is True
@@ -1,7 +1,7 @@
"""
E2E Online expansion tests for MiniCPM-o 4.5 covering modality combinations,
in-process token2wav vocoder behavior, max_num_seqs=1 serial execution,
and edge-case validation.
"""E2E online expansion tests for MiniCPM-o 4.5.
These cover modality combinations, separate Code2Wav behavior, sequential
request isolation, longer audio output, and Chinese speech.
"""
import os
@@ -94,10 +94,8 @@ def test_text_video_to_text_001(omni_server, openai_client) -> None:
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_sequential_requests_independent(omni_server, openai_client) -> None:
"""
Verify that sequential requests produce independent results and state
does not leak between requests. MiniCPM-o 4.5 has max_num_seqs=1 so
requests are processed one-at-a-time; the second request must not
receive the first request's audio or token state.
Verify that sequential requests produce independent results and that the
second request does not receive the first request's audio or token state.
Deploy Setting: default 2GPU
Input Modal: text (two different prompts)
Output Modal: text + audio (both)
@@ -138,7 +136,7 @@ def test_sequential_requests_independent(omni_server, openai_client) -> None:
def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
"""
Test text input generating a longer audio output to exercise the
token2wav decoder across multiple frames.
Code2Wav stage across multiple frames.
Deploy Setting: default 2GPU
Input Modal: text (longer prompt)
Output Modal: text + audio
@@ -146,14 +144,14 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
content_text="Tell me a short story about a cat in about 50 words.",
content_text="What is the capital of China? Answer in 40 words.",
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"key_words": {"audio": ["cat"]},
"key_words": {"audio": ["Beijing"]},
}
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
@@ -12,14 +12,11 @@ import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
DEMO_PATH = Path(__file__).resolve().parents[1] / "e2e/online_serving/minicpmo_realtime_duplex_scenarios.py"
MULTI_DEMO_PATH = (
Path(__file__).resolve().parents[1] / "e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py"
)
PAIR_DEMO_PATH = Path(__file__).resolve().parents[1] / "e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py"
SOFT_INTERRUPT_DEMO_PATH = (
Path(__file__).resolve().parents[1] / "e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py"
)
DRIVER_DIR = Path(__file__).resolve().parent
DEMO_PATH = DRIVER_DIR / "minicpmo_realtime_duplex_scenarios.py"
MULTI_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_multi_session.py"
PAIR_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_demo_pair.py"
SOFT_INTERRUPT_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_soft_interrupt.py"
def _load_demo_module():
@@ -338,14 +335,7 @@ def test_realtime_duplex_soft_interrupt_accepts_explicit_ref_audio(monkeypatch):
assert args.validation_mode == "model-policy"
@pytest.mark.parametrize(
"extra_args",
[
["--expect-second-response-substring", "一加一等于二"],
["--input-sha256", "a" * 64],
],
)
def test_realtime_duplex_soft_interrupt_response_required_needs_bound_fixture(monkeypatch, extra_args):
def test_realtime_duplex_soft_interrupt_response_required_needs_bound_fixture(monkeypatch):
demo = _load_soft_interrupt_demo_module()
monkeypatch.setattr(
demo.sys,
@@ -360,7 +350,6 @@ def test_realtime_duplex_soft_interrupt_response_required_needs_bound_fixture(mo
"ref.wav",
"--validation-mode",
"response-required",
*extra_args,
],
)
@@ -412,7 +401,6 @@ def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp
"_client_received_at_s": 2.6,
},
{"type": "response.done", "response_id": first_response_id, "_client_received_at_s": 3.0},
{"type": "response.listen", "_client_received_at_s": 3.2},
{"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 4.0},
{"type": "response.speak", "response_id": second_response_id, "_client_received_at_s": 4.0},
{
@@ -458,10 +446,11 @@ def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp
validation_mode="response-required",
min_responses=2,
min_audio_deltas_per_response=2,
expect_second_response_substring="一加一等于二",
expect_followup_response_substring="一加一等于二",
)
assert summary["ok"] is True
assert summary["listen_between_responses"] is False
assert summary["second_response_before_final_commit"] is True
assert summary["final_listen_after_commit"] is True
assert summary["response_audio_contract_ok"] is True
@@ -498,7 +487,7 @@ def test_realtime_duplex_soft_interrupt_model_policy_accepts_single_response(tmp
validation_mode="model-policy",
min_responses=2,
min_audio_deltas_per_response=2,
expect_second_response_substring=None,
expect_followup_response_substring=None,
)
assert summary["ok"] is True
@@ -537,14 +526,14 @@ def test_realtime_duplex_soft_interrupt_response_required_rejects_single_respons
validation_mode="response-required",
min_responses=2,
min_audio_deltas_per_response=2,
expect_second_response_substring="一加一等于二",
expect_followup_response_substring="一加一等于二",
)
assert summary["ok"] is False
assert summary["enough_responses"] is False
def test_realtime_duplex_soft_interrupt_checks_expected_text_in_second_response(tmp_path):
def test_realtime_duplex_soft_interrupt_reports_text_expectation_without_gating(tmp_path):
demo = _load_soft_interrupt_demo_module()
output = tmp_path / "second_response_text"
output.mkdir()
@@ -611,11 +600,32 @@ def test_realtime_duplex_soft_interrupt_checks_expected_text_in_second_response(
validation_mode="response-required",
min_responses=2,
min_audio_deltas_per_response=2,
expect_second_response_substring="一加一等于二",
expect_followup_response_substring="一加一等于二",
)
assert summary["ok"] is False
assert summary["second_response_transcript_expectation_ok"] is False
assert summary["ok"] is True
assert summary["followup_response_transcript_ok"] is True
assert summary["followup_response_transcript_expectation_ok"] is False
events = [
event
for event in events
if event.get("response_id") != second_response_id or event["type"] != "response.audio_transcript.delta"
]
(output / "events.jsonl").write_text(
"".join(demo.json.dumps(event) + "\n" for event in events),
encoding="utf-8",
)
missing_transcript_summary = demo.summarize_artifacts(
output_dir=output,
validation_mode="response-required",
min_responses=2,
min_audio_deltas_per_response=2,
expect_followup_response_substring="一加一等于二",
)
assert missing_transcript_summary["ok"] is False
assert missing_transcript_summary["followup_response_transcript_ok"] is False
def test_realtime_duplex_multi_session_resume_url_disables_autostart():
@@ -1523,7 +1533,7 @@ def test_realtime_duplex_demo_waits_at_each_model_unit_and_stops_after_speak():
assert len(ws.messages) == 10
def test_realtime_duplex_demo_listen_only_overlap_sends_next_turn_before_first_done(monkeypatch):
def test_realtime_duplex_demo_listen_only_overlap_accepts_silence_unit_before_first_done(monkeypatch):
demo = _load_demo_module()
state = demo.DemoState()
send_calls = 0
@@ -1548,7 +1558,6 @@ def test_realtime_duplex_demo_listen_only_overlap_sends_next_turn_before_first_d
del args
assert not state.response_done("resp-first")
kwargs["on_model_unit_ready"]()
state.add({"type": "input_audio_buffer.speech_started", "turn": 2})
_add_response_transcript(state, "resp-first", transcript="第一轮完成", audio=False)
state.add({"type": "response.created", "response": {"id": "resp-second"}})
_add_response_transcript(state, "resp-second", transcript="第二轮完成", audio=True)
@@ -1568,7 +1577,7 @@ def test_realtime_duplex_demo_listen_only_overlap_sends_next_turn_before_first_d
ws,
state,
b"\x00\x00",
b"\x01\x00",
b"\x00\x00",
transcripts=("first", "second"),
durations_ms=(None, None),
chunk_ms=200,
@@ -18,6 +18,7 @@ def test_chunk_accumulation_policy_replaces_snapshots_and_drains_delta_state():
accumulated = MultimodalPayload.from_dict(
{
"audio": torch.tensor([1.0]),
"meta.segment_end": torch.tensor([0]),
"meta.tts_is_last_chunk": torch.tensor([0]),
"meta.turn_end": torch.tensor([0]),
"meta.stable_request_value": "keep",
@@ -26,6 +27,7 @@ def test_chunk_accumulation_policy_replaces_snapshots_and_drains_delta_state():
incoming = MultimodalPayload.from_dict(
{
"audio": torch.tensor([2.0]),
"meta.segment_end": torch.tensor([1]),
"meta.tts_is_last_chunk": torch.tensor([1]),
"meta.turn_end": torch.tensor([1]),
}
@@ -41,6 +43,7 @@ def test_chunk_accumulation_policy_replaces_snapshots_and_drains_delta_state():
drain_delta_payload(merged)
assert "audio" not in merged
assert "meta.segment_end" not in merged
assert "meta.tts_is_last_chunk" not in merged
assert "meta.turn_end" not in merged
assert merged.metadata["meta.stable_request_value"] == "keep"
+50
View File
@@ -1813,6 +1813,56 @@ async def test_stage_pool_submit_update_refreshes_output_processor_state() -> No
@pytest.mark.asyncio
async def test_handle_streaming_update_unknown_request_is_dropped() -> None:
"""A streaming update for an untracked request must not resurrect it.
Updates always follow the first-chunk add_request through the same
ordered submission queue, so an unknown id can only mean the request
already finished or was aborted (e.g. the client disconnected).
Falling back to add_request created headless sessions that kept
cycling through the stages (issue #4271).
"""
class RecordingPool:
def __init__(self) -> None:
self.calls: list[tuple[str, Any]] = []
async def submit_update(self, request_id, req_state, request, *, prompt_text=None) -> int:
self.calls.append((request_id, prompt_text))
return 0
pool = RecordingPool()
orchestrator = object.__new__(Orchestrator)
orchestrator.async_chunk = False
orchestrator.request_states = {}
orchestrator.stage_pools = [pool]
add_request_calls: list[str] = []
async def record_add_request(msg) -> None:
add_request_calls.append(msg.request_id)
orchestrator._handle_add_request = record_add_request
await orchestrator._handle_streaming_update(
StageSubmissionMessage(
type="streaming_update",
request_id="req-unknown",
prompt=SimpleNamespace(request_id="req-unknown", prompt_token_ids=[1], resumable=True),
original_prompt={"prompt": "segment-2"},
output_prompt_text="segment-2",
sampling_params_list=[_sampling_params()],
final_stage_id=0,
preprocess_ms=0.0,
request_timestamp=time.time(),
enqueue_ts=time.perf_counter(),
)
)
assert pool.calls == []
assert add_request_calls == []
assert "req-unknown" not in orchestrator.request_states
async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> None:
class RecordingPool:
def __init__(self) -> None:
@@ -7,15 +7,24 @@ import asyncio
import queue
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock
from unittest.mock import ANY, AsyncMock, MagicMock
import janus
import pytest
from vllm.outputs import CompletionOutput, RequestOutput
from vllm.sampling_params import SamplingParams
from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
from vllm_omni.engine.orchestrator import (
Orchestrator,
OrchestratorRequestState,
_OrchestratorDuplexStagePort,
)
from vllm_omni.engine.stage_pool import StagePool
from vllm_omni.experimental.fullduplex.engine.contracts import (
DuplexStageRequestContext,
DuplexStageSubmission,
)
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -107,11 +116,50 @@ class FakePrewarmPool:
async def submit_initial(self, _request_id, _req_state, request, prompt_text=None):
self.submitted.append(request)
return 0
def get_bound_replica_id(self, _request_id):
return 0
def _duplex_stage_port_submission():
stage_pools = []
for stage_id in range(3):
pool = SimpleNamespace(
stage_client=SimpleNamespace(default_sampling_params=SamplingParams(max_tokens=1)),
stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
submit_initial=AsyncMock(return_value=stage_id + 10),
submit_update=AsyncMock(return_value=stage_id + 20),
)
stage_pools.append(pool)
request_states: dict[str, OrchestratorRequestState] = {}
prewarm = AsyncMock()
port = _OrchestratorDuplexStagePort(
stage_pools=stage_pools,
request_states=request_states,
running_counter=None,
cleanup_request_ids=AsyncMock(),
async_chunk=True,
prewarm_async_chunk_stages=prewarm,
)
context = DuplexStageRequestContext(
request_id="req-duplex",
session_id="session-duplex",
fence=DuplexFence("session-duplex"),
stage_id=0,
final_stage_id=2,
config_generation=0,
sampling_params=tuple(SamplingParams(max_tokens=1) for _ in range(3)),
)
port.ensure_request(context)
submission = DuplexStageSubmission(
context=context,
prompt={"prompt_token_ids": [1, 2]},
already_submitted=False,
)
return port, stage_pools, request_states, prewarm, submission
def _request_output(request_id: str) -> RequestOutput:
completion = CompletionOutput(
index=0,
@@ -199,11 +247,13 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None:
stage2 = FakePrewarmPool("receiver")
orchestrator.stage_pools = [stage0, stage1, stage2]
orchestrator._emit_tx_edge = lambda **_kwargs: None
orchestrator._record_duplex_stage_submission = MagicMock()
req_state = OrchestratorRequestState(
request_id="req-prewarm",
prompt={"prompt_token_ids": [1, 2]},
sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)],
final_stage_id=2,
duplex_identity=SimpleNamespace(),
)
await orchestrator._prewarm_async_chunk_stages(
@@ -216,6 +266,23 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None:
assert len(stage2.submitted) == 1
assert 1 not in req_state.stage_submit_ts
assert 2 in req_state.stage_submit_ts
orchestrator._record_duplex_stage_submission.assert_called_once_with(
2,
"req-prewarm",
0,
req_state,
)
@pytest.mark.asyncio
async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None:
port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission()
result = await port.submit(submission)
assert result.stage_id == 0
stage_pools[0].submit_initial.assert_awaited_once()
prewarm.assert_awaited_once_with("req-duplex", ANY, request_states["req-duplex"])
@pytest.mark.asyncio
+122 -1
View File
@@ -14,7 +14,7 @@ from vllm.v1.engine import FinishReason
from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
from vllm_omni.outputs import output_processor
from vllm_omni.outputs.output_modality import OutputModalityNames
from vllm_omni.outputs.output_modality import OutputModality, OutputModalityNames
from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -452,6 +452,127 @@ def test_no_detokenizer_make_request_output():
assert AUDIO in result.outputs[0].multimodal_output
def test_no_detokenizer_process_outputs_returns_nonterminal_audio_chunk(monkeypatch):
"""Generation-stage MM-only chunks must reach the orchestrator.
StagePool registers generation requests without a per-request collector,
so the processor return value is the only path from Stage2 to
Orchestrator._handle_processed_outputs().
"""
state = _make_no_detok_state(RequestOutputKind.DELTA)
processor = object.__new__(MultimodalOutputProcessor)
processor.output_modality = OutputModality.AUDIO
processor.request_states = {"r": state}
upstream_result = SimpleNamespace(request_outputs=[], reqs_to_abort=[])
monkeypatch.setattr(
VLLMOutputProcessor,
"process_outputs",
lambda *_args, **_kwargs: upstream_result,
)
engine_output = SimpleNamespace(
request_id="r",
multimodal_output={
"model_outputs": torch.arange(8, dtype=torch.float32),
"sr": torch.tensor(24000, dtype=torch.int32),
},
output_type="audio",
pooling_output=None,
new_token_ids=[],
finish_reason=None,
stop_reason=None,
kv_transfer_params=None,
routed_experts=None,
num_cached_tokens=0,
)
result = processor.process_outputs([engine_output])
assert len(result.request_outputs) == 1
assert result.request_outputs[0].request_id == "r"
assert AUDIO in result.request_outputs[0].outputs[0].multimodal_output
def _make_mm_only_output_processor(monkeypatch):
processor = object.__new__(MultimodalOutputProcessor)
processor.output_modality = OutputModality.AUDIO
processor.request_states = {"r": _make_no_detok_state(RequestOutputKind.DELTA)}
monkeypatch.setattr(
VLLMOutputProcessor,
"process_outputs",
lambda *_args, **_kwargs: SimpleNamespace(
request_outputs=[],
reqs_to_abort=[],
),
)
def finish_request(req_state):
processor.request_states.pop(req_state.request_id, None)
processor._finish_request = finish_request
return processor
def _audio_engine_output(*, is_segment_finished: bool, is_last_chunk: bool):
return SimpleNamespace(
request_id="r",
multimodal_output={
"model_outputs": torch.arange(8, dtype=torch.float32),
"meta.tts_is_last_chunk": torch.tensor([int(is_last_chunk)]),
},
output_type="audio",
pooling_output=None,
new_token_ids=[],
finish_reason=FinishReason.STOP,
stop_reason=None,
kv_transfer_params=None,
routed_experts=None,
num_cached_tokens=0,
is_segment_finished=is_segment_finished,
)
def test_mm_only_segment_finish_retains_request_state_for_next_audio_chunk(monkeypatch):
processor = _make_mm_only_output_processor(monkeypatch)
first = processor.process_outputs([_audio_engine_output(is_segment_finished=True, is_last_chunk=False)])
assert "r" in processor.request_states
assert len(first.request_outputs) == 1
assert first.request_outputs[0].finished is False
second = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)])
assert len(second.request_outputs) == 1
assert second.request_outputs[0].finished is True
assert "r" not in processor.request_states
def test_mm_only_nonfinal_audio_retains_state_without_segment_marker(monkeypatch):
processor = _make_mm_only_output_processor(monkeypatch)
first = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=False)])
assert "r" in processor.request_states
assert len(first.request_outputs) == 1
assert first.request_outputs[0].finished is False
second = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)])
assert len(second.request_outputs) == 1
assert second.request_outputs[0].finished is True
assert "r" not in processor.request_states
def test_mm_only_terminal_finish_removes_request_state(monkeypatch):
processor = _make_mm_only_output_processor(monkeypatch)
result = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)])
assert len(result.request_outputs) == 1
assert result.request_outputs[0].finished is True
assert "r" not in processor.request_states
def test_no_detokenizer_make_request_output_with_routed_experts():
"""make_request_output accepts the routed_experts arg that the multimodal
output channel (_process_mm_only_outputs) passes positionally, and attaches
+49
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import os
import socket
from contextlib import contextmanager
from types import SimpleNamespace
from typing import Any
@@ -11,6 +12,7 @@ import pytest
from pytest_mock import MockerFixture
from vllm.v1.engine.utils import EngineZmqAddresses
from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
from vllm_omni.engine.stage_engine_startup import (
@@ -189,6 +191,34 @@ class TestOmniMasterServerAllocation:
output_port = int(alloc.output_bind_address.split(":")[-1])
assert len({handshake_port, input_port, output_port}) == 3
def test_route_ports_remain_os_reserved_until_master_releases_them(self):
server = OmniMasterServer(master_address="127.0.0.1", master_port=15005, stage_ids=[0])
alloc = server.get_allocation(0)
ports = [
int(alloc.handshake_bind_address.rsplit(":", 1)[-1]),
int(alloc.input_bind_address.rsplit(":", 1)[-1]),
int(alloc.output_bind_address.rsplit(":", 1)[-1]),
]
for port in ports:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
with pytest.raises(OSError):
probe.bind(("127.0.0.1", port))
server.release_route_port_reservations(0, handshake=True, data=False)
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
probe.bind(("127.0.0.1", ports[0]))
for port in ports[1:]:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
with pytest.raises(OSError):
probe.bind(("127.0.0.1", port))
server.stop()
for port in ports[1:]:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
probe.bind(("127.0.0.1", port))
def test_get_zmq_addresses_returns_bind_addresses(self):
server = OmniMasterServer(master_address="127.0.0.1", master_port=15002, stage_ids=[0])
alloc = server.get_allocation(0)
@@ -397,6 +427,25 @@ class TestSingleStageModeDetection:
assert engine.single_stage_mode is False
assert engine._single_stage_id_filter is None
def test_stage_configs_path_loads_duplex_runtime_config(self, mocker: MockerFixture):
duplex_session = DuplexSessionRuntimeConfig(max_sessions=2)
mocker.patch(
"vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config",
return_value=None,
)
load_deploy_config = mocker.patch(
"vllm_omni.engine.async_omni_engine.load_deploy_config",
return_value=SimpleNamespace(duplex_session=duplex_session),
)
engine = self._make_engine_no_thread(
mocker,
stage_configs_path="/fake/duplex.yaml",
)
load_deploy_config.assert_called_once_with("/fake/duplex.yaml")
assert engine.duplex_session_config is duplex_session
def test_single_stage_mode_without_stage_id_has_no_filter(self, mocker: MockerFixture):
engine = self._make_engine_no_thread(
mocker,
@@ -2520,6 +2520,52 @@ def test_duplex_auto_response_tts_scheduler_eos_fallback():
assert results[0]["abort_data_plane_request"] is True
def test_duplex_auto_response_rearms_tts_boundary_for_next_segment_in_same_turn():
data_plane = _test_data_plane()
request_id = "duplex-sid-next-segment-same-turn-e0-stage0"
session = DuplexSession(
session_id="sid-next-segment-same-turn",
config=DuplexSessionConfig(extra_body={"auto_response": True}),
)
session.bind_response_turn(0)
data_plane.begin_request(request_id)
first_segment = _duplex_tts_output(
request_id=request_id,
samples=24000,
finished=False,
tts_is_last_chunk=True,
turn_id=0,
)
first_results = _project_data_plane(
data_plane,
{"data_plane_outputs": [first_segment]},
session=session,
)
assert len(first_results) == 1
assert first_results[0]["abort_data_plane_request"] is True
data_plane.begin_request(request_id)
second_segment = _duplex_tts_output(
request_id=request_id,
samples=48000,
finished=False,
text="hello again",
tts_is_last_chunk=True,
turn_id=0,
)
second_results = _project_data_plane(
data_plane,
{"data_plane_outputs": [second_segment]},
session=session,
)
assert len(second_results) == 1
assert second_results[0]["audio_data"] == "wav-24000"
assert second_results[0]["abort_data_plane_request"] is True
def test_duplex_turn_end_is_not_swallowed_by_finished_segment_fallback():
data_plane = _test_data_plane()
request_id = "duplex-sid-turn-end-after-segment-e0-stage0"
@@ -6202,19 +6248,27 @@ async def test_minicpmo_native_auto_response_keeps_request_bound_for_segment_con
config_timeout_s=0.1,
idle_timeout_s=1,
)
ws = TimedWebSocket()
def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None:
if data.get("type") != "session.created":
return
session = handler._registry.get("sid-native-auto-continuation")
assert session is not None
session.capabilities.chunk_period_ms = 50
ws.put(
{
"type": "input_audio_buffer.append",
"audio": _pcm_f32_b64(800),
"format": "pcm_f32le",
"sample_rate_hz": 16000,
}
)
ws.put({"type": "input_audio_buffer.commit", "final": True})
ws = TimedWebSocket(on_send=on_send)
event = _native_session_create("sid-native-auto-continuation")
event["session"]["extra_body"]["auto_response"] = True
ws.put(event)
ws.put(
{
"type": "input_audio_buffer.append",
"audio": _pcm_f32_b64(16000),
"format": "pcm_f32le",
"sample_rate_hz": 16000,
}
)
ws.put({"type": "input_audio_buffer.commit", "final": True})
await handler.handle_session(ws)
@@ -6225,6 +6279,112 @@ async def test_minicpmo_native_auto_response_keeps_request_bound_for_segment_con
assert final is False
@pytest.mark.asyncio
async def test_minicpmo_native_auto_response_real_input_preempts_unsent_silence():
request_id = "duplex-sid-native-real-before-silence-e0-stage0"
class SequencedAppendEngine(FakeEngineClient):
def __init__(self) -> None:
control_result = {
"operation": "append",
"session_id": "sid-native-real-before-silence",
"ok": True,
"unsupported_count": 0,
"error_count": 0,
"stage_results": [
{
"result": {
"data_plane_append": True,
"request_id": request_id,
"response_stage_id": 1,
}
}
],
}
terminal_segment = _duplex_tts_output(
request_id=request_id,
samples=0,
finished=True,
tts_is_last_chunk=True,
token_ids=[151645],
)
super().__init__(
control_result=control_result,
collect_outputs=[[terminal_segment], [], []],
collect_delay_s=0.01,
)
self.append_sequence: list[str] = []
self.first_real_started = asyncio.Event()
self.second_real_started = asyncio.Event()
self.real_count = 0
async def append_duplex_input_async(self, session_id: str, **kwargs):
kwargs.pop("expected_epoch", None)
result = await super().append_duplex_input_async(session_id, **kwargs)
payload = kwargs.get("payload")
is_silence = (
isinstance(payload, dict)
and payload.get("audio") == OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO
)
if is_silence:
self.append_sequence.append("silence")
else:
self.real_count += 1
self.append_sequence.append(f"real-{self.real_count}")
if self.real_count == 1:
self.first_real_started.set()
elif self.real_count == 2:
self.second_real_started.set()
return result
session_created = asyncio.Event()
engine = SequencedAppendEngine()
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
config_timeout_s=0.1,
idle_timeout_s=1,
)
def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None:
if data.get("type") != "session.created":
return
session = handler._registry.get("sid-native-real-before-silence")
assert session is not None
session.capabilities.chunk_period_ms = 50
ws.put(
{
"type": "input_audio_buffer.append",
"audio": _pcm_f32_b64(800),
"format": "pcm_f32le",
"sample_rate_hz": 16000,
}
)
session_created.set()
ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0)
event = _native_session_create("sid-native-real-before-silence")
event["session"]["extra_body"]["auto_response"] = True
ws.put(event)
handler_task = asyncio.create_task(handler.handle_session(ws))
await asyncio.wait_for(session_created.wait(), timeout=1)
await asyncio.wait_for(engine.first_real_started.wait(), timeout=1)
await asyncio.sleep(0.02)
ws.put(
{
"type": "input_audio_buffer.append",
"audio": _pcm_f32_b64(800, value=0.07),
"format": "pcm_f32le",
"sample_rate_hz": 16000,
}
)
await asyncio.wait_for(engine.second_real_started.wait(), timeout=1)
ws.put({"type": "session.close"})
await asyncio.wait_for(handler_task, timeout=2)
assert engine.append_sequence[:2] == ["real-1", "real-2"]
@pytest.mark.asyncio
async def test_minicpmo_native_auto_response_real_input_waits_for_submitted_silence_tail():
request_id = "duplex-sid-native-silence-tail-e0-stage0"
@@ -6294,18 +6454,26 @@ async def test_minicpmo_native_auto_response_real_input_waits_for_submitted_sile
config_timeout_s=0.1,
idle_timeout_s=1,
)
ws = TimedWebSocket(receive_timeout_s=2.0)
def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None:
if data.get("type") != "session.created":
return
session = handler._registry.get("sid-native-silence-tail")
assert session is not None
session.capabilities.chunk_period_ms = 50
ws.put(
{
"type": "input_audio_buffer.append",
"audio": _pcm_f32_b64(800, value=0.05),
"format": "pcm_f32le",
"sample_rate_hz": 16000,
}
)
ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0)
event = _native_session_create("sid-native-silence-tail")
event["session"]["extra_body"]["auto_response"] = True
ws.put(event)
ws.put(
{
"type": "input_audio_buffer.append",
"audio": _pcm_f32_b64(16000, value=0.05),
"format": "pcm_f32le",
"sample_rate_hz": 16000,
}
)
handler_task = asyncio.create_task(handler.handle_session(ws))
await asyncio.wait_for(engine.silence_append_started.wait(), timeout=1)
@@ -6330,6 +6498,135 @@ async def test_minicpmo_native_auto_response_real_input_waits_for_submitted_sile
assert engine.append_sequence == ["real-1", "silence-started", "silence-returned", "real-2"]
@pytest.mark.asyncio
async def test_minicpmo_native_auto_response_preserves_silence_continuations_across_append_tails():
request_id = "duplex-sid-native-continuation-tails-e0-stage0"
class SequencedAppendEngine(FakeEngineClient):
def __init__(self) -> None:
control_result = {
"operation": "append",
"session_id": "sid-native-continuation-tails",
"ok": True,
"unsupported_count": 0,
"error_count": 0,
"stage_results": [
{
"result": {
"data_plane_append": True,
"request_id": request_id,
"response_stage_id": 1,
}
}
],
}
super().__init__(control_result=control_result)
self.real_append_started = asyncio.Event()
self.release_real_append = asyncio.Event()
self.first_silence_started = asyncio.Event()
self.release_first_silence = asyncio.Event()
self.second_silence_started = asyncio.Event()
self.release_collect = asyncio.Event()
self.silence_count = 0
async def append_duplex_input_async(self, session_id: str, **kwargs):
kwargs.pop("expected_epoch", None)
payload = kwargs.get("payload")
assert isinstance(payload, dict)
result = await super().append_duplex_input_async(session_id, **kwargs)
if payload.get("audio") != OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO:
self.real_append_started.set()
await self.release_real_append.wait()
return result
self.silence_count += 1
if self.silence_count == 1:
self.first_silence_started.set()
await self.release_first_silence.wait()
else:
self.second_silence_started.set()
return result
async def collect_duplex_data_plane_outputs_async(
self,
request_id: str,
*,
response_stage_id: int | None = None,
timeout: float | None = None,
) -> list[object]:
del timeout
self.collected.append((request_id, response_stage_id))
await self.release_collect.wait()
return []
session_created = asyncio.Event()
def on_send(_ws: TimedWebSocket, data: dict[str, Any]) -> None:
if data.get("type") == "session.created":
session_created.set()
engine = SequencedAppendEngine()
handler = OmniDuplexSessionHandler(
chat_service=FakeChatService(engine),
config_timeout_s=0.1,
idle_timeout_s=1,
)
ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0)
event = _native_session_create("sid-native-continuation-tails")
event["session"]["extra_body"]["auto_response"] = True
ws.put(event)
ws.put(
{
"type": "input_audio_buffer.append",
"audio": _pcm_f32_b64(16000, value=0.05),
"format": "pcm_f32le",
"sample_rate_hz": 16000,
}
)
handler_task = asyncio.create_task(handler.handle_session(ws))
await asyncio.wait_for(session_created.wait(), timeout=1)
await asyncio.wait_for(engine.real_append_started.wait(), timeout=1)
session = handler._registry.get("sid-native-continuation-tails")
assert session is not None
session.capabilities.chunk_period_ms = 50
session.bind_request(request_id)
response_id = session.begin_response(turn_id=session.turn_id)
native = handler._minicpmo_session_state(session)
scheduler = native.silence_continuation_scheduler
assert scheduler is not None
payload = handler._native_silence_unit_payload()
payload["duplex_turn_id"] = session.turn_id
kwargs = {
"request_id": request_id,
"owner_id": f"response:{response_id}",
"response_id": response_id,
"response_owned": True,
"expected_epoch": session.epoch,
"expected_incarnation": session.incarnation,
"expected_model_turn_id": None,
"send_json": ws.send_json,
}
assert await scheduler(payload, **kwargs) is True
next_schedule = asyncio.create_task(scheduler(payload, **kwargs))
await asyncio.sleep(0.05)
assert not next_schedule.done()
engine.release_real_append.set()
await asyncio.wait_for(engine.first_silence_started.wait(), timeout=1)
assert not next_schedule.done()
engine.release_first_silence.set()
assert await asyncio.wait_for(next_schedule, timeout=1) is True
await asyncio.wait_for(engine.second_silence_started.wait(), timeout=1)
ws.put({"type": "session.close"})
await asyncio.wait_for(handler_task, timeout=2)
assert engine.silence_count == 2
@pytest.mark.asyncio
async def test_minicpmo_native_duplex_uses_segment_text_metadata_for_transcript_cursor():
request_id = "duplex-sid-native-repeat-text-e0-stage0"
+6 -1
View File
@@ -222,6 +222,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
stage_cfg = _make_stage_cfg(0, stage_type="llm")
stage_cfg.engine_args["async_chunk"] = True
parallel_config = SimpleNamespace(
data_parallel_size_local=1,
data_parallel_rank=0,
@@ -241,7 +242,10 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
"vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
return_value=(None, None, None),
)
mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
mock_connector_spec = mocker.patch(
"vllm_omni.engine.stage_init_utils.get_stage_connector_spec",
return_value={},
)
mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
mocker.patch(
"vllm_omni.engine.stage_init_utils.build_vllm_config",
@@ -277,6 +281,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
assert kwargs["omni_stage_config"] is stage_cfg
assert kwargs["replica_id"] is None
assert "socket_ownership" not in kwargs
assert mock_connector_spec.call_args.kwargs["async_chunk"] is True
assert mock_manager_cls.call_count == 1
mgr_kwargs = mock_manager_cls.call_args.kwargs
@@ -207,3 +207,68 @@ def test_playback_worklet_waits_before_playing_and_rebuffers_after_underrun():
capture_output=True,
text=True,
)
def test_playback_worklet_fades_terminal_drain_to_zero():
node = shutil.which("node")
if node is None:
pytest.skip("node is required for the AudioWorklet regression test")
script = textwrap.dedent(
"""
const fs = require('fs');
const vm = require('vm');
global.sampleRate = 1000;
const messages = [];
global.AudioWorkletProcessor = class {
constructor() {
this.port = {
onmessage: null,
postMessage: (message) => messages.push(message),
};
}
};
let Processor = null;
global.registerProcessor = (_name, processor) => { Processor = processor; };
vm.runInThisContext(fs.readFileSync(process.argv[1], 'utf8'));
const processor = new Processor();
const pcm = new Int16Array(100);
pcm.fill(16384);
processor.handleMessage({
type: 'audio',
pcm,
responseId: 'response-terminal',
initialBufferMs: 400,
});
processor.handleMessage({
type: 'drain',
responseId: 'response-terminal',
});
const output = new Float32Array(100);
processor.process([], [[output]]);
const assert = (condition, message) => {
if (!condition) throw new Error(message);
};
assert(output.some((sample) => sample !== 0), 'terminal audio must still play');
assert(output[output.length - 1] === 0, 'terminal drain must fade the final sample to zero');
const tail = Array.from(output.slice(-5), Math.abs);
assert(
tail.every((sample, index) => index === 0 || sample <= tail[index - 1]),
'terminal drain fade must decrease monotonically',
);
assert(
messages.filter((message) => message.type === 'playback-drained').length === 1,
'terminal drain must be reported exactly once after playback',
);
"""
)
subprocess.run(
[node, "-e", script, str(STATIC_ROOT / "playback_worklet.js")],
check=True,
capture_output=True,
text=True,
)
+48
View File
@@ -43,6 +43,54 @@ def test_append_emits_one_model_unit_when_multiple_units_are_buffered():
assert buffer.pending_byte_count == 16_000 * 4
def test_speech_marker_does_not_leak_across_irregular_chunk_boundaries():
buffer = MiniCPMO45PcmAppendBuffer()
assert buffer.append(pcm_payload(15_000), chunk_period_ms=1_000) is None
speech_unit = buffer.append(pcm_payload(2_000, speech=False), chunk_period_ms=1_000)
silence_unit = buffer.append(pcm_payload(15_000, speech=False), chunk_period_ms=1_000)
assert speech_unit is not None
assert speech_unit["is_speech"] is True
assert silence_unit is not None
assert silence_unit["is_speech"] is False
def test_force_listen_marker_does_not_leak_across_irregular_chunk_boundaries():
buffer = MiniCPMO45PcmAppendBuffer()
forced = pcm_payload(15_000, speech=False)
forced["force_listen"] = True
assert buffer.append(forced, chunk_period_ms=1_000) is None
forced_unit = buffer.append(pcm_payload(2_000, speech=False), chunk_period_ms=1_000)
unforced_unit = buffer.append(pcm_payload(15_000, speech=False), chunk_period_ms=1_000)
assert forced_unit is not None
assert forced_unit["force_listen"] is True
assert unforced_unit is not None
assert unforced_unit["force_listen"] is False
def test_pcm_append_rollback_restores_per_span_speech_markers():
buffer = MiniCPMO45PcmAppendBuffer()
assert buffer.append(pcm_payload(15_000), chunk_period_ms=1_000) is None
reservation = buffer.prepare_append(
pcm_payload(2_000, speech=False),
operation_id="mixed-unit",
chunk_period_ms=1_000,
)
assert reservation is not None
reservation.rollback()
first = buffer.append(pcm_payload(15_000, speech=False), chunk_period_ms=1_000)
second = buffer.flush(chunk_period_ms=1_000)
assert first is not None
assert first["is_speech"] is True
assert second is not None
assert second["is_speech"] is False
def test_commit_without_speech_does_not_synthesize_terminal_audio():
buffer = MiniCPMO45PcmAppendBuffer()
buffer.append(pcm_payload(8_000, speech=False), chunk_period_ms=1_000)
+179
View File
@@ -0,0 +1,179 @@
"""Shared fixtures and argument builders for MiniCPM-o 4.5 duplex E2E tests."""
from __future__ import annotations
import hashlib
import uuid
import wave
from pathlib import Path
from types import SimpleNamespace
import pytest
from huggingface_hub import snapshot_download
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
MODEL = "openbmb/MiniCPM-o-4_5"
DEPLOY_CONFIG = modify_stage_config(
get_deploy_config_path("minicpmo_4_5_duplex.yaml"),
updates={
"base_config": get_deploy_config_path("minicpmo_4_5.yaml"),
"stages": {
0: {"kv_cache_memory_bytes": 6 * 1024 * 1024 * 1024},
1: {"kv_cache_memory_bytes": 512 * 1024 * 1024},
2: {"kv_cache_memory_bytes": 256 * 1024 * 1024},
},
},
)
CORE_DEPLOY_CONFIG = modify_stage_config(
DEPLOY_CONFIG,
updates={
"stages": {
0: {"enforce_eager": True},
1: {"enforce_eager": True},
}
},
)
ASSET_DIR = Path(__file__).resolve().parents[1] / "assets" / "minicpmo_4_5"
RESPONSE_REQUIRED_WAV = ASSET_DIR / "response_required_16k.wav"
RESPONSE_REQUIRED_SHA256 = "2e5fd4eb3ee434ce107ee3a0591fa624a33f7683c7462f45fe651c443c9af941"
SOFT_INTERRUPT_WAV = ASSET_DIR / "soft_interrupt_16k.wav"
SOFT_INTERRUPT_SHA256 = "cadae6d0ddc510310f16f8775d6379f30a5195369ccb54ff10a8e3f9a1f4a2ea"
REF_AUDIO_RELATIVE_PATH = Path("assets") / "HT_ref_audio.wav"
SERVER_PARAMS = [
pytest.param(
OmniServerParams(
model=MODEL,
stage_config_path=DEPLOY_CONFIG,
use_stage_cli=True,
server_args=["--trust-remote-code"],
),
id="three-stage-single-gpu",
)
]
CORE_SERVER_PARAMS = [
pytest.param(
OmniServerParams(
model=MODEL,
stage_config_path=CORE_DEPLOY_CONFIG,
use_stage_cli=True,
server_args=["--trust-remote-code"],
),
id="three-stage-single-gpu",
)
]
def validated_wav(path: Path, expected_sha256: str) -> Path:
actual_sha256 = hashlib.sha256(path.read_bytes()).hexdigest()
if actual_sha256 != expected_sha256:
raise AssertionError(
f"MiniCPM-o duplex fixture SHA256 mismatch: expected {expected_sha256}, got {actual_sha256}"
)
with wave.open(str(path), "rb") as wav_file:
actual_format = (
wav_file.getframerate(),
wav_file.getnchannels(),
wav_file.getsampwidth(),
wav_file.getcomptype(),
)
expected_format = (16_000, 1, 2, "NONE")
if actual_format != expected_format:
raise AssertionError(f"MiniCPM-o duplex fixture must be 16 kHz mono PCM16, got {actual_format}")
return path
def validated_input_wav() -> Path:
return validated_wav(RESPONSE_REQUIRED_WAV, RESPONSE_REQUIRED_SHA256)
def validated_soft_interrupt_wav() -> Path:
return validated_wav(SOFT_INTERRUPT_WAV, SOFT_INTERRUPT_SHA256)
def resolve_ref_audio(model_prefix: str) -> Path:
if model_prefix:
model_root = Path(model_prefix) / MODEL
else:
model_root = Path(snapshot_download(MODEL, local_files_only=True))
ref_audio = model_root / REF_AUDIO_RELATIVE_PATH
if not ref_audio.is_file():
raise FileNotFoundError(f"MiniCPM-o checkpoint ref audio is missing: {ref_audio}")
return ref_audio
def realtime_url(omni_server) -> str:
return f"ws://{omni_server.host}:{omni_server.port}/v1/realtime?duplex=1"
def demo_args(
*,
omni_server,
input_wav: Path,
ref_audio: Path,
output_dir: Path,
) -> SimpleNamespace:
return SimpleNamespace(
url=realtime_url(omni_server),
model=omni_server.model,
session_id=f"duplex-ci-single-{uuid.uuid4().hex}",
input_wav=str(input_wav),
ref_audio=str(ref_audio),
turn_input_wav=[],
output_dir=str(output_dir),
output_audio_format="pcm16",
chunk_ms=200,
realtime_input=True,
first_turn_ms=1400,
turn_duration_ms=[],
first_turn_transcript="duplex CI speech",
omit_transcript_hints=False,
validation_mode="response-required",
temperature=0.0,
scenario="sequential",
require_audio=True,
require_distinct_inputs=False,
expect_empty_turn=[],
short_ack_ms=350,
turns=1,
timeout_s=120.0,
model_policy_settle_ms=2000,
)
def multi_session_args(
*,
omni_server,
input_wav: Path,
ref_audio: Path,
output_dir: Path,
response_required: bool,
) -> SimpleNamespace:
return SimpleNamespace(
url=realtime_url(omni_server),
model=omni_server.model,
sessions=2,
input_wav=str(input_wav),
ref_audio=str(ref_audio),
session_input_wav=[],
session_expected_token=[],
turn_input_wav=[],
output_dir=str(output_dir),
realtime_input=True,
chunk_ms=200,
turns=1,
first_turn_ms=1400,
turn_duration_ms=[],
response_required=response_required,
temperature=0.0 if response_required else None,
disconnect_session_index=0,
takeover_session_index=1,
resume_after_ms=100,
expire_session_index=None,
expire_after_s=40.0,
verify_admission_limit=None,
model_policy_settle_ms=2000,
timeout_s=180.0,
)
@@ -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
@@ -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()
@@ -43,6 +43,27 @@ def _delta(*codes: int):
}
def _duplex_delta(
*codes: int,
epoch: int = 3,
turn_id: int = 7,
text: str = "segment",
turn_end: bool = False,
):
text_utf8 = torch.tensor(list(text.encode("utf-8")), dtype=torch.uint8)
return {
"codes": {"audio": torch.tensor(codes, dtype=torch.long).reshape(-1, 1)},
"meta": {
"finished": torch.tensor(False),
"native_duplex": torch.tensor(True),
"duplex_epoch": torch.tensor(epoch),
"duplex_turn_id": torch.tensor(turn_id),
"llm_output_text_utf8": text_utf8,
"turn_end": torch.tensor(turn_end),
},
}
def _codes(payload) -> list[int]:
assert payload.codes is not None
assert isinstance(payload.codes.audio, torch.Tensor)
@@ -108,29 +129,57 @@ def test_short_final_flushes_silence_prefix_and_tail() -> None:
assert final.meta.finished.item() is True
def test_duplex_turn_end_waits_for_terminal_codec_flush() -> None:
manager = _manager()
request = _request("req-duplex")
body = tts2code2wav_async_chunk(
manager,
_duplex_delta(*range(25), turn_end=True),
request,
False,
)
final = tts2code2wav_async_chunk(
manager,
_duplex_delta(turn_end=True),
request,
True,
)
assert body is not None
assert body.meta.last_chunk is False
assert body.meta.turn_end is False
assert final is not None
assert _codes(final) == [22, 23, 24]
assert final.meta.last_chunk is True
assert final.meta.turn_end is True
def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> 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},
"meta": {"ref_audio_sr": 16000},
}
payload = tts2code2wav_async_chunk(manager, _delta(*range(7)), request, True)
payload = tts2code2wav_async_chunk(
manager,
_duplex_delta(*range(7), text="hello", turn_end=True),
request,
True,
)
assert payload is not None
assert payload.codes.ref.tolist() == pytest.approx([0.1, -0.1])
assert payload.meta.ref_audio_sr == 16000
assert payload.meta.native_duplex_segment_text == "hello"
torch.testing.assert_close(
payload.meta.llm_output_text_utf8,
torch.tensor(list(b"hello"), dtype=torch.uint8),
)
assert payload.meta.duplex_epoch == 3
assert payload.meta.duplex_turn_id == 7
assert payload.meta.segment_end is True
assert payload.meta.tts_is_last_chunk is True
assert payload.meta.turn_end is True
@@ -211,6 +260,186 @@ def test_empty_final_releases_wait_gate_once() -> None:
assert duplicate is None
def test_empty_duplex_boundary_uses_zero_length_transport_placeholder() -> None:
manager = _manager()
boundary = tts2code2wav_async_chunk(
manager,
_duplex_delta(text="boundary"),
_request("req-duplex"),
True,
)
assert boundary is not None
assert _codes(boundary) == [0]
assert boundary.meta.code_flat_numel == 0
assert boundary.meta.last_chunk is False
assert boundary.meta.is_segment_finished.item() is False
torch.testing.assert_close(
boundary.meta.llm_output_text_utf8,
torch.tensor(list(b"boundary"), dtype=torch.uint8),
)
def test_duplex_segments_preserve_stream_state_without_closing_turn() -> None:
manager = _manager()
request = _request("req-duplex")
request.additional_information = {
"codes": {"ref": [0.1, 0.2, 0.3]},
"meta": {"ref_audio_sr": 16000},
}
first = tts2code2wav_async_chunk(
manager,
_duplex_delta(10, 11, text="first"),
request,
True,
)
second = tts2code2wav_async_chunk(
manager,
_duplex_delta(12, 13, text="second"),
request,
True,
)
assert first is not None
assert second is not None
assert first.meta.last_chunk is False
assert second.meta.last_chunk is False
assert first.codes is not None
assert torch.allclose(first.codes.ref, torch.tensor([0.1, 0.2, 0.3]))
assert first.meta.ref_audio_sr == 16000
assert second.codes is not None
assert second.codes.ref is None
assert second.meta.cache_epoch == first.meta.cache_epoch
assert second.meta.chunk_seq == first.meta.chunk_seq + 1
assert second.meta.duplex_epoch == 3
assert second.meta.duplex_turn_id == 7
torch.testing.assert_close(
second.meta.llm_output_text_utf8,
torch.tensor(list(b"second"), dtype=torch.uint8),
)
assert second.meta.tts_is_last_chunk is True
assert second.meta.turn_end is False
assert first.meta.is_segment_finished.item() is False
assert second.meta.is_segment_finished.item() is False
def test_duplex_short_units_wait_for_minimum_stream_body() -> None:
manager = _manager()
request = _request("req-duplex")
first = tts2code2wav_async_chunk(
manager,
_duplex_delta(10, 11, 12, text="first"),
request,
True,
)
second = tts2code2wav_async_chunk(
manager,
_duplex_delta(13, 14, text="first"),
request,
True,
)
assert first is not None
assert _codes(first) == [0]
assert first.meta.code_flat_numel == 0
assert first.meta.last_chunk is False
assert first.meta.tts_is_last_chunk is True
assert second is not None
assert _codes(second) == [4218, 4218, 4218, 10, 11, 12, 13, 14]
assert second.meta.code_flat_numel == 8
torch.testing.assert_close(
second.meta.llm_output_text_utf8,
torch.tensor(list(b"firstfirst"), dtype=torch.uint8),
)
def test_duplex_empty_finish_callback_does_not_replay_previous_text() -> None:
manager = _manager()
request = _request("req-duplex")
first = tts2code2wav_async_chunk(
manager,
_duplex_delta(*range(25), text="first"),
request,
False,
)
boundary = tts2code2wav_async_chunk(
manager,
_duplex_delta(text="first"),
request,
True,
)
second = tts2code2wav_async_chunk(
manager,
_duplex_delta(*range(25, 50), text="second"),
request,
False,
)
assert first is not None
assert boundary is not None
assert boundary.meta.code_flat_numel == 0
assert second is not None
torch.testing.assert_close(
second.meta.llm_output_text_utf8,
torch.tensor(list(b"second"), dtype=torch.uint8),
)
def test_duplex_short_tail_does_not_replay_previous_segment_text() -> None:
manager = _manager()
request = _request("req-duplex")
def chunk(codes, text: str, finished: bool):
return tts2code2wav_async_chunk(
manager,
_duplex_delta(*codes, text=text),
request,
finished,
)
first = chunk(range(25), "和上海之间", False)
first_tail = chunk([25, 26], "和上海之间", True)
assert first is not None
assert first.meta.llm_output_text_utf8.tolist() == list("和上海之间".encode())
assert first_tail is not None
assert first_tail.meta.code_flat_numel == 0
assert chunk([27, 28, 29], "的距离大约是", False) is None
next_flush = chunk([], "的距离大约是", True)
assert next_flush is not None
assert next_flush.meta.llm_output_text_utf8.tolist() == list("的距离大约是".encode())
def test_duplex_turn_end_closes_epoch_and_next_turn_restarts_sequence() -> None:
manager = _manager()
request = _request("req-duplex")
turn_end = tts2code2wav_async_chunk(
manager,
_duplex_delta(14, turn_id=7, turn_end=True),
request,
True,
)
next_turn = tts2code2wav_async_chunk(
manager,
_duplex_delta(20, turn_id=8),
request,
True,
)
assert turn_end is not None
assert next_turn is not None
assert turn_end.meta.last_chunk is True
assert turn_end.meta.turn_end is True
assert turn_end.meta.is_segment_finished.item() is True
assert next_turn.meta.cache_epoch == turn_end.meta.cache_epoch + 1
assert next_turn.meta.chunk_seq == 0
assert next_turn.meta.last_chunk is False
def test_staggered_requests_keep_accumulators_isolated() -> None:
manager = _manager()
req_a = _request("a")
@@ -128,13 +128,125 @@ def test_native_duplex_speak_segment_reaches_split_talker() -> None:
info = converted["model_intermediate_buffer"]
assert info["native_duplex"] is True
assert info["ids"]["tts"] == [21, 22]
assert converted["prompt_token_ids"] == [0, 0, 0, 0]
assert converted["prompt_token_ids"] == [0, 0, 0]
assert info["meta"]["replace_streaming_prompt"] is True
assert info["meta"]["next_stage_prompt_len"] == 3
assert info["meta"]["turn_start"] is True
assert info["meta"]["segment_end"] is True
assert info["duplex"]["epoch"] == 3
assert info["duplex"]["turn_id"] == 7
def test_native_duplex_continuation_appends_only_new_talker_condition() -> None:
prompt_ids = [101, 102]
token_ids = {
"tts_bos_token_id": 9301,
"tts_eos_token_id": 9302,
"listen_token_id": 9303,
"speak_token_id": 9304,
"chunk_eos_token_id": 9308,
"chunk_tts_eos_token_id": 9309,
"turn_eos_token_id": 9310,
}
context = SimpleNamespace(
bridge_states={
"duplex": {
"epoch": 3,
"model_turn_id": 7,
}
}
)
first_ids = [9304, 21, 22, 9308]
first = _output(
prompt_ids=prompt_ids,
output_ids=first_ids,
latent=torch.arange(24, dtype=torch.float32).reshape(6, 4),
multimodal_output={
"duplex_prompt_token_ids": prompt_ids,
"meta": token_ids,
},
)
second_ids = [*first_ids, 9304, 23, 24, 9308]
second = _output(
prompt_ids=prompt_ids,
output_ids=second_ids,
latent=torch.arange(40, dtype=torch.float32).reshape(10, 4),
multimodal_output={
"duplex_prompt_token_ids": prompt_ids,
"meta": token_ids,
},
)
third_ids = [*second_ids, 9304, 25, 26, 9308]
third = _output(
prompt_ids=prompt_ids,
output_ids=third_ids,
latent=torch.arange(56, dtype=torch.float32).reshape(14, 4),
multimodal_output={
"duplex_prompt_token_ids": prompt_ids,
"meta": token_ids,
},
)
first_input = llm2tts([first], prompt=[{}], _streaming_context=context)[0]
second_input = llm2tts([second], prompt=[{}], _streaming_context=context)[0]
context.bridge_states["duplex"]["model_turn_id"] = 8
third_input = llm2tts([third], prompt=[{}], _streaming_context=context)[0]
assert first_input["model_intermediate_buffer"]["ids"]["tts"] == [21, 22]
assert second_input["model_intermediate_buffer"]["ids"]["tts"] == [23, 24]
assert third_input["model_intermediate_buffer"]["ids"]["tts"] == [25, 26]
assert first_input["model_intermediate_buffer"]["meta"]["turn_start"] is True
assert second_input["model_intermediate_buffer"]["meta"]["turn_start"] is False
assert third_input["model_intermediate_buffer"]["meta"]["turn_start"] is True
assert first_input["model_intermediate_buffer"]["meta"]["replace_streaming_prompt"] is True
assert "replace_streaming_prompt" not in second_input["model_intermediate_buffer"]["meta"]
assert third_input["model_intermediate_buffer"]["meta"]["replace_streaming_prompt"] is True
assert second_input["model_intermediate_buffer"]["meta"]["next_stage_prompt_len"] == 3
assert second_input["prompt_token_ids"] == [0, 0, 0]
def test_native_duplex_transcript_decodes_the_talker_condition_slice() -> None:
prompt_ids = [101, 102]
metadata = {
"tts_bos_token_id": 9301,
"tts_eos_token_id": 9302,
"listen_token_id": 9303,
"speak_token_id": 9304,
"chunk_eos_token_id": 9308,
"chunk_tts_eos_token_id": 9309,
"turn_eos_token_id": 9310,
}
token_text = {21: "杭州", 22: "和", 23: "州和", 24: "上海之间", 25: "大"}
def output(token_ids: list[int]):
return _output(
prompt_ids=prompt_ids,
output_ids=token_ids,
latent=torch.zeros((len(prompt_ids) + len(token_ids), 1)),
multimodal_output={
"duplex_prompt_token_ids": prompt_ids,
"meta": metadata,
},
)
context = SimpleNamespace(
bridge_states={"duplex": {"epoch": 3, "model_turn_id": 7}},
source_token_decoder=lambda ids, **_: "".join(token_text.get(int(token_id), "") for token_id in ids),
)
first_ids = [9304, 21, 22, 9308]
first_info = llm2tts([output(first_ids)], prompt=[{}], _streaming_context=context)[0]["model_intermediate_buffer"]
second_info = llm2tts(
[output([*first_ids, 23, 24, 25, 9308])],
prompt=[{}],
_streaming_context=context,
)[0]["model_intermediate_buffer"]
assert first_info["meta"]["native_duplex_segment_text"] == "杭州和"
assert second_info["ids"]["tts"] == [24, 25]
assert second_info["meta"]["native_duplex_segment_text"] == "上海之间大"
def test_native_duplex_requires_tokenizer_boundary_metadata() -> None:
latent = torch.zeros((3, 4))
source = _output(
@@ -4,6 +4,7 @@
from __future__ import annotations
from collections import defaultdict
from types import SimpleNamespace
import pytest
@@ -227,6 +228,47 @@ def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
assert "code_predictor_codes" not in payload
def test_talker2code2wav_async_chunk_flushes_cached_tail_on_stop_token() -> None:
request_id = "codec_tail"
stop_token_id = 999
transfer_manager = SimpleNamespace(
code_prompt_token_ids=defaultdict(list),
put_req_chunk=defaultdict(int, {request_id: 2}),
connector=SimpleNamespace(
config={
"extra": {
"initial_codec_chunk_frames": 4,
"codec_chunk_frames": 25,
"codec_left_context_frames": 25,
}
}
),
)
transfer_manager.code_prompt_token_ids[request_id] = [
torch.tensor([[frame, frame + 100]], dtype=torch.long) for frame in range(50)
]
request = SimpleNamespace(
external_req_id=request_id,
sampling_params=SimpleNamespace(
stop_token_ids=[stop_token_id],
stop_token_id=None,
),
)
payload = q3.talker2code2wav_async_chunk(
transfer_manager,
{"codes": {"audio": torch.tensor([[stop_token_id, 0]], dtype=torch.long)}},
request,
is_finished=True,
)
assert payload is not None
assert payload.meta.finished.item() is True
# 4 initial frames and one 25-frame chunk were already emitted. The final
# payload contains 25 frames of left context plus the remaining 21 frames.
assert payload.codes.audio.numel() == (25 + 21) * 2
def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
"""Full-payload path drops the terminal thinker row before talker prefill."""
request = SimpleNamespace(
+2
View File
@@ -1160,8 +1160,10 @@ class TestDeployConfigLoading:
assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False, False]
assert all("Async" not in (stage.scheduler_cls or "") for stage in stages)
assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"]
assert deploy.stages[1].enforce_eager is False
assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 2048
assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 0
assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [1]
@pytest.mark.parametrize(
+270 -15
View File
@@ -276,16 +276,16 @@ def test_minicpmo_model_hook_same_speech_row_second_step_does_not_rearm_pending_
next_logits[0, 7] = 30.0
next_logits[0, 8] = -2.0
next_logits[0, 10] = 20.0
original_next_logits = next_logits.clone()
model.prepare_duplex_sampling(next_logits, SimpleNamespace(), (row,))
assert state.pending_speech_context is False
assert state.current_turn_ended is False
assert torch.isneginf(next_logits[0, 7])
assert next_logits[0, 8].item() == 30.0
assert torch.equal(next_logits, original_next_logits)
def test_minicpmo_model_hook_mid_turn_speech_redirects_listen_to_tts_bos():
def test_minicpmo_model_hook_mid_turn_speech_preserves_logits_for_post_sample_redirect():
state = SimpleNamespace(
current_turn_ended=False,
last_terminator_token=None,
@@ -296,12 +296,11 @@ def test_minicpmo_model_hook_mid_turn_speech_redirects_listen_to_tts_bos():
logits[0, 7] = 30.0
logits[0, 8] = -2.0
logits[0, 10] = 20.0
original_logits = logits.clone()
model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))
assert torch.isneginf(logits[0, 7])
assert logits[0, 8].item() == 30.0
assert logits[0, 10].item() == 20.0
assert torch.equal(logits, original_logits)
def test_minicpmo_model_hook_ignores_serving_new_user_turn_marker():
@@ -317,12 +316,12 @@ def test_minicpmo_model_hook_ignores_serving_new_user_turn_marker():
logits = torch.zeros((1, 16), dtype=torch.float32)
logits[0, 7] = 5.0
logits[0, 10] = 20.0
original_logits = logits.clone()
model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))
assert state.current_turn_ended is False
assert state.last_terminator_token is None
assert torch.isneginf(logits[0, 7])
assert logits[0, 8].item() == 5.0
assert torch.equal(logits, original_logits)
def test_generic_ar_runner_builds_typed_duplex_sampling_rows():
@@ -613,6 +612,180 @@ def test_minicpmo_stage0_rejects_unknown_special_token_fallbacks():
runtime._require_special_token_ids()
def _minicpmo45_tokenizer_stub():
token_ids = {
token: index
for index, token in enumerate(
(
"<unit>",
"</unit>",
"<|listen|>",
"<|speak|>",
"<|tts_bos|>",
"<|tts_eos|>",
"<|tts_pad|>",
"<|chunk_eos|>",
"<|chunk_tts_eos|>",
"<|turn_eos|>",
),
start=101,
)
}
return SimpleNamespace(
unk_token_id=0,
convert_tokens_to_ids=lambda token: token_ids.get(token, 0),
encode=lambda token, add_special_tokens=False: (
[token_ids[token]] if not add_special_tokens and token in token_ids else [0]
),
)
def _transformers_processor_stub(auto_processor):
class _AutoImageProcessor:
@staticmethod
def register(*_args, **_kwargs):
return None
return SimpleNamespace(
AutoImageProcessor=_AutoImageProcessor,
AutoProcessor=auto_processor,
)
def test_minicpmo_stage0_loads_processor_from_hf_id(monkeypatch):
import sys
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
MiniCPMO45Stage0DuplexRuntime,
)
model_id = "openbmb/MiniCPM-o-4_5"
processor = SimpleNamespace(tokenizer=_minicpmo45_tokenizer_stub())
load_calls = []
class _AutoImageProcessor:
@staticmethod
def register(config_class, image_processor_class):
del image_processor_class
if isinstance(config_class, str):
raise AttributeError("'str' object has no attribute '__module__'")
class _AutoProcessor:
@classmethod
def from_pretrained(cls, model_path, *, trust_remote_code):
load_calls.append((model_path, trust_remote_code))
_AutoImageProcessor.register("MiniCPMVImageProcessor", object)
return processor
monkeypatch.setitem(
sys.modules,
"transformers",
SimpleNamespace(
AutoImageProcessor=_AutoImageProcessor,
AutoProcessor=_AutoProcessor,
),
)
loaded = MiniCPMO45Stage0DuplexRuntime._load_processor_from_path(model_id)
assert loaded is processor
assert load_calls == [(model_id, True)]
with pytest.raises(AttributeError, match="__module__"):
_AutoImageProcessor.register("MiniCPMVImageProcessor", object)
def test_minicpmo_stage0_processor_load_preserves_original_exception(monkeypatch):
import sys
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
MiniCPMO45Stage0DuplexRuntime,
)
original = ValueError("invalid processor config")
class _AutoProcessor:
@classmethod
def from_pretrained(cls, model_path, *, trust_remote_code):
del cls, model_path, trust_remote_code
raise original
monkeypatch.setitem(
sys.modules,
"transformers",
_transformers_processor_stub(_AutoProcessor),
)
with pytest.raises(RuntimeError, match="Failed to load MiniCPM-o duplex processor") as exc_info:
MiniCPMO45Stage0DuplexRuntime._load_processor_from_path("openbmb/MiniCPM-o-4_5")
assert exc_info.value.__cause__ is original
def test_minicpmo_stage0_processor_requires_tokenizer(monkeypatch):
import sys
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
MiniCPMO45Stage0DuplexRuntime,
)
class _AutoProcessor:
@classmethod
def from_pretrained(cls, model_path, *, trust_remote_code):
del cls, model_path, trust_remote_code
return object()
monkeypatch.setitem(
sys.modules,
"transformers",
_transformers_processor_stub(_AutoProcessor),
)
with pytest.raises(RuntimeError, match="does not expose a tokenizer"):
MiniCPMO45Stage0DuplexRuntime._load_processor_from_path("openbmb/MiniCPM-o-4_5")
def test_minicpmo_stage0_loaded_processor_validates_special_tokens(monkeypatch):
import sys
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
MiniCPMO45Stage0DuplexRuntime,
)
processor = SimpleNamespace(tokenizer=_minicpmo45_tokenizer_stub())
class _AutoProcessor:
@classmethod
def from_pretrained(cls, model_path, *, trust_remote_code):
del cls, model_path, trust_remote_code
return processor
monkeypatch.setitem(
sys.modules,
"transformers",
_transformers_processor_stub(_AutoProcessor),
)
runtime = MiniCPMO45Stage0DuplexRuntime(
SimpleNamespace(),
model_path="openbmb/MiniCPM-o-4_5",
device="cpu",
)
assert runtime.processor is processor
assert set(runtime._special_token_ids()) == {
"unit_token_id",
"unit_end_token_id",
"listen_token_id",
"speak_token_id",
"tts_bos_token_id",
"tts_eos_token_id",
"tts_pad_token_id",
"chunk_eos_token_id",
"chunk_tts_eos_token_id",
"turn_eos_token_id",
}
def test_minicpmo_stage0_data_plane_prefill_matches_official_unit_format():
import torch
@@ -1150,6 +1323,88 @@ def test_minicpmo_stage0_data_plane_final_first_chunk_does_not_add_silence_unit(
)
def test_minicpmo_stage0_final_does_not_promote_first_chunk_alignment_tail_to_unit():
import torch
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
MiniCPMO45Stage0DuplexRuntime,
_MiniCPMO45Stage0SessionState,
)
class _StageModel(torch.nn.Module):
first_chunk_ms = 10
chunk_ms = 8
sample_rate = 1000
def __init__(self):
super().__init__()
self.embed = torch.nn.Embedding(256, 2)
def get_input_embeddings(self):
return self.embed
def get_audio_hidden_states(self, _data):
return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]
class _MelProcessor:
sample_rate = 1000
is_first = True
def get_config(self):
return {"effective_first_chunk_ms": 9}
class _Processor:
def __init__(self):
self._streaming_mel_processor = _MelProcessor()
def get_streaming_chunk_size(self):
return 9 if self._streaming_mel_processor.is_first else 8
def process_audio_streaming(self, audio, **_kwargs):
self._streaming_mel_processor.is_first = False
return {"audio_features": audio, "audio_feature_lens": [[len(audio)]]}
runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
runtime.stage_model = _StageModel()
runtime.thinker = runtime.stage_model
runtime.tokenizer = SimpleNamespace(
unk_token_id=0,
convert_tokens_to_ids=lambda token: {
"<unit>": 1,
"</unit>": 2,
"<|listen|>": 3,
"<|speak|>": 4,
"<|tts_bos|>": 5,
"<|tts_eos|>": 6,
"<|tts_pad|>": 7,
"<|chunk_eos|>": 8,
"<|chunk_tts_eos|>": 9,
"<|turn_eos|>": 10,
"<|audio|>": 11,
}.get(token, 0),
encode=lambda text, add_special_tokens=False: [],
)
runtime.processor = _Processor()
runtime.device = "cpu"
runtime._init_token_ids()
state = _MiniCPMO45Stage0SessionState(session_id="sid-first-chunk-alignment-tail")
first = runtime._stage_prefill_embeddings_only(state, np.zeros(8, dtype=np.float32), seq=1)
second = runtime._stage_prefill_embeddings_only(state, np.zeros(8, dtype=np.float32), seq=2)
final = runtime._stage_prefill_embeddings_only(
state,
np.zeros(8, dtype=np.float32),
seq=3,
final=True,
)
assert first["input_token_ids"].count(1) == 1
assert second["input_token_ids"].count(1) == 1
assert final["input_token_ids"].count(1) == 1
assert state.audio_chunk_idx == 3
assert len(state.audio_buffer) == 1
def test_minicpmo_stage0_runtime_uses_loaded_vllm_embed_tokens_when_get_input_embeddings_is_broken():
import torch
@@ -1319,7 +1574,7 @@ def test_minicpmo_stage0_native_sampler_penalizes_text_from_prior_chunk():
model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
session_key = ("sid-cross-chunk-repetition", 0)
state = _MiniCPMO45Stage0SessionState(session_id=session_key[0])
state.generated_text_tokens = [198] * 8
state.generated_tokens = [198] * 8
model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state})
model._minicpmo45_duplex_row_sessions = {0: session_key}
@@ -1381,7 +1636,7 @@ def test_minicpmo_stage0_native_sampler_matches_official_negative_logit_penalty(
state = _MiniCPMO45Stage0SessionState(session_id=session_key[0])
repeated = 198
alternative = 1234
state.generated_text_tokens = [repeated]
state.generated_tokens = [repeated]
model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state})
model._minicpmo45_duplex_row_sessions = {0: session_key}
@@ -1405,7 +1660,7 @@ def test_minicpmo_stage0_native_sampler_matches_official_negative_logit_penalty(
assert sampled.sampled_token_ids.tolist() == [[repeated]]
def test_minicpmo_stage0_records_only_bounded_text_repetition_history():
def test_minicpmo_stage0_records_bounded_model_policy_history():
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
_MiniCPMO45Stage0SessionState,
)
@@ -1432,15 +1687,15 @@ def test_minicpmo_stage0_records_only_bounded_text_repetition_history():
}
for sampled in range(1000, 1520):
model._record_minicpmo45_duplex_terminator(0, sampled, token_ids)
model._record_minicpmo45_duplex_generation_token(0, sampled)
expected_history = list(range(1008, 1520))
assert state.generated_text_tokens == expected_history
assert state.generated_tokens == expected_history
for sampled in token_ids.values():
model._record_minicpmo45_duplex_terminator(0, sampled, token_ids)
model._record_minicpmo45_duplex_generation_token(0, sampled)
assert state.generated_text_tokens == expected_history
assert state.generated_tokens == [*expected_history[len(token_ids) :], *token_ids.values()]
def test_minicpmo_stage0_native_sampler_does_not_override_model_at_punctuation():
+11 -3
View File
@@ -731,9 +731,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
if self.log_stats:
session.record_event(EngineCoreEventType.QUEUED)
return
update_info = getattr(update, "additional_information", None)
update_meta = update_info.get("meta") if isinstance(update_info, dict) else None
if isinstance(update_meta, dict) and update_meta.get("replace_streaming_prompt") is True:
update_infos = (
getattr(update, "model_intermediate_buffer", None),
getattr(update, "additional_information", None),
)
replace_streaming_prompt = any(
isinstance(info, dict)
and isinstance(info.get("meta"), dict)
and info["meta"].get("replace_streaming_prompt") is True
for info in update_infos
)
if replace_streaming_prompt:
self._replace_streaming_session(session, update)
return
super()._update_request_as_session(session, update)
@@ -54,6 +54,21 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
)
self._latest_omni_connector_output: OmniConnectorOutput | None = None
def _handle_stopped_request(self, request: Request) -> bool:
if (
request.resumable
and not request.streaming_queue
and self.chunk_transfer_adapter is not None
and self.chunk_transfer_adapter.receives_chunks
):
# Downstream async-chunk stages receive the next segment from the
# connector, not from an API StreamingUpdate. Enqueue them as
# schedulable before the base class can park them in skipped_waiting.
request.status = RequestStatus.WAITING
self._enqueue_waiting_request(request)
return False
return super()._handle_stopped_request(request)
def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
"""One-shot generation fast path:
- Feed all input tokens of the request at once
@@ -63,6 +63,11 @@ class OmniSchedulerMixin:
session.num_computed_tokens = 0
session.prompt_token_ids = new_prompt
session.additional_information = update.additional_information or None
session.model_intermediate_buffer = getattr(
update,
"model_intermediate_buffer",
None,
)
session.update_block_hashes()
session.num_prompt_tokens = len(new_prompt)
session.arrival_time = update.arrival_time
+2 -1
View File
@@ -3,7 +3,7 @@
# The registered minicpmo_4_5 pipeline owns the fixed
# Thinker -> Talker -> Code2Wav topology. This overlay changes session
# lifetime and bounds all three stages to two live duplex sessions.
base_config: minicpmo_4_5_batching.yaml
base_config: minicpmo_4_5.yaml
pipeline: minicpmo_4_5
session_mode: duplex
active_stream_window: 1
@@ -31,6 +31,7 @@ stages:
# segment and turn boundaries travel as explicit output metadata.
default_sampling_params:
max_tokens: 2048
min_tokens: 0
stop_token_ids: [1]
- stage_id: 2
@@ -231,18 +231,18 @@ def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], re
ids = payload_data.get("ids", {})
meta = payload_data.get("meta", {})
replace_prompt = isinstance(meta, dict) and meta.get("replace_streaming_prompt") is True
condition_len = meta.get("next_stage_prompt_len") if replace_prompt else None
if isinstance(condition_len, int) and condition_len > 0:
next_stage_prompt_len = meta.get("next_stage_prompt_len") if isinstance(meta, dict) else None
if replace_prompt and isinstance(next_stage_prompt_len, int) and next_stage_prompt_len > 0:
# Some downstream stages consume complete, independently conditioned
# segments instead of extending an existing KV prefix. The producer
# declares that transport behavior explicitly in payload metadata.
new_prompt = [0] * condition_len
new_prompt = [0] * next_stage_prompt_len
request._output_token_ids.clear()
request._all_token_ids.clear()
request._all_token_ids.extend(new_prompt)
request.prompt_token_ids = new_prompt
request.num_computed_tokens = 0
request.num_prompt_tokens = condition_len
request.num_prompt_tokens = next_stage_prompt_len
request.update_block_hashes()
return
prompt_token_ids = ids.get("prompt", None)
@@ -255,7 +255,10 @@ def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], re
assert request.prompt_token_ids is not None
# Extend prompt with kept output tokens.
request.prompt_token_ids.extend(kept_output_tokens)
next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
if isinstance(next_stage_prompt_len, int) and next_stage_prompt_len > 0:
next_prompt_len = next_stage_prompt_len
else:
next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
new_prompt = [0] * next_prompt_len
request._all_token_ids.extend(new_prompt or ())
request.prompt_token_ids.extend(new_prompt or ())
@@ -283,7 +283,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
# Empty chunk with more data expected: keep polling.
has_new_ids = bool(new_ids.numel()) if use_tensor_codes else bool(new_ids)
if not has_new_ids and not payload_finished:
if not has_new_ids and payload_segment_finished:
# Preserve an explicit scheduler boundary even when it
# contains no new codec frames.
request.prompt_token_ids = [0]
if not has_new_ids and not payload_finished and not payload_segment_finished:
# The base recv loop treats False as "not ready yet" and
# requeues the request. Do not mark an empty non-terminal
# chunk as ready, otherwise Stage1 can consume before the
@@ -317,7 +321,13 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
multimodal_output=multimodal_output,
request=request,
# Existing processors use is_finished as a flush signal.
is_finished=is_segment_finished,
# Terminal stops no longer count as segment boundaries
# (is_segment_finished is False when the request finishes,
# see #5383), but the processor must still flush its
# accumulated tail on the terminal chunk — otherwise the
# downstream stage receives the finished marker without
# the final payload (#5413).
is_finished=is_segment_finished or is_finished,
)
except Exception as e:
@@ -332,7 +342,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
if payload_data.meta is None:
payload_data.meta = MetaStruct()
payload_data.meta.finished = torch.tensor(is_finished, dtype=torch.bool)
payload_data.meta.is_segment_finished = torch.tensor(is_segment_finished, dtype=torch.bool)
if payload_data.meta.is_segment_finished is None:
payload_data.meta.is_segment_finished = torch.tensor(is_segment_finished, dtype=torch.bool)
success, size, metadata = self.connector.put(
from_stage=str(stage_id),
@@ -725,6 +736,14 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
queue_snapshot = list(queue)
for request in queue_snapshot:
if not self._ensure_active_stream(request):
if target_status == RequestStatus.WAITING:
# A non-active placeholder must not remain visible to the
# scheduler: it has no connector payload yet, so running
# it would execute the downstream model with empty
# additional_information. Park it until restore_queues()
# and retry admission on the next scheduler tick.
queue.remove(request)
waiting_for_chunk_list.append(request)
continue
if request.status != RequestStatus.WAITING_FOR_CHUNK:
if request.request_id in self.requests_with_ready_chunks:
+4 -4
View File
@@ -209,10 +209,10 @@ class AsyncOmniEngine:
self._omni_master_port,
)
# Stage resolution pops deploy_config, so get the pipeline endpoint
# restriction beforehand. TODO (Alex) make this cleaner and refactor
# stage config resolution to remove kwargs hacks.
deploy_config_path = kwargs.get("deploy_config")
# Stage resolution pops deploy_config, so get pipeline-wide settings
# beforehand. The stage CLI exposes the same deploy YAML through
# stage_configs_path.
deploy_config_path = kwargs.get("deploy_config") or kwargs.get("stage_configs_path")
pipeline_config = StageConfigFactory.get_pipeline_config(
model=model,
trust_remote_code=trust_remote_code,
+31 -18
View File
@@ -13,7 +13,7 @@ from __future__ import annotations
import asyncio
import time as _time
from collections.abc import Callable
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
@@ -218,11 +218,18 @@ class _OrchestratorDuplexStagePort:
request_states: dict[str, OrchestratorRequestState],
running_counter: OmniRequestCounter | None,
cleanup_request_ids: Callable[..., Any],
async_chunk: bool,
prewarm_async_chunk_stages: Callable[
[str, Any, OrchestratorRequestState],
Awaitable[None],
],
) -> None:
self._stage_pools = stage_pools
self._request_states = request_states
self._running_counter = running_counter
self._cleanup_request_ids = cleanup_request_ids
self._async_chunk = async_chunk
self._prewarm_async_chunk_stages = prewarm_async_chunk_stages
@property
def stage_count(self) -> int:
@@ -299,6 +306,12 @@ class _OrchestratorDuplexStagePort:
replica_id = await pool.submit_update(context.request_id, request_state, request)
else:
replica_id = await pool.submit_initial(context.request_id, request_state, request, prompt_text=None)
if self._async_chunk and context.stage_id == 0:
await self._prewarm_async_chunk_stages(
context.request_id,
request,
request_state,
)
request_state.duplex_stage_fences[context.stage_id] = context.fence
request_state.stage_submit_ts[context.stage_id] = _time.time()
if not request_state.running_counter_registered and self._running_counter is not None:
@@ -388,6 +401,8 @@ class Orchestrator:
request_states=self.request_states,
running_counter=self._running_counter,
cleanup_request_ids=self._cleanup_request_ids,
async_chunk=self.async_chunk,
prewarm_async_chunk_stages=self._prewarm_async_chunk_stages,
),
result_sink=self.rpc_async_queue,
lifecycle_sink=self.output_async_queue,
@@ -681,24 +696,16 @@ class Orchestrator:
final_stage_id = msg.final_stage_id
req_state = self.request_states.get(request_id)
if req_state is None:
# Streaming updates always follow the first-chunk add_request
# through the same ordered submission queue, so an unknown id
# here can only mean the request already finished or was
# aborted (e.g. the client disconnected). Re-adding it would
# resurrect a headless session that keeps cycling through the
# stages with nobody consuming its outputs (issue #4271).
logger.warning(
"[Orchestrator] streaming_update for unknown req=%s, falling back to add_request",
"[Orchestrator] streaming_update for unknown req=%s; dropping (request finished or aborted)",
request_id,
)
fallback_msg = StageSubmissionMessage(
type="add_request",
request_id=msg.request_id,
prompt=msg.prompt,
original_prompt=msg.original_prompt,
output_prompt_text=msg.output_prompt_text,
sampling_params_list=msg.sampling_params_list,
final_stage_id=msg.final_stage_id,
final_output_stage_ids=msg.final_output_stage_ids,
preprocess_ms=msg.preprocess_ms,
request_timestamp=msg.request_timestamp,
enqueue_ts=msg.enqueue_ts,
)
await self._handle_add_request(fallback_msg)
return
if msg.sampling_params_list:
@@ -2025,7 +2032,7 @@ class Orchestrator:
_t_submit_start = _time.perf_counter()
if next_pool.stage_type == "diffusion":
await next_pool.submit_initial(
replica_id = await next_pool.submit_initial(
request_id,
req_state,
req_state.prompt,
@@ -2064,12 +2071,18 @@ class Orchestrator:
resumable=downstream_resumable,
)
request.external_req_id = request.request_id
await next_pool.submit_initial(
replica_id = await next_pool.submit_initial(
request_id,
req_state,
request,
prompt_text=None,
)
self._record_duplex_stage_submission(
next_stage_id,
request_id,
replica_id,
req_state,
)
# async_chunk pre-submit fires per stage edge (N-1 -> N). Source
# replica is stage 0's bound replica (single-replica thinker in
+175 -69
View File
@@ -183,6 +183,10 @@ class OmniMasterServer:
# seeded with the registration port (and coordinator ROUTER port) that
# are already bound on the same host.
self._allocated_ports: set[int] = set()
# Keep route ports bound until their ZMQ owners are ready to bind them.
# Recording only the integer leaves a race with unrelated startup
# consumers (notably NCCL's own ephemeral-port selection).
self._port_reservations: dict[int, socket.socket] = {}
if master_port:
self._allocated_ports.add(int(master_port))
# Coordinator ROUTER address echoed back in every registration reply
@@ -283,9 +287,11 @@ class OmniMasterServer:
``get_open_ports_list`` dedups only within its own call, so we redraw any
port that collides with one already recorded in ``self._allocated_ports``
(registration/coordinator ports plus every previously-allocated route
port) and register the winners before returning. Callers already hold the
relevant allocation context (``__init__`` is single-threaded;
registration holds ``self._alloc_lock``).
port). Each winner remains OS-bound until its ZMQ owner is ready, so
another startup subsystem cannot claim it in the allocation-to-bind
window. Callers already hold the relevant allocation context
(``__init__`` is single-threaded; registration holds
``self._alloc_lock``).
"""
picked: list[int] = []
# Bounded retry budget so a pathological host (near port exhaustion)
@@ -294,7 +300,17 @@ class OmniMasterServer:
for port in get_open_ports_list(count=count - len(picked)):
if port in self._allocated_ports:
continue
reservation = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
# Reserve on every local interface. The advertised address
# can be a routable interface IP, but wildcard binding also
# protects the route from local NCCL/socket allocations.
reservation.bind(("", port))
except OSError:
reservation.close()
continue
self._allocated_ports.add(port)
self._port_reservations[port] = reservation
picked.append(port)
if len(picked) == count:
return picked
@@ -303,6 +319,33 @@ class OmniMasterServer:
f"(host={self._address}); {len(self._allocated_ports)} already in use."
)
def release_route_port_reservations(
self,
stage_id: int,
replica_id: int = 0,
*,
handshake: bool = True,
data: bool = True,
) -> None:
"""Release reserved route ports immediately before their ZMQ binds."""
alloc = self.get_allocation(stage_id, replica_id)
ports: list[int] = []
if handshake:
handshake_port = _port_from_zmq_address(alloc.handshake_bind_address)
if handshake_port is not None:
ports.append(handshake_port)
if data:
for address in (alloc.input_bind_address, alloc.output_bind_address):
port = _port_from_zmq_address(address)
if port is not None:
ports.append(port)
with self._alloc_lock:
for port in ports:
reservation = self._port_reservations.pop(port, None)
if reservation is not None:
reservation.close()
def _next_free_replica_id(self, stage_id: int) -> int:
"""Return the next replica id to assign for an auto-assign registration.
@@ -437,6 +480,11 @@ class OmniMasterServer:
self._stop_event.set()
if self._thread is not None:
self._thread.join(timeout=10)
with self._alloc_lock:
reservations = list(self._port_reservations.values())
self._port_reservations.clear()
for reservation in reservations:
reservation.close()
# ------------------------------------------------------------------
# Internal server logic
@@ -771,20 +819,34 @@ def connect_remote_engine_cores(
handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address
with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
yield StageReplicaResources(
coordinator=coordinator,
addresses=addresses,
)
wait_for_engine_startup(
handshake_socket,
addresses,
engines_to_handshake,
vllm_config.parallel_config,
False, # coordinated_dp
vllm_config.cache_config,
None, # proc_manager (remote — no local procs)
None, # coord_process
omni_master_server.release_route_port_reservations(
stage_id,
replica_id,
handshake=True,
data=False,
)
try:
with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
yield StageReplicaResources(
coordinator=coordinator,
addresses=addresses,
)
wait_for_engine_startup(
handshake_socket,
addresses,
engines_to_handshake,
vllm_config.parallel_config,
False, # coordinated_dp
vllm_config.cache_config,
None, # proc_manager (remote — no local procs)
None, # coord_process
)
finally:
omni_master_server.release_route_port_reservations(
stage_id,
replica_id,
handshake=False,
data=True,
)
@@ -802,17 +864,31 @@ def connect_remote_diffusion_proc(
).handshake_bind_address
logger.info("Waiting for remote diffusion proc for stage %d replica %d", stage_id, replica_id)
with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
yield StageReplicaResources(addresses=addresses)
wait_for_engine_startup(
handshake_socket,
addresses,
[CoreEngine(index=0, local=False)],
_single_diffusion_parallel_config(local_client=False),
False,
None,
None,
None,
omni_master_server.release_route_port_reservations(
stage_id,
replica_id,
handshake=True,
data=False,
)
try:
with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
yield StageReplicaResources(addresses=addresses)
wait_for_engine_startup(
handshake_socket,
addresses,
[CoreEngine(index=0, local=False)],
_single_diffusion_parallel_config(local_client=False),
False,
None,
None,
None,
)
finally:
omni_master_server.release_route_port_reservations(
stage_id,
replica_id,
handshake=False,
data=True,
)
@@ -918,49 +994,63 @@ def _launch_omni_core_engines(
# Bind the pre-allocated handshake socket for this stage.
handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address
with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
if omni_coordinator_address is not None:
# Use the omni subclass so each spawned subprocess instantiates
# an OmniCoordClientForStage and heartbeats to the coordinator.
from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager
omni_master_server.release_route_port_reservations(
stage_id,
replica_id,
handshake=True,
data=False,
)
try:
with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
if omni_coordinator_address is not None:
# Use the omni subclass so each spawned subprocess instantiates
# an OmniCoordClientForStage and heartbeats to the coordinator.
from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager
with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock):
local_engine_manager: CoreEngineProcManager = StageEngineCoreProcManager(
local_engine_count=local_engine_count,
start_index=start_index,
local_start_index=local_start_index,
vllm_config=vllm_config,
local_client=True,
handshake_address=handshake_address,
executor_class=executor_class,
log_stats=log_stats,
omni_stage_id=stage_id,
omni_coordinator_address=omni_coordinator_address,
omni_replica_base_id=replica_id,
)
else:
with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock):
local_engine_manager = CoreEngineProcManager(
local_engine_count=local_engine_count,
start_index=start_index,
local_start_index=local_start_index,
vllm_config=vllm_config,
local_client=True,
handshake_address=handshake_address,
executor_class=executor_class,
log_stats=log_stats,
)
with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock):
local_engine_manager: CoreEngineProcManager = StageEngineCoreProcManager(
local_engine_count=local_engine_count,
start_index=start_index,
local_start_index=local_start_index,
vllm_config=vllm_config,
local_client=True,
handshake_address=handshake_address,
executor_class=executor_class,
log_stats=log_stats,
omni_stage_id=stage_id,
omni_coordinator_address=omni_coordinator_address,
omni_replica_base_id=replica_id,
)
else:
with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock):
local_engine_manager = CoreEngineProcManager(
local_engine_count=local_engine_count,
start_index=start_index,
local_start_index=local_start_index,
vllm_config=vllm_config,
local_client=True,
handshake_address=handshake_address,
executor_class=executor_class,
log_stats=log_stats,
)
yield local_engine_manager, coordinator, addresses
wait_for_engine_startup(
handshake_socket,
addresses,
engines_to_handshake,
parallel_config,
parallel_config.data_parallel_size > 1 and vllm_config.model_config.is_moe,
vllm_config.cache_config,
local_engine_manager,
coordinator.proc if coordinator else None,
yield local_engine_manager, coordinator, addresses
wait_for_engine_startup(
handshake_socket,
addresses,
engines_to_handshake,
parallel_config,
parallel_config.data_parallel_size > 1 and vllm_config.model_config.is_moe,
vllm_config.cache_config,
local_engine_manager,
coordinator.proc if coordinator else None,
)
finally:
omni_master_server.release_route_port_reservations(
stage_id,
replica_id,
handshake=False,
data=True,
)
@@ -1448,6 +1538,12 @@ def launch_diffusion_stage_replica(
omni_stage_config=stage_config,
replica_id=replica_id,
)
omni_master_server.release_route_port_reservations(
metadata.stage_id,
replica_id,
handshake=True,
data=False,
)
proc_manager = stage_diffusion_proc.StageDiffusionProcManager(
model=model,
od_config=od_config,
@@ -1461,6 +1557,12 @@ def launch_diffusion_stage_replica(
omni_stage_id=metadata.stage_id,
omni_replica_id=replica_id,
)
omni_master_server.release_route_port_reservations(
metadata.stage_id,
replica_id,
handshake=False,
data=True,
)
client = StageDiffusionClient.from_addresses(
metadata,
request_address=proc_manager.addresses.inputs[0],
@@ -1474,6 +1576,10 @@ def launch_diffusion_stage_replica(
lock_fds=lock_fds,
)
except Exception:
omni_master_server.release_route_port_reservations(
metadata.stage_id,
replica_id,
)
if proc_manager is not None:
proc_manager.shutdown()
if lock_fds:
+11 -1
View File
@@ -628,6 +628,13 @@ class AsyncOmni(EngineClient, OmniBase):
logger.debug(f"[AsyncOmni] Request {request_id} completed")
# The input pump can outlive a normally-completed generation
# (e.g. a realtime client that keeps streaming after the final
# output); a live pump keeps issuing streaming updates for a
# request the orchestrator no longer tracks.
if input_stream_task is not None and not input_stream_task.done():
input_stream_task.cancel()
self._log_summary_and_cleanup(request_id)
except (asyncio.CancelledError, GeneratorExit):
@@ -1059,7 +1066,10 @@ class AsyncOmni(EngineClient, OmniBase):
"""Submit request IDs to be aborted to the engine."""
await self.engine.abort_async(request_ids)
for rid in request_ids:
self.request_states.pop(rid, None)
state = self.request_states.pop(rid, None)
input_stream_task = getattr(state, "input_stream_task", None)
if input_stream_task is not None and not input_stream_task.done():
input_stream_task.cancel()
if self.log_stats:
logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids))
+1 -1
View File
@@ -892,7 +892,7 @@ def run_headless(args: TrackingNamespace) -> None:
stage_connector_spec = get_stage_connector_spec(
omni_transfer_config=omni_transfer_config,
stage_id=stage_id,
async_chunk=False,
async_chunk=bool(stage_cfg.engine_args.get("async_chunk", False)),
)
# ``runtime_cfg`` is mostly inherited from the parent's
@@ -138,6 +138,7 @@ class RealtimeConnection(VllmRealtimeConnection):
is_streaming=True,
)
result_gen = None
try:
result_gen = self.engine.generate(
prompt=streaming_input_gen,
@@ -196,6 +197,16 @@ class RealtimeConnection(VllmRealtimeConnection):
logger.exception("Error in generation: %s", e)
await self.send_error(str(e), "processing_error")
finally:
# Close the generator explicitly so AsyncOmni.generate's cleanup
# (input-pump cancellation and engine-side abort) runs now rather
# than whenever the event loop garbage-collects the async
# generator; the delay window is where a disconnected session
# keeps cycling through the stages (issue #4271).
if result_gen is not None:
try:
await result_gen.aclose()
except Exception:
logger.exception("Failed to close realtime result generator")
# Always send terminal event so clients don't hang forever.
if self._is_connected and not audio_done_sent:
try:
@@ -206,4 +217,10 @@ class RealtimeConnection(VllmRealtimeConnection):
self.audio_queue.get_nowait()
async def send_json(self, payload: dict):
await self.websocket.send_text(json.dumps(payload))
try:
await self.websocket.send_text(json.dumps(payload))
except Exception:
# A failed send means the client is gone; flag it so the
# generation loop stops instead of retrying into a dead socket.
self._is_connected = False
raise
@@ -65,12 +65,11 @@ class MiniCPMO45DataPlaneSession:
def begin_request(self, request_id: str) -> None:
state = self._requests.setdefault(request_id, _RequestState())
state.terminal = False
# Preserve turn-scoped cursors for resumable requests. The legacy path
# only reset request-scoped terminal markers on append.
request_turn = state.turns.get(None)
if request_turn is not None:
request_turn.tts_eos_done = False
request_turn.turn_eos_done = False
# A resumable append starts a new Talker segment, even when it
# continues the same model turn. Re-arm only the segment boundary;
# turn EOS remains deduplicated until the model advances the turn.
for turn_state in state.turns.values():
turn_state.tts_eos_done = False
def is_terminal(self, request_id: str | None) -> bool:
if request_id is None:
@@ -2,11 +2,19 @@ from __future__ import annotations
import base64
import binascii
from dataclasses import dataclass
from typing import Any
import numpy as np
@dataclass(frozen=True, slots=True)
class _PcmSpan:
byte_count: int
force_listen: bool
is_speech: bool
class MiniCPMO45PcmAppendReservation:
__slots__ = (
"_active",
@@ -15,6 +23,7 @@ class MiniCPMO45PcmAppendReservation:
"_owner",
"_raw",
"_sample_rate_hz",
"_spans",
"_turn_had_speech",
"_video_frames",
"operation_id",
@@ -31,6 +40,7 @@ class MiniCPMO45PcmAppendReservation:
sample_rate_hz: int,
force_listen: bool,
is_speech: bool,
spans: list[_PcmSpan] | None = None,
turn_had_speech: bool = False,
video_frames: list[str] | None = None,
) -> None:
@@ -41,6 +51,7 @@ class MiniCPMO45PcmAppendReservation:
self._sample_rate_hz = sample_rate_hz
self._force_listen = force_listen
self._is_speech = is_speech
self._spans = list(spans or [])
self._turn_had_speech = turn_had_speech
self._video_frames = list(video_frames or [])
self._active = True
@@ -97,9 +108,8 @@ class MiniCPMO45PcmAppendBuffer:
def __init__(self) -> None:
self._buffer = bytearray()
self._spans: list[_PcmSpan] = []
self._sample_rate_hz: int | None = None
self._force_listen = False
self._is_speech = False
self._turn_had_speech = False
self._reservation_seq = 0
self._reservations: list[MiniCPMO45PcmAppendReservation] = []
@@ -112,14 +122,13 @@ class MiniCPMO45PcmAppendBuffer:
reservation._active = False
self._reservations.clear()
self._buffer.clear()
self._spans.clear()
self._frame_queue.clear()
self._sample_rate_hz = None
self._force_listen = False
self._is_speech = False
self._turn_had_speech = False
def clear_force_listen(self) -> None:
self._force_listen = False
self._spans = [_PcmSpan(span.byte_count, False, span.is_speech) for span in self._spans]
def has_pending(self) -> bool:
return bool(self._buffer)
@@ -131,6 +140,56 @@ class MiniCPMO45PcmAppendBuffer:
def pending_byte_count(self) -> int:
return len(self._buffer)
def _append_span(self, span: _PcmSpan) -> None:
if span.byte_count <= 0:
return
if (
self._spans
and self._spans[-1].force_listen == span.force_listen
and self._spans[-1].is_speech == span.is_speech
):
previous = self._spans[-1]
self._spans[-1] = _PcmSpan(
previous.byte_count + span.byte_count,
span.force_listen,
span.is_speech,
)
return
self._spans.append(span)
def _consume_spans(self, byte_count: int) -> list[_PcmSpan]:
consumed: list[_PcmSpan] = []
remaining = byte_count
while remaining:
if not self._spans:
raise RuntimeError("MiniCPM-o PCM metadata is shorter than buffered audio")
span = self._spans.pop(0)
take = min(remaining, span.byte_count)
consumed.append(_PcmSpan(take, span.force_listen, span.is_speech))
if take < span.byte_count:
self._spans.insert(
0,
_PcmSpan(
span.byte_count - take,
span.force_listen,
span.is_speech,
),
)
remaining -= take
return consumed
def _prepend_spans(self, spans: list[_PcmSpan]) -> None:
current = self._spans
self._spans = []
for span in [*spans, *current]:
self._append_span(span)
def _buffer_force_listen(self) -> bool:
return any(span.force_listen for span in self._spans)
def _buffer_is_speech(self) -> bool:
return any(span.is_speech for span in self._spans)
def _reserve_passthrough(
self,
payload: dict[str, object],
@@ -180,12 +239,17 @@ class MiniCPMO45PcmAppendBuffer:
raise ValueError("MiniCPM-o native duplex audio append sample_rate_hz changed within a session")
self._sample_rate_hz = sample_rate_hz
self._buffer.extend(raw)
self._append_span(
_PcmSpan(
len(raw),
bool(payload.get("force_listen", False)),
bool(payload.get("is_speech", False)),
)
)
frames_in = payload.get("video_frames")
if isinstance(frames_in, list):
self._frame_queue.extend(frame for frame in frames_in if isinstance(frame, str) and frame)
self._turn_had_speech = self._turn_had_speech or bool(payload.get("is_speech", False))
self._force_listen = self._force_listen or bool(payload.get("force_listen", False))
self._is_speech = self._is_speech or bool(payload.get("is_speech", False))
if not allow_emit:
return None
@@ -208,6 +272,7 @@ class MiniCPMO45PcmAppendBuffer:
pad_samples = 0
emit_bytes = emit_samples * 4
reserved_raw = bytes(self._buffer[:emit_bytes])
reserved_spans = self._consume_spans(emit_bytes)
emit_raw = reserved_raw + b"\x00" * (pad_samples * 4)
del self._buffer[:emit_bytes]
@@ -227,11 +292,8 @@ class MiniCPMO45PcmAppendBuffer:
if emit_samples + pad_samples >= min_samples and self._frame_queue:
attached_frames = [self._frame_queue.pop(0)]
out["video_frames"] = attached_frames
out["force_listen"] = self._force_listen
out["is_speech"] = self._is_speech
if not self._buffer:
self._force_listen = False
self._is_speech = False
out["force_listen"] = any(span.force_listen for span in reserved_spans)
out["is_speech"] = any(span.is_speech for span in reserved_spans)
reservation = MiniCPMO45PcmAppendReservation(
owner=self,
operation_id=operation_id,
@@ -240,6 +302,7 @@ class MiniCPMO45PcmAppendBuffer:
sample_rate_hz=sample_rate_hz,
force_listen=bool(out.get("force_listen", False)),
is_speech=bool(out.get("is_speech", False)),
spans=reserved_spans,
turn_had_speech=self._turn_had_speech,
video_frames=attached_frames,
)
@@ -261,8 +324,8 @@ class MiniCPMO45PcmAppendBuffer:
"audio": "",
"format": "pcm_f32le",
"sample_rate_hz": self._sample_rate_hz or 16000,
"force_listen": self._force_listen,
"is_speech": self._is_speech,
"force_listen": self._buffer_force_listen(),
"is_speech": self._buffer_is_speech(),
}
reservation = self.prepare_append(
payload,
@@ -280,8 +343,8 @@ class MiniCPMO45PcmAppendBuffer:
payload=None,
raw=b"",
sample_rate_hz=self._sample_rate_hz or 0,
force_listen=self._force_listen,
is_speech=self._is_speech,
force_listen=self._buffer_force_listen(),
is_speech=self._buffer_is_speech(),
turn_had_speech=had_speech,
)
self._reservations.append(reservation)
@@ -289,8 +352,6 @@ class MiniCPMO45PcmAppendBuffer:
# Open the next physical input generation without touching already
# reserved appends. A rollback restores this generation's metadata.
self._sample_rate_hz = None
self._force_listen = False
self._is_speech = False
self._turn_had_speech = False
return reservation
@@ -335,12 +396,11 @@ class MiniCPMO45PcmAppendBuffer:
rolled_back = self._reservations[index:]
restored = b"".join(item._raw for item in rolled_back)
self._buffer[:0] = restored
self._prepend_spans([span for item in rolled_back for span in item._spans])
restored_frames = [frame for item in rolled_back for frame in item._video_frames]
if restored_frames:
self._frame_queue[:0] = restored_frames
self._sample_rate_hz = self._sample_rate_hz or reservation._sample_rate_hz
self._force_listen = self._force_listen or any(item._force_listen for item in rolled_back)
self._is_speech = self._is_speech or any(item._is_speech for item in rolled_back)
self._turn_had_speech = self._turn_had_speech or any(item._turn_had_speech for item in rolled_back)
for item in rolled_back:
item._active = False
@@ -354,8 +414,8 @@ class MiniCPMO45PcmAppendBuffer:
"audio": "",
"format": "pcm_f32le",
"sample_rate_hz": self._sample_rate_hz or 16000,
"force_listen": self._force_listen,
"is_speech": self._is_speech,
"force_listen": self._buffer_force_listen(),
"is_speech": self._buffer_is_speech(),
}
return self.append(payload, chunk_period_ms=chunk_period_ms, flush=True)
@@ -2,10 +2,10 @@ from __future__ import annotations
import base64
import copy
import sys
import time
from contextlib import suppress
from dataclasses import dataclass, field
from threading import Lock
from typing import Any
import numpy as np
@@ -14,6 +14,7 @@ from vllm_omni.experimental.fullduplex.minicpmo45.policy import MiniCPMO45Duplex
_MINICPMO45_SPECIAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS
_MINICPMO45_OPTIONAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS
_MINICPMO45_PROCESSOR_LOAD_LOCK = Lock()
@dataclass
@@ -35,7 +36,7 @@ class _MiniCPMO45Stage0SessionState:
pending_speech_context: bool = False
pending_speech_append_identity: tuple[int | None, int] | None = None
pending_speech_response_open: bool = False
generated_text_tokens: list[int] = field(default_factory=list)
generated_tokens: list[int] = field(default_factory=list)
class MiniCPMO45Stage0DuplexRuntime:
@@ -58,6 +59,8 @@ class MiniCPMO45Stage0DuplexRuntime:
else getattr(stage_model, "tokenizer", None)
)
self._init_token_ids()
if self.tokenizer is not None:
self._require_special_token_ids()
def _stage_runtime_ready(self) -> bool:
return self.processor is not None and self.tokenizer is not None and self.thinker is not None
@@ -196,19 +199,15 @@ class MiniCPMO45Stage0DuplexRuntime:
embed_parts.extend(state.context_embeds)
token_ids.extend(state.context_token_ids)
# Consume every complete chunk in the buffer so the appended span and
# the scheduler's slot reservation for this append agree exactly. A
# final append may zero-pad a real residual chunk, but it must not add
# a whole silence unit after all input was already consumed: official
# duplex generation runs once per microphone unit and client commit is
# not an additional model decision.
# Consume every complete processor chunk in the buffer so the appended
# span and the scheduler's slot reservation agree exactly. The serving
# PCM buffer pads a final real residual before it reaches Stage0. Do not
# pad again here: the first processor chunk is hop-aligned below 1035ms
# and intentionally leaves a small carry that is not another model unit.
units_built = 0
while True:
if len(state.audio_buffer) < chunk_size:
if not final or len(state.audio_buffer) == 0:
break
pad = np.zeros(chunk_size - len(state.audio_buffer), dtype=np.float32)
state.audio_buffer = np.concatenate([state.audio_buffer, pad])
break
audio_chunk = state.audio_buffer[:chunk_size]
batch_feature = self._process_streaming_audio(audio_chunk, state.audio_chunk_idx, processor=processor)
for name, value in (
@@ -224,6 +223,11 @@ class MiniCPMO45Stage0DuplexRuntime:
if units_built == 0:
return self._stage_prefill_result(False, start_time, "streaming audio embedding returned empty")
break
consumed_samples = self._consumed_audio_samples(
state.audio_chunk_idx,
chunk_size,
processor=processor,
)
if state.audio_chunk_idx > 0:
# Official duplex closes every unit (finalize_unit feeds the
# sampled terminator + </unit>) before the next <unit> opens.
@@ -255,9 +259,7 @@ class MiniCPMO45Stage0DuplexRuntime:
token_ids.extend(
[self._audio_embedding_placeholder_token_id()] * int(self._as_2d_tensor(audio_embeds).shape[0])
)
state.audio_buffer = state.audio_buffer[
self._consumed_audio_samples(state.audio_chunk_idx, chunk_size, processor=processor) :
]
state.audio_buffer = state.audio_buffer[consumed_samples:]
state.audio_chunk_idx += 1
units_built += 1
chunk_size = self._streaming_chunk_size(processor)
@@ -610,14 +612,33 @@ class MiniCPMO45Stage0DuplexRuntime:
def _load_processor_from_path(model_path: str | None) -> Any | None:
if not model_path:
return None
if model_path not in sys.path:
sys.path.insert(0, model_path)
try:
from processing_minicpmo import MiniCPMOProcessor
from transformers import AutoImageProcessor, AutoProcessor
return MiniCPMOProcessor.from_pretrained(model_path, trust_remote_code=True)
except Exception:
return None
original_register = AutoImageProcessor.register
def register_image_processor(config_class, *args, **kwargs):
# The checkpoint's auto_map already loads this class. Its
# legacy string registration is incompatible with some
# Transformers versions and is otherwise redundant.
if config_class == "MiniCPMVImageProcessor":
return None
return original_register(config_class, *args, **kwargs)
with _MINICPMO45_PROCESSOR_LOAD_LOCK:
AutoImageProcessor.register = staticmethod(register_image_processor)
try:
processor = AutoProcessor.from_pretrained(
model_path,
trust_remote_code=True,
)
finally:
AutoImageProcessor.register = staticmethod(original_register)
except Exception as exc:
raise RuntimeError(f"Failed to load MiniCPM-o duplex processor from {model_path!r}") from exc
if getattr(processor, "tokenizer", None) is None:
raise RuntimeError(f"MiniCPM-o duplex processor loaded from {model_path!r} does not expose a tokenizer")
return processor
def _init_token_ids(self) -> None:
if self.tokenizer is None:
@@ -503,11 +503,45 @@ class DuplexSessionRunnerMixin:
clear_completed_pending_silence()
pending_silence = native.pending_silence_task
if pending_silence is not None and not pending_silence.done():
if pending_silence is asyncio.current_task():
return False
try:
if not await pending_silence:
return False
except asyncio.CancelledError:
current = asyncio.current_task()
if current is not None and current.cancelling():
raise
return False
except Exception:
return False
clear_completed_pending_silence()
pending_silence = native.pending_silence_task
if pending_silence is not None and not pending_silence.done():
return False
append_tail = actor.native_append_tail
if (append_tail is None or append_tail.done()) and real_native_input_waiting():
return False
if actor.native_append_tail is not None and not actor.native_append_tail.done():
return False
if real_native_input_waiting():
return False
continuation_delay_s = max(
0.0,
float(session.capabilities.chunk_period_ms or 1000) / 1000.0,
)
if continuation_delay_s > 0:
await asyncio.sleep(continuation_delay_s)
if (
actor.native_append_tail is not append_tail
or ((append_tail is None or append_tail.done()) and real_native_input_waiting())
or self._native_silence_continuation_is_stale(
session,
request_id=request_id,
response_id=response_id,
response_owned=response_owned,
expected_epoch=expected_epoch,
expected_incarnation=expected_incarnation,
expected_model_turn_id=expected_model_turn_id,
)
):
return False
def _still_valid() -> bool:
return not real_native_input_waiting() and not self._native_silence_continuation_is_stale(
@@ -348,9 +348,16 @@ class BatchedToken2Wav(nn.Module):
previous: torch.Tensor,
window: torch.Tensor,
) -> torch.Tensor:
overlap = int(window.shape[0] // 2)
overlap = min(
int(window.shape[0] // 2),
int(speech.shape[-1]),
int(previous.shape[-1]),
)
result = speech.clone()
result[..., :overlap] = result[..., :overlap] * window[:overlap] + previous[..., -overlap:] * window[overlap:]
if overlap > 0:
result[..., :overlap] = (
result[..., :overlap] * window[:overlap] + previous[..., -overlap:] * window[-overlap:]
)
return result
def decode_batch(
@@ -360,6 +367,7 @@ class BatchedToken2Wav(nn.Module):
states: list[BatchedToken2WavState],
*,
last_chunk: bool,
flush_encoder: bool = False,
) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]:
batch_size = int(tokens.shape[0])
if batch_size != len(states):
@@ -382,7 +390,7 @@ class BatchedToken2Wav(nn.Module):
with self._autocast(tokens.device):
hidden, conformer_cnn, conformer_att = self._encode_chunk(
tokens,
last_chunk=last_chunk,
last_chunk=last_chunk or flush_encoder,
cnn_cache=flow_cache["conformer_cnn_cache"],
att_cache=flow_cache["conformer_att_cache"],
)
@@ -4,12 +4,12 @@
from __future__ import annotations
import hashlib
import json
import os
import tempfile
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from hashlib import sha256
from pathlib import Path
from typing import Any
@@ -77,6 +77,13 @@ class _RequestState:
token2wav: BatchedToken2WavState
@dataclass
class _RuntimePrompt:
cache_id: str
path: str
owners: set[str]
@dataclass(frozen=True)
class _WorkItem:
output_index: int
@@ -89,9 +96,11 @@ class _WorkItem:
last_chunk: bool
tokens: torch.Tensor
previous: _RequestState | None
segment_text: str
duplex_turn_id: int | None
duplex_epoch: int | None
runtime_prompt_key: str | None
duplex_epoch: int
duplex_turn_id: int
segment_text_utf8: torch.Tensor
tts_is_last_chunk: bool
segment_end: bool
turn_end: bool
has_payload: bool = True
@@ -120,12 +129,17 @@ class MiniCPMO45Code2Wav(nn.Module):
self.model_path = str(vllm_config.model_config.model)
self.backend: BatchedToken2Wav | None = None
self._states: dict[str, _RequestState] = {}
self._owned_prompt_wavs: dict[str, tuple[str, str]] = {}
self._runtime_prompts: dict[str, _RuntimePrompt] = {}
self._request_prompt_keys: dict[str, str] = {}
self._runtime_prompt_dir = tempfile.TemporaryDirectory(
prefix="minicpmo45-runtime-prompts-",
)
extra = self._extra_config()
self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1))
if self._min_batch_size < 1:
raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1")
self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio"))
self._prompt_wav_explicit = "prompt_wav" in extra
self._default_prompt_wav = str(
extra.get(
"prompt_wav",
@@ -133,6 +147,34 @@ class MiniCPMO45Code2Wav(nn.Module):
)
)
def _resolve_model_root(self) -> Path:
model_root = Path(self.model_path)
if model_root.is_dir():
return model_root
from vllm_omni.model_executor.model_loader.weight_utils import (
download_weights_from_hf_specific,
)
model_config = self.vllm_config.model_config
load_config = getattr(self.vllm_config, "load_config", None)
model_root = Path(
download_weights_from_hf_specific(
self.model_path,
getattr(load_config, "download_dir", None),
allow_patterns=[
"assets/HT_ref_audio.wav",
"assets/token2wav/*",
],
revision=getattr(model_config, "revision", None),
require_all=True,
)
)
self.model_path = str(model_root)
if not self._prompt_wav_explicit:
self._default_prompt_wav = str(model_root / "assets" / "HT_ref_audio.wav")
return model_root
def _extra_config(self) -> dict[str, Any]:
model_config = getattr(self.vllm_config, "model_config", None)
connector = getattr(model_config, "stage_connector_config", None)
@@ -142,63 +184,122 @@ class MiniCPMO45Code2Wav(nn.Module):
extra = getattr(connector, "extra", None)
return dict(extra) if isinstance(extra, Mapping) else {}
def _release_prompt(self, state_id: str) -> None:
owned = self._owned_prompt_wavs.pop(state_id, None)
if owned is None:
return
prompt_cache_id, prompt_wav = owned
if self.backend is not None:
self.backend.evict_prompt(prompt_cache_id, prompt_wav)
try:
os.unlink(prompt_wav)
except FileNotFoundError:
pass
def _write_reference_prompt(
self,
state_id: str,
reference: object,
sample_rate: object,
) -> tuple[str, str]:
waveform = torch.as_tensor(reference, dtype=torch.float32).reshape(-1).cpu()
if waveform.numel() == 0:
return self._default_prompt_id, self._default_prompt_wav
sample_rate_hz = int(_scalar(sample_rate, 24000))
if sample_rate_hz <= 0:
raise _batch_error(
"invalid_ref_audio_sample_rate",
request_id=state_id,
sample_rate=sample_rate_hz,
)
digest = hashlib.sha256()
digest.update(str(sample_rate_hz).encode("ascii"))
digest.update(waveform.numpy().tobytes())
prompt_cache_id = f"{state_id}:{digest.hexdigest()}"
current = self._owned_prompt_wavs.get(state_id)
if current is not None and current[0] == prompt_cache_id and os.path.exists(current[1]):
return current
self._release_prompt(state_id)
handle = tempfile.NamedTemporaryFile(prefix="minicpmo45-ref-", suffix=".wav", delete=False)
prompt_wav = handle.name
handle.close()
try:
sf.write(prompt_wav, waveform.numpy(), sample_rate_hz, format="WAV")
except Exception:
try:
os.unlink(prompt_wav)
except FileNotFoundError:
pass
raise
self._owned_prompt_wavs[state_id] = (prompt_cache_id, prompt_wav)
return prompt_cache_id, prompt_wav
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
return torch.zeros((input_ids.numel(), 1), device=input_ids.device, dtype=torch.float32)
def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None:
return None
def _materialize_runtime_prompt(
self,
ref_audio: Any,
sample_rate: Any,
) -> tuple[str, _RuntimePrompt]:
sample_rate_hz = int(_scalar(sample_rate, 0))
waveform = torch.as_tensor(ref_audio, dtype=torch.float32).reshape(-1).cpu().contiguous()
if sample_rate_hz <= 0:
raise _batch_error("invalid_ref_audio_sample_rate", sample_rate=sample_rate_hz)
if waveform.numel() == 0:
raise _batch_error("empty_ref_audio")
if not bool(torch.isfinite(waveform).all().item()):
raise _batch_error("non_finite_ref_audio")
digest = sha256()
digest.update(waveform.numpy().tobytes())
digest.update(str(sample_rate_hz).encode())
cache_key = digest.hexdigest()
cache_id = f"runtime-ref-{cache_key[:24]}-{sample_rate_hz}"
path = str(Path(self._runtime_prompt_dir.name) / f"minicpmo45_ref_{cache_key[:24]}_{sample_rate_hz}.wav")
entry = self._runtime_prompts.get(cache_key)
if entry is None:
entry = _RuntimePrompt(cache_id=cache_id, path=path, owners=set())
self._runtime_prompts[cache_key] = entry
prompt_path = Path(entry.path)
if not prompt_path.is_file():
with tempfile.NamedTemporaryFile(
dir=prompt_path.parent,
prefix=f".{prompt_path.stem}-",
suffix=".wav",
delete=False,
) as handle:
temporary_path = Path(handle.name)
try:
sf.write(
temporary_path,
waveform.numpy(),
sample_rate_hz,
format="WAV",
)
os.replace(temporary_path, prompt_path)
finally:
temporary_path.unlink(missing_ok=True)
return cache_key, entry
def _resolve_prompt(
self,
state_id: str,
info: Mapping[str, Any],
meta: Mapping[str, Any],
previous: _RequestState | None,
) -> tuple[str, str, str | None]:
codes = info.get("codes")
ref_audio = codes.get("ref") if isinstance(codes, Mapping) else None
if ref_audio is not None:
cache_key, entry = self._materialize_runtime_prompt(
ref_audio,
meta.get("ref_audio_sr"),
)
return entry.cache_id, entry.path, cache_key
if previous is not None:
return previous.prompt_cache_id, previous.prompt_wav, self._request_prompt_keys.get(state_id)
cache_key = self._request_prompt_keys.get(state_id)
entry = self._runtime_prompts.get(cache_key) if cache_key is not None else None
if entry is not None:
return entry.cache_id, entry.path, cache_key
return (
str(_scalar(meta.get("prompt_cache_id"), self._default_prompt_id)),
str(_scalar(meta.get("prompt_wav"), self._default_prompt_wav)),
None,
)
def _release_request_prompt(self, state_id: str) -> None:
cache_key = self._request_prompt_keys.pop(state_id, None)
entry = self._runtime_prompts.get(cache_key) if cache_key is not None else None
if entry is None:
return
entry.owners.discard(state_id)
if entry.owners:
return
if self.backend is not None:
self.backend.evict_prompt(entry.cache_id, entry.path)
Path(entry.path).unlink(missing_ok=True)
self._runtime_prompts.pop(cache_key, None)
def _commit_runtime_prompt_owners(self, items: list[_WorkItem]) -> None:
for item in items:
cache_key = item.runtime_prompt_key
if cache_key is None:
continue
previous_key = self._request_prompt_keys.get(item.state_id)
if previous_key != cache_key:
self._release_request_prompt(item.state_id)
entry = self._runtime_prompts.get(cache_key)
if entry is not None:
entry.owners.add(item.state_id)
self._request_prompt_keys[item.state_id] = cache_key
def _prune_unowned_runtime_prompts(self) -> None:
for cache_key, entry in list(self._runtime_prompts.items()):
if entry.owners:
continue
if self.backend is not None:
self.backend.evict_prompt(entry.cache_id, entry.path)
Path(entry.path).unlink(missing_ok=True)
self._runtime_prompts.pop(cache_key, None)
@staticmethod
def _split_segments(input_ids: torch.Tensor, counts: Any) -> list[torch.Tensor]:
flat = input_ids.reshape(-1)
@@ -247,9 +348,11 @@ class MiniCPMO45Code2Wav(nn.Module):
last_chunk=False,
tokens=segment.new_empty(0, dtype=torch.long),
previous=None,
segment_text="",
duplex_turn_id=None,
duplex_epoch=None,
runtime_prompt_key=None,
duplex_epoch=-1,
duplex_turn_id=-1,
segment_text_utf8=torch.empty(0, dtype=torch.uint8),
tts_is_last_chunk=False,
segment_end=False,
turn_end=False,
has_payload=False,
@@ -266,14 +369,15 @@ class MiniCPMO45Code2Wav(nn.Module):
chunk_seq=chunk_seq,
)
last_chunk = bool(_scalar(meta.get("last_chunk"), False))
tts_is_last_chunk = bool(_scalar(meta.get("tts_is_last_chunk"), False))
codes = info.get("codes")
audio = codes.get("audio") if isinstance(codes, Mapping) else None
reference = codes.get("ref") if isinstance(codes, Mapping) else None
tokens = _codec_tensor(audio, segment)
if last_chunk and int(_scalar(meta.get("code_flat_numel"), tokens.numel())) == 0:
if int(_scalar(meta.get("code_flat_numel"), tokens.numel())) == 0:
# The generation scheduler reserves one placeholder token for an
# empty terminal chunk. The producer's explicit length is the
# authority, so do not decode that placeholder as codec data.
# empty terminal or segment-boundary chunk. The producer's
# explicit length is the authority, so do not decode that
# placeholder as codec data.
tokens = segment.new_empty(0, dtype=torch.long)
previous = self._states.get(state_id)
if previous is None:
@@ -299,7 +403,6 @@ class MiniCPMO45Code2Wav(nn.Module):
cache_epoch=cache_epoch,
chunk_seq=chunk_seq,
)
self._release_prompt(state_id)
previous = None
elif chunk_seq != previous.chunk_seq + 1:
raise _batch_error(
@@ -308,22 +411,12 @@ class MiniCPMO45Code2Wav(nn.Module):
expected=previous.chunk_seq + 1,
actual=chunk_seq,
)
explicit_prompt_id = _scalar(meta.get("prompt_cache_id"))
explicit_prompt_wav = _scalar(meta.get("prompt_wav"))
if previous is not None and explicit_prompt_id is None and explicit_prompt_wav is None:
prompt_cache_id = previous.prompt_cache_id
prompt_wav = previous.prompt_wav
elif reference is not None:
prompt_cache_id, prompt_wav = self._write_reference_prompt(
state_id,
reference,
meta.get("ref_audio_sr"),
)
else:
prompt_cache_id = str(explicit_prompt_id or self._default_prompt_id)
prompt_wav = str(explicit_prompt_wav or self._default_prompt_wav)
prompt_cache_id, prompt_wav, runtime_prompt_key = self._resolve_prompt(
state_id,
info,
meta,
previous,
)
if previous is not None and prompt_cache_id != previous.prompt_cache_id:
raise _batch_error(
"prompt_changed_midstream",
@@ -338,8 +431,9 @@ class MiniCPMO45Code2Wav(nn.Module):
expected=previous.prompt_wav,
actual=prompt_wav,
)
duplex_turn_id = _scalar(meta.get("duplex_turn_id"))
duplex_epoch = _scalar(meta.get("duplex_epoch"))
segment_text_utf8 = meta.get("llm_output_text_utf8")
if not isinstance(segment_text_utf8, torch.Tensor):
segment_text_utf8 = torch.empty(0, dtype=torch.uint8)
return _WorkItem(
output_index=index,
state_id=state_id,
@@ -351,9 +445,11 @@ class MiniCPMO45Code2Wav(nn.Module):
last_chunk=last_chunk,
tokens=tokens,
previous=previous,
segment_text=str(_scalar(meta.get("native_duplex_segment_text"), "")),
duplex_turn_id=int(duplex_turn_id) if duplex_turn_id is not None else None,
duplex_epoch=int(duplex_epoch) if duplex_epoch is not None else None,
runtime_prompt_key=runtime_prompt_key,
duplex_epoch=int(_scalar(meta.get("duplex_epoch"), -1)),
duplex_turn_id=int(_scalar(meta.get("duplex_turn_id"), -1)),
segment_text_utf8=segment_text_utf8,
tts_is_last_chunk=tts_is_last_chunk,
segment_end=bool(_scalar(meta.get("segment_end"), False)),
turn_end=bool(_scalar(meta.get("turn_end"), False)),
)
@@ -371,6 +467,7 @@ class MiniCPMO45Code2Wav(nn.Module):
int(item.tokens.numel()),
cache_signature,
item.last_chunk,
item.tts_is_last_chunk,
item.cache_epoch,
)
@@ -435,26 +532,37 @@ class MiniCPMO45Code2Wav(nn.Module):
request_ids=len(state_ids),
)
items: list[_WorkItem] = []
for index, (state_id, segment, info) in enumerate(
zip(state_ids, segments, runtime_additional_information, strict=True)
):
if not isinstance(info, Mapping):
raise _batch_error(
"invalid_runtime_info",
output_index=index,
value_type=type(info).__name__,
)
items.append(self._parse_item(index, str(state_id), segment, info))
try:
for index, (state_id, segment, info) in enumerate(
zip(state_ids, segments, runtime_additional_information, strict=True)
):
if not isinstance(info, Mapping):
raise _batch_error(
"invalid_runtime_info",
output_index=index,
value_type=type(info).__name__,
)
items.append(self._parse_item(index, str(state_id), segment, info))
except Exception:
self._prune_unowned_runtime_prompts()
raise
state_ids = [item.state_id for item in items]
if len(state_ids) != len(set(state_ids)):
self._prune_unowned_runtime_prompts()
raise _batch_error("duplicate_request_in_forward", request_ids=state_ids)
outputs = [empty for _ in segments]
sentinels = [item for item in items if item.last_chunk and item.tokens.numel() == 0]
segment_markers = [
item for item in items if not item.last_chunk and item.tts_is_last_chunk and item.tokens.numel() == 0
]
compute_items = [item for item in items if item.tokens.numel() > 0]
invalid_empty = [
item.request_id for item in items if item.has_payload and not item.last_chunk and item.tokens.numel() == 0
item.request_id
for item in items
if item.has_payload and not item.last_chunk and not item.tts_is_last_chunk and item.tokens.numel() == 0
]
if invalid_empty:
self._prune_unowned_runtime_prompts()
raise _batch_error("empty_nonfinal_chunk", request_ids=invalid_empty)
buckets: dict[tuple[Any, ...], list[_WorkItem]] = {}
@@ -470,6 +578,7 @@ class MiniCPMO45Code2Wav(nn.Module):
if len(bucket) < self._min_batch_size
]
if undersized:
self._prune_unowned_runtime_prompts()
raise _batch_error(
"exact_shape_bucket_below_minimum",
minimum=self._min_batch_size,
@@ -477,6 +586,58 @@ class MiniCPMO45Code2Wav(nn.Module):
)
pending: dict[str, _RequestState | None] = {item.state_id: None for item in sentinels}
pending.update(
{
item.state_id: _RequestState(
cache_epoch=item.cache_epoch,
chunk_seq=item.chunk_seq,
prompt_cache_id=item.prompt_cache_id,
prompt_wav=item.prompt_wav,
token2wav=item.previous.token2wav,
)
for item in segment_markers
if item.previous is not None
}
)
initial_marker_buckets: dict[tuple[str, str], list[_WorkItem]] = {}
for item in segment_markers:
if item.previous is None:
initial_marker_buckets.setdefault(
(item.prompt_cache_id, item.prompt_wav),
[],
).append(item)
for bucket in initial_marker_buckets.values():
try:
features = self.backend.prepare_prompt(
bucket[0].prompt_cache_id,
bucket[0].prompt_wav,
)
states = self.backend.setup_batch(features, len(bucket))
except Exception as exc:
self._prune_unowned_runtime_prompts()
if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "):
raise
raise _batch_error(
"backend_unsupported_or_failed",
request_ids=[item.request_id for item in bucket],
error_type=type(exc).__name__,
error=str(exc),
) from exc
if len(states) != len(bucket):
self._prune_unowned_runtime_prompts()
raise _batch_error(
"backend_result_size_mismatch",
expected=len(bucket),
states=len(states),
)
for item, state in zip(bucket, states, strict=True):
pending[item.state_id] = _RequestState(
cache_epoch=item.cache_epoch,
chunk_seq=item.chunk_seq,
prompt_cache_id=item.prompt_cache_id,
prompt_wav=item.prompt_wav,
token2wav=state,
)
for bucket in buckets.values():
batch_size = len(bucket)
try:
@@ -496,6 +657,7 @@ class MiniCPMO45Code2Wav(nn.Module):
last_chunk=bucket[0].last_chunk,
)
except Exception as exc:
self._prune_unowned_runtime_prompts()
if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "):
raise
raise _batch_error(
@@ -505,6 +667,7 @@ class MiniCPMO45Code2Wav(nn.Module):
error=str(exc),
) from exc
if len(audios) != batch_size or len(next_states) != batch_size:
self._prune_unowned_runtime_prompts()
raise _batch_error(
"backend_result_size_mismatch",
expected=batch_size,
@@ -525,10 +688,10 @@ class MiniCPMO45Code2Wav(nn.Module):
)
)
self._commit_runtime_prompt_owners(items)
for request_id, state in pending.items():
if state is None:
self._states.pop(request_id, None)
self._release_prompt(request_id)
else:
self._states[request_id] = state
sample_rate_tensor = torch.as_tensor(sample_rate, dtype=torch.int32)
@@ -537,24 +700,15 @@ class MiniCPMO45Code2Wav(nn.Module):
multimodal_outputs={
"model_outputs": outputs,
"sr": [sample_rate_tensor.clone() for _ in outputs],
"meta.llm_output_text_utf8": [
torch.tensor(list(item.segment_text.encode("utf-8")), dtype=torch.uint8) for item in items
],
"meta.tts_is_last_chunk": [torch.tensor(item.last_chunk, dtype=torch.bool) for item in items],
# Generation runner wire payloads are flat and tensor-only.
# Dotted metadata keys are unflattened again by the output
# processor before the full-duplex data plane consumes them.
"meta.duplex_epoch": [torch.tensor(item.duplex_epoch, dtype=torch.int32) for item in items],
"meta.duplex_turn_id": [torch.tensor(item.duplex_turn_id, dtype=torch.int32) for item in items],
"meta.llm_output_text_utf8": [item.segment_text_utf8 for item in items],
"meta.tts_is_last_chunk": [torch.tensor(item.tts_is_last_chunk, dtype=torch.bool) for item in items],
"meta.segment_end": [torch.tensor(item.segment_end, dtype=torch.bool) for item in items],
"meta.turn_end": [torch.tensor(item.turn_end, dtype=torch.bool) for item in items],
"meta.duplex_turn_id": [
torch.tensor(item.duplex_turn_id, dtype=torch.int64)
if item.duplex_turn_id is not None
else torch.empty(0, dtype=torch.int64)
for item in items
],
"meta.duplex_epoch": [
torch.tensor(item.duplex_epoch, dtype=torch.int64)
if item.duplex_epoch is not None
else torch.empty(0, dtype=torch.int64)
for item in items
],
},
)
@@ -562,7 +716,7 @@ class MiniCPMO45Code2Wav(nn.Module):
for request_id in finished_req_ids:
state_id = str(request_id)
self._states.pop(state_id, None)
self._release_prompt(state_id)
self._release_request_prompt(state_id)
def make_omni_output(self, model_outputs: Any, **_: Any) -> OmniOutput:
if isinstance(model_outputs, OmniOutput):
@@ -600,10 +754,11 @@ class MiniCPMO45Code2Wav(nn.Module):
from stepaudio2.token2wav import Token2wav
extra = self._extra_config()
model_root = self._resolve_model_root()
prompt_path = Path(self._default_prompt_wav)
if not prompt_path.is_file():
raise FileNotFoundError(f"MiniCPM-o Code2Wav prompt audio not found: {prompt_path}")
token2wav_path = Path(self.model_path) / "assets" / "token2wav"
token2wav_path = model_root / "assets" / "token2wav"
if not token2wav_path.is_dir():
raise FileNotFoundError(f"MiniCPM-o Code2Wav assets not found: {token2wav_path}")
use_float16 = bool(extra.get("token2wav_float16", False))
@@ -175,7 +175,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
token_ids = self._minicpmo45_native_duplex_token_ids()
listen_id = int(token_ids.get("listen_token_id", -1))
tts_bos_id = int(token_ids.get("tts_bos_token_id", -1))
turn_eos_id = int(token_ids.get("turn_eos_token_id", -1))
if listen_id < 0 or listen_id >= logits.shape[-1]:
return
@@ -199,7 +198,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
continue
force_listen = payload.get("force_listen") is True
is_speech = payload.get("is_speech")
redirect_listen = False
segment_key = (row.request_id, row.seq if row.seq is not None else -1)
session_key = (row.session_id, row.incarnation) if row.session_id is not None else None
if turn_eos_id >= 0 and session_key is not None:
@@ -209,8 +207,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
)
if is_speech is True:
if state is not None:
if not getattr(state, "current_turn_ended", True):
redirect_listen = True
with suppress(Exception):
state.last_terminator_token = None
else:
@@ -218,7 +214,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
if turn_ended and not pending_speech_context:
force_listen = True
if not force_listen and not redirect_listen:
if not force_listen:
continue
if force_listen and segment_key in force_listen_segments:
continue
@@ -226,10 +222,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
logits[row_idx, :] = float("-inf")
logits[row_idx, listen_id] = 0.0
force_listen_segments.add(segment_key)
elif 0 <= tts_bos_id < logits.shape[-1]:
if logits[row_idx, listen_id] > logits[row_idx, tts_bos_id]:
logits[row_idx, tts_bos_id] = logits[row_idx, listen_id]
logits[row_idx, listen_id] = float("-inf")
# -------------------- Device utilities --------------------
@staticmethod
@@ -730,6 +722,10 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or []
raw_recent_tokens = output_token_ids[row_idx] if row_idx < len(output_token_ids) else []
recent_tokens = [int(token_id) for token_id in raw_recent_tokens if isinstance(token_id, int) and token_id >= 0]
temperature = float(self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7))
top_k = int(self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100))
top_p = float(self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8))
state = self._minicpmo45_duplex_state_for_row(row_idx)
if chunk_eos_id >= 0 and chunk_eos_id < logits.shape[-1]:
max_speak_tokens = int(
getattr(
@@ -765,23 +761,19 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
if valid_forbidden:
logits[:, valid_forbidden] = float("-inf")
special_ids = self._minicpmo45_native_special_token_ids(token_ids)
state = self._minicpmo45_duplex_state_for_row(row_idx)
generated_text_tokens = getattr(state, "generated_text_tokens", None)
repetition_tokens = generated_text_tokens if generated_text_tokens else recent_tokens
generated_tokens = getattr(state, "generated_tokens", None)
repetition_tokens = generated_tokens if generated_tokens else recent_tokens
repetition_penalty = 1.05
if repetition_penalty != 1.0 and repetition_tokens:
history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE
for token_id in set(repetition_tokens[-history_size:]):
if token_id in special_ids or token_id < 0 or token_id >= logits.shape[-1]:
if token_id < 0 or token_id >= logits.shape[-1]:
continue
logits[0, token_id] /= repetition_penalty
temperature = float(self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7))
top_k = int(self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100))
top_p = float(self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8))
if getattr(sampling_metadata, "all_greedy", False) or temperature <= 0:
sampled = int(torch.argmax(logits, dim=-1).item())
self._record_minicpmo45_duplex_generation_token(row_idx, sampled)
sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk(
sampled,
recent_tokens,
@@ -793,6 +785,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
logits = self._top_k_top_p_filter(logits, top_k=top_k, top_p=top_p)
probs = F.softmax(logits, dim=-1)
sampled = int(torch.multinomial(probs, num_samples=1, generator=generator).item())
self._record_minicpmo45_duplex_generation_token(row_idx, sampled)
sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk(
sampled,
recent_tokens,
@@ -901,6 +894,28 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
return int(tts_bos_id)
return int(sampled)
def _record_minicpmo45_duplex_generation_token(self, row_idx: int, sampled: int) -> None:
"""Track tokens returned by the model-policy decoder.
The released StreamDecoder is constructed without special_token_ids, so
every normally decoded token participates in repetition penalty. Forced
listen bypasses that decoder, while chunk_eos boundary decisions return
before this method is called.
"""
state = self._minicpmo45_duplex_state_for_row(row_idx)
if state is None:
return
payload = self._minicpmo45_duplex_payload_for_row(row_idx)
if isinstance(payload, dict) and payload.get("force_listen") is True:
return
generated_tokens = getattr(state, "generated_tokens", None)
if not isinstance(generated_tokens, list):
generated_tokens = []
state.generated_tokens = generated_tokens
generated_tokens.append(int(sampled))
history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE
del generated_tokens[:-history_size]
def _record_minicpmo45_duplex_terminator(self, row_idx: int, sampled: int, token_ids: dict[str, int]) -> None:
"""Remember sampled unit state for the next append.
@@ -928,15 +943,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
with suppress(Exception):
state.pending_speech_response_open = False
return
special_ids = self._minicpmo45_native_special_token_ids(token_ids)
if sampled not in special_ids:
generated_text_tokens = getattr(state, "generated_text_tokens", None)
if not isinstance(generated_text_tokens, list):
generated_text_tokens = []
state.generated_text_tokens = generated_text_tokens
generated_text_tokens.append(int(sampled))
history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE
del generated_text_tokens[:-history_size]
if (
sampled == tts_bos_id
and getattr(state, "current_turn_ended", True)
@@ -43,6 +43,7 @@ _CODEC_TOP_K = 25
_CODEC_TOP_P = 0.85
_CODEC_REPETITION_PENALTY = 1.05
_CODEC_MIN_TOKENS = 50
_DUPLEX_CODEC_TOKENS_PER_CHUNK = 26
def _max_audio_tokens(condition_tokens: int) -> int:
@@ -61,12 +62,7 @@ def _max_audio_tokens(condition_tokens: int) -> int:
def _restore_weight_norm_weight(weight_g: torch.Tensor, weight_v: torch.Tensor) -> torch.Tensor:
"""Materialize ``weight_norm(..., dim=0)`` checkpoint parameters."""
row_norm = torch.linalg.vector_norm(
weight_v.float(),
dim=1,
keepdim=True,
).clamp_min(torch.finfo(torch.float32).tiny)
return weight_v * (weight_g.to(dtype=weight_v.dtype) / row_norm.to(dtype=weight_v.dtype))
return torch._weight_norm(weight_v, weight_g, dim=0)
def _apply_repetition_penalty(
@@ -218,6 +214,8 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
self,
tts_token_ids: torch.Tensor,
tts_hidden_states: torch.Tensor,
*,
native_duplex: bool = False,
) -> torch.Tensor:
if tts_token_ids.numel() == 0 or tts_hidden_states.numel() == 0:
# The thinker can legally emit an empty speech segment (<|tts_bos|>
@@ -239,7 +237,12 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
hidden_embeds = self.projector_semantic(hidden)
if self._normalize:
hidden_embeds = F.normalize(hidden_embeds, p=2, dim=-1)
return torch.cat([text_embeds + hidden_embeds, self._boundary_embeddings()], dim=0)
audio_bos = self.emb_text(torch.tensor([self._tts_bos_id], device=device, dtype=torch.long))
condition = text_embeds + hidden_embeds
if native_duplex:
# Match MiniCPMTTS.generate_chunk's streaming condition.
return torch.cat([condition, audio_bos], dim=0)
return torch.cat([condition, self._boundary_embeddings()], dim=0)
def preprocess(
self,
@@ -281,18 +284,22 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
"MiniCPM-o Talker received an empty condition (request %s); this request produces no audio.",
info_dict.get("request_id"),
)
full_embeds = self._build_condition_embeddings(token_ids, hidden_states)
native_duplex = bool(info_dict.get("native_duplex", False))
full_embeds = self._build_condition_embeddings(
token_ids,
hidden_states,
native_duplex=native_duplex,
)
offset = int(info_dict.get("_omni_num_computed_tokens", 0))
if offset == 0 and span_len > full_embeds.shape[0]:
# Async-chunk prewarms Stage 1 with placeholder token IDs. Two
# Thinker handoffs can arrive before the first one executes,
# so the scheduler prompt may already include the newer
# placeholder prefix while additional_information still
# carries the preceding condition. Materialize that prefix as
# the embeddings of its actual zero token IDs and keep the
# complete Talker condition tail-aligned (audio BOS remains
# the final prefill token).
prefix_len = span_len - full_embeds.shape[0]
request_id = str(info_dict.get("request_id", "0"))
meta = info_dict.get("meta")
# The handoff rebuilds only the tail-aligned Talker condition.
# Materialize zero-token embeddings for any scheduler prompt
# prefix so chunked prefill can slice from a non-zero offset.
prompt_len = info_dict.get("_omni_prompt_len")
target_len = int(prompt_len) if prompt_len is not None else offset + span_len
prefix_len = target_len - full_embeds.shape[0]
if prefix_len > 0:
placeholder_ids = torch.zeros(
prefix_len,
dtype=torch.long,
@@ -308,13 +315,21 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
f"tts_ids={token_ids.shape[0]} tts_hidden={hidden_states.shape[0]} "
f"prompt_len={info_dict.get('_omni_prompt_len')}"
)
max_tokens = _max_audio_tokens(int(token_ids.numel()))
duplex_boundary = isinstance(meta, dict) and (
bool(meta.get("turn_start", False)) or bool(meta.get("turn_end", False))
)
if native_duplex:
max_tokens = _DUPLEX_CODEC_TOKENS_PER_CHUNK
min_tokens = 0 if duplex_boundary else _DUPLEX_CODEC_TOKENS_PER_CHUNK
else:
max_tokens = _max_audio_tokens(int(token_ids.numel()))
min_tokens = self._codec_min_tokens
state = {
"step": 0,
"max_tokens": max_tokens,
"min_tokens": min_tokens,
"finished": empty_condition,
}
request_id = str(info_dict.get("request_id", "0"))
request_states = getattr(self, "_request_audio_states", None)
if request_states is None:
request_states = {}
@@ -370,7 +385,12 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
penalty=self._codec_repetition_penalty,
window_size=_REPETITION_WINDOW,
)
if step < self._codec_min_tokens:
request_states = getattr(self, "_request_audio_states", {})
state = request_states.get(request_id)
min_tokens = (
int(state.get("min_tokens", self._codec_min_tokens)) if isinstance(state, dict) else self._codec_min_tokens
)
if step < min_tokens:
logits[..., eos_id] = float("-inf")
logits = _apply_top_k_top_p(
logits,
@@ -404,12 +424,61 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
raise RuntimeError(
f"MiniCPM-o continuous Talker received {len(sample_eligible)} sampling flags for {len(infos)} requests"
)
emit_duplex_metadata = any(isinstance(info, dict) and info.get("native_duplex") is True for info in infos)
stop_rows: list[torch.Tensor] = []
codec_deltas: list[torch.Tensor] = []
terminal_flags: list[torch.Tensor] = []
native_duplex_flags: list[torch.Tensor] = []
duplex_epochs: list[torch.Tensor] = []
duplex_turn_ids: list[torch.Tensor] = []
segment_texts_utf8: list[torch.Tensor] = []
turn_end_flags: list[torch.Tensor] = []
empty_delta = hidden.new_empty((0, 1), dtype=torch.long)
for index, info in enumerate(infos):
info_dict = info if isinstance(info, dict) else {}
native_duplex = info_dict.get("native_duplex") is True
if emit_duplex_metadata:
duplex_info = info_dict.get("duplex")
if not isinstance(duplex_info, dict):
duplex_info = {}
epoch = duplex_info.get("epoch", -1)
turn_id = duplex_info.get("turn_id", -1)
if native_duplex and not all(
isinstance(value, int) and not isinstance(value, bool) and value >= 0 for value in (epoch, turn_id)
):
raise RuntimeError(
"MiniCPM-o native duplex Talker requires non-negative integer "
f"epoch and turn_id, got epoch={epoch!r}, turn_id={turn_id!r}"
)
meta_info = info_dict.get("meta")
if not isinstance(meta_info, dict):
meta_info = {}
segment_text = meta_info.get("native_duplex_segment_text", "") if native_duplex else ""
if not isinstance(segment_text, str):
segment_text = ""
turn_eos_id = meta_info.get("turn_eos_token_id")
ids_info = info_dict.get("ids")
tts_ids = ids_info.get("tts") if native_duplex and isinstance(ids_info, dict) else None
if isinstance(tts_ids, torch.Tensor):
contains_turn_eos = isinstance(turn_eos_id, int) and bool(
torch.any(tts_ids.reshape(-1) == turn_eos_id).item()
)
elif isinstance(tts_ids, (list, tuple)):
contains_turn_eos = isinstance(turn_eos_id, int) and turn_eos_id in tts_ids
else:
contains_turn_eos = False
native_duplex_flags.append(torch.tensor(native_duplex, dtype=torch.bool))
duplex_epochs.append(torch.tensor(epoch if isinstance(epoch, int) else -1, dtype=torch.long))
duplex_turn_ids.append(torch.tensor(turn_id if isinstance(turn_id, int) else -1, dtype=torch.long))
segment_texts_utf8.append(
torch.tensor(
list(segment_text.encode("utf-8")),
dtype=torch.uint8,
)
)
turn_end_flags.append(torch.tensor(native_duplex and contains_turn_eos, dtype=torch.bool))
if not isinstance(info, dict):
stop_rows.append(hidden.new_tensor([0.0, float("-inf")]))
codec_deltas.append(empty_delta)
@@ -453,12 +522,15 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
codes = codes.to(device=hidden.device, dtype=torch.long).reshape(-1)
step = int(state.get("step", 0))
sampled = self._sample_audio_code(hidden[end - 1 : end], codes, request_id, step)
is_eos = int(sampled.item()) == self._num_audio_tokens - 1
sampled_id = int(sampled.item())
is_eos = sampled_id == self._num_audio_tokens - 1
state["step"] = int(state.get("step", 0)) + 1
reached_limit = int(state["step"]) >= int(state.get("max_tokens", 2048))
finished = is_eos or reached_limit
state["finished"] = finished
if not is_eos:
# MiniCPMTTS.generate_chunk consumes the boundary sample but
# returns only codes that were fed into the retained KV state.
if not is_eos and not reached_limit:
codes = torch.cat([codes[-(_REPETITION_WINDOW - 1) :], sampled.reshape(1)])
delta = sampled.reshape(1, 1)
else:
@@ -476,9 +548,20 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
self._batch_stop_logits = torch.stack(stop_rows, dim=0) if stop_rows else hidden.new_empty((0, 2))
# Lists are deliberate: the runner routes element i to request i,
# preserving compaction alignment while emitting only this step's code.
meta_outputs = {"finished": terminal_flags}
if emit_duplex_metadata:
meta_outputs.update(
{
"native_duplex": native_duplex_flags,
"duplex_epoch": duplex_epochs,
"duplex_turn_id": duplex_turn_ids,
"llm_output_text_utf8": segment_texts_utf8,
"turn_end": turn_end_flags,
}
)
multimodal_outputs: dict[str, Any] = {
"codes": {"audio": codec_deltas},
"meta": {"finished": terminal_flags},
"meta": meta_outputs,
}
return OmniOutput(
text_hidden_states=hidden,
@@ -21,6 +21,7 @@ logger = logging.getLogger(__name__)
_MINICPMO45_ASYNC_STATE = "_minicpmo45_async_codec_state"
_MINICPMO45_STREAM_RECORD = "_minicpmo45_async_stream_record"
_MINICPMO45_SILENCE_CODE = 4218
_MINICPMO45_MIN_STREAM_BODY_FRAMES = 5
class _MiniCPMO45MetaStruct(MetaStruct):
@@ -28,6 +29,7 @@ class _MiniCPMO45MetaStruct(MetaStruct):
ref_audio_sr: int | None = None
native_duplex_segment_text: str | None = None
llm_output_text_utf8: torch.Tensor | None = None
duplex_turn_id: int | None = None
duplex_epoch: int | None = None
segment_end: bool | None = None
@@ -212,6 +214,16 @@ def tts2code2wav_async_chunk(
internal_id = getattr(request, "request_id", None)
request_id = str(external_id if external_id is not None else internal_id)
internal_id = str(internal_id if internal_id is not None else request_id)
output_meta = multimodal_output.get("meta") if isinstance(multimodal_output, Mapping) else None
output_meta = output_meta if isinstance(output_meta, Mapping) else {}
native_duplex = bool(_coerce_int(output_meta.get("native_duplex")))
duplex_epoch = _coerce_int(output_meta.get("duplex_epoch"))
duplex_turn_id = _coerce_int(output_meta.get("duplex_turn_id"))
segment_text_utf8 = output_meta.get("llm_output_text_utf8")
if not isinstance(segment_text_utf8, torch.Tensor):
segment_text_utf8 = None
turn_end = bool(_coerce_int(output_meta.get("turn_end")))
duplex_turn_key = (duplex_epoch, duplex_turn_id)
request_payload = getattr(transfer_manager, "request_payload", None)
if request_payload is None:
@@ -227,7 +239,9 @@ def tts2code2wav_async_chunk(
record = {
"internal_id": internal_id,
"cache_epoch": 0,
"chunk_seq": 0,
"retired_internal_ids": set(),
"last_terminal_turn": None,
}
container[_MINICPMO45_STREAM_RECORD] = record
elif internal_id in record["retired_internal_ids"]:
@@ -236,6 +250,8 @@ def tts2code2wav_async_chunk(
record["retired_internal_ids"].add(record["internal_id"])
record["internal_id"] = internal_id
record["cache_epoch"] = int(record["cache_epoch"]) + 1
record["chunk_seq"] = 0
record["last_terminal_turn"] = None
_drop_codec_state(transfer_manager, request_id)
if _is_aborted(request):
@@ -243,44 +259,46 @@ def tts2code2wav_async_chunk(
_drop_codec_state(transfer_manager, request_id)
return None
if native_duplex and turn_end and record.get("last_terminal_turn") == duplex_turn_key:
return None
state = container.get(_MINICPMO45_ASYNC_STATE)
if not isinstance(state, dict):
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 = {}
state = {
"internal_id": internal_id,
"pending": [],
"pending_text_utf8": [],
"segment_text_recorded": False,
"left_context": [],
"codec_end": 0,
"ref_audio": codes_info.get("ref"),
"ref_audio_sr": meta_info.get("ref_audio_sr"),
"segment_text": meta_info.get("native_duplex_segment_text"),
"segment_end": bool(meta_info.get("segment_end", False)),
"turn_end": bool(meta_info.get("turn_end", False)),
"duplex_turn_id": duplex_info.get("model_turn_id", duplex_info.get("turn_id")),
"duplex_epoch": duplex_info.get("epoch"),
}
container[_MINICPMO45_ASYNC_STATE] = state
pending = state["pending"]
pending.extend(_extract_codec_delta(multimodal_output, request_id))
pending_text_utf8 = state.setdefault("pending_text_utf8", [])
current_text_utf8 = (
segment_text_utf8.detach().to(device="cpu", dtype=torch.uint8).reshape(-1).tolist()
if isinstance(segment_text_utf8, torch.Tensor)
else []
)
if native_duplex and current_text_utf8 and not state.get("segment_text_recorded", False):
# Talker repeats the unit text on every codec step. Queue it once for
# the first Code2Wav payload that can carry audio for this unit.
pending_text_utf8.extend(current_text_utf8)
state["segment_text_recorded"] = True
request_finished = getattr(request, "is_finished", None)
finished = bool(is_finished or (callable(request_finished) and request_finished()))
chunk_frames, left_context_frames = _codec_config(transfer_manager)
if not finished and len(pending) < chunk_frames:
flush_pending = finished
last_chunk = bool(flush_pending and (not native_duplex or turn_end))
if not flush_pending and len(pending) < chunk_frames:
return None
new_token_count = len(pending) if finished else chunk_frames
hold_short_unit = (
native_duplex and flush_pending and not last_chunk and 0 < len(pending) < _MINICPMO45_MIN_STREAM_BODY_FRAMES
)
new_token_count = 0 if hold_short_unit else (len(pending) if flush_pending else chunk_frames)
new_codes = pending[:new_token_count]
del pending[:new_token_count]
codec_start = int(state["codec_end"])
@@ -292,53 +310,79 @@ def tts2code2wav_async_chunk(
else:
context = list(state["left_context"])
output_codes = [*context, *new_codes]
history = [*state["left_context"], *new_codes]
history = output_codes
state["left_context"] = history[-left_context_frames:] if left_context_frames else []
elif finished and codec_start > 0 and state["left_context"]:
elif last_chunk and codec_start > 0 and state["left_context"]:
context = list(state["left_context"])
output_codes = context
else:
context = []
output_codes = []
state["codec_end"] = codec_end
code_flat_numel = len(output_codes)
if native_duplex and code_flat_numel > 0:
segment_text_utf8 = torch.tensor(pending_text_utf8, dtype=torch.uint8)
pending_text_utf8.clear()
if native_duplex and finished:
state["segment_text_recorded"] = False
if flush_pending and not last_chunk and code_flat_numel == 0:
# Keep the generic generation connector model-agnostic: a real token
# makes this control-only TTS boundary schedulable, while the explicit
# zero length tells Code2Wav to discard the placeholder.
output_codes = [0]
last_chunk = bool(finished and not pending)
if last_chunk:
if last_chunk and not native_duplex:
record["retired_internal_ids"].add(internal_id)
_drop_codec_state(transfer_manager, request_id)
chunk_seq = int(getattr(transfer_manager, "put_req_chunk", {}).get(request_id, 0))
chunk_seq = int(record["chunk_seq"])
record["chunk_seq"] = chunk_seq + 1
ref_audio = None
ref_audio_sr = None
if int(record["cache_epoch"]) == 0 and chunk_seq == 0:
request_info = getattr(request, "additional_information", None)
if isinstance(request_info, Mapping):
codes_info = request_info.get("codes")
meta_info = request_info.get("meta")
raw_ref_audio = codes_info.get("ref") if isinstance(codes_info, Mapping) else None
raw_ref_audio_sr = meta_info.get("ref_audio_sr") if isinstance(meta_info, Mapping) else None
ref_audio_sr = _coerce_int(raw_ref_audio_sr)
if raw_ref_audio is not None:
ref_audio = torch.as_tensor(raw_ref_audio, dtype=torch.float32).reshape(-1).cpu()
finished_tensor = torch.tensor(last_chunk, dtype=torch.bool)
ref_audio = state.get("ref_audio") if codec_start == 0 else None
return OmniPayloadStruct(
payload = 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,
ref=ref_audio,
),
meta=_MiniCPMO45MetaStruct(
request_id=request_id,
chunk_seq=chunk_seq,
cache_epoch=int(record["cache_epoch"]),
code_flat_numel=len(output_codes),
code_flat_numel=code_flat_numel,
codec_chunk_frames=new_token_count,
codec_left_context_frames=len(context),
left_context_size=len(context),
last_chunk=last_chunk,
stream_finished=finished_tensor,
finished=finished_tensor,
is_segment_finished=finished_tensor,
req_id=[request_id],
ref_audio_sr=_coerce_int(state.get("ref_audio_sr")),
native_duplex_segment_text=(
str(state["segment_text"]) if isinstance(state.get("segment_text"), str) else None
),
duplex_turn_id=_coerce_int(state.get("duplex_turn_id")),
duplex_epoch=_coerce_int(state.get("duplex_epoch")),
segment_end=bool(state.get("segment_end", False)),
turn_end=bool(state.get("turn_end", False)),
tts_is_last_chunk=last_chunk,
duplex_epoch=duplex_epoch,
duplex_turn_id=duplex_turn_id,
llm_output_text_utf8=segment_text_utf8,
tts_is_last_chunk=flush_pending,
turn_end=turn_end and last_chunk,
ref_audio_sr=ref_audio_sr,
),
request_id=request_id,
)
if last_chunk and native_duplex:
record["last_terminal_turn"] = duplex_turn_key
record["cache_epoch"] = int(record["cache_epoch"]) + 1
record["chunk_seq"] = 0
_drop_codec_state(transfer_manager, request_id)
return payload
def tts2code2wav_full_payload(
@@ -496,13 +540,33 @@ def _native_tts_boundary_token_ids(special_token_ids):
}
def _decode_native_duplex_token_ids(
token_ids: list[int],
streaming_context,
*,
request_id: str,
) -> str | None:
decode_token_ids = getattr(streaming_context, "source_token_decoder", None)
if not isinstance(decode_token_ids, Callable):
return None
decode_ids = [int(token_id) for token_id in token_ids]
try:
try:
return str(decode_token_ids(decode_ids, skip_special_tokens=True))
except TypeError:
return str(decode_token_ids(decode_ids))
except Exception:
logger.exception("Failed to decode MiniCPM-o duplex token delta for request_id=%s", request_id)
return ""
def _native_duplex_segment_output_ids(
output_ids: list[int],
output_text: str,
streaming_context,
*,
request_id: str,
) -> tuple[list[int], str]:
) -> tuple[list[int], str, bool]:
"""Slice the cumulative thinker output down to the current segment.
Tracks how many output tokens were already handed to the talker in the
@@ -511,7 +575,7 @@ def _native_duplex_segment_output_ids(
"""
bridge_states = getattr(streaming_context, "bridge_states", None)
if not isinstance(bridge_states, dict):
return output_ids, output_text
return output_ids, output_text, True
state = bridge_states.setdefault("minicpmo45_tts_handoff", {})
duplex_state = bridge_states.get("duplex")
turn_id = duplex_state.get("model_turn_id", duplex_state.get("turn_id")) if isinstance(duplex_state, dict) else None
@@ -521,17 +585,12 @@ def _native_duplex_segment_output_ids(
state["request_id"] = request_id
state["sent_output_len"] = 0
state["sent_output_ids"] = []
state["acc_tts_ids"] = []
state["acc_tts_hidden"] = []
previous_turn_id = state.get("turn_id")
sent_len = state.get("sent_output_len", 0)
prev_output_ids = state.get("sent_output_ids", [])
prev_turn_id = state.get("turn_id")
if not isinstance(sent_len, int) or sent_len < 0 or sent_len > len(output_ids):
# Shrunken cumulative output = epoch reset after barge-in: the talker
# condition history is stale too.
# Shrunken cumulative output = epoch reset after barge-in.
sent_len = 0
state["acc_tts_ids"] = []
state["acc_tts_hidden"] = []
elif sent_len and isinstance(prev_output_ids, list):
prev_prefix = prev_output_ids[:sent_len]
current_prefix = output_ids[:sent_len]
@@ -540,27 +599,15 @@ def _native_duplex_segment_output_ids(
# others keep returning cumulative ids. Detect restart from the
# token prefix instead of clearing the cursor at turn_eos.
sent_len = 0
state["acc_tts_ids"] = []
state["acc_tts_hidden"] = []
if isinstance(prev_turn_id, int) and isinstance(turn_id, int) and prev_turn_id != turn_id:
# The thinker output list is cumulative across clean turns, so the
# output cursor must stay put. The talker condition is per assistant
# turn, though; carrying it across turn_id boundaries makes stage 1
# replay the previous turn after its consumed cursor has been reset.
state["acc_tts_ids"] = []
state["acc_tts_hidden"] = []
turn_start = sent_len == 0 or previous_turn_id != turn_id
segment_ids = output_ids[sent_len:]
decode_token_ids = getattr(streaming_context, "source_token_decoder", None)
if isinstance(decode_token_ids, Callable):
decode_ids = [int(token_id) for token_id in segment_ids]
try:
try:
segment_text = str(decode_token_ids(decode_ids, skip_special_tokens=True))
except TypeError:
segment_text = str(decode_token_ids(decode_ids))
except Exception:
logger.exception("Failed to decode MiniCPM-o duplex token delta for request_id=%s", request_id)
segment_text = ""
decoded_segment_text = _decode_native_duplex_token_ids(
segment_ids,
streaming_context,
request_id=request_id,
)
if decoded_segment_text is not None:
segment_text = decoded_segment_text
elif sent_len == 0:
# Non-orchestrator unit tests and legacy callers may not install a
# decoder. Never slice cumulative text with a stale character cursor.
@@ -575,22 +622,7 @@ def _native_duplex_segment_output_ids(
state["sent_output_len"] = len(output_ids)
state["sent_output_ids"] = list(output_ids)
state["turn_id"] = turn_id
return segment_ids, segment_text
def _reset_native_tts_handoff(streaming_context) -> None:
bridge_states = getattr(streaming_context, "bridge_states", None)
if not isinstance(bridge_states, dict):
return
state = bridge_states.get("minicpmo45_tts_handoff")
if not isinstance(state, dict):
return
# Keep the output cursor: the thinker can keep reporting cumulative
# output after turn_eos on the same resumable request. If a runtime really
# restarts token ids, _native_duplex_segment_output_ids detects the prefix
# mismatch and resets the slice cursor there.
state["acc_tts_ids"] = []
state["acc_tts_hidden"] = []
return segment_ids, segment_text, turn_start
def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] | None:
@@ -623,32 +655,6 @@ def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] |
return metadata
def _accumulate_native_tts_handoff(streaming_context, new_ids, new_hidden):
"""Hand the talker the FULL accumulated condition on every handoff.
The runner's streaming buffer update is not merge-safe for a resumable
stage-1 request: in-place updates merge sub-keys, but a resume prefill
REPLACES the buffer, silently dropping every earlier segment's tts
tokens/hiddens (observed losing alternating reply segments — the talker
vocalized text it never saw between islands it did). Accumulate here so
the latest handoff always carries the complete history and downstream
replace semantics are lossless; the talker consumes by cursor.
"""
bridge_states = getattr(streaming_context, "bridge_states", None)
if not isinstance(bridge_states, dict):
return new_ids, new_hidden
state = bridge_states.setdefault("minicpmo45_tts_handoff", {})
acc_ids = state.setdefault("acc_tts_ids", [])
acc_hidden = state.setdefault("acc_tts_hidden", [])
if new_ids:
acc_ids.extend(int(t) for t in new_ids)
if new_hidden:
acc_hidden.extend(new_hidden)
if not acc_ids:
return None, None
return list(acc_ids), list(acc_hidden)
def _build_tts_scheduler_prompt_token_ids(
tts_token_ids: torch.Tensor | None,
llm_output_ids: list[int],
@@ -723,6 +729,7 @@ def llm2tts(
# buffer overflows.
llm_output_ids = list(llm_output_ids)
thinker_text = getattr(output, "text", "") or ""
native_turn_start = False
if _has_native_duplex_prompt_metadata(mm_output):
# The thinker's resumable duplex request reports cumulative
# output ids/text, but earlier segments are already folded into
@@ -732,7 +739,7 @@ def llm2tts(
# alignment holds, the talker prompt grows linearly with new
# tokens instead of quadratically, and downstream transcripts
# carry per-unit deltas instead of re-sending the whole reply.
llm_output_ids, thinker_text = _native_duplex_segment_output_ids(
llm_output_ids, thinker_text, native_turn_start = _native_duplex_segment_output_ids(
llm_output_ids,
thinker_text,
_streaming_context,
@@ -875,6 +882,18 @@ def llm2tts(
.to(torch.float32)
.contiguous()
)
handoff_ids = _coerce_token_id_list(tts_token_ids_slice) if tts_token_ids_slice is not None else None
if is_native_duplex_handoff and handoff_ids:
handoff_text = _decode_native_duplex_token_ids(
handoff_ids,
_streaming_context,
request_id=str(llm_output.request_id),
)
if handoff_text is not None:
# Match the released streaming_generate contract: transcript
# text comes from total_ids_in_unit, the same slice that
# conditions the Talker.
thinker_text = handoff_text
model_intermediate_buffer = build_duplex_intermediate_buffer(
request_id=str(llm_output.request_id),
prompt_token_ids=prompt_token_ids,
@@ -900,6 +919,7 @@ def llm2tts(
# The talker detects turn end from <|turn_eos|> inside the
# handed condition (official conditions on its embedding).
meta["turn_eos_token_id"] = int(turn_eos_id)
meta["turn_start"] = native_turn_start
if native_segment_end:
meta["segment_end"] = True
req_mm_data = multi_modal_data.get(llm_output.request_id)
@@ -911,30 +931,27 @@ def llm2tts(
if ref_audio is not None:
ref_waveform, ref_sr = ref_audio
set_ref_audio(model_intermediate_buffer, _to_transport_list(ref_waveform), ref_sr)
handoff_ids = _coerce_token_id_list(tts_token_ids_slice) if tts_token_ids_slice is not None else None
handoff_hidden = _to_transport_list(tts_hidden_slice) if tts_hidden_slice is not None else None
native_turn_end_handoff = False
if is_native_duplex_handoff:
turn_eos_id = special_token_ids.get("turn_eos_token_id")
native_turn_end_handoff = turn_eos_id is not None and handoff_ids is not None and turn_eos_id in handoff_ids
handoff_ids, handoff_hidden = _accumulate_native_tts_handoff(
_streaming_context,
handoff_ids,
handoff_hidden,
)
if not handoff_ids:
continue
set_tts_handoff(model_intermediate_buffer, handoff_ids, handoff_hidden)
if native_turn_end_handoff:
model_intermediate_buffer.setdefault("meta", {})["turn_end"] = True
_reset_native_tts_handoff(_streaming_context)
if handoff_ids is not None and handoff_hidden is not None:
condition_length = max(len(handoff_ids), len(handoff_hidden)) + 2
condition_suffix_length = 1 if is_native_duplex_handoff else 2
condition_length = max(len(handoff_ids), len(handoff_hidden)) + condition_suffix_length
scheduler_prompt_token_ids = [0] * condition_length
handoff_meta = model_intermediate_buffer.setdefault("meta", {})
handoff_meta["replace_streaming_prompt"] = True
handoff_meta["next_stage_prompt_len"] = condition_length
# Native duplex resumes one Talker request within a turn, but a new
# assistant turn must discard the previous turn's prompt and KV.
if not is_native_duplex_handoff or native_turn_start:
handoff_meta["replace_streaming_prompt"] = True
else:
scheduler_prompt_token_ids = _build_tts_scheduler_prompt_token_ids(
tts_token_ids_slice,
@@ -685,19 +685,32 @@ def talker2code2wav_async_chunk(
"""
Multimodal output version.
"""
if not isinstance(multimodal_output, Mapping):
return None
talker_codes = multimodal_output.get("codes", {})
if not isinstance(talker_codes, dict):
return None
code_predictor_codes = talker_codes.get("audio")
if code_predictor_codes is None:
return None
request_id = request.external_req_id
code_predictor_codes = None
if isinstance(multimodal_output, Mapping):
talker_codes = multimodal_output.get("codes", {})
if isinstance(talker_codes, dict):
code_predictor_codes = talker_codes.get("audio")
if code_predictor_codes.numel() == 0:
return None
sampling_params = getattr(request, "sampling_params", None)
stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or [])
stop_token_id = getattr(sampling_params, "stop_token_id", None)
if stop_token_id is not None:
stop_token_ids.add(stop_token_id)
if not code_predictor_codes.any():
append_codes = (
isinstance(code_predictor_codes, torch.Tensor)
and code_predictor_codes.numel() > 0
and bool(code_predictor_codes.any())
)
if append_codes:
first_codebook = int(code_predictor_codes[0, 0].item())
if first_codebook in stop_token_ids:
logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook)
append_codes = False
if append_codes:
transfer_manager.code_prompt_token_ids[request_id].append(code_predictor_codes)
elif not is_finished:
return None
connector = getattr(transfer_manager, "connector", None)
@@ -707,20 +720,10 @@ def talker2code2wav_async_chunk(
left_context_size_config = int(cfg.get("codec_left_context_frames", 25))
configured_initial_chunk_size = int(cfg.get("initial_codec_chunk_frames") or 0)
sampling_params = getattr(request, "sampling_params", None)
stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or [])
stop_token_id = getattr(sampling_params, "stop_token_id", None)
if stop_token_id is not None:
stop_token_ids.add(stop_token_id)
first_codebook = int(code_predictor_codes[0, 0].item())
if first_codebook in stop_token_ids:
logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook)
return None
request_id = request.external_req_id
chunk_id = transfer_manager.put_req_chunk[request_id]
transfer_manager.code_prompt_token_ids[request_id].append(code_predictor_codes)
length = len(transfer_manager.code_prompt_token_ids[request_id])
if length <= 0:
return None
if configured_initial_chunk_size > 0:
if chunk_id == 0:
@@ -732,6 +735,9 @@ def talker2code2wav_async_chunk(
if chunk_length != 0 and not is_finished:
return None
if is_finished and not append_codes and chunk_length == 0:
return None
context_length = chunk_length if chunk_length != 0 else chunk_size_config
# ensure left context does not exceed available length
if configured_initial_chunk_size > 0 and chunk_id == 1:
@@ -11,6 +11,7 @@ CHUNK_METADATA_KEYS: frozenset[str] = frozenset(
"duplex_epoch",
"duplex_turn_id",
"llm_output_text_utf8",
"segment_end",
"tts_is_last_chunk",
"turn_end",
}
+16 -4
View File
@@ -523,24 +523,27 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
upstream_outputs.append(eco)
# Handle multimodal-only outputs (generation stages) locally.
self._process_mm_only_outputs(mm_only_outputs)
mm_request_outputs = self._process_mm_only_outputs(mm_only_outputs)
# Delegate text/pooling outputs to upstream.
return super().process_outputs(
processed = super().process_outputs(
upstream_outputs,
engine_core_timestamp=engine_core_timestamp,
iteration_stats=iteration_stats,
)
processed.request_outputs.extend(mm_request_outputs)
return processed
def _process_mm_only_outputs(
self,
engine_core_outputs: list[EngineCoreOutput],
) -> None:
) -> list[OmniRequestOutput | PoolingRequestOutput]:
"""Handle outputs from generation stages that have no detokenizer.
These cannot go through upstream process_outputs because it asserts
detokenizer is not None when pooling_output is None.
"""
request_outputs: list[OmniRequestOutput | PoolingRequestOutput] = []
for eco in engine_core_outputs:
req_state = self.request_states.get(eco.request_id)
if req_state is None or not isinstance(req_state, OmniRequestState):
@@ -554,6 +557,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
req_state.num_cached_tokens = eco.num_cached_tokens
req_state.is_prefilling = False
is_non_final_audio_chunk = (
finish_reason is not None
and req_state.output_kind == RequestOutputKind.DELTA
and is_non_final_delta_audio_chunk(req_state.mm_accumulated, req_state.mm_type)
)
if request_output := req_state.make_request_output(
new_token_ids,
None, # pooling_output
@@ -564,9 +572,13 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
):
if req_state.queue is not None:
req_state.queue.put(request_output)
else:
request_outputs.append(request_output)
if finish_reason is not None:
is_segment_finished = bool(getattr(eco, "is_segment_finished", False))
if finish_reason is not None and not is_segment_finished and not is_non_final_audio_chunk:
self._finish_request(req_state)
return request_outputs
def _update_stats_from_output(
self,