@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f54f85129e
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 14, 05:30 UTC

[Model][2/N] Improve all pooling task | Support multi-vector retrieval (#25370)

Signed-off-by: wang.yuqi <noooop@126.com>

wang.yuqi committed Oct 15, 2025 at 19:14 UTC f54f85129e4665c16f39b097463c3c350ef34210
41 files changed +786 -399
examples/offline_inference/pooling/README.md
+6
@@ -26,6 +26,12 @@ python examples/offline_inference/pooling/embed_jina_embeddings_v3.py
26 python examples/offline_inference/pooling/embed_matryoshka_fy.py
27 ```
28
29 +## Multi vector retrieval usage
30 +
31 +```bash
32 +python examples/offline_inference/pooling/multi_vector_retrieval.py
33 +```
34 +
35 ## Named Entity Recognition (NER) usage
36
37 ```bash
examples/offline_inference/pooling/multi_vector_retrieval.py new
+56
@@ -0,0 +1,56 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from argparse import Namespace
5 +
6 +from vllm import LLM, EngineArgs
7 +from vllm.utils import FlexibleArgumentParser
8 +
9 +
10 +def parse_args():
11 + parser = FlexibleArgumentParser()
12 + parser = EngineArgs.add_cli_args(parser)
13 + # Set example specific arguments
14 + parser.set_defaults(
15 + model="BAAI/bge-m3",
16 + runner="pooling",
17 + enforce_eager=True,
18 + )
19 + return parser.parse_args()
20 +
21 +
22 +def main(args: Namespace):
23 + # Sample prompts.
24 + prompts = [
25 + "Hello, my name is",
26 + "The president of the United States is",
27 + "The capital of France is",
28 + "The future of AI is",
29 + ]
30 +
31 + # Create an LLM.
32 + # You should pass runner="pooling" for embedding models
33 + llm = LLM(**vars(args))
34 +
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 +
44 + # Generate embedding for each token. The output is a list of PoolingRequestOutput.
45 + outputs = llm.encode(prompts, pooling_task="token_embed")
46 +
47 + # Print the outputs.
48 + print("\nGenerated Outputs:\n" + "-" * 60)
49 + for prompt, output in zip(prompts, outputs):
50 + multi_vector = output.outputs.data
51 + print(multi_vector.shape)
52 +
53 +
54 +if __name__ == "__main__":
55 + args = parse_args()
56 + main(args)
examples/offline_inference/prithvi_geospatial_mae_io_processor.py
+1 -1
@@ -40,7 +40,7 @@ def main():
40 model_impl="terratorch",
41 )
42
43 - pooling_params = PoolingParams(task="encode", softmax=False)
43 + pooling_params = PoolingParams(task="token_classify", activation=False)
44 pooler_output = llm.encode(
45 img_prompt,
46 pooling_params=pooling_params,
examples/online_serving/pooling/README.md
+6
@@ -18,6 +18,12 @@ python examples/online_serving/pooling/embedding_embed_dtype_client.py
18 python examples/online_serving/pooling/jinaai_rerank_client.py
19 ```
20
21 +## Multi vector retrieval usage
22 +
23 +```bash
24 +python examples/online_serving/pooling/multi_vector_retrieval_client.py
25 +```
26 +
27 ## Named Entity Recognition (NER) usage
28
29 ```bash
examples/online_serving/pooling/multi_vector_retrieval_client.py new
+54
@@ -0,0 +1,54 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""
5 +Example online usage of Pooling API for multi vector retrieval.
6 +
7 +Run `vllm serve <model> --runner pooling`
8 +to start up the server in vLLM. e.g.
9 +
10 +vllm serve BAAI/bge-m3
11 +"""
12 +
13 +import argparse
14 +
15 +import requests
16 +import torch
17 +
18 +
19 +def post_http_request(prompt: dict, api_url: str) -> requests.Response:
20 + headers = {"User-Agent": "Test Client"}
21 + response = requests.post(api_url, headers=headers, json=prompt)
22 + return response
23 +
24 +
25 +def parse_args():
26 + parser = argparse.ArgumentParser()
27 + parser.add_argument("--host", type=str, default="localhost")
28 + parser.add_argument("--port", type=int, default=8000)
29 + parser.add_argument("--model", type=str, default="BAAI/bge-m3")
30 +
31 + return parser.parse_args()
32 +
33 +
34 +def main(args):
35 + api_url = f"http://{args.host}:{args.port}/pooling"
36 + model_name = args.model
37 +
38 + prompts = [
39 + "Hello, my name is",
40 + "The president of the United States is",
41 + "The capital of France is",
42 + "The future of AI is",
43 + ]
44 + prompt = {"model": model_name, "input": prompts}
45 +
46 + pooling_response = post_http_request(prompt=prompt, api_url=api_url)
47 + for output in pooling_response.json()["data"]:
48 + multi_vector = torch.tensor(output["data"])
49 + print(multi_vector.shape)
50 +
51 +
52 +if __name__ == "__main__":
53 + args = parse_args()
54 + main(args)
tests/conftest.py
+6 -2
@@ -1011,8 +1011,12 @@ class VllmRunner:
1011 req_outputs = self.llm.embed(inputs, *args, **kwargs)
1012 return [req_output.outputs.embedding for req_output in req_outputs]
1013
1014 - def encode(self, prompts: list[str]) -> list[list[float]]:
1015 - req_outputs = self.llm.encode(prompts)
1014 + def token_embed(self, prompts: list[str]) -> list[list[float]]:
1015 + req_outputs = self.llm.encode(prompts, pooling_task="token_embed")
1016 + return [req_output.outputs.data for req_output in req_outputs]
1017 +
1018 + def token_classify(self, prompts: list[str]) -> list[list[float]]:
1019 + req_outputs = self.llm.encode(prompts, pooling_task="token_classify")
1020 return [req_output.outputs.data for req_output in req_outputs]
1021
1022 def reward(self, prompts: list[str]) -> list[list[float]]:
tests/entrypoints/pooling/llm/test_classify.py
+1 -1
@@ -63,7 +63,7 @@ def test_encode_api(llm: LLM):
63 # chunked prefill does not support all pooling
64 err_msg = "pooling_task must be one of.+"
65 with pytest.raises(ValueError, match=err_msg):
66 - llm.encode(prompts, use_tqdm=False)
66 + llm.encode(prompts, pooling_task="token_classify", use_tqdm=False)
67
68
69 def test_score_api(llm: LLM):
tests/entrypoints/pooling/llm/test_embedding.py
+7
@@ -35,6 +35,13 @@ def llm():
35 cleanup_dist_env_and_memory()
36
37
38 +@pytest.mark.skip_global_cleanup
39 +def test_encode_api(llm: LLM):
40 + outputs = llm.encode(prompts, pooling_task="token_embed", use_tqdm=False)
41 + multi_vector = outputs[0].outputs.data
42 + assert multi_vector.shape == (11, 384)
43 +
44 +
45 def test_pooling_params(llm: LLM):
46 def get_outputs(normalize):
47 outputs = llm.embed(
tests/entrypoints/pooling/llm/test_encode.py
+8 -4
@@ -57,20 +57,24 @@ def test_multiple_pooling_params(llm: LLM):
57 ]
58
59 # Multiple PoolingParams should be matched with each prompt
60 - outputs = llm.encode(PROMPTS, pooling_params=pooling_params)
60 + outputs = llm.encode(PROMPTS, pooling_params=pooling_params, pooling_task="embed")
61 assert len(PROMPTS) == len(outputs)
62
63 # Exception raised, if the size of params does not match the size of prompts
64 with pytest.raises(ValueError):
65 - outputs = llm.encode(PROMPTS, pooling_params=pooling_params[:3])
65 + outputs = llm.encode(
66 + PROMPTS, pooling_params=pooling_params[:3], pooling_task="embed"
67 + )
68
69 # Single PoolingParams should be applied to every prompt
70 single_pooling_params = PoolingParams()
69 - outputs = llm.encode(PROMPTS, pooling_params=single_pooling_params)
71 + outputs = llm.encode(
72 + PROMPTS, pooling_params=single_pooling_params, pooling_task="embed"
73 + )
74 assert len(PROMPTS) == len(outputs)
75
76 # pooling_params is None, default params should be applied
73 - outputs = llm.encode(PROMPTS, pooling_params=None)
77 + outputs = llm.encode(PROMPTS, pooling_params=None, pooling_task="embed")
78 assert len(PROMPTS) == len(outputs)
79
80
tests/entrypoints/pooling/llm/test_reward.py
+12 -11
@@ -36,22 +36,23 @@ def llm():
36 cleanup_dist_env_and_memory()
37
38
39 -@pytest.mark.skip_global_cleanup
39 def test_pooling_params(llm: LLM):
41 - def get_outputs(softmax):
40 + def get_outputs(activation):
41 outputs = llm.reward(
43 - prompts, pooling_params=PoolingParams(softmax=softmax), use_tqdm=False
42 + prompts, pooling_params=PoolingParams(activation=activation), use_tqdm=False
43 )
44 return torch.cat([x.outputs.data for x in outputs])
45
47 - default = get_outputs(softmax=None)
48 - w_softmax = get_outputs(softmax=True)
49 - wo_softmax = get_outputs(softmax=False)
46 + default = get_outputs(activation=None)
47 + w_activation = get_outputs(activation=True)
48 + wo_activation = get_outputs(activation=False)
49
51 - assert torch.allclose(default, w_softmax, atol=1e-2), "Default should use softmax."
52 - assert not torch.allclose(w_softmax, wo_softmax, atol=1e-2), (
53 - "wo_softmax should not use softmax."
50 + assert torch.allclose(default, w_activation, atol=1e-2), (
51 + "Default should use activation."
52 )
55 - assert torch.allclose(softmax(wo_softmax), w_softmax, atol=1e-2), (
56 - "w_softmax should be close to softmax(wo_softmax)."
53 + assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
54 + "wo_activation should not use activation."
55 + )
56 + assert torch.allclose(softmax(wo_activation), w_activation, atol=1e-2), (
57 + "w_activation should be close to activation(wo_activation)."
58 )
tests/entrypoints/pooling/openai/test_embedding.py
+18
@@ -17,6 +17,7 @@ from tests.utils import RemoteOpenAIServer
17 from vllm.entrypoints.openai.protocol import (
18 EMBED_DTYPE_TO_TORCH_DTYPE,
19 EmbeddingResponse,
20 + PoolingResponse,
21 )
22 from vllm.transformers_utils.tokenizer import get_tokenizer
23
@@ -509,3 +510,20 @@ async def test_normalize(server: RemoteOpenAIServer, model_name: str):
510 assert torch.allclose(w_normal, F.normalize(wo_normal, p=2, dim=-1), atol=1e-2), (
511 "w_normal should be close to normal(wo_normal)."
512 )
513 +
514 +
515 +@pytest.mark.asyncio
516 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
517 +async def test_pooling(server: RemoteOpenAIServer, model_name: str):
518 + input_text = ["The chef prepared a delicious meal."]
519 +
520 + response = requests.post(
521 + server.url_for("pooling"),
522 + json={"model": model_name, "input": input_text, "encoding_format": "float"},
523 + )
524 +
525 + poolings = PoolingResponse.model_validate(response.json())
526 +
527 + assert len(poolings.data) == 1
528 + assert len(poolings.data[0].data) == 11
529 + assert len(poolings.data[0].data[0]) == 384
tests/entrypoints/pooling/openai/test_rerank.py
+18 -1
@@ -7,7 +7,7 @@ import torch
7 import torch.nn.functional as F
8
9 from tests.utils import RemoteOpenAIServer
10 -from vllm.entrypoints.openai.protocol import RerankResponse
10 +from vllm.entrypoints.openai.protocol import PoolingResponse, RerankResponse
11
12 MODEL_NAME = "BAAI/bge-reranker-base"
13 DTYPE = "bfloat16"
@@ -159,3 +159,20 @@ async def test_activation(server: RemoteOpenAIServer, model_name: str):
159 assert torch.allclose(F.sigmoid(wo_activation), w_activation, atol=1e-2), (
160 "w_activation should be close to activation(wo_activation)."
161 )
162 +
163 +
164 +@pytest.mark.asyncio
165 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
166 +async def test_pooling(server: RemoteOpenAIServer, model_name: str):
167 + input_text = ["The chef prepared a delicious meal."]
168 +
169 + response = requests.post(
170 + server.url_for("pooling"),
171 + json={"model": model_name, "input": input_text, "encoding_format": "float"},
172 + )
173 +
174 + poolings = PoolingResponse.model_validate(response.json())
175 +
176 + assert len(poolings.data) == 1
177 + assert len(poolings.data[0].data) == 11
178 + assert len(poolings.data[0].data[0]) == 1
tests/models/language/pooling/test_multi_vector_retrieval.py new
+45
@@ -0,0 +1,45 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import pytest
4 +import torch
5 +from transformers import AutoModel
6 +
7 +from tests.models.utils import check_embeddings_close
8 +
9 +
10 +@pytest.mark.parametrize(
11 + "model",
12 + ["BAAI/bge-m3"],
13 +)
14 +@pytest.mark.parametrize("dtype", ["half"])
15 +@torch.inference_mode
16 +def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype: str):
17 + with vllm_runner(
18 + model,
19 + runner="pooling",
20 + max_model_len=None,
21 + ) as vllm_model:
22 + vllm_outputs = vllm_model.token_embed(example_prompts)
23 +
24 + with hf_runner(
25 + model,
26 + auto_cls=AutoModel,
27 + ) as hf_model:
28 + tokenizer = hf_model.tokenizer
29 + hf_outputs = []
30 + for prompt in example_prompts:
31 + inputs = tokenizer([prompt], return_tensors="pt")
32 + inputs = hf_model.wrap_device(inputs)
33 + output = hf_model.model(**inputs)
34 + embedding = output.last_hidden_state[0].float()
35 + # normal
36 + hf_outputs.append(embedding.cpu())
37 +
38 + for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
39 + check_embeddings_close(
40 + embeddings_0_lst=hf_output,
41 + embeddings_1_lst=vllm_output,
42 + name_0="hf",
43 + name_1="vllm",
44 + tol=1e-2,
45 + )
tests/models/language/pooling/test_pooler_config_init_behaviour.py
+50 -8
@@ -93,7 +93,7 @@ def test_embed_models_using_normalize(
93 ],
94 )
95 @pytest.mark.parametrize("dtype", ["half"])
96 -def test_reward_models_using_softmax(
96 +def test_reward_models_using_activation(
97 hf_runner,
98 vllm_runner,
99 example_prompts,
@@ -104,22 +104,64 @@ def test_reward_models_using_softmax(
104 model,
105 max_model_len=1024,
106 dtype=dtype,
107 - pooler_config=PoolerConfig(softmax=False),
107 + pooler_config=PoolerConfig(activation=False),
108 ) as vllm_model:
109 - wo_softmax = vllm_model.encode(example_prompts)
109 + wo_activation = vllm_model.reward(example_prompts)
110
111 with vllm_runner(
112 - model, max_model_len=1024, dtype=dtype, pooler_config=PoolerConfig(softmax=True)
112 + model,
113 + max_model_len=1024,
114 + dtype=dtype,
115 + pooler_config=PoolerConfig(activation=True),
116 ) as vllm_model:
114 - w_softmax = vllm_model.encode(example_prompts)
117 + w_activation = vllm_model.reward(example_prompts)
118
116 - for wo, w in zip(wo_softmax, w_softmax):
119 + for wo, w in zip(wo_activation, w_activation):
120 wo = torch.tensor(wo)
121 w = torch.tensor(w)
122
123 assert not torch.allclose(wo, w, atol=1e-2), (
121 - "pooler_config softmax is not working"
124 + "pooler_config activation is not working"
125 )
126 assert torch.allclose(softmax(wo), w, atol=1e-2), (
124 - "w_softmax should be close to softmax(wo_softmax)."
127 + "w_activation should be close to activation(wo_activation)."
128 + )
129 +
130 +
131 +@pytest.mark.parametrize(
132 + "model",
133 + [
134 + "intfloat/multilingual-e5-small",
135 + ],
136 +)
137 +@pytest.mark.parametrize("dtype", ["half"])
138 +def test_multi_vector_retrieval_models_using_normalize(
139 + hf_runner,
140 + vllm_runner,
141 + example_prompts,
142 + model: str,
143 + dtype: str,
144 +) -> None:
145 + with vllm_runner(
146 + model,
147 + max_model_len=512,
148 + dtype=dtype,
149 + pooler_config=PoolerConfig(normalize=False),
150 + ) as vllm_model:
151 + wo_normalize = vllm_model.token_embed(example_prompts)
152 +
153 + with vllm_runner(
154 + model,
155 + max_model_len=512,
156 + dtype=dtype,
157 + pooler_config=PoolerConfig(normalize=True),
158 + ) as vllm_model:
159 + w_normalize = vllm_model.token_embed(example_prompts)
160 +
161 + for wo, w in zip(wo_normalize, w_normalize):
162 + assert not torch.allclose(wo, w, atol=1e-2), (
163 + "pooler_config normalize is not working"
164 + )
165 + assert torch.allclose(F.normalize(wo, p=2, dim=-1), w, atol=1e-2), (
166 + "w_normal should be close to normal(wo_normal)."
167 )
tests/models/language/pooling/test_token_classification.py
+2 -2
@@ -19,7 +19,7 @@ def test_bert_models(
19 dtype: str,
20 ) -> None:
21 with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
22 - vllm_outputs = vllm_model.encode(example_prompts)
22 + vllm_outputs = vllm_model.token_classify(example_prompts)
23
24 with hf_runner(
25 model, dtype=dtype, auto_cls=AutoModelForTokenClassification
@@ -50,7 +50,7 @@ def test_modernbert_models(
50 dtype: str,
51 ) -> None:
52 with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
53 - vllm_outputs = vllm_model.encode(example_prompts)
53 + vllm_outputs = vllm_model.token_classify(example_prompts)
54
55 with hf_runner(
56 model, dtype=dtype, auto_cls=AutoModelForTokenClassification
tests/models/multimodal/pooling/test_prithvi_mae.py
+1 -1
@@ -39,7 +39,7 @@ def _run_test(
39 max_num_seqs=32,
40 default_torch_num_threads=1,
41 ) as vllm_model:
42 - vllm_model.encode(prompt)
42 + vllm_model.llm.encode(prompt, pooling_task="token_classify")
43
44
45 MODELS = ["mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]
tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_gemma_embedding.py
+1 -1
@@ -30,7 +30,7 @@ class MyGemma2Embedding(nn.Module):
30
31 self.pooler = DispatchPooler(
32 {
33 - "encode": Pooler.for_encode(pooler_config),
33 + "token_embed": Pooler.for_token_embed(pooler_config),
34 "embed": Pooler.for_embed(pooler_config),
35 }
36 )
tests/plugins_tests/test_io_processor_plugins.py
+2 -3
@@ -93,7 +93,7 @@ def test_prithvi_mae_plugin_offline(vllm_runner, model_name: str):
93 out_data_format="b64_json",
94 )
95
96 - pooling_params = PoolingParams(task="encode", softmax=False)
96 + pooling_params = PoolingParams(activation=False)
97
98 with vllm_runner(
99 model_name,
@@ -108,8 +108,7 @@ def test_prithvi_mae_plugin_offline(vllm_runner, model_name: str):
108 io_processor_plugin="prithvi_to_tiff",
109 ) as llm_runner:
110 pooler_output = llm_runner.get_llm().encode(
111 - img_prompt,
112 - pooling_params=pooling_params,
111 + img_prompt, pooling_params=pooling_params, pooling_task="token_classify"
112 )
113 output = pooler_output[0].outputs
114
tests/test_pooling_params.py
+71 -22
@@ -1,10 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from dataclasses import dataclass
4 +
5 import pytest
6
7 from tests.models.utils import EmbedModelInfo
8 from vllm import PoolingParams
7 -from vllm.config import ModelConfig
9 +from vllm.config import ModelConfig, PoolerConfig
10
11 EMBEDDING_MODELS = [
12 EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False),
@@ -15,6 +17,15 @@ EMBEDDING_MODELS = [
17 ),
18 ]
19
20 +classify_parameters = ["activation"]
21 +embed_parameters = ["dimensions", "normalize"]
22 +step_pooling_parameters = ["step_tag_id", "returned_token_ids"]
23 +
24 +
25 +@dataclass()
26 +class MockModelConfig:
27 + pooler_config: PoolerConfig
28 +
29
30 def test_task():
31 pooling_params = PoolingParams()
@@ -24,25 +35,27 @@ def test_task():
35 pooling_params.verify(task="score")
36
37 with pytest.raises(ValueError):
27 - pooling_params.verify(task="encode")
38 + pooling_params.verify(task="classify")
39
40
41 def test_embed():
42 task = "embed"
43 + model_config = MockModelConfig(pooler_config=PoolerConfig(pooling_type="CLS"))
44 +
45 pooling_params = PoolingParams(normalize=None)
33 - pooling_params.verify(task=task)
46 + pooling_params.verify(task=task, model_config=model_config)
47
48 pooling_params = PoolingParams(normalize=True)
36 - pooling_params.verify(task=task)
49 + pooling_params.verify(task=task, model_config=model_config)
50
51 pooling_params = PoolingParams(normalize=False)
39 - pooling_params.verify(task=task)
52 + pooling_params.verify(task=task, model_config=model_config)
53
41 - invalid_parameters = ["activation", "softmax"]
54 + invalid_parameters = classify_parameters + step_pooling_parameters
55 for p in invalid_parameters:
56 with pytest.raises(ValueError):
57 pooling_params = PoolingParams(**{p: True})
45 - pooling_params.verify(task=task)
58 + pooling_params.verify(task=task, model_config=model_config)
59
60
61 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
@@ -73,35 +86,71 @@ def test_embed_dimensions(model_info: EmbedModelInfo):
86
87 @pytest.mark.parametrize("task", ["score", "classify"])
88 def test_classify(task):
89 + model_config = MockModelConfig(pooler_config=PoolerConfig(pooling_type="CLS"))
90 +
91 pooling_params = PoolingParams(activation=None)
77 - pooling_params.verify(task=task)
92 + pooling_params.verify(task=task, model_config=model_config)
93
94 pooling_params = PoolingParams(activation=True)
80 - pooling_params.verify(task=task)
95 + pooling_params.verify(task=task, model_config=model_config)
96
97 pooling_params = PoolingParams(activation=False)
83 - pooling_params.verify(task=task)
98 + pooling_params.verify(task=task, model_config=model_config)
99 +
100 + invalid_parameters = embed_parameters + step_pooling_parameters
101 + for p in invalid_parameters:
102 + with pytest.raises(ValueError):
103 + pooling_params = PoolingParams(**{p: True})
104 + pooling_params.verify(task=task, model_config=model_config)
105 +
106 +
107 +@pytest.mark.parametrize("pooling_type", ["ALL", "STEP"])
108 +def test_token_embed(pooling_type: str):
109 + task = "token_embed"
110 + model_config = MockModelConfig(
111 + pooler_config=PoolerConfig(pooling_type=pooling_type)
112 + )
113 +
114 + pooling_params = PoolingParams(normalize=None)
115 + pooling_params.verify(task=task, model_config=model_config)
116 +
117 + pooling_params = PoolingParams(normalize=True)
118 + pooling_params.verify(task=task, model_config=model_config)
119 +
120 + pooling_params = PoolingParams(normalize=False)
121 + pooling_params.verify(task=task, model_config=model_config)
122 +
123 + invalid_parameters = classify_parameters
124 + if pooling_type != "STEP":
125 + invalid_parameters = classify_parameters + step_pooling_parameters
126
85 - invalid_parameters = ["dimensions", "normalize", "softmax"]
127 for p in invalid_parameters:
128 with pytest.raises(ValueError):
129 pooling_params = PoolingParams(**{p: True})
89 - pooling_params.verify(task=task)
130 + pooling_params.verify(task=task, model_config=model_config)
131
132
92 -def test_encode():
93 - task = "encode"
94 - pooling_params = PoolingParams(softmax=None)
95 - pooling_params.verify(task=task)
133 +@pytest.mark.parametrize("pooling_type", ["ALL", "STEP"])
134 +def test_token_classify(pooling_type: str):
135 + task = "token_classify"
136 + model_config = MockModelConfig(
137 + pooler_config=PoolerConfig(pooling_type=pooling_type)
138 + )
139
97 - pooling_params = PoolingParams(softmax=True)
98 - pooling_params.verify(task=task)
140 + pooling_params = PoolingParams(activation=None)
141 + pooling_params.verify(task=task, model_config=model_config)
142 +
143 + pooling_params = PoolingParams(activation=True)
144 + pooling_params.verify(task=task, model_config=model_config)
145 +
146 + pooling_params = PoolingParams(activation=False)
147 + pooling_params.verify(task=task, model_config=model_config)
148
100 - pooling_params = PoolingParams(softmax=False)
101 - pooling_params.verify(task=task)
149 + invalid_parameters = embed_parameters
150 + if pooling_type != "STEP":
151 + invalid_parameters = embed_parameters + step_pooling_parameters
152
103 - invalid_parameters = ["dimensions", "normalize", "activation"]
153 for p in invalid_parameters:
154 with pytest.raises(ValueError):
155 pooling_params = PoolingParams(**{p: True})
107 - pooling_params.verify(task=task)
156 + pooling_params.verify(task=task, model_config=model_config)
vllm/entrypoints/llm.py
+18 -19
@@ -951,7 +951,7 @@ class LLM:
951 truncate_prompt_tokens: int | None = None,
952 use_tqdm: bool | Callable[..., tqdm] = True,
953 lora_request: list[LoRARequest] | LoRARequest | None = None,
954 - pooling_task: PoolingTask = "encode",
954 + pooling_task: PoolingTask | None = None,
955 tokenization_kwargs: dict[str, Any] | None = None,
956 ) -> list[PoolingRequestOutput]:
957 """Apply pooling to the hidden states corresponding to the input
@@ -986,25 +986,24 @@ class LLM:
986 instead pass them via the `inputs` parameter.
987 """
988
989 - if self.supported_tasks == ["encode"] and pooling_task is None:
990 - pooling_task = "encode"
989 + error_str = (
990 + "pooling_task required for `LLM.encode`\n"
991 + "Please use one of the more specific methods or set the "
992 + "pooling_task when using `LLM.encode`:\n"
993 + " - For embeddings, use `LLM.embed(...)` "
994 + 'or `pooling_task="embed"`.\n'
995 + " - For classification logits, use `LLM.classify(...)` "
996 + 'or `pooling_task="classify"`.\n'
997 + " - For similarity scores, use `LLM.score(...)`.\n"
998 + " - For rewards, use `LLM.reward(...)` "
999 + 'or `pooling_task="token_classify"`\n'
1000 + " - For token classification, "
1001 + 'use `pooling_task="token_classify"`\n'
1002 + ' - For multi-vector retrieval, use `pooling_task="token_embed"`'
1003 + )
1004
1005 if pooling_task is None:
993 - pooling_task = "embed" if "embed" in self.supported_tasks else "encode"
994 -
995 - logger.warning_once(
996 - "`LLM.encode` is currently using `pooling_task = %s`.\n"
997 - "Please use one of the more specific methods or set the "
998 - "task directly when using `LLM.encode`:\n"
999 - " - For embeddings, use `LLM.embed(...)` "
1000 - 'or `pooling_task="embed"`.\n'
1001 - " - For classification logits, use `LLM.classify(...)` "
1002 - 'or `pooling_task="classify"`.\n'
1003 - " - For rewards, use `LLM.reward(...)` "
1004 - 'or `pooling_task="reward"`\n'
1005 - " - For similarity scores, use `LLM.score(...)`.",
1006 - pooling_task,
1007 - )
1006 + raise ValueError(error_str)
1007
1008 model_config = self.model_config
1009 runner_type = model_config.runner_type
@@ -1206,7 +1205,7 @@ class LLM:
1205 lora_request=lora_request,
1206 pooling_params=pooling_params,
1207 truncate_prompt_tokens=truncate_prompt_tokens,
1209 - pooling_task="encode",
1208 + pooling_task="token_classify",
1209 )
1210
1211 def _embedding_score(
vllm/entrypoints/openai/api_server.py
+12 -9
@@ -1748,16 +1748,19 @@ async def init_app_state(
1748 else None
1749 )
1750 state.openai_serving_pooling = (
1751 - OpenAIServingPooling(
1752 - engine_client,
1753 - state.openai_serving_models,
1754 - request_logger=request_logger,
1755 - chat_template=resolved_chat_template,
1756 - chat_template_content_format=args.chat_template_content_format,
1757 - trust_request_chat_template=args.trust_request_chat_template,
1758 - log_error_stack=args.log_error_stack,
1751 + (
1752 + OpenAIServingPooling(
1753 + engine_client,
1754 + state.openai_serving_models,
1755 + supported_tasks=supported_tasks,
1756 + request_logger=request_logger,
1757 + chat_template=resolved_chat_template,
1758 + chat_template_content_format=args.chat_template_content_format,
1759 + trust_request_chat_template=args.trust_request_chat_template,
1760 + log_error_stack=args.log_error_stack,
1761 + )
1762 )
1760 - if "encode" in supported_tasks
1763 + if ("token_embed" in supported_tasks or "token_classify" in supported_tasks)
1764 else None
1765 )
1766 state.openai_serving_embedding = (
vllm/entrypoints/openai/protocol.py
+2 -2
@@ -1682,7 +1682,7 @@ class IOProcessorRequest(OpenAIBaseModel, Generic[T]):
1682 When using plugins IOProcessor plugins, the actual input is processed
1683 by the plugin itself. Hence, we use a generic type for the request data
1684 """
1685 - softmax: bool = True
1685 + activation: bool = False
1686
1687 embed_dtype: str = Field(
1688 default="float32",
@@ -1693,7 +1693,7 @@ class IOProcessorRequest(OpenAIBaseModel, Generic[T]):
1693 )
1694
1695 def to_pooling_params(self):
1696 - return PoolingParams(task="encode", softmax=self.softmax)
1696 + return PoolingParams(task="token_classify", activation=self.activation)
1697
1698
1699 class IOProcessorResponse(OpenAIBaseModel, Generic[T]):
vllm/entrypoints/openai/serving_pooling.py
+13 -1
@@ -35,6 +35,7 @@ from vllm.entrypoints.renderer import RenderConfig
35 from vllm.entrypoints.utils import _validate_truncation_size
36 from vllm.logger import init_logger
37 from vllm.outputs import PoolingOutput, PoolingRequestOutput
38 +from vllm.tasks import SupportedTask
39 from vllm.utils import merge_async_iterators
40
41 logger = init_logger(__name__)
@@ -62,6 +63,7 @@ class OpenAIServingPooling(OpenAIServing):
63 engine_client: EngineClient,
64 models: OpenAIServingModels,
65 *,
66 + supported_tasks: tuple[SupportedTask, ...],
67 request_logger: RequestLogger | None,
68 chat_template: str | None,
69 chat_template_content_format: ChatTemplateContentFormatOption,
@@ -75,6 +77,7 @@ class OpenAIServingPooling(OpenAIServing):
77 log_error_stack=log_error_stack,
78 )
79
80 + self.supported_tasks = supported_tasks
81 self.chat_template = chat_template
82 self.chat_template_content_format: Final = chat_template_content_format
83 self.trust_request_chat_template = trust_request_chat_template
@@ -178,8 +181,17 @@ class OpenAIServingPooling(OpenAIServing):
181 try:
182 pooling_params = request.to_pooling_params()
183
184 + if "token_embed" in self.supported_tasks:
185 + pooling_task = "token_embed"
186 + elif "token_classify" in self.supported_tasks:
187 + pooling_task = "token_classify"
188 + else:
189 + return self.create_error_response(
190 + f"pooling_task must be one of {self.supported_tasks}."
191 + )
192 +
193 try:
182 - pooling_params.verify("encode", self.model_config)
194 + pooling_params.verify(pooling_task, self.model_config)
195 except ValueError as e:
196 return self.create_error_response(str(e))
197
vllm/model_executor/layers/pooler.py
+251 -171
@@ -64,66 +64,6 @@ class PoolingParamsUpdate:
64 params.requires_token_ids = self.requires_token_ids
65
66
67 -class Pooler(nn.Module, ABC):
68 - """The interface required for all poolers used in pooling models in vLLM."""
69 -
70 - @staticmethod
71 - def for_encode(pooler_config: PoolerConfig):
72 - if pooler_config.pooling_type == "STEP":
73 - return StepPooler()
74 -
75 - resolved_config = ResolvedPoolingConfig(
76 - task="encode", pooling_type=PoolingType.ALL
77 - )
78 -
79 - return SimplePooler.from_config(resolved_config)
80 -
81 - @staticmethod
82 - def for_embed(pooler_config: PoolerConfig):
83 - resolved_config = ResolvedPoolingConfig.from_config(
84 - task="embed",
85 - pooler_config=pooler_config,
86 - )
87 -
88 - return SimplePooler.from_config(resolved_config)
89 -
90 - @staticmethod
91 - def for_classify(
92 - pooler_config: PoolerConfig,
93 - classifier: ClassifierFn | None,
94 - ):
95 - resolved_config = ResolvedPoolingConfig.from_config(
96 - task="classify",
97 - pooler_config=pooler_config,
98 - )
99 -
100 - pooling = PoolingMethod.from_pooling_type(resolved_config.pooling_type)
101 -
102 - return ClassifierPooler(
103 - pooling=pooling,
104 - classifier=classifier,
105 - )
106 -
107 - @abstractmethod
108 - def get_supported_tasks(self) -> Set[PoolingTask]:
109 - """Determine which pooling tasks are supported."""
110 - raise NotImplementedError
111 -
112 - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
113 - """
114 - Construct the updated pooling parameters to use for a supported task.
115 - """
116 - return PoolingParamsUpdate()
117 -
118 - @abstractmethod
119 - def forward(
120 - self,
121 - hidden_states: list[torch.Tensor] | torch.Tensor,
122 - pooling_metadata: PoolingMetadata,
123 - ) -> PoolerOutput:
124 - raise NotImplementedError
125 -
126 -
67 def get_prompt_lens(
68 hidden_states: torch.Tensor | list[torch.Tensor],
69 pooling_metadata: PoolingMetadata,
@@ -237,7 +177,7 @@ class PoolingMethod(nn.Module, ABC):
177
178 class CLSPool(PoolingMethod):
179 def get_supported_tasks(self) -> Set[PoolingTask]:
240 - return {"encode", "embed", "classify", "score"}
180 + return {"token_embed", "token_classify", "embed", "classify", "score"}
181
182 def forward_all(
183 self,
@@ -253,7 +193,7 @@ class CLSPool(PoolingMethod):
193
194 class LastPool(PoolingMethod):
195 def get_supported_tasks(self) -> Set[PoolingTask]:
256 - return {"encode", "embed", "classify", "score"}
196 + return {"token_embed", "token_classify", "embed", "classify", "score"}
197
198 def forward_all(
199 self,
@@ -265,7 +205,7 @@ class LastPool(PoolingMethod):
205
206 class AllPool(PoolingMethod):
207 def get_supported_tasks(self) -> Set[PoolingTask]:
268 - return {"encode"}
208 + return {"token_embed", "token_classify"}
209
210 def forward_all(
211 self,
@@ -284,7 +224,7 @@ class AllPool(PoolingMethod):
224
225 class MeanPool(PoolingMethod):
226 def get_supported_tasks(self) -> Set[PoolingTask]:
287 - return {"encode", "embed", "classify", "score"}
227 + return {"token_embed", "token_classify", "embed", "classify", "score"}
228
229 def forward_all(
230 self,
@@ -398,6 +338,82 @@ class LambdaPoolerActivation(PoolerActivation):
338 return self.fn(pooled_data)
339
340
341 +class Pooler(nn.Module, ABC):
342 + """The interface required for all poolers used in pooling models in vLLM."""
343 +
344 + @staticmethod
345 + def for_token_embed(pooler_config: PoolerConfig):
346 + head = TokenEmbeddingPoolerHead()
347 +
348 + if pooler_config.pooling_type == "STEP":
349 + return StepPooler(head=head)
350 +
351 + return AllPooler(head=head)
352 +
353 + @staticmethod
354 + def for_token_classify(
355 + pooler_config: PoolerConfig,
356 + classifier: ClassifierFn | None = None,
357 + act_fn: PoolerActivation | str | None = None,
358 + ):
359 + head = TokenClassifierPoolerHead(classifier=classifier, act_fn=act_fn)
360 +
361 + if pooler_config.pooling_type == "STEP":
362 + return StepPooler(head=head)
363 +
364 + return AllPooler(head=head)
365 +
366 + @staticmethod
367 + def for_embed(pooler_config: PoolerConfig):
368 + resolved_config = ResolvedPoolingConfig.from_config(
369 + task="embed",
370 + pooler_config=pooler_config,
371 + )
372 +
373 + pooling = PoolingMethod.from_pooling_type(resolved_config.pooling_type)
374 + head = EmbeddingPoolerHead()
375 +
376 + return SimplePooler(pooling=pooling, head=head)
377 +
378 + @staticmethod
379 + def for_classify(
380 + pooler_config: PoolerConfig,
381 + classifier: ClassifierFn | None,
382 + act_fn: PoolerActivation | str | None = None,
383 + ):
384 + resolved_config = ResolvedPoolingConfig.from_config(
385 + task="classify",
386 + pooler_config=pooler_config,
387 + )
388 +
389 + pooling = PoolingMethod.from_pooling_type(resolved_config.pooling_type)
390 +
391 + return ClassifierPooler(
392 + pooling=pooling,
393 + classifier=classifier,
394 + act_fn=act_fn,
395 + )
396 +
397 + @abstractmethod
398 + def get_supported_tasks(self) -> Set[PoolingTask]:
399 + """Determine which pooling tasks are supported."""
400 + raise NotImplementedError
401 +
402 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
403 + """
404 + Construct the updated pooling parameters to use for a supported task.
405 + """
406 + return PoolingParamsUpdate()
407 +
408 + @abstractmethod
409 + def forward(
410 + self,
411 + hidden_states: list[torch.Tensor] | torch.Tensor,
412 + pooling_metadata: PoolingMetadata,
413 + ) -> PoolerOutput:
414 + raise NotImplementedError
415 +
416 +
417 class PoolerHead(nn.Module):
418 def __init__(self, activation: PoolerActivation) -> None:
419 super().__init__()
@@ -416,7 +432,6 @@ class EmbeddingPoolerHead(PoolerHead):
432 super().__init__(activation=PoolerNormalize())
433
434 # Load ST projector if available
419 -
435 vllm_config = get_current_vllm_config()
436 self.projector: nn.Module | None = (
437 _load_st_projector(vllm_config.model_config) if vllm_config else None
@@ -471,39 +486,6 @@ class EmbeddingPoolerHead(PoolerHead):
486 return pooled_data
487
488
474 -class RewardPoolerHead(PoolerHead):
475 - def __init__(self) -> None:
476 - super().__init__(activation=PoolerClassify(static_num_labels=False))
477 -
478 - vllm_config = get_current_vllm_config()
479 - self.head_dtype = vllm_config.model_config.head_dtype
480 -
481 - def forward(
482 - self,
483 - pooled_data: list[torch.Tensor] | torch.Tensor,
484 - pooling_metadata: PoolingMetadata,
485 - ):
486 - if isinstance(pooled_data, list):
487 - pooled_data = [p.to(self.head_dtype) for p in pooled_data]
488 - else:
489 - pooled_data = pooled_data.to(self.head_dtype)
490 -
491 - pooling_params = get_pooling_params(pooling_metadata)
492 -
493 - # for softmax
494 - flags = [p.softmax for p in pooling_params]
495 - if len(set(flags)) == 1:
496 - if flags[0]:
497 - pooled_data = self.activation(pooled_data)
498 - else:
499 - pooled_data = [
500 - self.activation(vecs) if f else vecs
501 - for vecs, f in zip(pooled_data, flags)
502 - ]
503 -
504 - return pooled_data
505 -
506 -
489 class SimplePooler(Pooler):
490 """A layer that pools specific information from hidden states.
491
@@ -513,20 +495,6 @@ class SimplePooler(Pooler):
495 3. Returns structured results as `PoolerOutput`.
496 """
497
516 - @classmethod
517 - def from_config(
518 - cls,
519 - pooler_config: ResolvedPoolingConfig,
520 - ) -> "SimplePooler":
521 - pooling = PoolingMethod.from_pooling_type(pooler_config.pooling_type)
522 - if pooler_config.task == "embed":
523 - head = EmbeddingPoolerHead()
524 - elif pooler_config.task == "encode":
525 - head = RewardPoolerHead()
526 - else:
527 - raise NotImplementedError(f"Unknown task: {pooler_config.task}")
528 - return cls(pooling, head)
529 -
498 def __init__(self, pooling: PoolingMethod, head: PoolerHead) -> None:
499 super().__init__()
500
@@ -549,58 +517,6 @@ class SimplePooler(Pooler):
517 return pooled_data
518
519
552 -class StepPooler(Pooler):
553 - def __init__(
554 - self,
555 - ) -> None:
556 - super().__init__()
557 -
558 - self.pooling = AllPool()
559 - self.head = RewardPoolerHead()
560 -
561 - def extract_states(
562 - self,
563 - hidden_states: torch.Tensor | list[torch.Tensor],
564 - pooling_metadata: PoolingMetadata,
565 - ) -> list[torch.Tensor] | torch.Tensor:
566 - pooled_data_lst = self.pooling(hidden_states, pooling_metadata)
567 - prompt_token_ids = get_prompt_token_ids(pooling_metadata)
568 -
569 - pooled_data = list[torch.Tensor]()
570 -
571 - pooling_params = get_pooling_params(pooling_metadata)
572 -
573 - for data, token_id, pooling_param in zip(
574 - pooled_data_lst, prompt_token_ids, pooling_params
575 - ):
576 - step_tag_id = pooling_param.step_tag_id
577 - returned_token_ids = pooling_param.returned_token_ids
578 -
579 - if returned_token_ids is not None and len(returned_token_ids) > 0:
580 - data = data[:, returned_token_ids]
581 -
582 - if step_tag_id is not None:
583 - data = data[token_id == step_tag_id]
584 - pooled_data.append(data)
585 -
586 - return pooled_data
587 -
588 - def get_supported_tasks(self) -> Set[PoolingTask]:
589 - return {"encode"}
590 -
591 - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
592 - return PoolingParamsUpdate(requires_token_ids=True)
593 -
594 - def forward(
595 - self,
596 - hidden_states: torch.Tensor | list[torch.Tensor],
597 - pooling_metadata: PoolingMetadata,
598 - ) -> PoolerOutput:
599 - pooled_data = self.extract_states(hidden_states, pooling_metadata)
600 - pooled_data = self.head(pooled_data, pooling_metadata)
601 - return pooled_data
602 -
603 -
520 class ClassifierPooler(Pooler):
521 """A pooling layer for classification tasks.
522
@@ -611,26 +527,46 @@ class ClassifierPooler(Pooler):
527 """
528
529 @staticmethod
614 - def act_fn_for_seq_cls(config: ModelConfig):
615 - return get_classification_activation_function(config.hf_config)
530 + def act_fn_for_seq_cls(model_config: ModelConfig):
531 + return get_classification_activation_function(model_config.hf_config)
532 +
533 + @staticmethod
534 + def act_fn_for_cross_encoder(model_config: ModelConfig):
535 + return get_cross_encoder_activation_function(model_config.hf_config)
536
537 @staticmethod
618 - def act_fn_for_cross_encoder(config: ModelConfig):
619 - return get_cross_encoder_activation_function(config.hf_config)
538 + def resolve_act_fn(
539 + model_config: ModelConfig,
540 + static_num_labels: bool = True,
541 + act_fn: PoolerActivation | str | None = None,
542 + ):
543 + if isinstance(act_fn, str):
544 + if act_fn == "classify":
545 + return ClassifierPooler.act_fn_for_seq_cls(model_config)
546 + elif act_fn == "score":
547 + return ClassifierPooler.act_fn_for_cross_encoder(model_config)
548 + else:
549 + raise ValueError(f"act_fn [{act_fn=}] not supported.")
550 + elif act_fn is None:
551 + return PoolerClassify(static_num_labels=static_num_labels)
552 + else:
553 + assert callable(act_fn)
554 + return act_fn
555
556 def __init__(
557 self,
558 pooling: PoolingFn,
559 classifier: ClassifierFn | None,
625 - act_fn: PoolerActivation | None = None,
560 + act_fn: PoolerActivation | str | None = None,
561 ) -> None:
562 super().__init__()
563
564 vllm_config = get_current_vllm_config()
630 -
565 self.pooling = pooling
566 self.classifier = classifier
633 - self.act_fn = act_fn or PoolerClassify()
567 + self.act_fn = self.resolve_act_fn(
568 + vllm_config.model_config, static_num_labels=True, act_fn=act_fn
569 + )
570 self.logit_bias: float | None = (
571 vllm_config.model_config.pooler_config.logit_bias
572 )
@@ -672,6 +608,150 @@ class ClassifierPooler(Pooler):
608 return scores
609
610
611 +class TokenEmbeddingPoolerHead(EmbeddingPoolerHead):
612 + def forward(
613 + self, pooled_data: torch.Tensor, pooling_param: PoolingParams
614 + ) -> torch.Tensor:
615 + pooled_data = pooled_data.to(self.head_dtype)
616 + # pooled_data shape: [n_tokens, hidden_dimension]
617 +
618 + # Apply ST projector
619 + if self.projector is not None:
620 + pooled_data = self.projector(pooled_data)
621 + # pooled_data shape: [n_tokens, embedding_dimension]
622 +
623 + # for matryoshka representation
624 + pooled_data = pooled_data[..., : pooling_param.dimensions]
625 +
626 + # for normalize
627 + if pooling_param.normalize:
628 + pooled_data = self.activation(pooled_data)
629 +
630 + # pooled_data shape: [n_tokens, embedding_dimension]
631 + return pooled_data
632 +
633 +
634 +class TokenClassifierPoolerHead(nn.Module):
635 + def __init__(
636 + self,
637 + classifier: ClassifierFn | None,
638 + act_fn: PoolerActivation | str | None = None,
639 + ) -> None:
640 + super().__init__()
641 + vllm_config = get_current_vllm_config()
642 +
643 + self.classifier = classifier
644 + self.act_fn = ClassifierPooler.resolve_act_fn(
645 + vllm_config.model_config, static_num_labels=False, act_fn=act_fn
646 + )
647 + self.logit_bias: float | None = (
648 + vllm_config.model_config.pooler_config.logit_bias
649 + )
650 + self.head_dtype = vllm_config.model_config.head_dtype
651 +
652 + def get_supported_tasks(self) -> Set[PoolingTask]:
653 + return {"token_classify"}
654 +
655 + def forward(
656 + self,
657 + hidden_states: torch.Tensor,
658 + pooling_param: PoolingParams,
659 + ) -> torch.Tensor:
660 + hidden_states = hidden_states.to(self.head_dtype)
661 + # hidden_states shape: [n_token, hidden_size]
662 +
663 + if self.classifier is not None:
664 + scores = self.classifier(hidden_states)
665 + else:
666 + scores = hidden_states
667 + # scores shape: [n_token, num_labels]
668 +
669 + if self.logit_bias is not None:
670 + scores -= self.logit_bias
671 +
672 + if pooling_param.activation:
673 + scores = self.act_fn(scores)
674 +
675 + # scores shape: [n_token, num_labels]
676 + return scores
677 +
678 +
679 +class AllPooler(Pooler):
680 + def __init__(self, head: nn.Module | PoolerHead) -> None:
681 + super().__init__()
682 +
683 + self.pooling = AllPool()
684 + self.head = head
685 +
686 + def get_supported_tasks(self) -> Set[PoolingTask]:
687 + return {"token_embed", "token_classify"}
688 +
689 + def forward(
690 + self,
691 + hidden_states: torch.Tensor,
692 + pooling_metadata: PoolingMetadata,
693 + ) -> PoolerOutput:
694 + pooled_data = self.pooling(hidden_states, pooling_metadata)
695 + pooling_params = get_pooling_params(pooling_metadata)
696 + assert len(pooled_data) == len(pooling_params)
697 +
698 + pooled_data = [self.head(d, p) for d, p in zip(pooled_data, pooling_params)]
699 + return pooled_data
700 +
701 +
702 +class StepPooler(Pooler):
703 + def __init__(self, head: nn.Module | PoolerHead) -> None:
704 + super().__init__()
705 +
706 + self.pooling = AllPool()
707 + self.head = head
708 +
709 + def extract_states(
710 + self,
711 + hidden_states: torch.Tensor | list[torch.Tensor],
712 + pooling_metadata: PoolingMetadata,
713 + ) -> torch.Tensor | list[torch.Tensor]:
714 + pooled_data_lst = self.pooling(hidden_states, pooling_metadata)
715 + prompt_token_ids = get_prompt_token_ids(pooling_metadata)
716 +
717 + pooled_data = list[torch.Tensor]()
718 +
719 + pooling_params = get_pooling_params(pooling_metadata)
720 +
721 + for data, token_id, pooling_param in zip(
722 + pooled_data_lst, prompt_token_ids, pooling_params
723 + ):
724 + step_tag_id = pooling_param.step_tag_id
725 + returned_token_ids = pooling_param.returned_token_ids
726 +
727 + if returned_token_ids is not None and len(returned_token_ids) > 0:
728 + data = data[:, returned_token_ids]
729 +
730 + if step_tag_id is not None:
731 + data = data[token_id == step_tag_id]
732 + pooled_data.append(data)
733 +
734 + return pooled_data
735 +
736 + def get_supported_tasks(self) -> Set[PoolingTask]:
737 + return {"token_embed", "token_classify"}
738 +
739 + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
740 + return PoolingParamsUpdate(requires_token_ids=True)
741 +
742 + def forward(
743 + self,
744 + hidden_states: torch.Tensor | list[torch.Tensor],
745 + pooling_metadata: PoolingMetadata,
746 + ) -> PoolerOutput:
747 + pooled_data = self.extract_states(hidden_states, pooling_metadata)
748 + pooling_params = get_pooling_params(pooling_metadata)
749 + assert len(pooled_data) == len(pooling_params)
750 +
751 + pooled_data = [self.head(d, p) for d, p in zip(pooled_data, pooling_params)]
752 + return pooled_data
753 +
754 +
755 class DispatchPooler(Pooler):
756 """Dispatches calls to a sub-pooler based on the pooling task."""
757
vllm/model_executor/models/adapters.py
+15 -27
@@ -250,7 +250,7 @@ def as_embedding_model(cls: _T) -> _T:
250
251 self.pooler = DispatchPooler(
252 {
253 - "encode": Pooler.for_encode(pooler_config),
253 + "token_embed": Pooler.for_token_embed(pooler_config),
254 "embed": Pooler.for_embed(pooler_config),
255 },
256 )
@@ -279,11 +279,8 @@ def as_seq_cls_model(cls: _T) -> _T:
279 # Lazy import
280 from vllm.model_executor.layers.linear import ReplicatedLinear
281 from vllm.model_executor.layers.pooler import (
282 - ClassifierPooler,
282 DispatchPooler,
283 Pooler,
285 - PoolingMethod,
286 - PoolingType,
284 )
285 from vllm.model_executor.models.interfaces import SupportsCrossEncoding
286 from vllm.sequence import IntermediateTensors
@@ -302,42 +299,29 @@ def as_seq_cls_model(cls: _T) -> _T:
299 model_config.hidden_size,
300 config.num_labels,
301 bias=False,
305 - params_dtype=torch.float32,
302 + params_dtype=vllm_config.model_config.head_dtype,
303 quant_config=quant_config,
304 + return_bias=False,
305 prefix=maybe_prefix(prefix, "score"),
306 )
307
308 pooler_config = vllm_config.model_config.pooler_config
309 assert pooler_config is not None
310
313 - pooling_type_str = pooler_config.pooling_type
314 - assert pooling_type_str is not None
315 - pooling_type = PoolingType[pooling_type_str]
316 -
311 self.pooler = DispatchPooler(
312 {
319 - "encode": Pooler.for_encode(pooler_config),
320 - "classify": ClassifierPooler(
321 - pooling=PoolingMethod.from_pooling_type(pooling_type),
322 - classifier=self._classifier,
323 - act_fn=ClassifierPooler.act_fn_for_seq_cls(
324 - vllm_config.model_config
325 - ),
313 + "token_classify": Pooler.for_token_classify(
314 + pooler_config, classifier=self.score
315 + ),
316 + "classify": Pooler.for_classify(
317 + pooler_config, classifier=self.score, act_fn="classify"
318 ),
327 - "score": ClassifierPooler(
328 - pooling=PoolingMethod.from_pooling_type(pooling_type),
329 - classifier=self._classifier,
330 - act_fn=ClassifierPooler.act_fn_for_cross_encoder(
331 - vllm_config.model_config
332 - ),
319 + "score": Pooler.for_classify(
320 + pooler_config, classifier=self.score, act_fn="score"
321 ),
322 }
323 )
324
337 - def _classifier(self, x: torch.Tensor):
338 - x, _ = self.score(x.float())
339 - return x
340 -
325 def forward(
326 self,
327 input_ids: torch.Tensor,
@@ -393,7 +377,11 @@ def as_reward_model(cls: _T) -> _T:
377 assert pooler_config is not None
378
379 self.pooler = DispatchPooler(
396 - {"encode": Pooler.for_encode(pooler_config)},
380 + {
381 + "token_classify": Pooler.for_token_classify(
382 + pooler_config=pooler_config
383 + )
384 + }
385 )
386
387 ModelForReward.__name__ = _get_pooling_model_name(cls.__name__, "ForReward")
vllm/model_executor/models/bert.py
+10 -12
@@ -521,7 +521,7 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):
521 def _build_pooler(self, pooler_config: PoolerConfig) -> Pooler:
522 return DispatchPooler(
523 {
524 - "encode": Pooler.for_encode(pooler_config),
524 + "token_embed": Pooler.for_token_embed(pooler_config),
525 "embed": Pooler.for_embed(pooler_config),
526 }
527 )
@@ -724,7 +724,7 @@ class BertSpladeSparseEmbeddingModel(BertEmbeddingModel):
724
725 return DispatchPooler(
726 {
727 - "encode": Pooler.for_encode(pooler_config),
727 + "token_embed": Pooler.for_token_embed(pooler_config),
728 "embed": SPLADESparsePooler(
729 mlm_head=self.mlm_head,
730 cls_token_id=cls_id,
@@ -821,20 +821,16 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding, SupportsQu
821
822 self.pooler = DispatchPooler(
823 {
824 - "encode": Pooler.for_encode(pooler_config),
824 + "token_classify": Pooler.for_token_classify(
825 + pooler_config, classifier=self.classifier
826 + ),
827 "classify": ClassifierPooler(
828 pooling=self.bert.pooler,
829 classifier=self.classifier,
828 - act_fn=ClassifierPooler.act_fn_for_seq_cls(
829 - vllm_config.model_config
830 - ),
830 + act_fn="classify",
831 ),
832 "score": ClassifierPooler(
833 - pooling=self.bert.pooler,
834 - classifier=self.classifier,
835 - act_fn=ClassifierPooler.act_fn_for_cross_encoder(
836 - vllm_config.model_config
837 - ),
833 + pooling=self.bert.pooler, classifier=self.classifier, act_fn="score"
834 ),
835 }
836 )
@@ -891,7 +887,9 @@ class BertForTokenClassification(nn.Module):
887
888 self.pooler = DispatchPooler(
889 {
894 - "encode": Pooler.for_encode(pooler_config),
890 + "token_classify": Pooler.for_token_classify(
891 + pooler_config=pooler_config
892 + ),
893 }
894 )
895
vllm/model_executor/models/bert_with_rope.py
+5 -9
@@ -695,20 +695,16 @@ class GteNewForSequenceClassification(nn.Module, SupportsCrossEncoding):
695
696 self.pooler = DispatchPooler(
697 {
698 - "encode": Pooler.for_encode(pooler_config),
698 + "token_classify": Pooler.for_token_classify(
699 + pooler_config, classifier=self.classifier
700 + ),
701 "classify": ClassifierPooler(
702 pooling=self.new.pooler,
703 classifier=self.classifier,
702 - act_fn=ClassifierPooler.act_fn_for_seq_cls(
703 - vllm_config.model_config
704 - ),
704 + act_fn="classify",
705 ),
706 "score": ClassifierPooler(
707 - pooling=self.new.pooler,
708 - classifier=self.classifier,
709 - act_fn=ClassifierPooler.act_fn_for_cross_encoder(
710 - vllm_config.model_config
711 - ),
707 + pooling=self.new.pooler, classifier=self.classifier, act_fn="score"
708 ),
709 }
710 )
vllm/model_executor/models/clip.py
+1 -1
@@ -837,7 +837,7 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
837
838 self.pooler = DispatchPooler(
839 {
840 - "encode": Pooler.for_encode(pooler_config),
840 + "token_embed": Pooler.for_token_embed(pooler_config),
841 "embed": Pooler.for_embed(pooler_config),
842 }
843 )
vllm/model_executor/models/gpt2.py
+9 -2
@@ -353,8 +353,15 @@ class GPT2ForSequenceClassification(nn.Module, SupportsCrossEncoding):
353
354 self.pooler = DispatchPooler(
355 {
356 - "encode": Pooler.for_encode(pooler_config),
357 - "classify": Pooler.for_classify(pooler_config, classifier=self.score),
356 + "token_classify": Pooler.for_token_classify(
357 + pooler_config, classifier=self.score
358 + ),
359 + "classify": Pooler.for_classify(
360 + pooler_config, classifier=self.score, act_fn="classify"
361 + ),
362 + "score": Pooler.for_classify(
363 + pooler_config, classifier=self.score, act_fn="score"
364 + ),
365 }
366 )
367
vllm/model_executor/models/gritlm.py
+1 -1
@@ -239,7 +239,7 @@ class GritLM(LlamaForCausalLM):
239 if pooler_config is not None:
240 self.pooler = DispatchPooler(
241 {
242 - "encode": Pooler.for_encode(pooler_config),
242 + "token_embed": Pooler.for_token_embed(pooler_config),
243 "embed": GritLMPooler(vllm_config.model_config),
244 }
245 )
vllm/model_executor/models/internlm2.py
+1 -1
@@ -444,7 +444,7 @@ class InternLM2ForRewardModel(InternLM2ForCausalLM):
444 assert pooler_config is not None
445
446 self.pooler = DispatchPooler(
447 - {"encode": Pooler.for_encode(pooler_config)},
447 + {"token_classify": Pooler.for_token_classify(pooler_config)}
448 )
449
450 def forward(
vllm/model_executor/models/jamba.py
+7 -3
@@ -604,10 +604,14 @@ class JambaForSequenceClassification(JambaForCausalLM):
604
605 self.pooler = DispatchPooler(
606 {
607 - "encode": Pooler.for_encode(pooler_config),
607 + "token_classify": Pooler.for_token_classify(
608 + pooler_config, classifier=self.score
609 + ),
610 "classify": Pooler.for_classify(
609 - pooler_config,
610 - classifier=self.score,
611 + pooler_config, classifier=self.score, act_fn="classify"
612 + ),
613 + "score": Pooler.for_classify(
614 + pooler_config, classifier=self.score, act_fn="score"
615 ),
616 }
617 )
vllm/model_executor/models/jina_vl.py
+9 -3
@@ -97,9 +97,15 @@ class JinaVLForSequenceClassification(
97 self.score = JinaVLScorer(vllm_config.model_config)
98 self.pooler = DispatchPooler(
99 {
100 - "encode": Pooler.for_encode(pooler_config),
101 - "classify": Pooler.for_classify(pooler_config, classifier=self.score),
102 - "score": Pooler.for_classify(pooler_config, classifier=self.score),
100 + "token_classify": Pooler.for_token_classify(
101 + pooler_config, classifier=self.score
102 + ),
103 + "classify": Pooler.for_classify(
104 + pooler_config, classifier=self.score, act_fn="classify"
105 + ),
106 + "score": Pooler.for_classify(
107 + pooler_config, classifier=self.score, act_fn="score"
108 + ),
109 }
110 )
111
vllm/model_executor/models/modernbert.py
+8 -12
@@ -322,20 +322,14 @@ class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):
322
323 self.pooler = DispatchPooler(
324 {
325 - "encode": Pooler.for_encode(pooler_config),
325 + "token_classify": Pooler.for_token_classify(
326 + pooler_config, classifier=self.classifier
327 + ),
328 "classify": ClassifierPooler(
327 - pooling=self.pooling,
328 - classifier=self.classifier,
329 - act_fn=ClassifierPooler.act_fn_for_seq_cls(
330 - vllm_config.model_config
331 - ),
329 + pooling=self.pooling, classifier=self.classifier, act_fn="classify"
330 ),
331 "score": ClassifierPooler(
334 - pooling=self.pooling,
335 - classifier=self.classifier,
336 - act_fn=ClassifierPooler.act_fn_for_cross_encoder(
337 - vllm_config.model_config
338 - ),
332 + pooling=self.pooling, classifier=self.classifier, act_fn="score"
333 ),
334 }
335 )
@@ -421,7 +415,9 @@ class ModernBertForTokenClassification(nn.Module):
415
416 self.pooler = DispatchPooler(
417 {
424 - "encode": Pooler.for_encode(pooler_config),
418 + "token_classify": Pooler.for_token_classify(
419 + pooler_config=pooler_config
420 + ),
421 }
422 )
423
vllm/model_executor/models/qwen2_rm.py
+4 -2
@@ -107,7 +107,7 @@ class Qwen2ForRewardModel(Qwen2RewardBaseModel):
107 assert pooler_config is not None
108
109 self.pooler = DispatchPooler(
110 - {"encode": Pooler.for_encode(pooler_config)},
110 + {"token_classify": Pooler.for_token_classify(pooler_config)}
111 )
112
113
@@ -120,4 +120,6 @@ class Qwen2ForProcessRewardModel(Qwen2RewardBaseModel):
120 pooler_config = vllm_config.model_config.pooler_config
121 assert pooler_config is not None
122
123 - self.pooler = DispatchPooler({"encode": Pooler.for_encode(pooler_config)})
123 + self.pooler = DispatchPooler(
124 + {"token_classify": Pooler.for_token_classify(pooler_config)}
125 + )
vllm/model_executor/models/roberta.py
+6 -20
@@ -105,15 +105,7 @@ class RobertaClassificationHead(nn.Module):
105
106 @default_pooling_type("CLS")
107 class RobertaEmbeddingModel(BertEmbeddingModel):
108 - """A model that uses Roberta to provide embedding functionalities.
109 -
110 - This class encapsulates the BertModel and provides an interface for
111 - embedding operations and customized pooling functions.
112 -
113 - Attributes:
114 - model: An instance of BertModel used for forward operations.
115 - _pooler: An instance of Pooler used for pooling operations.
116 - """
108 + """A model that uses Roberta to provide embedding functionalities."""
109
110 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
111 super().__init__(vllm_config=vllm_config, prefix=prefix)
@@ -212,20 +204,14 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):
204
205 self.pooler = DispatchPooler(
206 {
215 - "encode": Pooler.for_encode(pooler_config),
207 + "token_classify": Pooler.for_token_classify(
208 + pooler_config=pooler_config, classifier=self.classifier
209 + ),
210 "classify": ClassifierPooler(
217 - pooling=CLSPool(),
218 - classifier=self.classifier,
219 - act_fn=ClassifierPooler.act_fn_for_seq_cls(
220 - vllm_config.model_config
221 - ),
211 + pooling=CLSPool(), classifier=self.classifier, act_fn="classify"
212 ),
213 "score": ClassifierPooler(
224 - pooling=CLSPool(),
225 - classifier=self.classifier,
226 - act_fn=ClassifierPooler.act_fn_for_cross_encoder(
227 - vllm_config.model_config
228 - ),
214 + pooling=CLSPool(), classifier=self.classifier, act_fn="score"
215 ),
216 }
217 )
vllm/model_executor/models/terratorch.py
+1 -1
@@ -250,7 +250,7 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):
250 assert pooler_config is not None
251
252 self.pooler = DispatchPooler(
253 - {"encode": Pooler.for_encode(pooler_config)},
253 + {"token_classify": Pooler.for_token_classify(pooler_config)}
254 )
255
256 def get_input_embeddings(
vllm/model_executor/models/transformers_pooling.py
+6 -12
@@ -135,7 +135,7 @@ class TransformersEmbeddingModel(TransformersPoolingBase):
135
136 self.pooler = DispatchPooler(
137 {
138 - "encode": Pooler.for_encode(pooler_config),
138 + "token_embed": Pooler.for_token_embed(pooler_config),
139 "embed": Pooler.for_embed(pooler_config),
140 }
141 )
@@ -190,20 +190,14 @@ class TransformersForSequenceClassification(TransformersPoolingBase):
190
191 self.pooler = DispatchPooler(
192 {
193 - "encode": Pooler.for_encode(pooler_config),
193 + "token_classify": Pooler.for_token_classify(
194 + pooler_config, classifier=self.classifier
195 + ),
196 "classify": ClassifierPooler(
195 - pooling=CLSPool(),
196 - classifier=self.classifier,
197 - act_fn=ClassifierPooler.act_fn_for_seq_cls(
198 - vllm_config.model_config
199 - ),
197 + pooling=CLSPool(), classifier=self.classifier, act_fn="classify"
198 ),
199 "score": ClassifierPooler(
202 - pooling=CLSPool(),
203 - classifier=self.classifier,
204 - act_fn=ClassifierPooler.act_fn_for_cross_encoder(
205 - vllm_config.model_config
206 - ),
200 + pooling=CLSPool(), classifier=self.classifier, act_fn="score"
201 ),
202 }
203 )
vllm/pooling_params.py
+34 -27
@@ -10,7 +10,7 @@ from vllm.sampling_params import RequestOutputKind
10 from vllm.tasks import PoolingTask
11
12 if TYPE_CHECKING:
13 - from vllm.config import ModelConfig
13 + from vllm.config import ModelConfig, PoolerConfig
14
15
16 class PoolingParams(
@@ -30,7 +30,6 @@ class PoolingParams(
30 if model support matryoshka representation.
31 activation: Whether to apply activation function to
32 the classification outputs.
33 - softmax: Whether to apply softmax to the reward outputs.
33 """
34
35 # --8<-- [start:common-pooling-params]
@@ -48,32 +47,19 @@ class PoolingParams(
47 activation: bool | None = None
48 # --8<-- [end:classification-pooling-params]
49
51 - ## for reward models
52 - softmax: bool | None = None
50 + ## for step pooling models
51 step_tag_id: int | None = None
52 returned_token_ids: list[int] | None = None
53
54 + ## Internal use only
55 task: PoolingTask | None = None
57 - """Internal use only."""
58 -
56 requires_token_ids: bool = False
60 - """Internal use only."""
61 -
57 extra_kwargs: dict[str, Any] | None = None
63 - """Internal use only."""
64 -
58 output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY
59
60 @property
61 def all_parameters(self) -> list[str]:
69 - return [
70 - "dimensions",
71 - "normalize",
72 - "activation",
73 - "softmax",
74 - "step_tag_id",
75 - "returned_token_ids",
76 - ]
62 + return ["dimensions", "normalize", "activation"]
63
64 @property
65 def valid_parameters(self):
@@ -81,7 +67,8 @@ class PoolingParams(
67 "embed": ["dimensions", "normalize"],
68 "classify": ["activation"],
69 "score": ["activation"],
84 - "encode": ["softmax", "step_tag_id", "returned_token_ids"],
70 + "token_embed": ["dimensions", "normalize"],
71 + "token_classify": ["activation"],
72 }
73
74 def clone(self) -> "PoolingParams":
@@ -100,7 +87,6 @@ class PoolingParams(
87 # NOTE: Task validation needs to done against the model instance,
88 # which is not available in model config. So, it's not included
89 # in this method
103 -
90 self._merge_default_parameters(model_config)
91 self._set_default_parameters(model_config)
92 self._verify_valid_parameters()
@@ -125,8 +111,34 @@ class PoolingParams(
111 if getattr(self, k, None) is None:
112 setattr(self, k, getattr(pooler_config, k))
113
114 + self._verify_step_pooling(pooler_config, valid_parameters)
115 +
116 + def _verify_step_pooling(
117 + self, pooler_config: "PoolerConfig", valid_parameters: list[str]
118 + ):
119 + step_pooling_parameters = ["step_tag_id", "returned_token_ids"]
120 + if pooler_config.pooling_type != "STEP":
121 + invalid_parameters = []
122 + for k in step_pooling_parameters:
123 + if getattr(self, k, None) is not None:
124 + invalid_parameters.append(k)
125 +
126 + if invalid_parameters:
127 + raise ValueError(
128 + f"Task {self.task} only supports {valid_parameters} "
129 + f"parameters, does not support "
130 + f"{invalid_parameters} parameters"
131 + )
132 + else:
133 + for k in step_pooling_parameters:
134 + if getattr(pooler_config, k, None) is None:
135 + continue
136 +
137 + if getattr(self, k, None) is None:
138 + setattr(self, k, getattr(pooler_config, k))
139 +
140 def _set_default_parameters(self, model_config: Optional["ModelConfig"]):
129 - if self.task == "embed":
141 + if self.task in ["embed", "token_embed"]:
142 if self.normalize is None:
143 self.normalize = True
144
@@ -150,13 +162,9 @@ class PoolingParams(
162 elif self.dimensions < 1:
163 raise ValueError("Dimensions must be greater than 0")
164
153 - elif self.task in ["classify", "score"]:
165 + elif self.task in ["classify", "score", "token_classify"]:
166 if self.activation is None:
167 self.activation = True
156 -
157 - elif self.task == "encode":
158 - if self.softmax is None:
159 - self.softmax = True
168 else:
169 raise ValueError(f"Unknown pooling task: {self.task}")
170
@@ -185,7 +193,6 @@ class PoolingParams(
193 f"normalize={self.normalize}, "
194 f"dimensions={self.dimensions}, "
195 f"activation={self.activation}, "
188 - f"softmax={self.softmax}, "
196 f"step_tag_id={self.step_tag_id}, "
197 f"returned_token_ids={self.returned_token_ids}, "
198 f"requires_token_ids={self.requires_token_ids}, "
vllm/tasks.py
+1 -1
@@ -5,7 +5,7 @@ from typing import Literal, get_args
5 GenerationTask = Literal["generate", "transcription"]
6 GENERATION_TASKS = get_args(GenerationTask)
7
8 -PoolingTask = Literal["encode", "embed", "classify", "score"]
8 +PoolingTask = Literal["embed", "classify", "score", "token_embed", "token_classify"]
9 POOLING_TASKS = get_args(PoolingTask)
10
11 SupportedTask = Literal[GenerationTask, PoolingTask]
vllm/v1/worker/gpu_model_runner.py
+7 -6
@@ -1926,15 +1926,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1926
1927 supported_tasks = list(model.pooler.get_supported_tasks())
1928
1929 - if (
1930 - self.scheduler_config.chunked_prefill_enabled
1931 - and "encode" in supported_tasks
1932 - ):
1933 - supported_tasks.remove("encode")
1929 + if self.scheduler_config.chunked_prefill_enabled:
1930 + if "token_embed" in supported_tasks:
1931 + supported_tasks.remove("token_embed")
1932 + if "token_classify" in supported_tasks:
1933 + supported_tasks.remove("token_classify")
1934
1935 logger.debug_once(
1936 "Chunked prefill is not supported with "
1937 - "encode task which using ALL pooling. "
1937 + "token_embed and token_classify tasks "
1938 + "which using ALL pooling. "
1939 "Please turn off chunked prefill by "
1940 "`--no-enable-chunked-prefill` before using it."
1941 )