rms_norm and rotary embedding back to combined impl

DeepLink-org · Sep 11, 2024 · d8f0ddc · d8f0ddc
1 parent 2b1b705
commit d8f0ddc
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/deeplink_ext/internevo_ops/rotary_embedding.py b/deeplink_ext/internevo_ops/rotary_embedding.py
@@ -4,7 +4,8 @@
 
 platform_type = deeplink_ext_get_platform_type()
 if platform_type == PlatformType.TORCH_NPU:
-    from ._rotary_embedding_npu import ApplyRotaryEmb
+    # from ._rotary_embedding_npu import ApplyRotaryEmb
+    from .rotary_embedding_fallback import ApplyRotaryEmbTorch as ApplyRotaryEmb
 elif platform_type == PlatformType.TORCH_DIPU:
     from ._rotary_embedding_dipu import ApplyRotaryEmb
 else:

diff --git a/deeplink_ext/interntrain_ops/rms_norm.py b/deeplink_ext/interntrain_ops/rms_norm.py
@@ -4,7 +4,8 @@
 
 platform_type = deeplink_ext_get_platform_type()
 if platform_type == PlatformType.TORCH_NPU:
-    from ._mixed_rms_norm_npu import MixedFusedRMSNorm
+    # from ._mixed_rms_norm_npu import MixedFusedRMSNorm
+    from .rms_norm_fallback import MixedRMSNormTorch as MixedFusedRMSNorm
 elif platform_type == PlatformType.TORCH_DIPU:
     from ._mixed_rms_norm_dipu import MixedFusedRMSNorm
 else: