[Perf][Wan2.2] Add fused RMSNorm replace WanRMS_norm on npu (#3067)

Signed-off-by: lyj-jjj <liuyingjun5@huawei.com>

vllm_omni/diffusion/layers/norm.py

@@ -108,3 +108,81 @@ class RMSNorm(CustomOp):
         out = x * torch.rsqrt(variance + self.variance_epsilon)
         out = self.weight.to(torch.float32) * out
         return out.to(input_dtype)
+
+
+class RMSNormVAE(CustomOp):
+    """Root Mean Square Layer Normalization for Channel-First or Last"""
+
+    def __init__(
+        self,
+        dim: int,
+        channel_first: bool = True,
+        images: bool = True,
+        bias: bool = False,
+        epsilon: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else None
+        self.epsilon = epsilon
+
+        self.gamma_rmsnorm = None
+
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_hip(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_npu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        import torch_npu
+
+        if self.gamma_rmsnorm is None:
+            self.gamma_rmsnorm = self.gamma.reshape(-1)
+
+        if self.channel_first:
+            x = x.transpose(1, -1)
+            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0].transpose(1, -1)
+        else:
+            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0]
+
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        out = (
+            F.normalize(
+                x,
+                dim=(1 if self.channel_first else -1),
+                eps=self.epsilon,
+            )
+            * self.scale
+            * self.gamma
+        )
+        if self.bias is not None:
+            out = out + self.bias
+        return out

vllm_omni/diffusion/models/wan2_2/__init__.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from .patch_diffusers import patch_wan_rms_norm
 from .pipeline_wan2_2 import (
     Wan22Pipeline,
     WanT2VDMD2Pipeline,
@@ -51,3 +52,5 @@ __all__ = [
     "VaceWanTransformerBlock",
     "WanVACETransformer3DModel",
 ]
+
+patch_wan_rms_norm()

vllm_omni/diffusion/models/wan2_2/patch_diffusers.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import sys
+
+from vllm_omni.diffusion.layers.norm import RMSNormVAE
+
+
+def patch_wan_rms_norm():
+    """Patch diffusers Wan RMSNorm implementation with RMSNormVAE."""
+
+    for module_name, module in sys.modules.items():
+        if hasattr(module, "WanRMS_norm"):
+            setattr(module, "WanRMS_norm", RMSNormVAE)