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..e191ddb1e 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,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: diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index 07ea3af7f..7f3663a97 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/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(