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 7, 22:45 UTC
[New Model] Add ColModernVBERT (#34558)
Signed-off-by: Athrael Soju <athrael.soju@gmail.com> Signed-off-by: athrael-soju <athrael-soju@users.noreply.github.com>
Athrael Soju committed
Feb 22, 2026 at 04:23 UTC
970861ac0cfc93d8ebdeb2c0f5d664289eafb51c
9 files changed
+784
docs/models/supported_models.md
+1
@@ -821,6 +821,7 @@ The following table lists those that are tested in vLLM.
821
| Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
822
|--------------|--------|--------|-------------------|----------------------|---------------------------|
823
| `CLIPModel` | CLIP | T / I | `openai/clip-vit-base-patch32`, `openai/clip-vit-large-patch14`, etc. | | |
824
+| `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | |
825
| `LlavaNextForConditionalGeneration`<sup>C</sup> | LLaVA-NeXT-based | T / I | `royokong/e5-v` | | ✅︎ |
826
| `Phi3VForCausalLM`<sup>C</sup> | Phi-3-Vision-based | T + I | `TIGER-Lab/VLM2Vec-Full` | | ✅︎ |
827
| `Qwen3VLForConditionalGeneration`<sup>C</sup> | Qwen3-VL | T + I + V | `Qwen/Qwen3-VL-Embedding-2B`, etc. | ✅︎ | ✅︎ |
examples/pooling/score/colmodernvbert_rerank_online.py
new
+166
@@ -0,0 +1,166 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Example of using ColModernVBERT late interaction model for reranking.
5
+
6
+ColModernVBERT is a multi-modal ColBERT-style model combining a SigLIP
7
+vision encoder with a ModernBERT text encoder. It produces per-token
8
+embeddings and uses MaxSim scoring for retrieval and reranking.
9
+Supports both text and image inputs.
10
+
11
+Start the server with:
12
+ vllm serve ModernVBERT/colmodernvbert-merged --max-model-len 8192
13
+
14
+Then run this script:
15
+ python colmodernvbert_rerank_online.py
16
+"""
17
+
18
+import requests
19
+
20
+MODEL = "ModernVBERT/colmodernvbert-merged"
21
+BASE_URL = "http://127.0.0.1:8000"
22
+
23
+headers = {"accept": "application/json", "Content-Type": "application/json"}
24
+
25
+IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/PNG_transparency_demonstration_1.png/300px-PNG_transparency_demonstration_1.png" # noqa: E501
26
+
27
+
28
+def rerank_text():
29
+ """Text-only reranking via /rerank endpoint."""
30
+ print("=" * 60)
31
+ print("1. Text reranking (/rerank)")
32
+ print("=" * 60)
33
+
34
+ data = {
35
+ "model": MODEL,
36
+ "query": "What is machine learning?",
37
+ "documents": [
38
+ "Machine learning is a subset of artificial intelligence.",
39
+ "Python is a programming language.",
40
+ "Deep learning uses neural networks for complex tasks.",
41
+ "The weather today is sunny.",
42
+ ],
43
+ }
44
+
45
+ response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
46
+
47
+ if response.status_code == 200:
48
+ result = response.json()
49
+ print("\n Ranked documents (most relevant first):")
50
+ for item in result["results"]:
51
+ doc_idx = item["index"]
52
+ score = item["relevance_score"]
53
+ print(f" [{score:.4f}] {data['documents'][doc_idx]}")
54
+ else:
55
+ print(f" Request failed: {response.status_code}")
56
+ print(f" {response.text[:300]}")
57
+
58
+
59
+def score_text():
60
+ """Text-only scoring via /score endpoint."""
61
+ print()
62
+ print("=" * 60)
63
+ print("2. Text scoring (/score)")
64
+ print("=" * 60)
65
+
66
+ query = "What is the capital of France?"
67
+ documents = [
68
+ "The capital of France is Paris.",
69
+ "Berlin is the capital of Germany.",
70
+ "Python is a programming language.",
71
+ ]
72
+
73
+ data = {
74
+ "model": MODEL,
75
+ "text_1": query,
76
+ "text_2": documents,
77
+ }
78
+
79
+ response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
80
+
81
+ if response.status_code == 200:
82
+ result = response.json()
83
+ print(f"\n Query: {query}\n")
84
+ for item in result["data"]:
85
+ idx = item["index"]
86
+ score = item["score"]
87
+ print(f" Doc {idx} (score={score:.4f}): {documents[idx]}")
88
+ else:
89
+ print(f" Request failed: {response.status_code}")
90
+ print(f" {response.text[:300]}")
91
+
92
+
93
+def score_text_top_n():
94
+ """Text reranking with top_n filtering via /rerank endpoint."""
95
+ print()
96
+ print("=" * 60)
97
+ print("3. Text reranking with top_n=2 (/rerank)")
98
+ print("=" * 60)
99
+
100
+ data = {
101
+ "model": MODEL,
102
+ "query": "What is the capital of France?",
103
+ "documents": [
104
+ "The capital of France is Paris.",
105
+ "Berlin is the capital of Germany.",
106
+ "Python is a programming language.",
107
+ "The Eiffel Tower is in Paris.",
108
+ ],
109
+ "top_n": 2,
110
+ }
111
+
112
+ response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
113
+
114
+ if response.status_code == 200:
115
+ result = response.json()
116
+ print(f"\n Top {data['top_n']} results:")
117
+ for item in result["results"]:
118
+ doc_idx = item["index"]
119
+ score = item["relevance_score"]
120
+ print(f" [{score:.4f}] {data['documents'][doc_idx]}")
121
+ else:
122
+ print(f" Request failed: {response.status_code}")
123
+ print(f" {response.text[:300]}")
124
+
125
+
126
+def rerank_multimodal():
127
+ """Multimodal reranking with text and image documents via /rerank."""
128
+ print()
129
+ print("=" * 60)
130
+ print("4. Multimodal reranking: text query vs image document (/rerank)")
131
+ print("=" * 60)
132
+
133
+ data = {
134
+ "model": MODEL,
135
+ "query": "A colorful logo with transparency",
136
+ "documents": [
137
+ {"content": [{"type": "image_url", "image_url": {"url": IMAGE_URL}}]},
138
+ "Python is a programming language.",
139
+ "The weather today is sunny.",
140
+ ],
141
+ }
142
+
143
+ response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
144
+
145
+ if response.status_code == 200:
146
+ result = response.json()
147
+ print("\n Ranked documents (most relevant first):")
148
+ labels = ["[image]", "Python doc", "Weather doc"]
149
+ for item in result["results"]:
150
+ doc_idx = item["index"]
151
+ score = item["relevance_score"]
152
+ print(f" [{score:.4f}] {labels[doc_idx]}")
153
+ else:
154
+ print(f" Request failed: {response.status_code}")
155
+ print(f" {response.text[:300]}")
156
+
157
+
158
+def main():
159
+ rerank_text()
160
+ score_text()
161
+ score_text_top_n()
162
+ rerank_multimodal()
163
+
164
+
165
+if __name__ == "__main__":
166
+ main()
tests/models/multimodal/pooling/test_colmodernvbert.py
new
+115
@@ -0,0 +1,115 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for ColModernVBERT multimodal late-interaction model.
4
+
5
+ColModernVBERT combines SigLIP vision encoder + ModernBERT text encoder
6
+with a pixel shuffle connector and ColBERT-style 128-dim per-token
7
+embeddings for visual document retrieval.
8
+"""
9
+
10
+import pytest
11
+import torch
12
+
13
+from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
14
+
15
+MODEL_NAME = "ModernVBERT/colmodernvbert-merged"
16
+COLBERT_DIM = 128
17
+DTYPE = "half"
18
+
19
+
20
+# -----------------------------------------------------------------------
21
+# Text-only tests
22
+# -----------------------------------------------------------------------
23
+
24
+
25
+def test_colmodernvbert_text_token_embed(vllm_runner):
26
+ """Text query produces per-token embeddings with shape (seq_len, 128)."""
27
+ with vllm_runner(
28
+ MODEL_NAME,
29
+ runner="pooling",
30
+ dtype=DTYPE,
31
+ enforce_eager=True,
32
+ ) as vllm_model:
33
+ outputs = vllm_model.token_embed(["What is machine learning?"])
34
+
35
+ assert len(outputs) == 1
36
+ emb = torch.tensor(outputs[0])
37
+ assert emb.dim() == 2
38
+ assert emb.shape[1] == COLBERT_DIM
39
+ assert emb.shape[0] > 1
40
+
41
+
42
+def test_colmodernvbert_text_relevance_ordering(vllm_runner):
43
+ """Relevant documents score higher than irrelevant ones."""
44
+ query = "What is machine learning?"
45
+ documents = [
46
+ "Machine learning is a subset of artificial intelligence.",
47
+ "The weather in Paris is mild in spring.",
48
+ ]
49
+
50
+ with vllm_runner(
51
+ MODEL_NAME,
52
+ runner="pooling",
53
+ dtype=DTYPE,
54
+ enforce_eager=True,
55
+ ) as vllm_model:
56
+ scores = vllm_model.score(query, documents)
57
+
58
+ assert len(scores) == 2
59
+ assert scores[0] > scores[1], "ML doc should score higher than weather doc"
60
+
61
+
62
+def test_colmodernvbert_text_late_interaction(vllm_runner):
63
+ """MaxSim scoring via vLLM matches manual computation."""
64
+ query = "What is the capital of France?"
65
+ doc = "The capital of France is Paris."
66
+
67
+ with vllm_runner(
68
+ MODEL_NAME,
69
+ runner="pooling",
70
+ dtype=DTYPE,
71
+ enforce_eager=True,
72
+ ) as vllm_model:
73
+ q_out = vllm_model.token_embed([query])
74
+ d_out = vllm_model.token_embed([doc])
75
+
76
+ q_emb = torch.tensor(q_out[0])
77
+ d_emb = torch.tensor(d_out[0])
78
+ manual_score = compute_maxsim_score(q_emb, d_emb).item()
79
+
80
+ vllm_scores = vllm_model.score(query, doc)
81
+
82
+ assert len(vllm_scores) == 1
83
+ assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
84
+
85
+
86
+# -----------------------------------------------------------------------
87
+# Image tests
88
+# -----------------------------------------------------------------------
89
+
90
+
91
+def test_colmodernvbert_image_token_embed(vllm_runner, image_assets):
92
+ """Image input produces per-token embeddings including vision tokens."""
93
+ with vllm_runner(
94
+ MODEL_NAME,
95
+ runner="pooling",
96
+ dtype=DTYPE,
97
+ enforce_eager=True,
98
+ ) as vllm_model:
99
+ image = image_assets[0].pil_image
100
+ inputs = vllm_model.get_inputs(
101
+ [""],
102
+ images=[image],
103
+ )
104
+ req_outputs = vllm_model.llm.encode(
105
+ inputs,
106
+ pooling_task="token_embed",
107
+ )
108
+ outputs = [req_output.outputs.data for req_output in req_outputs]
109
+
110
+ assert len(outputs) == 1
111
+ emb = torch.tensor(outputs[0])
112
+ assert emb.dim() == 2
113
+ assert emb.shape[1] == COLBERT_DIM
114
+ # Should have at least the image tokens (64 after pixel shuffle)
115
+ assert emb.shape[0] >= 64
tests/models/registry.py
+3
@@ -592,6 +592,9 @@ _EMBEDDING_EXAMPLE_MODELS = {
592
),
593
# [Multimodal]
594
"CLIPModel": _HfExamplesInfo("openai/clip-vit-base-patch32"),
595
+ "ColModernVBertForRetrieval": _HfExamplesInfo(
596
+ "ModernVBERT/colmodernvbert-merged",
597
+ ),
598
"LlavaNextForConditionalGeneration": _HfExamplesInfo("royokong/e5-v"),
599
"Phi3VForCausalLM": _HfExamplesInfo(
600
"TIGER-Lab/VLM2Vec-Full", trust_remote_code=True
vllm/model_executor/models/colmodernvbert.py
new
+430
@@ -0,0 +1,430 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ColModernVBERT: multimodal late-interaction retrieval model.
4
+
5
+Combines SigLIP vision encoder + ModernBERT text encoder with a pixel
6
+shuffle connector and ColBERT-style 128-dim per-token embeddings.
7
+
8
+Reference: https://huggingface.co/ModernVBERT/colmodernvbert-merged
9
+"""
10
+
11
+from collections.abc import Iterable, Mapping, Sequence
12
+from typing import ClassVar, Literal
13
+
14
+import torch
15
+from torch import nn
16
+from transformers import BatchFeature
17
+
18
+from vllm.config import VllmConfig
19
+from vllm.config.multimodal import BaseDummyOptions
20
+from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
21
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
22
+from vllm.multimodal import MULTIMODAL_REGISTRY
23
+from vllm.multimodal.inputs import (
24
+ MultiModalDataDict,
25
+ MultiModalFieldConfig,
26
+ MultiModalKwargsItems,
27
+)
28
+from vllm.multimodal.parse import ImageSize, MultiModalDataItems
29
+from vllm.multimodal.processing import (
30
+ BaseDummyInputsBuilder,
31
+ BaseMultiModalProcessor,
32
+ BaseProcessingInfo,
33
+ PromptIndexTargets,
34
+ PromptReplacement,
35
+ PromptUpdate,
36
+)
37
+from vllm.sequence import IntermediateTensors
38
+from vllm.transformers_utils.configs.colmodernvbert import ColModernVBertConfig
39
+
40
+from .interfaces import MultiModalEmbeddings, SupportsMultiModal
41
+from .interfaces_base import default_pooling_type
42
+from .modernbert import ModernBertEmbeddings, ModernBertLayer
43
+from .siglip import SiglipVisionModel
44
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
45
+
46
+# ---------------------------------------------------------------------------
47
+# Connector: pixel shuffle + simple linear projection
48
+# ---------------------------------------------------------------------------
49
+
50
+
51
+class ColModernVBertConnector(nn.Module):
52
+ """Pixel shuffle spatial reduction followed by a linear projection.
53
+
54
+ Reduces the vision encoder's token count by ``factor^2`` via pixel-shuffle
55
+ spatial rearrangement, then projects the concatenated channels to the text
56
+ encoder's hidden size with a single bias-free linear layer.
57
+ """
58
+
59
+ def __init__(self, config: ColModernVBertConfig):
60
+ super().__init__()
61
+ self.pixel_shuffle_factor = config.pixel_shuffle_factor
62
+ vision_hidden_size = config.vision_config.hidden_size
63
+ input_size = vision_hidden_size * (self.pixel_shuffle_factor**2)
64
+ output_size = config.hidden_size
65
+ self.proj = nn.Linear(input_size, output_size, bias=False)
66
+
67
+ def pixel_shuffle(self, features: torch.Tensor) -> torch.Tensor:
68
+ """Spatial rearrangement that reduces seq length by factor^2."""
69
+ batch_size, seq_length, hidden_size = features.shape
70
+ height = width = int(seq_length**0.5)
71
+ factor = self.pixel_shuffle_factor
72
+
73
+ # Reshape to (B, H, W, C)
74
+ features = features.view(batch_size, height, width, hidden_size)
75
+
76
+ # Reshape to (B, H/f, f, W/f, f, C)
77
+ features = features.view(
78
+ batch_size, height // factor, factor, width // factor, factor, hidden_size
79
+ )
80
+
81
+ # Permute to (B, H/f, W/f, f, f, C)
82
+ features = features.permute(0, 1, 3, 2, 4, 5)
83
+
84
+ # Reshape to (B, H/f, W/f, C * f^2)
85
+ new_hidden_size = hidden_size * (factor**2)
86
+ features = features.reshape(
87
+ batch_size, height // factor, width // factor, new_hidden_size
88
+ )
89
+
90
+ return features
91
+
92
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
93
+ features = self.pixel_shuffle(features)
94
+ batch_size = features.shape[0]
95
+ features = features.reshape(batch_size, -1, features.shape[-1])
96
+ return self.proj(features)
97
+
98
+
99
+# ---------------------------------------------------------------------------
100
+# Multimodal processing
101
+# ---------------------------------------------------------------------------
102
+
103
+
104
+class ColModernVBertProcessingInfo(BaseProcessingInfo):
105
+ def get_hf_config(self) -> ColModernVBertConfig:
106
+ return self.ctx.get_hf_config(ColModernVBertConfig)
107
+
108
+ def get_supported_mm_limits(self) -> Mapping[str, int | None]:
109
+ return {"image": None}
110
+
111
+ def get_image_size_with_most_features(self) -> ImageSize:
112
+ config = self.get_hf_config()
113
+ size = config.vision_config.image_size
114
+ return ImageSize(width=size, height=size)
115
+
116
+ def get_num_image_tokens(
117
+ self,
118
+ *,
119
+ image_width: int,
120
+ image_height: int,
121
+ ) -> int:
122
+ return self.get_hf_config().image_seq_len
123
+
124
+
125
+class ColModernVBertDummyInputsBuilder(
126
+ BaseDummyInputsBuilder[ColModernVBertProcessingInfo],
127
+):
128
+ def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
129
+ return ""
130
+
131
+ def get_dummy_mm_data(
132
+ self,
133
+ seq_len: int,
134
+ mm_counts: Mapping[str, int],
135
+ mm_options: Mapping[str, BaseDummyOptions] | None = None,
136
+ mm_processor_kwargs: Mapping[str, object] | None = None,
137
+ ) -> MultiModalDataDict:
138
+ num_images = mm_counts.get("image", 0)
139
+ target_width, target_height = self.info.get_image_size_with_most_features()
140
+ image_overrides = mm_options.get("image") if mm_options else None
141
+ return {
142
+ "image": self._get_dummy_images(
143
+ width=target_width,
144
+ height=target_height,
145
+ num_images=num_images,
146
+ overrides=image_overrides,
147
+ )
148
+ }
149
+
150
+
151
+class ColModernVBertMultiModalProcessor(
152
+ BaseMultiModalProcessor[ColModernVBertProcessingInfo],
153
+):
154
+ def _call_hf_processor(
155
+ self,
156
+ prompt: str,
157
+ mm_data: Mapping[str, object],
158
+ mm_kwargs: Mapping[str, object],
159
+ tok_kwargs: Mapping[str, object],
160
+ ) -> BatchFeature:
161
+ tokenizer = self.info.get_tokenizer()
162
+ text_encoding = tokenizer(
163
+ prompt,
164
+ return_tensors="pt",
165
+ **tok_kwargs,
166
+ )
167
+ result = BatchFeature(data=dict(text_encoding))
168
+
169
+ images = mm_data.get("images")
170
+ if images:
171
+ from transformers import Idefics3ImageProcessor
172
+
173
+ image_processor = Idefics3ImageProcessor.from_pretrained(
174
+ self.info.ctx.model_config.model,
175
+ revision=self.info.ctx.model_config.revision,
176
+ )
177
+ image_outputs = image_processor(
178
+ images=images,
179
+ do_image_splitting=False,
180
+ return_tensors="pt",
181
+ )
182
+ result.update(image_outputs)
183
+
184
+ return result
185
+
186
+ def _hf_processor_applies_updates(
187
+ self,
188
+ prompt_text: str,
189
+ mm_items: MultiModalDataItems,
190
+ hf_processor_mm_kwargs: Mapping[str, object],
191
+ tokenization_kwargs: Mapping[str, object],
192
+ ) -> bool:
193
+ return False
194
+
195
+ def _get_mm_fields_config(
196
+ self,
197
+ hf_inputs: BatchFeature,
198
+ hf_processor_mm_kwargs: Mapping[str, object],
199
+ ) -> Mapping[str, MultiModalFieldConfig]:
200
+ return dict(
201
+ pixel_values=MultiModalFieldConfig.batched("image"),
202
+ )
203
+
204
+ def _get_prompt_updates(
205
+ self,
206
+ mm_items: MultiModalDataItems,
207
+ hf_processor_mm_kwargs: Mapping[str, object],
208
+ out_mm_kwargs: MultiModalKwargsItems,
209
+ ) -> Sequence[PromptUpdate]:
210
+ config = self.info.get_hf_config()
211
+ image_token_id = config.image_token_id
212
+ num_tokens = config.image_seq_len
213
+
214
+ def get_replacement(item_idx: int):
215
+ return [image_token_id] * num_tokens
216
+
217
+ return [
218
+ PromptReplacement(
219
+ modality="image",
220
+ target=PromptIndexTargets.start(),
221
+ replacement=get_replacement,
222
+ ),
223
+ ]
224
+
225
+
226
+# ---------------------------------------------------------------------------
227
+# Model
228
+# ---------------------------------------------------------------------------
229
+
230
+
231
+@MULTIMODAL_REGISTRY.register_processor(
232
+ ColModernVBertMultiModalProcessor,
233
+ info=ColModernVBertProcessingInfo,
234
+ dummy_inputs=ColModernVBertDummyInputsBuilder,
235
+)
236
+@default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL")
237
+class ColModernVBertForRetrieval(nn.Module, SupportsMultiModal):
238
+ """ColModernVBERT multimodal late-interaction retrieval model.
239
+
240
+ Architecture:
241
+ Image -> SiglipVisionModel -> ColModernVBertConnector
242
+ ↓
243
+ Text -> ModernBertEmbeddings → [merge] → ModernBertLayers → norm
244
+ ↓
245
+ custom_text_proj → L2 norm
246
+ ↓
247
+ per-token 128-d embeddings
248
+ """
249
+
250
+ is_pooling_model = True
251
+ supports_late_interaction: ClassVar[Literal[True]] = True
252
+
253
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
254
+ super().__init__()
255
+ config: ColModernVBertConfig = vllm_config.model_config.hf_config
256
+ self.config = config
257
+ text_config = config.text_config
258
+ quant_config = vllm_config.quant_config
259
+
260
+ # --- Vision encoder (reuses SiglipVisionModel from siglip.py) ---
261
+ self.vision_model = SiglipVisionModel(
262
+ config.vision_config,
263
+ quant_config,
264
+ prefix=maybe_prefix(prefix, "vision_model"),
265
+ )
266
+
267
+ # --- Connector (pixel shuffle + linear projection) ---
268
+ self.connector = ColModernVBertConnector(config)
269
+
270
+ # --- Text encoder (built from ModernBERT components directly) ---
271
+ # We build the components individually rather than wrapping
272
+ # ``ModernBertModel`` because ``ModernBertEncoderLayer`` reads
273
+ # ``vllm_config.model_config.hf_config`` which would be
274
+ # ``ColModernVBertConfig``, not ``ModernBertConfig``.
275
+ self.text_embeddings = ModernBertEmbeddings(text_config)
276
+ self.text_layers = nn.ModuleList(
277
+ [
278
+ ModernBertLayer(
279
+ config=text_config,
280
+ layer_id=i,
281
+ prefix=f"{prefix}.text_layers.{i}",
282
+ )
283
+ for i in range(text_config.num_hidden_layers)
284
+ ]
285
+ )
286
+ self.text_final_norm = nn.LayerNorm(
287
+ text_config.hidden_size,
288
+ eps=text_config.norm_eps,
289
+ bias=text_config.norm_bias,
290
+ )
291
+
292
+ # --- ColBERT projection (768 -> 128, with bias) ---
293
+ self.custom_text_proj = nn.Linear(
294
+ text_config.hidden_size,
295
+ config.embedding_dim,
296
+ bias=True,
297
+ dtype=vllm_config.model_config.head_dtype,
298
+ )
299
+
300
+ # --- Pooler (applies projection + L2 normalize) ---
301
+ pooler_config = vllm_config.model_config.pooler_config
302
+ assert pooler_config is not None
303
+ self.pooler = pooler_for_token_embed(
304
+ pooler_config,
305
+ projector=self.custom_text_proj,
306
+ )
307
+
308
+ # ---- multimodal ---------------------------------------------------------
309
+
310
+ def _get_image_features(
311
+ self,
312
+ pixel_values: torch.Tensor,
313
+ ) -> torch.Tensor:
314
+ # Idefics3ImageProcessor may return (batch, tiles, C, H, W);
315
+ # flatten to (batch*tiles, C, H, W) for SiglipVisionModel.
316
+ if pixel_values.dim() == 5:
317
+ b, t, c, h, w = pixel_values.shape
318
+ pixel_values = pixel_values.reshape(b * t, c, h, w)
319
+ vision_outputs = self.vision_model(
320
+ pixel_values.to(dtype=self.vision_model.dtype),
321
+ )
322
+ return self.connector(vision_outputs)
323
+
324
+ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
325
+ pixel_values = kwargs.pop("pixel_values", None)
326
+ if pixel_values is None:
327
+ return []
328
+ assert isinstance(pixel_values, torch.Tensor)
329
+ image_features = self._get_image_features(pixel_values)
330
+ return list(image_features)
331
+
332
+ # ---- forward ------------------------------------------------------------
333
+
334
+ def forward(
335
+ self,
336
+ input_ids: torch.Tensor,
337
+ positions: torch.Tensor,
338
+ intermediate_tensors: IntermediateTensors | None = None,
339
+ inputs_embeds: torch.Tensor | None = None,
340
+ ) -> torch.Tensor:
341
+ hidden_states = self.text_embeddings(input_ids, inputs_embeds=inputs_embeds)
342
+
343
+ for layer in self.text_layers:
344
+ hidden_states = layer(hidden_states, positions)
345
+
346
+ return self.text_final_norm(hidden_states)
347
+
348
+ # ---- weight loading -----------------------------------------------------
349
+
350
+ # Checkpoint prefix → vLLM param prefix.
351
+ # More-specific prefixes must appear before shorter ones.
352
+ hf_to_vllm_mapper = WeightsMapper(
353
+ orig_to_new_prefix={
354
+ "model.text_model.layers.": "text_layers.",
355
+ "model.text_model.embeddings.": "text_embeddings.",
356
+ "model.text_model.final_norm.": "text_final_norm.",
357
+ "model.connector.modality_projection.": "connector.",
358
+ "model.custom_text_proj.": "custom_text_proj.",
359
+ "model.vision_model.": "vision_model.vision_model.",
360
+ "model.": "",
361
+ },
362
+ )
363
+
364
+ # Checkpoint names for DecoupledEmbedding parts
365
+ _BASE_EMB = "model.text_model.embeddings.tok_embeddings.weight"
366
+ _EXTRA_EMB = (
367
+ "model.text_model.embeddings.tok_embeddings.additional_embedding.weight"
368
+ )
369
+
370
+ def load_weights(
371
+ self,
372
+ weights: Iterable[tuple[str, torch.Tensor]],
373
+ ) -> set[str]:
374
+ # DecoupledEmbedding requires concatenating base + additional
375
+ # embedding tensors before loading, so we extract them first.
376
+ base_embedding_weight: torch.Tensor | None = None
377
+ additional_embedding_weight: torch.Tensor | None = None
378
+ remaining: list[tuple[str, torch.Tensor]] = []
379
+
380
+ for name, tensor in weights:
381
+ if name == self._BASE_EMB:
382
+ base_embedding_weight = tensor
383
+ elif name == self._EXTRA_EMB:
384
+ additional_embedding_weight = tensor
385
+ else:
386
+ remaining.append((name, tensor))
387
+
388
+ # Load all non-embedding weights via AutoWeightsLoader
389
+ loader = AutoWeightsLoader(self)
390
+ loaded_params = loader.load_weights(
391
+ remaining,
392
+ mapper=self.hf_to_vllm_mapper,
393
+ )
394
+
395
+ # Concatenate and load DecoupledEmbedding weights
396
+ if base_embedding_weight is not None:
397
+ combined = base_embedding_weight
398
+ if additional_embedding_weight is not None:
399
+ combined = torch.cat(
400
+ [base_embedding_weight, additional_embedding_weight],
401
+ dim=0,
402
+ )
403
+ param_name = "text_embeddings.tok_embeddings.weight"
404
+ params_dict = dict(self.named_parameters())
405
+ if param_name in params_dict:
406
+ param = params_dict[param_name]
407
+ weight_loader = getattr(
408
+ param,
409
+ "weight_loader",
410
+ default_weight_loader,
411
+ )
412
+ weight_loader(param, combined)
413
+ loaded_params.add(param_name)
414
+ elif additional_embedding_weight is not None:
415
+ raise ValueError(
416
+ "Found 'text_model.embeddings.tok_embeddings"
417
+ ".additional_embedding.weight' but not "
418
+ "'text_model.embeddings.tok_embeddings.weight'"
419
+ )
420
+
421
+ # The pooler wraps ``custom_text_proj`` as its head projector.
422
+ # Mark those params as loaded under the pooler path too.
423
+ if hasattr(self, "pooler") and hasattr(self.pooler, "head"):
424
+ head = self.pooler.head
425
+ projector = getattr(head, "projector", None)
426
+ if projector is not None and isinstance(projector, nn.Module):
427
+ for pname, _ in projector.named_parameters():
428
+ loaded_params.add(f"pooler.head.projector.{pname}")
429
+
430
+ return loaded_params
vllm/model_executor/models/registry.py
+1
@@ -248,6 +248,7 @@ _EMBEDDING_MODELS = {
248
"BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"),
249
# [Multimodal]
250
"CLIPModel": ("clip", "CLIPEmbeddingModel"),
251
+ "ColModernVBertForRetrieval": ("colmodernvbert", "ColModernVBertForRetrieval"),
252
"LlavaNextForConditionalGeneration": (
253
"llava_next",
254
"LlavaNextForConditionalGeneration",
vllm/transformers_utils/config.py
+1
@@ -74,6 +74,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
74
afmoe="AfmoeConfig",
75
bagel="BagelConfig",
76
chatglm="ChatGLMConfig",
77
+ colmodernvbert="ColModernVBertConfig",
78
colqwen3="ColQwen3Config",
79
ops_colqwen3="OpsColQwen3Config",
80
qwen3_vl_nemotron_embed="Qwen3VLNemotronEmbedConfig",
vllm/transformers_utils/configs/__init__.py
+2
@@ -18,6 +18,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
18
"AfmoeConfig": "vllm.transformers_utils.configs.afmoe",
19
"BagelConfig": "vllm.transformers_utils.configs.bagel",
20
"ChatGLMConfig": "vllm.transformers_utils.configs.chatglm",
21
+ "ColModernVBertConfig": "vllm.transformers_utils.configs.colmodernvbert",
22
"ColQwen3Config": "vllm.transformers_utils.configs.colqwen3",
23
"OpsColQwen3Config": "vllm.transformers_utils.configs.colqwen3",
24
"Qwen3VLNemotronEmbedConfig": "vllm.transformers_utils.configs.colqwen3",
@@ -71,6 +72,7 @@ __all__ = [
72
"AfmoeConfig",
73
"BagelConfig",
74
"ChatGLMConfig",
75
+ "ColModernVBertConfig",
76
"ColQwen3Config",
77
"OpsColQwen3Config",
78
"Qwen3VLNemotronEmbedConfig",
vllm/transformers_utils/configs/colmodernvbert.py
new
+65
@@ -0,0 +1,65 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Configuration for ColModernVBERT visual document retrieval model.
4
+
5
+ColModernVBERT combines SigLIP vision encoder + ModernBERT text encoder
6
+with a pixel shuffle connector and ColBERT-style 128-dim per-token embeddings.
7
+
8
+Reference: https://huggingface.co/ModernVBERT/colmodernvbert-merged
9
+"""
10
+
11
+from transformers import ModernBertConfig, PretrainedConfig, SiglipVisionConfig
12
+
13
+
14
+class ColModernVBertConfig(PretrainedConfig):
15
+ model_type = "colmodernvbert"
16
+
17
+ def __init__(
18
+ self,
19
+ embedding_dim: int = 128,
20
+ vlm_config: dict | None = None,
21
+ **kwargs,
22
+ ):
23
+ super().__init__(**kwargs)
24
+ self.embedding_dim = embedding_dim
25
+
26
+ if vlm_config is None:
27
+ vlm_config = {}
28
+
29
+ # Top-level VLM fields
30
+ self.image_token_id = vlm_config.get("image_token_id", 50407)
31
+ self.pixel_shuffle_factor = vlm_config.get("pixel_shuffle_factor", 4)
32
+ self.hidden_size = vlm_config.get("hidden_size", 768)
33
+ additional_vocab_size = vlm_config.get("additional_vocab_size", 40)
34
+
35
+ # Text config (ModernBERT)
36
+ text_cfg = vlm_config.get("text_config", {})
37
+ base_vocab = text_cfg.get("vocab_size", 50368)
38
+ self.text_config = ModernBertConfig(
39
+ vocab_size=base_vocab + additional_vocab_size,
40
+ hidden_size=text_cfg.get("hidden_size", 768),
41
+ intermediate_size=text_cfg.get("intermediate_size", 1152),
42
+ num_hidden_layers=text_cfg.get("num_hidden_layers", 22),
43
+ num_attention_heads=text_cfg.get("num_attention_heads", 12),
44
+ mlp_bias=text_cfg.get("mlp_bias", False),
45
+ max_position_embeddings=vlm_config.get("max_position_embeddings", 8192),
46
+ )
47
+
48
+ # Vision config (SigLIP)
49
+ vis_cfg = vlm_config.get("vision_config", {})
50
+ self.vision_config = SiglipVisionConfig(
51
+ hidden_size=vis_cfg.get("embed_dim", 768),
52
+ image_size=vis_cfg.get("image_size", 512),
53
+ patch_size=vis_cfg.get("patch_size", 16),
54
+ num_hidden_layers=vis_cfg.get("num_hidden_layers", 12),
55
+ intermediate_size=vis_cfg.get("intermediate_size", 3072),
56
+ num_attention_heads=vis_cfg.get("num_attention_heads", 12),
57
+ )
58
+
59
+ @property
60
+ def image_seq_len(self) -> int:
61
+ ps = self.vision_config.image_size // self.vision_config.patch_size
62
+ return (ps * ps) // (self.pixel_shuffle_factor**2)
63
+
64
+ def get_text_config(self, **kwargs):
65
+ return self.text_config