mindspore 2.0.0rc1__cp38-cp38-manylinux1_x86_64.whl → 2.2.0__cp38-cp38-manylinux1_x86_64.whl

mindspore/nn/layer/dense.py

@@ -41,8 +41,8 @@ def check_dense_inputs_same_shape(input1, input2, prim_name=None):
 @constexpr(check=False)
 def _check_is_tensor(param_name, input_data, cls_name):
     """Internal function, used to check whether the input data is Tensor."""
-    if input_data is not None and not isinstance(P.typeof(input_data), mstype.tensor_type):
-        raise TypeError(f"For '{cls_name}', the '{param_name}' must be '{mstype.tensor_type}', "
+    if input_data is not None and not isinstance(P.typeof(input_data), mstype.TensorType):
+        raise TypeError(f"For '{cls_name}', the '{param_name}' must be '{mstype.TensorType}', "
                         f"but got '{P.typeof(input_data)}'")
 
 
@@ -66,17 +66,18 @@ class BiDense(Cell):
     where :math:`x_{1}` is the first input tensor, :math:`x_{2}` is the second input tensor
     , :math:`A` is a weight matrix with the same data type as the :math:`x_{*}` created by the layer
     , and :math:`b` is a bias vector with the same data type as the :math:`x_{*}` created by the layer
-    (only if has_bias is True).
+    (only if has_bias is ``True`` ).
 
     Args:
         in1_channels (int): The number of channels in the input1 space.
         in2_channels (int): The number of channels in the input2 space.
         out_channels (int): The number of channels in the output space.
         weight_init (Union[Tensor, str, Initializer, numbers.Number]): The trainable weight_init parameter.
-            The values of str refer to the function `initializer`. Default: None.
+            The values of str refer to the function `initializer`. Default: ``None`` .
         bias_init (Union[Tensor, str, Initializer, numbers.Number]): The trainable bias_init parameter.
-            The values of str refer to the function `initializer`. Default: None.
-        has_bias (bool): Specifies whether the layer uses :math:`\text{bias}` vector. Default: True.
+            The values of str refer to the function `initializer`. Default: ``None`` .
+        has_bias (bool): Specifies whether the layer uses :math:`\text{bias}` vector. Default: ``True`` .
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Shape:
         - **input1** - :math:`(*, H_{in1})` where :math:`H_{in1}=\text{in1_channels}` and
@@ -90,17 +91,17 @@ class BiDense(Cell):
           are the same shape as the inputs.
 
     Dtype:
-        - **input1** (Tensor) - The dtype must be float16 or float32 and be same as **input2**.
-        - **input1** (Tensor) - The dtype must be float16 or float32 and be same as **input1**.
+        - **input1** (Tensor) - The dtype must be float16 or float32 and be same as **input2** .
+        - **input2** (Tensor) - The dtype must be float16 or float32 and be same as **input1** .
         - **output** (Tensor) - With the same dtype as the inputs.
 
     Weights:
         - **weight** (Parameter) - The learnable weights with shape
           :math:`(\text{out_channels}, \text{in1_channels}, \text{in2_channels})`.
-          When `weight_init` is `None`, the values are initialized from
+          When `weight_init` is ``None`` , the values are initialized from
           :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where :math:`k = \frac{1}{\text{in1_channels}}`.
         - **bias** (Parameter) - The learnable bias of shape :math:`(\text{out_channels})`.
-          If `has_bias` is `True` and `bias_init` is `None`, the values are initialized from
+          If `has_bias` is ``True`` and `bias_init` is ``None`` , the values are initialized from
           :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where :math:`k = \frac{1}{\text{in1_channels}}`.
 
     Raises:
@@ -116,6 +117,9 @@ class BiDense(Cell):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, nn
+        >>> import numpy as np
         >>> x1 = Tensor(np.random.randn(128, 20), mindspore.float32)
         >>> x2 = Tensor(np.random.randn(128, 30), mindspore.float32)
         >>> net = nn.BiDense(20, 30, 40)
@@ -130,7 +134,8 @@ class BiDense(Cell):
                  out_channels,
                  weight_init=None,
                  bias_init=None,
-                 has_bias=True):
+                 has_bias=True,
+                 dtype=mstype.float32):
         super().__init__()
         self.in_channels = Validator.check_positive_int(in1_channels, "in1_channels", self.cls_name)
         self.in_channels = Validator.check_positive_int(in2_channels, "in2_channels", self.cls_name)
@@ -153,7 +158,8 @@ class BiDense(Cell):
                                  f"equal to 'in2_channels'. But got 'weight_init': {weight_init}, "
                                  f"'out_channels': {out_channels}, 'in_channels': {in1_channels}, "
                                  f"'in2_channels': {in2_channels}")
-        self.weight = Parameter(initializer(weight_init, (out_channels, in1_channels, in2_channels)), 'weight')
+        self.weight = Parameter(initializer(weight_init, (out_channels, in1_channels, in2_channels), dtype=dtype),
+                                'weight')
 
         if self.has_bias:
             if bias_init is None:
@@ -163,7 +169,7 @@ class BiDense(Cell):
                     raise ValueError(f"For '{self.cls_name}', bias init shape error. The ndim of 'bias_init' should "
                                      f"be equal to 1, and the first dim must be equal to 'out_channels'. But got "
                                      f"'bias_init': {bias_init}, 'out_channels': {out_channels}.")
-            self.bias = Parameter(initializer(bias_init, [out_channels]), name="bias")
+            self.bias = Parameter(initializer(bias_init, [out_channels], dtype=dtype), name="bias")
             self.bias_add = P.BiasAdd()
         self.matmul = P.MatMul()
 

mindspore/nn/layer/embedding.py

@@ -62,13 +62,15 @@ class Embedding(Cell):
     Args:
         vocab_size (int): Size of the dictionary of embeddings.
         embedding_size (int): The size of each embedding vector.
-        use_one_hot (bool): Specifies whether to apply one_hot encoding form. Default: False.
+        use_one_hot (bool): Specifies whether to apply one_hot encoding form. Default: ``False`` .
         embedding_table (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the embedding_table.
-            Refer to class `initializer` for the values of string when a string
-            is specified. Default: 'normal'.
-        dtype (:class:`mindspore.dtype`): Data type of `x`. Default: mindspore.float32.
+            Refer to class `mindspore.common.initializer
+            <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.common.initializer.html>`_
+            for the values of string when a string is specified. Default: ``'normal'`` .
+        dtype (:class:`mindspore.dtype`): Data type of `x`. Default: ``mstype.float32`` .
         padding_idx (int, None): When the padding_idx encounters index, the output embedding vector of this index
-                                 will be initialized to zero. Default: None. The feature is inactivated.
+                                 will be initialized to zero. Default: ``None`` . The feature is inactivated.
+
     Inputs:
         - **x** (Tensor) - Tensor of shape :math:`(\text{batch_size}, \text{x_length})`. The elements of
           the Tensor must be integer and not larger than vocab_size. Otherwise the corresponding embedding vector will
@@ -86,6 +88,9 @@ class Embedding(Cell):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, nn
+        >>> import numpy as np
         >>> net = nn.Embedding(20000, 768,  True)
         >>> x = Tensor(np.ones([8, 128]), mindspore.int32)
         >>> # Maps the input word IDs to word embedding.
@@ -126,13 +131,10 @@ class Embedding(Cell):
         self.array_mul = P.MatMul()
         self.reshape = P.Reshape()
         self.get_shp = P.Shape()
-        self.get_tensor_shp = P.TensorShape()
         self.concat = P.Concat()
 
     def construct(self, ids):
         out_shape = self.get_shp(ids) + (self.embedding_size,)
-        if F.is_sequence_value_unknown(self.get_shp(ids)):
-            out_shape = self.concat((self.get_tensor_shp(ids), Tensor([self.embedding_size])))
         flat_ids = self.reshape_flat(ids, self.shp_flat)
 
         if self.use_one_hot:
@@ -145,12 +147,11 @@ class Embedding(Cell):
         return output
 
     def extend_repr(self):
-        s = 'vocab_size={}, embedding_size={}, use_one_hot={}, embedding_table={}, dtype={}, padding_idx={}'.format(
-            self.vocab_size, self.embedding_size, self.use_one_hot, self.embedding_table, self.dtype, self.padding_idx)
-        return s
+        return f'vocab_size={self.vocab_size}, embedding_size={self.embedding_size}, use_one_hot={self.use_one_hot}, ' \
+            f'embedding_table={self.embedding_table}, dtype={self.dtype}, padding_idx={self.padding_idx}'
 
 
-@constexpr
+@_primexpr
 def _make_axis_range(start, end):
     axis = tuple(range(start, end))
     return axis
@@ -177,19 +178,20 @@ class EmbeddingLookup(Cell):
         embedding_size (int): The size of each embedding vector.
         param_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the embedding_table.
             Refer to class `initializer` for the values of string when a string
-            is specified. Default: 'normal'.
+            is specified. Default: ``'normal'`` .
         target (str): Specifies the target where the op is executed. The value must in
-            ['DEVICE', 'CPU']. Default: 'CPU'.
+            [ ``'DEVICE'`` , ``'CPU'`` ]. Default: ``'CPU'`` .
         slice_mode (str): The slicing way in semi_auto_parallel/auto_parallel. The value must get through
-            :class:`mindspore.nn.EmbeddingLookup`. Default: 'nn.EmbeddingLookup.BATCH_SLICE'.
-        manual_shapes (tuple): The accompaniment array in field slice mode. Default: None.
+            :class:`mindspore.nn.EmbeddingLookup`. Default: ``'batch_slice'`` .
+        manual_shapes (tuple): The accompaniment array in field slice mode. Default: ``None`` .
         max_norm (Union[float, None]): A maximum clipping value. The data type must be float16, float32
-                                       or None. Default: None
-        sparse (bool): Using sparse mode. When 'target' is set to 'CPU', 'sparse' has to be true. Default: True.
-        vocab_cache_size (int): Cache size of the dictionary of embeddings. Default: 0. It is valid only in
+                                       or None. Default: ``None`` .
+        sparse (bool): Using sparse mode. When 'target' is set to 'CPU', 'sparse' has to be true. Default: ``True`` .
+        vocab_cache_size (int): Cache size of the dictionary of embeddings. Default: ``0`` . It is valid only in
             parameter server trainning mode and 'DEVICE' target. And the moment parameter of corresponding
             optimizer will also be set to the cache size. In addition, it should be noted that it will cost the 'DEVICE'
             memory, so suggests setting a reasonable value to avoid insufficient memory.
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Inputs:
         - **input_indices** (Tensor) - The shape of tensor is :math:`(y_1, y_2, ..., y_S)`.
@@ -216,6 +218,9 @@ class EmbeddingLookup(Cell):
         ``Ascend`` ``GPU`` ``CPU``
 
     Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, nn
+        >>> import numpy as np
         >>> input_indices = Tensor(np.array([[1, 0], [3, 2]]), mindspore.int32)
         >>> result = nn.EmbeddingLookup(4,2)(input_indices)
         >>> print(result.shape)
@@ -228,7 +233,7 @@ class EmbeddingLookup(Cell):
 
     def __init__(self, vocab_size, embedding_size, param_init='normal',
                  target='CPU', slice_mode='batch_slice', manual_shapes=None,
-                 max_norm=None, sparse=True, vocab_cache_size=0):
+                 max_norm=None, sparse=True, vocab_cache_size=0, dtype=mstype.float32):
         """Initialize EmbeddingLookup."""
         super(EmbeddingLookup, self).__init__()
         Validator.check_value_type('sparse', sparse, [bool], self.cls_name)
@@ -252,8 +257,8 @@ class EmbeddingLookup(Cell):
         if enable_ps:
             self._process_vocab_cache(slice_mode)
         self.embedding_size = Validator.check_positive_int(embedding_size, 'embedding_size', self.cls_name)
-        self.embedding_table = Parameter(initializer(param_init, [self.vocab_size, self.embedding_size]),
-                                         name='embedding_table')
+        self.embedding_table = Parameter(initializer(param_init, [self.vocab_size, self.embedding_size],
+                                                     dtype=dtype), name='embedding_table')
         parallel_mode = _get_parallel_mode()
         is_auto_parallel = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
         self.gather_revert = P.Gather()
