[Bugfix][MiniCPM-o] Allow ragged audio_feature_lens across a batch (#7071)

Signed-off-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com>
Co-authored-by: Pujitha Paladugu <10557236+pujitha24@users.noreply.github.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py

@@ -10,6 +10,10 @@ Covers:
   - ``process_audios``: >30s audio unpads per chunk without ``TypeError``
   - multiple >30s audios unpad in chunk order
   - <=30s audio unpadding byte-identical to the pre-fix behavior
+  - ``MiniCPMOAudioFeatureInputs``: a batch mixing a >30s and a <=30s audio
+    (different slice counts) validates instead of raising, whether from a
+    single request with multiple audios of different lengths or from
+    multiple requests sharing an encoder batch
 """
 
 from __future__ import annotations
@@ -26,6 +30,7 @@ from vllm.multimodal.inputs import (
 
 from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import (
     MiniCPMO45OmniLLMMultiModalProcessor,
+    MiniCPMOAudioFeatureInputs,
     MiniCPMOMultiModalDataParser,
     _minicpmo_field_config,
 )
@@ -152,3 +157,43 @@ class TestProcessAudiosUnpadding:
         assert len(features) == 1
         assert features[0].shape == (_FEAT_DIM, 980)
         assert torch.equal(features[0], fake_outputs["audio_features"][0][:, :980])
+
+
+class TestMixedChunkCountBatchValidation:
+    def test_batch_mixes_different_slice_counts(self) -> None:
+        # A >30s audio (3 slices) and a <=30s audio (2 slices) scheduled in
+        # the same encoder batch: audio_feature_lens is a list of per-audio
+        # tensors with different lengths (3 vs 2). This is the shape produced
+        # both by a single request with two audios of different lengths and
+        # by two separate requests batched together. Before the fix this
+        # raised ValueError: "audio_feature_lens contains inconsistent shapes".
+        audio_feature_lens = [
+            torch.tensor([_CHUNK_FRAMES, _CHUNK_FRAMES, 1]),
+            torch.tensor([_CHUNK_FRAMES, 1]),
+        ]
+        audio_features = [torch.zeros(_FEAT_DIM, _CHUNK_FRAMES) for _ in range(5)]
+
+        inputs = MiniCPMOAudioFeatureInputs(
+            type="audio_features",
+            audio_features=audio_features,
+            audio_feature_lens=audio_feature_lens,
+        )
+
+        assert inputs.audio_feature_lens is audio_feature_lens
+
+    def test_single_slice_count_batch_still_validated(self) -> None:
+        # Same slice count per audio (the common case) must still validate,
+        # unaffected by marking "s" dynamic.
+        audio_feature_lens = [
+            torch.tensor([_CHUNK_FRAMES, 1]),
+            torch.tensor([_CHUNK_FRAMES, 1]),
+        ]
+        audio_features = [torch.zeros(_FEAT_DIM, _CHUNK_FRAMES) for _ in range(4)]
+
+        inputs = MiniCPMOAudioFeatureInputs(
+            type="audio_features",
+            audio_features=audio_features,
+            audio_feature_lens=audio_feature_lens,
+        )
+
+        assert inputs.audio_feature_lens is audio_feature_lens

vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py

@@ -3796,11 +3796,17 @@ class MiniCPMOAudioFeatureInputs(TensorSchema):
 
     audio_feature_lens: Annotated[
         torch.Tensor | list[torch.Tensor],
-        TensorShape("bn", "s"),
+        TensorShape("bn", "s", dynamic_dims={"s"}),
     ]
     """
     This should be feature length of each audio slice,
     which equals to `audio_features.shape[-1]`
+
+    Each audio in the batch may be split into a different number of slices
+    (e.g. audio >30s splits into multiple slices while shorter audio doesn't),
+    so `s` must be dynamic: batching audios with different slice counts is a
+    normal, valid input and is already handled below via `hstack` + per-audio
+    iteration, not a schema violation.
     """