# 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. # ============================================================================== # Adapted from: # https://github.com/vllm-project/vllm/blob/fb6af8bc086328ca6659e72d11ffd4309ce4de22/vllm/model_executor/models/deepseek_v2.py """Inference-only DeepseekV2 model.""" from __future__ import annotations import logging from contextlib import nullcontext from typing import Any, Dict, Iterable, List, Optional, Tuple, Union import torch import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig from sglang.srt.batch_overlap.single_batch_overlap import SboFlags, compute_overlap_args from sglang.srt.batch_overlap.two_batch_overlap import ( MaybeTboDeepEPDispatcher, model_forward_maybe_tbo, ) from sglang.srt.configs.model_config import ( compute_mla_mscale_scaling, get_nsa_index_head_dim, get_nsa_index_n_heads, get_nsa_index_topk, is_deepseek_nsa, ) from sglang.srt.distributed import ( divide, get_moe_expert_parallel_world_size, get_pp_group, get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather, tensor_model_parallel_all_reduce, ) from sglang.srt.environ import envs 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 from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.activation import SiluAndMul from sglang.srt.layers.amx_utils import PackWeightMethod from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer from sglang.srt.layers.attention.nsa.utils import ( can_cp_split, cp_all_gather_rerange_output, cp_split_and_rebuild_data, cp_split_and_rebuild_position, is_nsa_enable_prefill_cp, nsa_use_prefill_cp, prepare_input_dp_with_cp_dsa, ) from sglang.srt.layers.communicator import ( LayerCommunicator, LayerScatterModes, enable_moe_dense_fully_dp, get_attn_tp_context, ) from sglang.srt.layers.communicator_nsa_cp import NSACPLayerCommunicator from sglang.srt.layers.dp_attention import ( get_attention_cp_rank, get_attention_cp_size, get_attention_tp_rank, get_attention_tp_size, is_dp_attention_enabled, ) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, ReplicatedLinear, RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import ( get_moe_a2a_backend, get_moe_runner_backend, should_use_flashinfer_cutlass_moe_fp4_allgather, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod from sglang.srt.layers.moe.token_dispatcher.base import ( BaseDispatcher, CombineInput, DispatchOutput, ) from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat from sglang.srt.layers.moe.utils import ( RoutingMethodType, filter_moe_weight_param_global_expert, ) from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.fp8 import Fp8Config from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.rotary_embedding import get_rope_wrapper from sglang.srt.layers.utils import PPMissingLayer from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.models.deepseek_common.attention_backend_handler import ( AttentionBackendRegistry, ) from sglang.srt.models.deepseek_common.attention_forward_methods import ( AttnForwardMethod, DeepseekMHAForwardMixin, DeepseekMLACpuForwardMixin, DeepseekMLAForwardMixin, DeepseekMLARocmForwardMixin, ) from sglang.srt.models.deepseek_common.deepseek_weight_loader import ( DeepseekV2WeightLoaderMixin, ) from sglang.srt.models.deepseek_common.utils import ( _device_sm, _get_llama_4_scaling, _is_cpu, _is_cpu_amx_available, _is_cuda, _is_gfx95_supported, _is_hip, _is_npu, _use_aiter, _use_aiter_gfx95, ) from sglang.srt.server_args import get_global_server_args from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils import ( BumpAllocator, LazyValue, add_prefix, is_non_idle_and_non_empty, log_info_on_rank0, make_layers, use_intel_amx_backend, ) if _use_aiter_gfx95: from sglang.srt.layers.rocm_linear_utils import ( aiter_dsv3_router_gemm, get_dsv3_gemm_output_zero_allocator_size, ) if _use_aiter: pass if _is_cuda: from sgl_kernel import dsv3_fused_a_gemm, dsv3_router_gemm elif _is_npu: from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import ( forward_dsa_core_npu, forward_dsa_prepare_npu, forward_mha_core_npu, forward_mha_prepare_npu, forward_mla_core_npu, forward_mla_prepare_npu, ) else: pass logger = logging.getLogger(__name__) 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, prefix: str = "", tp_rank: Optional[int] = None, tp_size: Optional[int] = None, ) -> None: super().__init__() self.tp_size = tp_size self.gate_up_proj = MergedColumnParallelLinear( hidden_size, [intermediate_size] * 2, bias=False, quant_config=quant_config, prefix=add_prefix("gate_up_proj", prefix), tp_rank=tp_rank, tp_size=tp_size, ) self.down_proj = RowParallelLinear( intermediate_size, hidden_size, bias=False, quant_config=quant_config, reduce_results=reduce_results, prefix=add_prefix("down_proj", prefix), tp_rank=tp_rank, tp_size=tp_size, ) if not hasattr(self.gate_up_proj, "weight") and hasattr( self.gate_up_proj, "weight_packed" ): self.gate_up_proj.weight = self.gate_up_proj.weight_packed if not hasattr(self.down_proj, "weight") and hasattr( self.down_proj, "weight_packed" ): self.down_proj.weight = self.down_proj.weight_packed if hidden_act != "silu": raise ValueError( f"Unsupported activation: {hidden_act}. " "Only silu is supported for now." ) self.act_fn = SiluAndMul() def forward( self, x, forward_batch=None, should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, gemm_output_zero_allocator: BumpAllocator = None, ): if (self.tp_size == 1) and x.shape[0] == 0: return x if ( gemm_output_zero_allocator is not None and x.shape[0] <= 256 and self.gate_up_proj.weight.dtype == torch.uint8 ): 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) gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) x, _ = self.down_proj( x, skip_all_reduce=should_allreduce_fusion or use_reduce_scatter, ) return x class MoEGate(nn.Module): def __init__( self, config, quant_config, prefix: str = "", is_nextn: bool = False, ): super().__init__() self.is_nextn = is_nextn self.weight = nn.Parameter( torch.empty((config.n_routed_experts, config.hidden_size)) ) if config.topk_method == "noaux_tc": correction_bias_dtype = torch.float32 if quant_config is not None: if ( quant_config.get_name() == "modelopt_fp4" and get_moe_runner_backend().is_flashinfer_trtllm() ): correction_bias_dtype = torch.bfloat16 elif _use_aiter and quant_config.get_name() in ( "fp8", "compressed_tensors", "quark", ): correction_bias_dtype = torch.bfloat16 self.e_score_correction_bias = nn.Parameter( torch.empty((config.n_routed_experts), dtype=correction_bias_dtype) ) else: self.e_score_correction_bias = None if _is_cpu and _is_cpu_amx_available: self.quant_method = PackWeightMethod(weight_names=["weight"]) self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() def forward( self, hidden_states, gemm_output_zero_allocator: BumpAllocator = None, forward_batch: ForwardBatch = None, ): if use_intel_amx_backend(self): return torch.ops.sgl_kernel.weight_packed_linear( hidden_states, self.weight, None, # bias True, # is_vnni ) if get_global_server_args().enable_deterministic_inference: return F.linear(hidden_states, self.weight, None) if forward_batch is not None and nsa_use_prefill_cp(forward_batch): logits = F.linear(hidden_states, self.weight, None) else: # NOTE: For some unknown reason, router_gemm seems degrade accept length. if ( _is_cuda and hidden_states.shape[0] <= 16 and hidden_states.shape[1] == 7168 and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384) and _device_sm >= 90 ): # router gemm output float32 logits = dsv3_router_gemm( hidden_states, self.weight, out_dtype=torch.float32 ) elif _use_aiter_gfx95 and hidden_states.shape[0] <= 256: logits = aiter_dsv3_router_gemm( hidden_states, self.weight, gemm_output_zero_allocator ) else: logits = F.linear(hidden_states, self.weight, None) return logits class DeepseekV2MoE(nn.Module): def __init__( self, config: PretrainedConfig, layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, is_nextn: bool = False, ): super().__init__() self.tp_size = get_tensor_model_parallel_world_size() self.moe_ep_size = get_moe_expert_parallel_world_size() self.routed_scaling_factor = config.routed_scaling_factor self.n_shared_experts = config.n_shared_experts self.num_fused_shared_experts = ( 0 if get_global_server_args().disable_shared_experts_fusion else config.n_shared_experts ) self.config = config self.layer_id = layer_id self.alt_stream = alt_stream self.is_nextn = is_nextn 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." ) self.gate = MoEGate( config=config, quant_config=quant_config, prefix=add_prefix("gate", prefix), is_nextn=is_nextn, ) # scaling factor for fused shared experts on AMD-platform. fused_shared_experts_scaling_factor = None if self.moe_ep_size > 1 and self.num_fused_shared_experts > 0: # if enable_ep_moe tp_szie == ep_size, every gpu get shared experts gemm output # so we scale with 1 / self.moe_ep_size in ep mode which will make it equalation as in tp mode # with fused_shared_experts fused_shared_experts_scaling_factor = 1.0 / float(self.moe_ep_size) self.experts = get_moe_impl_class(quant_config)( num_experts=config.n_routed_experts + self.num_fused_shared_experts + get_global_server_args().ep_num_redundant_experts, num_fused_shared_experts=self.num_fused_shared_experts, top_k=config.num_experts_per_tok + self.num_fused_shared_experts, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, layer_id=self.layer_id, quant_config=quant_config, routed_scaling_factor=self.routed_scaling_factor, routing_method_type=getattr( config, "routing_method_type", RoutingMethodType.DeepSeekV3 ), prefix=add_prefix("experts", prefix), ) self.topk = TopK( top_k=config.num_experts_per_tok + self.num_fused_shared_experts, layer_id=self.layer_id, 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, correction_bias=self.gate.e_score_correction_bias, quant_config=quant_config, routed_scaling_factor=self.routed_scaling_factor, apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk, fused_shared_experts_scaling_factor=fused_shared_experts_scaling_factor, # 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 (except trtllm). We use quant_config to determine the output format. output_format=( TopKOutputFormat.STANDARD if (quant_config is None) and (not get_moe_runner_backend().is_flashinfer_trtllm()) else None ), ) self.shared_experts_is_int8 = False self.shared_experts_is_fp8 = False self.shared_experts_weight_block_size = None if config.n_shared_experts is not None and self.num_fused_shared_experts == 0: intermediate_size = config.moe_intermediate_size * config.n_shared_experts # disable tp for shared experts when enable deepep moe, or with fp4 allgather 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) if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_nixl() or get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_ascend_fuseep() or get_moe_a2a_backend().is_flashinfer() or should_use_flashinfer_cutlass_moe_fp4_allgather() else {} ), ) 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", "awq_marlin", "moe_wna16", } self.shared_experts_is_int8 = ( not is_packed_weight and self.shared_experts.gate_up_proj.weight.dtype == torch.int8 ) self.shared_experts_is_fp8 = ( not is_packed_weight and self.shared_experts.gate_up_proj.weight.dtype == torch.float8_e4m3fn ) if self.shared_experts_is_fp8: if ( _use_aiter and config.quantization_config.get("quant_method") == "compressed-tensors" ): # For compressed-tensors ptpc model, don't need to check the weight_block_size pass else: 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 ) self.top_k = config.num_experts_per_tok if ( get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_nixl() or get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_ascend_fuseep() ): # TODO: we will support tp < ep in the future self.ep_size = get_moe_expert_parallel_world_size() self.num_experts = ( config.n_routed_experts + get_global_server_args().ep_num_redundant_experts ) 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 ) self._enable_a2a_moe = ( get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake() or get_moe_a2a_backend().is_nixl() or get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_ascend_fuseep() or get_moe_a2a_backend().is_flashinfer() ) self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo() def get_moe_weights(self): return [ x.data for name, x in self.experts.named_parameters() if name not in ["correction_bias"] and filter_moe_weight_param_global_expert( name, x, self.experts.num_local_experts ) ] def forward( self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch] = None, should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, gemm_output_zero_allocator: BumpAllocator = None, ) -> torch.Tensor: if not self._enable_a2a_moe: if ( self.alt_stream is not None and self.num_fused_shared_experts == 0 and hidden_states.shape[0] > 0 and get_is_capture_mode() ): return self.forward_normal_dual_stream( hidden_states, should_allreduce_fusion, use_reduce_scatter, gemm_output_zero_allocator, ) else: return self.forward_normal( hidden_states, should_allreduce_fusion, use_reduce_scatter, gemm_output_zero_allocator, ) else: return self.forward_deepep(hidden_states, forward_batch) def forward_normal_dual_stream( self, hidden_states: torch.Tensor, should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, gemm_output_zero_allocator: BumpAllocator = None, ) -> torch.Tensor: current_stream = torch.cuda.current_stream() self.alt_stream.wait_stream(current_stream) shared_output = self._forward_shared_experts( hidden_states, gemm_output_zero_allocator ) server_args = get_global_server_args() dispatch_info = ( ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id) if server_args.enable_eplb else None ) with torch.cuda.stream(self.alt_stream): # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states, gemm_output_zero_allocator) topk_output = self.topk( hidden_states, router_logits, expert_location_dispatch_info=dispatch_info, ) final_hidden_states = self.experts(hidden_states, topk_output) if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod): final_hidden_states *= self.routed_scaling_factor current_stream.wait_stream(self.alt_stream) final_hidden_states += shared_output if ( self.tp_size > 1 and not should_allreduce_fusion and not use_reduce_scatter and not should_use_flashinfer_cutlass_moe_fp4_allgather() ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states def forward_normal( self, hidden_states: torch.Tensor, should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, gemm_output_zero_allocator: BumpAllocator = None, ) -> torch.Tensor: if hasattr(self, "shared_experts") and use_intel_amx_backend( self.shared_experts.gate_up_proj ): return self.forward_cpu(hidden_states, should_allreduce_fusion) server_args = get_global_server_args() dispatch_info = ( ExpertLocationDispatchInfo.init_new(layer_id=self.layer_id) if server_args.enable_eplb else None ) if hidden_states.shape[0] > 0: if ( not self._fuse_shared_experts_inside_sbo ): # TODO: check if it supports mtp shared_output = self._forward_shared_experts( hidden_states, gemm_output_zero_allocator ) # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states, gemm_output_zero_allocator) topk_output = self.topk( hidden_states, router_logits, expert_location_dispatch_info=dispatch_info, ) else: shared_output = None topk_output = self.topk.empty_topk_output(hidden_states.device) if self._fuse_shared_experts_inside_sbo: shared_output = None def _pre_combine_hook( dispatcher: BaseDispatcher, combine_input: CombineInput ): nonlocal shared_output self.alt_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(self.alt_stream): shared_output = self._forward_shared_experts( hidden_states, gemm_output_zero_allocator ) pre_combine_hook_handle.remove() def _post_combine_hook( dispatcher: BaseDispatcher, hidden_states: torch.Tensor ): nonlocal shared_output torch.cuda.current_stream().wait_stream(self.alt_stream) post_combine_hook_handle.remove() pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook( _pre_combine_hook ) post_combine_hook_handle = ( self.experts.dispatcher.register_post_combine_hook(_post_combine_hook) ) final_hidden_states = self.experts( hidden_states, topk_output, ) if ( not _is_cuda and not _use_aiter or isinstance(self.experts.quant_method, KTEPWrapperMethod) ): # fused in biased_grouped_topk so we can skip here final_hidden_states *= self.routed_scaling_factor if shared_output is not None: final_hidden_states += shared_output if ( self.tp_size > 1 and not should_allreduce_fusion and not use_reduce_scatter and not should_use_flashinfer_cutlass_moe_fp4_allgather() ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states def forward_cpu( self, hidden_states: torch.Tensor, should_allreduce_fusion: bool = False, ) -> torch.Tensor: # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states) topk_output = self.topk(hidden_states, router_logits) fused_experts_out = self.experts( hidden_states=hidden_states, topk_output=topk_output ) assert use_intel_amx_backend( self.shared_experts.gate_up_proj ) == use_intel_amx_backend(self.shared_experts.down_proj) # [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 True, # is_vnni ) if self.tp_size > 1 and not should_allreduce_fusion: final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states def forward_deepep( self, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: shared_output = None sbo_enabled_flag = self._fuse_shared_experts_inside_sbo and not self.is_nextn sbo_overlap_dispatch_flag = ( sbo_enabled_flag and SboFlags.enable_dispatch_shared_one_stream_overlap() ) sbo_overlap_combine_flag = ( sbo_enabled_flag and SboFlags.enable_combine_shared_two_stream_overlap() ) if hidden_states.shape[0] > 0: # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states, forward_batch=forward_batch) if not sbo_enabled_flag: if self.alt_stream is not None: self.alt_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(self.alt_stream): shared_output = self._forward_shared_experts(hidden_states) shared_output.record_stream(self.alt_stream) shared_event = self.alt_stream.record_event() else: shared_output = self._forward_shared_experts(hidden_states) topk_output = self.topk( hidden_states, router_logits, num_token_non_padded=forward_batch.num_token_non_padded, expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( layer_id=self.layer_id, ), ) else: topk_output = self.topk.empty_topk_output(hidden_states.device) if sbo_overlap_dispatch_flag: shared_output = None def _deepep_dispatch_hook(dispatcher: BaseDispatcher): nonlocal shared_output shared_output = self._forward_shared_experts(hidden_states) for handle in deepep_dispatch_hook_handle: handle.remove() def _post_dispatch_hook( dispatcher: BaseDispatcher, dispatch_output: DispatchOutput ): combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = ( compute_overlap_args(dispatch_output, self.alt_stream) ) dispatcher.set_overlap_args( combine_overlap_args=combine_overlap_args, meta_overlap_args=meta_overlap_args, ) self.experts.set_overlap_args( down_gemm_overlap_args=down_gemm_overlap_args, meta_overlap_args=meta_overlap_args, ) post_dispatch_hook_handle.remove() def _post_combine_hook( dispatcher: BaseDispatcher, hidden_states: torch.Tensor ): dispatcher.clear_overlap_args() self.experts.clear_overlap_args() post_combine_hook_handle.remove() assert isinstance(self.experts.dispatcher, MaybeTboDeepEPDispatcher) deepep_dispatch_hook_handle = ( self.experts.dispatcher.register_deepep_dispatch_hook( _deepep_dispatch_hook ) ) post_dispatch_hook_handle = ( self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook) ) post_combine_hook_handle = ( self.experts.dispatcher.register_post_combine_hook(_post_combine_hook) ) elif sbo_overlap_combine_flag: shared_output = None def _post_dispatch_hook( dispatcher: BaseDispatcher, dispatch_output: DispatchOutput ): combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = ( compute_overlap_args(dispatch_output, self.alt_stream) ) dispatcher.set_overlap_args( combine_overlap_args=combine_overlap_args, meta_overlap_args=meta_overlap_args, ) self.experts.set_overlap_args( down_gemm_overlap_args=down_gemm_overlap_args, meta_overlap_args=meta_overlap_args, ) post_dispatch_hook_handle.remove() def _pre_combine_hook( dispatcher: BaseDispatcher, combine_input: CombineInput ): nonlocal shared_output if ( e := dispatcher.meta_overlap_args.get("record_event_after_down") ) is not None: e.record() # TODO reduce sm for non-deepgemm with deep_gemm_wrapper.configure_deep_gemm_num_sms( dispatcher.meta_overlap_args["compute_num_sms"] ): shared_output = self._forward_shared_experts(hidden_states) pre_combine_hook_handle.remove() def _post_combine_hook( dispatcher: BaseDispatcher, hidden_states: torch.Tensor ): dispatcher.clear_overlap_args() self.experts.clear_overlap_args() post_combine_hook_handle.remove() post_dispatch_hook_handle = ( self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook) ) pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook( _pre_combine_hook ) post_combine_hook_handle = ( self.experts.dispatcher.register_post_combine_hook(_post_combine_hook) ) elif envs.SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO.get(): # On GB200: Shared experts overlapped on alt_stream, down gemm overlapped with DeepEP Combine def _post_dispatch_hook( dispatcher: BaseDispatcher, dispatch_output: DispatchOutput ): combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = ( compute_overlap_args(dispatch_output, self.alt_stream) ) dispatcher.set_overlap_args( combine_overlap_args=combine_overlap_args, meta_overlap_args=meta_overlap_args, ) self.experts.set_overlap_args( down_gemm_overlap_args=down_gemm_overlap_args, meta_overlap_args=meta_overlap_args, ) post_dispatch_hook_handle.remove() def _pre_combine_hook( dispatcher: BaseDispatcher, combine_input: CombineInput ): if ( e := dispatcher.meta_overlap_args.get("record_event_after_down") ) is not None: e.record() pre_combine_hook_handle.remove() def _post_combine_hook( dispatcher: BaseDispatcher, hidden_states: torch.Tensor ): dispatcher.clear_overlap_args() self.experts.clear_overlap_args() post_combine_hook_handle.remove() post_dispatch_hook_handle = ( self.experts.dispatcher.register_post_dispatch_hook(_post_dispatch_hook) ) pre_combine_hook_handle = self.experts.dispatcher.register_pre_combine_hook( _pre_combine_hook ) post_combine_hook_handle = ( self.experts.dispatcher.register_post_combine_hook(_post_combine_hook) ) final_hidden_states = self.experts( hidden_states=hidden_states, topk_output=topk_output, ) if ( hidden_states.shape[0] > 0 and not sbo_enabled_flag and self.alt_stream is not None ): torch.cuda.current_stream().wait_event(shared_event) if shared_output is not None: x = shared_output # aiter moe call will handle routed_scaling_factor in the function # so add _use_aiter condition to eliminate to use self.routed_scaling_factor in add_ call if self.experts.should_fuse_routed_scaling_factor_in_topk or _use_aiter: x.add_(final_hidden_states) else: x.add_(final_hidden_states, alpha=self.routed_scaling_factor) final_hidden_states = x else: if not ( self.experts.should_fuse_routed_scaling_factor_in_topk or _use_aiter ): final_hidden_states *= self.routed_scaling_factor return final_hidden_states def _forward_shared_experts( self, hidden_states, gemm_output_zero_allocator: BumpAllocator = None ): if (hidden_states.shape[0] > 0) and (self.num_fused_shared_experts == 0): return self.shared_experts( hidden_states, gemm_output_zero_allocator=gemm_output_zero_allocator ) else: return None def op_gate(self, state): if is_non_idle_and_non_empty( state.forward_batch.forward_mode, state.hidden_states_mlp_input ): # router_logits: (num_tokens, n_experts) state.router_logits = self.gate(state.hidden_states_mlp_input) else: state.router_logits = None def op_shared_experts(self, state): hidden_states_mlp_input = state.pop("hidden_states_mlp_input") if (self.num_fused_shared_experts == 0) and is_non_idle_and_non_empty( state.forward_batch.forward_mode, hidden_states_mlp_input ): state.shared_output = self.shared_experts(hidden_states_mlp_input) else: state.shared_output = None def op_select_experts(self, state): router_logits = state.pop("router_logits") hidden_states = state.hidden_states_mlp_input if router_logits is not None: with get_global_expert_distribution_recorder().with_current_layer( self.layer_id ): state.topk_output = self.topk( 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, ), ) else: state.topk_output = self.topk.empty_topk_output(hidden_states.device) def op_dispatch_a(self, state): if self.ep_size > 1: self.experts.dispatcher.dispatch_a( hidden_states=state.hidden_states_mlp_input, topk_output=state.pop("topk_output"), tbo_subbatch_index=state.get("tbo_subbatch_index"), ) def op_dispatch_b(self, state): if self.ep_size > 1: with get_global_expert_distribution_recorder().with_current_layer( self.layer_id ): state.dispatch_output = self.experts.dispatcher.dispatch_b( tbo_subbatch_index=state.get("tbo_subbatch_index"), ) def op_experts(self, state): state.combine_input = self.experts.run_moe_core( dispatch_output=state.dispatch_output, ) def op_combine_a(self, state): if self.ep_size > 1: self.experts.dispatcher.combine_a( combine_input=state.pop("combine_input"), tbo_subbatch_index=state.get("tbo_subbatch_index"), ) state.pop("dispatch_output") def op_combine_b(self, state): if self.ep_size > 1: state.hidden_states_after_combine = self.experts.dispatcher.combine_b( tbo_subbatch_index=state.get("tbo_subbatch_index"), ) def op_output(self, state): final_hidden_states = state.pop("hidden_states_after_combine") if get_moe_a2a_backend().is_mori(): num_tokens = state.pop("num_tokens") final_hidden_states = final_hidden_states[:num_tokens] if (shared_output := state.pop("shared_output")) is not None: x = shared_output if _use_aiter: x.add_(final_hidden_states) else: x.add_(final_hidden_states, alpha=self.routed_scaling_factor) final_hidden_states = x elif _use_aiter: # fused in aiter_biased_grouped_topk so we can skip here pass else: final_hidden_states *= self.routed_scaling_factor state.hidden_states_mlp_output = final_hidden_states class DeepseekV2AttentionMLA( nn.Module, DeepseekMHAForwardMixin, DeepseekMLAForwardMixin, DeepseekMLARocmForwardMixin, DeepseekMLACpuForwardMixin, ): 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, reduce_results: bool = True, layer_id: int = None, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, skip_rope: bool = False, ) -> 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 self.quant_config = quant_config attn_tp_rank = get_attention_tp_rank() attn_tp_size = get_attention_tp_size() self.use_nsa = is_deepseek_nsa(config) self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() if self.nsa_enable_prefill_cp: assert self.use_nsa, "CP currently only supports deepseek v3.2 model" # cp reuse the attn_tp comm group but need to duplicate the weights if self.nsa_enable_prefill_cp and self.use_nsa: self.cp_size = get_attention_cp_size() self.num_heads = num_heads assert num_heads % attn_tp_size == 0 self.num_local_heads = num_heads // attn_tp_size self.scaling = self.qk_head_dim**-0.5 self.rope_theta = rope_theta self.max_position_embeddings = max_position_embeddings self.kv_cache_dtype = get_global_server_args().kv_cache_dtype # 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" # For tensor parallel attention if self.q_lora_rank is not None: self.fused_qkv_a_proj_with_mqa = ReplicatedLinear( self.hidden_size, self.q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim, bias=False, quant_config=quant_config, prefix=add_prefix("fused_qkv_a_proj_with_mqa", prefix), ) self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) self.q_b_proj = ColumnParallelLinear( q_lora_rank, self.num_heads * self.qk_head_dim, bias=False, quant_config=self._get_q_b_proj_quant_config(quant_config), prefix=add_prefix("q_b_proj", prefix), tp_rank=attn_tp_rank, tp_size=attn_tp_size, ) else: self.q_proj = ColumnParallelLinear( self.hidden_size, self.num_heads * self.qk_head_dim, bias=False, quant_config=quant_config, prefix=add_prefix("q_proj", prefix), tp_rank=attn_tp_rank, tp_size=attn_tp_size, ) 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), ) if self.use_nsa: is_neox_style = not getattr(config, "indexer_rope_interleave", False) 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, is_neox_style=is_neox_style, prefix=add_prefix("indexer", prefix), quant_config=quant_config, layer_id=layer_id, alt_stream=alt_stream, ) 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, ) self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) if not skip_rope: is_neox_style = not getattr(config, "rope_interleave", True) self.rotary_emb = get_rope_wrapper( 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=is_neox_style, device=get_global_server_args().device, ) if rope_scaling and rope_scaling.get("apply_yarn_scaling", True): self.scaling = compute_mla_mscale_scaling(rope_scaling, self.scaling) else: self.rotary_emb = None self.use_deepseek_yarn_rope = rope_scaling is not None self.attn_mqa = RadixAttention( 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, quant_config=quant_config, prefix=add_prefix("attn_mqa", prefix), ) 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, quant_config=quant_config, prefix=add_prefix("attn_mha", prefix), ) self.alt_stream = alt_stream self.attn_mha.kv_b_proj = None self.w_kc = None self.w_vc = None self.w_scale = 1.0 self.w_scale_k = None self.w_scale_v = None self.use_deep_gemm_bmm = False self.current_attention_backend = ( None # Attention backend used by current forward batch ) self.has_fused_proj = hasattr(self, "fused_qkv_a_proj_with_mqa") self.is_packed_weight = ( self.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() in {"awq", "awq_marlin", "moe_wna16"} ) self.use_min_latency_fused_a_gemm = ( self.has_fused_proj and not self.is_packed_weight 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 and _is_cuda and 90 <= _device_sm < 120 ) self.init_mha_forward() self.init_mla_forward() self.init_mla_fused_rope_rocm_forward() self.init_mla_fused_rope_cpu_forward() def dispatch_attn_forward_method( self, forward_batch: ForwardBatch ) -> AttnForwardMethod: # Determine attention backend used by current forward batch if forward_batch.forward_mode.is_decode_or_idle(): attention_backend = get_global_server_args().decode_attention_backend 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) if get_global_server_args().speculative_attention_mode == "decode": attention_backend = get_global_server_args().decode_attention_backend else: # default to prefill attention_backend = get_global_server_args().prefill_attention_backend else: attention_backend = get_global_server_args().prefill_attention_backend self.current_attention_backend = attention_backend handler = AttentionBackendRegistry.get_handler(attention_backend) return handler(self, forward_batch) 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") ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, zero_allocator: BumpAllocator, layer_scatter_modes: LayerScatterModes = None, llama_4_scaling: Optional[torch.Tensor] = None, ): s = self.forward_prepare( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, zero_allocator=zero_allocator, layer_scatter_modes=layer_scatter_modes, llama_4_scaling=llama_4_scaling, ) return self.forward_core(s) def forward_prepare( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, zero_allocator: BumpAllocator, layer_scatter_modes: LayerScatterModes = None, llama_4_scaling: Optional[torch.Tensor] = None, ): if self.attn_mha.kv_b_proj is None: self.attn_mha.kv_b_proj = self.kv_b_proj # when hidden_states is a tuple of tensors, the tuple will include quantized weight and scale tensor if isinstance(hidden_states, tuple): if ( not get_attn_tp_context().input_scattered and 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 ( not get_attn_tp_context().input_scattered and 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 attn_forward_method = self.dispatch_attn_forward_method(forward_batch) if attn_forward_method == AttnForwardMethod.MHA: inner_state = self.forward_normal_prepare( positions, hidden_states, forward_batch, zero_allocator ) elif attn_forward_method == AttnForwardMethod.MHA_CHUNKED_KV: inner_state = self.forward_normal_chunked_kv_prepare( positions, hidden_states, forward_batch, zero_allocator ) elif attn_forward_method == AttnForwardMethod.MHA_ONE_SHOT: inner_state = self.forward_normal_one_shot_prepare( positions, hidden_states, forward_batch, zero_allocator ) elif attn_forward_method == AttnForwardMethod.MLA: inner_state = self.forward_absorb_prepare( positions, hidden_states, forward_batch, zero_allocator, llama_4_scaling ) elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_ROCM: inner_state = self.forward_absorb_fused_mla_rope_prepare( positions, hidden_states, forward_batch, zero_allocator ) 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 ) elif attn_forward_method == AttnForwardMethod.MHA_NPU: inner_state = forward_mha_prepare_npu( self, positions, hidden_states, forward_batch, zero_allocator, layer_scatter_modes, ) elif attn_forward_method == AttnForwardMethod.MLA_NPU: inner_state = forward_mla_prepare_npu( self, positions, hidden_states, forward_batch, zero_allocator, layer_scatter_modes, ) elif attn_forward_method == AttnForwardMethod.DSA_NPU: inner_state = forward_dsa_prepare_npu( self, positions, hidden_states, forward_batch, zero_allocator, layer_scatter_modes, ) else: raise NotImplementedError return None, attn_forward_method, forward_batch, inner_state 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.MHA_ONE_SHOT: return self.forward_normal_one_shot_core(*inner_state) elif attn_forward_method == AttnForwardMethod.MLA: return self.forward_absorb_core(*inner_state) elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_ROCM: return self.forward_absorb_fused_mla_rope_core(*inner_state) elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_CPU: return self.forward_absorb_fused_mla_rope_cpu_core(*inner_state) elif attn_forward_method == AttnForwardMethod.MHA_NPU: return forward_mha_core_npu(self, *inner_state) elif attn_forward_method == AttnForwardMethod.MLA_NPU: return forward_mla_core_npu(self, *inner_state) elif attn_forward_method == AttnForwardMethod.DSA_NPU: return forward_dsa_core_npu(self, *inner_state) else: raise NotImplementedError def prepare_qkv_latent( self, hidden_states: torch.Tensor, forward_batch: ForwardBatch ): assert self.q_lora_rank is not None if ( (not isinstance(hidden_states, tuple)) and hidden_states.shape[0] >= 1 and hidden_states.shape[0] <= 16 and self.use_min_latency_fused_a_gemm ): qkv_latent = dsv3_fused_a_gemm( hidden_states, self.fused_qkv_a_proj_with_mqa.weight.T ) else: qkv_latent = self.fused_qkv_a_proj_with_mqa(hidden_states)[0] return qkv_latent def rebuild_cp_kv_cache(self, latent_cache, forward_batch, k_nope, k_pe): # support allgather+rerrange latent_cache[..., : self.kv_lora_rank] = k_nope.squeeze(1) latent_cache[..., self.kv_lora_rank :] = k_pe.squeeze(1) latent_cache_output = cp_all_gather_rerange_output( latent_cache.contiguous(), self.cp_size, forward_batch, torch.cuda.current_stream(), ) k_nope = latent_cache_output[..., : self.kv_lora_rank].unsqueeze(1) k_pe = latent_cache_output[..., self.kv_lora_rank :].unsqueeze(1) return k_nope, k_pe @staticmethod def _get_q_b_proj_quant_config(quant_config): if envs.SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN.get(): # refer to real DeepSeek V3 quant config return Fp8Config( is_checkpoint_fp8_serialized=True, weight_block_size=[128, 128], ) else: return quant_config class DeepseekV2DecoderLayer(nn.Module): def __init__( self, config: PretrainedConfig, layer_id: int, quant_config: Optional[QuantizationConfig] = None, moe_quant_config_override: Optional[QuantizationConfig] = None, is_nextn: bool = False, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, ) -> None: super().__init__() self.hidden_size = config.hidden_size self.config = config if hasattr(config, "rope_parameters"): rope_theta = config.rope_parameters["rope_theta"] assert rope_theta is not None, f"rope_theta not found in config: {config}" rope_type = config.rope_parameters.get("rope_type") rope_scaling = config.rope_parameters if rope_type != "default" else None else: rope_theta = config.rope_theta rope_scaling = config.rope_scaling max_position_embeddings = config.max_position_embeddings self.speculative_algorithm = SpeculativeAlgorithm.from_string( get_global_server_args().speculative_algorithm ) self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() self.layer_id = layer_id self.is_nextn = is_nextn 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), alt_stream=alt_stream, ) if not hasattr(config, "q_lora_rank") and envs.SGLANG_USE_AG_AFTER_QLORA.get(): raise ValueError( "SGLANG_USE_AG_AFTER_QLORA only supports the model with q_lora_rank" ) 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) is_next_layer_sparse = self._is_layer_sparse(layer_id + 1, is_nextn=False) self.layer_scatter_modes = LayerScatterModes.init_new( layer_id=layer_id, num_layers=1 if is_nextn else config.num_hidden_layers, is_layer_sparse=self.is_layer_sparse, is_previous_layer_sparse=is_previous_layer_sparse, is_next_layer_sparse=is_next_layer_sparse, ) if self.is_layer_sparse: self.mlp = DeepseekV2MoE( config=config, quant_config=moe_quant_config_override or quant_config, prefix=add_prefix("mlp", prefix), layer_id=self.layer_id, alt_stream=alt_stream, is_nextn=is_nextn, ) else: if enable_moe_dense_fully_dp(): mlp_tp_rank, mlp_tp_size = 0, 1 else: mlp_tp_rank, mlp_tp_size = None, None self.mlp = DeepseekV2MLP( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, prefix=add_prefix("mlp", prefix), tp_rank=mlp_tp_rank, tp_size=mlp_tp_size, ) 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 ) if self.nsa_enable_prefill_cp: self.layer_communicator = NSACPLayerCommunicator( layer_scatter_modes=self.layer_scatter_modes, input_layernorm=self.input_layernorm, post_attention_layernorm=self.post_attention_layernorm, allow_reduce_scatter=True, is_last_layer=( is_nextn or (self.layer_id == self.config.num_hidden_layers - 1) ), qkv_latent_func=self.self_attn.prepare_qkv_latent, ) else: self.layer_communicator = LayerCommunicator( layer_scatter_modes=self.layer_scatter_modes, input_layernorm=self.input_layernorm, post_attention_layernorm=self.post_attention_layernorm, allow_reduce_scatter=True, is_last_layer=( is_nextn or (self.layer_id == self.config.num_hidden_layers - 1) ), qkv_latent_func=self.self_attn.prepare_qkv_latent, ) 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 ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, residual: Optional[torch.Tensor], zero_allocator: BumpAllocator, gemm_output_zero_allocator: BumpAllocator = None, llama_4_scaling: Optional[torch.Tensor] = None, ) -> torch.Tensor: quant_format = ( "mxfp4" if ( _is_gfx95_supported 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 ) else ( "fp8" if ( _is_gfx95_supported 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 == getattr(torch, "float8_e4m3fn", None) ) else "" ) ) hidden_states, residual = self.layer_communicator.prepare_attn( hidden_states, residual, forward_batch, quant_format, ) hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, zero_allocator=zero_allocator, llama_4_scaling=llama_4_scaling, layer_scatter_modes=self.layer_scatter_modes, ) hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual, forward_batch ) should_allreduce_fusion = ( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( forward_batch ) ) # 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 ) if isinstance(self.mlp, DeepseekV2MLP): gemm_output_zero_allocator = None hidden_states = self.mlp( hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter, gemm_output_zero_allocator, ) if not self.nsa_enable_prefill_cp and should_allreduce_fusion: hidden_states._sglang_needs_allreduce_fusion = True if not should_allreduce_fusion: hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual, forward_batch ) return hidden_states, residual def op_comm_prepare_attn( self, state, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, residual: Optional[torch.Tensor], zero_allocator: BumpAllocator, tbo_subbatch_index: Optional[int] = None, ): state.hidden_states_after_comm_pre_attn, state.residual_after_input_ln = ( self.layer_communicator.prepare_attn(hidden_states, residual, forward_batch) ) if get_moe_a2a_backend().is_mori(): state.num_tokens = hidden_states.shape[0] state.update( dict( forward_batch=forward_batch, positions=positions, zero_allocator=zero_allocator, tbo_subbatch_index=tbo_subbatch_index, ) ) 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, ) ) 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( hidden_states, state.forward_batch ) else: state.hidden_states_mlp_output = hidden_states def op_comm_postprocess_layer(self, state): hidden_states, residual = self.layer_communicator.postprocess_layer( state.pop("hidden_states_mlp_output"), state.pop("residual_after_comm_pre_mlp"), state.forward_batch, ) 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 class DeepseekV2Model(nn.Module): fall_back_to_pt_during_load = False def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.padding_id = config.pad_token_id self.vocab_size = config.vocab_size self.first_k_dense_replace = config.first_k_dense_replace self.pp_group = get_pp_group() self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() if self.nsa_enable_prefill_cp: self.cp_size = get_attention_cp_size() else: self.cp_size = None if self.pp_group.is_first_rank: self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, use_attn_tp_group=is_dp_attention_enabled(), ) else: self.embed_tokens = PPMissingLayer() self.alt_stream = ( torch.cuda.Stream() if _is_cuda or envs.SGLANG_NPU_USE_MULTI_STREAM.get() else None ) 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), 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", # only for nvfp4 *( [ "w13_blockscale_swizzled", "w2_blockscale_swizzled", ] if hasattr(module, "w13_blockscale_swizzled") else [] ), ] if isinstance(module, FusedMoE) else [] ), ), ) 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) 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): # tp_size = get_tensor_model_parallel_world_size() a2a_backend = get_moe_a2a_backend() is_a2a_moe = ( a2a_backend.is_deepep() or a2a_backend.is_mori() or a2a_backend.is_mooncake() ) tp_size = ( 1 if is_a2a_moe else get_tensor_model_parallel_world_size() ) intermediate_size = ( config.moe_intermediate_size * config.n_shared_experts ) share_expert_output_size_per_partition = divide( intermediate_size * 2, tp_size ) allocate_size = share_expert_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, ) ) self.layers_to_capture = [] if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake(): self.enable_a2a_moe = True else: self.enable_a2a_moe = False # llama_4_scaling: for supporting Mistral-Large-3 model self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None) def get_input_embeddings(self) -> torch.Tensor: return self.embed_tokens def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> Union[torch.Tensor, PPProxyTensors]: total_num_layers = self.end_layer - self.start_layer 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 else: assert pp_proxy_tensors is not None hidden_states = pp_proxy_tensors["hidden_states"] residual = pp_proxy_tensors["residual"] device = hidden_states.device zero_allocator = BumpAllocator( buffer_size=total_num_layers * 2 * (2 if forward_batch.can_run_tbo else 1), dtype=torch.float32, device=device, ) 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 ) if nsa_use_prefill_cp(forward_batch): if self.pp_group.is_first_rank: hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) positions = cp_split_and_rebuild_position(forward_batch, positions) # llama_4_scaling: for supporting Mistral-Large-3 model # Compute llama 4 scaling once per forward pass if enabled llama_4_scaling: Optional[torch.Tensor] = None if self.llama_4_scaling_config is not None: llama_4_scaling = _get_llama_4_scaling( original_max_position_embeddings=self.llama_4_scaling_config[ "original_max_position_embeddings" ], scaling_beta=self.llama_4_scaling_config["beta"], positions=positions, ) 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 aux_hidden_states = [] for i in range(normal_start_layer, normal_end_layer): # NOTE: torch dynamo does not support graph break in context manager ctx = ( nullcontext() if not get_global_server_args().disable_piecewise_cuda_graph else get_global_expert_distribution_recorder().with_current_layer(i) ) with ctx: if i in self.layers_to_capture: if self.enable_a2a_moe and i > self.first_k_dense_replace: aux_hidden_state = tensor_model_parallel_all_gather( hidden_states + residual, dim=0 ) aux_hidden_states.append(aux_hidden_state) else: aux_hidden_states.append(hidden_states + residual) layer = self.layers[i] hidden_states, residual = layer( positions, hidden_states, forward_batch, residual, zero_allocator, gemm_output_zero_allocator, llama_4_scaling, ) if normal_end_layer != self.end_layer: hidden_states, residual = model_forward_maybe_tbo( layers=self.layers[normal_end_layer : self.end_layer], enable_tbo=True, positions=positions, forward_batch=forward_batch, hidden_states=hidden_states, residual=residual, input_data_scatter_mode=self.layers[ normal_end_layer - 1 ].layer_scatter_modes.layer_output_mode, zero_allocator=zero_allocator, ) 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) if self.pp_group.is_last_rank and nsa_use_prefill_cp(forward_batch): # allgather + rerrange hidden_states = cp_all_gather_rerange_output( hidden_states, self.cp_size, forward_batch, torch.cuda.current_stream(), ) if len(aux_hidden_states) == 0: return hidden_states return hidden_states, aux_hidden_states class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin): # for quark model load packed_modules_mapping = {} def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() # 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", ] self.pp_group = get_pp_group() self.config = config self.tp_size = get_tensor_model_parallel_world_size() self.quant_config = quant_config self.determine_num_fused_shared_experts() self.use_nsa = is_deepseek_nsa(config) self.model = DeepseekV2Model( config, quant_config, prefix=add_prefix("model", prefix) ) if self.pp_group.is_last_rank: if self.pp_group.world_size == 1 and config.tie_word_embeddings: self.lm_head = self.model.embed_tokens else: self.lm_head = ParallelLMHead( config.vocab_size, config.hidden_size, quant_config=quant_config, prefix=add_prefix("lm_head", prefix), use_attn_tp_group=get_global_server_args().enable_dp_lm_head, ) else: # ranks other than the last rank will have a placeholder layer self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config) 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) } ) self.capture_aux_hidden_states = False self.nsa_enable_prefill_cp = is_nsa_enable_prefill_cp() if self.nsa_enable_prefill_cp: self.cp_rank = get_attention_cp_rank() self.cp_size = get_attention_cp_size() else: self.cp_rank = self.cp_size = None q_lora_rank = config.q_lora_rank if hasattr(config, "q_lora_rank") else None get_attn_tp_context().init_context(q_lora_rank, is_deepseek_nsa(config)) @property def routed_experts_weights_of_layer(self): return self._routed_experts_weights_of_layer.value def determine_num_fused_shared_experts( self, architecture: str = "DeepseekV3ForCausalLM" ): self.num_fused_shared_experts = 0 if get_global_server_args().disable_shared_experts_fusion: return # Only Deepseek V3/R1 can use shared experts fusion optimization now. disable_reason = None if ( self.config.architectures[0] != architecture or self.config.n_routed_experts != 256 or self.config.n_shared_experts != 1 ): disable_reason = "Config does not support fused shared expert(s)." elif (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and ( not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4) ): disable_reason = ( "Only Deepseek V3/R1 on NV-platform with capability >= 80 " "or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization." ) elif get_moe_expert_parallel_world_size() > 1 and ( not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4) ): disable_reason = "Only Deepseek V3/R1 on AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization under expert parallelism." elif disable_reason is None and ( get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori() ): disable_reason = "Deepseek V3/R1 cannot use shared experts fusion optimization under deepep expert parallelism." elif self.quant_config and self.quant_config.get_name() == "w4afp8": disable_reason = "Deepseek V3/R1 W4AFP8 model uses different quant method for routed experts and shared experts." if disable_reason is not None: get_global_server_args().disable_shared_experts_fusion = True self.num_fused_shared_experts = 0 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 def get_input_embeddings(self) -> nn.Embedding: return self.model.embed_tokens @torch.no_grad() def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: if self.nsa_enable_prefill_cp: if can_cp_split(len(input_ids), self.cp_size, self.use_nsa, forward_batch): forward_batch.nsa_cp_metadata = prepare_input_dp_with_cp_dsa( len(input_ids), self.cp_rank, self.cp_size, forward_batch.seq_lens_cpu.tolist(), ) with get_attn_tp_context().maybe_input_scattered(forward_batch): hidden_states = self.model( input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors ) aux_hidden_states = None if self.capture_aux_hidden_states: hidden_states, aux_hidden_states = hidden_states if self.pp_group.is_last_rank: return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states ) else: return hidden_states @property def start_layer(self): return self.model.start_layer @property def end_layer(self): return self.model.end_layer def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False): self.do_load_weights(weights, is_nextn) 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() @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, ) def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None): if not self.pp_group.is_last_rank: return if layer_ids is None: self.capture_aux_hidden_states = True num_layers = self.config.num_hidden_layers self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3] else: self.capture_aux_hidden_states = True # we plus 1 here because in sglang, for the ith layer, it takes the output # of the (i-1)th layer as aux hidden state self.model.layers_to_capture = [val + 1 for val in layer_ids] class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM): pass class DeepseekV32ForCausalLM(DeepseekV2ForCausalLM): pass EntryClass = [DeepseekV2ForCausalLM, DeepseekV3ForCausalLM, DeepseekV32ForCausalLM]