[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:
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)
|
||||
Reference in new issue
Block a user