deepseek_v2.py 137 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
# Copyright 2023-2024 SGLang Team
# 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.
# ==============================================================================
14

Liangsheng Yin's avatar
Liangsheng Yin committed
15
16
17
# Adapted from:
# https://github.com/vllm-project/vllm/blob/fb6af8bc086328ca6659e72d11ffd4309ce4de22/vllm/model_executor/models/deepseek_v2.py
"""Inference-only DeepseekV2 model."""
fzyzcjy's avatar
fzyzcjy committed
18
from __future__ import annotations
19

20
import concurrent.futures
21
import logging
22
import os
23
from enum import IntEnum, auto
24
from typing import Any, Dict, Iterable, Optional, Tuple, Union
Liangsheng Yin's avatar
Liangsheng Yin committed
25
26

import torch
Ke Bao's avatar
Ke Bao committed
27
import torch.nn.functional as F
Liangsheng Yin's avatar
Liangsheng Yin committed
28
29
from torch import nn
from transformers import PretrainedConfig
30

fzyzcjy's avatar
fzyzcjy committed
31
32
33
34
35
36
from sglang.srt.configs.model_config import (
    get_nsa_index_head_dim,
    get_nsa_index_n_heads,
    get_nsa_index_topk,
    is_deepseek_nsa,
)
37
from sglang.srt.distributed import (
38
    get_moe_expert_parallel_world_size,
39
    get_pp_group,
Liangsheng Yin's avatar
Liangsheng Yin committed
40
    get_tensor_model_parallel_world_size,
41
    parallel_state,
Liangsheng Yin's avatar
Liangsheng Yin committed
42
43
    tensor_model_parallel_all_reduce,
)
44
45
46
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
    use_symmetric_memory,
)
fzyzcjy's avatar
fzyzcjy committed
47
48
49
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
50
from sglang.srt.layers import deep_gemm_wrapper
51
from sglang.srt.layers.activation import SiluAndMul
52
from sglang.srt.layers.amx_utils import PackWeightMethod
53
54
55
56
from sglang.srt.layers.attention.npu_ops.mla_preprocess import (
    NPUFusedMLAPreprocess,
    is_mla_preprocess_enabled,
)
fzyzcjy's avatar
fzyzcjy committed
57
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
58
59
60
61
62
from sglang.srt.layers.communicator import (
    LayerCommunicator,
    LayerScatterModes,
    enable_moe_dense_fully_dp,
)
Lianmin Zheng's avatar
Lianmin Zheng committed
63
64
65
from sglang.srt.layers.dp_attention import (
    get_attention_tp_rank,
    get_attention_tp_size,
66
    is_dp_attention_enabled,
Lianmin Zheng's avatar
Lianmin Zheng committed
67
)
68
from sglang.srt.layers.layernorm import RMSNorm
69
70
71
72
73
74
from sglang.srt.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
)
Liangsheng Yin's avatar
Liangsheng Yin committed
75
from sglang.srt.layers.logits_processor import LogitsProcessor
76
77
78
79
from sglang.srt.layers.moe import (
    get_deepep_mode,
    get_moe_a2a_backend,
    should_use_flashinfer_cutlass_moe_fp4_allgather,
80
    should_use_flashinfer_trtllm_moe,
81
)
82
from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class
83
84
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat
85
from sglang.srt.layers.quantization.base_config import QuantizationConfig
86
from sglang.srt.layers.quantization.fp8_kernel import (
87
    is_fp8_fnuz,
88
    per_tensor_quant_mla_fp8,
89
    per_token_group_quant_mla_deep_gemm_masked_fp8,
90
)
HandH1998's avatar
HandH1998 committed
91
from sglang.srt.layers.quantization.fp8_utils import (
92
    block_quant_dequant,
HandH1998's avatar
HandH1998 committed
93
    block_quant_to_tensor_quant,
94
    channel_quant_to_tensor_quant,
95
    normalize_e4m3fn_to_e4m3fnuz,
96
    quant_weight_ue8m0,
97
    requant_weight_ue8m0_inplace,
98
    transform_scale_ue8m0_inplace,
HandH1998's avatar
HandH1998 committed
99
)
100
101
102
from sglang.srt.layers.quantization.int8_utils import (
    block_dequant as int8_block_dequant,
)
Liangsheng Yin's avatar
Liangsheng Yin committed
103
from sglang.srt.layers.radix_attention import RadixAttention
104
105
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
106
107
108
109
from sglang.srt.layers.vocab_parallel_embedding import (
    ParallelLMHead,
    VocabParallelEmbedding,
)
110
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
111
from sglang.srt.model_loader.weight_utils import default_weight_loader
112
from sglang.srt.server_args import get_global_server_args
113
from sglang.srt.single_batch_overlap import SboFlags
114
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
115
116
117
118
from sglang.srt.two_batch_overlap import (
    MaybeTboDeepEPDispatcher,
    model_forward_maybe_tbo,
)
119
120
from sglang.srt.utils import (
    BumpAllocator,
121
    LazyValue,
122
    add_prefix,
123
    bind_or_assign,
124
    cpu_has_amx_support,
125
    get_bool_env_var,
126
    get_device_sm,
127
    get_int_env_var,
128
    is_cpu,
129
    is_cuda,
130
    is_flashinfer_available,
131
    is_gfx95_supported,
132
    is_hip,
133
    is_non_idle_and_non_empty,
134
    is_npu,
135
    is_nvidia_cublas_cu12_version_ge_12_9,
136
    is_sm100_supported,
137
    log_info_on_rank0,
138
    make_layers,
139
    use_intel_amx_backend,
140
)
141

142
_is_hip = is_hip()
Yineng Zhang's avatar
Yineng Zhang committed
143
_is_cuda = is_cuda()
144
_is_npu = is_npu()
145
_is_fp8_fnuz = is_fp8_fnuz()
146
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
147
148
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu = is_cpu()
149
_device_sm = get_device_sm()
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
_is_gfx95_supported = is_gfx95_supported()

_use_aiter_gfx95 = _use_aiter and _is_gfx95_supported

if _use_aiter_gfx95:
    from sglang.srt.layers.quantization.quark.utils import quark_post_load_weights
    from sglang.srt.layers.quantization.rocm_mxfp4_utils import (
        batched_gemm_afp4wfp4_pre_quant,
        fused_flatten_mxfp4_quant,
        fused_rms_mxfp4_quant,
    )
    from sglang.srt.layers.rocm_linear_utils import (
        aiter_dsv3_router_gemm,
        fused_qk_rope_cat,
        get_dsv3_gemm_output_zero_allocator_size,
    )
166

Yineng Zhang's avatar
Yineng Zhang committed
167
if _is_cuda:
168
169
170
    from sgl_kernel import (
        awq_dequantize,
        bmm_fp8,
171
        concat_mla_k,
172
173
174
175
        dsv3_fused_a_gemm,
        dsv3_router_gemm,
        merge_state_v2,
    )
176
177
elif _is_cpu and _is_cpu_amx_available:
    pass
178
elif _is_hip:
fzyzcjy's avatar
fzyzcjy committed
179
180
181
    from sglang.srt.layers.attention.triton_ops.rocm_mla_decode_rope import (
        decode_attention_fwd_grouped_rope,
    )
182
183
184
    from sglang.srt.layers.quantization.awq_triton import (
        awq_dequantize_triton as awq_dequantize,
    )
fzyzcjy's avatar
fzyzcjy committed
185
elif _is_npu:
186
187
188
    import custom_ops  # noqa: F401
    import sgl_kernel_npu  # noqa: F401
    import torch_npu  # noqa: F401
Yineng Zhang's avatar
Yineng Zhang committed
189
else:
190
    pass
Liangsheng Yin's avatar
Liangsheng Yin committed
191

192
193
_is_flashinfer_available = is_flashinfer_available()
_is_sm100_supported = is_cuda() and is_sm100_supported()
194
_is_cublas_ge_129 = is_nvidia_cublas_cu12_version_ge_12_9()
195

196

197
198
logger = logging.getLogger(__name__)

199
200
201
202
203
204
205
206
207

def enable_nextn_moe_bf16_cast_to_fp8(quant_config):
    return (
        quant_config is not None
        and quant_config.get_name() == "modelopt_fp4"
        and get_moe_a2a_backend().is_deepep()
    )


208
209
FORWARD_ABSORB_CORE_ATTENTION_BACKENDS = [
    "fa3",
fzyzcjy's avatar
fzyzcjy committed
210
    "nsa",
211
212
213
214
215
216
217
218
219
220
221
222
    "flashinfer",
    "cutlass_mla",
    "trtllm_mla",
    "ascend",
]


def add_forward_absorb_core_attention_backend(backend_name):
    if backend_name not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS:
        FORWARD_ABSORB_CORE_ATTENTION_BACKENDS.append(backend_name)
        logger.info(f"Added {backend_name} to FORWARD_ABSORB_CORE_ATTENTION_BACKENDS.")

Liangsheng Yin's avatar
Liangsheng Yin committed
223

224
225
226
227
228
229
230
class AttnForwardMethod(IntEnum):
    # Use multi-head attention
    MHA = auto()

    # Use absorbed multi-latent attention
    MLA = auto()

fzyzcjy's avatar
fzyzcjy committed
231
232
233
    # Use Deepseek V3.2 sparse multi-latent attention
    NPU_MLA_SPARSE = auto()

234
235
236
237
    # Use multi-head attention, but with KV cache chunked.
    # This method can avoid OOM when prefix lengths are long.
    MHA_CHUNKED_KV = auto()

238
239
240
    # Use MLA but with fused RoPE
    MLA_FUSED_ROPE = auto()

241
242
243
    # Use MLA with fused RoPE kernel for CPU
    MLA_FUSED_ROPE_CPU = auto()

244

245
246
247
248
249
250
251
252
253
254
255
256
257
def _dispatch_mla_subtype(attn, forward_batch):
    if _is_hip:
        if attn.rocm_fused_decode_mla and forward_batch.forward_mode.is_decode():
            return AttnForwardMethod.MLA_FUSED_ROPE
        else:
            return AttnForwardMethod.MLA
    else:
        if hasattr(attn, "fused_qkv_a_proj_with_mqa") and use_intel_amx_backend(attn):
            return AttnForwardMethod.MLA_FUSED_ROPE_CPU
        else:
            return AttnForwardMethod.MLA


fzyzcjy's avatar
fzyzcjy committed
258
class AttentionBackendRegistry:
259
260
261
262
263
264
265
266
267
268
269
    _handlers = {}

    @classmethod
    def register(cls, backend_name, handler_func):
        cls._handlers[backend_name] = handler_func

    @classmethod
    def get_handler(cls, backend_name):
        return cls._handlers.get(backend_name, cls._handlers.get("triton"))


fzyzcjy's avatar
fzyzcjy committed
270
def handle_attention_ascend(attn, forward_batch):
271
272
273
274
275
    if (
        forward_batch.forward_mode.is_extend()
        and not forward_batch.forward_mode.is_target_verify()
        and not forward_batch.forward_mode.is_draft_extend()
    ):
fzyzcjy's avatar
fzyzcjy committed
276
277
278
279
        if hasattr(attn, "indexer"):
            return AttnForwardMethod.NPU_MLA_SPARSE
        else:
            return AttnForwardMethod.MHA
280
    else:
fzyzcjy's avatar
fzyzcjy committed
281
282
283
284
        if hasattr(attn, "indexer"):
            return AttnForwardMethod.NPU_MLA_SPARSE
        else:
            return AttnForwardMethod.MLA
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302


def _get_sum_extend_prefix_lens(forward_batch):
    return (
        sum(forward_batch.extend_prefix_lens_cpu)
        if forward_batch.extend_prefix_lens_cpu is not None
        else 0
    )


def _is_extend_without_speculative(forward_batch):
    return (
        forward_batch.forward_mode.is_extend()
        and not forward_batch.forward_mode.is_target_verify()
        and not forward_batch.forward_mode.is_draft_extend()
    )


fzyzcjy's avatar
fzyzcjy committed
303
304
305
def _handle_attention_backend(
    attn: DeepseekV2AttentionMLA, forward_batch, backend_name
):
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
    sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch)
    disable_ragged = (
        backend_name in ["flashinfer", "flashmla"]
    ) and attn.flashinfer_mla_disable_ragged

    if (
        not disable_ragged
        and _is_extend_without_speculative(forward_batch)
        and (
            (
                sum_extend_prefix_lens >= attn.chunked_prefix_cache_threshold
                and not attn.disable_chunked_prefix_cache
            )
            or sum_extend_prefix_lens == 0
        )
    ):
        return AttnForwardMethod.MHA_CHUNKED_KV
    else:
        return _dispatch_mla_subtype(attn, forward_batch)


fzyzcjy's avatar
fzyzcjy committed
327
328
def handle_attention_flashinfer(attn, forward_batch):
    return _handle_attention_backend(attn, forward_batch, "flashinfer")
329
330


fzyzcjy's avatar
fzyzcjy committed
331
332
def handle_attention_fa3(attn, forward_batch):
    return _handle_attention_backend(attn, forward_batch, "fa3")
333
334


fzyzcjy's avatar
fzyzcjy committed
335
336
def handle_attention_flashmla(attn, forward_batch):
    return _handle_attention_backend(attn, forward_batch, "flashmla")
337
338


fzyzcjy's avatar
fzyzcjy committed
339
340
def handle_attention_cutlass_mla(attn, forward_batch):
    return _handle_attention_backend(attn, forward_batch, "cutlass_mla")
341
342


fzyzcjy's avatar
fzyzcjy committed
343
def handle_attention_fa4(attn, forward_batch):
344
345
346
347
    # TODO(cicirori): use FA4 MHA for DeepSeekV3 for now
    return AttnForwardMethod.MHA_CHUNKED_KV


fzyzcjy's avatar
fzyzcjy committed
348
def handle_attention_trtllm_mla(attn, forward_batch):
349
350
351
352
353
354
355
356
357
    sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch)
    if _is_extend_without_speculative(forward_batch) and (
        not attn.disable_chunked_prefix_cache or sum_extend_prefix_lens == 0
    ):
        return AttnForwardMethod.MHA_CHUNKED_KV
    else:
        return _dispatch_mla_subtype(attn, forward_batch)


fzyzcjy's avatar
fzyzcjy committed
358
def handle_attention_aiter(attn, forward_batch):
359
360
361
362
363
364
365
366
367
368
369
370
    if _is_extend_without_speculative(forward_batch):
        if is_dp_attention_enabled():
            if sum(forward_batch.extend_prefix_lens_cpu) == 0:
                return AttnForwardMethod.MHA
            else:
                return AttnForwardMethod.MLA
        else:
            return AttnForwardMethod.MHA
    else:
        return AttnForwardMethod.MLA


fzyzcjy's avatar
fzyzcjy committed
371
372
373
374
def handle_attention_nsa(attn, forward_batch):
    return AttnForwardMethod.MLA


fzyzcjy's avatar
fzyzcjy committed
375
def handle_attention_triton(attn, forward_batch):
376
377
378
379
380
381
382
383
384
    if (
        _is_extend_without_speculative(forward_batch)
        and sum(forward_batch.extend_prefix_lens_cpu) == 0
    ):
        return AttnForwardMethod.MHA
    else:
        return _dispatch_mla_subtype(attn, forward_batch)


Liangsheng Yin's avatar
Liangsheng Yin committed
385
386
387
388
389
390
391
392
class DeepseekV2MLP(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
        quant_config: Optional[QuantizationConfig] = None,
        reduce_results: bool = True,
393
        prefix: str = "",
394
395
        tp_rank: Optional[int] = None,
        tp_size: Optional[int] = None,
Liangsheng Yin's avatar
Liangsheng Yin committed
396
397
    ) -> None:
        super().__init__()
398
399
        self.tp_size = tp_size

Liangsheng Yin's avatar
Liangsheng Yin committed
400
        self.gate_up_proj = MergedColumnParallelLinear(
401
402
403
404
405
            hidden_size,
            [intermediate_size] * 2,
            bias=False,
            quant_config=quant_config,
            prefix=add_prefix("gate_up_proj", prefix),
406
407
            tp_rank=tp_rank,
            tp_size=tp_size,
Liangsheng Yin's avatar
Liangsheng Yin committed
408
409
410
411
412
413
414
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
            bias=False,
            quant_config=quant_config,
            reduce_results=reduce_results,
415
            prefix=add_prefix("down_proj", prefix),
416
417
            tp_rank=tp_rank,
            tp_size=tp_size,
Liangsheng Yin's avatar
Liangsheng Yin committed
418
419
420
421
422
423
424
425
        )
        if hidden_act != "silu":
            raise ValueError(
                f"Unsupported activation: {hidden_act}. "
                "Only silu is supported for now."
            )
        self.act_fn = SiluAndMul()

426
427
428
429
    def forward(
        self,
        x,
        forward_batch=None,
430
        should_allreduce_fusion: bool = False,
431
        use_reduce_scatter: bool = False,
432
        gemm_output_zero_allocator: BumpAllocator = None,
433
    ):
434
435
436
        if (self.tp_size == 1) and x.shape[0] == 0:
            return x

437
438
439
440
441
        if (
            gemm_output_zero_allocator is not None
            and x.shape[0] <= 256
            and self.gate_up_proj.weight.dtype == torch.uint8
        ):
442
443
444
445
446
            y = gemm_output_zero_allocator.allocate(
                x.shape[0] * self.gate_up_proj.output_size_per_partition
            ).view(x.shape[0], self.gate_up_proj.output_size_per_partition)
            x = (x, None, y)

Liangsheng Yin's avatar
Liangsheng Yin committed
447
448
        gate_up, _ = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
449
        x, _ = self.down_proj(
450
            x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter
451
        )
Liangsheng Yin's avatar
Liangsheng Yin committed
452
453
454
        return x


Ke Bao's avatar
Ke Bao committed
455
class MoEGate(nn.Module):
456
457
458
    def __init__(
        self,
        config,
459
        quant_config,
460
        prefix: str = "",
461
        is_nextn: bool = False,
462
    ):
Ke Bao's avatar
Ke Bao committed
463
        super().__init__()
464
        self.is_nextn = is_nextn
Ke Bao's avatar
Ke Bao committed
465
466
467
468
        self.weight = nn.Parameter(
            torch.empty((config.n_routed_experts, config.hidden_size))
        )
        if config.topk_method == "noaux_tc":
469
470
471
472
473
474
475
            correction_bias_dtype = (
                torch.bfloat16
                if quant_config is not None
                and quant_config.get_name() == "modelopt_fp4"
                and should_use_flashinfer_trtllm_moe()
                else torch.float32
            )
Ke Bao's avatar
Ke Bao committed
476
            self.e_score_correction_bias = nn.Parameter(
477
                torch.empty((config.n_routed_experts), dtype=correction_bias_dtype)
Ke Bao's avatar
Ke Bao committed
478
479
480
            )
        else:
            self.e_score_correction_bias = None
481
482
        if _is_cpu and _is_cpu_amx_available:
            self.quant_method = PackWeightMethod(weight_names=["weight"])
Ke Bao's avatar
Ke Bao committed
483

484
    def forward(self, hidden_states, gemm_output_zero_allocator: BumpAllocator = None):
485
        if use_intel_amx_backend(self):
486
487
488
489
490
491
492
            return torch.ops.sgl_kernel.weight_packed_linear(
                hidden_states,
                self.weight,
                None,  # bias
                True,  # is_vnni
            )

493
        # NOTE: For some unknown reason, router_gemm seems degrade accept length.
494
        if (
495
            _is_cuda
496
            and hidden_states.shape[0] <= 16
497
            and hidden_states.shape[1] == 7168
498
            and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
499
500
            and _device_sm >= 90
        ):
501
            # router gemm output float32
502
503
504
            logits = dsv3_router_gemm(
                hidden_states, self.weight, out_dtype=torch.float32
            )
505
506
507
508
        elif _use_aiter_gfx95 and hidden_states.shape[0] <= 256:
            logits = aiter_dsv3_router_gemm(
                hidden_states, self.weight, gemm_output_zero_allocator
            )
509
510
511
        else:
            logits = F.linear(hidden_states, self.weight, None)

Ke Bao's avatar
Ke Bao committed
512
513
514
        return logits


Liangsheng Yin's avatar
Liangsheng Yin committed
515
516
517
518
519
class DeepseekV2MoE(nn.Module):

    def __init__(
        self,
        config: PretrainedConfig,
fzyzcjy's avatar
fzyzcjy committed
520
        layer_id: int,
Liangsheng Yin's avatar
Liangsheng Yin committed
521
        quant_config: Optional[QuantizationConfig] = None,
522
        prefix: str = "",
523
        alt_stream: Optional[torch.cuda.Stream] = None,
524
        is_nextn: bool = False,
Liangsheng Yin's avatar
Liangsheng Yin committed
525
526
527
528
529
    ):
        super().__init__()
        self.tp_size = get_tensor_model_parallel_world_size()
        self.routed_scaling_factor = config.routed_scaling_factor
        self.n_shared_experts = config.n_shared_experts
530
531
        self.num_fused_shared_experts = (
            0
532
            if get_global_server_args().disable_shared_experts_fusion
533
534
            else config.n_shared_experts
        )
535
        self.config = config
fzyzcjy's avatar
fzyzcjy committed
536
        self.layer_id = layer_id
537
        self.alt_stream = alt_stream
538
        self.is_nextn = is_nextn
539

Liangsheng Yin's avatar
Liangsheng Yin committed
540
541
542
543
544
545
546
547
548
549
550
551
        if self.tp_size > config.n_routed_experts:
            raise ValueError(
                f"Tensor parallel size {self.tp_size} is greater than "
                f"the number of experts {config.n_routed_experts}."
            )

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

552
        self.gate = MoEGate(
553
554
555
556
            config=config,
            quant_config=quant_config,
            prefix=add_prefix("gate", prefix),
            is_nextn=is_nextn,
557
        )
Ke Bao's avatar
Ke Bao committed
558

559
        self.experts = get_moe_impl_class(quant_config)(
560
            num_experts=config.n_routed_experts
561
            + self.num_fused_shared_experts
562
            + get_global_server_args().ep_num_redundant_experts,
Cheng Wan's avatar
Cheng Wan committed
563
            num_fused_shared_experts=self.num_fused_shared_experts,
564
            top_k=config.num_experts_per_tok + self.num_fused_shared_experts,
565
566
            hidden_size=config.hidden_size,
            intermediate_size=config.moe_intermediate_size,
fzyzcjy's avatar
fzyzcjy committed
567
            layer_id=self.layer_id,
568
            quant_config=quant_config,
569
            routed_scaling_factor=self.routed_scaling_factor,
570
571
            prefix=add_prefix("experts", prefix),
        )
Liangsheng Yin's avatar
Liangsheng Yin committed
572

573
574
575
576
577
578
579
        self.topk = TopK(
            top_k=config.num_experts_per_tok + self.num_fused_shared_experts,
            renormalize=config.norm_topk_prob,
            use_grouped_topk=True,
            num_expert_group=config.n_group,
            num_fused_shared_experts=self.num_fused_shared_experts,
            topk_group=config.topk_group,
580
581
            correction_bias=self.gate.e_score_correction_bias,
            quant_config=quant_config,
582
            routed_scaling_factor=self.routed_scaling_factor,
fzyzcjy's avatar
fzyzcjy committed
583
            apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk,
584
585
586
            # Some Fp4 MoE backends require the output format to be bypassed but the MTP layers are unquantized
            # and requires the output format to be standard. We use quant_config to determine the output format.
            output_format=TopKOutputFormat.STANDARD if quant_config is None else None,
587
588
        )

589
590
591
        self.shared_experts_is_int8 = False
        self.shared_experts_is_fp8 = False
        self.shared_experts_weight_block_size = None
592
        if config.n_shared_experts is not None and self.num_fused_shared_experts == 0:
Liangsheng Yin's avatar
Liangsheng Yin committed
593
            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
594
            # disable tp for shared experts when enable deepep moe, or with fp4 allgather
595
596
597
598
599
600
601
602
603
            self.shared_experts = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
                reduce_results=False,
                prefix=add_prefix("shared_experts", prefix),
                **(
                    dict(tp_rank=0, tp_size=1)
604
                    if get_moe_a2a_backend().is_deepep()
605
                    or get_moe_a2a_backend().is_mooncake()
606
                    or should_use_flashinfer_cutlass_moe_fp4_allgather()
607
608
609
                    else {}
                ),
            )
AniZpZ's avatar
AniZpZ committed
610
611
612
613
            is_packed_weight = hasattr(
                self.shared_experts.gate_up_proj.quant_method, "quant_config"
            ) and self.shared_experts.gate_up_proj.quant_method.quant_config.get_name() in {
                "awq",
614
                "awq_marlin",
AniZpZ's avatar
AniZpZ committed
615
616
                "moe_wna16",
            }
617
            self.shared_experts_is_int8 = (
618
619
                not is_packed_weight
                and self.shared_experts.gate_up_proj.weight.dtype == torch.int8
620
621
            )
            self.shared_experts_is_fp8 = (
622
623
                not is_packed_weight
                and self.shared_experts.gate_up_proj.weight.dtype == torch.float8_e4m3fn
624
625
626
627
628
629
630
631
632
            )
            if self.shared_experts_is_fp8:
                assert (
                    self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size
                    == self.shared_experts.down_proj.quant_method.quant_config.weight_block_size
                )
                self.shared_experts_weight_block_size = (
                    self.shared_experts.gate_up_proj.quant_method.quant_config.weight_block_size
                )
633

634
635
        self.top_k = config.num_experts_per_tok

636
        if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
637
            # TODO: we will support tp < ep in the future
638
            self.ep_size = get_moe_expert_parallel_world_size()
639
640
            self.num_experts = (
                config.n_routed_experts
641
                + get_global_server_args().ep_num_redundant_experts
642
            )
643
644
645
646
647
648
649
650
651
            self.renormalize = config.norm_topk_prob
            self.topk_group = config.topk_group
            self.num_expert_group = config.n_group
            self.correction_bias = (
                self.gate.e_score_correction_bias.data
                if self.gate.e_score_correction_bias is not None
                else None
            )

652
            self.deepep_dispatcher = MaybeTboDeepEPDispatcher(
653
654
655
                group=parallel_state.get_tp_group().device_group,
                router_topk=self.top_k,
                permute_fusion=True,
656
                num_experts=self.num_experts,
657
                num_local_experts=config.n_routed_experts // self.tp_size,
Liangsheng Yin's avatar
Liangsheng Yin committed
658
                hidden_size=config.hidden_size,
659
                params_dtype=config.torch_dtype,
660
                deepep_mode=get_deepep_mode(),
661
                async_finish=True,
662
                return_recv_hook=True,
Liangsheng Yin's avatar
Liangsheng Yin committed
663
664
            )

665
666
667
        self._enable_a2a_moe = (
            get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake()
        )
668
        self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo()
669

670
671
672
673
674
675
676
    def get_moe_weights(self):
        return [
            x.data
            for name, x in self.experts.named_parameters()
            if name not in ["correction_bias"]
        ]

677
    def forward(
678
679
680
        self,
        hidden_states: torch.Tensor,
        forward_batch: Optional[ForwardBatch] = None,
681
        should_allreduce_fusion: bool = False,
682
        use_reduce_scatter: bool = False,
683
        gemm_output_zero_allocator: BumpAllocator = None,
684
    ) -> torch.Tensor:
685
        if not self._enable_a2a_moe:
686
687
688
689
            DUAL_STREAM_TOKEN_THRESHOLD = 1024
            if (
                self.alt_stream is not None
                and self.num_fused_shared_experts == 0
690
                and hidden_states.shape[0] > 0
691
692
                and hidden_states.shape[0] <= DUAL_STREAM_TOKEN_THRESHOLD
            ):
693
                return self.forward_normal_dual_stream(
694
695
696
                    hidden_states,
                    should_allreduce_fusion,
                    use_reduce_scatter,
697
                    gemm_output_zero_allocator,
698
                )
699
            else:
700
                return self.forward_normal(
701
702
703
                    hidden_states,
                    should_allreduce_fusion,
                    use_reduce_scatter,
704
                    gemm_output_zero_allocator,
705
                )
706
707
708
        else:
            return self.forward_deepep(hidden_states, forward_batch)

709
    def forward_normal_dual_stream(
710
711
        self,
        hidden_states: torch.Tensor,
712
        should_allreduce_fusion: bool = False,
713
        use_reduce_scatter: bool = False,
714
        gemm_output_zero_allocator: BumpAllocator = None,
715
    ) -> torch.Tensor:
716

717
718
        current_stream = torch.cuda.current_stream()
        self.alt_stream.wait_stream(current_stream)
719
720
721
        shared_output = self._forward_shared_experts(
            hidden_states, gemm_output_zero_allocator
        )
722

723
        with torch.cuda.stream(self.alt_stream):
724
            # router_logits: (num_tokens, n_experts)
725
            router_logits = self.gate(hidden_states, gemm_output_zero_allocator)
Cheng Wan's avatar
Cheng Wan committed
726
727
            topk_output = self.topk(hidden_states, router_logits)
            final_hidden_states = self.experts(hidden_states, topk_output)
728
729
            if not _is_cuda:
                final_hidden_states *= self.routed_scaling_factor
Cheng Wan's avatar
Cheng Wan committed
730

731
        current_stream.wait_stream(self.alt_stream)
732
733
        with use_symmetric_memory(parallel_state.get_tp_group()) as sm:
            final_hidden_states_out = torch.empty_like(final_hidden_states)
Cheng Wan's avatar
Cheng Wan committed
734

735
736
737
        torch.add(final_hidden_states, shared_output, out=final_hidden_states_out)
        final_hidden_states = final_hidden_states_out
        sm.tag(final_hidden_states)
738
739
740
741
742
743
        if (
            self.tp_size > 1
            and not should_allreduce_fusion
            and not use_reduce_scatter
            and not should_use_flashinfer_cutlass_moe_fp4_allgather()
        ):
744
745
746
            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
        return final_hidden_states

747
    def forward_normal(
748
749
        self,
        hidden_states: torch.Tensor,
750
        should_allreduce_fusion: bool = False,
751
        use_reduce_scatter: bool = False,
752
        gemm_output_zero_allocator: BumpAllocator = None,
753
    ) -> torch.Tensor:
754
755
        if hasattr(self, "shared_experts") and use_intel_amx_backend(
            self.shared_experts.gate_up_proj
756
        ):
757
            return self.forward_cpu(hidden_states, should_allreduce_fusion)
758

759
        if hidden_states.shape[0] > 0:
760
761
762
763
            if not self._fuse_shared_experts_inside_sbo:
                shared_output = self._forward_shared_experts(
                    hidden_states, gemm_output_zero_allocator
                )
764
            # router_logits: (num_tokens, n_experts)
765
            router_logits = self.gate(hidden_states, gemm_output_zero_allocator)
766
767
768
769
            topk_output = self.topk(hidden_states, router_logits)
        else:
            shared_output = None
            topk_output = self.topk.empty_topk_output(hidden_states.device)
770

771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
        if self._fuse_shared_experts_inside_sbo:
            shared_output = None

            def _forward_shared_experts_and_put_results():
                nonlocal shared_output
                shared_output = self._forward_shared_experts(
                    hidden_states, gemm_output_zero_allocator
                )

        final_hidden_states = self.experts(
            hidden_states,
            topk_output,
            **(
                dict(
                    forward_shared_experts=_forward_shared_experts_and_put_results,
                    alt_stream=self.alt_stream,
                )
                if self._fuse_shared_experts_inside_sbo
                else {}
            ),
        )
792
793
        if not _is_cuda and not _use_aiter:
            # fused in biased_grouped_topk so we can skip here
794
            final_hidden_states *= self.routed_scaling_factor
795
        if shared_output is not None:
796
797
798
799
800
            with use_symmetric_memory(parallel_state.get_tp_group()) as sm:
                final_hidden_states_out = torch.empty_like(final_hidden_states)
            torch.add(final_hidden_states, shared_output, out=final_hidden_states_out)
            final_hidden_states = final_hidden_states_out
            sm.tag(final_hidden_states)
801
802
803
804
805
806
        if (
            self.tp_size > 1
            and not should_allreduce_fusion
            and not use_reduce_scatter
            and not should_use_flashinfer_cutlass_moe_fp4_allgather()
        ):
807
808
809
            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
        return final_hidden_states

810
    def forward_cpu(
811
812
813
        self,
        hidden_states: torch.Tensor,
        should_allreduce_fusion: bool = False,
814
    ) -> torch.Tensor:
815
816
        # router_logits: (num_tokens, n_experts)
        router_logits = self.gate(hidden_states)
817
        topk_output = self.topk(hidden_states, router_logits)
818
        fused_experts_out = self.experts(
819
            hidden_states=hidden_states, topk_output=topk_output
820
821
        )

822
823
824
        assert use_intel_amx_backend(
            self.shared_experts.gate_up_proj
        ) == use_intel_amx_backend(self.shared_experts.down_proj)
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
        # [Note] inplace should be False in fused_experts.
        # If inplace is True in fused_experts (self.experts), hidden_states will be changed after fused_experts
        # While hidden_states is still needed in shared_expert.
        final_hidden_states = torch.ops.sgl_kernel.shared_expert_cpu(
            hidden_states,
            self.shared_experts.gate_up_proj.weight,
            self.shared_experts.down_proj.weight,
            fused_experts_out,
            self.routed_scaling_factor,
            True,  # inplace
            self.shared_experts_is_int8,  # use_int8_w8a8
            self.shared_experts_is_fp8,  # use_fp8_w8a16
            (
                self.shared_experts.gate_up_proj.weight_scale
                if self.shared_experts_is_int8
                else (
                    self.shared_experts.gate_up_proj.weight_scale_inv
                    if self.shared_experts_is_fp8
                    else None
                )
            ),  # w1_scale
            (
                self.shared_experts.down_proj.weight_scale
                if self.shared_experts_is_int8
                else (
                    self.shared_experts.down_proj.weight_scale_inv
                    if self.shared_experts_is_fp8
                    else None
                )
            ),  # w2_scale
            (
                self.shared_experts_weight_block_size
                if self.shared_experts_is_fp8
                else None
            ),  # block_size
            None,  # a1_scale
            None,  # a2_scale
            True,  # is_vnni
        )
864
        if self.tp_size > 1 and not should_allreduce_fusion:
865
866
867
            final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
        return final_hidden_states

868
869
870
871
    def forward_deepep(
        self, hidden_states: torch.Tensor, forward_batch: ForwardBatch
    ) -> torch.Tensor:
        shared_output = None
Cheng Wan's avatar
Cheng Wan committed
872
        if hidden_states.shape[0] > 0:
873
874
            # router_logits: (num_tokens, n_experts)
            router_logits = self.gate(hidden_states)
875
            if not self._fuse_shared_experts_inside_sbo:
876
                shared_output = self._forward_shared_experts(hidden_states)
877
878
879
            topk_weights, topk_idx, _ = self.topk(
                hidden_states,
                router_logits,
880
                num_token_non_padded=forward_batch.num_token_non_padded,
881
882
883
                expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
                    layer_id=self.layer_id,
                ),
884
885
            )
        else:
886
887
            topk_weights, topk_idx, _ = self.topk.empty_topk_output(
                hidden_states.device
888
            )
889

890
891
892
893
894
895
896
897
        if self._fuse_shared_experts_inside_sbo:
            shared_output = None

            def _forward_shared_experts_and_put_results():
                nonlocal shared_output
                shared_output = self._forward_shared_experts(hidden_states)

        final_hidden_states = self.experts(
898
899
900
            hidden_states=hidden_states,
            topk_idx=topk_idx,
            topk_weights=topk_weights,
901
            forward_batch=forward_batch,
902
903
904
905
            **(
                dict(
                    forward_shared_experts=_forward_shared_experts_and_put_results,
                    alt_stream=self.alt_stream,
906
907
                    # SBO is not yet implemented for NextN
                    disable_sbo=self.is_nextn,
908
909
910
911
                )
                if self._fuse_shared_experts_inside_sbo
                else {}
            ),
912
913
914
        )

        if shared_output is not None:
915
            x = shared_output
fzyzcjy's avatar
fzyzcjy committed
916
            if self.experts.should_fuse_routed_scaling_factor_in_topk:
917
918
919
                x.add_(final_hidden_states)
            else:
                x.add_(final_hidden_states, alpha=self.routed_scaling_factor)
920
921
            final_hidden_states = x
        else:
fzyzcjy's avatar
fzyzcjy committed
922
            if not self.experts.should_fuse_routed_scaling_factor_in_topk:
923
                final_hidden_states *= self.routed_scaling_factor
924
925
926

        return final_hidden_states

927
928
929
    def _forward_shared_experts(
        self, hidden_states, gemm_output_zero_allocator: BumpAllocator = None
    ):
930
        if (hidden_states.shape[0] > 0) and (self.num_fused_shared_experts == 0):
931
932
933
            return self.shared_experts(
                hidden_states, gemm_output_zero_allocator=gemm_output_zero_allocator
            )
934
935
936
        else:
            return None

937
    def op_gate(self, state):
938
        if is_non_idle_and_non_empty(
939
            state.forward_batch.forward_mode, state.hidden_states_mlp_input
940
        ):
941
            # router_logits: (num_tokens, n_experts)
942
            state.router_logits = self.gate(state.hidden_states_mlp_input)
943
        else:
944
            state.router_logits = None
945

946
    def op_shared_experts(self, state):
947
        hidden_states_mlp_input = state.pop("hidden_states_mlp_input")
948
        if (self.num_fused_shared_experts == 0) and is_non_idle_and_non_empty(
949
            state.forward_batch.forward_mode, hidden_states_mlp_input
950
        ):
951
            state.shared_output = self.shared_experts(hidden_states_mlp_input)
952
        else:
953
            state.shared_output = None
954

955
    def op_select_experts(self, state):
956
        router_logits = state.pop("router_logits")
957
958
        hidden_states = state.hidden_states_mlp_input

959
        if router_logits is not None:
960
961
962
            with get_global_expert_distribution_recorder().with_current_layer(
                self.layer_id
            ):
963
                state.topk_weights_local, state.topk_idx_local, _ = self.topk(
964
965
966
967
968
969
970
                    hidden_states=hidden_states,
                    router_logits=router_logits,
                    num_token_non_padded=state.forward_batch.num_token_non_padded,
                    expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new(
                        layer_id=self.layer_id,
                    ),
                )
971
972
973
974
975
976
977
        else:
            state.topk_idx_local = torch.full(
                (0, self.top_k), -1, dtype=torch.int, device=hidden_states.device
            )
            state.topk_weights_local = torch.empty(
                (0, self.top_k), dtype=torch.float32, device=hidden_states.device
            )
978

979
    def op_dispatch_a(self, state):
980
        if self.ep_size > 1:
981
            self.experts.deepep_dispatcher.dispatch_a(
982
                hidden_states=state.hidden_states_mlp_input,
983
                input_global_scale=None,
984
985
                topk_idx=state.pop("topk_idx_local"),
                topk_weights=state.pop("topk_weights_local"),
986
                forward_batch=state.forward_batch,
987
                tbo_subbatch_index=state.get("tbo_subbatch_index"),
988
            )
989

990
    def op_dispatch_b(self, state):
991
992
993
994
        if self.ep_size > 1:
            with get_global_expert_distribution_recorder().with_current_layer(
                self.layer_id
            ):
995
                state.dispatch_output = self.experts.deepep_dispatcher.dispatch_b(
996
997
                    tbo_subbatch_index=state.get("tbo_subbatch_index"),
                )
998
999

    def op_experts(self, state):
1000
1001
        state.hidden_states_experts_output = self.experts.moe_impl(
            dispatch_output=state.dispatch_output,
1002
        )
1003

1004
    def op_combine_a(self, state):
1005
        if self.ep_size > 1:
1006
            self.experts.deepep_dispatcher.combine_a(
1007
                hidden_states=state.pop("hidden_states_experts_output"),
1008
1009
                topk_idx=state.dispatch_output.topk_idx,
                topk_weights=state.dispatch_output.topk_weights,
1010
                forward_batch=state.forward_batch,
1011
                tbo_subbatch_index=state.get("tbo_subbatch_index"),
1012
            )
1013
            state.pop("dispatch_output")
1014

1015
    def op_combine_b(self, state):
1016
        if self.ep_size > 1:
1017
1018
1019
1020
            state.hidden_states_after_combine = (
                self.experts.deepep_dispatcher.combine_b(
                    tbo_subbatch_index=state.get("tbo_subbatch_index"),
                )
1021
            )
1022
1023

    def op_output(self, state):
1024
        final_hidden_states = state.pop("hidden_states_after_combine")
1025
1026
1027
1028
1029
1030
1031

        if (shared_output := state.pop("shared_output")) is not None:
            x = shared_output
            x.add_(final_hidden_states, alpha=self.routed_scaling_factor)
            final_hidden_states = x
        else:
            final_hidden_states *= self.routed_scaling_factor
Liangsheng Yin's avatar
Liangsheng Yin committed
1032

1033
        state.hidden_states_mlp_output = final_hidden_states
1034

Liangsheng Yin's avatar
Liangsheng Yin committed
1035
1036
1037
1038
1039
1040
1041
1042
1043

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

    if scale <= 1:
        return 1.0
    return 0.1 * mscale * math.log(scale) + 1.0


1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
class DeepseekV2AttentionMLA(nn.Module):

    def __init__(
        self,
        config: PretrainedConfig,
        hidden_size: int,
        num_heads: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        v_head_dim: int,
        q_lora_rank: int,
        kv_lora_rank: int,
        rope_theta: float = 10000,
        rope_scaling: Optional[Dict[str, Any]] = None,
        max_position_embeddings: int = 8192,
        quant_config: Optional[QuantizationConfig] = None,
Lianmin Zheng's avatar
Lianmin Zheng committed
1060
1061
        reduce_results: bool = True,
        layer_id: int = None,
1062
        prefix: str = "",
1063
        alt_stream: Optional[torch.cuda.Stream] = None,
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
    ) -> None:
        super().__init__()
        self.layer_id = layer_id
        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
Lianmin Zheng's avatar
Lianmin Zheng committed
1074
1075
1076
        attn_tp_rank = get_attention_tp_rank()
        attn_tp_size = get_attention_tp_size()

1077
        self.num_heads = num_heads
Lianmin Zheng's avatar
Lianmin Zheng committed
1078
1079
        assert num_heads % attn_tp_size == 0
        self.num_local_heads = num_heads // attn_tp_size
1080
1081
1082
1083
        self.scaling = self.qk_head_dim**-0.5
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings

fzyzcjy's avatar
fzyzcjy committed
1084
1085
1086
1087
        # NOTE modification to rope_scaling must be done early enough, b/c e.g. Indexer needs it
        if rope_scaling:
            rope_scaling["rope_type"] = "deepseek_yarn"

Lianmin Zheng's avatar
Lianmin Zheng committed
1088
1089
        # For tensor parallel attention
        if self.q_lora_rank is not None:
1090
            self.fused_qkv_a_proj_with_mqa = ReplicatedLinear(
Ke Bao's avatar
Ke Bao committed
1091
                self.hidden_size,
1092
                self.q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim,
1093
1094
                bias=False,
                quant_config=quant_config,
1095
                prefix=add_prefix("fused_qkv_a_proj_with_mqa", prefix),
1096
            )
Lianmin Zheng's avatar
Lianmin Zheng committed
1097
1098
1099
1100
            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,
Ke Bao's avatar
Ke Bao committed
1101
                bias=False,
1102
                quant_config=self._get_q_b_proj_quant_config(quant_config),
Lianmin Zheng's avatar
Lianmin Zheng committed
1103
1104
1105
                prefix=add_prefix("q_b_proj", prefix),
                tp_rank=attn_tp_rank,
                tp_size=attn_tp_size,
Ke Bao's avatar
Ke Bao committed
1106
            )
Lianmin Zheng's avatar
Lianmin Zheng committed
1107
1108
        else:
            self.q_proj = ColumnParallelLinear(
1109
                self.hidden_size,
Lianmin Zheng's avatar
Lianmin Zheng committed
1110
                self.num_heads * self.qk_head_dim,
1111
1112
                bias=False,
                quant_config=quant_config,
Lianmin Zheng's avatar
Lianmin Zheng committed
1113
1114
1115
                prefix=add_prefix("q_proj", prefix),
                tp_rank=attn_tp_rank,
                tp_size=attn_tp_size,
1116
            )
1117
1118
1119
1120
1121
1122
1123
1124
            self.kv_a_proj_with_mqa = ReplicatedLinear(
                self.hidden_size,
                self.kv_lora_rank + self.qk_rope_head_dim,
                bias=False,
                quant_config=quant_config,
                prefix=add_prefix("kv_a_proj_with_mqa", prefix),
            )

fzyzcjy's avatar
fzyzcjy committed
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
        self.use_nsa = is_deepseek_nsa(config)
        if self.use_nsa:
            self.indexer = Indexer(
                hidden_size=hidden_size,
                index_n_heads=get_nsa_index_n_heads(config),
                index_head_dim=get_nsa_index_head_dim(config),
                rope_head_dim=qk_rope_head_dim,
                index_topk=get_nsa_index_topk(config),
                q_lora_rank=q_lora_rank,
                max_position_embeddings=max_position_embeddings,
                rope_theta=rope_theta,
                scale_fmt="ue8m0",
                block_size=128,
                rope_scaling=rope_scaling,
                prefix=add_prefix("indexer", prefix),
                quant_config=quant_config,
                layer_id=layer_id,
                alt_stream=alt_stream,
            )

Lianmin Zheng's avatar
Lianmin Zheng committed
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
            quant_config=quant_config,
            prefix=add_prefix("kv_b_proj", prefix),
            tp_rank=attn_tp_rank,
            tp_size=attn_tp_size,
        )
        # O projection.
        self.o_proj = RowParallelLinear(
            self.num_heads * self.v_head_dim,
            self.hidden_size,
            bias=False,
            quant_config=quant_config,
            reduce_results=reduce_results,
            prefix=add_prefix("o_proj", prefix),
            tp_rank=attn_tp_rank,
            tp_size=attn_tp_size,
        )
1165
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
Ke Bao's avatar
Ke Bao committed
1166

1167
        self.rotary_emb = get_rope_wrapper(
1168
1169
1170
1171
1172
1173
            qk_rope_head_dim,
            rotary_dim=qk_rope_head_dim,
            max_position=max_position_embeddings,
            base=rope_theta,
            rope_scaling=rope_scaling,
            is_neox_style=False,
1174
            device=get_global_server_args().device,
1175
1176
1177
1178
1179
1180
1181
        )

        if rope_scaling:
            mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
            scaling_factor = rope_scaling["factor"]
            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
            self.scaling = self.scaling * mscale * mscale
Ke Bao's avatar
Ke Bao committed
1182
1183
        else:
            self.rotary_emb.forward = self.rotary_emb.forward_native
1184

1185
        self.attn_mqa = RadixAttention(
1186
1187
1188
1189
1190
1191
            self.num_local_heads,
            self.kv_lora_rank + self.qk_rope_head_dim,
            self.scaling,
            num_kv_heads=1,
            layer_id=layer_id,
            v_head_dim=self.kv_lora_rank,
1192
            quant_config=quant_config,
1193
            prefix=add_prefix("attn_mqa", prefix),
1194
1195
        )

1196
1197
1198
1199
1200
1201
1202
        self.attn_mha = RadixAttention(
            self.num_local_heads,
            self.qk_nope_head_dim + self.qk_rope_head_dim,
            self.scaling,
            num_kv_heads=self.num_local_heads,
            layer_id=layer_id,
            v_head_dim=self.v_head_dim,
1203
            quant_config=quant_config,
1204
            prefix=add_prefix("attn_mha", prefix),
1205
1206
        )

1207
        self.alt_stream = alt_stream
1208
        self.attn_mha.kv_b_proj = None
1209

Ke Bao's avatar
Ke Bao committed
1210
1211
        self.w_kc = None
        self.w_vc = None
1212
        self.w_scale = 1.0
1213

1214
1215
1216
1217
        self.w_scale_k = None
        self.w_scale_v = None
        self.use_deep_gemm_bmm = False

1218
1219
1220
1221
1222
1223
        self.flashinfer_mla_disable_ragged = (
            get_global_server_args().flashinfer_mla_disable_ragged
        )
        self.disable_chunked_prefix_cache = (
            get_global_server_args().disable_chunked_prefix_cache
        )
1224
1225
1226
1227

        self.current_attention_backend = (
            None  # Attention backend used by current forward batch
        )
1228
1229
1230
        self.rocm_fused_decode_mla = get_bool_env_var(
            "SGLANG_ROCM_FUSED_DECODE_MLA", "false"
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
1231

1232
        # TODO: Design a finer way to determine the threshold
1233
1234
1235
        self.chunked_prefix_cache_threshold = get_int_env_var(
            "SGL_CHUNKED_PREFIX_CACHE_THRESHOLD", 8192
        )
1236

1237
1238
1239
        # If we have self.fused_qkv_a_proj_with_mqa and we're running on CPU, we will choose the torch.ops.sgl_kernel.qkv_proj_with_rope_fused_weight kernel
        # which requires self.w_kc and self.w_vc to be packed.
        # If not, we will use torch.bmm and weight shouldn't be packed in this case
AniZpZ's avatar
AniZpZ committed
1240
1241
        has_fused_proj = hasattr(self, "fused_qkv_a_proj_with_mqa")
        if has_fused_proj and _is_cpu and _is_cpu_amx_available:
1242
1243
1244
1245
            self.quant_method = PackWeightMethod(
                weight_names=["w_kc", "w_vc"], transpose_dims=[[1, 2], [1, 2]]
            )

1246
        is_packed_weight = (
AniZpZ's avatar
AniZpZ committed
1247
1248
1249
            has_fused_proj
            and hasattr(self.fused_qkv_a_proj_with_mqa.quant_method, "quant_config")
            and self.fused_qkv_a_proj_with_mqa.quant_method.quant_config.get_name()
1250
            in {"awq", "awq_marlin", "moe_wna16"}
1251
        )
1252
        self.use_min_latency_fused_a_gemm = (
AniZpZ's avatar
AniZpZ committed
1253
            has_fused_proj
1254
            and not is_packed_weight
1255
1256
1257
            and self.fused_qkv_a_proj_with_mqa.weight.dtype == torch.bfloat16
            and self.fused_qkv_a_proj_with_mqa.weight.shape[0] == 2112
            and self.fused_qkv_a_proj_with_mqa.weight.shape[1] == 7168
1258
            and _is_cuda
1259
            and _device_sm >= 90
1260
1261
        )

1262
        self.qkv_proj_with_rope_is_int8 = (
AniZpZ's avatar
AniZpZ committed
1263
            has_fused_proj
1264
            and not is_packed_weight
1265
1266
1267
            and self.fused_qkv_a_proj_with_mqa.weight.dtype == torch.int8
        )
        self.qkv_proj_with_rope_is_fp8 = (
AniZpZ's avatar
AniZpZ committed
1268
            has_fused_proj
1269
            and not is_packed_weight
1270
1271
1272
1273
            and self.fused_qkv_a_proj_with_mqa.weight.dtype == torch.float8_e4m3fn
        )

        self.weight_block_size = None
1274
1275
1276
1277
1278
1279
        if self.qkv_proj_with_rope_is_fp8 and _is_cpu and _is_cpu_amx_available:
            assert getattr(
                self.fused_qkv_a_proj_with_mqa.quant_method, "block_quant", False
            ) == getattr(self.q_b_proj.quant_method, "block_quant", False)
            use_block_quant = getattr(
                self.fused_qkv_a_proj_with_mqa.quant_method, "block_quant", False
1280
1281
            )

1282
1283
1284
1285
1286
1287
1288
1289
            if use_block_quant:
                assert (
                    self.fused_qkv_a_proj_with_mqa.quant_method.quant_config.weight_block_size
                    == self.q_b_proj.quant_method.quant_config.weight_block_size
                )
                self.weight_block_size = (
                    self.fused_qkv_a_proj_with_mqa.quant_method.quant_config.weight_block_size
                )
1290
1291
1292
        self.is_mla_preprocess_enabled = is_mla_preprocess_enabled()
        if self.is_mla_preprocess_enabled:
            assert (
fzyzcjy's avatar
fzyzcjy committed
1293
1294
                quant_config is None or quant_config.get_name() == "w8a8_int8"
            ), "MLA Preprocess only works with Unquant or W8A8Int8"
1295
            self.mla_preprocess = None
1296

1297
1298
1299
    def dispatch_attn_forward_method(
        self, forward_batch: ForwardBatch
    ) -> AttnForwardMethod:
1300
1301
        # Determine attention backend used by current forward batch
        if forward_batch.forward_mode.is_decode_or_idle():
1302
            attention_backend = get_global_server_args().decode_attention_backend
1303
1304
1305
1306
1307
        elif (
            forward_batch.forward_mode.is_target_verify()
            or forward_batch.forward_mode.is_draft_extend()
        ):
            # Use the specified backend for speculative operations (both verify and draft extend)
1308
1309
            if get_global_server_args().speculative_attention_mode == "decode":
                attention_backend = get_global_server_args().decode_attention_backend
1310
            else:  # default to prefill
1311
                attention_backend = get_global_server_args().prefill_attention_backend
1312
        else:
1313
            attention_backend = get_global_server_args().prefill_attention_backend
1314
1315
        self.current_attention_backend = attention_backend

fzyzcjy's avatar
fzyzcjy committed
1316
        handler = AttentionBackendRegistry.get_handler(attention_backend)
1317
        return handler(self, forward_batch)
Lianmin Zheng's avatar
Lianmin Zheng committed
1318

1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
    def op_prepare(self, state):
        state.attn_intermediate_state = self.forward_prepare(
            positions=state.positions,
            hidden_states=state.pop("hidden_states_after_comm_pre_attn"),
            forward_batch=state.forward_batch,
            zero_allocator=state.zero_allocator,
        )

    def op_core(self, state):
        state.hidden_states_after_attn = self.forward_core(
            state.pop("attn_intermediate_state")
        )

1332
1333
1334
1335
    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
1336
        forward_batch: ForwardBatch,
1337
        zero_allocator: BumpAllocator,
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
    ):
        s = self.forward_prepare(
            positions=positions,
            hidden_states=hidden_states,
            forward_batch=forward_batch,
            zero_allocator=zero_allocator,
        )
        return self.forward_core(s)

    def forward_prepare(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        zero_allocator: BumpAllocator,
    ):
1354
1355
1356
        if self.attn_mha.kv_b_proj is None:
            self.attn_mha.kv_b_proj = self.kv_b_proj

1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
        # when hidden_states is a tuple of tensors, the tuple will include quantized weight and scale tensor
        if isinstance(hidden_states, tuple):
            if hidden_states[0].shape[0] == 0:
                assert (
                    not self.o_proj.reduce_results
                ), "short-circuiting allreduce will lead to hangs"
                return hidden_states[0]
        else:
            if hidden_states.shape[0] == 0:
                assert (
                    not self.o_proj.reduce_results
                ), "short-circuiting allreduce will lead to hangs"
                return hidden_states, None, forward_batch, None
1370

1371
1372
        attn_forward_method = self.dispatch_attn_forward_method(forward_batch)
        if attn_forward_method == AttnForwardMethod.MHA:
1373
1374
1375
            inner_state = self.forward_normal_prepare(
                positions, hidden_states, forward_batch, zero_allocator
            )
1376
        elif attn_forward_method == AttnForwardMethod.MHA_CHUNKED_KV:
1377
1378
            inner_state = self.forward_normal_chunked_kv_prepare(
                positions, hidden_states, forward_batch, zero_allocator
1379
            )
1380
        elif attn_forward_method == AttnForwardMethod.MLA:
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
            if not self.is_mla_preprocess_enabled:
                inner_state = self.forward_absorb_prepare(
                    positions, hidden_states, forward_batch, zero_allocator
                )
            else:
                # TODO(iforgetmyname): to be separated as a standalone func
                if self.mla_preprocess is None:
                    self.mla_preprocess = NPUFusedMLAPreprocess(
                        self.fused_qkv_a_proj_with_mqa,
                        self.q_a_layernorm,
                        self.kv_a_layernorm,
                        self.q_b_proj,
                        self.w_kc,
                        self.rotary_emb,
                        self.layer_id,
                        self.num_local_heads,
                        self.qk_nope_head_dim,
                        self.qk_rope_head_dim,
                    )
                inner_state = self.mla_preprocess.forward(
                    positions, hidden_states, forward_batch, zero_allocator
                )
1403
                inner_state = (*inner_state, None)  # add a position for topk_indices
fzyzcjy's avatar
fzyzcjy committed
1404
1405
1406
1407
        elif attn_forward_method == AttnForwardMethod.NPU_MLA_SPARSE:
            inner_state = self.forward_npu_sparse_prepare(
                positions, hidden_states, forward_batch, zero_allocator
            )
1408
        elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE:
1409
1410
            inner_state = self.forward_absorb_fused_mla_rope_prepare(
                positions, hidden_states, forward_batch, zero_allocator
1411
            )
1412
1413
1414
1415
        elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_CPU:
            inner_state = self.forward_absorb_fused_mla_rope_cpu_prepare(
                positions, hidden_states, forward_batch, zero_allocator
            )
1416
        else:
1417
            raise NotImplementedError
1418
        return None, attn_forward_method, forward_batch, inner_state
1419

1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
    def forward_core(self, intermediate_state):
        hidden_states, attn_forward_method, forward_batch, inner_state = (
            intermediate_state
        )
        if inner_state is None:
            return hidden_states

        if attn_forward_method == AttnForwardMethod.MHA:
            return self.forward_normal_core(*inner_state)
        elif attn_forward_method == AttnForwardMethod.MHA_CHUNKED_KV:
            return self.forward_normal_chunked_kv_core(*inner_state)
        elif attn_forward_method == AttnForwardMethod.MLA:
            return self.forward_absorb_core(*inner_state)
fzyzcjy's avatar
fzyzcjy committed
1433
1434
        elif attn_forward_method == AttnForwardMethod.NPU_MLA_SPARSE:
            return self.forward_npu_sparse_core(*inner_state)
1435
1436
        elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE:
            return self.forward_absorb_fused_mla_rope_core(*inner_state)
1437
1438
        elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_CPU:
            return self.forward_absorb_fused_mla_rope_cpu_core(*inner_state)
1439
1440
1441
1442
        else:
            raise NotImplementedError

    def forward_normal_prepare(
1443
1444
1445
1446
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
1447
1448
        zero_allocator: BumpAllocator,
    ):
1449
        if self.q_lora_rank is not None:
1450
1451
1452
            q, latent_cache = self.fused_qkv_a_proj_with_mqa(hidden_states)[0].split(
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
            )
1453
1454
1455
1456
1457
1458
            q = self.q_a_layernorm(q)
            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
        else:
            q = self.q_proj(hidden_states)[0].view(
                -1, self.num_local_heads, self.qk_head_dim
            )
1459
1460
            latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]

1461
1462
1463
        _, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
        kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
        latent_cache = latent_cache.unsqueeze(1)
1464
        kv_a = self.kv_a_layernorm(kv_a)
1465
1466
1467
1468
1469
1470
1471
1472
        kv = self.kv_b_proj(kv_a)[0]
        kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
        k_nope = kv[..., : self.qk_nope_head_dim]
        v = kv[..., self.qk_nope_head_dim :]
        k_pe = latent_cache[:, :, self.kv_lora_rank :]
        q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
        q[..., self.qk_nope_head_dim :] = q_pe
        k = torch.empty_like(q)
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484

        # Temporary for DeepSeek V3/R1 only, but can generalize if needed
        if (
            _is_cuda
            and (self.num_local_heads == 128)
            and (self.qk_nope_head_dim == 128)
            and (self.qk_rope_head_dim == 64)
        ):
            concat_mla_k(k=k, k_nope=k_nope, k_rope=k_pe)
        else:
            k[..., : self.qk_nope_head_dim] = k_nope
            k[..., self.qk_nope_head_dim :] = k_pe
1485

1486
1487
1488
        if not _is_npu:
            latent_cache[:, :, : self.kv_lora_rank] = kv_a.unsqueeze(1)
            latent_cache[:, :, self.kv_lora_rank :] = k_pe
1489

1490
1491
1492
1493
1494
1495
1496
1497
1498
            # Save latent cache
            forward_batch.token_to_kv_pool.set_kv_buffer(
                self.attn_mha, forward_batch.out_cache_loc, latent_cache, None
            )
        else:
            # To reduce a time-costing split operation
            forward_batch.token_to_kv_pool.set_kv_buffer(
                self.attn_mha, forward_batch.out_cache_loc, kv_a.unsqueeze(1), k_pe
            )
1499
1500
1501
1502

        return q, k, v, forward_batch

    def forward_normal_core(self, q, k, v, forward_batch):
1503
1504
1505
1506
1507
        attn_output = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False)
        attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim)
        output, _ = self.o_proj(attn_output)
        return output

Faraz's avatar
Faraz committed
1508
1509
1510
1511
1512
1513
    def _fuse_rope_for_trtllm_mla(self, forward_batch: ForwardBatch) -> bool:
        """
        Check if we should skip rope and do fused rope+quantize for TRTLLM MLA decode in fp8_e4m3 path.
        """
        return (
            self.current_attention_backend == "trtllm_mla"
1514
1515
1516
1517
            and (
                forward_batch.forward_mode.is_decode_or_idle()
                or forward_batch.forward_mode.is_target_verify()
            )
Faraz's avatar
Faraz committed
1518
1519
1520
            and forward_batch.attn_backend.data_type == torch.float8_e4m3fn
        )

1521
    def forward_absorb_prepare(
1522
1523
1524
1525
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
1526
        zero_allocator: BumpAllocator,
1527
    ):
1528
        from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
1529

fzyzcjy's avatar
fzyzcjy committed
1530
        q_lora = None
1531
        if self.q_lora_rank is not None:
1532
1533
1534
1535
1536
            if (
                (not isinstance(hidden_states, tuple))
                and hidden_states.shape[0] <= 16
                and self.use_min_latency_fused_a_gemm
            ):
1537
1538
1539
1540
1541
1542
                fused_qkv_a_proj_out = dsv3_fused_a_gemm(
                    hidden_states, self.fused_qkv_a_proj_with_mqa.weight.T
                )
            else:
                fused_qkv_a_proj_out = self.fused_qkv_a_proj_with_mqa(hidden_states)[0]
            q, latent_cache = fused_qkv_a_proj_out.split(
1543
1544
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
            )
1545
1546
1547
            k_nope = latent_cache[..., : self.kv_lora_rank]

            # overlap qk norm
1548
            if self.alt_stream is not None and get_is_capture_mode():
1549
1550
1551
1552
1553
1554
1555
                current_stream = torch.cuda.current_stream()
                self.alt_stream.wait_stream(current_stream)
                q = self.q_a_layernorm(q)
                with torch.cuda.stream(self.alt_stream):
                    k_nope = self.kv_a_layernorm(k_nope)
                current_stream.wait_stream(self.alt_stream)
            else:
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
                if _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
                    q, k_nope = fused_rms_mxfp4_quant(
                        q,
                        self.q_a_layernorm.weight,
                        self.q_a_layernorm.variance_epsilon,
                        k_nope,
                        self.kv_a_layernorm.weight,
                        self.kv_a_layernorm.variance_epsilon,
                    )
                else:
                    q = self.q_a_layernorm(q)
                    k_nope = self.kv_a_layernorm(k_nope)
1568

fzyzcjy's avatar
fzyzcjy committed
1569
1570
1571
1572
            # q_lora needed by indexer
            if self.use_nsa:
                q_lora = q

1573
            k_nope = k_nope.unsqueeze(1)
1574
1575
1576
1577
1578
            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
        else:
            q = self.q_proj(hidden_states)[0].view(
                -1, self.num_local_heads, self.qk_head_dim
            )
1579
            latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
1580
1581
1582
            k_nope = latent_cache[..., : self.kv_lora_rank]
            k_nope = self.kv_a_layernorm(k_nope).unsqueeze(1)

1583
        q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
1584
        k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1)
1585

1586
1587
        if self.use_deep_gemm_bmm:
            q_nope_val, q_nope_scale, masked_m, expected_m, aligned_m = (
1588
                per_token_group_quant_mla_deep_gemm_masked_fp8(q_nope.transpose(0, 1))
1589
1590
1591
1592
            )
            q_nope_out = q_nope.new_empty(
                (self.num_local_heads, aligned_m, self.kv_lora_rank)
            )
1593
            deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
1594
1595
1596
1597
1598
1599
1600
                (q_nope_val, q_nope_scale),
                (self.w_kc, self.w_scale_k),
                q_nope_out,
                masked_m,
                expected_m,
            )
            q_nope_out = q_nope_out[:, :expected_m, :]
1601
1602
        elif _is_hip:
            # TODO(haishaw): add bmm_fp8 to ROCm
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
            if _use_aiter_gfx95 and self.w_kc.dtype == torch.uint8:
                x = q_nope.transpose(0, 1)
                q_nope_out = torch.empty(
                    x.shape[0],
                    x.shape[1],
                    self.w_kc.shape[2],
                    device=x.device,
                    dtype=torch.bfloat16,
                )
                batched_gemm_afp4wfp4_pre_quant(
                    x,
                    self.w_kc.transpose(-2, -1),
                    self.w_scale_k.transpose(-2, -1),
                    torch.bfloat16,
                    q_nope_out,
                )
            else:
                q_nope_out = torch.bmm(
                    q_nope.to(torch.bfloat16).transpose(0, 1),
                    self.w_kc.to(torch.bfloat16) * self.w_scale,
                )
1624
        elif self.w_kc.dtype == torch.float8_e4m3fn:
1625
1626
1627
1628
1629
1630
1631
1632
1633
            # fix bmm_fp8 error under cublas12.9 caused by bumpallocator, detail in pr#11612
            q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8(
                q_nope.transpose(0, 1),
                (
                    torch.zeros((1,), dtype=torch.float32, device=q_nope.device)
                    if _is_cublas_ge_129
                    else zero_allocator.allocate(1)
                ),
            )
1634
1635
1636
1637
1638
            q_nope_out = bmm_fp8(
                q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16
            )
        else:
            q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc)
1639
1640

        q_nope_out = q_nope_out.transpose(0, 1)
Faraz's avatar
Faraz committed
1641

1642
        if not self._fuse_rope_for_trtllm_mla(forward_batch) and (
fzyzcjy's avatar
fzyzcjy committed
1643
            not _use_aiter or not _is_gfx95_supported or self.use_nsa
1644
        ):
Faraz's avatar
Faraz committed
1645
            q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
1646

fzyzcjy's avatar
fzyzcjy committed
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
        topk_indices = None
        if q_lora is not None:
            topk_indices = self.indexer(
                x=hidden_states,
                q_lora=q_lora,
                positions=positions,
                forward_batch=forward_batch,
                layer_id=self.layer_id,
            )

        return (
            q_pe,
            k_pe,
            q_nope_out,
            k_nope,
            forward_batch,
            zero_allocator,
            positions,
            topk_indices,
        )
1667
1668

    def forward_absorb_core(
fzyzcjy's avatar
fzyzcjy committed
1669
1670
1671
1672
1673
1674
1675
1676
1677
        self,
        q_pe,
        k_pe,
        q_nope_out,
        k_nope,
        forward_batch,
        zero_allocator,
        positions,
        topk_indices,
1678
    ):
1679
        if self.current_attention_backend in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS:
Faraz's avatar
Faraz committed
1680
1681
1682
1683
1684
1685
            extra_args = {}
            if self._fuse_rope_for_trtllm_mla(forward_batch):
                extra_args = {
                    "cos_sin_cache": self.rotary_emb.cos_sin_cache,
                    "is_neox": self.rotary_emb.is_neox_style,
                }
fzyzcjy's avatar
fzyzcjy committed
1686

1687
            attn_output = self.attn_mqa(
Faraz's avatar
Faraz committed
1688
1689
1690
1691
1692
1693
1694
                q_nope_out,
                k_nope,
                k_nope,
                forward_batch,
                q_rope=q_pe,
                k_rope=k_pe,
                **extra_args,
fzyzcjy's avatar
fzyzcjy committed
1695
                **(dict(topk_indices=topk_indices) if topk_indices is not None else {}),
1696
1697
            )
        else:
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
            if _use_aiter_gfx95:
                cos = self.rotary_emb.cos_cache
                sin = self.rotary_emb.sin_cache
                q, k = fused_qk_rope_cat(
                    q_nope_out,
                    q_pe,
                    k_nope,
                    k_pe,
                    positions,
                    cos,
                    sin,
                    self.rotary_emb.is_neox_style,
                )
            else:
                q = torch.cat([q_nope_out, q_pe], dim=-1)
                k = torch.cat([k_nope, k_pe], dim=-1)

fzyzcjy's avatar
fzyzcjy committed
1715
1716
1717
1718
1719
1720
1721
            attn_output = self.attn_mqa(
                q,
                k,
                k_nope,
                forward_batch,
                **(dict(topk_indices=topk_indices) if topk_indices is not None else {}),
            )
1722
1723
        attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)

1724
1725
        if self.use_deep_gemm_bmm:
            attn_output_val, attn_output_scale, masked_m, expected_m, aligned_m = (
1726
1727
                per_token_group_quant_mla_deep_gemm_masked_fp8(
                    attn_output.transpose(0, 1)
1728
1729
1730
1731
1732
                )
            )
            attn_bmm_output = attn_output.new_empty(
                (self.num_local_heads, aligned_m, self.v_head_dim)
            )
1733
            deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
1734
1735
1736
1737
1738
1739
                (attn_output_val, attn_output_scale),
                (self.w_vc, self.w_scale_v),
                attn_bmm_output,
                masked_m,
                expected_m,
            )
Ke Bao's avatar
Ke Bao committed
1740
1741
1742
            attn_bmm_output = (
                attn_bmm_output[:, :expected_m, :].transpose(0, 1).flatten(1, 2)
            )
1743
1744
        elif _is_hip:
            # TODO(haishaw): add bmm_fp8 to ROCm
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
            if _use_aiter_gfx95 and self.w_vc.dtype == torch.uint8:
                x = attn_output.transpose(0, 1)
                attn_bmm_output = torch.empty(
                    x.shape[0],
                    x.shape[1],
                    self.w_vc.shape[2],
                    device=x.device,
                    dtype=torch.bfloat16,
                )
                batched_gemm_afp4wfp4_pre_quant(
                    x,
                    self.w_vc.transpose(-2, -1),
                    self.w_scale_v.transpose(-2, -1),
                    torch.bfloat16,
                    attn_bmm_output,
                )
            else:
                attn_bmm_output = torch.bmm(
                    attn_output.to(torch.bfloat16).transpose(0, 1),
                    self.w_vc.to(torch.bfloat16) * self.w_scale,
                )

            if self.o_proj.weight.dtype == torch.uint8:
                attn_bmm_output = attn_bmm_output.transpose(0, 1)
                attn_bmm_output = fused_flatten_mxfp4_quant(attn_bmm_output)
            else:
                attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)

1773
        elif self.w_vc.dtype == torch.float8_e4m3fn:
1774
1775
1776
1777
1778
1779
1780
1781
            attn_output_val, attn_output_scale = per_tensor_quant_mla_fp8(
                attn_output.transpose(0, 1),
                (
                    torch.zeros((1,), dtype=torch.float32, device=attn_output.device)
                    if _is_cublas_ge_129
                    else zero_allocator.allocate(1)
                ),
            )
1782
1783
1784
1785
1786
1787
1788
            attn_bmm_output = bmm_fp8(
                attn_output_val,
                self.w_vc,
                attn_output_scale,
                self.w_scale,
                torch.bfloat16,
            )
Ke Bao's avatar
Ke Bao committed
1789
            attn_bmm_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
1790
        else:
Ke Bao's avatar
Ke Bao committed
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
            attn_bmm_output = torch.empty(
                (attn_output.shape[0], self.num_local_heads * self.v_head_dim),
                dtype=attn_output.dtype,
                device=attn_output.device,
            )
            torch.bmm(
                attn_output.transpose(0, 1),
                self.w_vc,
                out=attn_bmm_output.view(
                    -1, self.num_local_heads, self.v_head_dim
                ).transpose(0, 1),
            )
        output, _ = self.o_proj(attn_bmm_output)
1804
1805
1806

        return output

fzyzcjy's avatar
fzyzcjy committed
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
    def forward_npu_sparse_prepare(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        zero_allocator: BumpAllocator,
    ):
        """
        Reuse `self.q_lora_rank is not None` branch from forward_absorb_prepare
        """
        if self.is_mla_preprocess_enabled and forward_batch.forward_mode.is_decode():
            if self.mla_preprocess is None:
                self.mla_preprocess = NPUFusedMLAPreprocess(
                    self.fused_qkv_a_proj_with_mqa,
                    self.q_a_layernorm,
                    self.kv_a_layernorm,
                    self.q_b_proj,
                    self.w_kc,
                    self.rotary_emb,
                    self.layer_id,
                    self.num_local_heads,
                    self.qk_nope_head_dim,
                    self.qk_rope_head_dim,
                )
            (
                q_pe,
                k_pe,
                q_nope_out,
                k_nope,
                forward_batch,
                zero_allocator,
                positions,
            ) = self.mla_preprocess.forward(
                positions, hidden_states, forward_batch, zero_allocator
            )

            fused_qkv_a_proj_out = self.fused_qkv_a_proj_with_mqa(hidden_states)[0]
            q, _ = fused_qkv_a_proj_out.split(
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
            )
            q_lora = self.q_a_layernorm(q)
        else:
            from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode

            if (
                (not isinstance(hidden_states, tuple))
                and hidden_states.shape[0] <= 16
                and self.use_min_latency_fused_a_gemm
            ):
                fused_qkv_a_proj_out = dsv3_fused_a_gemm(
                    hidden_states, self.fused_qkv_a_proj_with_mqa.weight.T
                )
            else:
                fused_qkv_a_proj_out = self.fused_qkv_a_proj_with_mqa(hidden_states)[0]
            q, latent_cache = fused_qkv_a_proj_out.split(
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
            )
            k_nope = latent_cache[..., : self.kv_lora_rank]

            # overlap qk norm
            if self.alt_stream is not None and get_is_capture_mode():
                current_stream = torch.cuda.current_stream()
                self.alt_stream.wait_stream(current_stream)
                q = self.q_a_layernorm(q)
                with torch.cuda.stream(self.alt_stream):
                    k_nope = self.kv_a_layernorm(k_nope)
                current_stream.wait_stream(self.alt_stream)
            else:
                if _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
                    q, k_nope = fused_rms_mxfp4_quant(
                        q,
                        self.q_a_layernorm.weight,
                        self.q_a_layernorm.variance_epsilon,
                        k_nope,
                        self.kv_a_layernorm.weight,
                        self.kv_a_layernorm.variance_epsilon,
                    )
                else:
                    q = self.q_a_layernorm(q)
                    k_nope = self.kv_a_layernorm(k_nope)

            q_lora = q.clone()  # required for topk_indices
            k_nope = k_nope.unsqueeze(1)
            q = self.q_b_proj(q)[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
            )
            k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1)

            if self.use_deep_gemm_bmm:
                q_nope_val, q_nope_scale, masked_m, expected_m, aligned_m = (
                    per_token_group_quant_mla_deep_gemm_masked_fp8(
                        q_nope.transpose(0, 1)
                    )
                )
                q_nope_out = q_nope.new_empty(
                    (self.num_local_heads, aligned_m, self.kv_lora_rank)
                )
                deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
                    (q_nope_val, q_nope_scale),
                    (self.w_kc, self.w_scale_k),
                    q_nope_out,
                    masked_m,
                    expected_m,
                )
                q_nope_out = q_nope_out[:, :expected_m, :]
            elif _is_hip:
                # TODO(haishaw): add bmm_fp8 to ROCm
                if _use_aiter_gfx95 and self.w_kc.dtype == torch.uint8:
                    x = q_nope.transpose(0, 1)
                    q_nope_out = torch.empty(
                        x.shape[0],
                        x.shape[1],
                        self.w_kc.shape[2],
                        device=x.device,
                        dtype=torch.bfloat16,
                    )
                    batched_gemm_afp4wfp4_pre_quant(
                        x,
                        self.w_kc.transpose(-2, -1),
                        self.w_scale_k.transpose(-2, -1),
                        torch.bfloat16,
                        q_nope_out,
                    )
                else:
                    q_nope_out = torch.bmm(
                        q_nope.to(torch.bfloat16).transpose(0, 1),
                        self.w_kc.to(torch.bfloat16) * self.w_scale,
                    )
            elif self.w_kc.dtype == torch.float8_e4m3fn:
                q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8(
                    q_nope.transpose(0, 1),
                    zero_allocator.allocate(1),
                )
                q_nope_out = bmm_fp8(
                    q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16
                )
            else:
                q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc)

            q_nope_out = q_nope_out.transpose(0, 1)

            if not self._fuse_rope_for_trtllm_mla(forward_batch) and (
                not _use_aiter or not _is_gfx95_supported
            ):
                q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)

        # TODO: multi-stream indexer
        topk_indices = self.indexer(
            hidden_states, q_lora, positions, forward_batch, self.layer_id
        )

        return (
            q_pe,
            k_pe,
            q_nope_out,
            k_nope,
            topk_indices,
            forward_batch,
            zero_allocator,
            positions,
        )

    def forward_npu_sparse_core(
        self,
        q_pe,
        k_pe,
        q_nope_out,
        k_nope,
        topk_indices,
        forward_batch,
        zero_allocator,
        positions,
    ):
        attn_output = self.attn_mqa(
            q_nope_out.contiguous(),
            k_nope.contiguous(),
            k_nope.contiguous(),
            forward_batch,
            save_kv_cache=True,  # False if forward_batch.forward_mode.is_extend() else True,
            q_rope=q_pe.contiguous(),
            k_rope=k_pe.contiguous(),
            topk_indices=topk_indices,
        )
        attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)

        attn_bmm_output = torch.empty(
            (attn_output.shape[0], self.num_local_heads, self.v_head_dim),
            dtype=attn_output.dtype,
            device=attn_output.device,
        )

        if not forward_batch.forward_mode.is_decode():
            attn_output = attn_output.transpose(0, 1)
            torch.bmm(
                attn_output,
                self.w_vc,
                out=attn_bmm_output.view(
                    -1, self.num_local_heads, self.v_head_dim
                ).transpose(0, 1),
            )
        else:
            attn_output = attn_output.contiguous()
            torch.ops.npu.batch_matmul_transpose(
                attn_output, self.w_vc, attn_bmm_output
            )

        attn_bmm_output = attn_bmm_output.reshape(
            -1, self.num_local_heads * self.v_head_dim
        )

        output, _ = self.o_proj(attn_bmm_output)
        return output

2022
    def forward_absorb_fused_mla_rope_prepare(
2023
2024
2025
2026
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
2027
        zero_allocator: BumpAllocator,
2028
    ):
2029
2030
2031
2032
2033
2034
2035
2036
        enable_rope_fusion = (
            os.getenv("SGLANG_FUSED_MLA_ENABLE_ROPE_FUSION", "1") == "1"
        )
        q_len = hidden_states.shape[0]
        q_input = hidden_states.new_empty(
            q_len, self.num_local_heads, self.kv_lora_rank + self.qk_rope_head_dim
        )
        if self.q_lora_rank is not None:
2037
2038
2039
            q, latent_cache = self.fused_qkv_a_proj_with_mqa(hidden_states)[0].split(
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
            )
2040
2041
2042
2043
2044
2045
            q = self.q_a_layernorm(q)
            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
        else:
            q = self.q_proj(hidden_states)[0].view(
                -1, self.num_local_heads, self.qk_head_dim
            )
2046
            latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
2047
2048
        q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

2049
2050
        if _is_hip:
            # TODO(haishaw): add bmm_fp8 to ROCm
2051
2052
2053
2054
2055
            q_nope_out = torch.bmm(
                q_nope.to(torch.bfloat16).transpose(0, 1),
                self.w_kc.to(torch.bfloat16) * self.w_scale,
            )
        elif self.w_kc.dtype == torch.float8_e4m3fn:
2056
            q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8(
2057
2058
2059
                q_nope.transpose(0, 1),
                zero_allocator.allocate(1),
                dtype=torch.float8_e4m3fn,
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
            )
            q_nope_out = bmm_fp8(
                q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16
            )
        else:
            q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc)
        q_input[..., : self.kv_lora_rank] = q_nope_out.transpose(0, 1)
        v_input = latent_cache[..., : self.kv_lora_rank]
        v_input = self.kv_a_layernorm(v_input.contiguous()).unsqueeze(1)
        k_input = latent_cache.unsqueeze(1)
        k_input[..., : self.kv_lora_rank] = v_input

        if not enable_rope_fusion:
            k_pe = k_input[..., self.kv_lora_rank :]
            q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
            q_input[..., self.kv_lora_rank :] = q_pe
            k_input[..., self.kv_lora_rank :] = k_pe
            k_pe_output = None
        else:
            k_pe_output = torch.empty_like(k_input[..., self.kv_lora_rank :])

        q_input[..., self.kv_lora_rank :] = q_pe

        # attn_output = self.attn_mqa(q_input, k_input, v_input, forward_batch)
        # Use Fused ROPE with use_rope=OFF.
        attn_output = torch.empty(
            (q_len, self.num_local_heads, self.kv_lora_rank),
            dtype=q.dtype,
            device=q.device,
        )
        attn_logits, _, kv_indptr, kv_indices, _, _, _ = (
            forward_batch.attn_backend.forward_metadata
        )
        cos_sin_cache = self.rotary_emb.cos_sin_cache
        num_kv_split = forward_batch.attn_backend.num_kv_splits
        sm_scale = self.attn_mqa.scaling
        if attn_logits is None:
            attn_logits = torch.empty(
                (
                    forward_batch.batch_size,
                    self.num_local_heads,
                    num_kv_split,
                    self.kv_lora_rank + 1,
                ),
                dtype=torch.float32,
                device=q.device,
            )

        # save current latent cache.
        forward_batch.token_to_kv_pool.set_kv_buffer(
            self.attn_mqa, forward_batch.out_cache_loc, k_input, None
        )
        key_cache_buf = forward_batch.token_to_kv_pool.get_key_buffer(
            self.attn_mqa.layer_id
        )
        val_cache_buf = key_cache_buf[..., : self.kv_lora_rank]

2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
        return (
            q_input,
            key_cache_buf,
            val_cache_buf,
            attn_output,
            kv_indptr,
            kv_indices,
            k_pe_output,
            cos_sin_cache,
            positions,
            attn_logits,
            num_kv_split,
            sm_scale,
            enable_rope_fusion,
            k_input,
            forward_batch,
            zero_allocator,
        )

2136
2137
2138
2139
2140
2141
2142
    def forward_absorb_fused_mla_rope_cpu_prepare(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        zero_allocator: BumpAllocator,
    ):
2143
2144
        assert self.q_lora_rank is not None and use_intel_amx_backend(
            self
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
        ), "forward_absorb_fused_mla_rope_cpu_prepare requires q_lora_rank is not None and use_intel_amx_backend"

        q_input, k_input, v_input = (
            torch.ops.sgl_kernel.qkv_proj_with_rope_fused_weight(
                hidden_states,
                self.fused_qkv_a_proj_with_mqa.weight,
                self.q_b_proj.weight,
                self.w_kc,
                self.q_a_layernorm.weight,
                self.kv_a_layernorm.weight,
                positions,
                self.rotary_emb.cos_sin_cache,
                self.kv_a_layernorm.variance_epsilon,
                self.qkv_proj_with_rope_is_int8,
                self.qkv_proj_with_rope_is_fp8,
                (
                    self.fused_qkv_a_proj_with_mqa.weight_scale
                    if self.qkv_proj_with_rope_is_int8
                    else (
                        self.fused_qkv_a_proj_with_mqa.weight_scale_inv
                        if self.qkv_proj_with_rope_is_fp8
                        else None
                    )
                ),
                (
                    self.q_b_proj.weight_scale
                    if self.qkv_proj_with_rope_is_int8
                    else (
                        self.q_b_proj.weight_scale_inv
                        if self.qkv_proj_with_rope_is_fp8
                        else None
                    )
                ),
                True,  # is_vnni
                self.weight_block_size,
                self.q_lora_rank,
                self.kv_lora_rank,
                self.qk_rope_head_dim,
            )
        )
        return (q_input, k_input, v_input, forward_batch, zero_allocator)

2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
    def forward_absorb_fused_mla_rope_core(
        self,
        q_input,
        key_cache_buf,
        val_cache_buf,
        attn_output,
        kv_indptr,
        kv_indices,
        k_pe_output,
        cos_sin_cache,
        positions,
        attn_logits,
        num_kv_split,
        sm_scale,
        enable_rope_fusion,
        k_input,
        forward_batch,
        zero_allocator,
    ):
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
        decode_attention_fwd_grouped_rope(
            q_input,
            key_cache_buf,
            val_cache_buf,
            attn_output,
            kv_indptr,
            kv_indices,
            k_pe_output,
            self.kv_lora_rank,
            self.rotary_emb.rotary_dim,
            cos_sin_cache,
            positions,
            attn_logits,
            num_kv_split,
            sm_scale,
            logit_cap=self.attn_mqa.logit_cap,
            use_rope=enable_rope_fusion,
            is_neox_style=self.rotary_emb.is_neox_style,
        )

        if enable_rope_fusion:
            k_input[..., self.kv_lora_rank :] = k_pe_output
            forward_batch.token_to_kv_pool.set_kv_buffer(
                self.attn_mqa, forward_batch.out_cache_loc, k_input, None
            )

        attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)

2234
2235
        if _is_hip:
            # TODO(haishaw): add bmm_fp8 to ROCm
2236
2237
2238
2239
2240
            attn_bmm_output = torch.bmm(
                attn_output.to(torch.bfloat16).transpose(0, 1),
                self.w_vc.to(torch.bfloat16) * self.w_scale,
            )
        elif self.w_vc.dtype == torch.float8_e4m3fn:
2241
            attn_output_val, attn_output_scale = per_tensor_quant_mla_fp8(
2242
2243
2244
                attn_output.transpose(0, 1),
                zero_allocator.allocate(1),
                dtype=torch.float8_e4m3fn,
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
            )
            attn_bmm_output = bmm_fp8(
                attn_output_val,
                self.w_vc,
                attn_output_scale,
                self.w_scale,
                torch.bfloat16,
            )
        else:
            attn_bmm_output = torch.bmm(attn_output.transpose(0, 1), self.w_vc)
        attn_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
2256
2257
2258
2259
        output, _ = self.o_proj(attn_output)

        return output

2260
2261
2262
    def forward_absorb_fused_mla_rope_cpu_core(
        self, q_input, k_input, v_input, forward_batch, zero_allocator
    ):
2263
2264
        assert self.q_lora_rank is not None and use_intel_amx_backend(
            self
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
        ), "forward_absorb_fused_mla_rope_cpu_core requires q_lora_rank is not None and use_intel_amx_backend"

        attn_output = self.attn_mqa(q_input, k_input, v_input, forward_batch)
        attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)

        # [Note] Align shapes of bmm inputs.
        # Shapes of inputs:
        #   q_nope: [M, B, K]
        #   original self.w_kc: [B, K, N]
        #   current self.w_kc (which has been converted in PackWeightMethod): [B, N, K]

        # Shapes of inputs to sgl_kernel.cpu.bmm:
        #   out: [B, M, N]
        #   mat1: [B, M, K]
        #   mat2: [B, N, K]
        B = self.w_vc.size(0)
        N = self.w_vc.size(1)
        M = attn_output.size(0)
        output = torch.empty([M, int(B * N)], dtype=attn_output.dtype)
        attn_bmm_output = output.view([M, B, N]).transpose_(0, 1)
        torch.ops.sgl_kernel.bmm_cpu(
            attn_bmm_output,
            attn_output.transpose(0, 1),
            self.w_vc,
            True,  # is_vnni
            None,  # scale
        )
        attn_output = output
        output, _ = self.o_proj(attn_output)

        return output

2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
    def _chunked_prefix_attn_mha(
        self,
        q: torch.Tensor,
        accum_output: torch.Tensor,
        accum_lse: torch.Tensor,
        forward_batch: ForwardBatch,
    ) -> torch.Tensor:

        assert forward_batch.num_prefix_chunks is not None
        for i in range(forward_batch.num_prefix_chunks):
            forward_batch.set_prefix_chunk_idx(i)

            # Fetch latent cache from memory pool with precomputed chunked kv indices
            latent_cache_buf = forward_batch.token_to_kv_pool.get_key_buffer(
                self.attn_mha.layer_id
            )
2313
2314
2315
2316
2317
            latent_cache = (
                latent_cache_buf[forward_batch.prefix_chunk_kv_indices[i]]
                .contiguous()
                .to(q.dtype)
            )
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346

            kv_a_normed, k_pe = latent_cache.split(
                [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
            )
            kv_a_normed = kv_a_normed.squeeze(1).contiguous()
            kv = self.kv_b_proj(kv_a_normed)[0]
            kv = kv.view(
                -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim
            )
            v = kv[..., self.qk_nope_head_dim :]
            k_nope = kv[..., : self.qk_nope_head_dim]

            k = torch.empty(
                (
                    k_nope.shape[0],
                    self.num_local_heads,
                    self.qk_nope_head_dim + self.qk_rope_head_dim,
                ),
                dtype=v.dtype,
                device=v.device,
            )
            k[..., : self.qk_nope_head_dim] = k_nope
            k[..., self.qk_nope_head_dim :] = k_pe

            output, lse = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False)
            tmp_output = torch.empty_like(accum_output)
            tmp_lse = torch.empty_like(accum_lse)
            merge_state_v2(output, lse, accum_output, accum_lse, tmp_output, tmp_lse)
            accum_output, accum_lse = tmp_output, tmp_lse
2347
            del kv, k, v, output, lse, tmp_output, tmp_lse
2348
2349
2350

        return accum_output

2351
    def forward_normal_chunked_kv_prepare(
2352
2353
2354
2355
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
2356
2357
        zero_allocator: BumpAllocator,
    ):
2358
2359
2360
2361
2362
2363
2364
        # In normal mha, the k and v tensors will become overly large when the prefix length is long.
        # To avoid this, we split the kv cache into chunks and process them one after another.
        # Since mha is compute friendly, the for loop induced here will not introduce significant overhead.
        # The top comments in https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/mla/common.py
        # will be helpful for understanding the purpose of this function.

        # First do normal mha forward to get output for extended part
2365
2366
        return self.forward_normal_prepare(
            positions, hidden_states, forward_batch, zero_allocator
2367
2368
        )

2369
    def forward_normal_chunked_kv_core(self, q, k, v, forward_batch):
2370
2371
2372
2373
2374
2375
2376
2377
        has_extend_prefix = any(forward_batch.extend_prefix_lens_cpu)
        # Only initialize the info once
        if has_extend_prefix and forward_batch.num_prefix_chunks is None:
            forward_batch.prepare_chunked_prefix_cache_info(q.device)
            if hasattr(forward_batch.attn_backend, "init_mha_chunk_metadata"):
                forward_batch.attn_backend.init_mha_chunk_metadata(forward_batch)

        forward_batch.mha_return_lse = has_extend_prefix
2378
2379
        # Do mha for extended part without prefix
        forward_batch.set_attn_attend_prefix_cache(False)
2380
        attn_output = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False)
2381
2382

        # Do mha attention with chunked prefix cache if there are any sequence with prefix
2383
2384
        if has_extend_prefix:
            attn_output, lse = attn_output
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
            forward_batch.set_attn_attend_prefix_cache(True)
            attn_output = self._chunked_prefix_attn_mha(
                q=q,
                accum_output=attn_output,
                accum_lse=lse,
                forward_batch=forward_batch,
            )

        attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim)
        output, _ = self.o_proj(attn_output)
        return output

2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
    @staticmethod
    def _get_q_b_proj_quant_config(quant_config):
        if get_bool_env_var("SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN"):
            # refer to real DeepSeek V3 quant config
            return Fp8Config(
                is_checkpoint_fp8_serialized=True,
                weight_block_size=[128, 128],
            )
        else:
            return quant_config

2408

Liangsheng Yin's avatar
Liangsheng Yin committed
2409
2410
2411
2412
2413
2414
2415
class DeepseekV2DecoderLayer(nn.Module):

    def __init__(
        self,
        config: PretrainedConfig,
        layer_id: int,
        quant_config: Optional[QuantizationConfig] = None,
2416
        moe_quant_config: Optional[QuantizationConfig] = None,
2417
        is_nextn: bool = False,
2418
        prefix: str = "",
2419
        alt_stream: Optional[torch.cuda.Stream] = None,
Liangsheng Yin's avatar
Liangsheng Yin committed
2420
2421
2422
    ) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
2423
        self.config = config
Liangsheng Yin's avatar
Liangsheng Yin committed
2424
2425
2426
        rope_theta = getattr(config, "rope_theta", 10000)
        rope_scaling = getattr(config, "rope_scaling", None)
        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
2427
2428
2429
        self.speculative_algorithm = SpeculativeAlgorithm.from_string(
            get_global_server_args().speculative_algorithm
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
2430
        self.layer_id = layer_id
2431
        self.is_nextn = is_nextn
Baizhou Zhang's avatar
Baizhou Zhang committed
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
        self.self_attn = DeepseekV2AttentionMLA(
            config=config,
            hidden_size=self.hidden_size,
            num_heads=config.num_attention_heads,
            qk_nope_head_dim=config.qk_nope_head_dim,
            qk_rope_head_dim=config.qk_rope_head_dim,
            v_head_dim=config.v_head_dim,
            q_lora_rank=(
                config.q_lora_rank if hasattr(config, "q_lora_rank") else None
            ),
            kv_lora_rank=config.kv_lora_rank,
            rope_theta=rope_theta,
            rope_scaling=rope_scaling,
            max_position_embeddings=max_position_embeddings,
            quant_config=quant_config,
            layer_id=layer_id,
            reduce_results=False,
            prefix=add_prefix("self_attn", prefix),
2450
            alt_stream=alt_stream,
Baizhou Zhang's avatar
Baizhou Zhang committed
2451
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
2452

2453
2454
2455
2456
2457
        self.is_layer_sparse = self._is_layer_sparse(layer_id, is_nextn=is_nextn)
        is_previous_layer_sparse = self._is_layer_sparse(layer_id - 1, is_nextn=False)

        self.layer_scatter_modes = LayerScatterModes.init_new(
            layer_id=layer_id,
2458
            num_layers=1 if is_nextn else config.num_hidden_layers,
2459
2460
            is_layer_sparse=self.is_layer_sparse,
            is_previous_layer_sparse=is_previous_layer_sparse,
2461
2462
        )

2463
        if self.is_layer_sparse:
2464
2465
            self.mlp = DeepseekV2MoE(
                config=config,
2466
                quant_config=moe_quant_config or quant_config,
2467
                prefix=add_prefix("mlp", prefix),
fzyzcjy's avatar
fzyzcjy committed
2468
                layer_id=self.layer_id,
2469
                alt_stream=alt_stream,
2470
                is_nextn=is_nextn,
2471
            )
Liangsheng Yin's avatar
Liangsheng Yin committed
2472
        else:
2473
            if enable_moe_dense_fully_dp():
2474
2475
2476
                mlp_tp_rank, mlp_tp_size = 0, 1
            else:
                mlp_tp_rank, mlp_tp_size = None, None
Liangsheng Yin's avatar
Liangsheng Yin committed
2477
2478
2479
2480
2481
            self.mlp = DeepseekV2MLP(
                hidden_size=config.hidden_size,
                intermediate_size=config.intermediate_size,
                hidden_act=config.hidden_act,
                quant_config=quant_config,
2482
                prefix=add_prefix("mlp", prefix),
2483
2484
                tp_rank=mlp_tp_rank,
                tp_size=mlp_tp_size,
Liangsheng Yin's avatar
Liangsheng Yin committed
2485
            )
2486

Liangsheng Yin's avatar
Liangsheng Yin committed
2487
2488
2489
2490
2491
        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
        )

2492
2493
2494
2495
        self.layer_communicator = LayerCommunicator(
            layer_scatter_modes=self.layer_scatter_modes,
            input_layernorm=self.input_layernorm,
            post_attention_layernorm=self.post_attention_layernorm,
2496
            allow_reduce_scatter=True,
2497
2498
2499
            is_last_layer=(
                is_nextn or (self.layer_id == self.config.num_hidden_layers - 1)
            ),
2500
        )
2501
2502
2503
2504
2505
2506

    def _is_layer_sparse(self, layer_id: int, is_nextn: bool) -> bool:
        return is_nextn or (
            self.config.n_routed_experts is not None
            and layer_id >= self.config.first_k_dense_replace
            and layer_id % self.config.moe_layer_freq == 0
2507
2508
        )

Liangsheng Yin's avatar
Liangsheng Yin committed
2509
2510
2511
2512
    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
2513
        forward_batch: ForwardBatch,
Liangsheng Yin's avatar
Liangsheng Yin committed
2514
        residual: Optional[torch.Tensor],
2515
        zero_allocator: BumpAllocator,
2516
        gemm_output_zero_allocator: BumpAllocator = None,
Liangsheng Yin's avatar
Liangsheng Yin committed
2517
    ) -> torch.Tensor:
2518
2519
2520
        quant_format = (
            "mxfp4"
            if _is_gfx95_supported
2521
2522
2523
2524
            and getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None) is not None
            and getattr(self.self_attn.fused_qkv_a_proj_with_mqa, "weight", None)
            is not None
            and self.self_attn.fused_qkv_a_proj_with_mqa.weight.dtype == torch.uint8
2525
2526
2527
            else ""
        )

2528
        hidden_states, residual = self.layer_communicator.prepare_attn(
2529
2530
2531
2532
            hidden_states,
            residual,
            forward_batch,
            quant_format,
2533
2534
        )

2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
        hidden_states = self.self_attn(
            positions=positions,
            hidden_states=hidden_states,
            forward_batch=forward_batch,
            zero_allocator=zero_allocator,
        )

        hidden_states, residual = self.layer_communicator.prepare_mlp(
            hidden_states, residual, forward_batch
        )

2546
        should_allreduce_fusion = (
2547
2548
            self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer(
                forward_batch
2549
            )
2550
2551
        )

2552
2553
2554
2555
        # For DP with padding, reduce scatter can be used instead of all-reduce.
        use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter(
            forward_batch
        )
2556
2557
2558
2559

        if isinstance(self.mlp, DeepseekV2MLP):
            gemm_output_zero_allocator = None

2560
        hidden_states = self.mlp(
2561
2562
2563
2564
2565
            hidden_states,
            forward_batch,
            should_allreduce_fusion,
            use_reduce_scatter,
            gemm_output_zero_allocator,
2566
        )
2567

2568
        if should_allreduce_fusion:
2569
2570
            hidden_states._sglang_needs_allreduce_fusion = True

2571
        if not should_allreduce_fusion:
2572
2573
2574
2575
            hidden_states, residual = self.layer_communicator.postprocess_layer(
                hidden_states, residual, forward_batch
            )

2576
2577
        return hidden_states, residual

2578
2579
2580
2581
2582
2583
2584
2585
    def op_comm_prepare_attn(
        self,
        state,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        residual: Optional[torch.Tensor],
        zero_allocator: BumpAllocator,
2586
        tbo_subbatch_index: Optional[int] = None,
2587
2588
    ):
        state.hidden_states_after_comm_pre_attn, state.residual_after_input_ln = (
fzyzcjy's avatar
fzyzcjy committed
2589
            self.layer_communicator.prepare_attn(hidden_states, residual, forward_batch)
2590
2591
2592
2593
2594
2595
        )
        state.update(
            dict(
                forward_batch=forward_batch,
                positions=positions,
                zero_allocator=zero_allocator,
2596
                tbo_subbatch_index=tbo_subbatch_index,
2597
            )
2598
        )
2599

2600
2601
2602
2603
2604
2605
2606
    def op_comm_prepare_mlp(self, state):
        state.hidden_states_mlp_input, state.residual_after_comm_pre_mlp = (
            self.layer_communicator.prepare_mlp(
                state.pop("hidden_states_after_attn"),
                state.pop("residual_after_input_ln"),
                state.forward_batch,
            )
2607
        )
2608

2609
2610
2611
2612
2613
2614
2615
2616
    def op_mlp(self, state):
        hidden_states = state.pop("hidden_states_mlp_input")
        if not (
            enable_moe_dense_fully_dp()
            and (not self.is_layer_sparse)
            and hidden_states.shape[0] == 0
        ):
            state.hidden_states_mlp_output = self.mlp(
2617
                hidden_states, state.forward_batch
2618
2619
2620
            )
        else:
            state.hidden_states_mlp_output = hidden_states
2621

2622
    def op_comm_postprocess_layer(self, state):
2623
        hidden_states, residual = self.layer_communicator.postprocess_layer(
2624
2625
2626
            state.pop("hidden_states_mlp_output"),
            state.pop("residual_after_comm_pre_mlp"),
            state.forward_batch,
2627
        )
2628

2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
        output = dict(
            positions=state.positions,
            hidden_states=hidden_states,
            residual=residual,
            forward_batch=state.forward_batch,
            zero_allocator=state.zero_allocator,
            tbo_subbatch_index=state.tbo_subbatch_index,
        )

        state.clear(
            expect_keys={
                "positions",
                "forward_batch",
                "zero_allocator",
                "tbo_subbatch_index",
            }
        )
        return output
2647

Liangsheng Yin's avatar
Liangsheng Yin committed
2648
2649
2650
2651
2652
2653
2654
2655

class DeepseekV2Model(nn.Module):
    fall_back_to_pt_during_load = False

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
2656
        prefix: str = "",
Liangsheng Yin's avatar
Liangsheng Yin committed
2657
2658
2659
2660
    ) -> None:
        super().__init__()
        self.padding_id = config.pad_token_id
        self.vocab_size = config.vocab_size
2661
        self.first_k_dense_replace = config.first_k_dense_replace
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
        self.pp_group = get_pp_group()

        if self.pp_group.is_first_rank:
            self.embed_tokens = VocabParallelEmbedding(
                config.vocab_size,
                config.hidden_size,
                enable_tp=not is_dp_attention_enabled(),
            )
        else:
            self.embed_tokens = PPMissingLayer()
Liangsheng Yin's avatar
Liangsheng Yin committed
2672

2673
        self.alt_stream = torch.cuda.Stream() if _is_cuda else None
2674
2675
2676
2677
2678
2679
2680
2681
2682
2683
2684
2685
        self.layers, self.start_layer, self.end_layer = make_layers(
            config.num_hidden_layers,
            lambda idx, prefix: DeepseekV2DecoderLayer(
                config=config,
                layer_id=idx,
                quant_config=quant_config,
                prefix=prefix,
                alt_stream=self.alt_stream,
            ),
            pp_rank=self.pp_group.rank_in_group,
            pp_size=self.pp_group.world_size,
            prefix=add_prefix("layers", prefix),
fzyzcjy's avatar
fzyzcjy committed
2686
2687
2688
2689
2690
2691
2692
2693
2694
2695
            offloader_kwargs=dict(
                submodule_accessor=lambda layer: (
                    layer.mlp.experts
                    if isinstance(layer.mlp, DeepseekV2MoE)
                    else layer.mlp
                ),
                whitelist_param_names_creator=lambda module: (
                    [
                        "w13_weight",
                        "w2_weight",
fzyzcjy's avatar
fzyzcjy committed
2696
2697
2698
2699
2700
2701
2702
2703
2704
                        # only for nvfp4
                        *(
                            [
                                "w13_blockscale_swizzled",
                                "w2_blockscale_swizzled",
                            ]
                            if hasattr(module, "w13_blockscale_swizzled")
                            else []
                        ),
fzyzcjy's avatar
fzyzcjy committed
2705
2706
2707
2708
2709
                    ]
                    if isinstance(module, FusedMoE)
                    else []
                ),
            ),
Liangsheng Yin's avatar
Liangsheng Yin committed
2710
        )
2711
2712
2713
2714
        if self.pp_group.is_last_rank:
            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = PPMissingLayer(return_tuple=True)
Liangsheng Yin's avatar
Liangsheng Yin committed
2715

2716
2717
2718
2719
2720
2721
2722
2723
2724
2725
2726
2727
2728
2729
2730
2731
2732
2733
2734
2735
2736
2737
2738
2739
2740
2741
2742
2743
2744
2745
2746
        self.gemm_output_zero_allocator_size = 0
        if (
            _use_aiter_gfx95
            and config.n_routed_experts == 256
            and self.embed_tokens.embedding_dim == 7168
        ):
            num_moe_layers = sum(
                [
                    1
                    for i in range(len(self.layers))
                    if isinstance(self.layers[i].mlp, DeepseekV2MoE)
                ]
            )

            allocate_size = 0
            for i in range(len(self.layers)):
                if isinstance(self.layers[i].mlp, DeepseekV2MoE):
                    allocate_size = self.layers[
                        i
                    ].mlp.shared_experts.gate_up_proj.output_size_per_partition
                    break

            self.gemm_output_zero_allocator_size = (
                get_dsv3_gemm_output_zero_allocator_size(
                    config.n_routed_experts,
                    num_moe_layers,
                    allocate_size,
                    self.embed_tokens.embedding_dim,
                )
            )

2747
2748
2749
    def get_input_embeddings(self) -> torch.Tensor:
        return self.embed_tokens

Liangsheng Yin's avatar
Liangsheng Yin committed
2750
2751
2752
2753
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
2754
        forward_batch: ForwardBatch,
2755
        input_embeds: torch.Tensor = None,
2756
2757
2758
        pp_proxy_tensors: Optional[PPProxyTensors] = None,
    ) -> Union[torch.Tensor, PPProxyTensors]:
        total_num_layers = self.end_layer - self.start_layer
2759
        device = input_embeds.device if input_embeds is not None else input_ids.device
2760
        zero_allocator = BumpAllocator(
2761
            buffer_size=total_num_layers * 2 * (2 if forward_batch.can_run_tbo else 1),
2762
            dtype=torch.float32,
2763
            device=device,
2764
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
2765

2766
2767
2768
2769
2770
2771
2772
2773
2774
2775
2776
2777
2778
2779
2780
        has_gemm_output_zero_allocator = hasattr(
            self, "gemm_output_zero_allocator_size"
        )

        gemm_output_zero_allocator = (
            BumpAllocator(
                buffer_size=self.gemm_output_zero_allocator_size,
                dtype=torch.float32,
                device=device,
            )
            if has_gemm_output_zero_allocator
            and self.gemm_output_zero_allocator_size > 0
            else None
        )

2781
2782
2783
2784
2785
2786
        if self.pp_group.is_first_rank:
            if input_embeds is None:
                hidden_states = self.embed_tokens(input_ids)
            else:
                hidden_states = input_embeds
            residual = None
2787
        else:
2788
2789
2790
            assert pp_proxy_tensors is not None
            hidden_states = pp_proxy_tensors["hidden_states"]
            residual = pp_proxy_tensors["residual"]
2791

2792
2793
2794
2795
2796
2797
2798
2799
2800
2801
        normal_start_layer = self.start_layer
        normal_end_layer = self.end_layer
        if forward_batch.can_run_tbo:
            if (
                self.first_k_dense_replace > normal_start_layer
                and self.first_k_dense_replace < normal_end_layer
            ):
                normal_end_layer = self.first_k_dense_replace
            elif self.first_k_dense_replace < normal_start_layer:
                normal_end_layer = normal_start_layer = 0
2802

2803
        for i in range(normal_start_layer, normal_end_layer):
2804
2805
2806
            with get_global_expert_distribution_recorder().with_current_layer(i):
                layer = self.layers[i]
                hidden_states, residual = layer(
2807
2808
2809
2810
2811
2812
                    positions,
                    hidden_states,
                    forward_batch,
                    residual,
                    zero_allocator,
                    gemm_output_zero_allocator,
2813
                )
2814

2815
        if normal_end_layer != self.end_layer:
2816
            hidden_states, residual = model_forward_maybe_tbo(
2817
                layers=self.layers[normal_end_layer : self.end_layer],
2818
2819
2820
2821
2822
                enable_tbo=True,
                positions=positions,
                forward_batch=forward_batch,
                hidden_states=hidden_states,
                residual=residual,
2823
                input_data_scatter_mode=self.layers[
2824
                    normal_end_layer - 1
2825
                ].layer_scatter_modes.layer_output_mode,
2826
2827
2828
                zero_allocator=zero_allocator,
            )

2829
2830
2831
2832
2833
2834
2835
2836
2837
2838
2839
2840
2841
        if not self.pp_group.is_last_rank:
            return PPProxyTensors(
                {
                    "hidden_states": hidden_states,
                    "residual": residual,
                }
            )
        else:
            if not forward_batch.forward_mode.is_idle():
                if residual is None:
                    hidden_states = self.norm(hidden_states)
                else:
                    hidden_states, _ = self.norm(hidden_states, residual)
Liangsheng Yin's avatar
Liangsheng Yin committed
2842
2843
2844
2845
        return hidden_states


class DeepseekV2ForCausalLM(nn.Module):
2846
2847
    # for quark model load
    packed_modules_mapping = {}
Liangsheng Yin's avatar
Liangsheng Yin committed
2848
2849
2850
2851
2852

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
2853
        prefix: str = "",
Liangsheng Yin's avatar
Liangsheng Yin committed
2854
2855
    ) -> None:
        super().__init__()
2856
2857
2858
2859
2860
2861
2862
2863
2864
2865
2866
2867

        # for quark model load
        # Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None
        self.fuse_qkv_a_proj = (
            hasattr(config, "q_lora_rank") and config.q_lora_rank is not None
        )
        if self.fuse_qkv_a_proj:
            self.packed_modules_mapping["fused_qkv_a_proj_with_mqa"] = [
                "q_a_proj",
                "kv_a_proj_with_mqa",
            ]

2868
        self.pp_group = get_pp_group()
Liangsheng Yin's avatar
Liangsheng Yin committed
2869
        self.config = config
2870
        self.tp_size = get_tensor_model_parallel_world_size()
Liangsheng Yin's avatar
Liangsheng Yin committed
2871
        self.quant_config = quant_config
2872
        self.determine_num_fused_shared_experts()
2873
2874
2875
2876
2877
2878
2879
2880
        self.model = DeepseekV2Model(
            config, quant_config, prefix=add_prefix("model", prefix)
        )
        self.lm_head = ParallelLMHead(
            config.vocab_size,
            config.hidden_size,
            quant_config=quant_config,
            prefix=add_prefix("lm_head", prefix),
2881
            use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
2882
2883
2884
        )
        self.logits_processor = LogitsProcessor(config)

2885
2886
2887
2888
2889
2890
2891
2892
2893
2894
2895
2896
        self._routed_experts_weights_of_layer = LazyValue(
            lambda: {
                layer_id: layer.mlp.get_moe_weights()
                for layer_id, layer in enumerate(self.model.layers)
                if isinstance(layer.mlp, DeepseekV2MoE)
            }
        )

    @property
    def routed_experts_weights_of_layer(self):
        return self._routed_experts_weights_of_layer.value

2897
    def determine_num_fused_shared_experts(
2898
2899
        self, architecture: str = "DeepseekV3ForCausalLM"
    ):
2900
        self.num_fused_shared_experts = 0
2901
        if get_global_server_args().disable_shared_experts_fusion:
2902
2903
2904
2905
2906
2907
            return

        # Only Deepseek V3/R1 can use shared experts fusion optimization now.
        disable_reason = None
        if (
            not _is_cuda
2908
            or torch.cuda.get_device_capability("cuda") < (8, 0)
2909
2910
2911
2912
            or self.config.architectures[0] != architecture
            or self.config.n_routed_experts != 256
            or self.config.n_shared_experts != 1
        ):
2913
            disable_reason = "Only Deepseek V3/R1 on NV-platform with capability >= 80 can use shared experts fusion optimization."
2914
2915
        elif get_moe_expert_parallel_world_size() > 1:
            disable_reason = "Deepseek V3/R1 can not use shared experts fusion optimization under expert parallelism."
2916
2917
        elif self.quant_config.get_name() == "w4afp8":
            disable_reason = "Deepseek V3/R1 W4AFP8 model uses different quant method for routed experts and shared experts."
2918
2919

        if disable_reason is not None:
2920
            get_global_server_args().disable_shared_experts_fusion = True
Cheng Wan's avatar
Cheng Wan committed
2921
            self.num_fused_shared_experts = 0
2922
2923
2924
2925
2926
2927
2928
            log_info_on_rank0(
                logger,
                f"{disable_reason} Shared experts fusion optimization is disabled.",
            )
            return

        self.num_fused_shared_experts = self.config.n_shared_experts
2929

Mick's avatar
Mick committed
2930
2931
2932
    def get_input_embeddings(self) -> nn.Embedding:
        return self.model.embed_tokens

2933
    @torch.no_grad()
Liangsheng Yin's avatar
Liangsheng Yin committed
2934
2935
2936
2937
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
2938
        forward_batch: ForwardBatch,
2939
        input_embeds: torch.Tensor = None,
2940
        pp_proxy_tensors: Optional[PPProxyTensors] = None,
Liangsheng Yin's avatar
Liangsheng Yin committed
2941
    ) -> torch.Tensor:
2942
2943
        hidden_states = self.model(
            input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
2944
        )
Liangsheng Yin's avatar
Liangsheng Yin committed
2945

2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
        if self.pp_group.is_last_rank:
            return self.logits_processor(
                input_ids, hidden_states, self.lm_head, forward_batch
            )
        else:
            return hidden_states

    @property
    def start_layer(self):
        return self.model.start_layer

    @property
    def end_layer(self):
        return self.model.end_layer

2961
    def post_load_weights(self, is_nextn=False, weight_names=None):
inkcherry's avatar
inkcherry committed
2962
2963

        # Perform post-processing after loading weights
2964
2965
2966
2967
        if is_nextn:
            layer_ids = [self.config.num_hidden_layers]
        else:
            if weight_names is None:
2968
                layer_ids = range(self.model.start_layer, self.model.end_layer)
2969
2970
2971
2972
2973
            else:
                layer_ids = set()
                for name in weight_names:
                    if "kv_b_proj" in name:
                        layer_id = int(name.split(".")[2])
2974
                        if layer_id < self.config.num_hidden_layers:
2975
2976
                            layer_ids.add(layer_id)

2977
2978
2979
2980
2981
2982
        for layer_id in layer_ids:
            self_attn = (
                self.model.layers[layer_id].self_attn
                if not is_nextn
                else self.model.decoder.self_attn
            )
Baizhou Zhang's avatar
Baizhou Zhang committed
2983
2984
            if hasattr(self_attn.kv_b_proj, "qweight"):
                # AWQ compatible
2985
                if _is_cuda or _is_hip:
Baizhou Zhang's avatar
Baizhou Zhang committed
2986
2987
2988
2989
2990
                    w = awq_dequantize(
                        self_attn.kv_b_proj.qweight,
                        self_attn.kv_b_proj.scales,
                        self_attn.kv_b_proj.qzeros,
                    ).T
inkcherry's avatar
inkcherry committed
2991
                else:
Baizhou Zhang's avatar
Baizhou Zhang committed
2992
2993
2994
2995
2996
2997
2998
2999
3000
3001
3002
3003
                    w = awq_dequantize(
                        self_attn.kv_b_proj.qweight,
                        self_attn.kv_b_proj.scales,
                        self_attn.kv_b_proj.qzeros,
                        0,
                        0,
                        0,
                    ).T
            else:
                w = self_attn.kv_b_proj.weight
            # NOTE(HandH1998): Since `bmm_fp8` only supports per-tensor scale, we have to requantize `self_attn.kv_b_proj`.
            # This may affect the accuracy of fp8 model.
3004
3005
3006
            # Fix deepseek v3 blockwise bmm by using deep_gemm
            use_deep_gemm_bmm = False

Baizhou Zhang's avatar
Baizhou Zhang committed
3007
3008
3009
3010
            if w.dtype in (
                torch.float8_e4m3fn,
                torch.float8_e4m3fnuz,
            ):
3011
3012
3013
3014
                if (
                    hasattr(self.quant_config, "weight_block_size")
                    and self.quant_config.weight_block_size is not None
                ):
Baizhou Zhang's avatar
Baizhou Zhang committed
3015
                    weight_block_size = self.quant_config.weight_block_size
3016
3017
3018
3019
3020
3021
3022
3023
3024
3025
3026
3027
3028
3029
3030
3031
                    assert hasattr(self_attn.kv_b_proj, "weight_scale_inv")
                    if _is_fp8_fnuz:
                        weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
                            weight=w,
                            weight_scale=self_attn.kv_b_proj.weight_scale_inv,
                            input_scale=None,
                        )
                    else:
                        weight = w
                        weight_scale = self_attn.kv_b_proj.weight_scale_inv

                    if (
                        _is_cuda
                        and weight_block_size[0] == 128
                        and weight_block_size[1] == 128
                    ):
3032
3033
3034
3035
                        if (
                            deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
                            and not deep_gemm_wrapper.DEEPGEMM_BLACKWELL
                            and get_bool_env_var("SGL_USE_DEEPGEMM_BMM", "false")
3036
                        ):
3037
3038
                            block_scale = weight_scale
                            use_deep_gemm_bmm = True
3039
                        else:
3040
3041
3042
3043
                            w = block_quant_dequant(
                                weight,
                                weight_scale,
                                weight_block_size,
3044
                                torch.bfloat16,
3045
                            )
3046
3047
3048
3049
3050
                    else:
                        w, scale = block_quant_to_tensor_quant(
                            weight, weight_scale, weight_block_size
                        )
                        self_attn.w_scale = scale
Baizhou Zhang's avatar
Baizhou Zhang committed
3051
                else:
3052
3053
3054
3055
3056
3057
3058
3059
3060
3061
                    if _is_fp8_fnuz:
                        weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
                            weight=w,
                            weight_scale=self_attn.kv_b_proj.weight_scale,
                            input_scale=None,
                        )
                    else:
                        weight = w
                        weight_scale = self_attn.kv_b_proj.weight_scale

Baizhou Zhang's avatar
Baizhou Zhang committed
3062
3063
3064
3065
3066
3067
3068
3069
3070
3071
3072
3073
3074
3075
3076
3077
3078
3079
3080
                    w, scale = channel_quant_to_tensor_quant(weight, weight_scale)
                    self_attn.w_scale = scale

            if w.dtype == torch.int8:
                if hasattr(self.quant_config, "weight_block_size"):
                    # block-wise int8 need it
                    weight_block_size = self.quant_config.weight_block_size
                    if weight_block_size is not None:
                        assert hasattr(self_attn.kv_b_proj, "weight_scale_inv")
                        weight = w
                        weight_scale = self_attn.kv_b_proj.weight_scale_inv
                        w = int8_block_dequant(
                            weight, weight_scale, weight_block_size
                        ).to(torch.bfloat16)
                else:
                    # channel-wise int8 need it
                    w = w.to(torch.bfloat16) * self_attn.kv_b_proj.weight_scale.to(
                        torch.bfloat16
                    )
3081

Baizhou Zhang's avatar
Baizhou Zhang committed
3082
3083
3084
            w_kc, w_vc = w.unflatten(
                0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
            ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)
3085

3086
3087
3088
3089
3090
            if (
                _use_aiter_gfx95
                and self.quant_config is not None
                and self.quant_config.get_name() == "quark"
            ):
3091
3092
3093
3094
                w_kc, self_attn.w_scale_k, w_vc, self_attn.w_scale_v = (
                    quark_post_load_weights(self_attn, w, "mxfp4")
                )

3095
            if not use_deep_gemm_bmm:
3096
3097
3098
3099
3100
3101
                self_attn.w_kc = bind_or_assign(
                    self_attn.w_kc, w_kc.transpose(1, 2).contiguous().transpose(1, 2)
                )
                self_attn.w_vc = bind_or_assign(
                    self_attn.w_vc, w_vc.contiguous().transpose(1, 2)
                )
3102
3103
3104
3105
                if (
                    hasattr(self_attn.kv_b_proj, "weight_scale")
                    and self_attn.w_scale is None
                ):
3106
3107
3108
                    self_attn.w_scale = bind_or_assign(
                        self_attn.w_scale, self_attn.kv_b_proj.weight_scale
                    )
3109
3110
                    if _is_hip:
                        self_attn.w_scale *= 2.0
3111
3112
3113
3114
3115
3116
3117
3118
                # TODO: remove this after adding FP8 support in bmm cpu kernel
                if _is_cpu and _is_cpu_amx_available and w.dtype == torch.float8_e4m3fn:
                    self_attn.w_kc = (
                        self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale
                    )
                    self_attn.w_vc = (
                        self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale
                    )
3119
3120
3121
3122
3123
3124
            else:
                num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1]
                num_tiles_n = self_attn.v_head_dim // weight_block_size[0]
                ws_kc, ws_vc = block_scale.unflatten(
                    0, (-1, (num_tiles_k + num_tiles_n))
                ).split([num_tiles_k, num_tiles_n], dim=1)
3125
3126
3127
3128
3129
3130
3131
3132
3133
3134
                self_attn.w_scale_k = bind_or_assign(
                    self_attn.w_scale_k, ws_kc.transpose(1, 2).contiguous()
                )
                self_attn.w_scale_v = bind_or_assign(
                    self_attn.w_scale_v, ws_vc.contiguous()
                )
                self_attn.w_kc = bind_or_assign(
                    self_attn.w_kc, w_kc.transpose(1, 2).contiguous()
                )
                self_attn.w_vc = bind_or_assign(self_attn.w_vc, w_vc.contiguous())
3135
                self_attn.use_deep_gemm_bmm = True
inkcherry's avatar
inkcherry committed
3136

3137
3138
3139
        if (
            deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
            and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
3140
3141
            and hasattr(self.quant_config, "weight_block_size")
            and self.quant_config.weight_block_size is not None
3142
        ):
3143
            self._weight_requant_ue8m0(is_nextn)
3144

3145
3146
3147
3148
3149
3150
3151
        # TODO can move weight_requant_ue8m0 and transform_scale_ue8m0 into Fp8LinearMethod.process_weights_after_loading
        if (
            deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
            and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
            and get_bool_env_var("SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN")
        ):
            self._transform_scale_ue8m0(is_nextn)
3152
3153
3154
        if is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config):
            self._transform_scale_nextn_moe_ue8m0()

3155
    def _weight_requant_ue8m0(self, is_nextn=False):
3156
3157
3158
3159
3160
3161
3162
3163
3164
3165
        weight_block_size = self.quant_config.weight_block_size

        moe_layers = list(
            range(
                self.config.first_k_dense_replace,
                self.config.num_hidden_layers,
                self.config.moe_layer_freq,
            )
        )

3166
        num_hidden_layers = 1 if is_nextn else self.config.num_hidden_layers
3167

3168
3169
3170
3171
3172
        for layer_id in range(num_hidden_layers):
            if is_nextn:
                layer = self.model.decoder
            else:
                layer = self.model.layers[layer_id]
3173

3174
            module_list = [
3175
3176
                layer.self_attn.kv_b_proj,
                layer.self_attn.o_proj,
3177
3178
3179
3180
3181
3182
3183
3184
3185
3186
            ]

            if self.config.q_lora_rank is not None:
                module_list.append(layer.self_attn.fused_qkv_a_proj_with_mqa)
                module_list.append(layer.self_attn.q_b_proj)
            else:
                module_list.append(layer.self_attn.kv_a_proj_with_mqa)
                module_list.append(layer.self_attn.q_proj)

            for module in module_list:
3187
3188
3189
3190
                requant_weight_ue8m0_inplace(
                    module.weight, module.weight_scale_inv, weight_block_size
                )

3191
            if layer_id in moe_layers or is_nextn:
3192
3193
3194
3195
3196
3197
3198
3199
3200
                shared_experts = getattr(layer.mlp, "shared_experts", None)
                if shared_experts is not None:
                    for module in [
                        shared_experts.gate_up_proj,
                        shared_experts.down_proj,
                    ]:
                        requant_weight_ue8m0_inplace(
                            module.weight, module.weight_scale_inv, weight_block_size
                        )