@@ -264,7 +269,7 @@ class EmbeddingLookup(Cell):
         if is_auto_parallel:
             self.unique = P.Unique().shard(((1,),))
         if self.cache_enable and enable_ps:
-            self._set_voacb_cache_enable_for_ps(vocab_cache_size, embedding_size, vocab_size, param_init)
+            self._set_voacb_cache_enable_for_ps(vocab_cache_size, embedding_size, vocab_size, param_init, dtype=dtype)
             if is_auto_parallel:
                 self.unique.add_prim_attr('cache_enable', True)
         indices_shape_size = 2
@@ -300,15 +305,15 @@ class EmbeddingLookup(Cell):
             self.embeddinglookup.shard(((1, get_group_size()), indices_strategy))
         elif slice_mode == "batch_slice" and is_auto_parallel:
             indices_strategy = [get_group_size()]
-            indices_strategy.extend([1]*(indices_shape_size - 1))
+            indices_strategy.extend([1] * (indices_shape_size - 1))
             indices_strategy = tuple(indices_strategy)
             self.gatherv2.shard(((1, 1), indices_strategy))
             self.embeddinglookup.shard(((1, 1), indices_strategy))
         else:
             if is_auto_parallel:
                 support_mode = ["field_slice", "table_row_slice", "table_column_slice", "batch_slice"]
-                raise ValueError("For '{}', the 'slice_mode' must be in {}, "
-                                 "but got \"{}\".".format(self.cls_name, support_mode, slice_mode))
+                raise ValueError(f"For '{self.cls_name}', the 'slice_mode' must be in {support_mode}, "
+                                 f"but got \"{slice_mode}\".")
         if self.cache_enable and not enable_ps:
             raise ValueError(f"For '{self.cls_name}', haven't supported cache enable for not ps mode.")
         self.embedding_table.unique = self.forward_unique
@@ -351,7 +356,8 @@ class EmbeddingLookup(Cell):
             if _is_role_worker():
                 self.vocab_size = self.vocab_cache_size
 
-    def _set_voacb_cache_enable_for_ps(self, vocab_cache_size, embedding_size, vocab_size, param_init):
+    def _set_voacb_cache_enable_for_ps(self, vocab_cache_size, embedding_size, vocab_size, param_init,
+                                       dtype=mstype.float32):
         """PS embeddingLookup cache enable set."""
         if self.sparse:
             self.forward_unique = True
@@ -365,10 +371,10 @@ class EmbeddingLookup(Cell):
         if _enable_distributed_mindrt():
             self.rank_id = get_rank()
             if self.is_ps_server:
-                self._slice_pserver_embeddings("zeros")
+                self._slice_pserver_embeddings("zeros", dtype=dtype)
                 self._set_cache_enable_and_key_for_pserver(param_key)
 
