A high-throughput and memory-efficient inference and serving engine for LLMs
[Misc][Model] add io processor for query/document embeddings from ColBERT (jinaai/jina-colbert-v2) (#45210)
Signed-off-by: thomas <thomas.varghese@columbia.edu>
xx-thomas committed
Jun 15, 2026 at 20:32 UTC
b00e76ff72b0600ba9f4e4b3e0ce3d681de26b13
7 files changed
+479
.buildkite/test-amd.yaml
+5
@@ -608,6 +608,11 @@ steps:
608
- pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
609
- pip uninstall bge_m3_sparse_plugin -y
610
# END: `bge_m3_sparse io_processor` test
611
+ # BEGIN: `colbert_query io_processor` test
612
+ - pip install -e ./plugins/colbert_query_plugin
613
+ - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py
614
+ - pip uninstall colbert_query_plugin -y
615
+ # END: `colbert_query io_processor` test
616
# BEGIN: `stat_logger` plugins test
617
- pip install -e ./plugins/vllm_add_dummy_stat_logger
618
- pytest -v -s plugins_tests/test_stats_logger_plugins.py
.buildkite/test_areas/plugins.yaml
+4
@@ -27,6 +27,10 @@ steps:
27
- pip install -e ./plugins/bge_m3_sparse_plugin
28
- pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
29
- pip uninstall bge_m3_sparse_plugin -y
30
+ # test colbert_query io_processor plugin
31
+ - pip install -e ./plugins/colbert_query_plugin
32
+ - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py
33
+ - pip uninstall colbert_query_plugin -y
34
# end io_processor plugins test
35
# begin stat_logger plugins test
36
- pip install -e ./plugins/vllm_add_dummy_stat_logger
tests/plugins/colbert_query_plugin/colbert_query_processor/__init__.py
new
+6
@@ -0,0 +1,6 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+
5
+def register_colbert_query_embedding_processor():
6
+ return "colbert_query_processor.query_embedding_processor.ColBERTQueryEmbeddingProcessor" # noqa: E501
tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py
new
+194
@@ -0,0 +1,194 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Iterator, Sequence
4
+from typing import cast
5
+
6
+from vllm.config import VllmConfig
7
+from vllm.entrypoints.openai.engine.protocol import UsageInfo
8
+from vllm.inputs import PromptType, TokensPrompt
9
+from vllm.outputs import PoolingRequestOutput
10
+from vllm.plugins.io_processors.interface import IOProcessor
11
+from vllm.pooling_params import PoolingParams
12
+from vllm.renderers import BaseRenderer
13
+from vllm.utils.collection_utils import is_list_of
14
+
15
+from .types import (
16
+ QUERY_MAXLEN,
17
+ ColBERTEmbeddingCompletionRequestMixin,
18
+ ColBERTEmbeddingResponse,
19
+ ColBERTEmbeddingResponseData,
20
+)
21
+
22
+QUERY_MARKER_TOKEN = "[QueryMarker]"
23
+DOCUMENT_MARKER_TOKEN = "[DocumentMarker]"
24
+
25
+
26
+class ColBERTQueryEmbeddingProcessor(
27
+ IOProcessor[ColBERTEmbeddingCompletionRequestMixin, ColBERTEmbeddingResponse]
28
+):
29
+ """This IO processor only supports the ColBERT-style model jinaai/jina-colbert-v2.
30
+ It does not support all ColBERT-style variants (e.g. colbert-ir/colbertv2.0).
31
+ """
32
+
33
+ def __init__(self, vllm_config: VllmConfig, renderer: BaseRenderer):
34
+ super().__init__(vllm_config, renderer)
35
+ self.requests_cache: dict[str, ColBERTEmbeddingCompletionRequestMixin] = {}
36
+ self.renderer: BaseRenderer = renderer
37
+ # Context window (8192 for jinaai/jina-colbert-v2); caps document
38
+ # content length minus the 3 special-token slots.
39
+ self.max_model_len = vllm_config.model_config.max_model_len
40
+ self._query_marker_id: int | None = None
41
+ self._document_marker_id: int | None = None
42
+
43
+ def __repr__(self) -> str:
44
+ return (
45
+ f"ColBERTQueryEmbeddingProcessor("
46
+ f"query_maxlen={QUERY_MAXLEN}, "
47
+ f"doc_maxlen={self.max_model_len}, "
48
+ f"query_marker_token={QUERY_MARKER_TOKEN!r}, "
49
+ f"document_marker_token={DOCUMENT_MARKER_TOKEN!r})"
50
+ )
51
+
52
+ def _resolve_marker_ids(self, tokenizer) -> tuple[int, int]:
53
+ if self._query_marker_id is not None and self._document_marker_id is not None:
54
+ return self._query_marker_id, self._document_marker_id
55
+
56
+ unk_id = getattr(tokenizer, "unk_token_id", None)
57
+ marker_ids: list[int] = []
58
+ for marker in (QUERY_MARKER_TOKEN, DOCUMENT_MARKER_TOKEN):
59
+ marker_id = tokenizer.convert_tokens_to_ids(marker)
60
+ if marker_id is None or marker_id == unk_id:
61
+ raise ValueError(
62
+ f"Marker token {marker!r} not found in the tokenizer "
63
+ "vocabulary. This plugin requires a ColBERT model whose "
64
+ "tokenizer defines both "
65
+ f"{QUERY_MARKER_TOKEN!r} and {DOCUMENT_MARKER_TOKEN!r} "
66
+ "(e.g. jinaai/jina-colbert-v2)."
67
+ )
68
+ marker_ids.append(marker_id)
69
+
70
+ self._query_marker_id, self._document_marker_id = marker_ids
71
+ return self._query_marker_id, self._document_marker_id
72
+
73
+ def _iter_content_token_ids(
74
+ self,
75
+ tokenizer,
76
+ request_input: list[int] | list[list[int]] | str | list[str],
77
+ ) -> Iterator[list[int]]:
78
+ if isinstance(request_input, str):
79
+ yield tokenizer.encode(request_input, add_special_tokens=False)
80
+ return
81
+
82
+ if not isinstance(request_input, list) or not request_input:
83
+ raise ValueError("input must be a non-empty string or list")
84
+
85
+ if is_list_of(request_input, int):
86
+ yield list(cast(list[int], request_input))
87
+ return
88
+
89
+ for item in request_input:
90
+ if isinstance(item, str):
91
+ yield tokenizer.encode(item, add_special_tokens=False)
92
+ else:
93
+ yield list(cast(list[int], item))
94
+
95
+ def _build_query_prompt(
96
+ self,
97
+ tokenizer,
98
+ content_ids: list[int],
99
+ ) -> TokensPrompt:
100
+ """[CLS] [QueryMarker] <tokens> [SEP] [MASK]... up to QUERY_MAXLEN."""
101
+ query_marker_id, _ = self._resolve_marker_ids(tokenizer)
102
+ mask_token_id = tokenizer.mask_token_id
103
+ if mask_token_id is None:
104
+ raise ValueError(
105
+ "Tokenizer has no mask token; cannot perform query expansion."
106
+ )
107
+
108
+ # [CLS], marker and [SEP] take 3 slots.
109
+ content_ids = content_ids[: QUERY_MAXLEN - 3]
110
+ token_ids = [
111
+ tokenizer.cls_token_id,
112
+ query_marker_id,
113
+ *content_ids,
114
+ tokenizer.sep_token_id,
115
+ ]
116
+ token_ids += [mask_token_id] * (QUERY_MAXLEN - len(token_ids))
117
+ return TokensPrompt(prompt_token_ids=token_ids)
118
+
119
+ def _build_document_prompt(
120
+ self,
121
+ tokenizer,
122
+ content_ids: list[int],
123
+ ) -> TokensPrompt:
124
+ """[CLS] [DocumentMarker] <tokens> [SEP]"""
125
+ _, document_marker_id = self._resolve_marker_ids(tokenizer)
126
+
127
+ content_ids = content_ids[: self.max_model_len - 3]
128
+ token_ids = [
129
+ tokenizer.cls_token_id,
130
+ document_marker_id,
131
+ *content_ids,
132
+ tokenizer.sep_token_id,
133
+ ]
134
+ return TokensPrompt(prompt_token_ids=token_ids)
135
+
136
+ def parse_data(self, data: object) -> ColBERTEmbeddingCompletionRequestMixin:
137
+ if isinstance(data, dict):
138
+ return ColBERTEmbeddingCompletionRequestMixin(**data)
139
+ raise TypeError("request data should be a dictionary")
140
+
141
+ def pre_process(
142
+ self,
143
+ prompt: ColBERTEmbeddingCompletionRequestMixin,
144
+ request_id: str | None = None,
145
+ **kwargs,
146
+ ) -> PromptType | Sequence[PromptType]:
147
+ cache_key = request_id or "offline"
148
+ assert cache_key not in self.requests_cache, "request_id duplicated"
149
+ self.requests_cache[cache_key] = prompt
150
+
151
+ tokenizer = self.renderer.get_tokenizer()
152
+ prompts: list[TokensPrompt] = []
153
+ for content_ids in self._iter_content_token_ids(tokenizer, prompt.input):
154
+ if prompt.input_type == "query":
155
+ prompts.append(self._build_query_prompt(tokenizer, content_ids))
156
+ else:
157
+ prompts.append(self._build_document_prompt(tokenizer, content_ids))
158
+ return prompts
159
+
160
+ def merge_pooling_params(
161
+ self,
162
+ params: PoolingParams | None = None,
163
+ ) -> PoolingParams:
164
+ if params is None:
165
+ params = PoolingParams()
166
+ params.task = "token_embed"
167
+ params.skip_reading_prefix_cache = True
168
+ return params
169
+
170
+ def post_process(
171
+ self,
172
+ model_output: Sequence[PoolingRequestOutput],
173
+ request_id: str | None = None,
174
+ **kwargs,
175
+ ) -> ColBERTEmbeddingResponse:
176
+ raw_request = self.requests_cache.pop(request_id or "offline")
177
+
178
+ num_prompt_tokens = 0
179
+ response_data: list[ColBERTEmbeddingResponseData] = []
180
+ for idx, output in enumerate(model_output):
181
+ num_prompt_tokens += len(output.prompt_token_ids)
182
+ response_data.append(
183
+ ColBERTEmbeddingResponseData(
184
+ index=idx,
185
+ input_type=raw_request.input_type,
186
+ embedding=output.outputs.data.tolist(),
187
+ )
188
+ )
189
+
190
+ usage = UsageInfo(
191
+ prompt_tokens=num_prompt_tokens,
192
+ total_tokens=num_prompt_tokens,
193
+ )
194
+ return ColBERTEmbeddingResponse(data=response_data, usage=usage)
tests/plugins/colbert_query_plugin/colbert_query_processor/types.py
new
+33
@@ -0,0 +1,33 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Literal, get_args
5
+
6
+from pydantic import BaseModel, Field
7
+
8
+from vllm.entrypoints.openai.engine.protocol import UsageInfo
9
+from vllm.entrypoints.pooling.base.protocol import CompletionRequestMixin
10
+
11
+InputType = Literal["query", "document"]
12
+INPUT_TYPES: tuple[InputType, ...] = get_args(InputType)
13
+QUERY_MAXLEN = 32
14
+
15
+
16
+class ColBERTEmbeddingCompletionRequestMixin(CompletionRequestMixin):
17
+ input_type: InputType = Field(
18
+ description="Whether to encode the input as a ColBERT 'query' "
19
+ f"(query marker + [mask] expansion to {QUERY_MAXLEN} tokens) or as a "
20
+ "'document' (document marker only). Required.",
21
+ )
22
+
23
+
24
+class ColBERTEmbeddingResponseData(BaseModel):
25
+ index: int
26
+ object: str = "embedding"
27
+ input_type: InputType
28
+ embedding: list[list[float]]
29
+
30
+
31
+class ColBERTEmbeddingResponse(BaseModel):
32
+ data: list[ColBERTEmbeddingResponseData]
33
+ usage: UsageInfo
tests/plugins/colbert_query_plugin/setup.py
new
+15
@@ -0,0 +1,15 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from setuptools import setup
5
+
6
+setup(
7
+ name="colbert-query-plugin",
8
+ version="0.1",
9
+ packages=["colbert_query_processor"],
10
+ entry_points={
11
+ "vllm.io_processor_plugins": [
12
+ "colbert_query_plugin = colbert_query_processor:register_colbert_query_embedding_processor", # noqa: E501
13
+ ]
14
+ },
15
+)
tests/plugins_tests/test_colbert_query_io_processor_plugins.py
new
+222
@@ -0,0 +1,222 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import json
5
+from typing import TypedDict
6
+
7
+import pytest
8
+import requests
9
+
10
+from tests.utils import RemoteOpenAIServer
11
+from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse
12
+
13
+
14
+# Test configuration for ColBERT query plugin
15
+class ModelConfig(TypedDict):
16
+ model_name: str
17
+ plugin: str
18
+ query_input: str
19
+ document_input: str
20
+ hf_overrides: str
21
+ embedding_dim: int
22
+ query_maxlen: int
23
+
24
+
25
+model_config: ModelConfig = {
26
+ "model_name": "jinaai/jina-colbert-v2",
27
+ "plugin": "colbert_query_plugin",
28
+ "query_input": "What is machine learning?",
29
+ "document_input": "Machine learning is a subset of artificial intelligence.",
30
+ "hf_overrides": json.dumps({"architectures": ["ColBERTJinaRobertaModel"]}),
31
+ "embedding_dim": 128,
32
+ "query_maxlen": 32,
33
+}
34
+
35
+
36
+def _get_attr_or_val(obj: object | dict, key: str):
37
+ if isinstance(obj, dict) and key in obj:
38
+ return obj[key]
39
+ return getattr(obj, key, None)
40
+
41
+
42
+def _check_token_embeddings(entry, expected_input_type: str):
43
+ assert _get_attr_or_val(entry, "object") == "embedding"
44
+ assert _get_attr_or_val(entry, "input_type") == expected_input_type
45
+
46
+ embedding = _get_attr_or_val(entry, "embedding")
47
+ assert isinstance(embedding, list) and len(embedding) > 0
48
+ for token_embedding in embedding:
49
+ assert isinstance(token_embedding, list)
50
+ assert len(token_embedding) == model_config["embedding_dim"]
51
+ return embedding
52
+
53
+
54
+@pytest.fixture(scope="module")
55
+def server():
56
+ args = [
57
+ "--runner",
58
+ "pooling",
59
+ "--enforce-eager",
60
+ "--max-num-seqs",
61
+ "32",
62
+ "--trust-remote-code",
63
+ "--hf_overrides",
64
+ model_config["hf_overrides"],
65
+ "--io-processor-plugin",
66
+ model_config["plugin"],
67
+ ]
68
+
69
+ with RemoteOpenAIServer(model_config["model_name"], args) as remote_server:
70
+ yield remote_server
71
+
72
+
73
+def _post_pooling(server: RemoteOpenAIServer, data: dict):
74
+ request_payload = {
75
+ "model": model_config["model_name"],
76
+ "task": "plugin",
77
+ "data": data,
78
+ }
79
+ ret = requests.post(server.url_for("pooling"), json=request_payload)
80
+ ret.raise_for_status()
81
+ response = ret.json()
82
+ parsed_response = IOProcessorResponse(**response).data
83
+ assert parsed_response
84
+ return parsed_response
85
+
86
+
87
+def test_colbert_query_plugin_query_online(server: RemoteOpenAIServer):
88
+ """Queries are expanded to exactly query_maxlen token vectors."""
89
+ parsed_response = _post_pooling(
90
+ server, {"input": model_config["query_input"], "input_type": "query"}
91
+ )
92
+
93
+ data = _get_attr_or_val(parsed_response, "data")
94
+ assert len(data) == 1
95
+
96
+ embedding = _check_token_embeddings(data[0], "query")
97
+ assert len(embedding) == model_config["query_maxlen"]
98
+
99
+ usage = _get_attr_or_val(parsed_response, "usage")
100
+ assert _get_attr_or_val(usage, "prompt_tokens") == model_config["query_maxlen"]
101
+
102
+
103
+def test_colbert_query_plugin_document_online(server: RemoteOpenAIServer):
104
+ """Documents return one vector per token, with no mask expansion."""
105
+ parsed_response = _post_pooling(
106
+ server, {"input": model_config["document_input"], "input_type": "document"}
107
+ )
108
+
109
+ data = _get_attr_or_val(parsed_response, "data")
110
+ assert len(data) == 1
111
+
112
+ embedding = _check_token_embeddings(data[0], "document")
113
+ # No query expansion: number of vectors tracks the input length.
114
+ assert len(embedding) != model_config["query_maxlen"]
115
+
116
+ usage = _get_attr_or_val(parsed_response, "usage")
117
+ assert _get_attr_or_val(usage, "prompt_tokens") == len(embedding)
118
+
119
+
120
+def test_colbert_query_plugin_missing_input_type_online(server: RemoteOpenAIServer):
121
+ """input_type is required; omitting it is rejected."""
122
+ request_payload = {
123
+ "model": model_config["model_name"],
124
+ "task": "plugin",
125
+ "data": {"input": model_config["document_input"]},
126
+ }
127
+ ret = requests.post(server.url_for("pooling"), json=request_payload)
128
+ assert ret.status_code == 400
129
+
130
+
131
+def test_colbert_query_plugin_batch_online(server: RemoteOpenAIServer):
132
+ """A list input returns one entry per prompt."""
133
+ queries = ["What is machine learning?", "What is deep learning?"]
134
+ parsed_response = _post_pooling(server, {"input": queries, "input_type": "query"})
135
+
136
+ data = _get_attr_or_val(parsed_response, "data")
137
+ assert len(data) == len(queries)
138
+ for i, entry in enumerate(data):
139
+ assert _get_attr_or_val(entry, "index") == i
140
+ embedding = _check_token_embeddings(entry, "query")
141
+ assert len(embedding) == model_config["query_maxlen"]
142
+
143
+
144
+@pytest.mark.parametrize("input_type", ["query", "document"])
145
+def test_colbert_query_plugin_offline(vllm_runner, input_type: str):
146
+ """Test the ColBERT query plugin in offline mode."""
147
+ input_text = (
148
+ model_config["query_input"]
149
+ if input_type == "query"
150
+ else model_config["document_input"]
151
+ )
152
+ prompt = {
153
+ "data": {
154
+ "input": input_text,
155
+ "input_type": input_type,
156
+ }
157
+ }
158
+
159
+ with vllm_runner(
160
+ model_config["model_name"],
161
+ runner="pooling",
162
+ enforce_eager=True,
163
+ max_num_seqs=32,
164
+ trust_remote_code=True,
165
+ io_processor_plugin=model_config["plugin"],
166
+ hf_overrides=json.loads(model_config["hf_overrides"]),
167
+ default_torch_num_threads=1,
168
+ ) as llm_runner:
169
+ llm = llm_runner.get_llm()
170
+ pooler_output = llm.encode(prompt, pooling_task="plugin")
171
+
172
+ response = pooler_output[0].outputs
173
+ assert len(response.data) == 1
174
+
175
+ embedding = _check_token_embeddings(response.data[0], input_type)
176
+ if input_type == "query":
177
+ assert len(embedding) == model_config["query_maxlen"]
178
+ else:
179
+ assert len(embedding) != model_config["query_maxlen"]
180
+
181
+ assert response.usage.prompt_tokens == len(embedding)
182
+ assert response.usage.total_tokens == response.usage.prompt_tokens
183
+
184
+
185
+def test_colbert_query_plugin_offline_multiple_inputs(vllm_runner):
186
+ """Test the ColBERT query plugin with multiple inputs in offline mode."""
187
+ queries = [
188
+ "What is machine learning?",
189
+ "What is deep learning?",
190
+ "Why?",
191
+ ]
192
+ prompts = {
193
+ "data": {
194
+ "input": queries,
195
+ "input_type": "query",
196
+ }
197
+ }
198
+
199
+ with vllm_runner(
200
+ model_config["model_name"],
201
+ runner="pooling",
202
+ enforce_eager=True,
203
+ max_num_seqs=32,
204
+ trust_remote_code=True,
205
+ io_processor_plugin=model_config["plugin"],
206
+ hf_overrides=json.loads(model_config["hf_overrides"]),
207
+ default_torch_num_threads=1,
208
+ ) as llm_runner:
209
+ llm = llm_runner.get_llm()
210
+ pooler_output = llm.encode(prompts, pooling_task="plugin")
211
+
212
+ response = pooler_output[0].outputs
213
+ assert len(response.data) == len(queries)
214
+
215
+ for i, entry in enumerate(response.data):
216
+ assert entry.index == i
217
+ embedding = _check_token_embeddings(entry, "query")
218
+ assert len(embedding) == model_config["query_maxlen"]
219
+
220
+ expected_tokens = model_config["query_maxlen"] * len(queries)
221
+ assert response.usage.prompt_tokens == expected_tokens
222
+ assert response.usage.total_tokens == response.usage.prompt_tokens