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 3, 19:00 UTC
[Model][Bugfix] Fix Step3-VL image_embeds input path (#42333)
Signed-off-by: Kaivalya Dabhadkar <kdabhadkar@nvidia.com>
KaivalyaMDabhadkar committed
May 13, 2026 at 00:17 UTC
67c89fe40ac50b9d232087f07748bfcd25b54885
2 files changed
+72
-10
tests/models/multimodal/processing/test_step3_vl_image_embeds.py
new
+58
@@ -0,0 +1,58 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for Step3-VL precomputed image embedding inputs."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.model_executor.models.step3_vl import (
9
+ Step3VLForConditionalGeneration,
10
+ Step3VLImageEmbeddingInputs,
11
+)
12
+
13
+
14
+class _FakeStep3VL:
15
+ @staticmethod
16
+ def _process_image_features(image_features: torch.Tensor) -> torch.Tensor:
17
+ return image_features
18
+
19
+
20
+def test_image_embedding_inputs_construction():
21
+ """Step3VLImageEmbeddingInputs should store embeddings in the data field."""
22
+ image_embeds = torch.randn(2, 16, 64)
23
+
24
+ inputs = Step3VLImageEmbeddingInputs(
25
+ type="image_embeds",
26
+ data=image_embeds,
27
+ )
28
+
29
+ assert inputs["type"] == "image_embeds"
30
+ assert torch.equal(inputs["data"], image_embeds)
31
+ assert torch.equal(inputs.data, image_embeds)
32
+
33
+
34
+def test_image_embedding_inputs_validation_rejects_wrong_rank():
35
+ """Validation should reject tensors with wrong rank."""
36
+ with pytest.raises(ValueError, match="rank"):
37
+ Step3VLImageEmbeddingInputs(
38
+ type="image_embeds",
39
+ data=torch.randn(16, 64),
40
+ )
41
+
42
+
43
+def test_process_image_embeds_does_not_require_pixel_input_fields():
44
+ """The image_embeds branch should not reference patch pixel metadata."""
45
+ image_embeds = torch.randn(2, 4, 8)
46
+ image_input = Step3VLImageEmbeddingInputs(
47
+ type="image_embeds",
48
+ data=image_embeds,
49
+ )
50
+
51
+ outputs = Step3VLForConditionalGeneration._process_image_input(
52
+ _FakeStep3VL(),
53
+ image_input,
54
+ )
55
+
56
+ assert len(outputs) == 2
57
+ assert torch.equal(outputs[0], image_embeds[0])
58
+ assert torch.equal(outputs[1], image_embeds[1])
vllm/model_executor/models/step3_vl.py
+14
-10
@@ -592,7 +592,7 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
592
if image_embeds is not None:
593
return Step3VLImageEmbeddingInputs(
594
type="image_embeds",
595
- image_embeds=image_embeds.to(self.dtype),
595
+ data=image_embeds.to(self.dtype),
596
)
597
598
raise AssertionError("This line should be unreachable.")
@@ -615,15 +615,19 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
615
self, image_input: Step3VLImageInputs
616
) -> tuple[torch.Tensor, ...]:
617
if image_input["type"] == "image_embeds":
618
- image_features = image_input["image_embeds"]
619
- else:
620
- image_features = self._get_vision_model_output(image_input["pixel_values"])
621
- patch_image_features = (
622
- self._get_vision_model_output(image_input["patch_pixel_values"])
623
- if len(image_input["patch_pixel_values"]) > 0
624
- else None
625
- )
626
- num_patches = image_input["num_patches"]
618
+ image_features = image_input["data"]
619
+ return [
620
+ image_features[i].view(-1, image_features.shape[-1])
621
+ for i in range(image_features.shape[0])
622
+ ]
623
+
624
+ image_features = self._get_vision_model_output(image_input["pixel_values"])
625
+ patch_image_features = (
626
+ self._get_vision_model_output(image_input["patch_pixel_values"])
627
+ if len(image_input["patch_pixel_values"]) > 0
628
+ else None
629
+ )
630
+ num_patches = image_input["num_patches"]
631
632
image_features = self._process_image_features(image_features)
633
patch_image_features = (