fix: preserve CosyVoice full payload and live conditioning (#8343)

Signed-off-by: andyluo7 <andy.luo@amd.com>

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

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

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:

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
 

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

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: