mindspore 2.4.1__cp39-none-any.whl → 2.5.0__cp39-none-any.whl

mindspore/lib/plugin/ascend/custom_ascendc_910/op_proto/inc/op_proto.h

@@ -0,0 +1,33 @@
+#ifndef OP_PROTO_H_
+#define OP_PROTO_H_
+
+#include "graph/operator_reg.h"
+#include "register/op_impl_registry.h"
+
+namespace ge {
+
+REG_OP(DecoderKvCache)
+    .INPUT(cache, ge::TensorType::ALL())
+    .INPUT(update, ge::TensorType::ALL())
+    .INPUT(valid_seq_len, ge::TensorType::ALL())
+    .INPUT(batch_index, ge::TensorType::ALL())
+    .INPUT(seq_len_axis, ge::TensorType::ALL())
+    .INPUT(new_max_seq_len, ge::TensorType::ALL())
+    .INPUT(cur_max_seq_len, ge::TensorType::ALL())
+    .OUTPUT(out, ge::TensorType::ALL())
+    .OP_END_FACTORY_REG(DecoderKvCache);
+
+REG_OP(PromptKvCache)
+    .INPUT(cache, ge::TensorType::ALL())
+    .INPUT(update, ge::TensorType::ALL())
+    .INPUT(valid_seq_len, ge::TensorType::ALL())
+    .INPUT(batch_index, ge::TensorType::ALL())
+    .INPUT(seq_len_axis, ge::TensorType::ALL())
+    .INPUT(new_max_seq_len, ge::TensorType::ALL())
+    .INPUT(cur_max_seq_len, ge::TensorType::ALL())
+    .OUTPUT(out, ge::TensorType::ALL())
+    .OP_END_FACTORY_REG(PromptKvCache);
+
+}
+
+#endif

mindspore/lib/plugin/ascend/custom_ascendc_910/version.info

@@ -0,0 +1 @@
+custom_opp_compiler_version=7.6.0.1.220

mindspore/lib/plugin/ascend/custom_ascendc_910b/framework/npu_supported_ops.json

@@ -0,0 +1,14 @@
+{
+  "AllFinite": {
+    "isGray": false,
+    "isHeavy": false
+  },
+  "DecoderKvCache": {
+    "isGray": false,
+    "isHeavy": false
+  },
+  "PromptKvCache": {
+    "isGray": false,
+    "isHeavy": false
+  }
+}

mindspore/lib/plugin/ascend/custom_ascendc_910b/op_api/include/aclnn_decoder_kv_cache.h

@@ -0,0 +1,59 @@
+
+/*
+ * calution: this file was generated automaticlly donot change it.
+*/
+
+#ifndef ACLNN_DECODER_KV_CACHE_H_
+#define ACLNN_DECODER_KV_CACHE_H_
+
+#include "aclnn/acl_meta.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/* funtion: aclnnDecoderKvCacheGetWorkspaceSize
+ * parameters :
+ * cache : required
+ * update : required
+ * validSeqLen : required
+ * batchIndex : required
+ * seqLenAxis : required
+ * newMaxSeqLen : required
+ * curMaxSeqLen : required
+ * out : required
+ * workspaceSize : size of workspace(output).
+ * executor : executor context(output).
+ */
+__attribute__((visibility("default")))
+aclnnStatus aclnnDecoderKvCacheGetWorkspaceSize(
+    const aclTensor *cache,
+    const aclTensor *update,
+    const aclTensor *validSeqLen,
+    const aclTensor *batchIndex,
+    const aclTensor *seqLenAxis,
+    const aclTensor *newMaxSeqLen,
+    const aclTensor *curMaxSeqLen,
+    const aclTensor *out,
+    uint64_t *workspaceSize,
+    aclOpExecutor **executor);
+
+/* funtion: aclnnDecoderKvCache
+ * parameters :
+ * workspace : workspace memory addr(input).
+ * workspaceSize : size of workspace(input).
+ * executor : executor context(input).
+ * stream : acl stream.
+ */
+__attribute__((visibility("default")))
+aclnnStatus aclnnDecoderKvCache(
+    void *workspace,
+    uint64_t workspaceSize,
+    aclOpExecutor *executor,
+    aclrtStream stream);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

mindspore/lib/plugin/ascend/custom_ascendc_910b/op_api/include/aclnn_prompt_kv_cache.h

