diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py index 6853d0ff4..cfd0751c4 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py @@ -928,4 +928,51 @@ def test_talker_output_contract(monkeypatch, mrv2, packed): model._sampling_eps = 1e-6 model.mrv2_custom_sampler(SimpleNamespace(penalties_state=None)) output = model.forward(torch.ones(2, dtype=torch.long), torch.arange(2), inputs_embeds=hidden) - assert output is (hidden if mrv2 or packed else sentinel) + assert output is hidden + + +def test_request_conditioning_normalizes_singletons_without_changing_schema(): + from vllm_omni.data_entry_keys import to_struct + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning + + token = torch.tensor([[7, 8]], dtype=torch.int32) + feat = torch.randn(1, 4, 80) + speaker = torch.randn(1, 192) + length = torch.tensor([[2]], dtype=torch.int32) + raw = {"embed": {"speech_token": token, "speech_feat": feat, "embedding": speaker, "speech_token_len": [length]}} + result = to_struct(_normalize_request_conditioning(raw)) + assert result.embed.speech_token is token + assert result.embed.speech_feat is feat + assert result.embed.embedding is speaker + assert result.embed.speech_token_len is length + assert raw["embed"]["speech_token_len"] == [length] + + +def test_request_conditioning_rejects_unsplit_voices(): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning + + with pytest.raises(ValueError, match="unsplit batch"): + _normalize_request_conditioning({"embed": {"embedding": [torch.ones(1, 192), torch.zeros(1, 192)]}}) + + +def test_request_conditioning_reaches_code2wav_with_true_prompt_length(): + model = _make_code2wav_model() + output = model.forward( + torch.tensor([1, 2]), + torch.arange(2), + model_intermediate_buffer=[ + { + "embed": { + "speech_token": torch.tensor([[3, 2, 0, 0]], dtype=torch.int32), + "speech_feat": torch.randn(1, 8, 80), + "embedding": torch.randn(1, 192), + "speech_token_len": [torch.tensor([[2]], dtype=torch.int32)], + }, + "meta": {"req_id": ["voice"]}, + } + ], + ) + call = model.code2wav.forward_calls[0] + assert call["prompt_token"].shape == (1, 2) + assert call["prompt_feat"].shape == (1, 4, 80) + assert output.multimodal_outputs 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 08090d954..30914ea74 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 @@ -13,6 +13,7 @@ from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import ( text2flow_full_payload, text2flow_token_only, ) +from vllm_omni.outputs.mm_outputs import MultimodalPayload pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -121,6 +122,24 @@ def test_text2flow_full_payload_does_not_send_codec_ids(): assert torch.equal(payload["embed"]["speech_token"], torch.tensor([[1, 2]], dtype=torch.long)) +def test_text2flow_full_payload_accepts_structured_multimodal_payload(): + speech_token = torch.tensor([[1, 2]], dtype=torch.long) + payload = text2flow_full_payload( + None, + MultimodalPayload( + tensors={ + "embed.speech_token": speech_token, + "codes.audio": torch.tensor([7, 8, 9], dtype=torch.long), + } + ), + SimpleNamespace(external_req_id="req-structured"), + ) + + assert payload is not None + assert "codes" not in payload + assert torch.equal(payload["embed"]["speech_token"], speech_token) + + def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset(): transfer_manager = _transfer_manager() request = SimpleNamespace( @@ -355,3 +374,36 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio( assert payload_final is not None assert payload_final.meta.finished.item() is True assert payload_final.codes.audio.tolist() == [] + + +def test_full_payload_retains_true_length_for_padded_conditioning(): + from vllm_omni.data_entry_keys import to_struct + + token = torch.tensor([[7, 8, 0, 0]], dtype=torch.int32) + feat = torch.randn(1, 8, 80) + speaker = torch.randn(1, 192) + length = torch.tensor([[2]], dtype=torch.int32) + result = text2flow_full_payload( + None, + { + "embed.speech_token": token, + "embed.speech_feat": feat, + "embed.embedding": speaker, + "embed.speech_token_len": length, + }, + SimpleNamespace(request_id="padded"), + ) + payload = to_struct(result) + assert payload.embed.speech_token_len is length + + +def test_reference_features_are_snapshots_under_output_accumulation(): + from vllm_omni.outputs.mm_outputs import MultimodalPayload + from vllm_omni.outputs.output_modality import OutputModality + + first = torch.ones(1, 174, 80) + latest = torch.zeros(1, 1500, 80) + payload = MultimodalPayload.from_dict({"embed.speech_feat": first}) + payload = payload.merged_with(MultimodalPayload.from_dict({"embed.speech_feat": latest})) + payload.consolidate_tensors(OutputModality.LATENT) + assert payload["embed.speech_feat"] is latest 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 f67fc7a58..9e84785eb 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 @@ -971,12 +971,14 @@ def test_cosyvoice3_text2flow_full_payload_nested_fallback() -> None: def test_cosyvoice3_full_payload_replace_keys_present() -> None: - """Confirm _FULL_PAYLOAD_REPLACE_KEYS lists the three embed.* keys.""" + """Reference tensors and their length travel as complete snapshots.""" from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import ( _FULL_PAYLOAD_REPLACE_KEYS, ) - assert _FULL_PAYLOAD_REPLACE_KEYS == frozenset({"embed.speech_token", "embed.speech_feat", "embed.embedding"}) + assert _FULL_PAYLOAD_REPLACE_KEYS == frozenset( + {"embed.speech_token", "embed.speech_token_len", "embed.speech_feat", "embed.embedding"} + ) def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None: diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 654675474..f47ff7556 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -74,6 +74,27 @@ logger = init_logger(__name__) _RUNTIME_COMPONENTS_CACHE: dict[str, dict] = {} +def _normalize_request_conditioning(payload: dict) -> dict: + """Unwrap singleton conditioning containers at the per-request boundary. + + A processor/prefix-cache passthrough can retain a one-item tensor list. + The code2wav schema consumes one request, so that wrapper carries no + batch dimension. A multi-item list here is an unsplit batch and must + fail instead of silently choosing another request's voice. + """ + embed = payload.get("embed") + if not isinstance(embed, Mapping): + return payload + normalized = dict(embed) + for key in ("speech_token", "speech_feat", "embedding", "speech_token_len"): + value = normalized.get(key) + if isinstance(value, (list, tuple)): + if len(value) != 1: + raise ValueError(f"CosyVoice3 per-request {key} contains an unsplit batch of {len(value)} items") + normalized[key] = value[0] + return {**payload, "embed": normalized} + + def _cosyvoice3_trt_enabled() -> bool: """COSYVOICE3_TRT env toggle (default on) for the optional TensorRT paths. @@ -1501,9 +1522,10 @@ class CosyVoice3Model( # [total_tokens, hidden] hidden_states = self.model.llm(inputs_embeds, positions) - if getattr(self, "_mrv2_tensor_output", False) or cosyvoice3_packed_inference_enabled(): - return hidden_states - return self.make_omni_output(hidden_states, **kwargs) + # Both runners attach the live step's conditioning after replay. + # Capturing an OmniOutput here would retain the dummy prefill lists + # and re-emit them during decode, replacing the real reference. + return hidden_states elif self.model_stage == "cosyvoice3_code2wav": # Lazily swap the flow-decoder estimator to a TensorRT engine on the # first code2wav step (after weights are loaded), gated by the same @@ -1539,7 +1561,7 @@ class CosyVoice3Model( for idx, req_ids in enumerate(request_ids_list): raw = runtime_info[idx] if idx < len(runtime_info) and isinstance(runtime_info[idx], dict) else {} - payload = to_struct(raw) + payload = to_struct(_normalize_request_conditioning(raw)) meta = payload.meta embed = payload.embed diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index 07ea3af7f..41c825dd5 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py @@ -292,7 +292,9 @@ def talker2code2wav_async_chunk( # CONCAT across the (already trivial) per-request accumulator history so a # regression where decode unexpectedly re-emits them does not silently # duplicate the prefill tensor. See mixin._FULL_PAYLOAD_REPLACE_KEYS. -_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset({"embed.speech_token", "embed.speech_feat", "embed.embedding"}) +_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset( + {"embed.speech_token", "embed.speech_feat", "embed.embedding", "embed.speech_token_len"} +) def text2flow_token_only( @@ -347,20 +349,20 @@ def text2flow_full_payload( """ del transfer_manager rid = getattr(request, "external_req_id", None) or getattr(request, "request_id", "?") - if not isinstance(pooling_output, dict): + if not isinstance(pooling_output, Mapping): logger.warning( - "cosyvoice3.text2flow_full_payload: pooling_output not a dict " + "cosyvoice3.text2flow_full_payload: pooling_output not a mapping " "(type=%s) for req=%s; consumer wait gate may hang.", type(pooling_output).__name__, rid, ) return None embed_out: dict[str, Any] = {} - for key in ("speech_token", "speech_feat", "embedding"): + for key in ("speech_token", "speech_feat", "embedding", "speech_token_len"): v = pooling_output.get(f"embed.{key}") if v is None: nested = pooling_output.get("embed") - if isinstance(nested, dict): + if isinstance(nested, Mapping): v = nested.get(key) if isinstance(v, torch.Tensor) and v.numel() > 0: embed_out[key] = v diff --git a/vllm_omni/outputs/output_modality.py b/vllm_omni/outputs/output_modality.py index 412e61619..1e8489e99 100644 --- a/vllm_omni/outputs/output_modality.py +++ b/vllm_omni/outputs/output_modality.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Output modality types for vLLM-Omni. This module defines the OutputModality enum and TensorAccumulationStrategy @@ -87,7 +90,7 @@ class OutputModality(Flag): try: result |= cls[p.upper()] except KeyError: - raise ValueError(f"Unknown modality: {p!r}. Supported: {[m.name.lower() for m in cls]}") + raise ValueError(f"Unknown modality: {p!r}. Supported: {[(m.name or '').lower() for m in cls]}") return result @property @@ -123,7 +126,14 @@ class TensorAccumulationStrategy(Enum): # re-emitted unchanged at every step. Pipelines that produce such keys # register the correct strategy here via ``register_key_accumulation_strategy`` # instead of forcing every tensor of the modality through one default. -_KEY_ACCUMULATION_OVERRIDES: dict[str, TensorAccumulationStrategy] = {} +# Schema-defined prompt conditioning is a complete reference snapshot. Keep +# these defaults in the output process too, which may never import the model's +# worker-side stage input processor. Feature lengths can change between +# snapshots and cannot be concatenated as generated content deltas. +_KEY_ACCUMULATION_OVERRIDES: dict[str, TensorAccumulationStrategy] = { + key: TensorAccumulationStrategy.REPLACE + for key in ("embed.speech_token", "embed.speech_feat", "embed.embedding", "embed.speech_token_len") +} def register_key_accumulation_strategy(key: str, strategy: TensorAccumulationStrategy) -> None: