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 6, 23:00 UTC
[Model] Consolidate score logic by introduce score_type (#36479)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
wang.yuqi committed
Mar 10, 2026 at 21:32 UTC
a3189a08b0d3de44dd6d49c5d883abf29ac1e6fa
14 files changed
+213
-194
tests/models/registry.py
+65
-52
@@ -546,15 +546,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
546
_EMBEDDING_EXAMPLE_MODELS = {
547
# [Text-only]
548
"BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
549
- "HF_ColBERT": _HfExamplesInfo("answerdotai/answerai-colbert-small-v1"),
550
- "ColBERTModernBertModel": _HfExamplesInfo(
551
- "lightonai/GTE-ModernColBERT-v1",
552
- hf_overrides={"architectures": ["ColBERTModernBertModel"]},
553
- ),
554
- "ColBERTJinaRobertaModel": _HfExamplesInfo(
555
- "jinaai/jina-colbert-v2",
556
- trust_remote_code=True,
557
- hf_overrides={"architectures": ["ColBERTJinaRobertaModel"]},
549
+ "BertSpladeSparseEmbeddingModel": _HfExamplesInfo(
550
+ "naver/splade-v3",
551
+ hf_overrides={"architectures": ["BertSpladeSparseEmbeddingModel"]},
552
),
553
"BgeM3EmbeddingModel": _HfExamplesInfo("BAAI/bge-m3"),
554
"Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"),
@@ -568,10 +562,6 @@ _EMBEDDING_EXAMPLE_MODELS = {
562
trust_remote_code=True,
563
hf_overrides={"architectures": ["GteNewModel"]},
564
),
571
- "InternLM2ForRewardModel": _HfExamplesInfo(
572
- "internlm/internlm2-1_8b-reward", trust_remote_code=True
573
- ),
574
- "JambaForSequenceClassification": _HfExamplesInfo("ai21labs/Jamba-tiny-reward-dev"),
565
"LlamaModel": _HfExamplesInfo("llama", is_available_online=False),
566
"LlamaBidirectionalModel": _HfExamplesInfo(
567
"nvidia/llama-nemotron-embed-1b-v2", trust_remote_code=True
@@ -584,35 +574,14 @@ _EMBEDDING_EXAMPLE_MODELS = {
574
"nomic-ai/nomic-embed-text-v2-moe", trust_remote_code=True
575
),
576
"Qwen2Model": _HfExamplesInfo("ssmits/Qwen2-7B-Instruct-embed-base"),
587
- "Qwen2ForRewardModel": _HfExamplesInfo(
588
- "Qwen/Qwen2.5-Math-RM-72B",
589
- max_transformers_version="4.53",
590
- transformers_version_reason={
591
- "hf": "HF model uses remote code that is not compatible with latest Transformers" # noqa: E501
592
- },
593
- ),
594
- "Qwen2ForProcessRewardModel": _HfExamplesInfo(
595
- "Qwen/Qwen2.5-Math-PRM-7B",
596
- max_transformers_version="4.53",
597
- transformers_version_reason={
598
- "hf": "HF model uses remote code that is not compatible with latest Transformers" # noqa: E501
599
- },
600
- ),
577
"RobertaModel": _HfExamplesInfo("sentence-transformers/stsb-roberta-base-v2"),
578
"RobertaForMaskedLM": _HfExamplesInfo("sentence-transformers/all-roberta-large-v1"),
579
"VoyageQwen3BidirectionalEmbedModel": _HfExamplesInfo(
580
"voyageai/voyage-4-nano", trust_remote_code=True
581
),
582
"XLMRobertaModel": _HfExamplesInfo("intfloat/multilingual-e5-small"),
607
- "BertSpladeSparseEmbeddingModel": _HfExamplesInfo(
608
- "naver/splade-v3",
609
- hf_overrides={"architectures": ["BertSpladeSparseEmbeddingModel"]},
610
- ),
583
# [Multimodal]
584
"CLIPModel": _HfExamplesInfo("openai/clip-vit-base-patch32"),
613
- "ColModernVBertForRetrieval": _HfExamplesInfo(
614
- "ModernVBERT/colmodernvbert-merged",
615
- ),
585
"LlamaNemotronVLModel": _HfExamplesInfo(
586
"nvidia/llama-nemotron-embed-vl-1b-v2", trust_remote_code=True
587
),
@@ -621,15 +590,6 @@ _EMBEDDING_EXAMPLE_MODELS = {
590
"TIGER-Lab/VLM2Vec-Full", trust_remote_code=True
591
),
592
"Qwen2VLForConditionalGeneration": _HfExamplesInfo("MrLight/dse-qwen2-2b-mrl-v1"),
624
- "ColQwen3": _HfExamplesInfo(
625
- "TomoroAI/tomoro-colqwen3-embed-4b", trust_remote_code=True
626
- ),
627
- "OpsColQwen3Model": _HfExamplesInfo(
628
- "OpenSearch-AI/Ops-Colqwen3-4B", trust_remote_code=True
629
- ),
630
- "Qwen3VLNemotronEmbedModel": _HfExamplesInfo(
631
- "nvidia/nemotron-colembed-vl-4b-v2",
632
- ),
593
"SiglipModel": _HfExamplesInfo("google/siglip-base-patch16-224"),
594
"PrithviGeoSpatialMAE": _HfExamplesInfo(
595
"ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
@@ -649,21 +609,74 @@ _EMBEDDING_EXAMPLE_MODELS = {
609
),
610
}
611
652
-_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
653
- # [Decoder-only]
654
- "GPT2ForSequenceClassification": _HfExamplesInfo(
655
- "nie3e/sentiment-polish-gpt2-small"
612
+_LATE_INTERACTION_EXAMPLE_MODELS = {
613
+ # [Text-only]
614
+ "HF_ColBERT": _HfExamplesInfo("answerdotai/answerai-colbert-small-v1"),
615
+ "ColBERTModernBertModel": _HfExamplesInfo(
616
+ "lightonai/GTE-ModernColBERT-v1",
617
+ hf_overrides={"architectures": ["ColBERTModernBertModel"]},
618
),
657
- # [Cross-encoder]
619
+ "ColBERTJinaRobertaModel": _HfExamplesInfo(
620
+ "jinaai/jina-colbert-v2",
621
+ trust_remote_code=True,
622
+ hf_overrides={"architectures": ["ColBERTJinaRobertaModel"]},
623
+ ),
624
+ # [Multimodal]
625
+ "ColModernVBertForRetrieval": _HfExamplesInfo(
626
+ "ModernVBERT/colmodernvbert-merged",
627
+ ),
628
+ "ColQwen3": _HfExamplesInfo(
629
+ "TomoroAI/tomoro-colqwen3-embed-4b", trust_remote_code=True
630
+ ),
631
+ "OpsColQwen3Model": _HfExamplesInfo(
632
+ "OpenSearch-AI/Ops-Colqwen3-4B", trust_remote_code=True
633
+ ),
634
+ "Qwen3VLNemotronEmbedModel": _HfExamplesInfo(
635
+ "nvidia/nemotron-colembed-vl-4b-v2",
636
+ ),
637
+}
638
+
639
+
640
+_REWARD_EXAMPLE_MODELS = {
641
+ "InternLM2ForRewardModel": _HfExamplesInfo(
642
+ "internlm/internlm2-1_8b-reward", trust_remote_code=True
643
+ ),
644
+ "Qwen2ForRewardModel": _HfExamplesInfo(
645
+ "Qwen/Qwen2.5-Math-RM-72B",
646
+ max_transformers_version="4.53",
647
+ transformers_version_reason={
648
+ "hf": "HF model uses remote code that is not compatible with latest Transformers" # noqa: E501
649
+ },
650
+ ),
651
+ "Qwen2ForProcessRewardModel": _HfExamplesInfo(
652
+ "Qwen/Qwen2.5-Math-PRM-7B",
653
+ max_transformers_version="4.53",
654
+ transformers_version_reason={
655
+ "hf": "HF model uses remote code that is not compatible with latest Transformers" # noqa: E501
656
+ },
657
+ ),
658
+}
659
+
660
+_TOKEN_CLASSIFICATION_EXAMPLE_MODELS = {
661
+ "BertForTokenClassification": _HfExamplesInfo("boltuix/NeuroBERT-NER"),
662
+ "ModernBertForTokenClassification": _HfExamplesInfo(
663
+ "disham993/electrical-ner-ModernBERT-base"
664
+ ),
665
+}
666
+
667
+_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
668
"BertForSequenceClassification": _HfExamplesInfo(
669
"cross-encoder/ms-marco-MiniLM-L-6-v2"
670
),
661
- "BertForTokenClassification": _HfExamplesInfo("boltuix/NeuroBERT-NER"),
671
+ "GPT2ForSequenceClassification": _HfExamplesInfo(
672
+ "nie3e/sentiment-polish-gpt2-small"
673
+ ),
674
"GteNewForSequenceClassification": _HfExamplesInfo(
675
"Alibaba-NLP/gte-multilingual-reranker-base",
676
trust_remote_code=True,
677
hf_overrides={"architectures": ["GteNewForSequenceClassification"]},
678
),
679
+ "JambaForSequenceClassification": _HfExamplesInfo("ai21labs/Jamba-tiny-reward-dev"),
680
"LlamaBidirectionalForSequenceClassification": _HfExamplesInfo(
681
"nvidia/llama-nemotron-rerank-1b-v2", trust_remote_code=True
682
),
@@ -673,9 +686,6 @@ _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
686
"ModernBertForSequenceClassification": _HfExamplesInfo(
687
"Alibaba-NLP/gte-reranker-modernbert-base"
688
),
676
- "ModernBertForTokenClassification": _HfExamplesInfo(
677
- "disham993/electrical-ner-ModernBERT-base"
678
- ),
689
"RobertaForSequenceClassification": _HfExamplesInfo(
690
"cross-encoder/quora-roberta-base"
691
),
@@ -1273,6 +1283,9 @@ _TRANSFORMERS_BACKEND_MODELS = {
1283
_EXAMPLE_MODELS = {
1284
**_TEXT_GENERATION_EXAMPLE_MODELS,
1285
**_EMBEDDING_EXAMPLE_MODELS,
1286
+ **_LATE_INTERACTION_EXAMPLE_MODELS,
1287
+ **_REWARD_EXAMPLE_MODELS,
1288
+ **_TOKEN_CLASSIFICATION_EXAMPLE_MODELS,
1289
**_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS,
1290
**_MULTIMODAL_EXAMPLE_MODELS,
1291
**_SPECULATIVE_DECODING_EXAMPLE_MODELS,
tests/models/test_registry.py
+11
-8
@@ -56,21 +56,24 @@ def test_registry_imports(model_arch):
56
57
@create_new_process_for_each_test()
58
@pytest.mark.parametrize(
59
- "model_arch,is_mm,init_cuda,is_ce",
59
+ "model_arch,is_mm,init_cuda,score_type",
60
[
61
- ("LlamaForCausalLM", False, False, False),
62
- ("LlavaForConditionalGeneration", True, True, False),
63
- ("BertForSequenceClassification", False, False, True),
64
- ("RobertaForSequenceClassification", False, False, True),
65
- ("XLMRobertaForSequenceClassification", False, False, True),
61
+ ("LlamaForCausalLM", False, False, "bi-encoder"),
62
+ ("LlavaForConditionalGeneration", True, True, "bi-encoder"),
63
+ ("BertForSequenceClassification", False, False, "cross-encoder"),
64
+ ("RobertaForSequenceClassification", False, False, "cross-encoder"),
65
+ ("XLMRobertaForSequenceClassification", False, False, "cross-encoder"),
66
+ ("GteNewModel", False, False, "bi-encoder"),
67
+ ("GteNewForSequenceClassification", False, False, "cross-encoder"),
68
+ ("HF_ColBERT", False, False, "late-interaction"),
69
],
70
)
68
-def test_registry_model_property(model_arch, is_mm, init_cuda, is_ce):
71
+def test_registry_model_property(model_arch, is_mm, init_cuda, score_type):
72
model_info = ModelRegistry._try_inspect_model_cls(model_arch)
73
assert model_info is not None
74
75
assert model_info.supports_multimodal is is_mm
73
- assert model_info.supports_cross_encoding is is_ce
76
+ assert model_info.score_type == score_type
77
78
if init_cuda and current_platform.is_cuda_alike():
79
assert not torch.cuda.is_initialized()
vllm/config/model.py
+15
-7
@@ -20,6 +20,7 @@ from vllm.config.scheduler import RunnerType
20
from vllm.config.utils import config, getattr_iter
21
from vllm.logger import init_logger
22
from vllm.platforms import current_platform
23
+from vllm.tasks import ScoreType
24
from vllm.transformers_utils.config import (
25
ConfigFormat,
26
get_config,
@@ -1412,16 +1413,23 @@ class ModelConfig:
1413
return self._model_info.requires_raw_input_tokens
1414
1415
@property
1415
- def is_cross_encoder(self) -> bool:
1416
+ def score_type(self) -> ScoreType:
1417
+ """
1418
+ Score API handles score/rerank for:
1419
+ - "score" task (score_type: cross-encoder models)
1420
+ - "embed" task (score_type: bi-encoder models)
1421
+ - "token_embed" task (score_type: late interaction models)
1422
+ """
1423
+ # fixme: self._model_info.score_type is the score type before
1424
+ # as_seq_cls_model, which is "bi-encoder", rather than the
1425
+ # score type after as_seq_cls_model, which is "cross-encoder".
1426
+ # Therefore, the following logic is required.
1427
return (
1417
- self._model_info.supports_cross_encoding or self.convert_type == "classify"
1428
+ "cross-encoder"
1429
+ if self.convert_type == "classify"
1430
+ else self._model_info.score_type
1431
)
1432
1420
- @property
1421
- def is_late_interaction(self) -> bool:
1422
- """Check if model uses late interaction (ColBERT-style) scoring."""
1423
- return self._model_info.supports_late_interaction
1424
-
1433
@property
1434
def is_pp_supported(self) -> bool:
1435
return self._model_info.supports_pp
vllm/entrypoints/llm.py
+7
-7
@@ -1584,8 +1584,11 @@ class LLM:
1584
)
1585
1586
supported_tasks = self.supported_tasks
1587
+ score_type = self.model_config.score_type
1588
+ is_late_interaction = score_type == "late-interaction"
1589
+ is_cross_encoder = score_type == "cross-encoder"
1590
+
1591
# Late interaction models (e.g., ColBERT) use token_embed for scoring
1588
- is_late_interaction = model_config.is_late_interaction
1592
if not is_late_interaction and all(
1593
t not in supported_tasks for t in ("embed", "classify")
1594
):
@@ -1595,13 +1598,10 @@ class LLM:
1598
"`--convert embed` or `--convert classify`."
1599
)
1600
1598
- if (
1599
- model_config.is_cross_encoder
1600
- and getattr(model_config.hf_config, "num_labels", 0) != 1
1601
- ):
1601
+ if is_cross_encoder and getattr(model_config.hf_config, "num_labels", 0) != 1:
1602
raise ValueError("Score API is only enabled for num_labels == 1.")
1603
1604
- if not model_config.is_cross_encoder and chat_template is not None:
1604
+ if not is_cross_encoder and chat_template is not None:
1605
raise ValueError(
1606
"chat_template is only supported for cross-encoder models."
1607
)
@@ -1622,7 +1622,7 @@ class LLM:
1622
)
1623
encode_kwargs = tok_params.get_encode_kwargs()
1624
1625
- if model_config.is_cross_encoder:
1625
+ if is_cross_encoder:
1626
return self._cross_encoding_score(
1627
score_data_1,
1628
score_data_2,
vllm/entrypoints/pooling/__init__.py
+8
-8
@@ -37,10 +37,10 @@ def register_pooling_api_routers(
37
38
app.include_router(embed_router)
39
40
- # Score/rerank endpoints are available for:
41
- # - "score" task (cross-encoder models)
42
- # - "embed" task (bi-encoder models)
43
- # - "token_embed" task (late interaction models like ColBERT)
40
+ # Score API handles score/rerank for:
41
+ # - "score" task (score_type: cross-encoder models)
42
+ # - "embed" task (score_type: bi-encoder models)
43
+ # - "token_embed" task (score_type: late interaction models)
44
if any(t in supported_tasks for t in ("score", "embed", "token_embed")):
45
from vllm.entrypoints.pooling.score.api_router import router as score_router
46
@@ -101,10 +101,10 @@ def init_pooling_state(
101
if "classify" in supported_tasks
102
else None
103
)
104
- # ServingScores handles score/rerank for:
105
- # - "score" task (cross-encoder models)
106
- # - "embed" task (bi-encoder models)
107
- # - "token_embed" task (late interaction models like ColBERT)
104
+ # Score API handles score/rerank for:
105
+ # - "score" task (score_type: cross-encoder models)
106
+ # - "embed" task (score_type: bi-encoder models)
107
+ # - "token_embed" task (score_type: late interaction models)
108
state.serving_scores = (
109
ServingScores(
110
engine_client,
vllm/entrypoints/pooling/score/serving.py
+5
-6
@@ -69,16 +69,15 @@ class ServingScores(OpenAIServing):
69
70
self._tokenizer_executor = ThreadPoolExecutor(max_workers=1)
71
72
- self.is_cross_encoder = self.model_config.is_cross_encoder
73
- self.is_multimodal_model = self.model_config.is_multimodal_model
72
+ self.score_type = self.model_config.score_type
73
self.architecture = self.model_config.architecture
75
- self.is_late_interaction = self.model_config.is_late_interaction
74
+ self.is_multimodal_model = self.model_config.is_multimodal_model
75
77
- if self.is_cross_encoder:
76
+ if self.score_type == "cross-encoder":
77
self._score_func = self._cross_encoding_score
79
- elif self.is_late_interaction:
78
+ elif self.score_type == "late-interaction":
79
self._score_func = self._late_interaction_score
81
- else:
80
+ else: # "bi-encoder"
81
self._score_func = self._embedding_score
82
83
async def _embedding_score(
vllm/lora/model_manager.py
+5
-2
@@ -30,8 +30,11 @@ from vllm.lora.utils import (
30
replace_submodule,
31
)
32
from vllm.model_executor.layers.fused_moe import FusedMoE
33
-from vllm.model_executor.models import SupportsLoRA, supports_multimodal
34
-from vllm.model_executor.models.interfaces import is_pooling_model
33
+from vllm.model_executor.models import (
34
+ SupportsLoRA,
35
+ is_pooling_model,
36
+ supports_multimodal,
37
+)
38
from vllm.model_executor.models.module_mapping import MultiModelKeys
39
from vllm.model_executor.models.utils import PPMissingLayer
40
from vllm.multimodal import MULTIMODAL_REGISTRY
vllm/model_executor/models/colbert.py
+2
-5
@@ -18,7 +18,6 @@ Reference: https://arxiv.org/abs/2004.12832
18
"""
19
20
from collections.abc import Iterable
21
-from typing import ClassVar, Literal
21
22
import torch
23
from torch import nn
@@ -28,16 +27,16 @@ from vllm.model_executor.layers.pooler import Pooler
27
from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
28
29
from .bert import BertEmbeddingModel, BertModel
30
+from .interfaces import SupportsLateInteraction
31
from .interfaces_base import default_pooling_type
32
33
34
-class ColBERTMixin:
34
+class ColBERTMixin(nn.Module, SupportsLateInteraction):
35
"""Mixin that adds ColBERT late interaction support to any embedding model.
36
37
ColBERT (Contextualized Late Interaction over BERT) uses per-token
38
embeddings with a linear projection layer. This mixin provides:
39
40
- - ``supports_late_interaction`` class-var
40
- ColBERT linear projection initialisation / lazy creation
41
- Weight loading helpers for the projection layer
42
- A builder for the token-embedding pooler
@@ -52,8 +51,6 @@ class ColBERTMixin:
51
the ColBERT projection weight, then delegate the rest to the backbone.
52
"""
53
55
- supports_late_interaction: ClassVar[Literal[True]] = True
56
-
54
# Set during _init_colbert_components
55
colbert_dim: int | None
56
colbert_linear: nn.Linear | None
vllm/model_executor/models/colmodernvbert.py
+8
-4
@@ -9,7 +9,6 @@ Reference: https://huggingface.co/ModernVBERT/colmodernvbert-merged
9
"""
10
11
from collections.abc import Iterable, Mapping, Sequence
12
-from typing import ClassVar, Literal
12
13
import torch
14
from torch import nn
@@ -37,7 +36,11 @@ from vllm.multimodal.processing import (
36
from vllm.sequence import IntermediateTensors
37
from vllm.transformers_utils.configs.colmodernvbert import ColModernVBertConfig
38
40
-from .interfaces import MultiModalEmbeddings, SupportsMultiModal
39
+from .interfaces import (
40
+ MultiModalEmbeddings,
41
+ SupportsLateInteraction,
42
+ SupportsMultiModal,
43
+)
44
from .interfaces_base import default_pooling_type
45
from .modernbert import ModernBertEmbeddings, ModernBertLayer
46
from .siglip import SiglipVisionModel
@@ -234,7 +237,9 @@ class ColModernVBertMultiModalProcessor(
237
dummy_inputs=ColModernVBertDummyInputsBuilder,
238
)
239
@default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL")
237
-class ColModernVBertForRetrieval(nn.Module, SupportsMultiModal):
240
+class ColModernVBertForRetrieval(
241
+ nn.Module, SupportsMultiModal, SupportsLateInteraction
242
+):
243
"""ColModernVBERT multimodal late-interaction retrieval model.
244
245
Architecture:
@@ -248,7 +253,6 @@ class ColModernVBertForRetrieval(nn.Module, SupportsMultiModal):
253
"""
254
255
is_pooling_model = True
251
- supports_late_interaction: ClassVar[Literal[True]] = True
256
257
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
258
super().__init__()
vllm/model_executor/models/colqwen3.py
+2
-9
@@ -20,7 +20,6 @@ Target models:
20
"""
21
22
from collections.abc import Iterable, Mapping
23
-from typing import ClassVar, Literal
23
24
import torch
25
import torch.nn as nn
@@ -31,6 +30,7 @@ from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
30
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
31
from vllm.multimodal import MULTIMODAL_REGISTRY
32
33
+from .interfaces import SupportsLateInteraction
34
from .interfaces_base import default_pooling_type
35
from .qwen2_vl import Qwen2VLMultiModalDataParser
36
from .qwen3_vl import (
@@ -113,9 +113,7 @@ class ColQwen3ProcessingInfo(Qwen3VLProcessingInfo):
113
info=ColQwen3ProcessingInfo,
114
dummy_inputs=Qwen3VLDummyInputsBuilder,
115
)
116
-class ColQwen3Model(
117
- Qwen3VLForConditionalGeneration,
118
-):
116
+class ColQwen3Model(Qwen3VLForConditionalGeneration, SupportsLateInteraction):
117
"""ColQwen3 late interaction model for multi-modal retrieval/reranking.
118
119
This model extends Qwen3VLForConditionalGeneration with a ColBERT-style
@@ -132,16 +130,11 @@ class ColQwen3Model(
130
131
Attributes:
132
custom_text_proj: Linear projection from hidden_size to embed_dim
135
- supports_late_interaction: Flag indicating this model uses late
136
- interaction scoring
133
"""
134
135
# Mark this as a pooling model so vLLM routes to pooler path
136
is_pooling_model = True
137
142
- # Mark this model as supporting late interaction scoring
143
- supports_late_interaction: ClassVar[Literal[True]] = True
144
-
138
# Override hf_to_vllm_mapper to handle ColQwen3 weight naming.
139
# NOTE: WeightsMapper applies ALL matching prefix rules sequentially
140
# (no early exit), so more-specific prefixes must come first.
vllm/model_executor/models/interfaces.py
+4
-47
@@ -34,10 +34,11 @@ from vllm.inputs.data import PromptType
34
from vllm.logger import init_logger
35
from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
36
from vllm.model_executor.layers.quantization import QuantizationConfig
37
+from vllm.tasks import ScoreType
38
from vllm.utils.collection_utils import common_prefix
39
from vllm.utils.func_utils import supports_kw
40
40
-from .interfaces_base import VllmModel, is_pooling_model
41
+from .interfaces_base import VllmModel
42
43
if TYPE_CHECKING:
44
from vllm.config import VllmConfig
@@ -969,29 +970,7 @@ def supports_mamba_prefix_caching(
970
class SupportsCrossEncoding(Protocol):
971
"""The interface required for all models that support cross encoding."""
972
972
- supports_cross_encoding: ClassVar[Literal[True]] = True
973
-
974
-
975
-@overload
976
-def supports_cross_encoding(
977
- model: type[object],
978
-) -> TypeIs[type[SupportsCrossEncoding]]: ...
979
-
980
-
981
-@overload
982
-def supports_cross_encoding(model: object) -> TypeIs[SupportsCrossEncoding]: ...
983
-
984
-
985
-def _supports_cross_encoding(
986
- model: type[object] | object,
987
-) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:
988
- return getattr(model, "supports_cross_encoding", False)
989
-
990
-
991
-def supports_cross_encoding(
992
- model: type[object] | object,
993
-) -> TypeIs[type[SupportsCrossEncoding]] | TypeIs[SupportsCrossEncoding]:
994
- return is_pooling_model(model) and _supports_cross_encoding(model)
973
+ score_type: ClassVar[ScoreType] = "cross-encoder"
974
975
976
@runtime_checkable
@@ -1003,29 +982,7 @@ class SupportsLateInteraction(Protocol):
982
MaxSim (max over document tokens, sum over query tokens).
983
"""
984
1006
- supports_late_interaction: ClassVar[Literal[True]] = True
1007
-
1008
-
1009
-@overload
1010
-def supports_late_interaction(
1011
- model: type[object],
1012
-) -> TypeIs[type[SupportsLateInteraction]]: ...
1013
-
1014
-
1015
-@overload
1016
-def supports_late_interaction(model: object) -> TypeIs[SupportsLateInteraction]: ...
1017
-
1018
-
1019
-def _supports_late_interaction(
1020
- model: type[object] | object,
1021
-) -> TypeIs[type[SupportsLateInteraction]] | TypeIs[SupportsLateInteraction]:
1022
- return getattr(model, "supports_late_interaction", False)
1023
-
1024
-
1025
-def supports_late_interaction(
1026
- model: type[object] | object,
1027
-) -> TypeIs[type[SupportsLateInteraction]] | TypeIs[SupportsLateInteraction]:
1028
- return is_pooling_model(model) and _supports_late_interaction(model)
985
+ score_type: ClassVar[ScoreType] = "late-interaction"
986
987
988
class SupportsQuant:
vllm/model_executor/models/interfaces_base.py
+31
@@ -15,6 +15,7 @@ import torch.nn as nn
15
from typing_extensions import TypeIs, TypeVar
16
17
from vllm.logger import init_logger
18
+from vllm.tasks import ScoreType
19
from vllm.utils.func_utils import supports_kw
20
21
if TYPE_CHECKING:
@@ -187,6 +188,26 @@ class VllmModelForPooling(VllmModel[T_co], Protocol[T_co]):
188
decorator to conveniently set this field.
189
"""
190
191
+ score_type: ClassVar[ScoreType] = "bi-encoder"
192
+ """
193
+ Indicates the
194
+ [vllm.config.model.ModelConfig.score_type][]
195
+ to use by default.
196
+
197
+ Score API handles score/rerank for:
198
+ - "score" task (score_type: cross-encoder models)
199
+ - "embed" task (score_type: bi-encoder models)
200
+ - "token_embed" task (score_type: late interaction models)
201
+
202
+ score_type defaults to bi-encoder, then the Score API uses the "embed" task.
203
+ If you set score_type to cross-encoder via
204
+ [vllm.model_executor.models.interfaces.SupportsCrossEncoding][],
205
+ then the Score API uses the "score" task.
206
+ If you set score_type to late-interaction via
207
+ [vllm.model_executor.models.interfaces.SupportsLateInteraction][],
208
+ then the Score API uses the "token_embed" task.
209
+ """
210
+
211
pooler: Pooler
212
"""The pooler is only called on TP rank 0."""
213
@@ -250,3 +271,13 @@ def attn_type(attn_type: AttnTypeStr):
271
272
def get_attn_type(model: type[object] | object) -> AttnTypeStr:
273
return getattr(model, "attn_type", "decoder")
274
+
275
+
276
+def get_score_type(model: type[object] | object) -> ScoreType:
277
+ score_types = set()
278
+ for m in model.__mro__:
279
+ score_type = getattr(m, "score_type", "bi-encoder")
280
+ if score_type != "bi-encoder":
281
+ score_types.add(score_type)
282
+ assert len(score_types) < 2
283
+ return "bi-encoder" if not score_types else list(score_types)[0]
vllm/model_executor/models/registry.py
+44
-39
@@ -30,6 +30,7 @@ from vllm.config import (
30
)
31
from vllm.logger import init_logger
32
from vllm.logging_utils import logtime
33
+from vllm.tasks import ScoreType
34
from vllm.transformers_utils.dynamic_module import try_get_class_from_dynamic_module
35
from vllm.utils.hashing import safe_hash
36
@@ -48,8 +49,6 @@ from .interfaces import (
49
is_attention_free,
50
is_hybrid,
51
requires_raw_input_tokens,
51
- supports_cross_encoding,
52
- supports_late_interaction,
52
supports_mamba_prefix_caching,
53
supports_multimodal,
54
supports_multimodal_encoder_tp_data,
@@ -61,6 +60,7 @@ from .interfaces_base import (
60
get_attn_type,
61
get_default_seq_pooling_type,
62
get_default_tok_pooling_type,
63
+ get_score_type,
64
is_pooling_model,
65
is_text_generation_model,
66
)
@@ -214,19 +214,14 @@ _EMBEDDING_MODELS = {
214
# [Text-only]
215
"BertModel": ("bert", "BertEmbeddingModel"),
216
"BertSpladeSparseEmbeddingModel": ("bert", "BertSpladeSparseEmbeddingModel"),
217
- "HF_ColBERT": ("colbert", "ColBERTModel"),
218
- "ColBERTModernBertModel": ("colbert", "ColBERTModernBertModel"),
219
- "ColBERTJinaRobertaModel": ("colbert", "ColBERTJinaRobertaModel"),
217
+ "BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"),
218
"DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
219
"Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
220
"Gemma3TextModel": ("gemma3", "Gemma3Model"),
221
"GlmForCausalLM": ("glm", "GlmForCausalLM"),
224
- "GPT2ForSequenceClassification": ("gpt2", "GPT2ForSequenceClassification"),
222
"GritLM": ("gritlm", "GritLM"),
223
"GteModel": ("bert_with_rope", "SnowflakeGteNewModel"),
224
"GteNewModel": ("bert_with_rope", "GteNewModel"),
228
- "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
229
- "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"), # noqa: E501
225
"LlamaBidirectionalModel": ("llama", "LlamaBidirectionalModel"),
226
"LlamaModel": ("llama", "LlamaForCausalLM"),
227
**{
@@ -241,8 +236,6 @@ _EMBEDDING_MODELS = {
236
"Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
237
"Qwen2Model": ("qwen2", "Qwen2ForCausalLM"),
238
"Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
244
- "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
245
- "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"),
239
"RobertaForMaskedLM": ("roberta", "RobertaEmbeddingModel"),
240
"RobertaModel": ("roberta", "RobertaEmbeddingModel"),
241
"TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
@@ -252,19 +245,14 @@ _EMBEDDING_MODELS = {
245
"VoyageQwen3BidirectionalEmbedModel",
246
),
247
"XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
255
- "BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"),
248
# [Multimodal]
249
"CLIPModel": ("clip", "CLIPEmbeddingModel"),
258
- "ColModernVBertForRetrieval": ("colmodernvbert", "ColModernVBertForRetrieval"),
250
"LlavaNextForConditionalGeneration": (
251
"llava_next",
252
"LlavaNextForConditionalGeneration",
253
),
254
"Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
255
"Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"), # noqa: E501
265
- "ColQwen3": ("colqwen3", "ColQwen3Model"),
266
- "OpsColQwen3Model": ("colqwen3", "ColQwen3Model"),
267
- "Qwen3VLNemotronEmbedModel": ("colqwen3", "ColQwen3Model"),
256
"SiglipModel": ("siglip", "SiglipEmbeddingModel"),
257
"LlamaNemotronVLModel": (
258
"nemotron_vl",
@@ -277,35 +265,59 @@ _EMBEDDING_MODELS = {
265
"Terratorch": ("terratorch", "Terratorch"),
266
}
267
280
-_CROSS_ENCODER_MODELS = {
281
- "BertForSequenceClassification": ("bert", "BertForSequenceClassification"),
268
+_LATE_INTERACTION_MODELS = {
269
+ # [Text-only]
270
+ "HF_ColBERT": ("colbert", "ColBERTModel"),
271
+ "ColBERTModernBertModel": ("colbert", "ColBERTModernBertModel"),
272
+ "ColBERTJinaRobertaModel": ("colbert", "ColBERTJinaRobertaModel"),
273
+ # [Multimodal]
274
+ "ColModernVBertForRetrieval": ("colmodernvbert", "ColModernVBertForRetrieval"),
275
+ "ColQwen3": ("colqwen3", "ColQwen3Model"),
276
+ "OpsColQwen3Model": ("colqwen3", "ColQwen3Model"),
277
+ "Qwen3VLNemotronEmbedModel": ("colqwen3", "ColQwen3Model"),
278
+}
279
+
280
+_REWARD_MODELS = {
281
+ "InternLM2ForRewardModel": ("internlm2", "InternLM2ForRewardModel"),
282
+ "Qwen2ForRewardModel": ("qwen2_rm", "Qwen2ForRewardModel"),
283
+ "Qwen2ForProcessRewardModel": ("qwen2_rm", "Qwen2ForProcessRewardModel"),
284
+}
285
+
286
+_TOKEN_CLASSIFICATION_MODELS = {
287
"BertForTokenClassification": ("bert", "BertForTokenClassification"),
288
+ "ModernBertForTokenClassification": (
289
+ "modernbert",
290
+ "ModernBertForTokenClassification",
291
+ ),
292
+}
293
+
294
+_SEQUENCE_CLASSIFICATION_MODELS = {
295
+ "BertForSequenceClassification": ("bert", "BertForSequenceClassification"),
296
+ "GPT2ForSequenceClassification": ("gpt2", "GPT2ForSequenceClassification"),
297
"GteNewForSequenceClassification": (
298
"bert_with_rope",
299
"GteNewForSequenceClassification",
300
),
287
- "JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"),
301
+ "JambaForSequenceClassification": ("jamba", "JambaForSequenceClassification"), # noqa: E501
302
"LlamaBidirectionalForSequenceClassification": (
303
"llama",
304
"LlamaBidirectionalForSequenceClassification",
305
),
292
- "LlamaNemotronVLForSequenceClassification": (
293
- "nemotron_vl",
294
- "LlamaNemotronVLForSequenceClassification",
295
- ),
306
"ModernBertForSequenceClassification": (
307
"modernbert",
308
"ModernBertForSequenceClassification",
309
),
300
- "ModernBertForTokenClassification": (
301
- "modernbert",
302
- "ModernBertForTokenClassification",
303
- ),
310
"RobertaForSequenceClassification": ("roberta", "RobertaForSequenceClassification"),
311
"XLMRobertaForSequenceClassification": (
312
"roberta",
313
"RobertaForSequenceClassification",
314
),
315
+ # [Multimodal]
316
+ "JinaVLForRanking": ("jina_vl", "JinaVLForSequenceClassification"),
317
+ "LlamaNemotronVLForSequenceClassification": (
318
+ "nemotron_vl",
319
+ "LlamaNemotronVLForSequenceClassification",
320
+ ),
321
}
322
323
_MULTIMODAL_MODELS = {
@@ -606,7 +618,10 @@ _TRANSFORMERS_BACKEND_MODELS = {
618
_VLLM_MODELS = {
619
**_TEXT_GENERATION_MODELS,
620
**_EMBEDDING_MODELS,
609
- **_CROSS_ENCODER_MODELS,
621
+ **_LATE_INTERACTION_MODELS,
622
+ **_REWARD_MODELS,
623
+ **_TOKEN_CLASSIFICATION_MODELS,
624
+ **_SEQUENCE_CLASSIFICATION_MODELS,
625
**_MULTIMODAL_MODELS,
626
**_SPECULATIVE_DECODING_MODELS,
627
**_TRANSFORMERS_SUPPORTED_MODELS,
@@ -643,8 +658,7 @@ class _ModelInfo:
658
attn_type: AttnTypeStr
659
default_seq_pooling_type: SequencePoolingType
660
default_tok_pooling_type: TokenPoolingType
646
- supports_cross_encoding: bool
647
- supports_late_interaction: bool
661
+ score_type: ScoreType
662
supports_multimodal: bool
663
supports_multimodal_raw_input_only: bool
664
requires_raw_input_tokens: bool
@@ -667,8 +681,7 @@ class _ModelInfo:
681
default_seq_pooling_type=get_default_seq_pooling_type(model),
682
default_tok_pooling_type=get_default_tok_pooling_type(model),
683
attn_type=get_attn_type(model),
670
- supports_cross_encoding=supports_cross_encoding(model),
671
- supports_late_interaction=supports_late_interaction(model),
684
+ score_type=get_score_type(model),
685
supports_multimodal=supports_multimodal(model),
686
supports_multimodal_raw_input_only=supports_multimodal_raw_input_only(
687
model
@@ -1166,14 +1179,6 @@ class _ModelRegistry:
1179
model_cls, _ = self.inspect_model_cls(architectures, model_config)
1180
return model_cls.is_pooling_model
1181
1169
- def is_cross_encoder_model(
1170
- self,
1171
- architectures: str | list[str],
1172
- model_config: ModelConfig,
1173
- ) -> bool:
1174
- model_cls, _ = self.inspect_model_cls(architectures, model_config)
1175
- return model_cls.supports_cross_encoding
1176
-
1182
def is_multimodal_model(
1183
self,
1184
architectures: str | list[str],
vllm/tasks.py
+6
@@ -10,6 +10,12 @@ PoolingTask = Literal[
10
]
11
POOLING_TASKS: tuple[PoolingTask, ...] = get_args(PoolingTask)
12
13
+# Score API handles score/rerank for:
14
+# - "score" task (score_type: cross-encoder models)
15
+# - "embed" task (score_type: bi-encoder models)
16
+# - "token_embed" task (score_type: late interaction models)
17
+ScoreType = Literal["bi-encoder", "cross-encoder", "late-interaction"]
18
+
19
FrontendTask = Literal["render"]
20
FRONTEND_TASKS: tuple[FrontendTask, ...] = get_args(FrontendTask)
21