@@ -0,0 +1,59 @@
+
+/*
+ * calution: this file was generated automaticlly donot change it.
+*/
+
+#ifndef ACLNN_PROMPT_KV_CACHE_H_
+#define ACLNN_PROMPT_KV_CACHE_H_
+
+#include "aclnn/acl_meta.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/* funtion: aclnnPromptKvCacheGetWorkspaceSize
+ * parameters :
+ * cache : required
+ * update : required
+ * validSeqLen : required
+ * batchIndex : required
+ * seqLenAxis : required
+ * newMaxSeqLen : required
+ * curMaxSeqLen : required
+ * out : required
+ * workspaceSize : size of workspace(output).
+ * executor : executor context(output).
+ */
+__attribute__((visibility("default")))
+aclnnStatus aclnnPromptKvCacheGetWorkspaceSize(
+    const aclTensor *cache,
+    const aclTensor *update,
+    const aclTensor *validSeqLen,
+    const aclTensor *batchIndex,
+    const aclTensor *seqLenAxis,
+    const aclTensor *newMaxSeqLen,
+    const aclTensor *curMaxSeqLen,
+    const aclTensor *out,
+    uint64_t *workspaceSize,
+    aclOpExecutor **executor);
+
+/* funtion: aclnnPromptKvCache
+ * parameters :
+ * workspace : workspace memory addr(input).
+ * workspaceSize : size of workspace(input).
+ * executor : executor context(input).
+ * stream : acl stream.
+ */
+__attribute__((visibility("default")))
+aclnnStatus aclnnPromptKvCache(
+    void *workspace,
+    uint64_t workspaceSize,
+    aclOpExecutor *executor,
+    aclrtStream stream);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

mindspore/lib/plugin/ascend/{custom_ascendc_ops/op_impl/ai_core/tbe/custom_ascendc_ops_impl → custom_ascendc_910b/op_impl/ai_core/tbe/custom_ascendc_910b_impl}/dynamic/all_finite.py

@@ -1,3 +1,8 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+"""
+Copyright (c) Huawei Technologies Co., Ltd. 2024-2025. All rights reserved.
+"""
 
 import os, sys
 import ctypes
@@ -6,6 +11,7 @@ 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.tikcpp.compile_op import CommonUtility, AscendCLogLevel
 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
@@ -28,6 +34,7 @@ DTYPE_MAP = {"float32": ["DT_FLOAT", "float"],
     "dual_sub_int8": ["DT_DUAL_SUB_INT8", "unknown"],
     "dual_sub_uint8": ["DT_DUAL_SUB_UINT8", "unknown"],
     "string": ["DT_STRING", "unknown"],
+    "complex32": ["DT_COMPLEX32", "unknown"],
     "complex64": ["DT_COMPLEX64", "unknown"],
     "complex128": ["DT_COMPLEX128", "unknown"],
     "qint8": ["DT_QINT8", "unknown"],
@@ -49,21 +56,23 @@ def add_dtype_fmt_option_single(x, x_n, is_ref: bool = False):
     options.append("-DORIG_DTYPE_{n}={ot}".format(n=x_n_in_kernel, ot=DTYPE_MAP.get(x_dtype)[0]))
     options.append("-DFORMAT_{n}=FORMAT_{f}".format(n=x_n_in_kernel, f=x_fmt))
     return options
- 
+
 def get_dtype_fmt_options(__inputs__, __outputs__):
     options = []
+    input_names = ['gradient']
+    output_names = ['is_finite']
     unique_param_name_set = set()
-    for x in __inputs__:
+    for idx, x in enumerate(__inputs__):
         if x is None:
             continue
-        x_n = x.get("param_name")[:-5].upper()
+        x_n = input_names[idx].upper()
         unique_param_name_set.add(x_n)
         options += add_dtype_fmt_option_single(x, x_n)
- 
-    for x in __outputs__:
+
+    for idx, x in enumerate(__outputs__):
         if x is None:
             continue
-        x_n = x.get("param_name")[:-5].upper()
+        x_n = output_names[idx].upper()
         if x_n in unique_param_name_set:
             options += add_dtype_fmt_option_single(x, x_n, True)
         else:
@@ -74,27 +83,40 @@ def load_dso(so_path):
     try:
         ctypes.CDLL(so_path)
     except OSError as error :
-        print(error)
+        CommonUtility.print_compile_log("", error, AscendCLogLevel.LOG_ERROR)
         raise RuntimeError("cannot open %s" %(so_path))
     else:
-        print("load so succ ", so_path)
+        msg = "load so succ " + so_path
+        CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO)
 
 def get_shortsoc_compile_option(compile_option_list: list, shortsoc:str):
     compile_options = []
     if shortsoc in compile_option_list:
-        compile_options = compile_option_list[shortsoc]
-    elif '__ALLSOC__' in compile_option_list:
-        compile_options = compile_option_list['__ALLSOC__']
+        compile_options.extend(compile_option_list[shortsoc])
+    if '__ALLSOC__' in compile_option_list:
+        compile_options.extend(compile_option_list['__ALLSOC__'])
     return compile_options
 
 def get_kernel_source(src_file, dir_snake, dir_ex):
     src_ex = os.path.join(PYF_PATH, "..", "ascendc", dir_ex, src_file)
     if os.path.exists(src_ex):
         return src_ex
+    src = os.environ.get('BUILD_KERNEL_SRC')
+    if src and os.path.exists(src):
+        return src
     src = os.path.join(PYF_PATH, "..", "ascendc", dir_snake, src_file)
     if os.path.exists(src):
         return src
     src = os.path.join(PYF_PATH, src_file)
+    if os.path.exists(src):
+        return src
+    src = os.path.join(PYF_PATH, "..", "ascendc", dir_snake, dir_snake + ".cpp")
+    if os.path.exists(src):
+        return src
+    src = os.path.join(PYF_PATH, "..", "ascendc", dir_ex, dir_ex + ".cpp")
+    if os.path.exists(src):
+        return src
+    src = os.path.join(PYF_PATH, "..", "ascendc", os.path.splitext(src_file)[0], src_file)
     if os.path.exists(src):
         return src
     return src_ex
@@ -109,6 +131,8 @@ def _build_args(gradient_in__, is_finite_out_):
                 __inputs__.append(arg[0])
             else:
                 __inputs__.append(arg)
+        else:
+            __inputs__.append(arg)
     __outputs__ = []
     for arg in [is_finite_out_]:
         if arg != None:
@@ -118,24 +142,30 @@ def _build_args(gradient_in__, is_finite_out_):
                 __outputs__.append(arg[0])
             else:
                 __outputs__.append(arg)
+        else:
+            __outputs__.append(arg)
     __attrs__ = []
     return __inputs__, __outputs__, __attrs__
 
 @tbe_register.register_operator("AllFinite", trans_bool_to_s8=False)
 @para_check.check_op_params(para_check.REQUIRED_INPUT, para_check.REQUIRED_OUTPUT, para_check.KERNEL_NAME)
 def all_finite(gradient_in__, is_finite_out_, kernel_name="all_finite", impl_mode=""):
+    # do ascendc build step
     if get_current_build_config("enable_op_prebuild"):
         return
     __inputs__, __outputs__, __attrs__ = _build_args(gradient_in__, is_finite_out_)
     options = get_dtype_fmt_options(__inputs__, __outputs__)
     options += ["-x", "cce"]
-    bisheng = shutil.which("bisheng")
+    bisheng = os.environ.get('BISHENG_REAL_PATH')
+    if bisheng is None:
+        bisheng = shutil.which("bisheng")
     if bisheng != None:
         bisheng_path = os.path.dirname(bisheng)
         tikcpp_path = os.path.realpath(os.path.join(bisheng_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, "..", "..", "include"))
     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"))
@@ -144,8 +174,10 @@ def all_finite(gradient_in__, is_finite_out_, kernel_name="all_finite", impl_mod
         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")
+    if get_current_build_config("enable_deterministic_mode") == 1:
+        options.append("-DDETERMINISTIC_MODE=1")
+    else:
+        options.append("-DDETERMINISTIC_MODE=0")
 
     custom_compile_options = {},
     custom_all_compile_options = {},
@@ -162,11 +194,14 @@ def all_finite(gradient_in__, is_finite_out_, kernel_name="all_finite", impl_mod
     ascendc_src_file = "all_finite.cpp"
     src = get_kernel_source(ascendc_src_file, ascendc_src_dir, ascendc_src_dir_ex)
 
-    print("start compile Ascend C operator AllFinite. kernel name is " + kernel_name)
+    msg = "start compile Acend C Operator AllFinite, kernel name is " + kernel_name
+    CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO)
     op_type = "AllFinite"
     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=[gradient_in__], origin_outputs = [is_finite_out_])
+        attrs = __attrs__, impl_mode = impl_mode, origin_inputs=[gradient_in__], origin_outputs = [is_finite_out_],\
+                param_type_dynamic = False, mc2_ctx = [], param_type_list = ['required', 'required'], init_value_list = [None],\
+                output_shape_depend_on_compute = [])
     compile_op(src, origin_func_name, op_info, options, code_channel, '{}')
 
 def op_select_format(gradient_in__, is_finite_out_, impl_mode=""):

mindspore/lib/plugin/ascend/custom_ascendc_910b/op_impl/ai_core/tbe/custom_ascendc_910b_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);
+  }
+}