From b8cfa02c013b0b40dce5a1bebaf55e7d1c06dcf8 Mon Sep 17 00:00:00 2001 From: liupeng374 Date: Wed, 10 Dec 2025 22:55:13 +0800 Subject: [PATCH] [NPU] bug fix for mtp and w4a8 (#14806) --- .../npu/attention/ascend_backend.py | 2 ++ .../npu/quantization/fused_moe_method_npu.py | 20 ++++++++++--------- python/sglang/srt/layers/moe/ep_moe/layer.py | 7 ++++++- 3 files changed, 19 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 22d9021f2..2c804058a 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -297,6 +297,8 @@ class AscendAttnBackend(AttentionBackend): if ( self.use_mla and forward_batch.forward_mode.is_extend() + and not forward_batch.forward_mode.is_draft_extend(include_v2=True) + and not forward_batch.forward_mode.is_target_verify() and sum(forward_batch.extend_prefix_lens_cpu) > 0 ): self.forward_metadata.prefix_lens = forward_batch.extend_prefix_lens.to( diff --git a/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py index 938314b0f..279700b3e 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py @@ -567,11 +567,13 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase): group_list, output_dtype, ): + from sgl_kernel_npu.activation.swiglu_quant import swiglu_quant + hidden_states = torch.ops.npu.npu_grouped_matmul( x=[hidden_states], - weight=[self.w13_weight], - scale=[self.w13_weight_scale], - bias=[self.w13_scale_bias], + weight=[layer.w13_weight], + scale=[layer.w13_weight_scale], + bias=[layer.w13_scale_bias], per_token_scale=[hidden_states_scale], group_list=group_list, split_item=2, @@ -580,15 +582,15 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase): output_dtype=output_dtype, )[0] - # act_fn: swiglu - hidden_states = torch.ops.npu.npu_swiglu(hidden_states) - hidden_states, swiglu_out_scale = torch.ops.npu.npu_dynamic_quant(hidden_states) + hidden_states, swiglu_out_scale = swiglu_quant( + hidden_states, group_list, group_list_type + ) hidden_states = torch.ops.npu.npu_grouped_matmul( x=[hidden_states], - weight=[self.w2_weight], - scale=[self.w2_weight_scale], - bias=[self.w2_scale_bias], + weight=[layer.w2_weight], + scale=[layer.w2_weight_scale], + bias=[layer.w2_scale_bias], per_token_scale=[swiglu_out_scale], group_list=group_list, split_item=2, diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 97e654331..53c832e1d 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -6,6 +6,9 @@ from typing import TYPE_CHECKING, Any, Dict, Optional, Union import torch from sglang.srt.environ import envs +from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( + NPUW4A16Int4DynamicMoEMethod, +) from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.moe import ( get_deepep_mode, @@ -347,7 +350,9 @@ class DeepEPMoE(FusedMoE): ) else: input_quant = get_bool_env_var("DEEP_NORMAL_MODE_USE_INT8_QUANT") - if not input_quant and self.w13_weight.dtype != torch.int32: + if not input_quant and not isinstance( + self.quant_method, NPUW4A16Int4DynamicMoEMethod + ): hidden_states, hidden_states_scale = torch_npu.npu_dynamic_quant( hidden_states )