deepseek_v2.py 68.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
from vllm import envs
zhuwenwen's avatar
zhuwenwen committed
29

30
31
import typing
from collections.abc import Callable, Iterable
32
from itertools import islice
wangding zeng's avatar
wangding zeng committed
33
34
35

import torch
from torch import nn
36
from transformers import DeepseekV2Config, DeepseekV3Config
wangding zeng's avatar
wangding zeng committed
37

38
from vllm._aiter_ops import rocm_aiter_ops
39
from vllm.attention.backends.abstract import AttentionBackend
40
from vllm.attention.layer import Attention
41
from vllm.attention.ops.common import pack_seq_triton, unpack_seq_triton
42
from vllm.compilation.decorators import support_torch_compile
43
44
45
46
47
48
49
50
from vllm.config import CacheConfig, ParallelConfig, VllmConfig, get_current_vllm_config
from vllm.distributed import (
    get_ep_group,
    get_pp_group,
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
    tensor_model_parallel_all_gather,
)
51
52
from vllm.forward_context import get_forward_context
from vllm.logger import init_logger
wangding zeng's avatar
wangding zeng committed
53
from vllm.model_executor.layers.activation import SiluAndMul
54
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
55
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
56
from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
57
58
59
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
60
    QKVParallelLinear,
61
62
63
    ReplicatedLinear,
    RowParallelLinear,
)
wangding zeng's avatar
wangding zeng committed
64
from vllm.model_executor.layers.logits_processor import LogitsProcessor
65
from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
66
from vllm.model_executor.layers.quantization import QuantizationConfig
67
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
68
69
    per_token_group_quant_fp8,
)
wangding zeng's avatar
wangding zeng committed
70
71
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.vocab_parallel_embedding import (
72
73
74
    ParallelLMHead,
    VocabParallelEmbedding,
)
75
from vllm.model_executor.model_loader.weight_utils import (
76
77
78
    default_weight_loader,
    maybe_remap_kv_scale_name,
)
79
from vllm.model_executor.models.utils import sequence_parallel_chunk
80
from vllm.platforms import current_platform
81
from vllm.sequence import IntermediateTensors
82
from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits
83
from vllm.utils.torch_utils import direct_register_custom_op
84
85
86
87
from vllm.v1.attention.backends.mla.indexer import (
    DeepseekV32IndexerBackend,
    DeepseekV32IndexerMetadata,
)
88
from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
wangding zeng's avatar
wangding zeng committed
89

90
from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsPP
91
92
93
94
95
96
97
from .utils import (
    PPMissingLayer,
    is_pp_missing_parameter,
    make_empty_intermediate_tensors_factory,
    make_layers,
    maybe_prefix,
)
98

99
100
101
102
103
if current_platform.is_cuda_alike():
    from vllm import _custom_ops as ops
elif current_platform.is_xpu():
    from vllm._ipex_ops import ipex_ops as ops

104

105
106
logger = init_logger(__name__)

wangding zeng's avatar
wangding zeng committed
107

108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
class DeepseekAttention(nn.Module):
    """Normal MHA implementation used by Deepseek v1."""

    def __init__(
        self,
        vllm_config: VllmConfig,
        config: DeepseekV2Config | DeepseekV3Config,
        hidden_size: int,
        num_heads: int,
        max_position_embeddings: int = 8192,
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
        **kwargs,
    ) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        tp_size = get_tensor_model_parallel_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = config.num_key_value_heads
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
        self.head_dim = hidden_size // self.total_num_heads
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5
        self.max_position_embeddings = max_position_embeddings

        self.qkv_proj = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=False,
            quant_config=quant_config,
        )

        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
            quant_config=quant_config,
        )

        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position_embeddings,
165
            rope_parameters=config.rope_parameters,
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
        )
        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            num_kv_heads=self.num_kv_heads,
            cache_config=cache_config,
            quant_config=quant_config,
            prefix=f"{prefix}.attn",
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q, k = self.rotary_emb(positions, q, k)
        attn_output = self.attn(q, k, v)
        output, _ = self.o_proj(attn_output)
        return output

wangding zeng's avatar
wangding zeng committed
189
190
191
192
193
194
195

class DeepseekV2MLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
196
        quant_config: QuantizationConfig | None = None,
wangding zeng's avatar
wangding zeng committed
197
        reduce_results: bool = True,
198
        is_sequence_parallel=False,
199
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
200
201
    ) -> None:
        super().__init__()
202
203
204
205
206

        # If is_sequence_parallel, the input and output tensors are sharded
        # across the ranks within the tp_group. In this case the weights are
        # replicated and no collective ops are needed.
        # Otherwise we use standard TP with an allreduce at the end.
wangding zeng's avatar
wangding zeng committed
207
        self.gate_up_proj = MergedColumnParallelLinear(
208
209
            hidden_size,
            [intermediate_size] * 2,
wangding zeng's avatar
wangding zeng committed
210
            bias=False,
211
            quant_config=quant_config,
212
            disable_tp=is_sequence_parallel,
213
214
215
216
217
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
wangding zeng's avatar
wangding zeng committed
218
            bias=False,
219
            quant_config=quant_config,
220
            reduce_results=reduce_results,
221
            disable_tp=is_sequence_parallel,
222
223
            prefix=f"{prefix}.down_proj",
        )
wangding zeng's avatar
wangding zeng committed
224
        if hidden_act != "silu":
225
226
227
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )
wangding zeng's avatar
wangding zeng committed
228
229
        self.act_fn = SiluAndMul()

230
231
232
233
234
235
236
    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)
237
238
239
240
241
            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)
242
243
244
245
246
247
            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
248
249
250
251
252


class DeepseekV2MoE(nn.Module):
    def __init__(
        self,
253
        config: DeepseekV2Config | DeepseekV3Config,
254
        parallel_config: ParallelConfig,
255
        quant_config: QuantizationConfig | None = None,
256
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
257
258
259
    ):
        super().__init__()
        self.tp_size = get_tensor_model_parallel_world_size()
260
261
        self.tp_rank = get_tensor_model_parallel_rank()

262
        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
263
264

        self.ep_group = get_ep_group().device_group
265
        self.ep_rank = get_ep_group().rank_in_group
266
267
268
        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
269

270
        self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
271

272
        if config.hidden_act != "silu":
273
274
275
276
277
278
279
280
281
282
283
284
            raise ValueError(
                f"Unsupported activation: {config.hidden_act}. "
                "Only silu is supported for now."
            )

        self.gate = ReplicatedLinear(
            config.hidden_size,
            config.n_routed_experts,
            bias=False,
            quant_config=None,
            prefix=f"{prefix}.gate",
        )
285
        if getattr(config, "topk_method", None) == "noaux_tc":
286
            self.gate.e_score_correction_bias = nn.Parameter(
287
288
                torch.empty(config.n_routed_experts, dtype=torch.float32)
            )
289
290
291
        else:
            self.gate.e_score_correction_bias = None

292
        # Load balancing settings.
293
294
        eplb_config = parallel_config.eplb_config
        self.enable_eplb = parallel_config.enable_eplb
295

296
        self.n_redundant_experts = eplb_config.num_redundant_experts
297
        self.n_logical_experts = self.n_routed_experts
298
        self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
299
300
        self.n_local_physical_experts = self.n_physical_experts // self.ep_size

301
302
303
304
        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
        )
305

306
        self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
307
308
309
310
        self.is_fusion_moe_shared_experts_enabled = (
            rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
        )
        if config.n_shared_experts is None or self.is_fusion_moe_shared_experts_enabled:
311
312
            self.shared_experts = None
        else:
313
            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
314

wangding zeng's avatar
wangding zeng committed
315
316
317
318
319
            self.shared_experts = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
320
                is_sequence_parallel=self.is_sequence_parallel,
321
                reduce_results=False,
322
                prefix=f"{prefix}.shared_experts",
wangding zeng's avatar
wangding zeng committed
323
324
            )

325
326
        self.experts = SharedFusedMoE(
            shared_experts=self.shared_experts,
327
            gate=self.gate,
328
329
330
331
332
333
334
335
            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,
336
337
            num_expert_group=getattr(config, "n_group", 1),
            topk_group=getattr(config, "topk_group", 1),
338
            prefix=f"{prefix}.experts",
339
            scoring_func=getattr(config, "scoring_func", "softmax"),
340
            # we do scaling outside, set factor to 1.0 to avoid double mul
341
342
            # aiter applies routed_scaling_factor internally
            routed_scaling_factor=1.0
343
            if not self.is_rocm_aiter_moe_enabled
344
            else self.routed_scaling_factor,
345
346
347
348
            e_score_correction_bias=self.gate.e_score_correction_bias,
            enable_eplb=self.enable_eplb,
            num_redundant_experts=self.n_redundant_experts,
            is_sequence_parallel=self.is_sequence_parallel,
349
            n_shared_experts=config.n_shared_experts
350
            if self.is_fusion_moe_shared_experts_enabled
351
            else None,
352
        )
353

354
355
356
357
    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
358
359
        num_tokens, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)
360
361
362
363
364
365

        # Chunk the hidden states so they aren't replicated across TP ranks.
        # This avoids duplicate computation in self.experts.
        # TODO: We can replace the all_reduce at the end of attn with a
        # reduce_scatter instead of chunking here.
        if self.is_sequence_parallel:
366
            hidden_states = sequence_parallel_chunk(hidden_states)
367

368
369
370
371
372
        if self.experts.is_internal_router:
            # In this case, the gate/router runs inside the FusedMoE class
            fused_moe_out = self.experts(
                hidden_states=hidden_states, router_logits=hidden_states
            )
373
        else:
374
375
376
377
378
            # router_logits: (num_tokens, n_experts)
            router_logits, _ = self.gate(hidden_states)
            fused_moe_out = self.experts(
                hidden_states=hidden_states, router_logits=router_logits
            )
379

380
381
382
        shared_output, final_hidden_states = fused_moe_out
        if self.shared_experts is None:
            assert shared_output is None
383
384
385

        # Fix FP16 overflow
        # See DeepseekV2DecoderLayer for more details.
zhuwenwen's avatar
zhuwenwen committed
386
        if hidden_states.dtype != torch.float16:
387
            if not self.is_rocm_aiter_moe_enabled:
388
                final_hidden_states *= self.routed_scaling_factor
389
390
        elif self.shared_experts is not None:
            assert shared_output is not None
391
            shared_output *= 1.0 / self.routed_scaling_factor
392
393
394
395

        if self.shared_experts is not None:
            assert shared_output is not None
            final_hidden_states += shared_output
396

397
398
        if self.is_sequence_parallel:
            final_hidden_states = tensor_model_parallel_all_gather(
399
400
                final_hidden_states, 0
            )
401
402
            final_hidden_states = final_hidden_states[:num_tokens]
        elif self.tp_size > 1:
403
404
405
            final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
                final_hidden_states
            )
wangding zeng's avatar
wangding zeng committed
406
407
408
409
410
411

        return final_hidden_states.view(num_tokens, hidden_dim)


def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
    import math
412

wangding zeng's avatar
wangding zeng committed
413
414
415
416
417
    if scale <= 1:
        return 1.0
    return 0.1 * mscale * math.log(scale) + 1.0


418
419
420
421
422
423
424
425
426
def _get_llama_4_scaling(
    original_max_position_embeddings: int, scaling_beta: float, positions: torch.Tensor
) -> torch.Tensor:
    scaling = 1 + scaling_beta * torch.log(
        1 + torch.floor(positions / original_max_position_embeddings)
    )
    # Broadcast over num_heads and head_dim
    return scaling[..., None, None]

wangding zeng's avatar
wangding zeng committed
427
428
429
430

class DeepseekV2Attention(nn.Module):
    def __init__(
        self,
431
        vllm_config: VllmConfig,
432
        config: DeepseekV2Config | DeepseekV3Config,
wangding zeng's avatar
wangding zeng committed
433
434
435
436
437
438
439
440
        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,
        max_position_embeddings: int = 8192,
441
442
443
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
        topk_indices_buffer: torch.Tensor | None = None,
444
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
    ) -> 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.max_position_embeddings = max_position_embeddings
460
461
        assert topk_indices_buffer is None, (
            "topk_indices_buffer is not \
462
        supported for DeepseekV2Attention"
463
        )
wangding zeng's avatar
wangding zeng committed
464
465

        if self.q_lora_rank is not None:
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
            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_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,
                quant_config=quant_config,
                prefix=f"{prefix}.q_b_proj",
            )
wangding zeng's avatar
wangding zeng committed
481
        else:
482
483
484
485
486
487
488
            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",
            )
wangding zeng's avatar
wangding zeng committed
489

490
491
492
493
494
        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,
495
496
497
            prefix=f"{prefix}.kv_a_proj_with_mqa",
        )
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
wangding zeng's avatar
wangding zeng committed
498
499
500
501
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
502
            quant_config=quant_config,
503
504
            prefix=f"{prefix}.kv_b_proj",
        )
wangding zeng's avatar
wangding zeng committed
505
        # O projection.
506
507
508
509
510
511
512
        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",
        )
513
        if config.rope_parameters["rope_type"] != "default":
514
515
516
517
518
            config.rope_parameters["rope_type"] = (
                "deepseek_yarn"
                if config.rope_parameters.get("apply_yarn_scaling", True)
                else "deepseek_llama_scaling"
            )
519

520
521
522
523
        self.rotary_emb = get_rope(
            qk_rope_head_dim,
            rotary_dim=qk_rope_head_dim,
            max_position=max_position_embeddings,
524
            rope_parameters=config.rope_parameters,
525
526
            is_neox_style=False,
        )
wangding zeng's avatar
wangding zeng committed
527

528
529
530
531
        if (
            config.rope_parameters["rope_type"] != "default"
            and config.rope_parameters["rope_type"] == "deepseek_yarn"
        ):
532
533
            mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
            scaling_factor = config.rope_parameters["factor"]
wangding zeng's avatar
wangding zeng committed
534
535
536
            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
            self.scaling = self.scaling * mscale * mscale

537
538
539
540
541
542
543
544
545
        self.attn = Attention(
            self.num_local_heads,
            self.qk_head_dim,
            self.scaling,
            num_kv_heads=self.num_local_heads,
            cache_config=cache_config,
            quant_config=quant_config,
            prefix=f"{prefix}.attn",
        )
wangding zeng's avatar
wangding zeng committed
546
547
548
549
550

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
551
        llama_4_scaling: torch.Tensor | None,
wangding zeng's avatar
wangding zeng committed
552
553
554
555
    ) -> torch.Tensor:
        if self.q_lora_rank is not None:
            q = self.q_a_proj(hidden_states)[0]
            q = self.q_a_layernorm(q)
556
            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
wangding zeng's avatar
wangding zeng committed
557
        else:
558
559
560
561
            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)
wangding zeng's avatar
wangding zeng committed
562
        latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
563
        kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
wangding zeng's avatar
wangding zeng committed
564
        latent_cache = latent_cache.unsqueeze(1)
565
        kv_a = self.kv_a_layernorm(kv_a)
wangding zeng's avatar
wangding zeng committed
566
        kv = self.kv_b_proj(kv_a)[0]
567
        kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
wangding zeng's avatar
wangding zeng committed
568
        k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
569
        k_pe = latent_cache[:, :, self.kv_lora_rank :]
570

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

573
        q[..., self.qk_nope_head_dim :] = q_pe
wangding zeng's avatar
wangding zeng committed
574
        k = torch.empty_like(q)
575
576
        k[..., : self.qk_nope_head_dim] = k_nope
        k[..., self.qk_nope_head_dim :] = k_pe
577
578
579
580
581

        # Apply llama 4 scaling if provided
        if llama_4_scaling is not None:
            q *= llama_4_scaling

582
583
        # padding value to qk_head_dim for alignment
        v = torch.nn.functional.pad(
584
585
            v, [0, self.qk_head_dim - self.v_head_dim], value=0
        ).view(-1, self.num_local_heads * self.qk_head_dim)
586
        attn_output = self.attn(q, k, v)
587
588
589
        attn_output = attn_output.view(-1, self.num_local_heads, self.qk_head_dim)[
            ..., : self.v_head_dim
        ].reshape(-1, self.num_local_heads * self.v_head_dim)
wangding zeng's avatar
wangding zeng committed
590
591
592
593
        output, _ = self.o_proj(attn_output)
        return output


594
class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
595
596
597
    def __init__(
        self, head_dim: int, dtype: torch.dtype, prefix: str, cache_config: CacheConfig
    ):
598
599
600
601
602
603
604
605
606
607
608
        super().__init__()
        self.kv_cache = [torch.tensor([])]
        self.head_dim = head_dim
        self.prefix = prefix
        self.cache_config = cache_config
        self.dtype = dtype
        compilation_config = get_current_vllm_config().compilation_config
        if prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self

609
    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
610
611
612
613
614
615
616
        return MLAAttentionSpec(  # Only has one vector instead of K + V
            block_size=self.cache_config.block_size,
            num_kv_heads=1,
            head_size=self.head_dim,
            dtype=self.dtype,
        )

617
    def forward(self): ...
618
619
620
621
622
623
624
625
626
627
628
629
630
631

    def get_attn_backend(self) -> AttentionBackend:
        return DeepseekV32IndexerBackend



def sparse_attn_indexer(
    hidden_states: torch.Tensor,
    k_cache_prefix: str,
    kv_cache: torch.Tensor,
    q_fp8: torch.Tensor,
    k: torch.Tensor,
    weights: torch.Tensor,
    quant_block_size: int,
632
    scale_fmt: str | None,
633
634
635
636
    topk_tokens: int,
    head_dim: int,
    max_model_len: int,
    total_seq_lens: int,
637
    topk_indices_buffer: torch.Tensor | None,
638
639
640
) -> torch.Tensor:
    # careful! this will be None in dummy run
    attn_metadata = get_forward_context().attn_metadata
641
    fp8_dtype = current_platform.fp8_dtype()
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
    # assert isinstance(attn_metadata, dict)
    if not isinstance(attn_metadata, dict):
        return sparse_attn_indexer_fake(
            hidden_states,
            k_cache_prefix,
            kv_cache,
            q_fp8,
            k,
            weights,
            quant_block_size,
            scale_fmt,
            topk_tokens,
            head_dim,
            max_model_len,
            total_seq_lens,
            topk_indices_buffer,
        )
    attn_metadata = attn_metadata[k_cache_prefix]
    assert isinstance(attn_metadata, DeepseekV32IndexerMetadata)
    slot_mapping = attn_metadata.slot_mapping
    has_decode = attn_metadata.num_decodes > 0
    has_prefill = attn_metadata.num_prefills > 0
    num_decode_tokens = attn_metadata.num_decode_tokens

    ops.indexer_k_quant_and_cache(
        k,
        kv_cache,
        slot_mapping,
        quant_block_size,
        scale_fmt,
    )

674
    topk_indices_buffer[: hidden_states.shape[0]] = -1
675
676
    if has_prefill:
        prefill_metadata = attn_metadata.prefill
677
        for chunk in prefill_metadata.chunks:
678
679
680
            k_fp8 = torch.empty(
                [chunk.total_seq_lens, head_dim],
                device=k.device,
681
                dtype=fp8_dtype,
682
683
            )
            k_scale = torch.empty(
684
685
686
                [chunk.total_seq_lens, 4],
                device=k.device,
                dtype=torch.uint8,
687
            )
688
            ops.cp_gather_indexer_k_quant_cache(
689
690
691
692
693
694
                kv_cache,
                k_fp8,
                k_scale,
                chunk.block_table,
                chunk.cu_seq_lens,
            )
695
696
697
698
699
700
            fp8_mqa_logits_func = fp8_mqa_logits
            if current_platform.is_rocm():
                from vllm.attention.ops.rocm_aiter_mla_sparse import rocm_fp8_mqa_logits

                fp8_mqa_logits_func = rocm_fp8_mqa_logits
            logits = fp8_mqa_logits_func(
701
                q_fp8[chunk.token_start : chunk.token_end],
702
                (k_fp8, k_scale.view(torch.float32)),
703
                weights[chunk.token_start : chunk.token_end],
704
705
706
                chunk.cu_seqlen_ks,
                chunk.cu_seqlen_ke,
            )
707
708
            num_rows = logits.shape[0]
            assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
709
710
711
            topk_indices = topk_indices_buffer[
                chunk.token_start : chunk.token_end, :topk_tokens
            ]
712
713
714
715
716
717
718
719
            torch.ops._C.top_k_per_row(
                logits,
                chunk.cu_seqlen_ks,
                chunk.cu_seqlen_ke,
                topk_indices,
                num_rows,
                logits.stride(0),
                logits.stride(1),
720
            )
721
722
723
724
725
726
727
728
729
730
731
732
733

    if has_decode:
        decode_metadata = attn_metadata.decode
        # kv_cache size requirement [num_block, block_size, n_head, head_dim],
        # we only have [num_block, block_size, head_dim],
        kv_cache = kv_cache.unsqueeze(-2)
        decode_lens = decode_metadata.decode_lens
        if decode_metadata.requires_padding:
            # pad in edge case where we have short chunked prefill length <
            # decode_threshold since we unstrictly split
            # prefill and decode by decode_threshold
            # (currently set to 1 + speculative tokens)
            padded_q_fp8_decode_tokens = pack_seq_triton(
734
735
                q_fp8[:num_decode_tokens], decode_lens
            )
736
737
        else:
            padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
738
739
                decode_lens.shape[0], -1, *q_fp8.shape[1:]
            )
740
741
742
743
744
        # TODO: move and optimize below logic with triton kernels
        batch_size = padded_q_fp8_decode_tokens.shape[0]
        next_n = padded_q_fp8_decode_tokens.shape[1]
        assert batch_size == decode_metadata.seq_lens.shape[0]
        num_padded_tokens = batch_size * next_n
745
746
747
748
749
750
751
752
        fp8_paged_mqa_logits_func = fp8_paged_mqa_logits
        if current_platform.is_rocm():
            from vllm.attention.ops.rocm_aiter_mla_sparse import (
                rocm_fp8_paged_mqa_logits,
            )

            fp8_paged_mqa_logits_func = rocm_fp8_paged_mqa_logits
        logits = fp8_paged_mqa_logits_func(
753
754
755
756
757
758
759
760
            padded_q_fp8_decode_tokens,
            kv_cache,
            weights[:num_padded_tokens],
            decode_metadata.seq_lens,
            decode_metadata.block_table,
            decode_metadata.schedule_metadata,
            max_model_len=max_model_len,
        )
761
762
        num_rows = logits.shape[0]
        assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
763
764
765
        topk_indices = topk_indices_buffer[:num_decode_tokens, :topk_tokens]

        torch.ops._C.top_k_per_row_decode(
766
            logits,
767
768
            next_n,
            decode_metadata.seq_lens,
769
770
771
772
773
            topk_indices,
            num_rows,
            logits.stride(0),
            logits.stride(1),
        )
774
775
776
777
778
        if decode_metadata.requires_padding:
            # if padded, we need to unpack
            # the topk indices removing padded tokens
            topk_indices = unpack_seq_triton(
                topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]),
779
780
                decode_lens,
            )
781
782
783
            topk_indices_buffer[:num_decode_tokens, : topk_indices.shape[-1]] = (
                topk_indices
            )
784
785
786
787
788
789
790
791
792
793
794
795

    return topk_indices_buffer


def sparse_attn_indexer_fake(
    hidden_states: torch.Tensor,
    k_cache_prefix: str,
    kv_cache: torch.Tensor,
    q_fp8: torch.Tensor,
    k: torch.Tensor,
    weights: torch.Tensor,
    quant_block_size: int,
796
    scale_fmt: str | None,
797
798
799
800
    topk_tokens: int,
    head_dim: int,
    max_model_len: int,
    total_seq_lens: int,
801
    topk_indices_buffer: torch.Tensor | None,
802
803
804
805
) -> torch.Tensor:
    # profile run
    # NOTE(Chen): create the max possible flattened_kv. So that
    # profile_run can get correct memory usage.
806
807
808
    _flattened_kv = torch.empty(
        [total_seq_lens, head_dim + 4], device=k.device, dtype=torch.uint8
    )
809
810
    fp8_dtype = current_platform.fp8_dtype()
    _k_fp8 = _flattened_kv[..., :head_dim].view(fp8_dtype).contiguous()
811
812
813
814
815
816
817
818
819
820
821
822
823
824
    _k_scale = _flattened_kv[..., head_dim:].view(torch.float32).contiguous()
    return topk_indices_buffer


direct_register_custom_op(
    op_name="sparse_attn_indexer",
    op_func=sparse_attn_indexer,
    mutates_args=["topk_indices_buffer"],
    fake_impl=sparse_attn_indexer_fake,
    dispatch_key=current_platform.dispatch_key,
)


class Indexer(nn.Module):
825
826
827
    def __init__(
        self,
        vllm_config: VllmConfig,
828
        config: DeepseekV2Config | DeepseekV3Config,
829
830
        hidden_size: int,
        q_lora_rank: int,
831
832
833
        quant_config: QuantizationConfig | None,
        cache_config: CacheConfig | None,
        topk_indices_buffer: torch.Tensor | None,
834
835
        prefix: str = "",
    ):
836
837
838
839
840
841
842
843
844
845
        super().__init__()
        self.vllm_config = vllm_config
        self.config = config
        # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
        self.topk_tokens = config.index_topk
        self.n_head = config.index_n_heads  # 64
        self.head_dim = config.index_head_dim  # 128
        self.rope_dim = config.qk_rope_head_dim  # 64
        self.q_lora_rank = q_lora_rank  # 1536
        # no tensor parallel, just replicated
846
847
848
849
850
851
852
853
854
855
856
857
858
859
        self.wq_b = ReplicatedLinear(
            self.q_lora_rank,
            self.head_dim * self.n_head,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.wq_b",
        )
        self.wk = ReplicatedLinear(
            hidden_size,
            self.head_dim,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.wk",
        )
860
        self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
861
862
863
        self.weights_proj = ReplicatedLinear(
            hidden_size, self.n_head, quant_config=None, prefix=f"{prefix}.weights_proj"
        )
864
865
866
867
868
869
870
871
872
873
        self.softmax_scale = self.head_dim**-0.5

        self.scale_fmt = "ue8m0"
        self.quant_block_size = 128  # TODO: get from config
        self.topk_indices_buffer = topk_indices_buffer

        # NOTE: (zyongye) we use fp8 naive cache,
        #       where we store value in fp8 and scale in fp32
        #       per self.quant_block_size element
        self.k_cache = DeepseekV32IndexerCache(
874
            head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,
875
876
            dtype=torch.uint8,
            prefix=f"{prefix}.k_cache",
877
878
            cache_config=cache_config,
        )
879
880
        self.max_model_len = vllm_config.model_config.max_model_len
        self.prefix = prefix
881
882
        from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size

883
884
        self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config)

885
886
887
    def forward(
        self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb
    ) -> torch.Tensor:
888
889
890
        q, _ = self.wq_b(qr)
        q = q.view(-1, self.n_head, self.head_dim)
        q_pe, q_nope = torch.split(
891
892
            q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
        )
893
894
895
896

        k, _ = self.wk(hidden_states)
        k = self.k_norm(k)
        k_pe, k_nope = torch.split(
897
898
            k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
        )
899
900

        q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
901
902
        q = torch.cat([q_pe.squeeze(0), q_nope], dim=-1)
        k = torch.cat([k_pe.squeeze((0, 2)), k_nope], dim=-1)
903
904
905

        # we only quant q here since k quant is fused with cache insertion
        q = q.view(-1, self.head_dim)
906
907
908
909
910
911
        q_fp8, q_scale = per_token_group_quant_fp8(
            q,
            self.quant_block_size,
            column_major_scales=False,
            use_ue8m0=self.scale_fmt is not None,
        )
912
913
914
915
        q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim)
        q_scale = q_scale.view(-1, self.n_head, 1)

        weights, _ = self.weights_proj(hidden_states)
916
917
918
        weights = (
            weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
        )
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
        weights = weights.squeeze(-1)

        return torch.ops.vllm.sparse_attn_indexer(
            hidden_states,
            self.k_cache.prefix,
            self.k_cache.kv_cache[0],
            q_fp8,
            k,
            weights,
            self.quant_block_size,
            self.scale_fmt,
            self.topk_tokens,
            self.head_dim,
            self.max_model_len,
            self.max_total_seq_len,
            self.topk_indices_buffer,
        )


938
939
940
941
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).
942

943
944
        For more info see MLACommonImpl in:
        vllm/v1/attention/backends/mla/utils.py
945
946
947
948
    """

    def __init__(
        self,
949
        vllm_config: VllmConfig,
950
        config: DeepseekV2Config | DeepseekV3Config,
951
952
953
954
955
        hidden_size: int,
        num_heads: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        v_head_dim: int,
956
        q_lora_rank: int | None,
957
958
        kv_lora_rank: int,
        max_position_embeddings: int = 8192,
959
960
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
961
        prefix: str = "",
962
        topk_indices_buffer: torch.Tensor | None = None,
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
    ) -> 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.max_position_embeddings = max_position_embeddings

        if self.q_lora_rank is not None:
983
            self.fused_qkv_a_proj = MergedColumnParallelLinear(
984
985
986
987
                self.hidden_size,
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                bias=False,
                quant_config=quant_config,
988
                prefix=f"{prefix}.fused_qkv_a_proj",
989
990
                disable_tp=True,
            )
991
992
993
994
995
996
        else:
            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,
997
998
                prefix=f"{prefix}.kv_a_proj_with_mqa",
            )
999
1000

        if self.q_lora_rank is not None:
1001
1002
1003
1004
1005
1006
1007
1008
            self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
            self.q_b_proj = ColumnParallelLinear(
                self.q_lora_rank,
                self.num_heads * self.qk_head_dim,
                bias=False,
                quant_config=quant_config,
                prefix=f"{prefix}.q_b_proj",
            )
1009
        else:
1010
1011
1012
1013
1014
1015
1016
1017
            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_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
1018
1019
1020
1021
1022
        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,
1023
1024
1025
1026
1027
1028
1029
1030
1031
            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",
        )
1032

1033
        if config.rope_parameters["rope_type"] != "default":
1034
1035
1036
1037
1038
1039
            config.rope_parameters["rope_type"] = (
                "deepseek_yarn"
                if config.rope_parameters.get("apply_yarn_scaling", True)
                else "deepseek_llama_scaling"
            )

1040
1041
1042
1043
        self.rotary_emb = get_rope(
            qk_rope_head_dim,
            rotary_dim=qk_rope_head_dim,
            max_position=max_position_embeddings,
1044
            rope_parameters=config.rope_parameters,
1045
1046
            is_neox_style=False,
        )
1047
1048
1049
1050
1051

        if (
            config.rope_parameters["rope_type"] != "default"
            and config.rope_parameters["rope_type"] == "deepseek_yarn"
        ):
1052
1053
            mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
            scaling_factor = config.rope_parameters["factor"]
1054
1055
1056
            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
            self.scaling = self.scaling * mscale * mscale

zhuwenwen's avatar
zhuwenwen committed
1057
        self.is_v32 = hasattr(config, "index_topk")
1058
1059

        if self.is_v32:
1060
1061
1062
1063
            self.indexer_rope_emb = get_rope(
                qk_rope_head_dim,
                rotary_dim=qk_rope_head_dim,
                max_position=max_position_embeddings,
1064
                rope_parameters=config.rope_parameters,
1065
1066
                is_neox_style=True,
            )
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
            self.indexer = Indexer(
                vllm_config,
                config,
                hidden_size,
                q_lora_rank,
                quant_config,
                cache_config,
                topk_indices_buffer,
                f"{prefix}.indexer",
            )
1077
        else:
1078
            self.indexer_rope_emb = None
1079
1080
            self.indexer = None

1081
1082
        mla_modules = MLAModules(
            kv_a_layernorm=self.kv_a_layernorm,
1083
            kv_b_proj=self.kv_b_proj,
1084
1085
1086
            rotary_emb=self.rotary_emb,
            o_proj=self.o_proj,
            fused_qkv_a_proj=self.fused_qkv_a_proj
1087
1088
            if self.q_lora_rank is not None
            else None,
1089
            kv_a_proj_with_mqa=self.kv_a_proj_with_mqa
1090
1091
1092
            if self.q_lora_rank is None
            else None,
            q_a_layernorm=self.q_a_layernorm if self.q_lora_rank is not None else None,
1093
1094
            q_b_proj=self.q_b_proj if self.q_lora_rank is not None else None,
            q_proj=self.q_proj if self.q_lora_rank is None else None,
1095
            indexer=self.indexer,
1096
            indexer_rotary_emb=self.indexer_rope_emb,
1097
1098
            is_sparse=self.is_v32,
            topk_indices_buffer=topk_indices_buffer,
1099
        )
1100

1101
        self.mla_attn = MultiHeadLatentAttentionWrapper(
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
            self.hidden_size,
            self.num_local_heads,
            self.scaling,
            self.qk_nope_head_dim,
            self.qk_rope_head_dim,
            self.v_head_dim,
            self.q_lora_rank,
            self.kv_lora_rank,
            mla_modules,
            cache_config,
            quant_config,
            prefix,
1114
1115
1116
1117
1118
1119
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
1120
        llama_4_scaling: torch.Tensor | None,
1121
    ) -> torch.Tensor:
1122
        return self.mla_attn(positions, hidden_states, llama_4_scaling)
1123
1124


wangding zeng's avatar
wangding zeng committed
1125
class DeepseekV2DecoderLayer(nn.Module):
1126
1127
1128
1129
    def __init__(
        self,
        vllm_config: VllmConfig,
        prefix: str,
1130
1131
        config: DeepseekV2Config | None = None,
        topk_indices_buffer: torch.Tensor | None = None,
1132
    ) -> None:
wangding zeng's avatar
wangding zeng committed
1133
        super().__init__()
1134

1135
1136
        if config is None:
            config = vllm_config.model_config.hf_config
1137
1138
1139
1140
1141
        model_config = vllm_config.model_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config
        parallel_config = vllm_config.parallel_config

wangding zeng's avatar
wangding zeng committed
1142
        self.hidden_size = config.hidden_size
1143
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
1144
        moe_layer_freq = getattr(config, "moe_layer_freq", 1)
1145
1146
        # DecoderLayers are created with `make_layers` which passes the prefix
        # with the layer's index.
1147
        layer_idx = int(prefix.split(sep=".")[-1])
1148
        self.layer_idx = layer_idx
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158

        # verify MLA attention specific fields
        qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
        qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
        v_head_dim = getattr(config, "v_head_dim", 0)
        kv_lora_rank = getattr(config, "kv_lora_rank", 0)
        use_mha = config.model_type == "deepseek" or all(
            dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)
        )

1159
1160
        self.use_mha = use_mha

1161
1162
1163
        if use_mha:
            attn_cls = DeepseekAttention
        elif model_config.use_mla:
1164
1165
1166
1167
            attn_cls = DeepseekV2MLAAttention
        else:
            attn_cls = DeepseekV2Attention
        self.self_attn = attn_cls(
1168
            vllm_config=vllm_config,
wangding zeng's avatar
wangding zeng committed
1169
1170
1171
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
1172
1173
1174
            qk_nope_head_dim=qk_nope_head_dim,
            qk_rope_head_dim=qk_rope_head_dim,
            v_head_dim=v_head_dim,
1175
            q_lora_rank=config.q_lora_rank if hasattr(config, "q_lora_rank") else None,
1176
            kv_lora_rank=kv_lora_rank,
wangding zeng's avatar
wangding zeng committed
1177
1178
1179
            max_position_embeddings=max_position_embeddings,
            cache_config=cache_config,
            quant_config=quant_config,
1180
            prefix=f"{prefix}.self_attn",
1181
            topk_indices_buffer=topk_indices_buffer,
wangding zeng's avatar
wangding zeng committed
1182
        )
1183

1184
1185
1186
        if (
            config.n_routed_experts is not None
            and layer_idx >= config.first_k_dense_replace
1187
            and layer_idx % moe_layer_freq == 0
1188
        ):
1189
1190
            self.mlp = DeepseekV2MoE(
                config=config,
1191
                parallel_config=parallel_config,
1192
1193
1194
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
            )
wangding zeng's avatar
wangding zeng committed
1195
1196
1197
1198
1199
1200
        else:
            self.mlp = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=config.intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
1201
                prefix=f"{prefix}.mlp",
wangding zeng's avatar
wangding zeng committed
1202
            )
1203
1204
1205
1206
        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
        )
1207
        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
wangding zeng's avatar
wangding zeng committed
1208
1209
1210
1211
1212

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
1213
        residual: torch.Tensor | None,
1214
        llama_4_scaling: torch.Tensor | None = None,
wangding zeng's avatar
wangding zeng committed
1215
    ) -> torch.Tensor:
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
        if envs.USE_FUSED_RMS_QUANT:
            # Fix residual FP16 overflow
            residual_fix_overflow = False
            
            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
                
zhuwenwen's avatar
zhuwenwen committed
1239
            if hidden_states.dtype == torch.float16:
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
                # 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,
zhuwenwen's avatar
zhuwenwen committed
1250
                        DeepseekV2MLP) and hidden_states.dtype == torch.float16:
1251
1252
1253
1254
1255
1256
1257
                # 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
1258
        else:
1259
1260
1261
1262
            # Self Attention
            # Fix residual FP16 overflow
            residual_fix_overflow = False
            if residual is None:
1263
                residual = hidden_states.clone()
1264
1265
1266
                hidden_states = self.input_layernorm(hidden_states)
                residual_fix_overflow = True
            else:
1267
1268
                hidden_states, residual = self.input_layernorm(hidden_states, residual)
            
1269

1270
1271
1272
1273
1274
1275
1276
        attn_kwargs = {
            "positions": positions,
            "hidden_states": hidden_states,
        }
        if not self.use_mha:
            attn_kwargs["llama_4_scaling"] = llama_4_scaling
        hidden_states = self.self_attn(**attn_kwargs)
wangding zeng's avatar
wangding zeng committed
1277

1278
1279
1280
1281
        if (
            not isinstance(self.self_attn, DeepseekAttention)
            and hidden_states.dtype == torch.float16
        ):
1282
1283
1284
            # Fix FP16 overflow
            # We scale both hidden_states and residual before
            # rmsnorm, and rmsnorm result would not affect by scale.
1285
            hidden_states *= 1.0 / self.routed_scaling_factor
1286
1287
1288
            if self.layer_idx == 0:
                # The residual is shared by all layers, we only scale it on
                # first layer.
1289
                residual *= 1.0 / self.routed_scaling_factor
1290
1291

        # Fully Connected
1292
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
wangding zeng's avatar
wangding zeng committed
1293
        hidden_states = self.mlp(hidden_states)
1294

1295
        if isinstance(self.mlp, DeepseekV2MLP) and hidden_states.dtype == torch.float16:
1296
1297
1298
1299
1300
            # 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
1301
            hidden_states *= 1.0 / self.routed_scaling_factor
1302

wangding zeng's avatar
wangding zeng committed
1303
1304
1305
        return hidden_states, residual


1306
@support_torch_compile
wangding zeng's avatar
wangding zeng committed
1307
1308
1309
class DeepseekV2Model(nn.Module):
    fall_back_to_pt_during_load = False

1310
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
1311
        super().__init__()
1312
1313
1314

        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
1315
        self.config = config
1316
        self.device = current_platform.device_type
1317

wangding zeng's avatar
wangding zeng committed
1318
        self.vocab_size = config.vocab_size
zhuwenwen's avatar
zhuwenwen committed
1319
        self.is_v32 = hasattr(config, "index_topk")
1320
1321
1322
1323
1324
1325
        if self.is_v32:
            topk_tokens = config.index_topk
            topk_indices_buffer = torch.empty(
                vllm_config.scheduler_config.max_num_batched_tokens,
                topk_tokens,
                dtype=torch.int32,
1326
                device=self.device,
1327
            )
1328
1329
        else:
            topk_indices_buffer = None
wangding zeng's avatar
wangding zeng committed
1330

1331
1332
1333
1334
        if get_pp_group().is_first_rank:
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
1335
                quant_config=quant_config,
1336
1337
                prefix=f"{prefix}.embed_tokens",
            )
1338
1339
1340
1341
        else:
            self.embed_tokens = PPMissingLayer()
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
1342
            lambda prefix: DeepseekV2DecoderLayer(
1343
                vllm_config, prefix, topk_indices_buffer=topk_indices_buffer
1344
1345
1346
            ),
            prefix=f"{prefix}.layers",
        )
1347
1348
1349
1350
1351

        if get_pp_group().is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer()
1352
1353
1354
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
wangding zeng's avatar
wangding zeng committed
1355

1356
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1357
1358
        return self.embed_tokens(input_ids)

wangding zeng's avatar
wangding zeng committed
1359
1360
1361
1362
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
1363
1364
1365
        intermediate_tensors: IntermediateTensors | None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
1366
        if get_pp_group().is_first_rank:
1367
1368
1369
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
1370
                hidden_states = self.embed_input_ids(input_ids)
1371
1372
1373
1374
1375
1376
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
        # Compute llama 4 scaling once per forward pass if enabled
        llama_4_scaling_config = getattr(self.config, "llama_4_scaling", None)
        llama_4_scaling: torch.Tensor | None
        if llama_4_scaling_config is not None:
            llama_4_scaling = _get_llama_4_scaling(
                original_max_position_embeddings=llama_4_scaling_config[
                    "original_max_position_embeddings"
                ],
                scaling_beta=llama_4_scaling_config["beta"],
                positions=positions,
            )
        else:
            llama_4_scaling = None

1391
        for layer in islice(self.layers, self.start_layer, self.end_layer):
1392
1393
1394
            hidden_states, residual = layer(
                positions, hidden_states, residual, llama_4_scaling
            )
1395
1396

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

wangding zeng's avatar
wangding zeng committed
1401
1402
1403
1404
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
class DeepseekV2MixtureOfExperts(MixtureOfExperts):
    moe_mlp_layers: list[DeepseekV2MoE]
    """
    List of MoE MLP layers in the model.
    """

    def extract_moe_parameters(self, example_moe: DeepseekV2MoE | None):
        if example_moe is None:
            self.num_moe_layers = 0
            self.num_expert_groups = 0
            self.num_logical_experts = 0
            self.num_physical_experts = 0
            self.num_local_physical_experts = 0
            self.num_routed_experts = 0
            self.num_shared_experts = 0
            self.num_redundant_experts = 0
            logger.warning("DeepSeekV2: No DeepseekV2MoE layer found in model.layers.")
        else:
            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

    def update_physical_experts_metadata(
        self,
        num_physical_experts: int,
        num_local_physical_experts: int,
    ) -> None:
        assert self.num_local_physical_experts == num_local_physical_experts
        self.num_physical_experts = num_physical_experts
        self.num_local_physical_experts = num_local_physical_experts
        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
        for moe in self.moe_mlp_layers:
            moe.n_local_physical_experts = num_local_physical_experts
            moe.n_physical_experts = num_physical_experts
            moe.n_redundant_experts = self.num_redundant_experts
            moe.experts.update_expert_map()


class DeepseekV2ForCausalLM(
1447
    nn.Module, SupportsPP, DeepseekV2MixtureOfExperts, SupportsLoRA, SupportsEagle
1448
):
1449
1450
1451
    packed_modules_mapping = {
        "gate_up_proj": ["gate_proj", "up_proj"],
    }
1452
    model_cls = DeepseekV2Model
wangding zeng's avatar
wangding zeng committed
1453

1454
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
1455
        super().__init__()
1456
1457
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
1458
1459
1460
1461
1462
1463
1464

        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'

1465
        self.use_w4a16_moe_sz = os.environ.get('AWQ_MOE_SZ') == '1'
wangding zeng's avatar
wangding zeng committed
1466
1467
        self.config = config
        self.quant_config = quant_config
王敏's avatar
王敏 committed
1468

1469
1470
1471
1472
1473
1474
1475
1476
1477
        qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
        qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
        self.use_mha = config.model_type == "deepseek" or all(
            dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)
        )

        if self.use_mha:
            self.packed_modules_mapping["qkv_proj"] = ["q_proj", "k_proj", "v_proj"]

1478
1479
1480
1481
        # `packed_modules_mapping` needs to be modified before
        # initializing DeepseekV2Model, as it is passed inplace to
        # quantization config init and may be used to select the
        # quant_method for relevant layers during initialization.
1482
1483
1484
        self.fuse_qkv_a_proj = (
            hasattr(config, "q_lora_rank") and config.q_lora_rank is not None
        )
1485
1486
1487
1488
1489
1490
        if self.fuse_qkv_a_proj:
            self.packed_modules_mapping["fused_qkv_a_proj"] = [
                "q_a_proj",
                "kv_a_proj_with_mqa",
            ]

1491
        self.model = self.model_cls(
1492
1493
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
1494
        if get_pp_group().is_last_rank:
1495
1496
1497
1498
1499
1500
            self.lm_head = ParallelLMHead(
                config.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "lm_head"),
            )
1501
1502
        else:
            self.lm_head = PPMissingLayer()
wangding zeng's avatar
wangding zeng committed
1503
        self.logits_processor = LogitsProcessor(config.vocab_size)
1504
        self.make_empty_intermediate_tensors = (
1505
1506
            self.model.make_empty_intermediate_tensors
        )
1507
1508
1509
1510
1511
1512
1513
        # Set MoE hyperparameters
        self.num_moe_layers = (
            self.config.num_hidden_layers - self.config.first_k_dense_replace
        )
        self.set_moe_parameters()

    def set_moe_parameters(self):
1514
1515
        self.expert_weights = []

1516
        self.num_expert_groups = getattr(self.config, "n_group", 1)
1517

1518
1519
        self.moe_layers = []
        self.moe_mlp_layers = []
1520
        example_moe = None
1521
        for layer in self.model.layers:
1522
1523
1524
            if isinstance(layer, PPMissingLayer):
                continue

1525
1526
            assert isinstance(layer, DeepseekV2DecoderLayer)
            if isinstance(layer.mlp, DeepseekV2MoE):
1527
1528
                # Pick last one layer since the first ones may be dense layers.
                example_moe = layer.mlp
1529
                self.moe_mlp_layers.append(layer.mlp)
1530
1531
                self.moe_layers.append(layer.mlp.experts)

1532
        self.extract_moe_parameters(example_moe)
王敏's avatar
王敏 committed
1533
        self.use_llama_nn = os.environ.get('LLAMA_NN') == '1'
1534
        self.use_awq_pad = os.environ.get('AWQ_PAD') == '1'
1535
1536
1537
        self.tritonsingleton= W8a8GetCacheJSON() 
        self.tritonsingleton.topk = config.num_experts_per_tok
        self.tritonsingleton.quant_method=self.quant_method 
1538

1539
1540
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
1541

wangding zeng's avatar
wangding zeng committed
1542
1543
1544
1545
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
1546
1547
1548
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
1549
1550
1551
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
wangding zeng's avatar
wangding zeng committed
1552
1553
        return hidden_states

1554
1555
1556
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
1557
    ) -> torch.Tensor | None:
1558
        logits = self.logits_processor(self.lm_head, hidden_states)
wangding zeng's avatar
wangding zeng committed
1559
1560
        return logits

1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
        return SharedFusedMoE.make_expert_params_mapping(
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
            num_experts=self.config.n_routed_experts,
            num_redundant_experts=0,
        )
1571

1572
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1573
1574
1575
        rocm_aiter_moe_shared_expert_enabled = (
            rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
        )
wangding zeng's avatar
wangding zeng committed
1576
1577
1578
1579
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
1580
1581
        ]
        mla_params_mapping = [
1582
1583
            ("fused_qkv_a_proj", "q_a_proj", 0),
            ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
wangding zeng's avatar
wangding zeng committed
1584
        ]
1585
1586
1587
1588
1589
1590
1591
1592
1593
        mha_params_mapping = [
            ("qkv_proj", "q_proj", "q"),
            ("qkv_proj", "k_proj", "k"),
            ("qkv_proj", "v_proj", "v"),
        ]
        if self.use_mha:
            stacked_params_mapping.extend(mha_params_mapping)
        else:
            stacked_params_mapping.extend(mla_params_mapping)
wangding zeng's avatar
wangding zeng committed
1594

1595
1596
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
1597
        expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
王敏's avatar
王敏 committed
1598
1599
1600
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
1601
1602
1603
            num_experts=self.config.n_routed_experts
            + (
                self.config.n_shared_experts
1604
                if rocm_aiter_moe_shared_expert_enabled
1605
1606
                else 0
            ),
1607
1608
            num_redundant_experts=self.num_redundant_experts,
        )
1609

wangding zeng's avatar
wangding zeng committed
1610
        params_dict = dict(self.named_parameters())
1611
        loaded_params: set[str] = set()
wangding zeng's avatar
wangding zeng committed
1612
1613
1614
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
1615

1616
1617
1618
            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
1619

1620
1621
            is_fusion_moe_shared_experts_layer = (
                rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
1622
1623
            )

1624
            for param_name, weight_name, shard_id in stacked_params_mapping:
1625
                # Skip non-stacked layers and experts (experts handled below).
wangding zeng's avatar
wangding zeng committed
1626
1627
                if weight_name not in name:
                    continue
1628
1629
1630
1631
1632
1633
                # 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.
1634
                if ("mlp.experts." in name) and name not in params_dict:
1635
                    continue
1636
                if is_fusion_moe_shared_experts_layer:
1637
                    continue
1638
                name_mapped = name.replace(weight_name, param_name)
1639
1640
1641

                # QKV fusion is optional, fall back to normal
                # weight loading if it's not enabled
1642
                # if go with fusion option, then update name
1643
1644
1645
                if (
                    param_name == "fused_qkv_a_proj"
                ) and name_mapped not in params_dict:
1646
                    continue
1647
1648
                else:
                    name = name_mapped
wangding zeng's avatar
wangding zeng committed
1649
1650
1651
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
1652
1653
1654
1655

                if is_pp_missing_parameter(name, self):
                    continue

wangding zeng's avatar
wangding zeng committed
1656
1657
1658
1659
1660
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
1661
1662
                is_expert_weight = False

1663
1664
1665
1666
1667
1668
1669
1670
1671
                # Special handling: when AITER fusion_shared_experts is enabled,
                # checkpoints may provide a single widened shared_experts tensor
                # without explicit expert indices
                # (e.g. ...mlp.shared_experts.gate_proj.weight).
                # For models with multiple shared experts, split that tensor
                # evenly into per-shared-expert slices and load them into
                # appended expert slots mlp.experts.{n_routed_experts + j}.*
                # accordingly.
                num_chunks = 1
1672
                if is_fusion_moe_shared_experts_layer:
1673
1674
1675
1676
1677
1678
1679
1680
1681
                    num_chunks = getattr(self.config, "n_shared_experts", 1) or 1
                    # Determine split axis based on op type
                    # gate/up: ColumnParallel → split along dim 0
                    # down: RowParallel → split along dim 1
                    split_dim = 1 if "down_proj.weight" in name else 0
                    total = loaded_weight.shape[split_dim]
                    assert total % num_chunks == 0, (
                        f"Shared expert weight dim {total} "
                        f"not divisible by num_chunks {num_chunks}"
1682
                    )
1683
1684
1685
1686
1687
1688
                    chunk_size = total // num_chunks

                for j in range(num_chunks):
                    chunk_name = name
                    weight_to_load = loaded_weight

1689
                    if is_fusion_moe_shared_experts_layer:
1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
                        if split_dim == 0:
                            weight_to_load = loaded_weight[
                                j * chunk_size : (j + 1) * chunk_size, :
                            ]
                        else:
                            weight_to_load = loaded_weight[
                                :, j * chunk_size : (j + 1) * chunk_size
                            ]
                        # Synthesize an expert-style name so expert mapping
                        # can route it
                        chunk_name = name.replace(
                            "mlp.shared_experts",
                            f"mlp.experts.{self.config.n_routed_experts + j}",
                        )

                    # Use expert_params_mapping to locate the destination
                    # param and delegate to its expert-aware weight_loader
                    # with expert_id.
                    for mapping in expert_params_mapping:
                        param_name, weight_name, expert_id, shard_id = mapping
                        if weight_name not in chunk_name:
                            continue

                        # 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 = chunk_name.replace(weight_name, param_name)

                        if is_pp_missing_parameter(name_mapped, self):
                            continue

                        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,
                            weight_to_load,
                            name_mapped,
                            shard_id=shard_id,
                            expert_id=expert_id,
                            return_success=True,
                        )
                        if success:
1740
                            if not is_fusion_moe_shared_experts_layer:
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
                                name = name_mapped
                            else:
                                loaded_params.add(name_mapped)
                            break
                    else:
                        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

                        # Skip loading extra bias for GPTQ models.
                        if name.endswith(".bias") and name not in params_dict:
                            continue

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

                        if is_pp_missing_parameter(name, self):
                            continue
1763

1764
                        param = params_dict[name]
1765
1766
1767
1768
                        weight_loader = getattr(
                            param, "weight_loader", default_weight_loader
                        )
                        weight_loader(param, loaded_weight)
1769
            if not is_fusion_moe_shared_experts_layer:
1770
                loaded_params.add(name)
1771
                
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
        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
1801
            
1802
        return loaded_params
1803
1804


1805
1806
1807
1808
class DeepseekForCausalLM(DeepseekV2ForCausalLM):
    pass


1809
1810
class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
    pass
1811
1812


1813
1814
# Compatibility with
# https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/blob/main/configuration_deepseek.py
1815
def get_spec_layer_idx_from_weight_name(
1816
1817
    config: DeepseekV2Config | DeepseekV3Config, weight_name: str
) -> int | None:
1818
1819
1820
1821
    if (
        hasattr(config, "num_nextn_predict_layers")
        and config.num_nextn_predict_layers > 0
    ):
1822
1823
        layer_idx = config.num_hidden_layers
        for i in range(config.num_nextn_predict_layers):
1824
            if weight_name.startswith(f"model.layers.{layer_idx + i}."):
1825
1826
                return layer_idx + i
    return None