[Model][1/N] Support multiple poolers at model level

[DarkLight1337]

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.

FIX #21277

Purpose

This PR enables models to support multiple poolers, by:

• Splitting up the task validation part of Pooler.get_pooling_updates into Pooler.get_supported_tasks.
• Introducing DispatchPooler which dispatches pooler calls to sub-poolers based on the pooling task.

Apart from enabling models to customize each pooling task, this would eventually also let us output intermediate results within the same request, e.g.:

• Hidden states via the encode pooler with ALL pooling.
• Logits by adapting LogitsProcessor into a pooler.

Future work:

• Multi-task support based on DispatchPooler is not enabled at the API level yet. We still check model_config.supported_tasks which is based on the value of --task. This will be addressed in the next PR.
• To avoid confusing with pooling task, we will rename --task to --convert and ask users to pass --runner pooling to use pooling models in the next PR, emitting a deprecation warning to users still using --task.
• Memory profiling is based on list(Pooler.get_supported_tasks())[0] instead of running profiling on every sub-pooler and getting the maximum memory usage. So, it is possible for OOM to occur during inference time. We can refine this in a later PR.
• Disabling chunked prefill is still based on the pooling_type in the pooler config instead of the actual poolers being used, because the poolers are only created after model construction. Can we disable chunked prefill inside the model runner after the model is initialized?
• For now, Sentence Transformers pooler config overrides the poolers for all tasks. We should only override the "main" pooler based on the model architecture.

Notes:

• The output of LLM.encode (Pooling API) for embedding and classification models (except for models with Sentence Transforming pooler config) is changed by this PR because the default encode pooler is different from the default embed and default classification pooler. While technically a breaking change, users aren't supposed to use this API for embedding and classification tasks in the first place anyway.

Other changes:

• Renamed --task pooling to --task encode. This was introduced very recently in [Core] Support multiple tasks per model #20771 and is not documented, so this change should barely affect any users.
• Split up BertModel into a pooling variant and a non pooling variant to keep the type checker happy.

cc @maxdebayser @noooop @22quinn

Test Plan

The existing tests should pass.

Test Result

(Optional) Documentation Update

Updated the Pooling Models page to explain the new mechanism.

[github-actions]

👋 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.

🚀

[gemini-code-assist]

Code Review

This pull request introduces a major refactoring to support multiple pooling tasks at the model level. It introduces a DispatchPooler to route requests to different sub-poolers based on the task. The Pooler interface is updated with get_supported_tasks and helper methods like for_encode, for_embed, and for_classify. Many model implementations are updated to use this new structure. The changes are well-structured and improve modularity. I've found one high-severity issue related to weight loading in the new BertPoolingModel.

[Isotr0py]

Overall LGTM! Just a nit.

[noooop]

• Disabling chunked prefill is still based on the pooling_type in the pooler config instead of the actual poolers being used, because the poolers are only created after model construction. Can we disable chunked prefill inside the model runner after the model is initialized?

disabling chunked prefill & auto prefix cache should be controlled by attn_type rather than pooling_type. For instance, Alibaba-NLP/gte-Qwen2-1.5B-instruct uses encoder-only attention, which should have chunked prefill & auto prefix cache disabled.

[maxdebayser]

LGTM