【NPU】【bugfix】fix server error when mtp unquant (#26389)

Co-authored-by: cen121212 <luochen23@huawei.com>
Co-authored-by: Even Zhou <even.y.zhou@outlook.com>

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."""
+        """TODO adapt different quantization schemes for base model and draft model on NPU"""
-        if self.use_fp8:
+        pass
-            os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "1"
-        else:
-            os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "0"
 
     def set_overlap_args(
         self, combine_overlap_args: CombineOverlapArgs, meta_overlap_args: dict

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(
+        exit_stack = ExitStack()
-            buffer_size=2,
+        if (
-            dtype=torch.float32,
+            _is_npu
-            device=(
+            and self.quant_config is None
-                input_embeds.device if input_embeds is not None else input_ids.device
+            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:
+        try:
-            hidden_states = self.embed_tokens(input_ids)
+            zero_allocator = BumpAllocator(
-        else:
+                buffer_size=2,
-            hidden_states = input_embeds
+                dtype=torch.float32,
-
+                device=(
-        if hidden_states.shape[0] > 0:
+                    input_embeds.device
-            eh_input = torch.cat(
+                    if input_embeds is not None
-                (
+                    else input_ids.device
-                    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):
-            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 input_embeds is None:
-            if residual is not None:
+                hidden_states = self.embed_tokens(input_ids)
-                hidden_states, _ = self.shared_head.norm(hidden_states, residual)
             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_split_and_rebuild_data(forward_batch, hidden_states)
-                hidden_states = cp_all_gather_rerange_output(
+                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
 
 

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,
     ):
-
+        exit_stack = ExitStack()
-        assert input_embeds is None
-        input_embeds = forward_batch.mm_input_embeds
         if (
-            forward_batch.forward_mode.is_extend()
+            is_npu()
-            and forward_batch.contains_mm_inputs()
+            and self.quant_config is None
-            and not forward_batch.forward_mode.is_draft_extend(include_v2=True)
+            and get_global_server_args().quantization is not None
         ):
-            assert input_embeds is not None
+            # ascend mtp unquant
-            input_embeds = torch.cat(
+            exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True))
-                [input_embeds[:-1], self.model.embed_tokens(input_ids[-1].unsqueeze(0))]
+            exit_stack.enter_context(
+                envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False)
             )
 
-        if input_embeds is None:
+        try:
-            input_embeds = self.model.embed_tokens(input_ids)
+            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():
+            hidden_states = forward_batch.spec_info.hidden_states
-            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 = 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.fc(hidden_states)
-            hidden_states = self.model(
+
-                input_ids,
+            with get_global_expert_distribution_recorder().disable_this_region():
-                positions,
+                hidden_states = self.model(
-                forward_batch,
+                    input_ids,
-                hidden_states,
+                    positions,
-            )
+                    forward_batch,
+                    hidden_states,
+                )
+        finally:
+            exit_stack.close()
 
         return self.logits_processor(
             input_ids, hidden_states, self.lm_head, forward_batch

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,
     ):
-
+        exit_stack = ExitStack()
-        if input_embeds is None:
+        if (
-            input_embeds = self.model.embed_tokens(input_ids)
+            is_npu()
-
+            and self.quant_config is None
-        hidden_states = forward_batch.spec_info.hidden_states
+            and get_global_server_args().quantization is not None
-        # Some idle batch has 0 batch size. GemmaRMSNorm.forward would fail due to bs=0.
+        ):
-        if not forward_batch.forward_mode.is_idle():
+            # ascend mtp unquant
-            input_embeds = self.pre_fc_norm_embedding(input_embeds)
+            exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True))
-            hidden_states = self.pre_fc_norm_hidden(hidden_states)
+            exit_stack.enter_context(
-        hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1))
+                envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False)
-
-        with get_global_expert_distribution_recorder().disable_this_region():
-            hidden_states = self.model(
-                input_ids,
-                positions,
-                forward_batch,
-                hidden_states,
             )
 
+        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
         )

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)