# 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. # ============================================================================== """Inference-only DeepSeek NextN Speculative Decoding.""" import logging from typing import Iterable, Optional, Tuple import torch from torch import nn from transformers import PretrainedConfig from sglang.srt.configs.model_config import is_deepseek_nsa from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size from sglang.srt.environ import envs from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers.attention.nsa.utils import ( can_cp_split, cp_all_gather_rerange_output, cp_split_and_rebuild_data, is_nsa_enable_prefill_cp, nsa_use_prefill_cp, prepare_input_dp_with_cp_dsa, ) from sglang.srt.layers.dp_attention import ( get_attention_cp_rank, get_attention_cp_size, is_dp_attention_enabled, ) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.quantization import Fp8Config from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.models.deepseek_common.utils import enable_nextn_moe_bf16_cast_to_fp8 from sglang.srt.models.deepseek_v2 import DeepseekV2DecoderLayer, DeepseekV3ForCausalLM from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import BumpAllocator, add_prefix, is_cuda, is_npu logger = logging.getLogger(__name__) _is_cuda = is_cuda() _is_npu = is_npu() class DeepseekModelNextN(nn.Module): def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() if enable_nextn_moe_bf16_cast_to_fp8(quant_config): # refer to real DeepSeek V3 quant config moe_quant_config_override = Fp8Config( is_checkpoint_fp8_serialized=True, weight_block_size=[128, 128], ) else: moe_quant_config_override = None if quant_config is not None and quant_config.get_name() == "modelopt_fp4": logger.warning( "Overriding DeepseekV3ForCausalLMNextN quant config for modelopt_fp4 Deepseek model." ) quant_config = None self.vocab_size = config.vocab_size self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, use_attn_tp_group=is_dp_attention_enabled(), prefix=add_prefix("embed_tokens", prefix), ) self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) self.alt_stream = ( torch.cuda.Stream() if _is_cuda or envs.SGLANG_NPU_USE_MULTI_STREAM.get() else None ) layer_name = "decoder" if _is_npu and ( get_global_server_args().speculative_draft_model_path == get_global_server_args().model_path ): layer_name = "layers." + str(config.num_hidden_layers) self.decoder = DeepseekV2DecoderLayer( config, 0, quant_config=quant_config, moe_quant_config_override=moe_quant_config_override, is_nextn=True, prefix=add_prefix(layer_name, prefix), alt_stream=self.alt_stream, ) self.shared_head = nn.Module() self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) 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 def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, ) -> torch.Tensor: zero_allocator = BumpAllocator( buffer_size=2, dtype=torch.float32, device=( input_embeds.device if input_embeds is not None else input_ids.device ), ) if input_embeds is None: hidden_states = self.embed_tokens(input_ids) else: hidden_states = input_embeds if hidden_states.shape[0] > 0: hidden_states = self.eh_proj( torch.cat( ( self.enorm(hidden_states), self.hnorm(forward_batch.spec_info.hidden_states), ), dim=-1, ) ) if nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp): hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) residual = None with get_global_expert_distribution_recorder().disable_this_region(): hidden_states, residual = self.decoder( positions, hidden_states, forward_batch, residual, zero_allocator, ) if not forward_batch.forward_mode.is_idle(): if residual is not None: hidden_states, _ = self.shared_head.norm(hidden_states, residual) else: hidden_states = self.shared_head.norm(hidden_states) if nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp): # allgather + rerrange hidden_states = cp_all_gather_rerange_output( hidden_states, self.cp_size, forward_batch, torch.cuda.current_stream(), ) return hidden_states class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM): def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: nn.Module.__init__(self) self.config = config self.tp_size = get_tensor_model_parallel_world_size() self.quant_config = quant_config # if not set, model load will be broken in DeepseekV3ForCausalLM load_weights() self.pp_group = get_pp_group() self.determine_num_fused_shared_experts("DeepseekV3ForCausalLMNextN") self.use_nsa = is_deepseek_nsa(config) 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 = None self.cp_size = None self.model = DeepseekModelNextN( 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("model.shared_head.head", prefix), use_attn_tp_group=get_global_server_args().enable_dp_lm_head, ) self.logits_processor = LogitsProcessor(config) @torch.no_grad() def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: # TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2 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(), ) hidden_states = self.model(input_ids, positions, forward_batch) return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch ) def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): super().load_weights(weights, is_nextn=True) EntryClass = [DeepseekV3ForCausalLMNextN]