[Model] ColQwen3.5: fix retrieval correctness (bias + bidirectional) (#46108)

Signed-off-by: Athrael Soju <[email protected]>
Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>

docs/models/pooling_models/token_embed.md

@@ -61,7 +61,7 @@ Models of any architecture can be converted into embedding models using `--conve
 | `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | |
 | `ColPaliForRetrieval` | ColPali | T / I | `vidore/colpali-v1.3-hf` | | |
 | `ColQwen3` | Qwen3-VL | T / I | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` | | |
-| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3` | | |
+| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3`, `vultr/VultronRetrieverPrime-Qwen3.5-8B` | | |
 | `OpsColQwen3Model` | Qwen3-VL | T / I | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` | | |
 | `Qwen3VLNemotronEmbedModel` | Qwen3-VL | T / I | `nvidia/nemotron-colembed-vl-4b-v2`, `nvidia/nemotron-colembed-vl-8b-v2` | ✅︎ | ✅︎ |
 | `*ForConditionalGeneration`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | \* | N/A | \* | \* |

examples/pooling/score/colqwen3_5_rerank_online.py

@@ -7,11 +7,27 @@ ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5.
 It produces per-token embeddings and uses MaxSim scoring for retrieval
 and reranking. Supports both text and image inputs.
 
+Works for any ColQwen3.5 checkpoint, e.g. `athrael-soju/colqwen3.5-4.5B-v3`
+or `vultr/VultronRetrieverPrime-Qwen3.5-8B`.
+
 Start the server with:
-    vllm serve athrael-soju/colqwen3.5-4.5B --max-model-len 4096
+    vllm serve athrael-soju/colqwen3.5-4.5B-v3 --max-model-len 4096 \
+        --mm-processor-kwargs '{"min_pixels": 65536, "max_pixels": 1835008}'
 
 Then run this script:
     python colqwen3_5_rerank_online.py
+
+Parity note (matching the native colpali ColQwen3_5Processor pipeline):
+  - Visual-token budget: ColQwen3_5Processor uses max_num_visual_tokens=1792,
+    i.e. max_pixels = 1792 * (patch_size*merge_size)^2 = 1792 * 32^2 = 1835008
+    (with min_pixels = shortest_edge = 65536). Pass these via --mm-processor-kwargs
+    as above; the default budget gives fewer visual tokens and lower retrieval ndcg.
+  - When you build prompts yourself (token_embed), reproduce the processor exactly:
+      image (document): wrap in the instruction template
+        "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+        "Describe the image.<|im_end|><|endoftext|>"
+      query: append the augmentation suffix  <text> + "<|endoftext|>" * 10
+    Omitting these reproduces a silent ~2.5 ndcg@10 drop vs the native pipeline.
 """
 
 import requests

tests/models/multimodal/pooling/test_colqwen3_5.py

@@ -152,3 +152,21 @@ def test_colqwen3_5_relevance_ordering(
     dtype: str,
 ) -> None:
     _run_relevance_test(vllm_runner, model, dtype=dtype)
+
+
+def test_colqwen3_5_config_enables_bidirectional_attention() -> None:
+    """ColQwen3.5 retrieval must be served BIDIRECTIONAL (is_causal=False) so the
+    full_attention layers build with AttentionType.ENCODER_ONLY. This guards the
+    silent-causal regression (no GPU / model load needed)."""
+    from types import SimpleNamespace
+
+    from vllm.model_executor.models.config import (
+        MODELS_CONFIG_MAP,
+        ColQwen3_5Config,
+    )
+
+    assert MODELS_CONFIG_MAP["ColQwen3_5"] is ColQwen3_5Config
+
+    model_config = SimpleNamespace(hf_config=SimpleNamespace())
+    ColQwen3_5Config.verify_and_update_model_config(model_config)
+    assert model_config.hf_config.is_causal is False

vllm/model_executor/layers/attention/attention.py

@@ -581,7 +581,14 @@ class Attention(nn.Module, AttentionLayerBase):
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
         # Block size may get updated after model loading, refresh it
         block_size = vllm_config.cache_config.block_size
-        # Should not be called for enc-dec or encoder-only attention.
+        # Encoder-only attention is prefill-only and keeps no autoregressive KV
+        # cache. In hybrid models (e.g. Qwen3.5 / ColQwen3.5: GatedDeltaNet
+        # linear_attention interleaved with full_attention) the runner iterates
+        # every attention module to build the KV-cache spec, so an ENCODER_ONLY
+        # full_attention layer reaches here; it contributes no KV cache group.
+        if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
+            return None
+        # Should not be called for enc-dec attention.
         assert self.attn_type == AttentionType.DECODER
         quant_mode = get_kv_quant_mode(self.kv_cache_dtype)
         if self.sliding_window is not None:

vllm/model_executor/models/colqwen3_5.py

@@ -15,6 +15,7 @@ Based on: Qwen3.5 backbone with custom text projection
 
 Target models:
 - athrael-soju/colqwen3.5-4.5B-v3
+- vultr/VultronRetrieverPrime-Qwen3.5-8B
 """
 
 from collections.abc import Iterable, Mapping
@@ -166,12 +167,19 @@ class ColQwen3_5Model(
             or 128  # default from reference implementation
         )
 
+        # ColPali defines `custom_text_proj = nn.Linear(hidden, dim)`, i.e.
+        # bias=True by default, and the trained ColQwen3.5 checkpoints ship a
+        # `custom_text_proj.bias`. Construct with a bias and zero-initialize it:
+        # a (legacy) bias-less checkpoint then behaves identically to bias=False,
+        # while load_weights() below picks up a trained bias instead of silently
+        # dropping it (which shifts every per-token vector and the MaxSim ranking).
         self.custom_text_proj = nn.Linear(
             hidden_size,
             self.embed_dim,
-            bias=False,
+            bias=True,
             dtype=head_dtype,
         )
+        nn.init.zeros_(self.custom_text_proj.bias)
 
         pooler_config = vllm_config.model_config.pooler_config
         assert pooler_config is not None

vllm/model_executor/models/config.py

@@ -627,6 +627,20 @@ class Qwen3_5ForConditionalGenerationConfig(VerifyAndUpdateConfig):
             )
 
 
+class ColQwen3_5Config(Qwen3_5ForConditionalGenerationConfig):
+    """ColQwen3.5 (late-interaction retrieval) inherits Qwen3.5's mamba cache
+    handling and additionally serves BIDIRECTIONAL attention: ColPali-style
+    document/query encoding attends over the whole sequence, not causally. Set
+    is_causal=False so Qwen3NextAttention builds its full_attention layers with
+    AttentionType.ENCODER_ONLY (the linear_attention GatedDeltaNet layers are
+    unaffected). Generation arches keep the parent (causal) and are untouched.
+    """
+
+    @staticmethod
+    def verify_and_update_model_config(model_config: "ModelConfig") -> None:
+        model_config.hf_config.is_causal = False
+
+
 class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):
     @staticmethod
     def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -656,7 +670,7 @@ class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig):
 
 MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
     "ColBERTJinaRobertaModel": JinaRobertaModelConfig,
-    "ColQwen3_5": Qwen3_5ForConditionalGenerationConfig,
+    "ColQwen3_5": ColQwen3_5Config,
     "DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig,
     "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
     "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig,  # noqa: E501

vllm/model_executor/models/qwen3_next.py

@@ -62,6 +62,7 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
+from vllm.v1.attention.backend import AttentionType
 
 from .interfaces import (
     EagleModelMixin,
@@ -267,6 +268,15 @@ class Qwen3NextAttention(nn.Module):
             dual_chunk_attention_config=self.dual_chunk_attention_config,
         )
 
+        # Late-interaction retrieval models (e.g. ColQwen3.5) run BIDIRECTIONAL
+        # attention on the full_attention layers; they set config.is_causal=False
+        # via a VerifyAndUpdateConfig handler. Generation models leave is_causal
+        # unset (-> causal/DECODER), so this is a no-op for them. Mirrors qwen3.py.
+        attn_type = (
+            AttentionType.DECODER
+            if getattr(config, "is_causal", True)
+            else AttentionType.ENCODER_ONLY
+        )
         self.attn = Attention(
             self.num_heads,
             self.head_dim,
@@ -275,6 +285,7 @@ class Qwen3NextAttention(nn.Module):
             cache_config=cache_config,
             quant_config=quant_config,
             prefix=f"{prefix}.attn",
+            attn_type=attn_type,
             **{
                 "layer_idx": extract_layer_index(prefix),
                 "dual_chunk_attention_config": self.dual_chunk_attention_config,