[Bugfix][Frontend] Let pooling requests set padding (#51157)

Signed-off-by: Hert4 <ductransa01@gmail.com>

examples/pooling/embed/vision_embedding_offline.py

@@ -154,7 +154,13 @@ def run_siglip(seed: int):
     )
 
     print("Text embedding output:")
-    outputs = llm.embed(text, use_tqdm=False)
+    # SigLIP is trained with padding to a fixed length and no attention mask, so
+    # text embeddings only line up with image embeddings when they are padded.
+    outputs = llm.embed(
+        text,
+        tokenization_kwargs={"padding": "max_length"},
+        use_tqdm=False,
+    )
     print_embeddings(outputs[0].outputs.embedding)
 
     print("Image embedding output:")

examples/pooling/embed/vision_embedding_online.py

@@ -35,6 +35,7 @@ def create_chat_embeddings(
     encoding_format: Literal["base64", "float"] | NotGiven = NOT_GIVEN,
     continue_final_message: bool = False,
     add_special_tokens: bool = False,
+    padding: str | NotGiven = NOT_GIVEN,
 ) -> CreateEmbeddingResponse:
     """
     Convenience function for accessing vLLM's Chat Embeddings API,
@@ -49,6 +50,7 @@ def create_chat_embeddings(
             "encoding_format": encoding_format,
             "continue_final_message": continue_final_message,
             "add_special_tokens": add_special_tokens,
+            "padding": padding,
         },
     )
 
@@ -288,6 +290,8 @@ def run_siglip(client: OpenAI, model: str):
 
     print("Image embedding output:", response.data[0].embedding)
 
+    # SigLIP is trained with padding to a fixed length and no attention mask, so
+    # text embeddings only line up with image embeddings when they are padded.
     response = create_chat_embeddings(
         client,
         messages=[
@@ -300,6 +304,7 @@ def run_siglip(client: OpenAI, model: str):
         ],
         model=model,
         encoding_format="float",
+        padding="max_length",
     )
 
     print("Text embedding output:", response.data[0].embedding)

tests/entrypoints/pooling/embed/test_online.py

@@ -295,6 +295,37 @@ async def test_truncate_prompt_tokens(client: openai.AsyncOpenAI, model_name: st
         )
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_padding(client: openai.AsyncOpenAI, model_name: str):
+    input_texts = ["Hello"]
+
+    async def prompt_tokens(**extra_body):
+        response = await client.embeddings.create(
+            model=model_name,
+            input=input_texts,
+            extra_body=extra_body,
+        )
+        embeddings = EmbeddingResponse.model_validate(response.model_dump(mode="json"))
+        assert len(embeddings.data) == 1
+        return embeddings.usage.prompt_tokens
+
+    unpadded = await prompt_tokens()
+    padded = await prompt_tokens(padding="max_length")
+
+    # Models trained with a fixed sequence length and no attention mask need
+    # this; without it their embeddings are not comparable.
+    assert padded > unpadded
+    assert await prompt_tokens(padding="do_not_pad") == unpadded
+
+    with pytest.raises(openai.BadRequestError):
+        await client.embeddings.create(
+            model=model_name,
+            input=input_texts,
+            extra_body={"padding": "longest"},
+        )
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
 async def test_chat_request(

vllm/entrypoints/pooling/base/protocol.py

@@ -39,6 +39,17 @@ class PoolingBasicRequestMixin(OpenAIBaseModel):
 
     # --8<-- [start:pooling-common-extra-params]
     truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
+    padding: Literal["max_length", "do_not_pad"] | None = Field(
+        default=None,
+        description=(
+            "Whether to pad the prompt, using the same names as the "
+            "Transformers tokenizer. 'max_length' pads to the maximum input "
+            "length, 'do_not_pad' leaves the prompt as is. Models trained "
+            "with a fixed sequence length and no attention mask, such as "
+            "SigLIP, need 'max_length' to match training, otherwise their "
+            "embeddings are not comparable."
+        ),
+    )
     truncation_side: Literal["left", "right"] | None = Field(
         default=None,
         description=(
@@ -94,7 +105,7 @@ class PoolingBasicRequestMixin(OpenAIBaseModel):
     ) -> TokenizeParams:
         encoder_config = model_config.encoder_config or {}
         if max_output_tokens_param is None:
-            return TokenizeParams(
+            tok_params = TokenizeParams(
                 max_total_tokens=max_total_tokens,
                 max_output_tokens=max_output_tokens,
                 truncate_prompt_tokens=self.truncate_prompt_tokens,
@@ -103,17 +114,22 @@ class PoolingBasicRequestMixin(OpenAIBaseModel):
                 add_special_tokens=add_special_tokens,
                 max_total_tokens_param=max_total_tokens_param,
             )
+        else:
+            tok_params = TokenizeParams(
+                max_total_tokens=max_total_tokens,
+                max_output_tokens=max_output_tokens,
+                truncate_prompt_tokens=self.truncate_prompt_tokens,
+                truncation_side=self.truncation_side,
+                do_lower_case=encoder_config.get("do_lower_case", False),
+                add_special_tokens=add_special_tokens,
+                max_total_tokens_param=max_total_tokens_param,
+                max_output_tokens_param=max_output_tokens_param,
+            )
 
-        return TokenizeParams(
-            max_total_tokens=max_total_tokens,
-            max_output_tokens=max_output_tokens,
-            truncate_prompt_tokens=self.truncate_prompt_tokens,
-            truncation_side=self.truncation_side,
-            do_lower_case=encoder_config.get("do_lower_case", False),
-            add_special_tokens=add_special_tokens,
-            max_total_tokens_param=max_total_tokens_param,
-            max_output_tokens_param=max_output_tokens_param,
-        )
+        if self.padding is None:
+            return tok_params
+
+        return tok_params.with_kwargs(padding=self.padding)
 
 
 class PoolingTokenizeParamsMixin: