@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 439afa4eea
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 8, 21:30 UTC

feat: Add ColBERT late interaction model support (#33686)

Signed-off-by: Ilya Boytsov <ilyaboytsov1805@gmail.com> Signed-off-by: Ilya Boytsov <boytsovpanamera@mail.ru> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>

Ilya Boytsov committed Feb 5, 2026 at 01:05 UTC 439afa4eea14db2be232a9ce78eacc2c7bbfac77
13 files changed +974 -3
docs/models/pooling_models.md
+56
@@ -307,6 +307,62 @@ An OpenAI client example can be found here: [examples/pooling/embed/openai_embed
307
308 ## Specific models
309
310 +### ColBERT Late Interaction Models
311 +
312 +[ColBERT](https://arxiv.org/abs/2004.12832) (Contextualized Late Interaction over BERT) is a retrieval model that uses per-token embeddings and MaxSim scoring for document ranking. Unlike single-vector embedding models, ColBERT retains token-level representations and computes relevance scores through late interaction, providing better accuracy while being more efficient than cross-encoders.
313 +
314 +vLLM supports ColBERT models for reranking tasks, automatically applying MaxSim scoring for query-document relevance:
315 +
316 +```shell
317 +vllm serve answerdotai/answerai-colbert-small-v1
318 +```
319 +
320 +Currently supports ColBERT models with standard BERT encoders (e.g., `answerdotai/answerai-colbert-small-v1`, `colbert-ir/colbertv2.0`).
321 +
322 +ColBERT models with modified encoder architectures are not yet supported, including BERT variants with rotary embeddings (e.g., `jinaai/jina-colbert-v2`) or other custom encoders (e.g., `LiquidAI/LFM2-ColBERT-350M`).
323 +
324 +If your standard BERT ColBERT model's config doesn't specify the architecture as `HF_ColBERT`, override it with:
325 +
326 +```shell
327 +vllm serve your-colbert-model --hf-overrides '{"architectures": ["HF_ColBERT"]}'
328 +```
329 +
330 +Then you can use the rerank endpoint:
331 +
332 +```shell
333 +curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
334 + "model": "answerdotai/answerai-colbert-small-v1",
335 + "query": "What is machine learning?",
336 + "documents": [
337 + "Machine learning is a subset of artificial intelligence.",
338 + "Python is a programming language.",
339 + "Deep learning uses neural networks."
340 + ]
341 +}'
342 +```
343 +
344 +Or the score endpoint:
345 +
346 +```shell
347 +curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
348 + "model": "answerdotai/answerai-colbert-small-v1",
349 + "text_1": "What is machine learning?",
350 + "text_2": ["Machine learning is a subset of AI.", "The weather is sunny."]
351 +}'
352 +```
353 +
354 +You can also get the raw token embeddings using the pooling endpoint with `token_embed` task:
355 +
356 +```shell
357 +curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
358 + "model": "answerdotai/answerai-colbert-small-v1",
359 + "input": "What is machine learning?",
360 + "task": "token_embed"
361 +}'
362 +```
363 +
364 +An example can be found here: [examples/pooling/score/colbert_rerank_online.py](../../examples/pooling/score/colbert_rerank_online.py)
365 +
366 ### BAAI/bge-m3
367
368 The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json`
examples/pooling/score/colbert_rerank_online.py new
+57
@@ -0,0 +1,57 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Example of using ColBERT late interaction model for reranking.
5 +
6 +ColBERT (Contextualized Late Interaction over BERT) uses per-token embeddings
7 +and MaxSim scoring for document reranking, providing better accuracy than
8 +single-vector models while being more efficient than cross-encoders.
9 +
10 +Start the server with:
11 + vllm serve answerdotai/answerai-colbert-small-v1
12 +
13 +Then run this script:
14 + python colbert_rerank_online.py
15 +"""
16 +
17 +import json
18 +
19 +import requests
20 +
21 +url = "http://127.0.0.1:8000/rerank"
22 +
23 +headers = {"accept": "application/json", "Content-Type": "application/json"}
24 +
25 +data = {
26 + "model": "answerdotai/answerai-colbert-small-v1",
27 + "query": "What is machine learning?",
28 + "documents": [
29 + "Machine learning is a subset of artificial intelligence.",
30 + "Python is a programming language.",
31 + "Deep learning uses neural networks for complex tasks.",
32 + "The weather today is sunny.",
33 + ],
34 +}
35 +
36 +
37 +def main():
38 + response = requests.post(url, headers=headers, json=data)
39 +
40 + if response.status_code == 200:
41 + print("ColBERT Rerank Request successful!")
42 + result = response.json()
43 + print(json.dumps(result, indent=2))
44 +
45 + # Show ranked results
46 + print("\nRanked documents (most relevant first):")
47 + for item in result["results"]:
48 + doc_idx = item["index"]
49 + score = item["relevance_score"]
50 + print(f" Score {score:.4f}: {data['documents'][doc_idx]}")
51 + else:
52 + print(f"Request failed with status code: {response.status_code}")
53 + print(response.text)
54 +
55 +
56 +if __name__ == "__main__":
57 + main()
tests/entrypoints/pooling/score/test_online_colbert.py new
+154
@@ -0,0 +1,154 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Online API tests for ColBERT late interaction scoring."""
4 +
5 +import pytest
6 +import requests
7 +
8 +from tests.utils import RemoteOpenAIServer
9 +from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse
10 +
11 +# ColBERT model - using answerai-colbert-small-v1 as it's a smaller model
12 +MODEL_NAME = "answerdotai/answerai-colbert-small-v1"
13 +COLBERT_DIM = 96 # This model uses 96-dimensional output
14 +DTYPE = "half"
15 +MAX_MODEL_LEN = 512
16 +
17 +
18 +@pytest.fixture(scope="module")
19 +def server():
20 + args = [
21 + "--max-model-len",
22 + str(MAX_MODEL_LEN),
23 + ]
24 +
25 + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
26 + yield remote_server
27 +
28 +
29 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
30 +def test_colbert_rerank(server: RemoteOpenAIServer, model_name: str):
31 + """Test ColBERT rerank endpoint."""
32 + query = "What is the capital of France?"
33 + documents = [
34 + "The capital of Brazil is Brasilia.",
35 + "The capital of France is Paris.",
36 + ]
37 +
38 + rerank_response = requests.post(
39 + server.url_for("rerank"),
40 + json={
41 + "model": model_name,
42 + "query": query,
43 + "documents": documents,
44 + },
45 + )
46 + rerank_response.raise_for_status()
47 + rerank = RerankResponse.model_validate(rerank_response.json())
48 +
49 + assert rerank.id is not None
50 + assert rerank.results is not None
51 + assert len(rerank.results) == 2
52 +
53 + # The relevant document (Paris) should have higher score
54 + paris_result = next(r for r in rerank.results if r.index == 1)
55 + brazil_result = next(r for r in rerank.results if r.index == 0)
56 +
57 + assert paris_result.relevance_score > brazil_result.relevance_score
58 +
59 +
60 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
61 +def test_colbert_rerank_top_n(server: RemoteOpenAIServer, model_name: str):
62 + """Test ColBERT rerank with top_n parameter."""
63 + query = "What is the capital of France?"
64 + documents = [
65 + "The capital of Brazil is Brasilia.",
66 + "The capital of France is Paris.",
67 + "Machine learning is a field of AI.",
68 + ]
69 +
70 + rerank_response = requests.post(
71 + server.url_for("rerank"),
72 + json={
73 + "model": model_name,
74 + "query": query,
75 + "documents": documents,
76 + "top_n": 2,
77 + },
78 + )
79 + rerank_response.raise_for_status()
80 + rerank = RerankResponse.model_validate(rerank_response.json())
81 +
82 + assert len(rerank.results) == 2
83 + # Top result should be about Paris (index 1)
84 + assert rerank.results[0].index == 1
85 +
86 +
87 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
88 +def test_colbert_score(server: RemoteOpenAIServer, model_name: str):
89 + """Test ColBERT score endpoint."""
90 + text_1 = "What is the capital of France?"
91 + text_2 = ["The capital of France is Paris.", "Python is a language."]
92 +
93 + score_response = requests.post(
94 + server.url_for("score"),
95 + json={
96 + "model": model_name,
97 + "text_1": text_1,
98 + "text_2": text_2,
99 + },
100 + )
101 + score_response.raise_for_status()
102 + score = ScoreResponse.model_validate(score_response.json())
103 +
104 + assert score.id is not None
105 + assert score.data is not None
106 + assert len(score.data) == 2
107 +
108 + # The relevant document should have higher score
109 + assert score.data[0].score > score.data[1].score
110 +
111 +
112 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
113 +def test_colbert_token_embed(server: RemoteOpenAIServer, model_name: str):
114 + """Test ColBERT token_embed task via pooling endpoint."""
115 + text = "What is the capital of France?"
116 +
117 + pooling_response = requests.post(
118 + server.url_for("pooling"),
119 + json={
120 + "model": model_name,
121 + "input": text,
122 + "task": "token_embed",
123 + },
124 + )
125 + pooling_response.raise_for_status()
126 + pooling = pooling_response.json()
127 +
128 + assert "data" in pooling
129 + assert len(pooling["data"]) == 1
130 +
131 + # Token embeddings should be 2D
132 + embeddings = pooling["data"][0]["data"]
133 + assert isinstance(embeddings, list)
134 + assert len(embeddings) > 0 # Should have tokens
135 + assert len(embeddings[0]) == COLBERT_DIM
136 +
137 +
138 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
139 +def test_colbert_embed_not_supported(server: RemoteOpenAIServer, model_name: str):
140 + """Test that ColBERT model does not support 'embed' task."""
141 + text = "What is the capital of France?"
142 +
143 + pooling_response = requests.post(
144 + server.url_for("pooling"),
145 + json={
146 + "model": model_name,
147 + "input": text,
148 + "task": "embed",
149 + },
150 + )
151 +
152 + # Should return error
153 + assert pooling_response.status_code == 400
154 + assert "Task embed is not supported" in pooling_response.text
tests/models/language/pooling/test_colbert.py new
+247
@@ -0,0 +1,247 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for ColBERT late interaction scoring."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
9 +
10 +# ColBERT model - using answerai-colbert-small-v1 as it's a smaller model
11 +# suitable for testing (based on BERT-base)
12 +COLBERT_MODEL = "answerdotai/answerai-colbert-small-v1"
13 +COLBERT_DIM = 96 # This model uses 96-dimensional output
14 +
15 +TEXTS_1 = [
16 + "What is the capital of France?",
17 + "What is the capital of Germany?",
18 +]
19 +
20 +TEXTS_2 = [
21 + "The capital of France is Paris.",
22 + "The capital of Germany is Berlin.",
23 +]
24 +
25 +DTYPE = "half"
26 +
27 +
28 +@pytest.fixture(scope="module")
29 +def colbert_model_name():
30 + return COLBERT_MODEL
31 +
32 +
33 +def test_colbert_token_embed(vllm_runner, colbert_model_name):
34 + """Test that ColBERT model produces token embeddings."""
35 + with vllm_runner(
36 + colbert_model_name,
37 + runner="pooling",
38 + dtype=DTYPE,
39 + max_model_len=512,
40 + enforce_eager=True,
41 + ) as vllm_model:
42 + # Get token embeddings for a single text
43 + outputs = vllm_model.token_embed([TEXTS_1[0]])
44 +
45 + assert len(outputs) == 1
46 + # Token embeddings should be 2D: [num_tokens, colbert_dim]
47 + emb = torch.tensor(outputs[0])
48 + assert emb.dim() == 2
49 + assert emb.shape[1] == COLBERT_DIM
50 + # Should have at least a few tokens
51 + assert emb.shape[0] > 1
52 +
53 +
54 +def test_colbert_late_interaction_1_to_1(vllm_runner, colbert_model_name):
55 + """Test ColBERT late interaction scoring with 1:1 query-document pair."""
56 + with vllm_runner(
57 + colbert_model_name,
58 + runner="pooling",
59 + dtype=DTYPE,
60 + max_model_len=512,
61 + enforce_eager=True,
62 + ) as vllm_model:
63 + # Get token embeddings
64 + q_outputs = vllm_model.token_embed([TEXTS_1[0]])
65 + d_outputs = vllm_model.token_embed([TEXTS_2[0]])
66 +
67 + q_emb = torch.tensor(q_outputs[0])
68 + d_emb = torch.tensor(d_outputs[0])
69 +
70 + # Compute MaxSim manually
71 + manual_score = compute_maxsim_score(q_emb, d_emb).item()
72 +
73 + # Use the score API (which should internally use _late_interaction_score)
74 + vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2[0])
75 +
76 + assert len(vllm_scores) == 1
77 + assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
78 +
79 +
80 +def test_colbert_late_interaction_1_to_N(vllm_runner, colbert_model_name):
81 + """Test ColBERT late interaction scoring with 1:N query-documents."""
82 + with vllm_runner(
83 + colbert_model_name,
84 + runner="pooling",
85 + dtype=DTYPE,
86 + max_model_len=512,
87 + enforce_eager=True,
88 + ) as vllm_model:
89 + # Get token embeddings
90 + q_outputs = vllm_model.token_embed([TEXTS_1[0]])
91 + d_outputs = vllm_model.token_embed(TEXTS_2)
92 +
93 + q_emb = torch.tensor(q_outputs[0])
94 +
95 + # Compute MaxSim manually for each document
96 + manual_scores = []
97 + for d_out in d_outputs:
98 + d_emb = torch.tensor(d_out)
99 + manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
100 +
101 + # Use the score API
102 + vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2)
103 +
104 + assert len(vllm_scores) == 2
105 + for i in range(2):
106 + assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
107 +
108 +
109 +def test_colbert_late_interaction_N_to_N(vllm_runner, colbert_model_name):
110 + """Test ColBERT late interaction scoring with N:N query-documents."""
111 + with vllm_runner(
112 + colbert_model_name,
113 + runner="pooling",
114 + dtype=DTYPE,
115 + max_model_len=512,
116 + enforce_eager=True,
117 + ) as vllm_model:
118 + # Get token embeddings
119 + q_outputs = vllm_model.token_embed(TEXTS_1)
120 + d_outputs = vllm_model.token_embed(TEXTS_2)
121 +
122 + # Compute MaxSim manually for each pair
123 + manual_scores = []
124 + for q_out, d_out in zip(q_outputs, d_outputs):
125 + q_emb = torch.tensor(q_out)
126 + d_emb = torch.tensor(d_out)
127 + manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
128 +
129 + # Use the score API
130 + vllm_scores = vllm_model.score(TEXTS_1, TEXTS_2)
131 +
132 + assert len(vllm_scores) == 2
133 + for i in range(2):
134 + assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
135 +
136 +
137 +def test_colbert_relevance_ordering(vllm_runner, colbert_model_name):
138 + """Test that ColBERT scores relevant documents higher than irrelevant ones."""
139 + query = "What is machine learning?"
140 + documents = [
141 + "Machine learning is a subset of artificial intelligence.",
142 + "Python is a programming language.",
143 + "Deep learning uses neural networks.",
144 + ]
145 +
146 + with vllm_runner(
147 + colbert_model_name,
148 + runner="pooling",
149 + dtype=DTYPE,
150 + max_model_len=512,
151 + enforce_eager=True,
152 + ) as vllm_model:
153 + scores = vllm_model.score(query, documents)
154 +
155 + assert len(scores) == 3
156 + # ML-related documents should score higher than unrelated Python doc
157 + # Document 0 (ML definition) should be most relevant
158 + # Document 2 (Deep learning) should also be relevant
159 + # Document 1 (Python) should be least relevant
160 + assert scores[0] > scores[1], "ML doc should score higher than Python doc"
161 + assert scores[2] > scores[1], "DL doc should score higher than Python doc"
162 +
163 +
164 +def test_colbert_embed_not_supported(vllm_runner, colbert_model_name):
165 + """Test that ColBERT model does not support 'embed' task."""
166 + with (
167 + vllm_runner(
168 + colbert_model_name,
169 + runner="pooling",
170 + dtype=DTYPE,
171 + max_model_len=512,
172 + enforce_eager=True,
173 + ) as vllm_model,
174 + pytest.raises(ValueError, match="Embedding API is not supported"),
175 + ):
176 + vllm_model.embed([TEXTS_1[0]])
177 +
178 +
179 +def test_colbert_hf_comparison(vllm_runner, colbert_model_name):
180 + """Test that vLLM ColBERT produces same embeddings as HuggingFace."""
181 + import torch.nn.functional as F
182 + from huggingface_hub import hf_hub_download
183 + from safetensors.torch import load_file
184 + from transformers import AutoTokenizer, BertModel
185 +
186 + test_texts = [TEXTS_1[0], TEXTS_2[0]]
187 +
188 + # Get vLLM embeddings first (to avoid GPU memory contention)
189 + # Use fp32 to match HuggingFace default precision for fair comparison
190 + with vllm_runner(
191 + colbert_model_name,
192 + runner="pooling",
193 + dtype="float32",
194 + max_model_len=512,
195 + enforce_eager=True,
196 + ) as vllm_model:
197 + vllm_outputs = vllm_model.token_embed(test_texts)
198 +
199 + # Get HuggingFace reference embeddings on CPU
200 + # Load the base BERT model and manually apply the ColBERT linear projection
201 + hf_tokenizer = AutoTokenizer.from_pretrained(colbert_model_name)
202 + hf_bert = BertModel.from_pretrained(colbert_model_name)
203 + hf_bert.eval()
204 +
205 + # Load the ColBERT linear weights from safetensors
206 + weights_path = hf_hub_download(colbert_model_name, filename="model.safetensors")
207 + weights = load_file(weights_path)
208 + linear_weight = weights["linear.weight"] # [96, 384]
209 +
210 + hf_embeddings = []
211 + for text in test_texts:
212 + inputs = hf_tokenizer(text, return_tensors="pt")
213 + with torch.no_grad():
214 + outputs = hf_bert(**inputs)
215 + # Get last hidden state: [1, seq_len, 384]
216 + hidden_states = outputs.last_hidden_state
217 + # Apply ColBERT linear projection: [1, seq_len, 96]
218 + token_emb = F.linear(hidden_states, linear_weight)
219 + # L2 normalize
220 + token_emb = F.normalize(token_emb, p=2, dim=-1)
221 + hf_embeddings.append(token_emb.squeeze(0).float())
222 +
223 + # Compare embeddings
224 + for i, (hf_emb, vllm_out) in enumerate(zip(hf_embeddings, vllm_outputs)):
225 + vllm_emb = torch.tensor(vllm_out).float()
226 +
227 + # Print first few components for debugging
228 + print(f"\n=== Text {i}: '{test_texts[i][:30]}...' ===")
229 + print(f"HF shape: {hf_emb.shape}, vLLM shape: {vllm_emb.shape}")
230 + print(f"HF first token, first 10 dims: {hf_emb[0, :10].tolist()}")
231 + print(f"vLLM first token, first 10 dims: {vllm_emb[0, :10].tolist()}")
232 + print(f"HF last token, first 10 dims: {hf_emb[-1, :10].tolist()}")
233 + print(f"vLLM last token, first 10 dims: {vllm_emb[-1, :10].tolist()}")
234 +
235 + # Should have same shape
236 + assert hf_emb.shape == vllm_emb.shape, (
237 + f"Shape mismatch for text {i}: HF {hf_emb.shape} vs vLLM {vllm_emb.shape}"
238 + )
239 +
240 + # Should have same values (with tolerance for fp16)
241 + torch.testing.assert_close(
242 + vllm_emb,
243 + hf_emb,
244 + rtol=1e-2,
245 + atol=1e-2,
246 + msg=f"Embedding mismatch for text {i}",
247 + )
tests/models/registry.py
+1
@@ -520,6 +520,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
520 _EMBEDDING_EXAMPLE_MODELS = {
521 # [Text-only]
522 "BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
523 + "HF_ColBERT": _HfExamplesInfo("answerdotai/answerai-colbert-small-v1"),
524 "BgeM3EmbeddingModel": _HfExamplesInfo("BAAI/bge-m3"),
525 "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"),
526 "Gemma3TextModel": _HfExamplesInfo("google/embeddinggemma-300m"),
vllm/config/model.py
+5
@@ -1411,6 +1411,11 @@ class ModelConfig:
1411 self._model_info.supports_cross_encoding or self.convert_type == "classify"
1412 )
1413
1414 + @property
1415 + def is_late_interaction(self) -> bool:
1416 + """Check if model uses late interaction (ColBERT-style) scoring."""
1417 + return self._model_info.supports_late_interaction
1418 +
1419 @property
1420 def is_pp_supported(self) -> bool:
1421 return self._model_info.supports_pp
vllm/entrypoints/llm.py
+96 -1
@@ -44,6 +44,7 @@ from vllm.entrypoints.pooling.score.utils import (
44 ScoreMultiModalParam,
45 _cosine_similarity,
46 compress_token_type_ids,
47 + compute_maxsim_score,
48 get_score_prompt,
49 validate_score_input,
50 )
@@ -1368,6 +1369,87 @@ class LLM:
1369 items = self.engine_class.validate_outputs(scores, PoolingRequestOutput)
1370 return [ScoringRequestOutput.from_base(item) for item in items]
1371
1372 + def _late_interaction_score(
1373 + self,
1374 + data_1: list[ScoreData],
1375 + data_2: list[ScoreData],
1376 + *,
1377 + use_tqdm: bool | Callable[..., tqdm],
1378 + pooling_params: PoolingParams | None,
1379 + lora_request: list[LoRARequest] | LoRARequest | None,
1380 + tokenization_kwargs: dict[str, Any],
1381 + ) -> list[ScoringRequestOutput]:
1382 + """
1383 + Late interaction scoring (ColBERT MaxSim).
1384 +
1385 + Encodes queries and documents into per-token embeddings, then computes
1386 + MaxSim: sum over query tokens of max similarity to any document token.
1387 + """
1388 + from vllm.outputs import PoolingOutput
1389 +
1390 + tokenizer = self.get_tokenizer()
1391 +
1392 + # Extract text from ScoreData
1393 + text_1: list[str] = []
1394 + for text in data_1:
1395 + if not isinstance(text, str):
1396 + raise NotImplementedError(
1397 + "Late interaction scores currently do not support multimodal input."
1398 + )
1399 + text_1.append(text)
1400 +
1401 + text_2: list[str] = []
1402 + for text in data_2:
1403 + if not isinstance(text, str):
1404 + raise NotImplementedError(
1405 + "Late interaction scores currently do not support multimodal input."
1406 + )
1407 + text_2.append(text)
1408 +
1409 + encoded_output: list[PoolingRequestOutput] = self.encode(
1410 + text_1 + text_2,
1411 + use_tqdm=use_tqdm,
1412 + lora_request=lora_request,
1413 + pooling_params=pooling_params,
1414 + pooling_task="token_embed",
1415 + tokenization_kwargs=tokenization_kwargs,
1416 + )
1417 +
1418 + encoded_output_1: list[PoolingRequestOutput] = encoded_output[0 : len(text_1)]
1419 + encoded_output_2: list[PoolingRequestOutput] = encoded_output[len(text_1) :]
1420 +
1421 + if len(encoded_output_1) == 1:
1422 + encoded_output_1 = encoded_output_1 * len(encoded_output_2)
1423 +
1424 + # Compute MaxSim scores
1425 + scores: list[PoolingRequestOutput] = []
1426 + padding: list[int] = []
1427 + if (pad_token_id := tokenizer.pad_token_id) is not None:
1428 + padding = [pad_token_id]
1429 +
1430 + for emb_1, emb_2 in zip(encoded_output_1, encoded_output_2):
1431 + # emb_1.outputs.data: [query_len, dim]
1432 + # emb_2.outputs.data: [doc_len, dim]
1433 + q_emb = emb_1.outputs.data
1434 + d_emb = emb_2.outputs.data
1435 +
1436 + maxsim_score = compute_maxsim_score(q_emb, d_emb)
1437 +
1438 + tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids
1439 +
1440 + scores.append(
1441 + PoolingRequestOutput(
1442 + request_id=f"{emb_1.request_id}_{emb_2.request_id}",
1443 + outputs=PoolingOutput(data=maxsim_score),
1444 + prompt_token_ids=tokens,
1445 + num_cached_tokens=emb_1.num_cached_tokens + emb_2.num_cached_tokens,
1446 + finished=True,
1447 + )
1448 + )
1449 +
1450 + items = self.engine_class.validate_outputs(scores, PoolingRequestOutput)
1451 + return [ScoringRequestOutput.from_base(item) for item in items]
1452 +
1453 def _cross_encoding_score(
1454 self,
1455 data_1: list[ScoreData],
@@ -1497,7 +1579,11 @@ class LLM:
1579 )
1580
1581 supported_tasks = self.supported_tasks
1500 - if all(t not in supported_tasks for t in ("embed", "classify")):
1582 + # Late interaction models (e.g., ColBERT) use token_embed for scoring
1583 + is_late_interaction = model_config.is_late_interaction
1584 + if not is_late_interaction and all(
1585 + t not in supported_tasks for t in ("embed", "classify")
1586 + ):
1587 raise ValueError(
1588 "Score API is not supported by this model. "
1589 "Try converting the model using "
@@ -1538,6 +1624,15 @@ class LLM:
1624 tokenization_kwargs=encode_kwargs,
1625 score_template=chat_template,
1626 )
1627 + elif is_late_interaction:
1628 + return self._late_interaction_score(
1629 + score_data_1,
1630 + score_data_2,
1631 + use_tqdm=use_tqdm,
1632 + pooling_params=pooling_params,
1633 + lora_request=lora_request,
1634 + tokenization_kwargs=encode_kwargs,
1635 + )
1636 else:
1637 return self._embedding_score(
1638 score_data_1,
vllm/entrypoints/pooling/__init__.py
+10 -2
@@ -37,7 +37,11 @@ def register_pooling_api_routers(
37
38 app.include_router(embed_router)
39
40 - if "score" in supported_tasks or "embed" in supported_tasks:
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)
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
47 app.include_router(score_router)
@@ -101,6 +105,10 @@ def init_pooling_state(
105 if "classify" in supported_tasks
106 else None
107 )
108 + # ServingScores handles score/rerank for:
109 + # - "score" task (cross-encoder models)
110 + # - "embed" task (bi-encoder models)
111 + # - "token_embed" task (late interaction models like ColBERT)
112 state.openai_serving_scores = (
113 ServingScores(
114 engine_client,
@@ -109,6 +117,6 @@ def init_pooling_state(
117 score_template=resolved_chat_template,
118 log_error_stack=args.log_error_stack,
119 )
112 - if ("embed" in supported_tasks or "score" in supported_tasks)
120 + if any(t in supported_tasks for t in ("embed", "score", "token_embed"))
121 else None
122 )
vllm/entrypoints/pooling/score/serving.py
+140
@@ -31,6 +31,7 @@ from vllm.entrypoints.pooling.score.utils import (
31 ScoreInputs,
32 _cosine_similarity,
33 compress_token_type_ids,
34 + compute_maxsim_score,
35 get_score_prompt,
36 validate_score_input,
37 )
@@ -68,9 +69,12 @@ class ServingScores(OpenAIServing):
69 self.is_cross_encoder = self.model_config.is_cross_encoder
70 self.is_multimodal_model = self.model_config.is_multimodal_model
71 self.architecture = self.model_config.architecture
72 + self.is_late_interaction = self.model_config.is_late_interaction
73
74 if self.is_cross_encoder:
75 self._score_func = self._cross_encoding_score
76 + elif self.is_late_interaction:
77 + self._score_func = self._late_interaction_score
78 else:
79 self._score_func = self._embedding_score
80
@@ -172,6 +176,142 @@ class ServingScores(OpenAIServing):
176
177 return final_res_batch
178
179 + async def _late_interaction_score(
180 + self,
181 + data_1: list[ScoreData],
182 + data_2: list[ScoreData],
183 + request: RerankRequest | ScoreRequest,
184 + request_id: str,
185 + lora_request: LoRARequest | None = None,
186 + trace_headers: Mapping[str, str] | None = None,
187 + ) -> list[PoolingRequestOutput] | ErrorResponse:
188 + """
189 + Late interaction scoring (ColBERT MaxSim).
190 +
191 + Encodes queries and documents into per-token embeddings, then computes
192 + MaxSim: sum over query tokens of max similarity to any document token.
193 + """
194 + input_texts: list[str] = []
195 + for text in data_1 + data_2:
196 + if not isinstance(text, str):
197 + raise NotImplementedError(
198 + "Late interaction scores currently do not support multimodal input."
199 + )
200 + input_texts.append(text)
201 +
202 + model_config = self.model_config
203 + tokenizer = self.renderer.get_tokenizer()
204 +
205 + encode_async = make_async(
206 + tokenizer.encode,
207 + executor=self._tokenizer_executor,
208 + )
209 +
210 + tokenization_kwargs = request.build_tok_params(model_config).get_encode_kwargs()
211 + tokenized_prompts = await asyncio.gather(
212 + *(encode_async(t, **tokenization_kwargs) for t in input_texts)
213 + )
214 +
215 + engine_prompts: list[TokensPrompt] = []
216 + for tok_result, input_text in zip(tokenized_prompts, input_texts):
217 + text_token_prompt = self._validate_input(request, tok_result, input_text)
218 +
219 + engine_prompts.append(
220 + TokensPrompt(prompt_token_ids=text_token_prompt["prompt_token_ids"])
221 + )
222 +
223 + # Schedule the request and get the result generator.
224 + generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
225 +
226 + # Use token_embed task for late interaction models
227 + from vllm import PoolingParams
228 +
229 + pooling_params = PoolingParams(
230 + task="token_embed",
231 + truncate_prompt_tokens=request.truncate_prompt_tokens,
232 + use_activation=request.use_activation,
233 + )
234 +
235 + try:
236 + pooling_params.verify("token_embed", self.model_config)
237 + except ValueError as e:
238 + return self.create_error_response(str(e))
239 +
240 + for i, engine_prompt in enumerate(engine_prompts):
241 + request_id_item = f"{request_id}-{i}"
242 +
243 + self._log_inputs(
244 + request_id_item,
245 + input_texts[i],
246 + params=pooling_params,
247 + lora_request=lora_request,
248 + )
249 +
250 + generators.append(
251 + self.engine_client.encode(
252 + engine_prompt,
253 + pooling_params,
254 + request_id_item,
255 + lora_request=lora_request,
256 + trace_headers=trace_headers,
257 + priority=request.priority,
258 + )
259 + )
260 +
261 + result_generator = merge_async_iterators(*generators)
262 +
263 + # Collect token embeddings
264 + embeddings: list[PoolingRequestOutput | None] = [None] * len(engine_prompts)
265 +
266 + async for i, res in result_generator:
267 + embeddings[i] = res
268 +
269 + # Split into query and document embeddings
270 + emb_data_1: list[PoolingRequestOutput] = []
271 + emb_data_2: list[PoolingRequestOutput] = []
272 +
273 + for i in range(0, len(data_1)):
274 + assert (emb := embeddings[i]) is not None
275 + emb_data_1.append(emb)
276 +
277 + for i in range(len(data_1), len(embeddings)):
278 + assert (emb := embeddings[i]) is not None
279 + emb_data_2.append(emb)
280 +
281 + # Expand queries if 1:N scoring
282 + if len(emb_data_1) == 1:
283 + emb_data_1 = emb_data_1 * len(emb_data_2)
284 +
285 + # Compute MaxSim scores
286 + from vllm.outputs import PoolingOutput
287 +
288 + scores: list[PoolingRequestOutput] = []
289 + padding: list[int] = []
290 + if (pad_token_id := tokenizer.pad_token_id) is not None:
291 + padding = [pad_token_id]
292 +
293 + for emb_1, emb_2 in zip(emb_data_1, emb_data_2):
294 + # emb_1.outputs.data: [query_len, dim]
295 + # emb_2.outputs.data: [doc_len, dim]
296 + q_emb = emb_1.outputs.data
297 + d_emb = emb_2.outputs.data
298 +
299 + maxsim_score = compute_maxsim_score(q_emb, d_emb)
300 +
301 + tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids
302 +
303 + scores.append(
304 + PoolingRequestOutput(
305 + request_id=f"{emb_1.request_id}_{emb_2.request_id}",
306 + outputs=PoolingOutput(data=maxsim_score),
307 + prompt_token_ids=tokens,
308 + num_cached_tokens=emb_1.num_cached_tokens + emb_2.num_cached_tokens,
309 + finished=True,
310 + )
311 + )
312 +
313 + return scores
314 +
315 async def _cross_encoding_score(
316 self,
317 data_1: list[ScoreData],
vllm/entrypoints/pooling/score/utils.py
+18
@@ -3,6 +3,7 @@
3 from collections.abc import Iterable
4 from typing import Any, TypeAlias, cast
5
6 +import torch
7 from torch.nn import CosineSimilarity
8 from typing_extensions import Required, TypedDict
9
@@ -34,6 +35,23 @@ ScoreContentPartParam: TypeAlias = (
35 )
36
37
38 +def compute_maxsim_score(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tensor:
39 + """
40 + Compute ColBERT MaxSim score.
41 +
42 + Args:
43 + q_emb: Query token embeddings [query_len, dim]
44 + d_emb: Document token embeddings [doc_len, dim]
45 +
46 + Returns:
47 + MaxSim score (sum over query tokens of max similarity to any doc token)
48 + """
49 + # [query_len, doc_len]
50 + token_scores = torch.matmul(q_emb, d_emb.T)
51 + # Max over document tokens, sum over query tokens
52 + return token_scores.amax(dim=-1).sum()
53 +
54 +
55 class ScoreMultiModalParam(TypedDict, total=False):
56 """
57 A specialized parameter type for scoring multimodal content
vllm/model_executor/models/colbert.py new
+152
@@ -0,0 +1,152 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +ColBERT late interaction model for retrieval and reranking.
5 +
6 +ColBERT uses per-token embeddings and late interaction (MaxSim) scoring
7 +instead of single-vector representations or cross-encoder concatenation.
8 +
9 +Reference: https://arxiv.org/abs/2004.12832
10 +"""
11 +
12 +from collections.abc import Iterable
13 +from typing import ClassVar, Literal
14 +
15 +import torch
16 +from torch import nn
17 +
18 +from vllm.config import PoolerConfig, VllmConfig
19 +from vllm.model_executor.layers.pooler import Pooler
20 +from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
21 +
22 +from .bert import BertEmbeddingModel, BertModel
23 +from .interfaces_base import default_pooling_type
24 +
25 +
26 +@default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL")
27 +class ColBERTModel(BertEmbeddingModel):
28 + """ColBERT late interaction model for retrieval/reranking.
29 +
30 + This model extends BertEmbeddingModel with a ColBERT-style linear
31 + projection layer for per-token embeddings. It supports only:
32 + - "token_embed" task: Per-token embeddings for late interaction
33 +
34 + ColBERT is fundamentally a per-token embedding model - the linear
35 + projection is trained for per-token representations, not for CLS
36 + pooling. Use a dedicated dense embedding model if you need single-
37 + vector representations.
38 +
39 + The ColBERT scoring (MaxSim) is computed externally, either client-side
40 + or via the late interaction scoring path in ServingScores.
41 +
42 + Attributes:
43 + colbert_linear: Linear projection from hidden_size to colbert_dim
44 + supports_late_interaction: Flag indicating this model uses late
45 + interaction scoring
46 + """
47 +
48 + # Mark this model as supporting late interaction scoring
49 + supports_late_interaction: ClassVar[Literal[True]] = True
50 +
51 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
52 + # Get config before calling super().__init__
53 + config = vllm_config.model_config.hf_config
54 + self.hidden_size = config.hidden_size
55 + self.head_dtype = vllm_config.model_config.head_dtype
56 +
57 + # ColBERT dimension - check various config field names used by different
58 + # ColBERT implementations. If not found in config, will be inferred
59 + # from loaded weights in load_weights()
60 + self.colbert_dim: int | None = (
61 + getattr(config, "colbert_dim", None)
62 + or getattr(config, "dim", None)
63 + or getattr(config, "projection_dim", None)
64 + )
65 +
66 + # Initialize parent (this will call _build_pooler)
67 + super().__init__(vllm_config=vllm_config, prefix=prefix)
68 +
69 + def _build_model(self, vllm_config: VllmConfig, prefix: str = "") -> BertModel:
70 + return BertModel(vllm_config=vllm_config, prefix=prefix)
71 +
72 + def _build_colbert_linear(self) -> nn.Linear:
73 + """Build the ColBERT linear projection layer."""
74 + if self.colbert_dim is None:
75 + raise ValueError("colbert_dim must be set before building the linear layer")
76 + return nn.Linear(
77 + self.hidden_size,
78 + self.colbert_dim,
79 + bias=False,
80 + dtype=self.head_dtype,
81 + )
82 +
83 + def _build_pooler(self, pooler_config: PoolerConfig) -> Pooler:
84 + # ColBERT linear projection: hidden_size -> colbert_dim
85 + # Original ColBERT uses bias=False
86 + # If colbert_dim is not set from config, it will be inferred during
87 + # load_weights and the linear layer will be created there
88 + if self.colbert_dim is not None:
89 + self.colbert_linear = self._build_colbert_linear()
90 + else:
91 + # Placeholder - will be created when weights are loaded
92 + self.colbert_linear = None
93 +
94 + # ColBERT only supports token_embed - it's fundamentally a per-token
95 + # embedding model.
96 + return pooler_for_token_embed(
97 + pooler_config,
98 + projector=self.colbert_linear,
99 + )
100 +
101 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
102 + def _strip(name: str) -> str:
103 + for p in ("model.", "bert."):
104 + if name.startswith(p):
105 + name = name[len(p) :]
106 + return name
107 +
108 + weights_list = list(weights)
109 + model_side: list[tuple[str, torch.Tensor]] = []
110 + colbert_side: list[tuple[str, torch.Tensor]] = []
111 +
112 + for name, weight in weights_list:
113 + stripped = _strip(name)
114 + # Handle different checkpoint naming conventions for ColBERT linear
115 + if stripped in ("linear.weight", "colbert_linear.weight"):
116 + colbert_side.append(("colbert_linear.weight", weight))
117 + elif stripped.startswith("linear.") or stripped.startswith(
118 + "colbert_linear."
119 + ):
120 + new_name = stripped.replace("linear.", "colbert_linear.")
121 + colbert_side.append((new_name, weight))
122 + else:
123 + model_side.append((stripped, weight))
124 +
125 + # Load base BERT weights using BertModel.load_weights which handles QKV fusion
126 + loaded: set[str] = set()
127 + loaded_model = self.model.load_weights(model_side)
128 + loaded.update({"model." + n for n in loaded_model})
129 +
130 + # Load ColBERT linear weights
131 + if colbert_side:
132 + for name, weight in colbert_side:
133 + if name == "colbert_linear.weight":
134 + # Infer colbert_dim from weights if not set in config
135 + if self.colbert_dim is None:
136 + # Weight shape is [colbert_dim, hidden_size]
137 + self.colbert_dim = weight.shape[0]
138 + # Create the linear layer now that we know the dimension
139 + self.colbert_linear = self._build_colbert_linear()
140 + # Move to the same device as the model's existing parameters
141 + device = next(self.model.parameters()).device
142 + self.colbert_linear.to(device)
143 + # Update the pooler's projector to use the new linear layer
144 + self.pooler.head.projector = self.colbert_linear
145 +
146 + # Load weights directly into the pooler's projector
147 + weight = weight.to(self.pooler.head.projector.weight.device)
148 + self.pooler.head.projector.weight.data.copy_(weight)
149 + loaded.add("pooler.head.projector.weight")
150 + break
151 +
152 + return loaded
vllm/model_executor/models/interfaces.py
+34
@@ -981,6 +981,40 @@ def supports_cross_encoding(
981 return is_pooling_model(model) and _supports_cross_encoding(model)
982
983
984 +@runtime_checkable
985 +class SupportsLateInteraction(Protocol):
986 + """The interface required for all models that support late interaction.
987 +
988 + Late interaction models (like ColBERT) encode queries and documents
989 + separately into per-token embeddings, then compute similarity via
990 + MaxSim (max over document tokens, sum over query tokens).
991 + """
992 +
993 + supports_late_interaction: ClassVar[Literal[True]] = True
994 +
995 +
996 +@overload
997 +def supports_late_interaction(
998 + model: type[object],
999 +) -> TypeIs[type[SupportsLateInteraction]]: ...
1000 +
1001 +
1002 +@overload
1003 +def supports_late_interaction(model: object) -> TypeIs[SupportsLateInteraction]: ...
1004 +
1005 +
1006 +def _supports_late_interaction(
1007 + model: type[object] | object,
1008 +) -> TypeIs[type[SupportsLateInteraction]] | TypeIs[SupportsLateInteraction]:
1009 + return getattr(model, "supports_late_interaction", False)
1010 +
1011 +
1012 +def supports_late_interaction(
1013 + model: type[object] | object,
1014 +) -> TypeIs[type[SupportsLateInteraction]] | TypeIs[SupportsLateInteraction]:
1015 + return is_pooling_model(model) and _supports_late_interaction(model)
1016 +
1017 +
1018 class SupportsQuant:
1019 """The interface required for all models that support quantization."""
1020
vllm/model_executor/models/registry.py
+4
@@ -49,6 +49,7 @@ from .interfaces import (
49 is_hybrid,
50 requires_raw_input_tokens,
51 supports_cross_encoding,
52 + supports_late_interaction,
53 supports_mamba_prefix_caching,
54 supports_multimodal,
55 supports_multimodal_encoder_tp_data,
@@ -205,6 +206,7 @@ _EMBEDDING_MODELS = {
206 # [Text-only]
207 "BertModel": ("bert", "BertEmbeddingModel"),
208 "BertSpladeSparseEmbeddingModel": ("bert", "BertSpladeSparseEmbeddingModel"),
209 + "HF_ColBERT": ("colbert", "ColBERTModel"),
210 "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
211 "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
212 "Gemma3TextModel": ("gemma3", "Gemma3Model"),
@@ -593,6 +595,7 @@ class _ModelInfo:
595 default_seq_pooling_type: SequencePoolingType
596 default_tok_pooling_type: TokenPoolingType
597 supports_cross_encoding: bool
598 + supports_late_interaction: bool
599 supports_multimodal: bool
600 supports_multimodal_raw_input_only: bool
601 requires_raw_input_tokens: bool
@@ -616,6 +619,7 @@ class _ModelInfo:
619 default_tok_pooling_type=get_default_tok_pooling_type(model),
620 attn_type=get_attn_type(model),
621 supports_cross_encoding=supports_cross_encoding(model),
622 + supports_late_interaction=supports_late_interaction(model),
623 supports_multimodal=supports_multimodal(model),
624 supports_multimodal_raw_input_only=supports_multimodal_raw_input_only(
625 model