[Bugfix][Frontend] Let pooling requests set padding (#51157)
Signed-off-by: Hert4 <ductransa01@gmail.com>
Hert4 committed
Aug 27, 2026 at 18:06 UTC
4a6a3272e8d75518efe0a6f9393eb504f3ed2ee0
4 files changed
+70
-12
examples/pooling/embed/vision_embedding_offline.py
+7
-1
@@ -154,7 +154,13 @@ def run_siglip(seed: int):
154
)
155
156
print("Text embedding output:")
157
- outputs = llm.embed(text, use_tqdm=False)
157
+ # SigLIP is trained with padding to a fixed length and no attention mask, so
158
+ # text embeddings only line up with image embeddings when they are padded.
159
+ outputs = llm.embed(
160
+ text,
161
+ tokenization_kwargs={"padding": "max_length"},
162
+ use_tqdm=False,
163
+ )
164
print_embeddings(outputs[0].outputs.embedding)
165
166
print("Image embedding output:")
examples/pooling/embed/vision_embedding_online.py
+5
@@ -35,6 +35,7 @@ def create_chat_embeddings(
35
encoding_format: Literal["base64", "float"] | NotGiven = NOT_GIVEN,
36
continue_final_message: bool = False,
37
add_special_tokens: bool = False,
38
+ padding: str | NotGiven = NOT_GIVEN,
39
) -> CreateEmbeddingResponse:
40
"""
41
Convenience function for accessing vLLM's Chat Embeddings API,
@@ -49,6 +50,7 @@ def create_chat_embeddings(
50
"encoding_format": encoding_format,
51
"continue_final_message": continue_final_message,
52
"add_special_tokens": add_special_tokens,
53
+ "padding": padding,
54
},
55
)
56
@@ -288,6 +290,8 @@ def run_siglip(client: OpenAI, model: str):
290
291
print("Image embedding output:", response.data[0].embedding)
292
293
+ # SigLIP is trained with padding to a fixed length and no attention mask, so
294
+ # text embeddings only line up with image embeddings when they are padded.
295
response = create_chat_embeddings(
296
client,
297
messages=[
@@ -300,6 +304,7 @@ def run_siglip(client: OpenAI, model: str):
304
],
305
model=model,
306
encoding_format="float",
307
+ padding="max_length",
308
)
309
310
print("Text embedding output:", response.data[0].embedding)
tests/entrypoints/pooling/embed/test_online.py
+31
@@ -295,6 +295,37 @@ async def test_truncate_prompt_tokens(client: openai.AsyncOpenAI, model_name: st
295
)
296
297
298
+@pytest.mark.asyncio
299
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
300
+async def test_padding(client: openai.AsyncOpenAI, model_name: str):
301
+ input_texts = ["Hello"]
302
+
303
+ async def prompt_tokens(**extra_body):
304
+ response = await client.embeddings.create(
305
+ model=model_name,
306
+ input=input_texts,
307
+ extra_body=extra_body,
308
+ )
309
+ embeddings = EmbeddingResponse.model_validate(response.model_dump(mode="json"))
310
+ assert len(embeddings.data) == 1
311
+ return embeddings.usage.prompt_tokens
312
+
313
+ unpadded = await prompt_tokens()
314
+ padded = await prompt_tokens(padding="max_length")
315
+
316
+ # Models trained with a fixed sequence length and no attention mask need
317
+ # this; without it their embeddings are not comparable.
318
+ assert padded > unpadded
319
+ assert await prompt_tokens(padding="do_not_pad") == unpadded
320
+
321
+ with pytest.raises(openai.BadRequestError):
322
+ await client.embeddings.create(
323
+ model=model_name,
324
+ input=input_texts,
325
+ extra_body={"padding": "longest"},
326
+ )
327
+
328
+
329
@pytest.mark.asyncio
330
@pytest.mark.parametrize("model_name", [MODEL_NAME])
331
async def test_chat_request(
vllm/entrypoints/pooling/base/protocol.py
+27
-11
@@ -39,6 +39,17 @@ class PoolingBasicRequestMixin(OpenAIBaseModel):
39
40
# --8<-- [start:pooling-common-extra-params]
41
truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
42
+ padding: Literal["max_length", "do_not_pad"] | None = Field(
43
+ default=None,
44
+ description=(
45
+ "Whether to pad the prompt, using the same names as the "
46
+ "Transformers tokenizer. 'max_length' pads to the maximum input "
47
+ "length, 'do_not_pad' leaves the prompt as is. Models trained "
48
+ "with a fixed sequence length and no attention mask, such as "
49
+ "SigLIP, need 'max_length' to match training, otherwise their "
50
+ "embeddings are not comparable."
51
+ ),
52
+ )
53
truncation_side: Literal["left", "right"] | None = Field(
54
default=None,
55
description=(
@@ -94,7 +105,17 @@ class PoolingBasicRequestMixin(OpenAIBaseModel):
105
) -> TokenizeParams:
106
encoder_config = model_config.encoder_config or {}
107
if max_output_tokens_param is None:
97
- return TokenizeParams(
108
+ tok_params = TokenizeParams(
109
+ max_total_tokens=max_total_tokens,
110
+ max_output_tokens=max_output_tokens,
111
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
112
+ truncation_side=self.truncation_side,
113
+ do_lower_case=encoder_config.get("do_lower_case", False),
114
+ add_special_tokens=add_special_tokens,
115
+ max_total_tokens_param=max_total_tokens_param,
116
+ )
117
+ else:
118
+ tok_params = TokenizeParams(
119
max_total_tokens=max_total_tokens,
120
max_output_tokens=max_output_tokens,
121
truncate_prompt_tokens=self.truncate_prompt_tokens,
@@ -102,18 +123,13 @@ class PoolingBasicRequestMixin(OpenAIBaseModel):
123
do_lower_case=encoder_config.get("do_lower_case", False),
124
add_special_tokens=add_special_tokens,
125
max_total_tokens_param=max_total_tokens_param,
126
+ max_output_tokens_param=max_output_tokens_param,
127
)
128
107
- return TokenizeParams(
108
- max_total_tokens=max_total_tokens,
109
- max_output_tokens=max_output_tokens,
110
- truncate_prompt_tokens=self.truncate_prompt_tokens,
111
- truncation_side=self.truncation_side,
112
- do_lower_case=encoder_config.get("do_lower_case", False),
113
- add_special_tokens=add_special_tokens,
114
- max_total_tokens_param=max_total_tokens_param,
115
- max_output_tokens_param=max_output_tokens_param,
116
- )
129
+ if self.padding is None:
130
+ return tok_params
131
+
132
+ return tok_params.with_kwargs(padding=self.padding)
133
134
135
class PoolingTokenizeParamsMixin: