deepseek_v2.py 63.9 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."""
26

27
28
import typing
from collections.abc import Callable, Iterable
29
from itertools import islice
wangding zeng's avatar
wangding zeng committed
30
31
32

import torch
from torch import nn
33
from transformers import DeepseekV2Config, DeepseekV3Config
wangding zeng's avatar
wangding zeng committed
34

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

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

97
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

logger = init_logger(__name__)

wangding zeng's avatar
wangding zeng committed
104

105
106
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
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,
            max_position=max_position_embeddings,
161
            rope_parameters=config.rope_parameters,
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
        )
        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
186
187
188
189
190
191
class DeepseekV2MLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
192
        quant_config: QuantizationConfig | None = None,
wangding zeng's avatar
wangding zeng committed
193
        reduce_results: bool = True,
194
        is_sequence_parallel=False,
195
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
196
197
    ) -> None:
        super().__init__()
198
199
200
201
202

        # 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
203
        self.gate_up_proj = MergedColumnParallelLinear(
204
205
            hidden_size,
            [intermediate_size] * 2,
wangding zeng's avatar
wangding zeng committed
206
            bias=False,
207
            quant_config=quant_config,
208
            disable_tp=is_sequence_parallel,
209
210
211
212
213
214
215
216
217
218
219
            prefix=f"{prefix}.gate_up_proj",
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
            bias=False,
            quant_config=quant_config,
            reduce_results=reduce_results,
            disable_tp=is_sequence_parallel,
            prefix=f"{prefix}.down_proj",
        )
wangding zeng's avatar
wangding zeng committed
220
        if hidden_act != "silu":
221
222
223
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )
wangding zeng's avatar
wangding zeng committed
224
225
226
227
228
229
230
231
232
233
234
235
        self.act_fn = SiluAndMul()

    def forward(self, x):
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x, _ = self.down_proj(x)
        return x


class DeepseekV2MoE(nn.Module):
    def __init__(
        self,
236
        config: DeepseekV2Config | DeepseekV3Config,
237
        parallel_config: ParallelConfig,
238
        quant_config: QuantizationConfig | None = None,
239
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
240
241
242
    ):
        super().__init__()
        self.tp_size = get_tensor_model_parallel_world_size()
243
244
        self.tp_rank = get_tensor_model_parallel_rank()

245
        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
246
247

        self.ep_group = get_ep_group().device_group
248
        self.ep_rank = get_ep_group().rank_in_group
249
250
251
        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
252

253
        self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
254

255
        if config.hidden_act != "silu":
256
257
258
259
260
261
262
263
264
265
266
267
            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",
        )
268
        if getattr(config, "topk_method", None) == "noaux_tc":
269
            self.gate.e_score_correction_bias = nn.Parameter(
270
271
                torch.empty(config.n_routed_experts, dtype=torch.float32)
            )
272
273
274
        else:
            self.gate.e_score_correction_bias = None

275
        # Load balancing settings.
276
277
        eplb_config = parallel_config.eplb_config
        self.enable_eplb = parallel_config.enable_eplb
278

279
        self.n_redundant_experts = eplb_config.num_redundant_experts
280
        self.n_logical_experts = self.n_routed_experts
281
        self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts
282
283
        self.n_local_physical_experts = self.n_physical_experts // self.ep_size

284
285
286
287
        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
        )
288

289
        self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
290
291
292
293
        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:
294
295
            self.shared_experts = None
        else:
296
            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
297

wangding zeng's avatar
wangding zeng committed
298
299
300
301
302
            self.shared_experts = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
303
                is_sequence_parallel=self.is_sequence_parallel,
304
                reduce_results=False,
305
                prefix=f"{prefix}.shared_experts",
wangding zeng's avatar
wangding zeng committed
306
307
            )

308
309
        self.experts = SharedFusedMoE(
            shared_experts=self.shared_experts,
310
            gate=self.gate,
311
312
313
314
315
316
317
318
            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,
319
320
            num_expert_group=getattr(config, "n_group", 1),
            topk_group=getattr(config, "topk_group", 1),
321
            prefix=f"{prefix}.experts",
322
            scoring_func=getattr(config, "scoring_func", "softmax"),
323
            # we do scaling outside, set factor to 1.0 to avoid double mul
324
325
            # aiter applies routed_scaling_factor internally
            routed_scaling_factor=1.0
326
            if not self.is_rocm_aiter_moe_enabled
327
            else self.routed_scaling_factor,
328
329
330
331
            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,
332
            n_shared_experts=config.n_shared_experts
333
            if self.is_fusion_moe_shared_experts_enabled
334
            else None,
335
        )
336

wangding zeng's avatar
wangding zeng committed
337
338
339
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        num_tokens, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)
340
341
342
343
344
345

        # 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:
346
            hidden_states = sequence_parallel_chunk(hidden_states)
347

348
349
350
351
352
353
354
355
356
357
358
        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
            )
        else:
            # 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
            )
359

360
361
362
        shared_output, final_hidden_states = fused_moe_out
        if self.shared_experts is None:
            assert shared_output is None
363
364
365
366

        # Fix FP16 overflow
        # See DeepseekV2DecoderLayer for more details.
        if hidden_states.dtype != torch.float16:
367
            if not self.is_rocm_aiter_moe_enabled:
368
                final_hidden_states *= self.routed_scaling_factor
369
370
        elif self.shared_experts is not None:
            assert shared_output is not None
371
            shared_output *= 1.0 / self.routed_scaling_factor
372
373
374
375

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

377
378
        if self.is_sequence_parallel:
            final_hidden_states = tensor_model_parallel_all_gather(
379
380
                final_hidden_states, 0
            )
381
382
            final_hidden_states = final_hidden_states[:num_tokens]
        elif self.tp_size > 1:
383
384
385
            final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
                final_hidden_states
            )
wangding zeng's avatar
wangding zeng committed
386
387
388
389
390
391

        return final_hidden_states.view(num_tokens, hidden_dim)


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

wangding zeng's avatar
wangding zeng committed
393
394
395
396
397
    if scale <= 1:
        return 1.0
    return 0.1 * mscale * math.log(scale) + 1.0


398
399
400
401
402
403
404
405
406
407
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
408
409
410
class DeepseekV2Attention(nn.Module):
    def __init__(
        self,
411
        vllm_config: VllmConfig,
412
        config: DeepseekV2Config | DeepseekV3Config,
wangding zeng's avatar
wangding zeng committed
413
414
415
416
417
418
419
420
        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,
421
422
423
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
        topk_indices_buffer: torch.Tensor | None = None,
424
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
    ) -> 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
440
441
        assert topk_indices_buffer is None, (
            "topk_indices_buffer is not \
442
        supported for DeepseekV2Attention"
443
        )
wangding zeng's avatar
wangding zeng committed
444
445

        if self.q_lora_rank is not None:
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
            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
461
        else:
462
463
464
465
466
467
468
            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
469

470
471
472
473
474
        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,
475
476
477
            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
478
479
480
481
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
482
            quant_config=quant_config,
483
484
            prefix=f"{prefix}.kv_b_proj",
        )
wangding zeng's avatar
wangding zeng committed
485
        # O projection.
486
487
488
489
490
491
492
        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",
        )
493
        if config.rope_parameters["rope_type"] != "default":
494
495
496
497
498
            config.rope_parameters["rope_type"] = (
                "deepseek_yarn"
                if config.rope_parameters.get("apply_yarn_scaling", True)
                else "deepseek_llama_scaling"
            )
499

500
501
502
        self.rotary_emb = get_rope(
            qk_rope_head_dim,
            max_position=max_position_embeddings,
503
            rope_parameters=config.rope_parameters,
504
505
            is_neox_style=False,
        )
wangding zeng's avatar
wangding zeng committed
506

507
508
509
510
        if (
            config.rope_parameters["rope_type"] != "default"
            and config.rope_parameters["rope_type"] == "deepseek_yarn"
        ):
511
512
            mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
            scaling_factor = config.rope_parameters["factor"]
wangding zeng's avatar
wangding zeng committed
513
514
515
            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
            self.scaling = self.scaling * mscale * mscale

516
517
518
519
520
521
522
523
524
        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
525
526
527
528
529

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
530
        llama_4_scaling: torch.Tensor | None,
wangding zeng's avatar
wangding zeng committed
531
532
533
534
    ) -> torch.Tensor:
        if self.q_lora_rank is not None:
            q = self.q_a_proj(hidden_states)[0]
            q = self.q_a_layernorm(q)
535
            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
wangding zeng's avatar
wangding zeng committed
536
        else:
537
538
539
540
            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
541
        latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
542
        kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
wangding zeng's avatar
wangding zeng committed
543
        latent_cache = latent_cache.unsqueeze(1)
544
        kv_a = self.kv_a_layernorm(kv_a)
wangding zeng's avatar
wangding zeng committed
545
        kv = self.kv_b_proj(kv_a)[0]
546
        kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
wangding zeng's avatar
wangding zeng committed
547
        k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
548
        k_pe = latent_cache[:, :, self.kv_lora_rank :]
549

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

552
        q[..., self.qk_nope_head_dim :] = q_pe
wangding zeng's avatar
wangding zeng committed
553
        k = torch.empty_like(q)
554
555
        k[..., : self.qk_nope_head_dim] = k_nope
        k[..., self.qk_nope_head_dim :] = k_pe
556
557
558
559
560

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

561
562
        # padding value to qk_head_dim for alignment
        v = torch.nn.functional.pad(
563
564
            v, [0, self.qk_head_dim - self.v_head_dim], value=0
        ).view(-1, self.num_local_heads * self.qk_head_dim)
565
        attn_output = self.attn(q, k, v)
566
567
568
        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
569
570
571
572
        output, _ = self.o_proj(attn_output)
        return output


573
class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
574
575
576
    def __init__(
        self, head_dim: int, dtype: torch.dtype, prefix: str, cache_config: CacheConfig
    ):
577
578
579
580
581
582
583
584
585
586
587
        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

588
    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
589
590
591
592
593
594
595
        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,
        )

596
    def forward(self): ...
597
598
599
600
601
602
603
604
605
606
607
608
609

    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,
610
    scale_fmt: str | None,
611
612
613
614
    topk_tokens: int,
    head_dim: int,
    max_model_len: int,
    total_seq_lens: int,
615
    topk_indices_buffer: torch.Tensor | None,
616
617
618
) -> torch.Tensor:
    # careful! this will be None in dummy run
    attn_metadata = get_forward_context().attn_metadata
619
    fp8_dtype = current_platform.fp8_dtype()
620

621
622
    # assert isinstance(attn_metadata, dict)
    if not isinstance(attn_metadata, dict):
623
624
625
626
627
628
        # Reserve workspace for indexer during profiling run
        current_workspace_manager().get_simultaneous(
            ((total_seq_lens, head_dim), torch.float8_e4m3fn),
            ((total_seq_lens, 4), torch.uint8),
        )

629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
        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,
    )

659
    topk_indices_buffer[: hidden_states.shape[0]] = -1
660
661
    if has_prefill:
        prefill_metadata = attn_metadata.prefill
662
663
664
665
666
667
668
669

        # Get the full shared workspace buffers once (will allocate on first use)
        workspace_manager = current_workspace_manager()
        k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(
            ((total_seq_lens, head_dim), fp8_dtype),
            ((total_seq_lens, 4), torch.uint8),
        )

670
        for chunk in prefill_metadata.chunks:
671
672
            k_fp8 = k_fp8_full[: chunk.total_seq_lens]
            k_scale = k_scale_full[: chunk.total_seq_lens]
673
            ops.cp_gather_indexer_k_quant_cache(
674
675
676
677
678
679
                kv_cache,
                k_fp8,
                k_scale,
                chunk.block_table,
                chunk.cu_seq_lens,
            )
680
681
682
683
684
685
            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(
686
                q_fp8[chunk.token_start : chunk.token_end],
687
                (k_fp8, k_scale.view(torch.float32)),
688
                weights[chunk.token_start : chunk.token_end],
689
690
691
                chunk.cu_seqlen_ks,
                chunk.cu_seqlen_ke,
            )
692
            num_rows = logits.shape[0]
693
694
695
            topk_indices = topk_indices_buffer[
                chunk.token_start : chunk.token_end, :topk_tokens
            ]
696
            torch.ops._C.top_k_per_row_prefill(
697
698
699
700
701
702
703
                logits,
                chunk.cu_seqlen_ks,
                chunk.cu_seqlen_ke,
                topk_indices,
                num_rows,
                logits.stride(0),
                logits.stride(1),
704
                topk_tokens,
705
            )
706
707
708
709
710
711
712
713
714
715
716
717
718

    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(
719
720
                q_fp8[:num_decode_tokens], decode_lens
            )
721
722
        else:
            padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
723
724
                decode_lens.shape[0], -1, *q_fp8.shape[1:]
            )
725
726
727
728
729
        # 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
730
731
732
733
734
735
736
737
        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(
738
739
740
741
742
743
744
745
            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,
        )
746
        num_rows = logits.shape[0]
747
748
749
        topk_indices = topk_indices_buffer[:num_decode_tokens, :topk_tokens]

        torch.ops._C.top_k_per_row_decode(
750
            logits,
751
752
            next_n,
            decode_metadata.seq_lens,
753
754
755
756
            topk_indices,
            num_rows,
            logits.stride(0),
            logits.stride(1),
757
            topk_tokens,
758
        )
759
760
761
762
763
        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]),
764
765
                decode_lens,
            )
766
767
768
            topk_indices_buffer[:num_decode_tokens, : topk_indices.shape[-1]] = (
                topk_indices
            )
769
770
771
772
773
774
775
776
777
778
779
780

    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,
781
    scale_fmt: str | None,
782
783
784
785
    topk_tokens: int,
    head_dim: int,
    max_model_len: int,
    total_seq_lens: int,
786
    topk_indices_buffer: torch.Tensor | None,
787
788
789
790
791
792
793
794
795
796
797
798
799
800
) -> torch.Tensor:
    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):
801
802
803
    def __init__(
        self,
        vllm_config: VllmConfig,
804
        config: DeepseekV2Config | DeepseekV3Config,
805
806
        hidden_size: int,
        q_lora_rank: int,
807
808
809
        quant_config: QuantizationConfig | None,
        cache_config: CacheConfig | None,
        topk_indices_buffer: torch.Tensor | None,
810
811
        prefix: str = "",
    ):
812
813
814
815
816
817
818
819
820
821
        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
822
823
824
825
826
827
828
829
830
831
832
833
834
835
        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",
        )
836
        self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
837
        self.weights_proj = ReplicatedLinear(
838
839
840
841
842
            hidden_size,
            self.n_head,
            bias=False,
            quant_config=None,
            prefix=f"{prefix}.weights_proj",
843
        )
844
845
846
847
848
849
850
851
852
853
        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(
854
            head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,
855
856
            dtype=torch.uint8,
            prefix=f"{prefix}.k_cache",
857
858
            cache_config=cache_config,
        )
859
860
        self.max_model_len = vllm_config.model_config.max_model_len
        self.prefix = prefix
861
862
        from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size

863
864
        self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config)

865
866
867
    def forward(
        self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb
    ) -> torch.Tensor:
868
869
870
        q, _ = self.wq_b(qr)
        q = q.view(-1, self.n_head, self.head_dim)
        q_pe, q_nope = torch.split(
871
872
            q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
        )
873
874
875
876

        k, _ = self.wk(hidden_states)
        k = self.k_norm(k)
        k_pe, k_nope = torch.split(
877
878
            k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
        )
879
880

        q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
881
882
        q = torch.cat([q_pe.squeeze(0), q_nope], dim=-1)
        k = torch.cat([k_pe.squeeze((0, 2)), k_nope], dim=-1)
883
884
885

        # we only quant q here since k quant is fused with cache insertion
        q = q.view(-1, self.head_dim)
886
887
888
889
890
891
        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,
        )
892
893
894
895
        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)
896
897
898
        weights = (
            weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
        )
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
        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,
        )


918
919
920
921
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).
922

923
924
        For more info see MLACommonImpl in:
        vllm/v1/attention/backends/mla/utils.py
925
926
927
928
    """

    def __init__(
        self,
929
        vllm_config: VllmConfig,
930
        config: DeepseekV2Config | DeepseekV3Config,
931
932
933
934
935
        hidden_size: int,
        num_heads: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        v_head_dim: int,
936
        q_lora_rank: int | None,
937
938
        kv_lora_rank: int,
        max_position_embeddings: int = 8192,
939
940
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
941
        prefix: str = "",
942
        topk_indices_buffer: torch.Tensor | None = None,
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
    ) -> 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:
963
            self.fused_qkv_a_proj = MergedColumnParallelLinear(
964
965
966
967
                self.hidden_size,
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                bias=False,
                quant_config=quant_config,
968
                prefix=f"{prefix}.fused_qkv_a_proj",
969
970
                disable_tp=True,
            )
971
972
973
974
975
976
        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,
977
978
                prefix=f"{prefix}.kv_a_proj_with_mqa",
            )
979
980

        if self.q_lora_rank is not None:
981
982
983
984
985
986
987
988
            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",
            )
989
        else:
990
991
992
993
994
995
996
997
            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)
998
999
1000
1001
1002
        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,
1003
1004
1005
1006
1007
1008
1009
1010
1011
            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",
        )
1012

1013
        if config.rope_parameters["rope_type"] != "default":
1014
1015
1016
1017
1018
1019
            config.rope_parameters["rope_type"] = (
                "deepseek_yarn"
                if config.rope_parameters.get("apply_yarn_scaling", True)
                else "deepseek_llama_scaling"
            )

1020
1021
1022
        self.rotary_emb = get_rope(
            qk_rope_head_dim,
            max_position=max_position_embeddings,
1023
            rope_parameters=config.rope_parameters,
1024
1025
            is_neox_style=False,
        )
1026
1027
1028
1029
1030

        if (
            config.rope_parameters["rope_type"] != "default"
            and config.rope_parameters["rope_type"] == "deepseek_yarn"
        ):
