3c8e49596c · tduncan/Nvidia-Optimization-Stack-VLLM

docs/models/pooling_models/token_embed.md

+1 -1

@@ -61,7 +61,7 @@ Models of any architecture can be converted into embedding models using `--conve

66 | `Qwen3VLNemotronEmbedModel` | Qwen3-VL | T / I | `nvidia/nemotron-colembed-vl-4b-v2`, `nvidia/nemotron-colembed-vl-8b-v2` | ✅︎ | ✅︎ |

67 | `*ForConditionalGeneration`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | \* | N/A | \* | \* |

examples/pooling/score/colqwen3_5_rerank_online.py

+17 -1

@@ -7,11 +7,27 @@ ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5.

7 It produces per-token embeddings and uses MaxSim scoring for retrieval

8 and reranking. Supports both text and image inputs.

9

10 +Works for any ColQwen3.5 checkpoint, e.g. `athrael-soju/colqwen3.5-4.5B-v3`

11 +or `vultr/VultronRetrieverPrime-Qwen3.5-8B`.

12 +

13 Start the server with:

11 - vllm serve athrael-soju/colqwen3.5-4.5B --max-model-len 4096

14 + vllm serve athrael-soju/colqwen3.5-4.5B-v3 --max-model-len 4096 \

15 + --mm-processor-kwargs '{"min_pixels": 65536, "max_pixels": 1835008}'

16

17 Then run this script:

18 python colqwen3_5_rerank_online.py

19 +

20 +Parity note (matching the native colpali ColQwen3_5Processor pipeline):

21 + - Visual-token budget: ColQwen3_5Processor uses max_num_visual_tokens=1792,

22 + i.e. max_pixels = 1792 * (patch_size*merge_size)^2 = 1792 * 32^2 = 1835008

23 + (with min_pixels = shortest_edge = 65536). Pass these via --mm-processor-kwargs

24 + as above; the default budget gives fewer visual tokens and lower retrieval ndcg.

25 + - When you build prompts yourself (token_embed), reproduce the processor exactly:

26 + image (document): wrap in the instruction template

27 + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"

28 + "Describe the image.<|im_end|><|endoftext|>"

29 + query: append the augmentation suffix <text> + "<|endoftext|>" * 10

30 + Omitting these reproduces a silent ~2.5 ndcg@10 drop vs the native pipeline.

31 """

32

33 import requests

tests/models/multimodal/pooling/test_colqwen3_5.py

+18

@@ -152,3 +152,21 @@ def test_colqwen3_5_relevance_ordering(

152 dtype: str,

153 ) -> None:

154 _run_relevance_test(vllm_runner, model, dtype=dtype)

155 +

156 +

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 +

163 + from vllm.model_executor.models.config import (

164 + MODELS_CONFIG_MAP,

165 + ColQwen3_5Config,

166 + )

167 +

168 + assert MODELS_CONFIG_MAP["ColQwen3_5"] is ColQwen3_5Config

169 +

170 + model_config = SimpleNamespace(hf_config=SimpleNamespace())

171 + ColQwen3_5Config.verify_and_update_model_config(model_config)

172 + assert model_config.hf_config.is_causal is False

vllm/model_executor/layers/attention/attention.py

+8 -1

@@ -581,7 +581,14 @@ class Attention(nn.Module, AttentionLayerBase):

581 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:

582 # Block size may get updated after model loading, refresh it

583 block_size = vllm_config.cache_config.block_size

584 - # Should not be called for enc-dec or encoder-only attention.

584 + # Encoder-only attention is prefill-only and keeps no autoregressive KV

585 + # cache. In hybrid models (e.g. Qwen3.5 / ColQwen3.5: GatedDeltaNet

586 + # linear_attention interleaved with full_attention) the runner iterates

587 + # every attention module to build the KV-cache spec, so an ENCODER_ONLY

588 + # full_attention layer reaches here; it contributes no KV cache group.

589 + if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):

590 + return None

591 + # Should not be called for enc-dec attention.

592 assert self.attn_type == AttentionType.DECODER

593 quant_mode = get_kv_quant_mode(self.kv_cache_dtype)

594 if self.sliding_window is not None:

vllm/model_executor/models/colqwen3_5.py

+9 -1

@@ -15,6 +15,7 @@ Based on: Qwen3.5 backbone with custom text projection

15

16 Target models:

17 - athrael-soju/colqwen3.5-4.5B-v3

18 +- vultr/VultronRetrieverPrime-Qwen3.5-8B

19 """

20

21 from collections.abc import Iterable, Mapping

@@ -166,12 +167,19 @@ class ColQwen3_5Model(

167 or 128 # default from reference implementation

168 )

169

170 + # ColPali defines `custom_text_proj = nn.Linear(hidden, dim)`, i.e.

171 + # bias=True by default, and the trained ColQwen3.5 checkpoints ship a

172 + # `custom_text_proj.bias`. Construct with a bias and zero-initialize it:

173 + # a (legacy) bias-less checkpoint then behaves identically to bias=False,

174 + # while load_weights() below picks up a trained bias instead of silently

175 + # dropping it (which shifts every per-token vector and the MaxSim ranking).

176 self.custom_text_proj = nn.Linear(

177 hidden_size,

178 self.embed_dim,

172 - bias=False,

179 + bias=True,

180 dtype=head_dtype,

181 )

182 + nn.init.zeros_(self.custom_text_proj.bias)

183

184 pooler_config = vllm_config.model_config.pooler_config

185 assert pooler_config is not None

vllm/model_executor/models/config.py

+15 -1

@@ -627,6 +627,20 @@ class Qwen3_5ForConditionalGenerationConfig(VerifyAndUpdateConfig):

627 )

628

629

630 +class ColQwen3_5Config(Qwen3_5ForConditionalGenerationConfig):

631 + """ColQwen3.5 (late-interaction retrieval) inherits Qwen3.5's mamba cache

632 + handling and additionally serves BIDIRECTIONAL attention: ColPali-style

633 + document/query encoding attends over the whole sequence, not causally. Set

634 + is_causal=False so Qwen3NextAttention builds its full_attention layers with

635 + AttentionType.ENCODER_ONLY (the linear_attention GatedDeltaNet layers are

636 + unaffected). Generation arches keep the parent (causal) and are untouched.

637 + """

638 +

639 + @staticmethod

640 + def verify_and_update_model_config(model_config: "ModelConfig") -> None:

641 + model_config.hf_config.is_causal = False

642 +

643 +

644 class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig):

645 @staticmethod

646 def verify_and_update_model_config(model_config: "ModelConfig") -> None:

@@ -656,7 +670,7 @@ class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig):

670

671 MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {

672 "ColBERTJinaRobertaModel": JinaRobertaModelConfig,

659 - "ColQwen3_5": Qwen3_5ForConditionalGenerationConfig,

673 + "ColQwen3_5": ColQwen3_5Config,

674 "DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig,

675 "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,

676 "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501

vllm/model_executor/models/qwen3_next.py

+11

@@ -62,6 +62,7 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk

62 from vllm.platforms import current_platform

63 from vllm.sequence import IntermediateTensors

64 from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig

65 +from vllm.v1.attention.backend import AttentionType

66

67 from .interfaces import (

68 EagleModelMixin,

@@ -267,6 +268,15 @@ class Qwen3NextAttention(nn.Module):

268 dual_chunk_attention_config=self.dual_chunk_attention_config,

269 )

270

271 + # Late-interaction retrieval models (e.g. ColQwen3.5) run BIDIRECTIONAL

272 + # attention on the full_attention layers; they set config.is_causal=False

273 + # via a VerifyAndUpdateConfig handler. Generation models leave is_causal

274 + # unset (-> causal/DECODER), so this is a no-op for them. Mirrors qwen3.py.

275 + attn_type = (

276 + AttentionType.DECODER

277 + if getattr(config, "is_causal", True)

278 + else AttentionType.ENCODER_ONLY

279 + )

280 self.attn = Attention(

281 self.num_heads,

282 self.head_dim,

@@ -275,6 +285,7 @@ class Qwen3NextAttention(nn.Module):

285 cache_config=cache_config,

286 quant_config=quant_config,

287 prefix=f"{prefix}.attn",

288 + attn_type=attn_type,

289 **{

290 "layer_idx": extract_layer_index(prefix),

291 "dual_chunk_attention_config": self.dual_chunk_attention_config,

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