[PERF] Wan2.2 support rmsnorm fused op (#2583)

Signed-off-by: fan2956 <zhoufan53@huawei.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: gcanlin <canlinguosdu@gmail.com>

vllm_omni/diffusion/layers/adalayernorm.py

@@ -7,6 +7,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.linear import ReplicatedLinear
 
 from vllm_omni.diffusion.layers.custom_op import CustomOp
+from vllm_omni.diffusion.layers.norm import LayerNorm
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
@@ -27,7 +28,7 @@ class AdaLayerNorm(CustomOp):
         self.eps = eps
         self.elementwise_affine = elementwise_affine
         self.hidden_size = hidden_size
-        self.layernorm = nn.LayerNorm(self.hidden_size, elementwise_affine=self.elementwise_affine, eps=self.eps)
+        self.layernorm = LayerNorm(self.hidden_size, elementwise_affine=self.elementwise_affine, eps=self.eps)
 
     def forward_cuda(
         self,

vllm_omni/diffusion/layers/norm.py

@@ -0,0 +1,110 @@
+from importlib.util import find_spec
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.layers.custom_op import CustomOp
+
+logger = init_logger(__name__)
+
+_HAS_MINDIESD = find_spec("mindiesd") is not None
+
+
+class LayerNorm(nn.LayerNorm, CustomOp):
+    """
+    LayerNorm implementation that inherits from both ``nn.LayerNorm`` and ``CustomOp``.
+    NPU:
+        Uses ``mindiesd.fast_layernorm(self, x)`` when MindIE-SD is installed.
+    CUDA / HIP / XPU / native:
+        Falls back to FP32 nn.LayerNorm implementation.
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True):
+        super().__init__(normalized_shape=dim, eps=eps, elementwise_affine=elementwise_affine)
+        # CustomOp.__init__ cannot be called here because it would re-run
+        # nn.Module initialization and clear LayerNorm parameters.
+        self._forward_method = CustomOp.dispatch_forward(self)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self._forward_method(x)
+
+    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_xpu(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_npu(self, x: torch.Tensor) -> torch.Tensor:
+        if _HAS_MINDIESD:
+            try:
+                from mindiesd import fast_layernorm
+
+                return fast_layernorm(self, x)
+            except ImportError as e:
+                logger.warning_once(
+                    "mindiesd.fast_layernorm import failed, falling back to FP32 layer_norm: %s",
+                    e,
+                )
+
+        return self.forward_native(x)
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        origin_dtype = x.dtype
+        return F.layer_norm(
+            x.float(),
+            self.normalized_shape,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        ).to(origin_dtype)
+
+
+class RMSNorm(CustomOp):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    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
+
+        output = torch_npu.npu_rms_norm(x, gamma=self.weight, epsilon=self.variance_epsilon)[0]
+
+        return output
+
+    def forward_xpu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        input_dtype = x.dtype
+        x = x.to(torch.float32)
+        variance = x.pow(2).mean(-1, keepdim=True)
+        out = x * torch.rsqrt(variance + self.variance_epsilon)
+        out = self.weight.to(torch.float32) * out
+        return out.to(input_dtype)