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 8, 21:30 UTC
[Refactor] Move `task` outside of `PoolingParams.verify` (#33796)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
Cyrus Leung committed
Feb 5, 2026 at 17:33 UTC
038914b7c891c0b5b2853ec0574062dc3bea8073
24 files changed
+186
-216
.buildkite/test-amd.yaml
+2
@@ -70,6 +70,7 @@ steps:
70
- vllm/
71
- tests/test_inputs.py
72
- tests/test_outputs.py
73
+ - tests/test_pooling_params.py
74
- tests/multimodal
75
- tests/renderers
76
- tests/standalone_tests/lazy_imports.py
@@ -82,6 +83,7 @@ steps:
83
- python3 standalone_tests/lazy_imports.py
84
- pytest -v -s test_inputs.py
85
- pytest -v -s test_outputs.py
86
+ - pytest -v -s test_pooling_params.py
87
- pytest -v -s -m 'cpu_test' multimodal
88
- pytest -v -s renderers
89
- pytest -v -s tokenizers_
.buildkite/test-pipeline.yaml
+2
@@ -63,6 +63,7 @@ steps:
63
- vllm/
64
- tests/test_inputs.py
65
- tests/test_outputs.py
66
+ - tests/test_pooling_params.py
67
- tests/multimodal
68
- tests/renderers
69
- tests/standalone_tests/lazy_imports.py
@@ -75,6 +76,7 @@ steps:
76
- python3 standalone_tests/lazy_imports.py
77
- pytest -v -s test_inputs.py
78
- pytest -v -s test_outputs.py
79
+ - pytest -v -s test_pooling_params.py
80
- pytest -v -s -m 'cpu_test' multimodal
81
- pytest -v -s renderers
82
- pytest -v -s tokenizers_
.buildkite/test_areas/misc.yaml
+2
@@ -122,6 +122,7 @@ steps:
122
- vllm/
123
- tests/test_inputs.py
124
- tests/test_outputs.py
125
+ - tests/test_pooling_params.py
126
- tests/multimodal
127
- tests/renderers
128
- tests/standalone_tests/lazy_imports.py
@@ -134,6 +135,7 @@ steps:
135
- python3 standalone_tests/lazy_imports.py
136
- pytest -v -s test_inputs.py
137
- pytest -v -s test_outputs.py
138
+ - pytest -v -s test_pooling_params.py
139
- pytest -v -s -m 'cpu_test' multimodal
140
- pytest -v -s renderers
141
- pytest -v -s tokenizers_
tests/entrypoints/pooling/classify/test_online.py
+1
-3
@@ -469,6 +469,4 @@ async def test_pooling_not_supported(
469
},
470
)
471
assert response.json()["error"]["type"] == "BadRequestError"
472
- assert response.json()["error"]["message"].startswith(
473
- f"Task {task} is not supported"
474
- )
472
+ assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")
tests/entrypoints/pooling/embed/test_online.py
+1
-3
@@ -757,6 +757,4 @@ async def test_pooling_not_supported(
757
},
758
)
759
assert response.json()["error"]["type"] == "BadRequestError"
760
- assert response.json()["error"]["message"].startswith(
761
- f"Task {task} is not supported"
762
- )
760
+ assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")
tests/entrypoints/pooling/score/test_online_colbert.py
+5
-5
@@ -138,17 +138,17 @@ def test_colbert_token_embed(server: RemoteOpenAIServer, model_name: str):
138
@pytest.mark.parametrize("model_name", [MODEL_NAME])
139
def test_colbert_embed_not_supported(server: RemoteOpenAIServer, model_name: str):
140
"""Test that ColBERT model does not support 'embed' task."""
141
+ task = "embed"
142
text = "What is the capital of France?"
143
143
- pooling_response = requests.post(
144
+ response = requests.post(
145
server.url_for("pooling"),
146
json={
147
"model": model_name,
148
"input": text,
148
- "task": "embed",
149
+ "task": task,
150
},
151
)
152
152
- # Should return error
153
- assert pooling_response.status_code == 400
154
- assert "Task embed is not supported" in pooling_response.text
153
+ assert response.json()["error"]["type"] == "BadRequestError"
154
+ assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")
tests/entrypoints/pooling/score/test_online_rerank.py
+1
-3
@@ -232,6 +232,4 @@ async def test_pooling_not_supported(
232
},
233
)
234
assert response.json()["error"]["type"] == "BadRequestError"
235
- assert response.json()["error"]["message"].startswith(
236
- f"Task {task} is not supported"
237
- )
235
+ assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")
tests/test_pooling_params.py
+40
-50
@@ -27,35 +27,24 @@ class MockModelConfig:
27
pooler_config: PoolerConfig
28
29
30
-def test_task():
31
- pooling_params = PoolingParams()
32
- pooling_params.verify(task="score")
33
-
34
- pooling_params = PoolingParams(task="score")
35
- pooling_params.verify(task="score")
36
-
37
- with pytest.raises(ValueError):
38
- pooling_params.verify(task="classify")
39
-
40
-
30
def test_embed():
31
task = "embed"
32
model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
33
45
- pooling_params = PoolingParams(use_activation=None)
46
- pooling_params.verify(task=task, model_config=model_config)
34
+ pooling_params = PoolingParams(task=task, use_activation=None)
35
+ pooling_params.verify(model_config)
36
48
- pooling_params = PoolingParams(use_activation=True)
49
- pooling_params.verify(task=task, model_config=model_config)
37
+ pooling_params = PoolingParams(task=task, use_activation=True)
38
+ pooling_params.verify(model_config)
39
51
- pooling_params = PoolingParams(use_activation=False)
52
- pooling_params.verify(task=task, model_config=model_config)
40
+ pooling_params = PoolingParams(task=task, use_activation=False)
41
+ pooling_params.verify(model_config)
42
43
invalid_parameters = classify_parameters + step_pooling_parameters
44
for p in set(invalid_parameters) - set(embed_parameters):
45
with pytest.raises(ValueError):
57
- pooling_params = PoolingParams(**{p: True})
58
- pooling_params.verify(task=task, model_config=model_config)
46
+ pooling_params = PoolingParams(task=task, **{p: True})
47
+ pooling_params.verify(model_config)
48
49
50
@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
@@ -63,7 +52,6 @@ def test_embed_dimensions(model_info: EmbedModelInfo):
52
task = "embed"
53
model_config = ModelConfig(
54
model_info.name,
66
- task="auto",
55
tokenizer=model_info.name,
56
tokenizer_mode="auto",
57
trust_remote_code=False,
@@ -71,37 +59,39 @@ def test_embed_dimensions(model_info: EmbedModelInfo):
59
dtype="float16",
60
)
61
74
- pooling_params = PoolingParams(dimensions=None)
75
- pooling_params.verify(task=task, model_config=model_config)
62
+ pooling_params = PoolingParams(task=task, dimensions=None)
63
+ pooling_params.verify(model_config)
64
65
with pytest.raises(ValueError):
78
- pooling_params = PoolingParams(dimensions=1)
79
- pooling_params.verify(task=task, model_config=model_config)
66
+ pooling_params = PoolingParams(task=task, dimensions=1)
67
+ pooling_params.verify(model_config)
68
69
if model_info.is_matryoshka:
70
assert model_info.matryoshka_dimensions is not None
83
- pooling_params = PoolingParams(dimensions=model_info.matryoshka_dimensions[0])
84
- pooling_params.verify(task=task, model_config=model_config)
71
+ pooling_params = PoolingParams(
72
+ task=task, dimensions=model_info.matryoshka_dimensions[0]
73
+ )
74
+ pooling_params.verify(model_config)
75
76
77
@pytest.mark.parametrize("task", ["score", "classify"])
78
def test_classify(task):
79
model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
80
91
- pooling_params = PoolingParams(use_activation=None)
92
- pooling_params.verify(task=task, model_config=model_config)
81
+ pooling_params = PoolingParams(task=task, use_activation=None)
82
+ pooling_params.verify(model_config)
83
94
- pooling_params = PoolingParams(use_activation=True)
95
- pooling_params.verify(task=task, model_config=model_config)
84
+ pooling_params = PoolingParams(task=task, use_activation=True)
85
+ pooling_params.verify(model_config)
86
97
- pooling_params = PoolingParams(use_activation=False)
98
- pooling_params.verify(task=task, model_config=model_config)
87
+ pooling_params = PoolingParams(task=task, use_activation=False)
88
+ pooling_params.verify(model_config)
89
90
invalid_parameters = embed_parameters + step_pooling_parameters
91
for p in set(invalid_parameters) - set(classify_parameters):
92
with pytest.raises(ValueError):
103
- pooling_params = PoolingParams(**{p: True})
104
- pooling_params.verify(task=task, model_config=model_config)
93
+ pooling_params = PoolingParams(task=task, **{p: True})
94
+ pooling_params.verify(model_config)
95
96
97
@pytest.mark.parametrize("pooling_type", ["ALL", "STEP"])
@@ -111,14 +101,14 @@ def test_token_embed(pooling_type: str):
101
pooler_config=PoolerConfig(tok_pooling_type=pooling_type)
102
)
103
114
- pooling_params = PoolingParams(use_activation=None)
115
- pooling_params.verify(task=task, model_config=model_config)
104
+ pooling_params = PoolingParams(task=task, use_activation=None)
105
+ pooling_params.verify(model_config)
106
117
- pooling_params = PoolingParams(use_activation=True)
118
- pooling_params.verify(task=task, model_config=model_config)
107
+ pooling_params = PoolingParams(task=task, use_activation=True)
108
+ pooling_params.verify(model_config)
109
120
- pooling_params = PoolingParams(use_activation=False)
121
- pooling_params.verify(task=task, model_config=model_config)
110
+ pooling_params = PoolingParams(task=task, use_activation=False)
111
+ pooling_params.verify(model_config)
112
113
invalid_parameters = classify_parameters
114
if pooling_type != "STEP":
@@ -126,8 +116,8 @@ def test_token_embed(pooling_type: str):
116
117
for p in set(invalid_parameters) - set(embed_parameters):
118
with pytest.raises(ValueError):
129
- pooling_params = PoolingParams(**{p: True})
130
- pooling_params.verify(task=task, model_config=model_config)
119
+ pooling_params = PoolingParams(task=task, **{p: True})
120
+ pooling_params.verify(model_config)
121
122
123
@pytest.mark.parametrize("pooling_type", ["ALL", "STEP"])
@@ -137,14 +127,14 @@ def test_token_classify(pooling_type: str):
127
pooler_config=PoolerConfig(tok_pooling_type=pooling_type)
128
)
129
140
- pooling_params = PoolingParams(use_activation=None)
141
- pooling_params.verify(task=task, model_config=model_config)
130
+ pooling_params = PoolingParams(task=task, use_activation=None)
131
+ pooling_params.verify(model_config)
132
143
- pooling_params = PoolingParams(use_activation=True)
144
- pooling_params.verify(task=task, model_config=model_config)
133
+ pooling_params = PoolingParams(task=task, use_activation=True)
134
+ pooling_params.verify(model_config)
135
146
- pooling_params = PoolingParams(use_activation=False)
147
- pooling_params.verify(task=task, model_config=model_config)
136
+ pooling_params = PoolingParams(task=task, use_activation=False)
137
+ pooling_params.verify(model_config)
138
139
invalid_parameters = embed_parameters
140
if pooling_type != "STEP":
@@ -152,5 +142,5 @@ def test_token_classify(pooling_type: str):
142
143
for p in set(invalid_parameters) - set(classify_parameters):
144
with pytest.raises(ValueError):
155
- pooling_params = PoolingParams(**{p: True})
156
- pooling_params.verify(task=task, model_config=model_config)
145
+ pooling_params = PoolingParams(task=task, **{p: True})
146
+ pooling_params.verify(model_config)
vllm/entrypoints/llm.py
+8
-5
@@ -1135,11 +1135,12 @@ class LLM:
1135
# Use default pooling params.
1136
pooling_params = PoolingParams()
1137
1138
- if pooling_task not in self.supported_tasks:
1139
- raise ValueError(f"pooling_task must be one of {self.supported_tasks}.")
1140
-
1138
for param in as_iter(pooling_params):
1142
- param.verify(pooling_task, model_config)
1139
+ if param.task is None:
1140
+ param.task = pooling_task
1141
+ elif param.task != pooling_task:
1142
+ msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
1143
+ raise ValueError(msg)
1144
1145
self._validate_and_add_requests(
1146
prompts=prompts,
@@ -1472,8 +1473,9 @@ class LLM:
1473
1474
if pooling_params is None:
1475
pooling_params = PoolingParams(task="score")
1476
+ elif pooling_params.task is None:
1477
+ pooling_params.task = "score"
1478
1476
- pooling_params.verify("score", model_config)
1479
pooling_params_list = list[PoolingParams]()
1480
1481
prompts = list[PromptType]()
@@ -1836,6 +1838,7 @@ class LLM:
1838
lora_request=lora_request,
1839
tokenization_kwargs=tokenization_kwargs,
1840
priority=priority,
1841
+ supported_tasks=self.supported_tasks,
1842
)
1843
1844
self.llm_engine.add_request(
vllm/entrypoints/pooling/__init__.py
+1
-2
@@ -68,7 +68,6 @@ def init_pooling_state(
68
OpenAIServingPooling(
69
engine_client,
70
state.openai_serving_models,
71
- supported_tasks=supported_tasks,
71
request_logger=request_logger,
72
chat_template=resolved_chat_template,
73
chat_template_content_format=args.chat_template_content_format,
@@ -76,7 +75,7 @@ def init_pooling_state(
75
log_error_stack=args.log_error_stack,
76
)
77
)
79
- if any(task in POOLING_TASKS for task in supported_tasks)
78
+ if any(t in supported_tasks for t in POOLING_TASKS)
79
else None
80
)
81
state.openai_serving_embedding = (
vllm/entrypoints/pooling/base/protocol.py
-24
@@ -6,19 +6,15 @@ from typing import Annotated, Any
6
7
from pydantic import Field, model_validator
8
9
-from vllm import PoolingParams
9
from vllm.entrypoints.chat_utils import (
10
ChatCompletionMessageParam,
11
ChatTemplateContentFormatOption,
12
)
13
from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel
15
-from vllm.logger import init_logger
14
from vllm.renderers import ChatParams, merge_kwargs
15
from vllm.utils import random_uuid
16
from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness
17
20
-logger = init_logger(__name__)
21
-
18
19
class PoolingBasicRequestMixin(OpenAIBaseModel):
20
# --8<-- [start:pooling-common-params]
@@ -185,20 +181,6 @@ class EmbedRequestMixin(EncodingRequestMixin):
181
)
182
# --8<-- [end:embed-extra-params]
183
188
- def to_pooling_params(self):
189
- if self.normalize is not None:
190
- logger.warning_once(
191
- "`normalize` is deprecated and will be removed in v0.17. "
192
- "Please pass `use_activation` instead."
193
- )
194
- self.use_activation = self.normalize
195
-
196
- return PoolingParams(
197
- dimensions=self.dimensions,
198
- use_activation=self.use_activation,
199
- truncate_prompt_tokens=getattr(self, "truncate_prompt_tokens", None),
200
- )
201
-
184
185
class ClassifyRequestMixin(OpenAIBaseModel):
186
# --8<-- [start:classify-extra-params]
@@ -208,9 +190,3 @@ class ClassifyRequestMixin(OpenAIBaseModel):
190
"`None` uses the pooler's default, which is `True` in most cases.",
191
)
192
# --8<-- [end:classify-extra-params]
211
-
212
- def to_pooling_params(self):
213
- return PoolingParams(
214
- use_activation=self.use_activation,
215
- truncate_prompt_tokens=getattr(self, "truncate_prompt_tokens", None),
216
- )
vllm/entrypoints/pooling/classify/protocol.py
+18
@@ -6,6 +6,7 @@ from typing import Any, TypeAlias
6
7
from pydantic import Field
8
9
+from vllm import PoolingParams
10
from vllm.config import ModelConfig
11
from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
12
from vllm.entrypoints.pooling.base.protocol import (
@@ -14,9 +15,12 @@ from vllm.entrypoints.pooling.base.protocol import (
15
CompletionRequestMixin,
16
PoolingBasicRequestMixin,
17
)
18
+from vllm.logger import init_logger
19
from vllm.renderers import TokenizeParams
20
from vllm.utils import random_uuid
21
22
+logger = init_logger(__name__)
23
+
24
25
class ClassificationCompletionRequest(
26
PoolingBasicRequestMixin, CompletionRequestMixin, ClassifyRequestMixin
@@ -33,6 +37,13 @@ class ClassificationCompletionRequest(
37
max_total_tokens_param="max_model_len",
38
)
39
40
+ def to_pooling_params(self):
41
+ return PoolingParams(
42
+ task="classify",
43
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
44
+ use_activation=self.use_activation,
45
+ )
46
+
47
48
class ClassificationChatRequest(
49
PoolingBasicRequestMixin, ChatRequestMixin, ClassifyRequestMixin
@@ -55,6 +66,13 @@ class ClassificationChatRequest(
66
max_total_tokens_param="max_model_len",
67
)
68
69
+ def to_pooling_params(self):
70
+ return PoolingParams(
71
+ task="classify",
72
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
73
+ use_activation=self.use_activation,
74
+ )
75
+
76
77
ClassificationRequest: TypeAlias = (
78
ClassificationCompletionRequest | ClassificationChatRequest
vllm/entrypoints/pooling/classify/serving.py
-16
@@ -22,7 +22,6 @@ from vllm.entrypoints.pooling.classify.protocol import (
22
)
23
from vllm.logger import init_logger
24
from vllm.outputs import ClassificationOutput
25
-from vllm.pooling_params import PoolingParams
25
26
logger = init_logger(__name__)
27
@@ -159,18 +158,3 @@ class ServingClassification(OpenAIServing):
158
)
159
160
return await self.handle(ctx) # type: ignore[return-value]
162
-
163
- def _create_pooling_params(
164
- self,
165
- ctx: ClassificationServeContext,
166
- ) -> PoolingParams | ErrorResponse:
167
- pooling_params = super()._create_pooling_params(ctx)
168
- if isinstance(pooling_params, ErrorResponse):
169
- return pooling_params
170
-
171
- try:
172
- pooling_params.verify("classify", self.model_config)
173
- except ValueError as e:
174
- return self.create_error_response(str(e))
175
-
176
- return pooling_params
vllm/entrypoints/pooling/embed/protocol.py
+34
@@ -5,6 +5,7 @@ from typing import Any, TypeAlias
5
6
from pydantic import Field
7
8
+from vllm import PoolingParams
9
from vllm.config import ModelConfig
10
from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
11
from vllm.entrypoints.pooling.base.protocol import (
@@ -13,9 +14,12 @@ from vllm.entrypoints.pooling.base.protocol import (
14
EmbedRequestMixin,
15
PoolingBasicRequestMixin,
16
)
17
+from vllm.logger import init_logger
18
from vllm.renderers import TokenizeParams
19
from vllm.utils import random_uuid
20
21
+logger = init_logger(__name__)
22
+
23
24
def _get_max_total_output_tokens(
25
model_config: ModelConfig,
@@ -55,6 +59,21 @@ class EmbeddingCompletionRequest(
59
max_output_tokens_param="max_model_len - max_embed_len",
60
)
61
62
+ def to_pooling_params(self):
63
+ if self.normalize is not None:
64
+ logger.warning_once(
65
+ "`normalize` is deprecated and will be removed in v0.17. "
66
+ "Please pass `use_activation` instead."
67
+ )
68
+ self.use_activation = self.normalize
69
+
70
+ return PoolingParams(
71
+ task="embed",
72
+ dimensions=self.dimensions,
73
+ use_activation=self.use_activation,
74
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
75
+ )
76
+
77
78
class EmbeddingChatRequest(
79
PoolingBasicRequestMixin, ChatRequestMixin, EmbedRequestMixin
@@ -82,6 +101,21 @@ class EmbeddingChatRequest(
101
max_output_tokens_param="max_model_len - max_embed_len",
102
)
103
104
+ def to_pooling_params(self):
105
+ if self.normalize is not None:
106
+ logger.warning_once(
107
+ "`normalize` is deprecated and will be removed in v0.17. "
108
+ "Please pass `use_activation` instead."
109
+ )
110
+ self.use_activation = self.normalize
111
+
112
+ return PoolingParams(
113
+ task="embed",
114
+ dimensions=self.dimensions,
115
+ use_activation=self.use_activation,
116
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
117
+ )
118
+
119
120
EmbeddingRequest: TypeAlias = EmbeddingCompletionRequest | EmbeddingChatRequest
121
vllm/entrypoints/pooling/embed/serving.py
+2
-24
@@ -424,12 +424,6 @@ class OpenAIServingEmbedding(OpenAIServing):
424
if isinstance(pooling_params, ErrorResponse):
425
return pooling_params
426
427
- # Verify and set the task for pooling params
428
- try:
429
- pooling_params.verify("embed", self.model_config)
430
- except ValueError as e:
431
- return self.create_error_response(str(e))
432
-
427
if ctx.engine_prompts is None:
428
return self.create_error_response("Engine prompts not available")
429
@@ -463,8 +457,7 @@ class OpenAIServingEmbedding(OpenAIServing):
457
return None
458
459
except Exception as e:
466
- # TODO: Use a vllm-specific Validation Error
467
- return self.create_error_response(str(e))
460
+ return self.create_error_response(e)
461
462
async def _collect_batch(
463
self,
@@ -634,7 +627,7 @@ class OpenAIServingEmbedding(OpenAIServing):
627
return None
628
629
except Exception as e:
637
- return self.create_error_response(str(e))
630
+ return self.create_error_response(e)
631
632
async def create_embedding(
633
self,
@@ -661,18 +654,3 @@ class OpenAIServingEmbedding(OpenAIServing):
654
)
655
656
return await self.handle(ctx) # type: ignore[return-value]
664
-
665
- def _create_pooling_params(
666
- self,
667
- ctx: EmbeddingServeContext,
668
- ) -> PoolingParams | ErrorResponse:
669
- pooling_params = super()._create_pooling_params(ctx)
670
- if isinstance(pooling_params, ErrorResponse):
671
- return pooling_params
672
-
673
- try:
674
- pooling_params.verify("embed", self.model_config)
675
- except ValueError as e:
676
- return self.create_error_response(str(e))
677
-
678
- return pooling_params
vllm/entrypoints/pooling/pooling/protocol.py
+3
-1
@@ -53,6 +53,7 @@ class PoolingCompletionRequest(
53
self.use_activation = self.normalize
54
55
return PoolingParams(
56
+ task=self.task,
57
truncate_prompt_tokens=self.truncate_prompt_tokens,
58
use_activation=self.use_activation,
59
dimensions=self.dimensions,
@@ -90,6 +91,7 @@ class PoolingChatRequest(
91
self.use_activation = self.normalize
92
93
return PoolingParams(
94
+ task=self.task,
95
truncate_prompt_tokens=self.truncate_prompt_tokens,
96
use_activation=self.use_activation,
97
dimensions=self.dimensions,
@@ -104,7 +106,7 @@ class IOProcessorRequest(PoolingBasicRequestMixin, EncodingRequestMixin, Generic
106
task: PoolingTask = "plugin"
107
108
def to_pooling_params(self):
107
- return PoolingParams()
109
+ return PoolingParams(task=self.task)
110
111
112
class IOProcessorResponse(OpenAIBaseModel, Generic[T]):
vllm/entrypoints/pooling/pooling/serving.py
+2
-33
@@ -35,7 +35,6 @@ from vllm.entrypoints.pooling.utils import (
35
)
36
from vllm.logger import init_logger
37
from vllm.outputs import PoolingRequestOutput
38
-from vllm.tasks import PoolingTask, SupportedTask
38
from vllm.utils.async_utils import merge_async_iterators
39
from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness
40
@@ -48,7 +47,6 @@ class OpenAIServingPooling(OpenAIServing):
47
engine_client: EngineClient,
48
models: OpenAIServingModels,
49
*,
51
- supported_tasks: tuple[SupportedTask, ...],
50
request_logger: RequestLogger | None,
51
chat_template: str | None,
52
chat_template_content_format: ChatTemplateContentFormatOption,
@@ -62,7 +60,6 @@ class OpenAIServingPooling(OpenAIServing):
60
log_error_stack=log_error_stack,
61
)
62
65
- self.supported_tasks = supported_tasks
63
self.chat_template = chat_template
64
self.chat_template_content_format: Final = chat_template_content_format
65
self.trust_request_chat_template = trust_request_chat_template
@@ -152,32 +149,6 @@ class OpenAIServingPooling(OpenAIServing):
149
else:
150
pooling_params = request.to_pooling_params()
151
155
- pooling_task: PoolingTask
156
- if request.task is None:
157
- if "token_embed" in self.supported_tasks:
158
- pooling_task = "token_embed"
159
- elif "token_classify" in self.supported_tasks:
160
- pooling_task = "token_classify"
161
- elif "plugin" in self.supported_tasks:
162
- pooling_task = "plugin"
163
- else:
164
- return self.create_error_response(
165
- f"pooling_task must be one of {self.supported_tasks}."
166
- )
167
- else:
168
- pooling_task = request.task
169
-
170
- if pooling_task not in self.supported_tasks:
171
- return self.create_error_response(
172
- f"Task {pooling_task} is not supported, it"
173
- f" must be one of {self.supported_tasks}."
174
- )
175
-
176
- try:
177
- pooling_params.verify(pooling_task, self.model_config)
178
- except ValueError as e:
179
- return self.create_error_response(str(e))
180
-
152
for i, engine_prompt in enumerate(engine_prompts):
153
request_id_item = f"{request_id}-{i}"
154
@@ -212,8 +183,7 @@ class OpenAIServingPooling(OpenAIServing):
183
184
generators.append(generator)
185
except ValueError as e:
215
- # TODO: Use a vllm-specific Validation Error
216
- return self.create_error_response(str(e))
186
+ return self.create_error_response(e)
187
188
result_generator = merge_async_iterators(*generators)
189
@@ -251,8 +221,7 @@ class OpenAIServingPooling(OpenAIServing):
221
except asyncio.CancelledError:
222
return self.create_error_response("Client disconnected")
223
except ValueError as e:
254
- # TODO: Use a vllm-specific Validation Error
255
- return self.create_error_response(str(e))
224
+ return self.create_error_response(e)
225
226
return response
227
vllm/entrypoints/pooling/score/protocol.py
+10
-1
@@ -18,6 +18,7 @@ from vllm.entrypoints.pooling.score.utils import (
18
ScoreInputs,
19
)
20
from vllm.renderers import TokenizeParams
21
+from vllm.tasks import PoolingTask
22
from vllm.utils import random_uuid
23
24
@@ -40,8 +41,9 @@ class ScoreRequestMixin(PoolingBasicRequestMixin, ClassifyRequestMixin):
41
max_total_tokens_param="max_model_len",
42
)
43
43
- def to_pooling_params(self):
44
+ def to_pooling_params(self, task: PoolingTask = "score"):
45
return PoolingParams(
46
+ task=task,
47
truncate_prompt_tokens=self.truncate_prompt_tokens,
48
use_activation=self.use_activation,
49
)
@@ -122,6 +124,13 @@ class RerankRequest(PoolingBasicRequestMixin, ClassifyRequestMixin):
124
max_total_tokens_param="max_model_len",
125
)
126
127
+ def to_pooling_params(self, task: PoolingTask = "score"):
128
+ return PoolingParams(
129
+ task=task,
130
+ truncate_prompt_tokens=self.truncate_prompt_tokens,
131
+ use_activation=self.use_activation,
132
+ )
133
+
134
135
class RerankDocument(BaseModel):
136
text: str | None = None
vllm/entrypoints/pooling/score/serving.py
+5
-29
@@ -118,12 +118,7 @@ class ServingScores(OpenAIServing):
118
119
# Schedule the request and get the result generator.
120
generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
121
- pooling_params = request.to_pooling_params()
122
-
123
- try:
124
- pooling_params.verify("embed", self.model_config)
125
- except ValueError as e:
126
- return self.create_error_response(str(e))
121
+ pooling_params = request.to_pooling_params("embed")
122
123
for i, engine_prompt in enumerate(engine_prompts):
124
request_id_item = f"{request_id}-{i}"
@@ -223,19 +218,7 @@ class ServingScores(OpenAIServing):
218
# Schedule the request and get the result generator.
219
generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
220
226
- # Use token_embed task for late interaction models
227
- from vllm import PoolingParams
228
-
229
- pooling_params = PoolingParams(
230
- task="token_embed",
231
- truncate_prompt_tokens=request.truncate_prompt_tokens,
232
- use_activation=request.use_activation,
233
- )
234
-
235
- try:
236
- pooling_params.verify("token_embed", self.model_config)
237
- except ValueError as e:
238
- return self.create_error_response(str(e))
221
+ pooling_params = request.to_pooling_params("token_embed")
222
223
for i, engine_prompt in enumerate(engine_prompts):
224
request_id_item = f"{request_id}-{i}"
@@ -358,12 +341,7 @@ class ServingScores(OpenAIServing):
341
# Schedule the request and get the result generator.
342
generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
343
361
- default_pooling_params = request.to_pooling_params()
362
-
363
- try:
364
- default_pooling_params.verify("score", self.model_config)
365
- except ValueError as e:
366
- return self.create_error_response(str(e))
344
+ default_pooling_params = request.to_pooling_params("score")
345
346
for i, engine_prompt in enumerate(engine_prompts):
347
request_id_item = f"{request_id}-{i}"
@@ -497,8 +475,7 @@ class ServingScores(OpenAIServing):
475
except asyncio.CancelledError:
476
return self.create_error_response("Client disconnected")
477
except ValueError as e:
500
- # TODO: Use a vllm-specific Validation Error
501
- return self.create_error_response(str(e))
478
+ return self.create_error_response(e)
479
480
async def do_rerank(
481
self, request: RerankRequest, raw_request: Request | None = None
@@ -542,8 +519,7 @@ class ServingScores(OpenAIServing):
519
except asyncio.CancelledError:
520
return self.create_error_response("Client disconnected")
521
except ValueError as e:
545
- # TODO: Use a vllm-specific Validation Error
546
- return self.create_error_response(str(e))
522
+ return self.create_error_response(e)
523
524
def request_output_to_score_response(
525
self,
vllm/pooling_params.py
+4
-12
@@ -72,15 +72,7 @@ class PoolingParams(
72
"""Returns a deep copy of the PoolingParams instance."""
73
return deepcopy(self)
74
75
- def verify(
76
- self, task: PoolingTask, model_config: "ModelConfig | None" = None
77
- ) -> None:
78
- if self.task is None:
79
- self.task = task
80
- elif self.task != task:
81
- msg = f"You cannot overwrite {self.task=!r} with {task=!r}!"
82
- raise ValueError(msg)
83
-
75
+ def verify(self, model_config: "ModelConfig") -> None:
76
# plugin task uses io_processor.parse_request to verify inputs,
77
# skipping PoolingParams verify
78
if self.task == "plugin":
@@ -167,7 +159,7 @@ class PoolingParams(
159
if mds is not None:
160
if self.dimensions not in mds:
161
raise ValueError(
170
- f'Model "{model_config.served_model_name}" '
162
+ f"Model {model_config.served_model_name!r} "
163
f"only supports {str(mds)} matryoshka dimensions, "
164
f"use other output dimensions will "
165
f"lead to poor results."
@@ -179,7 +171,7 @@ class PoolingParams(
171
if self.use_activation is None:
172
self.use_activation = True
173
else:
182
- raise ValueError(f"Unknown pooling task: {self.task}")
174
+ raise ValueError(f"Unknown pooling task: {self.task!r}")
175
176
def _verify_valid_parameters(self):
177
assert self.task is not None, "task must be set"
@@ -194,7 +186,7 @@ class PoolingParams(
186
187
if invalid_parameters:
188
raise ValueError(
197
- f"Task {self.task} only supports {valid_parameters} "
189
+ f"Task {self.task!r} only supports {valid_parameters} "
190
f"parameters, does not support "
191
f"{invalid_parameters} parameters"
192
)
vllm/v1/engine/async_llm.py
+6
-1
@@ -269,7 +269,11 @@ class AsyncLLM(EngineClient):
269
cancel_task_threadsafe(handler)
270
271
async def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
272
- return await self.engine_core.get_supported_tasks_async()
272
+ if not hasattr(self, "_supported_tasks"):
273
+ # Cache the result
274
+ self._supported_tasks = await self.engine_core.get_supported_tasks_async()
275
+
276
+ return self._supported_tasks
277
278
async def add_request(
279
self,
@@ -355,6 +359,7 @@ class AsyncLLM(EngineClient):
359
trace_headers=trace_headers,
360
priority=priority,
361
data_parallel_rank=data_parallel_rank,
362
+ supported_tasks=await self.get_supported_tasks(),
363
)
364
prompt_text = get_prompt_text(prompt)
365
vllm/v1/engine/input_processor.py
+32
-2
@@ -31,6 +31,7 @@ from vllm.multimodal.utils import argsort_mm_positions
31
from vllm.pooling_params import PoolingParams
32
from vllm.renderers import BaseRenderer
33
from vllm.sampling_params import _SAMPLING_EPS, SamplingParams
34
+from vllm.tasks import POOLING_TASKS, SupportedTask
35
from vllm.tokenizers import TokenizerLike
36
from vllm.tokenizers.mistral import MistralTokenizer
37
from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid
@@ -196,13 +197,41 @@ class InputProcessor:
197
def _validate_params(
198
self,
199
params: SamplingParams | PoolingParams,
200
+ # TODO: Validate generation tasks as well once `supported_tasks`
201
+ # is passed to all `process_inputs` calls
202
+ supported_tasks: tuple[SupportedTask, ...] | None,
203
):
204
"""
205
Validate supported SamplingParam.
206
Should raise ValueError if unsupported for API Server.
207
"""
204
-
208
if isinstance(params, PoolingParams):
209
+ if supported_tasks is None:
210
+ raise RuntimeError("`supported_tasks` must be passed for pooling")
211
+
212
+ supported_pooling_tasks = [
213
+ task for task in supported_tasks if task in POOLING_TASKS
214
+ ]
215
+
216
+ if params.task is None:
217
+ if not supported_pooling_tasks:
218
+ raise ValueError("Pooling tasks are not supported")
219
+
220
+ if "token_embed" in supported_pooling_tasks:
221
+ params.task = "token_embed"
222
+ elif "token_classify" in supported_pooling_tasks:
223
+ params.task = "token_classify"
224
+ elif "plugin" in supported_pooling_tasks:
225
+ params.task = "plugin"
226
+
227
+ if params.task not in supported_pooling_tasks:
228
+ raise ValueError(
229
+ f"Unsupported task: {params.task!r} "
230
+ f"Supported tasks: {supported_pooling_tasks}"
231
+ )
232
+
233
+ params.verify(self.model_config)
234
+
235
return
236
237
self._validate_logprobs(params)
@@ -498,10 +527,11 @@ class InputProcessor:
527
trace_headers: Mapping[str, str] | None = None,
528
priority: int = 0,
529
data_parallel_rank: int | None = None,
530
+ supported_tasks: tuple[SupportedTask, ...] | None = None,
531
resumable: bool = False,
532
) -> EngineCoreRequest:
533
self._validate_lora(lora_request)
504
- self._validate_params(params)
534
+ self._validate_params(params, supported_tasks)
535
536
parallel_config = self.vllm_config.parallel_config
537
dp_size = parallel_config.data_parallel_size
vllm/v1/engine/llm_engine.py
+6
-1
@@ -201,7 +201,11 @@ class LLMEngine:
201
return outputs
202
203
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
204
- return self.engine_core.get_supported_tasks()
204
+ if not hasattr(self, "_supported_tasks"):
205
+ # Cache the result
206
+ self._supported_tasks = self.engine_core.get_supported_tasks()
207
+
208
+ return self._supported_tasks
209
210
def abort_request(self, request_ids: list[str], internal: bool = False) -> None:
211
"""Remove request_ids from EngineCore and Detokenizer."""
@@ -245,6 +249,7 @@ class LLMEngine:
249
tokenization_kwargs,
250
trace_headers,
251
priority,
252
+ supported_tasks=self.get_supported_tasks(),
253
)
254
prompt_text = get_prompt_text(prompt)
255
vllm/v1/worker/gpu_model_runner.py
+1
-1
@@ -5037,7 +5037,7 @@ class GPUModelRunner(
5037
5038
model = cast(VllmModelForPooling, self.get_model())
5039
dummy_pooling_params = PoolingParams(task=task)
5040
- dummy_pooling_params.verify(task=task, model_config=self.model_config)
5040
+ dummy_pooling_params.verify(self.model_config)
5041
to_update = model.pooler.get_pooling_updates(task)
5042
to_update.apply(dummy_pooling_params)
5043