1031
1032
            mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
            scaling_factor = config.rope_parameters["factor"]
1033
1034
1035
            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
            self.scaling = self.scaling * mscale * mscale

1036
1037
1038
        self.is_v32 = hasattr(config, "index_topk")

        if self.is_v32:
1039
1040
1041
            self.indexer_rope_emb = get_rope(
                qk_rope_head_dim,
                max_position=max_position_embeddings,
1042
                rope_parameters=config.rope_parameters,
1043
1044
                is_neox_style=True,
            )
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
            self.indexer = Indexer(
                vllm_config,
                config,
                hidden_size,
                q_lora_rank,
                quant_config,
                cache_config,
                topk_indices_buffer,
                f"{prefix}.indexer",
            )
1055
        else:
1056
            self.indexer_rope_emb = None
1057
1058
            self.indexer = None

1059
1060
        mla_modules = MLAModules(
            kv_a_layernorm=self.kv_a_layernorm,
1061
            kv_b_proj=self.kv_b_proj,
1062
1063
1064
            rotary_emb=self.rotary_emb,
            o_proj=self.o_proj,
            fused_qkv_a_proj=self.fused_qkv_a_proj
1065
1066
            if self.q_lora_rank is not None
            else None,
1067
            kv_a_proj_with_mqa=self.kv_a_proj_with_mqa
1068
1069
1070
            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,
1071
1072
            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,
1073
            indexer=self.indexer,
1074
            indexer_rotary_emb=self.indexer_rope_emb,
1075
1076
            is_sparse=self.is_v32,
            topk_indices_buffer=topk_indices_buffer,
1077
        )
1078

1079
        self.mla_attn = MultiHeadLatentAttentionWrapper(
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
            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,
1092
1093
1094
1095
1096
1097
        )

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
1098
        llama_4_scaling: torch.Tensor | None,
1099
    ) -> torch.Tensor:
1100
        return self.mla_attn(positions, hidden_states, llama_4_scaling)
1101
1102


wangding zeng's avatar
wangding zeng committed
1103
class DeepseekV2DecoderLayer(nn.Module):
1104
1105
1106
1107
    def __init__(
        self,
        vllm_config: VllmConfig,
        prefix: str,
1108
1109
        config: DeepseekV2Config | None = None,
        topk_indices_buffer: torch.Tensor | None = None,
1110
    ) -> None:
wangding zeng's avatar
wangding zeng committed
1111
        super().__init__()
1112

1113
1114
        if config is None:
            config = vllm_config.model_config.hf_config
1115
1116
1117
1118
1119
        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
1120
        self.hidden_size = config.hidden_size
1121
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
1122
        moe_layer_freq = getattr(config, "moe_layer_freq", 1)
1123
1124
        # DecoderLayers are created with `make_layers` which passes the prefix
        # with the layer's index.
1125
        layer_idx = int(prefix.split(sep=".")[-1])
1126
        self.layer_idx = layer_idx
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136

        # 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)
        )

1137
1138
        self.use_mha = use_mha

1139
1140
1141
        if use_mha:
            attn_cls = DeepseekAttention
        elif model_config.use_mla:
1142
1143
1144
1145
            attn_cls = DeepseekV2MLAAttention
        else:
            attn_cls = DeepseekV2Attention
        self.self_attn = attn_cls(
1146
            vllm_config=vllm_config,
wangding zeng's avatar
wangding zeng committed
1147
1148
1149
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
1150
1151
1152
            qk_nope_head_dim=qk_nope_head_dim,
            qk_rope_head_dim=qk_rope_head_dim,
            v_head_dim=v_head_dim,
1153
            q_lora_rank=config.q_lora_rank if hasattr(config, "q_lora_rank") else None,
1154
            kv_lora_rank=kv_lora_rank,
wangding zeng's avatar
wangding zeng committed
1155
1156
1157
            max_position_embeddings=max_position_embeddings,
            cache_config=cache_config,
            quant_config=quant_config,
1158
            prefix=f"{prefix}.self_attn",
1159
            topk_indices_buffer=topk_indices_buffer,
wangding zeng's avatar
wangding zeng committed
1160
        )
1161

1162
1163
1164
        if (
            config.n_routed_experts is not None
            and layer_idx >= config.first_k_dense_replace
1165
            and layer_idx % moe_layer_freq == 0
1166
        ):
1167
1168
            self.mlp = DeepseekV2MoE(
                config=config,
1169
                parallel_config=parallel_config,
1170
1171
1172
                quant_config=quant_config,
                prefix=f"{prefix}.mlp",
            )
wangding zeng's avatar
wangding zeng committed
1173
1174
1175
1176
1177
1178
        else:
            self.mlp = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=config.intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
1179
                prefix=f"{prefix}.mlp",
wangding zeng's avatar
wangding zeng committed
1180
            )
1181
1182
1183
1184
        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
        )
1185
        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
wangding zeng's avatar
wangding zeng committed
1186
1187
1188
1189
1190

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
1191
        residual: torch.Tensor | None,
1192
        llama_4_scaling: torch.Tensor | None = None,
wangding zeng's avatar
wangding zeng committed
1193
1194
1195
    ) -> torch.Tensor:
        # Self Attention
        if residual is None:
1196
            residual = hidden_states.clone()
wangding zeng's avatar
wangding zeng committed
1197
1198
            hidden_states = self.input_layernorm(hidden_states)
        else:
1199
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
1200
1201
1202
1203
1204
1205
1206
1207

        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
1208

1209
1210
1211
1212
        if (
            not isinstance(self.self_attn, DeepseekAttention)
            and hidden_states.dtype == torch.float16
        ):
1213
1214
1215
            # Fix FP16 overflow
            # We scale both hidden_states and residual before
            # rmsnorm, and rmsnorm result would not affect by scale.
1216
            hidden_states *= 1.0 / self.routed_scaling_factor
1217
1218
1219
            if self.layer_idx == 0:
                # The residual is shared by all layers, we only scale it on
                # first layer.
1220
                residual *= 1.0 / self.routed_scaling_factor
1221
1222

        # Fully Connected
1223
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
wangding zeng's avatar
wangding zeng committed
1224
        hidden_states = self.mlp(hidden_states)
1225

1226
        if isinstance(self.mlp, DeepseekV2MLP) and hidden_states.dtype == torch.float16:
1227
1228
1229
1230
1231
            # 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
1232
            hidden_states *= 1.0 / self.routed_scaling_factor
1233

wangding zeng's avatar
wangding zeng committed
1234
1235
1236
        return hidden_states, residual


1237
@support_torch_compile
wangding zeng's avatar
wangding zeng committed
1238
1239
1240
class DeepseekV2Model(nn.Module):
    fall_back_to_pt_during_load = False

1241
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
wangding zeng's avatar
wangding zeng committed
1242
        super().__init__()
1243
1244
1245

        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
1246
        self.config = config
1247
        self.device = current_platform.device_type
1248

wangding zeng's avatar
wangding zeng committed
1249
        self.vocab_size = config.vocab_size
1250
1251
1252
1253
1254
1255
1256
        self.is_v32 = hasattr(config, "index_topk")
        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,
1257
                device=self.device,
1258
            )
1259
1260
        else:
            topk_indices_buffer = None
wangding zeng's avatar
wangding zeng committed
1261

1262
1263
1264
1265
        if get_pp_group().is_first_rank:
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
1266
                quant_config=quant_config,
1267
1268
                prefix=f"{prefix}.embed_tokens",
            )
1269
1270
1271
1272
        else:
            self.embed_tokens = PPMissingLayer()
        self.start_layer, self.end_layer, self.layers = make_layers(
            config.num_hidden_layers,
1273
            lambda prefix: DeepseekV2DecoderLayer(
1274
                vllm_config, prefix, topk_indices_buffer=topk_indices_buffer
1275
1276
1277
            ),
            prefix=f"{prefix}.layers",
        )
1278
1279
1280
1281
1282

        if get_pp_group().is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer()
1283
1284
1285
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], config.hidden_size
        )
wangding zeng's avatar
wangding zeng committed
1286

1287
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1288
1289
        return self.embed_tokens(input_ids)

wangding zeng's avatar
wangding zeng committed
1290
1291
1292
1293
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
1294
1295
1296
        intermediate_tensors: IntermediateTensors | None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
1297
        if get_pp_group().is_first_rank:
1298
1299
1300
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
1301
                hidden_states = self.embed_input_ids(input_ids)
1302
1303
1304
1305
1306
1307
            residual = None
        else:
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]

1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
        # 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

1322
        for layer in islice(self.layers, self.start_layer, self.end_layer):
1323
1324
1325
            hidden_states, residual = layer(
                positions, hidden_states, residual, llama_4_scaling
            )
1326
1327

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

wangding zeng's avatar
wangding zeng committed
1332
1333
1334
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states

1335

1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
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(
1378
    nn.Module, SupportsPP, DeepseekV2MixtureOfExperts, SupportsLoRA, SupportsEagle
1379
):
1380
1381
1382
    packed_modules_mapping = {
        "gate_up_proj": ["gate_proj", "up_proj"],
    }
1383
    model_cls = DeepseekV2Model
1384
1385
1386
1387
1388
1389
1390

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        self.config = config
        self.quant_config = quant_config
1391

1392
1393
1394
1395
1396
1397
1398
1399
1400
        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"]

1401
1402
1403
1404
        # `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.
1405
1406
1407
        self.fuse_qkv_a_proj = (
            hasattr(config, "q_lora_rank") and config.q_lora_rank is not None
        )
1408
1409
1410
1411
1412
1413
        if self.fuse_qkv_a_proj:
            self.packed_modules_mapping["fused_qkv_a_proj"] = [
                "q_a_proj",
                "kv_a_proj_with_mqa",
            ]

1414
        self.model = self.model_cls(
1415
1416
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
1417
        if get_pp_group().is_last_rank:
1418
1419
1420
1421
1422
1423
            self.lm_head = ParallelLMHead(
                config.vocab_size,
                config.hidden_size,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "lm_head"),
            )
1424
1425
1426
1427
        else:
            self.lm_head = PPMissingLayer()
        self.logits_processor = LogitsProcessor(config.vocab_size)
        self.make_empty_intermediate_tensors = (
1428
1429
            self.model.make_empty_intermediate_tensors
        )
1430
1431
1432
1433
1434
1435
1436
        # 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):
1437
1438
        self.expert_weights = []

1439
        self.num_expert_groups = getattr(self.config, "n_group", 1)
1440

1441
1442
        self.moe_layers = []
        self.moe_mlp_layers = []
1443
        example_moe = None
1444
        for layer in self.model.layers:
1445
1446
1447
            if isinstance(layer, PPMissingLayer):
                continue

1448
1449
            assert isinstance(layer, DeepseekV2DecoderLayer)
            if isinstance(layer.mlp, DeepseekV2MoE):
1450
1451
                # Pick last one layer since the first ones may be dense layers.
                example_moe = layer.mlp
1452
                self.moe_mlp_layers.append(layer.mlp)
1453
1454
                self.moe_layers.append(layer.mlp.experts)

1455
        self.extract_moe_parameters(example_moe)
1456

1457
1458
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
1459
1460
1461
1462
1463

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
1464
1465
1466
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
    ) -> torch.Tensor | IntermediateTensors:
1467
1468
1469
        hidden_states = self.model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
