mindspore 2.4.1__cp310-none-any.whl → 2.4.10__cp310-none-any.whl

mindspore/experimental/llm_boost/atb/llama_boost.py

@@ -15,10 +15,16 @@
 """llm boost"""
 import json
 import mindspore.common.dtype as mstype
-from mindspore.experimental.llm_boost.atb.boost_base import AtbBoostBase
+from mindspore.experimental.llm_boost.atb.boost_base import (
+    AtbBoostBase,
+    PositionEmbeddingType,
+    NormType,
+)
 from mindspore._c_expression import LlmBoostBinder
 from mindspore.experimental.llm_boost.register import LlmBoostRegister, LlmBoostType
 
+CPP_LLAMA_MODEL_CLASS_NAME = "llama_LlamaDecoderModel"
+
 
 @LlmBoostRegister.register(LlmBoostType.BUILDIN, "Llama")
 class LlamaBoost(AtbBoostBase):
@@ -30,14 +36,17 @@ class LlamaBoost(AtbBoostBase):
         self.acl_encoder_operation_inputs = [None] * self.in_tensor_length
         self.acl_decoder_operation_inputs = [None] * self.in_tensor_length
         self.atb_encoder_operation = LlmBoostBinder(
-            "ATB", "llama_parallel_DecoderModel")
+            self.backend_name, CPP_LLAMA_MODEL_CLASS_NAME
+        )
         self.atb_decoder_operation = LlmBoostBinder(
-            "ATB", "llama_parallel_DecoderModel")
+            self.backend_name, CPP_LLAMA_MODEL_CLASS_NAME
+        )
 
     def init(self):
         """set param"""
         coder_param = {
-            "rmsNormEps": self.config.rms_norm_eps,
+            "normEps": self.config.rms_norm_eps,
+            "normType": NormType.RMS_NORM,
             "numAttentionHeadsPerRank": self.config.num_heads // self.device_num,
             "hiddenSizePerAttentionHead": self.head_dim,
             "numHiddenLayers": self.num_layers,
@@ -46,32 +55,41 @@ class LlamaBoost(AtbBoostBase):
             "isFA": False,
             "isBF16": self.dtype == mstype.bfloat16,
             "packQuantType": [[1, 1] for _ in range(self.num_layers)],
-            "linearQuantType": [[0, -1, -1, 0, 0, -1, 0] for _ in range(self.num_layers)],
-            "linearTransposeType": [[1, -1, -1, 1, 1, -1, 1] for i in range(self.num_layers)],
+            "linearQuantType": [
+                [0, -1, -1, 0, 0, -1, 0] for _ in range(self.num_layers)
+            ],
+            "linearTransposeType": [
+                [1, -1, -1, 1, 1, -1, 1] for i in range(self.num_layers)
+            ],
             "isEmbeddingParallel": False,
             "isLmHeadParallel": not self.config.parallel_config.vocab_emb_dp,
             "lmHeadTransposeType": 1,
-            "supportSwiGLU": True,
-            "kvQuant": self.kv_quant is not None,
+            "enableSwiGLU": True,
+            "enablekvQuant": self.kv_quant is not None,
             "rank": self.rank_id,
             "worldSize": self.device_num,
-            "backend": "lccl",
+            "backend": self.config.communication_backend,
             "rankTableFile": "",
-            "positionEmbeddingType": self.position_embedding_type,
+            "positionEmbeddingType": PositionEmbeddingType.ROPE,
             "hiddenSize": self.config.hidden_size,
             "gemma": False,
-            "enableAddNorm": True,
-            "supportCompressHead": False,
+            "enableAddNorm": False,
+            "enableCompressHead": False,
+            "isUnpadInputs": True,
         }
         encoder_param = {
-            **coder_param, "isPrefill": True,
-            "supportLcoc": True,
-            "supportSpeculate": False,
-            "skipWordEmbedding": False
+            **coder_param,
+            "isPrefill": True,
+            "enableLcoc": True,
+            "enableSpeculate": False,
+            "skipWordEmbedding": False,
+            "enableSplitFuse": False,
         }
         decoder_param = {
-            **coder_param, "isPrefill": False, "supportLcoc": False,
-            "supportSpeculate": False
+            **coder_param,
+            "isPrefill": False,
+            "enableLcoc": False,
+            "enableSpeculate": False,
         }
         self.atb_encoder_operation.init(json.dumps({**encoder_param}))
         self.atb_decoder_operation.init(json.dumps({**decoder_param}))
@@ -92,14 +110,15 @@ class LlamaBoost(AtbBoostBase):
             **kwargs
     ):
         """prepare inputs"""
-        self.acl_param = json.dumps({
-            "seqLen": seqLen,
-        })
-        self.acl_decoder_operation_inputs[0] = self.cast(
-            input_ids, mstype.int64)
+        self.acl_param = json.dumps(
+            {
+                "seqLen": seqLen,
+            }
+        )
+
+        self.acl_decoder_operation_inputs[0] = input_ids
         self.acl_decoder_operation_inputs[1] = self.placeholder
-        self.acl_decoder_operation_inputs[2] = self.cast(
-            position_ids, mstype.int32)
+        self.acl_decoder_operation_inputs[2] = position_ids
         self.acl_decoder_operation_inputs[3] = cos_embed
         self.acl_decoder_operation_inputs[4] = sin_embed
         self.acl_decoder_operation_inputs[5] = attention_mask
@@ -108,8 +127,6 @@ class LlamaBoost(AtbBoostBase):
         self.acl_decoder_operation_inputs[8] = self.placeholder
         self.acl_decoder_operation_inputs[9] = self.placeholder
         self.acl_decoder_operation_inputs[10] = self.placeholder
-        self.acl_decoder_operation_inputs[11] = self.cast(
-            input_lengths, mstype.int32)
-        self.acl_decoder_operation_inputs[12] = self.cast(
-            lm_head_indices, mstype.int64)
+        self.acl_decoder_operation_inputs[11] = input_lengths
+        self.acl_decoder_operation_inputs[12] = lm_head_indices
         return self.acl_decoder_operation_inputs, self.acl_param

mindspore/experimental/llm_boost/atb/qwen_boost.py

@@ -15,11 +15,14 @@
 """llm boost"""
 import json
 import mindspore.common.dtype as mstype
-from mindspore.experimental.llm_boost.atb.boost_base import AtbBoostBase
+from mindspore.experimental.llm_boost.atb.boost_base import AtbBoostBase, NormType
 from mindspore._c_expression import LlmBoostBinder
 from mindspore.experimental.llm_boost.register import LlmBoostRegister, LlmBoostType
 
 
+CPP_QWEN_MODEL_CLASS_NAME = "qwen_QwenDecoderModel"
+
+
 @LlmBoostRegister.register(LlmBoostType.BUILDIN, "Qwen")
 class QwenBoost(AtbBoostBase):
     """QwenBoost class"""
@@ -30,9 +33,11 @@ class QwenBoost(AtbBoostBase):
         self.acl_encoder_operation_inputs = [None] * self.in_tensor_length
         self.acl_decoder_operation_inputs = [None] * self.in_tensor_length
         self.atb_encoder_operation = LlmBoostBinder(
-            "ATB", "qwen_DecoderModel")
+            self.backend_name, CPP_QWEN_MODEL_CLASS_NAME
+        )
         self.atb_decoder_operation = LlmBoostBinder(
-            "ATB", "qwen_DecoderModel")
+            self.backend_name, CPP_QWEN_MODEL_CLASS_NAME
+        )
 
     def init(self):
         """set param"""
@@ -42,24 +47,43 @@ class QwenBoost(AtbBoostBase):
             "withEmbedding": True,
             "isEmbeddingParallel": True,
             "isLmHeadParallel": True,
-            "linearTransposeType": [[1, -1, -1, 1, 1, -1, 1] for i in range(self.num_layers)],
+            "linearTransposeType": [
+                [1, -1, -1, 1, 1, -1, 1] for i in range(self.num_layers)
+            ],
             "lmHeadTransposeType": 1,
-            "supportSwiGLU": not self.need_nz,
-            "rmsNormEps": self.config.rms_norm_eps,
+            "enableSwiGLU": not self.need_nz,
+            "normEps": self.config.rms_norm_eps,
+            "normType": NormType.RMS_NORM,
             "numAttentionHeadsPerRank": self.config.num_heads // self.device_num,
             "hiddenSizePerAttentionHead": self.head_dim,
             "numHiddenLayers": self.num_layers,
             "numKeyValueHeadsPerRank": self.n_kv_heads // self.device_num,
             "rank": self.rank_id,
             "worldSize": self.device_num,
