[NPU] add new fusion operator DispatchFFNCombine (#20245)

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)

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)

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,18 +528,57 @@ 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:
+        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.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
+            )
+
+            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()
-        layer.w13_weight.data = npu_format_cast(layer.w13_weight.data)
-
-        layer.w2_weight.data = npu_format_cast(layer.w2_weight.data)
-
             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(

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)
 

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(