deepseek_v2.py 54.7 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

wangding zeng's avatar
wangding zeng committed
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/llama/modeling_llama.py
# Copyright 2023 The vLLM team.
# Copyright 2023 DeepSeek-AI and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
25
"""Inference-only DeepseekV2/DeepseekV3 model."""
王敏's avatar
王敏 committed
26
27
import os
import re
28
import vllm.envs as envs
zhuwenwen's avatar
zhuwenwen committed
29

30
31
import typing
from collections.abc import Callable, Iterable
32
from typing import Any, Optional, Union
wangding zeng's avatar
wangding zeng committed
33
34
35
36
37

import torch
from torch import nn
from transformers import PretrainedConfig

38
from vllm.attention import Attention
39
from vllm.compilation.decorators import support_torch_compile
40
41
from vllm.config import (CacheConfig, ModelConfig, VllmConfig,
                         get_current_vllm_config)
王敏's avatar
王敏 committed
42
from vllm.distributed import (get_ep_group, get_pp_group, get_dp_group,
43
44
45
46
47
                              get_tensor_model_parallel_world_size,
                              tensor_model_parallel_reduce_scatter,
                              tensor_model_parallel_all_gather,
                              tensor_model_parallel_all_reduce,
                              get_tensor_model_parallel_rank)
wangding zeng's avatar
wangding zeng committed
48
from vllm.model_executor.layers.activation import SiluAndMul
49
50
from vllm.model_executor.layers.fused_moe import FusedMoE, SharedFusedMoE

王敏's avatar
王敏 committed
51
52
from vllm.model_executor.layers.fused_moe.mori_moe.layer import MoriMoE
from vllm.model_executor.layers.fused_moe.mori_moe.ep_moe_utlis import EPSharedExperts
wangding zeng's avatar
wangding zeng committed
53
54
55
56
57
58
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               MergedColumnParallelLinear,
                                               ReplicatedLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
59
from vllm.model_executor.layers.quantization import QuantizationConfig
wangding zeng's avatar
wangding zeng committed
60
61
62
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.vocab_parallel_embedding import (
    ParallelLMHead, VocabParallelEmbedding)
63
64
from vllm.model_executor.model_loader.weight_utils import (
    default_weight_loader, maybe_remap_kv_scale_name)
wangding zeng's avatar
wangding zeng committed
65
from vllm.model_executor.sampling_metadata import SamplingMetadata
66
from vllm.sequence import IntermediateTensors
wangding zeng's avatar
wangding zeng committed
67

68
from .interfaces import MixtureOfExperts, SupportsPP
69
from .utils import (PPMissingLayer, is_pp_missing_parameter,
70
71
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
王敏's avatar
王敏 committed
72
from vllm import _custom_ops as ops
73
from vllm.utils import W8a8GetCacheJSON
74

wangding zeng's avatar
wangding zeng committed
75
76
77
78
79
80
81
82
83
class DeepseekV2MLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
        quant_config: Optional[QuantizationConfig] = None,
        reduce_results: bool = True,
84
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
85
86
87
88
89
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size, [intermediate_size] * 2,
            bias=False,
90
91
            quant_config=quant_config,
            prefix=f"{prefix}.gate_up_proj")
wangding zeng's avatar
wangding zeng committed
92
93
94
95
        self.down_proj = RowParallelLinear(intermediate_size,
                                           hidden_size,
                                           bias=False,
                                           quant_config=quant_config,
96
97
                                           reduce_results=reduce_results,
                                           prefix=f"{prefix}.down_proj")
wangding zeng's avatar
wangding zeng committed
98
99
100
101
102
        if hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {hidden_act}. "
                             "Only silu is supported for now.")
        self.act_fn = SiluAndMul()

103
104
105
106
107
108
109
    def forward(self, x,
                rms_weight: Optional[torch.Tensor] = None,
                residual: Optional[torch.Tensor] = None,
                update_hd: Optional[bool] = False
                ):
        if envs.USE_FUSED_RMS_QUANT:
            gate_up, new_resi, _  = self.gate_up_proj(x, rms_weight, residual, update_hd=update_hd)
王敏's avatar
王敏 committed
110
111
112
113
114
115
            if envs.USE_FUSED_SILU_MUL_QUANT:
                x, _ = self.down_proj(gate_up, use_fused_silu_mul_quant=True)
            else:
                x = self.act_fn(gate_up)
                x, _ = self.down_proj(x)

116
117
118
119
120
121
            return x, new_resi
        else:
            gate_up, _ = self.gate_up_proj(x)
            x = self.act_fn(gate_up)
            x, _ = self.down_proj(x)
            return x
wangding zeng's avatar
wangding zeng committed
122
123
124
125
126
127
128
129


class DeepseekV2MoE(nn.Module):

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
130
        prefix: str = "",
131
        enable_eplb: bool = False,
wangding zeng's avatar
wangding zeng committed
132
133
134
135
    ):
        super().__init__()
        self.tp_size = get_tensor_model_parallel_world_size()
        self.routed_scaling_factor = config.routed_scaling_factor
136
137
138
139
140
141

        self.ep_group = get_ep_group().device_group
        self.ep_rank = self.ep_group.rank()
        self.ep_size = self.ep_group.size()
        self.n_routed_experts: int = config.n_routed_experts
        self.n_shared_experts: int = config.n_shared_experts
142
143
144
145
146

        if config.hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {config.hidden_act}. "
                             "Only silu is supported for now.")

wangding zeng's avatar
wangding zeng committed
147
        self.gate = ReplicatedLinear(config.hidden_size,
148
                                     config.n_routed_experts,
wangding zeng's avatar
wangding zeng committed
149
                                     bias=False,
150
151
                                     quant_config=None,
                                     prefix=f"{prefix}.gate")
152
153
154
155
156
157
        if config.topk_method == "noaux_tc":
            self.gate.e_score_correction_bias = nn.Parameter(
                torch.empty(config.n_routed_experts))
        else:
            self.gate.e_score_correction_bias = None

158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
        # Load balancing settings.
        vllm_config = get_current_vllm_config()
        parallel_config = vllm_config.parallel_config
        self.enable_eplb = enable_eplb

        self.n_redundant_experts = parallel_config.num_redundant_experts
        self.n_logical_experts = self.n_routed_experts
        self.n_physical_experts = (self.n_logical_experts +
                                   self.n_redundant_experts)
        self.n_local_physical_experts = self.n_physical_experts // self.ep_size

        self.physical_expert_start = (self.ep_rank *
                                      self.n_local_physical_experts)
        self.physical_expert_end = (self.physical_expert_start +
                                    self.n_local_physical_experts)

王敏's avatar
王敏 committed
174
        dp_size = get_dp_group().world_size
王敏's avatar
王敏 committed
175
        self.use_mori_ep = parallel_config.enable_expert_parallel and dp_size > 1 and envs.VLLM_ALL2ALL_BACKEND == 'mori'
176
        self.enable_expert_parallel = parallel_config.enable_expert_parallel
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
        self.use_deepep_ll = dp_size > 1 and parallel_config.enable_expert_parallel and \
            envs.VLLM_ALL2ALL_BACKEND == "deepep_low_latency"

        if not self.use_deepep_ll:
            moe_cls = FusedMoE if not self.use_mori_ep else MoriMoE
            self.experts = moe_cls(
                num_experts=config.n_routed_experts,
                top_k=config.num_experts_per_tok,
                hidden_size=config.hidden_size,
                intermediate_size=config.moe_intermediate_size,
                reduce_results=False,
                renormalize=config.norm_topk_prob,
                quant_config=quant_config,
                use_grouped_topk=True,
                num_expert_group=config.n_group,
                topk_group=config.topk_group,
                prefix=f"{prefix}.experts",
                scoring_func=config.scoring_func,
                e_score_correction_bias=self.gate.e_score_correction_bias,
                enable_eplb=self.enable_eplb,
                num_redundant_experts=self.n_redundant_experts,
                routed_scaling_factor=self.routed_scaling_factor)
            
            if config.n_shared_experts is not None:
                intermediate_size = (config.moe_intermediate_size *
                                    config.n_shared_experts)
                shared_expert_cls = DeepseekV2MLP if not self.use_mori_ep else EPSharedExperts
                self.shared_experts = shared_expert_cls(
                    hidden_size=config.hidden_size,
                    intermediate_size=intermediate_size,
                    hidden_act=config.hidden_act,
                    quant_config=quant_config,
                    reduce_results=self.experts.must_reduce_shared_expert_outputs(),
                    prefix=f"{prefix}.shared_experts",
                )
        else:
            if config.n_shared_experts is not None:
                intermediate_size = (config.moe_intermediate_size *
                                    config.n_shared_experts)
216
                self.shared_experts = EPSharedExperts(
217
218
219
220
221
222
223
224
225
226
                    hidden_size=config.hidden_size,
                    intermediate_size=intermediate_size,
                    hidden_act=config.hidden_act,
                    quant_config=quant_config,
                    reduce_results=dp_size != self.ep_size,
                    prefix=f"{prefix}.shared_experts",
                )
            self.experts = SharedFusedMoE(
                num_experts=config.n_routed_experts,
                top_k=config.num_experts_per_tok,
wangding zeng's avatar
wangding zeng committed
227
                hidden_size=config.hidden_size,
228
229
230
                intermediate_size=config.moe_intermediate_size,
                reduce_results=False,
                renormalize=config.norm_topk_prob,
wangding zeng's avatar
wangding zeng committed
231
                quant_config=quant_config,
232
233
234
235
236
237
238
239
240
241
                use_grouped_topk=True,
                num_expert_group=config.n_group,
                topk_group=config.topk_group,
                prefix=f"{prefix}.experts",
                scoring_func=config.scoring_func,
                e_score_correction_bias=self.gate.e_score_correction_bias,
                enable_eplb=self.enable_eplb,
                num_redundant_experts=self.n_redundant_experts,
                routed_scaling_factor=self.routed_scaling_factor,
                shared_experts=self.shared_experts)
242

243
        from vllm.two_batch_overlap.two_batch_overlap import tbo_all_reduce
lizhigong's avatar
lizhigong committed
244
        self.tbo_all_reduce = tbo_all_reduce
wangding zeng's avatar
wangding zeng committed
245

246
247
248
249
    def forward(self, hidden_states: torch.Tensor,
                rms_weight: Optional[torch.Tensor] = None,
                residual: Optional[torch.Tensor] = None
                ) -> torch.Tensor:
wangding zeng's avatar
wangding zeng committed
250
251
        num_tokens, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)
王敏's avatar
王敏 committed
252

253
        if not self.use_mori_ep and not self.use_deepep_ll:
王敏's avatar
王敏 committed
254
            if self.n_shared_experts is not None:
王敏's avatar
王敏 committed
255
256
257
258
                if envs.USE_FUSED_RMS_QUANT:
                    shared_output, new_resi = self.shared_experts(hidden_states, rms_weight, residual, update_hd=True)
                else:
                    shared_output = self.shared_experts(hidden_states)
259

wangding zeng's avatar
wangding zeng committed
260
261
        # router_logits: (num_tokens, n_experts)
        router_logits, _ = self.gate(hidden_states)
262

王敏's avatar
王敏 committed
263
        if not self.enable_expert_parallel:
zhuwenwen's avatar
zhuwenwen committed
264
            if envs.VLLM_USE_LIGHTOP_MOE_SUM_MUL_ADD:
265
                final_hidden_states = self.experts(
王敏's avatar
王敏 committed
266
267
268
                        hidden_states=hidden_states,
                        router_logits=router_logits,
                        shared_output=shared_output)
269
            else:
270
                if hidden_states.dtype != torch.float16:
王敏's avatar
王敏 committed
271
272
273
274
275
276
277
278
                    final_hidden_states = self.experts(
                        hidden_states=hidden_states,
                        router_logits=router_logits) * self.routed_scaling_factor
                else:
                    # Fix FP16 overflow
                    # See DeepseekV2DecoderLayer for more details.
                    final_hidden_states = self.experts(hidden_states=hidden_states,
                                                    router_logits=router_logits)
zhuwenwen's avatar
zhuwenwen committed
279
                if shared_output is not None:
280
                    if hidden_states.dtype != torch.float16:
zhuwenwen's avatar
zhuwenwen committed
281
282
283
284
285
286
                        final_hidden_states = final_hidden_states + shared_output
                    else:
                        # Fix FP16 overflow
                        # See DeepseekV2DecoderLayer for more details.
                        final_hidden_states = final_hidden_states + shared_output \
                            * (1. / self.routed_scaling_factor)
王敏's avatar
王敏 committed
287
        else:
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
            if self.use_deepep_ll:
                shared_output, final_hidden_states = self.experts(hidden_states=hidden_states, 
                                             router_logits=router_logits)

                if shared_output is not None:
                    if hidden_states.dtype != torch.float16:
                        final_hidden_states = final_hidden_states + shared_output
                    else:
                        # Fix FP16 overflow
                        # See DeepseekV2DecoderLayer for more details.
                        final_hidden_states = final_hidden_states + shared_output \
                            * (1. / self.routed_scaling_factor)
            elif self.use_mori_ep:
                final_hidden_states = self.experts(hidden_states=hidden_states,
                                                router_logits=router_logits)
            else:
王敏's avatar
王敏 committed
304
305
306
307
308
309
310
311
312
313
314
315
                final_hidden_states = self.experts(
                        hidden_states=hidden_states,
                        router_logits=router_logits)
                    
                if shared_output is not None:
                    if hidden_states.dtype != torch.float16:
                        final_hidden_states = final_hidden_states + shared_output
                    else:
                        # Fix FP16 overflow
                        # See DeepseekV2DecoderLayer for more details.
                        final_hidden_states = final_hidden_states + shared_output \
                            * (1. / self.routed_scaling_factor)
316

317
        if not self.use_mori_ep:
王敏's avatar
王敏 committed
318
319
320
321
322
323
324
            if self.tp_size > 1:
                if envs.VLLM_ENABLE_TBO:
                    final_hidden_states = self.tbo_all_reduce(final_hidden_states)
                else:
                    final_hidden_states = (
                        self.experts.maybe_all_reduce_tensor_model_parallel(
                            final_hidden_states))
325
326
327
328
        if envs.USE_FUSED_RMS_QUANT:
            return final_hidden_states.view(num_tokens, hidden_dim), new_resi
        else:
            return final_hidden_states.view(num_tokens, hidden_dim)
wangding zeng's avatar
wangding zeng committed
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350


def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
    import math
    if scale <= 1:
        return 1.0
    return 0.1 * mscale * math.log(scale) + 1.0


class DeepseekV2Attention(nn.Module):

    def __init__(
        self,
        config: PretrainedConfig,
        hidden_size: int,
        num_heads: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        v_head_dim: int,
        q_lora_rank: int,
        kv_lora_rank: int,
        rope_theta: float = 10000,
351
        rope_scaling: Optional[dict[str, Any]] = None,
wangding zeng's avatar
wangding zeng committed
352
353
354
        max_position_embeddings: int = 8192,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
355
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
        self.v_head_dim = v_head_dim
        self.q_lora_rank = q_lora_rank
        self.kv_lora_rank = kv_lora_rank
        self.num_heads = num_heads
        tp_size = get_tensor_model_parallel_world_size()
        assert num_heads % tp_size == 0
        self.num_local_heads = num_heads // tp_size
        self.scaling = self.qk_head_dim**-0.5
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings

        if self.q_lora_rank is not None:
            self.q_a_proj = ReplicatedLinear(self.hidden_size,
                                             self.q_lora_rank,
                                             bias=False,
377
378
                                             quant_config=quant_config,
                                             prefix=f"{prefix}.q_a_proj")
wangding zeng's avatar
wangding zeng committed
379
380
381
382
383
384
            self.q_a_layernorm = RMSNorm(self.q_lora_rank,
                                         eps=config.rms_norm_eps)
            self.q_b_proj = ColumnParallelLinear(q_lora_rank,
                                                 self.num_heads *
                                                 self.qk_head_dim,
                                                 bias=False,
385
386
                                                 quant_config=quant_config,
                                                 prefix=f"{prefix}.q_b_proj")
wangding zeng's avatar
wangding zeng committed
387
388
389
390
391
        else:
            self.q_proj = ColumnParallelLinear(self.hidden_size,
                                               self.num_heads *
                                               self.qk_head_dim,
                                               bias=False,
392
393
                                               quant_config=quant_config,
                                               prefix=f"{prefix}.q_proj")
wangding zeng's avatar
wangding zeng committed
394

395
396
397
398
399
400
        self.kv_a_proj_with_mqa = ReplicatedLinear(
            self.hidden_size,
            self.kv_lora_rank + self.qk_rope_head_dim,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.kv_a_proj_with_mqa")
wangding zeng's avatar
wangding zeng committed
401
402
403
404
405
406
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank,
                                      eps=config.rms_norm_eps)
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
407
408
            quant_config=quant_config,
            prefix=f"{prefix}.kv_b_proj")
wangding zeng's avatar
wangding zeng committed
409
410
411
412
        # O projection.
        self.o_proj = RowParallelLinear(self.num_heads * self.v_head_dim,
                                        self.hidden_size,
                                        bias=False,
413
414
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.o_proj")
415
416
        if rope_scaling:
            rope_scaling["rope_type"] = 'deepseek_yarn'
417

wangding zeng's avatar
wangding zeng committed
418
419
420
421
422
423
424
425
426
427
428
429
430
431
        self.rotary_emb = get_rope(qk_rope_head_dim,
                                   rotary_dim=qk_rope_head_dim,
                                   max_position=max_position_embeddings,
                                   base=rope_theta,
                                   rope_scaling=rope_scaling,
                                   is_neox_style=False)

        if rope_scaling:
            mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
            scaling_factor = rope_scaling["factor"]
            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
            self.scaling = self.scaling * mscale * mscale

        self.attn = Attention(self.num_local_heads,
432
                              self.qk_head_dim,
wangding zeng's avatar
wangding zeng committed
433
434
435
                              self.scaling,
                              num_kv_heads=self.num_local_heads,
                              cache_config=cache_config,
436
437
                              quant_config=quant_config,
                              prefix=f"{prefix}.attn")
wangding zeng's avatar
wangding zeng committed
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        if self.q_lora_rank is not None:
            q = self.q_a_proj(hidden_states)[0]
            q = self.q_a_layernorm(q)
            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads,
                                         self.qk_head_dim)
        else:
            q = self.q_proj(hidden_states)[0].view(-1, self.num_local_heads,
                                                   self.qk_head_dim)
        q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim],
                               dim=-1)
        latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
        kv_a, _ = latent_cache.split(
            [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
        latent_cache = latent_cache.unsqueeze(1)
        kv_a = self.kv_a_layernorm(kv_a.contiguous())
        kv = self.kv_b_proj(kv_a)[0]
        kv = kv.view(-1, self.num_local_heads,
                     self.qk_nope_head_dim + self.v_head_dim)
        k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
        k_pe = latent_cache[:, :, self.kv_lora_rank:]
464

wangding zeng's avatar
wangding zeng committed
465
        q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
466

wangding zeng's avatar
wangding zeng committed
467
468
469
470
        q[..., self.qk_nope_head_dim:] = q_pe
        k = torch.empty_like(q)
        k[..., :self.qk_nope_head_dim] = k_nope
        k[..., self.qk_nope_head_dim:] = k_pe
471
472
473
474
        # padding value to qk_head_dim for alignment
        v = torch.nn.functional.pad(
            v, [0, self.qk_head_dim - self.v_head_dim],
            value=0).view(-1, self.num_local_heads * self.qk_head_dim)
475
        attn_output = self.attn(q, k, v)
wangding zeng's avatar
wangding zeng committed
476
        attn_output = attn_output.view(
477
478
            -1, self.num_local_heads,
            self.qk_head_dim)[..., :self.v_head_dim].reshape(
wangding zeng's avatar
wangding zeng committed
479
480
481
482
483
                -1, self.num_local_heads * self.v_head_dim)
        output, _ = self.o_proj(attn_output)
        return output


484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
class DeepseekV2MLAAttention(nn.Module):
    """
    Main reference: DeepseekV2 paper, and FlashInfer Implementation
    (https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).
    
    For more info see MLACommonImpl in: vllm/attention/backends/mla/utils.py
    """

    def __init__(
        self,
        config: PretrainedConfig,
        hidden_size: int,
        num_heads: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        v_head_dim: int,
        q_lora_rank: Optional[int],
        kv_lora_rank: int,
        rope_theta: float = 10000,
503
        rope_scaling: Optional[dict[str, Any]] = None,
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
        max_position_embeddings: int = 8192,
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
        self.v_head_dim = v_head_dim

        self.q_lora_rank = q_lora_rank
        self.kv_lora_rank = kv_lora_rank

        self.num_heads = num_heads
        tp_size = get_tensor_model_parallel_world_size()
        assert num_heads % tp_size == 0
        self.num_local_heads = num_heads // tp_size

        self.scaling = self.qk_head_dim**-0.5
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings

        if self.q_lora_rank is not None:
529
530
            if envs.USE_FUSED_RMS_QUANT:
                self.q_a_proj = ReplicatedLinear(self.hidden_size,
531
532
533
                                             self.q_lora_rank,
                                             bias=False,
                                             quant_config=quant_config,
534
                                             eps=config.rms_norm_eps,
535
                                             prefix=f"{prefix}.q_a_proj")
536
                self.q_b_proj = ColumnParallelLinear(q_lora_rank,
537
538
539
540
                                                 self.num_heads *
                                                 self.qk_head_dim,
                                                 bias=False,
                                                 quant_config=quant_config,
541
                                                 eps=config.rms_norm_eps,
542
                                                 prefix=f"{prefix}.q_b_proj")
543
544
545
546
547
548
549
            else:
                self.q_a_proj = ReplicatedLinear(self.hidden_size,
                                             self.q_lora_rank,
                                             bias=False,
                                             quant_config=quant_config,
                                             prefix=f"{prefix}.q_a_proj")
                self.q_b_proj = ColumnParallelLinear(q_lora_rank,
550
551
552
553
554
                                                 self.num_heads *
                                                 self.qk_head_dim,
                                                 bias=False,
                                                 quant_config=quant_config,
                                                 prefix=f"{prefix}.q_b_proj")
王敏's avatar
王敏 committed
555

556
557
558
            self.q_a_layernorm = RMSNorm(self.q_lora_rank,
                                         eps=config.rms_norm_eps)

559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
        else:
            self.q_proj = ColumnParallelLinear(self.hidden_size,
                                               self.num_heads *
                                               self.qk_head_dim,
                                               bias=False,
                                               quant_config=quant_config,
                                               prefix=f"{prefix}.q_proj")

        self.kv_a_proj_with_mqa = ReplicatedLinear(
            self.hidden_size,
            self.kv_lora_rank + self.qk_rope_head_dim,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.kv_a_proj_with_mqa")
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank,
                                      eps=config.rms_norm_eps)
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.kv_b_proj")
        self.o_proj = RowParallelLinear(self.num_heads * self.v_head_dim,
                                        self.hidden_size,
                                        bias=False,
                                        quant_config=quant_config,
                                        prefix=f"{prefix}.o_proj")

587
588
        if rope_scaling:
            rope_scaling["rope_type"] = 'deepseek_yarn'
589
590
591
592
593
594
595
596
597
598
599
600
        self.rotary_emb = get_rope(qk_rope_head_dim,
                                   rotary_dim=qk_rope_head_dim,
                                   max_position=max_position_embeddings,
                                   base=rope_theta,
                                   rope_scaling=rope_scaling,
                                   is_neox_style=False)
        if rope_scaling:
            mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
            scaling_factor = rope_scaling["factor"]
            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
            self.scaling = self.scaling * mscale * mscale

601
602
603
604
605
606
        # In the MLA backend, kv_cache includes both k_c and
        # pe (i.e. decoupled position embeddings). In particular,
        # the concat_and_cache_mla op requires
        #     k_c.size(1) + k_pe.size(1) == kv_cache.size(2)
        # i.e.
        #     kv_lora_rank + qk_rope_head_dim == head_size
607
608
        self.mla_attn = Attention(
            num_heads=self.num_local_heads,
609
            head_size=self.kv_lora_rank + self.qk_rope_head_dim,
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
            scale=self.scaling,
            num_kv_heads=1,
            cache_config=cache_config,
            quant_config=quant_config,
            prefix=f"{prefix}.attn",
            use_mla=True,
            # MLA Args
            q_lora_rank=self.q_lora_rank,
            kv_lora_rank=self.kv_lora_rank,
            qk_nope_head_dim=self.qk_nope_head_dim,
            qk_rope_head_dim=self.qk_rope_head_dim,
            qk_head_dim=self.qk_head_dim,
            v_head_dim=self.v_head_dim,
            kv_b_proj=self.kv_b_proj,
        )

        self.prefix = prefix
        self.debug_layer_idx = int(self.prefix.split(".")[-2])

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
633
634
        rms_weight: Optional[torch.Tensor] = None,
        residual: Optional[torch.Tensor] = None
635
    ) -> torch.Tensor:
636
637
638
639
        if envs.USE_FUSED_RMS_QUANT and rms_weight is not None:
            if self.q_lora_rank is not None:
                q_c, new_residual, _, input_quant_args = self.q_a_proj(hidden_states, rms_weight=rms_weight, residual=residual, update_hd=False)
                q, _, _ = self.q_b_proj(q_c, rms_weight=self.q_a_layernorm.weight.data, residual=None, update_hd=False)
王敏's avatar
王敏 committed
640

641
642
643
644
            else:
                q = self.q_proj(hidden_states)[0]
            kv_c, k_pe = self.kv_a_proj_with_mqa(hidden_states, quant_args=input_quant_args, update_hd=False)[0].split(
                                [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
王敏's avatar
王敏 committed
645

王敏's avatar
王敏 committed
646
647
648
649
            if envs.VLLM_USE_LIGHTOP:
                kv_c_normed = self.kv_a_layernorm.forward_cuda_opt(kv_c)
            else:
                kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664

            q = q.view(-1, self.num_local_heads, self.qk_head_dim)
            # Add head dim of 1 to k_pe
            k_pe = k_pe.unsqueeze(1)

            q[..., self.qk_nope_head_dim:], k_pe = self.rotary_emb(
                positions, q[..., self.qk_nope_head_dim:], k_pe)

            attn_out = self.mla_attn(
                q,
                kv_c_normed,
                k_pe,
                output_shape=(hidden_states.shape[0],
                            self.num_local_heads * self.v_head_dim))
            return self.o_proj(attn_out)[0], new_residual
665
        else:
666
667
668
669
670
671
672
673
            if self.q_lora_rank is not None:
                q_c = self.q_a_proj(hidden_states)[0]
                q_c = self.q_a_layernorm(q_c)
                q = self.q_b_proj(q_c)[0]
            else:
                q = self.q_proj(hidden_states)[0]
            kv_c, k_pe = self.kv_a_proj_with_mqa(hidden_states)[0].split(
                [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
674
675
676
677
            if envs.VLLM_USE_LIGHTOP:
                kv_c_normed = self.kv_a_layernorm.forward_cuda_opt(kv_c)
            else:
                kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
678

679
680
681
            q = q.view(-1, self.num_local_heads, self.qk_head_dim)
            # Add head dim of 1 to k_pe
            k_pe = k_pe.unsqueeze(1)
682

683
684
            q[..., self.qk_nope_head_dim:], k_pe = self.rotary_emb(
                positions, q[..., self.qk_nope_head_dim:], k_pe)
685

686
687
688
689
690
691
692
            attn_out = self.mla_attn(
                q,
                kv_c_normed,
                k_pe,
                output_shape=(hidden_states.shape[0],
                            self.num_local_heads * self.v_head_dim))
            return self.o_proj(attn_out)[0]
693
694


wangding zeng's avatar
wangding zeng committed
695
696
697
698
699
class DeepseekV2DecoderLayer(nn.Module):

    def __init__(
        self,
        config: PretrainedConfig,
700
        prefix: str,
701
        model_config: ModelConfig,
wangding zeng's avatar
wangding zeng committed
702
703
        cache_config: Optional[CacheConfig] = None,
        quant_config: Optional[QuantizationConfig] = None,
704
        enable_eplb: bool = False,
wangding zeng's avatar
wangding zeng committed
705
706
707
708
709
710
711
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        rope_theta = getattr(config, "rope_theta", 10000)
        rope_scaling = getattr(config, "rope_scaling", None)
        max_position_embeddings = getattr(config, "max_position_embeddings",
                                          8192)
712
713
714
        # DecoderLayers are created with `make_layers` which passes the prefix
        # with the layer's index.
        layer_idx = int(prefix.split(sep='.')[-1])
715
        self.layer_idx = layer_idx
716
717
718
719

        self.dp_size = get_dp_group().world_size
        vllm_config = get_current_vllm_config()
        parallel_config = vllm_config.parallel_config
720
721
        self.use_deepep_ll = self.dp_size > 1 and parallel_config.enable_expert_parallel and \
            envs.VLLM_ALL2ALL_BACKEND == "deepep_low_latency"
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
        self.tp_size = get_tensor_model_parallel_world_size()

        if (config.n_routed_experts is not None
                and layer_idx >= config.first_k_dense_replace
                and layer_idx % config.moe_layer_freq == 0):
            self.mlp = DeepseekV2MoE(
                config=config,
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
                enable_eplb=enable_eplb,
            )
        else:
            self.mlp = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=config.intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
            )

742
743
744
745
746
        if model_config.use_mla:
            attn_cls = DeepseekV2MLAAttention
        else:
            attn_cls = DeepseekV2Attention
        self.self_attn = attn_cls(
wangding zeng's avatar
wangding zeng committed
747
748
749
750
751
752
753
754
755
756
757
758
759
760
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            qk_nope_head_dim=config.qk_nope_head_dim,
            qk_rope_head_dim=config.qk_rope_head_dim,
            v_head_dim=config.v_head_dim,
            q_lora_rank=config.q_lora_rank
            if hasattr(config, "q_lora_rank") else None,
            kv_lora_rank=config.kv_lora_rank,
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
            cache_config=cache_config,
            quant_config=quant_config,
761
            prefix=f"{prefix}.self_attn",
wangding zeng's avatar
wangding zeng committed
762
        )
763

wangding zeng's avatar
wangding zeng committed
764
765
766
767
        self.input_layernorm = RMSNorm(config.hidden_size,
                                       eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size,
                                                eps=config.rms_norm_eps)
768
        self.routed_scaling_factor = config.routed_scaling_factor
wangding zeng's avatar
wangding zeng committed
769
770
771
772
773
774
775

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: Optional[torch.Tensor],
    ) -> torch.Tensor:
776
777
778
        if envs.USE_FUSED_RMS_QUANT:
            # Fix residual FP16 overflow
            residual_fix_overflow = False
王敏's avatar
王敏 committed
779

780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
            assert self.input_layernorm.has_weight is True
            if residual is None:
                residual = hidden_states
                hidden_states, _ = self.self_attn(
                    positions = positions,
                    hidden_states = hidden_states,
                    rms_weight = self.input_layernorm.weight.data,
                    residual = None
                )
                residual_fix_overflow = True
            else:
                hidden_states, new_residual = self.self_attn(
                    positions = positions,
                    hidden_states = hidden_states,
                    rms_weight = self.input_layernorm.weight.data,
                    residual = residual
                )
                residual = new_residual
798
               
799
            if hidden_states.dtype == torch.float16:
800
801
802
803
804
805
806
807
808
809
                # rmsnorm, and rmsnorm result would not affect by scale.
                hidden_states *= 1. / self.routed_scaling_factor
                if self.layer_idx == 0 or residual_fix_overflow:
                    # The residual is shared by all layers, we only scale it on
                    # first layer.
                    residual *= 1. / self.routed_scaling_factor

            hidden_states, new_resi = self.mlp(hidden_states, self.post_attention_layernorm.weight.data, residual)

            if isinstance(self.mlp,
810
                        DeepseekV2MLP) and hidden_states.dtype == torch.float16:
811
812
813
814
815
816
817
818
                # Fix FP16 overflow
                # Scaling the DeepseekV2MLP output, it is the input of
                # input_layernorm of next decoder layer.
                # The scaling of DeepseekV2MOE output would be done in the forward
                # of DeepseekV2MOE
                hidden_states *= 1. / self.routed_scaling_factor
            return hidden_states, new_resi

wangding zeng's avatar
wangding zeng committed
819
        else:
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
            # Self Attention
            # Fix residual FP16 overflow
            residual_fix_overflow = False
            if residual is None:
                residual = hidden_states
                hidden_states = self.input_layernorm(hidden_states)
                residual_fix_overflow = True
            else:
                hidden_states, residual = self.input_layernorm(
                    hidden_states, residual)
            hidden_states = self.self_attn(
                positions=positions,
                hidden_states=hidden_states,
            )

835
            if hidden_states.dtype == torch.float16:
836
837
838
839
840
841
842
843
844
845
846
                # Fix FP16 overflow
                # We scale both hidden_states and residual before
                # rmsnorm, and rmsnorm result would not affect by scale.
                hidden_states *= 1. / self.routed_scaling_factor
                if self.layer_idx == 0 or residual_fix_overflow:
                    # The residual is shared by all layers, we only scale it on
                    # first layer.
                    residual *= 1. / self.routed_scaling_factor

            # Fully Connected
            hidden_states, residual = self.post_attention_layernorm(
wangding zeng's avatar
wangding zeng committed
847
                hidden_states, residual)
848
849
            
            if isinstance(self.mlp,
850
                        DeepseekV2MoE) and self.use_deepep_ll and self.tp_size > 1:
851
852
853
854
855
856
857
858
859

                self.tp_rank = get_tensor_model_parallel_rank()
                ori_bs = hidden_states.shape[0]
                pad_size = (ori_bs + self.tp_size - 1) // self.tp_size * self.tp_size - ori_bs
                if pad_size > 0:
                    hidden_states = torch.nn.functional.pad(hidden_states.contiguous(), [0, 0, 0, pad_size], value=0).contiguous()
                new_bs = (ori_bs+pad_size) // self.tp_size
                hidden_states = hidden_states[self.tp_rank*new_bs: (self.tp_rank+1)*new_bs, :]

860
            hidden_states = self.mlp(hidden_states)
wangding zeng's avatar
wangding zeng committed
861

862
            if isinstance(self.mlp,
863
                        DeepseekV2MoE) and self.use_deepep_ll and self.tp_size > 1:
864
865
866
                hidden_states = tensor_model_parallel_all_gather(hidden_states, dim=0).contiguous()
                hidden_states = hidden_states[:ori_bs, :].contiguous()

867
            if isinstance(self.mlp,
868
                        DeepseekV2MLP) and hidden_states.dtype == torch.float16:
869
870
871
872
873
874
                # Fix FP16 overflow
                # Scaling the DeepseekV2MLP output, it is the input of
                # input_layernorm of next decoder layer.
                # The scaling of DeepseekV2MOE output would be done in the forward
                # of DeepseekV2MOE
                hidden_states *= 1. / self.routed_scaling_factor
875

876
            return hidden_states, residual
wangding zeng's avatar
wangding zeng committed
877
878


879
@support_torch_compile
wangding zeng's avatar
wangding zeng committed
880
881
882
883
class DeepseekV2Model(nn.Module):

    fall_back_to_pt_during_load = False

884
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
885
        super().__init__()
886
887

        config = vllm_config.model_config.hf_config
888
        model_config = vllm_config.model_config
889
890
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
891
        enable_eplb = vllm_config.parallel_config.enable_eplb
892
        self.config = config
893

wangding zeng's avatar
wangding zeng committed
894
895
        self.vocab_size = config.vocab_size

896
897
898
899
        if get_pp_group().is_first_rank:
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
900
901
                quant_config=quant_config,
                prefix=f"{prefix}.embed_tokens")
902
903
904
905
906
907
908
909
        else:
            self.embed_tokens = PPMissingLayer()

        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
            lambda prefix: DeepseekV2DecoderLayer(
                config,
                prefix,
910
                model_config=model_config,
911
912
                cache_config=cache_config,
                quant_config=quant_config,
913
                enable_eplb=enable_eplb,
914
915
916
917
918
919
920
            ),
            prefix=f"{prefix}.layers")

        if get_pp_group().is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer()
921
922
923
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(
                ["hidden_states", "residual"], config.hidden_size))
wangding zeng's avatar
wangding zeng committed
924

925
926
927
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)

wangding zeng's avatar
wangding zeng committed
928
929
930
931
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
932
        intermediate_tensors: Optional[IntermediateTensors],
933
        inputs_embeds: Optional[torch.Tensor] = None,
934
    ) -> Union[torch.Tensor, IntermediateTensors]:
935
        if get_pp_group().is_first_rank:
936
937
938
939
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
940
941
942
943
944
945
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

946
        for layer in self.layers[self.start_layer:self.end_layer]:
947
            hidden_states, residual = layer(positions, hidden_states, residual)
948
949
950
951
952
953
954

        if not get_pp_group().is_last_rank:
            return IntermediateTensors({
                "hidden_states": hidden_states,
                "residual": residual
            })

wangding zeng's avatar
wangding zeng committed
955
956
957
958
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


959
class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
wangding zeng's avatar
wangding zeng committed
960

961
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
962
        super().__init__()
963
964
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
965
966
967
968
969
970
971

        self.quant_method = None
        if quant_config is not None:
            self.quant_method = quant_config.get_name()
            os.environ['LLAMA_NN'] = '0'
            os.environ['LM_NN'] = '0'

972
        self.use_w4a16_moe_sz = os.environ.get('AWQ_MOE_SZ') == '1'
wangding zeng's avatar
wangding zeng committed
973
974
        self.config = config
        self.quant_config = quant_config
王敏's avatar
王敏 committed
975

976
        self.model = DeepseekV2Model(vllm_config=vllm_config,
977
                                     prefix=maybe_prefix(prefix, "model"))
978
979
980
981
982
983
        if get_pp_group().is_last_rank:
            self.lm_head = ParallelLMHead(config.vocab_size,
                                          config.hidden_size,
                                          quant_config=quant_config)
        else:
            self.lm_head = PPMissingLayer()
wangding zeng's avatar
wangding zeng committed
984
        self.logits_processor = LogitsProcessor(config.vocab_size)
985
986
        self.make_empty_intermediate_tensors = (
            self.model.make_empty_intermediate_tensors)
987
988
989
990
991
992
993
994
        self.expert_weights = []

        # Set MoE hyperparameters
        self.num_moe_layers = (config.num_hidden_layers -
                               config.first_k_dense_replace)
        self.num_expert_groups = config.n_group

        self.moe_layers: list[FusedMoE] = []
995
        example_moe = None
996
        for layer in self.model.layers:
997
998
999
            if isinstance(layer, PPMissingLayer):
                continue

1000
1001
            assert isinstance(layer, DeepseekV2DecoderLayer)
            if isinstance(layer.mlp, DeepseekV2MoE):
1002
                example_moe = layer.mlp
1003
1004
1005
1006
1007
1008
1009
1010
1011
                self.moe_layers.append(layer.mlp.experts)

        # Pick last one layer since the first ones may be dense layers.
        self.num_logical_experts = example_moe.n_logical_experts
        self.num_physical_experts = example_moe.n_physical_experts
        self.num_local_physical_experts = example_moe.n_local_physical_experts
        self.num_routed_experts = example_moe.n_routed_experts
        self.num_shared_experts = example_moe.n_shared_experts
        self.num_redundant_experts = example_moe.n_redundant_experts
zhuwenwen's avatar
zhuwenwen committed
1012
        
王敏's avatar
王敏 committed
1013
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
1014
        self.use_awq_pad = os.environ.get('AWQ_PAD') == '1'
1015
1016
        self.tritonsingleton= W8a8GetCacheJSON() 
        self.tritonsingleton.topk = config.num_experts_per_tok
王敏's avatar
王敏 committed
1017
        self.tritonsingleton.quant_method=self.quant_method
1018

1019
1020
        parallel_config = vllm_config.parallel_config
        dp_size = get_dp_group().world_size
1021
        self.use_mori_ep = envs.VLLM_ALL2ALL_BACKEND == 'mori' and dp_size > 1 and parallel_config.enable_expert_parallel
1022

1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
    def set_eplb_state(
        self,
        expert_load_view: torch.Tensor,
        logical_to_physical_map: torch.Tensor,
        logical_replica_count: torch.Tensor,
    ) -> None:
        for layer_idx, layer in enumerate(self.moe_layers):
            # Register the expert weights.
            self.expert_weights.append(layer.get_expert_weights())
            layer.set_eplb_state(
                moe_layer_idx=layer_idx,
                expert_load_view=expert_load_view,
                logical_to_physical_map=logical_to_physical_map,
                logical_replica_count=logical_replica_count,
            )
1038

1039
1040
1041
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.get_input_embeddings(input_ids)

wangding zeng's avatar
wangding zeng committed
1042
1043
1044
1045
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
1046
        intermediate_tensors: Optional[IntermediateTensors] = None,
1047
        inputs_embeds: Optional[torch.Tensor] = None,
1048
    ) -> Union[torch.Tensor, IntermediateTensors]:
1049
        hidden_states = self.model(input_ids, positions, intermediate_tensors,
1050
                                   inputs_embeds)
wangding zeng's avatar
wangding zeng committed
1051
1052
        return hidden_states

1053
1054
1055
1056
1057
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
1058
        logits = self.logits_processor(self.lm_head, hidden_states,
wangding zeng's avatar
wangding zeng committed
1059
1060
1061
                                       sampling_metadata)
        return logits

1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
    def make_empty_intermediate_tensors(
            self, batch_size: int, dtype: torch.dtype,
            device: torch.device) -> IntermediateTensors:
        return IntermediateTensors({
            "hidden_states":
            torch.zeros((batch_size, self.config.hidden_size),
                        dtype=dtype,
                        device=device),
            "residual":
            torch.zeros((batch_size, self.config.hidden_size),
                        dtype=dtype,
                        device=device),
        })
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
        
    def restore_qzeros_tensor(self, qzeros, qscales):

        low_bits = qzeros & 0x0F
        high_bits = qzeros >> 4
        
        zeors_tensor = torch.stack([low_bits, high_bits], dim=2).view(qzeros.shape[0], -1 , qzeros.shape[-1])
        zeors_int16 = zeors_tensor.to(torch.int16)
        assert zeors_int16.shape == qscales.shape

        uint16_tensor1 = zeors_int16.view(torch.uint16)
        uint16_tensor2 = qscales.view(torch.uint16)
        
        uint32_tensor1 = uint16_tensor1.to(torch.int32) << 16
        uint32_tensor2 = uint16_tensor2.to(torch.int32)
        
        result_tensor = uint32_tensor1 + uint32_tensor2
        result_tensor =result_tensor.view(torch.uint32)
        result_tensor = result_tensor.transpose(1, 2).contiguous()
        return result_tensor
1095

1096
1097
    def load_weights(self, weights: Iterable[tuple[str,
                                                   torch.Tensor]]) -> set[str]:
wangding zeng's avatar
wangding zeng committed
1098
1099
1100
1101
1102
1103
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]

1104
        if self.use_mori_ep:
1105
1106
1107
            ep_moe_shared_experts_keys = "mlp.shared_experts"
            ep_moe_shared_experts_mapping = {ep_moe_shared_experts_keys:"mlp.experts.shared_experts"}

1108
1109
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
王敏's avatar
王敏 committed
1110
1111
1112
1113
        expert_params_mapping = FusedMoE.make_expert_params_mapping(
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
1114
1115
            num_experts=self.config.n_routed_experts,
            num_redundant_experts=self.num_redundant_experts)
1116

wangding zeng's avatar
wangding zeng committed
1117
        params_dict = dict(self.named_parameters())
1118
        loaded_params: set[str] = set()
wangding zeng's avatar
wangding zeng committed
1119
1120
1121
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
1122

1123
1124
1125
            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
            if spec_layer is not None:
                continue  # skip spec decode layers for main model
1126

wangding zeng's avatar
wangding zeng committed
1127
            for (param_name, weight_name, shard_id) in stacked_params_mapping:
1128
                # Skip non-stacked layers and experts (experts handled below).
wangding zeng's avatar
wangding zeng committed
1129
1130
                if weight_name not in name:
                    continue
1131
1132
1133
1134
1135
1136
1137
1138
                # We have mlp.experts[0].gate_proj in the checkpoint.
                # Since we handle the experts below in expert_params_mapping,
                # we need to skip here BEFORE we update the name, otherwise
                # name will be updated to mlp.experts[0].gate_up_proj, which
                # will then be updated below in expert_params_mapping
                # for mlp.experts[0].gate_gate_up_proj, which breaks load.
                if (("mlp.experts." in name) and name not in params_dict):
                    continue
wangding zeng's avatar
wangding zeng committed
1139
                name = name.replace(weight_name, param_name)
1140

1141
                if self.use_mori_ep:
1142
1143
                    name = name.replace(ep_moe_shared_experts_keys, ep_moe_shared_experts_mapping[ep_moe_shared_experts_keys])

wangding zeng's avatar
wangding zeng committed
1144
1145
1146
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
1147
1148
1149
1150

                if is_pp_missing_parameter(name, self):
                    continue

wangding zeng's avatar
wangding zeng committed
1151
1152
1153
1154
1155
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
1156
                is_expert_weight = False
1157
1158
1159
1160
                for mapping in expert_params_mapping:
                    param_name, weight_name, expert_id, shard_id = mapping
                    if weight_name not in name:
                        continue
1161

1162
1163
1164
1165
1166
1167
1168
1169
                    # Anyway, this is an expert weight and should not be
                    # attempted to load as other weights later
                    is_expert_weight = True

                    # Do not modify `name` since the loop may continue here
                    # Instead, create a new variable
                    name_mapped = name.replace(weight_name, param_name)

1170
                    if self.use_mori_ep:
1171
1172
                        name_mapped = name_mapped.replace(ep_moe_shared_experts_keys, ep_moe_shared_experts_mapping[ep_moe_shared_experts_keys])

1173
                    if is_pp_missing_parameter(name_mapped, self):
1174
1175
                        continue

1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
                    param = params_dict[name_mapped]
                    # We should ask the weight loader to return success or not
                    # here since otherwise we may skip experts with other
                    # available replicas.
                    weight_loader = typing.cast(Callable[..., bool],
                                                param.weight_loader)
                    success = weight_loader(param,
                                            loaded_weight,
                                            name_mapped,
                                            shard_id=shard_id,
                                            expert_id=expert_id,
                                            return_success=True)
                    if success:
1189
                        name = name_mapped
1190
                        break
1191
                else:
1192
1193
1194
1195
1196
1197
                    if is_expert_weight:
                        # We've checked that this is an expert weight
                        # However it's not mapped locally to this rank
                        # So we simply skip it
                        continue

1198
                    if self.use_mori_ep:
1199
                        name = name.replace(ep_moe_shared_experts_keys, ep_moe_shared_experts_mapping[ep_moe_shared_experts_keys])
1200
1201
1202
1203
                    # Skip loading extra bias for GPTQ models.
                    if name.endswith(".bias") and name not in params_dict:
                        continue

1204
1205
1206
1207
1208
                    # Remapping the name of FP8 kv-scale.
                    name = maybe_remap_kv_scale_name(name, params_dict)
                    if name is None:
                        continue

1209
1210
1211
                    if is_pp_missing_parameter(name, self):
                        continue

zhuwenwen's avatar
zhuwenwen committed
1212
1213
1214
1215
                    try:
                        param = params_dict[name]
                    except Exception as e:
                        continue
1216
1217
1218
                    weight_loader = getattr(param, "weight_loader",
                                            default_weight_loader)
                    weight_loader(param, loaded_weight)
1219
            loaded_params.add(name)
王敏's avatar
王敏 committed
1220
            
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
        if self.use_llama_nn and self.quant_method is None:
            lay_key_words = [
                "self_attn.q_proj.weight",
                "self_attn.q_a_proj.weight",
                "self_attn.q_b_proj.weight",
                "self_attn.kv_a_proj_with_mqa.weight",
                "self_attn.kv_b_proj.weight",
                "self_attn.o_proj.weight",
                "mlp.gate_up_proj.weight",
                "mlp.down_proj.weight",
                "mlp.gate.weight",
                "shared_experts.gate_up_proj.weight",
                "shared_experts.down_proj.weight",
                "lm_head.weight",
            ]

            combined_words = "|".join(lay_key_words)
            
            for layername in loaded_params:
                weight = params_dict[layername]
                matches = re.findall(combined_words, layername)
                if matches:
                    _weight = torch.zeros_like(weight.data)
                    ori_shape =_weight.shape
                    
                    ops.trans_w16_gemm(_weight, weight.data, _weight.shape[0], _weight.shape[1])
                    weight.data.copy_(_weight)
                    
                    weight.data=weight.data.reshape(ori_shape[1],-1)
zhuwenwen's avatar
zhuwenwen committed
1250
            
1251
        return loaded_params
1252
1253
1254
1255


class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
    pass
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266


def get_spec_layer_idx_from_weight_name(config: PretrainedConfig,
                                        weight_name: str) -> Optional[int]:
    if hasattr(config,
               "num_nextn_predict_layers") and (config.num_nextn_predict_layers
                                                > 0):
        layer_idx = config.num_hidden_layers
        for i in range(config.num_nextn_predict_layers):
            if weight_name.startswith(f"model.layers.{layer_idx+i}."):
                return layer_idx + i
王敏's avatar
王敏 committed
1267
    return None