[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>
This commit is contained in:
Prudhvi Vuda
2026-08-24 10:45:42 -07:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 4f686e182a
commit 342b8ebd8b
5 changed files with 154 additions and 11 deletions
@@ -52,6 +52,31 @@ def _run_test(
vllm_model.embed(texts)
vllm_model.embed([""], images=images)
# Mixed image+text batch must not skip the text encoder (#53091).
mixed_outputs = vllm_outputs_per_case[2]
check_embeddings_close(
embeddings_0_lst=[vllm_outputs_per_case[0][0]],
embeddings_1_lst=[mixed_outputs[0]],
name_0="text_only",
name_1="mixed_text",
)
check_embeddings_close(
embeddings_0_lst=[vllm_outputs_per_case[1][0]],
embeddings_1_lst=[mixed_outputs[1]],
name_0="image_only",
name_1="mixed_image",
)
empty_text = vllm_model.embed([""])
empty_sim = torch.nn.functional.cosine_similarity(
torch.tensor(mixed_outputs[0]),
torch.tensor(empty_text[0]),
dim=0,
)
assert empty_sim < 0.99, (
"Mixed-batch text embedding collapsed to the empty-string vector "
f"(cosine={empty_sim:.4f})"
)
with hf_runner(model, dtype=dtype, auto_cls=CLIPModel) as hf_model:
hf_outputs_per_case = []
for input_texts, input_images in input_cases:
@@ -101,6 +126,7 @@ def test_models(
input_cases = [
(HF_TEXT_PROMPTS, text_images),
(HF_IMAGE_PROMPTS, images),
([HF_TEXT_PROMPTS[0], HF_IMAGE_PROMPTS[0]], [None, images[0]]),
]
_run_test(
@@ -0,0 +1,42 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from vllm.model_executor.models.clip import (
dual_encoder_has_text_tokens,
merge_dual_encoder_text_and_vision,
)
pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
def test_text_only_batch_runs_text_encoder():
assert dual_encoder_has_text_tokens(False, None)
assert dual_encoder_has_text_tokens(False, torch.zeros(4, dtype=torch.bool))
def test_vision_only_batch_skips_text_encoder():
is_multimodal = torch.ones(8, dtype=torch.bool)
assert not dual_encoder_has_text_tokens(True, is_multimodal)
def test_mixed_batch_runs_text_encoder():
# Image tokens followed by a text sequence: the old batch-wide flag
# treated this as vision-only and skipped the text encoder (#53091).
is_multimodal = torch.tensor([True, True, True, False, False, False])
assert dual_encoder_has_text_tokens(True, is_multimodal)
def test_missing_token_mask_with_mm_embeddings_is_vision_only():
assert not dual_encoder_has_text_tokens(True, None)
def test_merge_keeps_vision_on_mm_tokens():
text = torch.zeros(4, 2)
vision = torch.ones(4, 2)
is_multimodal = torch.tensor([True, True, False, False])
out = merge_dual_encoder_text_and_vision(text, vision, is_multimodal)
assert torch.equal(out[:2], vision[:2])
assert torch.equal(out[2:], text[2:])
@@ -63,6 +63,20 @@ def _run_test(
vllm_model.embed(texts)
vllm_model.embed([""], images=images)
mixed_outputs = vllm_outputs_per_case[2]
check_embeddings_close(
embeddings_0_lst=[vllm_outputs_per_case[0][0]],
embeddings_1_lst=[mixed_outputs[0]],
name_0="text_only",
name_1="mixed_text",
)
check_embeddings_close(
embeddings_0_lst=[vllm_outputs_per_case[1][0]],
embeddings_1_lst=[mixed_outputs[1]],
name_0="image_only",
name_1="mixed_image",
)
with hf_runner(model, dtype=dtype, auto_cls=SiglipModel) as hf_model:
hf_outputs_per_case = []
for input_texts, input_images, tokenization_kwargs in input_cases:
@@ -122,6 +136,14 @@ def test_models(
},
),
(HF_IMAGE_PROMPTS, images, {}),
(
[HF_TEXT_PROMPTS[0], HF_IMAGE_PROMPTS[0]],
[None, images[0]],
{
"padding": "max_length",
"max_length": 64,
},
),
]
_run_test(
+47 -6
View File
@@ -119,6 +119,38 @@ def _get_vision_feature_select_strategy(pooling_type: str):
) from None
def dual_encoder_has_text_tokens(
has_mm_embeddings: bool,
is_multimodal: torch.Tensor | None,
) -> bool:
"""Whether a dual-encoder pooling batch still contains text tokens.
``embed_input_ids`` used to treat "the batch has any image embeddings" as
"the whole batch is vision-only". Mixed image+text batches then skipped
the text encoder, so completion-style text embeddings collapsed (issue
#53091). Honor the per-token ``is_multimodal`` mask instead.
"""
if not has_mm_embeddings or is_multimodal is None:
return not has_mm_embeddings
return bool((~is_multimodal).any().item())
def merge_dual_encoder_text_and_vision(
text_features: torch.Tensor,
vision_embeds: torch.Tensor,
is_multimodal: torch.Tensor | None,
) -> torch.Tensor:
"""Restore vision embeddings on multimodal tokens after the text encoder.
The model runner builds ``is_multimodal`` on CPU; text and vision features
are on the compute device.
"""
if is_multimodal is None or not bool(is_multimodal.any().item()):
return text_features
mm_mask = is_multimodal.to(device=vision_embeds.device, non_blocking=True)
return torch.where(mm_mask.unsqueeze(-1), vision_embeds, text_features)
class CLIPProcessingInfo(BaseProcessingInfo):
def get_hf_config(self):
return self.ctx.get_hf_config(CLIPConfig)
@@ -809,8 +841,9 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
self.pooler = DispatchPooler.for_embedding(pooler_config)
# Assumes that self.forward is called after self.embed_input_ids
self._is_text_input = True
# Set in embed_input_ids; consumed by forward.
self._has_text_tokens = True
self._mm_token_mask: torch.Tensor | None = None
def get_text_features(
self,
@@ -908,8 +941,12 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
*,
is_multimodal: torch.Tensor | None = None,
) -> torch.Tensor:
self._is_text_input = (
multimodal_embeddings is None or len(multimodal_embeddings) == 0
has_mm_embeddings = (
multimodal_embeddings is not None and len(multimodal_embeddings) > 0
)
self._mm_token_mask = is_multimodal
self._has_text_tokens = dual_encoder_has_text_tokens(
has_mm_embeddings, is_multimodal
)
# This is to satisfy the type checker for each overload
@@ -943,9 +980,10 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
# Multimodal inputs
assert inputs_embeds is not None
if not self._is_text_input:
if not self._has_text_tokens:
return inputs_embeds
vision_embeds = inputs_embeds
# NOTE: inputs_embeds in model runner has size text_config.projection_dim
# (instead of text_config.hidden_size) to accommodate image embeddings
hidden_size = self.text_embed_dim
@@ -955,7 +993,10 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
# No need to handle this case for now
raise NotImplementedError
return self.get_text_features(input_ids, positions, inputs_embeds)
text_features = self.get_text_features(input_ids, positions, inputs_embeds)
return merge_dual_encoder_text_and_vision(
text_features, vision_embeds, self._mm_token_mask
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
loader = AutoWeightsLoader(
+17 -5
View File
@@ -57,6 +57,7 @@ from vllm.sequence import IntermediateTensors
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
from vllm.utils.tensor_schema import TensorSchema, TensorShape
from .clip import dual_encoder_has_text_tokens, merge_dual_encoder_text_and_vision
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant
from .interfaces_base import default_pooling_type
from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
@@ -943,7 +944,9 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
self.pooler = DispatchPooler.for_embedding(pooler_config)
self._is_text_input = True
# Set in embed_input_ids; consumed by forward.
self._has_text_tokens = True
self._mm_token_mask: torch.Tensor | None = None
def get_text_features(
self,
@@ -1088,8 +1091,12 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
*,
is_multimodal: torch.Tensor | None = None,
) -> torch.Tensor:
self._is_text_input = (
multimodal_embeddings is None or len(multimodal_embeddings) == 0
has_mm_embeddings = (
multimodal_embeddings is not None and len(multimodal_embeddings) > 0
)
self._mm_token_mask = is_multimodal
self._has_text_tokens = dual_encoder_has_text_tokens(
has_mm_embeddings, is_multimodal
)
if multimodal_embeddings is None or is_multimodal is None:
@@ -1121,9 +1128,11 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
raise RuntimeError("PP is not supported for this model")
# Multimodal inputs (image embeddings)
if not self._is_text_input:
assert inputs_embeds is not None
if not self._has_text_tokens:
return inputs_embeds
vision_embeds = inputs_embeds
# NOTE: inputs_embeds in model runner has size text_config.projection_size
# (instead of text_config.hidden_size) to accommodate image embeddings
hidden_size = self.text_embed_dim
@@ -1133,7 +1142,10 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
# No need to handle this case for now
raise NotImplementedError
return self.get_text_features(input_ids, positions, inputs_embeds)
text_features = self.get_text_features(input_ids, positions, inputs_embeds)
return merge_dual_encoder_text_and_vision(
text_features, vision_embeds, self._mm_token_mask
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
loader = AutoWeightsLoader(