[Ascend] Wx/fix rope and repeat (#1156)

* fix rope for ascend speed

* add dim check for rope

* optimize rope using cat instead of repeat

* reimpl repeat and fix for corner case

* update device config for sum on ascend

---------

Co-authored-by: yewentao <[email protected]>

Author: POI-WX

impl/ascend/device_configs.py

@@ -267,6 +267,12 @@
         ),
     ),
 
+    'reduce_partial_op': dict(
+        name=['sum'],
+        atol=1e-3,
+        rtol=1e-4,
+    ),
+
     'reduce_partial_op_1': dict(
         name=['std'],
         tensor_para=dict(

impl/ascend_npu/diopi_impl/functions_ext/rotary_embedding.cpp

@@ -35,6 +35,8 @@ DIOPI_API diopiError_t diopiRotaryEmbedding(diopiContextHandle_t ctx, diopiTenso
     }
 
     BEGIN_CALL_ACL_OP(out, x, cos, sin);
+    TORCH_CHECK(xAt.size(-1) == 2 * cosAt.size(-1) && xAt.size(-1) == 2 * sinAt.size(-1),
+                "The size of the last dimension of x must be twice the size of the corresponding dimensions of cos and sin!");
     if (xAt.numel() == 0) {
         END_CALL_ACL_OP();
     }
@@ -49,10 +51,12 @@ DIOPI_API diopiError_t diopiRotaryEmbedding(diopiContextHandle_t ctx, diopiTenso
     at::Tensor outView = viewAs4D(outAt);
     at::Tensor cosView = viewAs4D(cosAt);
     at::Tensor sinView = viewAs4D(sinAt);
-    at::Tensor cosRepeated = op_api::repeat(cosView, {1, 1, 1, 2});
-    at::Tensor sinRepeated = op_api::repeat(sinView, {1, 1, 1, 2});
+    // To meet the ascend kernel requirement: the last dimension size of cos and sin is the same as the dimension size corresponding to input, use cat op to
+    // concatenate in the last dimension.
+    at::Tensor cosCat = op_api::cat({cosView, cosView}, -1);
+    at::Tensor sinCat = op_api::cat({sinView, sinView}, -1);
     if (conj) {
-        op_api::neg_(sinRepeated);
+        op_api::neg_(sinCat);
     }
 
     // According to API document
@@ -63,7 +67,7 @@ DIOPI_API diopiError_t diopiRotaryEmbedding(diopiContextHandle_t ctx, diopiTenso
 
     std::vector<at::Tensor> chunkResult = xView.chunk(2, -1);
     at::Tensor xNew = op_api::cat({chunkResult[1] * (-1), chunkResult[0]}, -1);
-    at::Tensor result = op_api::mul(cosRepeated, xView) + op_api::mul(sinRepeated, xNew);
+    at::Tensor result = op_api::mul(cosCat, xView) + op_api::mul(sinCat, xNew);
     outView.copy_(result);
 
     END_CALL_ACL_OP();

impl/ascend_npu/diopi_impl/repeat.cpp

@@ -5,23 +5,24 @@
  */
 
 #include "helper.hpp"
-#include "op_plugin/AclOpsInterface.h"
+#include "op_plugin/OpApiInterface.h"
+#include "op_plugin/utils/op_api_common.h"
 
 extern "C" {
 diopiError_t diopiRepeat(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiConstTensorHandle_t input, diopiSize_t repeatSize) {
     BEGIN_CALL_ACL_OP(out, input, repeatSize);
-    std::vector<int64_t> inputShape(inputAt.sizes().cbegin(), inputAt.sizes().cend());
-
-    if (inputShape.size() < repeatSize.len) {
-        while (inputShape.size() < repeatSize.len) {
-            inputShape.insert(inputShape.begin(), 1);
-        }
-
-        inputAt = impl::aten::viewStorage(inputAt, inputShape);
+    TORCH_CHECK(inputAt.dim() <= repeatSize.len, "repeats size should not be smaller than input tensor dim on ascend!");
+    // When repeatSize.len is equal to 0, out is the same as input.
+    if (repeatSize.len == 0) {
+        outAt.copy_(inputAt);
+        END_CALL_ACL_OP();
     }
 
-    at_npu::native::OpPreparation::markAsOutputForApplyTensor(outAt);
-    outAt = acl_op::repeat(inputAt, repeatSizeAt);
+    std::vector<int64_t> inputShape = inputAt.sizes().vec();
+    inputShape.insert(inputShape.begin(), repeatSize.len - inputAt.dim(), 1);
+    inputAt = impl::aten::viewStorage(inputAt, inputShape);
+
+    EXEC_NPU_CMD(aclnnRepeat, inputAt, repeatSizeAt, outAt);
     END_CALL_ACL_OP();
 }