"vllm/v1/attention/backend.py" did not exist on "b3230e1ac0f5c6b46f6c099ef0b3806df8a14feb"
deepseek_v2.py 63.8 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
wangding zeng's avatar
wangding zeng committed
86

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

96
97
98
99
100
101
102
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
103

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

        # 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
202
        self.gate_up_proj = MergedColumnParallelLinear(
203
204
            hidden_size,
            [intermediate_size] * 2,
wangding zeng's avatar
wangding zeng committed
205
            bias=False,
206
            quant_config=quant_config,
207
            disable_tp=is_sequence_parallel,
208
209
210
211
212
213
214
215
216
217
218
            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
219
        if hidden_act != "silu":
220
221
222
            raise ValueError(
                f"Unsupported activation: {hidden_act}. Only silu is supported for now."
            )
wangding zeng's avatar
wangding zeng committed
223
224
225
226
227
228
229
230
231
232
233
234
        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,
235
        config: DeepseekV2Config | DeepseekV3Config,
236
        parallel_config: ParallelConfig,
237
        quant_config: QuantizationConfig | None = None,
238
        prefix: str = "",
wangding zeng's avatar
wangding zeng committed
239
240
241
    ):
        super().__init__()
        self.tp_size = get_tensor_model_parallel_world_size()
242
243
        self.tp_rank = get_tensor_model_parallel_rank()

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

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

252
        self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
253

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return final_hidden_states.view(num_tokens, hidden_dim)


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

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


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

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

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

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

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

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

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

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

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

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

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


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

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

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

    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,
609
    scale_fmt: str | None,
610
611
612
613
    topk_tokens: int,
    head_dim: int,
    max_model_len: int,
    total_seq_lens: int,
614
    topk_indices_buffer: torch.Tensor | None,
615
616
617
) -> torch.Tensor:
    # careful! this will be None in dummy run
    attn_metadata = get_forward_context().attn_metadata
618
    fp8_dtype = current_platform.fp8_dtype()
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
    # 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,
    )

651
    topk_indices_buffer[: hidden_states.shape[0]] = -1
652
653
    if has_prefill:
        prefill_metadata = attn_metadata.prefill
654
        for chunk in prefill_metadata.chunks:
655
656
657
            k_fp8 = torch.empty(
                [chunk.total_seq_lens, head_dim],
                device=k.device,
658
                dtype=fp8_dtype,
659
660
            )
            k_scale = torch.empty(
661
662
663
                [chunk.total_seq_lens, 4],
                device=k.device,
                dtype=torch.uint8,
664
            )
665
            ops.cp_gather_indexer_k_quant_cache(
666
667
668
669
670
671
                kv_cache,
                k_fp8,
                k_scale,
                chunk.block_table,
                chunk.cu_seq_lens,
            )
672
673
674
675
676
677
            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(
678
                q_fp8[chunk.token_start : chunk.token_end],
679
                (k_fp8, k_scale.view(torch.float32)),
680
                weights[chunk.token_start : chunk.token_end],
681
682
683
                chunk.cu_seqlen_ks,
                chunk.cu_seqlen_ke,
            )
684
            num_rows = logits.shape[0]
685
686
687
            topk_indices = topk_indices_buffer[
                chunk.token_start : chunk.token_end, :topk_tokens
            ]
688
            torch.ops._C.top_k_per_row_prefill(
689
690
691
692
693
694
695
                logits,
                chunk.cu_seqlen_ks,
                chunk.cu_seqlen_ke,
                topk_indices,
                num_rows,
                logits.stride(0),
                logits.stride(1),
696
                topk_tokens,
697
            )
698
699
700
701
702
703
704
705
706
707
708
709
710

    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(
711
712
                q_fp8[:num_decode_tokens], decode_lens
            )
713
714
        else:
            padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
715
716
                decode_lens.shape[0], -1, *q_fp8.shape[1:]
            )
717
718
719
720
721
        # 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
722
723
724
725
726
727
728
729
        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(
730
731
732
733
734
735
736
737
            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,
        )
738
        num_rows = logits.shape[0]
739
740
741
        topk_indices = topk_indices_buffer[:num_decode_tokens, :topk_tokens]

        torch.ops._C.top_k_per_row_decode(
742
            logits,
743
744
            next_n,
            decode_metadata.seq_lens,
745
746
747
748
            topk_indices,
            num_rows,
            logits.stride(0),
            logits.stride(1),
749
            topk_tokens,
750
        )
751
752
753
754
755
        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]),
756
757
                decode_lens,
            )
758
759
760
            topk_indices_buffer[:num_decode_tokens, : topk_indices.shape[-1]] = (
                topk_indices
            )
761
762
763
764
765
766
767
768
769
770
771
772

    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,
773
    scale_fmt: str | None,
774
775
776
777
    topk_tokens: int,
    head_dim: int,
    max_model_len: int,
    total_seq_lens: int,
778
    topk_indices_buffer: torch.Tensor | None,
779
780
781
782
) -> torch.Tensor:
    # profile run
    # NOTE(Chen): create the max possible flattened_kv. So that
    # profile_run can get correct memory usage.
783
784
785
    _flattened_kv = torch.empty(
        [total_seq_lens, head_dim + 4], device=k.device, dtype=torch.uint8
    )
786
787
    fp8_dtype = current_platform.fp8_dtype()
    _k_fp8 = _flattened_kv[..., :head_dim].view(fp8_dtype).contiguous()
788
789
790
791
792
793
794
795
796
797
798
799
800
801
    _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):
802
803
804
    def __init__(
        self,
        vllm_config: VllmConfig,
805
        config: DeepseekV2Config | DeepseekV3Config,
806
807
        hidden_size: int,
        q_lora_rank: int,
808
809
810
        quant_config: QuantizationConfig | None,
        cache_config: CacheConfig | None,
        topk_indices_buffer: torch.Tensor | None,
811
812
        prefix: str = "",
    ):
813
814
815
816
817
818
819
820
821
822
        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
823
824
825
826
827
828
829
830
831
832
833
834
835
836
        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",
        )
837
        self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
838
839
840
        self.weights_proj = ReplicatedLinear(
            hidden_size, self.n_head, quant_config=None, prefix=f"{prefix}.weights_proj"
        )
841
842
843
844
845
846
847
848
849
850
        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(
851
            head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,
852
853
            dtype=torch.uint8,
            prefix=f"{prefix}.k_cache",
854
855
            cache_config=cache_config,
        )
856
857
        self.max_model_len = vllm_config.model_config.max_model_len
        self.prefix = prefix
858
859
        from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size

860
861
        self.max_total_seq_len = get_max_prefill_buffer_size(vllm_config)

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

        k, _ = self.wk(hidden_states)
        k = self.k_norm(k)
        k_pe, k_nope = torch.split(
874
875
            k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
        )
876
877

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

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


915
916
917
918
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).
919

920
921
        For more info see MLACommonImpl in:
        vllm/v1/attention/backends/mla/utils.py
922
923
924
925
    """

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

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

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

1017
1018
1019
        self.rotary_emb = get_rope(
            qk_rope_head_dim,
            max_position=max_position_embeddings,
1020
            rope_parameters=config.rope_parameters,
1021
1022
            is_neox_style=False,
        )
1023
1024
1025
1026
1027

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

1033
1034
1035
        self.is_v32 = hasattr(config, "index_topk")

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

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

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

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
1095
        llama_4_scaling: torch.Tensor | None,
1096
    ) -> torch.Tensor:
1097
        return self.mla_attn(positions, hidden_states, llama_4_scaling)
1098
1099


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

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

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

1134
1135
        self.use_mha = use_mha

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

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

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

        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
1205

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

        # Fully Connected
1220
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
wangding zeng's avatar
wangding zeng committed
1221
        hidden_states = self.mlp(hidden_states)
1222

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

wangding zeng's avatar
wangding zeng committed
1231
1232
1233
        return hidden_states, residual


1234
@support_torch_compile
wangding zeng's avatar
wangding zeng committed
1235
1236
1237
class DeepseekV2Model(nn.Module):
    fall_back_to_pt_during_load = False

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

        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
1243
        self.config = config
1244
        self.device = current_platform.device_type
1245

wangding zeng's avatar
wangding zeng committed
1246
        self.vocab_size = config.vocab_size
1247
1248
1249
1250
1251
1252
1253
        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,
1254
                device=self.device,
1255
            )
1256
1257
        else:
            topk_indices_buffer = None
wangding zeng's avatar
wangding zeng committed
1258

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

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

1284
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1285
1286
        return self.embed_tokens(input_ids)

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

1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
        # 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

1319
        for layer in islice(self.layers, self.start_layer, self.end_layer):
1320
1321
1322
            hidden_states, residual = layer(
                positions, hidden_states, residual, llama_4_scaling
            )
1323
1324

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

wangding zeng's avatar
wangding zeng committed
1329
1330
1331
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states

1332

1333
1334
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
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(
1375
    nn.Module, SupportsPP, DeepseekV2MixtureOfExperts, SupportsLoRA, SupportsEagle
1376
):
1377
1378
1379
    packed_modules_mapping = {
        "gate_up_proj": ["gate_proj", "up_proj"],
    }
1380
    model_cls = DeepseekV2Model
1381
1382
1383
1384
1385
1386
1387

    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
1388

1389
1390
1391
1392
1393
1394
1395
1396
1397
        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"]

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

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

1436
        self.num_expert_groups = getattr(self.config, "n_group", 1)
1437

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

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

1452
        self.extract_moe_parameters(example_moe)
1453

1454
1455
    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.model.embed_input_ids(input_ids)
1456
1457
1458
1459
1460

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

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
1472
    ) -> torch.Tensor | None:
1473
        logits = self.logits_processor(self.lm_head, hidden_states)
1474
1475
        return logits

1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
    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,
        )

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

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

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

1531
1532
1533
            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
1534

1535
1536
            is_fusion_moe_shared_experts_layer = (
                rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
1537
1538
            )

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

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

                if is_pp_missing_parameter(name, self):
                    continue

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

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

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

1604
                    if is_fusion_moe_shared_experts_layer:
1605
1606
1607
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
                        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:
1655
                            if not is_fusion_moe_shared_experts_layer:
1656
1657
1658
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
                                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)
1684
            if not is_fusion_moe_shared_experts_layer:
1685
                loaded_params.add(name)
1686

1687
        return loaded_params
1688
1689


1690
1691
1692
1693
class DeepseekForCausalLM(DeepseekV2ForCausalLM):
    pass


1694
1695
class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
    pass
1696
1697


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