[NPU] [Features] [Bugfix] Support mindiesd adaln

vllm_omni/diffusion/layers/adalayernorm.py

@@ -1,3 +1,5 @@
+from importlib.util import find_spec
+
 import torch
 import torch.nn as nn
 from vllm.logger import init_logger
@@ -6,6 +8,8 @@
 
 logger = init_logger(__name__)
 
+_HAS_MINDIESD = find_spec("mindiesd") is not None
+
 
 class AdaLayerNorm(CustomOp):
     """
@@ -83,10 +87,21 @@ def forward_npu(
     ) -> torch.Tensor:
         shift_result, scale_result, gate_result = self.preprocess(mod_params, index)
 
+        if _HAS_MINDIESD:
+            try:
+                from mindiesd import layernorm_scale_shift
+
+                output = layernorm_scale_shift(self.layernorm, x, scale_result, shift_result, fused=True)
+
+                return output, gate_result
+            except ImportError as e:
+                logger.warning_once(f"mindiesd import failed, falling back to torch_npu: {e}")
+
         import torch_npu
 
-        output = torch_npu.npu_layer_norm_eval(
-            x, normalized_shape=[self.hidden_size], weight=(1 + scale_result), bias=shift_result, eps=self.eps
+        output = (
+            torch_npu.npu_layer_norm_eval(x, normalized_shape=[self.hidden_size], eps=self.eps) * (1 + scale_result)
+            + shift_result
         )
 
         return output, gate_result