@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 236 lines 7.07 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for ColQwen3.5 late interaction model for multi-modal retrieval.
4
5 ColQwen3.5 is a multi-vector retrieval model based on Qwen3.5 backbone with
6 ColBERT-style late interaction scoring (MaxSim). It produces per-token
7 embeddings for both text and image inputs.
8 """
9
10 from types import SimpleNamespace
11
12 import pytest
13 import torch
14
15 from ....conftest import VllmRunner
16
17 MODELS = [
18 "athrael-soju/colqwen3.5-4.5B-v3",
19 ]
20
21 EMBED_DIMS = {
22 "athrael-soju/colqwen3.5-4.5B-v3": 320,
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 @pytest.fixture(scope="module", params=MODELS)
39 def colqwen3_5_model(request, vllm_runner):
40 model = request.param
41 with vllm_runner(
42 model,
43 runner="pooling",
44 dtype=DTYPE,
45 max_model_len=4096,
46 enforce_eager=True,
47 ) as vllm_model:
48 yield model, vllm_model
49
50
51 def _run_token_embed_test(
52 vllm_model: VllmRunner,
53 model: str,
54 ) -> None:
55 """Verify per-token embedding shape and L2 normalization."""
56 outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
57
58 assert len(outputs) == 1
59 emb = torch.tensor(outputs[0])
60 # Token embeddings should be 2D: [num_tokens, embed_dim]
61 assert emb.dim() == 2
62 assert emb.shape[1] == EMBED_DIMS[model]
63 assert emb.shape[0] > 1
64
65 # Verify L2 normalization
66 norms = torch.norm(emb, p=2, dim=-1)
67 torch.testing.assert_close(
68 norms,
69 torch.ones_like(norms),
70 rtol=1e-2,
71 atol=1e-2,
72 )
73
74
75 def _run_late_interaction_test(
76 vllm_model: VllmRunner,
77 ) -> None:
78 """Verify MaxSim scoring matches manual computation."""
79 from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
80
81 q_outputs = vllm_model.token_embed([TEXT_QUERIES[0]])
82 d_outputs = vllm_model.token_embed([TEXT_DOCUMENTS[0]])
83
84 q_emb = torch.tensor(q_outputs[0])
85 d_emb = torch.tensor(d_outputs[0])
86
87 manual_score = compute_maxsim_score(q_emb, d_emb).item()
88
89 vllm_scores = vllm_model.score(TEXT_QUERIES[0], TEXT_DOCUMENTS[0])
90
91 assert len(vllm_scores) == 1
92 assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
93
94
95 def _run_relevance_test(
96 vllm_model: VllmRunner,
97 ) -> None:
98 """Verify that relevant documents score higher than irrelevant ones."""
99 query = "What is machine learning?"
100 documents = [
101 "Machine learning is a subset of artificial intelligence.",
102 "The weather forecast shows rain tomorrow.",
103 "Deep learning uses neural networks for complex tasks.",
104 ]
105
106 scores = vllm_model.score(query, documents)
107
108 assert len(scores) == 3
109 assert scores[0] > scores[1], "ML doc should score higher than weather doc"
110 assert scores[2] > scores[1], "DL doc should score higher than weather doc"
111
112
113 def test_colqwen3_5_token_embed(colqwen3_5_model) -> None:
114 model, vllm_model = colqwen3_5_model
115 _run_token_embed_test(vllm_model, model)
116
117
118 def test_colqwen3_5_late_interaction_scoring(colqwen3_5_model) -> None:
119 _, vllm_model = colqwen3_5_model
120 _run_late_interaction_test(vllm_model)
121
122
123 def test_colqwen3_5_relevance_ordering(colqwen3_5_model) -> None:
124 _, vllm_model = colqwen3_5_model
125 _run_relevance_test(vllm_model)
126
127
128 @pytest.mark.parametrize(
129 ("contract", "expected_is_causal"),
130 [("causal", True), ("bidirectional", False)],
131 )
132 def test_colqwen3_5_config_applies_declared_attention_contract(
133 contract: str,
134 expected_is_causal: bool,
135 ) -> None:
136 from vllm.model_executor.models.config import (
137 MODELS_CONFIG_MAP,
138 ColQwen3_5Config,
139 )
140
141 assert MODELS_CONFIG_MAP["ColQwen3_5"] is ColQwen3_5Config
142
143 hf_config = SimpleNamespace(retrieval_attention_contract=contract)
144 text_config = SimpleNamespace()
145 model_config = SimpleNamespace(
146 hf_config=hf_config,
147 hf_text_config=text_config,
148 )
149 ColQwen3_5Config.verify_and_update_model_config(model_config)
150 assert hf_config.is_causal is expected_is_causal
151 assert text_config.is_causal is expected_is_causal
152
153
154 @pytest.mark.parametrize(
155 "hf_config",
156 [
157 SimpleNamespace(),
158 SimpleNamespace(retrieval_attention_contract="unsupported"),
159 SimpleNamespace(
160 retrieval_attention_contract="causal",
161 text_config=SimpleNamespace(retrieval_attention_contract="bidirectional"),
162 ),
163 ],
164 )
165 def test_colqwen3_5_config_rejects_invalid_attention_contract(hf_config) -> None:
166 from vllm.model_executor.models.config import ColQwen3_5Config
167
168 text_config = getattr(hf_config, "text_config", SimpleNamespace())
169 model_config = SimpleNamespace(
170 hf_config=hf_config,
171 hf_text_config=text_config,
172 )
173 with pytest.raises(ValueError, match="retrieval_attention_contract"):
174 ColQwen3_5Config.verify_and_update_model_config(model_config)
175
176
177 def test_colqwen3_5_bidirectional_contract_builds_encoder_only_attention(
178 monkeypatch,
179 ) -> None:
180 from vllm.model_executor.models import qwen3_next
181 from vllm.model_executor.models.config import ColQwen3_5Config
182 from vllm.v1.attention.backend import AttentionType
183
184 hf_config = SimpleNamespace(retrieval_attention_contract="bidirectional")
185 text_config = SimpleNamespace(
186 hidden_size=256,
187 num_attention_heads=2,
188 num_key_value_heads=1,
189 head_dim=128,
190 max_position_embeddings=4096,
191 rope_parameters={},
192 rms_norm_eps=1e-6,
193 )
194 model_config = SimpleNamespace(
195 hf_config=hf_config,
196 hf_text_config=text_config,
197 )
198 ColQwen3_5Config.verify_and_update_model_config(model_config)
199
200 captured = {}
201
202 class FakeAttention(torch.nn.Module):
203 def __init__(self, *args, **kwargs) -> None:
204 super().__init__()
205 captured["attn_type"] = kwargs["attn_type"]
206
207 monkeypatch.setattr(qwen3_next, "get_tensor_model_parallel_world_size", lambda: 1)
208 monkeypatch.setattr(
209 qwen3_next, "QKVParallelLinear", lambda *args, **kwargs: torch.nn.Identity()
210 )
211 monkeypatch.setattr(
212 qwen3_next, "RowParallelLinear", lambda *args, **kwargs: torch.nn.Identity()
213 )
214 monkeypatch.setattr(
215 qwen3_next,
216 "get_rope",
217 lambda *args, **kwargs: SimpleNamespace(is_neox_style=False),
218 )
219 monkeypatch.setattr(
220 qwen3_next, "Qwen3NextRMSNorm", lambda *args, **kwargs: torch.nn.Identity()
221 )
222 monkeypatch.setattr(qwen3_next, "Attention", FakeAttention)
223
224 qwen3_next.Qwen3NextAttention(text_config)
225
226 assert captured["attn_type"] is AttentionType.ENCODER_ONLY
227
228
229 def test_colqwen3_5_encoder_only_attention_has_no_kv_cache_spec() -> None:
230 from vllm.model_executor.layers.attention import Attention
231 from vllm.v1.attention.backend import AttentionType
232
233 attention = SimpleNamespace(attn_type=AttentionType.ENCODER_ONLY)
234 vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16))
235
236 assert Attention.get_kv_cache_spec(attention, vllm_config) is None