@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 586f286789
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] Pooling model activation supports per request control by PoolingParams (#20538)

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

wang.yuqi committed Aug 5, 2025 at 15:37 UTC 586f286789a09f5616be74ee8bedde0a9f698a72
21 files changed +948 -173
tests/entrypoints/llm/test_classify.py new
+67
@@ -0,0 +1,67 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import weakref
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm import LLM, PoolingParams
10 +from vllm.distributed import cleanup_dist_env_and_memory
11 +
12 +from ...models.utils import softmax
13 +
14 +MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
15 +
16 +prompts = ["The chef prepared a delicious meal."]
17 +
18 +
19 +@pytest.fixture(autouse=True)
20 +def v1(run_with_both_engines):
21 + # Simple autouse wrapper to run both engines for each test
22 + # This can be promoted up to conftest.py to run for every
23 + # test in a package
24 + pass
25 +
26 +
27 +@pytest.fixture(scope="module")
28 +def llm():
29 + # pytest caches the fixture so we use weakref.proxy to
30 + # enable garbage collection
31 + llm = LLM(model=MODEL_NAME,
32 + max_num_batched_tokens=32768,
33 + tensor_parallel_size=1,
34 + gpu_memory_utilization=0.75,
35 + enforce_eager=True,
36 + seed=0)
37 +
38 + with llm.deprecate_legacy_api():
39 + yield weakref.proxy(llm)
40 +
41 + del llm
42 +
43 + cleanup_dist_env_and_memory()
44 +
45 +
46 +@pytest.mark.skip_global_cleanup
47 +def test_pooling_params(llm: LLM):
48 +
49 + def get_outputs(activation):
50 + outputs = llm.classify(
51 + prompts,
52 + pooling_params=PoolingParams(activation=activation),
53 + use_tqdm=False)
54 + return torch.tensor([x.outputs.probs for x in outputs])
55 +
56 + default = get_outputs(activation=None)
57 + w_activation = get_outputs(activation=True)
58 + wo_activation = get_outputs(activation=False)
59 +
60 + assert torch.allclose(default, w_activation,
61 + atol=1e-2), "Default should use activation."
62 + assert not torch.allclose(
63 + w_activation, wo_activation,
64 + atol=1e-2), "wo_activation should not use activation."
65 + assert torch.allclose(
66 + softmax(wo_activation), w_activation, atol=1e-2
67 + ), "w_activation should be close to activation(wo_activation)."
tests/entrypoints/llm/test_embedding.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 +import weakref
5 +
6 +import pytest
7 +import torch
8 +import torch.nn.functional as F
9 +
10 +from vllm import LLM, PoolingParams
11 +from vllm.distributed import cleanup_dist_env_and_memory
12 +
13 +MODEL_NAME = "intfloat/multilingual-e5-small"
14 +
15 +prompts = ["The chef prepared a delicious meal."]
16 +
17 +
18 +@pytest.fixture(scope="module")
19 +def llm():
20 + # pytest caches the fixture so we use weakref.proxy to
21 + # enable garbage collection
22 + llm = LLM(model=MODEL_NAME,
23 + max_num_batched_tokens=32768,
24 + tensor_parallel_size=1,
25 + gpu_memory_utilization=0.75,
26 + enforce_eager=True,
27 + seed=0)
28 +
29 + with llm.deprecate_legacy_api():
30 + yield weakref.proxy(llm)
31 +
32 + del llm
33 +
34 + cleanup_dist_env_and_memory()
35 +
36 +
37 +@pytest.mark.skip_global_cleanup
38 +def test_pooling_params(llm: LLM):
39 +
40 + def get_outputs(normalize):
41 + outputs = llm.embed(prompts,
42 + pooling_params=PoolingParams(normalize=normalize),
43 + use_tqdm=False)
44 + return torch.tensor([x.outputs.embedding for x in outputs])
45 +
46 + default = get_outputs(normalize=None)
47 + w_normal = get_outputs(normalize=True)
48 + wo_normal = get_outputs(normalize=False)
49 +
50 + assert torch.allclose(default, w_normal,
51 + atol=1e-2), "Default should use normal."
52 + assert not torch.allclose(w_normal, wo_normal,
53 + atol=1e-2), "wo_normal should not use normal."
54 + assert torch.allclose(
55 + w_normal, F.normalize(wo_normal, p=2, dim=-1),
56 + atol=1e-2), "w_normal should be close to normal(wo_normal)."
tests/entrypoints/llm/test_reward.py new
+66
@@ -0,0 +1,66 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import weakref
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm import LLM, PoolingParams
10 +from vllm.distributed import cleanup_dist_env_and_memory
11 +
12 +from ...models.utils import softmax
13 +
14 +MODEL_NAME = "internlm/internlm2-1_8b-reward"
15 +
16 +prompts = ["The chef prepared a delicious meal."]
17 +
18 +
19 +@pytest.fixture(autouse=True)
20 +def v1(run_with_both_engines):
21 + # Simple autouse wrapper to run both engines for each test
22 + # This can be promoted up to conftest.py to run for every
23 + # test in a package
24 + pass
25 +
26 +
27 +@pytest.fixture(scope="module")
28 +def llm():
29 + # pytest caches the fixture so we use weakref.proxy to
30 + # enable garbage collection
31 + llm = LLM(model=MODEL_NAME,
32 + max_num_batched_tokens=32768,
33 + tensor_parallel_size=1,
34 + gpu_memory_utilization=0.75,
35 + enforce_eager=True,
36 + trust_remote_code=True,
37 + seed=0)
38 +
39 + with llm.deprecate_legacy_api():
40 + yield weakref.proxy(llm)
41 +
42 + del llm
43 +
44 + cleanup_dist_env_and_memory()
45 +
46 +
47 +@pytest.mark.skip_global_cleanup
48 +def test_pooling_params(llm: LLM):
49 +
50 + def get_outputs(softmax):
51 + outputs = llm.reward(prompts,
52 + pooling_params=PoolingParams(softmax=softmax),
53 + use_tqdm=False)
54 + return torch.cat([x.outputs.data for x in outputs])
55 +
56 + default = get_outputs(softmax=None)
57 + w_softmax = get_outputs(softmax=True)
58 + wo_softmax = get_outputs(softmax=False)
59 +
60 + assert torch.allclose(default, w_softmax,
61 + atol=1e-2), "Default should use softmax."
62 + assert not torch.allclose(w_softmax, wo_softmax,
63 + atol=1e-2), "wo_softmax should not use softmax."
64 + assert torch.allclose(
65 + softmax(wo_softmax), w_softmax,
66 + atol=1e-2), "w_softmax should be close to softmax(wo_softmax)."
tests/entrypoints/llm/test_score.py new
+69
@@ -0,0 +1,69 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import weakref
5 +
6 +import pytest
7 +import torch
8 +
9 +from vllm import LLM, PoolingParams
10 +from vllm.distributed import cleanup_dist_env_and_memory
11 +
12 +from ...models.utils import softmax
13 +
14 +MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
15 +
16 +
17 +@pytest.fixture(autouse=True)
18 +def v1(run_with_both_engines):
19 + # Simple autouse wrapper to run both engines for each test
20 + # This can be promoted up to conftest.py to run for every
21 + # test in a package
22 + pass
23 +
24 +
25 +@pytest.fixture(scope="module")
26 +def llm():
27 + # pytest caches the fixture so we use weakref.proxy to
28 + # enable garbage collection
29 + llm = LLM(model=MODEL_NAME,
30 + max_num_batched_tokens=32768,
31 + tensor_parallel_size=1,
32 + gpu_memory_utilization=0.75,
33 + enforce_eager=True,
34 + seed=0)
35 +
36 + with llm.deprecate_legacy_api():
37 + yield weakref.proxy(llm)
38 +
39 + del llm
40 +
41 + cleanup_dist_env_and_memory()
42 +
43 +
44 +@pytest.mark.skip_global_cleanup
45 +def test_pooling_params(llm: LLM):
46 +
47 + def get_outputs(activation):
48 + text_1 = "What is the capital of France?"
49 + text_2 = "The capital of France is Paris."
50 +
51 + outputs = llm.score(
52 + text_1,
53 + text_2,
54 + pooling_params=PoolingParams(activation=activation),
55 + use_tqdm=False)
56 + return torch.tensor([x.outputs.score for x in outputs])
57 +
58 + default = get_outputs(activation=None)
59 + w_activation = get_outputs(activation=True)
60 + wo_activation = get_outputs(activation=False)
61 +
62 + assert torch.allclose(default, w_activation,
63 + atol=1e-2), "Default should use activation."
64 + assert not torch.allclose(
65 + w_activation, wo_activation,
66 + atol=1e-2), "wo_activation should not use activation."
67 + assert torch.allclose(
68 + softmax(wo_activation), w_activation, atol=1e-2
69 + ), "w_activation should be close to activation(wo_activation)."
tests/entrypoints/openai/test_classification.py
+31
@@ -3,6 +3,8 @@
3
4 import pytest
5 import requests
6 +import torch
7 +import torch.nn.functional as F
8
9 from vllm.entrypoints.openai.protocol import ClassificationResponse
10
@@ -181,3 +183,32 @@ async def test_invocations(server: RemoteOpenAIServer):
183 assert classification_data.keys() == invocation_data.keys()
184 assert classification_data["probs"] == pytest.approx(
185 invocation_data["probs"], rel=0.01)
186 +
187 +
188 +@pytest.mark.asyncio
189 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
190 +async def test_activation(server: RemoteOpenAIServer, model_name: str):
191 + input_text = ["This product was excellent and exceeded my expectations"]
192 +
193 + async def get_outputs(activation):
194 + response = requests.post(server.url_for("classify"),
195 + json={
196 + "model": model_name,
197 + "input": input_text,
198 + "activation": activation
199 + })
200 + outputs = response.json()
201 + return torch.tensor([x['probs'] for x in outputs["data"]])
202 +
203 + default = await get_outputs(activation=None)
204 + w_activation = await get_outputs(activation=True)
205 + wo_activation = await get_outputs(activation=False)
206 +
207 + assert torch.allclose(default, w_activation,
208 + atol=1e-2), "Default should use activation."
209 + assert not torch.allclose(
210 + w_activation, wo_activation,
211 + atol=1e-2), "wo_activation should not use activation."
212 + assert torch.allclose(
213 + F.softmax(wo_activation, dim=-1), w_activation, atol=1e-2
214 + ), "w_activation should be close to activation(wo_activation)."
tests/entrypoints/openai/test_embedding.py
+34
@@ -8,6 +8,8 @@ import openai
8 import pytest
9 import pytest_asyncio
10 import requests
11 +import torch
12 +import torch.nn.functional as F
13
14 from vllm.entrypoints.openai.protocol import EmbeddingResponse
15 from vllm.transformers_utils.tokenizer import get_tokenizer
@@ -369,3 +371,35 @@ async def test_invocations_conversation(server: RemoteOpenAIServer):
371 embeddings_1_lst=[invocation_data["embedding"]],
372 name_0="chat",
373 name_1="invocation")
374 +
375 +
376 +@pytest.mark.asyncio
377 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
378 +async def test_normalize(server: RemoteOpenAIServer, model_name: str):
379 + input_text = ["The chef prepared a delicious meal."]
380 +
381 + async def get_outputs(normalize):
382 + request_args = {
383 + "model": MODEL_NAME,
384 + "input": input_text,
385 + "encoding_format": "float",
386 + "normalize": normalize
387 + }
388 +
389 + response = requests.post(server.url_for("v1/embeddings"),
390 + json=request_args)
391 + outputs = response.json()
392 +
393 + return torch.tensor([x['embedding'] for x in outputs["data"]])
394 +
395 + default = await get_outputs(normalize=None)
396 + w_normal = await get_outputs(normalize=True)
397 + wo_normal = await get_outputs(normalize=False)
398 +
399 + assert torch.allclose(default, w_normal,
400 + atol=1e-2), "Default should use normal."
401 + assert not torch.allclose(w_normal, wo_normal,
402 + atol=1e-2), "wo_normal should not use normal."
403 + assert torch.allclose(
404 + w_normal, F.normalize(wo_normal, p=2, dim=-1),
405 + atol=1e-2), "w_normal should be close to normal(wo_normal)."
tests/entrypoints/openai/test_rerank.py
+38
@@ -3,6 +3,8 @@
3
4 import pytest
5 import requests
6 +import torch
7 +import torch.nn.functional as F
8
9 from vllm.entrypoints.openai.protocol import RerankResponse
10
@@ -125,3 +127,39 @@ def test_invocations(server: RemoteOpenAIServer):
127 assert rerank_result.keys() == invocations_result.keys()
128 assert rerank_result["relevance_score"] == pytest.approx(
129 invocations_result["relevance_score"], rel=0.01)
130 +
131 +
132 +@pytest.mark.asyncio
133 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
134 +async def test_activation(server: RemoteOpenAIServer, model_name: str):
135 +
136 + async def get_outputs(activation):
137 + query = "What is the capital of France?"
138 + documents = [
139 + "The capital of Brazil is Brasilia.",
140 + "The capital of France is Paris."
141 + ]
142 +
143 + response = requests.post(server.url_for("rerank"),
144 + json={
145 + "model": model_name,
146 + "query": query,
147 + "documents": documents,
148 + "activation": activation
149 + })
150 + outputs = response.json()
151 +
152 + return torch.tensor([x['relevance_score'] for x in outputs["results"]])
153 +
154 + default = await get_outputs(activation=None)
155 + w_activation = await get_outputs(activation=True)
156 + wo_activation = await get_outputs(activation=False)
157 +
158 + assert torch.allclose(default, w_activation,
159 + atol=1e-2), "Default should use activation."
160 + assert not torch.allclose(
161 + w_activation, wo_activation,
162 + atol=1e-2), "wo_activation should not use activation."
163 + assert torch.allclose(
164 + F.sigmoid(wo_activation), w_activation, atol=1e-2
165 + ), "w_activation should be close to activation(wo_activation)."
tests/entrypoints/openai/test_score.py
+41
@@ -4,6 +4,7 @@ from typing import Any
4
5 import pytest
6 import requests
7 +import torch
8 import torch.nn.functional as F
9 from torch import tensor
10
@@ -220,3 +221,43 @@ class TestModel:
221 assert score_data.keys() == invocation_data.keys()
222 assert score_data["score"] == pytest.approx(
223 invocation_data["score"], rel=0.01)
224 +
225 + def test_activation(self, server: RemoteOpenAIServer, model: dict[str,
226 + Any]):
227 +
228 + def get_outputs(activation):
229 + text_1 = "What is the capital of France?"
230 + text_2 = "The capital of France is Paris."
231 + response = requests.post(server.url_for("score"),
232 + json={
233 + "model": model["name"],
234 + "text_1": text_1,
235 + "text_2": text_2,
236 + "activation": activation
237 + })
238 + if response.status_code != 200:
239 + return response
240 +
241 + outputs = response.json()
242 + return torch.tensor([x['score'] for x in outputs["data"]])
243 +
244 + if model["is_cross_encoder"]:
245 +
246 + default = get_outputs(activation=None)
247 + w_activation = get_outputs(activation=True)
248 + wo_activation = get_outputs(activation=False)
249 +
250 + assert torch.allclose(default, w_activation,
251 + atol=1e-2), "Default should use activation."
252 + assert not torch.allclose(
253 + w_activation, wo_activation,
254 + atol=1e-2), "wo_activation should not use activation."
255 + assert torch.allclose(
256 + F.sigmoid(wo_activation), w_activation, atol=1e-2
257 + ), "w_activation should be close to activation(wo_activation)."
258 + else:
259 + get_outputs(activation=None)
260 +
261 + # The activation parameter only works for the is_cross_encoder model
262 + response = get_outputs(activation=True)
263 + assert response.status_code == 400
tests/models/language/pooling/test_override_pooler_config.py new
+127
@@ -0,0 +1,127 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import pytest
4 +import torch
5 +import torch.nn.functional as F
6 +
7 +from tests.models.utils import softmax
8 +from vllm.config import PoolerConfig
9 +
10 +
11 +@pytest.mark.parametrize(
12 + "model",
13 + [
14 + "jason9693/Qwen2.5-1.5B-apeach",
15 + "papluca/xlm-roberta-base-language-detection"
16 + ],
17 +)
18 +@pytest.mark.parametrize("dtype", ["half"])
19 +def test_classify_models_using_activation(
20 + hf_runner,
21 + vllm_runner,
22 + example_prompts,
23 + model: str,
24 + dtype: str,
25 +) -> None:
26 +
27 + with vllm_runner(model,
28 + max_model_len=512,
29 + dtype=dtype,
30 + override_pooler_config=PoolerConfig(
31 + activation=False)) as vllm_model:
32 + wo_activation_out = vllm_model.classify(example_prompts)
33 +
34 + with vllm_runner(model,
35 + max_model_len=512,
36 + dtype=dtype,
37 + override_pooler_config=PoolerConfig(
38 + activation=True)) as vllm_model:
39 + w_activation_out = vllm_model.classify(example_prompts)
40 +
41 + for wo_activation, w_activation in zip(wo_activation_out,
42 + w_activation_out):
43 + wo_activation = torch.tensor(wo_activation)
44 + w_activation = torch.tensor(w_activation)
45 +
46 + assert not torch.allclose(
47 + wo_activation, w_activation,
48 + atol=1e-2), "override_pooler_config is not working"
49 + assert torch.allclose(softmax(wo_activation), w_activation,
50 + 1e-3 if dtype == "float" else 1e-2)
51 +
52 +
53 +@pytest.mark.parametrize(
54 + "model",
55 + [
56 + "intfloat/multilingual-e5-small",
57 + ],
58 +)
59 +@pytest.mark.parametrize("dtype", ["half"])
60 +def test_embed_models_using_normalize(
61 + hf_runner,
62 + vllm_runner,
63 + example_prompts,
64 + model: str,
65 + dtype: str,
66 +) -> None:
67 +
68 + with vllm_runner(model,
69 + max_model_len=512,
70 + dtype=dtype,
71 + override_pooler_config=PoolerConfig(
72 + normalize=False)) as vllm_model:
73 + wo_normalize = torch.tensor(vllm_model.embed(example_prompts))
74 +
75 + with vllm_runner(
76 + model,
77 + max_model_len=512,
78 + dtype=dtype,
79 + override_pooler_config=PoolerConfig(normalize=True)) as vllm_model:
80 + w_normalize = torch.tensor(vllm_model.embed(example_prompts))
81 +
82 + assert not torch.allclose(
83 + wo_normalize, w_normalize,
84 + atol=1e-2), "override_pooler_config normalize is not working"
85 + assert torch.allclose(
86 + F.normalize(wo_normalize, p=2, dim=-1), w_normalize,
87 + atol=1e-2), "w_normal should be close to normal(wo_normal)."
88 +
89 +
90 +@pytest.mark.parametrize(
91 + "model",
92 + [
93 + "internlm/internlm2-1_8b-reward",
94 + ],
95 +)
96 +@pytest.mark.parametrize("dtype", ["half"])
97 +def test_reward_models_using_softmax(
98 + hf_runner,
99 + vllm_runner,
100 + example_prompts,
101 + model: str,
102 + dtype: str,
103 +) -> None:
104 +
105 + with vllm_runner(
106 + model,
107 + max_model_len=1024,
108 + dtype=dtype,
109 + override_pooler_config=PoolerConfig(softmax=False)) as vllm_model:
110 + wo_softmax = vllm_model.encode(example_prompts)
111 +
112 + with vllm_runner(
113 + model,
114 + max_model_len=1024,
115 + dtype=dtype,
116 + override_pooler_config=PoolerConfig(softmax=True)) as vllm_model:
117 + w_softmax = vllm_model.encode(example_prompts)
118 +
119 + for wo, w in zip(wo_softmax, w_softmax):
120 + wo = torch.tensor(wo)
121 + w = torch.tensor(w)
122 +
123 + assert not torch.allclose(
124 + wo, w, atol=1e-2), "override_pooler_config softmax is not working"
125 + assert torch.allclose(
126 + softmax(wo), w,
127 + atol=1e-2), "w_softmax should be close to softmax(wo_softmax)."
tests/models/language/pooling/test_reward.py
+2 -2
@@ -103,7 +103,7 @@ def test_prm_models(
103
104 # check logits difference
105 for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
106 - hf_output = torch.tensor(hf_output)
107 - vllm_output = torch.tensor(vllm_output)
106 + hf_output = torch.tensor(hf_output).float()
107 + vllm_output = torch.tensor(vllm_output).float()
108
109 assert torch.allclose(hf_output, vllm_output, 1.5e-2)
tests/models/utils.py
+7
@@ -330,6 +330,13 @@ def matryoshka_fy(tensor: torch.Tensor, dimensions: int):
330 return tensor
331
332
333 +def softmax(data):
334 + if data.shape[-1] == 1:
335 + return F.sigmoid(data)
336 + else:
337 + return F.softmax(data, dim=-1)
338 +
339 +
340 class EmbedModelInfo(NamedTuple):
341 name: str
342 is_matryoshka: bool = False
tests/test_pooling_params.py new
+106
@@ -0,0 +1,106 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import pytest
4 +
5 +from tests.models.utils import EmbedModelInfo
6 +from vllm import PoolingParams
7 +from vllm.config import ModelConfig
8 +
9 +EMBEDDING_MODELS = [
10 + EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False),
11 + EmbedModelInfo("Snowflake/snowflake-arctic-embed-m-v1.5",
12 + is_matryoshka=True,
13 + matryoshka_dimensions=[256]),
14 +]
15 +
16 +
17 +def test_task():
18 + pooling_params = PoolingParams()
19 + pooling_params.verify(task="score")
20 +
21 + pooling_params = PoolingParams(task="score")
22 + pooling_params.verify(task="score")
23 +
24 + with pytest.raises(ValueError):
25 + pooling_params.verify(task="encode")
26 +
27 +
28 +def test_embed():
29 + task = "embed"
30 + pooling_params = PoolingParams(normalize=None)
31 + pooling_params.verify(task=task)
32 +
33 + pooling_params = PoolingParams(normalize=True)
34 + pooling_params.verify(task=task)
35 +
36 + pooling_params = PoolingParams(normalize=False)
37 + pooling_params.verify(task=task)
38 +
39 + invalid_parameters = ["activation", "softmax"]
40 + for p in invalid_parameters:
41 + with pytest.raises(ValueError):
42 + pooling_params = PoolingParams(**{p: True})
43 + pooling_params.verify(task=task)
44 +
45 +
46 +@pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
47 +def test_embed_dimensions(model_info: EmbedModelInfo):
48 + task = "embed"
49 + model_config = ModelConfig(
50 + model_info.name,
51 + task="auto",
52 + tokenizer=model_info.name,
53 + tokenizer_mode="auto",
54 + trust_remote_code=False,
55 + seed=0,
56 + dtype="float16",
57 + )
58 +
59 + pooling_params = PoolingParams(dimensions=None)
60 + pooling_params.verify(task=task, model_config=model_config)
61 +
62 + with pytest.raises(ValueError):
63 + pooling_params = PoolingParams(dimensions=1)
64 + pooling_params.verify(task=task, model_config=model_config)
65 +
66 + if model_info.is_matryoshka:
67 + assert model_info.matryoshka_dimensions is not None
68 + pooling_params = PoolingParams(
69 + dimensions=model_info.matryoshka_dimensions[0])
70 + pooling_params.verify(task=task, model_config=model_config)
71 +
72 +
73 +@pytest.mark.parametrize("task", ["score", "classify"])
74 +def test_classify(task):
75 + pooling_params = PoolingParams(activation=None)
76 + pooling_params.verify(task=task)
77 +
78 + pooling_params = PoolingParams(activation=True)
79 + pooling_params.verify(task=task)
80 +
81 + pooling_params = PoolingParams(activation=False)
82 + pooling_params.verify(task=task)
83 +
84 + invalid_parameters = ["dimensions", "normalize", "softmax"]
85 + for p in invalid_parameters:
86 + with pytest.raises(ValueError):
87 + pooling_params = PoolingParams(**{p: True})
88 + pooling_params.verify(task=task)
89 +
90 +
91 +def test_encode():
92 + task = "encode"
93 + pooling_params = PoolingParams(softmax=None)
94 + pooling_params.verify(task=task)
95 +
96 + pooling_params = PoolingParams(softmax=True)
97 + pooling_params.verify(task=task)
98 +
99 + pooling_params = PoolingParams(softmax=False)
100 + pooling_params.verify(task=task)
101 +
102 + invalid_parameters = ["dimensions", "normalize", "activation"]
103 + for p in invalid_parameters:
104 + with pytest.raises(ValueError):
105 + pooling_params = PoolingParams(**{p: True})
106 + pooling_params.verify(task=task)
vllm/config.py
+15 -15
@@ -913,15 +913,6 @@ class ModelConfig:
913 if getattr(pooler_config, k) is None:
914 setattr(pooler_config, k, v)
915
916 - if self.is_matryoshka:
917 - if pooler_config.normalize is None:
918 - pooler_config.normalize = True
919 - elif not pooler_config.normalize:
920 - raise ValueError(
921 - "`normalize` must be enabled (set to True) "
922 - "for models that are compatible with "
923 - "Matryoshka Representation.")
924 -
916 return pooler_config
917
918 return None
@@ -3438,25 +3429,34 @@ class PoolerConfig:
3429 [`vllm.model_executor.layers.pooler.PoolingType`][].
3430 """
3431
3432 + ## for embeddings models
3433 normalize: Optional[bool] = None
3434 """
3443 - Whether to normalize the pooled outputs. Usually, this should be set to
3444 - ``True`` for embedding outputs.
3435 + Whether to normalize the embeddings outputs.
3436 + """
3437 + dimensions: Optional[int] = None
3438 + """
3439 + Reduce the dimensions of embeddings if model
3440 + support matryoshka representation.
3441 """
3442
3447 - softmax: Optional[bool] = None
3443 + ## for classification models
3444 + activation: Optional[bool] = None
3445 """
3449 - Whether to apply softmax to the pooled outputs. Usually, this should be set
3450 - to ``True`` for classification outputs.
3446 + Whether to apply activation function to the classification outputs.
3447 """
3448
3449 + ## for reward models
3450 + softmax: Optional[bool] = None
3451 + """
3452 + Whether to apply softmax to the reward outputs.
3453 + """
3454 step_tag_id: Optional[int] = None
3455 """
3456 If set, only the score corresponding to the ``step_tag_id`` in the
3457 generated sentence should be returned. Otherwise, the scores for all tokens
3458 are returned.
3459 """
3459 -
3460 returned_token_ids: Optional[list[int]] = None
3461 """
3462 A list of indices for the vocabulary dimensions to be extracted,
vllm/entrypoints/llm.py
+19 -3
@@ -1189,6 +1189,8 @@ class LLM:
1189 /,
1190 *,
1191 use_tqdm: Union[bool, Callable[..., tqdm]] = True,
1192 + pooling_params: Optional[Union[PoolingParams,
1193 + Sequence[PoolingParams]]] = None,
1194 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
1195 ) -> list[ClassificationRequestOutput]:
1196 """
@@ -1207,7 +1209,8 @@ class LLM:
1209 it is used to create the progress bar.
1210 If `False`, no progress bar is created.
1211 lora_request: LoRA request to use for generation, if any.
1210 -
1212 + pooling_params: The pooling parameters for pooling. If None, we
1213 + use the default pooling parameters.
1214 Returns:
1215 A list of `ClassificationRequestOutput` objects containing the
1216 embedding vectors in the same order as the input prompts.
@@ -1220,6 +1223,7 @@ class LLM:
1223 items = self.encode(
1224 prompts,
1225 use_tqdm=use_tqdm,
1226 + pooling_params=pooling_params,
1227 lora_request=lora_request,
1228 pooling_task="classify",
1229 )
@@ -1272,6 +1276,7 @@ class LLM:
1276 text_2: list[Union[str, TextPrompt, TokensPrompt]],
1277 truncate_prompt_tokens: Optional[int] = None,
1278 use_tqdm: Union[bool, Callable[..., tqdm]] = True,
1279 + pooling_params: Optional[PoolingParams] = None,
1280 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
1281 ) -> list[ScoringRequestOutput]:
1282
@@ -1280,6 +1285,7 @@ class LLM:
1285 truncate_prompt_tokens=truncate_prompt_tokens,
1286 use_tqdm=use_tqdm,
1287 lora_request=lora_request,
1288 + pooling_params=pooling_params,
1289 pooling_task="embed",
1290 )
1291
@@ -1306,6 +1312,7 @@ class LLM:
1312 data_2: Union[list[str], list[ScoreContentPartParam]],
1313 truncate_prompt_tokens: Optional[int] = None,
1314 use_tqdm: Union[bool, Callable[..., tqdm]] = True,
1315 + pooling_params: Optional[PoolingParams] = None,
1316 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
1317 ) -> list[ScoringRequestOutput]:
1318 model_config = self.llm_engine.model_config
@@ -1317,7 +1324,12 @@ class LLM:
1324 if len(data_1) == 1:
1325 data_1 = data_1 * len(data_2)
1326
1320 - pooling_params = PoolingParams(task="score")
1327 + if pooling_params is None:
1328 + pooling_params = PoolingParams(task="score")
1329 +
1330 + model_config = self.llm_engine.model_config
1331 + pooling_params.verify("score", model_config)
1332 +
1333 tokenization_kwargs: dict[str, Any] = {}
1334
1335 _validate_truncation_size(model_config.max_model_len,
@@ -1379,6 +1391,7 @@ class LLM:
1391 *,
1392 truncate_prompt_tokens: Optional[int] = None,
1393 use_tqdm: Union[bool, Callable[..., tqdm]] = True,
1394 + pooling_params: Optional[PoolingParams] = None,
1395 lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
1396 ) -> list[ScoringRequestOutput]:
1397 """Generate similarity scores for all pairs `<text,text_pair>` or
@@ -1410,7 +1423,8 @@ class LLM:
1423 it is used to create the progress bar.
1424 If `False`, no progress bar is created.
1425 lora_request: LoRA request to use for generation, if any.
1413 -
1426 + pooling_params: The pooling parameters for pooling. If None, we
1427 + use the default pooling parameters.
1428 Returns:
1429 A list of `ScoringRequestOutput` objects containing the
1430 generated scores in the same order as the input prompts.
@@ -1494,6 +1508,7 @@ class LLM:
1508 data_2, # type: ignore[arg-type]
1509 truncate_prompt_tokens,
1510 use_tqdm,
1511 + pooling_params,
1512 lora_request)
1513 else:
1514 return self._embedding_score(
@@ -1502,6 +1517,7 @@ class LLM:
1517 data_2, # type: ignore[arg-type]
1518 truncate_prompt_tokens,
1519 use_tqdm,
1520 + pooling_params,
1521 lora_request)
1522
1523 def start_profile(self) -> None:
vllm/entrypoints/openai/protocol.py
+15 -5
@@ -1274,11 +1274,13 @@ class EmbeddingCompletionRequest(OpenAIBaseModel):
1274 "not set it, a random_uuid will be generated. This id is used "
1275 "through out the inference process and return in response."),
1276 )
1277 + normalize: Optional[bool] = None
1278
1279 # --8<-- [end:embedding-extra-params]
1280
1281 def to_pooling_params(self):
1281 - return PoolingParams(dimensions=self.dimensions)
1282 + return PoolingParams(dimensions=self.dimensions,
1283 + normalize=self.normalize)
1284
1285
1286 class EmbeddingChatRequest(OpenAIBaseModel):
@@ -1332,6 +1334,7 @@ class EmbeddingChatRequest(OpenAIBaseModel):
1334 "not set it, a random_uuid will be generated. This id is used "
1335 "through out the inference process and return in response."),
1336 )
1337 + normalize: Optional[bool] = None
1338 # --8<-- [end:chat-embedding-extra-params]
1339
1340 @model_validator(mode="before")
@@ -1344,7 +1347,8 @@ class EmbeddingChatRequest(OpenAIBaseModel):
1347 return data
1348
1349 def to_pooling_params(self):
1347 - return PoolingParams(dimensions=self.dimensions)
1350 + return PoolingParams(dimensions=self.dimensions,
1351 + normalize=self.normalize)
1352
1353
1354 EmbeddingRequest = Union[EmbeddingCompletionRequest, EmbeddingChatRequest]
@@ -1375,10 +1379,12 @@ class ScoreRequest(OpenAIBaseModel):
1379 "if the served model does not use priority scheduling."),
1380 )
1381
1382 + activation: Optional[bool] = None
1383 +
1384 # --8<-- [end:score-extra-params]
1385
1386 def to_pooling_params(self):
1381 - return PoolingParams()
1387 + return PoolingParams(activation=self.activation)
1388
1389
1390 class RerankRequest(OpenAIBaseModel):
@@ -1403,10 +1409,12 @@ class RerankRequest(OpenAIBaseModel):
1409 "if the served model does not use priority scheduling."),
1410 )
1411
1412 + activation: Optional[bool] = None
1413 +
1414 # --8<-- [end:rerank-extra-params]
1415
1416 def to_pooling_params(self):
1409 - return PoolingParams()
1417 + return PoolingParams(activation=self.activation)
1418
1419
1420 class RerankDocument(BaseModel):
@@ -1553,10 +1561,12 @@ class ClassificationRequest(OpenAIBaseModel):
1561 "if the served model does not use priority scheduling."),
1562 )
1563
1564 + activation: Optional[bool] = None
1565 +
1566 # --8<-- [end:classification-extra-params]
1567
1568 def to_pooling_params(self):
1559 - return PoolingParams()
1569 + return PoolingParams(activation=self.activation)
1570
1571
1572 class ClassificationData(OpenAIBaseModel):
vllm/model_executor/layers/pooler.py
+106 -116
@@ -41,35 +41,18 @@ class PoolingType(IntEnum):
41 @dataclass(frozen=True)
42 class ResolvedPoolingConfig:
43 pooling_type: PoolingType
44 -
45 - normalize: bool
46 - softmax: bool
47 - step_tag_id: Optional[int]
48 - returned_token_ids: Optional[list[int]]
44 + task: PoolingTask
45
46 @classmethod
47 def from_config_with_defaults(
48 cls,
49 + task: PoolingTask,
50 pooler_config: PoolerConfig,
51 pooling_type: PoolingType,
55 - normalize: bool,
56 - softmax: bool,
57 - step_tag_id: Optional[int] = None,
58 - returned_token_ids: Optional[list[int]] = None,
52 ) -> "ResolvedPoolingConfig":
60 - return cls(
61 - pooling_type=PoolingType[pooler_config.pooling_type]
62 - if pooler_config.pooling_type is not None else pooling_type,
63 - normalize=pooler_config.normalize
64 - if pooler_config.normalize is not None else normalize,
65 - softmax=pooler_config.softmax
66 - if pooler_config.softmax is not None else softmax,
67 - step_tag_id=pooler_config.step_tag_id
68 - if pooler_config.step_tag_id is not None else step_tag_id,
69 - returned_token_ids=pooler_config.returned_token_ids
70 - if pooler_config.returned_token_ids is not None else
71 - returned_token_ids,
72 - )
53 + return cls(task=task,
54 + pooling_type=PoolingType[pooler_config.pooling_type]
55 + if pooler_config.pooling_type is not None else pooling_type)
56
57
58 @dataclass(frozen=True)
@@ -89,22 +72,15 @@ class Pooler(nn.Module, ABC):
72 pooler_config: PoolerConfig,
73 *,
74 default_pooling_type: PoolingType = PoolingType.ALL,
92 - default_normalize: bool = False,
93 - default_softmax: bool = False,
94 - default_step_tag_id: Optional[int] = None,
95 - default_returned_token_ids: Optional[list[int]] = None,
75 ):
76 resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
77 + task="encode",
78 pooler_config=pooler_config,
79 pooling_type=default_pooling_type,
100 - normalize=default_normalize,
101 - softmax=default_softmax,
102 - step_tag_id=default_step_tag_id,
103 - returned_token_ids=default_returned_token_ids,
80 )
81
82 if resolved_config.pooling_type == PoolingType.STEP:
107 - return StepPooler.from_config(resolved_config)
83 + return StepPooler()
84
85 return SimplePooler.from_config(resolved_config)
86
@@ -113,14 +89,11 @@ class Pooler(nn.Module, ABC):
89 pooler_config: PoolerConfig,
90 *,
91 default_pooling_type: PoolingType = PoolingType.LAST,
116 - default_normalize: bool = True,
117 - default_softmax: bool = False,
92 ):
93 resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
94 + task="embed",
95 pooler_config=pooler_config,
96 pooling_type=default_pooling_type,
122 - normalize=default_normalize,
123 - softmax=default_softmax,
97 )
98
99 return SimplePooler.from_config(resolved_config)
@@ -131,23 +104,18 @@ class Pooler(nn.Module, ABC):
104 classifier: Optional[ClassifierFn],
105 *,
106 default_pooling_type: PoolingType = PoolingType.LAST,
134 - default_normalize: bool = False,
135 - default_softmax: bool = True,
107 ):
108 resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
109 + task="classify",
110 pooler_config=pooler_config,
111 pooling_type=default_pooling_type,
140 - normalize=default_normalize,
141 - softmax=default_softmax,
112 )
143 - base_pooler = SimplePooler.from_config(resolved_config)
144 - if classifier is None:
145 - return base_pooler
113 +
114 + pooling = PoolingMethod.from_pooling_type(resolved_config.pooling_type)
115
116 return ClassifierPooler(
148 - pooling=base_pooler.pooling,
117 + pooling=pooling,
118 classifier=classifier,
150 - act_fn=base_pooler.head.activation,
119 )
120
121 @abstractmethod
@@ -198,11 +166,17 @@ def get_prompt_token_ids(
166 ]
167
168
201 -def get_tasks(pooling_metadata: PoolingMetadata) -> list[PoolingTask]:
169 +def get_pooling_params(
170 + pooling_metadata: PoolingMetadata) -> list[PoolingParams]:
171 if isinstance(pooling_metadata, V0PoolingMetadata):
172 pooling_params = [p for _, p in pooling_metadata.seq_groups]
173 else:
174 pooling_params = pooling_metadata.pooling_params
175 + return pooling_params
176 +
177 +
178 +def get_tasks(pooling_metadata: PoolingMetadata) -> list[PoolingTask]:
179 + pooling_params = get_pooling_params(pooling_metadata)
180
181 tasks: list[PoolingTask] = [
182 task for pooling_param in pooling_params
@@ -484,49 +458,30 @@ class LambdaPoolerActivation(PoolerActivation):
458
459 class PoolerHead(nn.Module):
460
487 - @classmethod
488 - def from_config(cls, pooler_config: ResolvedPoolingConfig) -> "PoolerHead":
489 - if pooler_config.normalize and pooler_config.softmax:
490 - raise ValueError("`normalize=True` and `softmax=True` should not "
491 - "be set together")
492 -
493 - activation: PoolerActivation
494 - if pooler_config.normalize:
495 - activation = PoolerNormalize()
496 - elif pooler_config.softmax:
497 - activation = PoolerClassify()
498 - else:
499 - activation = PoolerIdentity()
500 -
501 - return cls(activation)
502 -
461 def __init__(self, activation: PoolerActivation) -> None:
462 super().__init__()
505 -
463 self.activation = activation
464
465 def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
466 pooling_metadata: PoolingMetadata):
467
511 - # Using float32 in PoolerHead
512 - if isinstance(pooled_data, list):
513 - for i in range(len(pooled_data)):
514 - pooled_data[i] = pooled_data[i].to(torch.float32)
515 - else:
516 - pooled_data = pooled_data.to(torch.float32)
468 + return self.activation(pooled_data)
469 +
470 +
471 +class EmbeddingPoolerHead(PoolerHead):
472 +
473 + def __init__(self) -> None:
474 + super().__init__(activation=PoolerNormalize())
475 +
476 + def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
477 + pooling_metadata: PoolingMetadata):
478 +
479 + pooling_params = get_pooling_params(pooling_metadata)
480
481 # for matryoshka representation
519 - if isinstance(pooling_metadata, V0PoolingMetadata):
520 - dimensions_list = [
521 - pooling_param.dimensions
522 - for _, pooling_param in pooling_metadata.seq_groups
523 - ]
524 - else:
525 - assert isinstance(pooled_data, list)
526 - dimensions_list = [
527 - pooling_param.dimensions
528 - for pooling_param in pooling_metadata.pooling_params
529 - ]
482 + dimensions_list = [
483 + pooling_param.dimensions for pooling_param in pooling_params
484 + ]
485 if any(d is not None for d in dimensions_list):
486 # change the output dimension
487 assert len(pooled_data) == len(dimensions_list)
@@ -541,7 +496,41 @@ class PoolerHead(nn.Module):
496 for vecs, d in zip(pooled_data, dimensions_list)
497 ]
498
544 - return self.activation(pooled_data)
499 + # for normalize
500 + flags = [p.normalize for p in pooling_params]
501 + if len(set(flags)) == 1:
502 + if flags[0]:
503 + pooled_data = self.activation(pooled_data)
504 + else:
505 + pooled_data = [
506 + self.activation(vecs) if f else vecs
507 + for vecs, f in zip(pooled_data, flags)
508 + ]
509 +
510 + return pooled_data
511 +
512 +
513 +class RewardPoolerHead(PoolerHead):
514 +
515 + def __init__(self) -> None:
516 + super().__init__(activation=PoolerClassify())
517 +
518 + def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
519 + pooling_metadata: PoolingMetadata):
520 + pooling_params = get_pooling_params(pooling_metadata)
521 +
522 + # for softmax
523 + flags = [p.softmax for p in pooling_params]
524 + if len(set(flags)) == 1:
525 + if flags[0]:
526 + pooled_data = self.activation(pooled_data)
527 + else:
528 + pooled_data = [
529 + self.activation(vecs) if f else vecs
530 + for vecs, f in zip(pooled_data, flags)
531 + ]
532 +
533 + return pooled_data
534
535
536 class SimplePooler(Pooler):
@@ -559,8 +548,12 @@ class SimplePooler(Pooler):
548 pooler_config: ResolvedPoolingConfig,
549 ) -> "SimplePooler":
550 pooling = PoolingMethod.from_pooling_type(pooler_config.pooling_type)
562 - head = PoolerHead.from_config(pooler_config)
563 -
551 + if pooler_config.task == "embed":
552 + head = EmbeddingPoolerHead()
553 + elif pooler_config.task == "encode":
554 + head = RewardPoolerHead()
555 + else:
556 + raise NotImplementedError(f"Unknown task: {pooler_config.task}")
557 return cls(pooling, head)
558
559 def __init__(self, pooling: PoolingMethod, head: PoolerHead) -> None:
@@ -587,29 +580,11 @@ class SimplePooler(Pooler):
580
581 class StepPooler(Pooler):
582
590 - @classmethod
591 - def from_config(cls, pooler_config: ResolvedPoolingConfig) -> "StepPooler":
592 - assert pooler_config.pooling_type == PoolingType.STEP
593 -
594 - return cls(
595 - PoolerHead.from_config(pooler_config),
596 - step_tag_id=pooler_config.step_tag_id,
597 - returned_token_ids=pooler_config.returned_token_ids,
598 - )
599 -
600 - def __init__(
601 - self,
602 - head: PoolerHead,
603 - *,
604 - step_tag_id: Optional[int] = None,
605 - returned_token_ids: Optional[list[int]] = None,
606 - ) -> None:
583 + def __init__(self, ) -> None:
584 super().__init__()
585
586 self.pooling = AllPool()
610 - self.head = head
611 - self.step_tag_id = step_tag_id
612 - self.returned_token_ids = returned_token_ids
587 + self.head = RewardPoolerHead()
588
589 def extract_states(
590 self,
@@ -620,10 +595,15 @@ class StepPooler(Pooler):
595 prompt_token_ids = get_prompt_token_ids(pooling_metadata)
596
597 pooled_data = list[torch.Tensor]()
623 - returned_token_ids = self.returned_token_ids
624 - step_tag_id = self.step_tag_id
598
626 - for data, token_id in zip(pooled_data_lst, prompt_token_ids):
599 + pooling_params = get_pooling_params(pooling_metadata)
600 +
601 + for data, token_id, pooling_param in zip(pooled_data_lst,
602 + prompt_token_ids,
603 + pooling_params):
604 + step_tag_id = pooling_param.step_tag_id
605 + returned_token_ids = pooling_param.returned_token_ids
606 +
607 if returned_token_ids is not None and len(returned_token_ids) > 0:
608 data = data[:, returned_token_ids]
609
@@ -669,14 +649,14 @@ class ClassifierPooler(Pooler):
649 def __init__(
650 self,
651 pooling: PoolingFn,
672 - classifier: ClassifierFn,
673 - act_fn: PoolerActivation,
652 + classifier: Optional[ClassifierFn],
653 + act_fn: Optional[PoolerActivation] = None,
654 ) -> None:
655 super().__init__()
656
657 self.pooling = pooling
658 self.classifier = classifier
679 - self.act_fn = act_fn
659 + self.act_fn = act_fn or PoolerClassify()
660
661 def get_supported_tasks(self) -> Set[PoolingTask]:
662 return {"classify", "score"}
@@ -688,15 +668,25 @@ class ClassifierPooler(Pooler):
668 ) -> PoolerOutput:
669 pooled_data = self.pooling(hidden_states, pooling_metadata)
670
691 - # apply classifier once on the full batch if possible
692 - if isinstance(pooled_data, torch.Tensor):
693 - pooled_output = self.classifier(pooled_data)
694 - elif len({data.shape for data in pooled_data}) <= 1:
695 - pooled_output = self.classifier(torch.stack(pooled_data))
696 - else:
697 - pooled_output = [self.classifier(data) for data in pooled_data]
671 + if self.classifier is not None:
672 + # apply classifier once on the full batch if possible
673 + if isinstance(pooled_data, torch.Tensor):
674 + pooled_data = self.classifier(pooled_data)
675 + elif len({data.shape for data in pooled_data}) <= 1:
676 + pooled_data = self.classifier(torch.stack(pooled_data))
677 + else:
678 + pooled_data = [self.classifier(data) for data in pooled_data]
679
699 - scores = self.act_fn(pooled_output)
680 + pooling_params = get_pooling_params(pooling_metadata)
681 + flags = [p.activation for p in pooling_params]
682 +
683 + if len(set(flags)) == 1:
684 + scores = self.act_fn(pooled_data) if flags[0] else pooled_data
685 + else:
686 + scores = [
687 + self.act_fn(vecs) if f else vecs
688 + for vecs, f in zip(pooled_data, flags)
689 + ]
690
691 return build_output(scores)
692
vllm/model_executor/models/config.py
+32
@@ -44,6 +44,15 @@ class GteNewModelConfig(VerifyAndUpdateConfig):
44 }
45
46
47 +class JambaForSequenceClassificationConfig(VerifyAndUpdateConfig):
48 +
49 + @staticmethod
50 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:
51 + pooler_config = vllm_config.model_config.pooler_config
52 + if pooler_config.activation is None:
53 + pooler_config.activation = False
54 +
55 +
56 class JinaRobertaModelConfig(VerifyAndUpdateConfig):
57
58 @staticmethod
@@ -155,6 +164,26 @@ class NomicBertModelConfig(VerifyAndUpdateConfig):
164 vllm_config.recalculate_max_model_len(max_model_len)
165
166
167 +class Qwen2ForProcessRewardModelConfig(VerifyAndUpdateConfig):
168 +
169 + @staticmethod
170 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:
171 + pooler_config = vllm_config.model_config.pooler_config
172 +
173 + if pooler_config.step_tag_id is None:
174 + pooler_config.step_tag_id = 151651
175 +
176 +
177 +class Qwen2ForRewardModelConfig(VerifyAndUpdateConfig):
178 +
179 + @staticmethod
180 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:
181 + pooler_config = vllm_config.model_config.pooler_config
182 +
183 + if pooler_config.softmax is None:
184 + pooler_config.softmax = False
185 +
186 +
187 class Qwen3ForSequenceClassificationConfig(VerifyAndUpdateConfig):
188
189 @staticmethod
@@ -309,8 +338,11 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
338 "GteModel": SnowflakeGteNewModelConfig,
339 "GteNewModel": GteNewModelConfig,
340 "NomicBertModel": NomicBertModelConfig,
341 + "Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig,
342 + "Qwen2ForRewardModel": Qwen2ForRewardModelConfig,
343 "Qwen3ForSequenceClassification": Qwen3ForSequenceClassificationConfig,
344 "XLMRobertaModel": JinaRobertaModelConfig,
345 "JinaVLForRanking": JinaVLForSequenceClassificationConfig,
346 + "JambaForSequenceClassification": JambaForSequenceClassificationConfig,
347 "GraniteMoeHybridForCausalLM": GraniteMoeHybridModelConfig,
348 }
vllm/model_executor/models/jamba.py
-2
@@ -593,7 +593,5 @@ class JambaForSequenceClassification(JambaForCausalLM):
593 pooler_config,
594 classifier=self.score,
595 default_pooling_type=PoolingType.LAST,
596 - default_normalize=False,
597 - default_softmax=False,
596 ),
597 })
vllm/model_executor/models/jina_vl.py
+1 -4
@@ -90,15 +90,12 @@ class JinaVLForSequenceClassification(Qwen2VLForConditionalGeneration,
90 prefix=maybe_prefix(prefix, "qwen2_vl"))
91 config = vllm_config.model_config.hf_config
92 pooler_config = vllm_config.model_config.pooler_config
93 + assert pooler_config is not None
94
95 # logit bias for sigmoid normalization
96 self.LOGIT_BIAS = 2.65
97
98 self.score = JinaVLScorer(config)
98 -
99 - pooler_config = vllm_config.model_config.pooler_config
100 - assert pooler_config is not None
101 -
99 self.pooler = DispatchPooler({
100 "encode":
101 Pooler.for_encode(pooler_config),
vllm/model_executor/models/qwen2_rm.py
-3
@@ -117,8 +117,5 @@ class Qwen2ForProcessRewardModel(Qwen2RewardBaseModel):
117 Pooler.for_encode(
118 pooler_config,
119 default_pooling_type=PoolingType.STEP,
120 - default_normalize=False,
121 - default_softmax=True,
122 - default_step_tag_id=151651,
120 )
121 })
vllm/pooling_params.py
+116 -23
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from copy import deepcopy
5 from typing import TYPE_CHECKING, Optional
6
7 import msgspec
@@ -19,13 +20,25 @@ class PoolingParams(
20 """API parameters for pooling models.
21
22 Attributes:
23 + normalize: Whether to normalize the embeddings outputs.
24 dimensions: Reduce the dimensions of embeddings
25 if model support matryoshka representation.
26 + activation: Whether to apply activation function to
27 + the classification outputs.
28 + softmax: Whether to apply softmax to the reward outputs.
29 """
30
31 + ## for embeddings models
32 dimensions: Optional[int] = None
33 + normalize: Optional[bool] = None
34
28 - output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY
35 + ## for classification models
36 + activation: Optional[bool] = None
37 +
38 + ## for reward models
39 + softmax: Optional[bool] = None
40 + step_tag_id: Optional[int] = None
41 + returned_token_ids: Optional[list[int]] = None
42
43 task: Optional[PoolingTask] = None
44 """Internal use only."""
@@ -33,15 +46,32 @@ class PoolingParams(
46 requires_token_ids: bool = False
47 """Internal use only."""
48
49 + output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY
50 +
51 + @property
52 + def all_parameters(self) -> list[str]:
53 + return [
54 + "dimensions", "normalize", "activation", "softmax", "step_tag_id",
55 + "returned_token_ids"
56 + ]
57 +
58 + @property
59 + def valid_parameters(self):
60 + return {
61 + "embed": ["dimensions", "normalize"],
62 + "classify": ["activation"],
63 + "score": ["activation"],
64 + "encode": ["softmax", "step_tag_id", "returned_token_ids"],
65 + }
66 +
67 def clone(self) -> "PoolingParams":
68 """Returns a deep copy of the PoolingParams instance."""
38 - return PoolingParams(
39 - dimensions=self.dimensions,
40 - task=self.task,
41 - requires_token_ids=self.requires_token_ids,
42 - )
69 + return deepcopy(self)
70 +
71 + def verify(self,
72 + task: PoolingTask,
73 + model_config: Optional["ModelConfig"] = None) -> None:
74
44 - def verify(self, task: PoolingTask, model_config: "ModelConfig") -> None:
75 if self.task is None:
76 self.task = task
77 elif self.task != task:
@@ -52,28 +82,91 @@ class PoolingParams(
82 # which is not available in model config. So, it's not included
83 # in this method
84
55 - if self.dimensions is not None:
56 - if not model_config.is_matryoshka:
57 - raise ValueError(
58 - f'Model "{model_config.served_model_name}" does not '
59 - f'support matryoshka representation, '
60 - f'changing output dimensions will lead to poor results.')
85 + self._merge_default_parameters(model_config)
86 + self._set_default_parameters(model_config)
87 + self._verify_valid_parameters()
88 +
89 + def _merge_default_parameters(self,
90 + model_config: Optional["ModelConfig"] = None
91 + ) -> None:
92 +
93 + if model_config is None:
94 + return
95
62 - mds = model_config.matryoshka_dimensions
63 - if mds is not None:
64 - if self.dimensions not in mds:
96 + pooler_config = model_config.pooler_config
97 + if pooler_config is None:
98 + return
99 +
100 + assert self.task is not None, "task must be set"
101 + valid_parameters = self.valid_parameters[self.task]
102 +
103 + for k in valid_parameters:
104 + if getattr(pooler_config, k, None) is None:
105 + continue
106 +
107 + if getattr(self, k, None) is None:
108 + setattr(self, k, getattr(pooler_config, k))
109 +
110 + def _set_default_parameters(self, model_config: Optional["ModelConfig"]):
111 + if self.task == "embed":
112 + if self.normalize is None:
113 + self.normalize = True
114 +
115 + if self.dimensions is not None and model_config is not None:
116 + if not model_config.is_matryoshka:
117 raise ValueError(
66 - f'Model "{model_config.served_model_name}" '
67 - f'only supports {str(mds)} matryoshka dimensions, '
68 - f'use other output dimensions will '
69 - f'lead to poor results.')
70 - elif self.dimensions < 1:
71 - raise ValueError("Dimensions must be greater than 0")
118 + f'Model "{model_config.served_model_name}" does not '
119 + f'support matryoshka representation, '
120 + f'changing output dimensions will lead to poor results.'
121 + )
122 +
123 + mds = model_config.matryoshka_dimensions
124 + if mds is not None:
125 + if self.dimensions not in mds:
126 + raise ValueError(
127 + f'Model "{model_config.served_model_name}" '
128 + f'only supports {str(mds)} matryoshka dimensions, '
129 + f'use other output dimensions will '
130 + f'lead to poor results.')
131 + elif self.dimensions < 1:
132 + raise ValueError("Dimensions must be greater than 0")
133 +
134 + elif self.task in ["classify", "score"]:
135 + if self.activation is None:
136 + self.activation = True
137 +
138 + elif self.task == "encode":
139 + if self.softmax is None:
140 + self.softmax = True
141 + else:
142 + raise ValueError(f"Unknown pooling task: {self.task}")
143 +
144 + def _verify_valid_parameters(self):
145 + assert self.task is not None, "task must be set"
146 + valid_parameters = self.valid_parameters[self.task]
147 + invalid_parameters = []
148 + for k in self.all_parameters:
149 + if k in valid_parameters:
150 + continue
151 +
152 + if getattr(self, k, None) is not None:
153 + invalid_parameters.append(k)
154 +
155 + if invalid_parameters:
156 + raise ValueError(
157 + f"Task {self.task} only supports {valid_parameters} "
158 + f"parameters, does not support "
159 + f"{invalid_parameters} parameters")
160
161 def __repr__(self) -> str:
162 return (f"PoolingParams("
75 - f"dimensions={self.dimensions}, "
163 f"task={self.task}, "
164 + f"normalize={self.normalize}, "
165 + f"dimensions={self.dimensions}, "
166 + f"activation={self.activation}, "
167 + f"softmax={self.softmax}, "
168 + f"step_tag_id={self.step_tag_id}, "
169 + f"returned_token_ids={self.returned_token_ids}, "
170 f"requires_token_ids={self.requires_token_ids})")
171
172 def __post_init__(self) -> None: