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 7, 00:15 UTC
[Bug] Fix pooling model benchmark script (#36300)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Mar 9, 2026 at 11:17 UTC
be292b7c14e08e6e6883d5ebee79240d04814159
2 files changed
+13
-5
vllm/benchmarks/lib/endpoint_request_func.py
+11
@@ -795,6 +795,17 @@ ASYNC_REQUEST_FUNCS: dict[str, RequestFunc] = {
795
"vllm-rerank": async_request_vllm_rerank,
796
}
797
798
+POOLING_BACKENDS = {
799
+ "openai-embeddings",
800
+ "openai-embeddings-chat",
801
+ "openai-embeddings-clip",
802
+ "openai-embeddings-vlm2vec",
803
+ "infinity-embeddings",
804
+ "infinity-embeddings-clip",
805
+ "vllm-pooling",
806
+ "vllm-rerank",
807
+}
808
+
809
OPENAI_COMPATIBLE_BACKENDS = [
810
k
811
for k, v in ASYNC_REQUEST_FUNCS.items()
vllm/benchmarks/serve.py
+2
-5
@@ -45,6 +45,7 @@ from vllm.benchmarks.datasets import SampleRequest, add_dataset_parser, get_samp
45
from vllm.benchmarks.lib.endpoint_request_func import (
46
ASYNC_REQUEST_FUNCS,
47
OPENAI_COMPATIBLE_BACKENDS,
48
+ POOLING_BACKENDS,
49
RequestFuncInput,
50
RequestFuncOutput,
51
)
@@ -1721,11 +1722,7 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
1722
goodput_config_dict = check_goodput_args(args)
1723
1724
backend = args.backend
1724
- task_type = (
1725
- TaskType.POOLING
1726
- if "embeddings" in backend or "rerank" in backend
1727
- else TaskType.GENERATION
1728
- )
1725
+ task_type = TaskType.POOLING if backend in POOLING_BACKENDS else TaskType.GENERATION
1726
1727
# Collect the sampling parameters.
1728
if task_type == TaskType.GENERATION: