diff --git a/python/sglang/srt/models/glm4_moe_nextn.py b/python/sglang/srt/models/glm4_moe_nextn.py index a123b4924..149fec102 100644 --- a/python/sglang/srt/models/glm4_moe_nextn.py +++ b/python/sglang/srt/models/glm4_moe_nextn.py @@ -12,9 +12,8 @@ # limitations under the License. # ============================================================================== -"""Inference-only GLM-4.5, GLM-4.6 Speculative Decoding.""" +"""Inference-only GLM-4.5, GLM-4.6 and GLM-4.7 Speculative Decoding.""" -import contextlib import logging from typing import Iterable, Optional, Tuple @@ -23,7 +22,6 @@ from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import get_tensor_model_parallel_world_size -from sglang.srt.environ import temp_set_env from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.layers.layernorm import RMSNorm @@ -36,7 +34,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.models.glm4_moe import Glm4MoeDecoderLayer, Glm4MoeForCausalLM from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import add_prefix +from sglang.srt.utils import add_prefix, is_npu logger = logging.getLogger(__name__) @@ -51,7 +49,7 @@ class Glm4MoeModelNextN(nn.Module): super().__init__() if quant_config is not None and quant_config.get_name() == "modelopt_fp4": logger.warning( - "Overriding Glm4MoeForCausalLMNextN quant config for modelopt_fp4 GLM-4.5 / GLM-4.6 model." + "Overriding Glm4MoeForCausalLMNextN quant config for modelopt_fp4 GLM-4.5 / GLM-4.6 / GLM-4.7 model." ) quant_config = None @@ -128,11 +126,13 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM): nn.Module.__init__(self) self.config = config self.tp_size = get_tensor_model_parallel_world_size() - self.needs_quant_draft = ( - get_global_server_args().speculative_draft_model_quantization is not None - or quant_config is not None - ) - quant_config = quant_config if self.needs_quant_draft else None + if ( + is_npu() + and get_global_server_args().speculative_draft_model_quantization is None + ): + quant_config = None + self.quant_config = quant_config + self.model = Glm4MoeModelNextN( config, quant_config, prefix=add_prefix("model", prefix) ) @@ -156,21 +156,7 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM): positions: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: - # Support unquant speculative draft model - if self.needs_quant_draft: - cxt = contextlib.nullcontext() - else: - # SGLANG_DEEPEP_BF16_DISPATCH is deprecated, will need - # to be removed in the future and moved to a new - # --deepep-dispatcher-output-dtype server argument. - unquant_patch = { - "SGLANG_DEEPEP_BF16_DISPATCH": "1", - } - cxt = temp_set_env(allow_sglang=True, **unquant_patch) - - with cxt: - hidden_states = self.model(input_ids, positions, forward_batch) - + hidden_states = self.model(input_ids, positions, forward_batch) return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch )