mindspore 2.3.0__cp310-cp310-win_amd64.whl → 2.4.1__cp310-cp310-win_amd64.whl

mindspore/{nn/layer → experimental/es}/embedding_service_layer.py

@@ -19,8 +19,10 @@ import mindspore as ms
 from mindspore import nn, ops, Tensor, Parameter
 from mindspore.ops.auto_generate import init_partition_map, init_embedding_hashmap, embedding_table_find_and_init,\
     embedding_table_find, fake_remote_lookup_uniqued
-from mindspore.ops.operations.manually_defined import EmbeddingTableImport, EmbeddingTableExport, \
-    EmbeddingComputeVarImport, EmbeddingComputeVarExport
+from mindspore.ops.auto_generate import EmbeddingTableImport, EmbeddingTableExport, \
+    EmbeddingComputeVarImport, EmbeddingComputeVarExport, EmbeddingTableEvict, EmbeddingFeatureMappingV2, \
+        EmbeddingFeatureMappingTableSize, EmbeddingFeatureMappingFind, EmbeddingFeatureMappingExport, \
+            EmbeddingFeatureMappingFileSize, EmbeddingFeatureMappingImport, EmbeddingFeatureMappingInsert
 
 
 class CounterFilter:
@@ -55,12 +57,14 @@ def _get_backward_float_params(optimizer_mode):
           [beta1_power, beta2_power, lr, weight_decay, beta1, beta2, epsilon]
         - when the backward_mode is 'adagrad', it means [lr,]
     """
-    if optimizer_mode == "adagrad":
+    if optimizer_mode == "adagrad" or optimizer_mode == "sgd":
         return [0.001]
     if optimizer_mode == "adam":
         return [0.9, 0.99, 0.001, 0.9, 0.999, 1e-08]
     if optimizer_mode == "ftrl":
         return [0.001, -0.5, 0.0, 0.0]
+    if optimizer_mode == "rmsprop":
+        return [0.001, 0.9, 0.1, 1e-08]
     # adamw
     return [0.9, 0.99, 0.001, 0.01, 0.9, 0.999, 1e-08]
 
@@ -99,6 +103,9 @@ class ESInitLayer(nn.Cell):
         self.default_value = None
 
     def construct(self):
+        """
+        ESInitLayer construct: init embedding hashmap
+        """
         init_partition = init_partition_map(self.ps_num_tensor,
                                             self.ps_ids_tensor,
                                             _embedding_dim=self.embedding_dim,
@@ -145,9 +152,36 @@ class ESInitLayer(nn.Cell):
 
 
 class EsEmbeddingLookup(nn.Cell):
+    r"""
+    Look up a PS embedding.
+
+    .. warning::
+        This is an experimental EmbeddingService API that is subject to change.
+
+    Args:
+        table_id (int): The table id.
+        es_initializer (EsInitializer): The EsInitialize object for PS embedding with table_id,
+            which can be None when the inference is performed.
+        embedding_dim (int): The embedding dim of keys for PS embedding with table_id.
+        max_key_num (int): The num of keys when lookup.
+        optimizer_mode (str): The type of optimizer. Default is ``None``.
+        optimizer_params (tuple[float]): The parameters of optimizer. Default is ``None``.
+        es_filter (CounterFilter): The option of counter filter for PS embedding with table_id. Default is ``None``.
+        es_padding_key (PaddingParamsOption): The option of padding key for PS embedding with table_id.
+            Default is ``None``.
+        es_completion_key (CompletionKeyOption): The option of completion key for PS embedding with table_id.
+            Default is ``None``.
+
+    Inputs:
+            - **keys** (Tensor): The keys of each feature in PS embedding.
+            - **actual_keys_input** (Tensor): Tensor composed of all unique elements of keys.
+            - **unique_indices** (Tensor): The index value of each element in keys to actual_keys_input .
+            - **key_count** (Tensor): The count of each element in the actual_keys_input to keys.
+
+    Supported Platforms:
+        ``Atlas A2 training series products``
     """
-    EsEmbeddingLookup.
-    """
+
     def __init__(self, table_id, es_initializer, embedding_dim, max_key_num, optimizer_mode=None,
                  optimizer_params=None, es_filter=None, es_padding_key=None, es_completion_key=None):
         super(EsEmbeddingLookup, self).__init__()
