[model][NPU]:Add NPU fusion operator patch to Zimage model to improve performance

This commit is contained in:
feng0w0
2026-02-09 12:31:34 +08:00
parent 3f8468893a
commit e56a4d5730

View File

@@ -1,4 +1,4 @@
import torch, math import torch, math, warnings
from PIL import Image from PIL import Image
from typing import Union from typing import Union
from tqdm import tqdm from tqdm import tqdm
@@ -682,7 +682,7 @@ def apply_npu_patch(enable_npu_patch: bool=True):
rms_norm_forward_transformers_npu, rms_norm_forward_transformers_npu,
rotary_emb_Zimage_npu rotary_emb_Zimage_npu
) )
warnings.warn("Replacing RMSNorm and Rope with NPU fusion operators to improve the performance of the model on NPU.Set enable_npu_patch=False to disable this feature.")
RMSNorm.forward = rms_norm_forward_npu RMSNorm.forward = rms_norm_forward_npu
Qwen3RMSNorm.forward = rms_norm_forward_transformers_npu Qwen3RMSNorm.forward = rms_norm_forward_transformers_npu
Attention.apply_rotary_emb = rotary_emb_Zimage_npu Attention.apply_rotary_emb = rotary_emb_Zimage_npu