@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 75fd68c7a5
Fork 0

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 3, 14:04 UTC

[Entrypoints] Split the pooling offline API into PoolingOfflineMixin. (#42267)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>

wang.yuqi committed May 15, 2026 at 16:05 UTC 75fd68c7a57a2f9560919a06863ea739ab46554c
9 files changed +531 -439
docs/models/pooling_models/README.md
+4 -4
@@ -131,24 +131,24 @@ enabling the corresponding APIs.
131
132 ### `LLM.classify`
133
134 -The [classify][vllm.LLM.classify] method outputs a probability vector for each prompt.
134 +The [classify][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.classify] method outputs a probability vector for each prompt.
135 It is primarily designed for [classification models](classify.md).
136 For more information about `LLM.embed`, see [this page](classify.md#offline-inference).
137
138 ### `LLM.embed`
139
140 -The [embed][vllm.LLM.embed] method outputs an embedding vector for each prompt.
140 +The [embed][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.embed] method outputs an embedding vector for each prompt.
141 It is primarily designed for [embedding models](embed.md).
142 For more information about `LLM.embed`, see [this page](embed.md#offline-inference).
143
144 ### `LLM.score`
145
146 -The [score][vllm.LLM.score] method outputs similarity scores between sentence pairs.
146 +The [score][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.score] method outputs similarity scores between sentence pairs.
147 It is primarily designed for [score models](scoring.md).
148
149 ### `LLM.encode`
150
151 -The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
151 +The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.
152
153 Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).
154
docs/models/pooling_models/classify.md
+2 -2
@@ -77,7 +77,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.
77
78 ### `LLM.classify`
79
80 -The [classify][vllm.LLM.classify] method outputs a probability vector for each prompt.
80 +The [classify][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.classify] method outputs a probability vector for each prompt.
81
82 ```python
83 from vllm import LLM
@@ -93,7 +93,7 @@ A code example can be found here: [examples/basic/offline_inference/classify.py]
93
94 ### `LLM.encode`
95
96 -The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
96 +The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.
97
98 Set `pooling_task="classify"` when using `LLM.encode` for classification Models:
99
docs/models/pooling_models/embed.md
+3 -3
@@ -120,7 +120,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.
120
121 ### `LLM.embed`
122
123 -The [embed][vllm.LLM.embed] method outputs an embedding vector for each prompt.
123 +The [embed][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.embed] method outputs an embedding vector for each prompt.
124
125 ```python
126 from vllm import LLM
@@ -136,7 +136,7 @@ A code example can be found here: [examples/basic/offline_inference/embed.py](..
136
137 ### `LLM.encode`
138
139 -The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
139 +The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.
140
141 Set `pooling_task="embed"` when using `LLM.encode` for embedding Models:
142
@@ -152,7 +152,7 @@ print(f"Data: {data!r}")
152
153 ### `LLM.score`
154
155 -The [score][vllm.LLM.score] method outputs similarity scores between sentence pairs.
155 +The [score][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.score] method outputs similarity scores between sentence pairs.
156
157 All models that support embedding task also support using the score API to compute similarity scores by calculating the cosine similarity of two input prompt's embeddings.
158
docs/models/pooling_models/reward.md
+1 -1
@@ -87,7 +87,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.
87
88 ### `LLM.encode`
89
90 -The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
90 +The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.
91
92 - Reward Models
93
docs/models/pooling_models/scoring.md
+1 -1
@@ -130,7 +130,7 @@ The following [pooling parameters][vllm.PoolingParams] are only supported by cro
130
131 ### `LLM.score`
132
133 -The [score][vllm.LLM.score] method outputs similarity scores between sentence pairs.
133 +The [score][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.score] method outputs similarity scores between sentence pairs.
134
135 ```python
136 from vllm import LLM
docs/models/pooling_models/token_classify.md
+1 -1
@@ -87,7 +87,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.
87
88 ### `LLM.encode`
89
90 -The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
90 +The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.
91
92 Set `pooling_task="token_classify"` when using `LLM.encode` for token classification Models:
93
docs/models/pooling_models/token_embed.md
+2 -2
@@ -94,7 +94,7 @@ The following [pooling parameters][vllm.PoolingParams] are supported.
94
95 ### `LLM.encode`
96
97 -The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
97 +The [encode][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.encode] method is available to all pooling models in vLLM.
98
99 Set `pooling_task="token_embed"` when using `LLM.encode` for token embedding Models:
100
@@ -110,7 +110,7 @@ print(f"Data: {data!r}")
110
111 ### `LLM.score`
112
113 -The [score][vllm.LLM.score] method outputs similarity scores between sentence pairs.
113 +The [score][vllm.entrypoints.pooling.offline.PoolingOfflineMixin.score] method outputs similarity scores between sentence pairs.
114
115 All models that support token embedding task also support using the score API to compute similarity scores by calculating the late interaction of two input prompts.
116
vllm/entrypoints/llm.py
+7 -425
@@ -34,9 +34,7 @@ from vllm.config.model import (
34 RunnerOption,
35 TokenizerMode,
36 )
37 -from vllm.config.quantization import (
38 - QuantizationConfigArgs,
39 -)
37 +from vllm.config.quantization import QuantizationConfigArgs
38 from vllm.distributed.weight_transfer.base import (
39 WeightTransferInitRequest,
40 WeightTransferUpdateRequest,
@@ -44,17 +42,12 @@ from vllm.distributed.weight_transfer.base import (
42 from vllm.engine.arg_utils import EngineArgs
43 from vllm.entrypoints.chat_utils import (
44 ChatCompletionMessageParam,
47 - ChatTemplateConfig,
45 ChatTemplateContentFormatOption,
46 load_chat_template,
47 )
51 -from vllm.entrypoints.pooling.factories import init_pooling_io_processors
52 -from vllm.entrypoints.pooling.scoring.io_processor import ScoringIOProcessor
53 -from vllm.entrypoints.pooling.scoring.typing import ScoreInput
54 -from vllm.entrypoints.pooling.typing import OfflineInputsContext, OfflineOutputsContext
48 +from vllm.entrypoints.pooling.offline import PoolingOfflineMixin
49 from vllm.entrypoints.utils import log_non_default_args
50 from vllm.inputs import (
57 - DataPrompt,
51 EngineInput,
52 PromptType,
53 TextPrompt,
@@ -63,13 +56,7 @@ from vllm.inputs import (
56 from vllm.logger import init_logger
57 from vllm.lora.request import LoRARequest
58 from vllm.model_executor.layers.quantization import QuantizationMethods
66 -from vllm.outputs import (
67 - ClassificationRequestOutput,
68 - EmbeddingRequestOutput,
69 - PoolingRequestOutput,
70 - RequestOutput,
71 - ScoringRequestOutput,
72 -)
59 +from vllm.outputs import PoolingRequestOutput, RequestOutput
60 from vllm.platforms import current_platform
61 from vllm.pooling_params import PoolingParams
62 from vllm.renderers import ChatParams, merge_kwargs
@@ -79,7 +66,6 @@ from vllm.renderers.inputs.preprocess import (
66 prompt_to_seq,
67 )
68 from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
82 -from vllm.tasks import SCORE_TYPE_MAP, PoolingTask
69 from vllm.tokenizers import TokenizerLike
70 from vllm.usage.usage_lib import UsageContext
71 from vllm.utils.counter import Counter
@@ -103,7 +89,7 @@ _P = TypeVar("_P", bound=SamplingParams | PoolingParams | None)
89 _R = TypeVar("_R", default=Any)
90
91
106 -class LLM:
92 +class LLM(PoolingOfflineMixin):
93 """An LLM for generating texts from given prompts and sampling parameters.
94
95 This class includes a tokenizer, a language model (possibly distributed
@@ -387,21 +373,14 @@ class LLM:
373
374 supported_tasks = self.llm_engine.get_supported_tasks()
375 self.supported_tasks = supported_tasks
390 - self.pooling_task = self.model_config.get_pooling_task(supported_tasks)
391 - if self.pooling_task is not None:
392 - logger.info("Supported pooling task: %s", self.pooling_task)
376
377 self.runner_type = self.model_config.runner_type
378 self.renderer = self.llm_engine.renderer
379 self.chat_template = load_chat_template(chat_template)
380 self.input_processor = self.llm_engine.input_processor
398 - self.chat_template_config = ChatTemplateConfig(chat_template=self.chat_template)
399 - self.pooling_io_processors = init_pooling_io_processors(
400 - supported_tasks=supported_tasks,
401 - vllm_config=self.llm_engine.vllm_config,
402 - renderer=self.renderer,
403 - chat_template_config=self.chat_template_config,
404 - )
381 +
382 + PoolingOfflineMixin.__init__(self)
383 +
384 # Cache for __repr__ to avoid repeated collective_rpc calls
385 self._cached_repr: str | None = None
386
@@ -1070,403 +1049,6 @@ class LLM:
1049 mm_processor_kwargs=mm_processor_kwargs,
1050 )
1051
1073 - def encode(
1074 - self,
1075 - prompts: PromptType | Sequence[PromptType] | DataPrompt,
1076 - pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,
1077 - *,
1078 - use_tqdm: bool | Callable[..., tqdm] = True,
1079 - lora_request: list[LoRARequest] | LoRARequest | None = None,
1080 - pooling_task: PoolingTask | None = None,
1081 - tokenization_kwargs: dict[str, Any] | None = None,
1082 - ) -> list[PoolingRequestOutput]:
1083 - """Apply pooling to the hidden states corresponding to the input
1084 - prompts.
1085 -
1086 - This class automatically batches the given prompts, considering
1087 - the memory constraint. For the best performance, put all of your prompts
1088 - into a single list and pass it to this method.
1089 -
1090 - Args:
1091 - prompts: The prompts to the LLM. You may pass a sequence of prompts
1092 - for batch inference. See [PromptType][vllm.inputs.PromptType]
1093 - for more details about the format of each prompt.
1094 - pooling_params: The pooling parameters for pooling. If None, we
1095 - use the default pooling parameters.
1096 - use_tqdm: If `True`, shows a tqdm progress bar.
1097 - If a callable (e.g., `functools.partial(tqdm, leave=False)`),
1098 - it is used to create the progress bar.
1099 - If `False`, no progress bar is created.
1100 - lora_request: LoRA request to use for generation, if any.
1101 - pooling_task: Override the pooling task to use.
1102 - tokenization_kwargs: Overrides for `tokenizer.encode`.
1103 -
1104 - Returns:
1105 - A list of `PoolingRequestOutput` objects containing the
1106 - pooled hidden states in the same order as the input prompts.
1107 - """
1108 -
1109 - if isinstance(prompts, dict) and "data" in prompts and pooling_task != "plugin":
1110 - raise ValueError(
1111 - "The 'data' field is only supported for the 'plugin' pooling task."
1112 - )
1113 - self._verify_pooling_task(pooling_task)
1114 - assert pooling_task is not None and pooling_task in self.pooling_io_processors
1115 -
1116 - io_processor = self.pooling_io_processors[pooling_task]
1117 -
1118 - if pooling_params is None:
1119 - pooling_params = PoolingParams()
1120 -
1121 - ctx = OfflineInputsContext(
1122 - prompts=prompts,
1123 - pooling_params=pooling_params,
1124 - tokenization_kwargs=tokenization_kwargs,
1125 - )
1126 -
1127 - engine_inputs = io_processor.pre_process_offline(ctx)
1128 - n_inputs = len(engine_inputs)
1129 - assert ctx.pooling_params is not None
1130 -
1131 - params_seq = self._params_to_seq(ctx.pooling_params, n_inputs)
1132 -
1133 - for param in params_seq:
1134 - if param.task is None:
1135 - param.task = pooling_task
1136 - elif pooling_task == "plugin":
1137 - # `plugin` task uses io_processor.parse_request to verify inputs.
1138 - # We actually allow plugin to overwrite pooling_task.
1139 - pass
1140 - elif param.task != pooling_task:
1141 - msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
1142 - raise ValueError(msg)
1143 -
1144 - seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs)
1145 - seq_priority = self._priority_to_seq(None, n_inputs)
1146 -
1147 - self._render_and_add_requests(
1148 - prompts=engine_inputs,
1149 - params=params_seq,
1150 - lora_requests=seq_lora_requests,
1151 - priorities=seq_priority,
1152 - )
1153 -
1154 - outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput)
1155 - outputs = io_processor.post_process_offline(
1156 - ctx=OfflineOutputsContext(outputs=outputs)
1157 - )
1158 - return outputs
1159 -
1160 - def _verify_pooling_task(self, pooling_task: PoolingTask | None):
1161 - if self.runner_type != "pooling":
1162 - raise ValueError(
1163 - "LLM.encode() is only supported for pooling models. "
1164 - "Try passing `--runner pooling` to use the model as a "
1165 - "pooling model."
1166 - )
1167 -
1168 - if pooling_task is None:
1169 - raise ValueError(
1170 - """
1171 - pooling_task required for `LLM.encode`.
1172 - Please use one of the more specific methods or set the pooling_task when using `LLM.encode`:
1173 - - For embeddings, use `LLM.embed(...)` or `pooling_task="embed"`.
1174 - - For classification logits, use `LLM.classify(...)` or `pooling_task="classify"`.
1175 - - For similarity scores, use `LLM.score(...)`.
1176 - - For rewards, `pooling_task="classify"` or `pooling_task="token_classify"`.
1177 - - For token classification, use `pooling_task="token_classify"`.
1178 - - For multi-vector retrieval, use `pooling_task="token_embed"`.
1179 - """ # noqa: E501
1180 - )
1181 -
1182 - if (
1183 - pooling_task in ("embed", "token_embed")
1184 - and pooling_task not in self.supported_tasks
1185 - ):
1186 - raise ValueError(
1187 - "Embedding API is not supported by this model. "
1188 - "Try converting the model using `--convert embed`."
1189 - )
1190 -
1191 - if (
1192 - pooling_task in ("classify", "token_classify")
1193 - and pooling_task not in self.supported_tasks
1194 - ):
1195 - raise ValueError(
1196 - "Classification API is not supported by this model. "
1197 - "Try converting the model using `--convert classify`."
1198 - )
1199 -
1200 - # plugin task uses io_processor.parse_request to verify inputs
1201 - if pooling_task != "plugin" and pooling_task != self.pooling_task:
1202 - if pooling_task not in self.supported_tasks:
1203 - raise ValueError(
1204 - f"Unsupported task: {pooling_task!r} "
1205 - f"Supported tasks: {self.supported_tasks}"
1206 - )
1207 - else:
1208 - raise ValueError(
1209 - f"Try switching the model's pooling_task "
1210 - f'via `PoolerConfig(task="{pooling_task}")`'
1211 - )
1212 -
1213 - if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors:
1214 - raise ValueError(
1215 - "No IOProcessor plugin installed. Please refer "
1216 - "to the documentation and to the "
1217 - "'prithvi_geospatial_mae_io_processor' "
1218 - "offline inference example for more details."
1219 - )
1220 -
1221 - def embed(
1222 - self,
1223 - prompts: PromptType | Sequence[PromptType],
1224 - *,
1225 - use_tqdm: bool | Callable[..., tqdm] = True,
1226 - pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,
1227 - lora_request: list[LoRARequest] | LoRARequest | None = None,
1228 - tokenization_kwargs: dict[str, Any] | None = None,
1229 - ) -> list[EmbeddingRequestOutput]:
1230 - """
1231 - Generate an embedding vector for each prompt.
1232 -
1233 - This class automatically batches the given prompts, considering
1234 - the memory constraint. For the best performance, put all of your prompts
1235 - into a single list and pass it to this method.
1236 -
1237 - Args:
1238 - prompts: The prompts to the LLM. You may pass a sequence of prompts
1239 - for batch inference. See [PromptType][vllm.inputs.PromptType]
1240 - for more details about the format of each prompt.
1241 - pooling_params: The pooling parameters for pooling. If None, we
1242 - use the default pooling parameters.
1243 - use_tqdm: If `True`, shows a tqdm progress bar.
1244 - If a callable (e.g., `functools.partial(tqdm, leave=False)`),
1245 - it is used to create the progress bar.
1246 - If `False`, no progress bar is created.
1247 - lora_request: LoRA request to use for generation, if any.
1248 - tokenization_kwargs: Overrides for `tokenizer.encode`.
1249 -
1250 - Returns:
1251 - A list of `EmbeddingRequestOutput` objects containing the
1252 - embedding vectors in the same order as the input prompts.
1253 - """
1254 -
1255 - items = self.encode(
1256 - prompts,
1257 - use_tqdm=use_tqdm,
1258 - pooling_params=pooling_params,
1259 - lora_request=lora_request,
1260 - pooling_task="embed",
1261 - tokenization_kwargs=tokenization_kwargs,
1262 - )
1263 -
1264 - return [EmbeddingRequestOutput.from_base(item) for item in items]
1265 -
1266 - def classify(
1267 - self,
1268 - prompts: PromptType | Sequence[PromptType],
1269 - *,
1270 - pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,
1271 - use_tqdm: bool | Callable[..., tqdm] = True,
1272 - lora_request: list[LoRARequest] | LoRARequest | None = None,
1273 - tokenization_kwargs: dict[str, Any] | None = None,
1274 - ) -> list[ClassificationRequestOutput]:
1275 - """
1276 - Generate class logits for each prompt.
1277 -
1278 - This class automatically batches the given prompts, considering
1279 - the memory constraint. For the best performance, put all of your prompts
1280 - into a single list and pass it to this method.
1281 -
1282 - Args:
1283 - prompts: The prompts to the LLM. You may pass a sequence of prompts
1284 - for batch inference. See [PromptType][vllm.inputs.PromptType]
1285 - for more details about the format of each prompt.
1286 - pooling_params: The pooling parameters for pooling. If None, we
1287 - use the default pooling parameters.
1288 - use_tqdm: If `True`, shows a tqdm progress bar.
1289 - If a callable (e.g., `functools.partial(tqdm, leave=False)`),
1290 - it is used to create the progress bar.
1291 - If `False`, no progress bar is created.
1292 - lora_request: LoRA request to use for generation, if any.
1293 - tokenization_kwargs: Overrides for `tokenizer.encode`.
1294 -
1295 - Returns:
1296 - A list of `ClassificationRequestOutput` objects containing the
1297 - embedding vectors in the same order as the input prompts.
1298 - """
1299 -
1300 - items = self.encode(
1301 - prompts,
1302 - use_tqdm=use_tqdm,
1303 - pooling_params=pooling_params,
1304 - lora_request=lora_request,
1305 - pooling_task="classify",
1306 - tokenization_kwargs=tokenization_kwargs,
1307 - )
1308 -
1309 - return [ClassificationRequestOutput.from_base(item) for item in items]
1310 -
1311 - def reward(
1312 - self,
1313 - prompts: PromptType | Sequence[PromptType],
1314 - /,
1315 - *,
1316 - pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,
1317 - use_tqdm: bool | Callable[..., tqdm] = True,
1318 - lora_request: list[LoRARequest] | LoRARequest | None = None,
1319 - tokenization_kwargs: dict[str, Any] | None = None,
1320 - ) -> list[PoolingRequestOutput]:
1321 - """
1322 - Generate rewards for each prompt.
1323 -
1324 - Args:
1325 - prompts: The prompts to the LLM. You may pass a sequence of prompts
1326 - for batch inference. See [PromptType][vllm.inputs.PromptType]
1327 - for more details about the format of each prompt.
1328 - pooling_params: The pooling parameters for pooling. If None, we
1329 - use the default pooling parameters.
1330 - use_tqdm: If `True`, shows a tqdm progress bar.
1331 - If a callable (e.g., `functools.partial(tqdm, leave=False)`),
1332 - it is used to create the progress bar.
1333 - If `False`, no progress bar is created.
1334 - lora_request: LoRA request to use for generation, if any.
1335 - tokenization_kwargs: Overrides for `tokenizer.encode`.
1336 -
1337 - Returns:
1338 - A list of `PoolingRequestOutput` objects containing the
1339 - pooled hidden states in the same order as the input prompts.
1340 - """
1341 - logger.warning_once(
1342 - "`llm.reward` api is deprecated and will be removed in v0.23. "
1343 - 'Please use `LLM.encode` with `pooling_task="classify"` or '
1344 - '`pooling_task="token_classify"` instead.'
1345 - )
1346 - return self.encode(
1347 - prompts,
1348 - use_tqdm=use_tqdm,
1349 - lora_request=lora_request,
1350 - pooling_params=pooling_params,
1351 - pooling_task="token_classify",
1352 - tokenization_kwargs=tokenization_kwargs,
1353 - )
1354 -
1355 - def score(
1356 - self,
1357 - data_1: ScoreInput | list[ScoreInput],
1358 - data_2: ScoreInput | list[ScoreInput],
1359 - /,
1360 - *,
1361 - use_tqdm: bool | Callable[..., tqdm] = True,
1362 - pooling_params: PoolingParams | None = None,
1363 - lora_request: list[LoRARequest] | LoRARequest | None = None,
1364 - tokenization_kwargs: dict[str, Any] | None = None,
1365 - chat_template: str | None = None,
1366 - ) -> list[ScoringRequestOutput]:
1367 - """Generate similarity scores for all pairs `<text,text_pair>` or
1368 - `<multi-modal data, multi-modal data pair>`.
1369 -
1370 - The inputs can be `1 -> 1`, `1 -> N` or `N -> N`.
1371 - In the `1 - N` case the `data_1` input will be replicated `N`
1372 - times to pair with the `data_2` inputs.
1373 - The input pairs are used to build a list of prompts for the
1374 - cross encoder model. This class automatically batches the prompts,
1375 - considering the memory constraint. For the best performance, put all
1376 - of your inputs into a single list and pass it to this method.
1377 -
1378 - Supports both text and multi-modal data (images, etc.) when used with
1379 - appropriate multi-modal models. For multi-modal inputs, ensure the
1380 - prompt structure matches the model's expected input format.
1381 -
1382 - Args:
1383 - data_1: Can be a single prompt, a list of prompts or
1384 - `ScoreMultiModalParam`, which can contain either text or
1385 - multi-modal data. When a list, it must have the same length as
1386 - the `data_2` list.
1387 - data_2: The data to pair with the query to form the input to
1388 - the LLM. Can be text or multi-modal data. See [PromptType]
1389 - [vllm.inputs.PromptType] for more details about the format of
1390 - each prompt.
1391 - pooling_params: The pooling parameters for pooling. If None, we
1392 - use the default pooling parameters.
1393 - use_tqdm: If `True`, shows a tqdm progress bar.
1394 - If a callable (e.g., `functools.partial(tqdm, leave=False)`),
1395 - it is used to create the progress bar.
1396 - If `False`, no progress bar is created.
1397 - lora_request: LoRA request to use for generation, if any.
1398 - chat_template: The chat template to use for the scoring. If None, we
1399 - use the model's default chat template.
1400 - tokenization_kwargs: Overrides for `tokenizer.encode`.
1401 - Returns:
1402 - A list of `ScoringRequestOutput` objects containing the
1403 - generated scores in the same order as the input prompts.
1404 - """
1405 -
1406 - if self.runner_type != "pooling":
1407 - raise ValueError(
1408 - "LLM.score() is only supported for pooling models. "
1409 - "Try passing `--runner pooling` to use the model as a "
1410 - "pooling model."
1411 - )
1412 -
1413 - score_type: str | None = SCORE_TYPE_MAP.get(self.pooling_task, None) # type: ignore[arg-type]
1414 - if (
1415 - score_type == "cross-encoder"
1416 - and getattr(self.model_config.hf_config, "num_labels", 0) != 1
1417 - ):
1418 - raise ValueError("Scoring API is only enabled for num_labels == 1.")
1419 -
1420 - if score_type is None or score_type not in self.pooling_io_processors:
1421 - raise ValueError("This model does not support the Scoring API.")
1422 -
1423 - io_processor = self.pooling_io_processors[score_type]
1424 - assert isinstance(io_processor, ScoringIOProcessor)
1425 -
1426 - pooling_task = io_processor.pooling_task
1427 - scoring_data = io_processor.valid_inputs(data_1, data_2)
1428 - n_queries = len(scoring_data.data_1)
1429 -
1430 - if pooling_params is None:
1431 - pooling_params = PoolingParams()
1432 -
1433 - ctx = OfflineInputsContext(
1434 - prompts=scoring_data,
1435 - pooling_params=pooling_params,
1436 - tokenization_kwargs=tokenization_kwargs,
1437 - chat_template=chat_template,
1438 - n_queries=n_queries,
1439 - )
1440 -
1441 - engine_inputs = io_processor.pre_process_offline(ctx)
1442 - n_inputs = len(engine_inputs)
1443 -
1444 - seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs)
1445 - params_seq = self._params_to_seq(ctx.pooling_params, n_inputs)
1446 -
1447 - for param in params_seq:
1448 - if param.task is None:
1449 - param.task = pooling_task
1450 - elif param.task != pooling_task:
1451 - msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
1452 - raise ValueError(msg)
1453 -
1454 - seq_priority = self._priority_to_seq(None, n_inputs)
1455 -
1456 - self._render_and_add_requests(
1457 - prompts=engine_inputs,
1458 - params=params_seq,
1459 - lora_requests=seq_lora_requests,
1460 - priorities=seq_priority,
1461 - )
1462 -
1463 - outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput)
1464 - outputs = io_processor.post_process_offline(
1465 - ctx=OfflineOutputsContext(outputs=outputs, n_queries=n_queries),
1466 - )
1467 -
1468 - return [ScoringRequestOutput.from_base(item) for item in outputs]
1469 -
1052 def start_profile(self, profile_prefix: str | None = None) -> None:
1053 """Start profiling with optional custom trace prefix.
1054
vllm/entrypoints/pooling/offline.py new
+510
@@ -0,0 +1,510 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from abc import ABC, abstractmethod
5 +from collections.abc import Callable, Iterable, Sequence
6 +from typing import Any
7 +
8 +from tqdm.auto import tqdm
9 +from typing_extensions import TypeVar
10 +
11 +from vllm.config import ModelConfig
12 +from vllm.entrypoints.chat_utils import ChatTemplateConfig
13 +from vllm.inputs import (
14 + DataPrompt,
15 + EngineInput,
16 + PromptType,
17 +)
18 +from vllm.logger import init_logger
19 +from vllm.lora.request import LoRARequest
20 +from vllm.outputs import (
21 + ClassificationRequestOutput,
22 + EmbeddingRequestOutput,
23 + PoolingRequestOutput,
24 + RequestOutput,
25 + ScoringRequestOutput,
26 +)
27 +from vllm.pooling_params import PoolingParams
28 +from vllm.renderers import BaseRenderer
29 +from vllm.sampling_params import SamplingParams
30 +from vllm.tasks import SCORE_TYPE_MAP, PoolingTask, SupportedTask
31 +from vllm.v1.engine.llm_engine import LLMEngine
32 +
33 +from .factories import init_pooling_io_processors
34 +from .scoring.io_processor import ScoringIOProcessor
35 +from .scoring.typing import ScoreInput
36 +from .typing import OfflineInputsContext, OfflineOutputsContext
37 +
38 +logger = init_logger(__name__)
39 +
40 +_P = TypeVar("_P", bound=SamplingParams | PoolingParams | None)
41 +_O = TypeVar(
42 + "_O",
43 + bound=RequestOutput | PoolingRequestOutput,
44 + default=RequestOutput | PoolingRequestOutput,
45 +)
46 +
47 +
48 +class PoolingOfflineMixin(ABC):
49 + """Offline inference for pooling models"""
50 +
51 + renderer: BaseRenderer
52 + llm_engine: "LLMEngine"
53 + model_config: ModelConfig
54 + runner_type: str
55 + chat_template: str | None
56 + supported_tasks: tuple[SupportedTask, ...]
57 +
58 + def __init__(self):
59 + self.pooling_task = self.model_config.get_pooling_task(self.supported_tasks)
60 + if self.pooling_task is not None:
61 + logger.info("Supported pooling task: %s", self.pooling_task)
62 +
63 + self.chat_template_config = ChatTemplateConfig(chat_template=self.chat_template)
64 + self.pooling_io_processors = init_pooling_io_processors(
65 + supported_tasks=self.supported_tasks,
66 + vllm_config=self.llm_engine.vllm_config,
67 + renderer=self.renderer,
68 + chat_template_config=self.chat_template_config,
69 + )
70 +
71 + def encode(
72 + self,
73 + prompts: PromptType | Sequence[PromptType] | DataPrompt,
74 + pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,
75 + *,
76 + use_tqdm: bool | Callable[..., tqdm] = True,
77 + lora_request: list[LoRARequest] | LoRARequest | None = None,
78 + pooling_task: PoolingTask | None = None,
79 + tokenization_kwargs: dict[str, Any] | None = None,
80 + ) -> list[PoolingRequestOutput]:
81 + """Apply pooling to the hidden states corresponding to the input
82 + prompts.
83 +
84 + This class automatically batches the given prompts, considering
85 + the memory constraint. For the best performance, put all of your prompts
86 + into a single list and pass it to this method.
87 +
88 + Args:
89 + prompts: The prompts to the LLM. You may pass a sequence of prompts
90 + for batch inference. See [PromptType][vllm.inputs.PromptType]
91 + for more details about the format of each prompt.
92 + pooling_params: The pooling parameters for pooling. If None, we
93 + use the default pooling parameters.
94 + use_tqdm: If `True`, shows a tqdm progress bar.
95 + If a callable (e.g., `functools.partial(tqdm, leave=False)`),
96 + it is used to create the progress bar.
97 + If `False`, no progress bar is created.
98 + lora_request: LoRA request to use for generation, if any.
99 + pooling_task: Override the pooling task to use.
100 + tokenization_kwargs: Overrides for `tokenizer.encode`.
101 +
102 + Returns:
103 + A list of `PoolingRequestOutput` objects containing the
104 + pooled hidden states in the same order as the input prompts.
105 + """
106 +
107 + if isinstance(prompts, dict) and "data" in prompts and pooling_task != "plugin":
108 + raise ValueError(
109 + "The 'data' field is only supported for the 'plugin' pooling task."
110 + )
111 + self._verify_pooling_task(pooling_task)
112 + assert pooling_task is not None and pooling_task in self.pooling_io_processors
113 +
114 + io_processor = self.pooling_io_processors[pooling_task]
115 +
116 + if pooling_params is None:
117 + pooling_params = PoolingParams()
118 +
119 + ctx = OfflineInputsContext(
120 + prompts=prompts,
121 + pooling_params=pooling_params,
122 + tokenization_kwargs=tokenization_kwargs,
123 + )
124 +
125 + engine_inputs = io_processor.pre_process_offline(ctx)
126 + n_inputs = len(engine_inputs)
127 + assert ctx.pooling_params is not None
128 +
129 + params_seq = self._params_to_seq(ctx.pooling_params, n_inputs)
130 +
131 + for param in params_seq:
132 + if param.task is None:
133 + param.task = pooling_task
134 + elif pooling_task == "plugin":
135 + # `plugin` task uses io_processor.parse_request to verify inputs.
136 + # We actually allow plugin to overwrite pooling_task.
137 + pass
138 + elif param.task != pooling_task:
139 + msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
140 + raise ValueError(msg)
141 +
142 + seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs)
143 + seq_priority = self._priority_to_seq(None, n_inputs)
144 +
145 + self._render_and_add_requests(
146 + prompts=engine_inputs,
147 + params=params_seq,
148 + lora_requests=seq_lora_requests,
149 + priorities=seq_priority,
150 + )
151 +
152 + outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput)
153 + outputs = io_processor.post_process_offline(
154 + ctx=OfflineOutputsContext(outputs=outputs)
155 + )
156 + return outputs
157 +
158 + def _verify_pooling_task(self, pooling_task: PoolingTask | None):
159 + if self.runner_type != "pooling":
160 + raise ValueError(
161 + "LLM.encode() is only supported for pooling models. "
162 + "Try passing `--runner pooling` to use the model as a "
163 + "pooling model."
164 + )
165 +
166 + if pooling_task is None:
167 + raise ValueError(
168 + """
169 + pooling_task required for `LLM.encode`.
170 + Please use one of the more specific methods or set the pooling_task when using `LLM.encode`:
171 + - For embeddings, use `LLM.embed(...)` or `pooling_task="embed"`.
172 + - For classification logits, use `LLM.classify(...)` or `pooling_task="classify"`.
173 + - For similarity scores, use `LLM.score(...)`.
174 + - For rewards, `pooling_task="classify"` or `pooling_task="token_classify"`.
175 + - For token classification, use `pooling_task="token_classify"`.
176 + - For multi-vector retrieval, use `pooling_task="token_embed"`.
177 + """ # noqa: E501
178 + )
179 +
180 + if (
181 + pooling_task in ("embed", "token_embed")
182 + and pooling_task not in self.supported_tasks
183 + ):
184 + raise ValueError(
185 + "Embedding API is not supported by this model. "
186 + "Try converting the model using `--convert embed`."
187 + )
188 +
189 + if (
190 + pooling_task in ("classify", "token_classify")
191 + and pooling_task not in self.supported_tasks
192 + ):
193 + raise ValueError(
194 + "Classification API is not supported by this model. "
195 + "Try converting the model using `--convert classify`."
196 + )
197 +
198 + # plugin task uses io_processor.parse_request to verify inputs
199 + if pooling_task != "plugin" and pooling_task != self.pooling_task:
200 + if pooling_task not in self.supported_tasks:
201 + raise ValueError(
202 + f"Unsupported task: {pooling_task!r} "
203 + f"Supported tasks: {self.supported_tasks}"
204 + )
205 + else:
206 + raise ValueError(
207 + f"Try switching the model's pooling_task "
208 + f'via `PoolerConfig(task="{pooling_task}")`'
209 + )
210 +
211 + if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors:
212 + raise ValueError(
213 + "No IOProcessor plugin installed. Please refer "
214 + "to the documentation and to the "
215 + "'prithvi_geospatial_mae_io_processor' "
216 + "offline inference example for more details."
217 + )
218 +
219 + def embed(
220 + self,
221 + prompts: PromptType | Sequence[PromptType],
222 + *,
223 + use_tqdm: bool | Callable[..., tqdm] = True,
224 + pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,
225 + lora_request: list[LoRARequest] | LoRARequest | None = None,
226 + tokenization_kwargs: dict[str, Any] | None = None,
227 + ) -> list[EmbeddingRequestOutput]:
228 + """
229 + Generate an embedding vector for each prompt.
230 +
231 + This class automatically batches the given prompts, considering
232 + the memory constraint. For the best performance, put all of your prompts
233 + into a single list and pass it to this method.
234 +
235 + Args:
236 + prompts: The prompts to the LLM. You may pass a sequence of prompts
237 + for batch inference. See [PromptType][vllm.inputs.PromptType]
238 + for more details about the format of each prompt.
239 + pooling_params: The pooling parameters for pooling. If None, we
240 + use the default pooling parameters.
241 + use_tqdm: If `True`, shows a tqdm progress bar.
242 + If a callable (e.g., `functools.partial(tqdm, leave=False)`),
243 + it is used to create the progress bar.
244 + If `False`, no progress bar is created.
245 + lora_request: LoRA request to use for generation, if any.
246 + tokenization_kwargs: Overrides for `tokenizer.encode`.
247 +
248 + Returns:
249 + A list of `EmbeddingRequestOutput` objects containing the
250 + embedding vectors in the same order as the input prompts.
251 + """
252 +
253 + items = self.encode(
254 + prompts,
255 + use_tqdm=use_tqdm,
256 + pooling_params=pooling_params,
257 + lora_request=lora_request,
258 + pooling_task="embed",
259 + tokenization_kwargs=tokenization_kwargs,
260 + )
261 +
262 + return [EmbeddingRequestOutput.from_base(item) for item in items]
263 +
264 + def classify(
265 + self,
266 + prompts: PromptType | Sequence[PromptType],
267 + *,
268 + pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,
269 + use_tqdm: bool | Callable[..., tqdm] = True,
270 + lora_request: list[LoRARequest] | LoRARequest | None = None,
271 + tokenization_kwargs: dict[str, Any] | None = None,
272 + ) -> list[ClassificationRequestOutput]:
273 + """
274 + Generate class logits for each prompt.
275 +
276 + This class automatically batches the given prompts, considering
277 + the memory constraint. For the best performance, put all of your prompts
278 + into a single list and pass it to this method.
279 +
280 + Args:
281 + prompts: The prompts to the LLM. You may pass a sequence of prompts
282 + for batch inference. See [PromptType][vllm.inputs.PromptType]
283 + for more details about the format of each prompt.
284 + pooling_params: The pooling parameters for pooling. If None, we
285 + use the default pooling parameters.
286 + use_tqdm: If `True`, shows a tqdm progress bar.
287 + If a callable (e.g., `functools.partial(tqdm, leave=False)`),
288 + it is used to create the progress bar.
289 + If `False`, no progress bar is created.
290 + lora_request: LoRA request to use for generation, if any.
291 + tokenization_kwargs: Overrides for `tokenizer.encode`.
292 +
293 + Returns:
294 + A list of `ClassificationRequestOutput` objects containing the
295 + embedding vectors in the same order as the input prompts.
296 + """
297 +
298 + items = self.encode(
299 + prompts,
300 + use_tqdm=use_tqdm,
301 + pooling_params=pooling_params,
302 + lora_request=lora_request,
303 + pooling_task="classify",
304 + tokenization_kwargs=tokenization_kwargs,
305 + )
306 +
307 + return [ClassificationRequestOutput.from_base(item) for item in items]
308 +
309 + def reward(
310 + self,
311 + prompts: PromptType | Sequence[PromptType],
312 + /,
313 + *,
314 + pooling_params: PoolingParams | Sequence[PoolingParams] | None = None,
315 + use_tqdm: bool | Callable[..., tqdm] = True,
316 + lora_request: list[LoRARequest] | LoRARequest | None = None,
317 + tokenization_kwargs: dict[str, Any] | None = None,
318 + ) -> list[PoolingRequestOutput]:
319 + """
320 + Generate rewards for each prompt.
321 +
322 + Args:
323 + prompts: The prompts to the LLM. You may pass a sequence of prompts
324 + for batch inference. See [PromptType][vllm.inputs.PromptType]
325 + for more details about the format of each prompt.
326 + pooling_params: The pooling parameters for pooling. If None, we
327 + use the default pooling parameters.
328 + use_tqdm: If `True`, shows a tqdm progress bar.
329 + If a callable (e.g., `functools.partial(tqdm, leave=False)`),
330 + it is used to create the progress bar.
331 + If `False`, no progress bar is created.
332 + lora_request: LoRA request to use for generation, if any.
333 + tokenization_kwargs: Overrides for `tokenizer.encode`.
334 +
335 + Returns:
336 + A list of `PoolingRequestOutput` objects containing the
337 + pooled hidden states in the same order as the input prompts.
338 + """
339 + logger.warning_once(
340 + "`llm.reward` api is deprecated and will be removed in v0.23. "
341 + 'Please use `LLM.encode` with `pooling_task="classify"` or '
342 + '`pooling_task="token_classify"` instead.'
343 + )
344 + return self.encode(
345 + prompts,
346 + use_tqdm=use_tqdm,
347 + lora_request=lora_request,
348 + pooling_params=pooling_params,
349 + pooling_task="token_classify",
350 + tokenization_kwargs=tokenization_kwargs,
351 + )
352 +
353 + def score(
354 + self,
355 + data_1: ScoreInput | list[ScoreInput],
356 + data_2: ScoreInput | list[ScoreInput],
357 + /,
358 + *,
359 + use_tqdm: bool | Callable[..., tqdm] = True,
360 + pooling_params: PoolingParams | None = None,
361 + lora_request: list[LoRARequest] | LoRARequest | None = None,
362 + tokenization_kwargs: dict[str, Any] | None = None,
363 + chat_template: str | None = None,
364 + ) -> list[ScoringRequestOutput]:
365 + """Generate similarity scores for all pairs `<text,text_pair>` or
366 + `<multi-modal data, multi-modal data pair>`.
367 +
368 + The inputs can be `1 -> 1`, `1 -> N` or `N -> N`.
369 + In the `1 - N` case the `data_1` input will be replicated `N`
370 + times to pair with the `data_2` inputs.
371 + The input pairs are used to build a list of prompts for the
372 + cross encoder model. This class automatically batches the prompts,
373 + considering the memory constraint. For the best performance, put all
374 + of your inputs into a single list and pass it to this method.
375 +
376 + Supports both text and multi-modal data (images, etc.) when used with
377 + appropriate multi-modal models. For multi-modal inputs, ensure the
378 + prompt structure matches the model's expected input format.
379 +
380 + Args:
381 + data_1: Can be a single prompt, a list of prompts or
382 + `ScoreMultiModalParam`, which can contain either text or
383 + multi-modal data. When a list, it must have the same length as
384 + the `data_2` list.
385 + data_2: The data to pair with the query to form the input to
386 + the LLM. Can be text or multi-modal data. See [PromptType]
387 + [vllm.inputs.PromptType] for more details about the format of
388 + each prompt.
389 + pooling_params: The pooling parameters for pooling. If None, we
390 + use the default pooling parameters.
391 + use_tqdm: If `True`, shows a tqdm progress bar.
392 + If a callable (e.g., `functools.partial(tqdm, leave=False)`),
393 + it is used to create the progress bar.
394 + If `False`, no progress bar is created.
395 + lora_request: LoRA request to use for generation, if any.
396 + chat_template: The chat template to use for the scoring. If None, we
397 + use the model's default chat template.
398 + tokenization_kwargs: Overrides for `tokenizer.encode`.
399 + Returns:
400 + A list of `ScoringRequestOutput` objects containing the
401 + generated scores in the same order as the input prompts.
402 + """
403 +
404 + if self.runner_type != "pooling":
405 + raise ValueError(
406 + "LLM.score() is only supported for pooling models. "
407 + "Try passing `--runner pooling` to use the model as a "
408 + "pooling model."
409 + )
410 +
411 + score_type: str | None = SCORE_TYPE_MAP.get(self.pooling_task, None) # type: ignore[arg-type]
412 + if (
413 + score_type == "cross-encoder"
414 + and getattr(self.model_config.hf_config, "num_labels", 0) != 1
415 + ):
416 + raise ValueError("Scoring API is only enabled for num_labels == 1.")
417 +
418 + if score_type is None or score_type not in self.pooling_io_processors:
419 + raise ValueError("This model does not support the Scoring API.")
420 +
421 + io_processor = self.pooling_io_processors[score_type]
422 + assert isinstance(io_processor, ScoringIOProcessor)
423 +
424 + pooling_task = io_processor.pooling_task
425 + scoring_data = io_processor.valid_inputs(data_1, data_2)
426 + n_queries = len(scoring_data.data_1)
427 +
428 + if pooling_params is None:
429 + pooling_params = PoolingParams()
430 +
431 + ctx = OfflineInputsContext(
432 + prompts=scoring_data,
433 + pooling_params=pooling_params,
434 + tokenization_kwargs=tokenization_kwargs,
435 + chat_template=chat_template,
436 + n_queries=n_queries,
437 + )
438 +
439 + engine_inputs = io_processor.pre_process_offline(ctx)
440 + n_inputs = len(engine_inputs)
441 +
442 + seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs)
443 + params_seq = self._params_to_seq(ctx.pooling_params, n_inputs)
444 +
445 + for param in params_seq:
446 + if param.task is None:
447 + param.task = pooling_task
448 + elif param.task != pooling_task:
449 + msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!"
450 + raise ValueError(msg)
451 +
452 + seq_priority = self._priority_to_seq(None, n_inputs)
453 +
454 + self._render_and_add_requests(
455 + prompts=engine_inputs,
456 + params=params_seq,
457 + lora_requests=seq_lora_requests,
458 + priorities=seq_priority,
459 + )
460 +
461 + outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput)
462 + outputs = io_processor.post_process_offline(
463 + ctx=OfflineOutputsContext(outputs=outputs, n_queries=n_queries),
464 + )
465 +
466 + return [ScoringRequestOutput.from_base(item) for item in outputs]
467 +
468 + @abstractmethod
469 + def _params_to_seq(
470 + self,
471 + params: _P | Sequence[_P],
472 + num_requests: int,
473 + ) -> Sequence[_P]:
474 + raise NotImplementedError
475 +
476 + @abstractmethod
477 + def _lora_request_to_seq(
478 + self,
479 + lora_request: LoRARequest | None | Sequence[LoRARequest | None],
480 + num_requests: int,
481 + ) -> Sequence[LoRARequest | None]:
482 + raise NotImplementedError
483 +
484 + @abstractmethod
485 + def _priority_to_seq(
486 + self,
487 + priority: list[int] | None,
488 + num_requests: int,
489 + ) -> Sequence[int]:
490 + raise NotImplementedError
491 +
492 + @abstractmethod
493 + def _render_and_add_requests(
494 + self,
495 + prompts: Iterable[EngineInput],
496 + params: Sequence[SamplingParams | PoolingParams],
497 + *,
498 + lora_requests: Sequence[LoRARequest | None] | None = None,
499 + priorities: Sequence[int] | None = None,
500 + ) -> list[str]:
501 + raise NotImplementedError
502 +
503 + @abstractmethod
504 + def _run_engine(
505 + self,
506 + output_type: type[_O] | tuple[type[_O], ...],
507 + *,
508 + use_tqdm: bool | Callable[..., tqdm] = True,
509 + ) -> list[_O]:
510 + raise NotImplementedError