1470
1471
1472
1473
1474
        return hidden_states

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
1475
    ) -> torch.Tensor | None:
1476
        logits = self.logits_processor(self.lm_head, hidden_states)
1477
1478
        return logits

1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
    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,
        )

1490
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1491
1492
1493
        rocm_aiter_moe_shared_expert_enabled = (
            rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
        )
wangding zeng's avatar
wangding zeng committed
1494
1495
1496
1497
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
1498
1499
        ]
        mla_params_mapping = [
1500
1501
            ("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
1502
        ]
1503
1504
1505
1506
1507
1508
1509
1510
1511
        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
1512

1513
1514
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
1515
        expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
1516
1517
1518
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
1519
1520
1521
            num_experts=self.config.n_routed_experts
            + (
                self.config.n_shared_experts
1522
                if rocm_aiter_moe_shared_expert_enabled
1523
1524
                else 0
            ),
1525
1526
            num_redundant_experts=self.num_redundant_experts,
        )
1527

wangding zeng's avatar
wangding zeng committed
1528
        params_dict = dict(self.named_parameters())
1529
        loaded_params: set[str] = set()
wangding zeng's avatar
wangding zeng committed
1530
        for name, loaded_weight in weights:
1531
1532
1533
            if "rotary_emb.inv_freq" in name:
                continue

1534
1535
1536
            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
1537

1538
1539
            is_fusion_moe_shared_experts_layer = (
                rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
1540
1541
            )

1542
            for param_name, weight_name, shard_id in stacked_params_mapping:
1543
                # Skip non-stacked layers and experts (experts handled below).
wangding zeng's avatar
wangding zeng committed
1544
1545
                if weight_name not in name:
                    continue
1546
1547
1548
1549
1550
1551
                # 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.
1552
                if ("mlp.experts." in name) and name not in params_dict:
1553
                    continue
1554
                if is_fusion_moe_shared_experts_layer:
1555
                    continue
1556
                name_mapped = name.replace(weight_name, param_name)
1557
1558
1559

                # QKV fusion is optional, fall back to normal
                # weight loading if it's not enabled
1560
                # if go with fusion option, then update name
1561
1562
1563
                if (
                    param_name == "fused_qkv_a_proj"
                ) and name_mapped not in params_dict:
1564
                    continue
1565
1566
                else:
                    name = name_mapped
wangding zeng's avatar
wangding zeng committed
1567
1568
1569
                # Skip loading extra bias for GPTQ models.
                if name.endswith(".bias") and name not in params_dict:
                    continue
1570
1571
1572
1573

                if is_pp_missing_parameter(name, self):
                    continue

wangding zeng's avatar
wangding zeng committed
1574
1575
1576
1577
1578
                param = params_dict[name]
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
1579
                is_expert_weight = False
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589

                # 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
1590
                if is_fusion_moe_shared_experts_layer:
1591
1592
1593
1594
1595
1596
1597
1598
1599
                    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}"
1600
                    )
1601
1602
1603
1604
1605
1606
                    chunk_size = total // num_chunks

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

1607
                    if is_fusion_moe_shared_experts_layer:
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
                        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:
1658
                            if not is_fusion_moe_shared_experts_layer:
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
1686
                                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

                        param = params_dict[name]
                        weight_loader = getattr(
                            param, "weight_loader", default_weight_loader
                        )
                        weight_loader(param, loaded_weight)
1687
            if not is_fusion_moe_shared_experts_layer:
1688
                loaded_params.add(name)
1689

1690
        return loaded_params
1691
1692


1693
1694
1695
1696
class DeepseekForCausalLM(DeepseekV2ForCausalLM):
    pass


1697
1698
class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
    pass
1699
1700


1701
1702
# Compatibility with
# https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/blob/main/configuration_deepseek.py
1703
def get_spec_layer_idx_from_weight_name(
1704
1705
    config: DeepseekV2Config | DeepseekV3Config, weight_name: str
) -> int | None:
1706
1707
1708
1709
    if (
        hasattr(config, "num_nextn_predict_layers")
        and config.num_nextn_predict_layers > 0
    ):
1710
1711
        layer_idx = config.num_hidden_layers
        for i in range(config.num_nextn_predict_layers):
1712
            if weight_name.startswith(f"model.layers.{layer_idx + i}."):
1713
1714
                return layer_idx + i
    return None