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

[new model] add COLQwen3 code & Inference (#34398)

Signed-off-by: craftsangjae <craftsangjae@gmail.com> Signed-off-by: katacoder <craftsangjae@gmail.com>

Kata Coder committed Feb 14, 2026 at 13:15 UTC d1ea65d0a1c606ae041b73fd45ccd33980ca08e7
10 files changed +935
docs/models/pooling_models.md
+71
@@ -374,6 +374,77 @@ curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
374
375 An example can be found here: [examples/pooling/score/colbert_rerank_online.py](../../examples/pooling/score/colbert_rerank_online.py)
376
377 +### ColQwen3 Multi-Modal Late Interaction Models
378 +
379 +ColQwen3 is based on [ColPali](https://arxiv.org/abs/2407.01449), which extends ColBERT's late interaction approach to **multi-modal** inputs. While ColBERT operates on text-only token embeddings, ColPali/ColQwen3 can embed both **text and images** (e.g. PDF pages, screenshots, diagrams) into per-token L2-normalized vectors and compute relevance via MaxSim scoring. ColQwen3 specifically uses Qwen3-VL as its vision-language backbone.
380 +
381 +| Architecture | Backbone | Example HF Models |
382 +|---|---|---|
383 +| `ColQwen3` | Qwen3-VL | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` |
384 +| `OpsColQwen3Model` | Qwen3-VL | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` |
385 +
386 +Start the server:
387 +
388 +```shell
389 +vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 4096
390 +```
391 +
392 +Then you can use the rerank endpoint:
393 +
394 +```shell
395 +curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
396 + "model": "TomoroAI/tomoro-colqwen3-embed-4b",
397 + "query": "What is machine learning?",
398 + "documents": [
399 + "Machine learning is a subset of artificial intelligence.",
400 + "Python is a programming language.",
401 + "Deep learning uses neural networks."
402 + ]
403 +}'
404 +```
405 +
406 +Or the score endpoint:
407 +
408 +```shell
409 +curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
410 + "model": "TomoroAI/tomoro-colqwen3-embed-4b",
411 + "text_1": "What is the capital of France?",
412 + "text_2": ["The capital of France is Paris.", "Python is a programming language."]
413 +}'
414 +```
415 +
416 +You can also get the raw token embeddings using the pooling endpoint with `token_embed` task:
417 +
418 +```shell
419 +curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
420 + "model": "TomoroAI/tomoro-colqwen3-embed-4b",
421 + "input": "What is machine learning?",
422 + "task": "token_embed"
423 +}'
424 +```
425 +
426 +For **image inputs**, use the chat-style `messages` field so that the vLLM multimodal processor handles them correctly:
427 +
428 +```shell
429 +curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
430 + "model": "TomoroAI/tomoro-colqwen3-embed-4b",
431 + "messages": [
432 + {
433 + "role": "user",
434 + "content": [
435 + {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
436 + {"type": "text", "text": "Describe the image."}
437 + ]
438 + }
439 + ]
440 +}'
441 +```
442 +
443 +Examples can be found here:
444 +
445 +- Multi-vector retrieval: [examples/pooling/token_embed/colqwen3_token_embed_online.py](../../examples/pooling/token_embed/colqwen3_token_embed_online.py)
446 +- Reranking: [examples/pooling/score/colqwen3_rerank_online.py](../../examples/pooling/score/colqwen3_rerank_online.py)
447 +
448 ### BAAI/bge-m3
449
450 The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json`
examples/pooling/score/colqwen3_rerank_online.py new
+130
@@ -0,0 +1,130 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Example of using ColQwen3 late interaction model for reranking.
5 +
6 +ColQwen3 is a multi-modal ColBERT-style model based on Qwen3-VL.
7 +It produces per-token embeddings and uses MaxSim scoring for retrieval
8 +and reranking. Supports both text and image inputs.
9 +
10 +Start the server with:
11 + vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 50000
12 +
13 +Then run this script:
14 + python colqwen3_rerank_online.py
15 +"""
16 +
17 +import requests
18 +
19 +MODEL = "TomoroAI/tomoro-colqwen3-embed-4b"
20 +BASE_URL = "http://127.0.0.1:8000"
21 +
22 +headers = {"accept": "application/json", "Content-Type": "application/json"}
23 +
24 +
25 +def rerank_text():
26 + """Text-only reranking via /rerank endpoint."""
27 + print("=" * 60)
28 + print("1. Text reranking (/rerank)")
29 + print("=" * 60)
30 +
31 + data = {
32 + "model": MODEL,
33 + "query": "What is machine learning?",
34 + "documents": [
35 + "Machine learning is a subset of artificial intelligence.",
36 + "Python is a programming language.",
37 + "Deep learning uses neural networks for complex tasks.",
38 + "The weather today is sunny.",
39 + ],
40 + }
41 +
42 + response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
43 +
44 + if response.status_code == 200:
45 + result = response.json()
46 + print("\n Ranked documents (most relevant first):")
47 + for item in result["results"]:
48 + doc_idx = item["index"]
49 + score = item["relevance_score"]
50 + print(f" [{score:.4f}] {data['documents'][doc_idx]}")
51 + else:
52 + print(f" Request failed: {response.status_code}")
53 + print(f" {response.text[:300]}")
54 +
55 +
56 +def score_text():
57 + """Text-only scoring via /score endpoint."""
58 + print()
59 + print("=" * 60)
60 + print("2. Text scoring (/score)")
61 + print("=" * 60)
62 +
63 + query = "What is the capital of France?"
64 + documents = [
65 + "The capital of France is Paris.",
66 + "Berlin is the capital of Germany.",
67 + "Python is a programming language.",
68 + ]
69 +
70 + data = {
71 + "model": MODEL,
72 + "text_1": query,
73 + "text_2": documents,
74 + }
75 +
76 + response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
77 +
78 + if response.status_code == 200:
79 + result = response.json()
80 + print(f"\n Query: {query}\n")
81 + for item in result["data"]:
82 + idx = item["index"]
83 + score = item["score"]
84 + print(f" Doc {idx} (score={score:.4f}): {documents[idx]}")
85 + else:
86 + print(f" Request failed: {response.status_code}")
87 + print(f" {response.text[:300]}")
88 +
89 +
90 +def score_text_top_n():
91 + """Text reranking with top_n filtering via /rerank endpoint."""
92 + print()
93 + print("=" * 60)
94 + print("3. Text reranking with top_n=2 (/rerank)")
95 + print("=" * 60)
96 +
97 + data = {
98 + "model": MODEL,
99 + "query": "What is the capital of France?",
100 + "documents": [
101 + "The capital of France is Paris.",
102 + "Berlin is the capital of Germany.",
103 + "Python is a programming language.",
104 + "The Eiffel Tower is in Paris.",
105 + ],
106 + "top_n": 2,
107 + }
108 +
109 + response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
110 +
111 + if response.status_code == 200:
112 + result = response.json()
113 + print(f"\n Top {data['top_n']} results:")
114 + for item in result["results"]:
115 + doc_idx = item["index"]
116 + score = item["relevance_score"]
117 + print(f" [{score:.4f}] {data['documents'][doc_idx]}")
118 + else:
119 + print(f" Request failed: {response.status_code}")
120 + print(f" {response.text[:300]}")
121 +
122 +
123 +def main():
124 + rerank_text()
125 + score_text()
126 + score_text_top_n()
127 +
128 +
129 +if __name__ == "__main__":
130 + main()
examples/pooling/token_embed/colqwen3_token_embed_online.py new
+198
@@ -0,0 +1,198 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# ruff: noqa: E501
4 +
5 +"""
6 +Example online usage of Pooling API for ColQwen3 multi-vector retrieval.
7 +
8 +ColQwen3 is a multi-modal late interaction model based on Qwen3-VL that
9 +produces per-token embeddings (320-dim, L2-normalized) for both text and
10 +image inputs. Similarity is computed via MaxSim scoring.
11 +
12 +This example mirrors the official TomoroAI inference code
13 +(https://huggingface.co/TomoroAI/tomoro-colqwen3-embed-4b) but uses the
14 +vLLM serving API instead of local HuggingFace model loading.
15 +
16 +Start the server with:
17 + vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 4096
18 +
19 +Then run this script:
20 + python colqwen3_token_embed_online.py
21 +"""
22 +
23 +import argparse
24 +import base64
25 +from io import BytesIO
26 +
27 +import numpy as np
28 +import requests
29 +from PIL import Image
30 +
31 +# ── Helpers ─────────────────────────────────────────────────
32 +
33 +
34 +def post_http_request(payload: dict, api_url: str) -> requests.Response:
35 + headers = {"User-Agent": "Test Client"}
36 + return requests.post(api_url, headers=headers, json=payload)
37 +
38 +
39 +def load_image(url: str) -> Image.Image:
40 + """Download an image from URL (handles Wikimedia 403)."""
41 + for hdrs in ({}, {"User-Agent": "Mozilla/5.0 (compatible; ColQwen3-demo/1.0)"}):
42 + resp = requests.get(url, headers=hdrs, timeout=10)
43 + if resp.status_code == 403:
44 + continue
45 + resp.raise_for_status()
46 + return Image.open(BytesIO(resp.content)).convert("RGB")
47 + raise RuntimeError(f"Could not fetch image from {url}")
48 +
49 +
50 +def encode_image_base64(image: Image.Image) -> str:
51 + """Encode a PIL image to a base64 data URI."""
52 + buf = BytesIO()
53 + image.save(buf, format="PNG")
54 + return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
55 +
56 +
57 +def compute_maxsim(q_emb: np.ndarray, d_emb: np.ndarray) -> float:
58 + """Compute ColBERT-style MaxSim score between query and document."""
59 + sim = q_emb @ d_emb.T
60 + return float(sim.max(axis=-1).sum())
61 +
62 +
63 +# ── Encode functions ────────────────────────────────────────
64 +
65 +
66 +def encode_queries(texts: list[str], model: str, api_url: str) -> list[np.ndarray]:
67 + """Encode text queries → list of multi-vector embeddings."""
68 + resp = post_http_request({"model": model, "input": texts}, api_url)
69 + return [np.array(item["data"]) for item in resp.json()["data"]]
70 +
71 +
72 +def encode_images(image_urls: list[str], model: str, api_url: str) -> list[np.ndarray]:
73 + """Encode image documents → list of multi-vector embeddings.
74 +
75 + Images are sent via the chat-style `messages` field so that the
76 + vLLM multimodal processor handles them correctly.
77 + """
78 + embeddings = []
79 + for url in image_urls:
80 + print(f" Loading: {url.split('/')[-1]}...")
81 + image = load_image(url)
82 + image_uri = encode_image_base64(image)
83 + resp = post_http_request(
84 + {
85 + "model": model,
86 + "messages": [
87 + {
88 + "role": "user",
89 + "content": [
90 + {"type": "image_url", "image_url": {"url": image_uri}},
91 + {"type": "text", "text": "Describe the image."},
92 + ],
93 + }
94 + ],
95 + },
96 + api_url,
97 + )
98 + result = resp.json()
99 + if resp.status_code != 200 or "data" not in result:
100 + print(f" Error ({resp.status_code}): {str(result)[:200]}")
101 + continue
102 + embeddings.append(np.array(result["data"][0]["data"]))
103 + return embeddings
104 +
105 +
106 +# ── Main ────────────────────────────────────────────────────
107 +
108 +
109 +def parse_args():
110 + parser = argparse.ArgumentParser()
111 + parser.add_argument("--host", type=str, default="localhost")
112 + parser.add_argument("--port", type=int, default=8000)
113 + parser.add_argument(
114 + "--model",
115 + type=str,
116 + default="TomoroAI/tomoro-colqwen3-embed-4b",
117 + )
118 + return parser.parse_args()
119 +
120 +
121 +def main(args):
122 + pooling_url = f"http://{args.host}:{args.port}/pooling"
123 + score_url = f"http://{args.host}:{args.port}/score"
124 + model = args.model
125 +
126 + # Same sample data as the official TomoroAI example
127 + queries = [
128 + "Retrieve the city of Singapore",
129 + "Retrieve the city of Beijing",
130 + "Retrieve the city of London",
131 + ]
132 + image_urls = [
133 + "https://upload.wikimedia.org/wikipedia/commons/2/27/Singapore_skyline_2022.jpg",
134 + "https://upload.wikimedia.org/wikipedia/commons/6/61/Beijing_skyline_at_night.JPG",
135 + "https://upload.wikimedia.org/wikipedia/commons/4/49/London_skyline.jpg",
136 + ]
137 +
138 + # ── 1) Text query embeddings ────────────────────────────
139 + print("=" * 60)
140 + print("1. Encode text queries (multi-vector)")
141 + print("=" * 60)
142 + query_embeddings = encode_queries(queries, model, pooling_url)
143 + for i, emb in enumerate(query_embeddings):
144 + norm = float(np.linalg.norm(emb[0]))
145 + print(f' Query {i}: {emb.shape} (L2 norm: {norm:.4f}) "{queries[i]}"')
146 +
147 + # ── 2) Image document embeddings ────────────────────────
148 + print()
149 + print("=" * 60)
150 + print("2. Encode image documents (multi-vector)")
151 + print("=" * 60)
152 + doc_embeddings = encode_images(image_urls, model, pooling_url)
153 + for i, emb in enumerate(doc_embeddings):
154 + print(f" Doc {i}: {emb.shape} {image_urls[i].split('/')[-1]}")
155 +
156 + # ── 3) Cross-modal MaxSim scoring ───────────────────────
157 + if doc_embeddings:
158 + print()
159 + print("=" * 60)
160 + print("3. Cross-modal MaxSim scores (text queries × image docs)")
161 + print("=" * 60)
162 + # Header
163 + print(f"{'':>35s}", end="")
164 + for j in range(len(doc_embeddings)):
165 + print(f" Doc {j:>2d}", end="")
166 + print()
167 + # Score matrix
168 + for i, q_emb in enumerate(query_embeddings):
169 + print(f" {queries[i]:<33s}", end="")
170 + for j, d_emb in enumerate(doc_embeddings):
171 + score = compute_maxsim(q_emb, d_emb)
172 + print(f" {score:6.2f}", end="")
173 + print()
174 +
175 + # ── 4) Text-only /score endpoint ────────────────────────
176 + print()
177 + print("=" * 60)
178 + print("4. Text-only late interaction scoring (/score endpoint)")
179 + print("=" * 60)
180 + text_query = "What is the capital of France?"
181 + text_docs = [
182 + "The capital of France is Paris.",
183 + "Berlin is the capital of Germany.",
184 + "Python is a programming language.",
185 + ]
186 + resp = post_http_request(
187 + {"model": model, "text_1": text_query, "text_2": text_docs},
188 + score_url,
189 + )
190 + print(f' Query: "{text_query}"\n')
191 + for item in resp.json()["data"]:
192 + idx = item["index"]
193 + print(f" Doc {idx} (score={item['score']:.4f}): {text_docs[idx]}")
194 +
195 +
196 +if __name__ == "__main__":
197 + args = parse_args()
198 + main(args)
tests/models/multimodal/pooling/test_colqwen3.py new
+156
@@ -0,0 +1,156 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for ColQwen3 late interaction model for multi-modal retrieval.
4 +
5 +ColQwen3 is a multi-vector retrieval model based on Qwen3-VL backbone with
6 +ColBERT-style late interaction scoring (MaxSim). It produces per-token
7 +embeddings for both text and image inputs.
8 +"""
9 +
10 +import pytest
11 +import torch
12 +
13 +from ....conftest import VllmRunner
14 +
15 +MODELS = [
16 + "TomoroAI/tomoro-colqwen3-embed-4b",
17 + "OpenSearch-AI/Ops-Colqwen3-4B",
18 +]
19 +
20 +EMBED_DIMS = {
21 + "TomoroAI/tomoro-colqwen3-embed-4b": 320,
22 + "OpenSearch-AI/Ops-Colqwen3-4B": 2560,
23 +}
24 +
25 +TEXT_QUERIES = [
26 + "What is the capital of France?",
27 + "Describe the contents of the document.",
28 +]
29 +
30 +TEXT_DOCUMENTS = [
31 + "The capital of France is Paris.",
32 + "This document contains important financial data.",
33 +]
34 +
35 +DTYPE = "half"
36 +
37 +
38 +def _run_token_embed_test(
39 + vllm_runner: type[VllmRunner],
40 + model: str,
41 + *,
42 + dtype: str,
43 +) -> None:
44 + """Verify per-token embedding shape and L2 normalization."""
45 + with vllm_runner(
46 + model,
47 + runner="pooling",
48 + dtype=dtype,
49 + max_model_len=4096,
50 + enforce_eager=True,
51 + ) as vllm_model:
52 + outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
53 +
54 + assert len(outputs) == 1
55 + emb = torch.tensor(outputs[0])
56 + # Token embeddings should be 2D: [num_tokens, embed_dim]
57 + assert emb.dim() == 2
58 + assert emb.shape[1] == EMBED_DIMS[model]
59 + assert emb.shape[0] > 1
60 +
61 + # Verify L2 normalization
62 + norms = torch.norm(emb, p=2, dim=-1)
63 + torch.testing.assert_close(
64 + norms,
65 + torch.ones_like(norms),
66 + rtol=1e-2,
67 + atol=1e-2,
68 + )
69 +
70 +
71 +def _run_late_interaction_test(
72 + vllm_runner: type[VllmRunner],
73 + model: str,
74 + *,
75 + dtype: str,
76 +) -> None:
77 + """Verify MaxSim scoring matches manual computation."""
78 + from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
79 +
80 + with vllm_runner(
81 + model,
82 + runner="pooling",
83 + dtype=dtype,
84 + max_model_len=4096,
85 + enforce_eager=True,
86 + ) as vllm_model:
87 + q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
88 + d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
89 +
90 + q_emb = torch.tensor(q_outputs[0])
91 + d_emb = torch.tensor(d_outputs[0])
92 +
93 + manual_score = compute_maxsim_score(q_emb, d_emb).item()
94 +
95 + vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0])
96 +
97 + assert len(vllm_scores) == 1
98 + assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
99 +
100 +
101 +def _run_relevance_test(
102 + vllm_runner: type[VllmRunner],
103 + model: str,
104 + *,
105 + dtype: str,
106 +) -> None:
107 + """Verify that relevant documents score higher than irrelevant ones."""
108 + query = "What is machine learning?"
109 + documents = [
110 + "Machine learning is a subset of artificial intelligence.",
111 + "The weather forecast shows rain tomorrow.",
112 + "Deep learning uses neural networks for complex tasks.",
113 + ]
114 +
115 + with vllm_runner(
116 + model,
117 + runner="pooling",
118 + dtype=dtype,
119 + max_model_len=4096,
120 + enforce_eager=True,
121 + ) as vllm_model:
122 + scores = vllm_model.score(query, documents)
123 +
124 + assert len(scores) == 3
125 + assert scores[0] > scores[1], "ML doc should score higher than weather doc"
126 + assert scores[2] > scores[1], "DL doc should score higher than weather doc"
127 +
128 +
129 +@pytest.mark.parametrize("model", MODELS)
130 +@pytest.mark.parametrize("dtype", [DTYPE])
131 +def test_colqwen3_token_embed(
132 + vllm_runner,
133 + model: str,
134 + dtype: str,
135 +) -> None:
136 + _run_token_embed_test(vllm_runner, model, dtype=dtype)
137 +
138 +
139 +@pytest.mark.parametrize("model", MODELS)
140 +@pytest.mark.parametrize("dtype", [DTYPE])
141 +def test_colqwen3_late_interaction_scoring(
142 + vllm_runner,
143 + model: str,
144 + dtype: str,
145 +) -> None:
146 + _run_late_interaction_test(vllm_runner, model, dtype=dtype)
147 +
148 +
149 +@pytest.mark.parametrize("model", MODELS)
150 +@pytest.mark.parametrize("dtype", [DTYPE])
151 +def test_colqwen3_relevance_ordering(
152 + vllm_runner,
153 + model: str,
154 + dtype: str,
155 +) -> None:
156 + _run_relevance_test(vllm_runner, model, dtype=dtype)
tests/models/registry.py
+6
@@ -597,6 +597,12 @@ _EMBEDDING_EXAMPLE_MODELS = {
597 "TIGER-Lab/VLM2Vec-Full", trust_remote_code=True
598 ),
599 "Qwen2VLForConditionalGeneration": _HfExamplesInfo("MrLight/dse-qwen2-2b-mrl-v1"),
600 + "ColQwen3": _HfExamplesInfo(
601 + "TomoroAI/tomoro-colqwen3-embed-4b", trust_remote_code=True
602 + ),
603 + "OpsColQwen3Model": _HfExamplesInfo(
604 + "OpenSearch-AI/Ops-Colqwen3-4B", trust_remote_code=True
605 + ),
606 "SiglipModel": _HfExamplesInfo("google/siglip-base-patch16-224"),
607 "PrithviGeoSpatialMAE": _HfExamplesInfo(
608 "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
vllm/model_executor/models/colqwen3.py new
+306
@@ -0,0 +1,306 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +ColQwen3 late interaction model for multi-modal retrieval and reranking.
5 +
6 +ColQwen3 extends Qwen3-VL with a ColBERT-style late interaction head,
7 +producing per-token embeddings for both text and image inputs. It uses
8 +MaxSim scoring for retrieval/reranking tasks.
9 +
10 +This model supports the "token_embed" pooling task and is designed for
11 +multi-vector retrieval of documents containing both text and images.
12 +
13 +Reference: https://arxiv.org/abs/2407.01449 (ColPali)
14 +Based on: Qwen3-VL backbone with custom text projection
15 +
16 +Target models:
17 +- TomoroAI/tomoro-colqwen3-embed-8b
18 +- OpenSearch-AI/Ops-Colqwen3-4B
19 +"""
20 +
21 +from collections.abc import Iterable, Mapping
22 +from typing import ClassVar, Literal
23 +
24 +import torch
25 +import torch.nn as nn
26 +from transformers.models.qwen3_vl import Qwen3VLProcessor
27 +
28 +from vllm.config import VllmConfig
29 +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_base import default_pooling_type
34 +from .qwen2_vl import Qwen2VLMultiModalDataParser
35 +from .qwen3_vl import (
36 + Qwen3VLDummyInputsBuilder,
37 + Qwen3VLForConditionalGeneration,
38 + Qwen3VLMultiModalProcessor,
39 + Qwen3VLProcessingInfo,
40 +)
41 +from .utils import AutoWeightsLoader, WeightsMapper
42 +
43 +
44 +class ColQwen3ProcessingInfo(Qwen3VLProcessingInfo):
45 + """Processing info for ColQwen3 models.
46 +
47 + ColQwen3 models (TomoroAI, OpenSearch-AI, etc.) use custom HuggingFace
48 + configs (e.g. ColQwen3Config, OpsColQwen3Config) that are not instances
49 + of Qwen3VLConfig. We override get_hf_config() and get_hf_processor()
50 + to skip the strict type check, similar to OpenCUAProcessingInfo.
51 + """
52 +
53 + def get_hf_config(self):
54 + return self.ctx.get_hf_config()
55 +
56 + def get_hf_processor(self, **kwargs: object) -> Qwen3VLProcessor:
57 + # Force standard Qwen3VLProcessor even when trust_remote_code=True.
58 + # ColQwen3 custom processors (e.g. ColQwen3Processor) have
59 + # incompatible interfaces with vLLM's Qwen3VLMultiModalProcessor.
60 + # The standard Qwen3VLProcessor handles both text and image inputs
61 + # correctly for the Qwen3-VL backbone.
62 + return self.ctx.get_hf_processor(
63 + Qwen3VLProcessor,
64 + use_fast=kwargs.pop("use_fast", True),
65 + **kwargs,
66 + )
67 +
68 + @property
69 + def _supports_video(self) -> bool:
70 + """Check if the HF processor supports video inputs."""
71 + return hasattr(self.get_hf_processor(), "video_processor")
72 +
73 + def get_video_processor(self, **kwargs: object):
74 + if not self._supports_video:
75 + raise AttributeError(
76 + f"The processor for {self.ctx.model_config.model} does not "
77 + "support video inputs (no video_processor attribute)."
78 + )
79 + return self.get_hf_processor(**kwargs).video_processor # type: ignore[attr-defined]
80 +
81 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
82 + limits: dict[str, int | None] = {"image": None}
83 + if self._supports_video:
84 + limits["video"] = None
85 + return limits
86 +
87 + def get_mm_max_tokens_per_item(
88 + self,
89 + seq_len: int,
90 + mm_counts: Mapping[str, int],
91 + ) -> Mapping[str, int]:
92 + max_image_tokens = self.get_max_image_tokens()
93 + result: dict[str, int] = {"image": max_image_tokens}
94 + if self._supports_video:
95 + max_video_tokens = self.get_max_video_tokens(seq_len, mm_counts)
96 + result["video"] = max_video_tokens
97 + return result
98 +
99 + def get_data_parser(self):
100 + hf_config = self.get_hf_config()
101 + spatial_merge_size = hf_config.vision_config.spatial_merge_size
102 + return Qwen2VLMultiModalDataParser(
103 + spatial_merge_size,
104 + video_needs_metadata=self._supports_video,
105 + expected_hidden_size=self._get_expected_hidden_size(),
106 + )
107 +
108 +
109 +@default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL")
110 +@MULTIMODAL_REGISTRY.register_processor(
111 + Qwen3VLMultiModalProcessor,
112 + info=ColQwen3ProcessingInfo,
113 + dummy_inputs=Qwen3VLDummyInputsBuilder,
114 +)
115 +class ColQwen3Model(
116 + Qwen3VLForConditionalGeneration,
117 +):
118 + """ColQwen3 late interaction model for multi-modal retrieval/reranking.
119 +
120 + This model extends Qwen3VLForConditionalGeneration with a ColBERT-style
121 + linear projection layer for per-token embeddings. It supports:
122 + - "token_embed" task: Per-token embeddings for late interaction scoring
123 +
124 + The model produces L2-normalized per-token embeddings by:
125 + 1. Running the Qwen3-VL backbone (vision + language) to get hidden states
126 + 2. Projecting hidden states through a linear layer (hidden_size -> embed_dim)
127 + 3. L2-normalizing the projected embeddings
128 +
129 + ColBERT-style MaxSim scoring is computed externally, either client-side
130 + or via the late interaction scoring path in ServingScores.
131 +
132 + Attributes:
133 + custom_text_proj: Linear projection from hidden_size to embed_dim
134 + supports_late_interaction: Flag indicating this model uses late
135 + interaction scoring
136 + """
137 +
138 + # Mark this as a pooling model so vLLM routes to pooler path
139 + is_pooling_model = True
140 +
141 + # Mark this model as supporting late interaction scoring
142 + supports_late_interaction: ClassVar[Literal[True]] = True
143 +
144 + # Override hf_to_vllm_mapper to handle ColQwen3 weight naming.
145 + # NOTE: WeightsMapper applies ALL matching prefix rules sequentially
146 + # (no early exit), so more-specific prefixes must come first.
147 + # TomoroAI: "vlm.model.visual.", "vlm.model.language_model."
148 + # ColPali: "model.visual.", "model.language_model."
149 + # OpenSearch: "visual.", "language_model." (no outer prefix,
150 + # re-prefixed to "model.*" in load_weights)
151 + hf_to_vllm_mapper = WeightsMapper(
152 + orig_to_new_prefix={
153 + # TomoroAI naming convention (most specific first)
154 + "vlm.model.visual.": "visual.",
155 + "vlm.lm_head.": "language_model.lm_head.",
156 + "vlm.model.language_model.": "language_model.model.",
157 + # ColPali / nvidia naming convention
158 + "model.visual.": "visual.",
159 + "lm_head.": "language_model.lm_head.",
160 + # OpenSearch-AI: after re-prefix, "language_model.model.*"
161 + # becomes "model.language_model.model.*" — handle this before
162 + # the shorter "model.language_model." rule to avoid double map
163 + "model.language_model.model.": "language_model.model.",
164 + "model.language_model.": "language_model.model.",
165 + }
166 + )
167 +
168 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
169 + super().__init__(vllm_config=vllm_config, prefix=prefix)
170 +
171 + config = vllm_config.model_config.hf_config
172 + head_dtype = vllm_config.model_config.head_dtype
173 +
174 + hidden_size = getattr(config, "hidden_size", None)
175 + if hidden_size is None and hasattr(config, "text_config"):
176 + hidden_size = config.text_config.hidden_size
177 + if hidden_size is None:
178 + raise ValueError(
179 + "Unable to determine text hidden size from config. "
180 + "Expected 'hidden_size' or 'text_config.hidden_size'."
181 + )
182 + self._proj_hidden_size = hidden_size
183 +
184 + # (TomoroAI: embed_dim, OpenSearch: dims, ColPali: dim)
185 + self.embed_dim: int | None = (
186 + getattr(config, "embed_dim", None)
187 + or getattr(config, "dims", None)
188 + or getattr(config, "dim", None)
189 + or getattr(config, "projection_dim", None)
190 + or getattr(config, "colbert_dim", None)
191 + )
192 +
193 + # Build the projection layer if embed_dim is known
194 + if self.embed_dim is not None:
195 + self.custom_text_proj = nn.Linear(
196 + hidden_size,
197 + self.embed_dim,
198 + bias=False,
199 + dtype=head_dtype,
200 + )
201 + else:
202 + # Will be created during load_weights when dim is inferred
203 + self.custom_text_proj = None
204 +
205 + pooler_config = vllm_config.model_config.pooler_config
206 + assert pooler_config is not None
207 + self.pooler = pooler_for_token_embed(
208 + pooler_config,
209 + projector=None,
210 + )
211 +
212 + def forward(
213 + self,
214 + input_ids: torch.Tensor | None,
215 + positions: torch.Tensor,
216 + intermediate_tensors=None,
217 + inputs_embeds: torch.Tensor | None = None,
218 + **kwargs: object,
219 + ) -> torch.Tensor:
220 + """Run forward pass producing per-token embeddings."""
221 + hidden_states = super().forward(
222 + input_ids=input_ids,
223 + positions=positions,
224 + intermediate_tensors=intermediate_tensors,
225 + inputs_embeds=inputs_embeds,
226 + **kwargs,
227 + )
228 +
229 + if not isinstance(hidden_states, torch.Tensor):
230 + return hidden_states # type: ignore
231 +
232 + proj_dtype = self.custom_text_proj.weight.dtype # type: ignore
233 + if hidden_states.dtype != proj_dtype:
234 + hidden_states = hidden_states.to(proj_dtype)
235 +
236 + # Project to embedding dimension and L2 normalize
237 + proj = self.custom_text_proj(hidden_states) # type: ignore
238 + return torch.nn.functional.normalize(proj, p=2, dim=-1)
239 +
240 + # Names used for the projection layer across different ColQwen3 variants
241 + _PROJ_LAYER_NAMES = {
242 + "custom_text_proj", # ColPali naming
243 + "embedding_proj_layer", # TomoroAI naming
244 + }
245 +
246 + def _is_proj_weight(self, name: str) -> bool:
247 + """Check if a weight name belongs to the projection layer."""
248 + return any(proj_name in name for proj_name in self._PROJ_LAYER_NAMES)
249 +
250 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
251 + """Load weights with special handling for ColQwen3 projection layer."""
252 + weights_list = list(weights)
253 + proj_weights: list[tuple[str, torch.Tensor]] = []
254 + model_weights: list[tuple[str, torch.Tensor]] = []
255 +
256 + # Scan all weight names to determine if re-prefixing is needed.
257 + # OpenSearch-AI models have unprefixed weights ("language_model.*",
258 + # "visual.*") that need "model." added so hf_to_vllm_mapper can
259 + # process them. Only re-prefix if ALL backbone weights are
260 + # unprefixed (no "vlm." or "model." prefix found).
261 + has_unprefixed = any(
262 + name.startswith("language_model.") or name.startswith("visual.")
263 + for name, _ in weights_list
264 + )
265 + has_prefixed = any(
266 + name.startswith("vlm.") or name.startswith("model.")
267 + for name, _ in weights_list
268 + )
269 + needs_reprefix = has_unprefixed and not has_prefixed
270 +
271 + for name, weight in weights_list:
272 + if self._is_proj_weight(name):
273 + proj_weights.append((name, weight))
274 + else:
275 + if needs_reprefix and not self._is_proj_weight(name):
276 + name = "model." + name
277 + model_weights.append((name, weight))
278 +
279 + loader = AutoWeightsLoader(self)
280 + loaded = loader.load_weights(model_weights, mapper=self.hf_to_vllm_mapper)
281 +
282 + if proj_weights:
283 + model_dtype = next(self.language_model.parameters()).dtype
284 + model_device = next(self.language_model.parameters()).device
285 +
286 + for name, weight in proj_weights:
287 + if self.embed_dim is None and "weight" in name:
288 + self.embed_dim = weight.shape[0]
289 + has_bias = any("bias" in n for n, _ in proj_weights)
290 + self.custom_text_proj = nn.Linear(
291 + self._proj_hidden_size,
292 + self.embed_dim,
293 + bias=has_bias,
294 + dtype=model_dtype,
295 + )
296 + self.custom_text_proj.to(model_device)
297 +
298 + if self.custom_text_proj is not None:
299 + param_name = name.split(".")[-1]
300 + param = getattr(self.custom_text_proj, param_name, None)
301 + if param is not None:
302 + weight = weight.to(device=param.device, dtype=param.dtype)
303 + default_weight_loader(param, weight)
304 + loaded.add(f"custom_text_proj.{param_name}")
305 +
306 + return loaded
vllm/model_executor/models/registry.py
+2
@@ -254,6 +254,8 @@ _EMBEDDING_MODELS = {
254 ),
255 "Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
256 "Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"), # noqa: E501
257 + "ColQwen3": ("colqwen3", "ColQwen3Model"),
258 + "OpsColQwen3Model": ("colqwen3", "ColQwen3Model"),
259 "SiglipModel": ("siglip", "SiglipEmbeddingModel"),
260 # Technically Terratorch models work on images, both in
261 # input and output. I am adding it here because it piggy-backs on embedding
vllm/transformers_utils/config.py
+2
@@ -74,6 +74,8 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
74 afmoe="AfmoeConfig",
75 bagel="BagelConfig",
76 chatglm="ChatGLMConfig",
77 + colqwen3="ColQwen3Config",
78 + ops_colqwen3="OpsColQwen3Config",
79 deepseek_vl_v2="DeepseekVLV2Config",
80 deepseek_v32="DeepseekV3Config",
81 flex_olmo="FlexOlmoConfig",
vllm/transformers_utils/configs/__init__.py
+6
@@ -18,6 +18,9 @@ _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 + "ColQwen3Config": "vllm.transformers_utils.configs.colqwen3",
22 + "OpsColQwen3Config": "vllm.transformers_utils.configs.colqwen3",
23 + "Qwen3VLNemotronEmbedConfig": "vllm.transformers_utils.configs.colqwen3",
24 "DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2",
25 "DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr",
26 "EAGLEConfig": "vllm.transformers_utils.configs.eagle",
@@ -68,6 +71,9 @@ __all__ = [
71 "AfmoeConfig",
72 "BagelConfig",
73 "ChatGLMConfig",
74 + "ColQwen3Config",
75 + "OpsColQwen3Config",
76 + "Qwen3VLNemotronEmbedConfig",
77 "DeepseekVLV2Config",
78 "DeepseekV3Config",
79 "DotsOCRConfig",
vllm/transformers_utils/configs/colqwen3.py new
+58
@@ -0,0 +1,58 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +ColQwen3 configuration that extends Qwen3VLConfig with embedding projection
5 +fields. This allows ColQwen3 models to be loaded without trust_remote_code
6 +by mapping their custom model_type (colqwen3, ops_colqwen3, etc.) to a
7 +standard config class that vLLM understands.
8 +
9 +Supported model_types:
10 +- colqwen3 (TomoroAI/tomoro-colqwen3-embed-8b)
11 +- ops_colqwen3 (OpenSearch-AI/Ops-Colqwen3-4B)
12 +- qwen3_vl_nemotron_embed (nvidia/nemotron-colembed-vl-8b-v2)
13 +"""
14 +
15 +from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig
16 +
17 +
18 +class ColQwen3Config(Qwen3VLConfig):
19 + """Configuration class for ColQwen3 models.
20 +
21 + Extends Qwen3VLConfig with additional fields used by ColQwen3 variants
22 + for the embedding projection layer.
23 + """
24 +
25 + # Accept any ColQwen3 variant model_type
26 + model_type = "colqwen3"
27 +
28 + def __init__(
29 + self,
30 + embed_dim: int | None = None,
31 + dims: int | None = None,
32 + dim: int | None = None,
33 + projection_dim: int | None = None,
34 + colbert_dim: int | None = None,
35 + pooling: str | None = None,
36 + **kwargs,
37 + ):
38 + # Store embedding projection config fields
39 + self.embed_dim = embed_dim
40 + self.dims = dims
41 + self.dim = dim
42 + self.projection_dim = projection_dim
43 + self.colbert_dim = colbert_dim
44 + self.pooling = pooling
45 +
46 + super().__init__(**kwargs)
47 +
48 +
49 +class OpsColQwen3Config(ColQwen3Config):
50 + """Configuration for OpenSearch-AI ColQwen3 variants."""
51 +
52 + model_type = "ops_colqwen3"
53 +
54 +
55 +class Qwen3VLNemotronEmbedConfig(ColQwen3Config):
56 + """Configuration for NVIDIA Nemotron ColEmbed variants."""
57 +
58 + model_type = "qwen3_vl_nemotron_embed"