refactor python code

csrc/extensions.cpp

@@ -1,4 +1,4 @@
 // Copyright (c) 2023, DeepLink.
 
 #include <cstdint>
 #include <string>
@@ -60,10 +60,11 @@
             eps);
 }
 
-void extApplyRotary(at::Tensor& output, const at::Tensor& input,
+void extApplyRotary(const at::Tensor& input1, const at::Tensor& input2,
                     const at::Tensor& cos, const at::Tensor& sin,
-                    const bool conj, const bool interleaved) {
-  callDiopi(diopiRotaryEmbedding, output, input, cos, sin, conj, interleaved);
+                    at::Tensor& output1, at::Tensor& output2,
+                    const bool conj) {
+  callDiopi(diopiApplyRotary, output1, output2, input1, input2, cos, sin, conj, false);
 }
 
 auto extMultiHeadAttention(at::Tensor& q, at::Tensor& k, at::Tensor& v,
@@ -443,7 +444,7 @@
     m.def("rms_norm_backward", &extRmsNormBackward,
           "deeplink ext_rms_norm_backward");
   }
-  if (&diopiRotaryEmbedding != nullptr) {
+  if (&diopiApplyRotary != nullptr) {
     m.def("apply_rotary", &extApplyRotary, "deeplink ext_apply_rotary");
   }
   if (&diopiMultiHeadAttention != nullptr) {

deeplink_ext/__init__.py

@@ -9,9 +9,6 @@ def _init():
     platform_type = deeplink_ext_get_platform_type()
     if platform_type == PlatformType.TORCH_DIPU:
         import torch_dipu
-    elif platform_type == PlatformType.TORCH_NPU:
-        warnings.warn("DeepLinkExt using torch_npu ...", ImportWarning)
-        import torch_npu
     else:
         raise ImportError
 

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/internevo_ops/__init__.py

@@ -1,46 +1,25 @@
 # 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,
+    FlashCrossAttention,
+    FlashSelfAttention,
+)
 
-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, ApplyRotaryEmbQKV_, apply_rotary
 
 __all__ = [
     "AdamW",
@@ -52,4 +31,6 @@
     "flash_attn_varlen_func",
     "MixedFusedRMSNorm",
     "ApplyRotaryEmb",
+    "ApplyRotaryEmbQKV_",
+    "apply_rotary",
 ]