[Bugfix][Multimodal] Encode text in mixed CLIP/SigLIP pooling batches (#53165)
Signed-off-by: Prudhvivuda <prudhvi12042001@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Prudhvi Vuda committed
Aug 24, 2026 at 13:45 UTC
342b8ebd8bd4595826f29ff95dfc48679a03a95a
5 files changed
+154
-11
tests/models/multimodal/pooling/test_clip.py
+26
@@ -52,6 +52,31 @@ def _run_test(
52
vllm_model.embed(texts)
53
vllm_model.embed([""], images=images)
54
55
+ # Mixed image+text batch must not skip the text encoder (#53091).
56
+ mixed_outputs = vllm_outputs_per_case[2]
57
+ check_embeddings_close(
58
+ embeddings_0_lst=[vllm_outputs_per_case[0][0]],
59
+ embeddings_1_lst=[mixed_outputs[0]],
60
+ name_0="text_only",
61
+ name_1="mixed_text",
62
+ )
63
+ check_embeddings_close(
64
+ embeddings_0_lst=[vllm_outputs_per_case[1][0]],
65
+ embeddings_1_lst=[mixed_outputs[1]],
66
+ name_0="image_only",
67
+ name_1="mixed_image",
68
+ )
69
+ empty_text = vllm_model.embed([""])
70
+ empty_sim = torch.nn.functional.cosine_similarity(
71
+ torch.tensor(mixed_outputs[0]),
72
+ torch.tensor(empty_text[0]),
73
+ dim=0,
74
+ )
75
+ assert empty_sim < 0.99, (
76
+ "Mixed-batch text embedding collapsed to the empty-string vector "
77
+ f"(cosine={empty_sim:.4f})"
78
+ )
79
+
80
with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model:
81
hf_outputs_per_case = []
82
for input_texts, input_images in input_cases:
@@ -101,6 +126,7 @@ def test_models(
126
input_cases = [
127
(HF_TEXT_PROMPTS, text_images),
128
(HF_IMAGE_PROMPTS, images),
129
+ ([HF_TEXT_PROMPTS[0], HF_IMAGE_PROMPTS[0]], [None, images[0]]),
130
]
131
132
_run_test(
tests/models/multimodal/pooling/test_dual_encoder_routing.py
new
+42
@@ -0,0 +1,42 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from vllm.model_executor.models.clip import (
8
+ dual_encoder_has_text_tokens,
9
+ merge_dual_encoder_text_and_vision,
10
+)
11
+
12
+pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
13
+
14
+
15
+def test_text_only_batch_runs_text_encoder():
16
+ assert dual_encoder_has_text_tokens(False, None)
17
+ assert dual_encoder_has_text_tokens(False, torch.zeros(4, dtype=torch.bool))
18
+
19
+
20
+def test_vision_only_batch_skips_text_encoder():
21
+ is_multimodal = torch.ones(8, dtype=torch.bool)
22
+ assert not dual_encoder_has_text_tokens(True, is_multimodal)
23
+
24
+
25
+def test_mixed_batch_runs_text_encoder():
26
+ # Image tokens followed by a text sequence: the old batch-wide flag
27
+ # treated this as vision-only and skipped the text encoder (#53091).
28
+ is_multimodal = torch.tensor([True, True, True, False, False, False])
29
+ assert dual_encoder_has_text_tokens(True, is_multimodal)
30
+
31
+
32
+def test_missing_token_mask_with_mm_embeddings_is_vision_only():
33
+ assert not dual_encoder_has_text_tokens(True, None)
34
+
35
+
36
+def test_merge_keeps_vision_on_mm_tokens():
37
+ text = torch.zeros(4, 2)
38
+ vision = torch.ones(4, 2)
39
+ is_multimodal = torch.tensor([True, True, False, False])
40
+ out = merge_dual_encoder_text_and_vision(text, vision, is_multimodal)
41
+ assert torch.equal(out[:2], vision[:2])
42
+ assert torch.equal(out[2:], text[2:])
tests/models/multimodal/pooling/test_siglip.py
+22
@@ -63,6 +63,20 @@ def _run_test(
63
vllm_model.embed(texts)
64
vllm_model.embed([""], images=images)
65
66
+ mixed_outputs = vllm_outputs_per_case[2]
67
+ check_embeddings_close(
68
+ embeddings_0_lst=[vllm_outputs_per_case[0][0]],
69
+ embeddings_1_lst=[mixed_outputs[0]],
70
+ name_0="text_only",
71
+ name_1="mixed_text",
72
+ )
73
+ check_embeddings_close(
74
+ embeddings_0_lst=[vllm_outputs_per_case[1][0]],
75
+ embeddings_1_lst=[mixed_outputs[1]],
76
+ name_0="image_only",
77
+ name_1="mixed_image",
78
+ )
79
+
80
with hf_runner(model, dtype=dtype, auto_cls=SiglipModel) as hf_model:
81
hf_outputs_per_case = []
82
for input_texts, input_images, tokenization_kwargs in input_cases:
@@ -122,6 +136,14 @@ def test_models(
136
},
137
),
138
(HF_IMAGE_PROMPTS, images, {}),
139
+ (
140
+ [HF_TEXT_PROMPTS[0], HF_IMAGE_PROMPTS[0]],
141
+ [None, images[0]],
142
+ {
143
+ "padding": "max_length",
144
+ "max_length": 64,
145
+ },
146
+ ),
147
]
148
149
_run_test(
vllm/model_executor/models/clip.py
+47
-6
@@ -119,6 +119,38 @@ def _get_vision_feature_select_strategy(pooling_type: str):
119
) from None
120
121
122
+def dual_encoder_has_text_tokens(
123
+ has_mm_embeddings: bool,
124
+ is_multimodal: torch.Tensor | None,
125
+) -> bool:
126
+ """Whether a dual-encoder pooling batch still contains text tokens.
127
+
128
+ ``embed_input_ids`` used to treat "the batch has any image embeddings" as
129
+ "the whole batch is vision-only". Mixed image+text batches then skipped
130
+ the text encoder, so completion-style text embeddings collapsed (issue
131
+ #53091). Honor the per-token ``is_multimodal`` mask instead.
132
+ """
133
+ if not has_mm_embeddings or is_multimodal is None:
134
+ return not has_mm_embeddings
135
+ return bool((~is_multimodal).any().item())
136
+
137
+
138
+def merge_dual_encoder_text_and_vision(
139
+ text_features: torch.Tensor,
140
+ vision_embeds: torch.Tensor,
141
+ is_multimodal: torch.Tensor | None,
142
+) -> torch.Tensor:
143
+ """Restore vision embeddings on multimodal tokens after the text encoder.
144
+
145
+ The model runner builds ``is_multimodal`` on CPU; text and vision features
146
+ are on the compute device.
147
+ """
148
+ if is_multimodal is None or not bool(is_multimodal.any().item()):
149
+ return text_features
150
+ mm_mask = is_multimodal.to(device=vision_embeds.device, non_blocking=True)
151
+ return torch.where(mm_mask.unsqueeze(-1), vision_embeds, text_features)
152
+
153
+
154
class CLIPProcessingInfo(BaseProcessingInfo):
155
def get_hf_config(self):
156
return self.ctx.get_hf_config(CLIPConfig)
@@ -809,8 +841,9 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
841
842
self.pooler = DispatchPooler.for_embedding(pooler_config)
843
812
- # Assumes that self.forward is called after self.embed_input_ids
813
- self._is_text_input = True
844
+ # Set in embed_input_ids; consumed by forward.
845
+ self._has_text_tokens = True
846
+ self._mm_token_mask: torch.Tensor | None = None
847
848
def get_text_features(
849
self,
@@ -908,8 +941,12 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
941
*,
942
is_multimodal: torch.Tensor | None = None,
943
) -> torch.Tensor:
911
- self._is_text_input = (
912
- multimodal_embeddings is None or len(multimodal_embeddings) == 0
944
+ has_mm_embeddings = (
945
+ multimodal_embeddings is not None and len(multimodal_embeddings) > 0
946
+ )
947
+ self._mm_token_mask = is_multimodal
948
+ self._has_text_tokens = dual_encoder_has_text_tokens(
949
+ has_mm_embeddings, is_multimodal
950
)
951
952
# This is to satisfy the type checker for each overload
@@ -943,9 +980,10 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
980
981
# Multimodal inputs
982
assert inputs_embeds is not None
946
- if not self._is_text_input:
983
+ if not self._has_text_tokens:
984
return inputs_embeds
985
986
+ vision_embeds = inputs_embeds
987
# NOTE: inputs_embeds in model runner has size text_config.projection_dim
988
# (instead of text_config.hidden_size) to accommodate image embeddings
989
hidden_size = self.text_embed_dim
@@ -955,7 +993,10 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
993
# No need to handle this case for now
994
raise NotImplementedError
995
958
- return self.get_text_features(input_ids, positions, inputs_embeds)
996
+ text_features = self.get_text_features(input_ids, positions, inputs_embeds)
997
+ return merge_dual_encoder_text_and_vision(
998
+ text_features, vision_embeds, self._mm_token_mask
999
+ )
1000
1001
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
1002
loader = AutoWeightsLoader(
vllm/model_executor/models/siglip.py
+17
-5
@@ -57,6 +57,7 @@ from vllm.sequence import IntermediateTensors
57
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
58
from vllm.utils.tensor_schema import TensorSchema, TensorShape
59
60
+from .clip import dual_encoder_has_text_tokens, merge_dual_encoder_text_and_vision
61
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant
62
from .interfaces_base import default_pooling_type
63
from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
@@ -943,7 +944,9 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
944
945
self.pooler = DispatchPooler.for_embedding(pooler_config)
946
946
- self._is_text_input = True
947
+ # Set in embed_input_ids; consumed by forward.
948
+ self._has_text_tokens = True
949
+ self._mm_token_mask: torch.Tensor | None = None
950
951
def get_text_features(
952
self,
@@ -1088,8 +1091,12 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
1091
*,
1092
is_multimodal: torch.Tensor | None = None,
1093
) -> torch.Tensor:
1091
- self._is_text_input = (
1092
- multimodal_embeddings is None or len(multimodal_embeddings) == 0
1094
+ has_mm_embeddings = (
1095
+ multimodal_embeddings is not None and len(multimodal_embeddings) > 0
1096
+ )
1097
+ self._mm_token_mask = is_multimodal
1098
+ self._has_text_tokens = dual_encoder_has_text_tokens(
1099
+ has_mm_embeddings, is_multimodal
1100
)
1101
1102
if multimodal_embeddings is None or is_multimodal is None:
@@ -1121,9 +1128,11 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
1128
raise RuntimeError("PP is not supported for this model")
1129
1130
# Multimodal inputs (image embeddings)
1124
- if not self._is_text_input:
1131
+ assert inputs_embeds is not None
1132
+ if not self._has_text_tokens:
1133
return inputs_embeds
1134
1135
+ vision_embeds = inputs_embeds
1136
# NOTE: inputs_embeds in model runner has size text_config.projection_size
1137
# (instead of text_config.hidden_size) to accommodate image embeddings
1138
hidden_size = self.text_embed_dim
@@ -1133,7 +1142,10 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
1142
# No need to handle this case for now
1143
raise NotImplementedError
1144
1136
- return self.get_text_features(input_ids, positions, inputs_embeds)
1145
+ text_features = self.get_text_features(input_ids, positions, inputs_embeds)
1146
+ return merge_dual_encoder_text_and_vision(
1147
+ text_features, vision_embeds, self._mm_token_mask
1148
+ )
1149
1150
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
1151
loader = AutoWeightsLoader(