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

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

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,