A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 03:15 UTC
[Deprecate] Deprecate pooling multi task support. (#37956)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
wang.yuqi committed
Mar 24, 2026 at 22:07 UTC
1b6cb920e6ebcac57154e6154578c39d4892a16c
18 files changed
+566
-66
docs/models/pooling_models/README.md
+35
-10
@@ -1,7 +1,8 @@
1
# Pooling Models
2
3
!!! note
4
- We currently support pooling models primarily for convenience. This is not guaranteed to provide any performance improvements over using Hugging Face Transformers or Sentence Transformers directly.
4
+ We currently support pooling models primarily for convenience. This is not guaranteed to provide any performance
5
+improvements over using Hugging Face Transformers or Sentence Transformers directly.
6
7
We plan to optimize pooling models in vLLM. Please comment on <https://github.com/vllm-project/vllm/issues/21796> if you have any suggestions!
8
@@ -12,22 +13,38 @@ Natural Language Processing (NLP) can be primarily divided into the following tw
13
- Natural Language Understanding (NLU)
14
- Natural Language Generation (NLG)
15
15
-The generative models supported by vLLM cover a variety of task types, such as the large language models (LLMs) we are familiar with, multimodal models (VLM) that handle multimodal inputs like images, videos, and audio, speech-to-text transcription models, and real-time models that support streaming input. Their common feature is the ability to generate text. Taking it a step further, vLLM-Omni supports the generation of multimodal content, including images, videos, and audio.
16
+The generative models supported by vLLM cover a variety of task types, such as the large language models (LLMs) we are
17
+familiar with, multimodal models (VLM) that handle multimodal inputs like images, videos, and audio, speech-to-text
18
+transcription models, and real-time models that support streaming input. Their common feature is the ability to generate
19
+text. Taking it a step further, vLLM-Omni supports the generation of multimodal content, including images, videos, and audio.
20
17
-As the capabilities of generative models continue to improve, the boundaries of these models are also constantly expanding. However, certain application scenarios still require specialized small language models to efficiently complete specific tasks. These models typically have the following characteristics:
21
+As the capabilities of generative models continue to improve, the boundaries of these models are also constantly expanding.
22
+However, certain application scenarios still require specialized small language models to efficiently complete specific tasks.
23
+These models typically have the following characteristics:
24
25
- They do not require content generation.
26
- They only need to perform very limited functions, without requiring strong generalization, creativity, or high intelligence.
27
- They demand extremely low latency and may operate on cost-constrained hardware.
28
- Text-only models typically have fewer than 1 billion parameters, while multimodal models generally have fewer than 10 billion parameters.
29
24
-Although these models are relatively small in scale, they are still based on the Transformer architecture, similar or even identical to the most advanced large language models today. Many recently released pooling models are also fine-tuned from large language models, allowing them to benefit from the continuous improvements in large models. This architecture similarity enables them to reuse much of vLLM’s infrastructure. If compatible, we would be happy to help them leverage the latest features of vLLM as well.
30
+Although these models are relatively small in scale, they are still based on the Transformer architecture, similar or
31
+even identical to the most advanced large language models today. Many recently released pooling models are also fine-tuned
32
+from large language models, allowing them to benefit from the continuous improvements in large models. This architecture
33
+similarity enables them to reuse much of vLLM’s infrastructure. If compatible, we would be happy to help them leverage
34
+the latest features of vLLM as well.
35
36
### Sequence-wise Task and Token-wise Task
37
28
-The key distinction between sequence-wise task and token-wise task lies in their output granularity: sequence-wise task produces a single result for an entire input sequence, whereas token-wise task yields a result for each individual token within the sequence.
38
+The key distinction between sequence-wise task and token-wise task lies in their output granularity: sequence-wise task
39
+produces a single result for an entire input sequence, whereas token-wise task yields a result for each individual token
40
+within the sequence.
41
30
-Of course, we also have "plugin" tasks that allow users to customize input and output processors. For more information, please refer to [IO Processor Plugins](../../design/io_processor_plugins.md).
42
+Many Pooling models support both (sequence) task and token task. When the default pooling task (e.g. a sequence-wise task)
43
+is not what you want, you need to manually specify (e.g. a token-wise task) via `PoolerConfig(task=<task>)` offline or
44
+`--pooler-config.task <task>` online.
45
+
46
+Of course, we also have "plugin" tasks that allow users to customize input and output processors. For more information,
47
+please refer to [IO Processor Plugins](../../design/io_processor_plugins.md).
48
49
### Pooling Tasks
50
@@ -39,11 +56,13 @@ Of course, we also have "plugin" tasks that allow users to customize input and o
56
| `token_embed` | Token-wise | vector representations for each token |
57
58
!!! note
42
- Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
59
+ Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models
60
+are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
61
62
### Score Types
63
46
-The scoring models is designed to compute similarity scores between two input prompts. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`.
64
+The scoring models is designed to compute similarity scores between two input prompts. It supports three model types
65
+(aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`.
66
67
| Pooling Tasks | Granularity | Outputs | Score Types | scoring function |
68
|-----------------------|---------------|----------------------------------------------|--------------------|--------------------------|
@@ -250,11 +269,17 @@ We have split the `encode` task into two more specific token-wise tasks: `token_
269
- `token_embed` is the same as `embed`, using normalization as the activation.
270
- `token_classify` is the same as `classify`, by default using softmax as the activation.
271
253
-Pooling models now default support all pooling, you can use it without any settings.
272
+Pooling models now support token-wise task.
273
274
- Extracting hidden states prefers using `token_embed` task.
275
- Named Entity Recognition (NER) and reward models prefers using `token_classify` task.
276
277
### Score task
278
260
-`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
279
+`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
280
+classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
281
+
282
+### Pooling multitask support
283
+
284
+Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
285
+you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.
docs/models/pooling_models/token_classify.md
+6
@@ -13,6 +13,12 @@ The key distinction between (sequence) classification and token classification l
13
14
Many classification models support both (sequence) classification and token classification. For further details on (sequence) classification, please refer to [this page](classify.md).
15
16
+!!! note
17
+
18
+ Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task (classify) is not
19
+ what you want, you need to manually specify it via `PoolerConfig(task="token_classify")` offline or
20
+ `--pooler-config.task token_classify` online.
21
+
22
## Typical Use Cases
23
24
### Named Entity Recognition (NER)
docs/models/pooling_models/token_embed.md
+6
@@ -13,6 +13,12 @@ The difference between the (sequence) embedding task and the token embedding tas
13
14
Many embedding models support both (sequence) embedding and token embedding. For further details on (sequence) embedding, please refer to [this page](embed.md).
15
16
+!!! note
17
+
18
+ Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task (embed) is not
19
+ what you want, you need to manually specify it via via `PoolerConfig(task="token_embed")` offline or
20
+ `--pooler-config.task token_embed` online.
21
+
22
## Typical Use Cases
23
24
### Multi-Vector Retrieval
tests/entrypoints/pooling/classify/test_offline.py
+8
-5
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
+import logging
4
import weakref
5
6
import pytest
@@ -67,8 +67,11 @@ def test_list_prompts(llm: LLM):
67
68
69
@pytest.mark.skip_global_cleanup
70
-def test_token_classify(llm: LLM):
71
- outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
70
+def test_token_classify(llm: LLM, caplog_vllm):
71
+ with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
72
+ outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
73
+ assert "deprecated" in caplog_vllm.text
74
+
75
assert len(outputs) == 1
76
assert isinstance(outputs[0], PoolingRequestOutput)
77
assert outputs[0].prompt_token_ids == prompt_token_ids
@@ -107,8 +110,8 @@ def test_score_api(llm: LLM):
110
llm.score("ping", "pong", use_tqdm=False)
111
112
110
-@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
113
+@pytest.mark.parametrize("task", ["embed", "token_embed"])
114
def test_unsupported_tasks(llm: LLM, task: PoolingTask):
112
- err_msg = f"Unsupported task: '{task}' Supported tasks.+"
115
+ err_msg = "Embedding API is not supported by this model.+"
116
with pytest.raises(ValueError, match=err_msg):
117
llm.encode(prompt, pooling_task=task, use_tqdm=False)
tests/entrypoints/pooling/embed/test_offline.py
+48
-6
@@ -1,19 +1,22 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
+import logging
4
import weakref
5
6
import pytest
7
import torch
8
import torch.nn.functional as F
9
10
-from vllm import LLM, PoolingParams
10
+from vllm import LLM, EmbeddingRequestOutput, PoolingParams
11
from vllm.distributed import cleanup_dist_env_and_memory
12
from vllm.platforms import current_platform
13
+from vllm.tasks import PoolingTask
14
15
MODEL_NAME = "intfloat/multilingual-e5-small"
16
16
-prompts = ["The chef prepared a delicious meal."]
17
+prompt = "The chef prepared a delicious meal."
18
+prompt_token_ids = [0, 581, 21861, 133888, 10, 8, 150, 60744, 109911, 5, 2]
19
+embedding_size = 384
20
21
22
@pytest.fixture(scope="module")
@@ -44,16 +47,48 @@ def llm():
47
48
49
@pytest.mark.skip_global_cleanup
47
-def test_token_embed(llm: LLM):
48
- outputs = llm.encode(prompts, pooling_task="token_embed", use_tqdm=False)
50
+def test_str_prompts(llm: LLM):
51
+ outputs = llm.embed(prompt, use_tqdm=False)
52
+ assert len(outputs) == 1
53
+ assert isinstance(outputs[0], EmbeddingRequestOutput)
54
+ assert outputs[0].prompt_token_ids == prompt_token_ids
55
+ assert len(outputs[0].outputs.embedding) == embedding_size
56
+
57
+
58
+@pytest.mark.skip_global_cleanup
59
+def test_token_ids_prompts(llm: LLM):
60
+ outputs = llm.embed([prompt_token_ids], use_tqdm=False)
61
+ assert len(outputs) == 1
62
+ assert isinstance(outputs[0], EmbeddingRequestOutput)
63
+ assert outputs[0].prompt_token_ids == prompt_token_ids
64
+ assert len(outputs[0].outputs.embedding) == embedding_size
65
+
66
+
67
+@pytest.mark.skip_global_cleanup
68
+def test_list_prompts(llm: LLM):
69
+ outputs = llm.embed([prompt, prompt_token_ids], use_tqdm=False)
70
+ assert len(outputs) == 2
71
+ for i in range(len(outputs)):
72
+ assert isinstance(outputs[i], EmbeddingRequestOutput)
73
+ assert outputs[i].prompt_token_ids == prompt_token_ids
74
+ assert len(outputs[i].outputs.embedding) == embedding_size
75
+
76
+
77
+@pytest.mark.skip_global_cleanup
78
+def test_token_embed(llm: LLM, caplog_vllm):
79
+ with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
80
+ outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
81
+ assert "deprecated" in caplog_vllm.text
82
+
83
multi_vector = outputs[0].outputs.data
84
assert multi_vector.shape == (11, 384)
85
86
87
+@pytest.mark.skip_global_cleanup
88
def test_pooling_params(llm: LLM):
89
def get_outputs(normalize):
90
outputs = llm.embed(
56
- prompts,
91
+ [prompt],
92
pooling_params=PoolingParams(use_activation=normalize),
93
use_tqdm=False,
94
)
@@ -70,3 +105,10 @@ def test_pooling_params(llm: LLM):
105
assert torch.allclose(w_normal, F.normalize(wo_normal, p=2, dim=-1), atol=1e-2), (
106
"w_normal should be close to normal(wo_normal)."
107
)
108
+
109
+
110
+@pytest.mark.parametrize("task", ["token_classify", "classify"])
111
+def test_unsupported_tasks(llm: LLM, task: PoolingTask):
112
+ err_msg = "Classification API is not supported by this model.+"
113
+ with pytest.raises(ValueError, match=err_msg):
114
+ llm.encode(prompt, pooling_task=task, use_tqdm=False)
tests/entrypoints/pooling/score/test_online_rerank.py
+6
-1
@@ -206,7 +206,12 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
206
async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
207
response = requests.post(
208
server.url_for("pooling"),
209
- json={"model": model_name, "input": input_text, "encoding_format": "float"},
209
+ json={
210
+ "model": model_name,
211
+ "task": "token_classify",
212
+ "input": input_text,
213
+ "encoding_format": "float",
214
+ },
215
)
216
217
poolings = PoolingResponse.model_validate(response.json())
tests/entrypoints/pooling/token_classify/__init__.py
tests/entrypoints/pooling/token_classify/test_offline.py
new
+78
@@ -0,0 +1,78 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import logging
4
+import weakref
5
+
6
+import pytest
7
+
8
+from vllm import LLM, PoolingRequestOutput
9
+from vllm.config import PoolerConfig
10
+from vllm.distributed import cleanup_dist_env_and_memory
11
+from vllm.tasks import PoolingTask
12
+
13
+MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
14
+
15
+prompt = "The chef prepared a delicious meal."
16
+prompt_token_ids = [785, 29706, 10030, 264, 17923, 15145, 13]
17
+num_labels = 2
18
+
19
+
20
+@pytest.fixture(scope="module")
21
+def llm():
22
+ # pytest caches the fixture so we use weakref.proxy to
23
+ # enable garbage collection
24
+ llm = LLM(
25
+ model=MODEL_NAME,
26
+ pooler_config=PoolerConfig(task="token_classify"),
27
+ max_num_batched_tokens=32768,
28
+ tensor_parallel_size=1,
29
+ gpu_memory_utilization=0.75,
30
+ enforce_eager=True,
31
+ seed=0,
32
+ )
33
+
34
+ yield weakref.proxy(llm)
35
+
36
+ del llm
37
+
38
+ cleanup_dist_env_and_memory()
39
+
40
+
41
+@pytest.mark.skip_global_cleanup
42
+def test_str_prompts(llm: LLM):
43
+ outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
44
+ assert len(outputs) == 1
45
+ assert isinstance(outputs[0], PoolingRequestOutput)
46
+ assert outputs[0].prompt_token_ids == prompt_token_ids
47
+ assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
48
+
49
+
50
+@pytest.mark.skip_global_cleanup
51
+def test_token_ids_prompts(llm: LLM):
52
+ outputs = llm.encode(
53
+ [prompt_token_ids], pooling_task="token_classify", use_tqdm=False
54
+ )
55
+ assert len(outputs) == 1
56
+ assert isinstance(outputs[0], PoolingRequestOutput)
57
+ assert outputs[0].prompt_token_ids == prompt_token_ids
58
+ assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
59
+
60
+
61
+@pytest.mark.skip_global_cleanup
62
+def test_score_api(llm: LLM):
63
+ err_msg = "Score API is only enabled for num_labels == 1."
64
+ with pytest.raises(ValueError, match=err_msg):
65
+ llm.score("ping", "pong", use_tqdm=False)
66
+
67
+
68
+@pytest.mark.parametrize("task", ["classify", "embed", "token_embed"])
69
+def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
70
+ if task == "classify":
71
+ with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
72
+ llm.encode(prompt, pooling_task=task, use_tqdm=False)
73
+ assert "deprecated" in caplog_vllm.text
74
+ else:
75
+ err_msg = "Embedding API is not supported by this model.+"
76
+
77
+ with pytest.raises(ValueError, match=err_msg):
78
+ llm.encode(prompt, pooling_task=task, use_tqdm=False)
tests/entrypoints/pooling/token_classify/test_online.py
new
+70
@@ -0,0 +1,70 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import requests
6
+
7
+from tests.utils import RemoteOpenAIServer
8
+from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
9
+
10
+MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
11
+DTYPE = "float32" # Use float32 to avoid NaN issue
12
+input_text = "This product was excellent and exceeded my expectations"
13
+input_tokens = [1986, 1985, 572, 9073, 323, 33808, 847, 16665]
14
+
15
+
16
+@pytest.fixture(scope="module")
17
+def server():
18
+ args = [
19
+ "--enforce-eager",
20
+ "--max-model-len",
21
+ "512",
22
+ "--dtype",
23
+ DTYPE,
24
+ "--pooler-config.task",
25
+ "token_classify",
26
+ ]
27
+
28
+ with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
29
+ yield remote_server
30
+
31
+
32
+@pytest.mark.asyncio
33
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
34
+async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
35
+ task = "token_classify"
36
+ response = requests.post(
37
+ server.url_for("pooling"),
38
+ json={
39
+ "model": model_name,
40
+ "input": input_text,
41
+ "encoding_format": "float",
42
+ "task": task,
43
+ },
44
+ )
45
+ poolings = PoolingResponse.model_validate(response.json())
46
+ assert len(poolings.data) == 1
47
+ assert len(poolings.data[0].data) == 8
48
+ assert len(poolings.data[0].data[0]) == 2
49
+
50
+
51
+@pytest.mark.asyncio
52
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
53
+@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
54
+async def test_pooling_not_supported(
55
+ server: RemoteOpenAIServer, model_name: str, task: str
56
+):
57
+ response = requests.post(
58
+ server.url_for("pooling"),
59
+ json={
60
+ "model": model_name,
61
+ "input": input_text,
62
+ "encoding_format": "float",
63
+ "task": task,
64
+ },
65
+ )
66
+
67
+ if task != "classify":
68
+ assert response.json()["error"]["type"] == "BadRequestError"
69
+ err_msg = f"Unsupported task: {task!r}"
70
+ assert response.json()["error"]["message"].startswith(err_msg)
tests/entrypoints/pooling/token_embed/__init__.py
tests/entrypoints/pooling/token_embed/test_offline.py
new
+75
@@ -0,0 +1,75 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import logging
4
+import weakref
5
+
6
+import pytest
7
+
8
+from vllm import LLM, PoolingRequestOutput
9
+from vllm.config import PoolerConfig
10
+from vllm.distributed import cleanup_dist_env_and_memory
11
+from vllm.platforms import current_platform
12
+from vllm.tasks import PoolingTask
13
+
14
+MODEL_NAME = "intfloat/multilingual-e5-small"
15
+
16
+prompt = "The chef prepared a delicious meal."
17
+prompt_token_ids = [0, 581, 21861, 133888, 10, 8, 150, 60744, 109911, 5, 2]
18
+embedding_size = 384
19
+
20
+
21
+@pytest.fixture(scope="module")
22
+def llm():
23
+ # ROCm: Use FLEX_ATTENTION backend as it's the only attention backend
24
+ # that supports encoder-only models on ROCm.
25
+ attention_config = None
26
+ if current_platform.is_rocm():
27
+ attention_config = {"backend": "FLEX_ATTENTION"}
28
+
29
+ # pytest caches the fixture so we use weakref.proxy to
30
+ # enable garbage collection
31
+ llm = LLM(
32
+ model=MODEL_NAME,
33
+ pooler_config=PoolerConfig(task="token_embed"),
34
+ max_num_batched_tokens=32768,
35
+ tensor_parallel_size=1,
36
+ gpu_memory_utilization=0.75,
37
+ enforce_eager=True,
38
+ seed=0,
39
+ attention_config=attention_config,
40
+ )
41
+ assert embedding_size == llm.model_config.embedding_size
42
+
43
+ yield weakref.proxy(llm)
44
+
45
+ del llm
46
+ cleanup_dist_env_and_memory()
47
+
48
+
49
+@pytest.mark.skip_global_cleanup
50
+def test_str_prompts(llm: LLM):
51
+ outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
52
+ assert len(outputs) == 1
53
+ assert isinstance(outputs[0], PoolingRequestOutput)
54
+ assert outputs[0].outputs.data.shape == (11, 384)
55
+
56
+
57
+@pytest.mark.skip_global_cleanup
58
+def test_token_ids_prompts(llm: LLM):
59
+ outputs = llm.encode([prompt_token_ids], pooling_task="token_embed", use_tqdm=False)
60
+ assert len(outputs) == 1
61
+ assert isinstance(outputs[0], PoolingRequestOutput)
62
+ assert outputs[0].outputs.data.shape == (11, 384)
63
+
64
+
65
+@pytest.mark.parametrize("task", ["embed", "classify", "token_classify"])
66
+def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
67
+ if task == "embed":
68
+ with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
69
+ llm.encode(prompt, pooling_task=task, use_tqdm=False)
70
+ assert "deprecated" in caplog_vllm.text
71
+ else:
72
+ err_msg = "Classification API is not supported by this model.+"
73
+
74
+ with pytest.raises(ValueError, match=err_msg):
75
+ llm.encode(prompt, pooling_task=task, use_tqdm=False)
tests/entrypoints/pooling/token_embed/test_online.py
new
+93
@@ -0,0 +1,93 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+
5
+import pytest
6
+import requests
7
+
8
+from tests.utils import RemoteOpenAIServer
9
+from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
10
+
11
+MODEL_NAME = "intfloat/multilingual-e5-small"
12
+DTYPE = "bfloat16"
13
+input_text = "The best thing about vLLM is that it supports many different models"
14
+input_tokens = [
15
+ 0,
16
+ 581,
17
+ 2965,
18
+ 13580,
19
+ 1672,
20
+ 81,
21
+ 23708,
22
+ 594,
23
+ 83,
24
+ 450,
25
+ 442,
26
+ 8060,
27
+ 7,
28
+ 5941,
29
+ 12921,
30
+ 115774,
31
+ 2,
32
+]
33
+
34
+
35
+@pytest.fixture(scope="module")
36
+def server():
37
+ args = [
38
+ "--runner",
39
+ "pooling",
40
+ "--dtype",
41
+ DTYPE,
42
+ "--enforce-eager",
43
+ "--max-model-len",
44
+ "512",
45
+ "--pooler-config.task",
46
+ "token_embed",
47
+ ]
48
+
49
+ with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
50
+ yield remote_server
51
+
52
+
53
+@pytest.mark.asyncio
54
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
55
+async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
56
+ task = "token_embed"
57
+ response = requests.post(
58
+ server.url_for("pooling"),
59
+ json={
60
+ "model": model_name,
61
+ "input": input_text,
62
+ "encoding_format": "float",
63
+ "task": task,
64
+ },
65
+ )
66
+
67
+ poolings = PoolingResponse.model_validate(response.json())
68
+
69
+ assert len(poolings.data) == 1
70
+ assert len(poolings.data[0].data) == len(input_tokens)
71
+ assert len(poolings.data[0].data[0]) == 384
72
+
73
+
74
+@pytest.mark.asyncio
75
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
76
+@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
77
+async def test_pooling_not_supported(
78
+ server: RemoteOpenAIServer, model_name: str, task: str
79
+):
80
+ response = requests.post(
81
+ server.url_for("pooling"),
82
+ json={
83
+ "model": model_name,
84
+ "input": "test",
85
+ "encoding_format": "float",
86
+ "task": task,
87
+ },
88
+ )
89
+
90
+ if task != "embed":
91
+ assert response.json()["error"]["type"] == "BadRequestError"
92
+ err_msg = f"Unsupported task: {task!r}"
93
+ assert response.json()["error"]["message"].startswith(err_msg)
tests/plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
+3
-3
@@ -102,7 +102,7 @@ async def test_bge_m3_sparse_plugin_online(
102
"""Test BGE-M3 sparse plugin in online mode via API."""
103
request_payload = {
104
"model": model_config["model_name"],
105
- "task": "token_classify",
105
+ "task": "plugin",
106
"data": {"input": model_config["test_input"], "return_tokens": return_tokens},
107
}
108
@@ -166,7 +166,7 @@ def test_bge_m3_sparse_plugin_offline(vllm_runner, return_tokens: bool):
166
default_torch_num_threads=1,
167
) as llm_runner:
168
llm = llm_runner.get_llm()
169
- pooler_output = llm.encode(prompt, pooling_task="token_classify")
169
+ pooler_output = llm.encode(prompt, pooling_task="plugin")
170
171
outputs = pooler_output[0]
172
@@ -213,7 +213,7 @@ def test_bge_m3_sparse_plugin_offline_multiple_inputs(vllm_runner):
213
default_torch_num_threads=1,
214
) as llm_runner:
215
llm = llm_runner.get_llm()
216
- pooler_output = llm.encode(prompts, pooling_task="token_classify")
216
+ pooler_output = llm.encode(prompts, pooling_task="plugin")
217
218
outputs = pooler_output[0]
219
vllm/config/model.py
+36
-1
@@ -25,7 +25,7 @@ from vllm.config.scheduler import RunnerType
25
from vllm.config.utils import config, getattr_iter
26
from vllm.logger import init_logger
27
from vllm.platforms import current_platform
28
-from vllm.tasks import ScoreType
28
+from vllm.tasks import PoolingTask, ScoreType, SupportedTask
29
from vllm.transformers_utils.config import (
30
ConfigFormat,
31
get_config,
@@ -1409,6 +1409,41 @@ class ModelConfig: # type: ignore[misc]
1409
1410
return diff_sampling_param
1411
1412
+ def get_pooling_task(
1413
+ self, supported_tasks: tuple[SupportedTask, ...]
1414
+ ) -> PoolingTask | None:
1415
+ if self.pooler_config is None:
1416
+ return None
1417
+
1418
+ pooling_task = self.pooler_config.task
1419
+
1420
+ if pooling_task is not None:
1421
+ if self.pooler_config.task in supported_tasks:
1422
+ return self.pooler_config.task
1423
+ else:
1424
+ raise RuntimeError(
1425
+ f"Unsupported task: {pooling_task!r} "
1426
+ f"Supported tasks: {supported_tasks}"
1427
+ )
1428
+
1429
+ if "token_classify" in supported_tasks:
1430
+ for architecture in self.architectures:
1431
+ if "ForTokenClassification" in architecture:
1432
+ return "token_classify"
1433
+
1434
+ priority: list[PoolingTask] = [
1435
+ "embed&token_classify",
1436
+ "embed",
1437
+ "classify",
1438
+ "token_embed",
1439
+ "token_classify",
1440
+ "plugin",
1441
+ ]
1442
+ for task in priority:
1443
+ if task in supported_tasks:
1444
+ return task
1445
+ return None
1446
+
1447
@cached_property
1448
def is_encoder_decoder(self) -> bool:
1449
"""Extract the HF encoder/decoder model flag."""
vllm/config/pooler.py
+6
@@ -5,6 +5,7 @@ from typing import Any, Literal, get_args
5
6
from vllm.config.utils import config
7
from vllm.logger import init_logger
8
+from vllm.tasks import PoolingTask
9
from vllm.utils.hashing import safe_hash
10
11
logger = init_logger(__name__)
@@ -20,6 +21,11 @@ TOK_POOLING_TYPES: tuple[TokenPoolingType, ...] = get_args(TokenPoolingType)
21
class PoolerConfig:
22
"""Controls the behavior of output pooling in pooling models."""
23
24
+ task: PoolingTask | None = None
25
+ """
26
+ The task used for pooling.
27
+ """
28
+
29
pooling_type: SequencePoolingType | TokenPoolingType | None = None
30
"""
31
The pooling method used for pooling.
vllm/entrypoints/llm.py
+65
-37
@@ -382,16 +382,19 @@ class LLM:
382
self.llm_engine = LLMEngine.from_engine_args(
383
engine_args=engine_args, usage_context=UsageContext.LLM_CLASS
384
)
385
+ self.model_config = self.llm_engine.model_config
386
self.engine_class = type(self.llm_engine)
387
388
self.request_counter = Counter()
389
self.default_sampling_params: dict[str, Any] | None = None
390
391
supported_tasks = self.llm_engine.get_supported_tasks()
391
- logger.info("Supported tasks: %s", supported_tasks)
392
self.supported_tasks = supported_tasks
393
+ self.pooling_task = self.model_config.get_pooling_task(supported_tasks)
394
+ if self.pooling_task is not None:
395
+ logger.info("Supported pooling task: %s", self.pooling_task)
396
394
- self.model_config = self.llm_engine.model_config
397
+ self.runner_type = self.model_config.runner_type
398
self.renderer = self.llm_engine.renderer
399
self.chat_template = load_chat_template(chat_template)
400
self.io_processor = self.llm_engine.io_processor
@@ -1072,31 +1075,7 @@ class LLM:
1075
pooled hidden states in the same order as the input prompts.
1076
"""
1077
1075
- if pooling_task is None:
1076
- raise ValueError(
1077
- "pooling_task required for `LLM.encode`\n"
1078
- "Please use one of the more specific methods or set the "
1079
- "pooling_task when using `LLM.encode`:\n"
1080
- " - For embeddings, use `LLM.embed(...)` "
1081
- 'or `pooling_task="embed"`.\n'
1082
- " - For classification logits, use `LLM.classify(...)` "
1083
- 'or `pooling_task="classify"`.\n'
1084
- " - For similarity scores, use `LLM.score(...)`.\n"
1085
- " - For rewards, use `LLM.reward(...)` "
1086
- 'or `pooling_task="token_classify"`\n'
1087
- " - For token classification, "
1088
- 'use `pooling_task="token_classify"`\n'
1089
- ' - For multi-vector retrieval, use `pooling_task="token_embed"`'
1090
- )
1091
-
1092
- model_config = self.model_config
1093
- runner_type = model_config.runner_type
1094
- if runner_type != "pooling":
1095
- raise ValueError(
1096
- "LLM.encode() is only supported for pooling models. "
1097
- "Try passing `--runner pooling` to use the model as a "
1098
- "pooling model."
1099
- )
1078
+ self._verify_pooling_task(pooling_task)
1079
1080
if isinstance(prompts, dict) and "data" in prompts:
1081
if self.io_processor is None:
@@ -1206,6 +1185,65 @@ class LLM:
1185
)
1186
return outputs
1187
1188
+ def _verify_pooling_task(self, pooling_task: PoolingTask | None):
1189
+ if self.runner_type != "pooling":
1190
+ raise ValueError(
1191
+ "LLM.encode() is only supported for pooling models. "
1192
+ "Try passing `--runner pooling` to use the model as a "
1193
+ "pooling model."
1194
+ )
1195
+
1196
+ if pooling_task is None:
1197
+ raise ValueError(
1198
+ "pooling_task required for `LLM.encode`\n"
1199
+ "Please use one of the more specific methods or set the "
1200
+ "pooling_task when using `LLM.encode`:\n"
1201
+ " - For embeddings, use `LLM.embed(...)` "
1202
+ 'or `pooling_task="embed"`.\n'
1203
+ " - For classification logits, use `LLM.classify(...)` "
1204
+ 'or `pooling_task="classify"`.\n'
1205
+ " - For similarity scores, use `LLM.score(...)`.\n"
1206
+ " - For rewards, use `LLM.reward(...)` "
1207
+ 'or `pooling_task="token_classify"`\n'
1208
+ " - For token classification, "
1209
+ 'use `pooling_task="token_classify"`\n'
1210
+ ' - For multi-vector retrieval, use `pooling_task="token_embed"`'
1211
+ )
1212
+
1213
+ if (
1214
+ pooling_task in ("embed", "token_embed")
1215
+ and pooling_task not in self.supported_tasks
1216
+ ):
1217
+ raise ValueError(
1218
+ "Embedding API is not supported by this model. "
1219
+ "Try converting the model using `--convert embed`."
1220
+ )
1221
+
1222
+ if (
1223
+ pooling_task in ("classify", "token_classify")
1224
+ and pooling_task not in self.supported_tasks
1225
+ ):
1226
+ raise ValueError(
1227
+ "Classification API is not supported by this model. "
1228
+ "Try converting the model using `--convert classify`."
1229
+ )
1230
+
1231
+ # plugin task uses io_processor.parse_request to verify inputs
1232
+ if pooling_task != "plugin" and pooling_task != self.pooling_task:
1233
+ if pooling_task not in self.supported_tasks:
1234
+ raise ValueError(
1235
+ f"Unsupported task: {pooling_task!r} "
1236
+ f"Supported tasks: {self.supported_tasks}"
1237
+ )
1238
+ else:
1239
+ logger.warning_once(
1240
+ "Pooling multitask support is deprecated and will "
1241
+ "be removed in v0.20. When the default pooling task is "
1242
+ "not what you want, you need to manually specify it "
1243
+ 'via PoolerConfig(task="%s"). ',
1244
+ pooling_task,
1245
+ )
1246
+
1247
def embed(
1248
self,
1249
prompts: PromptType | Sequence[PromptType],
@@ -1239,11 +1277,6 @@ class LLM:
1277
A list of `EmbeddingRequestOutput` objects containing the
1278
embedding vectors in the same order as the input prompts.
1279
"""
1242
- if "embed" not in self.supported_tasks:
1243
- raise ValueError(
1244
- "Embedding API is not supported by this model. "
1245
- "Try converting the model using `--convert embed`."
1246
- )
1280
1281
items = self.encode(
1282
prompts,
@@ -1289,11 +1322,6 @@ class LLM:
1322
A list of `ClassificationRequestOutput` objects containing the
1323
embedding vectors in the same order as the input prompts.
1324
"""
1292
- if "classify" not in self.supported_tasks:
1293
- raise ValueError(
1294
- "Classification API is not supported by this model. "
1295
- "Try converting the model using `--convert classify`."
1296
- )
1325
1326
items = self.encode(
1327
prompts,
vllm/entrypoints/pooling/__init__.py
+9
-2
@@ -45,9 +45,15 @@ def register_pooling_api_routers(
45
supported_tasks: tuple["SupportedTask", ...],
46
model_config: ModelConfig | None = None,
47
):
48
- from vllm.entrypoints.pooling.pooling.api_router import router as pooling_router
48
+ if model_config is None:
49
+ return
50
50
- app.include_router(pooling_router)
51
+ pooling_task = model_config.get_pooling_task(supported_tasks)
52
+
53
+ if pooling_task is not None:
54
+ from vllm.entrypoints.pooling.pooling.api_router import router as pooling_router
55
+
56
+ app.include_router(pooling_router)
57
58
if "classify" in supported_tasks:
59
from vllm.entrypoints.pooling.classify.api_router import (
@@ -91,6 +97,7 @@ def init_pooling_state(
97
engine_client,
98
state.openai_serving_models,
99
state.openai_serving_render,
100
+ supported_tasks=supported_tasks,
101
request_logger=request_logger,
102
chat_template=resolved_chat_template,
103
chat_template_content_format=args.chat_template_content_format,
vllm/entrypoints/pooling/pooling/serving.py
+22
-1
@@ -37,6 +37,7 @@ from vllm.inputs import ProcessorInputs
37
from vllm.logger import init_logger
38
from vllm.outputs import PoolingRequestOutput
39
from vllm.renderers.inputs.preprocess import prompt_to_seq
40
+from vllm.tasks import SupportedTask
41
from vllm.utils.async_utils import merge_async_iterators
42
from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness
43
@@ -49,6 +50,7 @@ class OpenAIServingPooling(OpenAIServing):
50
engine_client: EngineClient,
51
models: OpenAIServingModels,
52
openai_serving_render: OpenAIServingRender,
53
+ supported_tasks: tuple[SupportedTask, ...],
54
*,
55
request_logger: RequestLogger | None,
56
chat_template: str | None,
@@ -60,7 +62,8 @@ class OpenAIServingPooling(OpenAIServing):
62
models=models,
63
request_logger=request_logger,
64
)
63
-
65
+ self.supported_tasks = supported_tasks
66
+ self.pooling_task = self.model_config.get_pooling_task(supported_tasks)
67
self.openai_serving_render = openai_serving_render
68
self.chat_template = chat_template
69
self.chat_template_content_format: Final = chat_template_content_format
@@ -86,9 +89,27 @@ class OpenAIServingPooling(OpenAIServing):
89
90
lora_request = self._maybe_get_adapters(request)
91
92
+ if request.task is None:
93
+ request.task = self.pooling_task
94
+
95
if getattr(request, "dimensions", None) is not None:
96
return self.create_error_response("dimensions is currently not supported")
97
98
+ # plugin task uses io_processor.parse_request to verify inputs
99
+ if request.task != "plugin" and request.task != self.pooling_task:
100
+ if request.task not in self.supported_tasks:
101
+ raise ValueError(
102
+ f"Unsupported task: {request.task!r} "
103
+ f"Supported tasks: {self.supported_tasks}"
104
+ )
105
+ else:
106
+ logger.warning_once(
107
+ "Pooling multitask support is deprecated and will be removed "
108
+ "in v0.20. When the default pooling task is not what you want, you "
109
+ 'need to manually specify it via --pooler-config.task "%s". ',
110
+ request.task,
111
+ )
112
+
113
engine_prompts: Sequence[ProcessorInputs]
114
if use_io_processor := isinstance(request, IOProcessorRequest):
115
if self.io_processor is None: