[Core] Set pooling params based on task and model #21128

DarkLight1337 · 2025-07-17T16:15:04Z

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

Set pooling_params.use_cross_encoder and pooling_params.logits_processing_needs_token_ids based on the task and model, replacing the static logic at API level. This provides more flexibility for models to respond differently based on the task being used.

Implements #21058 (comment)

Test Plan

The existing tests should pass

Test Result

(Optional) Documentation Update

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

github-actions · 2025-07-17T16:15:12Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

DarkLight1337 · 2025-07-17T16:16:04Z

vllm/v1/worker/gpu_input_batch.py

        # req_index -> bad_words_token_ids
        self.bad_words_token_ids: dict[int, list[list[int]]] = {}

+        self.logits_processing_needs_token_ids = np.zeros(max_num_reqs,


Now this can be set on a per-request basis instead of being fixed at startup time

gemini-code-assist

Code Review

This pull request refactors how pooling parameters are handled by moving the logic from the API layer to the model/worker layer, which is a great improvement for flexibility and encapsulation. The changes are well-implemented across the codebase.

I've found one critical issue in vllm/entrypoints/llm.py that will lead to a runtime ValueError due to a task mismatch during verification. Please see the specific comment for details on the issue and how to resolve it. Other than that, the changes look solid.

vllm/entrypoints/llm.py

DarkLight1337 · 2025-07-17T16:18:47Z

vllm/worker/model_runner_base.py

    def get_model(self) -> nn.Module:
        raise NotImplementedError

+    def get_supported_pooling_tasks(self) -> list[PoolingTask]:


This is not technically used in V0 but I put it here to avoid forgetting to include this method when we implement the base class for V1 model runners

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

vllm/v1/worker/gpu_model_runner.py

maxdebayser

I've left a suggestion, but otherwise it LGTM

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Isotr0py

Overall LGTM!

vllm/entrypoints/openai/serving_embedding.py

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

DarkLight1337 · 2025-07-18T06:36:30Z

vllm/model_executor/models/gritlm.py

        **kwargs,
    ) -> None:
-        # Use full attention for pooling
+        # Use full attention for pooling (this is why V1 is not supported yet)


After this refactoring, the pooler implementation of GritLM becomes compatible with V1 model runner now (note that ~~logits_processing_needs_token_ids=True~~ requires_token_ids=True). However none of the V1 attention backends support bidirectional attention yet, so this model is still marked V0 only

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com>

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Signed-off-by: Paul Pak <paulpak58@gmail.com>

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Signed-off-by: Diego-Castan <diego.castan@ibm.com>

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

[Core] Set pooling params based on task and model

c880f24

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

DarkLight1337 requested a review from Isotr0py July 17, 2025 16:15

DarkLight1337 requested review from WoosukKwon, alexm-redhat, comaniac, njhill, robertgshaw2-redhat and ywang96 as code owners July 17, 2025 16:15

DarkLight1337 added the ready ONLY add when PR is ready to merge/full CI is needed label Jul 17, 2025

DarkLight1337 requested review from aarnphm, youkaichao and zhuohan123 as code owners July 17, 2025 16:15

mergify bot added frontend v1 tpu Related to Google TPUs labels Jul 17, 2025

DarkLight1337 commented Jul 17, 2025

View reviewed changes

gemini-code-assist bot reviewed Jul 17, 2025

View reviewed changes

vllm/entrypoints/llm.py Show resolved Hide resolved

DarkLight1337 commented Jul 17, 2025

View reviewed changes

DarkLight1337 added 2 commits July 17, 2025 16:21

Fix mypy

7ce1ed5

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Fix

7b52f6e

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

maxdebayser reviewed Jul 17, 2025

View reviewed changes

vllm/v1/worker/gpu_model_runner.py Show resolved Hide resolved

maxdebayser reviewed Jul 17, 2025

View reviewed changes

vllm/v1/worker/gpu_model_runner.py Outdated Show resolved Hide resolved

maxdebayser approved these changes Jul 17, 2025

View reviewed changes

