From 5493e75877a9941d9026ec383579dab6939aff89 Mon Sep 17 00:00:00 2001 From: jiangmengyu18 <56633611+jiangmengyu18@users.noreply.github.com> Date: Thu, 5 Mar 2026 09:01:19 +0800 Subject: [PATCH] [NPU] [Features] [Bugfix] Support mindiesd adaln (#1537) Signed-off-by: jiangmengyu18 <451528648@qq.com> Signed-off-by: Hongsheng Liu Co-authored-by: Hongsheng Liu Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> --- vllm_omni/diffusion/layers/adalayernorm.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index 9a644e1d9..c2389cc15 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/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 @@ from vllm_omni.diffusion.layers.custom_op import CustomOp logger = init_logger(__name__) +_HAS_MINDIESD = find_spec("mindiesd") is not None + class AdaLayerNorm(CustomOp): """ @@ -83,10 +87,21 @@ class AdaLayerNorm(CustomOp): ) -> 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