-            "backend": "lccl",
+            "backend": self.config.communication_backend,
             "packQuantType": [[1, 1] for _ in range(self.num_layers)],
-            "linearQuantType": [[0, -1, -1, 0, 0, -1, 0] for _ in range(self.num_layers)],
-            "kvQuant": self.kv_quant is not None,
+            "linearQuantType": [
+                [0, -1, -1, 0, 0, -1, 0] for _ in range(self.num_layers)
+            ],
+            "linearHasBias": [[True, False, False, False]] * self.num_layers,
+            "enableKvQuant": self.kv_quant is not None,
+            "enableLora": False,
+            "isUnpadInputs": True,
+            "enableAddNorm": False,
+        }
+        encoder_param = {
+            **param_dict,
+            "isPrefill": True,
+            "enableLcoc": False,
+            "enableSplitFuse": False,
+        }
+        decoder_param = {
+            **param_dict,
+            "isPrefill": False,
+            "enableLcoc": False,
+            "enableSpeculate": False,
+            "enablePrefixCache": False,
         }
-        encoder_param = {**param_dict, "isPrefill": True, "supportLcoc": False}
-        decoder_param = {**param_dict, "isPrefill": False,
-                         "supportLcoc": False, "supportSpeculate": False}
         self.atb_encoder_operation.init(json.dumps({**encoder_param}))
         self.atb_decoder_operation.init(json.dumps({**decoder_param}))
 
@@ -79,13 +103,14 @@ class QwenBoost(AtbBoostBase):
             **kwargs
     ):
         """prepare inputs"""
-        self.acl_param = json.dumps({
-            "seqLen": seqLen,
-        })
-        self.acl_decoder_operation_inputs[0] = self.cast(
-            input_ids, mstype.int64)
-        self.acl_decoder_operation_inputs[1] = self.cast(
-            position_ids, mstype.int32)
+        self.acl_param = json.dumps(
+            {
+                "seqLen": seqLen,
+            }
+        )
+
+        self.acl_decoder_operation_inputs[0] = input_ids
+        self.acl_decoder_operation_inputs[1] = position_ids
         self.acl_decoder_operation_inputs[2] = cos_embed
         self.acl_decoder_operation_inputs[3] = sin_embed
         self.acl_decoder_operation_inputs[4] = attention_mask
@@ -93,9 +118,7 @@ class QwenBoost(AtbBoostBase):
         self.acl_decoder_operation_inputs[6] = slots
         self.acl_decoder_operation_inputs[7] = self.placeholder
         self.acl_decoder_operation_inputs[8] = self.placeholder
-        self.acl_decoder_operation_inputs[9] = self.cast(
-            input_lengths, mstype.int32)
-        self.acl_decoder_operation_inputs[10] = self.cast(
-            lm_head_indices, mstype.int64)
-        self.acl_decoder_operation_inputs[11] = self.placeholder
+        self.acl_decoder_operation_inputs[9] = self.placeholder
+        self.acl_decoder_operation_inputs[10] = input_lengths
+        self.acl_decoder_operation_inputs[11] = lm_head_indices
         return self.acl_decoder_operation_inputs, self.acl_param

mindspore/include/api/context.h

@@ -236,7 +236,7 @@ std::string DeviceInfoContext::GetProviderDevice() const { return CharToString(G
 void DeviceInfoContext::SetProviderDevice(const std::string &device) { SetProviderDevice(StringToChar(device)); }
 
 /// \brief Derived from DeviceInfoContext, The configuration of the model running auto on the Host Devices, include
-/// CPU/GPU/NPU/Ascend310/Ascend910. This option is only valid for MindSpore Lite.
+/// CPU/GPU/NPU/Ascend. This option is only valid for MindSpore Lite.
 class MS_API AutoDeviceInfo : public DeviceInfoContext {
  public:
   /// \brief Get the type of this DeviceInfoContext.

mindspore/include/dataset/constants.h

@@ -108,8 +108,8 @@ enum class DATASET_API ManualOffloadMode {
 enum class DATASET_API MapTargetDevice {
   kCpu = 0,     ///< CPU Device.
   kGpu,         ///< Gpu Device.
-  kAscend310,   ///< Ascend310 Device.
-  kAscend910B,  ///< Ascend910B Device.
+  kAscend310,   ///<
+  kAscend910B,  ///<
   kInvalid = 100
 };
 

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

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

mindspore/lib/plugin/ascend/{custom_ascendc_ops → custom_ascendc_910}/op_impl/ai_core/tbe/config/ascend910/aic-ascend910-ops-info.json

@@ -1,46 +1,4 @@
 {
-    "AllFinite":{
-        "dynamicCompileStatic":{
-            "flag":"true"
-        },
-        "dynamicFormat":{
-            "flag":"true"
-        },
-        "dynamicRankSupport":{
-            "flag":"true"
-        },
-        "dynamicShapeSupport":{
-            "flag":"true"
-        },
-        "input0":{
-            "dtype":"float16,float32,bfloat16",
-            "format":"ND,ND,ND",
-            "name":"gradient",
-            "paramType":"required",
-            "shape":"all",
-            "unknownshape_format":"ND,ND,ND"
-        },
-        "needCheckSupport":{
-            "flag":"false"
-        },
-        "opFile":{
-            "value":"all_finite"
-        },
-        "opInterface":{
-            "value":"all_finite"
-        },
-        "output0":{
-            "dtype":"bool,bool,bool",
-            "format":"ND,ND,ND",
-            "name":"is_finite",
-            "paramType":"required",
-            "shape":"all",
-            "unknownshape_format":"ND,ND,ND"
-        },
-        "precision_reduce":{
-            "flag":"true"
-        }
-    },
     "DecoderKvCache":{
         "dynamicCompileStatic":{
             "flag":"true"

mindspore/lib/plugin/ascend/{custom_ascendc_ops/op_impl/ai_core/tbe/custom_ascendc_ops_impl → custom_ascendc_910/op_impl/ai_core/tbe/custom_ascendc_910_impl}/dynamic/decoder_kv_cache.py

@@ -1,3 +1,8 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+"""
+Copyright (c) Huawei Technologies Co., Ltd. 2023-2024. 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 = ['cache', 'update', 'valid_seq_len', 'batch_index', 'seq_len_axis', 'new_max_seq_len', 'cur_max_seq_len']
+    output_names = ['out']
     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(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, s
                 __inputs__.append(arg[0])
             else:
                 __inputs__.append(arg)
+        else:
+            __inputs__.append(arg)
     __outputs__ = []
     for arg in [out_out_]:
         if arg != None:
@@ -118,24 +142,30 @@ def _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, s
                 __outputs__.append(arg[0])
             else:
                 __outputs__.append(arg)
+        else:
+            __outputs__.append(arg)
     __attrs__ = []
     return __inputs__, __outputs__, __attrs__
 
 @tbe_register.register_operator("DecoderKvCache", trans_bool_to_s8=False)
 @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_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_, kernel_name="decoder_kv_cache", impl_mode=""):
+    # do ascendc build step
     if get_current_build_config("enable_op_prebuild"):
         return
     __inputs__, __outputs__, __attrs__ = _build_args(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_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 decoder_kv_cache(cache_in__, update_in__, valid_seq_len_in__, batch_index_in
         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 decoder_kv_cache(cache_in__, update_in__, valid_seq_len_in__, batch_index_in
     ascendc_src_file = "decoder_kv_cache.cpp"
     src = get_kernel_source(ascendc_src_file, ascendc_src_dir, ascendc_src_dir_ex)
 
-    print("start compile Ascend C operator DecoderKvCache. kernel name is " + kernel_name)
+    msg = "start compile Acend C Operator DecoderKvCache, kernel name is " + kernel_name
+    CommonUtility.print_compile_log("", msg, AscendCLogLevel.LOG_INFO)
     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_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__], origin_outputs = [out_out_])
+        attrs = __attrs__, impl_mode = impl_mode, origin_inputs=[cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__], origin_outputs = [out_out_],\
+                param_type_dynamic = False, mc2_ctx = [], param_type_list = ['required', 'required', 'required', 'required', 'required', 'required', '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(cache_in__, update_in__, valid_seq_len_in__, batch_index_in__, seq_len_axis_in__, new_max_seq_len_in__, cur_max_seq_len_in__, out_out_, impl_mode=""):