minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: cen121212 <luochen23@huawei.com>
Co-authored-by: Even Zhou <even.y.zhou@outlook.com>
This commit is contained in:
cen121212
2026-05-30 15:01:19 +03:00
committed by GitHub
co-authored by cen121212 Even Zhou
parent 282c46133f
commit b421e60eed
5 changed files with 148 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/token_dispatcher/deepep.py
+2 -6
View File
@@ -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
python/sglang/srt/models/deepseek_nextn.py
+69 -51
View File
@@ -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
python/sglang/srt/models/qwen3_5_mtp.py
+43 -24
View File
@@ -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
python/sglang/srt/models/qwen3_next_mtp.py
+33 -17
View File
@@ -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
)
test/registered/ascend/basic_function/parallel_strategy/expert_parallelism/test_npu_deepep.py
+1
View File
@@ -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)