[Bugfix] Replace CosyVoice prompt-conditioning snapshots during accumulation

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -971,12 +971,40 @@ 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."""
+    """Confirm all prompt-conditioning fields replace accumulated values."""
     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_feat",
+            "embed.speech_token_len",
+            "embed.embedding",
+        }
+    )
+
+
+def test_cosyvoice3_conditioning_replaces_snapshot_while_audio_accumulates() -> None:
+    from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import _FULL_PAYLOAD_REPLACE_KEYS
+    from vllm_omni.outputs.mm_outputs import MultimodalPayload
+    from vllm_omni.outputs.output_modality import OutputModality
+
+    previous = {key: torch.ones(1, 174) for key in _FULL_PAYLOAD_REPLACE_KEYS}
+    padded = {key: torch.zeros(1, 1500) for key in _FULL_PAYLOAD_REPLACE_KEYS}
+    previous["embed.speech_token_len"] = torch.tensor([174])
+    padded["embed.speech_token_len"] = torch.tensor([174])
+    payload = MultimodalPayload.from_dict({**previous, "audio": torch.ones(2)})
+    incoming = MultimodalPayload.from_dict({**padded, "audio": torch.ones(3)})
+    assert payload is not None and incoming is not None
+
+    payload = payload.merged_with(incoming)
+    payload.consolidate_tensors(OutputModality.AUDIO)
+
+    for key in _FULL_PAYLOAD_REPLACE_KEYS:
+        assert torch.equal(payload[key], padded[key])
+    assert torch.equal(payload["audio"], torch.ones(5))
 
 
 def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None:

vllm_omni/model_executor/stage_input_processors/cosyvoice3.py

@@ -18,6 +18,10 @@ from vllm_omni.data_entry_keys import (
 from vllm_omni.engine.serialization import deserialize_additional_information
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.cosyvoice3.utils import unpad_prompt_conditioning
+from vllm_omni.outputs.output_modality import (
+    TensorAccumulationStrategy,
+    register_key_accumulation_strategy,
+)
 
 logger = init_logger(__name__)
 
@@ -288,11 +292,28 @@ def talker2code2wav_async_chunk(
 # keeps the orchestrator off the heavy-tensor path.
 # ============================================================================
 
-# All three embed tensors are emitted once at prefill and must REPLACE-not-
+# Prompt-conditioning fields are emitted once at prefill and must REPLACE-not-
 # 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.speech_token_len",
+        "embed.embedding",
+    }
+)
+
+# The input conditioning snapshot remains in the request payload and the
+# talker also emits its padded snapshot. Both describe the same prompt;
+# concatenating them creates incompatible shapes before speech_token_len can
+# unpad the emitted value for code2wav.
+for _conditioning_key in _FULL_PAYLOAD_REPLACE_KEYS:
+    register_key_accumulation_strategy(
+        _conditioning_key,
+        TensorAccumulationStrategy.REPLACE,
+    )
 
 
 def text2flow_token_only(