[Bugfix] Give exact-shape codec models an unpadded input_ids view (#6775)

Signed-off-by: Inesh Reddy <ineshreddy249@gmail.com>
Signed-off-by: IneshReddy249 <ineshreddy249@gmail.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py

@@ -143,6 +143,7 @@ class MiniCPMO45Code2Wav(nn.Module):
     replace_runtime_additional_information = True
     requires_raw_input_tokens = True
     requires_request_ids = True
+    requires_exact_input_shape = True
     has_preprocess = False
     has_postprocess = False
 

vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License").
@@ -226,6 +229,7 @@ class MossTTSCodecDecoder(nn.Module):
     has_postprocess: bool = False
     enable_update_additional_information: bool = True
     requires_raw_input_tokens: bool = True
+    requires_exact_input_shape: bool = True
 
     _OUTPUT_SAMPLE_RATE: int = 24_000
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright 2025 The Qwen team.
 """Inference-only Qwen3-Omni-Moe unified model (thinker + talker + code2wav)."""
 
@@ -263,6 +263,8 @@ class Qwen3OmniMoeForConditionalGeneration(
             )
             self.model = self.code2wav
             self.requires_raw_input_tokens = True
+            # torch.split below requires sum(counts) == input_ids.numel() (#6712)
+            self.requires_exact_input_shape = True
         else:
             raise ValueError(
                 f"Invalid model_stage: {self.model_stage}. Must be one of: 'thinker', 'talker', 'code2wav'"