[RFC]: Make any vLLM model a pooling model

[DarkLight1337]

Motivation.

Currently, we have to open new PRs to add pooling functionality for existing architectures supported in vLLM. Since the code involved is basically the same for each model, there is potential to automate away this boilerplate.

Proposed Change.

Implement a pooling adapter that can be applied to any existing text generation model in vLLM. To preserve features such as LoRA, PP and multimodality, the adapter simply creates a new subclass of the original model.

• Embedding model adapter: [Model] Replace embedding models with pooling adapter #10769
• Classification model adapter [Model] Automatic conversion of classification and reward models #11469
• Reward model adapter [Model] Automatic conversion of classification and reward models #11469

The pooling adapter to apply depends on the purpose of the model. To facilitate this, the embedding task will be split up, so that the user can specify which adapter to apply to the model via --task:

• [Misc] Split up pooling tasks #10820
• [Frontend] Online Pooling API #11457

Meanwhile, current embedding-related classes will be renamed to avoid confusion between embed and other pooling tasks:

• [Misc] Rename embedding classes to pooling #10801

Feedback Period.

1-2 weeks

CC List.

@youkaichao @mgoin @robertgshaw2-neuralmagic @maxdebayser

Any Other Things.

Note that we can still directly map to pooling models in the model registry. This is used when the model architecture has different pooling defaults (e.g. pooling_type=CLS for BERT) or additional modules (e.g. score in Qwen2ForRewardModel). For models that already support pooling, the adapter returns the original model without modifications.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[maxdebayser]

I think this makes sense, but we should add a verification to prevent models that are already loaded as embedding models to be wrapped again. Would this wrapper assume a default Pooling configuration or should we raise an error if the user doesn't provide this information on the command line?

[DarkLight1337]

I think this makes sense, but we should add a verification to prevent models that are already loaded as embedding models to be wrapped again. Would this wrapper assume a default Pooling configuration or should we raise an error if the user doesn't provide this information on the command line?

I'm thinking of removing the existing embedding model implementations and replacing them with the wrapper. The embedding wrapper's default pooler config is shown in the above example. I'm thinking of resolving the pooler config in the same way as current code (i.e.: user > sentence-transformers > model), so the user doesn't have to provide this information in most cases.

[maxdebayser]

What would happen if the user loads a BertModel without --task-embedding? Would there be a validation to prevent this from happening our would we allow the decoder-only model runner to try to run the model for text generation?

[DarkLight1337]

In is_text_generation_model, we check for the existence of compute_logits and sample methods. Since those methods aren't in BertModel, --task generation is not allowed for that model.

[maxdebayser]

Ok, I think this makes sense.

[DarkLight1337]

Update: I have expanded the scope of this RFC to cover all pooling-based models.