diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index 57fa4f27a..d28145e2e 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -1,7 +1,6 @@ from __future__ import annotations import logging -import os from contextlib import nullcontext from dataclasses import dataclass from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, Union @@ -437,11 +436,8 @@ class _DeepEPDispatcherImplBase: # NVFP4 is supported on GPU, no adjustment needed def _update_int8_quant_env(self) -> None: - """Update the DEEP_NORMAL_MODE_USE_INT8_QUANT environment variable.""" - if self.use_fp8: - os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "1" - else: - os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "0" + """TODO adapt different quantization schemes for base model and draft model on NPU""" + pass def set_overlap_args( self, combine_overlap_args: CombineOverlapArgs, meta_overlap_args: dict diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 56e661e80..8b668cf56 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -16,6 +16,7 @@ import logging import os +from contextlib import ExitStack from typing import Iterable, Optional, Tuple import torch @@ -169,70 +170,87 @@ class DeepseekModelNextN(nn.Module): 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 - ), - ) + exit_stack = ExitStack() + if ( + _is_npu + and self.quant_config is None + and get_global_server_args().quantization is not None + ): + # ascend mtp unquant + exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True)) + exit_stack.enter_context( + envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False) + ) - if input_embeds is None: - hidden_states = self.embed_tokens(input_ids) - else: - hidden_states = input_embeds - - if hidden_states.shape[0] > 0: - eh_input = torch.cat( - ( - self.enorm(hidden_states), - self.hnorm( - forward_batch.spec_info.hidden_states - if self.rot_weight is None - else torch.matmul( - forward_batch.spec_info.hidden_states, self.rot_weight - ) - ), + try: + zero_allocator = BumpAllocator( + buffer_size=2, + dtype=torch.float32, + device=( + input_embeds.device + if input_embeds is not None + else input_ids.device ), - dim=-1, - ) - if isinstance(self.eh_proj, ReplicatedLinear): - hidden_states, _ = self.eh_proj(eh_input) - else: - hidden_states = self.eh_proj(eh_input) - - if dsa_use_prefill_cp( - forward_batch, self.dsa_enable_prefill_cp - ) or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp): - hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) - positions = cp_split_and_rebuild_position(forward_batch, positions) - residual = None - with get_global_expert_distribution_recorder().disable_this_region(): - hidden_states, residual, topk_indices = 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) + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) else: - hidden_states = self.shared_head.norm(hidden_states) + hidden_states = input_embeds + + if hidden_states.shape[0] > 0: + eh_input = torch.cat( + ( + self.enorm(hidden_states), + self.hnorm( + forward_batch.spec_info.hidden_states + if self.rot_weight is None + else torch.matmul( + forward_batch.spec_info.hidden_states, self.rot_weight + ) + ), + ), + dim=-1, + ) + if isinstance(self.eh_proj, ReplicatedLinear): + hidden_states, _ = self.eh_proj(eh_input) + else: + hidden_states = self.eh_proj(eh_input) if dsa_use_prefill_cp( forward_batch, self.dsa_enable_prefill_cp ) or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp): - # allgather + rerrange - hidden_states = cp_all_gather_rerange_output( + hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) + positions = cp_split_and_rebuild_position(forward_batch, positions) + residual = None + with get_global_expert_distribution_recorder().disable_this_region(): + hidden_states, residual, topk_indices = self.decoder( + positions, hidden_states, - self.cp_size, forward_batch, - torch.cuda.current_stream(), + 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 dsa_use_prefill_cp( + forward_batch, self.dsa_enable_prefill_cp + ) or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp): + # allgather + rerrange + hidden_states = cp_all_gather_rerange_output( + hidden_states, + self.cp_size, + forward_batch, + torch.cuda.current_stream(), + ) + finally: + exit_stack.close() + return hidden_states diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 92a1f992f..6b39a590c 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -16,6 +16,7 @@ import copy import logging +from contextlib import ExitStack from typing import Iterable, Optional, Tuple import torch @@ -23,6 +24,7 @@ from torch import nn from transformers import PretrainedConfig 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.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.layers.layernorm import GemmaRMSNorm @@ -140,38 +142,55 @@ class Qwen3_5ForCausalLMMTP(nn.Module): input_embeds: Optional[torch.Tensor] = None, **kwargs, ): - - assert input_embeds is None - input_embeds = forward_batch.mm_input_embeds + exit_stack = ExitStack() if ( - forward_batch.forward_mode.is_extend() - and forward_batch.contains_mm_inputs() - and not forward_batch.forward_mode.is_draft_extend(include_v2=True) + is_npu() + and self.quant_config is None + and get_global_server_args().quantization is not None ): - assert input_embeds is not None - input_embeds = torch.cat( - [input_embeds[:-1], self.model.embed_tokens(input_ids[-1].unsqueeze(0))] + # ascend mtp unquant + exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True)) + exit_stack.enter_context( + envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False) ) - if input_embeds is None: - input_embeds = self.model.embed_tokens(input_ids) + try: + assert input_embeds is None + input_embeds = forward_batch.mm_input_embeds + if ( + forward_batch.forward_mode.is_extend() + and forward_batch.contains_mm_inputs() + and not forward_batch.forward_mode.is_draft_extend(include_v2=True) + ): + assert input_embeds is not None + input_embeds = torch.cat( + [ + input_embeds[:-1], + self.model.embed_tokens(input_ids[-1].unsqueeze(0)), + ] + ) - hidden_states = forward_batch.spec_info.hidden_states + if input_embeds is None: + input_embeds = self.model.embed_tokens(input_ids) - if not forward_batch.forward_mode.is_idle(): - input_embeds = self.pre_fc_norm_embedding(input_embeds) - hidden_states = self.pre_fc_norm_hidden(hidden_states) - hidden_states = torch.cat([input_embeds, hidden_states], dim=-1) + hidden_states = forward_batch.spec_info.hidden_states - hidden_states = self.fc(hidden_states) + if not forward_batch.forward_mode.is_idle(): + input_embeds = self.pre_fc_norm_embedding(input_embeds) + hidden_states = self.pre_fc_norm_hidden(hidden_states) + hidden_states = torch.cat([input_embeds, hidden_states], dim=-1) - with get_global_expert_distribution_recorder().disable_this_region(): - hidden_states = self.model( - input_ids, - positions, - forward_batch, - hidden_states, - ) + hidden_states = self.fc(hidden_states) + + with get_global_expert_distribution_recorder().disable_this_region(): + hidden_states = self.model( + input_ids, + positions, + forward_batch, + hidden_states, + ) + finally: + exit_stack.close() return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch diff --git a/python/sglang/srt/models/qwen3_next_mtp.py b/python/sglang/srt/models/qwen3_next_mtp.py index 6c4a3ed85..3bade2141 100644 --- a/python/sglang/srt/models/qwen3_next_mtp.py +++ b/python/sglang/srt/models/qwen3_next_mtp.py @@ -16,6 +16,7 @@ import copy import logging +from contextlib import ExitStack from typing import Iterable, Optional, Tuple import torch @@ -23,6 +24,7 @@ from torch import nn from transformers import PretrainedConfig 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.layernorm import GemmaRMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor @@ -94,25 +96,39 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM): input_embeds: Optional[torch.Tensor] = None, **kwargs, ): - - if input_embeds is None: - input_embeds = self.model.embed_tokens(input_ids) - - hidden_states = forward_batch.spec_info.hidden_states - # Some idle batch has 0 batch size. GemmaRMSNorm.forward would fail due to bs=0. - if not forward_batch.forward_mode.is_idle(): - input_embeds = self.pre_fc_norm_embedding(input_embeds) - hidden_states = self.pre_fc_norm_hidden(hidden_states) - hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1)) - - with get_global_expert_distribution_recorder().disable_this_region(): - hidden_states = self.model( - input_ids, - positions, - forward_batch, - hidden_states, + exit_stack = ExitStack() + if ( + is_npu() + and self.quant_config is None + and get_global_server_args().quantization is not None + ): + # ascend mtp unquant + exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True)) + exit_stack.enter_context( + envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False) ) + try: + if input_embeds is None: + input_embeds = self.model.embed_tokens(input_ids) + + hidden_states = forward_batch.spec_info.hidden_states + # Some idle batch has 0 batch size. GemmaRMSNorm.forward would fail due to bs=0. + if not forward_batch.forward_mode.is_idle(): + input_embeds = self.pre_fc_norm_embedding(input_embeds) + hidden_states = self.pre_fc_norm_hidden(hidden_states) + hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1)) + + with get_global_expert_distribution_recorder().disable_this_region(): + hidden_states = self.model( + input_ids, + positions, + forward_batch, + hidden_states, + ) + finally: + exit_stack.close() + return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch ) diff --git a/test/registered/ascend/basic_function/parallel_strategy/expert_parallelism/test_npu_deepep.py b/test/registered/ascend/basic_function/parallel_strategy/expert_parallelism/test_npu_deepep.py index 36fcdaace..a56b71300 100644 --- a/test/registered/ascend/basic_function/parallel_strategy/expert_parallelism/test_npu_deepep.py +++ b/test/registered/ascend/basic_function/parallel_strategy/expert_parallelism/test_npu_deepep.py @@ -60,6 +60,7 @@ class TestAscendDeepEP(CustomTestCase): "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "32", "SGLANG_NPU_USE_MLAPO": "1", "TRANSFORMERS_VERBOSITY": "error", + "DEEP_NORMAL_MODE_USE_INT8_QUANT": "1", } os.environ.update(cls.extra_envs)