@@ -182,7 +216,7 @@ class EsEmbeddingLookup(nn.Cell):
             self.filter_freq = 1
             self.default_key_or_value = 1
             self.default_key = 0
-            self.default_value = -1.0
+            self.default_value = 1.0
 
         self.global_step = 1
         if es_padding_key is not None:
@@ -193,7 +227,7 @@ class EsEmbeddingLookup(nn.Cell):
             self.mask_zero = 0
             self.padding_key = 0
             self.padding_key_mask = 1
-        if self.optimizer_mode in ["adam", "ftrl", "adagrad"]:
+        if self.optimizer_mode in ["adam", "ftrl", "adagrad", "sgd", "rmsprop"]:
             self.backward_int_params = ([self.global_step], [self.mask_zero],
                                         [self.padding_key], [self.padding_key_mask])
         else:
@@ -211,6 +245,9 @@ class EsEmbeddingLookup(nn.Cell):
         self.max_grad_norm = Tensor([1.0], ms.float32)
 
     def construct(self, keys, actual_keys_input=None, unique_indices=None, key_count=None):
+        """
+        Using the corresponding query method to calculate the PS embedding for each key.
+        """
         origin_shape = None
         if len(keys.shape) != 1:
             origin_shape = keys.shape
@@ -227,11 +264,9 @@ class EsEmbeddingLookup(nn.Cell):
                 key_count = keys
         if self.training:
             if use_host_unique:
-                output = fake_remote_lookup_uniqued(table_id=self.table_id,
-                                                    keys=keys,
+                output = fake_remote_lookup_uniqued(table_id=self.table_id, keys=keys,
                                                     actual_keys_num=actual_keys_input,
-                                                    unique_indices=unique_indices,
-                                                    key_count=key_count,
+                                                    unique_indices=unique_indices, key_count=key_count,
                                                     max_grad_norm=self.max_grad_norm,
                                                     embedding_dim=self.embedding_dim,
                                                     initializer_mode=self.es_initializer.initializer_mode,
@@ -250,8 +285,7 @@ class EsEmbeddingLookup(nn.Cell):
                                                     default_value=self.default_value,
                                                     optimizer_mode=self.optimizer_mode,
                                                     optimizer_params=self.optimizer_params,
-                                                    _max_key_num=self.max_key_num,
-                                                    _table_id=self._table_id,
+                                                    _max_key_num=self.max_key_num, _table_id=self._table_id,
                                                     _use_counter_filter=use_counter_filter,
                                                     backward_mode=self.optimizer_mode,
                                                     backward_int_params=self.backward_int_params,
@@ -280,8 +314,7 @@ class EsEmbeddingLookup(nn.Cell):
                                                        default_value=self.default_value,
                                                        optimizer_mode=self.optimizer_mode,
                                                        optimizer_params=self.optimizer_params,
-                                                       _max_key_num=self.max_key_num,
-                                                       _table_id=self._table_id,
+                                                       _max_key_num=self.max_key_num, _table_id=self._table_id,
                                                        _use_counter_filter=use_counter_filter,
                                                        backward_mode=self.optimizer_mode,
                                                        backward_int_params=self.backward_int_params,
@@ -290,14 +323,10 @@ class EsEmbeddingLookup(nn.Cell):
                                                        completion_key_mask=self.completion_key_mask,
                                                        parameter=self.b)
         else:
-            output = embedding_table_find(self.table_id, keys,
-                                          embedding_dim=self.embedding_dim,
+            output = embedding_table_find(self.table_id, keys, embedding_dim=self.embedding_dim,
                                           default_value=self.default_value,
-                                          _max_key_num=self.max_key_num,
-                                          _table_id=self._table_id,
+                                          _max_key_num=self.max_key_num, _table_id=self._table_id,
                                           _use_counter_filter=use_counter_filter)
-        # input 20480 2 ->41960
-        # output 41960 embedding_dim -> 20480 2 embedding_dim
         if origin_shape is not None:
             output = self.reshape(output, origin_shape + (-1,))
         return output
@@ -321,10 +350,10 @@ class ESEmbeddingCKPTExport(nn.Cell):
         self.table_id_tensor = Tensor(table_id_list, ms.int32)
         self.depend = ops.Depend()
 
-    def construct(self):
-        export_op1 = self.embedding_table_export(self.file_path, self.ps_id_tensor, self.table_id_tensor)
+    def construct(self, global_step):
+        export_op1 = self.embedding_table_export(self.file_path, self.ps_id_tensor, self.table_id_tensor, global_step)
         z = self.depend(self.file_path, export_op1)
-        export_op2 = self.embedding_compute_var_export(z, self.ps_id_tensor, self.table_id_tensor)
+        export_op2 = self.embedding_compute_var_export(z, self.ps_id_tensor, self.table_id_tensor, global_step)
         return export_op2
 
 
@@ -345,8 +374,31 @@ class ESEmbeddingTableExport(nn.Cell):
         self.ps_id_tensor = Tensor(0, ms.int32)
         self.table_id_tensor = Tensor(table_id_list, ms.int32)
 
-    def construct(self):
-        y = self.op(self.file_path, self.ps_id_tensor, self.table_id_tensor)
+    def construct(self, global_step):
+        y = self.op(self.file_path, self.ps_id_tensor, self.table_id_tensor, global_step)
+        return y
+
+
+class ESIncrementalEmbeddingTableExport(nn.Cell):
+    """
+    ESIncrementalEmbeddingTableExport.
+    """
+    def __init__(self, embedding_dim_list, value_total_len_list, table_name_list, table_id_list,
+                 file_path, steps_to_live_list):
+        super(ESIncrementalEmbeddingTableExport, self).__init__()
+        self.op = EmbeddingTableExport(
+            embedding_dim_list,
+            value_total_len_list,
+            table_name=table_name_list,
+            steps_to_live_list=steps_to_live_list,
+            export_mode="new",
+            only_var_flag=True)
+        self.file_path = Tensor(np.array(file_path))
+        self.ps_id_tensor = Tensor(0, ms.int32)
+        self.table_id_tensor = Tensor(table_id_list, ms.int32)
+
+    def construct(self, global_step):
+        y = self.op(self.file_path, self.ps_id_tensor, self.table_id_tensor, global_step)
         return y
 
 
@@ -366,10 +418,10 @@ class ESEmbeddingCKPTImport(nn.Cell):
         self.table_id_tensor = Tensor(table_id_list, ms.int32)
         self.depend = ops.Depend()
 
-    def construct(self):
-        export_op1 = self.embedding_table_import(self.file_path, self.ps_id_tensor, self.table_id_tensor)
+    def construct(self, global_step):
+        export_op1 = self.embedding_table_import(self.file_path, self.ps_id_tensor, self.table_id_tensor, global_step)
         z = self.depend(self.file_path, export_op1)
-        export_op2 = self.embedding_compute_var_import(z, self.ps_id_tensor, self.table_id_tensor)
+        export_op2 = self.embedding_compute_var_import(z, self.ps_id_tensor, self.table_id_tensor, global_step)
         return export_op2
 
 
@@ -388,6 +440,142 @@ class ESEmbeddingTableImport(nn.Cell):
         self.ps_id_tensor = Tensor(0, ms.int32)
         self.table_id_tensor = Tensor(table_id_list, ms.int32)
 
+    def construct(self, global_step):
+        y = self.op(self.file_path, self.ps_id_tensor, self.table_id_tensor, global_step)
+        return y
+
+
+class ESEmbeddingTableEvict(nn.Cell):
+    """
+    ESEmbeddingTableEvict.
+    """
+    def __init__(self, var_handle, global_step, steps_to_live):
+        super(ESEmbeddingTableEvict, self).__init__()
+        self.op = EmbeddingTableEvict()
+        self.var_handle = Tensor(var_handle, ms.int32)
+        self.global_step = global_step
+        self.steps_to_live = steps_to_live
+
     def construct(self):
-        y = self.op(self.file_path, self.ps_id_tensor, self.table_id_tensor)
+        y = self.op(self.var_handle, self.global_step, self.steps_to_live)
         return y
+
+
+class ESEmbeddingFeatureMappingExport(nn.Cell):
+    """
+    ESEmbeddingFeatureMappingExport.
+    """
+    def __init__(self, file_path, export_value, var, var_name, small_table_embedding_dim):
+        super(ESEmbeddingFeatureMappingExport, self).__init__()
+        self.embedding_feature_mapping_table_size = EmbeddingFeatureMappingTableSize()
+        self.embedding_feature_mapping_find = EmbeddingFeatureMappingFind()
+        self.embedding_feature_mapping_export = EmbeddingFeatureMappingExport()
+        self.file_path = file_path
+        self.export_value = export_value
+        self.gather = ops.Gather()
+        self.var = Tensor(var, ms.float32)
+        self.var_name = Tensor(np.array([var_name]))
+        self.small_table_embedding_dim = [small_table_embedding_dim]
+        self.global_step = Tensor([-1], ms.int64)
+
+    def construct(self):
+        """
+        ESEmbeddingFeatureMappingExport construct: export feature mapping for data_parallel embedding.
+        """
+        feature_size = self.embedding_feature_mapping_table_size(self.var_name)
+        feature_id, offset_id = self.embedding_feature_mapping_find(self.var_name, feature_size, 1)
+        values = self.gather(self.var, offset_id, 0)
+        if self.export_value:
+            embed_values = values
+        else:
+            embed_values = Tensor([0], ms.float32)
+        feature_mapping_export = self.embedding_feature_mapping_export(self.file_path, self.var_name, self.global_step,
+                                                                       embed_values, self.small_table_embedding_dim,
+                                                                       [feature_id], [offset_id])
+        return feature_mapping_export
+
+
+class ESEmbeddingFeatureMappingImport(nn.Cell):
+    """
+    ESEmbeddingFeatureMappingImport.
+    """
+    def __init__(self, file_path, small_table_name, small_table_embedding_dim, only_offset_flag):
+        super(ESEmbeddingFeatureMappingImport, self).__init__()
+        self.embedding_feature_mapping_file_size = EmbeddingFeatureMappingFileSize()
+        self.embedding_feature_mapping_import = EmbeddingFeatureMappingImport()
+        self.embedding_feature_mapping_insert = EmbeddingFeatureMappingInsert()
+        self.file_path = file_path
+        self.small_table_name = Tensor(np.array([small_table_name]))
+        self.small_table_embedding_dim = [small_table_embedding_dim]
+        self.only_offset_flag = only_offset_flag
+        self.global_step = Tensor([-1], ms.int64)
+
+    def construct(self):
+        """
+        ESEmbeddingFeatureMappingImport construct: import feature mapping for data_parallel embedding.
+        """
+        feature_size = self.embedding_feature_mapping_file_size(self.file_path,
+                                                                self.small_table_name,
+                                                                self.global_step,
+                                                                self.small_table_embedding_dim,
+                                                                self.only_offset_flag)
+        feature_id, offset_id = self.embedding_feature_mapping_import(self.file_path,
+                                                                      self.small_table_name,
+                                                                      feature_size, self.global_step,
+                                                                      self.small_table_embedding_dim,
+                                                                      self.only_offset_flag, 1)
+        feature_mapping_insert = self.embedding_feature_mapping_insert(self.small_table_name, 1,
+                                                                       [feature_id], [offset_id])
+        return feature_mapping_insert
+
+
+class ESEmbeddingSmallTableLookup(nn.Cell):
+    r"""
+    Look up a data_parallel embedding.
+
+    .. warning::
+        This is an experimental EmbeddingService API that is subject to change.
+
+    Args:
+        name (str): The data_parallel embedding name.
+        rank_id (int): The rank id when look up data_parallel embedding key.
+        rank_size (int): The rank size when look up data_parallel embedding key.
+        small_table_to_variable (dict[str, parameter]): The dict to restore data_parallel embedding information:
+            key is table name, value is parameter.
+
+    Inputs:
+        - **ids_list** (Tensor) - The keys of each feature in data_parallel embedding.
+
+    Supported Platforms:
+        ``Atlas A2 training series products``
+    """
+
+    def __init__(self, name, rank_id, rank_size, small_table_to_variable):
+        super(ESEmbeddingSmallTableLookup, self).__init__()
+        self.small_table_to_variable = small_table_to_variable[name]
+        self.small_table_to_variable.feature_name = name
+        self.allgather = ops.AllGather()
+        self.gather = ops.Gather()
+        self.embedding_feature_mapping_v2 = EmbeddingFeatureMappingV2()
+        self.name = name
+        self.rank_id = rank_id
+        self.rank_size = rank_size
+
+    def construct(self, ids_list):
+        """
+        Using the EmbeddingFeatureMappingV2 method to mapping hash key to non hash key, and then get embedding value.
+        """
+        hash_key_shape = ids_list.shape
+        if self.rank_size > 1 and (hash_key_shape[0] is not None):
+            hash_key = ops.stop_gradient(self.allgather(ids_list))
+            non_hash_key = self.embedding_feature_mapping_v2(self.name, hash_key, [1], [1])
+            recovery_matrix = []
+            for i in range(hash_key_shape[0]):
+                recovery_matrix.append(self.rank_id * hash_key_shape[0] + i)
+            local_non_hash_keys = self.gather(non_hash_key, Tensor(recovery_matrix), 0)
+        else:
+            hash_key = ids_list
+            local_non_hash_keys = self.embedding_feature_mapping_v2(self.name, hash_key, [1], [1])
+
+        embedding = self.gather(self.small_table_to_variable, local_non_hash_keys, 0)
+        return embedding

mindspore/experimental/llm_boost/__init__.py

@@ -0,0 +1,21 @@
+# Copyright 2024 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.
+# ============================================================================
+"""LlmBoost Register"""
+from __future__ import absolute_import
+
+from mindspore.experimental.llm_boost.atb import *
+from mindspore.experimental.llm_boost.register import LlmBoostRegister
+
+__all__ = ['LlmBoostRegister']

mindspore/{nn/extend/layer → experimental/llm_boost/atb}/__init__.py

@@ -13,15 +13,11 @@
 # limitations under the License.
 # ============================================================================
 """
-Layer.
-
-The high-level components(Cells) used to construct the neural network.
+Provide llm boost for inference, such as LlamaBoost.
 """
 from __future__ import absolute_import
 
-from mindspore.nn.extend.layer import normalization
-from mindspore.nn.extend.layer.normalization import *
-
-__all__ = []
+from mindspore.experimental.llm_boost.atb.llama_boost import LlamaBoost
+from mindspore.experimental.llm_boost.atb.qwen_boost import QwenBoost
 
-__all__.extend(normalization.__all__)
+__all__ = ['LlamaBoost', 'QwenBoost']

mindspore/experimental/llm_boost/atb/boost_base.py

@@ -0,0 +1,211 @@
+# Copyright 2024 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.
+# ============================================================================
+"""boost base class"""
+import numpy as np
+import mindspore as ms
+from mindspore import ops, Tensor
+from mindspore.ops import operations as P
+import mindspore.common.dtype as mstype
+from mindspore._c_expression import _set_format
+
+from mindspore.common.parameter import Parameter
+from mindspore.experimental.llm_boost.utils import get_real_rank, get_real_group_size
+from mindspore.common.initializer import Zero
+
+
+class AttentionMask:
+    """attention mask"""
+
+    @classmethod
+    def static(cls, max_seq_len, dtype=mstype.float16, need_nz=False):
+        """cache mask"""
+        bias_cache = Tensor(np.tril(np.ones((max_seq_len, max_seq_len), dtype=np.bool_))).reshape(max_seq_len,
+                                                                                                  max_seq_len)
+        bias_cache = ~bias_cache
+        if dtype == mstype.float16:
+            mask_value = Tensor(np.finfo(np.float32).min, mstype.float16)
+        else:
+            mask_value = Tensor(1)
+        attn_mask = ops.masked_fill(Tensor(np.zeros(
+            (max_seq_len, max_seq_len)), dtype=mstype.float16), bias_cache, mask_value)
+        if need_nz:
+            # ND -> NZ
+            attn_mask = ops.reshape(attn_mask, (1, max_seq_len, max_seq_len))
+            attn_mask = ops.reshape(
+                attn_mask, (1, max_seq_len, max_seq_len // 16, 16))
+            attn_mask = ops.transpose(attn_mask, (0, 2, 1, 3)).contiguous()
+            attn_mask = _set_format(attn_mask, "FRACTAL_NZ")
+        return attn_mask
+
+
+class AtbBoostBase():
+    """atb boost base class"""
+
+    def __init__(self, config):
+        super().__init__()
+        self.is_first_iteration = False
+        self.config = config
+        self.dtype = config.compute_dtype
+        self.num_heads = config.num_heads
+        self.num_kv_heads = config.n_kv_heads if config.n_kv_heads else self.num_heads
+        self.num_layers = config.num_layers
+        self.n_kv_heads = config.n_kv_heads if config.n_kv_heads else config.num_heads
+        self.head_dim = config.hidden_size // self.num_heads
+        self.need_nz = False
+        if hasattr(config, "need_nz"):
+            self.need_nz = config.need_nz
+        self.placeholder = Tensor(np.zeros(1), dtype=self.dtype)
+        self.lm_head_indices_fake = Tensor([0], dtype=mstype.int64)
+        self.position_embedding_type = "ROPE"
+        self.add_norm_enable = True
+        self.max_decode_length = self.config.max_decode_length
+        self.max_base_len = 128
+        self.attn_mask = AttentionMask.static(
+            self.max_base_len, dtype=self.dtype, need_nz=self.need_nz)
+
+        self.cast = P.Cast()
+        self.reshape = P.Reshape()
+        self.kv_quant = None
+        self.rank_id = get_real_rank()
+        self.device_num = get_real_group_size()
+
+    def _convert_tensor_format_and_dtype(self, tensor, dtype=mstype.float16):
+        tensor = self.cast(tensor, dtype=dtype)
+        if self.need_nz:
+            tensor = _set_format(tensor, "FRACTAL_NZ")
+        return tensor
+
+    def set_weights(self, parm_dict, dtype=mstype.float16):
+        """set weights for llm boost"""
+        embedding_weight_name = "model.tok_embeddings.embedding_weight"
+        attention_norm_name = "attention_norm"
+        qkv_name = "attention.w_qkv"
+        o_name = "attention.wo"
+        mlp_norm_name = "ffn_norm"
+        mlp_gate_name = "feed_forward.w_gate_hidden"
+        mlp_down_name = "feed_forward.w2"
+        norm_out_name = "model.norm_out"
+        lm_head_name = "lm_head"
+        placeholder = Parameter(Tensor(np.zeros(1), dtype=dtype))
+
+        ascend_weight = []
+        ascend_weight.append(
+            self.cast(parm_dict[embedding_weight_name], dtype))
+        for i in range(self.num_layers):
+            ascend_weight.append(self._convert_tensor_format_and_dtype(
+                parm_dict[f"model.layers.{i}.{attention_norm_name}.weight"], dtype))
+            ascend_weight.extend([placeholder] * 3)
+
+            ascend_weight.append(
+                self._convert_tensor_format_and_dtype(parm_dict[f"model.layers.{i}.{qkv_name}.weight"], dtype))
+            ascend_weight.append(self._convert_tensor_format_and_dtype(parm_dict.get(
+                f"model.layers.{i}.{qkv_name}.bias", placeholder), dtype))
+            ascend_weight.extend([placeholder] * 16)
+
+            ascend_weight.append(
+                self._convert_tensor_format_and_dtype(parm_dict[f"model.layers.{i}.{o_name}.weight"], dtype))
+            ascend_weight.append(self._convert_tensor_format_and_dtype(parm_dict.get(
+                f"model.layers.{i}.{o_name}.bias", placeholder), dtype))
+            ascend_weight.extend([placeholder] * 4)
+
+            ascend_weight.append(
+                self._convert_tensor_format_and_dtype(parm_dict[f"model.layers.{i}.{mlp_norm_name}.weight"], dtype))
+            ascend_weight.extend([placeholder] * 3)
+
+            ascend_weight.append(
+                self._convert_tensor_format_and_dtype(parm_dict[f"model.layers.{i}.{mlp_gate_name}.weight"], dtype))
+            ascend_weight.append(self._convert_tensor_format_and_dtype(parm_dict.get(
+                f"model.layers.{i}.{mlp_gate_name}.bias", placeholder), dtype))
+            ascend_weight.extend([placeholder] * 10)
+
+            ascend_weight.append(
+                self._convert_tensor_format_and_dtype(parm_dict[f"model.layers.{i}.{mlp_down_name}.weight"], dtype))
+            ascend_weight.append(self._convert_tensor_format_and_dtype(parm_dict.get(
+                f"model.layers.{i}.{mlp_down_name}.bias", placeholder), dtype))
+            ascend_weight.extend([placeholder] * 4)
+
+        ascend_weight.append(
+            self._convert_tensor_format_and_dtype(parm_dict[f"{norm_out_name}.weight"], dtype))
+        ascend_weight.append(
+            self._convert_tensor_format_and_dtype(parm_dict[f"{lm_head_name}.weight"], dtype))
+        self.atb_encoder_operation.set_weights(ascend_weight)
+        self.atb_decoder_operation.set_weights(ascend_weight)
+
+    def set_kvcache(self, k_caches=None, v_caches=None):
+        """set kv_cache for llm boost"""
+        if not k_caches or v_caches:
+            if self.need_nz:
+                kv_shape = (self.config.num_blocks, self.num_kv_heads*self.head_dim //
+                            self.device_num // 16, self.config.block_size, 16)
+                k_caches = [_set_format(Parameter(Tensor(
+                    shape=kv_shape, dtype=self.dtype, init=Zero())), "FRACTAL_NZ") for _ in range(self.num_layers)]
+                v_caches = [_set_format(Parameter(Tensor(
+                    shape=kv_shape, dtype=self.dtype, init=Zero())), "FRACTAL_NZ") for _ in range(self.num_layers)]
+            else:
+                kv_shape = (self.config.num_blocks, self.config.block_size,
+                            self.num_kv_heads // self.device_num, self.head_dim)
+                k_caches = [Parameter(Tensor(
+                    shape=kv_shape, dtype=self.dtype, init=Zero())) for _ in range(self.num_layers)]
+                v_caches = [Parameter(Tensor(
+                    shape=kv_shape, dtype=self.dtype, init=Zero())) for _ in range(self.num_layers)]
+
+        self.atb_encoder_operation.set_kvcache(k_caches, v_caches)
+        self.atb_decoder_operation.set_kvcache(k_caches, v_caches)
+
+    def add_flags(self, is_first_iteration):
+        """add_flags."""
+        self.is_first_iteration = is_first_iteration
+
+    def _execute_operator(self, acl_inputs, acl_param):
+        """execute operator."""
+        if self.is_first_iteration:
+            acl_model_out = self.atb_encoder_operation.forward(
+                acl_inputs, acl_param)
+        else:
+            acl_model_out = self.atb_decoder_operation.forward(
+                acl_inputs, acl_param)
+        acl_hidden_state = acl_model_out[0]
+        return acl_hidden_state
+
+    def forward(self, boost_inputs):
+        r"""
+        LlmBoost forward.
+        """
+        input_ids = boost_inputs["input_ids"]
+        position_ids = boost_inputs["position_ids"]
+        cos_embed = boost_inputs["cos_embed"]
+        sin_embed = boost_inputs["sin_embed"]
+        block_tables = boost_inputs["block_tables"]
+        slot_mapping = boost_inputs["slot_mapping"]
+        batch_valid_length = boost_inputs["batch_valid_length"]
+        lm_head_indices = boost_inputs["lm_head_indices"]
+        seqLen = boost_inputs["seq_lens"]
+        if self.is_first_iteration:
+            attention_mask = self.attn_mask
+        else:
+            position_ids = batch_valid_length - 1
+            attention_mask = self.placeholder
+            lm_head_indices = self.lm_head_indices_fake
+
+        acl_inputs, acl_param = self._prepare_inputs(prefill=self.is_first_iteration, input_ids=input_ids,
+                                                     position_ids=position_ids, cos_embed=cos_embed,
+                                                     sin_embed=sin_embed, attention_mask=attention_mask,
+                                                     block_tables=block_tables, slots=slot_mapping,
+                                                     input_lengths=batch_valid_length, lm_head_indices=lm_head_indices,
+                                                     seqLen=seqLen)
+        ms.hal.synchronize()
+        logits = self._execute_operator(acl_inputs, acl_param)
+        logits = self.cast(logits, mstype.float32)
+        return logits