[Perf][Wan2.2] Add fused RMSNorm replace WanRMS_norm on npu (#3067)

Signed-off-by: lyj-jjj <liuyingjun5@huawei.com>
This commit is contained in:
lyj-jjj authored and GitHub committed 2026-04-23 18:25:10 +00:00
1 parent c8efdbd671
commit 19980bb37c
3 files changed
+95

No files matched your search

@@ -0,0 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys
from vllm_omni.diffusion.layers.norm import RMSNormVAE
def patch_wan_rms_norm():
"""Patch diffusers Wan RMSNorm implementation with RMSNormVAE."""
for module_name, module in sys.modules.items():
if hasattr(module, "WanRMS_norm"):
setattr(module, "WanRMS_norm", RMSNormVAE)