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 13, 14:30 UTC
[Model] Consolidate Deepseek-MoE implementation with DeepSeek-v2 (#28101)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Isotr0py committed
Nov 8, 2025 at 13:01 UTC
934a9c3b79e6cb860a8d23b7f317a5f63adf0fae
6 files changed
+144
-548
tests/models/registry.py
+4
-1
@@ -219,7 +219,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
219
"nvidia/Llama-3_3-Nemotron-Super-49B-v1",
220
trust_remote_code=True,
221
),
222
- "DeepseekForCausalLM": _HfExamplesInfo("deepseek-ai/deepseek-llm-7b-chat"),
222
+ "DeepseekForCausalLM": _HfExamplesInfo(
223
+ "deepseek-ai/deepseek-moe-16b-base",
224
+ trust_remote_code=True,
225
+ ),
226
"DeepseekV2ForCausalLM": _HfExamplesInfo(
227
"deepseek-ai/DeepSeek-V2-Lite-Chat",
228
trust_remote_code=True,
vllm/model_executor/models/deepseek.py
deleted
-517
@@ -1,517 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-# Adapted from
5
-# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
6
-# Copyright 2023 The vLLM team.
7
-# Copyright 2023 DeepSeek-AI and the HuggingFace Inc. team. All rights reserved.
8
-#
9
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10
-# and OPT implementations in this library. It has been modified from its
11
-# original forms to accommodate minor architectural differences compared
12
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13
-#
14
-# Licensed under the Apache License, Version 2.0 (the "License");
15
-# you may not use this file except in compliance with the License.
16
-# You may obtain a copy of the License at
17
-#
18
-# http://www.apache.org/licenses/LICENSE-2.0
19
-#
20
-# Unless required by applicable law or agreed to in writing, software
21
-# distributed under the License is distributed on an "AS IS" BASIS,
22
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23
-# See the License for the specific language governing permissions and
24
-# limitations under the License.
25
-"""Inference-only Deepseek model."""
26
-
27
-from collections.abc import Iterable
28
-from itertools import islice
29
-from typing import Any
30
-
31
-import torch
32
-from torch import nn
33
-from transformers import PretrainedConfig
34
-
35
-from vllm.attention import Attention
36
-from vllm.config import CacheConfig, VllmConfig
37
-from vllm.distributed import (
38
- get_pp_group,
39
- get_tensor_model_parallel_rank,
40
- get_tensor_model_parallel_world_size,
41
- tensor_model_parallel_all_reduce,
42
-)
43
-from vllm.model_executor.layers.activation import SiluAndMul
44
-from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
45
-from vllm.model_executor.layers.layernorm import RMSNorm
46
-from vllm.model_executor.layers.linear import (
47
- MergedColumnParallelLinear,
48
- QKVParallelLinear,
49
- ReplicatedLinear,
50
- RowParallelLinear,
51
-)
52
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
53
-from vllm.model_executor.layers.quantization import QuantizationConfig
54
-from vllm.model_executor.layers.rotary_embedding import get_rope
55
-from vllm.model_executor.layers.vocab_parallel_embedding import (
56
- ParallelLMHead,
57
- VocabParallelEmbedding,
58
-)
59
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
60
-from vllm.sequence import IntermediateTensors
61
-
62
-from .interfaces import SupportsLoRA, SupportsPP
63
-from .utils import (
64
- AutoWeightsLoader,
65
- extract_layer_index,
66
- is_pp_missing_parameter,
67
- make_empty_intermediate_tensors_factory,
68
- make_layers,
69
- maybe_prefix,
70
-)
71
-
72
-
73
-class DeepseekMLP(nn.Module):
74
- def __init__(
75
- self,
76
- hidden_size: int,
77
- intermediate_size: int,
78
- hidden_act: str,
79
- quant_config: QuantizationConfig | None = None,
80
- reduce_results: bool = True,
81
- prefix: str = "",
82
- ) -> None:
83
- super().__init__()
84
- self.gate_up_proj = MergedColumnParallelLinear(
85
- hidden_size,
86
- [intermediate_size] * 2,
87
- bias=False,
88
- quant_config=quant_config,
89
- prefix=f"{prefix}.gate_up_proj",
90
- )
91
- self.down_proj = RowParallelLinear(
92
- intermediate_size,
93
- hidden_size,
94
- bias=False,
95
- quant_config=quant_config,
96
- reduce_results=reduce_results,
97
- prefix=f"{prefix}.down_proj",
98
- )
99
- if hidden_act != "silu":
100
- raise ValueError(
101
- f"Unsupported activation: {hidden_act}. Only silu is supported for now."
102
- )
103
- self.act_fn = SiluAndMul()
104
-
105
- def forward(self, x):
106
- gate_up, _ = self.gate_up_proj(x)
107
- x = self.act_fn(gate_up)
108
- x, _ = self.down_proj(x)
109
- return x
110
-
111
-
112
-class DeepseekMoE(nn.Module):
113
- def __init__(
114
- self,
115
- config: PretrainedConfig,
116
- quant_config: QuantizationConfig | None = None,
117
- prefix: str = "",
118
- ):
119
- super().__init__()
120
- self.config = config
121
- self.rank = get_tensor_model_parallel_rank()
122
- self.tp_size = get_tensor_model_parallel_world_size()
123
- self.n_routed_experts = config.n_routed_experts
124
- self.top_k = config.num_experts_per_tok
125
- if self.tp_size > self.n_routed_experts:
126
- raise ValueError(
127
- f"Tensor parallel size {self.tp_size} is greater than "
128
- f"the number of experts {self.n_routed_experts}."
129
- )
130
-
131
- self.experts = nn.ModuleList(
132
- [
133
- DeepseekMLP(
134
- hidden_size=config.hidden_size,
135
- intermediate_size=config.moe_intermediate_size,
136
- hidden_act=config.hidden_act,
137
- quant_config=quant_config,
138
- reduce_results=False,
139
- )
140
- for idx in range(self.n_routed_experts)
141
- ]
142
- )
143
- self.pack_params()
144
-
145
- self.gate = ReplicatedLinear(
146
- config.hidden_size, self.n_routed_experts, bias=False, quant_config=None
147
- )
148
-
149
- if config.n_shared_experts is not None:
150
- intermediate_size = config.moe_intermediate_size * config.n_shared_experts
151
- self.shared_experts = DeepseekMLP(
152
- hidden_size=config.hidden_size,
153
- intermediate_size=intermediate_size,
154
- hidden_act=config.hidden_act,
155
- quant_config=quant_config,
156
- reduce_results=False,
157
- )
158
-
159
- def pack_params(self):
160
- w1 = []
161
- w2 = []
162
- for expert in self.experts:
163
- w1.append(expert.gate_up_proj.weight)
164
- w2.append(expert.down_proj.weight)
165
- self.w1 = torch._utils._flatten_dense_tensors(w1)
166
- w1s = torch._utils._unflatten_dense_tensors(self.w1, w1)
167
- for data, param in zip(w1s, w1):
168
- param.data = data
169
- self.w1 = self.w1.view(len(w1), *w1s[0].shape)
170
-
171
- self.w2 = torch._utils._flatten_dense_tensors(w2)
172
- w2s = torch._utils._unflatten_dense_tensors(self.w2, w2)
173
- for data, param in zip(w2s, w2):
174
- param.data = data
175
-
176
- self.w2 = self.w2.view(len(w2), *w2s[0].shape)
177
-
178
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
179
- num_tokens, hidden_dim = hidden_states.shape
180
- hidden_states = hidden_states.view(-1, hidden_dim)
181
- if self.config.n_shared_experts is not None:
182
- shared_output = self.shared_experts(hidden_states)
183
- # router_logits: (num_tokens, n_experts)
184
- router_logits, _ = self.gate(hidden_states)
185
-
186
- topk_weights, topk_ids, _ = fused_topk(
187
- hidden_states,
188
- router_logits,
189
- self.top_k,
190
- renormalize=self.config.norm_topk_prob,
191
- )
192
-
193
- final_hidden_states = fused_experts(
194
- hidden_states, self.w1, self.w2, topk_weights, topk_ids, inplace=True
195
- )
196
-
197
- if self.config.n_shared_experts is not None:
198
- final_hidden_states = final_hidden_states + shared_output
199
- final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
200
-
201
- return final_hidden_states.view(num_tokens, hidden_dim)
202
-
203
-
204
-class DeepseekAttention(nn.Module):
205
- def __init__(
206
- self,
207
- hidden_size: int,
208
- num_heads: int,
209
- num_kv_heads: int,
210
- rope_theta: float = 10000,
211
- rope_scaling: dict[str, Any] | None = None,
212
- max_position_embeddings: int = 8192,
213
- cache_config: CacheConfig | None = None,
214
- quant_config: QuantizationConfig | None = None,
215
- prefix: str = "",
216
- ) -> None:
217
- super().__init__()
218
- self.hidden_size = hidden_size
219
- tp_size = get_tensor_model_parallel_world_size()
220
- self.total_num_heads = num_heads
221
- assert self.total_num_heads % tp_size == 0
222
- self.num_heads = self.total_num_heads // tp_size
223
- self.total_num_kv_heads = num_kv_heads
224
- if self.total_num_kv_heads >= tp_size:
225
- # Number of KV heads is greater than TP size, so we partition
226
- # the KV heads across multiple tensor parallel GPUs.
227
- assert self.total_num_kv_heads % tp_size == 0
228
- else:
229
- # Number of KV heads is less than TP size, so we replicate
230
- # the KV heads across multiple tensor parallel GPUs.
231
- assert tp_size % self.total_num_kv_heads == 0
232
- self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
233
- self.head_dim = hidden_size // self.total_num_heads
234
- self.q_size = self.num_heads * self.head_dim
235
- self.kv_size = self.num_kv_heads * self.head_dim
236
- self.scaling = self.head_dim**-0.5
237
- self.rope_theta = rope_theta
238
- self.max_position_embeddings = max_position_embeddings
239
-
240
- self.qkv_proj = QKVParallelLinear(
241
- hidden_size,
242
- self.head_dim,
243
- self.total_num_heads,
244
- self.total_num_kv_heads,
245
- bias=False,
246
- quant_config=quant_config,
247
- prefix=f"{prefix}.qkv_proj",
248
- )
249
-
250
- self.o_proj = RowParallelLinear(
251
- self.total_num_heads * self.head_dim,
252
- hidden_size,
253
- bias=False,
254
- quant_config=quant_config,
255
- prefix=f"{prefix}.o_proj",
256
- )
257
-
258
- self.rotary_emb = get_rope(
259
- self.head_dim,
260
- rotary_dim=self.head_dim,
261
- max_position=max_position_embeddings,
262
- base=rope_theta,
263
- rope_scaling=rope_scaling,
264
- )
265
- self.attn = Attention(
266
- self.num_heads,
267
- self.head_dim,
268
- self.scaling,
269
- num_kv_heads=self.num_kv_heads,
270
- cache_config=cache_config,
271
- quant_config=quant_config,
272
- prefix=f"{prefix}.attn",
273
- )
274
-
275
- def forward(
276
- self,
277
- positions: torch.Tensor,
278
- hidden_states: torch.Tensor,
279
- ) -> torch.Tensor:
280
- qkv, _ = self.qkv_proj(hidden_states)
281
- q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
282
- q, k = self.rotary_emb(positions, q, k)
283
- attn_output = self.attn(q, k, v)
284
- output, _ = self.o_proj(attn_output)
285
- return output
286
-
287
-
288
-class DeepseekDecoderLayer(nn.Module):
289
- def __init__(
290
- self,
291
- config: PretrainedConfig,
292
- cache_config: CacheConfig | None = None,
293
- quant_config: QuantizationConfig | None = None,
294
- prefix: str = "",
295
- ) -> None:
296
- super().__init__()
297
- layer_idx = extract_layer_index(prefix)
298
- self.hidden_size = config.hidden_size
299
- rope_theta = getattr(config, "rope_theta", 10000)
300
- rope_scaling = getattr(config, "rope_scaling", None)
301
- max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
302
- moe_layer_freq = getattr(config, "moe_layer_freq", 1)
303
- self.self_attn = DeepseekAttention(
304
- hidden_size=self.hidden_size,
305
- num_heads=config.num_attention_heads,
306
- num_kv_heads=config.num_key_value_heads,
307
- rope_theta=rope_theta,
308
- rope_scaling=rope_scaling,
309
- max_position_embeddings=max_position_embeddings,
310
- cache_config=cache_config,
311
- quant_config=quant_config,
312
- prefix=f"{prefix}.self_attn",
313
- )
314
- if (
315
- config.n_routed_experts is not None
316
- and layer_idx >= config.first_k_dense_replace
317
- and layer_idx % moe_layer_freq == 0
318
- ):
319
- self.mlp = DeepseekMoE(
320
- config=config, quant_config=quant_config, prefix=f"{prefix}.mlp"
321
- )
322
- else:
323
- self.mlp = DeepseekMLP(
324
- hidden_size=config.hidden_size,
325
- intermediate_size=config.intermediate_size,
326
- hidden_act=config.hidden_act,
327
- quant_config=quant_config,
328
- prefix=f"{prefix}.mlp",
329
- )
330
- self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
331
- self.post_attention_layernorm = RMSNorm(
332
- config.hidden_size, eps=config.rms_norm_eps
333
- )
334
-
335
- def forward(
336
- self,
337
- positions: torch.Tensor,
338
- hidden_states: torch.Tensor,
339
- residual: torch.Tensor | None,
340
- ) -> torch.Tensor:
341
- # Self Attention
342
- if residual is None:
343
- residual = hidden_states
344
- hidden_states = self.input_layernorm(hidden_states)
345
- else:
346
- hidden_states, residual = self.input_layernorm(hidden_states, residual)
347
- hidden_states = self.self_attn(
348
- positions=positions,
349
- hidden_states=hidden_states,
350
- )
351
-
352
- # Fully Connected
353
- hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
354
- hidden_states = self.mlp(hidden_states)
355
- return hidden_states, residual
356
-
357
-
358
-class DeepseekModel(nn.Module):
359
- fall_back_to_pt_during_load = False
360
-
361
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
362
- super().__init__()
363
-
364
- config = vllm_config.model_config.hf_config
365
- cache_config = vllm_config.cache_config
366
- quant_config = vllm_config.quant_config
367
-
368
- self.vocab_size = config.vocab_size
369
-
370
- self.embed_tokens = VocabParallelEmbedding(
371
- config.vocab_size,
372
- config.hidden_size,
373
- )
374
- self.start_layer, self.end_layer, self.layers = make_layers(
375
- config.num_hidden_layers,
376
- lambda prefix: DeepseekDecoderLayer(
377
- config, cache_config, quant_config=quant_config, prefix=prefix
378
- ),
379
- prefix=f"{prefix}.layers",
380
- )
381
- self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
382
- self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
383
- ["hidden_states", "residual"], config.hidden_size
384
- )
385
-
386
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
387
- return self.embed_tokens(input_ids)
388
-
389
- def forward(
390
- self,
391
- input_ids: torch.Tensor,
392
- positions: torch.Tensor,
393
- intermediate_tensors: IntermediateTensors | None,
394
- inputs_embeds: torch.Tensor | None = None,
395
- ) -> torch.Tensor | IntermediateTensors:
396
- if get_pp_group().is_first_rank:
397
- if inputs_embeds is not None:
398
- hidden_states = inputs_embeds
399
- else:
400
- hidden_states = self.get_input_embeddings(input_ids)
401
- residual = None
402
- else:
403
- hidden_states = intermediate_tensors["hidden_states"]
404
- residual = intermediate_tensors["residual"]
405
- for layer in islice(self.layers, self.start_layer, self.end_layer):
406
- hidden_states, residual = layer(positions, hidden_states, residual)
407
- if not get_pp_group().is_last_rank:
408
- return IntermediateTensors(
409
- {"hidden_states": hidden_states, "residual": residual}
410
- )
411
- hidden_states, _ = self.norm(hidden_states, residual)
412
- return hidden_states
413
-
414
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
415
- stacked_params_mapping = [
416
- # (param_name, shard_name, shard_id)
417
- ("qkv_proj", "q_proj", "q"),
418
- ("qkv_proj", "k_proj", "k"),
419
- ("qkv_proj", "v_proj", "v"),
420
- ("gate_up_proj", "gate_proj", 0),
421
- ("gate_up_proj", "up_proj", 1),
422
- ]
423
-
424
- params_dict = dict(self.named_parameters())
425
- loaded_params: set[str] = set()
426
- for name, loaded_weight in weights:
427
- if "rotary_emb.inv_freq" in name:
428
- continue
429
- for param_name, weight_name, shard_id in stacked_params_mapping:
430
- if weight_name not in name:
431
- continue
432
- name = name.replace(weight_name, param_name)
433
- # Skip loading extra bias for GPTQ models.
434
- if name.endswith(".bias") and name not in params_dict:
435
- continue
436
- # Skip experts that are not assigned to this worker.
437
- if (
438
- "mlp.experts." in name or "mlp.shared_experts." in name
439
- ) and name not in params_dict:
440
- continue
441
- if is_pp_missing_parameter(name, self):
442
- continue
443
- param = params_dict[name]
444
- weight_loader = param.weight_loader
445
- weight_loader(param, loaded_weight, shard_id)
446
- break
447
- else:
448
- # Skip loading extra bias for GPTQ models.
449
- if name.endswith(".bias") and name not in params_dict:
450
- continue
451
- # Skip experts that are not assigned to this worker.
452
- if (
453
- "mlp.experts." in name or "mlp.shared_experts." in name
454
- ) and name not in params_dict:
455
- continue
456
- if is_pp_missing_parameter(name, self):
457
- continue
458
- param = params_dict[name]
459
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
460
- weight_loader(param, loaded_weight)
461
- loaded_params.add(name)
462
- return loaded_params
463
-
464
-
465
-class DeepseekForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
466
- packed_modules_mapping = {
467
- "qkv_proj": ["q_proj", "k_proj", "v_proj"],
468
- "gate_up_proj": ["gate_proj", "up_proj"],
469
- }
470
-
471
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
472
- super().__init__()
473
- config = vllm_config.model_config.hf_config
474
- quant_config = vllm_config.quant_config
475
- self.config = config
476
- self.quant_config = quant_config
477
- self.model = DeepseekModel(
478
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
479
- )
480
- self.lm_head = ParallelLMHead(
481
- config.vocab_size,
482
- config.hidden_size,
483
- quant_config=quant_config,
484
- prefix=maybe_prefix(prefix, "lm_head"),
485
- )
486
- if self.config.tie_word_embeddings:
487
- self.lm_head.weight = self.model.embed_tokens.weight
488
- self.logits_processor = LogitsProcessor(config.vocab_size)
489
- self.make_empty_intermediate_tensors = (
490
- self.model.make_empty_intermediate_tensors
491
- )
492
-
493
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
494
- return self.model.get_input_embeddings(input_ids)
495
-
496
- def forward(
497
- self,
498
- input_ids: torch.Tensor,
499
- positions: torch.Tensor,
500
- intermediate_tensors: IntermediateTensors | None = None,
501
- inputs_embeds: torch.Tensor | None = None,
502
- ) -> torch.Tensor | IntermediateTensors:
503
- hidden_states = self.model(
504
- input_ids, positions, intermediate_tensors, inputs_embeds
505
- )
506
- return hidden_states
507
-
508
- def compute_logits(
509
- self,
510
- hidden_states: torch.Tensor,
511
- ) -> torch.Tensor | None:
512
- logits = self.logits_processor(self.lm_head, hidden_states)
513
- return logits
514
-
515
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
516
- loader = AutoWeightsLoader(self)
517
- return loader.load_weights(weights)
vllm/model_executor/models/deepseek_ocr.py
-8
@@ -417,18 +417,10 @@ class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
417
f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"
418
)
419
420
- if self.text_config.topk_method == "noaux_tc":
421
- architectures = ["DeepseekV3ForCausalLM"]
422
- elif not self.text_config.use_mla:
423
- architectures = ["DeepseekForCausalLM"]
424
- else:
425
- architectures = ["DeepseekV2ForCausalLM"]
426
-
420
self.language_model = init_vllm_registered_model(
421
vllm_config=vllm_config,
422
hf_config=self.text_config,
423
prefix=maybe_prefix(prefix, "language_model"),
431
- architectures=architectures,
424
)
425
426
self.make_empty_intermediate_tensors = (
vllm/model_executor/models/deepseek_v2.py
+139
-13
@@ -58,6 +58,7 @@ from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
58
from vllm.model_executor.layers.linear import (
59
ColumnParallelLinear,
60
MergedColumnParallelLinear,
61
+ QKVParallelLinear,
62
ReplicatedLinear,
63
RowParallelLinear,
64
)
@@ -104,6 +105,92 @@ elif current_platform.is_xpu():
105
logger = init_logger(__name__)
106
107
108
+class DeepseekAttention(nn.Module):
109
+ """Normal MHA implementation used by Deepseek v1."""
110
+
111
+ def __init__(
112
+ self,
113
+ vllm_config: VllmConfig,
114
+ config: DeepseekV2Config | DeepseekV3Config,
115
+ hidden_size: int,
116
+ num_heads: int,
117
+ rope_theta: float = 10000,
118
+ rope_scaling: dict[str, Any] | None = None,
119
+ max_position_embeddings: int = 8192,
120
+ cache_config: CacheConfig | None = None,
121
+ quant_config: QuantizationConfig | None = None,
122
+ prefix: str = "",
123
+ **kwargs,
124
+ ) -> None:
125
+ super().__init__()
126
+ self.hidden_size = hidden_size
127
+ tp_size = get_tensor_model_parallel_world_size()
128
+ self.total_num_heads = num_heads
129
+ assert self.total_num_heads % tp_size == 0
130
+ self.num_heads = self.total_num_heads // tp_size
131
+ self.total_num_kv_heads = config.num_key_value_heads
132
+ if self.total_num_kv_heads >= tp_size:
133
+ # Number of KV heads is greater than TP size, so we partition
134
+ # the KV heads across multiple tensor parallel GPUs.
135
+ assert self.total_num_kv_heads % tp_size == 0
136
+ else:
137
+ # Number of KV heads is less than TP size, so we replicate
138
+ # the KV heads across multiple tensor parallel GPUs.
139
+ assert tp_size % self.total_num_kv_heads == 0
140
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
141
+ self.head_dim = hidden_size // self.total_num_heads
142
+ self.q_size = self.num_heads * self.head_dim
143
+ self.kv_size = self.num_kv_heads * self.head_dim
144
+ self.scaling = self.head_dim**-0.5
145
+ self.rope_theta = rope_theta
146
+ self.max_position_embeddings = max_position_embeddings
147
+
148
+ self.qkv_proj = QKVParallelLinear(
149
+ hidden_size,
150
+ self.head_dim,
151
+ self.total_num_heads,
152
+ self.total_num_kv_heads,
153
+ bias=False,
154
+ quant_config=quant_config,
155
+ )
156
+
157
+ self.o_proj = RowParallelLinear(
158
+ self.total_num_heads * self.head_dim,
159
+ hidden_size,
160
+ bias=False,
161
+ quant_config=quant_config,
162
+ )
163
+
164
+ self.rotary_emb = get_rope(
165
+ self.head_dim,
166
+ rotary_dim=self.head_dim,
167
+ max_position=max_position_embeddings,
168
+ base=rope_theta,
169
+ rope_scaling=rope_scaling,
170
+ )
171
+ self.attn = Attention(
172
+ self.num_heads,
173
+ self.head_dim,
174
+ self.scaling,
175
+ num_kv_heads=self.num_kv_heads,
176
+ cache_config=cache_config,
177
+ quant_config=quant_config,
178
+ prefix=f"{prefix}.attn",
179
+ )
180
+
181
+ def forward(
182
+ self,
183
+ positions: torch.Tensor,
184
+ hidden_states: torch.Tensor,
185
+ ) -> torch.Tensor:
186
+ qkv, _ = self.qkv_proj(hidden_states)
187
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
188
+ q, k = self.rotary_emb(positions, q, k)
189
+ attn_output = self.attn(q, k, v)
190
+ output, _ = self.o_proj(attn_output)
191
+ return output
192
+
193
+
194
class DeepseekV2MLP(nn.Module):
195
def __init__(
196
self,
@@ -163,7 +250,7 @@ class DeepseekV2MoE(nn.Module):
250
self.tp_size = get_tensor_model_parallel_world_size()
251
self.tp_rank = get_tensor_model_parallel_rank()
252
166
- self.routed_scaling_factor = config.routed_scaling_factor
253
+ self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
254
255
self.ep_group = get_ep_group().device_group
256
self.ep_rank = get_ep_group().rank_in_group
@@ -186,7 +273,7 @@ class DeepseekV2MoE(nn.Module):
273
quant_config=None,
274
prefix=f"{prefix}.gate",
275
)
189
- if config.topk_method == "noaux_tc":
276
+ if getattr(config, "topk_method", None) == "noaux_tc":
277
self.gate.e_score_correction_bias = nn.Parameter(
278
torch.empty(config.n_routed_experts, dtype=torch.float32)
279
)
@@ -236,10 +323,10 @@ class DeepseekV2MoE(nn.Module):
323
renormalize=config.norm_topk_prob,
324
quant_config=quant_config,
325
use_grouped_topk=True,
239
- num_expert_group=config.n_group,
240
- topk_group=config.topk_group,
326
+ num_expert_group=getattr(config, "n_group", 1),
327
+ topk_group=getattr(config, "topk_group", 1),
328
prefix=f"{prefix}.experts",
242
- scoring_func=config.scoring_func,
329
+ scoring_func=getattr(config, "scoring_func", "softmax"),
330
# we do scaling outside, set factor to 1.0 to avoid double mul
331
# aiter applies routed_scaling_factor internally
332
routed_scaling_factor=1.0
@@ -999,7 +1086,19 @@ class DeepseekV2DecoderLayer(nn.Module):
1086
# with the layer's index.
1087
layer_idx = int(prefix.split(sep=".")[-1])
1088
self.layer_idx = layer_idx
1002
- if model_config.use_mla:
1089
+
1090
+ # verify MLA attention specific fields
1091
+ qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
1092
+ qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
1093
+ v_head_dim = getattr(config, "v_head_dim", 0)
1094
+ kv_lora_rank = getattr(config, "kv_lora_rank", 0)
1095
+ use_mha = config.model_type == "deepseek" or all(
1096
+ dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)
1097
+ )
1098
+
1099
+ if use_mha:
1100
+ attn_cls = DeepseekAttention
1101
+ elif model_config.use_mla:
1102
attn_cls = DeepseekV2MLAAttention
1103
else:
1104
attn_cls = DeepseekV2Attention
@@ -1008,11 +1107,11 @@ class DeepseekV2DecoderLayer(nn.Module):
1107
config=config,
1108
hidden_size=self.hidden_size,
1109
num_heads=config.num_attention_heads,
1011
- qk_nope_head_dim=config.qk_nope_head_dim,
1012
- qk_rope_head_dim=config.qk_rope_head_dim,
1013
- v_head_dim=config.v_head_dim,
1110
+ qk_nope_head_dim=qk_nope_head_dim,
1111
+ qk_rope_head_dim=qk_rope_head_dim,
1112
+ v_head_dim=v_head_dim,
1113
q_lora_rank=config.q_lora_rank if hasattr(config, "q_lora_rank") else None,
1015
- kv_lora_rank=config.kv_lora_rank,
1114
+ kv_lora_rank=kv_lora_rank,
1115
rope_theta=rope_theta,
1116
rope_scaling=rope_scaling,
1117
max_position_embeddings=max_position_embeddings,
@@ -1045,7 +1144,7 @@ class DeepseekV2DecoderLayer(nn.Module):
1144
self.post_attention_layernorm = RMSNorm(
1145
config.hidden_size, eps=config.rms_norm_eps
1146
)
1048
- self.routed_scaling_factor = config.routed_scaling_factor
1147
+ self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
1148
1149
def forward(
1150
self,
@@ -1064,7 +1163,10 @@ class DeepseekV2DecoderLayer(nn.Module):
1163
hidden_states=hidden_states,
1164
)
1165
1067
- if hidden_states.dtype == torch.float16:
1166
+ if (
1167
+ not isinstance(self.self_attn, DeepseekAttention)
1168
+ and hidden_states.dtype == torch.float16
1169
+ ):
1170
# Fix FP16 overflow
1171
# We scale both hidden_states and residual before
1172
# rmsnorm, and rmsnorm result would not affect by scale.
@@ -1227,6 +1329,15 @@ class DeepseekV2ForCausalLM(
1329
self.config = config
1330
self.quant_config = quant_config
1331
1332
+ qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
1333
+ qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
1334
+ self.use_mha = config.model_type == "deepseek" or all(
1335
+ dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)
1336
+ )
1337
+
1338
+ if self.use_mha:
1339
+ self.packed_modules_mapping["qkv_proj"] = ["q_proj", "k_proj", "v_proj"]
1340
+
1341
# `packed_modules_mapping` needs to be modified before
1342
# initializing DeepseekV2Model, as it is passed inplace to
1343
# quantization config init and may be used to select the
@@ -1265,7 +1376,7 @@ class DeepseekV2ForCausalLM(
1376
def set_moe_parameters(self):
1377
self.expert_weights = []
1378
1268
- self.num_expert_groups = self.config.n_group
1379
+ self.num_expert_groups = getattr(self.config, "n_group", 1)
1380
1381
self.moe_layers = []
1382
self.moe_mlp_layers = []
@@ -1321,9 +1432,20 @@ class DeepseekV2ForCausalLM(
1432
# (param_name, shard_name, shard_id)
1433
("gate_up_proj", "gate_proj", 0),
1434
("gate_up_proj", "up_proj", 1),
1435
+ ]
1436
+ mla_params_mapping = [
1437
("fused_qkv_a_proj", "q_a_proj", 0),
1438
("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
1439
]
1440
+ mha_params_mapping = [
1441
+ ("qkv_proj", "q_proj", "q"),
1442
+ ("qkv_proj", "k_proj", "k"),
1443
+ ("qkv_proj", "v_proj", "v"),
1444
+ ]
1445
+ if self.use_mha:
1446
+ stacked_params_mapping.extend(mha_params_mapping)
1447
+ else:
1448
+ stacked_params_mapping.extend(mla_params_mapping)
1449
1450
# Params for weights, fp8 weight scales, fp8 activation scales
1451
# (param_name, weight_name, expert_id, shard_id)
@@ -1506,6 +1628,10 @@ class DeepseekV2ForCausalLM(
1628
return loaded_params
1629
1630
1631
+class DeepseekForCausalLM(DeepseekV2ForCausalLM):
1632
+ pass
1633
+
1634
+
1635
class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
1636
pass
1637
vllm/model_executor/models/deepseek_vl2.py
-8
@@ -403,18 +403,10 @@ class DeepseekVLV2ForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
403
f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"
404
)
405
406
- if self.text_config.topk_method == "noaux_tc":
407
- architectures = ["DeepseekV3ForCausalLM"]
408
- elif not self.text_config.use_mla:
409
- architectures = ["DeepseekForCausalLM"]
410
- else:
411
- architectures = ["DeepseekV2ForCausalLM"]
412
-
406
self.language_model = init_vllm_registered_model(
407
vllm_config=vllm_config,
408
hf_config=self.text_config,
409
prefix=maybe_prefix(prefix, "language"),
417
- architectures=architectures,
410
)
411
412
self.make_empty_intermediate_tensors = (
vllm/model_executor/models/registry.py
+1
-1
@@ -76,7 +76,7 @@ _TEXT_GENERATION_MODELS = {
76
"CwmForCausalLM": ("llama", "LlamaForCausalLM"),
77
"DbrxForCausalLM": ("dbrx", "DbrxForCausalLM"),
78
"DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
79
- "DeepseekForCausalLM": ("deepseek", "DeepseekForCausalLM"),
79
+ "DeepseekForCausalLM": ("deepseek_v2", "DeepseekForCausalLM"),
80
"DeepseekV2ForCausalLM": ("deepseek_v2", "DeepseekV2ForCausalLM"),
81
"DeepseekV3ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),
82
"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),