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, 19:00 UTC
Onboard voyage-4-nano (#33720)
Signed-off-by: Chengcheng Pei <chengchengpei@outlook.com> Signed-off-by: chengchengpei <5881383+chengchengpei@users.noreply.github.com> Co-authored-by: chengchengpei <5881383+chengchengpei@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
chengchengpei committed
Feb 5, 2026 at 22:23 UTC
965525667b70dc23463d57295dce792eba1ac452
8 files changed
+216
-2
docs/models/supported_models.md
+1
@@ -519,6 +519,7 @@ These models primarily support the [`LLM.embed`](./pooling_models.md#llmembed) A
519
| `LlamaModel`<sup>C</sup>, `LlamaForCausalLM`<sup>C</sup>, `MistralModel`<sup>C</sup>, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ |
520
| `Qwen2Model`<sup>C</sup>, `Qwen2ForCausalLM`<sup>C</sup> | Qwen2-based | `ssmits/Qwen2-7B-Instruct-embed-base` (see note), `Alibaba-NLP/gte-Qwen2-7B-instruct` (see note), etc. | ✅︎ | ✅︎ |
521
| `Qwen3Model`<sup>C</sup>, `Qwen3ForCausalLM`<sup>C</sup> | Qwen3-based | `Qwen/Qwen3-Embedding-0.6B`, etc. | ✅︎ | ✅︎ |
522
+| `VoyageQwen3BidirectionalEmbedModel`<sup>C</sup> | Voyage Qwen3-based with bidirectional attention | `voyageai/voyage-4-nano`, etc. | ✅︎ | ✅︎ |
523
| `RobertaModel`, `RobertaForMaskedLM` | RoBERTa-based | `sentence-transformers/all-roberta-large-v1`, etc. | | |
524
| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* |
525
tests/models/language/pooling_mteb_test/test_voyage.py
new
+56
@@ -0,0 +1,56 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from tests.models.language.pooling.embed_utils import correctness_test_embed_models
7
+from tests.models.utils import EmbedModelInfo
8
+
9
+from .mteb_embed_utils import mteb_test_embed_models
10
+
11
+MODELS = [
12
+ EmbedModelInfo(
13
+ "voyageai/voyage-4-nano",
14
+ architecture="VoyageQwen3BidirectionalEmbedModel",
15
+ enable_test=True,
16
+ seq_pooling_type="MEAN",
17
+ attn_type="encoder_only",
18
+ is_prefix_caching_supported=False,
19
+ is_chunked_prefill_supported=False,
20
+ hf_overrides={
21
+ "architectures": ["VoyageQwen3BidirectionalEmbedModel"],
22
+ "num_labels": 2048,
23
+ },
24
+ mteb_score=0.7054,
25
+ # === MTEB Results ===
26
+ # STS12: 0.6613
27
+ # STS13: 0.6906
28
+ # STS14: 0.6556
29
+ # STS15: 0.7843
30
+ # STS16: 0.7340
31
+ # STSBenchmark: 0.7063
32
+ # Average score: 0.7054
33
+ ),
34
+]
35
+
36
+
37
+@pytest.mark.parametrize("model_info", MODELS)
38
+def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
39
+ # Encoder-only attention models need enforce_eager=True to avoid
40
+ # CUDA graph capture issues with piecewise compilation
41
+ mteb_test_embed_models(
42
+ hf_runner, vllm_runner, model_info, vllm_extra_kwargs={"enforce_eager": True}
43
+ )
44
+
45
+
46
+@pytest.mark.parametrize("model_info", MODELS)
47
+def test_embed_models_correctness(
48
+ hf_runner, vllm_runner, model_info: EmbedModelInfo, example_prompts
49
+) -> None:
50
+ correctness_test_embed_models(
51
+ hf_runner,
52
+ vllm_runner,
53
+ model_info,
54
+ example_prompts,
55
+ vllm_extra_kwargs={"enforce_eager": True},
56
+ )
tests/models/registry.py
+3
@@ -565,6 +565,9 @@ _EMBEDDING_EXAMPLE_MODELS = {
565
),
566
"RobertaModel": _HfExamplesInfo("sentence-transformers/stsb-roberta-base-v2"),
567
"RobertaForMaskedLM": _HfExamplesInfo("sentence-transformers/all-roberta-large-v1"),
568
+ "VoyageQwen3BidirectionalEmbedModel": _HfExamplesInfo(
569
+ "voyageai/voyage-4-nano", trust_remote_code=True
570
+ ),
571
"XLMRobertaModel": _HfExamplesInfo("intfloat/multilingual-e5-small"),
572
"BertSpladeSparseEmbeddingModel": _HfExamplesInfo(
573
"naver/splade-v3",
vllm/config/model.py
+4
@@ -1513,6 +1513,10 @@ class ModelConfig:
1513
1514
@property
1515
def embedding_size(self):
1516
+ # Check for embedding_size set by model config (e.g., Voyage models)
1517
+ override = getattr(self.hf_config, "embedding_size", None)
1518
+ if override is not None:
1519
+ return override
1520
dense_modules = try_get_dense_modules(self.model, revision=self.revision)
1521
if dense_modules is not None:
1522
return dense_modules[-1]["out_features"]
vllm/model_executor/models/config.py
+8
@@ -582,6 +582,13 @@ class NemotronHForCausalLMConfig(VerifyAndUpdateConfig):
582
cache_config.mamba_ssm_cache_dtype = mamba_ssm_cache_dtype
583
584
585
+class VoyageQwen3BidirectionalEmbedModelConfig(VerifyAndUpdateConfig):
586
+ @staticmethod
587
+ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
588
+ model_config.hf_config.is_causal = False
589
+ model_config.hf_config.embedding_size = model_config.hf_config.num_labels
590
+
591
+
592
MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
593
"GteModel": SnowflakeGteNewModelConfig,
594
"GteNewModel": GteNewModelConfig,
@@ -604,4 +611,5 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
611
"DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
612
"NemotronHForCausalLM": NemotronHForCausalLMConfig,
613
"NemotronHPuzzleForCausalLM": NemotronHForCausalLMConfig,
614
+ "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig,
615
}
vllm/model_executor/models/qwen3.py
+10
-2
@@ -34,7 +34,10 @@ from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36
from vllm.logger import init_logger
37
-from vllm.model_executor.layers.attention import Attention
37
+from vllm.model_executor.layers.attention.encoder_only_attention import (
38
+ Attention,
39
+ EncoderOnlyAttention,
40
+)
41
from vllm.model_executor.layers.layernorm import RMSNorm
42
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
43
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -115,7 +118,12 @@ class Qwen3Attention(nn.Module):
118
rope_parameters=rope_parameters,
119
dual_chunk_attention_config=dual_chunk_attention_config,
120
)
118
- self.attn = Attention(
121
+ attn_cls = (
122
+ EncoderOnlyAttention
123
+ if attn_type == AttentionType.ENCODER_ONLY
124
+ else Attention
125
+ )
126
+ self.attn = attn_cls(
127
self.num_heads,
128
self.head_dim,
129
self.scaling,
vllm/model_executor/models/registry.py
+4
@@ -237,6 +237,10 @@ _EMBEDDING_MODELS = {
237
"RobertaModel": ("roberta", "RobertaEmbeddingModel"),
238
"TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
239
"TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
240
+ "VoyageQwen3BidirectionalEmbedModel": (
241
+ "voyage",
242
+ "VoyageQwen3BidirectionalEmbedModel",
243
+ ),
244
"XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
245
"BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"),
246
# [Multimodal]
vllm/model_executor/models/voyage.py
new
+130
@@ -0,0 +1,130 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from __future__ import annotations
4
+
5
+import re
6
+from collections import defaultdict
7
+from collections.abc import Iterable
8
+
9
+import torch
10
+import torch.nn as nn
11
+
12
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
13
+from vllm.model_executor.models.qwen3 import Qwen3Model
14
+from vllm.model_executor.models.utils import WeightsMapper
15
+
16
+WeightItem = tuple[str, torch.Tensor]
17
+
18
+_LAYER_RE = re.compile(r"^layers\.(\d+)\.(.+)$")
19
+
20
+
21
+class VoyageQwen3BidirectionalEmbedModel(Qwen3Model):
22
+ """
23
+ Qwen3Model + Voyage embedding head + bidirectional attention.
24
+
25
+ Checkpoint conventions (HF):
26
+ - MLP: gate_proj + up_proj (unfused)
27
+ - Attn: q_proj + k_proj + v_proj (unfused)
28
+ - Linear head: linear.weight
29
+ - Weights prefixed with "model." (e.g., model.layers.0...)
30
+
31
+ vLLM Qwen3Model expects:
32
+ - mlp.gate_up_proj (fused)
33
+ - self_attn.qkv_proj (fused)
34
+ - No "model." prefix
35
+
36
+ We remap/fuse weights using generator pipeline and load directly
37
+ (bypassing parent's stacked_params_mapping which would cause
38
+ double-transformation like qkv_proj -> qkqkv_proj).
39
+ """
40
+
41
+ hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
42
+
43
+ def __init__(self, *args, **kwargs):
44
+ super().__init__(*args, **kwargs)
45
+
46
+ # Embedding head (hidden_size -> num_labels, bias=False)
47
+ self.linear = nn.Linear(
48
+ self.config.hidden_size,
49
+ self.config.num_labels,
50
+ bias=False,
51
+ )
52
+
53
+ def forward(self, *args, **kwargs):
54
+ out = super().forward(*args, **kwargs)
55
+ return self.linear(out)
56
+
57
+ def _fuse_qkv_proj(self, weights: Iterable[WeightItem]) -> Iterable[WeightItem]:
58
+ """Fuse q_proj, k_proj, v_proj into qkv_proj."""
59
+ qkv_buf: dict[int, dict[str, torch.Tensor]] = defaultdict(dict)
60
+ qkv_suffixes = {
61
+ "self_attn.q_proj.weight": "q",
62
+ "self_attn.k_proj.weight": "k",
63
+ "self_attn.v_proj.weight": "v",
64
+ }
65
+
66
+ for name, tensor in weights:
67
+ m = _LAYER_RE.match(name)
68
+ if m and m.group(2) in qkv_suffixes:
69
+ layer_idx = int(m.group(1))
70
+ qkv_buf[layer_idx][qkv_suffixes[m.group(2)]] = tensor
71
+ else:
72
+ yield name, tensor
73
+
74
+ # Yield fused QKV weights
75
+ for layer_idx in sorted(qkv_buf.keys()):
76
+ parts = qkv_buf[layer_idx]
77
+ if all(p in parts for p in ("q", "k", "v")):
78
+ fused = torch.cat([parts["q"], parts["k"], parts["v"]], dim=0)
79
+ yield f"layers.{layer_idx}.self_attn.qkv_proj.weight", fused
80
+ elif parts:
81
+ missing = [p for p in ("q", "k", "v") if p not in parts]
82
+ raise ValueError(f"Layer {layer_idx} missing QKV parts: {missing}")
83
+
84
+ def _fuse_gate_up_proj(self, weights: Iterable[WeightItem]) -> Iterable[WeightItem]:
85
+ """Fuse gate_proj and up_proj into gate_up_proj."""
86
+ mlp_buf: dict[int, dict[str, torch.Tensor]] = defaultdict(dict)
87
+ mlp_suffixes = {
88
+ "mlp.gate_proj.weight": "gate",
89
+ "mlp.up_proj.weight": "up",
90
+ }
91
+
92
+ for name, tensor in weights:
93
+ m = _LAYER_RE.match(name)
94
+ if m and m.group(2) in mlp_suffixes:
95
+ layer_idx = int(m.group(1))
96
+ mlp_buf[layer_idx][mlp_suffixes[m.group(2)]] = tensor
97
+ else:
98
+ yield name, tensor
99
+
100
+ # Yield fused gate_up weights
101
+ for layer_idx in sorted(mlp_buf.keys()):
102
+ parts = mlp_buf[layer_idx]
103
+ if all(p in parts for p in ("gate", "up")):
104
+ fused = torch.cat([parts["gate"], parts["up"]], dim=0)
105
+ yield f"layers.{layer_idx}.mlp.gate_up_proj.weight", fused
106
+ elif parts:
107
+ missing = [p for p in ("gate", "up") if p not in parts]
108
+ raise ValueError(f"Layer {layer_idx} missing MLP parts: {missing}")
109
+
110
+ def load_weights(self, weights: Iterable[WeightItem]) -> set[str]:
111
+ """Remap, fuse, and load weights using generator pipeline."""
112
+ # Chain weight transformations
113
+ weights = self.hf_to_vllm_mapper.apply(weights)
114
+ weights = self._fuse_qkv_proj(weights)
115
+ weights = self._fuse_gate_up_proj(weights)
116
+
117
+ # Load weights directly into model parameters
118
+ # (bypass parent's stacked_params_mapping)
119
+ params_dict = dict(self.named_parameters())
120
+ loaded_params: set[str] = set()
121
+
122
+ for name, loaded_weight in weights:
123
+ if name not in params_dict:
124
+ continue
125
+ param = params_dict[name]
126
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
127
+ weight_loader(param, loaded_weight)
128
+ loaded_params.add(name)
129
+
130
+ return loaded_params