diff --git a/tests/clients/test_duplex_client.py b/tests/clients/test_duplex_client.py index 63d4c282a..176932829 100644 --- a/tests/clients/test_duplex_client.py +++ b/tests/clients/test_duplex_client.py @@ -1486,8 +1486,8 @@ def test_pcm16_wav_round_trip(tmp_path): def test_duplex_unit_boundary_and_residual_math(): - assert duplex_unit_boundary_ms(0) == 1030 - assert duplex_unit_boundary_ms(2) == 3030 + assert duplex_unit_boundary_ms(0) == 1000 + assert duplex_unit_boundary_ms(2) == 3000 assert has_residual_model_unit(b"\x00" * 32_000, chunk_period_ms=1000) is False assert has_residual_model_unit(b"\x00" * 32_002, chunk_period_ms=1000) is True created = {"type": "session.created", "session": {"capabilities": {"chunk_period_ms": 500}}} @@ -1512,9 +1512,13 @@ async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame(): # A composite belongs to the unit it was captured in, so it rides the # same append as that unit's base frame; a unit without one sends the # base alone. Frame k rides the append that closes model unit k. - assert [event["video_frames"] for event in appends if "video_frames" in event] == [["f0", "s0"], ["f1"]] - assert [event["audio_end_ms"] for event in appends if "video_frames" in event] == [1200, 2200] - assert frames_sent == 2 + assert [event["video_frames"] for event in appends if "video_frames" in event] == [ + ["f0", "s0"], + ["f1"], + ["f1"], + ] + assert [event["audio_end_ms"] for event in appends if "video_frames" in event] == [1000, 2000, 3000] + assert frames_sent == 3 sock.feed(SESSION_CLOSED) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 919b9ab2b..bffa8a026 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -1646,6 +1646,37 @@ def test_load_poll_generation_segment_marker_replaces_previous_chunk(build_adapt assert request.request_id in adapter.segment_finished_requests +@pytest.mark.parametrize("cache_epoch", [0, 1]) +def test_generation_first_chunk_replaces_prewarm_runner_payload(build_adapter, cache_epoch): + from vllm_omni.core.sched.output import OmniNewRequestData + + adapter, connector = build_adapter(stage_id=2, model_mode="generation") + request = _req("req-first-chunk", RequestStatus.WAITING) + request.model_intermediate_buffer = {"duplex": {"data_plane": True}, "global_request_id": "session"} + connector.get.return_value = ( + { + "codes": {"audio": torch.tensor([7, 8, 9, 10])}, + "meta": { + "cache_epoch": cache_epoch, + "chunk_seq": 0, + "code_flat_numel": 4, + "replace_runtime_additional_information": True, + }, + }, + 1, + ) + + assert adapter._poll_single_request(_dequeue_load_entry(adapter, request)) is True + scheduled = OmniNewRequestData.from_request(request, block_ids=([0],)) + + # New/resumed requests prioritize model_intermediate_buffer in the runner; + # it must carry the same codec snapshot as cached-request updates. + assert scheduled.model_intermediate_buffer == scheduled.additional_information + assert scheduled.model_intermediate_buffer["meta"]["cache_epoch"] == cache_epoch + assert scheduled.model_intermediate_buffer["meta"]["chunk_seq"] == 0 + assert scheduled.prompt_token_ids == [7, 8, 9, 10] + + def test_load_poll_generation_empty_replacement_snapshot_is_ready(build_adapter): adapter, connector = build_adapter(stage_id=2, model_mode="generation") request = _req("req-empty-marker", RequestStatus.WAITING, external_req_id="external-empty-marker") diff --git a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py index a23ce73e8..15c528c14 100644 --- a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py +++ b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py @@ -2006,19 +2006,18 @@ def test_realtime_duplex_demo_streams_one_video_frame_per_model_unit(): messages = _send_seconds_of_audio(demo, seconds=3, frames=["f0", "f1", "f2", "f3"]) # Frame k rides the append that closes model unit k, so Stage0 can bind it - # to that unit's audio. 3 s of audio closes units at 1030 ms and 2030 ms and - # leaves a residual too short for a third, so only two frames go out. - assert _sent_video_frames(messages) == [["f0"], ["f1"]] + # to that unit's audio. 3 s of audio closes units at 1000, 2000, and 3000 ms. + assert _sent_video_frames(messages) == [["f0"], ["f1"], ["f2"]] frame_indices = [index for index, message in enumerate(messages) if "video_frames" in message] - assert frame_indices == [5, 10] + assert frame_indices == [4, 9, 14] # The carrying appends are exactly the ones that reach a unit boundary. - assert [messages[index]["audio_end_ms"] for index in frame_indices] == [1200, 2200] + assert [messages[index]["audio_end_ms"] for index in frame_indices] == [1000, 2000, 3000] def test_realtime_duplex_demo_never_sends_a_frame_before_its_unit_can_close(): """Regression: frame 0 used to ride the very first 200 ms append. - Stage0 cannot close a unit until 1030 ms of audio has arrived, and it does + Stage0 cannot close a unit until 1000 ms of audio has arrived, and it does not carry frames across appends, so a frame sent that early was silently dropped and every later frame ended up one unit ahead of its audio. """ @@ -2035,11 +2034,19 @@ def test_realtime_duplex_demo_never_sends_a_frame_before_its_unit_can_close(): def test_realtime_duplex_demo_holds_the_last_video_frame_when_audio_outlives_the_clip(): demo = _load_demo_module() - # 4 s of audio closes three units (1030/2030/3030 ms); the two-frame clip - # holds its last frame for the third. - assert _sent_video_frames(_send_seconds_of_audio(demo, seconds=4, frames=["a", "b"])) == [["a"], ["b"], ["b"]] + # 4 s of audio closes four units; the two-frame clip holds its last frame. + assert _sent_video_frames(_send_seconds_of_audio(demo, seconds=4, frames=["a", "b"])) == [ + ["a"], + ["b"], + ["b"], + ["b"], + ] # A still image is a one-element clip and therefore repeats every unit. - assert _sent_video_frames(_send_seconds_of_audio(demo, seconds=3, frames=["still"])) == [["still"], ["still"]] + assert _sent_video_frames(_send_seconds_of_audio(demo, seconds=3, frames=["still"])) == [ + ["still"], + ["still"], + ["still"], + ] def test_realtime_duplex_demo_sends_each_units_stacked_composite_next_to_its_base_frame(): @@ -2055,7 +2062,7 @@ def test_realtime_duplex_demo_sends_each_units_stacked_composite_next_to_its_bas # Official pairing is frame_list=[base, composite of that unit's interior], # so the composite belongs to the same unit as the base beside it -- never # to the previous one. A unit without interior sub-frames sends base alone. - assert _sent_video_frames(messages) == [["f0", "s0"], ["f1", "s1"], ["f2"]] + assert _sent_video_frames(messages) == [["f0", "s0"], ["f1", "s1"], ["f2"], ["f2"]] def test_minicpmo_duplex_camera_fixture_returns_flat_base_frame_track(tmp_path, monkeypatch): diff --git a/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py b/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py index b1cc48177..b3baf871d 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py +++ b/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_forward.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from types import SimpleNamespace from unittest.mock import MagicMock import pytest import torch +from vllm.v1.worker.gpu_model_runner import GPUModelRunner from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( MiniCPMO45OmniForConditionalGeneration, @@ -16,6 +18,49 @@ from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_multimodal_runner_preserves_forwarded_token_identities(monkeypatch): + stage_model = torch.nn.Module() + stage_model.make_empty_intermediate_tensors = lambda: None + monkeypatch.setattr( + "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni.init_vllm_registered_model", + lambda **kwargs: stage_model, + ) + monkeypatch.setattr( + "vllm_omni.model_executor.models.minicpmo_4_5.duplex.compat.patch_minicpmo_remote_config", + lambda config: None, + ) + config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=SimpleNamespace(), multimodal_config=None, model_stage="llm") + ) + wrapper = MiniCPMO45OmniForConditionalGeneration(vllm_config=config) + hidden = torch.randn(2, 4) + wrapper.thinker.forward = MagicMock(return_value=hidden) + token_buffer = torch.tensor([99692, 3837, 0, 0]) + embed_buffer = torch.randn(4, 4) + runner = SimpleNamespace( + model=wrapper, + input_ids=SimpleNamespace(gpu=token_buffer), + inputs_embeds=SimpleNamespace(gpu=embed_buffer), + ) + + # This same upstream path prepares normal multimodal forwards and graph + # capture inputs. Embeddings must not replace the real token identities. + input_ids, inputs_embeds = GPUModelRunner._prepare_mm_inputs(runner, 2) + out = wrapper.forward(input_ids=input_ids, positions=torch.tensor([10, 11]), inputs_embeds=inputs_embeds) + + torch.testing.assert_close(out.multimodal_outputs["latent_input_ids"], token_buffer[:2, None]) + torch.testing.assert_close(out.multimodal_outputs["latent_positions"], torch.tensor([[10], [11]])) + torch.testing.assert_close(wrapper.thinker.forward.call_args.kwargs["inputs_embeds"], embed_buffer[:2]) + + # The codec Talker has no text vocabulary for randomized raw-token + # profiling. It must retain the existing embeddings-only input path. + config.model_config.model_stage = "tts" + runner.model = MiniCPMO45OmniForConditionalGeneration(vllm_config=config) + talker_ids, talker_embeds = GPUModelRunner._prepare_mm_inputs(runner, 2) + assert talker_ids is None + torch.testing.assert_close(talker_embeds, embed_buffer[:2]) + + def test_thinker_forward_returns_bare_hidden_states(): """Verify that MiniCPMO45OmniLLMForConditionalGeneration.forward returns a bare hidden_states Tensor (not a tuple), which stock vLLM GPUModelRunner diff --git a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py index c5d4d0cde..47b0f32c1 100644 --- a/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py +++ b/tests/model_executor/stage_input_processors/test_minicpmo_4_5_omni.py @@ -25,6 +25,11 @@ def _output( ): mm_output = dict(multimodal_output or {}) mm_output["latent"] = latent + if "duplex_prompt_token_ids" in mm_output and "latent_input_ids" not in mm_output: + forwarded_ids = [*prompt_ids, *output_ids] + assert len(forwarded_ids) == latent.shape[0] + mm_output["latent_input_ids"] = torch.tensor(forwarded_ids).reshape(-1, 1) + mm_output["latent_positions"] = torch.arange(len(forwarded_ids)).reshape(-1, 1) completion = SimpleNamespace( token_ids=output_ids if token_list is None else token_list, text="hello", @@ -142,6 +147,246 @@ def test_native_duplex_speak_segment_reaches_split_talker() -> None: assert info["duplex"]["turn_id"] == 7 +def test_native_duplex_uses_forwarded_row_ledger_after_later_audio_prefill() -> None: + prompt_ids = [101, 102] + output_ids = [9304, 21, 22, 9308] + row_ids = [101, 102, 9304, 21, 22, 999, 999] + latent = torch.arange(len(row_ids) * 4, dtype=torch.float32).reshape(-1, 4) + 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, + } + source = _output( + prompt_ids=prompt_ids, + output_ids=output_ids, + latent=latent, + multimodal_output={ + "duplex_prompt_token_ids": prompt_ids, + "latent_input_ids": torch.tensor(row_ids).reshape(-1, 1), + "latent_positions": torch.arange(len(row_ids)).reshape(-1, 1), + "meta": metadata, + }, + ) + context = SimpleNamespace(bridge_states={"duplex": {"epoch": 3, "model_turn_id": 7}}) + + converted = llm2tts([source], prompt=[{}], _streaming_context=context)[0] + + info = converted["model_intermediate_buffer"] + assert info["ids"]["tts"] == [21, 22] + assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[3:5]) + + +def test_native_duplex_selects_latest_contiguous_forwarded_span() -> None: + prompt_ids = [101, 102] + output_ids = [9304, 21, 22, 9308] + row_ids = [101, 21, 22, 102, 9304, 21, 22, 999] + positions = [0, 1, 3, 4, 5, 6, 7, 8] + latent = torch.arange(len(row_ids) * 4, dtype=torch.float32).reshape(-1, 4) + source = _output( + prompt_ids=prompt_ids, + output_ids=output_ids, + latent=latent, + multimodal_output={ + "duplex_prompt_token_ids": prompt_ids, + "latent_input_ids": torch.tensor(row_ids).reshape(-1, 1), + "latent_positions": torch.tensor(positions).reshape(-1, 1), + "meta": { + "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}}) + + converted = llm2tts([source], prompt=[{}], _streaming_context=context)[0] + + info = converted["model_intermediate_buffer"] + assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[5:7]) + + +def test_native_duplex_rejects_missing_forwarded_terminal_token() -> None: + prompt_ids = [101, 102] + output_ids = [9304, 21, 9310, 9308] + row_ids = [101, 102, 9304, 21, 999, 999] + latent = torch.arange(len(row_ids) * 4, dtype=torch.float32).reshape(-1, 4) + source = _output( + prompt_ids=prompt_ids, + output_ids=output_ids, + latent=latent, + multimodal_output={ + "duplex_prompt_token_ids": prompt_ids, + "latent_input_ids": torch.tensor(row_ids).reshape(-1, 1), + "latent_positions": torch.arange(len(row_ids)).reshape(-1, 1), + "meta": { + "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}}) + + with pytest.raises(ValueError, match="missing own-token hidden states"): + llm2tts([source], prompt=[{}], _streaming_context=context) + + +_NATIVE_META = { + "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, +} + + +def _native_source(prompt_ids, output_ids, row_ids, positions=None): + latent = torch.arange(len(row_ids) * 4, dtype=torch.float32).reshape(-1, 4) + positions = list(range(len(row_ids))) if positions is None else positions + source = _output( + prompt_ids=prompt_ids, + output_ids=output_ids, + latent=latent, + multimodal_output={ + "duplex_prompt_token_ids": prompt_ids, + "latent_input_ids": torch.tensor(row_ids).reshape(-1, 1), + "latent_positions": torch.tensor(positions).reshape(-1, 1), + "meta": dict(_NATIVE_META), + }, + ) + return source, latent + + +def _native_handoff(source): + context = SimpleNamespace(bridge_states={"duplex": {"epoch": 3, "model_turn_id": 7}}) + return llm2tts([source], prompt=[{}], _streaming_context=context)[0]["model_intermediate_buffer"] + + +def test_native_duplex_ledger_tolerates_async_lookahead_row_after_terminator() -> None: + """An async scheduler forwards the sampled chunk terminator in one extra + frame before the segment stop lands; its row trails the unit's rows and + must not shift the Talker slice.""" + source, latent = _native_source( + prompt_ids=[101, 102], + output_ids=[9304, 21, 22, 9308], + row_ids=[101, 102, 9304, 21, 22, 9308], + ) + + info = _native_handoff(source) + + assert info["ids"]["tts"] == [21, 22] + assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[3:5]) + + +def test_native_duplex_ledger_without_forwarded_terminator() -> None: + """Synchronous scheduling never forwards the chunk terminator.""" + source, latent = _native_source( + prompt_ids=[101, 102], + output_ids=[9304, 21, 22, 9308], + row_ids=[101, 102, 9304, 21, 22], + ) + + info = _native_handoff(source) + + assert info["ids"]["tts"] == [21, 22] + assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[3:5]) + + +def test_native_duplex_turn_eos_row_is_handed_to_talker() -> None: + """<|turn_eos|> is forwarded (it is not a stop token) and the model keeps + sampling until a chunk terminator; the Talker gets the turn_eos row and + nothing after it.""" + source, latent = _native_source( + prompt_ids=[101, 102], + output_ids=[9304, 21, 9310, 77, 9308], + row_ids=[101, 102, 9304, 21, 9310, 77], + ) + + info = _native_handoff(source) + + assert info["ids"]["tts"] == [21, 9310] + assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[3:5]) + assert info["meta"]["turn_end"] is True + + +def test_native_duplex_ledger_anchors_on_unit_after_reinjected_listens() -> None: + """Listens from earlier units sit at the head of the segment delta but + were re-injected as prompt tokens, so they are not contiguous with this + unit's decode rows in the ledger. Anchor on the unit, not the segment.""" + unit_end, unit, audio = 2, 1, 11 + source, latent = _native_source( + prompt_ids=[101, 102], + output_ids=[9303, 9303, 9304, 21, 22, 9308], + row_ids=[9303, unit_end, unit, audio, 9303, unit_end, unit, audio, 9304, 21, 22], + ) + + info = _native_handoff(source) + + assert info["ids"]["tts"] == [21, 22] + assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[9:11]) + + +@pytest.mark.parametrize("previous_terminator_forwarded", [False, True]) +@pytest.mark.parametrize("current_terminator_forwarded", [False, True]) +def test_native_duplex_ledger_prefers_latest_repeat_of_the_unit( + previous_terminator_forwarded: bool, current_terminator_forwarded: bool +) -> None: + previous_ids = [9304, 21, 22] + ([9308] if previous_terminator_forwarded else []) + current_ids = [9304, 21, 22] + ([9308] if current_terminator_forwarded else []) + current_start = len(previous_ids) + 1 + source, latent = _native_source( + prompt_ids=[101, 102], + output_ids=[9304, 21, 22, 9308], + row_ids=[*previous_ids, 555, *current_ids], + positions=[*range(current_start), *range(10, 10 + len(current_ids))], + ) + + info = _native_handoff(source) + + assert info["ids"]["tts"] == [21, 22] + assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[current_start + 1 : current_start + 3]) + + +def test_native_duplex_ledger_rejects_non_contiguous_positions() -> None: + source, _ = _native_source( + prompt_ids=[101, 102], + output_ids=[9304, 21, 22, 9308], + row_ids=[101, 102, 9304, 21, 22], + positions=[0, 1, 2, 3, 9], + ) + + with pytest.raises(ValueError, match="missing own-token hidden states"): + _native_handoff(source) + + +def test_native_duplex_mid_turn_tts_bos_slices_after_boundary() -> None: + source, latent = _native_source( + prompt_ids=[101, 102], + output_ids=[9301, 21, 22, 9308], + row_ids=[101, 102, 9301, 21, 22], + ) + + info = _native_handoff(source) + + assert info["ids"]["tts"] == [21, 22] + assert torch.equal(torch.tensor(info["hidden_states"]["tts"]), latent[3:5]) + + @pytest.mark.parametrize("folded_decisions", [0, 1, 2]) @pytest.mark.parametrize("bos_in_prompt", [False, True]) def test_native_duplex_tts_bos_aligns_after_window_rebuild(folded_decisions, bos_in_prompt) -> None: @@ -152,12 +397,17 @@ def test_native_duplex_tts_bos_aligns_after_window_rebuild(folded_decisions, bos output_ids = [21, 22, 9308] rows = len(prompt_ids) + len(output_ids) - folded_decisions latent = torch.arange(rows * 4, dtype=torch.float32).reshape(rows, 4) + forwarded_ids = ( + prompt_ids[folded_decisions:] + output_ids if bos_in_prompt else prompt_ids + output_ids[folded_decisions:] + ) source = _output( prompt_ids=prompt_ids, output_ids=output_ids, latent=latent, multimodal_output={ "duplex_prompt_token_ids": prompt_ids, + "latent_input_ids": torch.tensor(forwarded_ids).reshape(-1, 1), + "latent_positions": torch.arange(len(forwarded_ids)).reshape(-1, 1), "meta": { "tts_bos_token_id": 9301, "tts_eos_token_id": 9302, diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py index d3da3bba0..6a4b551ab 100644 --- a/tests/worker/test_native_duplex_hooks.py +++ b/tests/worker/test_native_duplex_hooks.py @@ -2916,3 +2916,112 @@ def test_ar_runner_applies_duplex_sampling_before_the_model_sampler(class_name: f"{_MODEL_SAMPLER_CALL}() at line {sampler_lineno}. The hook masks the logits and publishes the " f"row -> session map the sampler reads, so running it afterwards is a silent no-op." ) + + +class _NativeDuplexTokenizer: + eos_token_id = 151705 + unk_token_id = -1 + bad_token_ids: list[int] = [] + all_special_ids: list[int] = [] + + def convert_tokens_to_ids(self, token): + return { + "": 151683, + "": 151684, + "<|listen|>": 151705, + "<|speak|>": 151706, + "<|tts_bos|>": 151703, + "<|tts_eos|>": 151704, + "<|tts_pad|>": 151722, + "<|chunk_eos|>": 151718, + "<|chunk_tts_eos|>": 151721, + "<|turn_eos|>": 151717, + }.get(token, -1) + + +def _native_duplex_model_with_state(): + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.stage0 import ( + _MiniCPMO45Stage0SessionState, + ) + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( + MiniCPMO45OmniForConditionalGeneration, + ) + + session_key = "sid-lookahead" + state = _MiniCPMO45Stage0SessionState(session_id=session_key, current_turn_ended=False) + model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) + model.model_stage = "llm" + model.thinker = SimpleNamespace(get_tokenizer=lambda: _NativeDuplexTokenizer()) + model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state}) + model._minicpmo45_duplex_row_sessions = {0: session_key} + model._minicpmo45_duplex_row_payloads = {0: {"is_speech": True}} + model._minicpmo45_active_duplex_rows = [0] + return model, state + + +def _native_duplex_sampling_metadata(output_token_ids: list[int]): + return SimpleNamespace( + all_greedy=True, + all_random=False, + temperature=torch.tensor([0.0]), + top_k=torch.tensor([1]), + top_p=torch.tensor([1.0]), + generators={}, + prompt_token_ids=torch.tensor([[151683] * 16]), + output_token_ids=[output_token_ids], + ) + + +@pytest.mark.parametrize("terminator", [151718, 151721, 151705]) +def test_minicpmo_stage0_async_lookahead_frame_after_chunk_terminator_is_frozen(terminator): + """The frame async scheduling runs after a sampled chunk terminator is + discarded by the scheduler; it must neither decide anything nor touch the + session state that the next append re-injects.""" + model, state = _native_duplex_model_with_state() + state.pending_terminator_token = terminator + state.last_terminator_token = terminator + state.generated_tokens = [200, 201] + logits = torch.full((1, 151723), -100.0) + logits[0, 1234] = 30.0 # what the stale frame would otherwise pick + + sampled = model.sample(logits, _native_duplex_sampling_metadata([151706, 200, 201, terminator, -1])) + + assert sampled is not None + assert sampled.sampled_token_ids.tolist() == [[terminator]] + assert state.pending_terminator_token == terminator + assert state.last_terminator_token == terminator + assert state.generated_tokens == [200, 201] + assert state.current_turn_ended is False + + +def test_minicpmo_stage0_turn_eos_is_forwarded_not_pending(): + """<|turn_eos|> is fed like text in the official unit loop; only the + chunk terminator sampled afterwards is re-injected on the next append.""" + model, state = _native_duplex_model_with_state() + logits = torch.full((1, 151723), -100.0) + logits[0, 151717] = 30.0 + + sampled = model.sample(logits, _native_duplex_sampling_metadata([151706, 200, 201])) + + assert sampled.sampled_token_ids.tolist() == [[151717]] + assert state.pending_terminator_token is None + assert state.last_terminator_token == 151717 + assert state.current_turn_ended is True + + # The frame that consumes <|turn_eos|> keeps sampling; the chunk + # terminator it picks becomes the pending token for the next append. + logits = torch.full((1, 151723), -100.0) + logits[0, 151718] = 30.0 + sampled = model.sample(logits, _native_duplex_sampling_metadata([151706, 200, 201, 151717])) + + assert sampled.sampled_token_ids.tolist() == [[151718]] + assert state.pending_terminator_token == 151718 + assert state.current_turn_ended is True + + +def test_minicpmo_stage0_stop_tokens_exclude_turn_eos(): + from vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin import _stage0_stop_token_ids + + stop_ids = _stage0_stop_token_ids(_NativeDuplexTokenizer()) + + assert stop_ids == [151718, 151721, 151705] diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index b05c6eab2..7e6258960 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -7,10 +7,12 @@ from unittest.mock import Mock import pytest import torch +from vllm.utils.torch_utils import weak_ref_tensors from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.cudagraph_dispatcher import CUDAGraphMode from vllm.v1.worker.gpu_model_runner import GPUModelRunner +from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner from vllm_omni.worker.gpu_model_runner import ( OmniGPUModelRunner, @@ -21,6 +23,44 @@ from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorMode pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_model_forward_preserves_omni_payload_after_graph_weak_ref(monkeypatch): + hidden = torch.arange(8, dtype=torch.float32).reshape(2, 4) + payload = { + "latent": hidden, + "latent_input_ids": torch.tensor([[21], [22]]), + "latent_positions": torch.tensor([[10], [11]]), + } + original = OmniOutput(text_hidden_states=hidden, multimodal_outputs=payload) + # CUDAGraphWrapper weak-references its outputs both at capture and replay. + # vLLM converts NamedTuple outputs into plain tuples along this path. + # Only the CUDA storage-alias primitive is replaced for this CPU test; + # the upstream container conversion is exercised unchanged. + monkeypatch.setattr("vllm.utils.torch_utils.weak_ref_tensor", lambda value: value) + replay = weak_ref_tensors(original) + runner = object.__new__(OmniGPUModelRunner) + runner.model = SimpleNamespace(have_multimodal_outputs=True) + runner._build_model_kwargs_extra = lambda: {} + monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *_args, **_kwargs: replay) + + output = runner._model_forward() + output_hidden, output_payload = runner.extract_multimodal_outputs(output) + + assert isinstance(output, OmniOutput) + torch.testing.assert_close(output_hidden, hidden) + assert output_payload is payload + + +def test_model_forward_keeps_auxiliary_hidden_tuple(monkeypatch): + hidden = torch.ones(2, 4) + auxiliary = (hidden, hidden.clone(), None, None) + runner = object.__new__(OmniGPUModelRunner) + runner.model = SimpleNamespace(have_multimodal_outputs=True) + runner._build_model_kwargs_extra = lambda: {} + monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *_args, **_kwargs: auxiliary) + + assert runner._model_forward() is auxiliary + + def _runner_for_talker_graph_init( *, talker_mtp_graph_safe: bool | None, diff --git a/vllm_omni/clients/duplex.py b/vllm_omni/clients/duplex.py index 3b7688270..d6e782e79 100644 --- a/vllm_omni/clients/duplex.py +++ b/vllm_omni/clients/duplex.py @@ -111,13 +111,10 @@ __all__ = [ PCM16_SAMPLE_RATE = 16_000 PCM16_BYTES_PER_SAMPLE = 2 -# Server-side model unit boundaries, in cumulative appended audio. -# Stage0 configures the streaming mel processor with first_chunk_ms=1035 and -# chunk_ms=1000; the processor aligns the first chunk down to a hop_length (160 -# samples) multiple, so unit 0 closes at 16480 samples and every later unit -# closes 16000 samples after it. Camera frames must ride the append that closes -# a unit, otherwise Stage0 cannot bind them to that unit's audio. -DUPLEX_FIRST_UNIT_MS = 1030 +# The native PCM reservation closes every 1000 ms, including the first unit. +# Stage0 adds the first mel window's padding internally; the camera frame must +# arrive with the reservation, before that internal padding is applied. +DUPLEX_FIRST_UNIT_MS = 1000 DUPLEX_UNIT_MS = 1000 @@ -806,10 +803,9 @@ class DuplexClientBase(ABC): model unit ``k`` (see :func:`duplex_unit_boundary_ms`), which reproduces the official ``streaming_prefill(audio_waveform=<1 s>, frame_list=[frame])`` pairing: a second of audio and the picture - captured during it enter the same unit. Sending on whole-second - boundaries instead would strand frame 0 on an append that cannot - close a unit yet, and shift every later frame one unit ahead of its - audio. + captured during it enter the same unit. Sending a frame before its + whole-second unit boundary would strand it on an append that cannot + close the unit yet. ``stacked_video_frames`` is the optional parallel track of composites (see ``vllm_omni.experimental.fullduplex.video_stacking``): entry ``k`` diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index a07e2824c..bd6d7abf1 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -653,6 +653,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): continue info[key] = value request.additional_information = info + if replace_snapshot: + # New/resumed requests forward this field with priority + # over additional_information. Replace the prewarm session + # payload too, so chunk zero retains its codec metadata. + request.model_intermediate_buffer = info request.num_computed_tokens = 0 # Empty chunk with more data expected: keep polling. diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py index 44ae9dd9a..4a6ac7606 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py @@ -435,11 +435,14 @@ def _stage0_stop_token_ids(tokenizer: PreTrainedTokenizerBase | None) -> list[in if tokenizer is None: return [] out: list[int] = [] + # ``turn_eos`` is deliberately not a stop token: the official Talker + # conditions on the hidden state produced *after* ``<|turn_eos|>`` is fed, + # so Stage 0 must forward it once and stop on the unit terminator that the + # policy forces on the following step. stop_token_fields = ( "chunk_eos_token_id", "chunk_tts_eos_token_id", "listen_token_id", - "turn_eos_token_id", ) for field in stop_token_fields: token = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS[field] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index 03024bbe5..16269d893 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -64,6 +64,8 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp codec-token deltas for the separate Code2Wav stage. """ + requires_raw_input_tokens = True + @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: if modality.startswith("image"): @@ -92,6 +94,9 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp patch_minicpmo_remote_config(config) self.model_stage = vllm_config.model_config.model_stage + # The Thinker's row ledger needs real token identities even when + # embeddings are supplied, including during CUDA graph capture/replay. + self.requires_raw_input_tokens = self.model_stage == "llm" if self.model_stage == "llm": # Initialize thinker model (image preprocessing + vision encoder + 3D resampler) @@ -562,6 +567,11 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp # Return hidden states with latent in multimodal_outputs for stage_input_processors multimodal_outputs = {"latent": text_hidden_states} + # Keep per-forward row identities alongside the latent payload. + if thinker_input_ids is not None and thinker_positions is not None: + multimodal_outputs["latent_input_ids"] = thinker_input_ids.reshape(-1, 1) + multimodal_outputs["latent_positions"] = thinker_positions.reshape(-1, 1) + runtime_info = kwargs.get("runtime_additional_information") if runtime_info and isinstance(runtime_info, list) and len(runtime_info) > 0: duplex_rows = [] @@ -697,8 +707,21 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp if not native_rows or len(native_rows) != logits.shape[0]: return None + chunk_terminators = self._minicpmo45_chunk_terminator_token_ids(token_ids) + output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or [] sampled_ids: list[int] = [] for row_idx in range(logits.shape[0]): + accepted = output_token_ids[row_idx] if row_idx < len(output_token_ids) else [] + last_accepted = next((int(t) for t in reversed(accepted) if isinstance(t, int) and t >= 0), None) + if last_accepted in chunk_terminators: + # Async scheduling runs one lookahead frame after the chunk + # terminator was sampled but before the scheduler observes + # the segment stop. The scheduler discards this frame's + # token, so decide nothing here: re-emit the terminator and + # leave the model-owned policy state exactly as the accepted + # history left it. The next append re-injects that terminator. + sampled_ids.append(last_accepted) + continue row_logits = logits[row_idx : row_idx + 1].clone() sampled = self._sample_minicpmo45_native_duplex_row( row_logits, @@ -927,7 +950,8 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp segment before the next streaming update, but the official duplex format feeds it (terminator + ) into the KV at every unit boundary, and the model's listen/speak policy depends on seeing its own - past decisions. Non-terminators clear the turn-ended latch.""" + past decisions. Text clears the turn-ended latch; <|turn_eos|> sets it + without becoming pending, because it was forwarded in this unit.""" state = self._minicpmo45_duplex_state_for_row(row_idx) if state is None: return @@ -935,18 +959,26 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp force_listen = isinstance(payload, dict) and payload.get("force_listen") is True listen_id = token_ids.get("listen_token_id", -1) tts_bos_id = token_ids.get("tts_bos_token_id", -1) - chunk_eos_id = token_ids.get("chunk_eos_token_id", -1) - chunk_tts_eos_id = token_ids.get("chunk_tts_eos_token_id", -1) turn_eos_id = token_ids.get("turn_eos_token_id", -1) - terminators = {listen_id, chunk_eos_id, chunk_tts_eos_id, turn_eos_id} - if sampled in terminators: + if sampled in self._minicpmo45_chunk_terminator_token_ids(token_ids): state.pending_terminator_token = int(sampled) state.last_terminator_token = int(sampled) - if sampled == turn_eos_id or (sampled == listen_id and force_listen): + if sampled == listen_id and force_listen: state.current_turn_ended = True with suppress(Exception): state.pending_speech_response_open = False return + if sampled == turn_eos_id: + # Official streaming_generate feeds <|turn_eos|> like text (its + # hidden state conditions the Talker) and keeps sampling until a + # chunk terminator, so nothing is pending for the next append; + # only the turn-ended latch flips. + state.pending_terminator_token = None + state.last_terminator_token = int(sampled) + state.current_turn_ended = True + with suppress(Exception): + state.pending_speech_response_open = False + return if ( sampled == tts_bos_id and getattr(state, "current_turn_ended", True) @@ -965,6 +997,23 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp state.last_terminator_token = None state.current_turn_ended = False + @staticmethod + def _minicpmo45_chunk_terminator_token_ids(token_ids: dict[str, int]) -> set[int]: + """Official ``chunk_terminator_token_ids``: the tokens that close a unit. + + <|turn_eos|> is deliberately absent. It ends the turn but not the + unit: the model forwards it and keeps sampling until one of these. + """ + return { + int(token_id) + for token_id in ( + token_ids.get("listen_token_id", -1), + token_ids.get("chunk_eos_token_id", -1), + token_ids.get("chunk_tts_eos_token_id", -1), + ) + if token_id is not None and int(token_id) >= 0 + } + def _minicpmo45_tokenizer(self): if hasattr(self, "_minicpmo45_tokenizer_cache"): return self._minicpmo45_tokenizer_cache diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index ad329c271..5855cf079 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -681,6 +681,75 @@ def _native_duplex_segment_output_ids( return segment_ids, segment_text, turn_start +def _native_duplex_forwarded_hidden_rows( + mm_output: Mapping[str, object], + thinker_hidden_states: torch.Tensor, + *, + unit_ids: Sequence[int], + slice_start: int, + slice_end: int, + request_id: str, +) -> torch.Tensor: + """Return the rows the Thinker produced by forwarding ``unit_ids[slice]``. + + Stage 0 tags every accumulated hidden row with the token id and position + that produced it (``latent_input_ids`` / ``latent_positions``). Rows are + appended in forward order, so the payload can carry rows that are not this + unit's decode rows: the audio prefill of the append, a later prefill, or + the extra frame an asynchronous scheduler runs after the chunk terminator. + Counting back from the tail therefore cannot identify the span. Instead + anchor on the unit itself: its fed tokens (every decision but possibly the + trailing chunk terminator, which is only forwarded by a lookahead frame) + form one contiguous, position-monotonic block in the ledger. Match that + whole block, take the latest occurrence, and slice inside it. The official + Talker conditions on the hidden state *after* each token is fed, so the + row for ``unit_ids[k]`` is the ledger row whose input is ``unit_ids[k]``. + """ + row_ids = mm_output.get("latent_input_ids") + row_positions = mm_output.get("latent_positions") + if not isinstance(row_ids, torch.Tensor) or not isinstance(row_positions, torch.Tensor): + raise ValueError(f"MiniCPM-o native duplex: missing latent row ledger for request_id={request_id}") + ids = [int(token_id) for token_id in row_ids.reshape(-1).tolist()] + positions = [int(position) for position in row_positions.reshape(-1).tolist()] + if len(ids) != int(thinker_hidden_states.shape[0]) or len(positions) != len(ids): + raise ValueError( + "MiniCPM-o native duplex: latent row ledger mismatch " + f"ids={len(ids)} positions={len(positions)} rows={int(thinker_hidden_states.shape[0])} " + f"for request_id={request_id}" + ) + unit = [int(token_id) for token_id in unit_ids] + if not 0 <= slice_start < slice_end <= len(unit): + raise ValueError( + f"MiniCPM-o native duplex: invalid handoff slice [{slice_start}, {slice_end}) of {len(unit)} tokens" + ) + + def _latest_block(width: int) -> int | None: + block = unit[:width] + for base in range(len(ids) - width, -1, -1): + if ids[base : base + width] != block: + continue + first = positions[base] + if all(positions[base + k] == first + k for k in range(width)): + return base + return None + + # Compare both widths: an older complete unit must not outrank the current + # unit just because its terminator has not been forwarded yet. + latest_base: int | None = None + for width in (len(unit), len(unit) - 1): + if width < slice_end: + continue + base = _latest_block(width) + if base is not None and (latest_base is None or base > latest_base): + latest_base = base + if latest_base is not None: + return thinker_hidden_states[latest_base + slice_start : latest_base + slice_end].to(torch.float32).contiguous() + raise ValueError( + "MiniCPM-o native duplex: missing own-token hidden states for " + f"{unit[slice_start:slice_end]} (unit={unit}, ledger tail={ids[-32:]}) request_id={request_id}" + ) + + def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] | None: bridge_states = getattr(streaming_context, "bridge_states", None) if not isinstance(bridge_states, dict): @@ -857,88 +926,37 @@ def llm2tts( tts_token_ids_slice = tts_hidden_slice = None native_segment_end = False - if tts_bos_idx is not None and thinker_hidden_states.shape[0] > tts_bos_idx: - end_idx = tts_eos_idx if tts_eos_idx is not None else thinker_hidden_states.shape[0] - if is_native_duplex_handoff and tts_eos_idx is not None: - boundary_token = full_token_ids[tts_eos_idx] - native_segment_end = boundary_token in { - special_token_ids.get("chunk_eos_token_id"), - special_token_ids.get("chunk_tts_eos_token_id"), - } - if is_native_duplex_handoff: - # Earlier unforwarded decisions may already be folded into - # the rebuilt prompt while still appearing in this delta. - # Align the explicit tts_bos path by the segment's end, just - # like the native speak/text paths below, rather than adding - # those decisions to the prompt a second time. - out_start = tts_bos_idx - prompt_token_ids_len - out_end = tts_eos_idx - prompt_token_ids_len if tts_eos_idx is not None else len(llm_output_ids) - hidden_base = int(thinker_hidden_states.shape[0]) - len(llm_output_ids) - if hidden_base >= 0 and out_end > out_start: - tts_token_ids_slice = torch.tensor(llm_output_ids[out_start:out_end], dtype=torch.long) - tts_hidden_slice = ( - thinker_hidden_states[hidden_base + out_start : hidden_base + out_end] - .to(torch.float32) - .contiguous() - ) - else: + if not is_native_duplex_handoff: + if tts_bos_idx is not None and thinker_hidden_states.shape[0] > tts_bos_idx: + end_idx = tts_eos_idx if tts_eos_idx is not None else thinker_hidden_states.shape[0] tts_token_ids_slice = torch.tensor(full_token_ids[tts_bos_idx:end_idx], dtype=torch.long) tts_hidden_slice = thinker_hidden_states[tts_bos_idx:end_idx].to(torch.float32).contiguous() - elif is_native_duplex_handoff: - # Official MiniCPM-o duplex does not prefill an assistant - # <|tts_bos|> boundary before generation. A segment delta can - # start with SEVERAL unit decisions (forced/model listens from + else: + # Mirror the official streaming_generate unit loop. A segment delta + # can start with SEVERAL unit decisions (forced/model listens from # chunks that produced no handoff accumulate ahead of the speak), - # so skip the leading listen run, then the speak decision itself; - # hidden states for TTS start after it and stop at the next - # chunk terminator. - listen_id = special_token_ids.get("listen_token_id") - speak_id = special_token_ids.get("speak_token_id") + # so skip the leading listen run; the unit's own decisions start + # there. The first decision (<|speak|>, a <|tts_bos|> boundary, or + # the first text token when no marker is emitted) is fed back but + # is not part of total_hidden_in_unit; TTS conditions on the + # tokens after it, up to the chunk terminator. <|turn_eos|> is not + # a chunk terminator: it is fed like text, its hidden state is + # handed to the Talker, and anything sampled after it is a stale + # tail that must not enter TTS. out_ids = llm_output_ids - j = 0 - while j < len(out_ids) and out_ids[j] == listen_id: - j += 1 - if j < len(out_ids) and out_ids[j] == speak_id: - out_start = j + 1 + listen_id = special_token_ids.get("listen_token_id") + turn_eos_id = special_token_ids.get("turn_eos_token_id") + unit_start = 0 + while unit_start < len(out_ids) and out_ids[unit_start] == listen_id: + unit_start += 1 + if tts_bos_idx is not None: + out_start = max(unit_start, tts_bos_idx - prompt_token_ids_len) + elif unit_start < len(out_ids) and out_ids[unit_start] not in tts_end_ids: + out_start = unit_start + 1 + else: + out_start = None + if out_start is not None: out_end = len(out_ids) - turn_eos_id = special_token_ids.get("turn_eos_token_id") - for idx_t in range(out_start, len(out_ids)): - token_id = out_ids[idx_t] - if turn_eos_id is not None and token_id == turn_eos_id: - # Native duplex trains the talker on the <|turn_eos|> - # embedding itself, but any text sampled after it - # belongs to a stale tail and must not enter TTS. - out_end = idx_t + 1 - break - if token_id in tts_end_ids: - out_end = idx_t - native_segment_end = token_id in { - special_token_ids.get("chunk_eos_token_id"), - special_token_ids.get("chunk_tts_eos_token_id"), - } - break - # Map output indices onto hidden rows by END alignment: the - # leading decision tokens of the delta may ALSO be folded into - # the resumable prompt (they belong to earlier non-forwarded - # segments), so prompt_len + delta over-counts them and - # front-aligned indexing truncates the slice. The hidden - # tensor's last len(out_ids) rows are the delta's rows. - hidden_base = int(thinker_hidden_states.shape[0]) - len(out_ids) - if hidden_base >= 0 and out_end > out_start: - tts_token_ids_slice = torch.tensor(out_ids[out_start:out_end], dtype=torch.long) - tts_hidden_slice = ( - thinker_hidden_states[hidden_base + out_start : hidden_base + out_end] - .to(torch.float32) - .contiguous() - ) - elif j < len(out_ids) and out_ids[j] not in tts_end_ids: - # HF streaming_generate does not require an explicit <|speak|> - # marker. If a unit starts directly with text, the first token - # is fed back into the LLM but is not included in - # total_hidden_in_unit; TTS starts from the following token. - out_start = j + 1 - out_end = len(out_ids) - turn_eos_id = special_token_ids.get("turn_eos_token_id") for idx_t in range(out_start, len(out_ids)): token_id = out_ids[idx_t] if turn_eos_id is not None and token_id == turn_eos_id: @@ -951,13 +969,15 @@ def llm2tts( special_token_ids.get("chunk_tts_eos_token_id"), } break - hidden_base = int(thinker_hidden_states.shape[0]) - len(out_ids) - if hidden_base >= 0 and out_end > out_start: + if out_end > out_start: tts_token_ids_slice = torch.tensor(out_ids[out_start:out_end], dtype=torch.long) - tts_hidden_slice = ( - thinker_hidden_states[hidden_base + out_start : hidden_base + out_end] - .to(torch.float32) - .contiguous() + tts_hidden_slice = _native_duplex_forwarded_hidden_rows( + mm_output, + thinker_hidden_states, + unit_ids=out_ids[unit_start:], + slice_start=out_start - unit_start, + slice_end=out_end - unit_start, + request_id=str(llm_output.request_id), ) 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: diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 80eaf3e67..f6d6173b8 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -2117,6 +2117,16 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner): **model_kwargs, **model_kwargs_extra, ) + # CUDAGraphWrapper's weak_ref_tensors preserves the fields but turns + # NamedTuple outputs into plain tuples. Restore the Omni envelope + # before a model adapter or extraction discards its multimodal fields. + if ( + getattr(self.model, "have_multimodal_outputs", False) + and type(model_output) is tuple + and len(model_output) == len(OmniOutput._fields) + and (model_output[1] is None or isinstance(model_output[1], dict)) + ): + model_output = OmniOutput(*model_output) if not isinstance(model_output, (OmniOutput, IntermediateTensors)) and hasattr(self.model, "make_omni_output"): model_output = self.model.make_omni_output(model_output, **model_kwargs, **model_kwargs_extra) # Cache model output so later sample_tokens can consume multimodal results.