diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index 9a644e1d935..c2389cc151a 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 @@ 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