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 dfbb9d91d..6d25a1b41 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,19 @@ 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_ming_flash_omni_thinker2talker_token_only_smoke() -> None: diff --git a/vllm_omni/inputs/mm_processor.py b/vllm_omni/inputs/mm_processor.py index 125af757b..297b4d13e 100644 --- a/vllm_omni/inputs/mm_processor.py +++ b/vllm_omni/inputs/mm_processor.py @@ -8,11 +8,39 @@ from typing import TypeVar from transformers import BatchFeature from vllm.multimodal.parse import MultiModalDataItems -from vllm.multimodal.processing import BaseMultiModalProcessor, BaseProcessingInfo +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + ProcessorInputs, + cached_encode, +) _I = TypeVar("_I", bound=BaseProcessingInfo) +class OmniDummyInputsBuilder(BaseDummyInputsBuilder[_I]): + """Restore Omni's processor-input hook on the vLLM 0.31 API.""" + + def get_dummy_processor_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, object], + ) -> ProcessorInputs: + dummy_text = self.get_dummy_text(mm_counts) + dummy_mm_data = self.get_dummy_mm_data(seq_len, mm_counts, mm_options) + dummy_mm_items = self.info.parse_mm_data(dummy_mm_data, validate=False) + + tokenizer = self.info.ctx.tokenizer + dummy_prompt = ( + [] + if tokenizer is None + else cached_encode(tokenizer, dummy_text, truncation=False) + ) + return ProcessorInputs(prompt=dummy_prompt, mm_data_items=dummy_mm_items) + + class OmniMultiModalProcessor(BaseMultiModalProcessor[_I]): """Share data/passthrough handling, leaving model processing in its owner. @@ -23,6 +51,18 @@ class OmniMultiModalProcessor(BaseMultiModalProcessor[_I]): _OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text" + def get_dummy_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, object], + ) -> ProcessorInputs: + if isinstance(self.dummy_inputs, OmniDummyInputsBuilder): + return self.dummy_inputs.get_dummy_processor_inputs( + seq_len, mm_counts, mm_options + ) + return super().get_dummy_inputs(seq_len, mm_counts, mm_options) + def _apply_hf_processor_main( self, mm_items: MultiModalDataItems, diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 654675474..5552c7360 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -42,7 +42,7 @@ from vllm.v1.sample.ops.topk_topp_sampler import random_sample from vllm.v1.sample.sampler import Sampler from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct -from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.inputs.mm_processor import OmniDummyInputsBuilder, OmniMultiModalProcessor from vllm_omni.model_executor.models.cosyvoice3.ras_sampler import MAX_FUSED_TOP_K, fused_ras_sample from vllm_omni.model_executor.models.cosyvoice3.runtime import ( cosyvoice3_batch_flow_debug, @@ -675,7 +675,7 @@ class CosyVoice3MultiModalProcessor(OmniMultiModalProcessor[CosyVoice3MultiModal ] -class CosyVoice3DummyInputsBuilder(BaseDummyInputsBuilder[CosyVoice3MultiModalProcessingInfo]): +class CosyVoice3DummyInputsBuilder(OmniDummyInputsBuilder[CosyVoice3MultiModalProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: return "Hello, this is a test of the CosyVoice3 system capability." diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts.py b/vllm_omni/model_executor/models/glm_tts/glm_tts.py index e98c227a4..2eef3aa44 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts.py @@ -50,7 +50,7 @@ from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler -from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.inputs.mm_processor import OmniDummyInputsBuilder, OmniMultiModalProcessor from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one as _ras_sample_one from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform @@ -674,7 +674,7 @@ class GLMTTSMultiModalProcessor(OmniMultiModalProcessor[GLMTTSMultiModalProcessi ] -class GLMTTSDummyInputsBuilder(BaseDummyInputsBuilder[GLMTTSMultiModalProcessingInfo]): +class GLMTTSDummyInputsBuilder(OmniDummyInputsBuilder[GLMTTSMultiModalProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: return "This is a test of the GLM-TTS voice cloning system." diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py index 9bd8246f3..f364ecf60 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py @@ -44,7 +44,7 @@ from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler -from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.inputs.mm_processor import OmniDummyInputsBuilder, OmniMultiModalProcessor from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import ( NO_INTERLEAVE_NEXT_TOKEN_ID, @@ -224,7 +224,7 @@ class MiMoAudioLLMProcessingInfo( return {"audio": 1} -class MiMoAudioLLMDummyInputsBuilder(BaseDummyInputsBuilder[MiMoAudioLLMProcessingInfo]): +class MiMoAudioLLMDummyInputsBuilder(OmniDummyInputsBuilder[MiMoAudioLLMProcessingInfo]): _processor_inputs_cache: LRUCache = LRUCache(capacity=1024) def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py index 3fc53e754..f73c6fac8 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py @@ -37,7 +37,7 @@ from vllm.multimodal.processing import ( ) from vllm.sequence import IntermediateTensors -from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.inputs.mm_processor import OmniDummyInputsBuilder, OmniMultiModalProcessor from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig @@ -241,7 +241,7 @@ class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalPr ] -class OmniVoiceDummyInputsBuilder(BaseDummyInputsBuilder[OmniVoiceMultiModalProcessingInfo]): +class OmniVoiceDummyInputsBuilder(OmniDummyInputsBuilder[OmniVoiceMultiModalProcessingInfo]): def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: return "Hello, this is a test of the OmniVoice system." diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py index 5eabb4dbc..f67edd30b 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py @@ -838,6 +838,14 @@ class VoxtralTTSDummyInputsBuilder(BaseDummyInputsBuilder[VoxtralTTSProcessingIn class VoxtralTTSMultiModalProcessor(BaseMultiModalProcessor[VoxtralTTSProcessingInfo]): + def get_dummy_inputs( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, Any], + ) -> ProcessorInputs: + return self.dummy_inputs.get_dummy_processor_inputs(seq_len, mm_counts, mm_options) + def _get_mm_fields_config( self, hf_inputs: Mapping[str, NestedTensors], diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index 07ea3af7f..3e8af2c0a 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", + } +) + +# vLLM 0.31 keeps the input conditioning snapshot 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(