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 15, 16:00 UTC
[Model] Add Qwen3 and Qwen3MoE (#15289)
Signed-off-by: YamPengLi <yampayne.lyp@alibaba-inc.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
YamPengLi committed
Apr 7, 2025 at 19:06 UTC
7699258ef01392926750825df6b254db45b0c4e8
6 files changed
+893
-5
docs/source/models/supported_models.md
+10
@@ -478,6 +478,16 @@ See [this page](#generative-models) for more information on how to use generativ
478
* `Qwen/Qwen1.5-MoE-A2.7B`, `Qwen/Qwen1.5-MoE-A2.7B-Chat`, etc.
479
*
480
* ✅︎
481
+- * `Qwen3ForCausalLM`
482
+ * Qwen3
483
+ * `Qwen/Qwen3-8B`, etc.
484
+ * ✅︎
485
+ * ✅︎
486
+- * `Qwen3MoeForCausalLM`
487
+ * Qwen3MoE
488
+ * `Qwen/Qwen3-MoE-15B-A2B`, etc.
489
+ * ✅︎
490
+ * ✅︎
491
- * `StableLmForCausalLM`
492
* StableLM
493
* `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc.
tests/models/registry.py
+10
@@ -202,6 +202,16 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
202
"Qwen2ForCausalLM": _HfExamplesInfo("Qwen/Qwen2-7B-Instruct",
203
extras={"2.5": "Qwen/Qwen2.5-7B-Instruct"}), # noqa: E501
204
"Qwen2MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen1.5-MoE-A2.7B-Chat"),
205
+ "Qwen3ForCausalLM": _HfExamplesInfo(
206
+ "Qwen/Qwen3-8B",
207
+ is_available_online=False,
208
+ min_transformers_version="4.51"
209
+ ),
210
+ "Qwen3MoeForCausalLM": _HfExamplesInfo(
211
+ "Qwen/Qwen3-MoE-15B-A2B",
212
+ is_available_online=False,
213
+ min_transformers_version="4.51"
214
+ ),
215
"RWForCausalLM": _HfExamplesInfo("tiiuae/falcon-40b",
216
is_available_online=False),
217
"StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b", # noqa: E501
vllm/model_executor/models/qwen2.py
+11
-5
@@ -263,7 +263,11 @@ class Qwen2DecoderLayer(nn.Module):
263
})
264
class Qwen2Model(nn.Module):
265
266
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
266
+ def __init__(self,
267
+ *,
268
+ vllm_config: VllmConfig,
269
+ prefix: str = "",
270
+ decoder_layer_type: type[nn.Module] = Qwen2DecoderLayer):
271
super().__init__()
272
273
config = vllm_config.model_config.hf_config
@@ -297,12 +301,14 @@ class Qwen2Model(nn.Module):
301
else:
302
self.embed_tokens = PPMissingLayer()
303
304
+ # Use the provided decoder layer type or default to Qwen2DecoderLayer
305
+ decoder_layer_type = decoder_layer_type or Qwen2DecoderLayer
306
self.start_layer, self.end_layer, self.layers = make_layers(
307
config.num_hidden_layers,
302
- lambda prefix: Qwen2DecoderLayer(config=config,
303
- cache_config=cache_config,
304
- quant_config=quant_config,
305
- prefix=prefix),
308
+ lambda prefix: decoder_layer_type(config=config,
309
+ cache_config=cache_config,
310
+ quant_config=quant_config,
311
+ prefix=prefix),
312
prefix=f"{prefix}.layers",
313
)
314
vllm/model_executor/models/qwen3.py
new
+329
@@ -0,0 +1,329 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+# Copyright 2024 The Qwen team.
4
+# Copyright 2023 The vLLM team.
5
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
6
+#
7
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
8
+# and OPT implementations in this library. It has been modified from its
9
+# original forms to accommodate minor architectural differences compared
10
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
11
+#
12
+# Licensed under the Apache License, Version 2.0 (the "License");
13
+# you may not use this file except in compliance with the License.
14
+# You may obtain a copy of the License at
15
+#
16
+# http://www.apache.org/licenses/LICENSE-2.0
17
+#
18
+# Unless required by applicable law or agreed to in writing, software
19
+# distributed under the License is distributed on an "AS IS" BASIS,
20
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
21
+# See the License for the specific language governing permissions and
22
+# limitations under the License.
23
+"""Inference-only Qwen3 model compatible with HuggingFace weights."""
24
+from typing import Iterable, Optional, Set, Tuple, Union
25
+
26
+import torch
27
+from torch import nn
28
+from transformers import Qwen3Config
29
+
30
+from vllm.attention import Attention, AttentionType
31
+from vllm.compilation.decorators import support_torch_compile
32
+from vllm.config import CacheConfig, VllmConfig
33
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
34
+from vllm.logger import init_logger
35
+from vllm.model_executor.layers.layernorm import RMSNorm
36
+from vllm.model_executor.layers.linear import (QKVParallelLinear,
37
+ RowParallelLinear)
38
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
39
+from vllm.model_executor.layers.quantization import QuantizationConfig
40
+from vllm.model_executor.layers.rotary_embedding import get_rope
41
+from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
42
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
43
+from vllm.model_executor.sampling_metadata import SamplingMetadata
44
+from vllm.sequence import IntermediateTensors
45
+
46
+from .interfaces import SupportsLoRA, SupportsPP
47
+from .qwen2 import Qwen2MLP as Qwen3MLP
48
+from .qwen2 import Qwen2Model
49
+from .utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix
50
+
51
+logger = init_logger(__name__)
52
+
53
+
54
+class Qwen3Attention(nn.Module):
55
+
56
+ def __init__(self,
57
+ hidden_size: int,
58
+ num_heads: int,
59
+ num_kv_heads: int,
60
+ max_position: int = 4096 * 32,
61
+ head_dim: Optional[int] = None,
62
+ rms_norm_eps: float = 1e-06,
63
+ qkv_bias: bool = False,
64
+ rope_theta: float = 10000,
65
+ cache_config: Optional[CacheConfig] = None,
66
+ quant_config: Optional[QuantizationConfig] = None,
67
+ rope_scaling: Optional[Tuple] = None,
68
+ prefix: str = "",
69
+ attn_type: str = AttentionType.DECODER) -> None:
70
+ super().__init__()
71
+ self.hidden_size = hidden_size
72
+ tp_size = get_tensor_model_parallel_world_size()
73
+ self.total_num_heads = num_heads
74
+ assert self.total_num_heads % tp_size == 0
75
+ self.num_heads = self.total_num_heads // tp_size
76
+ self.total_num_kv_heads = num_kv_heads
77
+ if self.total_num_kv_heads >= tp_size:
78
+ # Number of KV heads is greater than TP size, so we partition
79
+ # the KV heads across multiple tensor parallel GPUs.
80
+ assert self.total_num_kv_heads % tp_size == 0
81
+ else:
82
+ # Number of KV heads is less than TP size, so we replicate
83
+ # the KV heads across multiple tensor parallel GPUs.
84
+ assert tp_size % self.total_num_kv_heads == 0
85
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
86
+ self.head_dim = head_dim or hidden_size // self.total_num_heads
87
+ self.q_size = self.num_heads * self.head_dim
88
+ self.kv_size = self.num_kv_heads * self.head_dim
89
+ self.scaling = self.head_dim**-0.5
90
+ self.rope_theta = rope_theta
91
+
92
+ self.qkv_proj = QKVParallelLinear(
93
+ hidden_size,
94
+ self.head_dim,
95
+ self.total_num_heads,
96
+ self.total_num_kv_heads,
97
+ bias=qkv_bias,
98
+ quant_config=quant_config,
99
+ prefix=f"{prefix}.qkv_proj",
100
+ )
101
+ self.o_proj = RowParallelLinear(
102
+ self.total_num_heads * self.head_dim,
103
+ hidden_size,
104
+ bias=False,
105
+ quant_config=quant_config,
106
+ prefix=f"{prefix}.o_proj",
107
+ )
108
+
109
+ self.rotary_emb = get_rope(
110
+ self.head_dim,
111
+ rotary_dim=self.head_dim,
112
+ max_position=max_position,
113
+ base=self.rope_theta,
114
+ rope_scaling=rope_scaling,
115
+ )
116
+ self.attn = Attention(self.num_heads,
117
+ self.head_dim,
118
+ self.scaling,
119
+ num_kv_heads=self.num_kv_heads,
120
+ cache_config=cache_config,
121
+ quant_config=quant_config,
122
+ prefix=f"{prefix}.attn",
123
+ attn_type=attn_type)
124
+ self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
125
+ self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
126
+
127
+ def forward(
128
+ self,
129
+ positions: torch.Tensor,
130
+ hidden_states: torch.Tensor,
131
+ ) -> torch.Tensor:
132
+ qkv, _ = self.qkv_proj(hidden_states)
133
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
134
+ # Add qk-norm
135
+ q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim,
136
+ self.head_dim)
137
+ q_by_head = self.q_norm.forward_native(q_by_head)
138
+ q = q_by_head.view(q.shape)
139
+ k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim,
140
+ self.head_dim)
141
+ k_by_head = self.k_norm.forward_native(k_by_head)
142
+ k = k_by_head.view(k.shape)
143
+ q, k = self.rotary_emb(positions, q, k)
144
+ attn_output = self.attn(q, k, v)
145
+ output, _ = self.o_proj(attn_output)
146
+ return output
147
+
148
+
149
+class Qwen3DecoderLayer(nn.Module):
150
+
151
+ def __init__(
152
+ self,
153
+ config: Qwen3Config,
154
+ cache_config: Optional[CacheConfig] = None,
155
+ quant_config: Optional[QuantizationConfig] = None,
156
+ prefix: str = "",
157
+ ) -> None:
158
+ super().__init__()
159
+ self.hidden_size = config.hidden_size
160
+ # Requires transformers > 4.32.0
161
+ rope_theta = getattr(config, "rope_theta", 1000000)
162
+ rope_scaling = getattr(config, "rope_scaling", None)
163
+
164
+ # By default, Qwen3 uses causal attention as it is a decoder-only model.
165
+ # You can override the HF config with `is_causal=False` to enable
166
+ # bidirectional attention, which is used in some embedding models
167
+ # (e.g. Alibaba-NLP/gte-Qwen3-7B-instruct)
168
+ if getattr(config, "is_causal", True):
169
+ attn_type = AttentionType.DECODER
170
+ else:
171
+ attn_type = AttentionType.ENCODER_ONLY
172
+
173
+ self.self_attn = Qwen3Attention(
174
+ hidden_size=self.hidden_size,
175
+ num_heads=config.num_attention_heads,
176
+ max_position=config.max_position_embeddings,
177
+ num_kv_heads=config.num_key_value_heads,
178
+ rope_theta=rope_theta,
179
+ rms_norm_eps=config.rms_norm_eps,
180
+ qkv_bias=getattr(config, 'attention_bias', False),
181
+ head_dim=getattr(config, 'head_dim', None),
182
+ cache_config=cache_config,
183
+ quant_config=quant_config,
184
+ rope_scaling=rope_scaling,
185
+ prefix=f"{prefix}.self_attn",
186
+ attn_type=attn_type,
187
+ )
188
+ self.mlp = Qwen3MLP(
189
+ hidden_size=self.hidden_size,
190
+ intermediate_size=config.intermediate_size,
191
+ hidden_act=config.hidden_act,
192
+ quant_config=quant_config,
193
+ prefix=f"{prefix}.mlp",
194
+ )
195
+ self.input_layernorm = RMSNorm(config.hidden_size,
196
+ eps=config.rms_norm_eps)
197
+ self.post_attention_layernorm = RMSNorm(config.hidden_size,
198
+ eps=config.rms_norm_eps)
199
+
200
+ def forward(
201
+ self,
202
+ positions: torch.Tensor,
203
+ hidden_states: torch.Tensor,
204
+ residual: Optional[torch.Tensor],
205
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
206
+ # Self Attention
207
+ if residual is None:
208
+ residual = hidden_states
209
+ hidden_states = self.input_layernorm(hidden_states)
210
+ else:
211
+ hidden_states, residual = self.input_layernorm(
212
+ hidden_states, residual)
213
+ hidden_states = self.self_attn(
214
+ positions=positions,
215
+ hidden_states=hidden_states,
216
+ )
217
+
218
+ # Fully Connected
219
+ hidden_states, residual = self.post_attention_layernorm(
220
+ hidden_states, residual)
221
+ hidden_states = self.mlp(hidden_states)
222
+ return hidden_states, residual
223
+
224
+
225
+ALL_DECODER_LAYER_TYPES = {
226
+ "attention": Qwen3DecoderLayer,
227
+}
228
+
229
+
230
+@support_torch_compile(
231
+ dynamic_arg_dims={
232
+ "input_ids": 0,
233
+ # positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
234
+ # otherwise (seq_len, ).
235
+ "positions": -1,
236
+ "intermediate_tensors": 0,
237
+ "inputs_embeds": 0,
238
+ })
239
+class Qwen3Model(Qwen2Model):
240
+
241
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
242
+ super().__init__(vllm_config=vllm_config,
243
+ prefix=prefix,
244
+ decoder_layer_type=Qwen3DecoderLayer)
245
+
246
+
247
+class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
248
+ packed_modules_mapping = {
249
+ "qkv_proj": [
250
+ "q_proj",
251
+ "k_proj",
252
+ "v_proj",
253
+ ],
254
+ "gate_up_proj": [
255
+ "gate_proj",
256
+ "up_proj",
257
+ ],
258
+ }
259
+
260
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
261
+ super().__init__()
262
+ config = vllm_config.model_config.hf_config
263
+ quant_config = vllm_config.quant_config
264
+ lora_config = vllm_config.lora_config
265
+
266
+ self.config = config
267
+ self.lora_config = lora_config
268
+
269
+ self.quant_config = quant_config
270
+ self.model = Qwen3Model(vllm_config=vllm_config,
271
+ prefix=maybe_prefix(prefix, "model"))
272
+
273
+ if get_pp_group().is_last_rank:
274
+ if config.tie_word_embeddings:
275
+ self.lm_head = self.model.embed_tokens
276
+ else:
277
+ self.lm_head = ParallelLMHead(config.vocab_size,
278
+ config.hidden_size,
279
+ quant_config=quant_config,
280
+ prefix=maybe_prefix(
281
+ prefix, "lm_head"))
282
+ else:
283
+ self.lm_head = PPMissingLayer()
284
+
285
+ self.logits_processor = LogitsProcessor(config.vocab_size)
286
+ self.sampler = get_sampler()
287
+
288
+ self.make_empty_intermediate_tensors = (
289
+ self.model.make_empty_intermediate_tensors)
290
+
291
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
292
+ return self.model.get_input_embeddings(input_ids)
293
+
294
+ def forward(
295
+ self,
296
+ input_ids: torch.Tensor,
297
+ positions: torch.Tensor,
298
+ intermediate_tensors: Optional[IntermediateTensors] = None,
299
+ inputs_embeds: Optional[torch.Tensor] = None,
300
+ ) -> Union[torch.Tensor, IntermediateTensors]:
301
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
302
+ inputs_embeds)
303
+ return hidden_states
304
+
305
+ def compute_logits(
306
+ self,
307
+ hidden_states: torch.Tensor,
308
+ sampling_metadata: SamplingMetadata,
309
+ ) -> Optional[torch.Tensor]:
310
+ logits = self.logits_processor(self.lm_head, hidden_states,
311
+ sampling_metadata)
312
+ return logits
313
+
314
+ def sample(
315
+ self,
316
+ logits: torch.Tensor,
317
+ sampling_metadata: SamplingMetadata,
318
+ ) -> Optional[SamplerOutput]:
319
+ next_tokens = self.sampler(logits, sampling_metadata)
320
+ return next_tokens
321
+
322
+ def load_weights(self, weights: Iterable[Tuple[str,
323
+ torch.Tensor]]) -> Set[str]:
324
+ loader = AutoWeightsLoader(
325
+ self,
326
+ skip_prefixes=(["lm_head."]
327
+ if self.config.tie_word_embeddings else None),
328
+ )
329
+ return loader.load_weights(weights)
vllm/model_executor/models/qwen3_moe.py
new
+531
@@ -0,0 +1,531 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+# Copyright 2024 The Qwen team.
4
+# Copyright 2023 The vLLM team.
5
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
6
+#
7
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
8
+# and OPT implementations in this library. It has been modified from its
9
+# original forms to accommodate minor architectural differences compared
10
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
11
+#
12
+# Licensed under the Apache License, Version 2.0 (the "License");
13
+# you may not use this file except in compliance with the License.
14
+# You may obtain a copy of the License at
15
+#
16
+# http://www.apache.org/licenses/LICENSE-2.0
17
+#
18
+# Unless required by applicable law or agreed to in writing, software
19
+# distributed under the License is distributed on an "AS IS" BASIS,
20
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
21
+# See the License for the specific language governing permissions and
22
+# limitations under the License.
23
+"""Inference-only Qwen3MoE model compatible with HuggingFace weights."""
24
+from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union
25
+
26
+import torch
27
+from torch import nn
28
+from transformers import PretrainedConfig
29
+
30
+from vllm.attention import Attention
31
+from vllm.compilation.decorators import support_torch_compile
32
+from vllm.config import CacheConfig, VllmConfig
33
+from vllm.distributed import (get_pp_group,
34
+ get_tensor_model_parallel_world_size,
35
+ tensor_model_parallel_all_reduce)
36
+from vllm.logger import init_logger
37
+from vllm.model_executor.layers.activation import SiluAndMul
38
+from vllm.model_executor.layers.fused_moe import FusedMoE
39
+from vllm.model_executor.layers.layernorm import RMSNorm
40
+from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
41
+ QKVParallelLinear,
42
+ ReplicatedLinear,
43
+ RowParallelLinear)
44
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
45
+from vllm.model_executor.layers.quantization import QuantizationConfig
46
+from vllm.model_executor.layers.rotary_embedding import get_rope
47
+from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
48
+from vllm.model_executor.layers.vocab_parallel_embedding import (
49
+ ParallelLMHead, VocabParallelEmbedding)
50
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
51
+from vllm.model_executor.sampling_metadata import SamplingMetadata
52
+from vllm.sequence import IntermediateTensors
53
+
54
+from .interfaces import SupportsPP
55
+from .utils import (extract_layer_index, is_pp_missing_parameter,
56
+ make_empty_intermediate_tensors_factory, make_layers,
57
+ maybe_prefix)
58
+
59
+logger = init_logger(__name__)
60
+
61
+
62
+class Qwen3MoeMLP(nn.Module):
63
+
64
+ def __init__(
65
+ self,
66
+ hidden_size: int,
67
+ intermediate_size: int,
68
+ hidden_act: str,
69
+ quant_config: Optional[QuantizationConfig] = None,
70
+ reduce_results: bool = True,
71
+ prefix: str = "",
72
+ ) -> None:
73
+ super().__init__()
74
+ self.gate_up_proj = MergedColumnParallelLinear(
75
+ hidden_size, [intermediate_size] * 2,
76
+ bias=False,
77
+ quant_config=quant_config,
78
+ prefix=f"{prefix}.gate_up_proj")
79
+ self.down_proj = RowParallelLinear(intermediate_size,
80
+ hidden_size,
81
+ bias=False,
82
+ quant_config=quant_config,
83
+ reduce_results=reduce_results,
84
+ prefix=f"{prefix}.down_proj")
85
+ if hidden_act != "silu":
86
+ raise ValueError(f"Unsupported activation: {hidden_act}. "
87
+ "Only silu is supported for now.")
88
+ self.act_fn = SiluAndMul()
89
+
90
+ def forward(self, x):
91
+ gate_up, _ = self.gate_up_proj(x)
92
+ x = self.act_fn(gate_up)
93
+ x, _ = self.down_proj(x)
94
+ return x
95
+
96
+
97
+class Qwen3MoeSparseMoeBlock(nn.Module):
98
+
99
+ def __init__(
100
+ self,
101
+ config: PretrainedConfig,
102
+ quant_config: Optional[QuantizationConfig] = None,
103
+ prefix: str = "",
104
+ ):
105
+ super().__init__()
106
+ self.tp_size = get_tensor_model_parallel_world_size()
107
+
108
+ if self.tp_size > config.num_experts:
109
+ raise ValueError(
110
+ f"Tensor parallel size {self.tp_size} is greater than "
111
+ f"the number of experts {config.num_experts}.")
112
+
113
+ self.experts = FusedMoE(num_experts=config.num_experts,
114
+ top_k=config.num_experts_per_tok,
115
+ hidden_size=config.hidden_size,
116
+ intermediate_size=config.moe_intermediate_size,
117
+ reduce_results=False,
118
+ renormalize=config.norm_topk_prob,
119
+ quant_config=quant_config,
120
+ prefix=f"{prefix}.experts")
121
+
122
+ self.gate = ReplicatedLinear(config.hidden_size,
123
+ config.num_experts,
124
+ bias=False,
125
+ quant_config=None,
126
+ prefix=f"{prefix}.gate")
127
+
128
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
129
+ # NOTE: hidden_states can have either 1D or 2D shape.
130
+ orig_shape = hidden_states.shape
131
+ hidden_dim = hidden_states.shape[-1]
132
+ hidden_states = hidden_states.view(-1, hidden_dim)
133
+
134
+ # router_logits: (num_tokens, n_experts)
135
+ router_logits, _ = self.gate(hidden_states)
136
+ final_hidden_states = self.experts(hidden_states=hidden_states,
137
+ router_logits=router_logits)
138
+ final_hidden_states = final_hidden_states
139
+ if self.tp_size > 1:
140
+ final_hidden_states = tensor_model_parallel_all_reduce(
141
+ final_hidden_states)
142
+
143
+ return final_hidden_states.view(orig_shape)
144
+
145
+
146
+class Qwen3MoeAttention(nn.Module):
147
+
148
+ def __init__(
149
+ self,
150
+ hidden_size: int,
151
+ num_heads: int,
152
+ num_kv_heads: int,
153
+ rope_theta: float = 10000,
154
+ rope_scaling: Optional[Dict[str, Any]] = None,
155
+ max_position_embeddings: int = 8192,
156
+ head_dim: Optional[int] = None,
157
+ rms_norm_eps: float = 1e-06,
158
+ qkv_bias: bool = False,
159
+ cache_config: Optional[CacheConfig] = None,
160
+ quant_config: Optional[QuantizationConfig] = None,
161
+ prefix: str = "",
162
+ ) -> None:
163
+ super().__init__()
164
+ self.hidden_size = hidden_size
165
+ tp_size = get_tensor_model_parallel_world_size()
166
+ self.total_num_heads = num_heads
167
+ assert self.total_num_heads % tp_size == 0
168
+ self.num_heads = self.total_num_heads // tp_size
169
+ self.total_num_kv_heads = num_kv_heads
170
+ if self.total_num_kv_heads >= tp_size:
171
+ # Number of KV heads is greater than TP size, so we partition
172
+ # the KV heads across multiple tensor parallel GPUs.
173
+ assert self.total_num_kv_heads % tp_size == 0
174
+ else:
175
+ # Number of KV heads is less than TP size, so we replicate
176
+ # the KV heads across multiple tensor parallel GPUs.
177
+ assert tp_size % self.total_num_kv_heads == 0
178
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
179
+ self.head_dim = head_dim or (hidden_size // self.total_num_heads)
180
+ self.q_size = self.num_heads * self.head_dim
181
+ self.kv_size = self.num_kv_heads * self.head_dim
182
+ self.scaling = self.head_dim**-0.5
183
+ self.rope_theta = rope_theta
184
+ self.max_position_embeddings = max_position_embeddings
185
+
186
+ self.qkv_proj = QKVParallelLinear(hidden_size,
187
+ self.head_dim,
188
+ self.total_num_heads,
189
+ self.total_num_kv_heads,
190
+ bias=qkv_bias,
191
+ quant_config=quant_config,
192
+ prefix=f"{prefix}.qkv_proj")
193
+
194
+ self.o_proj = RowParallelLinear(self.total_num_heads * self.head_dim,
195
+ hidden_size,
196
+ bias=False,
197
+ quant_config=quant_config,
198
+ prefix=f"{prefix}.o_proj")
199
+
200
+ self.rotary_emb = get_rope(
201
+ self.head_dim,
202
+ rotary_dim=self.head_dim,
203
+ max_position=max_position_embeddings,
204
+ base=rope_theta,
205
+ rope_scaling=rope_scaling,
206
+ )
207
+ self.attn = Attention(self.num_heads,
208
+ self.head_dim,
209
+ self.scaling,
210
+ num_kv_heads=self.num_kv_heads,
211
+ cache_config=cache_config,
212
+ quant_config=quant_config,
213
+ prefix=f"{prefix}.attn")
214
+
215
+ self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
216
+ self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
217
+
218
+ def forward(
219
+ self,
220
+ positions: torch.Tensor,
221
+ hidden_states: torch.Tensor,
222
+ ) -> torch.Tensor:
223
+ qkv, _ = self.qkv_proj(hidden_states)
224
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
225
+ # Add qk-norm
226
+ q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim,
227
+ self.head_dim)
228
+ q_by_head = self.q_norm.forward_native(q_by_head)
229
+ q = q_by_head.view(q.shape)
230
+
231
+ k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim,
232
+ self.head_dim)
233
+ k_by_head = self.k_norm.forward_native(k_by_head)
234
+ k = k_by_head.view(k.shape)
235
+ q, k = self.rotary_emb(positions, q, k)
236
+ attn_output = self.attn(q, k, v)
237
+ output, _ = self.o_proj(attn_output)
238
+ return output
239
+
240
+
241
+class Qwen3MoeDecoderLayer(nn.Module):
242
+
243
+ def __init__(
244
+ self,
245
+ config: PretrainedConfig,
246
+ cache_config: Optional[CacheConfig] = None,
247
+ quant_config: Optional[QuantizationConfig] = None,
248
+ prefix: str = "",
249
+ ) -> None:
250
+ super().__init__()
251
+ self.hidden_size = config.hidden_size
252
+ rope_theta = getattr(config, "rope_theta", 10000)
253
+ rope_scaling = getattr(config, "rope_scaling", None)
254
+ max_position_embeddings = getattr(config, "max_position_embeddings",
255
+ 8192)
256
+ self.self_attn = Qwen3MoeAttention(
257
+ hidden_size=self.hidden_size,
258
+ num_heads=config.num_attention_heads,
259
+ num_kv_heads=config.num_key_value_heads,
260
+ rope_theta=rope_theta,
261
+ rope_scaling=rope_scaling,
262
+ max_position_embeddings=max_position_embeddings,
263
+ rms_norm_eps=config.rms_norm_eps,
264
+ qkv_bias=getattr(config, 'attention_bias', False),
265
+ head_dim=getattr(config, 'head_dim', None),
266
+ cache_config=cache_config,
267
+ quant_config=quant_config,
268
+ prefix=f"{prefix}.self_attn",
269
+ )
270
+
271
+ # `mlp_only_layers` in the config.
272
+ layer_idx = extract_layer_index(prefix)
273
+ mlp_only_layers = ([] if not hasattr(config, "mlp_only_layers") else
274
+ config.mlp_only_layers)
275
+ if (layer_idx not in mlp_only_layers) and (
276
+ config.num_experts > 0 and
277
+ (layer_idx + 1) % config.decoder_sparse_step == 0):
278
+ self.mlp = Qwen3MoeSparseMoeBlock(config=config,
279
+ quant_config=quant_config,
280
+ prefix=f"{prefix}.mlp")
281
+ else:
282
+ self.mlp = Qwen3MoeMLP(hidden_size=config.hidden_size,
283
+ intermediate_size=config.intermediate_size,
284
+ hidden_act=config.hidden_act,
285
+ quant_config=quant_config,
286
+ prefix=f"{prefix}.mlp")
287
+ self.input_layernorm = RMSNorm(config.hidden_size,
288
+ eps=config.rms_norm_eps)
289
+ self.post_attention_layernorm = RMSNorm(config.hidden_size,
290
+ eps=config.rms_norm_eps)
291
+
292
+ def forward(
293
+ self,
294
+ positions: torch.Tensor,
295
+ hidden_states: torch.Tensor,
296
+ residual: Optional[torch.Tensor],
297
+ ) -> torch.Tensor:
298
+ # Self Attention
299
+ if residual is None:
300
+ residual = hidden_states
301
+ hidden_states = self.input_layernorm(hidden_states)
302
+ else:
303
+ hidden_states, residual = self.input_layernorm(
304
+ hidden_states, residual)
305
+ hidden_states = self.self_attn(
306
+ positions=positions,
307
+ hidden_states=hidden_states,
308
+ )
309
+
310
+ # Fully Connected
311
+ hidden_states, residual = self.post_attention_layernorm(
312
+ hidden_states, residual)
313
+ hidden_states = self.mlp(hidden_states)
314
+ return hidden_states, residual
315
+
316
+
317
+@support_torch_compile
318
+class Qwen3MoeModel(nn.Module):
319
+
320
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
321
+ super().__init__()
322
+
323
+ config = vllm_config.model_config.hf_config
324
+ cache_config = vllm_config.cache_config
325
+ quant_config = vllm_config.quant_config
326
+
327
+ self.padding_idx = config.pad_token_id
328
+ self.vocab_size = config.vocab_size
329
+
330
+ self.embed_tokens = VocabParallelEmbedding(
331
+ config.vocab_size,
332
+ config.hidden_size,
333
+ prefix=f"{prefix}.embed_tokens")
334
+ self.start_layer, self.end_layer, self.layers = make_layers(
335
+ config.num_hidden_layers,
336
+ lambda prefix: Qwen3MoeDecoderLayer(config=config,
337
+ cache_config=cache_config,
338
+ quant_config=quant_config,
339
+ prefix=prefix),
340
+ prefix=f"{prefix}.layers",
341
+ )
342
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
343
+ self.make_empty_intermediate_tensors = (
344
+ make_empty_intermediate_tensors_factory(
345
+ ["hidden_states", "residual"], config.hidden_size))
346
+
347
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
348
+ return self.embed_tokens(input_ids)
349
+
350
+ def forward(
351
+ self,
352
+ input_ids: torch.Tensor,
353
+ positions: torch.Tensor,
354
+ intermediate_tensors: Optional[IntermediateTensors] = None,
355
+ inputs_embeds: Optional[torch.Tensor] = None,
356
+ ) -> Union[torch.Tensor, IntermediateTensors]:
357
+ if get_pp_group().is_first_rank:
358
+ if inputs_embeds is not None:
359
+ hidden_states = inputs_embeds
360
+ else:
361
+ hidden_states = self.get_input_embeddings(input_ids)
362
+ residual = None
363
+ else:
364
+ assert intermediate_tensors is not None
365
+ hidden_states = intermediate_tensors["hidden_states"]
366
+ residual = intermediate_tensors["residual"]
367
+ for i in range(self.start_layer, self.end_layer):
368
+ layer = self.layers[i]
369
+ hidden_states, residual = layer(positions, hidden_states, residual)
370
+ if not get_pp_group().is_last_rank:
371
+ return IntermediateTensors({
372
+ "hidden_states": hidden_states,
373
+ "residual": residual
374
+ })
375
+ hidden_states, _ = self.norm(hidden_states, residual)
376
+ return hidden_states
377
+
378
+
379
+class Qwen3MoeForCausalLM(nn.Module, SupportsPP):
380
+
381
+ fall_back_to_pt_during_load = False
382
+
383
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
384
+ super().__init__()
385
+ config = vllm_config.model_config.hf_config
386
+ quant_config = vllm_config.quant_config
387
+ self.config = config
388
+ self.quant_config = quant_config
389
+ self.model = Qwen3MoeModel(vllm_config=vllm_config,
390
+ prefix=maybe_prefix(prefix, "model"))
391
+ self.lm_head = ParallelLMHead(config.vocab_size,
392
+ config.hidden_size,
393
+ quant_config=quant_config)
394
+ if self.config.tie_word_embeddings:
395
+ self.lm_head.weight = self.model.embed_tokens.weight
396
+ self.logits_processor = LogitsProcessor(config.vocab_size)
397
+ self.sampler = get_sampler()
398
+ self.make_empty_intermediate_tensors = (
399
+ self.model.make_empty_intermediate_tensors)
400
+
401
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
402
+ return self.model.get_input_embeddings(input_ids)
403
+
404
+ def forward(
405
+ self,
406
+ input_ids: torch.Tensor,
407
+ positions: torch.Tensor,
408
+ intermediate_tensors: Optional[IntermediateTensors] = None,
409
+ inputs_embeds: Optional[torch.Tensor] = None,
410
+ ) -> Union[torch.Tensor, IntermediateTensors]:
411
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
412
+ inputs_embeds)
413
+ return hidden_states
414
+
415
+ def compute_logits(
416
+ self,
417
+ hidden_states: torch.Tensor,
418
+ sampling_metadata: SamplingMetadata,
419
+ ) -> Optional[torch.Tensor]:
420
+ logits = self.logits_processor(self.lm_head, hidden_states,
421
+ sampling_metadata)
422
+ return logits
423
+
424
+ def sample(
425
+ self,
426
+ logits: Optional[torch.Tensor],
427
+ sampling_metadata: SamplingMetadata,
428
+ ) -> Optional[SamplerOutput]:
429
+ next_tokens = self.sampler(logits, sampling_metadata)
430
+ return next_tokens
431
+
432
+ def load_weights(self, weights: Iterable[Tuple[str,
433
+ torch.Tensor]]) -> Set[str]:
434
+ stacked_params_mapping = [
435
+ # (param_name, shard_name, shard_id)
436
+ ("qkv_proj", "q_proj", "q"),
437
+ ("qkv_proj", "k_proj", "k"),
438
+ ("qkv_proj", "v_proj", "v"),
439
+ ("gate_up_proj", "gate_proj", 0),
440
+ ("gate_up_proj", "up_proj", 1),
441
+ ]
442
+
443
+ # Params for weights, fp8 weight scales, fp8 activation scales
444
+ # (param_name, weight_name, expert_id, shard_id)
445
+ expert_params_mapping = FusedMoE.make_expert_params_mapping(
446
+ ckpt_gate_proj_name="gate_proj",
447
+ ckpt_down_proj_name="down_proj",
448
+ ckpt_up_proj_name="up_proj",
449
+ num_experts=self.config.num_experts)
450
+
451
+ params_dict = dict(self.named_parameters())
452
+ loaded_params: Set[str] = set()
453
+ for name, loaded_weight in weights:
454
+ if "rotary_emb.inv_freq" in name:
455
+ continue
456
+ for (param_name, weight_name, shard_id) in stacked_params_mapping:
457
+ # Skip non-stacked layers and experts (experts handled below).
458
+ if weight_name not in name:
459
+ continue
460
+ # We have mlp.experts[0].gate_proj in the checkpoint.
461
+ # Since we handle the experts below in expert_params_mapping,
462
+ # we need to skip here BEFORE we update the name, otherwise
463
+ # name will be updated to mlp.experts[0].gate_up_proj, which
464
+ # will then be updated below in expert_params_mapping
465
+ # for mlp.experts[0].gate_gate_up_proj, which breaks load.
466
+ if "mlp.experts" in name:
467
+ continue
468
+ name = name.replace(weight_name, param_name)
469
+ # Skip loading extra bias for GPTQ models.
470
+ if ((name.endswith(".bias") or name.endswith("_bias"))
471
+ and name not in params_dict):
472
+ continue
473
+ # Skip layers on other devices.
474
+ if is_pp_missing_parameter(name, self):
475
+ continue
476
+ if name not in params_dict:
477
+ continue
478
+
479
+ param = params_dict[name]
480
+ weight_loader = param.weight_loader
481
+ weight_loader(param, loaded_weight, shard_id)
482
+ break
483
+ else:
484
+ for mapping in expert_params_mapping:
485
+ param_name, weight_name, expert_id, shard_id = mapping
486
+ if weight_name not in name:
487
+ continue
488
+ name = name.replace(weight_name, param_name)
489
+ # Skip layers on other devices.
490
+ if is_pp_missing_parameter(name, self):
491
+ continue
492
+ # Skip loading extra bias for GPTQ models.
493
+ if ((name.endswith(".bias") or name.endswith("_bias"))
494
+ and name not in params_dict):
495
+ continue
496
+ param = params_dict[name]
497
+ weight_loader = param.weight_loader
498
+ weight_loader(param,
499
+ loaded_weight,
500
+ name,
501
+ shard_id=shard_id,
502
+ expert_id=expert_id)
503
+ break
504
+ else:
505
+ # Skip loading extra bias for GPTQ models.
506
+ if ((name.endswith(".bias") or name.endswith("_bias"))
507
+ and name not in params_dict):
508
+ continue
509
+ # Skip layers on other devices.
510
+ if is_pp_missing_parameter(name, self):
511
+ continue
512
+ # Remapping the name of FP8 kv-scale.
513
+ if name.endswith("kv_scale"):
514
+ remapped_kv_scale_name = name.replace(
515
+ ".kv_scale", ".attn.kv_scale")
516
+ if remapped_kv_scale_name not in params_dict:
517
+ logger.warning_once(
518
+ "Found kv scale in the checkpoint "
519
+ f"(e.g. {name}), but not found the expected "
520
+ f"name in the model "
521
+ f"(e.g. {remapped_kv_scale_name}). "
522
+ "kv-scale is not loaded.")
523
+ continue
524
+ else:
525
+ name = remapped_kv_scale_name
526
+ param = params_dict[name]
527
+ weight_loader = getattr(param, "weight_loader",
528
+ default_weight_loader)
529
+ weight_loader(param, loaded_weight)
530
+ loaded_params.add(name)
531
+ return loaded_params
vllm/model_executor/models/registry.py
+2
@@ -100,6 +100,8 @@ _TEXT_GENERATION_MODELS = {
100
"QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
101
"Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
102
"Qwen2MoeForCausalLM": ("qwen2_moe", "Qwen2MoeForCausalLM"),
103
+ "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),
104
+ "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),
105
"RWForCausalLM": ("falcon", "FalconForCausalLM"),
106
"StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),
107
"StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),