@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d2e2e856ad
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 4, 21:15 UTC

[Frontend] Remove frontend pooling multi task support. (#37861)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

wang.yuqi committed Apr 21, 2026 at 20:27 UTC d2e2e856ad4daf4b36bb37acaac06f3f61ac84ce
26 files changed +255 -268
docs/models/pooling_models/README.md
+3 -3
@@ -286,10 +286,10 @@ Pooling models now support token-wise task.
286
287 ### Score task
288
289 -`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
290 -classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
289 +`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
290 +equal to 1 can it be used as a scoring model and have its scoring API enabled.
291
292 ### Pooling multitask support
293
294 -Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
294 +Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
295 you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.
examples/pooling/token_embed/jina_embeddings_v4_offline.py
+57 -51
@@ -4,68 +4,74 @@
4 import torch
5
6 from vllm import LLM
7 +from vllm.config import PoolerConfig
8 from vllm.inputs import TextPrompt
9 from vllm.multimodal.utils import fetch_image
10
10 -# Initialize model
11 -model = LLM(
12 - model="jinaai/jina-embeddings-v4-vllm-text-matching",
13 - runner="pooling",
14 - max_model_len=1024,
15 - gpu_memory_utilization=0.8,
16 -)
11
18 -# Create text prompts
19 -text1 = "Ein wunderschöner Sonnenuntergang am Strand"
20 -text1_prompt = TextPrompt(prompt=f"Query: {text1}")
12 +def main():
13 + # Initialize model
14 + model = LLM(
15 + model="jinaai/jina-embeddings-v4-vllm-text-matching",
16 + pooler_config=PoolerConfig(task="token_embed"),
17 + runner="pooling",
18 + max_model_len=1024,
19 + gpu_memory_utilization=0.8,
20 + )
21
22 -text2 = "浜辺に沈む美しい夕日"
23 -text2_prompt = TextPrompt(prompt=f"Query: {text2}")
22 + # Create text prompts
23 + text1 = "Ein wunderschöner Sonnenuntergang am Strand"
24 + text1_prompt = TextPrompt(prompt=f"Query: {text1}")
25
25 -# Create image prompt
26 -image = fetch_image(
27 - "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
28 -)
29 -image_prompt = TextPrompt(
30 - prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n", # noqa: E501
31 - multi_modal_data={"image": image},
32 -)
26 + text2 = "浜辺に沈む美しい夕日"
27 + text2_prompt = TextPrompt(prompt=f"Query: {text2}")
28
34 -# Encode all prompts
35 -prompts = [text1_prompt, text2_prompt, image_prompt]
36 -outputs = model.encode(prompts, pooling_task="token_embed")
29 + # Create image prompt
30 + image = fetch_image(
31 + "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
32 + )
33 + image_prompt = TextPrompt(
34 + prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n", # noqa: E501
35 + multi_modal_data={"image": image},
36 + )
37
38 + # Encode all prompts
39 + prompts = [text1_prompt, text2_prompt, image_prompt]
40 + outputs = model.encode(prompts, pooling_task="token_embed")
41
39 -def get_embeddings(outputs):
40 - VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
42 + def get_embeddings(outputs):
43 + VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
44
42 - embeddings = []
43 - for output in outputs:
44 - if VISION_START_TOKEN_ID in output.prompt_token_ids:
45 - # Gather only vision tokens
46 - img_start_pos = torch.where(
47 - torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
48 - )[0][0]
49 - img_end_pos = torch.where(
50 - torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
51 - )[0][0]
52 - embeddings_tensor = output.outputs.data.detach().clone()[
53 - img_start_pos : img_end_pos + 1
54 - ]
55 - else:
56 - # Use all tokens for text-only prompts
57 - embeddings_tensor = output.outputs.data.detach().clone()
45 + embeddings = []
46 + for output in outputs:
47 + if VISION_START_TOKEN_ID in output.prompt_token_ids:
48 + # Gather only vision tokens
49 + img_start_pos = torch.where(
50 + torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
51 + )[0][0]
52 + img_end_pos = torch.where(
53 + torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
54 + )[0][0]
55 + embeddings_tensor = output.outputs.data.detach().clone()[
56 + img_start_pos : img_end_pos + 1
57 + ]
58 + else:
59 + # Use all tokens for text-only prompts
60 + embeddings_tensor = output.outputs.data.detach().clone()
61
59 - # Pool and normalize embeddings
60 - pooled_output = (
61 - embeddings_tensor.sum(dim=0, dtype=torch.float32)
62 - / embeddings_tensor.shape[0]
63 - )
64 - embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
65 - return embeddings
62 + # Pool and normalize embeddings
63 + pooled_output = (
64 + embeddings_tensor.sum(dim=0, dtype=torch.float32)
65 + / embeddings_tensor.shape[0]
66 + )
67 + embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
68 + return embeddings
69
70 + embeddings = get_embeddings(outputs)
71
68 -embeddings = get_embeddings(outputs)
72 + for embedding in embeddings:
73 + print(embedding.shape)
74
70 -for embedding in embeddings:
71 - print(embedding.shape)
75 +
76 +if __name__ == "__main__":
77 + main()
examples/pooling/token_embed/multi_vector_retrieval_offline.py
+16 -9
@@ -4,6 +4,7 @@
4 from argparse import Namespace
5
6 from vllm import LLM, EngineArgs
7 +from vllm.config import PoolerConfig
8 from vllm.utils.argparse_utils import FlexibleArgumentParser
9
10
@@ -13,6 +14,7 @@ def parse_args():
14 # Set example specific arguments
15 parser.set_defaults(
16 model="BAAI/bge-m3",
17 + pooler_config=PoolerConfig(task="token_embed"),
18 runner="pooling",
19 enforce_eager=True,
20 )
@@ -32,15 +34,6 @@ def main(args: Namespace):
34 # You should pass runner="pooling" for embedding models
35 llm = LLM(**vars(args))
36
35 - # Generate embedding. The output is a list of EmbeddingRequestOutputs.
36 - outputs = llm.embed(prompts)
37 -
38 - # Print the outputs.
39 - print("\nGenerated Outputs:\n" + "-" * 60)
40 - for prompt, output in zip(prompts, outputs):
41 - embeds = output.outputs.embedding
42 - print(len(embeds))
43 -
37 # Generate embedding for each token. The output is a list of PoolingRequestOutput.
38 outputs = llm.encode(prompts, pooling_task="token_embed")
39
@@ -50,6 +43,20 @@ def main(args: Namespace):
43 multi_vector = output.outputs.data
44 print(multi_vector.shape)
45
46 + query = "What is the capital of France?"
47 + documents = [
48 + "The capital of Brazil is Brasilia.",
49 + "The capital of France is Paris.",
50 + ]
51 + # Generate scores.
52 + outputs = llm.score(query, documents)
53 + # Print the outputs.
54 + print("\nGenerated Outputs:\n" + "-" * 60)
55 + for document, output in zip(documents, outputs):
56 + score = output.outputs.score
57 + print(f"Pair: {[query, document]!r} \nScore: {score}")
58 + print("-" * 60)
59 +
60
61 if __name__ == "__main__":
62 args = parse_args()
examples/pooling/token_embed/multi_vector_retrieval_online.py
+17 -3
@@ -7,10 +7,11 @@ Example online usage of Pooling API for multi vector retrieval.
7 Run `vllm serve <model> --runner pooling`
8 to start up the server in vLLM. e.g.
9
10 -vllm serve BAAI/bge-m3
10 +vllm serve BAAI/bge-m3 --pooler-config.task token_embed
11 """
12
13 import argparse
14 +import pprint
15
16 import requests
17 import torch
@@ -32,7 +33,8 @@ def parse_args():
33
34
35 def main(args):
35 - api_url = f"http://{args.host}:{args.port}/pooling"
36 + pooling_url = f"http://{args.host}:{args.port}/pooling"
37 + score_url = f"http://{args.host}:{args.port}/score"
38 model_name = args.model
39
40 prompts = [
@@ -43,11 +45,23 @@ def main(args):
45 ]
46 prompt = {"model": model_name, "input": prompts}
47
46 - pooling_response = post_http_request(prompt=prompt, api_url=api_url)
48 + pooling_response = post_http_request(prompt=prompt, api_url=pooling_url)
49 for output in pooling_response.json()["data"]:
50 multi_vector = torch.tensor(output["data"])
51 print(multi_vector.shape)
52
53 + queries = "What is the capital of France?"
54 + documents = [
55 + "The capital of Brazil is Brasilia.",
56 + "The capital of France is Paris.",
57 + ]
58 + prompt = {"model": model_name, "queries": queries, "documents": documents}
59 + score_response = post_http_request(prompt=prompt, api_url=score_url)
60 + print("\nPrompt when queries is string and documents is a list:")
61 + pprint.pprint(prompt)
62 + print("\nScore Response:")
63 + pprint.pprint(score_response.json())
64 +
65
66 if __name__ == "__main__":
67 args = parse_args()
tests/entrypoints/pooling/classify/test_offline.py
+4 -15
@@ -1,13 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import logging
3 import weakref
4
5 import pytest
6 import torch
7
8 from tests.models.utils import softmax
10 -from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
9 +from vllm import LLM, ClassificationRequestOutput, PoolingParams
10 from vllm.distributed import cleanup_dist_env_and_memory
11 from vllm.tasks import PoolingTask
12
@@ -66,18 +65,6 @@ def test_list_prompts(llm: LLM):
65 assert len(outputs[i].outputs.probs) == num_labels
66
67
69 -@pytest.mark.skip_global_cleanup
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
78 - assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
79 -
80 -
68 @pytest.mark.skip_global_cleanup
69 def test_pooling_params(llm: LLM):
70 def get_outputs(use_activation):
@@ -110,10 +97,12 @@ def test_score_api(llm: LLM):
97 llm.score("ping", "pong", use_tqdm=False)
98
99
113 -@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
100 +@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
101 def test_unsupported_tasks(llm: LLM, task: PoolingTask):
102 if task == "plugin":
103 err_msg = "No IOProcessor plugin installed."
104 + elif task == "token_classify":
105 + err_msg = "Try switching the model's pooling_task via.+"
106 else:
107 err_msg = "Embedding API is not supported by this model.+"
108 with pytest.raises(ValueError, match=err_msg):
tests/entrypoints/pooling/classify/test_online.py
+4 -20
@@ -436,26 +436,7 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
436
437 @pytest.mark.asyncio
438 @pytest.mark.parametrize("model_name", [MODEL_NAME])
439 -async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
440 - task = "token_classify"
441 - response = requests.post(
442 - server.url_for("pooling"),
443 - json={
444 - "model": model_name,
445 - "input": input_text,
446 - "encoding_format": "float",
447 - "task": task,
448 - },
449 - )
450 - poolings = PoolingResponse.model_validate(response.json())
451 - assert len(poolings.data) == 1
452 - assert len(poolings.data[0].data) == 8
453 - assert len(poolings.data[0].data[0]) == 2
454 -
455 -
456 -@pytest.mark.asyncio
457 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
458 -@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
439 +@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
440 async def test_pooling_not_supported(
441 server: RemoteOpenAIServer, model_name: str, task: str
442 ):
@@ -469,8 +450,11 @@ async def test_pooling_not_supported(
450 },
451 )
452 assert response.json()["error"]["type"] == "BadRequestError"
453 +
454 if task == "plugin":
455 err_msg = "No IOProcessor plugin installed."
456 + elif task == "token_classify":
457 + err_msg = "Try switching the model's pooling_task via"
458 else:
459 err_msg = f"Unsupported task: {task!r}"
460 assert response.json()["error"]["message"].startswith(err_msg)
tests/entrypoints/pooling/embed/test_offline.py
+6 -13
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import logging
3 import weakref
4
5 import pytest
@@ -38,11 +37,11 @@ def llm():
37 seed=0,
38 attention_config=attention_config,
39 )
40 + assert embedding_size == llm.model_config.embedding_size
41
42 yield weakref.proxy(llm)
43
44 del llm
45 -
45 cleanup_dist_env_and_memory()
46
47
@@ -74,16 +73,6 @@ def test_list_prompts(llm: LLM):
73 assert len(outputs[i].outputs.embedding) == embedding_size
74
75
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 -
76 @pytest.mark.skip_global_cleanup
77 def test_pooling_params(llm: LLM):
78 def get_outputs(normalize):
@@ -107,10 +96,14 @@ def test_pooling_params(llm: LLM):
96 )
97
98
110 -@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"])
99 +@pytest.mark.parametrize(
100 + "task", ["token_classify", "classify", "token_embed", "plugin"]
101 +)
102 def test_unsupported_tasks(llm: LLM, task: PoolingTask):
103 if task == "plugin":
104 err_msg = "No IOProcessor plugin installed."
105 + elif task == "token_embed":
106 + err_msg = "Try switching the model's pooling_task via.+"
107 else:
108 err_msg = "Classification API is not supported by this model.+"
109 with pytest.raises(ValueError, match=err_msg):
tests/entrypoints/pooling/embed/test_online.py
+5 -22
@@ -732,28 +732,9 @@ async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
732
733 @pytest.mark.asyncio
734 @pytest.mark.parametrize("model_name", [MODEL_NAME])
735 -async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
736 - task = "token_embed"
737 - response = requests.post(
738 - server.url_for("pooling"),
739 - json={
740 - "model": model_name,
741 - "input": input_text,
742 - "encoding_format": "float",
743 - "task": task,
744 - },
745 - )
746 -
747 - poolings = PoolingResponse.model_validate(response.json())
748 -
749 - assert len(poolings.data) == 1
750 - assert len(poolings.data[0].data) == len(input_tokens)
751 - assert len(poolings.data[0].data[0]) == 384
752 -
753 -
754 -@pytest.mark.asyncio
755 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
756 -@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
735 +@pytest.mark.parametrize(
736 + "task", ["classify", "token_classify", "token_embed", "plugin"]
737 +)
738 async def test_pooling_not_supported(
739 server: RemoteOpenAIServer, model_name: str, task: str
740 ):
@@ -769,6 +750,8 @@ async def test_pooling_not_supported(
750 assert response.json()["error"]["type"] == "BadRequestError"
751 if task == "plugin":
752 err_msg = "No IOProcessor plugin installed."
753 + elif task == "token_embed":
754 + err_msg = "Try switching the model's pooling_task via"
755 else:
756 err_msg = f"Unsupported task: {task!r}"
757 assert response.json()["error"]["message"].startswith(err_msg)
tests/entrypoints/pooling/scoring/test_cross_encoder_online.py
+3 -20
@@ -452,25 +452,6 @@ async def test_pooling_classify(server: RemoteOpenAIServer):
452 assert len(poolings.data[0].data) == 1
453
454
455 -@pytest.mark.asyncio
456 -async def test_pooling_token_classify(server: RemoteOpenAIServer):
457 - response = requests.post(
458 - server.url_for("pooling"),
459 - json={
460 - "model": MODEL_NAME,
461 - "task": "token_classify",
462 - "input": input_text,
463 - "encoding_format": "float",
464 - },
465 - )
466 -
467 - poolings = PoolingResponse.model_validate(response.json())
468 -
469 - assert len(poolings.data) == 1
470 - assert len(poolings.data[0].data) == len(input_tokens)
471 - assert len(poolings.data[0].data[0]) == 1
472 -
473 -
455 @pytest.mark.asyncio
456 async def test_rerank_max_tokens_per_doc(
457 server: RemoteOpenAIServer,
@@ -544,7 +525,7 @@ async def test_rerank_max_tokens_per_doc_validation(
525
526
527 @pytest.mark.asyncio
547 -@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
528 +@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
529 async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
530 response = requests.post(
531 server.url_for("pooling"),
@@ -558,6 +539,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
539 assert response.json()["error"]["type"] == "BadRequestError"
540 if task == "plugin":
541 err_msg = "No IOProcessor plugin installed."
542 + elif task == "token_classify":
543 + err_msg = "Try switching the model's pooling_task via"
544 else:
545 err_msg = f"Unsupported task: {task!r}"
546 assert response.json()["error"]["message"].startswith(err_msg)
tests/entrypoints/pooling/token_classify/test_offline.py
+8 -12
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import logging
3 import weakref
4
5 import pytest
@@ -60,22 +59,19 @@ def test_token_ids_prompts(llm: LLM):
59
60 @pytest.mark.skip_global_cleanup
61 def test_score_api(llm: LLM):
63 - err_msg = "Scoring API is only enabled for num_labels == 1."
62 + err_msg = "This model does not support the Scoring API."
63 with pytest.raises(ValueError, match=err_msg):
64 llm.score("ping", "pong", use_tqdm=False)
65
66
67 @pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
68 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
69 + if task == "plugin":
70 + err_msg = "No IOProcessor plugin installed."
71 + elif task == "classify":
72 + err_msg = "Try switching the model's pooling_task via.+"
73 else:
75 - if task == "plugin":
76 - err_msg = "No IOProcessor plugin installed."
77 - else:
78 - err_msg = "Embedding API is not supported by this model.+"
74 + err_msg = "Embedding API is not supported by this model.+"
75
80 - with pytest.raises(ValueError, match=err_msg):
81 - llm.encode(prompt, pooling_task=task, use_tqdm=False)
76 + with pytest.raises(ValueError, match=err_msg):
77 + llm.encode(prompt, pooling_task=task, use_tqdm=False)
tests/entrypoints/pooling/token_classify/test_online.py
+4 -1
@@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st
50
51 @pytest.mark.asyncio
52 @pytest.mark.parametrize("model_name", [MODEL_NAME])
53 -@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
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 ):
@@ -63,9 +63,12 @@ async def test_pooling_not_supported(
63 "task": task,
64 },
65 )
66 + assert response.json()["error"]["type"] == "BadRequestError"
67
68 if task == "plugin":
69 err_msg = "No IOProcessor plugin installed."
70 + elif task == "classify":
71 + err_msg = "Try switching the model's pooling_task via"
72 else:
73 err_msg = f"Unsupported task: {task!r}"
74 assert response.json()["error"]["message"].startswith(err_msg)
tests/entrypoints/pooling/token_embed/test_offline.py
+7 -11
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import logging
3 import weakref
4
5 import pytest
@@ -64,15 +63,12 @@ def test_token_ids_prompts(llm: LLM):
63
64 @pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
65 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
66 + if task == "plugin":
67 + err_msg = "No IOProcessor plugin installed."
68 + elif task == "embed":
69 + err_msg = "Try switching the model's pooling_task via.+"
70 else:
72 - if task == "plugin":
73 - err_msg = "No IOProcessor plugin installed."
74 - else:
75 - err_msg = "Classification API is not supported by this model.+"
71 + err_msg = "Classification API is not supported by this model.+"
72
77 - with pytest.raises(ValueError, match=err_msg):
78 - llm.encode(prompt, pooling_task=task, use_tqdm=False)
73 + with pytest.raises(ValueError, match=err_msg):
74 + llm.encode(prompt, pooling_task=task, use_tqdm=False)
tests/entrypoints/pooling/token_embed/test_online.py
+4 -1
@@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
73
74 @pytest.mark.asyncio
75 @pytest.mark.parametrize("model_name", [MODEL_NAME])
76 -@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
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 ):
@@ -86,9 +86,12 @@ async def test_pooling_not_supported(
86 "task": task,
87 },
88 )
89 + assert response.json()["error"]["type"] == "BadRequestError"
90
91 if task == "plugin":
92 err_msg = "No IOProcessor plugin installed."
93 + elif task == "embed":
94 + err_msg = "Try switching the model's pooling_task via"
95 else:
96 err_msg = f"Unsupported task: {task!r}"
97 assert response.json()["error"]["message"].startswith(err_msg)
tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py
+2
@@ -6,6 +6,7 @@ from transformers import AutoModel
6
7 from tests.models.utils import check_embeddings_close
8 from vllm import TokensPrompt
9 +from vllm.config import PoolerConfig
10
11
12 @pytest.mark.parametrize(
@@ -21,6 +22,7 @@ def test_embed_models(hf_runner, vllm_runner, model: str):
22 with vllm_runner(
23 model,
24 runner="pooling",
25 + pooler_config=PoolerConfig(task="token_embed"),
26 max_model_len=128,
27 max_num_batched_tokens=chunk_size,
28 enforce_eager=True,
tests/models/language/pooling/test_bge_m3.py
+46 -14
@@ -3,7 +3,6 @@
3 import httpx
4 import openai
5 import pytest
6 -import pytest_asyncio
6 import torch
7
8 from ....utils import RemoteOpenAIServer
@@ -25,29 +24,42 @@ sentences_2 = [
24 similarity_reference = [[0.6259, 0.3474], [0.3309, 0.6734]]
25 lexical_score_reference = [0.19554901123046875, 0.0]
26 colbert_score_reference = [0.7797, 0.4620]
27 +SUPPORTED_TASKS = ["embed", "token_embed", "token_classify"]
28 +
29 +
30 +@pytest.fixture(scope="module", params=SUPPORTED_TASKS)
31 +def pooling_task(request):
32 + yield request.param
33
34
35 @pytest.fixture(scope="module")
31 -def server():
36 +def server(pooling_task):
37 args = [
38 "--max-model-len",
39 str(MAX_MODEL_LEN),
40 "--hf-overrides",
41 '{"architectures": ["BgeM3EmbeddingModel"]}',
42 + "--pooler-config.task",
43 + pooling_task,
44 ]
45
46 with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
47 yield remote_server
48
49
43 -@pytest_asyncio.fixture
44 -async def client(server):
45 - async with server.get_async_client() as async_client:
46 - yield async_client
47 -
48 -
50 @pytest.mark.asyncio
50 -async def test_bge_m3_api_server_embedding(client: openai.AsyncOpenAI):
51 +async def test_bge_m3_api_server_embedding(server, pooling_task):
52 + client = server.get_async_client()
53 +
54 + if pooling_task != "embed":
55 + with pytest.raises(openai.InternalServerError):
56 + await run_client_embeddings(
57 + client,
58 + MODEL_NAME,
59 + sentences_1,
60 + )
61 + return
62 +
63 embeddings_list_1 = await run_client_embeddings(
64 client,
65 MODEL_NAME,
@@ -117,7 +129,14 @@ def compute_lexical_matching_score(
129
130
131 @pytest.mark.asyncio
120 -async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
132 +async def test_bge_m3_api_server_sparse_embedding(server, pooling_task):
133 + client = server.get_async_client()
134 +
135 + if pooling_task != "token_classify":
136 + with pytest.raises(openai.BadRequestError):
137 + await sparse_embeddings(client, sentences_1)
138 + return
139 +
140 embeddings_1 = await sparse_embeddings(client, sentences_1)
141 embeddings_2 = await sparse_embeddings(client, sentences_2)
142
@@ -137,9 +156,11 @@ async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
156
157
158 @pytest.mark.asyncio
140 -async def test_bge_m3_api_server_sparse_embedding_corner_case(
141 - client: openai.AsyncOpenAI,
142 -):
159 +async def test_bge_m3_api_server_sparse_embedding_corner_case(server, pooling_task):
160 + if pooling_task != "token_classify":
161 + return
162 +
163 + client = server.get_async_client()
164 embeddings = await sparse_embeddings(client, ["Hi"])
165 assert len(embeddings) == 1
166 assert 2673 in embeddings[0]
@@ -155,7 +176,18 @@ def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:
176
177
178 @pytest.mark.asyncio
158 -async def test_bge_m3_api_server_multi_vector(client: openai.AsyncOpenAI):
179 +async def test_bge_m3_api_server_multi_vector(server, pooling_task):
180 + client = server.get_async_client()
181 +
182 + if pooling_task != "token_embed":
183 + with pytest.raises(openai.BadRequestError):
184 + await client.post(
185 + "../pooling",
186 + body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
187 + cast_to=httpx.Response,
188 + )
189 + return
190 +
191 result_1 = await client.post(
192 "../pooling",
193 body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
tests/models/language/pooling/test_extract_hidden_states.py
+2 -11
@@ -4,6 +4,7 @@ import pytest
4 import torch
5
6 from vllm import TokensPrompt
7 +from vllm.config import PoolerConfig
8
9
10 @pytest.mark.parametrize(
@@ -20,6 +21,7 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
21 max_model_len=128,
22 enforce_eager=True,
23 runner="pooling",
24 + pooler_config=PoolerConfig(task="token_embed"),
25 enable_prefix_caching=True,
26 ) as vllm_model:
27 pooling_outputs = vllm_model.llm.encode(
@@ -44,14 +46,3 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
46 assert len(output.prompt_token_ids) == n
47 assert len(output.outputs.data) == n
48 assert output.num_cached_tokens == 0
47 -
48 - # skip_reading_prefix_cache can still write to cache
49 - # to accelerate following requests
50 - pooling_outputs = vllm_model.llm.encode(
51 - [TokensPrompt(prompt_token_ids=t) for t in token_prompts],
52 - pooling_task="embed",
53 - )
54 -
55 - for n, output in zip(n_prompt_tokens, pooling_outputs):
56 - assert len(output.prompt_token_ids) == n
57 - assert output.num_cached_tokens > 0
tests/models/language/pooling/test_multi_vector_retrieval.py
+2
@@ -5,6 +5,7 @@ import torch
5 from transformers import AutoModel
6
7 from tests.models.utils import check_embeddings_close
8 +from vllm.config import PoolerConfig
9
10
11 @pytest.mark.parametrize(
@@ -17,6 +18,7 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype
18 with vllm_runner(
19 model,
20 runner="pooling",
21 + pooler_config=PoolerConfig(task="token_embed"),
22 max_model_len=None,
23 ) as vllm_model:
24 vllm_outputs = vllm_model.token_embed(example_prompts)
tests/models/language/pooling/test_pooler_config_init_behaviour.py
+2 -2
@@ -146,7 +146,7 @@ def test_multi_vector_retrieval_models_using_normalize(
146 model,
147 max_model_len=512,
148 dtype=dtype,
149 - pooler_config=PoolerConfig(use_activation=False),
149 + pooler_config=PoolerConfig(use_activation=False, task="token_embed"),
150 ) as vllm_model:
151 wo_normalize = vllm_model.token_embed(example_prompts)
152
@@ -154,7 +154,7 @@ def test_multi_vector_retrieval_models_using_normalize(
154 model,
155 max_model_len=512,
156 dtype=dtype,
157 - pooler_config=PoolerConfig(use_activation=True),
157 + pooler_config=PoolerConfig(use_activation=True, task="token_embed"),
158 ) as vllm_model:
159 w_normalize = vllm_model.token_embed(example_prompts)
160
vllm/entrypoints/llm.py
+5 -8
@@ -79,7 +79,7 @@ from vllm.renderers.inputs.preprocess import (
79 prompt_to_seq,
80 )
81 from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
82 -from vllm.tasks import PoolingTask
82 +from vllm.tasks import SCORE_TYPE_MAP, PoolingTask
83 from vllm.tokenizers import TokenizerLike
84 from vllm.usage.usage_lib import UsageContext
85 from vllm.utils.counter import Counter
@@ -1207,12 +1207,9 @@ class LLM:
1207 f"Supported tasks: {self.supported_tasks}"
1208 )
1209 else:
1210 - logger.warning_once(
1211 - "Pooling multitask support is deprecated and will "
1212 - "be removed in v0.20. When the default pooling task is "
1213 - "not what you want, you need to manually specify it "
1214 - 'via PoolerConfig(task="%s"). ',
1215 - pooling_task,
1210 + raise ValueError(
1211 + f"Try switching the model's pooling_task "
1212 + f'via `PoolerConfig(task="{pooling_task}")`'
1213 )
1214
1215 if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors:
@@ -1410,7 +1407,7 @@ class LLM:
1407 "pooling model."
1408 )
1409
1413 - score_type = self.model_config.score_type
1410 + score_type: str | None = SCORE_TYPE_MAP.get(self.pooling_task, None) # type: ignore[arg-type]
1411 if (
1412 score_type == "cross-encoder"
1413 and getattr(self.model_config.hf_config, "num_labels", 0) != 1
vllm/entrypoints/pooling/base/serving.py
+3 -12
@@ -15,10 +15,7 @@ from starlette.datastructures import Headers
15 from vllm import PoolingParams, PoolingRequestOutput, envs
16 from vllm.config import VllmConfig
17 from vllm.engine.protocol import EngineClient
18 -from vllm.entrypoints.chat_utils import (
19 - ChatTemplateConfig,
20 - ChatTemplateContentFormatOption,
21 -)
18 +from vllm.entrypoints.chat_utils import ChatTemplateConfig
19 from vllm.entrypoints.logger import RequestLogger
20 from vllm.entrypoints.openai.engine.protocol import ErrorResponse
21 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
@@ -48,9 +45,7 @@ class PoolingServingBase(ABC):
45 models: OpenAIServingModels,
46 *,
47 request_logger: RequestLogger | None,
51 - chat_template: str | None = None,
52 - chat_template_content_format: ChatTemplateContentFormatOption = "auto",
53 - trust_request_chat_template: bool = False,
48 + chat_template_config: ChatTemplateConfig,
49 return_tokens_as_token_ids: bool = False,
50 log_error_stack: bool = False,
51 ):
@@ -63,11 +58,7 @@ class PoolingServingBase(ABC):
58 self.request_logger = request_logger
59 self.return_tokens_as_token_ids = return_tokens_as_token_ids
60 self.log_error_stack = log_error_stack
66 - self.chat_template_config = ChatTemplateConfig(
67 - chat_template=chat_template,
68 - chat_template_content_format=chat_template_content_format,
69 - trust_request_chat_template=trust_request_chat_template,
70 - )
61 + self.chat_template_config = chat_template_config
62
63 # Shared thread pool executor for preprocessing and postprocessing.
64 self._executor: Executor = models.renderer._executor
vllm/entrypoints/pooling/factories.py
+33 -24
@@ -10,7 +10,7 @@ from vllm.entrypoints.chat_utils import ChatTemplateConfig
10 from vllm.logger import init_logger
11 from vllm.plugins.io_processors import has_io_processor
12 from vllm.renderers import BaseRenderer
13 -from vllm.tasks import POOLING_TASKS, SupportedTask
13 +from vllm.tasks import POOLING_TASKS, SCORE_TYPE_MAP, SupportedTask
14
15 from .base.io_processor import PoolingIOProcessor
16 from .utils import enable_scoring_api
@@ -43,23 +43,24 @@ def init_pooling_io_processors(
43 ) -> dict[str, PoolingIOProcessor]:
44 model_config = vllm_config.model_config
45 processors: dict[str, type[PoolingIOProcessor]] = {}
46 + pooling_task = model_config.get_pooling_task(supported_tasks)
47
47 - if "classify" in supported_tasks:
48 + if pooling_task == "classify":
49 from .classify.io_processor import ClassifyIOProcessor
50
51 processors["classify"] = ClassifyIOProcessor
52
52 - if "token_classify" in supported_tasks:
53 + if pooling_task == "token_classify":
54 from .classify.io_processor import TokenClassifyIOProcessor
55
56 processors["token_classify"] = TokenClassifyIOProcessor
57
57 - if "embed" in supported_tasks:
58 + if pooling_task == "embed":
59 from .embed.io_processor import EmbedIOProcessor
60
61 processors["embed"] = EmbedIOProcessor
62
62 - if "token_embed" in supported_tasks:
63 + if pooling_task == "token_embed":
64 from .embed.io_processor import TokenEmbedIOProcessor
65
66 processors["token_embed"] = TokenEmbedIOProcessor
@@ -71,15 +72,15 @@ def init_pooling_io_processors(
72 from .pooling.io_processor import PluginWithIOProcessorPlugins
73
74 processors["plugin"] = PluginWithIOProcessorPlugins
74 - elif "plugin" in supported_tasks:
75 + elif pooling_task == "plugin":
76 from .pooling.io_processor import PluginWithoutIOProcessorPlugins
77
78 processors["plugin"] = PluginWithoutIOProcessorPlugins
79
80 if enable_scoring_api(supported_tasks, model_config):
80 - score_type = model_config.score_type
81 from .scoring.io_processor import ScoringIOProcessors
82
83 + score_type: str | None = SCORE_TYPE_MAP.get(pooling_task, None) # type: ignore[arg-type]
84 if score_type is not None and score_type in ScoringIOProcessors:
85 processors[score_type] = ScoringIOProcessors[score_type]
86
@@ -140,6 +141,10 @@ def init_pooling_state(
141 request_logger: RequestLogger | None,
142 supported_tasks: tuple["SupportedTask", ...],
143 ):
144 + model_config = engine_client.model_config
145 + if model_config is None:
146 + return
147 +
148 from vllm.entrypoints.chat_utils import load_chat_template
149 from vllm.tasks import POOLING_TASKS
150
@@ -148,8 +153,14 @@ def init_pooling_state(
153 from .pooling.serving import ServingPooling
154 from .scoring.serving import ServingScores
155
151 - model_config = engine_client.model_config
156 resolved_chat_template = load_chat_template(args.chat_template)
157 + pooling_task = model_config.get_pooling_task(supported_tasks)
158 +
159 + chat_template_config = ChatTemplateConfig(
160 + chat_template=resolved_chat_template,
161 + chat_template_content_format=args.chat_template_content_format,
162 + trust_request_chat_template=args.trust_request_chat_template,
163 + )
164
165 state.serving_pooling = (
166 (
@@ -158,9 +169,7 @@ def init_pooling_state(
169 state.openai_serving_models,
170 supported_tasks=supported_tasks,
171 request_logger=request_logger,
161 - chat_template=resolved_chat_template,
162 - chat_template_content_format=args.chat_template_content_format,
163 - trust_request_chat_template=args.trust_request_chat_template,
172 + chat_template_config=chat_template_config,
173 )
174 )
175 if any(t in supported_tasks for t in POOLING_TASKS)
@@ -171,11 +180,9 @@ def init_pooling_state(
180 engine_client,
181 state.openai_serving_models,
182 request_logger=request_logger,
174 - chat_template=resolved_chat_template,
175 - chat_template_content_format=args.chat_template_content_format,
176 - trust_request_chat_template=args.trust_request_chat_template,
183 + chat_template_config=chat_template_config,
184 )
178 - if "embed" in supported_tasks
185 + if pooling_task == "embed"
186 else None
187 )
188 state.serving_classification = (
@@ -183,21 +190,18 @@ def init_pooling_state(
190 engine_client,
191 state.openai_serving_models,
192 request_logger=request_logger,
186 - chat_template=resolved_chat_template,
187 - chat_template_content_format=args.chat_template_content_format,
188 - trust_request_chat_template=args.trust_request_chat_template,
193 + chat_template_config=chat_template_config,
194 )
190 - if "classify" in supported_tasks
195 + if pooling_task == "classify"
196 else None
197 )
198 state.serving_scores = (
199 ServingScores(
200 engine_client,
201 state.openai_serving_models,
202 + supported_tasks=supported_tasks,
203 request_logger=request_logger,
198 - chat_template=resolved_chat_template,
199 - chat_template_content_format=args.chat_template_content_format,
200 - trust_request_chat_template=args.trust_request_chat_template,
204 + chat_template_config=chat_template_config,
205 enable_flash_late_interaction=getattr(
206 args, "enable_flash_late_interaction", True
207 ),
@@ -214,7 +218,12 @@ def get_pooling_invocation_types(
218 # NOTE: Items defined earlier take higher priority
219 invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []
220
217 - if "embed" in supported_tasks:
221 + if model_config is None:
222 + return invocation_types
223 +
224 + pooling_task = model_config.get_pooling_task(supported_tasks)
225 +
226 + if pooling_task == "embed":
227 from .embed.api_router import create_embedding, embedding
228 from .embed.protocol import EmbeddingRequest
229
@@ -222,7 +231,7 @@ def get_pooling_invocation_types(
231 (EmbeddingRequest, (embedding, create_embedding)),
232 ]
233
225 - if "classify" in supported_tasks:
234 + if pooling_task == "classify":
235 from .classify.api_router import classify, create_classify
236 from .classify.protocol import ClassificationRequest
237
vllm/entrypoints/pooling/pooling/serving.py
+4 -6
@@ -78,17 +78,15 @@ class ServingPooling(PoolingServingBase):
78
79 # plugin task uses io_processor.parse_request to verify inputs
80 if pooling_task != "plugin" and pooling_task != self.pooling_task:
81 - if pooling_task not in self.io_processors:
81 + if pooling_task not in self.supported_tasks:
82 raise ValueError(
83 f"Unsupported task: {pooling_task!r} "
84 f"Supported tasks: {self.supported_tasks}"
85 )
86 else:
87 - logger.warning_once(
88 - "Pooling multitask support is deprecated and will be removed "
89 - "in v0.20. When the default pooling task is not what you want, you "
90 - "need to manually specify it via --pooler-config.task %s. ",
91 - pooling_task,
87 + raise ValueError(
88 + "Try switching the model's pooling_task "
89 + f"via --pooler-config.task {request.task}."
90 )
91
92 if pooling_task == "plugin" and "plugin" not in self.io_processors:
vllm/entrypoints/pooling/scoring/serving.py
+7 -1
@@ -8,6 +8,7 @@ from vllm.engine.protocol import EngineClient
8 from vllm.entrypoints.openai.engine.protocol import UsageInfo
9 from vllm.logger import init_logger
10 from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput
11 +from vllm.tasks import SCORE_TYPE_MAP, SupportedTask
12 from vllm.v1.pool.late_interaction import (
13 build_late_interaction_doc_params,
14 build_late_interaction_query_params,
@@ -38,10 +39,15 @@ class ServingScores(PoolingServing):
39 self,
40 engine_client: EngineClient,
41 *args,
42 + supported_tasks: tuple[SupportedTask, ...],
43 enable_flash_late_interaction: bool = True,
44 **kwargs,
45 ):
44 - self.io_processor_name: str = engine_client.model_config.score_type
46 + pooling_task = engine_client.model_config.get_pooling_task(supported_tasks)
47 + score_type = SCORE_TYPE_MAP.get(pooling_task, None) # type: ignore[arg-type]
48 + assert score_type is not None
49 +
50 + self.io_processor_name: str = score_type
51 self.enable_flash_late_interaction = (
52 self.io_processor_name == "late-interaction"
53 and enable_flash_late_interaction
vllm/entrypoints/pooling/utils.py
+6 -2
@@ -141,10 +141,14 @@ def enable_scoring_api(
141 supported_tasks: tuple["SupportedTask", ...],
142 model_config: ModelConfig | None = None,
143 ) -> bool:
144 - if any(t in supported_tasks for t in ("embed", "token_embed")):
144 + if model_config is None:
145 + return False
146 +
147 + pooling_task = model_config.get_pooling_task(supported_tasks)
148 + if pooling_task in ("embed", "token_embed"):
149 return True
150
147 - if model_config is not None and "classify" in supported_tasks:
151 + if pooling_task == "classify":
152 num_labels = getattr(model_config.hf_config, "num_labels", 0)
153 if num_labels != 1:
154 logger.debug_once("Scoring API is only enabled for num_labels == 1.")
vllm/pooling_params.py
-7
@@ -87,13 +87,6 @@ class PoolingParams(
87 return deepcopy(self)
88
89 def verify(self, model_config: ModelConfig) -> None:
90 - if self.task == "score":
91 - logger.warning_once(
92 - "`score` task is deprecated and will be removed in v0.20. "
93 - "Please use `classify` instead."
94 - )
95 - self.task = "classify"
96 -
90 # plugin task uses io_processor.parse_request to verify inputs,
91 # skipping PoolingParams verify
92 if self.task == "plugin":
vllm/tasks.py
+5
@@ -16,6 +16,11 @@ PoolingTask = Literal[
16 POOLING_TASKS: tuple[PoolingTask, ...] = get_args(PoolingTask)
17
18 ScoreType = Literal["bi-encoder", "cross-encoder", "late-interaction"]
19 +SCORE_TYPE_MAP: dict[PoolingTask, ScoreType] = {
20 + "embed": "bi-encoder",
21 + "classify": "cross-encoder",
22 + "token_embed": "late-interaction",
23 +}
24
25 FrontendTask = Literal["render"]
26 FRONTEND_TASKS: tuple[FrontendTask, ...] = get_args(FrontendTask)