[WIP] wan2.2: replace the small operator's WanRMS_norm with the fused operator on npu

[lyj-jjj]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

replace the small operator's WanRMS_norm with the fused operator on npu

Test Plan

• script

export ASCEND_RT_VISIBLE_DEVICES=8,9,10,11
export PYTORCH_NPU_ALLOC_CONF="expandable_segments:True"

python image_to_video.py \
--model /weights/Wan2.2-I2V-A14B-LightX2V-Diffusers \
--image i2v_input.jpg \
--prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside" \
--height 1280 \
--width 720 \
--num-frames 61 \
--guidance-scale 1.0 \
--guidance-scale-high 1.0 \
--num-inference-steps 4 \
--cfg-parallel-size 1 \
--ulysses-degree 4 \
--boundary-ratio 0.875 \
--flow-shift 12.0 \
--fps 16 \
--output i2v_output_origin_weight_1.mp4 \
--vae-patch-parallel-size 4 \
--vae-use-tiling \

Test Result

• performance
  before: Encoder-2.1s, Decoder-3.4s
  after: Encoder-1.6s, Decoder-2.5s ---- reduce-1.4s, improve-34%

• Accuracy

  berfore
  https://github.com/user-attachments/assets/683596fe-7e83-4272-ba16-0475c464cf73
  after
  https://github.com/user-attachments/assets/4597a909-27a8-4930-95eb-2f9c1913f0ae

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[gcanlin]

Why not directly use RMSNorm introduced in #2583?

[gcanlin]

We will only need these lines. And it's unnecessary to create a new file.

Suggested change

	from __future__ import annotations
	import torch
	from torch import nn
	from vllm_omni.platforms import current_omni_platform
	def patch_wan_rms_norm():
	'''Replace small operators with fused operators'''
	if not current_omni_platform.is_npu():
	return
	import torch_npu
	class WanRMS_norm(nn.Module):
	def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bias: bool = False) -> None:
	super().__init__()
	broadcastable_dims = (1, 1, 1) if not images else (1, 1)
	shape = (dim, *broadcastable_dims) if channel_first else (dim,)
	self.channel_first = channel_first
	self.scale = dim ** 0.5
	self.gamma = nn.Parameter(torch.ones(shape))
	self.gamma_new = None
	self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
	def forward(self, x):
	x = x.transpose(1, -1)
	if self.gamma_new is None:
	self.gamma_new = self.gamma.transpose(0, -1).reshape(-1)
	x_out = torch_npu.npu_rms_norm(x, self.gamma_new, epsilon=1e-6)
	x_out = x_out[0].transpose(1, -1)
	return x_out
	import sys
	for module_name, module in sys.modules.items():
	if hasattr(module, 'WanRMS_norm'):
	setattr(module, 'WanRMS_norm', WanRMS_norm)
	import sys
	for module_name, module in sys.modules.items():
	if hasattr(module, 'WanRMS_norm'):
	setattr(module, 'WanRMS_norm', RMSNorm)