[Refactor][Phase 1]Remove redundant functions and logs (#4986)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>

docs/configuration/pd_disaggregation.md

@@ -123,7 +123,7 @@ After inserting the extra thinker stage, renumber the remaining stages:
     runtime:
       devices: "2"
     engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only
 
   - stage_id: 3
     runtime:
@@ -131,7 +131,7 @@ After inserting the extra thinker stage, renumber the remaining stages:
     engine_args:
       max_num_seqs: 1
     engine_input_source: [2]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+    sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only
 ```
 
 Compared with the default Qwen3-Omni config:

docs/contributing/model/adding_omni_model.md

@@ -475,92 +475,64 @@ thinker_hidden_states = output.multimodal_output["24"]
 
 ### Implementation Example
 
-Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`:
+Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`. Each inter-stage edge should provide a **coherent processor set** rather than a single monolithic function:
+
+| Suffix | Role | Runs when |
+|--------|------|-----------|
+| `*_full_payload` | Worker-side payload producer | `async_chunk=false`; accumulates tensors and ships via connector |
+| `*_async_chunk` | Scheduler-side streaming producer | `async_chunk=true`; emits per-chunk payloads |
+| `*_token_only` | Orchestrator placeholder builder | `async_chunk=false`; allocates downstream prompt slots only |
 
 ```python
-# qwen3_omni.py
+# qwen3_omni.py (Thinker → Talker, non-async path)
 
-def thinker2talker(
-    stage_list: list[Any],
-    engine_input_source: list[int],
+def thinker2talker_token_only(
+    source_outputs: list[Any],
     prompt: OmniTokensPrompt | TextPrompt | None = None,
     requires_multimodal_data: bool = False,
+    streaming_context: Any | None = None,
 ) -> list[OmniTokensPrompt]:
-    """
-    Process thinker outputs to create talker inputs.
-
-    Args:
-        stage_list: List of stage objects
-        engine_input_source: Source stage IDs (typically [0] for thinker)
-        prompt: Original prompt data
-
-    Returns:
-        List of OmniTokensPrompt for talker stage
-    """
-    source_stage_id = engine_input_source[0]
-    thinker_outputs = stage_list[source_stage_id].engine_outputs
-    talker_inputs = []
-
-    for thinker_output in thinker_outputs:
-        output = thinker_output.outputs[0]
-        # Extract thinker embeddings and hidden states
-        thinker_prefill_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
-        thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
-
-        info = {
-            "thinker_prefill_embeddings": thinker_prefill_embeddings,
-            "thinker_hidden_states": thinker_hidden_states,
-            "thinker_sequences": thinker_output.prompt_token_ids + output.token_ids,
-            "thinker_input_ids": thinker_output.prompt_token_ids,
-        }
-
-        talker_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=[0] * computed_length,
-                additional_information=info,
-                multi_modal_data=None,
-            )
-        )
-
-    return talker_inputs
+    """Allocate talker prefill slots; bulk tensors arrive via the connector."""
+    ...
 
 
-def talker2code2wav(
-    stage_list: list[Any],
-    engine_input_source: list[int],
-    prompt: OmniTokensPrompt | TextPrompt | None = None,
-    requires_multimodal_data: bool = False,
-) -> list[OmniTokensPrompt]:
-    """
-    Process talker outputs to create code2wav inputs.
-    """
-    source_stage_id = engine_input_source[0]
-    talker_outputs = stage_list[source_stage_id].engine_outputs
-    code2wav_inputs = []
+def thinker2talker_full_payload(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any],
+    request: OmniEngineCoreRequest,
+) -> dict[str, Any] | None:
+    """Pack accumulated thinker hidden states into OmniPayload for stage-1."""
+    ...
 
-    for talker_output in talker_outputs:
-        output = talker_output.outputs[0]
-        # Extract codec codes
-        codec_codes = (
-            output.multimodal_output["code_predictor_codes"]
-            .to(torch.long)
-            .transpose(0, 1)
-            .cpu()
-            .to(torch.long)
-            .reshape(-1)
-            .tolist()
-        )
 
-        code2wav_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=codec_codes,
-                multi_modal_data=None,
-            )
-        )
-
-    return code2wav_inputs
+def thinker2talker_async_chunk(
+    transfer_manager: Any,
+    multimodal_output: OmniPayload | dict[str, Any],
+    request: OmniEngineCoreRequest,
+    is_finished: bool = False,
+) -> OmniPayloadStruct | None:
+    """Stream thinker rows to talker while async_chunk is enabled."""
+    ...
 ```
 
+Wire these in `pipeline.py`:
+
+```python
+StagePipelineConfig(
+    stage_id=0,
+    custom_process_next_stage_input_func=f"{_PROC}.thinker2talker_full_payload",
+    async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2talker_async_chunk",
+    ...
+),
+StagePipelineConfig(
+    stage_id=1,
+    sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
+    ...
+),
+```
+
+Do **not** add a no-suffix `thinker2talker` when a `sync_process_input_func` is already declared — `_select_processor_funcs()` always prefers the `*_token_only` hook in non-async mode, so the bare function would never run. See `docs/design/rfc_stage_input_processors_refactor.md` for the full contract.
+
 ## Testing
 
 For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -8,7 +8,6 @@ import torch
 
 from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
     talker2code2wav_async_chunk,
-    text2flow,
     text2flow_full_payload,
     text2flow_token_only,
 )
@@ -49,13 +48,13 @@ def _transfer_manager(
     )
 
 
-def test_text2flow_supports_batched_source_outputs():
+def test_text2flow_token_only_supports_batched_source_outputs():
     source_outputs = [
         _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)
+    outputs = text2flow_token_only(source_outputs=source_outputs, prompt=None)
 
     assert len(outputs) == 2
     assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
@@ -64,16 +63,6 @@ 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(

tests/model_executor/stage_input_processors/test_indextts2.py

@@ -7,7 +7,6 @@ import torch
 
 from vllm_omni.model_executor.stage_input_processors.indextts2 import (
     _strip_stop_token,
-    talker2s2mel,
     talker2s2mel_full_payload,
     talker2s2mel_token_only,
 )
@@ -108,118 +107,6 @@ def test_strip_stop_token_all_stop():
     assert lat.shape[1] == 0
 
 
-# ── talker2s2mel ──
-
-
-def test_talker2s2mel_basic():
-    mel_codes = torch.tensor([10, 20, 30, STOP_MEL_TOKEN])
-    latent = torch.randn(4, LATENT_DIM)
-    s_ref = torch.randn(1, 128)
-    ref_mel = torch.randn(1, 80, 200)
-    style = torch.randn(1, 256)
-
-    outputs = talker2s2mel(
-        source_outputs=[
-            _make_talker_output(
-                mel_codes=mel_codes,
-                latent=latent,
-                s_ref=s_ref,
-                ref_mel=ref_mel,
-                style=style,
-            )
-        ]
-    )
-
-    assert len(outputs) == 1
-    info = outputs[0]["additional_information"]
-    assert info["mel_codes"].shape == (1, 3)
-    assert info["latent"].shape == (1, 3, LATENT_DIM)
-    assert info["code_lens"].tolist() == [3]
-    assert "S_ref" in info
-    assert "ref_mel" in info
-    assert "style" in info
-    assert outputs[0]["prompt_token_ids"] == [0]
-
-
-def test_talker2s2mel_skips_unfinished():
-    mel_codes = torch.tensor([10, 20])
-    latent = torch.randn(2, LATENT_DIM)
-
-    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, finished=False)])
-
-    assert len(outputs) == 0
-
-
-def test_talker2s2mel_skips_empty_mel():
-    outputs = talker2s2mel(
-        source_outputs=[
-            _make_talker_output(
-                mel_codes=torch.tensor([]),
-                latent=torch.randn(0, LATENT_DIM),
-            )
-        ]
-    )
-
-    assert len(outputs) == 0
-
-
-def test_talker2s2mel_missing_optional_meta():
-    mel_codes = torch.tensor([10, 20])
-    latent = torch.randn(2, LATENT_DIM)
-
-    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent)])
-
-    assert len(outputs) == 1
-    info = outputs[0]["additional_information"]
-    assert "S_ref" not in info
-    assert "ref_mel" not in info
-    assert "style" not in info
-
-
-def test_talker2s2mel_batch():
-    out0 = _make_talker_output(
-        mel_codes=torch.tensor([1, 2, STOP_MEL_TOKEN]),
-        latent=torch.randn(3, LATENT_DIM),
-        s_ref=torch.randn(1, 64),
-    )
-    out1 = _make_talker_output(
-        mel_codes=torch.tensor([5, 6, 7, 8]),
-        latent=torch.randn(4, LATENT_DIM),
-        ref_mel=torch.randn(1, 80, 100),
-    )
-
-    outputs = talker2s2mel(source_outputs=[out0, out1])
-
-    assert len(outputs) == 2
-    assert outputs[0]["additional_information"]["code_lens"].tolist() == [2]
-    assert outputs[1]["additional_information"]["code_lens"].tolist() == [4]
-
-
-def test_talker2s2mel_tensors_on_cpu():
-    device = torch.device("cpu")
-    mel_codes = torch.tensor([10, 20], device=device)
-    latent = torch.randn(2, LATENT_DIM, device=device)
-    s_ref = torch.randn(1, 64, device=device)
-
-    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, s_ref=s_ref)])
-
-    info = outputs[0]["additional_information"]
-    assert info["latent"].device.type == "cpu"
-    assert info["mel_codes"].device.type == "cpu"
-    assert info["S_ref"].device.type == "cpu"
-
-
-def test_talker2s2mel_no_multimodal_dict():
-    output = SimpleNamespace(
-        finished=True,
-        outputs=[SimpleNamespace(multimodal_output="not_a_dict")],
-    )
-
-    results = talker2s2mel(source_outputs=[output])
-
-    assert len(results) == 0
-
-
 # ── full-payload connector path ──
 
 

tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py

@@ -11,59 +11,40 @@ import torch
 from vllm_omni.model_executor.stage_input_processors import mimo_audio as sip
 from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
     MAX_CODE2WAV_TOKENS,
-    llm2code2wav,
+    llm2code2wav_full_payload,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def _make_source_outputs(codec_codes: torch.Tensor, request_id: str = "req-0"):
-    output = SimpleNamespace(multimodal_output={"codes": {"audio": codec_codes}})
-    """Build minimal source_outputs with one talker output carrying codec codes."""
-    talker_output = SimpleNamespace(outputs=[output], request_id=request_id)
-    return [talker_output]
-
-
-def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
-    """Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated, not passed through."""
-    # prepend_and_flatten_colmajor produces 36 ids per (8, 4) codec frame:
-    # pad adds one row -> (9, 4) per frame, permuted and flattened.
-    # Pick enough frames to comfortably exceed the cap.
+def test_llm2code2wav_full_payload_truncates_when_flat_exceeds_max(caplog):
+    """Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated."""
     frames = (MAX_CODE2WAV_TOKENS // 36) + 100
     codec_codes = torch.ones(frames, 1, 8, 4, dtype=torch.long)
+    request = SimpleNamespace(request_id="req-long")
 
-    source_outputs = _make_source_outputs(codec_codes, request_id="req-long")
-
-    # Attach caplog's handler directly to the module logger so the warning is
-    # captured regardless of propagation (vllm's logger configuration can
-    # interact badly with caplog.at_level's default root-handler path).
     target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.mimo_audio")
     target_logger.addHandler(caplog.handler)
     prev_level = target_logger.level
     target_logger.setLevel(logging.WARNING)
     try:
-        prompts = llm2code2wav(source_outputs)
+        payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request)
     finally:
         target_logger.removeHandler(caplog.handler)
         target_logger.setLevel(prev_level)
 
-    assert len(prompts) == 1
-    assert len(prompts[0]["prompt_token_ids"]) == MAX_CODE2WAV_TOKENS
-    assert any("truncating" in rec.getMessage() for rec in caplog.records), (
-        f"Expected a 'truncating' warning; captured records: {[r.getMessage() for r in caplog.records]}"
-    )
+    assert payload is not None
+    assert len(payload["codes"]["audio"]) == MAX_CODE2WAV_TOKENS
 
 
-def test_llm2code2wav_short_sequence_unchanged():
-    """Short codec sequences are returned without truncation."""
+def test_llm2code2wav_full_payload_short_sequence_unchanged():
     codec_codes = torch.ones(4, 1, 8, 4, dtype=torch.long)
-    source_outputs = _make_source_outputs(codec_codes, request_id="req-short")
+    request = SimpleNamespace(request_id="req-short")
 
-    prompts = llm2code2wav(source_outputs)
+    payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request)
 
-    assert len(prompts) == 1
-    # 4 frames + 1 pad row, flattened col-major → well below the cap
-    assert 0 < len(prompts[0]["prompt_token_ids"]) <= MAX_CODE2WAV_TOKENS
+    assert payload is not None
+    assert 0 < len(payload["codes"]["audio"]) <= MAX_CODE2WAV_TOKENS
 
 
 def test_llm2code2wav_truncation_boundary_constant_matches_yaml():

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -163,7 +163,7 @@ def test_talker2code2wav_full_payload_filters_by_output_token_ids() -> None:
 
     assert payload is not None
     assert payload["codes"]["audio"] == [10, 20, 11, 21, 12, 22]
-    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
+    assert "code_predictor_codes" not in payload
 
 
 def test_talker2code2wav_full_payload_drops_count_matched_terminal_row() -> None:
@@ -204,7 +204,7 @@ def test_talker2code2wav_full_payload_drops_rows_aligned_to_non_codec_ids() -> N
 
     assert payload is not None
     assert payload["codes"]["audio"] == [0, 0, 0]
-    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
+    assert "code_predictor_codes" not in payload
 
 
 def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
@@ -224,7 +224,7 @@ def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
 
     assert payload is not None
     assert payload["codes"]["audio"] == [0, 7, 0, 8, 0, 9]
-    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
+    assert "code_predictor_codes" not in payload
 
 
 def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
@@ -247,7 +247,7 @@ def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
     assert payload["ids"]["all"] == [151644, 872, 3]
     assert payload["embed"]["prefill"].device.type == "cpu"
     assert payload["hidden_states"]["output"].device.type == "cpu"
-    assert payload["next_stage_prompt_len"] > 0
+    assert payload["meta"]["next_stage_prompt_len"] > 0
     assert payload["embed"]["prefill"].shape[0] == 2
     assert payload["hidden_states"]["output"].shape[0] == 2
 

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -12,9 +12,9 @@ from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
     max_ic_for_chunk_size,
 )
 from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
-    talker2code2wav,
     talker2code2wav_async_chunk,
     talker2code2wav_full_payload,
+    talker2code2wav_token_only,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -368,7 +368,15 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
     assert rid in tm.request_payload
 
 
-def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
+def test_non_async_token_only_sizes_placeholder_for_ref_and_audio_frames():
+    """``talker2code2wav_token_only`` only allocates placeholder prompt slots.
+
+    After the connector refactor, actual codec flattening (ref prepend +
+    codebook-major layout) is performed by ``talker2code2wav_full_payload`` on
+    the worker data plane.  The orchestrator hook still derives
+    ``left_context_size`` from stage-0 multimodal_output so Code2Wav can trim
+    reference frames once the connector payload arrives.
+    """
     ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
     audio_codes = torch.tensor(
         [
@@ -387,12 +395,38 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
     )
 
-    prompts = talker2code2wav(stage.engine_outputs)
+    prompts = talker2code2wav_token_only(stage.engine_outputs)
 
     assert len(prompts) == 1
     prompt = prompts[0]
     assert prompt["additional_information"] == {"meta": {"left_context_size": 2}}
-    assert prompt["prompt_token_ids"] == [
+    # 2 ref frames + 2 valid audio frames (zero row filtered), 4 quantizers.
+    assert prompt["prompt_token_ids"] == [0] * (_Q * (2 + 2))
+
+
+def test_full_payload_prepends_ref_code_and_flattens_codebook_major():
+    """Worker producer is authoritative for ref prepend + codec flatten."""
+    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+    audio_codes = torch.tensor(
+        [
+            [0, 0, 0, 0],
+            [1, 2, 3, 4],
+            [5, 6, 7, 8],
+        ],
+        dtype=torch.long,
+    )
+    pooling_output = {
+        "codes.audio": audio_codes,
+        "codes.ref": ref_code,
+        "meta.ref_code_len": 2,
+    }
+    request = SimpleNamespace(request_id="r", output_token_ids=list(range(3)))
+
+    payload = talker2code2wav_full_payload(transfer_manager=None, pooling_output=pooling_output, request=request)
+
+    assert payload is not None
+    assert payload["meta"]["left_context_size"] == 2
+    assert payload["codes"]["audio"].tolist() == [
         9,
         8,
         1,
@@ -433,7 +467,7 @@ def test_non_async_processor_filters_out_of_range_codec_values():
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
     )
 
-    prompts = talker2code2wav(stage.engine_outputs)
+    prompts = talker2code2wav_token_only(stage.engine_outputs)
 
     assert len(prompts) == 1
     prompt = prompts[0]

tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py

@@ -8,61 +8,12 @@ import pytest
 import torch
 
 from vllm_omni.model_executor.stage_input_processors.voxtral_tts import (
-    generator2tokenizer,
     generator2tokenizer_async_chunk,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-# ---- Helpers for generator2tokenizer (non-streaming) ----
-
-
-def _make_stage(engine_outputs):
-    return SimpleNamespace(engine_outputs=engine_outputs)
-
-
-def _make_output(audio_tensors):
-    """Build a single generator output with multimodal_output codes.audio."""
-    return SimpleNamespace(
-        outputs=[SimpleNamespace(multimodal_output={"codes": {"audio": audio_tensors}})],
-    )
-
-
-# ---- Tests for generator2tokenizer ----
-
-
-class TestGenerator2Tokenizer:
-    def test_single_output_single_chunk(self):
-        """Single output with one audio tensor produces correct prompt_token_ids."""
-        audio = torch.tensor([10, 20, 30, 40])
-        stage = _make_stage([_make_output([audio])])
-
-        result = generator2tokenizer(stage.engine_outputs)
-
-        assert len(result) == 1
-        assert result[0]["prompt_token_ids"] == [10, 20, 30, 40]
-        assert result[0]["multi_modal_data"] is None
-
-    def test_2d_audio_tensors_flattened(self):
-        """2D audio tensors (e.g., multi-codebook frames) are flattened."""
-        audio = torch.tensor([[1, 2, 3], [4, 5, 6]])  # (2, 3)
-        stage = _make_stage([_make_output([audio])])
-
-        result = generator2tokenizer(stage.engine_outputs)
-
-        assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6]
-
-    def test_no_outputs_yet_raises(self):
-        """Stage with engine_outputs=None raises RuntimeError."""
-        stage = _make_stage(engine_outputs=None)
-        with pytest.raises(TypeError):
-            generator2tokenizer(stage.engine_outputs)
-
-
-# ---- Helpers for generator2tokenizer_async_chunk (streaming) ----
-
-
 def _req(external_req_id: str, *, finished: bool):
     return SimpleNamespace(
         external_req_id=external_req_id,

tests/test_config_factory.py

@@ -1166,7 +1166,10 @@ class TestQwen3OmniPipeline:
         assert isinstance(s, StagePipelineConfig)
         assert s.input_sources == (0,)
         assert s.sampling_constraints["stop_token_ids"] == [2150]
-        assert s.custom_process_input_func is not None
+        # thinker2talker was removed: sync_process_input_func always wins.
+        assert s.custom_process_input_func is None
+        assert s.sync_process_input_func is not None
+        assert s.sync_process_input_func.endswith("thinker2talker_token_only")
         assert s.custom_process_next_stage_input_func is not None
 
     def test_code2wav(self):
@@ -1180,7 +1183,7 @@ class TestQwen3OmniPipeline:
         assert isinstance(s, StagePipelineConfig)
         assert s.execution_type == StageExecutionType.LLM_GENERATION
         assert s.final_output_type == "audio"
-        assert s.custom_process_input_func is not None
+        assert s.engine_output_type == "audio"
 
 
 class TestQwen2_5OmniPipeline:

vllm_omni/deploy/qwen3_omni_moe.yaml

@@ -47,7 +47,6 @@ stages:
       temperature: 0.9
       top_k: 50
       max_tokens: 4096
-      seed: 42
       repetition_penalty: 1.05
 
   - stage_id: 2
@@ -67,7 +66,6 @@ stages:
       top_p: 1.0
       top_k: -1
       max_tokens: 65536
-      seed: 42
       repetition_penalty: 1.1
 
 platforms:

vllm_omni/engine/stage_engine_core_client.py

@@ -230,7 +230,7 @@ class StageEngineCoreClientBase(StageClientBase):
 
     async def add_request_async(self, request: EngineCoreRequest) -> None:
         """Add request to the stage engine core."""
-        logger.info(
+        logger.debug(
             "[%s] stage-%s [rep-%s] add request: %s",
             self.__class__.__name__,
             self.stage_id,

vllm_omni/model_executor/models/cosyvoice3/pipeline.py

@@ -5,9 +5,8 @@
 Stage 0: Talker   — text prompt → speech tokens (LLM autoregressive).
 Stage 1: Code2Wav — flow-matching decoder → acoustic features → waveform.
   * ``sync_process_input_func`` (``text2flow_token_only``) runs when
-    ``deploy.async_chunk=false``: stage 1 builds full-sequence flow input
-    from the emitted speech tokens. ``text2flow`` is the
-    ``custom_process_input_func`` variant.
+    ``deploy.async_chunk=false``: stage 1 allocates placeholder slots while
+    bulk tensors arrive via ``text2flow_full_payload`` on the connector.
   * ``async_chunk_process_next_stage_input_func`` runs when
     ``deploy.async_chunk=true``: stage 0 streams codec chunks to stage 1
     through the shared-memory connector.
@@ -47,7 +46,6 @@ COSYVOICE3_PIPELINE = PipelineConfig(
             final_output=True,
             final_output_type="audio",
             engine_output_type="latent",
-            custom_process_input_func=f"{_PROC}.text2flow",
             sync_process_input_func=f"{_PROC}.text2flow_token_only",
         ),
     ),

vllm_omni/model_executor/models/mimo_audio/pipeline.py

@@ -56,7 +56,6 @@ MIMO_AUDIO_PIPELINE = PipelineConfig(
             final_output=True,
             final_output_type="audio",
             engine_output_type="audio",
-            custom_process_input_func=f"{_PROC}.llm2code2wav",
             sync_process_input_func=f"{_PROC}.llm2code2wav_token_only",
             sampling_constraints={"detokenize": False},
         ),

vllm_omni/model_executor/models/qwen3_omni/pipeline.py

@@ -51,7 +51,6 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             input_sources=(0,),
             hf_config_name="talker_config",
             engine_output_type="latent",
-            custom_process_input_func=f"{_PROC}.thinker2talker",
             sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
             custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"),
             async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
@@ -69,7 +68,6 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             final_output_type="audio",
             hf_config_name="thinker_config",
             engine_output_type="audio",
-            custom_process_input_func=f"{_PROC}.talker2code2wav",
             sampling_constraints={"detokenize": True},
         ),
     ),

vllm_omni/model_executor/stage_input_processors/cosyvoice3.py

@@ -119,30 +119,6 @@ def _decode_additional_information(raw_info: Any) -> dict[str, Any]:
     return decoded
 
 
-def text2flow(
-    source_outputs: list[Any],
-    prompt: OmniTokensPrompt | TextPrompt = None,
-    _requires_multimodal_data: bool = True,
-):
-    """Build stage-1 inputs by prefixing stage-0 prompt ids to its outputs."""
-    engine_inputs: list[OmniTokensPrompt] = []
-    for source_output in source_outputs:
-        output = source_output.outputs[0]
-        multi_modal_data = output.multimodal_output
-        if multi_modal_data is None:
-            raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}")
-
-        prefix_ids = _ensure_list(source_output.prompt_token_ids)
-        raw_output_ids = _ensure_list(output.cumulative_token_ids)
-        prompt_speech_ids = _prompt_speech_token_ids(multi_modal_data)
-        output_ids = _strip_prompt_prefix(raw_output_ids, prefix_ids)
-        output_ids = _strip_prompt_prefix(output_ids, prompt_speech_ids)
-        additional_info = dict(multi_modal_data)
-        additional_info.setdefault("ids", {})["prompt"] = prefix_ids
-        engine_inputs.append(OmniTokensPrompt(prompt_token_ids=output_ids, additional_information=additional_info))
-    return engine_inputs
-
-
 def talker2code2wav_async_chunk(
     transfer_manager: Any,
     multimodal_output: dict[str, Any] | None,

vllm_omni/model_executor/stage_input_processors/fish_speech.py

@@ -114,37 +114,6 @@ def _extract_last_frame(multimodal_output: dict[str, Any]) -> torch.Tensor | Non
     raise ValueError(f"Invalid audio_codes shape for Fish Speech async_chunk: {tuple(audio_codes.shape)}")
 
 
-def slow_ar_to_dac_decoder(
-    source_outputs: list[Any],
-    _prompt: Any = None,
-    _requires_multimodal_data: bool = False,
-) -> list[Any]:
-    """Non-async processor: wait for Slow AR to finish, then pass all codes to DAC decoder."""
-    from vllm_omni.inputs.data import OmniTokensPrompt
-
-    slow_ar_outputs = source_outputs
-    dac_inputs: list[OmniTokensPrompt] = []
-
-    for output in slow_ar_outputs:
-        out = output.outputs[0]
-        # audio_codes shape: [num_frames, num_codebooks]
-        audio_codes = out.multimodal_output["audio_codes"].to(torch.long)
-        # Filter zero-padded frames.
-        valid_mask = audio_codes.any(dim=1)
-        audio_codes = audio_codes[valid_mask]
-        # Codebook-major flat: [num_codebooks * num_frames]
-        codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist()
-        dac_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=codec_codes,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-                additional_information=None,
-            )
-        )
-    return dac_inputs
-
-
 def slow_ar_to_dac_decoder_async_chunk(
     transfer_manager: Any,
     multimodal_output: dict[str, Any] | None,

vllm_omni/model_executor/stage_input_processors/indextts2.py

@@ -143,65 +143,6 @@ def _request_id(request: Any) -> str:
     return str(getattr(request, "external_req_id", None) or getattr(request, "request_id", "?"))
 
 
-def talker2s2mel(
-    source_outputs: list[Any],
-    prompt: Any = None,
-    _requires_multimodal_data: bool = False,
-) -> list[Any]:
-    """Legacy orchestrator path: collect all Stage 0 output, format Stage 1 input."""
-    _old_nt = torch.get_num_threads()
-    torch.set_num_threads(1)
-    try:
-        return _talker2s2mel_impl(source_outputs)
-    finally:
-        torch.set_num_threads(_old_nt)
-
-
-def _talker2s2mel_impl(source_outputs: list[Any]) -> list[Any]:
-    from vllm_omni.inputs.data import OmniTokensPrompt
-
-    s2mel_inputs: list[OmniTokensPrompt] = []
-
-    for i, talker_output in enumerate(source_outputs):
-        if not talker_output.finished:
-            continue
-
-        output = talker_output.outputs[0]
-        mm = output.multimodal_output
-        if not isinstance(mm, dict):
-            logger.warning("Talker output %d has no multimodal_output dict", i)
-            continue
-
-        codes_dict = mm.get("codes", {})
-        mel_codes = codes_dict.get("mel")
-        if not isinstance(mel_codes, torch.Tensor) or mel_codes.numel() == 0:
-            logger.warning("Talker output %d has empty mel_codes", i)
-            continue
-        mel_codes = mel_codes.to(torch.long)
-
-        hs_dict = mm.get("hidden_states", {})
-        latent = hs_dict.get("latent")
-        if not isinstance(latent, torch.Tensor) or latent.numel() == 0:
-            logger.warning("Talker output %d has empty latent", i)
-            continue
-
-        meta = mm.get("meta", {})
-        if not isinstance(meta, dict):
-            meta = {}
-
-        additional_information = _build_s2mel_additional_information(mel_codes, latent, meta, context="talker2s2mel")
-        s2mel_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=[0],
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-                additional_information=additional_information,
-            )
-        )
-
-    return s2mel_inputs
-
-
 def talker2s2mel_token_only(
     source_outputs: list[Any],
     prompt: Any = None,

vllm_omni/model_executor/stage_input_processors/mimo_audio.py

@@ -1,7 +1,6 @@
 from typing import Any
 
 import torch
-from vllm.inputs import TextPrompt
 from vllm.logger import init_logger
 
 from vllm_omni.data_entry_keys import (
@@ -10,7 +9,6 @@ from vllm_omni.data_entry_keys import (
     OmniPayload,
     OmniPayloadStruct,
 )
-from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_CODEC_PAD_TOKEN_ID
 
 logger = init_logger(__name__)
@@ -230,88 +228,6 @@ def llm2code2wav_async_chunk(
     )
 
 
-def llm2code2wav(
-    source_outputs: list[Any],
-    prompt: OmniTokensPrompt | TextPrompt | None = None,
-    requires_multimodal_data: bool = False,
-) -> list[OmniTokensPrompt]:
-    """
-    Process talker outputs to create code2wav inputs.
-
-    Workflow:
-    1. Extract talker's codec code outputs (8-layer RVQ codes)
-    2. Flatten codes for code2wav input
-    3. Package for code2wav stage
-
-    Args:
-        prompt: Original prompt data
-        requires_multimodal_data: Whether multimodal data is required
-
-    Returns:
-        List of OmniTokensPrompt for code2wav stage
-    """
-    talker_outputs = source_outputs
-    code2wav_inputs = []
-
-    # Process each talker output
-    for i, talker_output in enumerate(talker_outputs):
-        output = talker_output.outputs[0]
-
-        # Extract codec codes from talker output
-        # Expected shape: [8, seq_len] (8-layer RVQ codes)
-        mm = output.multimodal_output
-        mm_codes = mm.get("codes", {})
-        mm_hs = mm.get("hidden_states", {})
-        if "audio" in mm_codes:
-            codec_codes = mm_codes["audio"].to(torch.long)  # [seq_batch_size, 1, 8, 4]
-            is_all_zero = (codec_codes == 0).all(dim=(1, 2, 3))
-            non_zero_indices = (~is_all_zero).nonzero(as_tuple=True)[0]
-            if len(non_zero_indices) == 0:
-                # All codec codes are zero - skip this request with a warning
-                request_id = getattr(talker_output, "request_id", f"unknown_{i}")
-                logger.warning(
-                    "Skipping request %s: all codec codes are zero (empty output from Stage-0). "
-                    "This may indicate the model failed to generate valid audio codes.",
-                    request_id,
-                )
-            else:
-                if len(non_zero_indices) < codec_codes.shape[0]:
-                    codec_codes = codec_codes[non_zero_indices]
-        elif "output" in mm_hs and "audio" not in mm_codes:
-            codec_codes = torch.zeros(1, 1, 8, 4, dtype=torch.long)
-        else:
-            raise ValueError(f"Invalid multimodal_output: {output.multimodal_output}")
-
-        pad_vec = torch.tensor([TALKER_CODEC_PAD_TOKEN_ID] * 4)
-
-        code_final = prepend_and_flatten_colmajor(codec_codes, pad_vec)
-        code_final = code_final.tolist()
-
-        # Guard against flattened sequences longer than code2wav's max_model_len.
-        # Without this, add_request raises ``could not broadcast input array
-        # from shape (N,) into shape (max_model_len,)`` and kills the engine
-        # core (see issue #2683). Mirrors the offline end2end.py safeguard.
-        if len(code_final) > MAX_CODE2WAV_TOKENS:
-            request_id = getattr(talker_output, "request_id", f"unknown_{i}")
-            logger.warning(
-                "Request %s: code_final len=%d > MAX_CODE2WAV_TOKENS=%d, truncating.",
-                request_id,
-                len(code_final),
-                MAX_CODE2WAV_TOKENS,
-            )
-            code_final = code_final[:MAX_CODE2WAV_TOKENS]
-
-        code2wav_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=code_final,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-            )
-        )
-
-    return code2wav_inputs
-
-
 # ============================================================================
 # Worker-connector data plane (non-async-chunk path).
 # AR runner's `flatten_payload` converts the model emit
@@ -329,8 +245,9 @@ _FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset()
 def _filter_zero_codec_rows(codec_codes: torch.Tensor) -> torch.Tensor:
     """Drop zero-padded codec rows from a 4-D `[N, 1, 8, 4]` tensor.
 
-    Mirrors the zero-row filter in the orchestrator-path `llm2code2wav`
-    body (see this file's ``llm2code2wav`` around line 224).
+    Shared by the sync placeholder builder (``llm2code2wav_token_only``) and
+    the full-payload producer (``llm2code2wav_full_payload``) so both size the
+    downstream codec sequence off the same non-zero frames.
     """
     if codec_codes.ndim != 4 or codec_codes.numel() == 0:
         return codec_codes

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -10,7 +10,6 @@ from typing import Any
 
 import torch
 from vllm.inputs import TextPrompt
-from vllm.platforms import current_platform
 
 from vllm_omni.data_entry_keys import (
     CodesStruct,
@@ -25,9 +24,7 @@ from vllm_omni.data_entry_keys import (
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.stage_input_processors.tts_utils import (
-    extract_language_from_prompt,
     extract_language_from_request,
-    extract_speaker_from_prompt,
     extract_speaker_from_request,
 )
 
@@ -569,25 +566,31 @@ def thinker2talker_full_payload(
         output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or [])
         all_token_ids = list(prompt_token_ids) + list(output_token_ids)
 
-    # Trim the trailing stop-token row from the accumulated thinker output.
-    # The accumulator captures one hidden-state row per executed thinker
-    # forward (prefill + every decode step including the one that emitted
-    # the stop_token), so for a finished request thinker_emb has exactly one
-    # row more than the rows the talker should consume. async_chunk's
-    # chunk-0 path naturally captures only the prefill / non-stop portion,
-    # which is why the [async_chunk] parametrization passes while [default]
-    # over-generates one codec frame on short outputs (e.g.
-    # test_one_word_prompt_001[default]: audio extends "London" with
-    # spurious phonemes).
-    if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 0:
+    # Drop the terminal stop-token row only when more than one row was
+    # accumulated; trimming a single row would ship 0 conditioning tensors
+    # while ids still has tokens and break talker prefill alignment.
+    if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 1:
         thinker_emb_prefill = thinker_emb[:-1]
     else:
         thinker_emb_prefill = thinker_emb
-    if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 0:
+    if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 1:
         thinker_hid_prefill = thinker_hid[:-1]
     else:
         thinker_hid_prefill = thinker_hid
 
+    emb_rows = int(thinker_emb_prefill.shape[0]) if isinstance(thinker_emb_prefill, torch.Tensor) else 0
+    hid_rows = int(thinker_hid_prefill.shape[0]) if isinstance(thinker_hid_prefill, torch.Tensor) else 0
+    if len(all_token_ids) > 0 and (emb_rows == 0 or hid_rows == 0):
+        logger.warning(
+            "thinker2talker_full_payload: empty thinker conditioning for req=%s "
+            "(ids_len=%s embed_rows=%s hidden_rows=%s); withholding payload.",
+            rid,
+            len(all_token_ids),
+            emb_rows,
+            hid_rows,
+        )
+        return None
+
     payload: OmniPayload = {
         "embed": {
             "prefill": thinker_emb_prefill.detach().cpu(),
@@ -599,7 +602,6 @@ def thinker2talker_full_payload(
         "ids": {"all": list(all_token_ids), "prompt": list(prompt_token_ids)},
         "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
     }
-    payload["next_stage_prompt_len"] = _compute_talker_prompt_ids_length(payload, device="cpu")
     speaker = extract_speaker_from_request(request)
     if speaker is not None:
         payload["speaker"] = speaker
@@ -609,154 +611,34 @@ def thinker2talker_full_payload(
     return payload
 
 
-def thinker2talker(
-    source_outputs: list[Any],
-    prompt: OmniTokensPrompt | TextPrompt | None = None,
-    requires_multimodal_data: bool = False,
-    streaming_context: Any | None = None,
-) -> list[OmniTokensPrompt]:
-    """
-    Process thinker outputs to create talker inputs.
-
-    Workflow:
-    1. Extract thinker's text generation outputs (token IDs + hidden states)
-    2. Split hidden states into: prompt embeddings + generated embeddings
-    3. Package for talker with additional information
-
-    In PD disaggregation mode, merges prefill-stage prompt embeddings with
-    decode-stage generated embeddings before handing off to the talker.
-
-    Args:
-        prompt: Original prompt data
-        requires_multimodal_data: Whether multimodal data is required
-
-    Returns:
-        List of OmniTokensPrompt for talker stage
-    """
-    thinker_outputs = source_outputs
-    talker_inputs: list[OmniTokensPrompt] = []
-
-    device = torch.device(current_platform.device_type)
-
-    # Process each thinker output
-    for i, thinker_output in enumerate(thinker_outputs):
-        output = thinker_output.outputs[0]
-        req_id = str(getattr(thinker_output, "request_id", f"idx-{i}"))
-        prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids)
-        output_ids = _ensure_list(output.cumulative_token_ids)
-        is_streaming_session = bool(getattr(streaming_context, "enabled", False))
-        if is_streaming_session:
-            prompt_token_ids, output_ids = _get_streaming_talker_tokens(
-                req_id,
-                prompt_token_ids,
-                output_ids,
-                getattr(streaming_context, "new_prompt_len_snapshot", None),
-                streaming_context,
-                clear_state=bool(getattr(thinker_output, "finished", False)),
-            )
-        thinker_sequences = prompt_token_ids + output_ids
-        thinker_input_ids = prompt_token_ids
-        new_seq_length = len(prompt_token_ids + output_ids) - 1
-        thinker_mm_raw = getattr(output, "multimodal_output", None)
-        if not isinstance(thinker_mm_raw, Mapping):
-            logger.debug("thinker2talker: skip req=%s due to empty multimodal_output", req_id)
-            continue
-        thinker_mm: OmniPayload = thinker_mm_raw
-        mm_hs = thinker_mm.get("hidden_states", {})
-        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, Mapping) else {}
-        emb_layer = _layer_tensor(mm_layers, _EMBED_LAYER_KEY)
-        hid_layer = _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY)
-        if emb_layer is None or hid_layer is None:
-            logger.debug("thinker2talker: skip req=%s due to missing hidden-state layers", req_id)
-            continue
-        thinker_emb = emb_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:]
-        thinker_hid = hid_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:]
-
-        prefill_mm: dict[str, Any] | None = None
-        prefill_mm = _get_prefill_multimodal_output(req_id, streaming_context)
-
-        if prefill_mm is not None:
-            expected_total = len(prompt_token_ids) + len(output_ids)
-            try:
-                thinker_emb, thinker_hid = _merge_pd_embeddings(
-                    thinker_emb, thinker_hid, prefill_mm, device, expected_total=expected_total
-                )
-            except Exception as exc:
-                logger.warning("[PD] Could not merge prefill embeddings: %s", exc)
-
-        payload = OmniPayloadStruct(
-            embed=EmbeddingsStruct(
-                prefill=thinker_emb,
-                tts_bos=_resolve_tts_token_embedding(
-                    "tts_bos", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
-                ),
-                tts_eos=_resolve_tts_token_embedding(
-                    "tts_eos", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
-                ),
-                tts_pad=_resolve_tts_token_embedding(
-                    "tts_pad", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
-                ),
-            ),
-            hidden_states=HiddenStatesStruct(output=thinker_hid),
-            ids=IdsStruct(all=thinker_sequences, prompt=thinker_input_ids),
-            speaker=extract_speaker_from_prompt(prompt, index=i),
-            language=extract_language_from_prompt(prompt, index=i),
-        )
-        info = to_dict(payload)
-        prompt_len = _compute_talker_prompt_ids_length(info, device=device)
-
-        talker_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=[0] * prompt_len,
-                additional_information=info,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-            )
-        )
-
-    return talker_inputs
-
-
 def thinker2talker_token_only(
     source_outputs: list[Any],
     prompt: OmniTokensPrompt | TextPrompt | None = None,
     requires_multimodal_data: bool = False,
     streaming_context: Any | None = None,
 ) -> list[OmniTokensPrompt]:
-    """Non-async-chunk Stage-1 input builder for the connector data plane.
+    """Orchestrator-side placeholder builder for Stage-1 (Talker) when
+    ``async_chunk=False``.
 
-    The worker connector (Stage-0 ``thinker2talker_full_payload`` →
-    ``_sync_local_stage_payloads``) supplies the bulk talker conditioning
-    tensors (embed / hidden_states / ids) via ``model_intermediate_buffer``.
-    The orchestrator only needs to ship a placeholder prefill prompt of the
-    correct length so the scheduler can allocate KV-cache slots.
+    After the communication-layer refactor, this function only allocates a
+    placeholder ``prompt_token_ids`` of the correct length so the scheduler can
+    reserve KV-cache slots. It does **not** forward bulk tensors or voice
+    metadata.
 
-    Small per-request voice metadata (``speaker`` / ``language``) is forwarded
-    here from the user prompt so the worker's line-408 buffer seed picks it
-    up. The connector-side ``extract_speaker_from_request`` reads the
-    strongly-typed ``request.additional_information.entries["speaker"]`` which
-    currently does not always round-trip the user-supplied voice; until that
-    plumbing is normalized, providing the small fields directly preserves
-    voice selection (regression discovered on Buildkite 9668:
-    ``test_speaker_002[default]`` lost the preset voice).
+    Bulk talker conditioning (embed / hidden_states / ids) and per-request
+    voice metadata (``speaker`` / ``language``) are packed by Stage-0
+    ``thinker2talker_full_payload`` via ``extract_speaker_from_request`` /
+    ``extract_language_from_request``, sent over the connector, and merged into
+    the Talker worker's ``model_intermediate_buffer`` on recv. The orchestrator
+    path therefore sets ``additional_information=None``.
+
+    ``prompt`` / ``requires_multimodal_data`` are kept for call-site signature
+    compatibility with other orchestrator input processors; they are unused.
     """
     talker_inputs: list[OmniTokensPrompt] = []
     for i, thinker_output in enumerate(source_outputs):
         output = thinker_output.outputs[0]
         req_id = str(getattr(thinker_output, "request_id", f"idx-{i}"))
-        # Skip-on-missing parity with thinker2talker_full_payload: if the
-        # connector builder would drop this request (no MM dict or missing
-        # hidden-state layers), do the same here so the worker buffer
-        # presence agrees with the orchestrator's scheduling decision.
-        thinker_mm_raw = getattr(output, "multimodal_output", None)
-        if not isinstance(thinker_mm_raw, Mapping):
-            logger.debug("thinker2talker_token_only: skip req=%s due to empty multimodal_output", req_id)
-            continue
-        mm_hs = thinker_mm_raw.get("hidden_states", {})
-        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, Mapping) else {}
-        if _layer_tensor(mm_layers, _EMBED_LAYER_KEY) is None or _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY) is None:
-            logger.debug("thinker2talker_token_only: skip req=%s due to missing hidden-state layers", req_id)
-            continue
         prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids)
         output_ids = _ensure_list(output.cumulative_token_ids)
         is_streaming_session = bool(getattr(streaming_context, "enabled", False))
@@ -773,21 +655,10 @@ def thinker2talker_token_only(
         thinker_input_ids = prompt_token_ids
         info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}}
         prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu")
-
-        # Forward only small voice metadata; bulk tensors come from the
-        # connector path via _sync_local_stage_payloads.
-        small_info: dict[str, Any] = {}
-        speaker = extract_speaker_from_prompt(prompt, index=i)
-        if speaker is not None:
-            small_info["speaker"] = speaker
-        language = extract_language_from_prompt(prompt, index=i)
-        if language is not None:
-            small_info["language"] = language
-
         talker_inputs.append(
             OmniTokensPrompt(
                 prompt_token_ids=[0] * prompt_len,
-                additional_information=(small_info if small_info else None),
+                additional_information=None,
                 multi_modal_data=None,
                 mm_processor_kwargs=None,
             )
@@ -949,59 +820,5 @@ def talker2code2wav_full_payload(
     )
     return {
         "codes": {"audio": codec_codes},
-        "code_predictor_codes": codec_codes,
         "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
     }
-
-
-def talker2code2wav(
-    source_outputs: list[Any],
-    _prompt: OmniTokensPrompt | TextPrompt | None = None,
-    _requires_multimodal_data: bool = False,
-    streaming_context: Any | None = None,
-) -> list[OmniTokensPrompt]:
-    """
-    Process talker outputs to create code2wav inputs.
-
-    Workflow:
-    1. Extract talker's codec code outputs (8-layer RVQ codes)
-    2. Flatten codes for code2wav input
-    3. Package for code2wav stage
-
-    Args:
-    Returns:
-        List of OmniTokensPrompt for code2wav stage
-    """
-    talker_outputs = source_outputs
-    code2wav_inputs: list[OmniTokensPrompt] = []
-    # Process each talker output
-    for i, talker_output in enumerate(talker_outputs):
-        output = talker_output.outputs[0]
-        req_id = str(getattr(talker_output, "request_id", f"idx-{i}"))
-        cur_seq_len = len(output.cumulative_token_ids) - 1
-        seq_len = cur_seq_len
-        is_streaming_session = bool(getattr(streaming_context, "enabled", False))
-        if is_streaming_session:
-            seq_len = _get_streaming_codec_delta_len(cur_seq_len, req_id, talker_output, streaming_context)
-        mm_raw = getattr(output, "multimodal_output", None)
-        if not isinstance(mm_raw, Mapping):
-            logger.debug("talker2code2wav: skip req=%s due to empty multimodal_output", req_id)
-            continue
-        mm: OmniPayload = mm_raw
-        if "codes" not in mm or not isinstance(mm.get("codes"), dict) or "audio" not in mm["codes"]:
-            logger.debug("talker2code2wav: skip req=%s due to missing codes.audio", req_id)
-            continue
-        # Extract codec codes from talker output
-        # Expected shape: [8, seq_len] (8-layer RVQ codes)
-        codec_codes = (
-            mm["codes"]["audio"][-seq_len:].to(torch.long).transpose(0, 1).cpu().to(torch.long).reshape(-1).tolist()
-        )  # 16, seq_len
-        code2wav_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=codec_codes,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-            )
-        )
-
-    return code2wav_inputs

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -38,103 +38,6 @@ def _qwen3_tts_empty_finished_payload():
     }
 
 
-def talker2code2wav(
-    source_outputs: list[Any],
-    prompt: Any = None,
-    _requires_multimodal_data: bool = False,
-) -> list[Any]:
-    """Non-async: collect all talker codes, then pass to code2wav at once."""
-    from vllm_omni.inputs.data import OmniTokensPrompt
-
-    talker_outputs = source_outputs
-    code2wav_inputs: list[OmniTokensPrompt] = []
-    for i, talker_output in enumerate(talker_outputs):
-        if not talker_output.finished:
-            # Non-async decode should only run once, after talker has
-            # accumulated the final code sequence.
-            continue
-        output = talker_output.outputs[0]
-        mm = output.multimodal_output
-        mm_codes = mm.get("codes", {})
-
-        # audio_codes shape: [num_frames, Q] where Q=num_quantizers (16)
-        audio_codes = mm_codes["audio"].to(torch.long)
-        token_ids = output.cumulative_token_ids
-
-        # token_ids provides an upper bound on the newly generated codec span.
-        # audio_codes may still contain zero-padded / invalid rows, so trim only
-        # after filtering valid frames instead of trying to align EOS indices.
-        seq_len = max(len(token_ids) - 1, 0)
-        # Filter invalid frames: zero-padded (EOS), out-of-range values (e.g.
-        # stop_token_id=2150 exceeds codebook_size=2048), and negative
-        # sentinels (e.g. -1 padding).
-        _CODEBOOK_SIZE = 2048
-        valid_mask = (
-            (audio_codes >= 0).all(dim=1) & audio_codes.any(dim=1) & (audio_codes.max(dim=1).values < _CODEBOOK_SIZE)
-        )
-        audio_codes = audio_codes[valid_mask]
-        if seq_len > 0 and audio_codes.ndim == 2 and int(audio_codes.shape[0]) > seq_len:
-            audio_codes = audio_codes[-seq_len:]
-        ref_code = mm_codes.get("ref")
-        ref_code_len = mm.get("meta", {}).get("ref_code_len")
-        if isinstance(ref_code_len, torch.Tensor):
-            ref_code_len = int(ref_code_len.reshape(-1)[-1].item()) if ref_code_len.numel() > 0 else 0
-        elif ref_code_len is None:
-            ref_code_len = 0
-        else:
-            ref_code_len = int(ref_code_len)
-        if isinstance(ref_code, list):
-            ref_code = ref_code[0] if ref_code else None
-        if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
-            ref_code = ref_code.to(torch.long).cpu().contiguous()
-            if ref_code.ndim == 1:
-                num_quantizers = int(audio_codes.shape[1]) if audio_codes.ndim == 2 and audio_codes.shape[1] > 0 else 16
-                if ref_code.numel() % num_quantizers != 0:
-                    logger.warning(
-                        "Ignoring malformed ref_code with %d elements not divisible by num_quantizers=%d",
-                        ref_code.numel(),
-                        num_quantizers,
-                    )
-                    ref_code = None
-                else:
-                    ref_code = ref_code.reshape(-1, num_quantizers)
-            elif ref_code.ndim != 2:
-                logger.warning("Ignoring malformed ref_code shape %s", tuple(ref_code.shape))
-                ref_code = None
-            if isinstance(ref_code, torch.Tensor) and ref_code_len > 0 and int(ref_code.shape[0]) > ref_code_len:
-                logger.warning(
-                    "Trimming ref_code from %d frames to ref_code_len=%d before Code2Wav.",
-                    int(ref_code.shape[0]),
-                    ref_code_len,
-                )
-                ref_code = ref_code[:ref_code_len]
-            if not isinstance(ref_code, torch.Tensor):
-                ref_code_len = 0
-            else:
-                ref_code_len = int(ref_code.shape[0])
-                audio_codes = torch.cat([ref_code.to(audio_codes.device), audio_codes], dim=0)
-        else:
-            ref_code_len = 0
-        # Code2Wav expects codebook-major flat: [Q*num_frames]
-        codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist()
-        additional_information = to_dict(
-            OmniPayloadStruct(
-                meta=MetaStruct(left_context_size=ref_code_len) if ref_code_len > 0 else None,
-                speaker=extract_speaker_from_prompt(prompt, index=i),
-                language=extract_language_from_prompt(prompt, index=i),
-            )
-        )
-        code2wav_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=codec_codes,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-                additional_information=additional_information if additional_information else None,
-            )
-        )
-    return code2wav_inputs
-
-
 def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None:
     audio_codes = multimodal_output.get("codes", {}).get("audio")
     if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
@@ -351,10 +254,7 @@ _NUM_QUANTIZERS_DEFAULT = 16
 
 
 def _filter_audio_codes_qwen3_tts(audio_codes: torch.Tensor) -> torch.Tensor:
-    """Filter zero-padded, out-of-range, and negative-padded codec frames.
-
-    Mirrors the orchestrator-path body in `talker2code2wav` above.
-    """
+    """Filter zero-padded, out-of-range, and negative-padded codec frames."""
     if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
         return audio_codes
     if audio_codes.ndim != 2:
@@ -367,10 +267,8 @@ def _filter_audio_codes_qwen3_tts(audio_codes: torch.Tensor) -> torch.Tensor:
 
 def _coerce_ref_code_len(raw) -> int:
     """Coerce mm["meta"]["ref_code_len"] / pooling_output["meta.ref_code_len"]
-    raw value (Tensor | int | None) into a non-negative int.  Mirrors the
-    extraction inlined in the legacy ``talker2code2wav`` path; clamps any
-    negative input to 0 since downstream code treats this as a non-negative
-    frame count."""
+    raw value (Tensor | int | None) into a non-negative int; clamps any
+    negative input to 0 since downstream code treats this as a frame count."""
     if isinstance(raw, torch.Tensor):
         value = int(raw.reshape(-1)[-1].item()) if raw.numel() > 0 else 0
     elif raw is None:
@@ -381,7 +279,7 @@ def _coerce_ref_code_len(raw) -> int:
 
 
 def _normalize_ref_code(ref_code, num_quantizers: int, ref_code_len: int):
-    """Coerce ref_code into a [ref_len, Q] tensor or None.  Mirrors orchestrator path."""
+    """Coerce ref_code into a [ref_len, Q] tensor or None."""
     if isinstance(ref_code, list):
         ref_code = ref_code[0] if ref_code else None
     if not isinstance(ref_code, torch.Tensor) or ref_code.numel() == 0:
@@ -407,8 +305,7 @@ def talker2code2wav_token_only(
 
     Sized to the expected codec token count (codebook-major flat:
     Q * (ref_frames + audio_frames)).  Speaker / language metadata are
-    extracted from `prompt` and threaded via `additional_information`
-    (orchestrator-style; same as the legacy `talker2code2wav` builder).
+    extracted from `prompt` and threaded via `additional_information`.
     Actual codec ids are delivered via the worker connector payload built
     by `talker2code2wav_full_payload`.
     """
@@ -473,8 +370,8 @@ def talker2code2wav_full_payload(
     Reads accumulated codec from `pooling_output["codes.audio"]` (CONCAT
     across steps via flatten_payload), latest `pooling_output["codes.ref"]`
     (prefill-emitted), and latest `pooling_output["meta.ref_code_len"]`.
-    Replicates the orchestrator-path body of `talker2code2wav` (filter,
-    crop to seq_len, prepend ref, codebook-major flatten).
+    Filters invalid frames, crops to seq_len, prepends ref, and flattens
+    codebook-major for code2wav consumption.
     """
     del transfer_manager
     rid = getattr(request, "request_id", "?")

vllm_omni/model_executor/stage_input_processors/voxtral_tts.py

@@ -2,7 +2,6 @@ from collections.abc import Mapping
 from typing import Any
 
 import torch
-from vllm.inputs import TextPrompt
 from vllm.logger import init_logger
 
 from vllm_omni.data_entry_keys import (
@@ -11,7 +10,6 @@ from vllm_omni.data_entry_keys import (
     OmniPayload,
     OmniPayloadStruct,
 )
-from vllm_omni.inputs.data import OmniTokensPrompt
 
 logger = init_logger(__name__)
 
@@ -31,29 +29,6 @@ def _codec_audio_tensors(multimodal_output: OmniPayload | dict[str, Any]) -> lis
     return None
 
 
-def generator2tokenizer(
-    source_outputs: list[Any],
-    _prompt: OmniTokensPrompt | TextPrompt = None,
-    _requires_multimodal_data: bool = False,
-):
-    tokenizer_inputs = []
-    for generator_output in source_outputs:
-        output = generator_output.outputs[0]
-        audio_tensors = _codec_audio_tensors(output.multimodal_output)
-        if not audio_tensors:
-            raise ValueError("Missing codes.audio in generator multimodal_output")
-        audio_tokens = torch.cat(audio_tensors, dim=-1).flatten().detach().cpu().tolist()
-        tokenizer_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=audio_tokens,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-            )
-        )
-
-    return tokenizer_inputs
-
-
 def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None:
     audio_tensors = _codec_audio_tensors(multimodal_output)
     if not audio_tensors:

vllm_omni/worker/gpu_ar_model_runner.py

@@ -8,7 +8,7 @@ from __future__ import annotations
 
 import gc
 import threading
-from collections.abc import Callable, Mapping, Sequence
+from collections.abc import Callable, Sequence
 from contextlib import nullcontext
 from copy import copy
 from dataclasses import replace
@@ -703,12 +703,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         if self.omni_prefix_cache is not None and is_last_pp_rank:
             # If this happens, it generally means the model is not following the correct
             # interface yet and is therefore currently not compatible with prefix cache.
-            if multimodal_outputs is not None and not isinstance(multimodal_outputs, Mapping):
-                logger.warning_once(
-                    "prefix caching expects mm outputs to be a dict, but got %s",
-                    type(multimodal_outputs),
-                )
-
             hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None
             # FIX: The .cpu attribute of slot_mapping is stale (not updated by the Triton
             # _compute_slot_mapping_kernel which only writes to .gpu). We must use .gpu and

vllm_omni/worker/gpu_model_runner.py

@@ -601,9 +601,6 @@ class OmniGPUModelRunner(GPUModelRunner):
             # Decode additional_information payloads (dictionary)
             try:
                 if getattr(new_req_data, "additional_information", None) is not None:
-                    logger.warning_once(
-                        "additional_information on request data is deprecated, use model_intermediate_buffer"
-                    )
                     info_dict = deserialize_additional_information(new_req_data.additional_information)
                     if info_dict:
                         self.model_intermediate_buffer[req_id] = info_dict
@@ -1471,15 +1468,9 @@ class OmniGPUModelRunner(GPUModelRunner):
         for new_req in scheduler_output.scheduled_new_reqs:
             payload_info = getattr(new_req, "additional_information", None)
             if isinstance(payload_info, dict):
-                logger.warning_once(
-                    "additional_information on request data is deprecated, use model_intermediate_buffer"
-                )
                 self._update_intermediate_buffer(new_req.req_id, payload_info)
 
         if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"):
-            logger.warning_once(
-                "additional_information on scheduled_cached_reqs is deprecated, use model_intermediate_buffer"
-            )
             cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {})
             if isinstance(cached_infos, dict):
                 for req_id, req_infos in cached_infos.items():
@@ -1707,7 +1698,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 self.text_step.gpu[dst].copy_(text_step)
 
                 for req_id_b, update_dict_b in zip(req_ids_b, updates, strict=True):
-                    self._merge_additional_information_update(req_id_b, update_dict_b)
+                    self._update_intermediate_buffer(req_id_b, update_dict_b)
 
                 decode_req_ids.extend(req_ids_b)
                 decode_start_offsets.extend(start_offsets_b)
@@ -1762,7 +1753,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                     decode_start_offsets.append(s)
 
                 # TODO(Peiqi): the merge stage could move out from the critical path
-                self._merge_additional_information_update(req_id, update_dict)
+                self._update_intermediate_buffer(req_id, update_dict)
 
                 # update the inputs_embeds and input_ids
                 seg_len = min(span_len, req_embeds.shape[0])
@@ -1912,7 +1903,7 @@ class OmniGPUModelRunner(GPUModelRunner):
             inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
             if code_predictor_codes is not None:
                 update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
-                self._merge_additional_information_update(req_id, update_dict)
+                self._update_intermediate_buffer(req_id, update_dict)
 
     def _model_forward(
         self,
@@ -1986,10 +1977,6 @@ class OmniGPUModelRunner(GPUModelRunner):
         # Backward compatible: mirror to old setattr location
         setattr(req_state, "additional_information_cpu", existing)
 
-    def _merge_additional_information_update(self, req_id, upd):
-        logger.warning_once("_merge_additional_information_update is deprecated, use _update_intermediate_buffer")
-        return self._update_intermediate_buffer(req_id, upd)
-
     def _update_streaming_input_additional_info(self, req_id):
         # For streaming input prefill case only. Set num processed tokens = 0 for new segment input
         cached_additional_info = self.model_intermediate_buffer.get(req_id, {})

vllm_omni/worker/omni_connector_model_runner_mixin.py

@@ -919,7 +919,7 @@ class OmniConnectorModelRunnerMixin:
         if not (finished_req_ids & pending_req_ids):
             return
 
-        logger.info(
+        logger.debug(
             "[Stage-%s] flush_full_payload_outputs: finished_req_ids=%s, pending=%s",
             self._stage_id,
             finished_req_ids,