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 3, 01:45 UTC
[Model] Use `AutoWeightsLoader` for Voyage (#42972)
Signed-off-by: Furkan Fidan <dev@yufufi.com>
Furkan F committed
May 22, 2026 at 07:28 UTC
e746a2eebf09b1f99beb6b3c60a5ba9d2f8c4875
1 file changed
+27
-90
vllm/model_executor/models/voyage.py
+27
-90
@@ -2,23 +2,20 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from __future__ import annotations
4
5
-from collections import defaultdict
5
from collections.abc import Iterable
6
8
-import regex as re
7
import torch
8
import torch.nn as nn
9
12
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
10
+from vllm.config import VllmConfig
11
from vllm.model_executor.models.qwen3 import Qwen3Model
14
-from vllm.model_executor.models.utils import WeightsMapper
12
+from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix
13
+from vllm.sequence import IntermediateTensors
14
15
WeightItem = tuple[str, torch.Tensor]
16
18
-_LAYER_RE = re.compile(r"^layers\.(\d+)\.(.+)$")
17
20
-
21
-class VoyageQwen3BidirectionalEmbedModel(Qwen3Model):
18
+class VoyageQwen3BidirectionalEmbedModel(nn.Module):
19
"""
20
Qwen3Model + Voyage embedding head + bidirectional attention.
21
@@ -32,16 +29,14 @@ class VoyageQwen3BidirectionalEmbedModel(Qwen3Model):
29
- mlp.gate_up_proj (fused)
30
- self_attn.qkv_proj (fused)
31
- 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).
32
"""
33
41
- hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
42
-
43
- def __init__(self, *args, **kwargs):
44
- super().__init__(*args, **kwargs)
34
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
35
+ super().__init__()
36
+ self.config = vllm_config.model_config.hf_config
37
+ self.model = Qwen3Model(
38
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
39
+ )
40
41
# Embedding head (hidden_size -> num_labels, bias=False)
42
self.linear = nn.Linear(
@@ -50,81 +45,23 @@ class VoyageQwen3BidirectionalEmbedModel(Qwen3Model):
45
bias=False,
46
)
47
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
48
+ self.make_empty_intermediate_tensors = (
49
+ self.model.make_empty_intermediate_tensors
50
+ )
51
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}")
52
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
53
+ return self.model.embed_input_ids(input_ids)
54
+
55
+ def forward(
56
+ self,
57
+ input_ids: torch.Tensor | None,
58
+ positions: torch.Tensor,
59
+ intermediate_tensors: IntermediateTensors | None = None,
60
+ inputs_embeds: torch.Tensor | None = None,
61
+ ) -> torch.Tensor:
62
+ out = self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
63
+ return self.linear(out)
64
65
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
66
+ loader = AutoWeightsLoader(self)
67
+ return loader.load_weights(weights)