-    def _slice_pserver_embeddings(self, param_init):
+    def _slice_pserver_embeddings(self, param_init, dtype=mstype.float32):
         '''
         Method to slice embedding tables on Parameter Servers.
         It helps to train with a large scale embedding table and is used only in Parameter Server training mode.
@@ -396,7 +402,7 @@ class EmbeddingLookup(Cell):
         for i in range(server_num):
             self.embedding_table_list.append(Parameter(initializer(param_init,
                                                                    [self.embedding_table_vocab_dim_list[i],
-                                                                    self.embedding_size]),
+                                                                    self.embedding_size], dtype=dtype),
                                                        name="embedding_table_server_" + str(i)))
 
             self.embedding_offset.append(offset)
@@ -495,17 +501,20 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
         field_size (int): The field size of the final outputs.
         param_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the embedding_table.
             Refer to class `initializer` for the values of string when a string
-            is specified. Default: 'normal'.
+            is specified. Default: ``'normal'`` .
         target (str): Specifies the target where the op is executed. The value must in
-            ['DEVICE', 'CPU']. Default: 'CPU'.
+            [ ``'DEVICE'`` , ``'CPU'`` ]. Default: ``'CPU'`` .
         slice_mode (str): The slicing way in semi_auto_parallel/auto_parallel. The value must get through
-            :class:`mindspore.nn.EmbeddingLookup`. Default: 'nn.EmbeddingLookup.BATCH_SLICE'.
-        feature_num_list (tuple): The accompaniment array in field slice mode. This is unused currently. Default: None.
-        max_norm (Union[float, None]): A maximum clipping value. The data type must be float16, float32
-                                       or None. Default: None
-        sparse (bool): Using sparse mode. When 'target' is set to 'CPU', 'sparse' has to be true. Default: True.
-        operator (str): The pooling method for the features in one field. Support 'SUM', 'MEAN' and 'MAX'.
-            Default: 'SUM'.
+            :class:`mindspore.nn.EmbeddingLookup`. Default: ``'batch_slice'``.
+        feature_num_list (tuple): The accompaniment array in field slice mode. This is unused currently.
+            Default:  ``None`` .
+        max_norm (Union[float, None]): A maximum clipping value. The data type must be float16, float32.
+            Default: ``None`` .
+        sparse (bool): Using sparse mode. When 'target' is set to ``'CPU'`` , 'sparse' has to be true.
+            Default: ``True`` .
+        operator (str): The pooling method for the features in one field. Support ``'SUM'`` , ``'MEAN'`` and
+            ``'MAX'`` . Default: ``'SUM'`` .
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
 
     Inputs:
         - **input_indices** (Tensor) - The shape of tensor is :math:`(batch\_size, seq\_length)`.
@@ -524,17 +533,19 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
         TypeError: If `vocab_size` or `embedding_size` or `field_size` is not an int.
         TypeError: If `sparse` is not a bool or `feature_num_list` is not a tuple.
         ValueError: If `vocab_size` or `embedding_size` or `field_size` is less than 1.
-        ValueError: If `target` is neither 'CPU' nor 'DEVICE'.
-        ValueError: If `slice_mode` is not one of 'batch_slice', 'field_slice', 'table_row_slice',
-                    'table_column_slice'.
-        ValueError: If `sparse` is False and `target` is 'CPU'.
-        ValueError: If `slice_mode` is 'field_slice' and `feature_num_list` is None.
-        ValueError: If `operator` is not one of 'SUM', 'MAX', 'MEAN'.
+        ValueError: If `target` is neither ``'CPU'`` nor ``'DEVICE'``.
+        ValueError: If `slice_mode` is not one of ``'batch_slice'``, ``'field_slice'``, ``'table_row_slice'``,
+                    ``'table_column_slice'`` .
+        ValueError: If `sparse` is False and `target` is ``'CPU'`` .
+        ValueError: If `slice_mode` is ``'field_slice'`` and `feature_num_list` is None.
+        ValueError: If `operator` is not one of ``'SUM'``, ``'MAX'``, ``'MEAN'`` .
 
     Supported Platforms:
         ``Ascend`` ``GPU``
 
     Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, nn
         >>> input_indices = Tensor([[2, 4, 6, 0, 0], [1, 3, 5, 0, 0]], mindspore.int32)
         >>> input_values = Tensor([[1, 1, 1, 0, 0], [1, 1, 1, 0, 0]], mindspore.float32)
         >>> field_ids = Tensor([[0, 1, 1, 0, 0], [0, 0, 1, 0, 0]], mindspore.int32)
@@ -548,10 +559,11 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
     OPERATOR_MAX = 'MAX'
 
     def __init__(self, vocab_size, embedding_size, field_size, param_init='normal', target='CPU',
-                 slice_mode='batch_slice', feature_num_list=None, max_norm=None, sparse=True, operator='SUM'):
+                 slice_mode='batch_slice', feature_num_list=None, max_norm=None, sparse=True, operator='SUM',
+                 dtype=mstype.float32):
         """Initialize MultiFieldEmbeddingLookup."""
         super(MultiFieldEmbeddingLookup, self).__init__(vocab_size, embedding_size, param_init, target,
-                                                        slice_mode, feature_num_list, max_norm, sparse)
+                                                        slice_mode, feature_num_list, max_norm, sparse, dtype=dtype)
         self.field_size = Validator.check_positive_int(field_size, 'field_size', self.cls_name)
         self.operator = operator
 
@@ -615,8 +627,9 @@ class MultiFieldEmbeddingLookup(EmbeddingLookup):
                 self.inf_add.shard(((1, 1, get_group_size()), (1, 1, 1)))
         else:
             if is_auto_parallel:
-                raise ValueError("For '{}', the 'slice_mode' must be in ['table_row_slice', 'batch_slice' and \
-                                       'table_column_slice'], but got {}".format(self.cls_name, str(slice_mode)))
+                raise ValueError(
+                    f"For '{self.cls_name}', the 'slice_mode' must be in ['table_row_slice', 'batch_slice' "
+                    f"and 'table_column_slice'], but got {str(slice_mode)}.")
 
         # Min value for fp32
         self.negative_inf_value = -3.402823466E+38

mindspore/nn/layer/flash_attention.py

@@ -0,0 +1,264 @@
+# 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.
+# ============================================================================
+"""
+A FlashAttention Layer.
+"""
+import math
+
+import mindspore.common.dtype as mstype
+from mindspore.common.tensor import Tensor
+from mindspore import ops
+from mindspore.nn.cell import Cell
+from mindspore.ops._op_impl._custom_op.flash_attention.flash_attention_impl import get_flash_attention
+from mindspore.ops.operations.nn_ops import FlashAttentionScore
+from mindspore._c_expression import MSContext
+
+__all__ = ['FlashAttention']
+
+
+class FlashAttention(Cell):
+    """Flash Attention Layer.
+
+    This function contains the flash attention primitives used in FlashAttention (see paper)
+    `FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness <https://arxiv.org/pdf/2205.14135.pdf>`
+
+    Specifically, it includes the following:
+
+    1. An interface for calling flashattention operation.
+    2. Two configuration parameters for enabling local block sparse of flashattention.
+
+    Args:
+        head_dim(int): The hidden size of input.
+        dropout_rate(float): The dropout rate of the attention score. Default 0.0.
+        prev_block_num(int): A integer to define the number of blocks to look ahead for local block sparse attention.
+            Default 65536.
+        next_block_num(int): A integer to define the number of blocks to look behind for local block sparse attention.
+            Default 65536.
+        tiling_stgy_name(str): A str to define tiling strategy of flash attention.
+        dp(int): data parallel.
+            Default 1.
+        mp(int): model parallel.
+            Default 1.
+        high_precision(bool): This mode has higher precision but some performance loss.
+            Default False.
+        have_attention_mask_batch(bool): indicates whether attention_mask contains the batch dimension.
+            Default True
+        alibi(bool): This parameter indicates whether the flashattention supports the Alibi.
+            Default: False
+
+
+    Inputs:
+      - **query** (Tensor) - Tensor query (:class:`mstype.fp16` [batch_size, head_num, seq_length, head_dim])
+      - **key** (Tensor) - Tensor key (:class:`mstype.fp16` [batch_size, head_num, seq_length, head_dim])
+      - **value** (Tensor) - Tensor value (:class:`mstype.fp16` [batch_size, head_num, seq_length, head_dim])
+      - **attention_mask** (Tensor) - Float Tensor the mask of (:class:`mstype.fp16` [batch_size, seq_length,
+          seq_length]): A matrix to pass masked information.
+
+    Outputs:
+        A Tensor. The output of the attention with shape [batch_size, head_num, seq_length, head_dim]
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import numpy as np
+        >>> from mindspore import dtype as mstype
+        >>> from mindspore.nn.layer.flash_attention import FlashAttention
+        >>> from mindspore import Tensor
+        >>> model = FlashAttention(head_dim=128,
+        ...                        dropout_rate=0.1,
+        ...                        prev_block_num=7,
+        ...                        next_block_num=0
+        ...                        )
+        >>> query = Tensor(np.ones((2, 16, 4096, 128)), mstype.float16)
+        >>> key = Tensor(np.ones((2, 16, 4096, 128)), mstype.float16)
+        >>> value = Tensor(np.ones((2, 16, 4096, 128)), mstype.float16)
+        >>> attention_mask = Tensor(np.ones((2, 4096, 4096)), mstype.float16)
+        >>> output = model(query, key, value, attention_mask)
+        >>> print(output.shape)
+        (2, 16, 4096, 128)
+    """
+
+    def __init__(self,
+                 head_dim,
+                 head_num,
+                 dropout_rate=0.0,
+                 prev_block_num=65536,
+                 next_block_num=65536,
+                 tiling_stgy_name="sparse",
+                 dp=1,
+                 mp=1,
+                 high_precision=False,
+                 have_attention_mask_batch=True,
+                 alibi=False
+                 ):
+        super(FlashAttention, self).__init__()
+
+        scaling_constant = math.sqrt(head_dim)
+        if scaling_constant == 0:
+            raise ValueError("the scaling constant must not be 0.")
+        self.scale_factor = Tensor([1. / scaling_constant], dtype=mstype.float16)
+
+        self.is_910A = MSContext.get_instance().get_ascend_soc_version() == "Ascend910"
+        if self.is_910A:
+            self.flash_attention = get_flash_attention(
+                prev_block_num=prev_block_num,
+                next_block_num=next_block_num,
+                tiling_stgy_name=tiling_stgy_name,
+                high_precision=high_precision
+            )
+            self.flash_attention.add_prim_attr("primitive_target", "Ascend")
+        else:
+            if alibi:
+                raise ValueError(f"When soc_version is not Ascend910A, alibi must be False")
+            self.transpose_4d_pre = ops.Transpose().shard(((dp, mp, 1, 1),))
+            self.transpose_4d_post = ops.Transpose().shard(((dp, 1, mp, 1),))
+            self.reshape = ops.Reshape()
+            self.zeros_like = ops.ZerosLike().shard(((dp, mp, 1, 1),))
+            self.zeros = ops.Zeros()
+            self.attn_expand_dims = ops.ExpandDims().shard(((dp, 1, 1),))
+            fa_strategies = ((dp, 1, mp),
+                             (dp, 1, mp),
+                             (dp, 1, mp),
+                             (dp, 1, 1, 1))
+            if dropout_rate > 1e-5:
+                fa_strategies += ((dp, mp, 1, 1),)
+            self.flash_attention = FlashAttentionScore(head_num=head_num, pre_tokens=prev_block_num,
+                                                       next_tokens=next_block_num,
+                                                       keep_prob=1 - dropout_rate,
+                                                       scale_value=1.0,
+                                                       inner_precise=0 if high_precision else 1).shard(fa_strategies)
+
+        self.ones = ops.Ones()
+        self.dim_mask = Tensor([1 for _ in range(head_dim)], dtype=mstype.int8)
+        self.scale_mul = ops.Mul().shard(((dp, mp, 1, 1), (1,)))
+        self.dropout_rate = dropout_rate
+        self.have_attention_mask_batch = have_attention_mask_batch
+        self.alibi = alibi
+        if self.dropout_rate > 1e-5:
+            self.keep_prob = Tensor(1 - self.dropout_rate, dtype=mstype.float16)
+            self.fill_v2 = ops.FillV2().shard(((dp, mp, 1, 1), ()))
+            self.tensor_one = Tensor(1.0, mstype.float16)
+            self.drop_gen_mask = ops.DropoutGenMask()
+            self.do_dropout = ops.DropoutDoMask().shard(((dp, mp, 1, 1),))
+            self.depend = ops.Depend()
+
+    def shard(self, in_strategy=None, out_strategy=None):
+        """Distributed configuration of FlashAttention
+        :param in_strategy: Describe the split strategy of operator input. Default: None.
+        :param out_strategy: Describe the split strategy of operator output, it is only for certain operators,
+                                  such as MatMul. Default: None.
+        :return:
+        """
+        if in_strategy is None:
+            # default: dp=1, mp=1, construct inputs only contain query, key, value
+            in_strategy = (
+                (1, 1, 1, 1),
+                (1, 1, 1, 1),
+                (1, 1, 1, 1),
+            )
+        self.flash_attention.shard(in_strategy)
+        dp = in_strategy[0][0]
+        mp = in_strategy[0][1]
+        self.flash_attention.add_prim_attr("dev_matrix_shape", [dp, mp, 1, 1])
+        inputs_tensor_map = [
+            [3, 2, 1, 0],
+            [3, 2, 1, 0],
+            [3, 2, 1, 0],
+        ]
+        if self.have_attention_mask_batch:
+            inputs_tensor_map.append([3, 1, 0])
+        else:
+            inputs_tensor_map.append([-1, 1, 0])
+
+        input_empty_args_num = 2
+        # dropout_mask
+        if self.dropout_rate > 1e-5:
+            input_empty_args_num -= 1
+            inputs_tensor_map.append([3, 2, 1, 0])
+
+        if self.alibi:
+            input_empty_args_num -= 1
+            inputs_tensor_map.append([3, 2, 1, 0])
+
+        self.flash_attention.add_prim_attr("inputs_tensor_map", inputs_tensor_map)
+
+        self.flash_attention.add_prim_attr("outputs_tensor_map", [
+            [3, 2, 1, 0],  # O
+            [3, 2, 1],  # L
+            [3, 2, 1]  # M
+        ])
+        self.flash_attention.add_prim_attr("as_loss_divisor", 0)
+        self.flash_attention.add_prim_attr("empty_mirror_ops", input_empty_args_num)
+
+    def construct(self, query, key, value, attn_mask=None, alibi_mask=None):
+        """FlashAttention forward
+        :param query:           [bsz, head_num, seq_len, head_dim]
+        :param key:           [bsz, head_num, seq_len, head_dim]
+        :param value:           [bsz, head_num, seq_len, head_dim]
+        :param attn_mask:   [1 or bsz, seq_len, seq_len], if not None
+        :param alibi_mask: [bsz, head_num, 1, seq_len], if not None
+        :return: output          [bsz, head_num, seq_len, head_dim]
+        """
+        query = self.scale_mul(query, self.scale_factor)
+        bsz, head_num, seq_len, head_dim = query.shape
+        _, k_head_num, k_seq_len, _ = key.shape
+        _, v_head_num, v_seq_len, _ = value.shape
+        if head_num != k_head_num or head_num != v_head_num:
+            raise ValueError(
+                "the head_num of query, key and value must be the same, "
+                "If different head_num are used, users need to change themselves to be same by tile.")
+        if seq_len % 16 != 0 or k_seq_len % 16 != 0 or k_seq_len != v_seq_len:
+            raise ValueError(
+                "query, key, value seq_len must be a multiple of 16, and key seq_len, value seq_len must be the same.")
+
+        if head_dim > 304:
+            raise ValueError(
+                "the head_dim must be less than 304, otherwise the ub would be OOM.")
+
+        if self.is_910A:
+            # 910A -- FlashAttentionPrimtive
+            if self.dropout_rate > 1e-5:
+                drop_mask_bits = self.drop_gen_mask((bsz, head_num, seq_len, seq_len), self.keep_prob)
+                tensor_shape = Tensor((bsz, head_num, seq_len, seq_len), mstype.int32)
+                ones = self.fill_v2(tensor_shape, self.tensor_one)
+                ones = self.depend(ones, query)
+                drop_mask = self.do_dropout(ones, drop_mask_bits, self.keep_prob)
+            else:
+                drop_mask = None
+            output, _, _ = self.flash_attention(query, key, value, attn_mask, drop_mask, alibi_mask)
+        else:
+            # FlashAttentionScore
+            # Useless input, just for binary calls.
+            if self.dropout_rate > 1e-5:
+                drop_mask_bits = self.reshape(self.drop_gen_mask((bsz, head_num, seq_len, seq_len), self.keep_prob),
+                                              (bsz, head_num, seq_len, seq_len // 8))
+            else:
+                drop_mask_bits = None
+            # (B, N, S, D) -> (B, S, H)
+            query = self.reshape(self.transpose_4d_pre(query, (0, 2, 1, 3)), (bsz, seq_len, -1))
+            key = self.reshape(self.transpose_4d_pre(key, (0, 2, 1, 3)), (bsz, seq_len, -1))
+            value = self.reshape(self.transpose_4d_pre(value, (0, 2, 1, 3)), (bsz, seq_len, -1))
+            attn_mask = self.attn_expand_dims(attn_mask, 1)
+            output, _, _ = self.flash_attention(query,
+                                                key,
+                                                value,
+                                                attn_mask,
+                                                drop_mask_bits,
+                                                None,
+                                                None)
+            output = self.transpose_4d_post(self.reshape(output, (bsz, seq_len, head_num, head_dim)), (0, 2, 1, 3))
+
+        return output