[data][llm] Add pooling parameter

[jeffreywang-anyscale]

Description

The preprocessor forwards only sampling_params to the engine today. For task_type="embed", however, we should also allow forwarding pooling_params, enabling features such as truncating the input prompt to a fixed token budget via truncate_prompt_tokens or normalizing the output embedding via normalize. See https://docs.vllm.ai/en/latest/api/vllm/#vllm.PoolingParams for a comprehensive list of supported attributes.

Related issues

Resolves #57805

Additional information

• Similar to sampling parameters, we allow users to specify pooling parameter under the "pooling_params" column for embedding tasks.
• Add tests in test_vllm_engine_stage to validate that the pooling parameters are received by the engine.

[gemini-code-assist]

Code Review

This pull request adds support for pooling_params in vLLM embedding tasks, allowing users to specify parameters like token truncation and embedding normalization. The changes include updating the vLLM engine stage to handle these parameters and adding corresponding tests to validate the new functionality. My review includes a couple of suggestions for improving maintainability and aligning the implementation with the documented behavior.

[kouroshHakha]

also test None, and {} ?

[jeffreywang-anyscale]

Added a test case for empty dict. None is not a possible value as we default to empty dict if pooling_params is not provided.

[kouroshHakha]

Critical can we test on some other property?

you want to basically test whether truncation is applied or whether normalizatio is applied.
It is not sufficient to check for request.params values to match what was sent in.

Idea: We can test that on input x the answer will be different comparing truncation=None vs. truncation=2, similarly we can test normalize=False vs. normalize=True

[jeffreywang-anyscale]

Validating the difference of outputs is a good idea.

[jeffreywang-anyscale]

sampling_params do not apply to encode. encode is deterministic.

[jeffreywang-anyscale]

Added a couple more test cases:

• Compare truncation=None vs. truncation=3
• Compare normalize=False vs. normalize=True
• Validate that truncation is effective on long prompts

[jeffreywang-anyscale]

In the latest revision, introduced a small fix:

Pooling parameter's truncate_prompt_tokens is not respected by AsyncLLMEngine.encode(). I filed an vllm-project/vllm#31012 in vLLM and have a vllm-project/vllm#31013 for it.

As a temporary solution, prompt truncation is handled through the truncate_prompt_tokens argument passed to AsyncLLMEngine.encode. From Ray's users perspective, any value provided via pooling_params will be honored as expected.

[kouroshHakha]

lgtm