From 082acd5c6b41001c18679bdc17e05ca13a430fde Mon Sep 17 00:00:00 2001 From: jiangmengyu18 <451528648@qq.com> Date: Fri, 27 Feb 2026 09:18:51 +0800 Subject: [PATCH 1/9] support mindiesd adaln Signed-off-by: jiangmengyu18 <451528648@qq.com> --- vllm_omni/diffusion/layers/adalayernorm.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index 9a644e1d935..236abbcaf09 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 @@ -83,11 +85,12 @@ def forward_npu( ) -> torch.Tensor: shift_result, scale_result, gate_result = self.preprocess(mod_params, index) - 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 - ) + if find_spec("mindiesd"): + from mindiesd import layernorm_scale_shift + output = layernorm_scale_shift(self.layernorm, x, scale_result, shift_result, fused=True) + else: + import torch_npu + 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 From bd59cebb39169afe304e22b2145739792513e7d3 Mon Sep 17 00:00:00 2001 From: jiangmengyu18 <451528648@qq.com> Date: Fri, 27 Feb 2026 12:03:23 +0800 Subject: [PATCH 2/9] cleancode Signed-off-by: jiangmengyu18 <451528648@qq.com> --- vllm_omni/diffusion/layers/adalayernorm.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index 236abbcaf09..1f7f9458ee2 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -90,7 +90,10 @@ def forward_npu( output = layernorm_scale_shift(self.layernorm, x, scale_result, shift_result, fused=True) else: import torch_npu - output = torch_npu.npu_layer_norm_eval(x, normalized_shape=[self.hidden_size], eps=self.eps) * (1 + scale_result) + shift_result + 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 From 607436308cdc1ead08d1019d1c060b41c59a21b3 Mon Sep 17 00:00:00 2001 From: jiangmengyu18 <451528648@qq.com> Date: Fri, 27 Feb 2026 12:27:41 +0800 Subject: [PATCH 3/9] add error handling and eliminate the repaeted find_spec call Signed-off-by: jiangmengyu18 <451528648@qq.com> --- vllm_omni/diffusion/layers/adalayernorm.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index 1f7f9458ee2..c299e4f000f 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -8,6 +8,8 @@ logger = init_logger(__name__) +_HAS_MINDIESD = find_spec("mindiesd") is not None + class AdaLayerNorm(CustomOp): """ @@ -85,9 +87,17 @@ def forward_npu( ) -> torch.Tensor: shift_result, scale_result, gate_result = self.preprocess(mod_params, index) - if find_spec("mindiesd"): - from mindiesd import layernorm_scale_shift - output = layernorm_scale_shift(self.layernorm, x, scale_result, shift_result, fused=True) + if _HAS_MINDIESD: + try: + from mindiesd import layernorm_scale_shift + output = layernorm_scale_shift(self.layernorm, x, scale_result, shift_result, fused=True) + except Exception as e: + logger.warning(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], eps=self.eps) * (1 + scale_result) + + shift_result + ) else: import torch_npu output = ( From 6314f5826167e2887d866c8094489a3c5e117526 Mon Sep 17 00:00:00 2001 From: jiangmengyu18 <451528648@qq.com> Date: Fri, 27 Feb 2026 15:29:49 +0800 Subject: [PATCH 4/9] cleancode Signed-off-by: jiangmengyu18 <451528648@qq.com> --- vllm_omni/diffusion/layers/adalayernorm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index c299e4f000f..731145e44a8 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -95,8 +95,8 @@ def forward_npu( logger.warning(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], eps=self.eps) * (1 + scale_result) - + shift_result + torch_npu.npu_layer_norm_eval(x, normalized_shape=[self.hidden_size], eps=self.eps) + * (1 + scale_result) + shift_result ) else: import torch_npu From f85a4482da93c08d2810b8c412f18694943d1268 Mon Sep 17 00:00:00 2001 From: jiangmengyu18 <451528648@qq.com> Date: Fri, 27 Feb 2026 15:34:22 +0800 Subject: [PATCH 5/9] cleancode Signed-off-by: jiangmengyu18 <451528648@qq.com> --- vllm_omni/diffusion/layers/adalayernorm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index 731145e44a8..f34c07d7a32 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -96,7 +96,8 @@ def forward_npu( import torch_npu output = ( torch_npu.npu_layer_norm_eval(x, normalized_shape=[self.hidden_size], eps=self.eps) - * (1 + scale_result) + shift_result + * (1 + scale_result) + + shift_result ) else: import torch_npu From 0b52ccbc9521e91a651a63386f60407b6774c2e0 Mon Sep 17 00:00:00 2001 From: jiangmengyu18 <451528648@qq.com> Date: Fri, 27 Feb 2026 15:39:43 +0800 Subject: [PATCH 6/9] cleancode Signed-off-by: jiangmengyu18 <451528648@qq.com> --- vllm_omni/diffusion/layers/adalayernorm.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index f34c07d7a32..b1e2afb7e03 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -90,10 +90,12 @@ def forward_npu( if _HAS_MINDIESD: try: from mindiesd import layernorm_scale_shift + output = layernorm_scale_shift(self.layernorm, x, scale_result, shift_result, fused=True) except Exception as e: logger.warning(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], eps=self.eps) * (1 + scale_result) @@ -101,6 +103,7 @@ def forward_npu( ) else: import torch_npu + output = ( torch_npu.npu_layer_norm_eval(x, normalized_shape=[self.hidden_size], eps=self.eps) * (1 + scale_result) + shift_result From 5db8aa6d04271cadb20eefa89f68664c21de4396 Mon Sep 17 00:00:00 2001 From: jiangmengyu18 <451528648@qq.com> Date: Sat, 28 Feb 2026 08:50:05 +0800 Subject: [PATCH 7/9] remove redundant adaln calls implemented by torch_npu Signed-off-by: jiangmengyu18 <451528648@qq.com> --- vllm_omni/diffusion/layers/adalayernorm.py | 23 +++++++++------------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index b1e2afb7e03..9f5b89ff481 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -92,22 +92,17 @@ def forward_npu( from mindiesd import layernorm_scale_shift output = layernorm_scale_shift(self.layernorm, x, scale_result, shift_result, fused=True) + + return output, gate_result except Exception as e: - logger.warning(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], eps=self.eps) - * (1 + scale_result) - + shift_result - ) - else: - import torch_npu + 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], eps=self.eps) * (1 + scale_result) - + shift_result - ) + 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 From 8e006b0116510f8b29f78eb64dcabd50a0399c07 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Sat, 28 Feb 2026 11:49:11 +0800 Subject: [PATCH 8/9] Update vllm_omni/diffusion/layers/adalayernorm.py Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> Signed-off-by: Hongsheng Liu --- vllm_omni/diffusion/layers/adalayernorm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index 9f5b89ff481..3d9fb8ce412 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -94,7 +94,8 @@ def forward_npu( output = layernorm_scale_shift(self.layernorm, x, scale_result, shift_result, fused=True) return output, gate_result - except Exception as e: + except ImportError as e: + logger.warning_once(f"mindiesd import failed, falling back to torch_npu: {e}") logger.warning_once(f"mindiesd import failed, falling back to torch_npu: {e}") import torch_npu From e106fbe37aaa407637e289b522af79e9a9c56adc Mon Sep 17 00:00:00 2001 From: jiangmengyu18 <451528648@qq.com> Date: Sat, 28 Feb 2026 12:08:39 +0800 Subject: [PATCH 9/9] remove rebundant logger.warning Signed-off-by: jiangmengyu18 <451528648@qq.com> --- vllm_omni/diffusion/layers/adalayernorm.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index 3d9fb8ce412..c2389cc151a 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -96,7 +96,6 @@ def forward_npu( return output, gate_result except ImportError as e: logger.warning_once(f"mindiesd import failed, falling back to torch_npu: {e}") - logger.warning_once(f"mindiesd import failed, falling back to torch_npu: {e}") import torch_npu