[NPU] [Diffusion] Use fused operator to improve Wan model E2E performance. (#24028)

Co-authored-by: gengjinsong <gengjinsong@huawei.com>
Co-authored-by: sglang-npu-bot <sglangnpu@163.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: gengjinsong <904939979@qq.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>

python/sglang/multimodal_gen/runtime/layers/elementwise.py

@@ -38,3 +38,10 @@ class MulAdd(CustomOp):
         self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
     ):
         return self.forward_native(a, b, c, k=k)
+
+    def forward_npu(
+        self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
+    ):
+        from sgl_kernel_npu.norm.scale_shift import fused_scale_shift
+
+        return fused_scale_shift(a, b, c, scale_constant=k)

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -520,6 +520,38 @@ class _ScaleResidualNormScaleShift(CustomOp):
         modulated = fuse_scale_shift_kernel(normalized, scale, shift)
         return modulated, residual_output
 
+    def forward_npu(
+        self,
+        residual: torch.Tensor,
+        x: torch.Tensor,
+        gate: torch.Tensor | int,
+        shift: torch.Tensor,
+        scale: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        from sgl_kernel_npu.norm.scale_shift import fused_scale_shift
+
+        # x.shape: [batch_size, seq_len, inner_dim]
+        if isinstance(gate, int):
+            # used by cross-attention, should be 1
+            assert gate == 1
+            residual_output = residual + x
+        elif isinstance(gate, torch.Tensor):
+            if gate.dim() == 4:
+                # gate.shape: [batch_size, num_frames, 1, inner_dim]
+                num_frames = gate.shape[1]
+                frame_seqlen = x.shape[1] // num_frames
+                residual_output = residual + (
+                    x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate
+                ).flatten(1, 2)
+            else:
+                # gate.shape: [batch_size, 1, inner_dim]
+                residual_output = residual + x * gate
+        else:
+            raise ValueError(f"Gate type {type(gate)} not supported")
+        normalized = self.norm(residual_output)
+        modulated = fused_scale_shift(normalized, scale, shift)
+        return modulated, residual_output
+
 
 class ScaleResidualLayerNormScaleShift(_ScaleResidualNormScaleShift):
     norm_type = "layer"
@@ -606,6 +638,15 @@ class _NormScaleShift(CustomOp):
         modulated = fuse_scale_shift_kernel(normalized, scale, shift)
         return modulated.to(x.dtype)
 
+    def forward_npu(
+        self, x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor
+    ) -> torch.Tensor:
+        from sgl_kernel_npu.norm.scale_shift import fused_scale_shift
+
+        normalized = self.norm(x)
+        modulated = fused_scale_shift(normalized, scale, shift)
+        return modulated.to(x.dtype)
+
 
 class LayerNormScaleShift(_NormScaleShift):
     norm_type = "layer"
@@ -902,9 +943,19 @@ def tensor_parallel_rms_norm(x: torch.Tensor, norm: "RMSNorm") -> torch.Tensor:
     src_dtype = x.dtype
     weight = norm.weight.tensor_split(tp_size)[tp_rank].float()
     x_fp32 = x.float()
-    variance = x_fp32.pow(2).mean(dim=-1, keepdim=True)
+    if _is_npu:
+        from sgl_kernel_npu.norm.rmsnorm_split import fused_rsqrt_mul, fused_variance
+
+        variance = fused_variance(x_fp32)
+    else:
+        variance = x_fp32.pow(2).mean(dim=-1, keepdim=True)
+
     variance = get_tp_group().all_reduce(
         variance, op=torch._C._distributed_c10d.ReduceOp.AVG
     )
-    output = x_fp32 * torch.rsqrt(variance + norm.variance_epsilon) * weight
+
+    if _is_npu:
+        output = fused_rsqrt_mul(x_fp32, variance, weight, norm.variance_epsilon)
+    else:
+        output = x_fp32 * torch.rsqrt(variance + norm.variance_epsilon) * weight
     return output.to(dtype=src_dtype)