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 14, 05:30 UTC
[Performance] V1 Classify Models E2E Performance Optimization (#23541)
Signed-off-by: wang.yuqi <noooop@126.com>
wang.yuqi committed
Aug 29, 2025 at 18:12 UTC
d9e00dbd1fcf9e4b6b0b42a228d7bb26175cbba4
5 files changed
+80
-37
tests/entrypoints/llm/test_classify.py
+6
@@ -62,3 +62,9 @@ def test_encode_api(llm: LLM):
62
err_msg = "pooling_task must be one of.+"
63
with pytest.raises(ValueError, match=err_msg):
64
llm.encode(prompts, use_tqdm=False)
65
+
66
+
67
+def test_score_api(llm: LLM):
68
+ err_msg = "Score API is only enabled for num_labels == 1."
69
+ with pytest.raises(ValueError, match=err_msg):
70
+ llm.score("ping", "pong", use_tqdm=False)
tests/entrypoints/openai/test_classification.py
+30
@@ -226,3 +226,33 @@ def test_pooling(server: RemoteOpenAIServer, model_name: str):
226
},
227
)
228
assert response.json()["error"]["type"] == "BadRequestError"
229
+
230
+
231
+@pytest.mark.asyncio
232
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
233
+def test_score(server: RemoteOpenAIServer, model_name: str):
234
+ # score api is only enabled for num_labels == 1.
235
+ response = requests.post(
236
+ server.url_for("score"),
237
+ json={
238
+ "model": model_name,
239
+ "text_1": "ping",
240
+ "text_2": "pong",
241
+ },
242
+ )
243
+ assert response.json()["error"]["type"] == "BadRequestError"
244
+
245
+
246
+@pytest.mark.asyncio
247
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
248
+def test_rerank(server: RemoteOpenAIServer, model_name: str):
249
+ # rerank api is only enabled for num_labels == 1.
250
+ response = requests.post(
251
+ server.url_for("rerank"),
252
+ json={
253
+ "model": model_name,
254
+ "query": "ping",
255
+ "documents": ["pong"],
256
+ },
257
+ )
258
+ assert response.json()["error"]["type"] == "BadRequestError"
vllm/entrypoints/openai/api_server.py
+1
-5
@@ -1805,17 +1805,13 @@ async def init_app_state(
1805
request_logger=request_logger,
1806
log_error_stack=args.log_error_stack,
1807
) if "classify" in supported_tasks else None
1808
-
1809
- enable_serving_reranking = ("classify" in supported_tasks and getattr(
1810
- model_config.hf_config, "num_labels", 0) == 1)
1808
state.openai_serving_scores = ServingScores(
1809
engine_client,
1810
model_config,
1811
state.openai_serving_models,
1812
request_logger=request_logger,
1813
log_error_stack=args.log_error_stack,
1817
- ) if ("embed" in supported_tasks or enable_serving_reranking) else None
1818
-
1814
+ ) if ("embed" in supported_tasks or "score" in supported_tasks) else None
1815
state.openai_serving_tokenization = OpenAIServingTokenization(
1816
engine_client,
1817
model_config,
vllm/model_executor/layers/pooler.py
+32
-28
@@ -13,12 +13,15 @@ import torch.nn.functional as F
13
from transformers import PretrainedConfig
14
15
from vllm.config import ModelConfig, PoolerConfig
16
+from vllm.logger import init_logger
17
from vllm.pooling_params import PoolingParams
18
from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
19
from vllm.tasks import PoolingTask
20
from vllm.utils import current_stream, resolve_obj_by_qualname
21
from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata
22
23
+logger = init_logger(__name__)
24
+
25
PoolingFn = Callable[
26
[Union[torch.Tensor, list[torch.Tensor]], PoolingMetadata],
27
Union[torch.Tensor, list[torch.Tensor]]]
@@ -183,7 +186,7 @@ def get_cross_encoder_activation_function(config: PretrainedConfig):
186
fn = resolve_obj_by_qualname(function_name)()
187
return PoolerActivation.wraps(fn)
188
186
- return PoolerScore()
189
+ return PoolerClassify()
190
191
192
def build_output(
@@ -371,22 +374,29 @@ class PoolerMultiLabelClassify(PoolerActivation):
374
375
class PoolerClassify(PoolerActivation):
376
374
- def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
375
- num_labels = pooled_data.shape[-1]
376
- if num_labels < 2:
377
- return F.sigmoid(pooled_data.float()).to(pooled_data.dtype)
378
-
379
- return F.softmax(pooled_data.float(), dim=-1).to(pooled_data.dtype)
380
-
377
+ def __init__(self, *, static_num_labels: bool = True) -> None:
378
+ super().__init__()
379
382
-class PoolerScore(PoolerActivation):
380
+ if static_num_labels:
381
+ from vllm.config import get_current_vllm_config
382
+ vllm_config = get_current_vllm_config()
383
+ self.num_labels = getattr(vllm_config.model_config.hf_config,
384
+ "num_labels", 0)
385
+ if self.num_labels == 0:
386
+ logger.warning("num_labels should be > 0 for classification"
387
+ "models, falling back to softmax. "
388
+ "Please check if the configuration is correct.")
389
+ else:
390
+ self.num_labels = None
391
392
def forward_chunk(self, pooled_data: torch.Tensor) -> torch.Tensor:
385
- num_labels = pooled_data.shape[-1]
393
+ num_labels = (self.num_labels if self.num_labels is not None else
394
+ pooled_data.shape[-1])
395
+
396
if num_labels < 2:
397
return F.sigmoid(pooled_data.float()).to(pooled_data.dtype)
398
389
- return pooled_data
399
+ return F.softmax(pooled_data.float(), dim=-1).to(pooled_data.dtype)
400
401
402
class LambdaPoolerActivation(PoolerActivation):
@@ -428,6 +438,10 @@ class EmbeddingPoolerHead(PoolerHead):
438
def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
439
pooling_metadata: PoolingMetadata):
440
441
+ if isinstance(pooled_data, list):
442
+ pooled_data = torch.stack(pooled_data)
443
+ # pooled_data shape: [batchsize, hidden_dimension]
444
+
445
# Apply ST projector
446
if self.projector is not None:
447
projector = cast(nn.Module, self.projector)
@@ -437,17 +451,11 @@ class EmbeddingPoolerHead(PoolerHead):
451
y = projector(x.to(torch.float32))
452
return y.to(orig_dtype)
453
440
- if isinstance(pooled_data, torch.Tensor):
441
- pooled_data = _proj(pooled_data)
442
- else:
443
- pooled_data = [_proj(t) for t in pooled_data]
454
+ pooled_data = _proj(pooled_data)
455
+ # pooled_data shape: [batchsize, embedding_dimension]
456
457
pooling_params = get_pooling_params(pooling_metadata)
458
447
- if isinstance(pooled_data, list):
448
- pooled_data = torch.stack(pooled_data)
449
- # pooled_data shape: [batchsize, embedding_dimension]
450
-
459
# for matryoshka representation
460
dimensions_list = [
461
pooling_param.dimensions for pooling_param in pooling_params
@@ -477,13 +485,14 @@ class EmbeddingPoolerHead(PoolerHead):
485
for vecs, f in zip(pooled_data, flags)
486
]
487
488
+ # pooled_data shape: [batchsize, embedding_dimension]
489
return pooled_data
490
491
492
class RewardPoolerHead(PoolerHead):
493
494
def __init__(self) -> None:
486
- super().__init__(activation=PoolerClassify())
495
+ super().__init__(activation=PoolerClassify(static_num_labels=False))
496
497
def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
498
pooling_metadata: PoolingMetadata):
@@ -637,19 +646,13 @@ class ClassifierPooler(Pooler):
646
pooling_metadata: PoolingMetadata,
647
) -> PoolerOutput:
648
pooled_data = self.pooling(hidden_states, pooling_metadata)
640
-
649
if isinstance(pooled_data, list):
650
pooled_data = torch.stack(pooled_data)
651
# pooled_data shape: [batchsize, hidden_size]
652
653
if self.classifier is not None:
646
- # apply classifier once on the full batch if possible
647
- if isinstance(pooled_data, torch.Tensor):
648
- pooled_data = self.classifier(pooled_data)
649
- elif len({data.shape for data in pooled_data}) <= 1:
650
- pooled_data = self.classifier(torch.stack(pooled_data))
651
- else:
652
- pooled_data = [self.classifier(data) for data in pooled_data]
654
+ pooled_data = self.classifier(pooled_data)
655
+ # pooled_data shape: [batchsize, num_labels]
656
657
pooling_params = get_pooling_params(pooling_metadata)
658
flags = [p.activation for p in pooling_params]
@@ -662,6 +665,7 @@ class ClassifierPooler(Pooler):
665
for vecs, f in zip(pooled_data, flags)
666
]
667
668
+ # scores shape: [batchsize, num_labels]
669
return build_output(scores)
670
671
vllm/v1/worker/gpu_model_runner.py
+11
-4
@@ -1248,10 +1248,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1248
and "encode" in supported_tasks):
1249
supported_tasks.remove("encode")
1250
1251
- logger.info_once("Chunked prefill is not supported with "
1252
- "encode task which using ALL pooling. "
1253
- "Please turn off chunked prefill by "
1254
- "`--no-enable-chunked-prefill` before using it.")
1251
+ logger.debug_once("Chunked prefill is not supported with "
1252
+ "encode task which using ALL pooling. "
1253
+ "Please turn off chunked prefill by "
1254
+ "`--no-enable-chunked-prefill` before using it.")
1255
+
1256
+ if "score" in supported_tasks:
1257
+ num_labels = getattr(self.model_config.hf_config, "num_labels", 0)
1258
+ if num_labels != 1:
1259
+ supported_tasks.remove("score")
1260
+ logger.debug_once(
1261
+ "Score API is only enabled for num_labels == 1.")
1262
1263
return supported_tasks
1264