From 9a44e7e019a90c4700cd06244bdf11791825843a Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Fri, 10 Jul 2026 15:13:35 +0800 Subject: [PATCH] [Refactor][Phase 1]Remove redundant functions and logs (#4986) Signed-off-by: amy-why-3459 --- docs/configuration/pd_disaggregation.md | 4 +- docs/contributing/model/adding_omni_model.md | 120 ++++----- .../test_cosyvoice3_stage_input_processors.py | 15 +- .../stage_input_processors/test_indextts2.py | 113 -------- .../test_mimo_audio_llm2code2wav.py | 43 +-- .../test_qwen3_omni_streaming_helpers.py | 8 +- .../test_qwen3_tts_async_chunk.py | 44 ++- .../test_voxtral_tts_async_chunk.py | 49 ---- tests/test_config_factory.py | 7 +- vllm_omni/deploy/qwen3_omni_moe.yaml | 2 - vllm_omni/engine/stage_engine_core_client.py | 2 +- .../models/cosyvoice3/pipeline.py | 6 +- .../models/mimo_audio/pipeline.py | 1 - .../models/qwen3_omni/pipeline.py | 2 - .../stage_input_processors/cosyvoice3.py | 24 -- .../stage_input_processors/fish_speech.py | 31 --- .../stage_input_processors/indextts2.py | 59 ---- .../stage_input_processors/mimo_audio.py | 89 +------ .../stage_input_processors/qwen3_omni.py | 251 +++--------------- .../stage_input_processors/qwen3_tts.py | 117 +------- .../stage_input_processors/voxtral_tts.py | 25 -- vllm_omni/worker/gpu_ar_model_runner.py | 8 +- vllm_omni/worker/gpu_model_runner.py | 19 +- .../omni_connector_model_runner_mixin.py | 2 +- 24 files changed, 162 insertions(+), 879 deletions(-) diff --git a/docs/configuration/pd_disaggregation.md b/docs/configuration/pd_disaggregation.md index 9196bdb02..7deb0443a 100644 --- a/docs/configuration/pd_disaggregation.md +++ b/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: diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index 1eaff1059..c20e94eda 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/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). diff --git a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py index 83168b4f5..5309c7675 100644 --- a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py +++ b/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( diff --git a/tests/model_executor/stage_input_processors/test_indextts2.py b/tests/model_executor/stage_input_processors/test_indextts2.py index 59f14b414..11468603e 100644 --- a/tests/model_executor/stage_input_processors/test_indextts2.py +++ b/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 ── diff --git a/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py b/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py index d1cc0dbf7..ae1c8e2f9 100644 --- a/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py +++ b/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(): diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py index 69b612a05..afc71c476 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py +++ b/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 diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index 82f7da711..bc6239e29 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/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] diff --git a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py index 3fd5675ce..dc248f211 100644 --- a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py +++ b/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, diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 6cf288b4c..c1c76180a 100644 --- a/tests/test_config_factory.py +++ b/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: diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml index f91eea5a8..5bcbb0fe3 100644 --- a/vllm_omni/deploy/qwen3_omni_moe.yaml +++ b/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: diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 78c209b0b..5fbbbc46d 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/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, diff --git a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py index 32db8ef5b..a8d51fdfe 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py +++ b/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", ), ), diff --git a/vllm_omni/model_executor/models/mimo_audio/pipeline.py b/vllm_omni/model_executor/models/mimo_audio/pipeline.py index 126c90176..9de6575d9 100644 --- a/vllm_omni/model_executor/models/mimo_audio/pipeline.py +++ b/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}, ), diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index 137e4e29d..ca3a09f89 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/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}, ), ), diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index 0477ea4e1..2956be94f 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/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, diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py index 668ff2ed4..f428c34a4 100644 --- a/vllm_omni/model_executor/stage_input_processors/fish_speech.py +++ b/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, diff --git a/vllm_omni/model_executor/stage_input_processors/indextts2.py b/vllm_omni/model_executor/stage_input_processors/indextts2.py index 9d38c8f2a..39d9d011c 100644 --- a/vllm_omni/model_executor/stage_input_processors/indextts2.py +++ b/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, diff --git a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py index 7c0db4131..3a1a4441f 100644 --- a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index d624b4a49..6152bb82b 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index 6cebc4c28..597fc4db5 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/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", "?") diff --git a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py index d3c1219b7..af2cb0a25 100644 --- a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py +++ b/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: diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 59bfbb855..d84404508 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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 diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index e922146db..b816a36f7 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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, {}) diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index c1421a231..b173f81ac 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/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,