fix: preserve Omni processor conditioning on vLLM 0.31

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

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:

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,

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."
 

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."
 

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:

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."
 

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],

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(