@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 84cf78acee
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 models default to using chunked prefill & prefix caching if supported. (#20930)

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

wang.yuqi committed Aug 12, 2025 at 00:41 UTC 84cf78acee1e75bfa163863b3674aeb3ba266844
31 files changed +452 -261
tests/entrypoints/llm/test_classify.py
+6
@@ -65,3 +65,9 @@ def test_pooling_params(llm: LLM):
65 assert torch.allclose(
66 softmax(wo_activation), w_activation, atol=1e-2
67 ), "w_activation should be close to activation(wo_activation)."
68 +
69 +
70 +def test_encode_api(llm: LLM):
71 + err_msg = "pooling_task must be one of.+"
72 + with pytest.raises(ValueError, match=err_msg):
73 + llm.encode(prompts, use_tqdm=False)
tests/entrypoints/openai/test_classification.py
+15
@@ -211,3 +211,18 @@ async def test_activation(server: RemoteOpenAIServer, model_name: str):
211 assert torch.allclose(
212 F.softmax(wo_activation, dim=-1), w_activation, atol=1e-2
213 ), "w_activation should be close to activation(wo_activation)."
214 +
215 +
216 +@pytest.mark.asyncio
217 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
218 +def test_pooling(server: RemoteOpenAIServer, model_name: str):
219 + # pooling api uses ALL pooling, which does not support chunked prefill.
220 + response = requests.post(
221 + server.url_for("pooling"),
222 + json={
223 + "model": model_name,
224 + "input": "test",
225 + "encoding_format": "float"
226 + },
227 + )
228 + assert response.json()["error"]["type"] == "BadRequestError"
tests/models/language/pooling/mteb_utils.py
+10 -2
@@ -177,9 +177,12 @@ def mteb_test_embed_models(hf_runner,
177 max_model_len=None,
178 **vllm_extra_kwargs) as vllm_model:
179
180 + model_config = vllm_model.llm.llm_engine.model_config
181 +
182 if model_info.architecture:
181 - assert (model_info.architecture
182 - in vllm_model.llm.llm_engine.model_config.architectures)
183 + assert model_info.architecture in model_config.architectures
184 + assert (model_config._model_info.default_pooling_type ==
185 + model_info.default_pooling_type)
186
187 vllm_main_score = run_mteb_embed_task(VllmMtebEncoder(vllm_model),
188 MTEB_EMBED_TASKS)
@@ -286,7 +289,12 @@ def mteb_test_rerank_models(hf_runner,
289 **vllm_extra_kwargs) as vllm_model:
290
291 model_config = vllm_model.llm.llm_engine.model_config
292 +
293 + if model_info.architecture:
294 + assert (model_info.architecture in model_config.architectures)
295 assert model_config.hf_config.num_labels == 1
296 + assert (model_config._model_info.default_pooling_type ==
297 + model_info.default_pooling_type)
298
299 vllm_main_score = run_mteb_rerank(vllm_mteb_encoder(vllm_model),
300 tasks=MTEB_RERANK_TASKS,
tests/models/language/pooling/test_auto_prefix_cache_support.py new
+93
@@ -0,0 +1,93 @@
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 AutoModelForSequenceClassification
6 +
7 +from tests.models.language.pooling.embed_utils import (
8 + run_embedding_correctness_test)
9 +
10 +
11 +@pytest.mark.parametrize(
12 + "model",
13 + ["jason9693/Qwen2.5-1.5B-apeach"],
14 +)
15 +@pytest.mark.parametrize("dtype", ["half"])
16 +def test_classify_models(
17 + hf_runner,
18 + vllm_runner,
19 + example_prompts,
20 + model: str,
21 + dtype: str,
22 +) -> None:
23 +
24 + example_prompts = example_prompts * 2
25 +
26 + with vllm_runner(model,
27 + max_model_len=512,
28 + dtype=dtype,
29 + enable_prefix_caching=True) as vllm_model:
30 + cache_config = vllm_model.llm.llm_engine.cache_config
31 + assert cache_config.enable_prefix_caching
32 + vllm_outputs = vllm_model.classify(example_prompts)
33 +
34 + with hf_runner(model,
35 + dtype=dtype,
36 + auto_cls=AutoModelForSequenceClassification) as hf_model:
37 + hf_outputs = hf_model.classify(example_prompts)
38 +
39 + for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
40 + hf_output = torch.tensor(hf_output)
41 + vllm_output = torch.tensor(vllm_output)
42 +
43 + assert torch.allclose(hf_output, vllm_output,
44 + 1e-3 if dtype == "float" else 1e-2)
45 +
46 +
47 +@pytest.mark.parametrize(
48 + "model",
49 + ["Qwen/Qwen3-Embedding-0.6B"],
50 +)
51 +@pytest.mark.parametrize("dtype", ["half"])
52 +def test_embed_models(
53 + hf_runner,
54 + vllm_runner,
55 + example_prompts,
56 + model: str,
57 + dtype: str,
58 +):
59 + example_prompts = [str(s).strip() for s in example_prompts] * 2
60 +
61 + with vllm_runner(
62 + model,
63 + runner="pooling",
64 + max_model_len=None,
65 + enable_prefix_caching=True,
66 + ) as vllm_model:
67 + cache_config = vllm_model.llm.llm_engine.cache_config
68 + assert cache_config.enable_prefix_caching
69 + vllm_outputs = vllm_model.embed(example_prompts)
70 +
71 + with hf_runner(
72 + model,
73 + is_sentence_transformer=True,
74 + ) as hf_model:
75 + run_embedding_correctness_test(hf_model, example_prompts, vllm_outputs)
76 +
77 +
78 +@pytest.mark.parametrize(
79 + "model",
80 + [
81 + "intfloat/e5-small",
82 + "Alibaba-NLP/gte-Qwen2-1.5B-instruct", # is_causal == False
83 + "papluca/xlm-roberta-base-language-detection",
84 + ])
85 +@pytest.mark.parametrize("dtype", ["half"])
86 +def test_non_causal_models(hf_runner, vllm_runner, example_prompts, model: str,
87 + dtype: str) -> None:
88 + with vllm_runner(model,
89 + max_model_len=512,
90 + dtype=dtype,
91 + enable_prefix_caching=True) as vllm_model:
92 + cache_config = vllm_model.llm.llm_engine.cache_config
93 + assert not cache_config.enable_prefix_caching
tests/models/language/pooling/test_baai.py
+61 -56
@@ -2,73 +2,78 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4
5 -from ...utils import EmbedModelInfo, RerankModelInfo
5 +from ...utils import (CLSPoolingEmbedModelInfo, CLSPoolingRerankModelInfo,
6 + EmbedModelInfo, LASTPoolingEmbedModelInfo,
7 + RerankModelInfo)
8 from .embed_utils import correctness_test_embed_models
9 from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
10
11 MODELS = [
12 ########## BertModel
11 - EmbedModelInfo("BAAI/bge-base-en",
12 - architecture="BertModel",
13 - enable_test=True),
14 - EmbedModelInfo("BAAI/bge-base-zh",
15 - architecture="BertModel",
16 - enable_test=False),
17 - EmbedModelInfo("BAAI/bge-small-en",
18 - architecture="BertModel",
19 - enable_test=False),
20 - EmbedModelInfo("BAAI/bge-small-zh",
21 - architecture="BertModel",
22 - enable_test=False),
23 - EmbedModelInfo("BAAI/bge-large-en",
24 - architecture="BertModel",
25 - enable_test=False),
26 - EmbedModelInfo("BAAI/bge-large-zh",
27 - architecture="BertModel",
28 - enable_test=False),
29 - EmbedModelInfo("BAAI/bge-large-zh-noinstruct",
30 - architecture="BertModel",
31 - enable_test=False),
32 - EmbedModelInfo("BAAI/bge-base-en-v1.5",
33 - architecture="BertModel",
34 - enable_test=False),
35 - EmbedModelInfo("BAAI/bge-base-zh-v1.5",
36 - architecture="BertModel",
37 - enable_test=False),
38 - EmbedModelInfo("BAAI/bge-small-en-v1.5",
39 - architecture="BertModel",
40 - enable_test=False),
41 - EmbedModelInfo("BAAI/bge-small-zh-v1.5",
42 - architecture="BertModel",
43 - enable_test=False),
44 - EmbedModelInfo("BAAI/bge-large-en-v1.5",
45 - architecture="BertModel",
46 - enable_test=False),
47 - EmbedModelInfo("BAAI/bge-large-zh-v1.5",
48 - architecture="BertModel",
49 - enable_test=False),
13 + CLSPoolingEmbedModelInfo("BAAI/bge-base-en",
14 + architecture="BertModel",
15 + enable_test=True),
16 + CLSPoolingEmbedModelInfo("BAAI/bge-base-zh",
17 + architecture="BertModel",
18 + enable_test=False),
19 + CLSPoolingEmbedModelInfo("BAAI/bge-small-en",
20 + architecture="BertModel",
21 + enable_test=False),
22 + CLSPoolingEmbedModelInfo("BAAI/bge-small-zh",
23 + architecture="BertModel",
24 + enable_test=False),
25 + CLSPoolingEmbedModelInfo("BAAI/bge-large-en",
26 + architecture="BertModel",
27 + enable_test=False),
28 + CLSPoolingEmbedModelInfo("BAAI/bge-large-zh",
29 + architecture="BertModel",
30 + enable_test=False),
31 + CLSPoolingEmbedModelInfo("BAAI/bge-large-zh-noinstruct",
32 + architecture="BertModel",
33 + enable_test=False),
34 + CLSPoolingEmbedModelInfo("BAAI/bge-base-en-v1.5",
35 + architecture="BertModel",
36 + enable_test=False),
37 + CLSPoolingEmbedModelInfo("BAAI/bge-base-zh-v1.5",
38 + architecture="BertModel",
39 + enable_test=False),
40 + CLSPoolingEmbedModelInfo("BAAI/bge-small-en-v1.5",
41 + architecture="BertModel",
42 + enable_test=False),
43 + CLSPoolingEmbedModelInfo("BAAI/bge-small-zh-v1.5",
44 + architecture="BertModel",
45 + enable_test=False),
46 + CLSPoolingEmbedModelInfo("BAAI/bge-large-en-v1.5",
47 + architecture="BertModel",
48 + enable_test=False),
49 + CLSPoolingEmbedModelInfo("BAAI/bge-large-zh-v1.5",
50 + architecture="BertModel",
51 + enable_test=False),
52 ########## XLMRobertaModel
51 - EmbedModelInfo("BAAI/bge-m3",
52 - architecture="XLMRobertaModel",
53 - enable_test=True),
53 + CLSPoolingEmbedModelInfo("BAAI/bge-m3",
54 + architecture="XLMRobertaModel",
55 + enable_test=True),
56 ########## Qwen2Model
55 - EmbedModelInfo("BAAI/bge-code-v1",
56 - architecture="Qwen2Model",
57 - dtype="float32",
58 - enable_test=True),
57 + LASTPoolingEmbedModelInfo("BAAI/bge-code-v1",
58 + architecture="Qwen2Model",
59 + dtype="float32",
60 + enable_test=True),
61 ]
62
63 RERANK_MODELS = [
64 ########## XLMRobertaForSequenceClassification
63 - RerankModelInfo("BAAI/bge-reranker-base",
64 - architecture="XLMRobertaForSequenceClassification",
65 - enable_test=True),
66 - RerankModelInfo("BAAI/bge-reranker-large",
67 - architecture="XLMRobertaForSequenceClassification",
68 - enable_test=False),
69 - RerankModelInfo("BAAI/bge-reranker-v2-m3",
70 - architecture="XLMRobertaForSequenceClassification",
71 - enable_test=False)
65 + CLSPoolingRerankModelInfo(
66 + "BAAI/bge-reranker-base",
67 + architecture="XLMRobertaForSequenceClassification",
68 + enable_test=True),
69 + CLSPoolingRerankModelInfo(
70 + "BAAI/bge-reranker-large",
71 + architecture="XLMRobertaForSequenceClassification",
72 + enable_test=False),
73 + CLSPoolingRerankModelInfo(
74 + "BAAI/bge-reranker-v2-m3",
75 + architecture="XLMRobertaForSequenceClassification",
76 + enable_test=False)
77 ]
78
79
tests/models/language/pooling/test_bge_reranker_v2_gemma.py
+4 -4
@@ -8,12 +8,12 @@ import torch
8
9 from tests.conftest import HfRunner
10
11 -from .mteb_utils import (RerankModelInfo, VllmMtebEncoder,
12 - mteb_test_rerank_models)
11 +from ...utils import LASTPoolingRerankModelInfo, RerankModelInfo
12 +from .mteb_utils import VllmMtebEncoder, mteb_test_rerank_models
13
14 RERANK_MODELS = [
15 - RerankModelInfo("BAAI/bge-reranker-v2-gemma",
16 - architecture="GemmaForSequenceClassification"),
15 + LASTPoolingRerankModelInfo("BAAI/bge-reranker-v2-gemma",
16 + architecture="GemmaForSequenceClassification"),
17 ]
18
19 PROMPT = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'." # noqa: E501
tests/models/language/pooling/test_cross_encoder.py
+7 -5
@@ -2,13 +2,15 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4
5 -from .mteb_utils import RerankModelInfo, mteb_test_rerank_models
5 +from ...utils import (CLSPoolingRerankModelInfo, LASTPoolingRerankModelInfo,
6 + RerankModelInfo)
7 +from .mteb_utils import mteb_test_rerank_models
8
9 RERANK_MODELS = [
8 - RerankModelInfo("cross-encoder/ms-marco-TinyBERT-L-2-v2",
9 - architecture="BertForSequenceClassification"),
10 - RerankModelInfo("tomaarsen/Qwen3-Reranker-0.6B-seq-cls",
11 - architecture="Qwen3ForSequenceClassification")
10 + CLSPoolingRerankModelInfo("cross-encoder/ms-marco-TinyBERT-L-2-v2",
11 + architecture="BertForSequenceClassification"),
12 + LASTPoolingRerankModelInfo("tomaarsen/Qwen3-Reranker-0.6B-seq-cls",
13 + architecture="Qwen3ForSequenceClassification")
14 ]
15
16
tests/models/language/pooling/test_gte.py
+44 -43
@@ -4,57 +4,58 @@ from typing import Any
4
5 import pytest
6
7 -from ...utils import check_transformers_version
8 -from .embed_utils import EmbedModelInfo, correctness_test_embed_models
7 +from ...utils import (CLSPoolingEmbedModelInfo, EmbedModelInfo,
8 + LASTPoolingEmbedModelInfo, check_transformers_version)
9 +from .embed_utils import correctness_test_embed_models
10 from .mteb_utils import mteb_test_embed_models
11
12 MODELS = [
13 ########## BertModel
13 - EmbedModelInfo("thenlper/gte-large",
14 - architecture="BertModel",
15 - enable_test=True),
16 - EmbedModelInfo("thenlper/gte-base",
17 - architecture="BertModel",
18 - enable_test=False),
19 - EmbedModelInfo("thenlper/gte-small",
20 - architecture="BertModel",
21 - enable_test=False),
22 - EmbedModelInfo("thenlper/gte-large-zh",
23 - architecture="BertModel",
24 - enable_test=False),
25 - EmbedModelInfo("thenlper/gte-base-zh",
26 - architecture="BertModel",
27 - enable_test=False),
28 - EmbedModelInfo("thenlper/gte-small-zh",
29 - architecture="BertModel",
30 - enable_test=False),
14 + CLSPoolingEmbedModelInfo("thenlper/gte-large",
15 + architecture="BertModel",
16 + enable_test=True),
17 + CLSPoolingEmbedModelInfo("thenlper/gte-base",
18 + architecture="BertModel",
19 + enable_test=False),
20 + CLSPoolingEmbedModelInfo("thenlper/gte-small",
21 + architecture="BertModel",
22 + enable_test=False),
23 + CLSPoolingEmbedModelInfo("thenlper/gte-large-zh",
24 + architecture="BertModel",
25 + enable_test=False),
26 + CLSPoolingEmbedModelInfo("thenlper/gte-base-zh",
27 + architecture="BertModel",
28 + enable_test=False),
29 + CLSPoolingEmbedModelInfo("thenlper/gte-small-zh",
30 + architecture="BertModel",
31 + enable_test=False),
32 ########### NewModel
32 - EmbedModelInfo("Alibaba-NLP/gte-multilingual-base",
33 - architecture="GteNewModel",
34 - enable_test=True),
35 - EmbedModelInfo("Alibaba-NLP/gte-base-en-v1.5",
36 - architecture="GteNewModel",
37 - enable_test=True),
38 - EmbedModelInfo("Alibaba-NLP/gte-large-en-v1.5",
39 - architecture="GteNewModel",
40 - enable_test=True),
33 + CLSPoolingEmbedModelInfo("Alibaba-NLP/gte-multilingual-base",
34 + architecture="GteNewModel",
35 + enable_test=True),
36 + CLSPoolingEmbedModelInfo("Alibaba-NLP/gte-base-en-v1.5",
37 + architecture="GteNewModel",
38 + enable_test=True),
39 + CLSPoolingEmbedModelInfo("Alibaba-NLP/gte-large-en-v1.5",
40 + architecture="GteNewModel",
41 + enable_test=True),
42 ########### Qwen2ForCausalLM
42 - EmbedModelInfo("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
43 - architecture="Qwen2ForCausalLM",
44 - enable_test=True),
43 + LASTPoolingEmbedModelInfo("Alibaba-NLP/gte-Qwen2-1.5B-instruct",
44 + architecture="Qwen2ForCausalLM",
45 + enable_test=True),
46 ########## ModernBertModel
46 - EmbedModelInfo("Alibaba-NLP/gte-modernbert-base",
47 - architecture="ModernBertModel",
48 - enable_test=True),
47 + CLSPoolingEmbedModelInfo("Alibaba-NLP/gte-modernbert-base",
48 + architecture="ModernBertModel",
49 + enable_test=True),
50 ########## Qwen3ForCausalLM
50 - EmbedModelInfo("Qwen/Qwen3-Embedding-0.6B",
51 - architecture="Qwen3ForCausalLM",
52 - dtype="float32",
53 - enable_test=True),
54 - EmbedModelInfo("Qwen/Qwen3-Embedding-4B",
55 - architecture="Qwen3ForCausalLM",
56 - dtype="float32",
57 - enable_test=False),
51 + LASTPoolingEmbedModelInfo("Qwen/Qwen3-Embedding-0.6B",
52 + architecture="Qwen3ForCausalLM",
53 + dtype="float32",
54 + enable_test=True),
55 + LASTPoolingEmbedModelInfo("Qwen/Qwen3-Embedding-4B",
56 + architecture="Qwen3ForCausalLM",
57 + dtype="float32",
58 + enable_test=False),
59 ]
60
61
tests/models/language/pooling/test_intfloat.py
+22 -22
@@ -2,34 +2,34 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4
5 -from ...utils import EmbedModelInfo
5 +from ...utils import CLSPoolingEmbedModelInfo, EmbedModelInfo
6 from .embed_utils import correctness_test_embed_models
7 from .mteb_utils import mteb_test_embed_models
8
9 MODELS = [
10 ########## BertModel
11 - EmbedModelInfo("intfloat/e5-small",
12 - architecture="BertModel",
13 - enable_test=True),
14 - EmbedModelInfo("intfloat/e5-base",
15 - architecture="BertModel",
16 - enable_test=False),
17 - EmbedModelInfo("intfloat/e5-large",
18 - architecture="BertModel",
19 - enable_test=False),
20 - EmbedModelInfo("intfloat/multilingual-e5-small",
21 - architecture="BertModel",
22 - enable_test=False),
11 + CLSPoolingEmbedModelInfo("intfloat/e5-small",
12 + architecture="BertModel",
13 + enable_test=True),
14 + CLSPoolingEmbedModelInfo("intfloat/e5-base",
15 + architecture="BertModel",
16 + enable_test=False),
17 + CLSPoolingEmbedModelInfo("intfloat/e5-large",
18 + architecture="BertModel",
19 + enable_test=False),
20 + CLSPoolingEmbedModelInfo("intfloat/multilingual-e5-small",
21 + architecture="BertModel",
22 + enable_test=False),
23 ########## XLMRobertaModel
24 - EmbedModelInfo("intfloat/multilingual-e5-base",
25 - architecture="XLMRobertaModel",
26 - enable_test=True),
27 - EmbedModelInfo("intfloat/multilingual-e5-large",
28 - architecture="XLMRobertaModel",
29 - enable_test=False),
30 - EmbedModelInfo("intfloat/multilingual-e5-large-instruct",
31 - architecture="XLMRobertaModel",
32 - enable_test=False),
24 + CLSPoolingEmbedModelInfo("intfloat/multilingual-e5-base",
25 + architecture="XLMRobertaModel",
26 + enable_test=True),
27 + CLSPoolingEmbedModelInfo("intfloat/multilingual-e5-large",
28 + architecture="XLMRobertaModel",
29 + enable_test=False),
30 + CLSPoolingEmbedModelInfo("intfloat/multilingual-e5-large-instruct",
31 + architecture="XLMRobertaModel",
32 + enable_test=False),
33 ]
34
35
tests/models/language/pooling/test_jina.py
+8 -6
@@ -6,20 +6,22 @@ import pytest
6
7 from vllm import PoolingParams
8
9 -from ...utils import EmbedModelInfo, RerankModelInfo
9 +from ...utils import (CLSPoolingEmbedModelInfo, CLSPoolingRerankModelInfo,
10 + EmbedModelInfo, RerankModelInfo)
11 from .embed_utils import (check_embeddings_close,
12 correctness_test_embed_models, matryoshka_fy)
13 from .mteb_utils import mteb_test_embed_models, mteb_test_rerank_models
14
15 EMBEDDING_MODELS = [
15 - EmbedModelInfo("jinaai/jina-embeddings-v3",
16 - architecture="XLMRobertaModel",
17 - is_matryoshka=True)
16 + CLSPoolingEmbedModelInfo("jinaai/jina-embeddings-v3",
17 + architecture="XLMRobertaModel",
18 + is_matryoshka=True)
19 ]
20
21 RERANK_MODELS = [
21 - RerankModelInfo("jinaai/jina-reranker-v2-base-multilingual",
22 - architecture="XLMRobertaForSequenceClassification")
22 + CLSPoolingRerankModelInfo(
23 + "jinaai/jina-reranker-v2-base-multilingual",
24 + architecture="XLMRobertaForSequenceClassification")
25 ]
26
27
tests/models/language/pooling/test_mxbai_rerank.py
+8 -7
@@ -7,15 +7,16 @@ import torch
7
8 from tests.conftest import HfRunner
9
10 -from .mteb_utils import RerankModelInfo, mteb_test_rerank_models
10 +from ...utils import LASTPoolingRerankModelInfo, RerankModelInfo
11 +from .mteb_utils import mteb_test_rerank_models
12
13 RERANK_MODELS = [
13 - RerankModelInfo("mixedbread-ai/mxbai-rerank-base-v2",
14 - architecture="Qwen2ForSequenceClassification",
15 - enable_test=True),
16 - RerankModelInfo("mixedbread-ai/mxbai-rerank-large-v2",
17 - architecture="Qwen2ForSequenceClassification",
18 - enable_test=False)
14 + LASTPoolingRerankModelInfo("mixedbread-ai/mxbai-rerank-base-v2",
15 + architecture="Qwen2ForSequenceClassification",
16 + enable_test=True),
17 + LASTPoolingRerankModelInfo("mixedbread-ai/mxbai-rerank-large-v2",
18 + architecture="Qwen2ForSequenceClassification",
19 + enable_test=False)
20 ]
21
22
tests/models/language/pooling/test_nomic.py
+14 -13
@@ -3,22 +3,23 @@
3
4 import pytest
5
6 -from .embed_utils import EmbedModelInfo, correctness_test_embed_models
6 +from ...utils import CLSPoolingEmbedModelInfo, EmbedModelInfo
7 +from .embed_utils import correctness_test_embed_models
8 from .mteb_utils import mteb_test_embed_models
9
10 MODELS = [
10 - EmbedModelInfo("nomic-ai/nomic-embed-text-v1",
11 - architecture="NomicBertModel",
12 - enable_test=True),
13 - EmbedModelInfo("nomic-ai/nomic-embed-text-v1.5",
14 - architecture="NomicBertModel",
15 - enable_test=False),
16 - EmbedModelInfo("nomic-ai/CodeRankEmbed",
17 - architecture="NomicBertModel",
18 - enable_test=False),
19 - EmbedModelInfo("nomic-ai/nomic-embed-text-v2-moe",
20 - architecture="NomicBertModel",
21 - enable_test=True)
11 + CLSPoolingEmbedModelInfo("nomic-ai/nomic-embed-text-v1",
12 + architecture="NomicBertModel",
13 + enable_test=True),
14 + CLSPoolingEmbedModelInfo("nomic-ai/nomic-embed-text-v1.5",
15 + architecture="NomicBertModel",
16 + enable_test=False),
17 + CLSPoolingEmbedModelInfo("nomic-ai/CodeRankEmbed",
18 + architecture="NomicBertModel",
19 + enable_test=False),
20 + CLSPoolingEmbedModelInfo("nomic-ai/nomic-embed-text-v2-moe",
21 + architecture="NomicBertModel",
22 + enable_test=True)
23 ]
24
25
tests/models/language/pooling/test_qwen3_reranker.py
+8 -7
@@ -8,15 +8,16 @@ import torch
8 from tests.conftest import HfRunner
9 from tests.utils import multi_gpu_test
10
11 -from .mteb_utils import RerankModelInfo, mteb_test_rerank_models
11 +from ...utils import LASTPoolingRerankModelInfo, RerankModelInfo
12 +from .mteb_utils import mteb_test_rerank_models
13
14 RERANK_MODELS = [
14 - RerankModelInfo("Qwen/Qwen3-Reranker-0.6B",
15 - architecture="Qwen3ForSequenceClassification",
16 - enable_test=True),
17 - RerankModelInfo("Qwen/Qwen3-Reranker-4B",
18 - architecture="Qwen3ForSequenceClassification",
19 - enable_test=False)
15 + LASTPoolingRerankModelInfo("Qwen/Qwen3-Reranker-0.6B",
16 + architecture="Qwen3ForSequenceClassification",
17 + enable_test=True),
18 + LASTPoolingRerankModelInfo("Qwen/Qwen3-Reranker-4B",
19 + architecture="Qwen3ForSequenceClassification",
20 + enable_test=False)
21 ]
22
23
tests/models/language/pooling/test_snowflake_arctic_embed.py
+34 -33
@@ -3,42 +3,43 @@
3
4 import pytest
5
6 -from .embed_utils import EmbedModelInfo, correctness_test_embed_models
6 +from ...utils import CLSPoolingEmbedModelInfo, EmbedModelInfo
7 +from .embed_utils import correctness_test_embed_models
8 from .mteb_utils import mteb_test_embed_models
9
10 MODELS = [
10 - EmbedModelInfo("Snowflake/snowflake-arctic-embed-xs",
11 - is_matryoshka=False,
12 - architecture="BertModel",
13 - enable_test=True),
14 - EmbedModelInfo("Snowflake/snowflake-arctic-embed-s",
15 - is_matryoshka=False,
16 - architecture="BertModel",
17 - enable_test=False),
18 - EmbedModelInfo("Snowflake/snowflake-arctic-embed-m",
19 - is_matryoshka=False,
20 - architecture="BertModel",
21 - enable_test=False),
22 - EmbedModelInfo("Snowflake/snowflake-arctic-embed-m-long",
23 - is_matryoshka=False,
24 - architecture="NomicBertModel",
25 - enable_test=True),
26 - EmbedModelInfo("Snowflake/snowflake-arctic-embed-l",
27 - is_matryoshka=False,
28 - architecture="BertModel",
29 - enable_test=False),
30 - EmbedModelInfo("Snowflake/snowflake-arctic-embed-m-v1.5",
31 - is_matryoshka=True,
32 - architecture="BertModel",
33 - enable_test=True),
34 - EmbedModelInfo("Snowflake/snowflake-arctic-embed-l-v2.0",
35 - is_matryoshka=True,
36 - architecture="XLMRobertaModel",
37 - enable_test=True),
38 - EmbedModelInfo("Snowflake/snowflake-arctic-embed-m-v2.0",
39 - is_matryoshka=True,
40 - architecture="GteModel",
41 - enable_test=True),
11 + CLSPoolingEmbedModelInfo("Snowflake/snowflake-arctic-embed-xs",
12 + is_matryoshka=False,
13 + architecture="BertModel",
14 + enable_test=True),
15 + CLSPoolingEmbedModelInfo("Snowflake/snowflake-arctic-embed-s",
16 + is_matryoshka=False,
17 + architecture="BertModel",
18 + enable_test=False),
19 + CLSPoolingEmbedModelInfo("Snowflake/snowflake-arctic-embed-m",
20 + is_matryoshka=False,
21 + architecture="BertModel",
22 + enable_test=False),
23 + CLSPoolingEmbedModelInfo("Snowflake/snowflake-arctic-embed-m-long",
24 + is_matryoshka=False,
25 + architecture="NomicBertModel",
26 + enable_test=True),
27 + CLSPoolingEmbedModelInfo("Snowflake/snowflake-arctic-embed-l",
28 + is_matryoshka=False,
29 + architecture="BertModel",
30 + enable_test=False),
31 + CLSPoolingEmbedModelInfo("Snowflake/snowflake-arctic-embed-m-v1.5",
32 + is_matryoshka=True,
33 + architecture="BertModel",
34 + enable_test=True),
35 + CLSPoolingEmbedModelInfo("Snowflake/snowflake-arctic-embed-l-v2.0",
36 + is_matryoshka=True,
37 + architecture="XLMRobertaModel",
38 + enable_test=True),
39 + CLSPoolingEmbedModelInfo("Snowflake/snowflake-arctic-embed-m-v2.0",
40 + is_matryoshka=True,
41 + architecture="GteModel",
42 + enable_test=True),
43 ]
44
45
tests/models/utils.py
+18
@@ -345,16 +345,34 @@ class EmbedModelInfo(NamedTuple):
345 matryoshka_dimensions: Optional[list[int]] = None
346 architecture: str = ""
347 dtype: str = "auto"
348 + default_pooling_type: str = ""
349 enable_test: bool = True
350
351
352 +class CLSPoolingEmbedModelInfo(EmbedModelInfo):
353 + default_pooling_type: str = "CLS"
354 +
355 +
356 +class LASTPoolingEmbedModelInfo(EmbedModelInfo):
357 + default_pooling_type: str = "LAST"
358 +
359 +
360 class RerankModelInfo(NamedTuple):
361 name: str
362 architecture: str = ""
363 dtype: str = "auto"
364 + default_pooling_type: str = ""
365 enable_test: bool = True
366
367
368 +class CLSPoolingRerankModelInfo(RerankModelInfo):
369 + default_pooling_type: str = "CLS"
370 +
371 +
372 +class LASTPoolingRerankModelInfo(RerankModelInfo):
373 + default_pooling_type: str = "LAST"
374 +
375 +
376 def dummy_hf_overrides(
377 hf_config: PretrainedConfig,
378 *,
tests/test_config.py
+14
@@ -227,6 +227,20 @@ def test_get_pooling_config_from_args():
227 assert asdict(pooling_config) == asdict(override_pooler_config)
228
229
230 +@pytest.mark.parametrize(
231 + ("model_id", "default_pooling_type", "pooling_type"),
232 + [
233 + ("tomaarsen/Qwen3-Reranker-0.6B-seq-cls", "LAST", "LAST"), # LLM
234 + ("intfloat/e5-small", "CLS", "MEAN"), # BertModel
235 + ("Qwen/Qwen2.5-Math-RM-72B", "ALL", "ALL"), # reward
236 + ("Qwen/Qwen2.5-Math-PRM-7B", "STEP", "STEP") # step reward
237 + ])
238 +def test_default_pooling_type(model_id, default_pooling_type, pooling_type):
239 + model_config = ModelConfig(model_id)
240 + assert model_config._model_info.default_pooling_type == default_pooling_type
241 + assert model_config.pooler_config.pooling_type == pooling_type
242 +
243 +
244 @pytest.mark.skipif(current_platform.is_rocm(),
245 reason="Xformers backend is not supported on ROCm.")
246 def test_get_bert_tokenization_sentence_transformer_config():
vllm/config/__init__.py
+8
@@ -871,6 +871,10 @@ class ModelConfig:
871 if getattr(pooler_config, k) is None:
872 setattr(pooler_config, k, v)
873
874 + default_pooling_type = self._model_info.default_pooling_type
875 + if pooler_config.pooling_type is None:
876 + pooler_config.pooling_type = default_pooling_type
877 +
878 return pooler_config
879
880 return None
@@ -3844,6 +3848,10 @@ class VllmConfig:
3848 disable_chunked_prefill_reasons.append(
3849 "Only \"last\" pooling supports chunked "
3850 "prefill and prefix caching; disabling both.")
3851 + elif not getattr(self.model_config.hf_config, "is_causal", True):
3852 + disable_chunked_prefill_reasons.append(
3853 + "Only models using causal attention supports chunked "
3854 + "prefill and prefix caching; disabling both.")
3855
3856 if disable_chunked_prefill_reasons:
3857 for reason in disable_chunked_prefill_reasons:
vllm/engine/arg_utils.py
+4 -5
@@ -1600,11 +1600,10 @@ class EngineArgs:
1600 else:
1601
1602 pooling_type = model_config.pooler_config.pooling_type
1603 -
1604 - # TODO: when encoder models are supported we'll have to
1605 - # check for causal attention here.
1606 - incremental_prefill_supported = (pooling_type is not None and
1607 - pooling_type.lower() == "last")
1603 + is_causal = getattr(model_config.hf_config, "is_causal", True)
1604 + incremental_prefill_supported = (pooling_type is not None
1605 + and pooling_type.lower() == "last"
1606 + and is_causal)
1607
1608 action = "Enabling" if \
1609 incremental_prefill_supported else "Disabling"
vllm/entrypoints/llm.py
+4
@@ -1100,6 +1100,10 @@ class LLM:
1100 "Try passing `--runner pooling` to use the model as a "
1101 "pooling model.")
1102
1103 + if pooling_task not in self.supported_tasks:
1104 + raise ValueError(
1105 + f"pooling_task must be one of {self.supported_tasks}.")
1106 +
1107 if prompt_token_ids is not None:
1108 parsed_prompts = self._convert_v1_inputs(
1109 prompts=cast(Optional[Union[str, list[str]]], prompts),
vllm/model_executor/layers/pooler.py
+11 -27
@@ -44,15 +44,14 @@ class ResolvedPoolingConfig:
44 task: PoolingTask
45
46 @classmethod
47 - def from_config_with_defaults(
47 + def from_config(
48 cls,
49 task: PoolingTask,
50 pooler_config: PoolerConfig,
51 - pooling_type: PoolingType,
51 ) -> "ResolvedPoolingConfig":
52 + assert pooler_config.pooling_type is not None
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)
54 + pooling_type=PoolingType[pooler_config.pooling_type])
55
56
57 @dataclass(frozen=True)
@@ -68,32 +67,20 @@ class Pooler(nn.Module, ABC):
67 """The interface required for all poolers used in pooling models in vLLM."""
68
69 @staticmethod
71 - def for_encode(
72 - pooler_config: PoolerConfig,
73 - *,
74 - default_pooling_type: PoolingType = PoolingType.ALL,
75 - ):
76 - resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
77 - task="encode",
78 - pooler_config=pooler_config,
79 - pooling_type=default_pooling_type,
80 - )
81 -
82 - if resolved_config.pooling_type == PoolingType.STEP:
70 + def for_encode(pooler_config: PoolerConfig):
71 + if pooler_config.pooling_type == "STEP":
72 return StepPooler()
73
74 + resolved_config = ResolvedPoolingConfig(task="encode",
75 + pooling_type=PoolingType.ALL)
76 +
77 return SimplePooler.from_config(resolved_config)
78
79 @staticmethod
88 - def for_embed(
89 - pooler_config: PoolerConfig,
90 - *,
91 - default_pooling_type: PoolingType = PoolingType.LAST,
92 - ):
93 - resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
80 + def for_embed(pooler_config: PoolerConfig):
81 + resolved_config = ResolvedPoolingConfig.from_config(
82 task="embed",
83 pooler_config=pooler_config,
96 - pooling_type=default_pooling_type,
84 )
85
86 return SimplePooler.from_config(resolved_config)
@@ -102,13 +89,10 @@ class Pooler(nn.Module, ABC):
89 def for_classify(
90 pooler_config: PoolerConfig,
91 classifier: Optional[ClassifierFn],
105 - *,
106 - default_pooling_type: PoolingType = PoolingType.LAST,
92 ):
108 - resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
93 + resolved_config = ResolvedPoolingConfig.from_config(
94 task="classify",
95 pooler_config=pooler_config,
111 - pooling_type=default_pooling_type,
96 )
97
98 pooling = PoolingMethod.from_pooling_type(resolved_config.pooling_type)
vllm/model_executor/models/adapters.py
+2 -2
@@ -182,8 +182,8 @@ def as_seq_cls_model(cls: _T) -> _T:
182 assert pooler_config is not None
183
184 pooling_type_str = pooler_config.pooling_type
185 - pooling_type = (PoolingType.LAST if pooling_type_str is None else
186 - PoolingType[pooling_type_str])
185 + assert pooling_type_str is not None
186 + pooling_type = PoolingType[pooling_type_str]
187
188 self.pooler = DispatchPooler({
189 "encode":
vllm/model_executor/models/bert.py
+8 -8
@@ -28,7 +28,8 @@ from vllm.model_executor.pooling_metadata import PoolingMetadata
28 from vllm.sequence import IntermediateTensors
29 from vllm.tasks import PoolingTask
30
31 -from .interfaces import SupportsCrossEncoding, SupportsQuant
31 +from .interfaces import (SupportsCrossEncoding, SupportsQuant,
32 + default_pooling_type)
33 from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
34
35
@@ -327,6 +328,7 @@ class BertOutput(nn.Module):
328
329
330 @support_torch_compile
331 +@default_pooling_type("CLS")
332 class BertModel(nn.Module, SupportsQuant):
333
334 is_pooling_model = True
@@ -401,6 +403,7 @@ class BertModel(nn.Module, SupportsQuant):
403 return loaded_params
404
405
406 +@default_pooling_type("ALL")
407 class BertPoolingModel(BertModel):
408
409 is_pooling_model = True
@@ -431,6 +434,7 @@ class BertPoolingModel(BertModel):
434 return loaded_params
435
436
437 +@default_pooling_type("CLS")
438 class BertEmbeddingModel(nn.Module, SupportsQuant):
439 """A model that uses Bert to provide embedding functionalities.
440
@@ -486,13 +490,8 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):
490
491 def _build_pooler(self, pooler_config: PoolerConfig) -> Pooler:
492 return DispatchPooler({
489 - "encode":
490 - Pooler.for_encode(pooler_config),
491 - "embed":
492 - Pooler.for_embed(
493 - pooler_config,
494 - default_pooling_type=PoolingType.CLS,
495 - ),
493 + "encode": Pooler.for_encode(pooler_config),
494 + "embed": Pooler.for_embed(pooler_config),
495 })
496
497
@@ -541,6 +540,7 @@ def _decode_token_type_ids(input_ids: torch.Tensor) -> torch.Tensor:
540 return token_type_ids
541
542
543 +@default_pooling_type("CLS")
544 class BertForSequenceClassification(nn.Module, SupportsCrossEncoding,
545 SupportsQuant):
546 """A model that uses Bert to provide embedding functionalities.
vllm/model_executor/models/bert_with_rope.py
+3 -1
@@ -27,7 +27,8 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
27 from vllm.model_executor.layers.vocab_parallel_embedding import (
28 VocabParallelEmbedding)
29 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
30 -from vllm.model_executor.models.interfaces import SupportsQuant
30 +from vllm.model_executor.models.interfaces import (SupportsQuant,
31 + default_pooling_type)
32 from vllm.model_executor.models.utils import WeightsMapper
33 from vllm.model_executor.utils import set_weight_attrs
34 from vllm.platforms import current_platform
@@ -401,6 +402,7 @@ class BertWithRopeEncoder(nn.Module):
402
403
404 @support_torch_compile
405 +@default_pooling_type("CLS")
406 class BertWithRope(nn.Module, SupportsQuant):
407 hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
408
vllm/model_executor/models/interfaces.py
+14
@@ -641,6 +641,20 @@ def supports_cross_encoding(
641 return is_pooling_model(model) and _supports_cross_encoding(model)
642
643
644 +def default_pooling_type(pooling_type: str) -> object:
645 + """Set default_pooling_type decorator. """
646 +
647 + def func(model: object):
648 + model.default_pooling_type = pooling_type
649 + return model
650 +
651 + return func
652 +
653 +
654 +def get_default_pooling_type(model: Union[type[object], object]) -> str:
655 + return getattr(model, "default_pooling_type", "LAST")
656 +
657 +
658 class SupportsQuant:
659 """The interface required for all models that support quantization."""
660
vllm/model_executor/models/internlm2.py
+2 -1
@@ -31,7 +31,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
31 from vllm.model_executor.sampling_metadata import SamplingMetadata
32 from vllm.sequence import IntermediateTensors
33
34 -from .interfaces import SupportsLoRA, SupportsPP
34 +from .interfaces import SupportsLoRA, SupportsPP, default_pooling_type
35 from .utils import (is_pp_missing_parameter,
36 make_empty_intermediate_tensors_factory, make_layers,
37 maybe_prefix)
@@ -401,6 +401,7 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
401 return loaded_params
402
403
404 +@default_pooling_type("ALL")
405 class InternLM2ForRewardModel(InternLM2ForCausalLM):
406
407 is_pooling_model = True
vllm/model_executor/models/jamba.py
+1 -3
@@ -22,8 +22,7 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
22 from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer
23 from vllm.model_executor.layers.mamba.mamba_utils import (
24 MambaStateShapeCalculator)
25 -from vllm.model_executor.layers.pooler import (DispatchPooler, Pooler,
26 - PoolingType)
25 +from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
26 from vllm.model_executor.layers.quantization import QuantizationConfig
27 from vllm.model_executor.layers.vocab_parallel_embedding import (
28 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
@@ -604,6 +603,5 @@ class JambaForSequenceClassification(JambaForCausalLM):
603 Pooler.for_classify(
604 pooler_config,
605 classifier=self.score,
607 - default_pooling_type=PoolingType.LAST,
606 ),
607 })
vllm/model_executor/models/modernbert.py
+4 -2
@@ -26,7 +26,8 @@ from vllm.model_executor.pooling_metadata import PoolingMetadata
26 from vllm.sequence import IntermediateTensors
27 from vllm.tasks import PoolingTask
28
29 -from .interfaces import SupportsCrossEncoding, SupportsV0Only
29 +from .interfaces import (SupportsCrossEncoding, SupportsV0Only,
30 + default_pooling_type)
31 from .utils import WeightsMapper, maybe_prefix
32
33
@@ -201,6 +202,7 @@ class ModernBertEncoderLayer(nn.Module):
202
203
204 @support_torch_compile
205 +@default_pooling_type("CLS")
206 class ModernBertModel(nn.Module):
207 hf_to_vllm_mapper = WeightsMapper(
208 orig_to_new_prefix={"layers.": "encoder_layer.layers."})
@@ -264,7 +266,6 @@ class ModernBertPooler(Pooler):
266 self.pooling = PoolingMethod.from_pooling_type(pooling_type)
267 self.dense = nn.Linear(config.hidden_size, config.hidden_size,
268 config.classifier_bias)
267 - self.pooling_type = config.classifier_pooling
269 self.act = nn.GELU()
270 self.norm = nn.LayerNorm(config.hidden_size,
271 eps=config.norm_eps,
@@ -294,6 +295,7 @@ class ModernBertPooler(Pooler):
295 return pooled_output
296
297
298 +@default_pooling_type("CLS")
299 class ModernBertForSequenceClassification(nn.Module, SupportsV0Only,
300 SupportsCrossEncoding):
301
vllm/model_executor/models/qwen2_rm.py
+6 -10
@@ -15,11 +15,10 @@ from torch import nn
15 from vllm.config import VllmConfig
16 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
17 RowParallelLinear)
18 -from vllm.model_executor.layers.pooler import (DispatchPooler, Pooler,
19 - PoolingType)
18 +from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
19 from vllm.sequence import IntermediateTensors
20
22 -from .interfaces import SupportsLoRA, SupportsPP
21 +from .interfaces import SupportsLoRA, SupportsPP, default_pooling_type
22 from .qwen2 import Qwen2Model
23 from .utils import AutoWeightsLoader, maybe_prefix
24
@@ -90,6 +89,7 @@ class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):
89 return loader.load_weights(weights)
90
91
92 +@default_pooling_type("ALL")
93 class Qwen2ForRewardModel(Qwen2RewardBaseModel):
94
95 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -103,6 +103,7 @@ class Qwen2ForRewardModel(Qwen2RewardBaseModel):
103 {"encode": Pooler.for_encode(pooler_config)}, )
104
105
106 +@default_pooling_type("STEP")
107 class Qwen2ForProcessRewardModel(Qwen2RewardBaseModel):
108
109 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -112,10 +113,5 @@ class Qwen2ForProcessRewardModel(Qwen2RewardBaseModel):
113 pooler_config = vllm_config.model_config.pooler_config
114 assert pooler_config is not None
115
115 - self.pooler = DispatchPooler({
116 - "encode":
117 - Pooler.for_encode(
118 - pooler_config,
119 - default_pooling_type=PoolingType.STEP,
120 - )
121 - })
116 + self.pooler = DispatchPooler(
117 + {"encode": Pooler.for_encode(pooler_config)})
vllm/model_executor/models/registry.py
+4 -2
@@ -25,8 +25,8 @@ from vllm.logger import init_logger
25 from vllm.transformers_utils.dynamic_module import (
26 try_get_class_from_dynamic_module)
27
28 -from .interfaces import (has_inner_state, has_noops, is_attention_free,
29 - is_hybrid, supports_cross_encoding,
28 +from .interfaces import (get_default_pooling_type, has_inner_state, has_noops,
29 + is_attention_free, is_hybrid, supports_cross_encoding,
30 supports_multimodal, supports_multimodal_raw_input,
31 supports_pp, supports_transcription, supports_v0_only)
32 from .interfaces_base import is_pooling_model, is_text_generation_model
@@ -305,6 +305,7 @@ class _ModelInfo:
305 architecture: str
306 is_text_generation_model: bool
307 is_pooling_model: bool
308 + default_pooling_type: str
309 supports_cross_encoding: bool
310 supports_multimodal: bool
311 supports_multimodal_raw_input: bool
@@ -323,6 +324,7 @@ class _ModelInfo:
324 architecture=model.__name__,
325 is_text_generation_model=is_text_generation_model(model),
326 is_pooling_model=is_pooling_model(model),
327 + default_pooling_type=get_default_pooling_type(model),
328 supports_cross_encoding=supports_cross_encoding(model),
329 supports_multimodal=supports_multimodal(model),
330 supports_multimodal_raw_input=supports_multimodal_raw_input(model),
vllm/model_executor/models/roberta.py
+3 -1
@@ -23,7 +23,7 @@ from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
23 from vllm.sequence import IntermediateTensors
24
25 from .bert_with_rope import BertWithRope, JinaRobertaModel
26 -from .interfaces import SupportsCrossEncoding
26 +from .interfaces import SupportsCrossEncoding, default_pooling_type
27
28
29 class RobertaEmbedding(nn.Module):
@@ -86,6 +86,7 @@ class RobertaClassificationHead(nn.Module):
86 return x
87
88
89 +@default_pooling_type("CLS")
90 class RobertaEmbeddingModel(BertEmbeddingModel):
91 """A model that uses Roberta to provide embedding functionalities.
92
@@ -149,6 +150,7 @@ class RobertaEmbeddingModel(BertEmbeddingModel):
150 return loader.load_weights(weights_list, mapper=mapper)
151
152
153 +@default_pooling_type("CLS")
154 class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):
155 """A model that uses Roberta to provide embedding functionalities.
156
vllm/v1/worker/gpu_model_runner.py
+12 -1
@@ -1272,7 +1272,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1272 if not is_pooling_model(model):
1273 return []
1274
1275 - return list(model.pooler.get_supported_tasks())
1275 + supported_tasks = list(model.pooler.get_supported_tasks())
1276 +
1277 + if (self.scheduler_config.chunked_prefill_enabled
1278 + and "encode" in supported_tasks):
1279 + supported_tasks.remove("encode")
1280 +
1281 + logger.info_once("Chunked prefill is not supported with "
1282 + "encode task which using ALL pooling. "
1283 + "Please turn off chunked prefill by "
1284 + "`--no-enable-chunked-prefill` before using it.")
1285 +
1286 + return supported_tasks
1287
1288 def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
1289 tasks = list[SupportedTask]()