@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 583a90e005
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 10, 09:00 UTC

[Refactor] Separate sequence and token pooling types (#32026)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Jan 10, 2026 at 12:53 UTC 583a90e0055f3fa1fc464b2dd62ea26a7d389ae1
42 files changed +324 -204
tests/model_executor/test_model_load_with_params.py
+4 -2
@@ -46,7 +46,8 @@ def test_model_loading_with_params(vllm_runner, monkeypatch):
46 assert model_config.encoder_config["do_lower_case"]
47
48 # asserts on the pooling config files
49 - assert model_config.pooler_config.pooling_type == "CLS"
49 + assert model_config.pooler_config.seq_pooling_type == "CLS"
50 + assert model_config.pooler_config.tok_pooling_type == "ALL"
51 assert model_config.pooler_config.normalize
52
53 # asserts on the tokenizer loaded
@@ -90,7 +91,8 @@ def test_roberta_model_loading_with_params(vllm_runner, monkeypatch):
91 assert not model_config.encoder_config["do_lower_case"]
92
93 # asserts on the pooling config files
93 - assert model_config.pooler_config.pooling_type == "MEAN"
94 + assert model_config.pooler_config.seq_pooling_type == "MEAN"
95 + assert model_config.pooler_config.tok_pooling_type == "ALL"
96 assert model_config.pooler_config.normalize
97
98 # asserts on the tokenizer loaded
tests/models/language/pooling/test_embedding.py
+1 -1
@@ -54,7 +54,7 @@ def test_models(
54 vllm_extra_kwargs = {}
55 if model == "ssmits/Qwen2-7B-Instruct-embed-base":
56 vllm_extra_kwargs["pooler_config"] = PoolerConfig(
57 - pooling_type="MEAN", normalize=False
57 + seq_pooling_type="MEAN", normalize=False
58 )
59
60 max_model_len: int | None = 512
tests/models/language/pooling/test_mm_classifier_conversion.py
+1 -1
@@ -88,7 +88,7 @@ def test_gemma_multimodal(
88 convert="classify",
89 load_format="auto",
90 hf_overrides=update_config,
91 - pooler_config=PoolerConfig(pooling_type="LAST"),
91 + pooler_config=PoolerConfig(seq_pooling_type="LAST"),
92 max_model_len=512,
93 enforce_eager=True,
94 tensor_parallel_size=1,
tests/models/language/pooling_mteb_test/mteb_embed_utils.py
+5 -2
@@ -162,8 +162,11 @@ def mteb_test_embed_models(
162 assert model_info.architecture in model_config.architectures
163
164 # Confirm whether the important configs in model_config are correct.
165 - if model_info.pooling_type is not None:
166 - assert model_config.pooler_config.pooling_type == model_info.pooling_type
165 + pooler_config = model_config.pooler_config
166 + if model_info.seq_pooling_type is not None:
167 + assert pooler_config.seq_pooling_type == model_info.seq_pooling_type
168 + if model_info.tok_pooling_type is not None:
169 + assert pooler_config.tok_pooling_type == model_info.tok_pooling_type
170 if model_info.attn_type is not None:
171 assert model_config.attn_type == model_info.attn_type
172 if model_info.is_prefix_caching_supported is not None:
tests/models/language/pooling_mteb_test/mteb_score_utils.py
+5 -2
@@ -254,8 +254,11 @@ def mteb_test_rerank_models(
254 assert model_config.hf_config.num_labels == 1
255
256 # Confirm whether the important configs in model_config are correct.
257 - if model_info.pooling_type is not None:
258 - assert model_config.pooler_config.pooling_type == model_info.pooling_type
257 + pooler_config = model_config.pooler_config
258 + if model_info.seq_pooling_type is not None:
259 + assert pooler_config.seq_pooling_type == model_info.seq_pooling_type
260 + if model_info.tok_pooling_type is not None:
261 + assert pooler_config.tok_pooling_type == model_info.tok_pooling_type
262 if model_info.attn_type is not None:
263 assert model_config.attn_type == model_info.attn_type
264 if model_info.is_prefix_caching_supported is not None:
tests/models/language/pooling_mteb_test/test_baai.py
+4 -4
@@ -17,7 +17,7 @@ MODELS = [
17 "BAAI/bge-base-en",
18 architecture="BertModel",
19 mteb_score=0.779336792,
20 - pooling_type="CLS",
20 + seq_pooling_type="CLS",
21 attn_type="encoder_only",
22 is_prefix_caching_supported=False,
23 is_chunked_prefill_supported=False,
@@ -54,7 +54,7 @@ MODELS = [
54 "BAAI/bge-m3",
55 architecture="XLMRobertaModel",
56 mteb_score=0.787343078,
57 - pooling_type="CLS",
57 + seq_pooling_type="CLS",
58 attn_type="encoder_only",
59 is_prefix_caching_supported=False,
60 is_chunked_prefill_supported=False,
@@ -65,7 +65,7 @@ MODELS = [
65 "BAAI/bge-code-v1",
66 architecture="Qwen2Model",
67 mteb_score=0.75724465,
68 - pooling_type="LAST",
68 + seq_pooling_type="LAST",
69 attn_type="decoder",
70 is_prefix_caching_supported=True,
71 is_chunked_prefill_supported=True,
@@ -79,7 +79,7 @@ RERANK_MODELS = [
79 "BAAI/bge-reranker-base",
80 architecture="XLMRobertaForSequenceClassification",
81 mteb_score=0.32398,
82 - pooling_type="CLS",
82 + seq_pooling_type="CLS",
83 attn_type="encoder_only",
84 is_prefix_caching_supported=False,
85 is_chunked_prefill_supported=False,
tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py
+1 -1
@@ -26,7 +26,7 @@ RERANK_MODELS = [
26 "method": "no_post_processing",
27 },
28 mteb_score=0.33757,
29 - pooling_type="LAST",
29 + seq_pooling_type="LAST",
30 attn_type="decoder",
31 is_prefix_caching_supported=True,
32 is_chunked_prefill_supported=True,
tests/models/language/pooling_mteb_test/test_cross_encoder.py
+2 -2
@@ -12,7 +12,7 @@ RERANK_MODELS = [
12 RerankModelInfo(
13 "cross-encoder/ms-marco-TinyBERT-L-2-v2",
14 architecture="BertForSequenceClassification",
15 - pooling_type="CLS",
15 + seq_pooling_type="CLS",
16 attn_type="encoder_only",
17 is_prefix_caching_supported=False,
18 is_chunked_prefill_supported=False,
@@ -21,7 +21,7 @@ RERANK_MODELS = [
21 RerankModelInfo(
22 "tomaarsen/Qwen3-Reranker-0.6B-seq-cls",
23 architecture="Qwen3ForSequenceClassification",
24 - pooling_type="LAST",
24 + seq_pooling_type="LAST",
25 attn_type="decoder",
26 is_prefix_caching_supported=True,
27 is_chunked_prefill_supported=True,
tests/models/language/pooling_mteb_test/test_gte.py
+7 -7
@@ -18,7 +18,7 @@ MODELS = [
18 "thenlper/gte-large",
19 mteb_score=0.76807651,
20 architecture="BertModel",
21 - pooling_type="MEAN",
21 + seq_pooling_type="MEAN",
22 attn_type="encoder_only",
23 is_prefix_caching_supported=False,
24 is_chunked_prefill_supported=False,
@@ -44,7 +44,7 @@ MODELS = [
44 architecture="GteNewModel",
45 mteb_score=0.775074696,
46 hf_overrides={"architectures": ["GteNewModel"]},
47 - pooling_type="CLS",
47 + seq_pooling_type="CLS",
48 attn_type="encoder_only",
49 is_prefix_caching_supported=False,
50 is_chunked_prefill_supported=False,
@@ -67,7 +67,7 @@ MODELS = [
67 "Alibaba-NLP/gte-Qwen2-1.5B-instruct",
68 mteb_score=0.758473459018872,
69 architecture="Qwen2ForCausalLM",
70 - pooling_type="LAST",
70 + seq_pooling_type="LAST",
71 attn_type="encoder_only",
72 is_prefix_caching_supported=False,
73 is_chunked_prefill_supported=False,
@@ -78,7 +78,7 @@ MODELS = [
78 "Alibaba-NLP/gte-modernbert-base",
79 mteb_score=0.748193353,
80 architecture="ModernBertModel",
81 - pooling_type="CLS",
81 + seq_pooling_type="CLS",
82 attn_type="encoder_only",
83 is_prefix_caching_supported=False,
84 is_chunked_prefill_supported=False,
@@ -89,7 +89,7 @@ MODELS = [
89 "Qwen/Qwen3-Embedding-0.6B",
90 mteb_score=0.771163695,
91 architecture="Qwen3ForCausalLM",
92 - pooling_type="LAST",
92 + seq_pooling_type="LAST",
93 attn_type="decoder",
94 is_prefix_caching_supported=True,
95 is_chunked_prefill_supported=True,
@@ -108,7 +108,7 @@ RERANK_MODELS = [
108 "Alibaba-NLP/gte-reranker-modernbert-base",
109 mteb_score=0.33386,
110 architecture="ModernBertForSequenceClassification",
111 - pooling_type="CLS",
111 + seq_pooling_type="CLS",
112 attn_type="encoder_only",
113 is_prefix_caching_supported=False,
114 is_chunked_prefill_supported=False,
@@ -119,7 +119,7 @@ RERANK_MODELS = [
119 mteb_score=0.33062,
120 architecture="GteNewForSequenceClassification",
121 hf_overrides={"architectures": ["GteNewForSequenceClassification"]},
122 - pooling_type="CLS",
122 + seq_pooling_type="CLS",
123 attn_type="encoder_only",
124 is_prefix_caching_supported=False,
125 is_chunked_prefill_supported=False,
tests/models/language/pooling_mteb_test/test_intfloat.py
+2 -2
@@ -13,7 +13,7 @@ MODELS = [
13 "intfloat/e5-small",
14 architecture="BertModel",
15 mteb_score=0.742285423,
16 - pooling_type="MEAN",
16 + seq_pooling_type="MEAN",
17 attn_type="encoder_only",
18 is_prefix_caching_supported=False,
19 is_chunked_prefill_supported=False,
@@ -29,7 +29,7 @@ MODELS = [
29 "intfloat/multilingual-e5-base",
30 architecture="XLMRobertaModel",
31 mteb_score=0.779325955,
32 - pooling_type="MEAN",
32 + seq_pooling_type="MEAN",
33 attn_type="encoder_only",
34 is_prefix_caching_supported=False,
35 is_chunked_prefill_supported=False,
tests/models/language/pooling_mteb_test/test_jina.py
+2 -2
@@ -24,7 +24,7 @@ EMBEDDING_MODELS = [
24 mteb_score=0.824413164,
25 architecture="XLMRobertaModel",
26 is_matryoshka=True,
27 - pooling_type="MEAN",
27 + seq_pooling_type="MEAN",
28 attn_type="encoder_only",
29 is_prefix_caching_supported=False,
30 is_chunked_prefill_supported=False,
@@ -36,7 +36,7 @@ RERANK_MODELS = [
36 "jinaai/jina-reranker-v2-base-multilingual",
37 mteb_score=0.33643,
38 architecture="XLMRobertaForSequenceClassification",
39 - pooling_type="CLS",
39 + seq_pooling_type="CLS",
40 attn_type="encoder_only",
41 is_prefix_caching_supported=False,
42 is_chunked_prefill_supported=False,
tests/models/language/pooling_mteb_test/test_mxbai_rerank.py
+1 -1
@@ -24,7 +24,7 @@ RERANK_MODELS = [
24 "mixedbread-ai/mxbai-rerank-base-v2",
25 architecture="Qwen2ForSequenceClassification",
26 hf_overrides=mxbai_rerank_hf_overrides,
27 - pooling_type="LAST",
27 + seq_pooling_type="LAST",
28 attn_type="decoder",
29 is_prefix_caching_supported=True,
30 is_chunked_prefill_supported=True,
tests/models/language/pooling_mteb_test/test_nemotron.py
+2 -2
@@ -19,7 +19,7 @@ EMBEDDING_MODELS = [
19 "nvidia/llama-nemotron-embed-1b-v2",
20 architecture="LlamaBidirectionalModel",
21 mteb_score=0.689164662128673,
22 - pooling_type="MEAN",
22 + seq_pooling_type="MEAN",
23 attn_type="encoder_only",
24 is_prefix_caching_supported=False,
25 is_chunked_prefill_supported=False,
@@ -32,7 +32,7 @@ RERANK_MODELS = [
32 architecture="LlamaBidirectionalForSequenceClassification",
33 chat_template_name="nemotron-rerank.jinja",
34 mteb_score=0.33994,
35 - pooling_type="MEAN",
35 + seq_pooling_type="MEAN",
36 attn_type="encoder_only",
37 is_prefix_caching_supported=False,
38 is_chunked_prefill_supported=False,
tests/models/language/pooling_mteb_test/test_nomic.py
+2 -2
@@ -14,7 +14,7 @@ MODELS = [
14 architecture="NomicBertModel",
15 mteb_score=0.737568559,
16 enable_test=True,
17 - pooling_type="MEAN",
17 + seq_pooling_type="MEAN",
18 attn_type="encoder_only",
19 is_prefix_caching_supported=False,
20 is_chunked_prefill_supported=False,
@@ -32,7 +32,7 @@ MODELS = [
32 architecture="NomicBertModel",
33 mteb_score=0.715488912,
34 enable_test=True,
35 - pooling_type="MEAN",
35 + seq_pooling_type="MEAN",
36 attn_type="encoder_only",
37 is_prefix_caching_supported=False,
38 is_chunked_prefill_supported=False,
tests/models/language/pooling_mteb_test/test_qwen3_reranker.py
+1 -1
@@ -27,7 +27,7 @@ RERANK_MODELS = [
27 architecture="Qwen3ForSequenceClassification",
28 hf_overrides=qwen3_reranker_hf_overrides,
29 chat_template_name="qwen3_reranker.jinja",
30 - pooling_type="LAST",
30 + seq_pooling_type="LAST",
31 attn_type="decoder",
32 is_prefix_caching_supported=True,
33 is_chunked_prefill_supported=True,
tests/models/language/pooling_mteb_test/test_snowflake_arctic_embed.py
+5 -5
@@ -14,7 +14,7 @@ MODELS = [
14 is_matryoshka=False,
15 architecture="BertModel",
16 mteb_score=0.714927797,
17 - pooling_type="CLS",
17 + seq_pooling_type="CLS",
18 attn_type="encoder_only",
19 is_prefix_caching_supported=False,
20 is_chunked_prefill_supported=False,
@@ -37,7 +37,7 @@ MODELS = [
37 is_matryoshka=False,
38 architecture="NomicBertModel",
39 mteb_score=0.681146831,
40 - pooling_type="CLS",
40 + seq_pooling_type="CLS",
41 attn_type="encoder_only",
42 is_prefix_caching_supported=False,
43 is_chunked_prefill_supported=False,
@@ -54,7 +54,7 @@ MODELS = [
54 is_matryoshka=True,
55 architecture="BertModel",
56 mteb_score=0.649088363,
57 - pooling_type="CLS",
57 + seq_pooling_type="CLS",
58 attn_type="encoder_only",
59 is_prefix_caching_supported=False,
60 is_chunked_prefill_supported=False,
@@ -65,7 +65,7 @@ MODELS = [
65 is_matryoshka=True,
66 architecture="XLMRobertaModel",
67 mteb_score=0.712258299,
68 - pooling_type="CLS",
68 + seq_pooling_type="CLS",
69 attn_type="encoder_only",
70 is_prefix_caching_supported=False,
71 is_chunked_prefill_supported=False,
@@ -76,7 +76,7 @@ MODELS = [
76 is_matryoshka=True,
77 architecture="GteModel",
78 mteb_score=0.706622444,
79 - pooling_type="CLS",
79 + seq_pooling_type="CLS",
80 attn_type="encoder_only",
81 is_prefix_caching_supported=False,
82 is_chunked_prefill_supported=False,
tests/models/language/pooling_mteb_test/test_st_projector.py
+2 -2
@@ -14,7 +14,7 @@ ST_PROJECTOR_MODELS = [
14 "TencentBAC/Conan-embedding-v1",
15 architecture="BertModel",
16 mteb_score=0.688611955,
17 - pooling_type="MEAN",
17 + seq_pooling_type="MEAN",
18 attn_type="encoder_only",
19 is_prefix_caching_supported=False,
20 is_chunked_prefill_supported=False,
@@ -24,7 +24,7 @@ ST_PROJECTOR_MODELS = [
24 "google/embeddinggemma-300m",
25 architecture="Gemma3TextModel",
26 mteb_score=0.7473819294684156,
27 - pooling_type="MEAN",
27 + seq_pooling_type="MEAN",
28 attn_type="encoder_only",
29 is_prefix_caching_supported=False,
30 is_chunked_prefill_supported=False,
tests/models/utils.py
+3 -1
@@ -11,6 +11,7 @@ import torch.nn.functional as F
11 from transformers import PretrainedConfig
12
13 from vllm.config.model import AttnTypeStr, ModelConfig, ModelDType, RunnerOption
14 +from vllm.config.pooler import SequencePoolingType, TokenPoolingType
15 from vllm.logprobs import Logprob, PromptLogprobs, SampleLogprobs
16 from vllm.multimodal.processing import InputProcessingContext
17 from vllm.tokenizers import cached_tokenizer_from_config
@@ -379,7 +380,8 @@ class ModelInfo:
380 max_model_len: int | None = None
381 hf_dtype: str = "float32"
382 hf_overrides: dict[str, Any] | None = None
382 - pooling_type: str | None = None
383 + seq_pooling_type: SequencePoolingType | None = None
384 + tok_pooling_type: TokenPoolingType | None = None
385 attn_type: AttnTypeStr | None = None
386 is_prefix_caching_supported: bool | None = None
387 is_chunked_prefill_supported: bool | None = None
tests/test_config.py
+49 -37
@@ -161,7 +161,8 @@ def test_get_pooling_config():
161
162 assert model_config.pooler_config is not None
163 assert model_config.pooler_config.normalize
164 - assert model_config.pooler_config.pooling_type == "MEAN"
164 + assert model_config.pooler_config.seq_pooling_type == "MEAN"
165 + assert model_config.pooler_config.tok_pooling_type == "ALL"
166
167
168 @pytest.mark.skipif(
@@ -169,7 +170,7 @@ def test_get_pooling_config():
170 )
171 def test_get_pooling_config_from_args():
172 model_id = "sentence-transformers/all-MiniLM-L12-v2"
172 - pooler_config = PoolerConfig(pooling_type="CLS", normalize=True)
173 + pooler_config = PoolerConfig(seq_pooling_type="CLS", normalize=True)
174 model_config = ModelConfig(model_id, pooler_config=pooler_config)
175
176 assert asdict(model_config.pooler_config) == asdict(pooler_config)
@@ -180,14 +181,25 @@ def test_get_pooling_config_from_args():
181 [
182 ("tomaarsen/Qwen3-Reranker-0.6B-seq-cls", "LAST", "LAST"), # LLM
183 ("intfloat/e5-small", "CLS", "MEAN"), # BertModel
184 + ],
185 +)
186 +def test_default_seq_pooling_type(model_id, default_pooling_type, pooling_type):
187 + model_config = ModelConfig(model_id)
188 + assert model_config._model_info.default_seq_pooling_type == default_pooling_type
189 + assert model_config.pooler_config.seq_pooling_type == pooling_type
190 +
191 +
192 +@pytest.mark.parametrize(
193 + ("model_id", "default_pooling_type", "pooling_type"),
194 + [
195 ("Qwen/Qwen2.5-Math-RM-72B", "ALL", "ALL"), # reward
196 ("Qwen/Qwen2.5-Math-PRM-7B", "STEP", "STEP"), # step reward
197 ],
198 )
187 -def test_default_pooling_type(model_id, default_pooling_type, pooling_type):
199 +def test_default_tok_pooling_type(model_id, default_pooling_type, pooling_type):
200 model_config = ModelConfig(model_id)
189 - assert model_config._model_info.default_pooling_type == default_pooling_type
190 - assert model_config.pooler_config.pooling_type == pooling_type
201 + assert model_config._model_info.default_tok_pooling_type == default_pooling_type
202 + assert model_config.pooler_config.tok_pooling_type == pooling_type
203
204
205 @pytest.mark.parametrize(
@@ -554,100 +566,100 @@ def test_s3_url_different_models_create_different_directories(mock_pull_files):
566 "jason9693/Qwen2.5-1.5B-apeach",
567 "decoder",
568 True,
557 - "Pooling models with causal attn and last pooling support chunked prefill.",
569 + "Pooling models with causal attn and LAST/ALL pooling support chunked prefill.", # noqa: E501
570 ),
571 (
572 "Qwen/Qwen3-Embedding-0.6B",
573 "decoder",
574 True,
563 - "Pooling models with causal attn and last pooling support chunked prefill.",
575 + "Pooling models with causal attn and LAST/ALL pooling support chunked prefill.", # noqa: E501
576 ),
577 (
578 "Qwen/Qwen2.5-Math-PRM-7B",
579 "decoder",
580 False,
569 - "Pooling models with step pooling does not support chunked prefill.",
581 + "Pooling models with causal attn and LAST/STEP pooling do not support chunked prefill.", # noqa: E501
582 ),
583 (
584 "internlm/internlm2-1_8b-reward",
585 "decoder",
586 True,
575 - "Pooling models with causal attn and all pooling support chunked prefill.",
587 + "Pooling models with causal attn and LAST/ALL pooling support chunked prefill.", # noqa: E501
588 ),
589 (
590 "BAAI/bge-base-en",
591 "encoder_only",
592 False,
581 - "Pooling models with bidirectional attn does not support chunked prefill.",
593 + "Pooling models with bidirectional attn do not support chunked prefill.", # noqa: E501
594 ),
595 (
596 "boltuix/NeuroBERT-NER",
597 "encoder_only",
598 False,
587 - "Pooling models with bidirectional attn does not support chunked prefill.",
599 + "Pooling models with bidirectional attn do not support chunked prefill.", # noqa: E501
600 ),
601 (
602 "papluca/xlm-roberta-base-language-detection",
603 "encoder_only",
604 False,
593 - "Pooling models with bidirectional attn does not support chunked prefill.",
605 + "Pooling models with bidirectional attn do not support chunked prefill.", # noqa: E501
606 ),
607 (
608 "Alibaba-NLP/gte-Qwen2-1.5B-instruct",
609 "encoder_only",
610 False,
599 - "Pooling models with bidirectional attn does not support chunked prefill.",
611 + "Pooling models with bidirectional attn do not support chunked prefill.", # noqa: E501
612 ),
613 (
614 "intfloat/e5-small",
615 "encoder_only",
616 False,
605 - "Pooling models with bidirectional attn does not support chunked prefill.",
617 + "Pooling models with bidirectional attn do not support chunked prefill.", # noqa: E501
618 ),
619 # multimodal models
620 (
621 "openai/clip-vit-base-patch32",
622 "decoder",
623 True,
612 - "Pooling models with causal attn and last pooling support chunked prefill.",
624 + "Pooling models with causal attn and LAST/ALL pooling support chunked prefill.", # noqa: E501
625 ),
626 (
627 "google/siglip-base-patch16-224",
628 "encoder_only",
629 False,
618 - "Pooling models with bidirectional attn does not support chunked prefill.",
630 + "Pooling models with bidirectional attn do not support chunked prefill.", # noqa: E501
631 ),
632 # generate models
633 (
634 "Qwen/Qwen3-0.6B",
635 "decoder",
636 True,
625 - "Generative models support chunked prefill.",
637 + "Generative models support chunked prefill.", # noqa: E501
638 ),
639 (
640 "Qwen/Qwen3-Next-80B-A3B-Instruct",
641 "hybrid",
642 True,
631 - "Generative models support chunked prefill.",
643 + "Generative models support chunked prefill.", # noqa: E501
644 ),
645 (
646 "ibm-granite/granite-4.0-h-small",
647 "hybrid",
648 True,
637 - "Generative models support chunked prefill.",
649 + "Generative models support chunked prefill.", # noqa: E501
650 ),
651 (
652 "state-spaces/mamba-130m-hf",
653 "attention_free",
654 True,
643 - "Generative models support chunked prefill.",
655 + "Generative models support chunked prefill.", # noqa: E501
656 ),
657 # encoder_decoder models
658 (
659 "openai/whisper-small",
660 "encoder_decoder",
661 False,
650 - "Encoder decoder models does not support chunked prefill.",
662 + "Encoder decoder models do not support chunked prefill.", # noqa: E501
663 ),
664 ],
665 )
@@ -673,100 +685,100 @@ def test_is_chunked_prefill_supported(
685 "jason9693/Qwen2.5-1.5B-apeach",
686 "decoder",
687 True,
676 - "Pooling models with causal attn and last pooling support prefix caching.",
688 + "Pooling models with causal attn and LAST/ALL pooling support prefix caching.", # noqa: E501
689 ),
690 (
691 "Qwen/Qwen3-Embedding-0.6B",
692 "decoder",
693 True,
682 - "Pooling models with causal attn and last pooling support prefix caching.",
694 + "Pooling models with causal attn and LAST/ALL pooling support prefix caching.", # noqa: E501
695 ),
696 (
697 "Qwen/Qwen2.5-Math-PRM-7B",
698 "decoder",
699 False,
688 - "Pooling models with step pooling does not support prefix caching.",
700 + "Pooling models with causal attn and LAST/STEP pooling do not support prefix caching.", # noqa: E501
701 ),
702 (
703 "internlm/internlm2-1_8b-reward",
704 "decoder",
705 True,
694 - "Pooling models with causal attn and all pooling support prefix caching.",
706 + "Pooling models with causal attn and LAST/ALL pooling support prefix caching.", # noqa: E501
707 ),
708 (
709 "BAAI/bge-base-en",
710 "encoder_only",
711 False,
700 - "Pooling models with bidirectional attn does not support prefix caching.",
712 + "Pooling models with bidirectional attn do not support prefix caching.", # noqa: E501
713 ),
714 (
715 "boltuix/NeuroBERT-NER",
716 "encoder_only",
717 False,
706 - "Pooling models with bidirectional attn does not support prefix caching.",
718 + "Pooling models with bidirectional attn do not support prefix caching.", # noqa: E501
719 ),
720 (
721 "papluca/xlm-roberta-base-language-detection",
722 "encoder_only",
723 False,
712 - "Pooling models with bidirectional attn does not support prefix caching.",
724 + "Pooling models with bidirectional attn do not support prefix caching.", # noqa: E501
725 ),
726 (
727 "Alibaba-NLP/gte-Qwen2-1.5B-instruct",
728 "encoder_only",
729 False,
718 - "Pooling models with bidirectional attn does not support prefix caching.",
730 + "Pooling models with bidirectional attn do not support prefix caching.", # noqa: E501
731 ),
732 (
733 "intfloat/e5-small",
734 "encoder_only",
735 False,
724 - "Pooling models with bidirectional attn does not support prefix caching.",
736 + "Pooling models with bidirectional attn do not support prefix caching.", # noqa: E501
737 ),
738 # multimodal models
739 (
740 "openai/clip-vit-base-patch32",
741 "decoder",
742 True,
731 - "Pooling models with causal attn and last pooling support prefix caching.",
743 + "Pooling models with causal attn and LAST/ALL pooling support prefix caching.", # noqa: E501
744 ),
745 (
746 "google/siglip-base-patch16-224",
747 "encoder_only",
748 False,
737 - "Pooling models with bidirectional attn does not support prefix caching.",
749 + "Pooling models with bidirectional attn do not support prefix caching.", # noqa: E501
750 ),
751 # generate models
752 (
753 "Qwen/Qwen3-0.6B",
754 "decoder",
755 True,
744 - "Generative models support prefix caching.",
756 + "Generative models support prefix caching.", # noqa: E501
757 ),
758 (
759 "Qwen/Qwen3-Next-80B-A3B-Instruct",
760 "hybrid",
761 False,
750 - "Hybrid models does not support prefix caching since the feature is still experimental.", # noqa: E501
762 + "Hybrid models do not support prefix caching since the feature is still experimental.", # noqa: E501
763 ),
764 (
765 "ibm-granite/granite-4.0-h-small",
766 "hybrid",
767 False,
756 - "Hybrid models does not support prefix caching since the feature is still experimental.", # noqa: E501
768 + "Hybrid models do not support prefix caching since the feature is still experimental.", # noqa: E501
769 ),
770 (
771 "state-spaces/mamba-130m-hf",
772 "attention_free",
773 False,
762 - "Attention free models does not support prefix caching since the feature is still experimental.", # noqa: E501
774 + "Attention free models do not support prefix caching since the feature is still experimental.", # noqa: E501
775 ),
776 # encoder_decoder models
777 (
778 "openai/whisper-small",
779 "encoder_decoder",
780 False,
769 - "Encoder decoder models does not support prefix caching.",
781 + "Encoder decoder models do not support prefix caching.", # noqa: E501
782 ),
783 ],
784 )
tests/test_pooling_params.py
+4 -4
@@ -40,7 +40,7 @@ def test_task():
40
41 def test_embed():
42 task = "embed"
43 - model_config = MockModelConfig(pooler_config=PoolerConfig(pooling_type="CLS"))
43 + model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
44
45 pooling_params = PoolingParams(normalize=None)
46 pooling_params.verify(task=task, model_config=model_config)
@@ -86,7 +86,7 @@ def test_embed_dimensions(model_info: EmbedModelInfo):
86
87 @pytest.mark.parametrize("task", ["score", "classify"])
88 def test_classify(task):
89 - model_config = MockModelConfig(pooler_config=PoolerConfig(pooling_type="CLS"))
89 + model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
90
91 pooling_params = PoolingParams(use_activation=None)
92 pooling_params.verify(task=task, model_config=model_config)
@@ -108,7 +108,7 @@ def test_classify(task):
108 def test_token_embed(pooling_type: str):
109 task = "token_embed"
110 model_config = MockModelConfig(
111 - pooler_config=PoolerConfig(pooling_type=pooling_type)
111 + pooler_config=PoolerConfig(tok_pooling_type=pooling_type)
112 )
113
114 pooling_params = PoolingParams(normalize=None)
@@ -134,7 +134,7 @@ def test_token_embed(pooling_type: str):
134 def test_token_classify(pooling_type: str):
135 task = "token_classify"
136 model_config = MockModelConfig(
137 - pooler_config=PoolerConfig(pooling_type=pooling_type)
137 + pooler_config=PoolerConfig(tok_pooling_type=pooling_type)
138 )
139
140 pooling_params = PoolingParams(use_activation=None)
vllm/config/model.py
+55 -39
@@ -539,9 +539,12 @@ class ModelConfig:
539 if getattr(self.pooler_config, k) is None:
540 setattr(self.pooler_config, k, v)
541
542 - default_pooling_type = self._model_info.default_pooling_type
543 - if self.pooler_config.pooling_type is None:
544 - self.pooler_config.pooling_type = default_pooling_type
542 + default_seq_pooling_type = self._model_info.default_seq_pooling_type
543 + if self.pooler_config.seq_pooling_type is None:
544 + self.pooler_config.seq_pooling_type = default_seq_pooling_type
545 + default_tok_pooling_type = self._model_info.default_tok_pooling_type
546 + if self.pooler_config.tok_pooling_type is None:
547 + self.pooler_config.tok_pooling_type = default_tok_pooling_type
548
549 self.dtype: torch.dtype = _get_and_verify_dtype(
550 self.model,
@@ -1543,8 +1546,8 @@ class ModelConfig:
1546 @property
1547 def attn_type(self) -> AttnTypeStr:
1548 if self.pooler_config is not None:
1546 - pooling_type = self._model_info.default_pooling_type.lower()
1547 - if pooling_type == "cls":
1549 + seq_pooling_type = self._model_info.default_seq_pooling_type
1550 + if seq_pooling_type == "CLS":
1551 return "encoder_only"
1552 else:
1553 is_causal = getattr(self.hf_config, "is_causal", True)
@@ -1561,89 +1564,102 @@ class ModelConfig:
1564 @property
1565 def is_chunked_prefill_supported(self) -> bool:
1566 attn_type = self.attn_type
1564 - if self.pooler_config is not None:
1567 +
1568 + if pooler_config := self.pooler_config:
1569 # for pooling models
1570 if attn_type == "encoder_only":
1571 logger.debug(
1568 - "Pooling models with bidirectional attn does not support "
1569 - "chunked prefill."
1572 + "Pooling models with bidirectional attn "
1573 + "do not support chunked prefill."
1574 )
1575 return False
1572 - elif attn_type == "decoder":
1573 - pooling_type = self.pooler_config.pooling_type.lower()
1574 - if pooling_type in ["mean", "step", "cls"]:
1576 +
1577 + if attn_type == "decoder":
1578 + if (
1579 + pooler_config.seq_pooling_type in ("MEAN", "CLS")
1580 + or pooler_config.tok_pooling_type == "STEP"
1581 + ):
1582 logger.debug(
1576 - "Pooling models with %s pooling does not "
1577 - "support chunked prefill.",
1578 - pooling_type,
1583 + "Pooling models with causal attn and %s/%s pooling "
1584 + "do not support chunked prefill.",
1585 + pooler_config.seq_pooling_type,
1586 + pooler_config.tok_pooling_type,
1587 )
1588 return False
1581 - elif pooling_type in ["all", "last"]:
1589 + else:
1590 logger.debug(
1583 - "Pooling models with causal attn and %s pooling support "
1584 - "chunked prefill.",
1585 - pooling_type,
1591 + "Pooling models with causal attn and %s/%s pooling "
1592 + "support chunked prefill.",
1593 + pooler_config.seq_pooling_type,
1594 + pooler_config.tok_pooling_type,
1595 )
1596 return True
1588 - else:
1589 - raise ValueError(f"{pooling_type=} not supported.")
1597 +
1598 # vllm currently does not have pooling models using hybrid,
1599 # attention_free or encoder_decoder attn types.
1600 return attn_type != "encoder_decoder"
1601 else:
1602 + # for generative models
1603 if attn_type == "encoder_decoder":
1595 - logger.debug("Encoder decoder models does not support chunked prefill.")
1604 + logger.debug("Encoder decoder models do not support chunked prefill.")
1605 return False
1606 +
1607 logger.debug("Generative models support chunked prefill.")
1608 return True
1609
1610 @property
1611 def is_prefix_caching_supported(self) -> bool:
1612 attn_type = self.attn_type
1603 - if self.pooler_config is not None:
1613 +
1614 + if pooler_config := self.pooler_config:
1615 # for pooling models
1616 if attn_type == "encoder_only":
1617 logger.debug(
1607 - "Pooling models with bidirectional attn does not "
1608 - "support prefix caching."
1618 + "Pooling models with bidirectional attn "
1619 + "do not support prefix caching."
1620 )
1621 return False
1611 - elif attn_type == "decoder":
1612 - pooling_type = self.pooler_config.pooling_type.lower()
1613 - if pooling_type in ["mean", "step", "cls"]:
1622 +
1623 + if attn_type == "decoder":
1624 + if (
1625 + pooler_config.seq_pooling_type in ("MEAN", "CLS")
1626 + or pooler_config.tok_pooling_type == "STEP"
1627 + ):
1628 logger.debug(
1615 - "Pooling models with %s pooling does not "
1616 - "support prefix caching.",
1617 - pooling_type,
1629 + "Pooling models with causal attn and %s/%s pooling "
1630 + "do not support prefix caching.",
1631 + pooler_config.seq_pooling_type,
1632 + pooler_config.tok_pooling_type,
1633 )
1634 return False
1620 - elif pooling_type in ["all", "last"]:
1635 + else:
1636 logger.debug(
1622 - "Pooling models with causal attn and %s pooling support "
1623 - "prefix caching.",
1624 - pooling_type,
1637 + "Pooling models with causal attn and %s/%s pooling "
1638 + "support prefix caching.",
1639 + pooler_config.seq_pooling_type,
1640 + pooler_config.tok_pooling_type,
1641 )
1642 return True
1627 - else:
1628 - raise ValueError(f"{pooling_type=} not supported.")
1643 +
1644 # vllm currently does not have pooling models using hybrid,
1645 # attention_free or encoder_decoder attn types.
1646 return False
1647 else:
1648 + # for generative models
1649 if attn_type == "hybrid":
1650 logger.debug(
1635 - "Hybrid models does not support prefix caching since the feature "
1651 + "Hybrid models do not support prefix caching since the feature "
1652 "is still experimental."
1653 )
1654 return False
1655 elif attn_type == "attention_free":
1656 logger.debug(
1641 - "Attention free models does not support prefix caching since the "
1657 + "Attention free models do not support prefix caching since the "
1658 "feature is still experimental."
1659 )
1660 return False
1661 elif attn_type == "encoder_decoder":
1646 - logger.debug("Encoder decoder models does not support prefix caching.")
1662 + logger.debug("Encoder decoder models do not support prefix caching.")
1663 return False
1664 else: # attn_type == "decoder"
1665 logger.debug("Generative models support prefix caching.")
vllm/config/pooler.py
+59 -7
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import Any, Literal
4 +from typing import Any, Literal, get_args
5
6 from pydantic.dataclasses import dataclass
7
@@ -11,7 +11,11 @@ from vllm.utils.hashing import safe_hash
11
12 logger = init_logger(__name__)
13
14 -PoolingTypeStr = Literal["LAST", "ALL", "CLS", "STEP", "MEAN"]
14 +SequencePoolingType = Literal["CLS", "LAST", "MEAN"]
15 +SEQ_POOLING_TYPES: tuple[SequencePoolingType, ...] = get_args(SequencePoolingType)
16 +
17 +TokenPoolingType = Literal["ALL", "STEP"]
18 +TOK_POOLING_TYPES: tuple[TokenPoolingType, ...] = get_args(TokenPoolingType)
19
20
21 @config
@@ -19,9 +23,26 @@ PoolingTypeStr = Literal["LAST", "ALL", "CLS", "STEP", "MEAN"]
23 class PoolerConfig:
24 """Controls the behavior of output pooling in pooling models."""
25
22 - pooling_type: PoolingTypeStr | None = None
26 + pooling_type: SequencePoolingType | TokenPoolingType | None = None
27 + """
28 + The pooling method used for pooling.
29 +
30 + If set, `seq_pooling_type` or `tok_pooling_type` are automatically populated
31 + with this field. Alternatively, users can set `seq_pooling_type` and
32 + `tok_pooling_type` explicitly.
33 +
34 + This field is mainly for user convenience. Internal code should always use
35 + `seq_pooling_type` or `tok_pooling_type` instead of `pooling_type`.
36 + """
37 +
38 + seq_pooling_type: SequencePoolingType | None = None
39 + """
40 + The pooling method used for sequence pooling.
41 + """
42 +
43 + tok_pooling_type: TokenPoolingType | None = None
44 """
24 - The pooling method of the pooling model.
45 + The pooling method used for tokenwise pooling.
46 """
47
48 ## for embeddings models
@@ -88,9 +109,40 @@ class PoolerConfig:
109 # raise deprecated warning for softmax and activation
110 self.use_activation = get_use_activation(self)
111
91 - def get_pooling_type(self) -> PoolingTypeStr:
92 - assert self.pooling_type is not None, "Should be resolved by ModelConfig"
93 - return self.pooling_type
112 + if pooling_type := self.pooling_type:
113 + if self.seq_pooling_type is not None:
114 + raise ValueError(
115 + "Cannot set both `pooling_type` and `seq_pooling_type`"
116 + )
117 + if self.tok_pooling_type is not None:
118 + raise ValueError(
119 + "Cannot set both `pooling_type` and `tok_pooling_type`"
120 + )
121 +
122 + if pooling_type in SEQ_POOLING_TYPES:
123 + logger.debug(
124 + "Resolved `pooling_type=%r` to `seq_pooling_type=%r`.",
125 + pooling_type,
126 + pooling_type,
127 + )
128 + self.seq_pooling_type = pooling_type
129 + elif pooling_type in TOK_POOLING_TYPES:
130 + logger.debug(
131 + "Resolved `pooling_type=%r` to `tok_pooling_type=%r`.",
132 + pooling_type,
133 + pooling_type,
134 + )
135 + self.tok_pooling_type = pooling_type
136 + else:
137 + raise NotImplementedError(pooling_type)
138 +
139 + def get_seq_pooling_type(self) -> SequencePoolingType:
140 + assert self.seq_pooling_type is not None, "Should be resolved by ModelConfig"
141 + return self.seq_pooling_type
142 +
143 + def get_tok_pooling_type(self) -> TokenPoolingType:
144 + assert self.tok_pooling_type is not None, "Should be resolved by ModelConfig"
145 + return self.tok_pooling_type
146
147 def compute_hash(self) -> str:
148 """
vllm/entrypoints/llm.py
+1 -1
@@ -172,7 +172,7 @@ class LLM:
172 The available overrides depend on the model that is being run.
173 For example, for Phi-3-Vision: `{"num_crops": 4}`.
174 pooler_config: Initialize non-default pooling config for the pooling
175 - model. e.g. `PoolerConfig(pooling_type="mean", normalize=False)`.
175 + model. e.g. `PoolerConfig(seq_pooling_type="MEAN", normalize=False)`.
176 compilation_config: Either an integer or a dictionary. If it is an
177 integer, it is used as the mode of compilation optimization. If it
178 is a dictionary, it can specify the full compilation configuration.
vllm/model_executor/layers/pooler/seqwise/methods.py
+4 -4
@@ -7,7 +7,7 @@ from typing import TypeAlias
7 import torch
8 import torch.nn as nn
9
10 -from vllm.config.pooler import PoolingTypeStr
10 +from vllm.config.pooler import SequencePoolingType
11 from vllm.model_executor.layers.pooler import PoolingParamsUpdate
12 from vllm.tasks import PoolingTask
13 from vllm.v1.pool.metadata import PoolingMetadata
@@ -82,11 +82,11 @@ class MeanPool(SequencePoolingMethod):
82 ) / prompt_lens.unsqueeze(1)
83
84
85 -def get_seq_pooling_method(pooling_type: PoolingTypeStr | str):
86 - if pooling_type == "LAST":
87 - return LastPool()
85 +def get_seq_pooling_method(pooling_type: SequencePoolingType | str):
86 if pooling_type == "CLS":
87 return CLSPool()
88 + if pooling_type == "LAST":
89 + return LastPool()
90 if pooling_type == "MEAN":
91 return MeanPool()
92
vllm/model_executor/layers/pooler/seqwise/poolers.py
+2 -2
@@ -85,7 +85,7 @@ class SequencePooler(Pooler):
85
86
87 def pooler_for_embed(pooler_config: PoolerConfig):
88 - pooling = get_seq_pooling_method(pooler_config.get_pooling_type())
88 + pooling = get_seq_pooling_method(pooler_config.get_seq_pooling_type())
89 head = EmbeddingPoolerHead()
90
91 return SequencePooler(pooling=pooling, head=head)
@@ -99,7 +99,7 @@ def pooler_for_classify(
99 act_fn: PoolerActivation | str | None = None,
100 ):
101 if pooling is None:
102 - pooling = get_seq_pooling_method(pooler_config.get_pooling_type())
102 + pooling = get_seq_pooling_method(pooler_config.get_seq_pooling_type())
103
104 head = ClassifierPoolerHead(classifier=classifier, act_fn=act_fn)
105
vllm/model_executor/layers/pooler/tokwise/methods.py
+2 -4
@@ -8,7 +8,7 @@ import torch
8 import torch.nn as nn
9
10 from vllm.config import get_current_vllm_config
11 -from vllm.config.pooler import PoolingTypeStr
11 +from vllm.config.pooler import TokenPoolingType
12 from vllm.model_executor.layers.pooler import PoolingParamsUpdate
13 from vllm.tasks import PoolingTask
14 from vllm.v1.pool.metadata import PoolingMetadata
@@ -113,12 +113,10 @@ class StepPool(AllPool):
113 return pooled_data
114
115
116 -def get_tok_pooling_method(pooling_type: PoolingTypeStr | str):
116 +def get_tok_pooling_method(pooling_type: TokenPoolingType | str):
117 if pooling_type == "ALL":
118 return AllPool()
119 if pooling_type == "STEP":
120 return StepPool()
121
122 - # TODO: Separate seq and tok pooling types so we don't need this fallback
123 - return AllPool()
122 raise NotImplementedError(f"Unknown tokenwise pooling type: {pooling_type!r}")
vllm/model_executor/layers/pooler/tokwise/poolers.py
+2 -2
@@ -85,7 +85,7 @@ class TokenPooler(Pooler):
85
86
87 def pooler_for_token_embed(pooler_config: PoolerConfig):
88 - pooling = get_tok_pooling_method(pooler_config.get_pooling_type())
88 + pooling = get_tok_pooling_method(pooler_config.get_tok_pooling_type())
89 head = TokenEmbeddingPoolerHead()
90
91 return TokenPooler(pooling=pooling, head=head)
@@ -99,7 +99,7 @@ def pooler_for_token_classify(
99 act_fn: PoolerActivation | str | None = None,
100 ):
101 if pooling is None:
102 - pooling = get_tok_pooling_method(pooler_config.get_pooling_type())
102 + pooling = get_tok_pooling_method(pooler_config.get_tok_pooling_type())
103
104 head = TokenClassifierPoolerHead(classifier=classifier, act_fn=act_fn)
105
vllm/model_executor/models/bert.py
+5 -5
@@ -357,7 +357,7 @@ class BertOutput(nn.Module):
357
358
359 @support_torch_compile
360 -@default_pooling_type("CLS")
360 +@default_pooling_type(seq_pooling_type="CLS")
361 class BertModel(nn.Module, SupportsQuant):
362 is_pooling_model = True
363
@@ -461,7 +461,7 @@ class BertPoolingModel(BertModel):
461 return loaded_params
462
463
464 -@default_pooling_type("CLS")
464 +@default_pooling_type(seq_pooling_type="CLS")
465 class BertEmbeddingModel(nn.Module, SupportsQuant):
466 """A model that uses Bert to provide embedding functionalities.
467
@@ -675,7 +675,7 @@ class SPLADESparsePooler(Pooler):
675 return torch.stack(pooled_list, dim=0).contiguous()
676
677
678 -@default_pooling_type("CLS")
678 +@default_pooling_type(seq_pooling_type="CLS")
679 class BertSpladeSparseEmbeddingModel(BertEmbeddingModel):
680 """
681 BertEmbeddingModel + SPLADE sparse embedding.
@@ -780,7 +780,7 @@ class BertSpladeSparseEmbeddingModel(BertEmbeddingModel):
780 return loaded
781
782
783 -@default_pooling_type("CLS")
783 +@default_pooling_type(seq_pooling_type="CLS")
784 class BertForSequenceClassification(nn.Module, SupportsCrossEncoding, SupportsQuant):
785 """A model that uses Bert to provide embedding functionalities.
786
@@ -849,7 +849,7 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding, SupportsQu
849
850
851 @attn_type("encoder_only")
852 -@default_pooling_type("ALL")
852 +@default_pooling_type(tok_pooling_type="ALL")
853 class BertForTokenClassification(nn.Module):
854 is_pooling_model = True
855
vllm/model_executor/models/bert_with_rope.py
+2 -2
@@ -441,7 +441,7 @@ class BertWithRopeEncoder(nn.Module):
441
442
443 @support_torch_compile
444 -@default_pooling_type("CLS")
444 +@default_pooling_type(seq_pooling_type="CLS")
445 class BertWithRope(nn.Module, SupportsQuant):
446 hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
447
@@ -670,7 +670,7 @@ class JinaRobertaModel(BertWithRope):
670 return super().load_weights(weights)
671
672
673 -@default_pooling_type("CLS")
673 +@default_pooling_type(seq_pooling_type="CLS")
674 class GteNewForSequenceClassification(nn.Module, SupportsCrossEncoding):
675 is_pooling_model = True
676
vllm/model_executor/models/clip.py
+3 -3
@@ -145,7 +145,7 @@ class CLIPProcessingInfo(BaseProcessingInfo):
145 image_width=image_width,
146 image_height=image_height,
147 ),
148 - _get_vision_feature_select_strategy(pooler_config.pooling_type),
148 + _get_vision_feature_select_strategy(pooler_config.seq_pooling_type),
149 )
150
151 def get_image_size_with_most_features(self) -> ImageSize:
@@ -819,7 +819,7 @@ class CLIPVisionModel(nn.Module):
819
820
821 # Assume EOS token corresponds to LAST token in text model
822 -@default_pooling_type("LAST")
822 +@default_pooling_type(seq_pooling_type="LAST")
823 @MULTIMODAL_REGISTRY.register_processor(
824 CLIPMultiModalProcessor,
825 info=CLIPProcessingInfo,
@@ -908,7 +908,7 @@ class CLIPEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
908 ) -> torch.Tensor:
909 if feature_select_strategy is None:
910 feature_select_strategy = _get_vision_feature_select_strategy(
911 - self.pooler_config.pooling_type
911 + self.pooler_config.seq_pooling_type
912 )
913
914 pooled_output = self.vision_model(
vllm/model_executor/models/config.py
+5 -4
@@ -94,12 +94,12 @@ class JinaRobertaModelConfig(VerifyAndUpdateConfig):
94 class LlamaBidirectionalConfig(VerifyAndUpdateConfig):
95 @staticmethod
96 def verify_and_update_model_config(model_config: "ModelConfig") -> None:
97 - from vllm.config.pooler import PoolingTypeStr
97 + from vllm.config.pooler import SequencePoolingType
98
99 hf_config = model_config.hf_config
100 hf_config.is_causal = False
101
102 - pooling_type_map: dict[str, PoolingTypeStr] = {
102 + pooling_type_map: dict[str, SequencePoolingType] = {
103 "avg": "MEAN",
104 "cls": "CLS",
105 "last": "LAST",
@@ -107,8 +107,9 @@ class LlamaBidirectionalConfig(VerifyAndUpdateConfig):
107
108 pooling_type = pooling_type_map.get(hf_config.pooling, None)
109 if pooling_type is None:
110 - raise ValueError(f"pool_type {hf_config.pooling} not supported")
111 - model_config.pooler_config.pooling_type = pooling_type
110 + raise ValueError(f"pool_type {hf_config.pooling!r} not supported")
111 +
112 + model_config.pooler_config.seq_pooling_type = pooling_type
113
114
115 class NomicBertModelConfig(VerifyAndUpdateConfig):
vllm/model_executor/models/gritlm.py
+1 -1
@@ -193,7 +193,7 @@ class GritLMPooler(SequencePooler):
193 return self.activation(pooled_data)
194
195
196 -@default_pooling_type("MEAN")
196 +@default_pooling_type(seq_pooling_type="MEAN")
197 class GritLM(LlamaForCausalLM):
198 """This class implements the embedding model for parasail-ai/GritLM-7B-vllm.
199
vllm/model_executor/models/interfaces_base.py
+33 -9
@@ -20,12 +20,13 @@ from vllm.utils.func_utils import supports_kw
20 if TYPE_CHECKING:
21 from vllm.config import VllmConfig
22 from vllm.config.model import AttnTypeStr
23 - from vllm.config.pooler import PoolingTypeStr
23 + from vllm.config.pooler import SequencePoolingType, TokenPoolingType
24 from vllm.model_executor.layers.pooler import Pooler
25 else:
26 VllmConfig = Any
27 Pooler = Any
28 - PoolingTypeStr = Any
28 + SequencePoolingType = Any
29 + TokenPoolingType = Any
30 AttnTypeStr = Any
31
32 logger = init_logger(__name__)
@@ -155,9 +156,19 @@ class VllmModelForPooling(VllmModel[T_co], Protocol[T_co]):
156 MRO of your model class.
157 """
158
158 - default_pooling_type: ClassVar[PoolingTypeStr] = "LAST"
159 + default_seq_pooling_type: ClassVar[SequencePoolingType] = "LAST"
160 """
160 - Indicates the [vllm.config.pooler.PoolerConfig.pooling_type][]
161 + Indicates the [vllm.config.pooler.PoolerConfig.seq_pooling_type][]
162 + to use by default.
163 +
164 + You can use the
165 + [vllm.model_executor.models.interfaces_base.default_pooling_type][]
166 + decorator to conveniently set this field.
167 + """
168 +
169 + default_tok_pooling_type: ClassVar[TokenPoolingType] = "ALL"
170 + """
171 + Indicates the [vllm.config.pooler.PoolerConfig.tok_pooling_type][]
172 to use by default.
173
174 You can use the
@@ -200,18 +211,31 @@ def is_pooling_model(
211 _T = TypeVar("_T", bound=type[nn.Module])
212
213
203 -def default_pooling_type(pooling_type: PoolingTypeStr):
204 - """Decorator to set `VllmModelForPooling.default_pooling_type`."""
214 +def default_pooling_type(
215 + *,
216 + seq_pooling_type: SequencePoolingType = "LAST",
217 + tok_pooling_type: TokenPoolingType = "ALL",
218 +):
219 + """Decorator to set `VllmModelForPooling.default_*_pooling_type`."""
220
221 def func(model: _T) -> _T:
207 - model.default_pooling_type = pooling_type # type: ignore
222 + model.default_seq_pooling_type = seq_pooling_type # type: ignore
223 + model.default_tok_pooling_type = tok_pooling_type # type: ignore
224 return model
225
226 return func
227
228
213 -def get_default_pooling_type(model: type[object] | object) -> PoolingTypeStr:
214 - return getattr(model, "default_pooling_type", "LAST")
229 +def get_default_seq_pooling_type(
230 + model: type[object] | object,
231 +) -> SequencePoolingType:
232 + return getattr(model, "default_seq_pooling_type", "LAST")
233 +
234 +
235 +def get_default_tok_pooling_type(
236 + model: type[object] | object,
237 +) -> TokenPoolingType:
238 + return getattr(model, "default_tok_pooling_type", "ALL")
239
240
241 def attn_type(attn_type: AttnTypeStr):
vllm/model_executor/models/internlm2.py
+1 -1
@@ -402,7 +402,7 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
402 return loaded_params
403
404
405 -@default_pooling_type("ALL")
405 +@default_pooling_type(tok_pooling_type="ALL")
406 class InternLM2ForRewardModel(InternLM2ForCausalLM):
407 is_pooling_model = True
408
vllm/model_executor/models/modernbert.py
+3 -3
@@ -221,7 +221,7 @@ class ModernBertEncoderLayer(nn.Module):
221
222
223 @support_torch_compile
224 -@default_pooling_type("CLS")
224 +@default_pooling_type(seq_pooling_type="CLS")
225 class ModernBertModel(nn.Module):
226 hf_to_vllm_mapper = WeightsMapper(
227 orig_to_new_prefix={"layers.": "encoder_layer.layers."}
@@ -308,7 +308,7 @@ class ModernBertPooler(SequencePooler):
308 return self.norm(self.act(self.dense(pooled_data)))
309
310
311 -@default_pooling_type("CLS")
311 +@default_pooling_type(seq_pooling_type="CLS")
312 class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):
313 is_pooling_model = True
314
@@ -395,7 +395,7 @@ class ModernBertPredictionHead(nn.Module):
395
396
397 @attn_type("encoder_only")
398 -@default_pooling_type("ALL")
398 +@default_pooling_type(tok_pooling_type="ALL")
399 class ModernBertForTokenClassification(nn.Module):
400 is_pooling_model = True
401
vllm/model_executor/models/qwen2_rm.py
+2 -2
@@ -96,7 +96,7 @@ class Qwen2RewardBaseModel(nn.Module, SupportsLoRA, SupportsPP):
96 return loader.load_weights(weights)
97
98
99 -@default_pooling_type("ALL")
99 +@default_pooling_type(tok_pooling_type="ALL")
100 class Qwen2ForRewardModel(Qwen2RewardBaseModel):
101 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
102 vllm_config.model_config.hf_config.num_labels = 1
@@ -108,7 +108,7 @@ class Qwen2ForRewardModel(Qwen2RewardBaseModel):
108 self.pooler = pooler_for_token_classify(pooler_config)
109
110
111 -@default_pooling_type("STEP")
111 +@default_pooling_type(tok_pooling_type="STEP")
112 class Qwen2ForProcessRewardModel(Qwen2RewardBaseModel):
113 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
114 vllm_config.model_config.hf_config.num_labels = 2
vllm/model_executor/models/registry.py
+9 -5
@@ -35,10 +35,11 @@ from vllm.utils.hashing import safe_hash
35
36 if TYPE_CHECKING:
37 from vllm.config.model import AttnTypeStr
38 - from vllm.config.pooler import PoolingTypeStr
38 + from vllm.config.pooler import SequencePoolingType, TokenPoolingType
39 else:
40 AttnTypeStr = Any
41 - PoolingTypeStr = Any
41 + SequencePoolingType = Any
42 + TokenPoolingType = Any
43
44
45 from .interfaces import (
@@ -57,7 +58,8 @@ from .interfaces import (
58 )
59 from .interfaces_base import (
60 get_attn_type,
60 - get_default_pooling_type,
61 + get_default_seq_pooling_type,
62 + get_default_tok_pooling_type,
63 is_pooling_model,
64 is_text_generation_model,
65 )
@@ -548,7 +550,8 @@ class _ModelInfo:
550 is_text_generation_model: bool
551 is_pooling_model: bool
552 attn_type: AttnTypeStr
551 - default_pooling_type: PoolingTypeStr
553 + default_seq_pooling_type: SequencePoolingType
554 + default_tok_pooling_type: TokenPoolingType
555 supports_cross_encoding: bool
556 supports_multimodal: bool
557 supports_multimodal_raw_input_only: bool
@@ -569,7 +572,8 @@ class _ModelInfo:
572 architecture=model.__name__,
573 is_text_generation_model=is_text_generation_model(model),
574 is_pooling_model=is_pooling_model(model),
572 - default_pooling_type=get_default_pooling_type(model),
575 + default_seq_pooling_type=get_default_seq_pooling_type(model),
576 + default_tok_pooling_type=get_default_tok_pooling_type(model),
577 attn_type=get_attn_type(model),
578 supports_cross_encoding=supports_cross_encoding(model),
579 supports_multimodal=supports_multimodal(model),
vllm/model_executor/models/roberta.py
+2 -2
@@ -93,7 +93,7 @@ class RobertaClassificationHead(nn.Module):
93 return x
94
95
96 -@default_pooling_type("CLS")
96 +@default_pooling_type(seq_pooling_type="CLS")
97 class RobertaEmbeddingModel(BertEmbeddingModel):
98 """A model that uses Roberta to provide embedding functionalities."""
99
@@ -150,7 +150,7 @@ class RobertaEmbeddingModel(BertEmbeddingModel):
150 return loader.load_weights(weights_list, mapper=mapper)
151
152
153 -@default_pooling_type("CLS")
153 +@default_pooling_type(seq_pooling_type="CLS")
154 class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):
155 """A model that uses Roberta to provide embedding functionalities.
156
vllm/model_executor/models/siglip.py
+3 -3
@@ -129,7 +129,7 @@ class SiglipProcessingInfo(BaseProcessingInfo):
129 image_width=image_width,
130 image_height=image_height,
131 ),
132 - _get_vision_feature_select_strategy(pooler_config.pooling_type),
132 + _get_vision_feature_select_strategy(pooler_config.seq_pooling_type),
133 )
134
135 def get_image_size_with_most_features(self) -> ImageSize:
@@ -998,7 +998,7 @@ class SiglipTextEmbeddings(nn.Module):
998
999
1000 # Assume EOS token corresponds to CLS token in text model
1001 -@default_pooling_type("CLS")
1001 +@default_pooling_type(seq_pooling_type="CLS")
1002 @MULTIMODAL_REGISTRY.register_processor(
1003 SiglipMultiModalProcessor,
1004 info=SiglipProcessingInfo,
@@ -1125,7 +1125,7 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):
1125 ) -> torch.Tensor:
1126 if feature_select_strategy is None:
1127 feature_select_strategy = _get_vision_feature_select_strategy(
1128 - self.pooler_config.pooling_type
1128 + self.pooler_config.seq_pooling_type
1129 )
1130
1131 pooled_output = self.vision_model(
vllm/pooling_params.py
+1 -1
@@ -140,7 +140,7 @@ class PoolingParams(
140 self, pooler_config: "PoolerConfig", valid_parameters: list[str]
141 ):
142 step_pooling_parameters = ["step_tag_id", "returned_token_ids"]
143 - if pooler_config.pooling_type != "STEP":
143 + if pooler_config.tok_pooling_type != "STEP":
144 invalid_parameters = []
145 for k in step_pooling_parameters:
146 if getattr(self, k, None) is not None:
vllm/tasks.py
+2 -2
@@ -3,11 +3,11 @@
3 from typing import Literal, get_args
4
5 GenerationTask = Literal["generate", "transcription"]
6 -GENERATION_TASKS = get_args(GenerationTask)
6 +GENERATION_TASKS: tuple[GenerationTask, ...] = get_args(GenerationTask)
7
8 PoolingTask = Literal[
9 "embed", "classify", "score", "token_embed", "token_classify", "plugin"
10 ]
11 -POOLING_TASKS = get_args(PoolingTask)
11 +POOLING_TASKS: tuple[PoolingTask, ...] = get_args(PoolingTask)
12
13 SupportedTask = Literal[GenerationTask, PoolingTask]
vllm/transformers_utils/config.py
+24 -21
@@ -10,9 +10,7 @@ from pathlib import Path
10 from typing import Any, Literal, TypeAlias
11
12 import huggingface_hub
13 -from huggingface_hub import (
14 - get_safetensors_metadata,
15 -)
13 +from huggingface_hub import get_safetensors_metadata
14 from packaging.version import Version
15 from transformers import GenerationConfig, PretrainedConfig
16 from transformers.models.auto.image_processing_auto import get_image_processor_config
@@ -742,7 +740,10 @@ def get_config(
740
741
742 @cache
745 -def get_pooling_config(model: str, revision: str | None = "main") -> dict | None:
743 +def get_pooling_config(
744 + model: str,
745 + revision: str | None = "main",
746 +) -> dict[str, Any] | None:
747 """
748 This function gets the pooling and normalize
749 config from the model - only applies to
@@ -793,38 +794,40 @@ def get_pooling_config(model: str, revision: str | None = "main") -> dict | None
794 )
795
796 if pooling:
796 - pooling_file_name = "{}/config.json".format(pooling["path"])
797 - pooling_dict = get_hf_file_to_dict(pooling_file_name, model, revision)
798 - pooling_type_name = next(
799 - (item for item, val in pooling_dict.items() if val is True), None
800 - )
797 + from vllm.config.pooler import SEQ_POOLING_TYPES, TOK_POOLING_TYPES
798
802 - if pooling_type_name is not None:
803 - pooling_type_name = get_pooling_config_name(pooling_type_name)
799 + pooling_file_name = "{}/config.json".format(pooling["path"])
800 + pooling_dict = get_hf_file_to_dict(pooling_file_name, model, revision) or {}
801
802 logger.info("Found pooling configuration.")
806 - return {"pooling_type": pooling_type_name, "normalize": normalize}
803 +
804 + config: dict[str, Any] = {"normalize": normalize}
805 + for key, val in pooling_dict.items():
806 + if val is True:
807 + pooling_type = parse_pooling_type(key)
808 + if pooling_type in SEQ_POOLING_TYPES:
809 + config["seq_pooling_type"] = pooling_type
810 + elif pooling_type in TOK_POOLING_TYPES:
811 + config["tok_pooling_type"] = pooling_type
812 + else:
813 + logger.debug("Skipping unrelated field: %r=%r", key, val)
814 +
815 + return config
816
817 return None
818
819
811 -def get_pooling_config_name(pooling_name: str) -> str | None:
820 +def parse_pooling_type(pooling_name: str):
821 if "pooling_mode_" in pooling_name:
822 pooling_name = pooling_name.replace("pooling_mode_", "")
823
824 if "_" in pooling_name:
816 - pooling_name = pooling_name.split("_")[0]
825 + pooling_name = pooling_name.split("_", 1)[0]
826
827 if "lasttoken" in pooling_name:
828 pooling_name = "last"
829
821 - supported_pooling_types = ["LAST", "ALL", "CLS", "STEP", "MEAN"]
822 - pooling_type_name = pooling_name.upper()
823 -
824 - if pooling_type_name in supported_pooling_types:
825 - return pooling_type_name
826 -
827 - raise NotImplementedError(f"Pooling type {pooling_type_name} not supported")
830 + return pooling_name.upper()
831
832
833 @cache