[Bugfix] Fix MiniCPM-o 4.5 native duplex speech alignment (#7974)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: liuqihao <liuqihao970610@gmail.com>

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)
 
 

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")

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):

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

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,

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 {
+            "<unit>": 151683,
+            "</unit>": 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]

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,

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``

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.

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]

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 + </unit>) 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

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:

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.