[Bugfix] Respect declared attention contract for ColQwen3.5 retrievers (#49372)
Signed-off-by: Athrael Soju <athrael.soju@gmail.com>
Athrael Soju committed
Jul 26, 2026 at 06:08 UTC
1240c74c0a47473449cf0c3a9c2d87a1e159f73b
2 files changed
+135
-16
tests/models/multimodal/pooling/test_colqwen3_5.py
+103
-8
@@ -7,6 +7,8 @@ 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
@@ -154,12 +156,14 @@ def test_colqwen3_5_relevance_ordering(
156
_run_relevance_test(vllm_runner, model, dtype=dtype)
157
158
157
-def test_colqwen3_5_config_enables_bidirectional_attention() -> None:
158
- """ColQwen3.5 retrieval must be served BIDIRECTIONAL (is_causal=False) so the
159
- full_attention layers build with AttentionType.ENCODER_ONLY. This guards the
160
- silent-causal regression (no GPU / model load needed)."""
161
- from types import SimpleNamespace
162
-
159
+@pytest.mark.parametrize(
160
+ ("contract", "expected_is_causal"),
161
+ [("causal", True), ("bidirectional", False)],
162
+)
163
+def test_colqwen3_5_config_applies_declared_attention_contract(
164
+ contract: str,
165
+ expected_is_causal: bool,
166
+) -> None:
167
from vllm.model_executor.models.config import (
168
MODELS_CONFIG_MAP,
169
ColQwen3_5Config,
@@ -167,6 +171,97 @@ def test_colqwen3_5_config_enables_bidirectional_attention() -> None:
171
172
assert MODELS_CONFIG_MAP["ColQwen3_5"] is ColQwen3_5Config
173
170
- model_config = SimpleNamespace(hf_config=SimpleNamespace())
174
+ hf_config = SimpleNamespace(retrieval_attention_contract=contract)
175
+ text_config = SimpleNamespace()
176
+ model_config = SimpleNamespace(
177
+ hf_config=hf_config,
178
+ hf_text_config=text_config,
179
+ )
180
+ ColQwen3_5Config.verify_and_update_model_config(model_config)
181
+ assert hf_config.is_causal is expected_is_causal
182
+ assert text_config.is_causal is expected_is_causal
183
+
184
+
185
+@pytest.mark.parametrize(
186
+ "hf_config",
187
+ [
188
+ SimpleNamespace(),
189
+ SimpleNamespace(retrieval_attention_contract="unsupported"),
190
+ SimpleNamespace(
191
+ retrieval_attention_contract="causal",
192
+ text_config=SimpleNamespace(retrieval_attention_contract="bidirectional"),
193
+ ),
194
+ ],
195
+)
196
+def test_colqwen3_5_config_rejects_invalid_attention_contract(hf_config) -> None:
197
+ from vllm.model_executor.models.config import ColQwen3_5Config
198
+
199
+ text_config = getattr(hf_config, "text_config", SimpleNamespace())
200
+ model_config = SimpleNamespace(
201
+ hf_config=hf_config,
202
+ hf_text_config=text_config,
203
+ )
204
+ with pytest.raises(ValueError, match="retrieval_attention_contract"):
205
+ ColQwen3_5Config.verify_and_update_model_config(model_config)
206
+
207
+
208
+def test_colqwen3_5_bidirectional_contract_builds_encoder_only_attention(
209
+ monkeypatch,
210
+) -> None:
211
+ from vllm.model_executor.models import qwen3_next
212
+ from vllm.model_executor.models.config import ColQwen3_5Config
213
+ from vllm.v1.attention.backend import AttentionType
214
+
215
+ hf_config = SimpleNamespace(retrieval_attention_contract="bidirectional")
216
+ text_config = SimpleNamespace(
217
+ hidden_size=256,
218
+ num_attention_heads=2,
219
+ num_key_value_heads=1,
220
+ head_dim=128,
221
+ max_position_embeddings=4096,
222
+ rope_parameters={},
223
+ rms_norm_eps=1e-6,
224
+ )
225
+ model_config = SimpleNamespace(
226
+ hf_config=hf_config,
227
+ hf_text_config=text_config,
228
+ )
229
ColQwen3_5Config.verify_and_update_model_config(model_config)
172
- assert model_config.hf_config.is_causal is False
230
+
231
+ captured = {}
232
+
233
+ class FakeAttention(torch.nn.Module):
234
+ def __init__(self, *args, **kwargs) -> None:
235
+ super().__init__()
236
+ captured["attn_type"] = kwargs["attn_type"]
237
+
238
+ monkeypatch.setattr(qwen3_next, "get_tensor_model_parallel_world_size", lambda: 1)
239
+ monkeypatch.setattr(
240
+ qwen3_next, "QKVParallelLinear", lambda *args, **kwargs: torch.nn.Identity()
241
+ )
242
+ monkeypatch.setattr(
243
+ qwen3_next, "RowParallelLinear", lambda *args, **kwargs: torch.nn.Identity()
244
+ )
245
+ monkeypatch.setattr(
246
+ qwen3_next,
247
+ "get_rope",
248
+ lambda *args, **kwargs: SimpleNamespace(is_neox_style=False),
249
+ )
250
+ monkeypatch.setattr(
251
+ qwen3_next, "Qwen3NextRMSNorm", lambda *args, **kwargs: torch.nn.Identity()
252
+ )
253
+ monkeypatch.setattr(qwen3_next, "Attention", FakeAttention)
254
+
255
+ qwen3_next.Qwen3NextAttention(text_config)
256
+
257
+ assert captured["attn_type"] is AttentionType.ENCODER_ONLY
258
+
259
+
260
+def test_colqwen3_5_encoder_only_attention_has_no_kv_cache_spec() -> None:
261
+ from vllm.model_executor.layers.attention import Attention
262
+ from vllm.v1.attention.backend import AttentionType
263
+
264
+ attention = SimpleNamespace(attn_type=AttentionType.ENCODER_ONLY)
265
+ vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16))
266
+
267
+ assert Attention.get_kv_cache_spec(attention, vllm_config) is None
vllm/model_executor/models/config.py
+32
-8
@@ -769,17 +769,41 @@ class Qwen3_5ForConditionalGenerationConfig(VerifyAndUpdateConfig):
769
770
771
class ColQwen3_5Config(Qwen3_5ForConditionalGenerationConfig):
772
- """ColQwen3.5 (late-interaction retrieval) inherits Qwen3.5's mamba cache
773
- handling and additionally serves BIDIRECTIONAL attention: ColPali-style
774
- document/query encoding attends over the whole sequence, not causally. Set
775
- is_causal=False so Qwen3NextAttention builds its full_attention layers with
776
- AttentionType.ENCODER_ONLY (the linear_attention GatedDeltaNet layers are
777
- unaffected). Generation arches keep the parent (causal) and are untouched.
778
- """
772
+ """Apply the attention contract declared by a ColQwen3.5 checkpoint."""
773
774
@staticmethod
775
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
782
- model_config.hf_config.is_causal = False
776
+ configs = {
777
+ id(config): config
778
+ for config in (
779
+ model_config.hf_config,
780
+ model_config.hf_text_config,
781
+ )
782
+ }
783
+ declarations = [
784
+ contract
785
+ for config in configs.values()
786
+ if (contract := getattr(config, "retrieval_attention_contract", None))
787
+ is not None
788
+ ]
789
+ supported = {"causal", "bidirectional"}
790
+ if not declarations:
791
+ raise ValueError(
792
+ "ColQwen3.5 checkpoints must declare "
793
+ "retrieval_attention_contract as 'causal' or 'bidirectional'"
794
+ )
795
+ if (
796
+ any(not isinstance(contract, str) for contract in declarations)
797
+ or any(contract not in supported for contract in declarations)
798
+ or len(set(declarations)) != 1
799
+ ):
800
+ raise ValueError(
801
+ "unsupported or conflicting ColQwen3.5 "
802
+ f"retrieval_attention_contract declarations: {declarations!r}"
803
+ )
804
+ is_causal = declarations[0] == "causal"
805
+ for config in configs.values():
806
+ config.is_causal = is_causal
807
808
809
class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):