refactor python code

refactor python code by ops

Author: Wrench-Git

deeplink_ext/ascend_speed/_flash_attention_dipu.py

@@ -9,7 +9,6 @@
 
 
 class FlashSelfAttention(torch.autograd.Function):
-
     @staticmethod
     def forward(
         ctx, q, k, v, attention_mask, dropout_p, softmax_scale, head_num, input_layout

deeplink_ext/ascend_speed/_rms_norm_dipu.py

@@ -9,7 +9,6 @@
 
 
 class RMSNorm(torch.autograd.Function):
-
     @staticmethod
     def forward(ctx, hidden_states, weight, eps):
         output = torch.empty_like(hidden_states)

deeplink_ext/ascend_speed/_scaled_masked_softmax_dipu.py

@@ -11,7 +11,6 @@
 
 
 class ScaledMaskedSoftmax(torch.autograd.Function):
-
     @staticmethod
     def forward(ctx, input, mask, scale, fixed_triu_mask):
         out = torch.empty_like(input)

deeplink_ext/ascend_speed/_scaled_masked_softmax_npu.py

@@ -7,7 +7,6 @@
 
 
 class ScaledMaskedSoftmax(torch.autograd.Function):
-
     @staticmethod
     def forward(ctx, input, mask, scale, fixed_triu_mask):
         out = torch_npu.npu_scaled_masked_softmax(input, mask, scale, fixed_triu_mask)

deeplink_ext/easyllm_ops/__init__.py

@@ -3,40 +3,22 @@
 _not_impl = "[deeplink_ext] {op_name} is not implemented in diopi. Falling back to the slower torch implementation."
 
 try:
-    from .adamw import AdamW
+    from deeplink_ext.ops.adamw import AdamW
 except Exception as e:
     print(_not_impl.format(op_name="adamw"))
     from torch.optim import AdamW
 
-try:
-    from .flash_attention import (
-        flash_attn_qkvpacked_func,
-        flash_attn_kvpacked_func,
-        flash_attn_func,
-        flash_attn_varlen_qkvpacked_func,
-        flash_attn_varlen_kvpacked_func,
-        flash_attn_varlen_func,
-    )
-except Exception as e:
-    print(_not_impl.format(op_name="flash attention"))
-    from .flash_attention_fallback import (
-        flash_attn_qkvpacked_func_torch as flash_attn_qkvpacked_func,
-        flash_attn_kvpacked_func_torch as flash_attn_kvpacked_func,
-        flash_attn_func_torch as flash_attn_func,
-        flash_attn_varlen_qkvpacked_func_torch as flash_attn_varlen_qkvpacked_func,
-        flash_attn_varlen_kvpacked_func_torch as flash_attn_varlen_kvpacked_func,
-        flash_attn_varlen_func_torch as flash_attn_varlen_func,
-    )
-
-try:
-    from .rms_norm import rms_norm
-except:
-    print(
-        _not_impl.format(op_name="RMSNorm"),
-    )
-    from .rms_norm_fallback import rms_norm_torch as rms_norm
+from deeplink_ext.ops.flash_attention import (
+    flash_attn_qkvpacked_func,
+    flash_attn_kvpacked_func,
+    flash_attn_func,
+    flash_attn_varlen_qkvpacked_func,
+    flash_attn_varlen_kvpacked_func,
+    flash_attn_varlen_func,
+)
 
-from .bert_padding import pad_input, unpad_input, index_first_axis
+from deeplink_ext.ops.rms_norm import rms_norm
+from deeplink_ext.ops.bert_padding import pad_input, unpad_input, index_first_axis
 
 __all__ = [
     "AdamW",

deeplink_ext/easyllm_ops/adamw.py

@@ -1,46 +1,23 @@
 # Copyright (c) 2024, DeepLink.
 
-_not_impl = "[deeplink_ext] {op_name} is not implemented in diopi. Falling back to the slower torch implementation."
-
 try:
-    from .adamw import AdamW
+    from deeplink_ext.ops.adamw import AdamW
 except Exception as e:
     print(_not_impl.format(op_name="adamw"))
     from torch.optim import AdamW
 
-try:
-    from .flash_attention import (
-        flash_attn_qkvpacked_func,
-        flash_attn_kvpacked_func,
-        flash_attn_func,
-        flash_attn_varlen_qkvpacked_func,
-        flash_attn_varlen_kvpacked_func,
-        flash_attn_varlen_func,
-    )
-except Exception as e:
-    print(_not_impl.format(op_name="flash attention"))
-    from .flash_attention_fallback import (
-        flash_attn_qkvpacked_func_torch as flash_attn_qkvpacked_func,
-        flash_attn_kvpacked_func_torch as flash_attn_kvpacked_func,
-        flash_attn_func_torch as flash_attn_func,
-        flash_attn_varlen_qkvpacked_func_torch as flash_attn_varlen_qkvpacked_func,
-        flash_attn_varlen_kvpacked_func_torch as flash_attn_varlen_kvpacked_func,
-        flash_attn_varlen_func_torch as flash_attn_varlen_func,
-    )
+from deeplink_ext.ops.flash_attention import (
+    flash_attn_qkvpacked_func,
+    flash_attn_kvpacked_func,
+    flash_attn_func,
+    flash_attn_varlen_qkvpacked_func,
+    flash_attn_varlen_kvpacked_func,
+    flash_attn_varlen_func,
+)
 
-try:
-    from .rms_norm import MixedFusedRMSNorm
-except:
-    print(
-        _not_impl.format(op_name="RMSNorm"),
-    )
-    from .rms_norm_fallback import MixedRMSNormTorch as MixedFusedRMSNorm
+from deeplink_ext.ops.rms_norm import MixedFusedRMSNorm
 
-try:
-    from .rotary_embedding import ApplyRotaryEmb
-except:
-    print(_not_impl.format(op_name="rotary embedding"))
-    from .rotary_embedding_fallback import ApplyRotaryEmbTorch as ApplyRotaryEmb
+from deeplink_ext.ops.rotary_embedding import ApplyRotaryEmb
 
 __all__ = [
     "AdamW",