diff --git a/tests/diffusion/distributed/test_wan_vae_fastpath_install.py b/tests/diffusion/distributed/test_wan_vae_fastpath_install.py
index abf36f0926e..9f9ae1e63b6 100644
--- a/tests/diffusion/distributed/test_wan_vae_fastpath_install.py
+++ b/tests/diffusion/distributed/test_wan_vae_fastpath_install.py
@@ -535,7 +535,7 @@ def test_upsample_forward_only_fuses_nearest_2x() -> None:
 
 
 def test_rms_norm_vae_substitute_is_not_matched() -> None:
-    from vllm_omni.diffusion.layers.norm import RMSNormVAE
+    from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
 
     assert not fastpath_forwards.is_diffusers_rms_norm(RMSNormVAE(8, images=False))
     _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
diff --git a/tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py b/tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py
index 6ed3e5f813b..5ad9330a5a7 100644
--- a/tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py
+++ b/tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py
@@ -195,6 +195,17 @@ def test_no_propagation_when_tf_quant_config_is_none(self, mocker: MockerFixture
 class TestPatchWanRmsNorm:
     """Test that patch_wan_rms_norm doesn't raise on concurrent module registration."""
 
+    @pytest.fixture(autouse=True)
+    def restore_wan_rms_norm(self):
+        originals = [
+            (module, module.__dict__["WanRMS_norm"])
+            for module in list(sys.modules.values())
+            if getattr(module, "__dict__", None) is not None and "WanRMS_norm" in module.__dict__
+        ]
+        yield
+        for module, original in originals:
+            module.__dict__["WanRMS_norm"] = original
+
     def test_patches_modules_with_wan_rms_norm(self):
         from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
         from vllm_omni.diffusion.models.wan2_2.patch_diffusers import patch_wan_rms_norm
diff --git a/vllm_omni/diffusion/models/wan2_2/__init__.py b/vllm_omni/diffusion/models/wan2_2/__init__.py
index 4059c4fb568..61679b91643 100644
--- a/vllm_omni/diffusion/models/wan2_2/__init__.py
+++ b/vllm_omni/diffusion/models/wan2_2/__init__.py
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from vllm_omni.platforms import current_omni_platform
+
 from .patch_diffusers import patch_wan_rms_norm
 from .pipeline_wan2_2 import (
     Wan22Pipeline,
@@ -53,4 +55,5 @@
     "WanVACETransformer3DModel",
 ]
 
-patch_wan_rms_norm()
+if current_omni_platform.is_npu():
+    patch_wan_rms_norm()