DarkLight1337 added 2 commits July 18, 2025 03:33

Fix GritLM

cf0730c

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Fix API server

ff6dd3f

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Isotr0py approved these changes Jul 18, 2025

View reviewed changes

vllm/entrypoints/openai/serving_embedding.py Outdated Show resolved Hide resolved

DarkLight1337 added 3 commits July 18, 2025 04:17

Remove print

a1fde18

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Create a separate params class to apply updates

bb43104

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Remove redundant use_cross_encoder

21feebc

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

DarkLight1337 added 2 commits July 18, 2025 06:08

Refactor GritLM

409f07d

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Rename

2bdc4cc

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

DarkLight1337 enabled auto-merge (squash) July 18, 2025 06:21

DarkLight1337 commented Jul 18, 2025

View reviewed changes

DarkLight1337 added 4 commits July 18, 2025 06:46

Rename logits_processing_needs_token_ids to requires_token_ids

78c656a

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Add docstring

6a05da1

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Merge branch 'main' into update-pooling-params

79fca42

Fix dummy run

30a4b98

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

vllm-bot merged commit 45badd0 into vllm-project:main Jul 18, 2025
63 of 66 checks passed

DarkLight1337 deleted the update-pooling-params branch July 18, 2025 12:42

Potabk mentioned this pull request Jul 19, 2025

[Bugfix] Fix break CI vllm-project/vllm-ascend#1888

Closed

x22x22 pushed a commit to x22x22/vllm that referenced this pull request Aug 5, 2025

[Core] Set pooling params based on task and model (vllm-project#21128)

20a43f6

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Pradyun92 pushed a commit to Pradyun92/vllm that referenced this pull request Aug 6, 2025

[Core] Set pooling params based on task and model (vllm-project#21128)

03432f0

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

npanpaliya pushed a commit to odh-on-pz/vllm-upstream that referenced this pull request Aug 6, 2025

[Core] Set pooling params based on task and model (vllm-project#21128)

353cd2f

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

jinzhen-lin pushed a commit to jinzhen-lin/vllm that referenced this pull request Aug 9, 2025

[Core] Set pooling params based on task and model (vllm-project#21128)

0a3fa41

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com>

paulpak58 pushed a commit to paulpak58/vllm that referenced this pull request Aug 13, 2025

[Core] Set pooling params based on task and model (vllm-project#21128)

fe8790e

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Signed-off-by: Paul Pak <paulpak58@gmail.com>

diegocastanibm pushed a commit to diegocastanibm/vllm that referenced this pull request Aug 15, 2025

[Core] Set pooling params based on task and model (vllm-project#21128)

f5ee201

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Signed-off-by: Diego-Castan <diego.castan@ibm.com>

epwalsh pushed a commit to epwalsh/vllm that referenced this pull request Aug 27, 2025

[Core] Set pooling params based on task and model (vllm-project#21128)

58992d5

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

[Core] Set pooling params based on task and model #21128

[Core] Set pooling params based on task and model #21128

DarkLight1337 commented Jul 17, 2025 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Jul 17, 2025

Uh oh!

DarkLight1337 Jul 17, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

DarkLight1337 Jul 17, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

maxdebayser left a comment

Uh oh!

Isotr0py left a comment

Uh oh!

Uh oh!

DarkLight1337 Jul 18, 2025 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

[Core] Set pooling params based on task and model #21128

[Core] Set pooling params based on task and model #21128

Conversation

DarkLight1337 commented Jul 17, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Test Plan

Test Result

(Optional) Documentation Update

Uh oh!

github-actions bot commented Jul 17, 2025

Uh oh!

DarkLight1337 Jul 17, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

DarkLight1337 Jul 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

maxdebayser left a comment

Choose a reason for hiding this comment

Uh oh!

Isotr0py left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

DarkLight1337 Jul 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

DarkLight1337 commented Jul 17, 2025 •

edited by github-actions bot

Loading

DarkLight1337 Jul 17, 2025 •

edited

Loading

DarkLight1337 Jul 18, 2025 •

edited

Loading