[3/5][core]refactor communication layer: PR 3 of 5, all other models in non async mode (#3719)

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

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py

@@ -293,7 +293,13 @@ def test_code2wav_forward_finalizes_hift_tail():
     model = object.__new__(CosyVoice3Code2Wav)
     nn.Module.__init__(model)
     model.hift = DummyHiFT()
-    model._forward_mel = lambda **_: torch.ones((1, 80, 8), dtype=torch.float32)
+    forward_mel_calls = []
+
+    def fake_forward_mel(**kwargs):
+        forward_mel_calls.append(kwargs)
+        return torch.ones((1, 80, 8), dtype=torch.float32)
+
+    model._forward_mel = fake_forward_mel
 
     out = model.forward(
         token=torch.tensor([[1, 2, 3]], dtype=torch.int32),
@@ -304,3 +310,4 @@ def test_code2wav_forward_finalizes_hift_tail():
 
     assert out.shape == (1, 1, 8)
     assert model.hift.finalize_calls == [True]
+    assert forward_mel_calls[0]["token_offset_tokens"] == 0

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -138,27 +138,6 @@ def _make_sampling_metadata(
     )
 
 
-def test_split_request_ids_uses_seq_token_counts():
-    CosyVoice3Model, _ = _cosyvoice3_model_and_runner()
-    ids = torch.tensor([10, 11, 12, 13, 14], dtype=torch.long)
-    chunks = CosyVoice3Model._split_request_ids(ids, [2, 2, 2])
-    assert [c.tolist() for c in chunks] == [[10, 11], [12, 13], [14]]
-
-
-def test_split_request_ids_honors_single_request_seq_token_counts():
-    CosyVoice3Model, _ = _cosyvoice3_model_and_runner()
-    ids = torch.tensor([10, 11, 12, 13, 14], dtype=torch.long)
-    chunks = CosyVoice3Model._split_request_ids(ids, [3])
-    assert [c.tolist() for c in chunks] == [[10, 11, 12]]
-
-
-def test_sanitize_codec_tokens_filters_out_of_range():
-    model = _make_code2wav_model()
-    raw = torch.tensor([-1, 0, 3, 4, 99], dtype=torch.long)
-    clean = model._sanitize_codec_tokens(raw)
-    assert clean.tolist() == [0, 3]
-
-
 def test_forward_prefers_token_offset_when_present():
     model = _make_code2wav_model()
 
@@ -265,6 +244,31 @@ def test_forward_uses_non_stream_decode_without_chunk_metadata():
     assert len(model.code2wav.forward_streaming_calls) == 0
     call = model.code2wav.forward_calls[0]
     assert call["token"].tolist() == [[0, 1, 2]]
+    assert call["token_offset_tokens"] == 0
+
+
+def test_forward_uses_non_stream_talker_prefill_offset():
+    model = _make_code2wav_model()
+
+    runtime_info = [
+        {
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "meta": {"talker_prefill_offset": 3},
+        }
+    ]
+
+    model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    assert model.code2wav.forward_calls[0]["token_offset_tokens"] == 3
 
 
 def test_forward_reuses_streaming_cache_state_between_chunks():

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -6,13 +6,19 @@ from types import SimpleNamespace
 
 import torch
 
-from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import talker2code2wav_async_chunk, text2flow
+from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
+    talker2code2wav_async_chunk,
+    text2flow,
+    text2flow_full_payload,
+    text2flow_token_only,
+)
 
 
-def _source_output(request_id: str, prompt_ids: list[int], out_ids: list[int], mm: dict):
+def _source_output(request_id: str, prompt_ids: list[int], out_ids: list[int], mm: dict, finished: bool = True):
     return SimpleNamespace(
         request_id=request_id,
         prompt_token_ids=prompt_ids,
+        finished=finished,
         outputs=[SimpleNamespace(token_ids=out_ids, cumulative_token_ids=out_ids, multimodal_output=mm)],
     )
 
@@ -45,8 +51,8 @@ def _transfer_manager(
 
 def test_text2flow_supports_batched_source_outputs():
     source_outputs = [
-        _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[1, 2]])}),
-        _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[3, 4]])}),
+        _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}),
+        _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[6, 7]])}),
     ]
 
     outputs = text2flow(source_outputs=source_outputs, prompt=None)
@@ -58,6 +64,65 @@ def test_text2flow_supports_batched_source_outputs():
     assert outputs[1]["additional_information"]["ids"]["prompt"] == [20, 21]
 
 
+def test_text2flow_strips_reference_speech_prefix_from_cumulative_ids():
+    source_outputs = [
+        _source_output("req-0", [10, 11], [8, 9, 1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}),
+    ]
+
+    outputs = text2flow(source_outputs=source_outputs, prompt=None)
+
+    assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
+
+
+def test_text2flow_token_only_strips_reference_speech_prefix_from_cumulative_ids():
+    source_outputs = [
+        _source_output(
+            "req-strip",
+            [10, 11],
+            [4, 5, 1, 2, 3],
+            {"embed": {"speech_token": torch.tensor([[4, 5]])}},
+        )
+    ]
+
+    outputs = text2flow_token_only(source_outputs=source_outputs, prompt=None)
+
+    assert len(outputs) == 1
+    assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
+    assert outputs[0]["additional_information"]["ids"]["prompt"] == [10, 11]
+
+
+def test_text2flow_token_only_marks_prompt_trim_for_stop_token_completion():
+    source_outputs = [
+        _source_output(
+            "req-stop",
+            [10, 11],
+            [4, 5, 1, 2, 6562],
+            {"embed": {"speech_token": torch.tensor([[4, 5]])}},
+        )
+    ]
+
+    outputs = text2flow_token_only(source_outputs=source_outputs, prompt=None)
+
+    assert outputs[0]["prompt_token_ids"] == [1, 2, 6562]
+    assert outputs[0]["additional_information"]["meta"]["talker_prefill_offset"] == 2
+
+
+def test_text2flow_full_payload_does_not_send_codec_ids():
+    payload = text2flow_full_payload(
+        None,
+        {
+            "embed.speech_token": torch.tensor([[1, 2]], dtype=torch.long),
+            "codes.audio": torch.tensor([7, 8, 9], dtype=torch.long),
+        },
+        SimpleNamespace(),
+    )
+
+    assert payload is not None
+    assert "codes" not in payload
+    assert "next_stage_prompt_len" not in payload["meta"]
+    assert torch.equal(payload["embed"]["speech_token"], torch.tensor([[1, 2]], dtype=torch.long))
+
+
 def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
     transfer_manager = _transfer_manager()
     request = SimpleNamespace(

tests/model_executor/stage_input_processors/test_qwen2_5_omni_thinker2talker.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Light coverage for qwen2_5_omni.thinker2talker_full_payload.
+
+Covers the finish-reason-aware stop-row trim contract: when the request
+status is FINISHED_STOPPED, the builder must drop one row from the
+accumulated hidden states (vLLM v1 appends the sampled stop token to
+output_token_ids before check_stop, so the trailing hidden-state row
+corresponds to the stop emission and must not reach the talker).
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
+    thinker2talker_full_payload,
+)
+
+
+def _make_request(
+    prompt_token_ids,
+    output_token_ids,
+    status_name: str | None = "FINISHED_STOPPED",
+):
+    status = SimpleNamespace(name=status_name) if status_name else None
+    return SimpleNamespace(
+        request_id="r1",
+        prompt_token_ids=prompt_token_ids,
+        output_token_ids=output_token_ids,
+        all_token_ids=list(prompt_token_ids) + list(output_token_ids),
+        status=status,
+        sampling_params=None,
+    )
+
+
+def test_finished_stopped_trims_one_decode_row():
+    """FINISHED_STOPPED: drop trailing hidden-state row so talker does not
+    consume the stop-emission row.
+    """
+    prompt = [1, 2, 3]
+    output = [10, 11, 12]
+    request = _make_request(prompt, output, status_name="FINISHED_STOPPED")
+    # 6 prompt+output rows + 1 stop-emission row = 7 hidden rows total.
+    hidden = torch.arange(7 * 4, dtype=torch.float32).reshape(7, 4)
+    pooling = {"hidden": hidden}
+
+    payload = thinker2talker_full_payload(transfer_manager=None, pooling_output=pooling, request=request)
+
+    assert payload is not None
+    # ids.output had one trailing stop row dropped: 3 - 1 = 2 remaining.
+    assert payload["ids"]["output"] == output[:-1]
+    # embed.prefill must cover only the prompt rows.
+    assert payload["embed"]["prefill"].shape[0] == len(prompt)
+    # hidden_states.output covers the decode rows minus the dropped stop row.
+    assert payload["hidden_states"]["output"].shape[0] == len(output) - 1
+
+
+def test_finished_length_capped_keeps_all_rows():
+    """FINISHED_LENGTH_CAPPED: no row drop; hidden_states.output covers
+    all decode rows.
+    """
+    prompt = [1, 2, 3]
+    output = [10, 11, 12]
+    request = _make_request(prompt, output, status_name="FINISHED_LENGTH_CAPPED")
+    hidden = torch.arange(6 * 4, dtype=torch.float32).reshape(6, 4)
+    pooling = {"hidden": hidden}
+
+    payload = thinker2talker_full_payload(transfer_manager=None, pooling_output=pooling, request=request)
+
+    assert payload is not None
+    assert payload["ids"]["output"] == output
+    assert payload["embed"]["prefill"].shape[0] == len(prompt)
+    assert payload["hidden_states"]["output"].shape[0] == len(output)
+
+
+def test_missing_hidden_returns_none():
+    """Defensive: pooling_output without "hidden" returns None."""
+    request = _make_request([1, 2], [3], status_name="FINISHED_STOPPED")
+    assert thinker2talker_full_payload(transfer_manager=None, pooling_output={}, request=request) is None
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -168,6 +168,7 @@ def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
 
 
 def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
+    """Full-payload path drops the terminal thinker row before talker prefill."""
     request = SimpleNamespace(
         request_id="thinker",
         prompt_token_ids=[151644, 872],
@@ -187,3 +188,741 @@ def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
     assert payload["embed"]["prefill"].device.type == "cpu"
     assert payload["hidden_states"]["output"].device.type == "cpu"
     assert payload["next_stage_prompt_len"] > 0
+    assert payload["embed"]["prefill"].shape[0] == 2
+    assert payload["hidden_states"]["output"].shape[0] == 2
+
+
+def test_accumulator_replaces_keys_in_replace_set() -> None:
+    """REPLACE-key semantics: subsequent emissions of the same key replace, not append."""
+    from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
+
+    class _StubMixin(OmniConnectorModelRunnerMixin):
+        def __init__(self):
+            self._pending_full_payload_send = {}
+            self._full_payload_replace_keys_cached = frozenset({"model_outputs"})
+
+    stub = _StubMixin()
+    stub.accumulate_full_payload_output(
+        "req1",
+        {
+            "model_outputs": torch.tensor([[1.0, 2.0]]),
+            "hidden_states.output": torch.tensor([[10.0]]),
+        },
+        request=None,
+    )
+    stub.accumulate_full_payload_output(
+        "req1",
+        {
+            "model_outputs": torch.tensor([[3.0, 4.0]]),
+            "hidden_states.output": torch.tensor([[20.0]]),
+        },
+        request=None,
+    )
+    output, _ = stub._materialize_full_payload_entry(stub._pending_full_payload_send["req1"])
+    # model_outputs REPLACED (second value only):
+    assert torch.equal(output["model_outputs"], torch.tensor([[3.0, 4.0]]))
+    # hidden_states.output CONCATENATED:
+    assert torch.equal(output["hidden_states.output"], torch.tensor([[10.0], [20.0]]))
+
+
+def test_accumulator_concat_default_when_no_replace_keys() -> None:
+    """Default semantics: 2-D+ tensors concat across emissions when not in replace_keys."""
+    from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
+
+    class _StubMixin(OmniConnectorModelRunnerMixin):
+        def __init__(self):
+            self._pending_full_payload_send = {}
+            self._full_payload_replace_keys_cached = frozenset()
+
+    stub = _StubMixin()
+    stub.accumulate_full_payload_output(
+        "req1",
+        {"embed.prefill": torch.tensor([[1.0]])},
+        request=None,
+    )
+    stub.accumulate_full_payload_output(
+        "req1",
+        {"embed.prefill": torch.tensor([[2.0]])},
+        request=None,
+    )
+    output, _ = stub._materialize_full_payload_entry(stub._pending_full_payload_send["req1"])
+    assert torch.equal(output["embed.prefill"], torch.tensor([[1.0], [2.0]]))
+
+
+def test_covo_audio_llm2code2wav_token_only_smoke() -> None:
+    """Smoke: covo_audio token-only builder returns placeholder prompts sized to audio_codes count."""
+    # source_outputs is a list of objects with .outputs[0].token_ids
+    from vllm_omni.model_executor.models.covo_audio.config_covo_audio import COVO_AUDIO_TOKEN_INDEX
+    from vllm_omni.model_executor.stage_input_processors.covo_audio import (
+        llm2code2wav_token_only,
+    )
+
+    class _Out:
+        def __init__(self, tids):
+            self.token_ids = tids
+
+    class _Wrapper:
+        def __init__(self, tids):
+            self.outputs = [_Out(tids)]
+
+    # 3 codec tokens + 2 non-codec
+    src = [_Wrapper([COVO_AUDIO_TOKEN_INDEX + 0, COVO_AUDIO_TOKEN_INDEX + 1, COVO_AUDIO_TOKEN_INDEX + 2, 100, 200])]
+    out = llm2code2wav_token_only(src)
+    assert len(out) == 1
+    assert len(out[0]["prompt_token_ids"]) == 3
+    assert out[0]["additional_information"] is None
+
+
+def test_covo_audio_llm2code2wav_full_payload_smoke() -> None:
+    """Smoke: covo_audio producer-side payload builder returns audio_codes + finished."""
+    from types import SimpleNamespace
+
+    from vllm_omni.model_executor.models.covo_audio.config_covo_audio import COVO_AUDIO_TOKEN_INDEX
+    from vllm_omni.model_executor.stage_input_processors.covo_audio import (
+        llm2code2wav_full_payload,
+    )
+
+    req = SimpleNamespace(
+        output_token_ids=[COVO_AUDIO_TOKEN_INDEX + 5, COVO_AUDIO_TOKEN_INDEX + 6, 99],
+    )
+    payload = llm2code2wav_full_payload(None, {}, req)
+    assert payload is not None
+    assert payload["codes"]["audio"] == [5, 6]
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_dynin_omni_token_only_smoke() -> None:
+    """Smoke: dynin_omni token-only builders return placeholders."""
+    from vllm_omni.model_executor.stage_input_processors.dynin_omni import (
+        token2text_to_token2image_token_only,
+    )
+
+    class _Out:
+        def __init__(self, tids, mm=None):
+            self.token_ids = tids
+            self.multimodal_output = mm
+
+    class _Wrapper:
+        def __init__(self, tids, mm=None):
+            self.outputs = [_Out(tids, mm)]
+            self.request_id = "r0"
+
+    class _Stage:
+        def __init__(self, outs):
+            self.engine_outputs = outs
+
+    src = [_Wrapper([10, 11, 12])]
+    out = token2text_to_token2image_token_only([_Stage(src)], [0])
+    assert len(out) == 1
+    assert len(out[0]["prompt_token_ids"]) == 3
+    assert out[0]["additional_information"] is None
+
+
+def test_dynin_omni_full_payload_smoke() -> None:
+    """Smoke: dynin_omni producer-side payload builder returns nested OmniPayload + carries metadata."""
+    from types import SimpleNamespace
+
+    from vllm_omni.model_executor.stage_input_processors.dynin_omni import (
+        token2text_to_token2image_full_payload,
+    )
+
+    pooling = {"token_ids": [1, 2, 3]}
+    req = SimpleNamespace(output_token_ids=[], additional_information={"speaker": ["alice"]})
+    payload = token2text_to_token2image_full_payload(None, pooling, req)
+    assert payload is not None
+    assert payload["codes"]["audio"] == [1, 2, 3]
+    assert payload["meta"]["finished"].item() is True
+    # additional_information is normalized + carried forward (speaker stays list-wrapped).
+    assert payload.get("speaker") == ["alice"]
+
+
+def test_qwen2_5_omni_talker2code2wav_token_only_smoke() -> None:
+    """Smoke: qwen2_5_omni talker→code2wav token_only marker + boundary strip."""
+    from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
+        TALKER_CODEC_END_TOKEN_ID,
+        TALKER_CODEC_START_TOKEN_ID,
+        talker2code2wav_token_only,
+    )
+
+    class _Out:
+        def __init__(self, tids):
+            self.cumulative_token_ids = tids
+
+    class _Wrap:
+        def __init__(self, tids):
+            self.outputs = [_Out(tids)]
+
+    # 3 inner codes wrapped by START + END
+    src = [_Wrap([TALKER_CODEC_START_TOKEN_ID, 10, 11, 12, TALKER_CODEC_END_TOKEN_ID])]
+    out = talker2code2wav_token_only(src)
+    assert len(out) == 1
+    assert len(out[0]["prompt_token_ids"]) == 3
+    assert out[0]["additional_information"] is None
+
+
+def test_qwen2_5_omni_talker2code2wav_full_payload_smoke() -> None:
+    """Smoke: qwen2_5_omni producer-side payload builder strips boundaries."""
+    from types import SimpleNamespace
+
+    from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
+        TALKER_CODEC_END_TOKEN_ID,
+        TALKER_CODEC_START_TOKEN_ID,
+        talker2code2wav_full_payload,
+    )
+
+    req = SimpleNamespace(
+        output_token_ids=[TALKER_CODEC_START_TOKEN_ID, 5, 6, 7, TALKER_CODEC_END_TOKEN_ID],
+    )
+    payload = talker2code2wav_full_payload(None, {}, req)
+    assert payload is not None
+    assert payload["codes"]["audio"] == [5, 6, 7]
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_qwen2_5_omni_talker2code2wav_filters_control_tokens_and_placeholders() -> None:
+    """Qwen2.5 code2wav receives codec ids only, not talker prompt/control ids."""
+    from types import SimpleNamespace
+
+    from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
+        TALKER_CODEC_END_TOKEN_ID,
+        TALKER_CODEC_PAD_TOKEN_ID,
+        TALKER_CODEC_START_TOKEN_ID,
+        talker2code2wav_full_payload,
+        talker2code2wav_token_only,
+    )
+
+    class _Out:
+        def __init__(self, tids):
+            self.cumulative_token_ids = tids
+
+    class _Wrap:
+        def __init__(self, tids):
+            self.outputs = [_Out(tids)]
+
+    raw_ids = [
+        TALKER_CODEC_START_TOKEN_ID,
+        TALKER_CODEC_PAD_TOKEN_ID,
+        5,
+        6,
+        TALKER_CODEC_END_TOKEN_ID,
+        -1,
+        -1,
+    ]
+
+    token_only = talker2code2wav_token_only([_Wrap(raw_ids)])
+    assert len(token_only) == 1
+    assert len(token_only[0]["prompt_token_ids"]) == 4
+
+    payload = talker2code2wav_full_payload(None, {}, SimpleNamespace(output_token_ids=raw_ids))
+    assert payload is not None
+    assert payload["codes"]["audio"] == [5, 6, 6, 6]
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_mimo_audio_llm2code2wav_token_only_smoke() -> None:
+    """Smoke: mimo_audio token-only builder sizes prompt."""
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
+        llm2code2wav_token_only,
+    )
+
+    class _Out:
+        def __init__(self, mm):
+            self.multimodal_output = mm
+
+    class _Wrap:
+        def __init__(self, mm):
+            self.outputs = [_Out(mm)]
+
+    # 3 batch rows of [1, 8, 4]: prepend_and_flatten_colmajor → 3*1*4*9 = 108
+    codes = torch.arange(96, dtype=torch.long).reshape(3, 1, 8, 4)
+    codes = codes.clamp(min=1)  # ensure nonzero so zero-row filter doesn't drop them
+    src = [_Wrap({"codes": {"audio": codes}})]
+    out = llm2code2wav_token_only(src)
+    assert len(out) == 1
+    assert len(out[0]["prompt_token_ids"]) == 108
+    assert out[0]["additional_information"] is None
+
+
+def test_mimo_audio_llm2code2wav_full_payload_smoke() -> None:
+    """Smoke: mimo_audio producer-side payload builder reads flat codes.audio + flattens."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
+        TALKER_CODEC_PAD_TOKEN_ID,
+        llm2code2wav_full_payload,
+    )
+
+    # Simulate accumulator output: 2 steps of [1, 1, 8, 4] CONCAT'd → [2, 1, 8, 4]
+    audio = torch.arange(2 * 1 * 8 * 4, dtype=torch.long).reshape(2, 1, 8, 4)
+    audio = audio.clamp(min=1)  # avoid zero-row drop
+    pooling_output = {"codes.audio": audio}
+    req = SimpleNamespace(output_token_ids=[])
+    payload = llm2code2wav_full_payload(None, pooling_output, req)
+    assert payload is not None
+    assert "codes" in payload and "audio" in payload["codes"]
+    # Flattened length = numel + B*4 (per-batch pad_vec prepended by prepend_and_flatten_colmajor)
+    batch_size = int(audio.shape[0])
+    assert len(payload["codes"]["audio"]) == audio.numel() + batch_size * 4
+    # prepend_and_flatten_colmajor: PAD appears at column start in col-major flatten.
+    # For shape [B=2, 1, 9, 4], each column has 1 PAD then 8 codec vals → PAD at indices 0, 9, 18, 27.
+    out = payload["codes"]["audio"]
+    assert out[0] == TALKER_CODEC_PAD_TOKEN_ID
+    assert out[9] == TALKER_CODEC_PAD_TOKEN_ID
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_mimo_audio_full_payload_nested_fallback() -> None:
+    """Back-compat: full_payload still works if runtime returns nested codes.audio."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
+        llm2code2wav_full_payload,
+    )
+
+    audio = torch.arange(1 * 1 * 8 * 4, dtype=torch.long).reshape(1, 1, 8, 4)
+    audio = audio.clamp(min=1)
+    pooling_output = {"codes": {"audio": audio}}  # nested, not flat
+    req = SimpleNamespace(output_token_ids=[])
+    payload = llm2code2wav_full_payload(None, pooling_output, req)
+    assert payload is not None
+    assert len(payload["codes"]["audio"]) == audio.numel() + int(audio.shape[0]) * 4
+
+
+def test_qwen3_tts_talker2code2wav_token_only_smoke() -> None:
+    """Smoke: qwen3_tts token-only sizes placeholder."""
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+        talker2code2wav_token_only,
+    )
+
+    class _Out:
+        def __init__(self, mm, tids):
+            self.multimodal_output = mm
+            self.cumulative_token_ids = tids
+
+    class _Wrap:
+        def __init__(self, mm, tids):
+            self.outputs = [_Out(mm, tids)]
+            self.finished = True
+
+    # 3 valid codec frames Q=16; non-zero & under codebook size
+    audio = torch.arange(3 * 16, dtype=torch.long).reshape(3, 16) + 1
+    mm = {"codes": {"audio": audio}}
+    src = [_Wrap(mm, list(range(10)))]  # seq_len = 9; 3 < 9, no trim
+    out = talker2code2wav_token_only(src)
+    assert len(out) == 1
+    # Codebook-major flat: 16 * 3 = 48
+    assert len(out[0]["prompt_token_ids"]) == 48
+
+
+def test_qwen3_tts_talker2code2wav_full_payload_smoke() -> None:
+    """Smoke: qwen3_tts full_payload reads flat codes.audio + flattens codebook-major."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+        talker2code2wav_full_payload,
+    )
+
+    # 3 valid codec frames [3, 16] CONCAT'd from per-step emits via flatten
+    audio = torch.arange(3 * 16, dtype=torch.long).reshape(3, 16) + 1
+    pooling_output = {"codes.audio": audio}
+    req = SimpleNamespace(output_token_ids=list(range(10)))  # seq_len = 9
+    payload = talker2code2wav_full_payload(None, pooling_output, req)
+    assert payload is not None
+    assert "codes" in payload and "audio" in payload["codes"]
+    # codebook-major: shape [3, 16] -> [16, 3] -> flatten = 48 entries
+    assert isinstance(payload["codes"]["audio"], torch.Tensor)
+    assert payload["codes"]["audio"].shape == (48,)
+    expected = audio.transpose(0, 1).reshape(-1)
+    assert torch.equal(payload["codes"]["audio"], expected)
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_qwen3_tts_full_payload_with_ref_code() -> None:
+    """Exact: ref_code is prepended (not appended) to audio, ref_code_len trims
+    ref, and the flatten is codebook-major.  Protects against ref-append-position
+    regressions, ref_code_len-not-applied bugs, and flatten-order regressions."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+        talker2code2wav_full_payload,
+    )
+
+    # Audio: 3 frames [3, 16] (no filter drops these — all positive, in-range).
+    audio = torch.arange(3 * 16, dtype=torch.long).reshape(3, 16) + 1
+    # Ref code: 2 frames [2, 16] (already 2-D), distinct value range so we can
+    # detect the prepend ordering.
+    ref = torch.arange(2 * 16, dtype=torch.long).reshape(2, 16) + 100
+    pooling_output = {
+        "codes.audio": audio,
+        "codes.ref": [ref],
+        "meta.ref_code_len": torch.tensor([2], dtype=torch.int32),
+    }
+    req = SimpleNamespace(output_token_ids=list(range(10)))  # seq_len = 9 > 3, no audio crop
+    payload = talker2code2wav_full_payload(None, pooling_output, req)
+    assert payload is not None
+
+    # Exact expected: ref (prepended) + audio (no crop since seq_len > rows), then
+    # transpose [5, 16] -> [16, 5] and flatten row-major (codebook-major).
+    expected = torch.cat([ref, audio], dim=0).transpose(0, 1).reshape(-1)
+    assert torch.equal(payload["codes"]["audio"], expected), (
+        f"codec flatten mismatch -- got first 8 = {payload['codes']['audio'][:8].tolist()}, "
+        f"expected first 8 = {expected[:8].tolist()}"
+    )
+    assert payload["codes"]["audio"].shape == (80,)  # 16 quantizers * (2 ref + 3 audio) frames
+
+    # Sanity guards: first codebook-major column = [ref[0,0], ref[1,0], audio[0,0], ...],
+    # so the prepend order must put 100 before 1.
+    first_col = payload["codes"]["audio"][:5].tolist()
+    assert first_col == [100, 116, 1, 17, 33], (
+        f"first column wrong: {first_col} -- ref likely appended instead of prepended"
+    )
+
+
+def test_qwen3_tts_full_payload_nested_fallback() -> None:
+    """Back-compat: full_payload works if pooler returns un-flattened nested dict."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+        talker2code2wav_full_payload,
+    )
+
+    audio = torch.arange(2 * 16, dtype=torch.long).reshape(2, 16) + 1
+    pooling_output = {"codes": {"audio": audio}}  # nested, not flat
+    req = SimpleNamespace(output_token_ids=list(range(10)))
+    payload = talker2code2wav_full_payload(None, pooling_output, req)
+    assert payload is not None
+    assert isinstance(payload["codes"]["audio"], torch.Tensor)
+    assert payload["codes"]["audio"].shape == (32,)  # 16 * 2
+
+
+def test_qwen3_tts_code2wav_prefers_connector_tensor_payload() -> None:
+    """Code2Wav should consume connector codec tensor instead of placeholder zeros."""
+    import torch
+
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import (
+        _codec_ids_from_payload_or_input,
+    )
+
+    placeholder = torch.zeros(6, dtype=torch.long)
+    codec = torch.arange(12, dtype=torch.long)
+
+    out = _codec_ids_from_payload_or_input(
+        placeholder,
+        {"codes": {"audio": codec}},
+    )
+
+    assert torch.equal(out, codec)
+
+
+def test_qwen3_tts_code2wav_accepts_legacy_list_payload() -> None:
+    """Back-compat: old list full-payloads still override placeholder tokens."""
+    import torch
+
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import (
+        _codec_ids_from_payload_or_input,
+    )
+
+    placeholder = torch.zeros(6, dtype=torch.long)
+
+    out = _codec_ids_from_payload_or_input(
+        placeholder,
+        {"codes": {"audio": [1, 2, 3, 4]}},
+    )
+
+    assert torch.equal(out, torch.tensor([1, 2, 3, 4], dtype=torch.long))
+
+
+def test_qwen3_tts_code2wav_forward_decodes_connector_payload() -> None:
+    """Forward should decode real connector codes, not token-only placeholders."""
+    from collections import Counter
+
+    import torch
+
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import (
+        Qwen3TTSCode2Wav,
+    )
+
+    class _Decoder:
+        def __init__(self):
+            self.last_codes = None
+
+        def chunked_decode(self, codes, **_kwargs):
+            self.last_codes = codes.detach().clone()
+            return codes.sum(dim=1).to(torch.float32)
+
+    decoder = _Decoder()
+    model = Qwen3TTSCode2Wav.__new__(Qwen3TTSCode2Wav)
+    torch.nn.Module.__init__(model)
+    model.decoder = decoder
+    model._num_quantizers = 2
+    model._total_upsample = 1
+    model._output_sample_rate = 24000
+    model._decode_chunk_frames = 300
+    model._decode_left_context_frames = 25
+    model._decode_batch_bucket_frames = []
+    model._decode_batch_max_size = 0
+    model._decode_variable_chunk_batch_min_frames = 326
+    model._logged_codec_stats = True
+    model._logged_malformed_codec_lengths = set()
+    model._batch_stats_enabled = False
+    model._batch_stats_log_every = 0
+    model._batch_stats_forwards = 0
+    model._batch_stats_groups = 0
+    model._batch_stats_requests = 0
+    model._batch_stats_padded_frames = 0
+    model._batch_stats_decoded_frames = 0
+    model._batch_stats_actual_frames = Counter()
+    model._batch_stats_bucket_groups = Counter()
+
+    payload_codes = torch.tensor([1, 3, 2, 4], dtype=torch.long)
+    out = model.forward(
+        input_ids=torch.zeros(4, dtype=torch.long),
+        runtime_additional_information=[{"codes": {"audio": payload_codes}, "meta": {}}],
+    )
+
+    assert decoder.last_codes is not None
+    assert torch.equal(decoder.last_codes, torch.tensor([[[1, 3], [2, 4]]], dtype=torch.long))
+    assert torch.equal(out.multimodal_outputs["model_outputs"][0], torch.tensor([3.0, 7.0]))
+
+
+def test_qwen3_tts_codec_filter_and_crop_edge_cases() -> None:
+    """Regression gate for codec filter + seq_len crop on both token_only and full_payload.
+
+    Mixes valid / all-zero / negative / >=_CODEBOOK_SIZE rows.  Asserts:
+    - Token-only placeholder length matches Q * (#kept rows after crop).
+    - Full-payload codes.audio matches the exact codebook-major flatten
+      of the kept-and-cropped rows.
+
+    Protects against future cleanup reverting the codex P2 #3 (negative
+    codec filter) or the _CODEBOOK_SIZE upper bound.
+    """
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
+        _CODEBOOK_SIZE,
+        talker2code2wav_full_payload,
+        talker2code2wav_token_only,
+    )
+
+    Q = 4  # simulated num_quantizers (default is 16; small here for readability)
+    # 7 rows: valid / all-zero / negative / out-of-range / boundary-valid / valid / valid.
+    audio_rows = [
+        [10, 20, 30, 40],  # row 0: valid -> KEEP
+        [0, 0, 0, 0],  # row 1: all-zero -> DROP
+        [50, -1, 60, 70],  # row 2: negative -> DROP
+        [100, _CODEBOOK_SIZE, 110, 120],  # row 3: >= 2048 -> DROP
+        [200, _CODEBOOK_SIZE - 1, 210, 220],  # row 4: boundary 2047 -> KEEP
+        [300, 310, 320, 330],  # row 5: valid -> KEEP
+        [400, 410, 420, 430],  # row 6: valid -> KEEP
+    ]
+    audio = torch.tensor(audio_rows, dtype=torch.long)
+    kept = [audio_rows[i] for i in (0, 4, 5, 6)]  # 4 rows after filter
+
+    # === token_only path ===
+    # cumulative_token_ids of length 4 -> seq_len = 3 -> crop kept[-3:] = rows {4, 5, 6}
+    class _Out:
+        def __init__(self, ctids, mm):
+            self.cumulative_token_ids = ctids
+            self.multimodal_output = mm
+
+    class _Wrap:
+        def __init__(self, ctids, mm):
+            self.outputs = [_Out(ctids, mm)]
+            self.finished = True
+
+    mm = {"codes": {"audio": audio}, "meta": {}}
+    src = [_Wrap(ctids=[1, 2, 3, 4], mm=mm)]
+    out = talker2code2wav_token_only(src, prompt=None)
+    assert len(out) == 1
+    # No ref_code -> ref_frames = 0; expected prompt_len = Q * (#kept-after-crop) = 4 * 3 = 12
+    assert len(out[0]["prompt_token_ids"]) == Q * 3
+
+    # === full_payload path ===
+    pooling_output = {"codes.audio": audio}
+    req = SimpleNamespace(output_token_ids=[1, 2, 3, 4])  # seq_len = 3
+    payload = talker2code2wav_full_payload(None, pooling_output, req)
+    assert payload is not None
+    # After filter + crop, kept rows = [row4, row5, row6] = [[200,2047,210,220],[300,310,320,330],[400,410,420,430]]
+    # Codebook-major flatten: transpose [3, Q] -> [Q, 3] -> reshape(-1)
+    cropped = torch.tensor(kept[-3:], dtype=torch.long)
+    expected = cropped.transpose(0, 1).reshape(-1)
+    assert torch.equal(payload["codes"]["audio"], expected)
+    # Sanity: confirm the boundary-valid 2047 survived (codex P2 #3 regression guard).
+    assert _CODEBOOK_SIZE - 1 in payload["codes"]["audio"].tolist()
+    # Sanity: confirm no negative or >=_CODEBOOK_SIZE codec id leaked through.
+    assert bool(((payload["codes"]["audio"] >= 0) & (payload["codes"]["audio"] < _CODEBOOK_SIZE)).all())
+
+
+def test_cosyvoice3_text2flow_token_only_smoke() -> None:
+    """Smoke: cosyvoice3 token-only carries ids.prompt only."""
+    from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
+        text2flow_token_only,
+    )
+
+    class _Out:
+        def __init__(self, tids):
+            self.cumulative_token_ids = tids
+            self.multimodal_output = {}
+
+    class _Wrap:
+        def __init__(self, output_tids, prompt_tids):
+            self.outputs = [_Out(output_tids)]
+            self.prompt_token_ids = prompt_tids
+            self.finished = True
+
+    # multimodal_output has embed.* + we expect token_only to preserve it.
+    import torch
+
+    embed = {"speech_token": torch.zeros(2, 4)}
+    src = [_Wrap(output_tids=[10, 20, 30], prompt_tids=[1, 2, 3, 4])]
+    src[0].outputs[0].multimodal_output = {"embed": embed}
+    out = text2flow_token_only(src)
+    assert len(out) == 1
+    # prompt_token_ids is the talker's cumulative_token_ids (real codec tokens, not zeros).
+    assert out[0]["prompt_token_ids"] == [10, 20, 30]
+    # additional_information carries ids.prompt PLUS the original multimodal_output (embed.* still inline).
+    # Heavy embed.* removal pending the model_intermediate_buffer plumbing on the code2wav side.
+    assert out[0]["additional_information"]["ids"]["prompt"] == [1, 2, 3, 4]
+    assert "embed" in out[0]["additional_information"]
+
+
+def test_cosyvoice3_text2flow_full_payload_smoke() -> None:
+    """Smoke: cosyvoice3 producer-side reads flat embed.* keys."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
+        text2flow_full_payload,
+    )
+
+    speech_token = torch.randn(4, 8)
+    speech_feat = torch.randn(4, 16)
+    embedding = torch.randn(1, 32)
+    pooling_output = {
+        "embed.speech_token": speech_token,
+        "embed.speech_feat": speech_feat,
+        "embed.embedding": embedding,
+    }
+    req = SimpleNamespace(external_req_id="r-1")
+    payload = text2flow_full_payload(None, pooling_output, req)
+    assert payload is not None
+    assert "embed" in payload
+    assert torch.equal(payload["embed"]["speech_token"], speech_token)
+    assert torch.equal(payload["embed"]["speech_feat"], speech_feat)
+    assert torch.equal(payload["embed"]["embedding"], embedding)
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_cosyvoice3_text2flow_full_payload_nested_fallback() -> None:
+    """Back-compat: full_payload works if pooler returns un-flattened nested embed."""
+    from types import SimpleNamespace
+
+    import torch
+
+    from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
+        text2flow_full_payload,
+    )
+
+    speech_token = torch.randn(3, 8)
+    pooling_output = {"embed": {"speech_token": speech_token}}  # nested, not flat
+    req = SimpleNamespace(external_req_id="r-2")
+    payload = text2flow_full_payload(None, pooling_output, req)
+    assert payload is not None
+    assert "speech_token" in payload["embed"]
+    assert torch.equal(payload["embed"]["speech_token"], speech_token)
+
+
+def test_cosyvoice3_full_payload_replace_keys_present() -> None:
+    """Confirm _FULL_PAYLOAD_REPLACE_KEYS lists the three embed.* keys."""
+    from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
+        _FULL_PAYLOAD_REPLACE_KEYS,
+    )
+
+    assert _FULL_PAYLOAD_REPLACE_KEYS == frozenset({"embed.speech_token", "embed.speech_feat", "embed.embedding"})
+
+
+def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None:
+    """Smoke: ming_flash_omni token-only carries voice metadata."""
+    from vllm_omni.model_executor.stage_input_processors.ming_flash_omni import (
+        thinker2talker_token_only,
+    )
+
+    class _Out:
+        def __init__(self, text):
+            self.text = text
+
+    class _Wrap:
+        def __init__(self, text):
+            self.outputs = [_Out(text)]
+
+    class _Prompt:
+        def __init__(self, info):
+            self.additional_information = info
+
+    src = [_Wrap("hello world")]
+    prompt = _Prompt({"voice_name": "ZH_FEMALE", "prompt_text": "ref text"})
+    out = thinker2talker_token_only(src, prompt=prompt)
+    assert len(out) == 1
+    assert out[0]["prompt_token_ids"] == [0]  # talker self-tokenizes; dummy id
+    info = out[0]["additional_information"]
+    assert info["text"] == "hello world"
+    assert info["voice_name"] == "ZH_FEMALE"
+    assert info["prompt_text"] == "ref text"
+    assert info["ming_task"] == "omni"
+
+
+def test_qwen2_5_omni_thinker2talker_token_only_smoke() -> None:
+    """Smoke: qwen2_5_omni thinker token-only allocates prompt slots; bulk payload ships via connector."""
+    from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
+        TALKER_CODEC_END_TOKEN_ID,
+        TALKER_CODEC_PAD_TOKEN_ID,
+        TALKER_CODEC_START_TOKEN_ID,
+        thinker2talker_token_only,
+    )
+
+    class _Wrap:
+        def __init__(self, prompt_tids, rid):
+            self.outputs = [object()]
+            self.prompt_token_ids = prompt_tids
+            self.request_id = rid
+
+    class _Prompt(dict):
+        pass
+
+    src = [_Wrap(prompt_tids=[1, 2, 3, 4, 5], rid="r-1")]
+    prompt = [_Prompt(multi_modal_data=None)]
+    out = thinker2talker_token_only(src, prompt=prompt)
+    assert len(out) == 1
+    expected_prompt_len = 1 + len([1, 2, 3, 4, 5]) + 1
+    assert len(out[0]["prompt_token_ids"]) == expected_prompt_len
+    assert out[0]["prompt_token_ids"][0] == TALKER_CODEC_START_TOKEN_ID
+    assert out[0]["prompt_token_ids"][-1] == TALKER_CODEC_END_TOKEN_ID
+    assert all(t == TALKER_CODEC_PAD_TOKEN_ID for t in out[0]["prompt_token_ids"][1:-1])
+    assert out[0]["additional_information"] is None
+
+
+def test_qwen2_5_omni_thinker2talker_full_payload_noop() -> None:
+    """thinker2talker_full_payload returns None when pooling_output lacks the "hidden" key (defensive)."""
+    from vllm_omni.model_executor.stage_input_processors.qwen2_5_omni import (
+        thinker2talker_full_payload,
+    )
+
+    payload = thinker2talker_full_payload(None, {"any": "thing"}, None)
+    assert payload is None