mindspore 2.2.14__cp39-cp39-manylinux1_x86_64.whl → 2.3.0rc1__cp39-cp39-manylinux1_x86_64.whl

mindspore/lib/plugin/ascend/custom_ascendc_ops/op_impl/ai_core/tbe/custom_ascendc_ops_impl/dynamic/add_custom.cpp

@@ -0,0 +1,81 @@
+/*
+ * Copyright (c) Huawei Technologies Co., Ltd. 2022-2023. All rights reserved.
+ *
+ * Function : z = x + y
+ * This sample is a very basic sample that implements vector add on Ascend platform.
+ */
+#include "kernel_operator.h"
+using namespace AscendC;
+
+constexpr int32_t BUFFER_NUM = 2;
+
+class KernelAdd {
+ public:
+  __aicore__ inline KernelAdd() {}
+  __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength, uint32_t tileNum) {
+    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
+    this->blockLength = totalLength / GetBlockNum();
+    this->tileNum = tileNum;
+    ASSERT(tileNum != 0 && "tile num can not be zero!");
+    this->tileLength = this->blockLength / tileNum / BUFFER_NUM;
+
+    xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x + this->blockLength * GetBlockIdx(), this->blockLength);
+    yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y + this->blockLength * GetBlockIdx(), this->blockLength);
+    zGm.SetGlobalBuffer((__gm__ DTYPE_Z *)z + this->blockLength * GetBlockIdx(), this->blockLength);
+    pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
+    pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_Y));
+    pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(DTYPE_Z));
+  }
+  __aicore__ inline void Process() {
+    int32_t loopCount = this->tileNum * BUFFER_NUM;
+    for (int32_t i = 0; i < loopCount; i++) {
+      CopyIn(i);
+      Compute(i);
+      CopyOut(i);
+    }
+  }
+
+ private:
+  __aicore__ inline void CopyIn(int32_t progress) {
+    LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
+    LocalTensor<DTYPE_Y> yLocal = inQueueY.AllocTensor<DTYPE_Y>();
+    DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
+    DataCopy(yLocal, yGm[progress * this->tileLength], this->tileLength);
+    inQueueX.EnQue(xLocal);
+    inQueueY.EnQue(yLocal);
+  }
+  __aicore__ inline void Compute(int32_t progress) {
+    LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
+    LocalTensor<DTYPE_Y> yLocal = inQueueY.DeQue<DTYPE_Y>();
+    LocalTensor<DTYPE_Z> zLocal = outQueueZ.AllocTensor<DTYPE_Z>();
+    Add(zLocal, xLocal, yLocal, this->tileLength);
+    outQueueZ.EnQue<DTYPE_Z>(zLocal);
+    inQueueX.FreeTensor(xLocal);
+    inQueueY.FreeTensor(yLocal);
+  }
+  __aicore__ inline void CopyOut(int32_t progress) {
+    LocalTensor<DTYPE_Z> zLocal = outQueueZ.DeQue<DTYPE_Z>();
+    DataCopy(zGm[progress * this->tileLength], zLocal, this->tileLength);
+    outQueueZ.FreeTensor(zLocal);
+  }
+
+ private:
+  TPipe pipe;
+  TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
+  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
+  GlobalTensor<DTYPE_X> xGm;
+  GlobalTensor<DTYPE_Y> yGm;
+  GlobalTensor<DTYPE_Z> zGm;
+  uint32_t blockLength = 0;
+  uint32_t tileNum = 0;
+  uint32_t tileLength = 0;
+};
+
+extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {
+  GET_TILING_DATA(tilingData, tiling);
+  KernelAdd op;
+  op.Init(x, y, z, tilingData.totalLength, tilingData.tileNum);
+  if (TILING_KEY_IS(1)) {
+    op.Process();
+  }
+}

mindspore/lib/plugin/ascend/custom_ascendc_ops/op_impl/ai_core/tbe/custom_ascendc_ops_impl/dynamic/add_custom.py

@@ -0,0 +1,134 @@
+
+import os, sys
+import ctypes
+import json
+import shutil
+from tbe.common.platform import get_soc_spec
+from tbe.common.utils import para_check
+from tbe.tikcpp import compile_op, replay_op, check_op_cap, generalize_op_params, get_code_channel, OpInfo
+from tbe.common.buildcfg import get_default_build_config
+from impl.util.platform_adapter import tbe_register
+from tbe.common.buildcfg import get_current_build_config
+PYF_PATH = os.path.dirname(os.path.realpath(__file__))
+
+DTYPE_MAP = {"float32": ["DT_FLOAT", "float"],
+    "float16": ["DT_FLOAT16", "half"],
+    "int8": ["DT_INT8", "int8_t"],
+    "int16": ["DT_INT16", "int16_t"],
+    "int32": ["DT_INT32", "int32_t"],
+    "int64": ["DT_INT64", "int64_t"],
+    "uint1": ["DT_UINT1", "uint8_t"],
+    "uint8": ["DT_UINT8", "uint8_t"],
+    "uint16": ["DT_UINT16", "uint16_t"],
+    "uint32": ["DT_UINT32", "uint32_t"],
+    "uint64": ["DT_UINT64", "uint64_t"],
+    "bool": ["DT_BOOL", "bool"],
+    "double": ["DT_DOUBLE", "double"],
+    "dual": ["DT_DUAL", "unknown"],
+    "dual_sub_int8": ["DT_DUAL_SUB_INT8", "unknown"],
+    "dual_sub_uint8": ["DT_DUAL_SUB_UINT8", "unknown"],
+    "string": ["DT_STRING", "unknown"],
+    "complex64": ["DT_COMPLEX64", "unknown"],
+    "complex128": ["DT_COMPLEX128", "unknown"],
+    "qint8": ["DT_QINT8", "unknown"],
+    "qint16": ["DT_QINT16", "unknown"],
+    "qint32": ["DT_QINT32", "unknown"],
+    "quint8": ["DT_QUINT8", "unknown"],
+    "quint16": ["DT_QUINT16", "unknown"],
+    "resource": ["DT_RESOURCE", "unknown"],
+    "string_ref": ["DT_STRING_REF", "unknown"],
+    "int4": ["DT_INT4", "int8_t"],
+    "bfloat16": ["DT_BF16", "bfloat16_t"]}
+
+def get_dtype_fmt_options(__inputs__, __outputs__):
+    options = []
+    for x in __inputs__ + __outputs__:
+        x_n = x.get("param_name").upper()
+        x_fmt = x.get("format")
+        x_dtype = x.get("dtype")
+        options.append("-DDTYPE_{n}={t}".format(n=x_n, t=DTYPE_MAP.get(x_dtype)[1]))
+        options.append("-DORIG_DTYPE_{n}={ot}".format(n=x_n, ot=DTYPE_MAP.get(x_dtype)[0]))
+        options.append("-DFORMAT_{n}=FORMAT_{f}".format(n=x_n, f=x_fmt))
+    return options
+
+def load_dso(so_path):
+    try:
+        ctypes.CDLL(so_path)
+    except OSError as error :
+        print(error)
+        raise RuntimeError("cannot open %s" %(so_path))
+    else:
+        print("load so succ ", so_path)
+
+def _build_args(x, y, z):
+    __inputs__ = []
+    for arg in [x, y]:
+        if arg != None:
+            if type(arg) is list:
+                if len(arg) == 0:
+                    continue
+                __inputs__.append(arg[0])
+            else:
+                __inputs__.append(arg)
+    __outputs__ = []
+    for arg in [z]:
+        if arg != None:
+            if type(arg) is list:
+                if len(arg) == 0:
+                    continue
+                __outputs__.append(arg[0])
+            else:
+                __outputs__.append(arg)
+    __attrs__ = []
+    return __inputs__, __outputs__, __attrs__
+
+@tbe_register.register_operator("AddCustom")
+@para_check.check_op_params(para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_OUTPUT, para_check.KERNEL_NAME)
+def add_custom(x, y, z, kernel_name="add_custom", impl_mode=""):
+    if get_current_build_config("enable_op_prebuild"):
+        return
+    __inputs__, __outputs__, __attrs__ = _build_args(x, y, z)
+    options = get_dtype_fmt_options(__inputs__, __outputs__)
+    options += ["-x", "cce"]
+    ccec = os.environ.get('CCEC_REAL_PATH')
+    if ccec is None:
+        ccec = shutil.which("ccec")
+    if ccec != None:
+        ccec_path = os.path.dirname(ccec)
+        tikcpp_path = os.path.realpath(os.path.join(ccec_path, "..", "..", "tikcpp"))
+    else:
+        tikcpp_path = os.path.realpath("/usr/local/Ascend/latest/compiler/tikcpp")
+    options.append("-I" + tikcpp_path)
+    options.append("-I" + os.path.join(tikcpp_path, "tikcfw"))
+    options.append("-I" + os.path.join(tikcpp_path, "tikcfw", "impl"))
+    options.append("-I" + os.path.join(tikcpp_path, "tikcfw", "interface"))
+    options.append("-I" + os.path.join(PYF_PATH, "..", "ascendc", "common"))
+    if impl_mode == "high_performance":
+        options.append("-DHIGH_PERFORMANCE=1")
+    elif impl_mode == "high_precision":
+        options.append("-DHIGH_PRECISION=1")
+    if get_default_build_config("enable_deterministic_mode") == 1:
+        options.append("-DDETEMINISTIC_MODE=1")
+    origin_func_name = "add_custom"
+    ascendc_src_dir = "add_custom"
+    ascendc_src_file = "add_custom.cpp"
+    src = os.path.join(PYF_PATH, "..", "ascendc", ascendc_src_dir, ascendc_src_file)
+    if not os.path.exists(src):
+        src = os.path.join(PYF_PATH, ascendc_src_file)
+
+    print("start compile Ascend C operator AddCustom. kernel name is add_custom")
+    op_type = "AddCustom"
+    code_channel = get_code_channel(src, kernel_name, op_type, options)
+    op_info = OpInfo(kernel_name = kernel_name, op_type = op_type, inputs = __inputs__, outputs = __outputs__,\
+        attrs = __attrs__, impl_mode = impl_mode, origin_inputs=[x, y], origin_outputs = [z])
+    compile_op(src, origin_func_name, op_info, options, code_channel, '{}')
+
+def op_select_format(x, y, z, impl_mode=""):
+    __inputs__, __outputs__, __attrs__ = _build_args(x, y, z)
+    result = check_op_cap("op_select_format", "AddCustom", __inputs__, __outputs__, __attrs__)
+    return result.decode("utf-8")
+
+def get_op_specific_info(x, y, z, impl_mode=""):
+    __inputs__, __outputs__, __attrs__ = _build_args(x, y, z)
+    result = check_op_cap("get_op_specific_info", "AddCustom", __inputs__, __outputs__, __attrs__)
+    return result.decode("utf-8")

mindspore/lib/plugin/ascend/custom_ascendc_ops/op_impl/ai_core/tbe/custom_ascendc_ops_impl/dynamic/decoder_kv_cache.cpp

@@ -0,0 +1,192 @@
+/**
+ * Copyright 2023 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "kernel_operator.h"
+using namespace AscendC;
+namespace {
+constexpr int32_t kBufferNum = 2;
+const int64_t kDivisor = 4;
+static __aicore__ inline int64_t CeilRound(int64_t value, int64_t divisor) {
+  if (divisor == 0) {
+    return 0;
+  }
+  return (value + divisor - 1) / divisor * divisor;
+}
+}  // namespace
+
+template <typename T>
+class KernelDecoderKvCache {
+ public:
+  __aicore__ inline KernelDecoderKvCache() {}
+
+  __aicore__ inline void GetValidSeqLen(GM_ADDR valid_seq_len) {
+    int64_t valid_seq_len_ub_size = CeilRound(b_, kDivisor);
+    valid_seq_len_gm_.SetGlobalBuffer((__gm__ int64_t *)valid_seq_len, valid_seq_len_ub_size);
+    pipe_.InitBuffer(valid_seq_len_queue_, 1, valid_seq_len_ub_size * sizeof(int64_t));
+    valid_seq_len_tensor_ = valid_seq_len_queue_.AllocTensor<int64_t>();
+    pipe_barrier((pipe_t)PIPE_ALL);
+    DataCopy(valid_seq_len_tensor_, valid_seq_len_gm_, valid_seq_len_ub_size);
+    pipe_barrier((pipe_t)PIPE_ALL);
+  }
+
+  __aicore__ inline void SplitBh() {
+    if (core_idx_ != core_num_ - 1) {
+      split_bh_ = f_split_bh_;
+      former_block_bh_ = f_f_bh_;
+      tail_block_bh_ = f_t_bh_;
+    } else {
+      split_bh_ = t_split_bh_;
+      former_block_bh_ = t_f_bh_;
+      tail_block_bh_ = t_t_bh_;
+    }
+  }
+
+  __aicore__ inline void Update(GM_ADDR cache, GM_ADDR update, LocalTensor<T> update_in_local_tensor) {
+    for (int64_t i = 0; i < split_bh_; i++) {
+      int64_t block_bh;
+      if (i != split_bh_ - 1) {
+        block_bh = former_block_bh_;
+      } else {
+        block_bh = tail_block_bh_;
+      }
+      update_gm_.SetGlobalBuffer((__gm__ T *)update + core_idx_ * update_core_stride_ + i * former_block_bh_ * us_ * d_,
+                                 block_bh * us_ * d_);
+
+      DataCopy(update_in_local_tensor, update_gm_, block_bh * us_ * d_);
+      pipe_barrier((pipe_t)PIPE_ALL);
+
+      update_queue_.EnQue(update_in_local_tensor);
+      LocalTensor<T> update_in_local_tensor_out = update_queue_.DeQue<T>();
+
+      for (int64_t j = 0; j < block_bh; j++) {
+        int64_t bh_idx = core_idx_ * former_bh_ + i * former_block_bh_ + j;
+        auto b_idx = bh_idx / h_;
+        pipe_barrier((pipe_t)PIPE_ALL);
+        auto s_idx = valid_seq_len_tensor_.GetValue(b_idx);
+        pipe_barrier((pipe_t)PIPE_ALL);
+        if (s_idx < 0) {
+          continue;
+        }
+        if (s_idx >= s_) {
+          s_idx = s_idx % s_;
+        }
+        out_gm_.SetGlobalBuffer((__gm__ T *)cache + bh_idx * s_ * d_ + s_idx * d_, us_ * d_);
+        int64_t src_offset = j * us_ * d_;
+        pipe_barrier((pipe_t)PIPE_ALL);
+        DataCopy(out_gm_, update_in_local_tensor_out[src_offset], us_ * d_);
+      }
+    }
+  }
+
+  __aicore__ inline void InitParam(GM_ADDR tiling) {
+    GET_TILING_DATA(tiling_data, tiling);
+    core_num_ = tiling_data.core_num;
+    b_ = tiling_data.b;
+    h_ = tiling_data.h;
+    s_ = tiling_data.s;
+    d_ = tiling_data.d;
+    us_ = tiling_data.us;
+    former_bh_ = tiling_data.former_bh;
+    tail_bh_ = tiling_data.tail_bh;
+    f_split_bh_ = tiling_data.f_split_bh;
+    f_f_bh_ = tiling_data.f_f_bh;
+    f_t_bh_ = tiling_data.f_t_bh;
+    t_split_bh_ = tiling_data.t_split_bh;
+    t_f_bh_ = tiling_data.t_f_bh;
+    t_t_bh_ = tiling_data.t_t_bh;
+  }
+
+  __aicore__ inline void Process(GM_ADDR cache, GM_ADDR update, GM_ADDR valid_seq_len, GM_ADDR batch_index,
+                                 GM_ADDR new_max_seq_len, GM_ADDR cur_max_seq_len, GM_ADDR tiling) {
+    InitParam(tiling);
+    core_idx_ = GetBlockIdx();
+    if (core_idx_ >= core_num_) {
+      return;
+    }
+
+    GetValidSeqLen(valid_seq_len);
+    update_core_stride_ = former_bh_ * us_ * d_;
+    cache_core_stride_ = former_bh_ * s_ * d_;
+
+    SplitBh();
+
+    pipe_.InitBuffer(update_queue_, kBufferNum, former_block_bh_ * us_ * d_ * sizeof(T));
+    LocalTensor<T> update_in_local_tensor = update_queue_.AllocTensor<T>();
+
+    Update(cache, update, update_in_local_tensor);
+
+    valid_seq_len_queue_.FreeTensor(valid_seq_len_tensor_);
+    update_queue_.FreeTensor(update_in_local_tensor);
+  }
+
+ private:
+  // gm
+  GlobalTensor<T> update_gm_;
+  GlobalTensor<int64_t> valid_seq_len_gm_;
+  GlobalTensor<int64_t> new_max_seq_len_gm_;
+  GlobalTensor<T> out_gm_;
+
+  // local
+  LocalTensor<int64_t> valid_seq_len_tensor_;
+
+  TPipe pipe_;
+  // create queues for input, in this case depth is equal to buffer num
+  TQue<QuePosition::VECIN, 1> update_queue_;
+  TQue<QuePosition::VECIN, 1> valid_seq_len_queue_;
+  TQue<QuePosition::VECIN, 1> new_max_seq_len_queue_;
+
+  int64_t split_bh_ = 0;
+  int64_t former_block_bh_ = 0;
+  int64_t tail_block_bh_ = 0;
+
+  int64_t core_idx_ = 0;
+  int64_t cache_core_stride_ = 0;
+  int64_t cache_block_length = 0;
+  int64_t update_core_stride_ = 0;
+  int64_t update_block_length_ = 0;
+
+  int64_t core_num_ = 0;
+  int64_t b_ = 0;
+  int64_t h_ = 0;
+  int64_t s_ = 0;
+  int64_t d_ = 0;
+  int64_t us_ = 0;
+
+  int64_t former_bh_ = 0;
+  int64_t tail_bh_ = 0;
+  int64_t f_split_bh_ = 0;
+  int64_t f_f_bh_ = 0;
+  int64_t f_t_bh_ = 0;
+  int64_t t_split_bh_ = 0;
+  int64_t t_f_bh_ = 0;
+  int64_t t_t_bh_ = 0;
+};
+
+extern "C" __global__ __aicore__ void decoder_kv_cache(GM_ADDR cache, GM_ADDR update, GM_ADDR valid_seq_len,
+                                                       GM_ADDR batch_index, GM_ADDR seq_len_axis,
+                                                       GM_ADDR new_max_seq_len, GM_ADDR cur_max_seq_len, GM_ADDR out,
+                                                       GM_ADDR workspace, GM_ADDR tiling) {
+  if (TILING_KEY_IS(1)) {
+    KernelDecoderKvCache<int8_t> op;
+    op.Process(cache, update, valid_seq_len, batch_index, new_max_seq_len, cur_max_seq_len, tiling);
+  } else if (TILING_KEY_IS(2)) {
+    KernelDecoderKvCache<int16_t> op;
+    op.Process(cache, update, valid_seq_len, batch_index, new_max_seq_len, cur_max_seq_len, tiling);
+  } else if (TILING_KEY_IS(4)) {
+    KernelDecoderKvCache<int32_t> op;
+    op.Process(cache, update, valid_seq_len, batch_index, new_max_seq_len, cur_max_seq_len, tiling);
+  }
+}

mindspore/lib/plugin/ascend/custom_ascendc_ops/op_impl/ai_core/tbe/custom_ascendc_ops_impl/dynamic/decoder_kv_cache.py

@@ -0,0 +1,134 @@
+
+import os, sys
+import ctypes
+import json
+import shutil
+from tbe.common.platform import get_soc_spec
+from tbe.common.utils import para_check
+from tbe.tikcpp import compile_op, replay_op, check_op_cap, generalize_op_params, get_code_channel, OpInfo
+from tbe.common.buildcfg import get_default_build_config
+from impl.util.platform_adapter import tbe_register
+from tbe.common.buildcfg import get_current_build_config
+PYF_PATH = os.path.dirname(os.path.realpath(__file__))
+
+DTYPE_MAP = {"float32": ["DT_FLOAT", "float"],
+    "float16": ["DT_FLOAT16", "half"],
+    "int8": ["DT_INT8", "int8_t"],
+    "int16": ["DT_INT16", "int16_t"],
+    "int32": ["DT_INT32", "int32_t"],
+    "int64": ["DT_INT64", "int64_t"],
+    "uint1": ["DT_UINT1", "uint8_t"],
+    "uint8": ["DT_UINT8", "uint8_t"],
+    "uint16": ["DT_UINT16", "uint16_t"],
+    "uint32": ["DT_UINT32", "uint32_t"],
+    "uint64": ["DT_UINT64", "uint64_t"],
+    "bool": ["DT_BOOL", "bool"],
+    "double": ["DT_DOUBLE", "double"],
+    "dual": ["DT_DUAL", "unknown"],
+    "dual_sub_int8": ["DT_DUAL_SUB_INT8", "unknown"],
+    "dual_sub_uint8": ["DT_DUAL_SUB_UINT8", "unknown"],
+    "string": ["DT_STRING", "unknown"],
+    "complex64": ["DT_COMPLEX64", "unknown"],
+    "complex128": ["DT_COMPLEX128", "unknown"],
+    "qint8": ["DT_QINT8", "unknown"],
+    "qint16": ["DT_QINT16", "unknown"],
+    "qint32": ["DT_QINT32", "unknown"],
+    "quint8": ["DT_QUINT8", "unknown"],
+    "quint16": ["DT_QUINT16", "unknown"],
+    "resource": ["DT_RESOURCE", "unknown"],
+    "string_ref": ["DT_STRING_REF", "unknown"],
+    "int4": ["DT_INT4", "int8_t"],
+    "bfloat16": ["DT_BF16", "bfloat16_t"]}
+
+def get_dtype_fmt_options(__inputs__, __outputs__):
+    options = []
+    for x in __inputs__ + __outputs__:
+        x_n = x.get("param_name").upper()
+        x_fmt = x.get("format")
+        x_dtype = x.get("dtype")
+        options.append("-DDTYPE_{n}={t}".format(n=x_n, t=DTYPE_MAP.get(x_dtype)[1]))
+        options.append("-DORIG_DTYPE_{n}={ot}".format(n=x_n, ot=DTYPE_MAP.get(x_dtype)[0]))
+        options.append("-DFORMAT_{n}=FORMAT_{f}".format(n=x_n, f=x_fmt))
+    return options
+
+def load_dso(so_path):
+    try:
+        ctypes.CDLL(so_path)
+    except OSError as error :
+        print(error)
+        raise RuntimeError("cannot open %s" %(so_path))
+    else:
+        print("load so succ ", so_path)
+
+def _build_args(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out):
+    __inputs__ = []
+    for arg in [cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len]:
+        if arg != None:
+            if type(arg) is list:
+                if len(arg) == 0:
+                    continue
+                __inputs__.append(arg[0])
+            else:
+                __inputs__.append(arg)
+    __outputs__ = []
+    for arg in [out]:
+        if arg != None:
+            if type(arg) is list:
+                if len(arg) == 0:
+                    continue
+                __outputs__.append(arg[0])
+            else:
+                __outputs__.append(arg)
+    __attrs__ = []
+    return __inputs__, __outputs__, __attrs__
+
+@tbe_register.register_operator("DecoderKvCache")
+@para_check.check_op_params(para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT, para_check.REQUIRED_OUTPUT, para_check.KERNEL_NAME)
+def decoder_kv_cache(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out, kernel_name="decoder_kv_cache", impl_mode=""):
+    if get_current_build_config("enable_op_prebuild"):
+        return
+    __inputs__, __outputs__, __attrs__ = _build_args(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out)
+    options = get_dtype_fmt_options(__inputs__, __outputs__)
+    options += ["-x", "cce"]
+    ccec = os.environ.get('CCEC_REAL_PATH')
+    if ccec is None:
+        ccec = shutil.which("ccec")
+    if ccec != None:
+        ccec_path = os.path.dirname(ccec)
+        tikcpp_path = os.path.realpath(os.path.join(ccec_path, "..", "..", "tikcpp"))
+    else:
+        tikcpp_path = os.path.realpath("/usr/local/Ascend/latest/compiler/tikcpp")
+    options.append("-I" + tikcpp_path)
+    options.append("-I" + os.path.join(tikcpp_path, "tikcfw"))
+    options.append("-I" + os.path.join(tikcpp_path, "tikcfw", "impl"))
+    options.append("-I" + os.path.join(tikcpp_path, "tikcfw", "interface"))
+    options.append("-I" + os.path.join(PYF_PATH, "..", "ascendc", "common"))
+    if impl_mode == "high_performance":
+        options.append("-DHIGH_PERFORMANCE=1")
+    elif impl_mode == "high_precision":
+        options.append("-DHIGH_PRECISION=1")
+    if get_default_build_config("enable_deterministic_mode") == 1:
+        options.append("-DDETEMINISTIC_MODE=1")
+    origin_func_name = "decoder_kv_cache"
+    ascendc_src_dir = "decoder_kv_cache"
+    ascendc_src_file = "decoder_kv_cache.cpp"
+    src = os.path.join(PYF_PATH, "..", "ascendc", ascendc_src_dir, ascendc_src_file)
+    if not os.path.exists(src):
+        src = os.path.join(PYF_PATH, ascendc_src_file)
+
+    print("start compile Ascend C operator DecoderKvCache. kernel name is decoder_kv_cache")
+    op_type = "DecoderKvCache"
+    code_channel = get_code_channel(src, kernel_name, op_type, options)
+    op_info = OpInfo(kernel_name = kernel_name, op_type = op_type, inputs = __inputs__, outputs = __outputs__,\
+        attrs = __attrs__, impl_mode = impl_mode, origin_inputs=[cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len], origin_outputs = [out])
+    compile_op(src, origin_func_name, op_info, options, code_channel, '{}')
+
+def op_select_format(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out, impl_mode=""):
+    __inputs__, __outputs__, __attrs__ = _build_args(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out)
+    result = check_op_cap("op_select_format", "DecoderKvCache", __inputs__, __outputs__, __attrs__)
+    return result.decode("utf-8")
+
+def get_op_specific_info(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out, impl_mode=""):
+    __inputs__, __outputs__, __attrs__ = _build_args(cache, update, valid_seq_len, batch_index, seq_len_axis, new_max_seq_len, cur_max_seq_len, out)
+    result = check_op_cap("get_op_specific_info", "DecoderKvCache", __inputs__, __outputs__, __attrs__)
+    return result.decode("utf-8")