[Model][1/N] Support multiple poolers at model level (#21227)

Signed-off-by: DarkLight1337 <[email protected]>

Author: DarkLight1337

docs/models/pooling_models.md

@@ -11,26 +11,51 @@ before returning them.
     As shown in the [Compatibility Matrix](../features/compatibility_matrix.md), most vLLM features are not applicable to
     pooling models as they only work on the generation or decode stage, so performance may not improve as much.
 
-For pooling models, we support the following `--task` options.
-The selected option sets the default pooler used to extract the final hidden states:
+If the model doesn't implement this interface, you can set `--task` which tells vLLM
+to convert the model into a pooling model.
 
-| Task                            | Pooling Type   | Normalization   | Softmax   |
-|---------------------------------|----------------|-----------------|-----------|
-| Embedding (`embed`)             | `LAST`         | ✅︎              | ❌         |
-| Classification (`classify`)     | `LAST`         | ❌               | ✅︎        |
-| Sentence Pair Scoring (`score`) | \*             | \*              | \*        |
+| `--task`   | Model type           | Supported pooling tasks       |
+|------------|----------------------|-------------------------------|
+| `embed`    | Embedding model      | `encode`, `embed`             |
+| `classify` | Classification model | `encode`, `classify`, `score` |
+| `reward`   | Reward model         | `encode`                      |
 
-\*The default pooler is always defined by the model.
+## Pooling Tasks
 
-!!! note
-    If the model's implementation in vLLM defines its own pooler, the default pooler is set to that instead of the one specified in this table.
+In vLLM, we define the following pooling tasks and corresponding APIs:
+
+| Task       | APIs               |
+|------------|--------------------|
+| `encode`   | `encode`           |
+| `embed`    | `embed`, `score`\* |
+| `classify` | `classify`         |
+| `score`    | `score`            |
+
+\*The `score` API falls back to `embed` task if the model does not support `score` task.
+
+Each pooling model in vLLM supports one or more of these tasks according to [Pooler.get_supported_tasks][vllm.model_executor.layers.Pooler.get_supported_tasks].
+
+By default, the pooler assigned to each task has the following attributes:
+
+| Task       | Pooling Type   | Normalization | Softmax |
+|------------|----------------|---------------|---------|
+| `encode`   | `ALL`          | ❌            | ❌      |
+| `embed`    | `LAST`         | ✅︎            | ❌      |
+| `classify` | `LAST`         | ❌            | ✅︎      |
+
+These defaults may be overridden by the model's implementation in vLLM.
 
 When loading [Sentence Transformers](https://huggingface.co/sentence-transformers) models,
-we attempt to override the default pooler based on its Sentence Transformers configuration file (`modules.json`).
+we attempt to override the defaults based on its Sentence Transformers configuration file (`modules.json`),
+which takes priority over the model's defaults.
+
+You can further customize this via the `--override-pooler-config` option,
+which takes priority over both the model's and Sentence Transformers's defaults.
+
+!!! note
 
-!!! tip
-    You can customize the model's pooling method via the `--override-pooler-config` option,
-    which takes priority over both the model's and Sentence Transformers's defaults.
+    The above configuration may be disregarded if the model's implementation in vLLM defines its own pooler
+    that is not based on [PoolerConfig][vllm.config.PoolerConfig].
 
 ## Offline Inference
 

tests/models/test_transformers.py

@@ -144,7 +144,7 @@ def test_quantization(
     "model",
     ["jason9693/Qwen2.5-1.5B-apeach"],
 )
-@pytest.mark.parametrize("dtype", ["half"])
+@pytest.mark.parametrize("dtype", ["float"])
 def test_classify(
     hf_runner,
     vllm_runner,

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_gemma_embedding.py

@@ -8,7 +8,7 @@
 import torch.nn as nn
 
 from vllm.config import VllmConfig
-from vllm.model_executor.layers.pooler import Pooler, PoolingType
+from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
 from vllm.model_executor.models.gemma2 import Gemma2Model
 from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
 from vllm.sequence import IntermediateTensors
@@ -26,12 +26,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.model = Gemma2Model(vllm_config=vllm_config,
                                  prefix=maybe_prefix(prefix, "model"))
 
-        self.pooler = Pooler.from_config_with_defaults(
-            vllm_config.model_config.pooler_config,
-            pooling_type=PoolingType.LAST,
-            normalize=True,
-            softmax=False,
-        )
+        pooler_config = vllm_config.model_config.pooler_config
+        assert pooler_config is not None
+
+        self.pooler = DispatchPooler({
+            "encode": Pooler.for_encode(pooler_config),
+            "embed": Pooler.for_embed(pooler_config),
+        })
 
         self.make_empty_intermediate_tensors = (
             self.model.make_empty_intermediate_tensors)

vllm/config.py

@@ -94,7 +94,7 @@
 TaskOption = Literal["auto", "generate", "embedding", "embed", "classify",
                      "score", "reward", "transcription", "draft"]
 
-_ResolvedTask = Literal["generate", "transcription", "pooling", "embed",
+_ResolvedTask = Literal["generate", "transcription", "encode", "embed",
                         "classify", "reward", "draft"]
 
 RunnerOption = Literal["auto", "generate", "pooling", "draft"]
@@ -103,7 +103,7 @@
 
 _RUNNER_TASKS: dict[RunnerType, list[_ResolvedTask]] = {
     "generate": ["generate", "transcription"],
-    "pooling": ["pooling", "embed", "classify", "reward"],
+    "pooling": ["encode", "embed", "classify", "reward"],
     "draft": [],
 }
 
@@ -579,7 +579,7 @@ def __post_init__(self) -> None:
         # user-selected task
         if runner_type == "pooling" and self.task == "auto":
             selected_task = all_supported_tasks[runner_type][-1]
-            assert selected_task != "pooling"
+            assert selected_task != "encode"
             self.task = selected_task
         self.supported_runner_types = supported_runner_types
         self.runner_type = runner_type
@@ -884,7 +884,7 @@ def _get_supported_pooling_tasks(
 
         supported_tasks = list[_ResolvedTask]()
         if registry.is_pooling_model(architectures):
-            supported_tasks.append("pooling")
+            supported_tasks.append("encode")
 
             # For now, users must specify the task (other than "pooling")
             # to use for pooling models

vllm/entrypoints/openai/api_server.py

@@ -1668,7 +1668,7 @@ async def init_app_state(
         request_logger=request_logger,
         chat_template=resolved_chat_template,
         chat_template_content_format=args.chat_template_content_format,
-    ) if "pooling" in model_config.supported_tasks else None
+    ) if "encode" in model_config.supported_tasks else None
     state.openai_serving_embedding = OpenAIServingEmbedding(
         engine_client,
         model_config,