3201
3202
3203
3204
3205
3206
3207
3208
3209
3210
3211
3212
3213
3214
3215
3216
3217
3218
3219

                experts = layer.mlp.experts
                if isinstance(experts, DeepEPMoE):
                    for w in [
                        experts.w13_weight_fp8,
                        experts.w2_weight_fp8,
                    ]:
                        requant_weight_ue8m0_inplace(w[0], w[1], weight_block_size)
            else:
                mlp = layer.mlp
                assert isinstance(mlp, DeepseekV2MLP)
                for module in [
                    mlp.gate_up_proj,
                    mlp.down_proj,
                ]:
                    requant_weight_ue8m0_inplace(
                        module.weight, module.weight_scale_inv, weight_block_size
                    )

3220
3221
3222
3223
3224
3225
3226
3227
3228
3229
3230
3231
3232
3233
3234
3235
3236
3237
3238
    # TODO can move weight_requant_ue8m0 and transform_scale_ue8m0 into Fp8LinearMethod.process_weights_after_loading
    def _transform_scale_ue8m0(self, is_nextn=False):
        num_hidden_layers = 1 if is_nextn else self.config.num_hidden_layers

        for layer_id in range(num_hidden_layers):
            if is_nextn:
                layer = self.model.decoder
            else:
                layer = self.model.layers[layer_id]

            module_list = []
            if self.config.q_lora_rank is not None:
                module_list.append(layer.self_attn.q_b_proj)

            for module in module_list:
                transform_scale_ue8m0_inplace(
                    module.weight_scale_inv, mn=module.weight.shape[-2]
                )

3239
3240
3241
3242
3243
3244
3245
3246
3247
3248
3249
3250
3251
3252
3253
3254
3255
3256
3257
3258
3259
3260
    # TODO avoid code dup (currently combine from weight_requant_ue8m0 and transform_scale_ue8m0)
    def _transform_scale_nextn_moe_ue8m0(self):
        layer = self.model.decoder

        shared_experts = getattr(layer.mlp, "shared_experts", None)
        if shared_experts is not None:
            for module in [
                shared_experts.gate_up_proj,
                shared_experts.down_proj,
            ]:
                transform_scale_ue8m0_inplace(
                    module.weight_scale_inv, mn=module.weight.shape[-2]
                )

        experts = layer.mlp.experts
        if isinstance(experts, DeepEPMoE):
            for w in [
                experts.w13_weight_fp8,
                experts.w2_weight_fp8,
            ]:
                transform_scale_ue8m0_inplace(w[1], mn=w[0].shape[-2])

3261
    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False):
3262

3263
3264
3265
        if is_nextn:
            if hasattr(self.config, "num_nextn_predict_layers"):
                num_nextn_layers = self.config.num_nextn_predict_layers
3266
                assert num_nextn_layers == 1, "Only 1 nextn layer is supported"
3267
3268
3269
3270
3271
3272
3273
3274
3275
                # compatible with old design
                nextn_layer_id = (
                    0
                    if self.config.num_hidden_layers == 1
                    else self.config.num_hidden_layers
                )
            else:
                raise ValueError("num_nextn_predict_layers is not in the config")

3276
3277
        if get_bool_env_var("SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN"):
            weights = self._quant_attn_to_fp8_ue8m0(weights, is_nextn=is_nextn)
3278
3279
3280
3281
3282
        if is_nextn and enable_nextn_moe_bf16_cast_to_fp8(self.quant_config):
            weights = self._quant_nextn_moe_to_fp8_ue8m0(
                weights, nextn_layer_id=nextn_layer_id
            )

Liangsheng Yin's avatar
Liangsheng Yin committed
3283
3284
3285
3286
3287
3288
3289
3290
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            ("gate_up_proj", "gate_proj", 0),
            ("gate_up_proj", "up_proj", 1),
        ]

        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
3291
        expert_params_mapping = FusedMoE.make_expert_params_mapping(
Liangsheng Yin's avatar
Liangsheng Yin committed
3292
3293
3294
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
3295
            num_experts=self.config.n_routed_experts + self.num_fused_shared_experts,
Liangsheng Yin's avatar
Liangsheng Yin committed
3296
        )
3297
3298
3299
        # Params for special naming rules in mixed-precision models, for example:
        # model.layers.xx.mlp.experts.xx.w1.input_scale. For details,
        # see https://huggingface.co/Barrrrry/DeepSeek-R1-W4AFP8/blob/main.
3300
        if self.quant_config and self.quant_config.get_name() == "w4afp8":
3301
3302
            expert_params_mapping += FusedMoE.make_expert_input_scale_params_mapping(
                num_experts=self.config.n_routed_experts
3303
            )
Liangsheng Yin's avatar
Liangsheng Yin committed
3304

3305
3306
3307
3308
3309
3310
        # Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None
        fuse_qkv_a_proj = hasattr(self.config, "q_lora_rank") and (
            self.config.q_lora_rank is not None
        )
        cached_a_proj = {} if fuse_qkv_a_proj else None

3311
3312
3313
3314
3315
3316
3317
3318
3319
        if is_nextn:
            nextn_layer_prefix = f"model.layers.{nextn_layer_id}"
            nextn_spec_weight_names = [
                "shared_head.norm",
                "eh_proj",
                "enorm",
                "hnorm",
            ]

3320
3321
        if self.num_fused_shared_experts > 0:
            assert self.num_fused_shared_experts == 1
3322
            log_info_on_rank0(logger, "Shared experts fusion optimization enabled.")
3323

3324
3325
3326
3327
3328
        with concurrent.futures.ThreadPoolExecutor() as executor:
            futures = []
            params_dict = dict(self.named_parameters())
            weight_names = []
            for name, loaded_weight in weights:
3329
3330
3331
3332
3333
3334
3335
3336
3337
3338
                layer_id = get_layer_id(name)
                if (
                    layer_id is not None
                    and hasattr(self.model, "start_layer")
                    and (
                        layer_id < self.model.start_layer
                        or layer_id >= self.model.end_layer
                    )
                ):
                    continue
3339
3340
3341
3342
3343
                if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name:
                    name = name.replace(
                        "mlp.shared_experts",
                        f"mlp.experts.{self.config.n_routed_experts}",
                    )
3344

3345
                weight_names.append(name)
3346

3347
3348
3349
3350
3351
3352
3353
3354
3355
3356
3357
3358
3359
                if not is_nextn:
                    if hasattr(self.config, "num_nextn_predict_layers"):
                        num_nextn_layers = self.config.num_nextn_predict_layers
                        if num_nextn_layers > 0 and name.startswith("model.layers"):
                            name_list = name.split(".")
                            if (
                                len(name_list) >= 3
                                and int(name_list[2]) >= self.config.num_hidden_layers
                            ):
                                continue
                else:
                    if not name.startswith(nextn_layer_prefix):
                        continue
3360

3361
3362
3363
                    # Use shared head and embed weights from target model
                    if "shared_head.head" in name or "embed_tokens" in name:
                        continue
3364

3365
3366
3367
3368
3369
3370
3371
3372
3373
3374
3375
3376
                    is_decoder = True
                    # For nextn specific weights
                    for weight_name in nextn_spec_weight_names:
                        if weight_name in name:
                            name = name.replace(nextn_layer_prefix, "model")
                            is_decoder = False
                            break
                    # For decoder layer weights
                    if is_decoder:
                        name = name.replace(nextn_layer_prefix, "model.decoder")

                if "rotary_emb.inv_freq" in name:
Liangsheng Yin's avatar
Liangsheng Yin committed
3377
                    continue
3378
3379
                for param_name, weight_name, shard_id in stacked_params_mapping:
                    # Skip non-stacked layers and experts (experts handled below).
Liangsheng Yin's avatar
Liangsheng Yin committed
3380
3381
                    if weight_name not in name:
                        continue
3382
3383
3384
3385
3386
3387
3388
3389
                    # We have mlp.experts[0].gate_proj in the checkpoint.
                    # Since we handle the experts below in expert_params_mapping,
                    # we need to skip here BEFORE we update the name, otherwise
                    # name will be updated to mlp.experts[0].gate_up_proj, which
                    # will then be updated below in expert_params_mapping
                    # for mlp.experts[0].gate_gate_up_proj, which breaks load.
                    if ("mlp.experts." in name) and name not in params_dict:
                        continue
Liangsheng Yin's avatar
Liangsheng Yin committed
3390
                    name = name.replace(weight_name, param_name)
3391
3392
3393
                    # Skip loading extra bias for GPTQ models.
                    if name.endswith(".bias") and name not in params_dict:
                        continue
Liangsheng Yin's avatar
Liangsheng Yin committed
3394
3395
                    param = params_dict[name]
                    weight_loader = param.weight_loader
3396
3397
                    futures.append(
                        executor.submit(weight_loader, param, loaded_weight, shard_id)
Liangsheng Yin's avatar
Liangsheng Yin committed
3398
3399
3400
                    )
                    break
                else:
3401
3402
3403
3404
3405
3406
3407
3408
3409
3410
3411
3412
3413
3414
3415
3416
                    for mapping in expert_params_mapping:
                        param_name, weight_name, expert_id, shard_id = mapping
                        if weight_name not in name:
                            continue
                        name = name.replace(weight_name, param_name)
                        param = params_dict[name]
                        weight_loader = param.weight_loader
                        futures.append(
                            executor.submit(
                                weight_loader,
                                param,
                                loaded_weight,
                                name,
                                shard_id=shard_id,
                                expert_id=expert_id,
                            )
3417
                        )
3418
3419
3420
3421
3422
                        break
                    else:
                        # Skip loading extra bias for GPTQ models.
                        if name.endswith(".bias") and name not in params_dict:
                            continue
3423
3424
3425
3426
3427
3428
                        # Skip loading embed_tokens if not first rank in pipeline parallelism
                        if ".embed_tokens." in name and not self.pp_group.is_first_rank:
                            continue
                        # Skip loading norm if not last rank in pipeline parallelism
                        if ".norm." in name and not self.pp_group.is_last_rank:
                            continue
3429
3430
                        if fuse_qkv_a_proj and (
                            "q_a_proj" in name or "kv_a_proj_with_mqa" in name
3431
                        ):
3432
3433
3434
                            cached_a_proj[name] = loaded_weight
                            q_a_proj_name = (
                                name
3435
                                if "q_a_proj" in name
3436
3437
3438
3439
3440
3441
                                else name.replace("kv_a_proj_with_mqa", "q_a_proj")
                            )
                            kv_a_proj_name = (
                                name
                                if "kv_a_proj_with_mqa" in name
                                else name.replace("q_a_proj", "kv_a_proj_with_mqa")
3442
3443
                            )

3444
3445
3446
3447
3448
3449
3450
3451
3452
3453
                            # When both q_a_proj and kv_a_proj_with_mqa has been cached, load the fused weight to parameter
                            if (
                                q_a_proj_name in cached_a_proj
                                and kv_a_proj_name in cached_a_proj
                            ):
                                q_a_proj_weight = cached_a_proj[q_a_proj_name]
                                kv_a_proj_weight = cached_a_proj[kv_a_proj_name]
                                cat_dim = 0
                                if self.quant_config is not None and (
                                    self.quant_config.get_name() == "awq"
3454
                                    or self.quant_config.get_name() == "awq_marlin"
3455
3456
3457
3458
3459
3460
3461
3462
3463
3464
3465
3466
3467
3468
3469
3470
3471
3472
3473
3474
3475
3476
3477
3478
3479
3480
3481
3482
3483
3484
3485
3486
3487
3488
3489
3490
3491
3492
3493
3494
3495
3496
3497
3498
                                    or self.quant_config.get_name() == "moe_wna16"
                                ):
                                    cat_dim = 1
                                fused_weight = torch.cat(
                                    [q_a_proj_weight, kv_a_proj_weight], dim=cat_dim
                                )
                                param_name = (
                                    name.replace(
                                        "q_a_proj", "fused_qkv_a_proj_with_mqa"
                                    )
                                    if "q_a_proj" in name
                                    else name.replace(
                                        "kv_a_proj_with_mqa",
                                        "fused_qkv_a_proj_with_mqa",
                                    )
                                )
                                param = params_dict[param_name]

                                weight_loader = getattr(
                                    param, "weight_loader", default_weight_loader
                                )
                                futures.append(
                                    executor.submit(weight_loader, param, fused_weight)
                                )
                                cached_a_proj.pop(q_a_proj_name)
                                cached_a_proj.pop(kv_a_proj_name)
                        else:
                            if (
                                "k_scale" in name or "v_scale" in name
                            ) and name not in params_dict:
                                # modelopt attn kv scale is named differently
                                for scale in ["k_scale", "v_scale"]:
                                    if scale in name:
                                        name = name.replace(
                                            f"{scale[0]}_proj", "attn_mqa"
                                        )
                                        break
                            if name not in params_dict:
                                # modelopt ckpt contains not needed weights for MTP module:
                                # model.decoder.self_attn.attn_mqa.v_scale and
                                # model.decoder.self_attn.attn_mqa.k_scale
                                logger.warning(f"{name} not found in params_dict.")
                                continue
                            param = params_dict[name]
3499
3500
3501
                            weight_loader = getattr(
                                param, "weight_loader", default_weight_loader
                            )
3502
3503
3504
3505
3506
3507
3508
                            futures.append(
                                executor.submit(weight_loader, param, loaded_weight)
                            )

            # Wait for all tasks to complete and raise any exceptions.
            for future in concurrent.futures.as_completed(futures):
                future.result()
Liangsheng Yin's avatar
Liangsheng Yin committed
3509

3510
        self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names)
Ke Bao's avatar
Ke Bao committed
3511

3512
3513
3514
3515
3516
3517
3518
3519
3520
3521
3522
3523
3524
3525
3526
3527
3528
3529
3530
3531
3532
3533
3534
3535
    def _quant_attn_to_fp8_ue8m0(self, weights, is_nextn):
        weights_dict = dict(weights)

        # temporarily only support DeepSeek V3/R1
        weight_block_size = [128, 128]

        for layer_id in trange(
            self.config.num_hidden_layers + int(is_nextn),
            desc="quant attn to fp8 ue8m0",
        ):
            for stem in [
                # may put tensors like `o_proj` here for DeepSeek FP4 ckpt v1
                "q_b_proj",
            ]:
                partial_name = f"model.layers.{layer_id}.self_attn.{stem}"
                original_weight = weights_dict[f"{partial_name}.weight"]
                out_w, out_s = quant_weight_ue8m0(
                    original_weight, weight_block_size=weight_block_size
                )
                weights_dict[f"{partial_name}.weight"] = out_w
                weights_dict[f"{partial_name}.weight_scale_inv"] = out_s

        return list(weights_dict.items())

3536
3537
3538
3539
3540
3541
3542
3543
3544
3545
3546
3547
3548
3549
3550
3551
3552
3553
3554
3555
3556
3557
3558
3559
3560
3561
3562
3563
3564
3565
3566
3567
    # TODO avoid code dup
    def _quant_nextn_moe_to_fp8_ue8m0(self, weights, nextn_layer_id: int):
        weights_dict = dict(weights)

        # temporarily only support DeepSeek V3/R1
        weight_block_size = [128, 128]

        for layer_id in [nextn_layer_id]:
            for expert_sub_name in [
                "shared_experts",
                *[
                    f"experts.{expert_id}"
                    for expert_id in range(self.config.n_routed_experts)
                ],
            ]:
                for stem in [
                    "gate_proj",
                    "up_proj",
                    "down_proj",
                ]:
                    partial_name = (
                        f"model.layers.{layer_id}.mlp.{expert_sub_name}.{stem}"
                    )
                    original_weight = weights_dict[f"{partial_name}.weight"]
                    out_w, out_s = quant_weight_ue8m0(
                        original_weight, weight_block_size=weight_block_size
                    )
                    weights_dict[f"{partial_name}.weight"] = out_w
                    weights_dict[f"{partial_name}.weight_scale_inv"] = out_s

        return list(weights_dict.items())

3568
3569
3570
3571
3572
3573
3574
3575
3576
3577
3578
    def get_embed_and_head(self):
        return self.model.embed_tokens.weight, self.lm_head.weight

    def set_embed_and_head(self, embed, head):
        del self.model.embed_tokens.weight
        del self.lm_head.weight
        self.model.embed_tokens.weight = embed
        self.lm_head.weight = head
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

3579
3580
3581
3582
3583
3584
3585
3586
    @classmethod
    def get_model_config_for_expert_location(cls, config):
        return ModelConfigForExpertLocation(
            num_layers=config.num_hidden_layers,
            num_logical_experts=config.n_routed_experts,
            num_groups=config.n_group,
        )

Liangsheng Yin's avatar
Liangsheng Yin committed
3587

fzyzcjy's avatar
fzyzcjy committed
3588
3589
3590
3591
3592
3593
3594
3595
AttentionBackendRegistry.register("ascend", handle_attention_ascend)
AttentionBackendRegistry.register("flashinfer", handle_attention_flashinfer)
AttentionBackendRegistry.register("fa3", handle_attention_fa3)
AttentionBackendRegistry.register("flashmla", handle_attention_flashmla)
AttentionBackendRegistry.register("cutlass_mla", handle_attention_cutlass_mla)
AttentionBackendRegistry.register("fa4", handle_attention_fa4)
AttentionBackendRegistry.register("trtllm_mla", handle_attention_trtllm_mla)
AttentionBackendRegistry.register("aiter", handle_attention_aiter)
fzyzcjy's avatar
fzyzcjy committed
3596
AttentionBackendRegistry.register("nsa", handle_attention_nsa)
fzyzcjy's avatar
fzyzcjy committed
3597
AttentionBackendRegistry.register("triton", handle_attention_triton)
3598
3599


HandH1998's avatar
HandH1998 committed
3600
3601
3602
3603
class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
    pass


fzyzcjy's avatar
fzyzcjy committed
3604
3605
3606
3607
3608
class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM):
    pass


EntryClass = [DeepseekV2ForCausalLM, DeepseekV3ForCausalLM, DeepseekV32ForCausalLM]