A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 10:00 UTC
[Bug Fix] Handle variable-length tensors in MultiModalFlatField batching (#31751)
Signed-off-by: Andrii Pasternak <andriipasternak31@gmail.com> Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
andrii.pasternak committed
Jan 29, 2026 at 10:42 UTC
615e8033e567171b4af15024038b302f386d2ee4
2 files changed
+86
tests/models/multimodal/generation/test_ultravox.py
+45
@@ -156,6 +156,51 @@ def test_models_with_multiple_audios(
156
)
157
158
159
+@pytest.mark.core_model
160
+@pytest.mark.parametrize("dtype", ["half"])
161
+@pytest.mark.parametrize("max_tokens", [32])
162
+def test_variable_length_audio_batching(
163
+ vllm_runner,
164
+ audio_assets: AudioTestAssets,
165
+ dtype: str,
166
+ max_tokens: int,
167
+) -> None:
168
+ """Test batching of requests with different audio durations.
169
+
170
+ This exercises the variable-length tensor handling in
171
+ MultiModalFlatField._reduce_data() which was buggy before
172
+ https://github.com/vllm-project/vllm/issues/31658 was fixed.
173
+ """
174
+ model_info = HF_EXAMPLE_MODELS.find_hf_info(MODEL_NAME)
175
+ model_info.check_available_online(on_fail="skip")
176
+ model_info.check_transformers_version(on_fail="skip")
177
+
178
+ # Create prompts with single audio each (different durations)
179
+ prompts_and_audios = []
180
+ for audio, question in zip(audio_assets, AUDIO_PROMPTS):
181
+ prompt = _get_prompt(1, question, VLLM_PLACEHOLDER)
182
+ prompts_and_audios.append((prompt, [audio.audio_and_sample_rate]))
183
+
184
+ with vllm_runner(
185
+ MODEL_NAME,
186
+ dtype=dtype,
187
+ enforce_eager=True,
188
+ limit_mm_per_prompt={"audio": 1},
189
+ ) as vllm_model:
190
+ # Generate for all prompts in a single batch
191
+ # This triggers the variable-length batching code path
192
+ outputs = vllm_model.generate_greedy(
193
+ [prompt for prompt, _ in prompts_and_audios],
194
+ max_tokens,
195
+ audios=[audios for _, audios in prompts_and_audios],
196
+ )
197
+
198
+ # Verify outputs were generated for each request
199
+ assert len(outputs) == len(prompts_and_audios)
200
+ for output in outputs:
201
+ assert len(output[1]) > 0, "Expected non-empty output"
202
+
203
+
204
@pytest.mark.asyncio
205
async def test_online_serving(client, audio_assets: AudioTestAssets):
206
"""Exercises online serving with/without chunked prefill enabled."""
vllm/multimodal/inputs.py
+41
@@ -604,6 +604,47 @@ class MultiModalFlatField(BaseMultiModalField):
604
)
605
return torch.concat(batch, dim=self.dim, out=out)
606
607
+ # Variable-length case: non-concat dimensions differ
608
+ # (e.g., Ultravox with different audio durations).
609
+ # Use slice-assign approach (more efficient than padding).
610
+ # See: https://github.com/vllm-project/vllm/issues/31658
611
+
612
+ ndim = batch[0].ndim
613
+
614
+ # Step 1: Compute output shape
615
+ # - Non-concat dims: take max across batch
616
+ # - Concat dim: sum across batch
617
+ max_sizes: list[int] = []
618
+ for d in range(ndim):
619
+ if d == dim:
620
+ max_sizes.append(sum(t.shape[d] for t in batch))
621
+ else:
622
+ max_sizes.append(max(t.shape[d] for t in batch))
623
+
624
+ # Step 2: Create zero-initialized output tensor
625
+ out = torch.zeros(
626
+ max_sizes,
627
+ dtype=batch[0].dtype,
628
+ device=batch[0].device,
629
+ pin_memory=pin_memory,
630
+ )
631
+
632
+ # Step 3: Slice-assign each tensor to its proper position
633
+ concat_offset = 0
634
+ for tensor in batch:
635
+ slices: list[slice] = []
636
+ for d in range(ndim):
637
+ if d == dim:
638
+ slices.append(
639
+ slice(concat_offset, concat_offset + tensor.shape[d])
640
+ )
641
+ else:
642
+ slices.append(slice(0, tensor.shape[d]))
643
+ out[tuple(slices)] = tensor
644
+ concat_offset += tensor.shape[dim]
645
+
646
+ return out
647
+
648
assert self.dim == 0, "dim == 0 is required for nested list"
649
return [e for elem in batch for e in elem]
650