From 532470bcca5260e9c7054eaa4389528d2e76cb86 Mon Sep 17 00:00:00 2001 From: chenxu214 Date: Wed, 18 Mar 2026 15:22:04 +0800 Subject: [PATCH] [NPU] add new fusion operator DispatchFFNCombine (#20245) --- python/sglang/srt/environ.py | 1 + .../sglang/srt/hardware_backend/npu/utils.py | 5 ++ python/sglang/srt/layers/moe/ep_moe/layer.py | 67 +++++++++++++++---- .../srt/layers/moe/token_dispatcher/fuseep.py | 1 + python/sglang/srt/server_args.py | 9 +++ 5 files changed, 69 insertions(+), 14 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e9c6481ba..25b95270c 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -322,6 +322,7 @@ class Envs: SGLANG_NPU_FORWARD_NATIVE_GEMMA_RMS_NORM = EnvBool(False) # Delay all-gather after qlora for better performance for Deepseek v3.2 SGLANG_USE_AG_AFTER_QLORA = EnvBool(False) + SGLANG_NPU_FUSED_MOE_MODE = EnvInt(1) # Quantization SGLANG_INT4_WEIGHT = EnvBool(False) diff --git a/python/sglang/srt/hardware_backend/npu/utils.py b/python/sglang/srt/hardware_backend/npu/utils.py index a78de6130..cdc5ad2fd 100644 --- a/python/sglang/srt/hardware_backend/npu/utils.py +++ b/python/sglang/srt/hardware_backend/npu/utils.py @@ -22,6 +22,11 @@ class NPUACLFormat(IntEnum): ACL_FORMAT_FRACTAL_NZ = 29 +class FusedMoEMode(IntEnum): + FUSED_DEEP_MOE = 1 + DISPATCH_FFN_COMBINE = 2 + + def _call_once(fn: Callable): @functools.wraps(fn) diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index d507dad12..ef2225ecc 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -7,7 +7,7 @@ import torch from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph from sglang.srt.environ import envs -from sglang.srt.hardware_backend.npu.utils import npu_format_cast +from sglang.srt.hardware_backend.npu.utils import FusedMoEMode, npu_format_cast from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.moe import ( get_deepep_mode, @@ -528,23 +528,62 @@ class NpuFuseEPMoE(DeepEPMoE): return weight.view(*original_shape[:dim], -1, *original_shape[dim + 1 :]) + def release_weight_cache(self, weight: torch.Tensor): + # .contiguous() introduces additional memory overhead and needs to be released using resize_(0) + origin_weight = weight.data.transpose(1, 2) + new_weight = origin_weight.contiguous() + origin_weight.untyped_storage().resize_(0) + return new_weight + + def scale_from_float_to_int64(self, scale): + import numpy as np + + scale = torch.from_numpy( + np.frombuffer( + scale.cpu().to(torch.float32).numpy().tobytes(), dtype=np.int32 + ).astype(np.int64) + ).to(scale.device) + return torch.nn.Parameter(scale, requires_grad=False) + def _process_weights_after_loading(self, layer: torch.nn.Module) -> None: - cpu_w13 = layer.w13_weight.data.transpose(1, 2).cpu() - layer.w13_weight.data = self.reshape_w13_weight(cpu_w13, -1).npu() - layer.w13_weight.data = npu_format_cast(layer.w13_weight.data) + if ( + envs.SGLANG_NPU_FUSED_MOE_MODE.get() + == FusedMoEMode.DISPATCH_FFN_COMBINE.value + ): + w13_weight = self.release_weight_cache(layer.w13_weight) + layer.w13_weight.data = npu_format_cast(w13_weight) + w2_weight = self.release_weight_cache(layer.w2_weight) + layer.w2_weight.data = npu_format_cast(w2_weight) - layer.w2_weight.data = npu_format_cast(layer.w2_weight.data) + layer.w13_weight_scale.data = layer.w13_weight_scale.data.view( + layer.w13_weight_scale.data.shape[0], -1 + ) + w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous() + layer.w2_weight_scale = torch.nn.Parameter( + w2_scale.to(torch.float32), requires_grad=False + ) - w13_scale = layer.w13_weight_scale.data.squeeze(-1).contiguous() - w13_scale = self.permute_w13_weight_scale(w13_scale, 128) - layer.w13_weight_scale = torch.nn.Parameter( - w13_scale.to(torch.float32), requires_grad=False - ) + layer.w13_weight_scale = self.scale_from_float_to_int64( + layer.w13_weight_scale.data + ) + layer.w2_weight_scale = self.scale_from_float_to_int64( + layer.w2_weight_scale.data + ) + else: + cpu_w13 = layer.w13_weight.data.transpose(1, 2).cpu() + layer.w13_weight.data = self.reshape_w13_weight(cpu_w13, -1).npu() + w13_scale = layer.w13_weight_scale.data.squeeze(-1).contiguous() + w13_scale = self.permute_w13_weight_scale(w13_scale, 128) + layer.w13_weight_scale = torch.nn.Parameter( + w13_scale.to(torch.float32), requires_grad=False + ) + layer.w13_weight.data = npu_format_cast(layer.w13_weight.data) + layer.w2_weight.data = npu_format_cast(layer.w2_weight.data) - w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous() - layer.w2_weight_scale = torch.nn.Parameter( - w2_scale.to(torch.float32), requires_grad=False - ) + w2_scale = layer.w2_weight_scale.data.squeeze(-1).contiguous() + layer.w2_weight_scale = torch.nn.Parameter( + w2_scale.to(torch.float32), requires_grad=False + ) if hasattr(layer, "w13_weight_offset"): layer.w13_weight_offset = torch.nn.Parameter( diff --git a/python/sglang/srt/layers/moe/token_dispatcher/fuseep.py b/python/sglang/srt/layers/moe/token_dispatcher/fuseep.py index ba29e1360..c33c337e2 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/fuseep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/fuseep.py @@ -79,6 +79,7 @@ class NpuFuseEPDispatcher(BaseDispatcher): gmm2_weight_scale=kwargs["gmm2_weight_scale"], num_max_dispatch_tokens_per_rank=self.num_max_dispatch_tokens_per_rank, num_experts=self.num_experts, + fuse_mode=envs.SGLANG_NPU_FUSED_MOE_MODE.get(), ) return FuseEPDispatchOutput(hidden_states) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 247a90576..448c7d6c4 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2633,6 +2633,15 @@ class ServerArgs: logger.warning( f"Ascend fused EP MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]." ) + fuse_mode = os.environ.get("SGLANG_NPU_FUSED_MOE_MODE", None) + if fuse_mode not in ["1", "2"]: + raise ValueError( + f"Wrong value of {fuse_mode=}, the NPU only support 1 or 2." + ) + elif fuse_mode == "2": + assert ( + self.quantization == "modelslim" + ), "When fuse_mode is set to 2, the NPU supports only ModelSlim quantization." if self.moe_a2a_backend == "flashinfer": self.ep_size = self.tp_size logger.warning(